Skip to content

Make code safe for large (>2G) buffers #247

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
skhrg wants to merge 6 commits into
base: master
Choose a base branch
from
+49 −28
Draft

Make code safe for large (>2G) buffers #247

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ae046e6
fix: use 64 bit indices to avoid overflows
skhrg Jul 3, 2026
9a22bf5
fix: switch to size_t
skhrg Jul 9, 2026
2dc8e3b
Merge branch 'master' into index_overflow
skhrg Jul 9, 2026
f3f270f
Add failing test case
tskisner Jul 9, 2026
e56adbd
fix: switch to int64_t anywhere where we index with a product and ass…
skhrg Jul 10, 2026
a10a7e3
feat: add check that buffer has no axes greater than 32 bits
skhrg Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions include/numpy_assist.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#pragma once

#include <boost/python.hpp>
#include <cstdint>
#include <exception>

#include "exceptions.h"
Expand Down Expand Up @@ -108,7 +109,8 @@ inline int numpysafe_extract_int(const bp::object obj, const std::string argstr)
}


static std::string shape_string(std::vector<int> shape)
template <typename T>
static std::string shape_string(std::vector<T> shape)
Comment on lines +112 to +113

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason not to just change this to <int64_t>?

@skhrg skhrg Jul 14, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we change shape below to be int_64 then than is fine. I had avoided doing that because we had wanted to proceed with the assumption that the shape of the array should be limited s.t it is always int in each dimension. Hence size stays an int (and thus ndet and nsamp stay int) but vshape is allowed to be int64_t so that we can check it even if it is more than an int.

{
std::ostringstream s;
s << "(";
Expand Down Expand Up @@ -177,7 +179,7 @@ class BufferWrapper {
if (!check_buffer_type<T>(*view.get()))
throw dtype_exception(name, type_name<T>());

std::vector<int> vshape;
std::vector<int64_t> vshape;
for (int i=0; i<view->ndim; i++)
vshape.push_back(view->shape[i]);

Expand All @@ -190,13 +192,17 @@ class BufferWrapper {
if (shape[i] == -2) {
if (ellipsis_count++) {
std::ostringstream s;
s << "Invalid shape specifier " << shape_string(shape) << " (multiple elipses).";
s << "Invalid shape specifier " << shape_string<int>(shape) << " (multiple elipses).";
throw shape_exception(name, s.str());
}
// Ignore 0 or more leading entries.
j = vshape.size() - (shape.size() - i) + 1;
} else if (j >= vshape.size()) {
break;
} else if (vshape[j] > std::numeric_limits<int>::max()) {
std::ostringstream s;
s << "Dimension" << i << " is greater than 32 bits in length.";
throw shape_exception(name, s.str());
} else if (shape[i] == -1 || shape[i] == vshape[j]) {
// Match.
j++;
Expand All @@ -206,8 +212,8 @@ class BufferWrapper {
}
if (i != shape.size() || j != vshape.size()) {
std::ostringstream s;
s << "Expected " << shape_string(shape) << " but got " <<
shape_string(vshape) << ".";
s << "Expected " << shape_string<int>(shape) << " but got " <<
shape_string<int64_t>(vshape) << ".";
throw shape_exception(name, s.str());
}
}
Expand Down
46 changes: 23 additions & 23 deletions src/array_ops.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ void nmat_detvecs_apply(const bp::object & ft, const bp::object & bins, const bp
//#pragma omp parallel for
for(int di = 0; di < ndet; di++)
for(int i = b1; i < b2; i++)
ft_[di*nmode+i] *= biD[di]/norm;
ft_[(int64_t)(di*nmode+i)] *= biD[di]/norm;

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sufficient ... must cast one of the multiplicands directly; e.g.
ft_[(int64_t)di*n_mode+i] ...

Note also that (*ft_buf.ptr_2d(di, i)) would get this same reference.

// Do ft += s*iV[ndet,nvec] dot Q [nvec,nm]
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, ndet, nm, nvec, s/norm, biV, nvec, Q, nm, 1.0f, ft_+b1, nmode);
delete [] Q;
Expand Down Expand Up @@ -250,22 +250,22 @@ void pcut_full_tod2vals_helper(const vector<RangesInt32> & rangemat, T * tod, in
for(int di = 0; di < rangemat.size(); di++)
for (auto const &r: rangemat[di].segments)
for(int j = r.first; j < r.second; j++, i++)
vals[i] = tod[di*nsamp+j];
vals[i] = tod[(int64_t)(di*nsamp+j)];

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same -- apply the cast to di.

}
template <typename T>
void pcut_full_vals2tod_helper(const vector<RangesInt32> & rangemat, T * tod, int nsamp, T * vals) {
int i = 0;
for(int di = 0; di < rangemat.size(); di++)
for (auto const &r: rangemat[di].segments)
for(int j = r.first; j < r.second; j++, i++)
tod[di*nsamp+j] = vals[i];
tod[(int64_t)(di*nsamp+j)] = vals[i];

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same.

}

template <typename T>
void pcut_full_translate_helper(const vector<RangesInt32> & iranges, const vector<RangesInt32> & oranges, int insamp, T * ivals, int onsamp, T * ovals) {
// Translate cut degrees of freedom from one sample rate to another.
// Simple interpolation for upscaling, averaging when downscaling
int ioff = 0, ooff = 0;
int64_t ioff = 0, ooff = 0;
for(int di = 0; di < iranges.size(); di++) {
for(int ri = 0; ri < iranges[di].segments.size(); ri++) {
auto const & irange = iranges[di].segments[ri];
Expand Down Expand Up @@ -306,7 +306,7 @@ int pcut_poly_measure_helper(const vector<RangesInt32> & rangemat, int resolutio
}
template <typename T>
void pcut_poly_tod2vals_helper(const vector<RangesInt32> & rangemat, int resolution, int nmax, T * tod, int nsamp, T * vals) {
int i = 0;
int64_t i = 0;
for(int di = 0; di < rangemat.size(); di++) {
for (auto const &r: rangemat[di].segments) {
int np = get_npoly(r.second-r.first, resolution, nmax);
Expand All @@ -318,7 +318,7 @@ void pcut_poly_tod2vals_helper(const vector<RangesInt32> & rangemat, int resolut
for(int p = 0; p < np; p++) vals[i+p] = 0;
for(int s = r.first; s < r.second; s++) {
T x = -1 + 2*(s-r.first)/T(r.second-r.first-1);
T t = tod[di*nsamp+s];
T t = tod[(int64_t)(di*nsamp+s)];

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same

In all of these disjoint loops over detectors, I think it's simpler to compute a local pointer and use that. The compiler will optimize it to the same thing, but it keeps things cleaner. i.e.

for(int64_t di = 0; di < rangemat.size(); di++) {
    auto this_signal = tod + di*nsamp;
    ...

vals[i] += t;
if(np > 1) vals[i+1] += t*x;
if(np > 2) {
Expand All @@ -336,7 +336,7 @@ void pcut_poly_tod2vals_helper(const vector<RangesInt32> & rangemat, int resolut
}
template <typename T>
void pcut_poly_vals2tod_helper(const vector<RangesInt32> & rangemat, int resolution, int nmax, T * tod, int nsamp, T * vals) {
int i = 0;
int64_t i = 0;
for(int di = 0; di < rangemat.size(); di++) {
for (auto const &r: rangemat[di].segments) {
int np = get_npoly(r.second-r.first, resolution, nmax);
Expand Down Expand Up @@ -376,7 +376,7 @@ template <typename T>
void pcut_poly_translate_helper(const vector<RangesInt32> & iranges, const vector<RangesInt32> & oranges, int resolution, int nmax, int insamp, T * ivals, int onsamp, T * ovals) {
// Translate cut degrees of freedom from one sample rate to another.
// Zero-pad poly coeffs if upsampling, truncate if downsampling
int ioff = 0, ooff = 0;
int64_t ioff = 0, ooff = 0;
for(int di = 0; di < iranges.size(); di++) {
for(int ri = 0; ri < iranges[di].segments.size(); ri++) {
auto const & irange = iranges[di].segments[ri];
Expand Down Expand Up @@ -596,12 +596,12 @@ void _moment(T* data, T* output, int moment, bool central, int64_t start, int64_
val = center;
}
else {
for(int si = start; si < stop; si++) {
for(int64_t si = start; si < stop; si++) {
val = val + pow(data[si] - center, moment);
}
val = val / bsize;
}
for(int si = start; si < stop; si++) {
for(int64_t si = start; si < stop; si++) {
output[si] = val;
}

Expand All @@ -620,7 +620,7 @@ void _block_moment(T* tod_data, T* output, int bsize, int moment, bool central,
_moment(tod_data, output, moment, central, ioff, ioff+shift);
}

for(int bi = 0; bi < nblock; bi++) {
for(int64_t bi = 0; bi < nblock; bi++) {
int64_t start = (bi * bsize) + shift;
int64_t stop = min(start + bsize, (int64_t)nsamp);
_moment(tod_data, output, moment, central, ioff+start, ioff+stop);
Expand Down Expand Up @@ -659,7 +659,7 @@ void _minmax(T* data, T* output, int mode, int64_t start, int64_t stop)
auto max = std::max_element(data+start, data+stop);
val = *max - *min;
}
for(int si = start; si < stop; si++) {
for(int64_t si = start; si < stop; si++) {
output[si] = val;
}

Expand All @@ -678,7 +678,7 @@ void _block_minmax(T* tod_data, T* output, int bsize, int mode, int ndet, int ns
_minmax(tod_data, output, mode, ioff, ioff+shift);
}

for(int bi = 0; bi < nblock; bi++) {
for(int64_t bi = 0; bi < nblock; bi++) {
int64_t start = (bi * bsize) + shift;
int64_t stop = min(start + bsize, (int64_t)nsamp);
_minmax(tod_data, output, mode, ioff+start, ioff+stop);
Expand Down Expand Up @@ -715,7 +715,7 @@ void _clean_flag(int* flag_data, int width, int ndet, int nsamp)
if(det_flag[si]==0) {
// If this block was too small
if(count<width) {
for(int i = si - count; i < si; i++){
for(int64_t i = si - count; i < si; i++){

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Odd to change this type, but not si...

det_flag[i] = 0;
}
}
Expand Down Expand Up @@ -1018,9 +1018,9 @@ void _interp1d(const bp::object & x, const bp::object & y, const bp::object & x_
#pragma omp for
for (int row = 0; row < n_rows; ++row) {

int y_row_start = row * y_data_stride;
int y_row_end = y_row_start + n_x;
int y_interp_row_start = row * y_interp_data_stride;
int64_t y_row_start = row * y_data_stride;
int64_t y_row_end = y_row_start + n_x;
int64_t y_interp_row_start = row * y_interp_data_stride;

T* y_row = y_data + y_row_start;
T* y_interp_row = y_interp_data + y_interp_row_start;
Expand Down Expand Up @@ -1057,9 +1057,9 @@ void _interp1d(const bp::object & x, const bp::object & y, const bp::object & x_
#pragma omp for
for (int row = 0; row < n_rows; ++row) {

int y_row_start = row * y_data_stride;
int y_row_end = y_row_start + n_x;
int y_interp_row_start = row * y_interp_data_stride;
int64_t y_row_start = row * y_data_stride;
int64_t y_row_end = y_row_start + n_x;
int64_t y_interp_row_start = row * y_interp_data_stride;

// Transform y row to double array for gsl
double y_dbl[n_x];
Expand Down Expand Up @@ -1133,7 +1133,7 @@ void _detrend(T* data, const int ndets, const int nsamps, const int row_stride,
if (method == "mean") {
#pragma omp parallel for num_threads(nthreads)
for (int i = 0; i < ndets; ++i) {
int ioff = i * row_stride;
int64_t ioff = i * row_stride;

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ding


T* data_row = data + ioff;

Expand All @@ -1153,7 +1153,7 @@ void _detrend(T* data, const int ndets, const int nsamps, const int row_stride,
else if (method == "median") {
#pragma omp parallel for num_threads(nthreads)
for (int i = 0; i < ndets; ++i) {
int ioff = i * row_stride;
int64_t ioff = i * row_stride;

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ding


T* data_row = data + ioff;

Expand Down Expand Up @@ -1185,7 +1185,7 @@ void _detrend(T* data, const int ndets, const int nsamps, const int row_stride,

#pragma omp parallel for num_threads(nthreads)
for (int i = 0; i < ndets; ++i) {
int ioff = i * row_stride;
int64_t ioff = i * row_stride;

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ding


T* data_row = data + ioff;

Expand Down
15 changes: 15 additions & 0 deletions test/test_overflow_block.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
import unittest

import so3g
import numpy as np


class TestOverflowBlock(unittest.TestCase):

def test_overflow(self):
a = np.zeros((4050, 603260), dtype='float32')
so3g.block_minmax(a, a, 800, 2, 0)

@mhasself mhasself Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add a bunch of unit tests for the "large array limit" of all the array_ops functions? They could be disabled by default but it would be nice to run them, against this PR for example, or anytime we add new stuff.


if __name__ == '__main__':
unittest.main()