[mlir] Implement one-to-n structural conversion for ForOp

[sabauma]

Summary

Add the missing one-to-n structural type conversion pattern for the
scf.for operation.

[sabauma]

Code cleanup and commenting

[sabauma]

Delete implicitly created body block

[sabauma]

Fix issue where operations created after the signature conversion were generated
inside the loop body rather than after the loop.

[ingomueller-net]

Thanks a lot for this addition! Looks good modulo the minor comments.

The build failure seems unrelated. I guess rebasing should make them pass.

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

160

Nit: I'd prefer if the comment used the full line length and did not have a line break in between the two sentences.

190

I think that this guard (and the scope around it) isn't necessary.

210

Nit: Please order alphabetically.

[sabauma]

Address comments from @ingomueller-net

[sabauma]

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

190

With the current implementation of applySignatureConversion, the guard is necessary since the rewriter is always left focused on the block that the signature conversion was applied to. That will results in unrealized_conversion_casts being generated in the loop body for the loop outputs when the rewriter.replaceOp is called.

It is not clear to me that is the most desirable behavior of applySignatureConversion, but I opted not to change it, as it is difficult to tell who might be depending on that particular behavior.

[ingomueller-net]

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

190

I can't reproduce this: I have removed guard on my machine and the tests still pass. (Also, I don't see any unrealized cast if I run the pass outside of lit.)

[sabauma]

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

190

Is it possible you did not have the @for_tuple_ops test? I added it later after discovering this specific issue. When I remove the insertion guard, that test fails for me with the following IR (note the pair of make_tuple ops after the scf.yield):

"func.func"() <{function_type = (i1) -> i1, sym_name = "for_tuple_ops"}> ({
  ^bb0(%arg0: i1):
    %0 = "arith.constant"() <{value = 0 : index}> : () -> index
    %1 = "arith.constant"() <{value = 1 : index}> : () -> index
    %2 = "arith.constant"() <{value = 10 : index}> : () -> index
    %3 = "scf.for"(%0, %2, %1, %arg0) ({
    ^bb0(%arg1: index, %arg2: i1):
      %7 = "test.make_tuple"() : () -> tuple<>
      %8 = "test.make_tuple"(%7, %arg2) : (tuple<>, i1) -> tuple<tuple<>, i1>
      %9 = "test.op"(%8) : (tuple<tuple<>, i1>) -> tuple<tuple<>, i1>
      %10 = "test.get_tuple_element"(%9) <{index = 0 : i32}> : (tuple<tuple<>, i1>) -> tuple<>
      %11 = "test.get_tuple_element"(%9) <{index = 1 : i32}> : (tuple<tuple<>, i1>) -> i1
      "scf.yield"(%11) : (i1) -> ()
      %12 = "test.make_tuple"() : () -> tuple<>
      %13 = "test.make_tuple"(%12, %3) : (tuple<>, i1) -> tuple<tuple<>, i1>
    }) : (index, index, index, i1) -> i1
    %4 = "test.op"(%13) : (tuple<tuple<>, i1>) -> tuple<tuple<>, i1>
    %5 = "test.get_tuple_element"(%4) <{index = 0 : i32}> : (tuple<tuple<>, i1>) -> tuple<>
    %6 = "test.get_tuple_element"(%4) <{index = 1 : i32}> : (tuple<tuple<>, i1>) -> i1
    "func.return"(%6) : (i1) -> ()
  }) : () -> ()
}

Those tuple operations should be placed after the for loop, rather than in the loop body.

[ingomueller-net]

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

190

I just double-checked: the test is there in my version -- and it does not fail without the guard.

This isn't a huge deal, I guess. I am rather trying to reproduce the error in order to know how to improve the API. This behavior isn't expected and I guess I'd prefer for it not to require the guard in the usual case. (But for that, a test case were it doesn't behave that way would be helpful.)

[sabauma]

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

190

I agree that the guard is a bit ugly, though I'm now more concerned about the seeming non-determinism between systems.

I'll try uploading a patch with the guard removed and see if the pre-merge testing experiences the same issue I'm seeing.

[sabauma]

test removing insertion guard

[sabauma]

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

190

I'm able to reproduce this issue on the pre-merge testing machines as well:

https://buildkite.com/llvm-project/premerge-checks/builds/162767#01892bd8-a62f-42f2-9e20-9fd10789e81e

[sabauma]

Re-add the InsertionGuard

[ingomueller-net]

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

190

Thanks a lot for helping me investigate and for your patience! I was finally able to reproduce the behavior you observed (I think my build set-up was broken and I ran the wrong version) and to submit a fix: https://reviews.llvm.org/D154684.

[ingomueller-net]

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

190

The fix has landed, so rebasing should make the insertion guard here unnecessary. Thanks for raising the issue and bearing with me!

[sabauma]

Remove InsertionGuard now that https://reviews.llvm.org/D154684 has landed

[sabauma]

@ingomueller-net Would you be able to submit the change. I do not have commit access.

[ingomueller-net]

Sure, done! Thanks again!

mlir/lib/Dialect/SCF/Transforms/OneToNTypeConversion.cpp

matchAndRewrite(ConditionOp op, OpAdaptor adaptor,
     // Convert operands.
     rewriter.updateRootInPlace(
         op, [&] { op->setOperands(adaptor.getFlatOperands()); });
 
     return success();
   }
 };
 
+class ConvertTypesInSCFForOp final : public OneToNOpConversionPattern<ForOp> {
+public:
+  using OneToNOpConversionPattern<ForOp>::OneToNOpConversionPattern;
+
+  static void buildEmptyBody(OpBuilder &, Location, Value, ValueRange) {}
+
+  LogicalResult
+  matchAndRewrite(ForOp forOp, OpAdaptor adaptor,
+                  OneToNPatternRewriter &rewriter) const override {
+    const OneToNTypeMapping &operandMapping = adaptor.getOperandMapping();
+    const OneToNTypeMapping &resultMapping = adaptor.getResultMapping();
+
+    // Nothing to do if there is no non-identity conversion.
+    if (!operandMapping.hasNonIdentityConversion() &&
+        !resultMapping.hasNonIdentityConversion())
+      return failure();
+
+    // If the lower-bound, upper-bound, or step were expanded,
+    // abort the conversion.
+    // This conversion does not know what to do in such cases.
+    ValueRange lbs = adaptor.getLowerBound();
+    ValueRange ubs = adaptor.getUpperBound();
+    ValueRange steps = adaptor.getStep();
+    if (lbs.size() != 1 || ubs.size() != 1 || steps.size() != 1)
+      return rewriter.notifyMatchFailure(
+          forOp, "index operands converted to multiple values");
+
+    Location loc = forOp.getLoc();
+
+    Region *region = &forOp.getRegion();
+    Block *block = &region->front();
+
+    // Construct the new for-op with an empty body.
+    ValueRange newInits = adaptor.getFlatOperands().drop_front(3);
+    auto newOp =
+        rewriter.create<ForOp>(loc, lbs[0], ubs[0], steps[0], newInits);
+    newOp->setAttrs(forOp->getAttrs());
+
+    // Convert the signature of the body region.
+    OneToNTypeMapping bodyTypeMapping(block->getArgumentTypes());
+    if (failed(typeConverter->convertSignatureArgs(block->getArgumentTypes(),
+                                                   bodyTypeMapping)))
+      return failure();
+
+    // Perform signature conversion on the body block.
+    rewriter.applySignatureConversion(block, bodyTypeMapping);
+
+    // Splice the old body region into the new for-op.
+    Region &dstRegion = newOp.getBodyRegion();
+    rewriter.inlineRegionBefore(forOp.getRegion(), dstRegion,
+                                dstRegion.begin());
+
+    rewriter.replaceOp(forOp, newOp.getResults(), resultMapping);
+
+    return success();
+  }
+};
+
 namespace mlir {
 namespace scf {
 
 void populateSCFStructuralOneToNTypeConversions(TypeConverter &typeConverter,
                                                 RewritePatternSet &patterns) {
   patterns.add<
       // clang-format off
       ConvertTypesInSCFConditionOp,
       ConvertTypesInSCFIfOp,
+      ConvertTypesInSCFForOp,
       ConvertTypesInSCFWhileOp,
       ConvertTypesInSCFYieldOp
       // clang-format on
       >(typeConverter, patterns.getContext());
 }
 
 } // namespace scf
 } // namespace mlir

mlir/test/Conversion/OneToNTypeConversion/scf-structural-one-to-n-type-conversion.mlir

%0 = scf.while (%arg2 = %arg0) : (tuple<tuple<>, i1>) -> tuple<tuple<>, i1> {
     scf.condition(%arg1) %1 : tuple<tuple<>, i1>
   } do {
   ^bb0(%arg2: tuple<tuple<>, i1>):
     %1 = "test.source"() : () -> tuple<tuple<>, i1>
     scf.yield %1 : tuple<tuple<>, i1>
   }
   return %0 : tuple<tuple<>, i1>
 }
+
+// -----
+
+// Test case: Nested 1:N type conversion is carried through scf.for and scf.yield.
+
+// CHECK-LABEL: func.func @for_operands_results(
+// CHECK-SAME:                                    %[[ARG0:.*]]: i1,
+// CHECK-SAME:                                    %[[ARG1:.*]]: i2) -> (i1, i2) {
+// CHECK-NEXT:    %[[C0:.+]] = arith.constant 0 : index
+// CHECK-NEXT:    %[[C1:.+]] = arith.constant 1 : index
+// CHECK-NEXT:    %[[C10:.+]] = arith.constant 10 : index
+// CHECK-NEXT:    %[[OUT:.+]]:2 = scf.for %arg2 = %[[C0]] to %[[C10]] step %[[C1]] iter_args(%[[ITER0:.+]] = %[[ARG0]], %[[ITER1:.+]] = %[[ARG1]]) -> (i1, i2) {
+// CHECK-NEXT:     scf.yield %[[ITER0]], %[[ITER1]] : i1, i2
+// CHECK-NEXT:    }
+// CHECK-NEXT:    return %[[OUT]]#0, %[[OUT]]#1 : i1, i2
+
+func.func @for_operands_results(%arg0: tuple<tuple<>, i1, tuple<i2>>) -> tuple<tuple<>, i1, tuple<i2>> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+
+  %0 = scf.for %i = %c0 to %c10 step %c1 iter_args(%acc = %arg0) -> tuple<tuple<>, i1, tuple<i2>> {
+    scf.yield %acc : tuple<tuple<>, i1, tuple<i2>>
+  }
+
+  return %0 : tuple<tuple<>, i1, tuple<i2>>
+}
+
+// -----
+
+// Test case: Nested 1:N type conversion is carried through scf.for and scf.yield
+
+// CHECK-LABEL: func.func @for_tuple_ops(
+// CHECK-SAME:                             %[[ARG0:.+]]: i1) -> i1 {
+// CHECK-NEXT: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-NEXT: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-NEXT: %[[C10:.+]] = arith.constant 10 : index
+// CHECK-NEXT: %[[OUT:.+]] = scf.for %arg1 = %[[C0]] to %[[C10]] step %[[C1]] iter_args(%[[ITER:.+]] = %[[ARG0]]) -> (i1) {
+// CHECK-NEXT:   %[[V1:.+]] = "test.make_tuple"() : () -> tuple<>
+// CHECK-NEXT:   %[[V2:.+]] = "test.make_tuple"(%[[V1]], %[[ITER]]) : (tuple<>, i1) -> tuple<tuple<>, i1>
+// CHECK-NEXT:   %[[V3:.+]] = "test.op"(%[[V2]]) : (tuple<tuple<>, i1>) -> tuple<tuple<>, i1>
+// CHECK-NEXT:   %[[V4:.+]] = "test.get_tuple_element"(%[[V3]]) <{index = 0 : i32}> : (tuple<tuple<>, i1>) -> tuple<>
+// CHECK-NEXT:   %[[V5:.+]] = "test.get_tuple_element"(%[[V3]]) <{index = 1 : i32}> : (tuple<tuple<>, i1>) -> i1
+// CHECK-NEXT:   scf.yield %[[V5]] : i1
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[OUT]] : i1
+
+func.func @for_tuple_ops(%arg0: tuple<tuple<>, i1>) -> tuple<tuple<>, i1> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+
+  %0 = scf.for %i = %c0 to %c10 step %c1 iter_args(%acc = %arg0) -> tuple<tuple<>, i1> {
+    %1 = "test.op"(%acc) : (tuple<tuple<>, i1>) -> tuple<tuple<>, i1>
+    scf.yield %1 : tuple<tuple<>, i1>
+  } 
+  return %0 : tuple<tuple<>, i1>
+}