[MLIR] Fix affine loop unroll for zero trip count loops

[bondhugula]

Summary

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

[bondhugula]

Update commit title/summary.

[vinayaka-polymage]

Simple clarification.

mlir/lib/Transforms/Utils/LoopUtils.cpp

1056

return failure() here ?, since nothing is done. I see that the return value is being used only at one place, and failure() would be more appropriate.

[bondhugula]

This comment was removed by bondhugula.

mlir/lib/Transforms/Utils/LoopUtils.cpp

1056

The failure() is meant to indicate the utility failed to do the desired transformation (i.e., it could not because of some condition that wasn't supported, handled, or other inability). This isn't really a failure but a trivial success (it has been unrolled trivially so to say).

[dcaballe]

LGTM, thanks!

mlir/lib/Transforms/Utils/LoopUtils.cpp

void mlir::getTileableBands(FuncOp f,
   }
 }
 
 /// Unrolls this loop completely.
 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);
   }
   return failure();
 }
 
 /// Unrolls this loop by the specified factor or by the trip count (if constant)

mlir/test/Dialect/Affine/unroll.mlir

func @unroll_with_iter_args_and_promotion(%arg0 : f32, %arg1 : f32) -> f32 {
   // UNROLL-BY-4-NEXT:   %[[V3:.*]] = addf %[[V2]]
   // UNROLL-BY-4-NEXT:   %[[V4:.*]] = addf %[[V3]]
   // UNROLL-BY-4-NEXT:   affine.yield %[[V4]]
   // UNROLL-BY-4-NEXT: }
   // UNROLL-BY-4-NEXT: %[[RES:.*]] = addf %[[SUM]],
   // 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
+}