#12815: update mish op (#14270)

tenstorrent · Oct 30, 2024 · dec1dc2 · dec1dc2
1 parent 312d97b
commit dec1dc2
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 32 deletions.
diff --git a/tests/ttnn/unit_tests/operations/eltwise/test_composite.py b/tests/ttnn/unit_tests/operations/eltwise/test_composite.py
@@ -321,8 +321,7 @@ def test_unary_composite_log1p_ttnn(input_shapes, device):
     ),
 )
 def test_unary_composite_mish_ttnn(input_shapes, device):
-    in_data1, input_tensor1 = data_gen_with_range(input_shapes, -100, 100, device)
-
+    in_data1, input_tensor1 = data_gen_with_range(input_shapes, -20, 100, device)
     output_tensor = ttnn.mish(input_tensor1)
     golden_function = ttnn.get_golden_function(ttnn.mish)
     golden_tensor = golden_function(in_data1)

diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_composite_op.cpp b/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_composite_op.cpp
@@ -315,26 +315,19 @@ Tensor _log1p(const Tensor& x, const std::optional<MemoryConfig>& output_mem_con
     return result_log1p;
 }
 
+// log[exp[x] + 1] => softplus[x]
 // mish[x] = x*tanh[softplus[x]]
-// use transformation y = x*tanh[softplus[x]] by broadcast
-// Ref: https://krutikabapat.github.io/Swish-Vs-Mish-Latest-Activation-Functions/
-Tensor _mish(const Tensor& x, const std::optional<MemoryConfig>& output_mem_config) {
-    std::vector<Tensor> output_tensors = {Tensor(operation::get_workers_for_op_output({x}))};
-    operation::launch_op(
-        [output_mem_config](
-            const std::vector<Tensor>& input_tensors,
-            const std::vector<std::optional<const Tensor>>& optional_input_tensors,
-            const std::vector<std::optional<Tensor>>& optional_output_tensors) mutable -> std::vector<Tensor> {
-            const auto& x = input_tensors.at(0);
-            Tensor sp_x = ttnn::softplus(x, 1.0f, 20.0f, output_mem_config);
-            Tensor tanh_x = ttnn::tanh(sp_x, output_mem_config);
-            sp_x.deallocate();
-            Tensor mish_x = ttnn::multiply(x, tanh_x, std::nullopt, output_mem_config);
-            return {mish_x};
-        },
-        {x},
-        output_tensors);
-    return output_tensors.at(0);
+Tensor ExecuteMish::invoke(const Tensor& x, const std::optional<MemoryConfig>& output_mem_config) {
+    using ttnn::operations::unary::UnaryWithParam;
+    using ttnn::operations::unary::UnaryOpType;
+    std::vector<UnaryWithParam> ops_chain = {
+    UnaryWithParam{UnaryOpType::EXP, 1.0f},
+    UnaryWithParam{UnaryOpType::ADD_UNARY_SFPU, 1.0f},
+    UnaryWithParam{UnaryOpType::LOG},
+    UnaryWithParam{UnaryOpType::TANH} };
+    Tensor result = ttnn::unary_chain(x, ops_chain, output_mem_config);
+    Tensor mish_x = ttnn::multiply(x, result, std::nullopt, output_mem_config);
+    return mish_x;
 }
 
 // multivariate log-gamma function

diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_composite_op.hpp b/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_composite_op.hpp
@@ -25,7 +25,6 @@ enum class UnaryCompositeOpType {
     DIGAMMA,
     LGAMMA,
     LOG1P,
-    MISH,
     MULTIGAMMALN,
     SINH,
     SOFTSIGN,
@@ -66,7 +65,6 @@ Tensor _cosh(const Tensor&, const std::optional<MemoryConfig>&);
 Tensor _digamma(const Tensor&, const std::optional<MemoryConfig>&);
 Tensor _lgamma(const Tensor&,  const std::optional<MemoryConfig>&);
 Tensor _log1p(const Tensor&, const std::optional<MemoryConfig>&);
-Tensor _mish(const Tensor&, const std::optional<MemoryConfig>&);
 Tensor _multigammaln(const Tensor&, const std::optional<MemoryConfig>&);
 Tensor _sinh(const Tensor&, const std::optional<MemoryConfig>&);
 Tensor _softsign(const Tensor&, const std::optional<MemoryConfig>&);
@@ -192,13 +190,6 @@ struct OpHandler<UnaryCompositeOpType::LOG1P> {
     }
 };
 
-template <>
-struct OpHandler<UnaryCompositeOpType::MISH> {
-    static Tensor handle(const Tensor& t1, const std::optional<MemoryConfig>& mem_cfg ) {
-        return _mish(t1, mem_cfg);
-    }
-};
-
 template <>
 struct OpHandler<UnaryCompositeOpType::MULTIGAMMALN> {
     static Tensor handle(const Tensor& t1, const std::optional<MemoryConfig>& mem_cfg ) {

diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_composite.hpp b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_composite.hpp
@@ -153,6 +153,12 @@ struct ExecuteRdiv {
         std::optional<Tensor> optional_output_tensor = std::nullopt);
 };
 
+struct ExecuteMish {
+    static Tensor invoke(
+        const Tensor& input_tensor,
+        const std::optional<MemoryConfig>& memory_config = std::nullopt);
+};
+
 }  // namespace unary
 }  // namespace operations
 
@@ -240,7 +246,7 @@ constexpr auto log1p = ttnn::register_operation_with_auto_launch_op<
     operations::unary::ExecuteUnaryCompositeOp<operations::unary::UnaryCompositeOpType::LOG1P>>();
 constexpr auto mish = ttnn::register_operation_with_auto_launch_op<
     "ttnn::mish",
-    operations::unary::ExecuteUnaryCompositeOp<operations::unary::UnaryCompositeOpType::MISH>>();
+    operations::unary::ExecuteMish>();
 constexpr auto multigammaln = ttnn::register_operation_with_auto_launch_op<
     "ttnn::multigammaln",
     operations::unary::ExecuteUnaryCompositeOp<operations::unary::UnaryCompositeOpType::MULTIGAMMALN>>();

diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp
@@ -1844,7 +1844,7 @@ void py_module(py::module& module) {
            +----------------------------+---------------------------------+-------------------+
 
         )doc");
-    detail::bind_unary_composite(module, ttnn::mish, R"doc(Performs mish function on :attr:`input_tensor`, not supported for grayskull.)doc");
+    detail::bind_unary_composite(module, ttnn::mish, R"doc(Performs mish function on :attr:`input_tensor`, not supported for grayskull.)doc", "[supported range -20 to inf]");
     detail::bind_unary_composite(module, ttnn::multigammaln, R"doc(Performs multigammaln function on :attr:`input_tensor`.)doc", "[supported range 1.6 to inf]",
         R"doc(Supported dtypes, layouts, and ranks: