[MLIR][GPU] Add xevm-attach-target transform pass. #147372
Conversation
|
@llvm/pr-subscribers-mlir-gpu @llvm/pr-subscribers-mlir-llvm Author: Sang Ik Lee (silee2) ChangesCo-authored-by: Artem Kroviakov [email protected] Full diff: https://github.com/llvm/llvm-project/pull/147372.diff 4 Files Affected:
diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
index 3766eb16e9429..b6fc0a1375fd8 100644
--- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
@@ -258,4 +258,38 @@ def GpuSPIRVAttachTarget: Pass<"spirv-attach-target", ""> {
];
}
+def GpuXeVMAttachTarget : Pass<"xevm-attach-target", "mlir::gpu::GPUModuleOp"> {
+ 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=*/"\"bmg\"", "Target chip.">,
+ Option<"optLevel", "O", "unsigned",
+ /*default=*/"2", "Optimization level.">,
+ ListOption<"linkLibs", "l", "std::string",
+ "Extra bitcode libraries paths to link to.">,
+ Option<"cmdOptions", "cmd-options", "std::string",
+ /*default=*/[{""}],
+ "Command line options passed to downstream compiler">];
+}
#endif // MLIR_DIALECT_GPU_PASSES
diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt
index 4862d1f722785..6e7d622b910f2 100644
--- a/mlir/lib/Dialect/GPU/CMakeLists.txt
+++ b/mlir/lib/Dialect/GPU/CMakeLists.txt
@@ -44,6 +44,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
Transforms/ShuffleRewriter.cpp
Transforms/SubgroupIdRewriter.cpp
Transforms/SubgroupReduceLowering.cpp
+ Transforms/XeVMAttachTarget.cpp
OBJECT
diff --git a/mlir/test/Dialect/LLVMIR/attach-targets.mlir b/mlir/test/Dialect/LLVMIR/attach-targets.mlir
index 83733db400798..09dbcc5f28075 100644
--- a/mlir/test/Dialect/LLVMIR/attach-targets.mlir
+++ b/mlir/test/Dialect/LLVMIR/attach-targets.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt %s --nvvm-attach-target='module=nvvm.* O=3 chip=sm_90' --rocdl-attach-target='module=rocdl.* O=3 chip=gfx90a' | FileCheck %s
-// RUN: mlir-opt %s --nvvm-attach-target='module=options.* O=1 chip=sm_70 fast=true ftz=true' --rocdl-attach-target='module=options.* l=file1.bc,file2.bc wave64=false finite-only=true' | FileCheck %s --check-prefix=CHECK_OPTS
+// RUN: mlir-opt %s --nvvm-attach-target='module=nvvm.* O=3 chip=sm_90' --rocdl-attach-target='module=rocdl.* O=3 chip=gfx90a' --xevm-attach-target='module=xevm.* O=3 chip=pvc' | FileCheck %s
+// RUN: mlir-opt %s --nvvm-attach-target='module=options.* O=1 chip=sm_70 fast=true ftz=true' --rocdl-attach-target='module=options.* l=file1.bc,file2.bc wave64=false finite-only=true' --xevm-attach-target='module=options.* O=1 chip=pvc' | FileCheck %s --check-prefix=CHECK_OPTS
module attributes {gpu.container_module} {
// Verify the target is appended.
@@ -18,12 +18,16 @@ gpu.module @nvvm_module_3 [#nvvm.target<O = 3, chip = "sm_90">] {
// CHECK: @rocdl_module [#rocdl.target<O = 3, chip = "gfx90a">] {
gpu.module @rocdl_module {
}
+// Verify that other targets are not added as they fail to match the regex, but XeVM does get appended.
+// CHECK: @xevm_module [#xevm.target<O = 3, chip = "pvc">] {
+gpu.module @xevm_module {
+}
// Check the options were added.
-// CHECK_OPTS: @options_module_1 [#nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>] {
+// CHECK_OPTS: @options_module_1 [#nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>, #xevm.target<O = 1, chip = "pvc">] {
gpu.module @options_module_1 {
}
// Check the options were added and that the first target was preserved.
-// CHECK_OPTS: @options_module_2 [#nvvm.target<O = 3, chip = "sm_90">, #nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>] {
+// CHECK_OPTS: @options_module_2 [#nvvm.target<O = 3, chip = "sm_90">, #nvvm.target<O = 1, chip = "sm_70", flags = {fast, ftz}>, #rocdl.target<flags = {finite_only, no_wave64}, link = ["file1.bc", "file2.bc"]>, #xevm.target<O = 1, chip = "pvc">] {
gpu.module @options_module_2 [#nvvm.target<O = 3, chip = "sm_90">] {
}
}
diff --git a/mlir/test/lib/Dialect/GPU/CMakeLists.txt b/mlir/test/lib/Dialect/GPU/CMakeLists.txt
index 4ca5974ed5a49..418c884dc03b3 100644
--- a/mlir/test/lib/Dialect/GPU/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/GPU/CMakeLists.txt
@@ -29,6 +29,7 @@ set(LIBS
MLIRTranslateLib
MLIRVectorDialect
MLIRVectorToLLVMPass
+ MLIRXeVMDialect
)
add_mlir_library(MLIRGPUTestPasses
|
rengolin
left a comment
rengolin
left a comment
There was a problem hiding this comment.
A few comments but overall looks good to me, thanks!
I'll let others review and approve.
There was a problem hiding this comment.
If the name is empty, do you just attach to all GPU modules?
There was a problem hiding this comment.
I didn't get your question.
filesToLink is files to link, like external spir-v library files. Name would be empty in most cases.
Behavior is the same as nvvm-attach-target.
There was a problem hiding this comment.
If the name is empty, do you just attach to all GPU modules?
I think so.
When moduleMatcher is empty, nvvm info is attached to all GPU module.
There was a problem hiding this comment.
Now after updates, this should be the behavior here as well.
There was a problem hiding this comment.
Both NV and ROC similar methods iterate through regions and basic blocks, but your definition above specified mlir::gpu::GPUModuleOp. Is that because it's assumed that this will only run on the GPU modules (and thus needs no search)? Or are we just not covering the general case in the tests?
There was a problem hiding this comment.
Both NV and ROC also assume that the attachment logic will only be applied to GPU modules:
as indicated by their pass description itself:
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 below.
There was a problem hiding this comment.
Updated logic to match nvvm, rocdl and other attach-target passes.
There was a problem hiding this comment.
I believe no need to use explicit namespaces because they already imported?
There was a problem hiding this comment.
please mention the type and avoid using auto because output type is not obvious in the context (not a cast).
Refer to LLVM coding manuals for auto usage.
There was a problem hiding this comment.
Replaced auto with actual type std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
chencha3
left a comment
chencha3
left a comment
There was a problem hiding this comment.
LGTM. Thanks for your efforts
Co-authored-by: Artem Kroviakov [email protected]
…. Update comment.
silee2
force-pushed
the
addXeVMAttachTarget
branch
from
17f139a to
0728219
Compare
9 participants
Add xevm-attach-target transform pass and unit-tests.
Co-authored-by: Artem Kroviakov [email protected]