Inductor Front-End: Deep Dive

This page provides a detailed reference for the Torch-Spyre Inductor front-end compiler. For a high-level overview of the full compilation pipeline, see Compiler Architecture.

The Torch-Spyre compilation pipeline. The left end (green) is entirely upstream PyTorch: Dynamo/Autograd and Inductor. The right end (pink) is Torch-Spyre’s custom Inductor backend, which generates OpSpecs, SuperDSCs, and host code. Torch-Spyre also adds configurations and extensions to the upstream stages to tailor them for the Spyre device.

Inductor Backend Registration

At import time the Spyre backend registers three components with Inductor. Together they take the place of the Triton/CUDA codegen path on a GPU:

Component	Module	Role
SuperDSCScheduling	scheduler.py	Inductor backend scheduling class. Decides how to group and order operations on the LoopLevelIR. Replaces Triton scheduling. Inductor’s own node-pairwise fusion (can_fuse_vertical/can_fuse_horizontal) is permanently disabled here. Real op-to-kernel grouping happens later, in spyre_fuse_nodes (see CustomPostFusionPasses below).
SpyrePythonWrapperCodegen	wrapper.py	Inductor wrapper-codegen class. Generates the Python wrapper that allocates tiled buffers via spyre_empty_with_layout() and dispatches kernels via async_compile.sdsc().
SpyreDeviceOpOverrides	device/op_overrides.py	Device-specific op overrides surfaced to Inductor.

The Spyre-specific Inductor configuration (decompositions, lowerings, the mm_to_bmm_pass that rewrites 2D matmul into 3D bmm for better core utilization, fusion heuristics, and dataflow-friendly Inductor config overrides) is activated through a single context manager:

from torch_spyre._inductor.patches import enable_spyre_context

with enable_spyre_context(...):
    compiled = torch.compile(model)

enable_spyre_context is the central entry point that wires everything together. The three registrations above happen earlier, at package import time, in _inductor/__init__.py.

Extending Compilation

The front-end adds compilation passes into upstream Inductor via six extension points, all registered in passes.py:

Extension Point	Stage	Purpose
CustomPreGradPasses	Pre-grad FX graph	Reserved for graph rewrites before autograd partitioning. The pipeline is empty today.
CustomPrePasses	Post-grad FX graph (early)	collect_spyre_hints snapshots spyre_hint annotations so they survive AOT re-tracing.
CustomPostPasses	Post-grad FX graph (late)	Late post-grad rewrites: decompose_addmm, mm_to_bmm_pass, bmm_unflatten_pass. (recover_spyre_hints runs separately, via a GraphLowering monkey-patch in patches.py.)
CustomPreFusionPasses	LoopLevelIR (pre-fusion)	Pre-fusion scheduler passes: propagate_mutation_layouts, build_loop_scheduler_nodes.
CustomPostFusionPasses	LoopLevelIR (post-fusion)	Post-fusion scheduler passes: spyre_fuse_nodes, prepare_spyre_kernels, hbm_pool_planning, verify_carried_reduction_ownership. Each bundle’s kernel is prepared once, with the real finalization, while HBM fallback is still available; emission binds only what pooling decided. The final pass checks the carried-reduction stages after placement is complete.
CustomPreSchedulingPasses	LoopLevelIR (pre-scheduler)	The pre-scheduling pipeline that runs immediately before the Scheduler is constructed (wired in via a GraphLowering._update_scheduler monkey-patch in patches.py). The full step list is in LoopLevelIR Passes below.

FX Graph Passes

Transformations on the FX Graph tend to be simpler to implement, but happen before the layout of intermediate Tensors in device memory has been computed. Therefore they need to be layout-agnostic. Some examples of passes that are appropriate to perform at this level are:

• replacing constants with size 1 tensors

• normalizing 2D mm into 3D bmm (mm_to_bmm_pass)

LoopLevelIR Passes

Passes on the LoopLevelIR run late in compilation. CustomPreSchedulingPasses dispatches them in a fixed order. Each step takes the GraphLowering and mutates graph.operations in place. Steps marked “Gated” are skipped when their config flag is off.

Working-set reduction (WSR) runs in two separate slots rather than one: a hint-driven half runs immediately after dead-code elimination, before stickification, because it only needs host-side FixedLayout (size/stride) and loop-variable ranges; a span-overflow half stays after stickification because it needs FixedTiledLayout.device_layout (device size, stride map) to reason about physical span. Running the hint-driven half before stickification also dissolves a cross-phase contract that used to exist between insert_restickify and the hint-copy machinery (issue #3135).

The first pass, splice_while_loops, lowers each for_each_tile construct. A for_each_tile region reaches Inductor as a WhileLoop operation. The pass proves the loop is a bounded tile scan, splices its body into the enclosing graph, and coarse-tiles the spliced body immediately so the synthesized dimension hints survive the later propagate_named_dims and assign_dim_hints passes. It runs to a fixed point, so a nested for_each_tile whose inner WhileLoop only appears after the outer body is spliced is also lowered. A WhileLoop that does not prove out as a tile scan is left in place for the upstream default path.

splice_while_loops proves each for_each_tile WhileLoop, splices its body into the enclosing graph, and coarse-tiles the spliced body. The pass re-runs until no WhileLoop remains, so a nested for_each_tile is lowered once its enclosing body has been spliced. A WhileLoop that does not prove out as a tile scan is left for the upstream default path.

#	Pass	Module	Notes
1	splice_while_loops	wsr/for_each_tile_lowering.py	Splices each for_each_tile WhileLoop region into its enclosing graph and coarse-tiles the spliced body. Runs to a fixed point so nested for_each_tile loops resolve.
2	deadcode_elimination	deadcode_elimination.py	Drops unreachable ops.
3	propagate_named_dims	propagate_named_dims.py	Propagates name_tensor_dims annotations from inputs through the op graph. Runs pre-stickification, so it only needs host-side FixedLayout.
4	validate_named_dims	propagate_named_dims.py	Checks the propagated named-dimension annotations for consistency before hints are lowered.
5	assign_dim_hints	propagate_named_dims.py	Lowers each spyre_hint scope to a per-op DimHint list.
6	_maybe_reorder_unhinted_interlopers	passes.py	Gated by config.ignore_wsr_hints (SPYRE_INDUCTOR_IGNORE_HINTS). Runs reorder_unhinted_interlopers to move unhinted ops that interrupt a hint-group run.
7	_maybe_coarse_tile_hints	passes.py	Gated by config.ignore_wsr_hints (SPYRE_INDUCTOR_IGNORE_HINTS). Hint-driven half of coarse tiling: hints_to_coarse_tile_groups + validate_coarse_tile_groups + coarse_tile_pre_stickify. Runs on host-side FixedLayout only.
8	insert_bmm_padding	padding.py	Pads y’s K to a stick boundary for mm/bmm. Runs pre-stickification so the padded buffer is laid out (and restickified if the matmul reads it through a view such as a transposed nn.Linear weight) like any user-written F.pad.
9	split_multi_ops	split_multi_ops.py	Splits multi-op loop bodies (e.g. type conversion + arithmetic) into separate single-op buffers and materializes constant args as SpyreConstantFallback.
10	propagate_spyre_tensor_layouts	propagate_layouts.py	Stamps FixedTiledLayout on every ComputedBuffer.
11	reorder_nonstick_dims	nonstick_dim_order.py	Swaps the largest non-stick device dimension into the slot between the two stick dimensions on matmul inputs, improving matmul work division and LX co-optimization.
12	validate_ops	split_multi_ops.py	Checks that each op’s inputs share the same ElementArrangement. Runs after layout propagation, when the SpyreTensorLayouts are available.
13	optimize_restickify_locations	optimize_restickify.py	Moves restickify ops to better placements before the layout is finalized.
14	finalize_layouts	insert_restickify.py	Settles tile-structure decisions before any new restickify is inserted.
15	insert_restickify	insert_restickify.py	Adds explicit re-tile ops where adjacent ops disagree on layout.
16	validate_no_restickify_on_mutation_targets	insert_restickify.py	Checks that no restickify was inserted directly on a mutation target.
17	enforce_indirect_access_layout	enforce_indirect_access_layout.py	Constrains the layout of tensors consumed by indirect (gather-style) access so the indexed dimension is addressable.
18	insert_post_mutation_restickify	insert_restickify.py	Handles restickification for slice-mutation buffers.
19	insert_restickify_padding	padding.py	Pads restickify ops to satisfy hardware alignment.
20	dedup_and_promote_constants	dedup_constants.py	Deduplicates identical constants and promotes shared ones.
21	_maybe_coarse_tile_span_overflow	passes.py	Gated by config.ignore_span_overflow_hints (SPYRE_INDUCTOR_IGNORE_SPAN_OVERFLOW_HINTS, default on). Span-overflow half of coarse tiling: span_overflow_groups + validate_coarse_tile_groups + coarse_tile_post_stickify. Runs post-stickification, so it needs FixedTiledLayout.device_layout for physical span arithmetic.
22	span_reduction	work_division.py	Reduces per-core access spans to fit the hardware memory budget.
23	cost_model_matmul_division + work_distribution	work_division.py	The cost-model pass claims a subset of matmuls; work_distribution covers the rest.
24	scratchpad_planning	scratchpad/	Gated by config.lx_planning. Allocates the LX scratchpad.
25	elide_proven_read_copies	read_copy_elision.py	Removes copies whose direct-read form is proven equivalent, after physical planning.

Once stickification has run, every ComputedBuffer carries a FixedTiledLayout, so the later passes can take device layout into account when making decisions.

For deeper treatment of individual passes see Working Set Reduction, Coarse-Tiling Loops, Work Division Planning, and Scratchpad Planning.

Views and Index Translation

Real models lean on views heavily. Here is the RoPE block from Granite:

def rope(cached_freqs, q):
    q_ = q.view(2, 256, 32, 128).view(2, 256, 32, 2, 64)         # B L H 2 D/2
    mul_out = cached_freqs[:, :, None, :, :, :] * q_.unsqueeze(-3)  # B L H 2 2 D/2
    sum_out = mul_out.sum(4, keepdim=True)                        # B L H 2 1 D/2
    return sum_out.flatten(3)                                     # B L H D

Two view calls, an unsqueeze, a reduction, and a flatten, all on the hot path of inference. Materializing a tensor copy at every one of those view boundaries would erase any benefit from tiling. The rest of this section walks through how the compiler keeps that from happening.

PyTorch models are full of view operations: reshape, view, transpose, permute, flatten, unsqueeze, slicing, and so on. A single transformer block in Granite goes through dozens of them. On Spyre we want most of these views to cost nothing at runtime; materializing a copy every time a tensor is reshaped would defeat the point of tiling.

A worked end-to-end example. Two tensors with different PyTorch shapes (x is rank-3, y is rank-2) both flow through one shared Inductor index expression. The compiler then lifts that single expression into a distinct device-coordinate vector for each tensor, and finally co-simplifies the iteration space so the integer divisions and modulos collapse into ordinary loop variables. The two tensors end up with different per-argument dim orders, which SuperDSC handles natively.

Inductor gives us a useful starting point. When it lowers a graph that involves views, it normalizes everything onto a shared iteration space and emits a single per-output index expression. For example, this code:

x = torch.rand(50, 10, 200, dtype=torch.float16)
y = torch.rand(500, 200, dtype=torch.float16)

def f(x, y):
    return x.flatten(0, 1) + y

result = torch.compile(f)(x.to("spyre"), y.to("spyre")).cpu()

produces an Inductor body that looks roughly like:

var_ranges = {p0: 500, p1: 200}
index0 = 200*p0 + p1

def body(self, ops):
    get_index = self.get_index('index0')
    load   = ops.load('arg0_1', get_index)
    load_1 = ops.load('arg1_1', get_index)
    add    = ops.add(load, load_1)
    store  = ops.store('buf0', get_index, add, None)
    return store

Both tensors share index0 even though x was rank-3 and y was rank-2 in the original program: the flatten was absorbed into a single linear expression 200*p0 + p1. From here, the Spyre compiler has to turn that one expression into per-tensor device coordinates. There are three steps.

1. Lift index expressions to device coordinates. The host iteration variables (p0, p1) describe positions in the PyTorch shape. They are mapped into per-tensor device coordinate expressions that walk the tiled, padded device shape. Continuing the example:

Host vars	x: device shape [10, 4, 50, 64]	y: device shape [4, 500, 64]
{p: 500, s: 200}, index0 = 200*p + s	[p%10, s//64, p//10, s%64]	[s//64, p, s%64]

The stick dimension always comes out as a pair of expressions, one for the tile index (floor(s/64)) and one for the intra-stick offset (Mod(s, 64)).

2. Co-simplify the iteration space and the per-tensor coordinates. A naive translation leaves expensive integer divisions in place. The front-end factors the iteration space so the divisions and modulos disappear. Splitting p into (q, p) with q = p // 10, p = p % 10 gives:

Iteration space	x	y
{p: 10, q: 50, s: 200}	[p, s//64, q, s%64]	[s//64, q, p, s%64]

Each tensor is now indexed with the same iteration variables, but the expressions inside its coordinate vector are simple and the dim orders are different per tensor. That is exactly what the SuperDSC IR allows (layoutDimOrder_ is per-argument). All of this is held in torch_spyre/_inductor/views.py (compute_coordinates, align_tensors, normalize_coordinates).

3. Emit the OpSpec. The simplified iteration space and per-tensor device coordinates are dropped onto the OpSpec and TensorArgs. The “Example: an add OpSpec” section in the Back-End Compiler doc walks through what one of these artifacts looks like.

The net result for the user is what you would expect: ops on tensor views run without cloning whenever the compiler can express the new layout as a different read pattern over the same storage. When that is not feasible (for example when a downstream op forces a different stick dimension), the insert_restickify pass adds an explicit re-stick operation so the rest of the pipeline still sees a clean layout.

Code Generation

We do code generation in three stages.

1. LoopLevelIR nodes are fused together to form Kernels.

2. Each Kernel is processed by spyre_kernel.py to convert it to a list of OpSpec (op_spec.py).

3. Finally, the codegen/ package translates OpSpec into SuperDSC JSON, the input format for the DeepTools back-end compiler.

Our intent is that the OpSpec will capture all important semantic information about the operation in a more human readable form than the SuperDSC JSON. Therefore, the OpSpec should be the primary artifact used to understand the output of the front-end compiler. Inspecting the SuperDSC JSON should only be necessary when debugging problems in the codegen package of the front-end compiler.

Extending Operations

We extend Inductor to compile Spyre-specific operations by adding Custom Operations. We modify how existing operations are compiled by adding Spyre-specific decompositions and lowerings. See Adding Operations for a step-by-step guide.

Custom Operations

Spyre-specific operations with no ATen equivalent are defined in customops.py using @torch.library.custom_op. Each custom op requires:

1. A signature definition (@custom_op)

2. A fake/meta function (@opname.register_fake)

3. Either a lowering + SpyreOpFuncs entry, or a decomposition that removes it from the graph before lowering

Decompositions

Spyre-specific decompositions are registered with @register_spyre_decompositions in decompositions.py. Decompositions transform complex ATen operations into simpler primitives before the graph is lowered to loop-level IR.

Lowerings

Spyre-specific lowerings to Inductor’s LoopLevelIR are defined in lowering.py using the @register_spyre_lowering decorator. This mechanism supports both the replacement of upstream lowerings and the addition of new lowerings for Spyre-specific custom operations.

Module Reference

The headline modules above are the ones a contributor reaches for first. The front-end is also made up of a number of smaller modules; the table below names each and points to the source.

Module	Purpose
passes.py	The six extension-point classes. It renders the LoopLevelIR before and after the pre-scheduling pipeline via format_operations (defined in pass_utils.py).
temp_passes.py	Transitional FX-graph rewrites registered in CustomPostPasses: decompose_addmm, mm_to_bmm_pass, and bmm_unflatten_pass. The “temp” name reflects the plan to retire them as upstream Inductor evolves.
propagate_layouts.py	propagate_spyre_tensor_layouts and propagate_mutation_layouts. Assigns FixedTiledLayout to every ComputedBuffer.
split_multi_ops.py	split_multi_ops and validate_ops. Splits multi-op loop bodies into single-op buffers (materializing constant args as SpyreConstantFallback) and validates that op inputs share the same ElementArrangement.
optimize_restickify.py	Optimizes restickify operations inserted by layout propagation.
insert_restickify.py	finalize_layouts, insert_restickify, validate_no_restickify_on_mutation_targets, insert_post_mutation_restickify. Settles tile-structure decisions, adds explicit re-tile ops where adjacent ops disagree on layout, checks that no restickify lands on a mutation target, and handles restickification for slice-mutation buffers.
hbm_pool_planning.py	HBM-pool allocation for intermediates not in LX.
deadcode_elimination.py	deadcode_elimination for the pre-scheduling LoopLevelIR.
work_division.py	span_reduction, cost_model_matmul_division, work_distribution, divide_pointwise_op, divide_reduction_op, apply_splits. Three-pass work division: span reduction (mandatory), cost-model matmul division (claims a subset of matmuls), and work distribution (covers the rest).
scratchpad/	scratchpad_planning and the layout-solver framework (GreedyLayoutSolver, FirstFitLayoutSolver, BestFitLayoutSolver). LX scratchpad allocation.
propagate_hints.py	spyre_hint context manager, the DimHint dataclass, and collect_spyre_hints / recover_spyre_hints for surviving AOT re-tracing.
propagate_named_dims.py	declare_tensor_dim, name_tensor_dims, propagate_named_dims, and assign_dim_hints. Propagates named-dim metadata through the op graph and lowers spyre_hint scopes to per-op DimHint lists.
coarse_tile.py	coarse_tile_pre_stickify, coarse_tile_post_stickify, _plan_tiling_propagation, _insert_all_read_copy_ops, _insert_all_reduction_ops, _insert_all_write_copy_ops. Wraps each hint-derived group in nested counted loops, scales per-iteration ranges, and plans/executes buffer propagation across loop boundaries.
coarse_tile_hints.py	hints_to_coarse_tile_groups, reorder_unhinted_interlopers. Derives coarse-tile groups from spyre_hint scopes and orders unhinted ops that sit between hinted ones.
coarse_tile_span_overflow.py	span_overflow_groups. Derives coarse-tile groups for ops whose per-core span exceeds the hardware limit, using device layout.
dedup_constants.py	dedup_and_promote_constants. Deduplicates identical constants in the LoopLevelIR.
loop_info.py	CoarseTileInfo, copy_op_metadata. Per-op metadata stamped by coarse_tile() and consumed by the scheduler, kernel codegen, and buffer-propagation pass.
dtype_ops.py	DtypeOpTable. Lookup table mapping PyTorch dtype pairs to Spyre hardware dtype-conversion operators.
pass_utils.py	Shared helpers for the pre-scheduling pipeline, including symbol-keyed work-division ownership and final Scheduler-boundary splits_by_index_coeff / apply_splits_from_index_coeff transport helpers.
views.py	compute_coordinates, align_tensors, normalize_coordinates, Term, matching_dim. Coordinate computation for memory-dep expressions and tensor alignment for fused kernels. Used by spyre_kernel.py and work_division.py.
multi_dim_reduction_pass.py	decompose_multi_dim_reductions. Splits a multi-dim reduction into a sequence of single-dim reductions. Not currently registered in any pass pipeline.
op_spec.py	OpSpec, TensorArg, UnimplementedOp. The high-level per-operation artifact emitted by spyre_kernel.py.
provenance.py	build_debug_handle. Source-to-kernel provenance: builds DebugHandle / SourceLoc from ComputedBuffer.origins, consumed by spyre_kernel.py and serialized into SuperDSC as debug_handle_.
spyre_kernel.py	Converts a fused kernel into a list of OpSpecs.
padding.py	insert_restickify_padding, insert_bmm_padding. Pads restickify ops and mm/bmm operands to satisfy hardware alignment.
fusion.py	spyre_fuse_nodes. Post-fusion scheduler pass.
wrapper.py	SpyrePythonWrapperCodegen. The host-code generator that produces the Python wrapper around device kernels.
codegen/	superdsc.py (SDSCArgs, SDSCSpec, parse_op_spec, compile_op_spec, _get_padded_iteration_space, _get_op_dim_labels), compute_ops.py (generate_sdsc), bundle.py (generate_bundle, emits scf.for for coarse-tiling LoopSpec trees). Translates OpSpec into SuperDSC JSON for the back-end compiler.
core_mapping.py	core_to_slice_mapping. Maps core_id values to work-slice offsets during SuperDSC emission.
enforce_indirect_access_layout.py	enforce_indirect_access_layout. Constrains the layout of tensors consumed by indirect (gather-style) access.
config.py	Spyre-specific Inductor configuration. The module attributes sencores, lx_planning, and dxp_lx_frac_avail are populated from the SENCORES, LX_PLANNING, and DXP_LX_FRAC_AVAIL environment variables.
ir.py	FixedTiledLayout, SpyreConstantFallback, SpyreEmptyFallback, SpyreReduction. Core IR types used throughout the pre-scheduling pipeline.
choices.py	Spyre-specific InductorChoices overrides (e.g. reduction split heuristics).
errors.py	Unsupported exception class for ops or configurations not yet handled by the Spyre backend.
indirect_access.py	Indirect-access helpers for gather/scatter ops on OpSpec.
span_overflow_hint_analysis.py	plan_span_overflow_tile, SpanOverflowTilePlan. Generates coarse-tiling hints for ops whose per-core span exceeds the hardware limit.
constants.py	Shared constants: DEVICE_NAME, BATCH_MATMUL_OP, TOPK_OPS.
logging_utils.py	get_inductor_logger and Spyre Inductor logging setup.

torch.compile(..., dynamic=True) is supported through the static-binary path. Shapes are specialized at compile time and the resulting binary is reused across calls with the same input geometry.