Adds dealloc MLIR memory operation to StandardOps.

PiperOrigin-RevId: 208896071

mlir/include/mlir/IR/StandardOps.h

@@ -198,6 +198,34 @@ private:
   explicit ConstantAffineIntOp(const Operation *state) : ConstantOp(state) {}
 };
 
+/// The "dealloc" operation frees the region of memory referenced by a memref
+/// which was originally created by the "alloc" operation.
+/// The "dealloc" operation should not be called on memrefs which alias an
+//  alloc'd memref (i.e. memrefs returned by the "view" and "reshape"
+/// operations).
+///
+///   %0 = alloc() : memref<8x64xf32, (d0, d1) -> (d0, d1), 1>
+///
+///   dealloc %0 : memref<8x64xf32, (d0, d1) -> (d0, d1), 1>
+///
+class DeallocOp
+    : public OpBase<DeallocOp, OpTrait::OneOperand, OpTrait::ZeroResult> {
+public:
+  SSAValue *getMemRef() { return getOperand(); }
+  const SSAValue *getMemRef() const { return getOperand(); }
+
+  static StringRef getOperationName() { return "dealloc"; }
+
+  // Hooks to customize behavior of this op.
+  const char *verify() const;
+  static bool parse(OpAsmParser *parser, OperationState *result);
+  void print(OpAsmPrinter *p) const;
+
+private:
+  friend class Operation;
+  explicit DeallocOp(const Operation *state) : OpBase(state) {}
+};
+
 /// The "dim" operation takes a memref or tensor operand and returns an
 /// "affineint".  It requires a single integer attribute named "index".  It
 /// returns the size of the specified dimension.  For example:

mlir/lib/IR/StandardOps.cpp

@@ -267,6 +267,28 @@ OperationState ConstantAffineIntOp::build(Builder *builder, int64_t value) {
   return result;
 }
 
+//===----------------------------------------------------------------------===//
+// DeallocOp
+//===----------------------------------------------------------------------===//
+
+void DeallocOp::print(OpAsmPrinter *p) const {
+  *p << "dealloc " << *getMemRef() << " : " << *getMemRef()->getType();
+}
+
+bool DeallocOp::parse(OpAsmParser *parser, OperationState *result) {
+  OpAsmParser::OperandType memrefInfo;
+  MemRefType *type;
+
+  return parser->parseOperand(memrefInfo) || parser->parseColonType(type) ||
+         parser->resolveOperand(memrefInfo, type, result->operands);
+}
+
+const char *DeallocOp::verify() const {
+  if (!isa<MemRefType>(getMemRef()->getType()))
+    return "operand must be a memref";
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//
@@ -466,7 +488,7 @@ const char *StoreOp::verify() const {
 
 /// Install the standard operations in the specified operation set.
 void mlir::registerStandardOperations(OperationSet &opSet) {
-  opSet.addOperations<AddFOp, AffineApplyOp, AllocOp, ConstantOp, DimOp, LoadOp,
-                      ReturnOp, StoreOp>(
+  opSet.addOperations<AddFOp, AffineApplyOp, AllocOp, ConstantOp, DeallocOp,
+                      DimOp, LoadOp, ReturnOp, StoreOp>(
       /*prefix=*/"");
 }

mlir/test/IR/memory-ops.mlir

@@ -29,6 +29,17 @@ bb0:
   return
 }
 
+// CHECK-LABEL: cfgfunc @dealloc() {
+cfgfunc @dealloc() {
+bb0:
+  // CHECK: %0 = alloc() : memref<1024x64xf32, #map0>
+  %0 = alloc() : memref<1024x64xf32, (d0, d1) -> (d0, d1), 0>
+
+  // CHECK: dealloc %0 : memref<1024x64xf32, #map0>
+  dealloc %0 : memref<1024x64xf32, (d0, d1) -> (d0, d1), 0>
+  return
+}
+
 // CHECK-LABEL: cfgfunc @load_store
 cfgfunc @load_store() {
 bb0: