#ifndef CAFFE2_OPERATORS_FUSED_ROWWISE_8BIT_CONVERSION_OPS_H_
|
#define CAFFE2_OPERATORS_FUSED_ROWWISE_8BIT_CONVERSION_OPS_H_
|
|
#include "caffe2/core/context.h"
|
#include "caffe2/core/export_caffe2_op_to_c10.h"
|
#include "caffe2/core/logging.h"
|
#include "caffe2/core/operator.h"
|
#include "caffe2/operators/reducer_functors.h"
|
#include "caffe2/utils/eigen_utils.h"
|
#include "caffe2/utils/math.h"
|
|
C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(Fused8BitRowwiseQuantizedToFloat);
|
|
namespace caffe2 {
|
|
#define IS_LITTLE_ENDIAN \
|
[] { \
|
const int32_t kValue = 1; \
|
return reinterpret_cast<const uint8_t*>(&kValue)[0] == 1; \
|
}()
|
|
template <
|
typename T,
|
void (*convert)(float* dst, const T* src, size_t N),
|
class Context>
|
class FloatToFused8BitRowwiseQuantizedOp : public Operator<Context> {
|
public:
|
static constexpr float kEpsilon = 1e-8f;
|
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
USE_SIMPLE_CTOR_DTOR(FloatToFused8BitRowwiseQuantizedOp)
|
|
bool RunOnDevice() override {
|
CAFFE_ENFORCE(IS_LITTLE_ENDIAN, "Unsupported endianness");
|
|
const auto& input = Input(DATA_FLOAT);
|
|
const auto input_rows = input.size(0);
|
const auto input_columns = input.size(1);
|
CAFFE_ENFORCE_EQ(input.dim(), 2, "Expect input to be a matrix");
|
|
// The "fused" representation stores the scale and bias with the row-wise
|
// quantized data in one tensor. Since we quantize with 8 bits (1 byte) and
|
// represent the scale and bias with 32-bit floats, we'll use the last 8
|
// bytes of each row for scale (4 bytes) and bias (4 bytes).
|
// | ... int8 data ... | scale | bias |
|
// | number_of_columns | 4B | 4B |
|
const std::vector<int64_t> output_dimensions = {input_rows,
|
input_columns + 8};
|
auto* output = Output(
|
DATA_FUSED_SCALE_BIAS_INT8, output_dimensions, at::dtype<uint8_t>());
|
|
const auto* input_data = input.template data<T>();
|
auto* output_data = output->template mutable_data<uint8_t>();
|
const auto output_columns = output->size(1);
|
|
if (!std::is_same<T, float>::value && !std::is_same<T, at::Half>::value) {
|
CAFFE_THROW("Unsupported data type");
|
}
|
|
vector<float> tmp;
|
tmp.resize(input_columns, 0.0);
|
|
for (size_t row = 0; row < input_rows; ++row) {
|
convert(tmp.data(), input_data + row * input_columns, input_columns);
|
ConstEigenVectorArrayMap<float> input_row(tmp.data(), input_columns);
|
uint8_t* output_row = output_data + row * output_columns;
|
EigenVectorArrayMap<uint8_t> output_row_values(output_row, input_columns);
|
EigenVectorArrayMap<float> output_row_scale_bias(
|
reinterpret_cast<float*>(output_row + input_columns), 2);
|
|
const float minimum_element = input_row.minCoeff();
|
const float maximum_element = input_row.maxCoeff();
|
const float range = maximum_element - minimum_element;
|
|
output_row_scale_bias(0) = range / 255.0f;
|
output_row_scale_bias(1) = minimum_element;
|
const auto inverse_scale = 255.0f / (range + kEpsilon);
|
output_row_values = ((input_row - minimum_element) * inverse_scale)
|
.round()
|
.cast<uint8_t>();
|
}
|
|
return true;
|
}
|
|
private:
|
INPUT_TAGS(DATA_FLOAT);
|
OUTPUT_TAGS(DATA_FUSED_SCALE_BIAS_INT8);
|
};
|
|
template <
|
typename T,
|
void (*convert)(T* dst, const float* src, size_t N),
|
class Context>
|
class Fused8BitRowwiseQuantizedToFloatOp : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
USE_SIMPLE_CTOR_DTOR(Fused8BitRowwiseQuantizedToFloatOp)
|
|
bool RunOnDevice() override {
|
CAFFE_ENFORCE(IS_LITTLE_ENDIAN, "Unsupported endianness");
|
|
const auto& input = Input(DATA_FUSED_SCALE_BIAS_INT8);
|
|
const auto input_rows = input.size(0);
|
const auto input_columns = input.size(1);
|
CAFFE_ENFORCE_EQ(input.dim(), 2, "Expect input to be a matrix");
|
|
// The last 8 bytes per row are the scale and the bias. The rest of
|
// input_columns is the number of values in the original row.
|
const std::vector<int64_t> output_dimensions = {input_rows,
|
input_columns - 8};
|
auto* output = Output(DATA_FLOAT, output_dimensions, at::dtype<T>());
|
const auto output_columns = output->size(1);
|
|
const auto* input_data = input.template data<uint8_t>();
|
T* output_data = output->template mutable_data<T>();
|
|
vector<float> tmp;
|
tmp.resize(input_columns, 0.0);
|
|
for (size_t row = 0; row < input_rows; ++row) {
|
const uint8_t* input_row = input_data + row * input_columns;
|
ConstEigenVectorArrayMap<uint8_t> input_row_values(
|
input_row, output_columns);
|
ConstEigenVectorArrayMap<float> input_row_scale_bias(
|
reinterpret_cast<const float*>(input_row + output_columns), 2);
|
|
EigenVectorArrayMap<float> output_row(tmp.data(), output_columns);
|
output_row = input_row_values.cast<float>() * input_row_scale_bias(0) +
|
input_row_scale_bias(1);
|
|
convert(output_data + row * output_columns, tmp.data(), output_columns);
|
}
|
return true;
|
}
|
|
private:
|
INPUT_TAGS(DATA_FUSED_SCALE_BIAS_INT8);
|
OUTPUT_TAGS(DATA_FLOAT);
|
};
|
|
#undef IS_LITTLE_ENDIAN
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_FUSED_ROWWISE_8BIT_CONVERSION_OPS_H_
|
