Skip to content

Commit

Permalink
#10071: Merge floor_bw to TTNN
Browse files Browse the repository at this point in the history
  • Loading branch information
mouliraj-mcw committed Jul 10, 2024
1 parent a19f7e6 commit 08c4e64
Show file tree
Hide file tree
Showing 9 changed files with 12 additions and 33 deletions.
1 change: 1 addition & 0 deletions docs/source/ttnn/ttnn/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,7 @@ Pointwise Unary
ttnn/multigammaln_bw
ttnn/add_bw
ttnn/eq_bw
ttnn/floor_bw

Pointwise Binary
================
Expand Down
2 changes: 0 additions & 2 deletions docs/source/ttnn/ttnn/dependencies/tt_lib.rst
Original file line number Diff line number Diff line change
Expand Up @@ -966,8 +966,6 @@ Backward Operations

.. autofunction:: tt_lib.tensor.repeat_bw

.. autofunction:: tt_lib.tensor.floor_bw

.. autofunction:: tt_lib.tensor.round_bw

.. autofunction:: tt_lib.tensor.unary_div_no_nan_bw
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import torch
import pytest
import tt_lib
import ttnn
from tests.tt_eager.python_api_testing.unit_testing.backward_ops.utility_funcs import compare_pcc, data_gen_with_range


Expand All @@ -22,7 +23,7 @@ def test_bw_floor(input_shapes, device):

pyt_y = torch.floor(in_data)

tt_output_tensor_on_device = tt_lib.tensor.floor_bw(grad_tensor)
tt_output_tensor_on_device = ttnn.floor_bw(grad_tensor, input_tensor)

in_data.retain_grad()

Expand Down
10 changes: 0 additions & 10 deletions tt_eager/tt_dnn/op_library/backward/backward_ops.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2053,16 +2053,6 @@ std::vector<Tensor> repeat_bw(
return operation::decorate_as_composite(__func__, _repeat_bw)(grad, input, shape, output_mem_config);
}

std::vector<Tensor> _floor_bw(const Tensor& grad, const MemoryConfig& output_mem_config) {
std::vector<Tensor> grad_tensor;
Tensor t_zero = zeros_like(grad, output_mem_config);
grad_tensor.emplace_back(t_zero);
return grad_tensor;
}
std::vector<Tensor> floor_bw(const Tensor& grad, const MemoryConfig& output_mem_config) {
return operation::decorate_as_composite(__func__, _floor_bw)(grad, output_mem_config);
}

std::vector<Tensor> _round_bw(const Tensor& grad, const MemoryConfig& output_mem_config) {
std::vector<Tensor> grad_tensor;
Tensor t_zero = zeros_like(grad, output_mem_config);
Expand Down
3 changes: 0 additions & 3 deletions tt_eager/tt_dnn/op_library/backward/backward_ops.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -528,9 +528,6 @@ std::vector<Tensor> complex_sub_bw(
std::vector<Tensor> repeat_bw(
const Tensor& grad, const Tensor& input, const Shape& shape, const MemoryConfig& output_mem_config);

std::vector<Tensor> floor_bw(
const Tensor& grad, const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);

std::vector<Tensor> round_bw(
const Tensor& grad, const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);

Expand Down
17 changes: 0 additions & 17 deletions tt_eager/tt_lib/csrc/tt_lib_bindings_tensor_backward_ops.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1477,23 +1477,6 @@ namespace tt::tt_metal::detail{
"output_mem_config", "Layout of tensor in TT Accelerator device memory banks", "MemoryConfig", "Default is interleaved in DRAM", "No"
)doc");


m_tensor.def("floor_bw", &tt::tt_metal::floor_bw,
py::arg("grad").noconvert(), py::arg("output_mem_config").noconvert() = operation::DEFAULT_OUTPUT_MEMORY_CONFIG, R"doc(
Returns an tensor of zeros like ``grad`` tensor
Input tensor must have BFLOAT16 data type.
Output tensor will have BFLOAT16 data type.
.. csv-table::
:header: "Argument", "Description", "Data type", "Valid range", "Required"
"grad", "Gradient tensor", "Tensor", "Tensor of shape [W, Z, Y, X]", "Yes"
"output_mem_config", "Layout of tensor in TT Accelerator device memory banks", "MemoryConfig", "Default is interleaved in DRAM", "No"
)doc");


m_tensor.def("round_bw", &tt::tt_metal::round_bw,
py::arg("grad").noconvert(), py::arg("output_mem_config").noconvert() = operation::DEFAULT_OUTPUT_MEMORY_CONFIG, R"doc(
Returns an tensor of zeros like ``grad`` tensor
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,8 @@ std::function<std::vector<ttnn::Tensor>(const Tensor&, const Tensor&, const Memo
return _assign_bw;
case UnaryBackwardOpType::MULTIGAMMALN_BW:
return _multigammaln_bw;
case UnaryBackwardOpType::FLOOR_BW:
return _floor_bw;
default:
TT_ASSERT(false && "Undefined op type");
return 0;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,5 +72,7 @@ constexpr auto assign_bw = ttnn::register_operation<operations::unary_backward::
constexpr auto multigammaln_bw = ttnn::register_operation<operations::unary_backward::ExecuteUnaryBackward<operations::unary_backward::UnaryBackwardOpType::MULTIGAMMALN_BW>>("ttnn::multigammaln_bw");
constexpr auto add_bw = ttnn::register_operation<operations::unary_backward::ExecuteUnaryBackward<operations::unary_backward::UnaryBackwardOpType::ADD_BW>>("ttnn::add_bw");
constexpr auto eq_bw = ttnn::register_operation<operations::unary_backward::ExecuteUnaryBackward<operations::unary_backward::UnaryBackwardOpType::EQ_BW>>("ttnn::eq_bw");
constexpr auto floor_bw = ttnn::register_operation<operations::unary_backward::ExecuteUnaryBackward<operations::unary_backward::UnaryBackwardOpType::FLOOR_BW>>("ttnn::floor_bw");


} // namespace ttnn
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,11 @@ void py_module(py::module& module) {
R"doc(Performs backward operations for equal to comparison on :attr:`input_tensor`, :attr:`alpha` or attr:`input_tensor_a`, attr:`input_tensor_b` with given :attr:`grad_tensor`.
Returns an tensor of zeros like input tensors.)doc");

detail::bind_unary_backward(
module,
ttnn::floor_bw,
R"doc(Performs backward operations for floor on :attr:`input_tensor` or attr:`input_tensor_a`, attr:`input_tensor_b` with given :attr:`grad_tensor`.)doc");

}

} // namespace binary_backward
Expand Down

0 comments on commit 08c4e64

Please sign in to comment.