33. A Dynamic Tensor Compiler with Fusion and Planned Memory

.venv/bin/python -m tensorjit
.venv/bin/python -m pytest tests/test_tensorjit.py -q

"""Immutable tensor expressions with bounded symbolic dimensions.

Supported operations are float32 and pure: add/mul/ReLU, transpose, reduction,
and rank-two matrix multiplication. Symbols identify equal logical dimensions.
"""
from __future__ import annotations
from dataclasses import dataclass
import math


@dataclass(frozen=True)
class Dim:
  name: str
  lower: int = 1
  upper: int = 4096

  def __post_init__(self):
    if not isinstance(self.name, str) or not self.name.isidentifier() or type(self.lower) is not int or type(self.upper) is not int or self.lower < 1 or self.upper < self.lower or self.upper >= 2**31:
      raise ValueError("invalid dimension name or bounds")


@dataclass(frozen=True, eq=False)
class Expr:
  op: str
  shape: tuple[int | Dim, ...]
  sources: tuple[Expr, ...] = ()
  arg: object = None

  def __add__(self, other): return binary("add", self, other)
  def __mul__(self, other): return binary("mul", self, other)
  def __neg__(self): return binary("mul", -1, self)
  def __sub__(self, other): return binary("add", self, -(other if isinstance(other, Expr) else constant(other)))
  def __rsub__(self, other): return binary("add", other, -self)
  def __radd__(self, other): return binary("add", other, self)
  def __rmul__(self, other): return binary("mul", other, self)
  def relu(self): return Expr("relu", self.shape, (self,))

  def expand_dims(self, axis):
    if not -len(self.shape)-1 <= axis <= len(self.shape): raise ValueError("invalid inserted axis")
    axis %= len(self.shape)+1
    return Expr("expand_dims", self.shape[:axis]+(1,)+self.shape[axis:], (self,), axis)

  def broadcast_to(self, shape):
    shape = tuple(shape)
    if any(not isinstance(d, Dim) and (type(d) is not int or d < 1 or d >= 2**31) for d in shape):
      raise ValueError("invalid target shape")
    if broadcast(self.shape, shape) != shape: raise ValueError("cannot broadcast to target shape")
    return Expr("broadcast", shape, (self,))

  def transpose(self):
    if len(self.shape) != 2: raise ValueError("transpose requires a matrix")
    return Expr("transpose", self.shape[::-1], (self,))

  def sum(self, axis=-1):
    if not self.shape or not -len(self.shape) <= axis < len(self.shape): raise ValueError("invalid reduction axis")
    axis %= len(self.shape)
    return Expr("sum", self.shape[:axis] + self.shape[axis+1:], (self,), axis)

  def __matmul__(self, other):
    if not isinstance(other, Expr) or len(self.shape) != 2 or len(other.shape) != 2 or self.shape[1] != other.shape[0]:
      raise ValueError("matmul requires rank-two operands with equal reduction dimensions")
    return Expr("matmul", (self.shape[0], other.shape[1]), (self, other))


def input(name, shape):
  shape = tuple(shape)
  if not name.isidentifier() or any(not isinstance(d, Dim) and (type(d) is not int or d < 1 or d >= 2**31) for d in shape):
    raise ValueError("invalid input name or shape")
  return Expr("input", shape, arg=name)


def constant(value):
  if type(value) not in (int, float) or not math.isfinite(value): raise ValueError("constant must be finite")
  if abs(value) > 3.4028234663852886e38: raise ValueError("constant outside finite float32 range")
  return Expr("constant", (), arg=float(value))


def broadcast(a, b):
  output = []
  for i in range(1, max(len(a), len(b))+1):
    left, right = a[-i] if i <= len(a) else 1, b[-i] if i <= len(b) else 1
    if left == right: output.append(left)
    elif left == 1: output.append(right)
    elif right == 1: output.append(left)
    else: raise ValueError(f"incompatible broadcast dimensions: {left}, {right}")
  return tuple(reversed(output))


def binary(op, left, right):
  left = left if isinstance(left, Expr) else constant(left)
  right = right if isinstance(right, Expr) else constant(right)
  return Expr(op, broadcast(left.shape, right.shape), (left, right))


def toposort(outputs):
  seen, order = set(), []
  pending = [(node, False) for node in reversed(outputs)]
  while pending:
    node, expanded = pending.pop()
    if node in seen: continue
    if expanded:
      seen.add(node); order.append(node)
    else:
      pending.append((node, True))
      pending.extend((source, False) for source in reversed(node.sources))
  return order

"""Fuse a tensor DAG and render runtime-sized C loops.

No fast-math, SIMD schedule search, GPU lowering, or advanced GEMM optimization is
implied. C uses int64 address arithmetic; checked capacities bound every allocation.
"""
from collections import Counter
from dataclasses import dataclass
from tensorjit.graph import Dim, Expr, toposort


@dataclass(frozen=True)
class Kernel:
  node: Expr
  dependencies: tuple[Expr, ...]
  source: str


@dataclass
class Plan:
  inputs: tuple[Expr, ...]
  outputs: tuple[Expr, ...]
  kernels: tuple[Kernel, ...]
  symbols: tuple[Dim, ...]
  names: dict[Expr, str]
  source: str


def compile_plan(outputs, fuse=True):
  outputs = tuple(outputs)
  if not outputs or any(not isinstance(x, Expr) for x in outputs): raise ValueError("nonempty tensor outputs required")
  if len(set(outputs)) != len(outputs): raise ValueError("duplicate outputs are unsupported")
  order = toposort(outputs)
  inputs = tuple(node for node in order if node.op == "input")
  if len({node.arg for node in inputs}) != len(inputs): raise ValueError("input names must be unique")
  symbols = {}
  for node in order:
    for dim in node.shape:
      if isinstance(dim, Dim):
        if dim.name in symbols and symbols[dim.name] != dim: raise ValueError("same symbol name has inconsistent bounds")
        symbols[dim.name] = dim
  symbols = tuple(symbols[name] for name in sorted(symbols))
  dimension_names = {dim: f"dim_{i}" for i, dim in enumerate(symbols)}
  counts = Counter(source for node in order for source in node.sources)
  stages = {node for node in order if node.op not in ("input", "constant") and
            (not fuse or node.op in ("sum", "matmul") or counts[node] > 1 or node in outputs)}
  stages.update(node for node in outputs if node.op in ("input", "constant"))
  names = {node: f"buf_{i}" for i, node in enumerate(order)}
  def dim_code(dim): return dimension_names[dim] if isinstance(dim, Dim) else str(dim)
  def offset(shape, indices):
    if not shape: return "0"
    expression = f"(int64_t)({indices[0]})"
    for dim, index in zip(shape[1:], indices[1:]): expression = f"(({expression})*({dim_code(dim)})+({index}))"
    return expression
  def map_indices(node, indices):
    return ["0" if d == 1 else i for d, i in zip(node.shape, indices[len(indices)-len(node.shape):])]

  kernels = []
  for stage in (node for node in order if node in stages):
    dependencies, lines = set(), []
    def expression(node, indices):
      if node.op == "constant":
        # repr always includes a decimal/exponent for Python float.
        return f"{node.arg!r}f"
      if node.op == "input" or (node in stages and node is not stage):
        dependencies.add(node)
        return f"{names[node]}[{offset(node.shape, indices)}]"
      if node.op in ("add", "mul"):
        args = [expression(source, map_indices(source, indices)) for source in node.sources]
        return f"({args[0]} {'+' if node.op == 'add' else '*'} {args[1]})"
      if node.op == "relu": return f"fmaxf({expression(node.sources[0], indices)}, 0.0f)"
      if node.op == "transpose": return expression(node.sources[0], indices[::-1])
      if node.op == "expand_dims": return expression(node.sources[0], indices[:node.arg]+indices[node.arg+1:])
      if node.op == "broadcast": return expression(node.sources[0], map_indices(node.sources[0], indices))
      if node.op == "positive": return f"({expression(node.sources[0], indices)} > 0.0f ? 1.0f : 0.0f)"
      raise ValueError(f"cannot inline {node.op}; expected a materialization boundary")
    indices = [f"i_{i}" for i in range(len(stage.shape))]
    depth = 1
    for index, dim in zip(indices, stage.shape):
      lines.append("  "*depth + f"for (int32_t {index} = 0; {index} < {dim_code(dim)}; {index}++) {{")
      depth += 1
    indent = "  "*depth
    if stage.op in ("sum", "matmul"):
      if stage.op == "sum":
        source = stage.sources[0]
        source_indices = list(indices)
        source_indices.insert(stage.arg, "reduction")
        extent = source.shape[stage.arg]
        product = expression(source, source_indices)
      else:
        left, right = stage.sources
        extent = left.shape[1]
        product = f"({expression(left, [indices[0], 'reduction'])} * {expression(right, ['reduction', indices[1]])})"
      lines.extend([indent+"float accumulator = 0.0f;", indent+f"for (int32_t reduction = 0; reduction < {dim_code(extent)}; reduction++) {{",
                    indent+f"  accumulator += {product};", indent+"}"])
      value = "accumulator"
    else:
      value = expression(stage, indices)
    destination = f"out_{names[stage]}"
    lines.append(indent + f"{destination}[{offset(stage.shape, indices)}] = {value};")
    for i in reversed(range(len(stage.shape))): lines.append("  "*(i+1) + "}")
    ordered = tuple(node for node in order if node in dependencies)
    arguments = [f"const float *{names[node]}" for node in ordered] + [f"float *{destination}"]
    arguments.extend(f"int32_t {dimension_names[dim]}" for dim in symbols)
    kernel_name = f"kernel_{len(kernels)}"
    source = f"void {kernel_name}({', '.join(arguments)}) {{\n" + "\n".join(lines) + "\n}"
    kernels.append(Kernel(stage, ordered, source))
  source = "#include <stdint.h>\n#include <math.h>\n" + "\n\n".join(k.source for k in kernels)
  return Plan(inputs, outputs, tuple(kernels), symbols, names, source)

"""Guarded dynamic-shape runtime, aligned memory planning, and native execution."""
import ctypes
import math
import shutil
import subprocess
import threading
from pathlib import Path
from tempfile import TemporaryDirectory
import numpy as np
from labs.memory import Temporary, plan_memory, verify_plan
from tensorjit.compiler import compile_plan
from tensorjit.graph import Dim


class CompiledGraph:
  def __init__(self, outputs, *, fuse=True, max_bytes=256*1024*1024):
    if type(max_bytes) is not int or max_bytes < 1: raise ValueError("max_bytes must be a positive integer")
    self.plan = compile_plan(outputs, fuse=fuse)
    self.max_bytes = max_bytes
    self.directory = TemporaryDirectory(prefix="tensorjit-")
    self.source = self.plan.source
    path = Path(self.directory.name)
    (path/"kernels.c").write_text(self.source)
    compiler = shutil.which("gcc")
    if compiler is None: raise RuntimeError("gcc is required")
    subprocess.run([compiler, "-std=c11", "-shared", "-fPIC", "-O2", "-ffp-contract=off", str(path/"kernels.c"), "-lm", "-o", str(path/"kernels.so")],
                   check=True, capture_output=True, text=True)
    self.library = ctypes.CDLL(str(path/"kernels.so"))
    self.functions = []
    pointer = ctypes.POINTER(ctypes.c_float)
    for i, kernel in enumerate(self.plan.kernels):
      function = getattr(self.library, f"kernel_{i}")
      function.argtypes = [pointer] * (len(kernel.dependencies)+1) + [ctypes.c_int32]*len(self.plan.symbols)
      function.restype = None
      self.functions.append(function)
    self.compilations, self.calls = 1, 0
    self.last_memory = {}
    # Native calls release the GIL. Serialize calls so observable counters describe
    # one completed invocation; the actual arenas are already per-invocation.
    self.lock = threading.Lock()

  def bind(self, arguments, explicit=None):
    expected = {node.arg for node in self.plan.inputs}
    if set(arguments) != expected: raise ValueError(f"expected inputs {sorted(expected)}, received {sorted(arguments)}")
    bindings = dict(explicit or {})
    symbol_table = {dim.name: dim for dim in self.plan.symbols}
    if set(bindings) - set(symbol_table): raise ValueError("unknown explicit dimension")
    for name, value in bindings.items():
      dim = symbol_table[name]
      if type(value) is not int or not dim.lower <= value <= dim.upper:
        raise ValueError(f"{name}: invalid explicit dimension binding")
    total_input_bytes = 0
    for node in self.plan.inputs:
      array = arguments[node.arg]
      if not isinstance(array, np.ndarray) or array.dtype != np.float32 or array.ndim != len(node.shape) or not array.flags.c_contiguous or not array.flags.aligned:
        raise ValueError(f"{node.arg}: expected aligned contiguous float32 with rank {len(node.shape)}")
      total_input_bytes += array.nbytes
      if total_input_bytes > self.max_bytes: raise ValueError("input bytes exceed configured budget")
      for expected_dim, actual in zip(node.shape, array.shape):
        if isinstance(expected_dim, Dim):
          if not expected_dim.lower <= actual <= expected_dim.upper: raise ValueError(f"{expected_dim.name}: dimension outside bounds")
          if expected_dim.name in bindings and bindings[expected_dim.name] != actual: raise ValueError(f"{expected_dim.name}: inconsistent dimension binding")
          bindings[expected_dim.name] = actual
        elif expected_dim != actual: raise ValueError("static dimension mismatch")
    # Output symbols must all be grounded in input metadata, never guessed.
    if set(bindings) != {dim.name for dim in self.plan.symbols}: raise ValueError("unbound output dimension")
    return bindings

  def __call__(self, **arguments):
    with self.lock: return self._run(arguments)

  def run_with_bindings(self, bindings, **arguments):
    """Supply metadata for symbols absent from all data inputs (e.g. constant gradients)."""
    with self.lock: return self._run(arguments, bindings)

  def _run(self, arguments, explicit=None):
    bindings = self.bind(arguments, explicit)
    def shape(node): return tuple(bindings[d.name] if isinstance(d, Dim) else d for d in node.shape)
    def nbytes(node):
      size = 4*math.prod(shape(node))
      if size > min(self.max_bytes, 2**63-1): raise ValueError("tensor allocation exceeds configured budget")
      return size
    # Derive liveness from the actual fused schedule, not original graph edges.
    index = {kernel.node: i for i, kernel in enumerate(self.plan.kernels)}
    retained = set(self.plan.outputs)
    temporaries = []
    for kernel in self.plan.kernels:
      node = kernel.node
      if node in retained: continue
      last = max([index[node]] + [i for i, consumer in enumerate(self.plan.kernels) if node in consumer.dependencies])
      temporaries.append(Temporary(self.plan.names[node], nbytes(node), index[node], last))
    allocations, peak = plan_memory(temporaries)
    verify_plan(temporaries, allocations)
    output_bytes = sum(nbytes(node) for node in self.plan.outputs)
    input_bytes = sum(array.nbytes for array in arguments.values())
    if input_bytes+peak+output_bytes > self.max_bytes: raise ValueError("inputs, workspace, and outputs exceed configured budget")
    arena = np.empty(peak+63, dtype=np.uint8)
    shift = (-arena.ctypes.data) % 64
    arena = arena[shift:shift+peak]
    arrays = {node: arguments[node.arg] for node in self.plan.inputs}
    outputs = {node: np.empty(shape(node), dtype=np.float32) for node in self.plan.outputs}
    for node, array in outputs.items():
      if node.op != "input": arrays[node] = array
    for kernel in self.plan.kernels:
      if kernel.node not in retained:
        allocation = allocations[self.plan.names[kernel.node]]
        arrays[kernel.node] = np.ndarray(shape(kernel.node), dtype=np.float32, buffer=arena, offset=allocation.offset)
    ptr = ctypes.POINTER(ctypes.c_float)
    dimensions = [bindings[d.name] for d in self.plan.symbols]
    for kernel, function in zip(self.plan.kernels, self.functions):
      # An input requested as output is copied, preserving caller ownership.
      destination = outputs[kernel.node] if kernel.node in retained else arrays[kernel.node]
      function(*(arrays[node].ctypes.data_as(ptr) for node in kernel.dependencies), destination.ctypes.data_as(ptr), *dimensions)
    self.calls += 1
    self.last_memory = {"temporary_bytes": peak, "unreused_temporary_bytes": sum(t.nbytes for t in temporaries),
                        "output_bytes": output_bytes, "input_bytes": input_bytes, "allocations": allocations, "bindings": bindings,
                        "kernel_count": len(self.plan.kernels)}
    return tuple(outputs[node] for node in self.plan.outputs)

.venv/bin/python -m examples.bench_tensorjit --sizes 16 32 64 --repeat 10 --output book/outputs/tensorjit.json

"""Measure dynamic tensor compilation, fusion, and workspace reuse.

Run from root: python -m examples.bench_tensorjit --sizes 16 32 64 --repeat 10
Use --output path.json to retain raw samples and the environment. CPU only.
The NumPy reference is an end-to-end expression, not a matched low-level GEMM.
"""
import argparse
import json
import os
import platform
import subprocess
import time
from pathlib import Path
import numpy as np
from tensorjit import CompiledGraph, Dim, input


def samples(call, repeat):
  call()  # Warm up; NumPy may initialize its BLAS thread pool.
  result = []
  for _ in range(repeat):
    start = time.perf_counter()
    call()
    result.append((time.perf_counter()-start)*1000)
  return result


def benchmark(sizes, repeat):
  if repeat < 1 or not sizes or min(sizes) < 1 or max(sizes) > 512:
    raise ValueError("positive repeats and sizes in [1,512] required")
  n = Dim("N", 1, 512)
  x, a, b, c, bias = (input(name, shape) for name, shape in
                       (("x", (n, n)), ("a", (n, n)), ("b", (n, n)), ("c", (n, n)), ("bias", (n,))))
  output = ((((x @ a) @ b) @ c) + bias).relu().sum(1)
  variants, compile_ms = {}, {}
  for fuse in (False, True):
    key = "fused" if fuse else "unfused"
    start = time.perf_counter()
    variants[key] = CompiledGraph((output,), fuse=fuse)
    compile_ms[key] = (time.perf_counter()-start)*1000
  report = {"platform": platform.platform(), "processor": platform.processor(), "numpy": np.__version__,
            "compiler": subprocess.run(["gcc", "--version"], capture_output=True, text=True, check=True).stdout.splitlines()[0],
            "thread_environment": {key: os.environ.get(key) for key in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS")},
            "compile_ms": compile_ms, "cases": []}
  rng = np.random.default_rng(17)
  for size in sizes:
    arrays = {name: rng.normal(size=(size, size)).astype(np.float32)*np.float32(1/np.sqrt(size)) for name in ("x", "a", "b", "c")}
    arrays["bias"] = rng.normal(size=size).astype(np.float32)
    def reference(): return np.fmax(((arrays["x"] @ arrays["a"]) @ arrays["b"]) @ arrays["c"] + arrays["bias"], 0).sum(1)
    expected = reference()
    case = {"shape": [size, size], "numpy_ms": samples(reference, repeat)}
    for key, compiled in variants.items():
      actual, = compiled(**arrays)
      np.testing.assert_allclose(actual, expected, atol=3e-5, rtol=3e-5)
      timing = samples(lambda: compiled(**arrays), repeat)
      case[key] = {"samples_ms": timing, "median_ms": float(np.median(timing)),
                   "kernels": compiled.last_memory["kernel_count"],
                   "temporary_bytes": compiled.last_memory["temporary_bytes"],
                   "unreused_temporary_bytes": compiled.last_memory["unreused_temporary_bytes"],
                   "max_absolute_error": float(np.max(np.abs(actual-expected)))}
    report["cases"].append(case)
  return report


def main():
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("--sizes", type=int, nargs="+", default=[16, 32, 64])
  parser.add_argument("--repeat", type=int, default=10)
  parser.add_argument("--output", type=Path)
  args = parser.parse_args()
  report = benchmark(args.sizes, args.repeat)
  content = json.dumps(report, indent=2) + "\n"
  if args.output: args.output.write_text(content)
  print(content, end="")


if __name__ == "__main__": main()