[mlir] Refactor TypeConverter to add conversions without inheritance

[rriddle]

Summary

This revision refactors the TypeConverter class to not use inheritance to add type conversions. It instead moves to a registration based system, where conversion callbacks are added to the converter with addConversion. This method takes a conversion callback, which must be convertible to any of the following forms(where T is a class derived from Type:

• Optional<Type> (T type)
  • This form represents a 1-1 type conversion. It should return nullptr or llvm::None to signify failure. If llvm::None is returned, the converter is allowed to try another conversion function to perform the conversion.
• Optional<LogicalResult>(T type, SmallVectorImpl<Type> &results)
  • This form represents a 1-N type conversion. It should return failure or llvm::None to signify a failed conversion. If the new set of types is empty, the type is removed and any usages of the existing value are expected to be removed during conversion. If llvm::None is returned, the converter is allowed to try another conversion function to perform the conversion.

When attempting to convert a type, the TypeConverter walks each of the registered converters starting with the one registered most recently.

[ftynse]

Thanks River! I had prototyped a similar thing, but you've beaten me to it and added fancy templates :)

I don't fully understand the meaning of return values on the conversion functions, and how they relate to the return value of TypeConverter::convertType.

• When we have Optional<LogicalResult>, what is the difference between llvm::None and Failure?
• TypeConverter::convertType currently returns nullptr on error, and it looks like it will keep doing so with both llvm::None and Failure making them indistinguishable.

I can see use case for having Optional<Type>: Optional<null-type> means we want to erase the type, equivalent to returning an empty vector in the vector version, and llvm::None means there was an error, but arguably one should just use the vector version for 1->0 type conversion (the Type version would only support 1-1 mappings).

On a separate but relate subject (not for this commit), I have been repeatedly asking people to check the result of type conversion before using the type because it may fail and crash their code. Maybe we should replace Type TypeConverter::convertType(Type) with LogicalResult TypeConverter::convertType(Type, Type&) to enforce that at the API level. WDYT?

mlir/include/mlir/Transforms/DialectConversion.h

179

checking resultOp for being non-null is redundant here, the if the previous line ensured that

181

Same here, resultOpt always has value

[ftynse]

Could you also update https://mlir.llvm.org/docs/DialectConversion/#type-conversion ?

[rriddle]

In D74584#1875881, @ftynse wrote:

Thanks River! I had prototyped a similar thing, but you've beaten me to it and added fancy templates :)

I don't fully understand the meaning of return values on the conversion functions, and how they relate to the return value of TypeConverter::convertType.

• When we have Optional<LogicalResult>, what is the difference between llvm::None and Failure?
• TypeConverter::convertType currently returns nullptr on error, and it looks like it will keep doing so with both llvm::None and Failure making them indistinguishable.

TypeConverter is a bit different than ConversionTarget with how converters get registered. The result types are mapped as:

• success: The conversion was successful
• failure: The conversion failed, abort the conversion process.
• llvm::None: The conversion failed, but the TypeConverter is allowed to fallback to another converter.

There can be different conversions that operate within the same subset of types. For example, consider the following:

TypeConverter converter;

// A converter to handle TensorTypes.
converter.addConversion([](TensorType type) { ... });

// A converter to handle a subset of types that match 'somePredicate'.
converter.addConversion([](Type type) {
  if (somePredicate(type)) {
    ...
    return success();
  }

  // Returning failure means that no other converters can attempt to convert this type.
  return failure();

  // Returning llvm::None means that the TypeConverter can fallback to other converters.
  return llvm::None; 
});

Without being able to toggle when fallbacks are allowed, you could only effectively have one converter registered at a time. In the above, only the 'somePredicate' converter would ever fire. What I intended for llvm::None was to differentiate "This type is illegal/isn't supported" and "I can't handle this type, but another converter might be able to." Does that make sense? If it does, is there a way that you see to make this more clear?

I can see use case for having Optional<Type>: Optional<null-type> means we want to erase the type, equivalent to returning an empty vector in the vector version, and llvm::None means there was an error, but arguably one should just use the vector version for 1->0 type conversion (the Type version would only support 1-1 mappings).

On a separate but relate subject (not for this commit), I have been repeatedly asking people to check the result of type conversion before using the type because it may fail and crash their code. Maybe we should replace Type TypeConverter::convertType(Type) with LogicalResult TypeConverter::convertType(Type, Type&) to enforce that at the API level. WDYT?

I would be +1 for that. I can still see use cases for when you know for a fact the type is convertible and want a convenient API, we could likely have that constraint spelled out in the method name and assert internally that the conversion succeeds.

[rriddle]

Resolve comments.

mlir/include/mlir/Transforms/DialectConversion.h

179

Thanks! Merged two functions together and missed updating this.

[ftynse]

In D74584#1878076, @rriddle wrote:

TypeConverter is a bit different than ConversionTarget with how converters get registered. The result types are mapped as:

• success: The conversion was successful
• failure: The conversion failed, abort the conversion process.
• llvm::None: The conversion failed, but the TypeConverter is allowed to fallback to another converter.

There can be different conversions that operate within the same subset of types. For example, consider the following:

<...>

Without being able to toggle when fallbacks are allowed, you could only effectively have one converter registered at a time. In the above, only the 'somePredicate' converter would ever fire. What I intended for llvm::None was to differentiate "This type is illegal/isn't supported" and "I can't handle this type, but another converter might be able to." Does that make sense? If it does, is there a way that you see to make this more clear?

The documentation exactly as in the list above should be sufficient. I could also interpret the negative results in the opposite way: llvm::None - error happened, abort the process; failure() - failed to convert with the callback, try the next one, so I'd prefer stated explicitly. The only thing that is more clear is a special tri-state return type with Success, TryNext, Abort, but it's probably an overkill.

I don't immediately see the case where one would want to disallow other converters, but they may appear given the infra allows it.

I can see use case for having Optional<Type>: Optional<null-type> means we want to erase the type, equivalent to returning an empty vector in the vector version, and llvm::None means there was an error, but arguably one should just use the vector version for 1->0 type conversion (the Type version would only support 1-1 mappings).

On a separate but relate subject (not for this commit), I have been repeatedly asking people to check the result of type conversion before using the type because it may fail and crash their code. Maybe we should replace Type TypeConverter::convertType(Type) with LogicalResult TypeConverter::convertType(Type, Type&) to enforce that at the API level. WDYT?

I would be +1 for that. I can still see use cases for when you know for a fact the type is convertible and want a convenient API, we could likely have that constraint spelled out in the method name and assert internally that the conversion succeeds.

Makes sense to me!

[rriddle]

In D74584#1878629, @ftynse wrote:

In D74584#1878076, @rriddle wrote:

TypeConverter is a bit different than ConversionTarget with how converters get registered. The result types are mapped as:

• success: The conversion was successful
• failure: The conversion failed, abort the conversion process.
• llvm::None: The conversion failed, but the TypeConverter is allowed to fallback to another converter.

There can be different conversions that operate within the same subset of types. For example, consider the following:

<...>

Without being able to toggle when fallbacks are allowed, you could only effectively have one converter registered at a time. In the above, only the 'somePredicate' converter would ever fire. What I intended for llvm::None was to differentiate "This type is illegal/isn't supported" and "I can't handle this type, but another converter might be able to." Does that make sense? If it does, is there a way that you see to make this more clear?

The documentation exactly as in the list above should be sufficient. I could also interpret the negative results in the opposite way: llvm::None - error happened, abort the process; failure() - failed to convert with the callback, try the next one, so I'd prefer stated explicitly. The only thing that is more clear is a special tri-state return type with Success, TryNext, Abort, but it's probably an overkill.

SGTM, will leave as-is for now and add a tri-state later if the readability becomes an issue.

I don't immediately see the case where one would want to disallow other converters, but they may appear given the infra allows it.

The current cases I've seen are for example; I want to reuse the LLVMTypeConverter but override the conversion for a specific type(e.g. VectorType). There is a possibility that we can solve this in a better way, but I'll leave that for a followup when we have more concrete usages.

I can see use case for having Optional<Type>: Optional<null-type> means we want to erase the type, equivalent to returning an empty vector in the vector version, and llvm::None means there was an error, but arguably one should just use the vector version for 1->0 type conversion (the Type version would only support 1-1 mappings).

On a separate but relate subject (not for this commit), I have been repeatedly asking people to check the result of type conversion before using the type because it may fail and crash their code. Maybe we should replace Type TypeConverter::convertType(Type) with LogicalResult TypeConverter::convertType(Type, Type&) to enforce that at the API level. WDYT?

I would be +1 for that. I can still see use cases for when you know for a fact the type is convertible and want a convenient API, we could likely have that constraint spelled out in the method name and assert internally that the conversion succeeds.

Makes sense to me!

[rriddle]

Thanks for the review Alex!

[rriddle]

Rebase

mlir/docs/DialectConversion.md

 ### Type Converter
 
 As stated above, the `TypeConverter` contains several hooks for detailing how to
 convert types. Several of these hooks are detailed below:
 
 ```c++
 class TypeConverter {
  public:
-  /// This hook allows for converting a type. This function should return
-  /// failure if no valid conversion exists, success otherwise. If the new set
-  /// of types is empty, the type is removed and any usages of the existing
-  /// value are expected to be removed during conversion.
-  virtual LogicalResult convertType(Type t, SmallVectorImpl<Type> &results);
-
-  /// This hook simplifies defining 1-1 type conversions. This function returns
-  /// the type to convert to on success, and a null type on failure.
-  virtual Type convertType(Type t);
+  /// Register a conversion function. A conversion function must be convertible
+  /// to any of the following forms(where `T` is a class derived from `Type`:
+  ///   * Optional<Type>(T)
+  ///     - This form represents a 1-1 type conversion. It should return nullptr
+  ///       or `llvm::None` to signify failure. If `llvm::None` is returned, the
+  ///       converter is allowed to try another conversion function to perform
+  ///       the conversion.
+  ///   * Optional<LogicalResult>(T, SmallVectorImpl<Type> &)
+  ///     - This form represents a 1-N type conversion. It should return
+  ///       `failure` or `llvm::None` to signify a failed conversion. If the new
+  ///       set of types is empty, the type is removed and any usages of the
+  ///       existing value are expected to be removed during conversion. If
+  ///       `llvm::None` is returned, the converter is allowed to try another
+  ///       conversion function to perform the conversion.
+  ///
+  /// When attempting to convert a type, e.g. via `convertType`, the
+  /// `TypeConverter` will invoke each of the converters starting with the one
+  /// most recently registered.
+  template <typename ConversionFnT>
+  void addConversion(ConversionFnT &&callback);
 
   /// This hook allows for materializing a conversion from a set of types into
   /// one result type by generating a cast operation of some kind. The generated
   /// operation should produce one result, of 'resultType', with the provided
   /// 'inputs' as operands. This hook must be overridden when a type conversion
   /// results in more than one type, or if a type conversion may persist after
   /// the conversion has finished.
   virtual Operation *materializeConversion(PatternRewriter &rewriter,

mlir/include/mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h

 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 class MLIRContext;
 class ModuleOp;
 template <typename T> class OpPassBase;
 
-class LinalgTypeConverter : public LLVMTypeConverter {
-public:
-  using LLVMTypeConverter::LLVMTypeConverter;
-  Type convertType(Type t) override;
-};
-
 /// Populate the given list with patterns that convert from Linalg to LLVM.
-void populateLinalgToLLVMConversionPatterns(LinalgTypeConverter &converter,
+void populateLinalgToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                             OwningRewritePatternList &patterns,
                                             MLIRContext *ctx);
 
 /// Create a pass to convert Linalg operations to the LLVMIR dialect.
 std::unique_ptr<OpPassBase<ModuleOp>> createConvertLinalgToLLVMPass();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_LINALGTOLLVM_LINALGTOLLVM_H_

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

public:
   /// Create an LLVMTypeConverter using the default
   /// LLVMTypeConverterCustomization.
   LLVMTypeConverter(MLIRContext *ctx);
 
   /// Create an LLVMTypeConverter using 'custom' customizations.
   LLVMTypeConverter(MLIRContext *ctx,
                     const LLVMTypeConverterCustomization &custom);
 
-  /// Convert types to LLVM IR.  This calls `convertAdditionalType` to convert
-  /// non-standard or non-builtin types.
-  Type convertType(Type t) override;
-
   /// Convert a function type.  The arguments and results are converted one by
   /// one and results are packed into a wrapped LLVM IR structure type. `result`
   /// is populated with argument mapping.
   LLVM::LLVMType convertFunctionSignature(FunctionType type, bool isVariadic,
                                           SignatureConversion &result);
 
   /// Convert a non-empty list of types to be returned from a function into a
   /// supported LLVM IR type.  In particular, if more than one values is
Operation *materializeConversion(PatternRewriter &rewriter, Type type,
                                    Location loc) override;
 
 protected:
   /// LLVM IR module used to parse/create types.
   llvm::Module *module;
   LLVM::LLVMDialect *llvmDialect;
 
 private:
-  Type convertStandardType(Type type);
-
   // Convert a function type.  The arguments and results are converted one by
   // one.  Additionally, if the function returns more than one value, pack the
   // results into an LLVM IR structure type so that the converted function type
   // returns at most one result.
   Type convertFunctionType(FunctionType type);
 
   // Convert the index type.  Uses llvmModule data layout to create an integer
   // of the pointer bitwidth.

mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h

 
 /// Type conversion from standard types to SPIR-V types for shader interface.
 ///
 /// For composite types, this converter additionally performs type wrapping to
 /// satisfy shader interface requirements: shader interface types must be
 /// pointers to structs.
 class SPIRVTypeConverter : public TypeConverter {
 public:
-  using TypeConverter::TypeConverter;
-
-  /// Converts the given standard `type` to SPIR-V correspondence.
-  Type convertType(Type type) override;
+  SPIRVTypeConverter();
 
   /// Gets the SPIR-V correspondence for the standard index type.
   static Type getIndexType(MLIRContext *context);
 };
 
 /// Base class to define a conversion pattern to lower `SourceOp` into SPIR-V.
 template <typename SourceOp>
 class SPIRVOpLowering : public OpConversionPattern<SourceOp> {

mlir/include/mlir/Support/STLExtras.h

struct FunctionTraits<ReturnType (*)(Args...), false> {
 
   /// The result type of this function.
   using result_t = ReturnType;
 
   /// The type of an argument to this function.
   template <size_t i>
   using arg_t = typename std::tuple_element<i, std::tuple<Args...>>::type;
 };
+/// Overload for non-class function type references.
+template <typename ReturnType, typename... Args>
+struct FunctionTraits<ReturnType (&)(Args...), false>
+    : public FunctionTraits<ReturnType (*)(Args...)> {};
 } // end namespace mlir
 
 // Allow tuples to be usable as DenseMap keys.
 // TODO: Move this to upstream LLVM.
 
 /// Simplistic combination of 32-bit hash values into 32-bit hash values.
 /// This function is taken from llvm/ADT/DenseMapInfo.h.
 static inline unsigned llvm_combineHashValue(unsigned a, unsigned b) {

mlir/include/mlir/Transforms/DialectConversion.h

public:
   private:
     /// The remapping information for each of the original arguments.
     SmallVector<Optional<InputMapping>, 4> remappedInputs;
 
     /// The set of new argument types.
     SmallVector<Type, 4> argTypes;
   };
 
-  /// This hook allows for converting a type. This function should return
-  /// failure if no valid conversion exists, success otherwise. If the new set
-  /// of types is empty, the type is removed and any usages of the existing
-  /// value are expected to be removed during conversion.
-  virtual LogicalResult convertType(Type t, SmallVectorImpl<Type> &results);
+  /// Register a conversion function. A conversion function must be convertible
+  /// to any of the following forms(where `T` is a class derived from `Type`:
+  ///   * Optional<Type>(T)
+  ///     - This form represents a 1-1 type conversion. It should return nullptr
+  ///       or `llvm::None` to signify failure. If `llvm::None` is returned, the
+  ///       converter is allowed to try another conversion function to perform
+  ///       the conversion.
+  ///   * Optional<LogicalResult>(T, SmallVectorImpl<Type> &)
+  ///     - This form represents a 1-N type conversion. It should return
+  ///       `failure` or `llvm::None` to signify a failed conversion. If the new
+  ///       set of types is empty, the type is removed and any usages of the
+  ///       existing value are expected to be removed during conversion. If
+  ///       `llvm::None` is returned, the converter is allowed to try another
+  ///       conversion function to perform the conversion.
+  /// Note: When attempting to convert a type, e.g. via 'convertType', the
+  ///       mostly recently added conversions will be invoked first.
+  template <typename FnT,
+            typename T = typename FunctionTraits<FnT>::template arg_t<0>>
+  void addConversion(FnT &&callback) {
+    registerConversion(wrapCallback<T>(std::forward<FnT>(callback)));
+  }
+
+  /// Convert the given type. This function should return failure if no valid
+  /// conversion exists, success otherwise. If the new set of types is empty,
+  /// the type is removed and any usages of the existing value are expected to
+  /// be removed during conversion.
+  LogicalResult convertType(Type t, SmallVectorImpl<Type> &results);
 
   /// This hook simplifies defining 1-1 type conversions. This function returns
   /// the type to convert to on success, and a null type on failure.
-  virtual Type convertType(Type t) { return t; }
+  Type convertType(Type t);
 
   /// Convert the given set of types, filling 'results' as necessary. This
   /// 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
public:
   /// results in more than one type, or if a type conversion may persist after
   /// the conversion has finished.
   virtual Operation *materializeConversion(PatternRewriter &rewriter,
                                            Type resultType,
                                            ArrayRef<Value> inputs,
                                            Location loc) {
     llvm_unreachable("expected 'materializeConversion' to be overridden");
   }
+
+private:
+  /// The signature of the callback used to convert a type. If the new set of
+  /// types is empty, the type is removed and any usages of the existing value
+  /// are expected to be removed during conversion.
+  using ConversionCallbackFn =
+      std::function<Optional<LogicalResult>(Type, SmallVectorImpl<Type> &)>;
+
+  /// Generate a wrapper for the given callback. This allows for accepting
+  /// different callback forms, that all compose into a single version.
+  /// With callback of form: `Optional<Type>(T)`
+  template <typename T, typename FnT>
+  std::enable_if_t<is_invocable<FnT, T>::value, ConversionCallbackFn>
+  wrapCallback(FnT &&callback) {
+    return wrapCallback<T>([=](T type, SmallVectorImpl<Type> &results) {
+      if (Optional<Type> resultOpt = callback(type)) {
+        bool wasSuccess = static_cast<bool>(resultOpt.getValue());
+        if (wasSuccess)
+          results.push_back(resultOpt.getValue());
+        return Optional<LogicalResult>(success(wasSuccess));
+      }
+      return Optional<LogicalResult>();
+    });
+  }
+  /// With callback of form: `Optional<LogicalResult>(T, SmallVectorImpl<> &)`
+  template <typename T, typename FnT>
+  std::enable_if_t<!is_invocable<FnT, T>::value, ConversionCallbackFn>
+  wrapCallback(FnT &&callback) {
+    return [=](Type type,
+               SmallVectorImpl<Type> &results) -> Optional<LogicalResult> {
+      T derivedType = type.dyn_cast<T>();
+      if (!derivedType)
+        return llvm::None;
+      return callback(derivedType, results);
+    };
+  }
+
+  /// Register a type conversion.
+  void registerConversion(ConversionCallbackFn callback) {
+    conversions.emplace_back(std::move(callback));
+  }
+
+  /// The set of registered conversion functions.
+  SmallVector<ConversionCallbackFn, 4> conversions;
 };
 
 //===----------------------------------------------------------------------===//
 // 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/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

 
 #include "../GPUCommon/IndexIntrinsicsOpLowering.h"
 #include "../GPUCommon/OpToFuncCallLowering.h"
 
 using namespace mlir;
 
 namespace {
 
-/// Derived type converter for GPU to NVVM lowering. The GPU dialect uses memory
-/// space 5 for private memory attributions, but NVVM represents private
-/// memory allocations as local `alloca`s in the default address space. This
-/// converter drops the private memory space to support the use case above.
-class NVVMTypeConverter : public LLVMTypeConverter {
-public:
-  using LLVMTypeConverter::LLVMTypeConverter;
-
-  Type convertType(Type type) override {
-    auto memref = type.dyn_cast<MemRefType>();
-    if (memref &&
-        memref.getMemorySpace() == gpu::GPUDialect::getPrivateAddressSpace()) {
-      type = MemRefType::Builder(memref).setMemorySpace(0);
-    }
-
-    return LLVMTypeConverter::convertType(type);
-  }
-};
-
 /// Converts all_reduce op to LLVM/NVVM ops.
 struct GPUAllReduceOpLowering : public ConvertToLLVMPattern {
   using AccumulatorFactory =
       std::function<Value(Location, Value, Value, ConversionPatternRewriter &)>;
 
   explicit GPUAllReduceOpLowering(LLVMTypeConverter &lowering_)
       : ConvertToLLVMPattern(gpu::AllReduceOp::getOperationName(),
                              lowering_.getDialect()->getContext(), lowering_),

 ///
 /// This pass only handles device code and is not meant to be run on GPU host
 /// code.
 class LowerGpuOpsToNVVMOpsPass
     : public OperationPass<LowerGpuOpsToNVVMOpsPass, gpu::GPUModuleOp> {
 public:
   void runOnOperation() override {
     gpu::GPUModuleOp m = getOperation();
+
+    /// MemRef conversion for GPU to NVVM lowering. The GPU dialect uses memory
+    /// space 5 for private memory attributions, but NVVM represents private
+    /// memory allocations as local `alloca`s in the default address space. This
+    /// converter drops the private memory space to support the use case above.
+    LLVMTypeConverter converter(m.getContext());
+    converter.addConversion([&](MemRefType type) -> Optional<Type> {
+      if (type.getMemorySpace() != gpu::GPUDialect::getPrivateAddressSpace())
+        return llvm::None;
+      return converter.convertType(MemRefType::Builder(type).setMemorySpace(0));
+    });
+
     OwningRewritePatternList patterns;
-    NVVMTypeConverter converter(m.getContext());
     populateStdToLLVMConversionPatterns(converter, patterns);
     populateGpuToNVVMConversionPatterns(converter, patterns);
     ConversionTarget target(getContext());
     target.addIllegalDialect<gpu::GPUDialect>();
     target.addIllegalOp<LLVM::FAbsOp, LLVM::FCeilOp, LLVM::CosOp,
                         LLVM::ExpOp>();
     target.addIllegalOp<FuncOp>();
     target.addLegalDialect<LLVM::LLVMDialect>();

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

 template <typename T>
 static LLVMType getPtrToElementType(T containerType,
                                     LLVMTypeConverter &lowering) {
   return lowering.convertType(containerType.getElementType())
       .template cast<LLVMType>()
       .getPointerTo();
 }
 
-// Convert the given type to the LLVM IR Dialect type.  The following
-// conversions are supported:
-//   - an Index type is converted into an LLVM integer type with pointer
-//     bitwidth (analogous to intptr_t in C);
-//   - an Integer type is converted into an LLVM integer type of the same width;
-//   - an F32 type is converted into an LLVM float type
-//   - a Buffer, Range or View is converted into an LLVM structure type
-//     containing the respective dynamic values.
-static Type convertLinalgType(Type t, LLVMTypeConverter &lowering) {
+/// Convert the given range descriptor type to the LLVMIR dialect.
+/// Range descriptor contains the range bounds and the step as 64-bit integers.
+///
+/// struct {
+///   int64_t min;
+///   int64_t max;
+///   int64_t step;
+/// };
+static Type convertRangeType(RangeType t, LLVMTypeConverter &converter) {
   auto *context = t.getContext();
-  auto int64Ty = lowering.convertType(IntegerType::get(64, context))
+  auto int64Ty = converter.convertType(IntegerType::get(64, context))
                      .cast<LLVM::LLVMType>();
-
-  // Range descriptor contains the range bounds and the step as 64-bit integers.
-  //
-  // struct {
-  //   int64_t min;
-  //   int64_t max;
-  //   int64_t step;
-  // };
-  if (t.isa<RangeType>())
-    return LLVMType::getStructTy(int64Ty, int64Ty, int64Ty);
-
-  return Type();
 }
 
 namespace {
 /// EDSC-compatible wrapper for MemRefDescriptor.
 class BaseViewConversionHelper {
 public:
   BaseViewConversionHelper(Type type)
       : d(MemRefDescriptor::undef(rewriter(), loc(), type)) {}
public:
   explicit RangeOpConversion(MLIRContext *context, LLVMTypeConverter &lowering_)
       : ConvertToLLVMPattern(RangeOp::getOperationName(), context, lowering_) {}
 
   PatternMatchResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto rangeOp = cast<RangeOp>(op);
     auto rangeDescriptorTy =
-        convertLinalgType(rangeOp.getResult().getType(), typeConverter);
+        convertRangeType(rangeOp.getType().cast<RangeType>(), typeConverter);
 
     edsc::ScopedContext context(rewriter, op->getLoc());
 
     // Fill in an aggregate value of the descriptor.
     RangeOpOperandAdaptor adaptor(operands);
     Value desc = llvm_undef(rangeDescriptorTy);
     desc = llvm_insertvalue(desc, adaptor.min(), rewriter.getI64ArrayAttr(0));
     desc = llvm_insertvalue(desc, adaptor.max(), rewriter.getI64ArrayAttr(1));
static FlatSymbolRefAttr getLibraryCallSymbolRef(Operation *op,
   // Insert before module terminator.
   rewriter.setInsertionPoint(module.getBody(),
                              std::prev(module.getBody()->end()));
   rewriter.create<FuncOp>(op->getLoc(), fnNameAttr.getValue(), libFnType,
                           ArrayRef<NamedAttribute>{});
   return fnNameAttr;
 }
 
-Type LinalgTypeConverter::convertType(Type t) {
-  if (auto result = LLVMTypeConverter::convertType(t))
-    return result;
-  return convertLinalgType(t, *this);
-}
-
 namespace {
 
 // LinalgOpConversion<LinalgOp> creates a new call to the
 // `LinalgOp::getLibraryCallName()` function.
 // The implementation of the function can be either in the same module or in an
 // externally linked library.
 template <typename LinalgOp>
 class LinalgOpConversion : public OpRewritePattern<LinalgOp> {
patterns.insert<CopyTransposeConversion, LinalgOpConversion<ConvOp>,
                   LinalgOpConversion<MatmulOp>, LinalgOpConversion<MatvecOp>>(
       ctx);
 }
 
 } // namespace
 
 /// Populate the given list with patterns that convert from Linalg to LLVM.
 void mlir::populateLinalgToLLVMConversionPatterns(
-    LinalgTypeConverter &converter, OwningRewritePatternList &patterns,
+    LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     MLIRContext *ctx) {
   patterns.insert<RangeOpConversion, ReshapeOpConversion, SliceOpConversion,
                   TransposeOpConversion, YieldOpConversion>(ctx, converter);
+
+  // Populate the type conversions for the linalg types.
+  converter.addConversion(
+      [&](RangeType type) { return convertRangeType(type, converter); });
 }
 
 namespace {
 struct ConvertLinalgToLLVMPass : public ModulePass<ConvertLinalgToLLVMPass> {
   void runOnModule() override;
 };
 } // namespace
 
 void ConvertLinalgToLLVMPass::runOnModule() {
   auto module = getModule();
 
   // Convert to the LLVM IR dialect using the converter defined above.
   OwningRewritePatternList patterns;
-  LinalgTypeConverter converter(&getContext());
+  LLVMTypeConverter converter(&getContext());
   populateAffineToStdConversionPatterns(patterns, &getContext());
   populateLoopToStdConversionPatterns(patterns, &getContext());
   populateStdToLLVMConversionPatterns(converter, patterns, /*useAlloca=*/false,
                                       /*emitCWrappers=*/true);
   populateVectorToLLVMConversionPatterns(converter, patterns);
   populateLinalgToStandardConversionPatterns(patterns, &getContext());
   populateLinalgToLLVMConversionPatterns(converter, patterns, &getContext());
 

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

 //===- ConvertStandardToLLVM.cpp - Standard to LLVM dialect conversion-----===//
 //
 // Part of the LLVM Project, 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 implements a pass to convert MLIR standard and builtin dialects
 // into the LLVM IR dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
-
 #include "mlir/ADT/TypeSwitch.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"

 
 /// Create an LLVMTypeConverter using 'custom' customizations.
 LLVMTypeConverter::LLVMTypeConverter(
     MLIRContext *ctx, const LLVMTypeConverterCustomization &customs)
     : llvmDialect(ctx->getRegisteredDialect<LLVM::LLVMDialect>()),
       customizations(customs) {
   assert(llvmDialect && "LLVM IR dialect is not registered");
   module = &llvmDialect->getLLVMModule();
+
+  // Register conversions for the standard types.
+  addConversion([&](FloatType type) { return convertFloatType(type); });
+  addConversion([&](FunctionType type) { return convertFunctionType(type); });
+  addConversion([&](IndexType type) { return convertIndexType(type); });
+  addConversion([&](IntegerType type) { return convertIntegerType(type); });
+  addConversion([&](MemRefType type) { return convertMemRefType(type); });
+  addConversion(
+      [&](UnrankedMemRefType type) { return convertUnrankedMemRefType(type); });
+  addConversion([&](VectorType type) { return convertVectorType(type); });
+
+  // LLVMType is legal, so add a pass-through conversion.
+  addConversion([](LLVM::LLVMType type) { return type; });
 }
 
 /// Get the LLVM context.
 llvm::LLVMContext &LLVMTypeConverter::getLLVMContext() {
   return module->getContext();
 }
 
 LLVM::LLVMType LLVMTypeConverter::getIndexType() {
Type LLVMTypeConverter::convertVectorType(VectorType type) {
   auto vectorType =
       LLVM::LLVMType::getVectorTy(elementType, type.getShape().back());
   auto shape = type.getShape();
   for (int i = shape.size() - 2; i >= 0; --i)
     vectorType = LLVM::LLVMType::getArrayTy(vectorType, shape[i]);
   return vectorType;
 }
 
-// Dispatch based on the actual type.  Return null type on error.
-Type LLVMTypeConverter::convertStandardType(Type t) {
-  return TypeSwitch<Type, Type>(t)
-      .Case([&](FloatType type) { return convertFloatType(type); })
-      .Case([&](FunctionType type) { return convertFunctionType(type); })
-      .Case([&](IndexType type) { return convertIndexType(type); })
-      .Case([&](IntegerType type) { return convertIntegerType(type); })
-      .Case([&](MemRefType type) { return convertMemRefType(type); })
-      .Case([&](UnrankedMemRefType type) {
-        return convertUnrankedMemRefType(type);
-      })
-      .Case([&](VectorType type) { return convertVectorType(type); })
-      .Case([](LLVM::LLVMType type) { return type; })
-      .Default([](Type) { return Type(); });
-}
-
 ConvertToLLVMPattern::ConvertToLLVMPattern(StringRef rootOpName,
                                            MLIRContext *context,
                                            LLVMTypeConverter &typeConverter_,
                                            PatternBenefit benefit)
     : ConversionPattern(rootOpName, benefit, context),
       typeConverter(typeConverter_) {}
 
 /*============================================================================*/

 void mlir::populateStdToLLVMBarePtrConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
     bool useAlloca) {
   populateStdToLLVMBarePtrFuncOpConversionPattern(converter, patterns);
   populateStdToLLVMNonMemoryConversionPatterns(converter, patterns);
   populateStdToLLVMMemoryConversionPatters(converter, patterns, useAlloca);
 }
 
-// Convert types using the stored LLVM IR module.
-Type LLVMTypeConverter::convertType(Type t) { return convertStandardType(t); }
-
 // Create an LLVM IR structure type if there is more than one result.
 Type LLVMTypeConverter::packFunctionResults(ArrayRef<Type> types) {
   assert(!types.empty() && "expected non-empty list of type");
 
   if (types.size() == 1)
     return convertType(types.front());
 
   SmallVector<LLVM::LLVMType, 8> resultTypes;

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

for (auto shape : tensorType.getShape()) {
       size *= shape;
     }
     return size;
   }
   // TODO: Add size computation for other types.
   return llvm::None;
 }
 
-static Type convertStdType(Type type) {
+SPIRVTypeConverter::SPIRVTypeConverter() {
+  addConversion([](Type type) -> Optional<Type> {
     // If the type is already valid in SPIR-V, directly return.
-  if (spirv::SPIRVDialect::isValidType(type)) {
-    return type;
-  }
-
-  if (auto indexType = type.dyn_cast<IndexType>()) {
-    return SPIRVTypeConverter::getIndexType(type.getContext());
-  }
-
-  if (auto memRefType = type.dyn_cast<MemRefType>()) {
+    return spirv::SPIRVDialect::isValidType(type) ? type : Optional<Type>();
+  });
+  addConversion([](IndexType indexType) {
+    return SPIRVTypeConverter::getIndexType(indexType.getContext());
+  });
+  addConversion([this](MemRefType memRefType) -> Type {
     // TODO(ravishankarm): For now only support default memory space. The memory
     // space description is not set is stone within MLIR, i.e. it depends on the
     // context it is being used. To map this to SPIR-V storage classes, we
     // should rely on the ABI attributes, and not on the memory space. This is
     // still evolving, and needs to be revisited when there is more clarity.
-    if (memRefType.getMemorySpace()) {
+    if (memRefType.getMemorySpace())
       return Type();
-    }
 
-    auto elementType = convertStdType(memRefType.getElementType());
-    if (!elementType) {
+    auto elementType = convertType(memRefType.getElementType());
+    if (!elementType)
       return Type();
-    }
 
     auto elementSize = getTypeNumBytes(elementType);
-    if (!elementSize) {
+    if (!elementSize)
       return Type();
-    }
+
     // TODO(ravishankarm) : Handle dynamic shapes.
     if (memRefType.hasStaticShape()) {
       auto arraySize = getTypeNumBytes(memRefType);
-      if (!arraySize) {
+      if (!arraySize)
         return Type();
-      }
+
       auto arrayType = spirv::ArrayType::get(
           elementType, arraySize.getValue() / elementSize.getValue(),
           elementSize.getValue());
       auto structType = spirv::StructType::get(arrayType, 0);
       // For now initialize the storage class to StorageBuffer. This will be
       // updated later based on whats passed in w.r.t to the ABI attributes.
       return spirv::PointerType::get(structType,
                                      spirv::StorageClass::StorageBuffer);
     }
-  }
-
-  if (auto tensorType = type.dyn_cast<TensorType>()) {
+    return Type();
+  });
+  addConversion([this](TensorType tensorType) -> Type {
     // TODO(ravishankarm) : Handle dynamic shapes.
-    if (!tensorType.hasStaticShape()) {
+    if (!tensorType.hasStaticShape())
       return Type();
-    }
-    auto elementType = convertStdType(tensorType.getElementType());
-    if (!elementType) {
+
+    auto elementType = convertType(tensorType.getElementType());
+    if (!elementType)
       return Type();
-    }
+
     auto elementSize = getTypeNumBytes(elementType);
-    if (!elementSize) {
+    if (!elementSize)
       return Type();
-    }
+
     auto tensorSize = getTypeNumBytes(tensorType);
-    if (!tensorSize) {
+    if (!tensorSize)
       return Type();
-    }
+
     return spirv::ArrayType::get(elementType,
                                  tensorSize.getValue() / elementSize.getValue(),
                                  elementSize.getValue());
-  }
-  return Type();
-}
-
-Type SPIRVTypeConverter::convertType(Type type) { return convertStdType(type); }
+  });
 
 //===----------------------------------------------------------------------===//
 // FuncOp Conversion Patterns
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A pattern for rewriting function signature to convert arguments of functions
 /// to be of valid SPIR-V types.

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) {
-  if (auto newT = convertType(t)) {
-    results.push_back(newT);
-    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;
   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;
+  if (failed(convertType(t, results)))
+    return nullptr;
+
+  // Check to ensure that only one type was produced.
+  return results.size() == 1 ? results.front() : nullptr;
+}
+
 /// Convert the given set of types, filling 'results' as necessary. This
 /// returns failure if the conversion of any of the types fails, success
 /// otherwise.
 LogicalResult TypeConverter::convertTypes(ArrayRef<Type> types,
                                           SmallVectorImpl<Type> &results) {
   for (auto type : types)
     if (failed(convertType(type, results)))
       return failure();

mlir/test/lib/TestDialect/TestPatterns.cpp

PatternMatchResult matchAndRewrite(Operation *op,
     return matchSuccess();
   }
 };
 } // namespace
 
 namespace {
 struct TestTypeConverter : public TypeConverter {
   using TypeConverter::TypeConverter;
+  TestTypeConverter() { addConversion(convertType); }
 
-  LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) override {
+  static LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) {
     // Drop I16 types.
     if (t.isInteger(16))
       return success();
 
     // Convert I64 to F64.
     if (t.isInteger(64)) {
       results.push_back(FloatType::getF64(t.getContext()));
       return success();