Skip to content

Commit 4148713

Browse files
tokkiwatokkiwa
committed
feat: Add SIMD support for K=3 and dynamic K dispatch
- Add K=3 SIMD implementation using AVX with padding - Add runtime K dispatch functions for K=1 to K=8 - Pass K value from Python to C++ encode/decode functions
1 parent 8ce6715 commit 4148713

2 files changed

Lines changed: 115 additions & 31 deletions

File tree

‎compressai/cpp_exts/rans/rans_interface.cpp‎

Lines changed: 111 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -255,8 +255,9 @@ std::tuple<float, float> _fast_gmm_cdf(
255255
const std::array<float, K> &weights) {
256256
float cdf1 = 0.0, cdf2 = 0.0;
257257

258-
// Check if SIMD path is enabled AND applicable (K==4)
258+
// Check if SIMD path is enabled AND applicable (K==3 or K==4)
259259
if (use_simd_path() && K == 4) {
260+
// K=4: Use full 128-bit lanes (4 floats each)
260261
__m128 x1Simd = _mm_set1_ps(static_cast<float>(x1));
261262
__m128 x2Simd = _mm_set1_ps(static_cast<float>(x2));
262263
__m256 x1x2Simd = _mm256_set_m128(x1Simd, x2Simd);
@@ -282,7 +283,35 @@ std::tuple<float, float> _fast_gmm_cdf(
282283
cdf2Simd_parts = _mm_hadd_ps(cdf2Simd_parts, cdf2Simd_parts);
283284
cdf2 = _mm_cvtss_f32(cdf2Simd_parts);
284285

285-
} else { // Generic loop for non-SIMD path or K != 4
286+
} else if (use_simd_path() && K == 3) {
287+
// K=3: Load 3 values + 1 padding (set weight to 0 for 4th element)
288+
__m128 x1Simd = _mm_set1_ps(static_cast<float>(x1));
289+
__m128 x2Simd = _mm_set1_ps(static_cast<float>(x2));
290+
__m256 x1x2Simd = _mm256_set_m128(x1Simd, x2Simd);
291+
292+
// Load 3 values and pad with 0 for means/scales, 0 for weights
293+
__m128 meansHalf = _mm_set_ps(0.0f, means[2], means[1], means[0]);
294+
__m256 meansSimd = _mm256_set_m128(meansHalf, meansHalf);
295+
__m128 scalesHalf = _mm_set_ps(1.0f, scales[2], scales[1], scales[0]); // scale=1 to avoid div by 0
296+
__m256 scalesSimd = _mm256_set_m128(scalesHalf, scalesHalf);
297+
__m128 weightsHalf = _mm_set_ps(0.0f, weights[2], weights[1], weights[0]); // weight=0 so 4th doesn't contribute
298+
__m256 weightsSimd = _mm256_set_m128(weightsHalf, weightsHalf);
299+
300+
__m256 x1x2Normalized = _mm256_div_ps(_mm256_sub_ps(x1x2Simd, meansSimd), scalesSimd);
301+
__m256 cdfs = _mm256_mul_ps(weightsSimd, _fast_gaussian_cdf(x1x2Normalized));
302+
303+
__m128 cdf2Simd_parts = _mm256_castps256_ps128(cdfs);
304+
__m128 cdf1Simd_parts = _mm256_extractf128_ps(cdfs, 1);
305+
306+
cdf1Simd_parts = _mm_hadd_ps(cdf1Simd_parts, cdf1Simd_parts);
307+
cdf1Simd_parts = _mm_hadd_ps(cdf1Simd_parts, cdf1Simd_parts);
308+
cdf1 = _mm_cvtss_f32(cdf1Simd_parts);
309+
310+
cdf2Simd_parts = _mm_hadd_ps(cdf2Simd_parts, cdf2Simd_parts);
311+
cdf2Simd_parts = _mm_hadd_ps(cdf2Simd_parts, cdf2Simd_parts);
312+
cdf2 = _mm_cvtss_f32(cdf2Simd_parts);
313+
314+
} else { // Generic loop for non-SIMD path or other K values
286315
for (int i = 0; i < K; ++i){
287316
cdf1 += weights[i] * _fast_gaussian_cdf((x1 - means[i])/scales[i]);
288317
cdf2 += weights[i] * _fast_gaussian_cdf((x2 - means[i])/scales[i]);
@@ -1058,6 +1087,78 @@ torch::Tensor RansDecoder::decode_stream_gmm(
10581087
// K_gmm_default is already defined above
10591088
// constexpr int K = 4; // Replaced by K_gmm_default for pybind context
10601089

1090+
// Helper functions to dispatch based on runtime K value
1091+
void BufferedRansEncoder_encode_with_indexes_gmm_dispatch(
1092+
BufferedRansEncoder& self,
1093+
const torch::Tensor &symbols, const torch::Tensor &scales,
1094+
const torch::Tensor &means, const torch::Tensor &weights,
1095+
const int32_t max_value, const int32_t K) {
1096+
switch(K) {
1097+
case 1: self.encode_with_indexes_gmm<1>(symbols, scales, means, weights, max_value); break;
1098+
case 2: self.encode_with_indexes_gmm<2>(symbols, scales, means, weights, max_value); break;
1099+
case 3: self.encode_with_indexes_gmm<3>(symbols, scales, means, weights, max_value); break;
1100+
case 4: self.encode_with_indexes_gmm<4>(symbols, scales, means, weights, max_value); break;
1101+
case 5: self.encode_with_indexes_gmm<5>(symbols, scales, means, weights, max_value); break;
1102+
case 6: self.encode_with_indexes_gmm<6>(symbols, scales, means, weights, max_value); break;
1103+
case 7: self.encode_with_indexes_gmm<7>(symbols, scales, means, weights, max_value); break;
1104+
case 8: self.encode_with_indexes_gmm<8>(symbols, scales, means, weights, max_value); break;
1105+
default: throw std::runtime_error("Unsupported K value: " + std::to_string(K) + ". K must be 1-8.");
1106+
}
1107+
}
1108+
1109+
py::bytes RansEncoder_encode_with_indexes_gmm_dispatch(
1110+
RansEncoder& self,
1111+
const torch::Tensor &symbols, const torch::Tensor &scales,
1112+
const torch::Tensor &means, const torch::Tensor &weights,
1113+
const int32_t max_value, const int32_t K) {
1114+
switch(K) {
1115+
case 1: return self.encode_with_indexes_gmm<1>(symbols, scales, means, weights, max_value);
1116+
case 2: return self.encode_with_indexes_gmm<2>(symbols, scales, means, weights, max_value);
1117+
case 3: return self.encode_with_indexes_gmm<3>(symbols, scales, means, weights, max_value);
1118+
case 4: return self.encode_with_indexes_gmm<4>(symbols, scales, means, weights, max_value);
1119+
case 5: return self.encode_with_indexes_gmm<5>(symbols, scales, means, weights, max_value);
1120+
case 6: return self.encode_with_indexes_gmm<6>(symbols, scales, means, weights, max_value);
1121+
case 7: return self.encode_with_indexes_gmm<7>(symbols, scales, means, weights, max_value);
1122+
case 8: return self.encode_with_indexes_gmm<8>(symbols, scales, means, weights, max_value);
1123+
default: throw std::runtime_error("Unsupported K value: " + std::to_string(K) + ". K must be 1-8.");
1124+
}
1125+
}
1126+
1127+
torch::Tensor RansDecoder_decode_with_indexes_gmm_dispatch(
1128+
RansDecoder& self,
1129+
const std::string &encoded,
1130+
const torch::Tensor &scales, const torch::Tensor &means,
1131+
const torch::Tensor &weights, const int32_t max_bs_value, const int32_t K) {
1132+
switch(K) {
1133+
case 1: return self.decode_with_indexes_gmm<1>(encoded, scales, means, weights, max_bs_value);
1134+
case 2: return self.decode_with_indexes_gmm<2>(encoded, scales, means, weights, max_bs_value);
1135+
case 3: return self.decode_with_indexes_gmm<3>(encoded, scales, means, weights, max_bs_value);
1136+
case 4: return self.decode_with_indexes_gmm<4>(encoded, scales, means, weights, max_bs_value);
1137+
case 5: return self.decode_with_indexes_gmm<5>(encoded, scales, means, weights, max_bs_value);
1138+
case 6: return self.decode_with_indexes_gmm<6>(encoded, scales, means, weights, max_bs_value);
1139+
case 7: return self.decode_with_indexes_gmm<7>(encoded, scales, means, weights, max_bs_value);
1140+
case 8: return self.decode_with_indexes_gmm<8>(encoded, scales, means, weights, max_bs_value);
1141+
default: throw std::runtime_error("Unsupported K value: " + std::to_string(K) + ". K must be 1-8.");
1142+
}
1143+
}
1144+
1145+
torch::Tensor RansDecoder_decode_stream_gmm_dispatch(
1146+
RansDecoder& self,
1147+
const torch::Tensor &scales, const torch::Tensor &means,
1148+
const torch::Tensor &weights, const int32_t max_bs_value, const int32_t K) {
1149+
switch(K) {
1150+
case 1: return self.decode_stream_gmm<1>(scales, means, weights, max_bs_value);
1151+
case 2: return self.decode_stream_gmm<2>(scales, means, weights, max_bs_value);
1152+
case 3: return self.decode_stream_gmm<3>(scales, means, weights, max_bs_value);
1153+
case 4: return self.decode_stream_gmm<4>(scales, means, weights, max_bs_value);
1154+
case 5: return self.decode_stream_gmm<5>(scales, means, weights, max_bs_value);
1155+
case 6: return self.decode_stream_gmm<6>(scales, means, weights, max_bs_value);
1156+
case 7: return self.decode_stream_gmm<7>(scales, means, weights, max_bs_value);
1157+
case 8: return self.decode_stream_gmm<8>(scales, means, weights, max_bs_value);
1158+
default: throw std::runtime_error("Unsupported K value: " + std::to_string(K) + ". K must be 1-8.");
1159+
}
1160+
}
1161+
10611162
PYBIND11_MODULE(ans, m) {
10621163
m.attr("__name__") = "compressai.ans";
10631164
m.doc() = "range Asymmetric Numeral System python bindings";
@@ -1075,12 +1176,8 @@ PYBIND11_MODULE(ans, m) {
10751176
const std::vector<int32_t> &, const std::vector<float> &, const int32_t>(
10761177
&BufferedRansEncoder::encode_with_indexes))
10771178
.def("encode_with_indexes_gmm",
1078-
static_cast<void (BufferedRansEncoder::*) (
1079-
const torch::Tensor &, const torch::Tensor &, // Changed types
1080-
const torch::Tensor &, const torch::Tensor &, // Changed types
1081-
const int32_t)>(
1082-
&BufferedRansEncoder::encode_with_indexes_gmm<K_gmm_default>), // Use K_gmm_default
1083-
py::arg("symbols"), py::arg("scales"), py::arg("means"), py::arg("weights"), py::arg("max_value"))
1179+
&BufferedRansEncoder_encode_with_indexes_gmm_dispatch,
1180+
py::arg("symbols"), py::arg("scales"), py::arg("means"), py::arg("weights"), py::arg("max_value"), py::arg("K") = K_gmm_default)
10841181
.def("flush", &BufferedRansEncoder::flush);
10851182

10861183
py::class_<RansEncoder>(m, "RansEncoder")
@@ -1096,12 +1193,8 @@ PYBIND11_MODULE(ans, m) {
10961193
const std::vector<int32_t> &, const std::vector<float> &, const int32_t>(
10971194
&RansEncoder::encode_with_indexes))
10981195
.def("encode_with_indexes_gmm",
1099-
static_cast<py::bytes (RansEncoder::*) (
1100-
const torch::Tensor &, const torch::Tensor &, // Changed types
1101-
const torch::Tensor &, const torch::Tensor &, // Changed types
1102-
const int32_t)>(
1103-
&RansEncoder::encode_with_indexes_gmm<K_gmm_default>), // Use K_gmm_default
1104-
py::arg("symbols"), py::arg("scales"), py::arg("means"), py::arg("weights"), py::arg("max_value"));
1196+
&RansEncoder_encode_with_indexes_gmm_dispatch,
1197+
py::arg("symbols"), py::arg("scales"), py::arg("means"), py::arg("weights"), py::arg("max_value"), py::arg("K") = K_gmm_default);
11051198

11061199
py::class_<RansDecoder>(m, "RansDecoder")
11071200
.def(py::init<>())
@@ -1124,22 +1217,11 @@ PYBIND11_MODULE(ans, m) {
11241217
&RansDecoder::decode_with_indexes),
11251218
"Decode a string to a list of symbols")
11261219
.def("decode_with_indexes_gmm",
1127-
static_cast<torch::Tensor (RansDecoder::*) ( // Changed return type
1128-
const std::string &,
1129-
const torch::Tensor &, // Changed types
1130-
const torch::Tensor &, // Changed types
1131-
const torch::Tensor &, // Changed types
1132-
const int32_t)>(
1133-
&RansDecoder::decode_with_indexes_gmm<K_gmm_default>), // Use K_gmm_default
1220+
&RansDecoder_decode_with_indexes_gmm_dispatch,
11341221
"Decode a string to a tensor of symbols",
1135-
py::arg("encoded"), py::arg("scales"), py::arg("means"), py::arg("weights"), py::arg("max_bs_value"))
1222+
py::arg("encoded"), py::arg("scales"), py::arg("means"), py::arg("weights"), py::arg("max_bs_value"), py::arg("K") = K_gmm_default)
11361223
.def("decode_stream_gmm",
1137-
static_cast<torch::Tensor (RansDecoder::*) (
1138-
const torch::Tensor &,
1139-
const torch::Tensor &,
1140-
const torch::Tensor &,
1141-
const int32_t)>(
1142-
&RansDecoder::decode_stream_gmm<K_gmm_default>),
1224+
&RansDecoder_decode_stream_gmm_dispatch,
11431225
"Decode stream to a tensor of symbols using GMM",
1144-
py::arg("scales"), py::arg("means"), py::arg("weights"), py::arg("max_bs_value"));
1226+
py::arg("scales"), py::arg("means"), py::arg("weights"), py::arg("max_bs_value"), py::arg("K") = K_gmm_default);
11451227
}

‎compressai/entropy_models/entropy_models.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -861,7 +861,8 @@ def compress(self, y, scales, means, weights):
861861
scales.to("cpu"),
862862
means.to("cpu"),
863863
weights.to("cpu"),
864-
abs_max + 1
864+
abs_max + 1,
865+
self.K
865866
)
866867

867868
return (rv, abs_max, zero_bitmap), y_quantized
@@ -885,7 +886,8 @@ def decompress(self, strings, abs_max, zero_bitmap, scales, means, weights):
885886
scales.to("cpu"),
886887
means.to("cpu"),
887888
weights.to("cpu"),
888-
abs_max + 1
889+
abs_max + 1,
890+
self.K
889891
)
890892

891893
#symbols = torch.tensor(values)

0 commit comments

Comments
 (0)