[mlir] Fix scf.for single iteration canonicalization check

[antiagainst]

Summary

We should be check whether lb + step >= ub to determine
whether this is a single iteration. Previously we were
checking lb + lb >= ub.

[jpienaar]

Nit: could you add bug description in change description?

[ftynse]

mlir/test/Dialect/Linalg/fusion.mlir

668

Why is this change necessary?

mlir/test/Dialect/SCF/canonicalize.mlir

240

I'd rather have a new test that exercises the previously wrong behavior. This doesn't seem to do so precisely because it is changed to have step value == upper bound value.

mlir/test/mlir-cpu-runner/mlir_test_cblas_interface.cpp

74–92 ↗

(On Diff #319280)

This looks irrelevant

[antiagainst]

Address comments

mlir/test/Dialect/Linalg/fusion.mlir

668

Because otherwise one scf.for loop will be canonicalized away. I think the purpose of this test is to check the two scf.for and linalg.fill/linalg.conv nested inside of them. So without the canonicalization it makes things clearer. Also, most of the tests in this file is performing on unknown dims.

mlir/test/Dialect/SCF/canonicalize.mlir

240

I've created a new test for this and made the step != lb & step != ub. :)

mlir/test/mlir-cpu-runner/mlir_test_cblas_interface.cpp

74–92 ↗

(On Diff #319280)

It was needed because of the scf.for canonicalization kicked in so some shapes are not ? anymore. But Nicolas killed this entire file. :)

[antiagainst]

In D95440#2522701, @jpienaar wrote:

Nit: could you add bug description in change description?

Done.

[antiagainst]

Sorry I accidentally landed this together with a following patch (which actually does not have dependency).. Reverted.

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(ForOp op,
     }
 
     auto step = op.step().getDefiningOp<ConstantOp>();
     if (!step)
       return failure();
 
     // If the loop is known to have 1 iteration, inline its body and remove the
     // loop.
-    llvm::APInt stepValue = lb.getValue().cast<IntegerAttr>().getValue();
+    llvm::APInt stepValue = step.getValue().cast<IntegerAttr>().getValue();
     if ((lbValue + stepValue).sge(ubValue)) {
       SmallVector<Value, 4> blockArgs;
       blockArgs.reserve(op.getNumIterOperands() + 1);
       blockArgs.push_back(op.lowerBound());
       llvm::append_range(blockArgs, op.getIterOperands());
       replaceOpWithRegion(rewriter, op, op.getLoopBody(), blockArgs);
       return success();
     }

mlir/test/Dialect/Linalg/fusion.mlir

 // CHECK:         linalg.yield
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> (d0, d1 - d2)>
 #map1 = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3, s4] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3 * s4)>
 #map2 = affine_map<()[s0] -> (s0 + 3)>
 
-func @fill_and_conv(%arg0: memref<1x4x5x1xf32>, %arg1: memref<2x3x1x1xf32>, %arg2: memref<1x4x5x1xf32>) {
+func @fill_and_conv(%arg0: memref<?x?x?x?xf32>, %arg1: memref<2x3x1x1xf32>, %arg2: memref<?x?x?x?xf32>) {
   %cst = constant 0.000000e+00 : f32
-  linalg.fill(%arg2, %cst) : memref<1x4x5x1xf32>, f32
+  linalg.fill(%arg2, %cst) : memref<?x?x?x?xf32>, f32
 
   %c4 = constant 4 : index
   %c1 = constant 1 : index
   %c0 = constant 0 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %4 = dim %arg1, %c0 : memref<2x3x1x1xf32>
   %5 = dim %arg1, %c1 : memref<2x3x1x1xf32>
-  %6 = dim %arg0, %c0 : memref<1x4x5x1xf32>
-  %7 = dim %arg0, %c1 : memref<1x4x5x1xf32>
-  %8 = dim %arg0, %c3 : memref<1x4x5x1xf32>
-  %9 = dim %arg2, %c0 : memref<1x4x5x1xf32>
-  %10 = dim %arg2, %c1 : memref<1x4x5x1xf32>
-  %11 = dim %arg2, %c2 : memref<1x4x5x1xf32>
-  %12 = dim %arg2, %c3 : memref<1x4x5x1xf32>
+  %6 = dim %arg0, %c0 : memref<?x?x?x?xf32>
+  %7 = dim %arg0, %c1 : memref<?x?x?x?xf32>
+  %8 = dim %arg0, %c3 : memref<?x?x?x?xf32>
+  %9 = dim %arg2, %c0 : memref<?x?x?x?xf32>
+  %10 = dim %arg2, %c1 : memref<?x?x?x?xf32>
+  %11 = dim %arg2, %c2 : memref<?x?x?x?xf32>
+  %12 = dim %arg2, %c3 : memref<?x?x?x?xf32>
   %13 = linalg.range %c0 : %6 : %c2 : !linalg.range
   %14 = linalg.range %c0 : %10 : %c3 : !linalg.range
   scf.for %arg3 = %c0 to %6 step %c2 {
     scf.for %arg4 = %c0 to %10 step %c3 {
       %15 = affine.min #map0(%c2, %c1, %arg3)
       %16 = affine.apply #map2()[%7]
       %17 = affine.min #map0(%16, %c4, %arg4)
-      %18 = dim %arg0, %c2 : memref<1x4x5x1xf32>
-      %19 = dim %arg0, %c3 : memref<1x4x5x1xf32>
-      %20 = subview %arg0[%arg3, %arg4, %c0, %c0] [%15, %17, %18, %19] [%c1, %c1, %c1, %c1] : memref<1x4x5x1xf32> to memref<?x?x?x?xf32, #map1>
+      %18 = dim %arg0, %c2 : memref<?x?x?x?xf32>
+      %19 = dim %arg0, %c3 : memref<?x?x?x?xf32>
+      %20 = subview %arg0[%arg3, %arg4, %c0, %c0] [%15, %17, %18, %19] [%c1, %c1, %c1, %c1] : memref<?x?x?x?xf32> to memref<?x?x?x?xf32, #map1>
       %21 = affine.min #map0(%c2, %c1, %arg3)
       %22 = affine.min #map0(%c3, %c4, %arg4)
-      %23 = dim %arg2, %c2 : memref<1x4x5x1xf32>
-      %24 = dim %arg2, %c3 : memref<1x4x5x1xf32>
-      %25 = subview %arg2[%arg3, %arg4, %c0, %c0] [%21, %22, %23, %24] [%c1, %c1, %c1, %c1] : memref<1x4x5x1xf32> to memref<?x?x?x?xf32, #map1>
+      %23 = dim %arg2, %c2 : memref<?x?x?x?xf32>
+      %24 = dim %arg2, %c3 : memref<?x?x?x?xf32>
+      %25 = subview %arg2[%arg3, %arg4, %c0, %c0] [%21, %22, %23, %24] [%c1, %c1, %c1, %c1] : memref<?x?x?x?xf32> to memref<?x?x?x?xf32, #map1>
       linalg.conv(%arg1, %20, %25) {dilations = [1, 1], strides = [1, 1]} : memref<2x3x1x1xf32>, memref<?x?x?x?xf32, #map1>, memref<?x?x?x?xf32, #map1>
     }
   }
   return
 }
 // CHECK-LABEL: func @fill_and_conv
 // CHECK: scf.for
 // CHECK:   scf.for

mlir/test/Dialect/SCF/canonicalize.mlir

%0 = scf.for %i = %arg0 to %arg0 step %c2 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[INIT]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
-// CHECK-LABEL: @replace_single_iteration_loop
-func @replace_single_iteration_loop() {
+// CHECK-LABEL: @replace_single_iteration_loop_1
+func @replace_single_iteration_loop_1() {
   // CHECK: %[[LB:.*]] = constant 42
   %c42 = constant 42 : index
   %c43 = constant 43 : index
   %c1 = constant 1 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32
   // CHECK-NOT: scf.for
   // CHECK: %[[VAL:.*]] = "test.op"(%[[LB]], %[[INIT]])
   %0 = scf.for %i = %c42 to %c43 step %c1 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }
   // CHECK: "test.consume"(%[[VAL]])
   "test.consume"(%0) : (i32) -> ()
   return
 }
 
+// CHECK-LABEL: @replace_single_iteration_loop_2
+func @replace_single_iteration_loop_2() {
+  // CHECK: %[[LB:.*]] = constant 5
+  %c5 = constant 5 : index
+  %c6 = constant 6 : index
+  %c11 = constant 11 : index
+  // CHECK: %[[INIT:.*]] = "test.init"
+  %init = "test.init"() : () -> i32
+  // CHECK-NOT: scf.for
+  // CHECK: %[[VAL:.*]] = "test.op"(%[[LB]], %[[INIT]])
+  %0 = scf.for %i = %c5 to %c11 step %c6 iter_args(%arg = %init) -> (i32) {
+    %1 = "test.op"(%i, %arg) : (index, i32) -> i32
+    scf.yield %1 : i32
+  }
+  // CHECK: "test.consume"(%[[VAL]])
+  "test.consume"(%0) : (i32) -> ()
+  return
+}
+
+
 // CHECK-LABEL: @replace_single_iteration_loop_non_unit_step
 func @replace_single_iteration_loop_non_unit_step() {
   // CHECK: %[[LB:.*]] = constant 42
   %c42 = constant 42 : index
   %c47 = constant 47 : index
   %c5 = constant 5 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32