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

[Federico Bruzzone]

================
@@ -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{});
----------------
FedericoBruzzone wrote:

Instead of `IntegerAttr{}` we can pass `alloc.getAlignmentAttr()` to maintain the alignment. Right?

Currently, the following:
```mlir
%0 = memref.alloc() {alignment = 128 : i64} : memref<8xf32>
```
after `alloc_to_global`:
```mlir
memref.global "private" @alloc : memref<8xf32> 
```


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