[Mlir-commits] [mlir] b2217b3 - [MLIR] Fix affine loop unroll corner case for full unroll

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2021-10-11T10:22:24+05:30
New Revision: b2217b36fe436d84766be0a527a0560feda1a67a

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

LOG: [MLIR] Fix affine loop unroll corner case for full unroll

Fix affine loop unroll for zero trip count loops. Add missing check.

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

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/LoopUtils.cpp
    mlir/test/Dialect/Affine/unroll.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 7a59d505e8a51..7442813255239 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -1052,6 +1052,8 @@ LogicalResult mlir::loopUnrollFull(AffineForOp forOp) {
   Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   if (mayBeConstantTripCount.hasValue()) {
     uint64_t tripCount = mayBeConstantTripCount.getValue();
+    if (tripCount == 0)
+      return success();
     if (tripCount == 1)
       return promoteIfSingleIteration(forOp);
     return loopUnrollByFactor(forOp, tripCount);

diff  --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir
index f78c6c325ea29..2f48ed2a9b4e1 100644
--- a/mlir/test/Dialect/Affine/unroll.mlir
+++ b/mlir/test/Dialect/Affine/unroll.mlir
@@ -641,3 +641,11 @@ func @unroll_with_iter_args_and_promotion(%arg0 : f32, %arg1 : f32) -> f32 {
   // UNROLL-BY-4-NEXT: return %[[RES]]
   return %sum : f32
 }
+
+// UNROLL-FULL: func @unroll_zero_trip_count_case
+func @unroll_zero_trip_count_case() {
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 0
+  affine.for %i = 0 to 0 {
+  }
+  return
+}