Skip to content

Dynamic shape on the C++ trace path (elementwise add)#269

Open
YWHyuk wants to merge 4 commits into
feature/togsim-cpp-tracefrom
feature/togsim-dynamic-shape
Open

Dynamic shape on the C++ trace path (elementwise add)#269
YWHyuk wants to merge 4 commits into
feature/togsim-cpp-tracefrom
feature/togsim-dynamic-shape

Conversation

@YWHyuk

@YWHyuk YWHyuk commented Jun 23, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds torch.compile(dynamic=True) support on the C++ trace path. One compiled
trace producer .so serves any input size: the producer reads its loop bound
from shape_args at runtime, while the per-tile cost table (sampled once, since a
tile's cost is shape-invariant) keys the timing.

Verified end to end: a single dynamic a + b runs at N=1024 (2 tiles, 183
cycles) and N=2048 (4 tiles, 261 cycles) from the same .so, the trace
scaling with shape_args.

What each commit does

  1. Guard MLIR tile sizing against symbolic dims — neutralize the tile-fit
    heuristics for a sympy dim (they only minimize a known dim's tail); all gated
    on isinstance(.., sympy.Expr), so static is unchanged.
  2. Emit symbolic loop bounds and dynamic memref dims — size symbol becomes a
    scalar-int kernel arg; memref<?xf32>; affine.for ... to %<name>_bound with a
    top-of-function memref.load + index_cast prologue.
  3. Make the kernel meta import-safe — stringify sympy in arg_attributes so the
    generated wrapper imports (the extent already arrives as a runtime arg).
  4. Skip compile-time Spike validation for dynamic — its fixed-shape validation
    binary can't be instantiated for a runtime extent.
  5. Sample per-tile cycles on a one-tile copypin_loops_to_one_tile (general
    for static + dynamic) runs gem5 sampling on a one-tile copy; the symbolic IR is
    kept for the producer / cost table.
  6. Emit a dynamic-shape trace producerbuild_skeleton skips the loop_end
    serialization it never needed; lower_to_emitc re-sources each still-used size
    arg from shape_args[k] via emitc.subscript, so the producer loop reads
    for (iv=0; iv<shape_args[k]; iv+=step).
  7. Pass the runtime shape via the attribute file — the scalar size goes into the
    existing per-kernel attribute YAML (shape_args), run_standalone passes
    --attribute, and main.cc fills shape_args for run_producer.
  8. Testtests/ops/elementwise/test_dynamic_add.py.

Known limitations / follow-ups

  • Timing only: the Spike functional run (which fills the output tensor) is
    skipped for dynamic, so output values are zeros today; the test checks output
    shape. Functional output for dynamic is a follow-up.
  • Static cost sampling can use the (already general) pin_loops_to_one_tile once it
    is decoupled from run_tog; op coverage beyond contiguous 1D add (matmul /

    1D-dynamic: symbolic-stride addressing, multi-symbol shape_args) is future work.

  • One change to the loop-padding mlir-opt pass (skip a symbolic bound) lives in the
    LLVM fork and is not part of this PR.

Base: stacked on feature/togsim-cpp-trace (#267).

🤖 Generated with Claude Code

@YWHyuk YWHyuk force-pushed the feature/togsim-dynamic-shape branch from c2a243b to 6b39c94 Compare June 23, 2026 06:00
@YWHyuk YWHyuk force-pushed the feature/togsim-cpp-trace branch 2 times, most recently from b7c1ec4 to 9033945 Compare June 24, 2026 13:37
@YWHyuk YWHyuk force-pushed the feature/togsim-dynamic-shape branch from 8e9db1a to 82e9255 Compare June 24, 2026 14:43
@YWHyuk YWHyuk force-pushed the feature/togsim-cpp-trace branch from c166abd to ed5c747 Compare June 25, 2026 07:29
@YWHyuk YWHyuk force-pushed the feature/togsim-cpp-trace branch from 1e43e9a to cfcb7a9 Compare July 7, 2026 07:10
@YWHyuk YWHyuk force-pushed the feature/togsim-dynamic-shape branch 2 times, most recently from 9dbc257 to 908ea30 Compare July 7, 2026 07:37
YWHyuk added 4 commits July 7, 2026 19:37
…asked-DMA clamp

Make the MLIR backend emit valid IR under torch.compile(dynamic=True), where a
loop range is a sympy symbol (e.g. s52) instead of a concrete int.

- is_symbolic_dim: single predicate for "runtime dim"; the tile-fit heuristics
  (trim_large_tail, get_padding_ratio, make_choices) skip a symbolic dim rather
  than doing concrete-int arithmetic on it (kept fixed init tile; tail becomes a
  runtime remainder). get_mlir_shape emits a dynamic memref dim (memref<?>).
- LoopLevel emits a symbolic upper bound as an index SSA (%<name>_bound);
  codegen_loops materializes it (memref.load + index_cast) in the prologue.
- set_ranges registers each symbolic dim's runtime-extent SSA once
  (dyn_bound_cses), the single source of %<name>_bound (dyn_bound_name); the
  loop bound and the masked-DMA clamp both read that var instead of rebuilding
  the name at codegen time.
- Masked-DMA clamp: a symbolic tile axis emits the same per-dim [low, high)
  clamp as a static non-dividing axis, with the runtime extent as the affine.min
  symbol operand -- so a non-dividing dynamic shape is handled by one path and
  the last partial tile falls out at runtime.
- axis_split: aligned axis-split detection is symbolic-aware (+ unit test).

All guards gate on is_symbolic_dim, so the static path is unchanged.
…gnostic validation, producer shape_args

Drive a dynamic kernel through compile + timing + functional validation without
a concrete extent.

- cycle_table.pin_loops_to_one_tile: per-tile cost is shape-invariant, so sample
  cycles on a one-tile copy of the (symbolic) IR; extension_codecache sizes the
  sampling host buffers to one tile (_concretize_attrs_for_sampling).
- mlir_caller_codegen: the Spike validation binary is shape-agnostic -- it loads
  each size arg into N_<sym> and mallocs / builds the memref descriptors from N
  at runtime, so one binary serves any size. extension_codecache builds and runs
  it for a dynamic kernel too; the per-tile spad measurement stays valid, so the
  spad-overflow gate runs unconditionally (static path unchanged).
- lower_to_emitc: the trace producer re-sources its loop bound from shape_args,
  so one trace.so serves any size. build_tog/build_skeleton skip the bfs
  serialization that assumes a concrete loop_end.
- is_dynamic_shape is detected from the size-symbol arg (MLIR_ARGS_VAR), not an
  IR-text sniff.
…via the attribute file

The dynamic trace producer reads its loop bounds from shape_args; feed them at
simulation time through the existing per-kernel attribute YAML (which already
carries address_info), not a bespoke channel.

- simulator: a scalar input (a dynamic size arg, e.g. s52) is not a tensor
  address -- collect such scalars into a shape_args sequence in the YAML, in arg
  order (== the producer's shape_args[k]); pass --attribute alongside --trace_so.
  dump_args/write_arg write the size arg's runtime value (int64) so the
  shape-agnostic Spike binary can load its loop bound.
- main.cc: add --attribute; in the trace branch load the YAML and fill
  shape_args, passed to run_producer (was nullptr, 0).
A single torch.compile(dynamic=True) add compiles to one trace producer .so and
is simulated at several input sizes -- the producer reads its loop bound from
shape_args at runtime, so the same .so serves any size. Exercises the whole
dynamic-shape pipeline end to end, including non-dividing sizes (the masked-DMA
symbolic clamp handles the remainder tile).
@YWHyuk YWHyuk force-pushed the feature/togsim-dynamic-shape branch from 908ea30 to 0797c45 Compare July 7, 2026 10:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant