rns-engine 0.6.0__tar.gz → 0.8.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {rns_engine-0.6.0/src/rns_engine.egg-info → rns_engine-0.8.0}/PKG-INFO +61 -9
- {rns_engine-0.6.0 → rns_engine-0.8.0}/README.md +60 -8
- {rns_engine-0.6.0 → rns_engine-0.8.0}/pyproject.toml +1 -1
- {rns_engine-0.6.0 → rns_engine-0.8.0}/setup.py +16 -8
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine/__init__.py +16 -0
- rns_engine-0.8.0/src/rns_engine/_weighted.cpp +141 -0
- rns_engine-0.8.0/src/rns_engine/capacity.py +390 -0
- rns_engine-0.8.0/src/rns_engine/coefficients.py +179 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine/weighted.py +47 -38
- {rns_engine-0.6.0 → rns_engine-0.8.0/src/rns_engine.egg-info}/PKG-INFO +61 -9
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine.egg-info/SOURCES.txt +7 -1
- rns_engine-0.8.0/tests/test_capacity.py +161 -0
- rns_engine-0.8.0/tests/test_coefficients.py +110 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/tests/test_weighted.py +50 -0
- rns_engine-0.8.0/tests/test_weighted_benchmark.py +71 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/LICENSE +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/setup.cfg +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine/_core.cpp +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine/engine.py +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine/signed.py +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine.egg-info/dependency_links.txt +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine.egg-info/not-zip-safe +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine.egg-info/requires.txt +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/src/rns_engine.egg-info/top_level.txt +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/tests/test_packaging.py +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/tests/test_rns.py +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/tests/test_session.py +0 -0
- {rns_engine-0.6.0 → rns_engine-0.8.0}/tests/test_signed.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: rns_engine
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.8.0
|
|
4
4
|
Summary: Exact RNS engine with session/cache API.
|
|
5
5
|
Author: Evan Wesley
|
|
6
6
|
License-Expression: AGPL-3.0-only
|
|
@@ -194,7 +194,7 @@ out = session.decode_signed(session.add(a, b))
|
|
|
194
194
|
|
|
195
195
|
## Weighted INT32 partial accumulation
|
|
196
196
|
|
|
197
|
-
Version 0.
|
|
197
|
+
Version 0.7 introduced the fused direct bridge for digit-plane, chunked-GEMM, and Tensor Core pipelines.
|
|
198
198
|
|
|
199
199
|
Give the engine signed INT32 partial outputs with shape:
|
|
200
200
|
|
|
@@ -246,7 +246,7 @@ Important details:
|
|
|
246
246
|
- `decode_modular()` always returns the canonical residue in `[0, M)`.
|
|
247
247
|
- Empty term sets produce an exact zero result.
|
|
248
248
|
|
|
249
|
-
The
|
|
249
|
+
The implementation uses one fused native `_weighted` call to read signed INT32 partials, collect exact per-term magnitude bounds, apply positional weights modulo `M`, and accumulate all four rails. The pre-fusion staged encode → scale → add body remains available internally as an exact reference for A/B tests. Local AVX2 diagnostics measured roughly 2.1×–2.9× speedup at one million outputs and larger gains on tiny many-term cases; those figures are hardware-specific CPU evidence, not CUDA, Tensor Core, or universal performance claims.
|
|
250
250
|
|
|
251
251
|
A precomputed receipt can also be created without materializing partial arrays:
|
|
252
252
|
|
|
@@ -260,6 +260,45 @@ certificate.require_unique()
|
|
|
260
260
|
|
|
261
261
|
---
|
|
262
262
|
|
|
263
|
+
## Exact GEMM capacity planning
|
|
264
|
+
|
|
265
|
+
Version 0.8 separates two different safety questions that must not be conflated:
|
|
266
|
+
|
|
267
|
+
1. Does each grouped raw digit-plane coefficient fit the native signed accumulator?
|
|
268
|
+
2. Does the final radix-weighted result fit uniquely inside the global RNS modulus product?
|
|
269
|
+
|
|
270
|
+
```python
|
|
271
|
+
import rns_engine as rns
|
|
272
|
+
|
|
273
|
+
local = rns.plan_grouped_coefficient_capacity(
|
|
274
|
+
inner_dimension=5120,
|
|
275
|
+
left_digit_abs_bounds=[127] * 8,
|
|
276
|
+
right_digit_abs_bounds=[127] * 8,
|
|
277
|
+
accumulator_bits=32,
|
|
278
|
+
)
|
|
279
|
+
|
|
280
|
+
global_plan = rns.plan_digit_plane_gemm_capacity(
|
|
281
|
+
inner_dimension=5120,
|
|
282
|
+
radix=128,
|
|
283
|
+
left_digit_abs_bounds=[127] * 8,
|
|
284
|
+
right_digit_abs_bounds=[127] * 8,
|
|
285
|
+
)
|
|
286
|
+
|
|
287
|
+
assert local.max_abs_bound == 660_643_840
|
|
288
|
+
assert local.safe
|
|
289
|
+
assert local.minimum_signed_accumulator_bits == 31
|
|
290
|
+
|
|
291
|
+
assert global_plan.capacity.required_modulus_bits == 126
|
|
292
|
+
assert global_plan.capacity.additional_bits_required == 71
|
|
293
|
+
assert not global_plan.current_unique
|
|
294
|
+
```
|
|
295
|
+
|
|
296
|
+
That workload is locally safe in INT32 but globally modular-only in the current four-rail range. The planner can also validate explicit proposed CRT moduli and rejects candidates that are not pairwise coprime with the existing product.
|
|
297
|
+
|
|
298
|
+
See [`CAPACITY_PLANNING.md`](CAPACITY_PLANNING.md) for the exact laws, examples, and limits of what the receipts prove.
|
|
299
|
+
|
|
300
|
+
---
|
|
301
|
+
|
|
263
302
|
## Core API
|
|
264
303
|
|
|
265
304
|
### Rail operations
|
|
@@ -285,6 +324,16 @@ certificate.require_unique()
|
|
|
285
324
|
- `WeightedInt32Result.decode_signed(require_unique=True)`
|
|
286
325
|
- `WeightedInt32Result.decode_modular()`
|
|
287
326
|
|
|
327
|
+
### Capacity planning
|
|
328
|
+
|
|
329
|
+
- `plan_signed_capacity(max_abs_bound, additional_moduli=())`
|
|
330
|
+
- `plan_weighted_sum_capacity(weights, term_abs_bounds, addend_abs_bound=0, additional_moduli=())`
|
|
331
|
+
- `plan_digit_plane_gemm_capacity(inner_dimension, radix, left_digit_abs_bounds, right_digit_abs_bounds, ...)`
|
|
332
|
+
- `plan_grouped_coefficient_capacity(inner_dimension, left_digit_abs_bounds, right_digit_abs_bounds, accumulator_bits=32)`
|
|
333
|
+
- `SignedCapacityPlan`
|
|
334
|
+
- `DigitPlaneGemmCapacityPlan`
|
|
335
|
+
- `GroupedCoefficientCapacityPlan`
|
|
336
|
+
|
|
288
337
|
### Scalar-broadcast encoded operations
|
|
289
338
|
|
|
290
339
|
- `mul_u64(...)`
|
|
@@ -359,7 +408,7 @@ Workload: `affine_repeat_u64_io(x, 1_000_003, 7, iterations=1000)`
|
|
|
359
408
|
- OpenMP with one thread: **82.80 billion ops/sec**
|
|
360
409
|
- OpenMP with two threads: **94.86 billion ops/sec**
|
|
361
410
|
|
|
362
|
-
These measurements cover the existing affine kernels, not
|
|
411
|
+
These measurements cover the existing affine kernels, not weighted INT32 reconstruction or the pure-Python capacity planners.
|
|
363
412
|
|
|
364
413
|
---
|
|
365
414
|
|
|
@@ -404,12 +453,15 @@ rns.HAS_AVX2
|
|
|
404
453
|
|
|
405
454
|
## Current release
|
|
406
455
|
|
|
407
|
-
### v0.
|
|
456
|
+
### v0.8.0
|
|
408
457
|
|
|
409
|
-
-
|
|
410
|
-
- exact
|
|
411
|
-
-
|
|
412
|
-
-
|
|
458
|
+
- exact local grouped-coefficient accumulator planning
|
|
459
|
+
- exact global digit-plane GEMM RNS-capacity planning
|
|
460
|
+
- strict separation between local INT32 safety and global signed uniqueness
|
|
461
|
+
- candidate CRT-modulus validation with pairwise-coprime enforcement
|
|
462
|
+
- exact modulus shortfall and additional-bit receipts
|
|
463
|
+
- fused native weighted signed INT32 partial accumulation
|
|
464
|
+
- arbitrary-precision positional weights and no-wrap certificates
|
|
413
465
|
- canonical centered signed encoding and decoding
|
|
414
466
|
- exact dot-product and GEMM output-bound receipts
|
|
415
467
|
- `SignedSession` high-level API
|
|
@@ -171,7 +171,7 @@ out = session.decode_signed(session.add(a, b))
|
|
|
171
171
|
|
|
172
172
|
## Weighted INT32 partial accumulation
|
|
173
173
|
|
|
174
|
-
Version 0.
|
|
174
|
+
Version 0.7 introduced the fused direct bridge for digit-plane, chunked-GEMM, and Tensor Core pipelines.
|
|
175
175
|
|
|
176
176
|
Give the engine signed INT32 partial outputs with shape:
|
|
177
177
|
|
|
@@ -223,7 +223,7 @@ Important details:
|
|
|
223
223
|
- `decode_modular()` always returns the canonical residue in `[0, M)`.
|
|
224
224
|
- Empty term sets produce an exact zero result.
|
|
225
225
|
|
|
226
|
-
The
|
|
226
|
+
The implementation uses one fused native `_weighted` call to read signed INT32 partials, collect exact per-term magnitude bounds, apply positional weights modulo `M`, and accumulate all four rails. The pre-fusion staged encode → scale → add body remains available internally as an exact reference for A/B tests. Local AVX2 diagnostics measured roughly 2.1×–2.9× speedup at one million outputs and larger gains on tiny many-term cases; those figures are hardware-specific CPU evidence, not CUDA, Tensor Core, or universal performance claims.
|
|
227
227
|
|
|
228
228
|
A precomputed receipt can also be created without materializing partial arrays:
|
|
229
229
|
|
|
@@ -237,6 +237,45 @@ certificate.require_unique()
|
|
|
237
237
|
|
|
238
238
|
---
|
|
239
239
|
|
|
240
|
+
## Exact GEMM capacity planning
|
|
241
|
+
|
|
242
|
+
Version 0.8 separates two different safety questions that must not be conflated:
|
|
243
|
+
|
|
244
|
+
1. Does each grouped raw digit-plane coefficient fit the native signed accumulator?
|
|
245
|
+
2. Does the final radix-weighted result fit uniquely inside the global RNS modulus product?
|
|
246
|
+
|
|
247
|
+
```python
|
|
248
|
+
import rns_engine as rns
|
|
249
|
+
|
|
250
|
+
local = rns.plan_grouped_coefficient_capacity(
|
|
251
|
+
inner_dimension=5120,
|
|
252
|
+
left_digit_abs_bounds=[127] * 8,
|
|
253
|
+
right_digit_abs_bounds=[127] * 8,
|
|
254
|
+
accumulator_bits=32,
|
|
255
|
+
)
|
|
256
|
+
|
|
257
|
+
global_plan = rns.plan_digit_plane_gemm_capacity(
|
|
258
|
+
inner_dimension=5120,
|
|
259
|
+
radix=128,
|
|
260
|
+
left_digit_abs_bounds=[127] * 8,
|
|
261
|
+
right_digit_abs_bounds=[127] * 8,
|
|
262
|
+
)
|
|
263
|
+
|
|
264
|
+
assert local.max_abs_bound == 660_643_840
|
|
265
|
+
assert local.safe
|
|
266
|
+
assert local.minimum_signed_accumulator_bits == 31
|
|
267
|
+
|
|
268
|
+
assert global_plan.capacity.required_modulus_bits == 126
|
|
269
|
+
assert global_plan.capacity.additional_bits_required == 71
|
|
270
|
+
assert not global_plan.current_unique
|
|
271
|
+
```
|
|
272
|
+
|
|
273
|
+
That workload is locally safe in INT32 but globally modular-only in the current four-rail range. The planner can also validate explicit proposed CRT moduli and rejects candidates that are not pairwise coprime with the existing product.
|
|
274
|
+
|
|
275
|
+
See [`CAPACITY_PLANNING.md`](CAPACITY_PLANNING.md) for the exact laws, examples, and limits of what the receipts prove.
|
|
276
|
+
|
|
277
|
+
---
|
|
278
|
+
|
|
240
279
|
## Core API
|
|
241
280
|
|
|
242
281
|
### Rail operations
|
|
@@ -262,6 +301,16 @@ certificate.require_unique()
|
|
|
262
301
|
- `WeightedInt32Result.decode_signed(require_unique=True)`
|
|
263
302
|
- `WeightedInt32Result.decode_modular()`
|
|
264
303
|
|
|
304
|
+
### Capacity planning
|
|
305
|
+
|
|
306
|
+
- `plan_signed_capacity(max_abs_bound, additional_moduli=())`
|
|
307
|
+
- `plan_weighted_sum_capacity(weights, term_abs_bounds, addend_abs_bound=0, additional_moduli=())`
|
|
308
|
+
- `plan_digit_plane_gemm_capacity(inner_dimension, radix, left_digit_abs_bounds, right_digit_abs_bounds, ...)`
|
|
309
|
+
- `plan_grouped_coefficient_capacity(inner_dimension, left_digit_abs_bounds, right_digit_abs_bounds, accumulator_bits=32)`
|
|
310
|
+
- `SignedCapacityPlan`
|
|
311
|
+
- `DigitPlaneGemmCapacityPlan`
|
|
312
|
+
- `GroupedCoefficientCapacityPlan`
|
|
313
|
+
|
|
265
314
|
### Scalar-broadcast encoded operations
|
|
266
315
|
|
|
267
316
|
- `mul_u64(...)`
|
|
@@ -336,7 +385,7 @@ Workload: `affine_repeat_u64_io(x, 1_000_003, 7, iterations=1000)`
|
|
|
336
385
|
- OpenMP with one thread: **82.80 billion ops/sec**
|
|
337
386
|
- OpenMP with two threads: **94.86 billion ops/sec**
|
|
338
387
|
|
|
339
|
-
These measurements cover the existing affine kernels, not
|
|
388
|
+
These measurements cover the existing affine kernels, not weighted INT32 reconstruction or the pure-Python capacity planners.
|
|
340
389
|
|
|
341
390
|
---
|
|
342
391
|
|
|
@@ -381,12 +430,15 @@ rns.HAS_AVX2
|
|
|
381
430
|
|
|
382
431
|
## Current release
|
|
383
432
|
|
|
384
|
-
### v0.
|
|
433
|
+
### v0.8.0
|
|
385
434
|
|
|
386
|
-
-
|
|
387
|
-
- exact
|
|
388
|
-
-
|
|
389
|
-
-
|
|
435
|
+
- exact local grouped-coefficient accumulator planning
|
|
436
|
+
- exact global digit-plane GEMM RNS-capacity planning
|
|
437
|
+
- strict separation between local INT32 safety and global signed uniqueness
|
|
438
|
+
- candidate CRT-modulus validation with pairwise-coprime enforcement
|
|
439
|
+
- exact modulus shortfall and additional-bit receipts
|
|
440
|
+
- fused native weighted signed INT32 partial accumulation
|
|
441
|
+
- arbitrary-precision positional weights and no-wrap certificates
|
|
390
442
|
- canonical centered signed encoding and decoding
|
|
391
443
|
- exact dot-product and GEMM output-bound receipts
|
|
392
444
|
- `SignedSession` high-level API
|
|
@@ -39,19 +39,27 @@ def get_link_args():
|
|
|
39
39
|
return []
|
|
40
40
|
|
|
41
41
|
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
42
|
+
def native_extension(name, source):
|
|
43
|
+
return Extension(
|
|
44
|
+
name,
|
|
45
|
+
sources=[source],
|
|
46
|
+
include_dirs=[pybind11.get_include(), np.get_include()],
|
|
47
|
+
extra_compile_args=get_compile_args(),
|
|
48
|
+
extra_link_args=get_link_args(),
|
|
49
|
+
language="c++",
|
|
50
|
+
)
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
ext = native_extension("rns_engine._core", "src/rns_engine/_core.cpp")
|
|
54
|
+
weighted_ext = native_extension(
|
|
55
|
+
"rns_engine._weighted",
|
|
56
|
+
"src/rns_engine/_weighted.cpp",
|
|
49
57
|
)
|
|
50
58
|
|
|
51
59
|
# Project metadata and package discovery live in pyproject.toml. Keeping them
|
|
52
60
|
# out of setup.py prevents the two build entry points from drifting apart.
|
|
53
61
|
setup(
|
|
54
|
-
ext_modules=[ext],
|
|
62
|
+
ext_modules=[ext, weighted_ext],
|
|
55
63
|
include_package_data=True,
|
|
56
64
|
zip_safe=False,
|
|
57
65
|
)
|
|
@@ -57,6 +57,17 @@ from .weighted import (
|
|
|
57
57
|
accumulate_weighted_int32,
|
|
58
58
|
certify_weighted_sum_bound,
|
|
59
59
|
)
|
|
60
|
+
from .capacity import (
|
|
61
|
+
DigitPlaneGemmCapacityPlan,
|
|
62
|
+
SignedCapacityPlan,
|
|
63
|
+
plan_digit_plane_gemm_capacity,
|
|
64
|
+
plan_signed_capacity,
|
|
65
|
+
plan_weighted_sum_capacity,
|
|
66
|
+
)
|
|
67
|
+
from .coefficients import (
|
|
68
|
+
GroupedCoefficientCapacityPlan,
|
|
69
|
+
plan_grouped_coefficient_capacity,
|
|
70
|
+
)
|
|
60
71
|
|
|
61
72
|
try:
|
|
62
73
|
__version__ = _distribution_version("rns_engine")
|
|
@@ -83,6 +94,10 @@ __all__ = [
|
|
|
83
94
|
"EncodedArray", "SessionCache", "Session", "SignedSession",
|
|
84
95
|
"SignedRangeCertificate", "certify_signed_bound", "certify_signed_dot_bound",
|
|
85
96
|
"WeightedInt32Result", "accumulate_weighted_int32", "certify_weighted_sum_bound",
|
|
97
|
+
"SignedCapacityPlan", "DigitPlaneGemmCapacityPlan",
|
|
98
|
+
"GroupedCoefficientCapacityPlan",
|
|
99
|
+
"plan_signed_capacity", "plan_weighted_sum_capacity",
|
|
100
|
+
"plan_digit_plane_gemm_capacity", "plan_grouped_coefficient_capacity",
|
|
86
101
|
"info",
|
|
87
102
|
]
|
|
88
103
|
|
|
@@ -96,3 +111,4 @@ def info():
|
|
|
96
111
|
print(f" AVX2 : {'yes' if HAS_AVX2 else 'no'}")
|
|
97
112
|
print(" Core APIs : add/sub/mul/fma + raw/omp/auto scalar-broadcast family")
|
|
98
113
|
print(" Bridge APIs : weighted signed INT32 accumulation with range receipts")
|
|
114
|
+
print(" Planning APIs : local coefficient safety + global digit-plane RNS capacity")
|
|
@@ -0,0 +1,141 @@
|
|
|
1
|
+
#include <pybind11/pybind11.h>
|
|
2
|
+
#include <pybind11/numpy.h>
|
|
3
|
+
#include <stdint.h>
|
|
4
|
+
#include <stdexcept>
|
|
5
|
+
#include <string>
|
|
6
|
+
#include <cstring>
|
|
7
|
+
|
|
8
|
+
namespace py = pybind11;
|
|
9
|
+
|
|
10
|
+
using arr_i32 = py::array_t<int32_t, py::array::c_style>;
|
|
11
|
+
using arr16 = py::array_t<uint16_t, py::array::c_style>;
|
|
12
|
+
using arr32 = py::array_t<uint32_t, py::array::c_style>;
|
|
13
|
+
using arr64 = py::array_t<uint64_t, py::array::c_style | py::array::forcecast>;
|
|
14
|
+
|
|
15
|
+
static constexpr uint32_t M0 = 127;
|
|
16
|
+
static constexpr uint32_t M1 = 8191;
|
|
17
|
+
static constexpr uint32_t M2 = 65536;
|
|
18
|
+
static constexpr uint32_t M3 = 524287;
|
|
19
|
+
static constexpr uint32_t M3_MASK = 0x7FFFFu;
|
|
20
|
+
static constexpr uint64_t BM = (uint64_t)M0 * M1 * M2 * M3;
|
|
21
|
+
|
|
22
|
+
static int64_t require_1d_len(const py::array& a, const char* name) {
|
|
23
|
+
if (a.ndim() != 1) {
|
|
24
|
+
throw std::invalid_argument(std::string(name) + " must be a 1D array");
|
|
25
|
+
}
|
|
26
|
+
return (int64_t)a.shape(0);
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
static inline uint16_t r127s(uint32_t x) {
|
|
30
|
+
x = (x & 0x7F) + (x >> 7);
|
|
31
|
+
x = (x & 0x7F) + (x >> 7);
|
|
32
|
+
return x >= M0 ? x - M0 : (uint16_t)x;
|
|
33
|
+
}
|
|
34
|
+
|
|
35
|
+
static inline uint16_t r8191s(uint32_t x) {
|
|
36
|
+
x = (x & 0x1FFF) + (x >> 13);
|
|
37
|
+
x = (x & 0x1FFF) + (x >> 13);
|
|
38
|
+
return x >= M1 ? x - M1 : (uint16_t)x;
|
|
39
|
+
}
|
|
40
|
+
|
|
41
|
+
static inline uint32_t fma524287s(uint32_t a, uint32_t b, uint32_t c) {
|
|
42
|
+
uint64_t x = (uint64_t)a * b + c;
|
|
43
|
+
x = (x & M3_MASK) + (x >> 19);
|
|
44
|
+
x = (x & M3_MASK) + (x >> 19);
|
|
45
|
+
x = (x & M3_MASK) + (x >> 19);
|
|
46
|
+
uint32_t r = (uint32_t)x;
|
|
47
|
+
return r >= M3 ? (r - M3) : r;
|
|
48
|
+
}
|
|
49
|
+
|
|
50
|
+
static inline uint32_t mod524287_u64(uint64_t x) {
|
|
51
|
+
x = (x & M3_MASK) + (x >> 19);
|
|
52
|
+
x = (x & M3_MASK) + (x >> 19);
|
|
53
|
+
x = (x & M3_MASK) + (x >> 19);
|
|
54
|
+
uint32_t r = (uint32_t)x;
|
|
55
|
+
return r >= M3 ? (r - M3) : r;
|
|
56
|
+
}
|
|
57
|
+
|
|
58
|
+
py::tuple weighted_int32(const arr_i32& partials, const arr64& weights) {
|
|
59
|
+
if (partials.ndim() != 2) {
|
|
60
|
+
throw std::invalid_argument("partials must be a 2D int32 array");
|
|
61
|
+
}
|
|
62
|
+
|
|
63
|
+
const int64_t terms = (int64_t)partials.shape(0);
|
|
64
|
+
const int64_t n = (int64_t)partials.shape(1);
|
|
65
|
+
const int64_t weight_count = require_1d_len(weights, "weights");
|
|
66
|
+
if (weight_count != terms) {
|
|
67
|
+
throw std::invalid_argument(
|
|
68
|
+
"weights length must match partial term count (" +
|
|
69
|
+
std::to_string(weight_count) + " != " + std::to_string(terms) + ")");
|
|
70
|
+
}
|
|
71
|
+
|
|
72
|
+
arr16 o0({n}), o1({n}), o2({n});
|
|
73
|
+
arr32 o3({n});
|
|
74
|
+
arr64 bounds({terms});
|
|
75
|
+
|
|
76
|
+
uint16_t* p0 = o0.mutable_data();
|
|
77
|
+
uint16_t* p1 = o1.mutable_data();
|
|
78
|
+
uint16_t* p2 = o2.mutable_data();
|
|
79
|
+
uint32_t* p3 = o3.mutable_data();
|
|
80
|
+
uint64_t* pb = bounds.mutable_data();
|
|
81
|
+
|
|
82
|
+
if (n > 0) {
|
|
83
|
+
std::memset(p0, 0, (size_t)n * sizeof(uint16_t));
|
|
84
|
+
std::memset(p1, 0, (size_t)n * sizeof(uint16_t));
|
|
85
|
+
std::memset(p2, 0, (size_t)n * sizeof(uint16_t));
|
|
86
|
+
std::memset(p3, 0, (size_t)n * sizeof(uint32_t));
|
|
87
|
+
}
|
|
88
|
+
if (terms > 0) {
|
|
89
|
+
std::memset(pb, 0, (size_t)terms * sizeof(uint64_t));
|
|
90
|
+
}
|
|
91
|
+
|
|
92
|
+
const int32_t* x = partials.data();
|
|
93
|
+
const uint64_t* w = weights.data();
|
|
94
|
+
|
|
95
|
+
{
|
|
96
|
+
py::gil_scoped_release release;
|
|
97
|
+
for (int64_t term = 0; term < terms; ++term) {
|
|
98
|
+
const uint64_t weight = w[term] % BM;
|
|
99
|
+
const uint16_t w0 = (uint16_t)(weight % M0);
|
|
100
|
+
const uint16_t w1 = (uint16_t)(weight % M1);
|
|
101
|
+
const uint16_t w2 = (uint16_t)weight;
|
|
102
|
+
const uint32_t w3 = mod524287_u64(weight);
|
|
103
|
+
const int32_t* row = x + term * n;
|
|
104
|
+
uint64_t max_abs = 0;
|
|
105
|
+
|
|
106
|
+
for (int64_t i = 0; i < n; ++i) {
|
|
107
|
+
const int32_t signed_value = row[i];
|
|
108
|
+
const uint64_t magnitude = signed_value < 0
|
|
109
|
+
? (uint64_t)(-(int64_t)signed_value)
|
|
110
|
+
: (uint64_t)signed_value;
|
|
111
|
+
if (magnitude > max_abs) max_abs = magnitude;
|
|
112
|
+
|
|
113
|
+
if (weight == 0) continue;
|
|
114
|
+
|
|
115
|
+
const uint64_t canonical = signed_value < 0
|
|
116
|
+
? (uint64_t)((int64_t)BM + (int64_t)signed_value)
|
|
117
|
+
: (uint64_t)signed_value;
|
|
118
|
+
const uint16_t x0 = (uint16_t)(canonical % M0);
|
|
119
|
+
const uint16_t x1 = (uint16_t)(canonical % M1);
|
|
120
|
+
const uint16_t x2 = (uint16_t)canonical;
|
|
121
|
+
const uint32_t x3 = mod524287_u64(canonical);
|
|
122
|
+
|
|
123
|
+
p0[i] = r127s((uint32_t)p0[i] + (uint32_t)x0 * w0);
|
|
124
|
+
p1[i] = r8191s((uint32_t)p1[i] + (uint32_t)x1 * w1);
|
|
125
|
+
p2[i] = (uint16_t)((uint32_t)p2[i] + (uint32_t)x2 * w2);
|
|
126
|
+
p3[i] = fma524287s(x3, w3, p3[i]);
|
|
127
|
+
}
|
|
128
|
+
pb[term] = max_abs;
|
|
129
|
+
}
|
|
130
|
+
}
|
|
131
|
+
|
|
132
|
+
return py::make_tuple(o0, o1, o2, o3, bounds);
|
|
133
|
+
}
|
|
134
|
+
|
|
135
|
+
PYBIND11_MODULE(_weighted, m) {
|
|
136
|
+
m.doc() = "Fused weighted signed INT32 accumulation for rns_engine.";
|
|
137
|
+
m.def(
|
|
138
|
+
"weighted_int32",
|
|
139
|
+
&weighted_int32,
|
|
140
|
+
"Fused signed int32 weighting and four-rail accumulation with per-term bounds.");
|
|
141
|
+
}
|