Strix is a static profiler and timeline simulator for TPU Mosaic/Pallas LLO dumps (post-finalize-llo .txt/.mlir). It parses an LLO dump, runs a dual-clock (VPU + DMA) simulator, analyzes stalls and bottlenecks, and prints a console summary. It can also emit Chrome/Perfetto-compatible trace JSON and Graphviz DOT dataflow graphs.

• Python 3.12+

Core analysis has no third-party dependencies (standard library only).

The optional tpu extra (for the import subcommand) needs JAX, libtpu, protobuf, and google-cloud-storage:

pip install -e ".[tpu]"

# Implicit analyze (backward-compatible)
python -m strix path/to/post-finalize-llo.txt

# Explicit analyze
python -m strix analyze path/to/post-finalize-llo.txt

Both python -m strix and python -m strix.cli are equivalent entry points.

This prints a performance summary to stdout and writes trace.json in the current directory.

Strix provides three subcommands:

python -m strix analyze [OPTIONS] PATH

Examples:

# Basic analysis
python -m strix analyze dump.txt

# Override SSA args and disable trace
python -m strix analyze dump.txt --arg %arg0=128 --trace-output ''

# Generate dataflow graph
python -m strix analyze dump.txt --dataflow-output graph.dot
dot -Tsvg graph.dot -o graph.svg

# Debug: dump parsed tree with limited depth
python -m strix analyze dump.txt --dump-tree --tree-max-depth 5

The console output includes:

• Overview: total instructions, FLOPs, bytes, simulated time, VPU/DMA utilization
• Bottleneck: classification (Compute / Memory / Latency / Balanced) with stall ratio
• Instruction mix: breakdown by operation category

python -m strix analyze-bundles [OPTIONS] PATH

Examples:

# Console table
python -m strix analyze-bundles final_bundles.txt

# Filter to a specific source line
python -m strix analyze-bundles final_bundles.txt --line 42

# Show source code alongside bundle mappings
python -m strix analyze-bundles final_bundles.txt --source-root ./src

# Export as JSON
python -m strix analyze-bundles final_bundles.txt --json bundles.json

python -m strix import [OPTIONS] KERNEL

Requires the [tpu] extras and authenticated GKE access.

Example:

python -m strix import kernels.chunk_kda_fwd \
    --shape 1,2048,4,128,128 \
    --tpu-type v7x \
    --tpu-topology 2x2x1

This deploys a K8s Job to the GKE cluster, runs the kernel with the given shape, collects HLO/LLO/Mosaic IR dumps, and downloads them locally for further analysis with strix analyze.

python scripts/cross_compile.py [OPTIONS]

Cross-compiles the fused MoE kernel on a small TPU slice (e.g. 4-chip v7x) targeting a larger topology (e.g. 2x8x8 = 128 chips, 256 devices). Uses jax.experimental.topologies to create a virtual mesh, then compiles via libtpu to dump HLO/LLO/Mosaic IR — without needing all physical devices.

Requires running on a TPU VM with JAX >= 0.5.x and libtpu.

Example:

python scripts/cross_compile.py --topology 2x8x8

IR dumps are written to <output-dir>/<topology>/{hlo,llo,mosaic}/.

The --trace-output file (default: trace.json) is compatible with Perfetto and Chrome chrome://tracing. The trace has four tracks:

Open the file in Perfetto to inspect simulated compute/DMA timelines, stall regions, and loop nesting structure.

Use --dataflow-output to export a Graphviz DOT file showing the SSA dataflow dependency graph. The graph includes:

• Hardware clusters: VPU and DMA operations grouped by execution unit
• Loop subgraphs: nested loop blocks shown as clusters
• Time-bucket rank constraints: same-T constraints displayed as rank annotations
• Cross-stream edges: dependencies between VPU and DMA highlighted

Render with Graphviz:

dot -Tsvg dataflow.dot -o dataflow.svg

LLO dump (.txt/.mlir)
        │
        ▼
   [Parser]          regex-based → Instruction tree
        │
        ▼
   [Simulator]       dual-clock (VPU/DMA), SSA resolution, loop expansion
        │
        ▼
   [Analyzer]        bottleneck classification, stall ratio, instruction mix
        │
        ├──► Console summary (always)
        ├──► Chrome trace JSON (--trace-output)
        └──► Graphviz DOT (--dataflow-output)

The hardware model defaults to v6e-class TPU specs (918 TFLOPS BF16, 32 GB HBM, 1600 GB/s bandwidth) and can be customized in strix/hardware.py.