[mlir][scf] refactor scf structuralOpConversion to better support 1:N type conversion

[Peiming]

Summary

This patch moves the 1:N type mapping into its own classes to allow better code reuse in D137100.

[Peiming]

fix typo

[aartbik]

I already had commented on the other one, so this is also LGTM, but wait for Alex.

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

75

Packs

[ftynse]

LGTM with comments addressed.

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

35–39

Please document top-level classes, especially the expectation of what the Converter should provide. Looking at the uses below, it looks like Structural1ToNConversionPattern follows CRTP and Converter is the concrete derived class. If that's the case, it is preferable to reflect this in the name of the attribute, e.g., ConcretePattern.

73

Did you mean "operation"?

86–88

Nit: wrap this into brackets. We omit brackets when both the condition and the body are single-line.

106

Please document.

126–127

This reflow looks irrelevant. Is your editor configured to wrap at 79 cols rather than 80?

[Peiming]

Address comments from Alex.

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

106

I add a more detailed comment in the base class.

[Peiming]

minor fix.

[Peiming]

rebase.

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

if (cast.getInputs().size() != 1) {
       unpacked.append(cast.getInputs().begin(), cast.getInputs().end());
       return;
     }
   }
   // 1 : 1 type conversion.
   unpacked.push_back(v);
 }
 
-class ConvertForOpTypes : public OpConversionPattern<ForOp> {
+// CRTP
+// A base class that takes care of 1:N type conversion, which maps the converted
+// op results (computed by the derived class) and materializes 1:N conversion.
+template <typename SourceOp, typename ConcretePattern>
+class Structural1ToNConversionPattern : public OpConversionPattern<SourceOp> {
 public:
-  using OpConversionPattern::OpConversionPattern;
+  using OpConversionPattern<SourceOp>::typeConverter;
+  using OpConversionPattern<SourceOp>::OpConversionPattern;
+  using OpAdaptor = typename OpConversionPattern<SourceOp>::OpAdaptor;
+
+  //
+  // Derived classes should provide the following method which performs the
+  // actual conversion. It should return llvm::None upon conversion failure and
+  // return the converted operation upon success.
+  //
+  // Optional<SourceOp> convertSourceOp(SourceOp op, OpAdaptor adaptor,
+  //                                    ConversionPatternRewriter &rewriter,
+  //                                    TypeRange dstTypes) const;
+
   LogicalResult
-  matchAndRewrite(ForOp op, OpAdaptor adaptor,
+  matchAndRewrite(SourceOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    SmallVector<Type> newResultTypes;
+    SmallVector<Type> dstTypes;
     SmallVector<unsigned> offsets;
     offsets.push_back(0);
     // Do the type conversion and record the offsets.
     for (Type type : op.getResultTypes()) {
-      if (failed(typeConverter->convertTypes(type, newResultTypes)))
-        return rewriter.notifyMatchFailure(op, "could not convert result");
-      offsets.push_back(newResultTypes.size());
+      if (failed(typeConverter->convertTypes(type, dstTypes)))
+        return rewriter.notifyMatchFailure(op, "could not convert result type");
+      offsets.push_back(dstTypes.size());
     }
 
+    // Calls the actual converter implementation to convert the operation.
+    Optional<SourceOp> newOp =
+        static_cast<const ConcretePattern *>(this)->convertSourceOp(
+            op, adaptor, rewriter, dstTypes);
+
+    if (!newOp)
+      return rewriter.notifyMatchFailure(op, "could not convert operation");
+
+    // Packs the return value.
+    SmallVector<Value> packedRets;
+    for (unsigned i = 1, e = offsets.size(); i < e; i++) {
+      unsigned start = offsets[i - 1], end = offsets[i];
+      unsigned len = end - start;
+      ValueRange mappedValue = newOp->getResults().slice(start, len);
+      if (len != 1) {
+        // 1 : N type conversion.
+        Type origType = op.getResultTypes()[i - 1];
+        Value mat = typeConverter->materializeSourceConversion(
+            rewriter, op.getLoc(), origType, mappedValue);
+        if (!mat) {
+          return rewriter.notifyMatchFailure(
+              op, "Failed to materialize 1:N type conversion");
+        }
+        packedRets.push_back(mat);
+      } else {
+        // 1 : 1 type conversion.
+        packedRets.push_back(mappedValue.front());
+      }
+    }
+
+    rewriter.replaceOp(op, packedRets);
+    return success();
+  }
+};
+
+class ConvertForOpTypes
+    : public Structural1ToNConversionPattern<ForOp, ConvertForOpTypes> {
+public:
+  using Structural1ToNConversionPattern::Structural1ToNConversionPattern;
+
+  // The callback required by CRTP.
+  Optional<ForOp> convertSourceOp(ForOp op, OpAdaptor adaptor,
+                                  ConversionPatternRewriter &rewriter,
+                                  TypeRange dstTypes) const {
     // Create a empty new op and inline the regions from the old op.
     //
     // This is a little bit tricky. We have two concerns here:
     //
     // 1. We cannot update the op in place because the dialect conversion
     // framework does not track type changes for ops updated in place, so it
     // won't insert appropriate materializations on the changed result types.
     // PR47938 tracks this issue, but it seems hard to fix. Instead, we need
     // to clone the op.
     //
     // 2. We need to resue the original region instead of cloning it, otherwise
     // the dialect conversion framework thinks that we just inserted all the
     // cloned child ops. But what we want is to "take" the child regions and let
     // the dialect conversion framework continue recursively into ops inside
     // those regions (which are already in its worklist; inlining them into the
     // new op's regions doesn't remove the child ops from the worklist).
 
     // convertRegionTypes already takes care of 1:N conversion.
     if (failed(rewriter.convertRegionTypes(&op.getLoopBody(), *typeConverter)))
-      return failure();
+      return llvm::None;
 
     // Unpacked the iteration arguments.
     SmallVector<Value> flatArgs;
     for (Value arg : adaptor.getInitArgs())
       unpackUnrealizedConversionCast(arg, flatArgs);
 
-    // We can not do clone as the number of result types after conversion might
-    // be different.
+    // We can not do clone as the number of result types after conversion
+    // might be different.
     ForOp newOp = rewriter.create<ForOp>(op.getLoc(), adaptor.getLowerBound(),
                                          adaptor.getUpperBound(),
                                          adaptor.getStep(), flatArgs);
 
     // Reserve whatever attributes in the original op.
     newOp->setAttrs(op->getAttrs());
 
     // We do not need the empty block created by rewriter.
     rewriter.eraseBlock(newOp.getBody(0));
     // Inline the type converted region from the original operation.
     rewriter.inlineRegionBefore(op.getLoopBody(), newOp.getLoopBody(),
                                 newOp.getLoopBody().end());
 
-    // Pack the return value.
-    SmallVector<Value, 6> packedRets;
-    for (unsigned i = 1, e = offsets.size(); i < e; i++) {
-      unsigned start = offsets[i - 1], end = offsets[i];
-      unsigned len = end - start;
-      ValueRange mappedValue = newOp.getResults().slice(start, len);
-      if (len != 1) {
-        // 1 : N type conversion.
-        Type origType = op.getResultTypes()[i - 1];
-        Value mat = typeConverter->materializeSourceConversion(
-            rewriter, op.getLoc(), origType, mappedValue);
-        if (!mat)
-          return rewriter.notifyMatchFailure(
-              op, "Failed to materialize 1:N type conversion");
-        packedRets.push_back(mat);
-      } else {
-        // 1 : 1 type conversion.
-        packedRets.push_back(mappedValue.front());
-      }
-    }
-
-    rewriter.replaceOp(op, packedRets);
-    return success();
+    return newOp;
   }
 };
 } // namespace
 
 namespace {
 class ConvertIfOpTypes : public OpConversionPattern<IfOp> {
 public:
   using OpConversionPattern::OpConversionPattern;