jiggablend/internal/runner/sandbox/gpu.go

package sandbox

import (
	"os"
	"path/filepath"
	"strings"
)

// Bind describes a host path to expose inside the sandbox.
type Bind struct {
	Host string
	// Dest is the path inside the sandbox (usually same as Host for absolute paths).
	Dest string
	// ReadOnly is true for library trees and Blender installs.
	ReadOnly bool
	// Dev is true for device nodes / device trees (podman --device or -v for dirs).
	Dev bool
	// Optional means skip if host path missing (no error).
	Optional bool
}

// GPUMounts returns host paths to expose for Cycles GPU backends based on detection.
// forceCPU skips GPU-specific devices/libs (still may include generic /dev via backend).
func GPUMounts(hasAMD, hasNVIDIA, hasIntel, forceCPU bool) []Bind {
	if forceCPU {
		return nil
	}
	var binds []Bind

	// DRM render nodes used by AMD, Intel, and some NVIDIA EGL paths.
	if hasAMD || hasNVIDIA || hasIntel {
		binds = append(binds, Bind{Host: "/dev/dri", Dest: "/dev/dri", Dev: true, Optional: true})
	}

	if hasAMD {
		binds = append(binds, Bind{Host: "/dev/kfd", Dest: "/dev/kfd", Dev: true, Optional: true})
		// Common ROCm install layouts
		for _, p := range []string{"/opt/rocm", "/usr/share/libdrm"} {
			binds = append(binds, Bind{Host: p, Dest: p, ReadOnly: true, Optional: true})
		}
		// Distro ROCm / amdgpu userspace libs often live under multiarch paths
		for _, p := range rocmLibHints() {
			binds = append(binds, Bind{Host: p, Dest: p, ReadOnly: true, Optional: true})
		}
	}

	if hasNVIDIA {
		for _, name := range nvidiaDeviceNames() {
			p := filepath.Join("/dev", name)
			binds = append(binds, Bind{Host: p, Dest: p, Dev: true, Optional: true})
		}
		for _, p := range nvidiaLibHints() {
			binds = append(binds, Bind{Host: p, Dest: p, ReadOnly: true, Optional: true})
		}
	}

	if hasIntel {
		for _, p := range intelLibHints() {
			binds = append(binds, Bind{Host: p, Dest: p, ReadOnly: true, Optional: true})
		}
	}

	return uniqueBinds(binds)
}

func nvidiaDeviceNames() []string {
	// Fixed control nodes + any /dev/nvidiaN
	names := []string{"nvidiactl", "nvidia-uvm", "nvidia-uvm-tools", "nvidia-modeset"}
	entries, err := os.ReadDir("/dev")
	if err != nil {
		return names
	}
	for _, e := range entries {
		n := e.Name()
		if strings.HasPrefix(n, "nvidia") && !containsString(names, n) {
			names = append(names, n)
		}
	}
	return names
}

func nvidiaLibHints() []string {
	hints := []string{
		"/usr/lib/wsl/lib", // WSL NVIDIA
		"/usr/local/nvidia",
		"/usr/local/cuda",
		"/usr/lib/x86_64-linux-gnu",
		"/usr/lib64",
	}
	// Driver stores often under /usr/lib/libcuda* — parent dirs already covered.
	// Also scan common multiarch for libcuda.so
	for _, dir := range []string{"/usr/lib", "/usr/lib64", "/usr/lib/x86_64-linux-gnu", "/lib", "/lib64"} {
		if matchesAny(dir, "libcuda.so*") || matchesAny(dir, "libnvidia-*.so*") {
			hints = append(hints, dir)
		}
	}
	return hints
}

func rocmLibHints() []string {
	hints := []string{
		"/usr/lib/x86_64-linux-gnu",
		"/usr/lib64",
		"/opt/amdgpu",
	}
	// ROCm versioned trees under /opt/rocm-*
	if entries, err := os.ReadDir("/opt"); err == nil {
		for _, e := range entries {
			if e.IsDir() && strings.HasPrefix(e.Name(), "rocm") {
				hints = append(hints, filepath.Join("/opt", e.Name()))
			}
		}
	}
	return hints
}

func intelLibHints() []string {
	return []string{
		"/usr/lib/x86_64-linux-gnu",
		"/usr/lib64",
		"/usr/lib/intel-opencl",
		"/etc/OpenCL",
	}
}

func matchesAny(dir, glob string) bool {
	m, err := filepath.Glob(filepath.Join(dir, glob))
	return err == nil && len(m) > 0
}

func containsString(ss []string, s string) bool {
	for _, x := range ss {
		if x == s {
			return true
		}
	}
	return false
}

func uniqueBinds(in []Bind) []Bind {
	seen := make(map[string]bool)
	var out []Bind
	for _, b := range in {
		key := b.Host + "|" + b.Dest + "|"
		if b.Dev {
			key += "d"
		}
		if b.ReadOnly {
			key += "r"
		}
		if seen[key] {
			continue
		}
		seen[key] = true
		out = append(out, b)
	}
	return out
}

// ExistingBinds filters to binds whose host path exists (or non-optional always kept for error reporting).
func ExistingBinds(binds []Bind) []Bind {
	var out []Bind
	for _, b := range binds {
		if pathExists(b.Host) {
			out = append(out, b)
			continue
		}
		if !b.Optional {
			out = append(out, b) // caller may error
		}
	}
	return out
}