ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-11 01:27:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP

Browse Source
This commit is contained in:
Enrico Degregori
2026-07-07 08:33:48 +00:00
parent d408192bfe
commit 51134d0e66
5 changed files with 44 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp
View File

@@ -1247,6 +1247,13 @@ llvm_amdgcn_raw_buffer_store_i16x4(int16x4_t vdata,
index_t soffset,
index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4i16");
CK_TILE_DEVICE_EXTERN void
llvm_amdgcn_raw_buffer_store_i16x8(int16x8_t vdata,
int32x4_t rsrc,
index_t voffset,
index_t soffset,
index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v8i16");
// buffer store i32
CK_TILE_DEVICE_EXTERN void
llvm_amdgcn_raw_buffer_store_i32(int32_t vdata,
@@ -2246,19 +2253,12 @@ CK_TILE_DEVICE void amd_buffer_store_impl(const thread_buffer<T, N> src_thread_d
}
else if constexpr(N == 8)
{
llvm_amdgcn_raw_buffer_store_i16x4(
src_thread_data.template get_as<int16x4_t>()[number<0>{}],
llvm_amdgcn_raw_buffer_store_i16x8(
bit_cast<int16x8_t>(src_thread_data),
dst_wave_buffer_resource,
dst_thread_addr_offset,
dst_wave_addr_offset,
static_cast<index_t>(coherence));
llvm_amdgcn_raw_buffer_store_i16x4(
src_thread_data.template get_as<int16x4_t>()[number<1>{}],
dst_wave_buffer_resource,
dst_thread_addr_offset,
dst_wave_addr_offset + 4 * sizeof(bf16_t),
static_cast<index_t>(coherence));
}
}
else if constexpr(std::is_same<T, uint16_t>::value)

1
include/ck_tile/core/arch/amd_buffer_coherence.hpp
View File

@@ -139,6 +139,7 @@ enum struct amd_buffer_coherence_enum
// Other archs compatiblity
DEVICE_NT0 = 0,
SYSTEM_NT0 = 0,
GROUP_NT0 = 0,
DEVICE_NT1 = glc,
SYSTEM_NT1 = slc,
DEVICE = 0, // Required for template parsing (see GFX11 comment re: Two-Phase Lookup)

20
include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
View File

@@ -341,17 +341,17 @@ struct CShuffleEpilogue
"LDS row stride must be 4B-aligned for bank-word padding logic");
// calculate how many elements to pad to avoid bank conflict
#if defined(__gfx950__) || defined(__gfx125__)
#if defined(__gfx950__)
constexpr index_t ElemsPer4B = BytesPerBank / ck_tile::gcd(BytesPerBank, DataTypeSize);
constexpr auto ToWords = [](index_t elems) constexpr {
return (elems * DataTypeSize) / BytesPerBank;
};
constexpr index_t BaseWords = ToWords(BaseStrideElems);
constexpr index_t PadWords = ((BaseWords % 2) == 0) ? 1 : 0;
constexpr auto PaddingAmount = PadWords * ElemsPer4B;
#else
// #if defined(__gfx950__)
// constexpr index_t ElemsPer4B = BytesPerBank / ck_tile::gcd(BytesPerBank, DataTypeSize);
// constexpr auto ToWords = [](index_t elems) constexpr {
// return (elems * DataTypeSize) / BytesPerBank;
// };
// constexpr index_t BaseWords = ToWords(BaseStrideElems);
// constexpr index_t PadWords = ((BaseWords % 2) == 0) ? 1 : 0;
// constexpr auto PaddingAmount = PadWords * ElemsPer4B;
// #else
constexpr auto PaddingAmount = VectorLen;
#endif
// #endif
constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
make_tuple(number<MPerIterationShuffle / MLdsLayer>{},

6
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_eight_waves_base.hpp
View File

@@ -493,11 +493,11 @@ struct GemmPipelineAgBgCrEightWavesImplBase : public GemmPipelineAgBgCrImplBase<
if constexpr(HasHotLoop && TailNum == TailNumber::Even)
{
asm volatile(";; Even Tail Start ;;");
__builtin_amdgcn_s_barrier();
// __builtin_amdgcn_s_barrier();
main_body(I0, I1);
__builtin_amdgcn_s_barrier();
// __builtin_amdgcn_s_barrier();
asm volatile(";; Even Tail End ;;");
__builtin_amdgcn_s_barrier();
// __builtin_amdgcn_s_barrier();
}
constexpr int tic = HasHotLoop ? (TailNum == TailNumber::Odd ? 0 : 1) : 1 - N_LOOP % 2;

28
include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp
View File

@@ -520,7 +520,9 @@ struct QuantGemmMultiDKernel
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
return make_naive_tensor_view<address_space_enum::global,
memory_operation_enum::set,
amd_buffer_coherence_enum::GROUP_NT0>(
a_ptr,
make_tuple(k_size, kargs.M),
make_tuple(kargs.stride_A, 1),
@@ -679,7 +681,9 @@ struct QuantGemmMultiDKernel
}
else if constexpr(kQuantType == QuantType::RowColQuant)
{
return make_naive_tensor_view<address_space_enum::global>(
return make_naive_tensor_view<address_space_enum::global,
memory_operation_enum::set,
amd_buffer_coherence_enum::DEVICE_NT0>(
aq_ptr,
make_tuple(kargs.M, kargs.N),
make_tuple(1, 0), // broadcasting over n
@@ -829,7 +833,9 @@ struct QuantGemmMultiDKernel
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
return make_naive_tensor_view<address_space_enum::global,
memory_operation_enum::set,
amd_buffer_coherence_enum::GROUP_NT0>(
b_ptr,
make_tuple(kargs.N, k_size),
make_tuple(kargs.stride_B, 1),
@@ -906,7 +912,9 @@ struct QuantGemmMultiDKernel
const auto& bq_tensor_view = [&]() {
if constexpr(kQuantType == QuantType::RowColQuant)
{
return make_naive_tensor_view<address_space_enum::global>(
return make_naive_tensor_view<address_space_enum::global,
memory_operation_enum::set,
amd_buffer_coherence_enum::DEVICE_NT0>(
bq_ptr,
make_tuple(kargs.M, kargs.N),
make_tuple(0, 1), // broadcasting over m
@@ -1111,7 +1119,9 @@ struct QuantGemmMultiDKernel
const auto& ds_tensor_view = generate_tuple(
[&](auto i) {
using DDataType_ = remove_cvref_t<std::tuple_element_t<i.value, DsDataType>>;
return make_tensor_view<address_space_enum::global>(
return make_tensor_view<address_space_enum::global,
memory_operation_enum::set,
amd_buffer_coherence_enum::SYSTEM_NT1>(
static_cast<const DDataType_*>(ds_ptr[i]), ds_desc[i]);
},
number<NumDTensor>{});
@@ -1187,7 +1197,9 @@ struct QuantGemmMultiDKernel
const auto& c_tensor_view = [&]() {
if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
return make_naive_tensor_view<address_space_enum::global,
DstInMemOp,
amd_buffer_coherence_enum::SYSTEM_NT1>(
c_ptr,
make_tuple(kargs.M, kargs.N),
make_tuple(kargs.stride_C, 1),
@@ -1196,7 +1208,9 @@ struct QuantGemmMultiDKernel
}
else
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
return make_naive_tensor_view<address_space_enum::global,
DstInMemOp,
amd_buffer_coherence_enum::SYSTEM_NT1>(
c_ptr,
make_tuple(kargs.M, kargs.N),
make_tuple(1, kargs.stride_C),