add uuid in cudaDeviceProperties (#125083)

jeffdaily · pytorchmergebot · commit f35fe4eaf1e9 · 2024-05-07T01:26:01.000Z
Replaces #99967. Fixes #99903. Pull Request resolved: #125083 Approved by: https://github.com/pruthvistony, https://github.com/albanD, https://github.com/eqy
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -273,6 +273,11 @@ def test_set_per_process_memory_fraction(self):
         tensor.fill_(1)
         self.assertTrue((tensor == 1).all())
 
+    def test_uuid(self):
+        uuid = torch.cuda.get_device_properties(0).uuid
+        assert len(str(uuid)) == 36  # xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
+        assert len(uuid.bytes) == 16
+
     def test_copy_non_blocking(self):
         def _test_copy_non_blocking(a, b):
             event = torch.cuda.Event()
diff --git a/torch/_C/__init__.pyi.in b/torch/_C/__init__.pyi.in
@@ -1891,6 +1891,7 @@ class _CudaDeviceProperties:
     is_multi_gpu_board: _int
     max_threads_per_multi_processor: _int
     gcnArchName: str
+    uuid: str
 
 # Functions related to SDPA
 class _SDPAParams:
diff --git a/torch/csrc/cuda/Module.cpp b/torch/csrc/cuda/Module.cpp
@@ -909,6 +909,44 @@ PyObject* THCPModule_cudaGetSyncDebugMode(PyObject* self, PyObject* noargs) {
 static void registerCudaDeviceProperties(PyObject* module) {
   // Add _cudaDevicePropertires class to torch._C
   auto m = py::handle(module).cast<py::module>();
+  // CUuuid is defined in either cuda.h or driver_types.h
+  // hipified to hipUUID which is defined in hip_runtime_api.h
+  py::class_<CUuuid>(m, "_CUuuid")
+      .def_property_readonly(
+          "bytes", [](const CUuuid& uuid) { return py::bytes(uuid.bytes); })
+      .def("__repr__", [](const CUuuid& uuid) {
+        // UUIDs are a 128-bit label. CUDA and HIP store this as char[16].
+        // For string representation, the code here expands this to
+        // 8-4-4-4-12 hex format, so each byte becomes 2 hex characters.
+        // Size is 16x2 hex characters + 4 hyphens + 1 null byte.
+        constexpr size_t size = sizeof(CUuuid) * 2 + 4 + 1;
+        char device_path_str[size] = {0};
+        snprintf(
+            device_path_str,
+            sizeof(device_path_str),
+            "%02x%02x%02x%02x-"
+            "%02x%02x-"
+            "%02x%02x-"
+            "%02x%02x-"
+            "%02x%02x%02x%02x%02x%02x",
+            (uint8_t)uuid.bytes[0],
+            (uint8_t)uuid.bytes[1],
+            (uint8_t)uuid.bytes[2],
+            (uint8_t)uuid.bytes[3],
+            (uint8_t)uuid.bytes[4],
+            (uint8_t)uuid.bytes[5],
+            (uint8_t)uuid.bytes[6],
+            (uint8_t)uuid.bytes[7],
+            (uint8_t)uuid.bytes[8],
+            (uint8_t)uuid.bytes[9],
+            (uint8_t)uuid.bytes[10],
+            (uint8_t)uuid.bytes[11],
+            (uint8_t)uuid.bytes[12],
+            (uint8_t)uuid.bytes[13],
+            (uint8_t)uuid.bytes[14],
+            (uint8_t)uuid.bytes[15]);
+        return std::string(device_path_str);
+      });
   py::class_<cudaDeviceProp>(m, "_CudaDeviceProperties")
       .def_readonly("name", &cudaDeviceProp::name)
       .def_readonly("major", &cudaDeviceProp::major)
@@ -935,6 +973,7 @@ static void registerCudaDeviceProperties(PyObject* module) {
           &cudaDeviceProp::name
 #endif // USE_ROCM
           )
+      .def_readonly("uuid", &cudaDeviceProp::uuid)
       .def("__repr__", [](const cudaDeviceProp& prop) {
         std::ostringstream stream;
         stream << "_CudaDeviceProperties(name='" << prop.name
@@ -944,7 +983,7 @@ static void registerCudaDeviceProperties(PyObject* module) {
 #endif // USE_ROCM
                << ", total_memory=" << prop.totalGlobalMem / (1024ull * 1024)
                << "MB, multi_processor_count=" << prop.multiProcessorCount
-               << ")";
+               << ", uuid=" << std::string(prop.uuid.bytes, 16) << ")";
         return stream.str();
       });
 
