[RFC] Move LowerCompositeOps before ConvertTensorToTileOps

[georgebisbas]

Summary

Propose moving LowerCompositeOps (slot 14) to run before ConvertTensorToTileOps (slot 12), so composite collectives can emit clean tensor-level primitives like pld.tensor.get/pld.tensor.put and let the existing conversion registry handle flattening + VEC staging tile allocation uniformly.

Problem

LowerCompositeOps currently runs after ConvertTensorToTileOps. Because tensor-to-tile conversion has already run, composite lowerings cannot emit tensor-level ops — they must duplicate the conversion logic: compute [rows, cols] from N-D shapes, allocate a tile.create, and emit pld.tile.get/pld.tile.put directly.

This duplication bit us in feat/relax-get-dst: relaxing pld.tile.get's dst to accept plain Tensor required fixing the verifier, codegen, and DSL wrappers — but the reviewer pointed out the same fix was missing in op_conversion_registry.cpp. The duplicated flattening logic in LowerCompositeOps's LowerTensorBroadcastRule must now also be kept in sync with any future type-constraint changes.

Current pass ordering

Slot	Pass
12	ConvertTensorToTileOps
13	OptimizeOrchTensors
14	LowerCompositeOps

LowerCompositeOps has empty PassProperties — requires nothing, produces nothing — suggesting no hard dependency on the output of passes 12-13.

Proposed

Slot	Pass
~12	LowerCompositeOps ← moved earlier
12	ConvertTensorToTileOps ← runs after, picks up tensor-level get/put

With this, composites could emit pld.tensor.get(dst, peer, src) and let ConvertTensorToTileOps handle the tile.create + pld.tile.get uniformly.

Open questions

1. Does LowerCompositeOps depend on any IRProperty or IR shape from ConvertTensorToTileOps? The allreduce lowering emits tile.load/tile.store — those would need rewriting to tensor-level ops.

2. What does OptimizeOrchTensors do to orchestration-level variables? Does LowerCompositeOps need its output?

3. Composite ops currently use tile.remote_load directly. Full benefit requires rewriting all composite lowerings to emit tensor-level primitives.

References

• src/ir/transforms/lower_composite_ops_pass.cpp — all lowering rules
• src/ir/transforms/op_conversion_registry.cpp — pld.tensor.get/put → tile.create + pld.tile.get/put
• docs/en/dev/passes/12-convert_tensor_to_tile_ops.md
• docs/en/dev/passes/14-lower_composite_ops.md
• Notes: pypto-3.0-notes/pr_plans/36-move-lower-composite-ops-before-convert-tensor-to-tile.md

[georgebisbas]

Other thoughts with help from DeepSeek:

The duplication between op_conversion_registry.cpp's pld.tensor.get conversion and lower_composite_ops_pass.cpp's LowerTensorBroadcastRule is real and will cause bugs.

Previous idea says "move LowerCompositeOps (slot 14) before ConvertTensorToTileOps (slot 12)." But LowerCompositeOps is a mixed pass -- it handles two fundamentally different categories of ops that have different placement requirements:

Ops lowered	Current placement	Reason
tile.sin / tile.cos	After conversion (slot 14)	These are tile-level trig decompositions. ConvertTensorToTileOps must first convert tensor.sin/tensor.cos to tile.sin/tile.cos. Moving before conversion means the mutator never finds tile.sin calls.
pld.tensor.allreduce	After conversion (slot 14)	Emits tile.load, tile.store, pld.tile.remote_load -- tile-level ops. This works because ConvertTensorToTileOps already converted the params to TileType. Before conversion, tile.load(target) would see a DistributedTensor param, and the entry-load insertion logic in ConvertTensorToTileOps Phase 1 Step 2 could conflict with the allreduce's explicit load.
broadcast, allgather, barrier, reduce_scatter (PR #1782)	After conversion (slot 14)	Emit tile.create + pld.tile.get/put. This is the duplication you identified.

Moving the entire LowerCompositeOps pass earlier would break tile.sin/tile.cos and risk subtle issues with allreduce. The right approach is more targeted.

A better proposal: Three options, ranked by risk

Option A (Lowest risk -- recommend): Extract a shared utility function.

Instead of moving passes, factor the flattening + staging tile logic into a shared utility (e.g. distributed_utils::EmitVecStagingTile(shape, dtype, span)) that both op_conversion_registry.cpp and lower_composite_ops_pass.cpp call. This eliminates duplication without touching pass ordering.

File: new helper in src/ir/transforms/distributed_utils.cpp (or alongside tile_conversion_utils):

struct VecStage { VarPtr stage_var; AssignStmtPtr create_stmt; };

VecStage EmitVecStagingTile(
    const std::vector<ExprPtr>& transfer_shape, DataType dtype, const Span& span);

Both sites call this and get identical behavior by construction.

Option B (Medium risk): Split LowerCompositeOps into a tensor-level pass and a tile-level pass.

• LowerTensorCollectives (new, before slot 12): lowers broadcast/allgather/barrier/reduce_scatter to tensor-level ops (pld.tensor.get/put, pld.system.notify/wait). ConvertTensorToTileOps then handles the pld.tensor.get → tile.create + pld.tile.get uniformly.
• LowerCompositeOps (stays at slot 14): continues handling tile-level composites (tile.sin/tile.cos, pld.tensor.allreduce).

Option C (Highest risk -- the current proposal): Move the entire pass. Requires rewriting the allreduce lowering to emit tensor-level ops (tensor.load, tile.remote_load, tensor.store), and handling the fact that tile.sin/tile.cos decomposition must still happen after conversion (the ops won't exist yet).

Recommendation

Option A can be a the best now-- minimal change, zero risk of regressing the sin/cos pipeline or the allreduce path, and it solves the root cause (duplication). Option B is worth tracking as a future cleanup once the collective op surface stabilizes.