feat(cpp): add quantize_symmetric CPU kernel

For now only per-tensor quantization is supported.

Author: dacorvo

bench/library/benchmark.py

@@ -9,6 +9,17 @@
 from quanto.library import disable_extensions
 
 
+def get_quantize_symmetric_bench(src_dtype, dst_dtype, per_axis, device):
+    a = torch.rand([10240, 10240], dtype=src_dtype).to(device)
+    scale = torch.fill((10240,), 0.5) if per_axis else torch.tensor(0.5)
+    scale = scale.to(src_dtype).to(device)
+
+    def bench_fn():
+        return torch.ops.quanto.quantize_symmetric(a, scale, dst_dtype)
+
+    return bench_fn
+
+
 def get_unpack_bench(bits, device):
     qmax = 2**bits
     a = torch.randint(0, qmax, [10240, 10240], dtype=torch.uint8).to(device)
@@ -69,6 +80,9 @@ def elapsed_time(self, other):
 
 
 GET_BENCH_FUNCTIONS = {
+    "quantize_symmetric_fp32_int8_per_tensor": lambda device: get_quantize_symmetric_bench(
+        torch.float32, torch.int8, False, device
+    ),
     "unpack_2bit": lambda device: get_unpack_bench(2, device),
     "unpack_4bit": lambda device: get_unpack_bench(4, device),
 }
@@ -89,7 +103,7 @@ def main():
             device = torch.device("cpu")
     else:
         device = torch.device(args.device)
-    all_kernels = ["unpack_2bit", "unpack_4bit"]
+    all_kernels = GET_BENCH_FUNCTIONS.keys()
     kernels = all_kernels if args.kernel is None else [args.kernel]
     for kernel in kernels:
         get_bench_fn = GET_BENCH_FUNCTIONS[kernel]

quanto/library/ext/cpp/__init__.py

@@ -19,6 +19,7 @@ def ext():
         _ext = load(
             name="quanto_cpp",
             sources=[
+                f"{module_path}/quantize.cpp",
                 f"{module_path}/unpack.cpp",
                 f"{module_path}/pybind_module.cpp",
             ],
@@ -27,6 +28,11 @@ def ext():
     return _ext
 
 
+@torch.library.impl("quanto_ext::quantize_symmetric", ["CPU"])
+def quantize_symmetric_cpp(t: torch.Tensor, scale: torch.Tensor, dtype: torch.Tensor.dtype):
+    return ext().quantize_symmetric(t, scale, dtype)
+
+
 @impl("quanto_ext::unpack", ["CPU", "CUDA"])
 def unpack_cpp(t: torch.Tensor, bits: int):
     return ext().unpack(t, bits)

quanto/library/ext/cpp/pybind_module.cpp

@@ -1,7 +1,20 @@
 #include <torch/extension.h>
+#include "quantize.h"
 #include "unpack.h"
 
+// !IMPORTANT! Some python objects such as dtype, device, are not mapped to C++ types,
+// and need to be explicitly converted using dedicated helpers before calling a C++ method.
+// As a consequence, when an operation takes such an object as parameter, instead
+// of creating a binding directly to the C++ method, you must create a binding to a
+// lambda method that converts the unmapped types and calls the C++ method.
+// See the binding of quantize_symmetric for instance.
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("quantize_symmetric",
+        [](const torch::Tensor& t, const torch::Tensor& scale, py::object dtype) {
+          return quantize_symmetric(t,
+                                    scale,
+                                    torch::python::detail::py_object_to_dtype(dtype));
+        }, "quantize_symmetric");
   m.def("unpack", &unpack, "unpack");
 }

quanto/library/ext/cpp/quantize.cpp

@@ -0,0 +1,64 @@
+#include "quantize.h"
+#include <torch/extension.h>
+
+
+template <typename T>
+torch::Tensor quantize_symmetric_per_tensor(const torch::Tensor& input, const torch::Tensor& scale) {
+    torch::Tensor output = torch::empty_like(input, c10::TensorOptions(c10::kChar).dtype(), LEGACY_CONTIGUOUS_MEMORY_FORMAT);
+    auto qdata = reinterpret_cast<int8_t*>(output.data_ptr());
+    auto numel = input.numel();
+    const T* const data = input.data_ptr<T>();
+    float float_scale = scale.data_ptr<T>()[0];
+    float inv_scale = float_scale == 0 ? 1.0f : 1.0f / float_scale;
+    for (const auto i : c10::irange(numel)) {
+        int64_t qvalue = lrintf(std::nearbyint(data[i] * inv_scale));
+        qvalue = std::max(-127LL, std::min(qvalue, 127LL));
+        qdata[i] = static_cast<int8_t>(qvalue);
+    }
+    return output;
+}
+
+
+int get_scale_axis(const torch::Tensor& scale) {
+    int axis = -1;
+    auto scale_dims = scale.sizes();
+    for (int i = 0; i < scale_dims.size(); ++i) {
+        if (scale_dims[i] != 1) {
+            axis = i;
+        }
+    }
+    return axis;
+}
+
+
+torch::Tensor quantize_symmetric_char(const torch::Tensor& input,
+                                      const torch::Tensor& scale) {
+    int axis = get_scale_axis(scale);
+    if (axis == -1) {
+        auto scale_dtype = scale.dtype();
+        if (scale_dtype == at::ScalarType::Float) {
+            return quantize_symmetric_per_tensor<float>(input, scale);
+        }
+        if (scale_dtype == at::ScalarType::Half) {
+            return quantize_symmetric_per_tensor<at::Half>(input, scale);
+        }
+        TORCH_CHECK(false, "Unsupported scale dtype:", scale_dtype)
+    }
+    TORCH_CHECK(false, "symmetric per-axis is not supported")
+}
+
+
+torch::Tensor quantize_symmetric(const torch::Tensor& input,
+                                 const torch::Tensor& scale,
+                                 at::ScalarType dtype) {
+    bool scalar_scale = (scale.sizes().size() == 0);
+    bool broadcastable_scale = (input.sizes().size() == scale.sizes().size());
+    TORCH_CHECK(scalar_scale || broadcastable_scale,
+                "Quantization scale must be scalar or broadcastable to the base tensor.")
+    TORCH_CHECK((scale.dtype() == at::ScalarType::Float) || (scale.dtype() == at::ScalarType::Half),
+                "Quantization scale must be float or float16.")
+    if (dtype == at::ScalarType::Char) {
+        return quantize_symmetric_char(input, scale);
+    }
+    TORCH_CHECK_NOT_IMPLEMENTED(false, "quantize_symmetric not supported for ", dtype)
+}

quanto/library/ext/cpp/quantize.h

@@ -0,0 +1,5 @@
+#include <torch/extension.h>
+
+torch::Tensor quantize_symmetric(const torch::Tensor& input,
+                                 const torch::Tensor& scale,
+                                 at::ScalarType dtype);