address comments

Author: matthias-springer

mlir/include/mlir/Transforms/DialectConversion.h

@@ -543,6 +543,9 @@ class ConversionPattern : public RewritePattern {
   }
 
   /// Hook for derived classes to implement combined matching and rewriting.
+  /// This overload supports only 1:1 replacements. The 1:N overload is called
+  /// by the driver. By default, it calls this 1:1 overload or reports a fatal
+  /// error if 1:N replacements were found.
   virtual LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const {
@@ -551,6 +554,9 @@ class ConversionPattern : public RewritePattern {
     rewrite(op, operands, rewriter);
     return success();
   }
+
+  /// Hook for derived classes to implement combined matching and rewriting.
+  /// This overload supports 1:N replacements.
   virtual LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,
                   ConversionPatternRewriter &rewriter) const {

mlir/lib/Transforms/Utils/DialectConversion.cpp

@@ -1152,11 +1152,18 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
     Type origType = operand.getType();
     Location operandLoc = inputLoc ? *inputLoc : operand.getLoc();
 
+    // Find the most recently mapped value. Unpack all temporary N:1
+    // materializations. Such conversions are a workaround around missing
+    // 1:N support in the ConversionValueMapping. (The conversion patterns
+    // already support 1:N replacements.)
+    Value repl = mapping.lookupOrDefault(operand);
+    SmallVector<Value> unpacked = unpackNTo1Materialization(repl);
+
     if (!currentTypeConverter) {
       // The current pattern does not have a type converter. I.e., it does not
       // distinguish between legal and illegal types. For each operand, simply
       // pass through the most recently mapped value.
-      remapped.push_back({mapping.lookupOrDefault(operand)});
+      remapped.push_back(std::move(unpacked));
       continue;
     }
 
@@ -1178,12 +1185,8 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
 
     if (legalTypes.size() != 1) {
       // TODO: This is a 1:N conversion. The conversion value mapping does not
-      // support such conversions yet. It stores the result of an argument
-      // materialization (i.e., a conversion back into a single SSA value)
-      // instead. Unpack such "workaround" materializations and hand the
-      // original replacement values to the adaptor.
-      Value repl = mapping.lookupOrDefault(operand);
-      SmallVector<Value> unpacked = unpackNTo1Materialization(repl);
+      // store such materializations yet. If the types of the most recently
+      // mapped values do not match, build a target materialization.
       if (TypeRange(unpacked) == legalTypes) {
         remapped.push_back(std::move(unpacked));
         continue;
@@ -1193,7 +1196,7 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
       // different legalized types.
       ValueRange targetMat = buildUnresolvedMaterialization(
           MaterializationKind::Target, computeInsertPoint(repl), operandLoc,
-          /*inputs=*/repl, /*outputType=*/legalTypes,
+          /*inputs=*/unpacked, /*outputType=*/legalTypes,
           /*originalType=*/origType, currentTypeConverter);
       remapped.push_back(targetMat);
       continue;
@@ -1211,7 +1214,7 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
       Value castValue = buildUnresolvedMaterialization(
           MaterializationKind::Target, computeInsertPoint(newOperand),
           operandLoc,
-          /*inputs=*/newOperand, /*outputType=*/desiredType,
+          /*inputs=*/unpacked, /*outputType=*/desiredType,
           /*originalType=*/origType, currentTypeConverter);
       mapping.map(newOperand, castValue);
       newOperand = castValue;
@@ -1447,7 +1450,10 @@ ConversionPatternRewriterImpl::unpackNTo1Materialization(Value value) {
 
   SmallVector<Value> result;
   for (Value v : castOp.getOperands()) {
-    // Keep unpacking if possible.
+    // Keep unpacking if possible. This is needed because during block
+    // signature conversions and 1:N op replacements, the driver may have
+    // inserted two materializations back-to-back: first an argument
+    // materialization, then a target materialization.
     llvm::append_range(result, unpackNTo1Materialization(v));
   }
   return result;

mlir/test/Transforms/test-legalizer.mlir

@@ -463,3 +463,14 @@ func.func @circular_mapping() {
   %0 = "test.erase_op"() : () -> (i64)
   "test.drop_operands_and_replace_with_valid"(%0) : (i64) -> ()
 }
+
+// -----
+
+func.func @test_1_to_n_block_signature_conversion() {
+  "test.duplicate_block_args"() ({
+  ^bb0(%arg0: i64):
+    "test.repetitive_1_to_n_consumer"(%arg0) : (i64) -> ()
+  }) {} : () -> ()
+  "test.return"() : () -> ()
+}
+

mlir/test/lib/Dialect/Test/TestOps.td

@@ -1886,6 +1886,11 @@ def LegalOpC : TEST_Op<"legal_op_c">,
   Arguments<(ins I32)>, Results<(outs I32)>;
 def LegalOpD : TEST_Op<"legal_op_d">, Arguments<(ins AnyType)>;
 
+def DuplicateBlockArgsOp : TEST_Op<"duplicate_block_args", [SingleBlock]> {
+  let arguments = (ins UnitAttr:$is_legal);
+  let regions = (region SizedRegion<1>:$body);
+}
+
 // Check that the conversion infrastructure can properly undo the creation of
 // operations where an operation was created before its parent, in this case,
 // in the parent's builder.

mlir/test/lib/Dialect/Test/TestPatterns.cpp

@@ -982,9 +982,25 @@ struct TestPassthroughInvalidOp : public ConversionPattern {
   TestPassthroughInvalidOp(MLIRContext *ctx)
       : ConversionPattern("test.invalid", 1, ctx) {}
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,
                   ConversionPatternRewriter &rewriter) const final {
-    rewriter.replaceOpWithNewOp<TestValidOp>(op, std::nullopt, operands,
+    SmallVector<Value> flattened;
+    for (auto it : llvm::enumerate(operands)) {
+      ValueRange range = it.value();
+      if (range.size() == 1) {
+        flattened.push_back(range.front());
+        continue;
+      }
+
+      // This is a 1:N replacement. Insert a test.cast op. (That's what the
+      // argument materialization used to do.)
+      flattened.push_back(
+          rewriter
+              .create<TestCastOp>(op->getLoc(),
+                                  op->getOperand(it.index()).getType(), range)
+              .getResult());
+    }
+    rewriter.replaceOpWithNewOp<TestValidOp>(op, std::nullopt, flattened,
                                              std::nullopt);
     return success();
   }
@@ -1010,23 +1026,13 @@ struct TestSplitReturnType : public ConversionPattern {
   TestSplitReturnType(MLIRContext *ctx)
       : ConversionPattern("test.return", 1, ctx) {}
   LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,
                   ConversionPatternRewriter &rewriter) const final {
     // Check for a return of F32.
     if (op->getNumOperands() != 1 || !op->getOperand(0).getType().isF32())
       return failure();
-
-    // Check if the first operation is a cast operation, if it is we use the
-    // results directly.
-    auto *defOp = operands[0].getDefiningOp();
-    if (auto packerOp =
-            llvm::dyn_cast_or_null<UnrealizedConversionCastOp>(defOp)) {
-      rewriter.replaceOpWithNewOp<TestReturnOp>(op, packerOp.getOperands());
-      return success();
-    }
-
-    // Otherwise, fail to match.
-    return failure();
+    rewriter.replaceOpWithNewOp<TestReturnOp>(op, operands[0]);
+    return success();
   }
 };
 
@@ -1181,6 +1187,47 @@ class TestEraseOp : public ConversionPattern {
   }
 };
 
+/// This pattern matches a test.duplicate_block_args op and duplicates all
+/// block arguments.
+class TestDuplicateBlockArgs
+    : public OpConversionPattern<DuplicateBlockArgsOp> {
+  using OpConversionPattern<DuplicateBlockArgsOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(DuplicateBlockArgsOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (op.getIsLegal())
+      return failure();
+    rewriter.startOpModification(op);
+    Block *body = &op.getBody().front();
+    TypeConverter::SignatureConversion result(body->getNumArguments());
+    for (auto it : llvm::enumerate(body->getArgumentTypes()))
+      result.addInputs(it.index(), {it.value(), it.value()});
+    rewriter.applySignatureConversion(body, result, getTypeConverter());
+    op.setIsLegal(true);
+    rewriter.finalizeOpModification(op);
+    return success();
+  }
+};
+
+/// This pattern replaces test.repetitive_1_to_n_consumer ops with a test.valid
+/// op. The pattern supports 1:N replacements and forwards the replacement
+/// values of the single operand as test.valid operands.
+class TestRepetitive1ToNConsumer : public ConversionPattern {
+public:
+  TestRepetitive1ToNConsumer(MLIRContext *ctx)
+      : ConversionPattern("test.repetitive_1_to_n_consumer", 1, ctx) {}
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    // A single operand is expected.
+    if (op->getNumOperands() != 1)
+      return failure();
+    rewriter.replaceOpWithNewOp<TestValidOp>(op, operands.front());
+    return success();
+  }
+};
+
 } // namespace
 
 namespace {
@@ -1258,9 +1305,11 @@ struct TestLegalizePatternDriver
         TestUpdateConsumerType, TestNonRootReplacement,
         TestBoundedRecursiveRewrite, TestNestedOpCreationUndoRewrite,
         TestReplaceEraseOp, TestCreateUnregisteredOp, TestUndoMoveOpBefore,
-        TestUndoPropertiesModification, TestEraseOp>(&getContext());
+        TestUndoPropertiesModification, TestEraseOp,
+        TestRepetitive1ToNConsumer>(&getContext());
     patterns.add<TestDropOpSignatureConversion, TestDropAndReplaceInvalidOp>(
         &getContext(), converter);
+    patterns.add<TestDuplicateBlockArgs>(converter, &getContext());
     mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(patterns,
                                                               converter);
     mlir::populateCallOpTypeConversionPattern(patterns, converter);
@@ -1312,6 +1361,9 @@ struct TestLegalizePatternDriver
     target.addDynamicallyLegalOp<TestOpInPlaceSelfFold>(
         [](TestOpInPlaceSelfFold op) { return op.getFolded(); });
 
+    target.addDynamicallyLegalOp<DuplicateBlockArgsOp>(
+        [](DuplicateBlockArgsOp op) { return op.getIsLegal(); });
+
     // Handle a partial conversion.
     if (mode == ConversionMode::Partial) {
       DenseSet<Operation *> unlegalizedOps;