diff --git a/mlir/include/mlir/Dialect/LLVMIR/Transforms/AttachTarget.h b/mlir/include/mlir/Dialect/LLVMIR/Transforms/AttachTarget.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/LLVMIR/Transforms/AttachTarget.h @@ -0,0 +1,42 @@ +//===- AttachTarget.h - Attaches GPU targets to GPU Modules -----*- C++ -*-===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_LLVMIR_TRANSFORMS_ATTACHTARGET_H +#define MLIR_DIALECT_LLVMIR_TRANSFORMS_ATTACHTARGET_H + +#include "mlir/Pass/Pass.h" +#include + +namespace mlir { +namespace NVVM { +#define GEN_PASS_DECL_NVVMATTACHTARGET +#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc" + +/// Creates a pass that attaches NVVM targets to GPU Module Ops. +std::unique_ptr createNVVMAttachTarget(); + +/// Creates a pass that attaches NVVM targets to GPU Module Ops using `options`. +std::unique_ptr +createNVVMAttachTarget(const NVVMAttachTargetOptions &options); +} // namespace NVVM + +namespace ROCDL { +#define GEN_PASS_DECL_ROCDLATTACHTARGET +#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc" + +/// Creates a pass that attaches ROCDL targets to GPU Module Ops. +std::unique_ptr createROCDLAttachTarget(); + +/// Creates a pass that attaches ROCDL targets to GPU Module Ops using +/// `options`. +std::unique_ptr +createROCDLAttachTarget(const ROCDLAttachTargetOptions &options); +} // namespace ROCDL +} // namespace mlir + +#endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_ATTACHTARGET_H diff --git a/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.h b/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.h --- a/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.h @@ -9,6 +9,7 @@ #ifndef MLIR_DIALECT_LLVMIR_TRANSFORMS_PASSES_H #define MLIR_DIALECT_LLVMIR_TRANSFORMS_PASSES_H +#include "mlir/Dialect/LLVMIR/Transforms/AttachTarget.h" #include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h" #include "mlir/Dialect/LLVMIR/Transforms/OptimizeForNVVM.h" #include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h" diff --git a/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td b/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td --- a/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td @@ -68,4 +68,111 @@ let constructor = "mlir::LLVM::createDIScopeForLLVMFuncOpPass()"; } +def NVVMAttachTarget: Pass<"nvvm-attach-target", ""> { + let summary = "Attaches an NVVM target attribute to a GPU Module."; + let description = [{ + This pass searches for all GPU Modules in the immediate regions and attaches + an NVVM target if the module matches the name specified by the `module` argument. + + Example: + ``` + // File: in.mlir: + gpu.module @nvvm_module_1 {...} + gpu.module @nvvm_module_2 {...} + gpu.module @rocdl_module_1 {...} + // mlir-opt --nvvm-attach-target="module=nvvm.* chip=sm_90" in.mlir + gpu.module @nvvm_module_1 [#nvvm.target] {...} + gpu.module @nvvm_module_2 [#nvvm.target] {...} + gpu.module @rocdl_module_1 {...} + ``` + }]; + let options = [ + Option<"moduleMatcher", "module", "std::string", + /*default=*/ [{""}], + "Regex used to identify the modules to attach the target to.">, + Option<"triple", "triple", "std::string", + /*default=*/ "\"nvptx64-nvidia-cuda\"", + "Target triple.">, + Option<"chip", "chip", "std::string", + /*default=*/"\"sm_50\"", + "Target chip.">, + Option<"features", "features", "std::string", + /*default=*/"\"+ptx60\"", + "Target features.">, + Option<"optLevel", "O", "unsigned", + /*default=*/"2", + "Optimization level.">, + Option<"fastFlag", "fast", "bool", + /*default=*/"false", + "Enable fast math mode.">, + Option<"ftzFlag", "ftz", "bool", + /*default=*/"false", + "Enable flush to zero for denormals.">, + ListOption<"linkLibs", "l", "std::string", + "Extra bitcode libraries paths to link to.">, + ]; + let constructor = "mlir::NVVM::createNVVMAttachTarget()"; +} + +def ROCDLAttachTarget: Pass<"rocdl-attach-target", ""> { + let summary = "Attaches a ROCDL target attribute to a GPU Module."; + let description = [{ + This pass searches for all GPU Modules in the immediate regions and attaches + a ROCDL target if the module matches the name specified by the `module` argument. + + Example: + ``` + // File: in.mlir: + gpu.module @nvvm_module_1 {...} + gpu.module @nvvm_module_2 {...} + gpu.module @rocdl_module_1 {...} + // mlir-opt --nvvm-attach-target="module=rocdl.* chip=gfx90a" in.mlir + gpu.module @nvvm_module_1 {...} + gpu.module @nvvm_module_2 {...} + gpu.module @rocdl_module_1 [#rocdl.target] {...} + ``` + }]; + let options = [ + Option<"moduleMatcher", "module", "std::string", + /*default=*/ [{""}], + "Regex used to identify the modules to attach the target to.">, + Option<"triple", "triple", "std::string", + /*default=*/ "\"amdgcn-amd-amdhsa\"", + "Target triple.">, + Option<"chip", "chip", "std::string", + /*default=*/"\"gfx900\"", + "Target chip.">, + Option<"features", "features", "std::string", + /*default=*/"\"\"", + "Target features.">, + Option<"abiVersion", "abi", "std::string", + /*default=*/"\"400\"", + "Optimization level.">, + Option<"optLevel", "O", "unsigned", + /*default=*/"2", + "Optimization level.">, + Option<"wave64Flag", "wave64", "bool", + /*default=*/"true", + "Use Wave64 mode.">, + Option<"fastFlag", "fast", "bool", + /*default=*/"false", + "Enable fast relaxed math opt.">, + Option<"dazFlag", "daz", "bool", + /*default=*/"false", + "Enable denormals are zero opt.">, + Option<"finiteOnlyFlag", "finite-only", "bool", + /*default=*/"false", + "Enable finite only opt.">, + Option<"unsafeMathFlag", "unsafe-math", "bool", + /*default=*/"false", + "Enable unsafe math opt.">, + Option<"correctSqrtFlag", "correct-sqrt", "bool", + /*default=*/"true", + "Enable correct rounded sqrt.">, + ListOption<"linkLibs", "l", "std::string", + "Extra bitcode libraries paths to link to.">, + ]; + let constructor = "mlir::ROCDL::createROCDLAttachTarget()"; +} + #endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_PASSES diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt b/mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt --- a/mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt @@ -1,8 +1,12 @@ +add_subdirectory(Utils) + add_mlir_dialect_library(MLIRLLVMIRTransforms DIScopeForLLVMFuncOp.cpp LegalizeForExport.cpp + NVVMAttachTarget.cpp OptimizeForNVVM.cpp RequestCWrappers.cpp + ROCDLAttachTarget.cpp TypeConsistency.cpp DEPENDS @@ -11,8 +15,11 @@ LINK_LIBS PUBLIC MLIRIR MLIRFuncDialect + MLIRGPUDialect MLIRLLVMDialect MLIRPass MLIRTransforms MLIRNVVMDialect + MLIRNVVMTarget + MLIRROCDLTarget ) diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp --- a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp +++ b/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp @@ -23,57 +23,6 @@ using namespace mlir; -/// If the given block has the same successor with different arguments, -/// introduce dummy successor blocks so that all successors of the given block -/// are different. -static void ensureDistinctSuccessors(Block &bb) { - // Early exit if the block cannot have successors. - if (bb.empty() || !bb.back().mightHaveTrait()) - return; - - auto *terminator = bb.getTerminator(); - - // Find repeated successors with arguments. - llvm::SmallDenseMap> successorPositions; - for (int i = 0, e = terminator->getNumSuccessors(); i < e; ++i) { - Block *successor = terminator->getSuccessor(i); - // Blocks with no arguments are safe even if they appear multiple times - // because they don't need PHI nodes. - if (successor->getNumArguments() == 0) - continue; - successorPositions[successor].push_back(i); - } - - // If a successor appears for the second or more time in the terminator, - // create a new dummy block that unconditionally branches to the original - // destination, and retarget the terminator to branch to this new block. - // There is no need to pass arguments to the dummy block because it will be - // dominated by the original block and can therefore use any values defined in - // the original block. - OpBuilder builder(terminator->getContext()); - for (const auto &successor : successorPositions) { - // Start from the second occurrence of a block in the successor list. - for (int position : llvm::drop_begin(successor.second, 1)) { - Block *dummyBlock = builder.createBlock(bb.getParent()); - terminator->setSuccessor(dummyBlock, position); - for (BlockArgument arg : successor.first->getArguments()) - dummyBlock->addArgument(arg.getType(), arg.getLoc()); - builder.create(terminator->getLoc(), - dummyBlock->getArguments(), successor.first); - } - } -} - -void mlir::LLVM::ensureDistinctSuccessors(Operation *op) { - op->walk([](Operation *nested) { - for (Region ®ion : llvm::make_early_inc_range(nested->getRegions())) { - for (Block &block : llvm::make_early_inc_range(region)) { - ::ensureDistinctSuccessors(block); - } - } - }); -} - namespace { struct LegalizeForExportPass : public LLVM::impl::LLVMLegalizeForExportBase { diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/NVVMAttachTarget.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/NVVMAttachTarget.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/LLVMIR/Transforms/NVVMAttachTarget.cpp @@ -0,0 +1,92 @@ +//===- NVVMAttachTarget.cpp - Attach an NVVM target -----------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/LLVMIR/Transforms/AttachTarget.h" + +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/IR/Builders.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Target/LLVM/NVVM/Target.h" +#include "llvm/Support/Regex.h" + +namespace mlir { +namespace NVVM { +#define GEN_PASS_DEF_NVVMATTACHTARGET +#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc" +} // namespace NVVM +} // namespace mlir + +using namespace mlir; +using namespace mlir::NVVM; + +namespace { +struct NVVMAttachTarget + : public NVVM::impl::NVVMAttachTargetBase { + using Base::Base; + + DictionaryAttr getFlags(OpBuilder &builder) const; + + void runOnOperation() override; + + void getDependentDialects(DialectRegistry ®istry) const override { + registerNVVMTarget(registry); + } +}; +} // namespace + +DictionaryAttr NVVMAttachTarget::getFlags(OpBuilder &builder) const { + UnitAttr unitAttr = builder.getUnitAttr(); + SmallVector flags; + auto addFlag = [&](StringRef flag) { + flags.push_back(builder.getNamedAttr(flag, unitAttr)); + }; + if (fastFlag) + addFlag("fast"); + if (ftzFlag) + addFlag("ftz"); + if (flags.size()) + return builder.getDictionaryAttr(flags); + return nullptr; +} + +void NVVMAttachTarget::runOnOperation() { + OpBuilder builder(&getContext()); + ArrayRef libs(linkLibs); + SmallVector filesToLink(libs.begin(), libs.end()); + auto target = builder.getAttr( + optLevel, triple, chip, features, getFlags(builder), + filesToLink.size() ? builder.getStrArrayAttr(filesToLink) : nullptr); + llvm::Regex matcher(moduleMatcher); + for (Region ®ion : getOperation()->getRegions()) + for (Block &block : region.getBlocks()) + for (auto module : block.getOps()) { + // Check if the name of the module matches. + if (!moduleMatcher.empty() && !matcher.match(module.getName())) + continue; + // Create the target array. + SmallVector targets; + if (std::optional attrs = module.getTargets()) + targets.append(attrs->getValue().begin(), attrs->getValue().end()); + targets.push_back(target); + // Remove any duplicate targets. + targets.erase(std::unique(targets.begin(), targets.end()), + targets.end()); + // Update the target attribute array. + module.setTargetsAttr(builder.getArrayAttr(targets)); + } +} + +std::unique_ptr NVVM::createNVVMAttachTarget() { + return std::make_unique(); +} + +std::unique_ptr +NVVM::createNVVMAttachTarget(const NVVMAttachTargetOptions &options) { + return std::make_unique(options); +} diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/ROCDLAttachTarget.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/ROCDLAttachTarget.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/LLVMIR/Transforms/ROCDLAttachTarget.cpp @@ -0,0 +1,101 @@ +//===- ROCDLAttachTarget.cpp - Attach an ROCDL target +//-----------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/LLVMIR/Transforms/AttachTarget.h" + +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/LLVMIR/ROCDLDialect.h" +#include "mlir/IR/Builders.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Target/LLVM/ROCDL/Target.h" +#include "llvm/Support/Regex.h" + +namespace mlir { +namespace ROCDL { +#define GEN_PASS_DEF_ROCDLATTACHTARGET +#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc" +} // namespace ROCDL +} // namespace mlir + +using namespace mlir; +using namespace mlir::ROCDL; + +namespace { +struct ROCDLAttachTarget + : public ROCDL::impl::ROCDLAttachTargetBase { + using Base::Base; + + DictionaryAttr getFlags(OpBuilder &builder) const; + + void runOnOperation() override; + + void getDependentDialects(DialectRegistry ®istry) const override { + registerROCDLTarget(registry); + } +}; +} // namespace + +DictionaryAttr ROCDLAttachTarget::getFlags(OpBuilder &builder) const { + UnitAttr unitAttr = builder.getUnitAttr(); + SmallVector flags; + auto addFlag = [&](StringRef flag) { + flags.push_back(builder.getNamedAttr(flag, unitAttr)); + }; + if (!wave64Flag) + addFlag("no_wave64"); + if (fastFlag) + addFlag("fast"); + if (dazFlag) + addFlag("daz"); + if (finiteOnlyFlag) + addFlag("finite_only"); + if (unsafeMathFlag) + addFlag("unsafe_math"); + if (!correctSqrtFlag) + addFlag("unsafe_sqrt"); + if (flags.size()) + return builder.getDictionaryAttr(flags); + return nullptr; +} + +void ROCDLAttachTarget::runOnOperation() { + OpBuilder builder(&getContext()); + ArrayRef libs(linkLibs); + SmallVector filesToLink(libs.begin(), libs.end()); + auto target = builder.getAttr( + optLevel, triple, chip, features, abiVersion, getFlags(builder), + filesToLink.size() ? builder.getStrArrayAttr(filesToLink) : nullptr); + llvm::Regex matcher(moduleMatcher); + for (Region ®ion : getOperation()->getRegions()) + for (Block &block : region.getBlocks()) + for (auto module : block.getOps()) { + // Check if the name of the module matches. + if (!moduleMatcher.empty() && !matcher.match(module.getName())) + continue; + // Create the target array. + SmallVector targets; + if (std::optional attrs = module.getTargets()) + targets.append(attrs->getValue().begin(), attrs->getValue().end()); + targets.push_back(target); + // Remove any duplicate targets. + targets.erase(std::unique(targets.begin(), targets.end()), + targets.end()); + // Update the target attribute array. + module.setTargetsAttr(builder.getArrayAttr(targets)); + } +} + +std::unique_ptr ROCDL::createROCDLAttachTarget() { + return std::make_unique(); +} + +std::unique_ptr +ROCDL::createROCDLAttachTarget(const ROCDLAttachTargetOptions &options) { + return std::make_unique(options); +} diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/Utils/CMakeLists.txt b/mlir/lib/Dialect/LLVMIR/Transforms/Utils/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/LLVMIR/Transforms/Utils/CMakeLists.txt @@ -0,0 +1,10 @@ +add_mlir_dialect_library(MLIRLLVMIRTransformsUtils + EnsureDistinctSuccessors.cpp + + DEPENDS + MLIRLLVMPassIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRLLVMDialect + ) diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/Utils/EnsureDistinctSuccessors.cpp copy from mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp copy to mlir/lib/Dialect/LLVMIR/Transforms/Utils/EnsureDistinctSuccessors.cpp --- a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp +++ b/mlir/lib/Dialect/LLVMIR/Transforms/Utils/EnsureDistinctSuccessors.cpp @@ -1,4 +1,4 @@ -//===- LegalizeForExport.cpp - Prepare for translation to LLVM IR ---------===// +//===- EnsureDistinctSuccessors.cpp - Prepare for translation to LLVM IR --===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. @@ -12,14 +12,6 @@ #include "mlir/IR/Block.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" - -namespace mlir { -namespace LLVM { -#define GEN_PASS_DEF_LLVMLEGALIZEFOREXPORT -#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc" -} // namespace LLVM -} // namespace mlir using namespace mlir; @@ -73,16 +65,3 @@ } }); } - -namespace { -struct LegalizeForExportPass - : public LLVM::impl::LLVMLegalizeForExportBase { - void runOnOperation() override { - LLVM::ensureDistinctSuccessors(getOperation()); - } -}; -} // namespace - -std::unique_ptr LLVM::createLegalizeForExportPass() { - return std::make_unique(); -} diff --git a/mlir/lib/Target/LLVMIR/CMakeLists.txt b/mlir/lib/Target/LLVMIR/CMakeLists.txt --- a/mlir/lib/Target/LLVMIR/CMakeLists.txt +++ b/mlir/lib/Target/LLVMIR/CMakeLists.txt @@ -37,7 +37,8 @@ LINK_LIBS PUBLIC MLIRDLTIDialect MLIRLLVMDialect - MLIRLLVMIRTransforms + MLIRLLVMIRTransformsUtils + MLIRTransforms MLIRTranslateLib ) diff --git a/mlir/test/Dialect/LLVMIR/attach-targets.mlir b/mlir/test/Dialect/LLVMIR/attach-targets.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/LLVMIR/attach-targets.mlir @@ -0,0 +1,29 @@ +// RUN: mlir-opt %s --nvvm-attach-target='module=nvvm.* O=3 chip=sm_90' --rocdl-attach-target='module=rocdl.* O=3 chip=gfx90a' | FileCheck %s +// RUN: mlir-opt %s --nvvm-attach-target='module=options.* O=1 chip=sm_70 fast=true ftz=true' --rocdl-attach-target='module=options.* l=file1.bc,file2.bc wave64=false finite-only=true' | FileCheck %s --check-prefix=CHECK_OPTS + +module attributes {gpu.container_module} { +// Verify the target is appended. +// CHECK: @nvvm_module_1 [#nvvm.target] { +gpu.module @nvvm_module_1 { +} +// Verify the target is appended. +// CHECK: @nvvm_module_2 [#nvvm.target, #nvvm.target] { +gpu.module @nvvm_module_2 [#nvvm.target] { +} +// Verify the target is not added multiple times. +// CHECK: @nvvm_module_3 [#nvvm.target] { +gpu.module @nvvm_module_3 [#nvvm.target] { +} +// Verify the NVVM target is not added as it fails to match the regex, but the ROCDL does get appended. +// CHECK: @rocdl_module [#rocdl.target] { +gpu.module @rocdl_module { +} +// Check the options were added. +// CHECK_OPTS: @options_module_1 [#nvvm.target, #rocdl.target] { +gpu.module @options_module_1 { +} +// Check the options were added and that the first target was preserved. +// CHECK_OPTS: @options_module_2 [#nvvm.target, #nvvm.target, #rocdl.target] { +gpu.module @options_module_2 [#nvvm.target] { +} +}