64 lines
2.2 KiB
C++
64 lines
2.2 KiB
C++
|
#include "triton/codegen/extern_lib.h"
|
||
|
|
|
||
|
|
#include "llvm/IR/Constants.h"
|
||
|
|
#include "llvm/IR/LegacyPassManager.h"
|
||
|
|
#include "llvm/IR/Metadata.h"
|
||
|
|
#include "llvm/IR/Type.h"
|
||
|
|
#include "llvm/Linker/Linker.h"
|
||
|
|
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
|
||
|
|
#include "triton/codegen/pass.h"
|
||
|
|
|
||
|
|
namespace triton {
|
||
|
|
|
||
|
|
namespace codegen {
|
||
|
|
|
||
|
|
std::unique_ptr<llvm::Module> ExternLib::load(llvm::LLVMContext& ctx) {
|
||
|
|
llvm::SMDiagnostic err;
|
||
|
|
auto mod = llvm::parseIRFile(this->path_, err, ctx);
|
||
|
|
if (!mod) {
|
||
|
|
throw std::runtime_error("Failed to load extern lib " + this->name_ +
|
||
|
|
" at " + this->path_);
|
||
|
|
}
|
||
|
|
return mod;
|
||
|
|
}
|
||
|
|
|
||
|
|
void ExternLib::link(std::unique_ptr<llvm::Module>& llvm,
|
||
|
|
std::unique_ptr<llvm::Module>& mod) {
|
||
|
|
// Set triple and data layout to match the target module
|
||
|
|
mod->setTargetTriple(llvm->getTargetTriple());
|
||
|
|
mod->setDataLayout(llvm->getDataLayout());
|
||
|
|
if (llvm::Linker::linkModules(*llvm, std::move(mod))) {
|
||
|
|
throw std::runtime_error("Failed to link extern lib " + this->name_ +
|
||
|
|
" at " + this->path_);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void LibDevice::opt(llvm::LLVMContext& ctx, std::unique_ptr<llvm::Module>& llvm) {
|
||
|
|
// Add nvvm reflect flags to llvm module
|
||
|
|
// https://llvm.org/docs/LangRef.html#module-flags-metadata
|
||
|
|
// i32 4: Override the other module.
|
||
|
|
// i32 1: Emit an error
|
||
|
|
// If both modules specify Override, but the values differ, an error
|
||
|
|
// will be emitted.
|
||
|
|
llvm::Type* I32 = llvm::Type::getInt32Ty(ctx);
|
||
|
|
llvm::Metadata* md_four =
|
||
|
|
llvm::ConstantAsMetadata::get(llvm::ConstantInt::getSigned(I32, 4));
|
||
|
|
llvm::Metadata* md_name = llvm::MDString::get(ctx, "nvvm-reflect-ftz");
|
||
|
|
llvm::Metadata* md_one =
|
||
|
|
llvm::ConstantAsMetadata::get(llvm::ConstantInt::getSigned(I32, 1));
|
||
|
|
llvm::MDNode* reflect = llvm::MDNode::get(ctx, {md_four, md_name, md_one});
|
||
|
|
llvm->addModuleFlag(reflect);
|
||
|
|
}
|
||
|
|
|
||
|
|
std::unique_ptr<ExternLib> create_extern_lib(const std::string& lib_name,
|
||
|
|
const std::string& lib_path) {
|
||
|
|
if (lib_name == "libdevice") {
|
||
|
|
return std::make_unique<LibDevice>(lib_name, lib_path);
|
||
|
|
} else {
|
||
|
|
throw std::runtime_error("Unknown external library: " + lib_name);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace codegen
|
||
|
|
} // namespace triton
|