bitsqueeze_pytorch.cpp

#include <torch/extension.h>
#include <pybind11/stl.h>
#include <vector>
#include <string>
#include <optional>
#include <stdexcept>
#include <iostream>

extern "C" {
    #include "include/bitsqueeze.h"
}

bsq_method_t get_method_from_string(const std::string& name) {
    if (name == "Q8_0") return Q8_0;
    if (name == "Q4_0") return Q4_0;
    if (name == "Q2_K") return Q2_K;
    if (name == "TOPK") return TOPK;
    if (name == "BF16") return BF16;
    if (name == "FP16") return FP16;
    if (name == "FP8")  return FP8;
    if (name == "FP4")  return FP4;
    if (name == "MXFP8") return MXFP8;
    if (name == "MXFP4") return MXFP4;
    if (name == "NVFP4") return NVFP4;
    if (name == "NF4_DQ") return NF4_DQ;
    if (name == "NF4") return NF4;
    if (name == "IQ2_XXS") return IQ2_XXS;
    if (name == "IQ2_XS") return IQ2_XS;
    if (name == "IQ2_S") return IQ2_S;
    if (name == "Q2_K_FAST") return Q2_K_FAST;
    if (name == "TOPK_IM") return TOPK_IM;
    throw std::runtime_error("Unknown BitSqueeze method: " + name);
}

class BitSqueezeBufferWrapper {
public:
    bitsqueeze_buffer_t* buf = nullptr;
    std::vector<int64_t> original_shape;

    BitSqueezeBufferWrapper() {}

    ~BitSqueezeBufferWrapper() {
        if (buf) {
            bsq_free(buf);
            buf = nullptr;
        }
    }

    BitSqueezeBufferWrapper(const BitSqueezeBufferWrapper&) = delete;
    BitSqueezeBufferWrapper& operator=(const BitSqueezeBufferWrapper&) = delete;
    BitSqueezeBufferWrapper(BitSqueezeBufferWrapper&& other) noexcept : buf(other.buf) { other.buf = nullptr; }

    static std::unique_ptr<BitSqueezeBufferWrapper> compress(
        torch::Tensor input,
        std::string method_name,
        float sparse_ratio = 0.0f,
        std::optional<torch::Tensor> importance = std::nullopt) 
    {
        TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous");
        TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be Float32");
        TORCH_CHECK(input.device().is_cpu(), "Input must be on CPU");

        bsq_method_t method = get_method_from_string(method_name);
        auto wrapper = std::make_unique<BitSqueezeBufferWrapper>();
        wrapper->original_shape = input.sizes().vec();

        const float* src_ptr = input.data_ptr<float>();
        const float* im_ptr = nullptr;

        if (importance.has_value() && importance->defined()) {
            torch::Tensor& im = *importance;
            TORCH_CHECK(im.is_contiguous(), "Importance tensor must be contiguous");
            TORCH_CHECK(im.scalar_type() == torch::kFloat32, "Importance must be Float32");
            TORCH_CHECK(im.sizes() == input.sizes(), "Importance shape must match input");
            im_ptr = im.data_ptr<float>();
        }

        int ret = 0;

        if (method == TOPK || method == TOPK_IM) {
            TORCH_CHECK(input.dim() == 2, "Sparse methods (TOPK, TOPK_IM) require 2D input");
            uint16_t num_tokens = (uint16_t)input.size(0);
            uint16_t num_features = (uint16_t)input.size(1);
            
            ret = bsq_compress_2d(src_ptr, num_tokens, num_features, sparse_ratio, method, &wrapper->buf, im_ptr);
        } else {
            uint64_t num_elements = (uint64_t)input.numel();
            ret = bsq_compress_1d(src_ptr, num_elements, method, &wrapper->buf, im_ptr);
        }

        TORCH_CHECK(ret == 0, "BitSqueeze compression failed");
        return wrapper;
    }

    torch::Tensor decompress() {
        TORCH_CHECK(buf != nullptr, "Buffer is empty/null");

        std::vector<int64_t> shape;
        if (buf->method == TOPK || buf->method == TOPK_IM) {
            shape.push_back(buf->shape.num_tokens);
            shape.push_back(buf->shape.num_features);
        } else {
            shape.push_back(buf->shape.num_elements);
        }

        auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU);
        torch::Tensor output = torch::empty(shape, options);

        int ret = bsq_decompress(buf, output.data_ptr<float>(), (uint64_t)output.numel());
        TORCH_CHECK(ret == 0, "BitSqueeze decompression failed");

        if (!original_shape.empty()) {
            int64_t total_elements = 1;
            for(auto d : original_shape) total_elements *= d;
            
            if (total_elements == output.numel()) {
                return output.view(original_shape); 
            }
        }

        return output;
    }

    int64_t get_packed_size() {
        if (!buf) return 0;
        return (int64_t)bsq_get_packed_size(buf);
    }

    std::string get_method_name() {
        if (!buf) return "INVALID";
        return std::to_string(buf->method);
    }
};

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    py::class_<BitSqueezeBufferWrapper>(m, "BitSqueezeBuffer")
        .def(py::init<>())
        .def_static("compress", &BitSqueezeBufferWrapper::compress,
            py::arg("input"),
            py::arg("method"),
            py::arg("sparse_ratio") = 0.0f,
            py::arg("importance") = py::none(),
            "Compress a float32 tensor into a BitSqueezeBuffer")
        
        .def("decompress", &BitSqueezeBufferWrapper::decompress)
        .def_property_readonly("size", &BitSqueezeBufferWrapper::get_packed_size)
        .def_property_readonly("method", &BitSqueezeBufferWrapper::get_method_name);
}