[Mlir-commits] [mlir] 0f6a2cd - [mlir] Fix signed ceildiv, loop normalization.

[Hendrik Greving]

Author: Hendrik Greving
Date: 2022-08-30T09:55:04-07:00
New Revision: 0f6a2cd2aba6a115cf468c0de648f35d6a689836

URL: https://github.com/llvm/llvm-project/commit/0f6a2cd2aba6a115cf468c0de648f35d6a689836
DIFF: https://github.com/llvm/llvm-project/commit/0f6a2cd2aba6a115cf468c0de648f35d6a689836.diff

LOG: [mlir] Fix signed ceildiv, loop normalization.

Fixes using the signed ceildiv op instead of incorrectly assuming positive loop bounds.
Adjusts the tests for above.

Differential Revision: https://reviews.llvm.org/D132953

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Utils/Utils.cpp
    mlir/test/Dialect/Affine/loop-coalescing.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 1bcc48a32d7b..24b92d698bd5 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -555,14 +555,13 @@ static LoopParams normalizeLoop(OpBuilder &boundsBuilder,
   // Compute the number of iterations the loop executes: ceildiv(ub - lb, step)
   // assuming the step is strictly positive.  Update the bounds and the step
   // of the loop to go from 0 to the number of iterations, if necessary.
-  // TODO: introduce support for negative steps or emit dynamic asserts
-  // on step positivity, whatever gets implemented first.
   if (isZeroBased && isStepOne)
     return {/*lowerBound=*/lowerBound, /*upperBound=*/upperBound,
             /*step=*/step};
 
   Value 
diff  = boundsBuilder.create<arith::SubIOp>(loc, upperBound, lowerBound);
-  Value newUpperBound = ceilDivPositive(boundsBuilder, loc, 
diff , step);
+  Value newUpperBound =
+      boundsBuilder.create<arith::CeilDivSIOp>(loc, 
diff , step);
 
   Value newLowerBound =
       isZeroBased ? lowerBound

diff  --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index 5abc59734ddf..9c17fb24be69 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -88,10 +88,7 @@ func.func @unnormalized_loops() {
 
   // Number of iterations in the outer scf.
   // CHECK: %[[
diff _i:.*]] = arith.subi %[[orig_ub_i]], %[[orig_lb_i]]
-  // CHECK: %[[c1:.*]] = arith.constant 1
-  // CHECK: %[[step_minus_c1:.*]] = arith.subi %[[orig_step_i]], %[[c1]]
-  // CHECK: %[[dividend:.*]] = arith.addi %[[
diff _i]], %[[step_minus_c1]]
-  // CHECK: %[[numiter_i:.*]] = arith.divui %[[dividend]], %[[orig_step_i]]
+  // CHECK: %[[numiter_i:.*]] = arith.ceildivsi %[[
diff _i]], %[[orig_step_i]]
 
   // Normalized lower bound and step for the outer scf.
   // CHECK: %[[lb_i:.*]] = arith.constant 0
@@ -99,7 +96,7 @@ func.func @unnormalized_loops() {
 
   // Number of iterations in the inner loop, the pattern is the same as above,
   // only capture the final result.
-  // CHECK: %[[numiter_j:.*]] = arith.divui {{.*}}, %[[orig_step_j]]
+  // CHECK: %[[numiter_j:.*]] = arith.ceildivsi {{.*}}, %[[orig_step_j]]
 
   // New bounds of the outer scf.
   // CHECK: %[[range:.*]] = arith.muli %[[numiter_i]], %[[numiter_j]]
@@ -135,13 +132,9 @@ func.func @parametric(%lb1 : index, %ub1 : index, %step1 : index,
   // Compute the number of iterations for each of the loops and the total
   // number of iterations.
   // CHECK: %[[range1:.*]] = arith.subi %[[orig_ub1]], %[[orig_lb1]]
-  // CHECK: %[[orig_step1_minus_1:.*]] = arith.subi %[[orig_step1]], %c1
-  // CHECK: %[[dividend1:.*]] = arith.addi %[[range1]], %[[orig_step1_minus_1]]
-  // CHECK: %[[numiter1:.*]] = arith.divui %[[dividend1]], %[[orig_step1]]
+  // CHECK: %[[numiter1:.*]] = arith.ceildivsi %[[range1]], %[[orig_step1]]
   // CHECK: %[[range2:.*]] = arith.subi %[[orig_ub2]], %[[orig_lb2]]
-  // CHECK: %[[orig_step2_minus_1:.*]] = arith.subi %arg5, %c1
-  // CHECK: %[[dividend2:.*]] = arith.addi %[[range2]], %[[orig_step2_minus_1]]
-  // CHECK: %[[numiter2:.*]] = arith.divui %[[dividend2]], %[[orig_step2]]
+  // CHECK: %[[numiter2:.*]] = arith.ceildivsi %[[range2]], %[[orig_step2]]
   // CHECK: %[[range:.*]] = arith.muli %[[numiter1]], %[[numiter2]] : index
 
   // Check that the outer loop is updated.