Dataflow Analyses in PyTorch and Triton

Short answer: yes—both do, but in different ways.Triton (the kernel compiler used by PyTorch Inductor’s GPU backend) is built on MLIR, and it uses MLIR’s dataflow analysis framework directly—the one that’s literally built around lattices and transfer functions. You can see Triton passes using mlir::dataflow::Lattice in the source (e.g., Allocation.cpp), which is exactly the “semilattice + transfer” setup you asked about. github.com
MLIR’s docs also spell out the framework (lattices, join/meet, fixpoints, transfer functions over IR), which is what Triton is leaning on. mlir.llvm.org
And in Triton’s pipeline, the MLIR IR is lowered to NVVM/LLVM and then codegen’d, so LLVM/MLIR optimization and analyses run on the way down. (Triton is explicitly an LLVM/MLIR-based compiler; the MLIR GPU -> NVVM pipeline even exposes an opt-level knob.) github.commlir.llvm.orgPyTorch itself has two stories:

  • In the legacy TorchScript JIT, PyTorch implements classic, fixed‑point dataflow analyses over its SSA IR. There’s an explicit liveness analysis that iterates to a fixed point (“canonical fixed-point liveness”), and an alias analysis that builds an AliasDb/MemoryDAG to reason about mutation/aliasing for safety of transforms (CSE/DCE/etc.). These are text‑book dataflow (lattice = sets, join = union, backward/forward transfer). gitlab.maisondelasimulation.frgithub.com
    Constant propagation and DCE are separate passes and are alias‑aware (constant‑prop has variants that consult alias info). gitlab.maisondelasimulation.fr
  • In the newer torch.compile / TorchInductor path, most graph work is scheduling, fusion, and codegen; there isn’t a public, general “lattice framework” in FX/Inductor the way MLIR has one. Inductor’s GPU kernels are emitted to Triton, so you get MLIR/LLVM dataflow “via Triton.” On CPU, Inductor emits C++ and hands it to a host compiler (Clang/GCC), so any LLVM dataflow is applied to the lowered C++/LLVM IR, not to the FX/Torch IR. (Inductor’s IR does track aliasing/mutation via TensorBox/StorageBox, but that’s not a generic, MLIR‑style dataflow framework.) dev-discuss.pytorch.org If you want a ready-made semilattice + transfer‑function API to write your own analyses, Triton/MLIR is the place to look (e.g., mlir::dataflow::Lattice, DataFlowAnalysis), and Triton already ships analyses like Axis/alias/alloc decisions on top of it. mlir.llvm.orggithub.comIf you were hoping PyTorch had a similar reusable framework inside its IRs: the JIT has specific analyses (liveness/alias), but there isn’t a general “plug‑in a lattice” API there. If you take the torch‑mlir route, then you can use MLIR’s dataflow infra on models originating from PyTorch.Happy to point you at concrete pass entry points or help sketch a tiny MLIR dataflow analysis if you tell me the property you want to track (ranges, contiguity/strides, purity, etc.).