Add the structure for testing the Sinkhorn-Knopp kernel

test/ck_tile/CMakeLists.txt

@@ -40,3 +40,4 @@ add_subdirectory(fmha)
 add_subdirectory(gemm_tile_engine)
 add_subdirectory(pooling)
 add_subdirectory(grouped_conv)
+add_subdirectory(sinkhorn)

test/ck_tile/sinkhorn/CMakeLists.txt

@@ -0,0 +1,8 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+if(GPU_TARGETS MATCHES "gfx9|gfx11|gfx12")
+    add_gtest_executable(test_ck_tile_sinkhorn test_sinkhorn.cpp)
+    target_link_libraries(test_ck_tile_sinkhorn)
+endif()
+

test/ck_tile/sinkhorn/test_sinkhorn.cpp

@@ -0,0 +1,42 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include <gtest/gtest.h>
+#include <vector>
+#include <cmath>
+#include <tuple>
+#include <iostream>
+#include <cstring>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/ops/sinkhorn_knopp.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+
+#include "test_sinkhorn_impl.hpp"
+
+// Shape parameters for different test configurations
+using Shape1_BlockWarps = ck_tile::sequence<4, 1>;
+using Shape1_BlockTile  = ck_tile::sequence<128, 128>;
+using Shape1_WarpTile   = ck_tile::sequence<32, 128>;
+using Shape1_ThreadTile = ck_tile::sequence<8, 8>;
+
+// Test configurations for different data types and input size
+using TestConfig_F16 = std::tuple<
+                                      ck_tile::half_t, // XDataType
+                                      float,           // ComputeDataType
+                                      float,           // YDataType
+                                      Shape1_BlockWarps,
+                                      Shape1_BlockTile,
+                                      Shape1_WarpTile,
+                                      Shape1_ThreadTile>;
+
+
+using TestTypes = ::testing::Types<TestConfig_F16>;
+
+TYPED_TEST_SUITE(TestCkTileSinkHorn, TestTypes);
+
+TYPED_TEST(TestCkTileSinkHorn, Test_4x4)
+{
+    this->RunGenericTest({4, 4}, 10);
+}

test/ck_tile/sinkhorn/test_sinkhorn_impl.hpp

@@ -0,0 +1,97 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include <gtest/gtest.h>
+#include <vector>
+#include <cmath>
+#include <tuple>
+#include <iostream>
+#include <cstring>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/ops/sinkhorn_knopp.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+
+template <typename Tuple>
+class TestCkTileSinkHorn: public ::testing::Test
+{
+    protected:
+    using XDataType               = std::tuple_element_t<0, Tuple>;
+    using ComputeDataType         = std::tuple_element_t<1, Tuple>;
+    using YDataType               = std::tuple_element_t<2, Tuple>;
+    using BlockWarps_             = std::tuple_element_t<3, Tuple>;
+    using BlockTile_              = std::tuple_element_t<4, Tuple>;
+    using WarpTile_               = std::tuple_element_t<5, Tuple>;
+    using ThreadTile_             = std::tuple_element_t<6, Tuple>;
+
+    using TestSinkhornShape =
+        ck_tile::SinkhornKnoppShape<
+            BlockWarps_, 
+            BlockTile_,
+            WarpTile_,
+            ThreadTile_
+            >;
+
+    void RunGenericTest(const std::vector<ck_tile::index_t>& input_shape, const int max_iterations)
+    {
+
+        SinkhornKnoppArgs args{};
+        args.input_m       = static_cast<ck_tile::index_t>(input_shape[0]);
+        args.max_iterations = max_iterations;
+
+        auto default_stride = {args.input_m, 1};
+
+        ck_tile::HostTensor<XDataType> h_x(input_shape, default_stride);
+        ck_tile::HostTensor<YDataType> h_y(input_shape, default_stride);
+
+        ck_tile::FillUniformDistribution<XDataType>{-5.f, 5.f}(h_x);
+
+        auto buffer_size = h_xs.get_element_space_size_in_bytes();
+        ck_tile::DeviceMem d_x_mem(h_x.get_element_space_size_in_bytes());
+        ck_tile::DeviceMem d_y_mem(output_buffer_size);
+
+        args.p_x = static_cast<void*>(d_x_mem.GetDeviceBuffer());
+        args.out = static_cast<void*>(d_y_mem.GetDeviceBuffer());
+
+        d_x_mem.ToDevice(h_x.data());
+        d_y_mem.ToDevice(h_y.data());
+
+        using Problem = ck_tile::SinkhornKnoppProblem<XDataType,
+                                                      YDataType,
+                                                      TestSinkhornShape,
+                                                      ComputeDataType
+                                                     >;
+        using Kernel = ck_tile::SinkhornKnoppKernelDummyNonStochastic<
+            Problem,
+            ck_tile::SinkhornKnoppPolicy>;
+
+        // Launch configuration
+        const ck_tile::index_t kBlockSize      = Kernel::BlockSize();
+        constexpr ck_tile::index_t kBlockPerCu = 1;
+
+        ck_tile::index_t kGridSize = 1 // TODO
+
+        //TODO
+        // if(!Kernel::IsSupportedArgument())
+        // {
+        //     throw std::runtime_error("Wrong! Arguments not supported!\n");
+        // }
+
+        ck_tile::launch_kernel(
+            ck_tile::stream_config{nullptr, false, 0},
+            ck_tile::make_kernel<kBlockPerCu>(Kernel{},
+                                              kGridSize,
+                                              kBlockSize,
+                                              0,
+                                              args));
+
+        // Reference computation
+        // TODO
+
+        // Transfer data from device and check error for each operation
+        // TODO
+
+        EXPECT_TRUE(true); // TODO
+    } 
+};