Caffe2 - C++ API
A deep learning, cross platform ML framework
batch_matmul_dnnlowp_op.h
1 
17 #pragma once
18 
19 #include "caffe2/operators/batch_matmul_op.h"
20 #include "caffe2/quantization/server/dnnlowp_op.h"
21 #include "fbgemm/Fbgemm.h"
22 
23 namespace caffe2 {
24 
25 template <typename T>
27  : public DNNLowPOp<T, BatchMatMulOp<CPUContext>> {
28  public:
29  BatchMatMulDNNLowPOp(const OperatorDef& operator_def, Workspace* ws);
30  bool RunOnDevice() override;
31 
32  USE_OPERATOR_FUNCTIONS(CPUContext);
33  USE_DNNLOWP_OPERATOR_BASE_FUNCTIONS(T, BatchMatMulOp<CPUContext>);
34 
35  private:
36  bool trans_a_;
37  bool trans_b_;
38  bool broadcast_{false};
39  bool is_B_constant_{false};
40 
41  std::vector<std::int8_t> B_quantized_;
42  std::vector<std::unique_ptr<fbgemm::PackBMatrix<std::int8_t>>> Bq_packed_;
43  std::vector<std::uint8_t> A_pack_buf_;
44  std::vector<std::int32_t> row_offsets_, column_offsets_;
45 
46  std::vector<dnnlowp::TensorQuantizationParams> B_qparams_;
47  std::vector<dnnlowp::RequantizationParams> requantization_params_;
48 
49  std::vector<std::int32_t> Y_int32_;
50  bool first_invocation_{true};
51 }; // BatchMatMulDNNLowPOp
52 
53 } // namespace caffe2
The CPU Context, representing the bare minimum of what a Context class in Caffe2 should implement...
Definition: context.h:40
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:47
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
A convenient base class for C2 operators with DNNLOWP engine.
Definition: dnnlowp_op.h:77