-
Notifications
You must be signed in to change notification settings - Fork 87
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
832a92e
commit 90c51e3
Showing
6 changed files
with
148 additions
and
142 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
110 changes: 110 additions & 0 deletions
110
ttnn/cpp/ttnn/operations/pool/maxpool/max_pool2d_pybind.cpp
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,110 @@ | ||
| // SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. | ||
| // | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| #include "ttnn/cpp/pybind11/decorators.hpp" | ||
| #include "ttnn/types.hpp" | ||
|
|
||
| #include "ttnn/operations/pool/maxpool/max_pool2d.hpp" | ||
| #include "ttnn/operations/pool/maxpool/max_pool2d_pybind.hpp" | ||
|
|
||
|
|
||
| namespace ttnn::operations::pool { | ||
|
|
||
| void bind_max_pool2d_operation(py::module& module) { | ||
| bind_registered_operation( | ||
| module, | ||
| ttnn::max_pool2d_new, | ||
| R"doc( | ||
| Max Pool 2D | ||
| +-------------------+-------------------------------+---------------+-------------+----------+ | ||
| | Argument | Description | Data type | Valid range | Required | | ||
| +===================+===============================+===============+=============+==========+ | ||
| | input | Input activations tensor | Tensor | | Yes | | ||
| | in_n | Input nbatch | Tensor | | Yes | | ||
| | in_h | Input height | Tensor | | Yes | | ||
| | in_w | Input width | Tensor | | Yes | | ||
| | kernel_h | kernel window height | uint32_t | | Yes | | ||
| | kernel_w | kernel window width | uint32_t | | Yes | | ||
| | stride_h | stride in height dim | uint32_t | | No | | ||
| | stride_w | stride in width dim | uint32_t | | No | | ||
| | pad_h | padding in height dim | uint32_t | | No | | ||
| | pad_w | padding in width dim | uint32_t | | No | | ||
| | dilation_h | kernel dilation in height dim | uint32_t | | No | | ||
| | dilation_w | kernel dilation in width dim | uint32_t | | No | | ||
| | memory_config | Output memory config | MemoryConfig | | No | | ||
| +-------------------+-------------------------------+---------------+-------------+----------+ | ||
| )doc", | ||
| ttnn::pybind_overload_t{ | ||
| [](const decltype(ttnn::max_pool2d_new)& self, const ttnn::Tensor& input_tensor, | ||
| uint32_t batch_size, | ||
| uint32_t input_h, | ||
| uint32_t input_w, | ||
| uint32_t channels, | ||
| std::array<uint32_t, 2> kernel_size, | ||
| std::array<uint32_t, 2> stride, | ||
| std::array<uint32_t, 2> padding, | ||
| std::array<uint32_t, 2> dilation, | ||
| ttnn::Device* device, | ||
| const uint8_t& queue_id) | ||
| -> ttnn::Tensor { return self(queue_id, | ||
| input_tensor, | ||
| batch_size, | ||
| input_h, | ||
| input_w, | ||
| channels, | ||
| kernel_size, | ||
| stride, | ||
| padding, | ||
| dilation, | ||
| device); }, | ||
| py::arg("input_tensor"), | ||
| py::arg("batch_size"), | ||
| py::arg("input_h"), | ||
| py::arg("input_w"), | ||
| py::arg("channels"), | ||
| py::arg("kernel_size"), | ||
| py::arg("stride"), | ||
| py::arg("padding"), | ||
| py::arg("dilation"), | ||
| py::kw_only(), | ||
| py::arg("device"), | ||
| py::arg("queue_id") = 0}, | ||
| ttnn::pybind_overload_t{ | ||
| [](const decltype(ttnn::max_pool2d_new)& self, const ttnn::Tensor& input_tensor, | ||
| uint32_t batch_size, | ||
| uint32_t input_h, | ||
| uint32_t input_w, | ||
| uint32_t channels, | ||
| std::array<uint32_t, 2> kernel_size, | ||
| std::array<uint32_t, 2> stride, | ||
| std::array<uint32_t, 2> padding, | ||
| std::array<uint32_t, 2> dilation, | ||
| DeviceMesh* device, | ||
| const uint8_t& queue_id) | ||
| -> ttnn::Tensor { return self(queue_id, | ||
| input_tensor, | ||
| batch_size, | ||
| input_h, | ||
| input_w, | ||
| channels, | ||
| kernel_size, | ||
| stride, | ||
| padding, | ||
| dilation, | ||
| device); }, | ||
| py::arg("input_tensor"), | ||
| py::arg("batch_size"), | ||
| py::arg("input_h"), | ||
| py::arg("input_w"), | ||
| py::arg("channels"), | ||
| py::arg("kernel_size"), | ||
| py::arg("stride"), | ||
| py::arg("padding"), | ||
| py::arg("dilation"), | ||
| py::kw_only(), | ||
| py::arg("device"), | ||
| py::arg("queue_id") = 0}); | ||
| } | ||
|
|
||
| } // namespace ttnn::operations::pool |
Oops, something went wrong.