[mlir] support materialization for 1-1 type conversions

[ftynse]

Summary

Dialect conversion infrastructure supports 1->N type conversions by requiring
individual conversions to provide facilities to generate operations
retrofitting N values into 1 of the original type when N > 1. This
functionality can also be used to materialize explicit "cast"-like operations,
but it did not support 1->1 type conversions until now. Modify TypeConverter to
support materialization of cast operations for 1-1 conversions.

This also makes materialization specification more extensible following the
same pattern as type conversions. Instead of overloading a virtual function,
users or subclasses of TypeConversion can now register type-specific
materialization callbacks that will be called in order for the given type.

[ftynse]

Clean up and prepare for review

[rriddle]

(Sorry for the delay)

Looks really clean, thanks!

mlir/include/mlir/Transforms/DialectConversion.h

116

typo: convertibe

117

nit: form:

162

Can we remove the virtual here? That should enable removing the vtable completely.

mlir/lib/Transforms/DialectConversion.cpp

319

Are these duplicated from below?

387–396

nit: Can you just use an if for the initialization?

Value newArg;
if (typeConverter)
  newArg = ...

[ftynse]

rebase

[ftynse]

(Sorry for the delay)

No worries! Flang folks seemed to need it, but they did not insist in the end.

mlir/include/mlir/Transforms/DialectConversion.h

162

I don't see any in-tree overrides of this, but let's do it as a separate commit in case somebody depends on it.

mlir/lib/Transforms/DialectConversion.cpp

319

Indeed, removed.

mlir/docs/DialectConversion.md

public:
   ///       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
+  /// Register a materialization function, which must be convertibe to the
+  /// following form
+  ///   `Optional<Value>(PatternRewriter &, T, ValueRange, Location)`,
+  /// where `T` is any subclass of `Type`. This function is responsible for
+  /// creating an operation, using the PatternRewriter and Location provided,
+  /// that "casts" a range of values into a single value of the given type `T`.
+  /// It must return a Value of the converted type on success, an `llvm::None`
+  /// if it failed but other materialization can be attempted, and `nullptr` on
+  /// unrecoverable failure. It will only be called for (sub)types of `T`.
+  /// Materialization functions must be provided 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,
-                                           Type resultType,
-                                           ArrayRef<Value> inputs,
-                                           Location loc);
+  template <typename FnT>
+  void addMaterialization(FnT &&callback);
 };
 ```
 
 ### Conversion Patterns
 
 When type conversion comes into play, the general Rewrite Patterns can no longer
 be used. This is due to the fact that the operands of the operation being
 matched will not correspond with the operands of the correct type as determined

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

public:
   /// C/C++ ABI lowering related to struct argument passing.
   Value promoteOneMemRefDescriptor(Location loc, Value operand,
                                    OpBuilder &builder);
 
   /// Converts the function type to a C-compatible format, in particular using
   /// pointers to memref descriptors for arguments.
   LLVM::LLVMType convertFunctionTypeCWrapper(FunctionType type);
 
-  /// Creates descriptor structs from individual values constituting them.
-  Operation *materializeConversion(PatternRewriter &rewriter, Type type,
-                                   ArrayRef<Value> values,
-                                   Location loc) override;
-
   /// Gets the LLVM representation of the index type. The returned type is an
   /// integer type with the size configured for this type converter.
   LLVM::LLVMType getIndexType();
 
   /// Gets the bitwidth of the index type when converted to LLVM.
   unsigned getIndexTypeBitwidth() { return customizations.indexBitwidth; }
 
 protected:

mlir/include/mlir/Transforms/DialectConversion.h

public:
   /// 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 llvm::function_traits<FnT>::template arg_t<0>>
   void addConversion(FnT &&callback) {
     registerConversion(wrapCallback<T>(std::forward<FnT>(callback)));
   }
 
+  /// Register a materialization function, which must be convertible to the
+  /// following form:
+  ///   `Optional<Value>(PatternRewriter &, T, ValueRange, Location)`,
+  /// where `T` is any subclass of `Type`. This function is responsible for
+  /// creating an operation, using the PatternRewriter and Location provided,
+  /// that "casts" a range of values into a single value of the given type `T`.
+  /// It must return a Value of the converted type on success, an `llvm::None`
+  /// if it failed but other materialization can be attempted, and `nullptr` on
+  /// unrecoverable failure. It will only be called for (sub)types of `T`.
+  /// Materialization functions must be provided when a type conversion
+  /// results in more than one type, or if a type conversion may persist after
+  /// the conversion has finished.
+  template <typename FnT,
+            typename T = typename llvm::function_traits<FnT>::template arg_t<1>>
+  void addMaterialization(FnT &&callback) {
+    registerMaterialization(
+        wrapMaterialization<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.
public:
 
   /// Return true if the inputs and outputs of the given function type are
   /// legal.
   bool isSignatureLegal(FunctionType funcType);
 
   /// This hook allows for converting a specific argument of a signature. It
   /// takes as inputs the original argument input number, type.
   /// On success, this function should populate 'result' with any new mappings.
   virtual LogicalResult convertSignatureArg(unsigned inputNo, Type type,
                                             SignatureConversion &result);
 
   /// 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);
 
-  /// 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,
-                                           Type resultType,
-                                           ArrayRef<Value> inputs,
-                                           Location loc) {
-    llvm_unreachable("expected 'materializeConversion' to be overridden");
-  }
+  /// Materialize a conversion from a set of types into one result type by
+  /// generating a cast operation of some kind.
+  Value materializeConversion(PatternRewriter &rewriter, Location loc,
+                              Type resultType, ValueRange inputs);
 
 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> &)>;
 
+  using MaterializationCallbackFn = std::function<Optional<Value>(
+      PatternRewriter &, Type, ValueRange, Location)>;
+
   /// 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<llvm::is_invocable<FnT, T>::value, ConversionCallbackFn>
   wrapCallback(FnT &&callback) {
     return wrapCallback<T>([callback = std::forward<FnT>(callback)](
                                T type, SmallVectorImpl<Type> &results) {
wrapCallback(FnT &&callback) {
     };
   }
 
   /// Register a type conversion.
   void registerConversion(ConversionCallbackFn callback) {
     conversions.emplace_back(std::move(callback));
   }
 
+  /// 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)](
+               PatternRewriter &rewriter, Type resultType, ValueRange inputs,
+               Location loc) -> Optional<Value> {
+      if (T derivedType = resultType.dyn_cast<T>())
+        return callback(rewriter, derivedType, inputs, loc);
+      return llvm::None;
+    };
+  }
+
+  /// Register a materialization.
+  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;
 };
 
 //===----------------------------------------------------------------------===//
 // 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/StandardToLLVM/StandardToLLVM.cpp

LLVMTypeConverter::LLVMTypeConverter(
   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; });
+
+  // Materialization for memrefs creates descriptor structs from individual
+  // values constituting them, when descriptors are used, i.e. more than one
+  // value represents a memref.
+  addMaterialization([&](PatternRewriter &rewriter,
+                         UnrankedMemRefType resultType, ValueRange inputs,
+                         Location loc) -> Optional<Value> {
+    if (inputs.size() == 1)
+      return llvm::None;
+    return UnrankedMemRefDescriptor::pack(rewriter, loc, *this, resultType,
+                                          inputs);
+  });
+  addMaterialization([&](PatternRewriter &rewriter, MemRefType resultType,
+                         ValueRange inputs, Location loc) -> Optional<Value> {
+    if (inputs.size() == 1)
+      return llvm::None;
+    return MemRefDescriptor::pack(rewriter, loc, *this, resultType, inputs);
+  });
 }
 
 /// Returns the MLIR context.
 MLIRContext &LLVMTypeConverter::getContext() {
   return *getDialect()->getContext();
 }
 
 /// Get the LLVM context.
LLVM::LLVMType resultType =
           ? LLVM::LLVMType::getVoidTy(llvmDialect)
           : unwrap(packFunctionResults(type.getResults()));
   if (!resultType)
     return {};
 
   return LLVM::LLVMType::getFunctionTy(resultType, inputs, false);
 }
 
-/// Creates descriptor structs from individual values constituting them.
-Operation *LLVMTypeConverter::materializeConversion(PatternRewriter &rewriter,
-                                                    Type type,
-                                                    ArrayRef<Value> values,
-                                                    Location loc) {
-  if (auto unrankedMemRefType = type.dyn_cast<UnrankedMemRefType>())
-    return UnrankedMemRefDescriptor::pack(rewriter, loc, *this,
-                                          unrankedMemRefType, values)
-        .getDefiningOp();
-
-  auto memRefType = type.dyn_cast<MemRefType>();
-  assert(memRefType && "1->N conversion is only supported for memrefs");
-  return MemRefDescriptor::pack(rewriter, loc, *this, memRefType, values)
-      .getDefiningOp();
-}
-
 // Convert a MemRef to an LLVM type. The result is a MemRef descriptor which
 // contains:
 //   1. the pointer to the data buffer, followed by
 //   2.  a lowered `index`-type integer containing the distance between the
 //   beginning of the buffer and the first element to be accessed through the
 //   view, followed by
 //   3. an array containing as many `index`-type integers as the rank of the
 //   MemRef: the array represents the size, in number of elements, of the memref

mlir/lib/Transforms/DialectConversion.cpp

for (unsigned i = 0, e = origBlock->getNumArguments(); i != e; ++i) {
           origArg.replaceAllUsesWith(argReplacementValue);
           continue;
         }
         // If there are any dangling uses then replace the argument with one
         // generated by the type converter. This is necessary as the cast must
         // persist in the IR after conversion.
         if (!origArg.use_empty()) {
           rewriter.setInsertionPointToStart(newBlock);
-          auto *newOp = typeConverter->materializeConversion(
-              rewriter, origArg.getType(), llvm::None, loc);
-          origArg.replaceAllUsesWith(newOp->getResult(0));
+          Value newArg = typeConverter->materializeConversion(
+              rewriter, loc, origArg.getType(), llvm::None);
+          assert(newArg &&
+                 "Couldn't materialize a block argument after 1->0 conversion");
+          origArg.replaceAllUsesWith(newArg);
         }
         continue;
       }
 
-      // If mapping is 1-1, replace the remaining uses and drop the cast
-      // operation.
-      // FIXME(riverriddle) This should check that the result type and operand
-      // type are the same, otherwise it should force a conversion to be
-      // materialized.
-      if (argInfo->newArgSize == 1) {
-        origArg.replaceAllUsesWith(
-            mapping.lookupOrDefault(newBlock->getArgument(argInfo->newArgIdx)));
-        continue;
-      }
-
-      // Otherwise this is a 1->N value mapping.
+      // Otherwise this is a 1->1+ value mapping.
       Value castValue = argInfo->castValue;
-      assert(argInfo->newArgSize > 1 && castValue && "expected 1->N mapping");
+      assert(argInfo->newArgSize >= 1 && castValue && "expected 1->1+ mapping");
 
       // If the argument is still used, replace it with the generated cast.
       if (!origArg.use_empty())
         origArg.replaceAllUsesWith(mapping.lookupOrDefault(castValue));
 
       // If all users of the cast were removed, we can drop it. Otherwise, keep
       // the operation alive and let the user handle any remaining usages.
-      if (castValue.use_empty())
+      if (castValue.use_empty() && castValue.getDefiningOp())
         castValue.getDefiningOp()->erase();
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Conversion
 
for (unsigned i = 0; i != origArgCount; ++i) {
     if (inputMap->replacementValue) {
       assert(inputMap->size == 0 &&
              "invalid to provide a replacement value when the argument isn't "
              "dropped");
       mapping.map(origArg, inputMap->replacementValue);
       continue;
     }
 
-    // If this is a 1->1 mapping, then map the argument directly.
-    if (inputMap->size == 1) {
-      mapping.map(origArg, newArgs[inputMap->inputNo]);
-      info.argInfo[i] = ConvertedArgInfo(inputMap->inputNo, inputMap->size);
-      continue;
-    }
-
-    // Otherwise, this is a 1->N mapping. Call into the provided type converter
-    // to pack the new values.
+    // Otherwise, this is a 1->1+ mapping. Call into the provided type converter
+    // to pack the new values. For 1->1 mappings, if there is no materialization
+    // provided, use the argument directly instead.
     auto replArgs = newArgs.slice(inputMap->inputNo, inputMap->size);
-    Operation *cast = typeConverter->materializeConversion(
-        rewriter, origArg.getType(), replArgs, loc);
-    assert(cast->getNumResults() == 1);
-    mapping.map(origArg, cast->getResult(0));
+    Value newArg;
+    if (typeConverter)
+      newArg = typeConverter->materializeConversion(
+          rewriter, loc, origArg.getType(), replArgs);
+    if (!newArg) {
+      assert(replArgs.size() == 1 &&
+             "couldn't materialize the result of 1->N conversion");
+      newArg = replArgs.front();
+    }
+    mapping.map(origArg, newArg);
     info.argInfo[i] =
-        ConvertedArgInfo(inputMap->inputNo, inputMap->size, cast->getResult(0));
+        ConvertedArgInfo(inputMap->inputNo, inputMap->size, newArg);
   }
 
   // Remove the original block from the region and return the new one.
   insertConversion(newBlock, std::move(info));
   return newBlock;
 }
 
 void ArgConverter::insertConversion(Block *newBlock,
LogicalResult TypeConverter::convertSignatureArg(unsigned inputNo, Type type,
   if (convertedTypes.empty())
     return success();
 
   // Otherwise, add the new inputs.
   result.addInputs(inputNo, convertedTypes);
   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;
+}
+
 /// 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.

mlir/test/Transforms/test-legalizer.mlir

 // CHECK-LABEL: func @remap_input_1_to_N_remaining_use(%arg0: f16, %arg1: f16)
 func @remap_input_1_to_N_remaining_use(%arg0: f32) {
   // CHECK-NEXT: [[CAST:%.*]] = "test.cast"(%arg0, %arg1) : (f16, f16) -> f32
   // CHECK-NEXT: "work"([[CAST]]) : (f32) -> ()
   // expected-remark@+1 {{op 'work' is not legalizable}}
   "work"(%arg0) : (f32) -> ()
 }
 
+// CHECK-LABEL: func @remap_materialize_1_to_1(%{{.*}}: i43)
+func @remap_materialize_1_to_1(%arg0: i42) {
+  // CHECK: %[[V:.*]] = "test.cast"(%arg0) : (i43) -> i42
+  // CHECK: "test.return"(%[[V]])
+  "test.return"(%arg0) : (i42) -> ()
+}
+
 // CHECK-LABEL: func @remap_input_to_self
 func @remap_input_to_self(%arg0: index) {
   // CHECK-NOT: test.cast
   // CHECK: "work"
   // expected-remark@+1 {{op 'work' is not legalizable}}
   "work"(%arg0) : (index) -> ()
 }
 

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

LogicalResult matchAndRewrite(IllegalOpWithRegionAnchor op,
     return success();
   };
 };
 } // namespace
 
 namespace {
 struct TestTypeConverter : public TypeConverter {
   using TypeConverter::TypeConverter;
-  TestTypeConverter() { addConversion(convertType); }
+  TestTypeConverter() {
+    addConversion(convertType);
+    addMaterialization(materializeCast);
+    addMaterialization(materializeOneToOneCast);
+  }
 
   static LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) {
     // Drop I16 types.
     if (t.isSignlessInteger(16))
       return success();
 
     // Convert I64 to F64.
     if (t.isSignlessInteger(64)) {
       results.push_back(FloatType::getF64(t.getContext()));
       return success();
     }
 
+    // Convert I42 to I43.
+    if (t.isInteger(42)) {
+      results.push_back(IntegerType::get(43, t.getContext()));
+      return success();
+    }
+
     // Split F32 into F16,F16.
     if (t.isF32()) {
       results.assign(2, FloatType::getF16(t.getContext()));
       return success();
     }
 
     // Otherwise, convert the type directly.
     results.push_back(t);
     return success();
   }
 
-  /// Override the hook to materialize a conversion. This is necessary because
-  /// we generate 1->N type mappings.
-  Operation *materializeConversion(PatternRewriter &rewriter, Type resultType,
-                                   ArrayRef<Value> inputs,
-                                   Location loc) override {
-    return rewriter.create<TestCastOp>(loc, resultType, inputs);
+  /// Hook for materializing a conversion. This is necessary because we generate
+  /// 1->N type mappings.
+  static Optional<Value> materializeCast(PatternRewriter &rewriter,
+                                         Type resultType, ValueRange inputs,
+                                         Location loc) {
+    if (inputs.size() == 1)
+      return inputs[0];
+    return rewriter.create<TestCastOp>(loc, resultType, inputs).getResult();
+  }
+
+  /// Materialize the cast for one-to-one conversion from i64 to f64.
+  static Optional<Value> materializeOneToOneCast(PatternRewriter &rewriter,
+                                                 IntegerType resultType,
+                                                 ValueRange inputs,
+                                                 Location loc) {
+    if (resultType.getWidth() == 42 && inputs.size() == 1)
+      return rewriter.create<TestCastOp>(loc, resultType, inputs).getResult();
+    return llvm::None;
   }
 };
 
 struct TestLegalizePatternDriver
     : public PassWrapper<TestLegalizePatternDriver, OperationPass<ModuleOp>> {
   /// The mode of conversion to use with the driver.
   enum class ConversionMode { Analysis, Full, Partial };