Dynamic shape on the C++ trace path (elementwise add)#269
Open
YWHyuk wants to merge 4 commits into
Open
Conversation
c2a243b to
6b39c94
Compare
b7c1ec4 to
9033945
Compare
8e9db1a to
82e9255
Compare
c166abd to
ed5c747
Compare
1e43e9a to
cfcb7a9
Compare
9dbc257 to
908ea30
Compare
…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).
908ea30 to
0797c45
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
torch.compile(dynamic=True)support on the C++ trace path. One compiledtrace producer
.soserves any input size: the producer reads its loop boundfrom
shape_argsat runtime, while the per-tile cost table (sampled once, since atile's cost is shape-invariant) keys the timing.
Verified end to end: a single dynamic
a + bruns at N=1024 (2 tiles, 183cycles) and N=2048 (4 tiles, 261 cycles) from the same
.so, the tracescaling with
shape_args.What each commit does
heuristics for a sympy dim (they only minimize a known dim's tail); all gated
on
isinstance(.., sympy.Expr), so static is unchanged.scalar-int kernel arg;
memref<?xf32>;affine.for ... to %<name>_boundwith atop-of-function
memref.load+index_castprologue.arg_attributesso thegenerated wrapper imports (the extent already arrives as a runtime arg).
binary can't be instantiated for a runtime extent.
pin_loops_to_one_tile(generalfor static + dynamic) runs gem5 sampling on a one-tile copy; the symbolic IR is
kept for the producer / cost table.
build_skeletonskips the loop_endserialization it never needed;
lower_to_emitcre-sources each still-used sizearg from
shape_args[k]viaemitc.subscript, so the producer loop readsfor (iv=0; iv<shape_args[k]; iv+=step).existing per-kernel attribute YAML (
shape_args),run_standalonepasses--attribute, andmain.ccfillsshape_argsforrun_producer.tests/ops/elementwise/test_dynamic_add.py.Known limitations / follow-ups
skipped for dynamic, so output values are zeros today; the test checks output
shape. Functional output for dynamic is a follow-up.
pin_loops_to_one_tileonce itis decoupled from
run_tog; op coverage beyond contiguous 1D add (matmul /LLVM fork and is not part of this PR.
Base: stacked on
feature/togsim-cpp-trace(#267).🤖 Generated with Claude Code