sort/argsort/topk tie order matches upstream on libc++ platforms (#33)

.github/scripts/gate_suites.py

@@ -190,6 +190,7 @@
     "ops/test_rnn.py": BOTH,
     "ops/test_scatter.py": BOTH,
     "ops/test_signedzero.py": BOTH,
+    "ops/test_sortties.py": BOTH,
     "ops/test_strided.py": BOTH,
     "ops/test_tail1.py": BOTH,
     "ops/test_tail2.py": BOTH,

README.md

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

docs/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 은 입구에서 이름으로 거부됩니다. **90** 개(이 칸은 한때 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)) |

tests/_support/gen_sortties_oracle.py

@@ -0,0 +1,27 @@
+"""Record UPSTREAM torch's tie order for sort/argsort/topk (issue #33).
+
+Run with the spike-venv interpreter and WITHOUT the vendored tree on the path:
+    .caches/spike-venv/bin/python tests/_support/gen_sortties_oracle.py
+Writes tests/ops/sortties_oracle.json.
+"""
+import json
+import os
+import sys
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+sys.path = [p for p in sys.path if "torchnative/python" not in p]
+import torch  # noqa: E402
+import sortties_cells as sc  # noqa: E402
+
+assert not hasattr(torch._C, "_aten_implemented"), "this is the shim, not upstream"
+out = {"torch": torch.__version__, "cells": {}}
+for c in sc.cells():
+    try:
+        out["cells"][c["id"]] = sc.encode(sc.run_cell(torch, c))
+    except RuntimeError as e:
+        out["cells"][c["id"]] = {"error": str(e)}
+dest = os.path.join(HERE, "..", "ops", "sortties_oracle.json")
+with open(dest, "w") as f:
+    json.dump(out, f, separators=(",", ":"))
+print(len(out["cells"]), "cells", torch.__version__)

tests/_support/sortties_cells.py

@@ -0,0 +1,145 @@
+"""The cells of the sort/argsort/topk tie-order oracle (issue #33).
+
+Shared by `gen_sortties_oracle.py` (which runs UPSTREAM torch and records
+the indices) and `tests/ops/test_sortties.py` (which replays the same cells
+on the shim). A cell is a plain dict; its `id` is the oracle's key.
+"""
+
+import hashlib
+import random
+
+NAN = float("nan")
+DTYPES = ["float32", "float16", "bfloat16", "float64", "int32", "int64", "uint8", "bool"]
+FLOATS = {"float32", "float16", "bfloat16", "float64"}
+SIZES = [1, 2, 15, 16, 17, 64, 1000, 5000]
+PATTERNS = ["allsame", "twoval", "ties4", "signedzero", "nan", "distinct"]
+# (rank layout name, shape builder, dim)
+LAYOUTS = {
+    "1d": (lambda n: (n,), -1),
+    "2d_last": (lambda n: (3, n), -1),
+    "2d_dim0": (lambda n: (n, 3), 0),
+    "3d_last": (lambda n: (2, 3, n), -1),
+    "3d_mid": (lambda n: (2, n, 3), 1),
+    "3d_dim0": (lambda n: (n, 2, 3), 0),
+}
+
+
+def pattern_values(pattern, dtype, count, seed):
+    """`count` python numbers. None when the pattern does not exist for the dtype."""
+    rng = random.Random(seed)
+    isf = dtype in FLOATS
+    if pattern == "allsame":
+        return [1] * count
+    if pattern == "twoval":
+        return [rng.choice([0, 1]) for _ in range(count)]
+    if pattern == "ties4":
+        v = [rng.choice([0, 1, 2, 3]) for _ in range(count)]
+        return [x % 2 for x in v] if dtype == "bool" else v
+    if pattern == "signedzero":
+        if not isf:
+            return None
+        return [rng.choice([0.0, -0.0, 0.0, -0.0, 1.0]) for _ in range(count)]
+    if pattern == "nan":
+        if not isf:
+            return None
+        return [rng.choice([NAN, 1.0, 2.0, 1.0, NAN]) for _ in range(count)]
+    if pattern == "distinct":
+        v = list(range(count))
+        rng.shuffle(v)
+        if dtype == "uint8":
+            v = [x % 256 for x in v]
+        if dtype == "bool":
+            v = [x % 2 for x in v]
+        return v
+    raise ValueError(pattern)
+
+
+def _layouts_for(n, dtype):
+    out = ["1d"]
+    if n <= 64 or n == 1000:
+        out += ["2d_last", "2d_dim0"]
+    if n <= 64 and dtype in ("float32", "bfloat16", "int64"):
+        out += ["3d_last", "3d_mid", "3d_dim0"]
+    return out
+
+
+def cells():
+    """Every cell, in a stable order. Each: id, op, dtype, pattern, n, layout, shape, dim, args."""
+    out = []
+    for dtype in DTYPES:
+        for n in SIZES:
+            for layout in _layouts_for(n, dtype):
+                shape = LAYOUTS[layout][0](n)
+                dim = LAYOUTS[layout][1]
+                for pattern in PATTERNS:
+                    if pattern_values(pattern, dtype, 1, 0) is None:
+                        continue
+                    # the wide sweep only where it can show something new
+                    wide = layout == "1d" or dtype in ("float32", "int64")
+                    ops = []
+                    for stable in (None, False, True):
+                        for desc in (False, True):
+                            ops.append(("sort", dict(stable=stable, descending=desc)))
+                    for st in (None, False, True):
+                        for desc in (False, True):
+                            ops.append(("argsort", dict(stable=st, descending=desc)))
+                    ks = sorted({1, max(1, n // 2), n}) if wide else [n]
+                    for k in ks:
+                        for largest in (True, False):
+                            for srt in (True, False):
+                                ops.append(("topk", dict(k=k, largest=largest, sorted=srt)))
+                    for op, a in ops:
+                        tag = ",".join(f"{kk}={vv}" for kk, vv in a.items())
+                        cid = f"{op}|{dtype}|{pattern}|{layout}|n{n}|{tag}"
+                        out.append(dict(id=cid, op=op, dtype=dtype, pattern=pattern,
+                                        n=n, layout=layout, shape=list(shape),
+                                        dim=dim, args=a))
+    return out
+
+
+def make_input(torch, c):
+    numel = 1
+    for s in c["shape"]:
+        numel *= s
+    seed = hash_seed(c["dtype"], c["pattern"], c["layout"], c["n"])
+    vals = pattern_values(c["pattern"], c["dtype"], numel, seed)
+    dt = getattr(torch, c["dtype"])
+    if c["dtype"] == "bool":
+        vals = [bool(v) for v in vals]
+    return torch.tensor(vals, dtype=dt).reshape(c["shape"])
+
+
+def hash_seed(*parts):
+    return int(hashlib.sha1("|".join(map(str, parts)).encode()).hexdigest()[:8], 16)
+
+
+def run_cell(torch, c):
+    """Indices (flat python list) the given torch answers for the cell."""
+    x = make_input(torch, c)
+    a = c["args"]
+    dim = c["dim"]
+    if c["op"] == "sort":
+        if a["stable"] is None:
+            idx = torch.sort(x, dim=dim, descending=a["descending"])[1]
+        else:
+            idx = torch.sort(x, dim=dim, descending=a["descending"], stable=a["stable"])[1]
+    elif c["op"] == "argsort":
+        if a["stable"] is None:
+            idx = torch.argsort(x, dim=dim, descending=a["descending"])
+        else:
+            idx = torch.argsort(x, dim=dim, descending=a["descending"], stable=a["stable"])
+    else:
+        idx = torch.topk(x, a["k"], dim=dim, largest=a["largest"], sorted=a["sorted"])[1]
+    return [int(v) for v in idx.reshape(-1).tolist()]
+
+
+def encode(indices):
+    """Short lists verbatim; long ones as their length and sha1 (every index still compared)."""
+    if len(indices) <= 200:
+        return indices
+    h = hashlib.sha1(",".join(map(str, indices)).encode()).hexdigest()
+    return {"len": len(indices), "sha1": h}
+
+
+def matches(recorded, indices):
+    return encode(indices) == recorded