2022-08-18 01:42:48 +08:00
|
|
|
#include "triton/Target/PTX/PTXTranslation.h"
|
|
|
|
#include "triton/Target/LLVMIR/LLVMIRTranslation.h"
|
2022-09-26 16:38:06 -07:00
|
|
|
|
|
|
|
#include "llvm/IR/IRBuilder.h"
|
|
|
|
#include "llvm/IR/LegacyPassManager.h"
|
|
|
|
#include "llvm/IR/Module.h"
|
|
|
|
#include "llvm/IR/Verifier.h"
|
|
|
|
#include "llvm/MC/TargetRegistry.h"
|
|
|
|
#include "llvm/Support/TargetSelect.h"
|
|
|
|
#include "llvm/Target/TargetMachine.h"
|
2022-11-07 16:01:50 +08:00
|
|
|
#include <filesystem>
|
2022-08-18 01:42:48 +08:00
|
|
|
|
|
|
|
namespace triton {
|
|
|
|
|
2022-11-20 22:25:27 -08:00
|
|
|
static void initLLVM() {
|
2022-09-26 16:38:06 -07:00
|
|
|
LLVMInitializeNVPTXTargetInfo();
|
|
|
|
LLVMInitializeNVPTXTarget();
|
|
|
|
LLVMInitializeNVPTXTargetMC();
|
|
|
|
LLVMInitializeNVPTXAsmPrinter();
|
|
|
|
}
|
|
|
|
|
2022-11-20 22:25:27 -08:00
|
|
|
static bool findAndReplace(std::string &str, const std::string &begin,
|
|
|
|
const std::string &end, const std::string &target) {
|
|
|
|
size_t startReplace = str.find(begin);
|
|
|
|
if (startReplace == std::string::npos)
|
2022-09-26 16:38:06 -07:00
|
|
|
return false;
|
2022-11-20 22:25:27 -08:00
|
|
|
size_t endReplace = str.find(end, startReplace);
|
|
|
|
if (endReplace == std::string::npos)
|
2022-09-26 16:38:06 -07:00
|
|
|
return false;
|
2022-11-20 22:25:27 -08:00
|
|
|
str.replace(startReplace, endReplace + 1 - startReplace, target);
|
2022-09-26 16:38:06 -07:00
|
|
|
return true;
|
|
|
|
}
|
|
|
|
|
2022-11-20 22:25:27 -08:00
|
|
|
static void linkExternal(llvm::Module &module) {
|
2022-11-07 16:01:50 +08:00
|
|
|
bool hasExternal = false;
|
2022-11-20 22:25:27 -08:00
|
|
|
for (auto &func : module) {
|
2022-11-07 16:01:50 +08:00
|
|
|
if (func.hasExternalLinkage()) {
|
|
|
|
hasExternal = true;
|
|
|
|
break;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if (hasExternal) {
|
|
|
|
namespace fs = std::filesystem;
|
|
|
|
// [triton root dir]/python/triton/language/libdevice.10.bc
|
|
|
|
static const fs::path libdevice = fs::path(__FILE__)
|
|
|
|
.parent_path()
|
|
|
|
.parent_path()
|
|
|
|
.parent_path()
|
|
|
|
.parent_path() /
|
|
|
|
"python" / "triton" / "language" /
|
|
|
|
"libdevice.10.bc";
|
2022-11-20 22:25:27 -08:00
|
|
|
if (mlir::triton::linkExternLib(module, libdevice.string()))
|
2022-11-07 16:01:50 +08:00
|
|
|
llvm::errs() << "link failed for: " << libdevice.string();
|
|
|
|
|
2022-11-20 22:25:27 -08:00
|
|
|
// please check https://llvm.org/docs/NVPTXUsage.html#reflection-parameters
|
|
|
|
// this will enable fast math path in libdevice
|
|
|
|
// for example, when enable nvvm-reflect-ftz, sqrt.approx.f32 will change to
|
|
|
|
// sqrt.approx.ftz.f32
|
|
|
|
auto &ctx = module.getContext();
|
2022-11-07 16:01:50 +08:00
|
|
|
llvm::Type *I32 = llvm::Type::getInt32Ty(ctx);
|
|
|
|
llvm::Metadata *mdFour =
|
|
|
|
llvm::ConstantAsMetadata::get(llvm::ConstantInt::getSigned(I32, 4));
|
|
|
|
llvm::Metadata *mdName = llvm::MDString::get(ctx, "nvvm-reflect-ftz");
|
|
|
|
llvm::Metadata *mdOne =
|
|
|
|
llvm::ConstantAsMetadata::get(llvm::ConstantInt::getSigned(I32, 1));
|
|
|
|
llvm::MDNode *reflect = llvm::MDNode::get(ctx, {mdFour, mdName, mdOne});
|
2022-11-20 22:25:27 -08:00
|
|
|
module.addModuleFlag(reflect);
|
2022-11-07 16:01:50 +08:00
|
|
|
}
|
2022-11-20 22:25:27 -08:00
|
|
|
}
|
|
|
|
|
|
|
|
std::string translateLLVMIRToPTX(llvm::Module &module, int cc, int version) {
|
|
|
|
linkExternal(module);
|
|
|
|
|
2022-09-26 16:38:06 -07:00
|
|
|
// LLVM version in use may not officially support target hardware
|
2022-11-20 22:25:27 -08:00
|
|
|
int maxNNVMCC = 75;
|
2022-09-26 16:38:06 -07:00
|
|
|
// options
|
|
|
|
auto options = llvm::cl::getRegisteredOptions();
|
2022-11-20 22:25:27 -08:00
|
|
|
auto *shortPtr =
|
2022-09-26 16:38:06 -07:00
|
|
|
static_cast<llvm::cl::opt<bool> *>(options["nvptx-short-ptr"]);
|
2022-11-20 22:25:27 -08:00
|
|
|
assert(shortPtr);
|
|
|
|
shortPtr->setValue(true);
|
2022-09-26 16:38:06 -07:00
|
|
|
// compute capability
|
2022-11-20 22:25:27 -08:00
|
|
|
std::string sm = "sm_" + std::to_string(cc);
|
2022-09-26 16:38:06 -07:00
|
|
|
// max PTX version
|
2022-11-20 22:25:27 -08:00
|
|
|
int ptxMajor = version / 10;
|
|
|
|
int ptxMinor = version % 10;
|
2022-09-26 16:38:06 -07:00
|
|
|
// create
|
|
|
|
llvm::SmallVector<char, 0> buffer;
|
|
|
|
std::string triple = "nvptx64-nvidia-cuda";
|
2022-11-20 22:25:27 -08:00
|
|
|
std::string proc = "sm_" + std::to_string(std::min(cc, maxNNVMCC));
|
2022-09-26 16:38:06 -07:00
|
|
|
std::string layout = "";
|
|
|
|
std::string features = "";
|
|
|
|
// std::string features = "+ptx" + std::to_string(std::min(ptx,
|
|
|
|
// max_nvvm_ptx));
|
2022-11-20 22:25:27 -08:00
|
|
|
initLLVM();
|
2022-09-26 16:38:06 -07:00
|
|
|
// verify and store llvm
|
|
|
|
llvm::legacy::PassManager pm;
|
|
|
|
pm.add(llvm::createVerifierPass());
|
2022-11-20 22:25:27 -08:00
|
|
|
pm.run(module);
|
2022-09-26 16:38:06 -07:00
|
|
|
// module->print(llvm::outs(), nullptr);
|
|
|
|
|
|
|
|
// create machine
|
2022-11-20 22:25:27 -08:00
|
|
|
module.setTargetTriple(triple);
|
2022-09-26 16:38:06 -07:00
|
|
|
std::string error;
|
|
|
|
auto target =
|
2022-11-20 22:25:27 -08:00
|
|
|
llvm::TargetRegistry::lookupTarget(module.getTargetTriple(), error);
|
2022-09-26 16:38:06 -07:00
|
|
|
llvm::TargetOptions opt;
|
|
|
|
opt.AllowFPOpFusion = llvm::FPOpFusion::Fast;
|
|
|
|
opt.UnsafeFPMath = false;
|
|
|
|
opt.NoInfsFPMath = false;
|
|
|
|
opt.NoNaNsFPMath = true;
|
|
|
|
llvm::TargetMachine *machine = target->createTargetMachine(
|
2022-11-20 22:25:27 -08:00
|
|
|
module.getTargetTriple(), proc, features, opt, llvm::Reloc::PIC_,
|
2022-09-26 16:38:06 -07:00
|
|
|
llvm::None, llvm::CodeGenOpt::Aggressive);
|
|
|
|
// set data layout
|
|
|
|
if (layout.empty())
|
2022-11-20 22:25:27 -08:00
|
|
|
module.setDataLayout(machine->createDataLayout());
|
2022-09-26 16:38:06 -07:00
|
|
|
else
|
2022-11-20 22:25:27 -08:00
|
|
|
module.setDataLayout(layout);
|
2022-09-26 16:38:06 -07:00
|
|
|
// emit machine code
|
2022-11-20 22:25:27 -08:00
|
|
|
for (llvm::Function &f : module.functions())
|
2022-09-26 16:38:06 -07:00
|
|
|
f.addFnAttr(llvm::Attribute::AlwaysInline);
|
|
|
|
llvm::legacy::PassManager pass;
|
|
|
|
llvm::raw_svector_ostream stream(buffer);
|
|
|
|
// emit
|
|
|
|
machine->addPassesToEmitFile(pass, stream, nullptr,
|
|
|
|
llvm::CodeGenFileType::CGFT_AssemblyFile);
|
2022-11-20 22:25:27 -08:00
|
|
|
pass.run(module);
|
2022-09-26 16:38:06 -07:00
|
|
|
|
|
|
|
// post-process
|
|
|
|
std::string result(buffer.begin(), buffer.end());
|
2022-11-20 22:25:27 -08:00
|
|
|
findAndReplace(result, ".version", "\n",
|
|
|
|
".version " + std::to_string(ptxMajor) + "." +
|
|
|
|
std::to_string(ptxMinor) + "\n");
|
|
|
|
findAndReplace(result, ".target", "\n", ".target " + sm + "\n");
|
|
|
|
while (findAndReplace(result, "\t// begin inline asm", "\n", ""))
|
2022-09-26 16:38:06 -07:00
|
|
|
;
|
2022-11-20 22:25:27 -08:00
|
|
|
while (findAndReplace(result, "\t// end inline asm", "\n", ""))
|
2022-09-26 16:38:06 -07:00
|
|
|
;
|
|
|
|
return result;
|
2022-08-18 01:42:48 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
} // namespace triton
|