speed up markov with C++
9 changed files+228−42
zia_benchmark/pyproject.tomlmodified+32−6View file
@@ -1,8 +1,34 @@
11 [build-system]
2-requires = ["setuptools>=45", "wheel", "setuptools_scm>=6.2"]
3-build-backend = "setuptools.build_meta"
2+requires = ["scikit-build-core>=0.5.0", "pybind11>=2.11.1"]
3+build-backend = "scikit_build_core.build"
44
5-[tool.pytest.ini_options]
6-testpaths = ["tests"]
7-python_files = ["test_*.py"]
8-addopts = "-ra -q"
5+[project]
6+name = "zia_benchmark"
7+version = "0.1.0"
8+description = "Benchmarking compression methods for numeric arrays"
9+readme = "README.md"
10+requires-python = ">=3.8"
11+authors = [
12+ { name = "Jeremy Magland" }
13+]
14+dependencies = [
15+ "numpy",
16+ "scipy",
17+ "zstandard",
18+ "simple_ans",
19+ "requests",
20+ "lindi",
21+ "brotli",
22+ "click",
23+ "numba",
24+ "pybind11>=2.11.1"
25+]
26+
27+[tool.scikit-build]
28+cmake.minimum-version = "3.15"
29+cmake.source-dir = "src/zia_benchmark/algorithms/simple_ans"
30+cmake.build-type = "Release"
31+wheel.packages = ["src/zia_benchmark"]
32+
33+[project.scripts]
34+zia-benchmark = "zia_benchmark.cli:main"
zia_benchmark/setup.pymodified+35−1View file
@@ -1,4 +1,31 @@
11 from setuptools import setup, find_packages
2+from setuptools.command.build_ext import build_ext
3+from setuptools import Extension
4+import os
5+import sys
6+import subprocess
7+
8+class CMakeExtension(Extension):
9+ def __init__(self, name, sourcedir=""):
10+ Extension.__init__(self, name, sources=[])
11+ self.sourcedir = os.path.abspath(sourcedir)
12+
13+class CMakeBuild(build_ext):
14+ def build_extension(self, ext):
15+ extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
16+
17+ cmake_args = [
18+ f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",
19+ f"-DPYTHON_EXECUTABLE={sys.executable}"
20+ ]
21+
22+ build_args = []
23+
24+ if not os.path.exists(self.build_temp):
25+ os.makedirs(self.build_temp)
26+
27+ subprocess.check_call(["cmake", ext.sourcedir] + cmake_args, cwd=self.build_temp)
28+ subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=self.build_temp)
229
330 setup(
431 name="zia_benchmark",
@@ -14,8 +41,15 @@ setup(
1441 "lindi",
1542 "brotli",
1643 "click",
17- "numba"
44+ "pybind11>=2.11.1"
1845 ],
46+ ext_modules=[
47+ CMakeExtension("zia_benchmark.algorithms.simple_ans.markov_reconstruct_cpp_ext",
48+ sourcedir="src/zia_benchmark/algorithms/simple_ans")
49+ ],
50+ cmdclass={
51+ "build_ext": CMakeBuild,
52+ },
1953 python_requires=">=3.8",
2054 entry_points={
2155 "console_scripts": [
zia_benchmark/src/zia_benchmark/algorithms/simple_ans/CMakeLists.txtadded+21−0View file
@@ -0,0 +1,21 @@
1+cmake_minimum_required(VERSION 3.15)
2+project(markov_reconstruct_cpp)
3+
4+set(CMAKE_CXX_STANDARD 14)
5+set(CMAKE_CXX_STANDARD_REQUIRED ON)
6+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
7+
8+find_package(Python REQUIRED COMPONENTS Interpreter Development.Module)
9+
10+# Fetch and include pybind11
11+include(FetchContent)
12+FetchContent_Declare(
13+ pybind11
14+ GIT_REPOSITORY https://github.com/pybind/pybind11.git
15+ GIT_TAG v2.11.1
16+)
17+FetchContent_MakeAvailable(pybind11)
18+
19+pybind11_add_module(markov_reconstruct_cpp_ext markov_reconstruct.cpp)
20+
21+install(TARGETS markov_reconstruct_cpp_ext DESTINATION zia_benchmark/algorithms/simple_ans)
zia_benchmark/src/zia_benchmark/algorithms/simple_ans/__init__.pymodified+10−5View file
@@ -1,6 +1,8 @@
11 import numpy as np
2-from .markov import markov_predict, markov_reconstruct
3-
2+from zia_benchmark.algorithms.simple_ans.markov_reconstruct_wrapper import (
3+ markov_reconstruct as markov_reconstruct_cpp,
4+)
5+from .markov import markov_predict
46
57 SOURCE_FILE = "simple_ans/__init__.py"
68
@@ -150,7 +152,7 @@ def simple_ans_markov_encode(x: np.ndarray) -> bytes:
150152 from simple_ans import ans_encode
151153
152154 assert x.ndim == 1
153- coeffs, initial, resid = markov_predict(x, M=6)
155+ coeffs, initial, resid = markov_predict(x, M=10)
154156 # Encode just the differences
155157 encoded = ans_encode(resid)
156158 if x.dtype == np.uint8:
@@ -233,8 +235,11 @@ def simple_ans_markov_decode(x: bytes, dtype: str) -> np.ndarray:
233235 symbol_values=symbol_values.astype(dtype),
234236 bitstream=bitstream,
235237 )
238+ import time
239+
236240 resid = ans_decode(encoded)
237- return markov_reconstruct(coeffs, initial, resid)
241+ output = markov_reconstruct_cpp(coeffs, initial, resid)
242+ return output
238243
239244
240245 algorithms = [
@@ -257,7 +262,7 @@ algorithms = [
257262 },
258263 {
259264 "name": "simple-ans-markov",
260- "version": "4",
265+ "version": "5",
261266 "encode": lambda x: simple_ans_markov_encode(x),
262267 "decode": lambda x, dtype: simple_ans_markov_decode(x, dtype),
263268 "description": "ANS compression with Markov prediction for exploiting temporal correlations in the data.",
zia_benchmark/src/zia_benchmark/algorithms/simple_ans/markov.pymodified+26−26View file
@@ -1,5 +1,4 @@
11 import numpy as np
2-import numba
32 from zia_benchmark._analysis import linear_fit
43
54
@@ -35,32 +34,33 @@ def markov_predict(x: np.ndarray, M: int) -> tuple:
3534 return coeffs, initial, residuals
3635
3736
38-@numba.jit(nopython=True)
39-def markov_reconstruct(
40- coeffs: np.ndarray, initial: np.ndarray, resid: np.ndarray
41-) -> np.ndarray:
42- """Reconstruct signal from Markov model parameters and residuals.
37+# C++/pybind11 implementation is a lot faster than the Numba implementation
38+# @numba.jit(nopython=True)
39+# def markov_reconstruct(
40+# coeffs: np.ndarray, initial: np.ndarray, resid: np.ndarray
41+# ) -> np.ndarray:
42+# """Reconstruct signal from Markov model parameters and residuals.
4343
44- Args:
45- coeffs: Model coefficients from linear regression
46- initial: Initial values needed for prediction
47- resid: Prediction residuals
44+# Args:
45+# coeffs: Model coefficients from linear regression
46+# initial: Initial values needed for prediction
47+# resid: Prediction residuals
4848
49- Returns:
50- np.ndarray: Reconstructed signal
51- """
52- M = len(initial) + 1 # Number of samples used in prediction
53- output = np.zeros(len(resid) + len(initial), dtype=resid.dtype)
54- output[: len(initial)] = initial # Set initial values
49+# Returns:
50+# np.ndarray: Reconstructed signal
51+# """
52+# M = len(initial) + 1 # Number of samples used in prediction
53+# output = np.zeros(len(resid) + len(initial), dtype=resid.dtype)
54+# output[: len(initial)] = initial # Set initial values
5555
56- # Reconstruct signal iteratively
57- for i in range(len(resid)):
58- # Get previous M-1 values to make prediction
59- prev_values = output[i : i + M - 1]
60- # Make prediction using coefficients
61- prediction = np.sum(coeffs[:-1] * prev_values) + coeffs[-1]
62- prediction = np.round(prediction)
63- # Add residual to get actual value
64- output[i + M - 1] = prediction + resid[i]
56+# # Reconstruct signal iteratively
57+# for i in range(len(resid)):
58+# # Get previous M-1 values to make prediction
59+# prev_values = output[i : i + M - 1]
60+# # Make prediction using coefficients
61+# prediction = np.sum(coeffs[:-1] * prev_values) + coeffs[-1]
62+# prediction = np.round(prediction)
63+# # Add residual to get actual value
64+# output[i + M - 1] = prediction + resid[i]
6565
66- return output
66+# return output
zia_benchmark/src/zia_benchmark/algorithms/simple_ans/markov_reconstruct.cppadded+69−0View file
@@ -0,0 +1,69 @@
1+#include <pybind11/pybind11.h>
2+#include <pybind11/numpy.h>
3+#include <cmath>
4+#include <iostream>
5+
6+namespace py = pybind11;
7+
8+py::array_t<int16_t> markov_reconstruct_cpp(
9+ py::array_t<float> coeffs,
10+ py::array_t<int16_t> initial,
11+ py::array_t<int16_t> resid
12+) {
13+ // Get array buffers
14+ auto coeffs_buf = coeffs.request();
15+ auto initial_buf = initial.request();
16+ auto resid_buf = resid.request();
17+
18+ // Get raw pointers to data
19+ float* coeffs_ptr = static_cast<float*>(coeffs_buf.ptr);
20+ int16_t* initial_ptr = static_cast<int16_t*>(initial_buf.ptr);
21+ int16_t* resid_ptr = static_cast<int16_t*>(resid_buf.ptr);
22+
23+ // Calculate dimensions
24+ size_t M = initial_buf.shape[0] + 1; // Number of samples used in prediction
25+ size_t output_size = resid_buf.shape[0] + initial_buf.shape[0];
26+
27+ // Create output array with explicit shape and memory ownership
28+ std::vector<ssize_t> shape = {static_cast<ssize_t>(output_size)};
29+ py::array_t<int16_t> output(shape);
30+ py::buffer_info output_buf = output.request(true); // Request writable buffer
31+ int16_t* output_ptr = static_cast<int16_t*>(output_buf.ptr);
32+
33+ // Copy initial values with bounds check
34+ for (size_t i = 0; i < initial_buf.shape[0] && i < output_size; i++) {
35+ output_ptr[i] = initial_ptr[i];
36+ }
37+
38+ size_t resid_size = resid_buf.shape[0];
39+
40+ // Reconstruct signal iteratively
41+ for (size_t i = 0; i < resid_size; i++) {
42+ float prediction = 0.0f;
43+
44+ // Calculate prediction using coefficients (excluding bias term)
45+ for (size_t j = 0; j < M - 1; j++) {
46+ float term = coeffs_ptr[j] * static_cast<float>(output_ptr[i + j]);
47+ prediction += term;
48+ }
49+
50+ // Add bias term separately
51+ prediction += coeffs_ptr[coeffs_buf.shape[0] - 1];
52+
53+ // Round prediction to nearest integer
54+ float rounded_prediction = std::round(prediction);
55+
56+ // Add residual and store result
57+ int16_t final_value = static_cast<int16_t>(rounded_prediction + static_cast<float>(resid_ptr[i]));
58+ output_ptr[i + M - 1] = final_value;
59+ }
60+
61+ return output;
62+}
63+
64+PYBIND11_MODULE(markov_reconstruct_cpp_ext, m) {
65+ m.doc() = "C++ implementation of markov_reconstruct using pybind11";
66+ m.def("markov_reconstruct_cpp", &markov_reconstruct_cpp,
67+ "Reconstruct signal from Markov model parameters and residuals",
68+ py::arg("coeffs"), py::arg("initial"), py::arg("resid"));
69+}
zia_benchmark/src/zia_benchmark/algorithms/simple_ans/markov_reconstruct_wrapper.pyadded+19−0View file
@@ -0,0 +1,19 @@
1+import numpy as np
2+from .markov_reconstruct_cpp_ext import markov_reconstruct_cpp
3+
4+
5+def markov_reconstruct(coeffs, initial, resid):
6+ """Reconstruct signal from Markov model parameters and residuals using C++ implementation.
7+
8+ Args:
9+ coeffs: Model coefficients from linear regression (float32)
10+ initial: Initial values needed for prediction (int16)
11+ resid: Prediction residuals (int16)
12+
13+ Returns:
14+ np.ndarray: Reconstructed signal (int16)
15+ """
16+ # Call C++ implementation
17+ return markov_reconstruct_cpp(
18+ coeffs.astype(np.float32), initial.astype(np.int16), resid.astype(np.int16)
19+ )
zia_benchmark/src/zia_benchmark/run_benchmarks.pymodified+0−4View file
@@ -13,10 +13,6 @@ from ._memobin import (
1313 download_from_memobin,
1414 exists_in_memobin,
1515 )
16-from .algorithms.simple_ans.markov import markov_reconstruct
17-
18-# warm up the JIT
19-markov_reconstruct(np.array([1, 2, 3]), np.array([1, 2]), np.array([1]))
2016
2117
2218 system_version = "v5"
zia_benchmark/test_markov_reconstruct.pyadded+16−0View file
@@ -0,0 +1,16 @@
1+import numpy as np
2+from zia_benchmark.algorithms.simple_ans.markov import markov_predict, markov_reconstruct
3+from zia_benchmark.algorithms.simple_ans.markov_reconstruct_wrapper import markov_reconstruct as markov_reconstruct_cpp
4+
5+coeffs = np.array([1, 2, 3], dtype=np.float32)
6+initial = np.array([7, 5], dtype=np.int16)
7+resid = np.array([6, 7], dtype=np.int16)
8+a = markov_reconstruct(coeffs, initial, resid)
9+
10+b = markov_reconstruct_cpp(coeffs, initial, resid)
11+
12+print(a)
13+print(b)
14+print(type(b))
15+print(b.shape)
16+print(b.dtype)