From 4523b2e35ddf45dfe22caabe2fa1fc0810c5f1b2 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Mon, 10 Mar 2025 13:51:00 +0200
Subject: [PATCH] [no-relnote] Add function to filter nvcaps by GPU

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 internal/nvcaps/nvcaps.go | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/internal/nvcaps/nvcaps.go b/internal/nvcaps/nvcaps.go
index 48d98ccf..8f6e272a 100644
--- a/internal/nvcaps/nvcaps.go
+++ b/internal/nvcaps/nvcaps.go
@@ -39,7 +39,13 @@ const (
 // MigMinor represents the minor number of a MIG device
 type MigMinor int
 
-// MigCap represents the path to a MIG cap file
+// MigCap represents the path to a MIG cap file.
+// These are listed in /proc/driver/nvidia-caps/mig-minors and have one of the
+// follown forms:
+//   - config
+//   - monitor
+//   - gpu{{ .gpuIndex }}/gi{{ .gi }}/access
+//   - gpu{{ .gpuIndex }}/gi{{ .gi }}/ci {{ .ci }}/access
 type MigCap string
 
 // MigCaps stores a map of MIG cap file paths to MIG minors
@@ -57,6 +63,31 @@ func NewComputeInstanceCap(gpu, gi, ci int) MigCap {
 	return MigCap(fmt.Sprintf("gpu%d/gi%d/ci%d/access", gpu, gi, ci))
 }
 
+// FilterForGPU limits the MIG Caps to those associated with a particular GPU.
+func (m MigCaps) FilterForGPU(gpu int) MigCaps {
+	if m == nil {
+		return nil
+	}
+	filtered := make(MigCaps)
+	for gi := 0; ; gi++ {
+		giCap := NewGPUInstanceCap(gpu, gi)
+		giMinor, exist := m[giCap]
+		if !exist {
+			break
+		}
+		filtered[giCap] = giMinor
+		for ci := 0; ; ci++ {
+			ciCap := NewComputeInstanceCap(gpu, gi, ci)
+			ciMinor, exist := m[ciCap]
+			if !exist {
+				break
+			}
+			filtered[ciCap] = ciMinor
+		}
+	}
+	return filtered
+}
+
 // GetCapDevicePath returns the path to the cap device for the specified cap.
 // An error is returned if the cap is invalid.
 func (m MigCaps) GetCapDevicePath(cap MigCap) (string, error) {