Google Git Google Git
Sign in
opensecura/3p/openxla/iree/233bd5db3e05e2f57386ff4883a3189ac748e319/./compiler/plugins/target/ROCM/ROCMTarget.cpp
blob: 254f3a2cb8b3e3cd7348f7f38b20739df6975b7d [file]
// Copyright 2021 The IREE Authors
//
// Licensed 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
#include "ROCMTargetUtils.h"
#include <cstddef>
#include <cstdint>
#include "compiler/plugins/target/ROCM/Dialect/ROCM/IR/ROCMAttrs.h"
#include "compiler/plugins/target/ROCM/Dialect/ROCM/Transforms/Passes.h"
#include "compiler/plugins/target/ROCM/builtins/ukernel/iree_uk_amdgpu_bitcode.h"
#include "iree/compiler/Codegen/Common/Passes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
#include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
#include "iree/compiler/Codegen/Utils/CodegenOptions.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
#include "iree/compiler/Dialect/HAL/Utils/ExecutableDebugInfoUtils.h"
#include "iree/compiler/Dialect/HAL/Utils/LLVMLinkerUtils.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "iree/compiler/PluginAPI/Client.h"
#include "iree/compiler/Utils/EmbeddedDataDirectory.h"
#include "iree/compiler/Utils/FlatbufferUtils.h"
#include "iree/compiler/Utils/ToolUtils.h"
#include "iree/schemas/amdgpu_executable_def_builder.h"
#include "iree/schemas/hip_executable_def_builder.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/BinaryFormat/ELF.h"
#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/Frontend/Offloading/Utility.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Verifier.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/StandardInstrumentations.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Endian.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/TargetParser/AMDGPUTargetParser.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Export.h"
namespace mlir::iree_compiler::IREE::HAL {
namespace {
enum class ContainerType {
// Automatically detect the container type from the target ABI attribute.
Auto,
// HIP ExecutableDef flatbuffer.
HIP,
// AMDGPU ExecutableDef flatbuffer.
AMDGPU,
// Raw HSACO image (ELF).
HSACO,
};
enum class AMDGPUTargetFeatureMode {
// Feature mode is not specified by the target ID.
Any,
// Feature is explicitly disabled.
Off,
// Feature is explicitly enabled.
On,
};
struct AMDGPUTargetFeatureModes {
// SRAM ECC target feature mode.
AMDGPUTargetFeatureMode sramecc = AMDGPUTargetFeatureMode::Any;
// XNACK target feature mode.
AMDGPUTargetFeatureMode xnack = AMDGPUTargetFeatureMode::Any;
};
static LogicalResult
setAMDGPUTargetFeatureMode(Location loc, StringRef featureName,
AMDGPUTargetFeatureMode featureMode,
AMDGPUTargetFeatureMode &targetFeatureMode) {
if (targetFeatureMode != AMDGPUTargetFeatureMode::Any) {
return emitError(loc, "duplicate ROCM target feature '")
<< featureName << "'";
}
targetFeatureMode = featureMode;
return success();
}
static FailureOr<AMDGPUTargetFeatureModes>
parseAMDGPUTargetFeatureModes(Location loc, StringRef targetFeatures) {
AMDGPUTargetFeatureModes modes;
SmallVector<StringRef> features;
llvm::SplitString(targetFeatures, features, ",");
for (StringRef rawFeature : features) {
AMDGPUTargetFeatureMode featureMode = AMDGPUTargetFeatureMode::Any;
StringRef feature = rawFeature;
if (feature.consume_front("+")) {
featureMode = AMDGPUTargetFeatureMode::On;
} else if (feature.consume_front("-")) {
featureMode = AMDGPUTargetFeatureMode::Off;
} else {
emitError(loc, "ROCM target feature must be prefixed with '+' or '-'; "
"but seen '")
<< rawFeature << "'";
return failure();
}
if (feature == "sramecc") {
if (failed(setAMDGPUTargetFeatureMode(loc, feature, featureMode,
modes.sramecc))) {
return failure();
}
} else if (feature == "xnack") {
if (failed(setAMDGPUTargetFeatureMode(loc, feature, featureMode,
modes.xnack))) {
return failure();
}
} else {
// We only support these two features to be set explicitly. Features like
// wavefrontsize are controlled and tuned by the compiler.
emitError(loc,
"ROCM target feature can only be 'sramecc' or 'xnack'; but "
"seen '")
<< feature << "'";
return failure();
}
}
return modes;
}
static void appendAMDGPUTargetFeatureSuffix(std::string &targetID,
StringRef featureName,
AMDGPUTargetFeatureMode mode) {
switch (mode) {
case AMDGPUTargetFeatureMode::Any:
return;
case AMDGPUTargetFeatureMode::Off:
targetID += ":";
targetID += featureName;
targetID += "-";
return;
case AMDGPUTargetFeatureMode::On:
targetID += ":";
targetID += featureName;
targetID += "+";
return;
}
}
static llvm::Triple getAMDGPUTargetTriple(StringRef targetArch) {
llvm::Triple triple("amdgcn-amd-amdhsa");
triple.setArch(llvm::Triple::amdgpu,
llvm::AMDGPU::getSubArchFromGPUName(targetArch));
return triple;
}
static FailureOr<llvm::AMDGPU::TargetID>
parseAMDGPUTargetID(Location loc, StringRef targetArch,
StringRef targetFeatures) {
FailureOr<AMDGPUTargetFeatureModes> modes =
parseAMDGPUTargetFeatureModes(loc, targetFeatures);
if (failed(modes)) {
return failure();
}
std::string targetID = targetArch.str();
appendAMDGPUTargetFeatureSuffix(targetID, "sramecc", modes->sramecc);
appendAMDGPUTargetFeatureSuffix(targetID, "xnack", modes->xnack);
auto parsedTargetID = llvm::AMDGPU::TargetID::parse(
getAMDGPUTargetTriple(targetArch), targetID);
if (!parsedTargetID) {
return emitError(loc) << "invalid ROCM target ID '" << targetID << "'";
}
return *parsedTargetID;
}
struct ROCMOptions {
std::string target = "";
std::string targetFeatures = "";
ContainerType containerType = ContainerType::Auto;
std::string bitcodeDirectory = getDefaultBitcodeDirectory();
int wavesPerEu = 0;
std::string enableROCMUkernels = "none";
std::string encodingLayoutResolver =
GPU::kDataTilingEncodingLayoutResolverName;
bool slpVectorization = true;
bool globalISel = false;
bool specializeDispatches = false;
bool enableTensorUKernels = false;
IREE::Codegen::DenormalFpMath denormalFpMathF32 =
IREE::Codegen::DenormalFpMath::None;
bool enableRegSpillWarning = false;
bool debugSymbols = false;
bool useAmdgcnSpirv = false;
void bindOptions(OptionsBinder &binder) {
using namespace llvm;
static cl::OptionCategory category("ROCM HAL Target");
binder.opt<std::string>(
"iree-rocm-target", target, cl::cat(category),
cl::desc(
// clang-format off
"ROCM target as expected by LLVM AMDGPU backend; e.g., "
"'gfx90a'/'gfx942' for targeting MI250/MI300 GPUs. "
"Additionally this also supports architecture code names like "
"'cdna3'/'rdna3' or some product names like 'mi300x'/'rtx7900xtx' "
"for a better experience. See "
"https://iree.dev/guides/deployment-configurations/gpu-rocm/ "
"for more details."
// clang-format on
));
binder.opt<std::string>(
"iree-rocm-target-features", targetFeatures, cl::cat(category),
cl::desc("ROCM target features as expected by LLVM AMDGPU backend; "
"e.g., '+sramecc,+xnack'."));
binder.opt<ContainerType>(
"iree-rocm-container-type", containerType,
llvm::cl::desc("Serialized executable container type."),
llvm::cl::cat(category),
llvm::cl::values(clEnumValN(ContainerType::Auto, "auto",
"Automatically detect the container type "
"from the target ABI attribute."),
clEnumValN(ContainerType::HIP, "hip",
"HIP ExecutableDef flatbuffer."),
clEnumValN(ContainerType::AMDGPU, "amdgpu",
"AMDGPU ExecutableDef flatbuffer."),
clEnumValN(ContainerType::HSACO, "hsaco",
"Raw HSACO image (ELF).")));
binder.opt<std::string>("iree-rocm-bc-dir", bitcodeDirectory,
cl::cat(category),
cl::desc("Directory of ROCM Bitcode."));
binder.opt<int>("iree-rocm-waves-per-eu", wavesPerEu, cl::cat(category),
cl::desc("Optimization hint specifying minimum "
"number of waves per execution unit."));
binder.opt<std::string>(
"iree-rocm-enable-ukernels", enableROCMUkernels, cl::cat(category),
cl::desc("Enables microkernels in the ROCM compiler backend. May be "
"`default`, `none`, `all`, or a comma-separated list of "
"specific unprefixed microkernels to enable, e.g. `mmt4d`."));
binder.opt<std::string>(
"iree-rocm-encoding-layout-resolver", encodingLayoutResolver,
cl::cat(category),
cl::desc("Selects the way that encodings will be "
"resolved. Options are: `none` (resolve to "
"identity layout), `pad` (additional padding "
"on allocations to maximize cache bandwidth), "
"and `data-tiling` (enable data tiled layouts)"));
binder.opt<bool>("iree-rocm-llvm-slp-vec", slpVectorization,
cl::cat(category),
cl::desc("Enable slp vectorization in llvm opt."));
binder.opt<bool>("iree-rocm-llvm-global-isel", globalISel,
cl::cat(category),
cl::desc("Enable global instruction selection in llvm."));
binder.opt<bool>(
"iree-rocm-specialize-dispatches", specializeDispatches,
cl::cat(category),
cl::desc(
"Enable runtime specialization of dynamically shaped dispatches."));
binder.opt<bool>("iree-rocm-enable-tensor-ukernels", enableTensorUKernels,
cl::cat(category),
cl::desc("Enable MLIR-based ukernels."));
binder.opt<IREE::Codegen::DenormalFpMath>(
"iree-rocm-denormal-fp-math-f32", denormalFpMathF32, cl::cat(category),
cl::desc("Denormal floating point math mode for f32"),
cl::values(
clEnumValN(IREE::Codegen::DenormalFpMath::PreserveSign,
"preserve-sign",
"Convert denormals to zero while preserving sign"),
clEnumValN(IREE::Codegen::DenormalFpMath::PositiveZero,
"positive-zero", "Convert denormals to positive zero")));
binder.opt<bool>("iree-rocm-enable-register-spill-warning",
enableRegSpillWarning, cl::cat(category),
cl::desc("Report register spilling for AMD GPUs"));
binder.opt<bool>("iree-rocm-emit-debug-info", debugSymbols,
cl::cat(category),
cl::desc("Generate and embed debug information (DWARF)."));
binder.opt<bool>(
"iree-rocm-use-spirv", useAmdgcnSpirv, cl::cat(category),
cl::desc("Produce SPIR-V binary (amdgcnspirv) instead of native ISA. "
"The HIP runtime JIT-compiles the SPIR-V to native ISA."));
}
LogicalResult verify(mlir::Builder &builder) const {
if (target.empty()) {
return emitError(builder.getUnknownLoc())
<< "ROCM target not set; did you forget to pass "
"'--iree-rocm-target'?";
}
if (GPU::normalizeHIPTarget(target).empty()) {
return emitError(builder.getUnknownLoc(), "Unknown ROCM target '")
<< target << "'";
}
if (failed(parseAMDGPUTargetID(builder.getUnknownLoc(),
GPU::normalizeHIPTarget(target),
targetFeatures))) {
return failure();
}
return success();
}
private:
static std::string getDefaultBitcodeDirectory() {
return mlir::iree_compiler::findPlatformLibDirectory("rocm");
}
};
// Returns the ABI or an empty string if unspecified.
static StringRef getABI(IREE::HAL::ExecutableTargetAttr targetAttr) {
if (targetAttr) {
if (auto config = targetAttr.getConfiguration()) {
auto abiAttr = targetAttr.getConfiguration().getAs<StringAttr>("abi");
return abiAttr ? abiAttr.getValue() : "";
}
}
return "";
}
static void dumpModuleToPath(StringRef path, StringRef baseName,
StringRef suffix, StringRef extension,
llvm::Module &module, StringRef header = {}) {
llvm::SmallVector<char, 0> data;
llvm::raw_svector_ostream ostream(data);
ostream << header;
module.print(ostream, nullptr);
dumpDataToPath(path, baseName, suffix, extension,
StringRef(data.data(), data.size()));
}
static std::string translateModuleToObj(llvm::Module &module,
llvm::TargetMachine &targetMachine) {
std::string targetObj;
{
llvm::raw_string_ostream stream(targetObj);
llvm::buffer_ostream pstream(stream);
llvm::legacy::PassManager codegenPasses;
targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
llvm::CodeGenFileType::ObjectFile);
codegenPasses.run(module);
}
return targetObj;
}
static std::string translateModuleToISA(llvm::Module &module,
llvm::TargetMachine &targetMachine) {
std::string targetISA;
{
llvm::raw_string_ostream stream(targetISA);
llvm::buffer_ostream pstream(stream);
llvm::legacy::PassManager codegenPasses;
targetMachine.addPassesToEmitFile(codegenPasses, pstream, nullptr,
llvm::CodeGenFileType::AssemblyFile);
codegenPasses.run(module);
}
return targetISA;
}
static void checkRegisterSpilling(IREE::HAL::ExecutableVariantOp &variantOp,
StringRef obj) {
uint16_t abiVersion;
llvm::StringMap<llvm::offloading::amdgpu::AMDGPUKernelMetaData> infoMap;
if (!llvm::offloading::amdgpu::getAMDGPUMetaDataFromImage(
llvm::MemoryBufferRef(obj, ""), infoMap, abiVersion)) {
for (const auto &[dispatchName, metaData] : infoMap) {
if (metaData.SGPRSpillCount > 0 || metaData.VGPRSpillCount > 0) {
emitWarning(variantOp.getLoc())
<< "Register spill: " << "VGPRSpillCount: "
<< metaData.VGPRSpillCount
<< " / SGPRSpillCount: " << metaData.SGPRSpillCount
<< " / Dispatch: " << dispatchName;
}
}
}
}
class ROCMTargetBackend final : public TargetBackend {
public:
ROCMTargetBackend(const ROCMOptions &options, GPUCodegenOptions codegenOpts)
: targetOptions(options), codegenOptions(std::move(codegenOpts)) {}
std::string getLegacyDefaultDeviceID() const final { return "hip"; }
void getDefaultExecutableTargets(
MLIRContext *context, StringRef deviceID, DictionaryAttr deviceConfigAttr,
SmallVectorImpl<IREE::HAL::ExecutableTargetAttr> &executableTargetAttrs)
const final {
if (auto target = getExecutableTarget(deviceID, context)) {
executableTargetAttrs.push_back(target);
}
}
IREE::HAL::ExecutableTargetAttr
getExecutableTarget(StringRef deviceID, MLIRContext *context) const {
Builder b(context);
SmallVector<NamedAttribute, 4> configItems;
auto addConfig = [&](StringRef name, Attribute value) {
configItems.emplace_back(name, value);
};
if (failed(targetOptions.verify(b))) {
return nullptr;
}
addConfig("abi", b.getStringAttr(deviceID));
std::string format;
if (targetOptions.useAmdgcnSpirv) {
format = "rocm-spirv-fb";
} else if (deviceID == "amdgpu") {
auto targetID = parseAMDGPUTargetID(
b.getUnknownLoc(), GPU::normalizeHIPTarget(targetOptions.target),
targetOptions.targetFeatures);
if (failed(targetID)) {
return nullptr;
}
format = targetID->getCanonicalTargetIDString();
} else {
format = "rocm-hsaco-fb"; // legacy HIP
}
if (auto target = GPU::getHIPTargetDetails(
targetOptions.target, targetOptions.targetFeatures, context)) {
addConfigGPUTarget(context, target, configItems);
if (targetOptions.encodingLayoutResolver !=
GPU::kNoEncodingLayoutResolverName) {
if (Attribute encoding = GPU::getHIPTargetEncodingLayoutAttr(
target, targetOptions.encodingLayoutResolver)) {
addConfig(IREE::Encoding::kEncodingResolverAttrName, encoding);
}
}
// Look for a default tuning spec.
auto rocmDialect = context->getOrLoadDialect<IREE::ROCM::ROCMDialect>();
// First check for a spec based on the sku.
std::optional<std::string> maybeSpecName = std::nullopt;
if (IREE::GPU::TargetChipAttr chip = target.getChip()) {
if (StringAttr sku = chip.getSku()) {
std::string specName =
llvm::formatv("iree_default_tuning_spec_{}.mlir", sku.strref());
if (!rocmDialect->hasBuiltin(specName)) {
maybeSpecName = specName;
}
}
}
// Then, if none was found, look for one based on the target arch.
if (!maybeSpecName) {
std::string specName =
llvm::formatv("iree_default_tuning_spec_{}.mlir", target.getArch());
if (rocmDialect->hasBuiltin(specName)) {
maybeSpecName = specName;
}
}
if (maybeSpecName) {
addConfig("iree_codegen.default_tuning_spec",
IREE::ROCM::BuiltinTuningModuleAttr::get(
context, maybeSpecName.value()));
}
}
addConfig("ukernels", b.getStringAttr(targetOptions.enableROCMUkernels));
if (targetOptions.enableROCMUkernels != "none") {
addConfig("iree_codegen.ukernel_provider",
IREE::ROCM::UKernelProviderAttr::get(context));
}
if (targetOptions.wavesPerEu > 0) {
addConfigWavesPerEu(b.getContext(), targetOptions.wavesPerEu,
configItems);
}
if (targetOptions.denormalFpMathF32 !=
IREE::Codegen::DenormalFpMath::None) {
addConfigDenormalFpMathF32(b.getContext(),
targetOptions.denormalFpMathF32, configItems);
}
if (targetOptions.enableTensorUKernels) {
addConfig(kUKernelProviderName,
IREE::ROCM::TensorUKernelProviderAttr::get(context));
}
return b.getAttr<IREE::HAL::ExecutableTargetAttr>(
b.getStringAttr("rocm"), b.getStringAttr(format),
b.getDictionaryAttr(configItems));
}
void getDependentDialects(DialectRegistry &registry) const final {
mlir::registerBuiltinDialectTranslation(registry);
mlir::registerLLVMDialectTranslation(registry);
mlir::registerROCDLDialectTranslation(registry);
registry.insert<IREE::Codegen::IREECodegenDialect>();
registry.insert<IREE::VectorExt::IREEVectorExtDialect>();
registry.insert<IREE::GPU::IREEGPUDialect>();
registry.insert<IREE::ROCM::ROCMDialect>();
registry.insert<IREE::Util::UtilDialect>();
registry.insert<mlir::gpu::GPUDialect>();
// Configuration may load and manipulate transform dialect libraries.
registerTransformDialectTranslationDependentDialects(registry);
}
void
buildConfigurationPassPipeline(IREE::HAL::ExecutableTargetAttr targetAttr,
OpPassManager &passManager) final {
if (targetOptions.specializeDispatches) {
if (auto attr = getGPUTargetAttr(targetAttr.getContext(), targetAttr)) {
ROCM::ApplyBuiltinPDLPatternsPassOptions options;
options.enableSpecialization = true;
if (IREE::GPU::TargetChipAttr chip = attr.getChip()) {
if (StringAttr sku = chip.getSku()) {
options.targets.push_back(sku.str());
}
}
options.targets.push_back(attr.getArch().str());
OpPassManager &modulePassManager = passManager.nest<ModuleOp>();
FunctionLikeNest(modulePassManager).addPass([&]() {
return ROCM::createApplyBuiltinPDLPatternsPass(options);
});
}
}
if (targetOptions.enableTensorUKernels) {
if (auto attr = getGPUTargetAttr(targetAttr.getContext(), targetAttr)) {
ROCM::ApplyBuiltinPDLPatternsPassOptions options;
options.enableTensorUKernels = true;
if (IREE::GPU::TargetChipAttr chip = attr.getChip()) {
if (StringAttr sku = chip.getSku()) {
options.targets.push_back(sku.str());
}
}
options.targets.push_back(attr.getArch().str());
OpPassManager &modulePassManager = passManager.nest<ModuleOp>();
FunctionLikeNest(modulePassManager).addPass([&]() {
return ROCM::createApplyBuiltinPDLPatternsPass(options);
});
}
}
buildCodegenConfigurationPreProcessingPassPipeline(passManager);
OpPassManager &modulePassManager = passManager.nest<ModuleOp>();
buildLLVMGPUCodegenCommonConfigurationPassPipeline(modulePassManager);
if (targetOptions.enableTensorUKernels) {
modulePassManager.addPass(
IREE::ROCM::createApplyBuiltinPDLPatternsDriverPass());
}
{
MaterializeTuningSpecsPassOptions passOptions;
passOptions.tuningSpecPath = codegenOptions.tuningSpecPath;
modulePassManager.addPass(createMaterializeTuningSpecsPass(passOptions));
}
modulePassManager.addPass(createMaterializeUserConfigsPass());
modulePassManager.addPass(createLLVMGPUSelectLoweringStrategyPass(
LLVMGPUSelectLoweringStrategyPassOptions{codegenOptions}));
if (shouldEmitPipelineConstraints()) {
FunctionLikeNest funcPassManager(modulePassManager);
funcPassManager.addPass(createInsertSMTConstraintsPass);
if (shouldVerifyPipelineConstraints()) {
funcPassManager.addPass(createVerifySMTConstraintsPass);
}
}
}
void buildTranslationPassPipeline(IREE::HAL::ExecutableTargetAttr targetAttr,
OpPassManager &passManager) final {
// Derive SPIR-V mode from the target format rather than the CLI flag,
// so that programmatically-constructed target attrs work correctly.
bool useSPIRV = targetAttr.getFormat() == "rocm-spirv-fb";
buildLLVMGPUCodegenPassPipeline(
passManager.nest<ModuleOp>(), /*useROCM=*/true,
targetOptions.debugSymbols, /*includeLLVMLowering=*/true, useSPIRV);
buildCodegenTranslationPostProcessingPassPipeline(passManager);
}
void buildLinkingPassPipeline(OpPassManager &passManager) final {
buildLLVMGPULinkingPassPipeline(passManager, "rocm");
}
// Performs optimizations on |module| (including LTO-style whole-program
// ones). Inspired by code section in
// https://github.com/iree-org/iree/blob/main/compiler/plugins/target/CUDA/CUDATarget.cpp
static void optimizeModule(llvm::Module &module,
llvm::TargetMachine &targetMachine,
bool slpVectorization,
std::string &outPassesString) {
// Workaround for upstream LLVM PR llvm/llvm-project#194924 (commit
// b09174b41e7e, "[AMDGPU] Enable runtime loop unrolling") which set
// UP.PartialThreshold = UP.Threshold / 4 (75 for AMDGPU)
// in AMDGPUTTIImpl::getUnrollingPreferences, regressing partial unrolling
// of small constant-trip-count reduction loops. Restore the prior
// un-overridden LLVM default (150) by bumping the cl::opt occurrence
// count, which makes LoopUnrollPass apply the override (see
// LoopUnrollPass.cpp gatherUnrollingPreferences). Skipped when the user
// already set --unroll-partial-threshold so an explicit override still
// wins. Safe to remove once the upstream regression is fixed.
auto &registeredCLOpts = llvm::cl::getRegisteredOptions();
auto unrollOptIt = registeredCLOpts.find("unroll-partial-threshold");
if (unrollOptIt != registeredCLOpts.end() &&
unrollOptIt->second->getNumOccurrences() == 0) {
unrollOptIt->second->addOccurrence(/*pos=*/0,
/*ArgName=*/"unroll-partial-threshold",
/*Value=*/"150");
}
llvm::LoopAnalysisManager lam;
llvm::FunctionAnalysisManager fam;
llvm::CGSCCAnalysisManager cgam;
llvm::ModuleAnalysisManager mam;
fam.registerPass([&] { return targetMachine.getTargetIRAnalysis(); });
llvm::PipelineTuningOptions pto;
pto.SLPVectorization = slpVectorization;
llvm::PassInstrumentationCallbacks pic;
llvm::StandardInstrumentations si(module.getContext(), false);
si.registerCallbacks(pic, &mam);
llvm::PassBuilder pb(&targetMachine, pto, std::nullopt, &pic);
llvm::ModulePassManager mpm;
pb.registerModuleAnalyses(mam);
pb.registerCGSCCAnalyses(cgam);
pb.registerFunctionAnalyses(fam);
pb.registerLoopAnalyses(lam);
pb.crossRegisterProxies(lam, fam, cgam, mam);
llvm::OptimizationLevel ol = llvm::OptimizationLevel::O2;
mpm.addPass(llvm::VerifierPass());
mpm.addPass(pb.buildPerModuleDefaultPipeline(ol));
mpm.addPass(llvm::VerifierPass());
llvm::raw_string_ostream os(outPassesString);
mpm.printPipeline(os, [&pic](StringRef className) {
auto passName = pic.getPassNameForClassName(className);
return passName.empty() ? className : passName;
});
mpm.run(module, mam);
}
LogicalResult
validateFinalizedModule(IREE::HAL::ExecutableVariantOp variantOp,
llvm::Module &module, bool allowExternalDecls) {
// The LLVM backend reads the target triple from the module itself (not just
// from the TargetMachine) to decide how to emit the code object and HSA
// metadata. If it is left unset, code generation crashes deep inside the
// AMDGPU AsmPrinter with an opaque signal (a null HSAMetadataStream
// dereference; see llvm-project@00a6186128d3). Fail early here with an
// actionable message instead. The triple is set in serializeExecutable
// right after the data layout.
if (module.getTargetTriple().str().empty()) {
return variantOp.emitError()
<< "the finalized LLVM module has no target triple set; it must "
"be set before serialization (see ROCMTargetBackend::"
"serializeExecutable). Without it the LLVM AMDGPU backend "
"crashes while emitting the object file.";
}
for (llvm::Function &func : module.functions()) {
if (func.isDeclaration() && !func.isIntrinsic() && !func.use_empty()) {
// In SPIR-V mode, external declarations (e.g. __ocml_*, __ockl_*,
// amdgcn intrinsics) are expected — they are resolved at JIT time.
if (allowExternalDecls) {
continue;
}
llvm::User *liveUser = *func.user_begin();
return variantOp.emitError()
<< "found an unresolved external function '" << func.getName()
<< "' in the final bitcode. A remaining live user is\n"
<< llvm::formatv("{}", *liveUser);
}
}
return success();
}
LogicalResult
serializeExecutable(const SerializationOptions &serializationOptions,
IREE::HAL::ExecutableVariantOp variantOp,
OpBuilder &executableBuilder) final {
ModuleOp innerModuleOp = variantOp.getInnerModule();
auto targetAttr = variantOp.getTargetAttr();
bool useSPIRV = targetAttr.getFormat() == "rocm-spirv-fb";
if (targetOptions.useAmdgcnSpirv && !useSPIRV) {
return variantOp.emitError()
<< "--iree-rocm-use-spirv requires ROCm executable target format "
"'rocm-spirv-fb', but got '"
<< targetAttr.getFormat() << "'";
}
StringRef targetArch = targetOptions.target;
StringRef targetFeatures = targetOptions.targetFeatures;
uint32_t preferredSubgroupSize = 64;
if (auto attr = getGPUTargetAttr(variantOp.getContext(), targetAttr)) {
targetArch = attr.getArch();
targetFeatures = attr.getFeatures();
preferredSubgroupSize = attr.getPreferredSubgroupSize();
}
auto targetID =
parseAMDGPUTargetID(variantOp.getLoc(), targetArch, targetFeatures);
if (failed(targetID)) {
return failure();
}
// We name our files after the executable name so that they are easy to
// track both during compilation (logs/artifacts/etc), as outputs (final
// intermediate code/binary files), and at runtime (loaded
// libraries/symbols/etc).
const std::string libraryName =
variantOp->getParentOfType<IREE::HAL::ExecutableOp>().getName().str();
// Collect all the entry point names.
auto exportOps = llvm::to_vector_of<IREE::HAL::ExecutableExportOp>(
variantOp.getExportOps());
std::optional<uint32_t> subgroupSize;
for (IREE::HAL::ExecutableExportOp exportOp : exportOps) {
// TODO: put this either on the variant or propagate as a function
// attribute instead - today this *must* be consistent across all exports
// and it shouldn't need to be.
if (auto setSubgroupSize = exportOp.getSubgroupSizeAsUInt()) {
if (setSubgroupSize.value() != 32 && setSubgroupSize.value() != 64) {
return variantOp.emitError()
<< "invalid subgroup size " << setSubgroupSize.value();
}
if (subgroupSize.has_value() &&
setSubgroupSize.value() != subgroupSize.value()) {
return variantOp.emitError()
<< "multiple exports with different subgroup sizes; this is a "
"limitation of the IREE compilation process and should be "
"fixed";
}
subgroupSize = setSubgroupSize.value();
}
}
std::string targetHSACO;
if (variantOp.isExternal()) {
if (!variantOp.getObjects().has_value()) {
return variantOp.emitOpError()
<< "no objects defined for external variant";
}
if (variantOp.getObjects()->getValue().size() != 1) {
// For now we assume there will be exactly one object file.
// In the future we will want to perform a linking step here and ideally
// support _also_ linking in the codegen results.
return variantOp.emitOpError() << "only one object reference is "
"supported for external variants";
}
// Read the HSACO from the object file.
auto objectAttr = cast<IREE::HAL::ExecutableObjectAttr>(
variantOp.getObjects()->getValue().front());
if (auto data = objectAttr.loadData()) {
targetHSACO = data.value();
} else {
return variantOp.emitOpError()
<< "object file could not be loaded: " << objectAttr;
}
} else {
auto maybeChipset = amdgpu::Chipset::parse(targetArch);
if (failed(maybeChipset)) {
return variantOp.emitOpError()
<< "could not parse AMDGPU chipset name '" << targetArch << "'";
}
amdgpu::Chipset chipset = *maybeChipset;
// Perform the translation in a separate context to avoid any
// multi-threading issues.
llvm::LLVMContext context;
std::unique_ptr<llvm::Module> llvmModule =
mlir::translateModuleToLLVMIR(innerModuleOp, context, libraryName);
if (!llvmModule) {
return variantOp.emitError() << "failed to translate the MLIR LLVM "
"dialect to the native llvm::Module";
}
for (auto func : innerModuleOp.getOps<LLVM::LLVMFuncOp>()) {
llvm::Function *llvmFunc = llvmModule->getFunction(func.getName());
if (llvmFunc->isDeclaration()) {
continue;
}
// Override flags as given by target func attrs.
if (auto funcAttrs =
func->getAttrOfType<DictionaryAttr>("llvm_func_attrs")) {
for (NamedAttribute funcAttr : funcAttrs) {
auto value = dyn_cast<StringAttr>(funcAttr.getValue());
if (!value) {
return variantOp->emitError()
<< "llvm_func_attrs attribute must be a dictionary of "
"strings. Attribute "
<< funcAttr.getName() << " is not a StringAttr.";
}
llvmFunc->addFnAttr(funcAttr.getName(), value.getValue());
}
}
}
std::unique_ptr<llvm::TargetMachine> targetMachine;
bool isWave64 = true;
{
llvm::Triple triple = getAMDGPUTargetTriple(targetArch);
std::string error;
const llvm::Target *target =
llvm::TargetRegistry::lookupTarget("", triple, error);
if (!target) {
return variantOp.emitError() << "cannot initialize target triple";
}
llvm::TargetOptions opt;
// Be extra cautious while this is less tested, and prevent unknown
// fallbacks from global isel.
//
// When GlobalISelAbort is set, any failure of GlobalISel,
// whether due to being not yet implemented or incorrect IR will result
// in an immediate abortion of compilation. This disables the fallback
// path of AMDGPUPassConfig::addInstSelector and 2 legacy passes which
// might work around unimplemented cases or errors in GlobalISel
// resulting in a successful compilation but would make one assume
// results are with GlobalISel when they are not.
opt.EnableGlobalISel = targetOptions.globalISel;
opt.GlobalISelAbort = targetOptions.globalISel
? llvm::GlobalISelAbortMode::Enable
: llvm::GlobalISelAbortMode::Disable;
SmallVector<std::string> features;
if (chipset.majorVersion >= 10 && chipset.majorVersion <= 12) {
switch (subgroupSize.value_or(preferredSubgroupSize)) {
case 32:
isWave64 = false;
features.emplace_back("+wavefrontsize32");
break;
default:
case 64:
features.emplace_back("+wavefrontsize64");
break;
}
}
// Mixed precision fma instructions have complicated semantics on
// gf9+ GPUs and can lead to numeric issues as seen in
// https://github.com/iree-org/iree/issues/18746 so we disable this
// feature.
if (targetArch.starts_with("gfx9")) {
features.emplace_back("-fma-mix-insts");
}
std::string featureStr = llvm::join(features, ",");
targetMachine.reset(target->createTargetMachine(
triple, targetArch, featureStr, opt, llvm::Reloc::Model::PIC_,
std::nullopt, llvm::CodeGenOptLevel::Aggressive));
if (!targetMachine) {
return variantOp.emitError() << "cannot initialize target machine";
}
}
// Code object version * 100.
constexpr uint32_t abiVersion = 500;
// For SPIR-V mode, the data layout is already set by PrepareForSPIRVPass
// on the MLIR module and propagated through MLIR->LLVM IR translation.
if (!useSPIRV) {
llvmModule->setDataLayout(targetMachine->createDataLayout());
// The AMDGPU AsmPrinter now reads the target triple from the module
// (not the TargetMachine) to decide whether to set up HSA metadata
// streaming; without this the module triple is empty and the printer
// null-derefs HSAMetadataStream. See llvm-project@00a6186128d3.
llvmModule->setTargetTriple(targetMachine->getTargetTriple());
// Let the backend know what code object version we're compiling for.
// This insulates us from changes to the default code object version
// that our CI or users may not be prepared for.
llvmModule->addModuleFlag(llvm::Module::Error,
"amdhsa_code_object_version", abiVersion);
// XNACK and SRAM ECC describe the whole code object. LLVM requires
// module flags for explicit settings, not subtarget features. Leave
// unspecified/unsupported settings absent to preserve their defaults.
auto addTargetIDFlag = [&](StringRef name,
llvm::AMDGPU::TargetIDSetting setting) {
using llvm::AMDGPU::TargetIDSetting;
if (setting == TargetIDSetting::On ||
setting == TargetIDSetting::Off) {
llvmModule->addModuleFlag(llvm::Module::Error, name,
setting == TargetIDSetting::On);
}
};
addTargetIDFlag("amdgpu.xnack", targetID->getXnackSetting());
addTargetIDFlag("amdgpu.sramecc", targetID->getSramEccSetting());
// Set the buffer OOB mode to "relaxed", since IREE does its own, more
// precise, mitigations for partially OOB buffer reads of underalligned
// vectors.
llvmModule->addModuleFlag(llvm::Module::Max, "amdgpu.buffer.oob.mode",
1);
}
for (llvm::Function &f : llvmModule->functions()) {
f.addFnAttr(llvm::Attribute::AlwaysInline);
}
// Link user-provided modules.
llvm::Linker linker(*llvmModule);
if (failed(linkCmdlineBitcodeFiles(
variantOp.getLoc(), linker, llvm::Linker::OverrideFromSrc,
*targetMachine, llvmModule->getContext()))) {
return failure();
}
// Link bitcode (*.bc) object attrs specified by the input program.
// Note that this happens after the command-line files so that the command
// line ones override the symbols coming from the embedded files.
auto specializationCallback = [&](llvm::Module &userModule) {
// TODO: inject __nvvm_reflect-style functions/globals for
// bitcode specialization based on the targetMachine and configuration.
// These could use any information we have on the IREE side as well as
// the TargetMachine.
};
unsigned linkerFlags =
llvm::Linker::LinkOnlyNeeded | llvm::Linker::OverrideFromSrc;
if (failed(linkBitcodeObjects(variantOp.getLoc(), linker, linkerFlags,
*targetMachine, variantOp.getObjectsAttr(),
llvmModule->getContext(),
specializationCallback))) {
return mlir::emitError(variantOp.getLoc())
<< "failed linking in user objects for target triple '"
<< targetArch.str() << "'";
}
// Link module to HIP device library (skip for SPIR-V — resolved at JIT).
if (!useSPIRV) {
if (targetOptions.bitcodeDirectory.empty()) {
return variantOp.emitError()
<< "cannot find ROCM bitcode files. Check your installation "
"consistency and in the worst case, set "
"--iree-rocm-bc-dir= to a path on your system.";
}
if (failed(linkHIPBitcodeIfNeeded(variantOp.getLoc(), llvmModule.get(),
targetArch,
targetOptions.bitcodeDirectory))) {
return failure();
}
// Sets HIP platform globals based on the target architecture.
if (failed(setHIPGlobals(variantOp.getLoc(), llvmModule.get(), chipset,
isWave64, abiVersion))) {
return failure();
}
}
if (!serializationOptions.dumpIntermediatesPath.empty()) {
dumpModuleToPath(serializationOptions.dumpIntermediatesPath,
serializationOptions.dumpBaseName, variantOp.getName(),
".linked.ll", *llvmModule);
}
// For example 'gfx942'.
StringRef targetCPU = targetMachine->getTargetCPU();
// For example 'amdgpu9.42-amd-amdhsa'.
std::string targetTriple = targetMachine->getTargetTriple().str();
// Run LLVM optimization passes.
// For SPIR-V mode, skip AMDGPU-specific optimization which introduces
// buffer resource intrinsics (AS 7) and other constructs incompatible
// with the SPIR-V backend.
std::string passesString;
if (!useSPIRV) {
optimizeModule(*llvmModule, *targetMachine,
targetOptions.slpVectorization, passesString);
}
if (!serializationOptions.dumpIntermediatesPath.empty()) {
// Additional context on '-mcpu' flag in PR comments, see for example:
// https://github.com/iree-org/iree/pull/20716#issuecomment-2851650421
std::string header = llvm::formatv(
R"TXT(; To reproduce the .optimized.ll from the .linked.ll, run:
; opt -S -mtriple={} -mcpu={} --passes='{}' <.linked.ll>
; The flag '-S' is to emit LLVMIR.
; The behavior of some passes depends on '-mtriple' and '-mcpu'.
)TXT",
targetTriple, targetCPU, passesString);
dumpModuleToPath(serializationOptions.dumpIntermediatesPath,
serializationOptions.dumpBaseName, variantOp.getName(),
".optimized.ll", *llvmModule, header);
}
if (failed(validateFinalizedModule(variantOp, *llvmModule, useSPIRV))) {
return failure();
}
if (useSPIRV) {
// SPIR-V codegen path: create a SPIR-V TargetMachine and emit binary.
llvm::Triple spirvTriple("spirv64-amd-amdhsa");
std::string spirvError;
const llvm::Target *spirvTarget =
llvm::TargetRegistry::lookupTarget("", spirvTriple, spirvError);
if (!spirvTarget) {
return variantOp.emitError()
<< "cannot find SPIR-V target: " << spirvError;
}
auto spirvTM = std::unique_ptr<llvm::TargetMachine>(
spirvTarget->createTargetMachine(
spirvTriple, /*CPU=*/"", /*Features=*/"", llvm::TargetOptions{},
llvm::Reloc::PIC_, std::nullopt,
llvm::CodeGenOptLevel::Default));
if (!spirvTM) {
return variantOp.emitError() << "cannot create SPIR-V target machine";
}
// If requested, dump the final .ll just before conversion to
// SPIRV-binary
if (!serializationOptions.dumpIntermediatesPath.empty()) {
dumpModuleToPath(serializationOptions.dumpIntermediatesPath,
serializationOptions.dumpBaseName,
variantOp.getName(), ".final.ll", *llvmModule);
}
// Emit SPIR-V binary (no lld/hsaco needed).
// The MLIR PrepareForSPIRV pass set triple, data layout, calling
// conventions, stripped assumes, and removed AMDGPU attrs. We
// reassert the data layout here from the TM to guarantee it matches
// exactly (MachineFunction::init() asserts on mismatch).
llvmModule->setDataLayout(spirvTM->createDataLayout());
llvmModule->setTargetTriple(spirvTriple);
std::string spirvBinary = translateModuleToObj(*llvmModule, *spirvTM);
if (spirvBinary.empty()) {
return variantOp.emitError()
<< "SPIR-V binary translation produced empty output";
}
if (!serializationOptions.dumpIntermediatesPath.empty()) {
dumpDataToPath(serializationOptions.dumpIntermediatesPath,
serializationOptions.dumpBaseName, variantOp.getName(),
".spv", spirvBinary);
}
// Wrap the SPIR-V binary in a clang offload bundle. The HIP runtime
// expects this format — raw SPIR-V is not accepted by
// hipModuleLoadData.
// Note: Multi-byte fields use little-endian byte order, matching the
// Clang offload bundler binary format specification.
const std::string bundleTarget = "hip-spirv64-amd-amdhsa--amdgcnspirv";
const char bundleMagic[] = "__CLANG_OFFLOAD_BUNDLE__";
const uint64_t numEntries = 1;
const uint64_t headerSize = sizeof(bundleMagic) - 1 +
sizeof(numEntries) + 3 * sizeof(uint64_t) +
bundleTarget.size();
const uint64_t dataOffset = headerSize;
const uint64_t dataSize = spirvBinary.size();
std::string bundle;
bundle.reserve(headerSize + dataSize);
// Magic string (no null terminator).
bundle.append(bundleMagic, sizeof(bundleMagic) - 1);
// Number of entries.
bundle.append(reinterpret_cast<const char *>(&numEntries),
sizeof(numEntries));
// Entry: offset, size, target string length, target string.
bundle.append(reinterpret_cast<const char *>(&dataOffset),
sizeof(dataOffset));
bundle.append(reinterpret_cast<const char *>(&dataSize),
sizeof(dataSize));
uint64_t targetLen = bundleTarget.size();
bundle.append(reinterpret_cast<const char *>(&targetLen),
sizeof(targetLen));
bundle.append(bundleTarget);
// SPIR-V binary data.
bundle.append(spirvBinary);
targetHSACO = std::move(bundle);
} else {
// Native ISA codegen path.
// Dump the assembly output.
if (!serializationOptions.dumpIntermediatesPath.empty()) {
auto moduleCopy = llvm::CloneModule(*llvmModule);
if (!moduleCopy) {
llvm::errs() << "Error: cloning LLVM IR failed\n";
return failure();
}
std::string asmHeader = llvm::formatv(
R"TXT(; To reproduce the .rocmasm from .optimized.ll, run:
; llc -mtriple={} -mcpu={} -mattr='{}' -O3 <.optimized.ll> -o <out.rocmasm>
)TXT",
targetTriple, targetCPU, targetMachine->getTargetFeatureString());
std::string targetISA =
translateModuleToISA(*moduleCopy.get(), *targetMachine);
dumpDataToPath(serializationOptions.dumpIntermediatesPath,
serializationOptions.dumpBaseName, variantOp.getName(),
".rocmasm", asmHeader + targetISA);
}
// Serialize hsaco kernel into the binary that we will embed in the
// final FlatBuffer.
std::string targetObj =
translateModuleToObj(*llvmModule, *targetMachine);
targetHSACO = createHsaco(variantOp.getLoc(), targetObj, libraryName);
if (targetHSACO.empty()) {
return failure();
}
// LLVM a86b58566a25 started encoding gfx1250's hardwired XNACK as
// XNACK_ON in ELF e_flags. ROCm then looks for a gfx1250:xnack+
// target, which the runtime does not register. Keep the pre-existing
// encoding for this non-selectable feature until LLVM's ELF emission
// is corrected. This does not change the kernel's XNACK codegen.
if (targetArch == "gfx1250") {
char *flags =
targetHSACO.data() + offsetof(llvm::ELF::Elf64_Ehdr, e_flags);
llvm::support::endian::write32le(
flags, llvm::support::endian::read32le(flags) &
~llvm::ELF::EF_AMDGPU_FEATURE_XNACK_V4);
}
if (targetOptions.enableRegSpillWarning) {
checkRegisterSpilling(variantOp, targetObj);
}
}
}
if (!serializationOptions.dumpBinariesPath.empty()) {
StringRef ext = useSPIRV ? ".spv" : ".hsaco";
dumpDataToPath(serializationOptions.dumpBinariesPath,
serializationOptions.dumpBaseName, variantOp.getName(),
ext, targetHSACO);
}
// Determine container type from the target ABI attribute.
ContainerType containerType = targetOptions.containerType;
if (containerType == ContainerType::Auto) {
if (getABI(targetAttr) == "amdgpu") {
containerType = ContainerType::AMDGPU;
} else {
containerType = ContainerType::HIP;
}
}
// Wrap the HSACO ELF binary in the requested container type (if any).
StringAttr executableBinaryFormat = variantOp.getTarget().getFormat();
FailureOr<DenseIntElementsAttr> binaryContainer;
switch (containerType) {
case ContainerType::Auto: {
// Resolved above; unreachable. Fall-through to error case.
assert(false && "auto container type must have resolved earlier");
break;
}
case ContainerType::AMDGPU: {
std::string canonicalTargetID = targetID->getCanonicalTargetIDString();
executableBinaryFormat =
executableBuilder.getStringAttr(canonicalTargetID);
binaryContainer = serializeAMDGPUBinaryContainer(
serializationOptions, variantOp, exportOps, canonicalTargetID,
targetHSACO);
break;
}
case ContainerType::HIP: {
binaryContainer = serializeHIPBinaryContainer(
serializationOptions, variantOp, exportOps, targetHSACO);
break;
}
case ContainerType::HSACO: {
SmallVector<uint8_t> image;
image.resize(targetHSACO.size());
std::memcpy(image.data(), targetHSACO.data(), image.size());
binaryContainer = DenseIntElementsAttr::get(
VectorType::get({static_cast<int64_t>(targetHSACO.size())},
executableBuilder.getI8Type()),
image);
break;
}
}
if (failed(binaryContainer) || !binaryContainer.value()) {
return failure();
}
// Add the binary data to the target executable.
auto binaryOp = IREE::HAL::ExecutableBinaryOp::create(
executableBuilder, variantOp.getLoc(), variantOp.getSymName(),
executableBinaryFormat, binaryContainer.value());
binaryOp.setMimeTypeAttr(
executableBuilder.getStringAttr("application/x-flatbuffers"));
return success();
}
protected:
FailureOr<DenseIntElementsAttr> serializeAMDGPUBinaryContainer(
const SerializationOptions &serializationOptions,
IREE::HAL::ExecutableVariantOp variantOp,
ArrayRef<IREE::HAL::ExecutableExportOp> exportOps, StringRef targetID,
StringRef hsacoModule) {
iree_compiler::FlatbufferBuilder builder;
iree_hal_amdgpu_ExecutableDef_start_as_root(builder);
// Attach embedded source file contents.
auto sourceFilesRef = createSourceFilesVec(
serializationOptions.debugLevel, variantOp.getSourcesAttr(), builder);
// Only a single module today.
SmallVector<iree_hal_amdgpu_ModuleDef_ref_t> moduleRefs;
{
auto hsacoImageRef = flatbuffers_string_create(
builder, hsacoModule.data(), hsacoModule.size());
moduleRefs.push_back(
iree_hal_amdgpu_ModuleDef_create(builder, hsacoImageRef));
}
auto modulesRef = builder.createOffsetVecDestructive(moduleRefs);
// Generate optional per-export debug information.
// May be empty if no debug information was requested.
auto exportDebugInfos =
createExportDefs(serializationOptions.debugLevel, exportOps, builder);
SmallVector<iree_hal_amdgpu_ExportDef_ref_t> exportRefs;
exportRefs.resize(exportOps.size(), 0);
for (auto exportOp : exportOps) {
auto ordinalAttr = exportOp.getOrdinalAttr();
if (!ordinalAttr) {
return mlir::emitError(exportOp.getLoc())
<< "could not compile rocm binary: export op is missing ordinal";
}
int64_t ordinal = ordinalAttr.getInt();
// Symbol names include a `.kd` suffix as that's what HSA expects.
auto symbolNameKd = (exportOp.getName() + ".kd").str();
auto symbolNameRef = builder.createString(symbolNameKd);
iree_hal_amdgpu_Dims_t workgroupSize = {0};
if (auto workgroupSizeAttr = exportOp.getWorkgroupSize()) {
auto workgroupSizeDims = workgroupSizeAttr->getValue();
workgroupSize.x = cast<IntegerAttr>(workgroupSizeDims[0]).getInt();
workgroupSize.y = cast<IntegerAttr>(workgroupSizeDims[1]).getInt();
workgroupSize.z = cast<IntegerAttr>(workgroupSizeDims[2]).getInt();
}
auto layoutAttr = exportOp.getLayoutAttr();
uint32_t constantCount = static_cast<uint32_t>(layoutAttr.getConstants());
SmallVector<iree_hal_amdgpu_BindingBits_enum_t> bindingFlags;
for (auto bindingAttr : layoutAttr.getBindings()) {
iree_hal_amdgpu_BindingBits_enum_t flags = 0;
if (allEnumBitsSet(bindingAttr.getFlags(),
IREE::HAL::DescriptorFlags::ReadOnly)) {
flags |= iree_hal_amdgpu_BindingBits_READ_ONLY;
}
if (allEnumBitsSet(bindingAttr.getFlags(),
IREE::HAL::DescriptorFlags::Indirect)) {
flags |= iree_hal_amdgpu_BindingBits_INDIRECT;
}
bindingFlags.push_back(flags);
}
auto bindingFlagsRef = iree_hal_amdgpu_BindingBits_vec_create(
builder, bindingFlags.data(), bindingFlags.size());
iree_hal_amdgpu_ExportDef_start(builder);
iree_hal_amdgpu_ExportDef_symbol_name_add(builder, symbolNameRef);
iree_hal_amdgpu_ExportDef_workgroup_size_add(builder, &workgroupSize);
iree_hal_amdgpu_ExportDef_constant_count_add(builder, constantCount);
iree_hal_amdgpu_ExportDef_binding_flags_add(builder, bindingFlagsRef);
iree_hal_amdgpu_ExportDef_debug_info_add(builder,
exportDebugInfos[ordinal]);
exportRefs[ordinal] = iree_hal_amdgpu_ExportDef_end(builder);
}
auto exportsRef = builder.createOffsetVecDestructive(exportRefs);
auto isaRef = builder.createString(targetID);
iree_hal_amdgpu_ExecutableDef_isa_add(builder, isaRef);
iree_hal_amdgpu_ExecutableDef_exports_add(builder, exportsRef);
iree_hal_amdgpu_ExecutableDef_modules_add(builder, modulesRef);
iree_hal_amdgpu_ExecutableDef_source_files_add(builder, sourceFilesRef);
iree_hal_amdgpu_ExecutableDef_end_as_root(builder);
return builder.getHeaderPrefixedBufferAttr(
variantOp.getContext(),
/*magic=*/iree_hal_amdgpu_ExecutableDef_file_identifier,
/*version=*/0);
}
FailureOr<DenseIntElementsAttr>
serializeHIPBinaryContainer(const SerializationOptions &serializationOptions,
IREE::HAL::ExecutableVariantOp variantOp,
ArrayRef<IREE::HAL::ExecutableExportOp> exportOps,
StringRef hsacoModule) {
iree_compiler::FlatbufferBuilder builder;
iree_hal_hip_ExecutableDef_start_as_root(builder);
// Attach embedded source file contents.
auto sourceFilesRef = createSourceFilesVec(
serializationOptions.debugLevel, variantOp.getSourcesAttr(), builder);
// Only a single module today.
SmallVector<iree_hal_hip_ModuleDef_ref_t> moduleRefs;
{
auto hsacoImageRef = flatbuffers_string_create(
builder, hsacoModule.data(), hsacoModule.size());
moduleRefs.push_back(
iree_hal_hip_ModuleDef_create(builder, hsacoImageRef));
}
auto modulesRef = builder.createOffsetVecDestructive(moduleRefs);
// Generate optional per-export debug information.
// May be empty if no debug information was requested.
auto exportDebugInfos =
createExportDefs(serializationOptions.debugLevel, exportOps, builder);
SmallVector<iree_hal_hip_ExportDef_ref_t> exportRefs;
exportRefs.resize(exportOps.size(), 0);
for (auto exportOp : exportOps) {
auto ordinalAttr = exportOp.getOrdinalAttr();
if (!ordinalAttr) {
return mlir::emitError(exportOp.getLoc())
<< "could not compile rocm binary: export op is missing ordinal";
}
int64_t ordinal = ordinalAttr.getInt();
auto kernelNameRef = builder.createString(exportOp.getName());
iree_hal_hip_BlockDims_t blockDims = {0};
if (auto workgroupSizeAttr = exportOp.getWorkgroupSize()) {
auto workgroupSize = workgroupSizeAttr->getValue();
blockDims.x = cast<IntegerAttr>(workgroupSize[0]).getInt();
blockDims.y = cast<IntegerAttr>(workgroupSize[1]).getInt();
blockDims.z = cast<IntegerAttr>(workgroupSize[2]).getInt();
}
auto layoutAttr = exportOp.getLayoutAttr();
uint32_t constantCount = static_cast<uint32_t>(layoutAttr.getConstants());
SmallVector<iree_hal_hip_BindingBits_enum_t> bindingFlags;
for (auto bindingAttr : layoutAttr.getBindings()) {
iree_hal_hip_BindingBits_enum_t flags = 0;
if (allEnumBitsSet(bindingAttr.getFlags(),
IREE::HAL::DescriptorFlags::ReadOnly)) {
flags |= iree_hal_hip_BindingBits_READ_ONLY;
}
if (allEnumBitsSet(bindingAttr.getFlags(),
IREE::HAL::DescriptorFlags::Indirect)) {
flags |= iree_hal_hip_BindingBits_INDIRECT;
}
bindingFlags.push_back(flags);
}
auto bindingFlagsRef = iree_hal_hip_BindingBits_vec_create(
builder, bindingFlags.data(), bindingFlags.size());
iree_hal_hip_ExportDef_start(builder);
iree_hal_hip_ExportDef_module_ordinal_add(builder, 0); // always 0 today
iree_hal_hip_ExportDef_kernel_name_add(builder, kernelNameRef);
iree_hal_hip_ExportDef_block_dims_add(builder, &blockDims);
iree_hal_hip_ExportDef_constant_count_add(builder, constantCount);
iree_hal_hip_ExportDef_binding_flags_add(builder, bindingFlagsRef);
iree_hal_hip_ExportDef_debug_info_add(builder, exportDebugInfos[ordinal]);
exportRefs[ordinal] = iree_hal_hip_ExportDef_end(builder);
}
auto exportsRef = builder.createOffsetVecDestructive(exportRefs);
iree_hal_hip_ExecutableDef_exports_add(builder, exportsRef);
iree_hal_hip_ExecutableDef_modules_add(builder, modulesRef);
iree_hal_hip_ExecutableDef_source_files_add(builder, sourceFilesRef);
iree_hal_hip_ExecutableDef_end_as_root(builder);
return builder.getHeaderPrefixedBufferAttr(
variantOp.getContext(),
/*magic=*/iree_hal_hip_ExecutableDef_file_identifier,
/*version=*/0);
}
private:
const ROCMOptions &targetOptions;
const GPUCodegenOptions codegenOptions;
};
class AMDGPUTargetDevice final : public TargetDevice {
public:
AMDGPUTargetDevice(const ROCMOptions & /*options*/) {}
IREE::HAL::DeviceTargetAttr
getDefaultDeviceTarget(MLIRContext *context,
const TargetRegistry &targetRegistry) const final {
Builder b(context);
auto deviceConfigAttr = b.getDictionaryAttr({});
auto executableConfigAttr = b.getDictionaryAttr({});
// If we had multiple target environments we would generate one target attr
// per environment, with each setting its own environment attribute.
SmallVector<IREE::HAL::ExecutableTargetAttr> executableTargetAttrs;
targetRegistry.getTargetBackend("rocm")->getDefaultExecutableTargets(
context, "amdgpu", executableConfigAttr, executableTargetAttrs);
return IREE::HAL::DeviceTargetAttr::get(context, b.getStringAttr("amdgpu"),
deviceConfigAttr,
executableTargetAttrs);
}
};
class HIPTargetDevice final : public TargetDevice {
public:
HIPTargetDevice(const ROCMOptions & /*options*/) {}
IREE::HAL::DeviceTargetAttr
getDefaultDeviceTarget(MLIRContext *context,
const TargetRegistry &targetRegistry) const final {
Builder b(context);
auto deviceConfigAttr = b.getDictionaryAttr({});
auto executableConfigAttr = b.getDictionaryAttr({});
// If we had multiple target environments we would generate one target attr
// per environment, with each setting its own environment attribute.
SmallVector<IREE::HAL::ExecutableTargetAttr> executableTargetAttrs;
targetRegistry.getTargetBackend("rocm")->getDefaultExecutableTargets(
context, "hip", executableConfigAttr, executableTargetAttrs);
return IREE::HAL::DeviceTargetAttr::get(context, b.getStringAttr("hip"),
deviceConfigAttr,
executableTargetAttrs);
}
};
struct ROCMSession final
: PluginSession<ROCMSession, ROCMOptions,
PluginActivationPolicy::DefaultActivated> {
static void registerPasses() { IREE::ROCM::registerROCMTargetPasses(); }
void onRegisterDialects(DialectRegistry &registry) final {
registry.insert<IREE::ROCM::ROCMDialect>();
}
void populateHALTargetDevices(IREE::HAL::TargetDeviceList &targets) final {
// #hal.device.target<"amdgpu", ...
targets.add("amdgpu", [&]() {
return std::make_shared<AMDGPUTargetDevice>(options);
});
// #hal.device.target<"hip", ...
targets.add("hip",
[&]() { return std::make_shared<HIPTargetDevice>(options); });
}
void populateHALTargetBackends(IREE::HAL::TargetBackendList &targets) final {
// #hal.executable.target<"rocm", ...
targets.add("rocm", [&]() {
LLVMInitializeAMDGPUTarget();
LLVMInitializeAMDGPUTargetMC();
LLVMInitializeAMDGPUTargetInfo();
LLVMInitializeAMDGPUAsmParser();
LLVMInitializeAMDGPUAsmPrinter();
// Also initialize SPIR-V target for amdgcnspirv mode.
LLVMInitializeSPIRVTarget();
LLVMInitializeSPIRVTargetInfo();
LLVMInitializeSPIRVTargetMC();
LLVMInitializeSPIRVAsmPrinter();
return std::make_shared<ROCMTargetBackend>(options, codegenOptions);
});
}
// Override Registration to also bind GPUCodegenOptions to the session.
struct Registration : PluginSession::Registration {
using PluginSession::Registration::Registration;
std::unique_ptr<AbstractPluginSession>
createUninitializedSession(OptionsBinder &localOptionsBinder) final {
auto instance = std::make_unique<ROCMSession>();
// Bootstrap target options from global CLI if available.
if (globalCLIOptions) {
instance->options = *(*globalCLIOptions);
}
instance->options.bindOptions(localOptionsBinder);
// Bootstrap codegen options from global CLI if available.
if (globalCLICodegenOptions) {
instance->codegenOptions = *(*globalCLICodegenOptions);
}
instance->codegenOptions.bindOptions(localOptionsBinder);
return instance;
}
void initializeCLI() final {
PluginSession::Registration::initializeCLI();
globalCLICodegenOptions = &GPUCodegenOptions::FromFlags::get();
}
std::optional<GPUCodegenOptions *> globalCLICodegenOptions;
};
GPUCodegenOptions codegenOptions;
};
// Iterate over ukernel bitcode embedded-data files, and insert them into the
// EmbeddedDataDirectory singleton.
static void addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory() {
EmbeddedDataDirectory::withGlobal([](EmbeddedDataDirectory &dir) {
const iree_file_toc_t *toc = iree_uk_amdgpu_bitcode_create();
for (size_t i = 0; i < iree_uk_amdgpu_bitcode_size(); ++i) {
dir.addFile(toc[i].name, llvm::StringRef{toc[i].data, toc[i].size});
}
});
}
} // namespace
} // namespace mlir::iree_compiler::IREE::HAL
extern "C" bool iree_register_compiler_plugin_hal_target_rocm(
mlir::iree_compiler::PluginRegistrar *registrar) {
registrar->registerPlugin<mlir::iree_compiler::IREE::HAL::ROCMSession>(
"hal_target_rocm");
mlir::iree_compiler::IREE::HAL::
addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory();
return true;
}
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::IREE::HAL::ROCMOptions);