[mlir][bufferization] Canonicalizer to skip extract_strided_metadata when operand is already a base memref.

[maerhart]

Summary

The extract_strided_metadata will be heavily used by the new buffer deallocation pass to get the base memref and pass it to the deallocation operation. This commit factors out some simplification logic of the pass into a canonicalization pattern.

[springerm]

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

1001–1010

The combination of m_AllOf and m_AnyOf (which is just needed to capture the value?) is a bit confusing. I would write this without matchers:

auto extractStridedOp = memref.getDefiningOp<ExtractStridedMetadataOp>();
if (!extractStridedOp) return failure();
auto allocOp = extractStridedOp.getSource().getDefiningOp<MemoryEffectOpInterface>();
if (!allocOp) return failure();
if (allocOp.getEffectOnValue ...)

[maerhart]

Use regular getDefiningOp chain instead of pattern matching

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

 LogicalResult DeallocOp::verify() {
   if (getMemrefs().size() != getConditions().size())
     return emitOpError(
         "must have the same number of conditions as memrefs to deallocate");
   return success();
 }
 
 static LogicalResult updateDeallocIfChanged(DeallocOp deallocOp,
-                                            ArrayRef<Value> memrefs,
-                                            ArrayRef<Value> conditions,
+                                            ValueRange memrefs,
+                                            ValueRange conditions,
                                             PatternRewriter &rewriter) {
-  if (deallocOp.getMemrefs() == memrefs)
+  if (deallocOp.getMemrefs() == memrefs &&
+      deallocOp.getConditions() == conditions)
     return failure();
 
   rewriter.updateRootInPlace(deallocOp, [&]() {
     deallocOp.getMemrefsMutable().assign(memrefs);
     deallocOp.getConditionsMutable().assign(conditions);
   });
   return success();
 }
for (auto [memref, cond] :
       }
     }
 
     return updateDeallocIfChanged(deallocOp, newMemrefs, newConditions,
                                   rewriter);
   }
 };
 
+/// The `memref.extract_strided_metadata` is often inserted to get the base
+/// memref if the operand is not already guaranteed to be the result of a memref
+/// allocation operation. This canonicalization pattern removes this extraction
+/// operation if the operand is now produced by an allocation operation (e.g.,
+/// due to other canonicalizations simplifying the IR).
+///
+/// Example:
+/// ```mlir
+/// %alloc = memref.alloc() : memref<2xi32>
+/// %base_memref, %offset, %size, %stride = memref.extract_strided_metadata
+///   %alloc : memref<2xi32> -> memref<i32>, index, index, index
+/// bufferization.dealloc (%base_memref : memref<i32>) if (%cond)
+/// ```
+/// is canonicalized to
+/// ```mlir
+/// %alloc = memref.alloc() : memref<2xi32>
+/// bufferization.dealloc (%alloc : memref<2xi32>) if (%cond)
+/// ```
+struct SkipExtractMetadataOfAlloc : public OpRewritePattern<DeallocOp> {
+  using OpRewritePattern<DeallocOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DeallocOp deallocOp,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value> newMemrefs(
+        llvm::map_range(deallocOp.getMemrefs(), [&](Value memref) {
+          auto extractStridedOp =
+              memref.getDefiningOp<memref::ExtractStridedMetadataOp>();
+          if (!extractStridedOp)
+            return memref;
+          Value allocMemref = extractStridedOp.getOperand();
+          auto allocOp = allocMemref.getDefiningOp<MemoryEffectOpInterface>();
+          if (!allocOp)
+            return memref;
+          if (allocOp.getEffectOnValue<MemoryEffects::Allocate>(allocMemref))
+            return allocMemref;
+          return memref;
+        }));
+
+    return updateDeallocIfChanged(deallocOp, newMemrefs,
+                                  deallocOp.getConditions(), rewriter);
+  }
+};
+
 } // anonymous namespace
 
 void DeallocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
   results.add<DeallocRemoveDuplicateDeallocMemrefs,
               DeallocRemoveDuplicateRetainedMemrefs,
               DeallocRemoveDeallocMemrefsContainedInRetained, EraseEmptyDealloc,
-              EraseAlwaysFalseDealloc>(context);
+              EraseAlwaysFalseDealloc, SkipExtractMetadataOfAlloc>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc"

mlir/test/Dialect/Bufferization/canonicalize.mlir

func.func @dealloc_always_false_condition(%arg0: memref<2xi32>, %arg1: memref<2xi32>, %arg2: i1) {
   bufferization.dealloc (%arg0, %arg1 : memref<2xi32>, memref<2xi32>) if (%false, %arg2)
   return
 }
 
 // CHECK-LABEL: func @dealloc_always_false_condition
 //  CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: memref<2xi32>, [[ARG2:%.+]]: i1)
 //  CHECK-NEXT: bufferization.dealloc ([[ARG1]] : {{.*}}) if ([[ARG2]])
 //  CHECK-NEXT: return
+
+// -----
+
+func.func @dealloc_base_memref_extract_of_alloc(%arg0: memref<2xi32>, %arg1: i1, %arg2: i1, %arg3: memref<2xi32>) {
+  %alloc = memref.alloc() : memref<2xi32>
+  %base0, %size0, %stride0, %offset0 = memref.extract_strided_metadata %alloc : memref<2xi32> -> memref<i32>, index, index, index
+  %base1, %size1, %stride1, %offset1 = memref.extract_strided_metadata %arg3 : memref<2xi32> -> memref<i32>, index, index, index
+  bufferization.dealloc (%base0, %arg0, %base1 : memref<i32>, memref<2xi32>, memref<i32>) if (%arg1, %arg2, %arg2)
+  return
+}
+
+// CHECK-LABEL: func @dealloc_base_memref_extract_of_alloc
+//  CHECK-SAME: ([[ARG0:%.+]]: memref<2xi32>, [[ARG1:%.+]]: i1, [[ARG2:%.+]]: i1, [[ARG3:%.+]]: memref<2xi32>)
+//  CHECK-NEXT: [[ALLOC:%.+]] = memref.alloc() : memref<2xi32>
+//  CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[ARG3]] :
+//  CHECK-NEXT: bufferization.dealloc ([[ALLOC]], [[ARG0]], [[BASE]] : memref<2xi32>, memref<2xi32>, memref<i32>) if ([[ARG1]], [[ARG2]], [[ARG2]])
+//  CHECK-NEXT: return