add workspace size

example/32_batched_gemm_scale_softmax_gemm/grouped_multihead_attention_backward_fp16.cpp

@@ -316,7 +316,7 @@ int run(int argc, char* argv[])
     std::vector<DeviceMemPtr> ygrad_tensors_device;
     std::vector<DeviceMemPtr> kgrad_tensors_device;
     std::vector<DeviceMemPtr> vgrad_tensors_device;
-    std::size_t group_count = 1;
+    std::size_t group_count = 3;
     std::size_t flop = 0, num_byte = 0;
     for(std::size_t i=0; i<group_count; i++){
         // int M  = 128 * (rand() % 8 + 1);
@@ -538,6 +538,17 @@ int run(int argc, char* argv[])
         Scale{alpha},
         QKVElementOp{},
         YElementOp{});
+        
+    DeviceMem problem_desc_workspace(gemm.GetWorkSpaceSize(&argument));
+
+    gemm.SetWorkSpacePointer(&argument, problem_desc_workspace.GetDeviceBuffer());
+
+    if(!gemm.IsSupportedArgument(argument))
+    {
+        std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl;
+
+        return 0;
+    }
 
     if(!gemm.IsSupportedArgument(argument))
     {

include/ck/tensor_operation/gpu/device/impl/device_grouped_multihead_attention_backward_xdl_cshuffle.hpp

@@ -1062,6 +1062,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle
         return IsSupportedArgument(*dynamic_cast<const Argument*>(p_arg));
     }
 
+    size_t GetWorkSpaceSize(const BaseArgument* p_arg) const override
+    {
+        return dynamic_cast<const Argument*>(p_arg)->group_count_ * sizeof(GroupKernelArg);
+    }
+
     static auto MakeArgument(const std::vector<const DataType*>& p_As,
                              const std::vector<const DataType*>& p_Bs,
                              const std::vector<const DataType*>& p_B1s,