[Mlir-commits] [mlir] [mlir] [memref] [transform] Add alloc_to_global op. (PR #211141)

[Matthias Springer]

================
@@ -125,6 +125,64 @@ void transform::ApplyResolveRankedShapedTypeResultDimsPatternsOp::
   memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
 }
 
+//===----------------------------------------------------------------------===//
+// Alloc and alloca to global utilities
+//===----------------------------------------------------------------------===//
+
+/// Converts an allocation operation (`memref.alloca` or `memref.alloc`) to a
+/// `memref.global` operation in the nearest symbol table, and replaces the
+/// allocation with a `memref.get_global` operation. Any `memref.dealloc`
+/// operations referencing the allocation are erased.
+template <typename AllocLikeOp>
+static DiagnosedSilenceableFailure
+allocLikeToGlobal(transform::TransformRewriter &rewriter,
+                  AllocLikeOp allocLikeOp, StringRef globalName,
+                  memref::GlobalOp &globalOp,
+                  memref::GetGlobalOp &getGlobalOp) {
+  MemRefType memrefType = allocLikeOp.getType();
+  if (!memrefType.hasStaticShape()) {
+    return emitSilenceableFailure(allocLikeOp->getLoc())
+           << "global ops require statically shaped memrefs, but got "
+           << memrefType;
+  }
+
+  MLIRContext *ctx = rewriter.getContext();
+  Location loc = allocLikeOp->getLoc();
+
+  // Find nearest symbol table.
+  Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(allocLikeOp);
+  assert(symbolTableOp && "expected payload to be in symbol table");
+  SymbolTable symbolTable(symbolTableOp);
+
+  // Insert a `memref.global` into the symbol table.
+  Type resultType = allocLikeOp.getResult().getType();
+  OpBuilder builder(rewriter.getContext());
+  // TODO: Add a better builder for this.
+  globalOp = memref::GlobalOp::create(
+      builder, loc, StringAttr::get(ctx, globalName),
+      StringAttr::get(ctx, "private"), TypeAttr::get(resultType), Attribute{},
+      UnitAttr{}, IntegerAttr{});
+  symbolTable.insert(globalOp);
+
+  // Remove any `memref.dealloc` operations referencing this allocation.
+  // We assume that the allocation does not escape the current container
+  // (e.g., via return or interprocedural function calls), so any deallocation
+  // is a direct user of the allocation. This scopes complexity and avoids
+  // the need for interprocedural escape analysis.
+  for (Operation *user : llvm::make_early_inc_range(allocLikeOp->getUsers())) {
+    if (auto dealloc = dyn_cast<memref::DeallocOp>(user))
+      rewriter.eraseOp(dealloc);
----------------
matthias-springer wrote:

This won't be able to handle cases like:

```
%0 = alloc
%1 = alloc
%2 = select %c, %0, %1
%3 = select %c, %1, %0
dealloc %2
dealloc %3
```

That's fine, but the limitations should be documented.


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