[Bug Report] Untilizing crashes or hangs with tensor reshaped to non tile aligned shape about tt-metal

[ayerofieiev-tt]

Yes. We will add a proper check and throw an error. Thank you!

[ayerofieiev-tt]

import torch
import ttnn

device = ttnn.open_device(0)

a = torch.rand(1, 1, 2, 2048)

ttnn.enable_program_cache(device)

x = ttnn.from_torch(a, device=device, layout=ttnn.Layout.TILE, dtype=ttnn.DataType.BFLOAT16)

Now x shape is [1, 1, 32, 2048] - that 32 is a padded dimension.
I'd expect reshape into [1, 2, 8, 256] to throw as illegal in this case.

@arakhmati , wdyt, should we add more checks to reshape?

[marty1885]

@ayerofieiev-tt is there an efficient way to work around this in TTNN? This reshape is one of the needed operation during my bringup of Gemma 2B on GGML. (I will need to eliminate this reshape for performance in the future. But now I'm focused on getting things running).

Thanks to your comment I figured out I can do this. But it's really slow,

if(dst->ne[0] % tt::constants::TILE_WIDTH != 0 || dst->ne[1] % tt::constants::TILE_HEIGHT != 0) {
    // This path is SLOW.
    tt::tt_metal::Tensor row_major_tensor = tt::tt_metal::untilize(t);
    tt::tt_metal::Tensor reshaped = row_major_tensor.reshape(target_shape);
    tt::tt_metal::Tensor ret = tt::tt_metal::tilize_with_zero_padding(reshaped);
    *dst_meta = {
        .tensor = std::make_shared<tt::tt_metal::Tensor>(std::move(ret)),
        .ggtype = dst->type,
        .bufctx = ((TensorWithMetadata*)dst->src[0]->extra)->bufctx
    };
    return;
}

IMO anything disallowed should throw. And footguns should be very hard to call and always has to be done explicitly.

[arakhmati]

@marty1885 this reshape is not currently supported in tile layout. You would need to write a custom operation for this case

[marty1885]

I see. If that's the case can we make TTNN throw when a invalid reshape is issued?