src/pkg/gpubuild/gpubuild.go - cos/tools - Git at Google

 // Copyright 2025 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 // Package gpubuild provides functionality for compiling and packaging GPU
 // drivers.
 package gpubuild

 import (
 	"bytes"
 	"errors"
 	"fmt"
 	"io/fs"
 	"log"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
 	"strconv"
 	"strings"

 	cosfs "cos.googlesource.com/cos/tools.git/src/pkg/fs"
 	"cos.googlesource.com/cos/tools.git/src/pkg/utils"
 )

 // DriverPackage describes a tarball generated by the driver compilation
 // process. The tarball looks like this:
 //
 // /drivers
 // /drivers/nvidia-uvm.ko
 // /drivers/nv-p2p-dummy.ko
 // /drivers/gdrdrv.ko
 // /drivers/nvidia.ko
 // /drivers/nvidia-peermem.ko
 // /drivers/nvidia-modeset.ko
 // /drivers/nvidia-drm.ko
 // /firmware
 // /firmware/nvidia
 // /firmware/nvidia/580.95.05
 // /firmware/nvidia/580.95.05/gsp_tu10x.bin
 // /firmware/nvidia/580.95.05/gsp_ga10x.bin
 type DriverPackage struct {
 	// Path is the local file path to the driver package.
 	Path string
 	// Version is the Nvidia driver version of the package.
 	Version string
 }

 func nvidiaInstallerLSM(runfilePath string) (string, error) {
 	var outBytes bytes.Buffer
 	cmd := exec.Command(runfilePath, "--lsm")
 	cmd.Stdout = &outBytes
 	cmd.Stderr = os.Stderr
 	if err := cmd.Run(); err != nil {
 		return "", fmt.Errorf("failed to read LSM from installer, see stderr for details: %v", err)
 	}
 	return string(outBytes.Bytes()), nil
 }

 func getArch(lsm string) (string, error) {
 	lines := strings.Split(lsm, "\n")
 	for _, line := range lines {
 		if strings.Contains(line, "Title:") && strings.Contains(line, "Linux-x86_64") {
 			return "x86_64", nil
 		}
 		if strings.Contains(line, "Title:") && strings.Contains(line, "Linux-aarch64") {
 			return "aarch64", nil
 		}
 	}
 	return "", errors.New("could not determine architecture from LSM")
 }

 func getVersion(lsm string) (string, error) {
 	lines := strings.Split(lsm, "\n")
 	for _, line := range lines {
 		if strings.Contains(line, "Version:") {
 			fields := strings.Fields(line)
 			if len(fields) < 2 {
 				return "", errors.New("could not determine version from LSM")
 			}
 			return fields[1], nil
 		}
 	}
 	return "", errors.New("could not determine version from LSM")
 }

 func copyPackageFiles(installerDir, driverDir, firmwareDir string) error {
 	return filepath.WalkDir(installerDir, func(path string, _ fs.DirEntry, err error) error {
 		if err != nil {
 			return err
 		}
 		if filepath.Ext(path) == ".ko" {
 			baseName := filepath.Base(path)
 			dest := filepath.Join(driverDir, baseName)
 			log.Printf("Copying %q to %q", path, dest)
 			if err := utils.CopyFile(path, dest); err != nil {
 				return err
 			}
 			return nil
 		}
 		isFirmware, err := filepath.Match(filepath.Join(filepath.Dir(path), "gsp*.bin"), path)
 		if err != nil {
 			return err
 		}
 		if isFirmware {
 			baseName := filepath.Base(path)
 			dest := filepath.Join(firmwareDir, baseName)
 			log.Printf("Copying %q to %q", path, dest)
 			if err := utils.CopyFile(path, dest); err != nil {
 				return err
 			}
 			return nil
 		}
 		return nil
 	})
 }

 // Compile generates a COS GPU driver package from a runfile.
 //
 // toolchainPath is the local file path to an unpacked toolchain, including
 // both cross-compilers and kernel headers. Cross-compilers should be in bin/,
 // kernel headers should be in a directory that matches usr/src/linux-headers-*.
 //
 // Compiles open source drivers if available, and falls back to closed source
 // drivers if not available.
 //
 // The result driver package will go in outDir as a tarball that looks like
 // nvidia-drivers-${version}.tgz.
 func Compile(runfilePath, toolchainPath, outDir string) (*DriverPackage, error) {
 	tmpDir, err := os.MkdirTemp("", "nvidia-installer")
 	if err != nil {
 		return nil, err
 	}
 	defer os.RemoveAll(tmpDir)
 	runfileExtract := filepath.Join(tmpDir, "runfile")
 	lsm, err := nvidiaInstallerLSM(runfilePath)
 	if err != nil {
 		return nil, err
 	}
 	log.Printf("Compiling runfile with details:\n%v", lsm)
 	driverArch, err := getArch(lsm)
 	if err != nil {
 		return nil, err
 	}
 	driverVersion, err := getVersion(lsm)
 	if err != nil {
 		return nil, err
 	}
 	log.Printf("Extracting runfile to %q", runfileExtract)
 	cmd := exec.Command(runfilePath, "-x", "--target", runfileExtract)
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 	if err := cmd.Run(); err != nil {
 		return nil, err
 	}
 	toolchainBin := filepath.Join(toolchainPath, "bin")
 	log.Print("Locating kernel headers")
 	kernelMatches, err := filepath.Glob(filepath.Join(toolchainPath, "usr/src/linux-headers-*"))
 	if err != nil {
 		return nil, err
 	}
 	if len(kernelMatches) != 1 {
 		return nil, fmt.Errorf("kernel headers directory is ambiguous; found %d directories", len(kernelMatches))
 	}
 	kernelHeaders := kernelMatches[0]
 	log.Printf("Found kernel headers at %q", kernelHeaders)
 	var srcSubDir string
 	if _, err := os.Stat(filepath.Join(runfileExtract, "kernel-open")); !os.IsNotExist(err) {
 		log.Println("Found kernel-open subdir in runfile - compiling open source drivers")
 		srcSubDir = "kernel-open"
 	} else {
 		log.Println("Could not find kernel-open subdir in runfile - compiling closed source drivers")
 		srcSubDir = "kernel"
 	}
 	kernelArch := driverArch
 	if driverArch == "aarch64" {
 		kernelArch = "arm64"
 	}
 	makeArgs := []string{
 		"-j", strconv.Itoa(runtime.NumCPU()),
 		"ARCH=" + kernelArch, "CROSS_COMPILE=" + driverArch + "-cros-linux-gnu-",
 		"CC=" + driverArch + "-cros-linux-gnu-clang", "LD=" + driverArch + "-cros-linux-gnu-ld.lld",
 		"V=1", "SYSSRC=" + kernelHeaders, "IGNORE_CC_MISMATCH=1",
 	}
 	log.Printf("Make args: %v", makeArgs)
 	cmd = exec.Command("make", makeArgs...)
 	cmd.Env = append(cmd.Environ(), "PATH="+toolchainBin+":"+os.Getenv("PATH"))
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 	cmd.Dir = filepath.Join(runfileExtract, srcSubDir)
 	if err := cmd.Run(); err != nil {
 		return nil, fmt.Errorf("GPU driver compilation failed, see stdout/stderr for more info: %v", err)
 	}
 	tarDir, err := os.MkdirTemp("", "nvidia-driver-output")
 	if err != nil {
 		return nil, err
 	}
 	defer os.RemoveAll(tarDir)
 	outputPath := filepath.Join(outDir, "nvidia-drivers-"+driverVersion+".tgz")
 	log.Printf("Creating result tarball at %q", outputPath)
 	tarDriverDir := filepath.Join(tarDir, "drivers")
 	tarFirmwareDir := filepath.Join(tarDir, "firmware/nvidia", driverVersion)
 	if err := os.Mkdir(tarDriverDir, 0755); err != nil {
 		return nil, err
 	}
 	if err := os.MkdirAll(tarFirmwareDir, 0755); err != nil {
 		return nil, err
 	}
 	if err := copyPackageFiles(runfileExtract, tarDriverDir, tarFirmwareDir); err != nil {
 		return nil, err
 	}
 	if err := cosfs.TarDir(tarDir, outputPath); err != nil {
 		return nil, err
 	}
 	log.Println("Runfile compilation succeeded")
 	return &DriverPackage{
 		Path:    outputPath,
 		Version: driverVersion,
 	}, nil
 }
	// Copyright 2025 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	// Package gpubuild provides functionality for compiling and packaging GPU
	// drivers.
	package gpubuild

	import (
	"bytes"
	"errors"
	"fmt"
	"io/fs"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	cosfs "cos.googlesource.com/cos/tools.git/src/pkg/fs"
	"cos.googlesource.com/cos/tools.git/src/pkg/utils"
	)

	// DriverPackage describes a tarball generated by the driver compilation
	// process. The tarball looks like this:
	//
	// /drivers
	// /drivers/nvidia-uvm.ko
	// /drivers/nv-p2p-dummy.ko
	// /drivers/gdrdrv.ko
	// /drivers/nvidia.ko
	// /drivers/nvidia-peermem.ko
	// /drivers/nvidia-modeset.ko
	// /drivers/nvidia-drm.ko
	// /firmware
	// /firmware/nvidia
	// /firmware/nvidia/580.95.05
	// /firmware/nvidia/580.95.05/gsp_tu10x.bin
	// /firmware/nvidia/580.95.05/gsp_ga10x.bin
	type DriverPackage struct {
	// Path is the local file path to the driver package.
	Path string
	// Version is the Nvidia driver version of the package.
	Version string
	}

	func nvidiaInstallerLSM(runfilePath string) (string, error) {
	var outBytes bytes.Buffer
	cmd := exec.Command(runfilePath, "--lsm")
	cmd.Stdout = &outBytes
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
	return "", fmt.Errorf("failed to read LSM from installer, see stderr for details: %v", err)
	}
	return string(outBytes.Bytes()), nil
	}

	func getArch(lsm string) (string, error) {
	lines := strings.Split(lsm, "\n")
	for _, line := range lines {
	if strings.Contains(line, "Title:") && strings.Contains(line, "Linux-x86_64") {
	return "x86_64", nil
	}
	if strings.Contains(line, "Title:") && strings.Contains(line, "Linux-aarch64") {
	return "aarch64", nil
	}
	}
	return "", errors.New("could not determine architecture from LSM")
	}

	func getVersion(lsm string) (string, error) {
	lines := strings.Split(lsm, "\n")
	for _, line := range lines {
	if strings.Contains(line, "Version:") {
	fields := strings.Fields(line)
	if len(fields) < 2 {
	return "", errors.New("could not determine version from LSM")
	}
	return fields[1], nil
	}
	}
	return "", errors.New("could not determine version from LSM")
	}

	func copyPackageFiles(installerDir, driverDir, firmwareDir string) error {
	return filepath.WalkDir(installerDir, func(path string, _ fs.DirEntry, err error) error {
	if err != nil {
	return err
	}
	if filepath.Ext(path) == ".ko" {
	baseName := filepath.Base(path)
	dest := filepath.Join(driverDir, baseName)
	log.Printf("Copying %q to %q", path, dest)
	if err := utils.CopyFile(path, dest); err != nil {
	return err
	}
	return nil
	}
	isFirmware, err := filepath.Match(filepath.Join(filepath.Dir(path), "gsp*.bin"), path)
	if err != nil {
	return err
	}
	if isFirmware {
	baseName := filepath.Base(path)
	dest := filepath.Join(firmwareDir, baseName)
	log.Printf("Copying %q to %q", path, dest)
	if err := utils.CopyFile(path, dest); err != nil {
	return err
	}
	return nil
	}
	return nil
	})
	}

	// Compile generates a COS GPU driver package from a runfile.
	//
	// toolchainPath is the local file path to an unpacked toolchain, including
	// both cross-compilers and kernel headers. Cross-compilers should be in bin/,
	// kernel headers should be in a directory that matches usr/src/linux-headers-*.
	//
	// Compiles open source drivers if available, and falls back to closed source
	// drivers if not available.
	//
	// The result driver package will go in outDir as a tarball that looks like
	// nvidia-drivers-${version}.tgz.
	func Compile(runfilePath, toolchainPath, outDir string) (*DriverPackage, error) {
	tmpDir, err := os.MkdirTemp("", "nvidia-installer")
	if err != nil {
	return nil, err
	}
	defer os.RemoveAll(tmpDir)
	runfileExtract := filepath.Join(tmpDir, "runfile")
	lsm, err := nvidiaInstallerLSM(runfilePath)
	if err != nil {
	return nil, err
	}
	log.Printf("Compiling runfile with details:\n%v", lsm)
	driverArch, err := getArch(lsm)
	if err != nil {
	return nil, err
	}
	driverVersion, err := getVersion(lsm)
	if err != nil {
	return nil, err
	}
	log.Printf("Extracting runfile to %q", runfileExtract)
	cmd := exec.Command(runfilePath, "-x", "--target", runfileExtract)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
	return nil, err
	}
	toolchainBin := filepath.Join(toolchainPath, "bin")
	log.Print("Locating kernel headers")
	kernelMatches, err := filepath.Glob(filepath.Join(toolchainPath, "usr/src/linux-headers-*"))
	if err != nil {
	return nil, err
	}
	if len(kernelMatches) != 1 {
	return nil, fmt.Errorf("kernel headers directory is ambiguous; found %d directories", len(kernelMatches))
	}
	kernelHeaders := kernelMatches[0]
	log.Printf("Found kernel headers at %q", kernelHeaders)
	var srcSubDir string
	if _, err := os.Stat(filepath.Join(runfileExtract, "kernel-open")); !os.IsNotExist(err) {
	log.Println("Found kernel-open subdir in runfile - compiling open source drivers")
	srcSubDir = "kernel-open"
	} else {
	log.Println("Could not find kernel-open subdir in runfile - compiling closed source drivers")
	srcSubDir = "kernel"
	}
	kernelArch := driverArch
	if driverArch == "aarch64" {
	kernelArch = "arm64"
	}
	makeArgs := []string{
	"-j", strconv.Itoa(runtime.NumCPU()),
	"ARCH=" + kernelArch, "CROSS_COMPILE=" + driverArch + "-cros-linux-gnu-",
	"CC=" + driverArch + "-cros-linux-gnu-clang", "LD=" + driverArch + "-cros-linux-gnu-ld.lld",
	"V=1", "SYSSRC=" + kernelHeaders, "IGNORE_CC_MISMATCH=1",
	}
	log.Printf("Make args: %v", makeArgs)
	cmd = exec.Command("make", makeArgs...)
	cmd.Env = append(cmd.Environ(), "PATH="+toolchainBin+":"+os.Getenv("PATH"))
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = filepath.Join(runfileExtract, srcSubDir)
	if err := cmd.Run(); err != nil {
	return nil, fmt.Errorf("GPU driver compilation failed, see stdout/stderr for more info: %v", err)
	}
	tarDir, err := os.MkdirTemp("", "nvidia-driver-output")
	if err != nil {
	return nil, err
	}
	defer os.RemoveAll(tarDir)
	outputPath := filepath.Join(outDir, "nvidia-drivers-"+driverVersion+".tgz")
	log.Printf("Creating result tarball at %q", outputPath)
	tarDriverDir := filepath.Join(tarDir, "drivers")
	tarFirmwareDir := filepath.Join(tarDir, "firmware/nvidia", driverVersion)
	if err := os.Mkdir(tarDriverDir, 0755); err != nil {
	return nil, err
	}
	if err := os.MkdirAll(tarFirmwareDir, 0755); err != nil {
	return nil, err
	}
	if err := copyPackageFiles(runfileExtract, tarDriverDir, tarFirmwareDir); err != nil {
	return nil, err
	}
	if err := cosfs.TarDir(tarDir, outputPath); err != nil {
	return nil, err
	}
	log.Println("Runfile compilation succeeded")
	return &DriverPackage{
	Path: outputPath,
	Version: driverVersion,
	}, nil
	}