From 98df07c3cd71072655aa6b3ecaec35fa195273d2 Mon Sep 17 00:00:00 2001
From: Ayush Ranjan <ayushranjan@google.com>
Date: Mon, 8 Jan 2024 19:19:36 -0800
Subject: [PATCH] Add NV50_P2P allocation class support in nvproxy.

Updates #9827.

PiperOrigin-RevId: 596766912
---
 pkg/abi/nvgpu/classes.go              | 35 +++++++++++++++++++++++++++
 pkg/sentry/devices/nvproxy/version.go |  2 ++
 2 files changed, 37 insertions(+)

diff --git a/pkg/abi/nvgpu/classes.go b/pkg/abi/nvgpu/classes.go
index 253bbb747..304a37c65 100644
--- a/pkg/abi/nvgpu/classes.go
+++ b/pkg/abi/nvgpu/classes.go
@@ -25,6 +25,7 @@ const (
 	NV01_DEVICE_0                    = 0x00000080
 	NV_MEMORY_FABRIC                 = 0x000000f8
 	NV20_SUBDEVICE_0                 = 0x00002080
+	NV50_P2P                         = 0x0000503b
 	NV50_THIRD_PARTY_P2P             = 0x0000503c
 	GT200_DEBUGGER                   = 0x000083de
 	GF100_SUBDEVICE_MASTER           = 0x000090e6
@@ -85,6 +86,40 @@ type NV2080_ALLOC_PARAMETERS struct {
 	SubDeviceID uint32
 }
 
+// NV503B_ALLOC_PARAMETERS is the alloc params type for NV50_P2P, from
+// src/common/sdk/nvidia/inc/class/cl503b.h.
+//
+// +marshal
+type NV503B_ALLOC_PARAMETERS struct {
+	HSubDevice              Handle
+	HPeerSubDevice          Handle
+	SubDevicePeerIDMask     uint32
+	PeerSubDevicePeerIDMask uint32
+	MailboxBar1Addr         uint64
+	MailboxTotalSize        uint32
+	Flags                   uint32
+}
+
+// NV503B_BAR1_P2P_DMA_INFO from src/common/sdk/nvidia/inc/class/cl503b.h.
+//
+// +marshal
+type NV503B_BAR1_P2P_DMA_INFO struct {
+	DmaAddress uint64
+	DmaSize    uint64
+}
+
+// NV503B_ALLOC_PARAMETERS_V535 is the updated version of
+// NV503B_ALLOC_PARAMETERS since 535.43.02.
+//
+// +marshal
+type NV503B_ALLOC_PARAMETERS_V535 struct {
+	NV503B_ALLOC_PARAMETERS
+	SubDeviceEgmPeerIDMask     uint32
+	PeerSubDeviceEgmPeerIDMask uint32
+	L2pBar1P2PDmaInfo          NV503B_BAR1_P2P_DMA_INFO
+	P2lBar1P2PDmaInfo          NV503B_BAR1_P2P_DMA_INFO
+}
+
 // NV503C_ALLOC_PARAMETERS is the alloc params type for NV50_THIRD_PARTY_P2P,
 // from src/common/sdk/nvidia/inc/class/cl503c.h.
 //
diff --git a/pkg/sentry/devices/nvproxy/version.go b/pkg/sentry/devices/nvproxy/version.go
index 400fa216c..210ea774a 100644
--- a/pkg/sentry/devices/nvproxy/version.go
+++ b/pkg/sentry/devices/nvproxy/version.go
@@ -282,6 +282,7 @@ func Init() {
 					nvgpu.NV01_EVENT_OS_EVENT:     rmAllocEventOSEvent,
 					nvgpu.NV01_DEVICE_0:           rmAllocSimple[nvgpu.NV0080_ALLOC_PARAMETERS],
 					nvgpu.NV20_SUBDEVICE_0:        rmAllocSimple[nvgpu.NV2080_ALLOC_PARAMETERS],
+					nvgpu.NV50_P2P:                rmAllocSimple[nvgpu.NV503B_ALLOC_PARAMETERS],
 					nvgpu.NV50_THIRD_PARTY_P2P:    rmAllocSimple[nvgpu.NV503C_ALLOC_PARAMETERS],
 					nvgpu.GT200_DEBUGGER:          rmAllocSimple[nvgpu.NV83DE_ALLOC_PARAMETERS],
 					nvgpu.FERMI_CONTEXT_SHARE_A:   rmAllocSimple[nvgpu.NV_CTXSHARE_ALLOCATION_PARAMETERS],
@@ -322,6 +323,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.NV50_P2P] = rmAllocSimple[nvgpu.NV503B_ALLOC_PARAMETERS_V535]
 			abi.allocationClass[nvgpu.NV_MEMORY_FABRIC] = rmAllocSimple[nvgpu.NV00F8_ALLOCATION_PARAMETERS_V535]
 			abi.uvmIoctl[nvgpu.UVM_MM_INITIALIZE] = uvmMMInitialize
 			return abi