#598: Runtime support for dealloc op, slight refactor (#611)

tenstorrent · Sep 5, 2024 · 344d4ad · 344d4ad
1 parent 9c61739
commit 344d4ad
Show file tree

Hide file tree

Showing 3 changed files with 213 additions and 165 deletions.
diff --git a/include/ttmlir/Target/TTNN/program.fbs b/include/ttmlir/Target/TTNN/program.fbs
@@ -121,6 +121,10 @@ table Conv2dOp {
   groups: uint32;
 }
 
+table DeallocOp {
+  in: tt.target.TensorRef;
+}
+
 union OpType {
   OpenDeviceOp,
   CloseDeviceOp,
@@ -135,7 +139,8 @@ union OpType {
   TransposeOp,
   Conv2dOp,
   ConcatOp,
-  ReshapeOp
+  ReshapeOp,
+  DeallocOp
 }
 
 table Operation {

diff --git a/lib/Target/TTNN/TTNNToFlatbuffer.cpp b/lib/Target/TTNN/TTNNToFlatbuffer.cpp
@@ -271,6 +271,14 @@ createSoftmaxOp(FlatbufferObjectCache &cache, SoftmaxOp op) {
   return ::tt::target::ttnn::CreateSoftmaxOp(*cache.fbb, in, out, dimension);
 }
 
+template <typename DeallocOp>
+::flatbuffers::Offset<::tt::target::ttnn::DeallocOp>
+createDeallocOp(FlatbufferObjectCache &cache, DeallocOp op) {
+  auto in =
+      cache.at<::tt::target::TensorRef>(getOperandThroughDPSOps(op.getInput()));
+  return ::tt::target::ttnn::CreateDeallocOp(*cache.fbb, in);
+}
+
 ::flatbuffers::Offset<::tt::target::ttnn::Operation>
 emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op,
                   std::string const &debugString) {
@@ -357,6 +365,10 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op,
     return createOperation(cache, createReshapeOp(cache, reshapeOp),
                            debugString);
   }
+  if (auto deallocOp = dyn_cast<DeallocOp>(op); deallocOp) {
+    return createOperation(cache, createDeallocOp(cache, deallocOp),
+                           debugString);
+  }
 
   llvm_unreachable("unhandled op in emitTTNNOperation");
 }