Skip to content

Torch Code Generation (Torch Codegen)

Overview

torch_codegen lowers PyPTO IR into an executable Python/PyTorch script that can be run with exec() for debugging and numerical validation.

Unlike production codegen (PTO/Orchestration), torch_codegen is designed to:

  • quickly reproduce IR semantics,
  • expose intermediate behavior in Python,
  • provide executable references for pass debugging and system tests.

Source location: python/pypto/debug/torch_codegen.py

Public entry API:

torch_codegen(node: _ir.Program | _ir.Function, check_shapes: bool = False) -> str

Goals and Boundaries

Goals

  • Keep IR-to-Python expression/statement mapping readable.
  • Cover common tensor/tile operations.
  • Execute the default dense Vec Buffer transfers and registered vector/scalar arithmetic (including full), including destination writes and storage aliases, for kernel.lower() and CompiledProgram.validate_ir(). Structured control flow preserves simultaneous scalar yields and planned Tile carry copies; ordinary cube operations model fractal layouts and shared physical storage, including transpose and wider accumulator views. Buffer pass dumps retain readable Python text and a lossless .msgpack companion; validate_ir() replays the companion because Buffer text is not executable DSL.
  • Simulate concurrent cross-core tpush/tpop behavior.
  • Provide actionable diagnostics for suspicious inputs (timeouts, invalid split dimensions, etc.).

Addressed Vec and cube allocations share one byte pool per memory space and function invocation. Overlapping addresses therefore alias even across distinct allocation handles; addressless allocations remain independent. Window views retain their parent pitch and have independent valid metadata.

Non-goals

  • No performance optimization.
  • No cycle-accurate Ascend hardware timing/memory simulation.
  • Not a production backend for training/inference.

High-Level Architecture

Generated code is assembled from three parts:

  1. Runtime preamble (_PREAMBLE) Provides helpers, tile/tensor boundary handling, cross-core runtime, and mixed-kernel scheduler.
  2. Group metadata injection (_GROUP_META.update(...)) Injected only when input is a Program and Group/AIC/AIV pairs are detected.
  3. Function bodies Emitted by TorchCodegen(IRVisitor) function-by-function and statement-by-statement.

Data flow:

Program/Function IR -> (optional _build_group_meta) -> TorchCodegen visitor emission -> final script string

Expression and Statement Emission Model

Expression layer

  • _OP_MAP maps op_name to OpHandler(args, kwargs) -> str.
  • _visit_expr_str() forces nested Call nodes through Python-side visit_call, avoiding nested-call result loss in some C++ visitor paths.
  • Binary/unary IR nodes are emitted through _BINARY_OP_STR plus generic visitor adapters.

Statement layer

  • AssignStmt: var = expr
  • EvalStmt: emit expression as-is (for side-effect calls)
  • ReturnStmt: supports single and tuple returns
  • ScopeStmt: transparent passthrough
  • ForStmt/WhileStmt: lower SSA iter_args + yield into mutable variable updates
  • IfStmt: branch-local yield writes back to return_vars

Naming and Function Isolation

_unique_name() normalizes variables by:

  • replacing non-identifier characters with _,
  • collapsing repeated underscores,
  • avoiding Python keywords and digit-leading names.

visit_function() resets _var_names/_name_counter/_yield_targets per function to prevent cross-function name pollution when object IDs are reused.

Operation Mapping System (_OP_MAP)

Mappings are registered by category:

  • tensor/tile elementwise, broadcast, reduction, logic, bitwise
  • matmul/matmul_acc and tile variants
  • create/full/cast/slice/read/write/assemble/fillpad
  • cross-core pipe operations
  • system.* operations (no-op in debug codegen)

Accumulating matmuls and init_cond

tile.matmul_acc / tensor.matmul_acc take an optional trailing init_cond operand — a BOOL scalar carrying the MAD's cmatrixInit bit. Where the predicate holds, the op overwrites the accumulator with lhs @ rhs instead of accumulating into it, so a reference that ignored it would model an overwrite as an accumulate.

The predicate is a scalar, not an elementwise mask, so it emits a whole-op branch rather than a torch.where:

Call Emitted expression
matmul_acc(acc, a, b) (acc + torch.matmul(a, b).float())
matmul_acc(acc, a, b, k == 0) _acc_init(acc, torch.matmul(a, b).float(), (k == 0))
matmul_acc(acc, a, b, True) torch.matmul(a, b).float()
matmul_acc(acc, a, b, False) (acc + torch.matmul(a, b).float())

A literal predicate is folded to the arm it selects, mirroring how PTO codegen picks pto.tmatmul over pto.tmatmul.acc instead of emitting an scf.if on a compile-time constant. _acc_init is a preamble helper shared by every accumulating matmul, so each honours the predicate through the one handler: tile.matmul_acc, tensor.matmul_acc, tile.batch_matmul_acc, and tile.gemv_acc (GEMV is a matmul whose M is 1, on the same cube MAD).

Cross-core operation mappings

  • tile.tpush_to_aiv -> _cross_core_rt.push_to_aiv(tile, split[, lane_stride])
  • tile.tpush_to_aic -> _cross_core_rt.push_to_aic(tile, split)
  • tile.tpop_from_aic -> _cross_core_rt.pop_from_aic(split)
  • tile.tpop_from_aiv -> _cross_core_rt.pop_from_aiv(split)
  • tile.get_subblock_idx -> _get_subblock_idx()

split is normalized by _split_mode_to_int() and carries pto-isa's TileSplitAxis:

  • 0: NONE
  • 1 / 2: UP_DOWN / LEFT_RIGHT — the two lanes hold equal extents
  • 3 / 4: the same axes over an ODD extent — lane 0 holds one cell more
  • invalid/unparsable values: raise ValueError (fail fast)

The odd codes are validated against the tile's RUNTIME valid extent, not its physical box: an even box with an odd valid extent (16 with valid_shape 15) is exactly the case they exist for. lane_stride rides along only when LowerAutoVectorSplit rebalanced a ragged boundary; the runtime then cuts the lanes there instead of at the box half. Odd codes are rejected on push_to_aic / pop_from_aiv, mirroring PTO codegen — pto-isa has no odd Vector-to-Cube transport.

Tile/Tensor Boundary and Valid-Region Semantics

Helpers in _PREAMBLE define boundary behavior:

  • _tile_load: zero-pad to requested shape on OOB and attach valid region
  • _tile_store: store only the valid region
  • _tensor_slice: materialize requested shape even when slicing out-of-bounds
  • _fillpad: pad invalid area with zero/min/max
  • _assemble: write source valid region into target at offsets

Valid region is propagated via dynamic attributes:

  • _pypto_valid_shape
  • _pypto_full_shape

Cross-core queueing/split/merge paths preserve these attributes as well, so boundary-tile valid regions remain consistent through tpush/tpop.

Cross-Core Runtime Design

Structure

_CrossCoreRuntime uses threading.Condition to implement blocking queue semantics. It maintains:

  • normal channels:
  • to AIV: _to_aiv and _to_aiv_split[split][lane]
  • to AIC: _to_aic and _to_aic_split[split][lane]
  • no-split dual-dispatch channels:
  • to AIV: _to_aiv_dual_nosplit[lane]
  • to AIC: _to_aic_dual_nosplit[lane]

It also provides:

  • reset(no_split_dual_aiv_dispatch=False): clear channels and configure no-split dual-dispatch mode before each mixed-group call
  • snapshot(): queue-depth snapshot for timeout diagnostics, including no-split dual-dispatch mode and dual-nosplit queue depths

Push/Pop semantics

push_to_aiv(tile, split) / pop_from_aic(split)

  • split=0 and no_split_dual_aiv_dispatch=False: single queue
  • split=0 and no_split_dual_aiv_dispatch=True:
  • push_to_aiv broadcasts one tile copy to both lane queues
  • pop_from_aic consumes from the current lane queue (lane in {0,1})
  • split=1/2: push splits tile into lane0/lane1; pop consumes by current lane

push_to_aic(tile, split) / pop_from_aiv(split)

  • split=0 and no_split_dual_aiv_dispatch=False: single queue
  • split=0 and no_split_dual_aiv_dispatch=True:
  • push_to_aic enqueues by current lane (lane in {0,1})
  • pop_from_aiv waits for both lane queues, then consumes a pair and returns lane0 payload
  • split=1/2: push enqueues by current lane; pop waits until both lanes are ready, then merges

Split dimension semantics (mirroring pto-isa's TileSplitAxis):

  • split=1 (UP_DOWN) / split=3 (UP_DOWN_ODD): split/merge on dim=0
  • split=2 (LEFT_RIGHT) / split=4 (LEFT_RIGHT_ODD): split/merge on dim=1
  • the _ODD codes take the ceil half for lane 0 and the floor half for lane 1

Synchronization and timeouts

  • per-pop timeout: _PIPE_WAIT_TIMEOUT_SEC (default 10s)
  • mixed-kernel group timeout: _MIXED_KERNEL_TIMEOUT_SEC (default 30s)
  • on mixed-kernel timeout/failure, runtime signals cancellation and notifies all waiters to avoid stale blocked threads polluting subsequent runs

Timeout errors include:

  • operation name, split, lane
  • alive thread names (group timeout)
  • current pipe snapshot (queue depths)

Group Mixed-Kernel Concurrent Scheduler

Metadata construction (_build_group_meta)

For each FunctionType.Group function in a Program, codegen pairs by naming convention:

  • <group>_aic
  • <group>_aiv

Metadata fields:

  • aic / aiv: callee names
  • split: Group split first; fallback to AIV split when Group split is 0
  • dual_aiv_dispatch: read from AIV attrs; for split==0, this forces dual AIV-lane dispatch in debug runtime

Scheduler entry

In visit_call(), for GlobalVar calls:

  • if name exists in _group_meta, emit _run_group_call(group_name, *args)
  • otherwise emit a normal direct function call

_run_group_call() dispatches to _run_mixed_kernels() when metadata exists.

Thread model

_run_mixed_kernels(group_name, meta, *args) behavior:

  • exactly 1 AIC thread
  • AIV thread count:
  • split in (1,2) -> 2 lanes
  • split == 0 and dual_aiv_dispatch == True -> 2 lanes
  • split == 0 and dual_aiv_dispatch == False -> 1 lane
  • each thread writes thread-local subblock_idx (AIC uses 0, AIV uses lane id)
  • runtime mode switch:
  • scheduler calls reset(no_split_dual_aiv_dispatch=(split == 0 and dual_aiv_dispatch))
  • no-split dual-lane pipe semantics are enabled only when required by group metadata
  • return contract: only aiv lane0 return value is propagated as Group return; if other lanes/roles produce non-None returns, runtime raises a contract violation error

Shape/Type Checks (check_shapes)

When check_shapes=True:

  • for function parameters: dtype check only, shape check disabled
  • rationale: InCore parameters may be boundary tiles with partial data
  • for assignment targets: tensor type and dtype checks enabled; shape checks follow static/dynamic strategy

Dynamic-dimension strategy:

  • if dimensions are not all static ConstInt, codegen checks ndim plus each static dimension index.

Error Handling and Observability

Main error classes:

  • TypeError: invalid entry node type
  • ValueError: unsupported op, invalid split/lane, split dimension whose parity contradicts the split code (1 / 2 need an even extent, 3 / 4 an odd one)
  • RuntimeError: pipe wait timeout, mixed-kernel thread failure/timeout

For concurrency failures, codegen reports the first captured thread traceback to improve localization by function and lane.

Current Limitations

  • system.* ops are emitted as no-ops in debug mode
  • Group pairing relies on <group>_aic/_aiv naming convention
  • split=1/2 requires an even split dimension; split=3/4 requires an odd one
  • Group return is fixed to AIV lane0 output
  • Runtime is a Python semantic simulator, not equivalent to hardware pipeline timing

Recommended coverage dimensions:

  • all split modes: NONE/UP_DOWN/LEFT_RIGHT
  • all communication directions: V->C, C->V, bidirectional V<->C
  • lane semantics via tile.get_subblock_idx
  • no-split + dual_aiv_dispatch lane behavior (lane0 + lane1)
  • Group call rewrite path (_run_group_call)
  • timeout/error paths (for example, intentionally unpaired push/pop)

References:

  • tests/ut/debug/test_torch_codegen.py
  • tests/st/codegen/torch/test_torch_codegen_cross_core.py
  • tests/st/codegen/torch/test_torch_codegen_qwen3_decode_scope3_mixed.py

Relation to Other Codegen Docs

They are complementary: torch_codegen provides executable semantic reference, while PTO/Orchestration target production generation.