Skip to content

Use newer version of copy_atom in epilogue collective #573

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
+610 −68
Open

Use newer version of copy_atom in epilogue collective #573

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
578ae95
Add new atoms in collectiveEpilogue
Oct 22, 2025
aaf2685
Use new make_block_2d_copy_{C,D} APIs for loads/stores
Oct 22, 2025
407a875
Code Cleanup
Oct 22, 2025
45ac04e
Merge branch 'main' into anamikac/add-newatoms-collectiveEpilogue
anamikac-intel Oct 23, 2025
12282e8
Rename xe_epilogue_legacy.cpp to xe_epilogue_legacy.hpp
anamikac-intel Oct 23, 2025
6596ac8
Update xe_epilogue_legacy.hpp
anamikac-intel Oct 23, 2025
7ea69d5
Avoid register spills
Oct 23, 2025
6944d90
remove hardcode layout
Oct 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions examples/00_bmg_gemm/00_bmg_gemm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -350,6 +350,10 @@ int main(int argc, const char** argv)
// Refer https://github.com/intel/sycl-tla/blob/main/media/docs/cpp/xe_rearchitecture.md
using GmemTiledCopyA = void; //XE_LOAD_2D<16, 32, 32>;
using GmemTiledCopyB = void; //XE_LOAD_2D_VNNI<16, 32, 32>;
using GmemTiledCopyC = XE_LOAD_2D<32, 8, 16>;
using GmemTiledCopyD = XE_STORE_2D<32, 8, 16>;



// Workgroup-level tile
using TileShape = Shape<_256, _256, _32>;
Expand All @@ -369,9 +373,8 @@ int main(int argc, const char** argv)

// For Intel BMG, PipelineStages defines how many k-blocks ahead to prefetch from A and B.
constexpr int PipelineStages = 2;
// For older version of copy/mma atom, use cutlass::gemm::MainloopIntelXeXMX16 as dispatch policy
using GEMMDispatchPolicy = cutlass::gemm::MainloopXeL1Staged<PipelineStages>;
using EpilogueDispatchPolicy = cutlass::epilogue::IntelXeXMX16;
using EpilogueDispatchPolicy = cutlass::epilogue::IntelXeL1Staged;

// This is the 'default' epilogue operation (Linear Combination) which performs everything in:
// (D = alpha * (A*B) + beta * C)
Expand All @@ -394,9 +397,9 @@ int main(int argc, const char** argv)
ElementOutput,
cutlass::gemm::TagToStrideC_t<LayoutD>, // Converts CUTLASS 2.x to CUTLASS 3.x representation
FusionCallBacks,
XE_2D_U32x8x16_LD_N, // The copy atom used to load matrix C
GmemTiledCopyC, // The copy atom used to load matrix C
void, void,
XE_2D_U32x8x16_ST_N, // The copy atom used to store matrix D
GmemTiledCopyD, // The copy atom used to store matrix D
void, void>;

// GEMM Mainloop - iteration over blocks in K dimension
Expand Down
12 changes: 11 additions & 1 deletion include/cute/atom/copy_traits_xe_2d.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1143,12 +1143,22 @@ template <class CopyOp, class TiledMMA, class CTensor>
auto get_block_2d_copy_C(TiledMMA const& tiled_mma, CTensor const& c_tensor)
{
if constexpr (!std::is_void_v<CopyOp>) {
return make_block_2d_copy_C(CopyOp{}, tiled_mma, c_tensor);
return make_block_2d_copy_CD(CopyOp{}, tiled_mma, c_tensor);
} else {
return make_block_2d_copy_C(tiled_mma, c_tensor);
}
}

template <class CopyOp, class TiledMMA, class DTensor>
auto get_block_2d_copy_D(TiledMMA const& tiled_mma, DTensor const& d_tensor)
{
if constexpr (!std::is_void_v<CopyOp>) {
return make_block_2d_copy_CD(CopyOp{}, tiled_mma, d_tensor);
} else {
return make_block_2d_copy_D(tiled_mma, d_tensor);
}
}

//
// Display utilities
//
Expand Down
1 change: 1 addition & 0 deletions include/cutlass/epilogue/collective/collective_epilogue.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ class CollectiveEpilogue {
#include "sm100_epilogue_array_tma_warpspecialized.hpp"
#if defined (SYCL_INTEL_TARGET)
#include "xe_epilogue.hpp"
#include "xe_epilogue_legacy.hpp"
#include "xe_array_epilogue.hpp"
#endif
//
Expand Down
131 changes: 70 additions & 61 deletions include/cutlass/epilogue/collective/xe_epilogue.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ template <
class CopyOpR2S_
>
class CollectiveEpilogue<
IntelXeXMX16,
IntelXeL1Staged,
CtaTileMNK_,
ElementC_,
StrideC_,
Expand All @@ -86,7 +86,7 @@ class CollectiveEpilogue<
//
// Type Aliases
//
using DispatchPolicy = IntelXeXMX16;
using DispatchPolicy = IntelXeL1Staged;
using CtaTileMNK = CtaTileMNK_;
using FusionCallbacks = FusionCallbacks_;
using ElementC = ElementC_;
Expand All @@ -102,9 +102,6 @@ class CollectiveEpilogue<
using CopyOpR2S = CopyOpR2S_;

using ThreadEpilogueOp = typename fusion::FusionCallbacksTraits<FusionCallbacks>::Operation;
using GmemTiledCopyC = conditional_t<cute::is_void_v<CopyOpG2R>, XE_2D_U32x8x16_LD_N, CopyOpG2R>;
using GmemTiledCopyD = cute::conditional_t<not cute::is_void_v<ElementD> && not cute::is_void_v<CopyOpR2G>,
CopyOpR2G, XE_2D_U32x8x16_ST_N>;
using ElementOutput = ElementD;
using ElementCompute = ElementAccumulator;

Expand All @@ -119,19 +116,10 @@ class CollectiveEpilogue<
static_assert(std::is_same_v<SmemLayoutAtomC, void>, "Copy operation to shared memory is not supported");
static_assert(std::is_same_v<SmemLayoutAtomD, void>, "Copy operation to shared memory is not supported");

using CopyThreadShape = Shape<_1, Int<SubgroupSize>>;

using Trait_C = Copy_Traits<GmemTiledCopyC, StrideC>;
using val_layout_load_C = decltype(make_layout(shape_div(typename Trait_C::BlockShape{}, CopyThreadShape{})));
using XE_Copy_C = decltype(make_tiled_copy(Copy_Atom<Trait_C, ElementC>{}, Layout<CopyThreadShape>{}, val_layout_load_C{}));

using Trait_D = Copy_Traits<GmemTiledCopyD, StrideD>;
using val_layout_store_D = decltype(make_layout(shape_div(typename Trait_D::BlockShape{}, CopyThreadShape{})));
using XE_Copy_D = decltype(make_tiled_copy(Copy_Atom<Trait_D, ElementD>{}, Layout<CopyThreadShape>{}, val_layout_store_D{}));

//remember this PR https://github.com/intel/sycl-tla/pull/565/files
private:
constexpr static bool is_source_supported = not cute::is_void_v<ElementC> && not cute::is_void_v<CopyOpG2R>;
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jiyang1011 - The validation logic from PR #565 that sets is_source_supported to false when CopyOpG2R is void needs updating. With this PR's automatic ops selection, both CopyOpG2R and CopyOpR2G can now legitimately be void since make_block_2d_copy_* automatically selects appropriate operations.

Copy link

@jiyang1011 jiyang1011 Oct 23, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we set a default copy trait like XeCopyAuto or something else which will also call make_block_2d_copy_* ?

constexpr static bool is_destination_supported = not cute::is_void_v<ElementD> && not cute::is_void_v<CopyOpR2G>;
constexpr static bool is_source_supported = not cute::is_void_v<ElementC>;
constexpr static bool is_destination_supported = not cute::is_void_v<ElementD>;

constexpr static bool is_m_major_C = detail::is_m_major<StrideC>();
constexpr static bool is_m_major_D = detail::is_m_major<StrideD>();
Expand All @@ -154,6 +142,15 @@ class CollectiveEpilogue<
};
using TensorStorage = typename SharedStorage::TensorStorage;

// Helper to get tensor types
template<class Element, class Stride>
using TensorTypeC = decltype(make_tensor(make_gmem_ptr(static_cast<Element const*>(nullptr)),
make_layout(make_shape(int{}, int{}, int{}), Stride{})));

template<class Element, class Stride>
using TensorTypeD = decltype(make_tensor(make_gmem_ptr(static_cast<Element*>(nullptr)),
make_layout(make_shape(int{}, int{}, int{}), Stride{})));

// Host side epilogue arguments
struct Arguments {
typename FusionCallbacks::Arguments thread{};
Expand All @@ -166,8 +163,8 @@ class CollectiveEpilogue<
// Device side epilogue params
struct Params {
typename FusionCallbacks::Params thread{};
XE_Copy_C xe_load_c;
XE_Copy_D xe_store_d;
TensorTypeC<ElementC, StrideC> mC;
TensorTypeD<ElementD, StrideD> mD;
};

//
Expand All @@ -183,23 +180,13 @@ class CollectiveEpilogue<
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_MNKL;

XE_Copy_C xe_load_c = {};
if constexpr (is_source_supported) {
auto mC = make_tensor(make_gmem_ptr(args.ptr_C), make_layout(make_shape(M, N, L), args.dC));
xe_load_c = {xe_load_c.with(mC)};
}

XE_Copy_D xe_store_d = {};
if constexpr (is_destination_supported) {
auto mD = make_tensor(make_gmem_ptr(args.ptr_D), make_layout(make_shape(M, N, L), args.dD));
xe_store_d = {xe_store_d.with(mD)};
}
auto mC = make_tensor(make_gmem_ptr(args.ptr_C), make_layout(make_shape(M, N, L), args.dC));
auto mD = make_tensor(make_gmem_ptr(args.ptr_D), make_layout(make_shape(M, N, L), args.dD));

return {
FusionCallbacks::to_underlying_arguments(problem_shape, args.thread, workspace),
xe_load_c,
xe_store_d,
mC,
mD
};
}

Expand Down Expand Up @@ -270,6 +257,24 @@ class CollectiveEpilogue<
return fusion_callbacks.is_producer_load_needed();
}

template<typename Tensor>
CUTLASS_DEVICE auto reshape_into_smaller_fragments(Tensor&& tensor) {
using namespace cute;

auto target_stride = make_stride(
make_stride(cute::ScaledBasis<cute::Int<1>, 0>{}, _0{}),
cute::ScaledBasis<cute::Int<8>, 0>{},
cute::ScaledBasis<cute::Int<16>, 1>{}
);

auto target_layout = make_layout(
make_shape(make_shape(_8{}, _1{}), _4{}, _4{}),
target_stride
);

return make_tensor(tensor.data(), target_layout);
}

template<
class ProblemShapeMNKL,
class TileShapeMNK,
Expand All @@ -286,7 +291,6 @@ class CollectiveEpilogue<
TiledMma tiled_mma,
int thread_idx) {

(void) tiled_mma;
using namespace cute;

static_assert(cute::rank(CtaTileMNK{}) == 3, "CtaTileMNK must be rank-3: [CTA_M, CTA_N, CTA_K]");
Expand All @@ -297,12 +301,11 @@ class CollectiveEpilogue<
static constexpr auto BLK_M = get<0>(CtaTileMNK{});
static constexpr auto BLK_N = get<1>(CtaTileMNK{});
static constexpr auto BLK_K = get<2>(CtaTileMNK{});
// static_assert(is_same_v<typename TiledMma::ThrLayoutVMNK, int>, "assertation fail");
static constexpr auto ATOM_M = get<1>(typename TiledMma::ThrLayoutVMNK{}.shape());
static constexpr auto ATOM_N = get<2>(typename TiledMma::ThrLayoutVMNK{}.shape());
static constexpr auto ATOM_K = get<3>(typename TiledMma::ThrLayoutVMNK{}.shape());
static_assert(

static_assert(
BLK_M % ATOM_M == 0 &&
BLK_N % ATOM_N == 0 &&
BLK_K % ATOM_K == 0,
Expand All @@ -316,46 +319,49 @@ class CollectiveEpilogue<
static constexpr int FragsN = get<1>(SubgroupTileShape{}) / get<1>(MmaAtomShape()); // B frags per sub_group

static constexpr int FragmentSize = (get<0>(MmaAtomShape()) * get<1>(MmaAtomShape())) / SubgroupSize;

// Indexing variables
auto [M, N, K, L] = problem_shape_mnkl;
auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl;
auto m_sg = get_sub_group_id() / ATOM_N;
auto n_sg = get_sub_group_id() % ATOM_N;

auto mn_shape = shape(typename decltype(params.xe_store_d)::Tiler_MN{});

auto sg_local_m_coord = get_sub_group_id() / ATOM_N;
auto sg_local_n_coord = get_sub_group_id() % ATOM_N;

auto sg_m_coord = m_coord * ATOM_M + sg_local_m_coord;
auto sg_n_coord = n_coord * ATOM_N + sg_local_n_coord;
auto sg_coord = make_coord(sg_m_coord, sg_n_coord, k_coord, l_coord);


auto wg_coord = make_coord(m_coord, n_coord, k_coord, l_coord);
bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed();

auto batch_idx = get<3>(wg_coord);
auto copy_c = get_block_2d_copy_C<CopyOpG2R>(tiled_mma, params.mC(_,_,batch_idx));
auto copy_d = get_block_2d_copy_D<CopyOpR2G>(tiled_mma, params.mD(_,_,batch_idx));



// Represent the full output tensor
Tensor mD_mnl = cute::get_xe_tensor(make_shape(M,N,L));

// Tile the output tensor per WG and select the tile for current WG
Tensor g_wg_D = local_tile(mD_mnl, take<0,2>(CtaTileMNK{}), make_coord(m_coord,n_coord,l_coord)); // (BLK_M,BLK_N)

// Tile the output tensor per SG and select tile for the current SG
Tensor gD = local_tile(g_wg_D, take<0,2>(SubgroupTileShape{}), make_coord(m_sg,n_sg)); // (SG_M,SG_N)
// Tile the output tensor for the current workgroup
Tensor gD = local_tile(mD_mnl, take<0,2>(CtaTileMNK{}), remove<2>(wg_coord)); // (BLK_M,BLK_N)

auto thread_xe_load_c = params.xe_load_c.get_thread_slice(thread_idx);
// Get thread-level partitioning across the entire workgroup tile
auto thread_xe_load_c = copy_c.get_thread_slice(thread_idx);
Tensor tCgC = thread_xe_load_c.partition_S(gD);
Copy link
Author

@anamikac-intel anamikac-intel Oct 23, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@petercad - I was trying to fix register spills... reducing copy operations not helping much. Issue seems to coming from tiling done.. In legacy code after partition tCgC and tCgD was :
tCgC : ArithTuple(0,0,0) o ((_8,_1),_4,_4):((_1@0,_0),_8@0,_16@1)
tCgD : ArithTuple(0,0,0) o ((_8,_1),_4,_4):((_1@0,_0),_8@0,_16@1)

So we have 8 fragments of size 4 x 4

whereas in new code we have 128 fragments of size 1 x 1:

tCgC: ArithTuple(0,0,0) o ((_8,(_4,_4)),_1,_1):((_1@0,(_8@0,_16@1)),_0,_0)
tCgD: ArithTuple(0,0,0) o ((_8,(_4,_4)),_1,_1):((_1@0,(_8@0,_16@1)),_0,_0)

I tried titling further with SubgroupTileShape{}: (_32,_64,_32) but same result
g_wg_D: ArithTuple(0,0,0) o (_256,_256):(_1@0,_1@1)
gD: ArithTuple(0,0,0) o (_32,_64):(_1@0,_1@1)
tCgC: ArithTuple(0,0,0) o ((_8,(_4,_4)),_1,_1):((_1@0,(_8@0,_16@1)),_0,_0)
tCgD: ArithTuple(0,0,0) o ((_8,(_4,_4)),_1,_1):((_1@0,(_8@0,_16@1)),_0,_0)

This seems to be the actual issue.. so when I reshaped the layout to tCgC/tCgD ArithTuple(0,0,0) o ((_8,_1),_4,_4):((_1@0,_0),_8@0,_16@1) (8 fragments of 4x4) the perf drop is fixed. But re-layouting tCgC/tCgD might not be best option so can you please check.

image

Copy link

@sanchitintel sanchitintel Oct 23, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we have 8 fragments of size 4 x 4

No, the per-thread D or C fragment size is 128 elements in both cases, but layout is different.
The WG tile size is (256, 256, 32).
There are 32 subgroups, each with 16 threads.
Each SG fragment for C or D is (32, 64) spatially since the subgroup layout is 8x4 in the example.
Each C or D thread fragment is sized 128 elements.

Both ((_8,_1),_4,_4) and ((_8,(_4,_4)),_1,_1) have 128 elements.

This comment was marked as outdated.

Sign in to view

This comment was marked as duplicate.

Sign in to view

Copy link

@sanchitintel sanchitintel Oct 24, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @petercad, recent commits having a lot more changes than just the thread fragment layout (seems to be equivalent to the previous one) seem to suggest that there's a lurking factor that fixed the performance issues that were observed earlier in this PR, and that the thread-fragment layout of new C, D copy atoms isn't problematic.

Thanks!

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes this code close to legacy one except I am using new copy atoms with the reshaping layout... but only concern is it only works with ops that has 16 width × 8 height #573 (comment)

Copy link

@petercad petercad Oct 24, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@sanchitintel -- will look at it in more detail shortly. The underlying cause for earlier regressions seems to be twofold:

  • Code scheduling issues in IGC. It seems it is not moving loads/stores around sufficiently to reduce register pressure.
  • For the C loads, make_block_2d_copy_c will try to make the blocks as large as possible (because it's operating on the assumption that you're loading all of C at once) but that brings additional register pressure

The second point is not enough to explain the spills (there is plenty of register space even if you do load huge chunks of C), but it aggravates the first point.


auto thread_xe_store_d = params.xe_store_d.get_thread_slice(thread_idx);
auto thread_xe_store_d = copy_d.get_thread_slice(thread_idx);
Tensor tCgD = thread_xe_store_d.partition_D(gD);

auto tCgC_frag = reshape_into_smaller_fragments(tCgC);
auto tCgD_frag = reshape_into_smaller_fragments(tCgD);

Tensor trC = make_tensor<typename TiledMma::ValTypeC>(Shape<Int<FragmentSize>>{});
Tensor trD_compute = make_tensor<ElementCompute>(Shape<Int<FragmentSize>>{});

// Because Sm90 uses shared memory, they are not tied to using the same accumulator values
// for MMA and Epilogue. But because we are operating directly in the accumulators, we need to be
// sure that we are operating on the same values.
ThrCopy thread_g2r = params.xe_load_c.get_slice(thread_idx);
ThrCopy thread_g2r = copy_c.get_slice(thread_idx);
auto mn_shape = shape(typename decltype(copy_d)::Tiler_MN{});

// OOB predication for tile quantization "residue"
// Absolute coordinate tensors (dynamic)
Expand All @@ -364,7 +370,7 @@ class CollectiveEpilogue<
Tensor cD_mn = local_tile(mD_crd, take<0,2>(CtaTileMNK{}), make_coord(m_coord, n_coord)); // (CTA_M,CTA_N)
Tensor tRS_cD_mn = thread_g2r.partition_S(flat_divide(cD_mn, mn_shape)); // (G2R,G2R_M,G2R_N,EPI_M,EPI_N)

Tensor tRS_cD = make_coord_tensor(tRS_cD_mn.layout()); // (G2R,G2R_M,G2R_N,EPI_M,EPI_N)
Tensor tRS_cD = make_coord_tensor(tRS_cD_mn.layout());

// Get the fusion callbacks
// Arguments passed here relate to sub-group tiles, rather than CTA (work-group) tiles
Expand All @@ -376,7 +382,7 @@ class CollectiveEpilogue<
sg_coord,
tiled_mma,
mn_shape,
params.xe_store_d,
copy_d,
cD,
residue_mn,
tRS_cD,
Expand All @@ -398,7 +404,8 @@ class CollectiveEpilogue<
FragsM * FragsN * FragmentSize * SubgroupSize * ATOM_M * ATOM_N * ATOM_K;
constexpr int MN = get<0>(CtaTileMNK{}) * get<1>(CtaTileMNK{});
static_assert(ValuesLoaded == MN, "the total elements loaded by all threads should be the same as MxN" );



auto synchronize = [&] () {};
CUTLASS_PRAGMA_UNROLL
for (int epi_n = 0; epi_n < FragsN; epi_n++) {
Expand All @@ -407,7 +414,7 @@ class CollectiveEpilogue<
cst_callbacks.begin_loop(epi_m, epi_n);

if (is_C_load_needed) {
copy(params.xe_load_c, tCgC(_, epi_m, epi_n), trC);
copy(copy_c, tCgC_frag(_, epi_m, epi_n), trC);
}

cst_callbacks.previsit(epi_m, epi_n, 0, is_C_load_needed);
Expand All @@ -419,21 +426,23 @@ class CollectiveEpilogue<
trD_compute_frag(epi_v) = cst_callbacks.visit(acc_frag_mn(epi_v), epi_v, epi_m, epi_n);
}
cst_callbacks.reduce(nullptr, synchronize, epi_m, epi_n, (epi_m == FragsM - 1 && epi_n == FragsN - 1), trD_compute_frag);

if constexpr (is_destination_supported) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(trD_compute_frag); ++i) {
trD_frag(i) = cutlass::NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize>{}(trD_compute_frag(i));
}
copy(params.xe_store_d, trD, tCgD(_, epi_m, epi_n));
copy(copy_d, trD, tCgD_frag(_, epi_m, epi_n));
}

cst_callbacks.end_loop(epi_m, epi_n);

}
}

cst_callbacks.end();
}

}

private:
Params const& params;
Expand All @@ -447,4 +456,4 @@ class CollectiveEpilogue<
} // namespace epilogue
} // namespace cutlass

/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
Loading
Loading