[Mlir-commits] [mlir] [MLIR][GPU] Add xevm-attach-target transform pass. (PR #147372)

[Artem Kroviakov]

================
@@ -0,0 +1,88 @@
+//===-- XeVMAttachTarget.cpp - DESC -----------------------------*- C++ -*-===//
+//
+// This file is 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the `GpuXeVMAttachTarget` pass, attaching `#xevm.target`
+// attributes to GPU modules.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/Regex.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_GPUXEVMATTACHTARGET
+#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::xevm;
+
+namespace {
+struct XeVMAttachTarget
+    : public mlir::impl::GpuXeVMAttachTargetBase<XeVMAttachTarget> {
+  using Base::Base;
+
+  DictionaryAttr getFlags(OpBuilder &builder) const;
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<mlir::xevm::XeVMDialect>();
+  }
+};
+} // namespace
+
+DictionaryAttr XeVMAttachTarget::getFlags(OpBuilder &builder) const {
+  SmallVector<NamedAttribute, 3> flags;
+  // Tokenize and set the optional command line options.
+  if (!cmdOptions.empty()) {
+    auto options = gpu::TargetOptions::tokenizeCmdOptions(cmdOptions);
+    if (!options.second.empty()) {
+      llvm::SmallVector<mlir::Attribute> xevmOptionAttrs;
+      for (const char *opt : options.second) {
+        xevmOptionAttrs.emplace_back(
+            mlir::StringAttr::get(builder.getContext(), StringRef(opt)));
+      }
+      flags.push_back(builder.getNamedAttr(
+          "cmd-options",
+          mlir::ArrayAttr::get(builder.getContext(), xevmOptionAttrs)));
+    }
+  }
+
+  if (!flags.empty())
+    return builder.getDictionaryAttr(flags);
+  return nullptr;
+}
+
+void XeVMAttachTarget::runOnOperation() {
+  OpBuilder builder(&getContext());
+  ArrayRef<std::string> libs(linkLibs);
+  SmallVector<StringRef> filesToLink(libs);
+  auto target = builder.getAttr<mlir::xevm::XeVMTargetAttr>(
+      optLevel, triple, chip, getFlags(builder),
+      filesToLink.empty() ? nullptr : builder.getStrArrayAttr(filesToLink));
+  llvm::Regex matcher(moduleMatcher);
+  // Check if the name of the module matches.
+  auto gpuModule = cast<gpu::GPUModuleOp>(getOperation());
----------------
akroviakov wrote:

Both NV and ROC also assume that the attachment logic will only be applied to GPU modules:
https://github.com/llvm/llvm-project/blob/df6ae4507119b22ffa9e830d5103dfe5b45f99b0/mlir/lib/Dialect/GPU/Transforms/NVVMAttachTarget.cpp#L87

as indicated by their pass description itself:
https://github.com/llvm/llvm-project/blob/a1ea9e632b6db1def1407abcc4c4037144d11eed/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td#L112

This PR decided to rely on the MLIR pass infrastructure for traversal. 

If you think we should follow nv/roc, please consider responding to the [comment](https://github.com/llvm/llvm-project/pull/147372#discussion_r2191944264) below.

https://github.com/llvm/llvm-project/pull/147372