TTNN pad does not support last dimension

[esmalTT]

Summary

Pad on device is required to improve end-to-end performance of UNet Shallow. The sharded input tensor needs to be padded from 4 -> 16 channels.

Running a pad on a device tensor of shape {1, 1, 2 * 1056 * 160, 4} to {1, 1, 2 * 1056 * 160, 16} throws the following error:

sharded pad does not support pad on last dim currently as that will cause perf degradation

[jaykru-tt]

@esmalTT Can you give an example test? I have been working on this for a bit and I'm actually now confused how you obtained a sharded tensor of that shape in the first place :)

[esmalTT]

@esmalTT Can you give an example test? I have been working on this for a bit and I'm actually now confused how you obtained a sharded tensor of that shape in the first place :)

I’m away from my computer until next week (so sadly I can’t run this) but this is what I need for UNet

import pytest
import torch
import ttnn

def test_unet_pad(device, use_program_cache):
    x = ttnn.from_torch(torch.rand([1, 1, 337920, 4]), dtype=ttnn.bfloat16)

    sharded_memory_config = ttnn.create_sharded_memory_config(
        [1, 1, 337920, 4], ttnn.CoreGrid(x=8, y=8), ttnn.ShardStrategy.HEIGHT
    )

    x = ttnn.to_device(x, device, sharded_memory_config)

    x = ttnn.pad(x, ((0,0),(0,0),(0,0),(0,12))) # pad up to 16

[jaykru-tt]

This is perfect, thanks @esmalTT :)

[ntarafdar]

hey @esmalTT , @bbradelTT ran into this problem too. @jaykru-tt is actively looking at this and should have something for next week.
Do you have a workaround for this? if so could you comment on it so @bbradelTT can see it too.

[esmalTT]

hey @esmalTT , @bbradelTT ran into this problem too. @jaykru-tt is actively looking at this and should have something for next week. Do you have a workaround for this? if so could you comment on it so @bbradelTT can see it too.

@ntarafdar @bbradelTT I was never able to come up with a workaround, we ended up avoiding the need for this entirely by making moving to channels-first instead of channels-last.

[bbradelTT]

Thanks @esmalTT that is a workaround. We can just transpose back afterwards:

def test_pad_transpose(device):
    input_shape = (1, 1, 12, 8)
    a = torch.ones(input_shape) 
    b = ttnn.from_torch(a, dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=device)
    b2 = ttnn.to_layout(b, layout=ttnn.ROW_MAJOR_LAYOUT)
    b3 = ttnn.transpose(b2, 1, 3)
    c = ttnn.pad(b3, [1, 32, 32, 32], [0, 0, 0, 0], 3)
    c2 = ttnn.transpose(c, 1, 3)
    c3 = ttnn.slice(c2, [0, 0, 0, 0], [1, 1, 32, 32])
    c4 = ttnn.to_layout(c3, layout=ttnn.TILE_LAYOUT)
    print(f'b2 {b2}\nb3 {b3}\n c {c}\nc2 {c2}\nc3 {c3}\nc4 {c4}')

outputs

...
c4 ttnn.Tensor([[[[ 1.00000,  1.00000,  ...,  3.00391,  3.00391],
               [ 1.00000,  1.00000,  ...,  3.00391,  3.00391],
               ...,
               [ 3.00391,  3.00391,  ...,  3.00391,  3.00391],
               [ 3.00391,  3.00391,  ...,  3.00391,  3.00391]]]], shape=Shape([1, 1, 32, 32]), dtype=DataType::FLOAT32, layout=Layout::TILE)

[bbradelTT]

Update: turns out the last dimension needs to be a multiple of 2, therefore the workaround does not work in many cases.

[bbradelTT]

Since there's no workaround and this is blocking P0s (#13621 and #14530), along with #15511 moving this to a P0 as well.

[jaykru-tt]

A conversation with @yugaoTT yesterday which revealed that we will need new kernels to support last dim padding. I'm starting on those today and hope to be done by end of day on Friday.

[jaykru-tt]

Update on progress from yesterday: I've got the new stick-wise sharded reader written out on paper. I just need to transcribe it to code this morning. I'll also be adding a new program factory for these cases. Fortunately the writer I've already written should work just fine. After that, the rest of the week will hopefully be devoted to debugging. Thanks all for your patience! :)

[jaykru-tt]

Update on current progress: kernels and the new op path are complete, tested working on a few cases. I need to add a few more tests, do a few low hanging optimization, and will start off the PR process tomorrow. cc @bbradelTT

[jaykru-tt]

@esmalTT this is fixed with tests for a reduced version of your input in main. Going to close this now :)