[MLIR] NFC clean up loop tiling test cases

[bondhugula]

Summary

Improve the checks and drop redundant ones. Rename the outlined maps
better.

mlir/test/Dialect/Affine/loop-tiling.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-loop-tile="tile-size=32" | FileCheck %s
-// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-loop-tile="cache-size=512" | FileCheck %s --check-prefix=MODEL
-// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-loop-tile="tile-size=32 separate" | FileCheck %s --check-prefix=SEPARATE
+// RUN: mlir-opt %s -split-input-file -affine-loop-tile="tile-size=32" | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -affine-loop-tile="cache-size=512" | FileCheck %s --check-prefix=MODEL
+// RUN: mlir-opt %s -split-input-file -affine-loop-tile="tile-size=32 separate" | FileCheck %s --check-prefix=SEPARATE
 
 // -----
 
-// CHECK-DAG: [[MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32)>
-// CHECK-DAG: [[MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32, 50)>
+// CHECK-DAG: [[UB:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32)>
+// CHECK-DAG: [[UB_MIN:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32, 50)>
 // CHECK-DAG: [[ID:#map[0-9]+]] = affine_map<(d0) -> (d0)>
 // CHECK-DAG: [[ID_PLUS_21:#map[0-9]+]] = affine_map<(d0) -> (d0 + 21)>
 
 // CHECK-LABEL: func @loop_tiling()
 // CHECK-NEXT:   affine.for %{{.*}} = 0 to 256 step 32 {
 // CHECK-NEXT:     affine.for %{{.*}} = 0 to 512 step 32 {
 // CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 step 32 {
-// CHECK-NEXT:         affine.for %{{.*}} = [[ID]](%{{.*}}) to [[MAP0]](%{{.*}}) {
-// CHECK-NEXT:           affine.for %{{.*}} = [[ID]](%{{.*}}) to [[MAP0]](%{{.*}}) {
-// CHECK-NEXT:             affine.for %{{.*}} = [[ID]](%{{.*}}) to [[MAP0]](%{{.*}}) {
-// CHECK-NEXT:               "foo"(%{{.*}}, %{{.*}}, %{{.*}}) : (index, index, index) -> ()
+// CHECK-NEXT:         affine.for %[[I:.*]] = [[ID]](%{{.*}}) to [[UB]](%{{.*}}) {
+// CHECK-NEXT:           affine.for %[[J:.*]] = [[ID]](%{{.*}}) to [[UB]](%{{.*}}) {
+// CHECK-NEXT:             affine.for %[[K:.*]] = [[ID]](%{{.*}}) to [[UB]](%{{.*}}) {
+// CHECK-NEXT:               "test.foo"(%[[I]], %[[J]], %[[K]])
 // CHECK-NEXT:             }
 // CHECK-NEXT:           }
 // CHECK-NEXT:         }
 // CHECK-NEXT:       }
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
 // CHECK-NEXT:   affine.for %{{.*}} = 0 to 50 step 32 {
-// CHECK-NEXT:     affine.for %{{.*}} = [[ID]](%{{.*}}) to min [[MAP1]](%{{.*}}) {
-// CHECK-NEXT:       "bar"(%{{.*}}, %{{.*}}) : (index, index) -> ()
+// CHECK-NEXT:     affine.for %[[X:.*]] = [[ID]](%{{.*}}) to min [[UB_MIN]](%{{.*}}) {
+// CHECK-NEXT:       "test.bar"(%[[X]], %[[X]])
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
 // CHECK-NEXT: affine.for %[[I:.*]] = 0 to 21 step 32 {
-// CHECK-NEXT:    affine.for %{{.*}} = [[ID]](%[[I]]) to [[ID_PLUS_21]](%[[I]]) {
-// CHECK-NEXT:      "foobar"(%{{.*}}) : (index) -> ()
-// CHECK-NEXT:    }
-// CHECK-NEXT:  }
+// CHECK-NEXT:   affine.for %[[Y:.*]] = [[ID]](%[[I]]) to [[ID_PLUS_21]](%[[I]])  {
+// CHECK-NEXT:     "test.foobar"(%[[Y]])
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
 // CHECK-NEXT:  return
 func @loop_tiling() {
   affine.for %i = 0 to 256 {
     affine.for %j = 0 to 512 {
       affine.for %k = 0 to 1024 {
-        "foo"(%i, %j, %k) : (index, index, index) -> ()
+        "test.foo"(%i, %j, %k) : (index, index, index) -> ()
       }
     }
   }
 
   affine.for %x = 0 to 50 {
-    "bar"(%x, %x) : (index, index) -> ()
+    "test.bar"(%x, %x) : (index, index) -> ()
   }
 
   // Intra-tile loop won't need a min expression.
   affine.for %y = 0 to 21 {
-    "foobar"(%y) : (index) -> ()
+    "test.foobar"(%y) : (index) -> ()
   }
 
   return
 }
 
 // -----
 
 // CHECK-DAG: [[IDENTITY:#map[0-9]+]] = affine_map<(d0) -> (d0)>
 // CHECK-DAG: [[LB:#map[0-9]+]] = affine_map<()[s0] -> (0, s0)>
 // CHECK-DAG: [[UB:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>
 // CHECK-DAG: [[UB_INTRA_TILE:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> (d0 + 32, s0, 4096 floordiv s1)>
 
 #lb = affine_map<()[s0] -> (0, s0)>
 #ub = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>
 // CHECK-LABEL: func @loop_max_min_bound(%{{.*}}: memref<?xi32>, %{{.*}}: index, %{{.*}}: index) {
 func @loop_max_min_bound(%A : memref<? x i32>, %L : index, %U : index) {
   %M = dim %A, 0 : memref<? x i32>
-  affine.for %iTT = max #lb()[%L] to min #ub()[%M, %U] {
-      %out = affine.apply affine_map<(d0) -> (d0)> (%iTT)
+  affine.for %i = max #lb()[%L] to min #ub()[%M, %U] {
+    addi %i, %i : index
   }
   return
 // CHECK:       affine.for %{{.*}} = max [[LB]]()[%{{.*}}] to min [[UB]]()[%{{.*}}, %{{.*}}] step 32 {
-// CHECK-NEXT:    affine.for %{{.*}} = [[IDENTITY]](%{{.*}}) to min [[UB_INTRA_TILE]](%{{.*}})[%{{.*}}, %{{.*}}] {
-// CHECK-NEXT:      %{{.*}} = affine.apply [[IDENTITY]](%{{.*}})
+// CHECK-NEXT:    affine.for %[[I:.*]] = [[IDENTITY]](%{{.*}}) to min [[UB_INTRA_TILE]](%{{.*}})[%{{.*}}, %{{.*}}] {
+// CHECK-NEXT:      addi %[[I]], %[[I]]
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 }
 
 // -----
 
 // Cache size is set to 512 KiB. This loop nest accesses about 49 MiB, and the
 // tile sizes chosen would be 6 x 6 x 6. However, to avoid min/max, which is
affine.for %i1 = 0 to %0 {
         %5 = addf %4, %3 : f32
         affine.store %5, %arg2[%i0, %i1] : memref<?x?xf32>
       }
     }
   }
   return
 }
 
-// CHECK:       %{{.*}} = dim %{{.*}}, 0 : memref<?x?xf32>
+// CHECK:       dim %{{.*}}, 0 : memref<?x?xf32>
 // CHECK-NEXT:  affine.for %{{.*}} = 0 to %{{.*}} step 32 {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to %{{.*}} step 32 {
 // CHECK-NEXT:      affine.for %{{.*}} = #map3(%{{.*}}) to min [[UBMAP]](%{{.*}})[%{{.*}}] {
 // CHECK-NEXT:        affine.for %{{.*}} = #map3(%{{.*}}) to min [[UBMAP]](%{{.*}})[%{{.*}}] {
 // CHECK-NEXT:          affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
 // CHECK-NEXT:          affine.for %{{.*}} = 0 to %{{.*}} {
-// CHECK-NEXT:            %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
-// CHECK-NEXT:            %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
-// CHECK-NEXT:            %{{.*}} = mulf %{{.*}}, %{{.*}} : f32
-// CHECK-NEXT:            %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
-// CHECK-NEXT:            %{{.*}} = addf %{{.*}}, %{{.*}} : f32
-// CHECK-NEXT:            affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
+// CHECK-NEXT:            affine.load
+// CHECK-NEXT:            affine.load
+// CHECK-NEXT:            mulf
+// CHECK-NEXT:            affine.load
+// CHECK-NEXT:            addf
+// CHECK-NEXT:            affine.store
 // CHECK-NEXT:          }
 // CHECK-NEXT:        }
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 
 // -----
 
 // CHECK-DAG: [[MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0)>
 // CHECK-DAG: [[MAP1:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0 + s1)>
 // CHECK-DAG: [[UBMAP:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> (d0 + 32, s0 + s1)>
 
 func @tile_with_loop_upper_bounds_in_two_symbols(%arg0: memref<?xf32>, %limit: index) {
   %dim0 = dim %arg0, 0 : memref<?xf32>
   affine.for %i0 = 0 to affine_map<()[s0, s1] -> (s0 + s1)> ()[%dim0, %limit] {
     %v0 = affine.load %arg0[%i0] : memref<?xf32>
   }
   return
 }
 
-// CHECK:       %{{.*}} = dim %{{.*}}, 0 : memref<?xf32>
+// CHECK:       dim %{{.*}}, 0 : memref<?xf32>
 // CHECK-NEXT:  affine.for %{{.*}} = 0 to [[MAP1]]()[%{{.*}}, %{{.*}}] step 32 {
 // CHECK-NEXT:    affine.for %{{.*}} = [[MAP0]](%{{.*}}) to min [[UBMAP]](%{{.*}})[%{{.*}}, %{{.*}}] {
-// CHECK-NEXT:      %{{.*}} = affine.load %{{.*}}[%{{.*}}] : memref<?xf32>
+// CHECK-NEXT:      affine.load
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 
 // -----
 
 func @tile_size_larger_than_trip_count_symbolic_bound(%M: index, %N :  index) {
   affine.for %i = affine_map<(d0) -> (d0)>(%M) to affine_map<(d0) -> (d0 + 2)>(%M) {
     affine.for %j = affine_map<(d0) -> (d0)>(%N) to affine_map<(d0) -> (d0 + 4)>(%N) {

 // CHECK:      affine.for %[[I:.*]] = #[[ID]](%[[M]]) to #[[ID_PLUS_2]](%[[M]]) step 32
 // CHECK-NEXT:   affine.for %[[J:.*]] = #[[ID]](%[[N]]) to #[[ID_PLUS_4]](%[[N]]) step 32
 // CHECK-NEXT:     affine.for %arg4 = #[[ID]](%[[I]]) to #[[ID_PLUS_2]](%[[I]])
 // CHECK-NEXT:       affine.for %arg5 = #[[ID]](%[[J]]) to #[[ID_PLUS_4]](%[[J]])
 // CHECK-NEXT:         "test.foo"
 
 // -----
 
-// CHECK-LABEL: func @trip_count_1
-// SEPARATE-LABEL: func @trip_count_1
-func @trip_count_1(%arg0: memref<196608x1xf32>, %arg1: memref<196608x1xf32>)
+// CHECK-LABEL: func @trip_count_one
+// SEPARATE-LABEL: func @trip_count_one
+func @trip_count_one(%arg0: memref<196608x1xf32>, %arg1: memref<196608x1xf32>)
     -> memref<196608x1xf32> {
   affine.for %i1 = 0 to 196608 {
     affine.for %i3 = 0 to 1 {
       %4 = affine.load %arg0[%i1, %i3] : memref<196608x1xf32>
       affine.store %4, %arg1[%i1, %i3] : memref<196608x1xf32>
     }
   }
   // CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<196608x1xf32>
   return %arg1 : memref<196608x1xf32>
 }
+// To make sure SEPRATE-DAGs further below do not match with something above.
 // SEPARATE: return
 
 // -----
 
 func @separate_full_tile_2d(%M : index, %N : index) {
-  affine.for %i0 = 0 to %M {
-    affine.for %i1 = 0 to %N {
-      "foo"() : () -> ()
+  affine.for %i = 0 to %M {
+    affine.for %j = 0 to %N {
+      "test.foo"() : () -> ()
     }
   }
   return
 }
 
 // SEPARATE-DAG: #[[SEP_COND:.*]] = affine_set<(d0, d1)[s0, s1] : (-d0 + s0 - 32 >= 0, -d1 + s1 - 32 >= 0)>
 // SEPARATE-DAG: #[[LB:.*]] = affine_map<(d0) -> (d0)>
 // SEPARATE-DAG: #[[FULL_TILE_UB:.*]] = affine_map<(d0) -> (d0 + 32)>
 // SEPARATE-DAG: #[[PART_TILE_UB:.*]] = affine_map<(d0)[s0] -> (d0 + 32, s0)>
 
-// SEPARATE:       affine.for %arg2
-// SEPARATE-NEXT:    affine.for %arg3
+// SEPARATE-LABEL: func @separate_full_tile_2d(
+// SEPARATE: %[[M:.*]]: index, %[[N:.*]]: index
+
+// SEPARATE:       affine.for %[[I:.*]] =
+// SEPARATE-NEXT:    affine.for %[[J:.*]] =
 // SEPARATE-NEXT:      affine.if #[[SEP_COND]](%arg2, %arg3)[%arg0, %arg1] {
-// SEPARATE-NEXT:        affine.for %arg4 = #[[LB]](%arg2) to #[[FULL_TILE_UB]](%arg2) {
-// SEPARATE-NEXT:          affine.for %arg5 = #[[LB]](%arg3) to #[[FULL_TILE_UB]](%arg3) {
-// SEPARATE-NEXT:           "foo"
+// SEPARATE-NEXT:        affine.for %{{.*}} = #[[LB]](%[[I]]) to #[[FULL_TILE_UB]](%[[I]]) {
+// SEPARATE-NEXT:          affine.for %{{.*}} = #[[LB]](%[[J]]) to #[[FULL_TILE_UB]](%[[J]]) {
+// SEPARATE-NEXT:           "test.foo"
 // SEPARATE-NEXT:          }
 // SEPARATE-NEXT:        }
 // SEPARATE-NEXT:      } else {
-// SEPARATE-NEXT:        affine.for %arg4 = #[[LB]](%arg2) to min #[[PART_TILE_UB]](%arg2)[%arg0] {
-// SEPARATE-NEXT:          affine.for %arg5 = #[[LB]](%arg3) to min #[[PART_TILE_UB]](%arg3)[%arg1] {
-// SEPARATE-NEXT:           "foo"
+// SEPARATE-NEXT:        affine.for %{{.*}} = #[[LB]](%[[I]]) to min #[[PART_TILE_UB]](%[[I]])[%[[M]]] {
+// SEPARATE-NEXT:          affine.for %{{.*}} = #[[LB]](%[[J]]) to min #[[PART_TILE_UB]](%[[J]])[%[[N]]] {
+// SEPARATE-NEXT:           "test.foo"
 // SEPARATE-NEXT:          }
 // SEPARATE-NEXT:        }
 // SEPARATE-NEXT:      }
 // SEPARATE-NEXT:    }
 // SEPARATE-NEXT:  }
 // SEPARATE-NEXT:  return
 
 // -----