lower via upstream path

Author: akroviakov

include/gc/Conversion/XeVMToLLVM/XeVMToLLVM.h

@@ -17,7 +17,7 @@ class RewritePatternSet;
 class Pass;
 
 #define GEN_PASS_DECL_CONVERTXEVMTOLLVMPASS
-#include "mlir/Conversion/Passes.h.inc"
+#include "gc/Conversion/Passes.h.inc"
 
 void populateXeVMToLLVMConversionPatterns(RewritePatternSet &patterns);
 

include/gc/ExecutionEngine/Driver/Driver.h

@@ -18,7 +18,7 @@ namespace mlir {
 class DialectRegistry;
 namespace gc {
 
-const DialectRegistry &initCompilerAndGetDialects();
+DialectRegistry &initCompilerAndGetDialects();
 
 // the pointers to XXXMemRefType
 using GeneralMemrefPtr = void *;

include/gc/Transforms/Passes.td

@@ -261,4 +261,38 @@ def LowerToTileVector : Pass<"lower-to-tile-vector", "func::FuncOp"> {
   ];
 }
 
+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_DIALECT_GC_PASSES

lib/gc/ExecutionEngine/Driver/CMakeLists.txt

@@ -25,6 +25,7 @@ else()
     MLIRToLLVMIRTranslationRegistration
   )
 endif()
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 set(GC_PASSES GcInterface GcPasses)
 if(GC_ENABLE_IMEX)
@@ -38,6 +39,7 @@ gc_add_mlir_library(GcJitWrapper
     ${MLIR_LINK_COMPONENTS}
     ${dialect_libs}
     ${conversion_libs}
+    ${extension_libs}
     ${GC_PASSES}
     GcAnalysis
     MLIRXeVMToLLVMIRTranslation

lib/gc/ExecutionEngine/Driver/Driver.cpp

@@ -12,9 +12,11 @@
 #include "gc/Dialect/OneDNNGraph/OneDNNGraphDialect.h"
 #endif
 #include "gc/Conversion/Passes.h"
+#include "gc/Target/LLVM/XeVM/Target.h"
 #include "gc/Target/LLVMIR/Dialect/XeVM/XeVMToLLVMIRTranslation.h"
 #include "gc/Transforms/Passes.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllExtensions.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Target/LLVMIR/Dialect/All.h"
@@ -34,19 +36,23 @@ static DialectRegistry initDialects() {
   registry.insert<mlir::cpuruntime::CPURuntimeDialect>();
   mlir::registerAllDialects(registry);
   mlir::cpuruntime::registerConvertCPURuntimeToLLVMInterface(registry);
+  mlir::registerAllExtensions(registry);
+  // Adds missing `LLVMTranslationDialectInterface` registration for dialect for
+  // gpu.module op
+  mlir::registerAllToLLVMIRTranslations(registry);
   mlir::registerConvertXeVMToLLVMInterface(registry);
+  mlir::registerXeVMDialectTranslation(registry);
+  mlir::xevm::registerXeVMTargetInterfaceExternalModels(registry);
 #ifdef GC_HAS_ONEDNN_DIALECT
   registry.insert<mlir::onednn_graph::OneDNNGraphDialect>();
 #endif
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   llvm::InitializeNativeTargetAsmParser();
-  mlir::registerAllToLLVMIRTranslations(registry);
-  mlir::registerXeVMDialectTranslation(registry);
   return registry;
 }
 
-const DialectRegistry &initCompilerAndGetDialects() {
+DialectRegistry &initCompilerAndGetDialects() {
   static DialectRegistry reg = initDialects();
   return reg;
 }

lib/gc/Target/LLVM/CMakeLists.txt

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

lib/gc/Transforms/CMakeLists.txt

@@ -30,7 +30,8 @@ gc_add_mlir_library(GcPasses
 
   DEPENDS
     GraphCompilerPassIncGen
-
+    GCConversionPassIncGen
+    
   LINK_LIBS PUBLIC
     ${mlir_dialect_libs}
     ${mlir_conversion_libs}

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,82 @@
+//===- XeVMAttachTarget.cpp - Attach an XeVM 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
+//
+//===----------------------------------------------------------------------===//
+//
+// 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));
+      }
+}

src/gc-cpu-runner/CMakeLists.txt

@@ -32,5 +32,5 @@ if(GC_DEV_LINK_LLVM_DYLIB)
 endif()
 
 gc_add_mlir_tool(gc-cpu-runner gc-cpu-runner.cpp)
-target_link_libraries(gc-cpu-runner PRIVATE GcCpuRuntime)
+target_link_libraries(gc-cpu-runner PRIVATE GcJitWrapper GcCpuRuntime)
 mlir_check_all_link_libraries(gc-cpu-runner)