feat(dnn/cuda): add cutlass nchw4 convolution

GitOrigin-RevId: 93c9b212f4543e50ec56b5623e1b77bf7209a45b

dnn/scripts/Makefile

@@ -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
 	./$^ --type cuda $@
 
-../src/cuda/conv_bias/int8/kimpl: gen_cuda_conv_bias_kern_impls.py
-	./$^ --type dp4a $@
+../src/cuda/conv_bias/int8/kimpl: gen_cuda_conv_bias_kern_impls.py gen_cutlass_conv_bias_kern_impls.py
+	./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
 	./gen_cuda_conv_bias_kern_impls.py --type imma $@

dnn/src/cuda/conv_bias/algo.cpp

@@ -91,7 +91,10 @@ ConvBiasForwardImpl::AlgoPack::AlgoPack() {
     }
 #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);
     for (size_t i = all_algo_size; i < all_algos.size(); ++i) {
         non_cudnn_algos.push_back(all_algos[i]);
@@ -253,6 +256,20 @@ void ConvBiasForwardImpl::AlgoPack::fill_imma_algos() {
 }
 #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::AlgoPack::cudnn_conv_from_enum(

dnn/src/cuda/conv_bias/algo.h

@@ -386,18 +386,39 @@ public:
 class ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm final
         : public AlgoBase {
 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;
     size_t get_workspace_in_bytes(const SizeArgs& args) const override;
     void exec(const ExecArgs& args) const override;
-    const char* name() const override {
-        return "INT8_NCHW4_DOTPROD_IMPLICIT_GEMM";
-    }
+    const char* name() const override { return m_name.c_str(); }
     bool is_reproducible() const override { return true; }
 
 private:
     WorkspaceBundle get_workspace_bundle(dt_byte* raw_ptr,
                                          const SizeArgs& args) const;
+    AlgoParam m_algo_param;
+    std::string m_name;
 };
 
 #if CUDA_VERSION >= 10000
@@ -578,7 +599,7 @@ public:
     AlgoMatmul8x8x32 matmul8x8x32;
     AlgoBatchedMatmul batched_matmul;
     Algo1x1 a1x1;
-    AlgoInt8NCHW4DotProdImplicitGemm int8_nchw4_dotprod;
+    std::vector<AlgoInt8NCHW4DotProdImplicitGemm> int8_nchw4_dotprod;
     AlgoInt8CHWN4DotProdImplicitGemm int8_chwn4_dotprod;
 #if CUDA_VERSION >= 10000
     AlgoQUInt4x4x32WMMA wmma_quint4x4x32;
@@ -605,6 +626,7 @@ private:
     void fill_imma_algos();
 #endif
     void fill_cudnn_algos();
+    void fill_dp4a_algos();
 };
 
 }  // namespace cuda

dnn/src/cuda/conv_bias/cutlass_convolution_wrapper.cu

@@ -19,7 +19,6 @@
 #endif
 #include "src/common/opr_param_defs_enumv.cuh"
 #include "src/cuda/conv_bias/cutlass_convolution_wrapper.cuh"
-
 #pragma GCC diagnostic pop
 
 using namespace megdnn;
@@ -149,4 +148,130 @@ INST(true);
 INST(false);
 #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

dnn/src/cuda/conv_bias/cutlass_convolution_wrapper.cuh

@@ -37,6 +37,15 @@ void do_conv_bias_int8_implicit_gemm_imma_ncdiv32hw32(
         const GemmCoord& threadblock_shape, const GemmCoord& warp_shape,
         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 cuda
 }  // namespace megdnn

dnn/src/cuda/conv_bias/implicit_gemm_int8_nchw32_imma.cpp

@@ -57,30 +57,16 @@ bool ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::is_available(
     // only support sm_75 or later, platform should have tensorcore int8
     // support
     available &= is_compute_capability_required(7, 5);
-    if (fh == 1 && fw == 1)
-        return available;
-    // 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;
+    // FIXME: too large filter size is not supported now 
+    available &= fh * fw <= 49;
     return available;
 }
 
 WorkspaceBundle
 ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::get_workspace_bundle(
         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();
-    if (fh == 1 && fw == 1) {
-        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}};
+    return WorkspaceBundle{raw_ptr, {ws_filter}};
 }
 
 size_t
@@ -148,9 +134,9 @@ void ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::exec(
                 false>(args.src_tensor->compatible_ptr<int8_t>(),
                        reinterpret_cast<int8_t*>(ws_filter),
                        args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr,
-                       args.dst_tensor->compatible_ptr<int8_t>(),
-                       nullptr, kern_param, nonlinear_mode,
-                       alpha, beta, gamma, dst_scale,
+                       args.dst_tensor->compatible_ptr<int8_t>(), nullptr,
+                       kern_param, nonlinear_mode, alpha, beta, gamma,
+                       dst_scale,
                        cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m,
                                                   m_algo_param.threadblock_n,
                                                   m_algo_param.threadblock_k},
@@ -159,14 +145,12 @@ void ConvBiasForwardImpl::AlgoInt8NCHW32IMMAImplicitGemm::exec(
                                                   m_algo_param.warp_k},
                        stream);
     } else {
-        auto workspace = ws.get(1);
         cutlass_wrapper::do_conv_bias_int8_implicit_gemm_imma_ncdiv32hw32<true>(
                 args.src_tensor->compatible_ptr<int8_t>(),
                 reinterpret_cast<int8_t*>(ws_filter),
                 args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr,
-                args.dst_tensor->compatible_ptr<int8_t>(),
-                reinterpret_cast<int*>(workspace), kern_param, nonlinear_mode,
-                alpha, beta, gamma, dst_scale,
+                args.dst_tensor->compatible_ptr<int8_t>(), nullptr, kern_param,
+                nonlinear_mode, alpha, beta, gamma, dst_scale,
                 cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m,
                                            m_algo_param.threadblock_n,
                                            m_algo_param.threadblock_k},

dnn/src/cuda/conv_bias/implicit_gemm_int8_nchw4_dp4a.cpp

@@ -11,7 +11,8 @@
 
 #include "./algo.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 cuda;
@@ -53,21 +54,16 @@ bool ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::is_available(
     // only support sm_61 or later, platform should have fast native int8
     // support
     available &= is_compute_capability_required(6, 1);
+    // FIXME: too large filter size is not supported now 
+    available &= fh * fw <= 49;
     return available;
 }
 
 WorkspaceBundle
 ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::get_workspace_bundle(
         dt_byte* raw_ptr, const SizeArgs& args) const {
-    size_t ws_size_src = args.src_layout->span().dist_byte();
-    size_t ws_size_filter = args.filter_layout->span().dist_byte();
-    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 ws_filter = args.filter_layout->span().dist_byte();
+    return WorkspaceBundle{raw_ptr, {ws_filter}};
 }
 
 size_t
@@ -84,27 +80,9 @@ void ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::exec(
     UNPACK_CONV_BIAS_NCHW4_PARAM(*(args.src_layout), fm, *(args.dst_layout),
                                  param);
     auto ws = get_workspace_bundle(args.workspace.raw_ptr, args);
-    auto ws_src = ws.get(0);
-    auto ws_filter = ws.get(1);
-    auto ws_dst = ws.get(2);
+    auto ws_filter = ws.get(0);
     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
     {
         TensorLayout src{{co, ci / 4 * fh * fw}, dtype::Int32()};
@@ -136,53 +114,42 @@ void ConvBiasForwardImpl::AlgoInt8NCHW4DotProdImplicitGemm::exec(
           dst_scale = args.dst_layout->dtype.param<dtype::QuantizedS8>().scale;
     float alpha = src_scale * filter_scale / dst_scale,
           beta = bias_scale / dst_scale;
-    
-    // process z
     int8_t* z_dev_ptr = nullptr;
-    float gamma = 1.f;
+    float gamma = 0.0;
     if (args.z_layout->ndim > 0) {
-        auto ws_z = ws.get(3);
-
-        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);
+        z_dev_ptr = args.z_tensor->compatible_ptr<int8_t>();
         float z_scale = args.z_layout->dtype.param<dtype::QuantizedS8>().scale;
         gamma = z_scale / dst_scale;
     }
-
-    convolution::PerChannelBiasVisitor bias_visitor;
-    bias_visitor.bias = args.bias_tensor->compatible_ptr<int32_t>();
-    ConvBiasForwardImpl::AlgoInt8CHWN4DotProdImplicitGemm::
-            dispatch_nonlinear_mode<convolution::PerChannelBiasVisitor>(
-                    reinterpret_cast<int8_t*>(ws_src),
-                    reinterpret_cast<int8_t*>(ws_filter), bias_visitor,
-                    z_dev_ptr, reinterpret_cast<int8_t*>(ws_dst), kern_param,
-                    alpha, beta, gamma, dst_scale, stream, param.nonlineMode);
-
-    // reformat chwn4 to nchw4
-    {
-        TensorLayout src{{co / 4 * ho * wo, n}, 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 = ws_dst;
-        ts_src.layout = src;
-        ts_dst.raw_ptr = args.dst_tensor->raw_ptr;
-        ts_dst.layout = dst;
-        auto&& transpose =
-                args.opr->handle()->create_operator<RelayoutForward>();
-        transpose->exec(ts_src, ts_dst);
+    uint32_t nonlinear_mode = static_cast<uint32_t>(param.nonlineMode);
+    if (fh == 1 && fw == 1) {
+        cutlass_wrapper::do_conv_bias_int8_implicit_gemm_dp4a_ncdiv4hw4<false>(
+                args.src_tensor->compatible_ptr<int8_t>(),
+                reinterpret_cast<int8_t*>(ws_filter),
+                args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr,
+                args.dst_tensor->compatible_ptr<int8_t>(), nullptr, kern_param,
+                nonlinear_mode, alpha, beta, gamma, dst_scale,
+                cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m,
+                                           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);
+    } else {
+        cutlass_wrapper::do_conv_bias_int8_implicit_gemm_dp4a_ncdiv4hw4<true>(
+                args.src_tensor->compatible_ptr<int8_t>(),
+                reinterpret_cast<int8_t*>(ws_filter),
+                args.bias_tensor->compatible_ptr<int32_t>(), z_dev_ptr,
+                args.dst_tensor->compatible_ptr<int8_t>(), nullptr, kern_param,
+                nonlinear_mode, alpha, beta, gamma, dst_scale,
+                cutlass_wrapper::GemmCoord{m_algo_param.threadblock_m,
+                                           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);
     }
 }
 

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

 /**
  * \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")
  *
  * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.

dnn/test/common/conv_bias.cpp

@@ -748,7 +748,7 @@ void check_conv_bias(DType src_dtype, DType filter_dtype, DType bias_dtype,
         bias_rng = std::make_unique<UniformIntRNG>(-50, 50);
         checker.set_epsilon(1 + 1e-3)
                 .set_max_avg_error(1e-1)
-                .set_max_avg_biased_error(1e-1);
+                .set_max_avg_biased_error(1e-3);
     } else if (src_dtype.enumv() == DTypeEnum::Float16) {
         rng = std::make_unique<NormalRNG>(2.f);
         megdnn_assert(bias_dtype.enumv() == DTypeEnum::Float16);

dnn/test/cuda/conv_bias_int8.cpp

@@ -18,8 +18,6 @@
 #include "test/cuda/fixture.h"
 #include "test/cuda/utils.h"
 
-
-#define MEGDNN_WITH_BENCHMARK 1
 #define V1(x) #x
 #define V(x) V1(x)
 
@@ -1228,8 +1226,17 @@ TEST_F(CUDA, BENCHMARK_CUTLASS_CONV_BIAS_INT8_NCHW32) {
             param::ConvBias::Format::NCHW32);
 }
 #endif
-#endif
 
+TEST_F(CUDA, BENCHMARK_CUTLASS_CONV_BIAS_INT8_NCHW4) {
+    require_compute_capability(6, 1);
+    benchmark_target_algo(
+            handle_cuda(), get_resnet50_bench_args(64),
+            dtype::QuantizedS8{1.2f}, dtype::QuantizedS8{1.3f},
+            dtype::QuantizedS32{1.2f * 1.3f}, dtype::QuantizedS8{1.0f},
+            "INT8_NCHW4_DOTPROD_IMPLICIT_GEMM",
+            param::ConvBias::Format::NCHW4);
+}
+#endif
 }  // namespace test
 }  // namespace megdnn
 

src/opr/test/dnn/convolution.cpp

@@ -2031,4 +2031,96 @@ TEST(TestOprDNN, HeuristicReproducible) {
 #undef get_shp
 }
 
+#if MGB_CUDA
+TEST(TestOprDNN, ConvolutionMultiCompNode) {
+    REQUIRE_GPU(2);
+    auto cn0 = CompNode::load("gpu0:0"), cn1 = CompNode::load("gpu0:1");
+    cn0.activate();
+    auto&& prop = CompNodeEnv::from_comp_node(cn0).cuda_env().device_prop;
+    auto sm_ver = prop.major * 10 + prop.minor;
+    if (sm_ver < 61) {
+        printf("This testcast ignored due to insufficient cuda cap(got: %d, "
+               "expected: %d)\n",
+               sm_ver, 61);
+        return;
+    }
+
+    HostTensorGenerator<dtype::Int8> gen;
+    auto mkvar = [&gen](const char* name, const TensorShape& shp,
+                     const DType& dtype,
+                     std::shared_ptr<ComputingGraph> graph,
+                     const CompNode& cn) {
+        return opr::TypeCvt::make(
+                opr::Host2DeviceCopy::make(*graph, gen(shp, cn)).rename(name),
+                dtype);
+    };
+    auto mkcvar = [&gen](const char* name, const TensorShape& shp,
+                      const DType& dtype,
+                      std::shared_ptr<ComputingGraph> graph,
+                      const CompNode& cn) {
+        return opr::TypeCvt::make(
+                opr::SharedDeviceTensor::make(*graph, *gen(shp, cn))
+                        .rename(name),
+                dtype);
+    };
+
+    auto graph0 = ComputingGraph::make();
+    graph0->options().graph_opt_level = 0;
+    auto graph1 = ComputingGraph::make();
+    graph1->options().graph_opt_level = 0;
+    auto make_func = [&gen, &mkvar, &mkcvar](
+                             std::shared_ptr<ComputingGraph> graph,
+                             const CompNode& cn) {
+        using Policy = opr::ConvBias::ExecutionPolicy;
+        using S = Policy::Strategy;
+        auto x = mkvar("x", {64, 32, 28, 28, 4}, dtype::QuantizedS8(2.5f),
+                       graph, cn),
+             w1 = mkcvar("w1", {256, 32, 5, 5, 4}, dtype::QuantizedS8(2.5f),
+                         graph, cn),
+             b1 = mkcvar("b1", {1, 64, 1, 1, 4}, dtype::QuantizedS32(6.25f),
+                         graph, cn),
+             w2 = mkcvar("w2", {256, 64, 3, 3, 4}, dtype::QuantizedS8(2.5f),
+                         graph, cn),
+             b2 = mkcvar("b2", {1, 64, 1, 1, 4}, dtype::QuantizedS32(6.25f),
+                         graph, cn);
+        opr::ConvBias::Param param;
+        param.format = opr::ConvBias::Param::Format::NCHW4;
+        param.nonlineMode = opr::ConvBias::Param::NonlineMode::RELU;
+        param.stride_h = param.stride_w = 2;
+        param.pad_h = param.pad_w = 2;
+        Policy policy;
+        policy.strategy = S::PROFILE;
+
+        auto y = opr::ConvBias::make(
+                x, w1, b1, param, policy,
+                OperatorNodeConfig{dtype::QuantizedS8(2.5f)});
+        param.stride_h = param.stride_w = 1;
+        param.pad_h = param.pad_w = 1;
+        y = opr::ConvBias::make(y, w2, b2, param, policy,
+                                OperatorNodeConfig{dtype::QuantizedS8(2.5f)});
+        return y;
+    };
+    auto y0 = make_func(graph0, cn0);
+    auto y1 = make_func(graph1, cn1);
+    HostTensorND host_y0, host_y1;
+    auto func0 = graph0->compile({make_callback_copy(y0, host_y0)});
+    auto func1 = graph1->compile({make_callback_copy(y1, host_y1)});
+
+     auto worker = [&func0, &func1](int wid) {
+        static int const iter_num = 1000;
+         if (wid == 0) {
+             for (int i = 0; i < iter_num; ++i)
+                 func0->execute();
+         } else {
+             for (int i = 0; i < iter_num; ++i)
+                 func1->execute();
+         }
+     };
+     std::thread worker0(worker, 0);
+     std::thread worker1(worker, 1);
+     worker0.join();
+     worker1.join();
+}
+#endif
+
 // vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}