add bias_relu, bias fusion

f8b551da · carlushuang · bfa4c686 · f8b551da · f8b551da · f8b551da
Commit f8b551da authored Jun 14, 2022 by carlushuang
8 changed files
--- a/example/cpu_02_conv2d_fwd_bias_relu_add/cpu_conv2d_fwd_bias_relu_add.cpp
+++ b/example/cpu_02_conv2d_fwd_bias_relu_add/cpu_conv2d_fwd_bias_relu_add.cpp
--- a/include/ck/tensor_operation/cpu/device/device_convnd_fwd_bias_activation_add_avx2_nhwc_kyxc_nhwk.hpp
+++ b/include/ck/tensor_operation/cpu/device/device_convnd_fwd_bias_activation_add_avx2_nhwc_kyxc_nhwk.hpp
@@ -37,6 +37,8 @@ template <typename InDataType,
          bool UseALocalBuffer,
          bool UseBLocalBuffer,
          bool UseCLocalBuffer,
+          bool FuseBias,
+          bool FuseAdd,
          bool BiasAlongGemmM>
 struct DeviceConvNDFwdBiasActivationAddAvx2_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Output_N_Ho_Wo_K
    : public DeviceConvFwdBiasActivationAdd<InElementwiseOperation,
@@ -607,19 +609,51 @@ struct DeviceConvNDFwdBiasActivationAddAvx2_Input_N_Hi_Wi_C_Weight_K_Y_X_C_Outpu
            !UseBLocalBuffer,
            ConvForwardSpecialization>;

-    using CThreadwiseCopy =
-        ck::cpu::ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_Residual_MxN<
-            CDataType,
-            C0DataType,
-            C1DataType,
-            CDataType,
-            CGridDesc,
-            C0GridDesc,
-            C1GridDesc,
-            decltype(GetOutputBlockDescriptor()),
-            OutElementwiseOperation,
-            !UseCLocalBuffer,
-            BiasAlongGemmM>;
+    static constexpr auto GetCThreadwiseCopy()
+    {
+        constexpr ck::index_t C_nDim = CGridDesc::GetNumOfDimension();
+        if constexpr(FuseBias && FuseAdd)
+        {
+            return ck::cpu::
+                ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_Residual_MxN<
+                    CDataType,
+                    C0DataType,
+                    C1DataType,
+                    CDataType,
+                    CGridDesc,
+                    C0GridDesc,
+                    C1GridDesc,
+                    decltype(GetOutputBlockDescriptor()),
+                    OutElementwiseOperation,
+                    !UseCLocalBuffer,
+                    BiasAlongGemmM>(CGridDesc{},
+                                    ck::make_zero_multi_index<C_nDim>(),
+                                    GetOutputBlockDescriptor(),
+                                    ck::make_zero_multi_index<C_nDim>(),
+                                    OutElementwiseOperation{});
+        }
+        else if constexpr(FuseBias && !FuseAdd)
+        {
+            return ck::cpu::ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_MxN<
+                CDataType,
+                C0DataType,
+                C1DataType,
+                CDataType,
+                CGridDesc,
+                C0GridDesc,
+                C1GridDesc,
+                decltype(GetOutputBlockDescriptor()),
+                OutElementwiseOperation,
+                !UseCLocalBuffer,
+                BiasAlongGemmM>(CGridDesc{},
+                                ck::make_zero_multi_index<C_nDim>(),
+                                GetOutputBlockDescriptor(),
+                                ck::make_zero_multi_index<C_nDim>(),
+                                OutElementwiseOperation{});
+        }
+    }
+
+    using CThreadwiseCopy = decltype(GetCThreadwiseCopy());

    using GridwiseGemm = ck::cpu::GridwiseGemmBiasActivationAddAvx2_MxN<
        ADataType,               // InDataType,

--- a/include/ck/tensor_operation/cpu/device/device_convnd_fwd_bias_activation_add_avx2_nhwc_kyxck8_nhwk.hpp
+++ b/include/ck/tensor_operation/cpu/device/device_convnd_fwd_bias_activation_add_avx2_nhwc_kyxck8_nhwk.hpp
@@ -37,6 +37,8 @@ template <typename InDataType,
          bool UseALocalBuffer,
          bool UseBLocalBuffer,
          bool UseCLocalBuffer,
+          bool FuseBias,
+          bool FuseAdd,
          bool BiasAlongGemmM>
 struct DeviceConvNDFwdBiasActivationAddAvx2_Input_N_Hi_Wi_C_Weight_K_Y_X_C_K8_Output_N_Ho_Wo_K
    : public DeviceConvFwdBiasActivationAdd<InElementwiseOperation,
@@ -584,19 +586,51 @@ struct DeviceConvNDFwdBiasActivationAddAvx2_Input_N_Hi_Wi_C_Weight_K_Y_X_C_K8_Ou
            !UseBLocalBuffer,
            ConvForwardSpecialization>;

-    using CThreadwiseCopy =
-        ck::cpu::ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_Residual_MxN<
-            CDataType,
-            C0DataType,
-            C1DataType,
-            CDataType,
-            CGridDesc,
-            C0GridDesc,
-            C1GridDesc,
-            decltype(GetOutputBlockDescriptor()),
-            OutElementwiseOperation,
-            !UseCLocalBuffer,
-            BiasAlongGemmM>;
+    static constexpr auto GetCThreadwiseCopy()
+    {
+        constexpr ck::index_t C_nDim = CGridDesc::GetNumOfDimension();
+        if constexpr(FuseBias && FuseAdd)
+        {
+            return ck::cpu::
+                ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_Residual_MxN<
+                    CDataType,
+                    C0DataType,
+                    C1DataType,
+                    CDataType,
+                    CGridDesc,
+                    C0GridDesc,
+                    C1GridDesc,
+                    decltype(GetOutputBlockDescriptor()),
+                    OutElementwiseOperation,
+                    !UseCLocalBuffer,
+                    BiasAlongGemmM>(CGridDesc{},
+                                    ck::make_zero_multi_index<C_nDim>(),
+                                    GetOutputBlockDescriptor(),
+                                    ck::make_zero_multi_index<C_nDim>(),
+                                    OutElementwiseOperation{});
+        }
+        else if constexpr(FuseBias && !FuseAdd)
+        {
+            return ck::cpu::ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_MxN<
+                CDataType,
+                C0DataType,
+                C1DataType,
+                CDataType,
+                CGridDesc,
+                C0GridDesc,
+                C1GridDesc,
+                decltype(GetOutputBlockDescriptor()),
+                OutElementwiseOperation,
+                !UseCLocalBuffer,
+                BiasAlongGemmM>(CGridDesc{},
+                                ck::make_zero_multi_index<C_nDim>(),
+                                GetOutputBlockDescriptor(),
+                                ck::make_zero_multi_index<C_nDim>(),
+                                OutElementwiseOperation{});
+        }
+    }
+
+    using CThreadwiseCopy = decltype(GetCThreadwiseCopy());

    using GridwiseGemm = ck::cpu::GridwiseGemmBiasActivationAddAvx2_MxN<
        ADataType,               // InDataType,
@@ -954,7 +988,7 @@ struct DeviceConvNDFwdBiasActivationAddAvx2_Input_N_Hi_Wi_C_Weight_K_Y_X_C_K8_Ou
                return "G";
        };
        // clang-format off
-        str << "DeviceConv" << std::to_string(NumDimSpatial) 
+        str << "DeviceConv" << std::to_string(NumDimSpatial)
            << "DFwd_BAA_Avx2_NHWC_KYXCK8"
            <<"_FS"<< static_cast<int>(ConvForwardSpecialization)
            <<"_KS"<< static_cast<int>(gridwise_gemm.dynamic_tunable.gemm_k_spec)

--- a/include/ck/tensor_operation/cpu/device/device_convnd_fwd_bias_activation_add_avx2_nhwc_yxck_nhwk.hpp
+++ b/include/ck/tensor_operation/cpu/device/device_convnd_fwd_bias_activation_add_avx2_nhwc_yxck_nhwk.hpp
@@ -36,6 +36,8 @@ template <typename InDataType,
          bool UseALocalBuffer,
          bool UseBLocalBuffer,
          bool UseCLocalBuffer,
+          bool FuseBias,
+          bool FuseAdd,
          bool BiasAlongGemmM>
 struct DeviceConvNDFwdBiasActivationAddAvx2_Input_N_Hi_Wi_C_Weight_Y_X_C_K_Output_N_Ho_Wo_K
    : public DeviceConvFwdBiasActivationAdd<InElementwiseOperation,
@@ -580,19 +582,51 @@ struct DeviceConvNDFwdBiasActivationAddAvx2_Input_N_Hi_Wi_C_Weight_Y_X_C_K_Outpu
            !UseBLocalBuffer,
            ConvForwardSpecialization>;

-    using CThreadwiseCopy =
-        ck::cpu::ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_Residual_MxN<
-            CDataType,
-            C0DataType,
-            C1DataType,
-            CDataType,
-            CGridDesc,
-            C0GridDesc,
-            C1GridDesc,
-            decltype(GetOutputBlockDescriptor()),
-            OutElementwiseOperation,
-            !UseCLocalBuffer,
-            BiasAlongGemmM>;
+    static constexpr auto GetCThreadwiseCopy()
+    {
+        constexpr ck::index_t C_nDim = CGridDesc::GetNumOfDimension();
+        if constexpr(FuseBias && FuseAdd)
+        {
+            return ck::cpu::
+                ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_Residual_MxN<
+                    CDataType,
+                    C0DataType,
+                    C1DataType,
+                    CDataType,
+                    CGridDesc,
+                    C0GridDesc,
+                    C1GridDesc,
+                    decltype(GetOutputBlockDescriptor()),
+                    OutElementwiseOperation,
+                    !UseCLocalBuffer,
+                    BiasAlongGemmM>(CGridDesc{},
+                                    ck::make_zero_multi_index<C_nDim>(),
+                                    GetOutputBlockDescriptor(),
+                                    ck::make_zero_multi_index<C_nDim>(),
+                                    OutElementwiseOperation{});
+        }
+        else if constexpr(FuseBias && !FuseAdd)
+        {
+            return ck::cpu::ThreadwiseTensorSliceTransferAvx2Specialization_MatC_Store_Bias_MxN<
+                CDataType,
+                C0DataType,
+                C1DataType,
+                CDataType,
+                CGridDesc,
+                C0GridDesc,
+                C1GridDesc,
+                decltype(GetOutputBlockDescriptor()),
+                OutElementwiseOperation,
+                !UseCLocalBuffer,
+                BiasAlongGemmM>(CGridDesc{},
+                                ck::make_zero_multi_index<C_nDim>(),
+                                GetOutputBlockDescriptor(),
+                                ck::make_zero_multi_index<C_nDim>(),
+                                OutElementwiseOperation{});
+        }
+    }
+
+    using CThreadwiseCopy = decltype(GetCThreadwiseCopy());

    using GridwiseGemm = ck::cpu::GridwiseGemmBiasActivationAddAvx2_MxN<
        ADataType,               // InDataType,

--- a/include/ck/tensor_operation/cpu/thread/threadwise_tensor_slice_transfer_avx2_specialization.hpp
+++ b/include/ck/tensor_operation/cpu/thread/threadwise_tensor_slice_transfer_avx2_specialization.hpp
--- a/library/src/tensor_operation_instance/cpu/conv2d_fwd_bias_activation_add/device_conv2d_bias_activation_add_avx2_nhwc_kyxc_nhwk_instance.cpp
+++ b/library/src/tensor_operation_instance/cpu/conv2d_fwd_bias_activation_add/device_conv2d_bias_activation_add_avx2_nhwc_kyxc_nhwk_instance.cpp
--- a/library/src/tensor_operation_instance/cpu/conv2d_fwd_bias_activation_add/device_conv2d_bias_activation_add_avx2_nhwc_kyxck8_nhwk_instance.cpp
+++ b/library/src/tensor_operation_instance/cpu/conv2d_fwd_bias_activation_add/device_conv2d_bias_activation_add_avx2_nhwc_kyxck8_nhwk_instance.cpp
--- a/library/src/tensor_operation_instance/cpu/conv2d_fwd_bias_activation_add/device_conv2d_bias_activation_add_avx2_nhwc_yxck_nhwk_instance.cpp
+++ b/library/src/tensor_operation_instance/cpu/conv2d_fwd_bias_activation_add/device_conv2d_bias_activation_add_avx2_nhwc_yxck_nhwk_instance.cpp