add attach target pass

akroviakov · akroviakov · commit 1cbd14cdbc2b · 2024-12-17T09:54:51.000Z
diff --git a/include/gc/Transforms/Passes.td b/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 @xevm_module_1 {...}
+    gpu.module @xevm_module_2 {...}
+    gpu.module @xevm_module_1 {...}
+    // mlir-opt --xevm-attach-target="module=xevm.* chip=pvc" in.mlir
+    gpu.module @xevm_module_1 {...}
+    gpu.module @xevm_module_2 {...}
+    gpu.module @xevm_module_1 [#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",
diff --git a/lib/gc/Target/LLVM/CMakeLists.txt b/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
diff --git a/lib/gc/Transforms/GPU/CMakeLists.txt b/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}
diff --git a/lib/gc/Transforms/GPU/XeVMAttachTarget.cpp b/lib/gc/Transforms/GPU/XeVMAttachTarget.cpp
@@ -0,0 +1,82 @@
+//===-- 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"
+#include <iostream>
+
+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
+
+DictionaryAttr XeVMAttachTarget::getFlags(OpBuilder &builder) const {
+  UnitAttr unitAttr = builder.getUnitAttr();
+  SmallVector<NamedAttribute, 2> flags;
+  auto addFlag = [&](StringRef flag) {
+    flags.push_back(builder.getNamedAttr(flag, unitAttr));
+  };
+  if (!flags.empty())
+    return builder.getDictionaryAttr(flags);
+  return nullptr;
+}
+
+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));
+      }
+}
diff --git a/src/gc-opt/CMakeLists.txt b/src/gc-opt/CMakeLists.txt
@@ -29,6 +29,10 @@ if(GC_DEV_LINK_LLVM_DYLIB)
 else()
   set(MLIR_LINK_COMPONENTS
     MLIROptLib
+    MLIRBuiltinToLLVMIRTranslation
+    MLIRLLVMDialect
+    MLIRLLVMToLLVMIRTranslation
+    MLIRToLLVMIRTranslationRegistration
   )
   get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
   get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
@@ -38,9 +42,12 @@ add_llvm_executable(gc-opt gc-opt.cpp)
 llvm_update_compile_flags(gc-opt)
 mlir_check_all_link_libraries(gc-opt)
 
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
+
 target_link_libraries(gc-opt PUBLIC GcInterface)
 target_link_libraries(gc-opt PRIVATE
   ${dialect_libs}
+  ${extension_libs}
   ${conversion_libs}
   ${MLIR_LINK_COMPONENTS}
   GcPasses
diff --git a/src/gc-opt/gc-opt.cpp b/src/gc-opt/gc-opt.cpp
@@ -24,10 +24,14 @@
 #ifdef GC_HAS_ONEDNN_DIALECT
 #include "gc/Dialect/OneDNNGraph/OneDNNGraphDialect.h"
 #endif
+#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
@@ -66,11 +70,15 @@ int main(int argc, char *argv[]) {
   registry.insert<mlir::cpuruntime::CPURuntimeDialect>();
   registry.insert<mlir::linalgx::LinalgxDialect>();
   registry.insert<mlir::microkernel::MicrokernelDialect>();
-  registry.insert<mlir::xevm::XeVMDialect>();
   mlir::registerAllDialects(registry);
 #ifdef GC_USE_IMEX
-  registry.insert<::imex::xetile::XeTileDialect, ::imex::gpux::GPUXDialect>();
+  registry.insert<::imex::xetile::XeTileDialect, ::imex::gpux::GPUXDialect,
+                  mlir::xevm::XeVMDialect>();
+  mlir::registerXeVMDialectTranslation(registry);
+  mlir::xevm::registerXeVMTargetInterfaceExternalModels(registry);
 #endif
+  mlir::registerAllExtensions(registry);
+  mlir::registerAllToLLVMIRTranslations(registry);
   mlir::cpuruntime::registerConvertCPURuntimeToLLVMInterface(registry);
   return mlir::asMainReturnCode(mlir::MlirOptMain(
       argc, argv, "Graph Compiler modular optimizer driver\n", registry));
diff --git a/test/mlir/test/gc/Transforms/GPU/module-to-binary-xevm.mlir b/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
+    }
+  }
+}
+
diff --git a/test/mlir/test/gc/Transforms/GPU/xevm-attach-target.mlir b/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
+    }
+  }
+}
+
