diff --git a/.github/scripts/gate_suites.py b/.github/scripts/gate_suites.py index a21e33e..08fd57d 100644 --- a/.github/scripts/gate_suites.py +++ b/.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, diff --git a/README.md b/README.md index 918ffc9..7c32266 100644 --- a/README.md +++ b/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)) | diff --git a/docs/devices/matrix.md b/docs/devices/matrix.md index 8abf4e3..d662639 100644 --- a/docs/devices/matrix.md +++ b/docs/devices/matrix.md @@ -2568,3 +2568,45 @@ which needs a build (`.scratch/cbmps_i32.py` is the probe). + +### 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. + + + + + diff --git a/docs/locale/README_ko.md b/docs/locale/README_ko.md index a69c33f..ee2c158 100644 --- a/docs/locale/README_ko.md +++ b/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)) | diff --git a/docs/platform/STATUS.md b/docs/platform/STATUS.md index d613026..d3b5713 100644 --- a/docs/platform/STATUS.md +++ b/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)) | diff --git a/tests/_support/test_shim.py b/tests/_support/test_shim.py index 941012f..69cb99a 100644 --- a/tests/_support/test_shim.py +++ b/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 " diff --git a/tests/devices/mps/test_metalplace.py b/tests/devices/mps/test_metalplace.py index f4fe5c1..d9f0654 100644 --- a/tests/devices/mps/test_metalplace.py +++ b/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 " diff --git a/tests/devices/mps/test_mpsindex.py b/tests/devices/mps/test_mpsindex.py index e3af659..10bf72c 100644 --- a/tests/devices/mps/test_mpsindex.py +++ b/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"} diff --git a/tests/devices/mps/test_mpsindexput.py b/tests/devices/mps/test_mpsindexput.py new file mode 100644 index 0000000..a9a543d --- /dev/null +++ b/tests/devices/mps/test_mpsindexput.py @@ -0,0 +1,389 @@ +"""`aten.index_put_.default` on mps for integer index tensors. + +`index_put_` was refused on mps through `device.rs::MPS_HOST_READBACK_OPS`: +its kernel (`aten.rs::index_put_on_host`) reads `self`, the indices and the +values to the host and walks them there, and candle's Metal scatter and +index_add kernels do not bounds-check, so calling them directly would write +outside the buffer where upstream raises. + +What computes now (`aten.rs::index_put_on_device`): int64 index tensors, +negative indices wrapped on the device, broadcasting among index tensors, +adjacent indices spliced in place and separated ones fronted (the +`index.Tensor` placement rules), `values` broadcast to the indexing result, +a zero-dim CPU `values` (the scalar `x[idx] = 5` lifts to), `accumulate=False` +and `accumulate=True`. Dtypes float32/float16/bfloat16/int64/bool. + +How it computes. The indexed axes are moved to the front and `self` is viewed +as `(block, span)`, exactly as `index_tensor_on_device` does; the flat index +is one `i64` per broadcast element. `accumulate=True` is one candle +`index_add` on that view. `accumulate=False` finds, per destination row, the +last source row that names it (a candle `scatter` of the source positions), +then `index_select`s those rows and `where`s them over `self`. + +**The readback budget, asserted exactly.** The bounds check is #89's +(`index_bounds_on_device`): each index tensor's min and max reduced on the +device and read back as one short vector. **1 download per call** with a +non-empty write, **0** when the broadcast is empty, **2** on an out-of-range +call (the second names the first offender before raising). Nothing is +written before that check passes: the error rows assert `self` unchanged. + +Oracles (AGENTS.md §16): every agreement row is compared with upstream torch +on **mps** in a separate subprocess; upstream cpu is recorded and must agree +with it (a row where they disagree is reported, not graded). The error rows +use upstream **cpu**'s exception type and text, because upstream mps raises +asynchronously: measured on torch 2.13.0, `x[:, tensor([9])] = 1.0` on mps +returns without raising and the `AcceleratorError` surfaces later (at exit), +so it is not an oracle for when or whether the call raises. + +**Duplicate indices with `accumulate=False` have no fixed winner upstream.** +Measured on 2.13.0: `zeros(5)[[1, 1, 1, 3]] = [1, 2, 3, 4]` gives `3` at +position 1 on cpu (last wins) and `1` on mps (first wins); with 1000 +duplicates per position mps gave `992` and `1984` (neither first nor last). +So the duplicate row asserts only that each written element is one of the +values aimed at it. The sum (`accumulate=True`) is defined, and is graded. + +What still refuses, by name, with zero downloads: a boolean or uint8 mask +index that is not the `masked_fill_` lowering `__setitem__` already makes +(its write count is data-dependent). An `int32` index tensor computes: it is +cast to int64 on the device (PR #100's i32 cast kernels). + +Nullifications this file is meant to catch: + +* the bounds check removed (`index_bounds_on_device` call deleted from + `index_put_on_device`) -> every row of + test_out_of_range_indices_raise_and_leave_self_unchanged (no raise, a + write outside the buffer, or 0 downloads instead of 2) +* the negative wrap removed -> the negative-index rows +* the router sending an mps receiver to `index_put_on_host` -> every + agreement row's download count (the host kernel reads everything) +* `index_put_` put back on `MPS_HOST_READBACK_OPS` -> every agreement row + (refused) and test_index_put_left_the_readback_list +* the separated/adjacent placement swapped -> the two 3-d rows +""" + +import json +import math +import os +import subprocess +import sys + +os.environ.setdefault("TORCH_USE_RTLD_GLOBAL", "1") + +import test_shim # noqa: F401 (puts the shim's _C first, as the other mps suites do) +from test_shim import _C +import _skip + + +_REPO_ROOT = str(next(p for p in __import__("pathlib").Path(__file__).resolve().parents + if (p / "AGENTS.md").is_file())) +_VENDOR_DIR = os.path.join(_REPO_ROOT, "torchnative", "python") + + +def _mps_or_skip(what): + if not os.path.isfile(os.path.join(_VENDOR_DIR, "torch", "_C.abi3.so")): + _skip.skip(" (skipped %s: no vendored shim at %s -- run scripts/vendor/install_shim.sh)" + % (what, _VENDOR_DIR)) + return False + try: + _C._aten_dispatch("aten.ones.default", [1], device=_C.device("mps")) + except (NotImplementedError, RuntimeError) as e: + _skip.skip(" (skipped %s: no mps device on this machine -- %s)" + % (what, str(e).splitlines()[0])) + return False + return True + + +def _run(payload, upstream): + env = dict(os.environ) + env.pop("TORCH_C_ARTEFACT", None) + if upstream: + env.pop("PYTHONPATH", None) + env.pop("TORCH_USE_RTLD_GLOBAL", None) + else: + env["PYTHONPATH"] = _VENDOR_DIR + env["TORCH_USE_RTLD_GLOBAL"] = "1" + payload = dict(payload, shim=not upstream) + proc = subprocess.run([sys.executable, "-c", _SCRIPT, json.dumps(payload)], + capture_output=True, text=True, timeout=1800, env=env) + lines = proc.stdout.strip().splitlines() + if not lines: + raise AssertionError("child produced no JSON (rc=%d):\n%s" + % (proc.returncode, proc.stderr[-3000:])) + return json.loads(lines[-1]) + + +# One block per child (a candle Metal panic poisons the encoder's mutex for +# the rest of a process, test_mpsgen.py). `X` is `_GRID` in the row's dtype on +# the device, built before the bracket; the block leaves its answer in `R`. +# The counters bracket only the block and are recorded whether it answers or +# raises; on a raise, `X` after the block is recorded too. +_SCRIPT = r""" +import json, sys +A = json.loads(sys.argv[1]) +import torch +IS_SHIM = hasattr(torch._C, "_aten_implemented") +assert IS_SHIM == A["shim"], ("wrong torch in this child", IS_SHIM) +D = A["device"] +DT = getattr(torch, A["dtype"]) +def T(rows, dtype=None): + return torch.tensor(rows, dtype=dtype or DT).to(D) +L = torch.int64 +X = T(A["data"]) +ns = dict(torch=torch, X=X, T=T, DT=DT, D=D, L=L) +before = torch._C._metal_counters() if IS_SHIM and D == "mps" else None +try: + exec(A["code"], ns) + r = ns["R"] + out = {"shape": list(r.shape), "dtype": str(r.dtype), "device": str(r.device)} +except BaseException as exc: + out = {"error": "%s: %s" % (type(exc).__name__, str(exc)[:600])} +after = torch._C._metal_counters() if before is not None else None +if "error" not in out: + out["values"] = ns["R"].cpu().double().flatten().tolist() +else: + out["self_after"] = X.cpu().double().flatten().tolist() +if before is not None: + out["counters"] = {k: after[k] - before[k] for k in before if k != "built"} +print(json.dumps(out)) +""" + +_GRID = [[1.0, -2.0, 3.0, 0.5], [4.0, 0.0, -1.0, 2.0], [-3.0, 5.0, 1.5, -0.5]] +_FLAT = [v for row in _GRID for v in row] +_FLOATS = ["float32", "float16", "bfloat16"] +_ALL = _FLOATS + ["int64", "bool"] + +# (name, code, dtypes, downloads). `downloads` is the designed count: 1 for a +# non-empty write (the min/max vector of every index tensor, one read), 0 for +# an empty one. +_ROWS = [ + ("x[idx] = scalar", "X[T([0, 2], L)] = 4\nR = X", _ALL, 1), + # int32 indices cast on the device since PR #100 (issue #32). + ("x[int32 idx] = scalar", "X[T([0, 2], torch.int32)] = 4\nR = X", _ALL, 1), + ("negative indices wrap, values broadcast over rows", + "X[T([-1, 0], L)] = T([1, 2, 3, 4])\nR = X", _ALL, 1), + ("x[:, idx] = v", "X[:, T([3, -4], L)] = T([[1, 2], [3, 4], [5, 6]])\nR = X", _ALL, 1), + ("x[idx0, idx1] = v", "X[T([0, 2, -1], L), T([1, -1, 0], L)] = T([9, 8, 7])\nR = X", + _ALL, 1), + ("index tensors broadcast", "X[T([[0], [2]], L), T([1, 3], L)] = T([[1, 2], [3, 4]])\nR = X", + _FLOATS + ["int64"], 1), + ("multi-dim index", "X[T([[0, 2]], L)] = 6\nR = X", ["float32", "int64"], 1), + ("adjacent indices spliced at axis 1", + "Y = X.reshape(3, 2, 2)\nY[:, T([1, 0], L), T([0, 1], L)] = T([5, 6])\nR = Y", + _FLOATS + ["int64"], 1), + ("separated indices go to the front", + "Y = X.reshape(3, 2, 2)\nY[T([0, 2], L), :, T([1, -2], L)] = T([[1, 2], [3, 4]])\nR = Y", + _FLOATS + ["int64"], 1), + ("0-dim cpu value through the aten op", + "X.index_put_((T([1], L),), torch.tensor(5, dtype=DT))\nR = X", _ALL, 1), + ("accumulate=True sums duplicates", + "X.index_put_((T([0, 0, 2, 0], L),), T([1, 2, 3, 4]).reshape(4, 1), accumulate=True)\nR = X", + _ALL, 1), + ("accumulate=True with two index tensors", + "X.index_put_((T([1, 1, -1], L), T([2, 2, 0], L)), T([5, 6, 7]), accumulate=True)\nR = X", + _FLOATS + ["int64"], 1), + ("accumulate=True runs at the receiver's precision", + "R = torch.zeros(1, dtype=DT).to(D)\n" + "R.index_put_((torch.zeros(64, dtype=L).to(D),), torch.full((64,), 0.01, dtype=DT).to(D), " + "accumulate=True)", _FLOATS, 1), + ("write through a view taken before the call", + "V = X[1]\nV[T([0, 3], L)] = 6\nR = X", ["float32", "int64"], 1), + ("int64 rows not named stay put", + "X[T([1], L)] = T([[10, 11, 12, 13]])\nR = X", ["int64", "float32"], 1), + ("empty index writes nothing", "X[T([], L)] = 1\nR = X", ["float32", "int64"], 0), +] + +# (name, code). Shim mps against upstream cpu: same exception type and text, +# exactly 2 downloads, and `X` unchanged. +_ERROR_ROWS = [ + ("index 3 past the end", "X[T([0, 3], L)] = 1"), + ("index -4 before the start", "X[T([0, -4], L)] = 1"), + ("dimension names the index position", "X[:, T([4], L)] = 1"), + ("second index tensor out of range", "X[T([0, 1], L), T([9, 1], L)] = 1"), + ("accumulate=True is checked too", + "X.index_put_((T([0, 5], L),), T(1.0), accumulate=True)"), +] + + +def _same(a, b): + if math.isnan(a) or math.isnan(b): + return math.isnan(a) and math.isnan(b) + return a == b + + +def _counted(label, got, downloads, problems): + c = got["counters"] + if c["host_downloads"] != downloads: + problems.append("%s: %d host downloads, designed %d (%s)" + % (label, c["host_downloads"], downloads, c)) + if downloads and c["compute_encoders"] <= 0: + problems.append("%s: no compute encoder ran (%s)" % (label, c)) + + +def test_index_put_computes_on_mps_and_agrees(): + """Every row: the shim's mps result equals upstream mps exactly (which + equals upstream cpu), stays on mps, and downloads exactly the designed + count with at least one compute encoder.""" + if not _mps_or_skip("index_put_ on mps"): + return + problems = [] + for name, code, dtypes, downloads in _ROWS: + for dtype in dtypes: + payload = dict(data=_GRID, code=code, dtype=dtype) + oracle = _run(dict(payload, device="mps"), upstream=True) + host = _run(dict(payload, device="cpu"), upstream=True) + label = "%s [%s]" % (name, dtype) + assert "error" not in oracle, (label, "upstream mps", oracle) + assert "error" not in host, (label, "upstream cpu", host) + if not all(_same(a, b) for a, b in zip(oracle["values"], host["values"])): + problems.append("%s: upstream mps %s and cpu %s disagree; not graded" + % (label, oracle["values"], host["values"])) + continue + got = _run(dict(payload, device="mps"), upstream=False) + if "error" in got: + problems.append("%s: shim raised %s" % (label, got["error"])) + continue + if not got["device"].startswith("mps"): + problems.append("%s: result on %s" % (label, got["device"])) + if (got["shape"], got["dtype"]) != (oracle["shape"], oracle["dtype"]): + problems.append("%s: %s %s, upstream %s %s" % ( + label, got["shape"], got["dtype"], oracle["shape"], oracle["dtype"])) + elif not all(_same(a, b) for a, b in zip(got["values"], oracle["values"])): + problems.append("%s: %s, upstream mps %s" % (label, got["values"], oracle["values"])) + _counted(label, got, downloads, problems) + assert not problems, "\n".join(problems) + + +def test_duplicate_indices_without_accumulate_write_one_of_their_values(): + """`x[[1, 1, 1]] = rows` without accumulate: upstream fixes no winner + (cpu takes the last, mps the first, measured), so each element of row 1 + must be one of the three values aimed at it and every other row must be + untouched. One download, as for any non-empty write.""" + if not _mps_or_skip("duplicate index_put_ on mps"): + return + code = ("X[T([1, 1, 1], L)] = T([[10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33]])\n" + "R = X") + problems = [] + for dtype in ("float32", "int64"): + label = "duplicates [%s]" % dtype + got = _run(dict(data=_GRID, code=code, dtype=dtype, device="mps"), upstream=False) + if "error" in got: + problems.append("%s: shim raised %s" % (label, got["error"])) + continue + values = got["values"] + for k in range(4): + if values[4 + k] not in (10 + k, 20 + k, 30 + k): + problems.append("%s: column %d holds %r, none of the values aimed at it" + % (label, k, values[4 + k])) + # `_GRID` in this dtype: int64 truncates toward zero (0.5 -> 0, 1.5 -> 1, + # -0.5 -> 0), so the untouched rows are compared with that, not `_FLAT`. + start = _FLAT if dtype == "float32" else [float(int(v)) for v in _FLAT] + if values[:4] != start[:4] or values[8:] != start[8:]: + problems.append("%s: rows 0 and 2 changed: %s" % (label, values)) + _counted(label, got, 1, problems) + assert not problems, "\n".join(problems) + + +def test_out_of_range_indices_raise_and_leave_self_unchanged(): + """Out-of-range positive and negative indices raise upstream cpu's type + and exact text on mps, after exactly two downloads, and nothing was + written: candle's unchecked Metal scatter would otherwise write outside + the buffer.""" + if not _mps_or_skip("index_put_ bounds on mps"): + return + problems = [] + for name, code in _ERROR_ROWS: + payload = dict(data=_GRID, code=code + "\nR = X", dtype="float32") + want = _run(dict(payload, device="cpu"), upstream=True) + got = _run(dict(payload, device="mps"), upstream=False) + assert "error" in want, (name, want) + if "error" not in got: + problems.append("%s: answered on mps (%s), upstream cpu raised %s" + % (name, got.get("values"), want["error"])) + continue + if got["error"] != want["error"]: + problems.append("%s: shim %r, upstream %r" % (name, got["error"], want["error"])) + if got["self_after"] != _FLAT: + problems.append("%s: self changed before the raise: %s" % (name, got["self_after"])) + c = got["counters"] + if c["host_downloads"] != 2: + problems.append("%s: %d host downloads on the error path, designed 2 (%s)" + % (name, c["host_downloads"], c)) + assert not problems, "\n".join(problems) + + +_REFUSAL_ROWS = [ + ("bool mask not the receiver's shape", "X[T([True, False, True], torch.bool)] = 1\nR = X", + "mask index is not implemented for the mps device"), +] + + +def test_masks_refuse_by_name(): + """The paths not implemented on the device refuse naming + `aten.index_put_.default` and the reason, never candle's words, and read + nothing back first.""" + if not _mps_or_skip("index_put_ refusals on mps"): + return + problems = [] + for name, code, says in _REFUSAL_ROWS: + got = _run(dict(data=_GRID, code=code, dtype="float32", device="mps"), upstream=False) + if "error" not in got: + problems.append("%s: answered on mps; if a device kernel landed, move this row " + "to an agreement row with counters" % name) + continue + err = got["error"] + if not err.startswith("NotImplementedError: aten.index_put_.default:"): + problems.append("%s: not a refusal naming index_put_: %s" % (name, err[:300])) + elif says not in err or "candle:" in err: + problems.append("%s: refusal does not say %r: %s" % (name, says, err[:300])) + if got["counters"]["host_downloads"] != 0: + problems.append("%s: read back before refusing (%s)" % (name, got["counters"])) + assert not problems, "\n".join(problems) + + +def test_index_put_left_the_readback_list(): + """The artefact's table: `index_put_` is off it, nothing moved into the + excused column instead, and its neighbours that still read back are on + it.""" + refused = set(_C._shim_mps_host_readback_ops()) + assert "aten.index_put_.default" not in refused + for op in ("aten.index_add_.default", "aten.index_add.default", "aten.scatter.src", + "aten.masked_scatter.default"): + assert op in refused, op + assert set(_C._shim_mps_readback_but_allowed()) == { + "aten._local_scalar_dense.default", "aten.uniform_.default"} + + +def test_only_the_router_calls_the_host_half(): + """`index_put_on_host` is a barrier in test_shim.py's readback + derivation; that is honest only while its one caller is the router that + tests for `Device::Cpu`.""" + mods = test_shim._all_rs_modules() + if mods is None: + _skip.skip(" (skipped router claim: torchnative/rust/torch_c/src is not beside " + "this file -- installed rather than in-tree)") + return + callers = set() + for module, (bodies, _) in mods.items(): + for fn, body in bodies.items(): + if module == "aten" and "index_put_on_host" in test_shim._CALL_IN_MODULE.findall(body): + callers.add(fn) + for other, name in test_shim._CALL_QUALIFIED.findall(body): + if other == "aten" and name == "index_put_on_host": + callers.add("%s::%s" % (module, fn)) + bodies = mods["aten"][0] + assert "index_put_on_host" in bodies and "index_put_on_device" in bodies + assert callers - {"index_put_on_host"} == {"index_put_inplace"}, callers + assert "Device::Cpu" in bodies["index_put_inplace"] + assert "index_bounds_on_device" in bodies["index_put_on_device"] + assert "index_put_on_host" in test_shim._MPS_READBACK_EXEMPT + + +def _main(): + items = [(name, fn) for name, fn in sorted(globals().items()) + if name.startswith("test_") and callable(fn)] + failures = _skip.run_tests(items, suite="test_mpsindexput") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/tests/devices/mps/test_mpssetitem.py b/tests/devices/mps/test_mpssetitem.py index dc6fc24..79fd6d9 100644 --- a/tests/devices/mps/test_mpssetitem.py +++ b/tests/devices/mps/test_mpssetitem.py @@ -14,13 +14,13 @@ The zero-dim CPU value is now accepted there, as `fill_.Tensor`'s is (#32), so a non-zero-dim CPU value still gets upstream's own `Expected all tensors to be on the same device` message. `index_put_` - itself reads back to the host and is refused on mps by the readback gate, + was refused on mps by the readback gate when this file was written, so a **bool mask the shape of the receiver with a scalar or zero-dim value** is lowered by `__setitem__` to `masked_fill_.Scalar` (Python or CPU scalar) or `where.self` (zero-dim value on the receiver's device), - which compute on Metal. An **integer index tensor** with a scalar still - reaches `index_put_` and refuses by name: an on-device lowering would have - to bounds-check the index on the device, which this build does not do. + which compute on Metal. An **integer index tensor** with a scalar reaches + `index_put_`, which now computes on the device with a device bounds check + (`index_put_on_device`, test_mpsindexput.py): one validation download. Grades (AGENTS.md §16): every computed row *agrees* with upstream torch on mps, run in a separate subprocess that asserts it imported upstream; upstream @@ -35,10 +35,10 @@ * `copy_inplace`'s `to_device` of the source removed -> every `w[0] = ...` and `copy_ from cpu` row (candle's slice-set error) * the `index_put_` zero-dim CPU exception in `visit_for_device` removed -> - the int-index refusal row (raises the device-mismatch RuntimeError instead - of the named NotImplementedError) + the int-index row (raises the device-mismatch RuntimeError instead of + computing) * `__setitem__`'s mask lowering removed -> every `w[w > 1] = ...` row - (refused by the readback gate) + (refused by `index_put_on_device`'s mask refusal) * the narrowing in the mask lowering removed (`9.5` passed through to an int64 receiver) -> the int64 `= 9.5` mask row's value, if `masked_fill_` does not truncate itself @@ -206,25 +206,38 @@ def test_scalar_setitem_computes_on_mps(): _agree(_COMPUTED_ROWS, _ALL) -def test_int_index_with_scalar_refuses_by_name(): - """`w[int64_index_on_mps] = scalar`: upstream computes; this build has no - bounds-checked device scatter, so `index_put_` refuses by name through - the readback gate rather than raising the device-mismatch message (which - is what the zero-dim CPU scalar used to trigger) or candle's words.""" - if not _mps_or_skip("int-index refusal on mps"): +def test_int_index_with_scalar_computes_on_mps(): + """`w[int64_index_on_mps] = scalar`: this used to refuse by name through + the readback gate (no bounds-checked device scatter). `index_put_` now + has one (`aten.rs::index_put_on_device`, test_mpsindexput.py), so the + write equals upstream mps (which equals upstream cpu), stays on mps, runs + a compute encoder, and downloads exactly one vector: the on-device min + and max of the index, which is the bounds check and never the answer.""" + if not _mps_or_skip("int-index setitem on mps"): return problems = [] for dtype in _ALL: label = "w[I] = 4 [%s]" % dtype payload = dict(data=_VEC, code="W[I] = 4", dtype=dtype) oracle = _run(dict(payload, device="mps"), upstream=True) - assert "error" not in oracle, (label, oracle) + host = _run(dict(payload, device="cpu"), upstream=True) + assert "error" not in oracle and "error" not in host, (label, oracle, host) + assert oracle["values"] == host["values"], (label, oracle, host) got = _run(dict(payload, device="mps"), upstream=False) - err = got.get("error", "") - if not err.startswith("NotImplementedError: aten.index_put_.default"): - problems.append("%s: not index_put_'s named refusal: %r" % (label, got)) - elif "candle:" in err or "not implemented for the mps device" not in err: - problems.append("%s: refusal gives no reason: %s" % (label, err[:300])) + if "error" in got: + problems.append("%s: shim raised %s" % (label, got["error"])) + continue + if not got["device"].startswith("mps"): + problems.append("%s: result on %s" % (label, got["device"])) + if not all(_same(a, b) for a, b in zip(got["values"], oracle["values"])): + problems.append("%s: result is %s, upstream mps %s" + % (label, got["values"], oracle["values"])) + c = got["counters"] + if c["compute_encoders"] <= 0: + problems.append("%s: no compute encoder ran (%s)" % (label, c)) + if c["host_downloads"] != 1: + problems.append("%s: %d host downloads, designed 1 (the bounds vector) (%s)" + % (label, c["host_downloads"], c)) assert not problems, "\n".join(problems) diff --git a/torchnative/rust/torch_c/src/aten.rs b/torchnative/rust/torch_c/src/aten.rs index 54a5d65..6d130c0 100644 --- a/torchnative/rust/torch_c/src/aten.rs +++ b/torchnative/rust/torch_c/src/aten.rs @@ -19050,40 +19050,7 @@ fn index_tensor_on_device( } else { index_bounds_on_device(OP, &picks, &dims, &broadcast)?; - // Strides of the gathered sub-block, in indexed-axis order. - let mut sub = vec![1usize; picks.len()]; - for j in (0..picks.len().saturating_sub(1)).rev() { - sub[j] = sub[j + 1] * dims[indexed[j + 1]]; - } - let mut offsets: Option = None; - for (j, (at, index, own)) in picks.iter().enumerate() { - let extent = Tensor::new(dims[*at] as i64, &device).map_err(|e| candle_err(OP, e))?; - let negative = index - .lt(&index.zeros_like().map_err(|e| candle_err(OP, e))?) - .map_err(|e| candle_err(OP, e))?; - let lifted = index.broadcast_add(&extent).map_err(|e| candle_err(OP, e))?; - let wrapped = negative - .where_cond(&lifted, index) - .map_err(|e| candle_err(OP, e))?; - let mut aligned = vec![1usize; broadcast_rank - own.len()]; - aligned.extend(own.iter().copied()); - let stride = Tensor::new(sub[j] as i64, &device).map_err(|e| candle_err(OP, e))?; - let term = wrapped - .reshape(aligned) - .and_then(|t| t.broadcast_as(broadcast.clone())) - .and_then(|t| t.broadcast_mul(&stride)) - .map_err(|e| candle_err(OP, e))?; - offsets = Some(match offsets { - None => term, - Some(acc) => acc.broadcast_add(&term).map_err(|e| candle_err(OP, e))?, - }); - } - let flat = offsets - .expect("picks is non-empty") - .broadcast_as(broadcast.clone()) - .and_then(|t| t.contiguous()) - .and_then(|t| t.reshape(total)) - .map_err(|e| candle_err(OP, e))?; + let flat = index_flat_offsets(OP, &picks, &dims, &broadcast, &device)?; let mut order = indexed.clone(); order.extend(untouched.iter().copied()); @@ -19124,6 +19091,62 @@ fn index_tensor_on_device( finish(py, out, tag) } +/// The flat `i64` offset of every broadcast element into the `(block, span)` +/// view of `self` whose block is the indexed axes in `picks` order, all on +/// the device: each index wrapped with `where(idx < 0, idx + size, idx)`, +/// reshaped right-aligned, broadcast, scaled by its sub-block stride and +/// summed. Shared by `index_tensor_on_device` and `index_put_on_device` so +/// the read and the write address the same element for the same index. +/// The caller bounds-checks first; this wraps but does not check. +fn index_flat_offsets( + op: &str, + picks: &[(usize, Tensor, Vec)], + dims: &[usize], + broadcast: &[usize], + device: &Device, +) -> PyResult { + let broadcast_rank = broadcast.len(); + let total: usize = broadcast.iter().product(); + // Strides of the gathered sub-block, in indexed-axis order. + let mut sub = vec![1usize; picks.len()]; + for j in (0..picks.len().saturating_sub(1)).rev() { + sub[j] = sub[j + 1] * dims[picks[j + 1].0]; + } + let mut offsets: Option = None; + for (j, (at, index, own)) in picks.iter().enumerate() { + let extent = Tensor::new(dims[*at] as i64, device).map_err(|e| candle_err(op, e))?; + let negative = index + .lt(&index.zeros_like().map_err(|e| candle_err(op, e))?) + .map_err(|e| candle_err(op, e))?; + let lifted = index.broadcast_add(&extent).map_err(|e| candle_err(op, e))?; + let wrapped = negative + .where_cond(&lifted, index) + .map_err(|e| candle_err(op, e))?; + let mut aligned = vec![1usize; broadcast_rank - own.len()]; + aligned.extend(own.iter().copied()); + let stride = Tensor::new(sub[j] as i64, device).map_err(|e| candle_err(op, e))?; + let term = wrapped + .reshape(aligned) + .and_then(|t| t.broadcast_as(broadcast.to_vec())) + .and_then(|t| t.broadcast_mul(&stride)) + .map_err(|e| candle_err(op, e))?; + offsets = Some(match offsets { + None => term, + Some(acc) => acc.broadcast_add(&term).map_err(|e| candle_err(op, e))?, + }); + } + offsets + .ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "{op}: index_flat_offsets called with no index tensor" + )) + })? + .broadcast_as(broadcast.to_vec()) + .and_then(|t| t.contiguous()) + .and_then(|t| t.reshape(total)) + .map_err(|e| candle_err(op, e)) +} + /// The device bounds check for `index_tensor_on_device`; see there. fn index_bounds_on_device( op: &str, @@ -22860,6 +22883,318 @@ fn masked_fill_inplace( /// consumes. docs/kernels/VIEWS.md §6.1 has the argument that this shape was never /// the obstacle §4 recorded it as being. fn index_put_inplace( + py: Python<'_>, + args: &Bound<'_, PyTuple>, + kwargs: Option<&Bound<'_, PyDict>>, +) -> PyResult> { + const OP: &str = "aten.index_put_.default"; + let receiver = tensor_receiver(OP, args, kwargs)?; + let on_cpu = { + let borrowed = receiver.borrow(); + matches!(borrowed.tensor()?.device(), Device::Cpu) + }; + if on_cpu { + index_put_on_host(py, args, kwargs) + } else { + index_put_on_device(py, args, kwargs) + } +} + +/// `aten.index_put_.default` on an accelerator: integer index tensors only, +/// every step a candle device op, bounds checked on the device before any +/// write (issue #99's follow-up; `test_mpsindexput.py`). +/// +/// The placement is `index_tensor_on_device`'s, read backwards: the indexed +/// axes go to the front and `self` is viewed as `(block, span)`, where +/// `block` is the product of the indexed extents and `span` the rest; the +/// flat index (`index_flat_offsets`, the same function the read uses) names +/// one row of that view per broadcast element. `values` is broadcast to +/// upstream's indexing result shape (adjacent indices spliced in place, +/// separated ones fronted), permuted to the fronted layout and viewed as +/// `(total, span)`. +/// +/// * **Bounds:** `index_bounds_on_device`, #89's check: one download of +/// each index tensor's on-device min and max, a second only on the error +/// path to name the offender, then upstream's CPU `IndexError` text. It +/// runs before anything is written, so a raising call leaves `self` as it +/// was. (Upstream's mps raises asynchronously, measured 2.13.0, so its CPU +/// is the oracle for the error.) +/// * **`accumulate=True`:** one `index_add` on the view. candle's kernel +/// adds the rows in index order in the storage dtype, which is the +/// receiver's-precision running sum `index_put_on_host` documents +/// (bfloat16 `64 x 0.01` is `0.65234375` on both upstream devices). +/// `bool` adds in `i64` and is compared with zero afterwards: a logical +/// or, as upstream's `*dst += *src` on a C++ `bool` is. +/// * **`accumulate=False`:** candle's Metal `scatter` cannot carry an +/// `i64`/`u8` value, so the rows are not scattered. The **position** of +/// each source row is: a `float32` `scatter` of `0..total` into a `-1` +/// column gives, per destination row, the last source row aimed at it +/// (the kernel loops in index order). Those rows are `index_select`ed and +/// `where`d over `self`. Duplicates therefore resolve last-wins, which is +/// upstream cpu's answer; upstream mps answers differently (first-wins in +/// a small case, neither in a large one), so no winner is promised. +/// Positions are exact in `float32` below `2^24`, and a larger write +/// refuses by name rather than mis-addressing. +/// +/// Refused by name, with nothing read back: a boolean or uint8 mask (its +/// write count is data-dependent; `__setitem__` lowers the receiver-shaped +/// mask with a scalar to `masked_fill_` instead), an `int32` index on mps +/// (no int32 Metal kernel in this base), an all-`None` index list. +fn index_put_on_device( + _py: Python<'_>, + args: &Bound<'_, PyTuple>, + kwargs: Option<&Bound<'_, PyDict>>, +) -> PyResult> { + use candle_core::DType; + const OP: &str = "aten.index_put_.default"; + let receiver = tensor_receiver(OP, args, kwargs)?; + let raw_indices = required(OP, args, kwargs, 1, "indices")?; + let items: Vec> = raw_indices.extract()?; + let values = tensor_arg(OP, args, kwargs, 2, "values")?; + let accumulate = bool_arg(args, kwargs, 3, "accumulate")?.unwrap_or(false); + let (tag, source) = { + let borrowed = receiver.borrow(); + (borrowed.tag(), borrowed.tensor()?.clone()) + }; + let device = source.device().clone(); + let label = accelerator_label(&device); + let dims = source.dims().to_vec(); + let rank = dims.len(); + if items.len() > rank { + return Err(pyo3::exceptions::PyIndexError::new_err(format!( + "too many indices for tensor of dimension {rank} (got {})", + items.len() + ))); + } + if values.tag() != tag { + return Err(pyo3::exceptions::PyRuntimeError::new_err(format!( + "Index put requires the source and destination dtypes match, got {} for the \ + destination and {} for the source.", + scalar_type_name(tag), + scalar_type_name(values.tag()) + ))); + } + + // (axis, int64 index tensor on `device`, its shape), as in + // `index_tensor_on_device`. + let mut picks: Vec<(usize, Tensor, Vec)> = Vec::new(); + let mut axis = 0usize; + for item in items.iter() { + if item.is_none() { + axis += 1; + continue; + } + let tensor = item.extract::().map_err(|_| { + pyo3::exceptions::PyTypeError::new_err(format!( + "{OP}: indices must be tensors or None, got {}", + item.get_type().name().map(|n| n.to_string()).unwrap_or_default() + )) + })?; + let index = match tensor.tag() { + TorchDType::Int64 => tensor.tensor()?.clone(), + // int32 indices cast on the device (issue #32: the fork has the + // i32 cast kernels since PR #100), as index.Tensor's do. + TorchDType::Int32 => tensor + .tensor()? + .to_dtype(DType::I64) + .map_err(|e| candle_err(OP, e))?, + TorchDType::Bool | TorchDType::UInt8 => { + return Err(refuse_device_mask_index(OP, label)); + } + _ => { + return Err(pyo3::exceptions::PyIndexError::new_err( + "tensors used as indices must be long, int, byte or bool tensors", + )); + } + }; + let index = index + .to_device(&device) + .and_then(|t| t.contiguous()) + .map_err(|e| candle_err(OP, e))?; + let shape = index.dims().to_vec(); + picks.push((axis, index, shape)); + axis += 1; + } + if picks.is_empty() { + return Err(not_implemented(format!( + "{OP}: an all-None index list is not implemented in torch._C shim" + ))); + } + + let shapes: Vec> = picks.iter().map(|(_, _, s)| s.clone()).collect(); + let broadcast = index_broadcast_shape(&shapes)?; + let broadcast_rank = broadcast.len(); + let total: usize = broadcast.iter().product(); + let indexed: Vec = picks.iter().map(|(at, _, _)| *at).collect(); + let untouched: Vec = (0..rank).filter(|a| !indexed.contains(a)).collect(); + let trailing: Vec = untouched.iter().map(|&a| dims[a]).collect(); + let adjacent = indexed.windows(2).all(|pair| pair[1] == pair[0] + 1); + let at = indexed[0]; + let spliced = adjacent && at > 0; + let last = indexed[indexed.len() - 1]; + + // Upstream's "indexing result": adjacent indices spliced in at `at`, + // otherwise the broadcast shape in front of every untouched axis. + let result_shape: Vec = if spliced { + let mut shape = dims[..at].to_vec(); + shape.extend(broadcast.iter().copied()); + shape.extend(dims[last + 1..].iter().copied()); + shape + } else { + let mut shape = broadcast.clone(); + shape.extend(trailing.iter().copied()); + shape + }; + let result_rank = result_shape.len(); + + // `values` broadcasts right-aligned onto that shape, checked before the + // empty-write shortcut as upstream (and `index_put_on_host`) check it. + let value_dims = values.tensor()?.dims().to_vec(); + let fits = value_dims.len() <= result_rank + && value_dims.iter().enumerate().all(|(k, &extent)| { + extent == 1 || extent == result_shape[result_rank - value_dims.len() + k] + }); + if !fits { + return Err(pyo3::exceptions::PyRuntimeError::new_err(format!( + "shape mismatch: value tensor of shape {value_dims:?} cannot be broadcast \ + to indexing result of shape {result_shape:?}" + ))); + } + if total == 0 { + // Upstream iterates zero elements: it writes and validates nothing. + return Ok(receiver.into_any().unbind()); + } + + // Nothing is written before this returns. + index_bounds_on_device(OP, &picks, &dims, &broadcast)?; + + let span: usize = trailing.iter().product(); + if span == 0 { + return Ok(receiver.into_any().unbind()); + } + let block: usize = indexed.iter().map(|&a| dims[a]).product(); + let flat = index_flat_offsets(OP, &picks, &dims, &broadcast, &device)?; + + // `values` on the device in the fronted `(total, span)` layout. A + // zero-dim CPU value (the scalar `x[idx] = 5` lifts to, admitted by the + // dispatcher's device check) is an upload, not a readback. + let mut fronted = values + .tensor()? + .to_device(&device) + .and_then(|t| t.broadcast_as(result_shape.clone())) + .map_err(|e| candle_err(OP, e))?; + if spliced { + let mut permutation: Vec = (at..at + broadcast_rank).collect(); + permutation.extend(0..at); + permutation.extend(at + broadcast_rank..result_rank); + fronted = fronted.permute(permutation).map_err(|e| candle_err(OP, e))?; + } + let values2d = fronted + .contiguous() + .and_then(|t| t.reshape((total, span))) + .map_err(|e| candle_err(OP, e))?; + + let mut order = indexed.clone(); + order.extend(untouched.iter().copied()); + let self2d = source + .permute(order.clone()) + .and_then(|t| t.contiguous()) + .and_then(|t| t.reshape((block, span))) + .map_err(|e| candle_err(OP, e))?; + + let storage = source.dtype(); + let updated = if accumulate { + if tag == TorchDType::Bool { + let sum = self2d + .to_dtype(DType::I64) + .and_then(|base| { + values2d + .to_dtype(DType::I64) + .and_then(|rows| base.index_add(&flat, &rows, 0)) + }) + .map_err(|e| candle_err(OP, e))?; + sum.zeros_like() + .and_then(|zero| sum.ne(&zero)) + .map_err(|e| candle_err(OP, e))? + } else { + self2d + .index_add(&flat, &values2d, 0) + .map_err(|e| candle_err(OP, e))? + } + } else { + if total >= (1usize << 24) { + return Err(not_implemented(format!( + "{OP}: a write of {total} index positions without accumulate is not \ + implemented on the {label} device. The last-wins resolution carries each \ + position in float32, which is exact only below 2^24, and an inexact \ + position would write the wrong row. Split the write, or move the tensor \ + with .cpu()." + ))); + } + // candle-metal-kernels has no `is_i64_u8`: bool/uint8 rows are + // selected as `i64` and narrowed back (0..=255 round-trips exactly). + let widen = storage == DType::U8; + let (base, rows_from) = if widen { + ( + self2d.to_dtype(DType::I64).map_err(|e| candle_err(OP, e))?, + values2d.to_dtype(DType::I64).map_err(|e| candle_err(OP, e))?, + ) + } else { + (self2d, values2d) + }; + let positions = Tensor::arange(0f32, total as f32, &device) + .and_then(|t| t.reshape((total, 1))) + .map_err(|e| candle_err(OP, e))?; + let ids = flat.reshape((total, 1)).map_err(|e| candle_err(OP, e))?; + let winner = Tensor::full(-1f32, (block, 1), &device) + .and_then(|t| t.scatter(&ids, &positions, 0)) + .map_err(|e| candle_err(OP, e))?; + let zero = winner.zeros_like().map_err(|e| candle_err(OP, e))?; + let written = winner.ge(&zero).map_err(|e| candle_err(OP, e))?; + let pick = written + .where_cond(&winner, &zero) + .and_then(|t| t.to_dtype(DType::I64)) + .and_then(|t| t.reshape(block)) + .map_err(|e| candle_err(OP, e))?; + let rows = rows_from + .index_select(&pick, 0) + .map_err(|e| candle_err(OP, e))?; + let merged = written + .broadcast_as((block, span)) + .and_then(|t| t.contiguous()) + .and_then(|mask| mask.where_cond(&rows, &base)) + .map_err(|e| candle_err(OP, e))?; + if widen { + merged.to_dtype(storage).map_err(|e| candle_err(OP, e))? + } else { + merged + } + }; + + let fronted_dims: Vec = order.iter().map(|&a| dims[a]).collect(); + let mut inverse = vec![0usize; rank]; + for (k, &a) in order.iter().enumerate() { + inverse[a] = k; + } + let full = updated + .reshape(fronted_dims) + .and_then(|t| t.permute(inverse)) + .and_then(|t| t.contiguous()) + .map_err(|e| candle_err(OP, e))?; + let replacement = if tag == TorchDType::Bool { + PyTensorBase::boolean(full)? + } else { + PyTensorBase::new(full)? + }; + write_back(OP, &receiver, replacement)?; + Ok(receiver.into_any().unbind()) +} + +/// The CPU half of `aten.index_put_.default`, reached only from the router +/// `index_put_inplace` when `self` is on the CPU. It reads every operand to +/// the host, which is why it is a named barrier in the readback derivation +/// (`test_shim.py::_MPS_READBACK_EXEMPT`). +fn index_put_on_host( _py: Python<'_>, args: &Bound<'_, PyTuple>, kwargs: Option<&Bound<'_, PyDict>>, diff --git a/torchnative/rust/torch_c/src/device.rs b/torchnative/rust/torch_c/src/device.rs index 032df2b..19906b6 100644 --- a/torchnative/rust/torch_c/src/device.rs +++ b/torchnative/rust/torch_c/src/device.rs @@ -723,7 +723,7 @@ fn shim_same_device(left: PyDevice, right: PyDevice) -> bool { /// third category. It would create a state `_metal_counters()` cannot tell /// from a real fallback -- the instrument §7.14 was built on -- and /// `MPS_READBACK_BUT_ALLOWED` is deliberately two names for that reason. -pub const MPS_HOST_READBACK_OPS: [&str; 88] = [ +pub const MPS_HOST_READBACK_OPS: [&str; 87] = [ // **Ten left on 2026-10-03 by being rewritten onto the device** (issue // #29, greedy `generate` on mps): `argmax.default`, `max.default`, // `min.default`, `max.dim`, `min.dim` (through `first_extremum_index`, @@ -785,7 +785,12 @@ pub const MPS_HOST_READBACK_OPS: [&str; 88] = [ "aten.im2col.default", "aten.index_add.default", "aten.index_add_.default", - "aten.index_put_.default", + // `aten.index_put_.default` left on 2026-10-17 for its integer-index + // paths: `aten.rs::index_put_on_device` writes with candle device ops + // (`index_add` for accumulate, a last-wins position `scatter` plus + // `index_select`/`where` otherwise), bounds-checked first by the same + // one-download `index_bounds_on_device` as `index.Tensor`. Mask and mps + // int32 indices refuse inside that kernel, by name. test_mpsindexput.py. "aten.linalg_qr.default", // §7.16: linalg_vector_norm_default -> norm_pow_walk -> read_flat "aten.linalg_vector_norm.default",