Fix decompose call graph test

Author: matthias-springer

mlir/include/mlir/Dialect/Func/Transforms/DecomposeCallGraphTypes.h

@@ -23,70 +23,10 @@
 
 namespace mlir {
 
-/// This class provides a hook that expands one Value into multiple Value's,
-/// with a TypeConverter-inspired callback registration mechanism.
-///
-/// For folks that are familiar with the dialect conversion framework /
-/// TypeConverter, this is effectively the inverse of a source/argument
-/// materialization. A target materialization is not what we want here because
-/// it always produces a single Value, but in this case the whole point is to
-/// decompose a Value into multiple Value's.
-///
-/// The reason we need this inverse is easily understood by looking at what we
-/// need to do for decomposing types for a return op. When converting a return
-/// op, the dialect conversion framework will give the list of converted
-/// operands, and will ensure that each converted operand, even if it expanded
-/// into multiple types, is materialized as a single result. We then need to
-/// undo that materialization to a single result, which we do with the
-/// decomposeValue hooks registered on this object.
-///
-/// TODO: Eventually, the type conversion infra should have this hook built-in.
-/// See
-/// https://llvm.discourse.group/t/extending-type-conversion-infrastructure/779/2
-class ValueDecomposer {
-public:
-  /// This method tries to decompose a value of a certain type using provided
-  /// decompose callback functions. If it is unable to do so, the original value
-  /// is returned.
-  void decomposeValue(OpBuilder &, Location, Type, Value,
-                      SmallVectorImpl<Value> &);
-
-  /// This method registers a callback function that will be called to decompose
-  /// a value of a certain type into 0, 1, or multiple values.
-  template <typename FnT, typename T = typename llvm::function_traits<
-                              std::decay_t<FnT>>::template arg_t<2>>
-  void addDecomposeValueConversion(FnT &&callback) {
-    decomposeValueConversions.emplace_back(
-        wrapDecomposeValueConversionCallback<T>(std::forward<FnT>(callback)));
-  }
-
-private:
-  using DecomposeValueConversionCallFn =
-      std::function<std::optional<LogicalResult>(
-          OpBuilder &, Location, Type, Value, SmallVectorImpl<Value> &)>;
-
-  /// Generate a wrapper for the given decompose value conversion callback.
-  template <typename T, typename FnT>
-  DecomposeValueConversionCallFn
-  wrapDecomposeValueConversionCallback(FnT &&callback) {
-    return
-        [callback = std::forward<FnT>(callback)](
-            OpBuilder &builder, Location loc, Type type, Value value,
-            SmallVectorImpl<Value> &newValues) -> std::optional<LogicalResult> {
-          if (T derivedType = dyn_cast<T>(type))
-            return callback(builder, loc, derivedType, value, newValues);
-          return std::nullopt;
-        };
-  }
-
-  SmallVector<DecomposeValueConversionCallFn, 2> decomposeValueConversions;
-};
-
 /// Populates the patterns needed to drive the conversion process for
 /// decomposing call graph types with the given `ValueDecomposer`.
 void populateDecomposeCallGraphTypesPatterns(MLIRContext *context,
                                              const TypeConverter &typeConverter,
-                                             ValueDecomposer &decomposer,
                                              RewritePatternSet &patterns);
 
 } // namespace mlir

mlir/lib/Dialect/Func/Transforms/DecomposeCallGraphTypes.cpp

@@ -13,53 +13,15 @@
 using namespace mlir;
 using namespace mlir::func;
 
-//===----------------------------------------------------------------------===//
-// ValueDecomposer
-//===----------------------------------------------------------------------===//
-
-void ValueDecomposer::decomposeValue(OpBuilder &builder, Location loc,
-                                     Type type, Value value,
-                                     SmallVectorImpl<Value> &results) {
-  for (auto &conversion : decomposeValueConversions)
-    if (conversion(builder, loc, type, value, results))
-      return;
-  results.push_back(value);
-}
-
-//===----------------------------------------------------------------------===//
-// DecomposeCallGraphTypesOpConversionPattern
-//===----------------------------------------------------------------------===//
-
-namespace {
-/// Base OpConversionPattern class to make a ValueDecomposer available to
-/// inherited patterns.
-template <typename SourceOp>
-class DecomposeCallGraphTypesOpConversionPattern
-    : public OpConversionPattern<SourceOp> {
-public:
-  DecomposeCallGraphTypesOpConversionPattern(const TypeConverter &typeConverter,
-                                             MLIRContext *context,
-                                             ValueDecomposer &decomposer,
-                                             PatternBenefit benefit = 1)
-      : OpConversionPattern<SourceOp>(typeConverter, context, benefit),
-        decomposer(decomposer) {}
-
-protected:
-  ValueDecomposer &decomposer;
-};
-} // namespace
-
 //===----------------------------------------------------------------------===//
 // DecomposeCallGraphTypesForFuncArgs
 //===----------------------------------------------------------------------===//
 
 namespace {
-/// Expand function arguments according to the provided TypeConverter and
-/// ValueDecomposer.
+/// Expand function arguments according to the provided TypeConverter.
 struct DecomposeCallGraphTypesForFuncArgs
-    : public DecomposeCallGraphTypesOpConversionPattern<func::FuncOp> {
-  using DecomposeCallGraphTypesOpConversionPattern::
-      DecomposeCallGraphTypesOpConversionPattern;
+    : public OpConversionPattern<func::FuncOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(func::FuncOp op, OpAdaptor adaptor,
@@ -100,19 +62,17 @@ struct DecomposeCallGraphTypesForFuncArgs
 //===----------------------------------------------------------------------===//
 
 namespace {
-/// Expand return operands according to the provided TypeConverter and
-/// ValueDecomposer.
+/// Expand return operands according to the provided TypeConverter.
 struct DecomposeCallGraphTypesForReturnOp
-    : public DecomposeCallGraphTypesOpConversionPattern<ReturnOp> {
-  using DecomposeCallGraphTypesOpConversionPattern::
-      DecomposeCallGraphTypesOpConversionPattern;
+    : public OpConversionPattern<ReturnOp> {
+  using OpConversionPattern::OpConversionPattern;
+
   LogicalResult
-  matchAndRewrite(ReturnOp op, OpAdaptor adaptor,
+  matchAndRewrite(ReturnOp op, OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     SmallVector<Value, 2> newOperands;
-    for (Value operand : adaptor.getOperands())
-      decomposer.decomposeValue(rewriter, op.getLoc(), operand.getType(),
-                                operand, newOperands);
+    for (ValueRange operand : adaptor.getOperands())
+      llvm::append_range(newOperands, operand);
     rewriter.replaceOpWithNewOp<ReturnOp>(op, newOperands);
     return success();
   }
@@ -124,74 +84,53 @@ struct DecomposeCallGraphTypesForReturnOp
 //===----------------------------------------------------------------------===//
 
 namespace {
-/// Expand call op operands and results according to the provided TypeConverter
-/// and ValueDecomposer.
+/// Expand call op operands and results according to the provided TypeConverter.
 struct DecomposeCallGraphTypesForCallOp
-    : public DecomposeCallGraphTypesOpConversionPattern<CallOp> {
-  using DecomposeCallGraphTypesOpConversionPattern::
-      DecomposeCallGraphTypesOpConversionPattern;
+    : public OpConversionPattern<CallOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(CallOp op, OpAdaptor adaptor,
+  matchAndRewrite(CallOp op, OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
 
     // Create the operands list of the new `CallOp`.
     SmallVector<Value, 2> newOperands;
-    for (Value operand : adaptor.getOperands())
-      decomposer.decomposeValue(rewriter, op.getLoc(), operand.getType(),
-                                operand, newOperands);
-
-    // Create the new result types for the new `CallOp` and track the indices in
-    // the new call op's results that correspond to the old call op's results.
-    //
-    // expandedResultIndices[i] = "list of new result indices that old result i
-    // expanded to".
+    for (ValueRange operand : adaptor.getOperands())
+      llvm::append_range(newOperands, operand);
+
+    // Create the new result types for the new `CallOp` and track the number of
+    // replacement types for each original op result.
     SmallVector<Type, 2> newResultTypes;
-    SmallVector<SmallVector<unsigned, 2>, 4> expandedResultIndices;
+    SmallVector<unsigned> expandedResultSizes;
     for (Type resultType : op.getResultTypes()) {
       unsigned oldSize = newResultTypes.size();
       if (failed(typeConverter->convertType(resultType, newResultTypes)))
         return failure();
-      auto &resultMapping = expandedResultIndices.emplace_back();
-      for (unsigned i = oldSize, e = newResultTypes.size(); i < e; i++)
-        resultMapping.push_back(i);
+      expandedResultSizes.push_back(newResultTypes.size() - oldSize);
     }
 
     CallOp newCallOp = rewriter.create<CallOp>(op.getLoc(), op.getCalleeAttr(),
                                                newResultTypes, newOperands);
 
-    // Build a replacement value for each result to replace its uses. If a
-    // result has multiple mapping values, it needs to be materialized as a
-    // single value.
-    SmallVector<Value, 2> replacedValues;
+    // Build a replacement value for each result to replace its uses.
+    SmallVector<ValueRange> replacedValues;
     replacedValues.reserve(op.getNumResults());
+    unsigned startIdx = 0;
     for (unsigned i = 0, e = op.getNumResults(); i < e; ++i) {
-      auto decomposedValues = llvm::to_vector<6>(
-          llvm::map_range(expandedResultIndices[i],
-                          [&](unsigned i) { return newCallOp.getResult(i); }));
-      if (decomposedValues.empty()) {
-        // No replacement is required.
-        replacedValues.push_back(nullptr);
-      } else if (decomposedValues.size() == 1) {
-        replacedValues.push_back(decomposedValues.front());
-      } else {
-        // Materialize a single Value to replace the original Value.
-        Value materialized = getTypeConverter()->materializeArgumentConversion(
-            rewriter, op.getLoc(), op.getType(i), decomposedValues);
-        replacedValues.push_back(materialized);
-      }
+      ValueRange repl = newCallOp.getResults().slice(startIdx, expandedResultSizes[i]);
+      replacedValues.push_back(repl);
+      startIdx += expandedResultSizes[i];
     }
-    rewriter.replaceOp(op, replacedValues);
+    rewriter.replaceOpWithMultiple(op, replacedValues);
     return success();
   }
 };
 } // namespace
 
 void mlir::populateDecomposeCallGraphTypesPatterns(
     MLIRContext *context, const TypeConverter &typeConverter,
-    ValueDecomposer &decomposer, RewritePatternSet &patterns) {
+    RewritePatternSet &patterns) {
   patterns
       .add<DecomposeCallGraphTypesForCallOp, DecomposeCallGraphTypesForFuncArgs,
-           DecomposeCallGraphTypesForReturnOp>(typeConverter, context,
-                                               decomposer);
+           DecomposeCallGraphTypesForReturnOp>(typeConverter, context);
 }

mlir/lib/Transforms/Utils/DialectConversion.cpp

@@ -1277,6 +1277,11 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
     }
 
     // Try to find a mapped value with the desired type.
+    if (legalTypes.empty()) {
+      remapped.push_back({});
+      continue;
+    }
+
     SmallVector<Value, 1> mat = mapping.lookupOrDefault(operand, legalTypes);
     if (!mat.empty()) {
       // Mapped value has the correct type or there is an existing
@@ -2577,34 +2582,29 @@ legalizeUnresolvedMaterialization(RewriterBase &rewriter,
   assert(!op.use_empty() &&
          "expected that dead materializations have already been DCE'd");
   Operation::operand_range inputOperands = op.getOperands();
-  Type outputType = op.getResultTypes()[0];
 
   // Try to materialize the conversion.
   if (const TypeConverter *converter = rewrite->getConverter()) {
     rewriter.setInsertionPoint(op);
-    Value newMaterialization;
+    SmallVector<Value> newMaterialization;
     switch (rewrite->getMaterializationKind()) {
     case MaterializationKind::Argument:
-      // Try to materialize an argument conversion.
-      newMaterialization = converter->materializeArgumentConversion(
-          rewriter, op->getLoc(), outputType, inputOperands);
-      if (newMaterialization)
-        break;
-      // If an argument materialization failed, fallback to trying a target
-      // materialization.
-      [[fallthrough]];
+      llvm_unreachable("argument materializations have been removed");
     case MaterializationKind::Target:
       newMaterialization = converter->materializeTargetConversion(
-          rewriter, op->getLoc(), outputType, inputOperands,
+          rewriter, op->getLoc(), op.getResultTypes(), inputOperands,
           rewrite->getOriginalType());
       break;
     case MaterializationKind::Source:
-      newMaterialization = converter->materializeSourceConversion(
-          rewriter, op->getLoc(), outputType, inputOperands);
+      assert(op.getNumResults() == 1 && "*:N source materializations are not supported");
+      Value sourceMat = converter->materializeSourceConversion(
+          rewriter, op->getLoc(), op.getResultTypes().front(), inputOperands);
+      if (sourceMat)
+        newMaterialization.push_back(sourceMat);
       break;
     }
-    if (newMaterialization) {
-      assert(newMaterialization.getType() == outputType &&
+    if (!newMaterialization.empty()) {
+      assert(TypeRange(newMaterialization) == op.getResultTypes() &&
              "materialization callback produced value of incorrect type");
       rewriter.replaceOp(op, newMaterialization);
       return success();
@@ -2614,8 +2614,8 @@ legalizeUnresolvedMaterialization(RewriterBase &rewriter,
   InFlightDiagnostic diag = op->emitError()
                             << "failed to legalize unresolved materialization "
                                "from ("
-                            << inputOperands.getTypes() << ") to " << outputType
-                            << " that remained live after conversion";
+                            << inputOperands.getTypes() << ") to (" << op.getResultTypes()
+                            << ") that remained live after conversion";
   diag.attachNote(op->getUsers().begin()->getLoc())
       << "see existing live user here: " << *op->getUsers().begin();
   return failure();