ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 19:18:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

k padding to 256

Browse Source
This commit is contained in:
qin letao
2025-03-04 14:49:01 +00:00
parent c792807a00
commit d5d01c56b1
4 changed files with 36 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8_bpreshuffle_padding.cpp
View File

@@ -129,7 +129,7 @@ void preShuffleBuffer(const DataType* src, DataType* dst, int N, int K, int NXdl
}
}
}
int GetPreShufflePadded(int K) { return (K + 64 - 1) / 64 * 64; }
int GetPreShufflePadded(int K) { return (K + ShufflePadded - 1) / ShufflePadded * ShufflePadded; }
using PassThrough = ck::tensor_operation::element_wise::PassThrough;
using AElementOp = PassThrough;
@@ -140,38 +140,38 @@ static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecializatio
using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle
// clang-format off
// < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
// AElementOp, BElementOp, CDEElementOp, GemmSpec, 256,
// 256, 256, 128,
// 16, 16,
// 16, 16,
// 8, 8,
// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
// 1, 2, S<1, 32, 1, 8>, S<8, 8, 1>,
// ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, B0DataType>;
< Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
< Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
AElementOp, BElementOp, CDEElementOp, GemmSpec, 256,
128, 128, 128,
16, 16,
32, 32,
2, 2,
S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
1, 1, S<1, 32, 1, 8>, S<8, 8, 1>,
ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v2, B0DataType>;
256, 256, 128,
16, 16,
16, 16,
8, 8,
S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
1, 2, S<1, 32, 1, 8>, S<8, 8, 1>,
ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, B0DataType>;
// < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
// AElementOp, BElementOp, CDEElementOp, GemmSpec, 256,
// 128, 128, 128,
// 16, 16,
// 32, 32,
// 2, 2,
// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
// 1, 1, S<1, 32, 1, 8>, S<8, 8, 1>,
// ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v2, B0DataType>;
// < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType,
// AElementOp, BElementOp, CDEElementOp, GemmSpec, 256,
// 256, 128, 128,
// 16, 16,
// 32, 32,
// 4, 2,
// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
// 1, 1, S<1, 32, 1, 8>, S<8, 8, 1>,
// ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v2, B0DataType>;
// 128, 256, 128,
// 16, 16,
// 32, 32,
// 4, 2,
// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
// 1, 1, S<1, 32, 1, 8>, S<8, 8, 1>,
// ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, B0DataType>;
// clang-format on
@@ -268,7 +268,7 @@ int main(int argc, char* argv[])
}
};
auto Knew = GetPreShufflePadded(K);
auto Knew = ck::tensor_operation::device::GetPreShufflePadded(K);
auto StrideBnew = Knew;
Tensor<A0DataType> a0_m_k(f_host_tensor_descriptor(M, K, StrideA, A0Layout{}));
Tensor<B0DataType> b0_k_n(f_host_tensor_descriptor(K, N, StrideB, B0Layout{}));

1
include/ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp
View File

@@ -147,6 +147,7 @@ struct DeviceGemmMultipleDSplitKBPreShuffle : public BaseOperator
virtual int GetPreShuffleParameters() = 0;
};
#define ShufflePadded 256
} // namespace device
} // namespace tensor_operation
} // namespace ck

6
include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp
View File

@@ -214,7 +214,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
__host__ __device__ static auto CalculateBKShufflePadded(index_t K)
{
return (K + 64 - 1) / 64 * 64;
return (K + ShufflePadded - 1) / ShufflePadded * ShufflePadded;
}
__host__ __device__ static auto CalculateKPadded(index_t K)
@@ -605,7 +605,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
MBlock{CalculateMBlock(M_)},
NBlock{CalculateNBlock(N_)},
BN0Shuffled{CalculateBN0Shuffled(N_)},
BK0Shuffled{CalculateBK0Shuffled((K_ + 64 - 1) / 64 * 64)}
BK0Shuffled{CalculateBK0Shuffled(CalculateBKShufflePadded(K_))}
{
}
@@ -936,7 +936,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
return false;
}
if(karg.N % NPerBlock != 0 && NXdlPerWave == 1)
if(karg.N % NPerBlock != 0 && karg.N % (NPerXdl * 2) != 0)
{
return false;
}

3
profiler/include/profiler/profile_gemm_multiply_multiply_weight_preshuffle_impl.hpp
View File

@@ -56,7 +56,8 @@ void preShuffleBuffer(
}
}
}
int GetPreShufflePadded(int K) { return (K + 64 - 1) / 64 * 64; }
int GetPreShufflePadded(int K) { return (K + ShufflePadded - 1) / ShufflePadded * ShufflePadded; }
template <typename ADataType,
typename BDataType,