[mlir][DialectConversion] Cache type conversions and add a few useful helpers

[rriddle]

Summary

It is quite common for the same type to be converted many types throughout the conversion process, and there isn't any good reason why we aren't caching that result. Especially given that we currently use identity conversion to signify legality. This revision also adds a few additional helpers to TypeConverter.

[ftynse]

mlir/lib/Transforms/DialectConversion.cpp

1745

I may be missing something, but doesn't DenseMap::lookup return a default-constructed ValueTy on failure? This will be a null Type in our case. Then: (a) we ignore cached conversion failures because we never enter this block; (b) the checks on cachedConversion not being null inside the block are spurious.

[rriddle]

Update

[rriddle]

mlir/lib/Transforms/DialectConversion.cpp

1745

Oops, thanks for the catch. Forgot to update this after things changed.

mlir/include/mlir/Transforms/DialectConversion.h

public:
     /// Remap an input of the original signature with a new set of types. The
     /// new types are appended to the new signature conversion.
     void addInputs(unsigned origInputNo, ArrayRef<Type> types);
 
     /// Append new input types to the signature conversion, this should only be
     /// used if the new types are not intended to remap an existing input.
     void addInputs(ArrayRef<Type> types);
 
-    /// Remap an input of the original signature with a range of types in the
-    /// new signature.
-    void remapInput(unsigned origInputNo, unsigned newInputNo,
-                    unsigned newInputCount = 1);
-
     /// Remap an input of the original signature to another `replacement`
     /// value. This drops the original argument.
     void remapInput(unsigned origInputNo, Value replacement);
 
   private:
+    /// Remap an input of the original signature with a range of types in the
+    /// new signature.
+    void remapInput(unsigned origInputNo, unsigned newInputNo,
+                    unsigned newInputCount = 1);
+
     /// The remapping information for each of the original arguments.
     SmallVector<Optional<InputMapping>, 4> remappedInputs;
 
     /// The set of new argument types.
     SmallVector<Type, 4> argTypes;
   };
 
   /// Register a conversion function. A conversion function must be convertible
public:
   /// returns failure if the conversion of any of the types fails, success
   /// otherwise.
   LogicalResult convertTypes(ArrayRef<Type> types,
                              SmallVectorImpl<Type> &results);
 
   /// Return true if the given type is legal for this type converter, i.e. the
   /// type converts to itself.
   bool isLegal(Type type);
+  /// Return true if all of the given types are legal for this type converter.
+  template <typename RangeT>
+  std::enable_if_t<!std::is_convertible<RangeT, Type>::value &&
+                       !std::is_convertible<RangeT, Operation *>::value,
+                   bool>
+  isLegal(RangeT &&range) {
+    return llvm::all_of(range, [this](Type type) { return isLegal(type); });
+  }
+  /// Return true if the given operation has legal operand and result types.
+  bool isLegal(Operation *op);
 
   /// Return true if the inputs and outputs of the given function type are
   /// legal.
-  bool isSignatureLegal(FunctionType funcType);
+  bool isSignatureLegal(FunctionType ty);
 
   /// This method allows for converting a specific argument of a signature. It
   /// takes as inputs the original argument input number, type.
   /// On success, it populates 'result' with any new mappings.
   LogicalResult convertSignatureArg(unsigned inputNo, Type type,
                                     SignatureConversion &result);
+  LogicalResult convertSignatureArgs(TypeRange types,
+                                     SignatureConversion &result,
+                                     unsigned origInputOffset = 0);
 
   /// This function converts the type signature of the given block, by invoking
   /// 'convertSignatureArg' for each argument. This function should return a
   /// valid conversion for the signature on success, None otherwise.
   Optional<SignatureConversion> convertBlockSignature(Block *block);
 
   /// Materialize a conversion from a set of types into one result type by
   /// generating a cast operation of some kind.
return [callback = std::forward<FnT>(callback)](
         return llvm::None;
       return callback(derivedType, results);
     };
   }
 
   /// Register a type conversion.
   void registerConversion(ConversionCallbackFn callback) {
     conversions.emplace_back(std::move(callback));
+    cachedDirectConversions.clear();
+    cachedMultiConversions.clear();
   }
 
   /// Generate a wrapper for the given materialization callback. The callback
   /// may take any subclass of `Type` and the wrapper will check for the target
   /// type to be of the expected class before calling the callback.
   template <typename T, typename FnT>
   MaterializationCallbackFn wrapMaterialization(FnT &&callback) {
     return [callback = std::forward<FnT>(callback)](
void registerMaterialization(MaterializationCallbackFn &&callback) {
     materializations.emplace_back(std::move(callback));
   }
 
   /// The set of registered conversion functions.
   SmallVector<ConversionCallbackFn, 4> conversions;
 
   /// The list of registered materialization functions.
   SmallVector<MaterializationCallbackFn, 2> materializations;
+
+  /// A set of cached conversions to avoid recomputing in the common case.
+  /// Direct 1-1 conversions are the most common, so this cache stores the
+  /// successful 1-1 conversions as well as all failed conversions.
+  DenseMap<Type, Type> cachedDirectConversions;
+  /// This cache stores the successful 1->N conversions, where N != 1.
+  DenseMap<Type, SmallVector<Type, 2>> cachedMultiConversions;
 };
 
 //===----------------------------------------------------------------------===//
 // Conversion Patterns
 //===----------------------------------------------------------------------===//
 
 /// Base class for the conversion patterns that require type changes. Specific
 /// conversions must derive this class and implement least one `rewrite` method.

mlir/lib/Dialect/Linalg/Transforms/TensorsToBuffers.cpp

void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<StandardOpsDialect>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
-    auto isIllegalType = [&](Type type) { return !converter.isLegal(type); };
     auto isLegalOperation = [&](Operation *op) {
-      return llvm::none_of(op->getOperandTypes(), isIllegalType) &&
-             llvm::none_of(op->getResultTypes(), isIllegalType);
+      return converter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
         Optional<ConversionTarget::DynamicLegalityCallbackFn>(
             isLegalOperation));
 
     // Mark Standard Return operations illegal as long as one operand is tensor.
     target.addDynamicallyLegalOp<mlir::ReturnOp>([&](mlir::ReturnOp returnOp) {
-      return llvm::none_of(returnOp.getOperandTypes(), isIllegalType);
+      return converter.isLegal(returnOp.getOperandTypes());
     });
 
     // Mark the function operation illegal as long as an argument is tensor.
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
       return converter.isSignatureLegal(funcOp.getType()) &&
              llvm::none_of(funcOp.getType().getResults(),
                            [&](Type type) { return type.isa<MemRefType>(); });
     });

mlir/lib/Transforms/DialectConversion.cpp

void TypeConverter::SignatureConversion::remapInput(unsigned origInputNo,
   assert(!remappedInputs[origInputNo] && "input has already been remapped");
   remappedInputs[origInputNo] =
       InputMapping{origInputNo, /*size=*/0, replacementValue};
 }
 
 /// This hooks allows for converting a type.
 LogicalResult TypeConverter::convertType(Type t,
                                          SmallVectorImpl<Type> &results) {
+  auto existingIt = cachedDirectConversions.find(t);
+  if (existingIt != cachedDirectConversions.end()) {
+    if (existingIt->second)
+      results.push_back(existingIt->second);
+    return success(existingIt->second != nullptr);
+  }
+  auto multiIt = cachedMultiConversions.find(t);
+  if (multiIt != cachedMultiConversions.end()) {
+    results.append(multiIt->second.begin(), multiIt->second.end());
+    return success();
+  }
+
   // Walk the added converters in reverse order to apply the most recently
   // registered first.
-  for (ConversionCallbackFn &converter : llvm::reverse(conversions))
-    if (Optional<LogicalResult> result = converter(t, results))
-      return *result;
+  size_t currentCount = results.size();
+  for (ConversionCallbackFn &converter : llvm::reverse(conversions)) {
+    if (Optional<LogicalResult> result = converter(t, results)) {
+      if (!succeeded(*result)) {
+        cachedDirectConversions.try_emplace(t, nullptr);
+        return failure();
+      }
+      auto newTypes = ArrayRef<Type>(results).drop_front(currentCount);
+      if (newTypes.size() == 1)
+        cachedDirectConversions.try_emplace(t, newTypes.front());
+      else
+        cachedMultiConversions.try_emplace(t, llvm::to_vector<2>(newTypes));
+      return success();
+    }
+  }
   return failure();
 }
 
 /// This hook simplifies defining 1-1 type conversions. This function returns
 /// the type to convert to on success, and a null type on failure.
 Type TypeConverter::convertType(Type t) {
   // Use the multi-type result version to convert the type.
   SmallVector<Type, 1> results;
LogicalResult TypeConverter::convertTypes(ArrayRef<Type> types,
   for (auto type : types)
     if (failed(convertType(type, results)))
       return failure();
   return success();
 }
 
 /// Return true if the given type is legal for this type converter, i.e. the
 /// type converts to itself.
-bool TypeConverter::isLegal(Type type) {
-  SmallVector<Type, 1> results;
-  return succeeded(convertType(type, results)) && results.size() == 1 &&
-         results.front() == type;
+bool TypeConverter::isLegal(Type type) { return convertType(type) == type; }
+/// Return true if the given operation has legal operand and result types.
+bool TypeConverter::isLegal(Operation *op) {
+  return isLegal(op->getOperandTypes()) && isLegal(op->getResultTypes());
 }
 
 /// Return true if the inputs and outputs of the given function type are
 /// legal.
-bool TypeConverter::isSignatureLegal(FunctionType funcType) {
-  return llvm::all_of(
-      llvm::concat<const Type>(funcType.getInputs(), funcType.getResults()),
-      [this](Type type) { return isLegal(type); });
+bool TypeConverter::isSignatureLegal(FunctionType ty) {
+  return isLegal(llvm::concat<const Type>(ty.getInputs(), ty.getResults()));
 }
 
 /// This hook allows for converting a specific argument of a signature.
 LogicalResult TypeConverter::convertSignatureArg(unsigned inputNo, Type type,
                                                  SignatureConversion &result) {
   // Try to convert the given input type.
   SmallVector<Type, 1> convertedTypes;
   if (failed(convertType(type, convertedTypes)))
     return failure();
 
   // If this argument is being dropped, there is nothing left to do.
   if (convertedTypes.empty())
     return success();
 
   // Otherwise, add the new inputs.
   result.addInputs(inputNo, convertedTypes);
   return success();
 }
+LogicalResult TypeConverter::convertSignatureArgs(TypeRange types,
+                                                  SignatureConversion &result,
+                                                  unsigned origInputOffset) {
+  for (unsigned i = 0, e = types.size(); i != e; ++i)
+    if (failed(convertSignatureArg(origInputOffset + i, types[i], result)))
+      return failure();
+  return success();
+}
 
 Value TypeConverter::materializeConversion(PatternRewriter &rewriter,
                                            Location loc, Type resultType,
                                            ValueRange inputs) {
   for (MaterializationCallbackFn &fn : llvm::reverse(materializations))
     if (Optional<Value> result = fn(rewriter, resultType, inputs, loc))
       return result.getValue();
   return nullptr;
 }
 
+/// This function converts the type signature of the given block, by invoking
+/// 'convertSignatureArg' for each argument. This function should return a valid
+/// conversion for the signature on success, None otherwise.
+auto TypeConverter::convertBlockSignature(Block *block)
+    -> Optional<SignatureConversion> {
+  SignatureConversion conversion(block->getNumArguments());
+  if (failed(convertSignatureArgs(block->getArgumentTypes(), conversion)))
+    return llvm::None;
+  return conversion;
+}
+
 /// Create a default conversion pattern that rewrites the type signature of a
 /// FuncOp.
 namespace {
 struct FuncOpSignatureConversion : public OpConversionPattern<FuncOp> {
   FuncOpSignatureConversion(MLIRContext *ctx, TypeConverter &converter)
       : OpConversionPattern(ctx), converter(converter) {}
 
   /// Hook for derived classes to implement combined matching and rewriting.
   LogicalResult
   matchAndRewrite(FuncOp funcOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     FunctionType type = funcOp.getType();
 
-    // Convert the original function arguments.
+    // Convert the original function types.
     TypeConverter::SignatureConversion result(type.getNumInputs());
-    for (unsigned i = 0, e = type.getNumInputs(); i != e; ++i)
-      if (failed(converter.convertSignatureArg(i, type.getInput(i), result)))
-        return failure();
-
-    // Convert the original function results.
     SmallVector<Type, 1> convertedResults;
-    if (failed(converter.convertTypes(type.getResults(), convertedResults)))
+    if (failed(converter.convertSignatureArgs(type.getInputs(), result)) ||
+        failed(converter.convertTypes(type.getResults(), convertedResults)))
       return failure();
 
     // Update the function signature in-place.
     rewriter.updateRootInPlace(funcOp, [&] {
       funcOp.setType(FunctionType::get(result.getConvertedTypes(),
                                        convertedResults, funcOp.getContext()));
       rewriter.applySignatureConversion(&funcOp.getBody(), result);
     });
     return success();
   }
 
   /// The type converter to use when rewriting the signature.
   TypeConverter &converter;
 };
 } // end anonymous namespace
 
 void mlir::populateFuncOpTypeConversionPattern(
     OwningRewritePatternList &patterns, MLIRContext *ctx,
     TypeConverter &converter) {
   patterns.insert<FuncOpSignatureConversion>(ctx, converter);
 }
 
-/// This function converts the type signature of the given block, by invoking
-/// 'convertSignatureArg' for each argument. This function should return a valid
-/// conversion for the signature on success, None otherwise.
-auto TypeConverter::convertBlockSignature(Block *block)
-    -> Optional<SignatureConversion> {
-  SignatureConversion conversion(block->getNumArguments());
-  for (unsigned i = 0, e = block->getNumArguments(); i != e; ++i)
-    if (failed(convertSignatureArg(i, block->getArgument(i).getType(),
-                                   conversion)))
-      return llvm::None;
-  return conversion;
-}
-
 //===----------------------------------------------------------------------===//
 // ConversionTarget
 //===----------------------------------------------------------------------===//
 
 /// Register a legality action for the given operation.
 void ConversionTarget::setOpAction(OperationName op,
                                    LegalizationAction action) {
   legalOperations[op] = {action, /*isRecursivelyLegal=*/false, llvm::None};

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

struct TestDropOpSignatureConversion : public ConversionPattern {
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Region &region = op->getRegion(0);
     Block *entry = &region.front();
 
     // Convert the original entry arguments.
     TypeConverter::SignatureConversion result(entry->getNumArguments());
-    for (unsigned i = 0, e = entry->getNumArguments(); i != e; ++i)
-      if (failed(converter.convertSignatureArg(
-              i, entry->getArgument(i).getType(), result)))
-        return failure();
+    if (failed(
+            converter.convertSignatureArgs(entry->getArgumentTypes(), result)))
 
     // Convert the region signature and just drop the operation.
     rewriter.applySignatureConversion(&region, result);
     rewriter.eraseOp(op);
     return success();
   }
 
   /// The type converter to use when rewriting the signature.

mlir/test/lib/Transforms/TestBufferPlacement.cpp

void runOnOperation() override {
     MLIRContext &context = this->getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<StandardOpsDialect>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
-    auto isIllegalType = [&](Type type) { return !converter.isLegal(type); };
     auto isLegalOperation = [&](Operation *op) {
-      return llvm::none_of(op->getOperandTypes(), isIllegalType) &&
-             llvm::none_of(op->getResultTypes(), isIllegalType);
+      return converter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
         Optional<ConversionTarget::DynamicLegalityCallbackFn>(
             isLegalOperation));
 
     // Mark Standard Return operations illegal as long as one operand is tensor.
     target.addDynamicallyLegalOp<mlir::ReturnOp>([&](mlir::ReturnOp returnOp) {
-      return llvm::none_of(returnOp.getOperandTypes(), isIllegalType);
+      return converter.isLegal(returnOp.getOperandTypes());
     });
 
     // Mark Standard Call Operation illegal as long as it operates on tensor.
-    target.addDynamicallyLegalOp<mlir::CallOp>([&](mlir::CallOp callOp) {
-      return llvm::none_of(callOp.getOperandTypes(), isIllegalType) &&
-             llvm::none_of(callOp.getResultTypes(), isIllegalType);
-    });
+    target.addDynamicallyLegalOp<mlir::CallOp>(
+        [&](mlir::CallOp callOp) { return converter.isLegal(callOp); });
 
     // Mark the function whose arguments are in tensor-type illegal.
     target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
       return converter.isSignatureLegal(funcOp.getType());
     });
 
     // Walk over all the functions to apply buffer assignment.
     this->getOperation().walk([&](FuncOp function) -> WalkResult {