Merge conflicts

Author: Ferdinand Lemaire

mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h

@@ -53,12 +53,14 @@ void populateDynamicDimSizes(OpBuilder &b, Location loc, Value shapedValue,
 /// This function returns `failure()` in case of unsupported casts. E.g., casts
 /// with differing element types or memory spaces.
 FailureOr<Value> castOrReallocMemRefValue(OpBuilder &b, Value value,
-                                          MemRefType type);
+                                          MemRefType type,
+                                          const BufferizationOptions &options);
 
 /// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the
 /// to_memref op are different, a memref.cast is needed.
 LogicalResult foldToMemrefToTensorPair(RewriterBase &rewriter,
-                                       ToMemrefOp toMemref);
+                                       ToMemrefOp toMemref,
+                                       const BufferizationOptions &options);
 
 /// Add the canonicalization patterns for bufferization.dealloc to the given
 /// pattern set to make them available to other passes (such as

mlir/include/mlir/Dialect/Func/Transforms/Passes.td

@@ -51,4 +51,15 @@ def DuplicateFunctionEliminationPass : Pass<"duplicate-function-elimination",
   let constructor = "mlir::func::createDuplicateFunctionEliminationPass()";
 }
 
+def AnnotateFunctionType: Pass<"annotate-function-type", "func::FuncOp"> {
+  let summary = "Annotate the function type as type attributes";
+  let description = [{
+    Annotates all the inputs and outputs of func.func operators with a type
+    attribute. The type attribute mirrors the actual type of the inputs/outputs.
+
+    This pass can be used to trace back the original types of func.func
+    operators in case they need to be modified.
+  }];
+}
+
 #endif // MLIR_DIALECT_FUNC_TRANSFORMS_PASSES_TD

mlir/include/mlir/Dialect/PDL/IR/Builtins.h

@@ -26,6 +26,19 @@ namespace pdl {
 void registerBuiltins(PDLPatternModule &pdlPattern);
 
 namespace builtin {
+enum class BinaryOpKind {
+  add,
+  sub,
+  mul,
+  div,
+  mod,
+};
+
+enum class UnaryOpKind {
+  log2,
+  exp2,
+};
+
 LogicalResult createDictionaryAttr(PatternRewriter &rewriter,
                                    PDLResultList &results,
                                    ArrayRef<PDLValue> args);
@@ -35,8 +48,27 @@ LogicalResult addEntryToDictionaryAttr(PatternRewriter &rewriter,
 Attribute createArrayAttr(PatternRewriter &rewriter);
 Attribute addElemToArrayAttr(PatternRewriter &rewriter, Attribute attr,
                              Attribute element);
+template <BinaryOpKind T>
+LogicalResult binaryOp(PatternRewriter &rewriter, PDLResultList &results,
+                       llvm::ArrayRef<PDLValue> args);
+LogicalResult mul(PatternRewriter &rewriter, PDLResultList &results,
+                  llvm::ArrayRef<PDLValue> args);
+LogicalResult div(PatternRewriter &rewriter, PDLResultList &results,
+                  llvm::ArrayRef<PDLValue> args);
+LogicalResult mod(PatternRewriter &rewriter, PDLResultList &results,
+                  llvm::ArrayRef<PDLValue> args);
 LogicalResult add(PatternRewriter &rewriter, PDLResultList &results,
                   llvm::ArrayRef<PDLValue> args);
+LogicalResult sub(PatternRewriter &rewriter, PDLResultList &results,
+                  llvm::ArrayRef<PDLValue> args);
+LogicalResult log2(PatternRewriter &rewriter, PDLResultList &results,
+                   llvm::ArrayRef<PDLValue> args);
+LogicalResult exp2(PatternRewriter &rewriter, PDLResultList &results,
+                   llvm::ArrayRef<PDLValue> args);
+
+template <BinaryOpKind T>
+LogicalResult binaryOp(PatternRewriter &rewriter, PDLResultList &results,
+                       llvm::ArrayRef<PDLValue> args);
 } // namespace builtin
 } // namespace pdl
 } // namespace mlir

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

@@ -23,9 +23,9 @@ using namespace mlir::bufferization;
 // Helper functions
 //===----------------------------------------------------------------------===//
 
-FailureOr<Value>
-mlir::bufferization::castOrReallocMemRefValue(OpBuilder &b, Value value,
-                                              MemRefType destType) {
+FailureOr<Value> mlir::bufferization::castOrReallocMemRefValue(
+    OpBuilder &b, Value value, MemRefType destType,
+    const BufferizationOptions &options) {
   auto srcType = llvm::cast<MemRefType>(value.getType());
 
   // Element type, rank and memory space must match.
@@ -73,18 +73,23 @@ mlir::bufferization::castOrReallocMemRefValue(OpBuilder &b, Value value,
     Value size = b.create<memref::DimOp>(loc, value, i);
     dynamicOperands.push_back(size);
   }
-  // TODO: Use alloc/memcpy callback from BufferizationOptions if called via
-  // BufferizableOpInterface impl of ToMemrefOp.
-  Value copy = b.create<memref::AllocOp>(loc, destType, dynamicOperands);
-  b.create<memref::CopyOp>(loc, value, copy);
+
+  FailureOr<Value> copy =
+      options.createAlloc(b, loc, destType, dynamicOperands);
+  if (failed(copy)) {
+    return failure();
+  }
+  if (failed(options.createMemCpy(b, loc, value, *copy))) {
+    return failure();
+  }
   return copy;
 }
 
 /// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the
 /// to_memref op are different, a memref.cast is needed.
-LogicalResult
-mlir::bufferization::foldToMemrefToTensorPair(RewriterBase &rewriter,
-                                              ToMemrefOp toMemref) {
+LogicalResult mlir::bufferization::foldToMemrefToTensorPair(
+    RewriterBase &rewriter, ToMemrefOp toMemref,
+    const BufferizationOptions &options) {
   auto memrefToTensor = toMemref.getTensor().getDefiningOp<ToTensorOp>();
   if (!memrefToTensor)
     return failure();
@@ -105,7 +110,7 @@ mlir::bufferization::foldToMemrefToTensorPair(RewriterBase &rewriter,
   // Ranked memref -> Ranked memref cast.
   if (rankedSrcType && rankedDestType) {
     FailureOr<Value> replacement = castOrReallocMemRefValue(
-        rewriter, memrefToTensor.getMemref(), rankedDestType);
+        rewriter, memrefToTensor.getMemref(), rankedDestType, options);
     if (failed(replacement))
       return failure();
 
@@ -802,7 +807,7 @@ struct ToMemrefToTensorFolding : public OpRewritePattern<ToMemrefOp> {
 
   LogicalResult matchAndRewrite(ToMemrefOp toMemref,
                                 PatternRewriter &rewriter) const final {
-    return foldToMemrefToTensorPair(rewriter, toMemref);
+    return foldToMemrefToTensorPair(rewriter, toMemref, {});
   }
 };
 
@@ -850,7 +855,7 @@ void ToMemrefOp::getCanonicalizationPatterns(RewritePatternSet &results,
 LogicalResult ToMemrefOp::bufferize(RewriterBase &rewriter,
                                     const BufferizationOptions &options) {
   // Fold to_memref(to_tensor(x)) to x. Insert a cast if necessary.
-  (void)foldToMemrefToTensorPair(rewriter, *this);
+  (void)foldToMemrefToTensorPair(rewriter, *this, options);
   // Note: The return value of `bufferize` indicates whether there was an error
   // or not. (And not whether the pattern matched or not.)
   return success();

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

@@ -75,7 +75,7 @@ BufferizeTypeConverter::BufferizeTypeConverter() {
       if (!rankedDestType)
         return nullptr;
       FailureOr<Value> replacement =
-          castOrReallocMemRefValue(builder, inputs[0], rankedDestType);
+          castOrReallocMemRefValue(builder, inputs[0], rankedDestType, {});
       if (failed(replacement))
         return nullptr;
       return *replacement;
@@ -509,8 +509,8 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
   // Fold all to_memref(to_tensor(x)) pairs.
   for (Operation *op : toMemrefOps) {
     rewriter.setInsertionPoint(op);
-    (void)bufferization::foldToMemrefToTensorPair(rewriter,
-                                                  cast<ToMemrefOp>(op));
+    (void)bufferization::foldToMemrefToTensorPair(
+        rewriter, cast<ToMemrefOp>(op), options);
   }
 
   // Remove all dead to_tensor ops.

mlir/lib/Dialect/Func/Transforms/AnnotateFunctionType.cpp

@@ -0,0 +1,47 @@
+//===- AnnotateInputTypes.cpp - Type attribute annotation for func ops ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass that creates type attributes for func parameters,
+// that mirror the actual type. This is useful when the func op input types
+// might change.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/Transforms/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace mlir::func {
+#define GEN_PASS_DEF_ANNOTATEFUNCTIONTYPE
+#include "mlir/Dialect/Func/Transforms/Passes.h.inc"
+} // namespace mlir::func
+
+namespace {
+struct AnnotateFunctionTypePass
+    : public mlir::func::impl::AnnotateFunctionTypeBase<
+          AnnotateFunctionTypePass> {
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+    auto inputs = func.getArgumentTypes();
+    auto results = func.getResultTypes();
+
+    for (const auto [argNum, type] : llvm::enumerate(inputs)) {
+      func.setArgAttr(argNum, "func.orig_type", TypeAttr::get(type));
+    }
+
+    for (const auto [resultNum, type] : llvm::enumerate(results)) {
+      func.setResultAttr(resultNum, "func.orig_type", TypeAttr::get(type));
+    }
+  }
+};
+} // namespace

mlir/lib/Dialect/Func/Transforms/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_mlir_dialect_library(MLIRFuncTransforms
+  AnnotateFunctionType.cpp
   DecomposeCallGraphTypes.cpp
   DuplicateFunctionElimination.cpp
   FuncBufferize.cpp