[ROCm] Added unit test to test the cuda_pluggable allocator (#154135)

pytorchbot · amd-sriram · jithunnair-amd · web-flow · commit 941732c8298c · 2025-05-22T15:23:24.000-04:00
[ROCm] Added unit test to test the cuda_pluggable allocator (#154041) Added unit test to include the cuda_pluggable allocator and replicate the apex setup.py to build nccl_allocator extension This test to check if this commit #152179 helps to build the cuda pluggable allocator in Rocm/Apex Pull Request resolved: #154041 Approved by: https://github.com/atalman, https://github.com/jeffdaily (cherry picked from commit c2660d2) Co-authored-by: skishore <sriramkumar.kishorekumar@amd.com> Co-authored-by: Jithun Nair <jithun.nair@amd.com>
diff --git a/test/test_cpp_extensions_jit.py b/test/test_cpp_extensions_jit.py
@@ -1163,6 +1163,48 @@ def test_aoti_torch_call_dispatcher(self):
         self.assertEqual(abs_t, torch.abs(t))
         self.assertEqual(floor_t, torch.floor(t))
 
+    @unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
+    def test_cuda_pluggable_allocator_include(self):
+        """
+        This method creates a minimal example to replicate the apex setup.py to build nccl_allocator extension
+        """
+
+        # the cpp source includes CUDAPluggableAllocator and has an empty exported function
+        cpp_source = """
+                #include <torch/csrc/cuda/CUDAPluggableAllocator.h>
+                #include <torch/extension.h>
+                int get_nccl_allocator() {
+                    return 0;
+                }
+                PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+                    m.def("get_nccl_allocator", []() { return get_nccl_allocator(); });
+                }
+                """
+
+        build_dir = tempfile.mkdtemp()
+        src_path = os.path.join(build_dir, "NCCLAllocator.cpp")
+
+        with open(src_path, mode="w") as f:
+            f.write(cpp_source)
+
+        # initially success is false
+        success = False
+        try:
+            # try to build the module
+            torch.utils.cpp_extension.load(
+                name="nccl_allocator",
+                sources=src_path,
+                verbose=True,
+                with_cuda=True,
+            )
+            # set success as true if built successfully
+            success = True
+        except Exception as e:
+            print(f"Failed to load the module: {e}")
+
+        # test if build was successful
+        self.assertEqual(success, True)
+
 
 if __name__ == "__main__":
     common.run_tests()
