提交 76fa7157 编写于 作者: M Megvii Engine Team

feat(dnn/cuda): add cutlass nchw4 convolution

GitOrigin-RevId: 93c9b212f4543e50ec56b5623e1b77bf7209a45b
上级 1f3f4abc
...@@ -36,8 +36,9 @@ all: ${PARAM_DEFS} ${ELEMWISE_IMPL} ${CUDA_CONV_IMPL} ...@@ -36,8 +36,9 @@ all: ${PARAM_DEFS} ${ELEMWISE_IMPL} ${CUDA_CONV_IMPL}
../src/cuda/elemwise_multi_type/kimpl: gen_elemwise_multi_type_kern_impls.py ../src/cuda/elemwise_multi_type/kimpl: gen_elemwise_multi_type_kern_impls.py
./$^ --type cuda $@ ./$^ --type cuda $@
../src/cuda/conv_bias/int8/kimpl: gen_cuda_conv_bias_kern_impls.py ../src/cuda/conv_bias/int8/kimpl: gen_cuda_conv_bias_kern_impls.py gen_cutlass_conv_bias_kern_impls.py
./$^ --type dp4a $@ ./gen_cuda_conv_bias_kern_impls.py --type dp4a $@
./gen_cutlass_conv_bias_kern_impls.py --type dp4a $@
../src/cuda/conv_bias/int8_imma/kimpl: gen_cuda_conv_bias_kern_impls.py gen_cutlass_conv_bias_kern_impls.py ../src/cuda/conv_bias/int8_imma/kimpl: gen_cuda_conv_bias_kern_impls.py gen_cutlass_conv_bias_kern_impls.py
./gen_cuda_conv_bias_kern_impls.py --type imma $@ ./gen_cuda_conv_bias_kern_impls.py --type imma $@
......
...@@ -91,7 +91,10 @@ ConvBiasForwardImpl::AlgoPack::AlgoPack() { ...@@ -91,7 +91,10 @@ ConvBiasForwardImpl::AlgoPack::AlgoPack() {
} }
#endif #endif
#endif #endif
all_algos.push_back(&int8_nchw4_dotprod); fill_dp4a_algos();
for (auto&& algo : int8_nchw4_dotprod) {
all_algos.push_back(&algo);
}
all_algos.push_back(&int8_chwn4_dotprod); all_algos.push_back(&int8_chwn4_dotprod);
for (size_t i = all_algo_size; i < all_algos.size(); ++i) { for (size_t i = all_algo_size; i < all_algos.size(); ++i) {
non_cudnn_algos.push_back(all_algos[i]); non_cudnn_algos.push_back(all_algos[i]);
...@@ -253,6 +256,20 @@ void ConvBiasForwardImpl::AlgoPack::fill_imma_algos() { ...@@ -253,6 +256,20 @@ void ConvBiasForwardImpl::AlgoPack::fill_imma_algos() {
} }
#endif #endif
void ConvBiasForwardImpl::AlgoPack::fill_dp4a_algos() {
using AlgoParam = AlgoInt8NCHW4DotProdImplicitGemm::AlgoParam;
int8_nchw4_dotprod.emplace_back(AlgoParam{128, 128, 32, 64, 32, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{128, 64, 32, 64, 32, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{64, 128, 32, 64, 32, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{32, 128, 32, 32, 64, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{128, 32, 32, 64, 32, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{64, 64, 32, 64, 32, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{32, 64, 32, 32, 64, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{64, 32, 32, 64, 32, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{32, 32, 32, 32, 32, 32});
int8_nchw4_dotprod.emplace_back(AlgoParam{16, 64, 8, 16, 64, 8});
}
ConvBiasForwardImpl::AlgoBase* ConvBiasForwardImpl::AlgoBase*
ConvBiasForwardImpl::AlgoPack::cudnn_conv_from_enum( ConvBiasForwardImpl::AlgoPack::cudnn_conv_from_enum(
......
...@@ -386,18 +386,39 @@ public: ...@@ -386,18 +386,39 @@ public:
class ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm final class ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm final
: public AlgoBase { : public AlgoBase {
public: public:
AlgoInt8NCHW4DotProdImplicitGemm() = default; struct AlgoParam {
int threadblock_m;
int threadblock_n;
int threadblock_k;
int warp_m;
int warp_n;
int warp_k;
std::string to_string() {
/// default algorithm
if (threadblock_m == 128 && threadblock_n == 128 &&
threadblock_k == 32 && warp_m == 32 && warp_n == 64 &&
warp_k == 32) {
return "";
}
return ssprintf("_%dX%dX%d_%dX%dX%d", threadblock_m, threadblock_n,
threadblock_k, warp_m, warp_n, warp_k);
}
};
AlgoInt8NCHW4DotProdImplicitGemm(AlgoParam algo_param)
: m_algo_param{algo_param},
m_name{ssprintf("INT8_NCHW4_DOTPROD_IMPLICIT_GEMM%s",
m_algo_param.to_string().c_str())} {}
bool is_available(const SizeArgs& args) const override; bool is_available(const SizeArgs& args) const override;
size_t get_workspace_in_bytes(const SizeArgs& args) const override; size_t get_workspace_in_bytes(const SizeArgs& args) const override;
void exec(const ExecArgs& args) const override; void exec(const ExecArgs& args) const override;
const char* name() const override { const char* name() const override { return m_name.c_str(); }
return "INT8_NCHW4_DOTPROD_IMPLICIT_GEMM";
}
bool is_reproducible() const override { return true; } bool is_reproducible() const override { return true; }
private: private:
WorkspaceBundle get_workspace_bundle(dt_byte* raw_ptr, WorkspaceBundle get_workspace_bundle(dt_byte* raw_ptr,
const SizeArgs& args) const; const SizeArgs& args) const;
AlgoParam m_algo_param;
std::string m_name;
}; };
#if CUDA_VERSION >= 10000 #if CUDA_VERSION >= 10000
...@@ -578,7 +599,7 @@ public: ...@@ -578,7 +599,7 @@ public:
AlgoMatmul8x8x32 matmul8x8x32; AlgoMatmul8x8x32 matmul8x8x32;
AlgoBatchedMatmul batched_matmul; AlgoBatchedMatmul batched_matmul;
Algo1x1 a1x1; Algo1x1 a1x1;
AlgoInt8NCHW4DotProdImplicitGemm int8_nchw4_dotprod; std::vector<AlgoInt8NCHW4DotProdImplicitGemm> int8_nchw4_dotprod;
AlgoInt8CHWN4DotProdImplicitGemm int8_chwn4_dotprod; AlgoInt8CHWN4DotProdImplicitGemm int8_chwn4_dotprod;
#if CUDA_VERSION >= 10000 #if CUDA_VERSION >= 10000
AlgoQUInt4x4x32WMMA wmma_quint4x4x32; AlgoQUInt4x4x32WMMA wmma_quint4x4x32;
...@@ -605,6 +626,7 @@ private: ...@@ -605,6 +626,7 @@ private:
void fill_imma_algos(); void fill_imma_algos();
#endif #endif
void fill_cudnn_algos(); void fill_cudnn_algos();
void fill_dp4a_algos();
}; };
} // namespace cuda } // namespace cuda
......
...@@ -19,7 +19,6 @@ ...@@ -19,7 +19,6 @@
#endif #endif
#include "src/common/opr_param_defs_enumv.cuh" #include "src/common/opr_param_defs_enumv.cuh"
#include "src/cuda/conv_bias/cutlass_convolution_wrapper.cuh" #include "src/cuda/conv_bias/cutlass_convolution_wrapper.cuh"
#pragma GCC diagnostic pop #pragma GCC diagnostic pop
using namespace megdnn; using namespace megdnn;
...@@ -149,4 +148,130 @@ INST(true); ...@@ -149,4 +148,130 @@ INST(true);
INST(false); INST(false);
#undef INST #undef INST
#if MEGDNN_TEGRA_X1
template <bool NeedLoadFromConstMem>
void megdnn::cuda::cutlass_wrapper::
do_conv_bias_int8_implicit_gemm_dp4a_ncdiv4hw4(
const int8_t* /* d_src */, const int8_t* /* d_filter */,
const int32_t* /* d_bias */, const int8_t* /* d_z */,
int8_t* /* d_dst */, int* /* workspace */,
const convolution::ConvParam& /* param */,
uint32_t /* nonlinear_mode */, float /* alpha */,
float /* beta */, float /* gamma */, float /* scale */,
const GemmCoord& /* threadblock_shape */,
const GemmCoord& /* warp_shape */, cudaStream_t /* stream */) {}
#else
template <bool NeedLoadFromConstMem>
void megdnn::cuda::cutlass_wrapper::
do_conv_bias_int8_implicit_gemm_dp4a_ncdiv4hw4(
const int8_t* d_src, const int8_t* d_filter,
const int32_t* d_bias, const int8_t* d_z, int8_t* d_dst,
int* workspace, const convolution::ConvParam& param,
uint32_t nonlinear_mode, float alpha, float beta, float gamma,
float scale, const GemmCoord& threadblock_shape,
const GemmCoord& warp_shape, cudaStream_t stream) {
#define DISPATCH_KERNEL_WITH_TILE_SHAPE(threadblock_m_, threadblock_n_, \
threadblock_k_, warp_m_, warp_n_, \
warp_k_, aligned_) \
if (threadblock_shape.m() == threadblock_m_ && \
threadblock_shape.n() == threadblock_n_ && \
threadblock_shape.k() == threadblock_k_ && \
warp_shape.m() == warp_m_ && warp_shape.n() == warp_n_ && \
warp_shape.k() == warp_k_) { \
using ThreadBlockShape = \
cutlass::gemm::GemmShape<threadblock_m_, threadblock_n_, \
threadblock_k_>; \
using WarpShape = cutlass::gemm::GemmShape<warp_m_, warp_n_, warp_k_>; \
using InstructionShape = cutlass::gemm::GemmShape<1, 1, 4>; \
using Convolution = cutlass::convolution::device::Convolution< \
int8_t, cutlass::layout::TensorNCxHWx<4>, int8_t, \
cutlass::layout::TensorCxRSKx<4>, ElementOutput, \
cutlass::layout::TensorNCxHWx<4>, int32_t, \
cutlass::layout::TensorNCxHWx<4>, int32_t, \
cutlass::convolution::ConvType::kConvolution, \
cutlass::arch::OpClassSimt, cutlass::arch::Sm61, \
ThreadBlockShape, WarpShape, InstructionShape, EpilogueOp, \
cutlass::convolution::threadblock:: \
ConvolutionNCxHWxThreadblockSwizzle< \
cutlass::convolution::ConvType::kConvolution>, \
2, 4, aligned_, NeedLoadFromConstMem>; \
typename Convolution::ConvolutionParameter conv_param{ \
param.n, param.ci, param.co, param.hi, param.wi, \
param.fh, param.fw, param.ho, param.wo, param.sh, \
param.sw, param.ph, param.pw, 1, 1}; \
return cutlass_convolution_wrapper<Convolution>( \
d_src, d_filter, d_bias, d_z, d_dst, workspace, conv_param, \
epilogue, stream); \
}
#define DISPATCH_KERNEL \
DISPATCH_KERNEL_WITH_TILE_SHAPE(128, 128, 32, 64, 32, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(128, 64, 32, 64, 32, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(64, 128, 32, 64, 32, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(128, 32, 32, 64, 32, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(32, 128, 32, 32, 64, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(64, 64, 32, 64, 32, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(32, 64, 32, 32, 64, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(64, 32, 32, 64, 32, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(32, 32, 32, 32, 32, 32, 16); \
DISPATCH_KERNEL_WITH_TILE_SHAPE(16, 64, 8, 16, 64, 8, 4); \
megdnn_assert(false, \
"unsupported threadblock shape (%dx%dx%d) and warp shape " \
"(%dx%dx%d)", \
threadblock_shape.m(), threadblock_shape.n(), \
threadblock_shape.k(), warp_shape.m(), warp_shape.n(), \
warp_shape.k());
using ElementOutput = int8_t;
using ElementAccumulator = int32_t;
using ElementBias = int32_t;
using ElementCompute = float;
using NonlineMode = megdnn::param_enumv::ConvBias::NonlineMode;
switch (nonlinear_mode) {
case NonlineMode::IDENTITY: {
using EpilogueOp =
cutlass::epilogue::thread::BiasAddLinearCombinationClamp<
ElementOutput, 4, ElementAccumulator, ElementBias,
ElementCompute>;
typename EpilogueOp::Params epilogue{alpha, beta, gamma};
DISPATCH_KERNEL;
}
case NonlineMode::RELU: {
using EpilogueOp = cutlass::epilogue::thread::
BiasAddLinearCombinationReluClamp<
ElementOutput, 4, ElementAccumulator, ElementBias,
ElementCompute>;
typename EpilogueOp::Params epilogue{alpha, beta, gamma, 0};
DISPATCH_KERNEL;
}
case NonlineMode::H_SWISH: {
using EpilogueOp = cutlass::epilogue::thread::
BiasAddLinearCombinationHSwishClamp<
ElementOutput, 4, ElementAccumulator, ElementBias,
ElementCompute>;
typename EpilogueOp::Params epilogue{alpha, beta, gamma, scale};
DISPATCH_KERNEL;
}
default:
megdnn_assert(false,
"unsupported nonlinear mode for conv bias operator");
}
#undef DISPATCH_KERNEL_WITH_TILE_SHAPE
#undef DISPATCH_KERNEL
}
#endif
#define INST(need_load_from_const_mem) \
template void megdnn::cuda::cutlass_wrapper:: \
do_conv_bias_int8_implicit_gemm_dp4a_ncdiv4hw4< \
need_load_from_const_mem>( \
const int8_t* d_src, const int8_t* d_filter, \
const int32_t* d_bias, const int8_t* d_z, int8_t* d_dst, \
int* workspace, const convolution::ConvParam& param, \
uint32_t nonlinear_mode, float alpha, float beta, \
float gamma, float scale, \
const GemmCoord& threadblock_shape, \
const GemmCoord& warp_shape, cudaStream_t stream);
INST(true);
INST(false);
#undef INST
// vim: syntax=cuda.doxygen // vim: syntax=cuda.doxygen
...@@ -37,6 +37,15 @@ void do_conv_bias_int8_implicit_gemm_imma_ncdiv32hw32( ...@@ -37,6 +37,15 @@ void do_conv_bias_int8_implicit_gemm_imma_ncdiv32hw32(
const GemmCoord& threadblock_shape, const GemmCoord& warp_shape, const GemmCoord& threadblock_shape, const GemmCoord& warp_shape,
cudaStream_t stream); cudaStream_t stream);
template <bool NeedLoadFromConstMem>
void do_conv_bias_int8_implicit_gemm_dp4a_ncdiv4hw4(
const int8_t* d_src, const int8_t* d_filter, const int32_t* d_bias,
const int8_t* d_z, int8_t* d_dst, int* workspace,
const convolution::ConvParam& param, uint32_t nonlinear_mode,
float alpha, float beta, float gamma, float scale,
const GemmCoord& threadblock_shape, const GemmCoord& warp_shape,
cudaStream_t stream);
} // namespace cutlass_wrapper } // namespace cutlass_wrapper
} // namespace cuda } // namespace cuda
} // namespace megdnn } // namespace megdnn
......
...@@ -57,30 +57,16 @@ bool ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::is_available( ...@@ -57,30 +57,16 @@ bool ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::is_available(
// only support sm_75 or later, platform should have tensorcore int8 // only support sm_75 or later, platform should have tensorcore int8
// support // support
available &= is_compute_capability_required(7, 5); available &= is_compute_capability_required(7, 5);
if (fh == 1 && fw == 1) // FIXME: too large filter size is not supported now
return available; available &= fh * fw <= 49;
// for non 1x1 convolution, we have to check constant memory size
auto&& device_prop = current_device_prop();
// const mem size >= 64K
available &= device_prop.totalConstMem >= 65536;
size_t const_mem_usage = get_workspace_in_bytes(args) -
args.filter_layout->span().dist_byte();
available &= const_mem_usage <= device_prop.totalConstMem;
return available; return available;
} }
WorkspaceBundle WorkspaceBundle
ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::get_workspace_bundle( ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::get_workspace_bundle(
dt_byte* raw_ptr, const SizeArgs& args) const { dt_byte* raw_ptr, const SizeArgs& args) const {
size_t ci = args.filter_layout->operator[](1) * 32;
size_t fh = args.filter_layout->operator[](2);
size_t fw = args.filter_layout->operator[](3);
size_t ws_filter = args.filter_layout->span().dist_byte(); size_t ws_filter = args.filter_layout->span().dist_byte();
if (fh == 1 && fw == 1) { return WorkspaceBundle{raw_ptr, {ws_filter}};
return WorkspaceBundle{raw_ptr, {ws_filter}};
}
size_t ws_size = (ci / 32) * fh * fw * sizeof(int32_t) * 2;
return WorkspaceBundle{raw_ptr, {ws_filter, ws_size}};
} }
size_t size_t
...@@ -148,9 +134,9 @@ void ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::exec( ...@@ -148,9 +134,9 @@ void ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::exec(
false>(args.src_tensor->compatible_ptr<int8_t>(), false>(args.src_tensor->compatible_ptr<int8_t>(),
reinterpret_cast<int8_t*>(ws_filter), reinterpret_cast<int8_t*>(ws_filter),
args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr, args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr,
args.dst_tensor->compatible_ptr<int8_t>(), args.dst_tensor->compatible_ptr<int8_t>(), nullptr,
nullptr, kern_param, nonlinear_mode, kern_param, nonlinear_mode, alpha, beta, gamma,
alpha, beta, gamma, dst_scale, dst_scale,
cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m, cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m,
m_algo_param.threadblock_n, m_algo_param.threadblock_n,
m_algo_param.threadblock_k}, m_algo_param.threadblock_k},
...@@ -159,14 +145,12 @@ void ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::exec( ...@@ -159,14 +145,12 @@ void ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::exec(
m_algo_param.warp_k}, m_algo_param.warp_k},
stream); stream);
} else { } else {
auto workspace = ws.get(1);
cutlass_wrapper::do_conv_bias_int8_implicit_gemm_imma_ncdiv32hw32<true>( cutlass_wrapper::do_conv_bias_int8_implicit_gemm_imma_ncdiv32hw32<true>(
args.src_tensor->compatible_ptr<int8_t>(), args.src_tensor->compatible_ptr<int8_t>(),
reinterpret_cast<int8_t*>(ws_filter), reinterpret_cast<int8_t*>(ws_filter),
args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr, args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr,
args.dst_tensor->compatible_ptr<int8_t>(), args.dst_tensor->compatible_ptr<int8_t>(), nullptr, kern_param,
reinterpret_cast<int*>(workspace), kern_param, nonlinear_mode, nonlinear_mode, alpha, beta, gamma, dst_scale,
alpha, beta, gamma, dst_scale,
cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m, cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m,
m_algo_param.threadblock_n, m_algo_param.threadblock_n,
m_algo_param.threadblock_k}, m_algo_param.threadblock_k},
......
...@@ -11,7 +11,8 @@ ...@@ -11,7 +11,8 @@
#include "./algo.h" #include "./algo.h"
#include "src/cuda/utils.h" #include "src/cuda/utils.h"
#include "src/cuda/convolution_helper/bias_visitor.cuh" #include "src/cuda/convolution_helper/parameter.cuh"
#include "src/cuda/conv_bias/cutlass_convolution_wrapper.cuh"
using namespace megdnn; using namespace megdnn;
using namespace cuda; using namespace cuda;
...@@ -53,21 +54,16 @@ bool ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::is_available( ...@@ -53,21 +54,16 @@ bool ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::is_available(
// only support sm_61 or later, platform should have fast native int8 // only support sm_61 or later, platform should have fast native int8
// support // support
available &= is_compute_capability_required(6, 1); available &= is_compute_capability_required(6, 1);
// FIXME: too large filter size is not supported now
available &= fh * fw <= 49;
return available; return available;
} }
WorkspaceBundle WorkspaceBundle
ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::get_workspace_bundle( ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::get_workspace_bundle(
dt_byte* raw_ptr, const SizeArgs& args) const { dt_byte* raw_ptr, const SizeArgs& args) const {
size_t ws_size_src = args.src_layout->span().dist_byte(); size_t ws_filter = args.filter_layout->span().dist_byte();
size_t ws_size_filter = args.filter_layout->span().dist_byte(); return WorkspaceBundle{raw_ptr, {ws_filter}};
size_t ws_size_dst = args.dst_layout->span().dist_byte();
if (args.z_layout->ndim > 0) {
size_t ws_size_z = args.z_layout->span().dist_byte();
return WorkspaceBundle{
raw_ptr, {ws_size_src, ws_size_filter, ws_size_dst, ws_size_z}};
}
return WorkspaceBundle{raw_ptr, {ws_size_src, ws_size_filter, ws_size_dst}};
} }
size_t size_t
...@@ -84,27 +80,9 @@ void ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::exec( ...@@ -84,27 +80,9 @@ void ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::exec(
UNPACK_CONV_BIAS_NCHW4_PARAM(*(args.src_layout), fm, *(args.dst_layout), UNPACK_CONV_BIAS_NCHW4_PARAM(*(args.src_layout), fm, *(args.dst_layout),
param); param);
auto ws = get_workspace_bundle(args.workspace.raw_ptr, args); auto ws = get_workspace_bundle(args.workspace.raw_ptr, args);
auto ws_src = ws.get(0); auto ws_filter = ws.get(0);
auto ws_filter = ws.get(1);
auto ws_dst = ws.get(2);
auto&& stream = cuda_stream(args.opr->handle()); auto&& stream = cuda_stream(args.opr->handle());
// reformat src from nchw4 to chwn4
{
TensorLayout src{{n, ci / 4 * hi * wi}, dtype::Int32()};
src.init_contiguous_stride();
TensorLayout dst = src;
dst.stride[0] = 1, dst.stride[1] = dst[0];
TensorND ts_src, ts_dst;
ts_src.raw_ptr = args.src_tensor->raw_ptr;
ts_src.layout = src;
ts_dst.raw_ptr = ws_src;
ts_dst.layout = dst;
auto&& transpose =
args.opr->handle()->create_operator<RelayoutForward>();
transpose->exec(ts_src, ts_dst);
}
// reformat filter from nchw4 to chwn4 // reformat filter from nchw4 to chwn4
{ {
TensorLayout src{{co, ci / 4 * fh * fw}, dtype::Int32()}; TensorLayout src{{co, ci / 4 * fh * fw}, dtype::Int32()};
...@@ -136,53 +114,42 @@ void ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::exec( ...@@ -136,53 +114,42 @@ void ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::exec(
dst_scale = args.dst_layout->dtype.param<dtype::QuantizedS8>().scale; dst_scale = args.dst_layout->dtype.param<dtype::QuantizedS8>().scale;
float alpha = src_scale * filter_scale / dst_scale, float alpha = src_scale * filter_scale / dst_scale,
beta = bias_scale / dst_scale; beta = bias_scale / dst_scale;
// process z
int8_t* z_dev_ptr = nullptr; int8_t* z_dev_ptr = nullptr;
float gamma = 1.f; float gamma = 0.0;
if (args.z_layout->ndim > 0) { if (args.z_layout->ndim > 0) {
auto ws_z = ws.get(3); z_dev_ptr = args.z_tensor->compatible_ptr<int8_t>();
TensorLayout src{{n, co / 4 * ho * wo}, dtype::Int32()};
src.init_contiguous_stride();
TensorLayout dst = src;
dst.stride[0] = 1, dst.stride[1] = dst[0];
TensorND ts_src, ts_dst;
ts_src.raw_ptr = args.z_tensor->raw_ptr;
ts_src.layout = src;
ts_dst.raw_ptr = ws_z;
ts_dst.layout = dst;
auto&& transpose =
args.opr->handle()->create_operator<RelayoutForward>();
transpose->exec(ts_src, ts_dst);
z_dev_ptr = reinterpret_cast<int8_t*>(ws_z);
float z_scale = args.z_layout->dtype.param<dtype::QuantizedS8>().scale; float z_scale = args.z_layout->dtype.param<dtype::QuantizedS8>().scale;
gamma = z_scale / dst_scale; gamma = z_scale / dst_scale;
} }
uint32_t nonlinear_mode = static_cast<uint32_t>(param.nonlineMode);
convolution::PerChannelBiasVisitor bias_visitor; if (fh == 1 && fw == 1) {
bias_visitor.bias = args.bias_tensor->compatible_ptr<int32_t>(); cutlass_wrapper::do_conv_bias_int8_implicit_gemm_dp4a_ncdiv4hw4<false>(
ConvBiasForwardImpl::AlgoInt8CHWN4DotProdImplicitGemm:: args.src_tensor->compatible_ptr<int8_t>(),
dispatch_nonlinear_mode<convolution::PerChannelBiasVisitor>( reinterpret_cast<int8_t*>(ws_filter),
reinterpret_cast<int8_t*>(ws_src), args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr,
reinterpret_cast<int8_t*>(ws_filter), bias_visitor, args.dst_tensor->compatible_ptr<int8_t>(), nullptr, kern_param,
z_dev_ptr, reinterpret_cast<int8_t*>(ws_dst), kern_param, nonlinear_mode, alpha, beta, gamma, dst_scale,
alpha, beta, gamma, dst_scale, stream, param.nonlineMode); cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m,
m_algo_param.threadblock_n,
// reformat chwn4 to nchw4 m_algo_param.threadblock_k},
{ cutlass_wrapper::GemmCoord{m_algo_param.warp_m,
TensorLayout src{{co / 4 * ho * wo, n}, dtype::Int32()}; m_algo_param.warp_n,
src.init_contiguous_stride(); m_algo_param.warp_k},
TensorLayout dst = src; stream);
dst.stride[0] = 1, dst.stride[1] = dst[0]; } else {
TensorND ts_src, ts_dst; cutlass_wrapper::do_conv_bias_int8_implicit_gemm_dp4a_ncdiv4hw4<true>(
ts_src.raw_ptr = ws_dst; args.src_tensor->compatible_ptr<int8_t>(),
ts_src.layout = src; reinterpret_cast<int8_t*>(ws_filter),
ts_dst.raw_ptr = args.dst_tensor->raw_ptr; args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr,
ts_dst.layout = dst; args.dst_tensor->compatible_ptr<int8_t>(), nullptr, kern_param,
auto&& transpose = nonlinear_mode, alpha, beta, gamma, dst_scale,
args.opr->handle()->create_operator<RelayoutForward>(); cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m,
transpose->exec(ts_src, ts_dst); m_algo_param.threadblock_n,
m_algo_param.threadblock_k},
cutlass_wrapper::GemmCoord{m_algo_param.warp_m,
m_algo_param.warp_n,
m_algo_param.warp_k},
stream);
} }
} }
......
/** /**
* \file * \file
* dnn/src/cuda/conv_bias/int8_imma/conv_bias_int8_implicit_gemm_imma_ncdiv32hw32.cuinl * dnn/src/cuda/conv_bias/int8/conv_bias_int8_implicit_gemm_cutlass_wrapper.cuinl
* MegEngine is Licensed under the Apache License, Version 2.0 (the "License") * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
* *
* Copyright (c) 2014-2020 Megvii Inc. All rights reserved. * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
......
此差异已折叠。
此差异已折叠。
此差异已折叠。
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册