quantize_mkldnn_op.cc 4.1 KB
Newer Older
X
xiaoli.liu@intel.com 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2016 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. */

15 16
#include "paddle/fluid/operators/quantize_op.h"

17
#include "paddle/fluid/framework/data_layout_transform.h"
X
xiaoli.liu@intel.com 已提交
18 19
#include "paddle/fluid/framework/tensor.h"
#include "paddle/fluid/platform/mkldnn_helper.h"
X
xiaoli.liu@intel.com 已提交
20
#include "paddle/fluid/platform/mkldnn_reuse.h"
X
xiaoli.liu@intel.com 已提交
21 22 23 24

namespace paddle {
namespace operators {

25 26 27
using dnnl::memory;
using dnnl::primitive;
using dnnl::reorder;
X
xiaoli.liu@intel.com 已提交
28
using platform::to_void_cast;
29
using Tensor = phi::DenseTensor;
30
using dnnl::stream;
31
using framework::DataLayout;
X
xiaoli.liu@intel.com 已提交
32 33 34 35 36

template <typename T>
class QuantOpKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
37 38
    auto* x = ctx.Input<phi::DenseTensor>("Input");
    auto* out = ctx.Output<phi::DenseTensor>("Output");
39 40 41 42 43 44

    const auto quantization_scale = ctx.Attr<float>("Scale");
    const auto quantization_shift = ctx.Attr<float>("Shift");
    const bool with_scale = quantization_scale != 1.0f;
    const bool with_shift = quantization_shift != 0.0f;

45 46
    PADDLE_ENFORCE_NE(quantization_scale,
                      0.0f,
47 48
                      platform::errors::InvalidArgument(
                          "Quantization scale must be different than 0.0f"));
49 50 51 52 53
    PADDLE_ENFORCE(quantization_shift <= 255 && quantization_shift >= 0,
                   platform::errors::InvalidArgument(
                       "Quantization shift must be lower or equal to ",
                       "255 and greater or equal to 0, but got %f",
                       quantization_shift));
54

X
xiaoli.liu@intel.com 已提交
55 56 57
    auto& dev_ctx =
        ctx.template device_context<platform::MKLDNNDeviceContext>();

58
    auto x_tz = phi::vectorize<int64_t>(x->dims());
X
xiaoli.liu@intel.com 已提交
59

60 61
    const bool is_negative_input = ctx.Attr<bool>("is_negative_input");
    const bool bfloat16 = ctx.Attr<bool>("bfloat16");
X
xiaoli.liu@intel.com 已提交
62

63 64
    dnnl::primitive_attr attrs;
    static constexpr int32_t mask = 0;
65

66 67 68
    if (with_scale) {
      attrs.set_output_scales(mask, {quantization_scale});
    }
69 70

    if (with_shift) {
71 72
      attrs.set_zero_points(
          DNNL_ARG_DST, mask, {static_cast<int32_t>(quantization_shift)});
73 74
    }

75 76 77
    framework::proto::VarType::Type x_paddle_dtype =
        framework::TransToProtoVarType(x->dtype());
    framework::proto::VarType::Type out_paddle_dtype;
78 79

    if (bfloat16) {
80
      out_paddle_dtype = framework::proto::VarType::BF16;
81
    } else if (is_negative_input && !with_shift) {
82
      out_paddle_dtype = framework::proto::VarType::INT8;
X
xiaoli.liu@intel.com 已提交
83
    } else {
84
      out_paddle_dtype = framework::proto::VarType::UINT8;
X
xiaoli.liu@intel.com 已提交
85
    }
86 87

    platform::ReorderMKLDNNHandler reorder_handler(
88 89 90 91 92
        x_tz,
        x_paddle_dtype,
        framework::ToMKLDNNDataType(x_paddle_dtype),
        out_paddle_dtype,
        framework::ToMKLDNNDataType(out_paddle_dtype),
93 94 95 96 97 98 99 100 101
        dev_ctx.GetEngine());

    auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
        x->mem_desc(), platform::to_void_cast(x->data<T>()));
    auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
        out, x->mem_desc(), dev_ctx.GetPlace());

    auto reorder_p = reorder_handler.AcquireReorder(
        reorder_dst_memory_p, reorder_src_memory_p, attrs);
102

103
    auto& astream = platform::MKLDNNDeviceContext::tls().get_stream();
104
    reorder_p->execute(astream, *reorder_src_memory_p, *reorder_dst_memory_p);
105
    astream.wait();
A
Adam 已提交
106

107
    out->set_mem_desc(reorder_dst_memory_p->get_desc());
X
xiaoli.liu@intel.com 已提交
108 109 110 111 112 113
  }
};
}  // namespace operators
}  // namespace paddle
namespace ops = paddle::operators;

114 115 116
REGISTER_OP_KERNEL(quantize,
                   MKLDNN,
                   ::paddle::platform::CPUPlace,
X
xiaoli.liu@intel.com 已提交
117
                   ops::QuantOpKernel<float>);