#6285: Add backward support for floor round and div_no_nan

tenstorrent · Mar 12, 2024 · 48b150d · 48b150d
1 parent b76f05a
commit 48b150d
Show file tree

Hide file tree

Showing 7 changed files with 211 additions and 0 deletions.
diff --git a/docs/source/ttnn/dependencies/tt_lib.rst b/docs/source/ttnn/dependencies/tt_lib.rst
@@ -1040,6 +1040,12 @@ Backward Operations
 
 .. autofunction:: tt_lib.tensor.le_bw
 
+.. autofunction:: tt_lib.tensor.floor_bw
+
+.. autofunction:: tt_lib.tensor.round_bw
+
+.. autofunction:: tt_lib.tensor.unary_div_no_nan_bw
+
 Loss Functions
 ==============
 

diff --git a/tests/tt_eager/python_api_testing/unit_testing/backward_ops/test_backward_floor.py b/tests/tt_eager/python_api_testing/unit_testing/backward_ops/test_backward_floor.py
@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import pytest
+import tt_lib
+from tests.tt_eager.python_api_testing.unit_testing.backward_ops.utility_funcs import compare_results, data_gen_pt_tt
+
+
+@pytest.mark.parametrize(
+    "input_shapes",
+    (
+        (torch.Size([1, 1, 32, 32])),
+        (torch.Size([1, 1, 320, 384])),
+        (torch.Size([1, 3, 320, 384])),
+    ),
+)
+def test_bw_floor(input_shapes, device):
+    grad_data, grad_tensor = data_gen_pt_tt(input_shapes, device)
+    in_data = torch.Tensor(size=input_shapes).uniform_()
+    in_data.requires_grad = True
+    input_tensor = (
+        tt_lib.tensor.Tensor(in_data, tt_lib.tensor.DataType.BFLOAT16).to(tt_lib.tensor.Layout.TILE).to(device)
+    )
+
+    pyt_y = torch.floor(in_data)
+
+    tt_output_tensor_on_device = tt_lib.tensor.floor_bw(grad_tensor)
+
+    in_data.retain_grad()
+
+    pyt_y.backward(gradient=grad_data)
+
+    golden_tensor = [in_data.grad]
+    comp_pass = compare_results(tt_output_tensor_on_device, golden_tensor)
+    assert comp_pass
diff --git a/tests/tt_eager/python_api_testing/unit_testing/backward_ops/test_backward_round.py b/tests/tt_eager/python_api_testing/unit_testing/backward_ops/test_backward_round.py
@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import pytest
+import tt_lib
+from tests.tt_eager.python_api_testing.unit_testing.backward_ops.utility_funcs import compare_results, data_gen_pt_tt
+
+
+@pytest.mark.parametrize(
+    "input_shapes",
+    (
+        (torch.Size([1, 1, 32, 32])),
+        (torch.Size([1, 1, 320, 384])),
+        (torch.Size([1, 3, 320, 384])),
+    ),
+)
+def test_bw_round(input_shapes, device):
+    grad_data, grad_tensor = data_gen_pt_tt(input_shapes, device)
+    in_data = torch.Tensor(size=input_shapes).uniform_()
+    in_data.requires_grad = True
+    input_tensor = (
+        tt_lib.tensor.Tensor(in_data, tt_lib.tensor.DataType.BFLOAT16).to(tt_lib.tensor.Layout.TILE).to(device)
+    )
+
+    pyt_y = torch.round(in_data)
+
+    tt_output_tensor_on_device = tt_lib.tensor.round_bw(grad_tensor)
+
+    in_data.retain_grad()
+
+    pyt_y.backward(gradient=grad_data)
+
+    golden_tensor = [in_data.grad]
+    comp_pass = compare_results(tt_output_tensor_on_device, golden_tensor)
+    assert comp_pass
diff --git a/...s/tt_eager/python_api_testing/unit_testing/backward_ops/test_backward_unary_div_no_nan.py b/...s/tt_eager/python_api_testing/unit_testing/backward_ops/test_backward_unary_div_no_nan.py
@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import pytest
+import tt_lib
+from tests.tt_eager.python_api_testing.unit_testing.backward_ops.utility_funcs import data_gen_pt_tt, compare_results
+
+
+def torch_div_no_nan(input, scalar):
+    return torch.where(torch.tensor(scalar) == 0, torch.zeros_like(input), torch.div(input, scalar))
+
+
+@pytest.mark.parametrize(
+    "input_shapes",
+    (
+        (torch.Size([1, 1, 32, 32])),
+        (torch.Size([1, 1, 320, 384])),
+        (torch.Size([1, 3, 320, 384])),
+    ),
+)
+@pytest.mark.parametrize("scalar", [0.05, 0.0, -0.5, 5.12])
+def test_bw_unary_div_no_nan(input_shapes, scalar, device):
+    in_data, input_tensor = data_gen_pt_tt(input_shapes, device, True)
+    grad_data, grad_tensor = data_gen_pt_tt(input_shapes, device)
+
+    tt_output_tensor_on_device = tt_lib.tensor.unary_div_no_nan_bw(grad_tensor, input_tensor, scalar=scalar)
+
+    in_data.retain_grad()
+
+    pyt_y = torch_div_no_nan(in_data, scalar)
+
+    pyt_y.backward(gradient=grad_data)
+
+    golden_tensor = [in_data.grad]
+    golden_tensor[0] = torch.where(torch.isnan(golden_tensor[0]), torch.zeros_like(in_data), golden_tensor[0])
+
+    status = compare_results(tt_output_tensor_on_device, golden_tensor)
+    assert status
diff --git a/tt_eager/tt_dnn/op_library/backward/backward_ops.cpp b/tt_eager/tt_dnn/op_library/backward/backward_ops.cpp
@@ -1551,6 +1551,41 @@ std::vector<Tensor> le_bw(const Tensor& grad, const MemoryConfig& output_mem_con
     return operation::decorate_as_composite(__func__, _le_bw)(grad, 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);
+    grad_tensor.emplace_back(t_zero);
+    return grad_tensor;
+}
+std::vector<Tensor> round_bw(const Tensor& grad, const MemoryConfig& output_mem_config)
+{
+    return operation::decorate_as_composite(__func__, _round_bw)(grad, output_mem_config);
+}
+
+std::vector<Tensor> _unary_div_no_nan_bw(const Tensor& grad, const Tensor& input, float scalar, const MemoryConfig& output_mem_config) {
+    std::vector<Tensor> grad_tensor;
+    Tensor zeros = zeros_like(grad, output_mem_config);
+    Tensor val = full_like(input, scalar, output_mem_config);
+    Tensor result = where(eq_unary(val, 0, output_mem_config), zeros, mul_unary(grad, 1/scalar, output_mem_config), output_mem_config);
+    grad_tensor.emplace_back(result);
+    return grad_tensor;
+}
+std::vector<Tensor> unary_div_no_nan_bw(const Tensor& grad, const Tensor& input, float scalar, const MemoryConfig& output_mem_config)
+{
+    return operation::decorate_as_composite(__func__, _unary_div_no_nan_bw)(grad, input, scalar, output_mem_config);
+}
+
 }//namespace tt_metal
 
 }//namespace tt
diff --git a/tt_eager/tt_dnn/op_library/backward/backward_ops.hpp b/tt_eager/tt_dnn/op_library/backward/backward_ops.hpp
@@ -236,6 +236,13 @@ std::vector<Tensor> log2_bw(const Tensor& grad, const Tensor& input, const Memor
 std::vector<Tensor> ge_bw(const Tensor& grad, const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);
 
 std::vector<Tensor> le_bw(const Tensor& grad, const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_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);
+
+std::vector<Tensor> unary_div_no_nan_bw(const Tensor& grad, const Tensor& input, float scalar=1.0f, const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);
+
 } //namespace tt_metal
 
 } //namespace tt
diff --git a/tt_eager/tt_lib/csrc/tt_lib_bindings_tensor_backward_ops.cpp b/tt_eager/tt_lib/csrc/tt_lib_bindings_tensor_backward_ops.cpp
@@ -1811,5 +1811,54 @@ namespace tt::tt_metal::detail{
                 "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("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
+
+            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("unary_div_no_nan_bw", &tt::tt_metal::unary_div_no_nan_bw,
+            py::arg("grad").noconvert(), py::arg("input").noconvert(), py::arg("scalar") = 1.0f, py::arg("output_mem_config").noconvert() = operation::DEFAULT_OUTPUT_MEMORY_CONFIG, R"doc(
+            Performs backward operations for division with given ``grad`` and ``scalar`` with no nan.
+
+            Input tensors 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"
+                "input", "Tensor", "Tensor", "Tensor of shape [W, Z, Y, X]", "Yes"
+                "scalar", "Scalar value", "float", "default to 1.0f", "No"
+                "output_mem_config", "Layout of tensor in TT Accelerator device memory banks", "MemoryConfig", "Default is interleaved in DRAM", "No"
+        )doc");
     }
 }