diff --git a/tests/ttnn/unit_tests/gtests/test_add.cpp b/tests/ttnn/unit_tests/gtests/test_add.cpp index 53226c5cb49..44de525b997 100644 --- a/tests/ttnn/unit_tests/gtests/test_add.cpp +++ b/tests/ttnn/unit_tests/gtests/test_add.cpp @@ -3,11 +3,11 @@ // SPDX-License-Identifier: Apache-2.0 #include "tests/tt_metal/tt_metal/unit_tests_common/common/common_fixture.hpp" -#include "ttnn_test_fixtures.hpp" #include "ttnn/device.hpp" #include "ttnn/operations/binary.hpp" #include "ttnn/operations/core.hpp" #include "ttnn/operations/creation.hpp" +#include "ttnn_test_fixtures.hpp" namespace ttnn { namespace operations { @@ -26,13 +26,17 @@ class Add1DTensorAndScalarFixture : public TTNNFixture, TEST_P(Add1DTensorAndScalarFixture, AddsScalarCorrectly) { auto param = GetParam(); const auto device_id = 0; - auto &device = ttnn::open_device(device_id); + auto& device = ttnn::open_device(device_id); std::array dimensions = {param.h, param.w}; ttnn::Shape shape(dimensions); - const auto input_tensor = ttnn::zeros(shape, ttnn::bfloat16, ttnn::TILE_LAYOUT, device); - const auto output_tensor = input_tensor + param.scalar; - const auto expected_tensor = ttnn::full(shape, param.scalar, ttnn::bfloat16, ttnn::TILE_LAYOUT, device); - TT_FATAL(tt::numpy::allclose<::bfloat16>(ttnn::from_device(expected_tensor), ttnn::from_device(output_tensor))); + + { + const auto input_tensor = ttnn::zeros(shape, ttnn::bfloat16, ttnn::TILE_LAYOUT, device, std::nullopt); + const auto output_tensor = input_tensor + param.scalar; + const auto expected_tensor = ttnn::operations::creation::full( + shape, param.scalar, ttnn::bfloat16, ttnn::TILE_LAYOUT, device, std::nullopt); + TT_FATAL(tt::numpy::allclose<::bfloat16>(ttnn::from_device(expected_tensor), ttnn::from_device(output_tensor))); + } ttnn::close_device(device); } diff --git a/tests/ttnn/unit_tests/module.mk b/tests/ttnn/unit_tests/module.mk index 457ac8ff890..1bd79b20c2b 100644 --- a/tests/ttnn/unit_tests/module.mk +++ b/tests/ttnn/unit_tests/module.mk @@ -2,7 +2,7 @@ TTNN_UNIT_TESTS_HOME_DIR = $(TT_METAL_HOME)/tests/ttnn/unit_tests -TTNN_UNIT_TESTS_DIRS := $(TTNN_UNIT_TESTS_HOME_DIR) $(TTNN_UNIT_TESTS_HOME_DIR)/operations +TTNN_UNIT_TESTS_DIRS := $(TTNN_UNIT_TESTS_HOME_DIR) $(TTNN_UNIT_TESTS_HOME_DIR)/gtests TTNN_UNIT_TESTS_SRCS := $(foreach dir,$(TTNN_UNIT_TESTS_DIRS),$(wildcard $(dir)/*.cpp)) diff --git a/tt_eager/tt_numpy/functions.hpp b/tt_eager/tt_numpy/functions.hpp index c8f513e579c..3daf3c32afe 100644 --- a/tt_eager/tt_numpy/functions.hpp +++ b/tt_eager/tt_numpy/functions.hpp @@ -30,7 +30,11 @@ namespace detail { template constexpr static DataType get_data_type() { - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return DataType::UINT16; + } else if constexpr (std::is_same_v) { + return DataType::INT32; + } else if constexpr (std::is_same_v) { return DataType::UINT32; } else if constexpr (std::is_same_v) { return DataType::FLOAT32; @@ -92,7 +96,8 @@ static Tensor full( return detail::full(shape, float(value), layout, device, output_mem_config); } case DataType::BFLOAT16: { - return detail::full(shape, bfloat16(value), layout, device, output_mem_config); + return detail::full( + shape, bfloat16(static_cast(value)), layout, device, output_mem_config); } default: TT_THROW("Unsupported DataType!"); } @@ -168,9 +173,9 @@ static Tensor ones_like( template static Tensor arange( - int64_t start, - int64_t stop, - int64_t step, + const int64_t start, + const int64_t stop, + const int64_t step, const Layout layout = Layout::ROW_MAJOR, Device* device = nullptr, const MemoryConfig& output_mem_config = MemoryConfig{ @@ -180,7 +185,7 @@ static Tensor arange( TT_ASSERT(step > 0, "Step must be greater than 0"); TT_ASSERT(start < stop, "Start must be less than step"); auto size = div_up((stop - start), step); - if (size%2 != 0){ + if (size % 2 != 0) { size++; } auto owned_buffer = tt_metal::owned_buffer::create(size); @@ -200,36 +205,6 @@ static Tensor arange( return output; } -template -static Tensor arange( - const int64_t& start, - const int64_t& step, - const Shape& shape, - const Layout layout = Layout::ROW_MAJOR, - Device* device = nullptr, - const MemoryConfig& output_mem_config = MemoryConfig{ - .memory_layout = tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { - constexpr DataType data_type = detail::get_data_type(); - // Current implementation restrictions - TT_ASSERT(step > 0, "Step must be greater than 0"); - auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(shape)); - - auto value = start; - for (auto index = 0; index < owned_buffer.size(); index++) { - if constexpr (std::is_same_v) { - owned_buffer[index++] = T(static_cast(value)); - } else { - owned_buffer[index++] = static_cast(value); - } - value += step; - } - auto output = Tensor(OwnedStorage{owned_buffer}, shape, data_type, layout); - if (device != nullptr) { - output = output.to(device, output_mem_config); - } - return output; -} - template static Tensor index_trilu( const Shape& shape, @@ -259,7 +234,7 @@ static Tensor index_trilu( owned_buffer[index + y * shape[ultimate] + x] = static_cast(value); } } // dim X - } // dim Y + } // dim Y index += offset; } auto output = Tensor(OwnedStorage{owned_buffer}, shape, data_type, Layout::ROW_MAJOR).to(layout); @@ -277,7 +252,6 @@ static Tensor index_width( Device* device = nullptr, const MemoryConfig& output_mem_config = MemoryConfig{ .memory_layout = tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { - auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(shape)); std::fill(owned_buffer.begin(), owned_buffer.end(), -std::numeric_limits::infinity()); auto& up_shape = shape.without_padding(); @@ -294,11 +268,11 @@ static Tensor index_width( value = value + 1; } // dim W value = 0; - index = index + (shape[ultimate] - up_shape[ultimate]); - }// dim H - index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); - } //dim c - } // dim N + index = index + (shape[ultimate] - up_shape[ultimate]); + } // dim H + index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); + } // dim c + } // dim N auto output = Tensor(OwnedStorage{owned_buffer}, shape, data_type, Layout::ROW_MAJOR).to(layout); if (device != nullptr) { output = output.to(device, output_mem_config); @@ -329,12 +303,12 @@ static Tensor index_height( owned_buffer[index++] = T(static_cast(value)); } // dim W value = value + 1; - index = index + (shape[ultimate] - up_shape[ultimate]); - } // dim H + index = index + (shape[ultimate] - up_shape[ultimate]); + } // dim H value = 0; - index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); - } // dim C - } // dim N + index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); + } // dim C + } // dim N auto output = Tensor(OwnedStorage{owned_buffer}, shape, data_type, Layout::ROW_MAJOR).to(layout); if (device != nullptr) { output = output.to(device, output_mem_config); @@ -365,11 +339,11 @@ static Tensor index_all( owned_buffer[index++] = T(static_cast(value)); value = value + 1; } // dim W - index = index + (shape[ultimate] - up_shape[ultimate]); - }// dim H - index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); - } // dim C - } // dim N + index = index + (shape[ultimate] - up_shape[ultimate]); + } // dim H + index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); + } // dim C + } // dim N auto output = Tensor(OwnedStorage{owned_buffer}, shape, data_type, Layout::ROW_MAJOR).to(layout); if (device != nullptr) { output = output.to(device, output_mem_config); @@ -396,18 +370,16 @@ static Tensor mask_padded_input( for (uint32_t c = 0; c < padded_shape[rank - 3]; c++) { for (uint32_t y = 0; y < padded_shape[penultimate]; y++) { for (uint32_t x = 0; x < padded_shape[ultimate]; x++) { - if (b < unpadded_shape[rank - 4] && c < unpadded_shape[rank - 3] && y < unpadded_shape[penultimate] && x < unpadded_shape[ultimate]) - { + if (b < unpadded_shape[rank - 4] && c < unpadded_shape[rank - 3] && + y < unpadded_shape[penultimate] && x < unpadded_shape[ultimate]) { owned_buffer[index++] = T(static_cast(1.0)); - } - else - { + } else { owned_buffer[index++] = T(static_cast(0.0)); } } // dim W - } // dim H - } // dim C - } // dim N + } // dim H + } // dim C + } // dim N auto output = Tensor(OwnedStorage{owned_buffer}, padded_shape, data_type, Layout::ROW_MAJOR).to(layout); if (device != nullptr) { output = output.to(device, output_mem_config); @@ -415,25 +387,30 @@ static Tensor mask_padded_input( return output; } -template -static Tensor fill_first_val_into_tensor(const Tensor& input_tensor, DataType data_type, - const Layout layout , Device * device = nullptr, - const MemoryConfig& output_mem_config = MemoryConfig{.memory_layout=tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { +template +static Tensor fill_first_val_into_tensor( + const Tensor& input_tensor, + DataType data_type, + const Layout layout, + Device* device = nullptr, + const MemoryConfig& output_mem_config = MemoryConfig{ + .memory_layout = tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { const Shape& s_a = input_tensor.get_legacy_shape(); - auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(s_a)); //ouput + auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(s_a)); // ouput auto device_buffer = input_tensor.device_buffer(); uint32_t size_in_bytes = device_buffer->size(); vector data_vec; - const char *TT_METAL_SLOW_DISPATCH_MODE = std::getenv("TT_METAL_SLOW_DISPATCH_MODE"); + const char* TT_METAL_SLOW_DISPATCH_MODE = std::getenv("TT_METAL_SLOW_DISPATCH_MODE"); if (TT_METAL_SLOW_DISPATCH_MODE == nullptr) { data_vec.resize(size_in_bytes / sizeof(T)); - tt::tt_metal::tensor_impl::read_data_from_device_buffer(input_tensor.device()->command_queue(), device_buffer, data_vec.data(), true); + tt::tt_metal::tensor_impl::read_data_from_device_buffer( + input_tensor.device()->command_queue(), device_buffer, data_vec.data(), true); } else { tt::tt_metal::tensor_impl::read_data_from_device_buffer(device_buffer, data_vec); } auto input_buffer = owned_buffer::create(std::move(data_vec)); const Shape input_tensor_strides = input_tensor.strides(); - for(uint32_t i = 0; i < tt_metal::compute_volume(s_a); i++) { + for (uint32_t i = 0; i < tt_metal::compute_volume(s_a); i++) { owned_buffer[i] = input_buffer[0]; } auto output = Tensor(OwnedStorage{owned_buffer}, s_a, data_type, layout).to(layout); @@ -443,41 +420,48 @@ static Tensor fill_first_val_into_tensor(const Tensor& input_tensor, DataType da return output; } -template -static Tensor prod_result_computation_GS(const Tensor& input_tensor, DataType data_type, - const Layout layout , Device * device = nullptr, - const MemoryConfig& output_mem_config = MemoryConfig{.memory_layout=tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { +template +static Tensor prod_result_computation_GS( + const Tensor& input_tensor, + DataType data_type, + const Layout layout, + Device* device = nullptr, + const MemoryConfig& output_mem_config = MemoryConfig{ + .memory_layout = tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { const Shape& s_a = input_tensor.get_legacy_shape(); - auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(s_a)); //ouput + auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(s_a)); // ouput auto device_buffer = input_tensor.device_buffer(); uint32_t size_in_bytes = device_buffer->size(); vector data_vec; - const char *TT_METAL_SLOW_DISPATCH_MODE = std::getenv("TT_METAL_SLOW_DISPATCH_MODE"); + const char* TT_METAL_SLOW_DISPATCH_MODE = std::getenv("TT_METAL_SLOW_DISPATCH_MODE"); if (TT_METAL_SLOW_DISPATCH_MODE == nullptr) { data_vec.resize(size_in_bytes / sizeof(T)); - tt::tt_metal::tensor_impl::read_data_from_device_buffer(input_tensor.device()->command_queue(), device_buffer, data_vec.data(), true); + tt::tt_metal::tensor_impl::read_data_from_device_buffer( + input_tensor.device()->command_queue(), device_buffer, data_vec.data(), true); } else { tt::tt_metal::tensor_impl::read_data_from_device_buffer(device_buffer, data_vec); } auto input_buffer = owned_buffer::create(std::move(data_vec)); const Shape input_tensor_strides = input_tensor.strides(); auto result = static_cast(1.0f); - for(uint32_t i = s_a[0]-1; i < s_a[0]; i++) { - for(int32_t j = s_a[1]-1; j < s_a[1]; j++) { - for(int32_t k = s_a[2]-32; k < s_a[2]; k++) { //access last tile - for(int32_t l = s_a[3]-32; l < s_a[3]; l++) { - auto input_index = l + input_tensor_strides[2] * k + input_tensor_strides[1] * j + input_tensor_strides[0] * i; - if(k>=s_a[2]-2 && l>=s_a[3]-32){ //to access 2*32 in TILE layout + for (uint32_t i = s_a[0] - 1; i < s_a[0]; i++) { + for (int32_t j = s_a[1] - 1; j < s_a[1]; j++) { + for (int32_t k = s_a[2] - 32; k < s_a[2]; k++) { // access last tile + for (int32_t l = s_a[3] - 32; l < s_a[3]; l++) { + auto input_index = + l + input_tensor_strides[2] * k + input_tensor_strides[1] * j + input_tensor_strides[0] * i; + if (k >= s_a[2] - 2 && l >= s_a[3] - 32) { // to access 2*32 in TILE layout result = result * static_cast(input_buffer[input_index]); owned_buffer[input_index] = static_cast(0.0f); - }else{ + } else { owned_buffer[input_index] = static_cast(0.0f); } } } } } - owned_buffer[0] = result; //store the result at the first position of the tensor,and the rest of the values as 0.0f + owned_buffer[0] = result; // store the result at the first position of the tensor,and the rest of the values as + // 0.0f auto output = Tensor(OwnedStorage{owned_buffer}, s_a, data_type, layout).to(layout); if (device != nullptr) { output = output.to(device, output_mem_config); @@ -485,19 +469,24 @@ static Tensor prod_result_computation_GS(const Tensor& input_tensor, DataType da return output; } -template -static Tensor prod_result_computation_WH_B0(const Tensor& input_tensor, DataType data_type, - const Layout layout , Device * device = nullptr, - const MemoryConfig& output_mem_config = MemoryConfig{.memory_layout=tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { +template +static Tensor prod_result_computation_WH_B0( + const Tensor& input_tensor, + DataType data_type, + const Layout layout, + Device* device = nullptr, + const MemoryConfig& output_mem_config = MemoryConfig{ + .memory_layout = tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { const Shape& s_a = input_tensor.get_legacy_shape(); - auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(s_a)); //ouput + auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(s_a)); // ouput auto device_buffer = input_tensor.device_buffer(); uint32_t size_in_bytes = device_buffer->size(); vector data_vec; - const char *TT_METAL_SLOW_DISPATCH_MODE = std::getenv("TT_METAL_SLOW_DISPATCH_MODE"); + const char* TT_METAL_SLOW_DISPATCH_MODE = std::getenv("TT_METAL_SLOW_DISPATCH_MODE"); if (TT_METAL_SLOW_DISPATCH_MODE == nullptr) { data_vec.resize(size_in_bytes / sizeof(T)); - tt::tt_metal::tensor_impl::read_data_from_device_buffer(input_tensor.device()->command_queue(), device_buffer, data_vec.data(), true); + tt::tt_metal::tensor_impl::read_data_from_device_buffer( + input_tensor.device()->command_queue(), device_buffer, data_vec.data(), true); } else { tt::tt_metal::tensor_impl::read_data_from_device_buffer(device_buffer, data_vec); } @@ -505,22 +494,27 @@ static Tensor prod_result_computation_WH_B0(const Tensor& input_tensor, DataType const Shape input_tensor_strides = input_tensor.strides(); auto result = static_cast(1.0f); // need to access the last 4 rows and alternating columns of index 17 ,19, 21, 23, 25, 27, 29, 31 - for(uint32_t i = s_a[0]-1; i < s_a[0]; i++) { - for(int32_t j = s_a[1]-1; j < s_a[1]; j++) { - for(int32_t k = s_a[2]-32; k < s_a[2]; k++) { //access last tile - for(int32_t l = s_a[3]-32; l < s_a[3]; l++) { - auto input_index = l + input_tensor_strides[2] * k + input_tensor_strides[1] * j + input_tensor_strides[0] * i; - if(k>=s_a[2]-4 && (l==s_a[3]-15 || l==s_a[3]-13 || l==s_a[3]-11 || l==s_a[3]-9 || l==s_a[3]-7 || l==s_a[3]-5 || l==s_a[3]-3 || l==s_a[3]-1)){ //to access 4*16 elements placed alternatively starting from index 17W in TILE layout + for (uint32_t i = s_a[0] - 1; i < s_a[0]; i++) { + for (int32_t j = s_a[1] - 1; j < s_a[1]; j++) { + for (int32_t k = s_a[2] - 32; k < s_a[2]; k++) { // access last tile + for (int32_t l = s_a[3] - 32; l < s_a[3]; l++) { + auto input_index = + l + input_tensor_strides[2] * k + input_tensor_strides[1] * j + input_tensor_strides[0] * i; + if (k >= s_a[2] - 4 && (l == s_a[3] - 15 || l == s_a[3] - 13 || l == s_a[3] - 11 || + l == s_a[3] - 9 || l == s_a[3] - 7 || l == s_a[3] - 5 || l == s_a[3] - 3 || + l == s_a[3] - 1)) { // to access 4*16 elements placed alternatively + // starting from index 17W in TILE layout result = result * static_cast(input_buffer[input_index]); owned_buffer[input_index] = static_cast(0.0f); - }else{ + } else { owned_buffer[input_index] = static_cast(0.0f); } } } } } - owned_buffer[0] = result; //store the result at the first position of the tensor,and the rest of the values as 0.0f + owned_buffer[0] = result; // store the result at the first position of the tensor,and the rest of the values as + // 0.0f auto output = Tensor(OwnedStorage{owned_buffer}, s_a, data_type, layout).to(layout); if (device != nullptr) { output = output.to(device, output_mem_config); @@ -528,7 +522,6 @@ static Tensor prod_result_computation_WH_B0(const Tensor& input_tensor, DataType return output; } - template static Tensor index_channel( const Shape& shape, @@ -537,7 +530,6 @@ static Tensor index_channel( Device* device = nullptr, const MemoryConfig& output_mem_config = MemoryConfig{ .memory_layout = tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { - auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(shape)); std::fill(owned_buffer.begin(), owned_buffer.end(), -std::numeric_limits::infinity()); auto& up_shape = shape.without_padding(); @@ -553,12 +545,12 @@ static Tensor index_channel( owned_buffer[index++] = T(static_cast(value)); } // dim W index = index + (shape[ultimate] - up_shape[ultimate]); - } // dim H + } // dim H value = value + 1; - index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); - } // dim C + index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); + } // dim C value = 0; - } // dim N + } // dim N auto output = Tensor(OwnedStorage{owned_buffer}, shape, data_type, Layout::ROW_MAJOR).to(layout); if (device != nullptr) { output = output.to(device, output_mem_config); @@ -574,7 +566,6 @@ static Tensor index_batch( Device* device = nullptr, const MemoryConfig& output_mem_config = MemoryConfig{ .memory_layout = tt::tt_metal::TensorMemoryLayout::INTERLEAVED}) { - auto owned_buffer = tt_metal::owned_buffer::create(tt_metal::compute_volume(shape)); std::fill(owned_buffer.begin(), owned_buffer.end(), -std::numeric_limits::infinity()); auto& up_shape = shape.without_padding(); @@ -590,11 +581,11 @@ static Tensor index_batch( owned_buffer[index++] = T(static_cast(value)); } // dim W index = index + (shape[ultimate] - up_shape[ultimate]); - } // dim H - index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); - } // dim C + } // dim H + index = index + ((shape[penultimate] - up_shape[penultimate]) * tt::constants::TILE_WIDTH); + } // dim C value = value + 1; - } // dim N + } // dim N auto output = Tensor(OwnedStorage{owned_buffer}, shape, data_type, Layout::ROW_MAJOR).to(layout); if (device != nullptr) { output = output.to(device, output_mem_config); diff --git a/ttnn/cpp/pybind11/operations/__init__.hpp b/ttnn/cpp/pybind11/operations/__init__.hpp index 252a5dd4955..4d55fa5a80e 100644 --- a/ttnn/cpp/pybind11/operations/__init__.hpp +++ b/ttnn/cpp/pybind11/operations/__init__.hpp @@ -10,6 +10,7 @@ #include "ccl.hpp" #include "core.hpp" #include "conv2d.hpp" +#include "creation.hpp" #include "data_movement.hpp" #include "embedding.hpp" #include "kv_cache.hpp" @@ -38,6 +39,9 @@ void py_module(py::module& module) { auto m_core = module.def_submodule("core", "core operations"); core::py_module(m_core); + auto m_creation = module.def_submodule("creation", "creation operations"); + creation::py_module(m_creation); + auto m_embedding = module.def_submodule("embedding", "embedding operations"); embedding::py_module(m_embedding); diff --git a/ttnn/cpp/pybind11/operations/binary.hpp b/ttnn/cpp/pybind11/operations/binary.hpp index 095ac9b0708..2a324d54164 100644 --- a/ttnn/cpp/pybind11/operations/binary.hpp +++ b/ttnn/cpp/pybind11/operations/binary.hpp @@ -20,7 +20,7 @@ namespace binary { namespace detail { template -void bind_binary(py::module& module, const binary_operation_t& operation, const std::string& description) { +void bind_binary_operation(py::module& module, const binary_operation_t& operation, const std::string& description) { auto doc = fmt::format( R"doc({0}(input_tensor_a: ttnn.Tensor, input_tensor_b: Union[ttnn.Tensor, int, float], *, memory_config: Optional[ttnn.MemoryConfig] = None, dtype: Optional[ttnn.DataType] = None, activations: Optional[List[str]] = None) -> ttnn.Tensor @@ -84,86 +84,86 @@ void bind_binary(py::module& module, const binary_operation_t& operation, const } // namespace detail void py_module(py::module& module) { - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::add, R"doc(Adds :attr:`input_tensor_a` to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{ input\_tensor\_a }}_i + \mathrm{{ input\_tensor\_b }}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::add_, R"doc(Adds :attr:`input_tensor_a` to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` in-place .. math:: \mathrm{{input\_tensor\_a}}_i + \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::subtract, R"doc(Subtracts :attr:`input_tensor_b` from :attr:`input_tensor_a` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{ input\_tensor\_a }}_i - \mathrm{{ input\_tensor\_b }}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::subtract_, R"doc(Subtracts :attr:`input_tensor_b` from :attr:`input_tensor_a` and returns the tensor with the same layout as :attr:`input_tensor_a` in-place .. math:: \mathrm{{input\_tensor\_a}}_i - \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::multiply, R"doc(Multiplies :attr:`input_tensor_a` by :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{ input\_tensor\_a }}_i \times \mathrm{{ input\_tensor\_b }}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::multiply_, R"doc(Multiplies :attr:`input_tensor_a` by :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` in-place .. math:: \mathrm{{input\_tensor\_a}}_i \times \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::eq, R"doc(Compares if :attr:`input_tensor_a` is equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{input\_tensor\_a}}_i == \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::ne, R"doc(Compares if :attr:`input_tensor_a` is not equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{input\_tensor\_a}}_i != \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::lt, R"doc(Compares if :attr:`input_tensor_a` is less than :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{input\_tensor\_a}}_i < \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::le, R"doc(MCompares if :attr:`input_tensor_a` is less than or equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{input\_tensor\_a}}_i <= \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::gt, R"doc(Compares if :attr:`input_tensor_a` is greater than :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{input\_tensor\_a}}_i > \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::ge, R"doc(Compares if :attr:`input_tensor_a` is greater than or equal to :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{input\_tensor\_a}}_i >= \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::logical_and, R"doc(Compute logical AND of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` .. math:: \mathrm{{input\_tensor\_a}}_i && \mathrm{{input\_tensor\_b}}_i)doc"); - detail::bind_binary( + detail::bind_binary_operation( module, ttnn::logical_or, R"doc(Compute logical OR of :attr:`input_tensor_a` and :attr:`input_tensor_b` and returns the tensor with the same layout as :attr:`input_tensor_a` diff --git a/ttnn/cpp/pybind11/operations/creation.hpp b/ttnn/cpp/pybind11/operations/creation.hpp new file mode 100644 index 00000000000..977a71b1471 --- /dev/null +++ b/ttnn/cpp/pybind11/operations/creation.hpp @@ -0,0 +1,217 @@ +// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include + +#include "ttnn/cpp/pybind11/decorators.hpp" +#include "ttnn/operations/creation.hpp" + +namespace py = pybind11; + +namespace ttnn { +namespace operations { +namespace creation { + +namespace detail { + +template +void bind_full_operation(py::module& module, const creation_operation_t& operation) { + auto doc = fmt::format( + R"doc({0}(shape: ttnn.Shape, fill_value: Union[int, float], dtype: Optional[ttnn.DataType] = None, layout: Optional[ttnn.Layout] = None, device: Optional[ttnn.Device] = None, memory_config: Optional[ttnn.MemoryConfig] = None)doc", + operation.name()); + + bind_registered_operation( + module, + operation, + doc, + ttnn::pybind_overload_t{ + [](const creation_operation_t& self, + const std::vector& shape, + const float fill_value, + const std::optional& dtype, + const std::optional& layout, + const std::optional>& device, + const std::optional& memory_config) -> ttnn::Tensor { + return self(ttnn::Shape{tt::tt_metal::Shape{shape}}, fill_value, dtype, layout, device, memory_config); + }, + py::arg("shape"), + py::arg("fill_value"), + py::arg("dtype") = std::nullopt, + py::arg("layout") = std::nullopt, + py::arg("device") = std::nullopt, + py::arg("memory_config") = std::nullopt}, + ttnn::pybind_overload_t{ + [](const creation_operation_t& self, + const std::vector& shape, + const int fill_value, + const std::optional& dtype, + const std::optional& layout, + const std::optional>& device, + const std::optional& memory_config) -> ttnn::Tensor { + return self(ttnn::Shape{tt::tt_metal::Shape{shape}}, fill_value, dtype, layout, device, memory_config); + }, + py::arg("shape"), + py::arg("fill_value"), + py::arg("dtype") = std::nullopt, + py::arg("layout") = std::nullopt, + py::arg("device") = std::nullopt, + py::arg("memory_config") = std::nullopt}); +} + +template +void bind_full_operation_with_hard_coded_value(py::module& module, const creation_operation_t& operation) { + auto doc = fmt::format( + R"doc({0}(shape: ttnn.Shape, dtype: Optional[ttnn.DataType] = None, layout: Optional[ttnn.Layout] = None, device: Optional[ttnn.Device] = None, memory_config: Optional[ttnn.MemoryConfig] = None)doc", + operation.name()); + + bind_registered_operation( + module, + operation, + doc, + ttnn::pybind_overload_t{ + [](const creation_operation_t& self, + const std::vector& shape, + const std::optional& dtype, + const std::optional& layout, + const std::optional>& device, + const std::optional& memory_config) -> ttnn::Tensor { + return self(ttnn::Shape{tt::tt_metal::Shape{shape}}, dtype, layout, device, memory_config); + }, + py::arg("shape"), + py::arg("dtype") = std::nullopt, + py::arg("layout") = std::nullopt, + py::arg("device") = std::nullopt, + py::arg("memory_config") = std::nullopt}); +} + +template +void bind_full_like_operation(py::module& module, const creation_operation_t& operation) { + auto doc = fmt::format( + R"doc({0}(tensor: ttnn.Tensor, fill_value: Union[int, float], dtype: Optional[ttnn.DataType] = None, layout: Optional[ttnn.Layout] = None, device: Optional[ttnn.Device] = None, memory_config: Optional[ttnn.MemoryConfig] = None)doc", + operation.name()); + + bind_registered_operation( + module, + operation, + doc, + ttnn::pybind_overload_t{ + [](const creation_operation_t& self, + const ttnn::Tensor& tensor, + const float fill_value, + const std::optional& dtype, + const std::optional& layout, + const std::optional>& device, + const std::optional& memory_config) -> ttnn::Tensor { + return self(tensor, fill_value, dtype, layout, device, memory_config); + }, + py::arg("tensor"), + py::arg("fill_value"), + py::arg("dtype") = std::nullopt, + py::arg("layout") = std::nullopt, + py::arg("device") = std::nullopt, + py::arg("memory_config") = std::nullopt}, + ttnn::pybind_overload_t{ + [](const creation_operation_t& self, + const ttnn::Tensor& tensor, + const int fill_value, + const std::optional& dtype, + const std::optional& layout, + const std::optional>& device, + const std::optional& memory_config) -> ttnn::Tensor { + return self(tensor, fill_value, dtype, layout, device, memory_config); + }, + py::arg("tensor"), + py::arg("fill_value"), + py::arg("dtype") = std::nullopt, + py::arg("layout") = std::nullopt, + py::arg("device") = std::nullopt, + py::arg("memory_config") = std::nullopt}); +} + +template +void bind_full_like_operation_with_hard_coded_value(py::module& module, const creation_operation_t& operation) { + auto doc = fmt::format( + R"doc({0}(tensor: ttnn.Tensor, dtype: Optional[ttnn.DataType] = None, layout: Optional[ttnn.Layout] = None, device: Optional[ttnn.Device] = None, memory_config: Optional[ttnn.MemoryConfig] = None)doc", + operation.name()); + + bind_registered_operation( + module, + operation, + doc, + ttnn::pybind_overload_t{ + [](const creation_operation_t& self, + const ttnn::Tensor& tensor, + const std::optional& dtype, + const std::optional& layout, + const std::optional>& device, + const std::optional& memory_config) -> ttnn::Tensor { + return self(tensor, dtype, layout, device, memory_config); + }, + py::arg("tensor"), + py::arg("dtype") = std::nullopt, + py::arg("layout") = std::nullopt, + py::arg("device") = std::nullopt, + py::arg("memory_config") = std::nullopt}); +} + +template +void bind_arange_operation(py::module& module, const creation_operation_t& operation) { + auto doc = fmt::format( + R"doc({0}(start: int = 0, stop: int, step: int = 1, dtype: ttnn.DataType = ttnn.bfloat16, device: ttnn.Device = None, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG)doc", + operation.name()); + + bind_registered_operation( + module, + operation, + doc, + ttnn::pybind_overload_t{ + [](const creation_operation_t& self, + const int64_t stop, + const DataType dtype, + const std::optional>& device, + const MemoryConfig& memory_config) -> ttnn::Tensor { return self(stop, dtype, device, memory_config); }, + py::arg("stop"), + py::arg("dtype") = ttnn::bfloat16, + py::arg("device") = std::nullopt, + py::arg("memory_config") = ttnn::DRAM_MEMORY_CONFIG} // namespace detail + , + ttnn::pybind_overload_t{ + [](const creation_operation_t& self, + const int64_t start, + const int64_t stop, + const int64_t step, + const DataType dtype, + const std::optional>& device, + const MemoryConfig& memory_config) -> ttnn::Tensor { + return self(start, stop, step, dtype, device, memory_config); + }, + py::arg("start"), + py::arg("stop"), + py::arg("step") = 1, + py::arg("dtype") = ttnn::bfloat16, + py::arg("device") = std::nullopt, + py::arg("memory_config") = ttnn::DRAM_MEMORY_CONFIG}); +} // namespace creation +} // namespace detail + +void py_module(py::module& module) { + detail::bind_full_operation(module, ttnn::full); + detail::bind_full_operation_with_hard_coded_value(module, ttnn::zeros); + detail::bind_full_operation_with_hard_coded_value(module, ttnn::ones); + detail::bind_full_operation_with_hard_coded_value(module, ttnn::empty); + + detail::bind_full_like_operation(module, ttnn::full_like); + detail::bind_full_like_operation_with_hard_coded_value(module, ttnn::zeros_like); + detail::bind_full_like_operation_with_hard_coded_value(module, ttnn::ones_like); + detail::bind_full_like_operation_with_hard_coded_value(module, ttnn::empty_like); + + detail::bind_arange_operation(module, ttnn::arange); +} + +} // namespace creation +} // namespace operations +} // namespace ttnn diff --git a/ttnn/cpp/pybind11/operations/pool.hpp b/ttnn/cpp/pybind11/operations/pool.hpp index 547ff7e30a1..38ee2fcdab8 100644 --- a/ttnn/cpp/pybind11/operations/pool.hpp +++ b/ttnn/cpp/pybind11/operations/pool.hpp @@ -21,7 +21,7 @@ namespace detail { void bind_global_avg_pool2d(py::module& module) { auto doc = fmt::format( - R"doc({0}(input_tensor: ttnn.Tensor, *, memory_config: Optional[ttnn.MemoryConfig] = None, dtype: Optional[ttnn.DataType] = None) -> ttnn.Tensor + R"doc({0}(input_tensor: ttnn.Tensor, *, memory_config: Optional[ttnn.MemoryConfig] = None, dtype: Optional[ttnn.DataType] = None) -> ttnn.Tensor Applies {0} to :attr:`input_tensor` by performing a 2D adaptive average pooling over an input signal composed of several input planes. This operation computes the average of all elements in each channel across the entire spatial dimensions. @@ -43,12 +43,12 @@ void bind_global_avg_pool2d(py::module& module) { >>> tensor = ttnn.from_torch(torch.randn((10, 3, 32, 32), dtype=ttnn.bfloat16), device=device) >>> output = {1}(tensor) )doc", - ttnn::operations::pool::global_avg_pool2d.name(), - ttnn::operations::pool::global_avg_pool2d.python_fully_qualified_name()); + ttnn::global_avg_pool2d.name(), + ttnn::global_avg_pool2d.python_fully_qualified_name()); bind_registered_operation( module, - ttnn::operations::pool::global_avg_pool2d, + ttnn::global_avg_pool2d, doc, ttnn::pybind_arguments_t{ py::arg("input_tensor"), @@ -59,9 +59,7 @@ void bind_global_avg_pool2d(py::module& module) { } // namespace detail -void py_module(py::module& module) { - detail::bind_global_avg_pool2d(module); -} +void py_module(py::module& module) { detail::bind_global_avg_pool2d(module); } } // namespace pool } // namespace operations diff --git a/ttnn/cpp/pybind11/operations/reduction.hpp b/ttnn/cpp/pybind11/operations/reduction.hpp index c804774a0a7..3aa8065eb2b 100644 --- a/ttnn/cpp/pybind11/operations/reduction.hpp +++ b/ttnn/cpp/pybind11/operations/reduction.hpp @@ -18,7 +18,7 @@ namespace reduction { namespace detail { template -void bind_reduction(py::module& module, const reduction_operation_t& operation) { +void bind_reduction_operation(py::module& module, const reduction_operation_t& operation) { auto doc = fmt::format( R"doc({0}(input_tensor: ttnn.Tensor, dim: Optional[Union[int, Tuple[int]]] = None, keepdim: bool = True, memory_config: Optional[ttnn.MemoryConfig] = None) -> ttnn.Tensor)doc", operation.name()); @@ -36,12 +36,12 @@ void bind_reduction(py::module& module, const reduction_operation_t& operation) } // namespace detail void py_module(py::module& module) { - detail::bind_reduction(module, ttnn::sum); - detail::bind_reduction(module, ttnn::mean); - detail::bind_reduction(module, ttnn::max); - detail::bind_reduction(module, ttnn::min); - detail::bind_reduction(module, ttnn::std); - detail::bind_reduction(module, ttnn::var); + detail::bind_reduction_operation(module, ttnn::sum); + detail::bind_reduction_operation(module, ttnn::mean); + detail::bind_reduction_operation(module, ttnn::max); + detail::bind_reduction_operation(module, ttnn::min); + detail::bind_reduction_operation(module, ttnn::std); + detail::bind_reduction_operation(module, ttnn::var); } } // namespace reduction diff --git a/ttnn/cpp/pybind11/operations/unary.hpp b/ttnn/cpp/pybind11/operations/unary.hpp index de065afb7de..cd6e6a722a8 100644 --- a/ttnn/cpp/pybind11/operations/unary.hpp +++ b/ttnn/cpp/pybind11/operations/unary.hpp @@ -20,7 +20,7 @@ namespace unary { namespace detail { template -void bind_unary(py::module& module, const unary_operation_t& operation) { +void bind_unary_operation(py::module& module, const unary_operation_t& operation) { auto doc = fmt::format( R"doc({0}(input_tensor: ttnn.Tensor, *, memory_config: Optional[ttnn.MemoryConfig] = None) -> ttnn.Tensor @@ -51,7 +51,7 @@ void bind_unary(py::module& module, const unary_operation_t& operation) { } template -void bind_unary_with_fast_and_approximate_mode(py::module& module, const unary_operation_t& operation) { +void bind_unary_operation_with_fast_and_approximate_mode(py::module& module, const unary_operation_t& operation) { auto doc = fmt::format( R"doc({0}(input_tensor: ttnn.Tensor, *, fast_and_approximate_mode: bool = False, memory_config: Optional[ttnn.MemoryConfig] = None) -> ttnn.Tensor @@ -87,7 +87,7 @@ void bind_unary_with_fast_and_approximate_mode(py::module& module, const unary_o } template -void bind_unary_with_float_parameter( +void bind_unary_operation_with_float_parameter( py::module& module, const unary_operation_t& operation, const std::string& parameter_name, @@ -122,7 +122,10 @@ void bind_unary_with_float_parameter( operation, doc, ttnn::pybind_arguments_t{ - py::arg("input_tensor"), py::arg(parameter_name.c_str()), py::kw_only(), py::arg("memory_config") = std::nullopt}); + py::arg("input_tensor"), + py::arg(parameter_name.c_str()), + py::kw_only(), + py::arg("memory_config") = std::nullopt}); } void bind_softplus(py::module& module) { @@ -165,59 +168,60 @@ void bind_softplus(py::module& module) { } // namespace detail void py_module(py::module& module) { - detail::bind_unary(module, ttnn::abs); - detail::bind_unary(module, ttnn::acos); - detail::bind_unary(module, ttnn::asin); - detail::bind_unary(module, ttnn::atan); - detail::bind_unary(module, ttnn::cos); - detail::bind_unary(module, ttnn::erfinv); - detail::bind_unary(module, ttnn::exp2); - detail::bind_unary(module, ttnn::expm1); - detail::bind_unary(module, ttnn::eqz); - detail::bind_unary(module, ttnn::gez); - detail::bind_unary(module, ttnn::gtz); - detail::bind_unary(module, ttnn::i0); - detail::bind_unary(module, ttnn::isfinite); - detail::bind_unary(module, ttnn::isinf); - detail::bind_unary(module, ttnn::isnan); - detail::bind_unary(module, ttnn::isneginf); - detail::bind_unary(module, ttnn::isposinf); - detail::bind_unary(module, ttnn::lez); - detail::bind_unary(module, ttnn::log); - detail::bind_unary(module, ttnn::log10); - detail::bind_unary(module, ttnn::log2); - detail::bind_unary(module, ttnn::logical_not); - detail::bind_unary(module, ttnn::ltz); - detail::bind_unary(module, ttnn::neg); - detail::bind_unary(module, ttnn::nez); - detail::bind_unary(module, ttnn::reciprocal); - detail::bind_unary(module, ttnn::relu); - detail::bind_unary(module, ttnn::relu6); - detail::bind_unary(module, ttnn::sigmoid); - detail::bind_unary(module, ttnn::sign); - detail::bind_unary(module, ttnn::signbit); - detail::bind_unary(module, ttnn::silu); - detail::bind_unary(module, ttnn::sin); - detail::bind_unary(module, ttnn::sqrt); - detail::bind_unary(module, ttnn::square); - detail::bind_unary(module, ttnn::tan); - detail::bind_unary(module, ttnn::tanh); + detail::bind_unary_operation(module, ttnn::abs); + detail::bind_unary_operation(module, ttnn::acos); + detail::bind_unary_operation(module, ttnn::asin); + detail::bind_unary_operation(module, ttnn::atan); + detail::bind_unary_operation(module, ttnn::cos); + detail::bind_unary_operation(module, ttnn::erfinv); + detail::bind_unary_operation(module, ttnn::exp2); + detail::bind_unary_operation(module, ttnn::expm1); + detail::bind_unary_operation(module, ttnn::eqz); + detail::bind_unary_operation(module, ttnn::gez); + detail::bind_unary_operation(module, ttnn::gtz); + detail::bind_unary_operation(module, ttnn::i0); + detail::bind_unary_operation(module, ttnn::isfinite); + detail::bind_unary_operation(module, ttnn::isinf); + detail::bind_unary_operation(module, ttnn::isnan); + detail::bind_unary_operation(module, ttnn::isneginf); + detail::bind_unary_operation(module, ttnn::isposinf); + detail::bind_unary_operation(module, ttnn::lez); + detail::bind_unary_operation(module, ttnn::log); + detail::bind_unary_operation(module, ttnn::log10); + detail::bind_unary_operation(module, ttnn::log2); + detail::bind_unary_operation(module, ttnn::logical_not); + detail::bind_unary_operation(module, ttnn::ltz); + detail::bind_unary_operation(module, ttnn::neg); + detail::bind_unary_operation(module, ttnn::nez); + detail::bind_unary_operation(module, ttnn::reciprocal); + detail::bind_unary_operation(module, ttnn::relu); + detail::bind_unary_operation(module, ttnn::relu6); + detail::bind_unary_operation(module, ttnn::sigmoid); + detail::bind_unary_operation(module, ttnn::sign); + detail::bind_unary_operation(module, ttnn::signbit); + detail::bind_unary_operation(module, ttnn::silu); + detail::bind_unary_operation(module, ttnn::sin); + detail::bind_unary_operation(module, ttnn::sqrt); + detail::bind_unary_operation(module, ttnn::square); + detail::bind_unary_operation(module, ttnn::tan); + detail::bind_unary_operation(module, ttnn::tanh); // Unaries with fast_and_approximate_mode - detail::bind_unary_with_fast_and_approximate_mode(module, ttnn::exp); - detail::bind_unary_with_fast_and_approximate_mode(module, ttnn::erf); - detail::bind_unary_with_fast_and_approximate_mode(module, ttnn::erfc); - detail::bind_unary_with_fast_and_approximate_mode(module, ttnn::gelu); - detail::bind_unary_with_fast_and_approximate_mode(module, ttnn::rsqrt); + detail::bind_unary_operation_with_fast_and_approximate_mode(module, ttnn::exp); + detail::bind_unary_operation_with_fast_and_approximate_mode(module, ttnn::erf); + detail::bind_unary_operation_with_fast_and_approximate_mode(module, ttnn::erfc); + detail::bind_unary_operation_with_fast_and_approximate_mode(module, ttnn::gelu); + detail::bind_unary_operation_with_fast_and_approximate_mode(module, ttnn::rsqrt); // Unaries with float parameter - detail::bind_unary_with_float_parameter(module, ttnn::elu, "alpha", "The alpha parameter for the ELU function"); - detail::bind_unary_with_float_parameter( + detail::bind_unary_operation_with_float_parameter( + module, ttnn::elu, "alpha", "The alpha parameter for the ELU function"); + detail::bind_unary_operation_with_float_parameter( module, ttnn::heaviside, "value", "The value parameter for the Heaviside function"); - detail::bind_unary_with_float_parameter( + detail::bind_unary_operation_with_float_parameter( module, ttnn::leaky_relu, "slope", "The slope parameter for the Leaky ReLU function"); - // detail::bind_unary_with_float_parameter(module, ttnn::prelu, "weight", "The weight parameter for the PReLU - // function"); + // detail::bind_unary_operation_with_float_parameter(module, ttnn::prelu, "weight", "The weight parameter for the + // PReLU function"); // Other unaries (composite operations) detail::bind_softplus(module); diff --git a/ttnn/cpp/ttnn/operations/creation.hpp b/ttnn/cpp/ttnn/operations/creation.hpp index 10bf8a7a2d3..8e071b19fd5 100644 --- a/ttnn/cpp/ttnn/operations/creation.hpp +++ b/ttnn/cpp/ttnn/operations/creation.hpp @@ -9,6 +9,7 @@ #include "tt_eager/tensor/types.hpp" #include "tt_eager/tt_numpy/functions.hpp" #include "tt_metal/impl/dispatch/command_queue.hpp" +#include "ttnn/decorators.hpp" #include "ttnn/types.hpp" #include "ttnn/validation.hpp" @@ -19,37 +20,192 @@ namespace creation { template inline ttnn::Tensor full( const ttnn::Shape& shape, - const T value, - const DataType data_type, - const Layout layout, - Device& device, - const MemoryConfig& memory_config = ttnn::DRAM_MEMORY_CONFIG) { - return tt::numpy::full(shape.with_tile_padding().value(), value, data_type, layout, &device, memory_config); + const T fill_value, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device_arg = std::nullopt, + const std::optional& memory_config = std::nullopt) { + Device* device = device_arg.has_value() ? &(device_arg.value().get()) : nullptr; + return tt::numpy::full( + shape.value(), + fill_value, + dtype.value_or(ttnn::bfloat16), + layout.value_or(ttnn::ROW_MAJOR_LAYOUT), + device, + memory_config.value_or(ttnn::DRAM_MEMORY_CONFIG)); } inline ttnn::Tensor zeros( const ttnn::Shape& shape, - const DataType data_type, - const Layout layout, - Device& device, - const MemoryConfig& memory_config = ttnn::DRAM_MEMORY_CONFIG) { - return full(shape, 0.0f, data_type, layout, device, memory_config); + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full(shape, 0.0f, dtype, layout, device, memory_config); } inline ttnn::Tensor ones( const ttnn::Shape& shape, - const DataType data_type, - const Layout layout, - Device& device, - const MemoryConfig& memory_config = ttnn::DRAM_MEMORY_CONFIG) { - return full(shape, 1.0f, data_type, layout, device, memory_config); + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full(shape, 1.0f, dtype, layout, device, memory_config); } +inline ttnn::Tensor empty( + const ttnn::Shape& shape, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full(shape, 0.0f, dtype, layout, device, memory_config); +} + +template +inline ttnn::Tensor full_like( + const ttnn::Tensor& tensor, + const T fill_value, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + if (ttnn::is_tensor_on_device_or_multidevice(tensor)) { + return full( + tensor.get_shape(), + fill_value, + dtype.value_or(tensor.get_dtype()), + layout.value_or(tensor.get_layout()), + device.value_or(*tensor.device()), + memory_config.value_or(tensor.memory_config())); + } else { + return full( + tensor.get_shape(), + fill_value, + dtype.value_or(tensor.get_dtype()), + layout.value_or(tensor.get_layout()), + device, + memory_config); + } +} + +inline ttnn::Tensor zeros_like( + const ttnn::Tensor& tensor, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full_like(tensor, 0.0f, dtype, layout, device, memory_config); +} + +inline ttnn::Tensor ones_like( + const ttnn::Tensor& tensor, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full_like(tensor, 1.0f, dtype, layout, device, memory_config); +} + +inline ttnn::Tensor empty_like( + const ttnn::Tensor& tensor, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full_like(tensor, 0.0f, dtype, layout, device, memory_config); +} + +struct Full { + static ttnn::Tensor execute( + const ttnn::Shape& shape, + const float fill_value, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full(shape, fill_value, dtype, layout, device, memory_config); + } + + static ttnn::Tensor execute( + const ttnn::Shape& shape, + const int fill_value, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full(shape, fill_value, dtype, layout, device, memory_config); + } +}; + +struct FullLike { + static ttnn::Tensor execute( + const ttnn::Tensor& tensor, + const float fill_value, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full_like(tensor, fill_value, dtype, layout, device, memory_config); + } + + static ttnn::Tensor execute( + const ttnn::Tensor& tensor, + const int fill_value, + const std::optional& dtype = std::nullopt, + const std::optional& layout = std::nullopt, + const std::optional>& device = std::nullopt, + const std::optional& memory_config = std::nullopt) { + return full_like(tensor, fill_value, dtype, layout, device, memory_config); + } +}; + +struct Arange { + static ttnn::Tensor execute( + const int64_t stop, + const DataType dtype = ttnn::bfloat16, + const std::optional>& device = std::nullopt, + const MemoryConfig& memory_config = ttnn::DRAM_MEMORY_CONFIG) { + return Arange::execute(0, stop, 1, dtype, device, memory_config); + } + + static ttnn::Tensor execute( + const int64_t start, + const int64_t stop, + const int64_t step = 1, + const DataType dtype = ttnn::bfloat16, + const std::optional>& device_arg = std::nullopt, + const MemoryConfig& memory_config = ttnn::DRAM_MEMORY_CONFIG) { + Device* device = device_arg.has_value() ? &(device_arg.value().get()) : nullptr; + switch (dtype) { + case ttnn::bfloat16: + return tt::numpy::arange<::bfloat16>(start, stop, step, ttnn::ROW_MAJOR_LAYOUT, device, memory_config); + case ttnn::float32: + return tt::numpy::arange(start, stop, step, ttnn::ROW_MAJOR_LAYOUT, device, memory_config); + case ttnn::uint16: + return tt::numpy::arange(start, stop, step, ttnn::ROW_MAJOR_LAYOUT, device, memory_config); + case ttnn::uint32: + return tt::numpy::arange(start, stop, step, ttnn::ROW_MAJOR_LAYOUT, device, memory_config); + case ttnn::int32: + return tt::numpy::arange(start, stop, step, ttnn::ROW_MAJOR_LAYOUT, device, memory_config); + default: TT_THROW("Unsupported dtype"); + } + } +}; + } // namespace creation } // namespace operations -using operations::creation::full; -using operations::creation::ones; -using operations::creation::zeros; +constexpr auto full = ttnn::register_operation("ttnn::full"); +constexpr auto zeros = ttnn::register_operation("ttnn::zeros"); +constexpr auto ones = ttnn::register_operation("ttnn::ones"); +constexpr auto empty = ttnn::register_operation("ttnn::empty"); + +constexpr auto full_like = ttnn::register_operation("ttnn::full_like"); +constexpr auto zeros_like = ttnn::register_operation("ttnn::zeros_like"); +constexpr auto ones_like = ttnn::register_operation("ttnn::ones_like"); +constexpr auto empty_like = ttnn::register_operation("ttnn::empty_like"); + +constexpr auto arange = ttnn::register_operation("ttnn::arange"); } // namespace ttnn diff --git a/ttnn/cpp/ttnn/operations/pool.hpp b/ttnn/cpp/ttnn/operations/pool.hpp index b3d0d16eca3..905a3e2fbe1 100644 --- a/ttnn/cpp/ttnn/operations/pool.hpp +++ b/ttnn/cpp/ttnn/operations/pool.hpp @@ -44,8 +44,10 @@ struct GlobalAveragePool2D { return result; } }; -constexpr auto global_avg_pool2d = - ttnn::register_operation("ttnn::pool::global_avg_pool2d"); } // namespace pool } // namespace operations + +constexpr auto global_avg_pool2d = + ttnn::register_operation("ttnn::pool::global_avg_pool2d"); + } // namespace ttnn diff --git a/ttnn/ttnn/__init__.py b/ttnn/ttnn/__init__.py index 0d7f2db6a85..dd8d36e9d96 100644 --- a/ttnn/ttnn/__init__.py +++ b/ttnn/ttnn/__init__.py @@ -283,6 +283,7 @@ def manage_config(name, value): from ttnn.operations.creation import ( arange, empty, + empty_like, full, full_like, ones, diff --git a/ttnn/ttnn/operations/creation.py b/ttnn/ttnn/operations/creation.py index 38a7b1c7b5e..be6c7e174f0 100644 --- a/ttnn/ttnn/operations/creation.py +++ b/ttnn/ttnn/operations/creation.py @@ -8,63 +8,13 @@ import ttnn -def _zeros_like_validate_input_tensors(operation_name, input_tensor, *args, **kwargs): - ttnn.validate_input_tensor( - operation_name, - input_tensor, - ranks=(2, 3, 4), - dtypes=(ttnn.bfloat16, ttnn.bfloat8_b), - layouts=(ttnn.ROW_MAJOR_LAYOUT,), - can_be_on_device=True, - can_be_on_cpu=False, - ) - - def _golden_function(input_tensor: ttnn.Tensor, **_): import torch return torch.zeros_like(input_tensor) -@ttnn.register_operation( - name="ttnn.zeros_like", - validate_input_tensors=_zeros_like_validate_input_tensors, - golden_function=_golden_function, -) -def zeros_like( - input_tensor: ttnn.Tensor, - *, - memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG, -) -> ttnn.Tensor: - r""" - zeros_like(input_tensor: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor - - Returns a new tensor filled with zero by taking input tensor shape as reference. - - Args: - * :attr:`input_tensor`: the input tensor for reference shape - - Keyword Args: - * :attr:`memory_config`: the memory configuration for the output tensor - """ - - original_shape = input_tensor.shape - input_tensor = ttnn.unsqueeze_to_4D(input_tensor) - output_tensor = ttnn.experimental.tensor.zeros_like(input_tensor, output_mem_config=memory_config) - output_tensor = ttnn.reshape(output_tensor, original_shape) - return output_tensor - - -def _ones_like_validate_input_tensors(operation_name, input_tensor, *args, **kwargs): - ttnn.validate_input_tensor( - operation_name, - input_tensor, - ranks=(2, 3, 4), - dtypes=(ttnn.bfloat16, ttnn.bfloat8_b), - layouts=(ttnn.ROW_MAJOR_LAYOUT,), - can_be_on_device=True, - can_be_on_cpu=False, - ) +zeros_like = ttnn.register_operation(golden_function=_golden_function)(ttnn._ttnn.operations.creation.zeros_like) def _golden_function(input_tensor: ttnn.Tensor, **_): @@ -73,45 +23,7 @@ def _golden_function(input_tensor: ttnn.Tensor, **_): return torch.ones_like(input_tensor) -@ttnn.register_operation( - name="ttnn.ones_like", - validate_input_tensors=_ones_like_validate_input_tensors, - golden_function=_golden_function, -) -def ones_like( - input_tensor: ttnn.Tensor, - *, - memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG, -) -> ttnn.Tensor: - r""" - ones_like(input_tensor: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor - - Returns a new tensor filled with one by taking input tensor shape as reference. - - Args: - * :attr:`input_tensor`: the input tensor for reference shape - - Keyword Args: - * :attr:`memory_config`: the memory configuration for the output tensor - """ - - original_shape = input_tensor.shape - input_tensor = ttnn.unsqueeze_to_4D(input_tensor) - output_tensor = ttnn.experimental.tensor.ones_like(input_tensor, output_mem_config=memory_config) - output_tensor = ttnn.reshape(output_tensor, original_shape) - return output_tensor - - -def _full_like_validate_input_tensors(operation_name, input_tensor, *args, **kwargs): - ttnn.validate_input_tensor( - operation_name, - input_tensor, - ranks=(2, 3, 4), - dtypes=(ttnn.bfloat16, ttnn.bfloat8_b), - layouts=(ttnn.ROW_MAJOR_LAYOUT,), - can_be_on_device=True, - can_be_on_cpu=False, - ) +ones_like = ttnn.register_operation(golden_function=_golden_function)(ttnn._ttnn.operations.creation.ones_like) def _golden_function(input_tensor: ttnn.Tensor, *, fill_value: float, **_): @@ -120,36 +32,16 @@ def _golden_function(input_tensor: ttnn.Tensor, *, fill_value: float, **_): return torch.full_like(input_tensor, fill_value) -@ttnn.register_operation( - name="ttnn.full_like", - validate_input_tensors=_full_like_validate_input_tensors, - golden_function=_golden_function, -) -def full_like( - input_tensor: ttnn.Tensor, - *, - fill_value: float, - memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG, -) -> ttnn.Tensor: - r""" +full_like = ttnn.register_operation(golden_function=_golden_function)(ttnn._ttnn.operations.creation.full_like) - full_like(input_tensor: ttnn.Tensor, *, fill_value: float, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor - Returns a new tensor filled with scalar value by taking input tensor shape as reference. +def _golden_function(input_tensor: ttnn.Tensor, *, fill_value: float, **_): + import torch - Args: - * :attr:`input_tensor`: the input tensor for reference shape + return torch.empty_like(input_tensor, fill_value) - Keyword Args: - * :attr:`fill_value`: the value to be filled - * :attr:`memory_config`: the memory configuration for the output tensor - """ - original_shape = input_tensor.shape - input_tensor = ttnn.unsqueeze_to_4D(input_tensor) - output_tensor = ttnn.experimental.tensor.full_like(input_tensor, fill_value, output_mem_config=memory_config) - output_tensor = ttnn.reshape(output_tensor, original_shape) - return output_tensor +empty_like = ttnn.register_operation(golden_function=_golden_function)(ttnn._ttnn.operations.creation.empty_like) def _golden_function(input_shape: ttnn.Shape, **_): @@ -158,42 +50,7 @@ def _golden_function(input_shape: ttnn.Shape, **_): return torch.zeros(input_shape) -def _zeros_validate_input_tensors(operation_name, input_shape, *args, **kwargs): - ... - - -@ttnn.register_operation( - name="ttnn.zeros", - validate_input_tensors=_zeros_validate_input_tensors, - golden_function=_golden_function, -) -def zeros( - input_shape: ttnn.Shape, - *, - device: ttnn.Device, - dtype: Union[ttnn.DataType, str] = ttnn.bfloat16, - layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, - memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG, -) -> ttnn.Tensor: - r""" - - zeros(input_shape: ttnn.Shape, *, device: ttnn.Device, dtype: Union[ttnn.DataType, str] = ttnn.bfloat16, layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor - - Returns a new tensor filled with zeros by taking input tensor shape as reference. - - Args: - * :attr:`input_shape`: the input shape for reference - - Keyword Args: - * :attr:`device`: the device for the output tensor - * :attr:`dtype`: the data type for the output tensor - * :attr:`layout`: the layout for the output tensor - * :attr:`memory_config`: the memory configuration for the output tensor - """ - output_tensor = ttnn.experimental.tensor.zeros( - input_shape, data_type=dtype, layout=layout, device=device, output_mem_config=memory_config - ) - return output_tensor +zeros = ttnn.register_operation(golden_function=_golden_function)(ttnn._ttnn.operations.creation.zeros) def _golden_function(input_shape: ttnn.Shape, **_): @@ -202,43 +59,7 @@ def _golden_function(input_shape: ttnn.Shape, **_): return torch.ones(input_shape) -def _ones_validate_input_tensors(operation_name, input_shape, *args, **kwargs): - ... - - -@ttnn.register_operation( - name="ttnn.ones", - validate_input_tensors=_ones_validate_input_tensors, - golden_function=_golden_function, -) -def ones( - input_shape: ttnn.Shape, - *, - device: ttnn.Device, - dtype: Union[ttnn.DataType, str] = ttnn.bfloat16, - layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, - memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG, -) -> ttnn.Tensor: - r""" - - ones(input_shape: ttnn.Shape, *, device: ttnn.Device, dtype: Union[ttnn.DataType, str] = ttnn.bfloat16, layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor - - Returns a new tensor filled with ones by taking input tensor shape as reference. - - Args: - * :attr:`input_shape`: the input shape for reference - - Keyword Args: - * :attr:`device`: the device for the output tensor - * :attr:`dtype`: the data type for the output tensor - * :attr:`layout`: the layout for the output tensor - * :attr:`memory_config`: the memory configuration for the output tensor - """ - - output_tensor = ttnn.experimental.tensor.ones( - input_shape, data_type=dtype, layout=layout, device=device, output_mem_config=memory_config - ) - return output_tensor +ones = ttnn.register_operation(golden_function=_golden_function)(ttnn._ttnn.operations.creation.ones) def _golden_function_full(input_shape: ttnn.Shape, fill_value: float, **_): @@ -247,55 +68,16 @@ def _golden_function_full(input_shape: ttnn.Shape, fill_value: float, **_): return torch.full(input_shape, fill_value=fill_value) -def _full_validate_input_tensors(operation_name, input_shape, *args, **kwargs): - ... +full = ttnn.register_operation(golden_function=_golden_function)(ttnn._ttnn.operations.creation.full) -@ttnn.register_operation( - name="ttnn.full", - validate_input_tensors=_full_validate_input_tensors, - golden_function=_golden_function, -) -def full( - input_shape: ttnn.Shape, - *, - device: ttnn.Device, - dtype: Union[ttnn.DataType, str] = ttnn.bfloat16, - layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, - fill_value: float, - memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG, -) -> ttnn.Tensor: - r""" - - full(input_shape: ttnn.Shape, *, device: ttnn.Device, dtype: Union[ttnn.DataType, str] = ttnn.bfloat16, layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, fill_value: float, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor - - Returns a new tensor filled with fill_value by taking input tensor shape as reference. - - Args: - * :attr:`input_shape`: the input shape for reference - - Keyword Args: - * :attr:`device`: the device for the output tensor - * :attr:`dtype`: the data type for the output tensor - * :attr:`layout`: the layout for the output tensor - * :attr:`fill_value`: the value to be filled - * :attr:`memory_config`: the memory configuration for the output tensor - - """ +def _golden_function(input_shape: ttnn.Shape, **_): + import torch - output_tensor = ttnn.experimental.tensor.full( - input_shape, - fill_value=fill_value, - device=device, - data_type=dtype, - layout=layout, - output_mem_config=memory_config, - ) - return output_tensor + return torch.empty(input_shape) -def _is_int(value): - return isinstance(value, (int)) +empty = ttnn.register_operation(golden_function=_golden_function)(ttnn._ttnn.operations.creation.empty) def _golden_function(start: int, end: int, step: int, **_): @@ -331,47 +113,10 @@ def arange( * :attr:`end` * :attr:`step` """ - if not _is_int(start) or not _is_int(end) or not _is_int(step): - raise TypeError("Expected three arguments to be a int") output_tensor = ttnn.experimental.tensor.arange(start, end, step, device, output_mem_config=memory_config) return output_tensor -def _golden_function(input_shape: ttnn.Shape, **_): - import torch - - return torch.empty(input_shape) - - -def _empty_validate_input_tensors(operation_name, input_shape, *args, **kwargs): - ... - - -@ttnn.register_operation( - name="ttnn.empty", - validate_input_tensors=_empty_validate_input_tensors, - golden_function=_golden_function, -) -def empty( - input_shape: ttnn.Shape, - device: ttnn.Device, - memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG, -) -> ttnn.Tensor: - r""" - - empty(input_shape: ttnn.Shape, device: ttnn.Device, memory_config: ttnn.MemoryConfig = ttnn.DRAM_MEMORY_CONFIG) -> ttnn.Tensor - - Returns a new empty tensor by taking input shape as reference. - - Args: - * :attr:`input_shape`: the input shape for reference - """ - - output_tensor = ttnn.experimental.tensor.empty(input_shape, device=device, output_mem_config=memory_config) - - return output_tensor - - __all__ = []