src/operator/quantization/quantized_conv.cu - master - githubmirror/incubator-mxnet

src/operator/quantization/quantized_conv.cu
296 строк13 KB
Zhenghui Jin
[v2.0][LICENSE] Port #20493 (#20608)
28 сен 2021, 03:10
28 сен 2021, 03:10a720b15
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file quantized_conv.cu
 * \brief
 * \author Ziheng Jiang, Jun Wu
 */
#include "../nn/convolution-inl.h"
#include "./quantization_utils.h"
#include "../tensor/matrix_op-inl.h"

namespace mxnet {
namespace op {

// value + bias_value * (range1 / limit_range1) * (limit_range2 / range2)
struct QuantizedBiasAddKernel {
  MSHADOW_XINLINE static void Map(int i,
                                  size_t bias_size,
                                  int32_t* out,
                                  const int8_t* bias,
                                  const float* min_out,
                                  const float* max_out,
                                  const float* min_bias,
                                  const float* max_bias,
                                  const size_t spatial_size) {
    using mshadow::red::limits::MaxValue;
    using mshadow::red::limits::MinValue;
    float float_for_one_out_quant =
        MaxAbs(*min_out, *max_out) / static_cast<double>(MaxValue<int32_t>());
    float float_for_one_bias_quant =
        MaxAbs(*min_bias, *max_bias) / static_cast<double>(MaxValue<int8_t>());
    const size_t channel_id = (i / spatial_size) % bias_size;
    out[i] = (out[i] * float_for_one_out_quant + bias[channel_id] * float_for_one_bias_quant) /
             float_for_one_out_quant;
  }
};

#if MXNET_USE_CUDNN == 1 && CUDA_VERSION >= 8000
STATIC_ASSERT_CUDNN_VERSION_GE(6000);
template <typename SrcType, typename DstType, typename CmpType>
class QuantizedCuDNNConvOp {
 public:
  QuantizedCuDNNConvOp() {
    CUDNN_CALL(cudnnCreateConvolutionDescriptor(&conv_desc_));
    CUDNN_CALL(cudnnCreateTensorDescriptor(&data_desc_));
    CUDNN_CALL(cudnnCreateTensorDescriptor(&out_desc_));
    CUDNN_CALL(cudnnCreateFilterDescriptor(&filter_desc_));
  }

  void Init(const ConvolutionParam& param,
            const OpContext& ctx,
            const mxnet::ShapeVector& in_shape,
            const mxnet::ShapeVector& out_shape) {
    param_ = param;
    CHECK_EQ(param_.kernel.ndim(), 2U)
        << "QuantizedCuDNNConvOp only supports 2D convolution for now";
    if (param_.layout.has_value()) {
      CHECK_EQ(param_.layout.value(), mshadow::kNCHW)
          << "QuantizedConvOp only supports NCHW for now";
    }
    if (param_.stride.ndim() == 0U)
      param_.stride = mshadow::Shape2(1, 1);
    if (param_.dilate.ndim() == 0U)
      param_.dilate = mshadow::Shape2(1, 1);
    if (param_.pad.ndim() == 0U)
      param_.pad = mshadow::Shape2(0, 0);
    N = 0, H = 2, W = 3, C = 1;
    src_type_ = mshadow::DataType<SrcType>::kCudnnFlag;
    dst_type_ = mshadow::DataType<DstType>::kCudnnFlag;
    cmp_type_ = mshadow::DataType<CmpType>::kCudnnFlag;
    algo_     = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;
    format_   = CUDNN_TENSOR_NHWC;
    InitDescriptors(in_shape, out_shape);
    GetTempSize(ctx);
  }

  ~QuantizedCuDNNConvOp() {
    CUDNN_CALL(cudnnDestroyFilterDescriptor(filter_desc_));
    CUDNN_CALL(cudnnDestroyTensorDescriptor(data_desc_));
    CUDNN_CALL(cudnnDestroyTensorDescriptor(out_desc_));
    CUDNN_CALL(cudnnDestroyConvolutionDescriptor(conv_desc_));
  }

  void Forward(const OpContext& ctx,
               const std::vector<TBlob>& in_data,
               const std::vector<OpReqType>& req,
               const std::vector<TBlob>& out_data) {
    CHECK_EQ(param_.kernel.ndim(), 2U)
        << "QuantizedCuDNNConvOp only supports 2D convolution for now";
    using namespace mshadow;
    CHECK_EQ(in_data.size(), param_.no_bias ? 6U : 9U);
    CHECK_EQ(out_data.size(), 3U);
    Stream<gpu>* s = ctx.get_stream<gpu>();
    CHECK_EQ(s->dnn_handle_ownership_, Stream<gpu>::OwnHandle);

    const TBlob& data           = in_data[0];
    const TBlob& filter         = in_data[1];
    const TBlob& out            = out_data[0];
    const mxnet::TShape& dshape = data.shape_;
    const mxnet::TShape& fshape = filter.shape_;
    const mxnet::TShape& oshape = out.shape_;

    // allocate workspace
    const int dev_id   = ctx.run_ctx.ctx.dev_id;
    const int dev_mask = gpu::kDevMask;
    if (!param_.layout.has_value() || param_.layout.value() == mshadow::kNCHW) {
      const size_t data_size   = dshape.Size();
      const size_t weight_size = fshape.Size();
      const size_t output_size = oshape.Size();
      size_t total_temp_bytes  = (workspace_ + data_size + weight_size) * sizeof(SrcType) +
                                output_size * (sizeof(DstType) + sizeof(int32_t));
      Tensor<gpu, 1, char> temp_space =
          ctx.requested[0].get_space_typed<gpu, 1, char>(mshadow::Shape1(total_temp_bytes), s);
      char* temp_dptr = temp_space.dptr_;
      TBlob data_(reinterpret_cast<SrcType*>(temp_dptr),
                  mxnet::TShape({dshape[N], dshape[H], dshape[W], dshape[C]}),
                  dev_mask,
                  DataType<SrcType>::kFlag,
                  dev_id);
      temp_dptr += data_size * sizeof(SrcType);
      TBlob filter_(reinterpret_cast<SrcType*>(temp_dptr),
                    mxnet::TShape({fshape[N], fshape[H], fshape[W], fshape[C]}),
                    dev_mask,
                    DataType<SrcType>::kFlag,
                    dev_id);
      temp_dptr += weight_size * sizeof(SrcType);

      // input:  [NCHW] => [NHWC](batch, in_height, in_width, in_channels)
      // filter: [NCHW] => [NHWC](out_channels, filter_height, filter_width, in_channels)
      TransposeImpl<gpu>(ctx.run_ctx, data, data_, mxnet::TShape({N, H, W, C}));
      TransposeImpl<gpu>(ctx.run_ctx, filter, filter_, mxnet::TShape({N, H, W, C}));
      TBlob out_(reinterpret_cast<DstType*>(temp_dptr),
                 mxnet::TShape({oshape[N], oshape[H], oshape[W], oshape[C]}),
                 dev_mask,
                 DataType<DstType>::kFlag,
                 dev_id);
      temp_dptr += output_size * sizeof(DstType);
      TBlob out_tcast(reinterpret_cast<int32_t*>(temp_dptr),
                      mxnet::TShape({oshape[N], oshape[H], oshape[W], oshape[C]}),
                      dev_mask,
                      DataType<int32_t>::kFlag,
                      dev_id);
      temp_dptr += output_size * sizeof(int32_t);
      // input:  [NHWC](batch, in_height, in_width, in_channels)
      // filter: [HWNC](out_channels, filter_height, filter_width, in_channels)
      // output: [NHWC](batch, out_height, out_width, out_channels)

      CUDNN_CALL(cudnnConvolutionForward(s->dnn_handle_,
                                         &alpha_,
                                         data_desc_,
                                         data_.dptr_,
                                         filter_desc_,
                                         filter_.dptr_,
                                         conv_desc_,
                                         algo_,
                                         temp_dptr,
                                         workspace_byte_,
                                         &beta_,
                                         out_desc_,
                                         out_.dptr_));

      Tensor<gpu, 1, DstType> out_tensor       = out_.FlatTo1D<gpu, DstType>(s);
      Tensor<gpu, 1, int32_t> out_tcast_tensor = out_tcast.FlatTo1D<gpu, int32_t>(s);
      Assign(out_tcast_tensor, kWriteTo, mshadow::expr::tcast<int32_t>(out_tensor));
      // output: [NHWC](batch, out_height, out_width, out_channels) => [NCHW]
      TransposeImpl<gpu>(ctx.run_ctx, out_tcast, out, mxnet::TShape({0, 3, 1, 2}));
    } else {
      LOG(FATAL) << "quantized_conv only supports NCHW for now";
    }

    // calculate the min/max range for out_data as it's a multiplication
    // of in_data[0] and in_data[1]. Need to rescale the min/max range of out_data
    // based on the min/max ranges of in_data[0] and in_data[1].
    const size_t num_inputs = param_.no_bias ? 2 : 3;
    mxnet_op::Kernel<QuantizationRangeForS8S8MultiplicationStruct, gpu>::Launch(
        s,
        1,
        out_data[1].dptr<float>(),
        out_data[2].dptr<float>(),
        in_data[num_inputs].dptr<float>(),
        in_data[num_inputs + 1].dptr<float>(),
        in_data[num_inputs + 2].dptr<float>(),
        in_data[num_inputs + 3].dptr<float>());

    if (!param_.no_bias) {
      if (param_.layout.has_value()) {
        CHECK_EQ(param_.layout.value(), mshadow::kNCHW)
            << "quantized_conv only supports NCHW when there is a bias";
      }
      const TBlob& bias = in_data[2];
      mxnet_op::Kernel<QuantizedBiasAddKernel, gpu>::Launch(s,
                                                            out.Size(),
                                                            bias.Size(),
                                                            out.dptr<int32_t>(),
                                                            bias.dptr<int8_t>(),
                                                            out_data[1].dptr<float>(),
                                                            out_data[2].dptr<float>(),
                                                            in_data[7].dptr<float>(),
                                                            in_data[8].dptr<float>(),
                                                            oshape[2] * oshape[3]);
    }
  }

  void InitDescriptors(const mxnet::ShapeVector& in_shape, const mxnet::ShapeVector& out_shape) {
    const mxnet::TShape& dshape = in_shape[0];
    const mxnet::TShape& kshape = in_shape[1];
    const mxnet::TShape& oshape = out_shape[0];
    CUDNN_CALL(cudnnSetConvolution2dDescriptor(conv_desc_,
                                               param_.pad[0],
                                               param_.pad[1],
                                               param_.stride[0],
                                               param_.stride[1],
                                               1,
                                               1,
                                               CUDNN_CROSS_CORRELATION,
                                               cmp_type_));

    CUDNN_CALL(cudnnSetTensor4dDescriptor(
        data_desc_, format_, src_type_, dshape[N], dshape[C], dshape[H], dshape[W]));
    CUDNN_CALL(cudnnSetTensor4dDescriptor(
        out_desc_, format_, dst_type_, oshape[N], oshape[C], oshape[H], oshape[W]));
    CUDNN_CALL(cudnnSetFilter4dDescriptor(
        filter_desc_, src_type_, format_, kshape[N], kshape[C], kshape[H], kshape[W]));
  }

  void GetTempSize(const OpContext& ctx) {
    mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
    CUDNN_CALL(cudnnGetConvolutionForwardWorkspaceSize(
        s->dnn_handle_, data_desc_, filter_desc_, conv_desc_, out_desc_, algo_, &workspace_byte_));
    workspace_ = workspace_byte_ / sizeof(SrcType) + 1;
  }

 private:
  ConvolutionParam param_;
  size_t workspace_;
  size_t workspace_byte_;
  cudnnDataType_t src_type_;
  cudnnDataType_t dst_type_;
  cudnnDataType_t cmp_type_;
  cudnnTensorFormat_t format_;
  cudnnConvolutionDescriptor_t conv_desc_;
  cudnnTensorDescriptor_t data_desc_;
  cudnnFilterDescriptor_t filter_desc_;
  cudnnTensorDescriptor_t out_desc_;
  cudnnConvolutionFwdAlgo_t algo_;
  uint32_t N, H, W, C;
  float alpha_ = 1.0f;
  float beta_  = 0.0f;
};      // class QuantizedCuDNNConvOp
#endif  // MXNET_USE_CUDNN == 1 && CUDA_VERSION >= 8000

void QuantizedConvForwardGPU(const nnvm::NodeAttrs& attrs,
                             const OpContext& ctx,
                             const std::vector<TBlob>& inputs,
                             const std::vector<OpReqType>& req,
                             const std::vector<TBlob>& outputs) {
  const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
  CHECK_EQ(param.kernel.ndim(), 2U)
      << "QuantizedConvForward<gpu> only supports 2D convolution for now";
#if MXNET_USE_CUDNN == 1 && CUDA_VERSION >= 8000
  typedef QuantizedCuDNNConvOp<int8_t, float, int32_t> QuantizedConvOpInt8;
#if DMLC_CXX11_THREAD_LOCAL
  static thread_local QuantizedConvOpInt8 op;
#else
  static MX_THREAD_LOCAL QuantizedConvOpInt8 op;
#endif  // DMLC_CXX11_THREAD_LOCAL
  op.Init(param, ctx, {inputs[0].shape_, inputs[1].shape_}, {outputs[0].shape_});
  op.Forward(ctx, inputs, req, outputs);
#else
  LOG(FATAL) << "QuantizedConvForward<gpu> only supports cudnnConvolutionForward "
                "with CUDNN >= 6.0 and CUDA >= 8.0";
#endif  // MXNET_USE_CUDNN == 1 && CUDA_VERSION >= 8000
}

NNVM_REGISTER_OP(_contrib_quantized_conv)
    .set_attr<FCompute>("FCompute<gpu>", QuantizedConvForwardGPU);

}  // namespace op
}  // namespace mxnet