Bind DLDeviceType

Author: ax3l

src/Base/Array4.H

@@ -6,7 +6,7 @@
 #pragma once
 
 #include "pyAMReX.H"
-#include "dlpack.h"
+#include "dlpack/dlpack.h"
 
 #include <AMReX_Array4.H>
 #include <AMReX_BLassert.H>

src/CMakeLists.txt

@@ -2,6 +2,7 @@
 add_subdirectory(AmrCore)
 add_subdirectory(Base)
 #add_subdirectory(Boundary)
+add_subdirectory(dlpack)
 #add_subdirectory(EB)
 #add_subdirectory(Extern)
 #add_subdirectory(LinearSolvers)

src/dlpack/CMakeLists.txt

@@ -0,0 +1,6 @@
+foreach(D IN LISTS AMReX_SPACEDIM)
+    target_sources(pyAMReX_${D}d
+        PRIVATE
+            DLPack.cpp
+    )
+endforeach()

src/dlpack/DLPack.cpp

@@ -0,0 +1,34 @@
+#include "pyAMReX.H"
+
+#include "dlpack.h"
+
+
+void init_DLPack(py::module& m)
+{
+    using namespace amrex;
+
+    // register types only if not already present, e.g., from another library
+    // that also implements DLPack bindings and exposes the types
+
+    py::type pyDLDeviceType = py::type::of<DLDeviceType>();
+    if (!pyDLDeviceType) {
+        py::native_enum<DLDeviceType>(m, "DLDeviceType", "enum.IntEnum")
+            .value("kDLCPU", DLDeviceType::kDLCPU)
+            .value("kDLCUDA", DLDeviceType::kDLCUDA)
+            .value("kDLCUDAHost", DLDeviceType::kDLCUDAHost)
+            .value("kDLOpenCL", DLDeviceType::kDLOpenCL)
+            .value("kDLVulkan", DLDeviceType::kDLVulkan)
+            .value("kDLMetal", DLDeviceType::kDLMetal)
+            .value("kDLVPI", DLDeviceType::kDLVPI)
+            .value("kDLROCM", DLDeviceType::kDLROCM)
+            .value("kDLROCMHost", DLDeviceType::kDLROCMHost)
+            .value("kDLExtDev", DLDeviceType::kDLExtDev)
+            .value("kDLCUDAManaged", DLDeviceType::kDLCUDAManaged)
+            .value("kDLOneAPI", DLDeviceType::kDLOneAPI)
+            .value("kDLWebGPU", DLDeviceType::kDLWebGPU)
+            .value("kDLHexagon", DLDeviceType::kDLHexagon)
+            .value("kDLMAIA", DLDeviceType::kDLMAIA)
+        ;
+    }
+
+}

src/dlpack.h renamed to src/dlpack/dlpack.h

@@ -20,6 +20,7 @@ void init_Arena(py::module&);
 void init_Array4(py::module&);
 void init_BaseFab(py::module&);
 void init_Box(py::module &);
+void init_DLPack(py::module &);
 void init_RealBox(py::module &);
 void init_BoxArray(py::module &);
 void init_CoordSys(py::module&);
@@ -98,6 +99,7 @@ PYBIND11_MODULE(amrex_3d_pybind, m) {
 
     // note: order from parent to child classes and argument usage
 
+    init_DLPack(m);
     init_AMReX(m);
     init_Arena(m);
     init_Dim3(m);