[WIP][mlir] enable conversion materialization for 1-1 type conversions

[ftynse]

Summary

Change the conversion materialization hook in TypeConverter to support the materialization of 1-1 conversions. Instead of having the hook as a virtual function, keep it as a callback, (1) it would compose better with different type conversion extensions and (2) this makes the hook optional. Change the hook signature to return Value instead of operation, allowing it to return the original value unmodified when no materialization is necessary for 1-1 conversion.

This diff also includes the changes to the LLVM and GPU dialect conversion to illustrate how this new support makes them simpler.

[merge_guards_bot]

Unit tests: pass. 62332 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/include/mlir/Transforms/DialectConversion.h

46

Do we forsee instances where we may want multiple materializers registered? I want to move TypeConverter closer to how ConversionTarget works and allow for it to be more easily extended. I can imagine a situation where we may want to materialize differently depending on the type.

[rriddle]

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

47

Do you intend for this to apply to PHI/non-entry Block arguments as well?

[ftynse]

mlir/include/mlir/Transforms/DialectConversion.h

46

Good point. The usecase I have only needs to materialize memref conversion, hence the change to return values. I'll look into conversion target to make this similarly more extendable.

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

47

Actually, no. Is there an easy way to differentiate here? Look at the owner of pointer?

[rriddle]

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

47

I think we will likely need some kind of indicator on the context of the conversion, in this case a Region signature conversion for the given operation. Ideally the pattern doing the conversion would handle this, but I that doesn't really scale/isn't extensible.

[ftynse]

Superseded by https://reviews.llvm.org/D79729

mlir/include/mlir/Transforms/DialectConversion.h

 //===----------------------------------------------------------------------===//
 // Type Conversion
 //===----------------------------------------------------------------------===//
 
 /// Base class for type conversion interface. Specific converters must
 /// derive this class and implement the pure virtual functions.
 class TypeConverter {
 public:
+  /// Type of the hook for materializing cast operations.
+  using ConversionMaterializer =
+      std::function<Value(PatternRewriter &, Type, ValueRange, Location)>;
+
+  /// Construct a TypeConverter. Derived classes must provide an implementation
+  /// of the ConversionMaterializer hook 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 provided when a
+  /// type conversion results in more than one type, or if a type conversion may
+  /// persist after the conversion has finished.
+  explicit TypeConverter(ConversionMaterializer materializer = nullptr)
+      : conversionMaterializer(materializer) {}
   virtual ~TypeConverter() = default;
 
   /// This class provides all of the information necessary to convert a type
   /// signature.
   class SignatureConversion {
   public:
     SignatureConversion(unsigned numOrigInputs)
         : remappedInputs(numOrigInputs) {}
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");
-  }
+  /// one result type by generating a cast operation of some kind.
+  ConversionMaterializer conversionMaterializer;
 };
 
 //===----------------------------------------------------------------------===//
 // Conversion Patterns
 //===----------------------------------------------------------------------===//
 
 /// Base class for the conversion patterns that require type changes. Specific
 /// conversions must derive this class and implement least one `rewrite` method.
public:
 
   /// Apply a signature conversion to the entry block of the given region. This
   /// replaces the entry block with a new block containing the updated
   /// signature. The new entry block to the region is returned for convenience.
   Block *
   applySignatureConversion(Region *region,
                            TypeConverter::SignatureConversion &conversion);
 
-  /// Replace all the uses of the block argument `from` with value `to`.
-  void replaceUsesOfBlockArgument(BlockArgument from, Value to);
-
   /// Return the converted value that replaces 'key'. Return 'key' if there is
   /// no such a converted value.
   Value getRemappedValue(Value key);
 
   //===--------------------------------------------------------------------===//
   // PatternRewriter Hooks
   //===--------------------------------------------------------------------===//
 

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     }
 
     // Move the region to the new function, update the entry block signature.
     rewriter.inlineRegionBefore(gpuFuncOp.getBody(), llvmFuncOp.getBody(),
                                 llvmFuncOp.end());
     rewriter.applySignatureConversion(&llvmFuncOp.getBody(),
                                       signatureConversion);
 
-    {
-      // For memref-typed arguments, insert the relevant loads in the beginning
-      // of the block to comply with the LLVM dialect calling convention. This
-      // needs to be done after signature conversion to get the right types.
-      OpBuilder::InsertionGuard guard(rewriter);
-      Block &block = llvmFuncOp.front();
-      rewriter.setInsertionPointToStart(&block);
-
-      for (auto en : llvm::enumerate(gpuFuncOp.getType().getInputs())) {
-        if (!en.value().isa<MemRefType>() &&
-            !en.value().isa<UnrankedMemRefType>())
-          continue;
-
-        BlockArgument arg = block.getArgument(en.index());
-        Value loaded = rewriter.create<LLVM::LoadOp>(loc, arg);
-        rewriter.replaceUsesOfBlockArgument(arg, loaded);
-      }
-    }
-
     rewriter.eraseOp(gpuFuncOp);
     return matchSuccess();
   }
 };
 
 struct GPUReturnOpLowering : public LLVMOpLowering {
   GPUReturnOpLowering(LLVMTypeConverter &typeConverter)
       : LLVMOpLowering(gpu::ReturnOp::getOperationName(),

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

 static llvm::cl::OptionCategory
     clOptionsCategory("Standard to LLVM lowering options");
 
 static llvm::cl::opt<bool>
     clUseAlloca(PASS_NAME "-use-alloca",
                 llvm::cl::desc("Replace emission of malloc/free by alloca"),
                 llvm::cl::init(false));
 
+/// Materializes the "load" operation from a pointer to a memref descriptor.
+/// This is used as a hook for type conversion.
+Value materializeDescriptorLoad(OpBuilder &builder, Type type,
+                                ValueRange pointer, Location loc) {
+  assert(pointer.size() == 1 && "only 1-1 conversion is supported");
+  if (!type.isa<MemRefType>() && !type.isa<UnrankedMemRefType>())
+    return pointer[0];
+  return builder.create<LLVM::LoadOp>(loc, pointer[0]);
+}
+
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx)
-    : llvmDialect(ctx->getRegisteredDialect<LLVM::LLVMDialect>()) {
+    : TypeConverter(materializeDescriptorLoad),
+      llvmDialect(ctx->getRegisteredDialect<LLVM::LLVMDialect>()) {
   assert(llvmDialect && "LLVM IR dialect is not registered");
   module = &llvmDialect->getLLVMModule();
 }
 
 // Get the LLVM context.
 llvm::LLVMContext &LLVMTypeConverter::getLLVMContext() {
   return module->getContext();
 }
auto newFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
         op->getLoc(), funcOp.getName(), llvmType, LLVM::Linkage::External,
         attributes);
     rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
                                 newFuncOp.end());
 
     // Tell the rewriter to convert the region signature.
     rewriter.applySignatureConversion(&newFuncOp.getBody(), result);
 
-    // Insert loads from memref descriptor pointers in function bodies.
-    if (!newFuncOp.getBody().empty()) {
-      Block *firstBlock = &newFuncOp.getBody().front();
-      rewriter.setInsertionPoint(firstBlock, firstBlock->begin());
-      for (unsigned idx : promotedArgIndices) {
-        BlockArgument arg = firstBlock->getArgument(idx);
-        Value loaded = rewriter.create<LLVM::LoadOp>(funcOp.getLoc(), arg);
-        rewriter.replaceUsesOfBlockArgument(arg, loaded);
-      }
-    }
-
     rewriter.eraseOp(op);
     return matchSuccess();
   }
 };
 
 //////////////// Support for Lowering operations on n-D vectors ////////////////
 namespace {
 // Helper struct to "unroll" operations on n-D vectors in terms of operations on

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(
+          assert(typeConverter->conversionMaterializer &&
+                 "ConversionMaterializer must be provided if uses of a block "
+                 "argument persist after 1->0 conversion");
+          Value newArg = typeConverter->conversionMaterializer(
               rewriter, origArg.getType(), llvm::None, loc);
-          origArg.replaceAllUsesWith(newOp->getResult(0));
+          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 cast or pack the new values. For 1->1 mappings, the converter may be
+    // omitted, 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 &&
-           cast->getNumOperands() == replArgs.size());
-    mapping.map(origArg, cast->getResult(0));
+    bool hasMaterializer =
+        typeConverter && typeConverter->conversionMaterializer;
+    assert((inputMap->size == 1 || hasMaterializer) &&
+           "ConversionMaterializer must be provided for 1->N block argument "
+           "conversions");
+    Value newArg = hasMaterializer
+                       ? typeConverter->conversionMaterializer(
+                             rewriter, origArg.getType(), replArgs, loc)
+                       : 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,

 }
 
 /// Apply a signature conversion to the entry block of the given region.
 Block *ConversionPatternRewriter::applySignatureConversion(
     Region *region, TypeConverter::SignatureConversion &conversion) {
   return impl->applySignatureConversion(region, conversion);
 }
 
-void ConversionPatternRewriter::replaceUsesOfBlockArgument(BlockArgument from,
-                                                           Value to) {
-  for (auto &u : from.getUses()) {
-    if (u.getOwner() == to.getDefiningOp())
-      continue;
-    u.getOwner()->replaceUsesOfWith(from, to);
-  }
-  impl->mapping.map(impl->mapping.lookupOrDefault(from), to);
-}
-
 /// Return the converted value that replaces 'key'. Return 'key' if there is
 /// no such a converted value.
 Value ConversionPatternRewriter::getRemappedValue(Value key) {
   return impl->mapping.lookupOrDefault(key);
 }
 
 /// PatternRewriter hook for splitting a block into two parts.
 Block *ConversionPatternRewriter::splitBlock(Block *block,

mlir/test/Conversion/StandardToLLVM/convert-memref-ops.mlir

//  CHECK-NEXT:  llvm.mlir.constant(13 : index) : !llvm.i64
   %3 = dim %mixed, 3 : memref<42x?x?x13x?xf32>
 //  CHECK-NEXT:  llvm.extractvalue %[[ld]][3, 4] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
   %4 = dim %mixed, 4 : memref<42x?x?x13x?xf32>
   return
 }
 
 // CHECK-LABEL: func @static_memref_dim(%arg0: !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }*">) {
 func @static_memref_dim(%static : memref<42x32x15x13x27xf32>) {
-// CHECK-NEXT:  %[[ld:.*]] = llvm.load %{{.*}} : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }*">
 // CHECK-NEXT:  llvm.mlir.constant(42 : index) : !llvm.i64
   %0 = dim %static, 0 : memref<42x32x15x13x27xf32>
 // CHECK-NEXT:  llvm.mlir.constant(32 : index) : !llvm.i64
   %1 = dim %static, 1 : memref<42x32x15x13x27xf32>
 // CHECK-NEXT:  llvm.mlir.constant(15 : index) : !llvm.i64
   %2 = dim %static, 2 : memref<42x32x15x13x27xf32>
 // CHECK-NEXT:  llvm.mlir.constant(13 : index) : !llvm.i64
   %3 = dim %static, 3 : memref<42x32x15x13x27xf32>
 // CHECK-NEXT:  llvm.mlir.constant(27 : index) : !llvm.i64
   %4 = dim %static, 4 : memref<42x32x15x13x27xf32>
   return
 }
 

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

 // RUN: mlir-opt %s -convert-std-to-llvm -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: func @address_space(
 //       CHECK:   %{{.*}}: !llvm<"{ float addrspace(7)*, float addrspace(7)*, i64, [1 x i64], [1 x i64] }*">)
-//       CHECK:   llvm.load %{{.*}} : !llvm<"{ float addrspace(7)*, float addrspace(7)*, i64, [1 x i64], [1 x i64] }*">
 func @address_space(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>, 7>) {
   %0 = alloc() : memref<32xf32, affine_map<(d0) -> (d0)>, 5>
   %1 = constant 7 : index
   // CHECK: llvm.load %{{.*}} : !llvm<"float addrspace(5)*">
   %2 = load %0[%1] : memref<32xf32, affine_map<(d0) -> (d0)>, 5>
   std.return
 }
 

mlir/test/lib/TestDialect/TestPatterns.cpp

PatternMatchResult matchAndRewrite(Operation *op,
     rewriter.replaceOp(op, {illegalOp});
     return matchSuccess();
   }
 };
 } // namespace
 
 namespace {
 struct TestTypeConverter : public TypeConverter {
-  using TypeConverter::TypeConverter;
+  TestTypeConverter()
+      : TypeConverter(TestTypeConverter::materializeConversion) {}
 
   LogicalResult convertType(Type t, SmallVectorImpl<Type> &results) override {
     // Drop I16 types.
     if (t.isInteger(16))
       return success();
 
     // Convert I64 to F64.
     if (t.isInteger(64)) {
       results.push_back(FloatType::getF64(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
+  /// Implement 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 {
+  static Value materializeConversion(PatternRewriter &rewriter, Type resultType,
+                                     ValueRange inputs, Location loc) {
+    if (inputs.size() == 1)
+      return inputs[0];
     return rewriter.create<TestCastOp>(loc, resultType, inputs);
   }
 };
 
 struct TestLegalizePatternDriver
     : public ModulePass<TestLegalizePatternDriver> {
   /// The mode of conversion to use with the driver.
   enum class ConversionMode { Analysis, Full, Partial };