[mlir][gpu] Add passes to attach (NVVM|ROCDL) target attributes to GPU Modules

[fmorac]

Summary

Adds the passes nvvm-attach-target & rocdl-attach-target for attaching nvvm.target` & rocdl.target attributes to GPU Modules.

These passes search GPU Modules in the immediate region of the Op being acted on, attaching the target attribute to the module.
Modules can be selected using a regex string, allowing fine grain attachment of targets, see the test attach-target.mlir for an example.

Depends on D154153

[mehdi_amini]

mlir/lib/Dialect/LLVMIR/Transforms/NVVMAttachTarget.cpp

70

Nit, inverse the condition, use continue;, and reduce indentation

mlir/lib/Dialect/LLVMIR/Transforms/ROCDLAttachTarget.cpp

61

Don't we have a C++ FastMathFlag class with enums, the string interface is... not great in general.

[fmorac]

mlir/lib/Dialect/LLVMIR/Transforms/ROCDLAttachTarget.cpp

61

Yes, but these have a more direct meaning in AMDGCN, these are names of control variables

https://github.com/RadeonOpenCompute/ROCm-Device-Libs/blob/amd-stg-open/doc/OCML.md#controls

I'd need to change the syntax in the attribute, but I can change it to use the arith Fast Math enum.

[mehdi_amini]

mlir/lib/Dialect/LLVMIR/Transforms/ROCDLAttachTarget.cpp

61

Forget it, I misread the context here.

[fmorac]

Addressed the comments.

The extra changes to LLVM Transforms are because the shared lib build failed in
the previous version -there was a cyclic dependency.

The issue was caused by LLVMTranslation using LLVMTransforms and NVVMTarget using libs in LLVMTranslation.
The solution is to move out the function needed by LLVMTranslation to a separete utils lib.

[mehdi_amini]

mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt

21

It seems a bit surprising to me that LLVM transforms would depend on the GPU stuff.
Seems like this new pass should move elsewhere instead?

[fmorac]

mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt

21

I added it here because it was related to NVVM & ROCDL, but I'll move it to GPU.

[fmorac]

Moved passes to GPU.

mlir/include/mlir/Dialect/LLVMIR/Transforms/AttachTarget.h

+//===- 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 <memory>
+
+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<Pass> createNVVMAttachTarget();
+
+/// Creates a pass that attaches NVVM targets to GPU Module Ops using `options`.
+std::unique_ptr<Pass>
+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<Pass> createROCDLAttachTarget();
+
+/// Creates a pass that attaches ROCDL targets to GPU Module Ops using
+/// `options`.
+std::unique_ptr<Pass>
+createROCDLAttachTarget(const ROCDLAttachTargetOptions &options);
+} // namespace ROCDL
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_ATTACHTARGET_H

mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.h

 //===- Passes.h - LLVM Pass Construction and Registration -----------------===//
 //
 // 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_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"
 #include "mlir/Dialect/LLVMIR/Transforms/TypeConsistency.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 

mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td

let description = [{
     complete debug informations, however it is a convenient way to get line
     tables for debugging purposes. This allow to step trough in a debugger
     line-by-line or get a backtrace with line numbers.
   }];
 
   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<chip = "sm_90">] {...}
+    gpu.module @nvvm_module_2 [#nvvm.target<chip = "sm_90">] {...}
+    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<chip = "gfx90a">] {...}
+    ```
+  }];
+  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

mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLLVMIRTransforms
   DIScopeForLLVMFuncOp.cpp
   LegalizeForExport.cpp
+  NVVMAttachTarget.cpp
   OptimizeForNVVM.cpp
   RequestCWrappers.cpp
+  ROCDLAttachTarget.cpp
   TypeConsistency.cpp
 
   DEPENDS
   MLIRLLVMPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRFuncDialect
   MLIRLLVMDialect
   MLIRPass
   MLIRTransforms
   MLIRNVVMDialect
   )

mlir/lib/Dialect/LLVMIR/Transforms/NVVMAttachTarget.cpp

+//===- 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<NVVMAttachTarget> {
+  using Base::Base;
+
+  DictionaryAttr getFlags(OpBuilder &builder) const;
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registerNVVMTarget(registry);
+  }
+};
+} // namespace
+
+DictionaryAttr NVVMAttachTarget::getFlags(OpBuilder &builder) const {
+  UnitAttr unitAttr = builder.getUnitAttr();
+  SmallVector<NamedAttribute, 2> 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<std::string> libs(linkLibs);
+  SmallVector<StringRef> filesToLink(libs.begin(), libs.end());
+  auto target = builder.getAttr<NVVMTargetAttr>(
+      optLevel, triple, chip, features, getFlags(builder),
+      filesToLink.size() ? builder.getStrArrayAttr(filesToLink) : nullptr);
+  llvm::Regex matcher(moduleMatcher);
+  for (Region &region : getOperation()->getRegions())
+    for (Block &block : region.getBlocks())
+      for (auto module : block.getOps<gpu::GPUModuleOp>()) {
+        // Check if the name of the module matches.
+        if (moduleMatcher.empty() || matcher.match(module.getName())) {
+          // Create the target array.
+          SmallVector<Attribute> targets;
+          if (std::optional<ArrayAttr> 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 attray.
+          module.setTargetsAttr(builder.getArrayAttr(targets));
+        }
+      }
+}
+
+std::unique_ptr<Pass> NVVM::createNVVMAttachTarget() {
+  return std::make_unique<NVVMAttachTarget>();
+}
+
+std::unique_ptr<Pass>
+NVVM::createNVVMAttachTarget(const NVVMAttachTargetOptions &options) {
+  return std::make_unique<NVVMAttachTarget>(options);
+}

mlir/lib/Dialect/LLVMIR/Transforms/ROCDLAttachTarget.cpp

+//===- 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<ROCDLAttachTarget> {
+  using Base::Base;
+
+  DictionaryAttr getFlags(OpBuilder &builder) const;
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registerROCDLTarget(registry);
+  }
+};
+} // namespace
+
+DictionaryAttr ROCDLAttachTarget::getFlags(OpBuilder &builder) const {
+  UnitAttr unitAttr = builder.getUnitAttr();
+  SmallVector<NamedAttribute, 6> 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<std::string> libs(linkLibs);
+  SmallVector<StringRef> filesToLink(libs.begin(), libs.end());
+  auto target = builder.getAttr<ROCDLTargetAttr>(
+      optLevel, triple, chip, features, abiVersion, getFlags(builder),
+      filesToLink.size() ? builder.getStrArrayAttr(filesToLink) : nullptr);
+  llvm::Regex matcher(moduleMatcher);
+  for (Region &region : getOperation()->getRegions())
+    for (Block &block : region.getBlocks())
+      for (auto module : block.getOps<gpu::GPUModuleOp>()) {
+        // Check if the name of the module matches.
+        if (moduleMatcher.empty() || matcher.match(module.getName())) {
+          // Create the target array.
+          SmallVector<Attribute> targets;
+          if (std::optional<ArrayAttr> 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 attray.
+          module.setTargetsAttr(builder.getArrayAttr(targets));
+        }
+      }
+}
+
+std::unique_ptr<Pass> ROCDL::createROCDLAttachTarget() {
+  return std::make_unique<ROCDLAttachTarget>();
+}
+
+std::unique_ptr<Pass>
+ROCDL::createROCDLAttachTarget(const ROCDLAttachTargetOptions &options) {
+  return std::make_unique<ROCDLAttachTarget>(options);
+}

mlir/test/Dialect/LLVMIR/attach-targets.mlir

+// 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<O = 3, chip = "sm_90">] {
+gpu.module @nvvm_module_1 {
+}
+// Verify the target is appended.
+// CHECK: @nvvm_module_2 [#nvvm.target<chip = "sm_60">, #nvvm.target<O = 3, chip = "sm_90">] {
+gpu.module @nvvm_module_2 [#nvvm.target<chip = "sm_60">] {
+}
+// Verify the target is not added multiple times.
+// CHECK: @nvvm_module_3 [#nvvm.target<O = 3, chip = "sm_90">] {
+gpu.module @nvvm_module_3 [#nvvm.target<O = 3, chip = "sm_90">] {
+}
+// 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<O = 3, chip = "gfx90a">] {
+gpu.module @rocdl_module {
+}
+// Check the options were added.
+// CHECK_OPTS: @options_module_1 [#nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>]  {
+gpu.module @options_module_1 {
+}
+// Check the options were added and that the first target was preserved.
+// CHECK_OPTS: @options_module_2 [#nvvm.target<O = 3, chip = "sm_90">, #nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>]  {
+gpu.module @options_module_2 [#nvvm.target<O = 3, chip = "sm_90">] {
+}
+}