Support optional subsystem ID matching in device filters

Signed-off-by: Karthik Vetrivel <kvetrivel@nvidia.com>

Author: mattwittwer

api/spec/v1/helpers.go

@@ -38,7 +38,7 @@ func (ms *MigConfigSpec) MatchesDeviceFilter(deviceID types.DeviceID) bool {
 
 	for _, df := range deviceFilter {
 		newDeviceID, _ := types.NewDeviceIDFromString(df)
-		if newDeviceID == deviceID {
+		if newDeviceID.Matches(deviceID) {
 			return true
 		}
 	}

api/spec/v1/helpers_test.go

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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 v1
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/mig-parted/pkg/types"
+)
+
+func TestMigConfigSpecMatchesDeviceFilter(t *testing.T) {
+	h100 := types.NewDeviceIDWithSubsystem(0x2330, 0x10DE, 0x16C0, 0x10DE)
+
+	testCases := []struct {
+		name   string
+		filter interface{}
+		want   bool
+	}{
+		{
+			name:   "empty filter matches all devices",
+			filter: "",
+			want:   true,
+		},
+		{
+			name:   "primary-only filter matches subsystem-aware hardware",
+			filter: "0x233010DE",
+			want:   true,
+		},
+		{
+			name:   "matching subsystem filter matches hardware",
+			filter: "0x233010DE:0x16C010DE",
+			want:   true,
+		},
+		{
+			name:   "sibling subsystem filter does not match hardware",
+			filter: "0x233010DE:0x16C110DE",
+			want:   false,
+		},
+		{
+			name:   "multiple filters stop at a matching subsystem entry",
+			filter: []string{"0x233010DE:0x16C110DE", "0x233010DE:0x16C010DE"},
+			want:   true,
+		},
+		{
+			name:   "malformed filters with extra separators do not match",
+			filter: "0x233010DE:0x16C010DE:extra",
+			want:   false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			spec := MigConfigSpec{
+				DeviceFilter: tc.filter,
+				Devices:      "all",
+			}
+
+			require.Equal(t, tc.want, spec.MatchesDeviceFilter(h100))
+		})
+	}
+}

cmd/nvidia-mig-parted/util/device.go

@@ -67,7 +67,7 @@ func pciVisitGPUs(visit func(*nvpci.NvidiaPCIDevice) error) error {
 func pciGetGPUDeviceIDs() ([]types.DeviceID, error) {
 	var ids []types.DeviceID
 	err := pciVisitGPUs(func(gpu *nvpci.NvidiaPCIDevice) error {
-		ids = append(ids, types.NewDeviceID(gpu.Device, gpu.Vendor))
+		ids = append(ids, types.NewDeviceIDWithSubsystem(gpu.Device, gpu.Vendor, gpu.SubsystemDevice, gpu.SubsystemVendor))
 		return nil
 	})
 	if err != nil {
@@ -91,7 +91,7 @@ func nvmlGetGPUDeviceIDs() ([]types.DeviceID, error) {
 			return nil
 		}
 
-		ids = append(ids, types.NewDeviceID(gpu.Device, gpu.Vendor))
+		ids = append(ids, types.NewDeviceIDWithSubsystem(gpu.Device, gpu.Vendor, gpu.SubsystemDevice, gpu.SubsystemVendor))
 		return nil
 	})
 	if err != nil {

pkg/mig/config/known_configs.go

@@ -23,8 +23,8 @@ import (
 	"github.com/NVIDIA/mig-parted/pkg/types"
 )
 
-const (
-	A100_SXM4_40GB types.DeviceID = 0x20B010DE
+var (
+	A100_SXM4_40GB = types.NewDeviceID(0x20B0, 0x10DE)
 )
 
 type LoopControl int

pkg/mig/discovery/discovery.go

@@ -32,10 +32,10 @@ import (
 // could be discovered.
 var ErrNoProfilesDiscovered = errors.New("no MIG profiles discovered for MIG-capable GPUs")
 
-const (
+var (
 	// deviceIDA30 is the PCI device ID for A30-24GB GPUs.
 	// NVML has a bug that reports incorrect profiles for this GPU.
-	deviceIDA30 = 0x20B710DE
+	deviceIDA30 = types.NewDeviceID(0x20B7, 0x10DE)
 )
 
 // ProfileInfo represents a discovered MIG profile with its metadata
@@ -65,9 +65,9 @@ func isCIProfile(c, g int) bool {
 func getDeviceID(dev nvdev.Device) (types.DeviceID, error) {
 	pciInfo, ret := nvml.Device(dev).GetPciInfo()
 	if ret != nvml.SUCCESS {
-		return 0, fmt.Errorf("failed to get PCI info: %v", ret)
+		return types.DeviceID{}, fmt.Errorf("failed to get PCI info: %v", ret)
 	}
-	return types.DeviceID(pciInfo.PciDeviceId), nil
+	return types.NewDeviceID(uint16(pciInfo.PciDeviceId>>16), uint16(pciInfo.PciDeviceId)), nil
 }
 
 // getHardcodedA30Profiles returns hardcoded MIG profiles for A30 GPUs.
@@ -132,7 +132,7 @@ func (d *discoverer) discoverProfiles() (DeviceProfiles, error) {
 		// Check for A30 - use hardcoded profiles due to NVML bug where
 		// GetGpuInstanceProfileInfo returns incorrect InstanceCount values.
 		// The hardcoded values match the A30 MIG profiles from nvidia-smi.
-		if uint32(deviceID) == deviceIDA30 {
+		if deviceIDA30.Matches(deviceID) {
 			log.Infof("Device %d is A30 (DeviceID: %s), using hardcoded profiles due to NVML bug",
 				i, deviceID.String())
 			result[i] = getHardcodedA30Profiles(deviceID)

pkg/mig/discovery/discovery_test.go

@@ -23,8 +23,6 @@ import (
 	"github.com/NVIDIA/go-nvml/pkg/nvml/mock/dgxa100"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-
-	"github.com/NVIDIA/mig-parted/pkg/types"
 )
 
 func TestDiscoverProfiles(t *testing.T) {
@@ -73,7 +71,7 @@ func TestIsCIProfile(t *testing.T) {
 }
 
 func TestGetHardcodedA30Profiles(t *testing.T) {
-	deviceID := types.DeviceID(deviceIDA30)
+	deviceID := deviceIDA30
 	profiles := getHardcodedA30Profiles(deviceID)
 
 	require.Len(t, profiles, 5)

pkg/types/device.go

@@ -19,36 +19,112 @@ package types
 import (
 	"fmt"
 	"strconv"
+	"strings"
 )
 
 // DeviceID represents a GPU Device ID as read from a GPUs PCIe config space.
-type DeviceID uint32
+type DeviceID struct {
+	Device          uint16
+	Vendor          uint16
+	SubsystemDevice uint16
+	SubsystemVendor uint16
+	HasSubsystem    bool
+}
 
 // NewDeviceID constructs a new 'DeviceID' from the device and vendor values pulled from a GPUs PCIe config space.
 func NewDeviceID(device, vendor uint16) DeviceID {
-	return DeviceID((uint32(device) << 16) | uint32(vendor))
+	return DeviceID{
+		Device: device,
+		Vendor: vendor,
+	}
+}
+
+// NewDeviceIDWithSubsystem constructs a new 'DeviceID' with subsystem values.
+func NewDeviceIDWithSubsystem(device, vendor, subDevice, subVendor uint16) DeviceID {
+	return DeviceID{
+		Device:          device,
+		Vendor:          vendor,
+		SubsystemDevice: subDevice,
+		SubsystemVendor: subVendor,
+		HasSubsystem:    true,
+	}
 }
 
 // NewDeviceIDFromString constructs a 'DeviceID' from its string representation.
 func NewDeviceIDFromString(str string) (DeviceID, error) {
-	deviceID, err := strconv.ParseInt(str, 0, 32)
+	parts := strings.Split(str, ":")
+	if len(parts) > 2 {
+		return DeviceID{}, fmt.Errorf(
+			"invalid DeviceID format '%v': expected '<devicevendor>' or '<devicevendor>:<subdevicevendor>'",
+			str,
+		)
+	}
+
+	deviceIDRaw, err := strconv.ParseUint(parts[0], 0, 32)
 	if err != nil {
-		return 0, fmt.Errorf("unable to create DeviceID from string '%v': %v", str, err)
+		return DeviceID{}, fmt.Errorf("unable to create DeviceID from string '%v': %v", str, err)
+	}
+
+	device, vendor := splitRawDeviceID(deviceIDRaw)
+
+	deviceID := DeviceID{
+		Device: device,
+		Vendor: vendor,
+	}
+
+	if len(parts) == 2 {
+		subIDRaw, err := strconv.ParseUint(parts[1], 0, 32)
+		if err != nil {
+			return DeviceID{}, fmt.Errorf("unable to create Subsystem from string '%v': %v", str, err)
+		}
+
+		subsystemDevice, subsystemVendor := splitRawDeviceID(subIDRaw)
+
+		deviceID.SubsystemDevice = subsystemDevice
+		deviceID.SubsystemVendor = subsystemVendor
+		deviceID.HasSubsystem = true
 	}
-	return DeviceID(deviceID), nil
+
+	return deviceID, nil
 }
 
 // String returns a 'DeviceID' as a string.
 func (d DeviceID) String() string {
-	return fmt.Sprintf("0x%X", uint32(d))
+	primary := fmt.Sprintf("0x%04X%04X", d.Device, d.Vendor)
+	if d.HasSubsystem {
+		return fmt.Sprintf("%s:0x%04X%04X", primary, d.SubsystemDevice, d.SubsystemVendor)
+	}
+	return primary
 }
 
 // GetVendor returns the 'vendor' portion of a 'DeviceID'.
 func (d DeviceID) GetVendor() uint16 {
-	return uint16(d)
+	return d.Vendor
 }
 
 // GetDevice returns the 'device' portion of a 'DeviceID'.
 func (d DeviceID) GetDevice() uint16 {
-	return uint16(d >> 16)
+	return d.Device
+}
+
+// Matches checks if a hardware GPU matches the DeviceID filter.
+// If the filter has a subsystem defined, it requires an exact match on all 4 components.
+// Otherwise, it only matches on the primary device and vendor IDs.
+func (filter DeviceID) Matches(hardware DeviceID) bool {
+	if filter.Device != hardware.Device || filter.Vendor != hardware.Vendor {
+		return false
+	}
+
+	if filter.HasSubsystem {
+		if filter.SubsystemDevice != hardware.SubsystemDevice || filter.SubsystemVendor != hardware.SubsystemVendor {
+			return false
+		}
+	}
+
+	return true
+}
+
+func splitRawDeviceID(raw uint64) (uint16, uint16) {
+	// nolint:gosec // raw is parsed as a 32-bit value, so both halves fit in a uint16
+	return uint16(raw >> 16), uint16(raw & 0xFFFF)
 }