[draft] Dialect Conversion without Rollback

This commit adds a dialect conversion driver without rollback: `OneShotDialectConversionDriver`

The new driver reuses some functionality of the greedy pattern rewrite driver. Just a proof of concept, code is not polished yet.

`OneShotConversionPatternRewriter` is a rewriter that materializes all IR changes immediately.

Author: matthias-springer

mlir/include/mlir/Transforms/DialectConversion.h

@@ -247,7 +247,8 @@ class TypeConverter {
   /// Attempts a 1-1 type conversion, expecting the result type to be
   /// `TargetType`. Returns the converted type cast to `TargetType` on success,
   /// and a null type on conversion or cast failure.
-  template <typename TargetType> TargetType convertType(Type t) const {
+  template <typename TargetType>
+  TargetType convertType(Type t) const {
     return dyn_cast_or_null<TargetType>(convertType(t));
   }
 
@@ -657,7 +658,7 @@ struct ConversionPatternRewriterImpl;
 /// This class implements a pattern rewriter for use with ConversionPatterns. It
 /// extends the base PatternRewriter and provides special conversion specific
 /// hooks.
-class ConversionPatternRewriter final : public PatternRewriter {
+class ConversionPatternRewriter : public PatternRewriter {
 public:
   ~ConversionPatternRewriter() override;
 
@@ -711,6 +712,17 @@ class ConversionPatternRewriter final : public PatternRewriter {
   LogicalResult getRemappedValues(ValueRange keys,
                                   SmallVectorImpl<Value> &results);
 
+  virtual void setCurrentTypeConverter(const TypeConverter *converter);
+
+  virtual const TypeConverter *getCurrentTypeConverter() const;
+
+  /// Populate the operands that are used for constructing the adapter into
+  /// `remapped`.
+  virtual LogicalResult getAdapterOperands(StringRef valueDiagTag,
+                                           std::optional<Location> inputLoc,
+                                           ValueRange values,
+                                           SmallVector<Value> &remapped);
+
   //===--------------------------------------------------------------------===//
   // PatternRewriter Hooks
   //===--------------------------------------------------------------------===//
@@ -755,6 +767,14 @@ class ConversionPatternRewriter final : public PatternRewriter {
   /// Return a reference to the internal implementation.
   detail::ConversionPatternRewriterImpl &getImpl();
 
+protected:
+  /// Protected constructor for `OneShotConversionPatternRewriter`. Does not
+  /// initialize `impl`.
+  explicit ConversionPatternRewriter(MLIRContext *ctx);
+
+  // Hide unsupported pattern rewriter API.
+  using OpBuilder::setListener;
+
 private:
   // Allow OperationConverter to construct new rewriters.
   friend struct OperationConverter;
@@ -765,9 +785,6 @@ class ConversionPatternRewriter final : public PatternRewriter {
   explicit ConversionPatternRewriter(MLIRContext *ctx,
                                      const ConversionConfig &config);
 
-  // Hide unsupported pattern rewriter API.
-  using OpBuilder::setListener;
-
   std::unique_ptr<detail::ConversionPatternRewriterImpl> impl;
 };
 

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

@@ -18,6 +18,8 @@
 
 namespace mlir {
 
+class ConversionTarget;
+
 /// This enum controls which ops are put on the worklist during a greedy
 /// pattern rewrite.
 enum class GreedyRewriteStrictness {
@@ -188,6 +190,10 @@ applyOpPatternsAndFold(ArrayRef<Operation *> ops,
                        GreedyRewriteConfig config = GreedyRewriteConfig(),
                        bool *changed = nullptr, bool *allErased = nullptr);
 
+LogicalResult
+applyPartialOneShotConversion(Operation *op, const ConversionTarget &target,
+                              const FrozenRewritePatternSet &patterns);
+
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_

mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp

@@ -15,6 +15,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include <memory>
 #include <type_traits>
 
@@ -1346,8 +1347,8 @@ void ConvertComplexToStandardPass::runOnOperation() {
   ConversionTarget target(getContext());
   target.addLegalDialect<arith::ArithDialect, math::MathDialect>();
   target.addLegalOp<complex::CreateOp, complex::ImOp, complex::ReOp>();
-  if (failed(
-          applyPartialConversion(getOperation(), target, std::move(patterns))))
+  if (failed(applyPartialOneShotConversion(getOperation(), target,
+                                           std::move(patterns))))
     signalPassFailure();
 }
 } // namespace

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

@@ -25,6 +25,7 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
@@ -475,9 +476,10 @@ struct ConvertNVGPUToNVVMPass
     target.addLegalDialect<::mlir::NVVM::NVVMDialect>();
     mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
         converter, patterns, target);
-    if (failed(applyPartialConversion(getOperation(), target,
-                                      std::move(patterns))))
+    if (failed(applyPartialOneShotConversion(getOperation(), target,
+                                             std::move(patterns))))
       signalPassFailure();
+    // applyPartialConversion
   }
 };
 

mlir/lib/Transforms/Utils/DialectConversion.cpp

@@ -1633,6 +1633,9 @@ ConversionPatternRewriter::ConversionPatternRewriter(
   setListener(impl.get());
 }
 
+ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx)
+    : PatternRewriter(ctx), impl(nullptr) {}
+
 ConversionPatternRewriter::~ConversionPatternRewriter() = default;
 
 void ConversionPatternRewriter::replaceOp(Operation *op, Operation *newOp) {
@@ -1819,6 +1822,22 @@ detail::ConversionPatternRewriterImpl &ConversionPatternRewriter::getImpl() {
   return *impl;
 }
 
+void ConversionPatternRewriter::setCurrentTypeConverter(
+    const TypeConverter *converter) {
+  impl->currentTypeConverter = converter;
+}
+
+const TypeConverter *
+ConversionPatternRewriter::getCurrentTypeConverter() const {
+  return impl->currentTypeConverter;
+}
+
+LogicalResult ConversionPatternRewriter::getAdapterOperands(
+    StringRef valueDiagTag, std::optional<Location> inputLoc, ValueRange values,
+    SmallVector<Value> &remapped) {
+  return impl->remapValues(valueDiagTag, inputLoc, *this, values, remapped);
+}
+
 //===----------------------------------------------------------------------===//
 // ConversionPattern
 //===----------------------------------------------------------------------===//
@@ -1827,16 +1846,18 @@ LogicalResult
 ConversionPattern::matchAndRewrite(Operation *op,
                                    PatternRewriter &rewriter) const {
   auto &dialectRewriter = static_cast<ConversionPatternRewriter &>(rewriter);
-  auto &rewriterImpl = dialectRewriter.getImpl();
 
   // Track the current conversion pattern type converter in the rewriter.
-  llvm::SaveAndRestore currentConverterGuard(rewriterImpl.currentTypeConverter,
-                                             getTypeConverter());
+  const TypeConverter *currentTypeConverter =
+      dialectRewriter.getCurrentTypeConverter();
+  auto resetTypeConverter = llvm::make_scope_exit(
+      [&] { dialectRewriter.setCurrentTypeConverter(currentTypeConverter); });
+  dialectRewriter.setCurrentTypeConverter(getTypeConverter());
 
   // Remap the operands of the operation.
-  SmallVector<Value, 4> operands;
-  if (failed(rewriterImpl.remapValues("operand", op->getLoc(), rewriter,
-                                      op->getOperands(), operands))) {
+  SmallVector<Value> operands;
+  if (failed(dialectRewriter.getAdapterOperands("operand", op->getLoc(),
+                                                op->getOperands(), operands))) {
     return failure();
   }
   return matchAndRewrite(op, operands, dialectRewriter);