Need support for ttnn.max_pool2d to accept block and width sharded input.

[punithsekar]

Describe the bug
ttnn.max_pool2d supports only height_sharded input tensor. Need support for block_sharded and width_sharded input.

To Reproduce
Steps to reproduce the behavior:

1. Checkout to branch punith/maxpool_issue
   
2. Run command pytest tests/ttnn/integration_tests/yolov4/test_ttnn_neck.py

Expected behavior
To accept Block_sharded and width_sharded layout.

Screenshots

E       RuntimeError: TT_FATAL @ ../ttnn/cpp/ttnn/operations/pool/maxpool/max_pool2d.cpp:58: shard_scheme == TensorMemoryLayout::HEIGHT_SHARDED
E       info:
E       Only height sharded tensors are supported.
E       backtrace:
E        --- /home/ubuntu/punith/tt-metal/ttnn/ttnn/_ttnn.so(+0x458ec9) [0x7f090cd77ec9]

Please complete the following environment information:

• Device - WH-n150

Additional context
The input shape which we pass to maxpool is 1,10,10,512[NHWC]. Since the Channels is higher it should happen in width or block sharding to increasing the performance.

Current values when we use height sharding,

pool_1 = ttnn.max_pool2d(
            input_tensor=output_tensor,
            batch_size=1,
            input_h=10,
            input_w=10,
            channels=512,
            kernel_size=[5, 5],
            stride=[1, 1],
            padding=[2, 2],
            dilation=[1, 1],
            device=device,
        )

Attributes:
{'memory_config_':'MemoryConfig(memory_layout=TensorMemoryLayout::HEIGHT_SHARDED;buffer_type=BufferType::L1;shard_spec=ShardSpec(grid={[(x=0;y=0) - (x=3;y=0)]};shape={25; 0};orientation=ShardOrientation::ROW_MAJOR;halo=0))'; 'output_dtype_': 'DataType::BFLOAT16'; 'sliding_window_config_': 'SlidingWindowConfig(batch_size=1; input_hw=(10;10); window_hw=(5;5); stride_hw=(1;1); pad_hw=(2;2); dilation_hw=(1;1); num_cores_nhw=4; core_range_set_={[(x=0;y=0) - (x=3;y=0)]})'}

Core_count: 4

Kernel duration: 1077197 ns

[punithsekar]

fyi @saichandax

[dvartaniansTT]

@mywoodstock is there a plan to support this towards yolov4 optimization efforts? cc: @mbahnasTT

[mywoodstock]

@mywoodstock is there a plan to support this towards yolov4 optimization efforts? cc: @mbahnasTT

Yes, the PR is nearly ready to be merged

[mywoodstock]

This is now in main

[dvartaniansTT]

thanks for the update @mywoodstock ! great news! We will test this on yolov4 and once confirmed we can close this issue.
@punithsekar please test this asap and let's see how it improves perf for yolov4.

[punithsekar]

@mywoodstock @dvartaniansTT, I am able to pass block-sharded input to the maxpool, and the execution is happening without any issue. However, the PCC of output coming from maxpool is very low(~0.055). I have create separate issue #14206 for it.

[dvartaniansTT]

@mywoodstock is this on your radar?
@punithsekar does this mean we are running at almost 0 pcc end to end now?

[mywoodstock]

@dvartaniansTT Yes, its being worked on: #14249

[punithsekar]

@dvartaniansTT , Yes, we are getting almost 0 pcc. The bug is tracked in #14249 issue as Abhinav mentioned.