Repository: GridTools/gt4py
Pull request #1661 (2024-09-25 09:12:27)

fext[next]: HIP support
Introduces the HIP language in the CMake setup.
`gtfn_gpu` uses the implementation (CUDA or HIP) as detected by cupy.


Patch:
diff --git a/src/gt4py/next/otf/binding/nanobind.py b/src/gt4py/next/otf/binding/nanobind.py
index 27f788d2..24913a13 100644
--- a/src/gt4py/next/otf/binding/nanobind.py
+++ b/src/gt4py/next/otf/binding/nanobind.py
@@ -188,7 +188,7 @@ def create_bindings(
     program_source
         The program source for which the bindings are created
     """
-    if program_source.language not in [languages.Cpp, languages.Cuda]:
+    if program_source.language not in [languages.CPP, languages.CUDA, languages.HIP]:
         raise ValueError(
             f"Can only create bindings for C++ program sources, received '{program_source.language}'."
         )
diff --git a/src/gt4py/next/otf/compilation/build_systems/cmake.py b/src/gt4py/next/otf/compilation/build_systems/cmake.py
index 14377d1b..33212848 100644
--- a/src/gt4py/next/otf/compilation/build_systems/cmake.py
+++ b/src/gt4py/next/otf/compilation/build_systems/cmake.py
@@ -22,7 +22,9 @@ from gt4py.next.otf.compilation.build_systems import cmake_lists
 @dataclasses.dataclass
 class CMakeFactory(
     compiler.BuildSystemProjectGenerator[
-        languages.Cpp | languages.Cuda, languages.LanguageWithHeaderFilesSettings, languages.Python
+        languages.CPP | languages.CUDA | languages.HIP,
+        languages.LanguageWithHeaderFilesSettings,
+        languages.Python,
     ]
 ):
     """Create a CMakeProject from a ``CompilableSource`` stage object with given CMake settings."""
@@ -34,7 +36,7 @@ class CMakeFactory(
     def __call__(
         self,
         source: stages.CompilableSource[
-            languages.Cpp | languages.Cuda,
+            languages.CPP | languages.CUDA | languages.HIP,
             languages.LanguageWithHeaderFilesSettings,
             languages.Python,
         ],
@@ -48,8 +50,8 @@ class CMakeFactory(
         header_name = f"{name}.{source.program_source.language_settings.header_extension}"
         bindings_name = f"{name}_bindings.{source.program_source.language_settings.file_extension}"
         cmake_languages = [cmake_lists.Language(name="CXX")]
-        if source.program_source.language is languages.Cuda:
-            cmake_languages = [*cmake_languages, cmake_lists.Language(name="CUDA")]
+        if (src_lang := source.program_source.language) in [languages.CUDA, languages.HIP]:
+            cmake_languages = [*cmake_languages, cmake_lists.Language(name=src_lang.__name__)]
         cmake_lists_src = cmake_lists.generate_cmakelists_source(
             name, source.library_deps, [header_name, bindings_name], languages=cmake_languages
         )
@@ -70,7 +72,7 @@ class CMakeFactory(
 @dataclasses.dataclass
 class CMakeProject(
     stages.BuildSystemProject[
-        languages.Cpp, languages.LanguageWithHeaderFilesSettings, languages.Python
+        languages.CPP, languages.LanguageWithHeaderFilesSettings, languages.Python
     ]
 ):
     """
diff --git a/src/gt4py/next/otf/compilation/build_systems/cmake_lists.py b/src/gt4py/next/otf/compilation/build_systems/cmake_lists.py
index ae63d661..0533adac 100644
--- a/src/gt4py/next/otf/compilation/build_systems/cmake_lists.py
+++ b/src/gt4py/next/otf/compilation/build_systems/cmake_lists.py
@@ -48,6 +48,9 @@ class CMakeListsGenerator(eve.codegen.TemplatedGenerator):
         if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
             set(CMAKE_CUDA_ARCHITECTURES 60)
         endif()
+        if(NOT DEFINED CMAKE_HIP_ARCHITECTURES)
+            set(CMAKE_HIP_ARCHITECTURES gfx90a)
+        endif()
         {{"\\n".join(languages)}}
 
         # Paths
diff --git a/src/gt4py/next/otf/compilation/build_systems/compiledb.py b/src/gt4py/next/otf/compilation/build_systems/compiledb.py
index debf8217..6bacde49 100644
--- a/src/gt4py/next/otf/compilation/build_systems/compiledb.py
+++ b/src/gt4py/next/otf/compilation/build_systems/compiledb.py
@@ -255,8 +255,8 @@ def _cc_create_compiledb(
     prog_src_name = f"{name}.{header_ext}"
     binding_src_name = f"{name}.{src_ext}"
     cmake_languages = [cmake_lists.Language(name="CXX")]
-    if prototype_program_source.language is languages.Cuda:
-        cmake_languages = [*cmake_languages, cmake_lists.Language(name="CUDA")]
+    if (src_lang := prototype_program_source.language) in [languages.CUDA, languages.HIP]:
+        cmake_languages = [*cmake_languages, cmake_lists.Language(name=src_lang.__name__)]
 
     prototype_project = cmake.CMakeProject(
         generator_name="Ninja",
diff --git a/src/gt4py/next/otf/languages.py b/src/gt4py/next/otf/languages.py
index db1094db..1564a300 100644
--- a/src/gt4py/next/otf/languages.py
+++ b/src/gt4py/next/otf/languages.py
@@ -59,10 +59,14 @@ class SDFG(LanguageTag):
 class NanobindSrcL(LanguageTag): ...
 
 
-class Cpp(NanobindSrcL):
+class CPP(NanobindSrcL):
     settings_class = LanguageWithHeaderFilesSettings
     ...
 
 
-class Cuda(NanobindSrcL):
+class CUDA(NanobindSrcL):
+    settings_class = LanguageWithHeaderFilesSettings
+
+
+class HIP(NanobindSrcL):
     settings_class = LanguageWithHeaderFilesSettings
diff --git a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py
index ac5325aa..5349464e 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py
@@ -68,6 +68,13 @@ class GTFNTranslationStep(
                     file_extension="cu",
                     header_extension="cuh",
                 )
+            case core_defs.DeviceType.ROCM:
+                return languages.LanguageWithHeaderFilesSettings(
+                    formatter_key=cpp_interface.CPP_DEFAULT.formatter_key,
+                    formatter_style=cpp_interface.CPP_DEFAULT.formatter_style,
+                    file_extension="hip",
+                    header_extension="h",
+                )
             case core_defs.DeviceType.CPU:
                 return cpp_interface.CPP_DEFAULT
             case _:
@@ -263,7 +270,7 @@ class GTFNTranslationStep(
 
     def _backend_header(self) -> str:
         match self.device_type:
-            case core_defs.DeviceType.CUDA:
+            case core_defs.DeviceType.CUDA | core_defs.DeviceType.ROCM:
                 return "gridtools/fn/backend/gpu.hpp"
             case core_defs.DeviceType.CPU:
                 return "gridtools/fn/backend/naive.hpp"
@@ -272,7 +279,7 @@ class GTFNTranslationStep(
 
     def _backend_type(self) -> str:
         match self.device_type:
-            case core_defs.DeviceType.CUDA:
+            case core_defs.DeviceType.CUDA | core_defs.DeviceType.ROCM:
                 return "gridtools::fn::backend::gpu<generated::block_sizes_t>{}"
             case core_defs.DeviceType.CPU:
                 return "gridtools::fn::backend::naive{}"
@@ -282,9 +289,11 @@ class GTFNTranslationStep(
     def _language(self) -> type[languages.NanobindSrcL]:
         match self.device_type:
             case core_defs.DeviceType.CUDA:
-                return languages.Cuda
+                return languages.CUDA
+            case core_defs.DeviceType.ROCM:
+                return languages.HIP
             case core_defs.DeviceType.CPU:
-                return languages.Cpp
+                return languages.CPP
             case _:
                 raise self._not_implemented_for_device_type()
 
@@ -297,7 +306,7 @@ class GTFNTranslationStep(
 
     def _library_name(self) -> str:
         match self.device_type:
-            case core_defs.DeviceType.CUDA:
+            case core_defs.DeviceType.CUDA | core_defs.DeviceType.ROCM:
                 return "gridtools_gpu"
             case core_defs.DeviceType.CPU:
                 return "gridtools_cpu"
diff --git a/src/gt4py/next/program_processors/runners/dace_common/workflow.py b/src/gt4py/next/program_processors/runners/dace_common/workflow.py
index dbe2b70f..73437664 100644
--- a/src/gt4py/next/program_processors/runners/dace_common/workflow.py
+++ b/src/gt4py/next/program_processors/runners/dace_common/workflow.py
@@ -108,7 +108,7 @@ def convert_args(
 ) -> stages.CompiledProgram:
     sdfg_program = inp.sdfg_program
     sdfg = sdfg_program.sdfg
-    on_gpu = True if device == core_defs.DeviceType.CUDA else False
+    on_gpu = True if device in [core_defs.DeviceType.CUDA, core_defs.DeviceType.ROCM] else False
 
     def decorated_program(
         *args: Any,
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/workflow.py b/src/gt4py/next/program_processors/runners/dace_iterator/workflow.py
index 72cc15a4..0b23fde9 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/workflow.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/workflow.py
@@ -56,7 +56,11 @@ class DaCeTranslator(
         offset_provider: dict[str, common.Dimension | common.Connectivity],
         column_axis: Optional[common.Dimension],
     ) -> dace.SDFG:
-        on_gpu = True if self.device_type == core_defs.DeviceType.CUDA else False
+        on_gpu = (
+            True
+            if self.device_type in [core_defs.DeviceType.CUDA, core_defs.DeviceType.ROCM]
+            else False
+        )
 
         return build_sdfg_from_itir(
             program,
diff --git a/src/gt4py/next/program_processors/runners/gtfn.py b/src/gt4py/next/program_processors/runners/gtfn.py
index a2badd51..9c4e7352 100644
--- a/src/gt4py/next/program_processors/runners/gtfn.py
+++ b/src/gt4py/next/program_processors/runners/gtfn.py
@@ -65,7 +65,7 @@ def convert_args(
 def _ensure_is_on_device(
     connectivity_arg: npt.NDArray, device: core_defs.DeviceType
 ) -> npt.NDArray:
-    if device == core_defs.DeviceType.CUDA:
+    if device in [core_defs.DeviceType.CUDA, core_defs.DeviceType.ROCM]:
         import cupy as cp
 
         if not isinstance(connectivity_arg, cp.ndarray):
@@ -160,7 +160,7 @@ class GTFNBackendFactory(factory.Factory):
         name_postfix = ""
         gpu = factory.Trait(
             allocator=next_allocators.StandardGPUFieldBufferAllocator(),
-            device_type=core_defs.DeviceType.CUDA,
+            device_type=next_allocators.CUPY_DEVICE or core_defs.DeviceType.CUDA,
             name_device="gpu",
         )
         cached = factory.Trait(
diff --git a/src/gt4py/storage/cartesian/utils.py b/src/gt4py/storage/cartesian/utils.py
index 052238fe..50500e53 100644
--- a/src/gt4py/storage/cartesian/utils.py
+++ b/src/gt4py/storage/cartesian/utils.py
@@ -61,27 +61,6 @@ if cp:
             device_type=core_defs.DeviceType.ROCM,
             array_utils=allocators.cupy_array_utils,
         )
-
-        class CUDAArrayInterfaceNDArray(cp.ndarray):
-            def __new__(cls, input_array: "cp.ndarray") -> CUDAArrayInterfaceNDArray:
-                return (
-                    input_array
-                    if isinstance(input_array, CUDAArrayInterfaceNDArray)
-                    else cp.asarray(input_array).view(cls)
-                )
-
-            @property
-            def __cuda_array_interface__(self) -> dict:
-                return {
-                    "shape": self.shape,
-                    "typestr": self.dtype.descr[0][1],
-                    "descr": self.dtype.descr,
-                    "stream": 1,
-                    "version": 3,
-                    "strides": self.strides,
-                    "data": (self.data.ptr, False),
-                }
-
     else:
         raise ValueError("CuPy is available but no suitable device was found.")
 
@@ -282,7 +261,7 @@ def allocate_cpu(
     return buffer.buffer, cast(np.ndarray, buffer.ndarray)
 
 
-def allocate_gpu(
+def _allocate_gpu(
     shape: Sequence[int],
     layout_map: allocators.BufferLayoutMap,
     dtype: DTypeLike,
@@ -292,7 +271,8 @@ def allocate_gpu(
     assert cp is not None
     assert _GPUBufferAllocator is not None, "GPU allocation library or device not found"
     device = core_defs.Device(  # type: ignore[type-var]
-        (core_defs.DeviceType.ROCM if gt_config.GT4PY_USE_HIP else core_defs.DeviceType.CUDA), 0
+        (core_defs.DeviceType.ROCM if gt_config.GT4PY_USE_HIP else core_defs.DeviceType.CUDA),
+        0,
     )
     buffer = _GPUBufferAllocator.allocate(
         shape,
@@ -308,17 +288,41 @@ def allocate_gpu(
     return buffer.buffer, buffer_ndarray
 
 
+allocate_gpu = _allocate_gpu
+
 if CUPY_DEVICE == core_defs.DeviceType.ROCM:
 
-    @functools.wraps(allocate_gpu)
-    def allocate_gpu_rocm(
+    class CUDAArrayInterfaceNDArray(cp.ndarray):
+        def __new__(cls, input_array: "cp.ndarray") -> CUDAArrayInterfaceNDArray:
+            return (
+                input_array
+                if isinstance(input_array, CUDAArrayInterfaceNDArray)
+                else cp.asarray(input_array).view(cls)
+            )
+
+        @property
+        def __cuda_array_interface__(self) -> dict:
+            return {
+                "shape": self.shape,
+                "typestr": self.dtype.descr[0][1],
+                "descr": self.dtype.descr,
+                "stream": 1,
+                "version": 3,
+                "strides": self.strides,
+                "data": (self.data.ptr, False),
+            }
+
+        __hip_array_interface__ = __cuda_array_interface__
+
+    @functools.wraps(_allocate_gpu)
+    def _allocate_gpu_rocm(
         shape: Sequence[int],
         layout_map: allocators.BufferLayoutMap,
         dtype: DTypeLike,
         alignment_bytes: int,
         aligned_index: Optional[Sequence[int]],
     ) -> Tuple["cp.ndarray", "cp.ndarray"]:
-        buffer, ndarray = allocate_gpu(shape, layout_map, dtype, alignment_bytes, aligned_index)
+        buffer, ndarray = _allocate_gpu(shape, layout_map, dtype, alignment_bytes, aligned_index)
         return buffer, CUDAArrayInterfaceNDArray(ndarray)
 
-    allocate_gpu = allocate_gpu_rocm
+    allocate_gpu = _allocate_gpu_rocm


Test patch:
diff --git a/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py b/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py
index 79507f46..0848aba5 100644
--- a/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py
+++ b/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py
@@ -72,7 +72,7 @@ def make_program_source(name: str) -> stages.ProgramSource:
         entry_point=entry_point,
         source_code=src,
         library_deps=[interface.LibraryDependency("gridtools_cpu", "master")],
-        language=languages.Cpp,
+        language=languages.CPP,
         language_settings=cpp_interface.CPP_DEFAULT,
         implicit_domain=False,
     )
diff --git a/tests/next_tests/unit_tests/otf_tests/test_languages.py b/tests/next_tests/unit_tests/otf_tests/test_languages.py
index 98f64201..7a3fc0c0 100644
--- a/tests/next_tests/unit_tests/otf_tests/test_languages.py
+++ b/tests/next_tests/unit_tests/otf_tests/test_languages.py
@@ -18,7 +18,7 @@ def test_basic_settings_with_cpp_rejected():
             entry_point=interface.Function(name="basic_settings_with_cpp", parameters=[]),
             source_code="",
             library_deps=(),
-            language=languages.Cpp,
+            language=languages.CPP,
             language_settings=languages.LanguageSettings(
                 formatter_key="cpp", formatter_style="llvm", file_extension="cpp"
             ),
@@ -31,7 +31,7 @@ def test_header_files_settings_with_cpp_accepted():
         entry_point=interface.Function(name="basic_settings_with_cpp", parameters=[]),
         source_code="",
         library_deps=(),
-        language=languages.Cpp,
+        language=languages.CPP,
         language_settings=languages.LanguageWithHeaderFilesSettings(
             formatter_key="cpp",
             formatter_style="llvm",
diff --git a/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py b/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py
index bef82ad8..e5abcd7f 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py
@@ -68,4 +68,4 @@ def test_codegen(fencil_example):
     )
     assert module.entry_point.name == fencil.id
     assert any(d.name == "gridtools_cpu" for d in module.library_deps)
-    assert module.language is languages.Cpp
+    assert module.language is languages.CPP