From fd9845ccea9bccebd6dbd1716e3216653c64e493 Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushranjan@google.com>
Date: Sat, 28 Oct 2023 12:42:21 -0700
Subject: [PATCH] Fix NV_MEMORY_FABRIC for 535.43.02+ drivers.

Updates #9372

PiperOrigin-RevId: 577494512
---
 pkg/abi/nvgpu/classes.go              | 13 +++++++++++++
 pkg/sentry/devices/nvproxy/version.go |  1 +
 2 files changed, 14 insertions(+)

diff --git a/pkg/abi/nvgpu/classes.go b/pkg/abi/nvgpu/classes.go
index c526823b3..253bbb747 100644
--- a/pkg/abi/nvgpu/classes.go
+++ b/pkg/abi/nvgpu/classes.go
@@ -232,6 +232,19 @@ type NV00F8_ALLOCATION_PARAMETERS struct {
 	Map        nv00f8Map
 }
 
+// NV00F8_ALLOCATION_PARAMETERS_V535 is the updated version of
+// NV00F8_ALLOCATION_PARAMETERS since 535.43.02.
+//
+// +marshal
+type NV00F8_ALLOCATION_PARAMETERS_V535 struct {
+	Alignment  uint64
+	AllocSize  uint64
+	PageSize   uint64
+	AllocFlags uint32
+	_          uint32
+	Map        nv00f8Map
+}
+
 // NV_CONFIDENTIAL_COMPUTE_ALLOC_PARAMS is the alloc param type for
 // NV_CONFIDENTIAL_COMPUTE, from src/common/sdk/nvidia/inc/class/clcb33.h.
 //
diff --git a/pkg/sentry/devices/nvproxy/version.go b/pkg/sentry/devices/nvproxy/version.go
index b05e9518d..45f0e45d6 100644
--- a/pkg/sentry/devices/nvproxy/version.go
+++ b/pkg/sentry/devices/nvproxy/version.go
@@ -316,6 +316,7 @@ func Init() {
 			abi.useRmAllocParamsV535 = true
 			abi.controlCmd[nvgpu.NV_CONF_COMPUTE_CTRL_CMD_SYSTEM_GET_CAPABILITIES] = rmControlSimple
 			abi.allocationClass[nvgpu.NV_CONFIDENTIAL_COMPUTE] = rmAllocSimple[nvgpu.NV_CONFIDENTIAL_COMPUTE_ALLOC_PARAMS]
+			abi.allocationClass[nvgpu.NV_MEMORY_FABRIC] = rmAllocSimple[nvgpu.NV00F8_ALLOCATION_PARAMETERS_V535]
 			abi.uvmIoctl[nvgpu.UVM_MM_INITIALIZE] = uvmMMInitialize
 			return abi
 		}