ttnn.group_norm usage guide

[ppadjinTT]

I am creating this issue as a single thread to understand how to pick the right parameters for ttnn.group_norm operation. We (tt-forge) are trying to add ttnn.group_norm operation to the tt-forge compiler in the form of TTNN and TTIR dialects in tt-mlir. The tricky part is picking the right grid size for a given input tensor.

Problem symptoms

For a given tensor, in these examples (N, 1, H*W, C) = (1, 8 * 8, 480) If grid_size for a given tensor is too big (e.g. 8x8 grid in gn-repro1.py, then division by zero happens (gn-repro1.log). On the other hand, for the same input if the grid is too small (e.g. 1x1 grid in gn-repro2.py), then L1 OOM happens because one Tensix has insufficient L1 to fit the input (gn-repro2.log).

Choosing the grid size

Any advice for how to choose the grid size for a given input tensor would be much appreciated! The ideal scenario would be that the user of ttnn.group_norm doesn't have to think about optimal grid placement, but in the meantime, any heuristic for choosing grid size is very useful.

Choosing num_out_blocks

Additionally, num_out_blocks parameter is 1 by default. It seems it also needs some hand tuning, so any advice on this is useful too.

gn-error1.log
gn-error2.log
gn-repro1.py
gn-repro2.py

)

[bbradelTT]

@ppadjinTT until the usage guide can be completed, please use ttnn.determine_expected_group_norm_sharded_config_and_grid_size() to get information, which is what is mentioned in the documentaiton.
E.g. in the first example

    sharded_mem_config, grid_size = ttnn.determine_expected_group_norm_sharded_config_and_grid_size(
    device=device,
    num_channels=num_channels,
    num_groups=num_groups,
    input_nhw=1 * 8 * 8,
    is_height_sharded=False,#True,
    is_row_major=False,#True,
    )
    print(f"SMC {sharded_mem_config} GS {grid_size}")

produces

SMC MemoryConfig(memory_layout=TensorMemoryLayout::BLOCK_SHARDED,buffer_type=BufferType::L1,shard_spec=ShardSpec{grid=[{"start":{"x":0,"y":0},"end":{"x":1,"y":7}], shape=[32, 16], orientation=ShardOrientation::COL_MAJOR},nd_shard_spec=std::nullopt,created_with_nd_shard_spec=0) GS ttnn.CoreGrid(x=2, y=8)