From e674160bcfb8380799f8045d3c476f6e4db1c01d Mon Sep 17 00:00:00 2001 From: Jerry Wu Date: Fri, 25 Oct 2024 08:02:16 +0000 Subject: [PATCH] #12151: Remove interim avg_pool2d --- tests/tt_eager/ops/test_average_pool.cpp | 4 ++-- .../trace_testing/misc/test_average_pool.py | 2 +- .../unit_testing/misc/test_average_pool.py | 2 +- ttnn/CMakeLists.txt | 2 +- ttnn/cpp/pybind11/operations/__init__.hpp | 2 +- .../global_avg_pool.cpp} | 4 ++-- .../global_avg_pool.hpp} | 14 +++++++------- .../global_avg_pool_pybind.hpp} | 19 +------------------ ttnn/tt_lib/fused_ops/average_pool.py | 8 ++++---- ttnn/ttnn/__init__.py | 1 - ttnn/ttnn/operations/pool.py | 4 ---- 11 files changed, 20 insertions(+), 42 deletions(-) rename ttnn/cpp/ttnn/operations/pool/{avgpool/avg_pool.cpp => global_avg_pool/global_avg_pool.cpp} (89%) rename ttnn/cpp/ttnn/operations/pool/{avgpool/avg_pool.hpp => global_avg_pool/global_avg_pool.hpp} (72%) rename ttnn/cpp/ttnn/operations/pool/{avgpool/avg_pool_pybind.hpp => global_avg_pool/global_avg_pool_pybind.hpp} (67%) diff --git a/tests/tt_eager/ops/test_average_pool.cpp b/tests/tt_eager/ops/test_average_pool.cpp index 969f4941d50..1f0f8f11b30 100644 --- a/tests/tt_eager/ops/test_average_pool.cpp +++ b/tests/tt_eager/ops/test_average_pool.cpp @@ -2,7 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 -#include "ttnn/operations/pool/avgpool/avg_pool.hpp" +#include "ttnn/operations/pool/global_avg_pool/global_avg_pool.hpp" #include "ttnn/operations/experimental/auto_format/auto_format.hpp" #include "ttnn/operations/numpy/functions.hpp" @@ -23,7 +23,7 @@ Tensor run_avg_pool_2d_resnet(tt::tt_metal::LegacyShape& tensor_shape, Device* d if (!AutoFormat::check_input_tensor_format(input_tensor, padded_input_shape)) { padded_input_tensor = AutoFormat::format_input_tensor(input_tensor, device, padded_input_shape, 0, Layout::TILE); // pad with 0s } - auto device_output = avg_pool2d(padded_input_tensor); + auto device_output = global_avg_pool2d(padded_input_tensor); return device_output.cpu(); }; diff --git a/tests/tt_eager/python_api_testing/trace_testing/misc/test_average_pool.py b/tests/tt_eager/python_api_testing/trace_testing/misc/test_average_pool.py index 976d988cbaa..7784996ce58 100644 --- a/tests/tt_eager/python_api_testing/trace_testing/misc/test_average_pool.py +++ b/tests/tt_eager/python_api_testing/trace_testing/misc/test_average_pool.py @@ -64,7 +64,7 @@ def test_run_average_pool(act_shape, dtype, device, use_program_cache, enable_as ttact_res = ttact.to(device) def run_ops(ttact_res): - return ttnn.avg_pool2d(ttact_res) + return ttnn.global_avg_pool2d(ttact_res) # Compile run_ops(ttact_res) diff --git a/tests/tt_eager/python_api_testing/unit_testing/misc/test_average_pool.py b/tests/tt_eager/python_api_testing/unit_testing/misc/test_average_pool.py index 924dc16e04c..44a2de6d048 100644 --- a/tests/tt_eager/python_api_testing/unit_testing/misc/test_average_pool.py +++ b/tests/tt_eager/python_api_testing/unit_testing/misc/test_average_pool.py @@ -44,7 +44,7 @@ def test_run_average_pool(act_shape, dtype, device): ttact = ttact.pad_to_tile(0.0) ttact = ttact.to(device) - out = ttnn.avg_pool2d(ttact) + out = ttnn.global_avg_pool2d(ttact) out = out.cpu().to(ttnn.ROW_MAJOR_LAYOUT) out_shape = [batch_size, 1, 1, channels] diff --git a/ttnn/CMakeLists.txt b/ttnn/CMakeLists.txt index c3e2ce3e2cd..8257159d66b 100644 --- a/ttnn/CMakeLists.txt +++ b/ttnn/CMakeLists.txt @@ -296,11 +296,11 @@ set(ALL_TTNN_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/normalization/softmax/softmax_pybind.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/normalization/softmax/device/multi_core/softmax_op_multi_core.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/normalization/softmax/device/softmax_op.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/avgpool/avg_pool.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/downsample/device/downsample_op.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/downsample/device/downsample_program_factory.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/downsample/downsample.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/downsample/downsample_pybind.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/maxpool/device/max_pool2d_device_op.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/maxpool/device/max_pool2d_multi_core_program_factory.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/pool/maxpool/max_pool2d.cpp diff --git a/ttnn/cpp/pybind11/operations/__init__.hpp b/ttnn/cpp/pybind11/operations/__init__.hpp index 6c165efb027..d8b31ec6590 100644 --- a/ttnn/cpp/pybind11/operations/__init__.hpp +++ b/ttnn/cpp/pybind11/operations/__init__.hpp @@ -36,8 +36,8 @@ #include "ttnn/operations/matmul/matmul_pybind.hpp" #include "ttnn/operations/moreh/moreh_pybind.hpp" #include "ttnn/operations/normalization/normalization_pybind.hpp" -#include "ttnn/operations/pool/avgpool/avg_pool_pybind.hpp" #include "ttnn/operations/pool/downsample/downsample_pybind.hpp" +#include "ttnn/operations/pool/global_avg_pool/global_avg_pool_pybind.hpp" #include "ttnn/operations/pool/maxpool/max_pool2d_pybind.hpp" #include "ttnn/operations/pool/upsample/upsample_pybind.hpp" #include "ttnn/operations/reduction/reduction_pybind.hpp" diff --git a/ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool.cpp b/ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool.cpp similarity index 89% rename from ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool.cpp rename to ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool.cpp index d6b67a146ff..6b1ae1eea97 100644 --- a/ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool.cpp +++ b/ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool.cpp @@ -2,7 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 -#include "ttnn/operations/pool/avgpool/avg_pool.hpp" +#include "ttnn/operations/pool/global_avg_pool/global_avg_pool.hpp" #include "ttnn/operations/reduction/generic/generic_reductions.hpp" namespace tt { @@ -22,7 +22,7 @@ Tensor pool_2d(const Tensor& input, const MemoryConfig& memory_config, const std } } -Tensor avg_pool2d(const Tensor& input, const MemoryConfig& memory_config, const std::optional& output_dtype) { +Tensor global_avg_pool2d(const Tensor& input, const MemoryConfig& memory_config, const std::optional& output_dtype) { TT_FATAL(input.storage_type() == StorageType::DEVICE, "Input tensor needs to be on device"); auto output = input; diff --git a/ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool.hpp b/ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool.hpp similarity index 72% rename from ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool.hpp rename to ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool.hpp index 047c3bf7368..f4c22b5b805 100644 --- a/ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool.hpp +++ b/ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool.hpp @@ -12,17 +12,17 @@ namespace tt { namespace tt_metal { -enum class PoolType { - AVG -}; +enum class PoolType { AVG }; -Tensor avg_pool2d(const Tensor& input, const MemoryConfig& memory_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG, const std::optional& output_dtype = std::nullopt); +Tensor global_avg_pool2d( + const Tensor& input, + const MemoryConfig& memory_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG, + const std::optional& output_dtype = std::nullopt); } // namespace tt_metal } // namespace tt - -#include "ttnn/operations/pool/avgpool/avg_pool.hpp" +#include "ttnn/operations/pool/global_avg_pool/global_avg_pool.hpp" #include "ttnn/decorators.hpp" #include "ttnn/operations/core/core.hpp" @@ -36,7 +36,7 @@ struct GlobalAveragePool2D { const std::optional& memory_config_arg = std::nullopt, const std::optional& output_dtype = std::nullopt) { auto memory_config = memory_config_arg.value_or(input.memory_config()); - auto result = tt::tt_metal::avg_pool2d(input, memory_config, output_dtype); + auto result = tt::tt_metal::global_avg_pool2d(input, memory_config, output_dtype); return result; } }; diff --git a/ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool_pybind.hpp b/ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool_pybind.hpp similarity index 67% rename from ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool_pybind.hpp rename to ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool_pybind.hpp index 6e71f05120c..320ff58ecc7 100644 --- a/ttnn/cpp/ttnn/operations/pool/avgpool/avg_pool_pybind.hpp +++ b/ttnn/cpp/ttnn/operations/pool/global_avg_pool/global_avg_pool_pybind.hpp @@ -8,7 +8,7 @@ #include #include "ttnn/cpp/pybind11/decorators.hpp" -#include "ttnn/operations/pool/avgpool/avg_pool.hpp" +#include "ttnn/operations/pool/global_avg_pool/global_avg_pool.hpp" #include "ttnn/types.hpp" namespace py = pybind11; @@ -64,23 +64,6 @@ void bind_global_avg_pool2d(py::module& module) { void py_module(py::module& module) { detail::bind_global_avg_pool2d(module); - module.def( - "avg_pool2d", - &avg_pool2d, - py::arg().noconvert(), - py::kw_only(), - py::arg("memory_config").noconvert() = operation::DEFAULT_OUTPUT_MEMORY_CONFIG, - py::arg("dtype").noconvert() = std::nullopt, - R"doc( - Average Pool 2D - It operates on tensors that have channels as the last dimension. - - +----------+----------------------------+------------+-------------------------------+----------+ - | Argument | Description | Data type | Valid range | Required | - +==========+============================+============+===============================+==========+ - | act | Input activations tensor | Tensor | | Yes | - +----------+----------------------------+------------+-------------------------------+----------+ - )doc"); } } // namespace avgpool diff --git a/ttnn/tt_lib/fused_ops/average_pool.py b/ttnn/tt_lib/fused_ops/average_pool.py index 8c13bc472b0..7af3e9b2808 100644 --- a/ttnn/tt_lib/fused_ops/average_pool.py +++ b/ttnn/tt_lib/fused_ops/average_pool.py @@ -6,9 +6,9 @@ import ttnn -def run_avg_pool_on_device_wrapper(device): - def avg_pool2d(x, output_mem_config, output_dtype=None): - out = ttnn.avg_pool2d(x, memory_config=output_mem_config, dtype=output_dtype) +def run_global_avg_pool_on_device_wrapper(device): + def global_avg_pool2d(x, output_mem_config, output_dtype=None): + out = ttnn.global_avg_pool2d(x, memory_config=output_mem_config, dtype=output_dtype) return out - return avg_pool2d + return global_avg_pool2d diff --git a/ttnn/ttnn/__init__.py b/ttnn/ttnn/__init__.py index 1ef02bb3f08..f0b90944980 100644 --- a/ttnn/ttnn/__init__.py +++ b/ttnn/ttnn/__init__.py @@ -296,7 +296,6 @@ def prelu(*args, **kwargs): # Alias for leaky_relu. TODO(#8544): implement PReL ) from ttnn.operations.conv2d import Conv2dConfig, get_conv_padded_input_shape_and_mem_config, get_conv_output_dim -from ttnn.operations.pool import avg_pool2d from ttnn.operations.conv1d import Conv1d, Conv1dConfig from ttnn.operations.transformer import SDPAProgramConfig diff --git a/ttnn/ttnn/operations/pool.py b/ttnn/ttnn/operations/pool.py index bb4c3793a73..1fb97b768bf 100644 --- a/ttnn/ttnn/operations/pool.py +++ b/ttnn/ttnn/operations/pool.py @@ -43,7 +43,3 @@ def golden_global_avg_pool2d(input_tensor: ttnn.Tensor): ttnn.attach_golden_function(ttnn.global_avg_pool2d, golden_global_avg_pool2d) - -avg_pool2d = ttnn.register_python_operation(name="ttnn.avg_pool2d", golden_function=golden_global_avg_pool2d)( - ttnn._ttnn.operations.pool.avg_pool2d -)