Add Reduce ops workaround for keepDim=false

[mrakitaTT]

This PR adds TTNN workarounds for these Metal issues:

• [Feature Request] Reduce ops with keepdim=False are not supported tt-metal#13361 - By decomposing reduce(keepDim=false) into reduce(keepDim=true) + reshape
• [Bug Report] Reduce ops support reducing across more than 2 dims, but only if dims are not specified tt-metal#16118 - By annulling dimensions argument when all dims are being reduced

As part of this work I've also:

• Enabled conversion of stablehlo.reduce op with multiple reduce dimensions
• Added reduce ops verifiers in TTIR
• Added a separate function in TTNNWorkarounds to run rewrite patterns for decomposition and layout workarounds
• Added lots of unit tests for reduce ops to cover conversions and verifiers
• Added lots of silicon tests for reduce ops

Opened issue #1624 on myself to revert these workarounds once Metal issues are fixed.

Closes #805, #848

After implementing these workarounds and running tests, I've encountered another Metal issue, this time in reshape op. I've debugged it and I have a local fix, I will send a PR to fix it in Metal repo, confirmed with reshape op owners. I've opened myself an issue #1640 to enable Reduce ops silicon tests after this fix is uplifted.

Another issue that I've encountered while working on this is that after the workaround pass decompositions, if we are changing the shapes of the ops tensors, that means that their layout needs to be changed too, but layout pass is done before the workaround pass. I've managed to solve it by reusing the layout of the input tensor, but I am not sure if that is a good solution and maybe we need to repeat some of the layout logic again after workaround decompositions. FYI @sdjordjevicTT

Here is the example TTNN IR before the workarounds:

%3 = "ttnn.sum"(%2) <{dim_arg = [0: i32, 1 : i32, 2: i32], keep_dim = false}> : (tensor<128x32x4xf32, #ttnn_layout2>) -> tensor<1xf32, #ttnn_layout2>

and after the workarounds:

%3 = "ttnn.sum"(%2) <{keep_dim = true}> : (tensor<128x32x4xf32, #ttnn_layout2>) -> tensor<1x1x1xf32, #ttnn_layout2>
%4 = "ttnn.reshape"(%3) <{shape = [1 : i32]}> : (tensor<1x1x1xf32, #ttnn_layout2>) -> tensor<1xf32, #ttnn_layout3>

[mrakitaTT]

FYI: @AleksKnezevic @mmanzoorTT @vladimirjovanovicTT

[azecevicTT]

Thank you writing a great description of PR.

There are some points that repeat throughout this PR review, so I didn't want to repeat myself, mainly about SmallVector, SmallSet, const, ref and std::optional interface. I would like to hear opinions of others on this matter. I believe we agreed on some of those points, but some of them weren't discussed.

[mrakitaTT]

@azecevicTT Thank you for detailed review and pointers to llvm docs! I've learned a few new things 😄 I've left one comment unresolved where I had a different opinion. Please let me know if you disagree, and also if I missed to cover something.

[azecevicTT]

Thank you for addressing comments. Just go once more through changes, there are has_value, value on std::optional, std::vector, std::unordered_set, const and ref on MLIR types in the places where I didn't mark them.

Regarding ReduceOps verification it's okay on my end even if it stays as-is, because of limited scope of that function.

[LPanosTT]

@mrakitaTT It looks like @bbradelTT has a fix in metal to handle keepDim=False properly in this PR: tenstorrent/tt-metal#16163

It's not merged yet but we'll keep an eye on it.

[sdjordjevicTT]

Can you please attach MLIR IR before and after your change? It is a bit hard to imagine these kinds of changes by just looking just at the code :)

[mrakitaTT]

Can you please attach MLIR IR before and after your change? It is a bit hard to imagine these kinds of changes by just looking just at the code :)

@sdjordjevicTT Sure, added example IRs to the PR description, will do it for the future PRs too :)

[bbradelTT]

I just merged in tenstorrent/tt-metal#16163 to support keepdim=False

It should work for many inputs, with the caveat that you need to use a release build to avoid asserts in many cases.

[mrakitaTT]

I just merged in tenstorrent/tt-metal#16163 to support keepdim=False

It should work for many inputs, with the caveat that you need to use a release build to avoid asserts in many cases.

Thank you @bbradelTT I wish you had let me know that you are going to start working on it.

In any case, I've tested this workaround in debug mode with lots of shapes and all tests pass. I think I am still going to merge it and then once we uplift the new version of Metal I will test your fix and then if it also passes our tests I can remove that workaround (tracked by #1624). I'll let you know if some tests don't pass and with which shapes/parameters.