From e81e0c72a70b0ffe31dba1e5b5b3a12e18f21cad Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Sat, 23 Sep 2023 00:44:47 -0700
Subject: [PATCH] Add NV_MEMORY_FABRIC to nvproxy.

Fixes #9372

PiperOrigin-RevId: 567816690
---
 pkg/abi/nvgpu/classes.go               | 20 ++++++++++++++++++++
 pkg/sentry/devices/nvproxy/frontend.go |  2 ++
 2 files changed, 22 insertions(+)

diff --git a/pkg/abi/nvgpu/classes.go b/pkg/abi/nvgpu/classes.go
index d1bdf09ed3..da6a159359 100644
--- a/pkg/abi/nvgpu/classes.go
+++ b/pkg/abi/nvgpu/classes.go
@@ -23,6 +23,7 @@ const (
 	NV01_MEMORY_SYSTEM_OS_DESCRIPTOR = 0x00000071
 	NV01_EVENT_OS_EVENT              = 0x00000079
 	NV01_DEVICE_0                    = 0x00000080
+	NV_MEMORY_FABRIC                 = 0x000000f8
 	NV20_SUBDEVICE_0                 = 0x00002080
 	NV50_THIRD_PARTY_P2P             = 0x0000503c
 	GT200_DEBUGGER                   = 0x000083de
@@ -220,3 +221,22 @@ type NV_HOPPER_USERMODE_A_PARAMS struct {
 	Bar1Mapping uint8
 	Priv        uint8
 }
+
+// +marshal
+type nv00f8Map struct {
+	offset  uint64
+	hVidMem Handle
+	flags   uint32
+}
+
+// NV00F8_ALLOCATION_PARAMETERS is the alloc param type for NV_MEMORY_FABRIC,
+// from src/common/sdk/nvidia/inc/class/cl00f8.h
+//
+// +marshal
+type NV00F8_ALLOCATION_PARAMETERS struct {
+	Alignment  uint64
+	AllocSize  uint64
+	PageSize   uint32
+	AllocFlags uint32
+	Map        nv00f8Map
+}
diff --git a/pkg/sentry/devices/nvproxy/frontend.go b/pkg/sentry/devices/nvproxy/frontend.go
index e005f1ae9a..66497a69f0 100644
--- a/pkg/sentry/devices/nvproxy/frontend.go
+++ b/pkg/sentry/devices/nvproxy/frontend.go
@@ -782,6 +782,8 @@ func rmAlloc(fi *frontendIoctlState) (uintptr, error) {
 		return rmAllocSimple[nvgpu.NV_HOPPER_USERMODE_A_PARAMS](fi, &ioctlParams, isNVOS64)
 	case nvgpu.GF100_SUBDEVICE_MASTER, nvgpu.VOLTA_USERMODE_A, nvgpu.TURING_USERMODE_A:
 		return rmAllocNoParams(fi, &ioctlParams, isNVOS64)
+	case nvgpu.NV_MEMORY_FABRIC:
+		return rmAllocSimple[nvgpu.NV00F8_ALLOCATION_PARAMETERS](fi, &ioctlParams, isNVOS64)
 	default:
 		fi.ctx.Warningf("nvproxy: unknown allocation class %#08x", ioctlParams.HClass)
 		return 0, linuxerr.EINVAL