diff --git a/AGENTS.md b/AGENTS.md index e45d878..99de7f5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,4 +1,5 @@ * Always read the full README.md before doing anything +* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes. * Build commands: * `cmake --build ./build_release` * `cmake --build ./build_debug` diff --git a/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md b/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md new file mode 100644 index 0000000..95f0f87 --- /dev/null +++ b/invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md @@ -0,0 +1,362 @@ +# Graph Compute Batch Physical-Fragment Invariant + +## Status + +This document is **normative** for Raptor's Spatial graph IR. + +Every developer or coding agent modifying Spatial graph construction, graph +verification, Blueprint handling, or `MergeComputeNodes` must read this file +after `README.md` and `AGENTS.md`. + +`AGENTS.md` must contain this instruction: + +```text +* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes. +``` + +## Scope + +This invariant applies to: + +- `spat.graph_compute_batch`; +- graph-level values produced by it; +- `tensor.parallel_insert_slice` operations that publish its lane results; +- `spat.blueprint` operations that describe logical reconstruction; +- graph analyses and transformations that consume those values; +- the graph-to-scheduled transition in `MergeComputeNodes`. + +It does **not** impose the same representation on: + +- `spat.scheduled_compute`; +- `spat.scheduled_compute_batch`; +- `pim.core` or `pim.core_batch`; +- values whose cross-core movement is already represented by explicit + `spat.channel_send` and `spat.channel_receive` operations. + +Scheduled IR represents execution on assigned cores. Communication and value +availability there are defined by local SSA forwarding and explicit +send/receive operations, not by the graph physical-fragment invariant. + +## Core invariant + +For every result of a `spat.graph_compute_batch` with `N` graph lanes: + +1. Every graph lane produces exactly one fragment for that result. +2. All lanes produce fragments with the same exact ranked tensor type `F`. +3. The graph result is a physical collection of those fragments with type: + + ```text + tensor + ``` + + Conceptually, the result is `N × F`: one leading physical fragment-slot + dimension followed by the complete per-lane fragment shape. +4. Physical slot `i` identifies a fragment publication. It does not, by itself, + identify a row, column, channel, tile, or any other logical tensor position. +5. The result type carries no logical reconstruction order. + +The leading dimension is therefore a **physical fragment-slot dimension**, not +a logical tensor dimension. + +## Per-lane computation is unrestricted + +The invariant constrains the published result representation, not what a lane +may compute. + +A graph lane may: + +- read several input slices; +- perform reductions; +- add or combine multiple columns; +- execute matrix/vector operations; +- produce a fragment that corresponds to any logical region; +- participate in a multi-stage or logarithmic reduction tree implemented by + following `spat.graph_compute` or `spat.graph_compute_batch` operations. + +Arithmetic combination is graph computation. `spat.blueprint` is not an +arithmetic reduction operation. + +### Example: `16×4 -> 16×2` + +Two graph lanes may compute: + +```text +lane 0: input[:, 0] + input[:, 1] -> tensor<16x1> +lane 1: input[:, 2] + input[:, 3] -> tensor<16x1> +``` + +The physical graph result is: + +```text +tensor<2x16x1> +``` + +A Blueprint then maps: + +```text +physical slot 0 -> logical output[:, 0:1] +physical slot 1 -> logical output[:, 1:2] +``` + +and describes the logical result `tensor<16x2>`. + +For a larger reduction, following graph compute batches may reduce fragments in +`ceil(log2(N))` stages. Every intermediate batch still publishes a physical +`batch × fragment` collection. + +## Physical publication inside `spat.graph_compute_batch` + +The batch body must publish each lane's fragment into the physical result. + +For one result with fragment type `F`, the corresponding +`tensor.parallel_insert_slice` must insert the fragment into one slot of the +physical `N × F` destination: + +```text +physical offsets = [slot, 0, 0, ...] +physical sizes = [1, shape(F)...] +physical strides = [1, 1, 1, ...] +``` + +The slot may be the graph lane directly or a statically analyzable permutation +of it. The insertion describes physical slot placement only. It must not use a +logical output dimension as the physical batch dimension. + +For each graph result, the body must contain exactly one physical publication +per graph lane. Since the body executes once per lane, this normally means one +`tensor.parallel_insert_slice` operation targeting that result. + +## Logical reconstruction + +Logical reconstruction is separate from physical publication. + +The reconstruction descriptor defines, for every physical fragment slot: + +- which physical batch operand owns the fragment; +- which physical slot contains it; +- its destination offsets in the logical tensor; +- its destination sizes; +- its destination strides; +- coverage and conflict policy where relevant. + +The persistent owner of this information is `spat.blueprint` or an equivalent +explicit graph-level reconstruction operation. + +A logical consumer must not infer reconstruction from the physical tensor type +or assume that physical slot order equals logical order. + +The logical mapping may be arbitrary. For example: + +```text +physical slot 0 -> logical row 13 +physical slot 1 -> logical row 4 +physical slot 2 -> logical row 10 +``` + +The physical result remains a regular `batch × fragment` tensor. + +## Relationship between `parallel_insert_slice` and Blueprint + +During graph construction, an algorithm may naturally describe logical +placement with `tensor.parallel_insert_slice` geometry. Before the graph is in +its canonical form: + +1. that geometry must be separated from physical fragment publication; +2. the graph batch result must be normalized to `N × F`; +3. the logical insertion geometry must be transferred to a persistent + `spat.blueprint` reconstruction descriptor. + +After normalization: + +- `parallel_insert_slice` inside `spat.graph_compute_batch` publishes into + physical fragment slots; +- `spat.blueprint` describes reconstruction into the logical tensor. + +The original graph operation may be erased only after all reconstruction +information needed by later stages has a persistent owner. + +## Blueprint semantics + +Blueprint is placement/reconstruction metadata. It may: + +- concatenate fragments; +- reorder fragments; +- insert fragments into arbitrary disjoint logical regions; +- describe complete or partial logical coverage; +- expose a logical tensor view when materialization is required. + +Blueprint must not silently perform arithmetic such as addition, multiplication, +maximum, or reduction. Such transformations must be represented by following +`spat.graph_compute` or `spat.graph_compute_batch` operations. + +A Blueprint consuming a physical fragment batch must explicitly identify the +physical source slot for every logical fragment. It must not derive that slot +from operand order unless that convention is explicitly represented and +verified. + +## Multiple results + +A `spat.graph_compute_batch` may have several results. + +For each result `r` independently: + +- every lane produces one fragment of type `F_r`; +- the graph result type is `N × F_r`; +- its physical publication and logical reconstruction descriptor are verified + independently. + +Different results may use different fragment shapes. + +## Graph consumers + +A graph consumer of a batch result may: + +1. consume fragments directly as physical fragments; +2. select one or more physical slots in a `spat.deferred_communication` body; +3. use a Blueprint to obtain or describe a logical reconstruction; +4. feed fragments to following graph computes or graph compute batches. + +A consumer must not treat the leading physical slot dimension as a logical +model dimension unless an explicit graph operation intentionally performs such +an interpretation. + +All constant selection, slicing, reshaping, concatenation, and other +compile-time shaping needed for a scheduled consumer must be encoded inside the +corresponding `spat.deferred_communication` body. Phase 2 must not recover +missing graph semantics by inspecting consumers after the deferred operation. + +## Graph lane, scheduled lane, and physical core are different identities + +These concepts must never be conflated: + +- **graph lane**: the lane of the original `spat.graph_compute_batch`; +- **physical fragment slot**: the slot in the graph batch result; +- **scheduled lane**: one lane of a `spat.scheduled_compute_batch` equivalence + class; +- **physical core**: the core selected by PEFT. + +The graph batch body or its Blueprint defines graph-lane-to-fragment-slot and +fragment-slot-to-logical-region mappings. + +PEFT defines graph-instance-to-core placement. + +Scheduled communication defines how values move between cores. + +## Scheduled IR exclusion + +Do not add a verifier requiring `spat.scheduled_compute_batch` results to have +`laneCount` as their first dimension. + +Do not rewrite scheduled values merely to resemble graph physical fragment +collections. + +When lowering graph IR into scheduled IR: + +- resolve graph fragments and reconstruction metadata before erasing their + graph owners; +- create local forwarding or `spat.channel_send`/`spat.channel_receive` for + cross-core dependencies; +- allow scheduled result representation to follow the scheduled IR contract; +- preserve numerical and deadlock correctness. + +The graph invariant is an input contract for scheduling, not a scheduled-value +layout contract. + +## Required verifier properties + +`spat.graph_compute_batch` verification must establish, for every result: + +1. the result is a static or otherwise supported ranked tensor; +2. result rank is exactly `fragment rank + 1`; +3. result dimension 0 equals `laneCount`; +4. every lane publication source has the same exact fragment type; +5. the physical insertion targets the corresponding result block argument; +6. physical insertion offsets have the fragment slot in dimension 0; +7. all remaining physical offsets are zero; +8. physical sizes are `[1] + fragment shape`; +9. physical strides are unit; +10. exactly one publication is defined for each graph result in the per-lane + body. + +These checks apply only to `spat.graph_compute_batch`, not to +`spat.scheduled_compute_batch`. + +Blueprint verification must establish that every logical reconstruction entry: + +- references an existing physical batch operand; +- references a valid physical fragment slot; +- maps a fragment compatible with the declared logical slice; +- stays within logical bounds; +- follows the declared conflict and coverage policies. + +## Invalid representations + +The following are invariant violations. + +### Logical aggregate returned directly by graph batch + +```text +laneCount = 16 +result = tensor<1x4x16x16> +``` + +with each lane inserting into logical dimension 2. + +This is a logical assembly masquerading as a graph batch result. The graph +result must instead be `16 × per-row-fragment`, and a Blueprint must describe +placement into `tensor<1x4x16x16>`. + +### Physical storage derived from logical destination shape + +Code equivalent to: + +```cpp +shape = logicalDestinationType.getShape(); +shape[logicalInsertionDimension] = laneCount; +``` + +is invalid. + +Physical graph storage must be derived from the per-lane fragment type: + +```cpp +physicalShape = [laneCount] + fragmentType.getShape(); +``` + +### Reconstruction inferred from result type + +It is invalid to assume that physical slot `i` belongs at logical offset `i`. +The Blueprint or another explicit reconstruction descriptor must state the +mapping. + +### Blueprint used for arithmetic + +It is invalid to encode `fragment0 + fragment1` as Blueprint reconstruction. +Create a following graph compute or graph compute batch for the addition. + +## Ownership + +- ONNX-to-Spatial lowering owns creation of valid graph fragment batches. +- Graph canonicalization owns normalization of temporary logical-assembly forms + into physical graph batches plus Blueprints. +- `spat.graph_compute_batch` verifier rejects invalid physical publications. +- `spat.blueprint` owns persistent logical reconstruction metadata. +- Deferred communication Phase 1 owns complete consumer-side constant shaping. +- Merge scheduling consumes this graph contract and introduces explicit + communication. +- Scheduled IR verifiers validate scheduled execution and communication, not + the graph fragment representation. + +## No repair downstream + +If graph IR violates this invariant, fix the graph producer or graph +canonicalization. + +Do not repair an invalid graph batch by: + +- guessing a lane dimension in `MergeComputeNodes`; +- deriving physical storage from a logical destination tensor; +- inspecting deferred-result users; +- reconstructing omitted Blueprint data after graph erasure; +- weakening graph verifiers; +- imposing the graph representation on scheduled operations. diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp index da4f5b0..28cc660 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp @@ -346,6 +346,52 @@ inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewr mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides); } +inline void publishGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter, + mlir::Location loc, + mlir::Value fragment, + mlir::Value output, + mlir::Value physicalSlot) { + auto fragmentType = mlir::cast(fragment.getType()); + mlir::SmallVector offsets {physicalSlot}; + mlir::SmallVector sizes {rewriter.getIndexAttr(1)}; + mlir::SmallVector strides {rewriter.getIndexAttr(1)}; + for (int64_t dim : fragmentType.getShape()) { + offsets.push_back(rewriter.getIndexAttr(0)); + sizes.push_back(rewriter.getIndexAttr(dim)); + strides.push_back(rewriter.getIndexAttr(1)); + } + createParallelInsertSliceIntoBatchOutput(rewriter, loc, fragment, output, offsets, sizes, strides); +} + +inline mlir::FailureOr +extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter, + mlir::Location loc, + mlir::Value physicalBatch, + mlir::OpFoldResult slot, + mlir::RankedTensorType fragmentType) { + if (fragmentType.getRank() == 0) + return mlir::failure(); + auto physicalType = mlir::dyn_cast(physicalBatch.getType()); + if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1) + return mlir::failure(); + mlir::SmallVector selectedShape {1}; + llvm::append_range(selectedShape, fragmentType.getShape()); + auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding()); + mlir::SmallVector offsets {slot}; + mlir::SmallVector sizes {rewriter.getIndexAttr(1)}; + mlir::SmallVector strides {rewriter.getIndexAttr(1)}; + for (int64_t dim : fragmentType.getShape()) { + offsets.push_back(rewriter.getIndexAttr(0)); + sizes.push_back(rewriter.getIndexAttr(dim)); + strides.push_back(rewriter.getIndexAttr(1)); + } + mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides); + mlir::SmallVector reassociation {{0, 1}}; + for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim) + reassociation.push_back({dim}); + return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult(); +} + template mlir::Value materializeOrComputeUnary(mlir::Value input, mlir::RankedTensorType resultType, diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp index 0553d32..90d04a4 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp @@ -19,9 +19,7 @@ RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) { } RankedTensorType getRowStripStorageType(RankedTensorType logicalType) { - return RankedTensorType::get({logicalType.getDimSize(2), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)}, - logicalType.getElementType(), - logicalType.getEncoding()); + return spatial::getGraphBatchPhysicalResultType(logicalType.getDimSize(2), getRowStripFragmentType(logicalType)); } std::pair, SmallVector> buildRowStripMetadata(RankedTensorType type) { @@ -39,29 +37,12 @@ std::pair, SmallVector> buildRowStripMetadata(Rank return {offsets, sizes}; } -SmallVector buildRowStripFragmentOffsets(PatternRewriter& rewriter, OpFoldResult row) { - return {row, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; -} - -SmallVector buildRowStripFragmentSizes(PatternRewriter& rewriter, RankedTensorType logicalType) { - return {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(logicalType.getDimSize(1)), - rewriter.getIndexAttr(1), - rewriter.getIndexAttr(logicalType.getDimSize(3))}; -} - Value extractRowStripFragment(Value storage, RankedTensorType logicalType, OpFoldResult row, PatternRewriter& rewriter, Location loc) { - return tensor::ExtractSliceOp::create(rewriter, - loc, - getRowStripFragmentType(logicalType), - storage, - buildRowStripFragmentOffsets(rewriter, row), - buildRowStripFragmentSizes(rewriter, logicalType), - getUnitStrides(rewriter, 4)); + return *extractGraphBatchPhysicalFragment(rewriter, loc, storage, row, getRowStripFragmentType(logicalType)); } void insertRowStripFragment(Value fragment, @@ -70,13 +51,11 @@ void insertRowStripFragment(Value fragment, OpFoldResult row, PatternRewriter& rewriter, Location loc) { - createParallelInsertSliceIntoBatchOutput(rewriter, - loc, - fragment, - output, - buildRowStripFragmentOffsets(rewriter, row), - buildRowStripFragmentSizes(rewriter, logicalType), - getUnitStrides(rewriter, 4)); + assert(fragment.getType() == getRowStripFragmentType(logicalType)); + assert(output.getType() == getRowStripStorageType(logicalType)); + auto slot = dyn_cast(row); + assert(slot && "row-strip graph publication requires a dynamic physical slot"); + publishGraphBatchPhysicalFragment(rewriter, loc, fragment, output, slot); } FailureOr createPerChannelConstantFragment(DenseElementsAttr denseAttr, @@ -145,30 +124,23 @@ FailureOr createRowStripStorageFromRows(Value rows, } FailureOr -materializeRowStripStorageToDense(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) { +createRowStripAssemblyBlueprint(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) { auto storageType = dyn_cast(storage.getType()); if (!storageType || storageType != getRowStripStorageType(logicalType)) return failure(); - auto batchOp = createSpatComputeBatch( - rewriter, loc, TypeRange {logicalType}, logicalType.getDimSize(2), {}, ValueRange {storage}, - [&](detail::SpatComputeBatchBodyArgs args) { - Value fragment = extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc); - createParallelInsertSliceIntoBatchOutput(rewriter, - loc, - fragment, - args.outputs.front(), - SmallVector {rewriter.getIndexAttr(0), - rewriter.getIndexAttr(0), - args.lane, - rewriter.getIndexAttr(0)}, - buildRowStripFragmentSizes(rewriter, logicalType), - getUnitStrides(rewriter, 4)); - return success(); - }); - if (failed(batchOp)) - return failure(); - return batchOp->getResult(0); + auto [offsets, sizes] = buildRowStripMetadata(logicalType); + int64_t height = logicalType.getDimSize(2); + SmallVector operandIndices(height, 0), sourceSlots, sourceOffsets(height, 0), strides(height * 4, 1); + for (int64_t row = 0; row < height; ++row) + sourceSlots.push_back(row); + return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, storage, ValueRange {}, + rewriter.getStringAttr("nchw"), rewriter.getStringAttr("nchw_row_strip"), + rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes), + rewriter.getStringAttr("nchw_row_strip_fragments"), rewriter.getStringAttr("fragment_assembly"), + rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots), + rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides), + rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput(); } FailureOr diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp index 3a436ba..abb850a 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp @@ -50,7 +50,7 @@ mlir::FailureOr createRowStripStorageFromRows(mlir::Value rows, mlir::PatternRewriter& rewriter, mlir::Location loc); -mlir::FailureOr materializeRowStripStorageToDense(mlir::Value storage, +mlir::FailureOr createRowStripAssemblyBlueprint(mlir::Value storage, mlir::RankedTensorType logicalType, mlir::PatternRewriter& rewriter, mlir::Location loc); diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp index 2eb1cdb..6125212 100644 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -75,7 +75,7 @@ materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType); if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes)) return failure(); - return materializeRowStripStorageToDense(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc); + return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc); } struct LowerSpatialPlansPass final : PassWrapper> { @@ -277,6 +277,8 @@ struct LowerSpatialPlansPass final : PassWrapper(&op)) { + if (std::optional mode = blueprintOp.getMode(); mode && *mode == "fragment_assembly") + continue; if (blueprintOp.getPhysicalLayout() == kDenseLayout) { rewriter.replaceOp(blueprintOp, blueprintOp.getInput()); continue; @@ -366,11 +368,15 @@ struct LowerSpatialPlansPass final : PassWrapper(op)) return; - if (isa(op) + if (auto blueprint = dyn_cast(op)) { + if (std::optional mode = blueprint.getMode(); mode && *mode == "fragment_assembly") + return; + op->emitOpError("planning blueprint must not remain after LowerSpatialPlans"); + hasIllegalOps = true; + } else if (isa(op) || op->getDialect()->getNamespace() == "onnx") { op->emitOpError("operation must not remain after LowerSpatialPlans"); hasIllegalOps = true; diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp index dc2995e..20a5c6c 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -1448,12 +1448,10 @@ static FailureOr reconstructDepthwiseGemmRows(Value pieces, Value laneIndex = createOrFoldAffineApply( rewriter, tileLoc, (d0 * tiling.totalPatches) + d1, ValueRange {channelTileIndex, patchIndex}, anchorOp); auto rowTileType = RankedTensorType::get({1, tiling.tileOutputChannels}, piecesType.getElementType()); - SmallVector pieceOffsets {laneIndex, rewriter.getIndexAttr(0)}; - SmallVector pieceSizes {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(tiling.tileOutputChannels)}; - Value rowTile = tensor::ExtractSliceOp::create( - rewriter, tileLoc, rowTileType, piecesArg, pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2)); - Value rowNext = insertOutputTile(rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc); + FailureOr rowTile = extractGraphBatchPhysicalFragment(rewriter, tileLoc, piecesArg, laneIndex, rowTileType); + if (failed(rowTile)) + return failure(); + Value rowNext = insertOutputTile(*rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc); tileYielded.push_back(rowNext); return success(); }); @@ -1556,8 +1554,9 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& auto gemmOutType = RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType()); - auto piecesType = RankedTensorType::get({tiling->totalPatches * tiling->numChannelTiles, tiling->tileOutputChannels}, - state.outType.getElementType()); + auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType()); + auto piecesType = spatial::getGraphBatchPhysicalResultType( + tiling->totalPatches * tiling->numChannelTiles, rowTileType); auto paddedInputType = cast(paddedInput.getType()); auto inputTileType = RankedTensorType::get({1, tiling->channelsPerTile, state.wType.getDimSize(2), state.wType.getDimSize(3)}, @@ -1626,7 +1625,6 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& *tiling, rewriter, loc); - auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType()); Value rowTile = spatial::SpatVMMOp::create(rewriter, loc, rowTileType, weightTile, inputTile).getResult(); if (args.inputs.size() > 1) { Value biasArg = pickInputByRank(/*rank=*/2); @@ -1637,11 +1635,7 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& Value biasTile = tiling->numChannelTiles == 1 ? biasArg : createBiasTile(biasArg, channelTileIndex, *tiling, rewriter, loc); rowTile = spatial::SpatVAddOp::create(rewriter, loc, rowTileType, rowTile, biasTile).getResult(); } - SmallVector outputOffsets {args.lane, rewriter.getIndexAttr(0)}; - SmallVector outputSizes {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(tiling->tileOutputChannels)}; - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, rowTile, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 2)); + publishGraphBatchPhysicalFragment(rewriter, loc, rowTile, args.outputs.front(), args.lane); return success(); }); if (failed(batchOp)) @@ -1650,14 +1644,10 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& auto nhwcType = RankedTensorType::get( {state.xType.getDimSize(0), state.outType.getDimSize(2), state.outType.getDimSize(3), state.outType.getDimSize(1)}, state.outType.getElementType()); - Value collectedRows = batchOp->getResult(0); - if (tiling->numChannelTiles != 1) { - auto reconstructedRows = - reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc); - if (failed(reconstructedRows)) - return failure(); - collectedRows = *reconstructedRows; - } + auto reconstructedRows = reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc); + if (failed(reconstructedRows)) + return failure(); + Value collectedRows = *reconstructedRows; return createCollectedConvOutput(ValueRange {collectedRows}, state.outType, diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp index 00d5b2b..5605d62 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp @@ -251,10 +251,7 @@ static FailureOr createVmmBatch(Value a, rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides); Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult(); - SmallVector pieceOffsets {args.lane, rewriter.getIndexAttr(0)}; - SmallVector pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())}; - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, unitStrides); + publishGraphBatchPhysicalFragment(rewriter, loc, piece, args.outputs.front(), args.lane); }); if (failed(batchOp)) return failure(); @@ -401,11 +398,7 @@ static FailureOr createVvdmulBatch(Value a, Value bVector = extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc); Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult(); - SmallVector outputOffsets {args.lane, rewriter.getIndexAttr(0)}; - SmallVector scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; - SmallVector unitStrides = getUnitStrides(rewriter, 2); - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, unitStrides); + publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane); }); if (failed(batchOp)) return failure(); @@ -447,15 +440,14 @@ static FailureOr createDynamicGemmOutputCompute(Value scal Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, nestedLoc); Value column = onnx_mlir::affineModConst(rewriter, nestedLoc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp()); - SmallVector scalarOffsets {lane, rewriter.getIndexAttr(0)}; SmallVector scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; SmallVector unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; - Value scalar = tensor::ExtractSliceOp::create( - rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, unitStrides) - .getResult(); + FailureOr scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType); + if (failed(scalar)) + return failure(); if (alpha != 1.0f) { Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc); - scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, scalar, alphaTensor).getResult(); + *scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, *scalar, alphaTensor).getResult(); } if (biasArg) { Value biasScalar = @@ -465,11 +457,11 @@ static FailureOr createDynamicGemmOutputCompute(Value scal biasScalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).getResult(); } - scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, scalar, biasScalar).getResult(); + *scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, *scalar, biasScalar).getResult(); } SmallVector outputOffsets {row, column}; Value outputNext = - tensor::InsertSliceOp::create(rewriter, nestedLoc, scalar, outputAcc, outputOffsets, scalarSizes, unitStrides) + tensor::InsertSliceOp::create(rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, scalarSizes, unitStrides) .getResult(); yielded.push_back(outputNext); return success(); @@ -505,14 +497,13 @@ static Value extractReductionPiece(Value partialPiecesArg, int64_t numOutRows, ConversionPatternRewriter& rewriter, Location loc) { - SmallVector unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; - SmallVector pieceSizes {rewriter.getIndexAttr(numOutRows), - rewriter.getIndexAttr(crossbarSize.getValue())}; + SmallVector unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; + SmallVector pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())}; SmallVector pieceOffsets { - createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0)}; - return tensor::ExtractSliceOp::create( - rewriter, loc, pieceType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides) - .getResult(); + createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; + auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast(crossbarSize.getValue())}, pieceType.getElementType()); + Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides); + return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector {{0, 1}, {2}}); } static Value reducePartialPiecesForHSlice(Value partialPiecesArg, @@ -730,7 +721,7 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp, return failure(); } - auto scalarPiecesType = RankedTensorType::get({laneCount64, 1}, outType.getElementType()); + auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount64, RankedTensorType::get({1, 1}, outType.getElementType())); auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, rewriter, loc); if (failed(batchOp)) return failure(); @@ -802,8 +793,8 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp, return failure(); } - auto partialPiecesType = - RankedTensorType::get({laneCount64, static_cast(crossbarSize.getValue())}, outType.getElementType()); + auto partialPiecesType = spatial::getGraphBatchPhysicalResultType( + laneCount64, RankedTensorType::get({1, static_cast(crossbarSize.getValue())}, outType.getElementType())); auto batchOp = createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc); if (failed(batchOp)) diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp index 56e2090..c981042 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp @@ -398,10 +398,7 @@ static FailureOr createBatchedVmmBatch(Value a, args.weights.front(), bBatchShape, outputBatchShape, batch, kOffset, hOffset, bTileType, rewriter, loc); Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult(); - SmallVector pieceOffsets {args.lane, rewriter.getIndexAttr(0)}; - SmallVector pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())}; - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2)); + publishGraphBatchPhysicalFragment(rewriter, loc, piece, args.outputs.front(), args.lane); }); if (failed(batchOp)) return failure(); @@ -506,10 +503,7 @@ static FailureOr createBatchedVvdmulBatch(Value a, Value bVector = extractDynamicBatchedBColumn( args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc); Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult(); - SmallVector outputOffsets {args.lane, rewriter.getIndexAttr(0)}; - SmallVector scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, getUnitStrides(rewriter, 2)); + publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane); }); if (failed(batchOp)) return failure(); @@ -548,14 +542,13 @@ static FailureOr createBatchedDynamicOutputCompute(Value scalarPieces, Value batchLane = affineModConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp); Value row = affineFloorDivConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp); Value column = affineModConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp); - SmallVector scalarOffsets {lane, rewriter.getIndexAttr(0)}; - SmallVector scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; - Value scalar = tensor::ExtractSliceOp::create( - rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, getUnitStrides(rewriter, 2)); + FailureOr scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType); + if (failed(scalar)) + return failure(); Value expanded = tensor::ExpandShapeOp::create(rewriter, nestedLoc, outputScalarType, - scalar, + *scalar, SmallVector { {0}, {1, 2} @@ -596,10 +589,11 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg, Value kOffset = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), kSlice * numOutRows); Value batchAndHSlice = arith::AddIOp::create(rewriter, loc, batchOffset, hOffset); Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset); - SmallVector offsets {pieceOffset, rewriter.getIndexAttr(0)}; - SmallVector sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(crossbarSize.getValue())}; - return tensor::ExtractSliceOp::create( - rewriter, loc, pieceType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 2)); + SmallVector offsets {pieceOffset, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; + SmallVector sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())}; + auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast(crossbarSize.getValue())}, pieceType.getElementType()); + Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 3)); + return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector {{0, 1}, {2}}); } static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg, @@ -1019,8 +1013,8 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern { if (succeeded(paddedRhs)) { Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc); const int64_t laneCount = plan.batch * plan.m * numKSlices * numOutHSlices; - auto partialPiecesType = RankedTensorType::get({laneCount, static_cast(crossbarSize.getValue())}, - shapeInfo->outType.getElementType()); + auto partialPiecesType = spatial::getGraphBatchPhysicalResultType( + laneCount, RankedTensorType::get({1, static_cast(crossbarSize.getValue())}, shapeInfo->outType.getElementType())); auto batchOp = createBatchedVmmBatch(paddedLhs, *paddedRhs, paddedLhsType, @@ -1061,7 +1055,8 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern { } } const int64_t laneCount = plan.batch * plan.m * plan.n; - auto scalarPiecesType = RankedTensorType::get({laneCount, 1}, shapeInfo->outType.getElementType()); + auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType( + laneCount, RankedTensorType::get({1, 1}, shapeInfo->outType.getElementType())); auto batchOp = createBatchedVvdmulBatch(plan.lhs, plan.lhsBatchShape, plan.rhs, diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp index 200d7a3..bb31c68 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp @@ -139,9 +139,7 @@ static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef } static RankedTensorType getLanePackedKeepdimsType(int64_t laneCount, RankedTensorType leafType) { - SmallVector shape(leafType.getShape().begin(), leafType.getShape().end()); - shape.front() = laneCount; - return RankedTensorType::get(shape, leafType.getElementType(), leafType.getEncoding()); + return spatial::getGraphBatchPhysicalResultType(laneCount, leafType); } static SmallVector getKeptAxes(ArrayRef reducedAxes) { @@ -191,12 +189,9 @@ static FailureOr buildReduceMeanKeepdimsBatch(Value input, SmallVector sliceOffsets; SmallVector sliceSizes; - SmallVector insertOffsets; - SmallVector insertSizes(inputType.getRank(), rewriter.getIndexAttr(1)); SmallVector unitStrides = getUnitStrides(rewriter, inputType.getRank()); sliceOffsets.reserve(inputType.getRank()); sliceSizes.reserve(inputType.getRank()); - insertOffsets.reserve(inputType.getRank()); auto batchOp = createSpatComputeBatch(rewriter, @@ -209,7 +204,6 @@ static FailureOr buildReduceMeanKeepdimsBatch(Value input, size_t keptAxisIndex = 0; sliceOffsets.clear(); sliceSizes.clear(); - insertOffsets.clear(); for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) { if (isReduced) { sliceOffsets.push_back(rewriter.getIndexAttr(0)); @@ -224,14 +218,10 @@ static FailureOr buildReduceMeanKeepdimsBatch(Value input, sliceSizes.push_back(rewriter.getIndexAttr(1)); } - insertOffsets.push_back(args.lane); - insertOffsets.append(inputType.getRank() - 1, rewriter.getIndexAttr(0)); - Value slice = tensor::ExtractSliceOp::create( rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides); Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult(); - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, reduced, args.outputs.front(), insertOffsets, insertSizes, unitStrides); + publishGraphBatchPhysicalFragment(rewriter, loc, reduced, args.outputs.front(), args.lane); }); if (failed(batchOp)) return failure(); @@ -276,10 +266,11 @@ static Value buildKeepdimsFromLanePackedBatch(Value batchValue, if (!pendingLeadingReducedAxes.empty()) expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end()); + if (batchType.getNumElements() != batchType.getDimSize(0)) + return {}; auto reshapeCompute = createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) { - auto flatType = - RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding()); + auto flatType = RankedTensorType::get({batchType.getNumElements()}, batchType.getElementType(), batchType.getEncoding()); Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat); Value compact = flat; if (compactKeptType != flatType) diff --git a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp index 6e20e14..d3de0d6 100644 --- a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp @@ -134,6 +134,7 @@ static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Va nullptr, nullptr, nullptr, + nullptr, nullptr); } diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index ca51339..1dd4f88 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -10,6 +10,8 @@ add_pim_library(SpatialOps Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp + Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp + Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp Transforms/MergeComputeNodes/ScheduledComputeReport.cpp diff --git a/src/PIM/Dialect/Spatial/Spatial.td b/src/PIM/Dialect/Spatial/Spatial.td index 3d884dd..eb4c84a 100644 --- a/src/PIM/Dialect/Spatial/Spatial.td +++ b/src/PIM/Dialect/Spatial/Spatial.td @@ -183,23 +183,10 @@ def SpatBlockYieldOp : SpatOp<"block_yield", [ } def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]> { - let summary = "Temporary scheduled communication placeholder"; + let summary = "Temporary scheduled payload derivation placeholder"; let arguments = (ins - Variadic:$sources, - StrAttr:$exchangeId, - StrAttr:$logicalProducer, - StrAttr:$logicalConsumer, - I64Attr:$sourceClass, - I64Attr:$targetClass, - I64Attr:$sourceCore, - I64Attr:$targetCore, - I64Attr:$sourceLane, - I64Attr:$targetLane, - StrAttr:$transferKind, - I64Attr:$resultIndex, - DenseI64ArrayAttr:$projectedTransfer, - I64Attr:$hostOutputOwner + Variadic:$sources ); let results = (outs @@ -312,6 +299,7 @@ def SpatBlueprintOp : SpatOp<"blueprint", []> { StrAttr:$indexMap, OptionalAttr:$mode, OptionalAttr:$fragmentOperandIndices, + OptionalAttr:$fragmentSourceSlots, OptionalAttr:$fragmentSourceOffsets, OptionalAttr:$fragmentStrides, OptionalAttr:$conflictPolicy, diff --git a/src/PIM/Dialect/Spatial/SpatialOps.cpp b/src/PIM/Dialect/Spatial/SpatialOps.cpp index f27d900..6acd8b5 100644 --- a/src/PIM/Dialect/Spatial/SpatialOps.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOps.cpp @@ -10,6 +10,18 @@ using namespace mlir; namespace onnx_mlir { namespace spatial { + +RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, RankedTensorType fragmentType) { + SmallVector shape {laneCount}; + llvm::append_range(shape, fragmentType.getShape()); + return RankedTensorType::get(shape, fragmentType.getElementType(), fragmentType.getEncoding()); +} + +FailureOr getGraphBatchFragmentType(RankedTensorType physicalType, int64_t expectedLaneCount) { + if (!physicalType || physicalType.getRank() < 1 || physicalType.getDimSize(0) != expectedLaneCount) + return failure(); + return RankedTensorType::get(physicalType.getShape().drop_front(), physicalType.getElementType(), physicalType.getEncoding()); +} namespace { std::optional getBlockArgument(Region& body, unsigned argIdx) { diff --git a/src/PIM/Dialect/Spatial/SpatialOps.hpp b/src/PIM/Dialect/Spatial/SpatialOps.hpp index eff0eaa..222766e 100644 --- a/src/PIM/Dialect/Spatial/SpatialOps.hpp +++ b/src/PIM/Dialect/Spatial/SpatialOps.hpp @@ -30,6 +30,10 @@ namespace onnx_mlir { namespace spatial { +mlir::RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, mlir::RankedTensorType fragmentType); +mlir::FailureOr +getGraphBatchFragmentType(mlir::RankedTensorType physicalType, int64_t expectedLaneCount); + bool isGraphComputeLike(mlir::Operation* op); bool isGraphBatchComputeLike(mlir::Operation* op); bool isScheduledComputeLike(mlir::Operation* op); diff --git a/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp b/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp index 2df4361..d9350f3 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp @@ -576,6 +576,10 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) { printer << " operandIndices "; printCompressedIntegerList(printer, *operandIndices); } + if (std::optional> sourceSlots = getFragmentSourceSlots()) { + printer << " sourceSlots "; + printCompressedIntegerList(printer, *sourceSlots); + } if (std::optional> sourceOffsets = getFragmentSourceOffsets()) { printer << " sourceOffsets "; printCompressedIntegerList(printer, *sourceOffsets); @@ -597,6 +601,7 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) { getIndexMapAttrName().getValue(), getModeAttrName().getValue(), getFragmentOperandIndicesAttrName().getValue(), + getFragmentSourceSlotsAttrName().getValue(), getFragmentSourceOffsetsAttrName().getValue(), getFragmentStridesAttrName().getValue(), getConflictPolicyAttrName().getValue(), @@ -620,6 +625,7 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result) SmallVector fragmentOffsets; SmallVector fragmentSizes; SmallVector fragmentOperandIndices; + SmallVector fragmentSourceSlots; SmallVector fragmentSourceOffsets; SmallVector fragmentStrides; @@ -637,6 +643,9 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result) if (succeeded(parser.parseOptionalKeyword("operandIndices")) && parseCompressedIntegerList(parser, fragmentOperandIndices)) return failure(); + if (succeeded(parser.parseOptionalKeyword("sourceSlots")) + && parseCompressedIntegerList(parser, fragmentSourceSlots)) + return failure(); if (succeeded(parser.parseOptionalKeyword("sourceOffsets")) && parseCompressedIntegerList(parser, fragmentSourceOffsets)) return failure(); @@ -667,6 +676,8 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result) result.addAttribute("mode", mode); if (!fragmentOperandIndices.empty()) result.addAttribute("fragmentOperandIndices", builder.getDenseI64ArrayAttr(fragmentOperandIndices)); + if (!fragmentSourceSlots.empty()) + result.addAttribute("fragmentSourceSlots", builder.getDenseI64ArrayAttr(fragmentSourceSlots)); if (!fragmentSourceOffsets.empty()) result.addAttribute("fragmentSourceOffsets", builder.getDenseI64ArrayAttr(fragmentSourceOffsets)); if (!fragmentStrides.empty()) diff --git a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp index e84f163..758b015 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp @@ -574,20 +574,26 @@ LogicalResult SpatBlueprintOp::verify() { auto stridesAttr = getFragmentStridesAttr(); auto operandIndicesAttr = getFragmentOperandIndicesAttr(); + auto sourceSlotsAttr = getFragmentSourceSlotsAttr(); auto sourceOffsetsAttr = getFragmentSourceOffsetsAttr(); if (!operandIndicesAttr) return emitError("fragment assembly blueprint requires fragment operand indices"); + if (!sourceSlotsAttr) + return emitError("fragment assembly blueprint requires physical fragment source slots"); if (!sourceOffsetsAttr) return emitError("fragment assembly blueprint requires fragment source offsets"); if (!stridesAttr) return emitError("fragment assembly blueprint requires fragment strides"); ArrayRef operandIndices = operandIndicesAttr.asArrayRef(); + ArrayRef sourceSlots = sourceSlotsAttr.asArrayRef(); ArrayRef sourceOffsets = sourceOffsetsAttr.asArrayRef(); ArrayRef strides = stridesAttr.asArrayRef(); if (strides.size() != offsets.size()) return emitError("fragment stride and offset arrays must have the same length"); if (sourceOffsets.size() != operandIndices.size()) return emitError("fragment source offset count must match fragment operand index count"); + if (sourceSlots.size() != operandIndices.size()) + return emitError("fragment source slot count must match fragment operand index count"); if (!getConflictPolicyAttr() || !getCoveragePolicyAttr()) return emitError("fragment assembly blueprint requires conflict and coverage policies"); if (getConflictPolicy() != "disjoint") @@ -622,14 +628,19 @@ LogicalResult SpatBlueprintOp::verify() { int64_t operandIndex = operandIndices[fragmentIndex]; if (operandIndex < 0 || operandIndex >= operandCount) return emitError("fragment assembly operand index is out of range"); + if (sourceSlots[fragmentIndex] < 0) + return emitError("fragment assembly physical source slot must be nonnegative"); if (sourceOffsets[fragmentIndex] < 0) return emitError("fragment assembly source offsets must be nonnegative"); auto operandType = dyn_cast(operands[operandIndex].getType()); if (!operandType || !operandType.hasStaticShape()) return emitError("fragment assembly blueprint requires static ranked tensor operands"); - if (operandType.getRank() != rank) - return emitError("fragment assembly blueprint requires operand/result rank match"); + if (operandType.getRank() != rank + 1) + return emitError("fragment assembly physical operand must have one leading source-slot dimension"); + if (sourceSlots[fragmentIndex] >= operandType.getDimSize(0)) + return emitError("fragment assembly physical source slot is out of range"); + auto fragmentType = RankedTensorType::get(operandType.getShape().drop_front(), operandType.getElementType(), operandType.getEncoding()); SmallVector fragmentOffsets; SmallVector fragmentSizes; @@ -645,12 +656,12 @@ LogicalResult SpatBlueprintOp::verify() { int64_t fragmentElements = 1; for (int64_t dim = 0; dim < rank; ++dim) fragmentElements *= fragmentSizes[dim]; - if (sourceOffsets[fragmentIndex] + fragmentElements > operandType.getNumElements()) - return emitError("fragment assembly source offset exceeds the operand bounds"); + if (sourceOffsets[fragmentIndex] + fragmentElements > fragmentType.getNumElements()) + return emitError("fragment assembly source offset exceeds the selected physical fragment bounds"); SmallVector sourceSliceOffsets = - expandFlatElementIndex(sourceOffsets[fragmentIndex], operandType.getShape()); + expandFlatElementIndex(sourceOffsets[fragmentIndex], fragmentType.getShape()); for (int64_t dim = 0; dim < rank; ++dim) - if (sourceSliceOffsets[dim] + fragmentSizes[dim] > operandType.getDimSize(dim)) + if (sourceSliceOffsets[dim] + fragmentSizes[dim] > fragmentType.getDimSize(dim)) return emitError("fragment assembly source offset must describe a valid unit-stride slice"); for (const auto& [existingOffsets, existingSizes] : slices) { @@ -746,7 +757,8 @@ LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) { Operation* terminator = block.getTerminator(); if (auto yieldOp = dyn_cast_or_null(terminator)) { - if (isScheduled) + auto realized = compute->template getAttrOfType("scheduled.realized"); + if (isScheduled && (!realized || !realized.getValue() || !compute.getBody().hasOneBlock())) return compute.emitOpError("scheduled compute blocks must terminate with spat.block_yield"); llvm::append_range(yieldedTypes, yieldOp->getOperandTypes()); continue; @@ -820,6 +832,16 @@ LogicalResult SpatBlockYieldOp::verify() { LogicalResult SpatDeferredCommunicationOp::verify() { if (getSources().empty()) return emitOpError("requires at least one source"); + static constexpr StringLiteral staleAttributes[] = { + "exchangeId", "logicalProducer", "logicalConsumer", "sourceClass", "targetClass", "sourceCore", + "targetCore", "sourceLane", "targetLane", "transferKind", "resultIndex", "projectedTransfer", + "hostOutputOwner", "source_cpus", "source_classes", "source_lane_ranges", "target_cpus", + "target_classes", "target_lane_ranges", "batched", "source_operand_for_scheduled_lane", + "multi_source_payload"}; + for (StringLiteral name : staleAttributes) + if (getOperation()->hasAttr(name)) + return emitOpError() << "does not accept stale routing attribute '" << name + << "'; source selection and shaping belong in the body and routing is derived in Phase 2"; if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication"))) return failure(); return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication"); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp new file mode 100644 index 0000000..6d98876 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp @@ -0,0 +1,274 @@ +#include "DeferredCommunicationDeadlock.hpp" + +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/DenseSet.h" + +namespace onnx_mlir::spatial { + +using namespace mlir; + +namespace { + +enum class EventKind { Compute, Send, Receive }; + +struct Event { + EventKind kind = EventKind::Compute; + uint64_t exchangeId = 0; +}; + +static LogicalResult simulate(Operation *anchor, + ArrayRef> streams, + StringRef phase) { + SmallVector cursor(streams.size()); + while (true) { + bool allFinished = true; + bool progressed = false; + for (unsigned stream = 0; stream < streams.size(); ++stream) { + if (cursor[stream] == streams[stream].size()) + continue; + allFinished = false; + if (streams[stream][cursor[stream]].kind == EventKind::Compute) { + ++cursor[stream]; + progressed = true; + } + } + if (allFinished) + return success(); + + for (unsigned source = 0; source < streams.size(); ++source) { + if (cursor[source] == streams[source].size()) + continue; + const Event &send = streams[source][cursor[source]]; + if (send.kind != EventKind::Send) + continue; + for (unsigned target = 0; target < streams.size(); ++target) { + if (cursor[target] == streams[target].size()) + continue; + const Event &receive = streams[target][cursor[target]]; + if (receive.kind == EventKind::Receive && receive.exchangeId == send.exchangeId) { + ++cursor[source]; + ++cursor[target]; + progressed = true; + break; + } + } + } + if (!progressed) { + InFlightDiagnostic diagnostic = anchor->emitError() + << phase << " communication rendezvous simulation made no progress"; + unsigned reported = 0; + for (unsigned stream = 0; stream < streams.size() && reported < 8; ++stream) { + if (cursor[stream] == streams[stream].size()) + continue; + const Event &event = streams[stream][cursor[stream]]; + diagnostic << (reported == 0 ? "; blocked " : ", ") << "stream " << stream + << " at exchange " << event.exchangeId; + ++reported; + } + return failure(); + } + } +} + +static std::optional getI64Attr(Operation *op, StringRef name) { + if (auto attr = op->getAttrOfType(name)) + return attr.getInt(); + return std::nullopt; +} + +} // namespace + +LogicalResult verifyPlannedCommunicationDeadlockFree( + Operation *anchor, + unsigned streamCount, + ArrayRef stepCounts, + ArrayRef transfers) { + if (stepCounts.size() != streamCount) + return anchor->emitError("communication plan stream count does not match step counts"); + + SmallVector> streams(streamCount); + SmallVector>> atBoundary(streamCount); + for (unsigned stream = 0; stream < streamCount; ++stream) + atBoundary[stream].resize(stepCounts[stream] + 1); + for (const PlannedCommunicationTransfer &transfer : transfers) { + if (transfer.sourceStream >= streamCount || transfer.targetStream >= streamCount + || transfer.producerStep >= stepCounts[transfer.sourceStream] + || transfer.consumerStep >= stepCounts[transfer.targetStream] + || transfer.sourceInsertionStep > stepCounts[transfer.sourceStream] + || transfer.targetInsertionStep > stepCounts[transfer.targetStream] + || transfer.sourceInsertionStep <= transfer.producerStep + || transfer.targetInsertionStep > transfer.consumerStep) + return anchor->emitError("communication plan references an invalid stream step"); + atBoundary[transfer.sourceStream][transfer.sourceInsertionStep].push_back( + {EventKind::Send, transfer.exchangeId}); + atBoundary[transfer.targetStream][transfer.targetInsertionStep].push_back( + {EventKind::Receive, transfer.exchangeId}); + } + for (unsigned stream = 0; stream < streamCount; ++stream) { + for (unsigned step = 0; step < stepCounts[stream]; ++step) { + llvm::append_range(streams[stream], atBoundary[stream][step]); + streams[stream].push_back({EventKind::Compute, 0}); + } + llvm::append_range(streams[stream], atBoundary[stream].back()); + } + return simulate(anchor, streams, "planned"); +} + +LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) { + DenseMap> operationsByExchange; + struct ParentExchange { + std::optional expectedTransfers; + DenseSet channels; + }; + DenseMap parentExchanges; + DenseMap streamByCore; + SmallVector cores; + bool invalid = false; + funcOp.walk([&](Operation *op) { + if (!isa(op)) + return; + std::optional exchangeId = getI64Attr(op, "raptor.exchange_id"); + if (exchangeId) + operationsByExchange[*exchangeId].push_back(op); + if (auto channels = op->getAttrOfType("raptor.batch_channel_ids")) + for (int64_t channel : channels.asArrayRef()) + operationsByExchange[channel].push_back(op); + if (std::optional parent = getI64Attr(op, "raptor.parent_exchange_id")) { + ParentExchange &group = parentExchanges[*parent]; + std::optional expected = + getI64Attr(op, "raptor.parent_transfer_count"); + if (!expected || *expected <= 0 + || (group.expectedTransfers && group.expectedTransfers != expected)) { + op->emitOpError( + "realized parent exchange has missing or inconsistent transfer count metadata"); + invalid = true; + } else { + group.expectedTransfers = expected; + } + if (exchangeId) + group.channels.insert(*exchangeId); + if (auto channels = + op->getAttrOfType("raptor.batch_channel_ids")) + group.channels.insert(channels.asArrayRef().begin(), + channels.asArrayRef().end()); + } + for (StringRef attrName : {"raptor.source_core", "raptor.target_core"}) + if (std::optional core = getI64Attr(op, attrName); core && !llvm::is_contained(cores, *core)) + cores.push_back(*core); + for (StringRef attrName : {"raptor.batch_source_cores", "raptor.batch_target_cores"}) + if (auto batchCores = op->getAttrOfType(attrName)) + for (int64_t core : batchCores.asArrayRef()) + if (!llvm::is_contained(cores, core)) + cores.push_back(core); + }); + llvm::sort(cores); + for (auto [index, core] : llvm::enumerate(cores)) + streamByCore[core] = index; + + SmallVector> streams(cores.size()); + funcOp.walk([&](Operation *op) { + if (!isa(op)) + return; + auto exchangeId = getI64Attr(op, "raptor.exchange_id"); + if (!exchangeId) { + auto channels = op->getAttrOfType("raptor.batch_channel_ids"); + auto sourceCores = op->getAttrOfType("raptor.batch_source_cores"); + auto targetCores = op->getAttrOfType("raptor.batch_target_cores"); + if (!channels || !sourceCores || !targetCores + || channels.size() != sourceCores.size() || channels.size() != targetCores.size()) { + op->emitOpError("realized compact batch communication is missing channel/core metadata"); + invalid = true; + return; + } + if (isa(op)) + for (auto [channel, sourceCore] : llvm::zip(channels.asArrayRef(), sourceCores.asArrayRef())) + streams[streamByCore.lookup(sourceCore)].push_back( + {EventKind::Send, static_cast(channel)}); + else + for (auto [channel, targetCore] : llvm::zip(channels.asArrayRef(), targetCores.asArrayRef())) + streams[streamByCore.lookup(targetCore)].push_back( + {EventKind::Receive, static_cast(channel)}); + return; + } + auto sourceCore = getI64Attr(op, "raptor.source_core"); + auto targetCore = getI64Attr(op, "raptor.target_core"); + if (!sourceCore || !targetCore) { + op->emitOpError("realized communication is missing core metadata"); + invalid = true; + return; + } + if (isa(op)) + streams[streamByCore.lookup(*sourceCore)].push_back({EventKind::Send, static_cast(*exchangeId)}); + else if (isa(op)) + streams[streamByCore.lookup(*targetCore)].push_back({EventKind::Receive, static_cast(*exchangeId)}); + }); + if (invalid) + return failure(); + for (const auto &entry : parentExchanges) + if (!entry.second.expectedTransfers + || entry.second.channels.size() + != static_cast(*entry.second.expectedTransfers)) + return funcOp.emitOpError() + << "parent exchange " << entry.first + << " does not contain its declared lane transfer set"; + + for (const auto &entry : operationsByExchange) { + if (entry.second.size() != 2 || !isa(entry.second[0]) + == !isa(entry.second[1])) + return funcOp.emitOpError() << "exchange " << entry.first << " does not have exactly one send and one receive"; + auto send = dyn_cast(entry.second[0]); + auto receive = dyn_cast(entry.second[1]); + if (!send) { + send = cast(entry.second[1]); + receive = cast(entry.second[0]); + } + if (send.getInput().getType() != receive.getOutput().getType()) + return send.emitOpError("send and receive payload types do not match"); + int64_t sendSource = 0; + int64_t sendTarget = 0; + if (auto channels = send->getAttrOfType("raptor.batch_channel_ids")) { + auto channelIt = llvm::find(channels.asArrayRef(), entry.first); + auto sources = send->getAttrOfType("raptor.batch_source_cores"); + auto targets = send->getAttrOfType("raptor.batch_target_cores"); + if (channelIt == channels.asArrayRef().end() || !sources || !targets) + return send.emitOpError("batch send channel metadata is incomplete"); + size_t index = std::distance(channels.asArrayRef().begin(), channelIt); + sendSource = sources.asArrayRef()[index]; + sendTarget = targets.asArrayRef()[index]; + } else { + auto source = getI64Attr(send, "raptor.source_core"); + auto target = getI64Attr(send, "raptor.target_core"); + if (!source || !target) + return send.emitOpError("send core metadata is incomplete"); + sendSource = *source; + sendTarget = *target; + } + int64_t receiveSource = 0; + int64_t receiveTarget = 0; + if (auto channels = receive->getAttrOfType("raptor.batch_channel_ids")) { + auto channelIt = llvm::find(channels.asArrayRef(), entry.first); + auto sources = receive->getAttrOfType("raptor.batch_source_cores"); + auto targets = receive->getAttrOfType("raptor.batch_target_cores"); + if (channelIt == channels.asArrayRef().end() || !sources || !targets) + return receive.emitOpError("batch receive channel metadata is incomplete"); + size_t index = std::distance(channels.asArrayRef().begin(), channelIt); + receiveSource = sources.asArrayRef()[index]; + receiveTarget = targets.asArrayRef()[index]; + } else { + auto source = getI64Attr(receive, "raptor.source_core"); + auto target = getI64Attr(receive, "raptor.target_core"); + if (!source || !target) + return receive.emitOpError("receive core metadata is incomplete"); + receiveSource = *source; + receiveTarget = *target; + } + if (receiveSource != sendSource || receiveTarget != sendTarget) + return receive.emitOpError("receive core metadata does not match its send"); + } + return simulate(funcOp, streams, "realized"); +} + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp new file mode 100644 index 0000000..062af6b --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Support/LLVM.h" + +namespace onnx_mlir::spatial { + +struct PlannedCommunicationTransfer { + uint64_t exchangeId = 0; + uint64_t parentExchangeId = 0; + unsigned sourceStream = 0; + unsigned targetStream = 0; + unsigned producerStep = 0; + unsigned consumerStep = 0; + unsigned sourceInsertionStep = 0; + unsigned targetInsertionStep = 0; +}; + +mlir::LogicalResult verifyPlannedCommunicationDeadlockFree( + mlir::Operation *anchor, + unsigned streamCount, + mlir::ArrayRef stepCounts, + mlir::ArrayRef transfers); + +mlir::LogicalResult verifyRealizedCommunicationDeadlockFree(mlir::func::FuncOp funcOp); + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp index 0261cd2..d3f726c 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp @@ -1,21 +1,16 @@ #include "DeferredCommunicationPlanning.hpp" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/IRMapping.h" -namespace onnx_mlir { -namespace spatial { +#include "llvm/ADT/SmallPtrSet.h" +namespace onnx_mlir::spatial { using namespace mlir; - namespace { -struct LaneProducerInfo { - ProducerValueRef producer; - Value originalSource; - DeferredEndpoint producerEndpoint; - DeferredEndpoint consumerEndpoint; -}; - static Value getBlockOperand(Block &block, ValueRange operands, Value value, unsigned firstArgument = 0) { auto it = llvm::find(operands, value); assert(it != operands.end() && "missing scheduled operand"); @@ -29,398 +24,226 @@ static FailureOr getOriginalProducerValue(const ProducerValueRef &produce return outputs[producer.resultIndex]; } -static DeferredEndpoint getDeferredEndpoint(const ComputeInstance &instance, - const MergeScheduleResult &schedule) { - size_t cpu = getScheduledCpuForComputeInstance(instance, schedule); - return DeferredEndpoint { - instance, - cpu, - getCanonicalPeftClassId(cpu, schedule), - instance.laneStart, - instance.laneCount - }; +static bool isSupportedDeferredShapingOp(Operation *op) { + return isa(op); } -static void setDeferredTransferMetadata(SpatDeferredCommunicationOp transfer, - OpBuilder &builder, - const DeferredTransferMetadata &transferMetadata) { - SmallVector sourceCpus; - SmallVector sourceClasses; - SmallVector sourceLaneRanges; - SmallVector targetCpus; - SmallVector targetClasses; - SmallVector targetLaneRanges; - for (const DeferredEndpoint &endpoint : transferMetadata.producers) { - sourceCpus.push_back(static_cast(endpoint.cpu)); - sourceClasses.push_back(static_cast(endpoint.canonicalClassId)); - sourceLaneRanges.push_back(static_cast(endpoint.laneStart)); - sourceLaneRanges.push_back(static_cast(endpoint.laneCount)); - } - for (const DeferredEndpoint &endpoint : transferMetadata.consumers) { - targetCpus.push_back(static_cast(endpoint.cpu)); - targetClasses.push_back(static_cast(endpoint.canonicalClassId)); - targetLaneRanges.push_back(static_cast(endpoint.laneStart)); - targetLaneRanges.push_back(static_cast(endpoint.laneCount)); - } - transfer->setAttr("source_cpus", builder.getDenseI64ArrayAttr(sourceCpus)); - transfer->setAttr("source_classes", builder.getDenseI64ArrayAttr(sourceClasses)); - transfer->setAttr("source_lane_ranges", builder.getDenseI64ArrayAttr(sourceLaneRanges)); - transfer->setAttr("target_cpus", builder.getDenseI64ArrayAttr(targetCpus)); - transfer->setAttr("target_classes", builder.getDenseI64ArrayAttr(targetClasses)); - transfer->setAttr("target_lane_ranges", builder.getDenseI64ArrayAttr(targetLaneRanges)); - // `batched=true` means this deferred communication belongs to a - // spat.scheduled_compute_batch block and the array attrs are indexed by - // scheduled lane. The legacy scalar attrs on the op keep representative values - // for compatibility; Phase 2 must consume the arrays when batched=true. - transfer->setAttr("batched", builder.getBoolAttr(transferMetadata.isBatched)); - if (transferMetadata.isBatched) { - transfer->setAttr("source_operand_for_scheduled_lane", - builder.getDenseI64ArrayAttr(transferMetadata.sourceOperandForScheduledLane)); - transfer->setAttr("multi_source_payload", builder.getBoolAttr(transferMetadata.multiSourcePayload)); - } -} - -static FailureOr projectDeferredPayload(Value selectedSource, - Value consumerInput, - OpBuilder &builder, - Location loc) { - if (selectedSource.getType() == consumerInput.getType()) - return selectedSource; - - auto slice = consumerInput.getDefiningOp(); - if (!slice || slice.getSource().getType() != selectedSource.getType()) - return failure(); - - IRMapping mapper; - mapper.map(slice.getSource(), selectedSource); - return cast(builder.clone(*slice, mapper)).getResult(); -} - -static FailureOr buildSelectedDeferredSource(OpBuilder &builder, - Location loc, +static FailureOr buildSelectedDeferredSource(OpBuilder &builder, Location loc, SpatDeferredCommunicationOp transfer, Value scheduledLane, ValueRange sourceBlockArgs, - DenseI64ArrayAttr sourceOperandForScheduledLaneAttr) { - SmallVector sourceOperandForScheduledLane(sourceOperandForScheduledLaneAttr.asArrayRef().begin(), - sourceOperandForScheduledLaneAttr.asArrayRef().end()); - Value sourceOperandIndexTable = - createI64LookupTableConstant(builder, transfer.getOperation(), sourceOperandForScheduledLane); - Value sourceIndexI64 = - tensor::ExtractOp::create(builder, loc, sourceOperandIndexTable, ValueRange {scheduledLane}).getResult(); - Value sourceIndex = arith::IndexCastOp::create(builder, loc, builder.getIndexType(), sourceIndexI64).getResult(); - - Type elementType = sourceBlockArgs.front().getType(); - auto sourceTensorType = dyn_cast(elementType); - if (!sourceTensorType || !sourceTensorType.hasStaticShape()) - return transfer.emitOpError( - "scheduled batch deferred communication with multiple original graph SSA sources requires compatible static ranked tensor sources"), - failure(); - - for (Value source : sourceBlockArgs) { - auto rankedSourceType = dyn_cast(source.getType()); - if (!rankedSourceType || !rankedSourceType.hasStaticShape() || rankedSourceType != sourceTensorType) - return transfer.emitOpError( - "scheduled batch deferred communication with multiple original graph SSA sources requires compatible static ranked tensor sources"), - failure(); + ArrayRef sourceOperandForScheduledLane) { + if (sourceBlockArgs.size() == 1) + return sourceBlockArgs.front(); + if (!scheduledLane || sourceOperandForScheduledLane.empty()) + return transfer.emitOpError("multiple deferred sources require the enclosing scheduled lane"), failure(); + Value table = createI64LookupTableConstant(builder, transfer.getOperation(), sourceOperandForScheduledLane); + Value i64 = tensor::ExtractOp::create(builder, loc, table, ValueRange {scheduledLane}).getResult(); + Value index = arith::IndexCastOp::create(builder, loc, builder.getIndexType(), i64).getResult(); + auto type = dyn_cast(sourceBlockArgs.front().getType()); + if (!type || !type.hasStaticShape()) + return transfer.emitOpError("multiple deferred sources require static ranked tensors"), failure(); + for (Value source : sourceBlockArgs) + if (source.getType() != type) + return transfer.emitOpError("multiple deferred sources require identical tensor types"), failure(); + SmallVector shape {static_cast(sourceBlockArgs.size())}; + llvm::append_range(shape, type.getShape()); + auto stacked = createEmptyTensorForType(builder, loc, RankedTensorType::get(shape, type.getElementType())); + if (failed(stacked)) + return failure(); + Value value = *stacked; + SmallVector sizes, strides(shape.size(), builder.getIndexAttr(1)); + sizes.push_back(builder.getIndexAttr(1)); + for (int64_t dim : type.getShape()) sizes.push_back(builder.getIndexAttr(dim)); + for (auto [i, source] : llvm::enumerate(sourceBlockArgs)) { + SmallVector expandedShape {1}; llvm::append_range(expandedShape, type.getShape()); + SmallVector reassociation {{0, 1}}; + for (int64_t dim = 1; dim < type.getRank(); ++dim) reassociation.push_back({dim + 1}); + Value expanded = tensor::ExpandShapeOp::create(builder, loc, + RankedTensorType::get(expandedShape, type.getElementType()), source, reassociation).getResult(); + SmallVector offsets(shape.size(), builder.getIndexAttr(0)); + offsets[0] = builder.getIndexAttr(i); + value = tensor::InsertSliceOp::create(builder, loc, expanded, value, offsets, sizes, strides).getResult(); } - if (sourceTensorType.getRank() == 0) - return transfer.emitOpError( - "scheduled batch deferred communication with multiple original graph SSA sources cannot yet build compact tensor source selection"), - failure(); - - SmallVector stackedShape {static_cast(sourceBlockArgs.size())}; - llvm::append_range(stackedShape, sourceTensorType.getShape()); - RankedTensorType stackedType = RankedTensorType::get(stackedShape, sourceTensorType.getElementType()); - auto stackedEmpty = createEmptyTensorForType(builder, loc, stackedType); - if (failed(stackedEmpty)) - return transfer.emitOpError( - "scheduled batch deferred communication with multiple original graph SSA sources cannot yet build compact tensor source selection"), - failure(); - Value stackedSources = *stackedEmpty; - - SmallVector insertSizes; - SmallVector unitStrides(stackedType.getRank(), builder.getIndexAttr(1)); - insertSizes.push_back(builder.getIndexAttr(1)); - for (int64_t dim : sourceTensorType.getShape()) - insertSizes.push_back(builder.getIndexAttr(dim)); - - for (auto [sourceIndexValue, source] : llvm::enumerate(sourceBlockArgs)) { - SmallVector expandedShape {1}; - llvm::append_range(expandedShape, sourceTensorType.getShape()); - SmallVector reassociation; - reassociation.push_back({0, 1}); - for (int64_t dim = 1; dim < sourceTensorType.getRank(); ++dim) - reassociation.push_back({dim + 1}); - Value expandedSource = tensor::ExpandShapeOp::create(builder, - loc, - RankedTensorType::get(expandedShape, sourceTensorType.getElementType()), - source, - reassociation) - .getResult(); - SmallVector insertOffsets(stackedType.getRank(), builder.getIndexAttr(0)); - insertOffsets[0] = builder.getIndexAttr(sourceIndexValue); - stackedSources = tensor::InsertSliceOp::create( - builder, loc, expandedSource, stackedSources, insertOffsets, insertSizes, unitStrides) - .getResult(); - } - - SmallVector extractOffsets(stackedType.getRank(), builder.getIndexAttr(0)); - SmallVector extractSizes; - SmallVector extractStrides(stackedType.getRank(), builder.getIndexAttr(1)); - extractOffsets[0] = sourceIndex; - extractSizes.push_back(builder.getIndexAttr(1)); - for (int64_t dim : sourceTensorType.getShape()) - extractSizes.push_back(builder.getIndexAttr(dim)); - SmallVector selectedSliceShape {1}; - llvm::append_range(selectedSliceShape, sourceTensorType.getShape()); - RankedTensorType selectedSliceType = RankedTensorType::get(selectedSliceShape, sourceTensorType.getElementType()); - Value selectedSlice = tensor::ExtractSliceOp::create( - builder, - loc, - selectedSliceType, - stackedSources, - extractOffsets, - extractSizes, - extractStrides) - .getResult(); - - SmallVector collapseReassociation; - collapseReassociation.push_back({0, 1}); - for (int64_t dim = 1; dim < sourceTensorType.getRank(); ++dim) - collapseReassociation.push_back({dim + 1}); - return tensor::CollapseShapeOp::create(builder, loc, sourceTensorType, selectedSlice, collapseReassociation) - .getResult(); + SmallVector offsets(shape.size(), builder.getIndexAttr(0)); offsets[0] = index; + SmallVector sliceShape {1}; llvm::append_range(sliceShape, type.getShape()); + auto slice = tensor::ExtractSliceOp::create(builder, loc, + RankedTensorType::get(sliceShape, type.getElementType()), value, offsets, sizes, strides); + // extract has a leading unit dimension; remove it without changing the payload. + SmallVector reassociation {{0, 1}}; + for (int64_t dim = 1; dim < type.getRank(); ++dim) reassociation.push_back({dim + 1}); + return tensor::CollapseShapeOp::create(builder, loc, type, slice.getResult(), reassociation).getResult(); } -static LogicalResult finalizeDeferredCommunication(OpBuilder &builder, - Location loc, - Value consumerInput, - SpatDeferredCommunicationOp transfer, - Value &result) { - Block *block = builder.createBlock(&transfer.getBody(), - transfer.getBody().end(), - TypeRange {transfer.getSources().getTypes()}, - SmallVector(transfer.getSources().size(), loc)); - builder.setInsertionPointToStart(block); +static bool isTopLevelShaping(Operation *op, Block &body) { + return op->getBlock() == &body && isSupportedDeferredShapingOp(op); +} - Value selectedSource = block->getArgument(0); - if (auto parentScheduled = dyn_cast(transfer->getParentOp())) { - auto sourceOperandForScheduledLaneAttr = - transfer->getAttrOfType("source_operand_for_scheduled_lane"); - auto multiSourcePayloadAttr = transfer->getAttrOfType("multi_source_payload"); - if (sourceOperandForScheduledLaneAttr && sourceOperandForScheduledLaneAttr.size() == parentScheduled.getLaneCount()) { - if (multiSourcePayloadAttr && multiSourcePayloadAttr.getValue()) { - FailureOr multiSourceSelection = buildSelectedDeferredSource( - builder, - loc, - transfer, - transfer->getBlock()->getArgument(0), - block->getArguments(), - sourceOperandForScheduledLaneAttr); - if (failed(multiSourceSelection)) - return failure(); - selectedSource = *multiSourceSelection; - } else { - selectedSource = block->getArgument(sourceOperandForScheduledLaneAttr.asArrayRef().front()); - } +static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan, + llvm::SmallPtrSetImpl &seen) { + if (value == plan.graphInput || value == plan.graphLane || value == plan.scheduledLane) + return true; + auto arg = dyn_cast(value); + if (arg) + return false; + Operation *op = value.getDefiningOp(); + if (op && op->hasTrait()) + return true; + if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second) + return op && seen.contains(op); + return llvm::all_of(op->getOperands(), [&](Value operand) { return isEligible(operand, body, plan, seen); }); +} + +static FailureOr clonePayloadRoot(Value root, Block &body, const DeferredInputPlan &plan, + OpBuilder &builder, SpatDeferredCommunicationOp transfer, + Value selectedSource) { + IRMapping mapping; + mapping.map(plan.graphInput, selectedSource); + std::function(Value)> cloneScheduledLane = [&](Value value) -> FailureOr { + if (mapping.contains(value)) return mapping.lookup(value); + if (value == plan.scheduledLane) return value; + if (isa(value)) + return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure(); + Operation *op = value.getDefiningOp(); + if (!op || !isSupportedDeferredShapingOp(op)) + return transfer.emitOpError("phase 1 cannot clone the scheduled graph-lane expression"), failure(); + for (Value operand : op->getOperands()) if (failed(cloneScheduledLane(operand))) return failure(); + Operation *copy = builder.clone(*op, mapping); + for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair)); + return mapping.lookup(value); + }; + std::function(Value)> clone = [&](Value value) -> FailureOr { + if (mapping.contains(value)) return mapping.lookup(value); + if (value == plan.graphLane) { + auto mappedLane = cloneScheduledLane(plan.scheduledGraphLane); + if (failed(mappedLane)) return failure(); + mapping.map(value, *mappedLane); + return *mappedLane; } - } - - FailureOr yielded = projectDeferredPayload(selectedSource, consumerInput, builder, loc); - if (failed(yielded) || (*yielded).getType() != consumerInput.getType()) - return transfer.emitOpError("cannot derive deferred communication payload from original graph producer value"); - - SpatYieldOp::create(builder, loc, *yielded); - result = transfer.getOutput(); - builder.setInsertionPointAfter(transfer); - return success(); + if (isa(value)) + return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure(); + Operation *op = value.getDefiningOp(); + if (!op || (!isTopLevelShaping(op, body) && !op->hasTrait())) + return transfer.emitOpError("phase 1 payload shaping contains an unsupported operation"), failure(); + for (Value operand : op->getOperands()) if (failed(clone(operand))) return failure(); + Operation *copy = builder.clone(*op, mapping); + for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair)); + return mapping.lookup(value); + }; + return clone(root); } -static LogicalResult createSingleLaneDeferredCommunication(OpBuilder &builder, - Location loc, - Value consumerInput, - Value originalSource, - const ProducerValueRef &producer, - const ComputeInstance &consumer, - const MergeScheduleResult &schedule, - uint64_t exchangeId, - Value &result) { - DeferredTransferMetadata transferMetadata; - transferMetadata.producers.push_back(getDeferredEndpoint(producer.instance, schedule)); - transferMetadata.consumers.push_back(getDeferredEndpoint(consumer, schedule)); - auto transfer = SpatDeferredCommunicationOp::create( - builder, - loc, - consumerInput.getType(), - ValueRange {originalSource}, - builder.getStringAttr(("x" + std::to_string(exchangeId)).c_str()), - builder.getStringAttr(getInstanceName(producer.instance)), - builder.getStringAttr(getInstanceName(consumer)), - builder.getI64IntegerAttr(static_cast(transferMetadata.producers.front().canonicalClassId)), - builder.getI64IntegerAttr(static_cast(transferMetadata.consumers.front().canonicalClassId)), - builder.getI64IntegerAttr(static_cast(transferMetadata.producers.front().cpu)), - builder.getI64IntegerAttr(static_cast(transferMetadata.consumers.front().cpu)), - builder.getI64IntegerAttr(producer.instance.laneStart), - builder.getI64IntegerAttr(consumer.laneStart), - builder.getStringAttr(consumerInput.getDefiningOp() ? "projected" : "ordinary"), - builder.getI64IntegerAttr(static_cast(producer.resultIndex)), - builder.getDenseI64ArrayAttr({}), - builder.getI64IntegerAttr(-1)); - setDeferredTransferMetadata(transfer, builder, transferMetadata); - return finalizeDeferredCommunication(builder, loc, consumerInput, transfer, result); -} - -static LogicalResult createScheduledLaneDeferredCommunication(OpBuilder &builder, - Location loc, - Value consumerInput, - ValueRange originalSources, - const ProducerValueRef &producer, - const DeferredTransferMetadata &transferMetadata, - uint64_t exchangeId, - Value &result) { - assert(transferMetadata.isBatched && "expected scheduled-lane deferred communication metadata"); - assert(!transferMetadata.producers.empty() && !transferMetadata.consumers.empty() && "expected non-empty endpoints"); - auto transfer = SpatDeferredCommunicationOp::create( - builder, - loc, - consumerInput.getType(), - originalSources, - builder.getStringAttr(("x" + std::to_string(exchangeId)).c_str()), - builder.getStringAttr(getInstanceName(producer.instance)), - builder.getStringAttr(getInstanceName(transferMetadata.consumers.front().instance)), - builder.getI64IntegerAttr(static_cast(transferMetadata.producers.front().canonicalClassId)), - builder.getI64IntegerAttr(static_cast(transferMetadata.consumers.front().canonicalClassId)), - builder.getI64IntegerAttr(static_cast(transferMetadata.producers.front().cpu)), - builder.getI64IntegerAttr(static_cast(transferMetadata.consumers.front().cpu)), - builder.getI64IntegerAttr(producer.instance.laneStart), - builder.getI64IntegerAttr(transferMetadata.consumers.front().laneStart), - builder.getStringAttr(consumerInput.getDefiningOp() ? "projected" : "ordinary"), - builder.getI64IntegerAttr(static_cast(producer.resultIndex)), - builder.getDenseI64ArrayAttr({}), - builder.getI64IntegerAttr(-1)); - setDeferredTransferMetadata(transfer, builder, transferMetadata); - return finalizeDeferredCommunication(builder, loc, consumerInput, transfer, result); +static void collectClosure(Value value, Block &body, const DeferredInputPlan &plan, + llvm::SmallPtrSetImpl &ops) { + Operation *op = value.getDefiningOp(); + if (!op || !isTopLevelShaping(op, body) || !ops.insert(op).second) return; + for (Value operand : op->getOperands()) + if (operand != plan.graphInput && operand != plan.graphLane) collectClosure(operand, body, plan, ops); } } // namespace -LogicalResult mapSingleCpuInput(OpBuilder &builder, - Location loc, - Value input, - const ComputeInstance &consumerInstance, - const MergeScheduleResult &schedule, - ValueRange scheduledInputs, - Block &block, - unsigned firstInputArgument, - ArrayRef carriedKeys, - uint64_t &exchangeId, - Value &mapped) { - std::optional producer = getProducerValueRef(input, &consumerInstance); - if (!producer) { - mapped = getBlockOperand(block, scheduledInputs, input, firstInputArgument); +LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput, + const ComputeInstance &consumerInstance, const MergeScheduleResult &, + ValueRange scheduledInputs, Block &block, unsigned firstInputArgument, + ArrayRef carriedKeys, Value graphLane, Value scheduledGraphLane, + DeferredInputPlan &plan) { + plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}}; + auto producer = getProducerValueRef(input, &consumerInstance); + if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); } + ProducerValueKey key {producer->instance, producer->resultIndex}; + auto carried = llvm::find(carriedKeys, key); + if (carried != carriedKeys.end()) { + plan.availableValue = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carried)); return success(); } - - ProducerValueKey producerKey {producer->instance, producer->resultIndex}; - auto carriedIt = llvm::find(carriedKeys, producerKey); - if (carriedIt != carriedKeys.end()) { - mapped = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carriedIt)); - return success(); - } - - FailureOr originalSource = getOriginalProducerValue(*producer); - if (failed(originalSource)) - return emitError(loc) << "cannot resolve original graph producer value for " << getInstanceName(producer->instance); - return createSingleLaneDeferredCommunication( - builder, loc, input, *originalSource, *producer, consumerInstance, schedule, exchangeId++, mapped); + auto source = getOriginalProducerValue(*producer); + if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value"; + plan.originalSources.push_back(*source); + return success(); } -static FailureOr getComputeInstanceInputIndex(const ComputeInstance &instance, Value input) { - auto inputs = getComputeInstanceInputs(instance); +LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput, + const ComputeStepTuple &tuple, const PeftClassPlan &, + const MergeScheduleResult &, ValueRange scheduledInputs, Block &block, + unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane, Value scheduledLane, + DeferredInputPlan &plan) { + plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane}; + const ComputeInstance &representative = tuple.instances.front(); + auto producer = getProducerValueRef(input, &representative); + if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); } + auto inputs = getComputeInstanceInputs(representative); auto it = llvm::find(inputs, input); - if (it == inputs.end()) - return failure(); - return static_cast(std::distance(inputs.begin(), it)); + if (it == inputs.end()) return emitError(loc) << "cannot resolve scheduled batch step input"; + unsigned inputIndex = std::distance(inputs.begin(), it); + for (const ComputeInstance &instance : tuple.instances) { + auto laneInputs = getComputeInstanceInputs(instance); + if (inputIndex >= laneInputs.size()) return emitError(loc) << "scheduled batch step input out of range"; + auto laneProducer = getProducerValueRef(laneInputs[inputIndex], &instance); + if (!laneProducer) return emitError(loc) << "scheduled batch step mixes host and producer inputs"; + auto source = getOriginalProducerValue(*laneProducer); + if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value"; + auto sourceIt = llvm::find(plan.originalSources, *source); + if (sourceIt == plan.originalSources.end()) { plan.sourceOperandForScheduledLane.push_back(plan.originalSources.size()); plan.originalSources.push_back(*source); } + else plan.sourceOperandForScheduledLane.push_back(std::distance(plan.originalSources.begin(), sourceIt)); + } + return success(); } -LogicalResult mapMultiCpuTupleInput(OpBuilder &builder, - Location loc, - Value input, - const ComputeStepTuple &stepTuple, - const PeftClassPlan &peftClassPlan, - const MergeScheduleResult &schedule, - ValueRange scheduledInputs, - Block &block, - unsigned firstInputArgument, - uint64_t &exchangeId, - Value &mapped) { - assert(!stepTuple.instances.empty() && "expected non-empty step tuple"); - const ComputeInstance &representative = stepTuple.instances.front(); - std::optional representativeProducer = getProducerValueRef(input, &representative); - if (!representativeProducer) { - FailureOr inputIndex = getComputeInstanceInputIndex(representative, input); - if (failed(inputIndex)) - return emitError(loc) << "cannot resolve scheduled batch step host input index"; - for (const ComputeInstance &instance : stepTuple.instances) { - auto instanceInputs = getComputeInstanceInputs(instance); - if (*inputIndex >= instanceInputs.size()) - return emitError(loc) << "scheduled batch step host input index out of range"; - if (instanceInputs[*inputIndex] != input || getProducerValueRef(instanceInputs[*inputIndex], &instance)) - return emitError(loc) << "scheduled batch step requires identical host inputs across lanes"; +LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc, Block &body, + ArrayRef plans, IRMapping &mapper, + llvm::SmallPtrSetImpl &absorbed) { + for (const DeferredInputPlan &plan : plans) { + if (plan.availableValue) { mapper.map(plan.graphInput, plan.availableValue); continue; } + SmallVector roots; + bool needsIdentity = false; + SmallVector worklist {plan.graphInput}; + llvm::SmallDenseSet seen; + while (!worklist.empty()) { + Value value = worklist.pop_back_val(); + if (!seen.insert(value).second) continue; + for (OpOperand &use : value.getUses()) { + Operation *user = use.getOwner(); + if (!isTopLevelShaping(user, body)) { needsIdentity = true; continue; } + llvm::SmallPtrSet eligibility; + if (!isEligible(user->getResult(0), body, plan, eligibility)) { needsIdentity = true; continue; } + for (Value result : user->getResults()) { + bool hasShapingUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return isTopLevelShaping(next.getOwner(), body); }); + bool hasOtherUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return !isTopLevelShaping(next.getOwner(), body); }); + if (hasOtherUse) roots.push_back(result); + if (hasShapingUse) worklist.push_back(result); + } + } + } + if (needsIdentity) roots.push_back(plan.graphInput); + llvm::sort(roots, [](Value a, Value b) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); }); + roots.erase(std::unique(roots.begin(), roots.end()), roots.end()); + for (Value root : roots) { + auto transfer = SpatDeferredCommunicationOp::create(builder, loc, root.getType(), plan.originalSources); + Block *deferred = builder.createBlock(&transfer.getBody(), transfer.getBody().end(), + TypeRange {transfer.getSources().getTypes()}, SmallVector(transfer.getSources().size(), loc)); + builder.setInsertionPointToStart(deferred); + auto selected = buildSelectedDeferredSource(builder, loc, transfer, plan.scheduledLane, + deferred->getArguments(), plan.sourceOperandForScheduledLane); + if (failed(selected)) return failure(); + auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected); + if (failed(payload)) return failure(); + SpatYieldOp::create(builder, loc, *payload); + mapper.map(root, transfer.getOutput()); + collectClosure(root, body, plan, absorbed); + builder.setInsertionPointAfter(transfer); } - mapped = getBlockOperand(block, scheduledInputs, input, firstInputArgument); - return success(); } - - FailureOr inputIndex = getComputeInstanceInputIndex(representative, input); - if (failed(inputIndex)) - return emitError(loc) << "cannot resolve scheduled batch step input index"; - - (void)peftClassPlan; - DeferredTransferMetadata transferMetadata; - transferMetadata.isBatched = true; - SmallVector laneProducerInfos; - SmallVector uniqueOriginalSources; - SmallVector sourceOperandForScheduledLane; - for (const ComputeInstance &instance : stepTuple.instances) { - auto instanceInputs = getComputeInstanceInputs(instance); - if (*inputIndex >= instanceInputs.size()) - return emitError(loc) << "scheduled batch step input index out of range"; - Value laneInput = instanceInputs[*inputIndex]; - std::optional laneProducer = getProducerValueRef(laneInput, &instance); - if (!laneProducer) - return emitError(loc) << "scheduled batch step mixes host and producer-derived inputs"; - FailureOr laneSource = getOriginalProducerValue(*laneProducer); - if (failed(laneSource)) - return emitError(loc) << "cannot resolve original graph producer value for " << getInstanceName(laneProducer->instance); - - auto sourceIt = llvm::find(uniqueOriginalSources, *laneSource); - if (sourceIt == uniqueOriginalSources.end()) { - sourceOperandForScheduledLane.push_back(static_cast(uniqueOriginalSources.size())); - uniqueOriginalSources.push_back(*laneSource); - } else { - sourceOperandForScheduledLane.push_back(std::distance(uniqueOriginalSources.begin(), sourceIt)); - } - - laneProducerInfos.push_back(LaneProducerInfo { - *laneProducer, - *laneSource, - getDeferredEndpoint(laneProducer->instance, schedule), - getDeferredEndpoint(instance, schedule), + for (Operation *op : absorbed) { + bool allResultsMapped = llvm::all_of(op->getResults(), [&](Value result) { + return mapper.contains(result) || llvm::all_of(result.getUses(), [&](OpOperand &use) { return absorbed.contains(use.getOwner()); }); }); + if (!allResultsMapped) absorbed.erase(op); } - - for (const LaneProducerInfo &laneProducerInfo : laneProducerInfos) { - transferMetadata.producers.push_back(laneProducerInfo.producerEndpoint); - transferMetadata.consumers.push_back(laneProducerInfo.consumerEndpoint); - } - transferMetadata.sourceOperandForScheduledLane = sourceOperandForScheduledLane; - transferMetadata.multiSourcePayload = uniqueOriginalSources.size() > 1; - - return createScheduledLaneDeferredCommunication( - builder, loc, input, uniqueOriginalSources, *representativeProducer, transferMetadata, exchangeId++, mapped); + return success(); } -} // namespace spatial -} // namespace onnx_mlir +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp index e4598cf..4165b9d 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp @@ -5,29 +5,42 @@ namespace onnx_mlir { namespace spatial { -LogicalResult mapSingleCpuInput(OpBuilder &builder, - Location loc, - Value input, - const ComputeInstance &consumerInstance, - const MergeScheduleResult &schedule, - ValueRange scheduledInputs, - Block &block, - unsigned firstInputArgument, - ArrayRef carriedKeys, - uint64_t &exchangeId, - Value &mapped); +// A graph input is either already available in the scheduled block, or is a +// graph result whose individual deterministic payloads are materialized later. +struct DeferredInputPlan { + BlockArgument graphInput; + Value availableValue; + SmallVector originalSources; + SmallVector sourceOperandForScheduledLane; + Value graphLane; + Value scheduledGraphLane; + Value scheduledLane; +}; -LogicalResult mapMultiCpuTupleInput(OpBuilder &builder, - Location loc, - Value input, - const ComputeStepTuple &stepTuple, - const PeftClassPlan &peftClassPlan, +LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input, + BlockArgument graphInput, + const ComputeInstance &consumerInstance, const MergeScheduleResult &schedule, - ValueRange scheduledInputs, - Block &block, + ValueRange scheduledInputs, Block &block, unsigned firstInputArgument, - uint64_t &exchangeId, - Value &mapped); + ArrayRef carriedKeys, + Value graphLane, Value scheduledGraphLane, + DeferredInputPlan &plan); + +LogicalResult prepareMultiCpuTupleInput(OpBuilder &builder, Location loc, Value input, + BlockArgument graphInput, + const ComputeStepTuple &stepTuple, + const PeftClassPlan &peftClassPlan, + const MergeScheduleResult &schedule, + ValueRange scheduledInputs, Block &block, + unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane, + Value scheduledLane, DeferredInputPlan &plan); + +LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc, + Block &graphBody, + ArrayRef plans, + IRMapping &mapper, + llvm::SmallPtrSetImpl &absorbed); } // namespace spatial } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp new file mode 100644 index 0000000..6db9d78 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp @@ -0,0 +1,1308 @@ +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/PatternMatch.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/SmallPtrSet.h" + +#include + +#include "DeferredCommunicationDeadlock.hpp" +#include "DeferredCommunicationRealization.hpp" +#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +namespace onnx_mlir::spatial { + +using namespace mlir; + +namespace { + +struct ScheduledInfo; + +struct DeferredBodyAnalysis { + struct SourceUse { + unsigned sourceIndex = 0; + tensor::ExtractSliceOp leadingSlice; + std::optional staticSourceLane; + }; + + Value yieldedValue; + SmallVector sourceSideOps; + SmallVector sourceUses; +}; + +struct ProducedValue { + ScheduledInfo* scheduled = nullptr; + unsigned step = 0; + unsigned resultIndex = 0; + int64_t graphId = -1; + int64_t core = -1; + int64_t laneStart = 0; + int64_t laneCount = 1; + Value payload; + Value published; +}; + +struct ScheduledInfo { + Operation* op = nullptr; + SmallVector blocks; + SmallVector stepAnchors; + SmallVector cores; + SmallVector stepSourceIds; + SmallVector resultOffsets; + SmallVector resultCounts; + SmallVector produced; + SmallVector streamIds; + + bool isBatch() const { return isa(op); } +}; + +struct LaneTransfer { + ProducedValue* source = nullptr; + unsigned sourceStream = 0; + unsigned targetStream = 0; + int64_t sourceCore = -1; + int64_t targetCore = -1; + unsigned targetLane = 0; + unsigned sourceIndex = 0; + uint64_t channelId = 0; + unsigned sourceInsertionStep = 0; + unsigned targetInsertionStep = 0; + Type transportType; + Operation* bodyRoot = nullptr; + bool local = false; +}; + +struct NormalizedScalarExchange { + uint64_t id = 0; + SpatDeferredCommunicationOp deferred; + ScheduledInfo* target = nullptr; + unsigned consumerStep = 0; + DeferredBodyAnalysis body; + LaneTransfer transfer; +}; + +struct NormalizedBatchedExchange { + uint64_t id = 0; + SpatDeferredCommunicationOp deferred; + ScheduledInfo* target = nullptr; + unsigned consumerStep = 0; + DeferredBodyAnalysis body; + SmallVector lanes; +}; + +struct RealizationPlan { + SmallVector scheduled; + SmallVector> producedStorage; + DenseMap> producedByGraph; + SmallVector scalar; + SmallVector batched; + SmallVector external; + SmallVector stepCounts; +}; + +static FailureOr> getI64Array(Operation* op, StringRef name) { + auto attr = op->getAttrOfType(name); + if (!attr) + return op->emitOpError() << "phase 2 requires '" << name << "' metadata"; + return SmallVector(attr.asArrayRef().begin(), attr.asArrayRef().end()); +} + +static FailureOr> getLaneTable(Operation* op, StringRef name, size_t expected) { + if (auto array = op->getAttrOfType(name)) { + if (array.size() != static_cast(expected)) + return op->emitOpError() << "phase 2 metadata '" << name << "' has the wrong size"; + return SmallVector(array.asArrayRef().begin(), array.asArrayRef().end()); + } + auto elements = op->getAttrOfType(name); + if (!elements || elements.getNumElements() != static_cast(expected)) + return op->emitOpError() << "phase 2 requires a correctly-sized '" << name << "' lane table"; + SmallVector values; + for (APInt value : elements.getValues()) + values.push_back(value.getSExtValue()); + return values; +} + +static Block* getScheduledBlock(SpatDeferredCommunicationOp deferred, Operation* scheduled) { + Block* block = deferred->getBlock(); + while (block && block->getParentOp() != scheduled) { + Operation* parent = block->getParentOp(); + block = parent ? parent->getBlock() : nullptr; + } + return block; +} + +static FailureOr getStepIndex(ScheduledInfo& info, Block* block) { + auto it = llvm::find(info.blocks, block); + if (it == info.blocks.end()) + return failure(); + return static_cast(std::distance(info.blocks.begin(), it)); +} + +static FailureOr +getScalarStepResult(SpatScheduledCompute scheduled, Block& block, unsigned resultIndex, unsigned resultCount) { + auto yield = dyn_cast(block.getTerminator()); + if (!yield || resultIndex >= resultCount || yield.getOutputs().size() < resultCount) + return scheduled.emitOpError("phase 2 cannot recover a scalar scheduled step result"); + return yield.getOutputs()[yield.getOutputs().size() - resultCount + resultIndex]; +} + +static FailureOr +getBatchStepResult(SpatScheduledComputeBatch scheduled, Block& block, unsigned globalResultIndex) { + auto inParallel = dyn_cast(block.getTerminator()); + if (!inParallel) + return scheduled.emitOpError("phase 2 cannot recover a batched scheduled step result"); + unsigned resultBase = 1 + scheduled.getWeights().size() + scheduled.getInputs().size(); + for (Operation& op : inParallel.getRegion().front()) { + auto insert = dyn_cast(op); + if (!insert) + continue; + auto destination = dyn_cast(insert.getDest()); + if (destination && destination.getOwner() == &block && destination.getArgNumber() == resultBase + globalResultIndex) + return insert.getSource(); + } + return scheduled.emitOpError("phase 2 cannot find the batched result insertion for a scheduled step"); +} + +static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan) { + unsigned nextStream = 0; + for (Operation& op : funcOp.getOps()) { + if (!isa(op)) + continue; + ScheduledInfo info; + info.op = &op; + Region& body = isa(op) ? cast(op).getBody() + : cast(op).getBody(); + for (Block& block : body) { + info.blocks.push_back(&block); + info.stepAnchors.push_back(&block.front()); + } + auto sourceIds = getI64Array(&op, "scheduled.step_source_ids"); + auto offsets = getI64Array(&op, "scheduled.step_result_offsets"); + auto counts = getI64Array(&op, "scheduled.step_result_counts"); + if (failed(sourceIds) || failed(offsets) || failed(counts)) + return failure(); + info.stepSourceIds = std::move(*sourceIds); + info.resultOffsets = std::move(*offsets); + info.resultCounts = std::move(*counts); + if (info.blocks.size() != info.stepSourceIds.size() || info.blocks.size() != info.resultOffsets.size() + || info.blocks.size() != info.resultCounts.size()) + return op.emitOpError("phase 2 scheduled step metadata does not match the region block count"); + + if (auto scalar = dyn_cast(op)) { + auto core = scalar->getAttrOfType(kCoreIdAttrName); + if (!core) + return scalar.emitOpError("phase 2 requires scalar coreId metadata"); + info.cores.push_back(core.getInt()); + } + else { + auto cores = op.getAttrOfType(kCoreIdsAttrName); + if (!cores) + return op.emitOpError("phase 2 requires batch coreIds metadata"); + for (int32_t core : cores.asArrayRef()) + info.cores.push_back(core); + } + for (size_t lane = 0; lane < info.cores.size(); ++lane) + info.streamIds.push_back(nextStream++); + plan.scheduled.push_back(std::move(info)); + } + + plan.stepCounts.resize(nextStream); + for (ScheduledInfo& info : plan.scheduled) + for (unsigned stream : info.streamIds) + plan.stepCounts[stream] = info.blocks.size(); + + for (ScheduledInfo& info : plan.scheduled) { + SmallVector laneStarts; + SmallVector laneCounts; + if (info.isBatch()) { + auto starts = getLaneTable(info.op, "scheduled.source_lane_starts", info.blocks.size() * info.cores.size()); + auto counts = getLaneTable(info.op, "scheduled.source_lane_counts", info.blocks.size() * info.cores.size()); + if (failed(starts) || failed(counts)) + return failure(); + laneStarts = std::move(*starts); + laneCounts = std::move(*counts); + } + else { + auto starts = getI64Array(info.op, "scheduled.source_lane_starts"); + auto counts = getI64Array(info.op, "scheduled.source_lane_counts"); + if (failed(starts) || failed(counts)) + return failure(); + laneStarts = std::move(*starts); + laneCounts = std::move(*counts); + } + + for (unsigned step = 0; step < info.blocks.size(); ++step) { + if (info.resultOffsets[step] < 0 || info.resultCounts[step] < 0) + return info.op->emitOpError("phase 2 scheduled result metadata must be non-negative"); + for (unsigned result = 0; result < static_cast(info.resultCounts[step]); ++result) { + unsigned globalResult = info.resultOffsets[step] + result; + if (globalResult >= info.op->getNumResults()) + return info.op->emitOpError("phase 2 scheduled result metadata is out of range"); + if (!info.isBatch()) { + auto payload = getScalarStepResult( + cast(info.op), *info.blocks[step], result, info.resultCounts[step]); + if (failed(payload)) + return failure(); + auto produced = std::make_unique(ProducedValue { + &info, + step, + result, + info.stepSourceIds[step], + info.cores.front(), + laneStarts[step], + std::max(laneCounts[step], 1), + *payload, + info.op->getResult(globalResult), + }); + info.produced.push_back(produced.get()); + plan.producedByGraph[produced->graphId].push_back(produced.get()); + plan.producedStorage.push_back(std::move(produced)); + continue; + } + auto payload = getBatchStepResult(cast(info.op), *info.blocks[step], globalResult); + if (failed(payload)) + return failure(); + for (unsigned lane = 0; lane < info.cores.size(); ++lane) { + size_t laneIndex = step * info.cores.size() + lane; + auto produced = std::make_unique(ProducedValue { + &info, + step, + result, + info.stepSourceIds[step], + info.cores[lane], + laneStarts[laneIndex], + std::max(laneCounts[laneIndex], 1), + *payload, + info.op->getResult(globalResult), + }); + info.produced.push_back(produced.get()); + plan.producedByGraph[produced->graphId].push_back(produced.get()); + plan.producedStorage.push_back(std::move(produced)); + } + } + } + } + return success(); +} + +static bool isSupportedDeferredBodyOp(Operation* op) { + return isa(op); +} + +static LogicalResult collectDeferredDependencies(SpatDeferredCommunicationOp deferred, + Value value, + llvm::SmallDenseSet& sourceIndices, + llvm::SmallPtrSetImpl& requiredOps, + Value allowedExternalLane = {}) { + Block& body = deferred.getBody().front(); + if (auto argument = dyn_cast(value)) { + if (argument.getOwner() != &body && value == allowedExternalLane) + return success(); + if (argument.getArgNumber() >= deferred.getSources().size()) + return deferred.emitOpError() << "deferred body yielded value depends on unknown block argument " + << argument.getArgNumber(); + sourceIndices.insert(argument.getArgNumber()); + return success(); + } + Operation* definingOp = value.getDefiningOp(); + if (!definingOp || definingOp->getBlock() != &body) { + if (value == allowedExternalLane) + return success(); + if (definingOp && definingOp->hasTrait()) + return success(); + return deferred.emitOpError("deferred body payload derivation captures a value unavailable in the source stream"); + } + if (!isSupportedDeferredBodyOp(definingOp)) + return deferred.emitOpError() << "deferred body payload derivation contains unsupported op '" + << definingOp->getName() << "'"; + if (!requiredOps.insert(definingOp).second) + return success(); + for (Value operand : definingOp->getOperands()) + if (failed(collectDeferredDependencies(deferred, operand, sourceIndices, requiredOps, allowedExternalLane))) + return failure(); + return success(); +} + +static bool dependsOn(Value value, Value needle, llvm::SmallPtrSetImpl& visited) { + if (value == needle) + return true; + Operation* op = value.getDefiningOp(); + if (!op || !visited.insert(op).second) + return false; + return llvm::any_of(op->getOperands(), [&](Value operand) { return dependsOn(operand, needle, visited); }); +} + +static std::optional evaluateConstantInteger(Value value, Value lane = {}, std::optional laneValue = std::nullopt) { + if (!value) + return std::nullopt; + if (value == lane && laneValue) + return laneValue; + if (auto constant = value.getDefiningOp()) + if (auto integer = dyn_cast(constant.getValue())) + return integer.getInt(); + if (auto cast = value.getDefiningOp()) + return evaluateConstantInteger(cast.getIn(), lane, laneValue); + if (auto add = value.getDefiningOp()) { + auto lhs = evaluateConstantInteger(add.getLhs(), lane, laneValue); + auto rhs = evaluateConstantInteger(add.getRhs(), lane, laneValue); + if (lhs && rhs) + return *lhs + *rhs; + } + if (auto sub = value.getDefiningOp()) { + auto lhs = evaluateConstantInteger(sub.getLhs(), lane, laneValue); + auto rhs = evaluateConstantInteger(sub.getRhs(), lane, laneValue); + if (lhs && rhs) + return *lhs - *rhs; + } + if (auto mul = value.getDefiningOp()) { + auto lhs = evaluateConstantInteger(mul.getLhs(), lane, laneValue); + auto rhs = evaluateConstantInteger(mul.getRhs(), lane, laneValue); + if (lhs && rhs) + return *lhs * *rhs; + } + if (auto apply = value.getDefiningOp()) { + auto evaluated = evaluateAffineApply(apply, [lane, laneValue](Value operand) -> FailureOr { + auto constant = evaluateConstantInteger(operand, lane, laneValue); + return constant ? FailureOr(*constant) : FailureOr(failure()); + }); + if (succeeded(evaluated)) + return *evaluated; + } + if (auto extract = value.getDefiningOp()) { + auto constant = extract.getTensor().getDefiningOp(); + auto elements = constant ? dyn_cast(constant.getValue()) : DenseIntElementsAttr(); + auto type = elements ? dyn_cast(elements.getType()) : RankedTensorType(); + if (!elements || !type || extract.getIndices().size() != static_cast(type.getRank())) + return std::nullopt; + int64_t linearIndex = 0; + for (auto [index, dimension] : llvm::zip(extract.getIndices(), type.getShape())) { + auto constantIndex = evaluateConstantInteger(index, lane, laneValue); + if (!constantIndex || *constantIndex < 0 || *constantIndex >= dimension) + return std::nullopt; + linearIndex = linearIndex * dimension + *constantIndex; + } + auto values = elements.getValues(); + return values[linearIndex].getSExtValue(); + } + return std::nullopt; +} + +static FailureOr analyzeDeferredBody(SpatDeferredCommunicationOp deferred, + std::optional targetLane = std::nullopt) { + Block& body = deferred.getBody().front(); + auto yield = dyn_cast(body.getTerminator()); + if (!yield || yield.getOutputs().size() != 1) + return deferred.emitOpError("phase 2 requires a single deferred yielded value"); + + DeferredBodyAnalysis analysis; + analysis.yieldedValue = yield.getOutputs().front(); + llvm::SmallDenseSet sourceIndices; + llvm::SmallPtrSet requiredOps; + Value scheduledLane; + if (auto scheduled = deferred->getParentOfType()) + scheduledLane = scheduled.getLaneArgument().value(); + if (failed(collectDeferredDependencies(deferred, analysis.yieldedValue, sourceIndices, requiredOps, scheduledLane))) + return failure(); + for (Operation& op : body.without_terminator()) + if (requiredOps.contains(&op)) + analysis.sourceSideOps.push_back(&op); + + for (unsigned sourceIndex : sourceIndices) { + DeferredBodyAnalysis::SourceUse sourceUse; + sourceUse.sourceIndex = sourceIndex; + Value sourceArgument = body.getArgument(sourceIndex); + for (Operation* op : analysis.sourceSideOps) { + auto slice = dyn_cast(op); + if (!slice) + continue; + llvm::SmallPtrSet visited; + if (!dependsOn(slice.getSource(), sourceArgument, visited)) + continue; + auto sourceType = dyn_cast(slice.getSourceType()); + if (!sourceType || sourceType.getRank() == 0 || slice.getMixedSizes().empty()) + continue; + auto leadingSize = getConstantIntValue(slice.getMixedSizes().front()); + if (!leadingSize || *leadingSize != 1) + continue; + sourceUse.leadingSlice = slice; + if (auto offset = dyn_cast(slice.getMixedOffsets().front())) + sourceUse.staticSourceLane = cast(offset).getInt(); + else { + Value dynamicOffset = cast(slice.getMixedOffsets().front()); + sourceUse.staticSourceLane = evaluateConstantInteger(dynamicOffset, scheduledLane, + targetLane ? std::optional(*targetLane) : std::nullopt); + } + break; + } + analysis.sourceUses.push_back(sourceUse); + } + return analysis; +} + +static DeferredBodyAnalysis::SourceUse* findSourceUse(DeferredBodyAnalysis& analysis, unsigned sourceIndex) { + auto it = llvm::find_if(analysis.sourceUses, [&](const auto& use) { return use.sourceIndex == sourceIndex; }); + return it == analysis.sourceUses.end() ? nullptr : &*it; +} + +static bool isIdentityDeferredBody(SpatDeferredCommunicationOp deferred, + const DeferredBodyAnalysis& analysis, + unsigned sourceIndex) { + Block& block = deferred.getBody().front(); + return analysis.sourceSideOps.empty() && analysis.yieldedValue == block.getArgument(sourceIndex); +} + +static LogicalResult configurePayloadAndReplacement(SpatDeferredCommunicationOp deferred, + DeferredBodyAnalysis& analysis, + ProducedValue& source, + unsigned sourceIndex, + int64_t sourceLane, + LaneTransfer& lane) { + lane.sourceIndex = sourceIndex; + lane.transportType = analysis.yieldedValue.getType(); + if (isIdentityDeferredBody(deferred, analysis, sourceIndex) && source.payload.getType() == analysis.yieldedValue.getType()) + return success(); + + // A matching source argument lets Phase 2 clone only the body-owned shaping. + if (deferred.getBody().front().getArgument(sourceIndex).getType() == source.payload.getType()) + return success(); + + if (auto* sourceUse = findSourceUse(analysis, sourceIndex); + sourceUse && sourceUse->leadingSlice + && sourceUse->leadingSlice.getResult().getType() == source.payload.getType()) { + if (sourceUse->staticSourceLane && *sourceUse->staticSourceLane != sourceLane) + return deferred.emitOpError("deferred-body source lane disagrees with scheduled source metadata"); + lane.bodyRoot = sourceUse->leadingSlice; + return success(); + } + return deferred.emitOpError( + "phase 1 must place all graph-result shaping inside the deferred_communication body"); +} + +static FailureOr +findProducer(RealizationPlan& plan, SpatDeferredCommunicationOp deferred, int64_t graphId, unsigned resultIndex, + std::optional graphLane) { + ProducedValue* match = nullptr; + for (ProducedValue* produced : plan.producedByGraph.lookup(graphId)) { + if (produced->resultIndex != resultIndex) + continue; + if (graphLane && (*graphLane < produced->laneStart || *graphLane >= produced->laneStart + produced->laneCount)) + continue; + if (match) + return deferred.emitOpError("phase 2 cannot uniquely resolve graph publication ownership"), failure(); + match = produced; + } + if (!match) + return deferred.emitOpError("phase 2 cannot map a deferred source to a scheduled producer"), failure(); + return match; +} + +static LogicalResult buildLaneTransfer(RealizationPlan& plan, + SpatDeferredCommunicationOp deferred, + DeferredBodyAnalysis& body, + ScheduledInfo& target, + unsigned consumerStep, + unsigned targetLane, + unsigned sourceIndex, + LaneTransfer& lane) { + Value sourceValue = deferred.getSources()[sourceIndex]; + auto sourceResult = dyn_cast(sourceValue); + if (!sourceResult) + return deferred.emitOpError("phase 2 requires deferred sources to be original graph results"); + Operation* graphProducer = sourceResult.getOwner(); + if (!isa_and_nonnull(graphProducer)) + return deferred.emitOpError("phase 2 requires deferred sources to be graph compute results"); + auto graphId = graphProducer ? graphProducer->getAttrOfType("scheduled.graph_id") : IntegerAttr(); + if (!graphId) + return deferred.emitOpError("phase 2 cannot identify the serialized graph producer"); + auto sourceUse = findSourceUse(body, sourceIndex); + std::optional sourceLane = sourceUse ? sourceUse->staticSourceLane : std::nullopt; + if (isa(graphProducer) && !sourceLane) + return deferred.emitOpError("phase 2 requires a statically resolved physical source slot for graph batch payloads"); + auto source = findProducer(plan, deferred, graphId.getInt(), sourceResult.getResultNumber(), sourceLane); + if (failed(source)) + return failure(); + if (failed(configurePayloadAndReplacement(deferred, body, **source, sourceIndex, sourceLane.value_or(0), lane))) + return failure(); + + lane.source = *source; + lane.sourceCore = (*source)->core; + lane.targetCore = target.cores[target.isBatch() ? targetLane : 0]; + lane.targetLane = targetLane; + lane.sourceStream = + (*source)->scheduled->streamIds[(*source)->scheduled->isBatch() + ? llvm::find((*source)->scheduled->cores, (*source)->core) - (*source)->scheduled->cores.begin() + : 0]; + lane.targetStream = target.streamIds[target.isBatch() ? targetLane : 0]; + lane.local = (*source)->scheduled == &target && lane.sourceStream == lane.targetStream && (*source)->step < consumerStep; + return success(); +} + +static LogicalResult normalizeDeferred(func::FuncOp funcOp, RealizationPlan& plan) { + DenseMap infoByOp; + for (ScheduledInfo& info : plan.scheduled) + infoByOp[info.op] = &info; + + SmallVector deferredOps; + funcOp.walk([&](SpatDeferredCommunicationOp deferred) { deferredOps.push_back(deferred); }); + uint64_t nextExchangeId = 0; + for (SpatDeferredCommunicationOp deferred : deferredOps) { + Operation* targetOp = deferred->getParentOfType(); + if (!targetOp) + targetOp = deferred->getParentOfType(); + ScheduledInfo* target = infoByOp.lookup(targetOp); + Block* scheduledBlock = target ? getScheduledBlock(deferred, targetOp) : nullptr; + auto step = target && scheduledBlock ? getStepIndex(*target, scheduledBlock) : FailureOr(failure()); + if (!target || failed(step)) + return deferred.emitOpError("phase 2 cannot locate the deferred consumer scheduled step"); + auto body = analyzeDeferredBody(deferred); + if (failed(body)) + return failure(); + + bool batched = target->isBatch(); + if (!batched) { + NormalizedScalarExchange exchange; + exchange.id = nextExchangeId++; + exchange.deferred = deferred; + exchange.target = target; + exchange.consumerStep = *step; + exchange.body = *body; + if (failed(buildLaneTransfer(plan, + deferred, + exchange.body, + *target, + *step, + 0, + 0, + exchange.transfer))) + return failure(); + plan.scalar.push_back(std::move(exchange)); + continue; + } + + size_t lanes = target->cores.size(); + + NormalizedBatchedExchange exchange; + exchange.id = nextExchangeId++; + exchange.deferred = deferred; + exchange.target = target; + exchange.consumerStep = *step; + exchange.body = *body; + for (unsigned laneIndex = 0; laneIndex < lanes; ++laneIndex) { + auto laneBody = analyzeDeferredBody(deferred, laneIndex); + if (failed(laneBody)) + return failure(); + if (laneBody->sourceUses.size() != 1) + return deferred.emitOpError("phase 2 scheduled-batch specialization must resolve exactly one source"); + unsigned sourceIndex = laneBody->sourceUses.front().sourceIndex; + LaneTransfer lane; + if (failed(buildLaneTransfer(plan, + deferred, + *laneBody, + *target, + *step, + laneIndex, + sourceIndex, + lane))) + return failure(); + exchange.lanes.push_back(std::move(lane)); + } + plan.batched.push_back(std::move(exchange)); + } + return success(); +} + +static LogicalResult scheduleCommunication(func::FuncOp funcOp, RealizationPlan& plan) { + llvm::MapVector> groups; + for (auto [index, transfer] : llvm::enumerate(plan.external)) + groups[transfer.parentExchangeId].push_back(index); + + SmallVector completedSteps(plan.stepCounts.size()); + SmallVector completedTransfers(plan.external.size()); + SmallVector scheduled; + while (scheduled.size() != plan.external.size()) { + bool progressed = false; + for (unsigned stream = 0; stream < completedSteps.size(); ++stream) { + if (completedSteps[stream] == plan.stepCounts[stream]) + continue; + unsigned step = completedSteps[stream]; + bool inputsReady = llvm::all_of(llvm::enumerate(plan.external), [&](auto indexed) { + const PlannedCommunicationTransfer& transfer = indexed.value(); + return transfer.targetStream != stream || transfer.consumerStep != step || completedTransfers[indexed.index()]; + }); + if (inputsReady) { + ++completedSteps[stream]; + progressed = true; + } + } + + SmallVector readyGroups; + for (auto& entry : groups) { + if (completedTransfers[entry.second.front()]) + continue; + bool ready = llvm::all_of(entry.second, [&](unsigned index) { + const PlannedCommunicationTransfer& transfer = plan.external[index]; + return completedSteps[transfer.sourceStream] > transfer.producerStep; + }); + if (ready) + readyGroups.push_back(entry.first); + } + if (!readyGroups.empty()) { + llvm::stable_sort(readyGroups, [&](uint64_t lhs, uint64_t rhs) { + auto priority = [&](uint64_t group) { + ArrayRef indices = groups[group]; + unsigned consumer = std::numeric_limits::max(); + unsigned source = std::numeric_limits::max(); + unsigned target = std::numeric_limits::max(); + for (unsigned index : indices) { + const PlannedCommunicationTransfer& transfer = plan.external[index]; + consumer = std::min(consumer, transfer.consumerStep); + source = std::min(source, transfer.sourceStream); + target = std::min(target, transfer.targetStream); + } + return std::tuple(consumer, source, target, group); + }; + return priority(lhs) < priority(rhs); + }); + ArrayRef selected = groups[readyGroups.front()]; + SmallVector ordered(selected.begin(), selected.end()); + llvm::stable_sort(ordered, [&](unsigned lhs, unsigned rhs) { + const auto& left = plan.external[lhs]; + const auto& right = plan.external[rhs]; + return std::tie(left.sourceStream, left.targetStream, left.exchangeId) + < std::tie(right.sourceStream, right.targetStream, right.exchangeId); + }); + for (unsigned index : ordered) { + completedTransfers[index] = true; + PlannedCommunicationTransfer transfer = plan.external[index]; + transfer.sourceInsertionStep = completedSteps[transfer.sourceStream]; + transfer.targetInsertionStep = completedSteps[transfer.targetStream]; + scheduled.push_back(transfer); + } + progressed = true; + } + if (!progressed) + return funcOp.emitOpError("global communication scheduler made no progress before IR mutation"); + } + plan.external = std::move(scheduled); + return success(); +} + +static LogicalResult buildAndVerifyPlan(func::FuncOp funcOp, RealizationPlan& plan) { + for (Operation& op : funcOp.getOps()) { + if (!isa(op)) + continue; + auto graphId = op.getAttrOfType("scheduled.graph_id"); + if (!graphId) + continue; + for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) { + bool hasExternalUse = llvm::any_of(result.getUses(), [](OpOperand& use) { + return !isa(use.getOwner()); + }); + if (!hasExternalUse) + continue; + SmallVector exact; + for (ProducedValue* produced : plan.producedByGraph.lookup(graphId.getInt())) + if (produced->resultIndex == resultIndex && produced->published.getType() == result.getType() + && !llvm::is_contained(exact, produced->published)) + exact.push_back(produced->published); + if (exact.size() != 1) + return op.emitOpError("phase 2 cannot prove unique final publication ownership for a partitioned graph result"); + } + } + + uint64_t nextChannel = 0; + for (NormalizedScalarExchange& exchange : plan.scalar) { + if (exchange.transfer.local) + continue; + exchange.transfer.channelId = nextChannel++; + plan.external.push_back({ + exchange.transfer.channelId, + exchange.id, + exchange.transfer.sourceStream, + exchange.transfer.targetStream, + exchange.transfer.source->step, + exchange.consumerStep, + }); + } + for (NormalizedBatchedExchange& exchange : plan.batched) + for (LaneTransfer& lane : exchange.lanes) { + if (lane.local) + continue; + lane.channelId = nextChannel++; + plan.external.push_back({ + lane.channelId, + exchange.id, + lane.sourceStream, + lane.targetStream, + lane.source->step, + exchange.consumerStep, + }); + } + if (failed(scheduleCommunication(funcOp, plan))) + return failure(); + DenseMap plannedByChannel; + for (PlannedCommunicationTransfer& transfer : plan.external) + plannedByChannel[transfer.exchangeId] = &transfer; + auto applyInsertionSteps = [&](LaneTransfer& lane) { + if (lane.local) + return; + PlannedCommunicationTransfer* transfer = plannedByChannel.lookup(lane.channelId); + assert(transfer && "every external lane must have one planned transfer"); + lane.sourceInsertionStep = transfer->sourceInsertionStep; + lane.targetInsertionStep = transfer->targetInsertionStep; + }; + for (NormalizedScalarExchange& exchange : plan.scalar) + applyInsertionSteps(exchange.transfer); + for (NormalizedBatchedExchange& exchange : plan.batched) + for (LaneTransfer& lane : exchange.lanes) + applyInsertionSteps(lane); + return verifyPlannedCommunicationDeadlockFree(funcOp, plan.stepCounts.size(), plan.stepCounts, plan.external); +} + +static LogicalResult validateScalarLinearization(ScheduledInfo& info) { + auto scheduled = cast(info.op); + for (unsigned index = 1; index < info.blocks.size(); ++index) { + auto previousYield = dyn_cast(info.blocks[index - 1]->getTerminator()); + if (!previousYield || previousYield.getOutputs().size() != info.blocks[index]->getNumArguments()) + return scheduled.emitOpError("phase 2 cannot linearize malformed scalar scheduled blocks"); + for (auto [argument, value] : llvm::zip(info.blocks[index]->getArguments(), previousYield.getOutputs())) { + if (argument.getType() == value.getType()) + continue; + for (Operation* user : argument.getUsers()) + if (!isa(user)) + return scheduled.emitOpError( + "phase 2 cannot linearize a mismatched carried value used by scheduled computation"); + } + } + return success(); +} + +static LogicalResult linearizeScalar(ScheduledInfo& info, IRRewriter& rewriter) { + auto scheduled = cast(info.op); + if (failed(validateScalarLinearization(info))) + return failure(); + Block* first = info.blocks.front(); + SmallVector> incoming(info.blocks.size()); + for (unsigned index = 1; index < info.blocks.size(); ++index) { + auto previousYield = cast(info.blocks[index - 1]->getTerminator()); + incoming[index].assign(previousYield.getOutputs().begin(), previousYield.getOutputs().end()); + } + IRMapping carriedValues; + for (unsigned index = 1; index < info.blocks.size(); ++index) + for (auto [argument, value] : llvm::zip(info.blocks[index]->getArguments(), incoming[index])) { + Value resolved = carriedValues.lookupOrDefault(value); + if (argument.getType() == resolved.getType()) { + carriedValues.map(argument, resolved); + argument.replaceAllUsesWith(resolved); + } + } + + for (unsigned index = 1; index < info.blocks.size(); ++index) { + Block* block = info.blocks[index]; + for (Operation& op : llvm::make_early_inc_range(block->without_terminator())) + op.moveBefore(first->getTerminator()); + } + for (Block* block : info.blocks) + cast(block->getTerminator()).erase(); + + SmallVector outputs(scheduled.getNumResults()); + for (ProducedValue* produced : info.produced) { + unsigned globalResult = info.resultOffsets[produced->step] + produced->resultIndex; + outputs[globalResult] = produced->payload; + } + if (llvm::any_of(outputs, [](Value value) { return !value; })) + return scheduled.emitOpError("phase 2 cannot recover every scheduled scalar result"); + rewriter.setInsertionPointToEnd(first); + SpatYieldOp::create(rewriter, scheduled.getLoc(), outputs); + scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true)); + for (unsigned index = 1; index < info.blocks.size(); ++index) { + for (auto [argumentIndex, argument] : llvm::enumerate(info.blocks[index]->getArguments())) + if (!argument.use_empty()) { + InFlightDiagnostic diagnostic = + scheduled.emitOpError("phase 2 scalar linearization left a live carried block argument"); + diagnostic << " at block " << index << ", argument " << argumentIndex << ", first user " + << (*argument.getUsers().begin())->getName(); + return failure(); + } + info.blocks[index]->erase(); + } + info.blocks.assign(1, first); + return success(); +} + +static LogicalResult linearizeBatch(ScheduledInfo& info, IRRewriter& rewriter) { + auto scheduled = cast(info.op); + Block* first = info.blocks.front(); + SmallVector terminators; + for (Block* block : info.blocks) { + auto inParallel = dyn_cast(block->getTerminator()); + if (!inParallel) + return scheduled.emitOpError("phase 2 cannot linearize a batch block without spat.in_parallel"); + terminators.push_back(inParallel); + } + for (unsigned index = 1; index < info.blocks.size(); ++index) + for (auto [argument, firstArgument] : llvm::zip(info.blocks[index]->getArguments(), first->getArguments())) { + if (argument.getType() != firstArgument.getType()) + return scheduled.emitOpError("phase 2 cannot linearize incompatible batch block arguments"); + argument.replaceAllUsesWith(firstArgument); + } + + for (unsigned index = 1; index < info.blocks.size(); ++index) + for (Operation& op : llvm::make_early_inc_range(info.blocks[index]->without_terminator())) + op.moveBefore(first->getTerminator()); + rewriter.setInsertionPoint(first->getTerminator()); + auto combined = SpatInParallelOp::create(rewriter, scheduled.getLoc()); + Block& combinedBlock = combined.getRegion().front(); + for (SpatInParallelOp terminator : terminators) + for (Operation& op : llvm::make_early_inc_range(terminator.getRegion().front())) + op.moveBefore(&combinedBlock, combinedBlock.end()); + for (SpatInParallelOp terminator : terminators) + terminator.erase(); + scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true)); + for (unsigned index = 1; index < info.blocks.size(); ++index) { + for (auto [argumentIndex, argument] : llvm::enumerate(info.blocks[index]->getArguments())) + if (!argument.use_empty()) { + InFlightDiagnostic diagnostic = + scheduled.emitOpError("phase 2 batch linearization left a live carried block argument"); + diagnostic << " at block " << index << ", argument " << argumentIndex << ", first user " + << (*argument.getUsers().begin())->getName(); + return failure(); + } + info.blocks[index]->erase(); + } + info.blocks.assign(1, first); + return success(); +} + +static void +setTransferAttrs(Operation* op, uint64_t exchangeId, uint64_t channelId, int64_t sourceCore, int64_t targetCore) { + OpBuilder builder(op); + op->setAttr("raptor.exchange_id", builder.getI64IntegerAttr(exchangeId)); + op->setAttr("raptor.channel_id", builder.getI64IntegerAttr(channelId)); + op->setAttr("raptor.source_core", builder.getI64IntegerAttr(sourceCore)); + op->setAttr("raptor.target_core", builder.getI64IntegerAttr(targetCore)); +} + +static FailureOr materializeSourcePayload(DeferredBodyAnalysis& body, LaneTransfer& lane, IRRewriter& rewriter) { + if (body.sourceSideOps.empty()) + return lane.source->payload; + + Block& deferredBlock = + cast(body.yieldedValue.getParentBlock()->getParentOp()).getBody().front(); + IRMapping mapping; + BlockArgument sourceArgument = deferredBlock.getArgument(lane.sourceIndex); + if (sourceArgument.getType() == lane.source->payload.getType()) + mapping.map(sourceArgument, lane.source->payload); + if (lane.bodyRoot) + for (Value result : lane.bodyRoot->getResults()) + if (result.getType() == lane.source->payload.getType()) + mapping.map(result, lane.source->payload); + + for (Operation* op : body.sourceSideOps) { + if (op == lane.bodyRoot) + continue; + bool needed = llvm::any_of(op->getResults(), [&](Value result) { + if (result == body.yieldedValue) + return true; + return llvm::any_of(body.sourceSideOps, [&](Operation* candidate) { + return candidate != lane.bodyRoot && llvm::is_contained(candidate->getOperands(), result); + }); + }); + if (!needed) + continue; + for (Value operand : op->getOperands()) { + if (mapping.contains(operand) || operand.getParentBlock() != &deferredBlock) + continue; + return op->emitOpError("source-side deferred-body clone has an unmapped source operand"); + } + Operation* clone = rewriter.clone(*op, mapping); + for (auto [result, clonedResult] : llvm::zip(op->getResults(), clone->getResults())) + mapping.map(result, clonedResult); + } + Value payload = mapping.lookupOrDefault(body.yieldedValue); + if (!payload || payload.getType() != lane.transportType) + return failure(); + return payload; +} + +static void replaceDeferred(SpatDeferredCommunicationOp deferred, + Value replacement, + SmallVectorImpl& erase) { + deferred.getOutput().replaceAllUsesWith(replacement); + if (!llvm::is_contained(erase, deferred.getOperation())) + erase.push_back(deferred); +} + +static Value createI64TableValue(OpBuilder& builder, Location loc, ArrayRef values, Value index) { + RankedTensorType type = RankedTensorType::get({static_cast(values.size())}, builder.getI64Type()); + Value table = arith::ConstantOp::create(builder, loc, type, DenseElementsAttr::get(type, values)); + Value value = tensor::ExtractOp::create(builder, loc, table, ValueRange {index}); + return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), value); +} + +static void setInsertionAtBoundary(IRRewriter& rewriter, ScheduledInfo& scheduled, unsigned step) { + if (step < scheduled.stepAnchors.size() && scheduled.stepAnchors[step]->getBlock()) { + rewriter.setInsertionPoint(scheduled.stepAnchors[step]); + return; + } + rewriter.setInsertionPoint(scheduled.blocks.front()->getTerminator()); +} + +static unsigned countExternalLanes(const NormalizedBatchedExchange& exchange) { + return llvm::count_if(exchange.lanes, [](const LaneTransfer& lane) { return !lane.local; }); +} + +static LogicalResult emitGuardedBatchSends(NormalizedBatchedExchange& exchange, + ScheduledInfo& source, + ArrayRef lanes, + IRRewriter& rewriter) { + Value laneIndex = cast(source.op).getLaneArgument().value(); + unsigned transferCount = countExternalLanes(exchange); + for (LaneTransfer* lane : lanes) { + auto coreIt = llvm::find(source.cores, lane->sourceCore); + if (coreIt == source.cores.end()) + return exchange.deferred.emitOpError("phase 2 cannot locate a partial batch send source lane"); + setInsertionAtBoundary(rewriter, *lane->source->scheduled, lane->sourceInsertionStep); + Location loc = exchange.deferred.getLoc(); + Value sourceLane = arith::ConstantIndexOp::create(rewriter, loc, std::distance(source.cores.begin(), coreIt)); + Value selected = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, laneIndex, sourceLane); + auto ifOp = scf::IfOp::create(rewriter, loc, selected, false); + rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); + auto payload = materializeSourcePayload(exchange.body, *lane, rewriter); + if (failed(payload)) + return exchange.deferred.emitOpError("failed to materialize partial batch source payload"); + Value channel = arith::ConstantIndexOp::create(rewriter, loc, lane->channelId); + Value sourceCore = arith::ConstantIndexOp::create(rewriter, loc, lane->sourceCore); + Value targetCore = arith::ConstantIndexOp::create(rewriter, loc, lane->targetCore); + auto send = SpatChannelSendOp::create(rewriter, loc, channel, sourceCore, targetCore, *payload); + setTransferAttrs(send, lane->channelId, lane->channelId, lane->sourceCore, lane->targetCore); + send->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); + send->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(transferCount)); + } + return success(); +} + +static LogicalResult emitCompactBatchSends(NormalizedBatchedExchange& exchange, + IRRewriter& rewriter) { + llvm::MapVector> bySource; + for (LaneTransfer& lane : exchange.lanes) + if (!lane.local) + bySource[lane.source->scheduled].push_back(&lane); + + for (auto& entry : bySource) { + ScheduledInfo* source = entry.first; + ArrayRef lanes = entry.second; + if (!source->isBatch()) { + for (LaneTransfer* lane : lanes) { + setInsertionAtBoundary(rewriter, *lane->source->scheduled, lane->sourceInsertionStep); + auto payload = materializeSourcePayload(exchange.body, *lane, rewriter); + if (failed(payload)) + return exchange.deferred.emitOpError("failed to materialize scalar source payload"); + Location loc = exchange.deferred.getLoc(); + Value channel = arith::ConstantIndexOp::create(rewriter, loc, lane->channelId); + Value sourceCore = arith::ConstantIndexOp::create(rewriter, loc, lane->sourceCore); + Value targetCore = arith::ConstantIndexOp::create(rewriter, loc, lane->targetCore); + auto send = SpatChannelSendOp::create(rewriter, loc, channel, sourceCore, targetCore, *payload); + setTransferAttrs(send, lane->channelId, lane->channelId, lane->sourceCore, lane->targetCore); + send->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); + send->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(countExternalLanes(exchange))); + } + continue; + } + if (lanes.size() != source->cores.size()) { + if (failed(emitGuardedBatchSends(exchange, *source, lanes, rewriter))) + return failure(); + continue; + } + SmallVector channels(source->cores.size(), -1); + SmallVector sourceCores(source->cores.size(), -1); + SmallVector targetCores(source->cores.size(), -1); + Value sourcePayload = lanes.front()->source->payload; + for (LaneTransfer* lane : lanes) { + auto coreIt = llvm::find(source->cores, lane->sourceCore); + if (coreIt == source->cores.end() || lane->source->payload != sourcePayload + || lane->sourceIndex != lanes.front()->sourceIndex || lane->bodyRoot != lanes.front()->bodyRoot + || lane->transportType != lanes.front()->transportType) + return exchange.deferred.emitOpError("phase 2 cannot compact heterogeneous batch sends"); + unsigned sourceLane = std::distance(source->cores.begin(), coreIt); + channels[sourceLane] = lane->channelId; + sourceCores[sourceLane] = lane->sourceCore; + targetCores[sourceLane] = lane->targetCore; + } + if (llvm::is_contained(channels, -1)) + return exchange.deferred.emitOpError("phase 2 compact batch send has an unmapped physical lane"); + + setInsertionAtBoundary(rewriter, *lanes.front()->source->scheduled, lanes.front()->sourceInsertionStep); + auto payload = materializeSourcePayload(exchange.body, *lanes.front(), rewriter); + if (failed(payload)) + return exchange.deferred.emitOpError("failed to materialize compact batch source payload"); + Location loc = exchange.deferred.getLoc(); + Value laneIndex = cast(source->op).getLaneArgument().value(); + Value channel = createI64TableValue(rewriter, loc, channels, laneIndex); + Value sourceCore = createI64TableValue(rewriter, loc, sourceCores, laneIndex); + Value targetCore = createI64TableValue(rewriter, loc, targetCores, laneIndex); + auto send = SpatChannelSendOp::create(rewriter, loc, channel, sourceCore, targetCore, *payload); + send->setAttr("raptor.batch_channel_ids", rewriter.getDenseI64ArrayAttr(channels)); + send->setAttr("raptor.batch_source_cores", rewriter.getDenseI64ArrayAttr(sourceCores)); + send->setAttr("raptor.batch_target_cores", rewriter.getDenseI64ArrayAttr(targetCores)); + send->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); + send->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(countExternalLanes(exchange))); + } + return success(); +} + +static LogicalResult +emitCompactBatchReceive(NormalizedBatchedExchange& exchange, IRRewriter& rewriter, SmallVectorImpl& erase) { + ScheduledInfo& target = *exchange.target; + if (!target.isBatch() || exchange.lanes.size() != target.cores.size() + || llvm::any_of(exchange.lanes, [](const LaneTransfer& lane) { return lane.local; })) + return exchange.deferred.emitOpError( + "phase 2 compact batch receive requires one external transfer for every target lane"); + + SmallVector channels(target.cores.size(), -1); + SmallVector sourceCores(target.cores.size(), -1); + SmallVector targetCores(target.cores.size(), -1); + Type payloadType = exchange.lanes.front().transportType; + for (LaneTransfer& lane : exchange.lanes) { + if (lane.targetLane >= target.cores.size() || lane.transportType != payloadType) + return exchange.deferred.emitOpError("phase 2 cannot compact heterogeneous batch receives"); + channels[lane.targetLane] = lane.channelId; + sourceCores[lane.targetLane] = lane.sourceCore; + targetCores[lane.targetLane] = lane.targetCore; + } + if (llvm::is_contained(channels, -1)) + return exchange.deferred.emitOpError("phase 2 compact batch receive has an unmapped target lane"); + + setInsertionAtBoundary(rewriter, target, exchange.lanes.front().targetInsertionStep); + Location loc = exchange.deferred.getLoc(); + Value laneIndex = cast(target.op).getLaneArgument().value(); + Value channel = createI64TableValue(rewriter, loc, channels, laneIndex); + Value sourceCore = createI64TableValue(rewriter, loc, sourceCores, laneIndex); + Value targetCore = createI64TableValue(rewriter, loc, targetCores, laneIndex); + auto receive = SpatChannelReceiveOp::create(rewriter, loc, payloadType, channel, sourceCore, targetCore); + receive->setAttr("raptor.batch_channel_ids", rewriter.getDenseI64ArrayAttr(channels)); + receive->setAttr("raptor.batch_source_cores", rewriter.getDenseI64ArrayAttr(sourceCores)); + receive->setAttr("raptor.batch_target_cores", rewriter.getDenseI64ArrayAttr(targetCores)); + receive->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); + receive->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(countExternalLanes(exchange))); + replaceDeferred(exchange.deferred, receive.getOutput(), erase); + return success(); +} + +static LogicalResult realizeScalarExchange(NormalizedScalarExchange& exchange, + IRRewriter& rewriter, + SmallVectorImpl& erase) { + LaneTransfer& transfer = exchange.transfer; + if (transfer.local) + rewriter.setInsertionPointAfterValue(transfer.source->payload); + else + setInsertionAtBoundary(rewriter, *transfer.source->scheduled, transfer.sourceInsertionStep); + auto payload = materializeSourcePayload(exchange.body, transfer, rewriter); + if (failed(payload)) + return exchange.deferred.emitOpError("failed to materialize scalar source payload"); + if (transfer.local) { + replaceDeferred(exchange.deferred, *payload, erase); + return success(); + } + + Location loc = exchange.deferred.getLoc(); + Value channel = arith::ConstantIndexOp::create(rewriter, loc, transfer.channelId); + Value sourceCore = arith::ConstantIndexOp::create(rewriter, loc, transfer.sourceCore); + Value targetCore = arith::ConstantIndexOp::create(rewriter, loc, transfer.targetCore); + auto createSend = [&] { + auto send = SpatChannelSendOp::create(rewriter, loc, channel, sourceCore, targetCore, *payload); + setTransferAttrs(send, transfer.channelId, transfer.channelId, transfer.sourceCore, transfer.targetCore); + send->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); + send->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(1)); + }; + if (transfer.source->scheduled->isBatch()) { + auto batch = cast(transfer.source->scheduled->op); + auto it = llvm::find(transfer.source->scheduled->cores, transfer.sourceCore); + if (it == transfer.source->scheduled->cores.end()) + return exchange.deferred.emitOpError("phase 2 cannot locate scalar exchange source batch lane"); + Value lane = arith::ConstantIndexOp::create(rewriter, loc, std::distance(transfer.source->scheduled->cores.begin(), it)); + Value selected = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, batch.getLaneArgument().value(), lane); + auto ifOp = scf::IfOp::create(rewriter, loc, selected, false); + rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); + createSend(); + } else { + createSend(); + } + + setInsertionAtBoundary(rewriter, *exchange.target, transfer.targetInsertionStep); + channel = arith::ConstantIndexOp::create(rewriter, loc, transfer.channelId); + sourceCore = arith::ConstantIndexOp::create(rewriter, loc, transfer.sourceCore); + targetCore = arith::ConstantIndexOp::create(rewriter, loc, transfer.targetCore); + auto receive = SpatChannelReceiveOp::create(rewriter, loc, transfer.transportType, channel, sourceCore, targetCore); + setTransferAttrs(receive, transfer.channelId, transfer.channelId, transfer.sourceCore, transfer.targetCore); + receive->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); + receive->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(1)); + replaceDeferred(exchange.deferred, receive.getOutput(), erase); + return success(); +} + +static LogicalResult realizeBatchedExchange(NormalizedBatchedExchange& exchange, + IRRewriter& rewriter, + SmallVectorImpl& erase) { + bool allLocal = llvm::all_of(exchange.lanes, [](const LaneTransfer& lane) { return lane.local; }); + if (allLocal) { + LaneTransfer& first = exchange.lanes.front(); + for (LaneTransfer& lane : exchange.lanes) + if (lane.source->payload != first.source->payload || lane.sourceIndex != first.sourceIndex + || lane.bodyRoot != first.bodyRoot) + return exchange.deferred.emitOpError("phase 2 cannot compact heterogeneous local batch payloads"); + rewriter.setInsertionPointAfterValue(first.source->payload); + auto payload = materializeSourcePayload(exchange.body, first, rewriter); + if (failed(payload)) + return exchange.deferred.emitOpError("failed to materialize local compact batch payload"); + replaceDeferred(exchange.deferred, *payload, erase); + return success(); + } + if (llvm::any_of(exchange.lanes, [](const LaneTransfer& lane) { return lane.local; })) + return exchange.deferred.emitOpError( + "phase 2 cannot compact a batch exchange mixing local and external target lanes"); + if (failed(emitCompactBatchSends(exchange, rewriter)) + || failed(emitCompactBatchReceive(exchange, rewriter, erase))) + return failure(); + return success(); +} + +static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { + IRRewriter rewriter(funcOp.getContext()); + for (ScheduledInfo& info : plan.scheduled) { + if (info.isBatch()) { + if (failed(linearizeBatch(info, rewriter))) + return failure(); + } + else if (failed(linearizeScalar(info, rewriter))) { + return failure(); + } + } + + SmallVector erase; + for (NormalizedScalarExchange& exchange : plan.scalar) + if (exchange.transfer.local && failed(realizeScalarExchange(exchange, rewriter, erase))) + return failure(); + for (NormalizedBatchedExchange& exchange : plan.batched) + if (llvm::all_of(exchange.lanes, [](const LaneTransfer& lane) { return lane.local; }) + && failed(realizeBatchedExchange(exchange, rewriter, erase))) + return failure(); + + DenseMap orderByParent; + for (auto [order, transfer] : llvm::enumerate(plan.external)) + orderByParent.try_emplace(transfer.parentExchangeId, order); + struct ExternalExchange { + unsigned order; + NormalizedScalarExchange* scalar = nullptr; + NormalizedBatchedExchange* batched = nullptr; + }; + SmallVector external; + for (NormalizedScalarExchange& exchange : plan.scalar) + if (!exchange.transfer.local) + external.push_back({orderByParent.lookup(exchange.id), &exchange, nullptr}); + for (NormalizedBatchedExchange& exchange : plan.batched) + if (llvm::any_of(exchange.lanes, [](const LaneTransfer& lane) { return !lane.local; })) + external.push_back({orderByParent.lookup(exchange.id), nullptr, &exchange}); + llvm::sort(external, [](const ExternalExchange& lhs, const ExternalExchange& rhs) { return lhs.order < rhs.order; }); + for (ExternalExchange& exchange : external) { + LogicalResult result = exchange.scalar ? realizeScalarExchange(*exchange.scalar, rewriter, erase) + : realizeBatchedExchange(*exchange.batched, rewriter, erase); + if (failed(result)) + return failure(); + } + for (Operation* op : erase) + if (op && op->getBlock() && isa(op)) { + if (!op->use_empty()) + return op->emitOpError("phase 2 cannot erase deferred communication with live uses"); + rewriter.eraseOp(op); + } + + for (Operation& op : funcOp.getOps()) { + if (!isa(op)) + continue; + auto graphId = op.getAttrOfType("scheduled.graph_id"); + if (!graphId) + continue; + for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) { + SmallVector externalUses; + for (OpOperand& use : result.getUses()) + if (!isa(use.getOwner())) + externalUses.push_back(&use); + if (externalUses.empty()) + continue; + SmallVector exact; + for (ProducedValue* produced : plan.producedByGraph.lookup(graphId.getInt())) + if (produced->resultIndex == resultIndex && produced->published.getType() == result.getType() + && !llvm::is_contained(exact, produced->published)) + exact.push_back(produced->published); + if (exact.size() == 1) { + for (OpOperand* use : externalUses) + use->set(exact.front()); + continue; + } + return op.emitOpError("phase 2 final publication ownership changed after planning"); + } + } + + SmallVector graphOps; + for (Operation& op : funcOp.getOps()) + if (isa(op)) + graphOps.push_back(&op); + for (Operation* op : llvm::reverse(graphOps)) { + if (!op->use_empty()) { + InFlightDiagnostic diagnostic = + op->emitOpError("phase 2 cannot erase an old graph compute with an unreconstructed live result"); + for (auto [resultIndex, result] : llvm::enumerate(op->getResults())) + if (!result.use_empty()) { + diagnostic << "; result " << resultIndex << " first used by " << (*result.getUsers().begin())->getName(); + break; + } + return failure(); + } + rewriter.eraseOp(op); + } + return success(); +} + +} // namespace + +LogicalResult realizeDeferredCommunication(func::FuncOp funcOp) { + RealizationPlan plan; + if (failed(collectScheduled(funcOp, plan)) || failed(normalizeDeferred(funcOp, plan)) + || failed(buildAndVerifyPlan(funcOp, plan))) + return failure(); + if (failed(realizePlan(funcOp, plan))) + return failure(); + + bool hasDeferred = false; + funcOp.walk([&](SpatDeferredCommunicationOp deferred) { + deferred.emitOpError("phase 2 left an unrealized deferred communication"); + hasDeferred = true; + }); + if (hasDeferred) + return failure(); + return verifyRealizedCommunicationDeadlockFree(funcOp); +} + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp new file mode 100644 index 0000000..b7c3695 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp @@ -0,0 +1,9 @@ +#pragma once + +#include "mlir/Dialect/Func/IR/FuncOps.h" + +namespace onnx_mlir::spatial { + +mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp); + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp index 4e012b2..c725c05 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp @@ -1,6 +1,8 @@ #include "ScheduledComputeMaterialization.hpp" #include "ScheduledComputeReport.hpp" #include "ScheduledComputeVerification.hpp" +#include "DeferredCommunicationRealization.hpp" +#include "DeferredCommunicationDeadlock.hpp" #include "mlir/Pass/Pass.h" @@ -80,9 +82,20 @@ struct MergeComputeNodesPass final : PassWrapperpeftClassPlans, materialization->materializedSchedules); - moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; " - "Phase 2 communication realization is not implemented"); - signalPassFailure(); + if (failed(realizeDeferredCommunication(funcOp))) { + moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed"); + signalPassFailure(); + return; + } + if (failed(verifyRealizedCommunicationDeadlockFree(funcOp))) { + moduleOp.emitError("MergeComputeNodes final communication verification failed"); + signalPassFailure(); + return; + } + if (failed(verifyScheduledSpatialInvariants(funcOp))) { + moduleOp.emitError("scheduled Spatial phase 2 verification failed"); + signalPassFailure(); + } } }; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp index 515686c..fc856a9 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp @@ -10,6 +10,8 @@ #include +#include "llvm/ADT/SmallPtrSet.h" + namespace onnx_mlir { namespace spatial { @@ -19,10 +21,6 @@ namespace { struct BatchFragmentSpec { RankedTensorType resultType; RankedTensorType sourceSliceType; - unsigned laneDim = 0; - SmallVector offsets; - SmallVector sizes; - SmallVector strides; }; static SmallVector remapMixedOffsets(ArrayRef mixedOffsets, IRMapping &mapper) { @@ -128,42 +126,38 @@ static FailureOr getBatchFragmentSpec(SpatComputeBatch batch, RankedTensorType sourceType = insert.getSourceType(); if (!destType || !sourceType || !destType.hasStaticShape() || !sourceType.hasStaticShape()) return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); - if (destType.getRank() == 0 || sourceType.getRank() == 0 || destType.getRank() != sourceType.getRank()) - return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + if (destType.getRank() != sourceType.getRank() + 1 || destType.getDimSize(0) != batch.getLaneCount() + || destType.getElementType() != sourceType.getElementType()) + return batch.emitOpError("graph_compute_batch result must be a leading physical-slot dimension followed by its fragment"); + if (!llvm::equal(destType.getShape().drop_front(), sourceType.getShape())) + return batch.emitOpError("graph_compute_batch result trailing shape must match its published fragment"); if (!insert.hasUnitStride()) return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); - auto offsets = insert.getMixedOffsets(); auto sizes = insert.getMixedSizes(); auto strides = insert.getMixedStrides(); if (offsets.size() != static_cast(destType.getRank()) || sizes.size() != static_cast(destType.getRank()) || strides.size() != static_cast(destType.getRank())) return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); - - std::optional laneDim; - for (unsigned dim = 0; dim < offsets.size(); ++dim) { - bool dimLooksLaneLike = sourceType.getShape()[dim] != destType.getShape()[dim]; - if (auto value = dyn_cast(offsets[dim])) - dimLooksLaneLike = dimLooksLaneLike || valueTransitivelyDependsOn(value, *batch.getLaneArgument()); - if (dimLooksLaneLike) { - if (laneDim && *laneDim != dim) - return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); - laneDim = dim; - } + if (!isa(offsets.front()) || !valueTransitivelyDependsOn(cast(offsets.front()), *batch.getLaneArgument())) + return batch.emitOpError("graph_compute_batch publication must select its physical slot in dimension zero"); + for (unsigned dim = 1; dim < offsets.size(); ++dim) { + auto offset = dyn_cast(offsets[dim]); + auto integer = dyn_cast_or_null(offset); + if (!integer || integer.getInt() != 0) + return batch.emitOpError("graph_compute_batch publication must have zero trailing offsets"); } - if (!laneDim) - return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); - - SmallVector shape(destType.getShape().begin(), destType.getShape().end()); - shape[*laneDim] = fragmentLaneCount; - return BatchFragmentSpec { - RankedTensorType::get(shape, destType.getElementType()), - sourceType, - *laneDim, - SmallVector(offsets.begin(), offsets.end()), - SmallVector(sizes.begin(), sizes.end()), - SmallVector(strides.begin(), strides.end()) + auto staticIndex = [](OpFoldResult value) -> std::optional { + auto attr = dyn_cast(value); + auto integer = dyn_cast_or_null(attr); + return integer ? std::optional(integer.getInt()) : std::nullopt; }; + if (staticIndex(sizes.front()) != 1) + return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape"); + for (auto [size, dim] : llvm::zip_equal(ArrayRef(sizes).drop_front(), sourceType.getShape())) + if (staticIndex(size) != dim) + return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape"); + return BatchFragmentSpec {spatial::getGraphBatchPhysicalResultType(fragmentLaneCount, sourceType), sourceType}; } return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); @@ -347,9 +341,12 @@ static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClass return success(); } -static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper, SmallVectorImpl &yieldedValues) { +static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper, + SmallVectorImpl &yieldedValues, + const llvm::SmallPtrSetImpl &absorbed) { for (Operation &op : source.without_terminator()) - builder.clone(op, mapper); + if (!absorbed.contains(&op)) + builder.clone(op, mapper); auto yield = cast(source.getTerminator()); for (Value output : yield.getOutputs()) yieldedValues.push_back(mapper.lookup(output)); @@ -362,7 +359,6 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew ValueRange scheduledInputs, Block &block, const MergeScheduleResult &schedule, - uint64_t &exchangeId, SmallVectorImpl &yieldedValues) { SmallVector initResults; SmallVector fragmentSpecs; @@ -373,7 +369,7 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew fragmentSpecs.push_back(*spec); auto empty = createEmptyTensorForType(rewriter, batch.getLoc(), spec->resultType); if (failed(empty)) - return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + return batch.emitOpError("scheduled materialization requires a physical graph fragment publication"); initResults.push_back(*empty); } @@ -393,28 +389,35 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew mapper.map(*batch.getLaneArgument(), originalLane); for (auto [index, weight] : llvm::enumerate(batch.getWeights())) mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight)); + SmallVector inputPlans; for (auto [index, input] : llvm::enumerate(batch.getInputs())) { - Value mapped; - if (failed(mapSingleCpuInput(builder, + DeferredInputPlan plan; + if (failed(prepareSingleCpuInput(builder, input.getLoc(), input, + *batch.getInputArgument(index), instance, schedule, scheduledInputs, block, scheduledWeights.size(), ArrayRef {}, - exchangeId, - mapped))) + *batch.getLaneArgument(), + lower, + plan))) return failure(); - mapper.map(*batch.getInputArgument(index), mapped); + inputPlans.push_back(std::move(plan)); } for (auto [index, outputArg] : llvm::enumerate(batch.getOutputs())) (void)outputArg, mapper.map(*batch.getOutputArgument(index), iterArgs[index]); Block &source = batch.getBody().front(); + llvm::SmallPtrSet absorbed; + if (failed(materializeDeferredPayloadDemands(builder, bodyLoc, source, inputPlans, mapper, absorbed))) + return failure(); for (Operation &op : source.without_terminator()) - builder.clone(op, mapper); + if (!absorbed.contains(&op)) + builder.clone(op, mapper); auto inParallel = dyn_cast(source.getTerminator()); if (!inParallel) @@ -432,13 +435,13 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew for (Operation &op : inParallel.getRegion().front()) { auto insert = dyn_cast(&op); if (!insert) - return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + return batch.emitOpError("scheduled materialization requires a physical graph fragment publication"); auto oldDest = dyn_cast(insert.getDest()); if (!oldDest || !outputIndexByArg.count(oldDest)) - return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"), failure(); + return batch.emitOpError("scheduled materialization requires a physical graph fragment publication"), failure(); size_t resultIndex = outputIndexByArg.lookup(oldDest); SmallVector offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper); - offsets[fragmentSpecs[resultIndex].laneDim] = localLane; + offsets.front() = localLane; current[resultIndex] = tensor::InsertSliceOp::create(builder, insert.getLoc(), mapper.lookup(insert.getSource()), @@ -463,8 +466,7 @@ static LogicalResult materializeSingleCpuPeftClass( const PeftClassPlan &peftClassPlan, const MergeScheduleResult &schedule, DenseMap &graphComputeToBlockMap, - ScheduledMaterializationRecord &record, - uint64_t &exchangeId) { + ScheduledMaterializationRecord &record) { size_t cpu = peftClassPlan.cpus.front(); auto instancesIt = peftClassPlan.instancesByCpu.find(cpu); assert(instancesIt != peftClassPlan.instancesByCpu.end() && "missing single-cpu schedule"); @@ -487,23 +489,29 @@ static LogicalResult materializeSingleCpuPeftClass( IRMapping mapper; for (auto [index, weight] : llvm::enumerate(compute.getWeights())) mapper.map(*compute.getWeightArgument(index), getBlockOperand(*block, scheduled.getWeights(), weight)); + SmallVector inputPlans; for (auto [index, input] : llvm::enumerate(compute.getInputs())) { - Value mapped; - if (failed(mapSingleCpuInput(rewriter, + DeferredInputPlan plan; + if (failed(prepareSingleCpuInput(rewriter, input.getLoc(), input, + *compute.getInputArgument(index), instance, schedule, scheduled.getInputs(), *block, scheduled.getWeights().size(), carriedKeys, - exchangeId, - mapped))) + {}, + {}, + plan))) return failure(); - mapper.map(*compute.getInputArgument(index), mapped); + inputPlans.push_back(std::move(plan)); } - cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues); + llvm::SmallPtrSet absorbed; + if (failed(materializeDeferredPayloadDemands(rewriter, compute.getLoc(), compute.getBody().front(), inputPlans, mapper, absorbed))) + return failure(); + cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues, absorbed); } else { auto batch = cast(instance.op); if (failed(materializeResultfulBatchChunkAsScalar(rewriter, @@ -513,7 +521,6 @@ static LogicalResult materializeSingleCpuPeftClass( scheduled.getInputs(), *block, schedule, - exchangeId, yieldedValues))) return failure(); } @@ -556,8 +563,7 @@ static SmallVector buildScheduledOutputInsertOffsets(OpBuilder &bu Location loc, Value scheduledLane, int64_t lanesPerScheduledLane, - RankedTensorType localFragmentType, - unsigned laneDim) { + RankedTensorType localFragmentType) { SmallVector offsets; Value scheduledOutputLane = scheduledLane; if (lanesPerScheduledLane != 1) { @@ -570,8 +576,8 @@ static SmallVector buildScheduledOutputInsertOffsets(OpBuilder &bu ValueRange {scheduledLane}) .getResult(); } - for (unsigned dim = 0; dim < static_cast(localFragmentType.getRank()); ++dim) - offsets.push_back(dim == laneDim ? OpFoldResult(scheduledOutputLane) : OpFoldResult(builder.getIndexAttr(0))); + offsets.push_back(scheduledOutputLane); + offsets.append(localFragmentType.getRank() - 1, OpFoldResult(builder.getIndexAttr(0))); return offsets; } @@ -581,8 +587,7 @@ static LogicalResult materializeMultiCpuPeftClass( const PeftClassPlan &peftClassPlan, const MergeScheduleResult &schedule, DenseMap &graphComputeToBlockMap, - ScheduledMaterializationRecord &record, - uint64_t &exchangeId) { + ScheduledMaterializationRecord &record) { std::map, Value> laneStartTableCache; ArrayRef stepPlans = record.stepPlans; for (const ScheduledStepPlan &stepPlan : stepPlans) { @@ -615,31 +620,37 @@ static LogicalResult materializeMultiCpuPeftClass( Value scheduledLane = block->getArgument(0); const ComputeInstance &representative = stepTuple.instances.front(); SmallVector finalLocalFragments; - SmallVector finalLaneDims; if (auto compute = dyn_cast(representative.op)) { IRMapping mapper; for (auto [index, weight] : llvm::enumerate(compute.getWeights())) mapper.map(*compute.getWeightArgument(index), block->getArgument(1 + index)); unsigned firstInputArg = 1 + scheduled.getWeights().size(); + SmallVector inputPlans; for (auto [index, input] : llvm::enumerate(compute.getInputs())) { - Value mapped; - if (failed(mapMultiCpuTupleInput(rewriter, + DeferredInputPlan plan; + if (failed(prepareMultiCpuTupleInput(rewriter, input.getLoc(), input, + *compute.getInputArgument(index), stepTuple, peftClassPlan, schedule, scheduled.getInputs(), *block, firstInputArg, - exchangeId, - mapped))) + {}, + {}, + scheduledLane, + plan))) return failure(); - mapper.map(*compute.getInputArgument(index), mapped); + inputPlans.push_back(std::move(plan)); } SmallVector yieldedValues; - cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues); + llvm::SmallPtrSet absorbed; + if (failed(materializeDeferredPayloadDemands(rewriter, compute.getLoc(), compute.getBody().front(), inputPlans, mapper, absorbed))) + return failure(); + cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues, absorbed); for (Value yielded : yieldedValues) { auto tensorType = dyn_cast(yielded.getType()); if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0) @@ -656,7 +667,6 @@ static LogicalResult materializeMultiCpuPeftClass( yielded, reassociation) .getResult()); - finalLaneDims.push_back(0); } } else { auto batch = cast(representative.op); @@ -702,26 +712,34 @@ static LogicalResult materializeMultiCpuPeftClass( for (auto [index, weight] : llvm::enumerate(batch.getWeights())) mapper.map(*batch.getWeightArgument(index), block->getArgument(1 + index)); unsigned firstInputArg = 1 + scheduled.getWeights().size(); + SmallVector inputPlans; for (auto [index, input] : llvm::enumerate(batch.getInputs())) { - Value mapped; - if (failed(mapMultiCpuTupleInput(builder, + DeferredInputPlan plan; + if (failed(prepareMultiCpuTupleInput(builder, input.getLoc(), input, + *batch.getInputArgument(index), stepTuple, peftClassPlan, schedule, scheduled.getInputs(), *block, firstInputArg, - exchangeId, - mapped))) + *batch.getLaneArgument(), + sourceLane, + scheduledLane, + plan))) return failure(); - mapper.map(*batch.getInputArgument(index), mapped); + inputPlans.push_back(std::move(plan)); } for (unsigned index = 0; index < batch.getNumResults(); ++index) mapper.map(*batch.getOutputArgument(index), iterArgs[index]); + llvm::SmallPtrSet absorbed; + if (failed(materializeDeferredPayloadDemands(builder, bodyLoc, batch.getBody().front(), inputPlans, mapper, absorbed))) + return failure(); for (Operation &op : batch.getBody().front().without_terminator()) - builder.clone(op, mapper); + if (!absorbed.contains(&op)) + builder.clone(op, mapper); auto inParallel = dyn_cast(batch.getBody().front().getTerminator()); if (!inParallel) @@ -735,13 +753,13 @@ static LogicalResult materializeMultiCpuPeftClass( for (Operation &op : inParallel.getRegion().front()) { auto insert = dyn_cast(&op); if (!insert) - return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + return batch.emitOpError("scheduled materialization requires a physical graph fragment publication"); auto oldDest = dyn_cast(insert.getDest()); if (!oldDest || !outputIndexByArg.count(oldDest)) - return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"), failure(); + return batch.emitOpError("scheduled materialization requires a physical graph fragment publication"), failure(); size_t resultIndex = outputIndexByArg.lookup(oldDest); SmallVector offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper); - offsets[fragmentSpecs[resultIndex].laneDim] = innerLane; + offsets.front() = innerLane; current[resultIndex] = tensor::InsertSliceOp::create(builder, insert.getLoc(), mapper.lookup(insert.getSource()), @@ -757,8 +775,6 @@ static LogicalResult materializeMultiCpuPeftClass( if (failed(loop)) return failure(); finalLocalFragments.assign(loop->results.begin(), loop->results.end()); - for (const BatchFragmentSpec &spec : fragmentSpecs) - finalLaneDims.push_back(spec.laneDim); } auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc()); @@ -771,8 +787,7 @@ static LogicalResult materializeMultiCpuPeftClass( scheduled.getLoc(), scheduledLane, lanesPerScheduledLane, - localFragmentType, - finalLaneDims[resultIndex]); + localFragmentType); SmallVector sizes; SmallVector strides; for (int64_t dim : localFragmentType.getShape()) { @@ -800,6 +815,14 @@ FailureOr materializeScheduledCompute(func::FuncOp funcOp, const MergeScheduleResult &schedule, PatternRewriter &rewriter) { + DenseMap graphIds; + int64_t nextGraphId = 0; + for (Operation &op : funcOp.getOps()) + if (isa(op)) { + graphIds[&op] = nextGraphId; + op.setAttr("scheduled.graph_id", rewriter.getI64IntegerAttr(nextGraphId++)); + } + llvm::MapVector peftClassPlans; for (const ComputeInstance &instance : schedule.dominanceOrderCompute) { size_t cpu = schedule.computeToCpuMap.lookup(instance); @@ -829,7 +852,6 @@ materializeScheduledCompute(func::FuncOp funcOp, DenseMap scheduledComputeBatches; DenseMap classToRecordIndex; std::vector materializedSchedules; - uint64_t exchangeId = 0; for (auto &entry : peftClassPlans) { PeftClassPlan &peftClassPlan = entry.second; @@ -855,9 +877,11 @@ materializeScheduledCompute(func::FuncOp funcOp, SmallVector stepResultCounts; SmallVector sourceLaneStarts; SmallVector sourceLaneCounts; + SmallVector stepSourceIds; size_t resultOffset = 0; for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front())) { stepSources.push_back(rewriter.getStringAttr(getInstanceName(instance))); + stepSourceIds.push_back(graphIds.lookup(instance.op)); sourceLaneSelectors.push_back(rewriter.getStringAttr(isa(instance.op) ? "scalar" : "affine")); size_t resultCount = getComputeInstanceResultValueCount(instance); stepResultOffsets.push_back(static_cast(resultOffset)); @@ -872,6 +896,7 @@ materializeScheduledCompute(func::FuncOp funcOp, } } scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources)); + scheduled->setAttr("scheduled.step_source_ids", rewriter.getDenseI64ArrayAttr(stepSourceIds)); scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(stepResultOffsets)); scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(stepResultCounts)); scheduled->setAttr("scheduled.source_lane_starts", rewriter.getDenseI64ArrayAttr(sourceLaneStarts)); @@ -894,8 +919,10 @@ materializeScheduledCompute(func::FuncOp funcOp, SmallVector resultCounts; SmallVector sourceLaneStarts; SmallVector sourceLaneCounts; + SmallVector stepSourceIds; for (const ScheduledStepPlan &stepPlan : record.stepPlans) { stepSources.push_back(rewriter.getStringAttr(getInstanceName(stepPlan.stepTuple.instances.front()))); + stepSourceIds.push_back(graphIds.lookup(stepPlan.stepTuple.instances.front().op)); sourceLaneSelectors.push_back(rewriter.getStringAttr(usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table")); resultOffsets.push_back(static_cast(stepPlan.resultOffset)); resultCounts.push_back(static_cast(stepPlan.resultCount)); @@ -908,6 +935,7 @@ materializeScheduledCompute(func::FuncOp funcOp, {static_cast(record.stepPlans.size()), static_cast(peftClassPlan.cpus.size())}, rewriter.getI64Type()); scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources)); + scheduled->setAttr("scheduled.step_source_ids", rewriter.getDenseI64ArrayAttr(stepSourceIds)); scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(resultOffsets)); scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(resultCounts)); scheduled->setAttr("scheduled.source_lane_starts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef(sourceLaneStarts))); @@ -931,8 +959,7 @@ materializeScheduledCompute(func::FuncOp funcOp, peftClassPlan, schedule, graphComputeToBlockMap, - record, - exchangeId))) + record))) return failure(); } else { if (failed(materializeMultiCpuPeftClass(rewriter, @@ -940,8 +967,7 @@ materializeScheduledCompute(func::FuncOp funcOp, peftClassPlan, schedule, graphComputeToBlockMap, - record, - exchangeId))) + record))) return failure(); } } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp index cb19067..f9f48a4 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp @@ -80,22 +80,6 @@ struct SourceLaneSelector { SmallVector tableValues; }; -struct DeferredEndpoint { - ComputeInstance instance; - size_t cpu = 0; - size_t canonicalClassId = 0; - uint32_t laneStart = 0; - uint32_t laneCount = 1; -}; - -struct DeferredTransferMetadata { - SmallVector producers; - SmallVector consumers; - bool isBatched = false; - SmallVector sourceOperandForScheduledLane; - bool multiSourcePayload = false; -}; - struct ScheduledMaterializationRecord { Operation *scheduledOp = nullptr; size_t canonicalPeftClassId = 0; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp index 455deb5..097711b 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp @@ -81,59 +81,18 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) { pim::CappedDiagnosticReporter diagnostics; funcOp.walk([&](SpatDeferredCommunicationOp transfer) { for (Value source : transfer.getSources()) { - if (isa_and_nonnull(source.getDefiningOp())) { + auto result = dyn_cast(source); + if (!result || !isa(result.getOwner())) { diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { - illegalOp->emitOpError("phase-check deferred communication source operand must remain an original graph SSA producer value"); + illegalOp->emitOpError("phase-check deferred communication source operand must be an original graph SSA result"); }); } } - - if (auto scheduled = dyn_cast(transfer->getParentOp())) { - auto batchedAttr = transfer->getAttrOfType("batched"); - auto sourceOperandForScheduledLane = - transfer->getAttrOfType("source_operand_for_scheduled_lane"); - auto multiSourcePayload = transfer->getAttrOfType("multi_source_payload"); - auto sourceCpus = transfer->getAttrOfType("source_cpus"); - auto sourceClasses = transfer->getAttrOfType("source_classes"); - auto sourceLaneRanges = transfer->getAttrOfType("source_lane_ranges"); - auto targetCpus = transfer->getAttrOfType("target_cpus"); - auto targetClasses = transfer->getAttrOfType("target_classes"); - auto targetLaneRanges = transfer->getAttrOfType("target_lane_ranges"); - int64_t laneCount = scheduled.getLaneCount(); - auto hasExpectedSize = [&](DenseI64ArrayAttr attr, int64_t expected, StringRef name) { - if (!attr || attr.size() != expected) - diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { - illegalOp->emitOpError() << "phase-check scheduled batch deferred communication requires " << name - << " length " << expected; - }); - }; - if (!batchedAttr || !batchedAttr.getValue()) - diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { - illegalOp->emitOpError("phase-check scheduled batch deferred communication requires batched metadata"); - }); - hasExpectedSize(sourceCpus, laneCount, "source_cpus"); - hasExpectedSize(sourceClasses, laneCount, "source_classes"); - hasExpectedSize(targetCpus, laneCount, "target_cpus"); - hasExpectedSize(targetClasses, laneCount, "target_classes"); - hasExpectedSize(sourceLaneRanges, laneCount * 2, "source_lane_ranges"); - hasExpectedSize(targetLaneRanges, laneCount * 2, "target_lane_ranges"); - hasExpectedSize(sourceOperandForScheduledLane, laneCount, "source_operand_for_scheduled_lane"); - if (!multiSourcePayload || multiSourcePayload.getValue() != (transfer.getSources().size() > 1)) - diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { - illegalOp->emitOpError("phase-check scheduled batch deferred communication has inconsistent multi_source_payload metadata"); - }); - if (sourceOperandForScheduledLane) { - for (int64_t sourceOperandIndex : sourceOperandForScheduledLane.asArrayRef()) { - if (sourceOperandIndex < 0 || sourceOperandIndex >= static_cast(transfer.getSources().size())) { - diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { - illegalOp->emitOpError( - "phase-check scheduled batch deferred communication source_operand_for_scheduled_lane index is out of range"); - }); - break; - } - } - } - } + if (!transfer->getParentOfType() && + !transfer->getParentOfType()) + diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError("phase-check deferred communication must be inside a scheduled compute"); + }); }); diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed"); diff --git a/third_party/PIMCOMP-NN b/third_party/PIMCOMP-NN new file mode 160000 index 0000000..0cfbfa5 --- /dev/null +++ b/third_party/PIMCOMP-NN @@ -0,0 +1 @@ +Subproject commit 0cfbfa55cca22fcfebdb8dcfc75420dcbbbbdf77 diff --git a/tools/analyze_spatial_ir_cardinality.py b/tools/analyze_spatial_ir_cardinality.py new file mode 100644 index 0000000..90abdef --- /dev/null +++ b/tools/analyze_spatial_ir_cardinality.py @@ -0,0 +1,253 @@ +#!/usr/bin/env python3.13 + +import argparse +import re +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Iterable + + +OP_PATTERNS = { + "tensor.extract_slice": re.compile(r"\btensor\.extract_slice\b"), + "tensor.insert_slice": re.compile(r"\btensor\.insert_slice\b"), + "spat.channel_send": re.compile(r"\bspat\.channel_send\b"), + "spat.channel_receive": re.compile(r"\bspat\.channel_receive\b"), + "scf.for": re.compile(r"\bscf\.for\b"), + "tensor.empty": re.compile(r"\btensor\.empty\b"), +} + +VALUE_RE = re.compile(r"^\s*(%[\w.$-]+)\s*=\s*(.+)$") +TYPE_RE = re.compile(r":\s*([^:]+?)\s*(?:to|into|$)") +CHANNEL_RE = re.compile(r"channel\s+(%c[-\w.$]+)") +FROM_TO_RE = re.compile(r"from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)") +EXTRACT_SLICE_RE = re.compile( + r"^\s*(%[\w.$-]+)\s*=\s*tensor\.extract_slice\s+(%[\w.$-]+)\[(.*?)\]\s*\[(.*?)\]\s*\[(.*?)\]\s*:\s*(.*?)\s+to\s+(.*)$" +) +INSERT_SLICE_RE = re.compile( + r"^\s*(%[\w.$-]+)\s*=\s*tensor\.insert_slice\s+(%[\w.$-]+)\s+into\s+(%[\w.$-]+)\[(.*?)\]\s*\[(.*?)\]\s*\[(.*?)\]\s*:\s*(.*?)\s+into\s+(.*)$" +) +CHANNEL_RECEIVE_RE = re.compile( + r"^\s*(%[\w.$-]+)\s*=\s*spat\.channel_receive\s+channel\s+(%c[-\w.$]+)\s+from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)\s*:\s*(.*)$" +) +CHANNEL_SEND_RE = re.compile( + r"^\s*spat\.channel_send\s+(%[\w.$-]+)\s+channel\s+(%c[-\w.$]+)\s+from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)\s*:\s*(.*)$" +) +CONST_INDEX_RE = re.compile(r"^\s*(%c[\w.$-]+)\s*=\s*arith\.constant\s+(-?\d+)\s*:\s*index\b") + + +@dataclass +class ChainGroup: + kind: str + signature: str + count: int = 0 + first_line: int = 0 + last_line: int = 0 + fragment_type: str = "" + dest_type: str = "" + varying_dims: set[int] = field(default_factory=set) + rows: list[int | None] = field(default_factory=list) + channels: list[int | None] = field(default_factory=list) + sources: list[int | None] = field(default_factory=list) + targets: list[int | None] = field(default_factory=list) + + def add(self, + line_no: int, + fragment_type: str, + dest_type: str, + offsets: list[str], + channel: int | None = None, + source: int | None = None, + target: int | None = None) -> None: + self.count += 1 + if self.first_line == 0: + self.first_line = line_no + self.last_line = line_no + self.fragment_type = fragment_type or self.fragment_type + self.dest_type = dest_type or self.dest_type + numeric_offsets = [] + for idx, offset in enumerate(offsets): + try: + numeric_offsets.append(int(offset)) + except ValueError: + self.varying_dims.add(idx) + numeric_offsets.append(None) + if self.rows is not None: + row = numeric_offsets[2] if len(numeric_offsets) > 2 else None + self.rows.append(row) + if len(self.rows) >= 2 and self.rows[-1] != self.rows[-2]: + self.varying_dims.add(2) + self.channels.append(channel) + self.sources.append(source) + self.targets.append(target) + + +def parse_const_indices(lines: Iterable[str]) -> dict[str, int]: + constants: dict[str, int] = {} + for line in lines: + match = CONST_INDEX_RE.match(line) + if match: + constants[match.group(1)] = int(match.group(2)) + return constants + + +def split_index_list(value: str) -> list[str]: + return [piece.strip() for piece in value.split(",") if piece.strip()] + + +def decode_const_index(token: str, constants: dict[str, int]) -> int | None: + token = token.strip() + if token in constants: + return constants[token] + try: + return int(token) + except ValueError: + return None + + +def sequence_kind(values: list[int | None]) -> str: + concrete = [value for value in values if value is not None] + if not concrete: + return "dynamic" + if len(concrete) == len(values) and all(b - a == 1 for a, b in zip(concrete, concrete[1:])): + return "consecutive" + if len(set(concrete)) == 1: + return "constant" + return "table" + + +def analyze_file(path: Path) -> tuple[Counter, dict[tuple[str, str], ChainGroup]]: + text = path.read_text() + lines = text.splitlines() + consts = parse_const_indices(lines) + counts = Counter() + groups: dict[tuple[str, str], ChainGroup] = {} + + for line in lines: + for name, pattern in OP_PATTERNS.items(): + if pattern.search(line): + counts[name] += 1 + + value_defs: dict[str, tuple[str, int, re.Match[str] | None]] = {} + for line_no, line in enumerate(lines, start=1): + if match := CHANNEL_RECEIVE_RE.match(line): + value_defs[match.group(1)] = ("receive", line_no, match) + elif match := EXTRACT_SLICE_RE.match(line): + value_defs[match.group(1)] = ("extract", line_no, match) + elif match := INSERT_SLICE_RE.match(line): + source = match.group(2) + producer = value_defs.get(source) + if not producer: + continue + + offsets = split_index_list(match.group(4)) + sizes = split_index_list(match.group(5)) + strides = split_index_list(match.group(6)) + dest = match.group(3) + dest_type = match.group(8).strip() + + if producer[0] == "receive": + recv = producer[2] + assert recv is not None + channel = decode_const_index(recv.group(2), consts) + source_core = decode_const_index(recv.group(3), consts) + target_core = decode_const_index(recv.group(4), consts) + signature = f"recv_insert:{recv.group(5).strip()}->{dest_type}|sizes={','.join(sizes)}|strides={','.join(strides)}" + group = groups.setdefault( + ("receive_to_insert", signature), + ChainGroup("receive_to_insert", signature), + ) + group.add(match.start() and producer[1] or line_no, + recv.group(5).strip(), + dest_type, + offsets, + channel=channel, + source=source_core, + target=target_core) + elif producer[0] == "extract": + extract = producer[2] + assert extract is not None + extract_offsets = split_index_list(extract.group(3)) + signature = ( + f"extract_insert:{extract.group(6).strip()}->{dest_type}|" + f"extract_sizes={extract.group(4).strip()}|insert_sizes={','.join(sizes)}|" + f"src={extract.group(2)}" + ) + group = groups.setdefault( + ("extract_to_insert", signature), + ChainGroup("extract_to_insert", signature), + ) + group.add(producer[1], extract.group(7).strip(), dest_type, offsets) + if extract_offsets and decode_const_index(extract_offsets[0], consts) is None: + group.varying_dims.add(0) + + elif match := CHANNEL_SEND_RE.match(line): + source_value = match.group(1) + producer = value_defs.get(source_value) + if not producer or producer[0] != "extract": + continue + extract = producer[2] + assert extract is not None + signature = ( + f"extract_send:{extract.group(6).strip()}->{match.group(5).strip()}|" + f"extract_sizes={extract.group(4).strip()}|src={extract.group(2)}" + ) + group = groups.setdefault( + ("extract_to_send", signature), + ChainGroup("extract_to_send", signature), + ) + group.add( + producer[1], + extract.group(7).strip(), + match.group(5).strip(), + split_index_list(extract.group(3)), + channel=decode_const_index(match.group(2), consts), + source=decode_const_index(match.group(3), consts), + target=decode_const_index(match.group(4), consts), + ) + + return counts, groups + + +def print_report(path: Path, counts: Counter, groups: dict[tuple[str, str], ChainGroup], limit: int) -> None: + print(f"== {path} ==") + print("counts:") + for name in OP_PATTERNS: + print(f" {name}: {counts[name]}") + + ranked = sorted(groups.values(), key=lambda group: (-group.count, group.first_line)) + print("hot chains:") + for group in ranked[:limit]: + varying = ",".join(str(dim) for dim in sorted(group.varying_dims)) or "none" + print(f" - kind: {group.kind}") + print(f" lines: {group.first_line}-{group.last_line}") + print(f" fragments: {group.count}") + print(f" fragment_type: {group.fragment_type}") + print(f" dest_type: {group.dest_type}") + print(f" varying_dims: {varying}") + if group.rows: + print(f" row_sequence: {sequence_kind(group.rows)}") + if group.channels: + print(f" channel_ids: {sequence_kind(group.channels)}") + if group.sources: + print(f" source_ids: {sequence_kind(group.sources)}") + if group.targets: + print(f" target_ids: {sequence_kind(group.targets)}") + print(f" signature: {group.signature}") + print() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Analyze repeated Spatial/PIM tensor IR cardinality patterns.") + parser.add_argument("paths", nargs="+", help="MLIR files to analyze.") + parser.add_argument("--limit", type=int, default=12, help="Maximum number of hot chains to print per file.") + args = parser.parse_args() + + for path_arg in args.paths: + path = Path(path_arg) + counts, groups = analyze_file(path) + print_report(path, counts, groups, args.limit) + + +if __name__ == "__main__": + main()