De-couple intertwined compile passes for host/device code

Allow pipeline instantiations on the host side without having to constexpr conditionaly use them within the kernel. Also utilize the warning added to the host/unsupported amdgcn struct.

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

Author: chris-tsiaousis-hpc

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

@@ -124,11 +124,22 @@ struct MmaPipelineBase
     template <typename VecTA, typename VecTB, typename VecTC>
     CK_TILE_DEVICE static decltype(auto) exec(VecTA&& a, VecTB&& b, VecTC&& accum)
     {
-        // TODO: c++20: Call template functions with MmaPipelineOptionFlags directly
-        auto pre = Derived::template preApply<Flags_>(
-            std::forward<VecTA>(a), std::forward<VecTB>(b), std::forward<VecTC>(accum));
-        Derived::execImpl(pre);
-        return Derived::template postApply<Flags_>(std::move(pre));
+        if constexpr(MmaOpTraits<typename Derived::FragWiseMmaOp>::IsSupported)
+        {
+            // TODO: c++20: Call template functions with MmaPipelineOptionFlags directly
+            auto pre = Derived::template preApply<Flags_>(
+                std::forward<VecTA>(a), std::forward<VecTB>(b), std::forward<VecTC>(accum));
+            Derived::execImpl(pre);
+            return Derived::template postApply<Flags_>(std::move(pre));
+        }
+        else
+        {
+            // Return the unsupported exec. This should print a runtime warning. (amdgcn_mma.hpp)
+            // Code should not reach here, but HOST/DEVICE compile passes are
+            // weirdly intertwined and instead of having constexpr in the calling
+            // site (tests) we do this. See also changes by this commit.
+            return Derived::FragWiseMmaOp::exec({}, {}, {});
+        }
     }
 };
 

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

@@ -42,6 +42,8 @@ template <typename MmaOp, typename CompilerTarget, typename Enable = void>
 struct MmaTransformsDefaultSelector
 {
     using SelectedTransforms = MmaDefaultPassThroughTransforms;
+    static_assert(CompilerTarget::TARGET_ID == amdgcn_target_id::HOST,
+                  "Device code should use another specialization.");
 };
 
 #if CK_TILE_CONCEPTS

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp

@@ -37,11 +37,12 @@ template <typename ADataType,
 struct SparseMma : public MmaPipelineBase<static_cast<int>(MmaPipelineOptionFlag::COMPRESS_A), // TODO: c++20: use MmaPipelineOptionFlags directly
                                           SparseMma<ADataType, BDataType, CDataType, FragM, FragN, FragK, CompilerTarget, MmaOp, MmaTransforms>>
 {
-    static_assert(MmaOpTraits<MmaOp>::IsSupported && MmaOpTraits<MmaOp>::IsSparse);
     using Base = MmaPipelineBase<static_cast<int>(MmaPipelineOptionFlag::COMPRESS_A), // TODO: c++20: use MmaPipelineOptionFlags directly
                                  SparseMma<ADataType, BDataType, CDataType, FragM, FragN, FragK, CompilerTarget, MmaOp, MmaTransforms>>;
     // clang-format on
 
+    using FragWiseMmaOp = MmaOp; // Expose the selected MmaOp
+
     // Calculate the uncompressed A vector type
     struct InternalAVecCalculator
     {

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

@@ -185,30 +185,26 @@ __global__ void test_sparse_accum_over_k(void* a, void* b, void* c, void* out)
                                                        CompilerTarget,
                                                        MmaOpFamily::SPARSE>::SelectedOp;
 
-    using MmaTraits = MmaOpTraits<MmaOp>;
+    using Pipeline =
+        SparseMma<AType, BType, CType, WaveTileM, WaveTileN, WaveTileK, CompilerTarget>;
 
-    if constexpr(MmaTraits::IsSupported)
-    {
-        using Pipeline = SparseMma<AType, BType, CType, WaveTileM, WaveTileN, WaveTileK, CompilerTarget>;
-
-        using AVecType = typename Pipeline::AVecType;
-        using BVecType = typename Pipeline::BVecType;
-        using CVecType = typename Pipeline::CVecType;
+    using AVecType = typename Pipeline::AVecType;
+    using BVecType = typename Pipeline::BVecType;
+    using CVecType = typename Pipeline::CVecType;
 
-        static constexpr uint32_t kIters = WaveTileK / MmaOp::kK;
+    static constexpr uint32_t kIters = WaveTileK / MmaOp::kK;
 
-        // Initialize the accumulator
-        CVecType result = *reinterpret_cast<CVecType*>(c);
+    // Initialize the accumulator
+    CVecType result = *reinterpret_cast<CVecType*>(c);
 
-        // Accumulate input AxB over FragK/BlockK iterations
-        for(uint32_t i = 0; i < kIters; ++i)
-        {
-            result = Pipeline::exec(
-                *reinterpret_cast<AVecType*>(a), *reinterpret_cast<BVecType*>(b), result);
-        }
-
-        *reinterpret_cast<CVecType*>(out) = result;
+    // Accumulate input AxB over FragK/BlockK iterations
+    for(uint32_t i = 0; i < kIters; ++i)
+    {
+        result = Pipeline::exec(
+            *reinterpret_cast<AVecType*>(a), *reinterpret_cast<BVecType*>(b), result);
     }
+
+    *reinterpret_cast<CVecType*>(out) = result;
 }
 
 // Live test on real hardware for sparse selection and execution.

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

@@ -23,38 +23,24 @@ template <typename AType,
 __global__ void test_pipeline(void* a, void* b, void* c)
 {
     using CompilerTarget = decltype(get_compiler_target());
-    using MmaOp          = typename MmaDefaultSelector<AType, // TODO: c++20 MmaOpI MmaOp = typename
-                                                              // MmaDefaultSelector<ADataType,
-                                                       BType,
-                                                       CType,
-                                                       WaveTileM,
-                                                       WaveTileN,
-                                                       WaveTileK,
-                                                       CompilerTarget,
-                                                       MmaOpFamily::DENSE>::SelectedOp;
-
-    using MmaTraits = MmaOpTraits<MmaOp>;
-
-    if constexpr(MmaTraits::IsSupported)
-    {
-        using Pipeline = WaveWiseMma<AType,
-                                     BType,
-                                     CType,
-                                     WaveTileM,
-                                     WaveTileN,
-                                     WaveTileK,
-                                     MmaOpFamily::DENSE,
-                                     MmaAccumPolicy::ROW_MAJOR,
-                                     CompilerTarget>;
-
-        using AVecType = typename Pipeline::AVecType;
-        using BVecType = typename Pipeline::BVecType;
-        using CVecType = typename Pipeline::CVecType;
-
-        Pipeline::exec(*reinterpret_cast<AVecType(*)[Pipeline::FragsM][Pipeline::FragsK]>(a),
-                       *reinterpret_cast<BVecType(*)[Pipeline::FragsN][Pipeline::FragsK]>(b),
-                       *reinterpret_cast<CVecType(*)[Pipeline::FragsM][Pipeline::FragsN]>(c));
-    }
+
+    using Pipeline = WaveWiseMma<AType,
+                                 BType,
+                                 CType,
+                                 WaveTileM,
+                                 WaveTileN,
+                                 WaveTileK,
+                                 MmaOpFamily::DENSE,
+                                 MmaAccumPolicy::ROW_MAJOR,
+                                 CompilerTarget>;
+
+    using AVecType = typename Pipeline::AVecType;
+    using BVecType = typename Pipeline::BVecType;
+    using CVecType = typename Pipeline::CVecType;
+
+    Pipeline::exec(*reinterpret_cast<AVecType(*)[Pipeline::FragsM][Pipeline::FragsK]>(a),
+                   *reinterpret_cast<BVecType(*)[Pipeline::FragsN][Pipeline::FragsK]>(b),
+                   *reinterpret_cast<CVecType(*)[Pipeline::FragsM][Pipeline::FragsN]>(c));
 }
 
 TEST(WaveWiseMmaPipeline, testKIter)