Merge remote-tracking branch 'origin/develop' into ck_tile/layernorm_fusion

library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_dynamic_op/xdl/device_grouped_conv3d_fwd_xdl_dynamic_op_ndhwgc_gkzyxc_ndhwgk_f16_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_dynamic_op_instance.hpp"
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+void add_device_grouped_conv3d_fwd_xdl_dynamic_op_ndhwgc_gkzyxc_ndhwgk_f16_instances(
+    std::vector<std::unique_ptr<DeviceGroupedConvFwdMultipleABD<3,
+                                                                NDHWGC,
+                                                                GKZYXC,
+                                                                ck::Tuple<>,
+                                                                NDHWGK,
+                                                                F16,
+                                                                F16,
+                                                                ck::Tuple<>,
+                                                                F16,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                DynamicUnaryOp>>>& instances)
+{
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_f16_instances<3,
+                                                             NDHWGC,
+                                                             GKZYXC,
+                                                             Tuple<>,
+                                                             NDHWGK,
+                                                             ConvFwdDefault>{});
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_f16_instances<3,
+                                                             NDHWGC,
+                                                             GKZYXC,
+                                                             Tuple<>,
+                                                             NDHWGK,
+                                                             ConvFwd1x1P0>{});
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_f16_instances<3,
+                                                             NDHWGC,
+                                                             GKZYXC,
+                                                             Tuple<>,
+                                                             NDHWGK,
+                                                             ConvFwd1x1S1P0>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_dynamic_op/xdl/device_grouped_conv3d_fwd_xdl_dynamic_op_ndhwgc_gkzyxc_ndhwgk_f32_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_dynamic_op_instance.hpp"
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+void add_device_grouped_conv3d_fwd_xdl_dynamic_op_ndhwgc_gkzyxc_ndhwgk_f32_instances(
+    std::vector<std::unique_ptr<DeviceGroupedConvFwdMultipleABD<3,
+                                                                NDHWGC,
+                                                                GKZYXC,
+                                                                ck::Tuple<>,
+                                                                NDHWGK,
+                                                                F32,
+                                                                F32,
+                                                                ck::Tuple<>,
+                                                                F32,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                DynamicUnaryOp>>>& instances)
+{
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_f32_instances<3,
+                                                             NDHWGC,
+                                                             GKZYXC,
+                                                             Tuple<>,
+                                                             NDHWGK,
+                                                             ConvFwdDefault>{});
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_f32_instances<3,
+                                                             NDHWGC,
+                                                             GKZYXC,
+                                                             Tuple<>,
+                                                             NDHWGK,
+                                                             ConvFwd1x1P0>{});
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_f32_instances<3,
+                                                             NDHWGC,
+                                                             GKZYXC,
+                                                             Tuple<>,
+                                                             NDHWGK,
+                                                             ConvFwd1x1S1P0>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/grouped_conv3d_fwd_dynamic_op/xdl/device_grouped_conv3d_fwd_xdl_dynamic_op_ndhwgc_gkzyxc_ndhwgk_int8_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_dynamic_op_instance.hpp"
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+void add_device_grouped_conv3d_fwd_xdl_dynamic_op_ndhwgc_gkzyxc_ndhwgk_int8_instances(
+    std::vector<std::unique_ptr<DeviceGroupedConvFwdMultipleABD<3,
+                                                                NDHWGC,
+                                                                GKZYXC,
+                                                                ck::Tuple<>,
+                                                                NDHWGK,
+                                                                int8_t,
+                                                                int8_t,
+                                                                ck::Tuple<>,
+                                                                int8_t,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                DynamicUnaryOp>>>& instances)
+{
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_int8_instances<3,
+                                                              NDHWGC,
+                                                              GKZYXC,
+                                                              Tuple<>,
+                                                              NDHWGK,
+                                                              ConvFwdDefault>{});
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_int8_instances<3,
+                                                              NDHWGC,
+                                                              GKZYXC,
+                                                              Tuple<>,
+                                                              NDHWGK,
+                                                              ConvFwd1x1P0>{});
+    add_device_operation_instances(
+        instances,
+        device_grouped_conv_fwd_xdl_dynamic_op_int8_instances<3,
+                                                              NDHWGC,
+                                                              GKZYXC,
+                                                              Tuple<>,
+                                                              NDHWGK,
+                                                              ConvFwd1x1S1P0>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

profiler/include/profiler/profile_gemm_multiply_multiply_impl.hpp

@@ -271,10 +271,12 @@ bool profile_gemm_multiply_multiply_impl(int do_verification,
                           << " TFlops, " << gb_per_sec << " GB/s, " << op_name << ", KBatch "
                           << kbatch_curr << std::endl;
 
-#if defined CK_ENABLE_FP8
+#if defined CK_ENABLE_FP8 || defined CK_ENABLE_INT8
                 // set softer tolerances for fp8
-                if constexpr(is_same_v<ADataType, f8_t> || is_same_v<BDataType, f8_t> ||
-                             is_same_v<EDataType, f8_t>)
+                if constexpr((is_same_v<ADataType, f8_t> || is_same_v<BDataType, f8_t> ||
+                              is_same_v<EDataType, f8_t>) ||
+                             (is_same_v<ADataType, int8_t> || is_same_v<BDataType, int8_t> ||
+                              is_same_v<EDataType, int8_t>))
                 {
                     std::string msg = "Error: Incorrect results!";
                     double rtol     = 1e-1;
@@ -286,7 +288,7 @@ bool profile_gemm_multiply_multiply_impl(int do_verification,
                 {
 #endif
                     pass = pass & ck::utils::check_err(e_m_n_device_result, e_m_n_host_result);
-#if defined CK_ENABLE_FP8
+#if defined CK_ENABLE_FP8 || defined CK_ENABLE_INT8
                 }
 #endif
 

profiler/include/profiler/profile_pool3d_fwd_impl.hpp

@@ -102,11 +102,22 @@ bool profile_pool3d_fwd_impl(PoolFwdInputParams& in_params, PoolFwdKernelParams&
     Tensor<IndexDataType> out_indices_n_c_do_ho_wo_device(
         f_host_tensor_descriptor(N, C, Do, Ho, Wo));
 
+    constexpr int inDataRangeTensor1{1};
+    constexpr int inDataRangeTensor2{5};
+    constexpr double inDataRangeTensor3{0.5};
+
     switch(in_params.init_method)
     {
-    case 0: in_n_c_di_hi_wi.GenerateTensorValue(GeneratorTensor_1<InDataType>{}); break;
-    case 1: in_n_c_di_hi_wi.GenerateTensorValue(GeneratorTensor_2<InDataType>{-5, 5}); break;
-    default: in_n_c_di_hi_wi.GenerateTensorValue(GeneratorTensor_3<InDataType>{-0.5, 0.5});
+    case 0:
+        in_n_c_di_hi_wi.GenerateTensorValue(GeneratorTensor_1<InDataType>{inDataRangeTensor1});
+        break;
+    case 1:
+        in_n_c_di_hi_wi.GenerateTensorValue(
+            GeneratorTensor_2<InDataType>{-inDataRangeTensor2, inDataRangeTensor2});
+        break;
+    default:
+        in_n_c_di_hi_wi.GenerateTensorValue(
+            GeneratorTensor_3<InDataType>{-inDataRangeTensor3, inDataRangeTensor3});
     }
 
     DeviceMem in_device_buf(sizeof(InDataType) * in_n_c_di_hi_wi.mDesc.GetElementSpaceSize());
@@ -229,12 +240,25 @@ bool profile_pool3d_fwd_impl(PoolFwdInputParams& in_params, PoolFwdKernelParams&
         {
             out_device_buf.FromDevice(out_n_c_do_ho_wo_device.mData.data());
 
-            auto tolerance = 1e-3;
-            bool pass      = ck::utils::check_err(out_n_c_do_ho_wo_device.mData,
+            auto absolute_error_threshold = 1.0;
+            switch(in_params.init_method)
+            {
+            case 0: absolute_error_threshold = static_cast<double>(inDataRangeTensor1); break;
+            case 1: absolute_error_threshold = static_cast<double>(inDataRangeTensor2); break;
+            default: absolute_error_threshold = inDataRangeTensor3;
+            }
+
+            absolute_error_threshold =
+                ck::utils::get_absolute_threshold<ComputeDataType, OutDataType>(
+                    absolute_error_threshold);
+            auto relative_error_threshold =
+                ck::utils::get_relative_threshold<ComputeDataType, OutDataType>();
+
+            bool pass = ck::utils::check_err(out_n_c_do_ho_wo_device.mData,
                                              out_n_c_do_ho_wo_host.mData,
                                              "Error: Incorrect results",
-                                             tolerance,
-                                             tolerance);
+                                             relative_error_threshold,
+                                             absolute_error_threshold);
 
             if constexpr(OutputIndex)
             {

profiler/src/profile_gemm_multiply_multiply.cpp

@@ -27,6 +27,7 @@ enum struct GemmDataType
     F16_F8_F16,     // 5
     F16_F16_F16_F8, // 6
     F8_F8_BF16,     // 7
+    INT8_INT8_BF16, // 8
 };
 
 #define OP_NAME "gemm_multiply_multiply"
@@ -39,7 +40,7 @@ int profile_gemm_multiply_multiply(int argc, char* argv[])
         printf("arg1: tensor operation (" OP_NAME ": " OP_DESC ")\n");
         printf("arg2: data type (0: fp32; 1: fp16; 2: bf16; 3: int8; 4: f8@f16; 5: f16@f8; 6: "
                "f16->f8; 7: f8->bf16, "
-               "comp f8)\n");
+               "comp f8; 8: int8->bf16)\n");
         printf("arg3: matrix layout (0: A[m, k] * B[k, n] = C[m, n];\n");
         printf("                     1: A[m, k] * B[n, k] = C[m, n];\n");
         printf("                     2: A[k, m] * B[k, n] = C[m, n];\n");
@@ -89,6 +90,8 @@ int profile_gemm_multiply_multiply(int argc, char* argv[])
     using F32  = float;
     using BF16 = ck::bhalf_t;
     using F8   = ck::f8_t;
+    using I8   = int8_t;
+    using I32  = int;
 
     using Row = ck::tensor_layout::gemm::RowMajor;
     using Col = ck::tensor_layout::gemm::ColumnMajor;
@@ -162,6 +165,11 @@ int profile_gemm_multiply_multiply(int argc, char* argv[])
         return profile(
             F8{}, F8{}, F8{}, F32{}, F32{}, F32{}, BF16{}, Row{}, Col{}, Row{}, Col{}, Row{});
     }
+    else if(data_type == GemmDataType::INT8_INT8_BF16 && layout == GemmMatrixLayout::MK_NK_MN)
+    {
+        return profile(
+            I8{}, I8{}, I8{}, I32{}, F32{}, F32{}, BF16{}, Row{}, Col{}, Row{}, Col{}, Row{});
+    }
     else
     {
         std::cout << "this data_type & layout is not implemented" << std::endl;

profiler/src/profile_grouped_conv_bwd_weight.cpp

@@ -182,6 +182,10 @@ int profile_grouped_conv_bwd_weight(int argc, char* argv[])
             // fp32 atomic add is used for weight tensor in bf16 kernel
             return profile(I2, NHWGC{}, GKYXC{}, NHWGK{}, BF16{}, F32{}, BF16{}, BF16{}, BF16{});
         }
+        if(data_type == ConvDataType::BF16_BF16_BF16)
+        {
+            return profile(I2, NHWGC{}, GKYXC{}, NHWGK{}, BF16{}, BF16{}, BF16{}, BF16{}, BF16{});
+        }
     }
     else if(num_dim_spatial == 2 && layout == ConvLayout::NGCHW_GKYXC_NGKHW)
     {
@@ -210,11 +214,6 @@ int profile_grouped_conv_bwd_weight(int argc, char* argv[])
             // fp32 atomic add is used for weight tensor in bf16 kernel
             return profile(I3, GNDHWC{}, GKZYXC{}, GNDHWK{}, BF16{}, F32{}, BF16{}, BF16{}, BF16{});
         }
-        if(data_type == ConvDataType::BF16_BF16_BF16)
-        {
-            return profile(
-                I3, GNDHWC{}, GKZYXC{}, GNDHWK{}, BF16{}, BF16{}, BF16{}, BF16{}, BF16{});
-        }
         else if(data_type == ConvDataType::I8_I8_I8)
         {
             return profile(

python/ck4inductor/grouped_conv_fwd/gen_instances.py

+# SPDX-License-Identifier: MIT
+# Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+import logging
+import os
+import subprocess
+from dataclasses import replace
+from functools import lru_cache
+from typing import List
+
+from ..util import library_path
+
+from .op import CKGroupedConvFwdOp
+
+log = logging.getLogger(__name__)
+
+
+def _ck_conv_instances_path():
+    conv_instances_path = os.path.join(  # noqa: F821
+        library_path(),
+        "include",
+        "ck",
+        "library",
+        "tensor_operation_instance",
+        "gpu",
+        "grouped_conv_fwd",
+    )
+    if not os.path.exists(conv_instances_path):
+        log.error(
+            "CK library conv instances path %s does not exist", conv_instances_path
+        )
+        return None
+    return conv_instances_path
+
+
+def parse_instances(str_instances: List[str]) -> List[CKGroupedConvFwdOp]:
+    """
+    Parse the lines containing Grouped Convolution Forward template instances
+    into `CKGroupedConvFwdOp` instances
+    """
+
+    def maybe_int(s):
+        try:
+            return int(s)
+        except ValueError:
+            return s
+
+    op_instances = []
+    # TODO: maybe use libclang for parsing C++ code in the future
+    # to avoid this hacky parsing logic below ? :) - copilot
+    for line in str_instances:
+        s_template_args = line.split("DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3")[
+            -1
+        ].strip("<>, ")
+        template_args = []
+        i_current = 0
+        while i_current < len(s_template_args):
+            if s_template_args[i_current] == " ":
+                # skip whitespace
+                i_current += 1
+                continue
+            elif s_template_args[i_current : i_current + 2] == "S<":
+                # parse template S<Index...>
+                i_next = s_template_args.find(">", i_current)
+                template_args.append(
+                    tuple(map(int, s_template_args[i_current + 2 : i_next].split(",")))
+                )
+                i_current = i_next + 2
+            else:
+                # all string attributes must be either type aliases or global constants in C++
+                i_next = s_template_args.find(",", i_current)
+                template_args.append(
+                    maybe_int(
+                        s_template_args[i_current : i_next if i_next != -1 else None]
+                    )
+                )
+                if i_next != -1:
+                    i_current = i_next + 1
+            if i_next == -1:
+                break
+
+        template_args[0] = -1  # n_dim_spatial
+        template_args[3] = tuple()  # ds_layout
+        template_args[9] = tuple()  # ds_element_dtype
+
+        new_instance = CKGroupedConvFwdOp(
+            *template_args,  # type: ignore[arg-type]
+        )
+
+        op_instances.append(new_instance)
+    return op_instances
+
+
+@lru_cache(None)
+def gen_conv_ops_library() -> List[CKGroupedConvFwdOp]:
+    """
+    Parse the Grouped Convolution Forward instances
+    defined in the Composable Kernel library folder.
+    """
+    ck_library_dir = _ck_conv_instances_path()
+    if not ck_library_dir:
+        return []
+
+    grep_result = subprocess.run(
+        [
+            "grep",
+            "-inR",
+            "DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3",
+            ck_library_dir,
+        ],
+        capture_output=True,
+        text=True,
+    )
+
+    op_instances = parse_instances(grep_result.stdout.strip().split("\n"))
+
+    log.debug("ck instances from library: %d", len(op_instances))
+
+    schedulers = [
+        "BlockGemmPipelineScheduler::Intrawave",
+        "BlockGemmPipelineScheduler::Interwave",
+    ]
+    conv_specs = [
+        "ConvolutionForwardSpecialization::Default",
+        "ConvolutionForwardSpecialization::Filter1x1Pad0",
+        "ConvolutionForwardSpecialization::Filter1x1Stride1Pad0",
+        "ConvolutionForwardSpecialization::OddC",
+    ]
+
+    # substitute templated args by looping through their domains
+    substitute_instances = []
+    for instance in op_instances:
+        sub_scheduler = (
+            instance.block_gemm_pipeline_scheduler == "BlkGemmPipeSched"
+        )
+        sub_spec = instance.conv_forward_specialization == "ConvSpec"
+        schedulers_range = (
+            schedulers if sub_scheduler else [instance.block_gemm_pipeline_scheduler]
+        )
+        spec_range = conv_specs if sub_spec else [instance.conv_forward_specialization]
+        for scheduler in schedulers_range:
+            for spec in spec_range:
+                for channels_last in [True, False]:
+                    if channels_last:
+                        a_layout = "NHWGC"
+                        e_layout = "NHWGK"
+                    else:
+                        a_layout = "NGCHW"
+                        e_layout = "NGKHW"
+                    substitute_instances.append(
+                        replace(
+                            instance,
+                            block_gemm_pipeline_scheduler=scheduler,
+                            conv_forward_specialization=spec,
+                            gemm_specialization="GemmSpecialization::MNKPadding",
+                            n_dim_spatial=2,
+                            a_layout=a_layout,
+                            b_layout="GKYXC",
+                            e_layout=e_layout,
+                        )
+                    )
+
+    return substitute_instances
+
+
+if __name__ == "__main__":
+    print(gen_conv_ops_library())

python/ck4inductor/grouped_conv_fwd/op.py

+# SPDX-License-Identifier: MIT
+# Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+from dataclasses import asdict, dataclass
+from typing import Optional, Tuple
+
+
+@dataclass
+class CKGroupedConvFwdOp:
+    n_dim_spatial: int
+    a_layout: str
+    b_layout: str
+    ds_layout: Tuple[str]
+    e_layout: str
+    a_element_dtype: str
+    b_element_dtype: str
+    acc_dtype: str
+    c_shuffle_dtype: str
+    ds_element_dtype: Tuple[str]
+    e_element_dtype: str
+    a_elementwise_op: str
+    b_elementwise_op: str
+    cde_elementwise_op: str
+    conv_forward_specialization: str
+    gemm_specialization: str
+
+    block_size: int
+    m_per_block: int
+    n_per_block: int
+    k_per_block: int
+    ak1: int
+    bk1: int
+    m_per_xdl: int
+    n_per_xdl: int
+    m_xdl_per_wave: int
+    n_xdl_per_wave: int
+    a_block_transfer_thread_cluster_lengths_ak0_m_ak1: Tuple[int, int, int]
+    a_block_transfer_thread_cluster_arrange_order: Tuple[int, int, int]
+    a_block_transfer_src_access_order: Tuple[int, int, int]
+    a_block_transfer_src_vector_dim: int
+    a_block_transfer_src_scalar_per_vector: int
+    a_block_transfer_dst_scalar_per_vector_ak1: int
+    a_block_lds_extra_m: bool
+
+    b_block_transfer_thread_cluster_lengths_bk0_n_bk1: Tuple[int, int, int]
+    b_block_transfer_thread_cluster_arrange_order: Tuple[int, int, int]
+    b_block_transfer_src_access_order: Tuple[int, int, int]
+
+    b_block_transfer_src_vector_dim: int
+    b_block_transfer_src_scalar_per_vector: int
+    b_block_transfer_dst_scalar_per_vector_bk1: int
+    b_block_lds_extra_n: bool
+
+    c_shuffle_m_xdl_per_wave_per_shuffle: int
+    c_shuffle_n_xdl_per_wave_per_shuffle: int
+    cde_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block: Tuple[  # noqa
+        int,
+        int,
+        int,
+        int,
+    ]
+    cde_block_transfer_scalar_per_vector_n_per_block: int
+    block_gemm_pipeline_scheduler: str
+    block_gemm_pipeline_version: str
+
+    a_compute_dtype: Optional[str] = None
+    b_compute_dtype: Optional[str] = None
+
+    def name(self):
+        # cpp alias for template instance
+        return (
+            f"ck_device_grouped_convolution_fwd_multiple_abd_xdl_c_shuffle_v3_"
+            f"{self.key_name()}"
+        )
+
+    def key_name(self):
+        # TBD; must be unique per instance. Intended to use as dict key
+        return "_".join(
+            [
+                "K"
+                + field_name.replace("_", "").lower()
+                + "V"
+                + (
+                    "x".join(map(str, iter(field_value)))
+                    if isinstance(field_value, tuple)
+                    else str(field_value).replace(":", "")
+                )
+                for field_name, field_value in self.dict_items()
+            ]
+        )
+
+    def dict_items(self):
+        return asdict(self).items()

python/ck4inductor/universal_gemm/gen_instances.py

+# SPDX-License-Identifier: MIT
+# Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
 import logging
 import os
 import subprocess
-from dataclasses import fields, replace
+from dataclasses import replace
 from functools import lru_cache, partial
 from typing import List
 

python/ck4inductor/universal_gemm/op.py

+# SPDX-License-Identifier: MIT
+# Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
 from dataclasses import asdict, dataclass
 from typing import Optional, Tuple
 

python/ck4inductor/util.py

+# SPDX-License-Identifier: MIT
+# Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
 import functools
 import os
 
 
 @functools.lru_cache(None)
 def library_path():
-    return os.path.join(os.path.dirname(__file__), 'library')
+    return os.path.join(os.path.dirname(__file__), "library")