Tracking: kornia-side follow-ups from the MPS audit (device bugs, MPS-slow paths, launch-bound loops, MKD)

Tracking issue — kornia-side follow-ups from the MPS audit

Context: while auditing kornia on Apple MPS for the PyTorch MPS team (torch nightly 2.15.0.dev20260913, Apple M1 8 GB, macOS 26.5.1, PYTORCH_ENABLE_MPS_FALLBACK unset, kornia 633951bd7), every aten call of the CPU test suite was shadow-replayed on MPS (66k signatures) and 292 public ops were timed on CPU and MPS at several shapes (924 device pairs, fresh-process re-measurements, per-aten attribution). The PyTorch-side findings are in #4522. Everything that is kornia's to fix is split into the individual issues below; this issue only tracks them.

One number to keep in mind when reading MPS ratios: one MPS kernel plus torch.mps.synchronize() costs ~270 µs, amortising to ~17 µs per kernel in a chain of 100, versus ~3 µs per op on the CPU. A kornia op made of hundreds of scalar-sized aten calls therefore loses on MPS regardless of image size.

A. Device-placement bugs (fail on MPS and CUDA alike)

• RandomAffine(shear=…), RandomShear, and through them auto.AutoAugment / auto.RandAugment / auto.TrivialAugment — Passed CPU tensor to MPS op from the shear sampler: #4415, open PRs #4507 and #4519 (no new issue).
• RandomLinearIllumination — device mismatch in its generator (_extract_device_dtype on the constructor tuples): #4536 (sibling of #3704, which fixed the Gaussian and corner variants).

B. Drop-in fixes for paths that are slow on MPS (each verified byte-identical or within float32 rounding)

• color.rgb_to_yuv420 — 6-D unfold+mean → avg_pool2d, MPS 88.8 ms → 2.5 ms at 1024²: #4524
• morphology.* — engine="convolution" is 8x faster on MPS with identical output; per-device default or honest docstring: #4525
• losses.total_variation — sum(dim=(-2,-1)) → flatten(-2).sum(-1), MPS 11.7 ms → 3.0 ms: #4526
• enhance.equalize_clahe / RandomClahe — 192 per-tile histc calls → one batched histogram: #4527 (same loop shape as #3728)
• geometry.epipolar.find_essential / RANSAC("essential") — run the 10x10 eigvals on CPU when on MPS (the only missing MPS op kornia hits); empties the #4159 manifest: #4528

C. Launch-bound implementations (hundreds of scalar-sized aten calls per invocation)

• utils.draw_line — Python loops over 0-d tensors and .item() per pixel, 627 ms per line at 1024² on MPS: #4529 (batched draw_lines was #2351)
• augmentation.RandomRain — per-image and per-row loops, 204 add + 100 index_put_ at batch 8: #4530
• augmentation.RandomCrop (+ RandomJigsaw, RandomRGBShift, RandomChannelShuffle, Normalize as siblings) — ~100 scalar ops per call, slicing fast path: #4531 (history: #2175, #3842)
• losses.HausdorffERLoss — ~12 launches × k iterations, boolean index_put where a clamp would do: #4532
• contrib.KMeans — per-cluster nonzero/index_select/mean loop, 160 host syncs per fit: #4533
• feature.ScaleSpaceDetector at 256² (2.4x slower, ≈ 2700 tiny ops) and enhance.jpeg_codec_differentiable at 256² (2.9x): no issue filed — both win on MPS at 1024² / batch 8 (0.3–0.5x), and #4254 already covers the detector pipeline.

D. Device-independent performance

• feature.MKDDescriptor — 1.4 s for 2048 patches on CPU and MPS alike; 75 % in a broadcast emb1 * emb2 + sum(dim=(2,3)) that is a contraction: #4534

E. Withdrawn

• HardNet8 / SmallSR orthogonal init in __init__ (#4535) — not a bug: pretrained=True overriding every initialised weight is the contract. The MPS linalg.qr behaviour it would have exposed stays on the PyTorch side (#4522 §1).

Tooling

The measurements above come from mps_cpu_sweep.py (public-API timing), aten_attribution.py (replays each aten call of one invocation on both devices) and repro_candidates.py; the standalone kernel repros are inline in #4522. A minimal counter of the aten calls one kornia op issues, for anyone checking a fix:

import collections, torch
from torch.utils._python_dispatch import TorchDispatchMode

class Counter(TorchDispatchMode):
    def __init__(self):
        super().__init__(); self.calls = collections.Counter()
    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        self.calls[str(func)] += 1
        return func(*args, **(kwargs or {}))

def count(fn):
    fn()
    c = Counter()
    with c:
        fn()
    print(sum(c.calls.values()), "aten calls")
    for op, n in c.calls.most_common(12):
        print(f"{n:6d}  {op}")

Environment

kornia 633951bd7 (main 060d84614 + 2 test-only commits) | torch 2.15.0.dev20260913 | python 3.11.14 | macOS-26.5.1-arm64 | MPS=True | PYTORCH_ENABLE_MPS_FALLBACK unset

Posted on behalf of @ducha-aiki by Claude (Fable 5.1).