search

tenstorrent / tt-metal

:metal: TT-NN operator library, and TT-Metalium low level kernel programming model.
https://docs.tenstorrent.com/ttnn/latest/index.html
Apache License 2.0
484 stars 79 forks source link

Need ttnn.permute support for 5D #15165

Open keerthana-r-mcw opened 1 week ago

keerthana-r-mcw commented 1 week ago

Describe the bug Need support for 5D ttnn.permute for PETR Transformer module.

Getting the following error when ttnn.permute applied for 5D tensor, RuntimeError: TT_FATAL @ ../ttnn/cpp/ttnn/operations/data_movement/permute/device/permute_device_operation.cpp:22: attributes.dims.back() == tensor_args.input_tensor.get_logical_shape().rank() - 1

Please complete the following environment information:

jerrysky3 commented 1 week ago

I ran into a similar issue on ttnn.permute with > 4D input, but with a different error

To reproduce

import torch
import ttnn

def main(device):
    torch_tensor = torch.rand((1, 3, 16, 16, 16, 16), dtype=torch.bfloat16)
    input_tensor = ttnn.from_torch(torch_tensor, layout=ttnn.TILE_LAYOUT, device=device)
    output_tensor = ttnn.permute(input_tensor, (0, 2, 4, 3, 5, 1))
    print(output_tensor)

if __name__ == "__main__":
    try:
        device = ttnn.open_device(device_id=0)
        main(device)
    finally:
        ttnn.close_device(device)

Error log

Traceback (most recent call last):
  File "tests/ttnn/unit_tests/exp_permute.py", line 13, in <module>
    main(device)
  File "tests/ttnn/unit_tests/exp_permute.py", line 7, in main
    output_tensor = ttnn.permute(input_tensor, (0, 2, 4, 3, 5, 1))
  File "/home/jerry/tt-metal/ttnn/ttnn/decorators.py", line 329, in __call__
    return self.function(*function_args, **function_kwargs)
RuntimeError: TT_FATAL @ ../ttnn/cpp/ttnn/operations/core/work_split/work_split_tilize.hpp:181: logical_shape.rank() >= 2 && logical_shape.rank() <= 4
info:
Only 2D, 3D, and 4D tensors are supported. Shape: Error
backtrace:
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x6dcf0b) [0x7f8341475f0b]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x6dcd16) [0x7f8341475d16]
 --- ttnn::operations::data_movement::detail::untilize_with_unpadding_multi_core_interleaved(tt::tt_metal::Tensor const&, tt::tt_metal::Tensor&, bool, bool)
 --- ttnn::operations::data_movement::detail::untilize_with_unpadding_multi_core(tt::tt_metal::Tensor const&, tt::tt_metal::Tensor&, bool, bool)
 --- ttnn::operations::data_movement::UntilizeWithUnpadding::create_program(std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> const&, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>&) const
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x720ed5) [0x7f83414b9ed5]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x42b36d) [0x7f83411c436d]
 --- void ttnn::device_operation::detail::launch_on_worker_thread<tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>, unsigned char, long, tt::tt_metal::operation::DeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>, tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>::tensor_args_t, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>, tt::tt_metal::v0::Device*>(unsigned char, long, tt::tt_metal::operation::DeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>> const&, tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>::tensor_args_t const&, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>&, tt::tt_metal::v0::Device*&)
 --- tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>::tensor_return_value_t ttnn::device_operation::detail::launch_on_single_device<tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>>(unsigned char, tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>::operation_attributes_t const&, tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>::tensor_args_t const&)
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x41defa) [0x7f83411b6efa]
 --- tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>::tensor_return_value_t ttnn::device_operation::detail::invoke<tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>>(unsigned char, tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>::operation_attributes_t const&, tt::tt_metal::operation::OldInfraDeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>::tensor_args_t const&)
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x41cb4b) [0x7f83411b5b4b]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x418c88) [0x7f83411b1c88]
 --- std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> tt::tt_metal::operation::run<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>(tt::tt_metal::operation::DeviceOperation<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>&&, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor const>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor const>>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>> const&, unsigned char)
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x720afd) [0x7f83414b9afd]
 --- ttnn::operations::data_movement::ExecuteUntilizeWithUnpadding::invoke(unsigned char, tt::tt_metal::Tensor const&, tt::tt_metal::LegacyShape const&, std::__1::optional<tt::tt_metal::MemoryConfig> const&, bool, bool)
 --- ttnn::operations::data_movement::ExecuteUntilizeWithUnpadding::invoke(tt::tt_metal::Tensor const&, tt::tt_metal::LegacyShape const&, std::__1::optional<tt::tt_metal::MemoryConfig> const&, bool, bool)
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x5a9d0c) [0x7f8341342d0c]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x5aaba3) [0x7f8341343ba3]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x5ab1fb) [0x7f83413441fb]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x5a8e49) [0x7f8341341e49]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x5a7bb6) [0x7f8341340bb6]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x59ed68) [0x7f8341337d68]
 --- tt::tt_metal::Tensor ttnn::operations::core::detail::to_layout_impl<tt::tt_metal::v0::Device>(tt::tt_metal::Tensor const&, tt::tt_metal::Layout, std::__1::optional<tt::tt_metal::DataType> const&, std::__1::optional<tt::tt_metal::MemoryConfig> const&, tt::tt_metal::v0::Device*)
 --- ttnn::operations::core::ToLayout::invoke(tt::tt_metal::Tensor const&, tt::tt_metal::Layout, std::__1::optional<tt::tt_metal::DataType> const&, std::__1::optional<tt::tt_metal::MemoryConfig> const&, tt::tt_metal::v0::Device*)
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x64a1ac) [0x7f83413e31ac]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x649f45) [0x7f83413e2f45]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x64ae6b) [0x7f83413e3e6b]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x64b45b) [0x7f83413e445b]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x6497b9) [0x7f83413e27b9]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x648635) [0x7f83413e1635]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x6473ab) [0x7f83413e03ab]
 --- ttnn::operations::data_movement::detail::permute_impl(tt::tt_metal::Tensor const&, tt::tt_metal::SmallVector<unsigned int, 8ul> const&, tt::tt_metal::MemoryConfig const&)
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x64710b) [0x7f83413e010b]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x436935) [0x7f83411cf935]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x437a4f) [0x7f83411d0a4f]
 --- void tt::tt_metal::operation::launch_op<std::__1::function<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> (std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor const>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor const>>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>> const&)>, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>>(std::__1::function<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> (std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor const>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor const>>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>> const&)>&&, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor const>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor const>>>, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>>, bool)
 --- tt::tt_metal::operation::launch_with_autoformat(std::__1::function<std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> (std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor const>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor const>>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>> const&)>&&, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>, std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor const>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor const>>>, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>>)
 --- ttnn::operations::data_movement::detail::permute_launch(tt::tt_metal::Tensor const&, tt::stl::Span<long const, 18446744073709551615ul>, tt::tt_metal::MemoryConfig const&)
 --- ttnn::operations::data_movement::ExecutePermute::invoke(unsigned char, tt::tt_metal::Tensor const&, tt::stl::Span<long const, 18446744073709551615ul>, std::__1::optional<tt::tt_metal::MemoryConfig> const&, bool)
ayerofieiev-tt commented 1 week ago

@ntarafdar , can you please review with the team and share the timeline?

punithsekar commented 5 days ago

To add on, Even used in SinePositionalEncoding3D class of petr model.