Improve NVML wrapper generation.

Author: maleadt

res/wrap/libnvml_prologue.jl

@@ -26,3 +26,10 @@ end
 
     return
 end
+
+macro NVML_STRUCT_VERSION(typename, version)
+    struct_typename = Symbol("nvml$(String(typename))_v$(version)_t")
+    struct_type = getfield(__module__, struct_typename)
+    struct_version = UInt32(sizeof(struct_type)) | (UInt32(version) << 24)
+    return :($struct_version)
+end

res/wrap/nvml.toml

@@ -3,35 +3,7 @@ library_name = "libnvml()"
 output_file_path = "../../lib/nvml/libnvml.jl"
 prologue_file_path = "./libnvml_prologue.jl"
 
-output_ignorelist = [
-    # we can't handle NVML_STRUCT_VERSION
-    "^nvmlMemory_v2$",
-    "^nvmlGpuInstanceProfileInfo_v2$",
-    "^nvmlComputeInstanceProfileInfo_v2$",
-    "^nvmlPowerValue_v2$",
-    "^nvmlProcessDetailList_v1$",
-    "^nvmlC2cModeInfo_v1$",
-    "^nvmlGpuInstanceProfileInfo_v3$",
-    "^nvmlComputeInstanceProfileInfo_v3$",
-    "nvmlPciInfoExt_v1",
-    "nvmlVgpuHeterogeneousMode_v1",
-    "nvmlVgpuPlacementId_v1",
-    "nvmlVgpuPlacementList_v1",
-    "nvmlVgpuInstancesUtilizationInfo_v1",
-    "nvmlVgpuProcessesUtilizationInfo_v1",
-    "nvmlProcessesUtilizationInfo_v1",
-    "nvmlEccSramErrorStatus_v1",
-    "nvmlSystemConfComputeSettings_v1",
-    "nvmlConfComputeSetKeyRotationThresholdInfo_v1",
-    "nvmlConfComputeGetKeyRotationThresholdInfo_v1",
-    "nvmlGpuFabricInfo_v2",
-    "NVML_DEVICE_PCI_BUS_ID_LEGACY_FMT",
-    "NVML_DEVICE_PCI_BUS_ID_FMT",
-    "nvmlClockOffset_v1",
-    "nvmlDeviceCapabilities_v1",
-    "nvmlVgpuTypeBar1Info_v1",
-    "nvmlSystemDriverBranchInfo_v1",
-]
+output_ignorelist = []
 
 
 [codegen]
@@ -42,7 +14,7 @@ always_NUL_terminated_string = true
 [api]
 checked_rettypes = [ "nvmlReturn_t" ]
 
-[api.nvmlInit_v2]
+[api.nvmlInit]
 needs_context = false
 
 [api.nvmlInitWithFlags]

res/wrap/wrap.jl

@@ -7,9 +7,6 @@
 #
 # To update the types of arguments, add `api.<function>.argtypes` to a library's TOML file.
 
-# TODO
-# - deal with NVML's `NVML_STRUCT_VERSION` (workaround: we ignore these symbols)
-
 using Clang
 using Clang.Generators
 
@@ -126,7 +123,7 @@ function rewriter!(ctx, options)
                 lhs, rhs = expr.args
                 if rhs isa Expr && rhs.head == :call
                     name = string(rhs.args[1])
-                    if endswith(name, "STRUCT_SIZE")
+                    if endswith(name, "STRUCT_SIZE") || endswith(name, "STRUCT_VERSION")
                         rhs.head = :macrocall
                         rhs.args[1] = Symbol("@", name)
                         insert!(rhs.args, 2, nothing)
@@ -182,14 +179,20 @@ function rewriter!(ctx, options)
                     push!(names, fn[1:end-3])
                 end
 
+                # versioned functions are aliased to the unversioned name, so we can
+                # reuse the same type rewrites.
+                if occursin("_v", fn)
+                    push!(names, replace(fn, r"_v\d+" => ""))
+                end
+
                 # look for a template rewrite: many libraries have very similar functions,
                 # e.g., `cublas[SDHCZ]gemm`, for which we can use the same type rewrites
                 # registered as `cublas𝕏gemm` template with `T` and `S` placeholders.
                 for name in copy(names), (typcode,(T,S)) in ["S"=>("Cfloat","Cfloat"),
-                                                              "D"=>("Cdouble","Cdouble"),
-                                                              "H"=>("Float16","Float16"),
-                                                              "C"=>("cuComplex","Cfloat"),
-                                                              "Z"=>("cuDoubleComplex","Cdouble")]
+                                                             "D"=>("Cdouble","Cdouble"),
+                                                             "H"=>("Float16","Float16"),
+                                                             "C"=>("cuComplex","Cfloat"),
+                                                             "Z"=>("cuDoubleComplex","Cdouble")]
                     idx = findfirst(typcode, name)
                     while idx !== nothing
                         template_name = name[1:idx.start-1] * "𝕏" * name[idx.stop+1:end]
@@ -265,7 +268,8 @@ function main(name="all")
     end
 
     if name == "all" || name == "nvml"
-        wrap("nvml", ["$cuda/nvml.h"]; include_dirs=[cuda])
+        wrap("nvml", ["$cuda/nvml.h"]; include_dirs=[cuda],
+             defines=["NVML_NO_UNVERSIONED_FUNC_DEFS=0"])
     end
 
     if name == "all" || name == "cupti"