[mlir][scf] Fix builder of WhileOp with region builder arguments.

[ingomueller-net]

Summary

The overload of WhileOp::build with arguments for builder functions for
the regions of the op was broken: It did not compute correctly the types
(and locations) of the region arguments, which lead to failed assertions
when the result types were different from the operand types.
Specifically, it used the result types (and operand locations) for *both*
regions, instead of the operand types (and locations) for the 'before'
region and the result types (and loecations) for the 'after' region.

[ingomueller-net]

I am not sure about whether this is the right form of testing the op. Note that the unit test crashes with a failed assertion if SCF.cpp isn't changed. I don't think that a lit test with an .mlir file would do the trick because it doesn't exercise that builder. If anyone has a better idea, please let me know.

mlir/lib/Dialect/SCF/IR/SCF.cpp

2771

Note that I am not sure whether these are the right locations. Do the arguments originate from the *operands* of the while loop or from the while loop itself? (In the latter case, this would have to be odsState.location, I guess.)

[mehdi_amini]

mlir/unittests/Dialect/SCF/BuildersTest.cpp

74

Is there a test-pass that could get this code? (and as such run it as a lit test)

[ingomueller-net]

mlir/unittests/Dialect/SCF/BuildersTest.cpp

74

git grep WhileOp -- mlir/test/ does not yield any results, and the three test passes for SCF (test-scf-for-utils, test-scf-if-utils, and test-scf-pipelining) test specific utility functions. In both cases, I don't see how I could piggy back my check onto something existing.

I *could* write a test pass that exercises that builder, though, if people think that that's more elegant. What would it do? Walk over all WhileOps and replace each of them with an identical copy using the builder in question?

Another possibility is to modernize some (non-test) code to use that builder. One possible place is this one in the ForToWhileLoop transform. That has certainly the lowest testing footprint but also the risk of the builder not being tested anymore if that code is refactored. Also note that this usage of the builder does not seem to trigger the bug, so just using the builder somewhere isn't a guarantee to test it properly.

I guess I can mitigate both problems by modernizing all other places I can find, such as this one, this one, and up to seven more in two transforms of SparseTensor.

What do people think?

[mehdi_amini]

mlir/unittests/Dialect/SCF/BuildersTest.cpp

74

. What would it do? Walk over all WhileOps and replace each of them with an identical copy using the builder in question?

Technically, it could just do exactly the same thing as what you're doing here: take a module, build a while loop into it, and that's it.

The reason for preferring lit test is mostly uniformity with the codebase, and historically the fact that it involved linking a single mlir-opt binary instead of relinking many unit-tests binaries (saving time and disk space). I'm not sure if this tradeoff is still valid with modern machines though, I frequently felt that C++ unit-tests allowed for more "targeted" checks that were a bit convoluted with FileCheck. Our test passes sometimes are in this category as well: printing "breadcrumbs" that are intended to be matched by FileCheck.

Here is one reference: https://lists.llvm.org/pipermail/llvm-dev/2020-June/142932.html

[ingomueller-net]

Implement test as a test pass rather than googletest-based unittest.

This implements the suggestion by @mehdi_amini, which I think works nicely here.

[ingomueller-net]

@mehdi_amini, thanks a lot for the explanations and the suggestion!

[Mogball]

I think a unit test would've been more appropriate in this case, but LGTM

[mehdi_amini]

In D142952#4100394, @Mogball wrote:

I think a unit test would've been more appropriate in this case, but LGTM

I don't disagree, but we don't have precedent to unit-test the builders, do we? Maybe we should change our "policy" (more "existing practice") on unit-tests? Worth a discourse proposal if you're motivated :)

[ingomueller-net]

I don't have a strong opinion either way. I'll wait until Monday and then land, unless somebody tells me not to.

[ingomueller-net]

Build failed due to presumably unrelated problem. Rebasing to trigger a rebuild.

mlir/lib/Dialect/SCF/IR/SCF.cpp

void WhileOp::build(::mlir::OpBuilder &odsBuilder,
   assert(beforeBuilder && "the builder callback for 'before' must be present");
   assert(afterBuilder && "the builder callback for 'after' must be present");
 
   odsState.addOperands(operands);
   odsState.addTypes(resultTypes);
 
   OpBuilder::InsertionGuard guard(odsBuilder);
 
-  SmallVector<Location, 4> blockArgLocs;
+  // Build before region.
+  SmallVector<Location, 4> beforeArgLocs;
+  beforeArgLocs.reserve(operands.size());
   for (Value operand : operands) {
-    blockArgLocs.push_back(operand.getLoc());
+    beforeArgLocs.push_back(operand.getLoc());
   }
 
   Region *beforeRegion = odsState.addRegion();
-  Block *beforeBlock = odsBuilder.createBlock(beforeRegion, /*insertPt=*/{},
-                                              resultTypes, blockArgLocs);
+  Block *beforeBlock = odsBuilder.createBlock(
+      beforeRegion, /*insertPt=*/{}, operands.getTypes(), beforeArgLocs);
   beforeBuilder(odsBuilder, odsState.location, beforeBlock->getArguments());
 
+  // Build after region.
+  SmallVector<Location, 4> afterArgLocs(resultTypes.size(), odsState.location);
+
   Region *afterRegion = odsState.addRegion();
   Block *afterBlock = odsBuilder.createBlock(afterRegion, /*insertPt=*/{},
-                                             resultTypes, blockArgLocs);
+                                             resultTypes, afterArgLocs);
   afterBuilder(odsBuilder, odsState.location, afterBlock->getArguments());
 }
 
 OperandRange WhileOp::getSuccessorEntryOperands(std::optional<unsigned> index) {
   assert(index && *index == 0 &&
          "WhileOp is expected to branch only to the first region");
 
   return getInits();

mlir/unittests/Dialect/CMakeLists.txt

 add_mlir_unittest(MLIRDialectTests
   BroadcastShapeTest.cpp
 )
 target_link_libraries(MLIRDialectTests
   PRIVATE
   MLIRIR
   MLIRDialect)
 
 add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)
+add_subdirectory(SCF)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(Transform)
 add_subdirectory(Utils)

mlir/unittests/Dialect/SCF/BuildersTest.cpp

+//===- BuildersTest.cpp - unit tests for scf op builders ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::scf;
+using namespace mlir::arith;
+
+// Ensure that WhileOp::biuld with builder functions for its regions compiles
+// and produces IR that verifies.
+TEST(SCFBuildersTest, whileOpBuilderWithRegionBuilderArgs) {
+  MLIRContext ctx;
+  ctx.getOrLoadDialect<ArithDialect>();
+  ctx.getOrLoadDialect<SCFDialect>();
+
+  OpBuilder b(&ctx);
+  Location loc = UnknownLoc::get(&ctx);
+  Type indexType = b.getIndexType();
+
+  // Build the following IR:
+  //
+  // module {
+  //   %c0 = arith.constant 0 : index
+  //   %c1 = arith.constant 1 : index
+  //   %c10 = arith.constant 10 : index
+  //   %0:2 = scf.while (%arg0 = %c0) : (index) -> (index, index) {
+  //     %1 = arith.cmpi slt, %arg0, %c10 : index
+  //     scf.condition(%1) %arg0, %c1 : index, index
+  //   } do {
+  //   ^bb0(%arg0: index, %arg1: index):
+  //     %1 = arith.addi %arg0, %arg1 : index
+  //     scf.yield %1 : index
+  //   }
+  // }
+
+  OwningOpRef<ModuleOp> module = b.create<ModuleOp>(loc);
+  b.setInsertionPointToStart(&module->getBodyRegion().front());
+
+  Value zero = b.create<ConstantIndexOp>(loc, 0);
+  Value one = b.create<ConstantIndexOp>(loc, 1);
+  Value ten = b.create<ConstantIndexOp>(loc, 10);
+
+  // Make argument types of regions differnt to catch some possible bugs.
+  SmallVector<Type, 1> operandTypes = {indexType};
+  SmallVector<Type, 2> resultTypes = {indexType, indexType};
+
+  // Make sure the builder runs without crash.
+  auto whileOp = b.create<WhileOp>(
+      loc, resultTypes, /*operands=*/zero, /*beforeBuilder=*/
+      [&](OpBuilder &builder, Location loc, ValueRange args) {
+        Value cmp =
+            builder.create<CmpIOp>(loc, CmpIPredicate::slt, args[0], ten);
+        EXPECT_EQ(args.size(), operandTypes.size());
+        builder.create<ConditionOp>(loc, cmp, ValueRange{args[0], one});
+      },
+      /*afterBuilder=*/
+      [&](OpBuilder &builder, Location loc, ValueRange args) {
+        EXPECT_EQ(args.size(), resultTypes.size());
+        Value updatedIndex = builder.create<AddIOp>(loc, args[0], args[1]);
+        builder.create<YieldOp>(loc, updatedIndex);
+      });
+
+  // Make sure the op we built verifies.
+  EXPECT_TRUE(whileOp.verify().succeeded());
+}

mlir/unittests/Dialect/SCF/CMakeLists.txt

+add_mlir_unittest(MLIRSCFTests
+  BuildersTest.cpp
+)
+target_link_libraries(MLIRSCFTests
+  PRIVATE
+  MLIRSCFDialect
+  )