[MLIR][SCF] Fix normalizeForallOp helper function

Previously the `normalizeForallOp` function did not work properly, since the
newly created op was not being returned in addition to the op failing
verification.

This patch fixes the helper function and adds a unit test for it.

Author: Cubevoid

mlir/lib/Dialect/SCF/Utils/Utils.cpp

@@ -1482,30 +1482,31 @@ FailureOr<scf::ForallOp> mlir::normalizeForallOp(RewriterBase &rewriter,
   SmallVector<OpFoldResult> ubs = forallOp.getMixedUpperBound();
   SmallVector<OpFoldResult> steps = forallOp.getMixedStep();
 
-  if (llvm::all_of(
-          lbs, [](OpFoldResult ofr) { return isConstantIntValue(ofr, 0); }) &&
-      llvm::all_of(
-          steps, [](OpFoldResult ofr) { return isConstantIntValue(ofr, 1); })) {
+  if (forallOp.isNormalized())
     return forallOp;
-  }
 
-  SmallVector<OpFoldResult> newLbs, newUbs, newSteps;
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(forallOp);
+  SmallVector<OpFoldResult> newUbs;
   for (auto [lb, ub, step] : llvm::zip_equal(lbs, ubs, steps)) {
     Range normalizedLoopParams =
         emitNormalizedLoopBounds(rewriter, forallOp.getLoc(), lb, ub, step);
-    newLbs.push_back(normalizedLoopParams.offset);
     newUbs.push_back(normalizedLoopParams.size);
-    newSteps.push_back(normalizedLoopParams.stride);
   }
+  (void)foldDynamicIndexList(newUbs);
 
+  // Use the normalized builder since the lower bounds are always 0 and the
+  // steps are always 1.
   auto normalizedForallOp = rewriter.create<scf::ForallOp>(
-      forallOp.getLoc(), newLbs, newUbs, newSteps, forallOp.getOutputs(),
-      forallOp.getMapping(), [](OpBuilder &, Location, ValueRange) {});
+      forallOp.getLoc(), newUbs, forallOp.getOutputs(), forallOp.getMapping(),
+      [](OpBuilder &, Location, ValueRange) {});
 
   rewriter.inlineRegionBefore(forallOp.getBodyRegion(),
                               normalizedForallOp.getBodyRegion(),
                               normalizedForallOp.getBodyRegion().begin());
+  // Remove the original empty block in the new loop.
+  rewriter.eraseBlock(&normalizedForallOp.getBodyRegion().back());
 
-  rewriter.replaceAllOpUsesWith(forallOp, normalizedForallOp);
-  return success();
+  rewriter.replaceOp(forallOp, normalizedForallOp);
+  return normalizedForallOp;
 }

mlir/unittests/Dialect/SCF/CMakeLists.txt

@@ -5,4 +5,5 @@ mlir_target_link_libraries(MLIRSCFTests
   PRIVATE
   MLIRIR
   MLIRSCFDialect
+  MLIRSCFUtils
 )

mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp

@@ -6,11 +6,15 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
@@ -23,7 +27,7 @@ using namespace mlir::scf;
 class SCFLoopLikeTest : public ::testing::Test {
 protected:
   SCFLoopLikeTest() : b(&context), loc(UnknownLoc::get(&context)) {
-    context.loadDialect<arith::ArithDialect, scf::SCFDialect>();
+    context.loadDialect<affine::AffineDialect, arith::ArithDialect, scf::SCFDialect>();
   }
 
   void checkUnidimensional(LoopLikeOpInterface loopLikeOp) {
@@ -88,6 +92,24 @@ class SCFLoopLikeTest : public ::testing::Test {
     EXPECT_EQ((*maybeInductionVars).size(), 2u);
   }
 
+  void checkNormalized(LoopLikeOpInterface loopLikeOp) {
+    std::optional<SmallVector<OpFoldResult>> maybeLb =
+        loopLikeOp.getLoopLowerBounds();
+    ASSERT_TRUE(maybeLb.has_value());
+    std::optional<SmallVector<OpFoldResult>> maybeStep =
+        loopLikeOp.getLoopSteps();
+    ASSERT_TRUE(maybeStep.has_value());
+
+    auto allEqual = [](ArrayRef<OpFoldResult> results, int64_t val) {
+      return llvm::all_of(results, [&](OpFoldResult ofr) {
+        auto intValue = getConstantIntValue(ofr);
+        return intValue.has_value() && intValue == val;
+      });
+    };
+    EXPECT_TRUE(allEqual(*maybeLb, 0));
+    EXPECT_TRUE(allEqual(*maybeStep, 1));
+  }
+
   MLIRContext context;
   OpBuilder b;
   Location loc;
@@ -138,3 +160,23 @@ TEST_F(SCFLoopLikeTest, queryMultidimensionalLooplikes) {
       ValueRange({step->getResult(), step->getResult()}), ValueRange());
   checkMultidimensional(parallelOp.get());
 }
+
+TEST_F(SCFLoopLikeTest, testForallNormalize) {
+  OwningOpRef<arith::ConstantIndexOp> lb =
+      b.create<arith::ConstantIndexOp>(loc, 1);
+  OwningOpRef<arith::ConstantIndexOp> ub =
+      b.create<arith::ConstantIndexOp>(loc, 10);
+  OwningOpRef<arith::ConstantIndexOp> step =
+      b.create<arith::ConstantIndexOp>(loc, 3);
+
+  scf::ForallOp forallOp = b.create<scf::ForallOp>(
+      loc, ArrayRef<OpFoldResult>({lb->getResult(), lb->getResult()}),
+      ArrayRef<OpFoldResult>({ub->getResult(), ub->getResult()}),
+      ArrayRef<OpFoldResult>({step->getResult(), step->getResult()}),
+      ValueRange(), std::nullopt);
+  IRRewriter rewriter(b);
+  FailureOr<scf::ForallOp> maybeNormalizedForallOp = normalizeForallOp(rewriter, forallOp);
+  EXPECT_TRUE(succeeded(maybeNormalizedForallOp));
+  OwningOpRef<scf::ForallOp> normalizedForallOp(*maybeNormalizedForallOp);
+  checkNormalized(normalizedForallOp.get());
+}