mps: index_put_ (x[int_index] = v) on Metal with a device-side bounds check

.github/scripts/gate_suites.py

@@ -129,6 +129,7 @@
     "devices/mps/test_mpsextreme.py": BOTH,
     "devices/mps/test_mpsfwd.py": BOTH,
     "devices/mps/test_mpsindex.py": BOTH,
+    "devices/mps/test_mpsindexput.py": BOTH,
     "devices/mps/test_mpsint32.py": BOTH,
     "devices/mps/test_mpsinplace.py": BOTH,
     "devices/mps/test_mpsrefuse.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, **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)) |
+| `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, **87** 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

@@ -2568,3 +2568,45 @@ which needs a build (`.scratch/cbmps_i32.py` is the probe).
 <!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsint32.py test_int32_arithmetic_widens_on_the_device_and_agrees present -->
 <!-- DOCWATCH: symbol-in-file torchnative/rust/vendor/candle-metal-kernels/src/metal_src/cast.metal "init_cast_all(i32, int32_t)" present -->
 <!-- DOCWATCH: symbol-in-file torchnative/rust/vendor/candle-core/src/metal_backend/mod.rs cast_i64_i32 present -->
+
+### 7.21 `index_put_` on Metal: the write side of §7.19, same counted validation readback
+
+**What changed.** `aten.rs::index_put_inplace` is now a router. A CPU `self`
+goes to `index_put_on_host` (the kernel it always had); every other device
+goes to `index_put_on_device`, which places the indexed axes in front exactly
+as `index_tensor_on_device` does, views `self` as `(block, span)` and computes
+the flat index with the same `index_flat_offsets`. `accumulate=True` is one
+candle `index_add` (rows added in index order at the receiver's precision;
+bool through `i64` and a comparison with zero, a logical or).
+`accumulate=False` scatters each source row's *position* (float32, exact
+below 2^24, refused by name above) into a `-1` column, which gives the last
+row aimed at each destination, then `index_select`s those rows and `where`s
+them over `self`: candle's Metal `scatter` has no `i64` or `u8` value kernel,
+positions need neither. `MPS_HOST_READBACK_OPS` went from 88 to 87.
+
+**The readback budget** is §7.19's, because the bounds check is the same
+function (`index_bounds_on_device`) and it runs before anything is written:
+1 download for a non-empty write, 0 for an empty broadcast, 2 on an
+out-of-range call (which raises upstream's CPU `IndexError` and leaves `self`
+unchanged). Upstream's own mps raises asynchronously here (measured, torch
+2.13.0: `x[:, tensor([9])] = 1.0` on mps returns, and the `AcceleratorError`
+surfaces at exit), so upstream cpu is the oracle for the error rows.
+
+**Duplicates without accumulate have no fixed answer upstream.** Measured on
+2.13.0: `zeros(5)[[1, 1, 1, 3]] = [1, 2, 3, 4]` puts 3 at position 1 on cpu
+and 1 on mps; with 1000 duplicates per position mps gave 992 and 1984. This
+build resolves last-wins (cpu's answer) and the suite asserts only that each
+element holds one of the values aimed at it.
+
+**What still refuses, by name, with nothing read back.** A bool or uint8
+mask index that is not `__setitem__`'s receiver-shaped mask with a scalar
+(that one is lowered to `masked_fill_`, #99), and an `int32` index tensor on
+mps (no int32 Metal kernel in this base; PR #100's device widening is the
+follow-up). `index_put_` itself still refuses a `values` dtype different from
+`self`'s, as upstream does (measured: `Index put requires the source and
+destination dtypes match`); only `__setitem__`'s scalar lift narrows.
+
+<!-- DOCWATCH: symbol-in-file torchnative/rust/torch_c/src/aten.rs index_put_on_device present -->
+<!-- DOCWATCH: symbol-in-file torchnative/rust/torch_c/src/aten.rs index_flat_offsets present -->
+<!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsindexput.py test_out_of_range_indices_raise_and_leave_self_unchanged present -->
+<!-- DOCWATCH: symbol-in-file tests/devices/mps/test_mpsindexput.py test_only_the_router_calls_the_host_half 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 은 입구에서 이름으로 거부됩니다. **88** 개(이 칸은 한때 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 은 입구에서 이름으로 거부됩니다. **87** 개(이 칸은 한때 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/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, **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)) |
+| `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, **87** of them (88 until `aten.index_put_.default` moved onto the device for integer indices, [`docs/devices/matrix.md`](../devices/matrix.md) §7.21; 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)) |

tests/_support/test_shim.py

@@ -25772,6 +25772,12 @@ def test_mps_is_refused_by_name_where_it_is_not_compiled_in():
         "`index_tensor` when `self` is on the CPU; every other device goes "
         "to `index_tensor_on_device`. test_mpsindex.py checks the router is "
         "its only caller, and its counters turn red if an mps call reaches it.",
+    "index_put_on_host":
+        "the CPU half of aten.index_put_.default, reached only from the router "
+        "`index_put_inplace` when `self` is on the CPU; every other device goes "
+        "to `index_put_on_device`. test_mpsindexput.py checks the router is its "
+        "only caller and counts an mps call's downloads (1, the validation "
+        "vector).",
     "gather_on_host":
         "the CPU half of aten.gather.default, reached only from the router "
         "`gather_default` when `self` is on the CPU; test_mpsindex.py checks "

tests/devices/mps/test_metalplace.py

@@ -121,6 +121,12 @@
 
 _EXEMPT = ("aten._local_scalar_dense.default", "aten.uniform_.default")
 
+# In-place operators whose device kernel makes a designed, counted validation
+# readback (docs/devices/matrix.md §7.20): `index_put_` reads back each index
+# tensor's on-device min and max, one download per call, and nothing else.
+# Exactly that count is allowed; any other count is still a finding.
+_VALIDATION_DOWNLOADS = {"aten.index_put_.default": 1}
+
 
 def _counters():
     c = _C._metal_counters()
@@ -477,7 +483,8 @@ def test_every_in_place_operator_that_reaches_on_mps_reads_nothing_back():
         "operators reaching. Refusals:\n  %s" % (len(reached), "\n  ".join(refused)))
 
     unexpected = [(op, dt, d) for op, dt, d in downloaded
-                  if op not in _HOST_COMPUTED_ON_MPS and op not in _EXEMPT]
+                  if op not in _HOST_COMPUTED_ON_MPS and op not in _EXEMPT
+                  and d["host_downloads"] != _VALIDATION_DOWNLOADS.get(op)]
     assert not unexpected, (
         "these in-place operators reach on mps and compute on the host -- a "
         "new one, not among the eight docs/devices/matrix.md §7.14 names:\n  "

tests/devices/mps/test_mpsindex.py

@@ -310,7 +310,7 @@ def test_the_two_ops_left_the_readback_list_and_the_rest_did_not():
     for op in ("aten.index.Tensor", "aten.gather.default"):
         assert op not in refused, op
     for op in ("aten.scatter.src", "aten.sort.default", "aten.multinomial.default",
-               "aten.view.dtype", "aten.index_put_.default", "aten.masked_select.default"):
+               "aten.view.dtype", "aten.masked_select.default"):
         assert op in refused, op
     assert set(_C._shim_mps_readback_but_allowed()) == {
         "aten._local_scalar_dense.default", "aten.uniform_.default"}