Skip to content

Commit aa5d1af

Browse files
zkoopmansgvisor-bot
zkoopmans
authored and
gvisor-bot
committed
Update the gpu driver installer to support ARM drivers
PiperOrigin-RevId: 730617254
1 parent 7d4b4bd commit aa5d1af

File tree

5 files changed

+276
-113
lines changed

5 files changed

+276
-113
lines changed

‎pkg/sentry/devices/nvproxy/nvproxy_test.go

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,11 @@ func TestInit(t *testing.T) {
3838
func TestAllSupportedHashesPresent(t *testing.T) {
3939
Init()
4040
for version, abi := range abis {
41-
if abi.checksum == "" {
42-
t.Errorf("unexpected empty value for driver %q", version.String())
41+
if abi.checksumX86_64 == "" {
42+
t.Errorf("unexpected empty checksumX86_64 value for driver %q", version.String())
43+
}
44+
if abi.checksumARM64 == "" {
45+
t.Errorf("unexpected empty checksumARM64 value for driver %q", version.String())
4346
}
4447
}
4548
}

‎pkg/sentry/devices/nvproxy/version.go

Lines changed: 167 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -107,8 +107,9 @@ type driverABIStructsFunc func() *driverABIStructs
107107
// abiConAndChecksum couples the driver's abiConstructor to the SHA256 checksum of its linux .run
108108
// driver installer file from NVIDIA.
109109
type abiConAndChecksum struct {
110-
cons driverABIFunc
111-
checksum string
110+
cons driverABIFunc
111+
checksumX86_64 string
112+
checksumARM64 string
112113
}
113114

114115
// driverABI defines the Nvidia kernel driver ABI proxied at a given version.
@@ -157,17 +158,28 @@ type DriverStruct struct {
157158
var abis map[DriverVersion]abiConAndChecksum
158159
var abisOnce sync.Once
159160

160-
// Note: runfileChecksum is the checksum of the .run file of the driver installer for linux from
161+
type addDriverABIArgs struct {
162+
major, minor, patch int
163+
runfileChecksumX86_64 string
164+
runfileChecksumARM64 string
165+
cons driverABIFunc
166+
}
167+
168+
// Note: runfileChecksums are the checksum of the .run file of the driver installer for linux from
161169
// nvidia.
162170
// To add a new version, add in support as normal and add the "addDriverABI" call for your version.
163171
// Run `make sudo TARGETS=//tools/gpu:main ARGS="checksum --version={}"` to get checksum.
164-
func addDriverABI(major, minor, patch int, runfileChecksum string, cons driverABIFunc) driverABIFunc {
172+
func addDriverABI(args addDriverABIArgs) driverABIFunc {
165173
if abis == nil {
166174
abis = make(map[DriverVersion]abiConAndChecksum)
167175
}
168-
version := NewDriverVersion(major, minor, patch)
169-
abis[version] = abiConAndChecksum{cons: cons, checksum: runfileChecksum}
170-
return cons
176+
version := NewDriverVersion(args.major, args.minor, args.patch)
177+
abis[version] = abiConAndChecksum{
178+
cons: args.cons,
179+
checksumX86_64: args.runfileChecksumX86_64,
180+
checksumARM64: args.runfileChecksumARM64,
181+
}
182+
return args.cons
171183
}
172184

173185
// Init initializes abis global map.
@@ -702,10 +714,41 @@ func Init() {
702714

703715
// The following exist on the "535" branch. They branched off the main
704716
// branch at 535.113.01.
705-
v535_183_01 := addDriverABI(535, 183, 01, "f6707afbdda9407e3cbc2e5128e60bcbcdbf02fae29958c72fafb5d405e8b883", v535_113_01)
706-
v535_183_06 := addDriverABI(535, 183, 06, "c7bb0a0569c5347845479ed4e3e4d885c6ee3b8adf068c3401cdf754d5ba3d3b", v535_183_01)
707-
v535_216_01 := addDriverABI(535, 216, 01, "5ddea1147810012e33967c3181341bcd6624bd3d654c63f845df833b4ece6af7", v535_183_06)
708-
_ = addDriverABI(535, 230, 02, "20cca9118083fcc8083158466e9cb2b616a7922206bcb7296b1fa5cc9af2e0fd", v535_216_01)
717+
v535_183_01 := addDriverABI(
718+
addDriverABIArgs{
719+
major: 535,
720+
minor: 183,
721+
patch: 01,
722+
runfileChecksumX86_64: "f6707afbdda9407e3cbc2e5128e60bcbcdbf02fae29958c72fafb5d405e8b883",
723+
runfileChecksumARM64: "c9d13b6250d24b76ef87a49b179f234564184a9f6d6414184668958b7f6d21e6",
724+
cons: v535_113_01,
725+
})
726+
v535_183_06 := addDriverABI(
727+
addDriverABIArgs{
728+
major: 535,
729+
minor: 183,
730+
patch: 06,
731+
runfileChecksumX86_64: "c7bb0a0569c5347845479ed4e3e4d885c6ee3b8adf068c3401cdf754d5ba3d3b",
732+
runfileChecksumARM64: "af3f72f5e4906805987844636b87ad1132650d05116272824c76dcc3f816d8e9",
733+
cons: v535_183_01,
734+
})
735+
v535_216_01 := addDriverABI(addDriverABIArgs{
736+
major: 535,
737+
minor: 216,
738+
patch: 01,
739+
runfileChecksumX86_64: "5ddea1147810012e33967c3181341bcd6624bd3d654c63f845df833b4ece6af7",
740+
runfileChecksumARM64: "4869ae0345b5892b2a50aed566c8226d3e07813d1190aa466feba5e9e21b33b9",
741+
cons: v535_183_06,
742+
})
743+
744+
_ = addDriverABI(addDriverABIArgs{
745+
major: 535,
746+
minor: 230,
747+
patch: 02,
748+
runfileChecksumX86_64: "20cca9118083fcc8083158466e9cb2b616a7922206bcb7296b1fa5cc9af2e0fd",
749+
runfileChecksumARM64: "ea000e6ff481f55e9bfedbea93b739368c635fe4be6156fdad560524ac7f363b",
750+
cons: v535_216_01,
751+
})
709752

710753
// 545.23.06 is an intermediate unqualified version from the main branch.
711754
v545_23_06 := func() *driverABI {
@@ -774,47 +817,90 @@ func Init() {
774817
return abi
775818
}
776819

777-
v550_54_14 := addDriverABI(550, 54, 14, "8c497ff1cfc7c310fb875149bc30faa4fd26d2237b2cba6cd2e8b0780157cfe3", func() *driverABI {
778-
abi := v550_40_07()
779-
abi.uvmIoctl[nvgpu.UVM_ALLOC_SEMAPHORE_POOL] = uvmHandler(uvmIoctlSimple[nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS_V550], compUtil)
780-
abi.uvmIoctl[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION] = uvmHandler(uvmIoctlHasFrontendFD[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS_V550], compUtil)
820+
v550_54_14 := addDriverABI(
821+
addDriverABIArgs{
822+
major: 550,
823+
minor: 54,
824+
patch: 14,
825+
runfileChecksumX86_64: "8c497ff1cfc7c310fb875149bc30faa4fd26d2237b2cba6cd2e8b0780157cfe3",
826+
runfileChecksumARM64: "b0fae8061633885c24f6b0c047649b46249a3bb44cadffbf658af28f80642c1d",
827+
cons: func() *driverABI {
828+
abi := v550_40_07()
829+
abi.uvmIoctl[nvgpu.UVM_ALLOC_SEMAPHORE_POOL] = uvmHandler(uvmIoctlSimple[nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS_V550], compUtil)
830+
abi.uvmIoctl[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION] = uvmHandler(uvmIoctlHasFrontendFD[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS_V550], compUtil)
781831

782-
prevStructs := abi.getStructs
783-
abi.getStructs = func() *driverABIStructs {
784-
structs := prevStructs()
785-
structs.uvmStructs[nvgpu.UVM_ALLOC_SEMAPHORE_POOL] = driverStructWithName(nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS_V550{}, "UVM_ALLOC_SEMAPHORE_POOL_PARAMS")
786-
structs.uvmStructs[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION] = driverStructWithName(nvgpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS_V550{}, "UVM_MAP_EXTERNAL_ALLOCATION_PARAMS")
787-
return structs
788-
}
789-
790-
return abi
791-
})
832+
prevStructs := abi.getStructs
833+
abi.getStructs = func() *driverABIStructs {
834+
structs := prevStructs()
835+
structs.uvmStructs[nvgpu.UVM_ALLOC_SEMAPHORE_POOL] = driverStructWithName(nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS_V550{}, "UVM_ALLOC_SEMAPHORE_POOL_PARAMS")
836+
structs.uvmStructs[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION] = driverStructWithName(nvgpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS_V550{}, "UVM_MAP_EXTERNAL_ALLOCATION_PARAMS")
837+
return structs
838+
}
839+
return abi
840+
},
841+
},
842+
)
792843

793-
v550_54_15 := addDriverABI(550, 54, 15, "2e859ae5f912a9a47aaa9b2d40a94a14f6f486b5d3b67c0ddf8b72c1c9650385", v550_54_14)
844+
v550_54_15 := addDriverABI(
845+
addDriverABIArgs{
846+
major: 550,
847+
minor: 54,
848+
patch: 15,
849+
runfileChecksumX86_64: "2e859ae5f912a9a47aaa9b2d40a94a14f6f486b5d3b67c0ddf8b72c1c9650385",
850+
runfileChecksumARM64: "49072d0c36ed85c7d8046776d34886f9ede9a6e4f46d5c7d533e8a8921d94cc1",
851+
cons: v550_54_14,
852+
},
853+
)
794854

795-
v550_90_07 := addDriverABI(550, 90, 07, "51acf579d5a9884f573a1d3f522e7fafa5e7841e22a9cec0b4bbeae31b0b9733", func() *driverABI {
796-
abi := v550_54_15()
797-
abi.controlCmd[nvgpu.NV_CONF_COMPUTE_CTRL_CMD_GPU_GET_KEY_ROTATION_STATE] = ctrlHandler(rmControlSimple, compUtil)
855+
v550_90_07 := addDriverABI(
856+
addDriverABIArgs{
857+
major: 550,
858+
minor: 90,
859+
patch: 07,
860+
runfileChecksumX86_64: "51acf579d5a9884f573a1d3f522e7fafa5e7841e22a9cec0b4bbeae31b0b9733",
861+
runfileChecksumARM64: "b896b76ae465307afc5b269c40bd8ccb279e6ea7d3ecae95534a91ecb1971572",
862+
cons: func() *driverABI {
863+
abi := v550_54_15()
864+
abi.controlCmd[nvgpu.NV_CONF_COMPUTE_CTRL_CMD_GPU_GET_KEY_ROTATION_STATE] = ctrlHandler(rmControlSimple, compUtil)
798865

799-
prevStructs := abi.getStructs
800-
abi.getStructs = func() *driverABIStructs {
801-
structs := prevStructs()
802-
structs.controlStructs[nvgpu.NV_CONF_COMPUTE_CTRL_CMD_GPU_GET_KEY_ROTATION_STATE] = simpleDriverStruct("NV_CONF_COMPUTE_CTRL_CMD_GPU_GET_KEY_ROTATION_STATE_PARAMS")
803-
return structs
804-
}
866+
prevStructs := abi.getStructs
867+
abi.getStructs = func() *driverABIStructs {
868+
structs := prevStructs()
869+
structs.controlStructs[nvgpu.NV_CONF_COMPUTE_CTRL_CMD_GPU_GET_KEY_ROTATION_STATE] = simpleDriverStruct("NV_CONF_COMPUTE_CTRL_CMD_GPU_GET_KEY_ROTATION_STATE_PARAMS")
870+
return structs
871+
}
805872

806-
return abi
807-
})
873+
return abi
874+
},
875+
},
876+
)
808877

809878
// This version does not belong on any branch, but it is a child of 550.90.07.
810-
_ = addDriverABI(550, 90, 12, "391883846713b9e700af2ae87f8ac671f5527508ce3f9f60058deb363e05162a", v550_90_07)
879+
_ = addDriverABI(
880+
addDriverABIArgs{
881+
major: 550,
882+
minor: 90,
883+
patch: 12,
884+
runfileChecksumX86_64: "391883846713b9e700af2ae87f8ac671f5527508ce3f9f60058deb363e05162a",
885+
runfileChecksumARM64: "0c410aff85b772bdb411d749c23e12ef2658f997e3094c41de8a0495a9fab4b4",
886+
cons: v550_90_07,
887+
},
888+
)
811889

812890
// 550.100 is an intermediate unqualified version from the main branch.
813891
v550_100 := v550_90_07
814892

815893
// The following exist on the "550" branch. They branched off the main
816894
// branch at 550.100.
817-
_ = addDriverABI(550, 127, 05, "d384f34f5d2a896bd7536d3deb6a6d973d8094a3ad485a1c2ee3bf5192086ae9", v550_100)
895+
_ = addDriverABI(addDriverABIArgs{
896+
major: 550,
897+
minor: 127,
898+
patch: 05,
899+
runfileChecksumX86_64: "d384f34f5d2a896bd7536d3deb6a6d973d8094a3ad485a1c2ee3bf5192086ae9",
900+
runfileChecksumARM64: "df0b06a89bc37fc8a8e2a152a9ba5a7de1c70636dab0ae62fd6f94e937847816",
901+
cons: v550_100,
902+
},
903+
)
818904

819905
// 555.42.02 is an intermediate unqualified version.
820906
v555_42_02 := func() *driverABI {
@@ -844,24 +930,45 @@ func Init() {
844930
return abi
845931
}
846932

847-
v560_35_03 := addDriverABI(560, 35, 03, "f2932c92fadd43c5b2341be453fc4f73f0ad7185c26bb7a43fbde81ae29f1fe3", v560_28_03)
848-
v565_57_01 := addDriverABI(565, 57, 01, "6eebe94e585e385e8804f5a74152df414887bf819cc21bd95b72acd0fb182c7a", v560_35_03)
933+
v560_35_03 := addDriverABI(addDriverABIArgs{
934+
major: 560,
935+
minor: 35,
936+
patch: 03,
937+
runfileChecksumX86_64: "f2932c92fadd43c5b2341be453fc4f73f0ad7185c26bb7a43fbde81ae29f1fe3",
938+
runfileChecksumARM64: "b3c64054abd1357a63c5162a337139a2cb3915da96fadbf5a900b6a438df1beb",
939+
cons: v560_28_03,
940+
})
941+
v565_57_01 := addDriverABI(addDriverABIArgs{
942+
major: 565,
943+
minor: 57,
944+
patch: 01,
945+
runfileChecksumX86_64: "6eebe94e585e385e8804f5a74152df414887bf819cc21bd95b72acd0fb182c7a",
946+
runfileChecksumARM64: "68355cdec3531b83b7cbebca5bcee6c3e8bd02a5c2636f4656a108525b2f61f1",
947+
cons: v560_35_03,
948+
})
849949

850-
_ = addDriverABI(570, 86, 15, "87709c19c7401243136bc0ec9e7f147c6803070a11449ae8f0819dee7963f76b", func() *driverABI {
851-
abi := v565_57_01()
852-
abi.allocationClass[nvgpu.TURING_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
853-
abi.allocationClass[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
854-
abi.allocationClass[nvgpu.HOPPER_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
950+
_ = addDriverABI(addDriverABIArgs{
951+
major: 570,
952+
minor: 86,
953+
patch: 15,
954+
runfileChecksumX86_64: "87709c19c7401243136bc0ec9e7f147c6803070a11449ae8f0819dee7963f76b",
955+
runfileChecksumARM64: "a663f81873bafda8313abb5a09f36c593426bb94a8bcc3f2017c79c95bf32978",
956+
cons: func() *driverABI {
957+
abi := v565_57_01()
958+
abi.allocationClass[nvgpu.TURING_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
959+
abi.allocationClass[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
960+
abi.allocationClass[nvgpu.HOPPER_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
855961

856-
prevStructs := abi.getStructs
857-
abi.getStructs = func() *driverABIStructs {
858-
structs := prevStructs()
859-
structs.allocationStructs[nvgpu.TURING_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
860-
structs.allocationStructs[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
861-
structs.allocationStructs[nvgpu.HOPPER_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
862-
return structs
863-
}
864-
return abi
962+
prevStructs := abi.getStructs
963+
abi.getStructs = func() *driverABIStructs {
964+
structs := prevStructs()
965+
structs.allocationStructs[nvgpu.TURING_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
966+
structs.allocationStructs[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
967+
structs.allocationStructs[nvgpu.HOPPER_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
968+
return structs
969+
}
970+
return abi
971+
},
865972
})
866973
})
867974
}
@@ -908,9 +1015,9 @@ func newDriverStruct(paramType reflect.Type, name string) DriverStruct {
9081015

9091016
// ForEachSupportDriver calls f on all supported drivers.
9101017
// Precondition: Init() must have been called.
911-
func ForEachSupportDriver(f func(version DriverVersion, checksum string)) {
1018+
func ForEachSupportDriver(f func(version DriverVersion, checksum_X86_64, checksum_ARM64 string)) {
9121019
for version, abi := range abis {
913-
f(version, abi.checksum)
1020+
f(version, abi.checksumX86_64, abi.checksumARM64)
9141021
}
9151022
}
9161023

@@ -941,12 +1048,12 @@ func SupportedDrivers() []DriverVersion {
9411048

9421049
// ExpectedDriverChecksum returns the expected checksum for a given version.
9431050
// Precondition: Init() must have been called.
944-
func ExpectedDriverChecksum(version DriverVersion) (string, bool) {
1051+
func ExpectedDriverChecksum(version DriverVersion) (string, string, bool) {
9451052
abi, ok := abis[version]
9461053
if !ok {
947-
return "", false
1054+
return "", "", false
9481055
}
949-
return abi.checksum, true
1056+
return abi.checksumX86_64, abi.checksumARM64, true
9501057
}
9511058

9521059
// SupportedIoctls returns the ioctl numbers that are supported by nvproxy at

0 commit comments

Comments
 (0)