unroll and jam: fix order of jammed bodies

- bodies would earlier appear in the order (i, i+3, i+2, i+1) instead of
  (i, i+1, i+2, i+3) for example for factor 4.

- clean up hardcoded test cases

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes #170

COPYBARA_INTEGRATE_REVIEW=tensorflow/mlir#170 from bondhugula:ujam b66b405b2b1894a03b376952e32a9d0292042665
PiperOrigin-RevId: 273613131

Author: bondhugula

lib/Transforms/LoopUnrollAndJam.cpp

@@ -209,8 +209,8 @@ LogicalResult mlir::loopUnrollJamByFactor(AffineForOp forOp,
   forOp.setStep(step * unrollJamFactor);
 
   auto *forOpIV = forOp.getInductionVar();
-  // Unroll and jam (appends unrollJamFactor-1 additional copies).
-  for (unsigned i = 1; i < unrollJamFactor; i++) {
+  // Unroll and jam (appends unrollJamFactor - 1 additional copies).
+  for (unsigned i = unrollJamFactor - 1; i >= 1; --i) {
     // Operand map persists across all sub-blocks.
     BlockAndValueMapping operandMapping;
     for (auto &subBlock : subBlocks) {

test/Transforms/unroll-jam.mlir

@@ -1,9 +1,15 @@
 // RUN: mlir-opt %s -affine-loop-unroll-jam -unroll-jam-factor=2 | FileCheck %s
+// RUN: mlir-opt %s -affine-loop-unroll-jam -unroll-jam-factor=4 | FileCheck --check-prefix=UJAM-FOUR %s
 
 // CHECK-DAG: [[MAP_PLUS_1:#map[0-9]+]] = (d0) -> (d0 + 1)
-// CHECK-DAG: [[M1:#map[0-9]+]] = ()[s0] -> (s0 + 8)
 // CHECK-DAG: [[MAP_DIV_OFFSET:#map[0-9]+]] = ()[s0] -> (((s0 - 1) floordiv 2) * 2 + 1)
 // CHECK-DAG: [[MAP_MULTI_RES:#map[0-9]+]] = ()[s0, s1] -> ((s0 floordiv 2) * 2, (s1 floordiv 2) * 2, 1024)
+// CHECK-DAG: [[MAP_SYM_UB:#map[0-9]+]] = ()[s0, s1] -> (s0, s1, 1024)
+
+// UJAM-FOUR-DAG: [[UBMAP:#map[0-9]+]] = ()[s0] -> (s0 + 8)
+// UJAM-FOUR-DAG: [[MAP_PLUS_1:#map[0-9]+]] = (d0) -> (d0 + 1)
+// UJAM-FOUR-DAG: [[MAP_PLUS_2:#map[0-9]+]] = (d0) -> (d0 + 2)
+// UJAM-FOUR-DAG: [[MAP_PLUS_3:#map[0-9]+]] = (d0) -> (d0 + 3)
 
 // CHECK-LABEL: func @unroll_jam_imperfect_nest() {
 func @unroll_jam_imperfect_nest() {
@@ -29,33 +35,33 @@ func @unroll_jam_imperfect_nest() {
 // CHECK-NEXT:     "addi32"([[RES4]], [[RES4]]) : (i32, i32) -> i32
 // CHECK-NEXT:   }
 // CHECK:        "foo"([[IV0]], [[RES1]])
-// CHECK-NEXT:   {{.*}} = affine.apply [[MAP_PLUS_1]]([[IV0]])
+// CHECK-NEXT:   affine.apply [[MAP_PLUS_1]]([[IV0]])
 // CHECK-NEXT:   "foo"({{.*}}, [[RES2]])
 // CHECK:      }
 // Cleanup loop (single iteration).
-// CHECK:      %{{.*}} = "addi32"(%c100, %c100)
+// CHECK:      "addi32"(%c100, %c100)
 // CHECK-NEXT: affine.for [[IV0]] = 0 to 17 {
 // CHECK-NEXT:   [[RESC:%[0-9]+]] = "addi32"(%c100, %c100)
 // CHECK-NEXT:   "addi32"([[RESC]], [[RESC]]) : (i32, i32) -> i32
 // CHECK-NEXT: }
-// CHECK-NEXT: %{{.*}} = "foo"(%c100, %{{.*}})
+// CHECK-NEXT: "foo"(%c100, %{{.*}})
 // CHECK-NEXT: return
 
 // CHECK-LABEL: func @loop_nest_unknown_count_1
 // CHECK-SAME: [[N:arg[0-9]+]]: index
 func @loop_nest_unknown_count_1(%N : index) {
   // CHECK-NEXT: affine.for %{{.*}} = 1 to [[MAP_DIV_OFFSET]]()[%[[N]]] step 2 {
   // CHECK-NEXT:   affine.for %{{.*}} = 1 to 100 {
-  // CHECK-NEXT:     %{{.*}} = "foo"() : () -> i32
-  // CHECK-NEXT:     %{{.*}} = "foo"() : () -> i32
+  // CHECK-NEXT:     "foo"() : () -> i32
+  // CHECK-NEXT:     "foo"() : () -> i32
   // CHECK-NEXT:   }
   // CHECK-NEXT: }
   // A cleanup loop should be generated here.
   // CHECK-NEXT: affine.for %{{.*}} = [[MAP_DIV_OFFSET]]()[%[[N]]] to %[[N]] {
   // CHECK-NEXT:   affine.for %{{.*}} = 1 to 100 {
   // CHECK-NEXT:     "foo"() : () -> i32
-  // CHECK_NEXT:   }
-  // CHECK_NEXT: }
+  // CHECK-NEXT:   }
+  // CHECK-NEXT: }
   affine.for %i = 1 to %N {
     affine.for %j = 1 to 100 {
       %x = "foo"() : () -> i32
@@ -64,24 +70,28 @@ func @loop_nest_unknown_count_1(%N : index) {
   return
 }
 
-// CHECK-LABEL: func @loop_nest_unknown_count_2
-// CHECK-SAME: %[[N:arg[0-9]+]]: index
+// UJAM-FOUR-LABEL: func @loop_nest_unknown_count_2
+// UJAM-FOUR-SAME: %[[N:arg[0-9]+]]: index
 func @loop_nest_unknown_count_2(%N : index) {
-  // CHECK-NEXT: affine.for [[IV0:%arg[0-9]+]] = %[[N]] to  [[M1]]()[%[[N]]] step 2 {
-  // CHECK-NEXT:   affine.for [[IV1:%arg[0-9]+]] = 1 to 100 {
-  // CHECK-NEXT:     "foo"([[IV0]]) : (index) -> i32
-  // CHECK-NEXT:     [[RES:%[0-9]+]] = affine.apply #map{{[0-9]+}}([[IV0]])
-  // CHECK-NEXT:     "foo"([[RES]])
-  // CHECK-NEXT:   }
-  // CHECK-NEXT: }
+  // UJAM-FOUR-NEXT: affine.for [[IV0:%arg[0-9]+]] = %[[N]] to  [[UBMAP]]()[%[[N]]] step 4 {
+  // UJAM-FOUR-NEXT:   affine.for [[IV1:%arg[0-9]+]] = 1 to 100 {
+  // UJAM-FOUR-NEXT:     "foo"([[IV0]])
+  // UJAM-FOUR-NEXT:     [[IV_PLUS_1:%[0-9]+]] = affine.apply [[MAP_PLUS_1]]([[IV0]])
+  // UJAM-FOUR-NEXT:     "foo"([[IV_PLUS_1]])
+  // UJAM-FOUR-NEXT:     [[IV_PLUS_2:%[0-9]+]] = affine.apply [[MAP_PLUS_2]]([[IV0]])
+  // UJAM-FOUR-NEXT:     "foo"([[IV_PLUS_2]])
+  // UJAM-FOUR-NEXT:     [[IV_PLUS_3:%[0-9]+]] = affine.apply [[MAP_PLUS_3]]([[IV0]])
+  // UJAM-FOUR-NEXT:     "foo"([[IV_PLUS_3]])
+  // UJAM-FOUR-NEXT:   }
+  // UJAM-FOUR-NEXT: }
   // The cleanup loop is a single iteration one and is promoted.
-  // CHECK-NEXT: [[RES:%[0-9]+]] = affine.apply [[M1]]()[%[[N]]]
-  // CHECK-NEXT: affine.for [[IV0]] = 1 to 100 {
-  // CHECK-NEXT:   "foo"([[RES]])
-  // CHECK_NEXT: }
+  // UJAM-FOUR-NEXT: [[RES:%[0-9]+]] = affine.apply [[UBMAP]]()[%[[N]]]
+  // UJAM-FOUR-NEXT: affine.for [[IV0]] = 1 to 100 {
+  // UJAM-FOUR-NEXT:   "foo"([[RES]])
+  // UJAM-FOUR-NEXT: }
   affine.for %i = %N to ()[s0] -> (s0+9) ()[%N] {
     affine.for %j = 1 to 100 {
-      %x = "foo"(%i) : (index) -> i32
+      "foo"(%i) : (index) -> ()
     }
   }
   return
@@ -102,11 +112,11 @@ func @loop_nest_symbolic_and_min_upper_bound(%M : index, %N : index, %K : index)
 // CHECK-NEXT:  affine.for [[IV0:%arg[0-9]+]] = 0 to min [[MAP_MULTI_RES]]()[%[[M]], %[[N]]] step 2 {
 // CHECK-NEXT:    affine.for [[IV1:%arg[0-9]+]] = 0 to %[[K]] {
 // CHECK-NEXT:      "foo"([[IV0]], [[IV1]])
-// CHECK-NEXT:      [[RES:%[0-9]+]] = affine.apply #map0([[IV0]])
+// CHECK-NEXT:      [[RES:%[0-9]+]] = affine.apply [[MAP_PLUS_1]]([[IV0]])
 // CHECK-NEXT:      "foo"([[RES]], [[IV1]])
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
-// CHECK-NEXT:  affine.for [[IV0]] = max [[MAP_MULTI_RES]]()[%[[M]], %[[N]]] to min #map9()[%[[M]], %[[N]]] {
+// CHECK-NEXT:  affine.for [[IV0]] = max [[MAP_MULTI_RES]]()[%[[M]], %[[N]]] to min [[MAP_SYM_UB]]()[%[[M]], %[[N]]] {
 // CHECK-NEXT:    affine.for [[IV1]] = 0 to %[[K]] {
 // CHECK-NEXT:      "foo"([[IV0]], [[IV1]])
 // CHECK-NEXT:    }