[MLIR] Make the ROCM integration tests runnable

[krzysz00]

Summary

• Move the #define s to the GPU Transform library from GPU Ops so that

SerializeToHsaco is non-trivially compiled

• Add required includes to SerializeToHsaco

• Move MCSubtargetInfo creation to the correct point in the

compilation process

• Change mlir in ROCM tests to account for renamed/moved ops

mlir/lib/Dialect/GPU/CMakeLists.txt

if(MLIR_ENABLE_ROCM_RUNNER)
   set(CMAKE_MODULE_PATH "${HIP_PATH}/cmake" ${CMAKE_MODULE_PATH})
   find_package(HIP)
   if (NOT HIP_FOUND)
     message(SEND_ERROR "Building mlir with ROCm support requires a working ROCm and HIP install")
   else()
     message(STATUS "ROCm HIP version: ${HIP_VERSION}")
   endif()
 
-  target_compile_definitions(obj.MLIRGPUOps
+  target_compile_definitions(obj.MLIRGPUTransforms
     PRIVATE
     __HIP_PLATFORM_HCC__
     __ROCM_PATH__="${ROCM_PATH}"
     MLIR_GPU_TO_HSACO_PASS_ENABLE=1
   )
 
-  target_include_directories(obj.MLIRGPUOps
+  target_include_directories(obj.MLIRGPUTransforms
     PRIVATE
     ${MLIR_SOURCE_DIR}/../lld/include
     ${HIP_PATH}/include
     ${ROCM_PATH}/include
   )
 
   target_link_libraries(MLIRGPUOps
     PRIVATE

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

 //===- LowerGPUToHSACO.cpp - Convert GPU kernel to HSACO blob -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass that serializes a gpu module into HSAco blob and
 // adds that blob as a string attribute of the module.
 //
 //===----------------------------------------------------------------------===//
 #include "mlir/Dialect/GPU/Passes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
 
 #if MLIR_GPU_TO_HSACO_PASS_ENABLE
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
 
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCCodeEmitter.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCParser/MCTargetAsmParser.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/LineIterator.h"
 #include "llvm/Support/Program.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/WithColor.h"
+
+#include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 
 #include "lld/Common/Driver.h"
 
 #include "hip/hip_version.h"
 
 #include <mutex>
 
srcMgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(isa),
                             llvm::SMLoc());
 
   const llvm::MCTargetOptions mcOptions;
   std::unique_ptr<llvm::MCRegisterInfo> mri(
       target->createMCRegInfo(this->triple));
   std::unique_ptr<llvm::MCAsmInfo> mai(
       target->createMCAsmInfo(*mri, this->triple, mcOptions));
   mai->setRelaxELFRelocations(true);
+  std::unique_ptr<llvm::MCSubtargetInfo> sti(
+      target->createMCSubtargetInfo(this->triple, this->chip, this->features));
 
-  llvm::MCContext ctx(triple, mai.get(), mri.get(), &srcMgr, &mcOptions);
+  llvm::MCContext ctx(triple, mai.get(), mri.get(), sti.get(), &srcMgr,
+                      &mcOptions);
   std::unique_ptr<llvm::MCObjectFileInfo> mofi(target->createMCObjectFileInfo(
       ctx, /*PIC=*/false, /*LargeCodeModel=*/false));
   ctx.setObjectFileInfo(mofi.get());
 
   SmallString<128> cwd;
   if (!llvm::sys::fs::current_path(cwd))
     ctx.setCompilationDir(cwd);
 
   std::unique_ptr<llvm::MCStreamer> mcStreamer;
   std::unique_ptr<llvm::MCInstrInfo> mcii(target->createMCInstrInfo());
-  std::unique_ptr<llvm::MCSubtargetInfo> sti(
-      target->createMCSubtargetInfo(this->triple, this->chip, this->features));
 
   llvm::MCCodeEmitter *ce = target->createMCCodeEmitter(*mcii, *mri, ctx);
   llvm::MCAsmBackend *mab = target->createMCAsmBackend(*sti, *mri, mcOptions);
   mcStreamer.reset(target->createMCObjectStreamer(
       triple, ctx, std::unique_ptr<llvm::MCAsmBackend>(mab),
       mab->createObjectWriter(os), std::unique_ptr<llvm::MCCodeEmitter>(ce),
       *sti, mcOptions.MCRelaxAll, mcOptions.MCIncrementalLinkerCompatible,
       /*DWARFMustBeAtTheEnd*/ false));

mlir/lib/ExecutionEngine/CMakeLists.txt

target_compile_definitions(mlir_rocm_runtime
     PRIVATE
     __HIP_PLATFORM_HCC__
   )
   target_include_directories(mlir_rocm_runtime
     PRIVATE
     ${HIP_PATH}/include
     ${ROCM_PATH}/include
   )
+  set_property(TARGET mlir_rocm_runtime
+    PROPERTY INSTALL_RPATH_USE_LINK_PATH ON)
+
   target_link_libraries(mlir_rocm_runtime
-    PRIVATE
+    PUBLIC
     ${ROCM_RUNTIME_LIBRARY}
   )
 endif()

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{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>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
-  %block_dim = dim %arg1, %c0 : memref<?xf32>
+  %block_dim = memref.dim %arg1, %c0 : memref<?xf32>
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %block_dim, %block_y = %c1, %block_z = %c1) {
-    store %arg0, %arg1[%tx] : memref<?xf32>
+    memref.store %arg0, %arg1[%tx] : memref<?xf32>
     gpu.terminator
   }
   return
 }
 
 // CHECK: [1, 1, 1, 1, 1]
 func @main() {
-  %arg0 = alloc() : memref<5xf32>
+  %arg0 = memref.alloc() : memref<5xf32>
   %21 = arith.constant 5 : i32
-  %22 = memref_cast %arg0 : memref<5xf32> to memref<?xf32>
-  %cast = memref_cast %22 : memref<?xf32> to memref<*xf32>
+  %22 = memref.cast %arg0 : memref<5xf32> to memref<?xf32>
+  %cast = memref.cast %22 : memref<?xf32> to memref<*xf32>
   gpu.host_register %cast : memref<*xf32>
-  %23 = memref_cast %22 : memref<?xf32> to memref<*xf32>
+  %23 = memref.cast %22 : memref<?xf32> to memref<*xf32>
   call @print_memref_f32(%23) : (memref<*xf32>) -> ()
   %24 = arith.constant 1.0 : f32
   %25 = call @mgpuMemGetDeviceMemRef1dFloat(%22) : (memref<?xf32>) -> (memref<?xf32>)
   call @other_func(%24, %25) : (f32, memref<?xf32>) -> ()
   call @print_memref_f32(%23) : (memref<*xf32>) -> ()
   return
 }
 
 func private @mgpuMemGetDeviceMemRef1dFloat(%ptr : memref<?xf32>) -> (memref<?xf32>)
 func private @print_memref_f32(%ptr : memref<*xf32>)

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{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]
 func @main() {
-  %arg = alloc() : memref<13xi32>
-  %dst = memref_cast %arg : memref<13xi32> to memref<?xi32>
+  %arg = memref.alloc() : memref<13xi32>
+  %dst = memref.cast %arg : memref<13xi32> to memref<?xi32>
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
-  %sx = dim %dst, %c0 : memref<?xi32>
-  %cast_dst = memref_cast %dst : memref<?xi32> to memref<*xi32>
+  %sx = memref.dim %dst, %c0 : memref<?xi32>
+  %cast_dst = memref.cast %dst : memref<?xi32> to memref<*xi32>
   gpu.host_register %cast_dst : memref<*xi32>
   %dst_device = call @mgpuMemGetDeviceMemRef1dInt32(%dst) : (memref<?xi32>) -> (memref<?xi32>)
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %c1, %block_z = %c1) {
     %t0 = arith.index_cast %tx : index to i32
-    store %t0, %dst_device[%tx] : memref<?xi32>
+    memref.store %t0, %dst_device[%tx] : memref<?xi32>
     gpu.terminator
   }
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %c1, %block_z = %c1) {
     %t0 = arith.index_cast %tx : index to i32
-    store %t0, %dst_device[%tx] : memref<?xi32>
+    memref.store %t0, %dst_device[%tx] : memref<?xi32>
     gpu.terminator
   }
   call @print_memref_i32(%cast_dst) : (memref<*xi32>) -> ()
   return
 }
 
 func private @mgpuMemGetDeviceMemRef1dInt32(%ptr : memref<?xi32>) -> (memref<?xi32>)
 func private @print_memref_i32(%ptr : memref<*xi32>)

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{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>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
-  %block_dim = dim %arg0, %c0 : memref<?xf32>
+  %block_dim = memref.dim %arg0, %c0 : memref<?xf32>
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %block_dim, %block_y = %c1, %block_z = %c1) {
-    %a = load %arg0[%tx] : memref<?xf32>
-    %b = load %arg1[%tx] : memref<?xf32>
+    %a = memref.load %arg0[%tx] : memref<?xf32>
+    %b = memref.load %arg1[%tx] : memref<?xf32>
     %c = arith.addf %a, %b : f32
-    store %c, %arg2[%tx] : memref<?xf32>
+    memref.store %c, %arg2[%tx] : memref<?xf32>
     gpu.terminator
   }
   return
 }
 
 // CHECK: [2.46, 2.46, 2.46, 2.46, 2.46]
 func @main() {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c5 = arith.constant 5 : index
   %cf1dot23 = arith.constant 1.23 : f32
-  %0 = alloc() : memref<5xf32>
-  %1 = alloc() : memref<5xf32>
-  %2 = alloc() : memref<5xf32>
-  %3 = memref_cast %0 : memref<5xf32> to memref<?xf32>
-  %4 = memref_cast %1 : memref<5xf32> to memref<?xf32>
-  %5 = memref_cast %2 : memref<5xf32> to memref<?xf32>
+  %0 = memref.alloc() : memref<5xf32>
+  %1 = memref.alloc() : memref<5xf32>
+  %2 = memref.alloc() : memref<5xf32>
+  %3 = memref.cast %0 : memref<5xf32> to memref<?xf32>
+  %4 = memref.cast %1 : memref<5xf32> to memref<?xf32>
+  %5 = memref.cast %2 : memref<5xf32> to memref<?xf32>
   scf.for %i = %c0 to %c5 step %c1 {
-    store %cf1dot23, %3[%i] : memref<?xf32>
-    store %cf1dot23, %4[%i] : memref<?xf32>
+    memref.store %cf1dot23, %3[%i] : memref<?xf32>
+    memref.store %cf1dot23, %4[%i] : memref<?xf32>
   }
-  %6 = memref_cast %3 : memref<?xf32> to memref<*xf32>
-  %7 = memref_cast %4 : memref<?xf32> to memref<*xf32>
-  %8 = memref_cast %5 : memref<?xf32> to memref<*xf32>
+  %6 = memref.cast %3 : memref<?xf32> to memref<*xf32>
+  %7 = memref.cast %4 : memref<?xf32> to memref<*xf32>
+  %8 = memref.cast %5 : memref<?xf32> to memref<*xf32>
   gpu.host_register %6 : memref<*xf32>
   gpu.host_register %7 : memref<*xf32>
   gpu.host_register %8 : memref<*xf32>
   %9 = call @mgpuMemGetDeviceMemRef1dFloat(%3) : (memref<?xf32>) -> (memref<?xf32>)
   %10 = call @mgpuMemGetDeviceMemRef1dFloat(%4) : (memref<?xf32>) -> (memref<?xf32>)
   %11 = call @mgpuMemGetDeviceMemRef1dFloat(%5) : (memref<?xf32>) -> (memref<?xf32>)
 
   call @vecadd(%9, %10, %11) : (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
   call @print_memref_f32(%8) : (memref<*xf32>) -> ()
   return
 }
 
 func private @mgpuMemGetDeviceMemRef1dFloat(%ptr : memref<?xf32>) -> (memref<?xf32>)
 func private @print_memref_f32(%ptr : memref<*xf32>)

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

 
 func @main() {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c4 = arith.constant 4 : index
   %cf1 = arith.constant 1.0 : f32
   %cf1dot23 = arith.constant 1.23 : f32
 
-  %arg0 = alloc() : memref<4xf32>
-  %arg1 = alloc() : memref<4xf32>
+  %arg0 = memref.alloc() : memref<4xf32>
+  %arg1 = memref.alloc() : memref<4xf32>
 
-  %22 = memref_cast %arg0 : memref<4xf32> to memref<?xf32>
-  %23 = memref_cast %arg1 : memref<4xf32> to memref<?xf32>
+  %22 = memref.cast %arg0 : memref<4xf32> to memref<?xf32>
+  %23 = memref.cast %arg1 : memref<4xf32> to memref<?xf32>
 
   scf.for %i = %c0 to %c4 step %c1 {
-    store %cf1dot23, %22[%i] : memref<?xf32>
-    store %cf1dot23, %23[%i] : memref<?xf32>
+    memref.store %cf1dot23, %22[%i] : memref<?xf32>
+    memref.store %cf1dot23, %23[%i] : memref<?xf32>
   }
 
-  %cast0 = memref_cast %22 : memref<?xf32> to memref<*xf32>
-  %cast1 = memref_cast %23 : memref<?xf32> to memref<*xf32>
+  %cast0 = memref.cast %22 : memref<?xf32> to memref<*xf32>
+  %cast1 = memref.cast %23 : memref<?xf32> to memref<*xf32>
 
   gpu.host_register %cast0 : memref<*xf32>
   gpu.host_register %cast1 : memref<*xf32>
 
   %24 = call @mgpuMemGetDeviceMemRef1dFloat(%22) : (memref<?xf32>) -> (memref<?xf32>)
   %26 = call @mgpuMemGetDeviceMemRef1dFloat(%23) : (memref<?xf32>) -> (memref<?xf32>)
 
   // CHECK: [1.23, 2.46, 2.46, 1.23]