diff --git a/ttnn/cpp/ttnn/operations/eltwise/binary/binary.hpp b/ttnn/cpp/ttnn/operations/eltwise/binary/binary.hpp index 67e00b2a311..514b3df54ca 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/binary/binary.hpp +++ b/ttnn/cpp/ttnn/operations/eltwise/binary/binary.hpp @@ -249,22 +249,6 @@ constexpr auto ne_ = ttnn::register_operation_with_auto_launch_op< "ttnn::ne_", operations::binary::InplaceRelationalBinary>(); -constexpr auto rsub_binary = ttnn::register_operation_with_auto_launch_op< - "ttnn::rsub_binary", - operations::binary::BinaryOperation>(); -constexpr auto power_binary = ttnn::register_operation_with_auto_launch_op< - "ttnn::power_binary", - operations::binary::BinaryOperationSfpu>(); -constexpr auto bitwise_and_binary = ttnn::register_operation_with_auto_launch_op< - "ttnn::bitwise_and_binary", - operations::binary::BinaryOperationSfpu>(); -constexpr auto bitwise_or_binary = ttnn::register_operation_with_auto_launch_op< - "ttnn::bitwise_or_binary", - operations::binary::BinaryOperationSfpu>(); -constexpr auto bitwise_xor_binary = ttnn::register_operation_with_auto_launch_op< - "ttnn::bitwise_xor_binary", - operations::binary::BinaryOperationSfpu>(); - template ttnn::Tensor operator+(const ttnn::Tensor& input_tensor_a, InputBType scalar) { return add(input_tensor_a, scalar); diff --git a/ttnn/cpp/ttnn/operations/eltwise/binary/binary_composite.hpp b/ttnn/cpp/ttnn/operations/eltwise/binary/binary_composite.hpp index f366c1104f1..c89bf48fae6 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/binary/binary_composite.hpp +++ b/ttnn/cpp/ttnn/operations/eltwise/binary/binary_composite.hpp @@ -18,6 +18,69 @@ namespace operations { namespace binary { +/** + * @brief Performs element-wise power operation on the input with the exponent. + * When exponent is Tensor, the supported dtypes are float32 and bfloat16. + * The tested range for the input is (-30,30) and for the exponent is (-20, 20). + * + * @param input The input tensor, i.e the base. + * @param exponent The exponent + * @return The result tensor + */ +struct ExecutePower { + static Tensor invoke( + uint8_t queue_id, + const Tensor& input_tensor, + uint32_t exponent, + const std::optional& memory_config = std::nullopt, + const std::optional& optional_output_tensor = std::nullopt); + + static Tensor invoke( + const Tensor& input_tensor, + uint32_t exponent, + const std::optional& memory_config = std::nullopt, + const std::optional& optional_output_tensor = std::nullopt); + + static Tensor invoke( + uint8_t queue_id, + const Tensor& input_tensor, + float exponent, + const std::optional& memory_config = std::nullopt, + const std::optional& optional_output_tensor = std::nullopt); + + static Tensor invoke( + const Tensor& input_tensor, + float exponent, + const std::optional& memory_config = std::nullopt, + const std::optional& optional_output_tensor = std::nullopt); + + static Tensor invoke( + uint8_t queue_id, + float input_a, + const Tensor& exponent, + const std::optional& memory_config = std::nullopt, + const std::optional& optional_output_tensor = std::nullopt); + + static Tensor invoke( + float input_a, + const Tensor& exponent, + const std::optional& memory_config = std::nullopt, + const std::optional& optional_output_tensor = std::nullopt); + + static Tensor invoke( + uint8_t queue_id, + const Tensor& input_tensor, + const Tensor& exponent, + const std::optional& memory_config = std::nullopt, + const std::optional& optional_output_tensor = std::nullopt); + + static Tensor invoke( + const Tensor& input_tensor, + const Tensor& exponent, + const std::optional& memory_config = std::nullopt, + const std::optional& optional_output_tensor = std::nullopt); +}; + template struct ExecuteBinaryCompositeOps { static Tensor invoke( @@ -436,5 +499,6 @@ constexpr auto rsub = ttnn::register_operation_with_auto_launch_op<"ttnn::rsub", constexpr auto bitwise_and = ttnn::register_operation_with_auto_launch_op<"ttnn::bitwise_and", operations::binary::ExecuteBitwiseAnd>(); constexpr auto bitwise_or = ttnn::register_operation_with_auto_launch_op<"ttnn::bitwise_or", operations::binary::ExecuteBitwiseOr>(); constexpr auto bitwise_xor = ttnn::register_operation_with_auto_launch_op<"ttnn::bitwise_xor", operations::binary::ExecuteBitwiseXor>(); +constexpr auto pow = ttnn::register_operation_with_auto_launch_op<"ttnn::pow", operations::binary::ExecutePower>(); } // namespace ttnn diff --git a/ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp b/ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp index d244c63a83f..e053b6155ed 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp +++ b/ttnn/cpp/ttnn/operations/eltwise/binary/binary_pybind.hpp @@ -1265,6 +1265,124 @@ void bind_binary_inplace_operation( py::arg("activations") = std::nullopt, py::arg("input_tensor_a_activation") = std::nullopt}); } + +template +void bind_power(py::module& module, const binary_operation_t& operation, const std::string& note = "") { + auto doc = fmt::format( + R"doc( + Perform element-wise {0} operation on :attr:`input_tensor` with :attr:`exponent`. + + .. math:: + \mathrm{{output\_tensor}}_i = (\mathrm{{input\_tensor}}_i ** \mathrm{{exponent}}_i) + + Args: + input_tensor (ttnn.Tensor): the input tensor. + exponent (float, int): the exponent value. + + Keyword Args: + memory_config (ttnn.MemoryConfig, optional): memory configuration for the operation. Defaults to `None`. + output_tensor (ttnn.Tensor, optional): preallocated output tensor. Defaults to `None`. + queue_id (int, optional): command queue id. Defaults to `0`. + + Returns: + ttnn.Tensor: the output tensor. + + Note: + Supported dtypes, layouts, and ranks: + + .. list-table:: + :header-rows: 1 + + * - Dtypes + - Layouts + - Ranks + * - BFLOAT16, BFLOAT8_B + - TILE + - 2, 3, 4 + + {2} + + Example: + >>> tensor = ttnn.from_torch(torch.tensor([[1, 2], [3, 4]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device) + >>> exponent = 2 + >>> output = {1}(tensor, exponent) + )doc", + ttnn::pow.base_name(), + ttnn::pow.python_fully_qualified_name(), + note); + + bind_registered_operation( + module, + ttnn::pow, + doc, + // integer exponent + ttnn::pybind_overload_t{ + [](const binary_operation_t& self, + const Tensor& input_tensor, + uint32_t exponent, + const std::optional& memory_config, + const std::optional& output_tensor, + const uint8_t queue_id) -> ttnn::Tensor { + return self(queue_id, input_tensor, exponent, memory_config, output_tensor); + }, + py::arg("input_tensor"), + py::arg("exponent"), + py::kw_only(), + py::arg("memory_config") = std::nullopt, + py::arg("output_tensor") = std::nullopt, + py::arg("queue_id") = 0}, + + // float exponent + ttnn::pybind_overload_t{ + [](const binary_operation_t& self, + const Tensor& input_tensor, + float exponent, + const std::optional& memory_config, + std::optional output_tensor, + const uint8_t queue_id) -> ttnn::Tensor { + return self(queue_id, input_tensor, exponent, memory_config, output_tensor); + }, + py::arg("input_tensor"), + py::arg("exponent"), + py::kw_only(), + py::arg("memory_config") = std::nullopt, + py::arg("output_tensor") = std::nullopt, + py::arg("queue_id") = ttnn::DefaultQueueId}, + + // tensor exponent + ttnn::pybind_overload_t{ + [](const binary_operation_t& self, + const Tensor& input_tensor, + const Tensor& exponent, + const std::optional& memory_config, + std::optional output_tensor, + const uint8_t queue_id) -> ttnn::Tensor { + return self(queue_id, input_tensor, exponent, memory_config, output_tensor); + }, + py::arg("input_tensor"), + py::arg("exponent"), + py::kw_only(), + py::arg("memory_config") = std::nullopt, + py::arg("output_tensor") = std::nullopt, + py::arg("queue_id") = ttnn::DefaultQueueId}, + + // scalar input - tensor exponent + ttnn::pybind_overload_t{ + [](const binary_operation_t& self, + float input, + const Tensor& exponent, + const std::optional& memory_config, + std::optional output_tensor, + const uint8_t queue_id) -> ttnn::Tensor { + return self(queue_id, input, exponent, memory_config, output_tensor); + }, + py::arg("input"), + py::arg("exponent"), + py::kw_only(), + py::arg("memory_config") = std::nullopt, + py::arg("output_tensor") = std::nullopt, + py::arg("queue_id") = ttnn::DefaultQueueId}); +} } // namespace detail void py_module(py::module& module) { @@ -1689,6 +1807,9 @@ void py_module(py::module& module) { R"doc(Performs Not equal to in-place operation on :attr:`input_a` and :attr:`input_b` and returns the tensor with the same layout as :attr:`input_tensor`)doc", R"doc(\mathrm{{input\_tensor\_a}}\: != \mathrm{{input\_tensor\_b}})doc", R"doc(BFLOAT16, BFLOAT8_B)doc"); + + detail::bind_power( + module, ttnn::pow, R"doc(When :attr:`exponent` is a Tensor, supported dtypes are: BFLOAT16, FLOAT32)doc"); } } // namespace binary diff --git a/ttnn/cpp/ttnn/operations/eltwise/binary/device/binary_composite_op.cpp b/ttnn/cpp/ttnn/operations/eltwise/binary/device/binary_composite_op.cpp index f09d2b08e8a..27cdec97ee9 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/binary/device/binary_composite_op.cpp +++ b/ttnn/cpp/ttnn/operations/eltwise/binary/device/binary_composite_op.cpp @@ -131,7 +131,7 @@ Tensor ExecuteMaximum::invoke( Tensor ExecuteMaximum::invoke( const Tensor& input_a, float value, const std::optional& output_mem_config) { - Tensor t_diff = ttnn::rsub_unary(input_a, value, output_mem_config); + Tensor t_diff = ttnn::rsub(input_a, value, output_mem_config); Tensor result = ttnn::where(t_diff, value, input_a); return result; } @@ -526,6 +526,111 @@ Tensor ExecuteLCM::invoke( return ttnn::abs(result); } +// power - floating point exponent +Tensor ExecutePower::invoke( + uint8_t queue_id, + const Tensor& input_a, + float exponent, + const std::optional& output_mem_config, + const std::optional& output_tensor) { + TT_FATAL(exponent >= 0.0f, "works for positive exponents only"); + const uint32_t exponent_floor = static_cast(std::floor(exponent)); + if (static_cast(exponent_floor) == exponent) { + if (output_tensor.has_value()) { + ttnn::power(queue_id, input_a, exponent_floor, output_mem_config, output_tensor); + return output_tensor.value(); + } + return ttnn::power(queue_id, input_a, exponent_floor, output_mem_config); + } + const float exponent_trunc = exponent - static_cast(exponent_floor); + Tensor pow_trunc_log = ttnn::multiply( + queue_id, ttnn::log(queue_id, input_a, output_mem_config), exponent_trunc, std::nullopt, output_mem_config); + Tensor pow_frac = ttnn::exp(queue_id, pow_trunc_log, false, output_mem_config); + pow_trunc_log.deallocate(); + float t_nan = std::nanf(""); + Tensor result = ttnn::multiply( + queue_id, + ttnn::power(queue_id, input_a, exponent_floor, output_mem_config), + pow_frac, + std::nullopt, + output_mem_config); + // To handle negative inputs: + // in torch For -ve inputs with float exponent power returns nan + auto output_memory_config = output_tensor.has_value() ? output_tensor.value().memory_config() + : output_mem_config.value_or(input_a.memory_config()); + result = ttnn::where( + ttnn::ltz(queue_id, input_a, output_mem_config), t_nan, result, output_memory_config, output_tensor); + return result; +} + +// power - floating point exponent +Tensor ExecutePower::invoke( + const Tensor& input_a, + float exponent, + const std::optional& output_mem_config, + const std::optional& output_tensor) { + return ExecutePower::invoke(DefaultQueueId, input_a, exponent, output_mem_config, std::move(output_tensor)); +} + +// power - integer exponent +Tensor ExecutePower::invoke( + uint8_t queue_id, + const Tensor& input, + uint32_t exponent, + const std::optional& output_mem_config, + const std::optional& output_tensor) { + return ttnn::power(queue_id, input, exponent, output_mem_config, output_tensor); +} + +// power - integer exponent +Tensor ExecutePower::invoke( + const Tensor& input, + uint32_t exponent, + const std::optional& output_mem_config, + const std::optional& output_tensor) { + return ExecutePower::invoke(DefaultQueueId, input, exponent, output_mem_config, std::move(output_tensor)); +} + +// power - tensor exponent +Tensor ExecutePower::invoke( + uint8_t queue_id, + const Tensor& input, + const Tensor& exponent, + const std::optional& output_mem_config, + const std::optional& output_tensor) { + return BinaryOperationSfpu::invoke( + queue_id, input, exponent, std::nullopt, output_mem_config, output_tensor); +} + +// power - tensor exponent +Tensor ExecutePower::invoke( + const Tensor& input, + const Tensor& exponent, + const std::optional& output_mem_config, + const std::optional& output_tensor) { + return ExecutePower::invoke(DefaultQueueId, input, exponent, output_mem_config, std::move(output_tensor)); +} + +// power - scalar input +Tensor ExecutePower::invoke( + uint8_t queue_id, + float input_a, + const Tensor& exponent, + const std::optional& output_mem_config, + const std::optional& output_tensor) { + Tensor input = ttnn::full_like(exponent, input_a); + return ExecutePower::invoke(queue_id, input, exponent, output_mem_config, std::move(output_tensor)); +} + +// power - scalar input +Tensor ExecutePower::invoke( + float input_a, + const Tensor& exponent, + const std::optional& output_mem_config, + const std::optional& output_tensor) { + return ExecutePower::invoke(DefaultQueueId, input_a, exponent, output_mem_config, std::move(output_tensor)); +} + Tensor ExecuteRsub::invoke( uint8_t queue_id, const Tensor& input_tensor_a, @@ -535,8 +640,7 @@ Tensor ExecuteRsub::invoke( const std::optional& optional_output_tensor, const std::optional& activations, const std::optional& input_tensor_a_activation) { - - return ttnn::rsub_binary( + return BinaryOperation::invoke( queue_id, input_tensor_a, input_tensor_b, @@ -573,13 +677,8 @@ Tensor ExecuteRsub::invoke( const float input_b, const std::optional& memory_config, const std::optional& optional_output_tensor) { - - return ttnn::rsub_unary( - queue_id, - input_tensor_a, - input_b, - memory_config, - optional_output_tensor); + return ttnn::operations::unary::ExecuteUnaryWithFloatParameter::invoke( + queue_id, input_tensor_a, input_b, memory_config, optional_output_tensor); } Tensor ExecuteRsub::invoke( @@ -603,14 +702,8 @@ Tensor ExecuteBitwiseAnd::invoke( const Tensor& input_tensor_b, const std::optional& memory_config, const std::optional& optional_output_tensor) { - - return ttnn::bitwise_and_binary( - queue_id, - input_tensor_a, - input_tensor_b, - std::nullopt, - memory_config, - optional_output_tensor); + return BinaryOperationSfpu::invoke( + queue_id, input_tensor_a, input_tensor_b, std::nullopt, memory_config, optional_output_tensor); } Tensor ExecuteBitwiseAnd::invoke( @@ -633,13 +726,9 @@ Tensor ExecuteBitwiseAnd::invoke( const int32_t input_b, const std::optional& memory_config, const std::optional& optional_output_tensor) { - - return ttnn::bitwise_and_unary( - queue_id, - input_tensor_a, - input_b, - memory_config, - optional_output_tensor); + return ttnn::operations::unary:: + ExecuteUnaryWithIntegerParameter::invoke( + queue_id, input_tensor_a, input_b, memory_config, optional_output_tensor); } Tensor ExecuteBitwiseAnd::invoke( @@ -663,14 +752,8 @@ Tensor ExecuteBitwiseOr::invoke( const Tensor& input_tensor_b, const std::optional& memory_config, const std::optional& optional_output_tensor) { - - return ttnn::bitwise_or_binary( - queue_id, - input_tensor_a, - input_tensor_b, - std::nullopt, - memory_config, - optional_output_tensor); + return BinaryOperationSfpu::invoke( + queue_id, input_tensor_a, input_tensor_b, std::nullopt, memory_config, optional_output_tensor); } Tensor ExecuteBitwiseOr::invoke( @@ -693,13 +776,9 @@ Tensor ExecuteBitwiseOr::invoke( const int32_t input_b, const std::optional& memory_config, const std::optional& optional_output_tensor) { - - return ttnn::bitwise_or_unary( - queue_id, - input_tensor_a, - input_b, - memory_config, - optional_output_tensor); + return ttnn::operations::unary:: + ExecuteUnaryWithIntegerParameter::invoke( + queue_id, input_tensor_a, input_b, memory_config, optional_output_tensor); } Tensor ExecuteBitwiseOr::invoke( @@ -723,14 +802,8 @@ Tensor ExecuteBitwiseXor::invoke( const Tensor& input_tensor_b, const std::optional& memory_config, const std::optional& optional_output_tensor) { - - return ttnn::bitwise_xor_binary( - queue_id, - input_tensor_a, - input_tensor_b, - std::nullopt, - memory_config, - optional_output_tensor); + return BinaryOperationSfpu::invoke( + queue_id, input_tensor_a, input_tensor_b, std::nullopt, memory_config, optional_output_tensor); } Tensor ExecuteBitwiseXor::invoke( @@ -753,13 +826,9 @@ Tensor ExecuteBitwiseXor::invoke( const int32_t input_b, const std::optional& memory_config, const std::optional& optional_output_tensor) { - - return ttnn::bitwise_xor_unary( - queue_id, - input_tensor_a, - input_b, - memory_config, - optional_output_tensor); + return ttnn::operations::unary:: + ExecuteUnaryWithIntegerParameter::invoke( + queue_id, input_tensor_a, input_b, memory_config, optional_output_tensor); } Tensor ExecuteBitwiseXor::invoke( 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 index a170e3d9855..04baa0f9599 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_composite_op.cpp +++ b/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_composite_op.cpp @@ -42,110 +42,6 @@ Tensor _tanhshrink(const Tensor& x, const std::optional& output_me return result; } -// power - floating point exponent -Tensor ExecutePower::invoke( - uint8_t queue_id, - const Tensor& input_a, - float exponent, - const std::optional& output_mem_config, - std::optional output_tensor) { - TT_FATAL(exponent >= 0.0f, "works for positive exponents only"); - const uint32_t exponent_floor = static_cast(std::floor(exponent)); - if (static_cast(exponent_floor) == exponent) { - if (output_tensor.has_value()) { - ttnn::power(queue_id, input_a, exponent_floor, output_mem_config, output_tensor); - return output_tensor.value(); - } - return ttnn::power(queue_id, input_a, exponent_floor, output_mem_config); - } - const float exponent_trunc = exponent - static_cast(exponent_floor); - Tensor pow_trunc_log = ttnn::multiply( - queue_id, ttnn::log(queue_id, input_a, output_mem_config), exponent_trunc, std::nullopt, output_mem_config); - Tensor pow_frac = ttnn::exp(queue_id, pow_trunc_log, false, output_mem_config); - pow_trunc_log.deallocate(); - float t_nan = std::nanf(""); - Tensor result = ttnn::multiply( - queue_id, - ttnn::power(queue_id, input_a, exponent_floor, output_mem_config), - pow_frac, - std::nullopt, - output_mem_config); - // To handle negative inputs: - // in torch For -ve inputs with float exponent power returns nan - auto output_memory_config = output_tensor.has_value() ? output_tensor.value().memory_config() - : output_mem_config.value_or(input_a.memory_config()); - result = ttnn::where( - ttnn::ltz(queue_id, input_a, output_mem_config), t_nan, result, output_memory_config, output_tensor); - return result; -} - -// power - floating point exponent -Tensor ExecutePower::invoke( - const Tensor& input_a, - float exponent, - const std::optional& output_mem_config, - std::optional output_tensor) { - return ExecutePower::invoke(DefaultQueueId, input_a, exponent, output_mem_config, std::move(output_tensor)); -} - -// power - integer exponent -Tensor ExecutePower::invoke( - uint8_t queue_id, - const Tensor& input, - uint32_t exponent, - const std::optional& output_mem_config, - std::optional output_tensor) { - return ttnn::power(queue_id, input, exponent, output_mem_config, output_tensor); -} - -// power - integer exponent -Tensor ExecutePower::invoke( - const Tensor& input, - uint32_t exponent, - const std::optional& output_mem_config, - std::optional output_tensor) { - return ExecutePower::invoke(DefaultQueueId, input, exponent, output_mem_config, std::move(output_tensor)); -} - -// power - tensor exponent -Tensor ExecutePower::invoke( - uint8_t queue_id, - const Tensor& input, - const Tensor& exponent, - const std::optional& output_mem_config, - std::optional output_tensor) { - return ttnn::power_binary(queue_id, input, exponent, std::nullopt, output_mem_config, output_tensor); -} - -// power - tensor exponent -Tensor ExecutePower::invoke( - const Tensor& input, - const Tensor& exponent, - const std::optional& output_mem_config, - std::optional output_tensor) { - return ExecutePower::invoke(DefaultQueueId, input, exponent, output_mem_config, std::move(output_tensor)); -} - -// power - scalar input -Tensor ExecutePower::invoke( - uint8_t queue_id, - float input_a, - const Tensor& exponent, - const std::optional& output_mem_config, - std::optional output_tensor) { - Tensor input = ttnn::full_like(exponent, input_a); - return ExecutePower::invoke(queue_id, input, exponent, output_mem_config, std::move(output_tensor)); -} - -// power - scalar input -Tensor ExecutePower::invoke( - float input_a, - const Tensor& exponent, - const std::optional& output_mem_config, - std::optional output_tensor) { - return ExecutePower::invoke(DefaultQueueId, input_a, exponent, output_mem_config, std::move(output_tensor)); -} - // acosh(x) = log(x + sqrt(x^2 - 1)) Tensor _acosh(const Tensor& input_a, const std::optional& output_mem_config) { TT_FATAL(input_a.storage_type() == StorageType::DEVICE, "Unary operation requires input to be on Device."); diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/unary.hpp b/ttnn/cpp/ttnn/operations/eltwise/unary/unary.hpp index a5a8d89087d..957b1edf0e0 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/unary/unary.hpp +++ b/ttnn/cpp/ttnn/operations/eltwise/unary/unary.hpp @@ -341,7 +341,6 @@ REGISTER_UNARY_OPERATION_WITH_FAST_AND_APPROXIMATE_MODE(rsqrt, RSQRT); // Unaries with float parameter REGISTER_UNARY_OPERATION_WITH_FLOAT_PARAMETER(elu, ELU); -REGISTER_UNARY_OPERATION_WITH_FLOAT_PARAMETER(rsub_unary, RSUB); REGISTER_UNARY_OPERATION_WITH_FLOAT_PARAMETER(heaviside, HEAVISIDE); REGISTER_UNARY_OPERATION_WITH_FLOAT_PARAMETER(leaky_relu, LEAKY_RELU); REGISTER_UNARY_OPERATION_WITH_FLOAT_PARAMETER(relu_max, RELU_MAX); @@ -357,9 +356,6 @@ REGISTER_UNARY_OPERATION_WITH_FLOAT_PARAMETER(ne_unary, UNARY_NE); REGISTER_UNARY_OPERATION_WITH_INTEGER_PARAMETER(power, POWER, uint32_t); REGISTER_UNARY_OPERATION_WITH_INTEGER_PARAMETER(bitwise_left_shift, LEFT_SHIFT, int32_t); REGISTER_UNARY_OPERATION_WITH_INTEGER_PARAMETER(bitwise_right_shift, RIGHT_SHIFT, int32_t); -REGISTER_UNARY_OPERATION_WITH_INTEGER_PARAMETER(bitwise_and_unary, BITWISE_AND, int32_t); -REGISTER_UNARY_OPERATION_WITH_INTEGER_PARAMETER(bitwise_or_unary, BITWISE_OR, int32_t); -REGISTER_UNARY_OPERATION_WITH_INTEGER_PARAMETER(bitwise_xor_unary, BITWISE_XOR, int32_t); // Other unaries constexpr auto dropout = diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_composite.hpp b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_composite.hpp index a523e09cb30..b7c2535aa01 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_composite.hpp +++ b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_composite.hpp @@ -12,69 +12,6 @@ namespace ttnn { namespace operations { namespace unary { -/** - * @brief Performs element-wise power operation on the input with the exponent. - * When exponent is Tensor, the supported dtypes are float32 and bfloat16. - * The tested range for the input is (-30,30) and for the exponent is (-20, 20). - * - * @param input The input tensor, i.e the base. - * @param exponent The exponent - * @return The result tensor - */ -struct ExecutePower { - static Tensor invoke( - uint8_t queue_id, - const Tensor& input_tensor, - uint32_t exponent, - const std::optional& memory_config = std::nullopt, - std::optional optional_output_tensor = std::nullopt); - - static Tensor invoke( - const Tensor& input_tensor, - uint32_t exponent, - const std::optional& memory_config = std::nullopt, - std::optional optional_output_tensor = std::nullopt); - - static Tensor invoke( - uint8_t queue_id, - const Tensor& input_tensor, - float exponent, - const std::optional& memory_config = std::nullopt, - std::optional optional_output_tensor = std::nullopt); - - static Tensor invoke( - const Tensor& input_tensor, - float exponent, - const std::optional& memory_config = std::nullopt, - std::optional optional_output_tensor = std::nullopt); - - static Tensor invoke( - uint8_t queue_id, - float input_a, - const Tensor& exponent, - const std::optional& memory_config = std::nullopt, - std::optional optional_output_tensor = std::nullopt); - - static Tensor invoke( - float input_a, - const Tensor& exponent, - const std::optional& memory_config = std::nullopt, - std::optional optional_output_tensor = std::nullopt); - - static Tensor invoke( - uint8_t queue_id, - const Tensor& input_tensor, - const Tensor& exponent, - const std::optional& memory_config = std::nullopt, - std::optional optional_output_tensor = std::nullopt); - - static Tensor invoke( - const Tensor& input_tensor, - const Tensor& exponent, - const std::optional& memory_config = std::nullopt, - std::optional optional_output_tensor = std::nullopt); -}; - template struct ExecuteUnaryCompositeOp { static Tensor invoke(const Tensor& input_tensor, const std::optional& memory_config = std::nullopt) { @@ -235,7 +172,6 @@ auto transform_first_matching_arg(Lambda lambda, First&& first, Rest&&... rest) constexpr auto rdiv = ttnn::register_operation_with_auto_launch_op<"ttnn::rdiv", operations::unary::ExecuteRdiv>(); -constexpr auto pow = ttnn::register_operation_with_auto_launch_op<"ttnn::pow", operations::unary::ExecutePower>(); constexpr auto tanhshrink = ttnn::register_operation_with_auto_launch_op< "ttnn::tanhshrink", operations::unary::ExecuteUnaryCompositeOp>(); diff --git a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp index 5f0a6514b35..d343f42bb92 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp +++ b/ttnn/cpp/ttnn/operations/eltwise/unary/unary_pybind.hpp @@ -975,125 +975,6 @@ void bind_identity(py::module& module, const unary_operation_t& operation) { py::arg("queue_id") = 0}); } -template -void bind_power(py::module& module, const unary_operation_t& operation, const std::string& note = "") { - auto doc = fmt::format( - R"doc( - Perform element-wise {0} operation on :attr:`input_tensor` with :attr:`exponent`. - - .. math:: - \mathrm{{output\_tensor}}_i = \verb|{0}|(\mathrm{{input\_tensor}}_i ** \mathrm{{exponent}}_i) - - Args: - input_tensor (ttnn.Tensor): the input tensor. - exponent (float, int): the exponent value. - - Keyword Args: - memory_config (ttnn.MemoryConfig, optional): memory configuration for the operation. Defaults to `None`. - output_tensor (ttnn.Tensor, optional): preallocated output tensor. Defaults to `None`. - queue_id (int, optional): command queue id. Defaults to `0`. - - Returns: - ttnn.Tensor: the output tensor. - - Note: - Supported dtypes, layouts, and ranks: - - .. list-table:: - :header-rows: 1 - - * - Dtypes - - Layouts - - Ranks - * - BFLOAT16, BFLOAT8_B - - TILE - - 2, 3, 4 - - {2} - - Example: - >>> tensor = ttnn.from_torch(torch.tensor([[1, 2], [3, 4]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device) - >>> exponent = 2 - >>> output = {1}(tensor, exponent) - )doc", - ttnn::pow.base_name(), - ttnn::pow.python_fully_qualified_name(), - note); - - bind_registered_operation( - module, - ttnn::pow, - doc, - // integer exponent - ttnn::pybind_overload_t{ - [](const unary_operation_t& self, - const Tensor& input_tensor, - uint32_t exponent, - const std::optional& memory_config, - std::optional output_tensor, - const uint8_t queue_id) -> ttnn::Tensor { - return self(queue_id, input_tensor, exponent, memory_config, output_tensor); - }, - py::arg("input_tensor"), - py::arg("exponent"), - py::kw_only(), - py::arg("memory_config") = std::nullopt, - py::arg("output_tensor") = std::nullopt, - py::arg("queue_id") = 0}, - - // float exponent - ttnn::pybind_overload_t{ - [](const unary_operation_t& self, - const Tensor& input_tensor, - float exponent, - const std::optional& memory_config, - std::optional output_tensor, - const uint8_t queue_id) -> ttnn::Tensor { - return self(queue_id, input_tensor, exponent, memory_config, output_tensor); - }, - py::arg("input_tensor"), - py::arg("exponent"), - py::kw_only(), - py::arg("memory_config") = std::nullopt, - py::arg("output_tensor") = std::nullopt, - py::arg("queue_id") = ttnn::DefaultQueueId}, - - // tensor exponent - ttnn::pybind_overload_t{ - [](const unary_operation_t& self, - const Tensor& input_tensor, - const Tensor& exponent, - const std::optional& memory_config, - std::optional output_tensor, - const uint8_t queue_id) -> ttnn::Tensor { - return self(queue_id, input_tensor, exponent, memory_config, output_tensor); - }, - py::arg("input_tensor"), - py::arg("exponent"), - py::kw_only(), - py::arg("memory_config") = std::nullopt, - py::arg("output_tensor") = std::nullopt, - py::arg("queue_id") = ttnn::DefaultQueueId}, - - // scalar input - tensor exponent - ttnn::pybind_overload_t{ - [](const unary_operation_t& self, - float input, - const Tensor& exponent, - const std::optional& memory_config, - std::optional output_tensor, - const uint8_t queue_id) -> ttnn::Tensor { - return self(queue_id, input, exponent, memory_config, output_tensor); - }, - py::arg("input"), - py::arg("exponent"), - py::kw_only(), - py::arg("memory_config") = std::nullopt, - py::arg("output_tensor") = std::nullopt, - py::arg("queue_id") = ttnn::DefaultQueueId} - ); -} - template void bind_unary_composite( py::module& module, @@ -1923,7 +1804,6 @@ void py_module(py::module& module) { detail::bind_sigmoid_accurate(module, ttnn::sigmoid_accurate); detail::bind_unary_chain(module, ttnn::unary_chain); detail::bind_identity(module, ttnn::identity); - detail::bind_power(module, ttnn::pow, R"doc(When :attr:`exponent` is a Tensor, supported dtypes are: BFLOAT16, FLOAT32)doc"); // unary composite imported into ttnn detail::bind_unary_composite(module, ttnn::deg2rad, R"doc(Performs deg2rad function on :attr:`input_tensor`.)doc", "", R"doc(BFLOAT16, BFLOAT8_B)doc"); diff --git a/ttnn/cpp/ttnn/operations/reduction/generic/generic_reductions.cpp b/ttnn/cpp/ttnn/operations/reduction/generic/generic_reductions.cpp index 6cb44e93b09..6b5e63c9433 100644 --- a/ttnn/cpp/ttnn/operations/reduction/generic/generic_reductions.cpp +++ b/ttnn/cpp/ttnn/operations/reduction/generic/generic_reductions.cpp @@ -5,7 +5,7 @@ #include "ttnn/operations/reduction/generic/generic_reductions.hpp" #include "ttnn/operations/data_movement/transpose/transpose.hpp" #include "ttnn/operations/eltwise/unary/unary.hpp" -#include "ttnn/operations/eltwise/unary/unary_composite.hpp" +#include "ttnn/operations/eltwise/binary/binary_composite.hpp" #include "ttnn/operations/reduction/generic/device/reduce_op.hpp" #include "ttnn/operations/core/core.hpp" namespace ttnn {