Add ck mha fp16 solver

CMakeLists.txt

@@ -327,7 +327,7 @@ add_compile_definitions($<$<COMPILE_LANGUAGE:CXX>:HIP_COMPILER_FLAGS=${HIP_COMPI
 # HIP
 if( MIOPEN_BACKEND STREQUAL "HIP" OR MIOPEN_BACKEND STREQUAL "HIPOC" OR MIOPEN_BACKEND STREQUAL "HIPNOGPU")
     if(MIOPEN_USE_COMPOSABLEKERNEL)
-        find_package(composable_kernel 1.0.0 COMPONENTS device_other_operations device_gemm_operations device_conv_operations device_reduction_operations)
+        find_package(composable_kernel 1.0.0 COMPONENTS device_other_operations device_gemm_operations device_conv_operations device_reduction_operations device_mha_operations)
     endif()
     if( MIOPEN_BACKEND STREQUAL "HIPNOGPU")
         set(MIOPEN_MODE_NOGPU 1)

requirements.txt

@@ -7,5 +7,5 @@ nlohmann/[email protected] -DJSON_MultipleHeaders=ON -DJSON_BuildTests=Off
 ROCm/[email protected]
 ROCm/[email protected]
 ROCm/frugally-deep@9683d557eb672ee2304f80f6682c51242d748a50
-ROCm/composable_kernel@9c0811f39a2262dbe4d71b81898187951c1e11ba -DCMAKE_BUILD_TYPE=Release -DINSTANCES_ONLY=ON
+ROCm/composable_kernel@294cb823142a815170cf1faa63e01a431a557a04 -DCMAKE_BUILD_TYPE=Release -DINSTANCES_ONLY=ON
 google/[email protected]

src/CMakeLists.txt

@@ -300,6 +300,7 @@ set( MIOpen_Source
     solver/layernorm/forward_layernorm2d_ck.cpp
     solver/layernorm/forward_layernorm4d_ck.cpp
     solver/layernorm/forward_t5layernorm.cpp
+    solver/mha/mha_ck_fa_v2_solver_forward.cpp
     solver/mha/mha_solver_backward.cpp
     solver/mha/mha_solver_forward.cpp
     solver/pooling/forward2d.cpp
@@ -822,7 +823,7 @@ target_include_directories(MIOpen PUBLIC
 )
 
 if(MIOPEN_USE_COMPOSABLEKERNEL)
-set(MIOPEN_CK_LINK_FLAGS composable_kernel::device_other_operations composable_kernel::device_gemm_operations composable_kernel::device_conv_operations composable_kernel::device_reduction_operations hip::host)
+set(MIOPEN_CK_LINK_FLAGS composable_kernel::device_other_operations composable_kernel::device_gemm_operations composable_kernel::device_conv_operations composable_kernel::device_reduction_operations composable_kernel::device_mha_operations hip::host)
 endif()
 
 if(WIN32)

src/include/miopen/mha/solvers.hpp

@@ -77,6 +77,28 @@ struct MhaBackward final : MhaSolver
     MIOPEN_INTERNALS_EXPORT bool MayNeedWorkspace() const override;
 };
 
+struct MhaCKFlashAttentionV2Forward final : MhaSolver
+{
+    const std::string& SolverDbId() const override
+    {
+        return GetSolverDbId<MhaCKFlashAttentionV2Forward>();
+    }
+
+    MIOPEN_INTERNALS_EXPORT bool
+    IsApplicable(const ExecutionContext& context,
+                 const miopen::mha::ProblemDescription& problem) const override;
+
+    MIOPEN_INTERNALS_EXPORT ConvSolution
+    GetSolution(const ExecutionContext& context,
+                const miopen::mha::ProblemDescription& problem) const override;
+
+    MIOPEN_INTERNALS_EXPORT std::size_t
+    GetWorkspaceSize(const ExecutionContext& context,
+                     const miopen::mha::ProblemDescription& problem) const override;
+
+    MIOPEN_INTERNALS_EXPORT bool MayNeedWorkspace() const override;
+};
+
 } // namespace mha
 
 } // namespace solver

src/problem.cpp

@@ -610,10 +610,12 @@ Problem::FindSolutionsImpl(Handle& handle,
 
     const auto algo = AlgorithmName{"Mha"};
 
+    static solver::mha::MhaCKFlashAttentionV2Forward mhaCKFAForwardSolver;
     static solver::mha::MhaForward mhaForwardSolver;
     static solver::mha::MhaBackward mhaBackwardSolver;
 
-    std::vector<solver::mha::MhaSolver*> solvers = {&mhaForwardSolver, &mhaBackwardSolver};
+    std::vector<solver::mha::MhaSolver*> solvers = {
+        &mhaCKFAForwardSolver, &mhaForwardSolver, &mhaBackwardSolver};
 
     for(auto solver : solvers)
     {

src/solution.cpp

@@ -394,47 +394,12 @@ void Solution::RunImpl(Handle& handle,
         }
     }();
 
-    if(invoker)
-    {
-        (*invoker)(handle, invoke_ctx);
-        return;
-    }
-
-    solver::mha::MhaForward mhaForward;
-    solver::mha::MhaBackward mhaBackward;
-
-    if(!kernels.empty())
-    {
-        const auto ctx          = ExecutionContext{&handle};
-        const auto mha_solution = GetSolver() == mhaForward.SolverDbId()
-                                      ? mhaForward.GetSolution(ctx, problem_description)
-                                      : mhaBackward.GetSolution(ctx, problem_description);
-        auto kernel_handles     = std::vector<Kernel>{std::begin(kernels), std::end(kernels)};
-
-        invoker = (*mha_solution.invoker_factory)(kernel_handles);
-        (*invoker)(handle, invoke_ctx);
-        return;
-    }
-
-    const auto net_cfg = problem_description.MakeNetworkConfig();
-    invoker            = handle.GetInvoker(net_cfg, GetSolver());
-
-    if(invoker)
-    {
-        (*invoker)(handle, invoke_ctx);
-        return;
-    }
-
-    auto ctx = ExecutionContext{&handle};
-
-    const auto mha_solution = GetSolver() == mhaForward.SolverDbId()
-                                  ? mhaForward.GetSolution(ctx, problem_description)
-                                  : mhaBackward.GetSolution(ctx, problem_description);
+    const auto algo    = AlgorithmName{"MHA"};
+    const auto solvers = solver::SolverContainer<solver::mha::MhaCKFlashAttentionV2Forward,
+                                                 solver::mha::MhaForward,
+                                                 solver::mha::MhaBackward>{};
 
-    invoker =
-        handle.PrepareInvoker(*mha_solution.invoker_factory, mha_solution.construction_params);
-    handle.RegisterInvoker(*invoker, net_cfg, GetSolver().ToString());
-    (*invoker)(handle, invoke_ctx);
+    solvers.ExecutePrimitive(handle, problem_description, algo, invoke_ctx);
 }
 
 void Solution::RunImpl(Handle& handle,

src/solver.cpp

@@ -679,6 +679,8 @@ inline SolverRegistrar::SolverRegistrar(IdRegistryData& registry)
     Register(registry, ++id, Primitive::ReLU, prelu::MultiWeightsBackward{}.SolverDbId());
     Register(registry, ++id, Primitive::ReLU, prelu::SingleWeightBackward{}.SolverDbId());
 
+    Register(registry, ++id, Primitive::Mha, mha::MhaCKFlashAttentionV2Forward{}.SolverDbId());
+
     // IMPORTANT: New solvers should be added to the end of the function!
 }