[MLIR][GPU] Add target arguments to SerializeToHsaco

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.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D114107

diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt
index 14520ce..fee795f 100644 (file)
--- a/mlir/lib/Dialect/GPU/CMakeLists.txt
+++ b/mlir/lib/Dialect/GPU/CMakeLists.txt
@@ -13,6 +13,7 @@ if (MLIR_ENABLE_ROCM_CONVERSIONS)
     AMDGPUCodeGen
     AMDGPUDesc
     AMDGPUInfo
+    target
   )
 endif()
 

diff --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
index 547b7bc..ab81ffa 100644 (file)
--- a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
@@ -48,8 +48,7 @@ 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";
@@ -132,12 +131,11 @@ static void maybeSetOption(Pass::Option<std::string> &option,
     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(
@@ -281,7 +279,8 @@ void mlir::registerGpuSerializeToHsacoPass() {
         LLVMInitializeAMDGPUTargetInfo();
         LLVMInitializeAMDGPUTargetMC();
 
-        return std::make_unique<SerializeToHsacoPass>();
+        return std::make_unique<SerializeToHsacoPass>("amdgcn-amd-amdhsa", "",
+                                                      "");
       });
 }
 #else  // MLIR_GPU_TO_HSACO_PASS_ENABLE

diff --git a/mlir/lib/ExecutionEngine/CMakeLists.txt b/mlir/lib/ExecutionEngine/CMakeLists.txt
index d630a3c..0e8d184 100644 (file)
--- a/mlir/lib/ExecutionEngine/CMakeLists.txt
+++ b/mlir/lib/ExecutionEngine/CMakeLists.txt
@@ -165,6 +165,28 @@ if(MLIR_ENABLE_ROCM_RUNNER)
     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

diff --git a/mlir/test/Integration/GPU/ROCM/gpu-to-hsaco.mlir b/mlir/test/Integration/GPU/ROCM/gpu-to-hsaco.mlir
index 19c7a60..fef4e04 100644 (file)
--- a/mlir/test/Integration/GPU/ROCM/gpu-to-hsaco.mlir
+++ b/mlir/test/Integration/GPU/ROCM/gpu-to-hsaco.mlir
@@ -1,6 +1,6 @@
 // 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 \

diff --git a/mlir/test/Integration/GPU/ROCM/lit.local.cfg b/mlir/test/Integration/GPU/ROCM/lit.local.cfg
index 0ced069..b0d086f 100644 (file)
--- a/mlir/test/Integration/GPU/ROCM/lit.local.cfg
+++ b/mlir/test/Integration/GPU/ROCM/lit.local.cfg
@@ -1,2 +1,4 @@
-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))

diff --git a/mlir/test/Integration/GPU/ROCM/two-modules.mlir b/mlir/test/Integration/GPU/ROCM/two-modules.mlir
index ec05444..c865172 100644 (file)
--- a/mlir/test/Integration/GPU/ROCM/two-modules.mlir
+++ b/mlir/test/Integration/GPU/ROCM/two-modules.mlir
@@ -1,6 +1,6 @@
 // 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 \

diff --git a/mlir/test/Integration/GPU/ROCM/vecadd.mlir b/mlir/test/Integration/GPU/ROCM/vecadd.mlir
index fbe43fa..1feaf34 100644 (file)
--- a/mlir/test/Integration/GPU/ROCM/vecadd.mlir
+++ b/mlir/test/Integration/GPU/ROCM/vecadd.mlir
@@ -1,7 +1,7 @@
 // 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 \

diff --git a/mlir/test/Integration/GPU/ROCM/vector-transferops.mlir b/mlir/test/Integration/GPU/ROCM/vector-transferops.mlir
index f9729f9..c65b790 100644 (file)
--- a/mlir/test/Integration/GPU/ROCM/vector-transferops.mlir
+++ b/mlir/test/Integration/GPU/ROCM/vector-transferops.mlir
@@ -1,7 +1,7 @@
 // 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 \

diff --git a/mlir/test/lit.site.cfg.py.in b/mlir/test/lit.site.cfg.py.in
index ef9900a..b48d878 100644 (file)
--- a/mlir/test/lit.site.cfg.py.in
+++ b/mlir/test/lit.site.cfg.py.in
@@ -42,6 +42,7 @@ 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@"