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

Merge branch 'develop' of github.com:ROCmSoftwarePlatform/composable_kernel into barkocot/fix-instance-fwd-xdl-log

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

@@ -565,7 +565,7 @@ struct DeviceGroupedConvBwdWeight_Wmma_CShuffle
         Block2CTileMap block_2_ctile_map_;
 
         // for computing batch offset
-        ComputePtrOffsetOfStridedBatch<I0> compute_ptr_offset_of_batch_;
+        ComputePtrOffsetOfStridedBatch<> compute_ptr_offset_of_batch_;
 
         OutElementwiseOperation a_element_op_;
         InElementwiseOperation b_element_op_;
@@ -647,7 +647,7 @@ struct DeviceGroupedConvBwdWeight_Wmma_CShuffle
                     DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock,
                     CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock,
                     remove_reference_t<typename GridwiseGemm::DefaultBlock2CTileMap>,
-                    ComputePtrOffsetOfStridedBatch<I0>,
+                    ComputePtrOffsetOfStridedBatch<>,
                     has_main_loop>;
 
                 using EmptyTuple = Tuple<>;

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

@@ -1197,7 +1197,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle
         Block2CTileMap block_2_ctile_map_;
 
         // for computing batch offset
-        ComputePtrOffsetOfStridedBatch<I0> compute_ptr_offset_of_batch_;
+        ComputePtrOffsetOfStridedBatch<> compute_ptr_offset_of_batch_;
 
         index_t M01_;
         index_t N01_;
@@ -1276,7 +1276,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle
                     remove_reference_t<DeviceOp::BGridDesc_K0_N_K1>,
                     remove_reference_t<DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
                     remove_reference_t<DeviceOp::Block2CTileMap>,
-                    ComputePtrOffsetOfStridedBatch<I0>,
+                    ComputePtrOffsetOfStridedBatch<>,
                     has_main_loop>;
 
                 return launch_and_time_kernel(stream_config,

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

@@ -15,7 +15,7 @@
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/convolution_forward_specialization.hpp"
 #include "ck/tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp"
-#include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_d.hpp"
+#include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/matrix_padder.hpp"
 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_dl_multiple_d.hpp"
@@ -216,18 +216,18 @@ template <index_t NDimSpatial,
           index_t CThreadTransferSrcDstVectorDim,
           index_t CThreadTransferDstScalarPerVector>
 struct DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK
-    : public DeviceGroupedConvFwdMultipleD<NDimSpatial,
-                                           ALayout,
-                                           BLayout,
-                                           DsLayout,
-                                           ELayout,
-                                           ADataType,
-                                           BDataType,
-                                           DsDataType,
-                                           EDataType,
-                                           AElementwiseOperation,
-                                           BElementwiseOperation,
-                                           CDEElementwiseOperation>
+    : public DeviceGroupedConvFwdMultipleABD<NDimSpatial,
+                                             ALayout,
+                                             BLayout,
+                                             DsLayout,
+                                             ELayout,
+                                             ADataType,
+                                             BDataType,
+                                             DsDataType,
+                                             EDataType,
+                                             AElementwiseOperation,
+                                             BElementwiseOperation,
+                                             CDEElementwiseOperation>
 {
     using DeviceOp = DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK;
 
@@ -537,7 +537,7 @@ struct DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK
         DefaultBlock2CTileMap block_2_ctile_map_;
 
         // for computing batch offset
-        ComputePtrOffsetOfStridedBatch<NumDTensor> compute_ptr_offset_of_batch_;
+        ComputePtrOffsetOfStridedBatch<I1, I1, NumDTensor> compute_ptr_offset_of_batch_;
 
         // element-wise op
         AElementwiseOperation a_element_op_;
@@ -601,7 +601,7 @@ struct DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK
                     DeviceOp::DsGridDesc_M0_M10_M11_N0_N10_N11,
                     DeviceOp::CGridDesc_M0_M10_M11_N0_N10_N11,
                     DefaultBlock2CTileMap,
-                    ComputePtrOffsetOfStridedBatch<NumDTensor>,
+                    ComputePtrOffsetOfStridedBatch<I1, I1, NumDTensor>,
                     has_main_loop,
                     has_double_loop>;
 

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

@@ -1090,7 +1090,7 @@ struct DeviceGroupedConvFwdMultipleDMultipleR_Xdl_CShuffle
         auto str = std::stringstream();
 
         // clang-format off
-        str << "DeviceGroupedConvFwdMultipleD_Xdl_CShuffle"
+        str << "DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle"
             << "<"
             << BlockSize << ", "
             << MPerBlock << ", "

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

@@ -15,7 +15,7 @@
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/convolution_forward_specialization.hpp"
 #include "ck/tensor_operation/operator_transform/transform_conv_fwd_to_gemm.hpp"
-#include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_d.hpp"
+#include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/matrix_padder.hpp"
 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_multiple_d_wmma_cshuffle.hpp"
@@ -92,18 +92,18 @@ template <index_t NDimSpatial,
           LoopScheduler LoopSched         = make_default_loop_scheduler(),
           ck::PipelineVersion PipelineVer = ck::PipelineVersion::v1>
 struct DeviceGroupedConvFwdMultipleD_Wmma_CShuffle
-    : public DeviceGroupedConvFwdMultipleD<NDimSpatial,
-                                           ALayout,
-                                           BLayout,
-                                           DsLayout,
-                                           ELayout,
-                                           ADataType,
-                                           BDataType,
-                                           DsDataType,
-                                           EDataType,
-                                           AElementwiseOperation,
-                                           BElementwiseOperation,
-                                           CDEElementwiseOperation>
+    : public DeviceGroupedConvFwdMultipleABD<NDimSpatial,
+                                             ALayout,
+                                             BLayout,
+                                             DsLayout,
+                                             ELayout,
+                                             ADataType,
+                                             BDataType,
+                                             DsDataType,
+                                             EDataType,
+                                             AElementwiseOperation,
+                                             BElementwiseOperation,
+                                             CDEElementwiseOperation>
 {
     using DeviceOp = DeviceGroupedConvFwdMultipleD_Wmma_CShuffle;
 
@@ -428,7 +428,7 @@ struct DeviceGroupedConvFwdMultipleD_Wmma_CShuffle
         typename GridwiseOp::DefaultBlock2CTileMap block_2_etile_map_;
 
         // for computing batch offset
-        ComputePtrOffsetOfStridedBatch<NumDTensor> compute_ptr_offset_of_batch_;
+        ComputePtrOffsetOfStridedBatch<I1, I1, NumDTensor> compute_ptr_offset_of_batch_;
 
         // element-wise op
         AElementwiseOperation a_element_op_;
@@ -485,7 +485,7 @@ struct DeviceGroupedConvFwdMultipleD_Wmma_CShuffle
                     typename GridwiseOp::DsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock,
                     typename GridwiseOp::EGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock,
                     remove_reference_t<typename GridwiseOp::DefaultBlock2CTileMap>,
-                    ComputePtrOffsetOfStridedBatch<NumDTensor>,
+                    ComputePtrOffsetOfStridedBatch<I1, I1, NumDTensor>,
                     has_main_loop>;
 
                 return launch_and_time_kernel(stream_config,

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

@@ -9,8 +9,77 @@ namespace ck {
 namespace tensor_operation {
 namespace device {
 
-template <index_t NumDTensor>
+template <index_t NumATensor = 1, index_t NumBTensor = 1, index_t NumDTensor = 0, typename = void>
 struct ComputePtrOffsetOfStridedBatch
+{
+};
+
+template <index_t NumATensor, index_t NumBTensor, index_t NumDTensor>
+struct ComputePtrOffsetOfStridedBatch<NumATensor,
+                                      NumBTensor,
+                                      NumDTensor,
+                                      ck::enable_if_t<(NumATensor > 1 || NumBTensor > 1)>>
+{
+    ComputePtrOffsetOfStridedBatch() = default;
+
+    ComputePtrOffsetOfStridedBatch(Array<ck::index_t, NumATensor>& BatchStrideAs,
+                                   Array<ck::index_t, NumBTensor>& BatchStrideBs,
+                                   Array<ck::index_t, NumDTensor>& BatchStrideDs,
+                                   index_t BatchStrideE)
+        : BatchStrideA_(BatchStrideAs),
+          BatchStrideB_(BatchStrideBs),
+          BatchStrideDs_(BatchStrideDs),
+          BatchStrideE_(BatchStrideE)
+    {
+    }
+
+    __host__ __device__ constexpr auto GetAsPtrOffset(index_t g_idx) const
+    {
+        Array<long_index_t, NumATensor> as_offset;
+        static_for<0, NumATensor, 1>{}(
+            [&](auto i) { as_offset(i) = g_idx * static_cast<long_index_t>(BatchStrideA_[i]); });
+        return as_offset;
+    }
+
+    __host__ __device__ constexpr auto GetBsPtrOffset(index_t g_idx) const
+    {
+        Array<long_index_t, NumBTensor> bs_offset;
+        static_for<0, NumBTensor, 1>{}(
+            [&](auto i) { bs_offset(i) = g_idx * static_cast<long_index_t>(BatchStrideB_[i]); });
+        return bs_offset;
+    }
+
+    __host__ __device__ constexpr auto GetDsPtrOffset(index_t g_idx) const
+    {
+        Array<long_index_t, NumDTensor> ds_offset;
+        static_for<0, NumDTensor, 1>{}(
+            [&](auto i) { ds_offset(i) = g_idx * static_cast<long_index_t>(BatchStrideDs_[i]); });
+        return ds_offset;
+    }
+
+    [[maybe_unused]] __host__ __device__ constexpr long_index_t GetEPtrOffset(index_t g_idx) const
+    {
+        return g_idx * static_cast<long_index_t>(BatchStrideE_);
+    }
+
+    // alias for kernels without multiple D
+    [[maybe_unused]] __host__ __device__ constexpr long_index_t GetCPtrOffset(index_t g_idx) const
+    {
+        return g_idx * static_cast<long_index_t>(BatchStrideE_);
+    }
+
+    Array<ck::index_t, NumATensor> BatchStrideA_;
+    Array<ck::index_t, NumBTensor> BatchStrideB_;
+    Array<ck::index_t, NumDTensor> BatchStrideDs_;
+    index_t BatchStrideE_;
+    index_t& BatchStrideC_ = BatchStrideE_; // alias for kernels without multiple D
+};
+
+template <index_t NumATensor, index_t NumBTensor, index_t NumDTensor>
+struct ComputePtrOffsetOfStridedBatch<NumATensor,
+                                      NumBTensor,
+                                      NumDTensor,
+                                      ck::enable_if_t<(NumATensor == 1 && NumBTensor == 1)>>
 {
     ComputePtrOffsetOfStridedBatch() = default;
 
@@ -54,13 +123,67 @@ struct ComputePtrOffsetOfStridedBatch
         return g_idx * static_cast<long_index_t>(BatchStrideE_);
     }
 
-    index_t BatchStrideA_;
-    index_t BatchStrideB_;
+    ck::index_t BatchStrideA_;
+    ck::index_t BatchStrideB_;
     Array<ck::index_t, NumDTensor> BatchStrideDs_;
     index_t BatchStrideE_;
     index_t& BatchStrideC_ = BatchStrideE_; // alias for kernels without multiple D
 };
 
+template <bool isTuple, typename Tensors>
+constexpr static auto GetNumABTensors()
+{
+    if constexpr(isTuple)
+    {
+        return Number<Tensors::Size()>{};
+    }
+    else
+    {
+        return Number<1>{};
+    }
+}
+
+template <bool isTuple, typename GridwiseGemm, typename DataType>
+constexpr static auto GetAGridPointer()
+{
+    if constexpr(isTuple)
+    {
+        return typename GridwiseGemm::AsGridPointer{};
+    }
+    else
+    {
+        return Tuple<const DataType*>{};
+    }
+}
+
+template <bool isTuple, typename GridwiseGemm, typename DataType>
+constexpr static auto GetBGridPointer()
+{
+    if constexpr(isTuple)
+    {
+        return typename GridwiseGemm::BsGridPointer{};
+    }
+    else
+    {
+        return Tuple<const DataType*>{};
+    }
+}
+
+template <bool isTuple, typename Id, typename Type>
+constexpr static auto UnpackDataType()
+{
+    if constexpr(isTuple)
+    {
+        // unpack if tuple
+        return tuple_element_t<Id{}, Type>{};
+    }
+    else
+    {
+        // if no, return Type
+        return Type{};
+    }
+}
+
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck

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

@@ -142,19 +142,18 @@ struct DeviceImageToColumnImpl
         decltype(BlockToCTileMap_M00_N0_M01Adapt<MPerBlock, KPerBlock, OutputGridDesc>(
             OutputGridDesc{}))>;
 
-    using GridwiseTensorRearrangeKernel =
-        GridwiseTensorRearrange<InputGridDesc,
-                                InputDataType,
-                                OutputGridDesc,
-                                OutputDataType,
-                                BlockSize,
-                                MPerBlock,
-                                KPerBlock,
-                                ThreadClusterLengths,
-                                ScalarPerVector,
-                                InMemoryDataOperationEnum::Set,
-                                Block2ETileMap,
-                                ComputePtrOffsetOfStridedBatch<I0>>;
+    using GridwiseTensorRearrangeKernel = GridwiseTensorRearrange<InputGridDesc,
+                                                                  InputDataType,
+                                                                  OutputGridDesc,
+                                                                  OutputDataType,
+                                                                  BlockSize,
+                                                                  MPerBlock,
+                                                                  KPerBlock,
+                                                                  ThreadClusterLengths,
+                                                                  ScalarPerVector,
+                                                                  InMemoryDataOperationEnum::Set,
+                                                                  Block2ETileMap,
+                                                                  ComputePtrOffsetOfStridedBatch<>>;
 
     struct Argument : public BaseArgument
     {
@@ -224,7 +223,7 @@ struct DeviceImageToColumnImpl
         InputGridDesc in_grid_desc_m_k_;
         OutputGridDesc out_grid_desc_m_k_;
 
-        ComputePtrOffsetOfStridedBatch<I0> compute_ptr_offset_of_batch_;
+        ComputePtrOffsetOfStridedBatch<> compute_ptr_offset_of_batch_;
     };
 
     struct Invoker : public BaseInvoker
@@ -246,7 +245,7 @@ struct DeviceImageToColumnImpl
                                                         OutputGridDesc,
                                                         OutputDataType,
                                                         Block2ETileMap,
-                                                        ComputePtrOffsetOfStridedBatch<I0>,
+                                                        ComputePtrOffsetOfStridedBatch<>,
                                                         GridwiseTensorRearrangeKernel>;
 
             float elapsed_time = launch_and_time_kernel(stream_config,

include/ck/tensor_operation/gpu/element/binary_element_wise_operation.hpp

@@ -85,10 +85,13 @@ struct Add
 
 struct ScaleAdd
 {
-    __host__ __device__ ScaleAdd(float scale) : scale_(scale) {}
+    __host__ __device__ ScaleAdd(float scale = 1.f) : scale_(scale) {}
 
     template <typename Y, typename X0, typename X1>
-    __host__ __device__ constexpr void operator()(Y& y, const X0& x0, const X1& x1) const;
+    __host__ __device__ constexpr void operator()(Y& y, const X0& x0, const X1& x1) const
+    {
+        y = ck::type_convert<Y>(scale_ * ck::type_convert<float>(x0) + ck::type_convert<float>(x1));
+    }
 
     template <>
     __host__ __device__ void

include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp

@@ -355,8 +355,8 @@ struct UnarySquare
     template <typename T>
     __host__ __device__ void operator()(T& y, const T& x) const
     {
-        static_assert(is_same_v<T, float> || is_same_v<T, double> || is_same_v<T, int32_t> ||
-                          is_same_v<T, int8_t>
+        static_assert(is_same_v<T, float> || is_same_v<T, half_t> || is_same_v<T, double> ||
+                          is_same_v<T, int32_t> || is_same_v<T, int8_t>
 #ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
                           || is_same_v<T, int4_t>
 #endif