Author: lonely eagle
Date: 2025-02-17T00:50:04+08:00
New Revision: 6812fc02fbb81d679f95d5c3e15768ae11e1bad8
URL: https://github.com/llvm/llvm-project/commit/6812fc02fbb81d679f95d5c3e15768ae11e1bad8
DIFF: https://github.com/llvm/llvm-project/commit/6812fc02fbb81d679f95d5c3e15768ae11e1bad8.diff
LOG: [mlir][affine] make affine-loop-unroll-jam interface pass (#127402)
Made affine-loop-unroll-jam an interface pass, so it can run on the gpu
module.
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Passes.h
mlir/include/mlir/Dialect/Affine/Passes.td
mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
mlir/test/Dialect/Affine/unroll-jam.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h
index ea5034b60d8bd..96bd3c6a9a7bc 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.h
+++ b/mlir/include/mlir/Dialect/Affine/Passes.h
@@ -102,7 +102,7 @@ std::unique_ptr<InterfacePass<FunctionOpInterface>> createLoopUnrollPass(
/// Creates a loop unroll jam pass to unroll jam by the specified factor. A
/// factor of -1 lets the pass use the default factor or the one on the command
/// line if provided.
-std::unique_ptr<OperationPass<func::FuncOp>>
+std::unique_ptr<InterfacePass<FunctionOpInterface>>
createLoopUnrollAndJamPass(int unrollJamFactor = -1);
/// Creates a pass to pipeline explicit movement of data across levels of the
diff --git a/mlir/include/mlir/Dialect/Affine/Passes.td b/mlir/include/mlir/Dialect/Affine/Passes.td
index 5325d3b0a1d69..728b8d25efcf2 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Passes.td
@@ -219,7 +219,7 @@ def AffineLoopUnroll : InterfacePass<"affine-loop-unroll", "FunctionOpInterface"
];
}
-def AffineLoopUnrollAndJam : Pass<"affine-loop-unroll-jam", "func::FuncOp"> {
+def AffineLoopUnrollAndJam : InterfacePass<"affine-loop-unroll-jam", "FunctionOpInterface"> {
let summary = "Unroll and jam affine loops";
let constructor = "mlir::affine::createLoopUnrollAndJamPass()";
let options = [
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
index a79160df7575a..13640f085951e 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopUnrollAndJam.cpp
@@ -75,7 +75,7 @@ struct LoopUnrollAndJam
};
} // namespace
-std::unique_ptr<OperationPass<func::FuncOp>>
+std::unique_ptr<InterfacePass<FunctionOpInterface>>
mlir::affine::createLoopUnrollAndJamPass(int unrollJamFactor) {
return std::make_unique<LoopUnrollAndJam>(
unrollJamFactor == -1 ? std::nullopt
diff --git a/mlir/test/Dialect/Affine/unroll-jam.mlir b/mlir/test/Dialect/Affine/unroll-jam.mlir
index 7874580b1c39c..8ed7fccf7d251 100644
--- a/mlir/test/Dialect/Affine/unroll-jam.mlir
+++ b/mlir/test/Dialect/Affine/unroll-jam.mlir
@@ -1,5 +1,6 @@
-// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll-jam="unroll-jam-factor=2" | FileCheck %s
-// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll-jam="unroll-jam-factor=4" | FileCheck --check-prefix=UJAM-FOUR %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=2}))" | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=4}))" | FileCheck --check-prefix=UJAM-FOUR %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(gpu.module(gpu.func(affine-loop-unroll-jam{unroll-jam-factor=2})))" | FileCheck --check-prefix=GPU-HJAM %s
// CHECK-DAG: [[$MAP_PLUS_1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
// CHECK-DAG: [[$MAP_DIV_OFFSET:#map[0-9]*]] = affine_map<()[s0] -> (((s0 - 1) floordiv 2) * 2 + 1)>
@@ -10,6 +11,8 @@
// UJAM-FOUR-DAG: [[$MAP_PLUS_2:#map[0-9]*]] = affine_map<(d0) -> (d0 + 2)>
// UJAM-FOUR-DAG: [[$MAP_PLUS_3:#map[0-9]*]] = affine_map<(d0) -> (d0 + 3)>
+// GPU-HJAM-DAG: [[$MAP_PLUS_1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
+
// CHECK-LABEL: func @unroll_jam_imperfect_nest() {
func.func @unroll_jam_imperfect_nest() {
affine.for %i = 0 to 101 {
@@ -46,6 +49,44 @@ func.func @unroll_jam_imperfect_nest() {
// CHECK-NEXT: "foo"(%c100, %{{.*}})
// CHECK-NEXT: return
+gpu.module @unroll_jam {
+ // GPU-HJAM-LABEL: func @unroll_jam_imperfect_nest() {
+ gpu.func @unroll_jam_imperfect_nest() {
+ affine.for %i = 0 to 101 {
+ %x = "addi32"(%i, %i) : (index, index) -> i32
+ affine.for %j = 0 to 17 {
+ %y = "addi32"(%i, %i) : (index, index) -> i32
+ %z = "addi32"(%y, %y) : (i32, i32) -> i32
+ }
+ %w = "foo"(%i, %x) : (index, i32) -> i32
+ }
+ gpu.return
+ }
+ // GPU-HJAM: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {
+ // GPU-HJAM-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])
+ // GPU-HJAM-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])
+ // GPU-HJAM-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])
+ // GPU-HJAM-NEXT: affine.for %{{.*}} = 0 to 17 {
+ // GPU-HJAM-NEXT: [[RES3:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])
+ // GPU-HJAM-NEXT: "addi32"([[RES3]], [[RES3]]) : (i32, i32) -> i32
+ // GPU-HJAM-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])
+ // GPU-HJAM-NEXT: [[RES4:%[0-9]+]] = "addi32"([[INC1]], [[INC1]])
+ // GPU-HJAM-NEXT: "addi32"([[RES4]], [[RES4]]) : (i32, i32) -> i32
+ // GPU-HJAM-NEXT: }
+ // GPU-HJAM: "foo"([[IV0]], [[RES1]])
+ // GPU-HJAM-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])
+ // GPU-HJAM-NEXT: "foo"({{.*}}, [[RES2]])
+ // GPU-HJAM: }
+ // Cleanup loop (single iteration).
+ // GPU-HJAM: "addi32"(%c100, %c100)
+ // GPU-HJAM-NEXT: affine.for [[IV0]] = 0 to 17 {
+ // GPU-HJAM-NEXT: [[RESC:%[0-9]+]] = "addi32"(%c100, %c100)
+ // GPU-HJAM-NEXT: "addi32"([[RESC]], [[RESC]]) : (i32, i32) -> i32
+ // GPU-HJAM-NEXT: }
+ // GPU-HJAM-NEXT: "foo"(%c100, %{{.*}})
+ // GPU-HJAM-NEXT: return
+}
+
// CHECK-LABEL: func @loop_nest_unknown_count_1
// CHECK-SAME: [[N:arg[0-9]+]]: index
func.func @loop_nest_unknown_count_1(%N : index) {