From d7832e97f3b7e30fb855ec88e9b46657aff76675 Mon Sep 17 00:00:00 2001 From: tarafdarTT Date: Thu, 4 Jul 2024 18:20:24 +0000 Subject: [PATCH] #9745: move unpad to slice ttnn cpp references --- .../op_library/backward/backward_ops.cpp | 14 +- .../tt_dnn/op_library/complex/complex_ops.cpp | 6 +- .../unpad/multi_core/unpad_op_multi_core.cpp | 6 +- ..._unary_unpad_dims_interleaved_start_id.cpp | 0 ...ary_unpad_dims_rm_interleaved_start_id.cpp | 0 ...nary_stick_layout_interleaved_start_id.cpp | 0 .../operations/data_movement/slice/slice.hpp | 127 ++++++++++++++++++ .../device/binary_backward_op.cpp | 6 +- 8 files changed, 143 insertions(+), 16 deletions(-) rename {tt_eager/tt_dnn/op_library/unpad => ttnn/cpp/ttnn/operations/data_movement/slice}/kernels/dataflow/reader_unary_unpad_dims_interleaved_start_id.cpp (100%) rename {tt_eager/tt_dnn/op_library/unpad => ttnn/cpp/ttnn/operations/data_movement/slice}/kernels/dataflow/reader_unary_unpad_dims_rm_interleaved_start_id.cpp (100%) rename {tt_eager/tt_dnn/op_library/unpad => ttnn/cpp/ttnn/operations/data_movement/slice}/kernels/dataflow/writer_unary_stick_layout_interleaved_start_id.cpp (100%) create mode 100644 ttnn/cpp/ttnn/operations/data_movement/slice/slice.hpp diff --git a/tt_eager/tt_dnn/op_library/backward/backward_ops.cpp b/tt_eager/tt_dnn/op_library/backward/backward_ops.cpp index 9950ce27275..b9a0e127671 100644 --- a/tt_eager/tt_dnn/op_library/backward/backward_ops.cpp +++ b/tt_eager/tt_dnn/op_library/backward/backward_ops.cpp @@ -12,7 +12,7 @@ #include "tt_dnn/op_library/permute/permute_op.hpp" #include "tt_dnn/op_library/reduce/reduce_op.hpp" #include "tt_dnn/op_library/reshape/reshape_op.hpp" -#include "tt_dnn/op_library/unpad/unpad_op.hpp" +#include "ttnn/operations/data_movement/slice/slice.hpp" #include "tt_eager/tensor/tensor_utils.hpp" #include "ttnn/operations/data_movement/pad/pad.hpp" #include "tt_numpy/functions.hpp" @@ -1532,9 +1532,9 @@ std::vector _prod_bw( const Shape start_index = {0, 0, 0, 0}; const Shape end_index = { grad.get_legacy_shape()[0] - 1, 0, grad.get_legacy_shape()[1] - 1, grad.get_legacy_shape()[2] - 1}; - Tensor new_unpad_tensor = unpad(required, start_index, end_index); + Tensor new_slice_tensor = ttnn::slice(required, start_index, end_index, std::nullopt); after_permute_dims = {0, 2, 3, 1}; - updated_grad = permute(new_unpad_tensor, after_permute_dims, output_mem_config); + updated_grad = permute(new_slice_tensor, after_permute_dims, output_mem_config); Tensor pad_updated_grad = updated_grad.pad_to_tile(1.0f); Tensor pad_prod_result = prod_result.pad_to_tile(1.0f); pad_updated_grad = pad_updated_grad.to(Layout::TILE); @@ -1549,8 +1549,8 @@ std::vector _prod_bw( const Shape start_index = {0, 0, 0, 0}; const Shape end_index = { grad.get_legacy_shape()[0] - 1, 0, grad.get_legacy_shape()[1] - 1, grad.get_legacy_shape()[3] - 1}; - Tensor new_unpad_tensor = unpad(required, start_index, end_index); - updated_grad = permute(new_unpad_tensor, after_permute_dims, output_mem_config); + Tensor new_slice_tensor = ttnn::slice(required, start_index, end_index, std::nullopt); + updated_grad = permute(new_slice_tensor, after_permute_dims, output_mem_config); if(updated_grad.get_layout()==Layout::ROW_MAJOR){ updated_grad = tt::tt_metal::change_layout_to_tile(updated_grad, output_mem_config); } @@ -1599,7 +1599,7 @@ std::vector _prod_bw( input.get_legacy_shape()[1] - 1, input.get_legacy_shape()[2] - 1, input.get_legacy_shape()[3] - 1}; - grad_result = unpad(result, start_index, end_index); + grad_result = ttnn::slice(result, start_index, end_index, std::nullopt); } grad_tensor.emplace_back(grad_result); return grad_tensor; @@ -1633,7 +1633,7 @@ std::vector _prod_bw( input.get_legacy_shape()[1] - 1, input.get_legacy_shape()[2] - 1, input.get_legacy_shape()[3] - 1}; - grad_result = unpad(result, start_index, end_index); + grad_result = ttnn::slice(result, start_index, end_index, std::nullopt); } grad_tensor.emplace_back(grad_result); return grad_tensor; diff --git a/tt_eager/tt_dnn/op_library/complex/complex_ops.cpp b/tt_eager/tt_dnn/op_library/complex/complex_ops.cpp index c9282eb2b49..81864d3c78c 100644 --- a/tt_eager/tt_dnn/op_library/complex/complex_ops.cpp +++ b/tt_eager/tt_dnn/op_library/complex/complex_ops.cpp @@ -6,7 +6,7 @@ #include "tt_dnn/op_library/concat/concat_op.hpp" #include "tt_dnn/op_library/bmm/bmm_op.hpp" #include "tt_dnn/op_library/reshape/reshape_op.hpp" -#include "tt_dnn/op_library/unpad/unpad_op.hpp" +#include "ttnn/operations/data_movement/slice/slice.hpp" #include "tt_numpy/functions.hpp" #include "tt_eager/tensor/tensor_utils.hpp" @@ -33,7 +33,7 @@ Tensor get_real(const Tensor& input, const MemoryConfig& output_mem_config) { Shape t_Shape = input.get_legacy_shape(); Shape start = {0, 0, 0, 0} ; Shape end = {t_Shape[0] - 1,t_Shape[1] - 1 ,t_Shape[2] - 1, (t_Shape[3] / 2) - 1}; - Tensor r_tensor = unpad(input, start, end, output_mem_config); + Tensor r_tensor = ttnn::slice(input, start, end, output_mem_config); return r_tensor; } @@ -41,7 +41,7 @@ Tensor get_imag(const Tensor& input, const MemoryConfig& output_mem_config) { Shape t_Shape = input.get_legacy_shape(); Shape start = {0, 0, 0, (t_Shape[3] / 2)}; Shape end = {t_Shape[0] - 1,t_Shape[1] - 1 ,t_Shape[2] - 1, (t_Shape[3] - 1)}; - Tensor i_tensor = unpad(input, start, end, output_mem_config); + Tensor i_tensor = ttnn::slice(input, start, end, output_mem_config); return i_tensor; } diff --git a/tt_eager/tt_dnn/op_library/unpad/multi_core/unpad_op_multi_core.cpp b/tt_eager/tt_dnn/op_library/unpad/multi_core/unpad_op_multi_core.cpp index 15ffaf9fcf6..a5427189e02 100644 --- a/tt_eager/tt_dnn/op_library/unpad/multi_core/unpad_op_multi_core.cpp +++ b/tt_eager/tt_dnn/op_library/unpad/multi_core/unpad_op_multi_core.cpp @@ -161,13 +161,13 @@ operation::ProgramWithCallbacks unpad_rm_multi_core( tt_metal::KernelHandle unary_reader_kernel_id = tt_metal::CreateKernel( program, - "tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/reader_unary_unpad_dims_rm_interleaved_start_id.cpp", + "ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/reader_unary_unpad_dims_rm_interleaved_start_id.cpp", total_cores, tt_metal::ReaderDataMovementConfig(reader_compile_time_args_vec)); tt_metal::KernelHandle unary_writer_kernel_id = tt_metal::CreateKernel( program, - "tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/writer_unary_stick_layout_interleaved_start_id.cpp", + "ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/writer_unary_stick_layout_interleaved_start_id.cpp", total_cores, tt_metal::WriterDataMovementConfig(writer_compile_time_args_vec)); @@ -435,7 +435,7 @@ operation::ProgramWithCallbacks unpad_tile_multi_core( // Tilized reader tt_metal::KernelHandle unary_reader_kernel_id = tt_metal::CreateKernel( program, - "tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/reader_unary_unpad_dims_interleaved_start_id.cpp", + "ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/reader_unary_unpad_dims_interleaved_start_id.cpp", total_cores, tt_metal::ReaderDataMovementConfig(reader_compile_time_args)); diff --git a/tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/reader_unary_unpad_dims_interleaved_start_id.cpp b/ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/reader_unary_unpad_dims_interleaved_start_id.cpp similarity index 100% rename from tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/reader_unary_unpad_dims_interleaved_start_id.cpp rename to ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/reader_unary_unpad_dims_interleaved_start_id.cpp diff --git a/tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/reader_unary_unpad_dims_rm_interleaved_start_id.cpp b/ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/reader_unary_unpad_dims_rm_interleaved_start_id.cpp similarity index 100% rename from tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/reader_unary_unpad_dims_rm_interleaved_start_id.cpp rename to ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/reader_unary_unpad_dims_rm_interleaved_start_id.cpp diff --git a/tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/writer_unary_stick_layout_interleaved_start_id.cpp b/ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/writer_unary_stick_layout_interleaved_start_id.cpp similarity index 100% rename from tt_eager/tt_dnn/op_library/unpad/kernels/dataflow/writer_unary_stick_layout_interleaved_start_id.cpp rename to ttnn/cpp/ttnn/operations/data_movement/slice/kernels/dataflow/writer_unary_stick_layout_interleaved_start_id.cpp diff --git a/ttnn/cpp/ttnn/operations/data_movement/slice/slice.hpp b/ttnn/cpp/ttnn/operations/data_movement/slice/slice.hpp new file mode 100644 index 00000000000..acdcef2e128 --- /dev/null +++ b/ttnn/cpp/ttnn/operations/data_movement/slice/slice.hpp @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "tt_eager/tensor/types.hpp" +#include "ttnn/cpp/ttnn/operations/core.hpp" +#include "tt_eager/tt_dnn/op_library/unpad/unpad_op.hpp" + +#include + + +namespace ttnn { +namespace operations { +namespace data_movement { + + + +struct ExecuteSlice { + static inline const std::array input_tensor_schemas() { + return {ttnn::TensorSchema{ + 2, // min rank + 4, // max rank + {ttnn::bfloat16, ttnn::bfloat8_b, ttnn::uint16, ttnn::int32, ttnn::uint32}, + {ttnn::TILE_LAYOUT}, + true, // can_be_on_device + false, // can_be_on_cpu + false, // can_be_scalar + false // is_optional} + }}; + } + + template + static auto input_tensors_to_validate(const ttnn::Tensor& input_tensor, Args&&... args) { + return std::make_tuple(input_tensor); + } + + static ttnn::Tensor execute_on_worker_thread( + uint8_t queue_id, + const ttnn::Tensor& input_tensor, + tt::tt_metal::Shape output_tensor_start, + tt::tt_metal::Shape output_tensor_end, + const std::optional& memory_config_arg) { + + auto memory_config = memory_config_arg.value_or(input_tensor.memory_config()); + + auto input_tensor_shape = input_tensor.get_legacy_shape(); + std::vector output_tensor_shape = { + output_tensor_end[0] - output_tensor_start[0] + 1, + output_tensor_end[1] - output_tensor_start[1] + 1, + output_tensor_end[2] - output_tensor_start[2] + 1, + output_tensor_end[3] - output_tensor_start[3] + 1, + }; + auto output_tensor = operation::run( + tt::tt_metal::Unpad{ + .output_tensor_start=output_tensor_start, + .output_tensor_end=output_tensor_end, + .output_mem_config=memory_config, + .output_shape=output_tensor_shape, + .input_shape=input_tensor_shape + }, + {input_tensor}).front(); + + return output_tensor; + } + + + static ttnn::Tensor execute_on_worker_thread( + const ttnn::Tensor& input_tensor, + tt::tt_metal::Shape output_tensor_start, + tt::tt_metal::Shape output_tensor_end, + const std::optional& memory_config_arg + ) { + + + return execute_on_worker_thread( + 0, + input_tensor, + output_tensor_start, + output_tensor_end, + memory_config_arg + ); + + } + + + static ttnn::Tensor execute_on_worker_thread( + uint8_t queue_id, + const ttnn::Tensor& input_tensor, + std::vector output_tensor_start, + std::vector output_tensor_end, + const std::optional& memory_config_arg) { + + return execute_on_worker_thread(queue_id, + input_tensor, + tt::tt_metal::Shape(output_tensor_start), + tt::tt_metal::Shape(output_tensor_end), + memory_config_arg + ); + + } + + static ttnn::Tensor execute_on_worker_thread( + const ttnn::Tensor& input_tensor, + std::vector output_tensor_start, + std::vector output_tensor_end, + const std::optional& memory_config_arg) { + + return execute_on_worker_thread( + 0, + input_tensor, + output_tensor_start, + output_tensor_end, + memory_config_arg + ); + } + + +}; + +} // namespace data_movement +} // namespace operations + +constexpr auto slice = ttnn::register_operation("ttnn::slice"); + +} // namespace ttnn diff --git a/ttnn/cpp/ttnn/operations/eltwise/binary_backward/device/binary_backward_op.cpp b/ttnn/cpp/ttnn/operations/eltwise/binary_backward/device/binary_backward_op.cpp index b79e40ef878..f78a540e955 100644 --- a/ttnn/cpp/ttnn/operations/eltwise/binary_backward/device/binary_backward_op.cpp +++ b/ttnn/cpp/ttnn/operations/eltwise/binary_backward/device/binary_backward_op.cpp @@ -9,7 +9,7 @@ #include "tt_eager/tt_dnn/op_library/bcast/bcast_op.hpp" #include "tt_eager/tt_dnn/op_library/composite/composite_ops.hpp" #include "tt_eager/tt_dnn/op_library/eltwise_unary/eltwise_unary_op.hpp" -#include "tt_eager/tt_dnn/op_library/unpad/unpad_op.hpp" +#include "ttnn/operations/data_movement/slice/slice.hpp" #include "tt_metal/common/constants.hpp" #include "tt_metal/host_api.hpp" #include "tt_metal/tools/profiler/op_profiler.hpp" @@ -374,7 +374,7 @@ std::vector _concat_bw( input.get_legacy_shape()[2] - 1, input.get_legacy_shape()[3] - 1}; - Tensor grad_a = unpad(grad, start_index, end_index); + Tensor grad_a = ttnn::slice(grad, start_index, end_index, std::nullopt); grad_tensor.emplace_back(grad_a); tt::tt_metal::Shape start_index_2 = {0, 0, 0, 0}; @@ -393,7 +393,7 @@ std::vector _concat_bw( grad.get_legacy_shape()[1] - 1, grad.get_legacy_shape()[2] - 1, grad.get_legacy_shape()[3] - 1}; - Tensor grad_b = unpad(grad, start_index_2, end_index_2); + Tensor grad_b = ttnn::slice(grad, start_index_2, end_index_2, std::nullopt); grad_tensor.emplace_back(grad_b); return grad_tensor;