feat(torchwave): Stop a returned sym_size bordering its operand (#18867) - #18867
Open
oerling wants to merge 6 commits into
Open
feat(torchwave): Stop a returned sym_size bordering its operand (#18867)#18867oerling wants to merge 6 commits into
oerling wants to merge 6 commits into
Conversation
✅ Deploy Preview for meta-velox canceled.
|
Contributor
|
@oerling has exported this pull request. If you are a Meta employee, you can view the originating Diff in D118061301. |
oerling
added a commit
to oerling/velox-1
that referenced
this pull request
Sep 5, 2026
…bookincubator#18867) Summary: Pull Request resolved: facebookincubator#18867 A metadata getter (`sym_size` / `sym_numel`) whose only role is to be a graph output still counts as a use of its operand when the partitioner builds its levels. `makeLevelsInner` increments the operand's producer refCount, and `makeCseBorder` turns anything above one into a border: the producer moves to its own earlier layer and every other consumer of it follows a layer later than it needed to. The getter's own work is one scalar field read on the host; the whole cost is the layer split it forces. `WaveConfig::deferSizeOutputs`, OFF BY DEFAULT, drops those getters from `top` before `makeExprLevels` so they stop being a reference, and puts them back before the last layer is built so they still run and still fill their output slot. The predicate counts users through the set reachable from the output node, not through `users()`. torch.export leaves dead `_operator.ge` / `_operator.le` shape guards behind once the asserts are stripped -- 428 of them on the ROO preproc graph -- and they appear in a value's `users()` while contributing to no level. Counting them refuses every candidate; with them filtered out, 255 of the graph's 297 top-level exprs qualify and 248 of those have an operand whose only other reachable user is a single consumer. It is off by default because on that graph it is a net loss, and the measurement is the point of the commit rather than the feature. See the test plan. Reviewed By: Yuhta Differential Revision: D118061301
oerling
force-pushed
the
export-D118061301
branch
from
September 5, 2026 08:16
cd0898e to
ef1cbcd
Compare
Selective Build Plan
Selective build plan |
Yuhta
approved these changes
Sep 5, 2026
oerling
added a commit
to oerling/velox-1
that referenced
this pull request
Sep 5, 2026
…bookincubator#18867) Summary: Pull Request resolved: facebookincubator#18867 A metadata getter (`sym_size` / `sym_numel`) whose only role is to be a graph output still counts as a use of its operand when the partitioner builds its levels. `makeLevelsInner` increments the operand's producer refCount, and `makeCseBorder` turns anything above one into a border: the producer moves to its own earlier layer and every other consumer of it follows a layer later than it needed to. The getter's own work is one scalar field read on the host; the whole cost is the layer split it forces. `WaveConfig::deferSizeOutputs`, OFF BY DEFAULT, drops those getters from `top` before `makeExprLevels` so they stop being a reference, and puts them back before the last layer is built so they still run and still fill their output slot. The predicate counts users through the set reachable from the output node, not through `users()`. torch.export leaves dead `_operator.ge` / `_operator.le` shape guards behind once the asserts are stripped -- 428 of them on the ROO preproc graph -- and they appear in a value's `users()` while contributing to no level. Counting them refuses every candidate; with them filtered out, 255 of the graph's 297 top-level exprs qualify and 248 of those have an operand whose only other reachable user is a single consumer. It is off by default because on that graph it is a net loss, and the measurement is the point of the commit rather than the feature. See the test plan. Reviewed By: Yuhta Differential Revision: D118061301
oerling
force-pushed
the
export-D118061301
branch
from
September 5, 2026 18:06
ef1cbcd to
ea81142
Compare
oerling
added a commit
to oerling/velox-1
that referenced
this pull request
Sep 5, 2026
…bookincubator#18867) Summary: Pull Request resolved: facebookincubator#18867 A metadata getter (`sym_size` / `sym_numel`) whose only role is to be a graph output still counts as a use of its operand when the partitioner builds its levels. `makeLevelsInner` increments the operand's producer refCount, and `makeCseBorder` turns anything above one into a border: the producer moves to its own earlier layer and every other consumer of it follows a layer later than it needed to. The getter's own work is one scalar field read on the host; the whole cost is the layer split it forces. `WaveConfig::deferSizeOutputs`, OFF BY DEFAULT, drops those getters from `top` before `makeExprLevels` so they stop being a reference, and puts them back before the last layer is built so they still run and still fill their output slot. The predicate counts users through the set reachable from the output node, not through `users()`. torch.export leaves dead `_operator.ge` / `_operator.le` shape guards behind once the asserts are stripped -- 428 of them on the ROO preproc graph -- and they appear in a value's `users()` while contributing to no level. Counting them refuses every candidate; with them filtered out, 255 of the graph's 297 top-level exprs qualify and 248 of those have an operand whose only other reachable user is a single consumer. It is off by default because on that graph it is a net loss, and the measurement is the point of the commit rather than the feature. See the test plan. Reviewed By: Yuhta Differential Revision: D118061301
oerling
force-pushed
the
export-D118061301
branch
from
September 5, 2026 19:52
ea81142 to
6e73d34
Compare
Summary: Pull Request resolved: facebookincubator#18861 Fuses `torch.ops.fb.batch_flip_and_truncate_sparse` on wave by splitting it, in `maybeReplace`, into `tw.flip_and_truncate_head` and `tw.flip_and_truncate_final` (device functions in `Flip.cuh`). The head processes one feature per block: it computes each row's post-truncation length, inclusive-scans those into new offsets and the original lengths into old offsets, and after an op barrier accumulates the per-feature output totals. The final partitions the concatenated output into fixed 8192-element tiles and grid-strides over them, binary-searching once per tile for the owning feature and row, then walking rows to the end of the tile. Fixed tiles keep every block's work equal even though rows span roughly 300 to 10000 elements after truncation. Each thread resolves the same per-row copy descriptor, which hoists the flip-mode branch out of the element loop and leaves coalesced writes and forward or reversed coalesced reads. It matches the eager kernel across the four flip modes (none, flip_before, flip_after, flip_before_and_after), per-feature `max_lengths`, `pad_value`, and `adaptive_max_len` with a per-row `adaptive_mask`. On the four flip instances of an internal preproc graph the final reaches 75-76% of an A100's nominal bandwidth, against 89% for a `cudaMemcpy` of the same size and 53-66% for the eager kernel's thread mapping. The final stages its per-feature constants (the offset, value and output pointers, row count, cap and mode) in shared memory rather than holding them in registers across the row loop. That is worth an occupancy step: the register version needs 64 registers and fits 4 blocks per SM, the shared version needs 48 and fits 5, and on a bandwidth-bound copy the extra block is worth about 3%. Sizing that array needs the feature count as a compile-time constant, which is where `Metadata::templateParamFuncs` comes in: a list of functions run on the node at codegen time to append extra template arguments to the device function call, empty for every other op. `Metadata::gridSizeSumsInputs` is the second new hook. An op whose work spans a whole tensor list was sized by its largest member, so the head got `ceil(768 / blockSize)` blocks no matter how many features it had, and most features queued onto three blocks; the flag switches that op to the sum instead. Neither hook changes any op that does not set it. Three wave-core fixes in `Compile.cpp` support this and correct general bugs. First, an elementwise producer of a `TensorList` element that a fused op reads is now materialized to that element's memory buffer in all modes, so a list built inside the op (for example values from an add) is written where the consumer reads it. Second, a fused op that reads such a list emits a barrier when a same-op element producer is not already behind one, independent of the input's `randomAccess` flag. Third, subgraph deduplication now hashes int-list (ScalarList) attributes, whose element values are baked into the generated code; without this two otherwise-identical ops that differ only in such an attribute (for example `adaptive_mask`) would share a wrong kernel. Reviewed By: Yuhta Differential Revision: D111315788
oerling
added a commit
to oerling/velox-1
that referenced
this pull request
Sep 6, 2026
…bookincubator#18867) Summary: Pull Request resolved: facebookincubator#18867 A metadata getter (`sym_size` / `sym_numel`) whose only role is to be a graph output still counts as a use of its operand when the partitioner builds its levels. `makeLevelsInner` increments the operand's producer refCount, and `makeCseBorder` turns anything above one into a border: the producer moves to its own earlier layer and every other consumer of it follows a layer later than it needed to. The getter's own work is one scalar field read on the host; the whole cost is the layer split it forces. `WaveConfig::deferSizeOutputs`, OFF BY DEFAULT, drops those getters from `top` before `makeExprLevels` so they stop being a reference, and puts them back before the last layer is built so they still run and still fill their output slot. The predicate counts users through the set reachable from the output node, not through `users()`. torch.export leaves dead `_operator.ge` / `_operator.le` shape guards behind once the asserts are stripped -- 428 of them on the ROO preproc graph -- and they appear in a value's `users()` while contributing to no level. Counting them refuses every candidate; with them filtered out, 255 of the graph's 297 top-level exprs qualify and 248 of those have an operand whose only other reachable user is a single consumer. It is off by default because on that graph it is a net loss, and the measurement is the point of the commit rather than the feature. See the test plan. Reviewed By: Yuhta Differential Revision: D118061301
oerling
force-pushed
the
export-D118061301
branch
from
September 6, 2026 01:10
6e73d34 to
3553ef4
Compare
…rnel wrote
Summary:
An elementwise op's operands are all marked `isRegister` by
`registerElementwise`, and `callNeedsBarrier` skips register operands on
the grounds that they flow inline rather than through memory. That holds
for a value fused into the expression tree; it does not hold for a leaf,
which is a load like any other. The test was that the two are the same
thing, and they are not, so no pure-elementwise op ever asked for a
barrier.
What that costs is visible whenever the leaf does not map index to index.
The ROO preproc graph fuses
(%3274) = where.self(...)
(%11769) = transpose.int(%3274, dim0=0, dim1=1)
(%11800) = mul(%11771, %11769)
into one kernel with nothing but `__syncthreads()` between the write and
the read. Through the transpose the element thread i of the `mul` wants
was written by a different block, so `%11800` came out part stale and
part uninitialized -- and only in the multi-block modes, which is what
made it look like a miscompile rather than a race.
The operand test `callNeedsBarrier` already had is now
`valueNeedsBarrier`, and `generateElementwise` runs it over each
subgraph's own memory leaves.
Differential Revision: D118061298
Summary:
Decomposes list-producing ops into per-tensor nodes, so each column is a value
the compiler can see: consumer counting, aliasing, CSE and cost-based block
shares all work per column rather than per bundle. On top of that, a column
folds its producer's gather into its own, so a run of subset ops
(`batch_flip_and_truncate_sparse`, `grouped_masked_select_jagged_1d`,
`group_length_guard_sparse`) becomes one gather over the original source with
nothing materialized in between.
`Metadata::decompose` plus `decomposeListOps` give one traversal that calls each
op's own rewrite rule, after CSE (cheaper on the bundled form) and before
partitioning, which the per-column nodes exist to inform.
A per-column gather is spelled as a chain of steps -- `kStepRange` for a flip,
`kStepSelect` for a select -- carried in `step_descs` / `step_scalars` with each
step's own tensors. A column absorbing its producer takes that chain whole: its
own step becomes outermost and the producer's steps append after it with their
tensor and scalar bases rebased. Which of the three getters runs then depends on
whether the composed chain contains a range ANYWHERE, not on what its outermost
step is. A select over a flip starts with a select and still needs the row walk,
because a range is defined per row.
Absorbing is refused unless the rows line up, which is a kernel constraint and
not a formality: `prepareRowSteps` prepares every step at one row number taken
from the outer walk, so row r of the outer has to be row r of the inner. Two
columns of one flip head hold the same lengths when they were built from the
same input lengths under the same cap -- and because a run of flips feeds each
column's new lengths to the next flip as that column's input lengths, that test
recurses down the run instead of bottoming out at tensor identity. Identity
alone decides only the last link and refuses everything deeper.
`group_length_guard_sparse`'s final stage is a range step written twice. For
output element idx of row r it computes
colOffsets[r - 1] + (idx - resultOffsets[r - 1])
which is what `rangeRowDescriptor` computes for a range at mode 0 with no cap
over those same two offset arrays. Saying so -- one
`tw.chain_gather_guard_column` per column in place of the list-form final --
lets a consumer that already has a range absorb the guard as a SUB-LEVEL of its
own rather than as a second step, so the intermediate is never written and the
depth does not grow. The row proof is the head's: element 1 of its output list
is the offsets a chain over those rows walks, so element 0 is the lengths they
were scanned from, which is what `chainRowsOf` now returns. That arm goes ahead
of the flip head's, because N + 2 is a multiple of three for a seven-column
guard and only the flip arm's target check would otherwise keep it out.
The guard's getter is its own registration rather than a reuse of either
neighbour, for two reasons and one trap. Every column of one guard keeps the
same count -- the per-row minimum summed over rows -- which the head already
has as a scalar, where the flip reads a per-column slice of cumulative lengths.
And the head writes row ENDS, as an inclusive scan, so `start_offsets` is 0
where a select's offsets stage passes 1. The trap is `sizeOrdinal`: the select's
registration names the step tensors, which for a guard column are the row
offsets, so the grid would be sized by the ROW count -- three orders of
magnitude short of the output on a real batch. Naming nothing is right here,
because the default kMax over the inputs picks the source, which is at least as
long as the output.
`WaveConfig::foldSharedChains` folds a producer into every consumer that can
absorb one rather than only into a sole reader. Two foldable consumers of one
column can never satisfy the sole-reader rule: whichever rewrite runs first sees
the other as an outside reader and declines, and the second then finds a gather
already reading the buffer, so the column is materialized and NEITHER folds.
Deciding that needs readers counted through `prim.ListPack` / `prim.ListUnpack`
to the ops that really read the data, since a pack is a live user that reads
nothing and a pack immediately unpacked is the identity.
It defaults ON, which only measurement could settle: a reader that folds beside
one that then declines for its own reasons leaves the buffer AND duplicates the
work. On the ROO preproc it takes every column it is offered and leaves none
behind -- see the numbers below. `sharedChain` now forces it OFF rather than
taking the default, so the unfolded arm stays covered.
`RangeRowDesc` carries its within-row fields in 32 bits. Only the two origins
are offsets into a whole column and need 64; `outLen` and the valid window are
bounded by the step's maxLen, and `dir` is +1 or -1. At four steps per chain the
struct is copied per row, so at 48 bytes it was the largest single contributor
to the chain gathers' register and spill footprint, and it is now 24. The window
bounds narrow through `narrowBound`, which saturates rather than wraps: a bound
past the field's range can only make the window empty, and an empty window stays
empty at the clamp, whereas a wrap could turn it back into a live one and read
out of range.
Differential Revision: D116667979
Summary:
Two commits that shape the same object -- the grid a step is laid out on and the
launches it runs as -- plus the occupancy work that made the second safe under
the single-ops debug pass.
== Balanced kernel launches ==
A step runs as one kernel launch capped at roughly one wave, so once it has as many ops as the wave has blocks, every op gets exactly one block and the step's makespan becomes the largest op's alone. On the ROO preproc graph node 30 step 0 has 272 ops against a 108-block wave, and one of them -- a 1.2M-element strided slice copy -- holds 87% of the work on that single block. Its 10.0M thread-block clocks against everything else's 1.4M leave the step 0.6% utilized and cost 8.0 ms. Sizing every op against a common per-block quantum and emitting the result as three packed launches takes the longest block to 0.60M clocks and the step to 1.4 ms, and the graph's kernel time from 23.1 ms to 17.1 ms.
`makeGrid` now hands its per-op costs and block caps to a launch layout that groups ops by the occupancy their shared memory allows, packs each group into wave-sized launches, and returns them as `StepVectors::segments`. The launch site loops over the segments, each with its own slice of the block array and its own shared memory. An op with no barrier may straddle two launches, since its blocks find their slice from `blockInOp` and `numBlocksInOp` wherever they run; an op with an `opBarrier` is shrunk to fit one launch rather than split, because that barrier waits for all of its blocks and only a cooperative launch keeps them co-resident. Every step also gets a `GridStats` measurement -- how far its grid is from a balanced, fully occupied wave, how many ops are stuck on one block, and how much occupancy one shared-memory-hungry op is costing the rest -- reported per step under the new `kGrid` trace bit and summarized in the performance report.
`partitionLaunches` is on by default; `--partition_launches=false` emits exactly the grid the graph does today. `--order_blocks_by_cost` is the cheaper half of the same idea, emitting a single launch's blocks in descending projected latency so the long poles start first; it is off by default, and measured no effect on ROO because the op it would have to help is the one pinned to a single block.
Two other passes change behaviour here, which is what moves the cat tests below. `breakDeviceSizedProducers` becomes `breakUnmeasurableProducers` and widens from `setsSizeOnDevice(producer)` to `setsSizeOnDevice(producer) || sizeNeedsReserve(producer)`: a concat lays its result out at one point and needs every operand's extent there, and an operand only its own reserve can measure is as far out of reach as one the device sizes. Both now end their kernel first. The cost is that a cat no longer fuses `tw.masked_select_final`, one more kernel boundary; the gain is that the concat can place its result at all, which takes `cat_alloc_group_test` from two placed concats to four and from two unplaceable operands to none.
Those are the assertions this diff updates in `ExecutorTest`. `catTest` drops the `multiKernel.fuses({cat, masked_select_final})` expectation for the reason above, and `catAllocGroupTest` moves to 4 groups, 9 members and 1 group that carves nothing -- `scaled`, whose operands are sized where the concat's own kernel computes them, so the group owns the result's buffer and lays the regions out but has no write to redirect. Both were correct-value failures: the numbers the graph produces never changed, only the plan the test asserts.
== Sizing a step's blocks by a per-block work quantum ==
The pro-rata split divides one wave's worth of blocks among a step's ops, so
what an op gets depends on how many other ops the step happens to have. A step
with about as many ops as a wave has blocks cannot give any of them more than
one: every op takes a block off the top whatever it costs, and the cooperative
trim then takes what is left from the tallest. On the ROO preproc graph one step
spends 210 of 441 blocks on copies holding 0.1% of the work and leaves a
1.2M-element op on two.
`sizeByQuantum` asks the other question -- how many blocks of a given duration
is this op's work worth -- which is a property of the op alone. The total is
rounded up to a whole wave and the surplus given to the ops with the most work
left per block; when the work exceeds `maxLaunchWaves` the quantum is bisected
LONGER rather than trimming the tallest ops, which is what avoids the trim.
Two things the sizing needs that were not previously recorded. `StepVectors`
now keeps `measuredUtil` -- mean block clocks over the slowest, the same figure
the per-step balance line reports -- and `measuredBlocks`, both from the
previous execution. A step whose blocks already finish together is left at the
width it had: the quantum says how much work there is, not whether the machine
has room to absorb it in parallel, and expanding an already-full step turns one
wave into several serialized launches running the same work. Only the
measurement can tell those apart. `GridStats::skew` cannot, because it is
derived from the same costs the sizing uses, so a mis-costed step looks
perfectly balanced to it -- which is exactly how the 1.2M-element op above
escaped the existing gate.
OFF BY DEFAULT, behind `WaveConfig::quantumGrid` / `--quantum_grid`, because it
is not a general win. Measured on the ROO graph at 256 rows, cg, three runs per
point, kernel time:
- with the flip-and-truncate fusion ON, it pays: paired with
`--auto_adjust_cost` it takes 32.5 ms to 24.8 ms. The two are superadditive
(-9.3% and -7.7% alone, -30% together) because the feedback is what stops a
frozen wrong cost from dominating a relative scheme.
- with `--tw_fat=false`, it COSTS 7-9%, at every `--min_block_us` from 5 to
100. It emits 34-60% more blocks (54k -> 72-86k) and none of them help:
four steps already at 90-98% utilization were widened to the wave budget for
+1.2 ms, and the extra blocks cost another +0.7 ms of interpretation filling
BlockInfo.
So this is a fix for a grid pathology that turning the flip fusion off removes
more cheaply, and the two should not be combined. It is landed default-off with
the measurements above rather than left out, because the pathology is real where
it applies and the sizing is the only thing that addresses it directly.
Also note both halves need a previous execution to have been measured -- the
quantum from `clocksPerCost`, the gate from `measuredUtil` -- so neither does
anything for a single-shot workload.
== Cooperative capacity from the driver ==
A cooperative launch may be no wider than the device holds co-resident, and the
width was being checked against an estimate. `CompiledModule::occupancy` and
`CompiledKernel::occupancy` expose
`cuOccupancyMaxActiveBlocksPerMultiprocessor`, and `CompositeKernel::occupancy`
memoizes it per dynamic-shared size, so the packer asks the driver a handful of
times per step rather than once per op.
`coopBlocksPerSM` bounds a cooperative launch by that figure, and only ever
downwards: `min(classOccupancy, driver)`, falling back to `classOccupancy` when
there is no driver to ask. `blocksPerSM` deliberately keeps the estimate. The
occupancy classes, the block budget and the BlockInfo reservation all derive
from it, so raising it there lets a launch outgrow a reservation sized from the
same number, and it repartitions a graph with no GPU behind it.
`exceedsCooperativeCapacity` forces a split when a single-launch plan would
exceed that width, whatever the skew gate says: leaving the step whole is not a
slower plan but a broken one. The one-block-per-op floor reaches this on its
own, since the trim in `makeGrid` cannot go below one block per op.
`debugSingleOps` is what exposed it. That pass declines to partition, which
reads as "no segments", but a step whose grid the normal pass laid out keeps
that pass's `sv.segments`. It now walks them, launching each at its own
`firstBlock` / `numBlocks` / `dynamicShared` / `cooperative` the way
`launchSegment` does, and steps one op at a time within a launch -- a barrier op
whole, everything else one block at a time, which is what puts a failure on a
block rather than on an op.
The same widening reaches `fb.offsets_to_lengths.default` and
`fb.offsets_to_ranges.default`, and there it was breaking kernels it did not
need to. Both declare a `reserveShape` and neither declared `shapeFromInput`,
so `sizeNeedsReserve` read them as extents nothing upstream can compute and
ended their kernel -- which cost the concats they feed their fusion, and is what
`MetaExecutorTest.offsetsOpsTest` was asserting against. Both reserves read
input 0 and nothing else: one returns its shape, the other `{numel, 1, 2}` from
it. `shapeFromInput` asks WHEN the extent becomes computable, not what it is --
its only two readers, `sizeNeedsReserve` and `hasReserveShapeInChain`, use it
purely to exclude a reserve whose answer is reachable from an input -- so both
are `shapeFromInput = 0` and both stay fusable.
== SCORING A SPLIT STEP ==
Once a step can run as several launches, the balance figure it is judged by has
to know that. `util` was `totalClocks / (maxClocks * numBlocks)` -- one
rectangle, the slowest block anywhere in the step times every block in it. The
launches run back to back on one stream, so the step's makespan is the SUM of
their maxima, and a block is only idle relative to the launch it actually ran
in. Scoring the whole step against one global max reads a short launch queued
behind a long one as imbalance: one step here reported 70% utilisation while
BOTH of its launches were above 75%.
`LaunchMeta` now carries the step's `segments`, which are contiguous ranges of
the same block array the DebugInfo is indexed by, so the report scores each
launch on its own:
util = totalClocks / SUM over launches of (max_L * blocks_L)
makespan = SUM over launches of max_L
A step that was not split collapses to exactly the old arithmetic and prints
exactly the old line, so no existing number moves. A split one gains a
per-launch breakdown, and each op line gains `in L0=n L1=m` -- how many of its
blocks landed in which launch, which is what says whether an op is holding one
up, since an op is only measured against the others it actually ran beside.
`waves` comes from `GridStats::targetBlocks`, already one wave at the occupancy
the step launched with, and says whether a launch is also several hardware
waves -- the same effect one level down, invisible in the block count.
== A GETTER IS NOT WORTH A BLOCK ==
`aten.sym_size.int` and `aten.sym_numel.default` fused into a kernel cost a
whole thread block that reads one field and exits. That block is charged
against its launch's slowest block, so a handful of them sink a step's balance
and no block count can fix it -- there is no width at which a no-op op
balances. `metadataGetterIsAlone` runs such a getter as a host-side shortcut
standalone instead, and stays fused when the getter is a subexpression of a
single fusable consumer, where the value is wanted on the device anyway.
The predicate already existed and was disabled at its two registration sites.
It is now on, behind `WaveConfig::metadataGetterStandalone` /
`--metadata_getter_standalone`. On the 1k ROO graph it takes about 232 blocks
per execution out of the grid -- node 36 step 0 goes from 800 blocks at 43.3%
utilisation to 540 at 65.6% -- and is worth 0.8% end to end.
Differential Revision: D116667980
Summary: A fused `aten.cat` or `aten.stack` of more than two operands used to fill its result with a chain of copies inside its own kernel, each copy waiting on the byte offset the one before it advanced, so one block walked the whole list. On one production preprocessing graph that leaves 357 operands on the serial chain, 213 of them in a single cat. The result is now allocated before its operands are produced and every operand is handed the region of the result it occupies, so the operands fill it independently. All ten wide concats of that graph now report no chained operand at all. The mechanism is a concat allocation group. The host lays the result out at the point placement settles on, and every operand ends in one of three states: written in place by the concat's own kernel, carved from the group's buffer by a launch of its own, or moved there by a copy op whose destination is that band. The first of those three states is not reachable by default. The parallel fill pushes every operand into a kernel of its own a step before the concat, which is what makes it "already placed" when the carve is decided, and an already-placed operand cannot be given a band -- so it is copied. For an operand whose producer could write the band directly, the pushdown is what creates the copy. `CompileCtx::concatOperandFusesInPlace` skips that pushdown when the producer exists and is neither placed nor a kernel input, the band is writable by it (`concatOperandCopyCause` is `kNone`), it is not `setsSizeOnDevice` and does not need a reserve, every one of its inputs is already materialized, and the concat is its sole consumer -- `prim.ListPack` and `aten.sym_size.int` do not count as other consumers. `placeInput` then fuses the producer into the concat's kernel and `reserveConcatOutput` binds it to its band, which is the "written in place by the concat's own kernel" state above. On the 1k ROO graph that takes 413.7 MB per execution out of the copies. Gated on `WaveConfig::concatOperandsInPlace`, off by default, with `--concat_operands_in_place`. Two things come with that which are diagnostics, and inert by themselves: - `ConcatCopyCause` and `concatCopyCauseText` name which of six reasons forced a copy, so the carve report says why rather than only that. - The byte tally in `materializeConcatGroup` is three-way, carved / in place / copied. Counting in-place as copied reports a graph that copies nothing as copying everything. Three more are not inert. They are ungated, and change what runs on every execution whether `concatOperandsInPlace` is set or not: - `ConcatInputInfo::writerId`, and `writtenValueId()` over it, name the value a launch actually writes. An operand reached through a `prim.ListUnpack` names the same tensor as the value that was packed, but it is the packed value the launch writes -- so that is what the group binds, counts occurrences of, claims, and asks the wait of. Binding the operand's own slot instead lets the unpack overwrite the band with the producer's own buffer. - One group per step its members are sized at, rather than one group per concat. They share the single result: whichever runs first allocates it and the rest carve into that one. The split is abandoned when any operand is not measurable at the earliest step, which leaves the single group that was there before. - `shapeRealized` distinguishes a shape that is known from one that is merely unrecorded, with an explicit `known` flag -- schedule points fill in as placement progresses, so a nullopt answer varied by call time. That decision is taken once. It used to be taken twice, by two passes that could disagree -- placement asked whether the operand's producer could write a band, the allocation group later asked whether a wave kernel launch writes it. On the ROO graph 195 operands of one cat answered differently and fell back to the serial fill. Placement now decides alone and the group applies it, so the two agree by construction. Deciding needs to know when each operand is written, so every value a launch writes records two points: the step that writes it, and the step from which the host can read its dimensions. They differ by one rule -- an output the device sizes, and any standalone's output, is measurable one step later than it is written. Laying a group out reads every operand's extent, carved or not, so a group with an operand the device sizes cannot be laid out until that count is back. A step carves its groups twice, before the host waits and again after, and the pre-wait pass asked only whether every member was sized -- never whether the group was one of those flagged as waiting. A group that carves nothing is what breaks on that: it still owns the result and lays every band out, and "every member sized" over an empty member list is vacuously true, so it always went in the early pass and measured a device-set extent from whatever the frame last held. On `catTest`/`catTest2` that is a concat joining a `masked_select`, whose whole result comes out wrong while the concat beside it, sized entirely on the host, comes out right. The flag was computed, ordered on and printed in the report all along; it is now also asked. Off the outermost axis an operand's band is pitched rather than a contiguous run. `__copy` gains a form that decomposes both sides independently, the source read at its own layout and the destination written through its strides, and `aten.clone.default` declares `mayWriteStrided` on the strength of it. A contiguous destination keeps the plain indexing, which is what nearly every copy hits. The flag has to reach the group as well: an operand an earlier kernel wrote is a formal in the concat's subgraph with no producer to ask, so `concatFootprint` asks the actual instead. Without that every copy crossing a kernel boundary was refused a band -- precisely the set the group exists to carve. A clone is also costed at 20 rather than the default 1, being a load and a store per element. An operand sized by a reserve function is measured by calling it, rather than refusing the whole concat: that is how its size is computed everywhere else, including by the concat's own reservation a moment later. On the ROO graph that is 251 of the 307 operands the pass used to decline. The bindings the reserve reads the frame through are captured by value when the footprint is built, which is the last point that still has the invocation. The plan is settled while the graph compiles rather than on its first execution. It is a function of the compiled grids alone, and it has to be: the carve pass reads its decisions back while a concat's kernel is generated, which is over before any execution starts. Its report, and the per-concat carve verdicts, are rendered there and printed by the first execution that runs with `kTiming` set -- they used to go to stdout while the trace bits were still clear, so the only account of why a concat carved nothing reached nobody. Finally, the mode no longer requires `freeIntermediates`. Callers set that after loading the graph, so every decision taken while the graph compiled saw it off and all eight concat groups came out empty: the pass looked like it was running and was not. With the freeing off a group's buffer is never released, so the next run finds it in the frame, resizes it and carves the same views out of it. The mode also runs on the multi-kernel grid now. It was gated to the cooperative one on the reasoning that only that grid has the step boundaries a lifetime is expressed in fixed before the first execution. The multi-kernel grid is settled by the same compilation, so what the gate should ask is whether the choice between the two has been made -- `isCg` holding a value -- not which way it went. Left at auto an op could run either, and a step index names nothing. Which grid the plan indexes is now one function, `allocGroupGrid`, rather than `cgGrid().empty() ? grid() : cgGrid()` written out at each of the three places that needed it: the footprint walk, the step count, and the executor pinning the op before the step loop. Those have to agree, and open-coding the choice three times is how they would stop agreeing. The schedule points are what actually blocked it. They were recorded only under the cooperative grid, so on multi-kernel `writtenPoint` was null for every value and every operand of every concat was refused with "no kernel launch writes it" -- 0 of 214 on the ROO graph, the mode running and carving nothing. Recording them whenever the grid choice is fixed is enough; the per-grid maps are cleared per variant by `newGrid`, and the graph-wide map is still written from one grid. On that graph at 256 rows the multi-kernel plan forms 127 groups over 1705 of 1871 allocated values -- 1578 fewer allocator calls per execution -- and places 21 concat operands, against 47 under the cooperative grid. One crash came with it. `makeConcatCopyDestination` mints the destination id on the main graph and mirrors it onto the copy node, which is normally a variant graph that has never seen the name. A concat outside every variant chain stays on the main graph, so all three grid variants ask for the destination there and the second collides with the first -- and the minted value itself, named `cat_copy_<id>`, collides with the copy output that takes its id. The destination is now named apart from the minted value and looked up before it is built, so one copy node per occurrence per graph is what exists. What is deliberately not done: the runtime single-block switch still moves an op to a grid with a different number of steps without the plan being rebuilt. That seam predates this and the cooperative path has always run with it. Pinning the grid closes it and costs about 16% on the graphs that switch, which is more than the grouping saves, so it is documented at `allocGroupEnabled` instead. == ON BY DEFAULT == `concatOperandsInPlace` now defaults to ON, and the predicate it gates is corrected so that it can be. It was landed off because it could never be measured end to end: the 1k run aborted with "too many blocks in cooperative launch" partway through, losing a concat before it died, and `allocGroupEnabled()` was cooperative-grid-only so cg0 formed no groups and could not be compared at all. The only evidence was a bytes-saved figure taken from a run that did not finish. Those blockers are gone -- the cooperative-launch abort is fixed a commit below, and the mode runs on the multi-kernel grid -- so the flag can be measured properly, and it pays: 2.0% off the 1k ROO graph, 453 MB per execution that stops being copied. `concatOperandFusesInPlace` had a real soundness gap that turning it on exposed. It asked `sizeNeedsReserve` of the producer, which is NOT the question the allocation group later asks: `sizeNeedsReserve` returns early for an elementwise producer and for a multi-output one before it ever looks at the reserve, while the group's `hasReserveShapeInChain` scans every `returnMeta` with no such exclusion. An elementwise producer carrying a real reserve therefore passed placement and then trapped in the group, which refuses the whole concat and leaves the copies placement had already minted writing bands nothing bound -- the hard error `checkNoOrphanCopies` exists to name. It now asks the group's question directly. Asking it of the direct producer alone is enough because the sole-absorbed-level rule already requires every one of its inputs to be materialized, so once the subgraph is extracted they are its inputs and the group's walk stops on them. The trade the default now takes: the pushdown buys a parallel fill and pays three times the memory traffic for it -- the producer writes its own buffer, then a copy reads that and writes the band -- plus a buffer live across the step boundary. Fusing in place writes the band once. It is a per-operand decision, so a graph of few narrow concats may not see the win this one does. Differential Revision: D118061300
…bookincubator#18867) Summary: Pull Request resolved: facebookincubator#18867 A metadata getter (`sym_size` / `sym_numel`) whose only role is to be a graph output still counts as a use of its operand when the partitioner builds its levels. `makeLevelsInner` increments the operand's producer refCount, and `makeCseBorder` turns anything above one into a border: the producer moves to its own earlier layer and every other consumer of it follows a layer later than it needed to. The getter's own work is one scalar field read on the host; the whole cost is the layer split it forces. `WaveConfig::deferSizeOutputs`, OFF BY DEFAULT, drops those getters from `top` before `makeExprLevels` so they stop being a reference, and puts them back before the last layer is built so they still run and still fill their output slot. The predicate counts users through the set reachable from the output node, not through `users()`. torch.export leaves dead `_operator.ge` / `_operator.le` shape guards behind once the asserts are stripped -- 428 of them on the ROO preproc graph -- and they appear in a value's `users()` while contributing to no level. Counting them refuses every candidate; with them filtered out, 255 of the graph's 297 top-level exprs qualify and 248 of those have an operand whose only other reachable user is a single consumer. It is off by default because on that graph it is a net loss, and the measurement is the point of the commit rather than the feature. See the test plan. Reviewed By: Yuhta Differential Revision: D118061301
oerling
force-pushed
the
export-D118061301
branch
from
September 6, 2026 02:26
3553ef4 to
d9a119b
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:
A metadata getter (
sym_size/sym_numel) whose only role is to be a graphoutput still counts as a use of its operand when the partitioner builds its
levels.
makeLevelsInnerincrements the operand's producer refCount, andmakeCseBorderturns anything above one into a border: the producer moves toits own earlier layer and every other consumer of it follows a layer later than
it needed to. The getter's own work is one scalar field read on the host; the
whole cost is the layer split it forces.
WaveConfig::deferSizeOutputs, OFF BY DEFAULT, drops those getters fromtopbefore
makeExprLevelsso they stop being a reference, and puts them backbefore the last layer is built so they still run and still fill their output
slot.
The predicate counts users through the set reachable from the output node, not
through
users(). torch.export leaves dead_operator.ge/_operator.leshape guards behind once the asserts are stripped -- 428 of them on the ROO
preproc graph -- and they appear in a value's
users()while contributing to nolevel. Counting them refuses every candidate; with them filtered out, 255 of the
graph's 297 top-level exprs qualify and 248 of those have an operand whose only
other reachable user is a single consumer.
It is off by default because on that graph it is a net loss, and the measurement
is the point of the commit rather than the feature. See the test plan.
Reviewed By: Yuhta
Differential Revision: D118061301