Issue

This PR cleans up an edge case that wasn't handled by #137243. The existing tiling code assumes that node.get_ranges() is a reliable source of pointwise and reduction numels. This is true for pointwise kernels, but the situation is more complicated with reductions. Since reductions change the number of elements in a tensor, not all ops within a reduction kernel will have the same number of iterations. For example, var_mean fuses pointwise division with the output of reduction sum, and the division lacks the corresponding reduction ranges.

Fix

Instead of getting numels from node.get_ranges(), explicitly pass the global pointwise and reduction numels to the relevant tiling functions. In SIMDKernel.complete_partial_tiling, we solve for the missing numel by diving the global numel by the partial tiling's numel. This ensures all tilings have the correct global numel.

Also, in SIMDKernel.is_compatible, add the global reduction numel to node ranges that are missing it. For example, {"x": 8, "r0_": 8} is compatible with a node of ranges ([8], []) when we have reduction_numel=8.

Finally, this PR generalizes some of the existing codegen to handle multiple reduction dims. We already had code to ignore reduction splits for pointwise kernels, but it only worked for 1D reductions. Now it can handle ND.

Test plan

This PR parametrizes the existing CI test for var_mean to also run with tiled reductions. It also adds a new test checking that var_mean generates 2D tilings (with tiled reduction enabled). These new tests would fail on the current main branch.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov