Add support for Nvidia driver versions after 535.43.02.

Nvproxy now supports 535.54.03 and 535.104.05 versions.
Tested on T4 GPU for these versions.
```
$ docker run --runtime=runsc --rm --gpus=all nvcr.io/nvidia/k8s/cuda-sample:vectoradd-cuda11.7.1-ubi8
[Vector addition of 50000 elements]
Copy input data from the host memory to the CUDA device
CUDA kernel launch with 196 blocks of 256 threads
Copy output data from the CUDA device to the host memory
Test PASSED
Done
```

PiperOrigin-RevId: 572312006

pkg/abi/nvgpu/classes.go

@@ -43,6 +43,7 @@ const (
 	AMPERE_COMPUTE_B                 = 0x0000c7c0
 	HOPPER_DMA_COPY_A                = 0x0000c8b5
 	ADA_COMPUTE_A                    = 0x0000c9c0
+	NV_CONFIDENTIAL_COMPUTE          = 0x0000cb33
 	HOPPER_COMPUTE_A                 = 0x0000cbc0
 )
 
@@ -220,7 +221,7 @@ type nv00f8Map struct {
 }
 
 // NV00F8_ALLOCATION_PARAMETERS is the alloc param type for NV_MEMORY_FABRIC,
-// from src/common/sdk/nvidia/inc/class/cl00f8.h
+// from src/common/sdk/nvidia/inc/class/cl00f8.h.
 //
 // +marshal
 type NV00F8_ALLOCATION_PARAMETERS struct {
@@ -230,3 +231,12 @@ type NV00F8_ALLOCATION_PARAMETERS struct {
 	AllocFlags 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.
+//
+// +marshal
+type NV_CONFIDENTIAL_COMPUTE_ALLOC_PARAMS struct {
+	Handle Handle
+	_      uint32
+}

pkg/abi/nvgpu/ctrl.go

@@ -253,3 +253,8 @@ const (
 	NVA06C_CTRL_CMD_SET_TIMESLICE   = 0xa06c0103
 	NVA06C_CTRL_CMD_PREEMPT         = 0xa06c0105
 )
+
+// From src/common/sdk/nvidia/inc/ctrl/ctrlcb33.h:
+const (
+	NV_CONF_COMPUTE_CTRL_CMD_SYSTEM_GET_CAPABILITIES = 0xcb330101
+)

pkg/abi/nvgpu/frontend.go

@@ -143,18 +143,24 @@ type RmAllocParamType interface {
 	GetPRightsRequested() P64
 	SetPAllocParms(p P64)
 	SetPRightsRequested(p P64)
-	FromOS64(other NVOS64Parameters)
-	ToOS64() NVOS64Parameters
+	FromOS64V535(other NVOS64ParametersV535)
+	ToOS64V535() NVOS64ParametersV535
 	GetPointer() uintptr
 	marshal.Marshallable
 }
 
 // GetRmAllocParamObj returns the appropriate implementation of
 // RmAllocParamType based on passed parameters.
-func GetRmAllocParamObj(isNVOS64 bool) RmAllocParamType {
+func GetRmAllocParamObj(isNVOS64 bool, isV535 bool) RmAllocParamType {
 	if isNVOS64 {
+		if isV535 {
+			return &NVOS64ParametersV535{}
+		}
 		return &NVOS64Parameters{}
 	}
+	if isV535 {
+		return &NVOS21ParametersV535{}
+	}
 	return &NVOS21Parameters{}
 }
 
@@ -190,8 +196,8 @@ func (n *NVOS21Parameters) SetPRightsRequested(p P64) {
 	panic("impossible")
 }
 
-// FromOS64 implements RmAllocParamType.FromOS64.
-func (n *NVOS21Parameters) FromOS64(other NVOS64Parameters) {
+// FromOS64V535 implements RmAllocParamType.FromOS64V535.
+func (n *NVOS21Parameters) FromOS64V535(other NVOS64ParametersV535) {
 	n.HRoot = other.HRoot
 	n.HObjectParent = other.HObjectParent
 	n.HObjectNew = other.HObjectNew
@@ -200,9 +206,9 @@ func (n *NVOS21Parameters) FromOS64(other NVOS64Parameters) {
 	n.Status = other.Status
 }
 
-// ToOS64 implements RmAllocParamType.ToOS64.
-func (n *NVOS21Parameters) ToOS64() NVOS64Parameters {
-	return NVOS64Parameters{
+// ToOS64V535 implements RmAllocParamType.ToOS64V535.
+func (n *NVOS21Parameters) ToOS64V535() NVOS64ParametersV535 {
+	return NVOS64ParametersV535{
 		HRoot:         n.HRoot,
 		HObjectParent: n.HObjectParent,
 		HObjectNew:    n.HObjectNew,
@@ -212,6 +218,62 @@ func (n *NVOS21Parameters) ToOS64() NVOS64Parameters {
 	}
 }
 
+// NVOS21ParametersV535 is the updated version of NVOS21Parameters starting
+// from 535.43.02.
+//
+// +marshal
+type NVOS21ParametersV535 struct {
+	HRoot         Handle
+	HObjectParent Handle
+	HObjectNew    Handle
+	HClass        uint32
+	PAllocParms   P64
+	ParamsSize    uint32
+	Status        uint32
+}
+
+// GetPAllocParms implements RmAllocParamType.GetPAllocParms.
+func (n *NVOS21ParametersV535) GetPAllocParms() P64 {
+	return n.PAllocParms
+}
+
+// GetPRightsRequested implements RmAllocParamType.GetPRightsRequested.
+func (n *NVOS21ParametersV535) GetPRightsRequested() P64 {
+	return 0
+}
+
+// SetPAllocParms implements RmAllocParamType.SetPAllocParms.
+func (n *NVOS21ParametersV535) SetPAllocParms(p P64) { n.PAllocParms = p }
+
+// SetPRightsRequested implements RmAllocParamType.SetPRightsRequested.
+func (n *NVOS21ParametersV535) SetPRightsRequested(p P64) {
+	panic("impossible")
+}
+
+// FromOS64V535 implements RmAllocParamType.FromOS64V535.
+func (n *NVOS21ParametersV535) FromOS64V535(other NVOS64ParametersV535) {
+	n.HRoot = other.HRoot
+	n.HObjectParent = other.HObjectParent
+	n.HObjectNew = other.HObjectNew
+	n.HClass = other.HClass
+	n.PAllocParms = other.PAllocParms
+	n.ParamsSize = other.ParamsSize
+	n.Status = other.Status
+}
+
+// ToOS64V535 implements RmAllocParamType.ToOS64V535.
+func (n *NVOS21ParametersV535) ToOS64V535() NVOS64ParametersV535 {
+	return NVOS64ParametersV535{
+		HRoot:         n.HRoot,
+		HObjectParent: n.HObjectParent,
+		HObjectNew:    n.HObjectNew,
+		HClass:        n.HClass,
+		PAllocParms:   n.PAllocParms,
+		ParamsSize:    n.ParamsSize,
+		Status:        n.Status,
+	}
+}
+
 // NVOS55Parameters is NVOS55_PARAMETERS, the parameter type for
 // NV_ESC_RM_DUP_OBJECT.
 //
@@ -383,16 +445,71 @@ func (n *NVOS64Parameters) SetPAllocParms(p P64) { n.PAllocParms = p }
 // SetPRightsRequested implements RmAllocParamType.SetPRightsRequested.
 func (n *NVOS64Parameters) SetPRightsRequested(p P64) { n.PRightsRequested = p }
 
-// FromOS64 implements RmAllocParamType.FromOS64.
-func (n *NVOS64Parameters) FromOS64(other NVOS64Parameters) {
-	*n = other
+// FromOS64V535 implements RmAllocParamType.FromOS64V535.
+func (n *NVOS64Parameters) FromOS64V535(other NVOS64ParametersV535) {
+	n.HRoot = other.HRoot
+	n.HObjectParent = other.HObjectParent
+	n.HObjectNew = other.HObjectNew
+	n.HClass = other.HClass
+	n.PAllocParms = other.PAllocParms
+	n.PRightsRequested = other.PRightsRequested
+	n.Flags = other.Flags
+	n.Status = other.Status
 }
 
-// ToOS64 implements RmAllocParamType.ToOS64.
-func (n *NVOS64Parameters) ToOS64() NVOS64Parameters {
-	return *n
+// ToOS64V535 implements RmAllocParamType.ToOS64V535.
+func (n *NVOS64Parameters) ToOS64V535() NVOS64ParametersV535 {
+	return NVOS64ParametersV535{
+		HRoot:            n.HRoot,
+		HObjectParent:    n.HObjectParent,
+		HObjectNew:       n.HObjectNew,
+		HClass:           n.HClass,
+		PAllocParms:      n.PAllocParms,
+		PRightsRequested: n.PRightsRequested,
+		Flags:            n.Flags,
+		Status:           n.Status,
+	}
 }
 
+// NVOS64ParametersV535 is the updated version of NVOS64Parameters starting
+// from 535.43.02.
+//
+// +marshal
+type NVOS64ParametersV535 struct {
+	HRoot            Handle
+	HObjectParent    Handle
+	HObjectNew       Handle
+	HClass           uint32
+	PAllocParms      P64
+	PRightsRequested P64
+	ParamsSize       uint32
+	Flags            uint32
+	Status           uint32
+	_                uint32
+}
+
+// GetPAllocParms implements RmAllocParamType.GetPAllocParms.
+func (n *NVOS64ParametersV535) GetPAllocParms() P64 {
+	return n.PAllocParms
+}
+
+// GetPRightsRequested implements RmAllocParamType.GetPRightsRequested.
+func (n *NVOS64ParametersV535) GetPRightsRequested() P64 {
+	return n.PRightsRequested
+}
+
+// SetPAllocParms implements RmAllocParamType.SetPAllocParms.
+func (n *NVOS64ParametersV535) SetPAllocParms(p P64) { n.PAllocParms = p }
+
+// SetPRightsRequested implements RmAllocParamType.SetPRightsRequested.
+func (n *NVOS64ParametersV535) SetPRightsRequested(p P64) { n.PRightsRequested = p }
+
+// FromOS64V535 implements RmAllocParamType.FromOS64V535.
+func (n *NVOS64ParametersV535) FromOS64V535(other NVOS64ParametersV535) { *n = other }
+
+// ToOS64V535 implements RmAllocParamType.ToOS64V535.
+func (n *NVOS64ParametersV535) ToOS64V535() NVOS64ParametersV535 { return *n }
+
 // Frontend ioctl parameter struct sizes.
 var (
 	SizeofIoctlRegisterFD             = uint32((*IoctlRegisterFD)(nil).SizeBytes())
@@ -403,6 +520,7 @@ var (
 	SizeofIoctlNVOS02ParametersWithFD = uint32((*IoctlNVOS02ParametersWithFD)(nil).SizeBytes())
 	SizeofNVOS00Parameters            = uint32((*NVOS00Parameters)(nil).SizeBytes())
 	SizeofNVOS21Parameters            = uint32((*NVOS21Parameters)(nil).SizeBytes())
+	SizeofNVOS21ParametersV535        = uint32((*NVOS21ParametersV535)(nil).SizeBytes())
 	SizeofIoctlNVOS33ParametersWithFD = uint32((*IoctlNVOS33ParametersWithFD)(nil).SizeBytes())
 	SizeofNVOS55Parameters            = uint32((*NVOS55Parameters)(nil).SizeBytes())
 	SizeofNVOS57Parameters            = uint32((*NVOS57Parameters)(nil).SizeBytes())
@@ -411,4 +529,5 @@ var (
 	SizeofNVOS54Parameters            = uint32((*NVOS54Parameters)(nil).SizeBytes())
 	SizeofNVOS56Parameters            = uint32((*NVOS56Parameters)(nil).SizeBytes())
 	SizeofNVOS64Parameters            = uint32((*NVOS64Parameters)(nil).SizeBytes())
+	SizeofNVOS64ParametersV535        = uint32((*NVOS64ParametersV535)(nil).SizeBytes())
 )

pkg/abi/nvgpu/frontend_unsafe.go

@@ -21,7 +21,17 @@ func (n *NVOS21Parameters) GetPointer() uintptr {
 	return uintptr(unsafe.Pointer(n))
 }
 
+// GetPointer implements RmAllocParamType.GetPointer.
+func (n *NVOS21ParametersV535) GetPointer() uintptr {
+	return uintptr(unsafe.Pointer(n))
+}
+
 // GetPointer implements RmAllocParamType.GetPointer.
 func (n *NVOS64Parameters) GetPointer() uintptr {
 	return uintptr(unsafe.Pointer(n))
 }
+
+// GetPointer implements RmAllocParamType.GetPointer.
+func (n *NVOS64ParametersV535) GetPointer() uintptr {
+	return uintptr(unsafe.Pointer(n))
+}

pkg/abi/nvgpu/status.go

@@ -16,7 +16,8 @@ package nvgpu
 
 // Status codes, from src/common/sdk/nvidia/inc/nvstatuscodes.h.
 const (
-	NV_ERR_INVALID_ADDRESS = 0x0000001e
-	NV_ERR_INVALID_LIMIT   = 0x0000002e
-	NV_ERR_NOT_SUPPORTED   = 0x00000056
+	NV_ERR_INVALID_ADDRESS  = 0x0000001e
+	NV_ERR_INVALID_ARGUMENT = 0x0000001f
+	NV_ERR_INVALID_LIMIT    = 0x0000002e
+	NV_ERR_NOT_SUPPORTED    = 0x00000056
 )

pkg/abi/nvgpu/uvm.go

@@ -44,6 +44,7 @@ const (
 	UVM_ALLOC_SEMAPHORE_POOL           = 68
 	UVM_VALIDATE_VA_RANGE              = 72
 	UVM_CREATE_EXTERNAL_RANGE          = 73
+	UVM_MM_INITIALIZE                  = 75
 )
 
 // +marshal
@@ -221,6 +222,12 @@ type UVM_CREATE_EXTERNAL_RANGE_PARAMS struct {
 	Pad0     [4]byte
 }
 
+// +marshal
+type UVM_MM_INITIALIZE_PARAMS struct {
+	UvmFD  int32
+	Status uint32
+}
+
 // From kernel-open/nvidia-uvm/uvm_types.h:
 
 const UVM_MAX_GPUS = NV_MAX_DEVICES

pkg/sentry/devices/nvproxy/BUILD

@@ -59,4 +59,5 @@ go_test(
     name = "nvproxy_test",
     srcs = ["nvproxy_test.go"],
     library = ":nvproxy",
+    deps = ["//pkg/abi/nvgpu"],
 )

pkg/sentry/devices/nvproxy/frontend.go

@@ -144,6 +144,10 @@ func (fd *frontendFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr,
 		panic("Ioctl should be called from a task context")
 	}
 
+	if log.IsLogging(log.Debug) {
+		ctx.Debugf("nvproxy: frontend ioctl: nr = %#08x, argSize = %#08x", nr, argSize)
+	}
+
 	fi := frontendIoctlState{
 		fd:              fd,
 		ctx:             ctx,
@@ -154,8 +158,6 @@ func (fd *frontendFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr,
 	}
 
 	// nr determines the argument type.
-	// Don't log nr since it's already visible as the last byte of cmd in
-	// strace logging.
 	// Implementors:
 	// - To map nr to a symbol, look in
 	// src/nvidia/arch/nvalloc/unix/include/nv_escape.h,
@@ -606,20 +608,30 @@ func ctrlSubdevFIFODisableChannels(fi *frontendIoctlState, ioctlParams *nvgpu.NV
 
 func rmAlloc(fi *frontendIoctlState) (uintptr, error) {
 	var isNVOS64 bool
-	switch fi.ioctlParamsSize {
-	case nvgpu.SizeofNVOS21Parameters:
-	case nvgpu.SizeofNVOS64Parameters:
-		isNVOS64 = true
-	default:
-		return 0, linuxerr.EINVAL
+	if fi.fd.nvp.abi.useRmAllocParamsV535 {
+		switch fi.ioctlParamsSize {
+		case nvgpu.SizeofNVOS21ParametersV535:
+		case nvgpu.SizeofNVOS64ParametersV535:
+			isNVOS64 = true
+		default:
+			return 0, linuxerr.EINVAL
+		}
+	} else {
+		switch fi.ioctlParamsSize {
+		case nvgpu.SizeofNVOS21Parameters:
+		case nvgpu.SizeofNVOS64Parameters:
+			isNVOS64 = true
+		default:
+			return 0, linuxerr.EINVAL
+		}
 	}
-	// Copy in parameters and convert to NVOS64ParametersR535, which is a super
+	// Copy in parameters and convert to NVOS64ParametersV535, which is a super
 	// set of all parameter types we support.
-	buf := nvgpu.GetRmAllocParamObj(isNVOS64)
+	buf := nvgpu.GetRmAllocParamObj(isNVOS64, fi.fd.nvp.abi.useRmAllocParamsV535)
 	if _, err := buf.CopyIn(fi.t, fi.ioctlParamsAddr); err != nil {
 		return 0, err
 	}
-	ioctlParams := buf.ToOS64()
+	ioctlParams := buf.ToOS64V535()
 
 	// hClass determines the type of pAllocParms.
 	if log.IsLogging(log.Debug) {
@@ -643,7 +655,7 @@ func rmAlloc(fi *frontendIoctlState) (uintptr, error) {
 
 // Unlike frontendIoctlSimple and rmControlSimple, rmAllocSimple requires the
 // parameter type since the parameter's size is otherwise unknown.
-func rmAllocSimple[Params any, PParams marshalPtr[Params]](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
+func rmAllocSimple[Params any, PParams marshalPtr[Params]](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64ParametersV535, isNVOS64 bool) (uintptr, error) {
 	if ioctlParams.PAllocParms == 0 {
 		return rmAllocInvoke[byte](fi, ioctlParams, nil, isNVOS64)
 	}
@@ -662,11 +674,11 @@ func rmAllocSimple[Params any, PParams marshalPtr[Params]](fi *frontendIoctlStat
 	return n, nil
 }
 
-func rmAllocNoParams(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
+func rmAllocNoParams(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64ParametersV535, isNVOS64 bool) (uintptr, error) {
 	return rmAllocInvoke[byte](fi, ioctlParams, nil, isNVOS64)
 }
 
-func rmAllocEventOSEvent(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error) {
+func rmAllocEventOSEvent(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64ParametersV535, isNVOS64 bool) (uintptr, error) {
 	var allocParams nvgpu.NV0005_ALLOC_PARAMETERS
 	if _, err := allocParams.CopyIn(fi.t, addrFromP64(ioctlParams.PAllocParms)); err != nil {
 		return 0, err

pkg/sentry/devices/nvproxy/frontend_unsafe.go

@@ -158,11 +158,11 @@ func ctrlSubdevGRGetInfo(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parame
 	return n, nil
 }
 
-func rmAllocInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, allocParams *Params, isNVOS64 bool) (uintptr, error) {
+func rmAllocInvoke[Params any](fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64ParametersV535, allocParams *Params, isNVOS64 bool) (uintptr, error) {
 	defer runtime.KeepAlive(allocParams) // since we convert to non-pointer-typed P64
 
-	sentryIoctlParams := nvgpu.GetRmAllocParamObj(isNVOS64)
-	sentryIoctlParams.FromOS64(*ioctlParams)
+	sentryIoctlParams := nvgpu.GetRmAllocParamObj(isNVOS64, fi.fd.nvp.abi.useRmAllocParamsV535)
+	sentryIoctlParams.FromOS64V535(*ioctlParams)
 	sentryIoctlParams.SetPAllocParms(p64FromPtr(unsafe.Pointer(allocParams)))
 	var rightsRequested nvgpu.RS_ACCESS_MASK
 	if ioctlParams.PRightsRequested != 0 {

pkg/sentry/devices/nvproxy/nvproxy_test.go

@@ -16,6 +16,8 @@ package nvproxy
 
 import (
 	"testing"
+
+	"gvisor.dev/gvisor/pkg/abi/nvgpu"
 )
 
 func TestInit(t *testing.T) {
@@ -25,3 +27,12 @@ func TestInit(t *testing.T) {
 		cons()
 	}
 }
+
+func TestNVOS21ParamsSize(t *testing.T) {
+	if nvgpu.SizeofNVOS21ParametersV535 != nvgpu.SizeofNVOS21Parameters {
+		// We assume the size of NVOS21_PARAMETERS struct did not change between
+		// V525 and V535. If this turns out to be false, a separate seccomp entry
+		// needs to be added for the new size value.
+		t.Errorf("SizeofNVOS21ParametersV535(%#08x) != SizeofNVOS21Parameters(%#08x)", nvgpu.SizeofNVOS21ParametersV535, nvgpu.SizeofNVOS21Parameters)
+	}
+}

pkg/sentry/devices/nvproxy/seccomp_filters.go

@@ -77,10 +77,16 @@ func Filters() seccomp.SyscallRules {
 				nonNegativeFD,
 				seccomp.EqualTo(frontendIoctlCmd(nvgpu.NV_ESC_RM_ALLOC, nvgpu.SizeofNVOS21Parameters)),
 			},
+			// Note that we don't need to add one for NVOS21ParametersV535, because
+			// SizeofNVOS21ParametersV535 == SizeofNVOS21Parameters. We test this.
 			seccomp.PerArg{
 				nonNegativeFD,
 				seccomp.EqualTo(frontendIoctlCmd(nvgpu.NV_ESC_RM_ALLOC, nvgpu.SizeofNVOS64Parameters)),
 			},
+			seccomp.PerArg{
+				nonNegativeFD,
+				seccomp.EqualTo(frontendIoctlCmd(nvgpu.NV_ESC_RM_ALLOC, nvgpu.SizeofNVOS64ParametersV535)),
+			},
 			seccomp.PerArg{
 				nonNegativeFD,
 				seccomp.EqualTo(frontendIoctlCmd(nvgpu.NV_ESC_RM_DUP_OBJECT, nvgpu.SizeofNVOS55Parameters)),
@@ -109,6 +115,10 @@ func Filters() seccomp.SyscallRules {
 				nonNegativeFD,
 				seccomp.EqualTo(nvgpu.UVM_INITIALIZE),
 			},
+			seccomp.PerArg{
+				nonNegativeFD,
+				seccomp.EqualTo(nvgpu.UVM_MM_INITIALIZE),
+			},
 			seccomp.PerArg{
 				nonNegativeFD,
 				seccomp.EqualTo(nvgpu.UVM_DEINITIALIZE),

pkg/sentry/devices/nvproxy/uvm.go

@@ -23,6 +23,7 @@ import (
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/fdnotifier"
 	"gvisor.dev/gvisor/pkg/hostarch"
+	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/marshal"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
@@ -124,6 +125,10 @@ func (fd *uvmFD) Ioctl(ctx context.Context, uio usermem.IO, sysno uintptr, args
 		panic("Ioctl should be called from a task context")
 	}
 
+	if log.IsLogging(log.Debug) {
+		ctx.Debugf("nvproxy: frontend ioctl %#08x", cmd)
+	}
+
 	ui := uvmIoctlState{
 		fd:              fd,
 		ctx:             ctx,
@@ -190,6 +195,44 @@ func uvmInitialize(ui *uvmIoctlState) (uintptr, error) {
 	return n, nil
 }
 
+func uvmMMInitialize(ui *uvmIoctlState) (uintptr, error) {
+	var ioctlParams nvgpu.UVM_MM_INITIALIZE_PARAMS
+	if _, err := ioctlParams.CopyIn(ui.t, ui.ioctlParamsAddr); err != nil {
+		return 0, err
+	}
+
+	failWithStatus := func(status uint32) error {
+		outIoctlParams := ioctlParams
+		outIoctlParams.Status = status
+		_, err := outIoctlParams.CopyOut(ui.t, ui.ioctlParamsAddr)
+		return err
+	}
+
+	uvmFileGeneric, _ := ui.t.FDTable().Get(ioctlParams.UvmFD)
+	if uvmFileGeneric == nil {
+		return 0, failWithStatus(nvgpu.NV_ERR_INVALID_ARGUMENT)
+	}
+	defer uvmFileGeneric.DecRef(ui.ctx)
+	uvmFile, ok := uvmFileGeneric.Impl().(*uvmFD)
+	if !ok {
+		return 0, failWithStatus(nvgpu.NV_ERR_INVALID_ARGUMENT)
+	}
+
+	sentryIoctlParams := ioctlParams
+	sentryIoctlParams.UvmFD = uvmFile.hostFD
+	n, err := uvmIoctlInvoke(ui, &sentryIoctlParams)
+	if err != nil {
+		return n, err
+	}
+
+	outIoctlParams := sentryIoctlParams
+	outIoctlParams.UvmFD = ioctlParams.UvmFD
+	if _, err := outIoctlParams.CopyOut(ui.t, ui.ioctlParamsAddr); err != nil {
+		return n, err
+	}
+	return n, nil
+}
+
 type hasRMCtrlFDPtr[T any] interface {
 	*T
 	marshal.Marshallable

pkg/sentry/devices/nvproxy/version.go

@@ -59,7 +59,7 @@ func (v driverVersion) String() string {
 
 type frontendIoctlHandler func(fi *frontendIoctlState) (uintptr, error)
 type controlCmdHandler func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters) (uintptr, error)
-type allocationClassHandler func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64Parameters, isNVOS64 bool) (uintptr, error)
+type allocationClassHandler func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64ParametersV535, isNVOS64 bool) (uintptr, error)
 type uvmIoctlHandler func(ui *uvmIoctlState) (uintptr, error)
 
 // A driverABIFunc constructs and returns a driverABI.
@@ -83,6 +83,8 @@ type driverABI struct {
 	uvmIoctl        map[uint32]uvmIoctlHandler
 	controlCmd      map[uint32]controlCmdHandler
 	allocationClass map[uint32]allocationClassHandler
+
+	useRmAllocParamsV535 bool
 }
 
 // abis is a global map containing all supported Nvidia driver ABIs. This is
@@ -250,8 +252,25 @@ func Init() {
 			}
 		})
 
-		v525_105_17 := addDriverABI(525, 105, 17, v525_60_13)
+		// 525.89.02 is an intermediate unqualified version from the main branch.
+		v525_89_02 := v525_60_13
 
+		// The following versions do not exist on the main branch. They branched off
+		// the main branch at 525.89.02.
+		v525_105_17 := addDriverABI(525, 105, 17, v525_89_02)
 		_ = addDriverABI(525, 125, 06, v525_105_17)
+
+		// v535.43.02 is an intermediate unqualified version from the main branch.
+		v535_43_02 := func() *driverABI {
+			abi := v525_89_02()
+			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.uvmIoctl[nvgpu.UVM_MM_INITIALIZE] = uvmMMInitialize
+			return abi
+		}
+
+		v535_54_03 := addDriverABI(535, 54, 03, v535_43_02)
+		_ = addDriverABI(535, 104, 05, v535_54_03)
 	})
 }