LowerCompositeOps Pass

Decomposes composite tile / distributed ops into primitive operations so codegen does not need to emit their high-level forms. Today the pass handles tile.sin / tile.cos (FP32 Cody-Waite + Horner), tile.select (packed predicate mask + TSEL / TSELS), packed tile.tquant_mx, and pld.tensor.* distributed collectives (allreduce (mesh and ring), allgather, reduce_scatter, broadcast, barrier, all_to_all, all_to_all_v). Mesh and ring allreduce may also create a metadata-preserving tensor.view so tile load/remote/store operate on a 2D flattened target window.

Overview

LowerCompositeOps is a function-level pass that rewrites every var = Call(...) AssignStmt whose callee appears in the pass's lowering dispatch table. For tile.sin / tile.cos, the rule emits a fixed-shape primitive tile-op recipe using tile.muls, tile.adds, tile.add, tile.sub, tile.mul, and tile.cast: Cody-Waite range reduction (4-part π split) followed by a degree-9 odd Horner polynomial. tile.tquant_mx becomes value-returning tile.tquant_mx_raw + tile.tmov_x2zz (gather_compare-style SSA) plus explicit workspace tiles; public FP8 aliases are added via reinterpret_view / transpose_view. For pld.tensor.* distributed collectives, the rules emit the cross-rank recipes documented below; pld.tensor.allreduce remains explicit-signal in InCore/composite lowering. The original target Var is preserved as the LHS of the final AssignStmt, so downstream uses keep the same name and identity.

Host-orchestrator pld.tensor.allreduce calls are skipped by this pass: SynthesizeAllReduceSignals first normalizes optional-signal host calls to the explicit-signal form, MaterializeCommDomainScopes places the data and signal windows into comm domains, and then LowerHostTensorCollectives lowers them to internal builtin dispatches.

The tile.sin / tile.cos rules are FP32-only. Non-FP32 trig inputs are rejected at op-construction time by the shared DeduceTileFP32OnlyType deducer (see src/ir/op/tile_ops/unary.cpp:94), so those rules only see well-typed FP32 operands. Distributed rules have their own dtype contracts; allreduce supports FP16 and FP32 as documented below.

The pass is structural no-op on programs that contain no registered composite call such as tile.sin, tile.cos, tile.tquant_mx, or pld.tensor.* distributed collectives: every other statement passes through IRMutator::VisitStmt_. The decomposition emits only primitive tile ops, the internal value-returning tile.tquant_mx_raw / tile.tmov_x2zz forms, and distributed primitives, none of which the mutator rewrites — so the pass is also idempotent.

Requires: nothing.

Produces: nothing.

Invalidates: nothing.

The empty PassProperties contract (kLowerCompositeOpsProperties in include/pypto/ir/transforms/pass_properties.h) reflects that the lowering operates within the existing tile/distributed vocabulary plus metadata-only tensor views (tensor.view) used to expose canonical flattened windows; a partial-prefix ring view also carries a flattened valid_shape. The pass neither establishes nor breaks any IRProperty.

When It Runs

LowerCompositeOps is the first entry of tile_pto_passes in the Default pipeline (see python/pypto/ir/pass_manager.py), running immediately after ConvertTensorToTileOps (slot 12) and OptimizeOrchTensors (slot 13). At this point all tensor-level transcendental calls (tensor.sin, tensor.cos) have been rewritten to their tile equivalents (tile.sin, tile.cos) by the conversion registry, and the tile pipeline is about to start tile-shape canonicalisation. Lowering trig before FlattenTileNdTo2D keeps the decomposition independent of the 2D-flattening rules — every primitive tile op in the recipe (tile.muls, tile.adds, tile.add, tile.sub, tile.mul, tile.cast) has well-defined behaviour at any rank. Packed tile.tquant_mx is also lowered here (after FlattenCallExpr has stabilized tuple consumers).

Pass order vs memory space: LowerCompositeOps runs before InferTileMemorySpace. Scratch tiles created by the tile.tquant_mx rule therefore stamp MemorySpace::Vec explicitly on tile.create so later memory planning still sees addresses; this is intentional, not a signal that Infer should have already run.

Architecture

The pass keeps statement rewriting separate from its lowering recipes. All rules live under src/ir/transforms/lower_composite/; the existing pass entry point owns traversal, result binding, and tuple alias tracking.

src/ir/transforms/lower_composite/
  lower_composite_builder.{h,cpp}   — per-call statement builder and control flow
  lower_composite_common.{h,cpp}    — shared collective geometry and validation
  lower_composite_rules.h          — CompositeLoweringFn and internal rule declarations
  lower_composite_rules.cpp        — op-name -> rule dispatch table (kRules)
  lower_composite_math.cpp         — sin/cos and their shared polynomial recipe
  lower_composite_select.cpp       — select, scalar materialization and scratch binding
  lower_composite_mx.cpp           — tquant_mx and its workspace/view construction
  lower_composite_{allreduce,allgather,reduce_scatter,broadcast,barrier,
                   all_to_all,all_to_all_v}.cpp
                                  — one translation unit per collective

src/ir/transforms/lower_composite_ops_pass.cpp
  LowerCompositeOpsMutator         — statement rewriting and tuple/variable remapping
  TransformLowerCompositeOps       — per-function context and managed-collective deferral
  LowerCompositeOps                — pass factory

CompositeLoweringFn receives the original CallPtr, already visited operands, and a LoweringBuilder&, and returns the final result expression. Rules must use the visited operands so variable remapping is preserved. The builder accumulates intermediate statements and provides primitive tile operations and structured control flow (EmitFor, EmitForReduce, EmitIf, EmitIfExpr). Its temporary counter is borrowed from the function's mutator and shared by nested builders; barrier generations remain local to each top-level collective call.

To add a composite op:

1. Define Lower<Op>Rule(call, args, builder) in the appropriate module under lower_composite/. Keep recipe-specific helpers and constants file-local. Declare the entry point in lower_composite_rules.h; add any new .cpp to CMakeLists.txt.
2. Add a {"<op>", &Lower<Op>Rule} entry to kRules in lower_composite_rules.cpp. This immutable function-local table uses explicit function references, with no static self-registration or initialization-order dependency. An unlisted call follows the existing IRMutator path.
3. For a new managed collective, also update the shared op_predicates::IsManagedTensorCollective classification. The mutator defers these calls in orchestration functions to the HOST/CHIP collective passes; the rule table does not maintain another classification list.

Multi-result rules return MakeTuple. The mutator retains result and alias-chain tracking so projections expose the same destinations. Rules emit primitive ops outside the table, preserving idempotency. AssignStmt, EvalStmt, and direct ReturnStmt calls retain their existing rewrite behavior. Moving a rule changes neither its backend queries nor the pass order and properties.

Algorithm (tile.tquant_mx rule)

The lowering creates source-dtype max / scaling write-only workspace tiles, then binds a value-returning tile.tquant_mx_raw(src, max, scaling) whose TupleType{INT8 dst, UINT8 exp} elements are projected with TupleGetItem (same SSA pattern as tile.gather_compare). Codegen resolves those projections via ResolveTupleResultElements and emits pto.tquant.mx with four outs (dst, exp, max, scaling). Public group_axis matches PTOAS grpAxis: axis1 keeps [M,K] and returns scale [M,K/32]; axis0 first transposes [N,K] to [K,N] and returns scale [K/32,N]. Both forms then bind value-returning tile.tmov_x2zz(exp, tmp). Axis1 tmp capacity is 64 + ceil(rows/16)*cols bytes (33-byte-aligned); axis0 uses a minimal 32-byte Vec pad required by TMovDnTo2Zz. Axis0 follows pto-isa TMovDnTo2Zz (pin be5ccb76): DN [M̂,N] → ZZ [N,M̂] row/row, then a zero-copy tile.transpose_view yields the public [M̂,N] col/col scale. A zero-copy FP8 data alias and FP8E8M0 scale alias complete the public tuple. Write effects on the workspace args keep the Calls alive under DCE when a public result is unused.

On Ascend950, the public quant data and scale may feed matmul_mx in the same InCore mixed task. ExpandMixedKernel carries both results directly over V2C.

The mutator records composite-produced MakeTuples (and SSA aliases of those tuples) in a private composite_tuples_ map so TupleGetItem projections fold without abusing global var_remap_ for arbitrary v = (a, b) assignments. The internal tile.tquant_mx_raw / tile.tmov_x2zz ops are not registered as composite rules, preserving pass idempotency.

Algorithm (tile.select rule)

tile.select(cond, on_true, on_false) is the scratch-free selection surface; tile.sel / tile.sels remain the 1:1 PTO forms where the caller owns the scratch buffer.

The condition is always the packed predicate mask — the tile.cmp / tile.cmps result type for this result's geometry, [M, roundup(ceil(N/8), 32)] UINT8 with valid_shape[1] = ceil(valid_N/8). IsPackedPredicateMask (include/pypto/ir/type_inference.h) is the single check, shared by DeduceTileSelectType and this rule, so the mask arrives ready for both TSEL and TSELS and the rule synthesizes no compare at all.

A plain 0/1 value tile is deliberately not a second accepted form. A mask's valid extent counts carrier bytes while a value's counts elements, so the two are different units that can match by coincidence: a [M, 32] UINT8 mask over a fully valid source has valid [M, 4], exactly like a [M, 32] value tile with 4 valid columns. Any geometry-based classifier therefore has a case where it guesses wrong and re-compares the carrier bytes, silently discarding the packed predicate bits. Callers holding a value convert it explicitly with tile.cmps(value, 0, cmp_type=1).

Pick the PTO form from the two value branches:

on_true	on_false	Emitted	Why
tile	tile	tile.create(TSEL scratch) + tile.sel(mask, T, F, tmp)	TSEL takes two tile sources directly; nothing is materialized
tile	scalar	tile.create(TSELS scratch) + tile.sels(mask, T, tmp, F)	exactly TSELS's shape, dst = mask ? src : scalar, one instruction
scalar	tile	tile.full(T) + tile.create(TSEL scratch) + tile.sel(mask, full, F, tmp)	TSELS has no scalar-on-true form, so the scalar is materialized and TSEL used

The tile/scalar path queries Backend::PlanOpScratch("tile.sels", request) with the source dtype and physical column count. It uses TSELS only when the plan is Required, the source is rank 2 with static columns, and the scalar dtype can be paired with the source. Otherwise it tries TSEL after materializing the scalar. Both targets reject BF16 for TSELS but accept it for TSEL; A2/A3 rejects 8-bit integers for both forms, so the pass reports an unsupported select instead of emitting an invalid fallback. The registered scratch planner and TSELS validator share the dtype predicate in pto_ops_select.cpp.

tile.full is the only scalar broadcast the tile pipeline can lower: tile.expands has no PTO emitter, so pl.expands itself fails codegen. That constrains a scalar branch to a literal and the result to static extents, and DeduceTileSelectType enforces both so the diagnostic lands on the user's call rather than surfacing as a tile.full error from inside this pass.

tile.full also stamps its result fully valid, and tile.sel reads the result's valid extents from its lhs, so a materialized branch would widen a tail-narrowed select. The rule re-types the created call with the other branch's view — overriding the view only, so the rewritten call still round-trips through print → parse.

Scratch geometry is backend-owned. The registered Backend::PlanOpScratch rules supply it through PassContext::GetBackend(): A2/A3 uses UINT32 [1, 16] for TSEL and sizes TSELS against one complete physical source row in the source dtype (src_dtype [1, N]); A5 uses UINT8 [1, 32] for both. Standalone lowering without a configured backend retains the historical A5 geometry and ungated dtype behavior. These plans describe compiler-created scratch; explicit tile.sel / tile.sels scratch operands remain caller-owned. As with the tile.tquant_mx rule, the scratch tile.create stamps MemorySpace::Vec explicitly because this pass runs before InferTileMemorySpace.

ND operands take the TSEL path. DeduceTileSelsType hard-requires rank 2 — its mask-coverage arithmetic indexes valid_shape[0] and [1] directly — and this pass runs before FlattenTileNdTo2D (slot 13 vs 14), so an ND operand is still ND here. The rule therefore gates TSELS on rank 2 and lets ND fall through to tile.full + tile.sel, neither of which restricts rank; FlattenTileNdTo2D then collapses the result with everything else. Nothing is rejected for rank.

Two tile branches must agree on physical shape, valid extents, and dtype. TSEL reads its two sources element-wise, does not broadcast, and has a single element type for src0 / src1 / dst. Accepting a [1, N] against an [M, N] would promise semantics the instruction cannot deliver — and would make the result depend on which side the narrow operand sat, since the shape would broadcast to [M, N] while the valid extents came from whichever branch was consulted. Promoting an FP16/FP32 pair to an FP32 result would emit a TSEL ptoas rejects with expects src0, src1, and dst to have the same element type. tile.select rejects all three mismatches at deduction and points at the explicit broadcast and cast ops.

Algorithm (sin / cos rule)

LowerSinCos in src/ir/transforms/lower_composite/lower_composite_math.cpp is parameterised on is_cos. The mutator overrides VisitStmt_(const AssignStmtPtr&) (rather than VisitCall) because each trig op expands to ~33 statements and each statement needs a fresh temp Var. Working at the statement level lets the rule append directly to the surrounding sequence via the builder.

Range Reduction (Cody-Waite, 4-part π split)

The goal is to express x = k·π + t (sin) or x = k·π + π/2 + t (cos) with t ∈ [-π/2, π/2] and k an integer. FP32 cannot represent π exactly, so a single x - k·π_fp32 carries a relative error of ~1e-7 per multiplication, which range-reduction error inflates linearly with |k|. Cody-Waite splits π into a fast-rounding head plus three (here four) small corrections so the residual cancellation only loses bits at the finest scales:

π ≈ PI_V2 + PI_C1 + PI_C2 + PI_C3 + PI_C4

t is computed as a chain of subtractions, each consuming one part:

t0 = x  - k_f * PI_V2
t1 = t0 - k_f * PI_C1
t2 = t1 - k_f * PI_C2
t3 = t2 - k_f * PI_C3
t4 = t3 - k_f * PI_C4

For sin, k_f = float(round(x · PI_INV)) using tile.cast mode ROUND (round-to-nearest, ties away from zero). For cos, the rounding is shifted by 0.5 so k represents the multiple of π whose midpoint lies near x:

k_f = float(rint(x · PI_INV + 0.5))   ; mode RINT (round-half-to-even)

The cos path also adds π/2 mid-reduction, split as PI_HALF_HEAD + PI_HALF_TAIL (Cody-Waite again). PI_HALF_HEAD is folded between PI_C1 and PI_C2, PI_HALF_TAIL after PI_C4, so that each addition shares the magnitude scale of the surrounding subtractions and the catastrophic-cancellation regime is shared across all 5+2 corrections.

Sign Computation

Once k is known as a float, the sign is computed without any conditional:

sign = floor(k_f / 2) · 4 + k_f · (-2) + 1
     = (-1)^k

The identity floor(k/2)·4 - 2·k + 1 evaluates to +1 for even k and -1 for odd k. To see this, write k = 2m + r with r ∈ {0, 1}:

floor(k/2) = m
floor(k/2)·4 - 2·k + 1 = 4m - 2(2m + r) + 1 = 1 - 2r

which is +1 when r = 0 and -1 when r = 1. The pass implements this in 6 ops:

half_k     = k_f * 0.5
floor_hk_i = int32(floor(half_k))         ; tile.cast mode FLOOR
floor_hk_f = float(floor_hk_i)
floor_x4   = floor_hk_f * 4.0
neg2_k     = k_f * (-2.0)
sign_pre   = floor_x4 + neg2_k
sign       = sign_pre + 1.0

Horner Polynomial

sin(t) for t ∈ [-π/2, π/2] is approximated by a degree-9 odd polynomial t · P(t²), where:

P(u) = (((R0·u + R1)·u + R2)·u + R3)·u + 1

The leading 1 constant in P(u) corresponds to the t¹ coefficient of the Taylor series, and R3 ≈ -1/6, R2 ≈ 1/120, R1 ≈ -1/5040, R0 ≈ 1/362880 correspond to the higher odd-power coefficients tuned for minimax accuracy on [-π/2, π/2]. Implementation:

t2     = t * t
p_r0   = t2 * R0
p_r1   = p_r0 + R1
p_t2_r1= p_r1 * t2
p_r2   = p_t2_r1 + R2
p_t2_r2= p_r2 * t2
p_r3   = p_t2_r2 + R3
p_t2_r3= p_r3 * t2
p_one  = p_t2_r3 + 1.0
t_p    = t * p_one
out    = sign * t_p

The same polynomial is used for both sin and cos: the cos path differs only in the range reduction, so by the time t enters the polynomial it already lies in [-π/2, π/2] and the polynomial does not need separate coefficients.

Sin vs Cos at a Glance

Step	sin	cos
1. k rounding	round(x · 1/π) (mode ROUND)	rint(x · 1/π + 0.5) (mode RINT)
2. range reduction	x - k·π (4-part)	x - k·π + π/2 (4-part + 2-part π/2)
3. sign	(-1)^k	(-1)^k (same identity, different k)
4. Horner	t · P(t²)	t · P(t²) (same polynomial)
5. result	sign · t · P(t²)	sign · t · P(t²)

Constants

All constants are FP32 literals (the k* literals near the top of src/ir/transforms/lower_composite/lower_composite_math.cpp, matching the framework reference at gitcode.com/cann/pypto:framework/src/interface/tileop/vector/unary.h):

Symbol	C++ literal	Role
PI_INV	0.31830988732818603515625f	1/π (head)
PI_V2	3.140625f	π head (Cody-Waite part 1)
PI_C1	0.0009670257568359375f	π split-1
PI_C2	6.2771141529083251953125e-7f	π split-2
PI_C3	1.21644916362129151821e-10f	π split-3
PI_C4	-1.0290623200529979163e-13f	π split-4
PI_HALF_HEAD	1.57079637050628662109375f	π/2 head (cos only)
PI_HALF_TAIL	-4.371139000189375e-8f	π/2 tail (cos only)
HALF	0.5f	k-pre offset (cos), sign step
M4	4.0f	sign step
NEG2	-2.0f	sign step
ONE	1.0f	sign + Horner constant term
R0	2.604926501e-6f	Horner coeff (degree 9)
R1	-1.980894471e-4f	Horner coeff (degree 7)
R2	8.333049340e-3f	Horner coeff (degree 5)
R3	-1.666665792e-1f	Horner coeff (degree 3)

tile.cast round modes (mirrors src/ir/op/tile_ops/unary.cpp registration):

Symbol	Value	Meaning
kCastModeNone	0	no rounding (typically int → float)
kCastModeRint	1	round-half-to-even
kCastModeRound	2	round-half-away-from-zero
kCastModeFloor	3	round toward -∞

Numerical Properties

• Absolute error: ≤ ~1e-5 over |x| ≤ 2π · 1024 (validated against NumPy by tests/ut/ir/transforms/test_lower_composite_ops_numerical.py). Inside one period, max abs error observed is ~1 ulp ≈ 1.19e-7.
• Range-reduction breakdown: beyond |x| ≈ 2^17, the FP32 representation of x itself loses fractional precision, so range-reduction error dominates regardless of how many π-correction terms are used. The 4-part Cody-Waite split chosen here is the standard CANN/PyPTO recipe and matches the reference implementation's behaviour on every tested x magnitude.
• dtype: FP32-only. FP16, BF16, and integer inputs are rejected at op-construction time (well before the pass runs) — see tests/ut/ir/operators/test_tensor_ops.py (tensor.sin/cos rejection) and tests/ut/ir/operators/test_tile_ops.py (tile.sin/cos rejection) for the rejection cases.
• NaN/Inf: NaN inputs propagate to NaN output (the polynomial preserves NaN). Inf inputs produce indeterminate values because the range-reduction k = round(x/π) step overflows; this matches the documented |x| ≤ 2^17 validity range.

Idempotency

Running LowerCompositeOps twice produces identical IR after the first run: the recipes emit only primitive tile ops, internal tile.tquant_mx_raw / tile.tmov_x2zz, and the distributed primitives listed below. None of those results is itself a registered composite call, so the second invocation visits the body and changes nothing.

The tile.tquant_mx rule emits the non-composite tile.tquant_mx_raw and tile.tmov_x2zz, so MX lowering is idempotent as well.

pld.tensor.* distributed collectives

The pass also lowers the pld.tensor.* family of window-bound distributed collectives. Each collective is a single composite Call that expands into a notify / wait + data-movement recipe, plus a self-clearing epilogue. The data-movement primitive differs by op: allgather and ring allreduce use pld.tile.put (TPUT-based push, auto-chunking through a VEC staging tile), broadcast relocates window data with pld.tile.get (GM→GM copy), while mesh allreduce and reduce_scatter pull peer chunks into a UB tile with pld.tile.remote_load. Allreduce and reduce-scatter select tile.add, tile.maximum, tile.minimum, or tile.mul by ReduceOp. All seven rules share the same self-clearing credit-barrier protocol (shared notify/wait helpers + EmitEpilogueReset) — see Barrier-signal protocol below — so a signal buffer is reusable across back-to-back calls, and even inside for / while / if.

Barrier-signal protocol

Each call uses AtomicAdd(1) credits and call-local Wait(>= g) thresholds that restart at 1. EmitBarrier pairs an announcement with a full wait-all; InCore all_to_all_v instead announces readiness to all peers, waits on demand for each non-self destination receiving nonzero rows, then uses EmitWaitAll for final Ge(2) completion. EmitEpilogueReset retires the call's N credits per non-self cell with AtomicAdd(-N), preserving early next-call credits. With a zero-row all_to_all_v edge, a fast peer can post both next-call credits before a slow rank retires its current two (the cell may reach 4); subtracting 2 retains both. The signal is all-zero after all ranks retire this call only if no next-call credits have arrived.

N and loop-resident wait thresholds may be runtime scalars (pld.system.notify's value only requires ScalarType): mesh allreduce derives each chunk's generation from its loop index. EmitBarrier counts its top-level invocations at compile time; loop and split-readiness schedules emit call-local expectations explicitly.

kGe (not kEq) is the load-bearing choice for every wait predicate: a fast peer can advance a cell past the value a slow rank is looking for before that rank ever polls it, so equality would deadlock. For the same reason kSet must never be mixed with kAtomicAdd on the same cells — a set could clobber an already-advanced counter.

Mesh ([NR, 1], one cell per rank) and ring ([2*(NR-1), NR], one row per round) signals use incompatible cell addressing, so sharing one buffer between the two is rejected as a shape mismatch (ValidateMeshSignalShape) — the only restriction the protocol still imposes.

pld.tensor.allreduce

The allreduce rule starts with a cross-rank ready barrier on shared signal cells (generation 1). It then processes a fully-valid target in UB-sized chunks. Every chunk performs a peer reduction, then barriers on its own call-local generation (2, 3, ...) before storing its result — preventing a fast rank from overwriting bytes a slow rank has not remote-loaded yet. The epilogue then subtracts the call's total credit (1 + chunk_count, built as an IR expression since chunk_count may depend on a runtime extent) back out of every non-self row. The partial-valid single-rectangle path issues exactly two barriers (ready + post-reduce) and its epilogue subtracts 2.

For a fully-valid packed target, mesh lowering creates a logical [1, product(all dimensions)] view and traverses it with physical tiles of at most 16 KiB. A statically known extent smaller than the budget shrinks the chunk to the smallest 32-byte-aligned physical width that covers it, so small allreduces do not reserve a full 16-KiB tile while remaining legal PTO tiles. The tail carries valid_shape=[1, min(chunk, remaining)] through both tile.load and pld.tile.remote_load, so the allocation stays static while the read/store extent is exact. If an ND target carries a partial TensorView.valid_shape, the pass preserves and reduces the representable [rows, cols] rectangle through the established single-rectangle path. Constant valid rectangles use their compact shape; symbolic valid extents fall back to the source's physical rectangle when that statically bounded rectangle fits within one 16-KiB chunk. Oversized partial rectangles, strided targets, DN partial views, and partial boxes that cannot be represented by the leading-dimension collapse are rejected explicitly.

Ring lowering uses one packed 2D view for its reduce-scatter and allgather phases, and moves data with TPUT pushes (pld.tile.put, non-atomic) instead of remote loads. A fully valid target becomes [1, SIZE]; a contiguous partial prefix keeps physical shape [1, product(target.shape)] and carries logical TensorView.valid_shape=[1, product(target.valid_shape)]. FP32 retains balanced floor(i * SIZE / NR) segment boundaries. FP16 rounds each interior boundary up to 16 elements and caps it at SIZE; consequently every non-empty segment and every UB subchunk starts at a 32-byte-aligned address. Per subchunk, each rank first reads its OWN value of the receive slot into a register tile (the slot is stable — own value only — until the left neighbour's push lands), barriers on the ready generation (2k+1), then TPUTs its send subchunk into the RIGHT neighbour's slot of the same index; a push-done barrier (2k+2) precedes the local read + reduce + store of the receive slot. The single shared VEC staging tile is narrowed with tile.set_validshape to each transfer's exact valid_cols, and the push transfer carries that dynamic extent — PTOAS >= v0.55 accepts dynamic partition-view shapes for tput (hw-native-sys/PTOAS#1069), preserving ragged and FP16 tails without padding the window. The non-atomic push + local reduce keeps every ReduceOp (Sum/Max/Min/Prod) working — only a remote-atomic TPUT<AtomicAdd> would be Sum-only. This supports non-divisible inputs and SIZE < NR without inserting holes into the public tensor layout. Every subchunk of every round barriers on a call-local ready + push-done generation pair; the epilogue then subtracts 2 * chunk_count (uniform across rounds, since every round's subchunk loop shares the same bound) from every row of the [2*(NR-1), NR] signal.

Any symbolic target or partial-valid extent that survives lowering must be runtime-bound by a kernel scalar, loop variable, or physical tensor-shape parameter; a type-metadata-only symbol is rejected during PTO codegen. A fully dynamic physical target dimension is bound from that tensor parameter.

Signal buffers are reusable across back-to-back allreduce calls — including a mesh allreduce over a symbolic extent, whose credit total is simply a runtime-computed expression rather than a value the compiler must know. See Barrier-signal protocol above.

Mesh and ring lowering support FP16 and FP32 with ReduceOp::kSum, kMax, kMin, and kProd for arbitrary positive element counts.

pld.tensor.allgather

Signature: allgather(local_data, target, signal). local_data is this rank's chunk (plain Tensor [1, SIZE] on the InCore path), target is a window-bound DistributedTensor[NR, SIZE] staging area that also serves as the result, and signal is the INT32 barrier. Push-based decomposition:

• tile.create([1, stage_cols], dtype=..., target_memory=Vec) — allocate a VEC staging tile for pld.tile.put auto-chunking. For a static extent, stage_cols is min(SIZE, chunk_elements) floored to a 32-byte row boundary when it is below one full chunk; for example, FP32 SIZE=17 produces stage_cols=16, not 17. A transfer shorter than one alignment unit (8 FP32 elements, 32 packed-FP4 byte carriers) has no positive stage width that is both aligned and no larger than the transfer, and remains unsupported. Symbolic extents use the full chunk bound — the same static-UB contract as MakeTputStageShape / chunk_cols, so a runtime width of 17 still allocates 4096 FP32 elements. The geometry uses physical storage bits, so a 16-KiB stage holds 4096 FP32 or 16384 packed-FP4 byte carriers. The stage is a bounded bounce buffer, not a copy of the transfer: pto-isa reads the full extent from the partition views and 2-D-slides the transfer through it, so sizing the stage from SIZE would only waste UB (a [1, 65537] FP32 stage is 256 KiB and overflows the VEC budget). pld.tile.put reads directly from the local_data Tensor source — no explicit tile.load is emitted.
• Phase 1: for peer in 0..NR-1, pld.tile.put(target, peer, local_data, put_stage, [my_rank, 0], [0, 0], [1, SIZE]) — push this rank's chunk into every peer's window at row my_rank. Self-store (peer == my_rank) uses HCCL identity mapping. pld.tile.put auto-chunks when SIZE exceeds the staging-tile capacity
• Phase 2: barrier (generation 1) + epilogue (subtract 1 from every non-self cell)
• Return target — the window IS the gathered [NR, SIZE] result (window-as-result, DistributedTensor)

Compared to the original pull-based allgather (4-arg with a separate out tensor), this push-based variant drops the out parameter and the per-peer pld.tile.get gather loop. Total HBM drops from (NR+1)×SIZE to NR×SIZE, at the cost of the window remaining occupied until the caller consumes the result.

pld.tensor.reduce_scatter

Decomposes into the same phase shape as allreduce's rectangle path:

• Phase 2a/2b: ready barrier (generation 1)
• Phase 3: for each peer, remote_load chunk r from peer p and accumulate into a local scratch with the ReduceOp-matching tile op (tile.add / tile.maximum / tile.minimum / tile.mul)
• Phase 3.5a/3.5b: post-reduce barrier (generation 2)
• Phase 4: tile.store the reduced chunk r back into target[r, 0:SIZE]
• Epilogue: subtract 2 from every non-self cell

target has shape [NR, SIZE]; each rank stages all NR chunks before the call. After the call, rank r's row [r, 0:SIZE] holds the element-wise reduction of chunk r across all ranks (Sum/Max/Min/Prod by ReduceOp). The post-reduce barrier is required for the same WAR reason as allreduce.

All four ReduceOps are supported — kSum, kMax, kMin, and kProd — routed through the shared Reduce() helper, the same dispatch the allreduce rules use.

pld.tensor.broadcast

Decomposes into a 3-phase recipe:

• Phase 2: barrier (generation 1)
• Phase 3: tile.create (VEC staging tile, capped at one 16-KiB chunk like allgather's) + pld.tile.get(target, peer=root, target, stage) on every rank — each rank reads root's slice into its own target. For peer == root the HCCL identity mapping makes the get a local no-op, so root keeps its own data while non-root ranks receive root's. Because the stage no longer derives from the target's extent, a dynamic target shape is accepted — the dynamic dim simply takes the chunk bound.
• Epilogue: subtract 1 from every non-self cell

root is a static int kwarg known at compile time.

pld.tensor.all_to_all

Symmetric push: each rank writes input[dest, :] into every peer window at row my_rank via pld.tile.put, then barriers. The shared VEC stage is the same capped [stage_rows, stage_cols] bounce buffer as allgather.

pld.tensor.all_to_all_v

Variable-size push (MPI_Alltoallv pattern). Each destination transfers clamp(send_counts[dest], 0, MAX_RECV) rows through the same capped 2-D stage; target.shape[0] % NR == 0 stays load-bearing so the receiver can index sender s at row s * MAX_RECV. The runtime row count is a dynamic partition-view dim; the stage itself stays a static chunk bound.

Each receiver atomically resets its own recv_counts cells and announces readiness to every peer. In the destination loop, the sender waits for each non-self destination at generation 1 before publishing its positive clamped count with AtomicAdd and pushing payload. It need not wait for every destination before starting a send. Zero-row destinations retain the reset count without a publication or payload transfer. Generation 2 completes the exchange and the epilogue subtracts both credits. Atomic reset and publication avoid whole-cache-line scalar write-back overwriting adjacent counts or signals; destination readiness protects consumers when the same windows are reused.

pld.tensor.barrier

Pure synchronization — no data movement. Decomposes into a barrier (generation 1) plus its epilogue (subtract 1 from every non-self cell).

The returned expression is the same signal tensor, enabling the rebind idiom signal = pld.tensor.barrier(signal). Under the self-clearing protocol every call restarts at generation 1 regardless of how many times the signal has already been used, so the rebind no longer needs to chain onto any prior state.

Implementation Notes

The mutator overrides VisitStmt_(const AssignStmtPtr&) rather than VisitCall because the decomposition splices ~33 statements per trig op into the surrounding sequence. Doing the splice from inside VisitCall would require returning multiple expressions, which IRMutator does not support; doing it from VisitStmt_ lets LowerSinCos build a vector<StmtPtr> and return either a single bound AssignStmt or a fresh SeqStmts.

Each intermediate result is bound to a fresh Var named via auto_name::BuildName with the user's target name as the base. The mutator's temp_counter_ is shared (by reference, through each LoweringBuilder) across all trig ops in a function so distinct ops do not collide on temp names.

The cast modes RINT (cos), ROUND (sin), FLOOR (sign), and None (int↔float) come from the tile-op registry's enum (src/ir/op/tile_ops/unary.cpp). Choosing the correct mode is load-bearing: ROUND for sin's k keeps k symmetric around zero so the Horner polynomial sees evenly distributed t; RINT for cos's k matches the +0.5 shift and ensures even k corresponds to even multiples of π/2.

Related

• Issue: #1289 — Add FP32-only tile.sin / tile.cos and a lowering pass.
• Reference implementation: gitcode.com/cann/pypto:framework/src/interface/tileop/vector/unary.h — the upstream CANN/PyPTO recipe whose constants and op-sequence this pass mirrors verbatim.
• Op deducer: DeduceTileFP32OnlyType in src/ir/op/tile_ops/unary.cpp:94 — enforces FP32-only at op-construction time.
• Conversion registry: RegisterSimple("tensor.sin", "tile.sin") and the cos counterpart in src/ir/transforms/op_conversion_registry.cpp — the upstream tensor-to-tile rewrite that produces the tile.sin / tile.cos calls this pass consumes.
• Tests: tests/ut/ir/transforms/test_lower_composite_ops.py (structural), tests/ut/ir/transforms/test_lower_composite_ops_numerical.py (NumPy-reference numerical).
• MX quantization tests: tests/ut/codegen/test_quant_mx_codegen.py (tuple consumption and memory planning).