Transpose weights conservatively on a CUDA device

[roytsmart]

Lets transpose_weights_conservative() take weights built with device="cuda" and leave its result on the device, so a forward model and its transpose can both stay there.

Before

Device weights were brought back to the host. The function then returned host values beside device indices, which neither path of regrid_from_weights() can apply:

TypeError: can't unbox array from PyObject into native value.

Device path

regridding/_weights/_weights_transposed/_cuda.py:

• Normalization: a kernel multiplies each weight by a factor for its input cell and divides it by the volume of its output cell. The operations run in the same order as on the host, so the results agree bit for bit.
• Indices: the two index arrays are swapped without being copied, as on the host.
• Values: stored as float64, as on the host.
• Empty slots: they keep their index of -1, now on the output side. The device scatter already skips that side (since Add convolve_weights() to fold a point-spread function into a set of weights #67). Their values are set to zero rather than read.
• Factors: sent once, before being broadcast along the orthogonal axes; each element reads its own row on the device. _cuda.rows() also reduces any axis the caller already broadcast (a stride of zero) before sending. convolve_weights() now sends its kernels the same way, so a kernel the caller broadcast is still sent once, as on main.
• Corrupted weights: an index past the end of the grids the weights say they were built for is flagged rather than read, and raises ValueError.

transpose_weights() already worked on device weights. Its docstring now says so, and a test covers it.

Shared by host and device

• Grid check: the grids' cells are compared once, up front, with the shape the weights were built for. Grids with a different number of cells along any axis, swapped axes included, raise ValueError instead of being read at the wrong cells. The message names the likely causes: the wrong grids or axes, non-conservative weights such as multilinear ones (which address vertices), or weights transposed since without swapping the grids. Grids with more resampled axes than the weights have axes raise the same way. Grids whose orthogonal axes cannot broadcast against the weights' raise their own ValueError.
• Per-cell factor: the volume of each input cell is divided by its weights_input squared once, on the host, rather than for every weight that touches it. With a weights_input, host results therefore differ from main by up to about 4e-16 relative, from the order of operations; without one they are identical. A zero weights_input in a cell no weight touches doesn't warn, as on main.
• Volumes per grid: each grid's volumes are computed only along the orthogonal axes it varies along. Broadcast axes are detected from zero strides, so this also works when the caller has already broadcast the grids, as named-arrays does. Through ctis, the sensor grid's areas were computed for each of the 80 channels and wavelengths, and are now computed once.
• Units: _util._dimensionless() reduces a unitless quantity with a scale, such as 50 %, to the number it stands for, converting to float64 before scaling. An array the caller broadcast is converted at its own size and broadcast again, so no caller copies it out. It serves weights(), convolve_weights() and the transpose. _util._unbroadcast() is the one place that finds broadcast axes, for the transpose, _dimensionless() and _cuda.rows(). Before, the host transpose was wrong by the scale (100 for percent), and the device would have been wrong by its square. A unit with dimensions is dropped by the transpose, as before.
• Precision: weights_input is converted to float64 before it is squared, so a float32 one gives the same result on both paths. Its broadcast axes are found before that conversion, which would otherwise copy it out to its full shape.

Parallel cell areas

The sweep that computes the cell areas is a prange loop inlined into _cell_volume_2d, which was compiled without parallel=True, so it ran on one thread. Now each sweep is one parallel loop over every line of edges of every grid. Its areas are bit for bit those of grid_volume() on each grid alone. Many small grids cost no more to start than one large one, and a few large grids still use every thread.

cell areas	serial (before)	after
1 grid of 3000²	199 ms	42 ms
4 grids of 1024²	79 ms	17 ms
500 grids of 64²	11.3 ms	6.4 ms
5000 grids of 16²	11.0 ms	3.8 ms

_sweep.weights_conservative_2d calls the same inlined sweep without parallel=True, so the forward host build still computes its areas on one thread. Fixing that is left for a separate PR.

Timings

Warm, on an RTX 4090 and 48 host threads.

transpose	host before	host after	device
4.5M weights (512² cells × 4)	76 ms	68 ms	14 ms
17.9M (1024² × 4)	389 ms	202 ms	45 ms
35.3M (2048² × 2)	904 ms	388 ms	166 ms

MART end to end: ctis's MartInverter on the host, against the same algorithm on the device, using device weights and torch for the rest. The device result matches ctis to 1e-14, with the same iteration count and χ².

problem	weights	ctis per iteration	device per iteration	end to end
tutorial, no PSF	1.1M	44 ms	12 ms	2.7×
PSF tutorial (21×21 PSFs)	148M	190 ms	21 ms	6.4×
256² scene, no PSF	17M	498 ms	13 ms	18×

The device's time per iteration barely changes with size, because launching one kernel per orthogonal element dominates it. One launch for all elements would be the next lever.

Tests

_weights_transposed/_cuda_test.py (marked cuda). Its input grids differ in cell area from one element to the next, so a row read for the wrong element fails.

• Matches the host: exact equality with the host, with no weights_input, one shared by every element, one that varies per element, and a float32 one.
• Stays on the device: results are device arrays that share the weights' index arrays, and applying them gives the host's result.
• Narrow types: int32 indices are kept and values become float64.
• Single: weights with no orthogonal axes.
• Units and errors: a percentage weights_input gives the same result as its number; mismatched grids raise before anything is sent; corrupted weights raise.
• Sent once: the volumes are sent once, whether or not the caller broadcast the grids first, and so is a pre-broadcast float32 or percent weights_input. Convolved weights with more orthogonal elements than their grids match the host. In _weights_convolved/_cuda_test.py, a kernel the caller broadcast is sent once, in double or single precision or as a percentage.
• Edge cases: an element with no weights launches no kernel, and transpose_weights() works on device weights.

In _weights_transposed_test.py:

• Units: percentages in double and single precision match their numbers, and a unit with dimensions is dropped.
• Grid mismatches: swapped axes, a larger input grid and a smaller output grid all raise.
• Pre-broadcast grids: they compute each distinct grid's areas once and give the same result.
• Areas: they equal grid_volume() bit for bit for 1, 2 and 9 grids, and a stacked transpose matches each grid transposed alone.
• Other errors and edge cases: multilinear weights, grids with too many axes and non-broadcasting orthogonal axes raise; an empty orthogonal axis transposes into no weights; an untouched zero weights_input doesn't warn. In _weights_convolved_test.py, a kernel that isn't numbers raises numpy's own error.

Local runs:

• regridding/_tests/test_util.py covers _unbroadcast() and _dimensionless().
• The full suite passes (494 tests). The transpose and convolve tests also pass with NUMBA_DISABLE_JIT=1.
• black and ruff are clean. pyright 1.1.411 reports no errors in the new or changed code; the 19 errors in the existing host transpose tests are already on main.

🤖 Generated with Claude Code

https://claude.ai/code/session_01GqceVSp8s2msxj1Dr9MrGU

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 99.71%. Comparing base (ed218c2) to head (dd03516).

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #69      +/-   ##
==========================================
+ Coverage   99.68%   99.71%   +0.03%     
==========================================
  Files          62       65       +3     
  Lines        3835     4251     +416     
==========================================
+ Hits         3823     4239     +416     
  Misses         12       12              

Flag	Coverage Δ
unittests	99.71% <100.00%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.