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

todo: fix asyn_load

Browse Source
This commit is contained in:
ZheWang
2025-12-22 11:30:18 +00:00
parent 7723175746
commit e9baabdfc5
6 changed files with 85 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
include/ck_tile/core/arch/amd_buffer_addressing.hpp
View File

@@ -1468,7 +1468,7 @@ amd_buffer_load_impl_with_bytes(int32x4_t src_wave_buffer_resource,
}
else if constexpr(N == 12)
{
int32x3_t tmp = llvm_amdgcn_raw_buffer_load_i32x3(src_wave_buffer_resource,
auto tmp = llvm_amdgcn_raw_buffer_load_i32x3(src_wave_buffer_resource,
src_thread_addr_offset,
src_wave_addr_offset,
static_cast<index_t>(coherence));

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

@@ -1135,12 +1135,24 @@ llvm_amdgcn_raw_buffer_store_i32x2(int32x2_t vdata,
index_t soffset,
index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i32");
using int32x3 = int32_t __attribute__((ext_vector_type(3)));
CK_TILE_DEVICE_EXTERN void
llvm_amdgcn_raw_buffer_store_i32x3(int32x3_t vdata,
int32x4_t rsrc,
index_t voffset,
index_t soffset,
index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v3i32");
llvm_amdgcn_raw_buffer_store_i32x3_(int32x3 vdata,
int32x4_t rsrc,
index_t voffset,
index_t soffset,
index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v3i32");
CK_TILE_DEVICE_EXTERN void llvm_amdgcn_raw_buffer_store_i32x3(
int32x3_t vdata, int32x4_t rsrc, index_t voffset, index_t soffset, index_t glc_slc)
{
int32x3 v_reg;
v_reg[0] = vdata.data[0];
v_reg[1] = vdata.data[1];
v_reg[2] = vdata.data[2];
llvm_amdgcn_raw_buffer_store_i32x3_(v_reg, rsrc, voffset, soffset, 0);
};
CK_TILE_DEVICE_EXTERN void
llvm_amdgcn_raw_buffer_store_i32x4(int32x4_t vdata,
@@ -1873,7 +1885,9 @@ CK_TILE_DEVICE void amd_buffer_store_impl(const thread_buffer<T, N> src_thread_d
(std::is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
(std::is_same<T, uint16_t>::value &&
(N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
(std::is_same<T, uint8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
(std::is_same<T, uint8_t>::value &&
(N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
std::is_same<T, f6x16_pk_t>::value && (N == 1),
"wrong! not implemented");
if constexpr(std::is_same<T, float>::value) // fp32

5
include/ck_tile/core/numeric/vector_type.hpp
View File

@@ -49,10 +49,9 @@ struct ext_vector<T_, N_, std::enable_if_t<!std::is_class_v<typename native_t<T_
template <index_t N_>
struct ext_vector<f6x16_pk_t, N_>
{
// static_assert(N_ == -1);
static constexpr index_t N = N_;
using value_type = uint32_t;
using type = value_type __attribute__((ext_vector_type(N))); // this is danguous
using value_type = f6x16_pk_t;
using type = f6x16_pk_t; // this is danguous
};
template <typename T_, index_t N_>

11
include/ck_tile/core/tensor/tile_window.hpp
View File

@@ -20,6 +20,15 @@
namespace ck_tile {
template <auto... val>
[[deprecated("Help function to print value")]] inline constexpr void CK_PRINTT()
{
}
template <typename... type>
[[deprecated("Help function to print value")]] inline constexpr void CK_PRINTT()
{
}
/**
* @brief This class provides tile (windowed) view and access to the device memory.
*
@@ -551,7 +560,7 @@ struct tile_window_with_static_distribution
using vector_t = typename Traits::vector_t;
using SFC_Ys = typename Traits::SFC_Ys;
CK_PRINTT<vector_t>();
// Precompute invariant values outside loops
const auto window_origin = lds_tile.get_window_origin();
const auto& bottom_tensor_view = lds_tile.get_bottom_tensor_view();

20
test/ck_tile/memory_copy/test_copy.cpp
View File

@@ -103,6 +103,24 @@ class TestCkTileMemoryCopy : public ::testing::TestWithParam<std::tuple<int, int
// reference
y_buf.FromDevice(y_host_dev.mData.data());
for(int i = 0; i < m; i++)
{
for(int j = 0; j < n; j += pack_k)
{
// for(int k_ = 0; k_ < pack_k; k_++)
// {
// int value = m & 0x3f;
std::cout << x_host(i, j).data_[0] << " " << x_host(i, j).data_[1] << " "
<< x_host(i, j).data_[0] << std::endl;
printf("==============\n");
std::cout << y_host_dev(i, j).data_[0] << " " << y_host_dev(i, j).data_[1] << " "
<< y_host_dev(i, j).data_[0] << std::endl;
printf("~~~~~~~~~~~~~\n");
}
}
// bool pass = ck_tile::check_err(y_host_dev, x_host);
// EXPECT_TRUE(pass);
@@ -110,7 +128,7 @@ class TestCkTileMemoryCopy : public ::testing::TestWithParam<std::tuple<int, int
};
// class TestCkTileMemoryCopyF6x16 : public TestCkTileMemoryCopy<ck_tile::pk_fp4_t, false>
class TestCkTileMemoryCopyF6x16 : public TestCkTileMemoryCopy<ck_tile::f6x16_pk_t, false>
class TestCkTileMemoryCopyF6x16 : public TestCkTileMemoryCopy<ck_tile::f6x16_pk_t, true>
{
};

66
test/ck_tile/memory_copy/test_copy.hpp
View File

@@ -89,7 +89,7 @@ struct TileCopy
constexpr index_t Y1 =
S::Warp_M /
Y2; // number of iterations each warp needs to perform to cover the entire tile window.
printf("Y0: %d Y1: %d Y2: %d X0: %d X1: %d \n", Y0, Y1, Y2, X0, X1);
// printf("Y0: %d Y1: %d Y2: %d X0: %d X1: %d \n", Y0, Y1, Y2, X0, X1);
constexpr auto outer_encoding =
tile_distribution_encoding<sequence<S::WaveGroups>,
tuple<sequence<Y0, Y1, Y2>, sequence<X0, X1>>,
@@ -146,38 +146,38 @@ struct TileCopy
{
if(my_id == warp_id)
{
// if constexpr(AsyncCopy)
// {
// async_load_tile(x_block_lds_write_window, x_block_window);
// // We don't have prefetch here, wait the data back immediately.
// // Wait all asyncload insts complete.
// // Wait all waves synced
// s_waitcnt_barrier<async_copy_fence_cnt>();
// auto lds_tile = load_tile(x_block_lds_read_window);
// // store from registers to DRAM
// store_tile(y_block_window, lds_tile);
// }
// else
// {
// load from DRAM to registers
auto dram_tile = load_tile(x_block_window);
auto ptr = dram_tile.get_thread_buffer().data[0].data_;
printf("block_id: %d, thread_id: %d, value: %d %d %d\n",
get_block_id(),
get_thread_id(),
ptr[0],
ptr[1],
ptr[2]);
// store in lds
store_tile(x_block_lds_write_window, dram_tile);
// Wait all lds write insts complete
// Wait all waves synced
block_sync_lds();
// read from lds to registers
// auto lds_tile = load_tile(x_block_lds_read_window);
// store from registers to DRAM
// store_tile(y_block_window, lds_tile);
// }
if constexpr(AsyncCopy)
{
async_load_tile(x_block_lds_write_window, x_block_window);
// We don't have prefetch here, wait the data back immediately.
// Wait all asyncload insts complete.
// Wait all waves synced
s_waitcnt_barrier<async_copy_fence_cnt>();
auto lds_tile = load_tile(x_block_lds_read_window);
// store from registers to DRAM
store_tile(y_block_window, lds_tile);
}
else
{
// load from DRAM to registers
auto dram_tile = load_tile(x_block_window);
// auto ptr = dram_tile.get_thread_buffer().data[0].data_;
// printf("block_id: %d, thread_id: %d, value: %d %d %d\n",
// get_block_id(),
// get_thread_id(),
// ptr[0],
// ptr[1],
// ptr[2]);
// store in lds
store_tile(x_block_lds_write_window, dram_tile);
// Wait all lds write insts complete
// Wait all waves synced
block_sync_lds();
// read from lds to registers
auto lds_tile = load_tile(x_block_lds_read_window);
// store from registers to DRAM
store_tile(y_block_window, lds_tile);
}
}
move_tile_window(x_block_window, {0, S::Block_N});