/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_bwd_weight.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_bwd_weight.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_bwd_weight.hpp Source File
gridwise_gemm_xdlops_bwd_weight.hpp
Go to the documentation of this file.
1 // SPDX-License-Identifier: MIT
2 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
3 
4 #pragma once
5 
17 
18 namespace ck {
19 
20 // Implementation of "Merge" transformation primitive that uses division and mod. It is supposed to
21 // be used for low_lengths that are known at compile time and are power of 2, otherwise performance
22 // will be very bad
23 template <typename LowLengths>
25 {
26  static constexpr index_t NDimLow = LowLengths::Size();
27 
30 
32  decltype(container_reverse_exclusive_scan(LowLengths{}, math::multiplies{}, Number<1>{}));
33 
34  using UpLengths =
35  decltype(make_tuple(container_reduce(LowLengths{}, math::multiplies{}, Number<1>{})));
36 
37  LowLengths low_lengths_;
40 
41  __host__ __device__ constexpr Merge_v4_no_carry() = default;
42 
43  __host__ __device__ constexpr Merge_v4_no_carry(const LowLengths& low_lengths)
44  : low_lengths_{low_lengths},
46  container_reverse_exclusive_scan(low_lengths, math::multiplies{}, Number<1>{})},
47  up_lengths_{make_tuple(container_reduce(low_lengths, math::multiplies{}, Number<1>{}))}
48  {
49  static_assert(LowerIndex::Size() == NDimLow, "wrong!");
50  }
51 
52  __host__ __device__ static constexpr index_t GetNumOfLowerDimension() { return NDimLow; }
53 
54  __host__ __device__ static constexpr index_t GetNumOfUpperDimension() { return 1; }
55 
56  __host__ __device__ constexpr const auto& GetUpperLengths() const { return up_lengths_; }
57 
58  template <typename LowIdx, typename UpIdx>
59  __host__ __device__ constexpr void CalculateLowerIndex(LowIdx& idx_low,
60  const UpIdx& idx_up) const
61  {
62  static_assert(LowIdx::Size() == NDimLow && UpIdx::Size() == 1,
63  "wrong! inconsistent # of dimension");
64 
65  index_t tmp = idx_up[Number<0>{}];
66 
67  // division and mod
68  static_for<0, NDimLow - 1, 1>{}([&](auto i) {
69  idx_low(i) = tmp / this->low_lengths_scan_[i];
70  tmp %= this->low_lengths_scan_[i];
71  });
72 
73  idx_low(Number<NDimLow - 1>{}) = tmp;
74  }
75 
76  template <typename LowIdxDiff,
77  typename UpIdxDiff,
78  typename LowIdx,
79  typename UpIdx,
80  index_t Hack>
81  __host__ __device__ void UpdateLowerIndex(LowIdxDiff& idx_diff_low,
82  const UpIdxDiff& idx_up_diff,
83  LowIdx& idx_low,
84  const UpIdx& idx_up_new,
85  Number<Hack>) const
86  {
87  static_assert(LowIdxDiff::Size() == NDimLow && UpIdxDiff::Size() == 1 &&
88  LowIdx::Size() == NDimLow && UpIdx::Size() == 1,
89  "wrong! inconsistent # of dimension");
90 
91  constexpr auto I0 = Number<0>{};
92  constexpr auto INm1 = Number<NDimLow - 1>{};
93 
94  index_t tmp = idx_up_new[I0];
95 
96  idx_low(INm1) = tmp;
97  idx_diff_low(INm1) = idx_up_diff[I0];
98  }
99 
100  __host__ __device__ static constexpr bool IsLinearTransform() { return false; }
101 
102  __host__ __device__ static constexpr bool IsValidUpperIndexAlwaysMappedToValidLowerIndex()
103  {
104  return true;
105  }
106 
107  __host__ __device__ static constexpr bool IsKnownAtCompileTime()
108  {
112  }
113 
114  template <typename UpIdx>
115  __host__ __device__ static constexpr bool
116  IsValidUpperIndexMappedToValidLowerIndex(const UpIdx& /* idx_up */)
117  {
118  return true;
119  }
120 
121  __host__ __device__ void Print() const
122  {
123  printf("{");
124  printf("Merge_v3_direct_division_mod_wrw, ");
125  printf("low_lengths_ ");
126  print_multi_index(low_lengths_);
127  printf("low_lengths_scan_ ");
128  print_multi_index(low_lengths_scan_);
129  printf("up_lengths_ ");
130  print_multi_index(up_lengths_);
131  printf("}");
132  }
133 };
134 
135 template <typename LowLengths>
136 __host__ __device__ constexpr auto make_merge_transform_v4_no_carry(const LowLengths& low_lengths)
137 {
138  return Merge_v4_no_carry<LowLengths>{low_lengths};
139 }
140 
141 template <typename GridwiseGemm,
142  typename FloatA,
143  typename FloatB,
144  typename FloatC,
145  typename AGridDesc_B_K0_M_K1,
146  typename BGridDesc_B_K0_N_K1,
147  typename CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock,
148  typename AElementwiseOperation,
149  typename BElementwiseOperation,
150  typename CElementwiseOperation,
151  typename CBlockClusterAdaptor,
152  bool HasMainKBlockLoop>
153 __global__ void
154 #if CK_USE_LAUNCH_BOUNDS
156 #endif
157  kernel_gemm_xdlops_bwd_weight(const FloatA* __restrict__ p_a_grid,
158  const FloatB* __restrict__ p_b_grid,
159  FloatC* __restrict__ p_c_grid,
160  const AGridDesc_B_K0_M_K1 a_b_k0_m_k1_grid_desc,
161  const BGridDesc_B_K0_N_K1 b_b_k0_n_k1_grid_desc,
162  const CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock
163  c_grid_desc_mblock_mperblock_nblock_nperblock,
164  const AElementwiseOperation a_element_op,
165  const BElementwiseOperation b_element_op,
166  const CElementwiseOperation c_element_op,
167  const CBlockClusterAdaptor c_block_cluster_adaptor)
168 {
169 #if defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx94__) || defined(__gfx11__) || \
170  defined(__gfx12__)
171  if constexpr(GridwiseGemm::template IsValidCompilationParameter<>())
172  {
173  __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
174 
175  GridwiseGemm::template Run<HasMainKBlockLoop>(p_a_grid,
176  p_b_grid,
177  p_c_grid,
178  p_shared,
179  a_b_k0_m_k1_grid_desc,
180  b_b_k0_n_k1_grid_desc,
181  c_grid_desc_mblock_mperblock_nblock_nperblock,
182  a_element_op,
183  b_element_op,
184  c_element_op,
185  c_block_cluster_adaptor);
186  }
187 #else
188  ignore = p_a_grid;
189  ignore = p_b_grid;
190  ignore = p_c_grid;
191  ignore = a_b_k0_m_k1_grid_desc;
192  ignore = b_b_k0_n_k1_grid_desc;
193  ignore = c_grid_desc_mblock_mperblock_nblock_nperblock;
194  ignore = a_element_op;
195  ignore = b_element_op;
196  ignore = c_element_op;
197  ignore = c_block_cluster_adaptor;
198 #endif // end of if (defined(__gfx908__) || defined(__gfx90a__))
199 }
200 
201 template <index_t BlockSize,
202  typename FloatA,
203  typename FloatB,
204  typename FloatAcc,
205  typename FloatC,
206  InMemoryDataOperationEnum CGlobalMemoryDataOperation,
207  typename AGridDesc_B_K0_M_K1,
208  typename BGridDesc_B_K0_N_K1,
209  typename CMNGridDesc,
210  typename AElementwiseOperation,
211  typename BElementwiseOperation,
212  typename CElementwiseOperation,
213  index_t MPerBlock,
214  index_t NPerBlock,
215  index_t K0PerBlock,
216  index_t MPerXdl,
217  index_t NPerXdl,
218  index_t K1Value,
219  index_t MRepeat,
220  index_t NRepeat,
221  typename ABlockTransferThreadClusterLengths_K0_M_K1,
222  typename ABlockTransferThreadClusterArrangeOrder,
223  typename ABlockTransferSrcAccessOrder,
224  index_t ABlockTransferSrcVectorDim,
225  index_t ABlockTransferSrcScalarPerVector,
226  index_t ABlockTransferDstScalarPerVector_K1,
227  bool AThreadTransferSrcResetCoordinateAfterRun,
228  bool ABlockLdsExtraM,
229  index_t ABlockLdsM1PerBlock,
230  index_t ABlockLdsM0PerBlock,
231  index_t ABlockLdsM1Padding,
232  typename BBlockTransferThreadClusterLengths_K0_N_K1,
233  typename BBlockTransferThreadClusterArrangeOrder,
234  typename BBlockTransferSrcAccessOrder,
235  index_t BBlockTransferSrcVectorDim,
236  index_t BBlockTransferSrcScalarPerVector,
237  index_t BBlockTransferDstScalarPerVector_K1,
238  bool BThreadTransferSrcResetCoordinateAfterRun,
239  bool BBlockLdsExtraN,
240  index_t BBlockLdsN1PerBlock,
241  index_t BBlockLdsN0PerBlock,
242  index_t BBlockLdsN1Padding,
243  index_t CShuffleMRepeatPerShuffle,
244  index_t CShuffleNRepeatPerShuffle,
245  index_t CBlockTransferScalarPerVector_NWaveNPerXDL,
246  typename CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
247  bool ABlockLdsExtraM1Wrw = false,
248  bool BBlockLdsExtraN1Wrw = false,
249  index_t NumGemmKPrefetchStage = 1,
250  PipelineVersion PipelineVer = PipelineVersion::v1,
251  typename ComputeTypeA = FloatA,
252  typename ComputeTypeB = ComputeTypeA>
254 {
255  static constexpr auto I0 = Number<0>{};
256  static constexpr auto I1 = Number<1>{};
257  static constexpr auto I2 = Number<2>{};
258  static constexpr auto I3 = Number<3>{};
259  static constexpr auto I4 = Number<4>{};
260  static constexpr auto I5 = Number<5>{};
261  static constexpr auto I6 = Number<6>{};
262  static constexpr auto I7 = Number<7>{};
263 
264  // K1 should be Number<...>
265  static constexpr auto K1 = Number<K1Value>{};
266 
268 
270  decltype(GridwiseGemmPipeline_Selector<PipelineVer, NumGemmKPrefetchStage>())>;
271 
272  // denorm test fix, required to work around fp16 mfma issue
273  // we convert fp16->fp32->bf16 and execute bf16 mfma instruction
274  // when mfma if fixed, remove this section and update
275  // FloatAAdjusted -> ComputeTypeA, FloatBAdjusted -> ComputeTypeB,
276  // throughout this file
277 #if CK_GFX90A_DENORM_WORKAROUND
278  using FloatAAdjusted =
280  using FloatBAdjusted =
282 #else
283  using FloatAAdjusted = ComputeTypeA;
284  using FloatBAdjusted = ComputeTypeB;
285 #endif
286 
287  // M0/M1/M1Padding
288  static constexpr auto M1PerBlock = Number<ABlockLdsM1PerBlock>{};
289  static constexpr auto M0PerBlock = Number<ABlockLdsM0PerBlock>{};
290  static constexpr auto M1Padding = Number<ABlockLdsM1Padding>{};
291 
292  // N0/N1/N1Padding
293  static constexpr auto N1PerBlock = Number<BBlockLdsN1PerBlock>{};
294  static constexpr auto N0PerBlock = Number<BBlockLdsN0PerBlock>{};
295  static constexpr auto N1Padding = Number<BBlockLdsN1Padding>{};
296 
297  __host__ __device__ static constexpr auto GetABlockDescriptor_K0PerBlock_MPerBlock_K1()
298  {
299  constexpr auto max_lds_align = K1;
300 
301  // A matrix in LDS memory, dst of blockwise copy
302  constexpr auto a_block_desc_k0_m_k1 = [&]() {
303  if constexpr(ABlockLdsExtraM)
304  {
305  if constexpr(ABlockLdsExtraM1Wrw)
306  {
307  constexpr auto a_block_desc_k0_m0_m1_k1 = make_naive_tensor_descriptor(
308  make_tuple(
310  make_tuple(Number<M0PerBlock>{} * (Number<M1PerBlock>{} * K1 + M1Padding),
311  Number<M1PerBlock>{} * K1 + M1Padding,
312  K1,
313  I1));
314 
315  constexpr auto a_block_desc_k0_m_k1_tmp = transform_tensor_descriptor(
316  a_block_desc_k0_m0_m1_k1,
323 
324  return a_block_desc_k0_m_k1_tmp;
325  }
326  else
327  {
330  make_tuple(Number<MPerBlock + 1>{} * K1, K1, I1));
331  }
332  }
333  else
334  {
336  make_tuple(Number<K0PerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
337  }
338  }();
339 
340  return a_block_desc_k0_m_k1;
341  }
342 
343  __host__ __device__ static constexpr auto GetABlockDescriptor_Batch_K0PerBlock_MPerBlock_K1()
344  {
345  constexpr auto max_lds_align = K1;
346 
347  // A matrix in LDS memory, dst of blockwise copy
348  constexpr auto a_block_desc_b_k0_m_k1 = [&]() {
349  if constexpr(ABlockLdsExtraM)
350  {
351  if constexpr(ABlockLdsExtraM1Wrw)
352  {
353  constexpr auto a_block_desc_b_k0_m0_m1_k1 = make_naive_tensor_descriptor(
358  K1),
360  (Number<M1PerBlock>{} * K1 + M1Padding),
361  Number<M0PerBlock>{} * (Number<M1PerBlock>{} * K1 + M1Padding),
362  Number<M1PerBlock>{} * K1 + M1Padding,
363  K1,
364  I1));
365 
366  constexpr auto a_block_desc_b_k0_m_k1_tmp = transform_tensor_descriptor(
367  a_block_desc_b_k0_m0_m1_k1,
375 
376  return a_block_desc_b_k0_m_k1_tmp;
377  }
378  else
379  {
383  Number<MPerBlock + 1>{} * K1,
384  K1,
385  I1));
386  }
387  }
388  else
389  {
392  max_lds_align);
393  }
394  }();
395 
396  return a_block_desc_b_k0_m_k1;
397  }
398 
399  __host__ __device__ static constexpr auto GetBBlockDescriptor_K0PerBlock_NPerBlock_K1()
400  {
401  constexpr auto max_lds_align = K1;
402 
403  // B matrix in LDS memory, dst of blockwise copy
404  constexpr auto b_block_desc_k0_n_k1 = [&]() {
405  if constexpr(BBlockLdsExtraN)
406  {
407  if constexpr(BBlockLdsExtraN1Wrw)
408  {
409  constexpr auto b_block_desc_k0_n0_n1_k1 = make_naive_tensor_descriptor(
410  make_tuple(
412  make_tuple(Number<N0PerBlock>{} * (Number<N1PerBlock>{} * K1 + N1Padding),
413  Number<N1PerBlock>{} * K1 + N1Padding,
414  K1,
415  I1));
416 
417  constexpr auto b_block_desc_k0_n_k1_tmp = transform_tensor_descriptor(
418  b_block_desc_k0_n0_n1_k1,
425 
426  return b_block_desc_k0_n_k1_tmp;
427  }
428  else
429  {
432  make_tuple(Number<NPerBlock + 1>{} * K1, K1, I1));
433  }
434  }
435  else
436  {
438  make_tuple(Number<K0PerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
439  }
440  }();
441 
442  return b_block_desc_k0_n_k1;
443  }
444 
445  __host__ __device__ static constexpr auto GetBBlockDescriptor_Batch_K0PerBlock_NPerBlock_K1()
446  {
447  constexpr auto max_lds_align = K1;
448 
449  // B matrix in LDS memory, dst of blockwise copy
450  constexpr auto b_block_desc_b_k0_n_k1 = [&]() {
451  if constexpr(BBlockLdsExtraN)
452  {
453  if constexpr(BBlockLdsExtraN1Wrw)
454  {
455  constexpr auto b_block_desc_b_k0_n0_n1_k1 = make_naive_tensor_descriptor(
460  K1),
462  (Number<N1PerBlock>{} * K1 + N1Padding),
463  Number<N0PerBlock>{} * (Number<N1PerBlock>{} * K1 + N1Padding),
464  Number<N1PerBlock>{} * K1 + N1Padding,
465  K1,
466  I1));
467 
468  constexpr auto b_block_desc_b_k0_n_k1_tmp = transform_tensor_descriptor(
469  b_block_desc_b_k0_n0_n1_k1,
477 
478  return b_block_desc_b_k0_n_k1_tmp;
479  }
480  else
481  {
485  Number<NPerBlock + 1>{} * K1,
486  K1,
487  I1));
488  }
489  }
490  else
491  {
494  max_lds_align);
495  }
496  }();
497 
498  return b_block_desc_b_k0_n_k1;
499  }
500 
501  __host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte()
502  {
503  constexpr auto max_lds_align = K1;
504 
505  // A matrix in LDS memory, dst of blockwise copy
506  constexpr auto a_b_k0_m_k1_block_desc = GetABlockDescriptor_Batch_K0PerBlock_MPerBlock_K1();
507 
508  // B matrix in LDS memory, dst of blockwise copy
509  constexpr auto b_b_k0_n_k1_block_desc = GetBBlockDescriptor_Batch_K0PerBlock_NPerBlock_K1();
510 
511  // LDS allocation for A and B: be careful of alignment
512  constexpr auto a_block_space_size = math::integer_least_multiple(
513  a_b_k0_m_k1_block_desc.GetElementSpaceSize(), max_lds_align);
514 
515  constexpr auto b_block_space_size = math::integer_least_multiple(
516  b_b_k0_n_k1_block_desc.GetElementSpaceSize(), max_lds_align);
517 
518  constexpr auto c_block_size =
519  GetCBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock().GetElementSpaceSize();
520 
521  return math::max((a_block_space_size * sizeof(FloatAAdjusted) +
522  b_block_space_size * sizeof(FloatBAdjusted)),
523  c_block_size * sizeof(FloatC));
524  }
525 
526  template <
527  InMemoryDataOperationEnum CGlobalMemoryDataOperation_ = InMemoryDataOperationEnum::Set>
528  __device__ static bool constexpr IsValidCompilationParameter()
529  {
530  return ck::tensor_operation::device::IsValidGemmCompilationParameter<
531  BlockSize,
532  MPerBlock,
533  NPerBlock,
534  MPerXdl,
535  NPerXdl,
536  MRepeat,
537  NRepeat,
538  FloatC,
539  CGlobalMemoryDataOperation>();
540  }
541  // block_id to matrix tile idx (m0, n0) mapping are controlled by {M01, N01}
542  template <typename Block2CTileMap>
543  __host__ __device__ static constexpr bool
544  CheckValidity(const AGridDesc_B_K0_M_K1& a_b_k0_m_k1_grid_desc,
545  const BGridDesc_B_K0_N_K1& b_b_k0_n_k1_grid_desc,
546  const CMNGridDesc& c_m_n_grid_desc,
547  const Block2CTileMap& block_2_ctile_map)
548  {
549  static_assert(is_known_at_compile_time<remove_cv_t<decltype(K1)>>::value,
550  "wrong! K1 need to be known at compile-time");
551 
552  static_assert((MPerBlock % (MPerXdl * MRepeat) == 0) &&
553  (NPerBlock % (NRepeat * NPerXdl)) == 0,
554  "Invalid tuning param!");
555 
556  const auto M = a_b_k0_m_k1_grid_desc.GetLength(I2);
557  const auto N = b_b_k0_n_k1_grid_desc.GetLength(I2);
558  const auto K0 = a_b_k0_m_k1_grid_desc.GetLength(I1);
559  const auto KBatch = a_b_k0_m_k1_grid_desc.GetLength(I0);
560 
561  // check gridwise gemm pipeline
562  const auto num_k_loop = K0 / K0PerBlock;
563 
564  if(!GridwiseGemmPipe::IsSupported(num_k_loop))
565  {
566  return false;
567  }
568 
569  if(!(M == c_m_n_grid_desc.GetLength(I0) && N == c_m_n_grid_desc.GetLength(I1) &&
570  K0 == b_b_k0_n_k1_grid_desc.GetLength(I1) &&
571  K1 == a_b_k0_m_k1_grid_desc.GetLength(I3) &&
572  K1 == b_b_k0_n_k1_grid_desc.GetLength(I3) &&
573  KBatch == b_b_k0_n_k1_grid_desc.GetLength(I0)))
574  return false;
575 
576  if(!(M % MPerBlock == 0 && N % NPerBlock == 0 && K0 % K0PerBlock == 0))
577  return false;
578 
579  if(!block_2_ctile_map.CheckValidity(c_m_n_grid_desc))
580  {
581  return false;
582  }
583 
584  // TODO: also check validity of all components (blockwise-copy, threadwise-copy, etc)
585  return true;
586  }
587 
588  __host__ __device__ static constexpr bool CalculateHasMainK0BlockLoop(index_t K0)
589  {
590  // const bool has_main_k0_block_loop = K0 > K0PerBlock;
591  const index_t num_loop = K0 / K0PerBlock;
592 
593  return GridwiseGemmPipe::CalculateHasMainLoop(num_loop);
594 
595  // return has_main_k0_block_loop;
596  }
597 
598  __host__ __device__ static constexpr auto
599  MakeCGridDesc_MBlock_MPerBlock_NBlock_NPerBlock(const CMNGridDesc& c_m_n_grid_desc)
600  {
601  const auto M = c_m_n_grid_desc.GetLength(I0);
602  const auto N = c_m_n_grid_desc.GetLength(I1);
603 
604  const auto MBlock = M / MPerBlock;
605  const auto NBlock = N / NPerBlock;
606 
608  c_m_n_grid_desc,
613  }
614 
615  // return block_id to C matrix tile idx (m0, n0) mapping
616  __host__ __device__ static constexpr auto MakeCBlockClusterAdaptor(
617  const CMNGridDesc& c_m_n_grid_desc, index_t M01, index_t N01, index_t KBatch)
618  {
620  c_m_n_grid_desc, M01, N01, KBatch);
621  }
622 
623  __host__ __device__ static constexpr auto
625  {
626  constexpr index_t MWave = MPerBlock / (MRepeat * MPerXdl);
627  constexpr index_t NWave = NPerBlock / (NRepeat * NPerXdl);
628 
630  make_tuple(I1,
632  I1,
634  }
635 
637  decltype(MakeCGridDesc_MBlock_MPerBlock_NBlock_NPerBlock(CMNGridDesc{}));
638  using CBlockClusterAdaptor = decltype(MakeCBlockClusterAdaptor(CMNGridDesc{}, 1, 1, 1));
639 
640  template <bool HasMainKBlockLoop>
641  __device__ static void Run(const FloatA* __restrict__ p_a_grid,
642  const FloatB* __restrict__ p_b_grid,
643  FloatC* __restrict__ p_c_grid,
644  void* __restrict__ p_shared,
645  const AGridDesc_B_K0_M_K1& a_b_k0_m_k1_grid_desc,
646  const BGridDesc_B_K0_N_K1& b_b_k0_n_k1_grid_desc,
648  c_grid_desc_mblock_mperblock_nblock_nperblock,
649  const AElementwiseOperation& a_element_op,
650  const BElementwiseOperation& b_element_op,
651  const CElementwiseOperation& c_element_op,
652  const CBlockClusterAdaptor& c_block_cluster_adaptor)
653  {
654  const auto a_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
655  p_a_grid, a_b_k0_m_k1_grid_desc.GetElementSpaceSize());
656  const auto b_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
657  p_b_grid, b_b_k0_n_k1_grid_desc.GetElementSpaceSize());
658  auto c_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
659  p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());
660 
661  const auto K0 = a_b_k0_m_k1_grid_desc.GetLength(I1);
662 
663  // divide block work by [M, N]
664  const auto block_work_idx =
665  c_block_cluster_adaptor.CalculateBottomIndex(make_multi_index(get_block_1d_id()));
666 
667  const index_t k_batch_id = block_work_idx[I0];
668 
669  if(!c_block_cluster_adaptor.ValidCTileIndex(
670  make_tuple(block_work_idx[I1], block_work_idx[I2]),
671  make_tuple(c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I0),
672  c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I2))))
673  {
674  return;
675  }
676 
677  // HACK: this force m/n_block_data_idx_on_grid into SGPR
678  const index_t m_block_data_idx_on_grid =
679  __builtin_amdgcn_readfirstlane(block_work_idx[I1] * MPerBlock);
680 
681  const index_t n_block_data_idx_on_grid =
682  __builtin_amdgcn_readfirstlane(block_work_idx[I2] * NPerBlock);
683 
684  // lds max alignment
685  constexpr auto max_lds_align = K1;
686 
687  // A matrix in LDS memory, dst of blockwise copy
688  constexpr auto a_k0_m_k1_block_desc = GetABlockDescriptor_K0PerBlock_MPerBlock_K1();
689 
690  constexpr auto a_b_k0_m_k1_block_desc = GetABlockDescriptor_Batch_K0PerBlock_MPerBlock_K1();
691  // B matrix in LDS memory, dst of blockwise copy
692  constexpr auto b_k0_n_k1_block_desc = GetBBlockDescriptor_K0PerBlock_NPerBlock_K1();
693 
694  constexpr auto b_b_k0_n_k1_block_desc = GetBBlockDescriptor_Batch_K0PerBlock_NPerBlock_K1();
695  // A matrix blockwise copy
696  auto a_blockwise_copy =
698  AElementwiseOperation,
700  InMemoryDataOperationEnum::Set,
702  ABlockTransferThreadClusterLengths_K0_M_K1,
703  ABlockTransferThreadClusterArrangeOrder,
704  FloatA,
706  decltype(a_b_k0_m_k1_grid_desc),
707  decltype(a_b_k0_m_k1_block_desc),
708  ABlockTransferSrcAccessOrder,
710  ABlockTransferSrcVectorDim,
711  3,
712  ABlockTransferSrcScalarPerVector,
713  ABlockTransferDstScalarPerVector_K1,
714  1,
715  1,
716  AThreadTransferSrcResetCoordinateAfterRun,
717  true>(
718  a_b_k0_m_k1_grid_desc,
719  make_multi_index(k_batch_id, 0, m_block_data_idx_on_grid, 0),
720  a_element_op,
721  a_b_k0_m_k1_block_desc,
722  make_multi_index(0, 0, 0, 0),
724 
725  // B matrix blockwise copy
726  auto b_blockwise_copy =
728  BElementwiseOperation,
730  InMemoryDataOperationEnum::Set,
732  BBlockTransferThreadClusterLengths_K0_N_K1,
733  BBlockTransferThreadClusterArrangeOrder,
734  FloatB,
736  decltype(b_b_k0_n_k1_grid_desc),
737  decltype(b_b_k0_n_k1_block_desc),
738  BBlockTransferSrcAccessOrder,
740  BBlockTransferSrcVectorDim,
741  3,
742  BBlockTransferSrcScalarPerVector,
743  BBlockTransferDstScalarPerVector_K1,
744  1,
745  1,
746  BThreadTransferSrcResetCoordinateAfterRun,
747  true>(
748  b_b_k0_n_k1_grid_desc,
749  make_multi_index(k_batch_id, 0, n_block_data_idx_on_grid, 0),
750  b_element_op,
751  b_b_k0_n_k1_block_desc,
752  make_multi_index(0, 0, 0, 0),
754 
755  // GEMM definition
756  // c_mtx += transpose(a_mtx) * b_mtx
757  // a_mtx[K0PerBlock, MPerBlock] is in LDS
758  // b_mtx[K0PerBlock, NPerBlock] is in LDS
759  // c_mtx[MPerBlock, NPerBlock] is distributed among threads, and saved in
760  // register
761  // sanity check
762  constexpr bool is_single_rate_mfma =
764  K1 <= 4) ||
767  K1 < 32))
768  ? true
769  : false;
770  constexpr auto is_scale_mfma = false;
771  constexpr index_t KPack = math::max(K1,
773  MPerXdl,
774  NPerXdl,
776  is_single_rate_mfma,
777  is_scale_mfma>::selected_mfma.k_per_blk);
778 
779  auto blockwise_gemm =
783  FloatAcc,
784  decltype(a_k0_m_k1_block_desc),
785  decltype(b_k0_n_k1_block_desc),
786  MPerXdl,
787  NPerXdl,
788  MRepeat,
789  NRepeat,
790  KPack>{};
791 
792  auto c_thread_buf = blockwise_gemm.GetCThreadBuffer();
793 
794  // LDS allocation for A and B: be careful of alignment
795  constexpr auto a_block_space_size =
796  math::integer_least_multiple(a_k0_m_k1_block_desc.GetElementSpaceSize(), max_lds_align);
797 
798  constexpr auto a_block_slice_copy_step = make_multi_index(0, K0PerBlock, 0, 0);
799  constexpr auto b_block_slice_copy_step = make_multi_index(0, K0PerBlock, 0, 0);
800 
801  auto a_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
802  static_cast<FloatAAdjusted*>(p_shared), a_k0_m_k1_block_desc.GetElementSpaceSize());
803 
804  auto b_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
805  static_cast<FloatBAdjusted*>(p_shared) + a_block_space_size,
806  b_k0_n_k1_block_desc.GetElementSpaceSize());
807 
808  // gridwise GEMM pipeline
809  const index_t K0BlockMainLoop = __builtin_amdgcn_readfirstlane(K0 / K0PerBlock);
810 
811  GridwiseGemmPipe::template Run<HasMainKBlockLoop>(a_b_k0_m_k1_grid_desc,
812  a_b_k0_m_k1_block_desc,
813  a_blockwise_copy,
814  a_grid_buf,
815  a_block_buf,
816  a_block_slice_copy_step,
817  b_b_k0_n_k1_grid_desc,
818  b_b_k0_n_k1_block_desc,
819  b_blockwise_copy,
820  b_grid_buf,
821  b_block_buf,
822  b_block_slice_copy_step,
823  blockwise_gemm,
824  c_thread_buf,
825  K0BlockMainLoop);
826 
827  // output: register to global memory
828  {
829  constexpr index_t MWave = MPerBlock / (MRepeat * MPerXdl);
830  constexpr index_t NWave = NPerBlock / (NRepeat * NPerXdl);
831 
832  constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc =
833  blockwise_gemm.GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2();
834 
835  constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc =
836  blockwise_gemm.GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2();
837 
838  constexpr auto M0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I0);
839  constexpr auto N0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I1);
840  constexpr auto M1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I2);
841  constexpr auto N1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I3);
842  constexpr auto M2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I4);
843  constexpr auto M3 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I5);
844  constexpr auto M4 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I6);
845  constexpr auto N2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I7);
846 
847  constexpr auto c_block_desc_mblock_mperblock_nblock_nperblock =
848  GetCBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock();
849 
850  auto c_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
851  static_cast<FloatC*>(p_shared),
852  c_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());
853 
854  static_assert(M1 == MWave, "");
855  static_assert(N1 == NWave, "");
856  static_assert(M2 * M3 * M4 == MPerXdl, "");
857  static_assert(N2 == NPerXdl, "");
858 
859  constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2 = transform_tensor_descriptor(
860  c_block_desc_mblock_mperblock_nblock_nperblock,
861  make_tuple(
862  make_freeze_transform(I0), // freeze mblock
863  make_unmerge_transform(make_tuple(CShuffleMRepeatPerShuffle,
864  M1,
865  M2,
866  M3,
867  M4)), // M1 = MWave, M2 * M3 * M4 = MPerXdl
868  make_freeze_transform(I0), // freeze nblock
869  make_unmerge_transform(make_tuple(CShuffleNRepeatPerShuffle,
870  N1,
871  N2))), // M1 = MWave, M2 * M3 * M4 = MPerXdl
873  make_tuple(
875 
876  // calculate origin of thread output tensor on global memory
877  // blockwise GEMM c matrix starting index
878  const auto c_thread_mtx_on_block =
879  blockwise_gemm.CalculateCThreadOriginDataIndex(I0, I0, I0, I0);
880 
881  const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0];
882  const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1];
883 
884  const auto m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor =
886  make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4))),
889 
890  const auto m_thread_data_on_block_idx =
891  m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor.CalculateBottomIndex(
892  make_multi_index(m_thread_data_on_block));
893 
894  const auto n_thread_data_on_block_to_n0_n1_n2_adaptor =
899 
900  const auto n_thread_data_on_block_idx =
901  n_thread_data_on_block_to_n0_n1_n2_adaptor.CalculateBottomIndex(
902  make_multi_index(n_thread_data_on_block));
903 
904  // VGPR to LDS
905  auto c_thread_copy_vgpr_to_lds =
907  FloatC,
908  decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc),
909  decltype(c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2),
911  Sequence<CShuffleMRepeatPerShuffle,
912  CShuffleNRepeatPerShuffle,
913  I1,
914  I1,
915  M2,
916  I1,
917  M4,
918  I1>,
920  7,
921  1,
922  InMemoryDataOperationEnum::Set,
923  1,
924  true>{
925  c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2,
927  0,
928  m_thread_data_on_block_idx[I1],
929  n_thread_data_on_block_idx[I1],
930  m_thread_data_on_block_idx[I2],
931  m_thread_data_on_block_idx[I3],
932  m_thread_data_on_block_idx[I4],
933  n_thread_data_on_block_idx[I2]),
935 
936  // LDS to global
937  auto c_block_copy_lds_to_global = ThreadGroupTensorSliceTransfer_v6r1<
938  ThisThreadBlock, // index_t BlockSize,
939  CElementwiseOperation, // ElementwiseOperation,
940  CGlobalMemoryDataOperation, // DstInMemOp,
941  Sequence<1,
942  CShuffleMRepeatPerShuffle * MWave * MPerXdl,
943  1,
944  CShuffleNRepeatPerShuffle * NWave * NPerXdl>, // BlockSliceLengths,
945  CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
946  Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder,
947  FloatC, // typename SrcData,
948  FloatC, // typename DstData,
949  decltype(c_block_desc_mblock_mperblock_nblock_nperblock),
950  decltype(c_grid_desc_mblock_mperblock_nblock_nperblock),
951  Sequence<0, 1, 2, 3>, // typename DimAccessOrder,
952  3, // index_t VectorDim,
953  CBlockTransferScalarPerVector_NWaveNPerXDL, // index_t ScalarPerVector,
954  true, // bool ThreadTransferSrcResetCoordinateAfterRun,
955  false> // bool ThreadTransferDstResetCoordinateAfterRun
956  {c_block_desc_mblock_mperblock_nblock_nperblock,
957  make_multi_index(0, 0, 0, 0),
958  c_grid_desc_mblock_mperblock_nblock_nperblock,
959  make_multi_index(block_work_idx[I1], 0, block_work_idx[I2], 0),
960  c_element_op};
961 
962  constexpr auto mxdlperwave_forward_step =
963  make_multi_index(0, CShuffleMRepeatPerShuffle * MWave * MPerXdl, 0, 0);
964  constexpr auto nxdlperwave_forward_step =
965  make_multi_index(0, 0, 0, CShuffleNRepeatPerShuffle * NWave * NPerXdl);
966  constexpr auto nxdlperwave_backward_step =
967  make_multi_index(0, 0, 0, -CShuffleNRepeatPerShuffle * NWave * NPerXdl);
968 
969  static_for<0, MRepeat, CShuffleMRepeatPerShuffle>{}([&](auto mxdlperwave_iter) {
970  constexpr auto mxdlperwave = mxdlperwave_iter;
971 
972  static_for<0, NRepeat, CShuffleNRepeatPerShuffle>{}([&](auto nxdlperwave_iter) {
973  constexpr bool nxdlperwave_forward_sweep =
974  (mxdlperwave % (2 * CShuffleMRepeatPerShuffle) == 0);
975 
976  constexpr index_t nxdlperwave_value =
977  nxdlperwave_forward_sweep
978  ? nxdlperwave_iter
979  : (NRepeat - nxdlperwave_iter - CShuffleNRepeatPerShuffle);
980 
981  constexpr auto nxdlperwave = Number<nxdlperwave_value>{};
982 
983  // make sure it's safe to do ds_write
984  block_sync_lds();
985 
986  // VGPR to LDS
987  c_thread_copy_vgpr_to_lds.Run(
988  c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc,
989  make_tuple(mxdlperwave, nxdlperwave, I0, I0, I0, I0, I0, I0),
990  c_thread_buf,
991  c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2,
992  c_block_buf);
993 
994  // make sure it's safe to do ds_read
995  block_sync_lds();
996 
997  // LDS to global
998  c_block_copy_lds_to_global.Run(c_block_desc_mblock_mperblock_nblock_nperblock,
999  c_block_buf,
1000  c_grid_desc_mblock_mperblock_nblock_nperblock,
1001  c_grid_buf);
1002 
1003  // move on nxdlperwave dimension
1004  if constexpr(nxdlperwave_forward_sweep &&
1005  (nxdlperwave < NRepeat - CShuffleNRepeatPerShuffle))
1006  {
1007  c_block_copy_lds_to_global.MoveDstSliceWindow(
1008  c_grid_desc_mblock_mperblock_nblock_nperblock,
1009  nxdlperwave_forward_step);
1010  }
1011  else if constexpr((!nxdlperwave_forward_sweep) && (nxdlperwave > 0))
1012  {
1013  c_block_copy_lds_to_global.MoveDstSliceWindow(
1014  c_grid_desc_mblock_mperblock_nblock_nperblock,
1015  nxdlperwave_backward_step);
1016  }
1017  });
1018 
1019  // move on mxdlperwave dimension
1020  if constexpr(mxdlperwave < MRepeat - CShuffleMRepeatPerShuffle)
1021  {
1022  c_block_copy_lds_to_global.MoveDstSliceWindow(
1023  c_grid_desc_mblock_mperblock_nblock_nperblock, mxdlperwave_forward_step);
1024  }
1025  });
1026  }
1027  }
1028 }; // namespace ck
1029 
1030 } // namespace ck
CK_TILE_DEVICE void block_sync_lds()
Definition: arch.hpp:192
#define CK_MIN_BLOCK_PER_CU
Definition: ck.hpp:31
#define CK_MAX_THREAD_PER_BLOCK
Definition: ck.hpp:30
__host__ constexpr __device__ auto integer_least_multiple(X x, Y y)
Definition: math.hpp:78
__host__ constexpr __device__ T max(T x)
Definition: math.hpp:84
__host__ __device__ multiplies() -> multiplies< void, void >
FIXME: create macro to replace 'host device' and nothing more.
Definition: ck.hpp:268
__host__ constexpr __device__ auto make_multi_index(Xs &&... xs)
Definition: array_multi_index.hpp:15
__host__ constexpr __device__ auto make_naive_tensor_descriptor(const Tuple< Lengths... > &lengths, const Tuple< Strides... > &strides)
Definition: tensor_descriptor_helper.hpp:49
__global__ void kernel_gemm_xdlops_bwd_weight(const FloatA *__restrict__ p_a_grid, const FloatB *__restrict__ p_b_grid, FloatC *__restrict__ p_c_grid, const AGridDesc_B_K0_M_K1 a_b_k0_m_k1_grid_desc, const BGridDesc_B_K0_N_K1 b_b_k0_n_k1_grid_desc, const CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock c_grid_desc_mblock_mperblock_nblock_nperblock, const AElementwiseOperation a_element_op, const BElementwiseOperation b_element_op, const CElementwiseOperation c_element_op, const CBlockClusterAdaptor c_block_cluster_adaptor)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:157
InMemoryDataOperationEnum
Definition: ck.hpp:277
__host__ constexpr __device__ auto make_naive_tensor_descriptor_packed(const Tuple< Lengths... > &lengths)
Definition: tensor_descriptor_helper.hpp:101
__host__ constexpr __device__ auto make_merge_transform(const LowLengths &low_lengths)
Definition: multi_index_transform_helper.hpp:55
__host__ constexpr __device__ auto make_merge_transform_v3_division_mod(const LowLengths &low_lengths)
Definition: multi_index_transform_helper.hpp:84
__host__ constexpr __device__ auto make_naive_tensor_descriptor_aligned(const Tuple< Lengths... > &lengths, Align align)
Definition: tensor_descriptor_helper.hpp:132
__host__ constexpr __device__ auto make_single_stage_tensor_adaptor(const Transforms &transforms, LowerDimensionOldTopIdss, UpperDimensionNewTopIdss)
Definition: tensor_adaptor.hpp:425
ushort bhalf_t
Definition: data_type.hpp:30
__host__ constexpr __device__ auto make_freeze_transform(const LowerIndex &low_idx)
Definition: multi_index_transform_helper.hpp:151
constexpr detail::ignore_t ignore
Definition: ignore.hpp:20
__device__ index_t get_block_1d_id()
Definition: get_id.hpp:47
typename conditional< predicate, X, Y >::type conditional_t
Definition: functional.hpp:115
__host__ constexpr __device__ auto container_reverse_exclusive_scan(const Array< TData, NSize > &x, Reduce f, TData init)
Definition: container_helper.hpp:213
__host__ constexpr __device__ auto make_pass_through_transform(const LowLength &low_length)
Definition: multi_index_transform_helper.hpp:12
__host__ constexpr __device__ auto make_tuple(Xs &&... xs)
Definition: tuple.hpp:211
remove_cv_t< remove_reference_t< T > > remove_cvref_t
Definition: type.hpp:297
__host__ constexpr __device__ auto make_unmerge_transform(const UpLengths &up_lengths, integral_constant< bool, Use24BitIntegerCalculation >=integral_constant< bool, false >{})
Definition: multi_index_transform_helper.hpp:90
int32_t index_t
Definition: ck.hpp:299
__host__ constexpr __device__ auto container_reduce(const Container &x, Reduce reduce, Init init, Number< IBegin >=Number< 0 >{}, Number< IEnd >=Number< Container::Size()>{}, Number< IStep >=Number< 1 >{})
Definition: container_helper.hpp:111
__host__ constexpr __device__ auto transform_tensor_descriptor(const OldTensorDescriptor &old_tensor_desc, const NewTransforms &new_transforms, NewLowerDimensionOldVisibleIdss, NewUpperDimensionNewVisibleIdss)
Definition: tensor_descriptor.hpp:319
PipelineVersion
Definition: gridwise_gemm_pipeline_selector.hpp:18
__host__ __device__ void print_multi_index(const Tuple< Xs... > &x)
Definition: statically_indexed_array_multi_index.hpp:147
typename remove_cv< T >::type remove_cv_t
Definition: type.hpp:295
__host__ constexpr __device__ auto make_merge_transform_v4_no_carry(const LowLengths &low_lengths)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:136
const GenericPointer< typename T::ValueType > T2 value
Definition: pointer.h:1350
Definition: array.hpp:14
Definition: block_to_ctile_map.hpp:720
Definition: blockwise_gemm_smfmac_xdlops.hpp:44
__host__ constexpr __device__ auto & GetCThreadBuffer()
Definition: blockwise_gemm_smfmac_xdlops.hpp:78
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:254
ThisThreadBlock< BlockSize > ThisThreadBlock
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:267
__host__ static constexpr __device__ auto GetABlockDescriptor_Batch_K0PerBlock_MPerBlock_K1()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:343
__host__ static constexpr __device__ auto GetBBlockDescriptor_K0PerBlock_NPerBlock_K1()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:399
__host__ static constexpr __device__ auto GetABlockDescriptor_K0PerBlock_MPerBlock_K1()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:297
__host__ static constexpr __device__ bool CalculateHasMainK0BlockLoop(index_t K0)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:588
ComputeTypeA FloatAAdjusted
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:283
static __device__ void Run(const FloatA *__restrict__ p_a_grid, const FloatB *__restrict__ p_b_grid, FloatC *__restrict__ p_c_grid, void *__restrict__ p_shared, const AGridDesc_B_K0_M_K1 &a_b_k0_m_k1_grid_desc, const BGridDesc_B_K0_N_K1 &b_b_k0_n_k1_grid_desc, const CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock &c_grid_desc_mblock_mperblock_nblock_nperblock, const AElementwiseOperation &a_element_op, const BElementwiseOperation &b_element_op, const CElementwiseOperation &c_element_op, const CBlockClusterAdaptor &c_block_cluster_adaptor)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:641
__host__ static constexpr __device__ auto MakeCBlockClusterAdaptor(const CMNGridDesc &c_m_n_grid_desc, index_t M01, index_t N01, index_t KBatch)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:616
ComputeTypeB FloatBAdjusted
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:284
decltype(MakeCGridDesc_MBlock_MPerBlock_NBlock_NPerBlock(CMNGridDesc{})) CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:637
__host__ static constexpr __device__ auto GetBBlockDescriptor_Batch_K0PerBlock_NPerBlock_K1()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:445
decltype(MakeCBlockClusterAdaptor(CMNGridDesc{}, 1, 1, 1)) CBlockClusterAdaptor
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:638
__host__ static constexpr __device__ auto GetCBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:624
static __device__ constexpr bool IsValidCompilationParameter()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:528
__host__ static constexpr __device__ auto MakeCGridDesc_MBlock_MPerBlock_NBlock_NPerBlock(const CMNGridDesc &c_m_n_grid_desc)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:599
__host__ static constexpr __device__ bool CheckValidity(const AGridDesc_B_K0_M_K1 &a_b_k0_m_k1_grid_desc, const BGridDesc_B_K0_N_K1 &b_b_k0_n_k1_grid_desc, const CMNGridDesc &c_m_n_grid_desc, const Block2CTileMap &block_2_ctile_map)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:544
remove_cvref_t< decltype(GridwiseGemmPipeline_Selector< PipelineVer, NumGemmKPrefetchStage >())> GridwiseGemmPipe
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:270
__host__ static constexpr __device__ index_t GetSharedMemoryNumberOfByte()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:501
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:25
__host__ constexpr __device__ Merge_v4_no_carry(const LowLengths &low_lengths)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:43
LowLengthsScan low_lengths_scan_
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:38
__host__ constexpr __device__ Merge_v4_no_carry()=default
decltype(make_tuple(container_reduce(LowLengths{}, math::multiplies{}, Number< 1 >{}))) UpLengths
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:35
__host__ static constexpr __device__ bool IsValidUpperIndexMappedToValidLowerIndex(const UpIdx &)
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:116
__host__ __device__ void UpdateLowerIndex(LowIdxDiff &idx_diff_low, const UpIdxDiff &idx_up_diff, LowIdx &idx_low, const UpIdx &idx_up_new, Number< Hack >) const
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:81
static constexpr index_t NDimLow
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:26
__host__ static constexpr __device__ index_t GetNumOfLowerDimension()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:52
__host__ constexpr __device__ const auto & GetUpperLengths() const
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:56
__host__ constexpr __device__ void CalculateLowerIndex(LowIdx &idx_low, const UpIdx &idx_up) const
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:59
__host__ static constexpr __device__ bool IsKnownAtCompileTime()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:107
__host__ static constexpr __device__ bool IsLinearTransform()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:100
UpLengths up_lengths_
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:39
decltype(container_reverse_exclusive_scan(LowLengths{}, math::multiplies{}, Number< 1 >{})) LowLengthsScan
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:32
__host__ static constexpr __device__ index_t GetNumOfUpperDimension()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:54
__host__ static constexpr __device__ bool IsValidUpperIndexAlwaysMappedToValidLowerIndex()
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:102
LowLengths low_lengths_
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:37
__host__ __device__ void Print() const
Definition: gridwise_gemm_xdlops_bwd_weight.hpp:121
Selects the appropriate MFMA instruction type and configuration for given data types and tile sizes o...
Definition: xdlops_gemm.hpp:1208
Definition: sequence.hpp:43
Blockwise data transfer.
Definition: thread_group_tensor_slice_transfer_v4r1.hpp:46
Definition: thread_group_tensor_slice_transfer_v6r1.hpp:34
Definition: threadwise_tensor_slice_transfer.hpp:39
Definition: integral_constant.hpp:20
Definition: is_known_at_compile_time.hpp:14
Definition: type.hpp:177
Definition: math.hpp:34
Definition: functional2.hpp:33
Definition: unary_element_wise_operation.hpp:334