[XeVM] Attach target pass #422

This PR adds a pass that attaches an xevm target to a GPU module. This enables generating binaries similar to the upstream via gpu-module-to-binary.

Author: akroviakov

include/gc/Transforms/Passes.td

@@ -150,6 +150,41 @@ def GpuTilingAndFusion : Pass<"gpu-tiling", "func::FuncOp"> {
            "The maximum workgroup size.">
     ];
 }
+
+def GpuXeVMAttachTarget: Pass<"xevm-attach-target", ""> {
+  let summary = "Attaches a XeVM target attribute to a GPU Module.";
+  let description = [{
+    This pass searches for all GPU Modules in the immediate regions and attaches
+    a XeVM target if the module matches the name specified by the `module` argument.
+
+    Example:
+    ```
+    // File: in.mlir:
+    gpu.module @nvvm_module_1 {...}
+    gpu.module @rocdl_module_2 {...}
+    gpu.module @xevm_module_3 {...}
+    // mlir-opt --xevm-attach-target="module=xevm.* chip=pvc" in.mlir
+    gpu.module @nvvm_module_1 {...}
+    gpu.module @rocdl_module_2 {...}
+    gpu.module @xevm_module_3 [#xevm.target<chip = "pvc">] {...}
+    ```
+  }];
+  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=*/ "\"spirv64-unknown-unknown\"",
+           "Target triple.">,
+    Option<"chip", "chip", "std::string",
+           /*default=*/"\"pvc\"",
+           "Target chip.">,
+    Option<"optLevel", "O", "unsigned",
+           /*default=*/"2",
+           "Optimization level.">
+  ];
+}
+
 #endif // GC_USE_IMEX
 
 def IterativeTilingAndFusion : Pass<"iterative-tiling-and-fusion",

lib/gc/Target/LLVM/CMakeLists.txt

@@ -6,6 +6,9 @@ gc_add_mlir_dialect_library(MLIRXeVMTarget
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LLVMIR
   ${PROJECT_SOURCE_DIR}/include/gc/Dialect/LLVMIR
+  
+  LINK_COMPONENTS
+  SPIRVCodeGen
 
   LINK_LIBS PUBLIC
   MLIRIR

lib/gc/Transforms/GPU/CMakeLists.txt

@@ -17,6 +17,7 @@ gc_add_mlir_library(GcGpuPasses
   GpuToGpuOcl.cpp
   LinalgToXeGPU.cpp
   Pipeline.cpp
+  XeVMAttachTarget.cpp
 
   DEPENDS
     GraphCompilerPassIncGen
@@ -31,6 +32,7 @@ gc_add_mlir_library(GcGpuPasses
     MLIRMathToSPIRV
     MLIRControlFlowToSPIRV
     MLIRMemRefTransforms
+    MLIRXeVMToLLVMIRTranslation
     GcInterface
     GcUtilsIR
     ${IMEX_LIBS}

lib/gc/Transforms/GPU/XeVMAttachTarget.cpp

@@ -0,0 +1,70 @@
+//===-- 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 "gc/Dialect/LLVMIR/XeVMDialect.h"
+
+#include "gc/Target/LLVM/XeVM/Target.h"
+#include "gc/Transforms/Passes.h"
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/Regex.h"
+
+namespace mlir {
+namespace gc {
+#define GEN_PASS_DEF_GPUXEVMATTACHTARGET
+#include "gc/Transforms/Passes.h.inc"
+} // namespace gc
+} // namespace mlir
+
+using namespace mlir::xevm;
+using namespace mlir;
+
+namespace {
+struct XeVMAttachTarget
+    : public gc::impl::GpuXeVMAttachTargetBase<XeVMAttachTarget> {
+  using Base::Base;
+
+  // DictionaryAttr getFlags(OpBuilder &builder) const;
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<xevm::XeVMDialect>();
+  }
+};
+} // namespace
+
+void XeVMAttachTarget::runOnOperation() {
+  OpBuilder builder(&getContext());
+  auto target = builder.getAttr<XeVMTargetAttr>(optLevel, triple, chip);
+  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()))
+          continue;
+        // 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(llvm::unique(targets), targets.end());
+        // Update the target attribute array.
+        module.setTargetsAttr(builder.getArrayAttr(targets));
+      }
+}

src/gc-opt/gc-opt.cpp

@@ -27,11 +27,14 @@
 #include "gc/Conversion/Passes.h"
 #include "mlir/Target/LLVMIR/Dialect/All.h"
 
+#include "gc/Target/LLVM/XeVM/Target.h"
+#include "gc/Target/LLVMIR/Dialect/XeVM/XeVMToLLVMIRTranslation.h"
 #include "gc/Transforms/Microkernel/MicrokernelPasses.h"
 #include "gc/Transforms/Passes.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/InitAllExtensions.h"
 #include "mlir/InitAllPasses.h"
+#include "mlir/Target/LLVMIR/Dialect/All.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 
 #ifdef GC_USE_IMEX
@@ -75,7 +78,11 @@ int main(int argc, char *argv[]) {
   mlir::registerAllDialects(registry);
 #ifdef GC_USE_IMEX
   registry.insert<::imex::xetile::XeTileDialect, ::imex::gpux::GPUXDialect>();
+  mlir::registerXeVMDialectTranslation(registry);
+  mlir::xevm::registerXeVMTargetInterfaceExternalModels(registry);
 #endif
+  mlir::registerAllExtensions(registry);
+  mlir::registerAllToLLVMIRTranslations(registry);
   mlir::cpuruntime::registerConvertCPURuntimeToLLVMInterface(registry);
   mlir::registerAllExtensions(registry); // TODO: cleanup
   // Adds missing `LLVMTranslationDialectInterface` registration for dialect for

test/mlir/test/gc/Transforms/GPU/module-to-binary-xevm.mlir

@@ -0,0 +1,12 @@
+// RUN: gc-opt %s --gpu-to-llvm --convert-gpu-to-llvm-spv --gpu-module-to-binary | FileCheck %s
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL:gpu.binary @entry_kernel
+  // CHECK:[#gpu.object<#xevm.target,
+  gpu.module @entry_kernel [#xevm.target] {
+    gpu.func @entry_kernel(%arg0: index) kernel attributes {} {
+      gpu.return
+    }
+  }
+}
+

test/mlir/test/gc/Transforms/GPU/xevm-attach-target.mlir

@@ -0,0 +1,10 @@
+// RUN: gc-opt %s --xevm-attach-target | FileCheck %s
+module attributes {gpu.container_module} {
+  //CHECK:gpu.module @entry_kernel [#xevm.target]
+  gpu.module @entry_kernel {
+    gpu.func @entry_kernel(%arg0: index) kernel attributes {} {
+      gpu.return
+    }
+  }
+}
+