How real applications reach a simulated accelerator: graph capture with torch.fx / torch.export and torch.compile backends, out-of-tree PyTorch devices, ONNX Runtime execution providers, MLIR, and Google's HEIR compiler for fully homomorphic encryption.
Concepts used here, and where they are explained. Each links to its series glossary entry: a short explanation, then links to the slides that explain it in depth, in this series, the Local LLM Hosting and Key Publications decks, or the FHE Accelerator Simulators series.
A simulator that only runs hand-written workloads answers the questions its authors thought of. One that runs what users actually write (PyTorch models, ONNX graphs, FHE programs) answers the questions customers will ask. "Real-world applications executed on the simulator" means the simulator sits at the bottom of a compiler toolchain, exactly where the silicon will.
| Approach | How | Strengths | Weaknesses |
|---|---|---|---|
| Ahead-of-time graph | Export the model to a graph (torch.export, ONNX, MLIR), lower it, simulate the operator schedule | Whole-program view; compiler-style optimisation and fusion; no runtime needed | Dynamic control flow and data-dependent shapes are awkward |
| Runtime trace | Run the model (even on fake tensors), intercept every operator, record name and shapes, replay into the simulator | Works on almost any model; captures what really executes | One trace per input shape; no compiler-level transformations |
| Simulator as a device | Register the simulator as a framework backend; ops dispatched to it compute results and accumulate simulated time | Users change one line; functional correctness can be checked; drives software bring-up | Most engineering; speed limited by functional computation |
Most teams build all three over time: traces first (cheap), graphs next (for the compiler), and a device last (for bring-up and customers).
torch.export traces a model into an ExportedProgram: a single FX graph of ATen-level operators with shape metadata on every node. It is the modern, sound replacement for TorchScript as an export format.
import torch
from torch.export import export
ep = export(model, (torch.randn(1, 2048, 4096),))
ep = ep.run_decompositions() # lower to Core ATen: linear becomes permute + mm, etc.
ops = []
for node in ep.graph.nodes:
if node.op == "call_function":
ins = [a.meta["val"].shape for a in node.args if isinstance(a, torch.fx.Node)]
out = node.meta["val"]
ops.append((node.target, ins, getattr(out, "shape", None)))
# e.g. (aten.mm.default, [(2048, 4096), (4096, 14336)], (2048, 14336))
# without run_decompositions() you would see aten.linear with (14336, 4096) weights
schedule = mapper.lower(ops) # your hardware mapping
report = simulator.run(schedule) # your simulatorrun_decompositions() lowers the graph to the smaller Core ATen operator set, which shrinks the number of operators the cost model must support.torch.export.Dim, so one export can cover a range of sequence lengths.torch.compile (TorchDynamo) captures graphs from ordinary Python as it runs and hands each graph to a backend: any callable that takes an FX GraphModule and example inputs and returns a callable. That is a clean hook for a simulator.
def sim_backend(gm: torch.fx.GraphModule, example_inputs):
ops = [(n.target, n.meta["example_value"].shape) # Dynamo's fake-tensor metadata
for n in gm.graph.nodes if n.op == "call_function"]
est = simulator.estimate(ops) # latency, utilisation, hot-spots
log.info("graph of %d ops: %.2f ms simulated", len(ops), est.ms)
return gm.forward # functional result from eager PyTorch
model = torch.compile(model, backend=sim_backend)
model(x) # user code unchangedtorch._dynamo.explain lists them.gm.forward; the simulator then stands in for the device.torch.compile also supports AOTAutograd, so the same hook sees backward graphs for training workloads.Every PyTorch operator passes through the dispatcher. A TorchDispatchMode sees each ATen call with its arguments, which is enough to record an exact operator trace of any model. Combine it with the meta device (tensors with shapes but no storage) and you can trace a 70B model on a laptop, with no weights and no arithmetic.
from torch.utils._python_dispatch import TorchDispatchMode
class OpTrace(TorchDispatchMode):
def __init__(self):
super().__init__(); self.ops = []
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
out = func(*args, **(kwargs or {}))
shapes = [tuple(a.shape) for a in args if isinstance(a, torch.Tensor)]
self.ops.append((str(func), shapes))
return out
with torch.device("meta"):
model = LlamaForCausalLM(config) # 70B parameters, zero bytes allocated
x = torch.zeros(1, 2048, dtype=torch.long)
with OpTrace() as t:
model(x)
simulator.replay(t.ops) # a real op sequence, simulatedChecked with PyTorch 2.14 (CPU): a 70B-sized MLP block traced on the meta device yields aten.mm on (2048, 8192) × (8192, 28672) with nothing allocated. This is the cheapest route to "real models on the simulator", and the trace doubles as a regression input: if a library upgrade changes the operator sequence, the diff shows exactly how.
PyTorch reserves a dispatch key, PrivateUse1, for out-of-tree accelerators. A vendor renames it (for example torch.utils.rename_privateuse1_backend("photon")), registers kernels for it in C++, and users write model.to("photon"). Several production accelerators integrate this way.
at::Tensor photon_mm(const at::Tensor& a, const at::Tensor& b) {
sim::Stats::instance().add("mm", a.sizes(), b.sizes()); // timing model
return photon::functional_mm(a, b); // bit-accurate result (incl. precision effects)
}
TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
m.impl("mm", photon_mm);
}Optical and analogue compute have finite precision and noise. A device-level simulator can apply the hardware's numerical behaviour to real tensors, so users see whether a model's accuracy survives, not only how fast it runs.
Memory allocators, copies, streams and a long tail of operators. Start with a fallback to CPU for unsupported operators and grow coverage, tracking coverage as an engineering metric.
ONNX is a framework-neutral graph format: nodes with standard operator types, versioned opsets, typed tensors. ONNX Runtime (ORT) executes ONNX graphs through pluggable Execution Providers (EPs) such as CPU, CUDA, TensorRT, OpenVINO and QNN.
GetCapability).so = ort.SessionOptions()
so.enable_profiling = True
sess = ort.InferenceSession("model.onnx", so,
providers=["CPUExecutionProvider"])
sess.run(None, {"input": x})
path = sess.end_profiling() # Chrome-trace JSON
Same format as the simulator's traces (deck 06): baseline and simulated timelines open side by side in Perfetto.
import onnx
from onnx import shape_inference
m = shape_inference.infer_shapes(onnx.load("model.onnx"))
shapes = {v.name: [d.dim_value for d in v.type.tensor_type.shape.dim]
for v in [*m.graph.input, *m.graph.value_info, *m.graph.output]}
shapes.update({t.name: list(t.dims) for t in m.graph.initializer}) # weights
for node in m.graph.node:
attrs = {a.name: onnx.helper.get_attribute_value(a) for a in node.attribute}
simulator.add_op(node.op_type, [shapes.get(i) for i in node.input], attrs)MLIR (Multi-Level Intermediate Representation, part of LLVM) is a framework for building compilers out of dialects, families of operations at a given abstraction, connected by progressive lowering passes.
mlir-opt-style drivers run pass pipelines on textual IR, so every stage is inspectable and testable with FileCheck-based lit tests.FHE lets a server compute on encrypted data without decrypting it. The client encrypts, the server evaluates a function on ciphertexts, and only the client can decrypt the result. It is the strongest form of confidential computing, and very expensive.
HEIR (Homomorphic Encryption Intermediate Representation) is an open-source, MLIR-based FHE compiler toolchain led by Google. Its aim is to let developers write ordinary-looking programs, mark what is secret, and compile them to an FHE scheme and a backend: a software library or a hardware accelerator.
heir-opt runs pass pipelines (scheme selection, parameter choice, packing and layout, lowering) and heir-translate emits code for a backend library.A first-order look at why FHE accelerators are memory and data-movement machines. A ciphertext is 2 polynomials × N coefficients × L limbs × 8 bytes. This is a rough sizing aid, not a scheme-accurate cost model.
For the full treatment (bootstrapping, key traffic, scratchpad sizing and optical NTT engines, simulated), see the sister series FHE Accelerator Simulators. The same lesson as the LLM roofline in deck 03: a faster compute primitive moves the bottleneck to memory and data movement unless the architecture keeps data close to the engine. Key-switching keys (often far larger than a ciphertext) make this worse.
| Layer | AI path | FHE path | Owned by |
|---|---|---|---|
| Frontend | PyTorch (export, compile, dispatch), ONNX | HEIR frontends | Framework integration |
| IR and lowering | FX/ATen → MLIR (torch-mlir, StableHLO, linalg) | secret → ckks/bgv → polynomial | Compiler |
| Mapping | Tile GEMMs, schedule attention, place KV | Schedule NTTs, key switches, rotations | Compiler and architecture |
| Simulator | Timing (and optionally function) of the schedule | Same engine, FHE primitives | Simulation |
| Oracle | PyTorch eager outputs | OpenFHE or Lattigo outputs | Verification |
| Report | Latency, throughput, hot-spots, traces | Latency per operation, bootstrapping cost, data movement | Simulation |
The discrete-event core, metrics, traces and test ladder from decks 02, 06 and 08 do not care whether the operations are GEMMs or NTTs. Keep the cost models and workload frontends pluggable and one simulator serves AI and FHE markets alike.
Deck 10 is a reading list for going deeper on everything in this series.