[Mlir-commits] [mlir] [mlir][bufferize] Add hoist-dynamic-allocs-option to buffer-results-to-out-params (PR #160985)

[Matthias Springer]

================
@@ -43,6 +45,49 @@ static bool hasStaticIdentityLayout(MemRefType type) {
   return type.getLayout().isIdentity();
 }
 
+/// Return the dynamic shapes of the `memref` based on the defining op. If the
+/// complete dynamic shape fails to be captured, return an empty value.
+/// Currently, only function block arguments are supported for capturing.
+static SmallVector<Value> getDynamicSize(Value memref, func::FuncOp funcOp) {
+  Operation *defOp = memref.getDefiningOp();
+  if (!defOp)
+    return {};
+  auto operands = defOp->getOperands();
+  SmallVector<Value> dynamicSizes;
+  for (Value size : operands) {
+    BlockArgument sizeSrc = dyn_cast<BlockArgument>(size);
+    if (!sizeSrc)
+      return {};
+
+    auto iter = llvm::find(funcOp.getArguments(), sizeSrc);
+    if (!iter)
+      return {};
+    dynamicSizes.push_back(*iter);
+  }
+  return dynamicSizes;
+}
+
+/// Returns the dynamic sizes at the callee, through the call relationship
+/// between the caller and callee.
+static SmallVector<Value> mapDynamicSizeAtCaller(func::CallOp call,
+                                                 func::FuncOp callee,
+                                                 ValueRange dynamicSizes) {
+  SmallVector<Value> mappedDynamicSizes;
+  for (Value size : dynamicSizes) {
+    auto callOperands = call.getOperands();
+    for (size_t i = 0, e = callOperands.size(); i < e; ++i) {
----------------
matthias-springer wrote:

nit: You can also write something like: `for [src, dst] : llvm::zip_equal(call.getOperands(), callee.getArguments())`

https://github.com/llvm/llvm-project/pull/160985