optimize FloatToFused8BitRowwiseQuantized and Fused8BitRowwiseQuantizedToFloat (pytorch#31470)

Summary: Pull Request resolved: pytorch#31470 Optimize performance of these two operators. Additionally use nearbyint instead of round to be consistent with 4-bit embedding table quantization. Reviewed By: hyuen Differential Revision: D19072103 fbshipit-source-id: efe96f14aeff7958cceb453ed625d3fd693891ff

Author: jspark1105

caffe2/operators/fused_rowwise_8bit_conversion_ops.cc

@@ -4,10 +4,6 @@
 namespace caffe2 {
 
 namespace {
-void convertfp32fp32(float* dst, const float* src, size_t N) {
- memcpy(dst, src, sizeof(float) * N);
-}
-
 void convertfp16fp32(float* dst, const at::Half* src, size_t N) {
  for (size_t i = 0; i < N; i++) {
  dst[i] = src[i];
@@ -23,7 +19,11 @@ void convertfp32fp16(at::Half* dst, const float* src, size_t N) {
 
 REGISTER_CPU_OPERATOR(
  FloatToFused8BitRowwiseQuantized,
- FloatToFused8BitRowwiseQuantizedOp<float, convertfp32fp32, CPUContext>);
+ FloatToFused8BitRowwiseQuantizedOp<
+ float,
+ nullptr,
+ false, /* HAS_CONVERT */
+ CPUContext>);
 OPERATOR_SCHEMA(FloatToFused8BitRowwiseQuantized)
  .NumInputs(1)
  .NumOutputs(1)
@@ -52,7 +52,11 @@ NO_GRADIENT(FloatToFused8BitRowwiseQuantized);
 
 REGISTER_CPU_OPERATOR(
  HalfFloatToFused8BitRowwiseQuantized,
- FloatToFused8BitRowwiseQuantizedOp<at::Half, convertfp16fp32, CPUContext>);
+ FloatToFused8BitRowwiseQuantizedOp<
+ at::Half,
+ convertfp16fp32,
+ true, /* HAS_CONVERT*/
+ CPUContext>);
 OPERATOR_SCHEMA(HalfFloatToFused8BitRowwiseQuantized)
  .NumInputs(1)
  .NumOutputs(1)
@@ -81,7 +85,11 @@ NO_GRADIENT(HalfFloatToFused8BitRowwiseQuantized);
 
 REGISTER_CPU_OPERATOR(
  Fused8BitRowwiseQuantizedToFloat,
- Fused8BitRowwiseQuantizedToFloatOp<float, convertfp32fp32, CPUContext>);
+ Fused8BitRowwiseQuantizedToFloatOp<
+ float,
+ nullptr,
+ false, /* HAS_CONVERT */
+ CPUContext>);
 OPERATOR_SCHEMA(Fused8BitRowwiseQuantizedToFloat)
  .NumInputs(1)
  .NumOutputs(1)
@@ -114,7 +122,11 @@ NO_GRADIENT(Fused8BitRowwiseQuantizedToFloat);
 
 REGISTER_CPU_OPERATOR(
  Fused8BitRowwiseQuantizedToHalfFloat,
- Fused8BitRowwiseQuantizedToFloatOp<at::Half, convertfp32fp16, CPUContext>);
+ Fused8BitRowwiseQuantizedToFloatOp<
+ at::Half,
+ convertfp32fp16,
+ true, /* HAS_CONVERT */
+ CPUContext>);
 OPERATOR_SCHEMA(Fused8BitRowwiseQuantizedToHalfFloat)
  .NumInputs(1)
  .NumOutputs(1)
@@ -152,7 +164,8 @@ NO_GRADIENT(Fused8BitRowwiseQuantizedToHalfFloat);
 using Fused8BitRowwiseQuantizedToFloatCPUOp =
  caffe2::Fused8BitRowwiseQuantizedToFloatOp<
  float,
- caffe2::convertfp32fp32,
+ nullptr,
+ false, /* HAS_CONVERT */
  caffe2::CPUContext>;
 
 C10_EXPORT_CAFFE2_OP_TO_C10_CPU(

caffe2/operators/fused_rowwise_8bit_conversion_ops.h

@@ -6,22 +6,23 @@
 #include "caffe2/core/logging.h"
 #include "caffe2/core/operator.h"
 #include "caffe2/operators/reducer_functors.h"
-#include "caffe2/utils/eigen_utils.h"
+#include "caffe2/perfkernels/fused_8bit_rowwise_conversion.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; \
+#define IS_LITTLE_ENDIAN  \
+ [] {  \
+ const int32_t kValue = 1;  \
+ return reinterpret_cast<const std::uint8_t*>(&kValue)[0] == 1; \
  }()
 
 template <
  typename T,
  void (*convert)(float* dst, const T* src, size_t N),
+ bool HAS_CONVERT,
  class Context>
 class FloatToFused8BitRowwiseQuantizedOp : public Operator<Context> {
  public:
@@ -45,40 +46,38 @@ class FloatToFused8BitRowwiseQuantizedOp : public Operator<Context> {
  // 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};
+ const std::vector<std::int64_t> output_dimensions = {
+ input_rows,
+ input_columns + static_cast<std::int64_t>(2 * sizeof(float))};
  auto* output = Output(
- DATA_FUSED_SCALE_BIAS_INT8, output_dimensions, at::dtype<uint8_t>());
+ DATA_FUSED_SCALE_BIAS_INT8,
+ output_dimensions,
+ at::dtype<std::uint8_t>());
 
  const auto* input_data = input.template data<T>();
- auto* output_data = output->template mutable_data<uint8_t>();
+ auto* output_data = output->template mutable_data<std::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");
+ bool is_float = std::is_same<T, float>::value;
+ if (!HAS_CONVERT) {
+ CAFFE_ENFORCE(is_float, "convert can be nullptr only if T is float");
+ FloatToFused8BitRowwiseQuantized(
+ reinterpret_cast<const float*>(input_data),
+ input_rows,
+ input_columns,
+ output_data);
+ return true;
  }
 
- vector<float> tmp;
- tmp.resize(input_columns, 0.0);
+ bool is_half = std::is_same<T, at::Half>::value;
+ CAFFE_ENFORCE(is_float || is_half);
+
+ vector<float> tmp(input_columns);
 
  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>();
+ FloatToFused8BitRowwiseQuantized(
+ tmp.data(), 1, input_columns, output_data + row * output_columns);
  }
 
  return true;
@@ -92,6 +91,7 @@ class FloatToFused8BitRowwiseQuantizedOp : public Operator<Context> {
 template <
  typename T,
  void (*convert)(T* dst, const float* src, size_t N),
+ bool HAS_CONVERT,
  class Context>
 class Fused8BitRowwiseQuantizedToFloatOp : public Operator<Context> {
  public:
@@ -109,28 +109,35 @@ class Fused8BitRowwiseQuantizedToFloatOp : public Operator<Context> {
 
  // 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};
+ const std::vector<std::int64_t> output_dimensions = {
+ input_rows,
+ input_columns - static_cast<std::int64_t>(2 * sizeof(float))};
  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>();
+ const auto* input_data = input.template data<std::uint8_t>();
  T* output_data = output->template mutable_data<T>();
 
- vector<float> tmp;
- tmp.resize(input_columns, 0.0);
+ bool is_float = std::is_same<T, float>::value;
 
- 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);
+ if (!HAS_CONVERT) {
+ CAFFE_ENFORCE(is_float, "convert can be nullptr only if T is float");
+ Fused8BitRowwiseQuantizedToFloat(
+ input_data,
+ input_rows,
+ input_columns,
+ reinterpret_cast<float*>(output_data));
+ return true;
+ }
 
- 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);
+ bool is_half = std::is_same<T, at::Half>::value;
+ CAFFE_ENFORCE(is_float || is_half);
 
+ vector<float> tmp(input_columns);
+
+ for (size_t row = 0; row < input_rows; ++row) {
+ Fused8BitRowwiseQuantizedToFloat(
+ input_data + row * input_columns, 1, input_columns, tmp.data());
  convert(output_data + row * output_columns, tmp.data(), output_columns);
  }
  return true;

caffe2/perfkernels/fused_8bit_rowwise_conversion.cc

@@ -0,0 +1,102 @@
+#include "fused_8bit_rowwise_conversion.h"
+
+#include <algorithm>
+#include <cmath>
+
+#include "common.h"
+
+namespace caffe2 {
+
+void FloatToFused8BitRowwiseQuantized__base(
+ const float* input,
+ int input_rows,
+ int input_columns,
+ std::uint8_t* output) {
+ constexpr float kEpsilon = 1e-8f;
+
+ int output_columns = input_columns + 2 * sizeof(float);
+ for (std::size_t row = 0; row < input_rows; ++row) {
+ const float* input_row = input + row * input_columns;
+ std::uint8_t* output_row = output + row * output_columns;
+ float* output_row_scale_bias =
+ reinterpret_cast<float*>(output_row + input_columns);
+
+ float minimum_element =
+ *std::min_element(input_row, input_row + input_columns);
+ float maximum_element =
+ *std::max_element(input_row, input_row + input_columns);
+ 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);
+ for (std::size_t col = 0; col < input_columns; ++col) {
+ output_row[col] =
+ std::lrintf((input_row[col] - minimum_element) * inverse_scale);
+ }
+ }
+}
+
+void Fused8BitRowwiseQuantizedToFloat__base(
+ const std::uint8_t* input,
+ int input_rows,
+ int input_columns,
+ float* output) {
+ int output_columns = input_columns - 2 * sizeof(float);
+
+ for (std::size_t row = 0; row < input_rows; ++row) {
+ const std::uint8_t* input_row = input + row * input_columns;
+ const float* input_row_scale_bias =
+ reinterpret_cast<const float*>(input_row + output_columns);
+ float* output_row = output + row * output_columns;
+
+ for (std::size_t col = 0; col < output_columns; ++col) {
+ output_row[col] =
+ input_row[col] * input_row_scale_bias[0] + input_row_scale_bias[1];
+ }
+ }
+}
+
+decltype(FloatToFused8BitRowwiseQuantized__base)
+ FloatToFused8BitRowwiseQuantized__avx2_fma;
+void FloatToFused8BitRowwiseQuantized(
+ const float* input,
+ int input_rows,
+ int input_columns,
+ std::uint8_t* output) {
+ AVX2_FMA_DO(
+ FloatToFused8BitRowwiseQuantized,
+ input,
+ input_rows,
+ input_columns,
+ output);
+ BASE_DO(
+ FloatToFused8BitRowwiseQuantized,
+ input,
+ input_rows,
+ input_columns,
+ output);
+}
+
+decltype(Fused8BitRowwiseQuantizedToFloat__base)
+ Fused8BitRowwiseQuantizedToFloat__avx2_fma;
+void Fused8BitRowwiseQuantizedToFloat(
+ const std::uint8_t* input,
+ int input_rows,
+ int input_columns,
+ float* output) {
+ AVX2_FMA_DO(
+ Fused8BitRowwiseQuantizedToFloat,
+ input,
+ input_rows,
+ input_columns,
+ output);
+ BASE_DO(
+ Fused8BitRowwiseQuantizedToFloat,
+ input,
+ input_rows,
+ input_columns,
+ output);
+}
+
+} // namespace caffe2

caffe2/perfkernels/fused_8bit_rowwise_conversion.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cstdint>
+
+namespace caffe2 {
+
+void FloatToFused8BitRowwiseQuantized(
+ const float* input,
+ int input_rows,
+ int input_columns,
+ std::uint8_t* output);
+
+void Fused8BitRowwiseQuantizedToFloat(
+ const std::uint8_t* input,
+ int input_rows,
+ int input_columns,
+ float* output);
+
+} // namespace caffe2