do not build argument materializations anymore

Author: matthias-springer

mlir/include/mlir/Transforms/DialectConversion.h

@@ -583,8 +583,8 @@ class ConversionPattern : public RewritePattern {
       : RewritePattern(std::forward<Args>(args)...),
         typeConverter(&typeConverter) {}
 
-  static SmallVector<Value>
-  getOneToOneAdaptorOperands(ArrayRef<ArrayRef<Value>> operands);
+  SmallVector<Value>
+  getOneToOneAdaptorOperands(ArrayRef<ArrayRef<Value>> operands) const;
 
 protected:
   /// An optional type converter for use by this pattern.
@@ -858,6 +858,9 @@ class ConversionPatternRewriter final : public PatternRewriter {
   /// PatternRewriter hook for replacing an operation.
   void replaceOp(Operation *op, Operation *newOp) override;
 
+  /// PatternRewriter hook for replacing an operation.
+  void replaceOpWithMultiple(Operation *op, ArrayRef<ValueRange> newValues);
+
   /// PatternRewriter hook for erasing a dead operation. The uses of this
   /// operation *must* be made dead by the end of the conversion process,
   /// otherwise an assert will be issued.

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

@@ -153,6 +153,7 @@ LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
                                        type.isVarArg());
   });
 
+/*
   // Argument materializations convert from the new block argument types
   // (multiple SSA values that make up a memref descriptor) back to the
   // original block argument type. The dialect conversion framework will then
@@ -198,16 +199,62 @@ LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
     return builder.create<UnrealizedConversionCastOp>(loc, resultType, desc)
         .getResult(0);
   });
+
+*/
   // Add generic source and target materializations to handle cases where
   // non-LLVM types persist after an LLVM conversion.
   addSourceMaterialization([&](OpBuilder &builder, Type resultType,
                                ValueRange inputs, Location loc) {
-    if (inputs.size() != 1)
-      return Value();
+    //if (inputs.size() != 1)
+    //  return Value();
 
     return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
         .getResult(0);
   });
+  addSourceMaterialization([&](OpBuilder &builder, MemRefType resultType,
+                               ValueRange inputs, Location loc) {
+    if (inputs.size()== 1 && isa<LLVM::LLVMStructType>(inputs.front().getType())) return Value();
+
+    Value desc;
+    if (inputs.size() == 1 && isa<LLVM::LLVMPointerType>(inputs.front().getType())) {
+      // This is a bare pointer. We allow bare pointers only for function entry
+      // blocks.
+      BlockArgument barePtr = dyn_cast<BlockArgument>(inputs.front());
+      if (!barePtr)
+        return Value();
+      Block *block = barePtr.getOwner();
+      if (!block->isEntryBlock() ||
+          !isa<FunctionOpInterface>(block->getParentOp()))
+        return Value();
+      desc = MemRefDescriptor::fromStaticShape(builder, loc, *this, resultType,
+                                               inputs[0]);
+    } else {
+      //llvm::errs() << "pack elems: " << inputs.size() << "\n";
+      //llvm::errs() << inputs[0] << "\n";
+      desc = MemRefDescriptor::pack(builder, loc, *this, resultType, inputs);
+      //llvm::errs() << "done packing\n";
+    }
+    // An argument materialization must return a value of type `resultType`,
+    // so insert a cast from the memref descriptor type (!llvm.struct) to the
+    // original memref type.
+    return builder.create<UnrealizedConversionCastOp>(loc, resultType, desc)
+        .getResult(0);
+  });
+  addSourceMaterialization([&](OpBuilder &builder, UnrankedMemRefType resultType,
+                               ValueRange inputs, Location loc) {
+    if (inputs.size() == 1) {
+      // Bare pointers are not supported for unranked memrefs because a
+      // memref descriptor cannot be built just from a bare pointer.
+      return Value();
+    }
+    Value desc =
+        UnrankedMemRefDescriptor::pack(builder, loc, *this, resultType, inputs);
+    // An argument materialization must return a value of type
+    // `resultType`, so insert a cast from the memref descriptor type
+    // (!llvm.struct) to the original memref type.
+    return builder.create<UnrealizedConversionCastOp>(loc, resultType, desc)
+        .getResult(0);
+  });
   addTargetMaterialization([&](OpBuilder &builder, Type resultType,
                                ValueRange inputs, Location loc) {
     if (inputs.size() != 1)
@@ -216,6 +263,51 @@ LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
     return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
         .getResult(0);
   });
+  addTargetMaterialization([&](OpBuilder &builder, Type resultType,
+                               ValueRange inputs,
+                               Location loc, Type originalType) -> Value {
+    llvm::errs() << "TARGET MAT: -> " << resultType << "\n";
+    if (!originalType) {
+      llvm::errs() << " -- no orig\n";
+      return Value();
+    }
+    if (auto memrefType = dyn_cast<MemRefType>(originalType)) {
+      assert(isa<LLVM::LLVMStructType>(resultType) && "expected struct type");
+      if (inputs.size() == 1) {
+        Value input = inputs.front();
+        if (auto castOp =input.getDefiningOp<UnrealizedConversionCastOp>()) {
+          if (castOp.getInputs().size() == 1 && isa<LLVM::LLVMPointerType>(castOp.getInputs()[0].getType())) {
+            input = castOp.getInputs()[0];
+          }
+        }
+        if (!isa<LLVM::LLVMPointerType>(input.getType()))
+          return Value();
+        BlockArgument barePtr = dyn_cast<BlockArgument>(input);
+        if (!barePtr)
+          return Value();
+        Block *block = barePtr.getOwner();
+        if (!block->isEntryBlock() ||
+            !isa<FunctionOpInterface>(block->getParentOp()))
+          return Value();
+        // Bare ptr
+        return MemRefDescriptor::fromStaticShape(builder, loc, *this, memrefType,
+                                                 input);
+      }
+      return MemRefDescriptor::pack(builder, loc, *this, memrefType, inputs);
+    }
+    if (auto memrefType = dyn_cast<UnrankedMemRefType>(originalType)) {
+      assert(isa<LLVM::LLVMStructType>(resultType) && "expected struct type");
+      if (inputs.size() == 1) {
+          // Bare pointers are not supported for unranked memrefs because a
+          // memref descriptor cannot be built just from a bare pointer.
+          return Value();
+      }
+      return UnrankedMemRefDescriptor::pack(builder, loc, *this,
+                                                    memrefType, inputs);
+    }
+
+      return Value();
+  });
 
   // Integer memory spaces map to themselves.
   addTypeAttributeConversion(

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

@@ -16,20 +16,16 @@ using namespace mlir::scf;
 
 namespace {
 
-// Unpacks the single unrealized_conversion_cast using the list of inputs
-// e.g., return [%b, %c, %d] for %a = unrealized_conversion_cast(%b, %c, %d)
-static void unpackUnrealizedConversionCast(Value v,
-                                           SmallVectorImpl<Value> &unpacked) {
-  if (auto cast =
-          dyn_cast_or_null<UnrealizedConversionCastOp>(v.getDefiningOp())) {
-    if (cast.getInputs().size() != 1) {
-      // 1 : N type conversion.
-      unpacked.append(cast.getInputs().begin(), cast.getInputs().end());
-      return;
-    }
-  }
-  // 1 : 1 type conversion.
-  unpacked.push_back(v);
+static SmallVector<Value> flattenValues(ArrayRef<ArrayRef<Value>> values) {
+  SmallVector<Value> result;
+  for (ArrayRef<Value> v : values)
+    llvm::append_range(result, v);
+  return result;
+}
+
+static Value getSingleValue(ArrayRef<Value> values) {
+  assert(values.size() == 1 && "expected single value");
+  return values.front();
 }
 
 // CRTP
@@ -40,19 +36,21 @@ class Structural1ToNConversionPattern : public OpConversionPattern<SourceOp> {
 public:
   using OpConversionPattern<SourceOp>::typeConverter;
   using OpConversionPattern<SourceOp>::OpConversionPattern;
-  using OpAdaptor = typename OpConversionPattern<SourceOp>::OpAdaptor;
+  using OneToNOpAdaptor =
+      typename OpConversionPattern<SourceOp>::OneToNOpAdaptor;
 
   //
   // Derived classes should provide the following method which performs the
   // actual conversion. It should return std::nullopt upon conversion failure
   // and return the converted operation upon success.
   //
-  // std::optional<SourceOp> convertSourceOp(SourceOp op, OpAdaptor adaptor,
-  //                                    ConversionPatternRewriter &rewriter,
-  //                                    TypeRange dstTypes) const;
+  // std::optional<SourceOp> convertSourceOp(
+  //     SourceOp op, OneToNOpAdaptor adaptor,
+  //     ConversionPatternRewriter &rewriter,
+  //     TypeRange dstTypes) const;
 
   LogicalResult
-  matchAndRewrite(SourceOp op, OpAdaptor adaptor,
+  matchAndRewrite(SourceOp op, OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     SmallVector<Type> dstTypes;
     SmallVector<unsigned> offsets;
@@ -73,28 +71,15 @@ class Structural1ToNConversionPattern : public OpConversionPattern<SourceOp> {
       return rewriter.notifyMatchFailure(op, "could not convert operation");
 
     // Packs the return value.
-    SmallVector<Value> packedRets;
+    SmallVector<ValueRange> 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());
-      }
+      packedRets.push_back(mappedValue);
     }
 
-    rewriter.replaceOp(op, packedRets);
+    rewriter.replaceOpWithMultiple(op, packedRets);
     return success();
   }
 };
@@ -105,7 +90,7 @@ class ConvertForOpTypes
   using Structural1ToNConversionPattern::Structural1ToNConversionPattern;
 
   // The callback required by CRTP.
-  std::optional<ForOp> convertSourceOp(ForOp op, OpAdaptor adaptor,
+  std::optional<ForOp> convertSourceOp(ForOp op, OneToNOpAdaptor adaptor,
                                        ConversionPatternRewriter &rewriter,
                                        TypeRange dstTypes) const {
     // Create a empty new op and inline the regions from the old op.
@@ -129,16 +114,13 @@ class ConvertForOpTypes
     if (failed(rewriter.convertRegionTypes(&op.getRegion(), *typeConverter)))
       return std::nullopt;
 
-    // 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.
-    ForOp newOp = rewriter.create<ForOp>(op.getLoc(), adaptor.getLowerBound(),
-                                         adaptor.getUpperBound(),
-                                         adaptor.getStep(), flatArgs);
+    ForOp newOp = rewriter.create<ForOp>(
+        op.getLoc(), getSingleValue(adaptor.getLowerBound()),
+        getSingleValue(adaptor.getUpperBound()),
+        getSingleValue(adaptor.getStep()),
+        flattenValues(adaptor.getInitArgs()));
 
     // Reserve whatever attributes in the original op.
     newOp->setAttrs(op->getAttrs());
@@ -160,12 +142,12 @@ class ConvertIfOpTypes
 public:
   using Structural1ToNConversionPattern::Structural1ToNConversionPattern;
 
-  std::optional<IfOp> convertSourceOp(IfOp op, OpAdaptor adaptor,
+  std::optional<IfOp> convertSourceOp(IfOp op, OneToNOpAdaptor adaptor,
                                       ConversionPatternRewriter &rewriter,
                                       TypeRange dstTypes) const {
 
-    IfOp newOp = rewriter.create<IfOp>(op.getLoc(), dstTypes,
-                                       adaptor.getCondition(), true);
+    IfOp newOp = rewriter.create<IfOp>(
+        op.getLoc(), dstTypes, getSingleValue(adaptor.getCondition()), true);
     newOp->setAttrs(op->getAttrs());
 
     // We do not need the empty blocks created by rewriter.
@@ -189,15 +171,11 @@ class ConvertWhileOpTypes
 public:
   using Structural1ToNConversionPattern::Structural1ToNConversionPattern;
 
-  std::optional<WhileOp> convertSourceOp(WhileOp op, OpAdaptor adaptor,
+  std::optional<WhileOp> convertSourceOp(WhileOp op, OneToNOpAdaptor adaptor,
                                          ConversionPatternRewriter &rewriter,
                                          TypeRange dstTypes) const {
-    // Unpacked the iteration arguments.
-    SmallVector<Value> flatArgs;
-    for (Value arg : adaptor.getOperands())
-      unpackUnrealizedConversionCast(arg, flatArgs);
-
-    auto newOp = rewriter.create<WhileOp>(op.getLoc(), dstTypes, flatArgs);
+    auto newOp = rewriter.create<WhileOp>(op.getLoc(), dstTypes,
+                                          flattenValues(adaptor.getOperands()));
 
     for (auto i : {0u, 1u}) {
       if (failed(rewriter.convertRegionTypes(&op.getRegion(i), *typeConverter)))
@@ -218,13 +196,10 @@ class ConvertYieldOpTypes : public OpConversionPattern<scf::YieldOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
-  matchAndRewrite(scf::YieldOp op, OpAdaptor adaptor,
+  matchAndRewrite(scf::YieldOp op, OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    SmallVector<Value> unpackedYield;
-    for (Value operand : adaptor.getOperands())
-      unpackUnrealizedConversionCast(operand, unpackedYield);
-
-    rewriter.replaceOpWithNewOp<scf::YieldOp>(op, unpackedYield);
+    rewriter.replaceOpWithNewOp<scf::YieldOp>(
+        op, flattenValues(adaptor.getOperands()));
     return success();
   }
 };
@@ -235,13 +210,10 @@ class ConvertConditionOpTypes : public OpConversionPattern<ConditionOp> {
 public:
   using OpConversionPattern<ConditionOp>::OpConversionPattern;
   LogicalResult
-  matchAndRewrite(ConditionOp op, OpAdaptor adaptor,
+  matchAndRewrite(ConditionOp op, OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    SmallVector<Value> unpackedYield;
-    for (Value operand : adaptor.getOperands())
-      unpackUnrealizedConversionCast(operand, unpackedYield);
-
-    rewriter.modifyOpInPlace(op, [&]() { op->setOperands(unpackedYield); });
+    rewriter.modifyOpInPlace(
+        op, [&]() { op->setOperands(flattenValues(adaptor.getOperands())); });
     return success();
   }
 };