Python SID adapter: Add support for HIP/ROCm buffers (GridTools#1759)

Light-weight PR to make the Python bindings work with AMD GPU buffers. Complementary to GridTools/gt4py#1278.

---------

Co-authored-by: Hannes Vogt <hannes@havogt.de>

AUTHORS

@@ -16,3 +16,4 @@ Auriane Reverdell (aurianer), ETH Zurich (CSCS)
 Mikael Simberg (msimberg), ETH Zurich (CSCS)
 Till Ehrengruber (tehrengruber), ETH Zurich (CSCS)
 Péter Kardos (petiaccja), ETH Zurich (EXCLAIM)
+Stefano Ubbiali (stubbiali), ETH Zurich

include/gridtools/storage/adapter/python_sid_adapter.hpp

@@ -268,7 +268,13 @@ namespace gridtools {
             static_assert(std::is_trivially_copy_constructible_v<T>,
                 "as_cuda_sid should be instantiated with the trivially copyable type");
 
+#if defined(__HIP__)
+            // This is a custom property that has to be added manually (not provided by cupy).
+            // Should be replaced by `dlpack` for uniform array interface support.
+            auto iface = src.attr("__hip_array_interface__").cast<pybind11::dict>();
+#else
             auto iface = src.attr("__cuda_array_interface__").cast<pybind11::dict>();
+#endif
 
             // shape
             array<size_t, Dim> shape;

tests/regression/py_bindings/CMakeLists.txt

@@ -36,4 +36,4 @@ pybind11_add_module(py_implementation implementation.cpp)
 
 target_link_libraries(py_implementation PRIVATE gridtools)
 
-add_test(NAME py_bindings COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/driver.py)
+add_test(NAME py_bindings COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/driver.py ${GT_CUDA_TYPE})

tests/regression/py_bindings/driver.py

@@ -6,6 +6,8 @@
 import numpy as np
 import py_implementation as testee
 
+HIP = True if len(sys.argv) > 1 and sys.argv[1] == "HIPCC-AMDGPU" else False
+
 def test_3d():
     src = np.fromfunction(lambda i, j, k : i + j + k, (3, 4, 5), dtype=np.double)
     dst = np.zeros_like(src)
@@ -35,9 +37,15 @@ def test_cuda_sid():
     class Mock:
         def __init__(self, **kwargs):
             self.kwargs = kwargs
-        @property
-        def __cuda_array_interface__(self):
-            return self.kwargs
+
+        if not HIP:
+            @property
+            def __cuda_array_interface__(self):
+                return self.kwargs
+        else:
+            @property
+            def __hip_array_interface__(self):
+                return self.kwargs
 
     mock = Mock(
         shape=(3, 4, 5),