[MLIR][GPU] Add target arguments to SerializeToHsaco

[krzysz00]

Summary

Compiling code for AMD GPUs requires knowledge of which chipset is
being targeted, especially if the code uses chipset-specific
intrinsics (which is the case in a downstream convolution generator).
This commit adds target, chipset and features arguments to the
SerializeToHsaco constructor to enable passing in this required
information.

It also amends the ROCm integration tests to pass in the target
chipset, which is set to the chipset of the first GPU on the system
executing the tests.

[mehdi_amini]

Looks good overall, just one question inline

mlir/test/Integration/GPU/ROCM/lit.local.cfg

13

Can this be done in CMake and encoded in lit.site.cfg.in with the rest of the config, instead of doing this on every invocation of lit?

[krzysz00]

• Move the chipset detection into CMake

[krzysz00]

mlir/test/Integration/GPU/ROCM/lit.local.cfg

13

Yep, we can do that.

[jpienaar]

Nice, let us know if you need helping landing

mlir/lib/Dialect/GPU/CMakeLists.txt

 if (MLIR_ENABLE_CUDA_CONVERSIONS)
   set(NVPTX_LIBS
     NVPTXCodeGen
     NVPTXDesc
     NVPTXInfo
   )
 endif()
 
 if (MLIR_ENABLE_ROCM_CONVERSIONS)
   set(AMDGPU_LIBS
     MCParser
     AMDGPUAsmParser
     AMDGPUCodeGen
     AMDGPUDesc
     AMDGPUInfo
+    target
   )
 endif()
 
 add_mlir_dialect_library(MLIRGPUOps
   IR/GPUDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

 #include <mutex>
 
 using namespace mlir;
 
 namespace {
 class SerializeToHsacoPass
     : public PassWrapper<SerializeToHsacoPass, gpu::SerializeToBlobPass> {
 public:
-  SerializeToHsacoPass();
-
+  SerializeToHsacoPass(StringRef triple, StringRef arch, StringRef features);
   StringRef getArgument() const override { return "gpu-to-hsaco"; }
   StringRef getDescription() const override {
     return "Lower GPU kernel function to HSACO binary annotations";
   }
 
 private:
   void getDependentDialects(DialectRegistry &registry) const override;
 

 
 // Sets the 'option' to 'value' unless it already has a value.
 static void maybeSetOption(Pass::Option<std::string> &option,
                            function_ref<std::string()> getValue) {
   if (!option.hasValue())
     option = getValue();
 }
 
-SerializeToHsacoPass::SerializeToHsacoPass() {
-  maybeSetOption(this->triple, [] { return "amdgcn-amd-amdhsa"; });
-  maybeSetOption(this->chip, [] {
-    static auto chip = getDefaultChip();
-    return chip;
-  });
+SerializeToHsacoPass::SerializeToHsacoPass(StringRef triple, StringRef arch,
+                                           StringRef features) {
+  maybeSetOption(this->triple, [&triple] { return triple.str(); });
+  maybeSetOption(this->chip, [&arch] { return arch.str(); });
+  maybeSetOption(this->features, [&features] { return features.str(); });
 }
 
 void SerializeToHsacoPass::getDependentDialects(
     DialectRegistry &registry) const {
   registerROCDLDialectTranslation(registry);
   gpu::SerializeToBlobPass::getDependentDialects(registry);
 }
 
PassRegistration<SerializeToHsacoPass> registerSerializeToHSACO(
       [] {
         // Initialize LLVM AMDGPU backend.
         LLVMInitializeAMDGPUAsmParser();
         LLVMInitializeAMDGPUAsmPrinter();
         LLVMInitializeAMDGPUTarget();
         LLVMInitializeAMDGPUTargetInfo();
         LLVMInitializeAMDGPUTargetMC();
 
-        return std::make_unique<SerializeToHsacoPass>();
+        return std::make_unique<SerializeToHsacoPass>("amdgcn-amd-amdhsa", "",
+                                                      "");
       });
 }
 #else  // MLIR_GPU_TO_HSACO_PASS_ENABLE
 void mlir::registerGpuSerializeToHsacoPass() {}
 #endif // MLIR_GPU_TO_HSACO_PASS_ENABLE

mlir/lib/ExecutionEngine/CMakeLists.txt

if(MLIR_ENABLE_ROCM_RUNNER)
   find_library(ROCM_RUNTIME_LIBRARY amdhip64
                PATHS "${HIP_PATH}/lib")
   if (NOT ROCM_RUNTIME_LIBRARY)
     message(SEND_ERROR "Could not locate ROCm HIP runtime library")
   else()
     message(STATUS "ROCm HIP runtime lib: ${ROCM_RUNTIME_LIBRARY}")
   endif()
 
+  if (NOT DEFINED ROCM_TEST_CHIPSET)
+    execute_process(COMMAND "${ROCM_PATH}/bin/rocm_agent_enumerator"
+    OUTPUT_VARIABLE AGENTS_STRING
+    ERROR_VARIABLE AGENTS_STRING
+    RESULT_VARIABLE AGENT_ENUMERATOR_RESULT)
+
+    if (NOT AGENT_ENUMERATOR_RESULT EQUAL 0)
+      message(SEND_ERROR "Could not run rocm_agent_enumerator and ROCM_TEST_CHIPSET is not defined")
+      set(AGENTS_STRING "")
+    endif()
+    string(STRIP AGENTS_STRING ${AGENTS_STRING})
+    string(REPLACE "\n" ";" AGENTS_LIST ${AGENTS_STRING})
+    list(FILTER AGENTS_LIST EXCLUDE REGEX "gfx000")
+    if (AGENTS_LIST STREQUAL "")
+      message(SEND_ERROR "No non-CPU ROCm agents found on the system, and ROCM_TEST_CHIPSET is not defined")
+    else()
+      list(GET AGENTS_LIST 0 FIRST_AGENT)
+      set(ROCM_TEST_CHIPSET ${FIRST_AGENT} CACHE STRING "Chipset for which to compile ROCm integration tests")
+      message(STATUS "Compiling integration tests for ${ROCM_TEST_CHIPSET}")
+    endif()
+  endif()
+
   add_mlir_library(mlir_rocm_runtime
     SHARED
     RocmRuntimeWrappers.cpp
 
     EXCLUDE_FROM_LIBMLIR
   )
   target_compile_definitions(mlir_rocm_runtime
     PRIVATE

mlir/test/Integration/GPU/ROCM/gpu-to-hsaco.mlir

 // RUN: mlir-opt %s \
 // RUN:   -gpu-kernel-outlining \
-// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-rocdl,gpu-to-hsaco)' \
+// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-rocdl,gpu-to-hsaco{chip=%chip})' \
 // RUN:   -gpu-to-llvm \
 // RUN: | mlir-cpu-runner \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @other_func(%arg0 : f32, %arg1 : memref<?xf32>) {

mlir/test/Integration/GPU/ROCM/lit.local.cfg

-if not config.enable_rocm_runner:
+if not config.enable_rocm_runner or not config.rocm_test_chipset:
   config.unsupported = True
+
+config.substitutions.append(('%chip', config.rocm_test_chipset))

mlir/test/Integration/GPU/ROCM/two-modules.mlir

 // RUN: mlir-opt %s \
 // RUN:   -gpu-kernel-outlining \
-// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-rocdl,gpu-to-hsaco)' \
+// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-rocdl,gpu-to-hsaco{chip=%chip})' \
 // RUN:   -gpu-to-llvm \
 // RUN: | mlir-cpu-runner \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

mlir/test/Integration/GPU/ROCM/vecadd.mlir

 // RUN: mlir-opt %s \
 // RUN:   -convert-scf-to-std \
 // RUN:   -gpu-kernel-outlining \
-// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-rocdl,gpu-to-hsaco)' \
+// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-rocdl,gpu-to-hsaco{chip=%chip})' \
 // RUN:   -gpu-to-llvm \
 // RUN: | mlir-cpu-runner \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @vecadd(%arg0 : memref<?xf32>, %arg1 : memref<?xf32>, %arg2 : memref<?xf32>) {

mlir/test/Integration/GPU/ROCM/vector-transferops.mlir

 // RUN: mlir-opt %s \
 // RUN:   -convert-scf-to-std \
 // RUN:   -gpu-kernel-outlining \
-// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-rocdl,gpu-to-hsaco)' \
+// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-rocdl,gpu-to-hsaco{chip=%chip})' \
 // RUN:   -gpu-to-llvm \
 // RUN: | mlir-cpu-runner \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @vectransferx2(%arg0 : memref<?xf32>, %arg1 : memref<?xf32>) {

mlir/test/lit.site.cfg.py.in

 config.mlir_tools_dir = "@MLIR_TOOLS_DIR@"
 config.mlir_cmake_dir = "@MLIR_CMAKE_DIR@"
 config.linalg_test_lib_dir = "@MLIR_DIALECT_LINALG_INTEGRATION_TEST_LIB_DIR@"
 config.build_examples = @LLVM_BUILD_EXAMPLES@
 config.run_cuda_tests = @MLIR_ENABLE_CUDA_CONVERSIONS@
 config.enable_cuda_runner = @MLIR_ENABLE_CUDA_RUNNER@
 config.run_rocm_tests = @MLIR_ENABLE_ROCM_CONVERSIONS@
 config.enable_rocm_runner = @MLIR_ENABLE_ROCM_RUNNER@
+config.rocm_test_chipset = "@ROCM_TEST_CHIPSET@"
 config.spirv_wrapper_library_dir = "@MLIR_SPIRV_WRAPPER_LIBRARY_DIR@"
 config.enable_spirv_cpu_runner = @MLIR_ENABLE_SPIRV_CPU_RUNNER@
 config.vulkan_wrapper_library_dir = "@MLIR_VULKAN_WRAPPER_LIBRARY_DIR@"
 config.enable_vulkan_runner = @MLIR_ENABLE_VULKAN_RUNNER@
 config.enable_bindings_python = @MLIR_ENABLE_BINDINGS_PYTHON@
 config.mlir_integration_test_dir = "@MLIR_INTEGRATION_TEST_DIR@"
 config.intel_sde_executable = "@INTEL_SDE_EXECUTABLE@"
 config.mlir_run_amx_tests = @MLIR_RUN_AMX_TESTS@