add some code for bfp16

example/20_convnd_bwd_weight_xdl/convnd_bwd_weight_xdl.cpp

@@ -18,9 +18,9 @@
 #include "device_convnd_backward_weight_xdl_c_shuffle_nhwc_kyxc_nhwk.hpp"
 #include "reference_conv_backward_weight.hpp"
 
-using InDataType  = ck::half_t;
-using WeiDataType = ck::half_t;
-using OutDataType = ck::half_t;
+using InDataType  = ck::bhalf_t;
+using WeiDataType = ck::bhalf_t;
+using OutDataType = ck::bhalf_t;
 using AccDataType = float;
 
 template <ck::index_t... Is>

include/ck/tensor_operation/gpu/device/device_convnd_backward_weight_xdl_c_shuffle_nhwc_kyxc_nhwk.hpp

@@ -962,6 +962,8 @@ struct DeviceConvndBwdWeightXdl_C_Shuffle_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_
             if constexpr(std::is_same<InDataType, ck::bhalf_t>::value)
             {
                 if(has_main_k0_block_loop)
+                {
+                    if(kbatch == 1)
                     {
                         const auto kernel = kernel_gemm_xdlops_bwd_weight<
                             GridwiseGemm,
@@ -969,7 +971,26 @@ struct DeviceConvndBwdWeightXdl_C_Shuffle_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_
                             CDataType,
                             remove_reference_t<DeviceOp::AGridDesc_K0_M_K1>,
                             remove_reference_t<DeviceOp::BGridDesc_K0_N_K1>,
-                        remove_reference_t<DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
+                            remove_reference_t<
+                                DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
+                            OutElementwiseOperation,
+                            InElementwiseOperation,
+                            WeiElementwiseOperation,
+                            remove_reference_t<DeviceOp::Block2CTileMap>,
+                            true>;
+
+                        Run(kernel);
+                    }
+                    else
+                    {
+                        const auto kernel = kernel_gemm_xdlops_bwd_weight<
+                            GridwiseGemmAtomicAddFloatForBf16,
+                            ADataType, // TODO: distiguish A/B datatype
+                            CDataType,
+                            remove_reference_t<DeviceOp::AGridDesc_K0_M_K1>,
+                            remove_reference_t<DeviceOp::BGridDesc_K0_N_K1>,
+                            remove_reference_t<
+                                DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
                             OutElementwiseOperation,
                             InElementwiseOperation,
                             WeiElementwiseOperation,
@@ -978,7 +999,10 @@ struct DeviceConvndBwdWeightXdl_C_Shuffle_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_
 
                         Run(kernel);
                     }
+                }
                 else
+                {
+                    if(kbatch == 1)
                     {
                         const auto kernel = kernel_gemm_xdlops_bwd_weight<
                             GridwiseGemm,
@@ -986,7 +1010,26 @@ struct DeviceConvndBwdWeightXdl_C_Shuffle_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_
                             CDataType,
                             remove_reference_t<DeviceOp::AGridDesc_K0_M_K1>,
                             remove_reference_t<DeviceOp::BGridDesc_K0_N_K1>,
-                        remove_reference_t<DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
+                            remove_reference_t<
+                                DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
+                            OutElementwiseOperation,
+                            InElementwiseOperation,
+                            WeiElementwiseOperation,
+                            remove_reference_t<DeviceOp::Block2CTileMap>,
+                            false>;
+
+                        Run(kernel);
+                    }
+                    else
+                    {
+                        const auto kernel = kernel_gemm_xdlops_bwd_weight<
+                            GridwiseGemmAtomicAddFloatForBf16,
+                            ADataType, // TODO: distiguish A/B datatype
+                            CDataType,
+                            remove_reference_t<DeviceOp::AGridDesc_K0_M_K1>,
+                            remove_reference_t<DeviceOp::BGridDesc_K0_N_K1>,
+                            remove_reference_t<
+                                DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
                             OutElementwiseOperation,
                             InElementwiseOperation,
                             WeiElementwiseOperation,
@@ -996,6 +1039,7 @@ struct DeviceConvndBwdWeightXdl_C_Shuffle_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_
                         Run(kernel);
                     }
                 }
+            }
             else
             {
                 if(has_main_k0_block_loop)