mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
a9b0a921d5
Summary: As GoogleTest `TEST` macro is non-compliant with it as well as `DEFINE_DISPATCH` All changes but the ones to `.clang-tidy` are generated using following script: ``` for i in `find . -type f -iname "*.c*" -or -iname "*.h"|xargs grep cppcoreguidelines-avoid-non-const-global-variables|cut -f1 -d:|sort|uniq`; do sed -i "/\/\/ NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)/d" $i; done ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/62008 Reviewed By: driazati, r-barnes Differential Revision: D29838584 Pulled By: malfet fbshipit-source-id: 1b2f8602c945bd4ce50a9bfdd204755556e31d13
81 lines
2.3 KiB
C++
81 lines
2.3 KiB
C++
#include "utility_dnnlowp_ops.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
template <typename T>
|
|
GatherDNNLowPOp<T>::GatherDNNLowPOp(
|
|
const OperatorDef& operator_def,
|
|
Workspace* ws)
|
|
: GatherOp<CPUContext>(operator_def, ws),
|
|
qfactory_(dnnlowp::GetQuantizationFactoryOf(this)) {}
|
|
|
|
template <typename T>
|
|
GatherDNNLowPOp<T>::~GatherDNNLowPOp() {
|
|
if (measure_quantization_error_) {
|
|
dnnlowp::ReportQuantizationError(this, quantization_error_stats_);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
bool GatherDNNLowPOp<T>::RunOnDevice() {
|
|
using namespace dnnlowp;
|
|
|
|
if (!arguments_parsed_) {
|
|
dnnlowp::ParseDNNLowPOperatorArguments(
|
|
this, &dequantize_output_, &measure_quantization_error_);
|
|
arguments_parsed_ = true;
|
|
}
|
|
|
|
if (!InputIsType<int8::Int8TensorCPU>(DATA)) {
|
|
if (dequantize_output_) {
|
|
return GatherOp<CPUContext>::RunOnDevice();
|
|
} else {
|
|
// If input or output is float, delegate to fp32 op
|
|
Fp32Op_()->DequantizeInput();
|
|
// dequantize input if it's not already float
|
|
if (!Fp32Op_()->Get()->RunOnDevice()) {
|
|
return false;
|
|
}
|
|
|
|
int8::Int8TensorCPU* output =
|
|
Outputs()[0]->template GetMutable<int8::Int8TensorCPU>();
|
|
|
|
output->t.ResizeLike(*Fp32Op_()->Get()->Output(0));
|
|
T* out_data = output->t.template mutable_data<T>();
|
|
|
|
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
|
|
TensorQuantizationParams out_qparams;
|
|
if (HasStaticQuantization(this)) {
|
|
out_qparams = GetStaticQuantizationParamsOf(this, 0);
|
|
} else {
|
|
out_qparams = Fp32Op_()->GetOutputQuantizationParams(qfactory_.get());
|
|
}
|
|
|
|
fbgemm::Quantize<T>(
|
|
static_cast<const float*>(Fp32Op_()->Get()->Output(0)->raw_data()),
|
|
out_data,
|
|
output->t.numel(),
|
|
out_qparams);
|
|
|
|
PropagateOutputTensorQuantizationParams(this, 0, out_qparams);
|
|
}
|
|
} else {
|
|
DispatchHelper<TensorTypes<int32_t, int64_t>>::call(this, Input(INDICES));
|
|
|
|
TensorQuantizationParams in_qparams =
|
|
GetInputTensorQuantizationParamsOf(this, 0, qfactory_.get());
|
|
|
|
PropagateOutputTensorQuantizationParams(this, 0, in_qparams);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
REGISTER_CPU_OPERATOR_WITH_ENGINE(Gather, DNNLOWP, GatherDNNLowPOp<uint8_t>);
|
|
REGISTER_CPU_OPERATOR_WITH_ENGINE(
|
|
Int8Gather,
|
|
DNNLOWP,
|
|
GatherDNNLowPOp<uint8_t>);
|
|
|
|
} // namespace caffe2
|