nvproxy: don't depend on HVASpace for non-virtmem NVOS32_FUNCTION_ALL…

…OC_SIZE PiperOrigin-RevId: 704780681
google · Dec 10, 2024 · 01e30ab · 01e30ab
1 parent c19c9c3
commit 01e30ab
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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