Merge branch 'develop' into wavelet_model

include/ck/tensor_operation/gpu/warp/xdlops_gemm.hpp

@@ -819,7 +819,7 @@ struct XdlopsGemm
         index_t n_offset = blk_i * mfma_instr.n_per_blk + blk_td;
         index_t m_offset = xdlops_i * mfma_instr.m_per_blk + blk_id * mfma_instr.group_size;
 
-        return CIndex{m_offset, n_offset};
+        return TransposeC ? CIndex{n_offset, m_offset} : CIndex{m_offset, n_offset};
     }
 
     static constexpr auto mfma = MfmaSelector<base_type, MPerXdlops, NPerXdlops>{};

include/ck/tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck/utility/common_header.hpp"
+#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/convolution_forward_specialization.hpp"
+
+namespace ck {
+namespace tensor_operation {
+
+template <index_t NDimSpatial, device::ConvolutionForwardSpecialization ConvForwardSpecialization>
+struct TransformConvFwdToGemm
+{
+    static constexpr auto I1 = Number<1>{};
+
+    template <typename ALayout,
+              typename std::enable_if<NDimSpatial == 1 &&
+                                          is_same_v<ALayout, tensor_layout::convolution::GNWC>,
+                                      bool>::type = false>
+    static auto
+    MakeADescriptor_M_K(const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* a_g_n_c_wis_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* b_g_k_c_xs_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* c_g_n_k_wos_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)
+    {
+        const index_t N = a_g_n_c_wis_lengths[1];
+        const index_t C = a_g_n_c_wis_lengths[2];
+
+        const index_t Wi = a_g_n_c_wis_lengths[3];
+
+        const index_t Wo = c_g_n_k_wos_lengths[3];
+
+        const index_t ConvStrideW = conv_filter_strides[0];
+
+        if constexpr(ConvForwardSpecialization ==
+                     device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
+        {
+            const index_t NWo = N * std::accumulate(c_g_n_k_wos_lengths.begin() + 3,
+                                                    c_g_n_k_wos_lengths.begin() + 3 + NDimSpatial,
+                                                    index_t{1},
+                                                    std::multiplies<index_t>());
+
+            const auto in_gemmm_gemmk_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(NWo, C));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else if constexpr(ConvForwardSpecialization ==
+                          device::ConvolutionForwardSpecialization::Filter1x1Pad0)
+        {
+            const auto in_n_wi_c_desc = make_naive_tensor_descriptor_packed(make_tuple(N, Wi, C));
+
+            const auto in_n_wo_c_desc = transform_tensor_descriptor(
+                in_n_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_embed_transform(make_tuple(Wo), make_tuple(ConvStrideW)),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
+
+            const auto in_gemmm_gemmk_desc = transform_tensor_descriptor(
+                in_n_wo_c_desc,
+                make_tuple(make_merge_transform(make_tuple(N, Wo)), make_pass_through_transform(C)),
+                make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else
+        {
+            const index_t X             = b_g_k_c_xs_lengths[3];
+            const index_t ConvDilationW = conv_filter_dilations[0];
+            const index_t InLeftPadW    = input_left_pads[0];
+            const index_t InRightPadW   = input_right_pads[0];
+
+            const auto in_n_wi_c_desc = make_naive_tensor_descriptor_packed(make_tuple(N, Wi, C));
+
+            const auto in_n_wip_c_desc = transform_tensor_descriptor(
+                in_n_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_pad_transform(Wi, InLeftPadW, InRightPadW),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
+
+            const auto in_n_x_wo_c_desc = transform_tensor_descriptor(
+                in_n_wip_c_desc,
+                make_tuple(
+                    make_pass_through_transform(N),
+                    make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW)),
+                    make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}));
+
+            const auto in_gemmm_gemmk_desc =
+                transform_tensor_descriptor(in_n_x_wo_c_desc,
+                                            make_tuple(make_merge_transform(make_tuple(N, Wo)),
+                                                       make_merge_transform(make_tuple(X, C))),
+                                            make_tuple(Sequence<0, 2>{}, Sequence<1, 3>{}),
+                                            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+    }
+
+    template <typename ALayout,
+              typename std::enable_if<NDimSpatial == 2 &&
+                                          is_same_v<ALayout, tensor_layout::convolution::GNHWC>,
+                                      bool>::type = false>
+    static auto
+    MakeADescriptor_M_K(const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* a_g_n_c_wis_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* b_g_k_c_xs_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* c_g_n_k_wos_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)
+    {
+        const index_t N = a_g_n_c_wis_lengths[1];
+        const index_t C = a_g_n_c_wis_lengths[2];
+
+        const index_t Hi = a_g_n_c_wis_lengths[3];
+        const index_t Wi = a_g_n_c_wis_lengths[4];
+
+        const index_t Ho = c_g_n_k_wos_lengths[3];
+        const index_t Wo = c_g_n_k_wos_lengths[4];
+
+        const index_t ConvStrideH = conv_filter_strides[0];
+        const index_t ConvStrideW = conv_filter_strides[1];
+
+        if constexpr(ConvForwardSpecialization ==
+                     device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
+        {
+            const index_t NHoWo = N * std::accumulate(c_g_n_k_wos_lengths.begin() + 3,
+                                                      c_g_n_k_wos_lengths.begin() + 3 + NDimSpatial,
+                                                      index_t{1},
+                                                      std::multiplies<index_t>());
+
+            const auto in_gemmm_gemmk_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(NHoWo, C));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else if constexpr(ConvForwardSpecialization ==
+                          device::ConvolutionForwardSpecialization::Filter1x1Pad0)
+        {
+            const auto in_n_hi_wi_c_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(N, Hi, Wi, C));
+
+            const auto in_n_ho_wo_c_desc = transform_tensor_descriptor(
+                in_n_hi_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_embed_transform(make_tuple(Ho), make_tuple(ConvStrideH)),
+                           make_embed_transform(make_tuple(Wo), make_tuple(ConvStrideW)),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+            const auto in_gemmm_gemmk_desc =
+                transform_tensor_descriptor(in_n_ho_wo_c_desc,
+                                            make_tuple(make_merge_transform(make_tuple(N, Ho, Wo)),
+                                                       make_pass_through_transform(C)),
+                                            make_tuple(Sequence<0, 1, 2>{}, Sequence<3>{}),
+                                            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else
+        {
+            const index_t Y = b_g_k_c_xs_lengths[3];
+            const index_t X = b_g_k_c_xs_lengths[4];
+
+            const index_t ConvDilationH = conv_filter_dilations[0];
+            const index_t ConvDilationW = conv_filter_dilations[1];
+
+            const index_t InLeftPadH = input_left_pads[0];
+            const index_t InLeftPadW = input_left_pads[1];
+
+            const index_t InRightPadH = input_right_pads[0];
+            const index_t InRightPadW = input_right_pads[1];
+
+            const auto in_n_hi_wi_c_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(N, Hi, Wi, C));
+
+            const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
+                in_n_hi_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           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>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+            const auto in_n_y_ho_x_wo_c_desc = transform_tensor_descriptor(
+                in_n_hip_wip_c_desc,
+                make_tuple(
+                    make_pass_through_transform(N),
+                    make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                    make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW)),
+                    make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+                make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+            const auto in_gemmm_gemmk_desc =
+                transform_tensor_descriptor(in_n_y_ho_x_wo_c_desc,
+                                            make_tuple(make_merge_transform(make_tuple(N, Ho, Wo)),
+                                                       make_merge_transform(make_tuple(Y, X, C))),
+                                            make_tuple(Sequence<0, 2, 4>{}, Sequence<1, 3, 5>{}),
+                                            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+    }
+
+    template <typename ALayout,
+              typename std::enable_if<NDimSpatial == 3 &&
+                                          is_same_v<ALayout, tensor_layout::convolution::GNDHWC>,
+                                      bool>::type = false>
+    static auto
+    MakeADescriptor_M_K(const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* a_g_n_c_wis_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* b_g_k_c_xs_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* c_g_n_k_wos_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)
+    {
+        const index_t N = a_g_n_c_wis_lengths[1];
+        const index_t C = a_g_n_c_wis_lengths[2];
+
+        const index_t Di = a_g_n_c_wis_lengths[3];
+        const index_t Hi = a_g_n_c_wis_lengths[4];
+        const index_t Wi = a_g_n_c_wis_lengths[5];
+
+        const index_t Do = c_g_n_k_wos_lengths[3];
+        const index_t Ho = c_g_n_k_wos_lengths[4];
+        const index_t Wo = c_g_n_k_wos_lengths[5];
+
+        const index_t ConvStrideD = conv_filter_strides[0];
+        const index_t ConvStrideH = conv_filter_strides[1];
+        const index_t ConvStrideW = conv_filter_strides[2];
+
+        if constexpr(ConvForwardSpecialization ==
+                     device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
+        {
+            const index_t NDoHoWo =
+                N * std::accumulate(c_g_n_k_wos_lengths.begin() + 3,
+                                    c_g_n_k_wos_lengths.begin() + 3 + NDimSpatial,
+                                    index_t{1},
+                                    std::multiplies<index_t>());
+
+            const auto in_gemmm_gemmk_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(NDoHoWo, C));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else if constexpr(ConvForwardSpecialization ==
+                          device::ConvolutionForwardSpecialization::Filter1x1Pad0)
+        {
+            const auto in_n_di_hi_wi_c_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(N, Di, Hi, Wi, C));
+
+            const auto in_n_do_ho_wo_c_desc = transform_tensor_descriptor(
+                in_n_di_hi_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_embed_transform(make_tuple(Do), make_tuple(ConvStrideD)),
+                           make_embed_transform(make_tuple(Ho), make_tuple(ConvStrideH)),
+                           make_embed_transform(make_tuple(Wo), make_tuple(ConvStrideW)),
+                           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_gemmm_gemmk_desc = transform_tensor_descriptor(
+                in_n_do_ho_wo_c_desc,
+                make_tuple(make_merge_transform(make_tuple(N, Do, Ho, Wo)),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0, 1, 2, 3>{}, Sequence<4>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else
+        {
+            const index_t Z = b_g_k_c_xs_lengths[3];
+            const index_t Y = b_g_k_c_xs_lengths[4];
+            const index_t X = b_g_k_c_xs_lengths[5];
+
+            const index_t ConvDilationD = conv_filter_dilations[0];
+            const index_t ConvDilationH = conv_filter_dilations[1];
+            const index_t ConvDilationW = conv_filter_dilations[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 auto in_n_di_hi_wi_c_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(N, Di, Hi, Wi, C));
+
+            const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
+                in_n_di_hi_wi_c_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_z_do_y_ho_x_wo_c_desc = transform_tensor_descriptor(
+                in_n_hip_wip_c_desc,
+                make_tuple(
+                    make_pass_through_transform(N),
+                    make_embed_transform(make_tuple(Z, Do), make_tuple(ConvDilationD, ConvStrideD)),
+                    make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                    make_embed_transform(make_tuple(X, Wo), 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_gemmm_gemmk_desc = transform_tensor_descriptor(
+                in_n_z_do_y_ho_x_wo_c_desc,
+                make_tuple(make_merge_transform(make_tuple(N, Do, Ho, Wo)),
+                           make_merge_transform(make_tuple(Z, Y, X, C))),
+                make_tuple(Sequence<0, 2, 4, 6>{}, Sequence<1, 3, 5, 7>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+    }
+
+    // TODO: implement ck::tensor_layout::convolution that describe packed/strided dimemsion as
+    // properties
+    template <typename ALayout,
+              typename std::enable_if<NDimSpatial == 1 &&
+                                          (is_same_v<ALayout, tensor_layout::convolution::G_NW_C> ||
+                                           is_same_v<ALayout, tensor_layout::convolution::NWGC>),
+                                      bool>::type = false>
+    static auto
+    MakeADescriptor_M_K(const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_strides,
+                        const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* b_g_k_c_xs_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* c_g_n_k_wos_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)
+    {
+        const index_t N = a_g_n_c_wis_lengths[1];
+        const index_t C = a_g_n_c_wis_lengths[2];
+
+        const index_t Wi = a_g_n_c_wis_lengths[3];
+
+        const index_t Wo = c_g_n_k_wos_lengths[3];
+
+        const index_t ConvStrideW = conv_filter_strides[0];
+
+        if constexpr(ConvForwardSpecialization ==
+                     device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
+        {
+            const index_t NHoWo = N * std::accumulate(c_g_n_k_wos_lengths.begin() + 3,
+                                                      c_g_n_k_wos_lengths.begin() + 3 + NDimSpatial,
+                                                      index_t{1},
+                                                      std::multiplies<index_t>());
+
+            // This is different
+            const index_t WiStride = a_g_n_c_wis_strides[2 + NDimSpatial];
+            const auto CStride     = I1;
+
+            const auto in_gemmm_gemmk_desc =
+                make_naive_tensor_descriptor(make_tuple(NHoWo, C), make_tuple(WiStride, CStride));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else if constexpr(ConvForwardSpecialization ==
+                          device::ConvolutionForwardSpecialization::Filter1x1Pad0)
+        {
+            // This is different
+            const index_t NStride  = a_g_n_c_wis_strides[1];
+            const index_t WiStride = a_g_n_c_wis_strides[3];
+            const auto CStride     = I1;
+
+            const auto in_n_wi_c_desc = make_naive_tensor_descriptor(
+                make_tuple(N, Wi, C), make_tuple(NStride, WiStride, CStride));
+
+            const auto in_n_wo_c_desc = transform_tensor_descriptor(
+                in_n_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_embed_transform(make_tuple(Wo), make_tuple(ConvStrideW)),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
+
+            const auto in_gemmm_gemmk_desc = transform_tensor_descriptor(
+                in_n_wo_c_desc,
+                make_tuple(make_merge_transform(make_tuple(N, Wo)), make_pass_through_transform(C)),
+                make_tuple(Sequence<0, 1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else
+        {
+            const index_t X             = b_g_k_c_xs_lengths[3];
+            const index_t ConvDilationW = conv_filter_dilations[0];
+            const index_t InLeftPadW    = input_left_pads[0];
+            const index_t InRightPadW   = input_right_pads[0];
+
+            // This is different
+            const index_t NStride  = a_g_n_c_wis_strides[1];
+            const index_t WiStride = a_g_n_c_wis_strides[3];
+            const auto CStride     = I1;
+
+            const auto in_n_wi_c_desc = make_naive_tensor_descriptor(
+                make_tuple(N, Wi, C), make_tuple(NStride, WiStride, CStride));
+
+            const auto in_n_wip_c_desc = transform_tensor_descriptor(
+                in_n_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_pad_transform(Wi, InLeftPadW, InRightPadW),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
+
+            const auto in_n_x_wo_c_desc = transform_tensor_descriptor(
+                in_n_wip_c_desc,
+                make_tuple(
+                    make_pass_through_transform(N),
+                    make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW)),
+                    make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}));
+
+            const auto in_gemmm_gemmk_desc =
+                transform_tensor_descriptor(in_n_x_wo_c_desc,
+                                            make_tuple(make_merge_transform(make_tuple(N, Wo)),
+                                                       make_merge_transform(make_tuple(X, C))),
+                                            make_tuple(Sequence<0, 2>{}, Sequence<1, 3>{}),
+                                            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+    }
+
+    template <typename ALayout,
+              typename std::enable_if<
+                  NDimSpatial == 2 && (is_same_v<ALayout, tensor_layout::convolution::G_NHW_C> ||
+                                       is_same_v<ALayout, tensor_layout::convolution::NHWGC>),
+                  bool>::type = false>
+    static auto
+    MakeADescriptor_M_K(const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_strides,
+                        const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* b_g_k_c_xs_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* c_g_n_k_wos_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)
+    {
+        const index_t N = a_g_n_c_wis_lengths[1];
+        const index_t C = a_g_n_c_wis_lengths[2];
+
+        const index_t Hi = a_g_n_c_wis_lengths[3];
+        const index_t Wi = a_g_n_c_wis_lengths[4];
+
+        const index_t Ho = c_g_n_k_wos_lengths[3];
+        const index_t Wo = c_g_n_k_wos_lengths[4];
+
+        const index_t ConvStrideH = conv_filter_strides[0];
+        const index_t ConvStrideW = conv_filter_strides[1];
+
+        if constexpr(ConvForwardSpecialization ==
+                     device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
+        {
+            const index_t NHoWo = N * std::accumulate(c_g_n_k_wos_lengths.begin() + 3,
+                                                      c_g_n_k_wos_lengths.begin() + 3 + NDimSpatial,
+                                                      index_t{1},
+                                                      std::multiplies<index_t>());
+
+            // This is different
+            const index_t WiStride = a_g_n_c_wis_strides[2 + NDimSpatial];
+            const auto CStride     = I1;
+
+            const auto in_gemmm_gemmk_desc =
+                make_naive_tensor_descriptor(make_tuple(NHoWo, C), make_tuple(WiStride, CStride));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else if constexpr(ConvForwardSpecialization ==
+                          device::ConvolutionForwardSpecialization::Filter1x1Pad0)
+        {
+            // This is different
+            const index_t NStride  = a_g_n_c_wis_strides[1];
+            const index_t HiStride = a_g_n_c_wis_strides[3];
+            const index_t WiStride = a_g_n_c_wis_strides[4];
+            const auto CStride     = I1;
+
+            const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor(
+                make_tuple(N, Hi, Wi, C), make_tuple(NStride, HiStride, WiStride, CStride));
+
+            const auto in_n_ho_wo_c_desc = transform_tensor_descriptor(
+                in_n_hi_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_embed_transform(make_tuple(Ho), make_tuple(ConvStrideH)),
+                           make_embed_transform(make_tuple(Wo), make_tuple(ConvStrideW)),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+            const auto in_gemmm_gemmk_desc =
+                transform_tensor_descriptor(in_n_ho_wo_c_desc,
+                                            make_tuple(make_merge_transform(make_tuple(N, Ho, Wo)),
+                                                       make_pass_through_transform(C)),
+                                            make_tuple(Sequence<0, 1, 2>{}, Sequence<3>{}),
+                                            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else
+        {
+            const index_t Y = b_g_k_c_xs_lengths[3];
+            const index_t X = b_g_k_c_xs_lengths[4];
+
+            const index_t ConvDilationH = conv_filter_dilations[0];
+            const index_t ConvDilationW = conv_filter_dilations[1];
+
+            const index_t InLeftPadH = input_left_pads[0];
+            const index_t InLeftPadW = input_left_pads[1];
+
+            const index_t InRightPadH = input_right_pads[0];
+            const index_t InRightPadW = input_right_pads[1];
+
+            // This is different
+            const index_t NStride  = a_g_n_c_wis_strides[1];
+            const index_t HiStride = a_g_n_c_wis_strides[3];
+            const index_t WiStride = a_g_n_c_wis_strides[4];
+            const auto CStride     = I1;
+
+            const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor(
+                make_tuple(N, Hi, Wi, C), make_tuple(NStride, HiStride, WiStride, CStride));
+
+            const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
+                in_n_hi_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           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>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+            const auto in_n_y_ho_x_wo_c_desc = transform_tensor_descriptor(
+                in_n_hip_wip_c_desc,
+                make_tuple(
+                    make_pass_through_transform(N),
+                    make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                    make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW)),
+                    make_pass_through_transform(C)),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+                make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+            const auto in_gemmm_gemmk_desc =
+                transform_tensor_descriptor(in_n_y_ho_x_wo_c_desc,
+                                            make_tuple(make_merge_transform(make_tuple(N, Ho, Wo)),
+                                                       make_merge_transform(make_tuple(Y, X, C))),
+                                            make_tuple(Sequence<0, 2, 4>{}, Sequence<1, 3, 5>{}),
+                                            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+    }
+
+    template <typename ALayout,
+              typename std::enable_if<
+                  NDimSpatial == 3 && (is_same_v<ALayout, tensor_layout::convolution::G_NDHW_C> ||
+                                       is_same_v<ALayout, tensor_layout::convolution::NDHWGC>),
+                  bool>::type = false>
+    static auto
+    MakeADescriptor_M_K(const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& a_g_n_c_wis_strides,
+                        const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* b_g_k_c_xs_strides */,
+                        const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* c_g_n_k_wos_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)
+    {
+        const index_t N = a_g_n_c_wis_lengths[1];
+        const index_t C = a_g_n_c_wis_lengths[2];
+
+        const index_t Di = a_g_n_c_wis_lengths[3];
+        const index_t Hi = a_g_n_c_wis_lengths[4];
+        const index_t Wi = a_g_n_c_wis_lengths[5];
+
+        const index_t Do = c_g_n_k_wos_lengths[3];
+        const index_t Ho = c_g_n_k_wos_lengths[4];
+        const index_t Wo = c_g_n_k_wos_lengths[5];
+
+        const index_t ConvStrideD = conv_filter_strides[0];
+        const index_t ConvStrideH = conv_filter_strides[1];
+        const index_t ConvStrideW = conv_filter_strides[2];
+
+        if constexpr(ConvForwardSpecialization ==
+                     device::ConvolutionForwardSpecialization::Filter1x1Stride1Pad0)
+        {
+            const index_t NDoHoWo =
+                N * std::accumulate(c_g_n_k_wos_lengths.begin() + 3,
+                                    c_g_n_k_wos_lengths.begin() + 3 + NDimSpatial,
+                                    index_t{1},
+                                    std::multiplies<index_t>());
+
+            // This is different
+            const index_t WiStride = a_g_n_c_wis_strides[2 + NDimSpatial];
+            const auto CStride     = I1;
+
+            const auto in_gemmm_gemmk_desc =
+                make_naive_tensor_descriptor(make_tuple(NDoHoWo, C), make_tuple(WiStride, CStride));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else if constexpr(ConvForwardSpecialization ==
+                          device::ConvolutionForwardSpecialization::Filter1x1Pad0)
+        {
+            // This is different
+            const index_t NStride  = a_g_n_c_wis_strides[1];
+            const index_t DiStride = a_g_n_c_wis_strides[3];
+            const index_t HiStride = a_g_n_c_wis_strides[4];
+            const index_t WiStride = a_g_n_c_wis_strides[5];
+            const auto CStride     = I1;
+
+            const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor(
+                make_tuple(N, Di, Hi, Wi, C),
+                make_tuple(NStride, DiStride, HiStride, WiStride, CStride));
+
+            const auto in_n_do_ho_wo_c_desc = transform_tensor_descriptor(
+                in_n_di_hi_wi_c_desc,
+                make_tuple(make_pass_through_transform(N),
+                           make_embed_transform(make_tuple(Do), make_tuple(ConvStrideD)),
+                           make_embed_transform(make_tuple(Ho), make_tuple(ConvStrideH)),
+                           make_embed_transform(make_tuple(Wo), make_tuple(ConvStrideW)),
+                           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_gemmm_gemmk_desc = transform_tensor_descriptor(
+                in_n_do_ho_wo_c_desc,
+                make_tuple(make_merge_transform(make_tuple(N, Do, Ho, Wo)),
+                           make_pass_through_transform(C)),
+                make_tuple(Sequence<0, 1, 2, 3>{}, Sequence<4>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+        else
+        {
+            const index_t Z = b_g_k_c_xs_lengths[3];
+            const index_t Y = b_g_k_c_xs_lengths[4];
+            const index_t X = b_g_k_c_xs_lengths[5];
+
+            const index_t ConvDilationD = conv_filter_dilations[0];
+            const index_t ConvDilationH = conv_filter_dilations[1];
+            const index_t ConvDilationW = conv_filter_dilations[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];
+
+            // This is different
+            const index_t NStride  = a_g_n_c_wis_strides[1];
+            const index_t DiStride = a_g_n_c_wis_strides[3];
+            const index_t HiStride = a_g_n_c_wis_strides[4];
+            const index_t WiStride = a_g_n_c_wis_strides[5];
+            const auto CStride     = I1;
+
+            const auto in_n_di_hi_wi_c_desc = make_naive_tensor_descriptor(
+                make_tuple(N, Di, Hi, Wi, C),
+                make_tuple(NStride, DiStride, HiStride, WiStride, CStride));
+
+            const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
+                in_n_di_hi_wi_c_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_z_do_y_ho_x_wo_c_desc = transform_tensor_descriptor(
+                in_n_hip_wip_c_desc,
+                make_tuple(
+                    make_pass_through_transform(N),
+                    make_embed_transform(make_tuple(Z, Do), make_tuple(ConvDilationD, ConvStrideD)),
+                    make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                    make_embed_transform(make_tuple(X, Wo), 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_gemmm_gemmk_desc = transform_tensor_descriptor(
+                in_n_z_do_y_ho_x_wo_c_desc,
+                make_tuple(make_merge_transform(make_tuple(N, Do, Ho, Wo)),
+                           make_merge_transform(make_tuple(Z, Y, X, C))),
+                make_tuple(Sequence<0, 2, 4, 6>{}, Sequence<1, 3, 5, 7>{}),
+                make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+            return in_gemmm_gemmk_desc;
+        }
+    }
+
+    template <typename BLayout,
+              typename std::enable_if<is_same_v<BLayout, tensor_layout::convolution::GKXC> ||
+                                          is_same_v<BLayout, tensor_layout::convolution::GKYXC> ||
+                                          is_same_v<BLayout, tensor_layout::convolution::GKZYXC>,
+                                      bool>::type = false>
+    static auto
+    MakeBDescriptor_N_K(const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* b_g_k_c_xs_strides */)
+    {
+        const index_t K = b_g_k_c_xs_lengths[1];
+        const index_t C = b_g_k_c_xs_lengths[2];
+
+        const index_t YX = std::accumulate(b_g_k_c_xs_lengths.begin() + 3,
+                                           b_g_k_c_xs_lengths.begin() + 3 + NDimSpatial,
+                                           index_t{1},
+                                           std::multiplies<index_t>());
+
+        const auto wei_gemmn_gemmk_desc =
+            make_naive_tensor_descriptor_packed(make_tuple(K, YX * C));
+
+        return wei_gemmn_gemmk_desc;
+    }
+
+    template <
+        typename BLayout,
+        typename std::enable_if<is_same_v<BLayout, tensor_layout::convolution::G_K_X_C> ||
+                                    is_same_v<BLayout, tensor_layout::convolution::G_K_YX_C> ||
+                                    is_same_v<BLayout, tensor_layout::convolution::G_K_ZYX_C> ||
+                                    is_same_v<BLayout, tensor_layout::convolution::KXGC> ||
+                                    is_same_v<BLayout, tensor_layout::convolution::KYXGC> ||
+                                    is_same_v<BLayout, tensor_layout::convolution::KZYXGC>,
+                                bool>::type = false>
+    static auto MakeBDescriptor_N_K(const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
+                                    const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_strides)
+    {
+        const index_t K = b_g_k_c_xs_lengths[1];
+        const index_t C = b_g_k_c_xs_lengths[2];
+
+        const index_t YX = std::accumulate(b_g_k_c_xs_lengths.begin() + 3,
+                                           b_g_k_c_xs_lengths.begin() + 3 + NDimSpatial,
+                                           index_t{1},
+                                           std::multiplies<index_t>());
+
+        const index_t KStride = b_g_k_c_xs_strides[1];
+        const index_t XStride = b_g_k_c_xs_strides[2 + NDimSpatial];
+        const auto CStride    = I1;
+
+        const auto wei_k_yx_c_desc = make_naive_tensor_descriptor(
+            make_tuple(K, YX, C), make_tuple(KStride, XStride, CStride));
+
+        const auto wei_gemmn_gemmk_desc = transform_tensor_descriptor(
+            wei_k_yx_c_desc,
+            make_tuple(make_pass_through_transform(K), make_merge_transform(make_tuple(YX, C))),
+            make_tuple(Sequence<0>{}, Sequence<1, 2>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        return wei_gemmn_gemmk_desc;
+    }
+
+    template <typename CLayout,
+              typename std::enable_if<is_same_v<CLayout, tensor_layout::convolution::GNWK> ||
+                                          is_same_v<CLayout, tensor_layout::convolution::GNHWK> ||
+                                          is_same_v<CLayout, tensor_layout::convolution::GNDHWK>,
+                                      bool>::type = false>
+    static auto
+    MakeCDescriptor_M_N(const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_lengths,
+                        const std::array<index_t, NDimSpatial + 3>& /* c_g_n_k_wos_strides */)
+    {
+        const index_t N = c_g_n_k_wos_lengths[1];
+        const index_t K = c_g_n_k_wos_lengths[2];
+
+        const index_t NHoWo = N * std::accumulate(c_g_n_k_wos_lengths.begin() + 3,
+                                                  c_g_n_k_wos_lengths.begin() + 3 + NDimSpatial,
+                                                  index_t{1},
+                                                  std::multiplies<index_t>());
+
+        const auto out_gemmm_gemmn_desc = make_naive_tensor_descriptor_packed(make_tuple(NHoWo, K));
+
+        return out_gemmm_gemmn_desc;
+    }
+
+    template <
+        typename CLayout,
+        typename std::enable_if<is_same_v<CLayout, tensor_layout::convolution::G_NW_K> ||
+                                    is_same_v<CLayout, tensor_layout::convolution::G_NHW_K> ||
+                                    is_same_v<CLayout, tensor_layout::convolution::G_NDHW_K> ||
+                                    is_same_v<CLayout, tensor_layout::convolution::NWGK> ||
+                                    is_same_v<CLayout, tensor_layout::convolution::NHWGK> ||
+                                    is_same_v<CLayout, tensor_layout::convolution::NDHWGK>,
+                                bool>::type = false>
+    static auto MakeCDescriptor_M_N(const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_lengths,
+                                    const std::array<index_t, NDimSpatial + 3>& c_g_n_k_wos_strides)
+    {
+        const index_t N = c_g_n_k_wos_lengths[1];
+        const index_t K = c_g_n_k_wos_lengths[2];
+
+        const auto KStride     = I1;
+        const index_t WoStride = c_g_n_k_wos_strides[NDimSpatial + 2];
+
+        const index_t NHoWo = N * std::accumulate(c_g_n_k_wos_lengths.begin() + 3,
+                                                  c_g_n_k_wos_lengths.begin() + 3 + NDimSpatial,
+                                                  index_t{1},
+                                                  std::multiplies<index_t>());
+
+        const auto out_gemmm_gemmn_desc =
+            make_naive_tensor_descriptor(make_tuple(NHoWo, K), make_tuple(WoStride, KStride));
+
+        return out_gemmm_gemmn_desc;
+    }
+};
+
+} // namespace tensor_operation
+} // namespace ck

include/ck/utility/amd_buffer_addressing.hpp

@@ -34,6 +34,21 @@ __device__ int32x4_t make_wave_buffer_resource(T* p_wave, index_t element_space_
     return wave_buffer_resource.content;
 }
 
+template <typename T>
+__device__ int32x4_t make_wave_buffer_resource_with_default_range(T* p_wave)
+{
+    BufferResource<T> wave_buffer_resource;
+
+    // wavewise base address (64 bit)
+    wave_buffer_resource.address(Number<0>{}) = const_cast<remove_cv_t<T>*>(p_wave);
+    // wavewise range (32 bit)
+    wave_buffer_resource.range(Number<2>{}) = 0xffffffff; // max possible range
+    // wavewise setting (32 bit)
+    wave_buffer_resource.config(Number<3>{}) = CK_BUFFER_RESOURCE_3RD_DWORD;
+
+    return wave_buffer_resource.content;
+}
+
 // buffer load i8
 __device__ int8_t
 llvm_amdgcn_raw_buffer_load_i8(int32x4_t srsrc,

include/ck/utility/data_type.hpp

@@ -9,6 +9,9 @@ namespace ck {
 
 using bhalf_t = ushort;
 using half_t  = _Float16;
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+using int4_t = _BitInt(4);
+#endif
 
 // vector_type
 template <typename T, index_t N>
@@ -130,6 +133,15 @@ struct scalar_type<int8_t>
     static constexpr index_t vector_size = 1;
 };
 
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+template <>
+struct scalar_type<int4_t>
+{
+    using type                           = int4_t;
+    static constexpr index_t vector_size = 1;
+};
+#endif
+
 //
 template <typename T>
 struct vector_type<T, 1>
@@ -1011,6 +1023,8 @@ struct NumericLimits
     {
         return std::numeric_limits<T>::quiet_NaN();
     }
+
+    __host__ __device__ static constexpr T Infinity() { return std::numeric_limits<T>::infinity(); }
 };
 
 template <>
@@ -1030,4 +1044,16 @@ struct NumericLimits<half_t>
     __host__ __device__ static constexpr half_t QuietNaN() { return bit_cast<half_t>(binary_qnan); }
 };
 
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+template <>
+struct NumericLimits<int4_t>
+{
+    __host__ __device__ static constexpr int4_t Min() { return int4_t(-8); }
+
+    __host__ __device__ static constexpr int4_t Max() { return int4_t(7); }
+
+    __host__ __device__ static constexpr int4_t Lowest() { return int4_t(-8); }
+};
+#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+
 } // namespace ck

include/ck/utility/functional.hpp

@@ -114,4 +114,18 @@ struct conditional<false, X, Y>
 template <bool predicate, class X, class Y>
 using conditional_t = typename conditional<predicate, X, Y>::type;
 
+// z = predicate ? x : y
+template <bool predicate, typename X, typename Y>
+constexpr auto conditional_expr(X&& x, Y&& y)
+{
+    if constexpr(predicate)
+    {
+        return std::forward<X>(x);
+    }
+    else
+    {
+        return std::forward<Y>(y);
+    }
+}
+
 } // namespace ck

include/ck/utility/math_v2.hpp

@@ -42,6 +42,14 @@ static inline __host__ half_t abs(half_t x)
     return abs_x;
 };
 
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+static inline __host__ int4_t abs(int4_t x)
+{
+    int4_t sgn = x >> (4 - 1);
+    return (x ^ sgn) - sgn;
+}
+#endif
+
 static inline __host__ bool isnan(float x) { return std::isnan(x); };
 
 static inline __host__ bool isnan(double x) { return std::isnan(x); };
@@ -65,6 +73,14 @@ static inline __host__ bool isnan(half_t x)
     return (xx & 0x7FFF) > 0x7C00;
 };
 
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+static inline __host__ bool isnan(int4_t x)
+{
+    (void)x;
+    return false;
+};
+#endif
+
 static inline __host__ float sqrt(float x) { return std::sqrt(x); };
 
 static inline __host__ double sqrt(double x) { return std::sqrt(x); };
@@ -89,6 +105,15 @@ static inline __device__ int32_t abs(int32_t x)
     return (x ^ sgn) - sgn;
 };
 
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+static inline __device__ int4_t abs(int4_t x)
+{
+    int4_t sgn = x >> (4 - 1);
+
+    return (x ^ sgn) - sgn;
+};
+#endif
+
 static inline __device__ half_t abs(half_t x) { return ::__habs(x); };
 
 static inline __device__ bool isnan(float x) { return ::isnan(x); };
@@ -107,6 +132,14 @@ static inline __device__ bool isnan(int32_t x)
     return false;
 };
 
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+static inline __device__ bool isnan(int4_t x)
+{
+    (void)x;
+    return false;
+};
+#endif
+
 static inline __device__ bool isnan(half_t x) { return ::__hisnan(x); };
 
 static inline __device__ float sqrt(float x) { return ::sqrtf(x); };

include/ck/utility/reduction_operator.hpp

@@ -21,9 +21,9 @@ namespace reduce {
 //                    vector space
 //                    (http://pages.cs.wisc.edu/~matthewb/pages/notes/pdf/linearalgebra/VectorSpaces.pdf).
 // 2) IsCompatibleInMemoryDataOperation() -- return true if the reduction task corresponding to this
-// operator can use the InMemoryDataOperation to finalize, or else it return false 3) operator() --
-// the first argument of the operator must be both an input & output, and the corresponding variable
-// usually stores
+// operator can use the InMemoryDataOperation to finalize, or else it return false
+// 3) operator() -- the first argument of the operator must be both an input & output, and the
+//                  corresponding variable usually stores
 //                  the accumulated result of many operator() calls; the second argument is only an
 //                  input. For indexable binary
 //                  operator, the second version of operator() has third argument (which is an
@@ -79,7 +79,7 @@ struct SquaredAdd
         static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                           is_same<T, half_t>::value || is_same<T, int32_t>::value ||
                           is_same<T, int8_t>::value,
-                      "The data type is not supported by the Max accumulator!");
+                      "The data type is not supported by the SquaredAdd accumulator!");
 
         a = a + b * b;
     }

library/include/ck/library/reference_tensor_operation/cpu/reference_batched_gemm.hpp

@@ -83,8 +83,8 @@ struct ReferenceBatchedGemm : public device::BaseOperator
             make_ParallelTensorFunctor(f_gmk_gkn_gmn,
                                        arg.c_g_m_n_.mDesc.GetLengths()[0],
                                        arg.c_g_m_n_.mDesc.GetLengths()[1],
-                                       arg.c_g_m_n_.mDesc.GetLengths()[2])();
-
+                                       arg.c_g_m_n_.mDesc.GetLengths()[2])(
+                std::thread::hardware_concurrency());
             return 0;
         }
 

library/include/ck/library/reference_tensor_operation/cpu/reference_batchnorm_forward_nhwc_c.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <vector>
+#include <array>
+#include <algorithm>
+#include <thread>
+
+#include "ck/tensor_operation/gpu/device/device_batchnorm_forward.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace host {
+
+template <typename InOutDataType, typename AccDataType>
+struct ReferenceBatchNormFwd_Input_N_H_W_C_Output_C : public device::DeviceBatchNormFwd<4, 3>
+{
+    struct Argument : public device::BaseArgument
+    {
+        Argument(const std::array<index_t, 4> xyLengths,
+                 const std::array<index_t, 4> xStrides,
+                 const std::array<index_t, 4> yStrides,
+                 const std::array<index_t, 1> bnScaleBiasMeanVarLengths,
+                 const std::array<index_t, 1> bnScaleBiasMeanVarStrides,
+                 const InOutDataType* p_x,
+                 const AccDataType* bnScale,
+                 const AccDataType* bnBias,
+                 InOutDataType* p_y,
+                 double exponentialAverageFactor,
+                 AccDataType* resultRunningMean,
+                 AccDataType* resultRunningVariance,
+                 double epsilon,
+                 AccDataType* resultSaveMean,
+                 AccDataType* resultSaveInvVariance)
+            : p_x_(p_x),
+              bnScale_(bnScale),
+              bnBias_(bnBias),
+              p_y_(p_y),
+              resultRunningMean_(resultRunningMean),
+              resultRunningVariance_(resultRunningVariance),
+              resultSaveMean_(resultSaveMean),
+              resultSaveInvVariance_(resultSaveInvVariance),
+              exponentialAverageFactor_(exponentialAverageFactor),
+              epsilon_(epsilon)
+        {
+            (void)xStrides;
+            (void)yStrides;
+            (void)bnScaleBiasMeanVarStrides;
+
+            if(xyLengths.size() != 4 || bnScaleBiasMeanVarLengths.size() != 1 ||
+               bnScaleBiasMeanVarLengths[0] != xyLengths[3])
+                throw std::runtime_error("Invalid tensor dimensions!");
+
+            n = xyLengths[0];
+            h = xyLengths[1];
+            w = xyLengths[2];
+            c = xyLengths[3];
+
+            resultSave    = (resultSaveMean != nullptr && resultSaveInvVariance != nullptr);
+            resultRunning = (resultRunningMean != nullptr && resultRunningVariance != nullptr);
+        }
+
+        const InOutDataType* p_x_;
+        const AccDataType* bnScale_;
+        const AccDataType* bnBias_;
+        InOutDataType* p_y_;
+
+        AccDataType* resultRunningMean_;
+        AccDataType* resultRunningVariance_;
+        AccDataType* resultSaveMean_;
+        AccDataType* resultSaveInvVariance_;
+
+        bool resultSave, resultRunning;
+
+        index_t n, h, w, c;
+
+        double exponentialAverageFactor_;
+        double epsilon_;
+    };
+
+    struct Invoker : public device::BaseInvoker
+    {
+        float Run(const Argument& arg)
+        {
+            auto thread_reduce_func = [&](auto iC) {
+                AccDataType reduceSize = type_convert<AccDataType>(arg.n) *
+                                         type_convert<AccDataType>(arg.h) *
+                                         type_convert<AccDataType>(arg.w);
+                index_t offset_C       = iC;
+                AccDataType mean       = type_convert<AccDataType>(0.0f);
+                AccDataType meansquare = type_convert<AccDataType>(0.0f);
+
+                // compute mean, meanquare, variance, invVariance
+                for(index_t iN = 0; iN < arg.n; iN++)
+                {
+                    index_t offset_N = iN * arg.h * arg.w * arg.c;
+                    for(index_t iH = 0; iH < arg.h; iH++)
+                    {
+                        index_t offset_H = iH * arg.w * arg.c;
+                        for(index_t iW = 0; iW < arg.w; iW++)
+                        {
+                            index_t offset_W = iW * arg.c;
+
+                            auto offset = offset_N + offset_H + offset_W + offset_C;
+
+                            AccDataType x = type_convert<AccDataType>(arg.p_x_[offset]);
+
+                            mean += x;
+                            meansquare += x * x;
+                        };
+                    }
+                };
+
+                mean       = mean / reduceSize;
+                meansquare = meansquare / reduceSize;
+
+                AccDataType variance = meansquare - mean * mean;
+                AccDataType invVariance =
+                    type_convert<AccDataType>(1.0f) /
+                    std::sqrt(type_convert<AccDataType>(arg.epsilon_) + variance);
+
+                // save the mean/invVariance if required
+                if(arg.resultSave)
+                {
+                    arg.resultSaveMean_[iC]        = mean;
+                    arg.resultSaveInvVariance_[iC] = invVariance;
+                };
+
+                // update the moving average if required
+                if(arg.resultRunning)
+                {
+                    arg.resultRunningMean_[iC] =
+                        arg.resultRunningMean_[iC] *
+                            type_convert<AccDataType>(1.0 - arg.exponentialAverageFactor_) +
+                        mean * arg.exponentialAverageFactor_;
+                    arg.resultRunningVariance_[iC] =
+                        arg.resultRunningVariance_[iC] *
+                            type_convert<AccDataType>(1.0 - arg.exponentialAverageFactor_) +
+                        variance * arg.exponentialAverageFactor_;
+                };
+
+                // Normalization
+                for(index_t iN = 0; iN < arg.n; iN++)
+                {
+                    index_t offset_N = iN * arg.h * arg.w * arg.c;
+                    for(index_t iH = 0; iH < arg.h; iH++)
+                    {
+                        index_t offset_H = iH * arg.w * arg.c;
+                        for(index_t iW = 0; iW < arg.w; iW++)
+                        {
+                            index_t offset_W = iW * arg.c;
+
+                            auto offset = offset_N + offset_H + offset_W + offset_C;
+
+                            AccDataType x = type_convert<AccDataType>(arg.p_x_[offset]);
+
+                            AccDataType norm_x =
+                                arg.bnScale_[iC] * (x - mean) * invVariance + arg.bnBias_[iC];
+
+                            arg.p_y_[offset] = type_convert<InOutDataType>(norm_x);
+                        };
+                    }
+                };
+            };
+
+            std::size_t num_thread      = std::thread::hardware_concurrency();
+            std::size_t work_per_thread = (arg.c + num_thread - 1) / num_thread;
+
+            std::vector<joinable_thread> threads(num_thread);
+
+            for(std::size_t it = 0; it < num_thread; ++it)
+            {
+                std::size_t ic_begin = it * work_per_thread;
+                std::size_t ic_end = std::min(static_cast<int>((it + 1) * work_per_thread), arg.c);
+
+                auto f = [=] {
+                    for(std::size_t ic = ic_begin; ic < ic_end; ++ic)
+                    {
+                        thread_reduce_func(ic);
+                    }
+                };
+
+                threads[it] = joinable_thread(f);
+            }
+
+            return (0.0f);
+        };
+
+        float Run(const device::BaseArgument* p_arg,
+                  const StreamConfig& /*stream_config*/ = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg));
+        };
+    };
+
+    bool IsSupportedArgument(const device::BaseArgument* p_arg) override
+    {
+        (void)p_arg;
+
+        return (true);
+    };
+
+    std::unique_ptr<device::BaseArgument>
+    MakeArgumentPointer(const std::array<index_t, 4> xyLengths,
+                        const std::array<index_t, 4> xStrides,
+                        const std::array<index_t, 4> yStrides,
+                        const std::array<index_t, 1> bnScaleBiasMeanVarLengths,
+                        const std::array<index_t, 1> bnScaleBiasMeanVarStrides,
+                        const void* p_x,
+                        const void* bnScale,
+                        const void* bnBias,
+                        void* p_y,
+                        double exponentialAverageFactor,
+                        void* resultRunningMean,
+                        void* resultRunningVariance,
+                        double epsilon,
+                        void* resultSaveMean,
+                        void* resultSaveInvVariance) override
+    {
+        return std::make_unique<Argument>(xyLengths,
+                                          xStrides,
+                                          yStrides,
+                                          bnScaleBiasMeanVarLengths,
+                                          bnScaleBiasMeanVarStrides,
+                                          static_cast<const InOutDataType*>(p_x),
+                                          static_cast<const AccDataType*>(bnScale),
+                                          static_cast<const AccDataType*>(bnBias),
+                                          static_cast<InOutDataType*>(p_y),
+                                          exponentialAverageFactor,
+                                          static_cast<AccDataType*>(resultRunningMean),
+                                          static_cast<AccDataType*>(resultRunningVariance),
+                                          epsilon,
+                                          static_cast<AccDataType*>(resultSaveMean),
+                                          static_cast<AccDataType*>(resultSaveInvVariance));
+    };
+
+    std::unique_ptr<device::BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>();
+    };
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "Reference_BatchNorm_Forward_NHWC_C<" << std::endl;
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace host
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/reference_tensor_operation/cpu/reference_batchnorm_infer_nhwc_c.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <vector>
+#include <array>
+#include <algorithm>
+
+#include "ck/tensor_operation/gpu/device/device_batchnorm_infer.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace host {
+
+template <typename InOutDataType, typename AccDataType>
+struct ReferenceBatchNormInfer_Input_N_H_W_C_Output_C : public device::DeviceBatchNormInfer<4, 3>
+{
+    struct Argument : public device::BaseArgument
+    {
+        Argument(const std::array<index_t, 4> xyLengths,
+                 const std::array<index_t, 4> xStrides,
+                 const std::array<index_t, 4> yStrides,
+                 const std::array<index_t, 1> bnScaleBiasMeanVarLengths,
+                 const std::array<index_t, 1> bnScaleBiasMeanVarStrides,
+                 const InOutDataType* p_x,
+                 const AccDataType* bnScale,
+                 const AccDataType* bnBias,
+                 double epsilon,
+                 const AccDataType* estimatedMean,
+                 const AccDataType* estimatedVariance,
+                 InOutDataType* p_y)
+            : p_x_(p_x),
+              bnScale_(bnScale),
+              bnBias_(bnBias),
+              epsilon_(epsilon),
+              estimatedMean_(estimatedMean),
+              estimatedVariance_(estimatedVariance),
+              p_y_(p_y)
+        {
+            (void)xStrides;
+            (void)yStrides;
+            (void)bnScaleBiasMeanVarStrides;
+
+            if(xyLengths.size() != 4 || bnScaleBiasMeanVarLengths.size() != 1 ||
+               bnScaleBiasMeanVarLengths[0] != xyLengths[3])
+                throw std::runtime_error("Invalid tensor dimensions!");
+
+            n = xyLengths[0];
+            h = xyLengths[1];
+            w = xyLengths[2];
+            c = xyLengths[3];
+        }
+
+        const InOutDataType* p_x_;
+        const AccDataType* bnScale_;
+        const AccDataType* bnBias_;
+
+        double epsilon_;
+
+        const AccDataType* estimatedMean_;
+        const AccDataType* estimatedVariance_;
+
+        InOutDataType* p_y_;
+
+        index_t n, h, w, c;
+    };
+
+    struct Invoker : public device::BaseInvoker
+    {
+        float Run(const Argument& arg)
+        {
+            auto thread_reduce_func = [&](auto iC) {
+                index_t offset_C     = iC;
+                AccDataType mean     = arg.estimatedMean_[offset_C];
+                AccDataType variance = arg.estimatedVariance_[offset_C];
+
+                AccDataType invVariance =
+                    type_convert<AccDataType>(1.0f) /
+                    std::sqrt(type_convert<AccDataType>(arg.epsilon_) + variance);
+
+                // Normalization
+                for(index_t iN = 0; iN < arg.n; iN++)
+                {
+                    index_t offset_N = iN * arg.h * arg.w * arg.c;
+                    for(index_t iH = 0; iH < arg.h; iH++)
+                    {
+                        index_t offset_H = iH * arg.w * arg.c;
+                        for(index_t iW = 0; iW < arg.w; iW++)
+                        {
+                            index_t offset_W = iW * arg.c;
+
+                            auto offset = offset_N + offset_H + offset_W + offset_C;
+
+                            AccDataType x = type_convert<AccDataType>(arg.p_x_[offset]);
+
+                            AccDataType norm_x =
+                                arg.bnScale_[iC] * (x - mean) * invVariance + arg.bnBias_[iC];
+
+                            arg.p_y_[offset] = type_convert<InOutDataType>(norm_x);
+                        };
+                    }
+                };
+            };
+
+            std::size_t num_thread      = std::thread::hardware_concurrency();
+            std::size_t work_per_thread = (arg.c + num_thread - 1) / num_thread;
+
+            std::vector<joinable_thread> threads(num_thread);
+
+            for(std::size_t it = 0; it < num_thread; ++it)
+            {
+                std::size_t ic_begin = it * work_per_thread;
+                std::size_t ic_end = std::min(static_cast<int>((it + 1) * work_per_thread), arg.c);
+
+                auto f = [=] {
+                    for(std::size_t ic = ic_begin; ic < ic_end; ++ic)
+                    {
+                        thread_reduce_func(ic);
+                    }
+                };
+
+                threads[it] = joinable_thread(f);
+            }
+
+            return (0.0f);
+        };
+
+        float Run(const device::BaseArgument* p_arg,
+                  const StreamConfig& /*stream_config*/ = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg));
+        };
+    };
+
+    bool IsSupportedArgument(const device::BaseArgument* p_arg) override
+    {
+        (void)p_arg;
+
+        return (true);
+    };
+
+    std::unique_ptr<device::BaseArgument>
+    MakeArgumentPointer(const std::array<index_t, 4> xyLengths,
+                        const std::array<index_t, 4> xStrides,
+                        const std::array<index_t, 4> yStrides,
+                        const std::array<index_t, 1> bnScaleBiasMeanVarLengths,
+                        const std::array<index_t, 1> bnScaleBiasMeanVarStrides,
+                        const void* p_x,
+                        const void* bnScale,
+                        const void* bnBias,
+                        double epsilon,
+                        const void* estimatedMean,
+                        const void* estimatedVariance,
+                        void* p_y) override
+    {
+        return std::make_unique<Argument>(xyLengths,
+                                          xStrides,
+                                          yStrides,
+                                          bnScaleBiasMeanVarLengths,
+                                          bnScaleBiasMeanVarStrides,
+                                          static_cast<const InOutDataType*>(p_x),
+                                          static_cast<const AccDataType*>(bnScale),
+                                          static_cast<const AccDataType*>(bnBias),
+                                          epsilon,
+                                          static_cast<const AccDataType*>(estimatedMean),
+                                          static_cast<const AccDataType*>(estimatedVariance),
+                                          static_cast<InOutDataType*>(p_y));
+    };
+
+    std::unique_ptr<device::BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>();
+    };
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "Reference_BatchNorm_Forward_NHWC_C<" << std::endl;
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace host
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/reference_tensor_operation/cpu/reference_sparse_embedding3_forward_layernorm.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <sstream>
+#include <vector>
+#include <algorithm>
+
+#include "ck/tensor_operation/gpu/device/device_base.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace host {
+
+template <typename EmbType,
+          typename IndexType,
+          typename GammaDataType,
+          typename BetaDataType,
+          typename AccDataType,
+          typename OutType>
+struct ReferenceSparseEmbedding3ForwardLayernorm : public device::BaseOperator
+{
+    struct Argument : public device::BaseArgument
+    {
+        Argument(Tensor<OutType>& output,
+                 const Tensor<EmbType>& emb_a,
+                 const Tensor<EmbType>& emb_b,
+                 const Tensor<EmbType>& emb_c,
+                 const Tensor<IndexType>& index_a,
+                 const Tensor<IndexType>& index_b,
+                 const Tensor<IndexType>& index_c,
+                 const Tensor<GammaDataType>& gamma,
+                 const Tensor<BetaDataType>& beta,
+                 ck::index_t NumRows,
+                 ck::index_t EmbeddingDim,
+                 ck::index_t IndexLength,
+                 AccDataType epsilon)
+            : output_(output),
+              emb_a_(emb_a),
+              emb_b_(emb_b),
+              emb_c_(emb_c),
+              index_a_(index_a),
+              index_b_(index_b),
+              index_c_(index_c),
+              gamma_(gamma),
+              beta_(beta),
+              NumRows_(NumRows),
+              EmbeddingDim_(EmbeddingDim),
+              IndexLength_(IndexLength),
+              epsilon_(epsilon)
+        {
+        }
+        Tensor<OutType>& output_;
+        const Tensor<EmbType> emb_a_;
+        const Tensor<EmbType> emb_b_;
+        const Tensor<EmbType> emb_c_;
+        const Tensor<IndexType> index_a_;
+        const Tensor<IndexType> index_b_;
+        const Tensor<IndexType> index_c_;
+        const Tensor<GammaDataType> gamma_;
+        const Tensor<BetaDataType> beta_;
+        ck::index_t NumRows_;
+        ck::index_t EmbeddingDim_;
+        ck::index_t IndexLength_;
+        AccDataType epsilon_;
+    };
+
+    // Invoker
+    struct Invoker : public device::BaseInvoker
+    {
+        float Run(const Argument& arg)
+        {
+            ck::index_t D = arg.EmbeddingDim_;
+            ck::index_t L = arg.IndexLength_;
+            ck::index_t E = arg.NumRows_;
+
+            Tensor<AccDataType> accumulator({L, D});
+
+            Tensor<AccDataType> mean({L});
+            Tensor<AccDataType> var({L});
+
+            accumulator.SetZero();
+
+            auto f_emb_per_row = [&](auto idx) {
+                IndexType idx_a = arg.index_a_(idx);
+                IndexType idx_b = arg.index_b_(idx);
+                IndexType idx_c = arg.index_c_(idx);
+
+                if(!((idx_a < E) && (idx_b < E) && (idx_c < E)))
+                {
+                    throw(std::runtime_error("wrong! out of range"));
+                }
+
+                for(auto d = 0; d < D; d++)
+                {
+                    auto v_a = ck::type_convert<AccDataType>(arg.emb_a_(idx_a, d));
+                    auto v_b = ck::type_convert<AccDataType>(arg.emb_b_(idx_b, d));
+                    auto v_c = ck::type_convert<AccDataType>(arg.emb_c_(idx_c, d));
+
+                    accumulator(idx, d) += v_a + v_b + v_c;
+                }
+            };
+            make_ParallelTensorFunctor(f_emb_per_row, L)(std::thread::hardware_concurrency());
+
+            // layernorm
+            for(auto idx = 0; idx < L; ++idx)
+            {
+                mean(idx) = 0;
+                var(idx)  = 0;
+
+                for(auto d = 0; d < D; ++d)
+                {
+                    auto x_val = accumulator(idx, d);
+                    mean(idx) += x_val;
+                    var(idx) += x_val * x_val;
+                }
+
+                mean(idx) = mean(idx) / D;
+                var(idx)  = (var(idx) / D) - (mean(idx) * mean(idx));
+            }
+
+            for(auto idx = 0; idx < L; ++idx)
+            {
+                for(auto d = 0; d < D; ++d)
+                {
+                    auto x_val          = accumulator(idx, d);
+                    auto y_val          = (x_val - mean(idx)) / sqrt(var(idx) + arg.epsilon_);
+                    y_val               = (y_val * arg.gamma_(d)) + arg.beta_(d);
+                    arg.output_(idx, d) = ck::type_convert<OutType>(y_val);
+                }
+            }
+            return 0;
+        }
+
+        float Run(const device::BaseArgument* p_arg,
+                  const StreamConfig& /* stream_config */ = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg));
+        }
+    };
+
+    static constexpr bool IsValidCompilationParameter()
+    {
+        // TODO: properly implement this check
+        return true;
+    }
+
+    bool IsSupportedArgument(const device::BaseArgument*) override { return true; }
+
+    static auto MakeArgument(Tensor<OutType>& output,
+                             const Tensor<EmbType>& emb_a,
+                             const Tensor<EmbType>& emb_b,
+                             const Tensor<EmbType>& emb_c,
+                             const Tensor<IndexType>& index_a,
+                             const Tensor<IndexType>& index_b,
+                             const Tensor<IndexType>& index_c,
+                             const Tensor<GammaDataType>& gamma,
+                             const Tensor<BetaDataType>& beta,
+                             ck::index_t NumRows,
+                             ck::index_t EmbeddingDim,
+                             ck::index_t IndexLength,
+                             AccDataType epsilon)
+    {
+        return Argument(output,
+                        emb_a,
+                        emb_b,
+                        emb_c,
+                        index_a,
+                        index_b,
+                        index_c,
+                        gamma,
+                        beta,
+                        NumRows,
+                        EmbeddingDim,
+                        IndexLength,
+                        epsilon);
+    }
+
+    static auto MakeInvoker() { return Invoker{}; }
+
+    virtual std::unique_ptr<device::BaseInvoker> MakeInvokerPointer()
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "ReferenceSparseEmbedding3ForwardLayernorm"
+            << std::endl;
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace host
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_add_relu_gemm_add.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_batched_gemm_multiple_d_gemm_multiple_d.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
+
+using CDE0ElementOp = ck::tensor_operation::element_wise::AddRelu;
+using CDE1ElementOp = ck::tensor_operation::element_wise::Add;
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+void add_device_batched_gemm_add_relu_gemm_add_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gno_gmo_instance(
+    std::vector<std::unique_ptr<DeviceBatchedGemmMultipleDGemmMultipleD<Row,
+                                                                        Col,
+                                                                        ck::Tuple<Row>,
+                                                                        Row,
+                                                                        ck::Tuple<Row>,
+                                                                        Row,
+                                                                        F16,
+                                                                        F16,
+                                                                        ck::Tuple<F16>,
+                                                                        F16,
+                                                                        ck::Tuple<F16>,
+                                                                        F16,
+                                                                        PassThrough,
+                                                                        PassThrough,
+                                                                        CDE0ElementOp,
+                                                                        PassThrough,
+                                                                        CDE1ElementOp>>>&
+        instances);
+
+void add_device_batched_gemm_add_relu_gemm_add_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gon_gmo_instance(
+    std::vector<std::unique_ptr<DeviceBatchedGemmMultipleDGemmMultipleD<Row,
+                                                                        Col,
+                                                                        ck::Tuple<Row>,
+                                                                        Col,
+                                                                        ck::Tuple<Row>,
+                                                                        Row,
+                                                                        F16,
+                                                                        F16,
+                                                                        ck::Tuple<F16>,
+                                                                        F16,
+                                                                        ck::Tuple<F16>,
+                                                                        F16,
+                                                                        PassThrough,
+                                                                        PassThrough,
+                                                                        CDE0ElementOp,
+                                                                        PassThrough,
+                                                                        CDE1ElementOp>>>&
+        instances);
+
+template <typename A0Layout,
+          typename B0Layout,
+          typename D0sLayout,
+          typename B1Layout,
+          typename D1sLayout,
+          typename E1Layout,
+          typename A0DataType,
+          typename B0DataType,
+          typename D0sDataType,
+          typename B1DataType,
+          typename D1sDataType,
+          typename E1DataType>
+struct DeviceOperationInstanceFactory<
+    ck::tensor_operation::device::DeviceBatchedGemmMultipleDGemmMultipleD<A0Layout,
+                                                                          B0Layout,
+                                                                          D0sLayout,
+                                                                          B1Layout,
+                                                                          D1sLayout,
+                                                                          E1Layout,
+                                                                          A0DataType,
+                                                                          B0DataType,
+                                                                          D0sDataType,
+                                                                          B1DataType,
+                                                                          D1sDataType,
+                                                                          E1DataType,
+                                                                          PassThrough,
+                                                                          PassThrough,
+                                                                          CDE0ElementOp,
+                                                                          PassThrough,
+                                                                          CDE1ElementOp>>
+{
+    using DeviceOp = DeviceBatchedGemmMultipleDGemmMultipleD<A0Layout,
+                                                             B0Layout,
+                                                             D0sLayout,
+                                                             B1Layout,
+                                                             D1sLayout,
+                                                             E1Layout,
+                                                             A0DataType,
+                                                             B0DataType,
+                                                             D0sDataType,
+                                                             B1DataType,
+                                                             D1sDataType,
+                                                             E1DataType,
+                                                             PassThrough,
+                                                             PassThrough,
+                                                             CDE0ElementOp,
+                                                             PassThrough,
+                                                             CDE1ElementOp>;
+
+    static auto GetInstances()
+    {
+        std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
+
+        if constexpr(is_same_v<A0DataType, half_t> && is_same_v<B0DataType, half_t> &&
+                     is_same_v<B1DataType, half_t> && is_same_v<E1DataType, half_t>)
+        {
+            if constexpr(is_same_v<A0Layout, Row> && is_same_v<B0Layout, Col> &&
+                         is_same_v<B1Layout, Row> && is_same_v<E1Layout, Row>)
+            {
+                add_device_batched_gemm_add_relu_gemm_add_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gno_gmo_instance(
+                    op_ptrs);
+            }
+            else if constexpr(is_same_v<A0Layout, Row> && is_same_v<B0Layout, Col> &&
+                              is_same_v<B1Layout, Col> && is_same_v<E1Layout, Row>)
+            {
+                add_device_batched_gemm_add_relu_gemm_add_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gon_gmo_instance(
+                    op_ptrs);
+            }
+        }
+        return op_ptrs;
+    }
+};
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_gemm.hpp

@@ -32,6 +32,20 @@ void add_device_batched_gemm_gemm_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gno_gmo_i
                                                       PassThrough,
                                                       PassThrough>>>& instances);
 
+void add_device_batched_gemm_gemm_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gon_gmo_instance(
+    std::vector<std::unique_ptr<DeviceBatchedGemmGemm<Row,
+                                                      Col,
+                                                      Col,
+                                                      Row,
+                                                      F16,
+                                                      F16,
+                                                      F16,
+                                                      F16,
+                                                      PassThrough,
+                                                      PassThrough,
+                                                      PassThrough,
+                                                      PassThrough,
+                                                      PassThrough>>>& instances);
 template <typename ALayout,
           typename B0Layout,
           typename B1Layout,
@@ -82,6 +96,12 @@ struct DeviceOperationInstanceFactory<
                 add_device_batched_gemm_gemm_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gno_gmo_instance(
                     op_ptrs);
             }
+            else if constexpr(is_same_v<ALayout, Row> && is_same_v<B0Layout, Col> &&
+                              is_same_v<B1Layout, Col> && is_same_v<CLayout, Row>)
+            {
+                add_device_batched_gemm_gemm_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gon_gmo_instance(
+                    op_ptrs);
+            }
         }
         return op_ptrs;
     }

library/include/ck/library/tensor_operation_instance/gpu/device_elementwise_instance.hpp

@@ -17,9 +17,12 @@ namespace tensor_operation {
 namespace device {
 namespace instance {
 
-using Normalize = ck::tensor_operation::element_wise::Normalize;
-using DeviceNormalizeFromMeanMeanSquarePtr =
-    ck::tensor_operation::device::DeviceElementwisePtr<5, 1, 2, Normalize>;
+using Normalize                            = ck::tensor_operation::element_wise::Normalize;
+using DeviceNormalizeFromMeanMeanSquarePtr = ck::tensor_operation::device::DeviceElementwiseBasePtr<
+    Tuple<half_t, float, float, half_t, half_t>,
+    Tuple<half_t>,
+    Normalize,
+    2>;
 
 void add_device_normalize_from_mean_squaremean_f16_f32_f32_f16_f16_instances(
     std::vector<DeviceNormalizeFromMeanMeanSquarePtr>& instances);

library/include/ck/library/tensor_operation_instance/gpu/gemm.hpp

@@ -4,6 +4,8 @@
 #pragma once
 
 #include <cstdlib>
+#include <memory>
+#include <vector>
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"

library/include/ck/library/tensor_operation_instance/gpu/layernorm.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_normalization.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+void add_device_layernorm_f16_rank2_instances(
+    std::vector<DeviceLayernormPtr<F16, F16, F16, F32, F16, PassThrough, 2, 1>>&);
+
+void add_device_layernorm_f16_rank4_instances(
+    std::vector<DeviceLayernormPtr<F16, F16, F16, F32, F16, PassThrough, 4, 3>>&);
+
+void add_device_layernorm_f32_rank2_instances(
+    std::vector<DeviceLayernormPtr<F32, F32, F32, F32, F32, PassThrough, 2, 1>>&);
+
+void add_device_layernorm_f32_rank4_instances(
+    std::vector<DeviceLayernormPtr<F32, F32, F32, F32, F32, PassThrough, 4, 3>>&);
+
+template <typename XDataType,
+          typename GammaDataType,
+          typename BetaDataType,
+          typename YDataType,
+          index_t Rank,
+          index_t NumReduceDim>
+struct DeviceOperationInstanceFactory<
+    ck::tensor_operation::device::DeviceLayernorm<XDataType,
+                                                  GammaDataType,
+                                                  BetaDataType,
+                                                  F32,
+                                                  YDataType,
+                                                  ck::tensor_operation::element_wise::PassThrough,
+                                                  Rank,
+                                                  NumReduceDim>>
+{
+    using DeviceOp = DeviceLayernorm<XDataType,
+                                     GammaDataType,
+                                     BetaDataType,
+                                     F32,
+                                     YDataType,
+                                     ck::tensor_operation::element_wise::PassThrough,
+                                     Rank,
+                                     NumReduceDim>;
+
+    static auto GetInstances()
+    {
+        std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
+
+        if constexpr(is_same_v<XDataType, F16> && is_same_v<GammaDataType, F16> &&
+                     is_same_v<BetaDataType, F16> && is_same_v<YDataType, F16>)
+        {
+            if constexpr(Rank == 2 && NumReduceDim == 1)
+                add_device_layernorm_f16_rank2_instances(op_ptrs);
+            else if constexpr(Rank == 4 && NumReduceDim == 3)
+                add_device_layernorm_f16_rank4_instances(op_ptrs);
+        }
+        else if constexpr(is_same_v<XDataType, F32> && is_same_v<GammaDataType, F32> &&
+                          is_same_v<BetaDataType, F32> && is_same_v<YDataType, F32>)
+        {
+            if constexpr(Rank == 2 && NumReduceDim == 1)
+                add_device_layernorm_f32_rank2_instances(op_ptrs);
+            else if constexpr(Rank == 4 && NumReduceDim == 3)
+                add_device_layernorm_f32_rank4_instances(op_ptrs);
+        }
+
+        return op_ptrs;
+    }
+};
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/tensor_operation_instance/gpu/softmax.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <type_traits>
+#include <vector>
+
+#include "ck/ck.hpp"
+#include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_softmax.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#include "ck/utility/data_type.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+using F16         = ck::half_t;
+using F32         = float;
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+
+void add_device_softmax_f16_f16_rank3_instances(
+    std::vector<DeviceSoftmaxPtr<F16, F32, F16, PassThrough, PassThrough, 3>>&);
+void add_device_softmax_f16_f16_rank4_instances(
+    std::vector<DeviceSoftmaxPtr<F16, F32, F16, PassThrough, PassThrough, 4>>&);
+
+void add_device_softmax_f32_f32_rank3_instances(
+    std::vector<DeviceSoftmaxPtr<F32, F32, F32, PassThrough, PassThrough, 3>>&);
+void add_device_softmax_f32_f32_rank4_instances(
+    std::vector<DeviceSoftmaxPtr<F32, F32, F32, PassThrough, PassThrough, 4>>&);
+
+template <typename InDataType, typename AccDataType, typename OutDataType, index_t Rank>
+struct DeviceOperationInstanceFactory<
+    ck::tensor_operation::device::
+        DeviceSoftmax<InDataType, AccDataType, OutDataType, PassThrough, PassThrough, Rank>>
+{
+    using DeviceOp =
+        DeviceSoftmax<InDataType, AccDataType, OutDataType, PassThrough, PassThrough, Rank>;
+
+    static auto GetInstances()
+    {
+        std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
+
+        if constexpr(std::is_same_v<InDataType, F16> && std::is_same_v<AccDataType, F32> &&
+                     std::is_same_v<OutDataType, F16>)
+        {
+            if constexpr(Rank == 3)
+                add_device_softmax_f16_f16_rank3_instances(op_ptrs);
+            else if constexpr(Rank == 4)
+                add_device_softmax_f16_f16_rank4_instances(op_ptrs);
+        }
+        else if constexpr(std::is_same_v<InDataType, F32> && std::is_same_v<AccDataType, F32> &&
+                          std::is_same_v<OutDataType, F32>)
+        {
+            if constexpr(Rank == 3)
+                add_device_softmax_f32_f32_rank3_instances(op_ptrs);
+            else if constexpr(Rank == 4)
+                add_device_softmax_f32_f32_rank4_instances(op_ptrs);
+        }
+
+        return op_ptrs;
+    }
+};
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/utility/check_err.hpp

@@ -15,6 +15,7 @@
 
 #include "ck/ck.hpp"
 #include "ck/utility/data_type.hpp"
+#include "ck/utility/type.hpp"
 #include "ck/host_utility/io.hpp"
 
 namespace ck {
@@ -150,7 +151,12 @@ check_err(const std::vector<T>& out,
 }
 
 template <typename T>
-typename std::enable_if<std::is_integral<T>::value && !std::is_same<T, bhalf_t>::value, bool>::type
+std::enable_if_t<(std::is_integral_v<T> && !std::is_same_v<T, bhalf_t>)
+#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
+                     || std::is_same_v<T, int4_t>
+#endif
+                 ,
+                 bool>
 check_err(const std::vector<T>& out,
           const std::vector<T>& ref,
           const std::string& msg = "Error: Incorrect results!",
@@ -159,7 +165,7 @@ check_err(const std::vector<T>& out,
 {
     if(out.size() != ref.size())
     {
-        std::cout << msg << " out.size() != ref.size(), :" << out.size() << " != " << ref.size()
+        std::cerr << msg << " out.size() != ref.size(), :" << out.size() << " != " << ref.size()
                   << std::endl;
         return false;
     }
@@ -180,8 +186,7 @@ check_err(const std::vector<T>& out,
             err_count++;
             if(err_count < 5)
             {
-                std::cout << msg << " out[" << i << "] != ref[" << i
-                          << "]: " << static_cast<int>(out[i]) << " != " << static_cast<int>(ref[i])
+                std::cerr << msg << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r
                           << std::endl;
             }
             res = false;

library/include/ck/library/utility/convolution_host_tensor_descriptor_helper.hpp

@@ -7,6 +7,7 @@
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 
 #include "ck/library/utility/convolution_parameter.hpp"
+#include "ck/library/utility/host_tensor.hpp"
 
 namespace ck {
 namespace utils {

library/include/ck/library/utility/convolution_parameter.hpp

@@ -49,30 +49,47 @@ struct ConvParam
 
     std::size_t GetFlops() const;
 
-    template <typename InDataType, typename WeiDataType, typename OutDataType>
-    std::size_t GetByte() const
+    template <typename InDataType>
+    std::size_t GetInputByte() const
     {
         // sizeof(InDataType) * (G * N * C * <input spatial lengths product>) +
+        return sizeof(InDataType) *
+               (G_ * N_ * C_ *
+                std::accumulate(std::begin(input_spatial_lengths_),
+                                std::begin(input_spatial_lengths_) + num_dim_spatial_,
+                                static_cast<std::size_t>(1),
+                                std::multiplies<std::size_t>()));
+    }
+
+    template <typename WeiDataType>
+    std::size_t GetWeightByte() const
+    {
         // sizeof(WeiDataType) * (G * K * C * <filter spatial lengths product>) +
+        return sizeof(WeiDataType) *
+               (G_ * K_ * C_ *
+                std::accumulate(std::begin(filter_spatial_lengths_),
+                                std::begin(filter_spatial_lengths_) + num_dim_spatial_,
+                                static_cast<std::size_t>(1),
+                                std::multiplies<std::size_t>()));
+    }
+
+    template <typename OutDataType>
+    std::size_t GetOutputByte() const
+    {
         // sizeof(OutDataType) * (G * N * K * <output spatial lengths product>);
-        return sizeof(InDataType) *
-                   (G_ * N_ * C_ *
-                    std::accumulate(std::begin(input_spatial_lengths_),
-                                    std::begin(input_spatial_lengths_) + num_dim_spatial_,
-                                    static_cast<std::size_t>(1),
-                                    std::multiplies<std::size_t>())) +
-               sizeof(WeiDataType) *
-                   (G_ * K_ * C_ *
-                    std::accumulate(std::begin(filter_spatial_lengths_),
-                                    std::begin(filter_spatial_lengths_) + num_dim_spatial_,
-                                    static_cast<std::size_t>(1),
-                                    std::multiplies<std::size_t>())) +
-               sizeof(OutDataType) * (G_ * N_ * K_ *
+        return sizeof(OutDataType) * (G_ * N_ * K_ *
                                       std::accumulate(std::begin(output_spatial_lengths_),
                                                       std::end(output_spatial_lengths_),
                                                       static_cast<std::size_t>(1),
                                                       std::multiplies<std::size_t>()));
     }
+
+    template <typename InDataType, typename WeiDataType, typename OutDataType>
+    std::size_t GetByte() const
+    {
+        return GetInputByte<InDataType>() + GetWeightByte<WeiDataType>() +
+               GetOutputByte<OutDataType>();
+    }
 };
 
 std::string get_conv_param_parser_helper_msg();