Merge branch 'develop' of...

Merge branch 'develop' of https://github.com/ROCmSoftwarePlatform/composable_kernel into add_fp16_wmma_conv_instance

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

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <array>
+
+#include "device_grouped_gemm.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <index_t NumDTensor = 0>
+struct GroupedGemmKernelArgument
+{
+    const void* p_a_grid;
+    const void* p_b_grid;
+    std::array<const void*, NumDTensor> p_ds_grid;
+    void* p_e_grid;
+
+    index_t M;
+    index_t N;
+    index_t K;
+
+    index_t StrideA;
+    index_t StrideB;
+    std::array<index_t, NumDTensor> StrideDs;
+    index_t StrideE;
+};
+
+template <typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename EDataType,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CElementwiseOperation>
+struct DeviceGroupedGemmFixedNK : DeviceGroupedGemm<ALayout,
+                                                    BLayout,
+                                                    DsLayout,
+                                                    ELayout,
+                                                    ADataType,
+                                                    BDataType,
+                                                    DsDataType,
+                                                    EDataType,
+                                                    AElementwiseOperation,
+                                                    BElementwiseOperation,
+                                                    CElementwiseOperation>
+{
+    virtual void SetDeviceKernelArgs(BaseArgument* p_arg, const void* kernel_args) const = 0;
+    virtual size_t GetDeviceKernelArgSize(const BaseArgument* p_arg) const               = 0;
+    virtual void SetKBatch(BaseArgument* p_arg, index_t k_batch) const                   = 0;
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

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

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <array>
+
+#include "ck/tensor_operation/gpu/device/device_base.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+/**
+ * \brief Image to column.
+ *
+ * This Device operator converts image ([G, N, Di, Hi, Wi, C]) to the gemm
+ * problem([N * Do * Ho * Wo, Z *  Y * X * C]). G must be equal to 1.
+ *
+ * \tparam NDimSpatial Number of spatial dimensions.
+ * \tparam InputLayout Input Layout.
+ * \tparam InputDataType Input Data Type.
+ * \tparam OutputDataType Output Data Type.
+ */
+template <index_t NDimSpatial,
+          typename InputLayout,
+          typename InputDataType,
+          typename OutputDataType>
+struct DeviceImageToColumn : public BaseOperator
+{
+
+    /**
+     * \brief Make argument pointer for image to column.
+     *
+     * \param p_in A pointer to the device memory of the input image.
+     * \param p_out A pointer to the device memory of the output.
+     * \param N Convolution batch size.
+     * \param C Convolution number of channels.
+     * \param input_spatial_lengths Input spatial lengths.
+     * \param filter_spatial_lengths Filter spatial lengths.
+     * \param output_spatial_lengths Output spatial lengths.
+     * \param input_g_n_c_wis_strides Input strides in order [G, N, C, D, H, W].
+     * \param output_m_k_strides Output strides.
+     * \param conv_filter_strides Convolution filter strides.
+     * \param conv_filter_dilations Convolution filter dilations.
+     * \param input_left_pads Convolution left pads.
+     * \param input_right_pads Convolution right pads.
+     * \return Pointer to the argument.
+     */
+    virtual std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const void* p_in,
+                        void* p_out,
+                        const ck::index_t N,
+                        const ck::index_t C,
+                        const std::array<index_t, NDimSpatial>& input_spatial_lengths,
+                        const std::array<index_t, NDimSpatial>& filter_spatial_lengths,
+                        const std::array<index_t, NDimSpatial>& output_spatial_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& input_g_n_c_wis_strides,
+                        const std::array<index_t, 2>& output_m_k_strides,
+                        const std::array<index_t, NDimSpatial>& conv_filter_strides,
+                        const std::array<index_t, NDimSpatial>& conv_filter_dilations,
+                        const std::array<index_t, NDimSpatial>& input_left_pads,
+                        const std::array<index_t, NDimSpatial>& input_right_pads) = 0;
+
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

include/ck/tensor_operation/gpu/device/device_index_pool_bwd.hpp → include/ck/tensor_operation/gpu/device/device_max_pool_bwd.hpp

@@ -13,7 +13,7 @@ namespace device {
 
 // For pooling which used indexable operation, such as MaxPool, MinPool...etc
 template <typename DOutDataType, typename IndexDataType, typename DInDataType>
-struct DeviceIndexPoolBwd : public BaseOperator
+struct DeviceMaxPoolBwd : public BaseOperator
 {
     virtual std::unique_ptr<BaseArgument>
     MakeArgumentPointer(const void* p_dout,
@@ -22,7 +22,8 @@ struct DeviceIndexPoolBwd : public BaseOperator
                         index_t dout_length,
                         index_t din_length,
                         std::vector<ck::index_t> window_lengths,
-                        std::vector<ck::index_t> window_strides) = 0;
+                        std::vector<ck::index_t> window_strides,
+                        std::vector<ck::index_t> window_dilations) = 0;
 
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
 };

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

@@ -17,6 +17,8 @@ template <index_t InOutRank,
           typename InDataType,
           typename OutDataType,
           typename IndexDataType,
+          typename InLayout,
+          typename OutLayout,
           ReduceTensorOp ReduceOpId,
           bool OutputIndex>
 struct DevicePoolFwd : public BaseOperator
@@ -25,13 +27,14 @@ struct DevicePoolFwd : public BaseOperator
     MakeArgumentPointer(const void* p_in_dev,
                         void* p_out_dev,
                         void* p_out_indices_dev,
-                        std::vector<ck::index_t> input_lengths,
-                        std::vector<ck::index_t> window_lengths,
-                        std::vector<ck::index_t> output_lengths,
-                        std::vector<ck::index_t> input_stride,
-                        std::vector<ck::index_t> output_stride,
-                        std::vector<ck::index_t> indices_stride,
-                        std::vector<ck::index_t> window_strides,
+                        std::vector<ck::index_t> input_n_c_wis_lengths,
+                        std::vector<ck::index_t> window_xs_lengths,
+                        std::vector<ck::index_t> output_n_c_wos_lengths,
+                        std::vector<ck::index_t> input_n_c_wis_stride,
+                        std::vector<ck::index_t> output_n_c_wis_stride,
+                        std::vector<ck::index_t> indices_n_c_wis_stride,
+                        std::vector<ck::index_t> window_xs_strides,
+                        std::vector<ck::index_t> window_xs_dilations,
                         std::vector<ck::index_t> input_left_pads,
                         std::vector<ck::index_t> input_right_pads,
                         std::vector<ck::index_t> pooling_dims) = 0;

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

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 

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

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <sstream>
+
+#include "ck/tensor_description/tensor_descriptor.hpp"
+#include "ck/tensor_description/tensor_descriptor_helper.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
+#include "ck/tensor_operation/gpu/device/device_avgpool_bwd.hpp"
+#include "ck/tensor_operation/gpu/grid/gridwise_2d_reduction_threadwise.hpp"
+#include "ck/host_utility/device_prop.hpp"
+#include "ck/host_utility/kernel_launch.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+// In and Din = [N, C, Di, Hi, Wi]
+// Out and Dout = [N, C, Do, Ho, Wo]
+// Out = AvgPoolFwd(In)
+// Din = AvgPoolBwd(Dout)
+// Pooling dimension = D, H, W
+template <typename DOutDataType,
+          typename DInDataType,
+          typename ComputeDataType,
+          ck::index_t BlockSize,
+          ck::index_t MThreadClusterSize,
+          ck::index_t KThreadClusterSize,
+          ck::index_t MThreadSliceSize,
+          ck::index_t KThreadSliceSize,
+          ck::index_t InSrcOutDstVectorSize>
+struct DeviceAvgPool3dBwd_NDHWC_NDHWC : public DeviceAvgPoolBwd<3,
+                                                                DOutDataType,
+                                                                DInDataType,
+                                                                tensor_layout::convolution::NDHWC,
+                                                                tensor_layout::convolution::NDHWC>
+{
+    static constexpr ck::index_t NDimSpatial = 3;
+
+    static constexpr auto I0 = Number<0>{};
+    static constexpr auto I1 = Number<1>{};
+
+    static constexpr ck::index_t M_BlockTileSize = MThreadClusterSize * MThreadSliceSize;
+    static constexpr ck::index_t K_BlockTileSize = KThreadClusterSize * KThreadSliceSize;
+
+    static auto
+    Make3DGridDescriptor_Out_M_K_In_M(const std::vector<ck::index_t>& dout_n_c_wos_lengths,
+                                      const std::vector<ck::index_t>& din_n_c_wos_length,
+                                      const std::vector<ck::index_t>& dout_n_c_wos_strides,
+                                      const std::vector<ck::index_t>& din_n_c_wos_strides,
+                                      const std::vector<ck::index_t>& window_lengths,
+                                      const std::vector<ck::index_t>& window_strides,
+                                      const std::vector<ck::index_t>& window_dilations,
+                                      const std::vector<ck::index_t>& input_left_pads,
+                                      const std::vector<ck::index_t>& input_right_pads,
+                                      const std::vector<ck::index_t>& tildes)
+    {
+        index_t i_ztilde = tildes[0];
+        index_t i_ytilde = tildes[1];
+        index_t i_xtilde = tildes[2];
+
+        const index_t N = dout_n_c_wos_lengths[0];
+        const index_t C = dout_n_c_wos_lengths[1];
+
+        const index_t Di = din_n_c_wos_length[2];
+        const index_t Hi = din_n_c_wos_length[3];
+        const index_t Wi = din_n_c_wos_length[4];
+
+        const index_t Do = dout_n_c_wos_lengths[2];
+        const index_t Ho = dout_n_c_wos_lengths[3];
+        const index_t Wo = dout_n_c_wos_lengths[4];
+
+        const index_t Z = window_lengths[0];
+        const index_t Y = window_lengths[1];
+        const index_t X = window_lengths[2];
+
+        const index_t InLeftPadD = input_left_pads[0];
+        const index_t InLeftPadH = input_left_pads[1];
+        const index_t InLeftPadW = input_left_pads[2];
+
+        const index_t InRightPadD = input_right_pads[0];
+        const index_t InRightPadH = input_right_pads[1];
+        const index_t InRightPadW = input_right_pads[2];
+
+        const index_t ConvStrideD = window_strides[0];
+        const index_t ConvStrideH = window_strides[1];
+        const index_t ConvStrideW = window_strides[2];
+
+        const index_t ConvDilationD = window_dilations[0];
+        const index_t ConvDilationH = window_dilations[1];
+        const index_t ConvDilationW = window_dilations[2];
+
+        const auto out_n_do_ho_wo_c_grid_desc =
+            make_naive_tensor_descriptor(make_tuple(N, Do, Ho, Wo, C),
+                                         make_tuple(dout_n_c_wos_strides[0],
+                                                    dout_n_c_wos_strides[2],
+                                                    dout_n_c_wos_strides[3],
+                                                    dout_n_c_wos_strides[4],
+                                                    dout_n_c_wos_strides[1]));
+
+        const auto GcdStrideDilationD = math::gcd(ConvStrideD, ConvDilationD);
+        const auto GcdStrideDilationH = math::gcd(ConvStrideH, ConvDilationH);
+        const auto GcdStrideDilationW = math::gcd(ConvStrideW, ConvDilationW);
+
+        const auto ZTilde = ConvStrideD / GcdStrideDilationD;
+        const auto YTilde = ConvStrideH / GcdStrideDilationH;
+        const auto XTilde = ConvStrideW / GcdStrideDilationW;
+
+        const auto ZDot = math::integer_divide_ceil(Z, ZTilde);
+        const auto YDot = math::integer_divide_ceil(Y, YTilde);
+        const auto XDot = math::integer_divide_ceil(X, XTilde);
+
+        const auto DTilde = Do + math::integer_divide_ceil(ConvDilationD * (Z - I1), ConvStrideD);
+        const auto HTilde = Ho + math::integer_divide_ceil(ConvDilationH * (Y - I1), ConvStrideH);
+        const auto WTilde = Wo + math::integer_divide_ceil(ConvDilationW * (X - I1), ConvStrideW);
+
+        // only work on Tildes that contribute to non-padding area of input tensor
+        const auto IDTildeSliceBegin = math::integer_divide_floor(
+            math::max(I0, InLeftPadD - ConvDilationD * (ZTilde - I1)), ConvStrideD);
+        const auto IHTildeSliceBegin = math::integer_divide_floor(
+            math::max(I0, InLeftPadH - ConvDilationH * (YTilde - I1)), ConvStrideH);
+        const auto IWTildeSliceBegin = math::integer_divide_floor(
+            math::max(I0, InLeftPadW - ConvDilationW * (XTilde - I1)), ConvStrideW);
+
+        const auto IDTildeSliceEnd =
+            math::min(DTilde, math::integer_divide_ceil(InLeftPadD + Di - I1, ConvStrideD) + I1);
+        const auto IHTildeSliceEnd =
+            math::min(HTilde, math::integer_divide_ceil(InLeftPadH + Hi - I1, ConvStrideH) + I1);
+        const auto IWTildeSliceEnd =
+            math::min(WTilde, math::integer_divide_ceil(InLeftPadW + Wi - I1, ConvStrideW) + I1);
+
+        const auto DTildeSlice = IDTildeSliceEnd - IDTildeSliceBegin;
+        const auto HTildeSlice = IHTildeSliceEnd - IHTildeSliceBegin;
+        const auto WTildeSlice = IWTildeSliceEnd - IWTildeSliceBegin;
+
+        // ReduceK is different for each Reduce
+        const auto ZDotSlice = math::integer_divide_ceil(Z - i_ztilde, ZTilde);
+        const auto YDotSlice = math::integer_divide_ceil(Y - i_ytilde, YTilde);
+        const auto XDotSlice = math::integer_divide_ceil(X - i_xtilde, XTilde);
+
+        // Problem size of reduction kernel
+        const index_t MRaw = N * DTildeSlice * HTildeSlice * WTildeSlice * C;
+        const index_t MPad = math::integer_least_multiple(MRaw, M_BlockTileSize) - MRaw;
+
+        const index_t KRaw = ZDotSlice * YDotSlice * XDotSlice;
+        const index_t KPad = math::integer_least_multiple(KRaw, K_BlockTileSize) - KRaw;
+
+        // Out[ReduceM, ReduceK]
+        const auto out_n_dop_hop_wop_c_grid_desc = transform_tensor_descriptor(
+            out_n_do_ho_wo_c_grid_desc,
+            make_tuple(make_pass_through_transform(N),
+                       make_pad_transform(Do, I0, I0),
+                       make_pad_transform(Ho, I0, I0),
+                       make_pad_transform(Wo, I0, I0),
+                       make_pass_through_transform(C)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}));
+
+        const auto out_n_zdot_dtilde_ydot_htilde_xdot_wtilde_c_grid_desc =
+            transform_tensor_descriptor(
+                out_n_dop_hop_wop_c_grid_desc,
+                make_tuple(
+                    make_pass_through_transform(N),
+                    make_embed_transform(make_tuple(ZDot, DTilde),
+                                         make_tuple(-ConvDilationD / GcdStrideDilationD, I1)),
+                    make_embed_transform(make_tuple(YDot, HTilde),
+                                         make_tuple(-ConvDilationH / GcdStrideDilationH, I1)),
+                    make_embed_transform(make_tuple(XDot, WTilde),
+                                         make_tuple(-ConvDilationW / GcdStrideDilationW, I1)),
+                    make_pass_through_transform(C)),
+                make_tuple(
+                    Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
+                make_tuple(Sequence<0>{},
+                           Sequence<1, 2>{},
+                           Sequence<3, 4>{},
+                           Sequence<5, 6>{},
+                           Sequence<7>{}));
+
+        const auto
+            out_n_zdotslice_dtildeslice_ydotslice_htildeslice_xdotslice_wtildeslice_c_grid_desc =
+                transform_tensor_descriptor(
+                    out_n_zdot_dtilde_ydot_htilde_xdot_wtilde_c_grid_desc,
+                    make_tuple(make_pass_through_transform(N),
+                               make_slice_transform(ZDot, I0, ZDotSlice),
+                               make_slice_transform(DTilde, IDTildeSliceBegin, DTildeSlice),
+                               make_slice_transform(YDot, I0, YDotSlice),
+                               make_slice_transform(HTilde, IHTildeSliceBegin, HTildeSlice),
+                               make_slice_transform(XDot, I0, XDotSlice),
+                               make_slice_transform(WTilde, IWTildeSliceBegin, WTildeSlice),
+                               make_pass_through_transform(C)),
+                    make_tuple(Sequence<0>{},
+                               Sequence<1>{},
+                               Sequence<2>{},
+                               Sequence<3>{},
+                               Sequence<4>{},
+                               Sequence<5>{},
+                               Sequence<6>{},
+                               Sequence<7>{}),
+                    make_tuple(Sequence<0>{},
+                               Sequence<1>{},
+                               Sequence<2>{},
+                               Sequence<3>{},
+                               Sequence<4>{},
+                               Sequence<5>{},
+                               Sequence<6>{},
+                               Sequence<7>{}));
+
+        const auto out_grid_desc_reducemraw_reducekraw = transform_tensor_descriptor(
+            out_n_zdotslice_dtildeslice_ydotslice_htildeslice_xdotslice_wtildeslice_c_grid_desc,
+            make_tuple(
+                make_merge_transform(make_tuple(N, DTildeSlice, HTildeSlice, WTildeSlice, C)),
+                make_merge_transform(make_tuple(ZDotSlice, YDotSlice, XDotSlice))),
+            make_tuple(Sequence<0, 2, 4, 6, 7>{}, Sequence<1, 3, 5>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        const auto out_grid_desc_reducem_reducek = transform_tensor_descriptor(
+            out_grid_desc_reducemraw_reducekraw,
+            make_tuple(make_right_pad_transform(MRaw, MPad), make_right_pad_transform(KRaw, KPad)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        // In[ReduceM]
+        const auto in_n_di_hi_wi_c_grid_desc =
+            make_naive_tensor_descriptor(make_tuple(N, Di, Hi, Wi, C),
+                                         make_tuple(din_n_c_wos_strides[0],
+                                                    din_n_c_wos_strides[2],
+                                                    din_n_c_wos_strides[3],
+                                                    din_n_c_wos_strides[4],
+                                                    din_n_c_wos_strides[1]));
+
+        const auto in_n_dip_hip_wip_c_grid_desc = transform_tensor_descriptor(
+            in_n_di_hi_wi_c_grid_desc,
+            make_tuple(make_pass_through_transform(N),
+                       make_pad_transform(Di, InLeftPadD, InRightPadD),
+                       make_pad_transform(Hi, InLeftPadH, InRightPadH),
+                       make_pad_transform(Wi, InLeftPadW, InRightPadW),
+                       make_pass_through_transform(C)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}));
+
+        const auto in_n_ztilde_dtilde_ytilde_htilde_xtilde_wtilde_c_grid_desc =
+            transform_tensor_descriptor(
+                in_n_dip_hip_wip_c_grid_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_embed_transform(make_tuple(XTilde, DTilde),
+                                                make_tuple(ConvDilationD, ConvStrideD)),
+                           make_embed_transform(make_tuple(YTilde, HTilde),
+                                                make_tuple(ConvDilationH, ConvStrideH)),
+                           make_embed_transform(make_tuple(XTilde, WTilde),
+                                                make_tuple(ConvDilationW, ConvStrideW)),
+                           make_pass_through_transform(C)),
+                make_tuple(
+                    Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
+                make_tuple(Sequence<0>{},
+                           Sequence<1, 2>{},
+                           Sequence<3, 4>{},
+                           Sequence<5, 6>{},
+                           Sequence<7>{}));
+
+        const auto in_n_dtildeslice_htildeslice_wtildeslice_c_grid_desc =
+            transform_tensor_descriptor(
+                in_n_ztilde_dtilde_ytilde_htilde_xtilde_wtilde_c_grid_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_freeze_transform(i_ztilde),
+                           make_slice_transform(DTilde, IDTildeSliceBegin, DTildeSlice),
+                           make_freeze_transform(i_ytilde),
+                           make_slice_transform(HTilde, IHTildeSliceBegin, HTildeSlice),
+                           make_freeze_transform(i_xtilde),
+                           make_slice_transform(WTilde, IWTildeSliceBegin, WTildeSlice),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{},
+                           Sequence<1>{},
+                           Sequence<2>{},
+                           Sequence<3>{},
+                           Sequence<4>{},
+                           Sequence<5>{},
+                           Sequence<6>{},
+                           Sequence<7>{}),
+                make_tuple(Sequence<0>{},
+                           Sequence<>{},
+                           Sequence<1>{},
+                           Sequence<>{},
+                           Sequence<2>{},
+                           Sequence<>{},
+                           Sequence<3>{},
+                           Sequence<4>{}));
+
+        const auto in_grid_desc_reducemraw = transform_tensor_descriptor(
+            in_n_dtildeslice_htildeslice_wtildeslice_c_grid_desc,
+            make_tuple(
+                make_merge_transform(make_tuple(N, DTildeSlice, HTildeSlice, WTildeSlice, C))),
+            make_tuple(Sequence<0, 1, 2, 3, 4>{}),
+            make_tuple(Sequence<0>{}));
+
+        const auto in_grid_desc_reducem =
+            transform_tensor_descriptor(in_grid_desc_reducemraw,
+                                        make_tuple(make_right_pad_transform(MRaw, MPad)),
+                                        make_tuple(Sequence<0>{}),
+                                        make_tuple(Sequence<0>{}));
+
+        return make_tuple(out_grid_desc_reducem_reducek, in_grid_desc_reducem);
+    }
+
+    using DoutDinGridDesc = decltype(Make3DGridDescriptor_Out_M_K_In_M({0, 0, 0, 0, 0},
+                                                                       {0, 0, 0, 0, 0},
+                                                                       {0, 0, 0, 0, 0},
+                                                                       {0, 0, 0, 0, 0},
+                                                                       {0, 0, 0},
+                                                                       {0, 0, 0},
+                                                                       {0, 0, 0},
+                                                                       {0, 0, 0},
+                                                                       {0, 0, 0},
+                                                                       {0, 0, 0}));
+
+    using DoutGridDesc_M_K = remove_cvref_t<tuple_element_t<0, DoutDinGridDesc>>;
+    using DinGridDesc_M    = remove_cvref_t<tuple_element_t<1, DoutDinGridDesc>>;
+
+    // FIXME
+    // for NDHWC, the dim C is the fastest dimension, and is not reduced.
+    // Hence, it is in M dimension for reduction kernel.
+    static constexpr index_t OutSrcInDstVectorDim = 0; // 0: M, 1: K
+
+    using PassThrough = tensor_operation::element_wise::PassThrough;
+    using Div         = tensor_operation::element_wise::UnaryDivide;
+
+    using gridwise_reduce = GridwiseReduction_mk_to_m_threadwise<DOutDataType,
+                                                                 DInDataType,
+                                                                 ComputeDataType,
+                                                                 int,
+                                                                 DoutGridDesc_M_K,
+                                                                 DinGridDesc_M,
+                                                                 reduce::Add,
+                                                                 PassThrough,
+                                                                 Div,
+                                                                 InMemoryDataOperationEnum::Set,
+                                                                 false, // propagate_nan
+                                                                 BlockSize,
+                                                                 MThreadSliceSize,
+                                                                 KThreadSliceSize,
+                                                                 OutSrcInDstVectorDim,
+                                                                 InSrcOutDstVectorSize,
+                                                                 InSrcOutDstVectorSize>;
+
+    struct Argument : public BaseArgument
+    {
+        Argument(const DOutDataType* p_dout,
+                 DInDataType* p_din,
+                 std::vector<ck::index_t> dout_n_c_wos_lengths,
+                 std::vector<ck::index_t> din_n_c_wos_length,
+                 std::vector<ck::index_t> dout_n_c_wos_strides,
+                 std::vector<ck::index_t> din_n_c_wos_strides,
+                 std::vector<ck::index_t> window_lengths,
+                 std::vector<ck::index_t> window_strides,
+                 std::vector<ck::index_t> window_dilations,
+                 std::vector<ck::index_t> input_left_pads,
+                 std::vector<ck::index_t> input_right_pads)
+            : p_dout_grid_{p_dout},
+              p_din_grid_{p_din},
+              dout_n_c_wos_lengths_{dout_n_c_wos_lengths},
+              din_n_c_wos_length_{din_n_c_wos_length},
+              dout_n_c_wos_strides_{dout_n_c_wos_strides},
+              din_n_c_wos_strides_{din_n_c_wos_strides},
+              num_reduce_{1},
+              div_element_op_{window_lengths[0] * window_lengths[1] * window_lengths[2]}
+        {
+            std::vector<ck::index_t> Tildes(NDimSpatial);
+            for(int i = 0; i < NDimSpatial; ++i)
+            {
+                int GcdStrideDilation = math::gcd(window_strides[i], window_dilations[i]);
+                Tildes[i]             = window_strides[i] / GcdStrideDilation;
+                num_reduce_ *= Tildes[i];
+            }
+
+            for(index_t i_ztilde = 0; i_ztilde < Tildes[0]; ++i_ztilde)
+            {
+                for(index_t i_ytilde = 0; i_ytilde < Tildes[1]; ++i_ytilde)
+                {
+                    for(index_t i_xtilde = 0; i_xtilde < Tildes[2]; ++i_xtilde)
+                    {
+                        // check slice is valid
+                        const auto ZDotSlice =
+                            math::integer_divide_ceil(window_lengths[0] - i_ztilde, Tildes[0]);
+                        const auto YDotSlice =
+                            math::integer_divide_ceil(window_lengths[1] - i_ytilde, Tildes[1]);
+                        const auto XDotSlice =
+                            math::integer_divide_ceil(window_lengths[2] - i_xtilde, Tildes[2]);
+
+                        if(ZDotSlice * YDotSlice * XDotSlice <= 0)
+                        {
+                            continue;
+                        }
+
+                        const auto dout_din_grid_desc =
+                            Make3DGridDescriptor_Out_M_K_In_M(dout_n_c_wos_lengths,
+                                                              din_n_c_wos_length,
+                                                              dout_n_c_wos_strides,
+                                                              din_n_c_wos_strides,
+                                                              window_lengths,
+                                                              window_strides,
+                                                              window_dilations,
+                                                              input_left_pads,
+                                                              input_right_pads,
+                                                              {i_ztilde, i_ytilde, i_xtilde});
+
+                        dout_grid_desc_m_k_container_.push_back(dout_din_grid_desc[I0]);
+                        din_grid_desc_m_container_.push_back(dout_din_grid_desc[I1]);
+                    }
+                }
+            }
+        }
+
+        const DOutDataType* p_dout_grid_;
+        DInDataType* p_din_grid_;
+        std::vector<ck::index_t> dout_n_c_wos_lengths_;
+        std::vector<ck::index_t> din_n_c_wos_length_;
+        std::vector<ck::index_t> dout_n_c_wos_strides_;
+        std::vector<ck::index_t> din_n_c_wos_strides_;
+
+        int num_reduce_;
+        std::vector<DoutGridDesc_M_K> dout_grid_desc_m_k_container_;
+        std::vector<DinGridDesc_M> din_grid_desc_m_container_;
+
+        Div div_element_op_;
+    };
+
+    struct Invoker : public BaseInvoker
+    {
+        float Run(const Argument& arg, const StreamConfig& stream_config = StreamConfig{})
+        {
+            float ave_time = 0;
+
+            for(index_t i = 0; i < arg.num_reduce_; i++)
+            {
+                const auto kernel = kernel_reduce_threadwise<gridwise_reduce,
+                                                             false,
+                                                             false,
+                                                             false, // don't have index input
+                                                             DOutDataType,
+                                                             DInDataType,
+                                                             ComputeDataType,
+                                                             int,
+                                                             DoutGridDesc_M_K,
+                                                             DinGridDesc_M,
+                                                             PassThrough,
+                                                             Div>;
+
+                ck::index_t M           = arg.dout_grid_desc_m_k_container_[i].GetLength(I0);
+                const index_t grid_size = (M / M_BlockTileSize);
+
+                ave_time += launch_and_time_kernel(stream_config,
+                                                   kernel,
+                                                   dim3(grid_size),
+                                                   dim3(BlockSize),
+                                                   0,
+                                                   arg.dout_grid_desc_m_k_container_[i],
+                                                   arg.din_grid_desc_m_container_[i],
+                                                   PassThrough{},
+                                                   arg.div_element_op_,
+                                                   float(1),
+                                                   arg.p_dout_grid_,
+                                                   nullptr,
+                                                   float(0),
+                                                   arg.p_din_grid_,
+                                                   nullptr);
+            }
+
+            return ave_time;
+        }
+
+        float Run(const BaseArgument* p_arg,
+                  const StreamConfig& stream_config = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg), stream_config);
+        }
+    };
+
+    static bool IsSupportedArgument(const Argument& arg)
+    {
+        constexpr index_t Rank = NDimSpatial + 2;
+        int doutFastestDim     = -1;
+        int dinFastestDim      = -1;
+
+        for(int i = 0; i < Rank; ++i)
+        {
+            if(arg.dout_n_c_wos_strides_[i] == 1)
+                doutFastestDim = i;
+            if(arg.din_n_c_wos_strides_[i] == 1)
+                dinFastestDim = i;
+        }
+
+        if(doutFastestDim == -1 || dinFastestDim == -1)
+        {
+            if constexpr(InSrcOutDstVectorSize != 1)
+                return false;
+        }
+        else
+        {
+            if(arg.dout_n_c_wos_lengths_[doutFastestDim] % InSrcOutDstVectorSize != 0)
+                return false;
+            if(arg.din_n_c_wos_length_[dinFastestDim] % InSrcOutDstVectorSize != 0)
+                return false;
+        }
+
+        return true;
+    }
+
+    bool IsSupportedArgument(const BaseArgument* p_arg) override
+    {
+        return IsSupportedArgument(*dynamic_cast<const Argument*>(p_arg));
+    }
+
+    std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const void* p_dout,
+                        void* p_din,
+                        std::vector<ck::index_t> dout_n_c_wos_lengths,
+                        std::vector<ck::index_t> din_n_c_wos_length,
+                        std::vector<ck::index_t> dout_n_c_wos_strides,
+                        std::vector<ck::index_t> din_n_c_wos_strides,
+                        std::vector<ck::index_t> window_lengths,
+                        std::vector<ck::index_t> window_strides,
+                        std::vector<ck::index_t> window_dilations,
+                        std::vector<ck::index_t> input_left_pads,
+                        std::vector<ck::index_t> input_right_pads) override
+    {
+        constexpr index_t Rank = NDimSpatial + 2;
+
+        if(dout_n_c_wos_strides.size() != Rank || din_n_c_wos_strides.size() != Rank ||
+           dout_n_c_wos_lengths.size() != Rank || din_n_c_wos_length.size() != Rank)
+            throw std::runtime_error("dimension is incorrect");
+
+        if(window_lengths.size() != NDimSpatial || window_strides.size() != NDimSpatial ||
+           window_dilations.size() != NDimSpatial || input_left_pads.size() != NDimSpatial ||
+           input_right_pads.size() != NDimSpatial)
+            throw std::runtime_error("dimension is incorrect");
+
+        return std::make_unique<Argument>(static_cast<const DOutDataType*>(p_dout),
+                                          static_cast<DInDataType*>(p_din),
+                                          dout_n_c_wos_lengths,
+                                          din_n_c_wos_length,
+                                          dout_n_c_wos_strides,
+                                          din_n_c_wos_strides,
+                                          window_lengths,
+                                          window_strides,
+                                          window_dilations,
+                                          input_left_pads,
+                                          input_right_pads);
+    }
+
+    std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "DeviceAvgPool3dBwd<" << BlockSize << ",";
+        str << "M_C" << MThreadClusterSize << "_S" << MThreadSliceSize << ",";
+        str << "K_C" << KThreadClusterSize << "_S" << KThreadSliceSize << ",";
+        str <<"InSrcOutDstVectorSize_" << InSrcOutDstVectorSize << ">";
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

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

@@ -543,9 +543,13 @@ struct DeviceBatchedContractionMultipleD_Xdl_CShuffle
         EGridDesc_G_M_N e_grid_desc_g_m_n_;
     };
 
+    using ComputeDataType = ADataType;
+
     // GridwiseGemm
     using GridwiseGemm = GridwiseGemmMultipleD_xdl_cshuffle<
-        ADataType, // TODO: distinguish A/B datatype
+        ADataType,
+        BDataType,
+        ComputeDataType,
         AccDataType,
         CShuffleDataType,
         DsDataType,
@@ -588,14 +592,18 @@ struct DeviceBatchedContractionMultipleD_Xdl_CShuffle
         LoopSched>;
 
     // desc for blockwise copy
-    using AGridDesc_AK0_M_AK1                          = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultAGridDescriptor_AK0_M_AK1(AGridDesc_M_K{}))>;
-    using BGridDesc_BK0_N_BK1                          = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultBGridDescriptor_BK0_N_BK1(BGridDesc_N_K{}))>;
-    using DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(DsGridDesc_M_N{}))>;
-    using EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock  = remove_cvref_t<decltype(
-        GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(EGridDesc_M_N{}))>;
+    using AGridDesc_AK0_M_AK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultAGridDescriptor_AK0_M_AK1(
+            AGridDesc_M_K{}))>;
+    using BGridDesc_BK0_N_BK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultBGridDescriptor_BK0_N_BK1(
+            BGridDesc_N_K{}))>;
+    using DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock = remove_cvref_t<
+        decltype(GridwiseGemm::MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
+            DsGridDesc_M_N{}))>;
+    using EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock =
+        remove_cvref_t<decltype(GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
+            EGridDesc_M_N{}))>;
 
     // block-to-e-tile map
     using Block2ETileMap =
@@ -840,9 +848,7 @@ struct DeviceBatchedContractionMultipleD_Xdl_CShuffle
 
     static bool IsSupportedArgument(const Argument& arg)
     {
-        if(!(ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx90a" ||
-             ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx941" ||
-             ck::get_device_name() == "gfx942"))
+        if(!ck::is_xdl_supported())
         {
             return false;
         }

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

@@ -331,8 +331,13 @@ struct DeviceBatchedGemmEPermuteXdl : public DeviceBatchedGemmEPermute<ALayout,
         EGridDesc_G0_G1_M_N e_grid_desc_g0_g1_m_n_;
     };
 
+    using ComputeDataType = ADataType;
+
+    // GridwiseGemm
     using GridwiseGemm = GridwiseGemmMultipleD_xdl_cshuffle<
-        ADataType, // TODO: distinguish A/B datatype
+        ADataType,
+        BDataType,
+        ComputeDataType,
         AccDataType,
         CShuffleDataType,
         ck::Tuple<>, // DsDataType,
@@ -378,13 +383,16 @@ struct DeviceBatchedGemmEPermuteXdl : public DeviceBatchedGemmEPermute<ALayout,
         CDEBlockTransferScalarPerVector_NPerBlock,
         LoopSched>;
 
-    using AGridDesc_AK0_M_AK1 = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultAGridDescriptor_AK0_M_AK1(AGridDesc_M_K{}))>;
-    using BGridDesc_BK0_N_BK1 = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultBGridDescriptor_BK0_N_BK1(BGridDesc_N_K{}))>;
+    using AGridDesc_AK0_M_AK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultAGridDescriptor_AK0_M_AK1(
+            AGridDesc_M_K{}))>;
+    using BGridDesc_BK0_N_BK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultBGridDescriptor_BK0_N_BK1(
+            BGridDesc_N_K{}))>;
 
-    using EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock = decltype(
-        GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(EGridDesc_M_N{}));
+    using EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock =
+        decltype(GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
+            EGridDesc_M_N{}));
     using Block2ETileMap = typename GridwiseGemm::DefaultBlock2ETileMap;
 
     // Argument
@@ -571,6 +579,11 @@ struct DeviceBatchedGemmEPermuteXdl : public DeviceBatchedGemmEPermute<ALayout,
 
     static bool IsSupportedArgument(const Argument& arg)
     {
+        if(!ck::is_xdl_supported())
+        {
+            return false;
+        }
+
         return GridwiseGemm::CheckValidity(arg.a_grid_desc_m_k_,
                                            arg.b_grid_desc_n_k_,
                                            ck::Tuple<>{},

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

@@ -589,9 +589,7 @@ struct DeviceBatchedGemmGemm_Xdl_CShuffle : public DeviceBatchedGemmGemm<ALayout
 
     static bool IsSupportedArgument(const Argument& arg)
     {
-        if(!(ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx90a" ||
-             ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx941" ||
-             ck::get_device_name() == "gfx942"))
+        if(!ck::is_xdl_supported())
         {
             return false;
         }

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

@@ -324,8 +324,12 @@ struct DeviceBatchedGemmMultiD_Xdl : public DeviceBatchedGemmMultiD<ALayout,
         index_t BatchStrideE_;
     };
 
+    using ComputeDataType = ADataType;
+
     using GridwiseGemm = GridwiseGemmMultipleD_xdl_cshuffle<
         ADataType, // TODO: distinguish A/B datatype
+        BDataType,
+        ComputeDataType,
         AccDataType,
         CShuffleDataType,
         DsDataType,
@@ -368,14 +372,18 @@ struct DeviceBatchedGemmMultiD_Xdl : public DeviceBatchedGemmMultiD<ALayout,
         LoopSched>;
 
     // desc for blockwise copy
-    using AGridDesc_AK0_M_AK1                          = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultAGridDescriptor_AK0_M_AK1(AGridDesc_M_K{}))>;
-    using BGridDesc_BK0_N_BK1                          = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultBGridDescriptor_BK0_N_BK1(BGridDesc_N_K{}))>;
-    using DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(DsGridDesc_M_N{}))>;
-    using EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock  = remove_cvref_t<decltype(
-        GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(EGridDesc_M_N{}))>;
+    using AGridDesc_AK0_M_AK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultAGridDescriptor_AK0_M_AK1(
+            AGridDesc_M_K{}))>;
+    using BGridDesc_BK0_N_BK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultBGridDescriptor_BK0_N_BK1(
+            BGridDesc_N_K{}))>;
+    using DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock = remove_cvref_t<
+        decltype(GridwiseGemm::MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
+            DsGridDesc_M_N{}))>;
+    using EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock =
+        remove_cvref_t<decltype(GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
+            EGridDesc_M_N{}))>;
 
     // block-to-e-tile map
     using Block2ETileMap =
@@ -580,9 +588,7 @@ struct DeviceBatchedGemmMultiD_Xdl : public DeviceBatchedGemmMultiD<ALayout,
 
     static bool IsSupportedArgument(const Argument& arg)
     {
-        if(!(ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx90a" ||
-             ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx941" ||
-             ck::get_device_name() == "gfx942"))
+        if(!ck::is_xdl_supported())
         {
             return false;
         }

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

@@ -70,10 +70,9 @@ __global__ void
             const ComputePtrOffsetOfBatch compute_ptr_offset_of_batch,
             const Block2CTileMap block_2_ctile_map)
 {
-// TODO: Enable for gfx90a after complier fix
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx906__) || defined(__gfx908__) ||              \
-    defined(__gfx940__) || defined(__gfx1030__) || defined(__gfx1100__) || defined(__gfx1101__) || \
-    defined(__gfx1102__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx906__) || defined(__gfx908__) ||             \
+    defined(__gfx90a__) || defined(__gfx940__) || defined(__gfx1030__) || defined(__gfx1100__) || \
+    defined(__gfx1101__) || defined(__gfx1102__))
 
     const index_t num_blocks_per_batch =
         __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
@@ -650,10 +649,10 @@ struct DeviceBatchedGemmMultipleD_Dl : public DeviceBatchedGemmMultiD<ALayout,
     static bool IsSupportedArgument(const Argument& arg)
     {
         // TODO: Enable for gfx90a after complier fix
-        if(ck::get_device_name() == "gfx906" || ck::get_device_name() == "gfx908" ||
-           ck::get_device_name() == "gfx1030" || ck::get_device_name() == "gfx940" ||
-           ck::get_device_name() == "gfx1100" || ck::get_device_name() == "gfx1101" ||
-           ck::get_device_name() == "gfx1102")
+        if(ck::get_device_name() == "gfx906" || ck::get_device_name() == "gfx90a" ||
+           ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx1030" ||
+           ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx1100" ||
+           ck::get_device_name() == "gfx1101" || ck::get_device_name() == "gfx1102")
         {
             bool pass = true;
             pass      = pass && arg.K_ % K1 == 0;

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

@@ -196,6 +196,8 @@ template <typename A0Layout,
           index_t B0BlockTransferSrcScalarPerVector,
           index_t B0BlockTransferDstScalarPerVector_BK1,
           bool B0BlockLdsExtraN,
+          index_t CDE0BlockTransferSrcVectorDim,
+          index_t CDE0BlockTransferSrcScalaerPerVector,
           typename B1BlockTransferThreadClusterLengths_BK0_N_BK1,
           typename B1BlockTransferThreadClusterArrangeOrder,
           typename B1BlockTransferSrcAccessOrder,
@@ -492,6 +494,8 @@ struct DeviceBatchedGemmMultipleDGemmMultipleD_Xdl_CShuffle
         B0BlockTransferDstScalarPerVector_BK1,
         true,
         B0BlockLdsExtraN,
+        CDE0BlockTransferSrcVectorDim,
+        CDE0BlockTransferSrcScalaerPerVector,
         B1BlockTransferThreadClusterLengths_BK0_N_BK1,
         B1BlockTransferThreadClusterArrangeOrder,
         B1BlockTransferSrcAccessOrder,
@@ -506,12 +510,15 @@ struct DeviceBatchedGemmMultipleDGemmMultipleD_Xdl_CShuffle
         CDE1ShuffleBlockTransferScalarPerVector_NPerBlock,
         LoopSched>;
 
-    using A0GridDesc_AK0_M_AK1 = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultA0GridDescriptor_AK0_M_AK1(A0GridDesc_M_K{}))>;
-    using B0GridDesc_BK0_N_BK1 = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultB0GridDescriptor_BK0_N_BK1(B0GridDesc_N_K{}))>;
-    using B1GridDesc_BK0_N_BK1 = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultB1GridDescriptor_BK0_N_BK1(B1GridDesc_N_K{}))>;
+    using A0GridDesc_AK0_M_AK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultA0GridDescriptor_AK0_M_AK1(
+            A0GridDesc_M_K{}))>;
+    using B0GridDesc_BK0_N_BK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultB0GridDescriptor_BK0_N_BK1(
+            B0GridDesc_N_K{}))>;
+    using B1GridDesc_BK0_N_BK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultB1GridDescriptor_BK0_N_BK1(
+            B1GridDesc_N_K{}))>;
 
     // Argument
     struct Argument : public BaseArgument
@@ -805,9 +812,7 @@ struct DeviceBatchedGemmMultipleDGemmMultipleD_Xdl_CShuffle
 
     static bool IsSupportedArgument(const Argument& arg)
     {
-        if(!(ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx90a" ||
-             ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx941" ||
-             ck::get_device_name() == "gfx942"))
+        if(!ck::is_xdl_supported())
         {
             return false;
         }

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

@@ -801,6 +801,11 @@ struct DeviceBatchedGemmReduce_Xdl_CShuffle : public DeviceGemmReduce<0, ReduceO
 
     static bool IsSupportedArgument(const Argument& arg)
     {
+        if(!ck::is_xdl_supported())
+        {
+            return false;
+        }
+
         return GridwiseGemm::CheckValidity(arg.a_grid_desc_ak0_m_ak1_,
                                            arg.b_grid_desc_bk0_n_bk1_,
                                            arg.c_grid_desc_m_n_,

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

@@ -723,9 +723,7 @@ struct DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle
         arg.Print();
 #endif
 
-        if(!(ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx90a" ||
-             ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx941" ||
-             ck::get_device_name() == "gfx942"))
+        if(!ck::is_xdl_supported())
         {
             return false;
         }

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

@@ -613,9 +613,7 @@ struct DeviceBatchedGemmSoftmaxGemm_Xdl_CShuffle
 
     static bool IsSupportedArgument(const Argument& arg)
     {
-        if(!(ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx90a" ||
-             ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx941" ||
-             ck::get_device_name() == "gfx942"))
+        if(!ck::is_xdl_supported())
         {
             return false;
         }

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

@@ -185,7 +185,7 @@ struct DeviceBatchedGemmXdl : public DeviceBatchedGemm<ALayout,
         AElementwiseOperation,
         BElementwiseOperation,
         CElementwiseOperation,
-        GemmSpecialization::MNPadding,
+        GemmSpecialization::MNKPadding,
         MPerBlock,
         NPerBlock,
         K0PerBlock,
@@ -310,7 +310,7 @@ struct DeviceBatchedGemmXdl : public DeviceBatchedGemm<ALayout,
 
     static bool IsSupportedArgument(const Problem& problem)
     {
-        if(problem.K % K1 != 0)
+        if(!ck::is_xdl_supported())
         {
             return false;
         }
@@ -411,7 +411,12 @@ struct DeviceBatchedGemmXdl : public DeviceBatchedGemm<ALayout,
             << BlockSize << ", "
             << MPerBlock << ", "
             << NPerBlock << ", "
-            << K0PerBlock
+            << K0PerBlock << ", "
+            << K1 << ", "
+            << MPerXDL << ", "
+            << NPerXDL << ", "
+            << MXdlPerWave << ", "
+            << NXdlPerWave << ", "
             << ">"
             << " NumGemmKPrefetchStage: "
             << NumGemmKPrefetchStage << ", "

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

@@ -123,7 +123,8 @@ struct DeviceCGemm_4Gemm_Xdl_CShuffle
         ALayout,
         BLayout,
         CLayout,
-        ADataType, // TODO: distinguish A/B datatype
+        ADataType,
+        BDataType,
         GemmAccDataType,
         CShuffleDataType,
         CDataType,
@@ -284,8 +285,11 @@ struct DeviceCGemm_4Gemm_Xdl_CShuffle
 
             if(GridwiseGemm::CalculateHasMainKBlockLoop(K))
             {
-                const auto kernel =
-                    kernel_gemm_xdl_cshuffle_v1<GridwiseGemm, ADataType, CDataType, true>;
+                const auto kernel = kernel_gemm_xdl_cshuffle_v1<GridwiseGemm,
+                                                                ADataType,
+                                                                BDataType,
+                                                                CDataType,
+                                                                true>;
 
                 ave_time += launch_and_time_kernel(stream_config,
                                                    kernel,
@@ -357,8 +361,11 @@ struct DeviceCGemm_4Gemm_Xdl_CShuffle
             }
             else
             {
-                const auto kernel =
-                    kernel_gemm_xdl_cshuffle_v1<GridwiseGemm, ADataType, CDataType, false>;
+                const auto kernel = kernel_gemm_xdl_cshuffle_v1<GridwiseGemm,
+                                                                ADataType,
+                                                                BDataType,
+                                                                CDataType,
+                                                                false>;
 
                 ave_time += launch_and_time_kernel(stream_config,
                                                    kernel,
@@ -448,6 +455,11 @@ struct DeviceCGemm_4Gemm_Xdl_CShuffle
 
     static bool IsSupportedArgument(const Argument& arg)
     {
+        if(!ck::is_xdl_supported())
+        {
+            return false;
+        }
+
         return GridwiseGemm::CheckValidity(arg);
     }
 

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

@@ -310,9 +310,13 @@ struct DeviceContractionMultipleD_Xdl_CShuffle
     using DsGridDesc_M_N = remove_cvref_t<decltype(MakeDsGridDescriptor_M_N({{}}, {{}}))>;
     using EGridDesc_M_N  = decltype(MakeEGridDescriptor_M_N({}, {}));
 
+    using ComputeDataType = ADataType;
+
     // GridwiseGemm
     using GridwiseGemm = GridwiseGemmMultipleD_xdl_cshuffle<
         ADataType, // TODO: distinguish A/B datatype
+        BDataType,
+        ComputeDataType,
         AccDataType,
         CShuffleDataType,
         DsDataType,
@@ -355,14 +359,18 @@ struct DeviceContractionMultipleD_Xdl_CShuffle
         LoopSched>;
 
     // desc for blockwise copy
-    using AGridDesc_AK0_M_AK1                          = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultAGridDescriptor_AK0_M_AK1(AGridDesc_M_K{}))>;
-    using BGridDesc_BK0_N_BK1                          = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDefaultBGridDescriptor_BK0_N_BK1(BGridDesc_N_K{}))>;
-    using DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock = remove_cvref_t<decltype(
-        GridwiseGemm::MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(DsGridDesc_M_N{}))>;
-    using EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock  = remove_cvref_t<decltype(
-        GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(EGridDesc_M_N{}))>;
+    using AGridDesc_AK0_M_AK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultAGridDescriptor_AK0_M_AK1(
+            AGridDesc_M_K{}))>;
+    using BGridDesc_BK0_N_BK1 =
+        remove_cvref_t<decltype(GridwiseGemm::MakeDefaultBGridDescriptor_BK0_N_BK1(
+            BGridDesc_N_K{}))>;
+    using DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock = remove_cvref_t<
+        decltype(GridwiseGemm::MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
+            DsGridDesc_M_N{}))>;
+    using EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock =
+        remove_cvref_t<decltype(GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
+            EGridDesc_M_N{}))>;
 
     // block-to-e-tile map
     using Block2ETileMap =
@@ -582,9 +590,7 @@ struct DeviceContractionMultipleD_Xdl_CShuffle
 
     static bool IsSupportedArgument(const Argument& arg)
     {
-        if(!(ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx90a" ||
-             ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx941" ||
-             ck::get_device_name() == "gfx942"))
+        if(!ck::is_xdl_supported())
         {
             return false;
         }

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

@@ -532,11 +532,12 @@ struct DeviceConv2dBwdWeightXdl_C_Shuffle_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_
             float ave_time = 0;
 
             const auto Run = [&](const auto& kernel) {
-                hipGetErrorString(hipMemset(
+                hipGetErrorString(hipMemsetAsync(
                     arg.p_c_grid_,
                     0,
                     arg.c_grid_desc_mblock_mperblock_nblock_nperblock_.GetElementSpaceSize() *
-                        sizeof(CDataType)));
+                        sizeof(CDataType),
+                    stream_config.stream_id_));
 
                 ave_time =
                     launch_and_time_kernel(stream_config,
@@ -649,6 +650,11 @@ struct DeviceConv2dBwdWeightXdl_C_Shuffle_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_
 
     static bool IsSupportedArgument(const Argument& arg)
     {
+        if(!ck::is_xdl_supported())
+        {
+            return false;
+        }
+
         // vector load A/B matrix from global memory
         if(!(ABlockTransferSrcVectorDim == 2 && BBlockTransferSrcVectorDim == 2 &&
              arg.Conv_K_ % ABlockTransferSrcScalarPerVector == 0 &&

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

@@ -616,6 +616,11 @@ struct DeviceConv2dBwdDataXdl_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_N_Ho_Wo_K
 
     static bool IsSupportedArgument(const Argument& arg)
     {
+        if(!ck::is_xdl_supported())
+        {
+            return false;
+        }
+
         if constexpr(ConvBackwardDataSpecialization ==
                      ConvolutionBackwardDataSpecialization::Filter1x1Stride1Pad0)
         {