From 01e30ab09b71b98a202ace56d0a6f94d23195f52 Mon Sep 17 00:00:00 2001
From: Jamie Liu <jamieliu@google.com>
Date: Tue, 10 Dec 2024 11:16:48 -0800
Subject: [PATCH] nvproxy: don't depend on HVASpace for non-virtmem
 NVOS32_FUNCTION_ALLOC_SIZE

PiperOrigin-RevId: 704780681
---
 pkg/abi/nvgpu/classes.go                      |  1 +
 pkg/abi/nvgpu/frontend.go                     | 20 +++++++++++++++++++
 pkg/sentry/devices/nvproxy/frontend_unsafe.go | 15 ++++++++++++--
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/pkg/abi/nvgpu/classes.go b/pkg/abi/nvgpu/classes.go
index 6c2b314d7a..26fda37a75 100644
--- a/pkg/abi/nvgpu/classes.go
+++ b/pkg/abi/nvgpu/classes.go
@@ -38,6 +38,7 @@ const (
 	NV01_MEMORY_SYSTEM               = 0x0000003e
 	NV01_MEMORY_LOCAL_USER           = 0x00000040
 	NV01_ROOT_CLIENT                 = 0x00000041
+	NV_MEMORY_EXTENDED_USER          = 0x00000042
 	NV01_MEMORY_SYSTEM_OS_DESCRIPTOR = 0x00000071
 	NV01_EVENT_OS_EVENT              = 0x00000079
 	NV01_DEVICE_0                    = 0x00000080
diff --git a/pkg/abi/nvgpu/frontend.go b/pkg/abi/nvgpu/frontend.go
index 43787aa45b..4f0164ffac 100644
--- a/pkg/abi/nvgpu/frontend.go
+++ b/pkg/abi/nvgpu/frontend.go
@@ -312,6 +312,26 @@ type NVOS32AllocSize struct {
 	CtagOffset      uint32
 }
 
+// Flags in NVOS32AllocSize.Flags:
+const (
+	NVOS32_ALLOC_FLAGS_VIRTUAL = 0x00080000
+)
+
+// Bitfields in NVOS32AllocSize.Attr:
+const (
+	NVOS32_ATTR_LOCATION_SHIFT  = 25
+	NVOS32_ATTR_LOCATION_MASK   = 0x3
+	NVOS32_ATTR_LOCATION_VIDMEM = 0
+)
+
+// Bitfields in NVOS32AllocSize.Attr2:
+const (
+	NVOS32_ATTR2_USE_EGM_SHIFT = 24
+	NVOS32_ATTR2_USE_EGM_MASK  = 0x1
+	NVOS32_ATTR2_USE_EGM_FALSE = 0
+	NVOS32_ATTR2_USE_EGM_TRUE  = 1
+)
+
 // IoctlNVOS33ParametersWithFD is the parameter type for NV_ESC_RM_MAP_MEMORY,
 // from src/nvidia/arch/nvalloc/unix/include/nv-unix-nvos-params-wrappers.h.
 //
diff --git a/pkg/sentry/devices/nvproxy/frontend_unsafe.go b/pkg/sentry/devices/nvproxy/frontend_unsafe.go
index 63d48d4735..f7957b0a49 100644
--- a/pkg/sentry/devices/nvproxy/frontend_unsafe.go
+++ b/pkg/sentry/devices/nvproxy/frontend_unsafe.go
@@ -325,8 +325,19 @@ func rmVidHeapControlAllocSize(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS32
 	fi.fd.dev.nvp.objsLock()
 	n, err := frontendIoctlInvoke(fi, ioctlParams)
 	if err == nil && ioctlParams.Status == nvgpu.NV_OK {
-		// src/nvidia/src/kernel/mem_mgr/virtual_mem.c:virtmemConstruct_IMPL() => refAddDependant()
-		fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, allocSizeParams.HMemory, nvgpu.NV50_MEMORY_VIRTUAL, &miscObject{}, ioctlParams.HObjectParent, ioctlParams.HVASpace)
+		// src/nvidia/interface/deprecated/rmapi_deprecated_vidheapctrl.c:_rmVidHeapControlAllocCommon()
+		if allocSizeParams.Flags&nvgpu.NVOS32_ALLOC_FLAGS_VIRTUAL != 0 {
+			// src/nvidia/src/kernel/mem_mgr/virtual_mem.c:virtmemConstruct_IMPL() => refAddDependant()
+			fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, allocSizeParams.HMemory, nvgpu.NV50_MEMORY_VIRTUAL, &miscObject{}, ioctlParams.HObjectParent, ioctlParams.HVASpace)
+		} else {
+			classID := nvgpu.ClassID(nvgpu.NV01_MEMORY_SYSTEM)
+			if (allocSizeParams.Attr2>>nvgpu.NVOS32_ATTR2_USE_EGM_SHIFT)&nvgpu.NVOS32_ATTR2_USE_EGM_MASK == nvgpu.NVOS32_ATTR2_USE_EGM_TRUE {
+				classID = nvgpu.NV_MEMORY_EXTENDED_USER
+			} else if (allocSizeParams.Attr>>nvgpu.NVOS32_ATTR_LOCATION_SHIFT)&nvgpu.NVOS32_ATTR_LOCATION_MASK == nvgpu.NVOS32_ATTR_LOCATION_VIDMEM {
+				classID = nvgpu.NV01_MEMORY_LOCAL_USER
+			}
+			fi.fd.dev.nvp.objAdd(fi.ctx, ioctlParams.HRoot, allocSizeParams.HMemory, classID, &miscObject{}, ioctlParams.HObjectParent)
+		}
 	}
 	fi.fd.dev.nvp.objsUnlock()
 	allocSizeParams.Address = origAddress