| 1 | /*************************************************************************************************** |
| 2 | * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights |
| 3 | *reserved. SPDX-License-Identifier: BSD-3-Clause |
| 4 | * |
| 5 | * Redistribution and use in source and binary forms, with or without |
| 6 | * modification, are permitted provided that the following conditions are met: |
| 7 | * |
| 8 | * 1. Redistributions of source code must retain the above copyright notice, |
| 9 | *this list of conditions and the following disclaimer. |
| 10 | * |
| 11 | * 2. Redistributions in binary form must reproduce the above copyright notice, |
| 12 | * this list of conditions and the following disclaimer in the documentation |
| 13 | * and/or other materials provided with the distribution. |
| 14 | * |
| 15 | * 3. Neither the name of the copyright holder nor the names of its |
| 16 | * contributors may be used to endorse or promote products derived from |
| 17 | * this software without specific prior written permission. |
| 18 | * |
| 19 | * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" |
| 20 | * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE |
| 21 | * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE |
| 22 | *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE |
| 23 | *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR |
| 24 | *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF |
| 25 | *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS |
| 26 | *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN |
| 27 | *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) |
| 28 | *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE |
| 29 | *POSSIBILITY OF SUCH DAMAGE. |
| 30 | * |
| 31 | **************************************************************************************************/ |
| 32 | |
| 33 | // |
| 34 | // This file is a modified excerpt of |
| 35 | // include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp |
| 36 | // from https://github.com/NVIDIA/cutlass v3.5.0 |
| 37 | // It has been modified to support either row/column or scalar broadcasting |
| 38 | // where the tensor being loaded from is always passed in via a device pointer. |
| 39 | // This lets one compiled kernel handle all cases of per-tensor or |
| 40 | // per-channel/per-token quantization. |
| 41 | // |
| 42 | // This interface also allows the scales to be passed in as tensors that |
| 43 | // consistently reside on the device, which avoids an issue with a previous |
| 44 | // implementation where scalars needed to be on the CPU since they |
| 45 | // were passed in via float values. This created a potential performance hazard |
| 46 | // if scales were initially on the device, and caused torch.compile graphs |
| 47 | // breaks when moving scales to the CPU. |
| 48 | // |
| 49 | #pragma once |
| 50 | |
| 51 | // Turn off clang-format for the entire file to keep it close to upstream |
| 52 | // clang-format off |
| 53 | |
| 54 | #include "cutlass/cutlass.h" |
| 55 | #include "cutlass/arch/barrier.h" |
| 56 | |
| 57 | #include "cute/tensor.hpp" |
| 58 | #include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp" |
| 59 | |
| 60 | namespace cutlass::epilogue::fusion { |
| 61 | |
| 62 | using namespace cute; |
| 63 | using namespace detail; |
| 64 | |
| 65 | // Row vector broadcast |
| 66 | template< |
| 67 | int Stages, |
| 68 | class CtaTileShapeMNK, |
| 69 | class Element, |
| 70 | class StrideMNL = Stride<_0,_1,_0>, |
| 71 | int Alignment = 128 / sizeof_bits_v<Element> |
| 72 | > |
| 73 | struct Sm90RowOrScalarBroadcast { |
| 74 | static_assert(Stages == 0, "Row broadcast doesn't support smem usage"); |
| 75 | static_assert(is_static_v<decltype(take<0,2>(StrideMNL{}))>); // batch stride can be dynamic or static |
| 76 | static_assert(take<0,2>(StrideMNL{}) == Stride<_0,_1>{}); |
| 77 | |
| 78 | struct SharedStorage { |
| 79 | array_aligned<Element, size<1>(CtaTileShapeMNK{})> smem; |
| 80 | }; |
| 81 | |
| 82 | // This struct has been modified to have a bool indicating that ptr_row is a |
| 83 | // scalar that must be broadcast, instead of containing a scalar that is |
| 84 | // valid if ptr_row is null. |
| 85 | struct Arguments { |
| 86 | Element const* ptr_row = nullptr; |
| 87 | bool row_broadcast = true; |
| 88 | StrideMNL dRow = {}; |
| 89 | }; |
| 90 | |
| 91 | using Params = Arguments; |
| 92 | |
| 93 | template <class ProblemShape> |
| 94 | static constexpr Params |
| 95 | to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { |
| 96 | return args; |
| 97 | } |
| 98 | |
| 99 | template <class ProblemShape> |
| 100 | static bool |
| 101 | can_implement(ProblemShape const& problem_shape, Arguments const& args) { |
| 102 | return true; |
| 103 | } |
| 104 | |
| 105 | template <class ProblemShape> |
| 106 | static size_t |
| 107 | get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { |
| 108 | return 0; |
| 109 | } |
| 110 | |
| 111 | template <class ProblemShape> |
| 112 | static cutlass::Status |
| 113 | initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, |
| 114 | CudaHostAdapter* cuda_adapter = nullptr) { |
| 115 | return cutlass::Status::kSuccess; |
| 116 | } |
| 117 | |
| 118 | CUTLASS_HOST_DEVICE |
| 119 | Sm90RowOrScalarBroadcast() { } |
| 120 | |
| 121 | CUTLASS_HOST_DEVICE |
| 122 | Sm90RowOrScalarBroadcast(Params const& params, SharedStorage const& shared_storage) |
| 123 | : params(params) |
| 124 | , smem(const_cast<Element*>(shared_storage.smem.data())) { } |
| 125 | |
| 126 | Params params; |
| 127 | Element *smem = nullptr; |
| 128 | |
| 129 | CUTLASS_DEVICE bool |
| 130 | is_producer_load_needed() const { |
| 131 | return false; |
| 132 | } |
| 133 | |
| 134 | CUTLASS_DEVICE bool |
| 135 | is_C_load_needed() const { |
| 136 | return false; |
| 137 | } |
| 138 | |
| 139 | CUTLASS_DEVICE bool |
| 140 | is_zero() const { |
| 141 | return (!params.row_broadcast && *(params.ptr_row) == Element(0)); |
| 142 | } |
| 143 | |
| 144 | template <class... Args> |
| 145 | CUTLASS_DEVICE auto |
| 146 | get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) { |
| 147 | return EmptyProducerLoadCallbacks{}; |
| 148 | } |
| 149 | |
| 150 | template <class GS_GTensor, class GS_STensor, class GS_CTensor, class Tiled_G2S, class SR_STensor, class SR_RTensor, class CTensor, class ThrResidue, class ThrNum> |
| 151 | struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks { |
| 152 | CUTLASS_DEVICE |
| 153 | ConsumerStoreCallbacks( |
| 154 | GS_GTensor tGS_gRow_, GS_STensor tGS_sRow_, |
| 155 | GS_CTensor tGS_cRow_, Tiled_G2S tiled_g2s_, |
| 156 | SR_STensor tSR_sRow_, SR_RTensor tSR_rRow_, |
| 157 | CTensor tCcRow_, ThrResidue residue_tCcRow_, ThrNum thr_num_, Params const& params_) |
| 158 | : tGS_gRow(tGS_gRow_) |
| 159 | , tGS_sRow(tGS_sRow_) |
| 160 | , tGS_cRow(tGS_cRow_) |
| 161 | , tiled_G2S(tiled_g2s_) |
| 162 | , tSR_sRow(tSR_sRow_) |
| 163 | , tSR_rRow(tSR_rRow_) |
| 164 | , tCcRow(tCcRow_) |
| 165 | , residue_tCcRow(residue_tCcRow_) |
| 166 | , params(params_) {} |
| 167 | |
| 168 | GS_GTensor tGS_gRow; // (CPY,CPY_M,CPY_N) |
| 169 | GS_STensor tGS_sRow; // (CPY,CPY_M,CPY_N) |
| 170 | GS_CTensor tGS_cRow; // (CPY,CPY_M,CPY_N) |
| 171 | Tiled_G2S tiled_G2S; |
| 172 | |
| 173 | SR_STensor tSR_sRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) |
| 174 | SR_RTensor tSR_rRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) |
| 175 | |
| 176 | CTensor tCcRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) |
| 177 | ThrResidue residue_tCcRow; // (m, n) |
| 178 | ThrNum thr_num; |
| 179 | Params const& params; |
| 180 | |
| 181 | CUTLASS_DEVICE void |
| 182 | begin() { |
| 183 | if (!params.row_broadcast) { |
| 184 | fill(tSR_rRow, *(params.ptr_row)); |
| 185 | return; |
| 186 | } |
| 187 | |
| 188 | auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; |
| 189 | cute::Tensor tGS_gRow_flt = filter_zeros(tGS_gRow); |
| 190 | cute::Tensor tGS_sRow_flt = filter_zeros(tGS_sRow); |
| 191 | cute::Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride())); |
| 192 | |
| 193 | for (int i = 0; i < size(tGS_gRow_flt); ++i) { |
| 194 | if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) { |
| 195 | continue; // OOB of SMEM, |
| 196 | } |
| 197 | if (elem_less(tGS_cRow_flt(i), make_coord(get<0>(residue_tCcRow), get<1>(residue_tCcRow)))) { |
| 198 | tGS_sRow_flt(i) = tGS_gRow_flt(i); |
| 199 | } |
| 200 | else { |
| 201 | tGS_sRow_flt(i) = Element(0); // Set to Zero when OOB so LDS could be issue without any preds. |
| 202 | } |
| 203 | } |
| 204 | synchronize(); |
| 205 | } |
| 206 | |
| 207 | CUTLASS_DEVICE void |
| 208 | begin_loop(int epi_m, int epi_n) { |
| 209 | if (epi_m == 0) { // Assumes M-major subtile loop |
| 210 | if (!params.row_broadcast) return; // Do not issue LDS when row is scalar |
| 211 | cute::Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n)); |
| 212 | cute::Tensor tSR_rRow_flt = filter_zeros(tSR_rRow); |
| 213 | copy(tSR_sRow_flt, tSR_rRow_flt); |
| 214 | } |
| 215 | } |
| 216 | |
| 217 | template <typename ElementAccumulator, int FragmentSize> |
| 218 | CUTLASS_DEVICE Array<Element, FragmentSize> |
| 219 | visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) { |
| 220 | Array<Element, FragmentSize> frg_row; |
| 221 | |
| 222 | CUTLASS_PRAGMA_UNROLL |
| 223 | for (int i = 0; i < FragmentSize; ++i) { |
| 224 | frg_row[i] = tSR_rRow(epi_v * FragmentSize + i); |
| 225 | } |
| 226 | |
| 227 | return frg_row; |
| 228 | } |
| 229 | }; |
| 230 | |
| 231 | template < |
| 232 | bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy |
| 233 | class... Args |
| 234 | > |
| 235 | CUTLASS_DEVICE auto |
| 236 | get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) { |
| 237 | auto [M, N, K, L] = args.problem_shape_mnkl; |
| 238 | auto [m, n, k, l] = args.tile_coord_mnkl; |
| 239 | using ThreadCount = decltype(size(args.tiled_copy)); |
| 240 | |
| 241 | cute::Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow); |
| 242 | cute::Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N) |
| 243 | cute::Tensor sRow = make_tensor(make_smem_ptr(smem), |
| 244 | make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N) |
| 245 | //// G2S: Gmem to Smem |
| 246 | auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{}, |
| 247 | Layout< Shape<_1, ThreadCount>, |
| 248 | Stride<_0, _1>>{}, |
| 249 | Layout<_1>{}); |
| 250 | auto thr_g2s = tiled_g2s.get_slice(args.thread_idx); |
| 251 | cute::Tensor tGS_gRow = thr_g2s.partition_S(gRow); |
| 252 | cute::Tensor tGS_sRow = thr_g2s.partition_D(sRow); |
| 253 | |
| 254 | //// G2S: Coord |
| 255 | auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{}))); |
| 256 | cute::Tensor tGS_cRow = thr_g2s.partition_S(cRow); |
| 257 | |
| 258 | //// S2R: Smem to Reg |
| 259 | cute::Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx); |
| 260 | cute::Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N) |
| 261 | |
| 262 | return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>( |
| 263 | tGS_gRow, |
| 264 | tGS_sRow, |
| 265 | tGS_cRow, tiled_g2s, |
| 266 | tSR_sRow, |
| 267 | tSR_rRow, |
| 268 | args.tCcD, |
| 269 | args.residue_cD, |
| 270 | ThreadCount{}, |
| 271 | params); |
| 272 | } |
| 273 | }; |
| 274 | |
| 275 | ///////////////////////////////////////////////////////////////////////////////////////////////// |
| 276 | |
| 277 | // Column vector broadcast |
| 278 | template< |
| 279 | int Stages, |
| 280 | class CtaTileShapeMNK, |
| 281 | class Element, |
| 282 | class StrideMNL = Stride<_1,_0,_0>, |
| 283 | int Alignment = 128 / sizeof_bits_v<Element> |
| 284 | > |
| 285 | struct Sm90ColOrScalarBroadcast { |
| 286 | static_assert(Stages == 0, "Column broadcast doesn't support smem usage yet"); |
| 287 | static_assert(Alignment * sizeof_bits_v<Element> % 128 == 0, "sub-16B alignment not supported yet"); |
| 288 | static_assert( |
| 289 | (cute::is_same_v<StrideMNL, Stride<_1,_0, _0>>) || // col vector broadcast, e.g. per-row alpha/bias |
| 290 | (cute::is_same_v<StrideMNL, Stride<_1,_0,int>>)); // batched col vector broadcast, e.g. batched per-row bias |
| 291 | |
| 292 | // Accumulator distributes col elements evenly amongst threads so we can just directly load from gmem |
| 293 | struct SharedStorage { }; |
| 294 | |
| 295 | // This struct has been modified to have a bool indicating that ptr_col is a |
| 296 | // scalar that must be broadcast, instead of containing a scalar that is |
| 297 | // valid if ptr_col is null. |
| 298 | struct Arguments { |
| 299 | Element const* ptr_col = nullptr; |
| 300 | bool col_broadcast = true; |
| 301 | StrideMNL dCol = {}; |
| 302 | }; |
| 303 | |
| 304 | using Params = Arguments; |
| 305 | |
| 306 | template <class ProblemShape> |
| 307 | static constexpr Params |
| 308 | to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { |
| 309 | return args; |
| 310 | } |
| 311 | |
| 312 | template <class ProblemShape> |
| 313 | static bool |
| 314 | can_implement(ProblemShape const& problem_shape, Arguments const& args) { |
| 315 | return true; |
| 316 | } |
| 317 | |
| 318 | template <class ProblemShape> |
| 319 | static size_t |
| 320 | get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { |
| 321 | return 0; |
| 322 | } |
| 323 | |
| 324 | template <class ProblemShape> |
| 325 | static cutlass::Status |
| 326 | initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, |
| 327 | CudaHostAdapter* cuda_adapter = nullptr) { |
| 328 | return cutlass::Status::kSuccess; |
| 329 | } |
| 330 | |
| 331 | CUTLASS_DEVICE bool |
| 332 | is_producer_load_needed() const { |
| 333 | return false; |
| 334 | } |
| 335 | |
| 336 | CUTLASS_DEVICE bool |
| 337 | is_C_load_needed() const { |
| 338 | return false; |
| 339 | } |
| 340 | |
| 341 | CUTLASS_DEVICE bool |
| 342 | is_zero() const { |
| 343 | return (!params.col_broadcast && *(params.ptr_col) == Element(0)); |
| 344 | } |
| 345 | |
| 346 | CUTLASS_HOST_DEVICE |
| 347 | Sm90ColOrScalarBroadcast() { } |
| 348 | |
| 349 | CUTLASS_HOST_DEVICE |
| 350 | Sm90ColOrScalarBroadcast(Params const& params, SharedStorage const& shared_storage) |
| 351 | : params(params) { } |
| 352 | |
| 353 | Params params; |
| 354 | |
| 355 | template <class... Args> |
| 356 | CUTLASS_DEVICE auto |
| 357 | get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) { |
| 358 | return EmptyProducerLoadCallbacks{}; |
| 359 | } |
| 360 | |
| 361 | template<class GTensor, class RTensor, class CTensor, class ProblemShape> |
| 362 | struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks { |
| 363 | CUTLASS_DEVICE |
| 364 | ConsumerStoreCallbacks( |
| 365 | GTensor&& tCgCol, |
| 366 | RTensor&& tCrCol, |
| 367 | CTensor&& tCcCol, |
| 368 | ProblemShape problem_shape, |
| 369 | Params const& params |
| 370 | ): |
| 371 | tCgCol(cute::forward<GTensor>(tCgCol)), |
| 372 | tCrCol(cute::forward<RTensor>(tCrCol)), |
| 373 | tCcCol(cute::forward<CTensor>(tCcCol)), |
| 374 | m(get<0>(problem_shape)), |
| 375 | params(params) {} |
| 376 | |
| 377 | GTensor tCgCol; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) |
| 378 | RTensor tCrCol; |
| 379 | CTensor tCcCol; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) |
| 380 | Params const& params; |
| 381 | int m; |
| 382 | |
| 383 | CUTLASS_DEVICE void |
| 384 | begin() { |
| 385 | if (!params.col_broadcast) { |
| 386 | fill(tCrCol, *(params.ptr_col)); |
| 387 | return; |
| 388 | } |
| 389 | |
| 390 | // tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and |
| 391 | // EPI_N are stride-0 for the column broadcast. Slice those modes at |
| 392 | // index 0 to avoid redundant copies AND ensure pred/data consistency |
| 393 | static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast"); |
| 394 | static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast"); |
| 395 | |
| 396 | auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) |
| 397 | auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) |
| 398 | auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) |
| 399 | |
| 400 | cute::Tensor pred = make_tensor<bool>(shape(tCgCol_s)); |
| 401 | CUTLASS_PRAGMA_UNROLL |
| 402 | for (int i = 0; i < size(pred); ++i) { |
| 403 | pred(i) = get<0>(tCcCol_s(i)) < m; |
| 404 | } |
| 405 | |
| 406 | copy_if(pred, tCgCol_s, tCrCol_s); |
| 407 | } |
| 408 | |
| 409 | template <typename ElementAccumulator, int FragmentSize> |
| 410 | CUTLASS_DEVICE Array<Element, FragmentSize> |
| 411 | visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) { |
| 412 | Array<Element, FragmentSize> frg_col; |
| 413 | cute::Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n); |
| 414 | |
| 415 | CUTLASS_PRAGMA_UNROLL |
| 416 | for (int i = 0; i < FragmentSize; ++i) { |
| 417 | frg_col[i] = tCrCol_mn(epi_v * FragmentSize + i); |
| 418 | } |
| 419 | |
| 420 | return frg_col; |
| 421 | } |
| 422 | |
| 423 | }; |
| 424 | |
| 425 | template < |
| 426 | bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy |
| 427 | class... Args |
| 428 | > |
| 429 | CUTLASS_DEVICE auto |
| 430 | get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) { |
| 431 | |
| 432 | auto [M, N, K, L] = args.problem_shape_mnkl; |
| 433 | cute::Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol); |
| 434 | cute::Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) |
| 435 | mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); |
| 436 | cute::Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) |
| 437 | |
| 438 | // Generate an identity tensor matching the shape of the global tensor and |
| 439 | // partition the same way, this will be used to generate the predicate |
| 440 | // tensor for loading |
| 441 | cute::Tensor cCol = make_identity_tensor(mCol.shape()); |
| 442 | cute::Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N) |
| 443 | cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx); |
| 444 | |
| 445 | return ConsumerStoreCallbacks( |
| 446 | cute::move(tCgCol), |
| 447 | cute::move(tCrCol), |
| 448 | cute::move(tCcCol), |
| 449 | args.problem_shape_mnkl, |
| 450 | params |
| 451 | ); |
| 452 | } |
| 453 | }; |
| 454 | |
| 455 | } |