[Conv] Backport CK split_k reduction workspace for grouped bwd/wrw solvers

Backport of the workspace portion of PR #1426 (commit e290adb,
ALMIOPEN-706) to ROCm 7.0.2. MIOpen previously sized the CK workspace for
the grouped bwd, grouped wrw, and 3D grouped wrw solvers via
GetCKAlphaBetaWorkspace (the alpha/beta backward-weights workspace query).
That query does not account for CK's split_k reduction workspace, so once
a tuned split_k>1 configuration was selected the reduction ran against an
under-allocated workspace and could return silently wrong results.

This change:
- Adds NextCKSplitkValue with CK's split_k autodeduce value (-1).
- Adds GetCKSplitkWorkspaceSize / GetCKSplitkMaxWorkspaceSize and queries
  CK for the actual per-instance workspace size across the split_k sweep.
- Uses that CK workspace size in GetWorkspaceSize (falling back to
  GetCKAlphaBetaWorkspace only when CK reports 0) and binds a
  correctly-sized workspace pointer in the NCHW/NHWC invoker factories.
- Removes the generic-search assertion that workspace size must be
  independent of the PerformanceConfig.

Verified via an A/B on the grouped bwd numerical tests: workspace-requiring
split_k>1 CK configurations are now selected and produce correct results
on int32-safe shapes. Shapes whose total element count exceeds INT_MAX
additionally require the large-tensor CK instances; after this change they
fault rather than return silently wrong results, and are not fully
resolved here (see ALMIOPEN-2305).

Also folds in the follow-up fix from mainline commit 54f9acb
(SreecharanGundaboluAMD, #3414): GetCKSplitkMaxWorkspaceSize must
initialize the sweep at split_k=1, not CkSplitkAutoDeduce. Starting at
the autodeduce sentinel makes NextCKSplitkValue return true on the first
step, so the loop evaluates only the autodeduce workspace and never
sweeps split_k=1,2,...,128 -- under-reporting the max workspace and
reintroducing the under-allocation this change is meant to fix. Also
widens the workspace accumulator to std::size_t to avoid int truncation
for >2GB workspaces.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: JonathanLichtnerAMD

projects/miopen/CHANGELOG.md

@@ -21,9 +21,11 @@ Full documentation for MIOpen is available [here](https://rocm.docs.amd.com/proj
 * [Conv] Eliminated redundant clearing of output buffers
 * [RNN] Updated selection heuristics
 * Updated tuning for MI300
+* [Conv] Improved Composable Kernel (CK) kernel selection during tuning for grouped bwd/wrw solvers
 
 ### Resolved issues
 * Fixed a segmentation fault when user specifies workspace smaller than what is required
+* [Conv] Fixed wrong results in grouped CK xdlops bwd/wrw solvers when a tuned split_k configuration required a reduction workspace that was not allocated
 * Fixed a layout calculation logic error that returned incorrect results and enabled less restrictive layout selection
 * Fixed memory access faults in misa kernels due to out-of-bounds memory usage
 * Fixed performance drop on gfx950 due to transpose kernel use

projects/miopen/src/include/miopen/conv/solvers.hpp

@@ -4643,6 +4643,10 @@ struct ConvHipImplicitGemm3DGroupWrwXdlops final
 private:
     template <typename DataType>
     bool CheckCKApplicability(const miopen::conv::ProblemDescription&) const;
+
+    template <typename DataType>
+    std::size_t GetCKMaxWorkspaceSize(const miopen::conv::ProblemDescription&) const;
+    size_t GetCKMaxWorkspaceSize(const miopen::conv::ProblemDescription& problem) const;
 };
 
 struct PerformanceConfigHipImplicitGemm3DGroupBwdXdlops
@@ -4815,6 +4819,8 @@ struct ConvHipImplicitGemmGroupBwdXdlops final
 private:
     template <typename DataType>
     bool CheckCKApplicability(const miopen::conv::ProblemDescription&) const;
+
+    size_t GetCKMaxWorkspaceSize(const miopen::conv::ProblemDescription& problem) const;
 };
 
 struct PerformanceConfigHipImplicitGemmGroupWrwXdlops
@@ -4909,6 +4915,8 @@ struct ConvHipImplicitGemmGroupWrwXdlops final
 private:
     template <typename DataType>
     bool CheckCKApplicability(const miopen::conv::ProblemDescription&) const;
+
+    size_t GetCKMaxWorkspaceSize(const miopen::conv::ProblemDescription& problem) const;
 };
 
 } // namespace conv

projects/miopen/src/include/miopen/generic_search.hpp

@@ -498,15 +498,6 @@ auto GenericSearch(const Solver s,
 
             try
             {
-                if(default_solution.workspace_sz != current_solution.workspace_sz)
-                {
-                    ret = -2;
-                    MIOPEN_LOG_E('#' << n_current << " (" << n_runs_total << ") "
-                                     << "Workspace size should not depend on PerformanceConfig: "
-                                     << default_solution.workspace_sz
-                                     << " != " << current_solution.workspace_sz);
-                }
-
                 invoker = profile_h.PrepareInvoker(*current_solution.invoker_factory,
                                                    current_solution.construction_params);
 

projects/miopen/src/include/miopen/solver/implicitgemm_ck_util.hpp

@@ -33,6 +33,7 @@
 #include <miopen/tensor_ops.hpp>
 #include <miopen/miopen_internal.h>
 #include <miopen/fusion/fusion_invoke_params.hpp>
+#include <miopen/solver/implicitgemm_util.hpp>
 
 #if MIOPEN_BACKEND_HIP && MIOPEN_USE_COMPOSABLEKERNEL
 #include <ck/utility/data_type.hpp>
@@ -50,7 +51,30 @@ struct ProblemDescription;
 } // namespace conv
 
 namespace solver {
+
+static constexpr int CkSplitkAutoDeduce = -1;
+
+template <int L, int H>
+inline static bool NextCKSplitkValue(int& v)
+{
+    assert((IsTwoPower<L, H>(v) || v == CkSplitkAutoDeduce));
+    if(v == H)
+    {
+        v = CkSplitkAutoDeduce;
+        return false;
+    }
+    if(v == CkSplitkAutoDeduce)
+    {
+        v = L;
+        return true;
+    }
+
+    v *= 2;
+    return false;
+}
+
 #if MIOPEN_BACKEND_HIP && MIOPEN_USE_COMPOSABLEKERNEL
+
 namespace conv {
 template <typename DataType>
 using DeviceOpGWrw = ck::tensor_operation::device::DeviceGroupedConvBwdWeight<
@@ -337,6 +361,35 @@ bool IsCKApplicable(const ProblemDescriptionType& problem)
     });
 }
 
+template <typename DeviceOpType,
+          typename CKArgsType,
+          typename ProblemDescriptionType = miopen::conv::ProblemDescription>
+size_t GetCKSplitkMaxWorkspaceSize(const ProblemDescriptionType& problem)
+{
+    const auto args                = CKArgsType{problem};
+    std::size_t max_workspace_size = 0;
+
+    const auto ptrs = DeviceOpType::GetInstances();
+    for(auto& ptr : ptrs)
+    {
+        // Cycle `split_k` over {1,2,4,...,128} then `CkSplitkAutoDeduce`.
+        // The loop then restarts from 1 for the next conv instance.
+        auto split_k = 1;
+        do
+        {
+            if(args.IsSupportedBySplitK(ptr, split_k))
+            {
+                auto workspace_size = args.GetCKSplitkWorkspaceSize(ptr, split_k);
+                if(workspace_size > max_workspace_size)
+                    max_workspace_size = workspace_size;
+            }
+        } while(!NextCKSplitkValue<1, 128>(split_k));
+    }
+
+    MIOPEN_LOG_I("Max workspace size reported by CK: " << max_workspace_size);
+    return max_workspace_size;
+}
+
 #define WORKAROUND_CK_ISSUE_1184 1
 #if WORKAROUND_CK_ISSUE_1184
 using WorkAroundHipEventProfiler = HipEventProfiler;
@@ -781,13 +834,14 @@ inline bool CKWrwRequireWorkspace(
 }
 
 /// \todo move to a cpp file
-inline size_t GetWorkspaceSizeLayoutTransformConv(const miopen::conv::ProblemDescription& problem)
+inline size_t GetWorkspaceSizeLayoutTransformConv(const miopen::conv::ProblemDescription& problem,
+                                                  size_t ck_ws_size = 0)
 {
     if(problem.IsLayoutNHWC())
     {
         if(problem.GetDirection() == ::miopen::conv::Direction::BackwardWeights)
         {
-            return GetCKAlphaBetaWorkspace(problem);
+            return (ck_ws_size > 0) ? ck_ws_size : GetCKAlphaBetaWorkspace(problem);
         }
         return 0;
     }
@@ -796,10 +850,11 @@ inline size_t GetWorkspaceSizeLayoutTransformConv(const miopen::conv::ProblemDes
 
     if(problem.GetDirection() == ::miopen::conv::Direction::BackwardWeights)
     {
-        MultiBufferWorkspaceTraits wt({GetPackedSize(problem.GetIn()),
-                                       GetPackedSize(problem.GetWeights()),
-                                       GetPackedSize(problem.GetOut()),
-                                       GetCKAlphaBetaWorkspace(problem)});
+        MultiBufferWorkspaceTraits wt(
+            {GetPackedSize(problem.GetIn()),
+             GetPackedSize(problem.GetWeights()),
+             GetPackedSize(problem.GetOut()),
+             (ck_ws_size > 0) ? ck_ws_size : GetCKAlphaBetaWorkspace(problem)});
         return wt.GetSize();
     }
 
@@ -1116,18 +1171,24 @@ ConvSolution InitInvokerFactoryNCHW(const ExecutionContext& ctx,
 
     std::optional<CKBWDWeightBufferDescriptor> _ck_buff_des;
 
-    if(problem.IsDirectionBackwardWrW())
-    {
-        _ck_buff_des.emplace(GetCKAlphaBetaWorkspace(problem), 0);
-    }
-
     auto ptr_iter = FindConvPtrByID(conv_ptrs, id_string);
     if(ptr_iter == conv_ptrs.end())
     {
         MIOPEN_LOG_E("PerformanceConfig kernel '" + kernel_id + "' does not exist.");
         return {miopenStatusInvalidValue};
     }
 
+    if constexpr(std::is_same_v<CastType, miopen::conv::WrWInvokeParams>)
+    {
+        auto ck_ws_size = ck_args.GetCKSplitkWorkspaceSize(*ptr_iter, split_k.value_or(1));
+        _ck_buff_des.emplace(ck_ws_size, 0);
+        result.workspace_sz = GetWorkspaceSizeLayoutTransformConv(problem, ck_ws_size);
+    }
+    else
+    {
+        result.workspace_sz = GetWorkspaceSizeLayoutTransformConv(problem);
+    }
+
     auto [_input1_tr_inst, _input2_tr_inst, _output_tr_inst, _output_init_tr_inst] =
         internal::MakeTaggedTransposeInstances<CKArgsType>(
             result, ctx, problem, ck_args, input1_op, input2_op, output_op, _ck_buff_des);
@@ -1234,8 +1295,6 @@ ConvSolution InitInvokerFactoryNCHW(const ExecutionContext& ctx,
             output_tr_inst.ConvertTo(handle, kernels, conv_tensors);
         };
     };
-
-    result.workspace_sz = GetWorkspaceSizeLayoutTransformConv(problem);
 #endif
     return result;
 }
@@ -1272,8 +1331,9 @@ ConvSolution InitInvokerFactoryNHWC(const ExecutionContext&,
         ConvSolution result;
 #if MIOPEN_BACKEND_HIP && MIOPEN_USE_COMPOSABLEKERNEL
         miopenAlphaBetaCase_t alpha_beta_case = problem.GetAlphaBetaCase();
-        [[maybe_unused]] bool should_allocated_wrw_buffer =
-            ShouldAllocateWorkSpaceBufferForWRW(problem);
+        auto ck_args = CKArgsType{problem};
+        auto ck_ws_size = ck_args.GetCKSplitkWorkspaceSize(*ptr_iter, split_k.value_or(1));
+        [[maybe_unused]] bool should_allocated_wrw_buffer = ck_ws_size > 0;
 
         result.invoker_factory = [kernel_id                   = kernel_id,
                                   split_k                     = split_k,
@@ -1334,7 +1394,7 @@ ConvSolution InitInvokerFactoryNHWC(const ExecutionContext&,
                 }
             };
         };
-        result.workspace_sz = GetWorkspaceSizeLayoutTransformConv(problem);
+        result.workspace_sz = GetWorkspaceSizeLayoutTransformConv(problem, ck_ws_size);
 #endif
         return result;
     }

projects/miopen/src/solver/conv/conv_hip_implicit_gemm_3d_grouped_wrw_xdlops.cpp

@@ -308,29 +308,32 @@ struct CKArgs
     template <typename ConvPtr>
     bool IsSupportedBy(const ConvPtr& conv_ptr) const
     {
-        auto arg_ptr = MakeArgPtr(conv_ptr, nullptr, nullptr, nullptr, 1.0f, 0.0f, 1);
-        // Creat dummy workspace to pass the ck IsSupportedArgument check.
-
-        int dummy_var = 1;
-        conv_ptr->SetWorkSpacePointer(arg_ptr.get(), &dummy_var);
-
+        auto arg_ptr        = MakeArgPtr(conv_ptr, nullptr, nullptr, nullptr, 1.0f, 0.0f, 1);
+        auto workspace_size = conv_ptr->GetWorkSpaceSize(arg_ptr.get());
+        if(workspace_size != 0)
+            conv_ptr->SetWorkSpacePointer(arg_ptr.get(), &workspace_size);
         return conv_ptr->IsSupportedArgument(arg_ptr.get());
     }
 
     template <typename ConvPtr>
     bool IsSupportedBySplitK(const ConvPtr& conv_ptr, int split_k) const
     {
-        auto arg_ptr = MakeArgPtr(conv_ptr, nullptr, nullptr, nullptr, 1.0f, 0.0f, split_k);
-
-        if(CKWrwRequireWorkspace(G, C1, K1, data_type, alpha_beta_case))
+        auto arg_ptr        = MakeArgPtr(conv_ptr, nullptr, nullptr, nullptr, 1.0f, 0.0f, split_k);
+        auto workspace_size = conv_ptr->GetWorkSpaceSize(arg_ptr.get());
+        if(workspace_size != 0)
         {
-            // Creat dummy workspace to pass the ck IsSupportedArgument check.
-            int dummy_var = 1;
-            conv_ptr->SetWorkSpacePointer(arg_ptr.get(), &dummy_var);
+            conv_ptr->SetWorkSpacePointer(arg_ptr.get(), &workspace_size);
         }
         return conv_ptr->IsSupportedArgument(arg_ptr.get());
     }
 
+    template <typename ConvPtr>
+    std::size_t GetCKSplitkWorkspaceSize(const ConvPtr& conv_ptr, int split_k) const
+    {
+        auto arg_ptr = MakeArgPtr(conv_ptr, nullptr, nullptr, nullptr, 1.0f, 0.0f, split_k);
+        return conv_ptr->GetWorkSpaceSize(arg_ptr.get());
+    }
+
     // Dim members are int64 (and length/stride arrays use ck::long_index_t)
     // so the NCHW stride builder above (e.g. Di*Hi*Wi*G*C) does not silently
     // overflow on tensors whose contiguous stride exceeds INT_MAX. Argument
@@ -469,7 +472,7 @@ bool PerformanceConfigHipImplicitGemm3DGroupWrwXdlops::SetNextValue(
     }
     do
     {
-        bool flag = NextTwoPower<1, 128>(split_k);
+        bool flag = NextCKSplitkValue<1, 128>(split_k);
         if(!flag)
         {
             kernel_id = valid_kernels[index] + "+" + std::to_string(split_k);
@@ -536,11 +539,54 @@ bool ConvHipImplicitGemm3DGroupWrwXdlops::IsValidPerformanceConfig(
     return config.IsValid(problem);
 }
 
+template <typename DataType>
+size_t
+ConvHipImplicitGemm3DGroupWrwXdlops::GetCKMaxWorkspaceSize(const ProblemDescription& problem) const
+{
+#if MIOPEN_BACKEND_HIP && MIOPEN_USE_COMPOSABLEKERNEL
+    switch(problem.GetAlphaBetaCase())
+    {
+    case BILINEAR:
+        return GetCKSplitkMaxWorkspaceSize<DeviceOpGBwdWeightBilinearPtrs<DataType>,
+                                           CKArgs<DataType>>(problem);
+    case SCALE:
+        return GetCKSplitkMaxWorkspaceSize<DeviceOpGBwdWeightScalePtrs<DataType>, CKArgs<DataType>>(
+            problem);
+    default:
+        return GetCKSplitkMaxWorkspaceSize<DeviceOpGBwdWeightDefaultPtrs<DataType>,
+                                           CKArgs<DataType>>(problem);
+    }
+#else
+    return 0;
+#endif
+}
+
+size_t
+ConvHipImplicitGemm3DGroupWrwXdlops::GetCKMaxWorkspaceSize(const ProblemDescription& problem) const
+{
+#if MIOPEN_BACKEND_HIP && MIOPEN_USE_COMPOSABLEKERNEL
+    switch(problem.GetInDataType())
+    {
+    case miopenHalf: return GetCKMaxWorkspaceSize<ck::half_t>(problem);
+    case miopenFloat: return GetCKMaxWorkspaceSize<float>(problem);
+    case miopenInt8: return GetCKMaxWorkspaceSize<int8_t>(problem);
+    case miopenBFloat16: return GetCKMaxWorkspaceSize<ck::bhalf_t>(problem);
+    case miopenInt64:
+    case miopenInt32:
+    case miopenFloat8_fnuz:
+    case miopenBFloat8_fnuz:
+    case miopenDouble: break;
+    }
+#endif
+    return 0; // other types not applicable for this solver
+}
+
 size_t
 ConvHipImplicitGemm3DGroupWrwXdlops::GetWorkspaceSize(const ExecutionContext&,
                                                       const ProblemDescription& problem) const
 {
-    return GetWorkspaceSizeLayoutTransformConv(problem);
+    auto ck_ws_size = GetCKMaxWorkspaceSize(problem);
+    return GetWorkspaceSizeLayoutTransformConv(problem, ck_ws_size);
 }
 
 PerformanceConfigHipImplicitGemm3DGroupWrwXdlops