Merge branch 'device-capability' into 'master'

Add device compute capability info See merge request nvidia/container-toolkit/gpu-monitoring-tools!2
tkestack · Aug 1, 2019 · 89aebf4 · 89aebf4
2 parents 11d9296 + 8906475
commit 89aebf4
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 11 deletions.
diff --git a/bindings/go/nvml/bindings.go b/bindings/go/nvml/bindings.go
@@ -201,6 +201,20 @@ func deviceGetTopologyCommonAncestor(h1, h2 handle) (*uint, error) {
 	return uintPtr(C.uint(level)), errorString(r)
 }
 
+func (h handle) deviceGetCudaComputeCapability() (*int, *int, error) {
+	var major, minor C.int
+
+	r := C.nvmlDeviceGetCudaComputeCapability(h.dev, &major, &minor)
+	if r != C.NVML_SUCCESS {
+		return nil, nil, errorString(r)
+	}
+
+	intMajor := int(major)
+	intMinor := int(minor)
+
+	return &intMajor, &intMinor, errorString(r)
+}
+
 func (h handle) deviceGetName() (*string, error) {
 	var name [szName]C.char
 

diff --git a/bindings/go/nvml/nvml.go b/bindings/go/nvml/nvml.go
@@ -193,18 +193,24 @@ type PCIInfo struct {
 	Bandwidth *uint
 }
 
+type CudaComputeCapabilityInfo struct {
+	Major *int
+	Minor *int
+}
+
 type Device struct {
 	handle
 
-	UUID        string
-	Path        string
-	Model       *string
-	Power       *uint
-	Memory      *uint64
-	CPUAffinity *uint
-	PCI         PCIInfo
-	Clocks      ClockInfo
-	Topology    []P2PLink
+	UUID                  string
+	Path                  string
+	Model                 *string
+	Power                 *uint
+	Memory                *uint64
+	CPUAffinity           *uint
+	PCI                   PCIInfo
+	Clocks                ClockInfo
+	Topology              []P2PLink
+	CudaComputeCapability CudaComputeCapabilityInfo
 }
 
 type UtilizationInfo struct {
@@ -345,6 +351,8 @@ func NewDevice(idx uint) (device *Device, err error) {
 	assert(err)
 	ccore, cmem, err := h.deviceGetMaxClockInfo()
 	assert(err)
+	cccMajor, cccMinor, err := h.deviceGetCudaComputeCapability()
+	assert(err)
 
 	if minor == nil || busid == nil || uuid == nil {
 		return nil, ErrUnsupportedGPU
@@ -370,6 +378,10 @@ func NewDevice(idx uint) (device *Device, err error) {
 			Cores:  ccore, // MHz
 			Memory: cmem,  // MHz
 		},
+		CudaComputeCapability: CudaComputeCapabilityInfo{
+			Major: cccMajor,
+			Minor: cccMinor,
+		},
 	}
 	if power != nil {
 		*device.Power /= 1000 // W

diff --git a/bindings/go/samples/nvml/deviceInfo/main.go b/bindings/go/samples/nvml/deviceInfo/main.go
@@ -14,15 +14,16 @@ const (
 Model          : {{or .Model "N/A"}}
 Path           : {{.Path}}
 Power          : {{if .Power}}{{.Power}} W{{else}}N/A{{end}}
-Memory	       : {{if .Memory}}{{.Memory}} MiB{{else}}N/A{{end}}
+Memory         : {{if .Memory}}{{.Memory}} MiB{{else}}N/A{{end}}
+CudaComputeCap : {{if .CudaComputeCapability.Major}}{{.CudaComputeCapability.Major}}.{{.CudaComputeCapability.Minor}}{{else}}N/A{{end}}
 CPU Affinity   : {{if .CPUAffinity}}NUMA node{{.CPUAffinity}}{{else}}N/A{{end}}
 Bus ID         : {{.PCI.BusID}}
 BAR1           : {{if .PCI.BAR1}}{{.PCI.BAR1}} MiB{{else}}N/A{{end}}
 Bandwidth      : {{if .PCI.Bandwidth}}{{.PCI.Bandwidth}} MB/s{{else}}N/A{{end}}
 Cores          : {{if .Clocks.Cores}}{{.Clocks.Cores}} MHz{{else}}N/A{{end}}
 Memory         : {{if .Clocks.Memory}}{{.Clocks.Memory}} MHz{{else}}N/A{{end}}
 P2P Available  : {{if not .Topology}}None{{else}}{{range .Topology}}
-    	       	 {{.BusID}} - {{(.Link.String)}}{{end}}{{end}}
+                      {{.BusID}} - {{(.Link.String)}}{{end}}{{end}}
 ---------------------------------------------------------------------
 `
 )