Skip to content

Scope the narrowed index arrays, and stop leaking them on the throwing path#159

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

Merged
balbasty merged 2 commits into from
Aug 20, 2026
Merged

Scope the narrowed index arrays, and stop leaking them on the throwing path #159

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
50d7cab
fix(core): scope the narrowed index arrays, and stop leaking them
claude Aug 20, 2026
7c10367
style: drop the blank lines the free_if_needed removal left behind
claude Aug 20, 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
115 changes: 115 additions & 0 deletions include/fastfields/core/autocast.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
#pragma once
#include <cstddef>
#include <cstdint>
// `hostNew` / `hostDelete` throw std::runtime_error on the pinned-host
// (cudaMallocHost) path. This header never included <stdexcept> for them: it
// compiled only because every translation unit that includes it happens to
// include <stdexcept> first. A standalone .cu that includes just this header
// fails to compile without this line.
#include <stdexcept>
#include <fastfields/core/dlpack.h>
#include <fastfields/core/cuda_switch.h>

Expand DownExpand Up@@ -195,5 +201,114 @@ inline void free_if_needed(OutPointer ptr)
_copy_if_needed<OutPointer, InpPointer>::free(ptr);
}

/***********************************************************************
* THE SAME THING, BUT SCOPED *
***********************************************************************/

/**
* `IndexArray<offset_t>` is the RAII form of the `copy_if_needed` /
* `free_if_needed` pair above, for the one job that pair is actually used
* for: handing a `const int64_t *` shape or stride array to a kernel
* templated on `offset_t`.
*
* Two things it fixes.
*
* **The manual pair leaks whenever anything between the two throws**, and in
* every dispatch wrapper in this tree something between them can. The shape
* is always
*
* const offset_t * _size = copy_if_needed<offset_t*>(size, n); // allocates
* ... as_weights(), new reduce_t[], the impl call ... // can throw
* free_if_needed<int64_t*>(_size); // skipped
*
* and on the CUDA side the impl call throws by design -- every
* `FF_CUDA_LAUNCH` does, and so does every `copyToDevice`. Measured with
* ASan/LSan on the `reg_field` wrapper shape: 32 bytes in 2 allocations
* escape per throwing call on the narrow arm. It is a per-call leak on a
* path user code is expected to hit (an invalid argument), not a
* once-per-process one.
*
* **It does not allocate at all for the sizes that actually occur.** These
* arrays are `nbatch + ndim + 1` long -- single digits in every caller -- so
* the elements live in the object itself and the narrow arm costs no
* allocator traffic. Only a genuinely large rank falls back to `hostNew`.
* That matters most under `FF_AUTOCAST_PINNED_HOST`, where the fallback is
* `cudaMallocHost`: a page-locking syscall, on the per-call path, three to
* five times per launch.
*
* The wide arm (`offset_t == int64_t`) is specialised below to borrow the
* caller's array verbatim -- no copy, no storage -- exactly as
* `_copy_if_needed<T*, const T*>` already does for that case. So with
* `FF_INDEX32=0`, where both arms name `int64_t`, this whole class is a
* pointer copy.
*
* Deliberately not convertible to a *non*-const pointer: nothing in the
* dispatch layer writes through a shape or stride array, and the wide arm
* aliases the caller's memory.
*/

// How many elements live inside the object before it reaches for the heap.
// `nbatch + ndim + 1` with ndim <= 3; 8 covers every shape this tree builds
// and keeps the object at 32 bytes on the narrow arm.
#ifndef FF_INDEX_ARRAY_INLINE
# define FF_INDEX_ARRAY_INLINE 8
#endif

template <class offset_t>
class IndexArray
{
public:
// A null `src` yields a null array and copies nothing. An absent
// optional operand is passed as a null-data descriptor whose stride
// array is never read, and the call sites that have one spelled it
// `wgt ? copy_if_needed<offset_t*>(stride_wgt, n) : nullptr`.
IndexArray(const int64_t * src, size_t numel)
: _ptr(nullptr), _heap(nullptr)
{
if (!src) return;
offset_t * dst;
if (numel <= FF_INDEX_ARRAY_INLINE) {
dst = _inln;
} else {
dst = _heap = hostNew<offset_t>(numel);
}
for (size_t i = 0; i < numel; ++i)
dst[i] = static_cast<offset_t>(src[i]);
_ptr = dst;
}

~IndexArray() { if (_heap) hostDelete<offset_t>(_heap); }

operator const offset_t * () const { return _ptr; }
const offset_t * get() const { return _ptr; }

private:
IndexArray(const IndexArray &);
IndexArray & operator=(const IndexArray &);

const offset_t * _ptr;
offset_t * _heap;
offset_t _inln[FF_INDEX_ARRAY_INLINE];
};

// Wide arm: there is nothing to narrow to, so borrow the caller's array.
// Same no-op the `_copy_if_needed<T*, const T*>` specialisation provides, and
// the reason `FF_INDEX32=0` costs nothing here either.
template <>
class IndexArray<int64_t>
{
public:
IndexArray(const int64_t * src, size_t /* numel */) : _ptr(src) {}

operator const int64_t * () const { return _ptr; }
const int64_t * get() const { return _ptr; }

private:
IndexArray(const IndexArray &);
IndexArray & operator=(const IndexArray &);

const int64_t * _ptr;
};

FF_NAMESPACE_END(FF_DEVICE)
FF_NAMESPACE_END(FF_NS)
87 changes: 28 additions & 59 deletions src/lib-cpu/distance.cpp
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -51,14 +51,12 @@ inline void _dt_euclidean(
const int64_t * size , // [ndim] data shape == (*batch, n)
const int64_t * stride ) // [ndim] data strides
{
const offset_t * _size = copy_if_needed<offset_t *>(size, ndim);
const offset_t * _stride = copy_if_needed<offset_t *>(stride, ndim);
const IndexArray<offset_t> _size (size, ndim);
const IndexArray<offset_t> _stride (stride, ndim);
scalar_t * _f = static_cast<scalar_t *>(f);
const offset_t _ndim = static_cast<offset_t >(ndim);
const scalar_t _w = static_cast<scalar_t >(w);
distance_e::dt(_ndim, _f, _w, _size, _stride);
free_if_needed<int64_t *>(_size);
free_if_needed<int64_t *>(_stride);
distance_e::dt(_ndim, _f, _w, _size.get(), _stride.get());
}
}

Expand DownExpand Up@@ -92,14 +90,12 @@ inline void _dt_l1(
const int64_t * size , // [ndim] data shape == (*batch, n)
const int64_t * stride ) // [ndim] data strides
{
const offset_t * _size = copy_if_needed<offset_t *>(size, ndim);
const offset_t * _stride = copy_if_needed<offset_t *>(stride, ndim);
const IndexArray<offset_t> _size (size, ndim);
const IndexArray<offset_t> _stride (stride, ndim);
scalar_t * _f = static_cast<scalar_t *>(f);
const offset_t _ndim = static_cast<offset_t >(ndim);
const scalar_t _w = static_cast<scalar_t >(w);
distance_l1::dt(_ndim, _f, _w, _size, _stride);
free_if_needed<int64_t *>(_size);
free_if_needed<int64_t *>(_stride);
distance_l1::dt(_ndim, _f, _w, _size.get(), _stride.get());
}
}

Expand DownExpand Up@@ -190,12 +186,12 @@ inline void _dt_spline_table(
// stride_loc -> loc (*batch, ndim) == nbatch+1
// stride_coeff -> coeff (*batch, npoints, ndim) == nbatch+2
// stride_times -> times (*batch, ntimes) == nbatch+1
const offset_t * _size = copy_if_needed<offset_t *>(size, nbatch+2);
const offset_t * _int64_time = copy_if_needed<offset_t *>(int64_time, nbatch);
const offset_t * _stride_dist = copy_if_needed<offset_t *>(stride_dist, nbatch);
const offset_t * _stride_loc = copy_if_needed<offset_t *>(stride_loc, nbatch+1);
const offset_t * _stride_coeff= copy_if_needed<offset_t *>(stride_coeff, nbatch+2);
const offset_t * _int64_times= copy_if_needed<offset_t *>(int64_times, nbatch+1);
const IndexArray<offset_t> _size (size, nbatch+2);
const IndexArray<offset_t> _int64_time (int64_time, nbatch);
const IndexArray<offset_t> _stride_dist (stride_dist, nbatch);
const IndexArray<offset_t> _stride_loc (stride_loc, nbatch+1);
const IndexArray<offset_t> _stride_coeff(stride_coeff, nbatch+2);
const IndexArray<offset_t> _int64_times(int64_times, nbatch+1);
scalar_t * _time = static_cast< scalar_t *>(time);
scalar_t * _dist = static_cast< scalar_t *>(dist);
const scalar_t * _loc = static_cast<const scalar_t *>(loc);
Expand All@@ -212,12 +208,6 @@ inline void _dt_spline_table(
_size, _int64_time, _stride_dist, _stride_loc, _stride_coeff, _int64_times,
_spline, _bound
);
free_if_needed<int64_t *>(_size);
free_if_needed<int64_t *>(_int64_time);
free_if_needed<int64_t *>(_stride_dist);
free_if_needed<int64_t *>(_stride_loc);
free_if_needed<int64_t *>(_stride_coeff);
free_if_needed<int64_t *>(_int64_times);
}
}

Expand DownExpand Up@@ -312,11 +302,11 @@ inline void _dt_spline_brent(
// stride_dist -> dist (*batch) == nbatch
// stride_loc -> loc (*batch, ndim) == nbatch+1
// stride_coeff -> coeff (*batch, npoints, ndim) == nbatch+2
const offset_t * _size = copy_if_needed<offset_t *>(size, nbatch+2);
const offset_t * _int64_time = copy_if_needed<offset_t *>(int64_time, nbatch);
const offset_t * _stride_dist = copy_if_needed<offset_t *>(stride_dist, nbatch);
const offset_t * _stride_loc = copy_if_needed<offset_t *>(stride_loc, nbatch+1);
const offset_t * _stride_coeff= copy_if_needed<offset_t *>(stride_coeff, nbatch+2);
const IndexArray<offset_t> _size (size, nbatch+2);
const IndexArray<offset_t> _int64_time (int64_time, nbatch);
const IndexArray<offset_t> _stride_dist (stride_dist, nbatch);
const IndexArray<offset_t> _stride_loc (stride_loc, nbatch+1);
const IndexArray<offset_t> _stride_coeff(stride_coeff, nbatch+2);
scalar_t * _time = static_cast< scalar_t *>(time);
scalar_t * _dist = static_cast< scalar_t *>(dist);
const scalar_t * _loc = static_cast<const scalar_t *>(loc);
Expand All@@ -334,11 +324,6 @@ inline void _dt_spline_brent(
_size, _int64_time, _stride_dist, _stride_loc, _stride_coeff,
_max_iter, _tol, _step, _spline, _bound
);
free_if_needed<int64_t *>(_size);
free_if_needed<int64_t *>(_int64_time);
free_if_needed<int64_t *>(_stride_dist);
free_if_needed<int64_t *>(_stride_loc);
free_if_needed<int64_t *>(_stride_coeff);
}
}

Expand DownExpand Up@@ -428,11 +413,11 @@ inline void _dt_spline_gaussnewton(
// stride_dist -> dist (*batch) == nbatch
// stride_loc -> loc (*batch, ndim) == nbatch+1
// stride_coeff -> coeff (*batch, npoints, ndim) == nbatch+2
const offset_t * _size = copy_if_needed<offset_t *>(size, nbatch+2);
const offset_t * _int64_time = copy_if_needed<offset_t *>(int64_time, nbatch);
const offset_t * _stride_dist = copy_if_needed<offset_t *>(stride_dist, nbatch);
const offset_t * _stride_loc = copy_if_needed<offset_t *>(stride_loc, nbatch+1);
const offset_t * _stride_coeff= copy_if_needed<offset_t *>(stride_coeff, nbatch+2);
const IndexArray<offset_t> _size (size, nbatch+2);
const IndexArray<offset_t> _int64_time (int64_time, nbatch);
const IndexArray<offset_t> _stride_dist (stride_dist, nbatch);
const IndexArray<offset_t> _stride_loc (stride_loc, nbatch+1);
const IndexArray<offset_t> _stride_coeff(stride_coeff, nbatch+2);
scalar_t * _time = static_cast< scalar_t *>(time);
scalar_t * _dist = static_cast< scalar_t *>(dist);
const scalar_t * _loc = static_cast<const scalar_t *>(loc);
Expand All@@ -449,11 +434,6 @@ inline void _dt_spline_gaussnewton(
_size, _int64_time, _stride_dist, _stride_loc, _stride_coeff,
_max_iter, _tol, _spline, _bound
);
free_if_needed<int64_t *>(_size);
free_if_needed<int64_t *>(_int64_time);
free_if_needed<int64_t *>(_stride_dist);
free_if_needed<int64_t *>(_stride_loc);
free_if_needed<int64_t *>(_stride_coeff);
}
}

Expand DownExpand Up@@ -602,11 +582,11 @@ _dt_mesh(
{
// vertices (N, D) and faces (M, D) are always 2D (the impl only reads
// stride[0] and stride[1]); their length is independent of loc's batch rank.
const offset_t * _size = copy_if_needed<offset_t *>(size, nbatch);
const offset_t * _stride_dist = copy_if_needed<offset_t *>(stride_dist, nbatch);
const offset_t * _stride_coord = copy_if_needed<offset_t *>(stride_coord, nbatch+1);
const offset_t * _stride_vertices= copy_if_needed<offset_t *>(stride_vertices,2);
const offset_t * _stride_faces = copy_if_needed<offset_t *>(stride_faces, 2);
const IndexArray<offset_t> _size (size, nbatch);
const IndexArray<offset_t> _stride_dist (stride_dist, nbatch);
const IndexArray<offset_t> _stride_coord (stride_coord, nbatch+1);
const IndexArray<offset_t> _stride_vertices(stride_vertices,2);
const IndexArray<offset_t> _stride_faces (stride_faces, 2);
scalar_t * _dist = static_cast< scalar_t *>(dist);
index_t * _nearest_vertex = static_cast< index_t *>(nearest_vertex);
const scalar_t * _coord = static_cast<const scalar_t *>(coord);
Expand All@@ -616,11 +596,7 @@ _dt_mesh(
const offset_t _nb_faces = static_cast< offset_t >(nb_faces);
const offset_t _nb_vertices = static_cast< offset_t >(nb_vertices);

const offset_t * _stride_nearest = (
stride_nearest
? copy_if_needed<offset_t *>(stride_nearest, nbatch)
: nullptr
);
const IndexArray<offset_t> _stride_nearest(stride_nearest, nbatch);

distance_mesh::dt<
ndim, scalar_t, index_t, offset_t
Expand All@@ -630,13 +606,6 @@ _dt_mesh(
_stride_dist, _stride_nearest, _stride_coord, _stride_vertices, _stride_faces,
_signed, naive
);

free_if_needed<int64_t *>(_size);
free_if_needed<int64_t *>(_stride_dist);
free_if_needed<int64_t *>(_stride_coord);
free_if_needed<int64_t *>(_stride_vertices);
free_if_needed<int64_t *>(_stride_faces);
if (_stride_nearest) free_if_needed<int64_t *>(_stride_nearest);
}

void dt_mesh(
Expand Down
Loading
Loading