Expanded the MmaOpI concept to support an extra int arg to the exec function

Also added a test for this.

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Author: chris-tsiaousis-hpc

projects/composablekernel/include/ck_tile/core/arch/mma/amdgcn_mma.hpp

@@ -6,6 +6,8 @@
 #include "ck_tile/core/arch/arch.hpp"
 #include "ck_tile/core/arch/mma/mma_op_family.hpp"
 #include "ck_tile/core/config.hpp"
+#include "ck_tile/core/numeric/half.hpp"
+#include "ck_tile/core/numeric/pk_fp4.hpp"
 #include "ck_tile/core/numeric/vector_type.hpp"
 #include "ck_tile/core/utility/ignore.hpp"
 
@@ -23,6 +25,20 @@ struct Unsupported;
 #if CK_TILE_CONCEPTS && CK_TILE_CONCEPTS_HEADER
 
 #include <concepts>
+/**
+ * @concept HasExecSignature
+ * @brief  Helper concept for exec signature check.
+ */
+template <typename MmaOp, typename... ExecArgs>
+concept HasExecSignature = requires {
+    {
+        MmaOp::exec(typename MmaOp::AVecType{},
+                    typename MmaOp::BVecType{},
+                    typename MmaOp::CVecType{},
+                    std::declval<ExecArgs>()...)
+    } -> std::convertible_to<typename MmaOp::CVecType>;
+};
+
 /**
  * @concept MmaOpI
  * @brief  Expresses the meta-data interface required for each MmaOp policy.
@@ -48,13 +64,7 @@ concept MmaOpI = requires(MmaOp op) {
     { MmaOp::kCNLane } -> std::convertible_to<unsigned int>;
     { MmaOp::kCM0PerLane } -> std::convertible_to<unsigned int>;
     { MmaOp::kCM1PerLane } -> std::convertible_to<unsigned int>;
-
-    // Static exec function
-    {
-        MmaOp::exec(
-            typename MmaOp::AVecType{}, typename MmaOp::BVecType{}, typename MmaOp::CVecType{})
-    } -> std::convertible_to<typename MmaOp::CVecType>;
-};
+} && (HasExecSignature<MmaOp> || HasExecSignature<MmaOp, int>);
 
 #endif // CK_TILE_CONCEPTS && CK_TILE_CONCEPTS_HEADER
 

projects/composablekernel/test/ck_tile/core/arch/mma/test_amdgcn_sparse_mma.cpp

@@ -71,6 +71,24 @@ TEST(SparseMMATrait, MmaOpTraitsIntegration)
     std::cout << "MmaOpTraits correctly integrates sparse operations" << std::endl;
 }
 
+TEST(SparseMMATrait, TestConceptRequirements)
+{
+#if CK_TILE_CONCEPTS && CK_TILE_CONCEPTS_HEADER
+    using TestSparseMmma = amdgcn_mma<fp16_t,
+                                      fp16_t,
+                                      fp32_t,
+                                      16u,
+                                      16u,
+                                      32u,
+                                      DefaultSparseMfmaCtrlFlags,
+                                      CompilerTargetGfx950,
+                                      MmaOpFamily::SPARSE>;
+    static_assert(MmaOpI<TestSparseMmma>);
+#else
+    GTEST_SKIP() << "Not compiled with concepts. Skipping test.";
+#endif // CK_TILE_CONCEPTS && CK_TILE_CONCEPTS_HEADER
+}
+
 TEST(SparseMMATrait, DenseVsSparseDistinction)
 {
     // Dense MFMA from mfma/mfma_gfx9.hpp