mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp

   1 //===- SerializeROCDLTarget.cpp ---------------------------------*- C++ -*-===//
   2 //
   3 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
  10 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
  11 #include "mlir/IR/MLIRContext.h"
  12 #include "mlir/InitAllDialects.h"
  13 #include "mlir/Parser/Parser.h"
  14 #include "mlir/Target/LLVM/ROCDL/Target.h"
  15 #include "mlir/Target/LLVM/ROCDL/Utils.h"
  16 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
  17 #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
  18 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
  19 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
  20
  21 #include "llvm/IRReader/IRReader.h"
  22 #include "llvm/Support/FileSystem.h"
  23 #include "llvm/Support/MemoryBufferRef.h"
  24 #include "llvm/Support/Path.h"
  25 #include "llvm/Support/TargetSelect.h"
  26 #include "llvm/Support/raw_ostream.h"
  27 #include "llvm/TargetParser/Host.h"
  28
  29 #include "gmock/gmock.h"
  30
  31 using namespace mlir;
  32
  33 // Skip the test if the AMDGPU target was not built.
  34 #if MLIR_ENABLE_ROCM_CONVERSIONS
  35 #define SKIP_WITHOUT_AMDGPU(x) x
  36 #else
  37 #define SKIP_WITHOUT_AMDGPU(x) DISABLED_##x
  38 #endif
  39
  40 class MLIRTargetLLVMROCDL : public ::testing::Test {
  41 protected:
  42   void SetUp() override {
  43     registerBuiltinDialectTranslation(registry);
  44     registerLLVMDialectTranslation(registry);
  45     registerGPUDialectTranslation(registry);
  46     registerROCDLDialectTranslation(registry);
  47     ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
  48   }
  49
  50   // Checks if a ROCm installation is available.
  51   bool hasROCMTools() {
  52     StringRef rocmPath = ROCDL::getROCMPath();
  53     if (rocmPath.empty())
  54       return false;
  55     llvm::SmallString<128> lldPath(rocmPath);
  56     llvm::sys::path::append(lldPath, "llvm", "bin", "ld.lld");
  57     return llvm::sys::fs::can_execute(lldPath);
  58   }
  59
  60   // Dialect registry.
  61   DialectRegistry registry;
  62
  63   // MLIR module used for the tests.
  64   const std::string moduleStr = R"mlir(
  65       gpu.module @rocdl_test {
  66         llvm.func @rocdl_kernel(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
  67         llvm.return
  68       }
  69     })mlir";
  70 };
  71
  72 // Test ROCDL serialization to LLVM.
  73 TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToLLVM)) {
  74   MLIRContext context(registry);
  75
  76   OwningOpRef<ModuleOp> module =
  77       parseSourceString<ModuleOp>(moduleStr, &context);
  78   ASSERT_TRUE(!!module);
  79
  80   // Create a ROCDL target.
  81   ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
  82
  83   // Serialize the module.
  84   auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
  85   ASSERT_TRUE(!!serializer);
  86   gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Offload);
  87   for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
  88     std::optional<SmallVector<char, 0>> object =
  89         serializer.serializeToObject(gpuModule, options);
  90     // Check that the serializer was successful.
  91     ASSERT_TRUE(object != std::nullopt);
  92     ASSERT_TRUE(!object->empty());
  93
  94     // Read the serialized module.
  95     llvm::MemoryBufferRef buffer(StringRef(object->data(), object->size()),
  96                                  "module");
  97     llvm::LLVMContext llvmContext;
  98     llvm::Expected<std::unique_ptr<llvm::Module>> llvmModule =
  99         llvm::getLazyBitcodeModule(buffer, llvmContext);
 100     ASSERT_TRUE(!!llvmModule);
 101     ASSERT_TRUE(!!*llvmModule);
 102
 103     // Check that it has a function named `foo`.
 104     ASSERT_TRUE((*llvmModule)->getFunction("rocdl_kernel") != nullptr);
 105   }
 106 }
 107 // Test ROCDL serialization to ISA with default code object version.
 108 TEST_F(MLIRTargetLLVMROCDL,
 109        SKIP_WITHOUT_AMDGPU(SerializeROCDLToISAWithDefaultCOV)) {
 110   MLIRContext context(registry);
 111
 112   OwningOpRef<ModuleOp> module =
 113       parseSourceString<ModuleOp>(moduleStr, &context);
 114   ASSERT_TRUE(!!module);
 115
 116   // Create a ROCDL target.
 117   ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
 118
 119   // Serialize the module.
 120   auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
 121   ASSERT_TRUE(!!serializer);
 122   gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
 123   for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
 124     std::optional<SmallVector<char, 0>> object =
 125         serializer.serializeToObject(gpuModule, options);
 126     // Check that the serializer was successful.
 127     EXPECT_TRUE(StringRef(object->data(), object->size())
 128                     .contains(".amdhsa_code_object_version 5"));
 129   }
 130 }
 131
 132 // Test ROCDL serialization to ISA with non-default code object version.
 133 TEST_F(MLIRTargetLLVMROCDL,
 134        SKIP_WITHOUT_AMDGPU(SerializeROCDLToISAWithNonDefaultCOV)) {
 135   MLIRContext context(registry);
 136
 137   OwningOpRef<ModuleOp> module =
 138       parseSourceString<ModuleOp>(moduleStr, &context);
 139   ASSERT_TRUE(!!module);
 140
 141   // Create a ROCDL target.
 142   ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(
 143       &context, 2, "amdgcn-amd-amdhsa", "gfx900", "", "400");
 144
 145   // Serialize the module.
 146   auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
 147   ASSERT_TRUE(!!serializer);
 148   gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
 149   for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
 150     std::optional<SmallVector<char, 0>> object =
 151         serializer.serializeToObject(gpuModule, options);
 152     // Check that the serializer was successful.
 153     EXPECT_TRUE(StringRef(object->data(), object->size())
 154                     .contains(".amdhsa_code_object_version 4"));
 155   }
 156 }
 157
 158 // Test ROCDL serialization to PTX.
 159 TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToPTX)) {
 160   MLIRContext context(registry);
 161
 162   OwningOpRef<ModuleOp> module =
 163       parseSourceString<ModuleOp>(moduleStr, &context);
 164   ASSERT_TRUE(!!module);
 165
 166   // Create a ROCDL target.
 167   ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
 168
 169   // Serialize the module.
 170   auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
 171   ASSERT_TRUE(!!serializer);
 172   gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
 173   for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
 174     std::optional<SmallVector<char, 0>> object =
 175         serializer.serializeToObject(gpuModule, options);
 176     // Check that the serializer was successful.
 177     ASSERT_TRUE(object != std::nullopt);
 178     ASSERT_TRUE(!object->empty());
 179
 180     ASSERT_TRUE(
 181         StringRef(object->data(), object->size()).contains("rocdl_kernel"));
 182   }
 183 }
 184
 185 // Test ROCDL serialization to Binary.
 186 TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToBinary)) {
 187   if (!hasROCMTools())
 188     GTEST_SKIP() << "ROCm installation not found, skipping test.";
 189
 190   MLIRContext context(registry);
 191
 192   OwningOpRef<ModuleOp> module =
 193       parseSourceString<ModuleOp>(moduleStr, &context);
 194   ASSERT_TRUE(!!module);
 195
 196   // Create a ROCDL target.
 197   ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
 198
 199   // Serialize the module.
 200   auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
 201   ASSERT_TRUE(!!serializer);
 202   gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary);
 203   for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
 204     std::optional<SmallVector<char, 0>> object =
 205         serializer.serializeToObject(gpuModule, options);
 206     // Check that the serializer was successful.
 207     ASSERT_TRUE(object != std::nullopt);
 208     ASSERT_FALSE(object->empty());
 209   }
 210 }
 211
 212 // Test ROCDL metadata.
 213 TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(GetELFMetadata)) {
 214   if (!hasROCMTools())
 215     GTEST_SKIP() << "ROCm installation not found, skipping test.";
 216
 217   MLIRContext context(registry);
 218
 219   // MLIR module used for the tests.
 220   const std::string moduleStr = R"mlir(
 221     gpu.module @rocdl_test {
 222     llvm.func @rocdl_kernel_1(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
 223       llvm.return
 224     }
 225     llvm.func @rocdl_kernel_0(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
 226       llvm.return
 227     }
 228     llvm.func @rocdl_kernel_2(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
 229       llvm.return
 230     }
 231     llvm.func @a_kernel(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {
 232       llvm.return
 233     }
 234   })mlir";
 235
 236   OwningOpRef<ModuleOp> module =
 237       parseSourceString<ModuleOp>(moduleStr, &context);
 238   ASSERT_TRUE(!!module);
 239
 240   // Create a ROCDL target.
 241   ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
 242
 243   // Serialize the module.
 244   auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
 245   ASSERT_TRUE(!!serializer);
 246   gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary);
 247   for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
 248     std::optional<SmallVector<char, 0>> object =
 249         serializer.serializeToObject(gpuModule, options);
 250     // Check that the serializer was successful.
 251     ASSERT_TRUE(object != std::nullopt);
 252     ASSERT_FALSE(object->empty());
 253     if (!object)
 254       continue;
 255     // Get the metadata.
 256     gpu::KernelTableAttr metadata =
 257         ROCDL::getKernelMetadata(gpuModule, *object);
 258     ASSERT_TRUE(metadata != nullptr);
 259     // There should be 4 kernels.
 260     ASSERT_TRUE(metadata.size() == 4);
 261     // Check that the lookup method returns finds the kernel.
 262     ASSERT_TRUE(metadata.lookup("a_kernel") != nullptr);
 263     ASSERT_TRUE(metadata.lookup("rocdl_kernel_0") != nullptr);
 264     // Check that the kernel doesn't exist.
 265     ASSERT_TRUE(metadata.lookup("not_existent_kernel") == nullptr);
 266     // Test the `KernelMetadataAttr` iterators.
 267     for (gpu::KernelMetadataAttr kernel : metadata) {
 268       // Check that the ELF metadata is present.
 269       ASSERT_TRUE(kernel.getMetadata() != nullptr);
 270       // Verify that `sgpr_count` is present and it is an integer attribute.
 271       ASSERT_TRUE(kernel.getAttr<IntegerAttr>("sgpr_count") != nullptr);
 272       // Verify that `vgpr_count` is present and it is an integer attribute.
 273       ASSERT_TRUE(kernel.getAttr<IntegerAttr>("vgpr_count") != nullptr);
 274     }
 275   }
 276 }