#0: upsample fix for undefined behavior on multi-devices run.

tenstorrent · Sep 19, 2024 · ebb2c97 · ebb2c97
1 parent b4e309e
commit ebb2c97
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 12 deletions.
diff --git a/models/experimental/functional_unet/tests/test_unet_trace.py b/models/experimental/functional_unet/tests/test_unet_trace.py
@@ -27,7 +27,7 @@
 
 @skip_for_grayskull("UNet not currently supported on GS")
 @pytest.mark.models_performance_bare_metal
-@pytest.mark.parametrize("device_params", [{"l1_small_size": 68864, "trace_region_size": 423936}], indirect=True)
+@pytest.mark.parametrize("device_params", [{"l1_small_size": 68864, "trace_region_size": 444416}], indirect=True)
 @pytest.mark.parametrize(
     "batch, groups, iterations",
     ((2, 1, 16),),
@@ -83,7 +83,7 @@ def test_unet_trace(
 @skip_for_grayskull("UNet not currently supported on GS")
 @pytest.mark.models_performance_bare_metal
 @pytest.mark.parametrize(
-    "device_params", [{"l1_small_size": 68864, "trace_region_size": 423936, "num_command_queues": 2}], indirect=True
+    "device_params", [{"l1_small_size": 68864, "trace_region_size": 442368, "num_command_queues": 2}], indirect=True
 )
 @pytest.mark.parametrize(
     "batch, groups, iterations",
@@ -202,7 +202,7 @@ def buffer_address(tensor):
 @pytest.mark.models_performance_bare_metal
 @pytest.mark.parametrize("enable_async_mode", (True,), indirect=True)
 @pytest.mark.parametrize(
-    "device_params", [{"l1_small_size": 68864, "trace_region_size": 423936, "num_command_queues": 2}], indirect=True
+    "device_params", [{"l1_small_size": 68864, "trace_region_size": 442368, "num_command_queues": 2}], indirect=True
 )
 @pytest.mark.parametrize(
     "batch, groups, iterations",

diff --git a/ttnn/cpp/ttnn/operations/pool/upsample/upsample.cpp b/ttnn/cpp/ttnn/operations/pool/upsample/upsample.cpp
@@ -11,15 +11,13 @@ namespace ttnn::operations::upsample {
 
 ttnn::Tensor ExecuteUpSample::invoke(const ttnn::Tensor& input_tensor,
     std::variant<int, tt::tt_metal::Array2D, tt::tt_metal::Array3D, tt::tt_metal::Array4D> scale_factor,
-    std::string mode,
-    std::optional<MemoryConfig> output_mem_config,
-    std::optional<DeviceComputeKernelConfig> compute_kernel_config) {
+    const std::string &mode,
+    const std::optional<MemoryConfig>& output_mem_config,
+    const std::optional<DeviceComputeKernelConfig>& compute_kernel_config) {
         MemoryConfig mem_config = output_mem_config.value_or(input_tensor.memory_config());
         ttnn::DeviceComputeKernelConfig config = compute_kernel_config.value_or(
             ttnn::init_device_compute_kernel_config(input_tensor.device()->arch(), std::nullopt, MathFidelity::HiFi4));
-        if(mode.empty()) {
-            mode = "nearest";
-        }
+
         int scale_h = 1;
         int scale_w = 1;
         std::visit(

diff --git a/ttnn/cpp/ttnn/operations/pool/upsample/upsample.hpp b/ttnn/cpp/ttnn/operations/pool/upsample/upsample.hpp
@@ -16,9 +16,9 @@ struct ExecuteUpSample {
     static ttnn::Tensor invoke(
         const ttnn::Tensor& input_tensor,
         std::variant<int, tt::tt_metal::Array2D, tt::tt_metal::Array3D, tt::tt_metal::Array4D> scale_factor,
-        std::string mode="nearest",
-        std::optional<MemoryConfig> output_mem_config = std::nullopt,
-        std::optional<DeviceComputeKernelConfig> compute_kernel_config = std::nullopt);
+        const std::string& mode=std::string("nearest"),
+        const std::optional<MemoryConfig>& output_mem_config = std::nullopt,
+        const std::optional<DeviceComputeKernelConfig>& compute_kernel_config = std::nullopt);
 };
 } // upsample
 } // operations