llvm-project/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

Ignoring revisions in .git-blame-ignore-revs. Click here to bypass and see the normal blame view.

333 lines
12 KiB
C++
Raw Normal View History

//===- mlir-rocm-runner.cpp - MLIR ROCM Execution Driver-------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This is a command line utility that executes an MLIR file on the GPU by
// translating MLIR to ROCDL/LLVM IR before JIT-compiling and executing the
// latter.
//
//===----------------------------------------------------------------------===//
#include "llvm/ADT/STLExtras.h"
#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
#include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
#include "mlir/Dialect/GPU/GPUDialect.h"
#include "mlir/Dialect/GPU/Passes.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
#include "mlir/ExecutionEngine/JitRunner.h"
#include "mlir/ExecutionEngine/OptUtils.h"
#include "mlir/IR/Function.h"
#include "mlir/IR/Module.h"
#include "mlir/InitAllDialects.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Support/FileUtilities.h"
#include "mlir/Target/ROCDLIR.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/Passes.h"
#include "llvm/Support/ErrorOr.h"
#include "llvm/Support/FileUtilities.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/LineIterator.h"
#include "llvm/Support/Program.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/TargetRegistry.h"
#include "llvm/Support/TargetSelect.h"
// MC headers.
#include "llvm/MC/MCAsmBackend.h"
#include "llvm/MC/MCAsmInfo.h"
#include "llvm/MC/MCCodeEmitter.h"
#include "llvm/MC/MCContext.h"
#include "llvm/MC/MCInstPrinter.h"
#include "llvm/MC/MCInstrInfo.h"
#include "llvm/MC/MCObjectFileInfo.h"
#include "llvm/MC/MCObjectWriter.h"
#include "llvm/MC/MCParser/AsmLexer.h"
#include "llvm/MC/MCParser/MCTargetAsmParser.h"
#include "llvm/MC/MCRegisterInfo.h"
#include "llvm/MC/MCStreamer.h"
#include "llvm/MC/MCSubtargetInfo.h"
#include "llvm/MC/MCTargetOptionsCommandFlags.h"
// lld headers.
#include "lld/Common/Driver.h"
// HIP headers.
#include "hip/hip_version.h"
using namespace mlir;
using namespace llvm;
using Blob = SmallVector<char, 0>;
static cl::opt<std::string> tripleName("triple", cl::desc("target triple"),
cl::value_desc("triple string"),
cl::init("amdgcn-amd-amdhsa"));
static cl::opt<std::string> targetChip("target", cl::desc("target chip"),
cl::value_desc("AMDGPU ISA version"),
cl::init(""));
static cl::opt<std::string> features("feature", cl::desc("target features"),
cl::value_desc("AMDGPU target features"),
cl::init(""));
static constexpr const char kRunnerProgram[] = "mlir-rocm-runner";
static constexpr const char kRocmAgentEnumerator[] = "rocm_agent_enumerator";
static constexpr const char kDefaultTargetChip[] = "gfx900";
static LogicalResult assembleIsa(const std::string isa, StringRef name,
Blob &result) {
raw_svector_ostream os(result);
std::string error;
Triple theTriple(Triple::normalize(tripleName));
const Target *theTarget =
TargetRegistry::lookupTarget(theTriple.normalize(), error);
if (!theTarget) {
WithColor::error(errs(), name) << error;
return failure();
}
SourceMgr srcMgr;
srcMgr.AddNewSourceBuffer(MemoryBuffer::getMemBuffer(isa), SMLoc());
const MCTargetOptions mcOptions;
std::unique_ptr<MCRegisterInfo> mri(theTarget->createMCRegInfo(tripleName));
std::unique_ptr<MCAsmInfo> mai(
theTarget->createMCAsmInfo(*mri, tripleName, mcOptions));
mai->setRelaxELFRelocations(true);
MCObjectFileInfo mofi;
MCContext ctx(mai.get(), mri.get(), &mofi, &srcMgr, &mcOptions);
mofi.InitMCObjectFileInfo(theTriple, false, ctx, false);
SmallString<128> cwd;
if (!sys::fs::current_path(cwd))
ctx.setCompilationDir(cwd);
std::unique_ptr<MCStreamer> mcStreamer;
std::unique_ptr<MCInstrInfo> mcii(theTarget->createMCInstrInfo());
std::unique_ptr<MCSubtargetInfo> sti(
theTarget->createMCSubtargetInfo(tripleName, targetChip, features));
MCCodeEmitter *ce = theTarget->createMCCodeEmitter(*mcii, *mri, ctx);
MCAsmBackend *mab = theTarget->createMCAsmBackend(*sti, *mri, mcOptions);
mcStreamer.reset(theTarget->createMCObjectStreamer(
theTriple, ctx, std::unique_ptr<MCAsmBackend>(mab),
mab->createObjectWriter(os), std::unique_ptr<MCCodeEmitter>(ce), *sti,
mcOptions.MCRelaxAll, mcOptions.MCIncrementalLinkerCompatible,
/*DWARFMustBeAtTheEnd*/ false));
mcStreamer->setUseAssemblerInfoForParsing(true);
std::unique_ptr<MCAsmParser> parser(
createMCAsmParser(srcMgr, ctx, *mcStreamer, *mai));
std::unique_ptr<MCTargetAsmParser> tap(
theTarget->createMCAsmParser(*sti, *parser, *mcii, mcOptions));
if (!tap) {
WithColor::error(errs(), name) << "assembler initialization error.\n";
return failure();
}
parser->setTargetParser(*tap);
parser->Run(false);
return success();
}
static LogicalResult createHsaco(const Blob &isaBlob, StringRef name,
Blob &hsacoBlob) {
// Save the ISA binary to a temp file.
int tempIsaBinaryFd = -1;
SmallString<128> tempIsaBinaryFilename;
std::error_code ec = sys::fs::createTemporaryFile(
"kernel", "o", tempIsaBinaryFd, tempIsaBinaryFilename);
if (ec) {
WithColor::error(errs(), name)
<< "temporary file for ISA binary creation error.\n";
return failure();
}
FileRemover cleanupIsaBinary(tempIsaBinaryFilename);
raw_fd_ostream tempIsaBinaryOs(tempIsaBinaryFd, true);
tempIsaBinaryOs << isaBlob;
tempIsaBinaryOs.close();
// Create a temp file for HSA code object.
int tempHsacoFD = -1;
SmallString<128> tempHsacoFilename;
ec = sys::fs::createTemporaryFile("kernel", "hsaco", tempHsacoFD,
tempHsacoFilename);
if (ec) {
WithColor::error(errs(), name)
<< "temporary file for HSA code object creation error.\n";
return failure();
}
FileRemover cleanupHsaco(tempHsacoFilename);
// Invoke lld. Expect a true return value from lld.
bool ret = lld::elf::link({"ld.lld", "-shared", tempIsaBinaryFilename.c_str(),
"-o", tempHsacoFilename.c_str()},
/*canEarlyExit=*/false, llvm::outs(), llvm::errs());
if (!ret) {
WithColor::error(errs(), name) << "lld invocation error.\n";
return failure();
}
// Load the HSA code object.
auto hsacoFile = mlir::openInputFile(tempHsacoFilename);
if (!hsacoFile) {
WithColor::error(errs(), name)
<< "read HSA code object from temp file error.\n";
return failure();
}
hsacoBlob.assign(hsacoFile->getBuffer().begin(),
hsacoFile->getBuffer().end());
return success();
}
static std::unique_ptr<llvm::Module> compileModuleToROCDLIR(Operation *m) {
auto llvmModule = translateModuleToROCDLIR(m);
// TODO(whchung): Link with ROCm-Device-Libs in case needed (ex: the Module
// depends on math functions).
return llvmModule;
}
static OwnedBlob compileISAToHsaco(const std::string isa, Location loc,
StringRef name) {
// ISA -> ISA in binary form via MC.
// Use lld to create HSA code object.
Blob isaBlob;
Blob hsacoBlob;
if (succeeded(assembleIsa(isa, name, isaBlob)) &&
succeeded(createHsaco(isaBlob, name, hsacoBlob)))
return std::make_unique<std::vector<char>>(hsacoBlob.begin(),
hsacoBlob.end());
WithColor::error(errs(), name) << "producing HSA code object error.\n";
return {};
}
static void configTargetChip() {
// Set targetChip to default value first.
targetChip = kDefaultTargetChip;
// Locate rocm_agent_enumerator.
llvm::ErrorOr<std::string> rocmAgentEnumerator = llvm::sys::findProgramByName(
kRocmAgentEnumerator, {__ROCM_PATH__ "/bin"});
std::error_code ec;
if ((ec = rocmAgentEnumerator.getError())) {
WithColor::warning(errs(), kRunnerProgram)
<< kRocmAgentEnumerator << " couldn't be located under "
<< __ROCM_PATH__ << ", set target as " << kDefaultTargetChip << "\n";
return;
}
// Prepare temp file to hold the outputs.
int tempFd = -1;
SmallString<128> tempFilename;
ec = sys::fs::createTemporaryFile("rocm_agent", "txt", tempFd, tempFilename);
if (ec) {
WithColor::warning(errs(), kRunnerProgram)
<< "temporary file for " << kRocmAgentEnumerator
<< " creation error, set target as " << kDefaultTargetChip << "\n";
return;
}
FileRemover cleanup(tempFilename);
// Invoke rocm_agent_enumerator.
std::string errorMessage;
SmallVector<StringRef, 2> args{"-t", "GPU"};
Optional<StringRef> redirects[3] = {{""}, tempFilename.str(), {""}};
int result =
llvm::sys::ExecuteAndWait(rocmAgentEnumerator.get(), args, llvm::None,
redirects, 0, 0, &errorMessage);
if (result) {
WithColor::warning(errs(), kRunnerProgram)
<< kRocmAgentEnumerator << " invocation error: " << errorMessage
<< ", set target as " << kDefaultTargetChip << "\n";
return;
}
// Load and parse the result.
auto gfxIsaList = mlir::openInputFile(tempFilename);
if (!gfxIsaList) {
WithColor::error(errs(), kRunnerProgram)
<< "read ROCm agent list temp file error, set target as "
<< kDefaultTargetChip << "\n";
return;
}
for (line_iterator lines(*gfxIsaList); !lines.is_at_end(); ++lines) {
// Skip the line with content "gfx000".
if (*lines == "gfx000")
continue;
// Use the first ISA version found.
targetChip = lines->str();
break;
}
}
static void configTargetFeatures() {
if (features.size() > 0)
features += ",";
// After ROCm 3.5, adopt HSA code object V3.
if (HIP_VERSION_MAJOR >= 3 && HIP_VERSION_MINOR >= 5)
features += "+code-object-v3";
else
features += "-code-object-v3";
}
static LogicalResult runMLIRPasses(ModuleOp m) {
PassManager pm(m.getContext());
applyPassManagerCLOptions(pm);
// Configure target chip ISA version if it has not been specified.
if (!targetChip.size())
configTargetChip();
// Configure target features per ROCm / HIP version.
configTargetFeatures();
pm.addPass(createGpuKernelOutliningPass());
auto &kernelPm = pm.nest<gpu::GPUModuleOp>();
kernelPm.addPass(createStripDebugInfoPass());
kernelPm.addPass(createLowerGpuOpsToROCDLOpsPass());
kernelPm.addPass(createConvertGPUKernelToBlobPass(
compileModuleToROCDLIR, compileISAToHsaco, tripleName, targetChip,
features, /*gpuBinaryAnnotation=*/"rocdl.hsaco"));
pm.addPass(createLowerToLLVMPass());
pm.addPass(createConvertGpuLaunchFuncToGpuRuntimeCallsPass(
/*gpuBinaryAnnotation=*/"rocdl.hsaco"));
return pm.run(m);
}
int main(int argc, char **argv) {
registerPassManagerCLOptions();
mlir::registerAllDialects();
llvm::InitLLVM y(argc, argv);
llvm::InitializeAllTargetInfos();
llvm::InitializeAllTargetMCs();
llvm::InitializeAllAsmParsers();
// Initialize LLVM AMDGPU backend.
LLVMInitializeAMDGPUTarget();
LLVMInitializeAMDGPUTargetInfo();
LLVMInitializeAMDGPUTargetMC();
LLVMInitializeAMDGPUAsmPrinter();
mlir::initializeLLVMPasses();
return mlir::JitRunnerMain(argc, argv, &runMLIRPasses);
}