mps index.Tensor / gather on Metal with device-side bounds check (#32, stacked on #84)

.github/scripts/gate_suites.py

@@ -123,8 +123,10 @@
     "devices/mps/test_mpsconst.py": BOTH,
     "devices/mps/test_mpsnarrow.py": BOTH,
     "devices/mps/test_mpsgen.py": BOTH,
+    "devices/mps/test_mpscb.py": BOTH,
     "devices/mps/test_mpsextreme.py": BOTH,
     "devices/mps/test_mpsfwd.py": BOTH,
+    "devices/mps/test_mpsindex.py": BOTH,
     "devices/mps/test_mpsinplace.py": BOTH,
     "devices/mps/test_mpsrefuse.py": BOTH,
     "devices/mps/test_mpsresidency.py": BOTH,

README.md

@@ -227,7 +227,7 @@ row, is [`docs/platform/STATUS.md`](docs/platform/STATUS.md).
 |---|:--:|:--:|:--:|:--:|:--:|:--:|---|
 | `cpu` | ✅ | ✅ | ⚠️ | ⚠️ | ⚠️ | ✅ | the only device that holds a tensor |
 | `meta` | ✅ | ✅ | ⚠️ | ⚠️ | ⚠️ | 🔲 | shape and dtype, no storage |
-| `mps` | ✅ | n/a | 🔲 | n/a | n/a | n/a | candle's Metal backend, on. An op whose kernel would compute on the **CPU under an `mps` label** is refused at the door, naming the op, **90** of them (this cell said 54), listed at runtime by `_C._shim_mps_host_readback_ops()`. A transformer forwards here ([`MPSATTN.md`](docs/devices/MPSATTN.md)) |
+| `mps` | ✅ | n/a | 🔲 | n/a | n/a | n/a | candle's Metal backend, on. An op whose kernel would compute on the **CPU under an `mps` label** is refused at the door, naming the op, **88** of them (this cell said 54), listed at runtime by `_C._shim_mps_host_readback_ops()`. A transformer forwards here ([`MPSATTN.md`](docs/devices/MPSATTN.md)) |
 | `vulkan` | ✅ | ❌ | n/a | 🔲 | 🔲 | n/a | thirty-one ops through real `VkBuffer`s; a pretrained BERT forwards and agrees with upstream ([`VULKAN7.md`](docs/devices/VULKAN7.md)) |
 | NNAPI · CoreML | ✅ *CoreML* | ✅ *NNAPI* | 🔲 | n/a | n/a | n/a | both execute; CoreML reaches the Neural Engine, NNAPI has met only a CPU driver |
 | `cuda` | n/a | n/a | n/a | ⚠️ | ⚠️ | n/a | wired, never compiled or run ([`CUDA.md`](docs/devices/CUDA.md)) |

docs/devices/matrix.md

@@ -244,6 +244,15 @@ the device too. It was recorded BREAKS in the table below.
 **It is closed now, by refusal rather than by a kernel.** The verdict on the
 `mps` and `cuda` columns is REFUSES, naming the op, the device and the reason.
 
+**A counted validation readback is not a host computation** (issue #32,
+§7.19). `aten.index.Tensor` and `aten.gather.default` left
+`MPS_HOST_READBACK_OPS` for their integer-index paths. They compute with
+candle device ops, and the only bytes they read back are each index tensor's
+minimum and maximum, reduced on the device and read as one short vector to
+bounds-check candle's unchecked Metal kernels: **1 download per call**, 2 on an
+out-of-range call (the second names the offender, then the call raises).
+`test_mpsindex.py` asserts those counts exactly.
+
 #### Why it stayed open, and what that says about the derivation
 
 `device.rs::MPS_HOST_READBACK_OPS` is exactly the gate for this shape and it did
@@ -1454,6 +1463,15 @@ refuse. So the refusal stays, and it changed sentence: it now names the
 *layout*, says what does work, and tells the reader to `.contiguous()`.
 Naming the backend is what sent 112 cells to the wrong diagnosis.
 
+**Lifted by issue #32.** A non-contiguous receiver is now written on the
+device as the contiguous runs it is made of: `write_on_device` iterates the
+leading axes that break contiguity and gives each run its own `slice_set`
+blit, so `x[:, 0:2].zero_()` and `mask[:, :, :q, :k].fill_(min)` stay on
+Metal and read nothing back (`tests/devices/mps/test_mpscb.py`, counters
+asserted). A layout needing more than 65536 runs (a transpose of a large
+matrix, whose runs are single elements) refuses by name with the run count.
+The cells of the tables above were not re-measured for this change.
+
 **Placement evidence, and it is the same ceiling as §7.5.** There is still no
 Metal dispatch counter, so this round cannot prove "the blit ran on the GPU"
 the way a `_vulkan_counters()` round can. The available standard is
@@ -1594,7 +1612,10 @@ rebinding the wrapper, which is the obvious wrong way to do this.
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_a_view_taken_before_the_write_sees_it_on_mps present -->
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_writing_into_an_offset_slice_on_mps_touches_only_that_slice present -->
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_a_self_overlapping_copy_on_mps_is_not_half_overwritten present -->
-<!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_a_strided_receiver_on_mps_refuses_and_names_its_layout present -->
+<!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_a_strided_receiver_on_mps_refuses_and_names_its_layout absent -->
+<!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_a_strided_receiver_on_mps_writes_every_run_on_the_device present -->
+<!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_a_strided_receiver_with_too_many_runs_refuses_by_name present -->
+<!-- DOCWATCH: symbol-in-file torchnative/rust/torch_c/src/tensor.rs contiguous_split present -->
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_the_device_write_door_performs_no_host_readback present -->
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_clamp_propagates_nan_on_mps_exactly_as_upstream_does present -->
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsinplace.py test_fill_on_mps_agrees_with_upstream present -->
@@ -2420,3 +2441,61 @@ They are listed as unexercised now, with the measured refusal as the reason.
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsconst.py test_the_sites_this_round_converted_are_accounted_for present -->
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsconst.py test_a_zero_dim_device_tensor_can_stand_in_for_a_scalar present -->
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsconst.py test_the_integer_arm_narrows_on_the_host_too present -->
+
+### 7.19 `index.Tensor` and `gather` on Metal: a counted validation readback, not a host computation (issue #32)
+
+**What changed.** `aten.rs::index_tensor` and `gather_default` are now routers.
+A CPU `self` goes to the host kernels it always had (`index_tensor_on_host`,
+`gather_on_host`); every other device goes to `index_tensor_on_device` and
+`gather_on_device`, which are candle device ops: negative indices wrapped with
+`where(idx < 0, idx + size, idx)`, flat offsets summed from the broadcast index
+tensors, then one `index_select` (or `gather` on a `self` narrowed to the
+index's extent off-axis). `MPS_HOST_READBACK_OPS` went from 90 to 88 (89 to 87 on the branch, before issue #33 added `aten.sort.stable`).
+
+**The readback budget.** candle's Metal `index_select` and `gather` kernels
+do not bounds-check, so an out-of-range index would read outside the buffer
+where upstream raises. The device paths reduce each index tensor to its min
+and max on the device and read that vector back through
+`index_validation_readback`, the only function on these paths that reads
+device bytes:
+
+| Call | Host downloads |
+|---|---|
+| non-empty result, any number of index tensors | **1** |
+| empty broadcast (nothing to gather) | **0** |
+| out of range (raises upstream's CPU error) | **2**: the second reads the indices to name the first offender in upstream's order |
+| refused (mask index, int32 index on mps) | **0** |
+
+That readback validates; it never feeds the answer. It is a named barrier in
+the derivation (`test_shim.py::_MPS_READBACK_EXEMPT`, beside the two host
+halves, whose only callers are the routers), and the exact counts are what
+keep it honest: routing an mps call to a host half, or reading an operand,
+changes the count and turns `test_mpsindex.py` red.
+
+**What still refuses, by name.** Boolean and uint8 mask indices on a device
+(the result's length is the mask's true count, which this build can only get
+from the host). `int32` index tensors on mps (no int32 Metal kernel, not even
+the cast). `int32` values on mps stay refused by
+`device::mps_int_dtype_refusal` as before. Value dtypes with Metal kernels:
+float32, float16, bfloat16, int64, int8, uint8/bool (gather widens uint8 to
+u32 on the device, because candle-metal-kernels has no `gather_i64_u8`).
+
+**Errors.** Upstream's CPU type and text: `IndexError: index 5 is out of bounds
+for dimension 0 with size 3`, where the dimension is the index tensor's
+position among the index tensors, not the axis (measured: `x[:, i]` on a
+`(2,3,4,5)` tensor says `dimension 0 with size 3`). The host kernel used to
+name the axis; it was corrected to match in the same round. `gather` raises
+`RuntimeError` and does not wrap negatives. Upstream's own mps raises
+`torch.AcceleratorError`, a `RuntimeError` subclass, with the same `index`
+text and `gather: index 5 is out of bounds for dimension with size 4` for
+gather (torch 2.13.0, measured); the CPU wording is the oracle.
+
+**The table in §6 is not re-measured here.** Its `aten.gather.default` and
+`aten.index.Tensor` rows still read REFUSES in the `mps` and `cuda` columns;
+those cells are stale until the sweep is re-run on a built tree.
+
+<!-- DOCWATCH: symbol-in-file torchnative/rust/torch_c/src/aten.rs index_validation_readback present -->
+<!-- DOCWATCH: symbol-in-file torchnative/rust/torch_c/src/aten.rs index_tensor_on_device present -->
+<!-- DOCWATCH: symbol-in-file torchnative/rust/torch_c/src/aten.rs gather_on_device present -->
+<!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsindex.py test_out_of_range_indices_raise_upstreams_error_on_mps present -->
+<!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsindex.py test_only_the_routers_call_the_host_halves present -->

docs/locale/README_ko.md

@@ -204,7 +204,7 @@ flowchart TB
 |---|:--:|:--:|:--:|:--:|:--:|:--:|---|
 | `cpu` | ✅ | ✅ | ⚠️ | ⚠️ | ⚠️ | ✅ | 텐서를 담는 유일한 장치 |
 | `meta` | ✅ | ✅ | ⚠️ | ⚠️ | ⚠️ | 🔲 | 형상과 dtype, 저장소 없음 |
-| `mps` | ✅ | n/a | 🔲 | n/a | n/a | n/a | candle 의 Metal 백엔드, 켜짐. **`mps` 라벨 아래 CPU 에서** 계산할 커널을 가진 op 은 입구에서 이름으로 거부됩니다. **90** 개(이 칸은 한때 54 라고 했습니다), 런타임에 `_C._shim_mps_host_readback_ops()` 로 나열. 트랜스포머가 여기서 forward 합니다 ([`MPSATTN.md`](../devices/MPSATTN.md)) |
+| `mps` | ✅ | n/a | 🔲 | n/a | n/a | n/a | candle 의 Metal 백엔드, 켜짐. **`mps` 라벨 아래 CPU 에서** 계산할 커널을 가진 op 은 입구에서 이름으로 거부됩니다. **88** 개(이 칸은 한때 54 라고 했습니다), 런타임에 `_C._shim_mps_host_readback_ops()` 로 나열. 트랜스포머가 여기서 forward 합니다 ([`MPSATTN.md`](../devices/MPSATTN.md)) |
 | `vulkan` | ✅ | ❌ | n/a | 🔲 | 🔲 | n/a | 실제 `VkBuffer` 를 통한 31 개 op; 사전학습 BERT 가 forward 하고 upstream 과 일치 ([`VULKAN7.md`](../devices/VULKAN7.md)) |
 | NNAPI · CoreML | ✅ *CoreML* | ✅ *NNAPI* | 🔲 | n/a | n/a | n/a | 둘 다 실행; CoreML 은 Neural Engine 에 닿고, NNAPI 는 CPU 드라이버만 만남 |
 | `cuda` | n/a | n/a | n/a | ⚠️ | ⚠️ | n/a | 배선됨, 컴파일도 실행도 된 적 없음 ([`CUDA.md`](../devices/CUDA.md)) |

docs/numerics/DTYPEDEV.md

@@ -274,6 +274,11 @@ Two rows are worth reading carefully:
   `f64` first and Metal has no `f64`. It now computes in the operand dtype, and
   the three cells are graded by value against upstream by
   `test_the_dtype_device_matrix_agrees_with_upstream`.
+* **`cumsum` on `mps` for `float32`, `float16` and `bfloat16` joined
+  `_MPS_REACHES` with issue #32**, for the same reason: it accumulated in
+  `f64`. Off the CPU it now accumulates in `opmath_in` (`float32`), which is
+  upstream mps's accumulator; the CPU keeps `f64`, as upstream's CPU kernel
+  does.
 
 ### 3.2 Dtypes this build cannot store at all
 

docs/platform/STATUS.md

@@ -220,7 +220,7 @@ module does not have, and a target this repo cannot check would ship unchecked
 |---|:--:|:--:|:--:|:--:|:--:|:--:|---|
 | `cpu` | ✅ | ✅ | ⚠️ | ⚠️ | ⚠️ | ✅ | the only device that holds a tensor |
 | `meta` | ✅ | ✅ | ⚠️ | ⚠️ | ⚠️ | 🔲 | shape and dtype, no storage |
-| `mps` | ✅ | n/a | 🔲 | n/a | n/a | n/a | candle's Metal backend, on. An op whose kernel would compute on the **CPU under an `mps` label** is refused at the door, naming the op, **90** of them (89 until issue #33 added `aten.sort.stable`; this cell said 54; it was 99 until issue #29 computed `argmax`, `max.default`, `isin` and the `bitwise_*` ops on the device). `aten.view.dtype` joined on 2026-09-19: it answered `mps` dispatches from the host and returned a **cpu** tensor, and the derivation could not see it because its readback lives in `tensor.rs` ([`docs/devices/matrix.md`](../devices/matrix.md) §4.1). It was 90 until `aten.abs.default`/`aten.abs_.default` had their integral path rewritten as candle ops on the device and left the set ([`docs/devices/matrix.md`](../devices/matrix.md) §7.3), listed by `_C._shim_mps_host_readback_ops()`. It was 89 until 2026-09-20, when the derivation stopped matching six helper names and started following the call graph: ten kernels reach `read_flat` through an un-named hop in the same file (`order_along`, `argsort_core`, `floor_divide_impl`, `norm_pow_walk`) or through another operator's kernel, and **twelve cells published AGREES on `mps` were being computed on the host** ([`docs/devices/matrix.md`](../devices/matrix.md) §7.16). All twelve are integral dtypes; the float columns of those ops were already not running on Metal. `aten._softmax.default` is **not** among them and a transformer forwards here ([`docs/devices/MPSATTN.md`](../devices/MPSATTN.md), [`docs/devices/MPS.md`](../devices/MPS.md)) |
+| `mps` | ✅ | n/a | 🔲 | n/a | n/a | n/a | candle's Metal backend, on. An op whose kernel would compute on the **CPU under an `mps` label** is refused at the door, naming the op, **88** of them (90 until issue #32 moved `index.Tensor` and `gather` onto the device with a counted validation readback, [`docs/devices/matrix.md`](../devices/matrix.md) §7.19; 89 until issue #33 added `aten.sort.stable`; this cell said 54; it was 99 until issue #29 computed `argmax`, `max.default`, `isin` and the `bitwise_*` ops on the device). `aten.view.dtype` joined on 2026-09-19: it answered `mps` dispatches from the host and returned a **cpu** tensor, and the derivation could not see it because its readback lives in `tensor.rs` ([`docs/devices/matrix.md`](../devices/matrix.md) §4.1). It was 90 until `aten.abs.default`/`aten.abs_.default` had their integral path rewritten as candle ops on the device and left the set ([`docs/devices/matrix.md`](../devices/matrix.md) §7.3), listed by `_C._shim_mps_host_readback_ops()`. It was 89 until 2026-09-20, when the derivation stopped matching six helper names and started following the call graph: ten kernels reach `read_flat` through an un-named hop in the same file (`order_along`, `argsort_core`, `floor_divide_impl`, `norm_pow_walk`) or through another operator's kernel, and **twelve cells published AGREES on `mps` were being computed on the host** ([`docs/devices/matrix.md`](../devices/matrix.md) §7.16). All twelve are integral dtypes; the float columns of those ops were already not running on Metal. `aten._softmax.default` is **not** among them and a transformer forwards here ([`docs/devices/MPSATTN.md`](../devices/MPSATTN.md), [`docs/devices/MPS.md`](../devices/MPS.md)) |
 | `vulkan` | ✅ | ❌ | n/a | 🔲 | 🔲 | n/a | thirty-one ops by name through real `VkBuffer`s (twenty-one SPIR-V shaders); a pretrained BERT forwards (eager, no mask) and agrees with upstream ([`docs/devices/VULKAN7.md`](../devices/VULKAN7.md)); every other op refuses naming itself |
 | NNAPI · CoreML | ✅ *CoreML* | ✅ *NNAPI* | 🔲 | n/a | n/a | n/a | **Both execute.** The capture layer is built and a whole model lowers through it. CoreML compiles and runs, and a float16 graph runs on the **Neural Engine**; the NNAPI blob replays through `ANeuralNetworksModel` on emulator and on a physical Snapdragon 8 Gen 2, but that device offers only the `nnapi-reference` CPU driver, so NNAPI is execution, **not acceleration** ([`docs/graph/NPU2.md`](../graph/NPU2.md)) |
 | `cuda` | n/a | n/a | n/a | ⚠️ | ⚠️ | n/a | **Wired, never run.** One `resolve()` arm, because `Device::Cuda` is already a variant of candle's enum, the `mps` shape, not the `vulkan` one, so **no kernel of ours**. Off unless built with `--cfg torch_c_cuda`, which reaches Linux and Windows only. When unavailable it names **which** of `not_built` / `no_driver` / `no_device` / `wrong_arch` / `unclassified` it is. Nothing has compiled it (this machine has no `nvcc`) and nothing has run it ([`docs/devices/CUDA.md`](../devices/CUDA.md)) |