Revert CK FP16 MHA integration

CMakeLists.txt

@@ -328,7 +328,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 device_mha_operations)
+        find_package(composable_kernel 1.0.0 COMPONENTS device_other_operations device_gemm_operations device_conv_operations device_reduction_operations)
     endif()
     if( MIOPEN_BACKEND STREQUAL "HIPNOGPU")
         set(MIOPEN_MODE_NOGPU 1)

src/CMakeLists.txt

@@ -309,7 +309,6 @@ 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/multimarginloss/forward_multimarginloss.cpp
@@ -859,7 +858,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 composable_kernel::device_mha_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 hip::host)
 endif()
 
 if(WIN32)

src/include/miopen/mha/solvers.hpp

@@ -77,28 +77,6 @@ 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

@@ -613,12 +613,10 @@ Problem::FindSolutionsImpl(const 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 = {
-        &mhaCKFAForwardSolver, &mhaForwardSolver, &mhaBackwardSolver};
+    std::vector<solver::mha::MhaSolver*> solvers = {&mhaForwardSolver, &mhaBackwardSolver};
 
     for(auto solver : solvers)
     {

src/solution.cpp

@@ -400,32 +400,15 @@ void Solution::RunImpl(const Handle& handle,
         return;
     }
 
-    auto getSolution = [&](const ExecutionContext& ctx) {
-        auto solverId = GetSolver();
-        solver::mha::MhaForward mhaForward;
-        solver::mha::MhaBackward mhaBackward;
-        solver::mha::MhaCKFlashAttentionV2Forward ckMhaForward;
-
-        if(solverId == ckMhaForward.SolverDbId())
-        {
-            return ckMhaForward.GetSolution(ctx, problem_description);
-        }
-        else if(solverId == mhaForward.SolverDbId())
-        {
-            return mhaForward.GetSolution(ctx, problem_description);
-        }
-        else if(solverId == mhaBackward.SolverDbId())
-        {
-            return mhaBackward.GetSolution(ctx, problem_description);
-        }
-
-        MIOPEN_THROW("No MHA solver with matching SolverDbId of " + solverId.ToString());
-    };
+    solver::mha::MhaForward mhaForward;
+    solver::mha::MhaBackward mhaBackward;
 
     if(!kernels.empty())
     {
         const auto ctx          = ExecutionContext{&handle};
-        const auto mha_solution = getSolution(ctx);
+        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);
@@ -442,8 +425,11 @@ void Solution::RunImpl(const Handle& handle,
         return;
     }
 
-    auto ctx                = ExecutionContext{&handle};
-    const auto mha_solution = getSolution(ctx);
+    auto ctx = ExecutionContext{&handle};
+
+    const auto mha_solution = GetSolver() == mhaForward.SolverDbId()
+                                  ? mhaForward.GetSolution(ctx, problem_description)
+                                  : mhaBackward.GetSolution(ctx, problem_description);
 
     invoker =
         handle.PrepareInvoker(*mha_solution.invoker_factory, mha_solution.construction_params);

src/solver.cpp

@@ -705,7 +705,8 @@ inline SolverRegistrar::SolverRegistrar(IdRegistryData& registry)
              Primitive::MultiMarginLoss,
              multimarginloss::MultiMarginLossForward{}.SolverDbId());
 
-    Register(registry, ++id, Primitive::Mha, mha::MhaCKFlashAttentionV2Forward{}.SolverDbId());
+    // removed CK MHA solver
+    ++id;
 
     RegisterWithSolver(
         registry, ++id, conv::ConvWinoRageRxS<2, 3>{}, miopenConvolutionAlgoWinograd);