From 8906475be65c348e3bb49b03e80d063d0016237a Mon Sep 17 00:00:00 2001
From: Jordan Jacobelli <jjacobelli@nvidia.com>
Date: Wed, 31 Jul 2019 09:25:14 -0700
Subject: [PATCH] Add device compute capability info

Signed-off-by: Jordan Jacobelli <jjacobelli@nvidia.com>
---
 bindings/go/nvml/bindings.go                | 14 ++++++++++
 bindings/go/nvml/nvml.go                    | 30 ++++++++++++++-------
 bindings/go/samples/nvml/deviceInfo/main.go |  5 ++--
 3 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/bindings/go/nvml/bindings.go b/bindings/go/nvml/bindings.go
index 87961e2..e68d4d1 100644
--- 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
index ec10e11..05c397f 100644
--- 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
index 7d1aa16..ec254a0 100644
--- a/bindings/go/samples/nvml/deviceInfo/main.go
+++ b/bindings/go/samples/nvml/deviceInfo/main.go
@@ -14,7 +14,8 @@ 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}}
@@ -22,7 +23,7 @@ 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}}
 ---------------------------------------------------------------------
 `
 )