Add support for compilation to ROCM (#344)

Author: ivanradanov

CMakeLists.txt

@@ -2,7 +2,8 @@ cmake_minimum_required(VERSION 3.10)
 
 include(CheckCXXSourceCompiles)
 
-set(POLYGEIST_ENABLE_CUDA 0 CACHE BOOL "Enable CUDA compilation support")
+set(POLYGEIST_ENABLE_CUDA 0 CACHE BOOL "Enable CUDA frontend and backend")
+set(POLYGEIST_ENABLE_ROCM 0 CACHE BOOL "Enable ROCM backend")
 
 if(POLICY CMP0068)
   cmake_policy(SET CMP0068 NEW)

README.md

@@ -30,7 +30,11 @@ ninja
 ninja check-mlir
 ```
 
-To enable compilation to cuda add `-DMLIR_ENABLE_CUDA_RUNNER=1` and remove `-DLLVM_TARGETS_TO_BUILD="host"` from the cmake arguments. (You may need to specify `CUDACXX`, `CUDA_PATH`, or `CMAKE_CUDA_COMPILER`)
+To enable compilation to cuda add `-DMLIR_ENABLE_CUDA_RUNNER=1` and remove `-DLLVM_TARGETS_TO_BUILD="host"` from the cmake arguments. (You may need to specify `CUDACXX`, `CUDA_PATH`, and/or `-DCMAKE_CUDA_COMPILER`)
+
+To enable the ROCM backend add `-DMLIR_ENABLE_ROCM_RUNNER=1` and remove `-DLLVM_TARGETS_TO_BUILD="host"` from the cmake arguments. (You may need to specify `-DHIP_CLANG_INCLUDE_PATH`, and/or `ROCM_PATH`)
+
+For faster compilation we recommend using `-DLLVM_USE_LINKER=lld`.
 
 2. Build Polygeist:
 ```sh
@@ -48,6 +52,10 @@ ninja check-polygeist-opt && ninja check-cgeist
 
 To enable compilation to cuda add `-DPOLYGEIST_ENABLE_CUDA=1`
 
+To enable the ROCM backend add `-DPOLYGEIST_ENABLE_ROCM=1`
+
+For faster compilation we recommend using `-DPOLYGEIST_USE_LINKER=lld`.
+
 #### Option 2: Using unified LLVM, MLIR, Clang, and Polygeist build
 
 Polygeist can also be built as an external LLVM project using [LLVM_EXTERNAL_PROJECTS](https://llvm.org/docs/CMake.html#llvm-related-variables).

include/polygeist/Passes/Passes.h

@@ -34,21 +34,29 @@ std::unique_ptr<Pass> createRemoveTrivialUsePass();
 std::unique_ptr<Pass> createParallelLowerPass(
     bool wrapParallelOps = false,
     PolygeistGPUStructureMode gpuKernelStructureMode = PGSM_Discard);
-std::unique_ptr<Pass> createCudaRTLowerPass();
+std::unique_ptr<Pass> createConvertCudaRTtoCPUPass();
+std::unique_ptr<Pass> createConvertCudaRTtoGPUPass();
+std::unique_ptr<Pass> createConvertCudaRTtoHipRTPass();
 std::unique_ptr<Pass>
 createConvertPolygeistToLLVMPass(const LowerToLLVMOptions &options,
-                                 bool useCStyleMemRef, bool onlyGpuModules);
+                                 bool useCStyleMemRef, bool onlyGpuModules,
+                                 std::string gpuTarget);
 std::unique_ptr<Pass> createConvertPolygeistToLLVMPass();
 std::unique_ptr<Pass> createForBreakToWhilePass();
 std::unique_ptr<Pass>
 createConvertParallelToGPUPass1(bool useOriginalThreadNums = false);
 std::unique_ptr<Pass>
 createConvertParallelToGPUPass2(bool emitGPUKernelLaunchBounds = true);
 std::unique_ptr<Pass> createGpuSerializeToCubinPass(
-    StringRef triple, StringRef arch, StringRef features, int llvmOptLevel,
-    int ptxasOptLevel, std::string ptxasPath, std::string libDevicePath,
-    bool outputIntermediate);
+    StringRef arch, StringRef features, int llvmOptLevel, int ptxasOptLevel,
+    std::string ptxasPath, std::string libDevicePath, bool outputIntermediate);
+std::unique_ptr<Pass>
+createGpuSerializeToHsacoPass(StringRef arch, StringRef features,
+                              int llvmOptLevel, int hsaOptLevel,
+                              std::string rocmPath, bool outputIntermediate);
+
 void registerGpuSerializeToCubinPass();
+void registerGpuSerializeToHsacoPass();
 
 void populateForBreakToWhilePatterns(RewritePatternSet &patterns);
 } // namespace polygeist

include/polygeist/Passes/Passes.td

@@ -13,11 +13,25 @@ def Mem2Reg : Pass<"mem2reg"> {
   let constructor = "mlir::polygeist::createMem2RegPass()";
 }
 
-def CudaRTLower : Pass<"cudart-lower", "mlir::ModuleOp"> {
+def ConvertCudaRTtoCPU : Pass<"convert-cudart-to-cpu", "mlir::ModuleOp"> {
   let summary = "Lower cudart functions to cpu versions";
   let dependentDialects =
       ["memref::MemRefDialect", "func::FuncDialect", "LLVM::LLVMDialect"];
-  let constructor = "mlir::polygeist::createCudaRTLowerPass()";
+  let constructor = "mlir::polygeist::createConvertCudaRTtoCPUPass()";
+}
+
+def ConvertCudaRTtoGPU : Pass<"convert-cudart-to-gpu", "mlir::ModuleOp"> {
+  let summary = "Lower cudart functions to generic gpu versions";
+  let dependentDialects =
+      ["memref::MemRefDialect", "func::FuncDialect", "LLVM::LLVMDialect", "gpu::GPUDialect"];
+  let constructor = "mlir::polygeist::createConvertCudaRTtoGPUPass()";
+}
+
+def ConvertCudaRTtoHipRT : Pass<"convert-cudart-to-gpu", "mlir::ModuleOp"> {
+  let summary = "Lower cudart functions to generic gpu versions";
+  let dependentDialects =
+      ["memref::MemRefDialect", "func::FuncDialect", "LLVM::LLVMDialect", "gpu::GPUDialect"];
+  let constructor = "mlir::polygeist::createConvertCudaRTtoGPUPass()";
 }
 
 def ParallelLower : Pass<"parallel-lower", "mlir::ModuleOp"> {

include/polygeist/PolygeistOps.td

@@ -100,7 +100,7 @@ def GPUThreadOp : Polygeist_Op<"gpu_thread", [
       "Value":$threadIndexX, "Value":$threadIndexY, "Value":$threadIndexZ)>];
 }
 
-def GPUAlternativesOp : Polygeist_Op<"gpu_alternatives", [
+def AlternativesOp : Polygeist_Op<"alternatives", [
   RecursiveMemoryEffects]> {
   let summary = "Provides several alternatives kernels for gpu code";
   let regions = (region VariadicRegion<SizedRegion<1>>:$regions);
@@ -141,7 +141,7 @@ def GPUErrorOp : Polygeist_Op<"gpu_error", [
 }
 
 def PolygeistYieldOp : Polygeist_Op<"polygeist_yield", [Pure, ReturnLike, Terminator,
-    ParentOneOf<["GPUAlternativesOp", "GPUWrapperOp", "GPUErrorOp", "GPUBlockOp", "GPUThreadOp"]>]> {
+    ParentOneOf<["AlternativesOp", "GPUWrapperOp", "GPUErrorOp", "GPUBlockOp", "GPUThreadOp"]>]> {
   let summary = "Polygeist ops terminator";
 }
 

lib/polygeist/ExecutionEngine/CMakeLists.txt

@@ -4,30 +4,9 @@ if(POLYGEIST_ENABLE_CUDA)
   find_package(CUDA)
   enable_language(CUDA)
 
-  # Shared lib wrapper
-
   find_library(CUDA_RUNTIME_LIBRARY cuda)
 
-  # add_mlir_library(polygeist_cuda_runtime
-  #   SHARED
-  #   CudaRuntimeWrappers.cpp
-
-  #   EXCLUDE_FROM_LIBMLIR
-  # )
-
-  # set_property(TARGET polygeist_cuda_runtime PROPERTY CXX_STANDARD 14)
-  # target_include_directories(polygeist_cuda_runtime
-  #   PRIVATE
-  #   ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
-  # )
-  # target_link_libraries(polygeist_cuda_runtime
-  #   PRIVATE
-  #   ${CUDA_RUNTIME_LIBRARY}
-  # )
-
   # Bitcode lib wrapper
-  set(POLYGEIST_EXECUTION_ENGINE_BITCODE_DIR ${CMAKE_CURRENT_BINARY_DIR})
-
   find_program(CLANG_TOOL clang PATHS ${LLVM_TOOLS_BINARY_DIR} NO_DEFAULT_PATH)
 
   set(source_directory ${CMAKE_CURRENT_SOURCE_DIR})
@@ -74,3 +53,60 @@ if(POLYGEIST_ENABLE_CUDA)
 
 
 endif()
+if(POLYGEIST_ENABLE_ROCM)
+  if (NOT DEFINED ROCM_PATH)
+    if (NOT DEFINED ENV{ROCM_PATH})
+      set(ROCM_PATH "/opt/rocm" CACHE PATH "Path to which ROCm has been installed")
+    else()
+      set(ROCM_PATH $ENV{ROCM_PATH} CACHE PATH "Path to which ROCm has been installed")
+    endif()
+  endif()
+  list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH} "${ROCM_PATH}/hip")
+  find_package(hip REQUIRED)
+
+  # Bitcode lib wrapper
+  find_program(CLANG_TOOL clang PATHS ${LLVM_TOOLS_BINARY_DIR} NO_DEFAULT_PATH)
+
+  set(source_directory ${CMAKE_CURRENT_SOURCE_DIR})
+  set(src_files
+    ${source_directory}/RocmRuntimeWrappers.cpp
+    )
+
+  set(bc_flags -c -emit-llvm -std=c++17 -fvisibility=hidden
+    -O3
+    #-nocudalib
+    -D__HIP_PLATFORM_AMD__
+    -Xclang -no-opaque-pointers
+    -I${ROCM_PATH}/include
+    -DPOLYGEIST_PGO_DEFAULT_DATA_DIR="${POLYGEIST_PGO_DEFAULT_DATA_DIR}"
+    -DPOLYGEIST_PGO_ALTERNATIVE_ENV_VAR="${POLYGEIST_PGO_ALTERNATIVE_ENV_VAR}"
+    -DPOLYGEIST_PGO_DATA_DIR_ENV_VAR="${POLYGEIST_PGO_DATA_DIR_ENV_VAR}"
+    )
+
+  foreach(src ${src_files})
+    get_filename_component(infile ${src} ABSOLUTE)
+    get_filename_component(filename ${src} NAME)
+    set(inc_outfile "${filename}.bin.h")
+    set(bc_outfile "${filename}.bc")
+
+    add_custom_command(OUTPUT ${bc_outfile}
+      COMMAND ${CLANG_TOOL}
+      ${bc_flags}
+      ${infile} -o ${bc_outfile}
+      DEPENDS ${infile}
+      COMMENT "Building LLVM bitcode ${bc_outfile}"
+      VERBATIM
+    )
+    add_custom_target(${bc_outfile}_target DEPENDS ${bc_outfile})
+    add_custom_command(OUTPUT ${inc_outfile}
+      COMMAND ${XXD_BIN} -i ${bc_outfile} ${inc_outfile}
+      DEPENDS ${bc_outfile}
+      COMMENT "Generating C header ${inc_outfile}"
+      VERBATIM
+    )
+    add_custom_target(execution_engine_rocm_wrapper_binary_include DEPENDS ${inc_outfile})
+    set_property(DIRECTORY APPEND PROPERTY ADDITIONAL_MAKE_CLEAN_FILES ${bc_outfile})
+    set_property(DIRECTORY APPEND PROPERTY ADDITIONAL_MAKE_CLEAN_FILES ${inc_outfile})
+
+  endforeach()
+endif()