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 } 277