Merge commit 'aad4cf098511b3f58c5bd3c32e4534d438f7539c' into develop

test/CMakeLists.txt

@@ -258,6 +258,7 @@ add_subdirectory(conv_util)
 add_subdirectory(reference_conv_fwd)
 add_subdirectory(gemm)
 add_subdirectory(gemm_add)
+add_subdirectory(gemm_bias_add_reduce)
 add_subdirectory(gemm_blockscale_wp)
 add_subdirectory(gemm_layernorm)
 add_subdirectory(gemm_multi_abd)

test/gemm_bias_add_reduce/CMakeLists.txt

@@ -0,0 +1,9 @@
+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+if(GPU_TARGETS MATCHES "gfx9|gfx11|gfx12")
+    add_gtest_executable(test_gemm_bias_add_reduce_fp16 test_gemm_bias_add_reduce_fp16.cpp)
+    if(result EQUAL 0)
+        target_link_libraries(test_gemm_bias_add_reduce_fp16 PRIVATE utility device_gemm_bias_add_reduce_instance)
+    endif()
+endif()

test/gemm_bias_add_reduce/test_gemm_bias_add_reduce_fp16.cpp

@@ -0,0 +1,106 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include <tuple>
+
+#include "gtest/gtest.h"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "test_gemm_common.hpp"
+
+using F16 = ck::half_t;
+using F32 = float;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+namespace {
+
+template <typename X, typename Y>
+struct tuple_concat;
+
+template <typename... Xs, typename... Ys>
+struct tuple_concat<std::tuple<Xs...>, std::tuple<Ys...>>
+{
+    using type = std::tuple<Xs..., Ys...>;
+};
+
+} // namespace
+
+template <typename Tuple>
+class TestGemmBiasAddReduce_FP16_MK_NK
+    : public ck::test::TestGemmBiasAddReduceCommon<
+          typename tuple_concat<std::tuple<Row, Col>, Tuple>::type>
+{
+};
+
+template <typename Tuple>
+class TestGemmBiasAddReduce_FP16_MK_KN
+    : public ck::test::TestGemmBiasAddReduceCommon<
+          typename tuple_concat<std::tuple<Row, Row>, Tuple>::type>
+{
+};
+
+template <typename Tuple>
+class TestGemmBiasAddReduce_FP16_KM_KN
+    : public ck::test::TestGemmBiasAddReduceCommon<
+          typename tuple_concat<std::tuple<Col, Row>, Tuple>::type>
+{
+};
+
+template <typename Tuple>
+class TestGemmBiasAddReduce_FP16_KM_NK
+    : public ck::test::TestGemmBiasAddReduceCommon<
+          typename tuple_concat<std::tuple<Col, Col>, Tuple>::type>
+{
+};
+
+// clang-format off
+using KernelTypes = ::testing::Types<
+    std::tuple< F16, F16, F16, F16, F16, F32>
+   >;
+// clang-format on
+
+TYPED_TEST_SUITE(TestGemmBiasAddReduce_FP16_MK_NK, KernelTypes);
+TYPED_TEST_SUITE(TestGemmBiasAddReduce_FP16_MK_KN, KernelTypes);
+TYPED_TEST_SUITE(TestGemmBiasAddReduce_FP16_KM_KN, KernelTypes);
+TYPED_TEST_SUITE(TestGemmBiasAddReduce_FP16_KM_NK, KernelTypes);
+
+TYPED_TEST(TestGemmBiasAddReduce_FP16_MK_NK, Regular)
+{
+    std::vector<int> Ms{512};
+    constexpr int N = 512;
+    constexpr int K = 1024;
+
+    for(int M : Ms)
+        this->Run(M, N, K);
+}
+
+TYPED_TEST(TestGemmBiasAddReduce_FP16_MK_KN, Regular)
+{
+    std::vector<int> Ms{512};
+    constexpr int N = 1024;
+    constexpr int K = 1024;
+
+    for(int M : Ms)
+        this->Run(M, N, K);
+}
+
+TYPED_TEST(TestGemmBiasAddReduce_FP16_KM_KN, Regular)
+{
+    std::vector<int> Ms{256};
+    constexpr int N = 512;
+    constexpr int K = 1024;
+
+    for(int M : Ms)
+        this->Run(M, N, K);
+}
+
+TYPED_TEST(TestGemmBiasAddReduce_FP16_KM_NK, Regular)
+{
+    std::vector<int> Ms{256};
+    constexpr int N = 1024;
+    constexpr int K = 1024;
+
+    for(int M : Ms)
+        this->Run(M, N, K);
+}

test/gemm_bias_add_reduce/test_gemm_common.hpp

@@ -0,0 +1,61 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include "gtest/gtest.h"
+#include "ck/ck.hpp"
+#include "profiler/profile_gemm_bias_add_reduce_impl.hpp"
+
+namespace ck {
+namespace test {
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using F32 = float;
+
+template <typename Tuple>
+class TestGemmBiasAddReduceCommon : public ::testing::Test
+{
+    protected:
+    using ALayout        = std::tuple_element_t<0, Tuple>;
+    using BLayout        = std::tuple_element_t<1, Tuple>;
+    using CLayout        = Row;
+    using ADataType      = std::tuple_element_t<2, Tuple>;
+    using BDataType      = std::tuple_element_t<3, Tuple>;
+    using CDataType      = std::tuple_element_t<4, Tuple>;
+    using BiasDataType   = std::tuple_element_t<5, Tuple>;
+    using D0DataType     = std::tuple_element_t<6, Tuple>;
+    using ReduceDataType = std::tuple_element_t<7, Tuple>;
+
+    public:
+    static constexpr bool verify_     = true;
+    static constexpr int init_method_ = 1; // integer value initialization
+    static constexpr bool log_        = false;
+    static constexpr bool bench_      = false; // measure kernel performance
+
+    void Run(const int M, const int N, const int K)
+    {
+        bool all_success = true;
+
+        int StrideA  = std::is_same_v<remove_cvref_t<ALayout>, Row> ? K : M;
+        int StrideB  = std::is_same_v<remove_cvref_t<BLayout>, Row> ? N : K;
+        int StrideD0 = std::is_same_v<remove_cvref_t<CLayout>, Row> ? N : M;
+        int StrideC  = std::is_same_v<CLayout, Row> ? N : M;
+
+        all_success =
+            all_success &
+            ck::profiler::profile_gemm_bias_add_reduce_impl<ADataType,
+                                                            BDataType,
+                                                            CDataType,
+                                                            BiasDataType,
+                                                            D0DataType,
+                                                            ReduceDataType,
+                                                            ALayout,
+                                                            BLayout,
+                                                            CLayout>(
+                verify_, init_method_, log_, bench_, M, N, K, StrideA, StrideB, StrideC, StrideD0);
+
+        EXPECT_TRUE(all_success);
+    }
+};
+
+} // namespace test
+} // namespace ck

test/grouped_gemm/CMakeLists.txt

@@ -12,6 +12,12 @@ if (CK_USE_XDL OR CK_USE_WMMA)
         target_link_libraries(test_grouped_gemm_splitk PRIVATE utility device_grouped_gemm_instance)
         add_dependencies(test_grouped_gemm test_grouped_gemm_splitk)
     endif()
+
+    add_gtest_executable(test_grouped_gemm_fastgelu test_grouped_gemm_fastgelu.cpp)
+    if(result EQUAL 0)
+        target_link_libraries(test_grouped_gemm_fastgelu PRIVATE utility device_grouped_gemm_fastgelu_instance)
+        add_dependencies(test_grouped_gemm test_grouped_gemm_fastgelu)
+    endif()
 endif()
 
 add_gtest_executable(test_grouped_gemm_interface test_grouped_gemm_interface_xdl.cpp)

test/grouped_gemm/test_grouped_gemm_fastgelu.cpp

@@ -0,0 +1,62 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#include <tuple>
+#include <vector>
+
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
+#include "ck/utility/data_type.hpp"
+
+#include "gtest/gtest.h"
+#include "test_grouped_gemm_util.hpp"
+
+ck::index_t param_mask     = 0xffffff;
+ck::index_t instance_index = -1;
+
+using F16  = ck::half_t;
+using BF16 = ck::bhalf_t;
+using F8   = ck::f8_t;
+using I8   = int8_t;
+
+using AElementOp   = ck::tensor_operation::element_wise::PassThrough;
+using BElementOp   = ck::tensor_operation::element_wise::PassThrough;
+using CDEElementOp = ck::tensor_operation::element_wise::FastGelu;
+
+using Row = ck::tensor_layout::gemm::RowMajor;
+using Col = ck::tensor_layout::gemm::ColumnMajor;
+
+template <typename Tuple>
+class TestGroupedGemm : public ck::test::TestGroupedGemm<Tuple, true>
+{
+};
+
+// clang-format off
+using KernelTypes = ::testing::Types<
+    std::tuple<     Row, Row, Row, F16, F16, F16, AElementOp, BElementOp, CDEElementOp>,
+    std::tuple<     Row, Col, Row, F16, F16, F16, AElementOp, BElementOp, CDEElementOp>,
+    std::tuple<     Col, Row, Row, F16, F16, F16, AElementOp, BElementOp, CDEElementOp>,
+    std::tuple<     Col, Col, Row, F16, F16, F16, AElementOp, BElementOp, CDEElementOp>
+>;
+// clang-format on
+
+TYPED_TEST_SUITE(TestGroupedGemm, KernelTypes);
+
+#include "test_grouped_gemm_ut_cases.inc"
+int main(int argc, char** argv)
+{
+    testing::InitGoogleTest(&argc, argv);
+    if(argc == 1) {}
+    else if(argc == 3)
+    {
+        param_mask     = strtol(argv[1], nullptr, 0);
+        instance_index = atoi(argv[2]);
+    }
+    else
+    {
+        std::cout << "Usage of " << argv[0] << std::endl;
+        std::cout << "Arg1,2: param_mask instance_index(-1 means all)" << std::endl;
+    }
+    return RUN_ALL_TESTS();
+}

test/grouped_gemm/test_grouped_gemm_ut_cases.inc

@@ -65,12 +65,11 @@ TYPED_TEST(TestGroupedGemm, MNKPadded)
 
 TYPED_TEST(TestGroupedGemm, TestLargeKBatch)
 {
-    // gfx11 does not support split-K due to missing atomic add for fp16/bf16
-    // Technically, we could still run the tests for fp32, but we currently don't have instances for
-    // it so we disable it entirely
-    if(ck::is_gfx11_supported())
-        GTEST_SKIP() << "Split-K not supported for FP16/BF16 on GFX11 due to missing atomic add "
-                        "instructions";
+    // In some cases Split K is not supported. Running this test would fail since no instance will
+    // be supported, so we skip the test
+    if(!this->IsSplitKSupported())
+        GTEST_SKIP() << "Split-K not supported for for the current configuration (FP16/BF16 on "
+                        "GFX11, or using CDE element-wise operation)";
 
     const std::vector<int> Ms{188, 210};
     constexpr int N = 768;

test/grouped_gemm/test_grouped_gemm_util.hpp

@@ -7,11 +7,14 @@
 #include <string>
 #include <sstream>
 #include <tuple>
+#include <type_traits>
 #include <vector>
 #include <gtest/gtest.h>
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
 #include "profiler/profile_grouped_gemm_impl.hpp"
 
 extern ck::index_t param_mask;
@@ -32,16 +35,46 @@ std::string serialize_range(const Range& range)
     return std::string(str.begin(), str.end() - 2);
 }
 
+// Helper primary template (will be specialized on the boolean)
+template <std::size_t N,
+          typename Tuple,
+          typename Default,
+          bool InRange = (N < std::tuple_size_v<std::remove_reference_t<Tuple>>)>
+struct tuple_element_or_impl;
+
+// Specialization for the in-range case: use std::tuple_element_t
+template <std::size_t N, typename Tuple, typename Default>
+struct tuple_element_or_impl<N, Tuple, Default, true>
+{
+    using type = std::tuple_element_t<N, std::remove_reference_t<Tuple>>;
+};
+
+// Specialization for the out-of-range case: use Default
+template <std::size_t N, typename Tuple, typename Default>
+struct tuple_element_or_impl<N, Tuple, Default, false>
+{
+    using type = Default;
+};
+
+// User-facing alias
+template <std::size_t N, typename Tuple, typename Default>
+using tuple_element_or_t = typename tuple_element_or_impl<N, Tuple, Default>::type;
+
 template <typename Tuple, bool FailIfNoSupportedInstances = false>
 class TestGroupedGemm : public testing::Test
 {
     protected:
-    using ALayout   = std::tuple_element_t<0, Tuple>;
-    using BLayout   = std::tuple_element_t<1, Tuple>;
-    using ELayout   = std::tuple_element_t<2, Tuple>;
-    using ADataType = std::tuple_element_t<3, Tuple>;
-    using BDataType = std::tuple_element_t<4, Tuple>;
-    using EDataType = std::tuple_element_t<5, Tuple>;
+    using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+    using ALayout      = std::tuple_element_t<0, Tuple>;
+    using BLayout      = std::tuple_element_t<1, Tuple>;
+    using ELayout      = std::tuple_element_t<2, Tuple>;
+    using ADataType    = std::tuple_element_t<3, Tuple>;
+    using BDataType    = std::tuple_element_t<4, Tuple>;
+    using EDataType    = std::tuple_element_t<5, Tuple>;
+    using AElementOp   = tuple_element_or_t<6, Tuple, PassThrough>;
+    using BElementOp   = tuple_element_or_t<7, Tuple, PassThrough>;
+    using CDEElementOp = tuple_element_or_t<8, Tuple, PassThrough>;
 
     using Row = ck::tensor_layout::gemm::RowMajor;
     using Col = ck::tensor_layout::gemm::ColumnMajor;
@@ -57,15 +90,25 @@ class TestGroupedGemm : public testing::Test
     bool fail_if_no_supported_instances_ = FailIfNoSupportedInstances;
     std::vector<int> k_batches_;
 
-    void SetUp() override
+    bool IsSplitKSupported()
     {
+        // gfx11 does not support split-K due to missing atomic add for fp16/bf16
+        // Technically, we could still use split-K for fp32, but we currently don't have
+        // instances for it so we disable it entirely
         constexpr bool require_16bit_atomic_add =
             std::is_same_v<EDataType, ck::half_t> || std::is_same_v<EDataType, ck::bhalf_t>;
-        if(require_16bit_atomic_add && ck::is_gfx11_supported())
+        bool missing_atomic_add = require_16bit_atomic_add && ck::is_gfx11_supported();
+
+        // CDE element operators are not supported in combination with split K
+        constexpr bool has_cde_element_operator = !std::is_same_v<CDEElementOp, PassThrough>;
+
+        return !missing_atomic_add && !has_cde_element_operator;
+    }
+
+    void SetUp() override
+    {
+        if(!IsSplitKSupported())
         {
-            // gfx11 does not support split-K due to missing atomic add for fp16/bf16
-            // Technically, we could still use split-K for fp32, but we currently don't have
-            // instances for it so we disable it entirely
             k_batches_ = {1};
         }
         else
@@ -147,21 +190,24 @@ class TestGroupedGemm : public testing::Test
                                                     float,
                                                     ALayout,
                                                     BLayout,
-                                                    ELayout>(verify_,
-                                                             init_method_,
-                                                             log_,
-                                                             bench_,
-                                                             Ms,
-                                                             Ns,
-                                                             Ks,
-                                                             StrideAs,
-                                                             StrideBs,
-                                                             StrideCs,
-                                                             kbatches,
-                                                             n_warmup_,
-                                                             n_iter_,
-                                                             instance_index,
-                                                             fail_if_no_supported_instances_);
+                                                    ELayout,
+                                                    AElementOp,
+                                                    BElementOp,
+                                                    CDEElementOp>(verify_,
+                                                                  init_method_,
+                                                                  log_,
+                                                                  bench_,
+                                                                  Ms,
+                                                                  Ns,
+                                                                  Ks,
+                                                                  StrideAs,
+                                                                  StrideBs,
+                                                                  StrideCs,
+                                                                  kbatches,
+                                                                  n_warmup_,
+                                                                  n_iter_,
+                                                                  instance_index,
+                                                                  fail_if_no_supported_instances_);
         EXPECT_TRUE(pass);
     }
 };