Use newer version of copy_atom in epilogue collective #573
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -68,7 +68,7 @@ template < | |
| class CopyOpR2S_ | ||
| > | ||
| class CollectiveEpilogue< | ||
| IntelXeL1Staged, | ||
| CtaTileMNK_, | ||
| ElementC_, | ||
| StrideC_, | ||
|
|
@@ -86,7 +86,7 @@ class CollectiveEpilogue< | |
| // | ||
| // Type Aliases | ||
| // | ||
| using DispatchPolicy = IntelXeL1Staged; | ||
| using CtaTileMNK = CtaTileMNK_; | ||
| using FusionCallbacks = FusionCallbacks_; | ||
| using ElementC = ElementC_; | ||
|
|
@@ -102,9 +102,6 @@ class CollectiveEpilogue< | |
| using CopyOpR2S = CopyOpR2S_; | ||
|
|
||
| using ThreadEpilogueOp = typename fusion::FusionCallbacksTraits<FusionCallbacks>::Operation; | ||
| using ElementOutput = ElementD; | ||
| using ElementCompute = ElementAccumulator; | ||
|
|
||
|
|
@@ -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"); | ||
|
|
||
| //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>; | ||
| 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>(); | ||
|
|
@@ -166,8 +154,11 @@ class CollectiveEpilogue< | |
| // Device side epilogue params | ||
| struct Params { | ||
| typename FusionCallbacks::Params thread{}; | ||
| ElementC const* ptr_C; | ||
| ElementD* ptr_D; | ||
| int M, N, K, L; | ||
| StrideC dC; | ||
| StrideD dD; | ||
| }; | ||
|
|
||
| // | ||
|
|
@@ -184,22 +175,13 @@ class CollectiveEpilogue< | |
| auto problem_shape_MNKL = append<4>(problem_shape, 1); | ||
| auto [M, N, K, L] = problem_shape_MNKL; | ||
|
|
||
| return { | ||
| FusionCallbacks::to_underlying_arguments(problem_shape, args.thread, workspace), | ||
| args.ptr_C, | ||
| args.ptr_D, | ||
| M, N, K, L, | ||
| args.dC, | ||
| args.dD | ||
| }; | ||
| } | ||
|
|
||
|
|
@@ -286,7 +268,6 @@ class CollectiveEpilogue< | |
| TiledMma tiled_mma, | ||
| int thread_idx) { | ||
|
|
||
| using namespace cute; | ||
|
|
||
| static_assert(cute::rank(CtaTileMNK{}) == 3, "CtaTileMNK must be rank-3: [CTA_M, CTA_N, CTA_K]"); | ||
|
|
@@ -297,143 +278,149 @@ 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{}); | ||
|
|
||
| // Indexing variables | ||
| auto [M, N, K, L] = problem_shape_mnkl; | ||
| auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl; | ||
|
|
||
| // Workgroup coordinates (no subgroup indexing needed) | ||
| auto wg_coord = make_coord(m_coord, n_coord, k_coord, l_coord); | ||
| auto batch_idx = get<3>(wg_coord); | ||
|
|
||
| bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed(); | ||
| auto mC = make_tensor(make_gmem_ptr(params.ptr_C), make_layout(make_shape(params.M, params.N, params.L), params.dC)); | ||
| auto mD = make_tensor(make_gmem_ptr(params.ptr_D), make_layout(make_shape(params.M, params.N, params.L), params.dD)); | ||
|
|
||
| auto copy_c = [&]() { | ||
| if constexpr (!std::is_void_v<CopyOpG2R>) { | ||
| return make_block_2d_copy_CD(CopyOpG2R{}, tiled_mma, mC(_,_,batch_idx)); | ||
| } else { | ||
| return make_block_2d_copy_C(tiled_mma, mC(_,_,batch_idx)); | ||
| } | ||
| }(); | ||
| auto copy_d = [&]() { | ||
| if constexpr (!std::is_void_v<CopyOpR2G>) { | ||
| return make_block_2d_copy_CD(CopyOpR2G{}, tiled_mma, mD(_,_,batch_idx)); | ||
| } else { | ||
| return make_block_2d_copy_D(tiled_mma, mD(_,_,batch_idx)); | ||
| } | ||
| }(); | ||
|
|
||
|
|
||
| // Represent the full output tensor | ||
| Tensor mD_mnl = cute::get_xe_tensor(make_shape(M,N,L)); | ||
|
|
||
| // 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) // change made | ||
|
|
||
| // 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); | ||
|
||
|
|
||
| auto thread_xe_store_d = copy_d.get_thread_slice(thread_idx); | ||
| Tensor tCgD = thread_xe_store_d.partition_D(gD); | ||
|
|
||
| // Create tensors sized for workgroup-level operation | ||
| Tensor trC = make_tensor<typename TiledMma::ValTypeC>(tCgC.shape()); | ||
| Tensor trD_compute = make_tensor<ElementCompute>(tCgD.shape()); | ||
| Tensor trD = make_tensor<ElementOutput>(tCgD.shape()); | ||
|
|
||
| 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" | ||
| Tensor mD_crd = make_identity_tensor(make_shape(M,N)); | ||
| Tensor cD = local_tile(mD_crd, take<0,2>(CtaTileMNK{}), make_coord(m_coord, n_coord)); | ||
| Tensor tRS_cD = thread_g2r.partition_S(flat_divide(cD, mn_shape)); | ||
|
|
||
| Tensor tRS_cD_coord = make_coord_tensor(tRS_cD.layout()); | ||
|
|
||
| // Get fusion callbacks at workgroup level | ||
| constexpr bool RefSrc = true; | ||
| auto residue_mn = make_coord(M, N); | ||
| auto cst_args = cutlass::epilogue::fusion::detail::ConsumerStoreArgs{ | ||
| problem_shape_mnkl, | ||
| CtaTileMNK{}, // Use workgroup tile shape | ||
| wg_coord, // Use workgroup coordinates | ||
| tiled_mma, | ||
| mn_shape, | ||
| copy_d, | ||
| cD, | ||
| residue_mn, | ||
| tRS_cD_coord, | ||
| residue_mn, | ||
| trC, | ||
| thread_idx, | ||
| }; | ||
| auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks<RefSrc>(cst_args); | ||
| auto synchronize = [&] () {}; | ||
|
|
||
| cst_callbacks.begin(); | ||
| // Load C tile if needed (distributed across all threads in workgroup) | ||
| if (is_C_load_needed) { | ||
| copy(copy_c, tCgC, trC); | ||
| } | ||
|
|
||
| // Single previsit for entire workgroup tile | ||
| cst_callbacks.previsit(0, 0, 0, is_C_load_needed); | ||
|
|
||
| static constexpr int FragmentSize = get<0>(MmaAtomShape()) * get<1>(MmaAtomShape()); | ||
| constexpr int num_fragments = size(accumulators) / FragmentSize; | ||
|
|
||
| CUTLASS_PRAGMA_UNROLL | ||
| for (int epi_v = 0; epi_v < num_fragments; ++epi_v) { | ||
| // Extract fragment | ||
| Array<ElementAccumulator, FragmentSize> frg_acc; | ||
| CUTLASS_PRAGMA_UNROLL | ||
| for (int f = 0; f < FragmentSize; ++f) { | ||
| frg_acc[f] = accumulators(epi_v * FragmentSize + f); | ||
| } | ||
|
|
||
| // Process fragment | ||
| auto result_frg = cst_callbacks.visit(frg_acc, epi_v, 0, 0); | ||
|
|
||
| // Store results | ||
| CUTLASS_PRAGMA_UNROLL | ||
| for (int f = 0; f < FragmentSize; ++f) { | ||
| trD_compute(epi_v * FragmentSize + f) = result_frg[f]; | ||
| } | ||
| } | ||
|
|
||
| cst_callbacks.reduce(nullptr, synchronize, 0, 0, true, trD_compute); | ||
|
|
||
| if constexpr (is_destination_supported) { | ||
| // Convert fragments using NumericArrayConverter | ||
| constexpr int num_fragments_trD_compute = size(trD_compute) / FragmentSize; | ||
| using Converter = cutlass::NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize>; | ||
| Converter converter{}; | ||
|
|
||
| CUTLASS_PRAGMA_UNROLL | ||
| for (int epi_v = 0; epi_v < num_fragments_trD_compute; ++epi_v) { | ||
| // Extract compute fragment | ||
| Array<ElementCompute, FragmentSize> trD_compute_frag; | ||
| CUTLASS_PRAGMA_UNROLL | ||
| for (int f = 0; f < FragmentSize; ++f) { | ||
| trD_compute_frag[f] = trD_compute(epi_v * FragmentSize + f); | ||
| } | ||
|
|
||
| // Convert fragment | ||
| auto trD_frag = converter(trD_compute_frag); | ||
|
|
||
| // Store converted fragment | ||
| CUTLASS_PRAGMA_UNROLL | ||
| for (int f = 0; f < FragmentSize; ++f) { | ||
| trD(epi_v * FragmentSize + f) = trD_frag[f]; | ||
|
|
||
| } | ||
| } | ||
|
|
||
| copy(copy_d, trD, tCgD); | ||
| } | ||
|
|
||
| cst_callbacks.end(); | ||
|
|
||
| } | ||
|
|
||
| private: | ||
| Params const& params; | ||
|
|
@@ -447,4 +434,4 @@ class CollectiveEpilogue< | |
| } // namespace epilogue | ||
| } // namespace cutlass | ||
|
|
||
| ///////////////////////////////////////////////////////////////////////////////////////////////// | ||
There was a problem hiding this comment.
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.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
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_* ?