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-exception
6//
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/IRReader/IRReader.h"
22#include "llvm/Support/MemoryBufferRef.h"
23#include "llvm/Support/Process.h"
24#include "llvm/Support/TargetSelect.h"
25#include "llvm/Support/raw_ostream.h"
26#include "llvm/TargetParser/Host.h"
27
28#include "gmock/gmock.h"
29
30using namespace mlir;
31
32// Skip the test if the NVPTX target was not built.
33#if MLIR_ENABLE_CUDA_CONVERSIONS
34#define SKIP_WITHOUT_NVPTX(x) x
35#else
36#define SKIP_WITHOUT_NVPTX(x) DISABLED_##x
37#endif
38
39class MLIRTargetLLVMNVVM : public ::testing::Test {
40protected:
41 void SetUp() override {
42 registerBuiltinDialectTranslation(registry);
43 registerLLVMDialectTranslation(registry);
44 registerGPUDialectTranslation(registry);
45 registerNVVMDialectTranslation(registry);
46 NVVM::registerNVVMTargetInterfaceExternalModels(registry);
47 }
48
49 // Checks if PTXAS is in PATH.
50 bool hasPtxas() {
51 // Find the `ptxas` compiler.
52 std::optional<std::string> ptxasCompiler =
53 llvm::sys::Process::FindInEnvPath(EnvName: "PATH", FileName: "ptxas");
54 return ptxasCompiler.has_value();
55 }
56
57 // Dialect registry.
58 DialectRegistry registry;
59
60 // MLIR module used for the tests.
61 const std::string moduleStr = R"mlir(
62 gpu.module @nvvm_test {
63 llvm.func @nvvm_kernel(%arg0: f32) attributes {gpu.kernel, nvvm.kernel} {
64 llvm.return
65 }
66 })mlir";
67};
68
69// Test NVVM serialization to LLVM.
70TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMMToLLVM)) {
71 MLIRContext context(registry);
72
73 OwningOpRef<ModuleOp> module =
74 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
75 ASSERT_TRUE(!!module);
76
77 // Create an NVVM target.
78 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
79
80 // Serialize the module.
81 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
82 ASSERT_TRUE(!!serializer);
83 gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Offload);
84 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
85 std::optional<SmallVector<char, 0>> object =
86 serializer.serializeToObject(gpuModule, options);
87 // Check that the serializer was successful.
88 ASSERT_TRUE(object != std::nullopt);
89 ASSERT_TRUE(!object->empty());
90
91 // Read the serialized module.
92 llvm::MemoryBufferRef buffer(StringRef(object->data(), object->size()),
93 "module");
94 llvm::LLVMContext llvmContext;
95 llvm::Expected<std::unique_ptr<llvm::Module>> llvmModule =
96 llvm::getLazyBitcodeModule(buffer, llvmContext);
97 ASSERT_TRUE(!!llvmModule);
98 ASSERT_TRUE(!!*llvmModule);
99
100 // Check that it has a function named `foo`.
101 ASSERT_TRUE((*llvmModule)->getFunction("nvvm_kernel") != nullptr);
102 }
103}
104
105// Test NVVM serialization to PTX.
106TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToPTX)) {
107 MLIRContext context(registry);
108
109 OwningOpRef<ModuleOp> module =
110 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
111 ASSERT_TRUE(!!module);
112
113 // Create an NVVM target.
114 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
115
116 // Serialize the module.
117 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
118 ASSERT_TRUE(!!serializer);
119 gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Assembly);
120 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
121 std::optional<SmallVector<char, 0>> object =
122 serializer.serializeToObject(gpuModule, options);
123 // Check that the serializer was successful.
124 ASSERT_TRUE(object != std::nullopt);
125 ASSERT_TRUE(!object->empty());
126
127 ASSERT_TRUE(
128 StringRef(object->data(), object->size()).contains("nvvm_kernel"));
129 }
130}
131
132// Test NVVM serialization to Binary.
133TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) {
134 if (!hasPtxas())
135 GTEST_SKIP() << "PTXAS compiler not found, skipping test.";
136
137 MLIRContext context(registry);
138
139 OwningOpRef<ModuleOp> module =
140 parseSourceString<ModuleOp>(sourceStr: moduleStr, config: &context);
141 ASSERT_TRUE(!!module);
142
143 // Create an NVVM target.
144 NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
145
146 // Serialize the module.
147 auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
148 ASSERT_TRUE(!!serializer);
149 gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary);
150 for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
151 std::optional<SmallVector<char, 0>> object =
152 serializer.serializeToObject(gpuModule, options);
153 // Check that the serializer was successful.
154 ASSERT_TRUE(object != std::nullopt);
155 ASSERT_TRUE(!object->empty());
156 }
157}
158

source code of mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp