From 254b4dc5309c1a8756ea9996beacf1ee83c0493a Mon Sep 17 00:00:00 2001
From: Quentin Bourgerie <bourgerie.quentin@gmail.com>
Date: Fri, 25 Mar 2022 16:05:58 +0100
Subject: [PATCH] enhance(compiler): Enable deallocation pass to free temporary
 buffers

---
 compiler/include/concretelang/Dialect/RT/IR/RTOps.td | 4 +++-
 compiler/lib/Dialect/RT/IR/RTOps.cpp                 | 6 ++++++
 compiler/lib/Support/Pipeline.cpp                    | 3 +++
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/compiler/include/concretelang/Dialect/RT/IR/RTOps.td b/compiler/include/concretelang/Dialect/RT/IR/RTOps.td
index 5a4164fc9..50b27e980 100644
--- a/compiler/include/concretelang/Dialect/RT/IR/RTOps.td
+++ b/compiler/include/concretelang/Dialect/RT/IR/RTOps.td
@@ -12,7 +12,9 @@ include "concretelang/Dialect/RT/IR/RTTypes.td"
 class RT_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<RT_Dialect, mnemonic, traits>;
 
-def DataflowTaskOp : RT_Op<"dataflow_task", [SingleBlockImplicitTerminator<"DataflowYieldOp">]> {
+def DataflowTaskOp : RT_Op<"dataflow_task", [
+    DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+    SingleBlockImplicitTerminator<"DataflowYieldOp">]> {
     let arguments = (ins Variadic<AnyType>: $inputs);
     let results = (outs Variadic<AnyType>:$outputs);
 
diff --git a/compiler/lib/Dialect/RT/IR/RTOps.cpp b/compiler/lib/Dialect/RT/IR/RTOps.cpp
index c444b2eac..91b510252 100644
--- a/compiler/lib/Dialect/RT/IR/RTOps.cpp
+++ b/compiler/lib/Dialect/RT/IR/RTOps.cpp
@@ -29,3 +29,9 @@ void DataflowTaskOp::build(
   reg->push_back(body);
   result.addTypes(resultTypes);
 }
+
+void DataflowTaskOp::getSuccessorRegions(
+    Optional<unsigned> index, ArrayRef<Attribute> operands,
+    SmallVectorImpl<RegionSuccessor> &regions) {
+  regions.push_back(RegionSuccessor(&body()));
+}
diff --git a/compiler/lib/Support/Pipeline.cpp b/compiler/lib/Support/Pipeline.cpp
index 6eec827b5..38d5f0af8 100644
--- a/compiler/lib/Support/Pipeline.cpp
+++ b/compiler/lib/Support/Pipeline.cpp
@@ -237,6 +237,9 @@ lowerStdToLLVMDialect(mlir::MLIRContext &context, mlir::ModuleOp &module,
   addPotentiallyNestedPass(
       pm, mlir::concretelang::createFinalizingBufferizePass(), enablePass);
 
+  addPotentiallyNestedPass(pm, mlir::createBufferDeallocationPass(),
+                           enablePass);
+
   if (parallelizeLoops)
     addPotentiallyNestedPass(pm, mlir::createConvertSCFToOpenMPPass(),
                              enablePass);