277 lines · 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-exception6//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 31using namespace mlir;32 33// Skip the test if the AMDGPU target was not built.34#if MLIR_ENABLE_ROCM_CONVERSIONS35#define SKIP_WITHOUT_AMDGPU(x) x36#else37#define SKIP_WITHOUT_AMDGPU(x) DISABLED_##x38#endif39 40class MLIRTargetLLVMROCDL : public ::testing::Test {41protected: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.return68 }69 })mlir";70};71 72// Test ROCDL serialization to LLVM.73TEST_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.108TEST_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 6"));129 }130}131 132// Test ROCDL serialization to ISA with non-default code object version.133TEST_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.159TEST_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.186TEST_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.213TEST_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.return224 }225 llvm.func @rocdl_kernel_0(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {226 llvm.return227 }228 llvm.func @rocdl_kernel_2(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {229 llvm.return230 }231 llvm.func @a_kernel(%arg0: f32) attributes {gpu.kernel, rocdl.kernel} {232 llvm.return233 }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