[CK_TILE] add load tile async unit test

include/ck_tile/host/check_err.hpp

@@ -679,10 +679,12 @@ std::enable_if_t<(std::is_same_v<ranges::range_value_t<Range>, ranges::range_val
     auto update_err = [&](pk_fp4_raw_t o, pk_fp4_raw_t r, std::size_t index) {
         if(o != r)
         {
-            std::cerr << msg << " out[" << index << "] != ref[" << index
-                      << "]: " << type_convert<float>(pk_fp4_t{o})
-                      << " != " << type_convert<float>(pk_fp4_t{r}) << std::endl;
-            ++err_count;
+            if(err_count++ < ERROR_DETAIL_LIMIT)
+            {
+                std::cerr << msg << " out[" << index << "] != ref[" << index
+                          << "]: " << type_convert<float>(pk_fp4_t{o})
+                          << " != " << type_convert<float>(pk_fp4_t{r}) << std::endl;
+            }
         }
     };
 

test/ck_tile/CMakeLists.txt

@@ -32,3 +32,4 @@ add_subdirectory(atomic_add_op)
 add_subdirectory(fmha)
 add_subdirectory(gemm_tile_engine)
 add_subdirectory(pooling)
+add_subdirectory(async)

test/ck_tile/async/CMakeLists.txt

@@ -0,0 +1 @@
+add_test_executable(test_async_load async_load.cpp)

test/ck_tile/async/async_load.cpp

@@ -0,0 +1,5 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+#include "run_test.inc"
+
+int main() { return run_load_store_tile<ck_tile::pk_fp4_t>() ? EXIT_SUCCESS : EXIT_FAILURE; }

test/ck_tile/async/kernel.hpp

@@ -0,0 +1,167 @@
+#pragma once
+
+#include <iostream>
+#include <string>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/common.hpp"
+#include "ck_tile/host/concat.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+#include "ck_tile/host/stream_utils.hpp"
+#include "ck_tile/core/utility/env.hpp"
+#include "ck_tile/core/utility/type_traits.hpp"
+
+namespace ck_tile {
+
+template <index_t M_Tile_, index_t N_Tile_, index_t M_Warp_, index_t N_Warp_>
+struct TileShape
+{
+    static constexpr index_t M        = M_Tile_;
+    static constexpr index_t N        = N_Tile_;
+    static constexpr index_t Mw       = M_Warp_;
+    static constexpr index_t Nw       = N_Warp_;
+    static constexpr index_t NumWarps = Mw * Nw;
+};
+
+template <typename DataType_, typename TileShape_>
+struct AsyncLSPolicy
+{
+    using DataType = remove_cvref_t<DataType_>;
+    using Shape    = remove_cvref_t<TileShape_>;
+    using RawType =
+        std::conditional_t<std::is_class_v<DataType>, typename DataType::type, DataType>;
+
+    constexpr static index_t max_vector_size = 16;
+    constexpr static index_t warp_size       = 64;
+    constexpr static index_t PackedSize      = numeric_traits<DataType>::PackedSize;
+    constexpr static index_t kMPerBlock      = Shape::M;
+    constexpr static index_t kNPerBlock      = Shape::N / PackedSize;
+
+    CK_TILE_HOST_DEVICE static constexpr auto MakeLdsBlockDescriptor()
+    {
+
+        constexpr auto lds_block_desc_0 =
+            make_naive_tensor_descriptor(make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}),
+                                         make_tuple(number<kNPerBlock>{}, number<1>{}),
+                                         number<16>{},
+                                         number<1>{});
+
+        return lds_block_desc_0;
+    }
+
+    CK_TILE_HOST_DEVICE static constexpr auto MakeDRAMDistribution()
+    {
+        static_assert(Shape::Nw == 1);
+        constexpr index_t N1 = max_vector_size / sizeof(RawType);
+        constexpr index_t N0 = kNPerBlock / N1;
+        constexpr index_t M2 = warp_size / N0;
+        constexpr index_t M1 = Shape::Mw;
+        constexpr index_t M0 = Shape::M / (M1 * M2);
+
+        constexpr auto encoding =
+            tile_distribution_encoding<sequence<1>,
+                                       tuple<sequence<M0, M1, M2>, sequence<N0, N1>>,
+                                       tuple<sequence<1>, sequence<1, 2>>,
+                                       tuple<sequence<1>, sequence<2, 0>>,
+                                       sequence<1, 2>,
+                                       sequence<0, 1>>{};
+        return make_static_tile_distribution(encoding);
+    }
+
+    CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
+    {
+        return sizeof(RawType) * MakeLdsBlockDescriptor().get_element_space_size();
+    }
+    CK_TILE_HOST_DEVICE static constexpr index_t GetVectorSize()
+    {
+        return std::min(static_cast<int>(kNPerBlock),
+                        static_cast<int>(max_vector_size / sizeof(RawType)));
+    }
+};
+
+struct AsyncLSKernelArgs
+{
+    void* a_ptr;
+    void* b_ptr;
+    index_t M;
+    index_t N;
+    index_t stride_A;
+    index_t stride_B;
+};
+
+template <typename Policy_>
+struct AsyncLSKernel
+{
+    using Policy   = remove_cvref_t<Policy_>;
+    using Shape    = typename Policy::Shape;
+    using DataType = typename Policy::DataType;
+    using RawType  = typename Policy::RawType;
+
+    constexpr static int kBlockPerCu    = 1;
+    constexpr static index_t kBlockSize = Shape::NumWarps * get_warp_size();
+    constexpr static index_t PackedSize = Policy::PackedSize;
+
+    CK_TILE_HOST static dim3 BlockSize() { return dim3(kBlockSize); }
+    CK_TILE_HOST static dim3 GridSize(index_t M, index_t N)
+    {
+        const index_t GridDimX = (M + Shape::M - 1) / Shape::M;
+        const index_t GridDimY = (N + Shape::N - 1) / Shape::N;
+        return dim3(GridDimX, GridDimY, 1);
+    }
+
+    CK_TILE_DEVICE auto operator()(AsyncLSKernelArgs kargs) const -> void
+    {
+        const index_t i_m = amd_wave_read_first_lane(blockIdx.x * Policy::kMPerBlock);
+        const index_t i_n = amd_wave_read_first_lane(blockIdx.y * Policy::kNPerBlock);
+
+        RawType* a_ptr = static_cast<RawType*>(kargs.a_ptr);
+        RawType* b_ptr = static_cast<RawType*>(kargs.b_ptr);
+
+        // allocate LDS
+        __shared__ RawType smem_ptr_0[Policy::GetSmemSize()];
+
+        const auto& a_tensor_view = make_naive_tensor_view<address_space_enum::global>(
+            a_ptr,
+            make_tuple(kargs.M, kargs.N / PackedSize),
+            make_tuple(kargs.stride_A / PackedSize, 1),
+            number<Policy::GetVectorSize()>{},
+            number<1>{});
+
+        const auto& b_tensor_view =
+            make_naive_tensor_view<address_space_enum::global, memory_operation_enum::set>(
+                b_ptr,
+                make_tuple(kargs.M, kargs.N / PackedSize),
+                make_tuple(kargs.stride_B / PackedSize, 1),
+                number<Policy::GetVectorSize()>{},
+                number<1>{});
+
+        auto a_block_window =
+            make_tile_window(a_tensor_view,
+                             make_tuple(number<Policy::kMPerBlock>{}, number<Policy::kNPerBlock>{}),
+                             {i_m, i_n},
+                             Policy::MakeDRAMDistribution());
+
+        auto b_block_window =
+            make_tile_window(b_tensor_view,
+                             make_tuple(number<Policy::kMPerBlock>{}, number<Policy::kNPerBlock>{}),
+                             {i_m, i_n});
+
+        auto lds_0_tensor_view =
+            make_tensor_view<address_space_enum::lds>(smem_ptr_0, Policy::MakeLdsBlockDescriptor());
+        auto lds_0_window =
+            make_tile_window(lds_0_tensor_view,
+                             make_tuple(number<Policy::kMPerBlock>{}, number<Policy::kNPerBlock>{}),
+                             {0, 0},
+                             Policy::MakeDRAMDistribution());
+#if 0
+        auto dram_tile = load_tile(a_block_window);
+        store_tile(lds_0_window, dram_tile);
+#else
+        async_load_tile(lds_0_window, a_block_window);
+        block_sync_lds();
+#endif
+        auto lds_tile = load_tile(lds_0_window);
+        store_tile(b_block_window, lds_tile);
+    }
+};
+} // namespace ck_tile

test/ck_tile/async/run_test.inc

@@ -0,0 +1,83 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <hip/hip_runtime.h>
+
+#include <cstring>
+#include <iostream>
+#include <ostream>
+#include <string>
+#include <tuple>
+
+#include "ck_tile/host.hpp"
+#include "kernel.hpp"
+
+template <typename DataType>
+float load_store_tile(const ck_tile::AsyncLSKernelArgs& args, const ck_tile::stream_config& s)
+
+{
+    constexpr int kBlockPerCu = 1;
+
+    constexpr ck_tile::index_t M_Tile = 32;
+    constexpr ck_tile::index_t N_Tile = 256;
+
+    constexpr ck_tile::index_t M_Warp = 2;
+    constexpr ck_tile::index_t N_Warp = 1;
+
+    using Shape = ck_tile::TileShape<M_Tile, N_Tile, M_Warp, N_Warp>;
+
+    using Policy = ck_tile::AsyncLSPolicy<DataType, Shape>;
+
+    using Kernel = ck_tile::AsyncLSKernel<Policy>;
+
+    const dim3 grids  = Kernel::GridSize(args.M, args.N);
+    const dim3 blocks = Kernel::BlockSize();
+
+    float ave_time = ck_tile::launch_kernel(
+        s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, args));
+
+    std::cout << "Run Load_Store_Tile with kernel " << M_Tile << "x" << N_Tile << ", input "
+              << args.M << "x" << args.N << ": " << ave_time << " ms, \n";
+
+    return ave_time;
+}
+
+template <typename DataType>
+float invoke_load_store_tile(ck_tile::DeviceMem& a_dev_buf,
+                             ck_tile::DeviceMem& b_dev_buf,
+                             ck_tile::index_t M,
+                             ck_tile::index_t N)
+{
+    auto args = ck_tile::AsyncLSKernelArgs{
+        a_dev_buf.GetDeviceBuffer(), b_dev_buf.GetDeviceBuffer(), M, N, N, N};
+    auto sc        = ck_tile::stream_config{nullptr, true, 1, 0, 1, true, true, 1};
+    float ave_time = load_store_tile<DataType>(args, sc);
+
+    return ave_time;
+}
+
+template <typename DataType>
+bool run_load_store_tile()
+{
+    constexpr size_t m = 64;
+    constexpr size_t n = 512;
+    constexpr size_t s = 1;
+
+    ck_tile::HostTensor<DataType> a_m_n({m, n}, {n, s});
+    ck_tile::HostTensor<DataType> b_m_n({m, n}, {n, s});
+
+    ck_tile::FillMonotonicSeq<DataType>{}(a_m_n);
+    b_m_n.SetZero();
+
+    ck_tile::DeviceMem a_m_n_dev_buf(a_m_n.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem b_m_n_dev_buf(b_m_n.get_element_space_size_in_bytes());
+    a_m_n_dev_buf.ToDevice(a_m_n.data());
+
+    invoke_load_store_tile<DataType>(a_m_n_dev_buf, b_m_n_dev_buf, m, n);
+
+    b_m_n_dev_buf.FromDevice(b_m_n.data());
+
+    bool pass = ck_tile::check_err(b_m_n, a_m_n);
+
+    return pass;
+}