[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.

Author: linuxlonelyeagle

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

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 = [

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

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) {