trntensor 0.18.0__tar.gz → 0.19.0__tar.gz

{trntensor-0.18.0 → trntensor-0.19.0}/CHANGELOG.md

@@ -7,6 +7,42 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.19.0] — 2026-10-31
+
+Validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0, native TorchNeuron) — see trnsci/trnsci#35.
+
+Headline: new native MP2-energy and AO→MO kernels beat aten by 1.25–1.9× and 1.3–1.6× in the `auto` region; native NKI matmul/bmm kernels deprecated (#52).
+
+Resolves #51: on the native TorchNeuron path (Trn2/Trn3), the AO→MO and MP2-energy kernels are rewritten for Trn2 and now beat aten, and the NKI matmul/bmm kernels are retired in favour of aten.
+
+### Added
+
+- **`mp2_energy_row_kernel`, the native MP2-energy kernel (#51).**
+  - One launch per occupied row, using the pair symmetry `e_ij = e_ji`.
+  - Per pair, one GEMM fills `T` in SBUF; the `Tᵀ` blocks are Tensor Engine transposes. `T` never reaches device memory.
+  - Gridless launch: at LNC=2 each physical core takes half of the `j` window (`num_programs()` / `program_id`), and it also runs at LNC=1.
+  - The design follows trnblas `pair_energy_sym` (trnblas#55, #58).
+  - Native `"auto"` uses it for device-resident `B` once each row launch carries at least 3 GFLOP. On Trn2 it measured 1.25–1.9× faster than aten: for example 0.69 s against 1.30 s at `nocc=96, nvir=672, naux=2304`.
+  - Near a full device the kernel's input staging (about 1.5× the padded `B`) needs more memory than aten's 64 MiB chunks. `"auto"` keeps aten when live allocations plus that staging would leave less than 3 GiB of the core free, and remembers shapes whose launch failed.
+  - `TRNTENSOR_MP2_ENERGY=auto|rows|aten|legacy` overrides the choice.
+- **`ao_to_mo_trn2_kernel`, the native AO→MO kernel (#51).**
+  - Two plain GEMM stages with the contraction index on the partition axis, so there are no DMA transposes.
+  - The auxiliary index is split across both physical cores at LNC=2.
+  - Native `"auto"` uses it for device-resident `eri` with `nocc ≤ 128` and stage-1 work of at least 5e8 multiply-adds (about nbasis ≥ 256). There it measured 1.3–1.6× faster than the two aten einsums on Trn2, at nbasis 256–768. Smaller problems stay on aten.
+  - `TRNTENSOR_AO_TO_MO=auto|trn2|aten|legacy` overrides the choice.
+- `scripts/bench_native_kernels.py`: a kernel-level A/B harness (aten against the original and the Trn2 kernels), with `--hold-gb` for the memory-pressure regime of aws-neuron-sdk#1447.
+
+### Deprecated
+
+- **On the native path, the NKI `matmul` / `batched_matmul` kernels and the original `mp2_energy` / `ao_to_mo_transform` kernels (#51).** `set_backend("nki")` still runs the matmul/bmm kernels there, with a one-time `FutureWarning`; a later release will route them to aten.
+  - A Trn2 rewrite of matmul used both cores, Tensor Engine transposes and B-tile reuse. It was 1.2–4.3× faster than the original kernel, but still slower than aten at every size (4096³ fp32: 4.10 ms against 3.80 ms; bf16: 1.61 ms against 1.43 ms). It was not kept.
+  - Native `"auto"` already used aten for these.
+  - The XLA path (Trn1) keeps all four original kernels unchanged and does not warn.
+
+### Changed
+
+- Under `set_backend("nki")` on the native path, `mp2_energy` and `ao_to_mo_transform` run the Trn2 kernels instead of the original ones.
+
 ## [0.18.0] — 2026-10-30
 
 Native TorchNeuron dispatch (Trn2/Trn3) alongside the PyTorch/XLA path (Trn1, maintenance); validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0) — see trnsci/trnsci#35.

{trntensor-0.18.0 → trntensor-0.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trntensor
-Version: 0.18.0
+Version: 0.19.0
 Summary: Tensor contractions for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0

{trntensor-0.18.0 → trntensor-0.19.0}/docs/api/nki.md

@@ -118,6 +118,8 @@ supported in maintenance mode."
 | `TRNSCI_NKI_DISPATCH=auto\|native\|xla` | Suite-wide. Force the native or XLA dispatch path; `auto` (default) detects. Forcing an unavailable path raises. |
 | `TRNSCI_XLA_DEPRECATION=1` | Suite-wide, off by default. Emit the one-time XLA-path and `to_xla` / `from_xla` deprecation warnings. |
 | `NKI_ENABLE_TRACE_CACHE=0` | TorchNeuron setting (not trntensor's). Turns off the cross-process NKI kernel cache, which cannot version editable installs and may serve a stale kernel after a source edit. Set it with `pip install -e`; the test suite sets it by default. |
+| `TRNTENSOR_MP2_ENERGY=auto\|rows\|aten\|legacy` | Native path: force the `mp2_energy` executor (`rows` = the Trn2 row kernel, `legacy` = the original kernel, deprecated there). |
+| `TRNTENSOR_AO_TO_MO=auto\|trn2\|aten\|legacy` | Native path: force the `ao_to_mo_transform` executor (`trn2` = the Trn2 kernel, `legacy` = the original kernel, deprecated there). |
 | `TRNTENSOR_FORCE_BACKEND=pytorch\|nki\|auto` | Override the backend selection at runtime without calling `set_backend()`. Used by benchmarks to sweep both paths on the same machine. |
 | `TRNTENSOR_USE_SIMULATOR=1` | Route kernel dispatch through `nki.simulate(kernel)(numpy_args)` on CPU instead of XLA → NEFF → hardware. Catches Python-trace-level errors without AWS round-trips. MLIR verifier errors remain hardware-only. |
 
@@ -146,15 +148,26 @@ NeuronCore and two host threads, with `scripts/bench_native_einsum.py` and
   The default of 500 M sits between these. A host with more threads moves the
   crossover up. Re-measure on your host with `scripts/autotune_dispatch.py`.
 
-Under `"auto"` on the native path, the fused quantum kernels (`mp2_energy`,
-`ao_to_mo_transform`) also run as aten, because both measured slower than
-aten on the device:
-
-- AO→MO was 6–20× slower.
-- MP2 energy was 1.0–1.5× slower, reaching parity at best.
-
-`set_backend("nki")` still runs the NKI kernels. The kernels are unchanged,
-and the XLA path keeps them under `"auto"`.
+Kernel status on the native path after #51:
+
+- **matmul / bmm: retired in favour of aten.** A Trn2 rewrite used both
+  physical cores through the `program_id` split, Tensor Engine transposes of
+  A and B-tile reuse across four stationaries. It was 1.2–4.3× faster than
+  the original kernel and still slower than aten at every size: 4096³ fp32
+  took 4.10 ms against 3.80 ms, and 4096³ bf16 1.61 ms against 1.43 ms.
+  It was not kept. `set_backend("nki")` still runs the original kernels on
+  the native path, with a one-time `FutureWarning`. A later release will
+  route them to aten.
+- **AO→MO: kept, rewritten for Trn2** (`ao_to_mo_trn2_kernel`). It is
+  1.3–1.6× faster than the two aten einsums at nbasis 256–768, and native
+  `"auto"` uses it there for device-resident `eri`. Below nbasis ≈ 256
+  `"auto"` keeps aten.
+- **MP2 energy: kept, rewritten for Trn2** (`mp2_energy_row_kernel`). It is
+  1.25–1.9× faster than aten once each row launch carries at least 3 GFLOP,
+  and native `"auto"` uses it there.
+
+The original AO→MO and MP2 kernels remain the XLA-path (Trn1) kernels. See
+[Quantum-chemistry primitives](quantum.md) for the tables.
 
 ## Kernels (internal)
 

trntensor-0.19.0/docs/api/quantum.md

@@ -0,0 +1,152 @@
+# Quantum-chemistry primitives
+
+Domain-specific fused kernels where the whole computation — contraction, elementwise shaping, reduction — lives in a single NKI program. The architectural core of trntensor (see [Architecture](../architecture.md)).
+
+## `trntensor.mp2_energy(B, eps_occ, eps_vir) -> scalar`
+
+Density-fitted second-order Møller–Plesset correlation energy.
+
+```
+E_MP2 = Σ_{i,j,a,b} T_{i,j,a,b} (2 T_{i,j,a,b} - T_{i,j,b,a}) / Δ_{i,j,a,b}
+
+T_{i,j,a,b}  = Σ_P  B[i, a, P] B[j, b, P]
+Δ_{i,j,a,b}  = ε_i + ε_j - ε_a - ε_b
+```
+
+### Arguments
+
+- `B: (nocc, nvir, naux) tensor` — density-fitted ERI coefficients
+- `eps_occ: (nocc,) tensor` — occupied orbital energies
+- `eps_vir: (nvir,) tensor` — virtual orbital energies
+
+### Returns
+
+A 0-D tensor containing the correlation energy.
+
+### Example
+
+```python
+import torch
+import trntensor
+
+nocc, nvir, naux = 5, 19, 72
+B = torch.randn(nocc, nvir, naux) * 0.1
+eps_occ = -torch.sort(torch.rand(nocc))[0] - 0.5
+eps_vir = torch.sort(torch.rand(nvir))[0] + 0.1
+
+E = trntensor.mp2_energy(B, eps_occ, eps_vir)
+print(f"E_MP2 = {E.item():.6f}")
+```
+
+### Backend behaviour
+
+- **CPU, and small problems on native TorchNeuron**: aten ops on the
+  operands' device. `T[i, j, a, b]` is contracted with `torch.einsum` in
+  chunks over `i`, which bounds peak memory to about 64 MiB of `T` at a time.
+  The `T`-weighted sum is vectorized.
+- **Native TorchNeuron (Trn2/Trn3): the row kernel.** `mp2_energy_row_kernel`
+  is launched once per occupied row `i`. It covers `j = i + d (mod nocc)` for
+  `d = 0..nocc/2` and uses the pair symmetry `e_ij = e_ji`. Per pair, one
+  GEMM fills `T = B[i] B[j]ᵀ` in SBUF, and the `Tᵀ` blocks come from Tensor
+  Engine transposes. `T` never reaches device memory. At LNC=2 each physical
+  core takes half of the `j` window. The design follows trnblas
+  `pair_energy_sym` (trnblas#55, #58).
+  - `"auto"` uses the row kernel when `B` is on the device and each row
+    launch carries at least 3 GFLOP, counted on the padded sizes as
+    `2·(nocc//2+1)·nvir²·naux`. Below that, the ~nocc custom-op launch
+    floors lose to aten.
+  - `set_backend("nki")` uses it whenever its SBUF working set fits.
+  - `TRNTENSOR_MP2_ENERGY=auto|rows|aten|legacy` overrides the choice.
+    `legacy` is the original single-launch kernel, which is deprecated on
+    the native path.
+
+  Measured on Trn2 (fp32, one LNC=2 core, device-resident, warm, #51):
+
+  | nocc, nvir, naux | aten | row kernel |
+  |---|---|---|
+  | 16, 128, 512 | 1.5 ms | 5.1 ms |
+  | 32, 256, 768 | 12.1 ms | 11.6 ms |
+  | 32, 320, 1024 | 24.9 ms | 19.8 ms |
+  | 48, 384, 1152 | 75.7 ms | 39.2 ms |
+  | 64, 448, 1536 | 240 ms | 120 ms |
+  | 96, 672, 2304 | 1.30 s | 0.69 s |
+
+  Holding 12 or 18 GB of the core's 24 GB on the device did not slow either
+  path. The aten path's 64 MiB `T` chunks stay clear of the slowdown in
+  aws-neuron-sdk#1447. The row kernel's input staging is about 1.5× the
+  padded `B`, so near a full device it needs more memory than aten does.
+  `"auto"` therefore keeps aten when this process's live device allocations
+  plus that staging would leave less than 3 GiB free. A shape whose kernel
+  launch fails anyway falls back to aten and is remembered for the rest of
+  the process.
+- **PyTorch/XLA path (Trn1)**: the original single-launch `mp2_energy_kernel`.
+  It accumulates `T` in PSUM, builds `Δ` on the Vector Engine, and writes one
+  partial per `(i, j)` pair.
+
+## `trntensor.ao_to_mo_transform(eri, C_occ, C_vir) -> Tensor`
+
+Fused 4-index AO→MO integral transform. The canonical quantum-chemistry primitive.
+
+```
+B[i, a, P] = Σ_{μ, ν} C_occ[μ, i] · C_vir[ν, a] · eri[μ, ν, P]
+```
+
+### Arguments
+
+- `eri: (nbasis, nbasis, naux) tensor` — density-fitted two-electron integrals in the AO basis
+- `C_occ: (nbasis, nocc) tensor` — occupied molecular orbital coefficients
+- `C_vir: (nbasis, nvir) tensor` — virtual molecular orbital coefficients
+
+### Returns
+
+A `(nocc, nvir, naux)` tensor of transformed DF coefficients B_ia^P.
+
+### Example — full DF-MP2 pipeline
+
+```python
+import trntensor
+
+B = trntensor.ao_to_mo_transform(eri, C_occ, C_vir)
+E = trntensor.mp2_energy(B, eps_occ, eps_vir)
+```
+
+### Backend behaviour
+
+- **CPU**: two explicit aten steps on the operands' device:
+  `einsum("mi,mnP->inP")` then `einsum("na,inP->iaP")`.
+- **Native TorchNeuron (Trn2/Trn3): `ao_to_mo_trn2_kernel`.** Both stages are
+  plain GEMMs with the contraction index on the partition axis, so there are
+  no DMA transposes.
+  - Stage 1: `X[i, ν, P] = Σ_μ C_occ[μ, i] eri[μ, ν, P]`.
+  - Stage 2: `B[i, a, P] = Σ_ν C_vir[ν, a] X[i, ν, P]`.
+  - The auxiliary index is split across both physical cores at LNC=2.
+  - `"auto"` uses the kernel for device-resident `eri` with `nocc ≤ 128` when
+    its SBUF working set fits and the stage-1 work `nocc·nbasis²·naux` is at
+    least 5e8 multiply-adds (about nbasis ≥ 256). Smaller problems lose to
+    the launch floor and the padding copies; otherwise it runs the two aten
+    steps.
+    `TRNTENSOR_AO_TO_MO=auto|trn2|aten|legacy` overrides the choice.
+
+  Measured on Trn2 (fp32, one LNC=2 core, device-resident, warm, #51):
+
+  | nbasis, nocc, naux | two aten einsums | Trn2 kernel (op) | `ao_to_mo_transform`, `"auto"` | original kernel |
+  |---|---|---|---|---|
+  | 64, 8, 128 | 0.41 ms | — | 0.41 ms (aten) | 6.0 ms (#49) |
+  | 128, 16, 256 | 0.49 ms | 0.43 ms | 0.49 ms (aten) | 16.1 ms |
+  | 256, 32, 512 | 1.23 ms | 1.11 ms | 0.92 ms | 112 ms (#49) |
+  | 512, 32, 1024 | 7.30 ms | 4.56 ms | 4.72 ms | — |
+  | 768, 48, 1536 | 18.3 ms | 11.8 ms | 12.0 ms | — |
+
+  The op column takes tile-multiple inputs (the raw kernel). The `"auto"`
+  column includes the padding of `nvir` to 128 and is the shipped route. At
+  nbasis 64 and 128 it measured 0.79 and 0.57 ms with the kernel forced,
+  which is why `"auto"` keeps aten there.
+
+- **PyTorch/XLA path (Trn1)**: the original kernel. It loops over P, with the
+  `(i, ν)` intermediate round-tripping through kernel-scratch HBM. On the
+  native path it is deprecated.
+
+### Current limitations
+
+- Both NKI kernels need `nocc ≤ 128`; larger `nocc` runs the aten steps (native) or raises `NotImplementedError` (XLA).
+- The `(nocc, nbasis, naux)` intermediate goes through kernel-scratch HBM between the two stages, as it does between the two aten einsums. Stage 1 leaves `i` on the partition axis, and stage 2 needs `ν` there.

{trntensor-0.18.0 → trntensor-0.19.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trntensor"
-version = "0.18.0"
+version = "0.19.0"
 description = "Tensor contractions for AWS Trainium via NKI"
 readme = "README.md"
 license = "Apache-2.0"

trntensor-0.19.0/scripts/bench_native_kernels.py

@@ -0,0 +1,199 @@
+"""Kernel-level A/B on TorchNeuron: trntensor NKI kernels vs aten (#51).
+
+Times, on device-resident ``neuron`` tensors created on the host and copied
+over (aws-neuron-sdk#1439), the warm median latency of:
+
+    aten      the aten formulation trntensor's native ``auto`` uses
+    nki_trn1  the original kernel (the XLA-path kernel, as on main)
+    tt_auto   the public trntensor entry point under set_backend("auto")
+    nki_trn2  the native Trn2 kernel (program_id split, no DMA transposes);
+              AO→MO and MP2 only. The Trn2 matmul candidate measured in #51
+              lost to aten at every size and was not kept.
+
+Families: matmul, ao2mo, mp2. ``--hold-gb G`` keeps G GB allocated on the
+device during the run (the memory-pressure regime of aws-neuron-sdk#1447).
+
+    python scripts/bench_native_kernels.py --family mp2 --hold-gb 12
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import statistics
+import time
+
+os.environ.setdefault("NKI_ENABLE_TRACE_CACHE", "0")
+
+import torch  # noqa: E402
+
+from trntensor import quantum  # noqa: E402
+from trntensor.nki import dispatch  # noqa: E402
+
+DEV = "neuron"
+
+
+def _sync():
+    torch.neuron.synchronize()
+
+
+def _time(fn, warm, reps, budget_s=20.0):
+    for _ in range(warm):
+        fn()
+    _sync()
+    t0 = time.perf_counter()
+    fn()
+    _sync()
+    one = time.perf_counter() - t0
+    if one > 0.05:
+        reps = max(3, min(reps, int(budget_s / one)))
+    ts = []
+    for _ in range(reps):
+        t0 = time.perf_counter()
+        fn()
+        _sync()
+        ts.append(time.perf_counter() - t0)
+    return statistics.median(ts) * 1e3
+
+
+def _relerr(x, ref):
+    x = x.detach().cpu().double()
+    ref = ref.double()
+    return float((x - ref).norm() / ref.norm().clamp_min(1e-300))
+
+
+def _ops():
+    return dispatch._get_native_ops()
+
+
+def cases_matmul(dtype, sizes):
+    sizes = sizes or [(n, n, n) for n in (512, 1024, 2048, 4096)] + [(512, 2048, 512)]
+    for M, K, N in sizes:
+
+        def make(M=M, K=K, N=N):
+            g = torch.Generator().manual_seed(0)
+            a = torch.randn(M, K, generator=g).to(dtype)
+            b = torch.randn(K, N, generator=g).to(dtype)
+            return (a, b), (a.double() @ b.double())
+
+        yield (
+            f"matmul {M}x{K}x{N} {str(dtype)[6:]}",
+            make,
+            {
+                "aten": lambda a, b: torch.matmul(a, b),
+                "nki_trn1": lambda a, b: _ops()["matmul"](a, b),
+            },
+        )
+
+
+def cases_ao2mo(dtype, sizes):
+    # (nbasis, nocc, naux), nvir = nbasis - nocc; nbasis/nvir/naux tile-padded below.
+    sizes = sizes or [(128, 16, 256), (256, 32, 512), (512, 32, 1024), (768, 48, 1536)]
+    for nb, no, nx in sizes:
+        nv = nb - no
+
+        def make(nb=nb, no=no, nv=nv, nx=nx):
+            g = torch.Generator().manual_seed(0)
+            co = torch.randn(nb, no, generator=g).to(dtype)
+            cv = torch.randn(nb, nv, generator=g).to(dtype)
+            eri = torch.randn(nb, nb, nx, generator=g).to(dtype)
+            ref = torch.einsum(
+                "na,inP->iaP", cv.double(), torch.einsum("mi,mnP->inP", co.double(), eri.double())
+            )
+            return (eri, co, cv), ref
+
+        modes = {
+            "aten": lambda eri, co, cv: quantum._cpu_ao_to_mo_transform(eri, co, cv),
+            "nki_trn2": lambda eri, co, cv: dispatch._nki_ao_to_mo_trn2(eri, co, cv),
+            "tt_auto": lambda eri, co, cv: quantum.ao_to_mo_transform(eri, co, cv),
+        }
+        if nb <= 128:  # the original kernel compiles for many minutes beyond this
+            modes["nki_trn1"] = lambda eri, co, cv: dispatch._nki_ao_to_mo_transform(eri, co, cv)
+        yield (f"ao2mo nb={nb} no={no} nv={nv} nx={nx}", make, modes)
+
+
+def cases_mp2(dtype, sizes):
+    sizes = sizes or [
+        (16, 128, 512),
+        (32, 128, 512),
+        (32, 256, 768),
+        (48, 384, 1152),
+        (64, 448, 1536),
+        (96, 672, 2304),
+    ]
+    for no, nv, nx in sizes:
+
+        def make(no=no, nv=nv, nx=nx):
+            g = torch.Generator().manual_seed(0)
+            B = (torch.randn(no, nv, nx, generator=g) * 0.05).to(dtype)
+            eo = -torch.rand(no, generator=g) - 1.0
+            ev = torch.rand(nv, generator=g) + 0.5
+            ref = quantum._cpu_mp2_energy(B.double(), eo.double(), ev.double())
+            return (B, eo, ev), ref
+
+        modes = {
+            "aten": lambda B, eo, ev: quantum._cpu_mp2_energy(B, eo, ev),
+            "nki_trn2": lambda B, eo, ev: dispatch._nki_mp2_energy_rows(B, eo, ev),
+            "tt_auto": lambda B, eo, ev: quantum.mp2_energy(B, eo, ev),
+        }
+        if no <= 32:
+            modes["nki_trn1"] = lambda B, eo, ev: dispatch._nki_mp2_energy(B, eo, ev)
+        yield (f"mp2 no={no} nv={nv} nx={nx}", make, modes)
+
+
+FAMILIES = {"matmul": cases_matmul, "ao2mo": cases_ao2mo, "mp2": cases_mp2}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--family", required=True, choices=list(FAMILIES))
+    ap.add_argument("--dtype", default="float32")
+    ap.add_argument("--modes", default="aten,nki_trn1,nki_trn2")
+    ap.add_argument("--sizes", default="", help="semicolon list of comma tuples")
+    ap.add_argument("--hold-gb", type=float, default=0.0)
+    ap.add_argument("--warm", type=int, default=2)
+    ap.add_argument("--reps", type=int, default=10)
+    ap.add_argument("--json", default="")
+    args = ap.parse_args()
+    dtype = getattr(torch, args.dtype)
+    sizes = [tuple(int(x) for x in s.split(",")) for s in args.sizes.split(";") if s]
+    want = args.modes.split(",")
+
+    held = []
+    if args.hold_gb > 0:
+        # Hold the working set in 1 GiB host-made blocks (device randn is
+        # shape-sensitive, #1439; zeros are enough to occupy the memory).
+        for _ in range(int(args.hold_gb)):
+            held.append(torch.ones(256 * 1024 * 1024, dtype=torch.float32).to(DEV))
+        _sync()
+    print(f"# family={args.family} dtype={args.dtype} hold_gb={args.hold_gb}", flush=True)
+    for label, make, modes in FAMILIES[args.family](dtype, sizes):
+        host, ref = make()
+        dev = tuple(t.to(DEV) for t in host)
+        _sync()
+        row = {"case": label, "hold_gb": args.hold_gb}
+        for m in want:
+            if m not in modes:
+                continue
+            fn = lambda f=modes[m], o=dev: f(*o)  # noqa: E731
+            try:
+                t0 = time.perf_counter()
+                out = fn()
+                _sync()
+                row[f"{m}_first_s"] = round(time.perf_counter() - t0, 2)
+                row[f"{m}_relerr"] = _relerr(out, ref)
+                del out
+                row[m] = _time(fn, args.warm, args.reps)
+            except Exception as exc:  # noqa: BLE001
+                row[m] = None
+                row[f"{m}_err"] = repr(exc)[:400]
+        del dev
+        print(json.dumps(row), flush=True)
+        if args.json:
+            with open(args.json, "a") as f:
+                f.write(json.dumps(row) + "\n")
+
+
+if __name__ == "__main__":
+    main()