赛前稳定版

tools/inspect_contraction_tree.py

@@ -0,0 +1,208 @@
+"""Inspect cotengra contraction trees for dominant torch matmul shapes."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import math
+import pickle
+from collections import Counter, defaultdict
+from pathlib import Path
+
+
+def _prod(values):
+    out = 1
+    for value in values:
+        out *= int(value)
+    return out
+
+
+def _broadcast_batch(a_batch, b_batch):
+    if a_batch == b_batch:
+        return _prod(a_batch)
+    if not a_batch:
+        return _prod(b_batch)
+    if not b_batch:
+        return _prod(a_batch)
+
+    ndim = max(len(a_batch), len(b_batch))
+    a_batch = (1,) * (ndim - len(a_batch)) + tuple(a_batch)
+    b_batch = (1,) * (ndim - len(b_batch)) + tuple(b_batch)
+    return _prod(max(a, b) for a, b in zip(a_batch, b_batch))
+
+
+def _load_tree(path, index):
+    with Path(path).open("rb") as f:
+        payload = pickle.load(f)
+    trees = payload["trees"] if isinstance(payload, dict) else payload
+    if not isinstance(trees, (list, tuple)):
+        trees = [trees]
+    return trees[index]
+
+
+def _analyze_tree(tree):
+    contract_mod = importlib.import_module("cotengra.contract")
+    contractions = contract_mod.extract_contractions(tree)
+    size_dict = tree.size_dict
+    ops = []
+    counts = Counter()
+
+    for op_index, (parent, left, right, tdot, arg, perm) in enumerate(contractions):
+        if left is None and right is None:
+            counts["preprocess"] += 1
+            continue
+
+        left_inds = tree.get_inds(left)
+        right_inds = tree.get_inds(right)
+        parent_inds = tree.get_inds(parent)
+        left_shape = tuple(size_dict[ix] for ix in left_inds)
+        right_shape = tuple(size_dict[ix] for ix in right_inds)
+
+        if tdot:
+            parsed = contract_mod._parse_tensordot_axes_to_matmul(
+                arg,
+                left_shape,
+                right_shape,
+            )
+        else:
+            parsed = contract_mod._parse_eq_to_batch_matmul(
+                arg,
+                left_shape,
+                right_shape,
+            )
+
+        (
+            _eq_a,
+            _eq_b,
+            new_shape_a,
+            new_shape_b,
+            _new_shape_ab,
+            _perm_ab,
+            pure_multiplication,
+        ) = parsed
+
+        matmul_shape = None
+        matmul_flops = 0
+        if pure_multiplication:
+            kind = "mul"
+        else:
+            a_shape = tuple(new_shape_a or left_shape)
+            b_shape = tuple(new_shape_b or right_shape)
+            batch = _broadcast_batch(a_shape[:-2], b_shape[:-2])
+            m, k, n = int(a_shape[-2]), int(a_shape[-1]), int(b_shape[-1])
+            kind = "mm" if batch == 1 else "bmm"
+            matmul_shape = (batch, m, k, n)
+            matmul_flops = batch * m * k * n
+
+        tree_flops = int(tree.get_flops(parent))
+        out_size = int(tree.get_size(parent))
+        ops.append(
+            {
+                "index": op_index,
+                "kind": kind,
+                "matmul_shape": matmul_shape,
+                "matmul_flops": matmul_flops,
+                "tree_flops": tree_flops,
+                "out_size": out_size,
+                "left_shape": left_shape,
+                "right_shape": right_shape,
+                "left_rank": len(left_inds),
+                "right_rank": len(right_inds),
+                "out_rank": len(parent_inds),
+                "perm": perm,
+            }
+        )
+        counts[kind] += 1
+
+    return contractions, ops, counts
+
+
+def _format_log(value, base):
+    return "-inf" if value <= 0 else f"{math.log(value, base):.3f}"
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("tree", help="Pickle file containing one tree or {'trees': [...]}.")
+    parser.add_argument("--index", type=int, default=0, help="Tree index in the file.")
+    parser.add_argument("--top", type=int, default=20, help="Number of top ops to print.")
+    parser.add_argument(
+        "--dtype-bytes",
+        type=int,
+        default=8,
+        help="Bytes per element for memory estimates, for example 8 for complex64.",
+    )
+    args = parser.parse_args()
+
+    tree = _load_tree(args.tree, args.index)
+    contractions, ops, counts = _analyze_tree(tree)
+    nslices = int(getattr(tree, "multiplicity", 1))
+    per_slice_flops = sum(op["tree_flops"] for op in ops)
+    per_slice_write = sum(op["out_size"] for op in ops)
+    max_out = max((op["out_size"] for op in ops), default=0)
+    all_flops = per_slice_flops * nslices
+    all_write = per_slice_write * nslices
+
+    print(f"tree={args.tree} index={args.index}")
+    print(
+        "summary "
+        f"slices={nslices} contractions={len(contractions)} "
+        f"counts={dict(counts)}"
+    )
+    print(
+        "per_slice "
+        f"log10_flops={_format_log(per_slice_flops, 10)} "
+        f"log10_write={_format_log(per_slice_write, 10)} "
+        f"log2_max_output={_format_log(max_out, 2)} "
+        f"max_output_gib={max_out * args.dtype_bytes / 1024**3:.6g}"
+    )
+    print(
+        "all_slices "
+        f"log10_flops={_format_log(all_flops, 10)} "
+        f"log10_write={_format_log(all_write, 10)}"
+    )
+
+    print(f"\ntop_{args.top}_ops_by_flops")
+    for op in sorted(ops, key=lambda item: item["tree_flops"], reverse=True)[: args.top]:
+        print(
+            f"op={op['index']} kind={op['kind']} "
+            f"flops={op['tree_flops']:.6e} out={op['out_size']:.6e} "
+            f"matmul={op['matmul_shape']} "
+            f"ranks=({op['left_rank']},{op['right_rank']}->{op['out_rank']}) "
+            f"lhs={op['left_shape']} rhs={op['right_shape']}"
+        )
+
+    by_shape = defaultdict(lambda: [0, 0, 0])
+    for op in ops:
+        shape = op["matmul_shape"]
+        if shape is None:
+            continue
+        by_shape[shape][0] += 1
+        by_shape[shape][1] += op["tree_flops"]
+        by_shape[shape][2] += op["out_size"]
+
+    print(f"\ntop_{args.top}_matmul_shapes_by_flops")
+    for shape, (count, flops, out_size) in sorted(
+        by_shape.items(),
+        key=lambda item: item[1][1],
+        reverse=True,
+    )[: args.top]:
+        print(
+            f"shape={shape} count={count} "
+            f"flops={flops:.6e} output={out_size:.6e}"
+        )
+
+    print(f"\ntop_{args.top}_matmul_shapes_by_count")
+    for shape, (count, flops, out_size) in sorted(
+        by_shape.items(),
+        key=lambda item: item[1][0],
+        reverse=True,
+    )[: args.top]:
+        print(
+            f"shape={shape} count={count} "
+            f"flops={flops:.6e} output={out_size:.6e}"
+        )
+
+
+if __name__ == "__main__":
+    main()