300 lines · cpp
1//===- SerializeNVVMTarget.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/Config/mlir-config.h"10#include "mlir/Dialect/GPU/IR/GPUDialect.h"11#include "mlir/Dialect/LLVMIR/NVVMDialect.h"12#include "mlir/IR/MLIRContext.h"13#include "mlir/InitAllDialects.h"14#include "mlir/Parser/Parser.h"15#include "mlir/Target/LLVM/NVVM/Target.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/NVVM/NVVMToLLVMIRTranslation.h"20 21#include "llvm/Bitcode/BitcodeWriter.h"22#include "llvm/Config/Targets.h" // for LLVM_HAS_NVPTX_TARGET23#include "llvm/IRReader/IRReader.h"24#include "llvm/Support/MemoryBufferRef.h"25#include "llvm/Support/Process.h"26#include "llvm/Support/SourceMgr.h"27#include "llvm/Support/TargetSelect.h"28#include "llvm/Support/raw_ostream.h"29#include "llvm/TargetParser/Host.h"30 31#include "gmock/gmock.h"32#include <cstdint>33 34using namespace mlir;35 36// Skip the test if the NVPTX target was not built.37#if LLVM_HAS_NVPTX_TARGET38#define SKIP_WITHOUT_NVPTX(x) x39#else40#define SKIP_WITHOUT_NVPTX(x) DISABLED_##x41#endif42 43class MLIRTargetLLVMNVVM : public ::testing::Test {44protected:45 void SetUp() override {46 registerBuiltinDialectTranslation(registry);47 registerLLVMDialectTranslation(registry);48 registerGPUDialectTranslation(registry);49 registerNVVMDialectTranslation(registry);50 NVVM::registerNVVMTargetInterfaceExternalModels(registry);51 }52 53 // Checks if PTXAS is in PATH.54 bool hasPtxas() {55 // Find the `ptxas` compiler.56 std::optional<std::string> ptxasCompiler =57 llvm::sys::Process::FindInEnvPath("PATH", "ptxas");58 return ptxasCompiler.has_value();59 }60 61 // Dialect registry.62 DialectRegistry registry;63 64 // MLIR module used for the tests.65 const std::string moduleStr = R"mlir(66 gpu.module @nvvm_test {67 llvm.func @nvvm_kernel(%arg0: f32) attributes {gpu.kernel, nvvm.kernel} {68 llvm.return69 }70 })mlir";71};72 73// Test NVVM serialization to LLVM.74TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMMToLLVM)) {75 MLIRContext context(registry);76 77 OwningOpRef<ModuleOp> module =78 parseSourceString<ModuleOp>(moduleStr, &context);79 ASSERT_TRUE(!!module);80 81 // Create an NVVM target.82 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);83 84 // Serialize the module.85 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);86 ASSERT_TRUE(!!serializer);87 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Offload);88 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {89 std::optional<SmallVector<char, 0>> object =90 serializer.serializeToObject(gpuModule, options);91 // Check that the serializer was successful.92 ASSERT_TRUE(object != std::nullopt);93 ASSERT_TRUE(!object->empty());94 95 // Read the serialized module.96 llvm::MemoryBufferRef buffer(StringRef(object->data(), object->size()),97 "module");98 llvm::LLVMContext llvmContext;99 llvm::Expected<std::unique_ptr<llvm::Module>> llvmModule =100 llvm::getLazyBitcodeModule(buffer, llvmContext);101 ASSERT_TRUE(!!llvmModule);102 ASSERT_TRUE(!!*llvmModule);103 104 // Check that it has a function named `foo`.105 ASSERT_TRUE((*llvmModule)->getFunction("nvvm_kernel") != nullptr);106 }107}108 109// Test NVVM serialization to PTX.110TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) {111 MLIRContext context(registry);112 113 OwningOpRef<ModuleOp> module =114 parseSourceString<ModuleOp>(moduleStr, &context);115 ASSERT_TRUE(!!module);116 117 // Create an NVVM target.118 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);119 120 // Serialize the module.121 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);122 ASSERT_TRUE(!!serializer);123 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Assembly);124 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {125 std::optional<SmallVector<char, 0>> object =126 serializer.serializeToObject(gpuModule, options);127 // Check that the serializer was successful.128 ASSERT_TRUE(object != std::nullopt);129 ASSERT_TRUE(!object->empty());130 131 ASSERT_TRUE(132 StringRef(object->data(), object->size()).contains("nvvm_kernel"));133 ASSERT_TRUE(StringRef(object->data(), object->size()).count('\0') == 0);134 }135}136 137// Test NVVM serialization to Binary.138TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) {139 if (!hasPtxas())140 GTEST_SKIP() << "PTXAS compiler not found, skipping test.";141 142 MLIRContext context(registry);143 144 OwningOpRef<ModuleOp> module =145 parseSourceString<ModuleOp>(moduleStr, &context);146 ASSERT_TRUE(!!module);147 148 // Create an NVVM target.149 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);150 151 // Serialize the module.152 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);153 ASSERT_TRUE(!!serializer);154 gpu::TargetOptions options("", {}, "", "", gpu::CompilationTarget::Binary);155 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {156 std::optional<SmallVector<char, 0>> object =157 serializer.serializeToObject(gpuModule, options);158 // Check that the serializer was successful.159 ASSERT_TRUE(object != std::nullopt);160 ASSERT_TRUE(!object->empty());161 }162}163 164// Test callback functions invoked with LLVM IR and ISA.165TEST_F(MLIRTargetLLVMNVVM,166 SKIP_WITHOUT_NVPTX(CallbackInvokedWithLLVMIRAndISA)) {167 MLIRContext context(registry);168 169 OwningOpRef<ModuleOp> module =170 parseSourceString<ModuleOp>(moduleStr, &context);171 ASSERT_TRUE(!!module);172 173 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);174 175 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);176 ASSERT_TRUE(!!serializer);177 178 std::string initialLLVMIR;179 auto initialCallback = [&initialLLVMIR](llvm::Module &module) {180 llvm::raw_string_ostream ros(initialLLVMIR);181 module.print(ros, nullptr);182 };183 184 std::string linkedLLVMIR;185 auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {186 llvm::raw_string_ostream ros(linkedLLVMIR);187 module.print(ros, nullptr);188 };189 190 std::string optimizedLLVMIR;191 auto optimizedCallback = [&optimizedLLVMIR](llvm::Module &module) {192 llvm::raw_string_ostream ros(optimizedLLVMIR);193 module.print(ros, nullptr);194 };195 196 std::string isaResult;197 auto isaCallback = [&isaResult](llvm::StringRef isa) {198 isaResult = isa.str();199 };200 201 gpu::TargetOptions options({}, {}, {}, {}, gpu::CompilationTarget::Assembly,202 {}, initialCallback, linkedCallback,203 optimizedCallback, isaCallback);204 205 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {206 std::optional<SmallVector<char, 0>> object =207 serializer.serializeToObject(gpuModule, options);208 209 ASSERT_TRUE(object != std::nullopt);210 ASSERT_TRUE(!object->empty());211 ASSERT_TRUE(!initialLLVMIR.empty());212 ASSERT_TRUE(!linkedLLVMIR.empty());213 ASSERT_TRUE(!optimizedLLVMIR.empty());214 ASSERT_TRUE(!isaResult.empty());215 216 initialLLVMIR.clear();217 linkedLLVMIR.clear();218 optimizedLLVMIR.clear();219 isaResult.clear();220 }221}222 223// Test linking LLVM IR from a resource attribute.224TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(LinkedLLVMIRResource)) {225 MLIRContext context(registry);226 std::string moduleStr = R"mlir(227 gpu.module @nvvm_test {228 llvm.func @bar()229 llvm.func @nvvm_kernel(%arg0: f32) attributes {gpu.kernel, nvvm.kernel} {230 llvm.call @bar() : () -> ()231 llvm.return232 }233 }234 )mlir";235 // Provide the library to link as a serialized bitcode blob.236 SmallVector<char> bitcodeToLink;237 {238 std::string linkedLib = R"llvm(239 define void @bar() {240 ret void241 }242 )llvm";243 llvm::SMDiagnostic err;244 llvm::MemoryBufferRef buffer(linkedLib, "linkedLib");245 llvm::LLVMContext llvmCtx;246 std::unique_ptr<llvm::Module> module = llvm::parseIR(buffer, err, llvmCtx);247 ASSERT_TRUE(module) << " Can't parse IR: " << err.getMessage();248 {249 llvm::raw_svector_ostream os(bitcodeToLink);250 WriteBitcodeToFile(*module, os);251 }252 }253 254 OwningOpRef<ModuleOp> module =255 parseSourceString<ModuleOp>(moduleStr, &context);256 ASSERT_TRUE(!!module);257 Builder builder(&context);258 259 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);260 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);261 262 // Hook to intercept the LLVM IR after linking external libs.263 std::string linkedLLVMIR;264 auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {265 llvm::raw_string_ostream ros(linkedLLVMIR);266 module.print(ros, nullptr);267 };268 269 // Store the bitcode as a DenseI8ArrayAttr.270 SmallVector<Attribute> librariesToLink;271 librariesToLink.push_back(DenseI8ArrayAttr::get(272 &context,273 ArrayRef<int8_t>((int8_t *)bitcodeToLink.data(), bitcodeToLink.size())));274 gpu::TargetOptions options({}, librariesToLink, {}, {},275 gpu::CompilationTarget::Assembly, {}, {},276 linkedCallback);277 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {278 std::optional<SmallVector<char, 0>> object =279 serializer.serializeToObject(gpuModule, options);280 281 // Verify that we correctly linked in the library: the external call is282 // replaced by the definition.283 ASSERT_TRUE(!linkedLLVMIR.empty());284 {285 llvm::SMDiagnostic err;286 llvm::MemoryBufferRef buffer(linkedLLVMIR, "linkedLLVMIR");287 llvm::LLVMContext llvmCtx;288 std::unique_ptr<llvm::Module> module =289 llvm::parseIR(buffer, err, llvmCtx);290 ASSERT_TRUE(module) << " Can't parse linkedLLVMIR: " << err.getMessage()291 << " IR: \n\b" << linkedLLVMIR;292 llvm::Function *bar = module->getFunction("bar");293 ASSERT_TRUE(bar);294 ASSERT_FALSE(bar->empty());295 }296 ASSERT_TRUE(object != std::nullopt);297 ASSERT_TRUE(!object->empty());298 }299}300