revert 1

tenstorrent · Dec 15, 2024 · 11b8754 · 11b8754
1 parent 5097212
commit 11b8754
Showing 1 changed file with 16 additions and 10 deletions.
diff --git a/ttnn/cpp/ttnn/operations/conv/conv2d/conv2d_utils.cpp b/ttnn/cpp/ttnn/operations/conv/conv2d/conv2d_utils.cpp
@@ -698,25 +698,31 @@ std::tuple<ttnn::Tensor, ParallelConfig, ParallelConfig, bool, bool> shard_or_re
 
         // In case we are in auto sharded codepath and convolution maps to matmul
         // Skip sharding of the input tensor and run the matmul out of interleaved tensor.
-        bool auto_shard_mm = auto_shard && is_mm_conv;
         if (input_tensor_on_device) {
             if (is_mm_conv && input_tensor.layout() == Layout::ROW_MAJOR &&
                 parallel_config.shard_scheme != TensorMemoryLayout::HEIGHT_SHARDED) {
                 // Workaround #13979 ttnn::tilize doesn't support BLOCK_SHARDED layout
                 input_tensor =
                     ttnn::to_layout(input_tensor, Layout::TILE, std::nullopt, std::nullopt, input_tensor.device());
             }
-            if (!auto_shard_mm) {
-                auto resharded_input_tensor = ttnn::to_memory_config(
-                    input_tensor, input_tensor_sharded_memory_config, std::nullopt);
-                if (conv_config.deallocate_activation) {
-                    input_tensor.deallocate();
-                    resharded_input_tensor = ttnn::operations::core::reallocate(resharded_input_tensor, resharded_input_tensor.memory_config());
-                }
-                input_tensor = resharded_input_tensor;
+            auto resharded_input_tensor = ttnn::to_memory_config(
+                input_tensor, input_tensor_sharded_memory_config, std::nullopt);
+            if (conv_config.deallocate_activation) {
+                input_tensor.deallocate();
+                resharded_input_tensor = ttnn::operations::core::reallocate(resharded_input_tensor, resharded_input_tensor.memory_config());
             }
+            input_tensor = resharded_input_tensor;
         } else {
-            input_tensor = ttnn::to_device(input_tensor, device, (auto_shard_mm ? ttnn::DRAM_MEMORY_CONFIG : input_tensor_sharded_memory_config));
+            if (is_mm_conv && input_tensor.layout() == Layout::ROW_MAJOR &&
+                parallel_config.shard_scheme != TensorMemoryLayout::HEIGHT_SHARDED) {
+                // Workaround #13979 ttnn::tilize doesn't support BLOCK_SHARDED layout
+                input_tensor = ttnn::to_device(input_tensor, device, std::nullopt);
+                input_tensor =
+                    ttnn::to_layout(input_tensor, Layout::TILE, std::nullopt, std::nullopt, input_tensor.device());
+                input_tensor = ttnn::to_memory_config(input_tensor, input_tensor_sharded_memory_config, std::nullopt);
+            } else {
+                input_tensor = ttnn::to_device(input_tensor, device, input_tensor_sharded_memory_config);
+            }
         }
     }
     return {input_tensor, parallel_config, output_parallel_config, needs_shard_or_reshard, use_non_tile_height};