[mlir][Transforms] Support 1:N mappings in ConversionValueMapping

Author: matthias-springer

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

@@ -153,31 +153,42 @@ 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
-  // insert a target materialization from the original block argument type to
-  // a legal type.
-  addArgumentMaterialization([&](OpBuilder &builder,
-                                 UnrankedMemRefType resultType,
-                                 ValueRange inputs, Location loc) {
+  // 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) {
+    return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
+        .getResult(0);
+  });
+
+  // Source materializations convert the MemrRef descriptor elements
+  // (multiple SSA values that make up a MemrRef descriptor) back to the
+  // original MemRef type.
+  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
+    // A source 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);
   });
-  addArgumentMaterialization([&](OpBuilder &builder, MemRefType resultType,
-                                 ValueRange inputs, Location loc) {
+  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) {
+    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());
@@ -192,15 +203,13 @@ LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
     } else {
       desc = MemRefDescriptor::pack(builder, loc, *this, resultType, inputs);
     }
-    // An argument materialization must return a value of type `resultType`,
+    // A source 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);
   });
-  // Add generic source and target materializations to handle cases where
-  // non-LLVM types persist after an LLVM conversion.
-  addSourceMaterialization([&](OpBuilder &builder, Type resultType,
+  addTargetMaterialization([&](OpBuilder &builder, Type resultType,
                                ValueRange inputs, Location loc) {
     if (inputs.size() != 1)
       return Value();
@@ -209,12 +218,50 @@ LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
         .getResult(0);
   });
   addTargetMaterialization([&](OpBuilder &builder, Type resultType,
-                               ValueRange inputs, Location loc) {
-    if (inputs.size() != 1)
+                               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 builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
-        .getResult(0);
+    return Value();
   });
 
   // Integer memory spaces map to themselves.

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

@@ -61,7 +61,6 @@ BufferizeTypeConverter::BufferizeTypeConverter() {
   addConversion([](UnrankedTensorType type) -> Type {
     return UnrankedMemRefType::get(type.getElementType(), 0);
   });
-  addArgumentMaterialization(materializeToTensor);
   addSourceMaterialization(materializeToTensor);
   addTargetMaterialization([](OpBuilder &builder, BaseMemRefType type,
                               ValueRange inputs, Location loc) -> Value {

mlir/lib/Dialect/EmitC/Transforms/TypeConversions.cpp

@@ -33,7 +33,6 @@ void mlir::populateEmitCSizeTTypeConversions(TypeConverter &converter) {
 
   converter.addSourceMaterialization(materializeAsUnrealizedCast);
   converter.addTargetMaterialization(materializeAsUnrealizedCast);
-  converter.addArgumentMaterialization(materializeAsUnrealizedCast);
 }
 
 /// Get an unsigned integer or size data type corresponding to \p ty.

mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp

@@ -154,7 +154,6 @@ class DetensorizeTypeConverter : public TypeConverter {
     });
 
     addSourceMaterialization(sourceMaterializationCallback);
-    addArgumentMaterialization(sourceMaterializationCallback);
   }
 };
 

mlir/lib/Dialect/Quant/Transforms/StripFuncQuantTypes.cpp

@@ -36,9 +36,10 @@ class QuantizedTypeConverter : public TypeConverter {
   static Type convertQuantizedType(QuantizedType quantizedType) {
     return quantizedType.getStorageType();
   }
-  
+
   static Type convertTensorType(TensorType tensorType) {
-    if (auto quantizedType = dyn_cast<QuantizedType>(tensorType.getElementType()))
+    if (auto quantizedType =
+            dyn_cast<QuantizedType>(tensorType.getElementType()))
       return tensorType.clone(convertQuantizedType(quantizedType));
     return tensorType;
   }
@@ -50,20 +51,19 @@ class QuantizedTypeConverter : public TypeConverter {
   }
 
 public:
-
   explicit QuantizedTypeConverter() {
     addConversion([](Type type) { return type; });
     addConversion(convertQuantizedType);
     addConversion(convertTensorType);
 
-    addArgumentMaterialization(materializeConversion);
     addSourceMaterialization(materializeConversion);
     addTargetMaterialization(materializeConversion);
   }
 };
 
 // Conversion pass
-class StripFuncQuantTypes : public impl::StripFuncQuantTypesBase<StripFuncQuantTypes> {
+class StripFuncQuantTypes
+    : public impl::StripFuncQuantTypesBase<StripFuncQuantTypes> {
 
   // Return whether a type is considered legal when occurring in the header of
   // a function or as an operand to a 'return' op.
@@ -74,11 +74,10 @@ class StripFuncQuantTypes : public impl::StripFuncQuantTypesBase<StripFuncQuantT
   }
 
 public:
-
   void runOnOperation() override {
-    
+
     auto moduleOp = cast<ModuleOp>(getOperation());
-    auto* context = &getContext();
+    auto *context = &getContext();
 
     QuantizedTypeConverter typeConverter;
     ConversionTarget target(*context);
@@ -111,4 +110,3 @@ class StripFuncQuantTypes : public impl::StripFuncQuantTypesBase<StripFuncQuantT
 
 } // namespace quant
 } // namespace mlir
-

mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp

@@ -69,9 +69,6 @@ SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
 
   // Required by scf.for 1:N type conversion.
   addSourceMaterialization(materializeTuple);
-
-  // Required as a workaround until we have full 1:N support.
-  addArgumentMaterialization(materializeTuple);
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

@@ -481,7 +481,6 @@ void mlir::vector::populateVectorLinearizeTypeConversionsAndLegality(
 
     return builder.create<vector::ShapeCastOp>(loc, type, inputs.front());
   };
-  typeConverter.addArgumentMaterialization(materializeCast);
   typeConverter.addSourceMaterialization(materializeCast);
   typeConverter.addTargetMaterialization(materializeCast);
   target.markUnknownOpDynamicallyLegal(