419 lines · cpp
1//===- Target.cpp - MLIR LLVM XeVM target compilation -----------*- C++ -*-===//2//3// Part of the LLVM Project, 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// This files defines XeVM target related functions including registration10// calls for the `#xevm.target` compilation attribute.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Target/LLVM/XeVM/Target.h"15 16#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"17#include "mlir/Dialect/GPU/IR/GPUDialect.h"18#include "mlir/Dialect/LLVMIR/XeVMDialect.h"19#include "mlir/IR/BuiltinAttributeInterfaces.h"20#include "mlir/IR/BuiltinDialect.h"21#include "mlir/IR/BuiltinTypes.h"22#include "mlir/IR/DialectResourceBlobManager.h"23#include "mlir/Target/LLVM/XeVM/Utils.h"24#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"25#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"26#include "mlir/Target/LLVMIR/Dialect/XeVM/XeVMToLLVMIRTranslation.h"27#include "mlir/Target/LLVMIR/Export.h"28#include "llvm/IR/LegacyPassManager.h"29#include "llvm/Target/TargetMachine.h"30 31#include "llvm/Bitcode/BitcodeWriter.h"32#include "llvm/Config/Targets.h"33#include "llvm/Support/FileSystem.h"34#include "llvm/Support/FileUtilities.h"35#include "llvm/Support/FormatVariadic.h"36#include "llvm/Support/MemoryBuffer.h"37#include "llvm/Support/Path.h"38#include "llvm/Support/Process.h"39#include "llvm/Support/Program.h"40#include "llvm/Support/TargetSelect.h"41#include "llvm/Support/raw_ostream.h"42 43#include <cstdint>44#include <cstdlib>45 46using namespace mlir;47using namespace mlir::xevm;48 49namespace {50// XeVM implementation of the gpu:TargetAttrInterface.51class XeVMTargetAttrImpl52 : public gpu::TargetAttrInterface::FallbackModel<XeVMTargetAttrImpl> {53public:54 std::optional<SmallVector<char, 0>>55 serializeToObject(Attribute attribute, Operation *module,56 const gpu::TargetOptions &options) const;57 58 Attribute createObject(Attribute attribute, Operation *module,59 const SmallVector<char, 0> &object,60 const gpu::TargetOptions &options) const;61};62} // namespace63 64void mlir::xevm::registerXeVMTargetInterfaceExternalModels(65 DialectRegistry ®istry) {66 registry.addExtension(+[](MLIRContext *ctx, XeVMDialect *dialect) {67 XeVMTargetAttr::attachInterface<XeVMTargetAttrImpl>(*ctx);68 });69}70 71void mlir::xevm::registerXeVMTargetInterfaceExternalModels(72 MLIRContext &context) {73 DialectRegistry registry;74 registerXeVMTargetInterfaceExternalModels(registry);75 context.appendDialectRegistry(registry);76}77 78SerializeGPUModuleBase::SerializeGPUModuleBase(79 Operation &module, XeVMTargetAttr xeTarget,80 const gpu::TargetOptions &targetOptions)81 : ModuleToObject(module, xeTarget.getTriple(), "", {}, xeTarget.getO()),82 xeTarget(xeTarget), librariesToLink(targetOptions.getLibrariesToLink()),83 targetOptions(targetOptions) {84 if (xeTarget.getLinkFiles())85 librariesToLink.append(xeTarget.getLinkFiles().begin(),86 xeTarget.getLinkFiles().end());87}88 89XeVMTargetAttr SerializeGPUModuleBase::getTarget() const { return xeTarget; }90 91std::optional<SmallVector<std::unique_ptr<llvm::Module>>>92SerializeGPUModuleBase::loadBitcodeFiles(llvm::Module &module) {93 if (librariesToLink.empty())94 return SmallVector<std::unique_ptr<llvm::Module>>();95 SmallVector<std::unique_ptr<llvm::Module>> bcFiles;96 if (failed(loadBitcodeFilesFromList(module.getContext(), librariesToLink,97 bcFiles)))98 return std::nullopt;99 return std::move(bcFiles);100}101 102gpu::GPUModuleOp SerializeGPUModuleBase::getGPUModuleOp() {103 return dyn_cast<gpu::GPUModuleOp>(&SerializeGPUModuleBase::getOperation());104}105 106// There is 1 way to finalize IL to native code: IGC107// There are 2 ways to access IGC: AOT (ocloc) and JIT (L0 runtime).108// - L0 runtime consumes IL and is external to MLIR codebase (rt wrappers).109// - `ocloc` tool can be "queried" from within MLIR.110std::optional<SmallVector<char, 0>>111SerializeGPUModuleBase::compileToBinary(const std::string &asmStr,112 StringRef inputFormat) {113 using TmpFile = std::pair<llvm::SmallString<128>, llvm::FileRemover>;114 // Find the `ocloc` tool.115 std::optional<std::string> oclocCompiler = findTool("ocloc");116 if (!oclocCompiler)117 return std::nullopt;118 Location loc = getGPUModuleOp().getLoc();119 std::string basename = llvm::formatv(120 "mlir-{0}-{1}-{2}", getGPUModuleOp().getNameAttr().getValue(),121 getTarget().getTriple(), getTarget().getChip());122 123 auto createTemp = [&](StringRef name,124 StringRef suffix) -> std::optional<TmpFile> {125 llvm::SmallString<128> filePath;126 if (auto ec = llvm::sys::fs::createTemporaryFile(name, suffix, filePath)) {127 getGPUModuleOp().emitError()128 << "Couldn't create the temp file: `" << filePath129 << "`, error message: " << ec.message();130 return std::nullopt;131 }132 return TmpFile(filePath, llvm::FileRemover(filePath.c_str()));133 };134 // Create temp file135 std::optional<TmpFile> asmFile = createTemp(basename, "asm");136 std::optional<TmpFile> binFile = createTemp(basename, "");137 std::optional<TmpFile> logFile = createTemp(basename, "log");138 if (!logFile || !asmFile || !binFile)139 return std::nullopt;140 // Dump the assembly to a temp file141 std::error_code ec;142 {143 llvm::raw_fd_ostream asmStream(asmFile->first, ec);144 if (ec) {145 emitError(loc) << "Couldn't open the file: `" << asmFile->first146 << "`, error message: " << ec.message();147 return std::nullopt;148 }149 asmStream << asmStr;150 if (asmStream.has_error()) {151 emitError(loc) << "An error occurred while writing the assembly to: `"152 << asmFile->first << "`.";153 return std::nullopt;154 }155 asmStream.flush();156 }157 // Set cmd options158 std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>> cmdOpts =159 targetOptions.tokenizeCmdOptions();160 // Example: --gpu-module-to-binary="opts='opt1 opt2'"161 const std::string cmdOptsStr = "\"" + llvm::join(cmdOpts.second, " ") + "\"";162 SmallVector<StringRef, 12> oclocArgs(163 {"ocloc", "compile", "-file", asmFile->first, inputFormat, "-device",164 getTarget().getChip(), "-output", binFile->first, "-output_no_suffix",165 "-options", cmdOptsStr});166 167// Dump tool invocation commands.168#define DEBUG_TYPE "serialize-to-binary"169 LLVM_DEBUG({170 llvm::dbgs() << "Tool invocation for module: "171 << getGPUModuleOp().getNameAttr() << "\n";172 llvm::interleave(oclocArgs, llvm::dbgs(), " ");173 llvm::dbgs() << "\n";174 });175#undef DEBUG_TYPE176 // Helper function for printing tool error logs.177 std::string message;178 auto emitLogError =179 [&](StringRef toolName) -> std::optional<SmallVector<char, 0>> {180 if (message.empty()) {181 llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> toolStderr =182 llvm::MemoryBuffer::getFile(logFile->first);183 if (toolStderr)184 emitError(loc) << toolName << " invocation failed. Log:\n"185 << toolStderr->get()->getBuffer();186 else187 emitError(loc) << toolName << " invocation failed.";188 return std::nullopt;189 }190 emitError(loc) << toolName191 << " invocation failed, error message: " << message;192 return std::nullopt;193 };194 std::optional<StringRef> redirects[] = {195 std::nullopt,196 logFile->first,197 logFile->first,198 };199 // Invoke ocloc.200 if (llvm::sys::ExecuteAndWait(oclocCompiler.value(), oclocArgs, std::nullopt,201 redirects, 0, 0, &message))202 return emitLogError("`ocloc`");203 binFile->first.append(".bin");204 llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> binaryBuffer =205 llvm::MemoryBuffer::getFile(binFile->first);206 if (!binaryBuffer) {207 emitError(loc) << "Couldn't open the file: `" << binFile->first208 << "`, error message: " << binaryBuffer.getError().message();209 return std::nullopt;210 }211 StringRef bin = (*binaryBuffer)->getBuffer();212 return SmallVector<char, 0>(bin.begin(), bin.end());213}214 215std::optional<std::string> SerializeGPUModuleBase::findTool(StringRef tool) {216 // 1. Check the toolkit path given in the command line.217 StringRef pathRef = targetOptions.getToolkitPath();218 SmallVector<char, 256> path;219 if (!pathRef.empty()) {220 path.insert(path.begin(), pathRef.begin(), pathRef.end());221 llvm::sys::path::append(path, "bin", tool);222 if (llvm::sys::fs::can_execute(path))223 return StringRef(path.data(), path.size()).str();224 }225 // 2. Check PATH.226 if (std::optional<std::string> toolPath =227 llvm::sys::Process::FindInEnvPath("PATH", tool))228 return *toolPath;229 230 getGPUModuleOp().emitError()231 << "Couldn't find the `" << tool232 << "` binary. Please specify the toolkit "233 "path via GpuModuleToBinaryPass or add the compiler to $PATH`.";234 return std::nullopt;235}236 237namespace {238class SPIRVSerializer : public SerializeGPUModuleBase {239public:240 SPIRVSerializer(Operation &module, XeVMTargetAttr xeTarget,241 const gpu::TargetOptions &targetOptions)242 : SerializeGPUModuleBase(module, xeTarget, targetOptions) {}243 244 static void init();245 246 /// Serializes the LLVM module to an object format, depending on the247 /// compilation target selected in target options.248 std::optional<SmallVector<char, 0>>249 moduleToObject(llvm::Module &llvmModule) override;250 251private:252 /// Translates the LLVM module to SPIR-V binary using LLVM's253 /// SPIR-V target.254 std::optional<std::string>255 translateToSPIRVBinary(llvm::Module &llvmModule,256 llvm::TargetMachine &targetMachine);257};258} // namespace259 260void SPIRVSerializer::init() {261 static llvm::once_flag initializeBackendOnce;262 llvm::call_once(initializeBackendOnce, []() {263#if LLVM_HAS_SPIRV_TARGET264 LLVMInitializeSPIRVTarget();265 LLVMInitializeSPIRVTargetInfo();266 LLVMInitializeSPIRVTargetMC();267 LLVMInitializeSPIRVAsmPrinter();268#endif269 });270}271 272std::optional<SmallVector<char, 0>>273SPIRVSerializer::moduleToObject(llvm::Module &llvmModule) {274#define DEBUG_TYPE "serialize-to-llvm"275 LLVM_DEBUG({276 llvm::dbgs() << "LLVM IR for module: " << getGPUModuleOp().getNameAttr()277 << "\n";278 llvm::dbgs() << llvmModule << "\n";279 llvm::dbgs().flush();280 });281#undef DEBUG_TYPE282 283 // Return LLVM IR if the compilation target is `offload`.284 if (targetOptions.getCompilationTarget() == gpu::CompilationTarget::Offload)285 return SerializeGPUModuleBase::moduleToObject(llvmModule);286 287#if !LLVM_HAS_SPIRV_TARGET288 getGPUModuleOp()->emitError("The `SPIRV` target was not built. Please enable "289 "it when building LLVM.");290 return std::nullopt;291#endif // LLVM_HAS_SPIRV_TARGET292 293 std::optional<llvm::TargetMachine *> targetMachine =294 getOrCreateTargetMachine();295 if (!targetMachine) {296 getGPUModuleOp().emitError() << "Target Machine unavailable for triple "297 << triple << ", can't optimize with LLVM\n";298 return std::nullopt;299 }300 301 // Return SPIRV if the compilation target is `assembly`.302 if (targetOptions.getCompilationTarget() ==303 gpu::CompilationTarget::Assembly) {304 std::optional<std::string> serializedISA =305 translateToISA(llvmModule, **targetMachine);306 if (!serializedISA) {307 getGPUModuleOp().emitError() << "Failed translating the module to ISA."308 << triple << ", can't compile with LLVM\n";309 return std::nullopt;310 }311 312#define DEBUG_TYPE "serialize-to-isa"313 LLVM_DEBUG({314 llvm::dbgs() << "SPIR-V for module: " << getGPUModuleOp().getNameAttr()315 << "\n";316 llvm::dbgs() << *serializedISA << "\n";317 llvm::dbgs().flush();318 });319#undef DEBUG_TYPE320 321 // Make sure to include the null terminator.322 StringRef bin(serializedISA->c_str(), serializedISA->size() + 1);323 return SmallVector<char, 0>(bin.begin(), bin.end());324 }325 326 // Level zero runtime is set up to accept SPIR-V binary327 // translateToSPIRVBinary translates the LLVM module to SPIR-V binary328 // using LLVM's SPIRV target.329 // compileToBinary can be used in the future if level zero runtime330 // implementation switches to native XeVM binary format.331 std::optional<std::string> serializedSPIRVBinary =332 translateToSPIRVBinary(llvmModule, **targetMachine);333 if (!serializedSPIRVBinary) {334 getGPUModuleOp().emitError() << "Failed translating the module to Binary.";335 return std::nullopt;336 }337 if (serializedSPIRVBinary->size() % 4) {338 getGPUModuleOp().emitError() << "SPIRV code size must be a multiple of 4.";339 return std::nullopt;340 }341 StringRef bin(serializedSPIRVBinary->c_str(), serializedSPIRVBinary->size());342 return SmallVector<char, 0>(bin.begin(), bin.end());343}344 345std::optional<std::string>346SPIRVSerializer::translateToSPIRVBinary(llvm::Module &llvmModule,347 llvm::TargetMachine &targetMachine) {348 std::string targetISA;349 llvm::raw_string_ostream stream(targetISA);350 351 { // Drop pstream after this to prevent the ISA from being stuck buffering352 llvm::buffer_ostream pstream(stream);353 llvm::legacy::PassManager codegenPasses;354 if (targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,355 llvm::CodeGenFileType::ObjectFile))356 return std::nullopt;357 358 codegenPasses.run(llvmModule);359 }360 return targetISA;361}362 363std::optional<SmallVector<char, 0>>364XeVMTargetAttrImpl::serializeToObject(Attribute attribute, Operation *module,365 const gpu::TargetOptions &options) const {366 if (!module)367 return std::nullopt;368 auto gpuMod = dyn_cast<gpu::GPUModuleOp>(module);369 if (!gpuMod) {370 module->emitError("expected to be a gpu.module op");371 return std::nullopt;372 }373 auto xeTarget = cast<XeVMTargetAttr>(attribute);374 if (xeTarget.getTriple().starts_with("spirv")) {375 gpuMod.walk([&](LLVM::LLVMFuncOp funcOp) {376 if (funcOp->hasAttr(gpu::GPUDialect::getKernelFuncAttrName())) {377 funcOp.setIntelReqdSubGroupSize(16);378 return WalkResult::interrupt();379 }380 return WalkResult::advance();381 });382 383 SPIRVSerializer serializer(*module, cast<XeVMTargetAttr>(attribute),384 options);385 serializer.init();386 387#if !LLVM_HAS_SPIRV_TARGET388 module->emitError("Cannot run `TargetRegistry::lookupTarget()` for SPIRV "389 "without having the target built.");390#endif391 392 return serializer.run();393 }394 module->emitError("Unsupported XeVM target triple: ") << xeTarget.getTriple();395 return std::nullopt;396}397 398Attribute399XeVMTargetAttrImpl::createObject(Attribute attribute, Operation *module,400 const SmallVector<char, 0> &object,401 const gpu::TargetOptions &options) const {402 Builder builder(attribute.getContext());403 gpu::CompilationTarget format = options.getCompilationTarget();404 auto xeTarget = cast<XeVMTargetAttr>(attribute);405 SmallVector<NamedAttribute, 2> properties;406 if (format == gpu::CompilationTarget::Assembly)407 properties.push_back(408 builder.getNamedAttr("O", builder.getI32IntegerAttr(xeTarget.getO())));409 410 DictionaryAttr objectProps;411 if (!properties.empty())412 objectProps = builder.getDictionaryAttr(properties);413 414 return builder.getAttr<gpu::ObjectAttr>(415 attribute, format,416 builder.getStringAttr(StringRef(object.data(), object.size())),417 objectProps, /*kernels=*/nullptr);418}419