[MLIR] Don't remove memref allocation if stored into another allocation

[wsmoses]

Summary

A canonicalization accidentally will remove a memref allocation if it is only stored into. However, this is incorrect if the allocation is the value being stored, not the allocation being stored into.

[rriddle]

mlir/lib/IR/Verifier.cpp

255 ↗

(On Diff #354589)

This looks unrelated. Can you remove?

[wsmoses]

mlir/lib/IR/Verifier.cpp

255 ↗

(On Diff #354589)

Yes, absolutely -- not sure how this got caught up in the diff.

[wsmoses]

Remove unrelated verifier patch

[rriddle]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

213

Nit: Drop the else after a return

The SO should also be camelCase.

[wsmoses]

Update format and naming

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 

 
 /// Fold alloc operations with no users or only store and dealloc uses.
 template <typename T>
 struct SimplifyDeadAlloc : public OpRewritePattern<T> {
   using OpRewritePattern<T>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(T alloc,
                                 PatternRewriter &rewriter) const override {
-    if (llvm::any_of(alloc->getUsers(), [](Operation *op) {
-          return !isa<StoreOp, DeallocOp>(op);
+    if (llvm::any_of(alloc->getUsers(), [&](Operation *op) {
+        if (auto storeOp = dyn_cast<StoreOp>(op))
+          return storeOp.value() == alloc;
+        return !isa<DeallocOp>(op);
         }))
       return failure();
 
     for (Operation *user : llvm::make_early_inc_range(alloc->getUsers()))
       rewriter.eraseOp(user);
 
     rewriter.eraseOp(alloc);
     return success();

mlir/test/Dialect/MemRef/canonicalize.mlir

 //  CHECK: %[[mem2:.+]] = memref.cast %[[mem1]] : memref<1xi32, #map> to memref<?xi32, #map>
 //  CHECK: return %[[mem2]] : memref<?xi32, #map>
 #map0 = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 func @alloc_const_fold_with_symbols2() -> memref<?xi32, #map0> {
   %c1 = constant 1 : index
   %0 = memref.alloc(%c1)[%c1, %c1] : memref<?xi32, #map0>
   return %0 : memref<?xi32, #map0>
 }
+
+// -----
+// CHECK-LABEL: func @allocator
+// CHECK:   %[[alloc:.+]] = memref.alloc
+// CHECK:   memref.store %[[alloc:.+]], %arg0
+func @allocator(%arg0 : memref<memref<?xi32>>, %arg1 : index)  {
+  %0 = memref.alloc(%arg1) : memref<?xi32>
+  memref.store %0, %arg0[] : memref<memref<?xi32>>
+  return 
+}