#7436: Added deallocate and moved validate behind program cache

tenstorrent · Apr 16, 2024 · 927f56f · 927f56f
1 parent fd60892
commit 927f56f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 12 deletions.
diff --git a/tt_eager/tt_dnn/op_library/run_operation.cpp b/tt_eager/tt_dnn/op_library/run_operation.cpp
@@ -154,7 +154,7 @@ OutputTensors run_device_operation(
 
     tt::stl::hash::hash_t program_hash = 0;
     if (program_cache.is_enabled()) {
-        get_or_create_program = [&program_cache, &program_hash](
+        get_or_create_program = [&program_cache, &program_hash, &optional_output_tensors](
                                     const DeviceOperation<OutputTensors>& operation,
                                     const Tensors& input_tensors,
                                     const OptionalConstTensors& optional_input_tensors,
@@ -165,6 +165,7 @@ OutputTensors run_device_operation(
             bool cache_hit = program_ptr.has_value();
             log_debug(tt::LogOp, "Program Hash: {} ({})", program_hash, cache_hit ? "HIT" : "MISS");
             if (not cache_hit) {
+                operation.validate(input_tensors, optional_input_tensors, optional_output_tensors);
                 program_ptr = std::make_shared<operation::CacheableProgram<OutputTensors>>(operation.create_program(input_tensors, optional_input_tensors, output_tensors));
                 program_cache.insert(program_hash, program_ptr.value());
             }
@@ -194,16 +195,17 @@ OutputTensors run_device_operation(
             return program_with_callbacks.program;
         };
     } else {
-        get_or_create_program = [](const DeviceOperation<OutputTensors>& operation,
+        get_or_create_program = [&optional_output_tensors](const DeviceOperation<OutputTensors>& operation,
                                    const Tensors& input_tensors,
                                    const OptionalConstTensors& optional_input_tensors,
                                    OutputTensors& output_tensors) -> std::shared_ptr<Program> {
+            operation.validate(input_tensors, optional_input_tensors, optional_output_tensors);
             auto program_with_callbacks =
                 operation.create_program(input_tensors, optional_input_tensors, output_tensors);
             return std::make_shared<Program>(std::move(program_with_callbacks.program));
         };
     }
-    operation.validate(input_tensors, optional_input_tensors, optional_output_tensors);
+
     auto output_tensors = operation.create_output_tensors(input_tensors, optional_output_tensors);
     auto program = get_or_create_program(operation, input_tensors, optional_input_tensors, output_tensors);
     uint32_t device_id = detail::get_device(input_tensors, optional_input_tensors)->id();

diff --git a/ttnn/ttnn/operations/core.py b/ttnn/ttnn/operations/core.py
@@ -526,15 +526,7 @@ def deallocate(tensor: ttnn.Tensor, *, force=True) -> None:
         >>> tensor = ttnn.to_layout(tensor, layout=ttnn.TILE_LAYOUT)
         >>> ttnn.deallocate(tensor)
     """
-
-    if not ttnn.CONFIG.enable_fast_runtime_mode:
-
-        def impl(tensor):
-            tensor.deallocate(force=force)
-
-        ttl.tensor.decorate_external_operation(impl, function_name="(ttnn) deallocate")(tensor)
-    else:
-        tensor.deallocate(force=force)
+    tensor.deallocate(force=force)
 
 
 def _to_memory_config_validate_input_tensors(operation_name, input_tensor, *args, **kwargs):