diff --git a/.github/scripts/gate_suites.py b/.github/scripts/gate_suites.py index 7621a62..ede9d5a 100644 --- a/.github/scripts/gate_suites.py +++ b/.github/scripts/gate_suites.py @@ -120,6 +120,7 @@ "devices/mps/test_metalcount.py": BOTH, "devices/mps/test_metalplace.py": BOTH, "devices/mps/test_mpsattn.py": BOTH, + "devices/mps/test_mpsclone.py": BOTH, "devices/mps/test_mpsconst.py": BOTH, "devices/mps/test_mpsnarrow.py": BOTH, "devices/mps/test_mpsgen.py": BOTH, diff --git a/tests/devices/mps/test_mpsclone.py b/tests/devices/mps/test_mpsclone.py new file mode 100644 index 0000000..c5d2f6d --- /dev/null +++ b/tests/devices/mps/test_mpsclone.py @@ -0,0 +1,240 @@ +"""`clone()` on mps is a copy, not a second name for the same buffer (issue #97). + +candle's `MetalStorage::try_clone` returned `self.clone()`, which shares the +`Arc`. candle never writes a buffer in place, so that was safe for +candle. torch._C does (`add_`, `fill_`, setitem), so on mps + + x = torch.ones(3, device="mps") + y = x.clone(); y.add_(1) # x changed too, before the fix + +The fix lives in the vendored fork (`int8-candle-0.11.0-cpu.patch`): a fresh +buffer of the same element count, filled by one blit. + +What this file asserts, each case in a fresh subprocess that imports the +vendored tree and asserts `hasattr(torch._C, "_aten_implemented")`: + +* *agrees*: for `clone()`, `clone(memory_format=torch.contiguous_format)`, + `detach().clone()`, `x.t().clone()` and `x[:, 1:].clone()`, crossed with + `add_(1)`, `fill_(7)`, `mul_(0)` and `y[0] = `, in float32, + float16, bfloat16 and int64: the source still holds its original values and + the clone holds upstream's answer. Upstream's answer is computed here in + plain Python; the inputs are small integers, exact in every one of the four + dtypes, so no tolerance is involved. +* *placement*: `_C._metal_counters()` shows at least one blit or compute + encoder opened during the `clone` call itself. A shared-buffer clone opens + none. +* *negative control*: `x.detach().add_(1)` and `torch.ops.aten.alias(x).add_(1)` + DO change `x`, because upstream defines both as views. A fix that made every + aliasing op copy would pass the cases above and fail these. + +What this file cannot find: a clone that copies through the host (download +then upload) would still open a blit for the readback and pass the counter +check; the counter bracket only proves the GPU was asked to do something. +""" + +import json +import os +import subprocess +import sys + +os.environ.setdefault("TORCH_USE_RTLD_GLOBAL", "1") + +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") +_VENDOR_SHIM = os.path.join(_VENDOR_DIR, "torch", "_C.abi3.so") + +_DTYPES = ("float32", "float16", "bfloat16", "int64") +_SHAPE = (2, 3) +_VALUES = [[-3, -1, 0], [2, -7, 5]] +_ROW = 9 # the value setitem writes into row 0 + +_VARIANTS = ("clone", "clone_contiguous", "detach_clone", "t_clone", "slice_clone") +_OPS = ("add_", "fill_", "mul_", "setitem") + +_CHILD = r""" +import json, os, sys +req = json.loads(sys.argv[1]) +import torch +assert hasattr(torch._C, "_aten_implemented"), "the child imported upstream, not the shim" +out = {"no_mps": None, "results": {}} +try: + torch._C._aten_dispatch("aten.ones.default", [1], device=torch._C.device("mps")) +except (NotImplementedError, RuntimeError) as e: + out["no_mps"] = "no mps device -- %s" % str(e).splitlines()[0] + print(json.dumps(out)); raise SystemExit(0) + +def counters(): + c = torch._C._metal_counters() + return c["blit_encoders"] + c["compute_encoders"] + +def take(x, variant): + if variant == "clone": + return x.clone() + if variant == "clone_contiguous": + return x.clone(memory_format=torch.contiguous_format) + if variant == "detach_clone": + return x.detach().clone() + if variant == "t_clone": + return x.t().clone() + if variant == "slice_clone": + return x[:, 1:].clone() + if variant == "detach_view": + return x.detach() + if variant == "alias_view": + return torch.ops.aten.alias(x) + raise ValueError(variant) + +def mutate(y, op, dt): + if op == "add_": + y.add_(1) + elif op == "fill_": + y.fill_(7) + elif op == "mul_": + y.mul_(0) + elif op == "setitem": + y[0] = torch.full((y.shape[1],), req["row"], dtype=dt, device="mps") + else: + raise ValueError(op) + +for name in req["dtypes"]: + dt = getattr(torch, name) + rec = {} + try: + x = torch.tensor(req["values"], dtype=dt).to("mps") + before = counters() + y = take(x, req["variant"]) + rec["encoders"] = counters() - before + mutate(y, req["op"], dt) + rec["device"] = str(y.device) + rec["x"] = [[float(v) for v in r] for r in x.cpu().tolist()] + rec["y"] = [[float(v) for v in r] for r in y.cpu().tolist()] + except Exception as e: + rec["error"] = "%s: %s" % (type(e).__name__, e) + out["results"][name] = rec +print(json.dumps(out)) +""" + + +def _run(variant, op): + if not os.path.isfile(_VENDOR_SHIM): + raise _skip.Skip("vendored tree not built: no %s " + "(scripts/vendor/vendor_torch.sh + scripts/vendor/install_shim.sh)" + % _VENDOR_SHIM) + env = dict(os.environ) + env["PYTHONPATH"] = _VENDOR_DIR + env["TORCH_USE_RTLD_GLOBAL"] = "1" + req = {"variant": variant, "op": op, "dtypes": list(_DTYPES), + "values": _VALUES, "row": _ROW} + proc = subprocess.run([sys.executable, "-c", _CHILD, json.dumps(req)], + capture_output=True, text=True, timeout=600, env=env) + lines = [l for l in proc.stdout.splitlines() if l.startswith("{")] + assert proc.returncode == 0 and lines, ( + "the mps child (%s, %s) exited %d:\n%s" % (variant, op, proc.returncode, proc.stderr[-2000:])) + out = json.loads(lines[-1]) + if out["no_mps"]: + raise _skip.Skip(out["no_mps"]) + return out["results"] + + +def _view(variant): + """Upstream's value of the cloned view, in plain Python.""" + v = [list(r) for r in _VALUES] + if variant == "t_clone": + return [list(c) for c in zip(*v)] + if variant == "slice_clone": + return [r[1:] for r in v] + return v + + +def _apply(m, op): + if op == "add_": + return [[e + 1 for e in r] for r in m] + if op == "fill_": + return [[7 for _ in r] for r in m] + if op == "mul_": + return [[0 for _ in r] for r in m] + if op == "setitem": + return [[_ROW for _ in m[0]]] + [list(r) for r in m[1:]] + raise ValueError(op) + + +def _f(m): + return [[float(e) for e in r] for r in m] + + +def _check_clone(variant, op): + results = _run(variant, op) + expected_y = _f(_apply(_view(variant), op)) + expected_x = _f(_VALUES) + for name in _DTYPES: + rec = results[name] + tag = "%s().%s on mps, %s" % (variant, op, name) + # A strided clone (`x.t().clone()`, `x[:, 1:].clone()`) keeps the view's + # strides, and an in-place write through a strided receiver on mps is + # issue #32's, landing with PR #89. Until then that refusal is named + # and skipped, so this suite still runs and passes the moment the + # writes land; any other refusal stays a failure. + if ("error" in rec and variant in ("t_clone", "slice_clone") + and "writing through a non-contiguous view" in rec["error"]): + raise _skip.Skip("%s: in-place writes through a strided mps receiver are " + "refused on this tree (issue #32, PR #89)" % tag) + assert "error" not in rec, "%s refused: %s" % (tag, rec["error"]) + assert rec["device"].startswith("mps"), (tag, rec["device"]) + assert rec["x"] == expected_x, ( + "%s changed the source (the clone shares its buffer): x=%r, expected %r" + % (tag, rec["x"], expected_x)) + assert rec["y"] == expected_y, ( + "%s: clone holds %r, upstream says %r" % (tag, rec["y"], expected_y)) + assert rec["encoders"] >= 1, ( + "%s: clone opened %d blit/compute encoders; a copy opens at least one" + % (tag, rec["encoders"])) + return len(_DTYPES) + + +def _make(variant, op): + def test(): + assert _check_clone(variant, op) == len(_DTYPES) + test.__name__ = "test_%s_then_%s_leaves_the_source_alone_on_mps" % (variant, op.rstrip("_")) + return test + + +for _v in _VARIANTS: + for _o in _OPS: + _t = _make(_v, _o) + globals()[_t.__name__] = _t +del _v, _o, _t + + +def _check_view_writes_through(variant): + results = _run(variant, "add_") + expected = _f(_apply(_VALUES, "add_")) + for name in _DTYPES: + rec = results[name] + tag = "%s.add_(1) on mps, %s" % (variant, name) + assert "error" not in rec, "%s refused: %s" % (tag, rec["error"]) + assert rec["x"] == expected, ( + "%s did not change the source, but upstream defines it as a view: " + "x=%r, expected %r" % (tag, rec["x"], expected)) + assert rec["y"] == expected, (tag, rec["y"], expected) + + +def test_detach_is_still_a_view_on_mps(): + _check_view_writes_through("detach_view") + + +def test_aten_alias_is_still_a_view_on_mps(): + _check_view_writes_through("alias_view") + + +def _main(): + items = [(name, fn) for name, fn in sorted(globals().items()) + if name.startswith("test_")] + failures = _skip.run_tests(items, suite="test_mpsclone") + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/torchnative/rust/vendor/candle-core/src/metal_backend/mod.rs b/torchnative/rust/vendor/candle-core/src/metal_backend/mod.rs index d7880e3..49b65eb 100644 --- a/torchnative/rust/vendor/candle-core/src/metal_backend/mod.rs +++ b/torchnative/rust/vendor/candle-core/src/metal_backend/mod.rs @@ -181,7 +181,20 @@ impl BackendStorage for MetalStorage { type Device = MetalDevice; fn try_clone(&self, _: &Layout) -> Result { - Ok(self.clone()) + // torchnative: upstream returned `self.clone()`, which shares the + // `Arc`. candle never writes a buffer in place, so that was + // safe there; torch._C does (`add_`, `fill_`, setitem), and + // `x.clone().add_(1)` then changed `x` on mps. A real copy: a fresh + // buffer of the same element count, filled by one blit, so the + // caller's layout (offset and strides) stays valid unchanged. + let buffer = self.device.new_buffer(self.count, self.dtype, "try_clone")?; + let length = self.count * self.dtype.size_in_bytes(); + if length > 0 { + let mut blit = self.device.blit_command_encoder()?; + blit.set_label("try_clone"); + blit.copy_from_buffer(&self.buffer, 0, &buffer, 0, length); + } + Ok(Self::new(buffer, self.device.clone(), self.count, self.dtype)) } fn dtype(&self) -> DType { diff --git a/torchnative/rust/vendor/int8-candle-0.11.0-cpu.patch b/torchnative/rust/vendor/int8-candle-0.11.0-cpu.patch index 32c87c2..16404ee 100644 --- a/torchnative/rust/vendor/int8-candle-0.11.0-cpu.patch +++ b/torchnative/rust/vendor/int8-candle-0.11.0-cpu.patch @@ -1014,7 +1014,7 @@ index 9cae473..a238e45 100644 } diff --git a/src/metal_backend/mod.rs b/src/metal_backend/mod.rs -index b3e31f4..d7880e3 100644 +index b3e31f4..49b65eb 100644 --- a/src/metal_backend/mod.rs +++ b/src/metal_backend/mod.rs @@ -21,6 +21,96 @@ use std::sync::{Arc, Mutex, PoisonError, RwLock, TryLockError}; @@ -1114,7 +1114,29 @@ index b3e31f4..d7880e3 100644 pub fn buffer_o<'a>(buffer: &'a Buffer, l: &Layout, dtype: DType) -> BufferOffset<'a> { BufferOffset { buffer, -@@ -114,6 +204,7 @@ impl BackendStorage for MetalStorage { +@@ -91,7 +181,20 @@ impl BackendStorage for MetalStorage { + type Device = MetalDevice; + + fn try_clone(&self, _: &Layout) -> Result { +- Ok(self.clone()) ++ // torchnative: upstream returned `self.clone()`, which shares the ++ // `Arc`. candle never writes a buffer in place, so that was ++ // safe there; torch._C does (`add_`, `fill_`, setitem), and ++ // `x.clone().add_(1)` then changed `x` on mps. A real copy: a fresh ++ // buffer of the same element count, filled by one blit, so the ++ // caller's layout (offset and strides) stays valid unchanged. ++ let buffer = self.device.new_buffer(self.count, self.dtype, "try_clone")?; ++ let length = self.count * self.dtype.size_in_bytes(); ++ if length > 0 { ++ let mut blit = self.device.blit_command_encoder()?; ++ blit.set_label("try_clone"); ++ blit.copy_from_buffer(&self.buffer, 0, &buffer, 0, length); ++ } ++ Ok(Self::new(buffer, self.device.clone(), self.count, self.dtype)) + } + + fn dtype(&self) -> DType { +@@ -114,6 +217,7 @@ impl BackendStorage for MetalStorage { DType::F32 => Ok(CpuStorage::F32(self.to_cpu()?)), DType::F64 => Ok(CpuStorage::F64(self.to_cpu()?)), DType::F8E4M3 => Ok(CpuStorage::F8E4M3(self.to_cpu()?)), @@ -1122,7 +1144,7 @@ index b3e31f4..d7880e3 100644 DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => { Err(crate::Error::UnsupportedDTypeForOp(self.dtype, "to_cpu_storage").bt()) } -@@ -475,6 +566,7 @@ impl BackendStorage for MetalStorage { +@@ -475,6 +579,7 @@ impl BackendStorage for MetalStorage { DType::BF16 => contiguous::const_set::BFLOAT, DType::F32 => contiguous::const_set::FLOAT, DType::I64 => contiguous::const_set::I64, @@ -1130,7 +1152,7 @@ index b3e31f4..d7880e3 100644 DType::U32 => contiguous::const_set::U32, DType::U8 => contiguous::const_set::U8, DType::F8E4M3 => crate::bail!("unsupported const-set f8e4m3"), -@@ -506,6 +598,7 @@ impl BackendStorage for MetalStorage { +@@ -506,6 +611,7 @@ impl BackendStorage for MetalStorage { DType::BF16 => strided::const_set::BFLOAT, DType::F32 => strided::const_set::FLOAT, DType::I64 => strided::const_set::I64, @@ -1138,7 +1160,7 @@ index b3e31f4..d7880e3 100644 DType::U32 => strided::const_set::U32, DType::U8 => strided::const_set::U8, DType::F8E4M3 => crate::bail!("unsupported const-set f8e4m3"), -@@ -537,6 +630,7 @@ impl BackendStorage for MetalStorage { +@@ -537,6 +643,7 @@ impl BackendStorage for MetalStorage { (DType::U8, Scalar::U8(s)) => set(self, s, l), (DType::U32, Scalar::U32(s)) => set(self, s, l), (DType::I64, Scalar::I64(s)) => set(self, s, l), @@ -1146,7 +1168,7 @@ index b3e31f4..d7880e3 100644 (DType::F16, Scalar::F16(s)) => set(self, s, l), (DType::BF16, Scalar::BF16(s)) => set(self, s, l), (DType::F32, Scalar::F32(s)) => set(self, s, l), -@@ -581,6 +675,18 @@ impl BackendStorage for MetalStorage { +@@ -581,6 +688,18 @@ impl BackendStorage for MetalStorage { (DType::I64, DType::F32) => "cast_i64_f32", (DType::I64, DType::U32) => "cast_i64_u32", (DType::I64, DType::U8) => "cast_i64_u8", @@ -1165,7 +1187,7 @@ index b3e31f4..d7880e3 100644 (DType::F16, DType::BF16) => "cast_f16_bf16", (DType::F16, DType::F32) => "cast_f16_f32", -@@ -634,6 +740,18 @@ impl BackendStorage for MetalStorage { +@@ -634,6 +753,18 @@ impl BackendStorage for MetalStorage { (DType::I64, DType::F16) => "cast_i64_f16_strided", (DType::I64, DType::U32) => "cast_i64_u32_strided", (DType::I64, DType::U8) => "cast_i64_u8_strided", @@ -1184,7 +1206,7 @@ index b3e31f4..d7880e3 100644 (DType::U32, DType::BF16) => "cast_u32_bf16_strided", (DType::U32, DType::F16) => "cast_u32_f16_strided", -@@ -878,6 +996,9 @@ impl BackendStorage for MetalStorage { +@@ -878,6 +1009,9 @@ impl BackendStorage for MetalStorage { (DType::U8, DType::I64) => "where_u8_i64", (DType::U8, DType::U32) => "where_u8_u32", (DType::U8, DType::U8) => "where_u8_u8", @@ -1194,7 +1216,7 @@ index b3e31f4..d7880e3 100644 (left, right) => crate::bail!("Metal where_cond {left:?} {right:?} not implemented"), }; let src = buffer_o(&self.buffer, layout, self.dtype); -@@ -1483,6 +1604,9 @@ impl BackendStorage for MetalStorage { +@@ -1483,6 +1617,9 @@ impl BackendStorage for MetalStorage { (DType::I64, DType::BF16) => "gather_i64_bf16", (DType::I64, DType::U32) => "gather_i64_u32", (DType::I64, DType::I64) => "gather_i64_i64", @@ -1204,7 +1226,7 @@ index b3e31f4..d7880e3 100644 (left, right) => crate::bail!("Metal gather {left:?} {right:?} not implemented"), }; let encoder = self.device.command_encoder()?; -@@ -1527,6 +1651,9 @@ impl BackendStorage for MetalStorage { +@@ -1527,6 +1664,9 @@ impl BackendStorage for MetalStorage { (DType::I64, DType::F32) => "s_i64_f32", (DType::I64, DType::F16) => "s_i64_f16", (DType::I64, DType::BF16) => "s_i64_bf16", @@ -1214,7 +1236,7 @@ index b3e31f4..d7880e3 100644 _ => Err(MetalError::UnexpectedDType { msg: "scatter ids should be u8/u32/i64", expected: DType::U32, -@@ -1576,6 +1703,9 @@ impl BackendStorage for MetalStorage { +@@ -1576,6 +1716,9 @@ impl BackendStorage for MetalStorage { (DType::I64, DType::F32) => "sa_i64_f32", (DType::I64, DType::F16) => "sa_i64_f16", (DType::I64, DType::BF16) => "sa_i64_bf16", @@ -1224,7 +1246,7 @@ index b3e31f4..d7880e3 100644 _ => Err(MetalError::UnexpectedDType { msg: "scatter-add ids should be u8/u32/i64", expected: DType::U32, -@@ -1638,6 +1768,9 @@ impl BackendStorage for MetalStorage { +@@ -1638,6 +1781,9 @@ impl BackendStorage for MetalStorage { (DType::I64, DType::F32) => "is_i64_f32", (DType::I64, DType::F16) => "is_i64_f16", (DType::I64, DType::BF16) => "is_i64_bf16", @@ -1234,7 +1256,7 @@ index b3e31f4..d7880e3 100644 (left, right) => { crate::bail!("Metal contiguous index_select {left:?} {right:?} not implemented") -@@ -1686,6 +1819,9 @@ impl BackendStorage for MetalStorage { +@@ -1686,6 +1832,9 @@ impl BackendStorage for MetalStorage { (DType::I64, DType::I64) => "ia_i64_i64", (DType::I64, DType::U32) => "ia_i64_u32", (DType::I64, DType::U8) => "ia_i64_u8", @@ -1244,7 +1266,7 @@ index b3e31f4..d7880e3 100644 (DType::U32, DType::BF16) => "ia_u32_bf16", (DType::U32, DType::F16) => "ia_u32_f16", -@@ -1809,6 +1945,7 @@ impl BackendStorage for MetalStorage { +@@ -1809,6 +1958,7 @@ impl BackendStorage for MetalStorage { DType::F16 => candle_metal_kernels::copy2d::HALF, DType::BF16 => candle_metal_kernels::copy2d::BFLOAT, DType::I64 => candle_metal_kernels::copy2d::I64, @@ -1252,7 +1274,7 @@ index b3e31f4..d7880e3 100644 DType::I32 => candle_metal_kernels::copy2d::I32, DType::I16 => candle_metal_kernels::copy2d::I16, DType::U32 => candle_metal_kernels::copy2d::U32, -@@ -1878,6 +2015,7 @@ impl BackendStorage for MetalStorage { +@@ -1878,6 +2028,7 @@ impl BackendStorage for MetalStorage { DType::F16 => candle_metal_kernels::unary::strided::copy::HALF, DType::BF16 => candle_metal_kernels::unary::strided::copy::BFLOAT, DType::I64 => candle_metal_kernels::unary::strided::copy::I64, @@ -1260,7 +1282,7 @@ index b3e31f4..d7880e3 100644 DType::U32 => candle_metal_kernels::unary::strided::copy::U32, DType::U8 => candle_metal_kernels::unary::strided::copy::U8, dtype => crate::bail!("Metal copy_strided {dtype:?} not implemented"), -@@ -2015,6 +2153,8 @@ impl MetalStorage { +@@ -2015,6 +2166,8 @@ impl MetalStorage { blit.copy_from_buffer(&self.buffer, 0, &buffer, 0, size); } self.device.flush_and_wait_current()?; @@ -1269,7 +1291,7 @@ index b3e31f4..d7880e3 100644 Ok(read_to_vec(&buffer, self.count)) } } -@@ -2166,6 +2306,13 @@ impl BackendDevice for MetalDevice { +@@ -2166,6 +2319,13 @@ impl BackendDevice for MetalDevice { .with_label(label) .build(), ), @@ -1283,7 +1305,7 @@ index b3e31f4..d7880e3 100644 CpuStorageRef::F6E2M3(_) | CpuStorageRef::F6E3M2(_) | CpuStorageRef::F4(_) -@@ -2249,6 +2396,13 @@ impl BackendDevice for MetalDevice { +@@ -2249,6 +2409,13 @@ impl BackendDevice for MetalDevice { .with_label(label) .build(), ),