From e94744868b767480101cd258309e8ee81d00541d Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:17:42 +0900 Subject: [PATCH 1/2] mps: strided in-place writes, cumsum and bool copy on Metal; pow refuses by name (#32) - write_on_device splits a non-contiguous receiver into contiguous runs and slice_sets each on the device; past 65536 runs it refuses by name. - fill_.Tensor accepts a 0-dim cpu value beside an mps receiver (x[..., :k] = 0). - cumsum, any and copy_ into bool stay in their own/opmath dtype off the cpu instead of widening to f64 (candle has no Metal F32->F64 cast). - pow.Tensor_Scalar with an exponent other than 2 on Metal refuses by name: candle fast-math Metal powf gives 3.0 ** 3 = 26.99999237. - int32 compute and index/gather/scatter/sort/multinomial/view.dtype stay refused by name (no int32 Metal kernels; Metal gather/scatter unchecked). Red before: cumsum f32 raised Metal to_dtype F32 F64; zero_ on [:, 0:2] refused. Nullified: dropping the run split, or the pow refusal, turns the suites red. --- .github/scripts/gate_suites.py | 1 + docs/devices/matrix.md | 14 +- docs/numerics/DTYPEDEV.md | 5 + tests/_support/test_shim.py | 33 ++- tests/devices/mps/test_mpscb.py | 305 +++++++++++++++++++++++++ tests/devices/mps/test_mpsfwd.py | 7 +- tests/devices/mps/test_mpsinplace.py | 107 +++++---- tests/devices/test_dtypedev.py | 6 + torchnative/rust/torch_c/src/aten.rs | 62 ++++- torchnative/rust/torch_c/src/tensor.rs | 97 ++++++-- 10 files changed, 557 insertions(+), 80 deletions(-) create mode 100644 tests/devices/mps/test_mpscb.py diff --git a/.github/scripts/gate_suites.py b/.github/scripts/gate_suites.py index 9daf99a3..c50baf9e 100644 --- a/.github/scripts/gate_suites.py +++ b/.github/scripts/gate_suites.py @@ -120,6 +120,7 @@ "devices/mps/test_mpsattn.py": BOTH, "devices/mps/test_mpsconst.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_mpsinplace.py": BOTH, diff --git a/docs/devices/matrix.md b/docs/devices/matrix.md index 37de091d..0c071f31 100644 --- a/docs/devices/matrix.md +++ b/docs/devices/matrix.md @@ -1454,6 +1454,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 +1603,10 @@ rebinding the wrapper, which is the obvious wrong way to do this. - + + + + diff --git a/docs/numerics/DTYPEDEV.md b/docs/numerics/DTYPEDEV.md index e8c47026..2bd80271 100644 --- a/docs/numerics/DTYPEDEV.md +++ b/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 diff --git a/tests/_support/test_shim.py b/tests/_support/test_shim.py index ec92a3f0..0533dc5a 100644 --- a/tests/_support/test_shim.py +++ b/tests/_support/test_shim.py @@ -25625,19 +25625,32 @@ def test_an_op_mps_cannot_run_refuses_and_names_the_op(): # and candle is never reached -- which would make this test assert the # wrong refusal. `cumsum` is the replacement and it is the same shape: # implemented here, dispatched through to candle, and bounced by Metal's - # missing `F32 -> F64` because the kernel accumulates in double. + # missing `F32 -> F64` because the kernel accumulated in double. + # + # **Issue #32 removed the last two such ops.** A sweep of every op in + # `_aten_implemented()` outside the readback gate found exactly two that + # candle's Metal backend bounced: `cumsum` and `pow.Tensor_Scalar`. + # `cumsum` now accumulates in its opmath dtype on Metal and computes; + # `pow` with an exponent other than 2 refuses by name, because candle's + # fast-math Metal `powf` gives 3.0 ** 3 = 26.99999237. Both halves are + # asserted: the computed one agrees without a host download, and the + # refused one is this shim's NotImplementedError, not a raw candle error. + before = _C._metal_counters() + got = d("aten.cumsum.default", a, 0) + after = _C._metal_counters() + assert got.device.type == "mps", got.device + assert after["host_downloads"] == before["host_downloads"], (before, after) + assert got.cpu().tolist() == [3.0, 4.0, 6.0], got.cpu().tolist() try: - d("aten.cumsum.default", a, 0) - except RuntimeError as e: + d("aten.pow.Tensor_Scalar", a, 3) + except NotImplementedError as e: message = str(e) - assert "aten.cumsum.default" in message, message - assert "Metal" in message, message + assert "aten.pow.Tensor_Scalar" in message, message + assert "candle:" not in message, message else: - raise AssertionError("aten.cumsum.default must not silently run elsewhere") - # And it really is candle refusing, not this crate's readback gate wearing - # a different exception -- without this the substitution above could be - # satisfied by any op that happens to raise. - assert "aten.cumsum.default" not in set(_C._shim_mps_host_readback_ops()) + raise AssertionError("aten.pow.Tensor_Scalar ** 3 on mps must refuse by name") + for op in ("aten.cumsum.default", "aten.pow.Tensor_Scalar"): + assert op not in set(_C._shim_mps_host_readback_ops()), op # An op nothing implements, for the control: same shape of refusal, so the # one above is not merely the generic path in disguise. diff --git a/tests/devices/mps/test_mpscb.py b/tests/devices/mps/test_mpscb.py new file mode 100644 index 00000000..fd8df732 --- /dev/null +++ b/tests/devices/mps/test_mpscb.py @@ -0,0 +1,305 @@ +"""Continuous batching's walls on mps (issue #32). + +Found by #13 (transformers continuous batching on mps). Four groups, and this +file pins each one to either a computed answer or a refusal by name +(AGENTS.md §18): + +1. **candle-symbol errors.** `cumsum` on a float mps tensor, `bool.copy_`, + `any` of a float tensor, and the float comparisons died with candle's own + `Metal contiguous to_dtype F32 F64 not implemented` (or `U8 F64`): the + kernels widened to `f64` first, and Metal has no `f64`. The comparisons + were fixed by #49; `cumsum` now accumulates in `opmath_in` off the CPU, and + the two `!= 0` tests run in the operand's own dtype. All compute on Metal. +2. **int32 on mps** stays refused, by this build's own sentence + (`device.rs::mps_int_dtype_refusal`). candle-metal-kernels in + `torchnative/rust/vendor` instantiates no `int32_t` cast, binary, affine, + fill or index kernel (only `copy2d`), so there is no device kernel to + reach, and widening for the caller would be a host cast under an mps label. +3. **Strided in-place writes** compute on Metal. `tensor.rs::write_on_device` + splits a non-contiguous receiver into its contiguous runs and blits each + with candle's `slice_set`; `fill_.Tensor` accepts a zero-dim CPU value + beside an mps receiver, as upstream does, so `x[..., :k] = 0` works. +4. **Host-readback ops.** `argmax` already computes on Metal (#29) and is + pinned here by counters. `index.Tensor`, `gather`, `scatter.src`, `sort`, + `multinomial` and `view.dtype` have no device kernel in this build that + matches upstream's contract (candle's Metal `gather`/`scatter` do no bounds + check, so an out-of-range index reads or writes outside the buffer where + upstream raises), so they stay refused by the readback gate, by name. + +Grades (AGENTS.md §16): every *computed* row *agrees* with upstream torch on +the CPU, run in a separate subprocess that asserts it imported upstream, and +is *on the GPU* by `_metal_counters()` bracketing only the expression: at +least one compute or blit encoder and zero host downloads. Every *refused* +row raises `NotImplementedError` whose first word is the aten op and which +carries no `candle:` text. + +Nullifications this file is meant to catch: + +* `cumsum_default`'s off-CPU accumulator back to `F64` -> the cumsum rows +* `any_from`'s device branch removed -> the `any` rows +* `copy_inplace`'s bool device branch removed -> the bool `copy_` row +* `write_on_device`'s run split removed (back to the refusal) -> every + strided row; split wrong (one run only) -> their values +* the `fill_.Tensor` CPU-scalar exception in `visit_for_device` removed -> + the `[..., :1] = 0` row (cross-device RuntimeError) +* an int32 op starting to answer, or its refusal losing its name -> + test_int32_on_mps_refuses_by_name +* a readback op leaving `MPS_HOST_READBACK_OPS` without a device kernel -> + test_readback_ops_refuse_by_name_on_mps (and test_metalplace.py) +""" + +import json +import math +import os +import subprocess +import sys + +os.environ.setdefault("TORCH_USE_RTLD_GLOBAL", "1") + +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 statement block per child (a candle Metal panic poisons the encoder's +# mutex for the rest of a process, test_mpsgen.py). `X` is the input on the +# device; the block leaves its answer in `R`. Only the block is bracketed. +_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"] +inf, nan = float("inf"), float("nan") +DT = getattr(torch, A["dtype"]) +def T(rows, dtype=None): + return torch.tensor(rows, dtype=dtype or DT).to(D) +out = {} +try: + X = T(A["data"]) + ns = dict(torch=torch, X=X, T=T, DT=DT, D=D, inf=inf, nan=nan) + before = torch._C._metal_counters() if IS_SHIM and D == "mps" else None + exec(A["code"], ns) + after = torch._C._metal_counters() if before is not None else None + r = ns["R"] + parts = list(r) if isinstance(r, tuple) else [r] + out = {"parts": [{"shape": list(p.shape), "dtype": str(p.dtype), + "device": str(p.device), + "values": p.cpu().double().flatten().tolist()} for p in parts]} + if before is not None: + out["counters"] = {k: after[k] - before[k] for k in before if k != "built"} +except BaseException as exc: + out = {"error": "%s: %s" % (type(exc).__name__, str(exc)[:600])} +print(json.dumps(out)) +""" + +_INF, _NAN = float("inf"), float("nan") +_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]] +_FLOATS = ["float32", "float16", "bfloat16"] + +# (name, data, code, dtypes). Values are exact in every dtype listed, and the +# cumsum partial sums are too, so the comparison is exact (NaN == NaN). +_CANDLE_SYMBOL_ROWS = [ + ("cumsum dim=1", _GRID, "R = torch.cumsum(X, dim=1)", _FLOATS), + ("cumsum dim=0", _GRID, "R = X.cumsum(0)", _FLOATS), + ("cumsum to float32", _GRID, "R = X.cumsum(1, dtype=torch.float32)", ["float16", "bfloat16"]), + ("lt.Tensor", _GRID, "R = X < X.flip(1)", _FLOATS), + ("le.Tensor", _GRID, "R = X <= X.flip(0)", _FLOATS), + ("lt.Scalar", [[1.0, -_INF, 2.5, _NAN]], "R = X < 2.0", _FLOATS), + ("gt.Scalar", [[1.0, _INF, 2.5, _NAN]], "R = X > 1.0", _FLOATS), + ("bool copy_", _GRID, + "B = X > 0\nC = torch.zeros_like(B)\nC.copy_(B)\nR = C", _FLOATS), + ("bool copy_ into a view", _GRID, + "B = X > 0\nC = torch.zeros(3, 6, dtype=torch.bool, device=D)\nC[:, 1:5].copy_(B)\nR = C", + ["float32"]), + ("any of a float tensor, dim=1", [[0.0, -0.0, 0.0], [0.0, _NAN, 0.0], [0.0, 0.0, 2.0]], + "R = X.any(dim=1)", _FLOATS), +] + +_STRIDED_ROWS = [ + ("zero_ on [:, 0:2]", _GRID, "Y = X.clone()\nY[:, 0:2].zero_()\nR = Y", _FLOATS + ["int64"]), + ("bulk input [:2, :3].zero_()", _GRID, "Y = X.clone()\nY[:2, :3].zero_()\nR = Y", ["int64"]), + ("fill_ on [:, 1:3]", _GRID, + "Y = X.clone()\nY[:, 1:3].fill_(torch.finfo(DT).min)\nR = Y", _FLOATS), + ("x[..., :1] = 0", _GRID, "Y = X.clone()\nY[..., :1] = 0\nR = Y", _FLOATS + ["int64"]), + ("copy_ into [:, 1:]", _GRID, "Y = X.clone()\nY[:, 1:].copy_(X[:, :-1])\nR = Y", _FLOATS), + ("4D mask [:, :, :2, :3].fill_(min)", [[0.0]], + "M = torch.zeros(2, 2, 3, 4, dtype=DT, device=D)\n" + "M[:, :, :2, :3].fill_(torch.finfo(DT).min)\nR = M", _FLOATS), + ("add_ on a column", _GRID, "Y = X.clone()\nY[:, 2].add_(10)\nR = Y", _FLOATS + ["int64"]), + ("fill_ on a transposed view", _GRID, "Y = X.clone()\nY.t()[1:3].fill_(7)\nR = Y", ["float32"]), +] + +_ARGMAX_ROWS = [ + ("argmax dim=1 with ties and NaN", [[1.0, 3.0, 3.0], [_NAN, 1.0, _NAN], [-_INF, -_INF, -_INF]], + "R = X.argmax(dim=1)", _FLOATS), + ("argmax whole tensor", _GRID, "R = X.argmax()", _FLOATS), +] + + +def _same(a, b): + if math.isnan(a) or math.isnan(b): + return math.isnan(a) and math.isnan(b) + return a == b + + +def _agree(rows): + problems = [] + for name, data, code, dtypes in rows: + for dtype in dtypes: + payload = dict(data=data, code=code, dtype=dtype) + want = _run(dict(payload, device="cpu"), upstream=True) + got = _run(dict(payload, device="mps"), upstream=False) + label = "%s [%s]" % (name, dtype) + assert "error" not in want, (label, want) + if "error" in got: + problems.append("%s: shim raised %s" % (label, got["error"])) + continue + for i, (g, w) in enumerate(zip(got["parts"], want["parts"])): + if not g["device"].startswith("mps"): + problems.append("%s: result %d on %s" % (label, i, g["device"])) + if (g["shape"], g["dtype"]) != (w["shape"], w["dtype"]): + problems.append("%s: result %d is %s %s, upstream %s %s" + % (label, i, g["shape"], g["dtype"], w["shape"], w["dtype"])) + elif not all(_same(a, b) for a, b in zip(g["values"], w["values"])): + problems.append("%s: result %d is %s, upstream %s" + % (label, i, g["values"], w["values"])) + c = got["counters"] + if c["compute_encoders"] + c["blit_encoders"] <= 0: + problems.append("%s: no compute or blit encoder ran (%s)" % (label, c)) + if c["host_downloads"] != 0: + problems.append("%s: something was read back to the host (%s)" % (label, c)) + assert not problems, "\n".join(problems) + + +def test_candle_symbol_errors_now_compute_on_mps(): + """cumsum, lt/le.Tensor, lt/gt.Scalar, bool copy_ and any on float mps + tensors: they used to die with candle's `Metal contiguous to_dtype ... F64`. + Each now equals upstream exactly, stays on mps, and downloads nothing.""" + if not _mps_or_skip("candle-symbol rows on mps"): + return + _agree(_CANDLE_SYMBOL_ROWS) + + +def test_strided_in_place_writes_compute_on_mps(): + """zero_, fill_, copy_ and add_ into non-contiguous views (column slices, + a 4D mask slice, a transposed view, `x[..., :1] = 0`) on mps: the whole + base equals upstream's, and the bracket shows device encoders and no + download.""" + if not _mps_or_skip("strided writes on mps"): + return + _agree(_STRIDED_ROWS) + + +def test_argmax_computes_on_mps(): + """argmax stays a device computation (first maximal index, first NaN).""" + if not _mps_or_skip("argmax on mps"): + return + _agree(_ARGMAX_ROWS) + + +def _refused(rows, must_say): + problems = [] + for name, data, code, dtype in rows: + got = _run(dict(data=data, code=code, dtype=dtype, device="mps"), upstream=False) + if "error" not in got: + problems.append("%s: answered on mps (%s); if a device kernel landed, " + "move this row to an agreement test with counters" + % (name, got["parts"][0]["device"])) + continue + err = got["error"] + if not err.startswith("NotImplementedError: aten."): + problems.append("%s: not a refusal naming its op: %s" % (name, err[:300])) + elif "candle:" in err: + problems.append("%s: refusal still carries candle's words: %s" % (name, err[:300])) + elif not any(s in err for s in must_say): + problems.append("%s: refusal gives no reason (%s): %s" % (name, must_say, err[:300])) + assert not problems, "\n".join(problems) + + +_INT32_ROWS = [ + ("embedding with int32 ids", [[0, 2], [1, 0]], + "W = torch.arange(12, dtype=torch.float32).reshape(3, 4).to(D)\n" + "R = torch.nn.functional.embedding(X, W)", "int32"), + ("fill_ int32", [[1, 2], [3, 4]], "X.fill_(3)\nR = X", "int32"), + ("copy_ int32 <- int64", [[1, 2], [3, 4]], "X.copy_(X.long())\nR = X", "int32"), + ("_to_copy to int32", [[1, 2], [3, 4]], "R = X.long().to(torch.int32)", "int32"), + ("sub.Tensor int32", [[1, 2], [3, 4]], "R = X - X", "int32"), +] + + +def test_int32_on_mps_refuses_by_name(): + """int32 compute on mps refuses with the build's own sentence (dtype, + device, reason, both ways out), never candle's shader symbol. There is no + int32 Metal kernel in candle-metal-kernels to reach.""" + if not _mps_or_skip("int32 refusals on mps"): + return + _refused(_INT32_ROWS, ("not implemented for int32 tensors on the mps device",)) + + +_READBACK_ROWS = [ + ("index.Tensor", _GRID, "R = X[T([0, 2], torch.int64)]", "float32"), + ("gather", _GRID, "R = X.gather(1, T([[0], [1], [2]], torch.int64))", "float32"), + ("scatter.src", _GRID, + "R = X.scatter(1, T([[0], [1], [2]], torch.int64), X[:, :1].contiguous())", "float32"), + ("sort", _GRID, "R = X.sort(dim=1)", "float32"), + ("multinomial", [[0.2, 0.8]], "R = torch.multinomial(X, 1)", "float32"), + ("view.dtype", _GRID, "R = X.view(torch.int32)", "float32"), +] + + +def test_readback_ops_refuse_by_name_on_mps(): + """The host-readback ops with no device kernel matching upstream's + contract stay refused by the readback gate, by name.""" + if not _mps_or_skip("readback refusals on mps"): + return + _refused(_READBACK_ROWS, ("not implemented for the mps device",)) + + +def _main(): + items = [(name, fn) for name, fn in sorted(globals().items()) + if name.startswith("test_")] + failures = _skip.run_tests(items, suite="test_mpscb") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/tests/devices/mps/test_mpsfwd.py b/tests/devices/mps/test_mpsfwd.py index 8fd3f74d..d8cba3b9 100644 --- a/tests/devices/mps/test_mpsfwd.py +++ b/tests/devices/mps/test_mpsfwd.py @@ -234,8 +234,11 @@ def test_pow_with_a_general_exponent_on_mps_refuses_rather_than_lowering_precisi x_g = d("aten._to_copy.default", _f32([4.0, 9.0], [2]), device=mps) try: d("aten.pow.Tensor_Scalar", x_g, 0.5) - except RuntimeError as e: - assert "F64" in str(e) or "f64" in str(e), str(e) + except NotImplementedError as e: + # Named refusal since issue #32 (it used to be candle's raw + # `to_dtype F32 F64`): it names the op and the missing float64. + assert "aten.pow.Tensor_Scalar" in str(e) and "float64" in str(e), str(e) + assert "candle:" not in str(e), str(e) else: raise AssertionError( "pow with a non-square exponent must not quietly compute in f32 " diff --git a/tests/devices/mps/test_mpsinplace.py b/tests/devices/mps/test_mpsinplace.py index fd03c2f7..8f714b76 100644 --- a/tests/devices/mps/test_mpsinplace.py +++ b/tests/devices/mps/test_mpsinplace.py @@ -34,16 +34,14 @@ device path is candle's own code, through its published API, and the fork's "two inputs and nothing else" contract is untouched. -**What is NOT lifted, stated first so the rest of this file is read -correctly.** A *non-contiguous* receiver -- `x[:, 0:2]`, an `expand`, a -transposed view -- still refuses on a device, because writing a strided run -needs a scatter kernel that candle does not expose. The refusal for that case -is narrower and names the layout, rather than naming the backend. **It must -not fall through to the host path**: doing so would compute the write on the -CPU under an `mps` label, which is what docs/devices/MPS.md §2 exists to -refuse. `test_the_device_write_door_performs_no_host_readback` is the -structural guard on that, and it is the weak half of this file's evidence -- -see below. +**Strided receivers (issue #32).** A *non-contiguous* receiver -- +`x[:, 0:2]`, an `expand`, a transposed view -- used to refuse on a device. +It is now written as the contiguous runs it is made of, one `slice_set` +blit per run, and refuses by name only past 65536 runs. **It must not fall +through to the host path**: doing so would compute the write on the CPU +under an `mps` label, which is what docs/devices/MPS.md §2 exists to refuse. +`test_the_device_write_door_performs_no_host_readback` is the structural +guard on that, and the strided test brackets the Metal counters. **What this file proves, and at which grade.** @@ -84,8 +82,10 @@ -> test_a_view_taken_before_the_write_sees_it_on_mps * the `.copy()` that breaks source/destination storage sharing removed -> test_a_self_overlapping_copy_on_mps_is_not_half_overwritten -* the contiguity check removed, so a strided receiver writes the wrong cells - -> test_a_strided_receiver_on_mps_refuses_and_names_its_layout +* the strided-receiver run split broken (wrong run, or one run only) + -> test_a_strided_receiver_on_mps_writes_every_run_on_the_device +* the run cap removed, so a transpose issues one blit per element + -> test_a_strided_receiver_with_too_many_runs_refuses_by_name * the contiguity check removed *and* the host path allowed to catch it -> test_the_device_write_door_performs_no_host_readback """ @@ -369,45 +369,72 @@ def test_a_self_overlapping_copy_on_mps_is_not_half_overwritten(): assert why is None, "x[0:2].copy_(x[2:4]) on mps: %s" % why -def test_a_strided_receiver_on_mps_refuses_and_names_its_layout(): - """What is **not** lifted, pinned so it cannot be lifted by accident. +_STRIDED_ORACLE = r""" +import json, sys +import torch +assert not hasattr(torch._C, "_aten_implemented") +req = json.loads(sys.argv[1]) +x = torch.tensor(req["values"]).reshape(req["shape"]).to(torch.float32) +x[:, 0:2].add_(100.0) +print(json.dumps([float(v) for v in x.to(torch.float64).flatten().tolist()])) +""" - A non-contiguous receiver (`x[:, 0:2]`) addresses a strided run, and - candle exposes no device scatter for that. It must refuse -- and it must - not refuse by naming the *backend*, which was the sentence - docs/devices/matrix.md §7.2 found sending readers to the wrong fix. - **The dangerous failure here is success, not refusal.** A device receiver - that quietly fell through to the host path would compute the write on the - CPU under an `mps` label, which is the thing docs/devices/MPS.md §2 - exists to refuse; that path is what `_host_written` below detects, by - checking the values are *not* silently correct. +def test_a_strided_receiver_on_mps_writes_every_run_on_the_device(): + """A non-contiguous receiver (`x[:, 0:2]`) on mps is written (issue #32). + + It used to refuse: candle has no device scatter over a layout. It is now + written as its contiguous runs, one `slice_set` blit each. The oracle is + the *whole base* from upstream, so a run written at the wrong offset, a run + skipped, or the third column touched all show. Placement: the Metal + counters must show no download during the write (a host fallback through + `flat_storage` reads the replacement back). """ - mps = _mps_or_skip("strided receiver refusal on mps") + mps = _mps_or_skip("strided receiver write on mps") if mps is None: return - x = _C._tensor_from_flat([float(v) for v in range(12)], [4, 3], _C.float32, mps) + values = [float(v) for v in range(12)] + want = _oracle(_STRIDED_ORACLE, {"values": values, "shape": [4, 3]}) + x = _C._tensor_from_flat(list(values), [4, 3], _C.float32, mps) strided = _C._aten_dispatch("aten.slice.Tensor", x, 1, 0, 2, 1) + before = dict(_C._metal_counters()) + _C._aten_dispatch("aten.add_.Scalar", strided, 100.0) + after = dict(_C._metal_counters()) + assert after["host_downloads"] == before["host_downloads"], ( + "the strided write read bytes back to the host: %r -> %r" + % (before, after)) + # `[4, 2]` with strides `[3, 1]` is four runs of two: four blits. + assert after["blit_encoders"] - before["blit_encoders"] >= 4, ( + "the strided write issued %d blits, expected one per run (4)" + % (after["blit_encoders"] - before["blit_encoders"])) + tol = dt_utils.tolerance_for("float32") + why = _close(_host(x), want, tol.atol, tol.rtol) + assert why is None, "x[:, 0:2].add_(100) on mps: %s" % why + + +def test_a_strided_receiver_with_too_many_runs_refuses_by_name(): + """A transpose has single-element runs; past the cap it refuses by name. + + `[512, 256].t()` is 131072 one-element runs, above `MAX_DEVICE_WRITE_RUNS` + (65536). It must refuse, name the layout and the run count, and must not + fall through to the host path. + """ + mps = _mps_or_skip("strided receiver run cap on mps") + if mps is None: + return + x = _C._aten_dispatch("aten.zeros.default", [512, 256], dtype=_C.float32, + device=mps) + t = _C._aten_dispatch("aten.t.default", x) try: - _C._aten_dispatch("aten.add_.Scalar", strided, 100.0) + _C._aten_dispatch("aten.zero_.default", t) except NotImplementedError as e: msg = str(e).splitlines()[0] else: raise AssertionError( - "add_ on a NON-contiguous mps receiver succeeded. If a device " - "scatter landed, that is good news and this test should become an " - "agreement test -- but check first that the write did not go " - "through the host path in tensor.rs, which would be a CPU " - "computation under an mps label.") - assert "non-contiguous" in msg, ( - "the strided refusal does not name the layout: %r" % msg) - assert "writing through a view is implemented" not in msg, ( - "the strided refusal is still the old unqualified sentence -- the one " - "that told 112 cells of matrix.md the *backend* was the limit when the " - "limit is the layout: %r" % msg) - assert "receiver does write on the device" in msg, ( - "the refusal does not tell the reader what does work, which is the " - "half that sends them to the right fix: %r" % msg) + "zero_ on a transposed [256, 512] mps view succeeded: the run cap " + "is gone and the write issued one device copy per element") + assert "non-contiguous view" in msg and "131072" in msg, msg + assert "contiguous run" in msg, msg def test_the_device_write_door_performs_no_host_readback(): diff --git a/tests/devices/test_dtypedev.py b/tests/devices/test_dtypedev.py index 36061ba1..72f662bd 100644 --- a/tests/devices/test_dtypedev.py +++ b/tests/devices/test_dtypedev.py @@ -847,6 +847,12 @@ def test_every_dtype_this_build_cannot_store_refuses_by_name(): "bfloat16|max", "bool|argmax", "bool|max", + # `*|cumsum` for the float dtypes joined with issue #32: off the CPU a + # float cumsum accumulates in `opmath_in` instead of `f64`, which Metal + # does not have (it used to die on `Metal contiguous to_dtype F32 F64`). + "bfloat16|cumsum", + "float16|cumsum", + "float32|cumsum", "float16|argmax", "float16|max", "float32|argmax", diff --git a/torchnative/rust/torch_c/src/aten.rs b/torchnative/rust/torch_c/src/aten.rs index ea529abd..6b82730a 100644 --- a/torchnative/rust/torch_c/src/aten.rs +++ b/torchnative/rust/torch_c/src/aten.rs @@ -2225,6 +2225,18 @@ fn visit_for_device( if op == "aten.copy_.default" { return Ok(true); } + // `fill_.Tensor`'s value may be a zero-dim CPU tensor whatever the + // receiver's device (issue #32): upstream treats it as a scalar, and + // `x[..., :k] = 0` on mps arrives here as exactly that. `fill_inplace` + // reads the one CPU element as a scalar, which is the `fill_.Scalar` + // path, so nothing on the device is read back. + if op == "aten.fill_.Tensor" && matches!(&*first, Some(Where::Dense(_))) { + if let crate::tensor::Repr::Dense(inner) = borrowed.repr() { + if inner.device().is_cpu() && inner.rank() == 0 { + return Ok(true); + } + } + } let seen = first.as_ref().expect("a disagreement needs a first").label(); Err(pyo3::exceptions::PyRuntimeError::new_err(format!( "Expected all tensors to be on the same device, but found at least two \ @@ -9621,6 +9633,21 @@ fn pow_tensor_scalar( float_narrower(tag)(exponent.as_f64()) }; let storage = PyDtype::new(tag).storage(OP)?; + // Off the CPU this used to surface as a raw candle `to_dtype F32 F64` + // error (issue #32). Computing in f32 there is not an answer either: + // candle's Metal `powf` is compiled with fast math and gives + // `3.0 ** 3 = 26.99999237` where upstream gives 27 (measured on this + // Mac). So it refuses by name, with the reason. + if !source.device().is_cpu() { + return Err(not_implemented(format!( + "{OP}: a floating-point base raised to an exponent other than 2 on {:?} is not implemented in \ + torch._C shim: Metal has no float64 to compute it in as the CPU \ + does, and candle's fast-math Metal powf is not exact \ + (3.0 ** 3 gives 26.99999237, upstream 27). x ** 2 computes on \ + the device. docs/devices/matrix.md §7.7", + source.device() + ))); + } let out = widen_f64(source) .and_then(|t| t.powf(e)) .and_then(|t| t.fast_to(storage)) @@ -13710,13 +13737,23 @@ fn cumsum_default( // op staying on the device it was asked for. let dims = input.tensor()?.dims().to_vec(); let n = dims[dim]; + // **Off the CPU a float accumulates in `opmath_in`** (issue #32): Metal + // has no `f64`, so the widening died with `Metal contiguous to_dtype F32 + // F64 not implemented` on every float `cumsum` on mps. `float32` + // accumulates in `float32` and `float16`/`bfloat16` in `float32`, which + // is upstream mps's own accumulator; the CPU keeps its `f64` one. + let on_cpu = input.tensor()?.device().is_cpu(); let acc = if storage.is_int() { candle_core::DType::I64 - } else { + } else if on_cpu { candle_core::DType::F64 + } else { + opmath_in(storage) }; let source = if acc == candle_core::DType::F64 { widen_f64(input.tensor()?) + } else if acc.is_float() { + input.tensor()?.fast_to(acc) } else { input.tensor()?.to_dtype(acc) } @@ -14750,7 +14787,10 @@ fn any_from(op: &str, source: &Tensor) -> PyResult { // The scalar `0` costs nothing on the device either: candle converts a // scalar operand on the host and uploads it (`cmp!`'s `to_dtype` then // `to_device`). - if source.dtype().is_int() { + // Off the CPU a float is tested in its own dtype (issue #32): Metal has no + // `f64`, and `x != 0` does not depend on the width (NaN is non-zero at + // every width), the same reasoning `to_copy_default`'s bool arm uses. + if source.dtype().is_int() || !source.device().is_cpu() { return source.ne(0f64).map_err(|e| candle_err(op, e)); } widen_f64(source) @@ -17827,7 +17867,9 @@ fn local_scalar_dense( return Ok((value != 0).into_bound_py_any(py)?.unbind()); } if input.tag().is_floating_point() { - let value = widen_f64(&flat) + // `_local_scalar_dense` *is* the readback (`.item()`), so the one + // element is widened after it reaches the host: Metal has no `f64`. + let value = widen_f64_host(&flat) .and_then(|t| t.to_vec1::()) .map_err(|e| candle_err(OP, e))?[0]; return Ok(value.into_bound_py_any(py)?.unbind()); @@ -18481,11 +18523,15 @@ fn copy_inplace( .and_then(|t| t.contiguous()) .map_err(|e| candle_err(OP, e))?; let replacement = if tag == TorchDType::Bool { - PyTensorBase::boolean( - widen_f64(&widened) - .and_then(|t| t.ne(0f64)) - .map_err(|e| candle_err(OP, e))?, - )? + // Off the CPU the source is tested in its own dtype (issue #32): + // `bool.copy_(bool)` on mps died with `Metal contiguous to_dtype U8 + // F64 not implemented`. `!= 0` is width-independent, as in `any_from`. + let truth = if widened.dtype().is_int() || !widened.device().is_cpu() { + widened.ne(0f64) + } else { + widen_f64(&widened).and_then(|t| t.ne(0f64)) + }; + PyTensorBase::boolean(truth.map_err(|e| candle_err(OP, e))?)? } else { let storage = PyDtype::new(tag).storage(OP)?; PyTensorBase::new(widened.fast_to(storage).map_err(|e| candle_err(OP, e))?)? diff --git a/torchnative/rust/torch_c/src/tensor.rs b/torchnative/rust/torch_c/src/tensor.rs index 16708a6a..a002a018 100644 --- a/torchnative/rust/torch_c/src/tensor.rs +++ b/torchnative/rust/torch_c/src/tensor.rs @@ -1226,14 +1226,11 @@ impl PyTensorBase { /// person actually writes to -- a whole tensor, or a `detach()`, `view()`, /// `reshape()`, `unsqueeze()` or leading-axis `slice`/`select` of one. /// -/// A *strided* view (`x[:, 0:2]`, a transpose, an `expand`) needs a scatter -/// over a discontiguous set of positions, and candle exposes no public -/// device-side scatter-by-layout: `copy2d` handles one stride pair, not a rank- -/// n odometer, and `InplaceOp1::metal_fwd` hands back a `MetalStorage` whose -/// element count is private and against which we would have to write a Metal -/// kernel of our own. **Refusing is the answer, and the refusal names the -/// layout** -- naming the *backend*, as the old sentence did, is what sent 112 -/// cells of docs/devices/matrix.md to the wrong diagnosis (§7.2). +/// A *strided* view (`x[:, 0:2]`, a transpose, an `expand`) is not one run, +/// and candle exposes no public device-side scatter-by-layout. It is written +/// as the several contiguous runs it is made of, one `slice_set` each (issue +/// #32), and refused by name only when that would take more than +/// `MAX_DEVICE_WRITE_RUNS` copies. It never falls through to the host path. /// /// **Nothing here changes `torchnative/rust/vendor/candle-core`.** Every call below is on the /// published API of the pinned crate; the fork's contract is "two inputs and @@ -1250,18 +1247,37 @@ fn write_on_device( if dest.elem_count() == 0 { return Ok(()); } - if !dest.is_contiguous() { + // **A strided receiver is split into the contiguous runs it addresses** + // (issue #32). Continuous batching's first write is + // `_bulk_input_tensor[:7, :33].zero_()`, and its attention mask is + // `mask[:, :, :q, :k].fill_(min)`: both are views whose innermost axes + // are contiguous and whose outer axes are cut. Every such layout is a set + // of `runs` unbroken stretches of `run_len` elements, one per index of + // the leading `split` axes, and each stretch is itself a contiguous view + // (`narrow` to size 1 on the leading axes; candle's `is_contiguous` + // skips size-1 axes). So each stretch goes through the same `slice_set` + // blit as a contiguous receiver: every write is on the device, nothing + // is read back, and no kernel of ours is involved. + // + // The cost is one blit per run. A layout whose innermost axis is itself + // strided (a transpose) has `run_len == 1`, one blit per *element*; past + // `MAX_DEVICE_WRITE_RUNS` that is refused by name rather than issued. + let split = contiguous_split(dest.dims(), dest.layout().stride()); + let runs: usize = dest.dims()[..split].iter().product(); + if split > 0 && runs > MAX_DEVICE_WRITE_RUNS { return Err(not_implemented(format!( - "{op}: writing through a non-contiguous view is implemented for the \ - CPU backend only in torch._C shim; this tensor is on {label} with \ - dims {:?} and strides {:?}. A *contiguous* receiver does write on \ - the device -- a whole tensor, or a view(), reshape(), detach(), \ - unsqueeze() or leading-axis slice of one. This layout addresses a \ - broken run, and candle 0.11.0 publishes no device scatter for one \ - (`copy2d` takes a single stride pair). .contiguous() the receiver, \ - or write the base. docs/devices/matrix.md §7.7", + "{op}: writing through this non-contiguous view on {label} would \ + take {runs} separate device copies (dims {:?}, strides {:?}; \ + each contiguous run is {} elements), more than the {} this shim \ + issues for one write. candle 0.11.0 publishes no device scatter \ + over an arbitrary layout, so a strided receiver is written one \ + contiguous run at a time, and this layout's runs are too short. \ + .contiguous() the receiver, or write the base. \ + docs/devices/matrix.md §7.7", dest.dims(), - dest.layout().stride() + dest.layout().stride(), + dest.elem_count() / runs.max(1), + MAX_DEVICE_WRITE_RUNS ))); } @@ -1294,7 +1310,50 @@ fn write_on_device( (dest.clone(), src) }; - dest.slice_set(&src, 0, 0).map_err(|e| candle_err(op, e)) + if split == 0 { + return dest.slice_set(&src, 0, 0).map_err(|e| candle_err(op, e)); + } + // The odometer over the leading `split` axes. `src` is contiguous with + // `dest`'s shape, so the same narrows address the matching run of it. + let outer: Vec = dest.dims()[..split].to_vec(); + let mut index = vec![0usize; split]; + for _ in 0..runs { + let mut d = dest.clone(); + let mut s = src.clone(); + for (axis, &i) in index.iter().enumerate() { + d = d.narrow(axis, i, 1).map_err(|e| candle_err(op, e))?; + s = s.narrow(axis, i, 1).map_err(|e| candle_err(op, e))?; + } + d.slice_set(&s, 0, 0).map_err(|e| candle_err(op, e))?; + for axis in (0..split).rev() { + index[axis] += 1; + if index[axis] < outer[axis] { + break; + } + index[axis] = 0; + } + } + Ok(()) +} + +/// The most device copies one strided write may issue (issue #32). A +/// `[B, H, q, k]` mask slice is `B * H * q` runs; this bound admits every +/// shape continuous batching writes and refuses a transpose of a large +/// matrix, whose runs are single elements. +const MAX_DEVICE_WRITE_RUNS: usize = 1 << 16; + +/// How many leading axes must be iterated so that what remains is one +/// contiguous run: `0` for a contiguous layout. The trailing axes kept are +/// those whose stride is the running element count (or whose extent is 1), +/// which is candle's own `Shape::is_contiguous` rule read from the inside. +fn contiguous_split(dims: &[usize], stride: &[usize]) -> usize { + let mut acc = 1usize; + let mut split = dims.len(); + while split > 0 && (dims[split - 1] <= 1 || stride[split - 1] == acc) { + acc *= dims[split - 1]; + split -= 1; + } + split } /// Whether an in-place op may write into a destination that addresses the From ffda2c54f28e7e91b47741635c44a4fb328681d5 Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Fri, 16 Oct 2026 05:31:53 +0900 Subject: [PATCH 2/2] mps: index.Tensor and gather compute on Metal with a device-side bounds check (#32) Every transformers forward hits index.Tensor, and on mps it was refused by the readback gate: candle's Metal index_select/gather do not bounds-check. - index_tensor/gather route a cpu self to the host kernels and any other device to candle device ops: negative indices wrap on device, index tensors broadcast, adjacent/separated axes placed as upstream. - Bounds: each index tensor's min/max reduce on the device and one small vector is read back to validate; out of range raises upstream cpu's error type and text (a second readback names the first offender). - bool/uint8 values go through i64 (no is_i64_u8/u32 or gather_i64_u8 Metal kernels). Masks and int32 indices on mps refuse by name. - cpu index.Tensor's error now names the index position, as upstream does. - MPS_HOST_READBACK_OPS 89 -> 87; tests and docs updated. Red before: all six test_mpsindex tests. Nullified: skipping the bounds check returned out-of-range values (index 5 answered) and turned two tests red. --- .github/scripts/gate_suites.py | 1 + README.md | 2 +- docs/devices/matrix.md | 67 +++ docs/locale/README_ko.md | 2 +- docs/platform/STATUS.md | 2 +- tests/_support/test_shim.py | 38 +- tests/devices/mps/test_mpsattn.py | 3 +- tests/devices/mps/test_mpscb.py | 12 +- tests/devices/mps/test_mpsgen.py | 40 +- tests/devices/mps/test_mpsindex.py | 352 ++++++++++++++ torchnative/rust/torch_c/src/aten.rs | 633 ++++++++++++++++++++++--- torchnative/rust/torch_c/src/device.rs | 15 +- 12 files changed, 1066 insertions(+), 101 deletions(-) create mode 100644 tests/devices/mps/test_mpsindex.py diff --git a/.github/scripts/gate_suites.py b/.github/scripts/gate_suites.py index c50baf9e..06e62e27 100644 --- a/.github/scripts/gate_suites.py +++ b/.github/scripts/gate_suites.py @@ -123,6 +123,7 @@ "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_viewdtype.py": BOTH, diff --git a/README.md b/README.md index 5c82f7a8..008305ad 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, **89** 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 0c071f31..9ff7a231 100644 --- a/docs/devices/matrix.md +++ b/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 @@ -2432,3 +2441,61 @@ They are listed as unexercised now, with the measured refusal as the reason. + +### 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 89 to 87. + +**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. + + + + + + diff --git a/docs/locale/README_ko.md b/docs/locale/README_ko.md index c617344c..ee2c1588 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 은 입구에서 이름으로 거부됩니다. **89** 개(이 칸은 한때 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 7592e361..c91b1867 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, **89** of them (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 (this cell said 54; it was 89 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; 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 0533dc5a..d7d5f743 100644 --- a/tests/_support/test_shim.py +++ b/tests/_support/test_shim.py @@ -25753,6 +25753,23 @@ def test_mps_is_refused_by_name_where_it_is_not_compiled_in(): "from a constant the test itself built on the CPU. It is the Rust " "half of the float16 double-rounding assertion in " "test_mpsconst.py, and it touches no device (matrix.md 7.18).", + "index_validation_readback": + "reads only a short i64 vector the device index paths built " + "themselves: each index tensor's on-device min and max (and, on the " + "error path only, the indices, to name the offender before raising). " + "It is validation for candle's unchecked Metal index kernels, never " + "an operand of the answer. test_mpsindex.py asserts the download " + "count per call (1 on success, 2 on an out-of-range raise), so it " + "cannot grow into an operand readback unnoticed (issue #32).", + "index_tensor_on_host": + "the CPU half of aten.index.Tensor, reached only from the router " + "`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.", + "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 " + "the router is its only caller and counts an mps call's downloads.", "bf16_from_f64_is_two_roundings_everywhere": "a #[cfg(test)] proof -- issue #28's bfloat16 witnesses, read back " "from constants and a candle `affine` the test itself built on the " @@ -26140,9 +26157,10 @@ def test_an_mps_op_that_would_compute_on_the_cpu_is_refused_and_names_the_op(): """The gate, from the caller's side. `nonzero` is the op docs/devices/VULKAN3.md §3 caught returning a correct value - from the CPU under an `mps` label; `gather` is the second name here because - it is the one still standing between a plain `BertModel(input_ids)` and an - `mps` forward (docs/devices/MPSATTN.md §6). + from the CPU under an `mps` label; `scatter.src` is the second name here. + It replaced `gather`, which left the list in issue #32 by moving onto the + device (test_mpsindex.py pins that); `scatter.src` is still a host + readback. **`_softmax` used to be the second name and is deliberately not any more.** It was rewritten onto the device and left the list, and `test_mpsattn.py` @@ -26156,7 +26174,7 @@ def test_an_mps_op_that_would_compute_on_the_cpu_is_refused_and_names_the_op(): on the CPU is untouched -- the gate is scoped to Metal arguments, so a version of it that fired everywhere would pass the first two. - **Removing the gate puts `nonzero` and `gather` back to answering + **Removing the gate puts `nonzero` and `scatter.src` back to answering correctly off the CPU with `mps` on the label**, which is what this test fails to notice if the arm in `aten_dispatch` is deleted -- verified by deleting it, not by assuming. @@ -26170,9 +26188,9 @@ def test_an_mps_op_that_would_compute_on_the_cpu_is_refused_and_names_the_op(): index = d("aten._to_copy.default", _C._tensor_from_flat([0.0, 1.0, 1.0, 0.0], [2, 2], _C.int64), device=mps) - for op in ("aten.nonzero.default", "aten.gather.default"): + for op in ("aten.nonzero.default", "aten.scatter.src"): try: - d(op, a) if op == "aten.nonzero.default" else d(op, a, 1, index) + d(op, a) if op == "aten.nonzero.default" else d(op, a, 1, index, a) except NotImplementedError as e: message = str(e) assert op in message, message @@ -26201,9 +26219,13 @@ def test_the_mps_refusal_list_is_not_empty_and_covers_the_named_regressions(): assert len(ops) > 40, len(ops) assert len(set(ops)) == len(ops), "duplicate entries" assert ops == sorted(ops), "kept sorted so a diff is readable" - for named in ("aten.nonzero.default", "aten.gather.default", - "aten.where.default", "aten.index.Tensor"): + for named in ("aten.nonzero.default", "aten.scatter.src", + "aten.where.default", "aten.sort.default"): assert named in ops, named + # `gather` and `index.Tensor` left in issue #32 by moving onto the device + # (test_mpsindex.py); named for the reason the softmax pair is below. + for moved in ("aten.gather.default", "aten.index.Tensor"): + assert moved not in ops, moved # `_softmax` and `_safe_softmax` were on this list and are not any more: # they were rewritten onto the device (docs/devices/MPSATTN.md), which is the only # way off it. Named here rather than left as an absence so that putting a diff --git a/tests/devices/mps/test_mpsattn.py b/tests/devices/mps/test_mpsattn.py index 1dc68d5b..243bb9da 100644 --- a/tests/devices/mps/test_mpsattn.py +++ b/tests/devices/mps/test_mpsattn.py @@ -382,7 +382,8 @@ def test_a_float_literal_and_a_scalar_where_reach_the_device_at_all(): max_position_embeddings=32) _BERT_IDS = [[1, 2, 3, 4, 5, 6, 7, 8]] # Supplied rather than left to default: `BertEmbeddings` otherwise builds them -# with `torch.gather`, which is still refused on `mps` (docs/devices/MPSATTN.md §6). +# with `torch.gather`, which was refused on `mps` (docs/devices/MPSATTN.md §6) +# until issue #32 moved it onto the device (test_mpsindex.py). # That op is in the *embeddings*, not the attention block. _BERT_TT = [[0] * 8] diff --git a/tests/devices/mps/test_mpscb.py b/tests/devices/mps/test_mpscb.py index fd8df732..05d29c40 100644 --- a/tests/devices/mps/test_mpscb.py +++ b/tests/devices/mps/test_mpscb.py @@ -20,11 +20,13 @@ with candle's `slice_set`; `fill_.Tensor` accepts a zero-dim CPU value beside an mps receiver, as upstream does, so `x[..., :k] = 0` works. 4. **Host-readback ops.** `argmax` already computes on Metal (#29) and is - pinned here by counters. `index.Tensor`, `gather`, `scatter.src`, `sort`, + pinned here by counters. `index.Tensor` and `gather` compute on Metal for + integer indices, bounds-checked on the device (test_mpsindex.py pins + them, with their one validation download). `scatter.src`, `sort`, `multinomial` and `view.dtype` have no device kernel in this build that - matches upstream's contract (candle's Metal `gather`/`scatter` do no bounds - check, so an out-of-range index reads or writes outside the buffer where - upstream raises), so they stay refused by the readback gate, by name. + matches upstream's contract (candle's Metal `scatter` does no bounds + check, so an out-of-range index writes outside the buffer where upstream + raises), so they stay refused by the readback gate, by name. Grades (AGENTS.md §16): every *computed* row *agrees* with upstream torch on the CPU, run in a separate subprocess that asserts it imported upstream, and @@ -276,8 +278,6 @@ def test_int32_on_mps_refuses_by_name(): _READBACK_ROWS = [ - ("index.Tensor", _GRID, "R = X[T([0, 2], torch.int64)]", "float32"), - ("gather", _GRID, "R = X.gather(1, T([[0], [1], [2]], torch.int64))", "float32"), ("scatter.src", _GRID, "R = X.scatter(1, T([[0], [1], [2]], torch.int64), X[:, :1].contiguous())", "float32"), ("sort", _GRID, "R = X.sort(dim=1)", "float32"), diff --git a/tests/devices/mps/test_mpsgen.py b/tests/devices/mps/test_mpsgen.py index 594ea79c..f4676823 100644 --- a/tests/devices/mps/test_mpsgen.py +++ b/tests/devices/mps/test_mpsgen.py @@ -38,9 +38,10 @@ compute encoders > 0 and **zero** host downloads. A host-computed twin can keep every value right; it cannot launch a compute encoder, and it has to download its operand. -* *refused by name* -- `aten.index.Tensor` (eager attention's - `padding_mask[batch_idx, kv_idx]`) is still a host readback and is pinned - as one; it is the wall left for eager-mode generate. +* *past the old wall* -- `aten.index.Tensor` (eager attention's + `padding_mask[batch_idx, kv_idx]`) was a host readback and the wall for + eager-mode generate; it computes on Metal since issue #32 + (test_mpsindex.py), and the eager test now requires the run to get past it. Nullifications this file is meant to catch: @@ -58,6 +59,7 @@ import json import os +import re import subprocess import sys @@ -198,19 +200,27 @@ def test_greedy_generate_on_mps_agrees_with_upstream_token_for_token(): "the host" % (c["host_download_bytes"], bound), c) -def test_greedy_generate_on_mps_with_eager_attention_stops_at_index_by_name(): - """Eager attention is the one configuration still walled, and the wall is - named: `masking_utils.padding_mask_function` indexes a bool mask with two - index tensors, `aten.index.Tensor`, whose kernel is still a host readback. - When that op moves onto the device this test goes red; replace it with - the sdpa test above run on `attn="eager"`.""" +def test_greedy_generate_on_mps_with_eager_attention_gets_past_index(): + """Eager attention used to stop at `aten.index.Tensor`: + `masking_utils.padding_mask_function` indexes a bool mask with two int64 + index tensors, and the op was a host readback. It now computes on Metal + (issue #32, test_mpsindex.py), so the run must not stop there. If it + completes, its tokens must equal upstream's; if a later op stops it, that + op must refuse by its own name, and the round that lands it turns this + into the sdpa agreement test above run on `attn="eager"`. + + NOT MEASURED when written: the round that moved index.Tensor could not + build the extension, so which of the two branches holds is the gate's + first answer.""" if not _mps_or_skip("eager generate on mps"): return got = _generate(True, "mps", "eager") - assert "error" in got, ( - "eager generate on mps now completes -- aten.index.Tensor moved; " - "turn this into an agreement test") - assert "aten.index.Tensor: not implemented for the mps device" in got["error"], got["error"] + if "error" not in got: + ref = _generate(False, "cpu", "eager") + assert got["tokens"] == ref["tokens"], (got["tokens"], ref["tokens"]) + return + assert "aten.index.Tensor" not in got["error"], got["error"] + assert re.search(r"NotImplementedError: aten\.[A-Za-z0-9_.]+: ", got["error"]), got["error"] # -------------------------------------------------------------------------- @@ -385,14 +395,14 @@ def test_rewritten_ops_agree_on_cpu(): def test_the_ten_generate_ops_are_not_refused_on_mps(): """Off the list by being rewritten -- the derivation scan in test_shim.py is what proves the kernels hold no readback; this checks the artefact's - table, and that `index.Tensor` (the eager wall) is still on it.""" + table, and that `index.Tensor` (the eager wall until issue #32) is off it.""" refused = set(_C._shim_mps_host_readback_ops()) for op in ("aten.argmax.default", "aten.max.default", "aten.min.default", "aten.max.dim", "aten.min.dim", "aten.isin.Tensor_Tensor", "aten.bitwise_and.Tensor", "aten.bitwise_or.Tensor", "aten.bitwise_xor.Tensor", "aten.bitwise_not.default"): assert op not in refused, op - assert "aten.index.Tensor" in refused + assert "aten.index.Tensor" not in refused assert "aten.multinomial.default" in refused, ( "do_sample=True's sampler left the list -- measure sampling on mps") diff --git a/tests/devices/mps/test_mpsindex.py b/tests/devices/mps/test_mpsindex.py new file mode 100644 index 00000000..d47cafcb --- /dev/null +++ b/tests/devices/mps/test_mpsindex.py @@ -0,0 +1,352 @@ +"""`aten.index.Tensor` and `aten.gather.default` on mps (issue #32). + +Both were refused by the host-readback gate: their kernels read the indices +(and, for `gather`, the operand) to the host, and candle's Metal +`index_select`/`gather` kernels do not bounds-check, so simply calling them +would read outside the buffer where upstream raises. Every transformers +forward under continuous batching reaches `index.Tensor`, so this was #32's +wall. + +What computes now (`aten.rs::index_tensor_on_device`, `gather_on_device`): +integer index tensors (`int64` everywhere; `int32` on cuda), negative +indices wrapped on the device for `index.Tensor` (and refused for `gather`, +which upstream does not wrap), broadcasting among index tensors, slices and +separating axes, empty indices; values float32/float16/bfloat16/int64/bool. + +**The readback budget, asserted exactly.** The bounds check reduces each +index tensor to its min and max on the device and reads that one short +vector back (`index_validation_readback`): **1 download per call** with a +non-empty result, **0** when the broadcast is empty, and **2** on an +out-of-range call (the second reads the indices only to name the first +offender, as upstream does, before raising). That readback is validation, +never an operand of the answer, and docs/devices/matrix.md §7.19 records it. + +What still refuses, by name: boolean/uint8 mask indices on a device (the +result's length is data-dependent), and `int32` index tensors on mps (no +int32 Metal kernel, not even the cast). + +Errors: upstream's **CPU** type and text (`IndexError: index 5 is out of +bounds for dimension 0 with size 3`, `RuntimeError` for `gather`). Upstream's +mps raises `torch.AcceleratorError` (a `RuntimeError`), with the same text +for `index` and `gather: index 5 is out of bounds for dimension with size 4` +for gather (measured, torch 2.13.0); the CPU wording is the oracle here. + +Nullifications this file is meant to catch: + +* the device bounds check removed (`index_bounds_on_device` returning Ok, or + the `ends` test in `gather_on_device` deleted) -> the out-of-range rows + (no raise, or candle reads garbage) and their download count +* the negative wrap removed -> the negative-index rows (wrong values or an + out-of-range read) +* the router sending an mps tensor to `index_tensor_on_host`/`gather_on_host` + -> every agreement row's download count (the host kernels read the whole + index, and `gather_on_host` the operand too), and + test_only_the_routers_call_the_host_halves +* a second validation download per index tensor -> the multi-index rows (2 + index tensors still 1 download) +* the mask refusal removed -> test_masks_and_mps_int32_indices_refuse_by_name +""" + +import json +import math +import os +import re +import subprocess +import sys + +os.environ.setdefault("TORCH_USE_RTLD_GLOBAL", "1") + +import test_shim +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 expression per child (a candle Metal panic poisons the encoder's mutex +# for the rest of a process, test_mpsgen.py). The counters bracket only the +# block, and are recorded whether it answers or raises: the error rows assert +# their download count 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 +out = {} +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() +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]] +_ALL = ["float32", "float16", "bfloat16", "int64", "bool"] +_FLOATS = ["float32", "float16", "bfloat16"] + +# (name, code, dtypes, downloads). `X` is `_GRID` in the row's dtype on the +# device. `downloads` is the designed count: 1 for a non-empty result (the +# min/max vector of every index tensor, in one read), 0 for an empty one. +_INDEX_ROWS = [ + ("one index on dim 0", "R = X[T([0, 2, 1], L)]", _ALL, 1), + ("negative indices wrap", "R = X[T([-1, 0, -3], L)]", _ALL, 1), + ("index on dim 1 after a slice", "R = X[:, T([3, -4, 1], L)]", _ALL, 1), + ("repeated indices", "R = X[T([1, 1, 1, 1], L)]", ["float32", "int64"], 1), + ("two index tensors broadcast", "R = X[T([[0], [2]], L), T([1, -1, 0], L)]", _ALL, 1), + ("multi-dim index", "R = X[T([[0, 1], [2, 0]], L)]", _ALL, 1), + ("adjacent indices spliced at axis 1", + "Y = X.reshape(3, 2, 2)\nR = Y[:, T([[1], [0]], L), T([0, 1, 1], L)]", _FLOATS + ["int64"], 1), + ("separated indices go to the front", + "Y = X.reshape(3, 2, 2)\nR = Y[T([0, 2], L), :, T([1, -2], L)]", _FLOATS + ["int64"], 1), + ("index of a transposed view", "R = X.t()[T([3, 0], L)]", ["float32", "int64"], 1), + ("0-d index through the aten op", + "R = torch.ops.aten.index.Tensor(X, [T(1, L)])", ["float32"], 1), + ("empty index", "R = X[T([], L)]", _ALL, 0), + ("empty index on dim 1", "R = X[:, T([], L)]", ["float32", "int64"], 0), +] + +_GATHER_ROWS = [ + ("gather dim 1", "R = X.gather(1, T([[0, 3], [1, 1], [2, 0]], L))", _ALL, 1), + ("gather dim 0, index smaller off-axis", "R = X.gather(0, T([[2, 0], [1, 2]], L))", _ALL, 1), + ("gather dim 0, index longer along dim", + "R = X.gather(0, T([[0, 1, 2, 0], [2, 2, 1, 0], [1, 0, 0, 1], [0, 0, 0, 0]], L))", + ["float32", "int64"], 1), + ("gather from a transposed view", + "R = X.t().gather(1, T([[0, 2], [1, 0], [2, 2], [0, 1]], L))", _FLOATS + ["int64"], 1), + ("gather with an empty index", "R = X.gather(1, T([[], [], []], L))", ["float32", "int64"], 0), +] + +# (name, code). Shim mps against upstream cpu: same exception type, same +# text, and exactly 2 downloads (the min/max vector, then the indices to name +# the first offender). +_ERROR_ROWS = [ + ("index 5 past the end", "R = X[T([0, 5], L)]"), + ("index -4 before the start", "R = X[T([0, -4], L)]"), + ("first offender is named, not the extreme", "R = X[T([7, -9, 5], L)]"), + ("dimension names the index position", "R = X[:, T([4], L)]"), + ("second index tensor out of range", "R = X[T([0, 1], L), T([9, 1], L)]"), + ("broadcast order picks the offender", "R = X[T([[0], [5]], L), T([1, 7], L)]"), + ("gather past the end", "R = X.gather(1, T([[0], [5], [1]], L))"), + ("gather does not wrap -1", "R = X.gather(1, T([[0], [-1], [1]], L))"), + ("gather names the first offender", "R = X.gather(1, T([[7], [-1], [9]], L))"), +] + + +def _same(a, b): + if math.isnan(a) or math.isnan(b): + return math.isnan(a) and math.isnan(b) + return a == b + + +def _agree(rows): + problems = [] + for name, code, dtypes, downloads in rows: + for dtype in dtypes: + payload = dict(data=_GRID, code=code, dtype=dtype) + want = _run(dict(payload, device="cpu"), upstream=True) + got = _run(dict(payload, device="mps"), upstream=False) + label = "%s [%s]" % (name, dtype) + assert "error" not in want, (label, want) + 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"]) != (want["shape"], want["dtype"]): + problems.append("%s: %s %s, upstream %s %s" % ( + label, got["shape"], got["dtype"], want["shape"], want["dtype"])) + elif not all(_same(a, b) for a, b in zip(got["values"], want["values"])): + problems.append("%s: %s, upstream %s" % (label, got["values"], want["values"])) + 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)) + assert not problems, "\n".join(problems) + + +def test_index_tensor_computes_on_mps_and_agrees(): + """Integer `x[idx]` on mps equals upstream cpu exactly, in every listed + dtype, stays on mps, and downloads exactly the designed count.""" + if not _mps_or_skip("index.Tensor on mps"): + return + _agree(_INDEX_ROWS) + + +def test_gather_computes_on_mps_and_agrees(): + """`gather` on mps equals upstream cpu exactly (index smaller off-axis, + longer along dim, transposed self, empty), with the designed download + count.""" + if not _mps_or_skip("gather on mps"): + return + _agree(_GATHER_ROWS) + + +def test_out_of_range_indices_raise_upstreams_error_on_mps(): + """Out-of-range positive and negative indices raise upstream cpu's + exception type and exact text on mps, after exactly two downloads. This + is the row candle's unchecked Metal kernels would otherwise turn into an + out-of-buffer read.""" + if not _mps_or_skip("index bounds on mps"): + return + problems = [] + for name, code in _ERROR_ROWS: + payload = dict(data=_GRID, code=code, 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 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"])) + 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", "R = X[X > 0]", "aten.index.Tensor", "mask index is not implemented for the mps device"), + ("uint8 mask", "R = X[(X > 0).to(torch.uint8)]", "aten.index.Tensor", + "mask index is not implemented for the mps device"), + ("int32 index", "R = X[T([0, 2], torch.int32)]", "aten.index.Tensor", + "int32 index tensor is not implemented on the mps device"), + ("int32 gather index", "R = X.gather(1, T([[0], [1], [2]], torch.int32))", "aten.gather.default", + "int32 index tensor is not implemented on the mps device"), +] + + +def test_masks_and_mps_int32_indices_refuse_by_name(): + """The paths not implemented on the device refuse with the op's name and + the reason, never candle's words, and read nothing back first.""" + if not _mps_or_skip("index refusals on mps"): + return + problems = [] + for name, code, op, 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 test with counters" % name) + continue + err = got["error"] + if not err.startswith("NotImplementedError: %s:" % op): + problems.append("%s: not a refusal naming %s: %s" % (name, op, 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_the_two_ops_left_the_readback_list_and_the_rest_did_not(): + """The artefact's table: `index.Tensor` and `gather` are off it, nothing + moved into the excused column instead, and their neighbours that still + read back are still on it.""" + refused = set(_C._shim_mps_host_readback_ops()) + 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"): + 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_routers_call_the_host_halves(): + """`index_tensor_on_host`, `gather_on_host` and `index_validation_readback` + are barriers in test_shim.py's readback derivation. A barrier is only + honest while its callers are the ones its reason names, so the call graph + is read here: each host half is called only by its router, which tests + for `Device::Cpu`; the validation readback only by the device paths.""" + 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 = {} + for module, (bodies, _) in mods.items(): + for fn, body in bodies.items(): + for name in test_shim._CALL_IN_MODULE.findall(body): + if module == "aten": + callers.setdefault(name, set()).add(fn) + for other, name in test_shim._CALL_QUALIFIED.findall(body): + if other == "aten": + callers.setdefault(name, set()).add("%s::%s" % (module, fn)) + bodies = mods["aten"][0] + for host, router in (("index_tensor_on_host", "index_tensor"), + ("gather_on_host", "gather_default")): + assert host in bodies, host + assert callers.get(host, set()) - {host} == {router}, (host, callers.get(host)) + assert re.search(r"Device::Cpu", bodies[router]), router + assert callers.get("index_validation_readback", set()) <= { + "index_extremes", "index_bounds_on_device", "gather_on_device"}, \ + callers.get("index_validation_readback") + for name in ("index_tensor_on_host", "gather_on_host", "index_validation_readback"): + assert name in test_shim._MPS_READBACK_EXEMPT, name + + +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_mpsindex") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/torchnative/rust/torch_c/src/aten.rs b/torchnative/rust/torch_c/src/aten.rs index 6b82730a..82527dee 100644 --- a/torchnative/rust/torch_c/src/aten.rs +++ b/torchnative/rust/torch_c/src/aten.rs @@ -18075,7 +18075,11 @@ fn mask_to_indices( /// /// A list shorter than the tensor's rank leaves the trailing axes alone; /// upstream does not pad it out with `None`, and neither does this. -fn index_tensor( +/// +/// **This is the CPU half.** `index_tensor` below routes every non-CPU +/// `self` to `index_tensor_on_device`; this kernel reads its indices to the +/// host and is reached only from that router. +fn index_tensor_on_host( py: Python<'_>, args: &Bound<'_, PyTuple>, kwargs: Option<&Bound<'_, PyDict>>, @@ -18143,30 +18147,9 @@ fn index_tensor( } // The index tensors broadcast against each other, right-aligned. - let broadcast_rank = picks.iter().map(|(_, _, shape)| shape.len()).max().unwrap_or(0); - let mut broadcast: Vec = vec![1; broadcast_rank]; - for (_, _, shape) in &picks { - let offset = broadcast_rank - shape.len(); - for (k, &extent) in shape.iter().enumerate() { - let slot = &mut broadcast[offset + k]; - if *slot == extent || extent == 1 { - continue; - } - if *slot == 1 { - *slot = extent; - continue; - } - let named = picks - .iter() - .map(|(_, _, s)| format!("{s:?}")) - .collect::>() - .join(", "); - return Err(pyo3::exceptions::PyIndexError::new_err(format!( - "shape mismatch: indexing tensors could not be broadcast together \ - with shapes {named}" - ))); - } - } + let shapes: Vec> = picks.iter().map(|(_, _, s)| s.clone()).collect(); + let broadcast = index_broadcast_shape(&shapes)?; + let broadcast_rank = broadcast.len(); // Each index tensor's stride *through the broadcast shape*: zero where it // is being stretched, its own row-major stride where it is not. @@ -18210,10 +18193,12 @@ fn index_tensor( let extent = dims[*at] as i64; let resolved = if raw < 0 { raw + extent } else { raw }; if resolved < 0 || resolved >= extent { - // Upstream names the dimension being indexed and its size, - // not the index tensor's own shape. + // Upstream names the index tensor's position `j` among the + // index tensors (not the axis it indexes) and that axis's + // size: `x[:, i]` with `i` out of range on a (2,3,4,5) `x` + // says "dimension 0 with size 3" (measured, torch 2.13.0). return Err(pyo3::exceptions::PyIndexError::new_err(format!( - "index {raw} is out of bounds for dimension {at} with size {extent}" + "index {raw} is out of bounds for dimension {j} with size {extent}" ))); } offset += (resolved as usize) * sub[j]; @@ -18259,6 +18244,399 @@ fn index_tensor( finish(py, out, input.tag()) } +/// `aten.index.Tensor`'s router: the CPU keeps the host kernel above, every +/// other device takes `index_tensor_on_device` below. +/// +/// **Why a router and not one kernel.** The host kernel answers every index +/// kind, boolean masks included, by reading the indices to the host; that is +/// free on the CPU and is exactly what `mps_host_readback_gate` exists to +/// refuse on Metal (docs/devices/matrix.md §4). The device kernel answers +/// the integer cases with candle device ops and refuses masks by name, so the +/// op left `MPS_HOST_READBACK_OPS` for the paths it computes (issue #32). +/// +/// `index_tensor_on_host` is a named barrier in the readback derivation +/// (`test_shim.py::_MPS_READBACK_EXEMPT`) because this `match` is the only +/// way to reach it, and `test_mpsindex.py` holds both halves of that down: +/// a source check that this router is its only caller, and a counter check +/// that an mps call downloads exactly the one validation vector and no more. +/// Routing an mps tensor to the host kernel turns the second red. +fn index_tensor( + py: Python<'_>, + args: &Bound<'_, PyTuple>, + kwargs: Option<&Bound<'_, PyDict>>, +) -> PyResult> { + const OP: &str = "aten.index.Tensor"; + let input = tensor_arg(OP, args, kwargs, 0, "self")?; + let on_cpu = matches!(input.tensor()?.device(), Device::Cpu); + if on_cpu { + index_tensor_on_host(py, args, kwargs) + } else { + index_tensor_on_device(py, args, kwargs) + } +} + +/// The torch name of a non-CPU candle device, for a refusal's wording. +fn accelerator_label(device: &Device) -> &'static str { + if crate::device::is_metal(device) { + "mps" + } else if device.is_cuda() { + "cuda" + } else { + "cpu" + } +} + +/// The shape a list of index tensors broadcasts to, right-aligned, or +/// upstream's `IndexError` naming every shape. Shapes only: no data. +fn index_broadcast_shape(shapes: &[Vec]) -> PyResult> { + let rank = shapes.iter().map(|shape| shape.len()).max().unwrap_or(0); + let mut broadcast: Vec = vec![1; rank]; + for shape in shapes { + let offset = rank - shape.len(); + for (k, &extent) in shape.iter().enumerate() { + let slot = &mut broadcast[offset + k]; + if *slot == extent || extent == 1 { + continue; + } + if *slot == 1 { + *slot = extent; + continue; + } + let named = shapes + .iter() + .map(|s| format!("{s:?}")) + .collect::>() + .join(", "); + return Err(pyo3::exceptions::PyIndexError::new_err(format!( + "shape mismatch: indexing tensors could not be broadcast together \ + with shapes {named}" + ))); + } + } + Ok(broadcast) +} + +/// **The one readback the device index paths make, and it is validation, +/// not the answer.** +/// +/// candle's Metal `index_select`/`gather` kernels do not bounds-check: an +/// out-of-range index reads outside the buffer where upstream raises. So +/// before launching them, `index_tensor_on_device` and `gather_on_device` +/// reduce every index tensor to its minimum and maximum **on the device**, +/// concatenate those scalars into one short vector, and read that vector +/// here: one host download per call, two `i64` per index tensor. Only when +/// that check fails is a second vector read (the indices themselves), to +/// name the first offending value exactly as upstream does, and the call +/// then raises -- nothing is computed from it. +/// +/// It is a named exemption in the readback derivation +/// (`test_shim.py::_MPS_READBACK_EXEMPT`), and `test_mpsindex.py` asserts +/// the download count per call so this cannot quietly become a path that +/// reads the operand: docs/devices/matrix.md §4 records the budget. +fn index_validation_readback(op: &str, values: &Tensor) -> PyResult> { + values + .flatten_all() + .and_then(|t| t.to_dtype(candle_core::DType::I64)) + .and_then(|t| t.to_vec1::()) + .map_err(|e| candle_err(op, e)) +} + +/// `[min_0, max_0, min_1, max_1, ...]` of the index tensors, reduced on the +/// device and read back in one download. Every tensor must be non-empty. +fn index_extremes(op: &str, tensors: &[&Tensor]) -> PyResult> { + let mut parts: Vec = Vec::with_capacity(2 * tensors.len()); + for tensor in tensors { + let flat = tensor.flatten_all().map_err(|e| candle_err(op, e))?; + parts.push(flat.min_keepdim(0).map_err(|e| candle_err(op, e))?); + parts.push(flat.max_keepdim(0).map_err(|e| candle_err(op, e))?); + } + let ends = Tensor::cat(&parts, 0).map_err(|e| candle_err(op, e))?; + index_validation_readback(op, &ends) +} + +/// Refuse a boolean or `uint8` mask index on an accelerator, by name. +fn refuse_device_mask_index(op: &str, label: &str) -> PyErr { + not_implemented(format!( + "{op}: a boolean (or uint8) mask index is not implemented for the {label} device. \ + The result's length is the number of true elements, which this build can only \ + count by reading the mask back to host memory -- a readback under the {label} \ + label that the shim refuses rather than doing silently. Integer index tensors compute on \ + {label}; move the tensor with .cpu() to index it with a mask on the CPU." + )) +} + +/// Refuse an `int32` index tensor on mps, by name. candle's Metal backend has +/// no int32 kernel at all, not even the cast to int64 (device.rs, the int16 / +/// int32 comment), so widening it would itself be a host cast. +fn refuse_mps_int32_index(op: &str) -> PyErr { + not_implemented(format!( + "{op}: an int32 index tensor is not implemented on the mps device. \ + candle's Metal backend has no int32 kernel, not even the cast to int64, so \ + widening it would be a host cast under an mps label. Index with an int64 tensor \ + (call .long() before moving it to mps), or move the tensor with .cpu()." + )) +} + +/// `aten.index.Tensor` on an accelerator: integer index tensors only, every +/// step a candle device op, bounds checked on the device. +/// +/// The placement rules are the host kernel's (see `index_tensor_on_host`): +/// indexed axes to the front, one flat `index_select` over their block, +/// then adjacent indexed axes spliced back in place. What differs is where +/// the arithmetic runs: +/// +/// * **Bounds:** `index_extremes` -- one download of `2 * k` scalars for `k` +/// index tensors. Valid iff `-size <= min` and `max < size` per axis, the +/// same range the host loop accepts. When it fails, the indices are read +/// once more only to name the first offender in upstream's iteration +/// order (broadcast element, then index tensor), and the call raises +/// `IndexError` in upstream's CPU words. Upstream's *mps* raises +/// `torch.AcceleratorError` (a `RuntimeError`) with the same text +/// (measured 2.13.0); the CPU type is kept so one `except IndexError` +/// serves both devices here. +/// * **Negative indices:** `where(idx < 0, idx + size, idx)` on the device. +/// * **Flat offsets:** each wrapped index reshaped right-aligned, broadcast, +/// scaled by its sub-block stride and summed, all `i64` device tensors. +/// +/// An empty broadcast shape launches nothing and reads nothing: upstream +/// iterates zero elements, so it validates nothing either. +fn index_tensor_on_device( + py: Python<'_>, + args: &Bound<'_, PyTuple>, + kwargs: Option<&Bound<'_, PyDict>>, +) -> PyResult> { + use candle_core::DType; + const OP: &str = "aten.index.Tensor"; + let input = tensor_arg(OP, args, kwargs, 0, "self")?; + let source = input.tensor()?.clone(); + let tag = input.tag(); + let device = source.device().clone(); + let label = accelerator_label(&device); + let raw = required(OP, args, kwargs, 1, "indices")?; + let items: Vec> = raw.extract()?; + + 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() + ))); + } + + // (axis, int64 index tensor on `device`, its shape) + 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(), + TorchDType::Int32 if crate::device::is_metal(&device) => { + return Err(refuse_mps_int32_index(OP)); + } + 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", + )); + } + }; + // An index tensor left on the host beside a device `self` is an + // upload of data that already lives on the host, not a readback. + 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 finish(py, source, tag); + } + + 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 mut shape: Vec = broadcast.clone(); + shape.extend(trailing.iter().copied()); + + let mut out = if total == 0 || trailing.iter().product::() == 0 { + // Nothing to gather. Upstream validates nothing when the broadcast + // is empty (it iterates zero elements); when only the trailing + // block is empty it still validates, so do that first. + if total > 0 { + index_bounds_on_device(OP, &picks, &dims, &broadcast)?; + } + Tensor::zeros(shape.clone(), source.dtype(), &device).map_err(|e| candle_err(OP, e))? + } 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 mut order = indexed.clone(); + order.extend(untouched.iter().copied()); + let block: usize = indexed.iter().map(|&a| dims[a]).product(); + let span: usize = trailing.iter().product(); + // candle-metal-kernels has no `is_i64_u8` either (the same gap as + // `gather_i64_u8` below), and no `is_i64_u32`: index_select has i64 ids + // only for bf16, f16, f32, i64 and i8 values. So a bool or uint8 `self` + // is selected as `i64` and narrowed back; 0..=255 round-trips exactly. + let storage = source.dtype(); + let widen = storage == DType::U8; + let selected = source + .permute(order) + .and_then(|t| t.contiguous()) + .and_then(|t| if widen { t.to_dtype(DType::I64) } else { Ok(t) }) + .and_then(|t| t.reshape((block, span))) + .and_then(|t| t.index_select(&flat, 0)) + .and_then(|t| t.reshape(shape.clone())) + .map_err(|e| candle_err(OP, e))?; + if widen { + selected.to_dtype(storage).map_err(|e| candle_err(OP, e))? + } else { + selected + } + }; + + let adjacent = indexed.windows(2).all(|pair| pair[1] == pair[0] + 1); + let at = indexed[0]; + if adjacent && at > 0 { + let mut permutation: Vec = (0..at).map(|k| broadcast_rank + k).collect(); + permutation.extend(0..broadcast_rank); + permutation.extend((at..trailing.len()).map(|k| broadcast_rank + k)); + out = out + .permute(permutation) + .and_then(|t| t.contiguous()) + .map_err(|e| candle_err(OP, e))?; + } + finish(py, out, tag) +} + +/// The device bounds check for `index_tensor_on_device`; see there. +fn index_bounds_on_device( + op: &str, + picks: &[(usize, Tensor, Vec)], + dims: &[usize], + broadcast: &[usize], +) -> PyResult<()> { + let tensors: Vec<&Tensor> = picks.iter().map(|(_, t, _)| t).collect(); + let ends = index_extremes(op, &tensors)?; + let ok = picks.iter().enumerate().all(|(j, (at, _, _))| { + let extent = dims[*at] as i64; + ends[2 * j] >= -extent && ends[2 * j + 1] < extent + }); + if ok { + return Ok(()); + } + // The error path: read the indices once (one more download) only to name + // the first offender in upstream's order, then raise. + let flats: Vec = tensors + .iter() + .map(|t| t.flatten_all()) + .collect::>>() + .map_err(|e| candle_err(op, e))?; + let joined = Tensor::cat(&flats, 0).map_err(|e| candle_err(op, e))?; + let values = index_validation_readback(op, &joined)?; + let rank = broadcast.len(); + let mut starts = Vec::with_capacity(picks.len()); + let mut strides: Vec> = Vec::with_capacity(picks.len()); + let mut start = 0usize; + for (_, _, shape) in picks { + starts.push(start); + start += shape.iter().product::(); + let mut own = vec![1usize; shape.len()]; + for k in (0..shape.len().saturating_sub(1)).rev() { + own[k] = own[k + 1] * shape[k + 1]; + } + let offset = rank - shape.len(); + let mut over = vec![0usize; rank]; + for k in 0..shape.len() { + over[offset + k] = if shape[k] == 1 { 0 } else { own[k] }; + } + strides.push(over); + } + let total: usize = broadcast.iter().product(); + let mut coords = vec![0usize; rank]; + for linear in 0..total { + let mut rest = linear; + for k in (0..rank).rev() { + coords[k] = rest % broadcast[k]; + rest /= broadcast[k]; + } + for (j, (at, _, _)) in picks.iter().enumerate() { + let position: usize = (0..rank).map(|k| coords[k] * strides[j][k]).sum(); + let raw = values[starts[j] + position]; + let extent = dims[*at] as i64; + if raw < -extent || raw >= extent { + return Err(pyo3::exceptions::PyIndexError::new_err(format!( + "index {raw} is out of bounds for dimension {j} with size {extent}" + ))); + } + } + } + // The extremes said out of range and the walk found nothing: a read of + // the wrong bytes, not an index error. Say so rather than invent one. + Err(pyo3::exceptions::PyRuntimeError::new_err(format!( + "{op}: the device bounds check reported an index out of range that the \ + indices themselves do not contain -- an internal inconsistency" + ))) +} + // --------------------------------------------------------------------------- // In-place ops // @@ -25959,7 +26337,10 @@ fn scatter_reduce_two( /// `sparse_grad` is accepted and ignored -- it selects an autograd /// representation, and there is no autograd here. Upstream's forward answer is /// the same for either value (measured). -fn gather_default( +/// +/// **This is the CPU half**, reached only from the `gather_default` router +/// below; every non-CPU `self` goes to `gather_on_device`. +fn gather_on_host( py: Python<'_>, args: &Bound<'_, PyTuple>, kwargs: Option<&Bound<'_, PyDict>>, @@ -25970,41 +26351,9 @@ fn gather_default( let index = tensor_arg(OP, args, kwargs, 2, "index")?; let _sparse_grad = bool_arg(args, kwargs, 3, "sparse_grad")?; - if !matches!(index.tag(), TorchDType::Int64 | TorchDType::Int32) { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "gather(): Expected dtype int32/int64 for index", - )); - } - let tag = input.tag(); - // `ensure_nonempty_dim`/`ensure_nonempty_size`: a 0-d tensor counts as one - // dimension of extent 1 for the purposes of both the rank check and the - // bounds check. The *output* still gets the index's real shape. - let self_dims: Vec = if input.tensor()?.rank() == 0 { - vec![1] - } else { - input.tensor()?.dims().to_vec() - }; - let idx_shape = index.tensor()?.dims().to_vec(); - let idx_dims: Vec = if idx_shape.is_empty() { - vec![1] - } else { - idx_shape.clone() - }; - if idx_dims.len() != self_dims.len() { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "Index tensor must have the same number of dimensions as input tensor", - )); - } - let dim = normalise_dim(OP, dim_raw, input.tensor()?.rank())?; - for d in 0..self_dims.len() { - if d != dim && idx_dims[d] > self_dims[d] { - return Err(pyo3::exceptions::PyRuntimeError::new_err(format!( - "Size does not match at dimension {d} expected index {idx_dims:?} \ - to be no larger than self {self_dims:?} apart from dimension {dim}" - ))); - } - } + let (dim, self_dims, idx_shape, idx_dims) = + gather_checked_shapes(OP, input.tensor()?, index.tensor()?, index.tag(), dim_raw)?; let source = read_flat(OP, input.tensor()?, tag)?; let positions = match read_flat(OP, index.tensor()?, index.tag())? { @@ -26055,6 +26404,164 @@ fn gather_default( finish(py, tensor, tag) } +/// `aten.gather.default`'s router, for `index_tensor`'s reason: the CPU keeps +/// the host kernel, every other device takes `gather_on_device`. +/// `gather_on_host` is a named barrier in the readback derivation and this is +/// its only caller (`test_mpsindex.py` checks both). +fn gather_default( + py: Python<'_>, + args: &Bound<'_, PyTuple>, + kwargs: Option<&Bound<'_, PyDict>>, +) -> PyResult> { + const OP: &str = "aten.gather.default"; + let input = tensor_arg(OP, args, kwargs, 0, "self")?; + let on_cpu = matches!(input.tensor()?.device(), Device::Cpu); + if on_cpu { + gather_on_host(py, args, kwargs) + } else { + gather_on_device(py, args, kwargs) + } +} + +/// `gather`'s argument rules, shared by both paths: `(dim, self dims as +/// upstream's ensure_nonempty sees them, index shape, index dims likewise)`. +/// Shapes and dtypes only; no data is read. +fn gather_checked_shapes( + op: &str, + input: &Tensor, + index: &Tensor, + index_tag: TorchDType, + dim_raw: isize, +) -> PyResult<(usize, Vec, Vec, Vec)> { + if !matches!(index_tag, TorchDType::Int64 | TorchDType::Int32) { + return Err(pyo3::exceptions::PyRuntimeError::new_err( + "gather(): Expected dtype int32/int64 for index", + )); + } + let self_dims: Vec = if input.rank() == 0 { + vec![1] + } else { + input.dims().to_vec() + }; + let idx_shape = index.dims().to_vec(); + let idx_dims: Vec = if idx_shape.is_empty() { + vec![1] + } else { + idx_shape.clone() + }; + if idx_dims.len() != self_dims.len() { + return Err(pyo3::exceptions::PyRuntimeError::new_err( + "Index tensor must have the same number of dimensions as input tensor", + )); + } + let dim = normalise_dim(op, dim_raw, input.rank())?; + for d in 0..self_dims.len() { + if d != dim && idx_dims[d] > self_dims[d] { + return Err(pyo3::exceptions::PyRuntimeError::new_err(format!( + "Size does not match at dimension {d} expected index {idx_dims:?} \ + to be no larger than self {self_dims:?} apart from dimension {dim}" + ))); + } + } + Ok((dim, self_dims, idx_shape, idx_dims)) +} + +/// `aten.gather.default` on an accelerator: candle's `gather`, bounds checked +/// on the device first. +/// +/// candle's Metal `gather` kernel does not bounds-check, and does not accept +/// an index smaller than `self` off-axis the way upstream does, so: +/// +/// * **Bounds:** `index_extremes` -- one download of the index's minimum and +/// maximum. `gather` does **not** wrap negatives (upstream, measured), so +/// valid iff `0 <= min` and `max < size`. On failure the index is read once +/// more to name the first offender in flat order, and the call raises +/// `RuntimeError` in upstream's CPU words (`index -1 is out of bounds for +/// dimension 1 with size 4`). Upstream's mps says `gather: index -1 is out +/// of bounds for dimension with size 4` as a `torch.AcceleratorError`, a +/// `RuntimeError` subclass (measured 2.13.0); the CPU text is kept. +/// * **Shape:** `self` is narrowed to the index's extent off-axis and made +/// contiguous, then `gather` along `dim`; the result takes the index's +/// own shape (0-d included). +fn gather_on_device( + py: Python<'_>, + args: &Bound<'_, PyTuple>, + kwargs: Option<&Bound<'_, PyDict>>, +) -> PyResult> { + use candle_core::DType; + const OP: &str = "aten.gather.default"; + let input = tensor_arg(OP, args, kwargs, 0, "self")?; + let dim_raw = dim_arg(args, kwargs, 1, "dim")?.ok_or_else(|| missing(OP, "dim"))?; + let index_arg = tensor_arg(OP, args, kwargs, 2, "index")?; + let _sparse_grad = bool_arg(args, kwargs, 3, "sparse_grad")?; + let tag = input.tag(); + let source = input.tensor()?.clone(); + let device = source.device().clone(); + let (dim, self_dims, idx_shape, idx_dims) = + gather_checked_shapes(OP, &source, index_arg.tensor()?, index_arg.tag(), dim_raw)?; + + let index = match index_arg.tag() { + TorchDType::Int32 if crate::device::is_metal(&device) => { + return Err(refuse_mps_int32_index(OP)); + } + TorchDType::Int32 => index_arg + .tensor()? + .to_dtype(DType::I64) + .map_err(|e| candle_err(OP, e))?, + _ => index_arg.tensor()?.clone(), + }; + let index = index + .to_device(&device) + .and_then(|t| t.contiguous()) + .map_err(|e| candle_err(OP, e))?; + + let count: usize = idx_dims.iter().product(); + if count == 0 { + let out = Tensor::zeros(idx_shape, source.dtype(), &device).map_err(|e| candle_err(OP, e))?; + return finish(py, out, tag); + } + + let size = self_dims[dim] as i64; + let ends = index_extremes(OP, &[&index])?; + if ends[0] < 0 || ends[1] >= size { + // The error path: one more download, only to name the offender. + let values = index_validation_readback(OP, &index)?; + let target = values + .iter() + .copied() + .find(|&v| v < 0 || v >= size) + .unwrap_or(if ends[0] < 0 { ends[0] } else { ends[1] }); + return Err(pyo3::exceptions::PyRuntimeError::new_err(format!( + "index {target} is out of bounds for dimension {dim} with size {}", + self_dims[dim] + ))); + } + + let mut narrowed = source.reshape(self_dims.clone()).map_err(|e| candle_err(OP, e))?; + for d in 0..self_dims.len() { + if d != dim && idx_dims[d] < self_dims[d] { + narrowed = narrowed.narrow(d, 0, idx_dims[d]).map_err(|e| candle_err(OP, e))?; + } + } + // candle-metal-kernels instantiates no `gather_i64_u8`, so a bool or + // uint8 `self` is widened to `u32` and narrowed back, both device casts; + // the values are 0..=255 either way, so the round trip is exact. + let storage = narrowed.dtype(); + let widen = storage == DType::U8; + if widen { + narrowed = narrowed.to_dtype(DType::U32).map_err(|e| candle_err(OP, e))?; + } + let mut out = narrowed + .contiguous() + .and_then(|t| t.gather(&index.reshape(idx_dims.clone())?.contiguous()?, dim)) + .and_then(|t| t.reshape(idx_shape)) + .map_err(|e| candle_err(OP, e))?; + if widen { + out = out.to_dtype(storage).map_err(|e| candle_err(OP, e))?; + } + finish(py, out, tag) +} + /// Round a vector through the storage dtype and back, in one candle pass. /// /// `multinomial`'s fast path needs this twice, because upstream writes both diff --git a/torchnative/rust/torch_c/src/device.rs b/torchnative/rust/torch_c/src/device.rs index a8ed79f1..0db86260 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; 89] = [ +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`, @@ -773,11 +773,16 @@ pub const MPS_HOST_READBACK_OPS: [&str; 89] = [ "aten.floor_divide.default", "aten.fmod.Scalar", "aten.fmod.Tensor", - "aten.gather.default", + // `aten.gather.default` and `aten.index.Tensor` left on 2026-10-16 (issue + // #32) for their integer-index paths: `aten.rs::gather_on_device` and + // `index_tensor_on_device` are candle device ops, bounds-checked by + // reading back only each index tensor's on-device min and max (one + // download per call, `index_validation_readback`). Boolean-mask indices + // and int32 indices on mps refuse inside those kernels, by name. + // docs/devices/matrix.md §4, test_mpsindex.py. "aten.histc.default", "aten.i0.default", "aten.im2col.default", - "aten.index.Tensor", "aten.index_add.default", "aten.index_add_.default", "aten.index_put_.default", @@ -863,7 +868,7 @@ pub const MPS_HOST_READBACK_OPS: [&str; 89] = [ ]; /// The two ops that read device bytes back and are **not** refused, with the -/// reason each is different in kind from the eighty-nine above. +/// reason each is different in kind from the eighty-seven above. /// /// The scan finds these too, so leaving them out of `MPS_HOST_READBACK_OPS` /// without saying why would look like an oversight rather than a decision. @@ -908,7 +913,7 @@ pub const MPS_READBACK_BUT_ALLOWED: [&str; 2] = [ /// the build does turns the gate red rather than misdirecting a user. /// /// It covers the ten names the deepened derivation added -/// (docs/devices/matrix.md §7.16) and no others. The other eighty-nine +/// (docs/devices/matrix.md §7.16) and no others. The other entries /// predate the note mechanism and carry the generic wording; extending the /// table is a measurement each time, which is why it was not done in bulk. ///