diff --git a/include/gc/Transforms/Passes.td b/include/gc/Transforms/Passes.td index 1f30b877..b7013fba 100644 --- 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 @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] {...} + ``` + }]; + 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 index 81d4d563..4803f47d 100644 --- 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 index f4b286b9..f7d10b3d 100644 --- 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 new file mode 100644 index 00000000..9eff769c --- /dev/null +++ b/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 { + using Base::Base; + + // DictionaryAttr getFlags(OpBuilder &builder) const; + + void runOnOperation() override; + + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } +}; +} // namespace + +void XeVMAttachTarget::runOnOperation() { + OpBuilder builder(&getContext()); + auto target = builder.getAttr(optLevel, triple, chip); + llvm::Regex matcher(moduleMatcher); + for (Region ®ion : getOperation()->getRegions()) + for (Block &block : region.getBlocks()) + for (auto module : block.getOps()) { + // Check if the name of the module matches. + if (!moduleMatcher.empty() && !matcher.match(module.getName())) + continue; + // Create the target array. + SmallVector targets; + if (std::optional 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/gc-opt.cpp b/src/gc-opt/gc-opt.cpp index 51c17836..9c39ab4e 100644 --- a/src/gc-opt/gc-opt.cpp +++ b/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 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 new file mode 100644 index 00000000..3b3f4a26 --- /dev/null +++ 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 new file mode 100644 index 00000000..d9eecabb --- /dev/null +++ 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 + } + } +} +