From 68a26fbc594076c55475a2f0723f999fe8a6c356 Mon Sep 17 00:00:00 2001
From: Jackson Nie <jnie@tenstorrent.com>
Date: Wed, 18 Dec 2024 13:33:02 -0500
Subject: [PATCH] Remove nop check (#1622)

The new submit API does not need to check for nops as it generically
returns tensors based on global id, regardless of if it's an input
tensor or not.
---
 runtime/lib/ttnn/program.cpp | 24 ------------------------
 runtime/lib/ttnn/runtime.cpp |  3 +++
 2 files changed, 3 insertions(+), 24 deletions(-)

diff --git a/runtime/lib/ttnn/program.cpp b/runtime/lib/ttnn/program.cpp
index 354f03e6f..b1f95a8e8 100644
--- a/runtime/lib/ttnn/program.cpp
+++ b/runtime/lib/ttnn/program.cpp
@@ -235,25 +235,6 @@ void ProgramExecutor::runOperation(const ::tt::target::ttnn::Operation *op) {
   }
 }
 
-// Nop is single input, output tensor where input is returned as output.
-static bool isNopProgram(const ::tt::target::ttnn::Program *program) {
-  return program->inputs()->size() == 1 && program->outputs()->size() == 1 &&
-         program->inputs()->Get(0)->global_id() ==
-             program->outputs()->Get(0)->global_id();
-}
-
-static ::ttnn::Tensor
-handleNopProgram(::tt::target::ttnn::Program const *program,
-                 std::vector<::ttnn::Tensor *> const &inputs) {
-  const ::ttnn::Tensor &input = *inputs[0];
-  ::ttnn::Tensor output =
-      ::ttnn::zeros(input.get_shape(), input.get_dtype(), input.get_layout());
-  const void *src = ::tt::tt_metal::get_raw_host_data_ptr(input);
-  void *dst = ::tt::tt_metal::get_raw_host_data_ptr(output);
-  std::memcpy(dst, src, input.volume() * input.element_size());
-  return output;
-}
-
 namespace legacy {
 
 static bool handleNopProgram(::tt::target::ttnn::Program const *program,
@@ -331,11 +312,6 @@ std::vector<Tensor> runProgram(::ttnn::MeshDevice &meshDevice,
   ::tt::target::ttnn::TTNNBinary const &fbb = *getBinary(executableHandle);
   ::tt::target::ttnn::Program const *program =
       fbb.programs()->Get(programIndex);
-  if (isNopProgram(program)) {
-    Tensor out =
-        utils::createRuntimeTensorFromTTNN(handleNopProgram(program, inputs));
-    return {out};
-  }
   std::unordered_map<uint32_t, ::ttnn::Tensor *> liveTensors;
   std::vector<uint32_t> programInputs;
   int inputIndex = 0;
diff --git a/runtime/lib/ttnn/runtime.cpp b/runtime/lib/ttnn/runtime.cpp
index 7d98f0d59..18bb0d47c 100644
--- a/runtime/lib/ttnn/runtime.cpp
+++ b/runtime/lib/ttnn/runtime.cpp
@@ -327,6 +327,9 @@ void memcpy(Tensor dst, Tensor src) {
 
 void deallocateTensor(Tensor &tensor, bool force) {
   ::ttnn::Tensor &ttnnTensor = tensor.as<::ttnn::Tensor>(DeviceRuntime::TTNN);
+  if (ttnnTensor.storage_type() == ::tt::tt_metal::StorageType::BORROWED) {
+    return;
+  }
   ::ttnn::deallocate(ttnnTensor, force);
 }