gridwise_gemm_xdl_cshuffle_v3_mx.hpp Source File

• include
• ck
• tensor_operation
• gpu
• grid

gridwise_gemm_xdl_cshuffle_v3_mx.hpp

Go to the documentation of this file.

1// SPDX-License-Identifier: MIT

2// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.

3

4#pragma once

5

6#include "ck/tensor_description/multi_index_transform_helper.hpp"

7#include "ck/tensor_description/tensor_descriptor.hpp"

8#include "ck/tensor_description/tensor_descriptor_helper.hpp"

9#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_mx_selector.hpp"

10#include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v4r1.hpp"

11#include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r1.hpp"

12#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"

13#include "ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp"

14#include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp"

15#include "ck/utility/common_header.hpp"

16#include "ck/utility/env.hpp"

17#include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp"

18

19namespace ck {

20

21#ifndef KERNEL_GEMM_XDL_CSHUFFLE_V3_MX

22#define KERNEL_GEMM_XDL_CSHUFFLE_V3_MX

23// Currently we do not have a elegant way to put single lds buffer & double lds buffer pipe in same

24// kernel function Blockers:

25// 1. Two separted declaration of __shared__ pointer is the key to make sure data access operate on

26// two lds chunks.

27// 2. Occupied __shared__ won't release until whole shader end, a.k.a AB and C may not use same lds

28// buffer when we declare __shared__ inside blkgemmpipe

29template <bool Use2LDS,

30 typename GridwiseGemm,

31 bool HasMainKBlockLoop,

32 InMemoryDataOperationEnum CGlobalMemoryDataOperation,

33 index_t MinimumOccupancy = 1,

34 TailNumber TailNum = TailNumber::Full>

35__global__ enable_if_t<!Use2LDS, void>

36#if CK_USE_LAUNCH_BOUNDS

37__launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)

38#endif

39 // __attribute__((amdgpu_waves_per_eu(1, 1)))

40 kernel_gemm_xdl_cshuffle_v3_mx(typename GridwiseGemm::Argument karg)

41{

42#if defined(__gfx950__)

43 if constexpr(GridwiseGemm::template IsValidCompilationParameter<CGlobalMemoryDataOperation>())

44 {

45 __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];

46

47 auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, blockIdx.z);

48

49 GridwiseGemm::template Run<HasMainKBlockLoop, CGlobalMemoryDataOperation, TailNum>(

50 karg.p_a_grid + splitk_batch_offset.a_k_split_offset,

51 karg.p_a_scale_grid + splitk_batch_offset.a_scale_k_split_offset,

52 karg.p_b_grid + splitk_batch_offset.b_k_split_offset,

53 karg.p_b_scale_grid + splitk_batch_offset.b_scale_k_split_offset,

54 karg.p_c_grid + splitk_batch_offset.c_reduce_offset,

55 p_shared,

56 karg);

57 }

58#else

59 ignore = karg;

60#endif // end of if (defined(__gfx9__))

61}

62

63template <bool Use2LDS,

64 typename GridwiseGemm,

65 bool HasMainKBlockLoop,

66 InMemoryDataOperationEnum CGlobalMemoryDataOperation,

67 index_t MinimumOccupancy = 1,

68 TailNumber TailNum = TailNumber::Full>

69__global__ enable_if_t<Use2LDS, void>

70#if CK_USE_LAUNCH_BOUNDS

71__launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)

72#endif

73 // __attribute__((amdgpu_waves_per_eu(1, 1)))

74 kernel_gemm_xdl_cshuffle_v3_mx(typename GridwiseGemm::Argument karg)

75{

76#if defined(__gfx950__)

77 if constexpr(GridwiseGemm::template IsValidCompilationParameter<CGlobalMemoryDataOperation>())

78 {

79 // Pass two lds pointer is the key to tell compiler that ds_read/write

80 // operate on different lds chunk at same time without order dependecy

81 __shared__ char p_shared_0[GridwiseGemm::GetSharedMemoryNumberOfByte()];

82 __shared__ char p_shared_1[GridwiseGemm::GetSharedMemoryNumberOfByte()];

83

84 auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, blockIdx.z);

85

86 GridwiseGemm::template Run_2Lds<HasMainKBlockLoop, CGlobalMemoryDataOperation, TailNum>(

87 karg.p_a_grid + splitk_batch_offset.a_k_split_offset,

88 karg.p_a_scale_grid + splitk_batch_offset.a_scale_k_split_offset,

89 karg.p_b_grid + splitk_batch_offset.b_k_split_offset,

90 karg.p_b_scale_grid + splitk_batch_offset.b_scale_k_split_offset,

91 karg.p_c_grid + splitk_batch_offset.c_reduce_offset,

92 p_shared_0,

93 p_shared_1,

94 karg);

95 }

96#else

97 ignore = karg;

98#endif // end of if (defined(__gfx9__))

99}

100#endif

101

102template <typename ALayout,

103 typename BLayout,

104 typename CLayout,

105 typename ADataType,

106 typename AScaleDataType,

107 typename BDataType,

108 typename BScaleDataType,

109 typename AccDataType,

110 typename CShuffleDataType,

111 typename CDataType,

112 typename AElementwiseOperation,

113 typename BElementwiseOperation,

114 typename CElementwiseOperation,

115 tensor_operation::device::GemmSpecialization GemmSpec,

116 index_t ScaleBlockSize, // Scaling block size

117 index_t BlockSize, // Thread block size

118 index_t MPerBlock,

119 index_t NPerBlock,

120 index_t KPerBlock,

121 index_t AK1Value,

122 index_t BK1Value,

123 index_t MPerXdl,

124 index_t NPerXdl,

125 index_t MXdlPerWave,

126 index_t NXdlPerWave,

127 typename ABlockTransferThreadClusterLengths_AK0_M_AK1,

128 typename ABlockTransferThreadClusterArrangeOrder,

129 typename ABlockTransferSrcAccessOrder,

130 index_t ABlockTransferSrcVectorDim,

131 index_t ABlockTransferSrcScalarPerVector,

132 index_t ABlockTransferDstScalarPerVector_AK1,

133 bool AThreadTransferSrcResetCoordinateAfterRun,

134 index_t ABlockLdsExtraM,

135 typename BBlockTransferThreadClusterLengths_BK0_N_BK1,

136 typename BBlockTransferThreadClusterArrangeOrder,

137 typename BBlockTransferSrcAccessOrder,

138 index_t BBlockTransferSrcVectorDim,

139 index_t BBlockTransferSrcScalarPerVector,

140 index_t BBlockTransferDstScalarPerVector_BK1,

141 bool BThreadTransferSrcResetCoordinateAfterRun,

142 index_t BBlockLdsExtraN,

143 index_t CShuffleMXdlPerWavePerShuffle,

144 index_t CShuffleNXdlPerWavePerShuffle,

145 typename CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,

146 index_t CShuffleBlockTransferScalarPerVector_NPerBlock,

147 BlockGemmPipelineScheduler BlkGemmPipeSched = BlockGemmPipelineScheduler::Intrawave,

148 BlockGemmPipelineVersion BlkGemmPipelineVer = BlockGemmPipelineVersion::v3,

149 typename ComputeTypeA =

150 ADataType, // XXX: These should always be the same as ADataType and BDataType

151 typename ComputeTypeB =

152 BDataType, // TODO: Hardcode them and remove from the list of template parameters

153 bool PermuteA = false,

154 bool PermuteB = false>

155struct GridwiseGemmMX_xdl_cshuffle_v3

156{

157

158 static constexpr auto I0 = Number<0>{};

159 static constexpr auto I1 = Number<1>{};

160 static constexpr auto I2 = Number<2>{};

161 static constexpr auto I3 = Number<3>{};

162 static constexpr auto I4 = Number<4>{};

163 static constexpr auto I5 = Number<5>{};

164 static constexpr auto I6 = Number<6>{};

165 static constexpr auto I7 = Number<7>{};

166 static constexpr auto I8 = Number<8>{};

167 static constexpr auto I9 = Number<9>{};

168

169 // K1 should be Number<...>

170 static constexpr auto AK0Number = Number<KPerBlock / AK1Value>{};

171 static constexpr auto BK0Number = Number<KPerBlock / BK1Value>{};

172 static constexpr auto AK1Number = Number<AK1Value>{};

173 static constexpr auto BK1Number = Number<BK1Value>{};

174

175 static constexpr auto lcm_AK1_BK1 = math::lcm(AK1Number, BK1Number);

176 static constexpr bool is_single_rate_mfma = false;

177 static constexpr auto is_scale_mfma = true;

178

179 static constexpr auto MXdlPack = 2;

180 static constexpr auto NXdlPack = 2;

181 static constexpr auto KXdlPack = 2;

182

183 //> KPack is at least the k_per_blk of selected mfma

184 //

185 // Should be a multiple of k_per_blk.

186 // TODO: Move this to blockwise pipeline base

187 // KPack in packed data types for pk A/B

188

189 static constexpr index_t APackedSize = packed_size_v<ADataType>;

190 static constexpr index_t BPackedSize = packed_size_v<BDataType>;

191

192 static constexpr index_t KPack =

193 math::max(lcm_AK1_BK1,

194 MfmaSelector<ComputeTypeA,

195 MPerXdl,

196 NPerXdl,

197 ComputeTypeB,

198 is_single_rate_mfma,

199 is_scale_mfma>::selected_mfma.k_per_blk /

200 APackedSize);

201

202 using ThisThreadBlock = ThisThreadBlock<BlockSize>;

203

204 __host__ static auto CalculateGridSize(index_t M, index_t N, index_t KBatch)

205 {

206 return std::make_tuple(Block2CTileMap::CalculateGridSize(M, N), 1, KBatch);

207 }

208

209 __host__ static auto CalculateMPadded(index_t M)

210 {

211 return math::integer_least_multiple(M, MPerBlock);

212 }

213

214 __host__ static auto CalculateNPadded(index_t N)

215 {

216 return math::integer_least_multiple(N, NPerBlock);

217 }

218

219 __host__ static auto CalculateKPadded(index_t K)

220 {

221 return math::integer_divide_ceil(K, KPerBlock) * KPerBlock;

222 }

223

224 __host__ static auto CalculateAK0Padded(index_t K, index_t K_Batch = 1)

225 {

226 auto K_t = K_Batch * KPerBlock;

227 return (K + K_t - 1) / K_t * (KPerBlock / AK1Value);

228 }

229

230 __host__ static auto CalculateBK0Padded(index_t K, index_t K_Batch = 1)

231 {

232 auto K_t = K_Batch * KPerBlock;

233 return (K + K_t - 1) / K_t * (KPerBlock / BK1Value);

234 }

235

236 __host__ static auto CalculateKPadded(index_t K, index_t K_Batch = 1)

237 {

238 auto K_t = K_Batch * KPerBlock;

239 return (K + K_t - 1) / K_t * KPerBlock;

240 }

241

242 __host__ static auto CalculateKRead(index_t K, index_t K_Batch = 1)

243 {

244 constexpr auto KReadVec = math::lcm(AK1Number, BK1Number);

245 auto K_t = K_Batch * KReadVec;

246 return (K + K_t - 1) / K_t * KReadVec;

247 }

248

249 __host__ static auto CalculateMBlock(index_t M)

250 {

251 return math::integer_divide_ceil(M, MPerBlock);

252 }

253

254 __host__ static auto CalculateNBlock(index_t N)

255 {

256 return math::integer_divide_ceil(N, NPerBlock);

257 }

258

259 template <index_t MNXdlPerWave,

260 index_t MNWaves,

261 index_t MNXdlPack,

262 index_t MNPerXdl,

263 typename TileDesc_K0_MN_K1>

264 __host__ __device__ static constexpr auto MakeGemmMmaTileDescriptor(const TileDesc_K0_MN_K1&)

265 {

266 constexpr index_t K0 = TileDesc_K0_MN_K1{}.GetLength(Number<0>{});

267 constexpr index_t MN = TileDesc_K0_MN_K1{}.GetLength(Number<1>{});

268 constexpr index_t K1 = TileDesc_K0_MN_K1{}.GetLength(Number<2>{});

269

270 constexpr auto permuted_desc = transform_tensor_descriptor(

271 TileDesc_K0_MN_K1{},

272 make_tuple(make_xor_with_modulo_transform(make_tuple(Number<MN>{}, Number<K0>{})),

273 make_pass_through_transform(Number<K1>{})),

274 make_tuple(Sequence<1, 0>{}, Sequence<2>{}),

275 make_tuple(Sequence<1, 0>{}, Sequence<2>{}));

276

277 return transform_tensor_descriptor(

278 permuted_desc,

279 make_tuple(make_merge_transform_v3_division_mod(make_tuple(Number<K0>{}, Number<K1>{})),

280 make_unmerge_transform(make_tuple(Number<MNXdlPerWave / MNXdlPack>{},

281 Number<MNWaves>{},

282 Number<MNXdlPack>{},

283 Number<MNPerXdl>{}))),

284 make_tuple(Sequence<0, 2>{}, Sequence<1>{}),

285 make_tuple(Sequence<4>{}, Sequence<0, 1, 2, 3>{}));

286 }

287

288 __host__ __device__ static auto MakeAGridDescriptor_AK0_M_AK1(

289 index_t M, index_t MPad, index_t K, index_t KPad, index_t StrideA, index_t AK0)

290 {

291 const auto a_grid_desc_mraw_kraw = [&]() {

292 if constexpr(is_same_v<tensor_layout::gemm::RowMajor, ALayout>)

293 {

294 return make_naive_tensor_descriptor(make_tuple(M, K), make_tuple(StrideA, I1));

295 }

296 else if constexpr(is_same_v<tensor_layout::gemm::ColumnMajor, ALayout>)

297 {

298 return make_naive_tensor_descriptor(make_tuple(M, K), make_tuple(I1, StrideA));

299 }

300 }();

301

302 using GemmSpecialization = tensor_operation::device::GemmSpecialization;

303

304 if constexpr(GemmSpec == GemmSpecialization::MKPadding ||

305 GemmSpec == GemmSpecialization::MNKPadding)

306 {

307 // pad both M and K

308 const auto a_grid_desc_m_k =

309 transform_tensor_descriptor(a_grid_desc_mraw_kraw,

310 make_tuple(make_right_pad_transform(M, MPad - M),

311 make_right_pad_transform(K, KPad - K)),

312 make_tuple(Sequence<0>{}, Sequence<1>{}),

313 make_tuple(Sequence<0>{}, Sequence<1>{}));

314

315 const auto a_grid_desc_ak0_m_ak1 = transform_tensor_descriptor(

316 a_grid_desc_m_k,

317 make_tuple(make_unmerge_transform(make_tuple(AK0, AK1Value)),

318 make_pass_through_transform(MPad)),

319 make_tuple(Sequence<1>{}, Sequence<0>{}),

320 make_tuple(Sequence<0, 2>{}, Sequence<1>{}));

321

322 return a_grid_desc_ak0_m_ak1;

323 }

324 else if constexpr(GemmSpec == GemmSpecialization::MPadding ||

325 GemmSpec == GemmSpecialization::MNPadding)

326 {

327 // pad M, but not K

328 const auto a_grid_desc_ak0_m_ak1 = transform_tensor_descriptor(

329 a_grid_desc_mraw_kraw,

330 make_tuple(make_unmerge_transform(make_tuple(K / KPerBlock, AK0Number, AK1Value)),

331 make_right_pad_transform(M, MPad - M)),

332 make_tuple(Sequence<1>{}, Sequence<0>{}),

333 make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));

334

335 const auto a_grid_desc_permuted = transform_tensor_descriptor(

336 a_grid_desc_ak0_m_ak1,

337 make_tuple(make_pass_through_transform(K / KPerBlock),

338 make_xor_with_modulo_transform(make_tuple(MPad, AK0Number)),

339 make_pass_through_transform(AK1Value)),

340 make_tuple(Sequence<0>{}, Sequence<2, 1>{}, Sequence<3>{}),

341 make_tuple(Sequence<0>{}, Sequence<2, 1>{}, Sequence<3>{}));

342

343 const auto a_grid_desc = transform_tensor_descriptor(

344 a_grid_desc_permuted,

345 make_tuple(

346 make_merge_transform_v3_division_mod(make_tuple(K / KPerBlock, AK0Number)),

347 make_pass_through_transform(MPad),

348 make_pass_through_transform(AK1Value)),

349 make_tuple(Sequence<0, 1>{}, Sequence<2>{}, Sequence<3>{}),

350 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));

351 return a_grid_desc;

352 }

353 else if constexpr(GemmSpec == GemmSpecialization::KPadding ||

354 GemmSpec == GemmSpecialization::NKPadding)

355 {

356 // pad K, but not M

357 const auto a_grid_desc_m_k = transform_tensor_descriptor(

358 a_grid_desc_mraw_kraw,

359 make_tuple(make_pass_through_transform(M), make_right_pad_transform(K, KPad - K)),

360 make_tuple(Sequence<0>{}, Sequence<1>{}),

361 make_tuple(Sequence<0>{}, Sequence<1>{}));

362

363 const auto a_grid_desc_ak0_m_ak1 = transform_tensor_descriptor(

364 a_grid_desc_m_k,

365 make_tuple(make_unmerge_transform(make_tuple(AK0, AK1Value)),

366 make_pass_through_transform(M)),

367 make_tuple(Sequence<1>{}, Sequence<0>{}),

368 make_tuple(Sequence<0, 2>{}, Sequence<1>{}));

369

370 return a_grid_desc_ak0_m_ak1;

371 }

372 else

373 {

374 // not pad M or K

375 const auto a_grid_desc_ak0_m_ak1 = transform_tensor_descriptor(

376 a_grid_desc_mraw_kraw,

377 make_tuple(make_unmerge_transform(make_tuple(K / KPerBlock, AK0Number, AK1Value)),

378 make_pass_through_transform(M)),

379 make_tuple(Sequence<1>{}, Sequence<0>{}),

380 make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));

381

382 const auto a_grid_desc_permuted = transform_tensor_descriptor(

383 a_grid_desc_ak0_m_ak1,

384 make_tuple(make_pass_through_transform(K / KPerBlock),

385 make_xor_with_modulo_transform(make_tuple(M, AK0Number)),

386 make_pass_through_transform(AK1Value)),

387 make_tuple(Sequence<0>{}, Sequence<2, 1>{}, Sequence<3>{}),

388 make_tuple(Sequence<0>{}, Sequence<2, 1>{}, Sequence<3>{}));

389

390 const auto a_grid_desc = transform_tensor_descriptor(

391 a_grid_desc_permuted,

392 make_tuple(

393 make_merge_transform_v3_division_mod(make_tuple(K / KPerBlock, AK0Number)),

394 make_pass_through_transform(M),

395 make_pass_through_transform(AK1Value)),

396 make_tuple(Sequence<0, 1>{}, Sequence<2>{}, Sequence<3>{}),

397 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));

398

399 return a_grid_desc;

400 }

401 }

402

403 __host__ __device__ static auto MakeBGridDescriptor_BK0_N_BK1(

404 index_t K, index_t KPad, index_t N, index_t NPad, index_t StrideB, index_t BK0)

405 {

406 const auto b_grid_desc_nraw_kraw = [&]() {

407 if constexpr(is_same<tensor_layout::gemm::RowMajor, BLayout>::value)

408 {

409 return make_naive_tensor_descriptor(make_tuple(N, K), make_tuple(I1, StrideB));

410 }

411 else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, BLayout>::value)

412 {

413 return make_naive_tensor_descriptor(make_tuple(N, K), make_tuple(StrideB, I1));

414 }

415 }();

416

417 using GemmSpecialization = tensor_operation::device::GemmSpecialization;

418

419 static_assert(!(is_same_v<remove_cvref_t<ADataType>, pk_i4_t> &&

420 GemmSpec != GemmSpecialization::Default),

421 "pk_i4_t does not support padding");

422 static_assert(!(is_same_v<remove_cvref_t<ADataType>, f4x2_pk_t> &&

423 (GemmSpec != GemmSpecialization::Default &&

424 GemmSpec != GemmSpecialization::MPadding)),

425 "f4x2_pk_t does not support K padding");

426 static_assert(!((is_same_v<remove_cvref_t<ADataType>, f6x16_pk_t> ||

427 is_same_v<remove_cvref_t<ADataType>, bf6x16_pk_t> ||

428 is_same_v<remove_cvref_t<ADataType>, f6x32_pk_t> ||

429 is_same_v<remove_cvref_t<ADataType>, bf6x32_pk_t>) &&

430 GemmSpec != GemmSpecialization::Default),

431 "Packed F6 types do not support padding");

432

433 if constexpr(GemmSpec == GemmSpecialization::NKPadding ||

434 GemmSpec == GemmSpecialization::MNKPadding)

435 {

436 // pad both N and K

437 const auto b_grid_desc_n_k =

438 transform_tensor_descriptor(b_grid_desc_nraw_kraw,

439 make_tuple(make_right_pad_transform(N, NPad - N),

440 make_right_pad_transform(K, KPad - K)),

441 make_tuple(Sequence<0>{}, Sequence<1>{}),

442 make_tuple(Sequence<0>{}, Sequence<1>{}));

443

444 const auto b_grid_desc_bk0_n_bk1 = transform_tensor_descriptor(

445 b_grid_desc_n_k,

446 make_tuple(make_unmerge_transform(make_tuple(BK0, BK1Value)),

447 make_pass_through_transform(NPad)),

448 make_tuple(Sequence<1>{}, Sequence<0>{}),

449 make_tuple(Sequence<0, 2>{}, Sequence<1>{}));

450

451 return b_grid_desc_bk0_n_bk1;

452 }

453 else if constexpr(GemmSpec == GemmSpecialization::NPadding ||

454 GemmSpec == GemmSpecialization::MNPadding)

455 {

456 // pad N, but not K

457 const auto b_grid_desc_bk0_n_bk1 = transform_tensor_descriptor(

458 b_grid_desc_nraw_kraw,

459 make_tuple(make_unmerge_transform(make_tuple(BK0, BK1Value)),

460 make_right_pad_transform(N, NPad - N)),

461 make_tuple(Sequence<1>{}, Sequence<0>{}),

462 make_tuple(Sequence<0, 2>{}, Sequence<1>{}));

463

464 return b_grid_desc_bk0_n_bk1;

465 }

466 else if constexpr(GemmSpec == GemmSpecialization::KPadding ||

467 GemmSpec == GemmSpecialization::MKPadding)

468 {

469 // pad K, but not N

470 const auto b_grid_desc_n_k = transform_tensor_descriptor(

471 b_grid_desc_nraw_kraw,

472 make_tuple(make_pass_through_transform(N), make_right_pad_transform(K, KPad - K)),

473 make_tuple(Sequence<0>{}, Sequence<1>{}),

474 make_tuple(Sequence<0>{}, Sequence<1>{}));

475

476 const auto b_grid_desc_bk0_n_bk1 = transform_tensor_descriptor(

477 b_grid_desc_n_k,

478 make_tuple(make_unmerge_transform(make_tuple(BK0, BK1Value)),

479 make_pass_through_transform(N)),

480 make_tuple(Sequence<1>{}, Sequence<0>{}),

481 make_tuple(Sequence<0, 2>{}, Sequence<1>{}));

482

483 return b_grid_desc_bk0_n_bk1;

484 }

485 else

486 {

487 if constexpr(!PermuteB)

488 {

489 // not pad N or K

490 const auto b_grid_desc_bk0_n_bk1 = transform_tensor_descriptor(

491 b_grid_desc_nraw_kraw,

492 make_tuple(

493 make_unmerge_transform(make_tuple(K / KPerBlock, BK0Number, BK1Value)),

494 make_pass_through_transform(N)),

495 make_tuple(Sequence<1>{}, Sequence<0>{}),

496 make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));

497

498 const auto b_grid_desc_permuted = transform_tensor_descriptor(

499 b_grid_desc_bk0_n_bk1,

500 make_tuple(make_pass_through_transform(K / KPerBlock),

501 make_xor_with_modulo_transform(make_tuple(N, BK0Number)),

502 make_pass_through_transform(BK1Value)),

503 make_tuple(Sequence<0>{}, Sequence<2, 1>{}, Sequence<3>{}),

504 make_tuple(Sequence<0>{}, Sequence<2, 1>{}, Sequence<3>{}));

505

506 const auto b_grid_desc = transform_tensor_descriptor(

507 b_grid_desc_permuted,

508 make_tuple(

509 make_merge_transform_v3_division_mod(make_tuple(K / KPerBlock, BK0Number)),

510 make_pass_through_transform(N),

511 make_pass_through_transform(BK1Value)),

512 make_tuple(Sequence<0, 1>{}, Sequence<2>{}, Sequence<3>{}),

513 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));

514

515 return b_grid_desc;

516 }

517 else

518 {

519 // Weight Tile Permute

520 constexpr index_t BK01 = KPerBlock / BK1Value;

521 // const index_t BK00 = BK0 / BK01;

522 const index_t BK0_ = StrideB / BK1Value;

523 const index_t BK00 = BK0_ / BK01;

524

525 const auto b_grid_desc_bk00_n_bk01_bk1_permute =

526 make_naive_tensor_descriptor_packed(make_tuple(BK00, N, BK01, BK1Value));

527

528 const auto b_grid_desc_bk0_n_bk1_permute = transform_tensor_descriptor(

529 b_grid_desc_bk00_n_bk01_bk1_permute,

530 make_tuple(make_merge_transform(make_tuple(BK00, BK01)),

531 make_pass_through_transform(make_tuple(N)),

532 make_pass_through_transform(BK1Value)),

533 make_tuple(Sequence<0, 2>{}, Sequence<1>{}, Sequence<3>{}),

534 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));

535

536 return b_grid_desc_bk0_n_bk1_permute;

537 }

538 }

539 }

540

541 template <typename ABlockDesc_AK0_M_AK1>

542 __host__ __device__ static constexpr auto

543 MakeAMmaTileDescriptor_M0_M1_M2_M3_K(const ABlockDesc_AK0_M_AK1&)

544 {

545 constexpr index_t MWaves = MPerBlock / (MXdlPerWave * MPerXdl);

546

547 return MakeGemmMmaTileDescriptor<MXdlPerWave, MWaves, MXdlPack, MPerXdl>(

548 ABlockDesc_AK0_M_AK1{});

549 }

550

551 template <typename BBlockDesc_BK0_N_BK1>

552 __host__ __device__ static constexpr auto

553 MakeBMmaTileDescriptor_N0_N1_N2_N3_K(const BBlockDesc_BK0_N_BK1&)

554 {

555 constexpr index_t NWaves = NPerBlock / (NXdlPerWave * NPerXdl);

556

557 return MakeGemmMmaTileDescriptor<NXdlPerWave, NWaves, NXdlPack, NPerXdl>(

558 BBlockDesc_BK0_N_BK1{});

559 }

560

561 __host__ __device__ static auto

562 MakeCGridDescriptor_M_N(index_t M, index_t MPad, index_t N, index_t NPad, index_t StrideC)

563 {

564 const auto c_grid_desc_mraw_nraw = [&]() {

565 if constexpr(is_same<tensor_layout::gemm::RowMajor, CLayout>::value)

566 {

567 return make_naive_tensor_descriptor(make_tuple(M, N), make_tuple(StrideC, I1));

568 }

569 else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, CLayout>::value)

570 {

571 return make_naive_tensor_descriptor(make_tuple(M, N), make_tuple(I1, StrideC));

572 }

573 }();

574

575 // pad M and N

576 return transform_tensor_descriptor(c_grid_desc_mraw_nraw,

577 make_tuple(make_right_pad_transform(M, MPad - M),

578 make_right_pad_transform(N, NPad - N)),

579 make_tuple(Sequence<0>{}, Sequence<1>{}),

580 make_tuple(Sequence<0>{}, Sequence<1>{}));

581#if 0

582 using GemmSpecialization = tensor_operation::device::GemmSpecialization;

583

584 if constexpr(GemmSpec == GemmSpecialization::MNPadding ||

585 GemmSpec == GemmSpecialization::MNKPadding)

586 {

587 // pad M and N

588 return transform_tensor_descriptor(c_grid_desc_mraw_nraw,

589 make_tuple(make_right_pad_transform(M, MPad - M),

590 make_right_pad_transform(N, NPad - N)),

591 make_tuple(Sequence<0>{}, Sequence<1>{}),

592 make_tuple(Sequence<0>{}, Sequence<1>{}));

593 }

594 else if constexpr(GemmSpec == GemmSpecialization::MPadding ||

595 GemmSpec == GemmSpecialization::MKPadding)

596 {

597 // pad M, but not N

598 return transform_tensor_descriptor(

599 c_grid_desc_mraw_nraw,

600 make_tuple(make_right_pad_transform(M, MPad - M), make_pass_through_transform(N)),

601 make_tuple(Sequence<0>{}, Sequence<1>{}),

602 make_tuple(Sequence<0>{}, Sequence<1>{}));

603 }

604 else if constexpr(GemmSpec == GemmSpecialization::NPadding ||

605 GemmSpec == GemmSpecialization::NKPadding)

606 {

607 // pad N, but not M

608 return transform_tensor_descriptor(

609 c_grid_desc_mraw_nraw,

610 make_tuple(make_pass_through_transform(M), make_right_pad_transform(N, NPad - N)),

611 make_tuple(Sequence<0>{}, Sequence<1>{}),

612 make_tuple(Sequence<0>{}, Sequence<1>{}));

613 }

614 else

615 {

616 // not pad M or N

617 return c_grid_desc_mraw_nraw;

618 }

619#endif

620 }

621

622 struct Problem

623 {

624 __host__ Problem(index_t M_,

625 index_t N_,

626 index_t K_,

627 index_t StrideA_,

628 index_t StrideScaleA_,

629 index_t StrideB_,

630 index_t StrideScaleB_,

631 index_t StrideC_,

632 index_t KBatch_)

633 : M{M_},

634 N{N_},

635 K{K_},

636 StrideA{StrideA_},

637 StrideScaleA{StrideScaleA_},

638 StrideB{StrideB_},

639 StrideScaleB{StrideScaleB_},

640 StrideC{StrideC_},

641 KBatch{KBatch_},

642 MPadded{CalculateMPadded(M_)},

643 NPadded{CalculateNPadded(N_)},

644 KRead{CalculateKRead(K_, KBatch_)},

645 KPadded{CalculateKPadded(K_, KBatch_)},

646 AK0{CalculateAK0Padded(K_, KBatch_)},

647 BK0{CalculateBK0Padded(K_, KBatch_)},

648 MBlock{CalculateMBlock(M_)},

649 NBlock{CalculateNBlock(N_)}

650 {

651 }

652

653 __host__ void Print() const

654 {

655 std::cout << "problem {" << "M:" << M << ", " << "N:" << N << ", " << "K:" << K << ", "

656 << "SA:" << StrideA << ", " << "SScaleA:" << StrideScaleA << ", "

657 << "SB:" << StrideB << ", " << "SScaleB:" << StrideScaleB << ", "

658 << "SC:" << StrideC << ", " << "MP:" << MPadded << ", " << "NP:" << NPadded

659 << ", " << "KRead:" << KRead << ", " << "KP:" << KPadded << ", "

660 << "AK0:" << AK0 << ", " << "BK0:" << BK0 << ", " << "MBlock: " << MBlock

661 << ", " << "NBlock: " << NBlock << "}" << std::endl;

662 }

663

664 index_t M;

665 index_t N;

666 index_t K;

667 index_t StrideA;

668 index_t StrideScaleA;

669 index_t StrideB;

670 index_t StrideScaleB;

671 index_t StrideC;

672 index_t KBatch;

673 index_t MPadded;

674 index_t NPadded;

675 index_t KRead;

676 index_t KPadded;

677 index_t AK0;

678 index_t BK0;

679 index_t MBlock;

680 index_t NBlock;

681 };

682

683 // Argument

684 struct Argument : public tensor_operation::device::BaseArgument, public Problem

685 {

686 __host__ Argument(const ADataType* p_a_grid_,

687 const AScaleDataType* p_a_scale_grid_,

688 const BDataType* p_b_grid_,

689 const BScaleDataType* p_b_scale_grid_,

690 CDataType* p_c_grid_,

691 index_t M_,

692 index_t N_,

693 index_t K_,

694 index_t StrideA_,

695 index_t StrideScaleA_,

696 index_t StrideB_,

697 index_t StrideScaleB_,

698 index_t StrideC_,

699 index_t k_batch_,

700 AElementwiseOperation a_element_op_,

701 BElementwiseOperation b_element_op_,

702 CElementwiseOperation c_element_op_,

703 bool is_reduce_ = false)

704 : Problem{M_,

705 N_,

706 K_ / APackedSize,

707 StrideA_ / APackedSize,

708 StrideScaleA_,

709 StrideB_ / BPackedSize,

710 StrideScaleB_,

711 StrideC_,

712 k_batch_},

713 p_a_grid{p_a_grid_},

714 p_a_scale_grid{p_a_scale_grid_},

715 p_b_grid{p_b_grid_},

716 p_b_scale_grid{p_b_scale_grid_},

717 p_c_grid{p_c_grid_},

718 a_element_op{a_element_op_},

719 b_element_op{b_element_op_},

720 c_element_op{c_element_op_},

721 is_reduce(is_reduce_)

722 {

723 }

724

725 __host__ __device__ inline bool IsReduceAdd() const

726 {

727 return (Problem::KBatch > 1) && is_reduce;

728 }

729

730 __host__ __device__ inline bool IsAtomicAdd() const

731 {

732 return (Problem::KBatch > 1) && (!is_reduce);

733 }

734

735 const ADataType* p_a_grid;

736 const AScaleDataType* p_a_scale_grid;

737 const BDataType* p_b_grid;

738 const BScaleDataType* p_b_scale_grid;

739 CDataType* p_c_grid;

740

741 const AElementwiseOperation a_element_op;

742 const BElementwiseOperation b_element_op;

743 const CElementwiseOperation c_element_op;

744 bool is_reduce;

745 };

746

747 struct SplitKBatchOffset

748 {

749

750 __device__ SplitKBatchOffset(Argument& karg, index_t k_id)

751 {

752 if constexpr(is_same_v<tensor_layout::gemm::RowMajor, ALayout>)

753 {

754 a_k_split_offset = k_id * karg.KRead;

755 }

756 else if constexpr(is_same_v<tensor_layout::gemm::ColumnMajor, ALayout>)

757 {

758 a_k_split_offset = k_id * karg.KRead * karg.StrideA;

759 }

760

761 if constexpr(is_same_v<tensor_layout::gemm::RowMajor, BLayout>)

762 {

763 b_k_split_offset = k_id * karg.KRead * karg.StrideB;

764 }

765 else if constexpr(is_same_v<tensor_layout::gemm::ColumnMajor, BLayout>)

766 {

767 if constexpr(!PermuteB)

768 {

769 b_k_split_offset = k_id * karg.KRead;

770 }

771 else

772 {

773 const int k0_offset = karg.KRead * karg.N;

774 b_k_split_offset = k_id * k0_offset;

775 }

776 }

777

778 // Calculate A scale offset

779 a_scale_k_split_offset =

780 k_id * karg.KRead / (ScaleBlockSize / APackedSize) * MXdlPack * MPerXdl;

781

782 // Calculate B scale offset

783 b_scale_k_split_offset =

784 k_id * karg.KRead / (ScaleBlockSize / BPackedSize) * NXdlPack * NPerXdl;

785

786 if(k_id < (karg.KBatch - 1))

787 {

788 karg.K = karg.KRead;

789 }

790 else

791 {

792 karg.K = karg.K - karg.KRead * (karg.KBatch - 1);

793 }

794

795 if(karg.IsReduceAdd())

796 {

797 c_reduce_offset = k_id * karg.M * karg.N;

798 }

799 else

800 {

801 c_reduce_offset = 0;

802 }

803 }

804

805 index_t a_k_split_offset;

806 index_t b_k_split_offset;

807 index_t a_scale_k_split_offset; // New member for scale matrix offset

808 index_t b_scale_k_split_offset; // New member for scale matrix offset

809 index_t c_reduce_offset;

810 };

811

812 __device__ static constexpr auto GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1()

813 {

814 constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl);

815 constexpr index_t NWave = NPerBlock / (NXdlPerWave * NPerXdl);

816 constexpr index_t WaveSize = BlockSize / (MWave * NWave);

817

818 // A matrix in LDS memory, dst of blockwise copy

819 if constexpr(ABlockLdsExtraM || BlkGemmPipelineVer == BlockGemmPipelineVersion::v4)

820 {

821 // contiguous in LDS

822 return make_naive_tensor_descriptor(

823 make_tuple(AK0Number, Number<MPerBlock>{}, AK1Number),

824 make_tuple(AK1Number, Number<KPerBlock>{}, I1));

825 }

826 // xor tensor transformation request more unnecessary vgpr usage, would cause register spill

827 // in some cases.

828 else if constexpr(is_same<tensor_layout::gemm::RowMajor, ALayout>::value)

829 {

830 constexpr auto a_lds_block_desc =

831 make_naive_tensor_descriptor(make_tuple(AK0Number, Number<MPerBlock>{}, AK1Number),

832 make_tuple(AK1Number, Number<KPerBlock>{}, I1));

833

834 constexpr auto a_lds_block_desc_permuted = transform_tensor_descriptor(

835 a_lds_block_desc,

836 make_tuple(make_xor_with_modulo_transform(

837 make_tuple(Number<MPerBlock>{}, Number<AK0Number>{})),

838 make_pass_through_transform(AK1Number)),

839 make_tuple(Sequence<1, 0>{}, Sequence<2>{}),

840 make_tuple(Sequence<1, 0>{}, Sequence<2>{}));

841

842 return a_lds_block_desc_permuted;

843 }

844 else // ColumnMajor A

845 {

846 // kfold and mpair dimension is not always required.

847 // more dimension in merge_transform increase the difficulty of generating immarg offset

848 // for compiler.

849 constexpr auto M0 = ABlockTransferThreadClusterLengths_AK0_M_AK1{}.At(I1);

850 constexpr auto M1 = MPerBlock / M0;

851

852 constexpr auto KThreadWrite = ABlockTransferThreadClusterLengths_AK0_M_AK1{}.At(I0);

853 constexpr auto K0PerThreadWrite = AK0Number / KThreadWrite;

854 constexpr auto KThreadRead = WaveSize / MPerXdl;

855 constexpr auto K0PerThreadRead = AK0Number / KThreadRead;

856

857 constexpr auto kfold = (AK1Number * M0 * sizeof(ADataType) > 128)

858 ? 1

859 : 128 / (AK1Number * M0 * sizeof(ADataType));

860 constexpr auto KThreadReadPerm =

861 (kfold * K0PerThreadWrite / K0PerThreadRead) > 1

862 ? KThreadRead / (kfold * K0PerThreadWrite / K0PerThreadRead)

863 : KThreadRead;

864

865 // 1<=mpair<=n0

866 constexpr auto mpair = (AK1Number * MPerXdl * sizeof(ADataType) > 128)

867 ? 1

868 : ((128 / (AK1Number * MPerXdl * sizeof(ADataType))) > M0

869 ? M0

870 : 128 / (AK1Number * MPerXdl * sizeof(ADataType)));

871

872 constexpr auto a_lds_block_desc = make_naive_tensor_descriptor_packed(

873 make_tuple(Number<KThreadWrite / kfold / KThreadReadPerm>{},

874 Number<K0PerThreadWrite>{},

875 Number<KThreadReadPerm * M1>{},

876 Number<kfold * M0 / mpair>{},

877 Number<mpair>{},

878 AK1Number));

879

880 constexpr auto a_lds_block_desc_permuted = transform_tensor_descriptor(

881 a_lds_block_desc,

882 make_tuple(

883 make_pass_through_transform(Number<KThreadWrite / kfold / KThreadReadPerm>{}),

884 make_pass_through_transform(Number<K0PerThreadWrite>{}),

885 make_xor_with_modulo_transform(

886 make_tuple(Number<KThreadReadPerm * M1>{}, Number<kfold * M0 / mpair>{})),

887 make_pass_through_transform(Number<mpair>{}),

888 make_pass_through_transform(AK1Number)),

889 make_tuple(

890 Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4>{}, Sequence<5>{}),

891 make_tuple(

892 Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4>{}, Sequence<5>{}));

893

894 constexpr auto a_lds_block_desc_unmerged = transform_tensor_descriptor(

895 a_lds_block_desc_permuted,

896 make_tuple(

897 make_pass_through_transform(Number<KThreadWrite / kfold / KThreadReadPerm>{}),

898 make_pass_through_transform(Number<K0PerThreadWrite>{}),

899 make_unmerge_transform(make_tuple(Number<KThreadReadPerm>{}, Number<M1>{})),

900 make_unmerge_transform(make_tuple(Number<kfold>{}, Number<M0 / mpair>{})),

901 make_pass_through_transform(Number<mpair>{}),

902 make_pass_through_transform(AK1Number)),

903 make_tuple(Sequence<0>{},

904 Sequence<1>{},

905 Sequence<2>{},

906 Sequence<3>{},

907 Sequence<4>{},

908 Sequence<5>{}),

909 make_tuple(Sequence<1>{},

910 Sequence<2>{},

911 Sequence<0, 3>{},

912 Sequence<4, 5>{},

913 Sequence<6>{},

914 Sequence<7>{}));

915

916 constexpr auto a_lds_block_desc_ak0_m_ak1 = transform_tensor_descriptor(

917 a_lds_block_desc_unmerged,

918 make_tuple(make_merge_transform_v3_division_mod(

919 make_tuple(Number<KThreadReadPerm>{},

920 Number<KThreadWrite / kfold / KThreadReadPerm>{},

921 Number<kfold>{},

922 Number<K0PerThreadWrite>{})),

923 make_merge_transform_v3_division_mod(

924 make_tuple(Number<M0 / mpair>{}, Number<mpair>{}, Number<M1>{})),

925 make_pass_through_transform(AK1Number)),

926 make_tuple(Sequence<0, 1, 4, 2>{}, Sequence<5, 6, 3>{}, Sequence<7>{}),

927 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));

928

929 return a_lds_block_desc_ak0_m_ak1;

930 }

931 }

932

933 __device__ static constexpr auto GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1()

934 {

935 constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl);

936 constexpr index_t NWave = NPerBlock / (NXdlPerWave * NPerXdl);

937 constexpr index_t WaveSize = BlockSize / (MWave * NWave);

938 // B matrix in LDS memory, dst of blockwise copy

939 if constexpr(BBlockLdsExtraN || BlkGemmPipelineVer == BlockGemmPipelineVersion::v4)

940 {

941 // contiguous in lds

942 return make_naive_tensor_descriptor(

943 make_tuple(BK0Number, Number<NPerBlock>{}, BK1Number),

944 make_tuple(BK1Number, Number<KPerBlock>{}, I1));

945 }

946 else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, BLayout>::value)

947 {

948 // NLdsLayer * K0 as logical Bank

949 constexpr auto b_lds_block_desc =

950 make_naive_tensor_descriptor(make_tuple(BK0Number, Number<NPerBlock>{}, BK1Number),

951 make_tuple(BK1Number, Number<KPerBlock>{}, I1));

952

953 constexpr auto b_lds_block_desc_permuted = transform_tensor_descriptor(

954 b_lds_block_desc,

955 make_tuple(make_xor_with_modulo_transform(

956 make_tuple(Number<NPerBlock>{}, Number<BK0Number>{})),

957 make_pass_through_transform(BK1Number)),

958 make_tuple(Sequence<1, 0>{}, Sequence<2>{}),

959 make_tuple(Sequence<1, 0>{}, Sequence<2>{}));

960

961 return b_lds_block_desc_permuted;

962 }

963 else // RowMajor B

964 {

965 constexpr auto N0 = BBlockTransferThreadClusterLengths_BK0_N_BK1{}.At(I1);

966 constexpr auto N1 = NPerBlock / N0;

967

968 constexpr auto KThreadWrite = BBlockTransferThreadClusterLengths_BK0_N_BK1{}.At(I0);

969 constexpr auto K0PerThreadWrite = BK0Number / KThreadWrite;

970 constexpr auto KThreadRead = WaveSize / NPerXdl;

971 constexpr auto K0PerThreadRead = BK0Number / KThreadRead;

972

973 constexpr auto kfold = (BK1Number * N0 * sizeof(BDataType) > 128)

974 ? 1

975 : 128 / (BK1Number * N0 * sizeof(BDataType));

976 constexpr auto KThreadReadPerm =

977 (kfold * K0PerThreadWrite / K0PerThreadRead) > 1

978 ? KThreadRead / (kfold * K0PerThreadWrite / K0PerThreadRead)

979 : KThreadRead;

980

981 // 1<=npair<=n0

982 constexpr auto npair = (BK1Number * NPerXdl * sizeof(BDataType) > 128)

983 ? 1

984 : ((128 / (BK1Number * NPerXdl * sizeof(BDataType))) > N0

985 ? N0

986 : 128 / (BK1Number * NPerXdl * sizeof(BDataType)));

987

988 constexpr auto b_lds_block_desc = make_naive_tensor_descriptor_packed(

989 make_tuple(Number<KThreadWrite / kfold / KThreadReadPerm>{},

990 Number<K0PerThreadWrite>{},

991 Number<KThreadReadPerm * N1>{},

992 Number<kfold * N0 / npair>{},

993 Number<npair>{},

994 BK1Number));

995

996 constexpr auto b_lds_block_desc_permuted = transform_tensor_descriptor(

997 b_lds_block_desc,

998 make_tuple(

999 make_pass_through_transform(Number<KThreadWrite / kfold / KThreadReadPerm>{}),

1000 make_pass_through_transform(Number<K0PerThreadWrite>{}),

1001 make_xor_with_modulo_transform(

1002 make_tuple(Number<KThreadReadPerm * N1>{}, Number<kfold * N0 / npair>{})),

1003 make_pass_through_transform(Number<npair>{}),

1004 make_pass_through_transform(BK1Number)),

1005 make_tuple(

1006 Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4>{}, Sequence<5>{}),

1007 make_tuple(

1008 Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4>{}, Sequence<5>{}));

1009

1010 constexpr auto b_lds_block_desc_unmerged = transform_tensor_descriptor(

1011 b_lds_block_desc_permuted,

1012 make_tuple(

1013 make_pass_through_transform(Number<KThreadWrite / kfold / KThreadReadPerm>{}),

1014 make_pass_through_transform(Number<K0PerThreadWrite>{}),

1015 make_unmerge_transform(make_tuple(Number<KThreadReadPerm>{}, Number<N1>{})),

1016 make_unmerge_transform(make_tuple(Number<kfold>{}, Number<N0 / npair>{})),

1017 make_pass_through_transform(Number<npair>{}),

1018 make_pass_through_transform(BK1Number)),

1019 make_tuple(Sequence<0>{},

1020 Sequence<1>{},

1021 Sequence<2>{},

1022 Sequence<3>{},

1023 Sequence<4>{},

1024 Sequence<5>{}),

1025 make_tuple(Sequence<1>{},

1026 Sequence<2>{},

1027 Sequence<0, 3>{},

1028 Sequence<4, 5>{},

1029 Sequence<6>{},

1030 Sequence<7>{}));

1031

1032 constexpr auto b_lds_block_desc_bk0_n_bk1 = transform_tensor_descriptor(

1033 b_lds_block_desc_unmerged,

1034 make_tuple(make_merge_transform_v3_division_mod(

1035 make_tuple(Number<KThreadReadPerm>{},

1036 Number<KThreadWrite / kfold / KThreadReadPerm>{},

1037 Number<kfold>{},

1038 Number<K0PerThreadWrite>{})),

1039 make_merge_transform_v3_division_mod(

1040 make_tuple(Number<N0 / npair>{}, Number<npair>{}, Number<N1>{})),

1041 make_pass_through_transform(BK1Number)),

1042 make_tuple(Sequence<0, 1, 4, 2>{}, Sequence<5, 6, 3>{}, Sequence<7>{}),

1043 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));

1044

1045 return b_lds_block_desc_bk0_n_bk1;

1046 }

1047 }

1048

1049 __device__ static constexpr auto GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock()

1050 {

1051 constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl);

1052 constexpr index_t NWave = NPerBlock / (NXdlPerWave * NPerXdl);

1053

1054 constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock =

1055 make_naive_tensor_descriptor_packed(

1056 make_tuple(I1,

1057 Number<CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl>{},

1058 I1,

1059 Number<CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>{}));

1060

1061 return c_shuffle_block_desc_mblock_mperblock_nblock_nperblock;

1062 }

1063

1064 using BlockwiseGemmPipe =

1065 remove_cvref_t<decltype(BlockGemmMXPipeline_Selector<

1066 BlkGemmPipelineVer,

1067 BlkGemmPipeSched,

1068 BlockSize,

1069 ScaleBlockSize,

1070 ADataType,

1071 AScaleDataType,

1072 BDataType,

1073 BScaleDataType,

1074 ComputeTypeA,

1075 AccDataType,

1076 decltype(GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1()),

1077 decltype(GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1()),

1078 decltype(MakeAMmaTileDescriptor_M0_M1_M2_M3_K(

1079 GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1())),

1080 decltype(MakeBMmaTileDescriptor_N0_N1_N2_N3_K(

1081 GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1())),

1082 ABlockTransferSrcScalarPerVector,

1083 BBlockTransferSrcScalarPerVector,

1084 MPerBlock,

1085 NPerBlock,

1086 KPerBlock,

1087 MPerXdl,

1088 NPerXdl,

1089 MXdlPerWave,

1090 NXdlPerWave,

1091 KPack>())>;

1092

1093 __device__ static constexpr index_t GetSharedMemoryNumberOfByte()

1094 {

1095 // LDS allocation for A and B: be careful of alignment

1096 constexpr auto a_block_desc_ak0_m_ak1 = GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();

1097 constexpr auto b_block_desc_bk0_n_bk1 = GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();

1098

1099 // lds max alignment

1100 constexpr auto max_lds_align = math::lcm(AK1Number, BK1Number);

1101

1102 constexpr auto a_block_space_size_aligned = math::integer_least_multiple(

1103 a_block_desc_ak0_m_ak1.GetElementSpaceSize(), max_lds_align);

1104

1105 constexpr auto b_block_space_size_aligned = math::integer_least_multiple(

1106 b_block_desc_bk0_n_bk1.GetElementSpaceSize(), max_lds_align);

1107

1108 // LDS allocation for C shuffle in LDS

1109 constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock =

1110 GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock();

1111

1112 constexpr auto c_block_size =

1113 c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize();

1114

1115 return math::max((a_block_space_size_aligned * sizeof(ADataType) +

1116 b_block_space_size_aligned * sizeof(BDataType)),

1117 c_block_size * sizeof(CShuffleDataType));

1118 }

1119

1120 IS_VALID_COMPILATION_PARAMETER_IMPL(CDataType)

1121

1122 // block_id to matrix tile idx (m0, n0) mapping are controlled by {M01, N01}

1123 __host__ static constexpr bool CheckValidity(const Argument& karg)

1124 {

1125 static_assert((MPerBlock % (MPerXdl * MXdlPerWave) == 0) &&

1126 (NPerBlock % (NXdlPerWave * NPerXdl)) == 0,

1127 "Invalid tuning param!");

1128

1129 static_assert(KPerBlock % (ScaleBlockSize / BPackedSize) == 0,

1130 "KPerBlock should be multiple of ScaleBlockSize");

1131

1132 if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::MPadding ||

1133 GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding ||

1134 GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding ||

1135 GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) &&

1136 !(is_same<tensor_layout::gemm::RowMajor, ALayout>::value))

1137 {

1138 if(!(karg.M % MPerBlock == 0))

1139 {

1140 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1141 {

1142 std::cout << "Arg M value is not a multiple of MPerBlock! M: " << karg.M << " "

1143 << __FILE__ << ":" << __LINE__ << ", in function: " << __func__

1144 << std::endl;

1145 }

1146 return false;

1147 }

1148 }

1149

1150 if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::NPadding ||

1151 GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding ||

1152 GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding ||

1153 GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) &&

1154 (is_same<tensor_layout::gemm::RowMajor, BLayout>::value))

1155 {

1156 if(!(karg.N % NPerBlock == 0))

1157 {

1158 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1159 {

1160 std::cout << "Arg N value is not a multiple of NPerBlock! N: " << karg.N << " "

1161 << __FILE__ << ":" << __LINE__ << ", in function: " << __func__

1162 << std::endl;

1163 }

1164 return false;

1165 }

1166 }

1167

1168 if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::KPadding ||

1169 GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding ||

1170 GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding ||

1171 GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding))

1172 {

1173 auto K_t = karg.KBatch * KPerBlock;

1174 if(!(karg.K % K_t == 0))

1175 {

1176 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1177 {

1178 std::cout << "Arg K value is not a multiple of K_Batch * K0PerBlock * K1! K: "

1179 << karg.K << " " << __FILE__ << ":" << __LINE__

1180 << ", in function: " << __func__ << std::endl;

1181 }

1182 return false;

1183 }

1184 }

1185 else

1186 {

1187 constexpr auto KReadVec = math::lcm(AK1Number, BK1Number);

1188 auto K_t = karg.KBatch * KReadVec;

1189 auto KReadPadSplited = math::integer_divide_ceil(karg.K, K_t) * KReadVec;

1190 if((KReadPadSplited * (karg.KBatch - 1)) >= karg.K)

1191 {

1192 return false;

1193 }

1194 }

1195

1196 if constexpr(is_same<tensor_layout::gemm::RowMajor, ALayout>::value)

1197 {

1198 if(karg.K % ABlockTransferSrcScalarPerVector != 0)

1199 {

1200 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1201 {

1202 std::cout << "Arg K (" << karg.K

1203 << ") value is not a multiple of ABlockTransferSrcScalarPerVector ("

1204 << ABlockTransferSrcScalarPerVector << " )! " << __FILE__ << ":"

1205 << __LINE__ << ", in function: " << __func__ << std::endl;

1206 }

1207 return false;

1208 }

1209 }

1210 else

1211 {

1212 if(karg.M % ABlockTransferSrcScalarPerVector != 0)

1213 {

1214 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1215 {

1216 std::cout << "Arg M (" << karg.M

1217 << ") value is not a multiple of ABlockTransferSrcScalarPerVector ("

1218 << ABlockTransferSrcScalarPerVector << " )! " << __FILE__ << ":"

1219 << __LINE__ << ", in function: " << __func__ << std::endl;

1220 }

1221 return false;

1222 }

1223 }

1224

1225 if constexpr(is_same<tensor_layout::gemm::RowMajor, BLayout>::value)

1226 {

1227 if(karg.N % BBlockTransferSrcScalarPerVector != 0)

1228 {

1229 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1230 {

1231 std::cout << "Arg N (" << karg.N

1232 << ") value is not a multiple of BBlockTransferSrcScalarPerVector ("

1233 << BBlockTransferSrcScalarPerVector << " )! " << __FILE__ << ":"

1234 << __LINE__ << ", in function: " << __func__ << std::endl;

1235 }

1236 return false;

1237 }

1238 }

1239 else

1240 {

1241 if(karg.K % BBlockTransferSrcScalarPerVector != 0)

1242 {

1243 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1244 {

1245 std::cout << "Arg K (" << karg.K

1246 << ") value is not a multiple of BBlockTransferSrcScalarPerVector ("

1247 << BBlockTransferSrcScalarPerVector << " )! " << __FILE__ << ":"

1248 << __LINE__ << ", in function: " << __func__ << std::endl;

1249 }

1250 return false;

1251 }

1252 }

1253

1254 if constexpr(is_same<tensor_layout::gemm::RowMajor, CLayout>::value)

1255 {

1256 if(karg.N % CShuffleBlockTransferScalarPerVector_NPerBlock != 0)

1257 {

1258 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1259 {

1260 std::cout << "Arg N (" << karg.N

1261 << ") value is not a multiple of "

1262 "CShuffleBlockTransferScalarPerVector_NPerBlock ("

1263 << CShuffleBlockTransferScalarPerVector_NPerBlock << " )! "

1264 << __FILE__ << ":" << __LINE__ << ", in function: " << __func__

1265 << std::endl;

1266 }

1267 return false;

1268 }

1269 }

1270 else

1271 {

1272 if(karg.M % CShuffleBlockTransferScalarPerVector_NPerBlock != 0)

1273 {

1274 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1275 {

1276 std::cout << "Arg M (" << karg.M

1277 << ") value is not a multiple of "

1278 "CShuffleBlockTransferScalarPerVector_NPerBlock ("

1279 << CShuffleBlockTransferScalarPerVector_NPerBlock << " )! "

1280 << __FILE__ << ":" << __LINE__ << ", in function: " << __func__

1281 << std::endl;

1282 }

1283 return false;

1284 }

1285 }

1286

1287 if constexpr(!(is_same<remove_cvref_t<CDataType>, half_t>::value ||

1288 is_same<remove_cvref_t<CDataType>, float>::value ||

1289 is_same<remove_cvref_t<CDataType>, bhalf_t>::value ||

1290 is_same<remove_cvref_t<CDataType>, int32_t>::value))

1291 {

1292 if(!karg.IsReduceAdd())

1293 {

1294 if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))

1295 {

1296 std::cout << " KBatch: " << karg.KBatch << " > 1 is not support yet" << __FILE__

1297 << ":" << __LINE__ << ", in function: " << __func__ << std::endl;

1298 }

1299 if(karg.KBatch > 1)

1300 {

1301 return false;

1302 }

1303 }

1304 }

1305#if 0

1306 // check gridwise gemm pipeline

1307 const auto num_k_loop = karg.AK0 / (KPerBlock / AK1Value);

1308

1309 if constexpr(BlkGemmPipelineVer != BlockGemmPipelineVersion::v1)

1310 {

1311 if(num_k_loop <= BlockwiseGemmPipe::PrefetchStages)

1312 {

1313 return false;

1314 }

1315 }

1316#endif

1317 // TODO: also check validity of all components (blockwise-copy, threadwise-copy, etc)

1318 return true;

1319 }

1320

1321 __host__ static constexpr bool CalculateHasMainKBlockLoop(index_t K)

1322 {

1323 const index_t num_loop = K / KPerBlock;

1324

1325 return BlockwiseGemmPipe::BlockHasHotloop(num_loop);

1326 }

1327

1328 __host__ static constexpr TailNumber CalculateKBlockLoopTailNum(index_t K)

1329 {

1330 const index_t num_loop = K / KPerBlock;

1331

1332 return BlockwiseGemmPipe::BlockLoopTailNum(num_loop);

1333 }

1334

1335 template <typename CGridDesc>

1336 __host__ __device__ static constexpr auto MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(

1337 const CGridDesc& c_grid_desc_m_n, index_t MBlock, index_t NBlock)

1338 {

1339 const auto c_grid_desc_mblock_mperblock_nblock_nperblock = transform_tensor_descriptor(

1340 c_grid_desc_m_n,

1341 make_tuple(make_unmerge_transform(make_tuple(MBlock, Number<MPerBlock>{})),

1342 make_unmerge_transform(make_tuple(NBlock, Number<NPerBlock>{}))),

1343 make_tuple(Sequence<0>{}, Sequence<1>{}),

1344 make_tuple(Sequence<0, 1>{}, Sequence<2, 3>{}));

1345

1346 return c_grid_desc_mblock_mperblock_nblock_nperblock;

1347 }

1348

1349 // return block_id to C matrix tile idx (m0, n0) mapping

1350 // if arch = gfx942

1351 using Block2CTileMap = BlockToCTileMap_Grouped_M00_N0_M01Adapt<8, MPerBlock, NPerBlock>;

1352 // using Block2CTileMap = BlockToCTileMap_3DGrid_KSplit<MPerBlock, NPerBlock>;

1353

1354 using mx_scale_t = e8m0_bexp_t;

1355 static constexpr index_t scale_pack_size_a = sizeof(AScaleDataType) / sizeof(mx_scale_t);

1356 static constexpr index_t scale_pack_size_b = sizeof(BScaleDataType) / sizeof(mx_scale_t);

1357 static_assert(KXdlPack * MXdlPack % scale_pack_size_a == 0,

1358 "A scale pack data type too large!");

1359 static_assert(KXdlPack * NXdlPack % scale_pack_size_b == 0,

1360 "B scale pack data type too large!");

1361

1362 static_assert(is_same_v<AElementwiseOperation, tensor_operation::element_wise::PassThrough> &&

1363 is_same_v<BElementwiseOperation, tensor_operation::element_wise::PassThrough>,

1364 "A/B ElementwiseOperation should be PassThrough as load_to_lds is used!");

1365

1366 template <typename AGridDesc_AK0_M_K1,

1367 typename AScaleGridDesc_AM_AK,

1368 typename BGridDesc_BK0_N_K1,

1369 typename BScaleGridDesc_BN_AK,

1370 typename CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock,

1371 bool HasMainKBlockLoop,

1372 InMemoryDataOperationEnum CGlobalMemoryDataOperation,

1373 TailNumber TailNum = TailNumber::Odd>

1374 __device__ static void Run(const ADataType* p_a_grid,

1375 const AScaleDataType* p_a_scale_grid,

1376 const BDataType* p_b_grid,

1377 const BScaleDataType* p_b_scale_grid,

1378 CDataType* p_c_grid,

1379 void* p_shared,

1380 const Problem& problem,

1381 const AGridDesc_AK0_M_K1& a_grid_desc_ak0_m_ak1,

1382 const AScaleGridDesc_AM_AK& a_scale_grid_desc_am_ak,

1383 const BGridDesc_BK0_N_K1& b_grid_desc_bk0_n_bk1,

1384 const BScaleGridDesc_BN_AK& b_scale_grid_desc_bn_ak,

1385 const CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock&

1386 c_grid_desc_mblock_mperblock_nblock_nperblock)

1387 {

1388 const auto a_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1389 p_a_grid, a_grid_desc_ak0_m_ak1.GetElementSpaceSize());

1390 const auto b_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1391 p_b_grid, b_grid_desc_bk0_n_bk1.GetElementSpaceSize());

1392 auto c_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1393 p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());

1394

1395 // A Scale buffer

1396 const auto a_scale_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1397 p_a_scale_grid, a_scale_grid_desc_am_ak.GetElementSpaceSize());

1398

1399 // B Scale buffer

1400 const auto b_scale_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1401 p_b_scale_grid, b_scale_grid_desc_bn_ak.GetElementSpaceSize());

1402

1403 const CElementwiseOperation c_element_op{};

1404

1405 // divide block work by [M, N]

1406 const auto block_2_ctile_map = Block2CTileMap{problem.M, problem.N, 4};

1407

1408 const auto block_work_idx =

1409 block_2_ctile_map.CalculateBottomIndex(make_multi_index(get_block_1d_id()));

1410

1411 if(!block_2_ctile_map.ValidCTileIndex(

1412 block_work_idx,

1413 make_tuple(c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I0),

1414 c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I2))))

1415 {

1416 return;

1417 }

1418

1419 const index_t block_m_id = __builtin_amdgcn_readfirstlane(block_work_idx[I0]);

1420 const index_t block_n_id = __builtin_amdgcn_readfirstlane(block_work_idx[I1]);

1421

1422 // HACK: this force m/n_block_data_idx_on_grid into SGPR

1423 const index_t m_block_data_idx_on_grid =

1424 __builtin_amdgcn_readfirstlane(block_m_id * MPerBlock);

1425

1426 const index_t n_block_data_idx_on_grid =

1427 __builtin_amdgcn_readfirstlane(block_n_id * NPerBlock);

1428

1429 // lds max alignment

1430 constexpr auto max_lds_align = math::lcm(AK1Number, BK1Number);

1431

1432 // A matrix in LDS memory, dst of blockwise copy

1433 constexpr auto a_block_desc_ak0_m_ak1 = GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();

1434

1435 // B matrix in LDS memory, dst of blockwise copy

1436 constexpr auto b_block_desc_bk0_n_bk1 = GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();

1437

1438 auto a_blockwise_copy =

1439 ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,

1440 Sequence<AK0Number, MPerBlock, AK1Number>,

1441 ABlockTransferThreadClusterLengths_AK0_M_AK1,

1442 ABlockTransferThreadClusterArrangeOrder,

1443 ADataType,

1444 ADataType,

1445 decltype(a_grid_desc_ak0_m_ak1),

1446 decltype(a_block_desc_ak0_m_ak1),

1447 ABlockTransferSrcAccessOrder,

1448 ABlockTransferSrcVectorDim,

1449 2,

1450 ABlockTransferSrcScalarPerVector>(

1451 a_grid_desc_ak0_m_ak1,

1452 make_multi_index(0, m_block_data_idx_on_grid, 0),

1453 a_block_desc_ak0_m_ak1,

1454 make_multi_index(0, 0, 0));

1455

1456 // B matrix blockwise copy

1457 auto b_blockwise_copy =

1458 ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,

1459 Sequence<BK0Number, NPerBlock, BK1Number>,

1460 BBlockTransferThreadClusterLengths_BK0_N_BK1,

1461 BBlockTransferThreadClusterArrangeOrder,

1462 BDataType,

1463 BDataType,

1464 decltype(b_grid_desc_bk0_n_bk1),

1465 decltype(b_block_desc_bk0_n_bk1),

1466 BBlockTransferSrcAccessOrder,

1467 BBlockTransferSrcVectorDim,

1468 2,

1469 BBlockTransferSrcScalarPerVector>(

1470 b_grid_desc_bk0_n_bk1,

1471 make_multi_index(0, n_block_data_idx_on_grid, 0),

1472 b_block_desc_bk0_n_bk1,

1473 make_multi_index(0, 0, 0));

1474

1475 // LDS allocation for A and B: be careful of alignment

1476 constexpr auto a_block_space_size_aligned = math::integer_least_multiple(

1477 a_block_desc_ak0_m_ak1.GetElementSpaceSize(), max_lds_align);

1478

1479 // Cast after lds

1480 auto a_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(

1481 static_cast<ADataType*>(p_shared), a_block_desc_ak0_m_ak1.GetElementSpaceSize());

1482

1483 auto b_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(

1484 reinterpret_cast<BDataType*>(static_cast<char*>(p_shared) +

1485 a_block_space_size_aligned * sizeof(ADataType)),

1486 b_block_desc_bk0_n_bk1.GetElementSpaceSize());

1487

1488 constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock / AK1Number, 0, 0);

1489 constexpr auto b_block_slice_copy_step = make_multi_index(KPerBlock / BK1Number, 0, 0);

1490

1491 // Blockwise GEMM pipeline

1492 static_assert(std::is_default_constructible_v<BlockwiseGemmPipe>);

1493 auto blockwise_gemm_pipeline = BlockwiseGemmPipe{};

1494 auto c_thread_buf = blockwise_gemm_pipeline.GetCThreadBuffer();

1495

1496 const index_t num_k_block_main_loop = __builtin_amdgcn_readfirstlane(

1497 (a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2)) /

1498 KPerBlock);

1499

1500 // Initial thread mapping for:

1501 // BlockSize = 256

1502 // MPerXdl=NPerXdl=32 and MPerBlock=NPerBlock=128 MRepeat=NRepeat=2 MWaves=NWaves=2

1503 // For each [m0, n0] tile, there are 4 waves:

1504 // tId in [ 0, 63] m x n = [ 0, 31] x [ 0, 31] waveId = [0, 0]

1505 // tId in [ 64, 127] m x n = [ 0, 31] x [32, 63] waveId = [0, 1]

1506 // tId in [128, 191] m x n = [32, 63] x [ 0, 31] waveId = [1, 0]

1507 // tId in [192, 255] m x n = [32, 63] x [32, 63] waveId = [1, 1]

1508

1509 // BlockSize = 128

1510 // MPerXdl=NPerXdl=16 and MPerBlock=128 NPerBlock=16 MRepeat=4 NRepeat=1 MWaves=2 NWaves=1

1511 // For each [m0, n0] tile, there are 2 waves:

1512 // tId in [ 0, 63] m x n = [ 0, 15] x [0, 15] waveId = [0, 0]

1513 // tId in [ 64, 127] m x n = [16, 31] x [0, 15] waveId = [1, 0]

1514

1515 // TODO: Document initial thread mapping for more combinations of parameters

1516

1517 const auto wave_idx = BlockwiseGemmPipe::GetWaveIdx();

1518 const auto waveId_m = wave_idx[I0];

1519 const auto waveId_n = wave_idx[I1];

1520

1521 // static constexpr auto mfma = BlockwiseGemmPipe::xdlops_gemm.mfma;

1522

1523 // auto thread_offset_k = (get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize) /

1524 // mfma.selected_mfma.num_threads_per_blk;

1525

1526 // A wave access continuous memory

1527 auto thread_offset_shuffled =

1528 get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize * KXdlPack * MXdlPack;

1529

1530 auto a_thread_offset_m = waveId_m;

1531

1532 auto a_scale_thread_copy = ThreadwiseTensorSliceTransfer_v2<

1533 AScaleDataType,

1534 AScaleDataType,

1535 decltype(a_scale_grid_desc_am_ak),

1536 decltype(BlockwiseGemmPipe::a_scale_thread_desc),

1537 Sequence<1, 1, KXdlPack * MXdlPack / scale_pack_size_a>, // SliceLengths

1538 Sequence<0, 1, 2>, // DimAccessOrder

1539 2, // SrcVectorDim

1540 KXdlPack * MXdlPack / scale_pack_size_a, // SrcScalarPerVector

1541 1, // SrcScalarStrideInVector

1542 true>(a_scale_grid_desc_am_ak,

1543 make_multi_index(block_m_id * MPerBlock / MPerXdl / MXdlPack + a_thread_offset_m,

1544 0,

1545 thread_offset_shuffled / scale_pack_size_a));

1546

1547 auto b_thread_offset_n = waveId_n;

1548

1549 auto b_scale_thread_copy = ThreadwiseTensorSliceTransfer_v2<

1550 BScaleDataType,

1551 BScaleDataType,

1552 decltype(b_scale_grid_desc_bn_ak),

1553 decltype(BlockwiseGemmPipe::b_scale_thread_desc),

1554 Sequence<1, 1, KXdlPack * NXdlPack / scale_pack_size_b>, // SliceLengths

1555 Sequence<0, 1, 2>, // DimAccessOrder

1556 2, // SrcVectorDim

1557 KXdlPack * MXdlPack / scale_pack_size_b, // SrcScalarPerVector

1558 1, // SrcScalarStrideInVector

1559 true>(b_scale_grid_desc_bn_ak,

1560 make_multi_index(block_n_id * NPerBlock / NPerXdl / NXdlPack + b_thread_offset_n,

1561 0,

1562 thread_offset_shuffled / scale_pack_size_b));

1563

1564 blockwise_gemm_pipeline.template Run<HasMainKBlockLoop, TailNum>(a_grid_desc_ak0_m_ak1,

1565 a_block_desc_ak0_m_ak1,

1566 a_blockwise_copy,

1567 a_grid_buf,

1568 a_block_buf,

1569 a_block_slice_copy_step,

1570 b_grid_desc_bk0_n_bk1,

1571 b_block_desc_bk0_n_bk1,

1572 b_blockwise_copy,

1573 b_grid_buf,

1574 b_block_buf,

1575 b_block_slice_copy_step,

1576 c_thread_buf,

1577 a_scale_grid_desc_am_ak,

1578 a_scale_thread_copy,

1579 a_scale_grid_buf,

1580 b_scale_grid_desc_bn_ak,

1581 b_scale_thread_copy,

1582 b_scale_grid_buf,

1583 num_k_block_main_loop);

1584

1585 // shuffle C and write out

1586 {

1587 static_assert(MXdlPerWave % CShuffleMXdlPerWavePerShuffle == 0 &&

1588 NXdlPerWave % CShuffleNXdlPerWavePerShuffle == 0,

1589 "wrong!");

1590 static_assert(CShuffleMXdlPerWavePerShuffle % MXdlPack == 0 &&

1591 CShuffleNXdlPerWavePerShuffle % NXdlPack == 0,

1592 "wrong!");

1593

1594 constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl);

1595 constexpr index_t NWave = NPerBlock / (NXdlPerWave * NPerXdl);

1596

1597 // TODO: hacky, fix it!

1598 constexpr auto c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2 =

1599 blockwise_gemm_pipeline.GetCThreadDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3();

1600

1601 // TODO: hacky, fix it!

1602 // c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp is only used to get lengths

1603 constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp =

1604 blockwise_gemm_pipeline.GetCBlockDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3();

1605

1606 constexpr auto M0 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I0);

1607 constexpr auto N0 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I1);

1608 constexpr auto M1 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I2);

1609 constexpr auto N1 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I3);

1610 constexpr auto M2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I4);

1611 constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I5);

1612 constexpr auto M3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I6);

1613 constexpr auto M4 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7);

1614 constexpr auto M5 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I8);

1615 constexpr auto N3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I9);

1616

1617 constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock =

1618 GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock();

1619

1620 auto c_shuffle_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(

1621 static_cast<CShuffleDataType*>(p_shared),

1622 c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());

1623

1624 constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2 = transform_tensor_descriptor(

1625 c_shuffle_block_desc_mblock_mperblock_nblock_nperblock,

1626 make_tuple(

1627 make_freeze_transform(I0),

1628 make_unmerge_transform(make_tuple(

1629 Number<CShuffleMXdlPerWavePerShuffle / MXdlPack>{}, // M0 (MXdlPerWave) per

1630 // shuffle

1631 M1, // M1 = MWave

1632 M2, // M2 = MXdlPack

1633 M3, // M3 * M4 * M5 = MPerXdl

1634 M4,

1635 M5)),

1636 make_freeze_transform(I0),

1637 make_unmerge_transform(make_tuple(

1638 Number<CShuffleNXdlPerWavePerShuffle / NXdlPack>{}, // N0 (NXdlPerWave) per

1639 // shuffle

1640 N1, // N1 = NWave

1641 N2, // N2 = NXdlPack

1642 N3))), // N3 = NPerXdl

1643 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),

1644 make_tuple(Sequence<>{},

1645 Sequence<0, 2, 4, 6, 7, 8>{},

1646 Sequence<>{},

1647 Sequence<1, 3, 5, 9>{}));

1648

1649 // calculate origin of thread output tensor on global memory

1650 // blockwise GEMM c matrix starting index

1651 const auto c_thread_mtx_on_block =

1652 blockwise_gemm_pipeline.CalculateCThreadOriginDataIndex(I0, I0, I0, I0);

1653

1654 const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0];

1655 const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1];

1656

1657 const auto m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor =

1658 make_single_stage_tensor_adaptor(

1659 make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4, M5))),

1660 make_tuple(Sequence<0, 1, 2, 3, 4, 5>{}),

1661 make_tuple(Sequence<0>{}));

1662

1663 const auto m_thread_data_on_block_idx =

1664 m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor.CalculateBottomIndex(

1665 make_multi_index(m_thread_data_on_block));

1666

1667 const auto n_thread_data_on_block_to_n0_n1_n2_adaptor =

1668 make_single_stage_tensor_adaptor(

1669 make_tuple(make_merge_transform(make_tuple(N0, N1, N2, N3))),

1670 make_tuple(Sequence<0, 1, 2, 3>{}),

1671 make_tuple(Sequence<0>{}));

1672

1673 const auto n_thread_data_on_block_idx =

1674 n_thread_data_on_block_to_n0_n1_n2_adaptor.CalculateBottomIndex(

1675 make_multi_index(n_thread_data_on_block));

1676

1677 // shuffle: threadwise copy C from VGPR to LDS

1678 auto c_thread_copy_vgpr_to_lds = ThreadwiseTensorSliceTransfer_v1r3<

1679 AccDataType,

1680 CShuffleDataType,

1681 decltype(c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2),

1682 decltype(c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2),

1683 ck::tensor_operation::element_wise::PassThrough,

1684 Sequence<CShuffleMXdlPerWavePerShuffle / MXdlPack,

1685 CShuffleNXdlPerWavePerShuffle / NXdlPack,

1686 I1,

1687 I1,

1688 M2,

1689 N2,

1690 M3,

1691 I1,

1692 M5,

1693 I1>,

1694 Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>,

1695 9,

1696 1,

1697 InMemoryDataOperationEnum::Set,

1698 1,

1699 true>{c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2,

1700 make_multi_index(0,

1701 0,

1702 m_thread_data_on_block_idx[I1],

1703 n_thread_data_on_block_idx[I1],

1704 m_thread_data_on_block_idx[I2],

1705 n_thread_data_on_block_idx[I2],

1706 m_thread_data_on_block_idx[I3],

1707 m_thread_data_on_block_idx[I4],

1708 m_thread_data_on_block_idx[I5],

1709 n_thread_data_on_block_idx[I3]),

1710 ck::tensor_operation::element_wise::PassThrough{}};

1711

1712 // shuffle: blockwise copy C from LDS to global

1713 auto c_shuffle_block_copy_lds_to_global = ThreadGroupTensorSliceTransfer_v6r1<

1714 ThisThreadBlock, // ThreadGroup

1715 CElementwiseOperation, // ElementwiseOperation,

1716 CGlobalMemoryDataOperation, // DstInMemOp,

1717 Sequence<1,

1718 CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,

1719 1,

1720 CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths,

1721 CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,

1722 Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder,

1723 CShuffleDataType, // typename SrcData,

1724 CDataType, // typename DstData,

1725 decltype(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock),

1726 decltype(c_grid_desc_mblock_mperblock_nblock_nperblock),

1727 Sequence<0, 1, 2, 3>, // typename DimAccessOrder,

1728 3, // index_t VectorDim,

1729 CShuffleBlockTransferScalarPerVector_NPerBlock, // index_t ScalarPerVector,

1730 true, // bool ThreadTransferSrcResetCoordinateAfterRun,

1731 false> // bool ThreadTransferDstResetCoordinateAfterRun>

1732 {c_shuffle_block_desc_mblock_mperblock_nblock_nperblock,

1733 make_multi_index(0, 0, 0, 0),

1734 c_grid_desc_mblock_mperblock_nblock_nperblock,

1735 make_multi_index(block_m_id, 0, block_n_id, 0),

1736 c_element_op};

1737

1738 // space filling curve for threadwise C in VGPR

1739 constexpr auto sfc_c_vgpr =

1740 SpaceFillingCurve<Sequence<MXdlPerWave / MXdlPack,

1741 NXdlPerWave / NXdlPack,

1742 1,

1743 1,

1744 MXdlPack,

1745 NXdlPack,

1746 M2,

1747 1,

1748 M4,

1749 1>,

1750 Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>,

1751 Sequence<CShuffleMXdlPerWavePerShuffle / MXdlPack,

1752 CShuffleNXdlPerWavePerShuffle / NXdlPack,

1753 1,

1754 1,

1755 MXdlPack,

1756 NXdlPack,

1757 M2,

1758 1,

1759 M4,

1760 1>>{};

1761

1762 // space filling curve for shuffled blockwise C in global mem

1763 constexpr auto sfc_c_global =

1764 SpaceFillingCurve<Sequence<1, MPerBlock, 1, NPerBlock>,

1765 Sequence<0, 2, 1, 3>,

1766 Sequence<1,

1767 CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,

1768 1,

1769 CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>{};

1770

1771 constexpr index_t num_access = sfc_c_vgpr.GetNumOfAccess();

1772

1773 static_assert(num_access == sfc_c_global.GetNumOfAccess(), "wrong!");

1774

1775 static_for<0, num_access, 1>{}([&](auto access_id) {

1776 // make sure it's safe to write to LDS

1777 block_sync_lds();

1778

1779 // each thread write its data from VGPR to LDS

1780 c_thread_copy_vgpr_to_lds.Run(c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2,

1781 sfc_c_vgpr.GetIndexTupleOfNumber(access_id),

1782 c_thread_buf,

1783 c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2,

1784 c_shuffle_block_buf);

1785

1786 // make sure it's safe to read from LDS

1787 block_sync_lds();

1788

1789 // each block copy its data from LDS to global

1790 c_shuffle_block_copy_lds_to_global.Run(

1791 c_shuffle_block_desc_mblock_mperblock_nblock_nperblock,

1792 c_shuffle_block_buf,

1793 c_grid_desc_mblock_mperblock_nblock_nperblock,

1794 c_grid_buf);

1795

1796 if constexpr(access_id < num_access - 1)

1797 {

1798 constexpr auto c_global_step = sfc_c_global.GetForwardStep(access_id);

1799

1800 // move on C

1801 c_shuffle_block_copy_lds_to_global.MoveDstSliceWindow(

1802 c_grid_desc_mblock_mperblock_nblock_nperblock, c_global_step);

1803 }

1804 });

1805 }

1806 }

1807

1808 template <bool HasMainKBlockLoop,

1809 InMemoryDataOperationEnum CGlobalMemoryDataOperation,

1810 TailNumber TailNum = TailNumber::Odd>

1811 __device__ static void Run(const ADataType* p_a_grid,

1812 const AScaleDataType* p_a_scale_grid,

1813 const BDataType* p_b_grid,

1814 const BScaleDataType* p_b_scale_grid,

1815 CDataType* p_c_grid,

1816 void* p_shared,

1817 const Problem& problem)

1818 {

1819 const auto a_grid_desc_ak0_m_ak1 = MakeAGridDescriptor_AK0_M_AK1(

1820 problem.M, problem.MPadded, problem.K, problem.KPadded, problem.StrideA, problem.AK0);

1821 const auto b_grid_desc_bk0_n_bk1 = MakeBGridDescriptor_BK0_N_BK1(

1822 problem.K, problem.KPadded, problem.N, problem.NPadded, problem.StrideB, problem.BK0);

1823 const auto c_grid_desc_m_n = MakeCGridDescriptor_M_N(

1824 problem.M, problem.MPadded, problem.N, problem.NPadded, problem.StrideC);

1825 const auto c_grid_desc_mblock_mperblock_nblock_nperblock =

1826 MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(

1827 c_grid_desc_m_n, problem.MBlock, problem.NBlock);

1828

1829 // A/B shuffled scale for better 8-bit scale access pattern

1830 // MNRepeat -> KRepeat -> KThreadPerXdl -> MNThreadPerXdl -> KXdlPack -> MNXdlPack

1831 const auto Padded_Scale_M =

1832 math::integer_divide_ceil(problem.M, ScaleBlockSize) * ScaleBlockSize;

1833 const auto a_scale_grid_desc_am_ak = make_naive_tensor_descriptor(

1834 make_tuple(Padded_Scale_M / (MXdlPack * MPerXdl),

1835 math::integer_divide_ceil(problem.K, (ScaleBlockSize / APackedSize)) /

1836 (KXdlPack * 64 / MPerXdl),

1837 64 * KXdlPack * MXdlPack / scale_pack_size_a),

1838 make_tuple(math::integer_divide_ceil(problem.K * problem.KBatch,

1839 (ScaleBlockSize / APackedSize)) *

1840 MPerXdl * MXdlPack / scale_pack_size_a,

1841 64 * KXdlPack * MXdlPack / scale_pack_size_a,

1842 1));

1843

1844 const auto b_scale_grid_desc_bn_ak = make_naive_tensor_descriptor(

1845 make_tuple(problem.N / (NXdlPack * NPerXdl),

1846 math::integer_divide_ceil(problem.K, (ScaleBlockSize / BPackedSize)) /

1847 (KXdlPack * 64 / NPerXdl),

1848 64 * KXdlPack * NXdlPack / scale_pack_size_b),

1849 make_tuple(math::integer_divide_ceil(problem.K * problem.KBatch,

1850 (ScaleBlockSize / BPackedSize)) *

1851 NPerXdl * NXdlPack / scale_pack_size_b,

1852 64 * KXdlPack * NXdlPack / scale_pack_size_b,

1853 1));

1854

1855 Run<decltype(a_grid_desc_ak0_m_ak1),

1856 decltype(a_scale_grid_desc_am_ak),

1857 decltype(b_grid_desc_bk0_n_bk1),

1858 decltype(b_scale_grid_desc_bn_ak),

1859 decltype(c_grid_desc_mblock_mperblock_nblock_nperblock),

1860 HasMainKBlockLoop,

1861 CGlobalMemoryDataOperation,

1862 TailNum>(p_a_grid,

1863 p_a_scale_grid,

1864 p_b_grid,

1865 p_b_scale_grid,

1866 p_c_grid,

1867 p_shared,

1868 problem,

1869 a_grid_desc_ak0_m_ak1,

1870 a_scale_grid_desc_am_ak,

1871 b_grid_desc_bk0_n_bk1,

1872 b_scale_grid_desc_bn_ak,

1873 c_grid_desc_mblock_mperblock_nblock_nperblock);

1874 }

1875

1876 template <typename AGridDesc_AK0_M_K1,

1877 typename AScaleGridDesc_AM_AK,

1878 typename BGridDesc_BK0_N_K1,

1879 typename BScaleGridDesc_BN_AK,

1880 typename CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock,

1881 bool HasMainKBlockLoop,

1882 InMemoryDataOperationEnum CGlobalMemoryDataOperation,

1883 TailNumber TailNum = TailNumber::Odd>

1884 __device__ static void Run_2Lds(const ADataType* p_a_grid,

1885 const AScaleDataType* p_a_scale_grid,

1886 const BDataType* p_b_grid,

1887 const BScaleDataType* p_b_scale_grid,

1888 CDataType* p_c_grid,

1889 void* p_shared_0,

1890 void* p_shared_1,

1891 const Problem& problem,

1892 const AGridDesc_AK0_M_K1& a_grid_desc_ak0_m_ak1,

1893 const AScaleGridDesc_AM_AK& a_scale_grid_desc_am_ak,

1894 const BGridDesc_BK0_N_K1& b_grid_desc_bk0_n_bk1,

1895 const BScaleGridDesc_BN_AK& b_scale_grid_desc_bn_ak,

1896 const CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock&

1897 c_grid_desc_mblock_mperblock_nblock_nperblock)

1898 {

1899 const auto a_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1900 p_a_grid, a_grid_desc_ak0_m_ak1.GetElementSpaceSize());

1901 const auto b_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1902 p_b_grid, b_grid_desc_bk0_n_bk1.GetElementSpaceSize());

1903 auto c_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1904 p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());

1905

1906 // A Scale buffer

1907 const auto a_scale_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1908 p_a_scale_grid, a_scale_grid_desc_am_ak.GetElementSpaceSize());

1909

1910 // B Scale buffer

1911 const auto b_scale_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(

1912 p_b_scale_grid, b_scale_grid_desc_bn_ak.GetElementSpaceSize());

1913

1914 const CElementwiseOperation c_element_op{};

1915

1916 // divide block work by [M, N]

1917 const auto block_2_ctile_map = Block2CTileMap{problem.M, problem.N, 4};

1918

1919 const auto block_work_idx =

1920 block_2_ctile_map.CalculateBottomIndex(make_multi_index(get_block_1d_id()));

1921

1922 if(!block_2_ctile_map.ValidCTileIndex(

1923 block_work_idx,

1924 make_tuple(c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I0),

1925 c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I2))))

1926 {

1927 return;

1928 }

1929

1930 const index_t block_m_id = __builtin_amdgcn_readfirstlane(block_work_idx[I0]);

1931 const index_t block_n_id = __builtin_amdgcn_readfirstlane(block_work_idx[I1]);

1932

1933 // HACK: this force m/n_block_data_idx_on_grid into SGPR

1934 const index_t m_block_data_idx_on_grid =

1935 __builtin_amdgcn_readfirstlane(block_m_id * MPerBlock);

1936

1937 const index_t n_block_data_idx_on_grid =

1938 __builtin_amdgcn_readfirstlane(block_n_id * NPerBlock);

1939

1940 // lds max alignment

1941 constexpr auto max_lds_align = math::lcm(AK1Number, BK1Number);

1942

1943 // A matrix in LDS memory, dst of blockwise copy

1944 constexpr auto a_block_desc_ak0_m_ak1 = GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();

1945

1946 // B matrix in LDS memory, dst of blockwise copy

1947 constexpr auto b_block_desc_bk0_n_bk1 = GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();

1948

1949 auto a_blockwise_copy =

1950 ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,

1951 Sequence<AK0Number, MPerBlock, AK1Number>,

1952 ABlockTransferThreadClusterLengths_AK0_M_AK1,

1953 ABlockTransferThreadClusterArrangeOrder,

1954 ADataType,

1955 ADataType,

1956 decltype(a_grid_desc_ak0_m_ak1),

1957 decltype(a_block_desc_ak0_m_ak1),

1958 ABlockTransferSrcAccessOrder,

1959 ABlockTransferSrcVectorDim,

1960 2,

1961 ABlockTransferSrcScalarPerVector>(

1962 a_grid_desc_ak0_m_ak1,

1963 make_multi_index(0, m_block_data_idx_on_grid, 0),

1964 a_block_desc_ak0_m_ak1,

1965 make_multi_index(0, 0, 0));

1966

1967 // B matrix blockwise copy

1968 auto b_blockwise_copy =

1969 ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,

1970 Sequence<BK0Number, NPerBlock, BK1Number>,

1971 BBlockTransferThreadClusterLengths_BK0_N_BK1,

1972 BBlockTransferThreadClusterArrangeOrder,

1973 BDataType,

1974 BDataType,

1975 decltype(b_grid_desc_bk0_n_bk1),

1976 decltype(b_block_desc_bk0_n_bk1),

1977 BBlockTransferSrcAccessOrder,

1978 BBlockTransferSrcVectorDim,

1979 2,

1980 BBlockTransferSrcScalarPerVector>(

1981 b_grid_desc_bk0_n_bk1,

1982 make_multi_index(0, n_block_data_idx_on_grid, 0),

1983 b_block_desc_bk0_n_bk1,

1984 make_multi_index(0, 0, 0));

1985

1986 // LDS allocation for A and B: be careful of alignment

1987 constexpr auto a_block_space_size_aligned = math::integer_least_multiple(

1988 a_block_desc_ak0_m_ak1.GetElementSpaceSize(), max_lds_align);

1989

1990 auto a_block_buf_ping = make_dynamic_buffer<AddressSpaceEnum::Lds>(

1991 static_cast<ADataType*>(p_shared_0), a_block_desc_ak0_m_ak1.GetElementSpaceSize());

1992

1993 auto b_block_buf_ping = make_dynamic_buffer<AddressSpaceEnum::Lds>(

1994 bit_cast<BDataType*>(static_cast<char*>(p_shared_0) +

1995 a_block_space_size_aligned * sizeof(ADataType)),

1996 b_block_desc_bk0_n_bk1.GetElementSpaceSize());

1997

1998 auto a_block_buf_pong = make_dynamic_buffer<AddressSpaceEnum::Lds>(

1999 static_cast<ADataType*>(p_shared_1), a_block_desc_ak0_m_ak1.GetElementSpaceSize());

2000

2001 auto b_block_buf_pong = make_dynamic_buffer<AddressSpaceEnum::Lds>(

2002 bit_cast<BDataType*>(bit_cast<char*>(p_shared_1) +

2003 a_block_space_size_aligned * sizeof(ADataType)),

2004 b_block_desc_bk0_n_bk1.GetElementSpaceSize());

2005

2006 auto a_block_bufs = make_tuple(a_block_buf_ping, a_block_buf_pong);

2007 auto b_block_bufs = make_tuple(b_block_buf_ping, b_block_buf_pong);

2008

2009 constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock / AK1Number, 0, 0);

2010 constexpr auto b_block_slice_copy_step = make_multi_index(KPerBlock / BK1Number, 0, 0);

2011

2012 // Blockwise GEMM pipeline

2013 static_assert(std::is_default_constructible_v<BlockwiseGemmPipe>);

2014 auto blockwise_gemm_pipeline = BlockwiseGemmPipe{};

2015 auto c_thread_buf = blockwise_gemm_pipeline.GetCThreadBuffer();

2016

2017 const index_t num_k_block_main_loop = __builtin_amdgcn_readfirstlane(

2018 (a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2)) /

2019 KPerBlock);

2020

2021 // Initial thread mapping for:

2022 // BlockSize = 256

2023 // MPerXdl=NPerXdl=32 and MPerBlock=NPerBlock=128 MRepeat=NRepeat=2 MWaves=NWaves=2

2024 // For each [m0, n0] tile, there are 4 waves:

2025 // tId in [ 0, 63] m x n = [ 0, 31] x [ 0, 31] waveId = [0, 0]

2026 // tId in [ 64, 127] m x n = [ 0, 31] x [32, 63] waveId = [0, 1]

2027 // tId in [128, 191] m x n = [32, 63] x [ 0, 31] waveId = [1, 0]

2028 // tId in [192, 255] m x n = [32, 63] x [32, 63] waveId = [1, 1]

2029

2030 // BlockSize = 128

2031 // MPerXdl=NPerXdl=16 and MPerBlock=128 NPerBlock=16 MRepeat=4 NRepeat=1 MWaves=2 NWaves=1

2032 // For each [m0, n0] tile, there are 2 waves:

2033 // tId in [ 0, 63] m x n = [ 0, 15] x [0, 15] waveId = [0, 0]

2034 // tId in [ 64, 127] m x n = [16, 31] x [0, 15] waveId = [1, 0]

2035

2036 // TODO: Document initial thread mapping for more combinations of parameters

2037

2038 const auto wave_idx = BlockwiseGemmPipe::GetWaveIdx();

2039 const auto waveId_m = wave_idx[I0];

2040 const auto waveId_n = wave_idx[I1];

2041

2042 // static constexpr auto mfma = BlockwiseGemmPipe::xdlops_gemm.mfma;

2043

2044 // auto thread_offset_k = (get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize) /

2045 // mfma.selected_mfma.num_threads_per_blk;

2046

2047 // A wave access continuous memory

2048 auto thread_offset_shuffled =

2049 get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize * KXdlPack * MXdlPack;

2050

2051 auto a_thread_offset_m = waveId_m;

2052

2053 auto a_scale_thread_copy = ThreadwiseTensorSliceTransfer_v2<

2054 AScaleDataType,

2055 AScaleDataType,

2056 decltype(a_scale_grid_desc_am_ak),

2057 decltype(BlockwiseGemmPipe::a_scale_thread_desc),

2058 Sequence<1, 1, KXdlPack * MXdlPack / scale_pack_size_a>, // SliceLengths

2059 Sequence<0, 1, 2>, // DimAccessOrder

2060 2, // SrcVectorDim

2061 KXdlPack * MXdlPack / scale_pack_size_a, // SrcScalarPerVector

2062 1, // SrcScalarStrideInVector

2063 true>(a_scale_grid_desc_am_ak,

2064 make_multi_index(block_m_id * MPerBlock / MPerXdl / MXdlPack + a_thread_offset_m,

2065 0,

2066 thread_offset_shuffled / scale_pack_size_a));

2067

2068 auto b_thread_offset_n = waveId_n;

2069

2070 auto b_scale_thread_copy = ThreadwiseTensorSliceTransfer_v2<

2071 BScaleDataType,

2072 BScaleDataType,

2073 decltype(b_scale_grid_desc_bn_ak),

2074 decltype(BlockwiseGemmPipe::b_scale_thread_desc),

2075 Sequence<1, 1, KXdlPack * NXdlPack / scale_pack_size_b>, // SliceLengths

2076 Sequence<0, 1, 2>, // DimAccessOrder

2077 2, // SrcVectorDim

2078 KXdlPack * MXdlPack / scale_pack_size_b, // SrcScalarPerVector

2079 1, // SrcScalarStrideInVector

2080 true>(b_scale_grid_desc_bn_ak,

2081 make_multi_index(block_n_id * NPerBlock / NPerXdl / NXdlPack + b_thread_offset_n,

2082 0,

2083 thread_offset_shuffled / scale_pack_size_b));

2084

2085 blockwise_gemm_pipeline.template Run<HasMainKBlockLoop, TailNum>(a_grid_desc_ak0_m_ak1,

2086 a_block_desc_ak0_m_ak1,

2087 a_blockwise_copy,

2088 a_grid_buf,

2089 a_block_bufs,

2090 a_block_slice_copy_step,

2091 b_grid_desc_bk0_n_bk1,

2092 b_block_desc_bk0_n_bk1,

2093 b_blockwise_copy,

2094 b_grid_buf,

2095 b_block_bufs,

2096 b_block_slice_copy_step,

2097 c_thread_buf,

2098 a_scale_grid_desc_am_ak,

2099 a_scale_thread_copy,

2100 a_scale_grid_buf,

2101 b_scale_grid_desc_bn_ak,

2102 b_scale_thread_copy,

2103 b_scale_grid_buf,

2104 num_k_block_main_loop);

2105

2106 // shuffle C and write out

2107 {

2108 static_assert(MXdlPerWave % CShuffleMXdlPerWavePerShuffle == 0 &&

2109 NXdlPerWave % CShuffleNXdlPerWavePerShuffle == 0,

2110 "wrong!");

2111 static_assert(CShuffleMXdlPerWavePerShuffle % MXdlPack == 0 &&

2112 CShuffleNXdlPerWavePerShuffle % NXdlPack == 0,

2113 "wrong!");

2114

2115 constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl);

2116 constexpr index_t NWave = NPerBlock / (NXdlPerWave * NPerXdl);

2117

2118 // TODO: hacky, fix it!

2119 constexpr auto c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2 =

2120 blockwise_gemm_pipeline.GetCThreadDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3();

2121

2122 // TODO: hacky, fix it!

2123 // c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp is only used to get lengths

2124 constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp =

2125 blockwise_gemm_pipeline.GetCBlockDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3();

2126

2127 constexpr auto M0 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I0);

2128 constexpr auto N0 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I1);

2129 constexpr auto M1 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I2);

2130 constexpr auto N1 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I3);

2131 constexpr auto M2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I4);

2132 constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I5);

2133 constexpr auto M3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I6);

2134 constexpr auto M4 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7);

2135 constexpr auto M5 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I8);

2136 constexpr auto N3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I9);

2137

2138 constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock =

2139 GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock();

2140

2141 auto c_shuffle_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(

2142 static_cast<CShuffleDataType*>(p_shared_0),

2143 c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());

2144

2145 constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2 = transform_tensor_descriptor(

2146 c_shuffle_block_desc_mblock_mperblock_nblock_nperblock,

2147 make_tuple(

2148 make_freeze_transform(I0),

2149 make_unmerge_transform(make_tuple(

2150 Number<CShuffleMXdlPerWavePerShuffle / MXdlPack>{}, // M0 (MXdlPerWave) per

2151 // shuffle

2152 M1, // M1 = MWave

2153 M2, // M2 = MXdlPack

2154 M3, // M3 * M4 * M5 = MPerXdl

2155 M4,

2156 M5)),

2157 make_freeze_transform(I0),

2158 make_unmerge_transform(make_tuple(

2159 Number<CShuffleNXdlPerWavePerShuffle / NXdlPack>{}, // N0 (NXdlPerWave) per

2160 // shuffle

2161 N1, // N1 = NWave

2162 N2, // N2 = NXdlPack

2163 N3))), // N3 = NPerXdl

2164 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),

2165 make_tuple(Sequence<>{},

2166 Sequence<0, 2, 4, 6, 7, 8>{},

2167 Sequence<>{},

2168 Sequence<1, 3, 5, 9>{}));

2169

2170 // calculate origin of thread output tensor on global memory

2171 // blockwise GEMM c matrix starting index

2172 const auto c_thread_mtx_on_block =

2173 blockwise_gemm_pipeline.CalculateCThreadOriginDataIndex(I0, I0, I0, I0);

2174

2175 const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0];

2176 const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1];

2177

2178 const auto m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor =

2179 make_single_stage_tensor_adaptor(

2180 make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4, M5))),

2181 make_tuple(Sequence<0, 1, 2, 3, 4, 5>{}),

2182 make_tuple(Sequence<0>{}));

2183

2184 const auto m_thread_data_on_block_idx =

2185 m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor.CalculateBottomIndex(

2186 make_multi_index(m_thread_data_on_block));

2187

2188 const auto n_thread_data_on_block_to_n0_n1_n2_adaptor =

2189 make_single_stage_tensor_adaptor(

2190 make_tuple(make_merge_transform(make_tuple(N0, N1, N2, N3))),

2191 make_tuple(Sequence<0, 1, 2, 3>{}),

2192 make_tuple(Sequence<0>{}));

2193

2194 const auto n_thread_data_on_block_idx =

2195 n_thread_data_on_block_to_n0_n1_n2_adaptor.CalculateBottomIndex(

2196 make_multi_index(n_thread_data_on_block));

2197

2198 // shuffle: threadwise copy C from VGPR to LDS

2199 auto c_thread_copy_vgpr_to_lds = ThreadwiseTensorSliceTransfer_v1r3<

2200 AccDataType,

2201 CShuffleDataType,

2202 decltype(c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2),

2203 decltype(c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2),

2204 ck::tensor_operation::element_wise::PassThrough,

2205 Sequence<CShuffleMXdlPerWavePerShuffle / MXdlPack,

2206 CShuffleNXdlPerWavePerShuffle / NXdlPack,

2207 I1,

2208 I1,

2209 M2,

2210 N2,

2211 M3,

2212 I1,

2213 M5,

2214 I1>,

2215 Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>,

2216 9,

2217 1,

2218 InMemoryDataOperationEnum::Set,

2219 1,

2220 true>{c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2,

2221 make_multi_index(0,

2222 0,

2223 m_thread_data_on_block_idx[I1],

2224 n_thread_data_on_block_idx[I1],

2225 m_thread_data_on_block_idx[I2],

2226 n_thread_data_on_block_idx[I2],

2227 m_thread_data_on_block_idx[I3],

2228 m_thread_data_on_block_idx[I4],

2229 m_thread_data_on_block_idx[I5],

2230 n_thread_data_on_block_idx[I3]),

2231 ck::tensor_operation::element_wise::PassThrough{}};

2232

2233 // shuffle: blockwise copy C from LDS to global

2234 auto c_shuffle_block_copy_lds_to_global = ThreadGroupTensorSliceTransfer_v6r1<

2235 ThisThreadBlock, // ThreadGroup

2236 CElementwiseOperation, // ElementwiseOperation,

2237 CGlobalMemoryDataOperation, // DstInMemOp,

2238 Sequence<1,

2239 CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,

2240 1,

2241 CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths,

2242 CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,

2243 Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder,

2244 CShuffleDataType, // typename SrcData,

2245 CDataType, // typename DstData,

2246 decltype(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock),

2247 decltype(c_grid_desc_mblock_mperblock_nblock_nperblock),

2248 Sequence<0, 1, 2, 3>, // typename DimAccessOrder,

2249 3, // index_t VectorDim,

2250 CShuffleBlockTransferScalarPerVector_NPerBlock, // index_t ScalarPerVector,

2251 true, // bool ThreadTransferSrcResetCoordinateAfterRun,

2252 false> // bool ThreadTransferDstResetCoordinateAfterRun>

2253 {c_shuffle_block_desc_mblock_mperblock_nblock_nperblock,

2254 make_multi_index(0, 0, 0, 0),

2255 c_grid_desc_mblock_mperblock_nblock_nperblock,

2256 make_multi_index(block_m_id, 0, block_n_id, 0),

2257 c_element_op};

2258

2259 // space filling curve for threadwise C in VGPR

2260 constexpr auto sfc_c_vgpr =

2261 SpaceFillingCurve<Sequence<MXdlPerWave / MXdlPack,

2262 NXdlPerWave / NXdlPack,

2263 1,

2264 1,

2265 MXdlPack,

2266 NXdlPack,

2267 M2,

2268 1,

2269 M4,

2270 1>,

2271 Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>,

2272 Sequence<CShuffleMXdlPerWavePerShuffle / MXdlPack,

2273 CShuffleNXdlPerWavePerShuffle / NXdlPack,

2274 1,

2275 1,

2276 MXdlPack,

2277 NXdlPack,

2278 M2,

2279 1,

2280 M4,

2281 1>>{};

2282

2283 // space filling curve for shuffled blockwise C in global mem

2284 constexpr auto sfc_c_global =

2285 SpaceFillingCurve<Sequence<1, MPerBlock, 1, NPerBlock>,

2286 Sequence<0, 2, 1, 3>,

2287 Sequence<1,

2288 CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,

2289 1,

2290 CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>{};

2291

2292 constexpr index_t num_access = sfc_c_vgpr.GetNumOfAccess();

2293

2294 static_assert(num_access == sfc_c_global.GetNumOfAccess(), "wrong!");

2295

2296 static_for<0, num_access, 1>{}([&](auto access_id) {

2297 // make sure it's safe to write to LDS

2298 block_sync_lds();

2299

2300 // each thread write its data from VGPR to LDS

2301 c_thread_copy_vgpr_to_lds.Run(c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2,

2302 sfc_c_vgpr.GetIndexTupleOfNumber(access_id),

2303 c_thread_buf,

2304 c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2,

2305 c_shuffle_block_buf);

2306

2307 // make sure it's safe to read from LDS

2308 block_sync_lds();

2309

2310 // each block copy its data from LDS to global

2311 c_shuffle_block_copy_lds_to_global.Run(

2312 c_shuffle_block_desc_mblock_mperblock_nblock_nperblock,

2313 c_shuffle_block_buf,

2314 c_grid_desc_mblock_mperblock_nblock_nperblock,

2315 c_grid_buf);

2316

2317 if constexpr(access_id < num_access - 1)

2318 {

2319 constexpr auto c_global_step = sfc_c_global.GetForwardStep(access_id);

2320

2321 // move on C

2322 c_shuffle_block_copy_lds_to_global.MoveDstSliceWindow(

2323 c_grid_desc_mblock_mperblock_nblock_nperblock, c_global_step);

2324 }

2325 });

2326 }

2327 }

2328

2329 template <bool HasMainKBlockLoop,

2330 InMemoryDataOperationEnum CGlobalMemoryDataOperation,

2331 TailNumber TailNum = TailNumber::Odd>

2332 __device__ static void Run_2Lds(const ADataType* p_a_grid,

2333 const AScaleDataType* p_a_scale_grid,

2334 const BDataType* p_b_grid,

2335 const BScaleDataType* p_b_scale_grid,

2336 CDataType* p_c_grid,

2337 void* p_shared_0,

2338 void* p_shared_1,

2339 const Problem& problem)

2340 {

2341 const auto a_grid_desc_ak0_m_ak1 = MakeAGridDescriptor_AK0_M_AK1(

2342 problem.M, problem.MPadded, problem.K, problem.KPadded, problem.StrideA, problem.AK0);

2343 const auto b_grid_desc_bk0_n_bk1 = MakeBGridDescriptor_BK0_N_BK1(

2344 problem.K, problem.KPadded, problem.N, problem.NPadded, problem.StrideB, problem.BK0);

2345 const auto c_grid_desc_m_n = MakeCGridDescriptor_M_N(

2346 problem.M, problem.MPadded, problem.N, problem.NPadded, problem.StrideC);

2347 const auto c_grid_desc_mblock_mperblock_nblock_nperblock =

2348 MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(

2349 c_grid_desc_m_n, problem.MBlock, problem.NBlock);

2350

2351 // A/B shuffled scale for better 8-bit scale access pattern

2352 // MNRepeat -> KRepeat -> KThreadPerXdl -> MNThreadPerXdl -> KXdlPack -> MNXdlPack

2353 const auto Padded_Scale_M =

2354 math::integer_divide_ceil(problem.M, ScaleBlockSize) * ScaleBlockSize;

2355 const auto a_scale_grid_desc_am_ak = make_naive_tensor_descriptor(

2356 make_tuple(Padded_Scale_M / (MXdlPack * MPerXdl),

2357 math::integer_divide_ceil(problem.K, (ScaleBlockSize / APackedSize)) /

2358 (KXdlPack * 64 / MPerXdl),

2359 64 * KXdlPack * MXdlPack / scale_pack_size_a),

2360 make_tuple(math::integer_divide_ceil(problem.K * problem.KBatch,

2361 (ScaleBlockSize / APackedSize)) *

2362 MPerXdl * MXdlPack / scale_pack_size_a,

2363 64 * KXdlPack * MXdlPack / scale_pack_size_a,

2364 1));

2365

2366 const auto b_scale_grid_desc_bn_ak = make_naive_tensor_descriptor(

2367 make_tuple(problem.N / (NXdlPack * NPerXdl),

2368 math::integer_divide_ceil(problem.K, (ScaleBlockSize / BPackedSize)) /

2369 (KXdlPack * 64 / NPerXdl),

2370 64 * KXdlPack * NXdlPack / scale_pack_size_b),

2371 make_tuple(math::integer_divide_ceil(problem.K * problem.KBatch,

2372 (ScaleBlockSize / BPackedSize)) *

2373 NPerXdl * NXdlPack / scale_pack_size_b,

2374 64 * KXdlPack * NXdlPack / scale_pack_size_b,

2375 1));

2376

2377 Run_2Lds<decltype(a_grid_desc_ak0_m_ak1),

2378 decltype(a_scale_grid_desc_am_ak),

2379 decltype(b_grid_desc_bk0_n_bk1),

2380 decltype(b_scale_grid_desc_bn_ak),

2381 decltype(c_grid_desc_mblock_mperblock_nblock_nperblock),

2382 HasMainKBlockLoop,

2383 CGlobalMemoryDataOperation,

2384 TailNum>(p_a_grid,

2385 p_a_scale_grid,

2386 p_b_grid,

2387 p_b_scale_grid,

2388 p_c_grid,

2389 p_shared_0,

2390 p_shared_1,

2391 problem,

2392 a_grid_desc_ak0_m_ak1,

2393 a_scale_grid_desc_am_ak,

2394 b_grid_desc_bk0_n_bk1,

2395 b_scale_grid_desc_bn_ak,

2396 c_grid_desc_mblock_mperblock_nblock_nperblock);

2397 }

2398};

2399

2400} // namespace ck

block_to_ctile_map.hpp

blockwise_gemm_pipeline_xdlops_mx_selector.hpp

CK_MAX_THREAD_PER_BLOCK

#define CK_MAX_THREAD_PER_BLOCK

Definition ck.hpp:30

common_header.hpp

IS_VALID_COMPILATION_PARAMETER_IMPL

#define IS_VALID_COMPILATION_PARAMETER_IMPL(CDataType_)

Definition device_base.hpp:178

element_wise_operation.hpp

multi_index_transform_helper.hpp

ck::math::integer_least_multiple

__host__ __device__ constexpr auto integer_least_multiple(X x, Y y)

Definition utility/math.hpp:78

ck::math::max

__host__ __device__ constexpr T max(T x)

Definition utility/math.hpp:84

ck::math::integer_divide_ceil

__host__ __device__ constexpr auto integer_divide_ceil(X x, Y y)

Definition utility/math.hpp:72

ck::math::lcm

__host__ __device__ constexpr auto lcm(X x, Y y)

Definition utility/math.hpp:198

ck::tensor_operation::device::GemmSpecialization

GemmSpecialization

Definition gemm_specialization.hpp:11

ck::tensor_operation::device::GemmSpecialization::MKPadding

@ MKPadding

Definition gemm_specialization.hpp:18

ck::tensor_operation::device::GemmSpecialization::KPadding

@ KPadding

Definition gemm_specialization.hpp:16

ck::tensor_operation::device::GemmSpecialization::NPadding

@ NPadding

Definition gemm_specialization.hpp:15

ck::tensor_operation::device::GemmSpecialization::MPadding

@ MPadding

Definition gemm_specialization.hpp:14

ck::tensor_operation::device::GemmSpecialization::MNKPadding

@ MNKPadding

Definition gemm_specialization.hpp:20

ck::tensor_operation::device::GemmSpecialization::MNPadding

@ MNPadding

Definition gemm_specialization.hpp:17

ck::tensor_operation::device::GemmSpecialization::NKPadding

@ NKPadding

Definition gemm_specialization.hpp:19

ck

Definition ck.hpp:268

ck::bhalf_t

ushort bhalf_t

Definition data_type.hpp:30

ck::make_multi_index

__host__ __device__ constexpr auto make_multi_index(Xs &&... xs)

Definition array_multi_index.hpp:15

ck::f6x16_pk_t

f6_pk_t< f6_t, 16 > f6x16_pk_t

Definition data_type.hpp:180

ck::make_pass_through_transform

__host__ __device__ constexpr auto make_pass_through_transform(const LowLength &low_length)

Definition multi_index_transform_helper.hpp:12

ck::index_t

int32_t index_t

Definition ck.hpp:299

ck::BlockGemmMXPipeline_Selector

constexpr auto BlockGemmMXPipeline_Selector()

Definition blockwise_gemm_pipeline_xdlops_mx_moe_selector.hpp:36

ck::make_naive_tensor_descriptor

__host__ __device__ constexpr auto make_naive_tensor_descriptor(const Tuple< Lengths... > &lengths, const Tuple< Strides... > &strides)

Definition tensor_descriptor_helper.hpp:49

ck::InMemoryDataOperationEnum

InMemoryDataOperationEnum

Definition ck.hpp:277

ck::InMemoryDataOperationEnum::Set

@ Set

Definition ck.hpp:278

ck::make_single_stage_tensor_adaptor

__host__ __device__ constexpr auto make_single_stage_tensor_adaptor(const Transforms &transforms, LowerDimensionOldTopIdss, UpperDimensionNewTopIdss)

Definition tensor_description/tensor_adaptor.hpp:425

ck::remove_cvref_t

remove_cv_t< remove_reference_t< T > > remove_cvref_t

Definition type.hpp:297

ck::BlockGemmPipelineVersion

BlockGemmPipelineVersion

Definition blkgemmpipe_scheduler.hpp:12

ck::BlockGemmPipelineVersion::v3

@ v3

Definition blkgemmpipe_scheduler.hpp:16

ck::BlockGemmPipelineVersion::v4

@ v4

Definition blkgemmpipe_scheduler.hpp:17

ck::BlockGemmPipelineVersion::v1

@ v1

Definition blkgemmpipe_scheduler.hpp:14

ck::make_freeze_transform

__host__ __device__ constexpr auto make_freeze_transform(const LowerIndex &low_idx)

Definition multi_index_transform_helper.hpp:151

ck::make_right_pad_transform

__host__ __device__ constexpr auto make_right_pad_transform(const LowLength &low_length, const RightPadLength &right_pad, integral_constant< bool, SkipIsValidCheck >=integral_constant< bool, false >{})

Definition multi_index_transform_helper.hpp:37

ck::half_t

_Float16 half_t

Definition data_type.hpp:31

ck::make_xor_with_modulo_transform

__host__ __device__ constexpr auto make_xor_with_modulo_transform(const LowLengths &low_lengths)

Definition multi_index_transform_helper.hpp:185

ck::Number

integral_constant< index_t, N > Number

Definition number.hpp:12

ck::TailNumber

TailNumber

Definition blkgemmpipe_scheduler.hpp:31

ck::TailNumber::Odd

@ Odd

Definition blkgemmpipe_scheduler.hpp:33

ck::TailNumber::Full

@ Full

Definition blkgemmpipe_scheduler.hpp:49

ck::make_merge_transform

__host__ __device__ constexpr auto make_merge_transform(const LowLengths &low_lengths)

Definition multi_index_transform_helper.hpp:55

ck::bf6x32_pk_t

f6_pk_t< bf6_t, 32 > bf6x32_pk_t

Definition data_type.hpp:183

ck::ignore

constexpr detail::ignore_t ignore

Definition utility/ignore.hpp:20

ck::get_block_1d_id

__device__ index_t get_block_1d_id()

Definition get_id.hpp:47

ck::f6x32_pk_t

f6_pk_t< f6_t, 32 > f6x32_pk_t

Definition data_type.hpp:181

ck::EnvIsEnabled

bool EnvIsEnabled(EnvVar)

Definition utility/env.hpp:140

ck::is_same_v

constexpr bool is_same_v

Definition type.hpp:283

ck::make_merge_transform_v3_division_mod

__host__ __device__ constexpr auto make_merge_transform_v3_division_mod(const LowLengths &low_lengths)

Definition multi_index_transform_helper.hpp:84

ck::kernel_gemm_xdl_cshuffle_v3_mx

__global__ enable_if_t<!Use2LDS, void > kernel_gemm_xdl_cshuffle_v3_mx(typename GridwiseGemm::Argument karg)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:40

ck::BlockGemmPipelineScheduler

BlockGemmPipelineScheduler

Definition blkgemmpipe_scheduler.hpp:25

ck::BlockGemmPipelineScheduler::Intrawave

@ Intrawave

Definition blkgemmpipe_scheduler.hpp:26

ck::make_naive_tensor_descriptor_packed

__host__ __device__ constexpr auto make_naive_tensor_descriptor_packed(const Tuple< Lengths... > &lengths)

Definition tensor_descriptor_helper.hpp:101

ck::make_tuple

__host__ __device__ constexpr auto make_tuple(Xs &&... xs)

Definition utility/tuple.hpp:211

ck::get_thread_local_1d_id

__device__ index_t get_thread_local_1d_id()

Definition get_id.hpp:41

ck::transform_tensor_descriptor

__host__ __device__ constexpr auto transform_tensor_descriptor(const OldTensorDescriptor &old_tensor_desc, const NewTransforms &new_transforms, NewLowerDimensionOldVisibleIdss, NewUpperDimensionNewVisibleIdss)

Definition tensor_description/tensor_descriptor.hpp:319

ck::bit_cast

__host__ __device__ constexpr Y bit_cast(const X &x)

Definition type.hpp:306

ck::block_sync_lds

__device__ void block_sync_lds()

Definition synchronization.hpp:16

ck::make_unmerge_transform

__host__ __device__ constexpr auto make_unmerge_transform(const UpLengths &up_lengths, integral_constant< bool, Use24BitIntegerCalculation >=integral_constant< bool, false >{})

Definition multi_index_transform_helper.hpp:90

ck::packed_size_v

constexpr index_t packed_size_v

Definition data_type.hpp:411

ck::make_dynamic_buffer

__host__ __device__ constexpr auto make_dynamic_buffer(T *p, ElementSpaceSize element_space_size)

Definition dynamic_buffer.hpp:472

ck::enable_if_t

typename std::enable_if< B, T >::type enable_if_t

Definition enable_if.hpp:27

ck::bf6x16_pk_t

f6_pk_t< bf6_t, 16 > bf6x16_pk_t

Definition data_type.hpp:182

value

const GenericPointer< typename T::ValueType > T2 value

Definition pointer.h:1697

int32_t

signed int int32_t

Definition stdint.h:123

ck::BlockToCTileMap_Grouped_M00_N0_M01Adapt

Definition block_to_ctile_map.hpp:271

ck::BlockToCTileMap_Grouped_M00_N0_M01Adapt< 8, MPerBlock, NPerBlock >::CalculateGridSize

__host__ static __device__ constexpr index_t CalculateGridSize(index_t M, index_t N)

Definition block_to_ctile_map.hpp:283

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:685

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::a_element_op

const AElementwiseOperation a_element_op

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:741

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::Argument

__host__ Argument(const ADataType *p_a_grid_, const AScaleDataType *p_a_scale_grid_, const BDataType *p_b_grid_, const BScaleDataType *p_b_scale_grid_, CDataType *p_c_grid_, index_t M_, index_t N_, index_t K_, index_t StrideA_, index_t StrideScaleA_, index_t StrideB_, index_t StrideScaleB_, index_t StrideC_, index_t k_batch_, AElementwiseOperation a_element_op_, BElementwiseOperation b_element_op_, CElementwiseOperation c_element_op_, bool is_reduce_=false)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:686

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::IsAtomicAdd

__host__ __device__ bool IsAtomicAdd() const

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:730

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::IsReduceAdd

__host__ __device__ bool IsReduceAdd() const

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:725

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::p_a_scale_grid

const AScaleDataType * p_a_scale_grid

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:736

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::p_c_grid

CDataType * p_c_grid

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:739

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::c_element_op

const CElementwiseOperation c_element_op

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:743

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::is_reduce

bool is_reduce

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:744

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::p_a_grid

const ADataType * p_a_grid

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:735

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::p_b_scale_grid

const BScaleDataType * p_b_scale_grid

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:738

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::p_b_grid

const BDataType * p_b_grid

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:737

ck::GridwiseGemmMX_xdl_cshuffle_v3::Argument::b_element_op

const BElementwiseOperation b_element_op

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:742

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::Print

__host__ void Print() const

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:653

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::StrideB

index_t StrideB

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:669

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::StrideScaleA

index_t StrideScaleA

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:668

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::KPadded

index_t KPadded

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:676

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::StrideA

index_t StrideA

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:667

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::M

index_t M

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:664

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::MPadded

index_t MPadded

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:673

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::NBlock

index_t NBlock

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:680

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::N

index_t N

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:665

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::NPadded

index_t NPadded

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:674

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::Problem

__host__ Problem(index_t M_, index_t N_, index_t K_, index_t StrideA_, index_t StrideScaleA_, index_t StrideB_, index_t StrideScaleB_, index_t StrideC_, index_t KBatch_)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:624

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::BK0

index_t BK0

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:678

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::K

index_t K

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:666

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::KRead

index_t KRead

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:675

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::MBlock

index_t MBlock

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:679

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::StrideC

index_t StrideC

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:671

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::StrideScaleB

index_t StrideScaleB

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:670

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::KBatch

index_t KBatch

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:672

ck::GridwiseGemmMX_xdl_cshuffle_v3::Problem::AK0

index_t AK0

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:677

ck::GridwiseGemmMX_xdl_cshuffle_v3::SplitKBatchOffset::b_k_split_offset

index_t b_k_split_offset

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:806

ck::GridwiseGemmMX_xdl_cshuffle_v3::SplitKBatchOffset::SplitKBatchOffset

__device__ SplitKBatchOffset(Argument &karg, index_t k_id)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:750

ck::GridwiseGemmMX_xdl_cshuffle_v3::SplitKBatchOffset::a_scale_k_split_offset

index_t a_scale_k_split_offset

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:807

ck::GridwiseGemmMX_xdl_cshuffle_v3::SplitKBatchOffset::a_k_split_offset

index_t a_k_split_offset

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:805

ck::GridwiseGemmMX_xdl_cshuffle_v3::SplitKBatchOffset::b_scale_k_split_offset

index_t b_scale_k_split_offset

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:808

ck::GridwiseGemmMX_xdl_cshuffle_v3::SplitKBatchOffset::c_reduce_offset

index_t c_reduce_offset

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:809

ck::GridwiseGemmMX_xdl_cshuffle_v3

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:156

ck::GridwiseGemmMX_xdl_cshuffle_v3::Run

static __device__ void Run(const ADataType *p_a_grid, const AScaleDataType *p_a_scale_grid, const BDataType *p_b_grid, const BScaleDataType *p_b_scale_grid, CDataType *p_c_grid, void *p_shared, const Problem &problem)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1811

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateHasMainKBlockLoop

static __host__ constexpr bool CalculateHasMainKBlockLoop(index_t K)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1321

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::ThisThreadBlock

ThisThreadBlock< BlockSize > ThisThreadBlock

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:202

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::BPackedSize

static constexpr index_t BPackedSize

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:190

ck::GridwiseGemmMX_xdl_cshuffle_v3::MakeBGridDescriptor_BK0_N_BK1

__host__ static __device__ auto MakeBGridDescriptor_BK0_N_BK1(index_t K, index_t KPad, index_t N, index_t NPad, index_t StrideB, index_t BK0)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:403

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateBK0Padded

static __host__ auto CalculateBK0Padded(index_t K, index_t K_Batch=1)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:230

ck::GridwiseGemmMX_xdl_cshuffle_v3::Run

static __device__ void Run(const ADataType *p_a_grid, const AScaleDataType *p_a_scale_grid, const BDataType *p_b_grid, const BScaleDataType *p_b_scale_grid, CDataType *p_c_grid, void *p_shared, const Problem &problem, const AGridDesc_AK0_M_K1 &a_grid_desc_ak0_m_ak1, const AScaleGridDesc_AM_AK &a_scale_grid_desc_am_ak, const BGridDesc_BK0_N_K1 &b_grid_desc_bk0_n_bk1, const BScaleGridDesc_BN_AK &b_scale_grid_desc_bn_ak, const CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock &c_grid_desc_mblock_mperblock_nblock_nperblock)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1374

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::CheckValidity

static __host__ constexpr bool CheckValidity(const Argument &karg)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1123

ck::GridwiseGemmMX_xdl_cshuffle_v3::Run_2Lds

static __device__ void Run_2Lds(const ADataType *p_a_grid, const AScaleDataType *p_a_scale_grid, const BDataType *p_b_grid, const BScaleDataType *p_b_scale_grid, CDataType *p_c_grid, void *p_shared_0, void *p_shared_1, const Problem &problem)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:2332

ck::GridwiseGemmMX_xdl_cshuffle_v3::GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1

static __device__ constexpr auto GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1()

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:812

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::BK1Number

static constexpr auto BK1Number

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:173

ck::GridwiseGemmMX_xdl_cshuffle_v3::MakeGemmMmaTileDescriptor

__host__ static __device__ constexpr auto MakeGemmMmaTileDescriptor(const TileDesc_K0_MN_K1 &)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:264

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateAK0Padded

static __host__ auto CalculateAK0Padded(index_t K, index_t K_Batch=1)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:224

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateGridSize

static __host__ auto CalculateGridSize(index_t M, index_t N, index_t KBatch)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:204

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::is_single_rate_mfma

static constexpr bool is_single_rate_mfma

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:176

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateNBlock

static __host__ auto CalculateNBlock(index_t N)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:254

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateNPadded

static __host__ auto CalculateNPadded(index_t N)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:214

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I0

static constexpr auto I0

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:158

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I8

static constexpr auto I8

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:166

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I3

static constexpr auto I3

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:161

ck::GridwiseGemmMX_xdl_cshuffle_v3::GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock

static __device__ constexpr auto GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock()

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1049

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::MXdlPack

static constexpr auto MXdlPack

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:179

ck::GridwiseGemmMX_xdl_cshuffle_v3::MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock

__host__ static __device__ constexpr auto MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(const CGridDesc &c_grid_desc_m_n, index_t MBlock, index_t NBlock)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1336

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::KXdlPack

static constexpr auto KXdlPack

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:181

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateKPadded

static __host__ auto CalculateKPadded(index_t K)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:219

ck::GridwiseGemmMX_xdl_cshuffle_v3::GetSharedMemoryNumberOfByte

static __device__ constexpr index_t GetSharedMemoryNumberOfByte()

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1093

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateKRead

static __host__ auto CalculateKRead(index_t K, index_t K_Batch=1)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:242

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::BK0Number

static constexpr auto BK0Number

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:171

ck::GridwiseGemmMX_xdl_cshuffle_v3::Run_2Lds

static __device__ void Run_2Lds(const ADataType *p_a_grid, const AScaleDataType *p_a_scale_grid, const BDataType *p_b_grid, const BScaleDataType *p_b_scale_grid, CDataType *p_c_grid, void *p_shared_0, void *p_shared_1, const Problem &problem, const AGridDesc_AK0_M_K1 &a_grid_desc_ak0_m_ak1, const AScaleGridDesc_AM_AK &a_scale_grid_desc_am_ak, const BGridDesc_BK0_N_K1 &b_grid_desc_bk0_n_bk1, const BScaleGridDesc_BN_AK &b_scale_grid_desc_bn_ak, const CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock &c_grid_desc_mblock_mperblock_nblock_nperblock)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1884

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::AK0Number

static constexpr auto AK0Number

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:170

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::scale_pack_size_b

static constexpr index_t scale_pack_size_b

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1356

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I7

static constexpr auto I7

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:165

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::mx_scale_t

e8m0_bexp_t mx_scale_t

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1354

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I2

static constexpr auto I2

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:160

ck::GridwiseGemmMX_xdl_cshuffle_v3::MakeAMmaTileDescriptor_M0_M1_M2_M3_K

__host__ static __device__ constexpr auto MakeAMmaTileDescriptor_M0_M1_M2_M3_K(const ABlockDesc_AK0_M_AK1 &)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:543

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::lcm_AK1_BK1

static constexpr auto lcm_AK1_BK1

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:175

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::BlockwiseGemmPipe

remove_cvref_t< decltype(BlockGemmMXPipeline_Selector< BlkGemmPipelineVer, BlkGemmPipeSched, BlockSize, ScaleBlockSize, ADataType, AScaleDataType, BDataType, BScaleDataType, ComputeTypeA, GemmAccDataType, decltype(GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1()), decltype(GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1()), decltype(MakeAMmaTileDescriptor_M0_M1_M2_M3_K(GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1())), decltype(MakeBMmaTileDescriptor_N0_N1_N2_N3_K(GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1())), ABlockTransferSrcScalarPerVector, BBlockTransferSrcScalarPerVector, MPerBlock, NPerBlock, KPerBlock, MPerXdl, NPerXdl, MXdlPerWave, NXdlPerWave, KPack >())> BlockwiseGemmPipe

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1064

ck::GridwiseGemmMX_xdl_cshuffle_v3::MakeCGridDescriptor_M_N

__host__ static __device__ auto MakeCGridDescriptor_M_N(index_t M, index_t MPad, index_t N, index_t NPad, index_t StrideC)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:562

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I1

static constexpr auto I1

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:159

ck::GridwiseGemmMX_xdl_cshuffle_v3::MakeAGridDescriptor_AK0_M_AK1

__host__ static __device__ auto MakeAGridDescriptor_AK0_M_AK1(index_t M, index_t MPad, index_t K, index_t KPad, index_t StrideA, index_t AK0)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:288

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I6

static constexpr auto I6

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:164

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::scale_pack_size_a

static constexpr index_t scale_pack_size_a

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1355

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateMBlock

static __host__ auto CalculateMBlock(index_t M)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:249

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::APackedSize

static constexpr index_t APackedSize

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:189

ck::GridwiseGemmMX_xdl_cshuffle_v3::MakeBMmaTileDescriptor_N0_N1_N2_N3_K

__host__ static __device__ constexpr auto MakeBMmaTileDescriptor_N0_N1_N2_N3_K(const BBlockDesc_BK0_N_BK1 &)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:553

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I9

static constexpr auto I9

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:167

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I4

static constexpr auto I4

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:162

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateKBlockLoopTailNum

static __host__ constexpr TailNumber CalculateKBlockLoopTailNum(index_t K)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1328

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::NXdlPack

static constexpr auto NXdlPack

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:180

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::AK1Number

static constexpr auto AK1Number

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:172

ck::GridwiseGemmMX_xdl_cshuffle_v3::GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1

static __device__ constexpr auto GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1()

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:933

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateMPadded

static __host__ auto CalculateMPadded(index_t M)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:209

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::I5

static constexpr auto I5

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:163

ck::GridwiseGemmMX_xdl_cshuffle_v3::CalculateKPadded

static __host__ auto CalculateKPadded(index_t K, index_t K_Batch=1)

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:236

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::Block2CTileMap

BlockToCTileMap_Grouped_M00_N0_M01Adapt< 8, MPerBlock, NPerBlock > Block2CTileMap

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:1351

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::KPack

static constexpr index_t KPack

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:192

ck::GridwiseGemmMX_xdl_cshuffle_v3< ALayout, BLayout, CLayout, ADataType, AScaleDataType, BDataType, BScaleDataType, GemmAccDataType, CShuffleDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, GemmSpec, ScaleBlockSize, BlockSize, MPerBlock, NPerBlock, KPerBlock, AK1, BK1, MPerXDL, NPerXDL, MXdlPerWave, NXdlPerWave_, ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, CShuffleBlockTransferScalarPerVector_NPerBlock, BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, ComputeTypeB >::is_scale_mfma

static constexpr auto is_scale_mfma

Definition gridwise_gemm_xdl_cshuffle_v3_mx.hpp:177

ck::MfmaSelector

Selects the appropriate MFMA instruction type and configuration for given data types and tile sizes o...

Definition xdlops_gemm.hpp:1208

ck::Sequence

Definition utility/sequence.hpp:43

ck::SpaceFillingCurve

Definition tensor_space_filling_curve.hpp:20

ck::ThreadGroupTensorSliceTransfer_DirectLoad

Definition thread_group_tensor_slice_transfer_direct_load.hpp:55

ck::ThreadGroupTensorSliceTransfer_v6r1

Definition thread_group_tensor_slice_transfer_v6r1.hpp:34

ck::ThreadwiseTensorSliceTransfer_v1r3

Definition threadwise_tensor_slice_transfer.hpp:39

ck::ThreadwiseTensorSliceTransfer_v2

Helper structure that facilitates transfer of source (grid) data to destination threads.

Definition threadwise_tensor_slice_transfer.hpp:234

ck::constant::value

static constexpr value_type value

Definition utility/integral_constant.hpp:13

ck::e8m0_bexp_t

Unsigned representation of a conventional biased Float32 exponent.

Definition utility/e8m0.hpp:26

ck::f4x2_pk_t

Definition data_type.hpp:42

ck::is_same

Definition type.hpp:177

ck::pk_i4_t

Definition data_type.hpp:187

ck::static_for

Definition functional2.hpp:33

ck::tensor_operation::device::BaseArgument

Definition device_base.hpp:197

ck::tensor_operation::element_wise::PassThrough

Definition tensor_operation/gpu/element/unary_element_wise_operation.hpp:340

tensor_descriptor.hpp

tensor_descriptor_helper.hpp

thread_group_tensor_slice_transfer_direct_load.hpp

thread_group_tensor_slice_transfer_v4r1.hpp

thread_group_tensor_slice_transfer_v6r1.hpp

threadwise_tensor_slice_transfer.hpp

env.hpp

CK_ENV

#define CK_ENV(name)

Definition utility/env.hpp:129