Added __HIP_PLATFORM_AMD__=1 (microsoft#4570)

baodii · Nov 7, 2023 · 72a0401 · 72a0401
1 parent d67b21a
commit 72a0401
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/op_builder/builder.py b/op_builder/builder.py
@@ -506,6 +506,9 @@ def jit_load(self, verbose=True):
                 cxx_args.append("-DBF16_AVAILABLE")
                 nvcc_args.append("-DBF16_AVAILABLE")
 
+        if self.is_rocm_pytorch():
+            cxx_args.append("-D__HIP_PLATFORM_AMD__=1")
+
         op_module = load(name=self.name,
                          sources=self.strip_empty_entries(sources),
                          extra_include_paths=self.strip_empty_entries(extra_include_paths),