Rearchitecture: Xe epilogue

[petercad]

This PR builds on #573, adding a CollectiveEpilogue with support for the new block 2D copy atoms.

The existing epilogue implementation was mostly rewritten, as it had many hardcoded assumptions and limitations:

• Subgroups own a contiguous tile within the workgroup tile
• Subgroup tiles are laid out n-major within the workgroup tile
• C/D atoms have the same block size
• One copy atom of data is processed at a time
• C/D atoms must bring data in the exact same layout as the accumulator

The new implementation removes all these restrictions.

Its API is also somewhat different, mostly in ways that more closely match the SM90 epilogues:

• Configurable EpilogueTile template parameter controls the block size for epilogue computation.
• Fusion callbacks receive workgroup-scope tiling information, not subgroup-scope tiling information (because CuTe's TiledMMA is very flexible -- the subgroup "tile" may not be contiguous).
• Vectorization for the epilogue compute operations is configurable via the ComputeVectorLen constexpr variable. Currently this is set to operate on one MMA atom's worth of accumulator data at a time, but if we want to make it user-configurable like the NV epilogues (where it's a template parameter for the dispatch policy), that's possible.
• It receives the TiledMMA as a template parameter rather than an argument to operator().
• The S2R/R2S copy operation parameters are omitted (a difference vs. SM90) as they are irrelevant to both the old and new epilogue implementation.

The new implementation glues together C/D loads and compute with reorders, so it can support efficient data type and layout conversions outside of the epilogue computation.

[rolandschulz]

isn't this the same as the static_cast? why do you use static_cast in one case and this in the other?

[petercad]

Using a static_cast here induces a copy of the tensor data, rather than reusing the existing data, which is the intention here.

[rolandschulz]

why does this take a rvalue-ref?

[petercad]

The idea is to introduce two flavors of make_subgroup_tensor. Given an lvalue reference, it makes a view of an existing rmem Tensor. Given an rvalue reference, it assumes ownership of the incoming Tensor's data.

[rolandschulz]

Nit: Nvidia uses "_needed" instead of "_supported". I think that's a better name.

[petercad]

I agree for destination. For source, it also depends on the fusion -- that check happens in operator().

[taozha2]

This changed epilogue API, and bind the MMA to epilogue, I think epilogue is an independent component which is not only for MMA, right? and what will happen if "TiledMma == void" here?

[petercad]

The epilogue already requires the TiledMMA. Before it was an argument to operator() -- in this PR I made it a template parameter to the epilogue itself. Either way we must have the TiledMMA to understand what data the accumulator contains so we know where to write it to global memory.