wehub-resource-sync
244 lines
7.5 KiB
Python
244 lines
7.5 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import sys
|
|
|
|
dirname, filename = os.path.split(os.path.abspath(sys.argv[0]))
|
|
sys.path.append(dirname + "/../")
|
|
from util import SubstituteTemplate
|
|
|
|
# For beginners, these template parameters may be difficult to understand.
|
|
# Please refer to the conv-related demo of CUTLASS for better understanding.
|
|
# https://github.com/NVIDIA/cutlass/tree/master/examples
|
|
|
|
CommonCutlassConvKernelDeclare = """
|
|
cutlass::Status ${kernel_func_name}(const ConvAllParams& params) {
|
|
using kernel_base =
|
|
typename cutlass::conv::kernel::${conv_kind_name}<
|
|
${element_a},
|
|
${layout_a},
|
|
${element_b},
|
|
${layout_b},
|
|
${element_c},
|
|
${layout_c},
|
|
${element_accum},
|
|
${opcode_class},
|
|
${arch},
|
|
cutlass::gemm::GemmShape<${Tshape}>,
|
|
cutlass::gemm::GemmShape<${Wshape}>,
|
|
cutlass::gemm::GemmShape<${Ishape}>,
|
|
${epi_part},
|
|
${swizzling_functor},
|
|
${stages},
|
|
${math_operator},
|
|
${iterator_algorithm},
|
|
${stride_support},
|
|
${align_a},
|
|
${align_b}
|
|
>::Kernel;
|
|
|
|
using ImplicitGemm =
|
|
cutlass::conv::device::ImplicitGemmConvolution<kernel_base>;
|
|
|
|
${element_a} *input = (${element_a} *)(params.input);
|
|
${element_b} *weight = (${element_b} *)(params.weight);
|
|
${element_c} *bias = (${element_c} *)(params.bias);
|
|
${element_c} *output = (${element_c} *)(params.output);
|
|
// only used by conv2d_bias_residual
|
|
auto residual = (${element_c} *)(params.residual);
|
|
|
|
int batch = params.batch;
|
|
int ic = params.ic;
|
|
int ih = params.ih;
|
|
int iw = params.iw;
|
|
int kh = params.kh;
|
|
int kw = params.kw;
|
|
int oc = params.oc;
|
|
int pad_h0 = params.pad_h0;
|
|
int pad_w0 = params.pad_w0;
|
|
int stride_h = params.stride_h;
|
|
int stride_w = params.stride_w;
|
|
int groups = params.groups;
|
|
int kc = ic / groups;
|
|
|
|
int oh = params.oh;
|
|
int ow = params.ow;
|
|
int dilation_h = params.dilation_h;
|
|
int dilation_w = params.dilation_w;
|
|
int split_k_slices = ${split_k_slices};
|
|
|
|
cutlass::conv::Conv2dProblemSize problem_size({batch, ih, iw, ic},
|
|
{oc, kh, kw, ic / groups},
|
|
{pad_h0, 0, pad_w0, 0},
|
|
{stride_h, stride_w},
|
|
{dilation_h, dilation_w},
|
|
{batch, oh, ow, oc},
|
|
cutlass::conv::Mode::kCrossCorrelation,
|
|
split_k_slices,
|
|
groups);
|
|
"""
|
|
|
|
# This is the execution part of this cutlass conv kernel.
|
|
|
|
CommonCutlassConvKernelExecute = """
|
|
ImplicitGemm implicit_gemm_op;
|
|
size_t bytes = implicit_gemm_op.get_workspace_size(arguments);
|
|
|
|
auto stream = params.stream;
|
|
void *workspace = params.workspace;
|
|
|
|
cutlass::Status status = implicit_gemm_op.can_implement(arguments);
|
|
CUTLASS_CHECK(status);
|
|
status = implicit_gemm_op.initialize(arguments, workspace);
|
|
CUTLASS_CHECK(status);
|
|
status = implicit_gemm_op(stream);
|
|
CUTLASS_CHECK(status);
|
|
return status;
|
|
}
|
|
"""
|
|
|
|
# CommonConvFunction is a wrapper for many kernels
|
|
# a func_name is like conv2d_bias_silu_sm75
|
|
# it has many kernels, we should pick up a performance-best
|
|
# ${func_name} is like conv2d_bias_silu_sm75
|
|
# ${enum_op_name} is like CONV2D_BIAS_SILU
|
|
|
|
CommonConvFunction = """
|
|
|
|
${kernel_func_declare}
|
|
|
|
std::vector<std::function<cutlass::Status(const ConvAllParams)>>
|
|
${func_name}_all_func = {${all_kernel_func_name}};
|
|
|
|
std::map<std::vector<int>, int> map_problem_${func_name};
|
|
std::mutex ${func_name}_mutex;
|
|
|
|
bool ${func_name}(ConvAllParams params) {
|
|
int batch = params.batch;
|
|
int ic = params.ic;
|
|
int ih = params.ih;
|
|
int iw = params.iw;
|
|
int kh = params.kh;
|
|
int kw = params.kw;
|
|
int oc = params.oc;
|
|
//int pad_h0 = params.pad_h0;
|
|
//int pad_w0 = params.pad_w0;
|
|
int groups = params.groups;
|
|
int stride_h = params.stride_h;
|
|
int stride_w = params.stride_w;
|
|
|
|
std::vector<int> problem_size = {
|
|
batch, ic, ih, iw, kh, kw, oc, groups, stride_h, stride_w};
|
|
|
|
if (map_problem_${func_name}.count(problem_size)) {
|
|
${func_name}_all_func[map_problem_${func_name}.at(problem_size)](
|
|
params);
|
|
return true;
|
|
}
|
|
|
|
int best_config_index = ProfileToGetBestConfig(
|
|
${func_name}_all_func, params, ${enum_op_name});
|
|
|
|
std::lock_guard<std::mutex> guard(${func_name}_mutex);
|
|
|
|
map_problem_${func_name}[problem_size] = best_config_index;
|
|
${func_name}_all_func[best_config_index](params);
|
|
return true;
|
|
}
|
|
"""
|
|
|
|
|
|
# We should wrapper all op_name_with_sm_version into a function
|
|
# like : wrapper conv2d_bias_silu_sm75, conv2d_bias_silu_sm80, conv2d_bias_silu_sm86 into conv2d_bias_silu for phi kernel
|
|
# this function is invoked by phi kernel
|
|
|
|
CommonWrapperForPhi = """
|
|
bool ${op_name}(ConvAllParams params) {
|
|
${dispatch_body}
|
|
}
|
|
"""
|
|
|
|
|
|
def convert_c_data_type(dtype):
|
|
if dtype == "fp16":
|
|
return "Conv2dDataType::fp16"
|
|
elif dtype == "bf16":
|
|
return "Conv2dDataType::bf16"
|
|
elif dtype == "fp32":
|
|
return "Conv2dDataType::fp32"
|
|
else:
|
|
return None
|
|
|
|
|
|
CommonDispatchTemp = '''
|
|
if (params.sm_version == ${sm_code} && params.data_type == ${data_type})
|
|
{
|
|
return ${op_name_with_sm}(params);
|
|
}
|
|
'''
|
|
|
|
|
|
# this is a file's ending part
|
|
|
|
CommonTail = '''
|
|
} // namespace cutlass_internal
|
|
} // namespace fusion
|
|
} // namespace phi
|
|
'''
|
|
|
|
|
|
# Wrap different sm versions into a function called by phi
|
|
def GenerateFunctionForPhi(
|
|
sm_versions_and_types, support_epi_funcs, underscore_names, camel_names
|
|
):
|
|
generated_code = ""
|
|
for epi_func in support_epi_funcs:
|
|
dispatch_body = ""
|
|
for sm_version, data_type in sm_versions_and_types:
|
|
sm_dicts = {}
|
|
sm_dicts["sm_code"] = sm_version
|
|
sm_dicts["data_type"] = convert_c_data_type(data_type)
|
|
sm_dicts["op_name_with_sm"] = (
|
|
underscore_names[epi_func].lower()
|
|
+ "_sm"
|
|
+ sm_version
|
|
+ "_"
|
|
+ data_type
|
|
)
|
|
dispatch_body += SubstituteTemplate(CommonDispatchTemp, sm_dicts)
|
|
dispatch_body += ''' return false;'''
|
|
op_dicts = {}
|
|
op_dicts["dispatch_body"] = dispatch_body
|
|
op_dicts["op_name"] = camel_names[epi_func]
|
|
generated_code += SubstituteTemplate(CommonWrapperForPhi, op_dicts)
|
|
return generated_code
|
|
|
|
|
|
# We modify some template parameters based on CommonCutlassConvKernelDeclare.
|
|
CommonCutlassConv2dDepthwiseKernelDeclare = (
|
|
CommonCutlassConvKernelDeclare.replace(
|
|
"${align_a}", "cutlass::MatrixShape<${strided_shape}>"
|
|
)
|
|
.replace("${align_b}", "cutlass::MatrixShape<${dilation_shape}>")
|
|
.replace("ImplicitGemmConvolution", "DirectConvolution")
|
|
.replace(
|
|
"cutlass::gemm::GemmShape<${Tshape}>,",
|
|
'''cutlass::gemm::GemmShape<${Tshape}>,
|
|
cutlass::conv::TensorNHWCShape<${T_output_shape}>,
|
|
cutlass::MatrixShape<${filter_shape}>,
|
|
''',
|
|
)
|
|
)
|