Research/bf16 phase0 review

.gitattributes

@@ -1 +1,23 @@
 * text=auto
+
+# Phase 0 reproducibility artifacts: force canonical LF everywhere under
+# results/phase0 so the raw-byte sha256 hashes recorded in manifest.json inputs,
+# the numerical_validation case_binding, and the c2_checkpoint_manifest
+# artifact_hashes stay stable across EVERY checkout (Task 12). Without this, a
+# fresh Windows checkout smudges text to CRLF and the re-hash mismatches ->
+# NUMERICAL/C2 demoted to UNKNOWN. gitattributes has NO {a,b} brace expansion ->
+# one line per extension. (** matches zero-or-more dirs -> top-level + subdirs.)
+results/phase0/**/*.csv eol=lf
+results/phase0/**/*.json eol=lf
+results/phase0/**/*.hlo eol=lf
+results/phase0/**/*.txt eol=lf
+results/phase0/**/*.md eol=lf
+# Phase 0 producer scripts: pin LF so OneDrive sync on Windows never smudges
+# the working copy to CRLF (which would cause content-less M phantoms after a
+# suite run).  The csv/json/hlo/txt/md patterns above cover the artifacts; this
+# covers the .py generators that Black and the test suite read/write.
+results/_phase0/*.py eol=lf
+
+# F7: raw XLA/compiler dump evidence is exempt from blank-at-eof style
+# check (trailing newline is part of the captured output, not a style error).
+results/phase0/c1_optimized_hlo/**/*.hlo whitespace=-blank-at-eof

CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ### Added
 
+- add unified ContractionAlgebra interface (`set_contraction_algebra(...)` + boundary
+  encode/decode) with tropical (max-plus) and complex<bfloat16> pair-algebra
+  reference applications.
+
 - Add efficient `expectation_pss` method for `U1Circuit`.
 
 - Support MVP mode for timeevol methods.

applications/bcomplex32_algebra.py

@@ -0,0 +1,210 @@
+"""complex<bfloat16> pair-algebra — a reference APPLICATION of ContractionAlgebra.
+
+Pair repr: complex tensor split into PairTensor(re, im) of bf16. Contraction = 4 real
+bf16 matmuls (4M). Activated via ``cons.set_contraction_algebra(ComplexPairAlgebra())`` or
+the ``bcomplex32()`` CM. PairTensor keeps the pair axis off tn.Node (dodges the axis==edge wall).
+"""
+
+from typing import Any, Dict, Iterator, List, Tuple
+import contextlib
+
+import numpy as np
+
+import tensorcircuit.cons as cons
+from tensorcircuit.contraction_algebra import (
+    ContractionAlgebra,
+    PairTensor,
+    Representation,
+)
+
+Tensor = Any
+Backend = Any
+
+
+def _bf16_dtype(be: Backend) -> Any:
+    """Backend-native bf16 dtype. numpy needs ml_dtypes; GPU backends use native bf16
+    (their cast rejects ml_dtypes.bfloat16 — see pytorch_backend.cast)."""
+    name = getattr(be, "name", None)
+    if name == "numpy":
+        import ml_dtypes
+
+        return ml_dtypes.bfloat16
+    if name == "jax":
+        import jax.numpy as jnp
+
+        return jnp.bfloat16
+    if name == "pytorch":
+        import torch
+
+        return torch.bfloat16
+    if name == "tensorflow":
+        import tensorflow as tf
+
+        return tf.bfloat16
+    if name == "cupy":
+        # cupy has no ``cupy.bfloat16``; it accepts ml_dtypes.bfloat16 as a native
+        # dtype (astype + cuBLAS bf16 GEMM both work). Same resolution as numpy.
+        import ml_dtypes
+
+        return ml_dtypes.bfloat16
+    raise NotImplementedError(f"bf16 dtype unknown for backend {name!r}")
+
+
+def _complex_to_pair(be: Backend, t: Tensor) -> PairTensor:
+    """complex tensor -> PairTensor(re, im) of bf16."""
+    bf = _bf16_dtype(be)
+    re = be.cast(be.real(t), bf)
+    im = be.cast(be.imag(t), bf)
+    return PairTensor(re, im)
+
+
+def _pair_to_complex(be: Backend, pair: PairTensor) -> Tensor:
+    """PairTensor of bf16 -> complex tensor (recombine; no copy risk via cast)."""
+    re, im = pair.unpack()
+    re = be.cast(re, cons.rdtypestr)
+    im = be.cast(im, cons.rdtypestr)
+    # be.i() returns 1j as a backend tensor. tensorflow's eager mode cannot build a
+    # complex tensor from a Python ``1j`` applied to a float tensor (TypeError:
+    # "Cannot convert 1j to EagerTensor of dtype float"), unlike numpy/jax/torch/cupy.
+    # Found during L3 GPU testing on TF.
+    return be.cast(re, cons.dtypestr) + be.i(cons.dtypestr) * be.cast(im, cons.dtypestr)
+
+
+def _pair_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor:
+    """Complex tensordot = 4 real bf16 tensordots (4M). Uses be.tensordot (never patched)."""
+    ar, ai = PairTensor.unpack_pair(a)
+    br, bi = PairTensor.unpack_pair(b)
+    cr = be.tensordot(ar, br, axes) - be.tensordot(ai, bi, axes)
+    ci = be.tensordot(ar, bi, axes) + be.tensordot(ai, br, axes)
+    return PairTensor.pack_result(be, cr, ci, not isinstance(a, PairTensor))
+
+
+# Strategy note: on numpy, bf16 single-operand einsum manually decomposes into
+# diagonal → sum → transpose (staying in bf16 end-to-end) because numpy's
+# ``np.einsum`` rejects bfloat16 dtypes. GPU backends (jax/torch/tf/cupy) accept
+# bf16 natively and take ``be.einsum`` directly — see ``_einsum_single_operand_half``.
+# This is different from the tropical algebra's ``_tropical_einsum`` single-operand
+# path, which delegates to ``be.einsum`` for repeated-index resolution — the tropical
+# backend (float64) has no such dtype restriction. Both produce equivalent results
+# under their respective semirings, but the implementation strategy is dictated by
+# dtype constraints rather than algebraic differences.
+def _einsum_single_operand_half_numpy(
+    be: Backend, x: Tensor, lhs: str, out_subs: str
+) -> Tensor:
+    """numpy-only fallback: decompose 1-operand einsum into diagonal + sum + transpose
+    because numpy's einsum rejects bf16. Stays bf16 end-to-end."""
+    x_subs = list(lhs)
+    while True:
+        dup = next((c for c in set(x_subs) if x_subs.count(c) > 1), None)
+        if dup is None:
+            break
+        pos = [i for i, c in enumerate(x_subs) if c == dup]
+        x = np.diagonal(x, axis1=pos[0], axis2=pos[-1])
+        x_subs = [c for i, c in enumerate(x_subs) if i != pos[-1]] + [dup]
+    out_set = set(out_subs)
+    sum_indices = [c for c in x_subs if c not in out_set]
+    if sum_indices:
+        x = be.sum(x, axis=tuple(x_subs.index(c) for c in sum_indices))
+        x_subs = [c for c in x_subs if c in out_set]
+    if x_subs != list(out_subs):
+        perm = tuple(x_subs.index(c) for c in out_subs)
+        x = be.transpose(x, perm)
+    return x
+
+
+def _einsum_single_operand_half(
+    be: Backend, x: Tensor, lhs: str, out_subs: str
+) -> Tensor:
+    """Apply a 1-operand einsum to one bf16 half. GPU backends (jax/torch/tf/cupy)
+    take native ``be.einsum`` (XLA/cuBLAS-fused, accepts bf16). numpy rejects bf16,
+    so it falls back to the manual decomposition. ``np.diagonal`` never runs on GPU."""
+    name = getattr(be, "name", None)
+    if name == "numpy":
+        return _einsum_single_operand_half_numpy(be, x, lhs, out_subs)
+    if name == "cupy":
+        # tensornetwork's cupy backend has not implemented einsum (NotImplementedError),
+        # but cupy.einsum itself exists and runs bf16 on cuBLAS. Call it directly.
+        import cupy
+
+        return cupy.einsum(f"{lhs}->{out_subs}", x)
+    return be.einsum(f"{lhs}->{out_subs}", x)
+
+
+def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> PairTensor:
+    """Complex einsum = 4 real bf16 einsums (4M for 2 operands, 2 for 1)."""
+    if len(operands) == 1:
+        a = operands[0]
+        if "->" not in eq:
+            return a
+        lhs, out_subs = eq.split("->")
+        ar, ai = PairTensor.unpack_pair(a)
+        return PairTensor(
+            _einsum_single_operand_half(be, ar, lhs, out_subs),
+            _einsum_single_operand_half(be, ai, lhs, out_subs),
+        )
+
+    a, b = operands
+    ar, ai = PairTensor.unpack_pair(a)
+    br, bi = PairTensor.unpack_pair(b)
+
+    if "->" not in eq:
+        raise ValueError(
+            f"implicit-mode einsum {eq!r} not supported for bf16; use explicit '->'"
+        )
+    lhs, out_subs = eq.split("->")
+    a_subs, b_subs = lhs.split(",")
+
+    a_set: set[str] = set(a_subs)
+    b_set: set[str] = set(b_subs)
+    contracted = [c for c in a_subs if c in b_set]
+
+    a_free = [c for c in a_subs if c not in b_set]
+    b_free = [c for c in b_subs if c not in a_set]
+    out_order = list(out_subs)
+
+    def _contract(x: Tensor, y: Tensor) -> Tensor:
+        if contracted:
+            a_axes = [a_subs.index(c) for c in contracted]
+            b_axes = [b_subs.index(c) for c in contracted]
+            result = be.tensordot(x, y, axes=(a_axes, b_axes))
+        else:
+            result = be.tensordot(x, y, axes=0)
+        free_order = a_free + b_free
+        if free_order != out_order:
+            perm = [free_order.index(c) for c in out_order]
+            result = be.transpose(result, perm)
+        return result
+
+    cr = _contract(ar, br) - _contract(ai, bi)
+    ci = _contract(ar, bi) + _contract(ai, br)
+    return PairTensor.pack_result(be, cr, ci, not isinstance(a, PairTensor))
+
+
+class PairBf16Representation(Representation):
+    name = "pair_bf16"
+
+    def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]:
+        return [_complex_to_pair(be, t) for t in tensors]
+
+    def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]:
+        return _pair_to_complex(be, tensor), {}
+
+
+class ComplexPairAlgebra(ContractionAlgebra):
+    name = "bcomplex32_pair"
+    representation = PairBf16Representation()
+
+    def get_contractor_kwargs(self) -> Dict[str, Any]:
+        return {"prefer_einsum": True}
+
+    def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor:
+        return _pair_tensordot(be, a, b, axes)
+
+    def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor:
+        return _pair_einsum(be, eq, *operands)
+
+
+@contextlib.contextmanager
+def bcomplex32() -> Iterator[None]:
+    with cons.runtime_contraction_algebra(ComplexPairAlgebra()):
+        yield

applications/benchmarks/aggregate_results.py

@@ -0,0 +1,71 @@
+"""Aggregate bench_bf16_gpu CSV rows into a markdown report.
+
+For each (backend, circuit, n): report bf16 vs complex64 peak-mem ratio, speedup,
+and bf16 accuracy. For micro rows: report wall-time per backend.
+"""
+
+import argparse
+import csv
+from collections import defaultdict
+
+
+def _gib(b: float) -> str:
+    return f"{b / (1024 ** 3):.2f} GiB" if b else "-"
+
+
+def aggregate(csv_path: str) -> str:
+    e2e = defaultdict(dict)  # (backend, circuit, n) -> {dtype: row}
+    micro = []
+    with open(csv_path, newline="") as f:
+        for r in csv.DictReader(f):
+            if r.get("mode") == "micro":
+                micro.append(r)
+            else:
+                e2e[(r["backend"], r["circuit"], int(r["n"]))][r["dtype"]] = r
+
+    lines = ["# bf16 GPU benchmark results", ""]
+    lines += ["## End-to-end (bf16 vs complex64)", ""]
+    lines += [
+        "| backend | circuit | n | c64 mem | bf16 mem | mem ratio | c64 s | bf16 s | speedup | bf16 max-abs-err |"
+    ]
+    lines += ["|---|---|---|---|---|---|---|---|---|---|"]
+    for (backend, circuit, n), d in sorted(e2e.items()):
+        c64, bf = d.get("complex64"), d.get("bf16")
+        if not (c64 and bf):
+            continue
+        c64_mem = int(c64["peak_smi_bytes"] or 0)
+        bf_mem = int(bf["peak_smi_bytes"] or 0)
+        ratio = f"{c64_mem / bf_mem:.2f}x" if bf_mem else "-"
+        c64_s = float(c64["wall_s"])
+        bf_s = float(bf["wall_s"])
+        speedup = f"{c64_s / bf_s:.2f}x" if bf_s else "-"
+        err = bf.get("max_abs_err") or "-"
+        lines.append(
+            f"| {backend} | {circuit} | {n} | {_gib(c64_mem)} | {_gib(bf_mem)} | {ratio} "
+            f"| {c64_s:.2f} | {bf_s:.2f} | {speedup} | {err} |"
+        )
+
+    if micro:
+        lines += ["", "## Micro (4M bf16 GEMM, single matmul)", ""]
+        lines += ["| backend | m | wall s |", "|---|---|---|"]
+        for r in sorted(micro, key=lambda x: (x["backend"], int(x["n"]))):
+            lines.append(f"| {r['backend']} | {r['n']} | {float(r['wall_s']):.3f} |")
+
+    return "\n".join(lines) + "\n"
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("csv")
+    p.add_argument("-o", "--out", default=None)
+    args = p.parse_args()
+    md = aggregate(args.csv)
+    if args.out:
+        with open(args.out, "w") as f:
+            f.write(md)
+    else:
+        print(md)
+
+
+if __name__ == "__main__":
+    main()