From d16496d1956c5ba609c367ec686eee67cbc93841 Mon Sep 17 00:00:00 2001 From: Takuya Nakaike Date: Fri, 7 Aug 2026 07:24:04 +0000 Subject: [PATCH 01/10] docs(inter-tile): add gather as a first-class op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add ktdp.inter_tile_gather as a full first-class delivery-op section (new §6), alongside consume / reduce / reduce_scatter. Gather assembles the producers' partials into one tensor by ordered concatenation along gather_dimension — no fold. Design choices: - General consumer set: one tile per group is a plain gather, all tiles is an all-gather. - Optional producer_dependency_per_consumer for a partial/segmented gather over a declared producer subset. - No combiner region and no identity operand (pure positional assembly, ordered by within-group local index). Also: add a full-IR multi-group example (§9.5.1, 128x8x12x64), update the scope/motivation/coverage tables and count-prose (four ops -> five, three patterns -> four), extend the relationship note, and resolve the gather open question (removed from the open-questions list). Renumber the sections after gather down by one to make room for §6, fixing internal cross-references. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Takuya Nakaike --- docs/inter-tile-communication.md | 362 ++++++++++++++++++++++++------- 1 file changed, 286 insertions(+), 76 deletions(-) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index bb8b74f..0daa17d 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -1,9 +1,9 @@ # Inter-tile communications in KTIR -**Scope:** Four ops — `ktdp.inter_tile_produce`, `ktdp.inter_tile_consume`, -`ktdp.inter_tile_reduce`, and `ktdp.inter_tile_reduce_scatter` — -that together cover all three inter-tile communication patterns: broadcast, -all-reduce, and reduce-scatter. +**Scope:** Five ops — `ktdp.inter_tile_produce`, `ktdp.inter_tile_consume`, +`ktdp.inter_tile_reduce`, `ktdp.inter_tile_reduce_scatter`, and +`ktdp.inter_tile_gather` — that together cover all four inter-tile +communication patterns: broadcast, all-reduce, reduce-scatter, and gather. --- @@ -14,7 +14,8 @@ Inter-tile communication involves three orthogonal concerns: 1. **Production** — which tiles contribute data and what they contribute. 2. **Delivery** — how the contributed data is delivered to the receiving tiles: pass-through unchanged (broadcast), folded by a combiner - (reduce), or folded then scattered (reduce-scatter). + (reduce), folded then scattered (reduce-scatter), or assembled by + ordered concatenation of the producers' partials (gather). 3. **Synchronization granularity** — whether each consumer tile waits for *all* producer tiles in its group to complete (full-barrier mode), or only for the specific producers whose data it requires (per-tile @@ -22,7 +23,7 @@ Inter-tile communication involves three orthogonal concerns: individual dependencies are satisfied, reducing stall time when producers finish at different times. -Separating production and delivery into a unified production op plus three +Separating production and delivery into a unified production op plus four delivery ops keeps each op single-purpose and enables any combination: | Pattern | Production op | Delivery op | @@ -30,6 +31,7 @@ delivery ops keeps each op single-purpose and enables any combination: | Broadcast | `ktdp.inter_tile_produce` | `ktdp.inter_tile_consume` | | Reduce | `ktdp.inter_tile_produce` | `ktdp.inter_tile_reduce` | | Reduce-scatter | `ktdp.inter_tile_produce` | `ktdp.inter_tile_reduce_scatter` | +| Gather | `ktdp.inter_tile_produce` | `ktdp.inter_tile_gather` | `ktdp.inter_tile_produce` returns a `!ktdp.tile_future` SSA value. The group set `#groups` is carried as a parameter of the future @@ -43,7 +45,7 @@ that future as its operand. The def-use edge from production to delivery encodes the happens-before ordering with no explicit barriers in the IR. The synchronization granularity — full-barrier or per-tile — is controlled by the `producer_dependency_per_consumer` attribute on the delivery op -(§3.1, §4.1, §5.1). Corresponding production and delivery ops are expected +(§3.1, §4.1, §5.1, §6.1). Corresponding production and delivery ops are expected to be adjacent in a single basic block to avoid dead locks. --- @@ -356,7 +358,92 @@ from their respective independent reductions. --- -## 6. Synchronization model +## 6. `ktdp.inter_tile_gather` — assembling delivery op + +### 6.1 Operand and attributes + +**Operand:** `!ktdp.tile_future` returned by +`ktdp.inter_tile_produce`. The `#groups` parameter supplies the group set; +there is no separate `groups` attribute (§1). + +**`consumer_tiles_per_group`** — tiles that receive the assembled tensor. +The set is unrestricted: selecting one tile per group is a plain gather +(one tile assembles the group's full tensor); selecting all tiles is an +all-gather (every tile in the group holds the same assembled tensor). + +**`gather_dimension`** (i64) — axis of the partial type `T_p` along which +the producers' partials are concatenated. Partials are placed along this +axis in ascending within-group local-index order. + +**`producer_dependency_per_consumer`** *(optional)* — identical in form to +§3.1. When absent, every producer tile in the group is assembled (complete +gather) and the consumer waits for all of them. When present, consumer `c` +assembles only the partials from its declared producer tiles, concatenated +in ascending local-index order — a partial (segmented) gather over the +declared subset. The subset and coverage invariants of §3.1 apply. So that +the single op result type is well-formed, every consumer's declared +producer set must have the same cardinality; the verifier rejects unequal +cardinalities. + +**No combiner region and no `identity` operand.** Unlike +`ktdp.inter_tile_reduce` / `ktdp.inter_tile_reduce_scatter`, gather performs +no folding — it assembles slices by position. Like `ktdp.inter_tile_consume` +it therefore carries no region and no identity operand. + +### 6.2 Type rules + +For each role `i`, `T_g_i` is `T_p_i` with the size along `gather_dimension` +multiplied by `K`, where `K` is the number of producers assembled per +consumer: `|producer tiles per group|` when +`producer_dependency_per_consumer` is absent, or the (common) cardinality +of the per-consumer producer set when it is present. The same +`gather_dimension` and the same `K` apply to all roles. + +**Per-tile slice.** The producer with within-group local index `l` (its +ascending position among the assembled producers) occupies slice +`[l*chunk : (l+1)*chunk]` along `gather_dimension` in the output, where +`chunk = T_p[gather_dimension]` is the producer partial's own size along +that axis. Unlike the reduce combiner — which may be applied in any order — +this placement is deterministic and requires no commutativity. + +### 6.3 Op signature + +```mlir +%gathered_1, ..., %gathered_N = ktdp.inter_tile_gather(%future) + consumer_tiles_per_group = , + gather_dimension = , + producer_dependency_per_consumer = // optional; default: all producers + : !ktdp.tile_future -> T_g_1, ..., T_g_N +``` + +No block is needed — the assembled value is an SSA result consumed by +ordinary function-scope SPMD code, exactly as with `ktdp.inter_tile_consume` +(§3.2). + +### 6.4 Result semantics + +The op produces N variadic SSA values, one per partial-tensor role. The +values are *per-tile-valued* — each consumer tile holds its assembled +result when the op completes. Every consumer tile in a group holds the same +assembled tensor (its group's ordered concatenation of producer partials); +tiles in different groups hold their own group's assembly. A one-tile +consumer set is a plain gather; an all-tiles consumer set is an all-gather. + +**Non-participating tiles.** Results are undefined for tiles not in +`consumer_tiles_per_group`, as in §4.5 and §5.5. + +**Multi-tensor (variadic) gather.** N ≥ 1 partials are supported, following +the same structure as §4.5 — each role is concatenated independently along +`gather_dimension`. + +Synchronization follows the shared model in §7: the def-use edge from +`ktdp.inter_tile_produce` orders production before delivery, and +`producer_dependency_per_consumer` selects full-barrier (absent) or per-tile +(present) waiting. + +--- + +## 7. Synchronization model No explicit barriers appear in the IR. The `!ktdp.tile_future` SSA value carries **per-tile availability signals** rather than a monolithic @@ -410,22 +497,23 @@ mode. --- -## 7. Coverage of inter-core communication patterns +## 8. Coverage of inter-core communication patterns -These four ops are sufficient to express all three inter-core +These five ops are sufficient to express all four inter-core communication patterns: -| Pattern | `inter_tile_produce` | Delivery op | `scatter_dimension` | +| Pattern | `inter_tile_produce` | Delivery op | Split/assemble dim | |---------|---------------------|-------------|---------------------| -| Broadcast | one producer tile per group | `inter_tile_consume` | absent | -| Reduce | all tiles per group | `inter_tile_reduce` | absent | -| Reduce-scatter | all tiles per group | `inter_tile_reduce_scatter` | present | +| Broadcast | one producer tile per group | `inter_tile_consume` | — | +| Reduce | all tiles per group | `inter_tile_reduce` | — | +| Reduce-scatter | all tiles per group | `inter_tile_reduce_scatter` | `scatter_dimension` | +| Gather | all tiles per group | `inter_tile_gather` | `gather_dimension` | --- -## 8. Pattern instantiation +## 9. Pattern instantiation -### 8.1 Broadcast → `inter_tile_produce` + `inter_tile_consume` +### 9.1 Broadcast → `inter_tile_produce` + `inter_tile_consume` ```mlir // 4 tiles, 1 group: tile 0 loads W; all 4 tiles compute. @@ -455,7 +543,7 @@ communication patterns: ktdp.store %C, ... ``` -### 8.2 Reduce → `inter_tile_produce` + `inter_tile_reduce` +### 9.2 Reduce → `inter_tile_produce` + `inter_tile_reduce` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -483,7 +571,7 @@ ktdp.store %C, ... } ``` -#### 8.2.1 Full IR — single-group reduce (96×64) +#### 9.2.1 Full IR — single-group reduce (96×64) **Layout and partitioning.** `A` and `B` are `tensor<96x64xf16>` in global memory. The kernel computes the column-wise sum of `A + B`, producing a @@ -600,7 +688,7 @@ module { } ``` -#### 8.2.2 Full IR — multi-group reduce (128×8×12×64) +#### 9.2.2 Full IR — multi-group reduce (128×8×12×64) **Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in global memory. The four axes have distinct roles: @@ -764,7 +852,7 @@ module { } ``` -### 8.3 Reduce-scatter → `inter_tile_produce` + `inter_tile_reduce_scatter` +### 9.3 Reduce-scatter → `inter_tile_produce` + `inter_tile_reduce_scatter` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -795,7 +883,7 @@ module { // Each tile holds a different slice — ownership explicit via SSA result. ``` -#### 8.3.1 Full IR — multi-group reduce-scatter (128×8×12×64) +#### 9.3.1 Full IR — multi-group reduce-scatter (128×8×12×64) **Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in global memory. The four axes have distinct roles: @@ -810,7 +898,7 @@ in global memory. The four axes have distinct roles: 32 tiles, 8 groups of 4. `g = t / 4`, `l = t % 4`. Tile `(g, l)` reads slice `[*, g, l*3 : l*3+3, *]` — shape `<128x1x3x64>`. The per-tile pipeline through to `%partial_4d` (shape `<128x1x1x64>`) is identical -to §8.2.2. +to §9.2.2. The op reduces dim 2 (within-group tile axis, size 1) and scatters dim 0 (128 / 4 = 32 rows per tile). Tile `(g, l)` ends up with rows @@ -961,9 +1049,9 @@ module { } ``` -### 8.4 Per-tile synchronization → `inter_tile_consume` with `producer_dependency_per_consumer` +### 9.4 Per-tile synchronization → `inter_tile_consume` with `producer_dependency_per_consumer` -#### 8.4.1 Per-tile pairing within a single group +#### 9.4.1 Per-tile pairing within a single group Four tiles per group: tiles `4g` and `4g+1` are producers, tiles `4g+2` and `4g+3` are consumers. Each consumer depends on its dedicated producer @@ -1005,7 +1093,7 @@ both producers finish. With it, each consumer stalls only for its own producer, halving the worst-case wait when the two producers finish at different times. -#### 8.4.2 Butterfly mirror exchange across multiple groups +#### 9.4.2 Butterfly mirror exchange across multiple groups Eight groups of 4 tiles; all 4 tiles in each group both produce and consume. Tile `c = 4g + l` waits only for its mirror partner @@ -1057,9 +1145,165 @@ eliminated). : !ktdp.tile_future, #all_groups> -> tensor<64xf16> ``` +### 9.5 Gather → `inter_tile_produce` + `inter_tile_gather` + +```mlir +// 4 tiles per group, 8 groups (32 tiles total). +#all_group_tiles = affine_set<(i)[g] : (i - 4*g >= 0, -i + 4*g + 3 >= 0)> +#group_consumer = affine_set<(i)[g] : (i - 4*g == 0)> +#all_groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +// All tiles contribute a partial slab. +%partial_future = ktdp.inter_tile_produce + producer_tiles_per_group = #all_group_tiles + : tensor<128x1x3x64xf16> -> !ktdp.tile_future, #all_groups> +{ + ^bb0(%gid: index): + ktdp.yield_partial %partial_4d : tensor<128x1x3x64xf16> +} + +// Gather along dim 2; one consumer per group (tile 4g) assembles the four +// 3-wide slabs. No combiner, no identity — placement is by within-group +// local index. gather_dimension = 2 → 3 * 4 = 12; consumer gets <128x1x12x64>. +%assembled = ktdp.inter_tile_gather(%partial_future) + consumer_tiles_per_group = #group_consumer, + gather_dimension = 2 + : !ktdp.tile_future, #all_groups> -> tensor<128x1x12x64xf16> +// The consumer holds the full assembled tensor — ownership via SSA result. +``` + +#### 9.5.1 Full IR — multi-group gather (128×8×12×64) + +**Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in +HBM. The four axes have distinct roles: + +- Dim 0 (size 128): preserved through this op. +- Dim 1 (size 8): the **group axis** — 8 groups. +- Dim 2 (size 12): the **gather axis** — within each group, 4 tiles each own + a 3-wide slab that gather concatenates back into the full 12. +- Dim 3 (size 64): vector / stick axis, preserved. + +32 tiles, 8 groups of 4. `g = t / 4`, `l = t % 4`. Tile `(g, l)` reads +slice `[*, g, l*3 : l*3+3, *]` — shape `<128x1x3x64>`. Each tile's partial +is the summed slab `A + B` over its own columns (no reduction across tiles). +Gather along dim 2 places tile `(g, l)`'s slab at columns `[l*3 : l*3+3]` of +the assembled `<128x1x12x64>`, which one consumer per group (tile `4g`) +writes back to `E[*, g, *, *]`. + +```mlir +#A_view_set = affine_set<(d0, d1, d2, d3) : + (d0 >= 0, -d0 + 127 >= 0, + d1 >= 0, -d1 + 7 >= 0, + d2 >= 0, -d2 + 11 >= 0, + d3 >= 0, -d3 + 63 >= 0)> + +#AB_tile_set = affine_set<(d0, d1, d2, d3) : + (d0 >= 0, -d0 + 127 >= 0, + d1 == 0, + d2 >= 0, -d2 + 2 >= 0, + d3 >= 0, -d3 + 63 >= 0)> + +// E access tile for the consumer: 128x1x12x64 anchored at [0, g, 0, 0]. +#E_tile_set = affine_set<(d0, d1, d2, d3) : + (d0 >= 0, -d0 + 127 >= 0, + d1 == 0, + d2 >= 0, -d2 + 11 >= 0, + d3 >= 0, -d3 + 63 >= 0)> + +#identity_4d = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> + +#group_tiles = affine_set<(i)[g] : (i - 4*g >= 0, -i + 4*g + 3 >= 0)> +#group_consumer = affine_set<(i)[g] : (i - 4*g == 0)> +#all_groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +module { + func.func @inter_tile_gather_multi_group() { + %c0 = arith.constant 0 : index + %c4 = arith.constant 4 : index + %col_slab = arith.constant 3 : index // 12 / 4 + + %A_start = arith.constant 1024 : index + %B_start = arith.constant 12583936 : index + %E_start = arith.constant 25166848 : index + + %A_view = ktdp.construct_memory_view %A_start, sizes: [128, 8, 12, 64], + strides: [6144, 768, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.spyre_memory_space + } : memref<128x8x12x64xf16> + %B_view = ktdp.construct_memory_view %B_start, sizes: [128, 8, 12, 64], + strides: [6144, 768, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.spyre_memory_space + } : memref<128x8x12x64xf16> + + // Per-tile compute (function-scope SPMD). + %t = ktdp.get_compute_tile_id : index + %g = arith.divui %t, %c4 : index + %l = arith.remui %t, %c4 : index + %col_anchor = arith.muli %l, %col_slab : index + + %A_access = ktdp.construct_access_tile %A_view[%c0, %g, %col_anchor, %c0] { + access_tile_set = #AB_tile_set, access_tile_order = #identity_4d + } : memref<128x8x12x64xf16> -> !ktdp.access_tile<128x1x3x64xindex> + %B_access = ktdp.construct_access_tile %B_view[%c0, %g, %col_anchor, %c0] { + access_tile_set = #AB_tile_set, access_tile_order = #identity_4d + } : memref<128x8x12x64xf16> -> !ktdp.access_tile<128x1x3x64xindex> + + %A_tile = ktdp.load %A_access + : !ktdp.access_tile<128x1x3x64xindex> -> tensor<128x1x3x64xf16> + %B_tile = ktdp.load %B_access + : !ktdp.access_tile<128x1x3x64xindex> -> tensor<128x1x3x64xf16> + + // No reduction — the summed slab is this tile's partial; gather will + // concatenate the four slabs along dim 2. + %AB_init = tensor.empty() : tensor<128x1x3x64xf16> + %partial_4d = linalg.add ins(%A_tile, %B_tile + : tensor<128x1x3x64xf16>, tensor<128x1x3x64xf16>) + outs(%AB_init : tensor<128x1x3x64xf16>) + -> tensor<128x1x3x64xf16> + + // Produce: every tile contributes its 3-wide slab to the future. + %partial_future = ktdp.inter_tile_produce + producer_tiles_per_group = #group_tiles + : tensor<128x1x3x64xf16> + -> !ktdp.tile_future, #all_groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %partial_4d : tensor<128x1x3x64xf16> + } + + // Gather dim 2: 4 producers x 3 = 12. One consumer (tile 4g) per group + // assembles the full <128x1x12x64>. No combiner region, no identity. + %assembled = ktdp.inter_tile_gather(%partial_future) + consumer_tiles_per_group = #group_consumer, + gather_dimension = 2 + : !ktdp.tile_future, #all_groups> + -> tensor<128x1x12x64xf16> + + // Post-gather: the consumer tile 4g writes its group's full slab to + // E[*, g, *, *]. Ownership is explicit via the def-use chain of %assembled. + %E_view = ktdp.construct_memory_view %E_start, sizes: [128, 8, 12, 64], + strides: [6144, 768, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.spyre_memory_space + } : memref<128x8x12x64xf16> + + %E_access = ktdp.construct_access_tile %E_view[%c0, %g, %c0, %c0] { + access_tile_set = #E_tile_set, access_tile_order = #identity_4d + } : memref<128x8x12x64xf16> -> !ktdp.access_tile<128x1x12x64xindex> + + ktdp.store %assembled, %E_access + : tensor<128x1x12x64xf16>, !ktdp.access_tile<128x1x12x64xindex> + + return + } +} +``` + --- -## 9. Relationship to existing ops +## 10. Relationship to existing ops | Existing op | Maps to in this design | |-------------|------------------------| @@ -1068,19 +1312,22 @@ eliminated). | `inter_tile_reduce` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce` — producer block removed from the reduction op | | `inter_tile_reduce_scatter` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce_scatter` — producer block removed from the reduction op | +`ktdp.inter_tile_gather` (§6) has no pre-existing counterpart — it is new +in this design; the earlier ops offered no ordered-concatenation delivery. + The `!ktdp.tile_future` type is shared across all ops; its `#groups` parameter carries the group set (§1). The previous `ktdp.inter_tile` single op (Approach B draft) is replaced -by this four-op design: `ktdp.inter_tile` had producer and optional +by this five-op design: `ktdp.inter_tile` had producer and optional combiner regions in one op with `consumer_tiles_per_group` determining -delivery mode. The four-op design makes production and delivery explicitly +delivery mode. The five-op design makes production and delivery explicitly separate ops, with the delivery mode determined by which delivery op is chosen rather than by attribute combinations. --- -## 10. Open questions +## 11. Open questions **Q1. `consumer_tiles` ⊄ `producer_tiles`?** A tile in `consumer_tiles_per_group` but not in `producer_tiles_per_group` @@ -1099,12 +1346,7 @@ the produce block, four combiner arguments yielding two values, two delivery op results. Each result follows the same per-op type rules independently. -**Q3. Gather pattern.** -A fourth pattern — all tiles produce, one tile consumes the concatenation -of all partials (no combining) — is not covered. See §11.2 for a -candidate op design. - -**Q4. Consume placement.** +**Q3. Consume placement.** Whether the verifier should enforce that delivery ops appear only inside a guard matching `consumer_tiles_per_group`, or whether this is left to lowering. If the union of consumer sets equals the set of all executing @@ -1113,9 +1355,9 @@ reaches the delivery op would be a verifier error. --- -## 11. Possible extensions +## 12. Possible extensions -### 11.1 Multiple delivery ops per future +### 12.1 Multiple delivery ops per future The current spec restricts a `!ktdp.tile_future` value to exactly one delivery op use. A natural extension would allow multiple delivery @@ -1157,44 +1399,7 @@ requires separate `ktdp.inter_tile_produce` ops for separate delivery concerns. If future use cases demonstrate a clear need for the shared production pattern, this restriction can be relaxed. -### 11.2 Gather — `ktdp.inter_tile_gather` (new op) - -**Why the current design cannot represent gather.** Gather assembles -each producer tile's partial into a single tensor by concatenating slices -in deterministic order — tile with local index `l` contributes the `l`-th -slice. This ordered concatenation cannot be expressed with any existing -delivery op: - -- `inter_tile_reduce` requires the combiner to be pure and - associative-commutative, so the scheduler is free to combine in any - order. Concatenation is not commutative (tile `0 ‖ 1 ≠ 1 ‖ 0`), so it - is not a valid combiner. -- `inter_tile_consume` with multiple producers requires all producers to - hold the same value (broadcast semantics); it cannot assemble different - slices from different producers into one tensor. - -**Candidate op.** - -```mlir -%gathered = ktdp.inter_tile_gather(%future) - consumer_tiles_per_group = , - gather_dimension = - : !ktdp.tile_future -> T_g -``` - -**Type rule.** `T_g` is `T_p` with the size along `gather_dimension` -multiplied by `|producer tiles per group|`. Producer tile with -within-group local index `l` (ascending position in the group) occupies -slice `[l*chunk : (l+1)*chunk]` along `gather_dimension` in the -assembled output, where `chunk` is the size of `T_p` along -`gather_dimension`. - -**Ordering.** Unlike the reduce combiner — which may be applied in any -order — the placement of each producer's partial in the output is -determined by its within-group local index. The ordering is -deterministic and does not require commutativity. - -### 11.3 Scatter — `ktdp.inter_tile_scatter` (new op) +### 12.2 Scatter — `ktdp.inter_tile_scatter` (new op) **Why the current design can represent scatter, but awkwardly.** Scatter sends a different slice of a single producer's tensor to each @@ -1229,7 +1434,7 @@ local index `l` receives slice `[l*chunk : (l+1)*chunk]` along tiles per group|`. The size along `scatter_dimension` must be divisible by the per-group consumer-tile count. -### 11.4 Summary of operation coverage +### 12.3 Summary of operation coverage | Pattern | Producers per group | Delivery op | Result per consumer | |---------|--------------------|-----------------------------|---------------------| @@ -1238,3 +1443,8 @@ by the per-group consumer-tile count. | Reduce-scatter | N | `inter_tile_reduce_scatter` | 1/N slice of reduced | | Gather | N | `inter_tile_gather` | full assembled tensor | | Scatter | 1 | `inter_tile_scatter` | 1/N slice of full | + +All rows except **Scatter** are first-class ops in this design (§2–§6). +Scatter remains a candidate extension (§12.2): it is expressible today via +`inter_tile_reduce_scatter` with a single producer, so it earns a dedicated +op only if the ergonomic problems above prove worth a new op. From 72ff9ebe478150d458f07c4d9e1dcbead171b2cc Mon Sep 17 00:00:00 2001 From: Takuya Nakaike Date: Fri, 7 Aug 2026 07:43:06 +0000 Subject: [PATCH 02/10] docs(inter-tile): add scatter as a first-class op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add ktdp.inter_tile_scatter as a full first-class delivery-op section (new §7). Scatter splits a single producer's tensor into ordered slices, one per consumer, along scatter_dimension — the dual of gather. Design choices: - Single producer per group (verifier rejects multi-producer selection). - No producer_dependency_per_consumer: with one producer the full-barrier and per-tile modes collapse, so the attribute is omitted. - No combiner region and no identity operand (pure positional partition). - Slice type is the honest T_p with one axis divided (<128x1x64> -> <32x1x64>). Also: add a full-IR multi-group example (§10.6.1, 128x8x64), update the scope/motivation/coverage tables and count-prose (five ops -> six, four patterns -> five), extend the relationship note, and resolve the consumer-need-not-be-producer open question for scatter (removed from the open-questions list). Flip the coverage-summary note to "all rows are first-class ops (§2–§7)". Renumber the Synchronization, Coverage, Pattern instantiation, Relationship, Open questions, and Extensions sections down by one to make room for §7. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Takuya Nakaike --- docs/inter-tile-communication.md | 368 ++++++++++++++++++++++++------- 1 file changed, 285 insertions(+), 83 deletions(-) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index 0daa17d..39dba02 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -1,9 +1,10 @@ # Inter-tile communications in KTIR -**Scope:** Five ops — `ktdp.inter_tile_produce`, `ktdp.inter_tile_consume`, -`ktdp.inter_tile_reduce`, `ktdp.inter_tile_reduce_scatter`, and -`ktdp.inter_tile_gather` — that together cover all four inter-tile -communication patterns: broadcast, all-reduce, reduce-scatter, and gather. +**Scope:** Six ops — `ktdp.inter_tile_produce`, `ktdp.inter_tile_consume`, +`ktdp.inter_tile_reduce`, `ktdp.inter_tile_reduce_scatter`, +`ktdp.inter_tile_gather`, and `ktdp.inter_tile_scatter` — that together +cover all five inter-tile communication patterns: broadcast, all-reduce, +reduce-scatter, gather, and scatter. --- @@ -14,8 +15,9 @@ Inter-tile communication involves three orthogonal concerns: 1. **Production** — which tiles contribute data and what they contribute. 2. **Delivery** — how the contributed data is delivered to the receiving tiles: pass-through unchanged (broadcast), folded by a combiner - (reduce), folded then scattered (reduce-scatter), or assembled by - ordered concatenation of the producers' partials (gather). + (reduce), folded then scattered (reduce-scatter), assembled by + ordered concatenation of the producers' partials (gather), or split + by ordered partition of a single producer's tensor (scatter). 3. **Synchronization granularity** — whether each consumer tile waits for *all* producer tiles in its group to complete (full-barrier mode), or only for the specific producers whose data it requires (per-tile @@ -23,7 +25,7 @@ Inter-tile communication involves three orthogonal concerns: individual dependencies are satisfied, reducing stall time when producers finish at different times. -Separating production and delivery into a unified production op plus four +Separating production and delivery into a unified production op plus five delivery ops keeps each op single-purpose and enables any combination: | Pattern | Production op | Delivery op | @@ -32,6 +34,7 @@ delivery ops keeps each op single-purpose and enables any combination: | Reduce | `ktdp.inter_tile_produce` | `ktdp.inter_tile_reduce` | | Reduce-scatter | `ktdp.inter_tile_produce` | `ktdp.inter_tile_reduce_scatter` | | Gather | `ktdp.inter_tile_produce` | `ktdp.inter_tile_gather` | +| Scatter | `ktdp.inter_tile_produce` | `ktdp.inter_tile_scatter` | `ktdp.inter_tile_produce` returns a `!ktdp.tile_future` SSA value. The group set `#groups` is carried as a parameter of the future @@ -436,14 +439,106 @@ consumer set is a plain gather; an all-tiles consumer set is an all-gather. the same structure as §4.5 — each role is concatenated independently along `gather_dimension`. -Synchronization follows the shared model in §7: the def-use edge from +Synchronization follows the shared model in §8: the def-use edge from `ktdp.inter_tile_produce` orders production before delivery, and `producer_dependency_per_consumer` selects full-barrier (absent) or per-tile (present) waiting. --- -## 7. Synchronization model +## 7. `ktdp.inter_tile_scatter` — splitting delivery op + +### 7.1 Operand and attributes + +**Operand:** `!ktdp.tile_future` returned by +`ktdp.inter_tile_produce`. The `#groups` parameter supplies the group set; +there is no separate `groups` attribute (§1). Each group has exactly one +producer tile per role — the tile that holds the whole tensor to be split. +The verifier rejects a `producer_tiles_per_group` that selects more than +one tile per group. + +**`consumer_tiles_per_group`** — tiles that receive the slices. The +producer's tensor is partitioned into `|consumer tiles per group|` equal +chunks along `scatter_dimension`, one chunk delivered to each consumer in +ascending within-group local-index order. + +**`scatter_dimension`** (i64) — axis of the producer type `T_p` along +which the tensor is split. Its size must be divisible by the number of +consumers per group. + +**No `producer_dependency_per_consumer`.** With a single producer per +group there is exactly one producer to wait for, so full-barrier and +per-tile synchronization collapse to the same thing; the attribute would +be degenerate and is therefore omitted. + +**No combiner region and no `identity` operand.** Like +`ktdp.inter_tile_consume` and `ktdp.inter_tile_gather`, scatter performs no +folding — it partitions one tensor by position. It carries no region and +no identity operand. + +### 7.2 Type rules + +For each role `i`, `T_s_i` is `T_p_i` with the size along +`scatter_dimension` divided by `|consumer tiles per group|`. The same +`scatter_dimension` and the same divisor apply to all roles. + +**Per-tile slice.** The consumer with within-group local index `l` (its +ascending position among the consumers) receives slice +`[l*chunk : (l+1)*chunk]` along `scatter_dimension`, where +`chunk = T_p[scatter_dimension] / |consumers|` is the per-consumer slice +size. This placement is deterministic and requires no commutativity. + +**Advantage over reduce-scatter.** `ktdp.inter_tile_reduce_scatter` can +express a bare split only by folding a single-producer axis with a +meaningless combiner and identity, and it shrinks a *within-group tile +axis* — forcing the partial to carry an artificial unit dimension. Scatter +splits the natural data axis directly, so the slice type is the honest +`T_p` with one axis divided (e.g. `<128x1x64>` → `<32x1x64>`) rather than +`<1x...>`. The op name and signature match the pattern. + +### 7.3 Op signature + +```mlir +%scattered_1, ..., %scattered_N = ktdp.inter_tile_scatter(%future) + consumer_tiles_per_group = , + scatter_dimension = + : !ktdp.tile_future -> T_s_1, ..., T_s_N +``` + +No block is needed — the slice value is an SSA result consumed by ordinary +function-scope SPMD code, exactly as with `ktdp.inter_tile_consume` (§3.2). + +### 7.4 Result semantics + +The op produces N variadic SSA values, one per tensor role. The values are +*per-tile-valued* — each consumer tile holds its own slice when the op +completes. Consumers in a group receive disjoint, ordered slices that +together tile the producer's tensor along `scatter_dimension`; tiles in +different groups partition their own group's producer tensor. + +**Non-participating tiles.** Results are undefined for tiles not in +`consumer_tiles_per_group`, as in §4.5 and §5.5. + +**Multi-tensor (variadic) scatter.** N ≥ 1 tensors are supported, following +the same structure as §4.5 — each role is split independently along +`scatter_dimension`. + +**Consumers need not be producers.** A consumer tile that does not appear +in `producer_tiles_per_group` simply receives its slice; unlike a partial +gather or reduce there is nothing for a non-producing consumer to +contribute or miss, so no coverage obligation arises. For a pure split the +consumer set is therefore unconstrained relative to the producer set — +resolving, for scatter, the general question of whether a consumer must +also be a producer. + +Synchronization follows the shared model in §8: the def-use edge from +`ktdp.inter_tile_produce` orders production before delivery. With a single +producer per group the wait is unconditional — every consumer waits for +that one producer — so there is no per-tile mode to select. + +--- + +## 8. Synchronization model No explicit barriers appear in the IR. The `!ktdp.tile_future` SSA value carries **per-tile availability signals** rather than a monolithic @@ -497,9 +592,9 @@ mode. --- -## 8. Coverage of inter-core communication patterns +## 9. Coverage of inter-core communication patterns -These five ops are sufficient to express all four inter-core +These six ops are sufficient to express all five inter-core communication patterns: | Pattern | `inter_tile_produce` | Delivery op | Split/assemble dim | @@ -508,12 +603,13 @@ communication patterns: | Reduce | all tiles per group | `inter_tile_reduce` | — | | Reduce-scatter | all tiles per group | `inter_tile_reduce_scatter` | `scatter_dimension` | | Gather | all tiles per group | `inter_tile_gather` | `gather_dimension` | +| Scatter | one producer tile per group | `inter_tile_scatter` | `scatter_dimension` | --- -## 9. Pattern instantiation +## 10. Pattern instantiation -### 9.1 Broadcast → `inter_tile_produce` + `inter_tile_consume` +### 10.1 Broadcast → `inter_tile_produce` + `inter_tile_consume` ```mlir // 4 tiles, 1 group: tile 0 loads W; all 4 tiles compute. @@ -543,7 +639,7 @@ communication patterns: ktdp.store %C, ... ``` -### 9.2 Reduce → `inter_tile_produce` + `inter_tile_reduce` +### 10.2 Reduce → `inter_tile_produce` + `inter_tile_reduce` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -571,7 +667,7 @@ ktdp.store %C, ... } ``` -#### 9.2.1 Full IR — single-group reduce (96×64) +#### 10.2.1 Full IR — single-group reduce (96×64) **Layout and partitioning.** `A` and `B` are `tensor<96x64xf16>` in global memory. The kernel computes the column-wise sum of `A + B`, producing a @@ -688,7 +784,7 @@ module { } ``` -#### 9.2.2 Full IR — multi-group reduce (128×8×12×64) +#### 10.2.2 Full IR — multi-group reduce (128×8×12×64) **Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in global memory. The four axes have distinct roles: @@ -852,7 +948,7 @@ module { } ``` -### 9.3 Reduce-scatter → `inter_tile_produce` + `inter_tile_reduce_scatter` +### 10.3 Reduce-scatter → `inter_tile_produce` + `inter_tile_reduce_scatter` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -883,7 +979,7 @@ module { // Each tile holds a different slice — ownership explicit via SSA result. ``` -#### 9.3.1 Full IR — multi-group reduce-scatter (128×8×12×64) +#### 10.3.1 Full IR — multi-group reduce-scatter (128×8×12×64) **Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in global memory. The four axes have distinct roles: @@ -898,7 +994,7 @@ in global memory. The four axes have distinct roles: 32 tiles, 8 groups of 4. `g = t / 4`, `l = t % 4`. Tile `(g, l)` reads slice `[*, g, l*3 : l*3+3, *]` — shape `<128x1x3x64>`. The per-tile pipeline through to `%partial_4d` (shape `<128x1x1x64>`) is identical -to §9.2.2. +to §10.2.2. The op reduces dim 2 (within-group tile axis, size 1) and scatters dim 0 (128 / 4 = 32 rows per tile). Tile `(g, l)` ends up with rows @@ -1049,9 +1145,9 @@ module { } ``` -### 9.4 Per-tile synchronization → `inter_tile_consume` with `producer_dependency_per_consumer` +### 10.4 Per-tile synchronization → `inter_tile_consume` with `producer_dependency_per_consumer` -#### 9.4.1 Per-tile pairing within a single group +#### 10.4.1 Per-tile pairing within a single group Four tiles per group: tiles `4g` and `4g+1` are producers, tiles `4g+2` and `4g+3` are consumers. Each consumer depends on its dedicated producer @@ -1093,7 +1189,7 @@ both producers finish. With it, each consumer stalls only for its own producer, halving the worst-case wait when the two producers finish at different times. -#### 9.4.2 Butterfly mirror exchange across multiple groups +#### 10.4.2 Butterfly mirror exchange across multiple groups Eight groups of 4 tiles; all 4 tiles in each group both produce and consume. Tile `c = 4g + l` waits only for its mirror partner @@ -1145,7 +1241,7 @@ eliminated). : !ktdp.tile_future, #all_groups> -> tensor<64xf16> ``` -### 9.5 Gather → `inter_tile_produce` + `inter_tile_gather` +### 10.5 Gather → `inter_tile_produce` + `inter_tile_gather` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -1172,7 +1268,7 @@ eliminated). // The consumer holds the full assembled tensor — ownership via SSA result. ``` -#### 9.5.1 Full IR — multi-group gather (128×8×12×64) +#### 10.5.1 Full IR — multi-group gather (128×8×12×64) **Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in HBM. The four axes have distinct roles: @@ -1301,9 +1397,160 @@ module { } ``` +### 10.6 Scatter → `inter_tile_produce` + `inter_tile_scatter` + +```mlir +// 4 tiles per group, 8 groups (32 tiles total). +#group_producer = affine_set<(i)[g] : (i - 4*g == 0)> +#all_group_tiles = affine_set<(i)[g] : (i - 4*g >= 0, -i + 4*g + 3 >= 0)> +#all_groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +// One producer per group (tile 4g) holds the whole 128-row tensor. +%whole_future = ktdp.inter_tile_produce + producer_tiles_per_group = #group_producer + : tensor<128x1x64xf16> -> !ktdp.tile_future, #all_groups> +{ + ^bb0(%gid: index): + ktdp.yield_partial %whole : tensor<128x1x64xf16> +} + +// Scatter along dim 0; the four tiles per group each receive one 32-row +// chunk. No combiner, no identity — placement is by within-group local +// index. scatter_dimension = 0 → 128 / 4 = 32; each consumer gets <32x1x64>. +%chunk = ktdp.inter_tile_scatter(%whole_future) + consumer_tiles_per_group = #all_group_tiles, + scatter_dimension = 0 + : !ktdp.tile_future, #all_groups> -> tensor<32x1x64xf16> +// Each consumer holds its own 32-row slice — ownership via SSA result. +``` + +#### 10.6.1 Full IR — multi-group scatter (128×8×64) + +**Layout and partitioning.** `A` and `B` are `tensor<128x8x64xf16>` in +HBM. The three axes have distinct roles: + +- Dim 0 (size 128): the **scatter axis** — the producer's 128 rows are + split into 4 chunks of 32, one per consumer tile. +- Dim 1 (size 8): the **group axis** — 8 groups. +- Dim 2 (size 64): vector / stick axis, preserved. + +32 tiles, 8 groups of 4. `g = t / 4`, `l = t % 4`. Per group, the single +producer tile `4g` reads its group's whole slab `A[*, g, *]` / +`B[*, g, *]` — shape `<128x1x64>` — sums them, and produces the summed +tensor. Scatter along dim 0 delivers chunk `[l*32 : l*32+32, *, *]` to the +consumer with within-group local index `l`, which writes its `<32x1x64>` +slice back to `E[l*32 : l*32+32, g, *]`. + +```mlir +#A_view_set = affine_set<(d0, d1, d2) : + (d0 >= 0, -d0 + 127 >= 0, + d1 >= 0, -d1 + 7 >= 0, + d2 >= 0, -d2 + 63 >= 0)> + +// Producer partial: the whole 128-row slab of one group, anchored at [0, g, 0]. +#whole_tile_set = affine_set<(d0, d1, d2) : + (d0 >= 0, -d0 + 127 >= 0, + d1 == 0, + d2 >= 0, -d2 + 63 >= 0)> + +// Consumer chunk: 32 rows, anchored at [l*32, g, 0]. +#chunk_tile_set = affine_set<(d0, d1, d2) : + (d0 >= 0, -d0 + 31 >= 0, + d1 == 0, + d2 >= 0, -d2 + 63 >= 0)> + +#identity_3d = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + +#group_producer = affine_set<(i)[g] : (i - 4*g == 0)> +#all_group_tiles = affine_set<(i)[g] : (i - 4*g >= 0, -i + 4*g + 3 >= 0)> +#all_groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +module { + func.func @inter_tile_scatter_multi_group() { + %c0 = arith.constant 0 : index + %c4 = arith.constant 4 : index + %row_chunk = arith.constant 32 : index // 128 / 4 + + %A_start = arith.constant 1024 : index + %B_start = arith.constant 1049600 : index + %E_start = arith.constant 2098176 : index + + %A_view = ktdp.construct_memory_view %A_start, sizes: [128, 8, 64], + strides: [512, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.spyre_memory_space + } : memref<128x8x64xf16> + %B_view = ktdp.construct_memory_view %B_start, sizes: [128, 8, 64], + strides: [512, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.spyre_memory_space + } : memref<128x8x64xf16> + + %t = ktdp.get_compute_tile_id : index + %g = arith.divui %t, %c4 : index + %l = arith.remui %t, %c4 : index + + // Produce: only the group's producer tile (4g) runs this region; it + // reads and sums its group's whole 128-row slab. + %whole_future = ktdp.inter_tile_produce + producer_tiles_per_group = #group_producer + : tensor<128x1x64xf16> + -> !ktdp.tile_future, #all_groups> + { + ^bb0(%gid: index): + %A_access = ktdp.construct_access_tile %A_view[%c0, %gid, %c0] { + access_tile_set = #whole_tile_set, access_tile_order = #identity_3d + } : memref<128x8x64xf16> -> !ktdp.access_tile<128x1x64xindex> + %B_access = ktdp.construct_access_tile %B_view[%c0, %gid, %c0] { + access_tile_set = #whole_tile_set, access_tile_order = #identity_3d + } : memref<128x8x64xf16> -> !ktdp.access_tile<128x1x64xindex> + + %A_tile = ktdp.load %A_access + : !ktdp.access_tile<128x1x64xindex> -> tensor<128x1x64xf16> + %B_tile = ktdp.load %B_access + : !ktdp.access_tile<128x1x64xindex> -> tensor<128x1x64xf16> + + %AB_init = tensor.empty() : tensor<128x1x64xf16> + %whole = linalg.add ins(%A_tile, %B_tile + : tensor<128x1x64xf16>, tensor<128x1x64xf16>) + outs(%AB_init : tensor<128x1x64xf16>) + -> tensor<128x1x64xf16> + ktdp.yield_partial %whole : tensor<128x1x64xf16> + } + + // Scatter dim 0: 128 / 4 = 32. Each of the four consumer tiles per group + // receives one 32-row chunk. No combiner region, no identity. + %chunk = ktdp.inter_tile_scatter(%whole_future) + consumer_tiles_per_group = #all_group_tiles, + scatter_dimension = 0 + : !ktdp.tile_future, #all_groups> + -> tensor<32x1x64xf16> + + // Post-scatter: consumer (g, l) writes its 32-row chunk to + // E[l*32 : l*32+32, g, *]. Ownership is explicit via the def-use chain. + %row_anchor = arith.muli %l, %row_chunk : index + + %E_view = ktdp.construct_memory_view %E_start, sizes: [128, 8, 64], + strides: [512, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.spyre_memory_space + } : memref<128x8x64xf16> + + %E_access = ktdp.construct_access_tile %E_view[%row_anchor, %g, %c0] { + access_tile_set = #chunk_tile_set, access_tile_order = #identity_3d + } : memref<128x8x64xf16> -> !ktdp.access_tile<32x1x64xindex> + + ktdp.store %chunk, %E_access + : tensor<32x1x64xf16>, !ktdp.access_tile<32x1x64xindex> + + return + } +} +``` + --- -## 10. Relationship to existing ops +## 11. Relationship to existing ops | Existing op | Maps to in this design | |-------------|------------------------| @@ -1312,33 +1559,26 @@ module { | `inter_tile_reduce` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce` — producer block removed from the reduction op | | `inter_tile_reduce_scatter` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce_scatter` — producer block removed from the reduction op | -`ktdp.inter_tile_gather` (§6) has no pre-existing counterpart — it is new -in this design; the earlier ops offered no ordered-concatenation delivery. +`ktdp.inter_tile_gather` (§6) and `ktdp.inter_tile_scatter` (§7) have no +pre-existing counterparts — both are new in this design; the earlier ops +offered neither ordered-concatenation delivery (gather) nor +single-producer ordered-partition delivery (scatter). The `!ktdp.tile_future` type is shared across all ops; its `#groups` parameter carries the group set (§1). The previous `ktdp.inter_tile` single op (Approach B draft) is replaced -by this five-op design: `ktdp.inter_tile` had producer and optional +by this six-op design: `ktdp.inter_tile` had producer and optional combiner regions in one op with `consumer_tiles_per_group` determining -delivery mode. The five-op design makes production and delivery explicitly +delivery mode. The six-op design makes production and delivery explicitly separate ops, with the delivery mode determined by which delivery op is chosen rather than by attribute combinations. --- -## 11. Open questions +## 12. Open questions -**Q1. `consumer_tiles` ⊄ `producer_tiles`?** -A tile in `consumer_tiles_per_group` but not in `producer_tiles_per_group` -has no partial to contribute — it would receive the reduced result without -participating in the reduction, implying an implicit identity contribution. -Whether this is allowed and whether identity injection is implicit or must -be explicit in the producer block is still open. The inverse direction -(producer tile outside the consumer set) is explicitly supported: see the -reduce-to-one and reduce-to-subset cases in §4.1. - -**Q2. Multi-tensor generalization.** +**Q1. Multi-tensor generalization.** The existing ops support variadic partials (N ≥ 1 for argmax-style reductions). The ops here should carry the same variadic structure. For N = 2 (argmax): two identities, two `ktdp.yield_partial` operands in @@ -1346,7 +1586,7 @@ the produce block, four combiner arguments yielding two values, two delivery op results. Each result follows the same per-op type rules independently. -**Q3. Consume placement.** +**Q2. Consume placement.** Whether the verifier should enforce that delivery ops appear only inside a guard matching `consumer_tiles_per_group`, or whether this is left to lowering. If the union of consumer sets equals the set of all executing @@ -1355,9 +1595,9 @@ reaches the delivery op would be a verifier error. --- -## 12. Possible extensions +## 13. Possible extensions -### 12.1 Multiple delivery ops per future +### 13.1 Multiple delivery ops per future The current spec restricts a `!ktdp.tile_future` value to exactly one delivery op use. A natural extension would allow multiple delivery @@ -1399,42 +1639,7 @@ requires separate `ktdp.inter_tile_produce` ops for separate delivery concerns. If future use cases demonstrate a clear need for the shared production pattern, this restriction can be relaxed. -### 12.2 Scatter — `ktdp.inter_tile_scatter` (new op) - -**Why the current design can represent scatter, but awkwardly.** -Scatter sends a different slice of a single producer's tensor to each -consumer tile. This is expressible today using `inter_tile_reduce_scatter` -with a single producer per group and a trivial identity combiner: the -reduce over one partial is the partial itself, and the scatter then -distributes slices to consumers. However, this encoding has three -ergonomic problems: - -1. The partial type must carry an artificial unit within-group tile axis - (e.g., `tensor<1 x 128 x 64>` instead of `tensor<128 x 64>`) because - `inter_tile_reduce_scatter` expects a reducible tile axis to collapse. -2. A combiner region must be provided even though it is never meaningfully - invoked; any pure combiner with the correct types satisfies the - verifier, but the requirement is misleading. -3. The op name `inter_tile_reduce_scatter` suggests reduction is taking - place, obscuring the actual intent. - -**Candidate op.** - -```mlir -%my_slice = ktdp.inter_tile_scatter(%future) - consumer_tiles_per_group = , - scatter_dimension = - : !ktdp.tile_future -> T_s -``` - -**Type rule.** `T_s` is `T_p` with the size along `scatter_dimension` -divided by `|consumer tiles per group|`. Consumer tile with within-group -local index `l` receives slice `[l*chunk : (l+1)*chunk]` along -`scatter_dimension`, where `chunk = T_p[scatter_dimension] / |consumer -tiles per group|`. The size along `scatter_dimension` must be divisible -by the per-group consumer-tile count. - -### 12.3 Summary of operation coverage +### 13.2 Summary of operation coverage | Pattern | Producers per group | Delivery op | Result per consumer | |---------|--------------------|-----------------------------|---------------------| @@ -1444,7 +1649,4 @@ by the per-group consumer-tile count. | Gather | N | `inter_tile_gather` | full assembled tensor | | Scatter | 1 | `inter_tile_scatter` | 1/N slice of full | -All rows except **Scatter** are first-class ops in this design (§2–§6). -Scatter remains a candidate extension (§12.2): it is expressible today via -`inter_tile_reduce_scatter` with a single producer, so it earns a dedicated -op only if the ergonomic problems above prove worth a new op. +All rows are first-class ops in this design (§2–§7). From 7e7179245251e4eae082fe1b3b7a51456a77ad11 Mon Sep 17 00:00:00 2001 From: Takuya Nakaike Date: Fri, 21 Aug 2026 07:09:39 +0000 Subject: [PATCH 03/10] docs(inter-tile): reorganize, add all_to_all, rename dim attributes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restructures the RFC per the review of #53 so that the machinery shared by the delivery ops is stated once instead of repeated per op. - Reorganization: §1 gives the three-property decomposition (combine × placement × cardinality) and a semantics matrix; §3 states the shared delivery semantics (operand, consumer set, local index, dependency attribute, combiner, synchronization, result); §4 the placement algebra and type rules; §5 the verification rules R1-R14. §6 then reduces each op to what is only true of it. - Adds ktdp.inter_tile_all_to_all as a first-class delivery op (§6.5, worked example §7.6). It is the permute placement -- split and concat in one step -- and is not decomposable into gather + scatter, since gather delivers the same tensor to every consumer and scatter permits only one producer. - Renames gather_dimension/scatter_dimension to gather_dim/scatter_dim, consistent across all delivery ops. Naming the split axis uniformly makes the separate all_to_all divisibility rule redundant: it is now covered by R9. - Moves the verifier's current state into its own non-normative section (§8) so §5 is purely normative, and adds the two rules the pass already enforces for reduce: R13 (consumer set subset of producer set) and R14 (reduce mode gate, C == P or |C| == 1). - Reserves "axis" for tensor and tile axes, and "rank" for a tensor's number of dimensions; the within-group local index is a tile's "position". Co-authored-by: Yu Chin Fabian Lim Co-Authored-By: Claude Opus 5 Signed-off-by: Takuya Nakaike --- docs/inter-tile-communication.md | 1472 +++++++++++++++++++----------- 1 file changed, 946 insertions(+), 526 deletions(-) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index 39dba02..88a51e7 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -1,23 +1,36 @@ # Inter-tile communications in KTIR -**Scope:** Six ops — `ktdp.inter_tile_produce`, `ktdp.inter_tile_consume`, -`ktdp.inter_tile_reduce`, `ktdp.inter_tile_reduce_scatter`, -`ktdp.inter_tile_gather`, and `ktdp.inter_tile_scatter` — that together -cover all five inter-tile communication patterns: broadcast, all-reduce, -reduce-scatter, gather, and scatter. +**Scope:** Seven ops — one production op, `ktdp.inter_tile_produce`, and +six delivery ops: `ktdp.inter_tile_consume`, `ktdp.inter_tile_reduce`, +`ktdp.inter_tile_reduce_scatter`, `ktdp.inter_tile_gather`, +`ktdp.inter_tile_all_to_all`, and `ktdp.inter_tile_scatter`. Together they +cover the six inter-core communication patterns: broadcast, all-reduce, +reduce-scatter, gather, all-to-all, and scatter. + +**Organization.** The delivery ops share almost all of their machinery. +That machinery is stated once, in §3 (operand, consumer set, local index, +dependency attribute, combiner, synchronization, result semantics), §4 +(type rules), and §5 (verification rules). §6 then defines each op by +what is *only* true of it. + +**Rule numbering.** The verification rules are numbered R1–R14 and +collected in §5, which is their single point of definition. They are cited +as `(Rn)` at the place the attribute they constrain is introduced, so a +citation like `(R1)` in §2.1 means "§5 states this rule; here is the +attribute it applies to." + +Sections are normative except §8 (implementation status) and §9 (observed +backend patterns). --- -## 1. Motivation +## 1. Motivation and the three-property decomposition -Inter-tile communication involves three orthogonal concerns: +Inter-tile communication involves three separate concerns: 1. **Production** — which tiles contribute data and what they contribute. -2. **Delivery** — how the contributed data is delivered to the receiving - tiles: pass-through unchanged (broadcast), folded by a combiner - (reduce), folded then scattered (reduce-scatter), assembled by - ordered concatenation of the producers' partials (gather), or split - by ordered partition of a single producer's tensor (scatter). +2. **Delivery** — how the contributed data is mapped onto the receiving + tiles' results. 3. **Synchronization granularity** — whether each consumer tile waits for *all* producer tiles in its group to complete (full-barrier mode), or only for the specific producers whose data it requires (per-tile @@ -25,16 +38,71 @@ Inter-tile communication involves three orthogonal concerns: individual dependencies are satisfied, reducing stall time when producers finish at different times. -Separating production and delivery into a unified production op plus five -delivery ops keeps each op single-purpose and enables any combination: - -| Pattern | Production op | Delivery op | -|---------|--------------|-------------| -| Broadcast | `ktdp.inter_tile_produce` | `ktdp.inter_tile_consume` | -| Reduce | `ktdp.inter_tile_produce` | `ktdp.inter_tile_reduce` | -| Reduce-scatter | `ktdp.inter_tile_produce` | `ktdp.inter_tile_reduce_scatter` | -| Gather | `ktdp.inter_tile_produce` | `ktdp.inter_tile_gather` | -| Scatter | `ktdp.inter_tile_produce` | `ktdp.inter_tile_scatter` | +Separating production from delivery keeps each op single-purpose and +enables any combination: one production op plus a choice of delivery op. +The pairing is **one-to-one** — a production op is consumed by exactly one +delivery op (R2, §2.3). A pattern needing two deliveries therefore needs +two `ktdp.inter_tile_produce` ops; §10.3 discusses relaxing that. + +### 1.1 Semantics matrix + +The six delivery ops differ in exactly three independent properties. +"Property" rather than "axis" throughout: in this document *axis* always +means a tensor or tile axis. + +- **combine** — `none` | `fold` (combiner region + identity operand). +- **placement** — how producer contributions map onto consumer results: + `replicate` | `concat` | `permute` | `split`. +- **cardinality** — producer tiles per group × consumer tiles per group. + +| Op | combine | placement | producers/grp | consumers/grp | dim attrs | region | identity | +|---|---|---|---|---|---|---|---| +| `consume` | none | replicate | 1 | free | — | — | — | +| `reduce` | fold | replicate | all | free | — | combiner | yes | +| `reduce_scatter` | fold | split | all | free | `scatter_dim` | combiner | yes | +| `gather` | none | concat | all | free | `gather_dim` | — | — | +| `all_to_all` | none | permute | all | all | `scatter_dim`, `gather_dim` | — | — | +| `scatter` | none | split | 1 | free | `scatter_dim` | — | — | + +`all_to_all` is listed before `scatter` because it shares the +all-producers cardinality cell with `gather` and `reduce_scatter`, and +because its relationship to the two copy-only placements is structural: +**permute = split + concat in one step**, which is why it carries both dim +attributes and no new ones. + +Three things this matrix makes visible: + +- **`placement` takes only four values.** The per-op type rules are four + formulas (§4), not six. +- **The empty cells are principled.** `none` × `replicate` with all + producers is undefined (which producer's value wins?), and `fold` × + `concat` / `fold` × `permute` is meaningless (fold what, then shuffle + what?). +- **`all_to_all` is the fourth placement value, not a special case.** + +### 1.2 Pattern coverage + +The "all-" prefixed patterns are not separate ops: an op whose +`consumers/grp` cell is `free` already subsumes its all-tiles case by +widening `consumer_tiles_per_group`. The consumer set is therefore a +column here, since it is what distinguishes gather from all-gather and +all-to-all from scatter. + +| Pattern | Producers/grp | Consumers/grp | Delivery op | Result per consumer | +|---------|---------------|---------------|-------------|---------------------| +| Broadcast | 1 | free | `inter_tile_consume` | full copy | +| Reduce-to-one | all | 1 | `inter_tile_reduce` | fully reduced | +| All-reduce | all | all | `inter_tile_reduce` | fully reduced | +| Reduce-scatter | all | free | `inter_tile_reduce_scatter` | 1/C slice of reduced | +| Gather | all | 1 | `inter_tile_gather` | full assembled tensor | +| All-gather | all | all | `inter_tile_gather` | full assembled tensor | +| All-to-all | all | all | `inter_tile_all_to_all` | one slice from every producer | +| Scatter | 1 | free | `inter_tile_scatter` | 1/C slice of full | + +`inter_tile_scatter` and `inter_tile_consume` have no natural "all-" +variant: R8 (§5) limits them to one producer per group. + +### 1.3 The future value `ktdp.inter_tile_produce` returns a `!ktdp.tile_future` SSA value. The group set `#groups` is carried as a parameter of the future @@ -43,13 +111,14 @@ production and delivery ops. Each delivery op therefore infers the groups from its operand type, and a group mismatch between production and delivery is inexpressible — the def-use edge already requires the operand type to equal the result type, so the type system rejects it structurally -rather than a verifier catching it after the fact. Each delivery op takes -that future as its operand. The def-use edge from production to delivery -encodes the happens-before ordering with no explicit barriers in the IR. -The synchronization granularity — full-barrier or per-tile — is controlled -by the `producer_dependency_per_consumer` attribute on the delivery op -(§3.1, §4.1, §5.1, §6.1). Corresponding production and delivery ops are expected -to be adjacent in a single basic block to avoid dead locks. +rather than a verifier catching it after the fact. + +The def-use edge from production to delivery encodes the happens-before +ordering with no explicit barriers in the IR. The synchronization +granularity — full-barrier or per-tile — is controlled by the +`producer_dependency_per_consumer` attribute on the delivery op (§3.4). +Corresponding production and delivery ops are expected to be adjacent in +a single basic block to avoid deadlocks. --- @@ -63,15 +132,14 @@ id) and one symbol (`g`, the group index). For example, `affine_set<(i)[g] : (i - 4*g >= 0, -i + 4*g + 3 >= 0)>` selects tile ids `4g .. 4g+3` for any group index `g`. An enumerated form (a list of tile-id lists) is supported as a fallback when per-group membership is irregular. -- Broadcast: selects exactly one tile per group. -- Reduce / reduce-scatter: selects all tiles in the group. +Which cardinality each delivery op requires of this set is given by the +`producers/grp` column of §1.1 and enforced by R8 (§5). -**Disjointness invariant.** For any two distinct group indices `g_1 != g_2` -in `groups`, `producer_tiles_per_group(g_1)` and +**Disjointness invariant (R1).** For any two distinct group indices +`g_1 != g_2` in `groups`, `producer_tiles_per_group(g_1)` and `producer_tiles_per_group(g_2)` must be disjoint. Every producing tile is -in exactly one group. The verifier enforces this. The motivation is -unambiguous group membership: each tile contributes to exactly one group's -production. +in exactly one group. The motivation is unambiguous group membership: each +tile contributes to exactly one group's production. **`groups`** — affine integer set defining the range of valid group indices. For example, `affine_set<(g) : (g >= 0, -g + 7 >= 0)>` defines 8 groups, @@ -79,7 +147,7 @@ indexed `0..7`. It bounds the range of the `g` symbol used by `producer_tiles_per_group`. This set is **not** a standalone attribute: it is carried as the trailing parameter of the result `!ktdp.tile_future<..., #groups>` type, and every delivery op infers it -from its operand type (§1). +from its operand type (§1.3). ### 2.2 Producer region @@ -114,7 +182,11 @@ with the per-tile compute that produced `%my_partial` living at function scope (where it is naturally executed by every tile under SPMD). Richer bodies are allowed when the contribution is awkward to hoist — e.g., compute that depends only on `%gid`, or visibly local "contribution -preparation" the author wants to keep adjacent to the op. +preparation" the author wants to keep adjacent to the op. A +single-producer-per-group op (`consume`, `scatter`) is the case where a +richer body is normally *required*: the loads that feed the partial must +not run on the group's non-producing tiles, so they belong inside the +region rather than at function scope (§7.7.1). ### 2.3 Op signature @@ -132,61 +204,90 @@ preparation" the author wants to keep adjacent to the op. signals. Each producer tile's contribution becomes independently observable the moment that tile executes `ktdp.yield_partial`. -**Single-use invariant.** `%future` must have exactly one use — the -single delivery op that consumes it. A second use is a verifier error. -If two delivery ops need to communicate with the same set of producers, -they must each have their own `ktdp.inter_tile_produce`. +**Single-use invariant (R2).** `%future` must have exactly one use — the +single delivery op that consumes it. If two delivery ops need to +communicate with the same set of producers, they must each have their own +`ktdp.inter_tile_produce`. §10.2 discusses relaxing this. --- -## 3. `ktdp.inter_tile_consume` — plain delivery op (broadcast) +## 3. Shared delivery semantics + +Everything in this section holds for **every** delivery op. §6 states +only per-op deltas; where §6 is silent, this section governs. + +### 3.1 Notation -### 3.1 Operand and attributes +| Symbol | Meaning | +|---|---| +| `T_p_i` | the partial type of role `i`, as yielded by `ktdp.yield_partial` | +| `N` | number of partial-tensor roles (variadic arity), `N >= 1` | +| `P` | number of producer tiles a given consumer assembles from / waits on | +| `C` | number of consumer tiles per group, `\|consumer_tiles_per_group(g)\|` | +| `l` | within-group local index (§3.3) | + +`P` is `|producer_tiles_per_group(g)|` when +`producer_dependency_per_consumer` is absent, and the (common, by R6) +cardinality of the per-consumer producer set when it is present. + +### 3.2 Operand and consumer set **Operand:** `!ktdp.tile_future` — the future returned by the corresponding `ktdp.inter_tile_produce`. The def-use edge is the ordering constraint, and the `#groups` parameter of this type supplies the group set. There is no separate `groups` attribute; a group -mismatch with production is inexpressible (§1). - -**`consumer_tiles_per_group`** — tiles that receive the delivered value -per group. - -**`producer_dependency_per_consumer`** *(optional)* — affine integer set -`(p)[c, g]` over producer tile IDs `p`, parameterized by consumer tile -`c` and group index `g`. For consumer tile `c` in group `g`, only the -producer tiles satisfying this set are waited on and received by the -delivery op. If absent, the consumer waits for all producer tiles in the -group (full-barrier semantics). - -**Verifier invariants when present:** - -1. **Subset check.** The declared set must be a subset of - `producer_tiles_per_group`. Referencing a non-producer tile is a - verifier error. - ``` - { p | ∃ c, g : producer_dependency_per_consumer(p)[c, g] } - ⊆ - { p | ∃ g : p ∈ producer_tiles_per_group(g) } - ``` -2. **Coverage check.** For every group `g` and every producer tile `p` - in `producer_tiles_per_group(g)`, at least one consumer tile `c` in - `consumer_tiles_per_group(g)` must satisfy - `producer_dependency_per_consumer(p)[c, g]`. A producer tile absent - from the union is a verifier error — it would yield a value that no - consumer ever reads, risking a deadlock in push-based lowerings. - ``` - ∀ g, ∀ p ∈ producer_tiles_per_group(g) : - ∃ c ∈ consumer_tiles_per_group(g) : - producer_dependency_per_consumer(p)[c, g] - ``` - -Because a `%future` has exactly one delivery op (§2.3), these invariants -are checked against that single delivery op. Note that folding `groups` -into the future type removes only the group-match check: the subset and -coverage checks above reference `producer_tiles_per_group`, which lives on -the producing op, so the verifier still reads it across the def-use edge -(that set is not carried in the type). +mismatch with production is inexpressible (§1.3). + +**`consumer_tiles_per_group`** — affine integer set, of the same +`(i)[g]` form as `producer_tiles_per_group`, selecting the tiles that +receive a result per group. Its permitted cardinality per op is the +`consumers/grp` column of §1.1. + +The operations that use a delivery op's result are performed only by the +tiles in `consumer_tiles_per_group`. This ownership constraint is carried +by the def-use chain from the result: any use of the result is reachable +only by consumer tiles. No block is needed on any delivery op for +post-delivery computation — that is ordinary function-scope SPMD code +consuming the SSA value. + +### 3.3 Within-group local index — normative + +**`l` is a tile's position, counting from 0 in ascending tile-id order, +among the relevant set within its group** — the producer set for `concat` placement (and for +`permute`'s `gather_dim`), the consumer set for `split` placement (and +for `permute`'s `scatter_dim`). + +This definition is what makes ordered placement well-defined. Without it, +concatenation and split orders are pinned down only by contiguous-tile-id +coincidence and break silently under non-monotone tile assignments. Every +"ascending local-index order" in this document means exactly this +position — never a tile id, and never an offset in the textual order of an +enumerated set. ("Position" rather than "rank": in this document *rank* +always means a tensor's number of dimensions.) + +### 3.4 `producer_dependency_per_consumer` *(optional)* + +Affine integer set `(p)[c, g]` over producer tile IDs `p`, parameterized +by consumer tile `c` and group index `g`. For consumer tile `c` in group +`g`, only the producer tiles satisfying this set are waited on and +received. If absent, the consumer waits on and receives from **all** +producer tiles in the group (full-barrier semantics). + +The attribute has two distinct effects, depending on placement: + +- For `replicate` placement it selects *which* producer a consumer reads + and *when* it unblocks — a synchronization refinement only. +- For `concat` and `permute` placements it additionally narrows the set + of contributions assembled, yielding a partial (segmented) gather over + the declared subset; `P` and hence the result type follow from it. +- For `fold` placement it makes the result a partial reduction over the + declared subset: contributions from the remaining producers are treated + as the identity. + +`scatter` is the one op that does not accept the attribute (§6.6). + +Its verification obligations are R3–R7 (§5); they are stated there and +not restated here. Not every symbol needs to appear in a given instantiation: @@ -201,79 +302,267 @@ Not every symbol needs to appear in a given instantiation: `(p)[c, g] : (p + c - 8*g - 3 == 0)`, where the sum `p + c` differs for each group. -### 3.2 Semantics +### 3.5 Combiner region and `identity` — `fold` placement only -No combining occurs. The value produced by the producer tile(s) in each -group is delivered unchanged to every consumer tile in that group. -When `producer_tiles_per_group` selects one tile per group, every -consumer tile in that group receives the same value (broadcast semantics). +The two `fold` ops (`reduce`, `reduce_scatter`) carry a combiner region +and an `identity` operand list. The four copy-only ops carry neither: +they place contributions by position, so there is nothing to fold and no +identity element to supply. -The operations that use `%result` are performed only by the tiles specified -by `consumer_tiles_per_group`. This ownership constraint is enforced by -the def-use chain from `%result`: any use of `%result` is reachable only -by consumer tiles. +**Region.** A single block receiving `2N` arguments — +`%lhs_1, ..., %lhs_N, %rhs_1, ..., %rhs_N` with each `%lhs_i` and +`%rhs_i` of type `T_p_i` — terminated by +`ktdp.yield_reduced %val_1, ..., %val_N : T_p_1, ..., T_p_N`. -No block is needed — post-delivery computation is ordinary function-scope -SPMD code that uses the SSA value. +**Purity (R10).** The combiner must be pure — no memory effects, no calls +to side-effecting ops. Pure tensor ops (`tensor.empty`, `linalg` on +tensors, `arith.*`) are allowed. -### 3.3 Op signature +**Combine ordering.** The associative-commutative contract is by user +agreement; the scheduler is free to combine in tree, ring, linear, or any +hardware-native topology. Different groups' reductions are independent +and may be scheduled in parallel. This freedom is what distinguishes +`fold` from the copy-only placements, whose ordered placement by `l` +(§3.3) is deterministic and requires no commutativity. -```mlir -%result_1, ..., %result_N = ktdp.inter_tile_consume(%future) - consumer_tiles_per_group = , - producer_dependency_per_consumer = // optional; default: all producers - : !ktdp.tile_future -> T_p_1, ..., T_p_N -``` +**`identity` (R11).** `N` variadic SSA operands, one per role. Each +identity tensor's shape and element type must match the corresponding +partial type `T_p_i` — *not* the result type. The identities are hoisted +before the op and shared across all groups and all tiles. Combining any +identity with its corresponding partial yields that partial. + +### 3.6 Synchronization model + +No explicit barriers appear in the IR. The +`!ktdp.tile_future` SSA value carries **per-tile +availability signals** rather than a monolithic group barrier: + +1. Each producer tile's contribution becomes independently observable as + soon as that tile executes `ktdp.yield_partial` in the production + block. +2. A delivery op cannot use a producer tile's contribution until that + tile's signal is set in `%future`. +3. The producer tiles a given consumer tile waits for are declared by + `producer_dependency_per_consumer` (§3.4): + + - **Absent (default) — full-barrier mode:** consumer tile `c` in group + `g` waits for every producer tile in `producer_tiles_per_group(g)` + before the delivery op executes. This maps directly to a hardware + group barrier and preserves the simplest safety guarantee. + - **Present — per-tile mode:** consumer tile `c` waits only for the + producer tiles `p` satisfying + `producer_dependency_per_consumer(p)[c, g]`. The consumer unblocks + as soon as those specific tiles have completed, without waiting for + unrelated producers. Different consumer tiles may declare different + dependency sets, enabling fine-grained producer–consumer pipelining. + +A multi-producer wait is therefore a **per-consumer AND-join over +existing per-tile signals**, not a new primitive. This is why the +all-producers ops (`reduce`, `reduce_scatter`, `gather`, `all_to_all`) +introduce no synchronization machinery beyond what a single-producer op +already needs: they differ only in how many signals the join covers. + +In SPMD KTIR, a tile cannot observe other tiles' partials except through +a dialect-defined boundary. The `ktdp.inter_tile_produce` block is that +boundary — it names the per-tile contribution and exposes it via +`%future`. The delivery op's result tensor is an SSA value that cannot +materialize until the declared dependencies are satisfied; standard MLIR +dataflow ordering applies. + +Lowering inserts target-specific hardware synchronization: a group +barrier for full-barrier mode, and point-to-point ready/wait signals for +per-tile mode. + +### 3.7 Result semantics + +Every delivery op produces `N` variadic SSA values, one per +partial-tensor role. The values are **per-tile-valued**: each consumer +tile holds its own result value when the op completes. Whether tiles in +the same group hold the *same* value is a property of the placement: + +| placement | tiles in one group hold | tiles in different groups hold | +|---|---|---| +| `replicate` | the same value | their own group's value | +| `concat` | the same assembled tensor | their own group's assembly | +| `split` | disjoint ordered slices that tile the whole | slices of their own group's tensor | +| `permute` | different assemblies (one slice per producer) | their own group's exchange | + +**Non-participating tiles.** Results are undefined for tiles not in +`consumer_tiles_per_group`. + +**Multi-tensor (variadic) delivery.** `N >= 1` roles are supported by +every op, and all roles share the same attributes (`scatter_dim`, +`gather_dim`, `P`, `C`) — only the types differ. Argmax-style reductions, +where each contribution is a correlated tuple of tensors (values, +indices), use `N = 2`: two identities, two yielded partials, four +combiner arguments yielding two combined values, two op results. Each +role's result type follows the §4 rule independently. --- -## 4. `ktdp.inter_tile_reduce` — reduction delivery op +## 4. Placement algebra and type rules + +Result types are a function of the placement value alone. There are four +formulas, applied per role `i` to `T_p_i`: + +| placement | result type derived from `T_p` | +|---|---| +| `replicate` | within-group tile axes collapsed (`fold`) / `T_p` unchanged (`none`) | +| `concat` | extent along `gather_dim` multiplied by `P` | +| `split` | extent along `scatter_dim` divided by `C` | +| `permute` | extent along `scatter_dim` divided by `C`, **and** extent along `gather_dim` multiplied by `P` | + +`reduce_scatter` is `fold` + `split`: the within-group tile axes are +collapsed first, then the `split` formula applies to the collapsed type. + +**Which slice a tile gets.** For `split`, the consumer with local index +`l` (§3.3) receives `[l*chunk : (l+1)*chunk]` along `scatter_dim`, where +`chunk = T_p[scatter_dim] / C`. For `concat`, the producer with local +index `l` occupies `[l*chunk : (l+1)*chunk]` along `gather_dim`, where +`chunk = T_p[gather_dim]` is that producer's own extent along the axis. +For `permute`, both hold simultaneously: consumer `l_c` receives, from +each producer `l_p`, that producer's `scatter_dim` slice `l_c`, placed at +`gather_dim` position `l_p`. + +**Conservation in the square case.** Whenever `P == C`, the `permute` +result has the same element count as `T_p` — one axis is divided and +another multiplied by the same factor — so a square all-to-all is a pure +redistribution of ownership. If additionally `scatter_dim == gather_dim`, +the result *type* equals `T_p`: the distributed transpose, which is the +uniform one-to-one shuffle the SDSC backend emits today (§9). + +**Why `split` divides an honest data axis.** `reduce_scatter` collapses a +*within-group tile axis* and then splits, so its partial must carry an +artificial unit dimension for the collapse to consume. `scatter`, +`gather`, and `all_to_all` split or grow a natural data axis directly, so +their types stay honest: `<128x1x64>` → `<32x1x64>` rather than +`<1x...>`. -### 4.1 Operand and attributes +--- -**Operand:** `!ktdp.tile_future` returned by -`ktdp.inter_tile_produce`. The `#groups` parameter supplies the group set; -there is no separate `groups` attribute (§1). +## 5. Verification rules + +Principle: **each rule has exactly one owner and one statement; +applicability is a column, not a restatement.** "Owner" is the op that +carries the attribute the rule constrains. + +| Rule | Owner | consume | reduce | red_scat | gather | all_to_all | scatter | +|---|---|---|---|---|---|---|---| +| R1 group disjointness (§2.1) | produce | y | y | y | y | y | y | +| R2 single-use future (§2.3) | produce | y | y | y | y | y | y | +| R3 dep set subset of producers | delivery | y | y | y | y | y | n/a | +| R4 every producer covered by some consumer | delivery | y | y | y | y | y | n/a | +| R5 dep sets pairwise disjoint | delivery | — | — | — | y | y | n/a | +| R6 uniform dep-set cardinality | delivery | — | — | — | y | y | n/a | +| R7 uniform producer cardinality across groups | delivery | — | — | — | y | y | n/a | +| R8 producers per group = 1 | delivery | y | — | — | — | — | y | +| R9 `scatter_dim` extent divisible by `C` | delivery | — | — | y | — | y | y | +| R10 combiner purity (§3.5) | delivery | — | y | y | — | — | — | +| R11 identity shape matches `T_p` (§3.5) | delivery | — | y | y | — | — | — | +| R12 `gather_dim` extent × `P` well-defined | delivery | — | — | — | y | y | — | +| R13 consumer set subset of producer set | delivery | — | y | ? | ? | ? | n | +| R14 reduce mode gate: `C == P` or `\|C\| == 1` | delivery | — | y | ? | — | — | — | + +Statements: + +- **R3 — subset.** The declared dependency set must be a subset of + `producer_tiles_per_group`; referencing a non-producer tile is an + error. + + ```text + { p | ∃ c, g : producer_dependency_per_consumer(p)[c, g] } + ⊆ + { p | ∃ g : p ∈ producer_tiles_per_group(g) } + ``` + +- **R4 — coverage.** For every group `g` and every producer `p` in + `producer_tiles_per_group(g)`, at least one consumer `c` in + `consumer_tiles_per_group(g)` must satisfy + `producer_dependency_per_consumer(p)[c, g]`. An uncovered producer + yields a value no consumer reads, risking deadlock in push-based + lowerings. + + ```text + ∀ g, ∀ p ∈ producer_tiles_per_group(g) : + ∃ c ∈ consumer_tiles_per_group(g) : + producer_dependency_per_consumer(p)[c, g] + ``` + +- **R5 — pairwise disjointness.** For the assembling placements + (`concat`, `permute`), distinct consumers' declared dependency sets + must be disjoint. R4 alone requires only that each producer be claimed + by *at least one* consumer, which combined with R6 admits declared sets + that double-count producers — and a double-counted producer has no + well-defined position in the assembly. +- **R6 — uniform dep-set cardinality.** All consumers in a group must + declare the same number of producers, so `P` is a single number and the + op has one static result type. +- **R7 — uniform producer cardinality across groups.** + `producer_tiles_per_group` is a parameterized affine set over `g` and + nothing otherwise requires equal cardinality per group. Since the op + result is a single static tensor type, unequal groups yield no + expressible result type for the assembling placements. +- **R8 — single producer.** Exactly one producer tile per group, for the + ops whose `producers/grp` cell is `1`. +- **R9 — split divisibility.** `T_p[scatter_dim] % C == 0` (for + `reduce_scatter`, the post-collapse extent). One rule covers all three + splitting ops because they share the `scatter_dim` attribute; a separate + divisibility rule for `all_to_all` would only be needed if its split + axis had its own attribute name. +- **R12 — concat well-definedness.** The result extent along `gather_dim` + is `P × T_p[gather_dim]`, which requires every assembled producer to + contribute the same extent along that axis. For the square + `all_to_all` case this follows from R7 + R9, but it must be stated + independently for the non-square case. +- **R13 — consumer set subset of producer set.** Every consumer tile in a + group must also be a producer in that group, i.e. + `consumer_tiles_per_group(g) ⊆ producer_tiles_per_group(g)`. Whether + this should hold is §10.1; it is currently enforced for `reduce` only. +- **R14 — reduce mode gate.** For `reduce`, the consumer set must either + equal the producer set (all-reduce) or be a single tile + (reduce-to-one); a strict multi-tile subset — reduce-to-subset — is + rejected. This is a present implementation restriction, not a design + conclusion (§10.1). -**`consumer_tiles_per_group`** — tiles that receive the reduced result. +--- -**`producer_dependency_per_consumer`** *(optional)* — identical in form -to §3.1. When present, only the partials from the specified producer -tiles are combined; contributions from the remaining producer tiles are -treated as the identity. The result is therefore a partial reduction over -the declared subset. If absent, all producer tiles contribute -(full-barrier semantics). +## 6. The delivery ops -**`identity`** — N variadic SSA operands, one per partial-tensor role. -Each identity tensor's shape and element type must match the corresponding -partial type `T_p_i` (not the rank-reduced result type `T_r_i`). The -identities are hoisted before the op and shared across all groups and all -tiles. Combining any identity with its corresponding partial yields that -partial. +Each subsection states only what is specific to that op: its cells from +§1.1, its result type from §4, its signature, and any op-specific +argument. Shared machinery is §3; rules are §5. -### 4.2 Reducer region +### 6.1 `ktdp.inter_tile_consume` — broadcast -The op has a single region with a block that receives `2N` arguments — -`%lhs_1, ..., %lhs_N, %rhs_1, ..., %rhs_N` with each `%lhs_i` and -`%rhs_i` of type `T_p_i` — and terminates with -`ktdp.yield_reduced %val_1, ..., %val_N : T_p_1, ..., T_p_N`. +`combine = none`, `placement = replicate`, one producer per group, +consumer set free, no dim attribute, no region, no identity. -**Purity.** The combiner must be pure — no memory effects, no calls to -side-effecting ops. Pure tensor ops (`tensor.empty`, `linalg` on tensors, -`arith.*`) are allowed. The verifier rejects combiners containing ops with -side effects. +**Result type.** `T_p_i` unchanged (§4, `replicate` + `none`). -**Combine ordering.** The associative-commutative contract is by user -agreement; the scheduler is free to combine in tree, ring, linear, or any -hardware-native topology. Different groups' reductions are independent and -may be scheduled in parallel. +**Semantics.** No combining occurs. The value produced by the group's +producer tile is delivered unchanged to every consumer tile in that +group — broadcast. + +```mlir +%result_1, ..., %result_N = ktdp.inter_tile_consume(%future) + consumer_tiles_per_group = , + producer_dependency_per_consumer = // optional; default: all producers + : !ktdp.tile_future -> T_p_1, ..., T_p_N +``` -### 4.3 Type rules +With one producer per group, `producer_dependency_per_consumer` is a pure +synchronization refinement (§3.4): it changes when each consumer +unblocks, never what it receives. That is what makes `consume` also the +op for one-to-one permutation exchange — a bijective dependency set over +a multi-producer group (§7.4.2). -For each role `i`, `T_r_i` is `T_p_i` with the within-group tile axes -collapsed. The same set of axes is removed for all roles. +### 6.2 `ktdp.inter_tile_reduce` — reduction -### 4.4 Op signature +`combine = fold`, `placement = replicate`, all tiles produce, consumer +set free, no dim attribute, combiner region and `identity` per §3.5. + +**Result type.** `T_r_i` is `T_p_i` with the within-group tile axes +collapsed; the same axes are removed for all roles. ```mlir %r_1, ..., %r_N = ktdp.inter_tile_reduce(%future) @@ -288,57 +577,26 @@ collapsed. The same set of axes is removed for all roles. } ``` -### 4.5 Result semantics - -The op produces N variadic SSA values, one per partial-tensor role. The -values are *per-tile-valued* — each consumer tile holds a result value -when the op completes. Every consumer tile in a group holds the same -value — that group's fully reduced result. Tiles in different groups hold -different values (each its own group's reduction). - -**Non-participating tiles.** Results are undefined for tiles not in `consumer_tiles_per_group`. - -**Multi-tensor (variadic) reductions.** N ≥ 1 partials are supported. -Argmax-style reductions, where each partial is a correlated tuple of -tensors (values, indices), use N = 2: two identities, two yielded -partials, four reducer region arguments yielding two combined values, -two op results. The structure generalizes uniformly. - ---- - -## 5. `ktdp.inter_tile_reduce_scatter` — reduction + scatter delivery op - -### 5.1 Operand and attributes - -Identical to §4.1 , plus one additional attribute: - -**`scatter_dimension`** (i64) — axis of the post-reduction type along -which the result is split row-major across the consumer tiles. The size -along this dimension must be divisible by the per-group consumer-tile -count. +Consumer set = producer set is all-reduce; a single consumer per group is +reduce-to-one. Both are supported today; a strict multi-tile subset is +not (R14). -**Per-tile slice.** For a tile with within-group local index `l` (its -ascending position among the consumer tiles in the group), the slice -received is `reduced[l*chunk : (l+1)*chunk]` along `scatter_dimension`, -where `chunk = post_reduction_shape[scatter_dimension] / |consumer tiles per group|`. +### 6.3 `ktdp.inter_tile_reduce_scatter` — reduction then split -### 5.2 Reducer region +`combine = fold`, `placement = split`, all tiles produce, consumer set +free, `scatter_dim`, combiner region and `identity` per §3.5. -Identical to §4.2: pure, associative-commutative, combine ordering -unspecified. Different groups' reductions proceed independently. +**`scatter_dim`** (i64) — axis of the *post-collapse* type along which the +reduced result is split row-major across the consumer tiles (R9). -### 5.3 Type rules - -For each role `i`, `T_r_i` is `T_p_i` with the within-group tile axes -collapsed and then sliced along `scatter_dimension`. The same axes and the -same scatter split apply to all roles. - -### 5.4 Op signature +**Result type.** `T_r_i` is `T_p_i` with the within-group tile axes +collapsed and then divided by `C` along `scatter_dim`. The same axes and +the same split apply to all roles. ```mlir %chunk_1, ..., %chunk_N = ktdp.inter_tile_reduce_scatter(%future) consumer_tiles_per_group = , - scatter_dimension = , + scatter_dim = , producer_dependency_per_consumer = , // optional; default: all producers identity(%id_1 : T_p_1, ..., %id_N : T_p_N) : !ktdp.tile_future -> T_r_1, ..., T_r_N @@ -349,267 +607,135 @@ same scatter split apply to all roles. } ``` -### 5.5 Result semantics - -Each consumer tile in a group holds its own row-major slice of that -group's reduced result along `scatter_dimension`. Different tiles in the -same group hold different non-overlapping slices whose concatenation is -the group's full reduced result. Tiles in different groups hold results -from their respective independent reductions. - -**Non-participating tiles.** Same constraint as §4.5. - ---- - -## 6. `ktdp.inter_tile_gather` — assembling delivery op - -### 6.1 Operand and attributes +### 6.4 `ktdp.inter_tile_gather` — ordered assembly -**Operand:** `!ktdp.tile_future` returned by -`ktdp.inter_tile_produce`. The `#groups` parameter supplies the group set; -there is no separate `groups` attribute (§1). +`combine = none`, `placement = concat`, all tiles produce, consumer set +free, `gather_dim`, no region, no identity. -**`consumer_tiles_per_group`** — tiles that receive the assembled tensor. -The set is unrestricted: selecting one tile per group is a plain gather -(one tile assembles the group's full tensor); selecting all tiles is an -all-gather (every tile in the group holds the same assembled tensor). +**`gather_dim`** (i64) — axis of `T_p` along which the producers' +partials are concatenated, in ascending producer local-index order +(§3.3). -**`gather_dimension`** (i64) — axis of the partial type `T_p` along which -the producers' partials are concatenated. Partials are placed along this -axis in ascending within-group local-index order. - -**`producer_dependency_per_consumer`** *(optional)* — identical in form to -§3.1. When absent, every producer tile in the group is assembled (complete -gather) and the consumer waits for all of them. When present, consumer `c` -assembles only the partials from its declared producer tiles, concatenated -in ascending local-index order — a partial (segmented) gather over the -declared subset. The subset and coverage invariants of §3.1 apply. So that -the single op result type is well-formed, every consumer's declared -producer set must have the same cardinality; the verifier rejects unequal -cardinalities. - -**No combiner region and no `identity` operand.** Unlike -`ktdp.inter_tile_reduce` / `ktdp.inter_tile_reduce_scatter`, gather performs -no folding — it assembles slices by position. Like `ktdp.inter_tile_consume` -it therefore carries no region and no identity operand. - -### 6.2 Type rules - -For each role `i`, `T_g_i` is `T_p_i` with the size along `gather_dimension` -multiplied by `K`, where `K` is the number of producers assembled per -consumer: `|producer tiles per group|` when -`producer_dependency_per_consumer` is absent, or the (common) cardinality -of the per-consumer producer set when it is present. The same -`gather_dimension` and the same `K` apply to all roles. - -**Per-tile slice.** The producer with within-group local index `l` (its -ascending position among the assembled producers) occupies slice -`[l*chunk : (l+1)*chunk]` along `gather_dimension` in the output, where -`chunk = T_p[gather_dimension]` is the producer partial's own size along -that axis. Unlike the reduce combiner — which may be applied in any order — -this placement is deterministic and requires no commutativity. - -### 6.3 Op signature +**Result type.** `T_g_i` is `T_p_i` with the extent along `gather_dim` +multiplied by `P` (R12). ```mlir %gathered_1, ..., %gathered_N = ktdp.inter_tile_gather(%future) consumer_tiles_per_group = , - gather_dimension = , + gather_dim = , producer_dependency_per_consumer = // optional; default: all producers : !ktdp.tile_future -> T_g_1, ..., T_g_N ``` -No block is needed — the assembled value is an SSA result consumed by -ordinary function-scope SPMD code, exactly as with `ktdp.inter_tile_consume` -(§3.2). - -### 6.4 Result semantics - -The op produces N variadic SSA values, one per partial-tensor role. The -values are *per-tile-valued* — each consumer tile holds its assembled -result when the op completes. Every consumer tile in a group holds the same -assembled tensor (its group's ordered concatenation of producer partials); -tiles in different groups hold their own group's assembly. A one-tile -consumer set is a plain gather; an all-tiles consumer set is an all-gather. +One consumer per group is a plain gather; the full group as consumer set +is all-gather — the same op with a wider set (§1.2), not a separate op. +With `producer_dependency_per_consumer` present the assembly is a partial +(segmented) gather over each consumer's declared subset, subject to +R5–R7. -**Non-participating tiles.** Results are undefined for tiles not in -`consumer_tiles_per_group`, as in §4.5 and §5.5. - -**Multi-tensor (variadic) gather.** N ≥ 1 partials are supported, following -the same structure as §4.5 — each role is concatenated independently along -`gather_dimension`. - -Synchronization follows the shared model in §8: the def-use edge from -`ktdp.inter_tile_produce` orders production before delivery, and -`producer_dependency_per_consumer` selects full-barrier (absent) or per-tile -(present) waiting. - ---- - -## 7. `ktdp.inter_tile_scatter` — splitting delivery op - -### 7.1 Operand and attributes - -**Operand:** `!ktdp.tile_future` returned by -`ktdp.inter_tile_produce`. The `#groups` parameter supplies the group set; -there is no separate `groups` attribute (§1). Each group has exactly one -producer tile per role — the tile that holds the whole tensor to be split. -The verifier rejects a `producer_tiles_per_group` that selects more than -one tile per group. +### 6.5 `ktdp.inter_tile_all_to_all` — split and reassemble -**`consumer_tiles_per_group`** — tiles that receive the slices. The -producer's tensor is partitioned into `|consumer tiles per group|` equal -chunks along `scatter_dimension`, one chunk delivered to each consumer in -ascending within-group local-index order. +`combine = none`, `placement = permute`, all tiles produce, all tiles +consume, both `scatter_dim` and `gather_dim`, no region, no identity. -**`scatter_dimension`** (i64) — axis of the producer type `T_p` along -which the tensor is split. Its size must be divisible by the number of -consumers per group. +**Attributes.** `scatter_dim` (i64) — axis each producer splits into `C` +chunks (R9). `gather_dim` (i64) — axis along which each consumer +concatenates the chunks it received, in ascending producer local-index +order (R12). The two may be equal (pure ownership transpose along one +axis) or different (reshape-transpose, e.g. split heads and regather +sequence). -**No `producer_dependency_per_consumer`.** With a single producer per -group there is exactly one producer to wait for, so full-barrier and -per-tile synchronization collapse to the same thing; the attribute would -be degenerate and is therefore omitted. - -**No combiner region and no `identity` operand.** Like -`ktdp.inter_tile_consume` and `ktdp.inter_tile_gather`, scatter performs no -folding — it partitions one tensor by position. It carries no region and -no identity operand. +**Result type.** `T_c_i` is `T_p_i` with the `scatter_dim` extent divided +by `C` and the `gather_dim` extent multiplied by `P`. In the square case +(`P == C` and `scatter_dim == gather_dim`) `T_c_i == T_p_i` (§4). -### 7.2 Type rules +**Semantics.** Consumer with local index `l_c` receives, from each +producer `l_p`, the slice `[l_c*chunk : (l_c+1)*chunk]` of that +producer's tensor along `scatter_dim` (`chunk = T_p[scatter_dim] / C`), +and concatenates those `P` slices along `gather_dim` in ascending `l_p` +order. -For each role `i`, `T_s_i` is `T_p_i` with the size along -`scatter_dimension` divided by `|consumer tiles per group|`. The same -`scatter_dimension` and the same divisor apply to all roles. +```mlir +%out_1, ..., %out_N = ktdp.inter_tile_all_to_all(%future) + consumer_tiles_per_group = , + scatter_dim = , + gather_dim = , + producer_dependency_per_consumer = // optional; default: all producers + : !ktdp.tile_future -> T_c_1, ..., T_c_N +``` -**Per-tile slice.** The consumer with within-group local index `l` (its -ascending position among the consumers) receives slice -`[l*chunk : (l+1)*chunk]` along `scatter_dimension`, where -`chunk = T_p[scatter_dimension] / |consumers|` is the per-consumer slice -size. This placement is deterministic and requires no commutativity. +**Why it is a first-class op rather than a composition.** All-to-all +requires every tile to be simultaneously a producer of `C` distinct +slices and a consumer of `P` distinct slices: -**Advantage over reduce-scatter.** `ktdp.inter_tile_reduce_scatter` can -express a bare split only by folding a single-producer axis with a -meaningless combiner and identity, and it shrinks a *within-group tile -axis* — forcing the partial to carry an artificial unit dimension. Scatter -splits the natural data axis directly, so the slice type is the honest -`T_p` with one axis divided (e.g. `<128x1x64>` → `<32x1x64>`) rather than -`<1x...>`. The op name and signature match the pattern. +```text +tile 0: A[0][0..3] tile 0: A[0][0] A[1][0] A[2][0] A[3][0] +tile 1: A[1][0..3] tile 1: A[0][1] A[1][1] A[2][1] A[3][1] +tile 2: A[2][0..3] --> tile 2: A[0][2] A[1][2] A[2][2] A[3][2] +tile 3: A[3][0..3] tile 3: A[0][3] A[1][3] A[2][3] A[3][3] +``` -### 7.3 Op signature +Neither existing copy-only op admits this. `gather` delivers the *same* +assembled tensor to every consumer (§3.7) and cannot give consumers +different content. `scatter` permits exactly one producer per group (R8) +and cannot have every tile contribute. Composing them materializes the +full concatenation on every tile — wrong data volume and wrong +communication pattern. + +The only faithful composition is `C` separate `produce`+`scatter` pairs +(one per source tile, forced by R2) followed by a per-consumer `concat` +in ordinary SPMD code: `C×` the ops, `C×` the produce handles, and +reassembly pushed out of the inter-tile layer. That composition is the +useful *reference lowering* — it is why `all_to_all` needs no new +synchronization (§3.6) and no new verification beyond `scatter` ∪ +`gather` (R9 and R5–R7/R12 respectively) — but it is the wrong surface +form. Note also that **one-to-one permutation** of whole partials is +already expressible as `consume` + a bijective dependency set (§7.4.2); +`all_to_all` is only for the split-and-redistribute case, so the two +mechanisms do not overlap. + +**Generalizing `scatter` to `P > 1` is not an alternative.** Adding a +`gather_dim` and lifting R8 on `scatter` *is* `all_to_all` under another +name; it hides the multi-producer wait inside `scatter` and gives that op +two regimes. A separate op keeps one-op-one-pattern and leaves +`scatter`'s `P == 1` contract clean. + +### 6.6 `ktdp.inter_tile_scatter` — ordered split + +`combine = none`, `placement = split`, one producer per group (R8), +consumer set free, `scatter_dim`, no region, no identity. + +**`scatter_dim`** (i64) — axis of `T_p` along which the single producer's +tensor is partitioned into `C` equal chunks (R9), one per consumer in +ascending consumer local-index order (§3.3). + +**Result type.** `T_s_i` is `T_p_i` with the extent along `scatter_dim` +divided by `C`. ```mlir %scattered_1, ..., %scattered_N = ktdp.inter_tile_scatter(%future) consumer_tiles_per_group = , - scatter_dimension = + scatter_dim = : !ktdp.tile_future -> T_s_1, ..., T_s_N ``` -No block is needed — the slice value is an SSA result consumed by ordinary -function-scope SPMD code, exactly as with `ktdp.inter_tile_consume` (§3.2). - -### 7.4 Result semantics - -The op produces N variadic SSA values, one per tensor role. The values are -*per-tile-valued* — each consumer tile holds its own slice when the op -completes. Consumers in a group receive disjoint, ordered slices that -together tile the producer's tensor along `scatter_dimension`; tiles in -different groups partition their own group's producer tensor. - -**Non-participating tiles.** Results are undefined for tiles not in -`consumer_tiles_per_group`, as in §4.5 and §5.5. - -**Multi-tensor (variadic) scatter.** N ≥ 1 tensors are supported, following -the same structure as §4.5 — each role is split independently along -`scatter_dimension`. +**No `producer_dependency_per_consumer`.** With a single producer per +group there is exactly one producer to wait for, so full-barrier and +per-tile synchronization collapse to the same thing; the attribute would +be degenerate. R3–R7 are therefore `n/a` for this op (§5). **Consumers need not be producers.** A consumer tile that does not appear in `producer_tiles_per_group` simply receives its slice; unlike a partial -gather or reduce there is nothing for a non-producing consumer to -contribute or miss, so no coverage obligation arises. For a pure split the -consumer set is therefore unconstrained relative to the producer set — -resolving, for scatter, the general question of whether a consumer must -also be a producer. - -Synchronization follows the shared model in §8: the def-use edge from -`ktdp.inter_tile_produce` orders production before delivery. With a single -producer per group the wait is unconditional — every consumer waits for -that one producer — so there is no per-tile mode to select. - ---- - -## 8. Synchronization model - -No explicit barriers appear in the IR. The `!ktdp.tile_future` -SSA value carries **per-tile availability signals** rather than a monolithic -group barrier: - -1. Each producer tile's contribution becomes independently observable as - soon as that tile executes `ktdp.yield_partial` in the production - block. -2. A delivery op cannot use a producer tile's contribution until that - tile's signal is set in `%future`. -3. The producer tiles a given consumer tile waits for are declared by the - `producer_dependency_per_consumer` attribute on the delivery op: - - - **Absent (default) — full-barrier mode:** consumer tile `c` in group - `g` waits for every producer tile in `producer_tiles_per_group(g)` - before the delivery op executes. This maps directly to a hardware - group barrier and preserves the simplest safety guarantee. - - **Present — per-tile mode:** consumer tile `c` waits only for the - producer tiles `p` satisfying `producer_dependency_per_consumer(p)[c, - g]`. The consumer unblocks as soon as those specific tiles have - completed, without waiting for unrelated producers. Different consumer - tiles may declare different dependency sets, enabling fine-grained - producer–consumer pipelining. - -**Single-use invariant.** Each `%future` value has exactly one delivery -op use. A second use is a verifier error (§2.3). - -**Subset invariant.** `producer_dependency_per_consumer`, when present, -must be a subset of `producer_tiles_per_group`. Referencing a -non-producer tile is a verifier error. - -**Coverage invariant.** When `producer_dependency_per_consumer` is -present, every producer tile must be declared as a dependency by at -least one consumer tile in the same group. Formally, for every group -`g` and every producer `p` in `producer_tiles_per_group(g)`, there must -exist a consumer `c` in `consumer_tiles_per_group(g)` satisfying -`producer_dependency_per_consumer(p)[c, g]`. A producer not covered by -any consumer is a verifier error — it yields a value that no consumer -reads, which risks a deadlock in push-based lowerings. - -In SPMD KTIR, a tile cannot observe other tiles' partials except through a -dialect-defined boundary. The `ktdp.inter_tile_produce` block is that -boundary — it names the per-tile contribution and exposes it via -`%future`. The delivery op's result tensor is an SSA value that cannot -materialize until the declared dependencies are satisfied; standard MLIR -dataflow ordering applies. - -Lowering inserts target-specific hardware synchronization: a group barrier -for full-barrier mode, and point-to-point ready/wait signals for per-tile -mode. +gather or a reduce there is nothing for a non-producing consumer to +contribute or miss, so no coverage obligation arises. For a pure split +the consumer set is unconstrained relative to the producer set — which +resolves §10.1 for `scatter`, and only for `scatter`. --- -## 9. Coverage of inter-core communication patterns +## 7. Pattern instantiation -These six ops are sufficient to express all five inter-core -communication patterns: - -| Pattern | `inter_tile_produce` | Delivery op | Split/assemble dim | -|---------|---------------------|-------------|---------------------| -| Broadcast | one producer tile per group | `inter_tile_consume` | — | -| Reduce | all tiles per group | `inter_tile_reduce` | — | -| Reduce-scatter | all tiles per group | `inter_tile_reduce_scatter` | `scatter_dimension` | -| Gather | all tiles per group | `inter_tile_gather` | `gather_dimension` | -| Scatter | one producer tile per group | `inter_tile_scatter` | `scatter_dimension` | - ---- - -## 10. Pattern instantiation - -### 10.1 Broadcast → `inter_tile_produce` + `inter_tile_consume` +### 7.1 Broadcast → `inter_tile_produce` + `inter_tile_consume` ```mlir // 4 tiles, 1 group: tile 0 loads W; all 4 tiles compute. @@ -639,7 +765,7 @@ communication patterns: ktdp.store %C, ... ``` -### 10.2 Reduce → `inter_tile_produce` + `inter_tile_reduce` +### 7.2 Reduce → `inter_tile_produce` + `inter_tile_reduce` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -667,7 +793,7 @@ ktdp.store %C, ... } ``` -#### 10.2.1 Full IR — single-group reduce (96×64) +#### 7.2.1 Full IR — single-group reduce (96×64) **Layout and partitioning.** `A` and `B` are `tensor<96x64xf16>` in global memory. The kernel computes the column-wise sum of `A + B`, producing a @@ -784,7 +910,7 @@ module { } ``` -#### 10.2.2 Full IR — multi-group reduce (128×8×12×64) +#### 7.2.2 Full IR — multi-group reduce (128×8×12×64) **Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in global memory. The four axes have distinct roles: @@ -948,7 +1074,7 @@ module { } ``` -### 10.3 Reduce-scatter → `inter_tile_produce` + `inter_tile_reduce_scatter` +### 7.3 Reduce-scatter → `inter_tile_produce` + `inter_tile_reduce_scatter` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -965,10 +1091,10 @@ module { } // Reduce and scatter; each tile receives its own slice along dim 0. -// scatter_dimension = 0 → 128-row axis split across 4 tiles; each gets <32x1x64>. +// scatter_dim = 0 → 128-row axis split across 4 tiles; each gets <32x1x64>. %my_chunk = ktdp.inter_tile_reduce_scatter(%partial_future) consumer_tiles_per_group = #all_group_tiles, - scatter_dimension = 0, + scatter_dim = 0, identity(%add_id : tensor<128x1x1x64xf16>) : !ktdp.tile_future, #all_groups> -> tensor<32x1x64xf16> { @@ -979,7 +1105,7 @@ module { // Each tile holds a different slice — ownership explicit via SSA result. ``` -#### 10.3.1 Full IR — multi-group reduce-scatter (128×8×12×64) +#### 7.3.1 Full IR — multi-group reduce-scatter (128×8×12×64) **Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in global memory. The four axes have distinct roles: @@ -994,7 +1120,7 @@ in global memory. The four axes have distinct roles: 32 tiles, 8 groups of 4. `g = t / 4`, `l = t % 4`. Tile `(g, l)` reads slice `[*, g, l*3 : l*3+3, *]` — shape `<128x1x3x64>`. The per-tile pipeline through to `%partial_4d` (shape `<128x1x1x64>`) is identical -to §10.2.2. +to §7.2.2. The op reduces dim 2 (within-group tile axis, size 1) and scatters dim 0 (128 / 4 = 32 rows per tile). Tile `(g, l)` ends up with rows @@ -1110,7 +1236,7 @@ module { // Group axis (dim 1) preserved. Each tile receives <32x1x64>. %my_chunk = ktdp.inter_tile_reduce_scatter(%partial_future) consumer_tiles_per_group = #group_tiles, - scatter_dimension = 0, + scatter_dim = 0, identity(%add_id : tensor<128x1x1x64xf16>) : !ktdp.tile_future, #all_groups> -> tensor<32x1x64xf16> @@ -1145,9 +1271,9 @@ module { } ``` -### 10.4 Per-tile synchronization → `inter_tile_consume` with `producer_dependency_per_consumer` +### 7.4 Per-tile synchronization → `inter_tile_consume` with `producer_dependency_per_consumer` -#### 10.4.1 Per-tile pairing within a single group +#### 7.4.1 Per-tile pairing within a single group Four tiles per group: tiles `4g` and `4g+1` are producers, tiles `4g+2` and `4g+3` are consumers. Each consumer depends on its dedicated producer @@ -1189,7 +1315,7 @@ both producers finish. With it, each consumer stalls only for its own producer, halving the worst-case wait when the two producers finish at different times. -#### 10.4.2 Butterfly mirror exchange across multiple groups +#### 7.4.2 Butterfly mirror exchange across multiple groups Eight groups of 4 tiles; all 4 tiles in each group both produce and consume. Tile `c = 4g + l` waits only for its mirror partner @@ -1241,7 +1367,7 @@ eliminated). : !ktdp.tile_future, #all_groups> -> tensor<64xf16> ``` -### 10.5 Gather → `inter_tile_produce` + `inter_tile_gather` +### 7.5 Gather → `inter_tile_produce` + `inter_tile_gather` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -1260,18 +1386,18 @@ eliminated). // Gather along dim 2; one consumer per group (tile 4g) assembles the four // 3-wide slabs. No combiner, no identity — placement is by within-group -// local index. gather_dimension = 2 → 3 * 4 = 12; consumer gets <128x1x12x64>. +// local index. gather_dim = 2 → 3 * 4 = 12; consumer gets <128x1x12x64>. %assembled = ktdp.inter_tile_gather(%partial_future) consumer_tiles_per_group = #group_consumer, - gather_dimension = 2 + gather_dim = 2 : !ktdp.tile_future, #all_groups> -> tensor<128x1x12x64xf16> // The consumer holds the full assembled tensor — ownership via SSA result. ``` -#### 10.5.1 Full IR — multi-group gather (128×8×12×64) +#### 7.5.1 Full IR — multi-group gather (128×8×12×64) -**Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in -HBM. The four axes have distinct roles: +**Layout and partitioning.** `A` and `B` are `tensor<128x8x12x64xf16>` in global +memory. The four axes have distinct roles: - Dim 0 (size 128): preserved through this op. - Dim 1 (size 8): the **group axis** — 8 groups. @@ -1325,12 +1451,12 @@ module { %A_view = ktdp.construct_memory_view %A_start, sizes: [128, 8, 12, 64], strides: [6144, 768, 64, 1] { coordinate_set = #A_view_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x8x12x64xf16> %B_view = ktdp.construct_memory_view %B_start, sizes: [128, 8, 12, 64], strides: [6144, 768, 64, 1] { coordinate_set = #A_view_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x8x12x64xf16> // Per-tile compute (function-scope SPMD). @@ -1373,7 +1499,7 @@ module { // assembles the full <128x1x12x64>. No combiner region, no identity. %assembled = ktdp.inter_tile_gather(%partial_future) consumer_tiles_per_group = #group_consumer, - gather_dimension = 2 + gather_dim = 2 : !ktdp.tile_future, #all_groups> -> tensor<128x1x12x64xf16> @@ -1382,7 +1508,7 @@ module { %E_view = ktdp.construct_memory_view %E_start, sizes: [128, 8, 12, 64], strides: [6144, 768, 64, 1] { coordinate_set = #A_view_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x8x12x64xf16> %E_access = ktdp.construct_access_tile %E_view[%c0, %g, %c0, %c0] { @@ -1397,7 +1523,173 @@ module { } ``` -### 10.6 Scatter → `inter_tile_produce` + `inter_tile_scatter` +### 7.6 All-to-all → `inter_tile_produce` + `inter_tile_all_to_all` + +```mlir +// 4 tiles per group, 8 groups (32 tiles total). +#all_group_tiles = affine_set<(i)[g] : (i - 4*g >= 0, -i + 4*g + 3 >= 0)> +#all_groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +// Sequence-parallel production: every tile owns a 128-row shard of all 4 heads. +%partial_future = ktdp.inter_tile_produce + producer_tiles_per_group = #all_group_tiles + : tensor<128x1x4x64xf16> -> !ktdp.tile_future, #all_groups> +{ + ^bb0(%gid: index): + ktdp.yield_partial %partial_4d : tensor<128x1x4x64xf16> +} + +// Head-parallel consumption: split dim 2 (heads) across the 4 consumers, +// regather dim 0 (sequence) from the 4 producers. +// scatter_dim = 2 → 4 / 4 = 1; gather_dim = 0 → 128 * 4 = 512. +%relaid = ktdp.inter_tile_all_to_all(%partial_future) + consumer_tiles_per_group = #all_group_tiles, + scatter_dim = 2, + gather_dim = 0 + : !ktdp.tile_future, #all_groups> -> tensor<512x1x1x64xf16> +// Every tile is both producer and consumer; P == C == 4, so the element count +// is conserved (128*4 = 512*1) even though the type changes. +``` + +#### 7.6.1 Full IR — sequence-parallel to head-parallel (512×8×4×64) + +**Layout and partitioning.** `A`, `B`, and `E` are `tensor<512x8x4x64xf16>` +in global memory. The four axes have distinct roles: + +- Dim 0 (size 512): the **gather axis** — sequence. Sharded 4 ways before + the op, whole after it. +- Dim 1 (size 8): the **group axis** — 8 groups. +- Dim 2 (size 4): the **scatter axis** — heads. Whole before the op, + sharded 4 ways after it. +- Dim 3 (size 64): vector / stick axis, preserved. + +32 tiles, 8 groups of 4. `g = t / 4`, `l = t % 4`. Before the op, tile +`(g, l)` owns sequence shard `l`: it reads `[l*128 : l*128+128, g, *, *]`, +shape `<128x1x4x64>`, and its partial is `A + B` over those rows. After the +op, tile `(g, l)` owns head `l` for the whole sequence, shape +`<512x1x1x64>`, and writes it back to `E[*, g, l, *]`. + +This is the pattern a sequence-parallel prefill hands to a head-parallel +attention: the ownership axis moves from dim 0 to dim 2 in one collective, +with no tile ever holding more than its `1/4` share. + +```mlir +#A_view_set = affine_set<(d0, d1, d2, d3) : + (d0 >= 0, -d0 + 511 >= 0, + d1 >= 0, -d1 + 7 >= 0, + d2 >= 0, -d2 + 3 >= 0, + d3 >= 0, -d3 + 63 >= 0)> + +// A/B access tile for the producer: 128x1x4x64 anchored at [l*128, g, 0, 0]. +#AB_tile_set = affine_set<(d0, d1, d2, d3) : + (d0 >= 0, -d0 + 127 >= 0, + d1 == 0, + d2 >= 0, -d2 + 3 >= 0, + d3 >= 0, -d3 + 63 >= 0)> + +// E access tile for the consumer: 512x1x1x64 anchored at [0, g, l, 0]. +#E_tile_set = affine_set<(d0, d1, d2, d3) : + (d0 >= 0, -d0 + 511 >= 0, + d1 == 0, + d2 == 0, + d3 >= 0, -d3 + 63 >= 0)> + +#identity_4d = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> + +#group_tiles = affine_set<(i)[g] : (i - 4*g >= 0, -i + 4*g + 3 >= 0)> +#all_groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +module { + func.func @inter_tile_all_to_all_relayout() { + %c0 = arith.constant 0 : index + %c4 = arith.constant 4 : index + %row_shard = arith.constant 128 : index // 512 / 4 + + %A_start = arith.constant 1024 : index + %B_start = arith.constant 2098176 : index + %E_start = arith.constant 4195328 : index + + %A_view = ktdp.construct_memory_view %A_start, sizes: [512, 8, 4, 64], + strides: [2048, 256, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.memory_space + } : memref<512x8x4x64xf16> + %B_view = ktdp.construct_memory_view %B_start, sizes: [512, 8, 4, 64], + strides: [2048, 256, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.memory_space + } : memref<512x8x4x64xf16> + + // Per-tile compute (function-scope SPMD). + %t = ktdp.get_compute_tile_id : index + %g = arith.divui %t, %c4 : index + %l = arith.remui %t, %c4 : index + %row_anchor = arith.muli %l, %row_shard : index + + %A_access = ktdp.construct_access_tile %A_view[%row_anchor, %g, %c0, %c0] { + access_tile_set = #AB_tile_set, access_tile_order = #identity_4d + } : memref<512x8x4x64xf16> -> !ktdp.access_tile<128x1x4x64xindex> + %B_access = ktdp.construct_access_tile %B_view[%row_anchor, %g, %c0, %c0] { + access_tile_set = #AB_tile_set, access_tile_order = #identity_4d + } : memref<512x8x4x64xf16> -> !ktdp.access_tile<128x1x4x64xindex> + + %A_tile = ktdp.load %A_access + : !ktdp.access_tile<128x1x4x64xindex> -> tensor<128x1x4x64xf16> + %B_tile = ktdp.load %B_access + : !ktdp.access_tile<128x1x4x64xindex> -> tensor<128x1x4x64xf16> + + // No reduction — the summed sequence shard is this tile's partial; the + // all-to-all redistributes it from sequence-sharded to head-sharded. + %AB_init = tensor.empty() : tensor<128x1x4x64xf16> + %partial_4d = linalg.add ins(%A_tile, %B_tile + : tensor<128x1x4x64xf16>, tensor<128x1x4x64xf16>) + outs(%AB_init : tensor<128x1x4x64xf16>) + -> tensor<128x1x4x64xf16> + + // Produce: every tile contributes its sequence shard to the future. + %partial_future = ktdp.inter_tile_produce + producer_tiles_per_group = #group_tiles + : tensor<128x1x4x64xf16> + -> !ktdp.tile_future, #all_groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %partial_4d : tensor<128x1x4x64xf16> + } + + // All-to-all: consumer l takes head slice l (scatter_dim = 2, 4 / 4 = 1) + // from each of the 4 producers, and concatenates them along the sequence + // axis (gather_dim = 0, 128 * 4 = 512) in ascending producer local-index + // order. The producer's local index picks the destination row block; + // the consumer's local index picks the head. No combiner, no identity. + %relaid = ktdp.inter_tile_all_to_all(%partial_future) + consumer_tiles_per_group = #group_tiles, + scatter_dim = 2, + gather_dim = 0 + : !ktdp.tile_future, #all_groups> + -> tensor<512x1x1x64xf16> + + // Post-exchange: tile (g, l) now owns head l for the whole sequence and + // writes it to E[*, g, l, *]. Every tile is a consumer, so unlike the + // gather example there is no idle tile after the collective. + %E_view = ktdp.construct_memory_view %E_start, sizes: [512, 8, 4, 64], + strides: [2048, 256, 64, 1] { + coordinate_set = #A_view_set, + memory_space = #ktdp.memory_space + } : memref<512x8x4x64xf16> + + %E_access = ktdp.construct_access_tile %E_view[%c0, %g, %l, %c0] { + access_tile_set = #E_tile_set, access_tile_order = #identity_4d + } : memref<512x8x4x64xf16> -> !ktdp.access_tile<512x1x1x64xindex> + + ktdp.store %relaid, %E_access + : tensor<512x1x1x64xf16>, !ktdp.access_tile<512x1x1x64xindex> + + return + } +} +``` + +### 7.7 Scatter → `inter_tile_produce` + `inter_tile_scatter` ```mlir // 4 tiles per group, 8 groups (32 tiles total). @@ -1416,18 +1708,18 @@ module { // Scatter along dim 0; the four tiles per group each receive one 32-row // chunk. No combiner, no identity — placement is by within-group local -// index. scatter_dimension = 0 → 128 / 4 = 32; each consumer gets <32x1x64>. +// index. scatter_dim = 0 → 128 / 4 = 32; each consumer gets <32x1x64>. %chunk = ktdp.inter_tile_scatter(%whole_future) consumer_tiles_per_group = #all_group_tiles, - scatter_dimension = 0 + scatter_dim = 0 : !ktdp.tile_future, #all_groups> -> tensor<32x1x64xf16> // Each consumer holds its own 32-row slice — ownership via SSA result. ``` -#### 10.6.1 Full IR — multi-group scatter (128×8×64) +#### 7.7.1 Full IR — multi-group scatter (128×8×64) -**Layout and partitioning.** `A` and `B` are `tensor<128x8x64xf16>` in -HBM. The three axes have distinct roles: +**Layout and partitioning.** `A` and `B` are `tensor<128x8x64xf16>` in global +memory. The three axes have distinct roles: - Dim 0 (size 128): the **scatter axis** — the producer's 128 rows are split into 4 chunks of 32, one per consumer tile. @@ -1441,6 +1733,14 @@ tensor. Scatter along dim 0 delivers chunk `[l*32 : l*32+32, *, *]` to the consumer with within-group local index `l`, which writes its `<32x1x64>` slice back to `E[l*32 : l*32+32, g, *]`. +**Why the loads live inside the produce region.** Unlike the other full-IR +examples, which hoist their `ktdp.load`s to function scope, this one keeps +them inside `ktdp.inter_tile_produce`. That is deliberate, and it follows +from single-producer cardinality (§2.2): only tile `4g` may read the +group's whole slab, so hoisting the loads would make every tile in the +group execute them. The other examples have every tile produce, so +function-scope loads are correct there. + ```mlir #A_view_set = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 127 >= 0, @@ -1478,12 +1778,12 @@ module { %A_view = ktdp.construct_memory_view %A_start, sizes: [128, 8, 64], strides: [512, 64, 1] { coordinate_set = #A_view_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x8x64xf16> %B_view = ktdp.construct_memory_view %B_start, sizes: [128, 8, 64], strides: [512, 64, 1] { coordinate_set = #A_view_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x8x64xf16> %t = ktdp.get_compute_tile_id : index @@ -1522,7 +1822,7 @@ module { // receives one 32-row chunk. No combiner region, no identity. %chunk = ktdp.inter_tile_scatter(%whole_future) consumer_tiles_per_group = #all_group_tiles, - scatter_dimension = 0 + scatter_dim = 0 : !ktdp.tile_future, #all_groups> -> tensor<32x1x64xf16> @@ -1533,7 +1833,7 @@ module { %E_view = ktdp.construct_memory_view %E_start, sizes: [128, 8, 64], strides: [512, 64, 1] { coordinate_set = #A_view_set, - memory_space = #ktdp.spyre_memory_space + memory_space = #ktdp.memory_space } : memref<128x8x64xf16> %E_access = ktdp.construct_access_tile %E_view[%row_anchor, %g, %c0] { @@ -1550,103 +1850,223 @@ module { --- -## 11. Relationship to existing ops +## 8. Implementation status + +Where the rules of §5 stand in the verifier today. Non-normative: this +section records the current state, not an obligation. + +The legality pass (`lib/Conversion/ConvertToKTIR/KTIRCheckLegality.cpp`, +182 lines) currently walks only `InterTileProduceOp` and +`InterTileReduceOp`: + +| Rule | Op | Check | Location | +|---|---|---|---| +| R2 | `inter_tile_produce` | `future.hasOneUse()` | `KTIRCheckLegality.cpp:80–85` | +| R13 | `inter_tile_reduce` | `C ⊆ P` per group | `KTIRCheckLegality.cpp:107–117` | +| R14 | `inter_tile_reduce` | `C == P` or `\|C\| == 1` | `KTIRCheckLegality.cpp:119–128` | +| R3 | `inter_tile_reduce` | declared dep `p ∈ P(g)` | `KTIRCheckLegality.cpp:151–160` | +| R4 | `inter_tile_reduce` | every `p` covered by some dep | `KTIRCheckLegality.cpp:163–174` | + +**Not yet implemented:** R1, R5, R6, R7, R8, R9, R10, R11, R12, and +R3/R4/R13/R14 for every op other than `reduce`. R5 and R7 are enforced in +the Torch-Spyre SDSC planner (`_compatible_partitions`) but are absent +from the KTIR verifier entirely — the gap exists at both the spec and the +implementation level. + +**Two asymmetries the §5 matrix forces into the open.** + +1. R8 is stated as a verifier obligation for `scatter` but is merely + conventional for `consume`, even though both ops have the same + single-producer cardinality. Neither is implemented yet, so this is a + spec asymmetry to decide rather than inherit. +2. R13 and R14 are implemented for `reduce` only, and R13 is the + implementation of open question §10.1 (must a consumer also be a + producer?) for that one op. The `?` cells in the §5 matrix are exactly + that question, unresolved: for `scatter` the answer is **no** (§6.6), for + `reduce` the current answer is **yes** (enforced), and for + `reduce_scatter` / `gather` / `all_to_all` it is undecided. R14's + mode gate is likewise a current implementation restriction, not a + design conclusion. -| Existing op | Maps to in this design | -|-------------|------------------------| -| `inter_tile_produce` | `ktdp.inter_tile_produce` — `consumer_tiles_per_group` moved to the delivery op; `producer_tile_per_group` → `producer_tiles_per_group` (generalized to multi-producer) | -| `inter_tile_consume` | `ktdp.inter_tile_consume` — unchanged semantics | -| `inter_tile_reduce` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce` — producer block removed from the reduction op | -| `inter_tile_reduce_scatter` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce_scatter` — producer block removed from the reduction op | +--- -`ktdp.inter_tile_gather` (§6) and `ktdp.inter_tile_scatter` (§7) have no -pre-existing counterparts — both are new in this design; the earlier ops -offered neither ordered-concatenation delivery (gather) nor -single-producer ordered-partition delivery (scatter). +## 9. Backend pattern catalogue — non-normative + +This section records relayout patterns observed in the Torch-Spyre backend +(`scratchpad/lx_relayout.py`) and how they map onto the ops above. It is +descriptive, not normative: nothing here constrains the op definitions, +and the mapping column is a proposal for lowering rather than a +guarantee. Rows whose split axes are unknown are omitted. + +All implemented patterns emit a **single** SDSC `opfunc = "shuffle"` whose +entire payload is two `coreIdToWkSlice_` tables — one per tensor in +`coordinates_` — describing per-core ownership before and after the +movement. Classification uses +`gathered_dims = src_syms − dst_syms`, +`scattered_dims = dst_syms − src_syms`, +`factor = num_cores // prod(dst splits)`. + +| ID | SenDNN op | Src `work_div` | Dst view | gathered | scattered | factor | KTIR op | Backend status | +|---|---|---|---|---|---|---|---|---| +| P01 | all-gather | `{H:8, Lk:4}` | replicated (all cores) | H, Lk | — | N/A | `inter_tile_gather` (all consumers) | missing — replication not supported | +| P03 | grouped all-gather | `{H:8, Lk:4}` | `{H:8}` | Lk | — | 4 | `inter_tile_gather` | #3440 (open PR) | +| P04 | grouped all-gather | `{Lk:32}` | `{H:8}` | Lk | H | 4 | `inter_tile_gather` | #3440 (open PR) | +| P06 | all-to-all | `{H:8, Lk:4}` | `{Lk:32}` | H | — | 1 | `inter_tile_all_to_all` | main (#3439) | +| P08 | all-to-all (axis transpose) | `{A:4, B:8}` | `{A:8, B:4}` | — | — | 1 | `inter_tile_all_to_all` + explicit coord map | missing — axis swap inexpressible | +| P14 | all-to-all | `{H:8, Lq:4}` | `{Lq:32}` | H | — | 1 | `inter_tile_all_to_all` | main (#3439) | + +- **P03 / P04** show that `gathered_dims ≠ ∅` with `factor > 1` maps to + `inter_tile_gather`. For P04, `scattered_dims = {H}` (H is introduced at + the destination) — a combined gather+scatter in one shuffle step, so the + KTIR op must express both axes. +- **P06 / P14** show `inter_tile_all_to_all` at `factor = 1`: + `gathered_dims = {H}`, `scattered_dims = ∅`, with H contracted into the + 1-D destination axis. +- **P08 is the case `all_to_all` with dim attributes cannot express.** Both + sides split the same two dims with swapped counts (`[4,8] → [8,4]`). The + `coreIdToWkSlice_` tables differ because the mixed-radix odometer + ordering changes, but no single `scatter_dim`/`gather_dim` pair captures + the transformation. Marked out of scope for `all_to_all`; a verifier + should reject non-decomposable transpositions rather than silently + mis-lower them. §10.4 (Option B) is the eventual fix. +- **P01** (full replication) is blocked because `_compatible_partitions` + requires distinct slices per destination core. It maps to + `inter_tile_gather` with `consumer_tiles_per_group = all`, but needs new + backend support for non-bijective shuffle. + +**Classification decision rule** (from `TensorArg.work_division` on a +relayout identity OpSpec): + +| `gathered_dims` | `scattered_dims` | `factor` | `slot_exprs_differ` | KTIR op | +|---|---|---|---|---| +| ∅ | ∅ | 1 | false | no-op | +| ∅ | ∅ | 1 | true | `inter_tile_all_to_all` + explicit coord map | +| non-∅ | ∅ or non-∅ | > 1 | — | `inter_tile_gather` | +| non-∅ | ∅ | 1 | — | `inter_tile_all_to_all` | +| ∅ | non-∅ | 1 | — | `inter_tile_scatter` | +| ∅ | ∅ | 1 | — (replication) | `inter_tile_gather` (all consumers) | + +**Fused relayout is deferred.** Relayout stays a separate preceding op and +fusion is a lowering concern. The backend structurally cannot fuse them +today: restickify is a separate pass that runs before LX planning, and +restickified weights are explicitly barred as shuffle sources. -The `!ktdp.tile_future` type is shared across all ops; its -`#groups` parameter carries the group set (§1). +--- -The previous `ktdp.inter_tile` single op (Approach B draft) is replaced -by this six-op design: `ktdp.inter_tile` had producer and optional -combiner regions in one op with `consumer_tiles_per_group` determining -delivery mode. The six-op design makes production and delivery explicitly -separate ops, with the delivery mode determined by which delivery op is -chosen rather than by attribute combinations. +## 10. Open questions and extensions ---- +### 10.1 Must a consumer also be a producer? + +Open for `consume`, `reduce`, `reduce_scatter`, `gather`, and +`all_to_all`; **resolved for `scatter`** — no (§6.6). + +The current implementation answers *yes* for `reduce` and enforces it +(R13, `KTIRCheckLegality.cpp:107–117`), whose error text names this +question explicitly. That is one op's implementation choice, not a design +conclusion for the family. The related R14 mode gate — `reduce` supports +all-reduce (`C == P`) and reduce-to-one (`|C| == 1`) but rejects a strict +multi-tile consumer subset — is likewise a present restriction awaiting a +decision. -## 12. Open questions +Deciding this per op is what the `?` cells in §5 record. -**Q1. Multi-tensor generalization.** -The existing ops support variadic partials (N ≥ 1 for argmax-style -reductions). The ops here should carry the same variadic structure. -For N = 2 (argmax): two identities, two `ktdp.yield_partial` operands in -the produce block, four combiner arguments yielding two values, two -delivery op results. Each result follows the same per-op type rules -independently. +### 10.2 Delivery-op placement -**Q2. Consume placement.** -Whether the verifier should enforce that delivery ops appear only inside -a guard matching `consumer_tiles_per_group`, or whether this is left to +Whether the verifier should enforce that a delivery op appears only inside +a guard matching `consumer_tiles_per_group`, or whether that is left to lowering. If the union of consumer sets equals the set of all executing -tiles, no guard is needed; otherwise a tile not in the consumer set that +tiles, no guard is needed; otherwise a tile outside the consumer set that reaches the delivery op would be a verifier error. +### 10.3 Multiple delivery ops per future + +R2 restricts a `!ktdp.tile_future<...>` value to exactly one delivery use. +A natural extension would allow several delivery ops to consume the same +future, each declaring its own `producer_dependency_per_consumer` — one +`ktdp.inter_tile_produce` serving two independent deliveries, e.g. one +waiting on the first half of the producers and another on the second half. + +**Expressiveness gain.** Patterns that today need two separate +`ktdp.inter_tile_produce` ops with identical producer regions collapse to +one produce plus two delivery ops, removing redundant producer-side code +and making the shared production explicit in the IR. + +**Verification cost.** Single-use keeps R4 (coverage) local: the verifier +inspects one delivery op to confirm every producer tile is covered. With +multiple uses, coverage becomes global — for every group `g` and producer +`p ∈ producer_tiles_per_group(g)`, at least one consumer `c` across *any* +delivery op must satisfy `producer_dependency_per_consumer(p)[c, g]`. That +requires collecting and unioning the dependency sets from all uses of the +SSA value before checking, a def-use traversal rather than a local per-op +check. R5 (pairwise disjointness) would have to become cross-op too. + +**Lowering cost.** Each delivery op declaring a dependency set introduces +its own point-to-point signals. A producer `p` may then need to signal +multiple consumers across different delivery ops, and lowering must emit +each signal exactly once and receive it exactly once per dependent +consumer. In full-barrier mode, multiple delivery ops on one future also +require handling duplicate barrier waits: a producer-side barrier cannot +be issued until every dependent delivery op is ready to receive. + +Given that cost, the current design requires separate +`ktdp.inter_tile_produce` ops for separate delivery concerns. If real use +cases demand shared production, the restriction can be relaxed. + +### 10.4 Option B — explicit coordinate map + +Replace the `scatter_dim`/`gather_dim` attribute pair with a single +source-to-destination affine map, subsuming all four placement values in +one mechanism. This is the shape interface-specs PR 14 already uses +(`SHUFFLE` as source/destination coordinate sets), and the backend already +speaks per-core partitionings (`coreIdToWkSlice_` tables) rather than dim +attributes — so Option B is arguably closer to the existing contract, and +it is the only form that expresses P08 (§9). + +**Direction: dim attributes first, Option B recorded as the long-term +target.** The dim-attribute form is additive and reviewable on its own, +and it covers every pattern the backend implements today. Note the honest +counter-argument: because none of the six delivery ops is built yet, the +usual "Option A is cheaper because it is incremental" argument is weaker +here than normal. + +### 10.5 Post-v1 extensions to `all_to_all` + +- **`all_to_all_v`** — uneven shard sizes, i.e. per-consumer split extents + instead of a uniform `T_p[scatter_dim] / C`. This relaxes R9 into a + per-consumer size list and needs a variadic size attribute; no current + backend pattern requires it. +- **Multi-axis relayout** — splitting or gathering along more than one + axis in a single op. Expressible today only as a sequence of + `all_to_all` ops; Option B (§10.4) is the natural home for it. +- **`inter_tile_shuffle` as a naming alias** — the SDSC backend calls the + primitive `shuffle`. `all_to_all` is kept as the op name because it is + the established collective term and because `shuffle` is the *lowering* + of several patterns, not just this one (§9). + --- -## 13. Possible extensions - -### 13.1 Multiple delivery ops per future - -The current spec restricts a `!ktdp.tile_future` value to -exactly one delivery op use. A natural extension would allow multiple delivery -ops to consume the same future, with each declaring its own -`producer_dependency_per_consumer`. This would let one -`ktdp.inter_tile_produce` serve two independent delivery operations — -for example, one delivery op waits for the first half of the producer -tiles while another waits for the second half. - -**Expressiveness gain.** Patterns that currently require two separate -`ktdp.inter_tile_produce` ops (with identical producer regions) could be -expressed with a single produce op and two delivery ops. This reduces -redundant producer-side code and makes the shared production explicit in -the IR. - -**Verification complexity.** The single-use restriction keeps the -coverage check local: the verifier inspects only the one delivery op to -confirm every producer tile is covered. With multiple delivery ops, the -coverage invariant must be checked globally across all uses of the -future: for every group `g` and every producer `p` in -`producer_tiles_per_group(g)`, at least one consumer tile `c` across -any of the delivery ops must satisfy `producer_dependency_per_consumer(p)[c, g]`. -This requires the verifier to collect and union the dependency sets from -all uses of the SSA value before checking coverage — a cross-op, -def-use-traversal analysis rather than a local per-op check. - -**Lowering complexity.** Each delivery op that declares a -`producer_dependency_per_consumer` introduces its own set of -point-to-point synchronization signals. A producer tile `p` may now -need to signal multiple consumers across different delivery ops, and the -lowering must ensure that all signals are emitted exactly once by `p` -and received exactly once by each dependent consumer. In full-barrier -mode, multiple delivery ops on the same future also require the lowering -to handle duplicate barrier waits — a producer-side barrier cannot be -issued until all delivery ops that depend on it are ready to receive. - -Given this added verification and lowering complexity, the current design -requires separate `ktdp.inter_tile_produce` ops for separate delivery -concerns. If future use cases demonstrate a clear need for the shared -production pattern, this restriction can be relaxed. - -### 13.2 Summary of operation coverage - -| Pattern | Producers per group | Delivery op | Result per consumer | -|---------|--------------------|-----------------------------|---------------------| -| Broadcast | 1 | `inter_tile_consume` | full copy | -| Reduce | N | `inter_tile_reduce` | fully reduced | -| Reduce-scatter | N | `inter_tile_reduce_scatter` | 1/N slice of reduced | -| Gather | N | `inter_tile_gather` | full assembled tensor | -| Scatter | 1 | `inter_tile_scatter` | 1/N slice of full | - -All rows are first-class ops in this design (§2–§7). +## Appendix A. Relationship to the pre-existing ops + +| Existing op | Maps to in this design | +|-------------|------------------------| +| `inter_tile_produce` | `ktdp.inter_tile_produce` — `consumer_tiles_per_group` moved to the delivery op; `producer_tile_per_group` → `producer_tiles_per_group` (generalized to multi-producer) | +| `inter_tile_consume` | `ktdp.inter_tile_consume` — unchanged semantics | +| `inter_tile_reduce` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce` — producer block removed from the reduction op | +| `inter_tile_reduce_scatter` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce_scatter` — producer block removed from the reduction op | + +`ktdp.inter_tile_gather` (§6.4), `ktdp.inter_tile_all_to_all` (§6.5), and +`ktdp.inter_tile_scatter` (§6.6) have no pre-existing counterparts. The +earlier ops offered none of ordered-concatenation delivery (gather), +split-and-reassemble delivery (all-to-all), or single-producer +ordered-partition delivery (scatter). + +The `!ktdp.tile_future` type is shared across all ops; its +`#groups` parameter carries the group set (§1.3). + +The previous `ktdp.inter_tile` single op (Approach B draft) is replaced by +this seven-op design. `ktdp.inter_tile` carried producer and optional +combiner regions in one op, with `consumer_tiles_per_group` determining the +delivery mode. Splitting production from delivery makes the mode a choice +of op rather than an inference over attribute combinations — which is what +lets §3 state the shared machinery once and §6 reduce each op to its own +cells. From a01c23c4839690adb77d123e586caf99814a28e9 Mon Sep 17 00:00:00 2001 From: Yu Chin Fabian Lim Date: Thu, 27 Aug 2026 20:09:27 -0400 Subject: [PATCH 04/10] incorporate data from torch-spyre Signed-off-by: Yu Chin Fabian Lim --- docs/inter-tile-communication.md | 878 ++++++++++++++++++++----------- 1 file changed, 583 insertions(+), 295 deletions(-) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index 88a51e7..e56009d 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -19,8 +19,9 @@ as `(Rn)` at the place the attribute they constrain is introduced, so a citation like `(R1)` in §2.1 means "§5 states this rule; here is the attribute it applies to." -Sections are normative except §8 (implementation status) and §9 (observed -backend patterns). +Sections are normative except §8 (implementation status) and §9 (measured +backend patterns). §9 doubles as the evidence for *which* ops a lowering +must actually emit: §9.3 reads the requirement off 51 measured relayouts. --- @@ -42,7 +43,11 @@ Separating production from delivery keeps each op single-purpose and enables any combination: one production op plus a choice of delivery op. The pairing is **one-to-one** — a production op is consumed by exactly one delivery op (R2, §2.3). A pattern needing two deliveries therefore needs -two `ktdp.inter_tile_produce` ops; §10.3 discusses relaxing that. +two `ktdp.inter_tile_produce` ops. Allowing several deliveries per future +would let them share one production, but it makes R4 (coverage) and R5 +(disjointness) non-local — they would have to union the dependency sets across +every use of the SSA value — so the restriction stands until a use case needs +it. ### 1.1 Semantics matrix @@ -59,16 +64,26 @@ means a tensor or tile axis. |---|---|---|---|---|---|---|---| | `consume` | none | replicate | 1 | free | — | — | — | | `reduce` | fold | replicate | all | free | — | combiner | yes | -| `reduce_scatter` | fold | split | all | free | `scatter_dim` | combiner | yes | -| `gather` | none | concat | all | free | `gather_dim` | — | — | -| `all_to_all` | none | permute | all | all | `scatter_dim`, `gather_dim` | — | — | -| `scatter` | none | split | 1 | free | `scatter_dim` | — | — | +| `reduce_scatter` | fold | split | all | free | `scatter_dimensions` | combiner | yes | +| `gather` | none | concat | all | free | `gather_dimensions` | — | — | +| `all_to_all` | none | permute | all | all | `split_dimensions`, `concat_dimensions` | — | — | +| `scatter` | none | split | 1 | free | `scatter_dimensions` | — | — | `all_to_all` is listed before `scatter` because it shares the all-producers cardinality cell with `gather` and `reduce_scatter`, and because its relationship to the two copy-only placements is structural: **permute = split + concat in one step**, which is why it carries both dim -attributes and no new ones. +attributes and no new ones. `all_to_all` names them `split_dimensions` and +`concat_dimensions` — the same two roles `scatter_dimensions` and +`gather_dimensions` play on the single-role ops, renamed because on this +op both are present at once and *scatter* / *gather* would then name +neither the op nor a unique role. + +**Every dim attribute is a list of axis indices** into `T_p` — an +`i64` array, not a single `i64` — flattened in list order per §4. §9.3 +contains a measured pattern whose concat is three axes wide, so the +list-valued form is required by a named pattern rather than reserved for +a corner case. Three things this matrix makes visible: @@ -207,7 +222,7 @@ observable the moment that tile executes `ktdp.yield_partial`. **Single-use invariant (R2).** `%future` must have exactly one use — the single delivery op that consumes it. If two delivery ops need to communicate with the same set of producers, they must each have their own -`ktdp.inter_tile_produce`. §10.2 discusses relaxing this. +`ktdp.inter_tile_produce` (see §1). --- @@ -254,8 +269,8 @@ consuming the SSA value. **`l` is a tile's position, counting from 0 in ascending tile-id order, among the relevant set within its group** — the producer set for `concat` placement (and for -`permute`'s `gather_dim`), the consumer set for `split` placement (and -for `permute`'s `scatter_dim`). +`permute`'s `concat_dimensions`), the consumer set for `split` placement +(and for `permute`'s `split_dimensions`). This definition is what makes ordered placement well-defined. Without it, concatenation and split orders are pinned down only by contiguous-tile-id @@ -391,8 +406,8 @@ the same group hold the *same* value is a property of the placement: `consumer_tiles_per_group`. **Multi-tensor (variadic) delivery.** `N >= 1` roles are supported by -every op, and all roles share the same attributes (`scatter_dim`, -`gather_dim`, `P`, `C`) — only the types differ. Argmax-style reductions, +every op, and all roles share the same attributes (`scatter_dimensions`, +`gather_dimensions`, `P`, `C`) — only the types differ. Argmax-style reductions, where each contribution is a correlated tuple of tensors (values, indices), use `N = 2`: two identities, two yielded partials, four combiner arguments yielding two combined values, two op results. Each @@ -407,36 +422,98 @@ formulas, applied per role `i` to `T_p_i`: | placement | result type derived from `T_p` | |---|---| -| `replicate` | within-group tile axes collapsed (`fold`) / `T_p` unchanged (`none`) | -| `concat` | extent along `gather_dim` multiplied by `P` | -| `split` | extent along `scatter_dim` divided by `C` | -| `permute` | extent along `scatter_dim` divided by `C`, **and** extent along `gather_dim` multiplied by `P` | - -`reduce_scatter` is `fold` + `split`: the within-group tile axes are -collapsed first, then the `split` formula applies to the collapsed type. +| `replicate` | `T_p` unchanged — no rank reduction | +| `concat` | extents along `gather_dimensions` multiplied by `P` in total | +| `split` | extents along `scatter_dimensions` divided by `C` in total | +| `permute` | extents along `scatter_dimensions` divided by `C` in total, **and** extents along `gather_dimensions` multiplied by `P` in total | + +`reduce_scatter` is `fold` + `split`: the `split` formula applies directly +to `T_p`, with no collapse first. + +**No rank reduction anywhere.** All four formulas keep `T_p`'s rank. The +same reasoning that settled it for `reduce` (§6.2) applies to `concat` and +`split`: the axis the op concatenates along or splits is an axis `T_p` +already has, so there is nothing to collapse and no rank to restore. A +`concat` result differs from `T_p` only in the extent along the listed axes, +a `split` result likewise — never in rank. This keeps every result type a +per-axis extent rewrite of the partial, which is what lets §10.3 reason +about layout transparency one axis at a time. + +**Axis sets and flattening — normative.** Each dim attribute is a *list* +of axis indices into `T_p`, not a single axis. A list of length `n > 1` +denotes the product space of those axes, linearized as a row-major +(mixed-radix odometer) order over the listed extents: **the first entry is +the slowest-varying and the last is the fastest-varying**. Write +`E(D) = prod(T_p[d] for d in D)` for the flattened extent of axis set `D`. +The single-axis case is `n == 1`, where `E(D) = T_p[d]` and every formula +below reduces to its familiar form; `n == 0` is invalid for an op that +carries the attribute. + +**The list is in ascending numerical order (R9).** Entries must ascend, so +the slowest-to-fastest flattening above coincides with ascending axis index +and the attribute has exactly one legal spelling per axis set. Two reasons +this is a rule and not a convention. It removes a silent-miscompile class: +`[2, 0]` and `[0, 2]` are both "valid, distinct, non-empty" and would flatten +to *different* data orders, so a reversed list passes every other check while +meaning something else. And it makes attribute equality a list comparison — +which §4's conservation case below depends on, since `all_to_all` decides +whether `T_c == T_p` by testing `split_dimensions == concat_dimensions`. + +Entries need not be **adjacent**: `[0, 2]` over a rank-3 partial is legal and +is exactly what physicalization produces (§10.3). + +**Split and concat apply to the floordiv axis — normative.** When a listed +axis is a **sticked** axis — one that a stick layout has split into a +`floordiv` (chunk-count) axis and a `mod` (within-stick) axis — the `÷ C` or +`× P` applies to the **floordiv axis only**. The `mod` axis is invariant: its +extent is the stick size, and changing it would redefine what a stick is. + +This settles what "`E(D)` divided by `C`" alone leaves open, since a flattened +extent does not say which listed axis absorbs the factor. For a partial +`[2, 16, 32]` (logical `[16, 64]`, stick 32) with `gather_dimensions = [0, 2]` +and `P = 4`, the result is `[8, 16, 32]` — the chunk count goes `2 → 8` and +the stick axis stays `32`, which is exactly the physicalization of the logical +result `[16, 256]`. Absorbing into the `mod` axis instead would give +`[2, 16, 128]`: the same flattened extent, the wrong tensor. + +A useful consequence: **R9 applied to the floordiv axis is the stick-multiple +check.** `E(floordiv) % C == 0` holds exactly when the logical result extent +is a whole multiple of the stick, so a split that would drive the result +sub-stick fails R9 rather than needing a rule of its own. On the partial +above, `C = 2` gives `2 % 2 == 0` and a result of `[1, 16, 32]`; `C = 4` gives +`2 % 4 ≠ 0` and is rejected — correctly, since the logical result `[16, 16]` +is half a stick and unrepresentable in that layout. + +Fixing this order is a requirement, not a convenience: §9.3 contains a +measured three-axis concat, so the flattening must be well-defined over +more than two axes for a *named* pattern rather than only a corner case. **Which slice a tile gets.** For `split`, the consumer with local index -`l` (§3.3) receives `[l*chunk : (l+1)*chunk]` along `scatter_dim`, where -`chunk = T_p[scatter_dim] / C`. For `concat`, the producer with local -index `l` occupies `[l*chunk : (l+1)*chunk]` along `gather_dim`, where -`chunk = T_p[gather_dim]` is that producer's own extent along the axis. -For `permute`, both hold simultaneously: consumer `l_c` receives, from -each producer `l_p`, that producer's `scatter_dim` slice `l_c`, placed at -`gather_dim` position `l_p`. +`l` (§3.3) receives `[l*chunk : (l+1)*chunk]` of the flattened +`scatter_dimensions` space, where `chunk = E(scatter_dimensions) / C`. +For `concat`, the producer with local index `l` occupies +`[l*chunk : (l+1)*chunk]` of the flattened `gather_dimensions` space, +where `chunk = E(gather_dimensions)` is that producer's own flattened +extent over those axes. For `permute`, both hold simultaneously: consumer +`l_c` receives, from each producer `l_p`, that producer's +`scatter_dimensions` slice `l_c`, placed at `gather_dimensions` position +`l_p`. + +A multi-axis split or concat therefore needs no special case in the +decision procedure of §9.1: the dimension attributes are the axis *sets* +themselves, in this order. **Conservation in the square case.** Whenever `P == C`, the `permute` result has the same element count as `T_p` — one axis is divided and another multiplied by the same factor — so a square all-to-all is a pure -redistribution of ownership. If additionally `scatter_dim == gather_dim`, +redistribution of ownership. If additionally `split_dimensions == concat_dimensions`, the result *type* equals `T_p`: the distributed transpose, which is the uniform one-to-one shuffle the SDSC backend emits today (§9). -**Why `split` divides an honest data axis.** `reduce_scatter` collapses a -*within-group tile axis* and then splits, so its partial must carry an -artificial unit dimension for the collapse to consume. `scatter`, -`gather`, and `all_to_all` split or grow a natural data axis directly, so -their types stay honest: `<128x1x64>` → `<32x1x64>` rather than -`<1x...>`. +**Why `split` divides an honest data axis.** Every splitting op divides an +extent of an axis the partial already has, so the types stay honest: +`<128x1x64>` → `<32x1x64>`, never `<1x...>`. No op manufactures a unit +dimension for a collapse to consume, and none removes one. --- @@ -456,10 +533,10 @@ carries the attribute the rule constrains. | R6 uniform dep-set cardinality | delivery | — | — | — | y | y | n/a | | R7 uniform producer cardinality across groups | delivery | — | — | — | y | y | n/a | | R8 producers per group = 1 | delivery | y | — | — | — | — | y | -| R9 `scatter_dim` extent divisible by `C` | delivery | — | — | y | — | y | y | +| R9 flattened split extent divisible by `C` | delivery | — | — | y | — | y | y | | R10 combiner purity (§3.5) | delivery | — | y | y | — | — | — | | R11 identity shape matches `T_p` (§3.5) | delivery | — | y | y | — | — | — | -| R12 `gather_dim` extent × `P` well-defined | delivery | — | — | — | y | y | — | +| R12 flattened concat extent × `P` well-defined | delivery | — | — | — | y | y | — | | R13 consumer set subset of producer set | delivery | — | y | ? | ? | ? | n | | R14 reduce mode gate: `C == P` or `\|C\| == 1` | delivery | — | y | ? | — | — | — | @@ -504,16 +581,37 @@ Statements: expressible result type for the assembling placements. - **R8 — single producer.** Exactly one producer tile per group, for the ops whose `producers/grp` cell is `1`. -- **R9 — split divisibility.** `T_p[scatter_dim] % C == 0` (for - `reduce_scatter`, the post-collapse extent). One rule covers all three - splitting ops because they share the `scatter_dim` attribute; a separate - divisibility rule for `all_to_all` would only be needed if its split - axis had its own attribute name. -- **R12 — concat well-definedness.** The result extent along `gather_dim` - is `P × T_p[gather_dim]`, which requires every assembled producer to - contribute the same extent along that axis. For the square - `all_to_all` case this follows from R7 + R9, but it must be stated - independently for the non-square case. +- **R9 — split divisibility.** `E(D_split) % C == 0`, where `D_split` is + the op's split axis set (`scatter_dimensions`, or `split_dimensions` for + `all_to_all`) and `E` is the flattened extent of §4. Stating the rule on the + *flattened* extent is what lets one rule cover all three splitting ops + and every arity: a multi-axis split need only divide in the product, + not axis by axis. + + Every axis index in the list must be a valid, distinct axis of `T_p`; the + list must be non-empty; and the entries must be in **ascending numerical + order** (§4). Repeated indices would double-count an extent in `E`, and an + out-of-order list would silently denote a different flattening. + Repeated indices would double-count an extent in `E`. +- **R11 and the shipped constraint.** R11 pins `identity` to `T_p`, while + the implemented `reduce` ties it to *results* (`KTDP.td:172-174`). With no + rank reduction (§4) these coincide for `reduce`, since its result *is* + `T_p`. They diverge for `reduce_scatter`, whose result is `T_p` split by + `C`: R11's `T_p` is the correct one there, since the identity is combined + with partials before the split. A verifier generalizing the shipped + constraint to `reduce_scatter` must therefore retarget it from results to + the future's partial types (§10.3). + +- **R12 — concat well-definedness.** The result flattened extent over the + concat axis set `D_concat` (`gather_dimensions`, or `concat_dimensions` + for `all_to_all`) is `P × E(D_concat)`, which requires every assembled + producer to contribute the same extent along *each* listed axis — equal + products alone would not give a well-defined multi-axis assembly, since + the flattening of §4 depends on the individual extents. The same + validity conditions as R9 apply to the list. For the square + `all_to_all` case the divisibility follows from R7 + R9, but it must be + stated independently for the non-square case, which §9.3 shows is + measured and not hypothetical. - **R13 — consumer set subset of producer set.** Every consumer tile in a group must also be a producer in that group, i.e. `consumer_tiles_per_group(g) ⊆ producer_tiles_per_group(g)`. Whether @@ -561,8 +659,13 @@ a multi-producer group (§7.4.2). `combine = fold`, `placement = replicate`, all tiles produce, consumer set free, no dim attribute, combiner region and `identity` per §3.5. -**Result type.** `T_r_i` is `T_p_i` with the within-group tile axes -collapsed; the same axes are removed for all roles. +**Result type.** `T_r_i == T_p_i` — no rank reduction. An earlier draft +collapsed the within-group tile axes; the implementation deliberately does +not (`KTDP.td:197`), because the partial already carries that axis and +keeping it makes the op simpler: result, partial and `identity` are then one +type, tied declaratively (`KTDP.td:168-174`) rather than by a shape +computation. This is also what makes `reduce` transparent under +physicalization (§10.3). ```mlir %r_1, ..., %r_N = ktdp.inter_tile_reduce(%future) @@ -584,19 +687,19 @@ not (R14). ### 6.3 `ktdp.inter_tile_reduce_scatter` — reduction then split `combine = fold`, `placement = split`, all tiles produce, consumer set -free, `scatter_dim`, combiner region and `identity` per §3.5. +free, `scatter_dimensions`, combiner region and `identity` per §3.5. -**`scatter_dim`** (i64) — axis of the *post-collapse* type along which the +**`scatter_dimensions`** (`i64` array) — axes of `T_p` along which the reduced result is split row-major across the consumer tiles (R9). -**Result type.** `T_r_i` is `T_p_i` with the within-group tile axes -collapsed and then divided by `C` along `scatter_dim`. The same axes and -the same split apply to all roles. +**Result type.** `T_r_i` is `T_p_i` with the flattened extent over +`scatter_dimensions` divided by `C` — no rank reduction (§4). The same axes +and the same split apply to all roles. ```mlir %chunk_1, ..., %chunk_N = ktdp.inter_tile_reduce_scatter(%future) consumer_tiles_per_group = , - scatter_dim = , + scatter_dimensions = , producer_dependency_per_consumer = , // optional; default: all producers identity(%id_1 : T_p_1, ..., %id_N : T_p_N) : !ktdp.tile_future -> T_r_1, ..., T_r_N @@ -610,19 +713,21 @@ the same split apply to all roles. ### 6.4 `ktdp.inter_tile_gather` — ordered assembly `combine = none`, `placement = concat`, all tiles produce, consumer set -free, `gather_dim`, no region, no identity. +free, `gather_dimensions`, no region, no identity. -**`gather_dim`** (i64) — axis of `T_p` along which the producers' -partials are concatenated, in ascending producer local-index order -(§3.3). +**`gather_dimensions`** (`i64` array) — axes of `T_p` along which the +producers' partials are concatenated, in ascending producer local-index +order (§3.3). A multi-axis set concatenates in the flattened space of §4, +listed axes ordered slowest- to fastest-varying; §9.3's all-gather +patterns supply a measured three-axis case. -**Result type.** `T_g_i` is `T_p_i` with the extent along `gather_dim` -multiplied by `P` (R12). +**Result type.** `T_g_i` is `T_p_i` with the flattened extent over +`gather_dimensions` multiplied by `P` (R12). ```mlir %gathered_1, ..., %gathered_N = ktdp.inter_tile_gather(%future) consumer_tiles_per_group = , - gather_dim = , + gather_dimensions = , producer_dependency_per_consumer = // optional; default: all producers : !ktdp.tile_future -> T_g_1, ..., T_g_N ``` @@ -636,30 +741,33 @@ R5–R7. ### 6.5 `ktdp.inter_tile_all_to_all` — split and reassemble `combine = none`, `placement = permute`, all tiles produce, all tiles -consume, both `scatter_dim` and `gather_dim`, no region, no identity. - -**Attributes.** `scatter_dim` (i64) — axis each producer splits into `C` -chunks (R9). `gather_dim` (i64) — axis along which each consumer -concatenates the chunks it received, in ascending producer local-index -order (R12). The two may be equal (pure ownership transpose along one -axis) or different (reshape-transpose, e.g. split heads and regather -sequence). - -**Result type.** `T_c_i` is `T_p_i` with the `scatter_dim` extent divided -by `C` and the `gather_dim` extent multiplied by `P`. In the square case -(`P == C` and `scatter_dim == gather_dim`) `T_c_i == T_p_i` (§4). +consume, both `split_dimensions` and `concat_dimensions`, no region, no +identity. + +**Attributes.** `split_dimensions` (`i64` array) — axes each producer +splits into `C` chunks (R9). `concat_dimensions` (`i64` array) — axes +along which each consumer concatenates the chunks it received, in +ascending producer local-index order (R12). Both are flattened in list +order per §4. The two sets may be equal (pure ownership transpose along +one axis set) or different (reshape-transpose, e.g. split heads and +regather sequence); §9.3 shows both readings occur in measurement. + +**Result type.** `T_c_i` is `T_p_i` with the flattened `split_dimensions` +extent divided by `C` and the flattened `concat_dimensions` extent +multiplied by `P`. In the square case (`P == C` and +`split_dimensions == concat_dimensions`) `T_c_i == T_p_i` (§4). **Semantics.** Consumer with local index `l_c` receives, from each producer `l_p`, the slice `[l_c*chunk : (l_c+1)*chunk]` of that -producer's tensor along `scatter_dim` (`chunk = T_p[scatter_dim] / C`), -and concatenates those `P` slices along `gather_dim` in ascending `l_p` -order. +producer's tensor in the flattened `split_dimensions` space +(`chunk = E(split_dimensions) / C`), and concatenates those `P` slices in +the flattened `concat_dimensions` space in ascending `l_p` order. ```mlir %out_1, ..., %out_N = ktdp.inter_tile_all_to_all(%future) consumer_tiles_per_group = , - scatter_dim = , - gather_dim = , + split_dimensions = , + concat_dimensions = , producer_dependency_per_consumer = // optional; default: all producers : !ktdp.tile_future -> T_c_1, ..., T_c_N ``` @@ -695,27 +803,32 @@ already expressible as `consume` + a bijective dependency set (§7.4.2); mechanisms do not overlap. **Generalizing `scatter` to `P > 1` is not an alternative.** Adding a -`gather_dim` and lifting R8 on `scatter` *is* `all_to_all` under another +concat axis set and lifting R8 on `scatter` *is* `all_to_all` under another name; it hides the multi-producer wait inside `scatter` and gives that op two regimes. A separate op keeps one-op-one-pattern and leaves `scatter`'s `P == 1` contract clean. +**Why not `inter_tile_shuffle`.** The SDSC backend calls the primitive +`shuffle`, but `shuffle` is the *lowering* of several patterns rather than this +one alone (§9), so `all_to_all` is used as the established collective term. + ### 6.6 `ktdp.inter_tile_scatter` — ordered split `combine = none`, `placement = split`, one producer per group (R8), -consumer set free, `scatter_dim`, no region, no identity. +consumer set free, `scatter_dimensions`, no region, no identity. -**`scatter_dim`** (i64) — axis of `T_p` along which the single producer's -tensor is partitioned into `C` equal chunks (R9), one per consumer in -ascending consumer local-index order (§3.3). +**`scatter_dimensions`** (`i64` array) — axes of `T_p` along which the +single producer's tensor is partitioned into `C` equal chunks (R9), one +per consumer in ascending consumer local-index order (§3.3), flattened in +list order per §4. -**Result type.** `T_s_i` is `T_p_i` with the extent along `scatter_dim` -divided by `C`. +**Result type.** `T_s_i` is `T_p_i` with the flattened extent over +`scatter_dimensions` divided by `C`. ```mlir %scattered_1, ..., %scattered_N = ktdp.inter_tile_scatter(%future) consumer_tiles_per_group = , - scatter_dim = + scatter_dimensions = : !ktdp.tile_future -> T_s_1, ..., T_s_N ``` @@ -797,13 +910,15 @@ ktdp.store %C, ... **Layout and partitioning.** `A` and `B` are `tensor<96x64xf16>` in global memory. The kernel computes the column-wise sum of `A + B`, producing a -1-D `tensor<64xf16>`. +`tensor<1x64xf16>` (the leading unit dim is the within-group tile axis, +preserved by the op per §4). The 32 compute tiles form a single group. Tile `t` owns rows `t*3 .. t*3+2` of `A` and `B` — a 3×64 slab each. The per-tile contribution is the row-reduced partial expanded to `tensor<1x64xf16>`, -where the leading unit dimension is the within-group tile axis the op -collapses. Every tile holds the same `%reduced : tensor<64xf16>` (all-reduce case: consumer set = producer set). +whose leading unit dimension is the within-group tile axis. The op preserves +it (§4), so every tile holds the same `%reduced : tensor<1x64xf16>` +(all-reduce case: consumer set = producer set). ```mlir #A_view_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 95 >= 0, d1 >= 0, -d1 + 63 >= 0)> @@ -876,12 +991,12 @@ module { ktdp.yield_partial %partial_2d : tensor<1x64xf16> } - // Reduce: unit dim 0 is the within-group tile axis; the op collapses it. - // Every tile holds the same %reduced : tensor<64xf16> (all-reduce case). + // Reduce: unit dim 0 is the within-group tile axis; the op preserves it. + // Every tile holds the same %reduced : tensor<1x64xf16> (all-reduce case). %reduced = ktdp.inter_tile_reduce(%partial_future) consumer_tiles_per_group = #group_tiles, identity(%add_id : tensor<1x64xf16>) - : !ktdp.tile_future, #all_groups> -> tensor<64xf16> + : !ktdp.tile_future, #all_groups> -> tensor<1x64xf16> { ^bb0(%lhs: tensor<1x64xf16>, %rhs: tensor<1x64xf16>): %init = tensor.empty() : tensor<1x64xf16> @@ -891,8 +1006,7 @@ module { } // Post-reduction: every tile redundantly writes the same value. - %reduced_2d = tensor.expand_shape %reduced [[0, 1]] output_shape [1, 64] - : tensor<64xf16> into tensor<1x64xf16> + // No expand_shape needed — the result already carries the unit dim. %E_view = ktdp.construct_memory_view %E_start, sizes: [1, 64], strides: [64, 1] { coordinate_set = #E_view_set, @@ -902,7 +1016,7 @@ module { access_tile_set = #E_tile_set, access_tile_order = #identity_2d } : memref<1x64xf16> -> !ktdp.access_tile<1x64xindex> - ktdp.store %reduced_2d, %E_access + ktdp.store %reduced, %E_access : tensor<1x64xf16>, !ktdp.access_tile<1x64xindex> return @@ -925,10 +1039,10 @@ There are 32 compute tiles forming 8 groups of 4. For tile `t`, `g = t / 4` and `l = t % 4`. Tile `(g, l)` reads slice `[*, g, l*3 : l*3+3, *]` of `A` and `B` — shape `<128x1x3x64>` each. -The partial is `<128x1x1x64>`: dim 1 is the group axis (preserved), dim 2 -is the within-group tile axis (collapsed by the op to `<128x1x64>`). All -four tiles in a group hold identical values; different groups hold -different values. +The partial is `<128x1x1x64>`: dim 1 is the group axis and dim 2 the +within-group tile axis, both preserved, so the result is `<128x1x1x64>` +too (§4). All four tiles in a group hold identical values; different groups +hold different values. ```mlir #A_view_set = affine_set<(d0, d1, d2, d3) : @@ -1034,13 +1148,13 @@ module { ktdp.yield_partial %partial_4d : tensor<128x1x1x64xf16> } - // Multi-group reduce: dim 2 (within-group tile axis) collapsed. - // Dim 1 (group axis) preserved. Each tile gets its group's <128x1x64>. + // Multi-group reduce: no rank reduction — dims 1 and 2 both preserved. + // Each tile gets its group's <128x1x1x64>. %my_group_result = ktdp.inter_tile_reduce(%partial_future) consumer_tiles_per_group = #group_tiles, identity(%add_id : tensor<128x1x1x64xf16>) : !ktdp.tile_future, #all_groups> - -> tensor<128x1x64xf16> + -> tensor<128x1x1x64xf16> { ^bb0(%lhs: tensor<128x1x1x64xf16>, %rhs: tensor<128x1x1x64xf16>): %init = tensor.empty() : tensor<128x1x1x64xf16> @@ -1052,9 +1166,7 @@ module { } // Post-reduction: each tile writes its group's result to slice [*, g, l, *]. - %my_result_4d = tensor.expand_shape %my_group_result [[0], [1, 2], [3]] - output_shape [128, 1, 1, 64] - : tensor<128x1x64xf16> into tensor<128x1x1x64xf16> + // No expand_shape needed — the result already carries both unit dims. %E_view = ktdp.construct_memory_view %E_start, sizes: [128, 8, 4, 64], strides: [2048, 256, 64, 1] { @@ -1066,7 +1178,7 @@ module { access_tile_set = #E_tile_set, access_tile_order = #identity_4d } : memref<128x8x4x64xf16> -> !ktdp.access_tile<128x1x1x64xindex> - ktdp.store %my_result_4d, %E_access + ktdp.store %my_group_result, %E_access : tensor<128x1x1x64xf16>, !ktdp.access_tile<128x1x1x64xindex> return @@ -1091,12 +1203,14 @@ module { } // Reduce and scatter; each tile receives its own slice along dim 0. -// scatter_dim = 0 → 128-row axis split across 4 tiles; each gets <32x1x64>. +// scatter_dimensions = [0] → 128-row axis split across 4 tiles; each gets +// <32x1x1x64> (rank preserved, §4). %my_chunk = ktdp.inter_tile_reduce_scatter(%partial_future) consumer_tiles_per_group = #all_group_tiles, - scatter_dim = 0, + scatter_dimensions = [0], identity(%add_id : tensor<128x1x1x64xf16>) - : !ktdp.tile_future, #all_groups> -> tensor<32x1x64xf16> + : !ktdp.tile_future, #all_groups> + -> tensor<32x1x1x64xf16> { ^bb0(%lhs: tensor<128x1x1x64xf16>, %rhs: tensor<128x1x1x64xf16>): %sum = linalg.add ins(%lhs, %rhs ...) ... @@ -1122,9 +1236,9 @@ slice `[*, g, l*3 : l*3+3, *]` — shape `<128x1x3x64>`. The per-tile pipeline through to `%partial_4d` (shape `<128x1x1x64>`) is identical to §7.2.2. -The op reduces dim 2 (within-group tile axis, size 1) and scatters dim 0 -(128 / 4 = 32 rows per tile). Tile `(g, l)` ends up with rows -`[l*32 : (l+1)*32]` of group `g`'s reduced `<128x1x64>`. +The op reduces across the group and scatters dim 0 (128 / 4 = 32 rows per +tile), preserving rank (§4). Tile `(g, l)` ends up with rows +`[l*32 : (l+1)*32]` of group `g`'s reduced `<128x1x1x64>`. ```mlir #A_view_set = affine_set<(d0, d1, d2, d3) : @@ -1145,7 +1259,7 @@ The op reduces dim 2 (within-group tile axis, size 1) and scatters dim 0 d1 >= 0, -d1 + 7 >= 0, d2 >= 0, -d2 + 63 >= 0)> -// E access tile per writer: 32x1x64 anchored at [l*32, g, 0]. +// E access tile per writer: 32x1x64 in E's 3-D memref, anchored at [l*32, g, 0]. #E_tile_set = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 31 >= 0, d1 == 0, @@ -1232,14 +1346,14 @@ module { ktdp.yield_partial %partial_4d : tensor<128x1x1x64xf16> } - // Reduce dim 2 (within-group tile axis). Scatter dim 0 (chunk = 32). - // Group axis (dim 1) preserved. Each tile receives <32x1x64>. + // Reduce across the group, then scatter dim 0 (chunk = 32). + // No rank reduction: dims 1 and 2 preserved. Each tile receives <32x1x1x64>. %my_chunk = ktdp.inter_tile_reduce_scatter(%partial_future) consumer_tiles_per_group = #group_tiles, - scatter_dim = 0, + scatter_dimensions = [0], identity(%add_id : tensor<128x1x1x64xf16>) : !ktdp.tile_future, #all_groups> - -> tensor<32x1x64xf16> + -> tensor<32x1x1x64xf16> { ^bb0(%lhs: tensor<128x1x1x64xf16>, %rhs: tensor<128x1x1x64xf16>): %init = tensor.empty() : tensor<128x1x1x64xf16> @@ -1263,7 +1377,12 @@ module { access_tile_set = #E_tile_set, access_tile_order = #identity_3d } : memref<128x8x64xf16> -> !ktdp.access_tile<32x1x64xindex> - ktdp.store %my_chunk, %E_access + // Rank reduction now lives in ordinary code, not the op: collapse the + // within-group tile axis to match E's 3-D layout. + %my_chunk_3d = tensor.collapse_shape %my_chunk [[0], [1, 2], [3]] + : tensor<32x1x1x64xf16> into tensor<32x1x64xf16> + + ktdp.store %my_chunk_3d, %E_access : tensor<32x1x64xf16>, !ktdp.access_tile<32x1x64xindex> return @@ -1386,10 +1505,10 @@ eliminated). // Gather along dim 2; one consumer per group (tile 4g) assembles the four // 3-wide slabs. No combiner, no identity — placement is by within-group -// local index. gather_dim = 2 → 3 * 4 = 12; consumer gets <128x1x12x64>. +// local index. gather_dimensions = [2] → 3 * 4 = 12; consumer gets <128x1x12x64>. %assembled = ktdp.inter_tile_gather(%partial_future) consumer_tiles_per_group = #group_consumer, - gather_dim = 2 + gather_dimensions = [2] : !ktdp.tile_future, #all_groups> -> tensor<128x1x12x64xf16> // The consumer holds the full assembled tensor — ownership via SSA result. ``` @@ -1499,7 +1618,7 @@ module { // assembles the full <128x1x12x64>. No combiner region, no identity. %assembled = ktdp.inter_tile_gather(%partial_future) consumer_tiles_per_group = #group_consumer, - gather_dim = 2 + gather_dimensions = [2] : !ktdp.tile_future, #all_groups> -> tensor<128x1x12x64xf16> @@ -1541,11 +1660,11 @@ module { // Head-parallel consumption: split dim 2 (heads) across the 4 consumers, // regather dim 0 (sequence) from the 4 producers. -// scatter_dim = 2 → 4 / 4 = 1; gather_dim = 0 → 128 * 4 = 512. +// split_dimensions = [2] → 4 / 4 = 1; concat_dimensions = [0] → 128 * 4 = 512. %relaid = ktdp.inter_tile_all_to_all(%partial_future) consumer_tiles_per_group = #all_group_tiles, - scatter_dim = 2, - gather_dim = 0 + split_dimensions = [2], + concat_dimensions = [0] : !ktdp.tile_future, #all_groups> -> tensor<512x1x1x64xf16> // Every tile is both producer and consumer; P == C == 4, so the element count // is conserved (128*4 = 512*1) even though the type changes. @@ -1656,15 +1775,15 @@ module { ktdp.yield_partial %partial_4d : tensor<128x1x4x64xf16> } - // All-to-all: consumer l takes head slice l (scatter_dim = 2, 4 / 4 = 1) + // All-to-all: consumer l takes head slice l (split_dimensions = [2], 4 / 4 = 1) // from each of the 4 producers, and concatenates them along the sequence - // axis (gather_dim = 0, 128 * 4 = 512) in ascending producer local-index + // axis (concat_dimensions = [0], 128 * 4 = 512) in ascending producer local-index // order. The producer's local index picks the destination row block; // the consumer's local index picks the head. No combiner, no identity. %relaid = ktdp.inter_tile_all_to_all(%partial_future) consumer_tiles_per_group = #group_tiles, - scatter_dim = 2, - gather_dim = 0 + split_dimensions = [2], + concat_dimensions = [0] : !ktdp.tile_future, #all_groups> -> tensor<512x1x1x64xf16> @@ -1708,10 +1827,10 @@ module { // Scatter along dim 0; the four tiles per group each receive one 32-row // chunk. No combiner, no identity — placement is by within-group local -// index. scatter_dim = 0 → 128 / 4 = 32; each consumer gets <32x1x64>. +// index. scatter_dimensions = [0] → 128 / 4 = 32; each consumer gets <32x1x64>. %chunk = ktdp.inter_tile_scatter(%whole_future) consumer_tiles_per_group = #all_group_tiles, - scatter_dim = 0 + scatter_dimensions = [0] : !ktdp.tile_future, #all_groups> -> tensor<32x1x64xf16> // Each consumer holds its own 32-row slice — ownership via SSA result. ``` @@ -1822,7 +1941,7 @@ module { // receives one 32-row chunk. No combiner region, no identity. %chunk = ktdp.inter_tile_scatter(%whole_future) consumer_tiles_per_group = #all_group_tiles, - scatter_dim = 0 + scatter_dimensions = [0] : !ktdp.tile_future, #all_groups> -> tensor<32x1x64xf16> @@ -1868,7 +1987,11 @@ The legality pass (`lib/Conversion/ConvertToKTIR/KTIRCheckLegality.cpp`, | R4 | `inter_tile_reduce` | every `p` covered by some dep | `KTIRCheckLegality.cpp:163–174` | **Not yet implemented:** R1, R5, R6, R7, R8, R9, R10, R11, R12, and -R3/R4/R13/R14 for every op other than `reduce`. R5 and R7 are enforced in +R3/R4/R13/R14 for every op other than `reduce`. Of the ops §9.3 shows the +backend requires — `gather`, `all_to_all`, `scatter`, `consume` — none has +a verifier today, and R9/R12 in particular are stated over *flattened* +multi-axis extents (§4), so implementing them means validating an axis +list, not a single index. R5 and R7 are enforced in the Torch-Spyre SDSC planner (`_compatible_partitions`) but are absent from the KTIR verifier entirely — the gap exists at both the spec and the implementation level. @@ -1892,59 +2015,172 @@ implementation level. ## 9. Backend pattern catalogue — non-normative -This section records relayout patterns observed in the Torch-Spyre backend -(`scratchpad/lx_relayout.py`) and how they map onto the ops above. It is -descriptive, not normative: nothing here constrains the op definitions, -and the mapping column is a proposal for lowering rather than a -guarantee. Rows whose split axes are unknown are omitted. - -All implemented patterns emit a **single** SDSC `opfunc = "shuffle"` whose -entire payload is two `coreIdToWkSlice_` tables — one per tensor in -`coordinates_` — describing per-core ownership before and after the -movement. Classification uses -`gathered_dims = src_syms − dst_syms`, -`scattered_dims = dst_syms − src_syms`, -`factor = num_cores // prod(dst splits)`. - -| ID | SenDNN op | Src `work_div` | Dst view | gathered | scattered | factor | KTIR op | Backend status | -|---|---|---|---|---|---|---|---|---| -| P01 | all-gather | `{H:8, Lk:4}` | replicated (all cores) | H, Lk | — | N/A | `inter_tile_gather` (all consumers) | missing — replication not supported | -| P03 | grouped all-gather | `{H:8, Lk:4}` | `{H:8}` | Lk | — | 4 | `inter_tile_gather` | #3440 (open PR) | -| P04 | grouped all-gather | `{Lk:32}` | `{H:8}` | Lk | H | 4 | `inter_tile_gather` | #3440 (open PR) | -| P06 | all-to-all | `{H:8, Lk:4}` | `{Lk:32}` | H | — | 1 | `inter_tile_all_to_all` | main (#3439) | -| P08 | all-to-all (axis transpose) | `{A:4, B:8}` | `{A:8, B:4}` | — | — | 1 | `inter_tile_all_to_all` + explicit coord map | missing — axis swap inexpressible | -| P14 | all-to-all | `{H:8, Lq:4}` | `{Lq:32}` | H | — | 1 | `inter_tile_all_to_all` | main (#3439) | - -- **P03 / P04** show that `gathered_dims ≠ ∅` with `factor > 1` maps to - `inter_tile_gather`. For P04, `scattered_dims = {H}` (H is introduced at - the destination) — a combined gather+scatter in one shuffle step, so the - KTIR op must express both axes. -- **P06 / P14** show `inter_tile_all_to_all` at `factor = 1`: - `gathered_dims = {H}`, `scattered_dims = ∅`, with H contracted into the - 1-D destination axis. -- **P08 is the case `all_to_all` with dim attributes cannot express.** Both - sides split the same two dims with swapped counts (`[4,8] → [8,4]`). The - `coreIdToWkSlice_` tables differ because the mixed-radix odometer - ordering changes, but no single `scatter_dim`/`gather_dim` pair captures - the transformation. Marked out of scope for `all_to_all`; a verifier - should reject non-decomposable transpositions rather than silently - mis-lower them. §10.4 (Option B) is the eventual fix. -- **P01** (full replication) is blocked because `_compatible_partitions` - requires distinct slices per destination core. It maps to - `inter_tile_gather` with `consumer_tiles_per_group = all`, but needs new - backend support for non-bijective shuffle. - -**Classification decision rule** (from `TensorArg.work_division` on a -relayout identity OpSpec): - -| `gathered_dims` | `scattered_dims` | `factor` | `slot_exprs_differ` | KTIR op | +Relayout patterns measured in the Torch-Spyre backend, and which op of §6 +each one needs. Descriptive, not normative: its purpose is to establish +**which ops a lowering must actually emit**, and with what attribute +arity. + +Every relayout compiles to one SDSC entry with `opfunc = "shuffle"`, whose +payload is a pair of per-core ownership tables — what each core owns before +the movement and after. A table maps `core_id → {axis: slice_index}`; a +core's **region** is the intersection of its per-axis ranges, and an axis +absent from an entry is uncut. Classifying a pattern means deciding, from +those two tables, which delivery op expresses the same movement. + +**Coarsened and refined.** Let `Ns(a)` and `Nd(a)` be the number of slices +axis `a` is cut into at the source and destination (1 if absent). Then `a` +is **coarsened** when `Ns(a) > Nd(a)` — fewer, larger pieces after the +move, so data must be *assembled* along it — and **refined** when +`Nd(a) > Ns(a)`, so data must be *split* along it. Write + +``` +C = {a : Ns(a) > Nd(a)} # coarsened — assemble +R = {a : Nd(a) > Ns(a)} # refined — split +``` + +These two sets are what select the op. Everything else the classification +needs is a product down the axes: `len(src_regions) = prod(Ns(a))`, +`len(dst_regions) = prod(Nd(a))`, and `components = prod(gcd(Ns(a), Nd(a)))` +— a **component** being a maximal set of source and destination regions +that exchange only among themselves, which for `all_to_all` is exactly a +permute group. + +Two cautions. Region counts classify; **core** counts mislead, since +several cores may hold one region and a region count need not divide the +core count. And axes must be aligned by **physical axis, not by label** — +labels differ between the two sides, and this is the only place a wrong +answer can enter. + +**Axis names versus axis indices.** Axis sets are written below with the +backend's symbol names (`mb`, `in`, `out`, …), the vocabulary the ownership +tables speak. The op attributes of §6 are `i64` arrays of *axis indices* +into `T_p`; a lowering resolves each named axis to its position in the +producer tile type, preserving list order, which §4 fixes as slowest- to +fastest-varying. + +### 9.1 The decision table + +| # | condition | result | +|---|---|---| +| — | `prod(Ns(a)) != len(src_regions)` or `prod(Nd(a)) != len(dst_regions)` | **not a work-division pair** — enumerate regions instead. Check first | +| — | any axis ragged (non-uniform overlap) | *insufficient information* — no single op has uniform dependency-set cardinality (R6) | +| 1 | `C = ∅` and `R = ∅` | regions identical: `no op needed` if every core's region is its own, else `inter_tile_consume` — a **relocation**, or a **broadcast** where destination regions are shared | +| 2 | `C = ∅`, `R ≠ ∅` | `inter_tile_scatter`, `scatter_dimensions = R` | +| 3 | `C ≠ ∅`, `R ≠ ∅`, `prod(Nd(a)) == num_cores` | `inter_tile_all_to_all`, `split_dimensions = R`, `concat_dimensions = C`; one group per component | +| 4 | `C ≠ ∅` otherwise | `inter_tile_gather`, `gather_dimensions = C`; consumers per group = the destination region's holders | + +The dimension attributes are **the axis sets themselves**, in the order §4 +fixes — which is why they must be list-valued (§1.1). Rows 2–4 return +*insufficient information* when a source region has several holders, since +R8 admits one producer per group and the tables do not say which transmits +(§10.2). + +Row 3's `prod(Nd(a)) == num_cores` is the one irreducibly global test: +holding the division fixed and varying the core count changes the op, so no +per-axis quantity can see it. It is also weaker than asking whether any +destination region is shared — the two part company on every row-4 output +under an idle-core reading (§10.2), which is why the guard rows run first. + +For `all_to_all`, group sizes follow from the component count — each tile +contributes `M` slices and receives `K`: + +```python +M = len(dst_regions) / components # consumers per group +K = len(src_regions) / components # producers per group +``` + +Note the sides: `M` counts *destination* regions. Inverting them is +invisible on a square exchange and wrong on every other. + +### 9.2 Worked example + +**A real all-to-all** — `{mb:8, out:4} → {mb:32}` on 32 cores. + +| axis | `Ns(a)` | `Nd(a)` | relation | `gcd` | |---|---|---|---|---| -| ∅ | ∅ | 1 | false | no-op | -| ∅ | ∅ | 1 | true | `inter_tile_all_to_all` + explicit coord map | -| non-∅ | ∅ or non-∅ | > 1 | — | `inter_tile_gather` | -| non-∅ | ∅ | 1 | — | `inter_tile_all_to_all` | -| ∅ | non-∅ | 1 | — | `inter_tile_scatter` | -| ∅ | ∅ | 1 | — (replication) | `inter_tile_gather` (all consumers) | +| `mb` | 8 | 32 | refined | 8 | +| `out` | 4 | 1 | coarsened | 1 | + +32 regions a side, `components = 8`. Both sets non-empty and +`prod(Nd(a)) = 32 = num_cores`, so row 3: **`inter_tile_all_to_all`, +`split_dimensions = [mb]`, `concat_dimensions = [out]`**, with **8 groups +and `M = K = 4`** — eight independent 4-way exchanges, not one 32-way. + +By contrast, source and destination cutting *different* axes, +`{Lk:32} → {H:8}`, makes both sets non-empty yet `prod(Nd(a)) = 8 ≠ 32`, so +row 4 gathers — the only shape where the global comparison does the work, +and unattested in the measurements. + +### 9.3 Measured use cases + +51 measured relayouts, each `opfunc = "shuffle"` on 32 cores, carrying +explicit regions and per-region core sets — so replication versus idleness, +which a work division can never settle, is read directly. + +| use case | n | src | dst | C | R | KTIR op | +|---|---|---|---|---|---|---| +| all-gather to every core | 3 | `{in:2, out:8, x:2}` | `{}` | `in`,`out`,`x` | — | `inter_tile_gather`, `gather_dimensions = [in, out, x]` | +| all-gather, 4 cores idle | 1 | `{in:32}` | `{}` | `in` | — | `inter_tile_gather`, `gather_dimensions = [in]` | +| grouped gather, drop an axis | 6 | `{mb:8, in:4}` | `{mb:8}` | `in` | — | `inter_tile_gather`, `gather_dimensions = [in]` | +| grouped gather, coarsen one axis | 15 | `{mb:32}` | `{mb:8}` | `mb` | — | `inter_tile_gather`, `gather_dimensions = [mb]` | +| grouped gather, coarsen one axis | 3 | `{mb:16}` | `{mb:8}` | `mb` | — | `inter_tile_gather`, `gather_dimensions = [mb]` | +| all-to-all, square | 6 | `{mb:8, out:4}` | `{mb:32}` | `out` | `mb` | `inter_tile_all_to_all`, split `[mb]` / concat `[out]`; 8 groups, `M=K=4` | +| all-to-all, square | 3 | `{x:8, mb:4}` | `{x:32}` | `mb` | `x` | `inter_tile_all_to_all`, split `[x]` / concat `[mb]`; 8 groups, `M=K=4` | +| all-to-all, **non-square** | 1 | `{mb:16}` | `{mb:8, out:4}` | `mb` | `out` | `inter_tile_all_to_all`, split `[out]` / concat `[mb]`; 8 groups, **`M=4, K=2`** | +| pure split | 12 | `{y:16}` | `{y:32}` | — | `y` | `inter_tile_scatter`, `scatter_dimensions = [y]` | +| broadcast | — | `{h:8}` on 8 cores | `{h:8}` × 4 cores | — | — | `inter_tile_consume`, consumer set widened to the 4 holders | +| selection, not a partition | 1 | `{mb:8, out:4}` | *selection* | — | — | **not a work-division pair** — guard row | + +Divisions are the **measured** ones, in `layoutDimOrder_` order. The +broadcast row comes from separate broadcast work (PR #4061), not the 51. + +**Which ops this requires.** Four of the six delivery ops, with these +arities: + +| Op | measured files | arity needed | +|---|---|---| +| `inter_tile_gather` | 28 | up to **3 axes** | +| `inter_tile_all_to_all` | 10 | 1 axis each side, but **non-square** `M ≠ K` | +| `inter_tile_scatter` | 12 | 1 axis | +| `inter_tile_consume` | broadcast work | — | + +`inter_tile_reduce` and `inter_tile_reduce_scatter` are exercised by none +of the 51 — expected, since a relayout moves ownership without combining +values. They stay required by §7.2 and §7.3, which are not relayouts. + +Two consequences for implementation order: `gather` carries the most +measured weight *and* the widest arity, so its multi-axis path cannot be +deferred; and `all_to_all`'s non-square case is measured, not +hypothetical, so `P == C` is not a safe simplifying assumption. + +**What the measurements also establish.** + +- **A three-axis concat exists**, so §4's flattening order must be fixed + over three axes — list-valued attributes are a requirement of a named + pattern, not a corner case. +- **Idleness is the norm, replicated sources are rare.** Every source + region in the 51 has one holder; 16 files have fewer source regions than + cores and all resolve to single holders plus idle cores. One destination + region is held by 28 cores with 4 idle. +- **The contiguity assumption is false** — four-core destination groups are + contiguous in 9 files and strided in 15, so the core-to-region map is not + a function of the division. This is why §3.3 defines position by + ascending tile id within the set rather than by contiguity, and why + `producer_tiles_per_group` / `consumer_tiles_per_group` must come from + the tables rather than the axis counts. +- **The stick-level assumption holds** — 20 files coarsen or refine the + stick axis and every piece size on it is an exact stick multiple. +- **One file is a selection, not a partition** — 1/512 coverage, which is + what the §9.1 validity guard is for. A selection is not a delivery: it + needs a select op before one. + +**Still unmeasured.** Two recorded patterns match no file: one needs a side +with 8 active cores (measured counts are 1, 16, 28, 32), and one is an axis +transpose `{A:4, B:8} → {A:8, B:4}`. The transpose classifies under row 3 — +`C = {B}`, `R = {A}`, 16 components, `M = K = 2` — but with nothing coarsened +or refined in the *region* sense it could equally be read as row 1, and only +divisions keyed by physical axis settle which. Uniformity holds on all 51 +measured files, so R6/R7 are so far confirmed rather than assumed. **Fused relayout is deferred.** Relayout stays a separate preceding op and fusion is a lowering concern. The backend structurally cannot fuse them @@ -1953,120 +2189,172 @@ restickified weights are explicitly barred as shuffle sources. --- -## 10. Open questions and extensions +## 10. Open questions ### 10.1 Must a consumer also be a producer? -Open for `consume`, `reduce`, `reduce_scatter`, `gather`, and -`all_to_all`; **resolved for `scatter`** — no (§6.6). - -The current implementation answers *yes* for `reduce` and enforces it -(R13, `KTIRCheckLegality.cpp:107–117`), whose error text names this -question explicitly. That is one op's implementation choice, not a design -conclusion for the family. The related R14 mode gate — `reduce` supports -all-reduce (`C == P`) and reduce-to-one (`|C| == 1`) but rejects a strict -multi-tile consumer subset — is likewise a present restriction awaiting a -decision. - -Deciding this per op is what the `?` cells in §5 record. - -### 10.2 Delivery-op placement - -Whether the verifier should enforce that a delivery op appears only inside -a guard matching `consumer_tiles_per_group`, or whether that is left to -lowering. If the union of consumer sets equals the set of all executing -tiles, no guard is needed; otherwise a tile outside the consumer set that -reaches the delivery op would be a verifier error. - -### 10.3 Multiple delivery ops per future - -R2 restricts a `!ktdp.tile_future<...>` value to exactly one delivery use. -A natural extension would allow several delivery ops to consume the same -future, each declaring its own `producer_dependency_per_consumer` — one -`ktdp.inter_tile_produce` serving two independent deliveries, e.g. one -waiting on the first half of the producers and another on the second half. - -**Expressiveness gain.** Patterns that today need two separate -`ktdp.inter_tile_produce` ops with identical producer regions collapse to -one produce plus two delivery ops, removing redundant producer-side code -and making the shared production explicit in the IR. - -**Verification cost.** Single-use keeps R4 (coverage) local: the verifier -inspects one delivery op to confirm every producer tile is covered. With -multiple uses, coverage becomes global — for every group `g` and producer -`p ∈ producer_tiles_per_group(g)`, at least one consumer `c` across *any* -delivery op must satisfy `producer_dependency_per_consumer(p)[c, g]`. That -requires collecting and unioning the dependency sets from all uses of the -SSA value before checking, a def-use traversal rather than a local per-op -check. R5 (pairwise disjointness) would have to become cross-op too. - -**Lowering cost.** Each delivery op declaring a dependency set introduces -its own point-to-point signals. A producer `p` may then need to signal -multiple consumers across different delivery ops, and lowering must emit -each signal exactly once and receive it exactly once per dependent -consumer. In full-barrier mode, multiple delivery ops on one future also -require handling duplicate barrier waits: a producer-side barrier cannot -be issued until every dependent delivery op is ready to receive. - -Given that cost, the current design requires separate -`ktdp.inter_tile_produce` ops for separate delivery concerns. If real use -cases demand shared production, the restriction can be relaxed. - -### 10.4 Option B — explicit coordinate map - -Replace the `scatter_dim`/`gather_dim` attribute pair with a single -source-to-destination affine map, subsuming all four placement values in -one mechanism. This is the shape interface-specs PR 14 already uses -(`SHUFFLE` as source/destination coordinate sets), and the backend already -speaks per-core partitionings (`coreIdToWkSlice_` tables) rather than dim -attributes — so Option B is arguably closer to the existing contract, and -it is the only form that expresses P08 (§9). - -**Direction: dim attributes first, Option B recorded as the long-term -target.** The dim-attribute form is additive and reviewable on its own, -and it covers every pattern the backend implements today. Note the honest -counter-argument: because none of the six delivery ops is built yet, the -usual "Option A is cheaper because it is incremental" argument is weaker -here than normal. - -### 10.5 Post-v1 extensions to `all_to_all` - -- **`all_to_all_v`** — uneven shard sizes, i.e. per-consumer split extents - instead of a uniform `T_p[scatter_dim] / C`. This relaxes R9 into a - per-consumer size list and needs a variadic size attribute; no current - backend pattern requires it. -- **Multi-axis relayout** — splitting or gathering along more than one - axis in a single op. Expressible today only as a sequence of - `all_to_all` ops; Option B (§10.4) is the natural home for it. -- **`inter_tile_shuffle` as a naming alias** — the SDSC backend calls the - primitive `shuffle`. `all_to_all` is kept as the op name because it is - the established collective term and because `shuffle` is the *lowering* - of several patterns, not just this one (§9). +**What turns on it:** whether the verifier rejects a delivery op whose consumer +set is not contained in its producer set. That check exists and runs today — R13 +for `reduce` (`KTIRCheckLegality.cpp:107–117`) — so whoever implements +`gather`, `all_to_all` or `reduce_scatter` must decide whether to extend it, +and the answer changes which programs are legal. + +**Why it is unresolved:** `reduce`'s *yes* is one op's implementation choice, +made when it was the only delivery op. `scatter`'s *no* is settled and +argued (§6.6) — a consumer that receives a slice contributes nothing, so +there is nothing to miss. Neither generalizes: `gather` and `all_to_all` +assemble, so a non-producing consumer is coherent for them in a way it is not +for a reduction. The `?` cells in §5 are exactly the ops still to decide, and +R14's mode gate (all-reduce or reduce-to-one, no strict multi-tile subset) is +a present restriction on `reduce` awaiting the same call. + +### 10.2 Two things a work division cannot settle + +Both are §9.1 escape hatches, and both need the per-region core sets rather +than the division. + +**Producer election.** Rows 2–4 return *insufficient information* when a +source region has several holders: R8 admits one producer per group, and the +tables record who *holds* a region, not who *transmits* it. Unforced by +measurement — every source region across the 51 files has a single holder — so +the choice between electing a canonical producer (lowest tile id), requiring +the frontend to pick, and rejecting replicated sources can wait. + +**Replication versus idleness.** Row 3's `prod(Nd(a)) == num_cores` is weaker +than asking whether a destination region is shared, and the two part company on +every row-4 output: a region held by several cores is either genuine +replication or one consumer plus idle cores. Both occur in measurement — one +region held by 28 cores with 4 idle, and the broadcast work genuinely +replicating — so the distinction is real. What is open is whether the op +surface should mark it, or whether `consumer_tiles_per_group` naming the actual +holders suffices. This is why §9.1's membership step runs before the +core-count test. + +### 10.3 Physicalization: which ops are layout-transparent + +Raised by Triton issue #92. **Physicalization** rewrites a tensor to a stick +layout, splitting one axis by the stick size with the chunk count at the front +and the within-stick extent at the back: + +``` +logical [16, 64], stick on the 64 axis, stick = 32 + → physical [64/32, 16, 32] = [2, 16, 32] +``` + +Rank grows by one and the logical stick axis becomes **two non-adjacent +physical axes**. Nothing in this repository represents a stick layout today, so +what follows is a design obligation, not current behaviour. + +**Why today's `reduce` is transparent.** It carries no axis-index attribute and +pins results to partials (`KTDP.td:168-171`), so physicalizing the input carries +the result along with no op knowledge — the "elementwise" property. Issue #92's +failure is adjacent: the `identity` operand is tied to results +(`KTDP.td:172-174`) but materialized at logical rank before any layout pass +runs. That is a *propagation* bug, and since the identity is a splat, +re-materializing it at the right type is shape-agnostic by construction. + +**The split follows §1.1 exactly**, because §4 makes result type a function of +`placement` alone and `replicate` is the only placement naming no axis set: + +| Op | placement | Axis attrs | Result vs partial | Transparent? | +|---|---|---|---|---| +| `consume` | replicate | — | identical | **yes** | +| `reduce` | replicate | — | identical | **yes** | +| `reduce_scatter` | split | `scatter_dimensions` | ÷ `C` | no — attrs, shape, identity | +| `gather` | concat | `gather_dimensions` | × `P` | no — attrs, shape | +| `all_to_all` | permute | `split_`/`concat_dimensions` | ÷ `C` and × `P` | no — attrs, shape | +| `scatter` | split | `scatter_dimensions` | ÷ `C` | no — attrs, shape | + +`consume` joins `reduce`. `scatter` does **not**, despite being copy-only: it +divides an extent and names the axis it divides. No rank reduction (§4) is +load-bearing here — a collapse is an axis-*position* operation, so a `reduce` +that collapsed would not be transparent either. + +**What §4's rules already settle.** A dim attribute naming a sticked axis +becomes *two* indices (`[1]` → `[0, 2]`), which only the list-valued form can +express, and §4's slowest-to-fastest order is exactly what the stick layout +produces — physical `(c, m, s)` holds logical `n = c*32 + s`. The floordiv rule +fixes which axis absorbs the ×`P` or ÷`C`, and R9 on the floordiv axis is then +precisely the stick-multiple check: `scatter` with `C = 4` on a 2-chunk axis +fails `2 % 4`, correctly rejecting a logical result of `[16,16]` that is half a +stick. `E(D)` itself is invariant (`2 × 32 = 64`), so R9 and R12 cannot change +verdict on the flattened extent — provided a rewrite lists *both* halves of a +split axis; listing one half is simply the wrong rewrite, and R9 catches it. + +**Axis indices shift, and physicalization is where that is handled.** The +chunk-count axis is inserted at the *front*, so logical axis 0 of `[16,64]` +becomes physical axis 1: no dim attribute survives untouched, including one +naming an axis physicalization never split. Left unshifted, +`gather_dimensions = [0]` names the chunk axis instead — a valid, distinct +index, so R9/R12 pass and the op is silently wrong. + +The remedy needs no new mechanism. Physicalization **is** the logical-to-physical +mapping, so the pass that applies it already knows which logical axis was split, +the stick size, and where every logical axis landed — exactly the information a +dim attribute needs. The attributes name logical axes as authored, and the pass +rewrites them in the same step it retypes the tensors: `[1]` → `[0, 2]` for the +split axis, `[0]` → `[1]` for the shifted one. Nothing downstream re-derives it, +and the ops stay layout-agnostic, which matches §9's framing where axes are +backend symbol names until lowering. + +This is not the shape of issue #92. There the `identity` was missed because +`retypeChain` walks forward along operand 0 and never reaches a sibling +operand — an incompleteness in *which values* the pass visits. Attributes sit on +the op the pass is already rewriting, so they are in reach by construction; what +is required is that the mapping be applied to them, not that it be discovered +somewhere else. + +**What is still open.** + +1. **R12's per-axis clause gains teeth.** Single-axis lists make it trivial; + `[0,2]` makes it two checks. Stick size depends on element type (32 for f32, + 64 for f16), so variadic roles with mixed types can have equal products and + unequal per-axis extents — reachable, since §3.7 requires all roles to share + one axis set. +2. **`reduce_scatter`'s identity.** Its identity must match `T_p` while its + result is `T_p` split by `C`, so issue #92's fix is needed there in a harder + form — and the shipped constraint must be retargeted from results to partials + (R11, §5). +3. **The floordiv rule against a sticked multi-axis pattern**, once one is + measured. §9.3's three-axis concat has no sticked axis, so it does not test it. --- -## Appendix A. Relationship to the pre-existing ops - -| Existing op | Maps to in this design | -|-------------|------------------------| -| `inter_tile_produce` | `ktdp.inter_tile_produce` — `consumer_tiles_per_group` moved to the delivery op; `producer_tile_per_group` → `producer_tiles_per_group` (generalized to multi-producer) | -| `inter_tile_consume` | `ktdp.inter_tile_consume` — unchanged semantics | -| `inter_tile_reduce` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce` — producer block removed from the reduction op | -| `inter_tile_reduce_scatter` | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce_scatter` — producer block removed from the reduction op | - -`ktdp.inter_tile_gather` (§6.4), `ktdp.inter_tile_all_to_all` (§6.5), and -`ktdp.inter_tile_scatter` (§6.6) have no pre-existing counterparts. The -earlier ops offered none of ordered-concatenation delivery (gather), -split-and-reassemble delivery (all-to-all), or single-producer -ordered-partition delivery (scatter). - -The `!ktdp.tile_future` type is shared across all ops; its -`#groups` parameter carries the group set (§1.3). - -The previous `ktdp.inter_tile` single op (Approach B draft) is replaced by -this seven-op design. `ktdp.inter_tile` carried producer and optional -combiner regions in one op, with `consumer_tiles_per_group` determining the -delivery mode. Splitting production from delivery makes the mode a choice -of op rather than an inference over attribute combinations — which is what -lets §3 state the shared machinery once and §6 reduce each op to its own -cells. +## Appendix A. Relationship to what exists today + +**Two of the seven ops exist.** `include/ktir/Dialect/KTDP/KTDP.td` defines +`ktdp.inter_tile_produce` and `ktdp.inter_tile_reduce`, plus the +`ktdp.yield_partial` and `ktdp.yield_reduced` terminators. That is all — +the other five delivery ops are new work, not revisions of existing ops. + +| Op | Status today | This design | +|---|---|---| +| `inter_tile_produce` | exists (`KTDP.td:107`) | already matches: carries `producer_tiles_per_group` and no consumer set, returns a future | +| `inter_tile_reduce` | exists (`KTDP.td:165`) | already matches: consumes the future, carries `consumer_tiles_per_group` and a reducer region only | +| `inter_tile_consume` | **not implemented** | new (§6.1) | +| `inter_tile_reduce_scatter` | **not implemented** | new (§6.3) | +| `inter_tile_gather` | **not implemented** | new (§6.4) | +| `inter_tile_all_to_all` | **not implemented** | new (§6.5) | +| `inter_tile_scatter` | **not implemented** | new (§6.6) | + +`inter_tile_consume` and `inter_tile_reduce_scatter` appear in the current +tree only as prose: `KTDP.td:70` names them as unbuilt future work, and +`KTDPTypes.td:235` lists them among the delivery ops the future type is +*intended* to serve. Neither has an op definition, so this document is a +specification for five new ops rather than a restructuring of existing +ones. + +Cross-checking against §9.3: of the five unbuilt ops, `gather`, +`all_to_all`, `scatter` and `consume` are required by measured relayouts, +while `reduce_scatter` is required only by the reduction patterns of §7.3. + +**The `!ktdp.tile_future` type** already exists +(`KTDPTypes.td:231`) and is shared across all ops; its `#groups` parameter +carries the group set (§1.3). + +**The earlier single-op draft.** A `ktdp.inter_tile` op carrying producer +and optional combiner regions in one op, with `consumer_tiles_per_group` +determining the delivery mode, was drafted but never landed. Splitting +production from delivery makes the mode a choice of op rather than an +inference over attribute combinations — which is what lets §3 state the +shared machinery once and §6 reduce each op to its own cells. From d5f87c344fc9aaf30cbf34d8eb342e6c7251ba15 Mon Sep 17 00:00:00 2001 From: Takuya Nakaike Date: Fri, 28 Aug 2026 01:23:39 +0000 Subject: [PATCH 05/10] =?UTF-8?q?docs(inter-tile):=20name=20every=20table;?= =?UTF-8?q?=20fix=20the=20=C2=A74=20claim=20about=20SDSC=20shuffles?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two small follow-ups on top of the measured-data rework. Table names. Every table gets an explicit "**Bold name.** One row per ..." line so prose can refer to it by name instead of "the table": semantics matrix (§1.1), coverage table (§1.2), symbol table (§3.1), sharing table (§3.7), type-rule table (§4), verification matrix (§5), the two dependency tables (§7.4.1, §7.4.2), implemented-rule table (§8) and migration table (Appendix A). §8's two references to "the §5 matrix" follow the new name. §4 conservation case. The equal-type all-to-all — the distributed transpose, where split_dimensions == concat_dimensions — was described as "the uniform one-to-one shuffle the SDSC backend emits today". The measured data contradicts that: every measured all-to-all splits and concats *different* axes (§9.3), so P == C conserves the element count while the type still changes, and the equal-type case is recorded but unmeasured. Co-Authored-By: Claude Opus 5 Signed-off-by: Takuya Nakaike --- docs/inter-tile-communication.md | 35 ++++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index e56009d..119b804 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -60,6 +60,8 @@ means a tensor or tile axis. `replicate` | `concat` | `permute` | `split`. - **cardinality** — producer tiles per group × consumer tiles per group. +**Semantics matrix.** One row per delivery op. + | Op | combine | placement | producers/grp | consumers/grp | dim attrs | region | identity | |---|---|---|---|---|---|---|---| | `consume` | none | replicate | 1 | free | — | — | — | @@ -103,6 +105,8 @@ widening `consumer_tiles_per_group`. The consumer set is therefore a column here, since it is what distinguishes gather from all-gather and all-to-all from scatter. +**Coverage table.** One row per named collective pattern. + | Pattern | Producers/grp | Consumers/grp | Delivery op | Result per consumer | |---------|---------------|---------------|-------------|---------------------| | Broadcast | 1 | free | `inter_tile_consume` | full copy | @@ -233,6 +237,8 @@ only per-op deltas; where §6 is silent, this section governs. ### 3.1 Notation +**Symbol table.** These names are used unqualified throughout. + | Symbol | Meaning | |---|---| | `T_p_i` | the partial type of role `i`, as yielded by `ktdp.yield_partial` | @@ -393,7 +399,9 @@ per-tile mode. Every delivery op produces `N` variadic SSA values, one per partial-tensor role. The values are **per-tile-valued**: each consumer tile holds its own result value when the op completes. Whether tiles in -the same group hold the *same* value is a property of the placement: +the same group hold the *same* value is a property of the placement. + +**Sharing table.** One row per placement value. | placement | tiles in one group hold | tiles in different groups hold | |---|---|---| @@ -418,7 +426,9 @@ role's result type follows the §4 rule independently. ## 4. Placement algebra and type rules Result types are a function of the placement value alone. There are four -formulas, applied per role `i` to `T_p_i`: +formulas, applied per role `i` to `T_p_i`. + +**Type-rule table.** One row per placement value. | placement | result type derived from `T_p` | |---|---| @@ -507,8 +517,10 @@ themselves, in this order. result has the same element count as `T_p` — one axis is divided and another multiplied by the same factor — so a square all-to-all is a pure redistribution of ownership. If additionally `split_dimensions == concat_dimensions`, -the result *type* equals `T_p`: the distributed transpose, which is the -uniform one-to-one shuffle the SDSC backend emits today (§9). +the result *type* equals `T_p`: the distributed transpose. That equal-type +case is not what the SDSC backend emits today — every measured all-to-all +splits and concats *different* axes (§9.3), so `P == C` conserves the +element count while the type still changes. **Why `split` divides an honest data axis.** Every splitting op divides an extent of an axis the partial already has, so the types stay honest: @@ -523,6 +535,8 @@ Principle: **each rule has exactly one owner and one statement; applicability is a column, not a restatement.** "Owner" is the op that carries the attribute the rule constrains. +**Verification matrix.** One row per rule, one column per delivery op. + | Rule | Owner | consume | reduce | red_scat | gather | all_to_all | scatter | |---|---|---|---|---|---|---|---| | R1 group disjointness (§2.1) | produce | y | y | y | y | y | y | @@ -1399,6 +1413,8 @@ and `4g+3` are consumers. Each consumer depends on its dedicated producer (`4g+2` ← `4g`, `4g+3` ← `4g+1`), so the pairing is `p = c - 2` — a constant relative offset that does not depend on the group index `g`. +**Dependency table** for group 0: + | group | producer | consumer | |-------|----------|----------| | 0 | 0 | 2 | @@ -1447,6 +1463,8 @@ asking (different consumers within the group have different mirrors), and `11`, `19`, ... for groups `0`, `1`, `2`, ..., so `g` cannot be eliminated). +**Dependency table**, first two groups: + | group | producer | consumer | |-------|----------|----------| | 0 | 0 | 3 | @@ -1978,6 +1996,8 @@ The legality pass (`lib/Conversion/ConvertToKTIR/KTIRCheckLegality.cpp`, 182 lines) currently walks only `InterTileProduceOp` and `InterTileReduceOp`: +**Implemented-rule table.** One row per check that exists today. + | Rule | Op | Check | Location | |---|---|---|---| | R2 | `inter_tile_produce` | `future.hasOneUse()` | `KTIRCheckLegality.cpp:80–85` | @@ -1996,7 +2016,7 @@ the Torch-Spyre SDSC planner (`_compatible_partitions`) but are absent from the KTIR verifier entirely — the gap exists at both the spec and the implementation level. -**Two asymmetries the §5 matrix forces into the open.** +**Two asymmetries the verification matrix (§5) forces into the open.** 1. R8 is stated as a verifier obligation for `scatter` but is merely conventional for `consume`, even though both ops have the same @@ -2004,7 +2024,7 @@ implementation level. spec asymmetry to decide rather than inherit. 2. R13 and R14 are implemented for `reduce` only, and R13 is the implementation of open question §10.1 (must a consumer also be a - producer?) for that one op. The `?` cells in the §5 matrix are exactly + producer?) for that one op. The `?` cells in that matrix are exactly that question, unresolved: for `scatter` the answer is **no** (§6.6), for `reduce` the current answer is **yes** (enforced), and for `reduce_scatter` / `gather` / `all_to_all` it is undecided. R14's @@ -2327,6 +2347,9 @@ somewhere else. `ktdp.yield_partial` and `ktdp.yield_reduced` terminators. That is all — the other five delivery ops are new work, not revisions of existing ops. +**Migration table.** One row per op of this design, with its state in +`KTDP.td` today. + | Op | Status today | This design | |---|---|---| | `inter_tile_produce` | exists (`KTDP.td:107`) | already matches: carries `producer_tiles_per_group` and no consumer set, returns a future | From bbcc83059f7f26db351c3b1ca888ffc3a42a42a2 Mon Sep 17 00:00:00 2001 From: Takuya Nakaike Date: Fri, 28 Aug 2026 03:00:21 +0000 Subject: [PATCH 06/10] =?UTF-8?q?docs(inter-tile):=20pin=20the=20Ns/Nd=20c?= =?UTF-8?q?ounting=20rule=20and=20state=20=C2=A79.1's=20coverage=20clause?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit §9 already decides the one measured file that is a selection rather than a partition: the guard row's verdict is "not a work-division pair — enumerate regions instead". Nothing here changes that. What a reader could not do was reach that verdict from the stated clauses, for two reasons. First, the preamble defined `Ns(a)`/`Nd(a)` as "the number of slices axis `a` is cut into", which admits two readings — count the distinct slices an ownership table contains, or divide the axis extent by one slice's extent. They agree on 50 of the 51 measured files; on the selection the destination gives `Nd(mb) = 1` counted, since every core names `mb[511:512]`, and `512` divided. Counting is the reading intended: dividing counts 511 pieces no core owns, and would report `mb` as refined when the movement assembles along it. So pin the definition to counting, and add that file's two ownership tables to the preamble so the rule can be read off real data rather than asserted. §9.2 works a pattern forward from slice counts already given; this pair shows where the counts come from. Second, under that rule both region-count clauses of the guard row pass — `prod(Ns(a)) = 8 × 4 = 32` and `prod(Nd(a)) = 1 × 32 = 32`, matching the region counts on either side. What disqualifies the file is the coverage §9.3 already reports: `32 × 64 = 2048` elements against `512 × 32 × 64 = 1048576`, one 512th of it. No count clause can see that, since a selection is a perfectly consistent set of regions that happens not to add up to the tensor. So write the condition the verdict already implies as a third clause on the same row, and explain its two qualifications: it sums a side's *distinct* regions, or a per-core sum overshoots by design, and it measures against the value being delivered, or the select-then-deliver that repairs a selection would trip the guard it was meant to satisfy. Signed-off-by: Takuya Nakaike --- docs/inter-tile-communication.md | 56 ++++++++++++++++++++++++++++---- 1 file changed, 50 insertions(+), 6 deletions(-) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index 119b804..f268560 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -2047,11 +2047,14 @@ core's **region** is the intersection of its per-axis ranges, and an axis absent from an entry is uncut. Classifying a pattern means deciding, from those two tables, which delivery op expresses the same movement. -**Coarsened and refined.** Let `Ns(a)` and `Nd(a)` be the number of slices -axis `a` is cut into at the source and destination (1 if absent). Then `a` -is **coarsened** when `Ns(a) > Nd(a)` — fewer, larger pieces after the -move, so data must be *assembled* along it — and **refined** when -`Nd(a) > Ns(a)`, so data must be *split* along it. Write +**Coarsened and refined.** Let `Ns(a)` and `Nd(a)` be the number of +**distinct slices** axis `a` carries in the source and destination tables — +counted from the slices a table actually contains, never from slices it does +not: dividing the axis extent by one slice's extent counts pieces no core +owns. An axis absent from an entry counts 1. Then `a` is **coarsened** when +`Ns(a) > Nd(a)` — fewer, larger pieces after the move, so data must be +*assembled* along it — and **refined** when `Nd(a) > Ns(a)`, so data must be +*split* along it. Write ``` C = {a : Ns(a) > Nd(a)} # coarsened — assemble @@ -2071,6 +2074,29 @@ core count. And axes must be aligned by **physical axis, not by label** — labels differ between the two sides, and this is the only place a wrong answer can enter. +**Where the two readings differ.** §9.2 works a pattern forward from slice +counts already given; this pair shows where those counts come from, on the +one measured file where counting and dividing disagree. The tensor is +`512 × 32 × 64` elements, `out` counting sticks. The source divides it 8 +ways on `mb` and 4 ways on `out`; the destination keeps **one** `mb` index — +the last, `mb[511]` — and spreads that single row over all 32 cores, one +stick each. + +**Ownership tables.** One row per side; core ids run row-major with `mb` +outermost. + +| | slice counts | core 0 | core 28 | per-core type | +|---|---|---|---|---| +| src | `{mb:8, out:4}` | `mb[0:64] × out[0:8]` | `mb[448:512] × out[0:8]` | `tensor<64x8x64xf16>` | +| dst | `{mb:1, out:32}` | `mb[511:512] × out[0:1]` | `mb[511:512] × out[28:29]` | `tensor<1x1x64xf16>` | + +Read off the destination table, `mb` carries one distinct slice — every core +names `mb[511:512]` — so `Nd(mb) = 1` against `Ns(mb) = 8`, and `mb` is +coarsened; `out` goes the other way, `4` against `32`, so it is refined. +Divide instead, extent `512` by the slice's extent `1`, and `Nd(mb) = 512`: +`mb` would come out *refined*, on the strength of 511 pieces the table never +mentions. Only the first reading is a fact about the tables. + **Axis names versus axis indices.** Axis sets are written below with the backend's symbol names (`mb`, `in`, `out`, …), the vocabulary the ownership tables speak. The op attributes of §6 are `i64` arrays of *axis indices* @@ -2082,7 +2108,7 @@ fastest-varying. | # | condition | result | |---|---|---| -| — | `prod(Ns(a)) != len(src_regions)` or `prod(Nd(a)) != len(dst_regions)` | **not a work-division pair** — enumerate regions instead. Check first | +| — | `prod(Ns(a)) != len(src_regions)` or `prod(Nd(a)) != len(dst_regions)`, or either side's distinct regions do not cover the tensor | **not a work-division pair** — enumerate regions instead. Check first | | — | any axis ragged (non-uniform overlap) | *insufficient information* — no single op has uniform dependency-set cardinality (R6) | | 1 | `C = ∅` and `R = ∅` | regions identical: `no op needed` if every core's region is its own, else `inter_tile_consume` — a **relocation**, or a **broadcast** where destination regions are shared | | 2 | `C = ∅`, `R ≠ ∅` | `inter_tile_scatter`, `scatter_dimensions = R` | @@ -2095,6 +2121,24 @@ fixes — which is why they must be list-valued (§1.1). Rows 2–4 return R8 admits one producer per group and the tables do not say which transmits (§10.2). +**The coverage clause of the first guard row.** Its other two clauses are +region counts; this one is a volume — each **distinct** region's element +count, summed over a side, against the element count of the value being +delivered. Both qualifications matter. **Distinct**, because summed per +*core* an all-gather exceeds the tensor by design: the file whose one +destination region is held by 28 cores (§9.3) would overshoot 28 times over. +And **the value delivered**, because after a select that is the selected +sub-tensor and not the original — otherwise the select-then-deliver that +repairs a selection would trip the guard it was meant to satisfy. + +The clause is also the only test in §9 that reads slice **sizes** and the +tensor shape rather than slice counts, which is why the count clauses cannot +replace it. On the pair tabled above, `prod(Nd(a)) = 1 × 32 = 32` equals the +destination region count and `prod(Ns(a)) = 8 × 4 = 32` equals the source's, +so **both count clauses pass**. Coverage is what fails: the 32 distinct +destination regions hold one stick each, `32 × 64 = 2048` elements against +the tensor's `512 × 32 × 64 = 1048576` — a 512th of it. + Row 3's `prod(Nd(a)) == num_cores` is the one irreducibly global test: holding the division fixed and varying the core count changes the op, so no per-axis quantity can see it. It is also weaker than asking whether any From 4a0890c6e574bc228035dd8aef5df3c7633fbf67 Mon Sep 17 00:00:00 2001 From: Yu Chin Fabian Lim Date: Fri, 28 Aug 2026 18:15:49 -0400 Subject: [PATCH 07/10] docs(inter-tile): restate R8 as single-source delivery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit R8 read "exactly one producer tile per group", which contradicted the document's own `consume` examples: the per-tile pairing example declares two producers per group and the butterfly exchange declares four. §6.1 also claimed the dependency attribute never changes what a consumer receives, while §3.4 says that for `replicate` placement it selects which producer a consumer reads. Both cannot hold. What `replicate` and `split` actually require is that each *consumer tile* have a single source -- their results have room for one contribution and no combiner to fold a second -- so state R8 that way: forall g, forall c in consumer_tiles_per_group(g) : |dep(c, g)| == 1 One producer per group is then the common case that satisfies it without an attribute (broadcast). A multi-producer group stays legal for `consume` when the attribute pairs each consumer tile with exactly one producer: that is a routing pattern, several point-to-point deliveries sharing one produce op, not an all-producers delivery. Absent the attribute, a multi-producer group is rejected -- receiving from every producer is §1.1's undefined cell. A producer may serve several consumer tiles; it may not serve none, which R4 already requires. `scatter` takes no dependency attribute, so for it R8 reduces to its original form. Also records that the one-symbol dependency spelling is now implemented, and drops a duplicated sentence in R9. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Yu Chin Fabian Lim --- docs/inter-tile-communication.md | 94 ++++++++++++++++++++++++-------- 1 file changed, 72 insertions(+), 22 deletions(-) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index f268560..59bee64 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -64,12 +64,12 @@ means a tensor or tile axis. | Op | combine | placement | producers/grp | consumers/grp | dim attrs | region | identity | |---|---|---|---|---|---|---|---| -| `consume` | none | replicate | 1 | free | — | — | — | +| `consume` | none | replicate | 1 per consumer | free | — | — | — | | `reduce` | fold | replicate | all | free | — | combiner | yes | | `reduce_scatter` | fold | split | all | free | `scatter_dimensions` | combiner | yes | | `gather` | none | concat | all | free | `gather_dimensions` | — | — | | `all_to_all` | none | permute | all | all | `split_dimensions`, `concat_dimensions` | — | — | -| `scatter` | none | split | 1 | free | `scatter_dimensions` | — | — | +| `scatter` | none | split | 1 per group | free | `scatter_dimensions` | — | — | `all_to_all` is listed before `scatter` because it shares the all-producers cardinality cell with `gather` and `reduce_scatter`, and @@ -119,7 +119,11 @@ all-to-all from scatter. | Scatter | 1 | free | `inter_tile_scatter` | 1/C slice of full | `inter_tile_scatter` and `inter_tile_consume` have no natural "all-" -variant: R8 (§5) limits them to one producer per group. +variant: R8 (§5) gives each consumer tile exactly one source, so there is no +all-producers case to widen to. `consume` still admits a group holding several +producers, but only as a **routing** pattern in which the dependency attribute +pairs each consumer tile with one of them (§6.1) — several point-to-point +deliveries sharing one `produce`, not an all-producers delivery. ### 1.3 The future value @@ -546,7 +550,7 @@ carries the attribute the rule constrains. | R5 dep sets pairwise disjoint | delivery | — | — | — | y | y | n/a | | R6 uniform dep-set cardinality | delivery | — | — | — | y | y | n/a | | R7 uniform producer cardinality across groups | delivery | — | — | — | y | y | n/a | -| R8 producers per group = 1 | delivery | y | — | — | — | — | y | +| R8 single-source delivery | delivery | y | — | — | — | — | y | | R9 flattened split extent divisible by `C` | delivery | — | — | y | — | y | y | | R10 combiner purity (§3.5) | delivery | — | y | y | — | — | — | | R11 identity shape matches `T_p` (§3.5) | delivery | — | y | y | — | — | — | @@ -593,8 +597,40 @@ Statements: nothing otherwise requires equal cardinality per group. Since the op result is a single static tensor type, unequal groups yield no expressible result type for the assembling placements. -- **R8 — single producer.** Exactly one producer tile per group, for the - ops whose `producers/grp` cell is `1`. +- **R8 — single-source delivery.** For the ops whose `producers/grp` cell is + `1` — `inter_tile_consume` and `inter_tile_scatter` — every consumer tile + must receive from exactly one producer tile: + + ```text + ∀ g, ∀ c ∈ consumer_tiles_per_group(g) : |dep(c, g)| == 1 + ``` + + where `dep(c, g)` is the producer set `producer_dependency_per_consumer` + declares for consumer tile `c` (§3.4). `inter_tile_scatter` takes no such + attribute (§6.6), so for it the rule reduces to its simplest form: exactly + one producer tile per group. + + **Why per consumer tile rather than per group.** These two ops deliver into a + result no larger than one contribution — unchanged for `replicate`, a `1/C` + slice for `split` — with no combiner and nowhere to put a second value. A + consumer tile holding two contributions is therefore the undefined cell of + §1.1, "which producer's value wins?". What the op needs is not that the + *group* hold one producer, but that each *consumer tile* have a single + source. The two coincide when `|P(g)| == 1`, the common case (broadcast, + §7.1), which needs no attribute at all. + + **For `inter_tile_consume` with `|P(g)| > 1` the attribute is required.** + There is no meaningful default, because receiving from every producer is + exactly that undefined cell. With the attribute, such a group is a + **routing** pattern — several independent point-to-point deliveries sharing + one `produce` op, as in §7.4.1 and §7.4.2 — and no consumer tile ever sees + two values. A group with `|P(g)| > 1` and no attribute is rejected. + + A producer **may** serve several consumer tiles (multicast within the + group); it may not serve none, which R4 already requires. So `dep` need not + be injective — only single-valued per consumer tile, and total over + producers. + - **R9 — split divisibility.** `E(D_split) % C == 0`, where `D_split` is the op's split axis set (`scatter_dimensions`, or `split_dimensions` for `all_to_all`) and `E` is the flattened extent of §4. Stating the rule on the @@ -606,7 +642,6 @@ Statements: list must be non-empty; and the entries must be in **ascending numerical order** (§4). Repeated indices would double-count an extent in `E`, and an out-of-order list would silently denote a different flattening. - Repeated indices would double-count an extent in `E`. - **R11 and the shipped constraint.** R11 pins `identity` to `T_p`, while the implemented `reduce` ties it to *results* (`KTDP.td:172-174`). With no rank reduction (§4) these coincide for `reduce`, since its result *is* @@ -646,8 +681,8 @@ argument. Shared machinery is §3; rules are §5. ### 6.1 `ktdp.inter_tile_consume` — broadcast -`combine = none`, `placement = replicate`, one producer per group, -consumer set free, no dim attribute, no region, no identity. +`combine = none`, `placement = replicate`, one producer per consumer tile +(R8), consumer set free, no dim attribute, no region, no identity. **Result type.** `T_p_i` unchanged (§4, `replicate` + `none`). @@ -662,11 +697,17 @@ group — broadcast. : !ktdp.tile_future -> T_p_1, ..., T_p_N ``` -With one producer per group, `producer_dependency_per_consumer` is a pure -synchronization refinement (§3.4): it changes when each consumer -unblocks, never what it receives. That is what makes `consume` also the -op for one-to-one permutation exchange — a bijective dependency set over -a multi-producer group (§7.4.2). +With one producer per group the attribute is a pure synchronization +refinement (§3.4): there is only one value to receive, so it changes when +each consumer unblocks and nothing else. With **several** producers per +group it also names the sender, and R8 then requires it — each consumer +tile must be paired with exactly one producer. Such a group is a +**routing** pattern rather than a broadcast: several independent +point-to-point deliveries sharing one `produce` op, which is what lets +`consume` express per-tile pairing (§7.4.1) and one-to-one permutation +exchange (§7.4.2). Delivering to a consumer tile from more than one +producer is never legal here — with no combiner and a result the size of +one contribution, there would be nowhere to put the second value (§1.1). ### 6.2 `ktdp.inter_tile_reduce` — reduction @@ -2016,12 +2057,20 @@ the Torch-Spyre SDSC planner (`_compatible_partitions`) but are absent from the KTIR verifier entirely — the gap exists at both the spec and the implementation level. +**Dependency-set arity.** §3.4's one-symbol spelling `(p)[c]` is accepted as +of `KTDPInterTileHelpers.cpp:69–100` and `KTIRCheckLegality.cpp:135–142`. +Before that, `depTilesOf` always bound two symbols and the pass rejected any +set whose symbol count was not exactly 2, so the group-independent form this +document documents — and uses in §7.4.1 — was unusable in practice. The symbol +count now selects how many values are bound, and 3-or-more is diagnosed. + **Two asymmetries the verification matrix (§5) forces into the open.** -1. R8 is stated as a verifier obligation for `scatter` but is merely - conventional for `consume`, even though both ops have the same - single-producer cardinality. Neither is implemented yet, so this is a - spec asymmetry to decide rather than inherit. +1. R8 is a verifier obligation for both `consume` and `scatter`, but it + bites differently: `scatter` takes no dependency attribute, so one + producer per group is the whole rule, whereas `consume` admits a + multi-producer group whenever the attribute pairs each consumer tile with + exactly one producer (§5). Neither is implemented yet. 2. R13 and R14 are implemented for `reduce` only, and R13 is the implementation of open question §10.1 (must a consumer also be a producer?) for that one op. The `?` cells in that matrix are exactly @@ -2118,8 +2167,8 @@ fastest-varying. The dimension attributes are **the axis sets themselves**, in the order §4 fixes — which is why they must be list-valued (§1.1). Rows 2–4 return *insufficient information* when a source region has several holders, since -R8 admits one producer per group and the tables do not say which transmits -(§10.2). +R8 gives each consumer tile exactly one source and the tables do not say which +holder transmits (§10.2). **The coverage clause of the first guard row.** Its other two clauses are region counts; this one is a volume — each **distinct** region's element @@ -2278,8 +2327,9 @@ Both are §9.1 escape hatches, and both need the per-region core sets rather than the division. **Producer election.** Rows 2–4 return *insufficient information* when a -source region has several holders: R8 admits one producer per group, and the -tables record who *holds* a region, not who *transmits* it. Unforced by +source region has several holders: R8 requires each consumer tile to have +exactly one source, and the tables record who *holds* a region, not who +*transmits* it. Unforced by measurement — every source region across the 51 files has a single holder — so the choice between electing a canonical producer (lowest tile id), requiring the frontend to pick, and rejecting replicated sources can wait. From 1f9797de04055d32f5ae1c754ad255846e8cbf77 Mon Sep 17 00:00:00 2001 From: Yu Chin Fabian Lim Date: Mon, 31 Aug 2026 17:36:08 -0400 Subject: [PATCH 08/10] docs(inter-tile): say where a bounded extent can be expressed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit §9 says three times that a selection is not a delivery and that coverage is measured against the value delivered, but never says where a bounded extent lives. The ownership tables raise it directly: the destination owns mb[511:512], and nothing in the delivery ops can say that. The tile sets name which tiles participate, the dimension attributes name which axes split or concatenate, and §4's type rules are extent arithmetic with no base coordinate -- none of the six ops carries an offset. So a bounded extent is a property of T_p, which gets it from the access tile the partial was loaded through. The snippet shows that chain -- access tile, load, yield_partial, delivery -- in §6's placeholder style rather than with one measured file's coordinates, since what generalizes is that the extent and offset are arguments to construct_access_tile and appear nowhere downstream: selecting a different sub-tensor changes only T_p. Whether a delivery is needed at all is then the ordinary classification question, which for this pair §9.1 row 1 answers with `no op needed`. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Yu Chin Fabian Lim --- docs/inter-tile-communication.md | 42 ++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index 59bee64..fe0e391 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -2146,6 +2146,48 @@ Divide instead, extent `512` by the slice's extent `1`, and `Nd(mb) = 512`: `mb` would come out *refined*, on the strength of 511 pieces the table never mentions. Only the first reading is a fact about the tables. +**How a bounded extent is expressed.** The destination above owns +`mb[511:512]`, and no delivery op can say that. The tile sets name which +*tiles* participate, the dimension attributes name which *axes* split or +concatenate, and §4's type rules are extent arithmetic with no base coordinate: +none of the six ops carries an offset. + +A bounded extent is therefore a property of `T_p`, and `T_p` gets it from the +access tile the partial was loaded through. The chain is the same whatever the +bound: + +```mlir +// The extent is the access tile's shape and the offset is its anchor. Both are +// arguments here, and neither appears again downstream. +%access = ktdp.construct_access_tile %view[] { + access_tile_set = , access_tile_order = +} : memref<...> -> !ktdp.access_tile<...xindex> +%partial = ktdp.load %access : !ktdp.access_tile<...xindex> -> T_p + +// From here the bound is invisible: the delivery op sees a T_p and an axis +// set, never a coordinate. Selecting a different sub-tensor changes only T_p. +%future = ktdp.inter_tile_produce producer_tiles_per_group = + : T_p -> !ktdp.tile_future +{ ^bb0(%gid: index): ktdp.yield_partial %partial : T_p } +%result = ktdp.inter_tile_consume(%future) + consumer_tiles_per_group = + : !ktdp.tile_future -> T_p +``` + +`ktdp.construct_access_tile` fixes the coordinates, `ktdp.load` yields the +value, `ktdp.yield_partial` only names it, and the delivery carries whatever the +partial turned out to be. That is why a selection is not a delivery, and why the +coverage clause measures the value *delivered*: read as a delivery of the whole +tensor the pair above covers 1/512 and trips the guard, while the sub-tensor it +should have been is covered exactly. + +Whether a delivery is then needed at all is the ordinary classification +question. Here it is not — each core owns a different stick of the selected row, +so nothing crosses cores and row 1 of §9.1 gives `no op needed`. Had they all +needed the same stick it would be a broadcast, and the load would move inside +the producer region, since one producer per group (R8) means the non-producing +tiles must not run it (§2.2, and §7.7.1 for the same reason on `scatter`). + **Axis names versus axis indices.** Axis sets are written below with the backend's symbol names (`mb`, `in`, `out`, …), the vocabulary the ownership tables speak. The op attributes of §6 are `i64` arrays of *axis indices* From 1b1520f527b5c92af86598a6c4c59f3c062589d9 Mon Sep 17 00:00:00 2001 From: Takuya Nakaike Date: Mon, 7 Sep 2026 05:21:38 +0000 Subject: [PATCH 09/10] docs(inter-tile): state the producer-consumer relation per delivery op MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Coverage was being read off operation names, so the two tile sets a delivery relates were only recoverable by tracing R4-R8 and R13 across sections. State the relation directly and settle what was left open. - §1.2: a relation table, one row per delivery op, answering whether a consumer must also be a producer, whether one producer may serve several consumers, and whether the two sets coincide. - §3.3: pin, normatively, that when producer_dependency_per_consumer is present an assembling consumer's positions come from its own declared set, not from the group's producer set. §3.4 already implies this by making P follow from the subset: a P-chunk result needs P consecutive positions, which only the declared subset supplies. - §5 R13: y for reduce, n for the other five. §6.6's argument turns on combine = none, which all four copy-only ops share, not on scatter's single source; and §9.3's non-square all_to_all measures K = 2 producers against M = 4 consumers per group, so a receiving-only tile is a measured shape there. - §5 R5: scope the disjointness obligation to a partitioning use. Well-definedness comes from §3.3, so identical or overlapping dependency sets - an all-gather, a multicast source - must not be rejected as such. - §6.5: "all tiles produce, all tiles consume" means the whole declared set on each side, not that the two sets are the same tiles. - §8, §10.1: record the above; what remains open is reduce alone, where dropping R13 also means revisiting R14's mode gate. Signed-off-by: Takuya Nakaike --- docs/inter-tile-communication.md | 151 +++++++++++++++++++++++-------- 1 file changed, 113 insertions(+), 38 deletions(-) diff --git a/docs/inter-tile-communication.md b/docs/inter-tile-communication.md index fe0e391..abd4566 100644 --- a/docs/inter-tile-communication.md +++ b/docs/inter-tile-communication.md @@ -125,6 +125,36 @@ producers, but only as a **routing** pattern in which the dependency attribute pairs each consumer tile with one of them (§6.1) — several point-to-point deliveries sharing one `produce`, not an all-producers delivery. +**Producer–consumer relation.** One row per delivery op. The two cardinality +columns of §1.1 say how *many* tiles take part; this table says how the two +sets *relate*, which is what a lowering and a test fixture need. "Set" means +`producer_tiles_per_group(g)` and `consumer_tiles_per_group(g)`. + +| Op | Must a consumer also be a producer? | May one producer serve several consumers? | Are the two sets equal? | +|---|---|---|---| +| `consume` | no | **yes** — multicast within the group (R8) | no — the common case is one producer, many consumers (§7.1) | +| `reduce` | **yes**, current restriction only (R13, §10.1) | yes — every consumer takes the whole fold | only in the all-reduce case; `reduce_to_one` is the single-consumer case (R14) | +| `reduce_scatter` | no | yes | not required | +| `gather` | no | **yes** — the same producer piece appears in several consumers' assemblies | no — plain gather has one consumer, all-gather the whole group (§1.2) | +| `all_to_all` | no | yes | **not required** — §9.3's measured non-square case has `K = 2` producers and `M = 4` consumers per group (§9.1) | +| `scatter` | no — argued in §6.6 | yes | no | + +Two consequences worth stating once rather than per op. + +**A source may be shared by several consumers, for every delivery op.** R8 +puts it directly for the single-source ops ("a producer **may** serve several +consumer tiles (multicast within the group)"), and it holds for the assembling +ops too: two consumers whose declared dependency sets name the same producer +each place that producer's contribution in their own assembly, by their own +`P` and their own `l` order (§3.3). R5 is not a prohibition on sharing; see +its statement in §5. + +**Only `reduce` requires its consumers to be producers.** For the four +copy-only ops the question is settled by §6.6's argument — a consumer that +receives contributes nothing, so there is nothing for it to miss — and for +`all_to_all` the measured non-square case settles it by example. `reduce`'s +`yes` is the restriction §10.1 describes, not a property of folding. + ### 1.3 The future value `ktdp.inter_tile_produce` returns a `!ktdp.tile_future` SSA @@ -290,6 +320,22 @@ position — never a tile id, and never an offset in the textual order of an enumerated set. ("Position" rather than "rank": in this document *rank* always means a tensor's number of dimensions.) +**Which set, when `producer_dependency_per_consumer` is present — normative.** +For an assembling consumer the relevant producer set is **that consumer's own +declared dependency set** `dep(c, g)`, not the whole +`producer_tiles_per_group(g)`. So consumer `c` assembles `|dep(c, g)|` pieces +at positions `0 .. |dep(c, g)| - 1`, ordered by ascending tile id within +`dep(c, g)`. This is the reading §3.4 already implies by making `P` — and +hence the result type — follow from the declared subset: a result of `P` +chunks needs `P` consecutive positions, which only the declared subset +supplies. Taking positions from the group's producer set instead would leave a +non-prefix subset with no defined slots, independently of anything R5 says. + +A consequence worth stating: because each consumer's positions come from its +own set, two consumers may declare **overlapping or identical** sets and each +assembly stays well-defined — the shared-source case of §1.2, and what R5 is +not about. + ### 3.4 `producer_dependency_per_consumer` *(optional)* Affine integer set `(p)[c, g]` over producer tile IDs `p`, parameterized @@ -555,7 +601,7 @@ carries the attribute the rule constrains. | R10 combiner purity (§3.5) | delivery | — | y | y | — | — | — | | R11 identity shape matches `T_p` (§3.5) | delivery | — | y | y | — | — | — | | R12 flattened concat extent × `P` well-defined | delivery | — | — | — | y | y | — | -| R13 consumer set subset of producer set | delivery | — | y | ? | ? | ? | n | +| R13 consumer set subset of producer set | delivery | — | y | n | n | n | n | | R14 reduce mode gate: `C == P` or `\|C\| == 1` | delivery | — | y | ? | — | — | — | Statements: @@ -583,12 +629,24 @@ Statements: producer_dependency_per_consumer(p)[c, g] ``` -- **R5 — pairwise disjointness.** For the assembling placements - (`concat`, `permute`), distinct consumers' declared dependency sets - must be disjoint. R4 alone requires only that each producer be claimed - by *at least one* consumer, which combined with R6 admits declared sets - that double-count producers — and a double-counted producer has no - well-defined position in the assembly. +- **R5 — pairwise disjointness, for a partitioning use only.** For the + assembling placements (`concat`, `permute`), when the declared dependency + sets are meant to *split* the producers among the consumers — a segmented + assembly, §3.4 — they must be pairwise disjoint. R4 alone requires only that + each producer be claimed by *at least one* consumer, which combined with R6 + admits declared sets that double-count producers, and a producer claimed by + two segments of one partition has no single position in the assembly. + + **The rule does not forbid a shared source, and is not what makes an + assembly well-defined.** Distinct consumers may declare identical or + overlapping sets: that is what an all-gather is (§6.4, every consumer + names every producer) and what a multicast producer is (R8). + Well-definedness comes from §3.3 instead — each consumer's positions are + taken from its *own* declared set, so its assembly is fixed whatever other + consumers declare. What R5 adds is only the partitioning intent: when the + sets are meant to split the producers, a producer in two of them belongs + to two segments at once. A verifier must therefore not reject overlap as + such. - **R6 — uniform dep-set cardinality.** All consumers in a group must declare the same number of producers, so `P` is a single number and the op has one static result type. @@ -661,10 +719,17 @@ Statements: `all_to_all` case the divisibility follows from R7 + R9, but it must be stated independently for the non-square case, which §9.3 shows is measured and not hypothetical. -- **R13 — consumer set subset of producer set.** Every consumer tile in a - group must also be a producer in that group, i.e. - `consumer_tiles_per_group(g) ⊆ producer_tiles_per_group(g)`. Whether - this should hold is §10.1; it is currently enforced for `reduce` only. +- **R13 — consumer set subset of producer set.** Where the rule applies, + every consumer tile in a group must also be a producer in that group, i.e. + `consumer_tiles_per_group(g) ⊆ producer_tiles_per_group(g)`. **It applies + to `reduce` only.** For the four copy-only ops a receiving-only tile is + legal: §6.6's argument — a consumer that receives contributes nothing, so + there is nothing for it to miss — turns on `combine = none`, which all + four share, and not on `scatter`'s single source. `all_to_all` is settled + by measurement rather than by argument: §9.3's non-square case has + `K = 2` producers against `M = 4` consumers per group, so two of each + group's consumers are not producers. `reduce`'s `y` is the implementation + restriction described in §10.1, kept because R14's mode gate assumes it. - **R14 — reduce mode gate.** For `reduce`, the consumer set must either equal the producer set (all-reduce) or be a single tile (reduce-to-one); a strict multi-tile subset — reduce-to-subset — is @@ -795,9 +860,16 @@ R5–R7. ### 6.5 `ktdp.inter_tile_all_to_all` — split and reassemble -`combine = none`, `placement = permute`, all tiles produce, all tiles -consume, both `split_dimensions` and `concat_dimensions`, no region, no -identity. +`combine = none`, `placement = permute`, every declared producer contributes +and every declared consumer receives, both `split_dimensions` and +`concat_dimensions`, no region, no identity. + +**The two sets need not coincide.** "All produce, all consume" means each +delivery draws from the whole producer set and delivers to the whole consumer +set — not that the two sets are the same tiles. §9.3's measured non-square +case has `K = 2` producers against `M = 4` consumers per group, so a consumer +that is not a producer is a measured shape here, and R13 is `n` for this op +(§5). **Attributes.** `split_dimensions` (`i64` array) — axes each producer splits into `C` chunks (R9). `concat_dimensions` (`i64` array) — axes @@ -2071,14 +2143,12 @@ count now selects how many values are bound, and 3-or-more is diagnosed. producer per group is the whole rule, whereas `consume` admits a multi-producer group whenever the attribute pairs each consumer tile with exactly one producer (§5). Neither is implemented yet. -2. R13 and R14 are implemented for `reduce` only, and R13 is the - implementation of open question §10.1 (must a consumer also be a - producer?) for that one op. The `?` cells in that matrix are exactly - that question, unresolved: for `scatter` the answer is **no** (§6.6), for - `reduce` the current answer is **yes** (enforced), and for - `reduce_scatter` / `gather` / `all_to_all` it is undecided. R14's - mode gate is likewise a current implementation restriction, not a - design conclusion. +2. R13 and R14 are implemented for `reduce` only. R13 now reads `y` for + `reduce` and `n` for every other op (§10.1), so what remains is an + implementation asymmetry rather than an open question: the enforced check + is the one op where the rule applies, and the four copy-only ops need no + check at all. R14's mode gate stays a current implementation restriction, + not a design conclusion. --- @@ -2346,22 +2416,27 @@ restickified weights are explicitly barred as shuffle sources. ## 10. Open questions -### 10.1 Must a consumer also be a producer? - -**What turns on it:** whether the verifier rejects a delivery op whose consumer -set is not contained in its producer set. That check exists and runs today — R13 -for `reduce` (`KTIRCheckLegality.cpp:107–117`) — so whoever implements -`gather`, `all_to_all` or `reduce_scatter` must decide whether to extend it, -and the answer changes which programs are legal. - -**Why it is unresolved:** `reduce`'s *yes* is one op's implementation choice, -made when it was the only delivery op. `scatter`'s *no* is settled and -argued (§6.6) — a consumer that receives a slice contributes nothing, so -there is nothing to miss. Neither generalizes: `gather` and `all_to_all` -assemble, so a non-producing consumer is coherent for them in a way it is not -for a reduction. The `?` cells in §5 are exactly the ops still to decide, and -R14's mode gate (all-reduce or reduce-to-one, no strict multi-tile subset) is -a present restriction on `reduce` awaiting the same call. +### 10.1 Must a consumer also be a producer? — resolved for the copy-only ops + +**Resolved: no, except for `reduce`.** R13 is now `n` for `consume`, +`reduce_scatter`, `gather`, `all_to_all` and `scatter`, and `y` for `reduce` +alone (§5, and the relation table in §1.2). + +**What settled it.** §6.6's argument does not depend on `scatter`'s single +source — it turns on `combine = none`: a consumer that receives contributes +nothing, so there is nothing for it to miss and no coverage obligation arises. +All four copy-only ops share `combine = none`, so the argument carries to them +unchanged. For `all_to_all` there is also a measured instance: §9.3's +non-square case has `K = 2` producers against `M = 4` consumers per group, so +two consumers per group are not producers. + +**What is still open, and it is only about `reduce`.** `reduce`'s `y` is the +implementation choice made when it was the only delivery op, and R14's mode +gate (all-reduce or reduce-to-one, no strict multi-tile subset) assumes it. +Dropping R13 for `reduce` means deciding what a fold delivered to a +non-contributing tile means — in effect a reduce composed with a broadcast — +and revisiting R14 in the same change. Nothing in this document requires that +today. ### 10.2 Two things a work division cannot settle From 56e119452b65ad732ca46fb85a563169f00f1c9f Mon Sep 17 00:00:00 2001 From: Takuya Nakaike Date: Mon, 7 Sep 2026 09:06:46 +0000 Subject: [PATCH 10/10] docs(inter-tile): minimal examples for the cases #4300 asks first MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit torch-spyre#4300 §7 names KT-01/02 (shared sources, receiving-only cores), KT-05 (selection) and KT-07 (completed sums) as the cases that decide semantics rather than implementation polish. This adds the smallest KTIR form for each, plus KT-03 (logical order against tile-id order), KT-04 (multi-axis assembly), KT-06 (permuted core ownership) and KT-09 (many readers of one delivery, and scratch reuse) — 8 of the 10 KT cases. Each file carries its fixture, its operation sequence, its expected values, the rules it exercises, and its verification status stated rather than assumed. Fixtures come from the pinned ownership catalog, not from the inventory's summary rows: those rows cannot separate divisions that give the same piece counts, and their shape column is alphabetically sorted rather than in layoutDimOrder_. Two further traps in reading it are recorded in kt06. source_pieces is ordered lexicographically by its key string, so piece index is not owner order, and source_core_patterns is not a per-piece record - the owners field on each piece is the authoritative one, and deriving the map from index position yields a bijection that is not the transpose. And the byte fields are logical: prod(extents) * word_length equals logical_tensor_bytes even for the record whose SDSC run shows an extent-1 y carried as 64 padded lanes, so padding is invisible in the catalog for every record and no record can settle a layout question alone. Only produce and reduce exist in KTDP.td, so kt07a - an all-reduce - is the only delivery that parses and passes --ktir-check-legality end to end. The others parse up to their delivery op and stop there; each says which op it waits on. kt06's variant A also passes, for the opposite reason: it is the branch §9.1 row 1 resolves to "no op needed", so there is no delivery op in it to be missing. Every file states its operation sequence, collected as a table in the README. Side by side that view carries what no single file shows: kt05 is the only load inside a produce region (§2.2, since only one tile per group produces); kt06 variant A is one op long and the absent delivery is its claim; kt06 variant B loads and stores at the same address, so its permutation is entirely in the dependency attribute; kt09 half 1 is the only case whose wrong version has an identical shape and identical values, differing only in emitted transfer count; and 8 of the 10 sequences rely on a landing store the IR never names. What the examples establish, read off the catalog or the verifier rather than argued: - Receiving-only cores are the norm. 64 of 130 records have destination owners that are not source owners, so R13 = n for the copy-only ops is required by measurement. - The catalog holds no reduction at all: 130 of 130 are copy-only. So KT-07's halves come from different places, and the risk it guards against is one-sided - reading a copy as a fold. - A fold op cannot stand in for a copy op. Writing KT-07 (b) as a degenerate reduce over one producer is rejected by the shipped legality pass on R13, which is evidence rather than argument. - Overlapping-but-differing dependency sets occur in 0 of 130 records, so R5's scoping is forced by measurement while §3.3's "which set" clause is not. kt01 carries a synthetic second function for the latter, which also discriminates §3.1's P derivation; the measured form is blind to both. - Where a required order disagrees with the order the IR supplies, whether a verifier can see it depends only on where the ordering is allowed to live. kt06's permutation lives in a dependency set, where R3 through R8 reach it - and the grouping choice is what gives them anything to compare, since R1, R4 and R5 are obligations within a group: 8 groups of 4 make R1 and R3 bite, one group of 32 leaves both nearly trivial, and 32 singleton groups make R4 and R5 vacuous. kt03's reorder has to live in store bases, where no rule reaches it at all. - KT-04 needs no new op - it is gather with two entries in gather_dimensions - and its verdict splits. §4's normative flattening plus R9's ascending rule and R12's per-axis check close both spelling errors, a reversed list and a single axis where two are needed. They do not close the case where the producers' tile order and the required cell order are a transpose, which is the requirement's own toy and is kt03's finding in multi-axis form. - One delivery may have many readers, and R2 does not stand in the way: it constrains the tile_future, not the delivery's result. kt09's fixture makes the cost measured rather than argued - mean-LayerNormNorm_out is the source of relayouts[1, 8, 16], the Q/K/V projections off one layer norm, and the three records are identical at 12 MiB each, so a delivery per reader moves 36 MiB where 12 suffices. The catalog invites it by being indexed on (consumer, input); 8 of its 120 distinct tensors are shared this way. Measurement never forces the ordering problem. Sources contribute in ascending tile-id order in 130 of 130 records, GR-PF-052 and its four siblings agree with the transpose, and GR-PF-121 agrees with §4's odometer for all 16 producers. So kt03 and kt04's toy are synthetic, and the measured functions beside them need nothing beyond a plain delivery. KT-03 is satisfied as written rather than blocked: coverage §5 permits "local selection/reordering ... if fully expressed and stays on-chip", and two stores at swapped bases are exactly that. Nor do they cost anything - a copy delivery lands every received tile in LX before a compute unit can read it, so P landing stores are mandatory whatever the order is, and the two stores here are those landings aimed at swapped bases. What the case does cost is verifiability, since no rule relates the store bases to §3.3's assembly order. An ordering attribute on the delivery op was considered and rejected on that basis: an attribute can be checked for well-formedness but never for intent, so it would add a second statement of the ordering without moving the case into the verifier - and, given the mandatory landings, without saving a write either. The same landing rule is why kt09's two halves are a real choice rather than a good form and a bad one. Keep the delivered value a tensor and no store appears in the IR, so its readers hold no buffer and its lifetime is the backend's; store it into LX and the live range becomes ordinary MLIR memory effects, at no extra cost because the landing happens either way. Reduction is exempt from the rule - a compute unit can send over a different ring - which is why kt07a can leave its result unstored. kt05 also shows the §9.1 guard doing what it is for: it rejects the pre-select reading of GR-PF-055 on coverage, and after the select the post-select pair classifies as row 2 and yields the scatter in the file. The op is derived, not chosen. Group structure is stated the same way in every file - a table keyed on g, producers before consumers, the formula in the column head - because the affine sets all bind g, so the h and q two of the files used disagreed with the IR they described. One aside, independent of any case: RFC 0682 defines access_tile_order both as a dimension nesting order and as a sort key over an intermediate variable space, and only the second could reorder within a dimension. kt03 deliberately does not depend on the generous reading, and KTDP.td's op description carries neither sentence. Not written: KT-08 (stick indivisibility) and KT-10 (deliberate bad variants). KT-08 has 19 measured records - y 32->64 in relayouts[39..50], a sub-stick y 4->2 in [57, 64, 73] - but needs an SDSC run first, because the catalog's logical byte fields cannot say whether y is sticked for those shapes or at what size. This line of work is suspended here: the plan has moved to expressing these relayouts with ktdp.construct_distributed_memory_view instead of the produce/delivery pair. The examples are kept because the fixtures, the measured counts and the catalog-reading notes are independent of which form is chosen, and because they record what each case demands of any answer. Signed-off-by: Takuya Nakaike --- docs/inter-tile-examples/README.md | 314 ++++++++++++ .../kt01-shared-sources.mlir | 305 ++++++++++++ .../kt02-receiving-only-cores.mlir | 163 +++++++ .../kt03-logical-order.mlir | 219 +++++++++ .../kt04-multi-axis-assembly.mlir | 310 ++++++++++++ .../kt05-select-then-scatter.mlir | 295 ++++++++++++ .../kt06-owner-permutation.mlir | 410 ++++++++++++++++ .../kt07a-raw-contributions.mlir | 163 +++++++ .../kt07b-completed-sum.mlir | 176 +++++++ .../kt09-reuse-and-lifetime.mlir | 445 ++++++++++++++++++ 10 files changed, 2800 insertions(+) create mode 100644 docs/inter-tile-examples/README.md create mode 100644 docs/inter-tile-examples/kt01-shared-sources.mlir create mode 100644 docs/inter-tile-examples/kt02-receiving-only-cores.mlir create mode 100644 docs/inter-tile-examples/kt03-logical-order.mlir create mode 100644 docs/inter-tile-examples/kt04-multi-axis-assembly.mlir create mode 100644 docs/inter-tile-examples/kt05-select-then-scatter.mlir create mode 100644 docs/inter-tile-examples/kt06-owner-permutation.mlir create mode 100644 docs/inter-tile-examples/kt07a-raw-contributions.mlir create mode 100644 docs/inter-tile-examples/kt07b-completed-sum.mlir create mode 100644 docs/inter-tile-examples/kt09-reuse-and-lifetime.mlir diff --git a/docs/inter-tile-examples/README.md b/docs/inter-tile-examples/README.md new file mode 100644 index 0000000..36688fd --- /dev/null +++ b/docs/inter-tile-examples/README.md @@ -0,0 +1,314 @@ +# Minimal inter-tile examples + +Smallest KTIR forms for the cases +[torch-spyre#4300](https://github.com/torch-spyre/torch-spyre/pull/4300) asks +this PR to settle first. Its `lx_relayout_workload_coverage.md` §7 names KT-01/02 +(shared sources and receiving-only cores), KT-05 (selection) and KT-07 (completed +sums) as the ones that "decide semantics rather than implementation polish". + +Each file carries its fixture, its expected values, which rule it exercises, and +its verification status. Fixtures are read from the **pinned ownership catalog** +rather than from the inventory's summary rows, because those rows cannot +distinguish divisions that produce the same piece counts — `32 -> 8` with 4 +inputs per region matches two different §9.3 rows of +[`inter-tile-communication.md`](../inter-tile-communication.md). + +## Status + +| File | Case | Parses | Legality | Blocked on | +|---|---|:---:|:---:|---| +| [`kt07a-raw-contributions.mlir`](kt07a-raw-contributions.mlir) | KT-07 (a) | **yes** | **yes** | — | +| [`kt07b-completed-sum.mlir`](kt07b-completed-sum.mlir) | KT-07 (b) | no | — | `ktdp.inter_tile_consume` | +| [`kt01-shared-sources.mlir`](kt01-shared-sources.mlir) | KT-01 | no | — | `ktdp.inter_tile_gather` | +| [`kt02-receiving-only-cores.mlir`](kt02-receiving-only-cores.mlir) | KT-02 | no | — | `ktdp.inter_tile_gather` | +| [`kt05-select-then-scatter.mlir`](kt05-select-then-scatter.mlir) | KT-05 | no | — | `ktdp.inter_tile_scatter` | +| [`kt03-logical-order.mlir`](kt03-logical-order.mlir) | KT-03 | no | — | `ktdp.inter_tile_gather` | +| [`kt06-owner-permutation.mlir`](kt06-owner-permutation.mlir) | KT-06 | variant A **yes** | variant A **yes** | variant B: `ktdp.inter_tile_consume` | +| [`kt04-multi-axis-assembly.mlir`](kt04-multi-axis-assembly.mlir) | KT-04 | no | — | `ktdp.inter_tile_gather` | +| [`kt09-reuse-and-lifetime.mlir`](kt09-reuse-and-lifetime.mlir) | KT-09 | no | — | `ktdp.inter_tile_gather` | + +The four cases §7 asks to settle first are all here, plus KT-03, KT-04, KT-06 and +KT-09. `kt01`, `kt04`, `kt05`, `kt06` and `kt09` carry two functions each. In `kt01`, +`kt04` and `kt09` the second is one the measured data cannot supply, marked as such +in the file; in `kt06` both are measured and the pair *is* the requirement. `kt03` is +synthetic throughout. + +Only `inter_tile_produce` and `inter_tile_reduce` exist today, so KT-07 (a) — an +all-reduce — is the only *delivery* that can be checked end to end. Its `produce` +half is shared with the others, and that half does parse; the files that fail do so +at their delivery op and nowhere else. `kt06`'s variant A also passes legality, for +the opposite reason: it is the branch §9.1 row 1 resolves to "no op needed", so +there is no delivery op in it to be missing. + +## Operation sequences + +Each file states its own sequence in a header block. Collected here, because the +structural differences are what separate the cases: + +| File | Sequence | +|---|---| +| `kt01`, `kt02` | `LX-Load (a0) -> Produce -> Gather` | +| `kt03` | `LX-Load (a0) -> Produce -> Gather -> Slice -> LX-Store (a1, rows 32..63) -> Slice -> LX-Store (a1, rows 0..31)` | +| `kt04` | `LX-Load (a0) -> Produce -> Gather -> LX-Store (a1)` | +| `kt05` | `Produce[ LX-Load (a0) ] -> Scatter` | +| `kt06` A | `LX-Load (a0)` | +| `kt06` B | `LX-Load (a0) -> Produce -> Consume -> LX-Store (a0)` | +| `kt07a` | `LX-Load (a0) -> LocalReduce -> Expand -> Produce -> Reduce[ Add ]` | +| `kt07b` | `LX-Load (a0) -> LocalReduce -> Expand -> Produce -> Consume` | +| `kt09` 1 | `LX-Load (a0) -> Produce -> Gather -> Mul x 3 -> LX-Store x 3 (a1, a2, a3)` | +| `kt09` 2 | `Produce -> Gather -> LX-Store (a0) -> LX-Load x 2 (a0) -> Add -> LX-Store (a1)`, twice | + +Six things are visible only in this view: + +- **`kt05` is the only file whose load is *inside* the produce region.** Everywhere + else every tile produces, so the load can sit outside; there only one tile per + group produces and the group's eight consumers must not run it (§2.2). +- **`kt06` variant A is one op long, and the missing delivery is the claim** — §9.1 + row 1's "no op needed" branch, which is also why it passes legality. +- **`kt06` variant B loads and stores at the same address.** The permutation is + entirely in the dependency attribute; no local address moves. +- **`kt03`'s two stores go to one view at swapped base rows.** That is where its + reorder lives, and since a copy delivery must land every received tile in LX + anyway, those two stores are the mandatory landings rather than extra work. +- **`kt09` half 1 is the only sequence with one delivery and several readers**, and + the only one where a wrong version has an identical shape *and* identical values — + three deliveries differ from one only in emitted transfer count. +- **`kt09` half 2 is the only sequence that lands the delivered value and reads it + back.** Everywhere else the delivered value stays a `tensor`, so the landing store + the hardware performs is not named in the IR and its live range is not either. + +`Reduce[ Add ]` and `Produce[ … ]` denote a region. `LocalReduce` is ordinary +`linalg` over the tile's own rows, before any inter-tile op. + +Execution and application verdicts are not attempted. §1 of the coverage document +requires all three separately; only the expression column is addressed here. + +**Spelling.** These files use the *implemented* spelling: +`!ktdp.tile_future<(T_p), groups = affine_set<...>>` — partial types +parenthesised, `groups` a keyword — and `inter_tile_produce` with no operand +types before `->`. §7 of `inter-tile-communication.md` writes +`!ktdp.tile_future` and `produce ... : T_p -> ...`, which matches +neither `KTDPTypes.td:222` nor `KTDP.td`'s `assemblyFormat`. Checked: + +``` +$ ktir-opt spec_form.mlir +spec_form.mlir:5:63: error: expected '->' +``` + +So §7 as printed is not parser input. Worth reconciling separately from the +semantics this directory is about. + +## What the examples establish + +**Receiving-only cores are the measured norm, not a corner case.** 64 of the +catalog's 130 records have destination owners that are not source owners, across +three route classes (`grouped_all_gather_with_replication`, +`replicate_or_owner_remap`, `all_gather`). R13 = `n` for the copy-only ops +(§10.1) is required by measurement, not only by argument. `kt02` is +GR-PF-002 (`relayouts[1]`), where each group has producers `{4g, 4g+2}` and +consumers `{4g..4g+3}` — and the producers' tile ids are **not adjacent**, which +is precisely why §3.3 defines `l` as a position rather than a tile id. + +**The catalog contains no reduction at all.** All 130 records are `STCDPOpLx` +relayouts and every route class is copy-only: `all_gather` 26, +`grouped_all_gather_with_replication` 65, `replicate_or_owner_remap` 32, +`permutation` 6, `general_relayout` 1. So KT-07's two halves come from different +places: (b) is what every one of the 130 is, while (a) — a genuine cross-core +fold — comes from split-K matmul, the pattern `inter_tile_reduce` was implemented +for. The risk KT-07 guards against is therefore one-sided: the mistake to avoid +is reading a copy as a fold, and the measured consumer names (`mean_*`, +`_safe_softmax-Sum`, `mm-BMM_1`) invite exactly that. + +**A fold op cannot stand in for a copy op.** Writing KT-07 (b)'s broadcast as a +degenerate `inter_tile_reduce` over one producer is rejected by the shipped +legality pass: + +``` +error: consumer_tiles_per_group for group 0 is not a subset of + producer_tiles_per_group (a consumer tile that did not produce is + unsupported; see open question Q1) +``` + +That is R13 (`KTIRCheckLegality.cpp:107-117`); R14's mode gate would reject it +too. `consume` is needed, and this is evidence rather than argument. + +**The physical layout is not guessable from the inventory.** Its Source shape +column is alphabetically sorted, not `layoutDimOrder_` order. A real SDSC run of +the same logical shape (512 x 4096 x 1, fp16) reports + +``` +layoutDimOrder_ = ["mb", "out", "y"] stickDimOrder_ = ["y"] stickSize_ = [64] +device_size = [1, 4096, 512, 64] device_coordinates = [0, c1, c0, 0] +``` + +so `y` is both the innermost logical axis and the stick axis, an extent-1 `y` is +carried as 64 padded lanes, the physical form is one rank higher than the logical +one, and the two data axes appear in reverse order. `kt02` uses this; its one +remaining assumption is that the record's order is `["mb", "in", "y"]`, `in` +taking the position `out` held. + +**Overlapping-but-differing dependency sets do not occur in measurement.** 14 of +the 130 records have a source piece feeding several destination pieces, but in +**0 of 130** do two consumers declare different overlapping sets — a shared +source is always taken by consumers declaring the *same* set. So R5's scoping is +required by measurement while §3.3's "which set" clause is not yet forced by it. +`kt01` therefore has two functions: the measured form, which cannot tell §3.3's +two readings apart, and a synthetic one that can, since a shared producer lands +at different positions in two consumers' assemblies. The same function +discriminates the `P` derivation of §3.1, where the measured form again cannot: +its producer count and its `|dep|` are both 2. + +**The §9.1 guard forces a select rather than forbidding the record.** `kt05` is +GR-PF-055, whose destination regions are 32 slivers of row 511 — a 512th of the +tensor, so the coverage clause rejects the pre-select pair. After the select the +guard is satisfied, which §9.1 states itself ("the value delivered ... after a +select that is the selected sub-tensor and not the original"), and the post-select +pair classifies as row 2, `inter_tile_scatter` with `scatter_dimensions = R`. The +op in the file is derived, not chosen. This record also needs no layout +assumption: its axes are literally `mb`, `out`, `y`, matching the measured +`layoutDimOrder_`. + +**KT-03 is satisfied by two stores; what it costs is verifiability.** On the +gathered axis, the sources contributing to a destination piece are in ascending +tile-id order in **130 of 130** records, so §3.3's ordering rule is never +contradicted by measurement — `kt03` is synthetic for that reason. Where the order +does disagree, the reorder is expressible today because the assembled value has no +memory identity until it is stored: two stores at swapped base coordinates carry +it, with identity order inside each and no extra data movement. A single access +tile cannot, since it is a base plus a region relative to that base and offers no +way to permute positions *within* a dimension. Coverage §5 permits exactly this +form — "local selection/reordering is acceptable if fully expressed and stays +on-chip" — so **the requirement needs no new capability.** For a live intermediate +that is never stored, the reorder becomes an ordinary `tensor` permutation: still +local and on-chip, so still allowed, but no longer zero-copy. + +What it does cost is verification. **No rule relates the store bases to §3.3's +assembly order.** A verifier sees two well-formed stores that between them cover +the region exactly once, which is all it is asked to see; storing the assembly +verbatim has the same shape, the same element count, the same coverage, and the +wrong answer. So the correctness of the reorder falls **outside inter-tile +verification, and the numerical check is its only guarantee.** An ordering +attribute on the delivery op redefining §3.3's `l` was considered and rejected on +that basis: an attribute can be checked for well-formedness but never for intent, +so it would not move the case into the verifier — it would only add a second +statement of the ordering that can disagree with the stores. + +**One delivery can have many readers, and R2 does not stand in the way.** `kt09` is +the case where the catalog's own indexing invites the mistake. `mean-LayerNormNorm_out` +is the source tensor of **three** records — `relayouts[1, 8, 16]`, the Q/K/V +projections off one layer norm — and the three are **identical** in extents, pieces, +owner tables and route class, each recording 12 MiB of remote traffic. They are one +relayout that three consumers need. Since the catalog is indexed by (consumer, input), +a reader that walks records one at a time emits one transfer each and moves **36 MiB +where 12 suffices, exactly 3×**. The spelling that avoids it is the natural one: R2 +constrains the `tile_future`, not the delivery's *result*, so one `produce` + one +delivery + N readers of a plain `tensor` is legal and is what the file shows. Eight of +the catalog's 120 distinct tensors are shared this way. + +Note what separates right from wrong here: **nothing numerical.** Three deliveries +give the same values as one. Only the emitted transfer count distinguishes them, which +is coverage §1's "check the emitted memory accesses as well as numerical output". + +**Lifetime splits into an expressible half and an inexpressible one.** The *source* +buffer's anti-dependence is carried by ordinary MLIR: iteration 1's `ktdp.store` into +the scratch cannot be hoisted above iteration 0's `ktdp.load` from it, because those +are conflicting memory effects on one `memref`. No inter-tile rule is needed. But the +*delivered value's* buffer lifetime cannot be expressed at all — the delivery returns +a `tensor`, so its readers touch no memref and the memref read ended at `ktdp.load`, +before the produce. **At this level a reader cannot extend a live range because there +is no live range to extend**, so coverage §1's "buffers preserved until last reader" +lands on the emitted program. That is the same "no address" fact `kt03` used to put its +reorder in store bases and that Table 1 records for dist-mem-view intermediates — +three cases, one root cause. + +**Multi-axis assembly needs no new op, and its verdict splits.** `kt04` is +`inter_tile_gather` with two entries in `gather_dimensions` rather than one — §9.1 +row 4, the same row as `kt01` and `kt02`. What multi-axis adds is that §4's +flattening becomes load-bearing, and two of the three errors it invites are closed +by rule. `gather_dimensions = [1, 0]`, a column-major assembly, is rejected by R9's +ascending-order requirement, which §4 says exists for exactly that reason: `[2, 0]` +and `[0, 2]` "would flatten to *different* data orders, so a reversed list passes +every other check while meaning something else". `gather_dimensions = [1]`, one axis +where two are needed, is rejected by R12, whose per-axis obligation is explicit that +"equal products alone would not give a well-defined multi-axis assembly". R12 also +confirms the ascending rule reaches `gather_dimensions` and not only the split ops. + +The third error is closed by nothing, and the requirement's own toy is that fixture. +Read its ownership as (row, column) — the reading under which its stated wrong +answer `[0,2,1,3]` is what a gather actually produces — and `source1` holds cell +`(1,0)` while sitting at ascending position `l = 1`, which §4's odometer sends to +`(0,1)`. The assembly comes out `[[0,2],[1,3]]` where `[[0,1],[2,3]]` is required. +Every rule is satisfied: ascending list, both axes present, uniform extents, +`P = 4 = 2 × 2`, and the declared `tensor<2x2xf16>`. **Same shape, same element +count, transposed answer.** This is `kt03`'s finding in multi-axis form, and the fix +is again outside the delivery op — assign tile ids upstream to match the odometer, +store in pieces at swapped bases, or permute the assembled tensor locally. + +Measurement does not force the problem. `relayouts[120]` (`cat_1-kvCacheScatter`, +16 pieces → 1, `P = 16`, concat `mb` ×8 and `out` ×2) **agrees with the odometer for +all 16 producers** — ascending owner position `l` maps to `(mb = l/2, out = 64·(l mod +2))`, checked off the `owners` fields. So the measured function needs nothing beyond +a plain two-axis gather, exactly as `kt03`'s 130 of 130 never contradict §3.3. + +**An owner permutation is verifiable; a within-region reorder is not.** `kt06` is +the same family of problem as `kt03` — a required order disagreeing with the order +the IR supplies — with the opposite outcome, and the difference is only where the +ordering information is allowed to live. GR-PF-052 (`relayouts[51]`, plus four +siblings with identical geometry) has 32 pieces going to 32 pieces of **the same +size**, `P = 1`, a bijection. Reading the `owners` field of every piece: with +`J = j/2` and `M = mb/64`, the source owner is `8J + M` and the destination owner is +`J + 4M` — a transpose of a 4×8 grid, which is exactly what coverage §5 means by +"row-major versus column-major core order". Because the reorder is **across cores +rather than within a region**, it lives in `producer_dependency_per_consumer`, where +R3 through R8 all reach it: a dropped producer fails R4, a doubled one fails R5. +What the rules cannot catch is a bijection that is the *wrong* bijection — using +`π` where `π⁻¹` belongs — which is coverage §5's "wrong core map with identical +split counts", and only the values distinguish it. `kt03`, by contrast, has to put +its ordering in store bases where no rule reaches it at all. + +**The two variants are the two branches of §9.1 row 1, not two spellings.** With +`C = ∅` and `R = ∅`, that row resolves to "`no op needed` if every core's region is +its own, else `inter_tile_consume` — a **relocation**". Coverage §5's "preserve or +explicitly change ownership" names the same pair, and which branch applies is a +property of the *consumer*, not of the edge. So `kt06`'s variant A contains no +inter-tile op — and that absence is the claim, which is also why it is the one file +besides `kt07a` that passes legality today. Its precondition is unexpressed +anywhere in the IR, which is the failure coverage §5 warns about: same ops, same +types, same addresses, wrong answer on 30 of 32 cores, because only owners 0 and 31 +are fixed points of the transpose. + +**Two traps in reading the catalog, found while deriving that map.** First, +`source_pieces` is ordered lexicographically by its `key` string (`p0, p1, p10, +p11, … p2, p20`), so **piece index is not owner order**, and `source_core_patterns` +is a separate summary carrying no piece key. The authoritative owner is the +`owners` field on each piece; deriving the map from index position instead yields a +bijection that is *not* a transpose — an artifact of the ordering. Second, the +catalog's byte fields are **logical**: `prod(extents) × word_length == +logical_tensor_bytes` exactly, including for GR-PF-055 where an SDSC run showed a +`y` extent of 1 is physically 64 padded lanes. So padding is invisible in the +catalog by construction, and no record can answer a layout question on its own. +(`source_piece_bytes` is also the total over all pieces, not the size of one.) + +**An aside on `access_tile_order`, independent of the above.** RFC 0682 defines it +twice over, as "the rightmost dimension in the output space corresponds to the +innermost iteration dimension" and as "the enumeration of points in the +intermediate variable space". The first reads as a dimension nesting order, the +second as a sort key; only the second could reorder within a dimension. `kt03` +deliberately does not depend on the generous reading, and `KTDP.td`'s op +description carries neither sentence. Worth settling in the dialect, but nothing +here waits on it. + +## Not written yet + +| Case | What it needs | +|---|---| +| KT-08 | Stick indivisibility. Measured, in four shapes: `y` 32→64 in `relayouts[39..50]` (each source holds half a stick), `y` 4→2 in `relayouts[57,64,73]` (a split inside one stick), `y` 1→4 in `[56,63,71]`, `y` 24→192 in `[117]`. **Needs an SDSC run first** — the catalog's byte fields are logical, so it cannot say whether `y` is sticked here or at what size (see the layout note in `kt06`). | +| KT-10 | Deliberate bad variants: remove a required piece, duplicate a piece inside one destination, corrupt an owner, supply a wrong output shape. Only `kt07a`'s and `kt06` variant A's negatives can actually be run today. | + +The negative variants of KT-10 matter for every file here: coverage §5 requires +that reversing two fragments, omitting one, or reusing a wrong core map with +identical split counts all **fail**. None of that is written yet. + +The production-shaped four-way V case (source `mb:32`, destination +`mb:8, qpk:4`) is also absent: it needs the saved physical descriptors, which are +not in this repository. diff --git a/docs/inter-tile-examples/kt01-shared-sources.mlir b/docs/inter-tile-examples/kt01-shared-sources.mlir new file mode 100644 index 0000000..b8e3b3b --- /dev/null +++ b/docs/inter-tile-examples/kt01-shared-sources.mlir @@ -0,0 +1,305 @@ +// KT-01 — a source shared by several receivers, written as an explicit +// dependency. +// +// torch-spyre#4300 lx_relayout_workload_coverage.md §5 KT-01: "Four-way V toy +// above, default and explicit equal dependencies. **R5 must not reject valid +// sharing merely because different receiving cores use the same producer +// pieces.** This concern is about explicit dependencies; default gather already +// describes an all-producer assembly." And §4: "Check both the default +// all-producer dependency and an explicitly written equivalent dependency. They +// should describe the same values." +// +// OPERATION SEQUENCE. Both functions: +// +// LX-Load (addr_0) -> Produce -> Gather +// +// addr_0 = the producer's own piece in its LX. The delivered region is left +// unstored: this file is about the dependency attribute, not about placement. +// +// ============================================================================ +// SAME FIXTURE AS kt02-receiving-only-cores.mlir — ONLY THE ATTRIBUTE IS ADDED. +// +// GR-PF-002, i.e. `relayouts[1]` of the pinned ownership catalog +// (sendnn_sdsc_lx_replay_manifest.json, sha256 4c8aca2e…c747d4), consumer +// `mm-BMM_1` input 0, route_class grouped_all_gather_with_replication. +// +// Group structure, from the owner tables. Source piece k is owned by core 2k +// and starts at mb = 32k, so group g (destination mb 64g..64g+63) is: +// +// group g producers {4g, 4g+2} consumers {4g .. 4g+3} dep(c,g) non-prod +// g=0 {0, 2} {0, 1, 2, 3} {0, 2} 1, 3 +// g=1 {4, 6} {4, 5, 6, 7} {4, 6} 5, 7 +// g=2 {8, 10} {8, 9, 10, 11} {8, 10} 9, 11 +// ... 8 groups, regular +// g=7 {28, 30} {28, 29, 30, 31} {28, 30} 29, 31 +// +// Source piece {mb:32, in:4096, y:1}, one owner each; destination piece +// {mb:64, in:4096, y:1}, four owners each. +// +// kt02 leaves `producer_dependency_per_consumer` absent, which §3.4 defines as +// "the consumer waits on and receives from **all** producer tiles in the group". +// This file writes that same set out. The two must describe the same values. +// +// default form (kt02) attribute absent -> P = |{4g, 4g+2}| = 2 +// explicit form (here) dep(c, g) = {4g, 4g+2} -> P = 2, for every c +// +// WHY kt02 IS NOT ENOUGH. §4 asks to "check **both** the default all-producer +// dependency and an explicitly written equivalent dependency", so the pair is the +// requirement and neither file satisfies it alone. More than that: **R5 cannot +// be exercised by kt02 at all** — the default form carries no dependency +// attribute, so there is nothing for the rule to apply to. Coverage §5 says as +// much: "This concern is about explicit dependencies; default gather already +// describes an all-producer assembly." Writing the set is what makes R5 able to +// fire, and it also pins P explicitly rather than by default, so a verifier that +// mis-derived P from the producer set would show up here. +// ============================================================================ +// +// WHAT IT SETTLES. All four consumers of a group name **the same two +// producers**, so their declared dependency sets are identical — not disjoint. +// Read as a blanket rule, R5's pairwise disjointness would reject this, and with +// it every all-gather (§6.4: the full group as consumer set, every consumer +// naming every producer) and every multicast source (R8: "a producer **may** +// serve several consumer tiles"). §5 now scopes the obligation to a +// *partitioning* use and says so directly: "A verifier must therefore not reject +// overlap as such." +// +// Well-definedness does not come from disjointness. It comes from §3.3's rule +// that an assembling consumer's positions are taken from **its own** declared +// set: here that set is {4g, 4g+2}, so l(4g) = 0 and l(4g+2) = 1 and the two +// slabs land in mb order. Whatever the other three consumers declare cannot +// disturb it. +// +// The other dependency rules are satisfied and worth checking against: +// R3 dep(c, g) subset of producer_tiles_per_group(g) {4g,4g+2} ⊆ {4g,4g+2} +// R4 every producer named by some consumer both, by all four +// R6 uniform cardinality across consumers |dep| = 2 for all +// R13 consumers need not be producers 4g+1, 4g+3 are not +// +// Note the producers' tile ids are **not adjacent** — 4g and 4g+2. That is why +// §3.3 defines `l` as a position in ascending tile-id order rather than a tile +// id: taken literally, tile id 4g+2 would name slot 2 of a two-slot assembly. +// +// Geometry, expected values and the physical layout are as in kt02; see that +// file, which carries the SDSC-verified `layoutDimOrder_` / `stickDimOrder_` +// derivation and the one assumption left in it. +// +// ============================================================================ +// TWO FUNCTIONS, BECAUSE THE MEASURED FORM CANNOT DISCRIMINATE §3.3. +// +// @kt01_shared_sources — measured. GR-PF-002, explicit dependency naming both +// producers for every consumer. Operationally identical to kt02: the attribute +// restates the default, so the two files must deliver the same values, and that +// agreement is the test coverage §4 asks for. It exercises R5's scoping, since +// four identical sets are not disjoint. It does NOT exercise §3.3's "which set" +// clause: with identical sets a producer lands at the same position under either +// reading, so the clause is invisible here. +// +// @kt01_overlapping_sets — SYNTHETIC. Sets that overlap and differ, which is +// the only form that discriminates §3.3. Searched for in the catalog and not +// found: 14 of 130 records have a source piece feeding several destination +// pieces, but in **0 of 130** do two consumers declare different overlapping +// sets — a shared source is always taken by consumers that declare the *same* +// set. So §3.3's clause is currently unforced by measurement while R5's scoping +// is required by it, and both facts belong in the record. +// +// CONSISTENCY TO VERIFY. A written dependency set is a second, independent +// statement of something the producer set already says, so the two can disagree. +// Four checks, in the order a verifier would reach them: +// +// 1. dep(c, g) subset of producer_tiles_per_group(g) (R3) +// 2. every producer named by at least one consumer (R4) +// 3. |dep(c, g)| the same for every consumer of the group (R6) +// 4. **P taken from dep, not from the producer set**, and the declared +// result type following from that P (§3.1, R12) +// +// Check 4 is the one with two possible answers. §3.1 defines `P` as +// `|producer_tiles_per_group(g)|` when the attribute is **absent** and the +// per-consumer cardinality when it is **present**, so a verifier has to pick the +// right source. The two functions differ in exactly this: +// +// @kt01_shared_sources producers/group 2, |dep| 2 -> cannot discriminate +// @kt01_overlapping_sets producers/group 4, |dep| 2 -> discriminates +// +// In the second, taking P from the producer set gives 4 and would make mb go +// 32 -> 128, contradicting the declared tensor<1x4096x64x64xf16>. Taking it from +// dep gives 2 and 32 -> 64, which is what the file declares. So the synthetic +// function is the discriminating test for the P derivation as well as for §3.3's +// position derivation — the same structure of problem, and the measured form is +// blind to both. +// +// The fifth check is not local to one file: **the explicit-all form must deliver +// what the default form does**, which means running kt02 and +// @kt01_shared_sources on the same input and comparing their outputs to each +// other, not only each to a reference (§4). +// ============================================================================ +// +// NOT VERIFIED as a whole: `ktdp.inter_tile_gather` is specified (§6.4) but +// absent from KTDP.td. The `ktdp.inter_tile_produce` half of both functions +// parses and round-trips, as do all the dependency sets on their own. + +#producers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 2 >= 0, + i mod 2 == 0)> +#consumers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 3 >= 0)> +#groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +// producer_dependency_per_consumer: `(p)[c, g]` per §3.4. Every consumer `c` of +// group `g` names both producers, so `c` does not appear in the constraints — +// the mapping is the same for every consumer, which is exactly the sharing R5 +// must not reject. Written out rather than omitted so the two forms can be +// compared; omitting it is kt02. +#dep_all = affine_set<(p)[c, g] : (p - 4 * g >= 0, -p + 4 * g + 2 >= 0, + p mod 2 == 0)> + +// The whole of the producer's own piece: [1, 4096, 32, 64]. +#piece = affine_set<(d0, d1, d2, d3) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 4095 >= 0, + d2 >= 0, -d2 + 31 >= 0, d3 >= 0, -d3 + 63 >= 0)> +#ident4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> + +module { + func.func @kt01_shared_sources() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + + // The producer's own LX piece, at offset 0. `ct_local` with no ct_id is the + // executing tile's local memory, so no tile id enters the address. + %own = ktdp.construct_memory_view %base, + sizes: [1, 4096, 32, 64], strides: [8388608, 2048, 64, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<1x4096x32x64xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident4 + } : memref<1x4096x32x64xf16> -> !ktdp.access_tile<1x4096x32x64xindex> + + %piece_val = ktdp.load %own_access + : !ktdp.access_tile<1x4096x32x64xindex> -> tensor<1x4096x32x64xf16> + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producers + -> !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %piece_val : tensor<1x4096x32x64xf16> + } + + // The only difference from kt02: producer_dependency_per_consumer is + // written, naming both producers for every consumer. P is unchanged at 2, + // so the result type is unchanged and the delivered values are the same as + // the default form's. A verifier that rejected the identical sets here + // would reject all-gather too. + %region = ktdp.inter_tile_gather(%future) + consumer_tiles_per_group = #consumers, + gather_dimensions = [2], + producer_dependency_per_consumer = #dep_all + : !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + -> tensor<1x4096x64x64xf16> + + // Expected at cores 4g..4g+3: destination piece g, logical + // (mb = 64g..64g+63, in = 0..4095, y = 0), core 4g supplying the lower mb + // half and core 4g+2 the upper. All four hold the same assembled region + // (§3.7), byte for byte identical to what kt02 delivers. Undefined + // elsewhere. + // + // The pair is the test: run kt02 and this file on the same input and compare + // the two outputs, not just each against a reference. "They should describe + // the same values" is the requirement, so a difference between the two forms + // is itself a failure even if both look plausible alone. + return + } +} + +// --------------------------------------------------------------------------- +// SYNTHETIC. Overlapping, differing dependency sets. +// +// Not a measured shape: no record in the catalog has two consumers declaring +// different overlapping sets (0 of 130). Written because it is the only form +// that tells §3.3's two readings apart, so if the clause is ever load-bearing +// this is the shape that will show it. +// +// Deviation from GR-PF-002, stated so it is not mistaken for measurement: four +// producers per group instead of two, and three consumers instead of four, so +// that a two-wide sliding window fits inside a group. The piece type is +// unchanged. +// +// Group structure, same layout as the table above. +// +// group g producers {4g .. 4g+3} consumers {4g .. 4g+2} produces only +// g=0 {0, 1, 2, 3} {0, 1, 2} 3 +// g=1 {4, 5, 6, 7} {4, 5, 6} 7 +// g=2 {8, 9, 10, 11} {8, 9, 10} 11 +// ... 8 groups, regular +// g=7 {28, 29, 30, 31} {28, 29, 30} 31 +// +// dep(c, g) = {c, c+1}, so within group g the three consumers name +// 4g->{4g,4g+1}, 4g+1->{4g+1,4g+2}, 4g+2->{4g+2,4g+3}. Tile 4g+3 produces but +// never consumes. +// +// Positions within each consumer's own set (§3.3): +// +// consumer 4g+0 : {4g+0, 4g+1} l(4g+0)=0, l(4g+1)=1 +// consumer 4g+1 : {4g+1, 4g+2} l(4g+1)=0, l(4g+2)=1 +// consumer 4g+2 : {4g+2, 4g+3} l(4g+2)=0, l(4g+3)=1 +// +// PRODUCER 4g+1 LANDS AT SLOT 1 IN ONE ASSEMBLY AND SLOT 0 IN ANOTHER. That is +// possible only if positions come from the consumer's own declared set, which is +// what §3.3 now states. Taken from the group's producer set instead, 4g+1 would +// have one position and the two assemblies could not both be right. +// +// Rules, checked: R3 dep ⊆ producers; R4 all four producers named by some +// consumer (4g+0 by c=4g+0; 4g+3 by c=4g+2); R6 |dep| = 2 for every consumer, so +// P = 2 and the result type is single-valued; R5 scoped — the sets overlap and +// are not a partition, so the disjointness obligation does not apply. +// +// Expected values: the same bit-per-coordinate scheme as kt02. The check that +// matters is that a coordinate supplied by producer 4g+1 appears in the *upper* +// mb half of consumer 4g+0's region and the *lower* half of consumer 4g+1's. +// Swapping those two is the failure this example exists to catch. +// --------------------------------------------------------------------------- + +#producers4 = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 3 >= 0)> +#consumers3 = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 2 >= 0)> +// dep(c, g) = {c, c+1}. `g` is unused: the rule is the same in every group. +#dep_window = affine_set<(p)[c, g] : (p - c >= 0, -p + c + 1 >= 0)> + +module { + func.func @kt01_overlapping_sets() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + + %own = ktdp.construct_memory_view %base, + sizes: [1, 4096, 32, 64], strides: [8388608, 2048, 64, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<1x4096x32x64xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident4 + } : memref<1x4096x32x64xf16> -> !ktdp.access_tile<1x4096x32x64xindex> + + %piece_val = ktdp.load %own_access + : !ktdp.access_tile<1x4096x32x64xindex> -> tensor<1x4096x32x64xf16> + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producers4 + -> !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %piece_val : tensor<1x4096x32x64xf16> + } + + // P = 2 from the declared subsets (§3.4), not 4 from the producer set, so + // mb goes 32 -> 64 and each consumer assembles only its own two slabs. + %region = ktdp.inter_tile_gather(%future) + consumer_tiles_per_group = #consumers3, + gather_dimensions = [2], + producer_dependency_per_consumer = #dep_window + : !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + -> tensor<1x4096x64x64xf16> + + // Expected: consumer c holds producers c and c+1 in that order. Results are + // undefined at 4g+3, which consumes nothing here (§3.7). + return + } +} diff --git a/docs/inter-tile-examples/kt02-receiving-only-cores.mlir b/docs/inter-tile-examples/kt02-receiving-only-cores.mlir new file mode 100644 index 0000000..08bf1ce --- /dev/null +++ b/docs/inter-tile-examples/kt02-receiving-only-cores.mlir @@ -0,0 +1,163 @@ +// KT-02 — receiving-only cores, on a measured Granite relayout. +// +// OPERATION SEQUENCE. +// +// LX-Load (addr_0) -> Produce -> Gather +// +// addr_0 = the producer's own piece in its LX. The load is *outside* the produce +// region because every tile here is a producer; contrast kt05. +// +// ============================================================================ +// Fixture: GR-PF-002 of lx_relayout_granite_inventory.md, i.e. relayouts[1] of +// the pinned ownership catalog +// sendnn_sdsc_lx_replay_manifest.json, sha256 +// 4c8aca2e1989eefb76c4ee40b99a4a87800aac4134cbae68c41810ab98c747d4 +// read directly rather than from the summary row, which cannot distinguish the +// divisions that produce the same `32 -> 8` counts. +// +// consumer mm-BMM_1, input 0 (family P09) +// route_class grouped_all_gather_with_replication +// extents {mb: 512, in: 4096, y: 1} src == dst +// word_length 2 (fp16) +// remote_required true +// 16 source pieces {mb: 32, in: 4096, y: 1}, one owner each +// 8 destination pieces {mb: 64, in: 4096, y: 1}, four owners each +// +// Group structure, derived from the owner tables. Source piece k is owned by +// core 2k and starts at mb = 32k, so group g (destination mb 64g..64g+63) is: +// +// group g producers {4g, 4g+2} consumers {4g .. 4g+3} non-producers +// g=0 {0, 2} {0, 1, 2, 3} 1, 3 +// g=1 {4, 6} {4, 5, 6, 7} 5, 7 +// ... 8 groups, regular +// g=7 {28, 30} {28, 29, 30, 31} 29, 31 +// +// WHAT IT SETTLES. Cores 4g+1 and 4g+3 receive without producing, so R13 must +// be `n` for `gather` (§10.1 of ../inter-tile-communication.md). This is not a +// corner case: 64 of the catalog's 130 records have destination owners that are +// not source owners, across three route classes +// (grouped_all_gather_with_replication, replicate_or_owner_remap, all_gather). +// +// It also exercises §3.3 directly. The producers of a group are 4g and 4g+2 — +// NOT adjacent tile ids — so `l` cannot be the tile id. Taken as a position in +// ascending tile-id order, l(4g) = 0 and l(4g+2) = 1, and since core 4g holds +// the lower mb half, ascending l reproduces ascending mb. Using tile ids +// directly would leave slot 1 empty and slot 2 filled. +// +// ============================================================================ +// Physical layout. Taken from a real SDSC run of the same logical shape +// (512 x 4096 x 1, fp16), whose descriptor reports +// +// layoutDimOrder_ = ["mb", "out", "y"] stickDimOrder_ = ["y"] +// stickSize_ = [64] +// device_size = [1, 4096, 512, 64] +// device_coordinates = [0, c1, c0, 0] c0 = mb(512), c1 = out(4096) +// stride_map = [1, -1, 4096, -1] 4096<->stride 1, 512<->stride 4096 +// +// So: `y` is the innermost logical axis AND the stick axis, and an extent-1 `y` +// is carried as 64 padded lanes — only lane 0 holds data. The physical form is +// one rank higher than the logical one, the two logical data axes appear in +// reverse order, and the stick is last: +// +// physical = [1, , , 64] +// +// ASSUMPTION, the only one left: this record's layoutDimOrder_ is +// ["mb", "in", "y"]. Confirmed for the run above is ["mb", "out", "y"] — `y` +// innermost and sticked, `mb` first among the data axes — and `in` takes the +// position `out` held. Only the numeric index in `gather_dimensions` depends +// on it; if the order were ["in", "mb", "y"] the gathered axis would move from +// physical 2 to physical 1. +// +// whole tensor logical (mb=512, in=4096, y=1) -> [1, 4096, 512, 64] +// producer piece logical (mb=32, in=4096, y=1) -> [1, 4096, 32, 64] +// destination region logical (mb=64, in=4096, y=1) -> [1, 4096, 64, 64] +// +// The gathered axis is `mb`, physical dim 2. It is NOT the stick axis, so +// `gather_dimensions` is a single index and §4's floordiv rule does not apply: +// P = 2 multiplies the mb extent 32 -> 64 and nothing else moves. +// +// EXPECTED VALUES. 4096 x 512 fp16 cannot be checked with a ramp — a large +// float ramp rounds distinct coordinate IDs together (coverage §1). Drive one +// logical coordinate at a time with an exactly representable 1.0 against 0.0 +// and check it lands where the owner tables say. Lanes 1..63 are `y` padding +// and must not be compared. +// +// NOT VERIFIED as a whole: `ktdp.inter_tile_gather` is specified (§6.4) but +// absent from KTDP.td. The `ktdp.inter_tile_produce` half parses and +// round-trips through ktir-opt. +// ============================================================================ + +// Producers of group g: {4g, 4g+2}. Even tile ids in [4g, 4g+2]. +#producers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 2 >= 0, + i mod 2 == 0)> +// Consumers of group g: {4g, 4g+1, 4g+2, 4g+3}, the destination piece's owners. +#consumers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 3 >= 0)> +#groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +// The whole of the producer's own piece: [1, 4096, 32, 64]. +#piece = affine_set<(d0, d1, d2, d3) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 4095 >= 0, + d2 >= 0, -d2 + 31 >= 0, d3 >= 0, -d3 + 63 >= 0)> +#ident4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> + +module { + func.func @kt02_receiving_only_cores() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + + // `ct_local` with no ct_id is the executing tile's own LX. The pieces + // start distributed one per producer (source_owner_group_sizes = [1]), so a + // fresh HBM load of the assembled region at one core would not reproduce + // this starting ownership. `remote_required` is true in the record. + // Dense in its own buffer: lane 1, mb 64, in 32*64 = 2048, and the extent-1 + // axis takes the remaining product. + %own = ktdp.construct_memory_view %base, + sizes: [1, 4096, 32, 64], strides: [8388608, 2048, 64, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<1x4096x32x64xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident4 + } : memref<1x4096x32x64xf16> -> !ktdp.access_tile<1x4096x32x64xindex> + + // Function scope: every producer reads only its own piece at offset 0, so + // no tile id enters the address and the load need not sit inside the + // region (§2.2). + %piece_val = ktdp.load %own_access + : !ktdp.access_tile<1x4096x32x64xindex> -> tensor<1x4096x32x64xf16> + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producers + -> !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %piece_val : tensor<1x4096x32x64xf16> + } + + // combine = none, placement = concat, gathered axis = mb = physical dim 2. + // P = 2, so mb goes 32 -> 64 (R12). All four consumers of the group hold + // the same assembled region (§3.7), which is the `with_replication` half of + // the record's route_class. + %region = ktdp.inter_tile_gather(%future) + consumer_tiles_per_group = #consumers, + gather_dimensions = [2] + : !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + -> tensor<1x4096x64x64xf16> + + // Expected at cores 4g..4g+3: destination piece g, logical + // (mb = 64g..64g+63, in = 0..4095, y = 0), with core 4g supplying the lower + // mb half and core 4g+2 the upper. Undefined elsewhere (§3.7). + // + // All four consumers hold the SAME assembled region, not a quarter each: + // the record carries one destination rectangle with four owners, and §3.7 + // gives `concat` "the same assembled tensor" for tiles in one group. So the + // check runs per consumer, and all four must agree — including on order, + // which is what makes coverage §5's "reverse two fragments" check bite. + // + // Counted per core the destination covers 8 * 64 * 4 = 2048 mb rows against + // a 512-row tensor, 4x over. That is why §9.1's coverage clause is stated on + // *distinct* regions: the 8 distinct ones sum to 512 and pass the guard. + return + } +} diff --git a/docs/inter-tile-examples/kt03-logical-order.mlir b/docs/inter-tile-examples/kt03-logical-order.mlir new file mode 100644 index 0000000..b6eab73 --- /dev/null +++ b/docs/inter-tile-examples/kt03-logical-order.mlir @@ -0,0 +1,219 @@ +// KT-03 — logical order disagreeing with tile-id order. +// +// torch-spyre#4300 lx_relayout_workload_coverage.md §5 KT-03: "Source core0 holds +// `[2,3]`, core1 holds `[0,1]`; receiver needs `[0,1,2,3]`. **Show how logical +// order is preserved when it differs from core-ID order.** Local +// selection/reordering is acceptable if fully expressed and stays on-chip." +// +// OPERATION SEQUENCE. +// +// LX-Load (addr_0) -> Produce -> Gather +// -> Slice -> LX-Store (addr_1, rows 32..63) +// -> Slice -> LX-Store (addr_1, rows 0..31) +// +// addr_0 = own piece, addr_1 = the assembled region. **The two stores are the +// whole point**: same destination view, swapped base rows, so the reorder is +// carried by which base each slice goes to. Both stores are landings the transfer +// requires anyway, so the swap costs nothing. +// +// ============================================================================ +// SYNTHETIC, WITH ONE THING CHANGED. +// +// Geometry, groups and piece types are GR-PF-002's, the same as kt01 and kt02: +// +// group g producers {4g, 4g+2} consumers {4g .. 4g+3} non-producers +// g=0 {0, 2} {0, 1, 2, 3} 1, 3 +// g=1 {4, 6} {4, 5, 6, 7} 5, 7 +// ... 8 groups, regular +// g=7 {28, 30} {28, 29, 30, 31} 29, 31 +// +// Only which core holds which half differs: +// +// measured (kt02) core 4g holds mb 64g .. 64g+31 (lower) +// core 4g+2 holds mb 64g+32.. 64g+63 (upper) +// here core 4g holds the UPPER half +// core 4g+2 holds the LOWER half +// +// §3.3 orders an assembly by ascending tile id, so l(4g) = 0 and l(4g+2) = 1 and +// the gather produces [upper, lower]. The consumer needs [lower, upper]. No +// dependency set fixes this: a set carries no order, and the order comes from the +// tile id, which the work division fixed upstream. +// +// This shape does not occur in measurement. Checked across all 130 records of +// the pinned catalog: on the gathered axis, the sources contributing to a +// destination piece are in ascending tile-id order in **130 of 130**. So §3.3's +// rule is never contradicted by the measured data. The case is still worth +// answering, because the answer is not where one looks for it: the reorder lives +// at the store, not on the delivery op. +// ============================================================================ +// +// HOW THE REORDER IS EXPRESSED, AND WHAT IT COSTS. +// +// The gathered value is a `tensor` — no address, no memory space, and +// `!ktdp.tile_future` carries `tensor` rather than `memref`. So the assembly is +// not "in LX in the wrong order": it is not in LX at all. Memory identity begins +// at `ktdp.store`, and that is where the reorder goes. +// +// **A single access tile cannot do it.** An access tile is a base coordinate +// plus a region relative to that base, so there is no way to permute positions +// *within* a dimension. `access_tile_order` does not help: RFC 0682 says "the +// rightmost dimension in the output space corresponds to the innermost iteration +// dimension", which reads as a dimension nesting order — outer to inner — and a +// nesting order cannot reorder positions inside one dimension. (Its next +// sentence, "the enumeration of points in the intermediate variable space", reads +// more generously, and the two are not obviously the same thing. This file does +// not depend on the generous reading; see the aside at the end.) +// +// **Two stores do it, and avoid the problem rather than solving it.** Each store +// writes a contiguous run at its own base, with identity order inside, and the +// swap is carried entirely by *which base each slice goes to*: +// +// tensor rows 0..31 (upper, from core 4g) -> memory rows 32..63 +// tensor rows 32..63 (lower, from core 4g+2) -> memory rows 0..31 +// +// Cost: **none.** Not "two stores instead of one" — two stores instead of two. For +// a copy delivery this hardware lands every received tile in LX before a compute unit +// can read it, so P landing stores are mandatory whatever the order is. The two +// stores here *are* those landings, aimed at swapped bases. Every element is written +// exactly once and no pass over the region is added. +// +// So a permutation of P pieces costing "up to P stores" is not a cost at all: P is +// exactly how many landings the transfer already requires, and P is 2, 4 or 8 in the +// measured divisions. This is also why the rejected ordering attribute below would +// have bought no performance — there is no redundant write for it to remove. +// +// **The zero-copy case and the general case.** Two stores are zero-copy, and +// that depends on the value being stored. If the assembled value feeds `linalg` +// directly, as a live intermediate, there is no access tile to carry the reorder +// and it becomes an ordinary `tensor` permutation — still local, still on-chip, +// so still inside what the requirement allows, but with real data movement. +// Either way one intent is spread over P ops: the permutation becomes visible +// only after reading which slice goes to which base. +// +// THE REQUIREMENT IS SATISFIED. Coverage §5 asks to "show how logical order is +// preserved when it differs from core-ID order", and states that "local +// selection/reordering is acceptable if fully expressed and stays on-chip". Two +// stores are exactly that: fully expressed — every element's destination is read +// off the `%c32` / `%c0` bases — on-chip, and zero-copy. **So this case needs no +// new capability.** +// +// WHAT IT COSTS IS VERIFICATION. §3.3 fixes the assembly order, the two stores +// place the halves at swapped bases, and **no rule relates the second to the +// first.** A verifier sees two well-formed stores that between them cover the +// region exactly once, which is all it is asked to see. Storing the assembly +// verbatim instead has the same shape, the same element count, the same coverage +// — and the wrong answer. So the correctness of the reorder is **outside +// inter-tile verification**, and the numerical check is its only guarantee. That +// is why the expected values at the end of this file are its point, and the op +// list is not. +// +// AN ORDERING ATTRIBUTE WAS CONSIDERED AND REJECTED. A `gather_order` affine map +// on the delivery op, redefining §3.3's `l`, would put the intent in one place +// instead of spreading it over P stores. It is not worth it. An attribute can be +// checked for well-formedness — is it a permutation of `0..P-1`? — but never for +// intent, so it would not move this case out of "numerical test only" and into the +// verifier. It would only add a second statement of the ordering that can +// disagree with the stores. Verifiability is the reason the explicit form is +// preferred at all, so a construct that adds surface without adding checkable +// content is a loss. +// +// Aside, independent of this case: RFC 0682 defines `access_tile_order` twice over +// — "the rightmost dimension in the output space corresponds to the innermost +// iteration dimension", and "the enumeration of points in the intermediate +// variable space". The first reads as a dimension nesting order, the second as a +// sort key; only the second could reorder within a dimension. This file +// deliberately does not depend on the generous reading, and `KTDP.td`'s op +// description carries neither sentence. Worth settling in the dialect, but not a +// prerequisite for anything here. +// +// NOT VERIFIED as a whole: `ktdp.inter_tile_gather` is specified (§6.4) but absent +// from KTDP.td. Everything else here — the produce, both slices and both +// stores — parses and round-trips. +// ============================================================================ + +#producers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 2 >= 0, + i mod 2 == 0)> +#consumers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 3 >= 0)> +#groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +#piece = affine_set<(d0, d1, d2, d3) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 4095 >= 0, + d2 >= 0, -d2 + 31 >= 0, d3 >= 0, -d3 + 63 >= 0)> +#region = affine_set<(d0, d1, d2, d3) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 4095 >= 0, + d2 >= 0, -d2 + 63 >= 0, d3 >= 0, -d3 + 63 >= 0)> +#ident4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> + +module { + func.func @kt03_logical_order() { + %c0 = arith.constant 0 : index + %c32 = arith.constant 32 : index + %base = arith.constant 0 : index + // Past the end of the producer's own piece (1*4096*32*64 = 8388608 elements), so + // the assembly is not written over its own source. + %dest = arith.constant 8388608 : index + + %own = ktdp.construct_memory_view %base, + sizes: [1, 4096, 32, 64], strides: [8388608, 2048, 64, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<1x4096x32x64xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident4 + } : memref<1x4096x32x64xf16> -> !ktdp.access_tile<1x4096x32x64xindex> + + %piece_val = ktdp.load %own_access + : !ktdp.access_tile<1x4096x32x64xindex> -> tensor<1x4096x32x64xf16> + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producers + -> !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %piece_val : tensor<1x4096x32x64xf16> + } + + // Assembles [upper, lower] — ascending tile id, per §3.3. Nothing on this op + // can say otherwise. + %region_val = ktdp.inter_tile_gather(%future) + consumer_tiles_per_group = #consumers, + gather_dimensions = [2] + : !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + -> tensor<1x4096x64x64xf16> + + %out = ktdp.construct_memory_view %dest, + sizes: [1, 4096, 64, 64], strides: [16777216, 4096, 64, 1] { + coordinate_set = #region, + memory_space = #ktdp.memory_space + } : memref<1x4096x64x64xf16> + + // Upper half, tensor rows 0..31, to memory rows 32..63. Note the base: + // %c32. The order inside is identity; nothing is permuted within a + // dimension. + %upper = tensor.extract_slice %region_val[0, 0, 0, 0] [1, 4096, 32, 64] + [1, 1, 1, 1] + : tensor<1x4096x64x64xf16> to tensor<1x4096x32x64xf16> + %upper_at = ktdp.construct_access_tile %out[%c0, %c0, %c32, %c0] { + access_tile_set = #piece, access_tile_order = #ident4 + } : memref<1x4096x64x64xf16> -> !ktdp.access_tile<1x4096x32x64xindex> + ktdp.store %upper, %upper_at + : tensor<1x4096x32x64xf16>, !ktdp.access_tile<1x4096x32x64xindex> + + // Lower half, tensor rows 32..63, to memory rows 0..31. Base %c0. + %lower = tensor.extract_slice %region_val[0, 0, 32, 0] [1, 4096, 32, 64] + [1, 1, 1, 1] + : tensor<1x4096x64x64xf16> to tensor<1x4096x32x64xf16> + %lower_at = ktdp.construct_access_tile %out[%c0, %c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident4 + } : memref<1x4096x64x64xf16> -> !ktdp.access_tile<1x4096x32x64xindex> + ktdp.store %lower, %lower_at + : tensor<1x4096x32x64xf16>, !ktdp.access_tile<1x4096x32x64xindex> + + // Expected in memory at cores 4g..4g+3: mb 64g..64g+63 in logical order — + // core 4g+2's half at rows 0..31 and core 4g's at 32..63. The failure to + // catch is storing the assembly verbatim: same shape, same element count, + // wrong answer. That is coverage §5's "reverse two fragments" check. + return + } +} diff --git a/docs/inter-tile-examples/kt04-multi-axis-assembly.mlir b/docs/inter-tile-examples/kt04-multi-axis-assembly.mlir new file mode 100644 index 0000000..ed678d6 --- /dev/null +++ b/docs/inter-tile-examples/kt04-multi-axis-assembly.mlir @@ -0,0 +1,310 @@ +// KT-04 — assembly along more than one axis. +// +// torch-spyre#4300 lx_relayout_workload_coverage.md §5 KT-04: "Four 1x1 tiles of a +// 2x2 logical matrix: source0=(0,0), source1=(1,0), source2=(0,1), source3=(1,1). +// Receiver needs `[[0,1],[2,3]]` when values encode row-major coordinates. **Show +// multi-axis assembly, including exact result dimensions and placement. A flat +// concatenation `[0,2,1,3]` is wrong. Add the full multi-axis Granite shapes +// afterward.**" +// +// No new op: this is `inter_tile_gather` with **two** entries in +// `gather_dimensions` instead of one. §9.1 row 4, the same row as kt01 and kt02. +// What multi-axis adds is that §4's flattening becomes load-bearing, and the two +// functions below disagree about whether that is enough. +// +// OPERATION SEQUENCE. +// +// @kt04_measured LX-Load (addr_0) -> Produce -> Gather -> LX-Store (addr_1) +// @kt04_toy_2x2 LX-Load (addr_0) -> Produce -> Gather +// +// addr_0 = own piece, addr_1 = the assembled region. One Gather with two entries +// in `gather_dimensions` — the sequence is kt02's, and only the attribute differs. +// The toy leaves its result unstored; its point is what the Gather assembles. +// +// ============================================================================ +// WHAT §4 FIXES, NORMATIVELY. +// +// "A list of length n > 1 denotes the product space of those axes, linearized as +// a row-major (mixed-radix odometer) order over the listed extents: **the first +// entry is the slowest-varying and the last is the fastest-varying**." +// +// "**The list is in ascending numerical order (R9).** Entries must ascend ... It +// removes a silent-miscompile class: `[2, 0]` and `[0, 2]` are both 'valid, +// distinct, non-empty' and would flatten to *different* data orders, so a +// reversed list passes every other check while meaning something else." +// +// R12 adds the per-axis obligation: the result flattened extent over the concat set +// is `P * E(D_concat)`, and that "requires every assembled producer to contribute +// the same extent along *each* listed axis — equal products alone would not give a +// well-defined multi-axis assembly, since the flattening of §4 depends on the +// individual extents". "The same validity conditions as R9 apply to the list", so +// the ascending rule covers `gather_dimensions` and not only the split ops. +// +// Together those close the two *spelling* errors this case invites: +// +// gather_dimensions = [1, 0] descending, i.e. a column-major assembly +// -> rejected by R9's ascending rule +// gather_dimensions = [1] one axis multiplies where two must +// -> rejected by R12: 2 * E != the declared extent +// +// They do **not** close a third case, which is what the toy below is. +// ============================================================================ + +// --------------------------------------------------------------------------- +// MEASURED FIXTURE: GR-PF-121, `relayouts[120]` of the pinned catalog +// (sendnn_sdsc_lx_replay_manifest.json, sha256 4c8aca2e…c747d4). +// +// consumer cat_1-kvCacheScatter, input 0 +// route_class all_gather +// extents {mb: 8, out: 128, x: 1, y: 1} src == dst +// word_length 2 (fp16) +// 16 source pieces {mb: 1, out: 64, x: 1, y: 1}, one owner each +// 1 destination piece {mb: 8, out: 128, x: 1, y: 1}, owned by core 0 +// source_fragments_per_destination_piece = 16 -> P = 16 +// destination_pieces_per_source_piece = 1 +// +// Two axes grow, so both are concat axes and neither is split: +// +// mb 1 -> 8 x8 +// out 64 -> 128 x2 8 * 2 = 16 = P +// +// GROUP STRUCTURE. One group — there is one destination region. Producers are the +// **even** cores, consumer is core 0: +// +// group g producers {2k, k=0..15} consumers {0} idle +// g=0 {0, 2, 4, ..., 30} {0} odd cores 1, 3, ..., 31 +// ... 1 group only (g = 0) +// +// Core 0 is both a producer and the consumer, so R13 is satisfied here and this +// case does not depend on §10.1 — unlike kt02, which is the receiving-only case. +// The odd cores are neither, which is §10.2's "Replication versus idleness": with +// `prod(Nd(a)) = 1` against 32 cores the §9.1 guard sends this to row 4 (`gather`) +// rather than row 3, and the consumer set is the destination region's holders, `{0}`. +// +// THE OWNER ORDER AGREES WITH §4's ODOMETER, which is the point of using this +// record. Sorting the 16 producers by ascending tile id and taking `l` as the +// position (§3.3): +// +// l owner holds odometer (mb = l/2, out = 64*(l mod 2)) +// 0 0 mb=0 out= 0 (0, 0) +// 1 2 mb=0 out= 64 (0, 64) +// 2 4 mb=1 out= 0 (1, 0) +// 3 6 mb=1 out= 64 (1, 64) +// ... ... ... ... +// 14 28 mb=7 out= 0 (7, 0) +// 15 30 mb=7 out= 64 (7, 64) +// +// Checked for all 16: `l -> (mb = l floordiv 2, out = 64 * (l mod 2))` matches the +// ascending-owner order exactly, reading the `owners` field of each piece. So `mb` +// is the slow axis and `out` the fast one, which is what `gather_dimensions = [0, 1]` +// means under §4 — ascending, first entry slowest. **The measured data confirms +// §4's flattening rather than contradicting it**, so this function needs nothing +// beyond a plain two-axis gather. +// +// LAYOUT. As in kt06, the piece is written in the record's logical axes; `x` and +// `y` are extent-1 and carry no data, and the catalog cannot settle the physical +// form because its byte fields are logical (see kt06's layout note). Nothing here +// depends on it: both concat axes are data axes, and §10.3's physicalization +// inserts its chunk axis at the front of source and destination alike. +// --------------------------------------------------------------------------- + +#producers = affine_set<(i)[g] : (i >= 0, -i + 31 >= 0, i mod 2 == 0)> +#consumer = affine_set<(i)[g] : (i == 0)> +#groups = affine_set<(g) : (g == 0)> + +// One producer's piece, (mb, out) = (1, 64). +#piece = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 63 >= 0)> +#region = affine_set<(d0, d1) : ( + d0 >= 0, -d0 + 7 >= 0, d1 >= 0, -d1 + 127 >= 0)> +#ident2 = affine_map<(d0, d1) -> (d0, d1)> + +module { + func.func @kt04_measured() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + %base_out = arith.constant 1024 : index + + // The producer's own piece. `ct_local` with no ct_id is the executing tile's + // own LX, so no tile id enters the address. + %own = ktdp.construct_memory_view %base, + sizes: [1, 64], strides: [64, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<1x64xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident2 + } : memref<1x64xf16> -> !ktdp.access_tile<1x64xindex> + + %piece_val = ktdp.load %own_access + : !ktdp.access_tile<1x64xindex> -> tensor<1x64xf16> + + // The load sits outside the region, as in kt01 and kt02: every producer runs + // it and the single consumer is itself a producer, so there is no + // non-producing tile that must be kept out of it (§2.2). + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producers + -> !ktdp.tile_future<(tensor<1x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %piece_val : tensor<1x64xf16> + } + + // combine = none, placement = concat, consumers = one (§1.2's `gather`). + // TWO axes, ascending: axis 0 (`mb`) is the slow one, axis 1 (`out`) the fast + // one, per §4. R12: P * E(D) = 16 * (1*64) = 1024 = 8 * 128, and per axis + // mb x8 and out x2 with 8*2 = 16 = P. + %region_val = ktdp.inter_tile_gather(%future) + consumer_tiles_per_group = #consumer, + gather_dimensions = [0, 1] + : !ktdp.tile_future<(tensor<1x64xf16>), groups = #groups> + -> tensor<8x128xf16> + + // Base 1024, not 0: the producer's own piece occupies elements 0..63, and the + // assembled region must not be written over its own source. + %out = ktdp.construct_memory_view %base_out, + sizes: [8, 128], strides: [128, 1] { + coordinate_set = #region, + memory_space = #ktdp.memory_space + } : memref<8x128xf16> + + %out_at = ktdp.construct_access_tile %out[%c0, %c0] { + access_tile_set = #region, access_tile_order = #ident2 + } : memref<8x128xf16> -> !ktdp.access_tile<8x128xindex> + + ktdp.store %region_val, %out_at + : tensor<8x128xf16>, !ktdp.access_tile<8x128xindex> + + // Expected at core 0: the whole {mb: 8, out: 128} region, with core `2*(2m + n)` + // supplying `mb = m, out = 64n..64n+63`. Undefined at every other tile (§3.7). + // + // Values: identify each element by its (mb, out) coordinate one bit at a time, + // per coverage §1's rule for low-precision tensors, so a misplaced piece names + // the core it came from. + // + // Failures to catch, and which rule reaches each. (1) `gather_dimensions = [1]` + // — one axis where two are needed: R12, since 2 * 64 = 128 accounts for `out` + // but leaves `mb` at 1 against the declared 8. (2) `gather_dimensions = [1, 0]` + // — R9's ascending rule, which exists precisely so this cannot be spelled. + // (3) Swapping which core supplies which cell while keeping the piece count: + // **no rule reaches it**, and it is the toy below. + return + } +} + +// --------------------------------------------------------------------------- +// THE 2x2 TOY — SYNTHETIC, AND IT DOES NOT REDUCE TO A PLAIN GATHER. +// +// The requirement states the toy's ownership explicitly, and reading it as (row, +// column) — the reading under which its own wrong answer `[0,2,1,3]` is the one a +// gather actually produces — the coordinates are a transpose: +// +// group g producers {0, 1, 2, 3} consumers {0} assembled +// g=0 {0, 1, 2, 3} {0} 2x2, l -> (l/2, l mod 2) +// ... 1 group only (g = 0) +// +// producer holds cell row-major value at that cell +// source0 (0, 0) 0 +// source1 (1, 0) 2 +// source2 (0, 1) 1 +// source3 (1, 1) 3 +// +// The receiver needs `[[0,1],[2,3]]`, i.e. value 1 at cell (0,1) and value 2 at +// cell (1,0). +// +// Now apply §4 with `gather_dimensions = [0, 1]` over factors (2, 2). Ascending +// tile order gives `l = 0,1,2,3` for source0..3, and the odometer sends +// `l -> (r = l floordiv 2, c = l mod 2)`: +// +// l producer value odometer cell required cell +// 0 source0 0 (0, 0) (0, 0) ok +// 1 source1 2 (0, 1) (1, 0) WRONG +// 2 source2 1 (1, 0) (0, 1) WRONG +// 3 source3 3 (1, 1) (1, 1) ok +// +// The assembled result is `[[0,2],[1,3]]` — and read out in flat order that is +// `0,2,1,3`, the sequence the requirement names as wrong. So "a flat concatenation +// `[0,2,1,3]` is wrong" has a second reading beyond the obvious one: it is not only +// that the result must be 2x2 rather than 1x4 (R12 gives that), it is that **the +// producers' tile order and the required cell order are a transpose of each other**, +// and the odometer cannot know. +// +// WHY NO RULE CATCHES IT. Compare with the two spellings above. Here +// `gather_dimensions = [0, 1]` is ascending, both axes are listed, every producer +// contributes the same 1x1 extent, `P = 4 = 2 * 2`, and the result type is +// `tensor<2x2xf16>`. R9, R12, R5, R6 and R8 are all satisfied — by the correct +// assembly and by the transposed one alike, because they constrain extents and +// cardinalities and never the identity of what lands where. Same shape, same +// element count, wrong answer. +// +// This is kt03's finding in its multi-axis form, and it should be read together +// with that file: an assembly's order comes from §3.3's ascending tile id, no +// attribute overrides it, and where the required order disagrees the fix is not in +// the delivery op. Three routes, all outside `gather`: +// +// (a) Assign tile ids upstream so that ascending order matches the odometer. +// This is the real answer for a *planner*, and it is why the measured record +// above has no problem: the work division already agrees with §4. +// (b) Store in two or four pieces at swapped bases, as kt03 does — correct, +// zero-copy, and invisible to the verifier. +// (c) Permute the assembled `tensor` locally before use, for a live intermediate. +// +// So KT-04's verdict splits. The two *spelling* errors are closed by R9 and R12, +// which is a stronger position than kt03's. The *placement* error is not closed at +// all, and the toy is the fixture that shows it. Measurement does not force the +// problem — `relayouts[120]` agrees with the odometer — exactly as kt03's 130 of 130 +// do not contradict §3.3. +// +// NOT VERIFIED as a whole: `ktdp.inter_tile_gather` is specified (§6.4) but absent +// from KTDP.td, in both functions. The `ktdp.inter_tile_produce` halves parse and +// round-trip, as do all the affine sets on their own. +// --------------------------------------------------------------------------- + +#four = affine_set<(i)[g] : (i >= 0, -i + 3 >= 0)> +#cell = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 >= 0)> +#matrix = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 1 >= 0)> + +module { + func.func @kt04_toy_2x2() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + + // One 1x1 cell per producer. Which logical cell it is comes from the work + // division, not from anything in this function — which is the whole problem. + %own = ktdp.construct_memory_view %base, + sizes: [1, 1], strides: [1, 1] { + coordinate_set = #cell, + memory_space = #ktdp.memory_space + } : memref<1x1xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0] { + access_tile_set = #cell, access_tile_order = #ident2 + } : memref<1x1xf16> -> !ktdp.access_tile<1x1xindex> + + %cell_val = ktdp.load %own_access + : !ktdp.access_tile<1x1xindex> -> tensor<1x1xf16> + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #four + -> !ktdp.tile_future<(tensor<1x1xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %cell_val : tensor<1x1xf16> + } + + // Well-formed and still wrong: ascending list, both axes, uniform extents, + // P = 4 = 2 * 2, result type exactly as §4 requires. It assembles + // [[0,2],[1,3]] from the toy's ownership, where [[0,1],[2,3]] is required. + %matrix_val = ktdp.inter_tile_gather(%future) + consumer_tiles_per_group = #consumer, + gather_dimensions = [0, 1] + : !ktdp.tile_future<(tensor<1x1xf16>), groups = #groups> + -> tensor<2x2xf16> + + // Expected, if the toy's ownership is taken as given: [[0,1],[2,3]]. + // Produced by this op: [[0,2],[1,3]]. The difference is a transpose and no + // rule in §5 sees it — see the discussion above for the three places the fix + // can live, none of which is this op. + return + } +} diff --git a/docs/inter-tile-examples/kt05-select-then-scatter.mlir b/docs/inter-tile-examples/kt05-select-then-scatter.mlir new file mode 100644 index 0000000..6328de1 --- /dev/null +++ b/docs/inter-tile-examples/kt05-select-then-scatter.mlir @@ -0,0 +1,295 @@ +// KT-05 — select a row locally, then split it across the receivers. +// +// torch-spyre#4300 lx_relayout_workload_coverage.md §5 KT-05: "Proposal's +// `[512,32,64]` select: source core `4m+h` owns rows `64m:64m+64`, sticks +// `8h:8h+8`. Destination core `d` needs row511, stick `d`. ... Source +// `28+floor(d/8)` supplies destination `d`. Select locally, then split/send the +// selected row. The original layout is not a no-op; fresh HBM loads at +// destinations would change the starting contract." And §5's closing note: +// "Test the proposal's 32-stick example and Granite P14's 4096-column example +// separately; they have the same relation but different widths." +// +// Both are here: @kt05_granite_p14 is the measured one, @kt05_proposal_32_sticks +// the narrower one the proposal states. +// +// OPERATION SEQUENCE. Both functions: +// +// Produce[ LX-Load (addr_0) ] -> Scatter +// +// addr_0 = the producer's own slab. **The load is inside the produce region**, +// unlike every other file here, because only one tile per group produces and the +// group's eight consumers must not run it (§2.2). The select is the access tile +// that load uses, so it needs no op of its own. +// +// ============================================================================ +// MEASURED FIXTURE: GR-PF-055, `relayouts[54]` of the pinned catalog +// (sendnn_sdsc_lx_replay_manifest.json, sha256 4c8aca2e…c747d4). +// +// consumer slice_161-Stcdp, input 0 (family P14) +// tensor mean_80-LayerNormNorm_out +// extents {mb: 512, out: 4096, y: 1} src == dst +// word_length 2 (fp16) +// remote_required true +// 32 source pieces {mb:64, out:1024, y:1}, one owner each +// 32 destination pieces {mb:1, out:128, y:1}, one owner each +// +// Source piece k is owned by core k: rows 64*(k/4) upward, out columns 1024*(k%4) +// upward. So the 32 source pieces form an 8 x 4 grid — eight row blocks of 64 by +// four out-slabs of 1024. Every destination piece is a single row, **mb = 511**, +// and 128 out columns starting at out = 128*d, owned by core d. +// +// WHICH CORES SEND, AND TO WHOM. Two observations settle it. +// +// 1. Every destination wants row 511, and row 511 lies in the **last row block**, +// rows 448..511. Only four cores hold that block: 28, 29, 30, 31. The other +// 28 send nothing, which is what the catalog records as +// `destination_pieces_per_source_piece = 0` for pieces 0..27. +// +// 2. Those four hold row 511 in four consecutive out-slabs of 1024: +// +// core 28 out 0..1023 core 30 out 2048..3071 +// core 29 out 1024..2047 core 31 out 3072..4095 +// +// A destination is 128 columns wide, so **each slab is 1024/128 = 8 +// destinations wide**, and those destinations are consecutive in d: +// +// core 28 -> d = 0.. 7 core 30 -> d = 16..23 +// core 29 -> d = 8..15 core 31 -> d = 24..31 +// +// That is the whole structure: one producer per slab, eight consumers each. +// Coverage §5 states it as "source 28 + floor(d/8) supplies destination d". +// +// Group structure. A group is one slab being distributed, so group g is core 28+g +// serving the eight consumers {8g .. 8g+7}. Checked against the record: d = 0..7 +// come from core 28, d = 8..15 from core 29, d = 16..23 from core 30, d = 24..31 +// from core 31. +// +// The 8 in `8g+7` is also R9's `C` below — the slab's fan-out and the split count +// are one number, because the slab is what gets split. +// +// group g producers {28+g} consumers {8g .. 8g+7} out columns of row 511 +// g=0 {28} {0, 1, ..., 7} 0..1023 -> 128 each +// g=1 {29} {8, 9, ..., 15} 1024..2047 +// g=2 {30} {16, ..., 23} 2048..3071 +// g=3 {31} {24, ..., 31} 3072..4095 +// +// WHY THIS IS NOT A NO-OP. Core 0 needs out 0..127 of row 511, and row 511 is +// owned by cores 28..31 — core 0 owns rows 0..63. Coverage §5 makes the point +// with one coordinate: "Core0's required `(511,0,:)` starts on **source28**. +// That one coordinate is enough to disprove a no-communication interpretation." +// The record's `remote_required` is true. +// +// WHY A SELECT COMES FIRST — AND WHY THE GUARD THEN STOPS FIRING. +// +// Read as one relayout of the whole tensor — 32 owner slabs on one side, 32 +// slivers of row 511 on the other — the pair fails §9.1's first guard row. Not +// on the counts: `prod(Ns)` and `prod(Nd)` both equal their region counts. It is +// the *coverage* clause that catches it, since the 32 distinct destination +// regions hold 32 * 128 = 4096 elements against the tensor's +// 512 * 4096 = 2097152 — a 512th of it. §9.3 records this as the one +// "selection, not a partition" file. +// +// **After the select the guard is satisfied**, and §9.1 says so itself: the +// coverage clause is measured against "the value delivered ... after a select +// that is the selected sub-tensor and not the original — otherwise the +// select-then-deliver that repairs a selection would trip the guard it was meant +// to satisfy." The value delivered here is the selected row, mb = 1 by +// out = 4096 — which is what `yield_partial` hands over below: +// +// source regions 4 shards of out 1024 4 * 1024 = 4096 covers the row +// destinations 32 chunks of out 128 32 * 128 = 4096 covers the row +// Ns = {out: 4}, Nd = {out: 32} -> C = {}, R = {out} +// +// That is §9.1 row 2, `inter_tile_scatter` with `scatter_dimensions = R`. So the +// guard's job is to reject the pre-select reading and force the select, not to +// say the record is inexpressible: the post-select pair classifies normally and +// yields exactly the op below, with `[1]` naming `out`. +// +// The selection itself needs no new op — an access tile over the selected +// sub-rectangle expresses it, which is what the produce region does. +// +// The catalog labels the record `route_class = permutation`. That describes the +// pre-select pair (32 pieces to 32 pieces, one fragment each) and is not a KTIR +// op choice; how the label is computed is not in the JSON. The op to emit comes +// from §9.1 applied to the post-select pair, as above. +// +// LAYOUT. No assumption is needed for this record — unlike kt02. A real SDSC +// run of this very logical shape (512 x 4096 x 1, fp16) reports +// +// layoutDimOrder_ = ["mb", "out", "y"] stickDimOrder_ = ["y"] +// stickSize_ = [64] device_size = [1, 4096, 512, 64] +// +// and this record's axes are literally `mb`, `out`, `y`, so the measured order +// applies directly. `y` is innermost and sticked, an extent-1 `y` is carried as +// 64 padded lanes, and the physical form is [1, , , 64]: +// +// producer's own piece logical (mb=64, out=1024, y=1) -> [1, 1024, 64, 64] +// after the select logical (mb=1, out=1024, y=1) -> [1, 1024, 1, 64] +// per consumer logical (mb=1, out=128, y=1) -> [1, 128, 1, 64] +// +// The split axis is `out`, physical dim 1 — not the stick axis, so a single +// index and no floordiv rule. R9: E = 1024, C = 8, 1024 % 8 == 0. In sticks +// that is 16 % 8 == 0, two sticks per consumer, so the split stays on stick +// boundaries (§4's stick-multiple reading of R9). +// +// R8 holds with one producer per group, so `scatter` takes no +// `producer_dependency_per_consumer` (§6.6) and full-barrier and per-tile +// synchronization coincide. R13 is `n` for `scatter` (§6.6, §10.1) and it is +// needed: in group g=0 the producer is core 28 and the consumers are 0..7, which +// are disjoint from it. Ordering is by consumer local index `l` (§3.3): +// consumer 8g+j takes chunk j. +// +// NOT VERIFIED as a whole: `ktdp.inter_tile_scatter` is specified (§6.6) but +// absent from KTDP.td. The `ktdp.inter_tile_produce` half parses, including the +// in-region select. +// ============================================================================ + +// Producer of group g: core 28+g. +#producer = affine_set<(i)[g] : (i - g - 28 == 0)> +// Consumers of group g: 8g .. 8g+7. +#consumers = affine_set<(i)[g] : (i - 8 * g >= 0, -i + 8 * g + 7 >= 0)> +#groups = affine_set<(g) : (g >= 0, -g + 3 >= 0)> + +// The producer's whole slab, and the single row selected out of it. Both sets +// are relative to the access tile's base coordinate. +#slab = affine_set<(d0, d1, d2, d3) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 1023 >= 0, + d2 >= 0, -d2 + 63 >= 0, d3 >= 0, -d3 + 63 >= 0)> +#last_row = affine_set<(d0, d1, d2, d3) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 1023 >= 0, + d2 >= 0, -d2 >= 0, d3 >= 0, -d3 + 63 >= 0)> +#ident4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> + +module { + func.func @kt05_granite_p14() { + %c0 = arith.constant 0 : index + %c63 = arith.constant 63 : index + %base = arith.constant 0 : index + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producer + -> !ktdp.tile_future<(tensor<1x1024x1x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + // The loads belong INSIDE the region: with one producer per group they + // must not run on the group's eight non-producing consumers (§2.2, which + // names this as the case where a richer body is normally required). + // + // `ct_local` with no ct_id is the executing tile's own LX. Dense: + // lane 1, mb 64, out 64*64 = 4096, and the extent-1 axis takes the rest. + %own = ktdp.construct_memory_view %base, + sizes: [1, 1024, 64, 64], strides: [4194304, 4096, 64, 1] { + coordinate_set = #slab, + memory_space = #ktdp.memory_space + } : memref<1x1024x64x64xf16> + + // THE SELECT. Global row 511 is local row 63 of this producer's slab + // (it owns rows 448..511), so the access tile is anchored at mb = 63 with + // extent 1 there and full extent elsewhere. No new op: an access tile + // over a sub-rectangle is the selection §9.3 asks for. + %sel = ktdp.construct_access_tile %own[%c0, %c0, %c63, %c0] { + access_tile_set = #last_row, access_tile_order = #ident4 + } : memref<1x1024x64x64xf16> -> !ktdp.access_tile<1x1024x1x64xindex> + + %row = ktdp.load %sel + : !ktdp.access_tile<1x1024x1x64xindex> -> tensor<1x1024x1x64xf16> + + ktdp.yield_partial %row : tensor<1x1024x1x64xf16> + } + + // combine = none, placement = split. The selected row is cut along `out` + // into C = 8 chunks of 128, one per consumer in ascending local index (§3.3). + // No combiner, no identity, and no dependency attribute (§6.6). + %chunk = ktdp.inter_tile_scatter(%future) + consumer_tiles_per_group = #consumers, + scatter_dimensions = [1] + : !ktdp.tile_future<(tensor<1x1024x1x64xf16>), groups = #groups> + -> tensor<1x128x1x64xf16> + + // Expected at core d: row 511, out columns 128d..128d+127, from source + // 28 + d/8. The consumers of a group hold **disjoint** ordered slices that + // tile the selected row (§3.7, `split`) — unlike the gather cases, where all + // consumers hold the same thing. Undefined at no tile: every core consumes. + // + // Failure to catch: give consumer 8g+j chunk j+1, or read the selection from + // HBM at each destination. The second is the one coverage §5 warns about — + // it would produce the right values while changing the starting ownership, + // so the emitted memory accesses have to be checked as well. + return + } +} + +// --------------------------------------------------------------------------- +// The proposal's narrower example, stated in coverage §5 as `[512, 32, 64]`: +// 512 rows, 32 sticks, 64 lanes. Same core relationship, different width. +// +// Same structure as above, one slab narrower. Row 511 is again in the last row +// block, so again only cores 28..31 send; they hold it in four slabs of 8 sticks, +// and a destination is one stick, so each slab is 8 destinations wide: +// +// core 28 sticks 0.. 7 -> d = 0.. 7 core 30 sticks 16..23 -> d = 16..23 +// core 29 sticks 8..15 -> d = 8..15 core 31 sticks 24..31 -> d = 24..31 +// +// group g producers {28+g} consumers {8g .. 8g+7} sticks of row 511 +// g=0 {28} {0, ..., 7} 0..7 -> 1 each +// g=1 {29} {8, ..., 15} 8..15 +// g=2 {30} {16, ..., 23} 16..23 +// g=3 {31} {24, ..., 31} 24..31 +// +// WHY IT IS WORTH TESTING SEPARATELY. The width is what differs, and it lands +// R9 on its boundary: 8 sticks over 8 consumers is **one stick each**, where the +// Granite record gives two. A split that went one step further would be +// sub-stick and R9 would have to reject it (§4: "a split that would drive the +// result sub-stick fails R9 rather than needing a rule of its own"). So this +// function is the tight case and the measured one is the slack case. +// +// Shapes are the proposal's own rank-3 spelling (rows, sticks, lanes) rather +// than the SDSC physical form used above; the two are not mixed in one function. +// --------------------------------------------------------------------------- + +#slab3 = affine_set<(d0, d1, d2) : ( + d0 >= 0, -d0 + 63 >= 0, d1 >= 0, -d1 + 7 >= 0, d2 >= 0, -d2 + 63 >= 0)> +#last_row3 = affine_set<(d0, d1, d2) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 7 >= 0, d2 >= 0, -d2 + 63 >= 0)> +#ident3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + +module { + func.func @kt05_proposal_32_sticks() { + %c0 = arith.constant 0 : index + %c63 = arith.constant 63 : index + %base = arith.constant 0 : index + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producer + -> !ktdp.tile_future<(tensor<1x8x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + %own = ktdp.construct_memory_view %base, + sizes: [64, 8, 64], strides: [512, 64, 1] { + coordinate_set = #slab3, + memory_space = #ktdp.memory_space + } : memref<64x8x64xf16> + + // Global row 511 is local row 63 again. + %sel = ktdp.construct_access_tile %own[%c63, %c0, %c0] { + access_tile_set = #last_row3, access_tile_order = #ident3 + } : memref<64x8x64xf16> -> !ktdp.access_tile<1x8x64xindex> + + %row = ktdp.load %sel + : !ktdp.access_tile<1x8x64xindex> -> tensor<1x8x64xf16> + + ktdp.yield_partial %row : tensor<1x8x64xf16> + } + + // Split the stick axis 8 ways: E = 8, C = 8, one stick per consumer. This + // is R9 at its boundary. + %chunk = ktdp.inter_tile_scatter(%future) + consumer_tiles_per_group = #consumers, + scatter_dimensions = [1] + : !ktdp.tile_future<(tensor<1x8x64xf16>), groups = #groups> + -> tensor<1x1x64xf16> + + // Expected at core d: row 511, stick d, from source 28 + d/8. + return + } +} diff --git a/docs/inter-tile-examples/kt06-owner-permutation.mlir b/docs/inter-tile-examples/kt06-owner-permutation.mlir new file mode 100644 index 0000000..356fb74 --- /dev/null +++ b/docs/inter-tile-examples/kt06-owner-permutation.mlir @@ -0,0 +1,410 @@ +// KT-06 — equal split counts, permuted core ownership. +// +// torch-spyre#4300 lx_relayout_workload_coverage.md §5 KT-06: "Gemma row-major +// versus column-major core order, described above. **Equal split counts are +// insufficient. Preserve or explicitly change ownership.**" +// +// And from the coverage document's Gemma discussion: "Both variants are required: +// if the next operation adopts the source order, zero-copy is correct; reading the +// old addresses with the new interpretation is not." +// +// Both variants are here. They are not two spellings of one transfer — they are +// the two branches §9.1 row 1 already names, and which branch applies is a +// property of the *consumer*, not of this edge. +// +// OPERATION SEQUENCE. +// +// @kt06_preserve_ownership LX-Load (addr_0) +// @kt06_relocate LX-Load (addr_0) -> Produce -> Consume +// -> LX-Store (addr_0) +// +// Variant A is one op long, and **the absence of a delivery is its claim** — +// §9.1 row 1's "no op needed" branch. Variant B loads and stores at the *same* +// address: the permutation is entirely in the dependency attribute, and no local +// address moves. +// +// ============================================================================ +// MEASURED FIXTURE: GR-PF-052, `relayouts[51]` of the pinned catalog +// (sendnn_sdsc_lx_replay_manifest.json, sha256 4c8aca2e…c747d4). +// +// consumer bmm-wtAttnHeadBreak-VirtualReshape-Output-Restickify, input 0 +// route_class permutation +// extents {j: 8, mb: 512, out: 128, x: 1, y: 1} src == dst +// word_length 2 (fp16) +// 32 source pieces {j: 2, mb: 64, out: 128, x: 1, y: 1}, one owner each +// 32 destination pieces, **the same size**, one owner each +// source_fragments_per_destination_piece = 1 -> P = 1 +// destination_pieces_per_source_piece = 1 -> a bijection +// +// Two sibling records have the identical geometry and owner map: `relayouts[52]` +// (`bmm_2-…`, fold_factor 38, so the folded layer body) and `relayouts[53]` +// (`bmm_78-…`). `relayouts[118]` and `[119]` are the decode counterparts. So +// this edge occurs once per layer, 40 layers, and one KTIR form covers all five +// records. +// +// THE OWNER MAP. Let J = j/2 in [0,4) and M = mb/64 in [0,8), so a region is +// (J, M) and there are 4*8 = 32 of them. Reading the `owners` field of every +// piece on both sides: +// +// source owner k = 8*J + M (M varies fastest — row-major in (J,M)) +// destination owner m = J + 4*M (J varies fastest — column-major) +// +// Verified against all 32 pieces on both sides: `k == 8*(j/2) + (mb/64)` holds for +// every source piece and `m == (j/2) + 4*(mb/64)` for every destination piece. +// +// That is a transpose of a 4x8 index grid, which is exactly what coverage §5 calls +// "row-major versus column-major core order". So +// +// pi(k) = (k floordiv 8) + 4 * (k mod 8) +// pi^-1(m) = 8 * (m mod 4) + (m floordiv 4) +// +// pi = [0, 4, 8,12,16,20,24,28, 1, 5, 9,13,17,21,25,29, +// 2, 6,10,14,18,22,26,30, 3, 7,11,15,19,23,27,31] +// +// It is a bijection, and **only owners 0 and 31 are fixed points**. That matters +// for the test: an implementation that emits nothing at all is wrong on 30 of the +// 32 cores, so the identity is not a plausible near-miss here. +// +// Group structure, derived from the owner tables. The transpose decomposes into +// **8 groups of 4 cores**. Writing a consumer as `c = 4g + l` with `g = c/4` and +// `l = c mod 4`, its source is +// +// pi^-1(4g + l) = 8*l + g +// +// so group `g` has consumers `{4g, 4g+1, 4g+2, 4g+3}` — four *consecutive* cores — +// drawing from producers `{g, g+8, g+16, g+24}` — four cores **spaced 8 apart**, +// i.e. exactly the cores congruent to `g` mod 8: +// +// group g producers {g, g+8, g+16, g+24} consumers {4g .. 4g+3} +// g=0 {0, 8, 16, 24} {0, 1, 2, 3} +// g=1 {1, 9, 17, 25} {4, 5, 6, 7} +// g=2 {2, 10, 18, 26} {8, 9, 10, 11} +// ... 8 groups, regular +// g=7 {7, 15, 23, 31} {28, 29, 30, 31} +// +// The pairing inside a group is `4g + l <- g + 8l` for l = 0..3, so group 0 is +// 0<-0, 1<-8, 2<-16, 3<-24 and group 1 is 4<-1, 5<-9, 6<-17, 7<-25. +// +// Checked against the record and not against the formula: consumer `4g + l` +// receives exactly what producer `g + 8*l` holds, for all 32 pairs, with both sides +// read out of the `owners` field of each piece. So `#dep_transpose` below is +// consistent with the measured tables, not merely with the transpose that +// summarises them. +// +// Both families partition `0..31` — the consumer sets because they are consecutive +// blocks, the producer sets because they are the residue classes mod 8 — so R1 holds +// and is not vacuous. +// +// This is the semantically right grouping, and not only a tidier one: see "WHY 8 +// GROUPS OF 4" below, where it turns out to catch strictly more than the +// alternatives. +// +// WHY 8 GROUPS OF 4. Three groupings express this same permutation, and they are +// not equally checkable, because **R1, R4 and R5 are all obligations relative to a +// group.** The grouping decides how much of the map a verifier ever compares. +// +// 32 singleton groups one producer and one consumer each. `pi^-1` is affine, +// so it is expressible — and **worthless**: with one +// consumer per group, R4 (every producer named) and R5 +// (dependency sets disjoint) are vacuously true, and no +// rule ever relates the 32 pairs to each other. +// +// one group of 32 R4 and R5 become surjectivity and injectivity of +// `pi^-1`, so the bijection is checked. But R1 is trivial +// with a single group, and R3 (dep within the producer set) +// is nearly so, since every tile is in it. +// +// 8 groups of 4 R4 and R5 check four 4-element bijections instead of one +// 32-element one — equally binding — and **R1 and R3 start +// working too.** R1 has two partitions of `0..31` to +// verify; R3 rejects any dep that names a producer outside +// `{g, g+8, g+16, g+24}`. +// +// The last is strictly the strongest. Take the error "consumer 0 reads core 1 +// instead of core 0": under 8 groups it fails **R3** immediately, because 1 is not +// congruent to 0 mod 8 and so is not a producer of group 0. Under one group of 32 +// it slips past R3 and is caught only later and less directly, by R5, once some +// other consumer is found to name core 1 as well. Under singleton groups it is not +// caught at all. +// +// What no grouping catches is a permutation *within* a group — swapping which of +// `{0, 8, 16, 24}` feeds consumers 0 and 1 keeps every rule satisfied. So the +// rules narrow the error to "a valid permutation of the right four sources, but not +// this one", and the numerical check closes the rest. +// +// WHY NOT `all_to_all`. A four-into-four exchange between whole pieces invites it, +// and §6.5 rules it out in a sentence: "one-to-one permutation of whole partials is +// already expressible as `consume` + a bijective dependency set (§7.4.2); +// `all_to_all` is only for the split-and-redistribute case, so the two mechanisms +// do not overlap." +// +// The reason is that §9.1 picks the op from two bits — does the edge split, does it +// concatenate. `C = empty` and `R = empty` here, so it is row 1 (`consume`); +// `all_to_all` is row 3, which needs *both* non-empty. Mechanically it also cannot +// be spelled: §6.5 requires both `split_dimensions` and `concat_dimensions`, and the +// result type is `T_p` with the split extent divided by `C` and the concat extent +// multiplied by `P`. This case needs `T_c == T_p`, which forces `C = P = 1` and so +// 32 singleton groups — and then R9 still requires a non-empty split list, so one +// would be declaring a split into one chunk that does not happen, with `placement = +// permute` on a delivery that hands over an unmodified piece. +// +// The distinguishing test: `all_to_all` would be right if each consumer needed **a +// quarter of each of the four sources** instead of **the whole of one**. §6.5 does +// allow what it calls a "pure ownership transpose along one axis set", but that +// transposes *within* the data and cuts every piece; this record cuts nothing. +// +// The near miss is in the same fixture family. GR-PF-001 (`relayouts[0]`) also has +// `inputs per region = 2`, but `destination_pieces_per_source_piece = 4`, so its +// sources *are* cut — `mb` concatenated while `out` is split. That one is a genuine +// `all_to_all`, and it is the KT-04 candidate. In the catalog the two families are +// disjoint by construction: 37 records have neither axis set non-empty, 12 have both. +// +// ---------------------------------------------------------------------------- +// HOW TO READ THE CATALOG FOR THIS RECORD — one trap, worth recording. +// +// `relayouts[*].source_pieces` is ordered lexicographically by its `key` string: +// p0, p1, p10, p11, ..., p19, p2, p20, ... So **piece index is not owner order.** +// And `source_core_patterns` is a separate summary of the form +// `{cores: [...], pieces: N}`; it happens to have 32 entries for this record but +// it is not a per-piece record and carries no piece key. +// +// The authoritative owner is the **`owners` field on each piece**. Deriving the +// map from `source_core_patterns[k]` alongside `source_pieces[k]` instead yields a +// table that is a bijection but not a transpose — an artifact of the lexicographic +// ordering, not a property of the record. The transpose above only appears once +// the `owners` field is used. +// ---------------------------------------------------------------------------- +// +// LAYOUT, AND WHY THIS CASE DOES NOT DEPEND ON IT. +// +// No SDSC run exists for this shape, unlike GR-PF-055 (see kt05). And the catalog +// cannot supply one, because **its byte fields are logical**: for this record +// prod(extents) * word_length == logical_tensor_bytes exactly (1048576), and the +// same identity holds for GR-PF-055, where an SDSC run showed that a `y` extent of +// 1 is physically carried as 64 padded lanes. Padding is therefore invisible in +// the catalog by construction, for every record. +// +// (A second naming trap: `source_piece_bytes` is the total over all pieces, not +// the size of one. Here it is 1048576 = 32 * 32768.) +// +// This case is decidable anyway. The permutation is a statement about **which +// core** holds a region, and §10.3's physicalization inserts the chunk-count axis +// at the front of *both* sides identically — it cannot change an owner map. So +// the piece type below is written in the record's logical axes (`j`, `mb`, `out`, +// with the extent-1 `x` and `y` carrying no data), and this file makes no claim +// about the physical form. It is not mixed with the SDSC physical spelling used +// in kt02 and kt05. **KT-08 is where the layout question actually bites**, and it +// needs the SDSC run before it can be written. +// ============================================================================ +// +// WHY THE PERMUTATION IS VERIFIABLE HERE, UNLIKE KT-03. +// +// kt03 is the same family of problem — a required order disagreeing with the order +// the IR supplies by default — and there the reorder had to go into the *store +// bases*, where no rule reaches it, so only a numerical check catches a mistake. +// +// Here the reorder is **across cores rather than within a region**, so it lives in +// `producer_dependency_per_consumer`, and the rules do reach it: +// +// R1 groups pairwise exclusive over tiles both families partition 0..31 +// R3 dep(c) subset of the group's producers `g + 8*(c mod 4)` is in g mod 8 +// R4 every producer named by some consumer surjective onto the 4 producers +// R5 dependency sets pairwise disjoint injective over the 4 consumers +// R6 uniform cardinality |dep(c)| = 1 for every c +// R8 exactly one source per consumer tile dep is a function of (c, g) +// +// R5 applies here in the strong sense: this is a *partitioning* use — every producer +// is claimed by exactly one consumer of its group — so unlike kt01 the disjointness +// obligation is live and satisfied. A dropped producer fails R4, a doubled one +// fails R5, and a source taken from the wrong group fails R3. All of that depends +// on the grouping; see "WHY 8 GROUPS OF 4" above. +// +// What is *not* caught is a permutation within one group — swapping which of +// `{g, g+8, g+16, g+24}` feeds two of the group's consumers satisfies R1 through R8. +// So the rules confine the error to "a valid permutation of the right four sources, +// but not this one", and the numerical check closes the rest. That is a strictly +// better position than kt03's, and the difference is only where the ordering +// information was allowed to live. +// +// NOT VERIFIED as a whole: `ktdp.inter_tile_consume` is specified (§6.1) but absent +// from KTDP.td. The `ktdp.inter_tile_produce` half parses and round-trips, as do +// all three affine sets on their own — including the `floordiv` / `mod` in +// `#dep_transpose`, checked with ktir-opt. +// ============================================================================ + +// 8 groups of 4. Producers of group `g` are the cores congruent to `g` mod 8, +// `{g, g+8, g+16, g+24}`; consumers are the consecutive block `{4g .. 4g+3}`. Both +// families partition `0..31`, which is what makes R1 bite. +#producers = affine_set<(i)[g] : (i mod 8 - g == 0, i >= 0, -i + 31 >= 0)> +#consumers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 3 >= 0)> +#groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +// producer_dependency_per_consumer: `(p)[c, g]` per §3.4. Consumer `c = 4g + l` +// waits for exactly one producer, `g + 8*l`, i.e. `g + 8*(c mod 4)`. +// +// Both `c` and `g` are required, as in §7.4.2's butterfly: `c` because each of the +// four consumers in a group has a different source, and `g` because the producer is +// offset by the group index. +#dep_transpose = affine_set<(p)[c, g] : (p - g - 8 * (c mod 4) == 0)> + +// The whole of one owner's piece: (j, mb, out) = (2, 64, 128). +#piece = affine_set<(d0, d1, d2) : ( + d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 63 >= 0, d2 >= 0, -d2 + 127 >= 0)> +#ident3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + +// --------------------------------------------------------------------------- +// VARIANT A — the consumer adopts the source order. NO TRANSFER. +// +// §9.1 row 1 states this branch itself: with `C = empty` and `R = empty`, the +// regions are identical, and it is "`no op needed` if every core's region is its +// own". Here every core's region *is* its own — the piece it already holds — so +// the correct emission for this variant is nothing at all. Coverage §5's +// "preserve ... ownership" is this arm. +// +// The whole content of this function is therefore a load and the consumer, and +// **the absence of an inter-tile op is the claim being made.** +// +// THE PRECONDITION, WHICH IS NOT EXPRESSED ANYWHERE IN THE IR. This is only +// correct if the consumer indexes its slab as k = 8*(j/2) + (mb/64) — the source +// map. If it indexes as m = (j/2) + 4*(mb/64), the same addresses now name +// different data on 30 of the 32 cores, and nothing in this function changes: +// same ops, same types, same addresses, wrong answer. +// +// That is the failure coverage §5 warns about in the sentence "reading the old +// addresses with the new interpretation is not [correct]". It is invisible to a +// verifier here for the same structural reason as kt03: the owner map has no +// syntactic home in this variant. Which is the argument for making the consumer's +// expectation explicit somewhere — and, in variant B, it is. +// --------------------------------------------------------------------------- + +module { + func.func @kt06_preserve_ownership() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + + // The tile's own piece. `ct_local` with no ct_id is the executing tile's own + // LX, so no tile id enters the address. Dense: out 1, mb 128, j 64*128. + %own = ktdp.construct_memory_view %base, + sizes: [2, 64, 128], strides: [8192, 128, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<2x64x128xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident3 + } : memref<2x64x128xf16> -> !ktdp.access_tile<2x64x128xindex> + + %piece_val = ktdp.load %own_access + : !ktdp.access_tile<2x64x128xindex> -> tensor<2x64x128xf16> + + // The consumer runs here, on the tile's own data, in the source order. + // Expected at core k: region (J, M) with k = 8*J + M — that is, j in + // [2*(k/8), 2*(k/8)+2) and mb in [64*(k%8), 64*(k%8)+64). Unchanged from + // what the core already held. + return + } +} + +// --------------------------------------------------------------------------- +// VARIANT B — the consumer requires the destination order. A RELOCATION. +// +// Same §9.1 row 1, other branch: regions identical but a core's destination region +// is *not* its own, so the op is `inter_tile_consume` — the row calls this "a +// **relocation**". §6.1 calls the same shape a **routing** pattern rather than a +// broadcast: with several producers per group the dependency attribute names the +// sender, R8 requires it, and the group is "several independent point-to-point +// deliveries sharing one `produce` op, which is what lets `consume` express ... +// one-to-one permutation exchange (§7.4.2)". +// +// Coverage §5's "or explicitly change ownership" is this arm, and `#dep_transpose` +// is where the change is stated. +// +// LOCAL ADDRESSES DO NOT MOVE. Both the load and the store are at base 0 in the +// executing tile's own LX: core k reads its own buffer, core m writes its own +// buffer. What changes is only *which* core's value lands in which core's buffer. +// So the entire permutation is carried by the dependency set and nothing leaks +// into the addressing — the opposite of kt03, and the reason this case is +// verifiable while kt03 is not. +// --------------------------------------------------------------------------- + +module { + func.func @kt06_relocate() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + + %own = ktdp.construct_memory_view %base, + sizes: [2, 64, 128], strides: [8192, 128, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<2x64x128xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident3 + } : memref<2x64x128xf16> -> !ktdp.access_tile<2x64x128xindex> + + %piece_val = ktdp.load %own_access + : !ktdp.access_tile<2x64x128xindex> -> tensor<2x64x128xf16> + + // Four producers per group, `{g, g+8, g+16, g+24}`, and every tile is a producer + // of exactly one group. The load sits outside the region here, as in kt01 and + // kt02: no tile is a non-producer that must be kept from executing it (§2.2). + // kt05 is the case where it must go inside. + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producers + -> !ktdp.tile_future<(tensor<2x64x128xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %piece_val : tensor<2x64x128xf16> + } + + // combine = none, placement = replicate, |P(g)| = 4, |dep(c)| = 1 (§6.1). No + // dim attribute, no region, no identity — the result type is `T_p` unchanged, + // because a relocation moves a piece without reshaping it. With |P(g)| > 1 the + // dependency attribute is *required* (§3.4), and it is the one attribute that + // does the work: §6.1 calls this a routing pattern rather than a broadcast — + // four independent point-to-point deliveries sharing one `produce`. + %relocated = ktdp.inter_tile_consume(%future) + consumer_tiles_per_group = #consumers, + producer_dependency_per_consumer = #dep_transpose + : !ktdp.tile_future<(tensor<2x64x128xf16>), groups = #groups> + -> tensor<2x64x128xf16> + + // Store into the consumer's own buffer, at base 0 — the same address the load + // used. See "LOCAL ADDRESSES DO NOT MOVE" above. + %out = ktdp.construct_memory_view %base, + sizes: [2, 64, 128], strides: [8192, 128, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<2x64x128xf16> + + %out_at = ktdp.construct_access_tile %out[%c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident3 + } : memref<2x64x128xf16> -> !ktdp.access_tile<2x64x128xindex> + + ktdp.store %relocated, %out_at + : tensor<2x64x128xf16>, !ktdp.access_tile<2x64x128xindex> + + // Expected at core m: the region (J, M) with m = J + 4*M — that is, j in + // [2*(m mod 4), 2*(m mod 4)+2) and mb in [64*(m/4), 64*(m/4)+64) — supplied by + // core 8*(m mod 4) + (m floordiv 4). Concretely for the first few: + // + // core 0 <- core 0 j=0..1 mb= 0..63 (fixed point) + // core 1 <- core 8 j=2..3 mb= 0..63 + // core 2 <- core 16 j=4..5 mb= 0..63 + // core 3 <- core 24 j=6..7 mb= 0..63 + // core 4 <- core 1 j=0..1 mb= 64..127 + // core 31 <- core 31 j=6..7 mb=448..511 (fixed point) + // + // Values: identify each element by its (j, mb, out) coordinate one bit at a + // time, per coverage §1's rule for low-precision tensors, so that a + // misdelivered piece names the core it actually came from. + // + // Failures to catch. (1) Emitting nothing — variant A's answer given variant + // B's requirement — is wrong on 30 of 32 cores, since only 0 and 31 are fixed + // points. (2) Using pi instead of pi^-1: the dependency set must map a + // *consumer* to its source, and pi maps the other way. Both are bijections, + // so R3 through R8 accept either and only the values distinguish them. This is + // the "wrong core map with identical split counts" check of coverage §5. + return + } +} diff --git a/docs/inter-tile-examples/kt07a-raw-contributions.mlir b/docs/inter-tile-examples/kt07a-raw-contributions.mlir new file mode 100644 index 0000000..bc81480 --- /dev/null +++ b/docs/inter-tile-examples/kt07a-raw-contributions.mlir @@ -0,0 +1,163 @@ +// RUN: ktir-opt %s --ktir-check-legality | ktir-opt | FileCheck %s +// +// KT-07, half (a) — four raw contributions folded once. +// +// torch-spyre#4300 lx_relayout_workload_coverage.md §5 KT-07: "Raw +// contributions 1, 2, 4, 8 versus already-completed value 15. Distinguish +// reduction from delivery of a completed sum". Half (b) is the completed-sum +// delivery, in kt07b-completed-sum.mlir. §1 of that document states the +// reference: "sum the independent contributions once. If the producing matmul +// has already completed that sum, the following copy must not sum it again." +// +// OPERATION SEQUENCE. +// +// LX-Load (addr_0) -> LocalReduce -> Expand -> Produce -> Reduce[ Add ] +// +// addr_0 = the tile's own contribution slab. `LocalReduce` is ordinary `linalg` +// over the 128 rows, before any inter-tile op; `Reduce[ Add ]` is the cross-tile +// fold with its combiner region. The result is left unstored — a fold may travel +// a different ring and need not land in LX per tile, unlike a copy delivery. +// +// ============================================================================ +// TWO REDUCTIONS, AND ONLY ONE OF THEM IS THE INTER-TILE OP. +// +// tensor<1x128x64xf16> each tile's own slab +// | linalg.reduce over dim 1 <- LOCAL, along an axis +// tensor<1x64xf16> +// | tensor.expand_shape +// tensor<1x1x64xf16> the partial +// | ktdp.inter_tile_reduce <- CROSS-TILE, elementwise +// tensor<1x1x64xf16> the result, same type +// +// The axis reduced is dim 1, extent 128 — a plain data axis, **not** the +// trailing stick-sized 64, which is carried through untouched. That is the +// point of writing it this way: it shows what collapses and what does not. +// +// The inter-tile op reduces no axis at all. `placement = replicate`, so by §4 +// the result type equals the partial type; the fold runs across *tiles* and is +// applied elementwise at each of the 64 coordinates, which are never combined +// with one another. Layout therefore plays no part in this case; +// kt02-receiving-only-cores.mlir is the example that carries the +// layout-verified physical form. +// ============================================================================ +// +// Group structure. One group, four tiles, every tile both producer and consumer. +// +// group g producers {4g .. 4g+3} consumers {4g .. 4g+3} contributions +// g=0 {0, 1, 2, 3} {0, 1, 2, 3} 1, 2, 4, 8 -> 15 +// ... 1 group only (g = 0) +// +// Fixture. +// The harness writes 2^t / 128 into every element of core t's own LX slab, so +// the local sum over the 128 rows is exactly 2^t and the four contributions +// entering the fold are 1, 2, 4, 8. The fully reduced value is 15. +// +// core 0 -> 1 expected at every core after the all-reduce: 15 +// core 1 -> 2 +// core 2 -> 4 a missing contribution gives 14, 13, 11 or 7 +// core 3 -> 8 a double-counted one gives 16 or more +// +// That is why the values are powers of two rather than a ramp: the result names +// exactly which contributions arrived, which is the bit-at-a-time scheme +// coverage §1 asks for. Every value here is exact in fp16 — 2^-7 .. 2^-4, the +// partial sums k * 2^(t-7) for k <= 128, and 15. +// +// Why this half has no measured instance. All 130 records of the pinned +// ownership catalog are `STCDPOpLx` relayouts, and every route class is +// copy-only: all_gather 26, grouped_all_gather_with_replication 65, +// replicate_or_owner_remap 32, permutation 6, general_relayout 1. Not one is a +// reduction, which is what §9.3 of ../inter-tile-communication.md says of the 51 +// and holds for all 130. A genuine cross-core fold comes from split-K matmul +// instead — the pattern `inter_tile_reduce` was implemented for. Half (b), by +// contrast, is what every one of those 130 records is. +// +// R13 and R14 are both satisfied: consumer set == producer set, so the consumers +// are producers (R13, `y` for `reduce`) and this is the all-reduce arm of R14. + +#tiles = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 3 >= 0)> +#one_group = affine_set<(g) : (g == 0)> + +// The whole of the tile's own slab. +#slab = affine_set<(d0, d1, d2) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 127 >= 0, d2 >= 0, -d2 + 63 >= 0)> +#ident3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + +// CHECK-LABEL: func.func @kt07a_raw_contributions +module { + func.func @kt07a_raw_contributions() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + %zero = arith.constant 0.0 : f16 + + // The contributions are input data, read from the tile's own LX. + // `ct_local` with no ct_id is the executing tile's local memory, and the + // slab is that tile's whole buffer, so no tile id enters the address. + // Dense: lane 1, dim1 64, dim0 128*64 = 8192. + %own = ktdp.construct_memory_view %base, + sizes: [1, 128, 64], strides: [8192, 64, 1] { + coordinate_set = #slab, + memory_space = #ktdp.memory_space + } : memref<1x128x64xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0] { + access_tile_set = #slab, access_tile_order = #ident3 + } : memref<1x128x64xf16> -> !ktdp.access_tile<1x128x64xindex> + + %slab = ktdp.load %own_access + : !ktdp.access_tile<1x128x64xindex> -> tensor<1x128x64xf16> + + // LOCAL reduction, along dim 1 (128). Ordinary linalg; the inter-tile ops + // are not involved. The trailing 64 survives. + %red_init = tensor.empty() : tensor<1x64xf16> + %red_zero = linalg.fill ins(%zero : f16) outs(%red_init : tensor<1x64xf16>) + -> tensor<1x64xf16> + %local = linalg.reduce { arith.addf } + ins(%slab : tensor<1x128x64xf16>) + outs(%red_zero : tensor<1x64xf16>) + dimensions = [1] + + // Restore the reduced axis as a unit dim so the partial keeps the rank of + // the slab: dim 1 is now 1, the "already folded" axis. + %partial = tensor.expand_shape %local [[0], [1, 2]] + output_shape [1, 1, 64] + : tensor<1x64xf16> into tensor<1x1x64xf16> + + // identity for the fold: 0.0, shaped as the partial type (R11). + %id_init = tensor.empty() : tensor<1x1x64xf16> + %identity = linalg.fill ins(%zero : f16) outs(%id_init : tensor<1x1x64xf16>) + -> tensor<1x1x64xf16> + + // CHECK: ktdp.inter_tile_produce + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #tiles + -> !ktdp.tile_future<(tensor<1x1x64xf16>), groups = #one_group> + { + ^bb0(%gid: index): + ktdp.yield_partial %partial : tensor<1x1x64xf16> + } + + // CROSS-TILE fold. combine = fold, placement = replicate: no axis changes. + // Writing the combiner region is what grants the scheduler permission to + // re-associate (§3.5, "the associative-commutative contract is by user + // agreement"), so tree, ring or linear are all legal. Every consumer holds + // the same value (§3.7). + // + // CHECK: ktdp.inter_tile_reduce + %reduced = ktdp.inter_tile_reduce(%future) + consumer_tiles_per_group = #tiles, + identity(%identity : tensor<1x1x64xf16>) + : !ktdp.tile_future<(tensor<1x1x64xf16>), groups = #one_group> + -> tensor<1x1x64xf16> + { + ^bb0(%lhs: tensor<1x1x64xf16>, %rhs: tensor<1x1x64xf16>): + %acc = tensor.empty() : tensor<1x1x64xf16> + %sum = linalg.add + ins(%lhs, %rhs : tensor<1x1x64xf16>, tensor<1x1x64xf16>) + outs(%acc : tensor<1x1x64xf16>) -> tensor<1x1x64xf16> + ktdp.yield_reduced %sum : tensor<1x1x64xf16> + } + + // Expected: 15.0 at each of the 64 trailing coordinates, on all four cores. + return + } +} diff --git a/docs/inter-tile-examples/kt07b-completed-sum.mlir b/docs/inter-tile-examples/kt07b-completed-sum.mlir new file mode 100644 index 0000000..fa5d02d --- /dev/null +++ b/docs/inter-tile-examples/kt07b-completed-sum.mlir @@ -0,0 +1,176 @@ +// KT-07, half (b) — an already-completed sum is moved, not folded again. +// +// torch-spyre#4300 lx_relayout_workload_coverage.md §1: "Reduction tests have a +// different reference: sum the independent contributions once. **If the +// producing matmul has already completed that sum, the following copy must not +// sum it again.**" Half (a), the genuine fold, is in +// kt07a-raw-contributions.mlir and it verifies. +// +// OPERATION SEQUENCE. +// +// LX-Load (addr_0) -> LocalReduce -> Expand -> Produce -> Consume +// +// Identical to kt07a up to the delivery, which is the point: the two differ only +// in `Reduce[ Add ]` versus `Consume`, and the shapes do not separate them. +// +// ============================================================================ +// THE ONLY DIFFERENCE FROM HALF (a) IS THE DELIVERY OP. +// +// Both halves read a slab from the tile's own LX and reduce it locally over dim +// 1 to a tensor<1x1x64xf16> partial. That block is textually identical here and +// in half (a). What differs: +// +// half (a) 4 producers, partials 1, 2, 4, 8 ktdp.inter_tile_reduce -> 15 +// half (b) 1 producer, partial 15 ktdp.inter_tile_consume -> 15 +// +// The producer set has to differ too — R8 gives a broadcast one producer per +// group — but nothing else does. Reading this edge as half (a) would fold a +// completed sum a second time. +// +// WHAT 15 IS, AND WHAT 60 IS. 15 = 1 + 2 + 4 + 8, the value half (a) computes; +// it is used here so both halves have the same answer and can be read side by +// side. With one producer the bit-per-tile property of half (a) does not carry +// over — 15 identifies nothing here, it is simply the completed sum. +// +// 60 is NOT a possible outcome of this fixture: only core 4g holds a value, so +// there is nothing to fold four of. 60 is what the *wrong model* of this edge +// produces — reading it as a fold over four holders — which is why the case +// exists. For a single-producer set the shipped verifier rejects that model +// outright (see STATUS below), so here the mistake is caught rather than +// mis-valued. On an edge where several cores really do hold the completed sum +// it would return 60 silently instead, and only the recorded expected values +// would catch it. +// +// Note the shapes do not separate the two readings. `reduce` keeps the partial +// type (§4, replicate), so a completed sum being copied and a fresh fold being +// computed have the same result type. Only the expected values separate them, +// which is why the coverage document's test contract records, per edge, whether +// the inputs are "independent tensor elements, unfinished sums, or +// already-completed sums". +// ============================================================================ +// +// Group structure. One group, four tiles, but only one of them produces. +// +// group g producers {4g} consumers {4g .. 4g+3} non-producers +// g=0 {0} {0, 1, 2, 3} 1, 2, 3 +// ... 1 group only (g = 0) +// +// Fixture. The harness writes 15 / 128 into every +// element of core 4g's own LX slab, and the local reduction below sums the 128 +// rows to exactly 15. That local reduction stands in for the producing matmul: +// it is where the sum is completed, and the delivery must not repeat it. Cores +// 4g+1, 4g+2 and 4g+3 hold nothing and must still end with 15. 15/128 and +// every partial sum k * 15 * 2^-7 for k <= 128 are exact in fp16. +// +// THIS HALF IS WHAT THE MEASURED CATALOG IS. All 130 records of the pinned +// ownership catalog are `STCDPOpLx` relayouts and every route class is +// copy-only: all_gather 26, grouped_all_gather_with_replication 65, +// replicate_or_owner_remap 32, permutation 6, general_relayout 1. Not one is a +// reduction. So the risk this case guards against is one-sided: the mistake to +// avoid is reading a copy as a fold, and 130 of the 130 measured edges are +// copies whose consumer names (`mean_*`, `_safe_softmax-Sum`, `mm-BMM_1`) invite +// exactly that reading. +// +// ============================================================================ +// STATUS: this form does NOT parse. `ktdp.inter_tile_consume` is specified +// (§6.1 of ../inter-tile-communication.md) but absent from KTDP.td, which +// defines `inter_tile_produce` and `inter_tile_reduce` only. +// +// And it cannot be worked around with a degenerate `reduce` over one producer. +// Writing this delivery as `inter_tile_reduce` with +// `producer_tiles_per_group = {4g}` and `consumer_tiles_per_group = {4g..4g+3}` +// is rejected by the shipped legality pass: +// +// $ ktir-opt broadcast_as_degenerate_reduce.mlir --ktir-check-legality +// error: consumer_tiles_per_group for group 0 is not a subset of +// producer_tiles_per_group (a consumer tile that did not produce is +// unsupported; see open question Q1) +// +// That is R13 (`KTIRCheckLegality.cpp:107-117`), and R14's mode gate would +// reject it too: the consumer set neither equals the producer set nor is a +// single tile. So half (b) needs `consume`; the fold op cannot stand in for it. +// The R13 message names the open question, which §10.1 now resolves as `n` for +// every op except `reduce` — but for `reduce` itself the check is correct to +// fire, because folding into a non-contributing tile is the case §10.1 leaves +// open. +// ============================================================================ + +#tiles = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 3 >= 0)> +// The completion owner: core 4g only. Declaring it is the "declare completion +// owners" half of KT-07's requirement. +#completed = affine_set<(i)[g] : (i - 4 * g == 0)> +#one_group = affine_set<(g) : (g == 0)> + +// The whole of the tile's own slab. +#slab = affine_set<(d0, d1, d2) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 127 >= 0, d2 >= 0, -d2 + 63 >= 0)> +#ident3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> + +module { + func.func @kt07b_completed_sum() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + %zero = arith.constant 0.0 : f16 + + // ---- identical to half (a) from here ---- + // + // One producer per group, so strictly these loads belong inside the produce + // region: they must not run on the cores that hold nothing (§2.2). They are + // at function scope here so the block stays textually identical to half (a), + // which is the comparison this pair is for. A production form moves them + // into the region, as §7.7.1 does for `scatter`. + %own = ktdp.construct_memory_view %base, + sizes: [1, 128, 64], strides: [8192, 64, 1] { + coordinate_set = #slab, + memory_space = #ktdp.memory_space + } : memref<1x128x64xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0] { + access_tile_set = #slab, access_tile_order = #ident3 + } : memref<1x128x64xf16> -> !ktdp.access_tile<1x128x64xindex> + + %slab = ktdp.load %own_access + : !ktdp.access_tile<1x128x64xindex> -> tensor<1x128x64xf16> + + // LOCAL reduction, along dim 1 (128). This is where the sum is completed. + %red_init = tensor.empty() : tensor<1x64xf16> + %red_zero = linalg.fill ins(%zero : f16) outs(%red_init : tensor<1x64xf16>) + -> tensor<1x64xf16> + %local = linalg.reduce { arith.addf } + ins(%slab : tensor<1x128x64xf16>) + outs(%red_zero : tensor<1x64xf16>) + dimensions = [1] + + %partial = tensor.expand_shape %local [[0], [1, 2]] + output_shape [1, 1, 64] + : tensor<1x64xf16> into tensor<1x1x64xf16> + // ---- identical to half (a) up to here ---- + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #completed + -> !ktdp.tile_future<(tensor<1x1x64xf16>), groups = #one_group> + { + ^bb0(%gid: index): + ktdp.yield_partial %partial : tensor<1x1x64xf16> + } + + // combine = none, placement = replicate. No combiner region and no + // identity: there is nothing to fold, which is the whole point. Half (a) + // has both at this position. R8 is satisfied with one producer per group, + // so no dependency attribute is needed and full-barrier and per-tile + // synchronization coincide (§6.1). + // + // Ordering is trivial for `replicate`: every consumer receives the whole + // value, so there is no `l` and no assembly order to get wrong — the other + // half of KT-07's "declare completion owners and ordering". + %value = ktdp.inter_tile_consume(%future) + consumer_tiles_per_group = #tiles + : !ktdp.tile_future<(tensor<1x1x64xf16>), groups = #one_group> + -> tensor<1x1x64xf16> + + // Expected: 15.0 at each of the 64 trailing coordinates, on all four cores. + // 60.0 is what a fold over four holders would give; see the note above on + // why this fixture cannot produce it and where it could. + return + } +} diff --git a/docs/inter-tile-examples/kt09-reuse-and-lifetime.mlir b/docs/inter-tile-examples/kt09-reuse-and-lifetime.mlir new file mode 100644 index 0000000..822a8aa --- /dev/null +++ b/docs/inter-tile-examples/kt09-reuse-and-lifetime.mlir @@ -0,0 +1,445 @@ +// KT-09 — many readers of one delivered value, and scratch reuse across iterations. +// +// torch-spyre#4300 lx_relayout_workload_coverage.md §5 KT-09: "Multiple consumers +// read one delivered X; then two expert/page iterations reuse scratch only after +// reads complete. **Demonstrate reuse and lifetime without making every arithmetic +// reader initiate another transfer. This does not require multiple delivery users +// of one future.**" +// +// That last sentence is the whole answer to the first half, and it is worth stating +// why: **R2 constrains the `!ktdp.tile_future`, not the delivery's result.** The +// future must have exactly one use — the delivery op (`KTIRCheckLegality.cpp:80-85`). +// The delivery's *result* is an ordinary `tensor` SSA value with no use restriction, +// so any number of arithmetic readers may consume it and none of them is a transfer. +// One `produce` + one delivery + N readers is the natural spelling, not a workaround. +// +// OPERATION SEQUENCE. +// +// @kt09_three_readers LX-Load (addr_0) -> Produce -> Gather -> Mul x 3 +// -> LX-Store x 3 (addr_1, addr_2, addr_3) +// +// @kt09_scratch_reuse Produce -> Gather -> LX-Store (addr_0) +// -> LX-Load x 2 (addr_0) -> Add -> LX-Store (addr_1) +// Produce -> Gather -> LX-Store (addr_0) +// -> LX-Load x 2 (addr_0) -> Add -> LX-Store (addr_2) +// +// Half 1: addr_0 = own piece, addr_1..3 = the Q/K/V outputs. One Gather feeds +// three Muls — the delivered value stays a tensor and is never landed by the IR. +// +// Half 2: addr_0 = the reused area, addr_1/addr_2 = per-iteration outputs. The +// second `LX-Store (addr_0)` is the reuse, and it is a write-after-read against +// the two `LX-Load (addr_0)` above it. No LX-Load before the first Produce: the +// partials are computed, so addr_0 is the only area in the function. +// +// ============================================================================ +// MEASURED FIXTURE: the same relayout recorded three times, once per reader. +// +// `mean-LayerNormNorm_out` (prefill) is the source tensor of **three** catalog +// records, all input 0 of a matmul: +// +// relayouts[ 1] consumer mm-BMM_1 (this is GR-PF-002, kt01's and kt02's) +// relayouts[ 8] consumer mm_1-BMM_1 +// relayouts[16] consumer mm_2-BMM_1 +// +// One layer-norm output feeding three projections — Q, K and V. Checked: the three +// records are **identical** in extents, source and destination pieces, owner tables, +// route class and fragment counts. They are not three relayouts; they are one +// relayout that three consumers need. +// +// extents {in: 4096, mb: 512, y: 1} +// 16 source pieces {in: 4096, mb: 32, y: 1}, one owner each +// 8 destination pieces {in: 4096, mb: 64, y: 1}, four owners each +// remote_destination_bytes 12582912 (12 MiB), the same in all three records +// +// THE COST OF GETTING IT WRONG IS MEASURED, NOT ARGUED. Emitting one delivery per +// reader moves 3 * 12 MiB = 36 MiB where 12 MiB suffices — **exactly 3x** the ring +// traffic for the same result. The catalog invites the mistake by construction: it +// is indexed by (consumer, input), so a shared source appears as several records and +// a reader that walks the records one at a time will emit one transfer each. Eight +// of the 120 distinct tensors in the catalog are shared this way. +// +// Group structure — GR-PF-002's, the same as kt01 and kt02. Source piece k is owned +// by core 2k and starts at mb = 32k, so group g (destination mb 64g..64g+63) is: +// +// group g producers {4g, 4g+2} consumers {4g .. 4g+3} non-producers +// g=0 {0, 2} {0, 1, 2, 3} 1, 3 +// g=1 {4, 6} {4, 5, 6, 7} 5, 7 +// ... 8 groups, regular +// g=7 {28, 30} {28, 29, 30, 31} 29, 31 +// +// LAYOUT. As kt02 derives it: an SDSC run of this logical shape reports +// `layoutDimOrder_ = ["mb", "out", "y"]`, `stickDimOrder_ = ["y"]`, `stickSize_ = 64`, +// so `y` is innermost and sticked, an extent-1 `y` is 64 padded lanes, and the +// physical form is one rank higher with the data axes reversed. The one assumption +// kt02 leaves standing — that this record's order is `["mb", "in", "y"]` — is +// unchanged here and does not affect what this file is about. +// ============================================================================ +// +// WHAT THE TWO HALVES OF KT-09 COST, AND WHY THEY DIFFER. +// +// **Half 1, many readers: expressible, and checked.** R2 is a rule and it is +// implemented, so a spelling that gave each reader its own future would be *rejected* +// if it tried to reuse one produce — and if it duplicated the produce instead, the +// duplication is visible as N `inter_tile_produce` ops in the IR. Either way the +// mistake is not silent. This is the one half of one KT case so far that the +// verifier and the eye both reach. +// +// **Half 2, lifetime: expressible, but only if the buffer is named.** Two things are +// being asked for and KTIR treats them differently. +// +// (a) A named area's ordering IS expressed, and needs no inter-tile rule. +// Iteration 1 must not store into `%shared` before iteration 0's readers have +// loaded from it. Both are memory effects on one `memref`, so ordinary MLIR +// carries it: a write-after-read on one memref is a dependence no pass may +// reorder. Half 2 is that, and nothing more. +// +// (b) An *unnamed* delivered value has no live range to reason about. The +// delivery returns a `tensor` — no address, no memory space (§2.2's future +// carries `tensor`, not `memref`). So half 1's three readers read no buffer; +// the memref read ended at `ktdp.load`, before the produce, and **those +// readers cannot extend a live range because there is no live range to +// extend.** For a copy delivery the buffer nevertheless exists physically — +// the hardware lands every received tile in LX — so the question is never +// whether there is one, only whether the IR says so. Name it and the lifetime +// is checkable; leave it unnamed and coverage §1's "buffers preserved until +// last reader" lands on the emitted program instead. +// +// The root cause of (b) is the same "no address" fact kt03 leaned on to put its +// reorder in store bases, and the same one behind Table 1's note that a +// dist-mem-view cannot declare ownership for an intermediate. Three cases, one +// root cause — but note the sign differs. For kt03 the absence of an address is +// what makes the reorder free; here it is what makes the lifetime invisible. +// +// NOT VERIFIED as a whole: `ktdp.inter_tile_gather` is specified (§6.4) but absent +// from KTDP.td, in both functions. The `ktdp.inter_tile_produce` halves, the three +// readers and the scratch ordering all parse and round-trip. +// ============================================================================ + +#producers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 2 >= 0, + i mod 2 == 0)> +#consumers = affine_set<(i)[g] : (i - 4 * g >= 0, -i + 4 * g + 3 >= 0)> +#groups = affine_set<(g) : (g >= 0, -g + 7 >= 0)> + +#piece = affine_set<(d0, d1, d2, d3) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 4095 >= 0, + d2 >= 0, -d2 + 31 >= 0, d3 >= 0, -d3 + 63 >= 0)> +#region = affine_set<(d0, d1, d2, d3) : ( + d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + 4095 >= 0, + d2 >= 0, -d2 + 63 >= 0, d3 >= 0, -d3 + 63 >= 0)> +#ident4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> + +// --------------------------------------------------------------------------- +// HALF 1 — one delivery, three readers. +// +// The three `linalg.mul` ops stand for the Q, K and V projections; what they compute +// is not the point and a real kernel would put matmuls here. The point is what sits +// *above* them: one `inter_tile_produce`, one `inter_tile_gather`, and `%region_val` +// used three times. Reading a `tensor` three times is free — there is no second +// transfer to elide, because a transfer is an op and there is only one. +// --------------------------------------------------------------------------- + +module { + func.func @kt09_three_readers() { + %c0 = arith.constant 0 : index + %base = arith.constant 0 : index + %one = arith.constant 1.0 : f16 + %two = arith.constant 2.0 : f16 + %four = arith.constant 4.0 : f16 + + // Three distinct destination buffers. The source piece occupies + // 1*4096*32*64 = 8388608 elements from 0, and each assembled region is + // 1*4096*64*64 = 16777216, so the three bases are laid end to end after it. + %base_q = arith.constant 8388608 : index + %base_k = arith.constant 25165824 : index + %base_v = arith.constant 41943040 : index + + %own = ktdp.construct_memory_view %base, + sizes: [1, 4096, 32, 64], strides: [8388608, 2048, 64, 1] { + coordinate_set = #piece, + memory_space = #ktdp.memory_space + } : memref<1x4096x32x64xf16> + + %own_access = ktdp.construct_access_tile %own[%c0, %c0, %c0, %c0] { + access_tile_set = #piece, access_tile_order = #ident4 + } : memref<1x4096x32x64xf16> -> !ktdp.access_tile<1x4096x32x64xindex> + + // The memref read ends HERE. Everything downstream is value semantics. + %piece_val = ktdp.load %own_access + : !ktdp.access_tile<1x4096x32x64xindex> -> tensor<1x4096x32x64xf16> + + %future = ktdp.inter_tile_produce + producer_tiles_per_group = #producers + -> !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %piece_val : tensor<1x4096x32x64xf16> + } + + // ONE delivery. `%future` has exactly one use, which is R2 satisfied. + %region_val = ktdp.inter_tile_gather(%future) + consumer_tiles_per_group = #consumers, + gather_dimensions = [2] + : !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + -> tensor<1x4096x64x64xf16> + + // THREE readers of the one result. R2 says nothing about this — it is a rule on + // the future, and `%region_val` is a plain tensor. + %wq_e = tensor.empty() : tensor<1x4096x64x64xf16> + %wq = linalg.fill ins(%one : f16) outs(%wq_e : tensor<1x4096x64x64xf16>) + -> tensor<1x4096x64x64xf16> + %q_e = tensor.empty() : tensor<1x4096x64x64xf16> + %q = linalg.mul ins(%region_val, %wq : tensor<1x4096x64x64xf16>, + tensor<1x4096x64x64xf16>) + outs(%q_e : tensor<1x4096x64x64xf16>) -> tensor<1x4096x64x64xf16> + + %wk_e = tensor.empty() : tensor<1x4096x64x64xf16> + %wk = linalg.fill ins(%two : f16) outs(%wk_e : tensor<1x4096x64x64xf16>) + -> tensor<1x4096x64x64xf16> + %k_e = tensor.empty() : tensor<1x4096x64x64xf16> + %k = linalg.mul ins(%region_val, %wk : tensor<1x4096x64x64xf16>, + tensor<1x4096x64x64xf16>) + outs(%k_e : tensor<1x4096x64x64xf16>) -> tensor<1x4096x64x64xf16> + + %wv_e = tensor.empty() : tensor<1x4096x64x64xf16> + %wv = linalg.fill ins(%four : f16) outs(%wv_e : tensor<1x4096x64x64xf16>) + -> tensor<1x4096x64x64xf16> + %v_e = tensor.empty() : tensor<1x4096x64x64xf16> + %v = linalg.mul ins(%region_val, %wv : tensor<1x4096x64x64xf16>, + tensor<1x4096x64x64xf16>) + outs(%v_e : tensor<1x4096x64x64xf16>) -> tensor<1x4096x64x64xf16> + + // Store the three results to three distinct destinations, so the readers are + // visibly independent rather than a chain. The three `coordinate_set`s are the + // same `#region` on purpose: Q, K and V are three *different* tensors and each + // core holds the same region of each, so what has to differ is the base address, + // not the coordinate set. + %out_q = ktdp.construct_memory_view %base_q, + sizes: [1, 4096, 64, 64], strides: [16777216, 4096, 64, 1] { + coordinate_set = #region, + memory_space = #ktdp.memory_space + } : memref<1x4096x64x64xf16> + %at_q = ktdp.construct_access_tile %out_q[%c0, %c0, %c0, %c0] { + access_tile_set = #region, access_tile_order = #ident4 + } : memref<1x4096x64x64xf16> -> !ktdp.access_tile<1x4096x64x64xindex> + ktdp.store %q, %at_q + : tensor<1x4096x64x64xf16>, !ktdp.access_tile<1x4096x64x64xindex> + + %out_k = ktdp.construct_memory_view %base_k, + sizes: [1, 4096, 64, 64], strides: [16777216, 4096, 64, 1] { + coordinate_set = #region, + memory_space = #ktdp.memory_space + } : memref<1x4096x64x64xf16> + %at_k = ktdp.construct_access_tile %out_k[%c0, %c0, %c0, %c0] { + access_tile_set = #region, access_tile_order = #ident4 + } : memref<1x4096x64x64xf16> -> !ktdp.access_tile<1x4096x64x64xindex> + ktdp.store %k, %at_k + : tensor<1x4096x64x64xf16>, !ktdp.access_tile<1x4096x64x64xindex> + + %out_v = ktdp.construct_memory_view %base_v, + sizes: [1, 4096, 64, 64], strides: [16777216, 4096, 64, 1] { + coordinate_set = #region, + memory_space = #ktdp.memory_space + } : memref<1x4096x64x64xf16> + %at_v = ktdp.construct_access_tile %out_v[%c0, %c0, %c0, %c0] { + access_tile_set = #region, access_tile_order = #ident4 + } : memref<1x4096x64x64xf16> -> !ktdp.access_tile<1x4096x64x64xindex> + ktdp.store %v, %at_v + : tensor<1x4096x64x64xf16>, !ktdp.access_tile<1x4096x64x64xindex> + + // Expected: `%q`, `%k` and `%v` each hold the assembled region scaled by 1, 2 + // and 4, so a reader that received the wrong region shows up as the wrong + // multiple. The three destination extents are 16777216 elements apart, so the + // stores do not alias and a later reader of Q cannot see V's result. + // + // Failures to catch. (1) Three deliveries instead of one: the result is + // identical, so **only the emitted transfer count separates them** — 36 MiB + // against 12 MiB. Coverage §1's "check the emitted memory accesses as well as + // numerical output" is the check that bites, and no numerical test can. + // (2) Duplicating the produce as well as the delivery: visible in the IR as + // three `inter_tile_produce` ops, and each core then computes its partial three + // times. + return + } +} + +// --------------------------------------------------------------------------- +// HALF 2 — one LX area reused by two iterations, with several readers in between. +// +// SYNTHETIC in its iteration structure: Granite 3.3 is dense, so it has no expert +// loop, and the pinned catalog holds no two-iteration reuse fixture. A page loop +// over `cat_*-kvCacheScatter` would be the measured shape to attach here once its +// descriptors are available. The delivery itself is the measured one above. +// +// **The received tile is stored, and that store is not overhead.** For an LX-to-LX +// transfer this hardware lands every received tile in LX before a compute unit can +// read it, so a landing store exists whether or not the IR names one. Half 1 does +// not name it — the delivered `tensor` feeds `linalg` directly and the backend +// invents the buffer. Half 2 names it, and that is the whole difference between +// them: naming costs nothing extra and buys a live range that the rules can see. +// +// **This is a fact about copy delivery, not about `reduce`.** A reduction is not +// confined to the LX-to-LX path: a compute unit can send a tile out over a different +// ring, so a fold need not land each partial in LX on the way. So the landing store +// is implied by `consume`, `gather`, `scatter` and `all_to_all` — the ops that move +// data unchanged — and not by `inter_tile_reduce`, which is why kt07a can leave its +// result unstored without that being a shortcut. +// +// So the two halves are a genuine choice, not a good form and a bad one: +// +// delivered value kept as a `tensor` no store in the IR; the staging buffer is +// (half 1, and kt03/kt04/kt06) the backend's, and its lifetime with it +// +// delivered value stored into LX the landing is explicit; readers read the +// (half 2) memref, so the live range is in the IR +// +// WHAT THIS DEMONSTRATES. `%shared` is the reused area. Iteration 0 stores the +// delivered region into it and two readers load from it; iteration 1 then stores its +// own delivered region into the same address. That second store is the reuse, and it +// is a **write-after-read** against both of iteration 0's loads — conflicting memory +// effects on one `memref`, which no pass may reorder. Coverage §5's "reuse scratch +// only after reads complete" is therefore discharged by ordinary MLIR memory +// semantics, and needs no inter-tile rule at all. +// +// Note which readers hold the buffer. In half 1 the three readers consume a +// `tensor` and hold nothing, so they cannot keep a buffer alive. Here they consume +// the `memref`, and they can. That is the same distinction from the other side. +// +// The partials are computed rather than loaded, so `%shared` is the only memref in +// this function and the reuse cannot be confused with a source buffer's traffic. +// --------------------------------------------------------------------------- + +module { + func.func @kt09_scratch_reuse() { + %c0 = arith.constant 0 : index + %v0 = arith.constant 1.0 : f16 + %v1 = arith.constant 8.0 : f16 + + // The reused LX area, and the two per-iteration outputs after it. One region is + // 1*4096*64*64 = 16777216 elements. + %base_shared = arith.constant 0 : index + %base_out0 = arith.constant 16777216 : index + %base_out1 = arith.constant 33554432 : index + + %shared = ktdp.construct_memory_view %base_shared, + sizes: [1, 4096, 64, 64], strides: [16777216, 4096, 64, 1] { + coordinate_set = #region, + memory_space = #ktdp.memory_space + } : memref<1x4096x64x64xf16> + %shared_at = ktdp.construct_access_tile %shared[%c0, %c0, %c0, %c0] { + access_tile_set = #region, access_tile_order = #ident4 + } : memref<1x4096x64x64xf16> -> !ktdp.access_tile<1x4096x64x64xindex> + + // ---- iteration 0 ---- + %e0 = tensor.empty() : tensor<1x4096x32x64xf16> + %p0 = linalg.fill ins(%v0 : f16) outs(%e0 : tensor<1x4096x32x64xf16>) + -> tensor<1x4096x32x64xf16> + + %fut0 = ktdp.inter_tile_produce + producer_tiles_per_group = #producers + -> !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %p0 : tensor<1x4096x32x64xf16> + } + + %region_0 = ktdp.inter_tile_gather(%fut0) + consumer_tiles_per_group = #consumers, + gather_dimensions = [2] + : !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + -> tensor<1x4096x64x64xf16> + + // THE LANDING, named. On hardware this store happens either way; writing it + // here is what gives `%shared` a live range the IR can reason about. + ktdp.store %region_0, %shared_at + : tensor<1x4096x64x64xf16>, !ktdp.access_tile<1x4096x64x64xindex> + + // TWO readers, both reading the LX area rather than a tensor. These are what + // iteration 1's store must wait for. + %a0 = ktdp.load %shared_at + : !ktdp.access_tile<1x4096x64x64xindex> -> tensor<1x4096x64x64xf16> + %b0 = ktdp.load %shared_at + : !ktdp.access_tile<1x4096x64x64xindex> -> tensor<1x4096x64x64xf16> + + %s0 = tensor.empty() : tensor<1x4096x64x64xf16> + %c0v = linalg.add ins(%a0, %b0 : tensor<1x4096x64x64xf16>, + tensor<1x4096x64x64xf16>) + outs(%s0 : tensor<1x4096x64x64xf16>) -> tensor<1x4096x64x64xf16> + + %out0 = ktdp.construct_memory_view %base_out0, + sizes: [1, 4096, 64, 64], strides: [16777216, 4096, 64, 1] { + coordinate_set = #region, + memory_space = #ktdp.memory_space + } : memref<1x4096x64x64xf16> + %at0 = ktdp.construct_access_tile %out0[%c0, %c0, %c0, %c0] { + access_tile_set = #region, access_tile_order = #ident4 + } : memref<1x4096x64x64xf16> -> !ktdp.access_tile<1x4096x64x64xindex> + ktdp.store %c0v, %at0 + : tensor<1x4096x64x64xf16>, !ktdp.access_tile<1x4096x64x64xindex> + + // ---- iteration 1, reusing %shared ---- + %e1 = tensor.empty() : tensor<1x4096x32x64xf16> + %p1 = linalg.fill ins(%v1 : f16) outs(%e1 : tensor<1x4096x32x64xf16>) + -> tensor<1x4096x32x64xf16> + + %fut1 = ktdp.inter_tile_produce + producer_tiles_per_group = #producers + -> !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + { + ^bb0(%gid: index): + ktdp.yield_partial %p1 : tensor<1x4096x32x64xf16> + } + + %region_1 = ktdp.inter_tile_gather(%fut1) + consumer_tiles_per_group = #consumers, + gather_dimensions = [2] + : !ktdp.tile_future<(tensor<1x4096x32x64xf16>), groups = #groups> + -> tensor<1x4096x64x64xf16> + + // THE REUSE. This store overwrites iteration 0's landing, and it is a + // write-after-read against `%a0` and `%b0` above. Hoisting it over either of + // them would give iteration 0's readers iteration 1's data. + ktdp.store %region_1, %shared_at + : tensor<1x4096x64x64xf16>, !ktdp.access_tile<1x4096x64x64xindex> + + %a1 = ktdp.load %shared_at + : !ktdp.access_tile<1x4096x64x64xindex> -> tensor<1x4096x64x64xf16> + %b1 = ktdp.load %shared_at + : !ktdp.access_tile<1x4096x64x64xindex> -> tensor<1x4096x64x64xf16> + + %s1 = tensor.empty() : tensor<1x4096x64x64xf16> + %c1v = linalg.add ins(%a1, %b1 : tensor<1x4096x64x64xf16>, + tensor<1x4096x64x64xf16>) + outs(%s1 : tensor<1x4096x64x64xf16>) -> tensor<1x4096x64x64xf16> + + %out1 = ktdp.construct_memory_view %base_out1, + sizes: [1, 4096, 64, 64], strides: [16777216, 4096, 64, 1] { + coordinate_set = #region, + memory_space = #ktdp.memory_space + } : memref<1x4096x64x64xf16> + %at1 = ktdp.construct_access_tile %out1[%c0, %c0, %c0, %c0] { + access_tile_set = #region, access_tile_order = #ident4 + } : memref<1x4096x64x64xf16> -> !ktdp.access_tile<1x4096x64x64xindex> + ktdp.store %c1v, %at1 + : tensor<1x4096x64x64xf16>, !ktdp.access_tile<1x4096x64x64xindex> + + // Expected. Every producer contributes 1.0 in iteration 0, so the assembled + // region is 1.0 everywhere and `%out0` holds 1.0 + 1.0 = 2.0. Iteration 1 + // contributes 8.0, so `%out1` holds 16.0. Both exact in fp16. + // + // The failures this fixture exists to catch: + // + // `%out0` reading 16.0 iteration 1's store was hoisted over `%a0` or `%b0` — + // the reuse happened before the reads completed, which + // is exactly coverage §5's requirement. + // `%out1` reading 2.0 iteration 1's store was elided as dead, since a naive + // pass may see `%shared` written twice and keep the + // first. + // either reading 1.0 one of the two loads was folded away, so the add ran + // against a single reader and the buffer's live range + // was shorter than the source says. + // + // Coverage §1's "repeat after poison" applies directly: poison `%shared` between + // the iterations and `%out1` must still be 16.0. + return + } +}