[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 ↗

(On Diff #493488)

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 ↗

(On Diff #493488)

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 ↗

(On Diff #493488)

. 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/test/Dialect/SCF/while-op-builder.mlir

+// RUN: mlir-opt %s -test-scf-while-op-builder | FileCheck %s
+
+// CHECK-LABEL: @testMatchingTypes
+func.func @testMatchingTypes(%arg0 : i32) {
+  %0 = scf.while (%arg1 = %arg0) : (i32) -> (i32) {
+    %c10 = arith.constant 10 : i32
+    %1 = arith.cmpi slt, %arg1, %c10 : i32
+    scf.condition(%1) %arg1 : i32
+  } do {
+  ^bb0(%arg1: i32):
+    scf.yield %arg1 : i32
+  }
+  // Expect the same loop twice (the dummy added by the test pass and the
+  // original one).
+  // CHECK: %[[V0:.*]] = scf.while (%[[arg1:.*]] = %[[arg0:.*]]) : (i32) -> i32 {
+  // CHECK: %[[V1:.*]] = scf.while (%[[arg2:.*]] = %[[arg0]]) : (i32) -> i32 {
+  return
+}
+
+// CHECK-LABEL: @testNonMatchingTypes
+func.func @testNonMatchingTypes(%arg0 : i32) {
+  %c1 = arith.constant 1 : i32
+  %c10 = arith.constant 10 : i32
+  %0:2 = scf.while (%arg1 = %arg0) : (i32) -> (i32, i32) {
+    %1 = arith.cmpi slt, %arg1, %c10 : i32
+    scf.condition(%1) %arg1, %c1 : i32, i32
+  } do {
+  ^bb0(%arg1: i32, %arg2: i32):
+    %1 = arith.addi %arg1, %arg2 : i32
+    scf.yield %1 : i32
+  }
+  // Expect the same loop twice (the dummy added by the test pass and the
+  // original one).
+  // CHECK: %[[V0:.*]] = scf.while (%[[arg1:.*]] = %[[arg0:.*]]) : (i32) -> (i32, i32) {
+  // CHECK: %[[V1:.*]] = scf.while (%[[arg2:.*]] = %[[arg0]]) : (i32) -> (i32, i32) {
+  return
+}

mlir/test/lib/Dialect/SCF/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRSCFTestPasses
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
   TestSCFUtils.cpp
+  TestWhileOpBuilder.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
   MLIRSCFTransforms
   MLIRSCFUtils
   MLIRTransformUtils
   )

mlir/test/lib/Dialect/SCF/TestWhileOpBuilder.cpp

+//===- TestWhileOpBuilder.cpp - Pass to test WhileOp::build ---------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to test some builder functions of WhileOp. It
+// tests the regression explained in https://reviews.llvm.org/D142952, where
+// a WhileOp::build overload crashed when fed with operands of different types
+// than the result types.
+//
+// To test the build function, the pass copies each WhileOp found in the body
+// of a FuncOp and adds an additional WhileOp with the same operands and result
+// types (but dummy computations) using the builder in question.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+using namespace mlir::arith;
+using namespace mlir::scf;
+
+namespace {
+struct TestSCFWhileOpBuilderPass
+    : public PassWrapper<TestSCFWhileOpBuilderPass,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSCFWhileOpBuilderPass)
+
+  StringRef getArgument() const final { return "test-scf-while-op-builder"; }
+  StringRef getDescription() const final {
+    return "test build functions of scf.while";
+  }
+  explicit TestSCFWhileOpBuilderPass() = default;
+  TestSCFWhileOpBuilderPass(const TestSCFWhileOpBuilderPass &pass) = default;
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+    func.walk([&](WhileOp whileOp) {
+      Location loc = whileOp->getLoc();
+      ImplicitLocOpBuilder builder(loc, whileOp);
+
+      // Create a WhileOp with the same operands and result types.
+      TypeRange resultTypes = whileOp->getResultTypes();
+      ValueRange operands = whileOp->getOperands();
+      builder.create<WhileOp>(
+          loc, resultTypes, operands, /*beforeBuilder=*/
+          [&](OpBuilder &b, Location loc, ValueRange args) {
+            // Just cast the before args into the right types for condition.
+            ImplicitLocOpBuilder builder(loc, b);
+            auto castOp =
+                builder.create<UnrealizedConversionCastOp>(resultTypes, args);
+            auto cmp = builder.create<ConstantIntOp>(/*value=*/1, /*width=*/1);
+            builder.create<ConditionOp>(cmp, castOp->getResults());
+          },
+          /*afterBuilder=*/
+          [&](OpBuilder &b, Location loc, ValueRange args) {
+            // Just cast the after args into the right types for yield.
+            ImplicitLocOpBuilder builder(loc, b);
+            auto castOp = builder.create<UnrealizedConversionCastOp>(
+                operands.getTypes(), args);
+            builder.create<YieldOp>(castOp->getResults());
+          });
+    });
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestSCFWhileOpBuilderPass() {
+  PassRegistration<TestSCFWhileOpBuilderPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestMemRefStrideCalculation();
 void registerTestOpaqueLoc();
 void registerTestPadFusion();
 void registerTestPDLByteCodePass();
 void registerTestPDLLPasses();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
+void registerTestSCFWhileOpBuilderPass();
 void registerTestShapeMappingPass();
 void registerTestSliceAnalysisPass();
 void registerTestTensorCopyInsertionPass();
 void registerTestTensorTransforms();
 void registerTestTilingInterface();
 void registerTestTopologicalSortAnalysisPass();
 void registerTestTransformDialectEraseSchedulePass();
 void registerTestTransformDialectInterpreterPass();
#endif
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPadFusion();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestPDLLPasses();
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();
+  mlir::test::registerTestSCFWhileOpBuilderPass();
   mlir::test::registerTestShapeMappingPass();
   mlir::test::registerTestSliceAnalysisPass();
   mlir::test::registerTestTensorCopyInsertionPass();
   mlir::test::registerTestTensorTransforms();
   mlir::test::registerTestTilingInterface();
   mlir::test::registerTestTopologicalSortAnalysisPass();
   mlir::test::registerTestTransformDialectEraseSchedulePass();
   mlir::test::registerTestTransformDialectInterpreterPass();