[mlir] Implement pass utils for 1:N type conversions.

The current dialect conversion does not support 1:N type conversions.
This commit implements a (poor-man's) dialect conversion pass that does
just that. To keep the pass independent of the "real" dialect conversion
infrastructure, it provides a specialization of the TypeConverter class
that allows for N:1 target materializations, a specialization of the
RewritePattern and PatternRewriter classes that automatically add
appropriate unrealized casts supporting 1:N type conversions and provide
converted operands for implementing subclasses, and a conversion driver
that applies the provided patterns and replaces the unrealized casts
that haven't folded away with user-provided materializations.

The current pass is powerful enough to express many existing manual
solutions for 1:N type conversions or extend transforms that previously
didn't support them, out of which this patch implements call graph type
decomposition (which is currently implemented with a ValueDecomposer
that is only used there).

The goal of this pass is to illustrate the effect that 1:N type
conversions could have, gain experience in how patterns should be
written that achieve that effect, and get feedback on how the APIs of
the dialect conversion should be extended or changed to support such
patterns. The hope is that the "real" dialect conversion eventually
supports such patterns, at which point, this pass could be removed
again.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D144469

Author: ingomueller-net

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

@@ -0,0 +1,26 @@
+//===- OneToNTypeFuncConversions.h - 1:N type conv. for Func ----*- C++ -*-===//
+//
+// Licensed 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_FUNC_TRANSFORMS_ONETONTYPEFUNCCONVERSIONS_H
+#define MLIR_DIALECT_FUNC_TRANSFORMS_ONETONTYPEFUNCCONVERSIONS_H
+
+namespace mlir {
+class TypeConverter;
+class RewritePatternSet;
+} // namespace mlir
+
+namespace mlir {
+
+// Populates the provided pattern set with patterns that do 1:N type conversions
+// on func ops. This is intended to be used with `applyPartialOneToNConversion`.
+void populateFuncTypeConversionPatterns(TypeConverter &typeConverter,
+                                        RewritePatternSet &patterns);
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_FUNC_TRANSFORMS_ONETONTYPEFUNCCONVERSIONS_H

mlir/include/mlir/Transforms/OneToNTypeConversion.h

@@ -0,0 +1,256 @@
+//===-- OneToNTypeConversion.h - Utils for 1:N type conversion --*- C++ -*-===//
+//
+// Licensed 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 provides utils for implementing (poor-man's) dialect conversion
+// passes with 1:N type conversions.
+//
+// The main function, `applyPartialOneToNConversion`, first applies a set of
+// `RewritePattern`s, which produce unrealized casts to convert the operands and
+// results from and to the source types, and then replaces all newly added
+// unrealized casts by user-provided materializations. For this to work, the
+// main function requires a special `TypeConverter`, a special
+// `PatternRewriter`, and special RewritePattern`s, which extend their
+// respective base classes for 1:N type converions.
+//
+// Note that this is much more simple-minded than the "real" dialect conversion,
+// which checks for legality before applying patterns and does probably many
+// other additional things. Ideally, some of the extensions here could be
+// integrated there.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_ONETONTYPECONVERSION_H
+#define MLIR_TRANSFORMS_ONETONTYPECONVERSION_H
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+
+/// Extends `TypeConverter` with 1:N target materializations. Such
+/// materializations have to provide the "reverse" of 1:N type conversions,
+/// i.e., they need to materialize N values with target types into one value
+/// with a source type (which isn't possible in the base class currently).
+class OneToNTypeConverter : public TypeConverter {
+public:
+  /// Callback that expresses user-provided materialization logic from the given
+  /// value to N values of the given types. This is useful for expressing target
+  /// materializations for 1:N type conversions, which materialize one value in
+  /// a source type as N values in target types.
+  using OneToNMaterializationCallbackFn =
+      std::function<std::optional<SmallVector<Value>>(OpBuilder &, TypeRange,
+                                                      Value, Location)>;
+
+  /// Creates the mapping of the given range of original types to target types
+  /// of the conversion and stores that mapping in the given (signature)
+  /// conversion. This function simply calls
+  /// `TypeConverter::convertSignatureArgs` and exists here with a different
+  /// name to reflect the broader semantic.
+  LogicalResult computeTypeMapping(TypeRange types,
+                                   SignatureConversion &result) {
+    return convertSignatureArgs(types, result);
+  }
+
+  /// Applies one of the user-provided 1:N target materializations. If several
+  /// exists, they are tried out in the reverse order in which they have been
+  /// added until the first one succeeds. If none succeeds, the functions
+  /// returns `std::nullopt`.
+  std::optional<SmallVector<Value>>
+  materializeTargetConversion(OpBuilder &builder, Location loc,
+                              TypeRange resultTypes, Value input) const;
+
+  /// Adds a 1:N target materialization to the converter. Such materializations
+  /// build IR that converts N values with target types into 1 value of the
+  /// source type.
+  void addTargetMaterialization(OneToNMaterializationCallbackFn &&callback) {
+    oneToNTargetMaterializations.emplace_back(std::move(callback));
+  }
+
+private:
+  SmallVector<OneToNMaterializationCallbackFn> oneToNTargetMaterializations;
+};
+
+/// Stores a 1:N mapping of types and provides several useful accessors. This
+/// class extends `SignatureConversion`, which already supports 1:N type
+/// mappings but lacks some accessors into the mapping as well as access to the
+/// original types.
+class OneToNTypeMapping : public TypeConverter::SignatureConversion {
+public:
+  OneToNTypeMapping(TypeRange originalTypes)
+      : TypeConverter::SignatureConversion(originalTypes.size()),
+        originalTypes(originalTypes) {}
+
+  using TypeConverter::SignatureConversion::getConvertedTypes;
+
+  /// Returns the list of types that corresponds to the original type at the
+  /// given index.
+  TypeRange getConvertedTypes(unsigned originalTypeNo) const;
+
+  /// Returns the list of original types.
+  TypeRange getOriginalTypes() const { return originalTypes; }
+
+  /// Returns the slice of converted values that corresponds the original value
+  /// at the given index.
+  ValueRange getConvertedValues(ValueRange convertedValues,
+                                unsigned originalValueNo) const;
+
+  /// Fills the given result vector with as many copies of the location of the
+  /// original value as the number of values it is converted to.
+  void convertLocation(Value originalValue, unsigned originalValueNo,
+                       llvm::SmallVectorImpl<Location> &result) const;
+
+  /// Fills the given result vector with as many copies of the lociation of each
+  /// original value as the number of values they are respectively converted to.
+  void convertLocations(ValueRange originalValues,
+                        llvm::SmallVectorImpl<Location> &result) const;
+
+  /// Returns true iff at least one type conversion maps an input type to a type
+  /// that is different from itself.
+  bool hasNonIdentityConversion() const;
+
+private:
+  llvm::SmallVector<Type> originalTypes;
+};
+
+/// Extends the basic `RewritePattern` class with a type converter member and
+/// some accessors to it. This is useful for patterns that are not
+/// `ConversionPattern`s but still require access to a type converter.
+class RewritePatternWithConverter : public mlir::RewritePattern {
+public:
+  /// Construct a conversion pattern with the given converter, and forward the
+  /// remaining arguments to RewritePattern.
+  template <typename... Args>
+  RewritePatternWithConverter(TypeConverter &typeConverter, Args &&...args)
+      : RewritePattern(std::forward<Args>(args)...),
+        typeConverter(&typeConverter) {}
+
+  /// Return the type converter held by this pattern, or nullptr if the pattern
+  /// does not require type conversion.
+  TypeConverter *getTypeConverter() const { return typeConverter; }
+
+  template <typename ConverterTy>
+  std::enable_if_t<std::is_base_of<TypeConverter, ConverterTy>::value,
+                   ConverterTy *>
+  getTypeConverter() const {
+    return static_cast<ConverterTy *>(typeConverter);
+  }
+
+protected:
+  /// A type converter for use by this pattern.
+  TypeConverter *const typeConverter;
+};
+
+/// Specialization of `PatternRewriter` that `OneToNConversionPattern`s use. The
+/// class provides additional rewrite methods that are specific to 1:N type
+/// conversions.
+class OneToNPatternRewriter : public PatternRewriter {
+public:
+  OneToNPatternRewriter(MLIRContext *context) : PatternRewriter(context) {}
+
+  /// Replaces the results of the operation with the specified list of values
+  /// mapped back to the original types as specified in the provided type
+  /// mapping. That type mapping must match the replaced op (i.e., the original
+  /// types must be the same as the result types of the op) and the new values
+  /// (i.e., the converted types must be the same as the types of the new
+  /// values).
+  void replaceOp(Operation *op, ValueRange newValues,
+                 const OneToNTypeMapping &resultMapping);
+  using PatternRewriter::replaceOp;
+
+  /// Applies the given argument conversion to the given block. This consists of
+  /// replacing each original argument with N arguments as specified in the
+  /// argument conversion and inserting unrealized casts from the converted
+  /// values to the original types, which are then used in lieu of the original
+  /// ones. (Eventually, `applyPartialOneToNConversion` replaces these casts
+  /// with a user-provided argument materialization if necessary.) This is
+  /// similar to `ArgConverter::applySignatureConversion` but (1) handles 1:N
+  /// type conversion properly and probably (2) doesn't handle many other edge
+  /// cases.
+  Block *applySignatureConversion(Block *block,
+                                  OneToNTypeMapping &argumentConversion);
+};
+
+/// Base class for patterns with 1:N type conversions. Derived classes have to
+/// overwrite the `matchAndRewrite` overlaod that provides additional
+/// information for 1:N type conversions.
+class OneToNConversionPattern : public RewritePatternWithConverter {
+public:
+  using RewritePatternWithConverter::RewritePatternWithConverter;
+
+  /// This function has to be implemented by base classes and is called from the
+  /// usual overloads. Like in "normal" `DialectConversion`, the function is
+  /// provided with the converted operands (which thus have target types). Since
+  /// 1:N conversion are supported, there is usually no 1:1 relationship between
+  /// the original and the converted operands. Instead, the provided
+  /// `operandMapping` can be used to access the converted operands that
+  /// correspond to a particular original operand. Similarly, `resultMapping`
+  /// is provided to help with assembling the result values, which may have 1:N
+  /// correspondences as well. In that case, the original op should be replaced
+  /// with the overload of `replaceOp` that takes the provided `resultMapping`
+  /// in order to deal with the mapping of converted result values to their
+  /// usages in the original types correctly.
+  virtual LogicalResult matchAndRewrite(Operation *op,
+                                        OneToNPatternRewriter &rewriter,
+                                        const OneToNTypeMapping &operandMapping,
+                                        const OneToNTypeMapping &resultMapping,
+                                        ValueRange convertedOperands) const = 0;
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const final;
+};
+
+/// This class is a wrapper around `OneToNConversionPattern` for matching
+/// against instances of a particular op class.
+template <typename SourceOp>
+class OneToNOpConversionPattern : public OneToNConversionPattern {
+public:
+  OneToNOpConversionPattern(TypeConverter &typeConverter, MLIRContext *context,
+                            PatternBenefit benefit = 1,
+                            ArrayRef<StringRef> generatedNames = {})
+      : OneToNConversionPattern(typeConverter, SourceOp::getOperationName(),
+                                benefit, context, generatedNames) {}
+
+  using OneToNConversionPattern::matchAndRewrite;
+
+  /// Overload that derived classes have to override for their op type.
+  virtual LogicalResult matchAndRewrite(SourceOp op,
+                                        OneToNPatternRewriter &rewriter,
+                                        const OneToNTypeMapping &operandMapping,
+                                        const OneToNTypeMapping &resultMapping,
+                                        ValueRange convertedOperands) const = 0;
+
+  LogicalResult matchAndRewrite(Operation *op, OneToNPatternRewriter &rewriter,
+                                const OneToNTypeMapping &operandMapping,
+                                const OneToNTypeMapping &resultMapping,
+                                ValueRange convertedOperands) const final {
+    return matchAndRewrite(cast<SourceOp>(op), rewriter, operandMapping,
+                           resultMapping, convertedOperands);
+  }
+};
+
+/// Applies the given set of patterns recursively on the given op and adds user
+/// materializations where necessary. The patterns are expected to be
+/// `OneToNConversionPattern`, which help converting the types of the operands
+/// and results of the matched ops. The provided type converter is used to
+/// convert the operands of matched ops from their original types to operands
+/// with different types. Unlike in `DialectConversion`, this supports 1:N type
+/// conversions. Those conversions at the "boundary" of the pattern application,
+/// where converted results are not consumed by replaced ops that expect the
+/// converted operands or vice versa, the function inserts user materializations
+/// from the type converter. Also unlike `DialectConversion`, there are no legal
+/// or illegal types; the function simply applies the given patterns and does
+/// not fail if some ops or types remain unconverted (i.e., the conversion is
+/// only "partial").
+LogicalResult
+applyPartialOneToNConversion(Operation *op, OneToNTypeConverter &typeConverter,
+                             const FrozenRewritePatternSet &patterns);
+
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_ONETONTYPECONVERSION_H

mlir/lib/Dialect/Func/Transforms/CMakeLists.txt

@@ -3,6 +3,7 @@ add_mlir_dialect_library(MLIRFuncTransforms
   DuplicateFunctionElimination.cpp
   FuncBufferize.cpp
   FuncConversions.cpp
+  OneToNFuncConversions.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Func/Transforms