Adds dealloc MLIR memory operation to StandardOps.
PiperOrigin-RevId: 208896071
diff --git a/include/mlir/IR/StandardOps.h b/include/mlir/IR/StandardOps.h
index ae60174..460576f 100644
--- a/include/mlir/IR/StandardOps.h
+++ b/include/mlir/IR/StandardOps.h
@@ -198,6 +198,34 @@
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:
diff --git a/lib/IR/StandardOps.cpp b/lib/IR/StandardOps.cpp
index e78290e..a1603ba 100644
--- a/lib/IR/StandardOps.cpp
+++ b/lib/IR/StandardOps.cpp
@@ -268,6 +268,28 @@
}
//===----------------------------------------------------------------------===//
+// 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 @@
/// 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=*/"");
}
diff --git a/test/IR/memory-ops.mlir b/test/IR/memory-ops.mlir
index c8ab294..1dd281b 100644
--- a/test/IR/memory-ops.mlir
+++ b/test/IR/memory-ops.mlir
@@ -29,6 +29,17 @@
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: