Add HIP to pybind11

Author: jakurzak

pybind_interface/decide/decide.cpp

@@ -60,7 +60,7 @@ int detect_instructions() {
 }
 
 enum GPUCapabilities {
-    CUDA = 0, CUSTATEVEC = 1, NO_GPU = 10, NO_CUSTATEVEC = 11 };
+    CUDA = 0, CUSTATEVEC = 1, HIP = 2, NO_GPU = 10, NO_CUSTATEVEC = 11 };
 
 // For now, GPU detection is performed at compile time, as our wheels are
 // generated on Github Actions runners which do not have GPU support.
@@ -70,6 +70,8 @@ enum GPUCapabilities {
 int detect_gpu() {
   #ifdef __NVCC__
   GPUCapabilities gpu = CUDA;
+  #elif __HIP__
+  GPUCapabilities gpu = HIP;
   #else
   GPUCapabilities gpu = NO_GPU;
   #endif

pybind_interface/hip/CMakeLists.txt

@@ -0,0 +1,28 @@
+cmake_minimum_required(VERSION 3.18)
+project(qsim LANGUAGES CXX HIP)
+
+IF (WIN32)
+    set(CMAKE_CXX_FLAGS "/O2 /openmp")
+ELSE()
+    set(CMAKE_CXX_FLAGS "-O3 -fopenmp")
+ENDIF()
+
+if(APPLE)
+    set(CMAKE_CXX_STANDARD 14)
+    include_directories("/usr/local/include" "/usr/local/opt/llvm/include")
+    link_directories("/usr/local/lib" "/usr/local/opt/llvm/lib")
+endif()
+
+INCLUDE(../GetPybind11.cmake)
+find_package(PythonLibs 3.7 REQUIRED)
+find_package(HIP REQUIRED)
+
+include_directories(${PYTHON_INCLUDE_DIRS} ${pybind11_SOURCE_DIR}/include)
+
+hip_add_library(qsim_hip MODULE pybind_main_hip.cpp)
+
+set_target_properties(qsim_hip PROPERTIES
+       PREFIX "${PYTHON_MODULE_PREFIX}"
+       SUFFIX "${PYTHON_MODULE_EXTENSION}"
+)
+set_source_files_properties(pybind_main_hip.cpp PROPERTIES LANGUAGE HIP)

pybind_interface/hip/pybind_main_hip.cpp

@@ -0,0 +1,47 @@
+// Copyright 2019 Google LLC. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "pybind_main_hip.h"
+
+#include "../../lib/simulator_cuda.h"
+
+namespace qsim {
+  using Simulator = SimulatorCUDA<float>;
+
+  struct Factory {
+    using Simulator = qsim::Simulator;
+    using StateSpace = Simulator::StateSpace;
+
+    Factory(
+      unsigned num_sim_threads,
+      unsigned num_state_threads,
+      unsigned num_dblocks
+    ) : ss_params{num_state_threads, num_dblocks} {}
+
+    StateSpace CreateStateSpace() const {
+      return StateSpace(ss_params);
+    }
+
+    Simulator CreateSimulator() const {
+      return Simulator();
+    }
+
+    StateSpace::Parameter ss_params;
+  };
+
+  inline void SetFlushToZeroAndDenormalsAreZeros() {}
+  inline void ClearFlushToZeroAndDenormalsAreZeros() {}
+}
+
+#include "../pybind_main.cpp"

pybind_interface/hip/pybind_main_hip.h

@@ -0,0 +1,17 @@
+// Copyright 2019 Google LLC. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "../pybind_main.h"
+
+PYBIND11_MODULE(qsim_hip, m) { GPU_MODULE_BINDINGS }

qsimcirq/__init__.py

@@ -34,6 +34,8 @@ def _load_qsim_gpu():
     instr = qsim_decide.detect_gpu()
     if instr == 0:
         qsim_gpu = importlib.import_module("qsimcirq.qsim_cuda")
+    elif instr == 2:
+        qsim_gpu = importlib.import_module("qsimcirq.qsim_hip")
     else:
         qsim_gpu = None
     return qsim_gpu

setup.py

@@ -111,6 +111,7 @@ def build_extension(self, ext):
         CMakeExtension("qsimcirq/qsim_cuda"),
         CMakeExtension("qsimcirq/qsim_custatevec"),
         CMakeExtension("qsimcirq/qsim_decide"),
+        CMakeExtension("qsimcirq/qsim_hip"),
     ],
     cmdclass=dict(build_ext=CMakeBuild),
     zip_safe=False,