ttnn.conv2d fails with a special shape

[jerrysky3]

During testing I found this special shape will fail ttnn.conv2d with the error below:

RuntimeError: TT_FATAL @ ../ttnn/cpp/ttnn/operations/conv/conv2d/device/conv2d_op_width_sharded_program_factory.cpp:61: a.get_layout() == Layout::ROW_MAJOR
info:
Conv activation should be in row major layout
backtrace:
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(+0x3f5c38) [0x7fd42beaac38]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(_ZN4ttnn10operations4conv6conv2d47multi_core_optimized_conv_width_sharded_v2_implERN2tt8tt_metal2v07ProgramERKNS4_6TensorESA_RKNS_5types5ShapeENSt3__18optionalIS9_EESH_NS0_14sliding_window19SlidingWindowConfigEjjbbbRKNS2_34OptimizedConvParallelizationConfigERKNS2_24OptimizedConvBlockConfigEbbRS8_NSF_7variantIJNS_28GrayskullComputeKernelConfigENS_27WormholeComputeKernelConfigEEEEbbb+0x63f1) [0x7fd42c033751]
 --- /home/jerry/tt-metal/ttnn/ttnn/_ttnn.so(_ZN4ttnn10operations4conv6conv2d40multi_core_optimized_conv_sharded_v2_newERKN2tt8tt_metal6TensorES7_NSt3__18optionalIS6_EENS0_14sliding_window19SlidingWindowConfigEjjbb12MathFidelityRKNS2_34OptimizedConvParallelizationConfigERKNS2_24OptimizedConvBlockConfigENS4_8DataTypeENS8_5arrayIjLm4EEEbNS9_IKNS8_7variantIJNS_28GrayskullComputeKernelConfigENS_27WormholeComputeKernelConfigEEEEEERS5_bbbbb+0x5bc) [0x7fd42c028eac]
 --- ttnn::operations::conv::conv2d::OptimizedConvNew::create_program(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<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>>&) const

To reproduce

Run the example below:

import torch
import ttnn


def main(device):
    torch_input = torch.rand((1, 1, 7, 7), dtype=torch.bfloat16)
    torch_weight = torch.rand((16, 1, 2, 2), dtype=torch.bfloat16)

    stride = 2
    padding = 0
    dilation = 1
    
    input_tensor = ttnn.from_torch(torch_input, layout=ttnn.TILE_LAYOUT, device=device)
    input_tensor = ttnn.permute(input_tensor, (0, 2, 3, 1))
    weight_tensor = ttnn.from_torch(torch_weight, layout=ttnn.TILE_LAYOUT, device=device)
    weight_tensor = ttnn.from_device(weight_tensor)
    weight_tensor = ttnn.to_layout(weight_tensor, layout=ttnn.ROW_MAJOR_LAYOUT)

    output_tensor, _, _, _, _ = ttnn.conv2d(
        input_tensor=input_tensor,
        weight_tensor=weight_tensor,
        batch_size=1,
        in_channels=1,
        out_channels=16,
        input_height=7,
        input_width=7,
        kernel_size=(2, 2),
        stride=(stride, stride),
        padding=(padding, padding),
        dilation=(dilation, dilation),
        groups=1,
        device=device,
    )


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

[mywoodstock]

Unable to repro as I'm hitting that reshape error before it reaches conv :(

[mywoodstock]

OK, ignoring that reshape assert, the example code runs to completion for me.

[jerrysky3]

OK, ignoring that reshape assert, the example code runs to completion for me.

Thanks for looking into this. So looks like this issue won't exist anymore once tenstorrent/pytorch2.0_ttnn#429 is resolved?