address comments

Signed-off-by: Max Dawkins <[email protected]>

Author: Max191

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

@@ -221,24 +221,23 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
         std::optional<SmallVector<OpFoldResult>> ubs =
             loopLike.getLoopUpperBounds();
         // If loop bounds cannot be found, assume possibly zero trip count.
-        if (!lbs || !ubs) {
+        if (!lbs || !ubs)
           return;
-        }
+
         // Otherwise, use ValueBounds to find the maximum lower bound and
         // minimum upper bound. If the bounds are found, and maxLb is less
         // than the minUb, then the loop will not have zero trip count.
         for (auto [lb, ub] : llvm::zip_equal(lbs.value(), ubs.value())) {
           FailureOr<int64_t> maxLb =
               ValueBoundsConstraintSet::computeConstantBound(
-                  presburger::BoundType::UB, /*var=*/lb,
+                  presburger::BoundType::UB, lb,
                   /*stopCondition=*/nullptr, /*closedUB=*/true);
           if (failed(maxLb)) {
             return;
           }
           FailureOr<int64_t> minUb =
               ValueBoundsConstraintSet::computeConstantBound(
-                  presburger::BoundType::LB, /*var=*/ub,
-                  /*stopCondition=*/nullptr);
+                  presburger::BoundType::LB, ub);
           if (failed(minUb)) {
             return;
           }

mlir/test/Dialect/Linalg/hoisting.mlir

@@ -308,20 +308,40 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// CHECK-LABEL:  func.func @no_hoisting_zero_trip_loop
-func.func @no_hoisting_zero_trip_loop(%arg0: memref<20xi32>, %lb: index, %ub: index) {
+// CHECK-LABEL:  func.func @no_hoisting_unknown_bound_loop
+func.func @no_hoisting_unknown_bound_loop(%memref0: memref<20xi32>, %lb: index, %ub: index) {
   %c0_i32 = arith.constant 0 : i32
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
-  // %lb and %ub are unbounded, so do not hoist.
 
+  // %lb and %ub are unbounded, so do not hoist.
   // CHECK:       scf.for {{.*}} {
   // CHECK-NEXT:    vector.transfer_read
-  // CHECK-NEXT:    "prevent.dce"
+  // CHECK-NEXT:    "test.some_use"
   scf.for %arg2 = %lb to %ub step %c1 {
-    %read = vector.transfer_read %arg0[%c0], %c0_i32 {in_bounds = [true]} : memref<20xi32>, vector<4xi32>
-    "prevent.dce"(%read) : (vector<4xi32>) ->()
+    %read = vector.transfer_read %memref0[%c0], %c0_i32 {in_bounds = [true]} : memref<20xi32>, vector<4xi32>
+    "test.some_use"(%read) : (vector<4xi32>) ->()
+  }
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["func.func"]} in %arg1
+      : (!transform.any_op) -> !transform.any_op
+    transform.structured.hoist_redundant_vector_transfers %0 { verify_non_zero_trip }
+      : (!transform.any_op) -> !transform.any_op
+    transform.yield
   }
+}
+
+// -----
+
+// CHECK-LABEL:  func.func @no_hoisting_possibly_zero_trip_loop
+func.func @no_hoisting_possibly_zero_trip_loop(%memref0: memref<20xi32>, %lb: index, %ub: index) {
+  %c0_i32 = arith.constant 0 : i32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
 
   // %lb_0 is in range [%lb, 8], and %ub_0 is in range [4, %ub].
   // Since %lb_0 could be greater than %ub_0, do not hoist.
@@ -330,23 +350,43 @@ func.func @no_hoisting_zero_trip_loop(%arg0: memref<20xi32>, %lb: index, %ub: in
 
   // CHECK:       scf.for {{.*}} {
   // CHECK-NEXT:    vector.transfer_read
-  // CHECK-NEXT:    "prevent.dce"
+  // CHECK-NEXT:    "test.some_use"
   scf.for %arg2 = %lb_0 to %ub_0 step %c1 {
-    %read = vector.transfer_read %arg0[%c0], %c0_i32 {in_bounds = [true]} : memref<20xi32>, vector<4xi32>
-    "prevent.dce"(%read) : (vector<4xi32>) ->()
+    %read = vector.transfer_read %memref0[%c0], %c0_i32 {in_bounds = [true]} : memref<20xi32>, vector<4xi32>
+    "test.some_use"(%read) : (vector<4xi32>) ->()
   }
+  return
+}
 
-  // %lb_1 is in range [%lb, 4], and %ub_1 is in range [8, %ub].
-  // Since %lb_1 is guaranteed to be less than %ub_1, hoisting is possible.
-  %lb_1 = affine.min affine_map<(d0) -> (d0, 4)>(%lb)
-  %ub_1 = affine.max affine_map<(d0) -> (d0, 8)>(%ub)
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["func.func"]} in %arg1
+      : (!transform.any_op) -> !transform.any_op
+    transform.structured.hoist_redundant_vector_transfers %0 { verify_non_zero_trip }
+      : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL:  func.func @hoisting_non_zero_trip_loop
+func.func @hoisting_non_zero_trip_loop(%memref0: memref<20xi32>, %lb: index, %ub: index) {
+  %c0_i32 = arith.constant 0 : i32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+
+  // %lb_0 is in range [%lb, 4], and %ub_0 is in range [8, %ub].
+  // Since %lb_0 is guaranteed to be less than %ub_0, hoisting is possible.
+  %lb_0 = affine.min affine_map<(d0) -> (d0, 4)>(%lb)
+  %ub_0 = affine.max affine_map<(d0) -> (d0, 8)>(%ub)
 
   // CHECK:       vector.transfer_read
   // CHECK:       scf.for {{.*}} {
-  // CHECK-NEXT:    "prevent.dce"
-  scf.for %arg2 = %lb_1 to %ub_1 step %c1 {
-    %read = vector.transfer_read %arg0[%c0], %c0_i32 {in_bounds = [true]} : memref<20xi32>, vector<4xi32>
-    "prevent.dce"(%read) : (vector<4xi32>) ->()
+  // CHECK-NEXT:    "test.some_use"
+  scf.for %arg2 = %lb_0 to %ub_0 step %c1 {
+    %read = vector.transfer_read %memref0[%c0], %c0_i32 {in_bounds = [true]} : memref<20xi32>, vector<4xi32>
+    "test.some_use"(%read) : (vector<4xi32>) ->()
   }
   return
 }
@@ -421,7 +461,7 @@ func.func @no_hoisting_collapse_shape_2(%vec: vector<1x12x1xi32>) {
     %collapse_shape = memref.collapse_shape %alloca [[0, 1, 2]] : memref<1x12x1xi32> into memref<12xi32>
     vector.transfer_write %vec, %alloca[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<1x12x1xi32>, memref<1x12x1xi32>
     %read = vector.transfer_read %collapse_shape[%c0], %c0_i32 {in_bounds = [true]} : memref<12xi32>, vector<12xi32>
-    "prevent.dce"(%read) : (vector<12xi32>) ->()
+    "test.some_use"(%read) : (vector<12xi32>) ->()
   }
   return
 }