build: make the 32-bit index axis a per-backend option (FF_INDEX32), default unchanged

.github/workflows/ci.yml

@@ -282,6 +282,23 @@ jobs:
  boundflags: "-DFF_STATIC_BOUNDS=0 -DFF_STATIC_BOUND_DCT2=1 -DFF_STATIC_BOUND_DST2=1"
  splineflags: "-DFF_STATIC_SPLINES=0 -DFF_STATIC_SPLINE_NEAREST=1 -DFF_STATIC_SPLINE_LINEAR=1 -DFF_STATIC_SPLINE_QUADRATIC=1 -DFF_STATIC_SPLINE_CUBIC=1"
  jobs: "2"
+ # The third axis, in its non-default position: FF_INDEX32=0 collapses
+ # the int32_t/int64_t offset dispatch onto one instantiation (see
+ # INDEXFLAGS in src/lib-cpu/Makefile). This leg is why the knob is a
+ # knob: an option that has only ever been built in its default
+ # position is a claim, not a configuration. All-Dynamic bounds and
+ # splines because they are the cheapest policy that still runs the
+ # whole suite -- what this leg is asking about is the index axis, and
+ # cross-policy agreement is the other legs' job. Same 59,886 checks:
+ # the offset width is an internal representation choice and no test
+ # result may depend on it.
+ - name: clang-index64
+ cxx: clang++
+ apt: clang
+ boundflags: "-DFF_STATIC_BOUNDS=0"
+ splineflags: "-DFF_STATIC_SPLINES=0"
+ indexflags: "-DFF_INDEX32=0"
+ jobs: "2"
  # g++ at -O1, not the suite's -O3, and deliberately so: g++ was
  # measured needing ~6.3 GB RSS on reg_field.cpp at -O3, so two such
  # objects in parallel would sit right at the edge of a 16 GB runner.
@@ -336,7 +353,8 @@ jobs:
  CXX="ccache ${{ matrix.cxx }}" \
  ${{ matrix.cxxflags && format('CXXFLAGS="{0}"', matrix.cxxflags) || '' }} \
  BOUNDFLAGS="${{ matrix.boundflags }}" \
- SPLINEFLAGS="${{ matrix.splineflags }}"
+ SPLINEFLAGS="${{ matrix.splineflags }}" \
+ ${{ matrix.indexflags && format('INDEXFLAGS="{0}"', matrix.indexflags) || '' }}
  - name: Show ccache stats
  if: always()
  run: ccache -s
@@ -575,12 +593,32 @@ jobs:
  # where posdef/resize/restrict/splinc were absent from MODULES and eleven
  # FF_CUDA:: symbols the hub calls unconditionally were undefined, with this
  # job green throughout. Do not reduce this job back to `make cuda`.
+ #
+ # Two legs, over the 32-bit index axis (INDEXFLAGS / FF_INDEX32, see
+ # src/lib-cuda/Makefile). `index32` is the shipping default and is the gate;
+ # `index64` builds the same library with the offset dispatch collapsed onto
+ # int64_t. The second leg exists because CUDA is where that axis is
+ # expensive -- it is x2 device code, x2 SASS and x2 ptxas memory -- and
+ # because a build option nobody has ever compiled in its non-default
+ # position is not an option. It is also strictly the *cheaper* of the two
+ # legs (roughly half the instantiations), and it makes the FFMEM table below
+ # a direct measurement of what the axis costs on this backend rather than an
+ # extrapolation from the host compiler. Both legs run the full compile + hub
+ # link + ldd check; neither is allowed to fail.
  build-cuda:
- name: build-cuda (compile + link)
+ name: build-cuda (${{ matrix.name }})
  needs: changes
  if: needs.changes.outputs.cuda == 'true'
  runs-on: ubuntu-latest
  timeout-minutes: 120
+ strategy:
+ fail-fast: false
+ matrix:
+ include:
+ - name: index32
+ indexflags: ""
+ - name: index64
+ indexflags: "-DFF_INDEX32=0"
  steps:
  - uses: actions/checkout@v5
  - name: Install CUDA toolkit and clang
@@ -607,6 +645,7 @@ jobs:
  run: |
  set -o pipefail
  make -C . cuda -j2 CXX=clang++ CXXFLAGS="-std=c++14 -O1" \
+ ${{ matrix.indexflags && format('INDEXFLAGS="{0}"', matrix.indexflags) || '' }} \
  NVCC="/usr/bin/time -f 'FFMEM %M kB %e s %C' nvcc" \
  2>&1 | tee /tmp/build-cuda.log
  - name: Link the hub against the CUDA backend
@@ -634,6 +673,7 @@ jobs:
  # Cross-policy correctness is the test-cpu matrix's job.
  run: |
  make -C . lib -j2 USE_CUDA=1 CXX=clang++ \
+ ${{ matrix.indexflags && format('INDEXFLAGS="{0}"', matrix.indexflags) || '' }} \
  BOUNDFLAGS="-DFF_STATIC_BOUNDS=0" \
  SPLINEFLAGS="-DFF_STATIC_SPLINES=0"
  - name: Confirm the hub resolves every backend symbol at load time
@@ -692,7 +732,7 @@ jobs:
  | sed -E 's#^FFMEM ([0-9]+) kB +([0-9.]+) s .*-o ([^ ]+).*#\1 kB \2 s \3#' \
  | sort -rn | head -40)
  echo "$table"
- { echo '### Peak nvcc RSS per module'; echo; echo '```'; \
+ { echo "### Peak nvcc RSS per module (${{ matrix.name }})"; echo; echo '```'; \
  echo "$table"; echo '```'; } >> "$GITHUB_STEP_SUMMARY"
 
  # No table at all means the compile died before any nvcc finished

CLAUDE.md

@@ -111,9 +111,10 @@ the CPU path is the tested source of truth and CUDA is **compile+link only**.
 ## CI
 
 `.github/workflows/ci.yml`, path-filtered. `codespell` always; `test-cpu` (a
-3-leg `BOUNDFLAGS`/`SPLINEFLAGS` matrix + a g++ leg), `sanitize` (ASan+UBSan)
-and `tsan` on kernels/cpu/hub changes; `test-hub` on hub changes; `build-cuda`
-and `compile-probe-cuda` on kernels/cuda changes.
+3-leg `BOUNDFLAGS`/`SPLINEFLAGS` matrix + an `INDEXFLAGS` leg + a g++ leg),
+`sanitize` (ASan+UBSan) and `tsan` on kernels/cpu/hub changes; `test-hub` on
+hub changes; `build-cuda` (two legs, one per `FF_INDEX32` position) and
+`compile-probe-cuda` on kernels/cuda changes.
 
 **The `tsan` leg is the only one that runs anything in parallel.** With the
 shipping `GRAIN_SIZE` (32768) every workload in `tests/lib-cpu/` is below the
@@ -161,6 +162,18 @@ pushpull's fully-static order×bound compile is nightly
  `BOUNDFLAGS` / `SPLINEFLAGS`. These live **outside** `CXXFLAGS` on purpose so
  that a `CXXFLAGS=` override (as CUDA CI does, to force `-O1`) cannot silently
  drop the policy.
+- **So is the 32-bit index axis**, via `INDEXFLAGS` / `FF_INDEX32`
+ (`core/dispatch.h`) — the third member of that family and the most expensive
+ of the three: every templated kernel is templated on `offset_t`, whose two
+ values are chosen per call by `canUse32BitIndexMath`, so the narrow path is
+ exactly ×2 instantiations of everything below the dispatch layer.
+ `INDEXFLAGS="-DFF_INDEX32=0"` collapses both arms onto `int64_t`. Same
+ outside-`CXXFLAGS` rule, and **the default (on) is set separately in
+ `src/lib-cpu/Makefile` and `src/lib-cuda/Makefile`** — that per-library
+ default is what makes it a per-backend option, so do not hoist it into
+ `make/common.mk`. The narrow path is an inherited ATen register-pressure
+ optimisation that has never been benchmarked here (no GPU in CI); the
+ default does not move without one.
 - **CUDA memory limits are measured, not guessed** — and the numbers that used
  to be recorded here were wrong. CUDA CI forces `-O1` and `-j2` because
  `ptxas` was OOM-killed at ~16 GB on `reg_field.cpp`, and `src/lib-cuda`'s

MIGRATION.md

@@ -197,6 +197,39 @@ are identical either way. Applied to the regularisers only so far —
 `resize`/`restrict`/`splinc`/`pushpull` still use the static `bound::utils<B>`.
 See fastfields-lib#43.
 
+## The 32-bit index axis (`FF_INDEX32`)
+
+The third build-time axis, and the widest. Every templated kernel below the
+dispatch layer is templated on `offset_t`, and `offset_t` has exactly two
+values, chosen per call by `canUse32BitIndexMath`: `int32_t` when every
+operand's largest element offset fits in 32 bits, `int64_t` otherwise. So the
+narrow path costs exactly **x2 instantiations of everything** underneath —
+x2 device code, x2 SASS and x2 ptxas memory on the CUDA side. `core/autocast.h`
+(`copy_if_needed` / `free_if_needed`) exists solely to feed that second
+instantiation with narrowed shape and stride arrays.
+
+`FF_INDEX32=0` (`core/dispatch.h`) makes the narrow arm of every index
+dispatch name `int64_t` too: the two arms become one instantiation, the axis
+collapses, and the `canUse32BitIndexMath` probe folds away with it. Results are
+identical; code size, compile cost and per-voxel speed move.
+
+Which position each backend takes is a **per-library** default —
+`INDEXFLAGS` in `src/lib-cpu/Makefile` and `src/lib-cuda/Makefile`, exactly as
+`BOUNDFLAGS`/`SPLINEFLAGS` are — so the CPU library can drop the axis while the
+CUDA library keeps it, or the reverse. **Both default to on**, i.e. unchanged
+behaviour, and CI builds both positions on both backends.
+
+The open question is whether the narrow path is worth its cost. It is an ATen
+inheritance (register pressure) and has never been benchmarked in this project,
+because there is no GPU in CI. Note the shape of the trade before flipping
+anything: the backend that keeps the axis is CUDA, and CUDA is also the one
+with no build headroom left (`reg_flow` at 12.98 GB of a 16 GB runner). On the
+narrow path CUDA additionally pays a `cudaMallocHost`/`cudaFreeHost` per array
+per call, while 365 of 365 relevant `impl/cuda` upload sites use the
+*synchronous* `copyToDevice` — so the pinning has no async copy to enable and
+may partly offset the register-pressure win. Measured numbers and the exact
+benchmark that would settle it: fastfields-lib#94 and the follow-up issue.
+
 ## Porting pattern (per module)
 
 Use `distance.{h,cpp}` at each level as the template.

include/fastfields/api/cpu/pushpull_dispatch.h

@@ -37,9 +37,9 @@ typedef double reduce_t;
 #define FF_PP_DTYPE(D, I, B, FN, args...) \
  switch (code) { \
  case kDLFloat: switch (bits) { \
- case 32: return (use_32bits ? FN<D,I,B,float, int32_t>(args) \
+ case 32: return (use_32bits ? FN<D,I,B,float, off32_t>(args) \
  : FN<D,I,B,float, int64_t>(args)); \
- case 64: return (use_32bits ? FN<D,I,B,double,int32_t>(args) \
+ case 64: return (use_32bits ? FN<D,I,B,double,off32_t>(args) \
  : FN<D,I,B,double,int64_t>(args)); \
  default: break; \
  }; default: break; \

include/fastfields/api/cuda/pushpull_dispatch.h

@@ -37,9 +37,9 @@ typedef double reduce_t;
 #define FF_PP_DTYPE(D, I, B, FN, args...) \
  switch (code) { \
  case kDLFloat: switch (bits) { \
- case 32: return (use_32bits ? FN<D,I,B,float, int32_t>(args) \
+ case 32: return (use_32bits ? FN<D,I,B,float, off32_t>(args) \
  : FN<D,I,B,float, int64_t>(args)); \
- case 64: return (use_32bits ? FN<D,I,B,double,int32_t>(args) \
+ case 64: return (use_32bits ? FN<D,I,B,double,off32_t>(args) \
  : FN<D,I,B,double,int64_t>(args)); \
  default: break; \
  }; default: break; \

include/fastfields/core/dispatch.h

@@ -64,9 +64,92 @@
 // at the call site instead of hiding in one file's private prologue.
 #define FF_CVOIDPTR_OR_NULL(x) (x.data ? FF_CVOIDPTR(x) : nullptr)
 
+/***********************************************************************
+ * THE 32-BIT INDEX AXIS *
+ ***********************************************************************/
+
+/**
+ * `FF_INDEX32` -- compile-time policy for the 32-bit index (`offset_t`) axis,
+ * the exact analogue of `FF_STATIC_BOUNDS` (impl/kernels/bounds.h) and
+ * `FF_STATIC_SPLINES` (impl/kernels/spline.h) one axis further out.
+ *
+ * Every templated kernel below the dispatch layer is templated on `offset_t`,
+ * and `offset_t` has exactly two values chosen per call by
+ * `canUse32BitIndexMath`: `int32_t` when every operand's largest element
+ * offset fits in 32 bits, `int64_t` otherwise. The narrow one exists to cut
+ * register pressure -- an optimisation inherited from ATen -- and it costs
+ * exactly x2 instantiations of everything underneath, hence (on CUDA) x2
+ * device code and x2 ptxas memory. Measured on `reg_flow`, the module that
+ * peaks at 12.98 GB of a 16 GB runner: dropping the axis is -50.3%
+ * instantiations and -44.6% peak RSS (fastfields-lib#94).
+ *
+ * FF_INDEX32=1 (default, and today's behaviour on both backends)
+ * both arms exist; `canUse32BitIndexMath` picks per call.
+ * FF_INDEX32=0 the narrow arm names `int64_t` too, so the two arms are
+ * the same instantiation and the axis collapses. The
+ * `canUse32BitIndexMath` call folds away with it. Results
+ * are identical either way; only code size, compile cost and
+ * per-voxel speed move.
+ *
+ * It is ONE switch rather than an `FF_INDEX32_CPU` / `FF_INDEX32_CUDA` pair on
+ * purpose. `core/` is backend-agnostic by contract -- `src/lib-cpu` (host
+ * compiler) and `src/lib-cuda` (nvcc) compile this same header -- so a
+ * per-backend *name* would force the header to branch on `__CUDACC__`, which
+ * puts the policy in the source instead of in the build and would have to be
+ * repeated by any third consumer. BOUNDFLAGS/SPLINEFLAGS already establish the
+ * alternative and this follows it exactly: one macro, and the *per-library
+ * Makefile* chooses the default. That is what makes the option per-backend,
+ * and it is why the CPU library can drop the axis while the CUDA library keeps
+ * it (or the reverse) without either one knowing about the other. See
+ * INDEXFLAGS in src/lib-cpu/Makefile and src/lib-cuda/Makefile.
+ *
+ * NB the axis is genuinely unbenchmarked here: there is no GPU in CI, so the
+ * register-pressure win the narrow path is meant to buy has never been
+ * measured in this project. The default therefore stays where it has always
+ * been (on); this is a knob, not a decision.
+ */
+#ifndef FF_INDEX32
+# define FF_INDEX32 1
+#endif
+
+#if (FF_INDEX32 != 0) && (FF_INDEX32 != 1)
+# error "FF_INDEX32 must be 0 or 1 (see include/fastfields/core/dispatch.h)"
+#endif
+
 // Can this tensor's shape/stride arithmetic be narrowed to 32-bit offsets?
-// See core/autocast.h.
-#define FF_CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides))
+// See core/autocast.h. With the axis off there is nothing to narrow to, so
+// this is a compile-time `false` and the O(ndim) probe disappears from every
+// dispatch site.
+#if FF_INDEX32
+# define FF_CANUSE32BITS(x) (canUse32BitIndexMath(x.ndim, x.shape, x.strides))
+#else
+# define FF_CANUSE32BITS(x) (false)
+#endif
+
+FF_NAMESPACE_BEGIN(FF_NS)
+
+/**
+ * The offset type the *narrow* arm of every index dispatch names, i.e. the
+ * `offset_t` template argument on the `use_32bits ? f<..,off32_t>(a)
+ * : f<..,int64_t>(a)` sites.
+ *
+ * A typedef and not a macro, per the `FF_`-prefix rule in CLAUDE.md: a name in
+ * `ff::` is collision-safe with no prefix at all, and the dispatch sources are
+ * all inside `ff::cpu` / `ff::cuda`, so it resolves unqualified exactly where
+ * `int32_t` used to be spelled.
+ *
+ * When `FF_INDEX32` is 0 this is `int64_t` -- deliberately the same type as
+ * the wide arm, which is precisely how the axis collapses: both arms then name
+ * one instantiation, `use_32bits` is a compile-time `false`, and the ternary
+ * has nothing left to choose between.
+ */
+#if FF_INDEX32
+typedef int32_t off32_t;
+#else
+typedef int64_t off32_t;
+#endif
+
+FF_NAMESPACE_END(FF_NS)
 
 /***********************************************************************
  * CHECKS *

src/lib-cpu/Makefile

@@ -31,6 +31,29 @@ BOUNDFLAGS ?=
 # Kept out of CXXFLAGS for the same reason as BOUNDFLAGS.
 SPLINEFLAGS ?=
 
+# 32-bit index compile policy (see core/dispatch.h) -- the same idea as
+# BOUNDFLAGS/SPLINEFLAGS, one axis further out again, and the most expensive of
+# the three: EVERY templated kernel below this layer is templated on `offset_t`,
+# whose two values (int32_t / int64_t) are chosen per call by
+# canUse32BitIndexMath. So the narrow path costs exactly x2 instantiations of
+# everything, measured at -50.3% instantiations / -44.6% peak compiler RSS on
+# reg_flow when dropped (fastfields-lib#94).
+#
+# make INDEXFLAGS="-DFF_INDEX32=0" # one offset type (int64_t) everywhere
+#
+# Results are identical either way; code size, compile cost and per-voxel speed
+# move. The narrow path is an inherited register-pressure optimisation that has
+# never been benchmarked in this project, and it matters least here: the CPU
+# backend is not the one whose build is memory-constrained, and it has 64-bit
+# integer units. It stays ON by default all the same -- flipping it is a
+# performance decision and wants a benchmark, not a build tidy-up.
+#
+# This is a SEPARATE default from src/lib-cuda/Makefile's, which is the whole
+# point of putting it here rather than in make/common.mk: the two backends can
+# take the axis independently. Kept out of CXXFLAGS for the same reason as
+# BOUNDFLAGS/SPLINEFLAGS -- a CXXFLAGS= override must not silently drop policy.
+INDEXFLAGS ?= -DFF_INDEX32=1
+
 # `test` target only: a sparser *default* BOUNDFLAGS/SPLINEFLAGS than the
 # library's (empty = fully static). pushpull's own dispatch used to carry a
 # second, hand-duplicated switch behind `-DFF_TEST_SPARSE` purely to keep a
@@ -88,7 +111,7 @@ $(TARGET): $(OBJECTS) | $(LIBDIR)
 # header (this is a header-only codebase) rebuilds the affected library object
 # instead of leaving a stale binary.
 $(OBJDIR)/%.$(MOSUF): %.cpp | $(OBJDIR)
-	$(CXX) $(CXXFLAGS) $(BOUNDFLAGS) $(SPLINEFLAGS) $(INCLUDES) $(PICFLAG) -MMD -MP -c -o $@ $<
+	$(CXX) $(CXXFLAGS) $(BOUNDFLAGS) $(SPLINEFLAGS) $(INDEXFLAGS) $(INCLUDES) $(PICFLAG) -MMD -MP -c -o $@ $<
 
 ########################################################################
 #	Tests
@@ -123,7 +146,7 @@ TESTDRVOBJ = $(patsubst $(TESTSRCDIR)/%.cpp,$(TESTOBJDIR)/%.$(MOSUF),$(TESTSRC))
 .SECONDARY: $(TESTMODOBJ) $(TESTDRVOBJ)
 
 # -MMD -MP emit header dependency files (*.d) so header edits trigger rebuilds.
-TESTCPPFLAGS = $(CXXFLAGS) $(BOUNDFLAGS) $(SPLINEFLAGS) -DFF_TEST_SPARSE $(INCLUDES) -MMD -MP
+TESTCPPFLAGS = $(CXXFLAGS) $(BOUNDFLAGS) $(SPLINEFLAGS) $(INDEXFLAGS) -DFF_TEST_SPARSE $(INCLUDES) -MMD -MP
 
 $(TESTOBJDIR):
 	$(MKDIR) $(TESTOBJDIR)