Skip to content

Update the gpu driver installer to support ARM drivers #11517

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions pkg/sentry/devices/nvproxy/nvproxy_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,11 @@ func TestInit(t *testing.T) {
func TestAllSupportedHashesPresent(t *testing.T) {
Init()
for version, abi := range abis {
if abi.checksum == "" {
t.Errorf("unexpected empty value for driver %q", version.String())
if abi.checksumX86_64 == "" {
t.Errorf("unexpected empty checksumX86_64 value for driver %q", version.String())
}
if abi.checksumARM64 == "" {
t.Errorf("unexpected empty checksumARM64 value for driver %q", version.String())
}
}
}
Expand Down
227 changes: 167 additions & 60 deletions pkg/sentry/devices/nvproxy/version.go
Original file line number Diff line number Diff line change
Expand Up @@ -107,8 +107,9 @@ type driverABIStructsFunc func() *driverABIStructs
// abiConAndChecksum couples the driver's abiConstructor to the SHA256 checksum of its linux .run
// driver installer file from NVIDIA.
type abiConAndChecksum struct {
cons driverABIFunc
checksum string
cons driverABIFunc
checksumX86_64 string
checksumARM64 string
}

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

// Note: runfileChecksum is the checksum of the .run file of the driver installer for linux from
type addDriverABIArgs struct {
major, minor, patch int
runfileChecksumX86_64 string
runfileChecksumARM64 string
cons driverABIFunc
}

// Note: runfileChecksums are the checksum of the .run file of the driver installer for linux from
// nvidia.
// To add a new version, add in support as normal and add the "addDriverABI" call for your version.
// Run `make sudo TARGETS=//tools/gpu:main ARGS="checksum --version={}"` to get checksum.
func addDriverABI(major, minor, patch int, runfileChecksum string, cons driverABIFunc) driverABIFunc {
func addDriverABI(args addDriverABIArgs) driverABIFunc {
if abis == nil {
abis = make(map[DriverVersion]abiConAndChecksum)
}
version := NewDriverVersion(major, minor, patch)
abis[version] = abiConAndChecksum{cons: cons, checksum: runfileChecksum}
return cons
version := NewDriverVersion(args.major, args.minor, args.patch)
abis[version] = abiConAndChecksum{
cons: args.cons,
checksumX86_64: args.runfileChecksumX86_64,
checksumARM64: args.runfileChecksumARM64,
}
return args.cons
}

// Init initializes abis global map.
Expand Down Expand Up @@ -702,10 +714,41 @@ func Init() {

// The following exist on the "535" branch. They branched off the main
// branch at 535.113.01.
v535_183_01 := addDriverABI(535, 183, 01, "f6707afbdda9407e3cbc2e5128e60bcbcdbf02fae29958c72fafb5d405e8b883", v535_113_01)
v535_183_06 := addDriverABI(535, 183, 06, "c7bb0a0569c5347845479ed4e3e4d885c6ee3b8adf068c3401cdf754d5ba3d3b", v535_183_01)
v535_216_01 := addDriverABI(535, 216, 01, "5ddea1147810012e33967c3181341bcd6624bd3d654c63f845df833b4ece6af7", v535_183_06)
_ = addDriverABI(535, 230, 02, "20cca9118083fcc8083158466e9cb2b616a7922206bcb7296b1fa5cc9af2e0fd", v535_216_01)
v535_183_01 := addDriverABI(
addDriverABIArgs{
major: 535,
minor: 183,
patch: 01,
runfileChecksumX86_64: "f6707afbdda9407e3cbc2e5128e60bcbcdbf02fae29958c72fafb5d405e8b883",
runfileChecksumARM64: "c9d13b6250d24b76ef87a49b179f234564184a9f6d6414184668958b7f6d21e6",
cons: v535_113_01,
})
v535_183_06 := addDriverABI(
addDriverABIArgs{
major: 535,
minor: 183,
patch: 06,
runfileChecksumX86_64: "c7bb0a0569c5347845479ed4e3e4d885c6ee3b8adf068c3401cdf754d5ba3d3b",
runfileChecksumARM64: "af3f72f5e4906805987844636b87ad1132650d05116272824c76dcc3f816d8e9",
cons: v535_183_01,
})
v535_216_01 := addDriverABI(addDriverABIArgs{
major: 535,
minor: 216,
patch: 01,
runfileChecksumX86_64: "5ddea1147810012e33967c3181341bcd6624bd3d654c63f845df833b4ece6af7",
runfileChecksumARM64: "4869ae0345b5892b2a50aed566c8226d3e07813d1190aa466feba5e9e21b33b9",
cons: v535_183_06,
})

_ = addDriverABI(addDriverABIArgs{
major: 535,
minor: 230,
patch: 02,
runfileChecksumX86_64: "20cca9118083fcc8083158466e9cb2b616a7922206bcb7296b1fa5cc9af2e0fd",
runfileChecksumARM64: "ea000e6ff481f55e9bfedbea93b739368c635fe4be6156fdad560524ac7f363b",
cons: v535_216_01,
})

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

v550_54_14 := addDriverABI(550, 54, 14, "8c497ff1cfc7c310fb875149bc30faa4fd26d2237b2cba6cd2e8b0780157cfe3", func() *driverABI {
abi := v550_40_07()
abi.uvmIoctl[nvgpu.UVM_ALLOC_SEMAPHORE_POOL] = uvmHandler(uvmIoctlSimple[nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS_V550], compUtil)
abi.uvmIoctl[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION] = uvmHandler(uvmIoctlHasFrontendFD[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS_V550], compUtil)
v550_54_14 := addDriverABI(
addDriverABIArgs{
major: 550,
minor: 54,
patch: 14,
runfileChecksumX86_64: "8c497ff1cfc7c310fb875149bc30faa4fd26d2237b2cba6cd2e8b0780157cfe3",
runfileChecksumARM64: "b0fae8061633885c24f6b0c047649b46249a3bb44cadffbf658af28f80642c1d",
cons: func() *driverABI {
abi := v550_40_07()
abi.uvmIoctl[nvgpu.UVM_ALLOC_SEMAPHORE_POOL] = uvmHandler(uvmIoctlSimple[nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS_V550], compUtil)
abi.uvmIoctl[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION] = uvmHandler(uvmIoctlHasFrontendFD[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS_V550], compUtil)

prevStructs := abi.getStructs
abi.getStructs = func() *driverABIStructs {
structs := prevStructs()
structs.uvmStructs[nvgpu.UVM_ALLOC_SEMAPHORE_POOL] = driverStructWithName(nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS_V550{}, "UVM_ALLOC_SEMAPHORE_POOL_PARAMS")
structs.uvmStructs[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION] = driverStructWithName(nvgpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS_V550{}, "UVM_MAP_EXTERNAL_ALLOCATION_PARAMS")
return structs
}

return abi
})
prevStructs := abi.getStructs
abi.getStructs = func() *driverABIStructs {
structs := prevStructs()
structs.uvmStructs[nvgpu.UVM_ALLOC_SEMAPHORE_POOL] = driverStructWithName(nvgpu.UVM_ALLOC_SEMAPHORE_POOL_PARAMS_V550{}, "UVM_ALLOC_SEMAPHORE_POOL_PARAMS")
structs.uvmStructs[nvgpu.UVM_MAP_EXTERNAL_ALLOCATION] = driverStructWithName(nvgpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS_V550{}, "UVM_MAP_EXTERNAL_ALLOCATION_PARAMS")
return structs
}
return abi
},
},
)

v550_54_15 := addDriverABI(550, 54, 15, "2e859ae5f912a9a47aaa9b2d40a94a14f6f486b5d3b67c0ddf8b72c1c9650385", v550_54_14)
v550_54_15 := addDriverABI(
addDriverABIArgs{
major: 550,
minor: 54,
patch: 15,
runfileChecksumX86_64: "2e859ae5f912a9a47aaa9b2d40a94a14f6f486b5d3b67c0ddf8b72c1c9650385",
runfileChecksumARM64: "49072d0c36ed85c7d8046776d34886f9ede9a6e4f46d5c7d533e8a8921d94cc1",
cons: v550_54_14,
},
)

v550_90_07 := addDriverABI(550, 90, 07, "51acf579d5a9884f573a1d3f522e7fafa5e7841e22a9cec0b4bbeae31b0b9733", func() *driverABI {
abi := v550_54_15()
abi.controlCmd[nvgpu.NV_CONF_COMPUTE_CTRL_CMD_GPU_GET_KEY_ROTATION_STATE] = ctrlHandler(rmControlSimple, compUtil)
v550_90_07 := addDriverABI(
addDriverABIArgs{
major: 550,
minor: 90,
patch: 07,
runfileChecksumX86_64: "51acf579d5a9884f573a1d3f522e7fafa5e7841e22a9cec0b4bbeae31b0b9733",
runfileChecksumARM64: "b896b76ae465307afc5b269c40bd8ccb279e6ea7d3ecae95534a91ecb1971572",
cons: func() *driverABI {
abi := v550_54_15()
abi.controlCmd[nvgpu.NV_CONF_COMPUTE_CTRL_CMD_GPU_GET_KEY_ROTATION_STATE] = ctrlHandler(rmControlSimple, compUtil)

prevStructs := abi.getStructs
abi.getStructs = func() *driverABIStructs {
structs := prevStructs()
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")
return structs
}
prevStructs := abi.getStructs
abi.getStructs = func() *driverABIStructs {
structs := prevStructs()
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")
return structs
}

return abi
})
return abi
},
},
)

// This version does not belong on any branch, but it is a child of 550.90.07.
_ = addDriverABI(550, 90, 12, "391883846713b9e700af2ae87f8ac671f5527508ce3f9f60058deb363e05162a", v550_90_07)
_ = addDriverABI(
addDriverABIArgs{
major: 550,
minor: 90,
patch: 12,
runfileChecksumX86_64: "391883846713b9e700af2ae87f8ac671f5527508ce3f9f60058deb363e05162a",
runfileChecksumARM64: "0c410aff85b772bdb411d749c23e12ef2658f997e3094c41de8a0495a9fab4b4",
cons: v550_90_07,
},
)

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

// The following exist on the "550" branch. They branched off the main
// branch at 550.100.
_ = addDriverABI(550, 127, 05, "d384f34f5d2a896bd7536d3deb6a6d973d8094a3ad485a1c2ee3bf5192086ae9", v550_100)
_ = addDriverABI(addDriverABIArgs{
major: 550,
minor: 127,
patch: 05,
runfileChecksumX86_64: "d384f34f5d2a896bd7536d3deb6a6d973d8094a3ad485a1c2ee3bf5192086ae9",
runfileChecksumARM64: "df0b06a89bc37fc8a8e2a152a9ba5a7de1c70636dab0ae62fd6f94e937847816",
cons: v550_100,
},
)

// 555.42.02 is an intermediate unqualified version.
v555_42_02 := func() *driverABI {
Expand Down Expand Up @@ -844,24 +930,45 @@ func Init() {
return abi
}

v560_35_03 := addDriverABI(560, 35, 03, "f2932c92fadd43c5b2341be453fc4f73f0ad7185c26bb7a43fbde81ae29f1fe3", v560_28_03)
v565_57_01 := addDriverABI(565, 57, 01, "6eebe94e585e385e8804f5a74152df414887bf819cc21bd95b72acd0fb182c7a", v560_35_03)
v560_35_03 := addDriverABI(addDriverABIArgs{
major: 560,
minor: 35,
patch: 03,
runfileChecksumX86_64: "f2932c92fadd43c5b2341be453fc4f73f0ad7185c26bb7a43fbde81ae29f1fe3",
runfileChecksumARM64: "b3c64054abd1357a63c5162a337139a2cb3915da96fadbf5a900b6a438df1beb",
cons: v560_28_03,
})
v565_57_01 := addDriverABI(addDriverABIArgs{
major: 565,
minor: 57,
patch: 01,
runfileChecksumX86_64: "6eebe94e585e385e8804f5a74152df414887bf819cc21bd95b72acd0fb182c7a",
runfileChecksumARM64: "68355cdec3531b83b7cbebca5bcee6c3e8bd02a5c2636f4656a108525b2f61f1",
cons: v560_35_03,
})

_ = addDriverABI(570, 86, 15, "87709c19c7401243136bc0ec9e7f147c6803070a11449ae8f0819dee7963f76b", func() *driverABI {
abi := v565_57_01()
abi.allocationClass[nvgpu.TURING_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
abi.allocationClass[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
abi.allocationClass[nvgpu.HOPPER_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
_ = addDriverABI(addDriverABIArgs{
major: 570,
minor: 86,
patch: 15,
runfileChecksumX86_64: "87709c19c7401243136bc0ec9e7f147c6803070a11449ae8f0819dee7963f76b",
runfileChecksumARM64: "a663f81873bafda8313abb5a09f36c593426bb94a8bcc3f2017c79c95bf32978",
cons: func() *driverABI {
abi := v565_57_01()
abi.allocationClass[nvgpu.TURING_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
abi.allocationClass[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)
abi.allocationClass[nvgpu.HOPPER_CHANNEL_GPFIFO_A] = allocHandler(rmAllocChannelV570, compUtil)

prevStructs := abi.getStructs
abi.getStructs = func() *driverABIStructs {
structs := prevStructs()
structs.allocationStructs[nvgpu.TURING_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
structs.allocationStructs[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
structs.allocationStructs[nvgpu.HOPPER_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
return structs
}
return abi
prevStructs := abi.getStructs
abi.getStructs = func() *driverABIStructs {
structs := prevStructs()
structs.allocationStructs[nvgpu.TURING_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
structs.allocationStructs[nvgpu.AMPERE_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
structs.allocationStructs[nvgpu.HOPPER_CHANNEL_GPFIFO_A] = driverStructWithName(nvgpu.NV_CHANNEL_ALLOC_PARAMS_V570{}, "NV_CHANNEL_ALLOC_PARAMS")
return structs
}
return abi
},
})
})
}
Expand Down Expand Up @@ -908,9 +1015,9 @@ func newDriverStruct(paramType reflect.Type, name string) DriverStruct {

// ForEachSupportDriver calls f on all supported drivers.
// Precondition: Init() must have been called.
func ForEachSupportDriver(f func(version DriverVersion, checksum string)) {
func ForEachSupportDriver(f func(version DriverVersion, checksum_X86_64, checksum_ARM64 string)) {
for version, abi := range abis {
f(version, abi.checksum)
f(version, abi.checksumX86_64, abi.checksumARM64)
}
}

Expand Down Expand Up @@ -941,12 +1048,12 @@ func SupportedDrivers() []DriverVersion {

// ExpectedDriverChecksum returns the expected checksum for a given version.
// Precondition: Init() must have been called.
func ExpectedDriverChecksum(version DriverVersion) (string, bool) {
func ExpectedDriverChecksum(version DriverVersion) (string, string, bool) {
abi, ok := abis[version]
if !ok {
return "", false
return "", "", false
}
return abi.checksum, true
return abi.checksumX86_64, abi.checksumARM64, true
}

// SupportedIoctls returns the ioctl numbers that are supported by nvproxy at
Expand Down
Loading