gtest to test atomic_add for a tensor (#2716)

* Code drop for gtest to test atomic_add for a tensor * Adding additional test cases * Fix clang errors in CI pipeline * Updated test cases * Fix the Navi card atomic add problem * solved the define problem * add more print out traces * Fix the float4 missing case * solved the gfx9 errors * Address the comment --------- Co-authored-by: Khushbu <khuagarw@amd.com> Co-authored-by: Thomas Ning <Thomas.Ning@amd.com>
2026-05-02 12:41:26 +00:00 · 2025-09-03 15:32:54 -07:00
parent 0282d98412
commit 80ce6a573b
5 changed files with 529 additions and 0 deletions
--- a/test/ck_tile/atomic_add_op/test_atomic.hpp
+++ b/test/ck_tile/atomic_add_op/test_atomic.hpp
@@ -0,0 +1,115 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/common.hpp"
+#include "ck_tile/ops/common/tensor_layout.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+
+namespace ck_tile {
+
+template <typename BlockWaves, typename BlockTile, typename WaveTile, typename Vector>
+struct AtomicKernelShape
+{
+    static constexpr index_t MWarps = BlockWaves::at(number<0>{});
+    static constexpr index_t NWarps = BlockWaves::at(number<1>{});
+
+    static constexpr index_t Block_M = BlockTile::at(number<0>{});
+    static constexpr index_t Block_N = BlockTile::at(number<1>{});
+
+    static constexpr index_t Warp_M = WaveTile::at(number<0>{});
+    static constexpr index_t Warp_N = WaveTile::at(number<1>{});
+
+    static constexpr index_t Vector_M = Vector::at(number<0>{});
+    static constexpr index_t Vector_N = Vector::at(number<1>{});
+
+    static constexpr index_t ThreadPerWarp_M = Warp_M / Vector_M;
+    static constexpr index_t ThreadPerWarp_N = Warp_N / Vector_N;
+
+    static constexpr index_t WarpPerBlock_M = MWarps;
+    static constexpr index_t WarpPerBlock_N = NWarps;
+
+    static constexpr index_t Repeat_M = Block_M / (WarpPerBlock_M * Warp_M);
+    static constexpr index_t Repeat_N = Block_N / (WarpPerBlock_N * Warp_N);
+
+    static constexpr index_t WaveNum = reduce_on_sequence(BlockWaves{}, multiplies{}, number<1>{});
+
+    static constexpr index_t BlockSize = get_warp_size() * WaveNum;
+};
+
+template <typename XDataType_, typename BlockShape_>
+struct AtomicKernelProblem
+{
+    using XDataType  = remove_cvref_t<XDataType_>;
+    using BlockShape = remove_cvref_t<BlockShape_>;
+};
+
+template <typename Problem_>
+struct AtomicKernel
+{
+    using Problem   = remove_cvref_t<Problem_>;
+    using XDataType = typename Problem::XDataType;
+
+    static constexpr index_t kBlockSize = Problem::BlockShape::BlockSize;
+
+    template <typename Problem>
+    CK_TILE_DEVICE static constexpr auto MakeTileDistribution()
+    {
+        using S = typename Problem::BlockShape;
+
+        constexpr index_t warp_size = get_warp_size();
+
+        constexpr index_t X0 = S::ThreadPerWarp_N;
+        constexpr index_t X1 = S::Vector_N;
+
+        constexpr index_t Y0 = S::WaveNum;
+        constexpr index_t Y2 = warp_size / X0;
+        constexpr index_t Y1 = S::Warp_M / Y2;
+
+        constexpr auto encoding =
+            tile_distribution_encoding<sequence<1>,
+                                       tuple<sequence<Y0, Y1, Y2>, sequence<X0, X1>>,
+                                       tuple<sequence<0, 1>, sequence<1, 2>>,
+                                       tuple<sequence<0, 0>, sequence<2, 0>>,
+                                       sequence<1, 2>,
+                                       sequence<1, 1>>{};
+
+        return make_static_tile_distribution(encoding);
+    }
+
+    CK_TILE_DEVICE void operator()(XDataType* input, index_t M, index_t N) const
+    {
+        using S = typename Problem::BlockShape;
+
+        constexpr auto block_dims = make_tuple(number<S::Block_M>{}, number<S::Block_N>{});
+
+        const index_t iM = __builtin_amdgcn_readfirstlane(get_block_id() * S::Block_M);
+
+        const auto input_view =
+            make_naive_tensor_view<address_space_enum::global, memory_operation_enum::atomic_add>(
+                input, make_tuple(M, N), make_tuple(N, 1), number<S::Vector_N>{}, number<1>{});
+        auto input_window = make_tile_window(input_view, block_dims, {iM, 0});
+
+        const index_t num_iterations =
+            __builtin_amdgcn_readfirstlane(integer_divide_ceil(N, S::Block_N));
+        using tmp_tile =
+            decltype(make_static_distributed_tensor<XDataType>(MakeTileDistribution<Problem>()));
+
+        for(int iN = __builtin_amdgcn_readfirstlane(0); iN < num_iterations; iN++)
+        {
+            tmp_tile add_value_tile;
+            tile_elementwise_inout([](auto& c) { c = static_cast<XDataType>(1.0f); },
+                                   add_value_tile);
+
+            update_tile(input_window, add_value_tile);
+            __syncthreads();
+
+            move_tile_window(input_window, {0, S::Block_N});
+        }
+    }
+};
+
+} // namespace ck_tile