From 51fdb830e59c5273a11ba825cfe30abd5070eb30 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Tue, 14 Jul 2026 16:57:58 +0200 Subject: [PATCH] Cose belle --- README.md | 6 +- src/PIM/Common/CMakeLists.txt | 1 + src/PIM/Common/IR/AffineUtils.cpp | 58 +- src/PIM/Common/IR/AffineUtils.hpp | 16 +- src/PIM/Common/IR/ConstantUtils.cpp | 20 +- src/PIM/Common/IR/ConstantUtils.hpp | 7 +- src/PIM/Common/IR/CoreBlockUtils.cpp | 49 +- src/PIM/Common/IR/ShapingUtils.cpp | 39 + src/PIM/Common/IR/ShapingUtils.hpp | 13 + src/PIM/Common/IR/TensorSliceUtils.cpp | 35 + src/PIM/Common/IR/TensorSliceUtils.hpp | 6 + src/PIM/Compiler/PimCodeGen.cpp | 93 +- src/PIM/Compiler/PimCompilerOptions.cpp | 10 +- src/PIM/Compiler/PimCompilerOptions.hpp | 7 +- src/PIM/Compiler/PimMemoryLiveness.cpp | 21 +- .../Conversion/ONNXToSpatial/CompileTime.cpp | 94 +- .../Conversion/ONNXToSpatial/CompileTime.hpp | 5 + .../ONNXToSpatial/LowerSpatialPlansPass.cpp | 38 +- .../Patterns/Math/ReduceMean.cpp | 150 +- .../Patterns/Tensor/Transpose.cpp | 33 +- .../BatchCoreLoweringPatterns.cpp | 5 + src/PIM/Conversion/SpatialToPim/Common.cpp | 66 +- .../SpatialToPim/CoreLoweringPatterns.cpp | 84 +- .../Pim/Transforms/Bufferization/Common.cpp | 65 +- .../Bufferization/ContiguityPatterns.cpp | 249 +- src/PIM/Dialect/Spatial/CMakeLists.txt | 1 + .../Spatial/SpatialOpsCanonicalization.cpp | 3 +- .../DeferredCommunicationDeadlock.cpp | 376 ++- .../DeferredCommunicationPlanning.cpp | 197 +- .../DeferredCommunicationRealization.cpp | 2916 +++++++++++++++-- .../DeferredProjectionAnalysis.cpp | 574 +++- .../DeferredProjectionAnalysis.hpp | 16 + .../MergeComputeNodesPass.cpp | 14 + .../ScheduledComputeMaterialization.cpp | 60 +- .../ScheduledComputePlan.hpp | 5 +- .../ScheduledComputeVerification.cpp | 21 +- .../SpatialDataflowCsvExporter.cpp | 898 +++++ .../SpatialDataflowCsvExporter.hpp | 28 + 38 files changed, 5365 insertions(+), 914 deletions(-) create mode 100644 src/PIM/Common/IR/ShapingUtils.cpp create mode 100644 src/PIM/Common/IR/ShapingUtils.hpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp diff --git a/README.md b/README.md index e6aeae2..169a723 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,9 @@ Pass these to `onnx-mlir` when compiling for PIM: the codegen tail. - `--pim-emit-json` - also emit `core_*.json` instruction files alongside `core_*.pim`. +- `--pim-export-spatial-dataflow=` - control Spatial + dataflow CSV reports. The default `all` emits graph, scheduled, and realized + snapshots under `reports/`. - `--use-experimental-conv-impl` - use the alternate convolution lowering. - `--ignore-concat-error` - soft-fail a ConcatOp corner case. @@ -167,7 +170,8 @@ Each validation run writes artifacts in the model workspace, for example under - `simulation/out.bin` - raw simulator output used for comparison. The compiler currently dumps dialect snapshots such as `spatial0.mlir`, -`spatial1_dcp_merged.mlir`, `pim0.mlir`, `pim1_buff.mlir`, +`spatial1_graph.mlir`, `spatial2_scheduled_no_comm.mlir`, +`spatial3_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`, `pim2_coalesced.mlir`, and `pim3_folded.mlir` when an output directory is available. diff --git a/src/PIM/Common/CMakeLists.txt b/src/PIM/Common/CMakeLists.txt index 9151cc4..2056a13 100644 --- a/src/PIM/Common/CMakeLists.txt +++ b/src/PIM/Common/CMakeLists.txt @@ -8,6 +8,7 @@ add_pim_library(OMPimCommon IR/IndexingUtils.cpp IR/LoopUtils.cpp IR/ShapeUtils.cpp + IR/ShapingUtils.cpp IR/SubviewUtils.cpp IR/TensorSliceUtils.cpp IR/WeightUtils.cpp diff --git a/src/PIM/Common/IR/AffineUtils.cpp b/src/PIM/Common/IR/AffineUtils.cpp index c9dc72d..8f1e570 100644 --- a/src/PIM/Common/IR/AffineUtils.cpp +++ b/src/PIM/Common/IR/AffineUtils.cpp @@ -31,7 +31,7 @@ static FailureOr ceilDivSigned(int64_t lhs, int64_t rhs) { } Value createOrFoldAffineApply( - RewriterBase& rewriter, Location loc, AffineMap map, ValueRange operands, Operation* constantAnchor) { + OpBuilder& builder, Location loc, AffineMap map, ValueRange operands, Operation* constantAnchor) { assert(constantAnchor && "expected a valid constant anchor"); assert(map.getNumResults() == 1 && "affine.apply expects a single-result affine map"); @@ -40,91 +40,91 @@ Value createOrFoldAffineApply( for (Value operand : operands) { std::optional constantValue = matchConstantIndexValue(operand); if (!constantValue) - return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult(); - operandConstants.push_back(rewriter.getIndexAttr(*constantValue)); + return affine::AffineApplyOp::create(builder, loc, map, operands).getResult(); + operandConstants.push_back(builder.getIndexAttr(*constantValue)); } SmallVector foldedResults; if (succeeded(map.constantFold(operandConstants, foldedResults)) && foldedResults.size() == 1) if (auto constantResult = dyn_cast(foldedResults.front())) - return getOrCreateIndexConstant(rewriter, constantAnchor, constantResult.getInt()); + return getOrCreateIndexConstant(builder, constantAnchor, constantResult.getInt()); - return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult(); + return affine::AffineApplyOp::create(builder, loc, map, operands).getResult(); } Value createOrFoldAffineApply( - RewriterBase& rewriter, Location loc, AffineExpr expr, ValueRange dims, Operation* constantAnchor) { + OpBuilder& builder, Location loc, AffineExpr expr, ValueRange dims, Operation* constantAnchor) { AffineMap map = AffineMap::get(/*dimCount=*/dims.size(), /*symbolCount=*/0, expr); - return createOrFoldAffineApply(rewriter, loc, map, dims, constantAnchor); + return createOrFoldAffineApply(builder, loc, map, dims, constantAnchor); } -Value affineMulConst(RewriterBase& rewriter, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) { +Value affineMulConst(OpBuilder& builder, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) { assert(constantAnchor && "expected a valid constant anchor"); if (multiplier == 0) - return getOrCreateIndexConstant(rewriter, constantAnchor, 0); + return getOrCreateIndexConstant(builder, constantAnchor, 0); if (multiplier == 1) return value; - AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext()); - return createOrFoldAffineApply(rewriter, loc, d0 * multiplier, ValueRange {value}, constantAnchor); + AffineExpr d0 = getAffineDimExpr(0, builder.getContext()); + return createOrFoldAffineApply(builder, loc, d0 * multiplier, ValueRange {value}, constantAnchor); } -Value affineAddConst(RewriterBase& rewriter, Location loc, Value value, int64_t offset, Operation* constantAnchor) { +Value affineAddConst(OpBuilder& builder, Location loc, Value value, int64_t offset, Operation* constantAnchor) { assert(constantAnchor && "expected a valid constant anchor"); if (offset == 0) return value; - AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext()); - return createOrFoldAffineApply(rewriter, loc, d0 + offset, ValueRange {value}, constantAnchor); + AffineExpr d0 = getAffineDimExpr(0, builder.getContext()); + return createOrFoldAffineApply(builder, loc, d0 + offset, ValueRange {value}, constantAnchor); } -Value affineModConst(RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) { +Value affineModConst(OpBuilder& builder, Location loc, Value value, int64_t divisor, Operation* constantAnchor) { assert(constantAnchor && "expected a valid constant anchor"); assert(divisor > 0 && "expected a positive affine.mod divisor"); if (divisor == 1) - return getOrCreateIndexConstant(rewriter, constantAnchor, 0); + return getOrCreateIndexConstant(builder, constantAnchor, 0); - AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext()); - return createOrFoldAffineApply(rewriter, loc, d0 % divisor, ValueRange {value}, constantAnchor); + AffineExpr d0 = getAffineDimExpr(0, builder.getContext()); + return createOrFoldAffineApply(builder, loc, d0 % divisor, ValueRange {value}, constantAnchor); } Value affineFloorDivConst( - RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) { + OpBuilder& builder, Location loc, Value value, int64_t divisor, Operation* constantAnchor) { assert(constantAnchor && "expected a valid constant anchor"); assert(divisor > 0 && "expected a positive affine.floor_div divisor"); if (divisor == 1) return value; - AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext()); - return createOrFoldAffineApply(rewriter, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor); + AffineExpr d0 = getAffineDimExpr(0, builder.getContext()); + return createOrFoldAffineApply(builder, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor); } Value affineAddModConst( - RewriterBase& rewriter, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) { + OpBuilder& builder, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) { assert(constantAnchor && "expected a valid constant anchor"); assert(divisor > 0 && "expected a positive affine.mod divisor"); if (divisor == 1) - return getOrCreateIndexConstant(rewriter, constantAnchor, 0); + return getOrCreateIndexConstant(builder, constantAnchor, 0); - AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext()); + AffineExpr d0 = getAffineDimExpr(0, builder.getContext()); AffineExpr expr = d0; if (offset != 0) expr = expr + offset; - return createOrFoldAffineApply(rewriter, loc, expr % divisor, ValueRange {value}, constantAnchor); + return createOrFoldAffineApply(builder, loc, expr % divisor, ValueRange {value}, constantAnchor); } Value affineAddFloorDivConst( - RewriterBase& rewriter, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) { + OpBuilder& builder, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) { assert(constantAnchor && "expected a valid constant anchor"); assert(divisor > 0 && "expected a positive affine.floor_div divisor"); if (divisor == 1) - return offset == 0 ? value : affineAddConst(rewriter, loc, value, offset, constantAnchor); + return offset == 0 ? value : affineAddConst(builder, loc, value, offset, constantAnchor); - AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext()); + AffineExpr d0 = getAffineDimExpr(0, builder.getContext()); AffineExpr expr = d0; if (offset != 0) expr = expr + offset; - return createOrFoldAffineApply(rewriter, loc, expr.floorDiv(divisor), ValueRange {value}, constantAnchor); + return createOrFoldAffineApply(builder, loc, expr.floorDiv(divisor), ValueRange {value}, constantAnchor); } FailureOr evaluateAffineExpr(AffineExpr expr, ArrayRef dims, ArrayRef symbols) { diff --git a/src/PIM/Common/IR/AffineUtils.hpp b/src/PIM/Common/IR/AffineUtils.hpp index d9d3f58..b96f000 100644 --- a/src/PIM/Common/IR/AffineUtils.hpp +++ b/src/PIM/Common/IR/AffineUtils.hpp @@ -11,50 +11,50 @@ namespace onnx_mlir { using IndexValueResolver = llvm::function_ref(mlir::Value)>; -mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter, +mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder, mlir::Location loc, mlir::AffineMap map, mlir::ValueRange operands, mlir::Operation* constantAnchor); -mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter, +mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder, mlir::Location loc, mlir::AffineExpr expr, mlir::ValueRange dims, mlir::Operation* constantAnchor); -mlir::Value affineMulConst(mlir::RewriterBase& rewriter, +mlir::Value affineMulConst(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, int64_t multiplier, mlir::Operation* constantAnchor); -mlir::Value affineAddConst(mlir::RewriterBase& rewriter, +mlir::Value affineAddConst(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, int64_t offset, mlir::Operation* constantAnchor); -mlir::Value affineModConst(mlir::RewriterBase& rewriter, +mlir::Value affineModConst(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, int64_t divisor, mlir::Operation* constantAnchor); -mlir::Value affineFloorDivConst(mlir::RewriterBase& rewriter, +mlir::Value affineFloorDivConst(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, int64_t divisor, mlir::Operation* constantAnchor); -mlir::Value affineAddModConst(mlir::RewriterBase& rewriter, +mlir::Value affineAddModConst(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, int64_t offset, int64_t divisor, mlir::Operation* constantAnchor); -mlir::Value affineAddFloorDivConst(mlir::RewriterBase& rewriter, +mlir::Value affineAddFloorDivConst(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, int64_t offset, diff --git a/src/PIM/Common/IR/ConstantUtils.cpp b/src/PIM/Common/IR/ConstantUtils.cpp index 47fd66a..8f7c5b7 100644 --- a/src/PIM/Common/IR/ConstantUtils.cpp +++ b/src/PIM/Common/IR/ConstantUtils.cpp @@ -49,7 +49,7 @@ Value getOrCreateConstant(OperationFolder& folder, Operation* anchorOp, Attribut return folder.getOrCreateConstant(hostBlock, arithDialect, value, type); } -Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute value, Type type) { +Value getOrCreateConstant(OpBuilder& builder, Operation* anchorOp, Attribute value, Type type) { assert(anchorOp && "expected a valid anchor operation"); Block* hostBlock = getConstantInsertionBlock(anchorOp); for (Operation& op : *hostBlock) { @@ -59,9 +59,16 @@ Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute return constantOp.getResult(); } - OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPointToStart(hostBlock); - return arith::ConstantOp::create(rewriter, anchorOp->getLoc(), type, cast(value)).getResult(); + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(hostBlock); + return arith::ConstantOp::create(builder, anchorOp->getLoc(), type, cast(value)).getResult(); +} + +Value createConstantAtHostBlockStart(OpBuilder& builder, Operation* anchorOp, TypedAttr value) { + assert(anchorOp && "expected a valid anchor operation"); + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(getConstantInsertionBlock(anchorOp)); + return arith::ConstantOp::create(builder, anchorOp->getLoc(), value).getResult(); } Value getOrCreateConstantLike(OperationFolder& folder, arith::ConstantOp constantOp) { @@ -73,9 +80,8 @@ Value getOrCreateIndexConstant(OperationFolder& folder, Operation* anchorOp, int return getOrCreateConstant(folder, anchorOp, builder.getIndexAttr(value), builder.getIndexType()); } -Value getOrCreateIndexConstant(RewriterBase& rewriter, Operation* anchorOp, int64_t value) { - Builder builder(anchorOp->getContext()); - return getOrCreateConstant(rewriter, anchorOp, builder.getIndexAttr(value), builder.getIndexType()); +Value getOrCreateIndexConstant(OpBuilder& builder, Operation* anchorOp, int64_t value) { + return getOrCreateConstant(builder, anchorOp, builder.getIndexAttr(value), builder.getIndexType()); } void hoistAndUniquifyIndexConstants(func::FuncOp funcOp, RewriterBase& rewriter) { diff --git a/src/PIM/Common/IR/ConstantUtils.hpp b/src/PIM/Common/IR/ConstantUtils.hpp index fc8be72..595738d 100644 --- a/src/PIM/Common/IR/ConstantUtils.hpp +++ b/src/PIM/Common/IR/ConstantUtils.hpp @@ -16,13 +16,16 @@ mlir::Value getOrCreateConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type); mlir::Value -getOrCreateConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type); +getOrCreateConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type); + +mlir::Value +createConstantAtHostBlockStart(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::TypedAttr value); mlir::Value getOrCreateConstantLike(mlir::OperationFolder& folder, mlir::arith::ConstantOp constantOp); mlir::Value getOrCreateIndexConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, int64_t value); -mlir::Value getOrCreateIndexConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, int64_t value); +mlir::Value getOrCreateIndexConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, int64_t value); void hoistAndUniquifyIndexConstants(mlir::func::FuncOp funcOp, mlir::RewriterBase& rewriter); diff --git a/src/PIM/Common/IR/CoreBlockUtils.cpp b/src/PIM/Common/IR/CoreBlockUtils.cpp index 42a6699..68547d8 100644 --- a/src/PIM/Common/IR/CoreBlockUtils.cpp +++ b/src/PIM/Common/IR/CoreBlockUtils.cpp @@ -36,9 +36,10 @@ bool isCoreStaticAddressOp(mlir::Operation* op) { mlir::LogicalResult walkPimCoreBlock(mlir::Block& block, - const StaticValueKnowledge& knowledge, + const StaticValueKnowledge& initialKnowledge, llvm::function_ref callback) { bool hasFailure = false; + StaticValueKnowledge knowledge = initialKnowledge; for (mlir::Operation& op : block) { if (mlir::isa(op) || isCoreStaticAddressOp(&op)) continue; @@ -89,6 +90,27 @@ walkPimCoreBlock(mlir::Block& block, continue; } + if (auto switchOp = mlir::dyn_cast(op)) { + auto selector = resolveIndexValue(switchOp.getArg(), knowledge); + if (failed(selector)) { + switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen"); + hasFailure = true; + continue; + } + mlir::Region* selected = &switchOp.getDefaultRegion(); + for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions())) + if (caseValue == *selector) { + selected = &caseRegion; + break; + } + if (failed(walkPimCoreBlock(selected->front(), knowledge, callback))) + hasFailure = true; + auto yield = mlir::cast(selected->front().getTerminator()); + for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands())) + knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge); + continue; + } + if (failed(callback(op, knowledge))) hasFailure = true; } @@ -97,9 +119,10 @@ walkPimCoreBlock(mlir::Block& block, mlir::LogicalResult walkPimCoreBlockStructurally( mlir::Block& block, - const StaticValueKnowledge& knowledge, + const StaticValueKnowledge& initialKnowledge, llvm::function_ref callback) { bool hasFailure = false; + StaticValueKnowledge knowledge = initialKnowledge; for (mlir::Operation& op : block) { if (mlir::isa(op) || isCoreStaticAddressOp(&op)) continue; @@ -159,6 +182,28 @@ mlir::LogicalResult walkPimCoreBlockStructurally( continue; } + if (auto switchOp = mlir::dyn_cast(op)) { + auto selector = resolveIndexValue(switchOp.getArg(), knowledge); + if (failed(selector)) { + switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM verification"); + hasFailure = true; + continue; + } + mlir::Region* selected = &switchOp.getDefaultRegion(); + for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions())) + if (caseValue == *selector) { + selected = &caseRegion; + break; + } + for (mlir::Region& region : switchOp->getRegions()) + if (failed(walkPimCoreBlockStructurally(region.front(), knowledge, callback))) + hasFailure = true; + auto yield = mlir::cast(selected->front().getTerminator()); + for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands())) + knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge); + continue; + } + if (failed(callback(op, knowledge))) hasFailure = true; } diff --git a/src/PIM/Common/IR/ShapingUtils.cpp b/src/PIM/Common/IR/ShapingUtils.cpp new file mode 100644 index 0000000..f91b3f4 --- /dev/null +++ b/src/PIM/Common/IR/ShapingUtils.cpp @@ -0,0 +1,39 @@ +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" + +#include "ShapingUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Dialect/ONNX/ONNXOps.hpp" + +using namespace mlir; + +namespace onnx_mlir { + +bool isShapingOnlyOp(Operation *op) { + return isa(op); +} + +bool isPureIndexComputationOp(Operation *op) { + if (op->getNumRegions() != 0 || op->getNumResults() == 0 || op->hasTrait() + || !isMemoryEffectFree(op)) + return false; + auto isIndexOrInteger = [](Type type) { return type.isIndex() || isa(type); }; + return llvm::all_of(op->getOperandTypes(), isIndexOrInteger) + && llvm::all_of(op->getResultTypes(), isIndexOrInteger); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Common/IR/ShapingUtils.hpp b/src/PIM/Common/IR/ShapingUtils.hpp new file mode 100644 index 0000000..75b5d9b --- /dev/null +++ b/src/PIM/Common/IR/ShapingUtils.hpp @@ -0,0 +1,13 @@ +#pragma once + +namespace mlir { +class Operation; +} + +namespace onnx_mlir { + +bool isShapingOnlyOp(mlir::Operation *op); + +bool isPureIndexComputationOp(mlir::Operation *op); + +} // namespace onnx_mlir diff --git a/src/PIM/Common/IR/TensorSliceUtils.cpp b/src/PIM/Common/IR/TensorSliceUtils.cpp index 447036e..dfdbc99 100644 --- a/src/PIM/Common/IR/TensorSliceUtils.cpp +++ b/src/PIM/Common/IR/TensorSliceUtils.cpp @@ -68,4 +68,39 @@ Value insertStaticSlice( .getResult(); } +FailureOr addLeadingUnitTensorDimension(OpBuilder& builder, Location loc, Value value) { + auto type = dyn_cast(value.getType()); + if (!type || !type.hasStaticShape()) + return failure(); + SmallVector shape {1}; + llvm::append_range(shape, type.getShape()); + auto resultType = RankedTensorType::get(shape, type.getElementType(), type.getEncoding()); + SmallVector reassociation; + if (type.getRank() != 0) { + reassociation.push_back({0, 1}); + for (int64_t dim = 1; dim < type.getRank(); ++dim) + reassociation.push_back({dim + 1}); + } + return tensor::ExpandShapeOp::create(builder, loc, resultType, value, reassociation).getResult(); +} + +FailureOr removeLeadingUnitTensorDimension( + OpBuilder& builder, Location loc, Value value, RankedTensorType resultType) { + if (value.getType() == resultType) + return value; + auto type = dyn_cast(value.getType()); + if (!type || !resultType || !type.hasStaticShape() || !resultType.hasStaticShape() + || type.getRank() != resultType.getRank() + 1 || type.getDimSize(0) != 1 + || type.getElementType() != resultType.getElementType() + || !llvm::equal(type.getShape().drop_front(), resultType.getShape())) + return failure(); + SmallVector reassociation; + if (resultType.getRank() != 0) { + reassociation.push_back({0, 1}); + for (int64_t dim = 1; dim < resultType.getRank(); ++dim) + reassociation.push_back({dim + 1}); + } + return tensor::CollapseShapeOp::create(builder, loc, resultType, value, reassociation).getResult(); +} + } // namespace onnx_mlir diff --git a/src/PIM/Common/IR/TensorSliceUtils.hpp b/src/PIM/Common/IR/TensorSliceUtils.hpp index 3bc31d1..92b5dd0 100644 --- a/src/PIM/Common/IR/TensorSliceUtils.hpp +++ b/src/PIM/Common/IR/TensorSliceUtils.hpp @@ -25,4 +25,10 @@ mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter, mlir::Value dest, llvm::ArrayRef offsets); +mlir::FailureOr +addLeadingUnitTensorDimension(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value); + +mlir::FailureOr removeLeadingUnitTensorDimension( + mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, mlir::RankedTensorType resultType); + } // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index b949eb7..012e8d9 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -1119,7 +1119,8 @@ struct CompiledCoreNode { enum class Kind : uint8_t { Op, Loop, - If + If, + IndexSwitch }; Kind kind = Kind::Op; @@ -1132,6 +1133,9 @@ struct CompiledCoreNode { std::unique_ptr> loopBody; std::unique_ptr> thenBody; std::unique_ptr> elseBody; + llvm::SmallVector caseValues; + llvm::SmallVector>> caseBodies; + std::unique_ptr> defaultBody; }; static FailureOr classifyCompiledCoreOpKind(Operation& op) { @@ -1231,6 +1235,31 @@ compileCoreEmissionPlan(Block& block, Operation* weightOwner, llvm::SmallVectorI continue; } + if (auto switchOp = dyn_cast(op)) { + auto selector = compileIndexExpr(switchOp.getArg()); + if (failed(selector)) { + switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen"); + return failure(); + } + CompiledCoreNode switchNode; + switchNode.kind = CompiledCoreNode::Kind::IndexSwitch; + switchNode.op = switchOp.getOperation(); + switchNode.condition = *selector; + llvm::append_range(switchNode.caseValues, switchOp.getCases()); + for (mlir::Region& region : switchOp.getCaseRegions()) { + auto body = std::make_unique>(); + if (failed(compileCoreEmissionPlan(region.front(), weightOwner, *body))) + return failure(); + switchNode.caseBodies.push_back(std::move(body)); + } + switchNode.defaultBody = std::make_unique>(); + if (failed(compileCoreEmissionPlan( + switchOp.getDefaultRegion().front(), weightOwner, *switchNode.defaultBody))) + return failure(); + plan.push_back(std::move(switchNode)); + continue; + } + auto opKind = classifyCompiledCoreOpKind(op); if (failed(opKind)) { InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'"; @@ -1313,6 +1342,31 @@ static LogicalResult executeCompiledCorePlan( continue; } + if (node.kind == CompiledCoreNode::Kind::IndexSwitch) { + auto selector = node.condition.evaluate(knowledge); + auto switchOp = cast(node.op); + if (failed(selector)) { + switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen"); + return failure(); + } + const llvm::SmallVectorImpl* selectedBody = node.defaultBody.get(); + mlir::Region* selectedRegion = &switchOp.getDefaultRegion(); + for (auto [index, caseValue] : llvm::enumerate(node.caseValues)) + if (caseValue == *selector) { + selectedBody = node.caseBodies[index].get(); + selectedRegion = &switchOp.getCaseRegions()[index]; + break; + } + if (failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge, + resolveWeightSlot, processedOperations, + batchLane, batchLaneCount))) + return failure(); + auto yield = cast(selectedRegion->front().getTerminator()); + for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands())) + knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge); + continue; + } + switch (node.opKind) { case CompiledCoreOpKind::Load: coreCodeGen.codeGenLoadOp(cast(node.op), knowledge); @@ -1413,6 +1467,36 @@ static int64_t codeGenCoreOps( return failed(result) ? -1 : static_cast(processedOperations); } +static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) { + std::string outputCorePath = + (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str(); + std::error_code errorCode; + raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None); + if (errorCode) { + errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n'; + return InvalidOutputFileAccess; + } + + pim_binary::writeHeader(coreBinaryStream); + pim_binary::patchInstructionCount(coreBinaryStream, 0); + coreBinaryStream.close(); + + if (!pimEmitJson.getValue()) + return CompilerSuccess; + + std::string outputCoreJsonPath = + (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str(); + errorCode = std::error_code(); + raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode); + if (errorCode) { + errs() << "Error while opening core json file `" << outputCoreJsonPath << "`: " << errorCode.message() << '\n'; + return InvalidOutputFileAccess; + } + coreJsonStream << "[]"; + coreJsonStream.close(); + return CompilerSuccess; +} + OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::string& outputDirPath) { if (!outputDirPath.empty()) { if (auto error = sys::fs::create_directory(outputDirPath)) { @@ -1657,6 +1741,13 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std:: if (jobResults[jobIndex].status != CompilerSuccess) return jobResults[jobIndex].status; + if (jobs.empty()) { + if (auto err = emitEmptyCoreArtifacts(outputDirPath, 0)) + return err; + xbarsPerArrayGroup["core0"] = json::Array {}; + memory.recordCoreReport(0, MemoryReportRow {}); + } + llvm::SmallVector weightRequests; weightRequests.reserve(jobs.size()); for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) { diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index 5e73d58..5153481 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -59,13 +59,15 @@ llvm::cl::opt pimConvLowering( llvm::cl::opt pimExportSpatialDataflow( "pim-export-spatial-dataflow", - llvm::cl::desc("Emit Gephi-importable CSV dataflow reports around MergeComputeNodes materialization"), + llvm::cl::desc("Emit Gephi-importable CSV dataflow reports for Spatial pipeline snapshots"), llvm::cl::values(clEnumValN(SpatialDataflowExportNone, "none", "Do not emit Spatial dataflow CSV reports")), - llvm::cl::values(clEnumValN(SpatialDataflowExportPre, "pre", "Emit pre-materialization Spatial dataflow CSV reports")), llvm::cl::values( - clEnumValN(SpatialDataflowExportPost, "post", "Emit post-materialization Spatial dataflow CSV reports")), + clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit spatial1 graph dataflow CSV reports")), llvm::cl::values( - clEnumValN(SpatialDataflowExportBoth, "both", "Emit both pre- and post-materialization Spatial dataflow CSV reports")), + clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 scheduled dataflow CSV reports")), + llvm::cl::values( + clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 realized dataflow CSV reports")), + llvm::cl::values(clEnumValN(SpatialDataflowExportAll, "all", "Emit all Spatial dataflow CSV reports")), llvm::cl::init(SpatialDataflowExportNone), llvm::cl::cat(OnnxMlirOptions)); diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index b5d931d..51525a5 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -44,9 +44,10 @@ typedef enum { typedef enum { SpatialDataflowExportNone = 0, - SpatialDataflowExportPre = 1, - SpatialDataflowExportPost = 2, - SpatialDataflowExportBoth = 3, + SpatialDataflowExportSpatial1 = 1, + SpatialDataflowExportSpatial2 = 2, + SpatialDataflowExportSpatial3 = 3, + SpatialDataflowExportAll = 4, } PimSpatialDataflowExportType; extern llvm::cl::OptionCategory OnnxMlirOptions; diff --git a/src/PIM/Compiler/PimMemoryLiveness.cpp b/src/PIM/Compiler/PimMemoryLiveness.cpp index decf990..9e91ac9 100644 --- a/src/PIM/Compiler/PimMemoryLiveness.cpp +++ b/src/PIM/Compiler/PimMemoryLiveness.cpp @@ -291,7 +291,26 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord if (auto yieldOp = dyn_cast(user)) { auto forOp = dyn_cast(yieldOp->getParentOp()); - if (!forOp) { + auto ifOp = dyn_cast(yieldOp->getParentOp()); + auto indexSwitch = dyn_cast(yieldOp->getParentOp()); + if (ifOp) { + for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) { + if (operand != value) + continue; + pendingValues.push_back(ifOp.getResult(index)); + appendAliasDescription(interval.aliasesFollowed, ifOp.getResult(index)); + } + } + else if (indexSwitch) { + for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) { + if (operand != value) + continue; + pendingValues.push_back(indexSwitch.getResult(index)); + appendAliasDescription(interval.aliasesFollowed, + indexSwitch.getResult(index)); + } + } + else if (!forOp) { addFallbackReason(interval.fallbackReason, "yield without scf.for parent"); } else { diff --git a/src/PIM/Conversion/ONNXToSpatial/CompileTime.cpp b/src/PIM/Conversion/ONNXToSpatial/CompileTime.cpp index f1e9771..5fa76d7 100644 --- a/src/PIM/Conversion/ONNXToSpatial/CompileTime.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/CompileTime.cpp @@ -9,10 +9,12 @@ #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" +#include #include #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -21,24 +23,7 @@ using namespace mlir; namespace onnx_mlir { -namespace { - -static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) { - return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; }); -} - -static bool hasConstantIndices(tensor::ExtractOp extractOp) { - return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); }); -} - -static bool isStaticTensorResult(Operation* op) { - return llvm::all_of(op->getResultTypes(), [](Type type) { - auto shapedType = dyn_cast(type); - return shapedType && shapedType.hasStaticShape(); - }); -} - -static FailureOr transposeDenseElements(DenseElementsAttr denseAttr, ArrayRef perms) { +FailureOr transposeDenseElementsAttr(DenseElementsAttr denseAttr, ArrayRef perms) { auto tensorType = dyn_cast(denseAttr.getType()); if (!tensorType) return failure(); @@ -59,7 +44,45 @@ static FailureOr transposeDenseElements(DenseElementsAttr den auto transposedType = RankedTensorType::get(transposedShape, tensorType.getElementType(), tensorType.getEncoding()); if (denseAttr.isSplat()) - return DenseElementsAttr::get(transposedType, denseAttr.getSplatValue()); + return DenseElementsAttr::getFromRawBuffer(transposedType, denseAttr.getRawData()); + + const unsigned elementBitWidth = tensorType.getElementTypeBitWidth(); + const ArrayRef inputData = denseAttr.getRawData(); + if (elementBitWidth % 8 == 0) { + const size_t elementBytes = elementBitWidth / 8; + const size_t expectedBytes = denseAttr.getNumElements() * elementBytes; + if (inputData.size() == expectedBytes) { + SmallVector transposedData(expectedBytes); + if (rank == 2 && perms[0] == 1 && perms[1] == 0) { + const int64_t rows = tensorType.getDimSize(0); + const int64_t columns = tensorType.getDimSize(1); + for (int64_t row = 0; row < rows; ++row) + for (int64_t column = 0; column < columns; ++column) + std::memcpy(transposedData.data() + (column * rows + row) * elementBytes, + inputData.data() + (row * columns + column) * elementBytes, + elementBytes); + return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData); + } + + SmallVector originalStrides = computeRowMajorStrides(tensorType.getShape()); + SmallVector transposedStrides = computeRowMajorStrides(transposedShape); + SmallVector originalIndices(rank); + for (int64_t linearIndex = 0; linearIndex < tensorType.getNumElements(); ++linearIndex) { + int64_t remaining = linearIndex; + for (int64_t dim = 0; dim < rank; ++dim) { + originalIndices[dim] = originalStrides.empty() ? 0 : remaining / originalStrides[dim]; + remaining = originalStrides.empty() ? 0 : remaining % originalStrides[dim]; + } + int64_t transposedLinearIndex = 0; + for (int64_t dim = 0; dim < rank; ++dim) + transposedLinearIndex += originalIndices[perms[dim]] * transposedStrides[dim]; + std::memcpy(transposedData.data() + transposedLinearIndex * elementBytes, + inputData.data() + linearIndex * elementBytes, + elementBytes); + } + return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData); + } + } SmallVector originalValues(denseAttr.getValues()); SmallVector transposedValues(originalValues.size()); @@ -84,16 +107,30 @@ static FailureOr transposeDenseElements(DenseElementsAttr den return DenseElementsAttr::get(transposedType, transposedValues); } +namespace { + +static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) { + return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; }); +} + +static bool hasConstantIndices(tensor::ExtractOp extractOp) { + return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); }); +} + +static bool isStaticTensorResult(Operation* op) { + return llvm::all_of(op->getResultTypes(), [](Type type) { + auto shapedType = dyn_cast(type); + return shapedType && shapedType.hasStaticShape(); + }); +} + static FailureOr reshapeDenseElements(DenseElementsAttr denseAttr, RankedTensorType resultType) { auto sourceType = dyn_cast(denseAttr.getType()); - if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements()) + if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements() + || sourceType.getElementType() != resultType.getElementType()) return failure(); - if (denseAttr.isSplat()) - return DenseElementsAttr::get(resultType, denseAttr.getSplatValue()); - - SmallVector values(denseAttr.getValues()); - return DenseElementsAttr::get(resultType, values); + return DenseElementsAttr::getFromRawBuffer(resultType, denseAttr.getRawData()); } static FailureOr extractSliceDenseElements(DenseElementsAttr denseAttr, @@ -161,7 +198,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm: perm.reserve(transposeOp.getPermAttr().size()); for (IntegerAttr attr : transposeOp.getPermAttr().getAsRange()) perm.push_back(attr.getInt()); - auto transposedAttr = transposeDenseElements(inputAttr, perm); + auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm); return succeeded(transposedAttr) ? *transposedAttr : nullptr; } @@ -171,7 +208,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm: return nullptr; SmallVector perm(transposeOp.getPermutation().begin(), transposeOp.getPermutation().end()); - auto transposedAttr = transposeDenseElements(inputAttr, perm); + auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm); return succeeded(transposedAttr) ? *transposedAttr : nullptr; } @@ -219,6 +256,9 @@ getCompileTimeSourceImpl(Operation* op, llvm::SmallPtrSetImpl& visit chainLength += 1; + if (!isShapingOnlyOp(op)) + return std::nullopt; + if (auto extractOp = dyn_cast(op)) return hasConstantIndices(extractOp) ? getCompileTimeSourceImpl(extractOp.getTensor().getDefiningOp(), visited, chainLength) diff --git a/src/PIM/Conversion/ONNXToSpatial/CompileTime.hpp b/src/PIM/Conversion/ONNXToSpatial/CompileTime.hpp index c9ad809..b5e6795 100644 --- a/src/PIM/Conversion/ONNXToSpatial/CompileTime.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/CompileTime.hpp @@ -4,6 +4,8 @@ #include "mlir/IR/Operation.h" #include "mlir/IR/Value.h" +#include "llvm/ADT/ArrayRef.h" + namespace onnx_mlir { struct CompileTimeSource { @@ -19,4 +21,7 @@ bool isCompileTimeOp(mlir::Operation* op); mlir::DenseElementsAttr getHostConstDenseElementsAttr(mlir::Value value); +mlir::FailureOr transposeDenseElementsAttr( + mlir::DenseElementsAttr denseAttr, llvm::ArrayRef permutation); + } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp index 6125212..e60c2ef 100644 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -20,6 +20,7 @@ #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp" #include "src/Accelerators/PIM/Pass/PIMPasses.h" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -330,17 +331,25 @@ struct LowerSpatialPlansPass final : PassWrapper topLevelHelperOps; + funcOp.walk([&](Operation* op) { + if (isa(op)) + return WalkResult::skip(); + if (isa(op)) + topLevelHelperOps.push_back(op); + return WalkResult::advance(); + }); + for (Operation *helper : topLevelHelperOps) { + if (failed(applyPartialConversion( + helper, helperTarget, frozenHelperPatterns))) { + moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering"); + signalPassFailure(); + return; + } } - FrozenRewritePatternSet nestedHelperPatterns([&] { - RewritePatternSet patterns(ctx); - populateGemmPatterns(patterns, ctx); - populateTransposePatterns(patterns, ctx); - return patterns; - }()); ConversionTarget nestedHelperTarget(*ctx); nestedHelperTarget.addLegalDialectemitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering"); signalPassFailure(); return; @@ -392,6 +402,12 @@ struct LowerSpatialPlansPass final : PassWrapper -#include #include #include @@ -122,14 +121,6 @@ static RankedTensorType getKeepdimsType(RankedTensorType inputType, Type element return RankedTensorType::get(shape, elementType, inputType.getEncoding()); } -static RankedTensorType getCompactKeptType(RankedTensorType inputType, Type elementType, ArrayRef reducedAxes) { - SmallVector shape; - for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes)) - if (!isReduced) - shape.push_back(dim); - return RankedTensorType::get(shape, elementType, inputType.getEncoding()); -} - static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef reducedAxes) { SmallVector shape; shape.reserve(inputType.getRank()); @@ -228,59 +219,80 @@ static FailureOr buildReduceMeanKeepdimsBatch(Value input, return (*batchOp).getResult(0); } -static Value buildKeepdimsFromLanePackedBatch(Value batchValue, - RankedTensorType keepdimsType, - RankedTensorType compactKeptType, - ArrayRef reducedAxes, - ConversionPatternRewriter& rewriter, - Location loc) { - auto batchType = cast(batchValue.getType()); - if (batchType == keepdimsType) - return batchValue; +static FailureOr buildReduceMeanKeepdimsBlueprint( + Value batchValue, RankedTensorType keepdimsType, + ArrayRef reducedAxes, ConversionPatternRewriter& rewriter, + Location loc) { + auto batchType = dyn_cast(batchValue.getType()); + int64_t rank = keepdimsType.getRank(); + if (!batchType || !batchType.hasStaticShape() + || !keepdimsType.hasStaticShape() + || static_cast(reducedAxes.size()) != rank + || batchType.getRank() != rank + 1 + || batchType.getElementType() != keepdimsType.getElementType()) + return failure(); - SmallVector collapseToFlat {{}}; - for (int64_t axis = 0; axis < batchType.getRank(); ++axis) - collapseToFlat.front().push_back(axis); - - SmallVector expandFlatToCompact(1); - for (int64_t axis = 0; axis < compactKeptType.getRank(); ++axis) - expandFlatToCompact.front().push_back(axis); - - SmallVector expandCompactToKeepdims; - ReassociationIndices pendingLeadingReducedAxes; + int64_t laneCount = 1; + SmallVector keptAxes; + SmallVector keptAxisStrides; for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) { - if (isReduced) { - if (expandCompactToKeepdims.empty()) - pendingLeadingReducedAxes.push_back(axis); - else - expandCompactToKeepdims.back().push_back(axis); - continue; - } - - expandCompactToKeepdims.emplace_back(); - auto& group = expandCompactToKeepdims.back(); - group.append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end()); - pendingLeadingReducedAxes.clear(); - group.push_back(axis); + int64_t dim = keepdimsType.getDimSize(axis); + if (dim <= 0 || (isReduced && dim != 1)) + return failure(); + if (!isReduced) + keptAxes.push_back(axis); } - if (!pendingLeadingReducedAxes.empty()) - expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end()); + keptAxisStrides.resize(keptAxes.size(), 1); + for (int64_t index = static_cast(keptAxes.size()) - 1; + index >= 0; --index) { + keptAxisStrides[index] = laneCount; + int64_t dim = keepdimsType.getDimSize(keptAxes[index]); + if (laneCount > std::numeric_limits::max() / dim) + return failure(); + laneCount *= dim; + } + if (batchType.getDimSize(0) != laneCount + || llvm::any_of(batchType.getShape().drop_front(), + [](int64_t dim) { return dim != 1; })) + return failure(); - if (batchType.getNumElements() != batchType.getDimSize(0)) - return {}; - auto reshapeCompute = - createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) { - 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) - compact = tensor::ExpandShapeOp::create(rewriter, loc, compactKeptType, flat, expandFlatToCompact); - Value keepdims = compact; - if (keepdimsType != compactKeptType) - keepdims = tensor::ExpandShapeOp::create(rewriter, loc, keepdimsType, compact, expandCompactToKeepdims); - spatial::SpatYieldOp::create(rewriter, loc, keepdims); - }); - return reshapeCompute.getResult(0); + SmallVector operandIndices(laneCount, 0); + SmallVector sourceSlots; + SmallVector sourceOffsets(laneCount, 0); + SmallVector fragmentOffsets; + sourceSlots.reserve(laneCount); + fragmentOffsets.reserve(laneCount * rank); + for (int64_t lane = 0; lane < laneCount; ++lane) { + sourceSlots.push_back(lane); + size_t keptAxisIndex = 0; + for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) { + if (isReduced) { + fragmentOffsets.push_back(0); + continue; + } + int64_t dim = keepdimsType.getDimSize(axis); + fragmentOffsets.push_back( + (lane / keptAxisStrides[keptAxisIndex]) % dim); + ++keptAxisIndex; + } + } + SmallVector fragmentSizes(fragmentOffsets.size(), 1); + SmallVector fragmentStrides(fragmentOffsets.size(), 1); + return spatial::SpatBlueprintOp::create( + rewriter, loc, keepdimsType, batchValue, ValueRange {}, + rewriter.getStringAttr("nchw"), + rewriter.getStringAttr("fragmented"), + rewriter.getDenseI64ArrayAttr(fragmentOffsets), + rewriter.getDenseI64ArrayAttr(fragmentSizes), + rewriter.getStringAttr("reduce_mean_keepdims_fragments"), + rewriter.getStringAttr("fragment_assembly"), + rewriter.getDenseI64ArrayAttr(operandIndices), + rewriter.getDenseI64ArrayAttr(sourceSlots), + rewriter.getDenseI64ArrayAttr(sourceOffsets), + rewriter.getDenseI64ArrayAttr(fragmentStrides), + rewriter.getStringAttr("disjoint"), + rewriter.getStringAttr("complete")) + .getOutput(); } static SmallVector buildCollapseReassociation(ArrayRef reducedAxes) { @@ -357,26 +369,36 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern { Location loc = reduceMeanOp.getLoc(); RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType()); - RankedTensorType compactKeptType = getCompactKeptType(inputType, resultType.getElementType(), reducedAxes); RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes); int64_t laneCount = 1; - for (int64_t dim : compactKeptType.getShape()) + for (auto [dim, isReduced] : llvm::zip_equal(keepdimsType.getShape(), reducedAxes)) { + if (isReduced) + continue; + if (dim <= 0 || laneCount > std::numeric_limits::max() / dim) + return rewriter.notifyMatchFailure( + reduceMeanOp, "ReduceMean physical lane count is not representable"); laneCount *= dim; + } RankedTensorType batchType = getLanePackedKeepdimsType(laneCount, leafType); auto lanePackedKeepdims = buildReduceMeanKeepdimsBatch(adaptor.getData(), reducedAxes, batchType, leafType, rewriter, loc); if (failed(lanePackedKeepdims)) return failure(); - Value reducedKeepdims = - buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc); + auto reducedKeepdims = buildReduceMeanKeepdimsBlueprint( + *lanePackedKeepdims, keepdimsType, reducedAxes, rewriter, loc); + if (failed(reducedKeepdims)) + return rewriter.notifyMatchFailure( + reduceMeanOp, + "cannot build physical-fragment ReduceMean keepdims reconstruction"); if (semantics->keepdims != 0) { - rewriter.replaceOp(reduceMeanOp, reducedKeepdims); + rewriter.replaceOp(reduceMeanOp, *reducedKeepdims); return success(); } - Value reduced = squeezeReducedAxes(reducedKeepdims, resultType, reducedAxes, rewriter, loc); + Value reduced = squeezeReducedAxes( + *reducedKeepdims, resultType, reducedAxes, rewriter, loc); rewriter.replaceOp(reduceMeanOp, reduced); return success(); } diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp index 2ed6957..4a37d9e 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp @@ -5,7 +5,7 @@ #include "llvm/ADT/SmallVector.h" -#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -52,35 +52,12 @@ static FailureOr materializeTransposedConstant(Value input, return failure(); } - if (denseAttr.isSplat()) - return getOrCreateConstant(rewriter, - rewriter.getInsertionBlock()->getParentOp(), - DenseElementsAttr::get(resultType, denseAttr.getSplatValue()), - resultType); - - SmallVector inputValues(denseAttr.getValues()); - SmallVector resultValues(inputValues.size()); - SmallVector inputStrides = computeRowMajorStrides(inputType.getShape()); - SmallVector resultStrides = computeRowMajorStrides(resultType.getShape()); - SmallVector inputIndices(inputType.getRank(), 0); - - for (auto [linearIndex, value] : llvm::enumerate(inputValues)) { - int64_t remaining = static_cast(linearIndex); - for (int64_t dim = 0; dim < inputType.getRank(); ++dim) { - inputIndices[dim] = inputStrides.empty() ? 0 : remaining / inputStrides[dim]; - remaining = inputStrides.empty() ? 0 : remaining % inputStrides[dim]; - } - - int64_t resultLinearIndex = 0; - for (int64_t dim = 0; dim < resultType.getRank(); ++dim) - resultLinearIndex += inputIndices[permutation[dim]] * resultStrides[dim]; - - resultValues[resultLinearIndex] = value; - } - + auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation); + if (failed(transposedAttr) || transposedAttr->getType() != resultType) + return failure(); return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), - DenseElementsAttr::get(resultType, resultValues), + *transposedAttr, resultType); } diff --git a/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp b/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp index f30f73a..5461e63 100644 --- a/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp @@ -400,6 +400,11 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul if (isa(op)) continue; + // Cloning a region-bearing operation may leave the rewriter inside that + // region. Every old-block operation is lowered at the core-batch body + // boundary. + rewriter.setInsertionPointToEnd(newBlock); + if (auto blueprint = dyn_cast(op)) { std::optional modeAttr = blueprint.getMode(); if (modeAttr && *modeAttr == "fragment_assembly") { diff --git a/src/PIM/Conversion/SpatialToPim/Common.cpp b/src/PIM/Conversion/SpatialToPim/Common.cpp index f5859d2..d34912b 100644 --- a/src/PIM/Conversion/SpatialToPim/Common.cpp +++ b/src/PIM/Conversion/SpatialToPim/Common.cpp @@ -8,6 +8,8 @@ #include #include "Common.hpp" +#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" @@ -192,21 +194,23 @@ forEachContiguousDestinationChunk(ArrayRef destShape, } static mlir::Value -createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index, int64_t stepBytes) { +createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index, + int64_t stepBytes, Operation *constantAnchor) { if (stepBytes == 0) return start; - mlir::Value step = arith::ConstantIndexOp::create(builder, loc, stepBytes); - mlir::Value scaled = arith::MulIOp::create(builder, loc, index, step).getResult(); - return arith::AddIOp::create(builder, loc, start, scaled).getResult(); + return createOrFoldAffineApply( + builder, loc, builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1) * stepBytes, + ValueRange {start, index}, constantAnchor); } static mlir::Value createIndexedOffset(OpBuilder& builder, Location loc, mlir::Value indexArg, - ArrayRef values) { + ArrayRef values, + Operation *constantAnchor) { assert(!values.empty() && "expected lane-indexed values"); if (llvm::all_of(values.drop_front(), [&](int64_t value) { return value == values.front(); })) - return arith::ConstantIndexOp::create(builder, loc, values.front()); + return getOrCreateIndexConstant(builder, constantAnchor, values.front()); if (values.size() >= 2) { int64_t step = values[1] - values[0]; @@ -214,21 +218,18 @@ static mlir::Value createIndexedOffset(OpBuilder& builder, return values[index] == values.front() + static_cast(index) * step; }); if (arithmetic) { - mlir::Value base = arith::ConstantIndexOp::create(builder, loc, values.front()); - mlir::Value stepValue = arith::ConstantIndexOp::create(builder, loc, step); - mlir::Value scaledIndex = arith::MulIOp::create(builder, loc, indexArg, stepValue).getResult(); - return arith::AddIOp::create(builder, loc, base, scaledIndex).getResult(); + return createOrFoldAffineApply( + builder, loc, builder.getAffineDimExpr(0) * step + values.front(), + ValueRange {indexArg}, constantAnchor); } } - mlir::Value selected = arith::ConstantIndexOp::create(builder, loc, values.front()); - for (auto [lane, value] : llvm::enumerate(values.drop_front())) { - mlir::Value indexValue = arith::ConstantIndexOp::create(builder, loc, static_cast(lane + 1)); - mlir::Value cmp = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::eq, indexArg, indexValue); - mlir::Value candidate = arith::ConstantIndexOp::create(builder, loc, value); - selected = arith::SelectOp::create(builder, loc, cmp, candidate, selected); - } - return selected; + RankedTensorType tableType = RankedTensorType::get( + {static_cast(values.size())}, builder.getI64Type()); + DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values); + mlir::Value table = getOrCreateConstant(builder, constantAnchor, tableAttr, tableType); + mlir::Value selected = tensor::ExtractOp::create(builder, loc, table, ValueRange {indexArg}); + return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), selected).getResult(); } struct FragmentAssemblyCopyRunFamily { @@ -433,11 +434,11 @@ static FailureOr emitFragmentAssemblyCopyRun(OpBuilder& builder, mlir::Value hostStart; mlir::Value sourceStart; if (laneArg) { - hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane); - sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane); + hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane, anchor); + sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane, anchor); } else { - hostStart = arith::ConstantIndexOp::create(builder, loc, run.hostStartBytesByLane.front()); - sourceStart = arith::ConstantIndexOp::create(builder, loc, run.sourceStartBytesByLane.front()); + hostStart = getOrCreateIndexConstant(builder, anchor, run.hostStartBytesByLane.front()); + sourceStart = getOrCreateIndexConstant(builder, anchor, run.sourceStartBytesByLane.front()); } if (hostRunStartDelta) @@ -459,9 +460,9 @@ static FailureOr emitFragmentAssemblyCopyRun(OpBuilder& builder, .getOutput(); } - mlir::Value lowerBound = arith::ConstantIndexOp::create(builder, loc, 0); - mlir::Value upperBound = arith::ConstantIndexOp::create(builder, loc, run.count); - mlir::Value step = arith::ConstantIndexOp::create(builder, loc, 1); + mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0); + mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, run.count); + mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1); FailureOr loop = buildNormalizedScfFor( builder, loc, @@ -474,9 +475,10 @@ static FailureOr emitFragmentAssemblyCopyRun(OpBuilder& builder, mlir::Value flatIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { - mlir::Value hostOffset = createSteppedOffset(loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes); + mlir::Value hostOffset = createSteppedOffset( + loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes, anchor); mlir::Value sourceOffset = - createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes); + createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes, anchor); mlir::Value copied = pim::PimMemCopyDevToHostOp::create(loopBuilder, bodyLoc, @@ -506,9 +508,9 @@ static FailureOr emitFragmentAssemblyCopyRunFamily(OpBuilder& build return emitFragmentAssemblyCopyRun( builder, loc, family.prototype, hostTarget, anchor, laneArg, baseHostOffset); - mlir::Value lowerBound = arith::ConstantIndexOp::create(builder, loc, 0); - mlir::Value upperBound = arith::ConstantIndexOp::create(builder, loc, family.sourceRunStartDeltas.size()); - mlir::Value step = arith::ConstantIndexOp::create(builder, loc, 1); + mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0); + mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, family.sourceRunStartDeltas.size()); + mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1); FailureOr outerLoop = buildNormalizedScfFor( builder, loc, @@ -522,9 +524,9 @@ static FailureOr emitFragmentAssemblyCopyRunFamily(OpBuilder& build ValueRange iterArgs, SmallVectorImpl& yielded) { mlir::Value sourceRunStartDelta = - createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas); + createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas, anchor); mlir::Value hostRunStartDelta = - createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas); + createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas, anchor); FailureOr copied = emitFragmentAssemblyCopyRun(loopBuilder, bodyLoc, family.prototype, diff --git a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp b/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp index fba107e..d4e0648 100644 --- a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp @@ -10,7 +10,9 @@ #include "Conversion/ONNXToSpatial/Common/Common.hpp" #include "Conversion/SpatialToPim/SpatialToPimPass.hpp" #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp" @@ -180,16 +182,79 @@ static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute com return success(); } +static bool isHostMaterializableHelperOp(Operation* op) { + if (isa(op)) + return true; + if (isa(op) || op->hasTrait()) + return true; + if (auto blueprint = dyn_cast(op)) { + std::optional mode = blueprint.getMode(); + return mode && *mode == "fragment_assembly"; + } + return isShapingOnlyOp(op) || isPureIndexComputationOp(op); +} + +static FailureOr> +analyzeHostMaterializableHelper(spatial::SpatScheduledCompute computeOp) { + DenseMap folded; + for (auto [weightIndex, weight] : llvm::enumerate(computeOp.getWeights())) { + auto argument = computeOp.getWeightArgument(weightIndex); + if (!argument) + return failure(); + Attribute constant; + if (matchPattern(weight, m_Constant(&constant))) + folded[*argument] = constant; + } + Block& block = computeOp.getBody().front(); + for (Operation& op : block) { + if (!isHostMaterializableHelperOp(&op)) + return failure(); + if (isa(op) + || (isShapingOnlyOp(&op) && !isPureIndexComputationOp(&op))) + continue; + if (isa(op) || op.hasTrait()) { + for (Value result : op.getResults()) { + Attribute constant; + if (!matchPattern(result, m_Constant(&constant))) + return failure(); + folded[result] = constant; + } + continue; + } + if (!isPureIndexComputationOp(&op) || op.getNumRegions() != 0) + return failure(); + SmallVector operands; + for (Value operand : op.getOperands()) { + auto it = folded.find(operand); + if (it == folded.end()) + return failure(); + operands.push_back(it->second); + } + SmallVector results; + if (failed(op.fold(operands, results)) + || results.size() != op.getNumResults()) + return failure(); + for (auto [result, foldResult] : llvm::zip(op.getResults(), results)) { + auto attribute = dyn_cast(foldResult); + if (!attribute) + return failure(); + folded[result] = attribute; + } + } + return folded; +} + static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatScheduledCompute computeOp, IRRewriter& rewriter, OperationFolder& constantFolder) { if (!computeOp.getInputs().empty() || computeOp.getNumResults() != 1) return false; + if (computeOp.getResult(0).use_empty()) + return false; if (!llvm::all_of(computeOp.getResult(0).getUsers(), [](Operation* user) { return isa(user); })) return false; - Block& block = computeOp.getBody().front(); if (block.getNumArguments() != computeOp.getWeights().size()) return false; @@ -197,6 +262,9 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule auto yieldOp = dyn_cast(block.getTerminator()); if (!yieldOp || yieldOp.getNumOperands() != 1) return false; + auto folded = analyzeHostMaterializableHelper(computeOp); + if (failed(folded)) + return false; rewriter.setInsertionPoint(computeOp); IRMapping mapping; @@ -218,6 +286,20 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule } } + if (isa(op) || op.hasTrait() + || isPureIndexComputationOp(&op)) { + for (Value result : op.getResults()) { + auto it = folded->find(result); + if (it == folded->end()) + return false; + mapping.map( + result, + getOrCreateConstant(constantFolder, computeOp, it->second, + result.getType())); + } + continue; + } + cloneMappedHelperOperands(&op, mapping, rewriter, constantFolder); Operation* clonedOp = rewriter.clone(op, mapping); for (auto [originalResult, newResult] : llvm::zip(op.getResults(), clonedOp->getResults())) diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/Common.cpp b/src/PIM/Dialect/Pim/Transforms/Bufferization/Common.cpp index 129078d..6b752b5 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/Common.cpp +++ b/src/PIM/Dialect/Pim/Transforms/Bufferization/Common.cpp @@ -1,10 +1,23 @@ #include "Dialect/Pim/Transforms/Bufferization/Common.hpp" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" using namespace mlir; +static SmallVector getSelectionRegions(OpResult result) { + SmallVector regions; + if (auto selection = dyn_cast(result.getOwner())) + for (Region ®ion : selection->getRegions()) + regions.push_back(®ion); + else if (auto selection = dyn_cast(result.getOwner())) { + regions.push_back(&selection.getThenRegion()); + regions.push_back(&selection.getElseRegion()); + } + return regions; +} + static bool isCoreBatchInputArgument(Value value) { auto blockArg = dyn_cast(value); if (!blockArg) @@ -92,20 +105,46 @@ FailureOr onnx_mlir::pim::getPimAddressBase(Value value, const StaticValu } bool onnx_mlir::pim::isHostBackedPimAddress(Value value, const StaticValueKnowledge& knowledge) { - auto base = getPimStorageBase(value, knowledge); - if (failed(base)) - return false; - - if (isCoreBatchInputArgument(*base)) - return true; - - return isa_and_nonnull(base->getDefiningOp()); + llvm::SmallPtrSet visited; + std::function isHost = [&](Value current) { + auto base = getPimStorageBase(current, knowledge); + if (failed(base) || !visited.insert(*base).second) + return false; + bool resultIsHost = isCoreBatchInputArgument(*base) + || isa_and_nonnull(base->getDefiningOp()); + auto result = dyn_cast(*base); + SmallVector regions = result ? getSelectionRegions(result) + : SmallVector(); + if (!resultIsHost && !regions.empty()) + resultIsHost = llvm::all_of(regions, [&](Region *region) { + auto yield = dyn_cast(region->front().getTerminator()); + return yield && result.getResultNumber() < yield.getNumOperands() + && isHost(yield.getOperand(result.getResultNumber())); + }); + visited.erase(*base); + return resultIsHost; + }; + return isHost(value); } bool onnx_mlir::pim::isDeviceLocalPimAddress(Value value, const StaticValueKnowledge& knowledge) { - auto base = getPimStorageBase(value, knowledge); - if (failed(base)) - return false; - - return isa_and_nonnull(base->getDefiningOp()); + llvm::SmallPtrSet visited; + std::function isDevice = [&](Value current) { + auto base = getPimStorageBase(current, knowledge); + if (failed(base) || !visited.insert(*base).second) + return false; + bool resultIsDevice = isa_and_nonnull(base->getDefiningOp()); + auto result = dyn_cast(*base); + SmallVector regions = result ? getSelectionRegions(result) + : SmallVector(); + if (!resultIsDevice && !regions.empty()) + resultIsDevice = llvm::all_of(regions, [&](Region *region) { + auto yield = dyn_cast(region->front().getTerminator()); + return yield && result.getResultNumber() < yield.getNumOperands() + && isDevice(yield.getOperand(result.getResultNumber())); + }); + visited.erase(*base); + return resultIsDevice; + }; + return isDevice(value); } diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp b/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp index 2df4402..fc0cfa2 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp +++ b/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp @@ -2,6 +2,8 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" +#include "llvm/Support/MathExtras.h" + #include "ContiguityPatterns.hpp" #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" @@ -33,6 +35,7 @@ struct CopyEndpointPlan { struct CopyLoopPlan { SmallVector outerShape; + int64_t outerElements = 0; int64_t chunkBytes = 0; ByteOffsetExpr targetBaseOffset; ByteOffsetExpr sourceBaseOffset; @@ -74,6 +77,24 @@ static void appendTerm(ByteOffsetExpr& expr, Value value, int64_t scale) { expr.terms.push_back(ByteOffsetTerm {value, scale}); } +static FailureOr checkedPositiveMul(int64_t lhs, int64_t rhs) { + int64_t result = 0; + if (lhs < 0 || rhs < 0 || llvm::MulOverflow(lhs, rhs, result)) + return failure(); + return result; +} + +static FailureOr checkedPositiveProduct(ArrayRef values) { + int64_t result = 1; + for (int64_t value : values) { + auto product = checkedPositiveMul(result, value); + if (failed(product)) + return failure(); + result = *product; + } + return result; +} + static FailureOr> getStaticMemRefStrides(MemRefType type) { SmallVector strides; int64_t offset = 0; @@ -84,6 +105,165 @@ static FailureOr> getStaticMemRefStrides(MemRefType type) { return strides; } +static FailureOr> getProvenMemRefStrides(Value value) { + llvm::SmallPtrSet visiting; + std::function>(Value)> prove = + [&](Value current) -> FailureOr> { + auto type = dyn_cast(current.getType()); + if (!type || !visiting.insert(current).second) + return failure(); + if (auto strides = getStaticMemRefStrides(type); succeeded(strides)) { + visiting.erase(current); + return strides; + } + if (auto castOp = current.getDefiningOp()) { + auto strides = prove(castOp.getSource()); + visiting.erase(current); + return strides; + } + if (auto subview = current.getDefiningOp()) { + auto sourceStrides = prove(subview.getSource()); + if (failed(sourceStrides) || subview.getSourceType().getRank() != subview.getType().getRank()) { + visiting.erase(current); + return failure(); + } + SmallVector strides; + for (auto [sourceStride, viewStride] : + llvm::zip_equal(*sourceStrides, subview.getStaticStrides())) { + if (ShapedType::isDynamic(viewStride) || viewStride < 0) { + visiting.erase(current); + return failure(); + } + auto stride = checkedPositiveMul(sourceStride, viewStride); + if (failed(stride)) { + visiting.erase(current); + return failure(); + } + strides.push_back(*stride); + } + visiting.erase(current); + return strides; + } + if (auto expand = current.getDefiningOp()) { + auto sourceStrides = prove(expand.getSrc()); + auto resultType = dyn_cast(expand.getResult().getType()); + auto sourceType = dyn_cast(expand.getSrc().getType()); + if (failed(sourceStrides) || !sourceType || !resultType + || !resultType.hasStaticShape() + || sourceStrides->size() != static_cast(sourceType.getRank()) + || llvm::any_of(resultType.getShape(), [](int64_t dim) { + return dim <= 0; + })) { + visiting.erase(current); + return failure(); + } + SmallVector strides(resultType.getRank()); + SmallVector assigned(resultType.getRank(), false); + for (auto [sourceDim, group] : + llvm::enumerate(expand.getReassociationIndices())) { + if (sourceDim >= sourceStrides->size() || group.empty()) { + visiting.erase(current); + return failure(); + } + int64_t stride = (*sourceStrides)[sourceDim]; + for (int64_t resultDim : llvm::reverse(group)) { + if (resultDim < 0 || resultDim >= resultType.getRank() + || assigned[resultDim]) { + visiting.erase(current); + return failure(); + } + strides[resultDim] = stride; + assigned[resultDim] = true; + auto nextStride = checkedPositiveMul( + stride, resultType.getDimSize(resultDim)); + if (failed(nextStride)) { + visiting.erase(current); + return failure(); + } + stride = *nextStride; + } + } + if (llvm::is_contained(assigned, false)) { + visiting.erase(current); + return failure(); + } + visiting.erase(current); + return strides; + } + if (auto collapse = current.getDefiningOp()) { + auto sourceStrides = prove(collapse.getSrc()); + auto sourceType = dyn_cast(collapse.getSrc().getType()); + if (failed(sourceStrides) || !sourceType + || !sourceType.hasStaticShape() + || sourceStrides->size() != static_cast(sourceType.getRank())) { + visiting.erase(current); + return failure(); + } + SmallVector strides; + for (ArrayRef group : collapse.getReassociationIndices()) { + if (group.empty()) { + visiting.erase(current); + return failure(); + } + for (int64_t dim : group) + if (dim < 0 || dim >= sourceType.getRank() + || sourceType.getDimSize(dim) <= 0 + || (*sourceStrides)[dim] < 0) { + visiting.erase(current); + return failure(); + } + for (auto pair : llvm::zip(group.drop_back(), group.drop_front())) { + int64_t outer = std::get<0>(pair); + int64_t inner = std::get<1>(pair); + auto expectedOuterStride = checkedPositiveMul( + (*sourceStrides)[inner], sourceType.getDimSize(inner)); + if (failed(expectedOuterStride) + || (*sourceStrides)[outer] != *expectedOuterStride) { + visiting.erase(current); + return failure(); + } + } + strides.push_back((*sourceStrides)[group.back()]); + } + visiting.erase(current); + return strides; + } + auto result = dyn_cast(current); + SmallVector regions; + if (result) { + if (auto selection = dyn_cast(result.getOwner())) + for (Region ®ion : selection->getRegions()) + regions.push_back(®ion); + else if (auto selection = dyn_cast(result.getOwner())) { + regions.push_back(&selection.getThenRegion()); + regions.push_back(&selection.getElseRegion()); + } + } + if (regions.empty()) { + visiting.erase(current); + return failure(); + } + std::optional> common; + for (Region *region : regions) { + auto yield = dyn_cast(region->front().getTerminator()); + if (!yield || result.getResultNumber() >= yield.getNumOperands()) { + visiting.erase(current); + return failure(); + } + auto strides = prove(yield.getOperand(result.getResultNumber())); + if (failed(strides) || (common && *common != *strides)) { + visiting.erase(current); + return failure(); + } + common = std::move(*strides); + } + visiting.erase(current); + return common ? FailureOr>(std::move(*common)) + : FailureOr>(failure()); + }; + return prove(value); +} + static FailureOr getShapedByteSize(MemRefType type) { if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())) return failure(); @@ -119,12 +299,15 @@ inferLogicalCopyShape(MemRefType targetType, MemRefType sourceType, int64_t size return failure(); } -static FailureOr getContiguousSuffixRank(MemRefType type, ArrayRef copyShape) { - if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType()) +static FailureOr getContiguousSuffixRank(Value value, ArrayRef copyShape) { + auto type = dyn_cast(value.getType()); + if (!type || !type.hasStaticShape() || !hasByteSizedElementType(type.getElementType()) || type.getRank() != static_cast(copyShape.size())) return failure(); + if (llvm::any_of(copyShape, [](int64_t dim) { return dim <= 0; })) + return failure(); - auto strides = getStaticMemRefStrides(type); + auto strides = getProvenMemRefStrides(value); if (failed(strides)) return failure(); @@ -134,7 +317,10 @@ static FailureOr getContiguousSuffixRank(MemRefType type, ArrayRef analyzeCopyEndpoint(Value value, Value initia if (!sourceType || !sourceType.hasStaticShape() || !hasByteSizedElementType(sourceType.getElementType())) return failure(); - auto sourceStrides = getStaticMemRefStrides(sourceType); + auto sourceStrides = getProvenMemRefStrides(subviewOp.getSource()); if (failed(sourceStrides)) return failure(); int64_t elementByteWidth = static_cast(getElementTypeSizeInBytes(sourceType.getElementType())); for (auto [offset, stride] : llvm::zip_equal(subviewOp.getMixedOffsets(), *sourceStrides)) { - int64_t byteScale = stride * elementByteWidth; + auto byteScale = checkedPositiveMul(stride, elementByteWidth); + if (failed(byteScale)) + return failure(); if (auto attr = dyn_cast(offset)) { - endpoint.offset.constant += cast(attr).getInt() * byteScale; + auto constantOffset = checkedPositiveMul( + cast(attr).getInt(), *byteScale); + if (failed(constantOffset) + || llvm::AddOverflow(endpoint.offset.constant, *constantOffset, + endpoint.offset.constant)) + return failure(); continue; } - appendTerm(endpoint.offset, cast(offset), byteScale); + appendTerm(endpoint.offset, cast(offset), *byteScale); } endpoint.base = subviewOp.getSource(); @@ -213,8 +406,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO auto sourceBytes = getShapedByteSize(sourceType); if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes) && *targetBytes == size && *sourceBytes == size) { - auto targetSuffixRank = getContiguousSuffixRank(targetType, targetType.getShape()); - auto sourceSuffixRank = getContiguousSuffixRank(sourceType, sourceType.getShape()); + auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape()); + auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape()); if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank) && *targetSuffixRank == targetType.getRank() && *sourceSuffixRank == sourceType.getRank()) { CopyRewritePlan plan; @@ -230,8 +423,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO if (failed(logicalCopyShape)) return failure(); - auto targetSuffixRank = getContiguousSuffixRank(targetType, *logicalCopyShape); - auto sourceSuffixRank = getContiguousSuffixRank(sourceType, *logicalCopyShape); + auto targetSuffixRank = getContiguousSuffixRank(target, *logicalCopyShape); + auto sourceSuffixRank = getContiguousSuffixRank(source, *logicalCopyShape); if (failed(targetSuffixRank) || failed(sourceSuffixRank)) return failure(); @@ -246,8 +439,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO return plan; } - auto targetStrides = getStaticMemRefStrides(targetType); - auto sourceStrides = getStaticMemRefStrides(sourceType); + auto targetStrides = getProvenMemRefStrides(target); + auto sourceStrides = getProvenMemRefStrides(source); if (failed(targetStrides) || failed(sourceStrides)) return failure(); @@ -257,11 +450,27 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO plan.loop.sourceBaseOffset = plan.source.offset; plan.loop.outerShape.assign(logicalCopyShape->begin(), logicalCopyShape->end() - contiguousSuffixRank); SmallVector chunkShape(logicalCopyShape->end() - contiguousSuffixRank, logicalCopyShape->end()); - plan.loop.chunkBytes = getNumElements(chunkShape) * elementByteWidth; - for (int64_t stride : ArrayRef(*targetStrides).take_front(plan.loop.outerShape.size())) - plan.loop.targetOuterByteStrides.push_back(stride * elementByteWidth); - for (int64_t stride : ArrayRef(*sourceStrides).take_front(plan.loop.outerShape.size())) - plan.loop.sourceOuterByteStrides.push_back(stride * elementByteWidth); + auto outerElements = checkedPositiveProduct(plan.loop.outerShape); + auto chunkElements = checkedPositiveProduct(chunkShape); + auto chunkBytes = failed(chunkElements) + ? FailureOr(failure()) + : checkedPositiveMul(*chunkElements, elementByteWidth); + if (failed(outerElements) || failed(chunkBytes)) + return failure(); + plan.loop.outerElements = *outerElements; + plan.loop.chunkBytes = *chunkBytes; + for (int64_t stride : ArrayRef(*targetStrides).take_front(plan.loop.outerShape.size())) { + auto byteStride = checkedPositiveMul(stride, elementByteWidth); + if (failed(byteStride)) + return failure(); + plan.loop.targetOuterByteStrides.push_back(*byteStride); + } + for (int64_t stride : ArrayRef(*sourceStrides).take_front(plan.loop.outerShape.size())) { + auto byteStride = checkedPositiveMul(stride, elementByteWidth); + if (failed(byteStride)) + return failure(); + plan.loop.sourceOuterByteStrides.push_back(*byteStride); + } if (plan.loop.chunkBytes <= 0) return failure(); return plan; @@ -361,7 +570,7 @@ static LogicalResult rewriteCopyLikeOp(CopyOp copyOp, } Value c0 = createIndexConstant(rewriter, anchorOp, 0); - Value cUpper = createIndexConstant(rewriter, anchorOp, getNumElements(plan->loop.outerShape)); + Value cUpper = createIndexConstant(rewriter, anchorOp, plan->loop.outerElements); Value cStep = createIndexConstant(rewriter, anchorOp, 1); auto loop = buildNormalizedScfFor( rewriter, diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index 33f02a9..1c2b303 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -17,6 +17,7 @@ add_pim_library(SpatialOps Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp Transforms/MergeComputeNodes/ScheduledComputeReport.cpp Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp + Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp diff --git a/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp b/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp index 7b3bb5b..ec6944b 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp @@ -8,6 +8,7 @@ #include "llvm/Support/LogicalResult.h" #include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" using namespace mlir; @@ -130,7 +131,7 @@ struct CanonicalizeSingleLaneComputeBatchPattern : OpRewritePattern> 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(); + DenseMap headSends; + DenseMap headReceives; + SmallVector readyComputes; + SmallVector readyExchanges; + size_t computeCursor = 0; + size_t exchangeCursor = 0; + unsigned finishedStreams = 0; - 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; - } - } + auto registerHead = [&](unsigned stream) { + if (cursor[stream] == streams[stream].size()) { + ++finishedStreams; + return; } - 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(); + const Event &event = streams[stream][cursor[stream]]; + if (event.kind == EventKind::Compute) { + readyComputes.push_back(stream); + return; } + DenseMap &heads = + event.kind == EventKind::Send ? headSends : headReceives; + DenseMap &peers = + event.kind == EventKind::Send ? headReceives : headSends; + heads[event.exchangeId] = stream; + if (peers.contains(event.exchangeId)) + readyExchanges.push_back(event.exchangeId); + }; + for (unsigned stream = 0; stream < streams.size(); ++stream) + registerHead(stream); + + while (computeCursor != readyComputes.size() + || exchangeCursor != readyExchanges.size()) { + if (computeCursor != readyComputes.size()) { + unsigned stream = readyComputes[computeCursor++]; + ++cursor[stream]; + registerHead(stream); + continue; + } + + uint64_t exchange = readyExchanges[exchangeCursor++]; + auto send = headSends.find(exchange); + auto receive = headReceives.find(exchange); + if (send == headSends.end() || receive == headReceives.end()) + continue; + unsigned source = send->second; + unsigned target = receive->second; + headSends.erase(send); + headReceives.erase(receive); + ++cursor[source]; + ++cursor[target]; + registerHead(source); + registerHead(target); } + + if (finishedStreams == streams.size()) + return success(); + 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) { @@ -79,6 +97,98 @@ static std::optional getI64Attr(Operation *op, StringRef name) { return std::nullopt; } +static LogicalResult getI64ArrayAttr( + Operation *op, StringRef name, + std::optional> &values) { + Attribute attr = op->getAttr(name); + if (!attr) + return success(); + if (auto array = dyn_cast(attr)) { + values.emplace(array.asArrayRef()); + return success(); + } + auto elements = dyn_cast(attr); + auto type = elements ? dyn_cast(elements.getType()) + : RankedTensorType(); + if (!elements || !type || type.getRank() != 1 + || !type.getElementType().isInteger(64)) + return op->emitOpError() << "has invalid " << name << " metadata"; + values.emplace(); + values->reserve(elements.getNumElements()); + for (const APInt &value : elements.getValues()) + values->push_back(value.getSExtValue()); + return success(); +} + +struct RealizedLogicalTransfer { + int64_t channelId = -1; + int64_t parentExchangeId = -1; + int64_t parentTransferCount = 0; + int64_t sourceCore = -1; + int64_t targetCore = -1; +}; + +static LogicalResult forEachRealizedLogicalTransfer( + Operation *op, + function_ref callback) { + auto scalarChannel = getI64Attr(op, "raptor.channel_id"); + std::optional> batchChannels; + if (failed(getI64ArrayAttr( + op, "raptor.batch_channel_ids", batchChannels))) + return failure(); + if (scalarChannel && batchChannels) + return op->emitOpError( + "mixes scalar and compact logical transfer metadata"); + + if (scalarChannel) { + auto exchange = getI64Attr(op, "raptor.exchange_id"); + auto parent = getI64Attr(op, "raptor.parent_exchange_id"); + auto count = getI64Attr(op, "raptor.parent_transfer_count"); + auto source = getI64Attr(op, "raptor.source_core"); + auto target = getI64Attr(op, "raptor.target_core"); + if (!exchange || !parent || !count || !source || !target) + return op->emitOpError( + "is missing scalar logical transfer metadata"); + RealizedLogicalTransfer transfer { + *scalarChannel, *parent, *count, *source, *target}; + if (*exchange != transfer.channelId || transfer.channelId < 0 + || transfer.parentExchangeId < 0 || transfer.parentTransferCount <= 0 + || transfer.sourceCore < 0 || transfer.targetCore < 0) + return op->emitOpError("has invalid scalar logical transfer metadata"); + return callback(transfer); + } + + std::optional> sources, targets, parents, counts; + if (failed(getI64ArrayAttr(op, "raptor.batch_source_cores", sources)) + || failed(getI64ArrayAttr(op, "raptor.batch_target_cores", targets)) + || failed(getI64ArrayAttr( + op, "raptor.batch_parent_exchange_ids", parents)) + || failed(getI64ArrayAttr( + op, "raptor.batch_parent_transfer_counts", counts))) + return failure(); + if (!batchChannels || !sources || !targets || !parents || !counts) + return op->emitOpError( + "is missing compact logical transfer metadata"); + size_t size = batchChannels->size(); + if (size == 0 || sources->size() != size || targets->size() != size + || parents->size() != size || counts->size() != size) + return op->emitOpError( + "has non-parallel compact logical transfer metadata"); + for (auto values : llvm::zip_equal( + *batchChannels, *parents, *counts, *sources, *targets)) { + RealizedLogicalTransfer transfer { + std::get<0>(values), std::get<1>(values), std::get<2>(values), + std::get<3>(values), std::get<4>(values)}; + if (transfer.channelId < 0 || transfer.parentExchangeId < 0 + || transfer.parentTransferCount <= 0 || transfer.sourceCore < 0 + || transfer.targetCore < 0) + return op->emitOpError("has invalid compact logical transfer metadata"); + if (failed(callback(transfer))) + return failure(); + } + return success(); +} + } // namespace LogicalResult verifyPlannedCommunicationDeadlockFree( @@ -118,7 +228,11 @@ LogicalResult verifyPlannedCommunicationDeadlockFree( } LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) { - DenseMap> operationsByExchange; + struct LogicalOperation { + Operation *op = nullptr; + RealizedLogicalTransfer transfer; + }; + DenseMap> operationsByExchange; struct ParentExchange { std::optional expectedTransfers; DenseSet channels; @@ -130,39 +244,25 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) { 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); + if (failed(forEachRealizedLogicalTransfer( + op, [&](const RealizedLogicalTransfer &transfer) -> LogicalResult { + operationsByExchange[transfer.channelId].push_back( + {op, transfer}); + ParentExchange &parent = + parentExchanges[transfer.parentExchangeId]; + if (parent.expectedTransfers + && *parent.expectedTransfers + != transfer.parentTransferCount) + return op->emitOpError( + "declares an inconsistent parent transfer count"); + parent.expectedTransfers = transfer.parentTransferCount; + parent.channels.insert(transfer.channelId); + for (int64_t core : {transfer.sourceCore, transfer.targetCore}) + if (!llvm::is_contained(cores, core)) + cores.push_back(core); + return success(); + }))) + invalid = true; }); llvm::sort(cores); for (auto [index, core] : llvm::enumerate(cores)) @@ -172,38 +272,18 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) { 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"); + if (failed(forEachRealizedLogicalTransfer( + op, [&](const RealizedLogicalTransfer &transfer) { + unsigned stream = streamByCore.lookup( + isa(op) ? transfer.sourceCore + : transfer.targetCore); + streams[stream].push_back( + {isa(op) ? EventKind::Send + : EventKind::Receive, + static_cast(transfer.channelId)}); + return success(); + }))) 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(); @@ -216,56 +296,38 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) { << " 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 (entry.second.size() != 2 + || isa(entry.second[0].op) + == isa(entry.second[1].op)) { + return funcOp.emitOpError() + << "exchange " << entry.first + << " does not have exactly one send and one receive (sends=" + << llvm::count_if(entry.second, [](const LogicalOperation &item) { + return isa(item.op); + }) + << ", receives=" + << llvm::count_if(entry.second, [](const LogicalOperation &item) { + return isa(item.op); + }) + << ")"; } + const LogicalOperation &first = entry.second[0]; + const LogicalOperation &second = entry.second[1]; + const LogicalOperation &sendRecord = + isa(first.op) ? first : second; + const LogicalOperation &receiveRecord = + isa(first.op) ? first : second; + auto send = cast(sendRecord.op); + auto receive = cast(receiveRecord.op); 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) + if (receiveRecord.transfer.sourceCore != sendRecord.transfer.sourceCore + || receiveRecord.transfer.targetCore + != sendRecord.transfer.targetCore + || receiveRecord.transfer.parentExchangeId + != sendRecord.transfer.parentExchangeId + || receiveRecord.transfer.parentTransferCount + != sendRecord.transfer.parentTransferCount) return receive.emitOpError("receive core metadata does not match its send"); } return simulate(funcOp, streams, "realized"); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp index 6b08094..b0eb598 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp @@ -8,6 +8,10 @@ #include "llvm/ADT/SmallPtrSet.h" +#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" + namespace onnx_mlir::spatial { using namespace mlir; namespace { @@ -112,11 +116,34 @@ static FailureOr buildBlueprintReconstruction( return result; } -static bool isSupportedDeferredShapingOp(Operation *op) { - return isa(op); +static FailureOr buildIndexSwitchSelection(OpBuilder &builder, Location loc, + Value selector, ValueRange candidates, + Operation *diagnosticOwner) { + if (candidates.empty()) + return diagnosticOwner->emitOpError("direct selection requires at least one candidate"), failure(); + Type type = candidates.front().getType(); + if (llvm::any_of(candidates, [&](Value candidate) { return candidate.getType() != type; })) + return diagnosticOwner->emitOpError("direct selection requires identical candidate types"), failure(); + if (candidates.size() == 1) + return candidates.front(); + + SmallVector cases; + for (int64_t index = 0; index < static_cast(candidates.size()) - 1; ++index) + cases.push_back(index); + auto selection = scf::IndexSwitchOp::create( + builder, loc, TypeRange {type}, selector, cases, cases.size()); + auto buildYield = [&](Region ®ion, Value candidate) { + OpBuilder::InsertionGuard guard(builder); + Block *block = builder.createBlock(®ion); + builder.setInsertionPointToEnd(block); + scf::YieldOp::create(builder, loc, candidate); + }; + for (auto [region, candidate] : llvm::zip(selection.getCaseRegions(), candidates.drop_back())) + buildYield(region, candidate); + // The scheduled-lane verifier guarantees an in-range selector, so default is + // the final lane without an otherwise-unreachable extra branch. + buildYield(selection.getDefaultRegion(), candidates.back()); + return selection.getResult(0); } static FailureOr buildSelectedDeferredSource(OpBuilder &builder, Location loc, @@ -128,46 +155,28 @@ static FailureOr buildSelectedDeferredSource(OpBuilder &builder, Location 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(); + auto scheduled = transfer->getParentOfType(); + if (!scheduled || sourceOperandForScheduledLane.size() != static_cast(scheduled.getLaneCount())) + return transfer.emitOpError("deferred source mapping must cover every scheduled lane"), failure(); + SmallVector candidates; + candidates.reserve(sourceOperandForScheduledLane.size()); + for (int64_t sourceIndex : sourceOperandForScheduledLane) { + if (sourceIndex < 0 || sourceIndex >= static_cast(sourceBlockArgs.size())) + return transfer.emitOpError("deferred source mapping operand is out of range"), failure(); + candidates.push_back(sourceBlockArgs[sourceIndex]); } - 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(); + return buildIndexSwitchSelection(builder, loc, scheduledLane, candidates, transfer.getOperation()); } -static bool isTopLevelShaping(Operation *op, Block &body) { - return op->getBlock() == &body && isSupportedDeferredShapingOp(op); +static bool isDeferredPayloadCandidateOp(Operation *op) { + return isShapingOnlyOp(op) || isCompileTimeOp(op) || isPureIndexComputationOp(op); +} + +static bool isTopLevelDeferredOperation(Operation *op, Block &body, + const DeferredInputPlan &plan) { + (void)plan; + return op->getBlock() == &body + && (isDeferredPayloadCandidateOp(op) || isa(op)); } static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan, @@ -180,7 +189,7 @@ static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan, Operation *op = value.getDefiningOp(); if (op && op->hasTrait()) return true; - if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second) + if (!op || !isTopLevelDeferredOperation(op, body, plan) || !seen.insert(op).second) return op && seen.contains(op); return llvm::all_of(op->getOperands(), [&](Value operand) { return isEligible(operand, body, plan, seen); }); } @@ -196,7 +205,7 @@ static FailureOr clonePayloadRoot(Value root, Block &body, const Deferred if (isa(value)) return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure(); Operation *op = value.getDefiningOp(); - if (!op || !isSupportedDeferredShapingOp(op)) + if (!op || (!isDeferredPayloadCandidateOp(op) && !op->hasTrait())) 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); @@ -214,7 +223,7 @@ static FailureOr clonePayloadRoot(Value root, Block &body, const Deferred 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())) + if (!op || (!isTopLevelDeferredOperation(op, body, plan) && !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); @@ -225,65 +234,32 @@ static FailureOr clonePayloadRoot(Value root, Block &body, const Deferred } static bool dependsOnGraphLane(Value value, Value graphLane, Block &body, + const DeferredInputPlan &plan, llvm::SmallPtrSetImpl &seen) { if (value == graphLane) return true; Operation *op = value.getDefiningOp(); - if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second) + if (!op || !isTopLevelDeferredOperation(op, body, plan) || !seen.insert(op).second) return false; - return llvm::any_of(op->getOperands(), [&](Value operand) { - return dependsOnGraphLane(operand, graphLane, body, seen); - }); -} - -static FailureOr buildPayloadAggregate(OpBuilder &builder, Location loc, - ArrayRef payloads) { - auto payloadType = dyn_cast(payloads.front().getType()); - if (!payloadType || !payloadType.hasStaticShape()) - return failure(); - SmallVector shape {static_cast(payloads.size())}; - llvm::append_range(shape, payloadType.getShape()); - auto aggregateType = RankedTensorType::get(shape, payloadType.getElementType()); - auto empty = createEmptyTensorForType(builder, loc, aggregateType); - if (failed(empty)) return failure(); - Value aggregate = *empty; - SmallVector sizes, strides(shape.size(), builder.getIndexAttr(1)); - sizes.push_back(builder.getIndexAttr(1)); - for (int64_t dim : payloadType.getShape()) sizes.push_back(builder.getIndexAttr(dim)); - SmallVector reassociation {{0, 1}}; - for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1}); - SmallVector expandedShape {1}; llvm::append_range(expandedShape, payloadType.getShape()); - auto expandedType = RankedTensorType::get(expandedShape, payloadType.getElementType()); - for (auto [index, payload] : llvm::enumerate(payloads)) { - Value expanded = tensor::ExpandShapeOp::create(builder, loc, expandedType, payload, reassociation).getResult(); - SmallVector offsets(shape.size(), builder.getIndexAttr(0)); - offsets[0] = builder.getIndexAttr(index); - aggregate = tensor::InsertSliceOp::create(builder, loc, expanded, aggregate, offsets, sizes, strides).getResult(); + if (auto loop = dyn_cast(op)) { + bool depends = false; + loop.getRegion().walk([&](Operation *nested) { + depends |= llvm::is_contained(nested->getOperands(), graphLane); + }); + if (depends) + return true; } - return aggregate; -} - -static FailureOr selectPayloadAggregate(OpBuilder &builder, Location loc, Value aggregate, - Value localLane) { - auto aggregateType = cast(aggregate.getType()); - SmallVector payloadShape(aggregateType.getShape().begin() + 1, aggregateType.getShape().end()); - auto payloadType = RankedTensorType::get(payloadShape, aggregateType.getElementType()); - SmallVector offsets(aggregateType.getRank(), builder.getIndexAttr(0)); offsets[0] = localLane; - SmallVector sizes, strides(aggregateType.getRank(), builder.getIndexAttr(1)); - sizes.push_back(builder.getIndexAttr(1)); - for (int64_t dim : payloadShape) sizes.push_back(builder.getIndexAttr(dim)); - SmallVector unitShape {1}; llvm::append_range(unitShape, payloadShape); - Value unit = tensor::ExtractSliceOp::create(builder, loc, - RankedTensorType::get(unitShape, aggregateType.getElementType()), aggregate, offsets, sizes, strides).getResult(); - SmallVector reassociation {{0, 1}}; - for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1}); - return tensor::CollapseShapeOp::create(builder, loc, payloadType, unit, reassociation).getResult(); + return llvm::any_of(op->getOperands(), [&](Value operand) { + return dependsOnGraphLane(operand, graphLane, body, plan, seen); + }); } 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; + if (!op || !isTopLevelDeferredOperation(op, body, plan) || !ops.insert(op).second) return; + if (auto loop = dyn_cast(op)) + loop.getRegion().walk([&](Operation *nested) { ops.insert(nested); }); for (Value operand : op->getOperands()) if (operand != plan.graphInput && operand != plan.graphLane) collectClosure(operand, body, plan, ops); } @@ -371,12 +347,12 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc if (!seen.insert(value).second) continue; for (OpOperand &use : value.getUses()) { Operation *user = use.getOwner(); - if (!isTopLevelShaping(user, body)) { needsIdentity = true; continue; } + if (!isTopLevelDeferredOperation(user, body, plan)) { 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); }); + bool hasShapingUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return isTopLevelDeferredOperation(next.getOwner(), body, plan); }); + bool hasOtherUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return !isTopLevelDeferredOperation(next.getOwner(), body, plan); }); if (hasOtherUse) roots.push_back(result); if (hasShapingUse) worklist.push_back(result); } @@ -388,7 +364,7 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc for (Value root : roots) { llvm::SmallPtrSet laneDependencies; bool scalarize = plan.scalarizedGraphLaneBase - && dependsOnGraphLane(root, plan.graphLane, body, laneDependencies); + && dependsOnGraphLane(root, plan.graphLane, body, plan, laneDependencies); OpBuilder::InsertPoint restore = builder.saveInsertionPoint(); Operation *loop = nullptr; if (scalarize) { @@ -419,9 +395,8 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc if (failed(selected)) return failure(); Value boundGraphLane; if (scalarize) { - Value offsetValue = arith::ConstantIndexOp::create(builder, loc, offset); - boundGraphLane = offset ? arith::AddIOp::create(builder, loc, plan.scalarizedGraphLaneBase, offsetValue).getResult() - : plan.scalarizedGraphLaneBase; + boundGraphLane = affineAddConst( + builder, loc, plan.scalarizedGraphLaneBase, offset, transfer.getOperation()); } auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected, boundGraphLane); if (failed(payload)) return failure(); @@ -431,30 +406,26 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc } if (scalarize) { builder.restoreInsertionPoint(restore); - if (payloads.size() == 1) { - mapper.map(root, payloads.front()); - } else { - if (loop) builder.setInsertionPoint(loop); - else builder.setInsertionPointToEnd(plan.scalarizedHoistBlock); - auto aggregate = buildPayloadAggregate(builder, loc, payloads); - if (failed(aggregate)) return failure(); - builder.restoreInsertionPoint(restore); - auto selected = selectPayloadAggregate(builder, loc, *aggregate, plan.scalarizedLocalLane); - if (failed(selected)) return failure(); - mapper.map(root, *selected); - } + auto selected = buildIndexSwitchSelection( + builder, loc, plan.scalarizedLocalLane, payloads, root.getDefiningOp()); + if (failed(selected)) return failure(); + mapper.map(root, *selected); } else { mapper.map(root, payloads.front()); } collectClosure(root, body, plan, absorbed); } } + SmallVector notFullyAbsorbed; 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); + if (!allResultsMapped) + notFullyAbsorbed.push_back(op); } + for (Operation *op : notFullyAbsorbed) + absorbed.erase(op); return success(); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp index b65ddc3..d53b1c9 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp @@ -2,21 +2,26 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Dominance.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/SetVector.h" #include "llvm/ADT/SmallPtrSet.h" - #include +#include #include "DeferredCommunicationDeadlock.hpp" #include "DeferredProjectionAnalysis.hpp" #include "DeferredCommunicationRealization.hpp" +#include "ScheduledComputePlan.hpp" #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" @@ -37,6 +42,8 @@ struct ProducedValue { int64_t laneStart = 0; int64_t laneCount = 1; unsigned scheduledLane = 0; + int64_t publishedSlotStart = 0; + int64_t publishedSlotCount = 1; Value payload; Value published; }; @@ -95,6 +102,69 @@ struct NormalizedDeferredExchange { unsigned externalTransferCount = 0; }; +struct FragmentTransferRef { + NormalizedDeferredExchange *exchange = nullptr; + DeferredSpecialization *specialization = nullptr; + FragmentTransfer *transfer = nullptr; +}; + +struct LogicalTransferMetadata { + uint64_t channelId = 0; + uint64_t parentExchangeId = 0; + unsigned parentTransferCount = 0; + int64_t sourceCore = -1; + int64_t targetCore = -1; +}; + +static LogicalTransferMetadata getLogicalTransferMetadata( + const FragmentTransferRef &ref) { + return {ref.transfer->channelId, ref.exchange->exchangeId, + ref.exchange->externalTransferCount, ref.transfer->sourceCore, + ref.transfer->targetCore}; +} + +struct SendEmissionSignature { + ScheduledInfo *scheduled = nullptr; + Value payload; + Type fragmentType; + bool hasGraphLane = false; + bool sourceIsBatch = false; +}; + +static SendEmissionSignature getSendEmissionSignature( + const FragmentTransferRef &ref) { + const FragmentRequirement &requirement = + ref.specialization->requirements[ref.transfer->requirementIndex]; + return {requirement.producer->scheduled, + requirement.producer->payload, + requirement.publicationFragmentType, + requirement.graphLane.has_value(), + requirement.producer->scheduled->isBatch()}; +} + +static bool hasRepresentableSendOffset(const FragmentTransferRef &ref) { + const FragmentRequirement &requirement = + ref.specialization->requirements[ref.transfer->requirementIndex]; + if (!requirement.graphLane) + return true; + int64_t offset = + *requirement.graphLane - requirement.producer->laneStart; + return offset >= 0 && offset < requirement.producer->laneCount; +} + +static bool haveCompatibleSendEmission( + const FragmentTransferRef &lhs, const FragmentTransferRef &rhs, + unsigned lhsInsertionStep, unsigned rhsInsertionStep) { + SendEmissionSignature left = getSendEmissionSignature(lhs); + SendEmissionSignature right = getSendEmissionSignature(rhs); + return lhsInsertionStep == rhsInsertionStep + && left.scheduled == right.scheduled && left.payload == right.payload + && left.fragmentType == right.fragmentType + && left.hasGraphLane == right.hasGraphLane + && left.sourceIsBatch == right.sourceIsBatch + && hasRepresentableSendOffset(lhs) && hasRepresentableSendOffset(rhs); +} + struct RealizationPlan { SmallVector scheduled; SmallVector> producedStorage; @@ -102,9 +172,35 @@ struct RealizationPlan { GraphBatchPublicationCache publicationCache; SmallVector exchanges; SmallVector external; + DenseMap transferByChannel; SmallVector stepCounts; }; +enum class BoundaryActionKind { Send, Receive }; + +struct BoundaryAction { + BoundaryActionKind kind = BoundaryActionKind::Send; + FragmentTransferRef ref; + unsigned globalOrder = 0; +}; + +struct ScheduledBoundaryPlan { + ScheduledInfo *scheduled = nullptr; + unsigned insertionStep = 0; + SmallVector, 0> actionsByLane; + SmallVector targetExchanges; +}; + +struct BatchBoundaryClass { + SmallVector lanes; +}; + +using ReceiveCache = DenseMap; +using AssemblyCache = DenseMap; +using IndexConstantCache = DenseMap; +using DenseConstantCache = DenseMap; +using DeferredEraseSet = llvm::SetVector; + static FailureOr> getI64Array(Operation* op, StringRef name) { auto attr = op->getAttrOfType(name); if (!attr) @@ -151,7 +247,12 @@ getScalarStepResult(SpatScheduledCompute scheduled, Block& block, unsigned resul return yield.getOutputs()[yield.getOutputs().size() - resultCount + resultIndex]; } -static FailureOr +struct BatchStepResult { + Value payload; + tensor::ParallelInsertSliceOp publication; +}; + +static FailureOr getBatchStepResult(SpatScheduledComputeBatch scheduled, Block& block, unsigned globalResultIndex) { auto inParallel = dyn_cast(block.getTerminator()); if (!inParallel) @@ -163,7 +264,7 @@ getBatchStepResult(SpatScheduledComputeBatch scheduled, Block& block, unsigned g continue; auto destination = dyn_cast(insert.getDest()); if (destination && destination.getOwner() == &block && destination.getArgNumber() == resultBase + globalResultIndex) - return insert.getSource(); + return BatchStepResult {insert.getSource(), insert}; } return scheduled.emitOpError("phase 2 cannot find the batched result insertion for a scheduled step"); } @@ -248,6 +349,14 @@ static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan cast(info.op), *info.blocks[step], result, info.resultCounts[step]); if (failed(payload)) return failure(); + int64_t publishedSlotCount = 1; + if (laneCounts[step] > 1) { + auto publishedType = dyn_cast(info.op->getResult(globalResult).getType()); + if (!publishedType || !publishedType.hasStaticShape() || publishedType.getRank() == 0 + || publishedType.getDimSize(0) != laneCounts[step]) + return info.op->emitOpError("phase 2 scalar scheduled publication geometry is inconsistent"); + publishedSlotCount = laneCounts[step]; + } auto produced = std::make_unique(ProducedValue { &info, step, @@ -257,6 +366,8 @@ static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan laneStarts[step], std::max(laneCounts[step], 1), 0, + 0, + publishedSlotCount, *payload, info.op->getResult(globalResult), }); @@ -265,11 +376,26 @@ static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan plan.producedStorage.push_back(std::move(produced)); continue; } - auto payload = getBatchStepResult(cast(info.op), *info.blocks[step], globalResult); - if (failed(payload)) + auto stepResult = getBatchStepResult(cast(info.op), *info.blocks[step], globalResult); + if (failed(stepResult)) return failure(); + auto publishedType = dyn_cast(info.op->getResult(globalResult).getType()); + if (!publishedType || !publishedType.hasStaticShape() || publishedType.getRank() == 0) + return info.op->emitOpError("phase 2 batched scheduled publication must have a static leading slot dimension"); for (unsigned lane = 0; lane < info.cores.size(); ++lane) { size_t laneIndex = step * info.cores.size() + lane; + StaticIndexEnvironment environment; + environment.bindings[info.blocks[step]->getArgument(0)] = lane; + auto slotStart = evaluateDeferredIndex( + stepResult->publication.getMixedOffsets().front(), environment); + auto slotCount = evaluateDeferredIndex( + stepResult->publication.getMixedSizes().front(), environment); + auto slotStride = evaluateDeferredIndex( + stepResult->publication.getMixedStrides().front(), environment); + if (failed(slotStart) || failed(slotCount) || failed(slotStride) + || *slotStart < 0 || *slotCount <= 0 || *slotStride != 1 + || *slotStart > publishedType.getDimSize(0) - *slotCount) + return info.op->emitOpError("phase 2 batched scheduled publication geometry is invalid"); auto produced = std::make_unique(ProducedValue { &info, step, @@ -279,7 +405,9 @@ static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan laneStarts[laneIndex], std::max(laneCounts[laneIndex], 1), lane, - *payload, + *slotStart, + *slotCount, + stepResult->payload, info.op->getResult(globalResult), }); info.produced.push_back(produced.get()); @@ -359,9 +487,16 @@ static LogicalResult retargetBlueprintToScheduledPublications( it = std::prev(publications.end()); } scheduledOperandIndices.push_back(std::distance(publications.begin(), it)); - scheduledSourceSlots.push_back((*producer)->scheduled->isBatch() - ? (*producer)->scheduledLane - : graphLane - (*producer)->laneStart); + int64_t localGraphLane = graphLane - (*producer)->laneStart; + if (localGraphLane < 0 || localGraphLane >= (*producer)->laneCount + || localGraphLane >= (*producer)->publishedSlotCount) + return blueprint.emitOpError("phase 2 Blueprint graph lane is outside its scheduled publication"); + int64_t scheduledSourceSlot = (*producer)->publishedSlotStart + localGraphLane; + auto publishedType = dyn_cast((*producer)->published.getType()); + if (!publishedType || !publishedType.hasStaticShape() || publishedType.getRank() == 0 + || scheduledSourceSlot < 0 || scheduledSourceSlot >= publishedType.getDimSize(0)) + return blueprint.emitOpError("phase 2 Blueprint scheduled publication slot is out of range"); + scheduledSourceSlots.push_back(scheduledSourceSlot); Operation *scheduled = (*producer)->published.getDefiningOp(); if (blueprint->getBlock() == scheduled->getBlock() && blueprint->isBeforeInBlock(scheduled)) @@ -448,6 +583,24 @@ static LogicalResult buildDeferredSpecialization(RealizationPlan &plan, if (failed(append(physicalSlot, graphLane, (*publication)->publicationFragmentType, position))) return failure(); } } + if (specialization.program.insertAssembly) + for (DeferredInsertAssemblyEntry &entry : + specialization.program.insertAssembly->entries) { + std::optional requirementIndex; + for (auto [index, requirement] : + llvm::enumerate(specialization.requirements)) + if (requirement.leafIndex == entry.leafIndex + && requirement.selectedPosition == 0) { + if (requirementIndex) + return exchange.deferred.emitOpError( + "phase 2 insert assembly requirement is ambiguous"); + requirementIndex = index; + } + if (!requirementIndex) + return exchange.deferred.emitOpError( + "phase 2 insert assembly requirement is missing"); + entry.requirementIndex = *requirementIndex; + } exchange.specializations.push_back(std::move(specialization)); return success(); } @@ -488,97 +641,479 @@ static LogicalResult normalizeDeferred(func::FuncOp funcOp, RealizationPlan& pla 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); +struct GroupDependency { + unsigned stream = 0; + unsigned requiredCompletedStep = 0; +}; - 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]) +struct StaticEmissionCost { + unsigned logicalTransferCount = 0; + unsigned sendRunCount = 0; + unsigned transfersCoveredBySendRuns = 0; + unsigned largestSendRun = 0; + unsigned assemblyCoveredReceiveCount = 0; + unsigned estimatedScalarReceiveCount = 0; + unsigned estimatedBatchBehaviorClassCount = 0; +}; + +struct ParentTransferGroup { + uint64_t parentExchangeId = 0; + SmallVector originalTransferIndices; + SmallVector orderedTransferIndices; + SmallVector dependencies; + unsigned unsatisfiedDependencies = 0; + bool ready = false; + bool scheduled = false; + StaticEmissionCost cost; + std::tuple deterministicPriority; + std::optional firstSendSignature; +}; + +struct SchedulerStreamState { + unsigned completedStep = 0; + SmallVector pendingIncomingAtConsumerStep; + SmallVector> groupsUnlockedAtStep; + bool queued = false; +}; + +static bool haveSameSendEmissionSignature( + const SendEmissionSignature &lhs, const SendEmissionSignature &rhs) { + return lhs.scheduled == rhs.scheduled && lhs.payload == rhs.payload + && lhs.fragmentType == rhs.fragmentType + && lhs.hasGraphLane == rhs.hasGraphLane + && lhs.sourceIsBatch == rhs.sourceIsBatch; +} + +static size_t hashSendEmissionSignature( + const SendEmissionSignature &signature) { + return static_cast(llvm::hash_combine( + signature.scheduled, signature.payload.getAsOpaquePointer(), + signature.fragmentType.getAsOpaquePointer(), signature.hasGraphLane, + signature.sourceIsBatch)); +} + +static bool hasBetterStaticPriority(const ParentTransferGroup &lhs, + const ParentTransferGroup &rhs) { + auto addedOps = [](const ParentTransferGroup &group) { + return group.cost.sendRunCount + group.cost.estimatedScalarReceiveCount; + }; + auto lhsScore = std::tuple( + addedOps(lhs), lhs.cost.estimatedBatchBehaviorClassCount, + std::numeric_limits::max() + - lhs.cost.assemblyCoveredReceiveCount, + std::numeric_limits::max() + - lhs.cost.transfersCoveredBySendRuns, + std::numeric_limits::max() - lhs.cost.largestSendRun, + lhs.cost.sendRunCount, lhs.deterministicPriority); + auto rhsScore = std::tuple( + addedOps(rhs), rhs.cost.estimatedBatchBehaviorClassCount, + std::numeric_limits::max() + - rhs.cost.assemblyCoveredReceiveCount, + std::numeric_limits::max() + - rhs.cost.transfersCoveredBySendRuns, + std::numeric_limits::max() - rhs.cost.largestSendRun, + rhs.cost.sendRunCount, rhs.deterministicPriority); + return lhsScore < rhsScore; +} + +static void buildGroupTransferOrder(RealizationPlan &plan, + ParentTransferGroup &group) { + SmallVector semanticOrder; + llvm::SmallDenseSet orderedIndices; + DenseMap indexByChannel; + for (unsigned index : group.originalTransferIndices) + indexByChannel[plan.external[index].exchangeId] = index; + FragmentTransferRef groupRef = plan.transferByChannel.lookup( + plan.external[group.originalTransferIndices.front()].exchangeId); + for (DeferredSpecialization &specialization : + groupRef.exchange->specializations) { + if (!specialization.program.insertAssembly) + continue; + for (const DeferredInsertAssemblyEntry &entry : + specialization.program.insertAssembly->entries) { + if (entry.requirementIndex >= specialization.transfers.size()) 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()]; + FragmentTransfer &transfer = + specialization.transfers[entry.requirementIndex]; + if (transfer.local) + continue; + auto it = indexByChannel.find(transfer.channelId); + if (it != indexByChannel.end() + && orderedIndices.insert(it->second).second) + semanticOrder.push_back(it->second); + } + } + for (unsigned index : group.originalTransferIndices) + if (orderedIndices.insert(index).second) + semanticOrder.push_back(index); + + SmallVector> classes; + DenseMap> classesBySignature; + for (unsigned index : semanticOrder) { + FragmentTransferRef ref = + plan.transferByChannel.lookup(plan.external[index].exchangeId); + if (!hasRepresentableSendOffset(ref)) { + classes.push_back({index}); + continue; + } + SendEmissionSignature signature = getSendEmissionSignature(ref); + size_t signatureHash = hashSendEmissionSignature(signature); + SmallVector &candidates = classesBySignature[signatureHash]; + auto candidateIt = llvm::find_if(candidates, [&](unsigned classIndex) { + ArrayRef candidate = classes[classIndex]; + FragmentTransferRef first = plan.transferByChannel.lookup( + plan.external[candidate.front()].exchangeId); + return haveSameSendEmissionSignature( + getSendEmissionSignature(first), signature); + }); + if (candidateIt == candidates.end()) { + candidates.push_back(classes.size()); + classes.push_back({index}); + } else { + classes[*candidateIt].push_back(index); + } + } + llvm::stable_sort(classes, [&](ArrayRef lhs, + ArrayRef rhs) { + if (lhs.size() != rhs.size()) + return lhs.size() > rhs.size(); + const PlannedCommunicationTransfer &left = plan.external[lhs.front()]; + const PlannedCommunicationTransfer &right = plan.external[rhs.front()]; + return std::tie(left.sourceStream, left.targetStream, left.exchangeId) + < std::tie(right.sourceStream, right.targetStream, right.exchangeId); + }); + group.cost.logicalTransferCount = group.originalTransferIndices.size(); + group.cost.sendRunCount = classes.size(); + for (ArrayRef compactClass : classes) { + llvm::append_range(group.orderedTransferIndices, compactClass); + group.cost.largestSendRun = + std::max(group.cost.largestSendRun, compactClass.size()); + if (compactClass.size() >= 2) + group.cost.transfersCoveredBySendRuns += compactClass.size(); + } + if (!group.orderedTransferIndices.empty()) { + FragmentTransferRef first = plan.transferByChannel.lookup( + plan.external[group.orderedTransferIndices.front()].exchangeId); + group.firstSendSignature = getSendEmissionSignature(first); + } +} + +static void buildGroupStaticCost(RealizationPlan &plan, + ParentTransferGroup &group) { + if (group.originalTransferIndices.empty()) + return; + FragmentTransferRef ref = plan.transferByChannel.lookup( + plan.external[group.originalTransferIndices.front()].exchangeId); + unsigned assemblyCoverage = 0; + for (DeferredSpecialization &specialization : ref.exchange->specializations) + if (specialization.program.insertAssembly) { + unsigned externalEntries = llvm::count_if( + specialization.transfers, + [](const FragmentTransfer &transfer) { return !transfer.local; }); + assemblyCoverage = std::max(assemblyCoverage, externalEntries); + } + group.cost.assemblyCoveredReceiveCount = assemblyCoverage; + group.cost.estimatedScalarReceiveCount = + group.cost.logicalTransferCount - + std::min(group.cost.logicalTransferCount, assemblyCoverage); + group.cost.estimatedBatchBehaviorClassCount = + ref.exchange->target->isBatch() ? ref.exchange->specializations.size() : 1; +} + +static SmallVector buildParentTransferGroups( + RealizationPlan &plan) { + llvm::MapVector groupByParent; + SmallVector groups; + for (auto [index, transfer] : llvm::enumerate(plan.external)) { + auto inserted = groupByParent.insert({transfer.parentExchangeId, + groups.size()}); + if (inserted.second) { + ParentTransferGroup group; + group.parentExchangeId = transfer.parentExchangeId; + groups.push_back(std::move(group)); + } + groups[inserted.first->second].originalTransferIndices.push_back(index); + } + + for (ParentTransferGroup &group : groups) { + DenseMap thresholdByStream; + unsigned minConsumer = std::numeric_limits::max(); + unsigned minSource = std::numeric_limits::max(); + unsigned minTarget = std::numeric_limits::max(); + for (unsigned index : group.originalTransferIndices) { + const PlannedCommunicationTransfer &transfer = plan.external[index]; + thresholdByStream[transfer.sourceStream] = std::max( + thresholdByStream.lookup(transfer.sourceStream), + transfer.producerStep + 1); + minConsumer = std::min(minConsumer, transfer.consumerStep); + minSource = std::min(minSource, transfer.sourceStream); + minTarget = std::min(minTarget, transfer.targetStream); + } + FragmentTransferRef ref = plan.transferByChannel.lookup( + plan.external[group.originalTransferIndices.front()].exchangeId); + for (DeferredSpecialization &specialization : ref.exchange->specializations) + for (FragmentTransfer &transfer : specialization.transfers) + if (transfer.local) { + const FragmentRequirement &requirement = + specialization.requirements[transfer.requirementIndex]; + thresholdByStream[transfer.targetStream] = std::max( + thresholdByStream.lookup(transfer.targetStream), + requirement.producer->step + 1); + } + SmallVector dependencyStreams; + dependencyStreams.reserve(thresholdByStream.size()); + for (auto entry : thresholdByStream) + dependencyStreams.push_back(entry.first); + llvm::sort(dependencyStreams); + for (unsigned stream : dependencyStreams) + group.dependencies.push_back( + {stream, thresholdByStream.lookup(stream)}); + group.unsatisfiedDependencies = llvm::count_if( + group.dependencies, [](const GroupDependency &dependency) { + return dependency.requiredCompletedStep != 0; }); - if (inputsReady) { - ++completedSteps[stream]; - progressed = true; + group.deterministicPriority = + {minConsumer, minSource, minTarget, group.parentExchangeId}; + buildGroupTransferOrder(plan, group); + buildGroupStaticCost(plan, group); + } + return groups; +} + +static LogicalResult scheduleCommunication(func::FuncOp funcOp, + RealizationPlan &plan) { + SmallVector groups = + buildParentTransferGroups(plan); + SmallVector streams(plan.stepCounts.size()); + for (auto [stream, state] : llvm::enumerate(streams)) { + state.pendingIncomingAtConsumerStep.resize(plan.stepCounts[stream]); + state.groupsUnlockedAtStep.resize(plan.stepCounts[stream] + 1); + } + for (const PlannedCommunicationTransfer &transfer : plan.external) + ++streams[transfer.targetStream] + .pendingIncomingAtConsumerStep[transfer.consumerStep]; + for (auto [groupIndex, group] : llvm::enumerate(groups)) + for (const GroupDependency &dependency : group.dependencies) + streams[dependency.stream] + .groupsUnlockedAtStep[dependency.requiredCompletedStep] + .push_back(groupIndex); + + std::function heapCompare = + [&](unsigned lhs, unsigned rhs) { + return hasBetterStaticPriority(groups[rhs], groups[lhs]); + }; + std::priority_queue, + std::function> + readyByStaticPriority(heapCompare); + using ReadyHeap = std::priority_queue< + unsigned, std::vector, + std::function>; + DenseMap> + readyByFirstSendSignature; + auto addReady = [&](unsigned groupIndex) { + ParentTransferGroup &group = groups[groupIndex]; + if (group.ready || group.scheduled) + return; + group.ready = true; + readyByStaticPriority.push(groupIndex); + if (group.firstSendSignature) { + std::unique_ptr &bucket = readyByFirstSendSignature[ + hashSendEmissionSignature(*group.firstSendSignature)]; + if (!bucket) + bucket = std::make_unique(heapCompare); + bucket->push(groupIndex); + } + }; + for (auto [groupIndex, group] : llvm::enumerate(groups)) + if (group.unsatisfiedDependencies == 0) + addReady(groupIndex); + + std::queue advanceable; + auto enqueueIfAdvanceable = [&](unsigned stream) { + SchedulerStreamState &state = streams[stream]; + if (!state.queued && state.completedStep < plan.stepCounts[stream] + && state.pendingIncomingAtConsumerStep[state.completedStep] == 0) { + state.queued = true; + advanceable.push(stream); + } + }; + for (unsigned stream = 0; stream < streams.size(); ++stream) + enqueueIfAdvanceable(stream); + auto advanceStreams = [&] { + bool advanced = false; + while (!advanceable.empty()) { + unsigned stream = advanceable.front(); + advanceable.pop(); + SchedulerStreamState &state = streams[stream]; + state.queued = false; + if (state.completedStep == plan.stepCounts[stream] + || state.pendingIncomingAtConsumerStep[state.completedStep] != 0) + continue; + ++state.completedStep; + advanced = true; + for (unsigned groupIndex : + state.groupsUnlockedAtStep[state.completedStep]) { + ParentTransferGroup &group = groups[groupIndex]; + assert(group.unsatisfiedDependencies != 0); + if (--group.unsatisfiedDependencies == 0) + addReady(groupIndex); } + enqueueIfAdvanceable(stream); } + return advanced; + }; - SmallVector readyGroups; - for (auto &entry : groups) { - if (completedTransfers[entry.second.front()]) - continue; - if (llvm::all_of(entry.second, [&](unsigned index) { - const PlannedCommunicationTransfer &transfer = plan.external[index]; - return completedSteps[transfer.sourceStream] > transfer.producerStep; - })) - readyGroups.push_back(entry.first); + SmallVector scheduled; + scheduled.reserve(plan.external.size()); + unsigned scheduledGroups = 0; + while (scheduledGroups != groups.size()) { + bool progressed = advanceStreams(); + std::optional chosen; + unsigned chosenTailExtension = 0; + if (!scheduled.empty()) { + const PlannedCommunicationTransfer &tail = scheduled.back(); + FragmentTransferRef tailRef = + plan.transferByChannel.lookup(tail.exchangeId); + size_t signatureHash = + hashSendEmissionSignature(getSendEmissionSignature(tailRef)); + auto bucketIt = readyByFirstSendSignature.find(signatureHash); + SmallVector inspected; + ReadyHeap *bucket = bucketIt == readyByFirstSendSignature.end() + ? nullptr : bucketIt->second.get(); + while (bucket && !bucket->empty() && inspected.size() != 32) { + unsigned groupIndex = bucket->top(); + bucket->pop(); + ParentTransferGroup &group = groups[groupIndex]; + if (!group.ready || group.scheduled) + continue; + inspected.push_back(groupIndex); + if (!group.firstSendSignature + || !haveSameSendEmissionSignature( + getSendEmissionSignature(tailRef), + *group.firstSendSignature)) + continue; + unsigned extension = 0; + for (unsigned transferIndex : group.orderedTransferIndices) { + const PlannedCommunicationTransfer &transfer = + plan.external[transferIndex]; + FragmentTransferRef ref = + plan.transferByChannel.lookup(transfer.exchangeId); + if (!haveCompatibleSendEmission( + tailRef, ref, tail.sourceInsertionStep, + streams[transfer.sourceStream].completedStep)) + break; + ++extension; + } + if (extension == 0) + continue; + if (!chosen || extension > chosenTailExtension + || (extension == chosenTailExtension + && hasBetterStaticPriority(group, groups[*chosen]))) { + chosen = groupIndex; + chosenTailExtension = extension; + } + } + for (unsigned groupIndex : inspected) + if (!chosen || groupIndex != *chosen) + bucket->push(groupIndex); } - 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); - }); - SmallVector ordered(groups[readyGroups.front()].begin(), - groups[readyGroups.front()].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; + while (!chosen && !readyByStaticPriority.empty()) { + unsigned candidate = readyByStaticPriority.top(); + readyByStaticPriority.pop(); + if (groups[candidate].ready && !groups[candidate].scheduled) + chosen = candidate; + } + if (chosen) { + ParentTransferGroup &group = groups[*chosen]; + group.ready = false; + group.scheduled = true; + ++scheduledGroups; + for (unsigned index : group.orderedTransferIndices) { PlannedCommunicationTransfer transfer = plan.external[index]; - transfer.sourceInsertionStep = completedSteps[transfer.sourceStream]; - transfer.targetInsertionStep = completedSteps[transfer.targetStream]; + transfer.sourceInsertionStep = + streams[transfer.sourceStream].completedStep; + transfer.targetInsertionStep = + streams[transfer.targetStream].completedStep; scheduled.push_back(transfer); + unsigned &pending = streams[transfer.targetStream] + .pendingIncomingAtConsumerStep[transfer.consumerStep]; + assert(pending != 0); + --pending; + if (pending == 0 + && streams[transfer.targetStream].completedStep + == transfer.consumerStep) + enqueueIfAdvanceable(transfer.targetStream); } progressed = true; } - if (!progressed) { - InFlightDiagnostic diagnostic = funcOp.emitOpError( - "global communication scheduler made no progress before IR mutation"); - unsigned reported = 0; - for (auto [index, transfer] : llvm::enumerate(plan.external)) { - if (completedTransfers[index] || reported == 8) - continue; - diagnostic << (reported++ == 0 ? "; blocked " : ", ") - << "channel " << transfer.exchangeId << " source stream " - << transfer.sourceStream << " at step " - << completedSteps[transfer.sourceStream] << " after producer " - << transfer.producerStep << " target stream " - << transfer.targetStream << " at step " - << completedSteps[transfer.targetStream] << " before consumer " - << transfer.consumerStep; - } - return failure(); + if (progressed) + continue; + + InFlightDiagnostic diagnostic = funcOp.emitOpError( + "global communication scheduler made no progress before IR mutation"); + unsigned reported = 0; + for (const ParentTransferGroup &group : groups) { + if (group.scheduled || reported == 8) + continue; + diagnostic << (reported++ == 0 ? "; blocked " : ", ") + << "parent " << group.parentExchangeId; + auto dependency = llvm::find_if( + group.dependencies, [&](const GroupDependency &candidate) { + return streams[candidate.stream].completedStep + < candidate.requiredCompletedStep; + }); + if (dependency != group.dependencies.end()) + diagnostic << " stream " << dependency->stream << " requires " + << dependency->requiredCompletedStep << " has " + << streams[dependency->stream].completedStep; + const PlannedCommunicationTransfer &first = + plan.external[group.originalTransferIndices.front()]; + diagnostic << " consumer stream " << first.targetStream << " step " + << first.consumerStep; } + return failure(); } plan.external = std::move(scheduled); return success(); } +static FailureOr getSpecializationTargetInsertionStep( + NormalizedDeferredExchange &exchange, DeferredSpecialization &specialization) { + std::optional targetStep; + for (FragmentTransfer &transfer : specialization.transfers) { + if (transfer.local) + continue; + if (targetStep && *targetStep != transfer.targetInsertionStep) + return exchange.deferred.emitOpError() + << "exchange " << exchange.exchangeId << " scheduled lane " + << specialization.targetScheduledLane.value_or(0) << " channel " + << transfer.channelId << " has target insertion step " + << transfer.targetInsertionStep << " instead of " << *targetStep, + failure(); + targetStep = transfer.targetInsertionStep; + } + return targetStep.value_or(exchange.consumerStep); +} + +static LogicalResult verifyExchangeTargetInsertionSteps(RealizationPlan &plan) { + for (NormalizedDeferredExchange &exchange : plan.exchanges) { + std::optional sharedStep; + for (DeferredSpecialization &specialization : exchange.specializations) { + auto step = getSpecializationTargetInsertionStep(exchange, specialization); + if (failed(step)) + return failure(); + if (exchange.target->isBatch() && sharedStep && *sharedStep != *step) + return exchange.deferred.emitOpError() + << "exchange " << exchange.exchangeId << " scheduled lane " + << specialization.targetScheduledLane.value_or(0) + << " has target insertion step " << *step << " instead of shared step " + << *sharedStep; + sharedStep = *step; + } + } + return success(); +} + static LogicalResult buildAndVerifyPlan(func::FuncOp funcOp, RealizationPlan& plan) { uint64_t nextChannel = 0; for (NormalizedDeferredExchange &exchange : plan.exchanges) @@ -588,6 +1123,9 @@ static LogicalResult buildAndVerifyPlan(func::FuncOp funcOp, RealizationPlan& pl continue; transfer.channelId = nextChannel++; ++exchange.externalTransferCount; + if (!plan.transferByChannel.try_emplace( + transfer.channelId, FragmentTransferRef {&exchange, &specialization, &transfer}).second) + return exchange.deferred.emitOpError("phase 2 assigned a duplicate communication channel"); const FragmentRequirement &requirement = specialization.requirements[transfer.requirementIndex]; plan.external.push_back({transfer.channelId, exchange.exchangeId, transfer.sourceStream, transfer.targetStream, @@ -595,28 +1133,134 @@ static LogicalResult buildAndVerifyPlan(func::FuncOp funcOp, RealizationPlan& pl } if (failed(scheduleCommunication(funcOp, plan))) return failure(); - DenseMap plannedByChannel; - DenseMap orderByChannel; for (auto [order, transfer] : llvm::enumerate(plan.external)) { - plannedByChannel[transfer.exchangeId] = &transfer; - orderByChannel[transfer.exchangeId] = order; + FragmentTransferRef ref = plan.transferByChannel.lookup(transfer.exchangeId); + if (!ref.exchange || !ref.specialization || !ref.transfer) + return funcOp.emitOpError("phase 2 cannot resolve a planned communication channel"); + ref.transfer->globalOrder = order; + ref.transfer->sourceInsertionStep = transfer.sourceInsertionStep; + ref.transfer->targetInsertionStep = transfer.targetInsertionStep; } - auto applyInsertionSteps = [&](FragmentTransfer& transfer) { - if (transfer.local) - return; - PlannedCommunicationTransfer* planned = plannedByChannel.lookup(transfer.channelId); - assert(planned && "every external fragment must have one planned transfer"); - transfer.globalOrder = orderByChannel.lookup(transfer.channelId); - transfer.sourceInsertionStep = planned->sourceInsertionStep; - transfer.targetInsertionStep = planned->targetInsertionStep; - }; - for (NormalizedDeferredExchange &exchange : plan.exchanges) - for (DeferredSpecialization &specialization : exchange.specializations) - for (FragmentTransfer &transfer : specialization.transfers) - applyInsertionSteps(transfer); + if (failed(verifyExchangeTargetInsertionSteps(plan))) + return failure(); return verifyPlannedCommunicationDeadlockFree(funcOp, plan.stepCounts.size(), plan.stepCounts, plan.external); } +static FailureOr> +buildScheduledBoundaryPlans(RealizationPlan &plan) { + SmallVector boundaries; + DenseMap, unsigned> boundaryIndex; + SmallVector> actionCounts; + auto getBoundary = [&](ScheduledInfo *scheduled, + unsigned insertionStep) -> ScheduledBoundaryPlan & { + auto key = std::make_pair(scheduled, insertionStep); + if (auto it = boundaryIndex.find(key); it != boundaryIndex.end()) + return boundaries[it->second]; + ScheduledBoundaryPlan boundary; + boundary.scheduled = scheduled; + boundary.insertionStep = insertionStep; + boundary.actionsByLane.resize(scheduled->cores.size()); + boundaryIndex[key] = boundaries.size(); + boundaries.push_back(std::move(boundary)); + actionCounts.emplace_back(scheduled->cores.size()); + return boundaries.back(); + }; + + for (const PlannedCommunicationTransfer &transfer : plan.external) { + FragmentTransferRef ref = + plan.transferByChannel.lookup(transfer.exchangeId); + if (!ref.exchange || !ref.specialization || !ref.transfer) + return failure(); + FragmentRequirement &requirement = + ref.specialization->requirements[ref.transfer->requirementIndex]; + getBoundary(requirement.producer->scheduled, + ref.transfer->sourceInsertionStep); + getBoundary(ref.exchange->target, ref.transfer->targetInsertionStep); + } + for (NormalizedDeferredExchange &exchange : plan.exchanges) { + unsigned targetStep = exchange.consumerStep; + if (exchange.externalTransferCount != 0) { + auto step = getSpecializationTargetInsertionStep( + exchange, exchange.specializations.front()); + if (failed(step)) + return failure(); + targetStep = *step; + } + getBoundary(exchange.target, targetStep); + } + + for (const PlannedCommunicationTransfer &transfer : plan.external) { + FragmentTransferRef ref = + plan.transferByChannel.lookup(transfer.exchangeId); + FragmentRequirement &requirement = + ref.specialization->requirements[ref.transfer->requirementIndex]; + unsigned sourceBoundary = boundaryIndex.lookup( + {requirement.producer->scheduled, + ref.transfer->sourceInsertionStep}); + unsigned targetBoundary = boundaryIndex.lookup( + {ref.exchange->target, ref.transfer->targetInsertionStep}); + ++actionCounts[sourceBoundary][requirement.producer->scheduledLane]; + ++actionCounts[targetBoundary][ + ref.specialization->targetScheduledLane.value_or(0)]; + } + for (auto [boundary, counts] : llvm::zip(boundaries, actionCounts)) + for (auto [actions, count] : llvm::zip(boundary.actionsByLane, counts)) + actions.reserve(count); + + for (auto [globalOrder, transfer] : llvm::enumerate(plan.external)) { + FragmentTransferRef ref = + plan.transferByChannel.lookup(transfer.exchangeId); + if (!ref.exchange || !ref.specialization || !ref.transfer) + return failure(); + FragmentRequirement &requirement = + ref.specialization->requirements[ref.transfer->requirementIndex]; + getBoundary(requirement.producer->scheduled, + ref.transfer->sourceInsertionStep) + .actionsByLane[requirement.producer->scheduledLane] + .push_back({BoundaryActionKind::Send, ref, + static_cast(globalOrder)}); + unsigned targetLane = + ref.specialization->targetScheduledLane.value_or(0); + getBoundary(ref.exchange->target, ref.transfer->targetInsertionStep) + .actionsByLane[targetLane] + .push_back({BoundaryActionKind::Receive, ref, + static_cast(globalOrder)}); + } + + for (NormalizedDeferredExchange &exchange : plan.exchanges) { + unsigned targetStep = exchange.consumerStep; + if (exchange.externalTransferCount != 0) { + auto step = getSpecializationTargetInsertionStep( + exchange, exchange.specializations.front()); + if (failed(step)) + return failure(); + targetStep = *step; + } + getBoundary(exchange.target, targetStep).targetExchanges.push_back(&exchange); + } + + DenseMap scheduledOrder; + for (auto [index, scheduled] : llvm::enumerate(plan.scheduled)) + scheduledOrder[&scheduled] = index; + llvm::stable_sort(boundaries, [&](const ScheduledBoundaryPlan &lhs, + const ScheduledBoundaryPlan &rhs) { + return std::tie(scheduledOrder[lhs.scheduled], lhs.insertionStep) + < std::tie(scheduledOrder[rhs.scheduled], rhs.insertionStep); + }); + for (ScheduledBoundaryPlan &boundary : boundaries) { + for (ArrayRef actions : boundary.actionsByLane) { + if (actions.size() < 2) + continue; + for (auto pair : llvm::zip(actions.drop_back(), actions.drop_front())) + if (std::get<0>(pair).globalOrder >= std::get<1>(pair).globalOrder) + return boundary.scheduled->op->emitOpError( + "phase 2 boundary actions are not in committed order"), + failure(); + } + } + return boundaries; +} + static LogicalResult validateScalarLinearization(ScheduledInfo& info) { auto scheduled = cast(info.op); for (unsigned index = 1; index < info.blocks.size(); ++index) { @@ -743,10 +1387,31 @@ setTransferAttrs(Operation* op, uint64_t exchangeId, uint64_t channelId, int64_t static void replaceDeferred(SpatDeferredCommunicationOp deferred, Value replacement, - SmallVectorImpl& erase) { + DeferredEraseSet &erase) { deferred.getOutput().replaceAllUsesWith(replacement); - if (!llvm::is_contained(erase, deferred.getOperation())) - erase.push_back(deferred); + erase.insert(deferred); +} + +static Value getOrCreateCachedIndexConstant( + IndexConstantCache &cache, IRRewriter &rewriter, Operation *anchor, + int64_t value) { + if (Value constant = cache.lookup(value)) + return constant; + Value constant = createConstantAtHostBlockStart( + rewriter, anchor, rewriter.getIndexAttr(value)); + cache[value] = constant; + return constant; +} + +static Value getOrCreateCachedTypedConstant( + DenseConstantCache &cache, IRRewriter &rewriter, Operation *anchor, + TypedAttr value) { + if (Value constant = cache.lookup(value)) + return constant; + Value constant = createConstantAtHostBlockStart( + rewriter, anchor, value); + cache[value] = constant; + return constant; } static void setInsertionAtBoundary(IRRewriter& rewriter, ScheduledInfo& scheduled, unsigned step) { @@ -784,9 +1449,34 @@ static FailureOr materializeProducedFragment(const FragmentRequirement &r SmallVector unitShape {1}; llvm::append_range(unitShape, fragmentType.getShape()); auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType()); Value unit = tensor::ExtractSliceOp::create(rewriter, loc, unitType, payload, offsets, sizes, strides); - SmallVector reassociation {{0, 1}}; - for (int64_t dim = 1; dim < fragmentType.getRank(); ++dim) reassociation.push_back({dim + 1}); - return tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, unit, reassociation).getResult(); + return removeLeadingUnitTensorDimension(rewriter, loc, unit, fragmentType); +} + +static FailureOr materializeProducedFragment( + const FragmentRequirement &requirement, Value localOffset, + IRRewriter &rewriter, Location loc) { + Value payload = requirement.producer->payload; + if (payload.getType() == requirement.publicationFragmentType) + return payload; + auto payloadType = dyn_cast(payload.getType()); + auto fragmentType = dyn_cast(requirement.publicationFragmentType); + if (!requirement.graphLane || !payloadType || !fragmentType + || payloadType.getRank() != fragmentType.getRank() + 1 + || payloadType.getDimSize(0) != requirement.producer->laneCount + || !llvm::equal(payloadType.getShape().drop_front(), fragmentType.getShape())) + return failure(); + SmallVector offsets(payloadType.getRank(), rewriter.getIndexAttr(0)); + SmallVector sizes, strides(payloadType.getRank(), rewriter.getIndexAttr(1)); + offsets[0] = localOffset; + sizes.push_back(rewriter.getIndexAttr(1)); + for (int64_t dim : fragmentType.getShape()) + sizes.push_back(rewriter.getIndexAttr(dim)); + SmallVector unitShape {1}; + llvm::append_range(unitShape, fragmentType.getShape()); + auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType()); + Value unit = tensor::ExtractSliceOp::create( + rewriter, loc, unitType, payload, offsets, sizes, strides); + return removeLeadingUnitTensorDimension(rewriter, loc, unit, fragmentType); } static FailureOr applyInnerFragmentGeometry(Value fragment, const DeferredProjectionLeaf &leaf, @@ -823,42 +1513,65 @@ static FailureOr reconstructProjectionLeaf(const DeferredProjectionLeaf & } if (llvm::any_of(selected, [](Value value) { return !value; })) return failure(); if (leaf.kind == DeferredLeafKind::ScalarSource) return selected.front(); - if (selected.size() == 1 && selected.front().getType() == leaf.reconstructedType) return selected.front(); - Value result = tensor::EmptyOp::create(rewriter, loc, leaf.reconstructedType.getShape(), - leaf.reconstructedType.getElementType()); - auto outputType = leaf.reconstructedType; - for (auto [position, fragment] : llvm::enumerate(selected)) { - auto shaped = applyInnerFragmentGeometry(fragment, leaf, rewriter, loc); - if (failed(shaped)) return failure(); - auto fragmentType = cast((*shaped).getType()); - SmallVector unitShape {1}; llvm::append_range(unitShape, fragmentType.getShape()); - auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType()); - SmallVector reassociation {{0, 1}}; - for (int64_t dim = 1; dim < fragmentType.getRank(); ++dim) reassociation.push_back({dim + 1}); - Value unit = tensor::ExpandShapeOp::create(rewriter, loc, unitType, *shaped, reassociation); - SmallVector offsets(outputType.getRank(), rewriter.getIndexAttr(0)); - offsets[0] = rewriter.getIndexAttr(position); - SmallVector sizes, strides(outputType.getRank(), rewriter.getIndexAttr(1)); - sizes.push_back(rewriter.getIndexAttr(1)); - for (int64_t dim : fragmentType.getShape()) sizes.push_back(rewriter.getIndexAttr(dim)); - result = tensor::InsertSliceOp::create(rewriter, loc, unit, result, offsets, sizes, strides); + SmallVector shaped; + shaped.reserve(selected.size()); + for (Value fragment : selected) { + auto fragmentResult = applyInnerFragmentGeometry(fragment, leaf, rewriter, loc); + if (failed(fragmentResult)) + return failure(); + shaped.push_back(*fragmentResult); } - return result; + if (shaped.size() == 1 && shaped.front().getType() == leaf.reconstructedType) + return shaped.front(); + SmallVector expanded; + expanded.reserve(shaped.size()); + for (Value fragment : shaped) { + auto unit = addLeadingUnitTensorDimension(rewriter, loc, fragment); + if (failed(unit)) + return failure(); + expanded.push_back(*unit); + } + if (expanded.size() == 1 + && expanded.front().getType() == leaf.reconstructedType) + return expanded.front(); + Value result = tensor::ConcatOp::create( + rewriter, loc, leaf.reconstructedType, 0, expanded).getResult(); + return result.getType() == leaf.reconstructedType ? FailureOr(result) + : FailureOr(failure()); } static FailureOr materializeDeferredSpecializationResult(DeferredSpecialization &specialization, ArrayRef available, + IndexConstantCache &indexConstants, + DenseConstantCache &typedConstants, IRRewriter &rewriter, Location loc) { IRMapping mapping; + for (auto [value, staticValue] : specialization.program.staticValues) { + if (mapping.contains(value)) + continue; + Type type = value.getType(); + if (!type.isIndex() && !isa(type)) + return failure(); + Value constant = type.isIndex() + ? getOrCreateCachedIndexConstant( + indexConstants, rewriter, specialization.program.deferred, + staticValue) + : getOrCreateCachedTypedConstant( + typedConstants, rewriter, specialization.program.deferred, + rewriter.getIntegerAttr(type, staticValue)); + mapping.map(value, constant); + } for (auto [index, leaf] : llvm::enumerate(specialization.program.leaves)) { auto reconstructed = reconstructProjectionLeaf(leaf, index, specialization, available, rewriter, loc); if (failed(reconstructed)) return failure(); mapping.map(leaf.replacementRoot, *reconstructed); } - if (specialization.targetScheduledLane) { + if (specialization.targetScheduledLane && !mapping.contains(specialization.program.scheduledLane)) { mapping.map(specialization.program.scheduledLane, - arith::ConstantIndexOp::create(rewriter, loc, - *specialization.targetScheduledLane)); + getOrCreateCachedIndexConstant( + indexConstants, rewriter, + specialization.program.deferred, + *specialization.targetScheduledLane)); } for (Operation *op : specialization.program.residualOps) { Operation *copy = rewriter.clone(*op, mapping); @@ -872,16 +1585,19 @@ static FailureOr materializeDeferredSpecializationResult(DeferredSpeciali static LogicalResult emitFragmentSendHere( NormalizedDeferredExchange &exchange, DeferredSpecialization &specialization, FragmentTransfer &transfer, - IRRewriter &rewriter) { + IndexConstantCache &indexConstants, IRRewriter &rewriter) { Location loc = exchange.deferred.getLoc(); auto payload = materializeProducedFragment( specialization.requirements[transfer.requirementIndex], rewriter, loc); if (failed(payload)) return failure(); auto send = SpatChannelSendOp::create( rewriter, loc, - arith::ConstantIndexOp::create(rewriter, loc, transfer.channelId), - arith::ConstantIndexOp::create(rewriter, loc, transfer.sourceCore), - arith::ConstantIndexOp::create(rewriter, loc, transfer.targetCore), + getOrCreateCachedIndexConstant(indexConstants, rewriter, + exchange.deferred, transfer.channelId), + getOrCreateCachedIndexConstant(indexConstants, rewriter, + exchange.deferred, transfer.sourceCore), + getOrCreateCachedIndexConstant(indexConstants, rewriter, + exchange.deferred, transfer.targetCore), *payload); setTransferAttrs(send, transfer.channelId, transfer.channelId, transfer.sourceCore, transfer.targetCore); @@ -892,145 +1608,1785 @@ static LogicalResult emitFragmentSendHere( return success(); } -static LogicalResult emitFragmentSend(NormalizedDeferredExchange &exchange, DeferredSpecialization &specialization, - FragmentTransfer &transfer, IRRewriter &rewriter) { - if (transfer.local) return success(); - setInsertionAtBoundary(rewriter, *specialization.requirements[transfer.requirementIndex].producer->scheduled, - transfer.sourceInsertionStep); - Location loc = exchange.deferred.getLoc(); - ProducedValue *producer = specialization.requirements[transfer.requirementIndex].producer; - if (!producer->scheduled->isBatch()) - return emitFragmentSendHere(exchange, specialization, transfer, rewriter); - auto batch = cast(producer->scheduled->op); - Value lane = arith::ConstantIndexOp::create(rewriter, loc, producer->scheduledLane); - Value selected = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, *batch.getLaneArgument(), lane); - auto ifOp = scf::IfOp::create(rewriter, loc, selected, false); - rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); - return emitFragmentSendHere(exchange, specialization, transfer, rewriter); +static FailureOr emitFragmentReceiveHere( + const FragmentTransferRef &ref, ReceiveCache &receiveCache, + IndexConstantCache &indexConstants, IRRewriter &rewriter) { + FragmentTransfer &transfer = *ref.transfer; + FragmentRequirement &requirement = + ref.specialization->requirements[transfer.requirementIndex]; + Location loc = ref.exchange->deferred.getLoc(); + auto receive = SpatChannelReceiveOp::create( + rewriter, loc, requirement.publicationFragmentType, + getOrCreateCachedIndexConstant(indexConstants, rewriter, + ref.exchange->deferred, + transfer.channelId), + getOrCreateCachedIndexConstant(indexConstants, rewriter, + ref.exchange->deferred, + transfer.sourceCore), + getOrCreateCachedIndexConstant(indexConstants, rewriter, + ref.exchange->deferred, + transfer.targetCore)); + setTransferAttrs(receive, transfer.channelId, transfer.channelId, + transfer.sourceCore, transfer.targetCore); + receive->setAttr("raptor.parent_exchange_id", + rewriter.getI64IntegerAttr(ref.exchange->exchangeId)); + receive->setAttr("raptor.parent_transfer_count", + rewriter.getI64IntegerAttr( + ref.exchange->externalTransferCount)); + receiveCache[&transfer] = receive.getOutput(); + return receive.getOutput(); } static FailureOr receiveFragment(NormalizedDeferredExchange &exchange, DeferredSpecialization &specialization, - FragmentTransfer &transfer, IRRewriter &rewriter) { + FragmentTransfer &transfer, + ReceiveCache &receiveCache, + IRRewriter &rewriter) { if (transfer.local) return materializeProducedFragment(specialization.requirements[transfer.requirementIndex], rewriter, exchange.deferred.getLoc()); - Location loc = exchange.deferred.getLoc(); - auto receive = SpatChannelReceiveOp::create(rewriter, loc, - specialization.requirements[transfer.requirementIndex].publicationFragmentType, - arith::ConstantIndexOp::create(rewriter, loc, transfer.channelId), - arith::ConstantIndexOp::create(rewriter, loc, transfer.sourceCore), - arith::ConstantIndexOp::create(rewriter, loc, transfer.targetCore)); - setTransferAttrs(receive, transfer.channelId, transfer.channelId, transfer.sourceCore, transfer.targetCore); - receive->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.exchangeId)); - receive->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(exchange.externalTransferCount)); - return receive.getOutput(); + if (Value received = receiveCache.lookup(&transfer)) + return received; + return exchange.deferred.emitOpError( + "phase 2 boundary did not materialize a planned receive"), + failure(); } +static FailureOr> tryRealizeAnalyzedInsertAssembly( + NormalizedDeferredExchange &exchange, + DeferredSpecialization &specialization, IRRewriter &rewriter); + static FailureOr realizeSpecialization(NormalizedDeferredExchange &exchange, - DeferredSpecialization &specialization, IRRewriter &rewriter) { + DeferredSpecialization &specialization, + ReceiveCache &receiveCache, + AssemblyCache &assemblyCache, + IndexConstantCache &indexConstants, + DenseConstantCache &typedConstants, + IRRewriter &rewriter) { + if (Value assembled = assemblyCache.lookup(&specialization)) + return assembled; SmallVector available(specialization.requirements.size()); for (FragmentTransfer &transfer : specialization.transfers) if (transfer.local) { - auto value = receiveFragment(exchange, specialization, transfer, rewriter); + auto value = receiveFragment( + exchange, specialization, transfer, receiveCache, rewriter); if (failed(value)) return failure(); available[transfer.requirementIndex] = *value; } - struct OrderedTransfer { - DeferredSpecialization *specialization; - FragmentTransfer *transfer; - }; - SmallVector ordered; - for (DeferredSpecialization &owner : exchange.specializations) - for (FragmentTransfer &transfer : owner.transfers) - if (!transfer.local) - ordered.push_back({&owner, &transfer}); - llvm::sort(ordered, [](const OrderedTransfer &lhs, const OrderedTransfer &rhs) { - return lhs.transfer->globalOrder < rhs.transfer->globalOrder; + SmallVector ordered; + for (FragmentTransfer &transfer : specialization.transfers) + if (!transfer.local) + ordered.push_back(&transfer); + llvm::stable_sort(ordered, [](FragmentTransfer *lhs, FragmentTransfer *rhs) { + return lhs->globalOrder < rhs->globalOrder; }); - unsigned targetScheduledLane = specialization.targetScheduledLane.value_or(0); - for (OrderedTransfer event : ordered) { - FragmentRequirement &requirement = - event.specialization->requirements[event.transfer->requirementIndex]; - if (requirement.producer->scheduled == exchange.target && - requirement.producer->scheduledLane == targetScheduledLane && - failed(emitFragmentSendHere(exchange, *event.specialization, - *event.transfer, rewriter))) - return failure(); - if (event.specialization == &specialization) { - auto value = receiveFragment(exchange, specialization, *event.transfer, - rewriter); - if (failed(value)) return failure(); - available[event.transfer->requirementIndex] = *value; - } + for (FragmentTransfer *transfer : ordered) { + auto value = receiveFragment( + exchange, specialization, *transfer, receiveCache, rewriter); + if (failed(value)) return failure(); + available[transfer->requirementIndex] = *value; } - return materializeDeferredSpecializationResult(specialization, available, rewriter, exchange.deferred.getLoc()); + return materializeDeferredSpecializationResult( + specialization, available, indexConstants, typedConstants, rewriter, + exchange.deferred.getLoc()); } -static LogicalResult realizeFragmentExchange(NormalizedDeferredExchange &exchange, IRRewriter &rewriter, - SmallVectorImpl &erase) { - if (exchange.specializations.size() != (exchange.target->isBatch() ? exchange.target->cores.size() : 1)) - return exchange.deferred.emitOpError("phase 2 specialization count does not match target lanes"); - if (!exchange.target->isBatch()) { - unsigned targetStep = exchange.consumerStep; - for (const FragmentTransfer &transfer : exchange.specializations.front().transfers) - if (!transfer.local) { - targetStep = transfer.targetInsertionStep; - break; - } - setInsertionAtBoundary(rewriter, *exchange.target, targetStep); - auto result = realizeSpecialization(exchange, exchange.specializations.front(), rewriter); - if (failed(result)) return failure(); - replaceDeferred(exchange.deferred, *result, erase); - return success(); +static Value buildLaneLookup(ArrayRef values, Value lane, + Operation *anchor, IRRewriter &rewriter, + Location loc) { + Value table = createI64LookupTableConstant(rewriter, anchor, values); + Value selected = tensor::ExtractOp::create( + rewriter, loc, table, ValueRange {lane}); + return arith::IndexCastOp::create( + rewriter, loc, rewriter.getIndexType(), selected); +} + +static Value buildLaneLookupCached( + ArrayRef values, Value lane, Operation *anchor, + DenseConstantCache &denseConstants, IRRewriter &rewriter, Location loc) { + auto type = RankedTensorType::get( + {static_cast(values.size())}, rewriter.getI64Type()); + Value table = getOrCreateCachedTypedConstant( + denseConstants, rewriter, anchor, + DenseElementsAttr::get(type, values)); + Value selected = tensor::ExtractOp::create( + rewriter, loc, table, ValueRange {lane}); + return arith::IndexCastOp::create( + rewriter, loc, rewriter.getIndexType(), selected); +} + +static void setBatchTransferAttrs( + Operation *op, ArrayRef refs, + RewriterBase &rewriter) { + SmallVector channels, parents, counts, sources, targets; + for (const FragmentTransferRef &ref : refs) { + LogicalTransferMetadata metadata = getLogicalTransferMetadata(ref); + channels.push_back(metadata.channelId); + parents.push_back(metadata.parentExchangeId); + counts.push_back(metadata.parentTransferCount); + sources.push_back(metadata.sourceCore); + targets.push_back(metadata.targetCore); } - // Each branch specializes the deferred program to the explicit scheduled lane. - auto batch = cast(exchange.target->op); - auto resultType = dyn_cast(exchange.deferred.getOutput().getType()); - if (!resultType || !resultType.hasStaticShape()) - return exchange.deferred.emitOpError( - "phase 2 scheduled-batch lane dispatch requires a static ranked result"); - Location loc = exchange.deferred.getLoc(); - setInsertionAtBoundary(rewriter, *exchange.target, exchange.consumerStep); - Value dispatchBuffer = tensor::EmptyOp::create( - rewriter, loc, resultType.getShape(), resultType.getElementType()); - auto storeResult = [&](Value value) -> Value { - SmallVector offsets(resultType.getRank(), - rewriter.getIndexAttr(0)); - SmallVector sizes; - SmallVector strides(resultType.getRank(), - rewriter.getIndexAttr(1)); - for (int64_t dimension : resultType.getShape()) - sizes.push_back(rewriter.getIndexAttr(dimension)); - return tensor::InsertSliceOp::create(rewriter, loc, value, - dispatchBuffer, offsets, sizes, - strides); + auto dense = [&](ArrayRef values) -> DenseIntElementsAttr { + auto type = RankedTensorType::get( + {static_cast(values.size())}, rewriter.getI64Type()); + if (llvm::all_equal(values)) + values = values.take_front(); + return DenseIntElementsAttr::get(type, values); }; - Value result = dispatchBuffer; - for (unsigned lane = 0; lane < exchange.specializations.size(); ++lane) { - Value required = arith::ConstantIndexOp::create(rewriter, loc, lane); - Value selected = arith::CmpIOp::create( - rewriter, loc, arith::CmpIPredicate::eq, *batch.getLaneArgument(), - required); - auto ifOp = scf::IfOp::create(rewriter, loc, TypeRange {resultType}, - selected, true); - rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); - auto value = realizeSpecialization(exchange, exchange.specializations[lane], - rewriter); + op->setAttr("raptor.batch_channel_ids", dense(channels)); + op->setAttr("raptor.batch_source_cores", dense(sources)); + op->setAttr("raptor.batch_target_cores", dense(targets)); + op->setAttr("raptor.batch_parent_exchange_ids", dense(parents)); + op->setAttr("raptor.batch_parent_transfer_counts", dense(counts)); +} + +static FailureOr> tryRealizeAnalyzedInsertAssembly( + NormalizedDeferredExchange &exchange, + DeferredSpecialization &specialization, IRRewriter &rewriter) { + const std::optional &assembly = + specialization.program.insertAssembly; + if (!assembly || assembly->entries.empty()) + return std::optional(); + + SmallVector ordered; + DenseMap entryByRequirement; + ordered.reserve(assembly->entries.size()); + for (const DeferredInsertAssemblyEntry &entry : assembly->entries) { + if (entry.requirementIndex >= specialization.transfers.size()) + return std::optional(); + FragmentTransfer &transfer = + specialization.transfers[entry.requirementIndex]; + if (transfer.local + || transfer.requirementIndex != entry.requirementIndex) + return std::optional(); + ordered.push_back(&transfer); + entryByRequirement[entry.requirementIndex] = &entry; + } + llvm::stable_sort(ordered, [](FragmentTransfer *lhs, + FragmentTransfer *rhs) { + return lhs->globalOrder < rhs->globalOrder; + }); + FragmentRequirement &reference = specialization.requirements[ + ordered.front()->requirementIndex]; + for (auto [index, transfer] : llvm::enumerate(ordered)) { + const FragmentRequirement &requirement = + specialization.requirements[transfer->requirementIndex]; + if (requirement.publicationFragmentType + != reference.publicationFragmentType + || transfer->globalOrder != ordered.front()->globalOrder + index + || transfer->targetInsertionStep + != ordered.front()->targetInsertionStep) + return std::optional(); + } + if (ordered.size() != specialization.requirements.size()) + return std::optional(); + + Location loc = exchange.deferred.getLoc(); + Operation *clonedInitial = rewriter.clone(*assembly->initialValue); + Value assembled = clonedInitial->getResult(0); + SmallVector channels, sources, targets; + SmallVector metadataRefs; + for (FragmentTransfer *transfer : ordered) { + channels.push_back(transfer->channelId); + sources.push_back(transfer->sourceCore); + targets.push_back(transfer->targetCore); + metadataRefs.push_back({&exchange, &specialization, transfer}); + } + auto insertReceived = [&](Value index, Value current) -> FailureOr { + Value channel = buildLaneLookup( + channels, index, exchange.deferred, rewriter, loc); + Value source = buildLaneLookup( + sources, index, exchange.deferred, rewriter, loc); + Value target = buildLaneLookup( + targets, index, exchange.deferred, rewriter, loc); + auto receive = SpatChannelReceiveOp::create( + rewriter, loc, reference.publicationFragmentType, + channel, source, target); + setBatchTransferAttrs(receive, metadataRefs, rewriter); + + auto buildGeometry = [&](auto member, + unsigned dimension) -> OpFoldResult { + SmallVector values; + for (FragmentTransfer *transfer : ordered) { + const DeferredInsertAssemblyEntry *entry = + entryByRequirement.lookup(transfer->requirementIndex); + values.push_back((entry->targetGeometry.*member)[dimension]); + } + if (llvm::all_equal(values)) + return rewriter.getIndexAttr(values.front()); + return buildLaneLookup( + values, index, exchange.deferred, rewriter, loc); + }; + unsigned rank = assembly->resultType.getRank(); + SmallVector offsets, sizes, strides; + for (unsigned dimension = 0; dimension < rank; ++dimension) { + offsets.push_back(buildGeometry( + &StaticSliceGeometry::offsets, dimension)); + sizes.push_back(buildGeometry( + &StaticSliceGeometry::sizes, dimension)); + strides.push_back(buildGeometry( + &StaticSliceGeometry::strides, dimension)); + } + return tensor::InsertSliceOp::create( + rewriter, loc, receive.getOutput(), current, offsets, sizes, strides) + .getResult(); + }; + if (ordered.size() == 1) { + Value zero = getOrCreateIndexConstant( + rewriter, exchange.deferred.getOperation(), 0); + auto result = insertReceived(zero, assembled); + if (failed(result)) + return failure(); + return std::optional(*result); + } + Value lower = getOrCreateIndexConstant( + rewriter, exchange.deferred.getOperation(), 0); + Value upper = getOrCreateIndexConstant( + rewriter, exchange.deferred.getOperation(), ordered.size()); + Value step = getOrCreateIndexConstant( + rewriter, exchange.deferred.getOperation(), 1); + auto loop = buildNormalizedScfFor( + rewriter, loc, lower, upper, step, ValueRange {assembled}, + [&](OpBuilder &, Location, Value index, ValueRange iterArgs, + SmallVectorImpl &yielded) -> LogicalResult { + auto result = insertReceived(index, iterArgs.front()); + if (failed(result)) + return failure(); + yielded.push_back(*result); + return success(); + }); + if (failed(loop)) + return failure(); + return std::optional(loop->results.front()); +} + +static FailureOr materializeBatchLocalFragment( + ArrayRef requirements, Value lane, + SpatDeferredCommunicationOp deferred, IRRewriter &rewriter, + ArrayRef transfers = {}) { + FragmentRequirement *reference = nullptr; + SmallVector offsets; + offsets.reserve(requirements.size()); + bool commonPayload = true; + for (auto [index, requirement] : llvm::enumerate(requirements)) { + bool active = transfers.empty() || transfers[index]->local; + if (!active) { + offsets.push_back(0); + continue; + } + if (!reference) + reference = requirement; + if (requirement->producer->payload != reference->producer->payload + || requirement->publicationFragmentType + != reference->publicationFragmentType + || requirement->graphLane.has_value() != reference->graphLane.has_value()) + commonPayload = false; + offsets.push_back(requirement->graphLane + ? *requirement->graphLane + - requirement->producer->laneStart + : 0); + } + if (!reference) + return failure(); + if (commonPayload) { + if (!reference->graphLane) + return materializeProducedFragment(*reference, rewriter, + deferred.getLoc()); + int64_t placeholder = *reference->graphLane - reference->producer->laneStart; + for (auto [index, transfer] : llvm::enumerate(transfers)) + if (!transfer->local) + offsets[index] = placeholder; + Value offset = buildLaneLookup(offsets, lane, deferred.getOperation(), + rewriter, deferred.getLoc()); + return materializeProducedFragment(*reference, offset, rewriter, + deferred.getLoc()); + } + + struct LocalClass { + FragmentRequirement *representative = nullptr; + SmallVector lanes; + }; + SmallVector classes; + for (auto [laneIndex, requirementValue] : llvm::enumerate(requirements)) { + FragmentRequirement *requirement = requirementValue; + if (!transfers.empty() && !transfers[laneIndex]->local) + continue; + auto sameBehavior = [&](const LocalClass &candidate) { + FragmentRequirement *other = candidate.representative; + return requirement->producer->payload == other->producer->payload + && requirement->publicationFragmentType + == other->publicationFragmentType + && requirement->graphLane.has_value() + == other->graphLane.has_value(); + }; + auto classIt = llvm::find_if(classes, sameBehavior); + if (classIt == classes.end()) { + classes.push_back({requirement, {static_cast(laneIndex)}}); + } else { + classIt->lanes.push_back(laneIndex); + } + } + if (classes.empty() + || llvm::any_of(classes, [&](const LocalClass &localClass) { + return localClass.representative->publicationFragmentType + != reference->publicationFragmentType; + })) + return failure(); + + auto materializeClass = [&](const LocalClass &localClass) + -> FailureOr { + FragmentRequirement *representative = localClass.representative; + if (!representative->graphLane) + return materializeProducedFragment( + *representative, rewriter, deferred.getLoc()); + int64_t defaultOffset = + *representative->graphLane - representative->producer->laneStart; + SmallVector classOffsets(requirements.size(), defaultOffset); + for (unsigned laneIndex : localClass.lanes) { + FragmentRequirement *requirement = requirements[laneIndex]; + int64_t offset = + *requirement->graphLane - requirement->producer->laneStart; + if (offset < 0 || offset >= requirement->producer->laneCount) + return failure(); + classOffsets[laneIndex] = offset; + } + Value offset = buildLaneLookup( + classOffsets, lane, deferred.getOperation(), rewriter, + deferred.getLoc()); + return materializeProducedFragment( + *representative, offset, rewriter, deferred.getLoc()); + }; + if (classes.size() == 1) + return materializeClass(classes.front()); + + auto buildMembership = [&](ArrayRef members) -> Value { + if (members.size() == 1) { + Value member = getOrCreateIndexConstant( + rewriter, deferred.getOperation(), members.front()); + return arith::CmpIOp::create( + rewriter, deferred.getLoc(), arith::CmpIPredicate::eq, lane, + member); + } + SmallVector membership(requirements.size()); + for (unsigned member : members) + membership[member] = 1; + Value selected = buildLaneLookup( + membership, lane, deferred.getOperation(), rewriter, + deferred.getLoc()); + Value zero = getOrCreateIndexConstant( + rewriter, deferred.getOperation(), 0); + return arith::CmpIOp::create( + rewriter, deferred.getLoc(), arith::CmpIPredicate::ne, selected, + zero); + }; + if (classes.size() == 2) { + unsigned thenIndex = classes[0].lanes.size() <= classes[1].lanes.size() + ? 0 : 1; + unsigned elseIndex = 1 - thenIndex; + Value condition = buildMembership(classes[thenIndex].lanes); + auto selection = scf::IfOp::create( + rewriter, deferred.getLoc(), + TypeRange {reference->publicationFragmentType}, condition, true); + auto buildBranch = [&](Region ®ion, const LocalClass &localClass) + -> LogicalResult { + OpBuilder::InsertionGuard guard(rewriter); + auto existingYield = dyn_cast_or_null( + region.front().empty() ? nullptr : ®ion.front().back()); + if (existingYield) + rewriter.setInsertionPoint(existingYield); + else + rewriter.setInsertionPointToStart(®ion.front()); + auto value = materializeClass(localClass); + if (failed(value)) + return failure(); + if (existingYield) + existingYield->setOperands(*value); + else + scf::YieldOp::create(rewriter, deferred.getLoc(), *value); + return success(); + }; + if (failed(buildBranch(selection.getThenRegion(), classes[thenIndex])) + || failed(buildBranch(selection.getElseRegion(), classes[elseIndex]))) + return failure(); + return selection.getResult(0); + } + + llvm::stable_sort(classes, [](const LocalClass &lhs, + const LocalClass &rhs) { + return lhs.lanes.front() < rhs.lanes.front(); + }); + SmallVector classByLane(requirements.size()); + for (auto [classIndex, localClass] : llvm::enumerate(classes)) + for (unsigned laneIndex : localClass.lanes) + classByLane[laneIndex] = classIndex; + Value classId = buildLaneLookup(classByLane, lane, deferred.getOperation(), + rewriter, deferred.getLoc()); + SmallVector cases; + for (int64_t classIndex = 0; + classIndex < static_cast(classes.size()) - 1; + ++classIndex) + cases.push_back(classIndex); + auto selection = scf::IndexSwitchOp::create( + rewriter, deferred.getLoc(), + TypeRange {reference->publicationFragmentType}, classId, cases, + cases.size()); + auto buildRegion = [&](Region ®ion, const LocalClass &localClass) + -> LogicalResult { + OpBuilder::InsertionGuard guard(rewriter); + Block *block = rewriter.createBlock(®ion); + rewriter.setInsertionPointToEnd(block); + auto value = materializeClass(localClass); if (failed(value)) return failure(); - scf::YieldOp::create(rewriter, loc, storeResult(*value)); - rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front()); - scf::YieldOp::create(rewriter, loc, result); - rewriter.setInsertionPointAfter(ifOp); - result = ifOp.getResult(0); + scf::YieldOp::create(rewriter, deferred.getLoc(), *value); + return success(); + }; + for (auto [region, localClass] : + llvm::zip(selection.getCaseRegions(), ArrayRef(classes).drop_back())) + if (failed(buildRegion(region, localClass))) + return failure(); + if (failed(buildRegion(selection.getDefaultRegion(), classes.back()))) + return failure(); + return selection.getResult(0); +} + +struct DynamicSliceGeometry { + SmallVector offsets; + SmallVector sizes; + SmallVector strides; + RankedTensorType resultType; + bool identity = false; +}; + +static FailureOr buildBatchInnerGeometry( + ArrayRef specializations, unsigned leafIndex, + ArrayRef classLanes, Value lane, Operation *anchor, + IRRewriter &rewriter, Location loc) { + if (classLanes.empty()) + return failure(); + const DeferredProjectionLeaf &reference = + specializations[classLanes.front()].program.leaves[leafIndex]; + unsigned rank = reference.innerGeometry.offsets.size(); + if (reference.innerGeometry.sizes.size() != rank + || reference.innerGeometry.strides.size() != rank) + return failure(); + + DynamicSliceGeometry geometry; + geometry.identity = true; + auto buildDimension = [&](auto member, unsigned dimension) + -> FailureOr { + SmallVector values(specializations.size(), + (reference.innerGeometry.*member)[dimension]); + SmallVector activeValues; + for (unsigned classLane : classLanes) { + if (classLane >= specializations.size() + || leafIndex + >= specializations[classLane].program.leaves.size()) + return failure(); + const DeferredProjectionLeaf &leaf = + specializations[classLane].program.leaves[leafIndex]; + if ((leaf.innerGeometry.*member).size() != rank + || leaf.reconstructedType != reference.reconstructedType) + return failure(); + values[classLane] = (leaf.innerGeometry.*member)[dimension]; + activeValues.push_back(values[classLane]); + } + if (llvm::all_equal(activeValues)) + return OpFoldResult(rewriter.getIndexAttr(activeValues.front())); + return OpFoldResult(buildLaneLookup(values, lane, anchor, rewriter, loc)); + }; + for (unsigned dimension = 0; dimension < rank; ++dimension) { + auto offset = buildDimension(&StaticSliceGeometry::offsets, dimension); + auto size = buildDimension(&StaticSliceGeometry::sizes, dimension); + auto stride = buildDimension(&StaticSliceGeometry::strides, dimension); + if (failed(offset) || failed(size) || failed(stride)) + return failure(); + geometry.offsets.push_back(*offset); + geometry.sizes.push_back(*size); + geometry.strides.push_back(*stride); } - replaceDeferred(exchange.deferred, result, erase); + if (reference.kind != DeferredLeafKind::GraphBatchProjection) { + geometry.resultType = reference.reconstructedType; + return geometry; + } + if (reference.reconstructedType.getRank() != rank + 1) + return failure(); + geometry.resultType = RankedTensorType::get( + reference.reconstructedType.getShape().drop_front(), + reference.reconstructedType.getElementType()); + for (unsigned classLane : classLanes) { + const DeferredProjectionLeaf &leaf = + specializations[classLane].program.leaves[leafIndex]; + for (unsigned dimension = 0; dimension < rank; ++dimension) + geometry.identity &= leaf.innerGeometry.offsets[dimension] == 0 + && leaf.innerGeometry.sizes[dimension] + == geometry.resultType.getDimSize(dimension) + && leaf.innerGeometry.strides[dimension] == 1; + } + return geometry; +} + +static FailureOr reconstructBatchProjectionLeaf( + NormalizedDeferredExchange &exchange, unsigned leafIndex, + const BatchBoundaryClass &group, ArrayRef available, Value lane, + IRRewriter &rewriter) { + unsigned referenceLane = group.lanes.front(); + DeferredSpecialization &reference = + exchange.specializations[referenceLane]; + const DeferredProjectionLeaf &leaf = reference.program.leaves[leafIndex]; + SmallVector selected( + leaf.kind == DeferredLeafKind::ScalarSource ? 1 + : leaf.physicalSlots.size()); + for (auto [requirementIndex, requirement] : + llvm::enumerate(reference.requirements)) + if (requirement.leafIndex == leafIndex + && requirement.selectedPosition < selected.size()) { + if (selected[requirement.selectedPosition]) + return failure(); + selected[requirement.selectedPosition] = available[requirementIndex]; + } + if (llvm::any_of(selected, [](Value value) { return !value; })) + return failure(); + if (leaf.kind == DeferredLeafKind::ScalarSource) + return selected.front(); + + auto geometry = buildBatchInnerGeometry( + exchange.specializations, leafIndex, group.lanes, lane, + exchange.deferred, rewriter, exchange.deferred.getLoc()); + if (failed(geometry)) + return failure(); + SmallVector shaped; + for (Value fragment : selected) { + if (geometry->identity) { + shaped.push_back(fragment); + continue; + } + shaped.push_back(tensor::ExtractSliceOp::create( + rewriter, exchange.deferred.getLoc(), geometry->resultType, fragment, + geometry->offsets, geometry->sizes, geometry->strides)); + } + if (shaped.size() == 1 && shaped.front().getType() == leaf.reconstructedType) + return shaped.front(); + SmallVector expanded; + for (Value fragment : shaped) { + auto unit = addLeadingUnitTensorDimension( + rewriter, exchange.deferred.getLoc(), fragment); + if (failed(unit)) + return failure(); + expanded.push_back(*unit); + } + if (expanded.size() == 1 + && expanded.front().getType() == leaf.reconstructedType) + return expanded.front(); + return tensor::ConcatOp::create( + rewriter, exchange.deferred.getLoc(), leaf.reconstructedType, 0, + expanded).getResult(); +} + +static FailureOr materializeBatchResult( + NormalizedDeferredExchange &exchange, ArrayRef available, + const BatchBoundaryClass &group, Value lane, IRRewriter &rewriter) { + unsigned referenceLane = group.lanes.front(); + DeferredSpecialization &reference = exchange.specializations[referenceLane]; + IRMapping mapping; + if (reference.program.scheduledLane) + mapping.map(reference.program.scheduledLane, lane); + for (auto [value, staticValue] : reference.program.staticValues) { + if (mapping.contains(value)) + continue; + SmallVector values; + values.reserve(exchange.specializations.size()); + for (DeferredSpecialization &specialization : exchange.specializations) { + auto it = specialization.program.staticValues.find(value); + if (it == specialization.program.staticValues.end()) + return failure(); + values.push_back(it->second); + } + Value selected = buildLaneLookup(values, lane, exchange.deferred, + rewriter, exchange.deferred.getLoc()); + if (!value.getType().isIndex()) + selected = arith::IndexCastOp::create( + rewriter, exchange.deferred.getLoc(), value.getType(), selected); + mapping.map(value, selected); + } + for (auto [index, leaf] : llvm::enumerate(reference.program.leaves)) { + auto reconstructed = reconstructBatchProjectionLeaf( + exchange, index, group, available, lane, rewriter); + if (failed(reconstructed)) + return failure(); + mapping.map(leaf.replacementRoot, *reconstructed); + } + for (Operation *op : reference.program.residualOps) { + Operation *copy = rewriter.clone(*op, mapping); + for (auto [original, result] : llvm::zip(op->getResults(), + copy->getResults())) + mapping.map(original, result); + } + Value result = mapping.lookupOrDefault(reference.program.yieldedValue); + return result && result.getType() == exchange.deferred.getOutput().getType() + ? FailureOr(result) + : FailureOr(failure()); +} + +static FailureOr tryEmitLoopedScalarSends( + ArrayRef sends, IRRewriter &rewriter) { + if (sends.size() < 2) + return false; + FragmentTransferRef first = sends.front(); + FragmentRequirement &reference = first.specialization->requirements[ + first.transfer->requirementIndex]; + ProducedValue *producer = reference.producer; + if (producer->scheduled->isBatch()) + return false; + + SmallVector channels, sources, targets, offsets; + for (FragmentTransferRef ref : sends) { + FragmentRequirement &requirement = ref.specialization->requirements[ + ref.transfer->requirementIndex]; + if (requirement.producer->scheduled != producer->scheduled + || ref.transfer->sourceInsertionStep + != first.transfer->sourceInsertionStep + || requirement.producer->payload != producer->payload + || requirement.publicationFragmentType + != reference.publicationFragmentType + || requirement.graphLane.has_value() != reference.graphLane.has_value()) + return false; + int64_t offset = requirement.graphLane + ? *requirement.graphLane + - requirement.producer->laneStart + : 0; + if (offset < 0 || offset >= requirement.producer->laneCount) + return failure(); + channels.push_back(ref.transfer->channelId); + sources.push_back(ref.transfer->sourceCore); + targets.push_back(ref.transfer->targetCore); + offsets.push_back(offset); + } + + Location loc = first.exchange->deferred.getLoc(); + Value lower = getOrCreateIndexConstant( + rewriter, first.exchange->deferred.getOperation(), 0); + Value upper = getOrCreateIndexConstant( + rewriter, first.exchange->deferred.getOperation(), sends.size()); + Value step = getOrCreateIndexConstant( + rewriter, first.exchange->deferred.getOperation(), 1); + auto loop = buildNormalizedScfFor( + rewriter, loc, lower, upper, step, ValueRange {}, + [&](OpBuilder &, Location, Value index, ValueRange, + SmallVectorImpl &) -> LogicalResult { + Value channel = buildLaneLookup( + channels, index, first.exchange->deferred, rewriter, loc); + Value source = buildLaneLookup( + sources, index, first.exchange->deferred, rewriter, loc); + Value target = buildLaneLookup( + targets, index, first.exchange->deferred, rewriter, loc); + Value payload; + if (reference.graphLane) { + Value offset = buildLaneLookup( + offsets, index, first.exchange->deferred, rewriter, loc); + auto materialized = materializeProducedFragment( + reference, offset, rewriter, loc); + if (failed(materialized)) + return failure(); + payload = *materialized; + } + else { + auto materialized = materializeProducedFragment( + reference, rewriter, loc); + if (failed(materialized)) + return failure(); + payload = *materialized; + } + auto send = SpatChannelSendOp::create( + rewriter, loc, channel, source, target, payload); + setBatchTransferAttrs(send, sends, rewriter); + return success(); + }); + if (failed(loop)) + return failure(); + return true; +} + +static FailureOr tryEmitConsecutiveAnalyzedAssembly( + ArrayRef actions, size_t start, + AssemblyCache &assemblyCache, IRRewriter &rewriter) { + const BoundaryAction &first = actions[start]; + DeferredSpecialization *specialization = first.ref.specialization; + if (first.kind != BoundaryActionKind::Receive) + return 0; + size_t transferCount = llvm::count_if( + specialization->transfers, + [](const FragmentTransfer &transfer) { return !transfer.local; }); + if (transferCount == 0 || start + transferCount > actions.size()) + return 0; + llvm::SmallDenseSet expected; + for (FragmentTransfer &transfer : specialization->transfers) + if (!transfer.local) + expected.insert(&transfer); + for (const BoundaryAction &action : actions.slice(start, transferCount)) + if (action.kind != BoundaryActionKind::Receive + || action.ref.specialization != specialization + || !expected.erase(action.ref.transfer)) + return 0; + if (!expected.empty()) + return 0; + FailureOr> assembled = + tryRealizeAnalyzedInsertAssembly( + *first.ref.exchange, *specialization, rewriter); + if (failed(assembled)) + return failure(); + if (!*assembled) + return 0; + assemblyCache[specialization] = **assembled; + return transferCount; +} + +static LogicalResult realizeScalarBoundary( + ScheduledBoundaryPlan &boundary, ReceiveCache &receiveCache, + AssemblyCache &assemblyCache, IndexConstantCache &indexConstants, + DenseConstantCache &typedConstants, IRRewriter &rewriter, + DeferredEraseSet &erase) { + setInsertionAtBoundary( + rewriter, *boundary.scheduled, boundary.insertionStep); + ArrayRef actions = boundary.actionsByLane.front(); + for (size_t index = 0; index < actions.size();) { + const BoundaryAction &action = actions[index]; + if (action.kind == BoundaryActionKind::Receive) { + auto compacted = tryEmitConsecutiveAnalyzedAssembly( + actions, index, assemblyCache, rewriter); + if (failed(compacted)) + return failure(); + if (*compacted != 0) { + index += *compacted; + continue; + } + if (failed(emitFragmentReceiveHere( + action.ref, receiveCache, indexConstants, rewriter))) + return failure(); + ++index; + continue; + } + + size_t end = index + 1; + while (end < actions.size() + && actions[end].kind == BoundaryActionKind::Send + && haveCompatibleSendEmission( + action.ref, actions[end].ref, boundary.insertionStep, + boundary.insertionStep)) + ++end; + SmallVector sends; + for (const BoundaryAction &send : actions.slice(index, end - index)) + sends.push_back(send.ref); + bool compacted = false; + if (sends.size() >= 2) { + auto result = tryEmitLoopedScalarSends(sends, rewriter); + if (failed(result)) + return failure(); + compacted = *result; + } + if (!compacted) + for (FragmentTransferRef ref : sends) + if (failed(emitFragmentSendHere( + *ref.exchange, *ref.specialization, *ref.transfer, + indexConstants, rewriter))) + return failure(); + index = end; + } + + for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) { + auto result = realizeSpecialization( + *exchange, exchange->specializations.front(), receiveCache, + assemblyCache, indexConstants, typedConstants, rewriter); + if (failed(result)) + return failure(); + replaceDeferred(exchange->deferred, *result, erase); + } + return success(); +} + +struct BatchActionSemanticKey { + BoundaryActionKind kind = BoundaryActionKind::Send; + uint64_t parentExchangeId = 0; + unsigned requirementIndex = 0; + unsigned leafIndex = 0; + unsigned selectedPosition = 0; + Type publicationFragmentType; + std::optional sendSignature; + + bool operator==(const BatchActionSemanticKey &other) const { + if (kind != other.kind + || publicationFragmentType != other.publicationFragmentType) + return false; + if (kind == BoundaryActionKind::Send) + return sendSignature && other.sendSignature + && haveSameSendEmissionSignature(*sendSignature, + *other.sendSignature); + return parentExchangeId == other.parentExchangeId + && requirementIndex == other.requirementIndex + && leafIndex == other.leafIndex + && selectedPosition == other.selectedPosition + && !sendSignature && !other.sendSignature; + } +}; + +struct DeferredLeafStructuralKey { + DeferredLeafKind kind = DeferredLeafKind::ScalarSource; + RankedTensorType reconstructedType; + unsigned selectedPositionCount = 0; + unsigned innerGeometryRank = 0; + bool identityGeometry = false; + + bool operator==(const DeferredLeafStructuralKey &other) const { + return kind == other.kind && reconstructedType == other.reconstructedType + && selectedPositionCount == other.selectedPositionCount + && innerGeometryRank == other.innerGeometryRank + && identityGeometry == other.identityGeometry; + } +}; + +struct DeferredInsertAssemblyStructuralKey { + RankedTensorType resultType; + SmallVector requirementIndices; + SmallVector geometryRanks; + + bool operator==(const DeferredInsertAssemblyStructuralKey &other) const { + return resultType == other.resultType + && requirementIndices == other.requirementIndices + && geometryRanks == other.geometryRanks; + } +}; + +struct BatchExchangeSemanticKey { + uint64_t exchangeId = 0; + SmallVector residualOps; + Value yieldedValue; + SmallVector leaves; + SmallVector localByRequirement; + std::optional assembly; + + bool operator==(const BatchExchangeSemanticKey &other) const { + return exchangeId == other.exchangeId && residualOps == other.residualOps + && yieldedValue == other.yieldedValue && leaves == other.leaves + && localByRequirement == other.localByRequirement + && assembly == other.assembly; + } +}; + +struct BatchLaneSemanticKey { + SmallVector actions; + SmallVector exchanges; + + bool operator==(const BatchLaneSemanticKey &other) const { + return actions == other.actions && exchanges == other.exchanges; + } +}; + +static BatchLaneSemanticKey buildBatchLaneSemanticKey( + const ScheduledBoundaryPlan &boundary, unsigned lane) { + BatchLaneSemanticKey key; + for (const BoundaryAction &action : boundary.actionsByLane[lane]) { + const FragmentRequirement &requirement = + action.ref.specialization->requirements[ + action.ref.transfer->requirementIndex]; + BatchActionSemanticKey actionKey { + action.kind, action.ref.exchange->exchangeId, + action.ref.transfer->requirementIndex, requirement.leafIndex, + requirement.selectedPosition, requirement.publicationFragmentType, + std::nullopt}; + if (action.kind == BoundaryActionKind::Send) + actionKey.sendSignature = getSendEmissionSignature(action.ref); + key.actions.push_back(std::move(actionKey)); + } + for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) { + DeferredSpecialization &specialization = exchange->specializations[lane]; + BatchExchangeSemanticKey exchangeKey; + exchangeKey.exchangeId = exchange->exchangeId; + exchangeKey.residualOps = specialization.program.residualOps; + exchangeKey.yieldedValue = specialization.program.yieldedValue; + for (const DeferredProjectionLeaf &leaf : specialization.program.leaves) { + bool identity = leaf.kind != DeferredLeafKind::GraphBatchProjection; + if (!identity) { + identity = true; + for (auto [offset, stride] : llvm::zip( + leaf.innerGeometry.offsets, leaf.innerGeometry.strides)) + identity &= offset == 0 && stride == 1; + } + exchangeKey.leaves.push_back( + {leaf.kind, leaf.reconstructedType, + static_cast(leaf.kind == DeferredLeafKind::ScalarSource + ? 1 + : leaf.physicalSlots.size()), + static_cast(leaf.innerGeometry.offsets.size()), identity}); + } + for (const FragmentTransfer &transfer : specialization.transfers) + exchangeKey.localByRequirement.push_back(transfer.local); + if (specialization.program.insertAssembly) { + DeferredInsertAssemblyStructuralKey assemblyKey; + assemblyKey.resultType = + specialization.program.insertAssembly->resultType; + for (const DeferredInsertAssemblyEntry &entry : + specialization.program.insertAssembly->entries) { + assemblyKey.requirementIndices.push_back(entry.requirementIndex); + assemblyKey.geometryRanks.push_back(entry.targetGeometry.offsets.size()); + } + exchangeKey.assembly = std::move(assemblyKey); + } + key.exchanges.push_back(std::move(exchangeKey)); + } + return key; +} + +static size_t hashBatchLaneSemanticKey(const BatchLaneSemanticKey &key) { + llvm::hash_code hash = llvm::hash_value(key.actions.size()); + for (const BatchActionSemanticKey &action : key.actions) { + hash = llvm::hash_combine( + hash, static_cast(action.kind), + action.publicationFragmentType.getAsOpaquePointer()); + if (action.sendSignature) + hash = llvm::hash_combine( + hash, hashSendEmissionSignature(*action.sendSignature)); + else + hash = llvm::hash_combine( + hash, action.parentExchangeId, action.requirementIndex, + action.leafIndex, action.selectedPosition); + } + for (const BatchExchangeSemanticKey &exchange : key.exchanges) { + hash = llvm::hash_combine( + hash, exchange.exchangeId, exchange.yieldedValue.getAsOpaquePointer(), + exchange.residualOps.size(), exchange.leaves.size(), + exchange.localByRequirement.size(), exchange.assembly.has_value()); + for (Operation *op : exchange.residualOps) + hash = llvm::hash_combine(hash, op); + for (const DeferredLeafStructuralKey &leaf : exchange.leaves) + hash = llvm::hash_combine( + hash, static_cast(leaf.kind), + leaf.reconstructedType.getAsOpaquePointer(), + leaf.selectedPositionCount, leaf.innerGeometryRank, + leaf.identityGeometry); + for (bool local : exchange.localByRequirement) + hash = llvm::hash_combine(hash, local); + if (exchange.assembly) { + hash = llvm::hash_combine( + hash, exchange.assembly->resultType.getAsOpaquePointer()); + for (auto values : llvm::zip_equal( + exchange.assembly->requirementIndices, + exchange.assembly->geometryRanks)) + hash = llvm::hash_combine( + hash, std::get<0>(values), std::get<1>(values)); + } + } + return static_cast(hash); +} + +static SmallVector classifyBatchBoundary( + const ScheduledBoundaryPlan &boundary) { + SmallVector classes; + SmallVector classKeys; + DenseMap> classesByHash; + for (unsigned lane = 0; lane < boundary.actionsByLane.size(); ++lane) { + BatchLaneSemanticKey key = buildBatchLaneSemanticKey(boundary, lane); + size_t hash = hashBatchLaneSemanticKey(key); + std::optional matchingClass; + for (unsigned classIndex : classesByHash.lookup(hash)) + if (classKeys[classIndex] == key) { + matchingClass = classIndex; + break; + } + if (!matchingClass) { + classesByHash[hash].push_back(classes.size()); + classes.push_back({{lane}}); + classKeys.push_back(std::move(key)); + } else { + classes[*matchingClass].lanes.push_back(lane); + } + } + return classes; +} + +static Value buildBatchActionTableIndex( + Value lane, Value actionIndex, unsigned laneCount, Operation *anchor, + IRRewriter &rewriter, Location loc) { + // Runtime tables are action-major over every lane; compact metadata is + // action-major over active class lanes only. + Value width = getOrCreateIndexConstant(rewriter, anchor, laneCount); + Value actionBase = arith::MulIOp::create( + rewriter, loc, actionIndex, width); + return arith::AddIOp::create(rewriter, loc, actionBase, lane); +} + +static LogicalResult emitBatchSendRun( + const ScheduledBoundaryPlan &boundary, const BatchBoundaryClass &group, + size_t start, size_t runLength, Value lane, + DenseConstantCache &denseConstants, IRRewriter &rewriter) { + unsigned laneCount = boundary.actionsByLane.size(); + const BoundaryAction &firstAction = + boundary.actionsByLane[group.lanes.front()][start]; + FragmentRequirement &reference = + firstAction.ref.specialization->requirements[ + firstAction.ref.transfer->requirementIndex]; + size_t tableSize = static_cast(laneCount) * runLength; + SmallVector channels(tableSize, + firstAction.ref.transfer->channelId); + SmallVector sources(tableSize, + firstAction.ref.transfer->sourceCore); + SmallVector targets(tableSize, + firstAction.ref.transfer->targetCore); + int64_t defaultOffset = reference.graphLane + ? *reference.graphLane - reference.producer->laneStart : 0; + SmallVector offsets(tableSize, defaultOffset); + SmallVector metadataRefs; + for (size_t actionOffset = 0; actionOffset < runLength; + ++actionOffset) + for (unsigned classLane : group.lanes) { + const BoundaryAction &action = + boundary.actionsByLane[classLane][start + actionOffset]; + FragmentRequirement &requirement = + action.ref.specialization->requirements[ + action.ref.transfer->requirementIndex]; + if (action.kind != BoundaryActionKind::Send + || !haveCompatibleSendEmission( + firstAction.ref, action.ref, boundary.insertionStep, + boundary.insertionStep)) + return failure(); + size_t tableIndex = actionOffset * laneCount + classLane; + channels[tableIndex] = action.ref.transfer->channelId; + sources[tableIndex] = action.ref.transfer->sourceCore; + targets[tableIndex] = action.ref.transfer->targetCore; + if (requirement.graphLane) + offsets[tableIndex] = + *requirement.graphLane - requirement.producer->laneStart; + metadataRefs.push_back(action.ref); + } + + auto emitOne = [&](Value tableIndex) -> LogicalResult { + Location loc = firstAction.ref.exchange->deferred.getLoc(); + Value payload; + if (reference.graphLane) { + Value offset = buildLaneLookupCached( + offsets, tableIndex, firstAction.ref.exchange->deferred, + denseConstants, rewriter, loc); + auto materialized = materializeProducedFragment( + reference, offset, rewriter, loc); + if (failed(materialized)) + return failure(); + payload = *materialized; + } else { + auto materialized = materializeProducedFragment( + reference, rewriter, loc); + if (failed(materialized)) + return failure(); + payload = *materialized; + } + auto send = SpatChannelSendOp::create( + rewriter, loc, + buildLaneLookupCached(channels, tableIndex, + firstAction.ref.exchange->deferred, + denseConstants, rewriter, loc), + buildLaneLookupCached(sources, tableIndex, + firstAction.ref.exchange->deferred, + denseConstants, rewriter, loc), + buildLaneLookupCached(targets, tableIndex, + firstAction.ref.exchange->deferred, + denseConstants, rewriter, loc), + payload); + setBatchTransferAttrs(send, metadataRefs, rewriter); + return success(); + }; + if (runLength == 1) + return emitOne(lane); + Location loc = firstAction.ref.exchange->deferred.getLoc(); + Value lower = getOrCreateIndexConstant( + rewriter, firstAction.ref.exchange->deferred, 0); + Value upper = getOrCreateIndexConstant( + rewriter, firstAction.ref.exchange->deferred, runLength); + Value step = getOrCreateIndexConstant( + rewriter, firstAction.ref.exchange->deferred, 1); + auto loop = buildNormalizedScfFor( + rewriter, loc, lower, upper, step, ValueRange {}, + [&](OpBuilder &, Location, Value actionIndex, ValueRange, + SmallVectorImpl &) -> LogicalResult { + Value tableIndex = buildBatchActionTableIndex( + lane, actionIndex, laneCount, + firstAction.ref.exchange->deferred, rewriter, loc); + return emitOne(tableIndex); + }); + return success(succeeded(loop)); +} + +static LogicalResult emitBatchReceiveAction( + const ScheduledBoundaryPlan &boundary, const BatchBoundaryClass &group, + size_t actionIndex, Value lane, ReceiveCache &receiveCache, + DenseConstantCache &denseConstants, IRRewriter &rewriter) { + unsigned laneCount = boundary.actionsByLane.size(); + const BoundaryAction &first = + boundary.actionsByLane[group.lanes.front()][actionIndex]; + FragmentRequirement &reference = first.ref.specialization->requirements[ + first.ref.transfer->requirementIndex]; + SmallVector channels(laneCount, first.ref.transfer->channelId); + SmallVector sources(laneCount, first.ref.transfer->sourceCore); + SmallVector targets(laneCount, first.ref.transfer->targetCore); + SmallVector refs; + for (unsigned classLane : group.lanes) { + const BoundaryAction &action = + boundary.actionsByLane[classLane][actionIndex]; + FragmentRequirement &requirement = + action.ref.specialization->requirements[ + action.ref.transfer->requirementIndex]; + if (action.kind != BoundaryActionKind::Receive + || action.ref.exchange != first.ref.exchange + || requirement.publicationFragmentType + != reference.publicationFragmentType) + return failure(); + channels[classLane] = action.ref.transfer->channelId; + sources[classLane] = action.ref.transfer->sourceCore; + targets[classLane] = action.ref.transfer->targetCore; + refs.push_back(action.ref); + } + Location loc = first.ref.exchange->deferred.getLoc(); + auto receive = SpatChannelReceiveOp::create( + rewriter, loc, reference.publicationFragmentType, + buildLaneLookupCached(channels, lane, first.ref.exchange->deferred, + denseConstants, rewriter, loc), + buildLaneLookupCached(sources, lane, first.ref.exchange->deferred, + denseConstants, rewriter, loc), + buildLaneLookupCached(targets, lane, first.ref.exchange->deferred, + denseConstants, rewriter, loc)); + setBatchTransferAttrs(receive, refs, rewriter); + for (const FragmentTransferRef &ref : refs) + receiveCache[ref.transfer] = receive.getOutput(); + return success(); +} + +static FailureOr> tryEmitBatchAnalyzedAssembly( + ScheduledBoundaryPlan &boundary, const BatchBoundaryClass &group, + NormalizedDeferredExchange &exchange, size_t startAction, Value lane, + ReceiveCache &receiveCache, DenseConstantCache &denseConstants, + IRRewriter &rewriter) { + (void)receiveCache; + unsigned referenceLane = group.lanes.front(); + DeferredSpecialization &reference = + exchange.specializations[referenceLane]; + const std::optional &assembly = + reference.program.insertAssembly; + if (!assembly || assembly->entries.empty()) + return std::optional(); + size_t entryCount = assembly->entries.size(); + if (startAction + entryCount + > boundary.actionsByLane[referenceLane].size()) + return std::optional(); + + Type fragmentType; + for (unsigned classLane : group.lanes) { + DeferredSpecialization &specialization = + exchange.specializations[classLane]; + if (!specialization.program.insertAssembly + || specialization.program.insertAssembly->entries.size() + != entryCount) + return std::optional(); + for (auto [entryIndex, entry] : llvm::enumerate( + specialization.program.insertAssembly->entries)) { + if (entry.requirementIndex >= specialization.requirements.size()) + return std::optional(); + const BoundaryAction &action = + boundary.actionsByLane[classLane][startAction + entryIndex]; + const FragmentRequirement &requirement = + specialization.requirements[entry.requirementIndex]; + if (action.kind != BoundaryActionKind::Receive + || action.ref.exchange != &exchange + || action.ref.transfer->requirementIndex != entry.requirementIndex + || action.ref.specialization != &specialization + || (fragmentType + && requirement.publicationFragmentType != fragmentType)) + return std::optional(); + fragmentType = requirement.publicationFragmentType; + } + } + if (!fragmentType) + return std::optional(); + + unsigned laneCount = boundary.actionsByLane.size(); + size_t tableSize = entryCount * laneCount; + const BoundaryAction &first = + boundary.actionsByLane[referenceLane][startAction]; + SmallVector channels(tableSize, first.ref.transfer->channelId); + SmallVector sources(tableSize, first.ref.transfer->sourceCore); + SmallVector targets(tableSize, first.ref.transfer->targetCore); + unsigned rank = assembly->resultType.getRank(); + SmallVector> offsets( + rank, SmallVector(tableSize)); + SmallVector> sizes( + rank, SmallVector(tableSize, 1)); + SmallVector> strides( + rank, SmallVector(tableSize, 1)); + SmallVector metadataRefs; + metadataRefs.reserve(entryCount * group.lanes.size()); + for (size_t entryIndex = 0; entryIndex < entryCount; ++entryIndex) + for (unsigned classLane : group.lanes) { + DeferredSpecialization &specialization = + exchange.specializations[classLane]; + const DeferredInsertAssemblyEntry &entry = + specialization.program.insertAssembly->entries[entryIndex]; + const BoundaryAction &action = + boundary.actionsByLane[classLane][startAction + entryIndex]; + size_t tableIndex = entryIndex * laneCount + classLane; + channels[tableIndex] = action.ref.transfer->channelId; + sources[tableIndex] = action.ref.transfer->sourceCore; + targets[tableIndex] = action.ref.transfer->targetCore; + for (unsigned dimension = 0; dimension < rank; ++dimension) { + offsets[dimension][tableIndex] = entry.targetGeometry.offsets[dimension]; + sizes[dimension][tableIndex] = entry.targetGeometry.sizes[dimension]; + strides[dimension][tableIndex] = entry.targetGeometry.strides[dimension]; + } + metadataRefs.push_back(action.ref); + } + + Operation *initial = rewriter.clone(*assembly->initialValue); + Value assembled = initial->getResult(0); + Location loc = exchange.deferred.getLoc(); + auto emitEntry = [&](Value tableIndex, Value current) -> FailureOr { + auto receive = SpatChannelReceiveOp::create( + rewriter, loc, fragmentType, + buildLaneLookupCached(channels, tableIndex, exchange.deferred, + denseConstants, rewriter, loc), + buildLaneLookupCached(sources, tableIndex, exchange.deferred, + denseConstants, rewriter, loc), + buildLaneLookupCached(targets, tableIndex, exchange.deferred, + denseConstants, rewriter, loc)); + setBatchTransferAttrs(receive, metadataRefs, rewriter); + SmallVector targetOffsets, targetSizes, targetStrides; + for (unsigned dimension = 0; dimension < rank; ++dimension) { + targetOffsets.push_back(buildLaneLookupCached( + offsets[dimension], tableIndex, exchange.deferred, denseConstants, + rewriter, loc)); + targetSizes.push_back(buildLaneLookupCached( + sizes[dimension], tableIndex, exchange.deferred, denseConstants, + rewriter, loc)); + targetStrides.push_back(buildLaneLookupCached( + strides[dimension], tableIndex, exchange.deferred, denseConstants, + rewriter, loc)); + } + return tensor::InsertSliceOp::create( + rewriter, loc, receive.getOutput(), current, targetOffsets, + targetSizes, targetStrides).getResult(); + }; + if (entryCount == 1) { + auto result = emitEntry(lane, assembled); + if (failed(result)) + return failure(); + return std::optional(*result); + } + + Value lower = getOrCreateIndexConstant(rewriter, exchange.deferred, 0); + Value upper = getOrCreateIndexConstant( + rewriter, exchange.deferred, entryCount); + Value step = getOrCreateIndexConstant(rewriter, exchange.deferred, 1); + auto loop = buildNormalizedScfFor( + rewriter, loc, lower, upper, step, ValueRange {assembled}, + [&](OpBuilder &, Location, Value actionIndex, ValueRange iterArgs, + SmallVectorImpl &yielded) -> LogicalResult { + Value tableIndex = buildBatchActionTableIndex( + lane, actionIndex, laneCount, exchange.deferred, rewriter, loc); + auto result = emitEntry(tableIndex, iterArgs.front()); + if (failed(result)) + return failure(); + yielded.push_back(*result); + return success(); + }); + if (failed(loop)) + return failure(); + return std::optional(loop->results.front()); +} + +static FailureOr> emitBatchBoundaryClass( + ScheduledBoundaryPlan &boundary, const BatchBoundaryClass &group, + Value lane, DenseConstantCache &denseConstants, IRRewriter &rewriter) { + ReceiveCache receiveCache; + AssemblyCache assemblyCache; + ArrayRef representative = + boundary.actionsByLane[group.lanes.front()]; + for (size_t index = 0; index < representative.size();) { + if (representative[index].kind == BoundaryActionKind::Receive) { + NormalizedDeferredExchange &exchange = + *representative[index].ref.exchange; + auto assembled = tryEmitBatchAnalyzedAssembly( + boundary, group, exchange, index, lane, receiveCache, + denseConstants, rewriter); + if (failed(assembled)) + return failure(); + if (*assembled) { + DeferredSpecialization &specialization = + exchange.specializations[group.lanes.front()]; + assemblyCache[&specialization] = **assembled; + index += specialization.program.insertAssembly->entries.size(); + continue; + } + if (failed(emitBatchReceiveAction( + boundary, group, index, lane, receiveCache, denseConstants, + rewriter))) + return failure(); + ++index; + continue; + } + size_t end = index + 1; + while (end < representative.size() + && representative[end].kind == BoundaryActionKind::Send + && haveCompatibleSendEmission( + representative[index].ref, representative[end].ref, + boundary.insertionStep, boundary.insertionStep)) + ++end; + if (failed(emitBatchSendRun( + boundary, group, index, end - index, lane, denseConstants, + rewriter))) + return failure(); + index = end; + } + + SmallVector results; + for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) { + unsigned referenceLane = group.lanes.front(); + DeferredSpecialization &specialization = + exchange->specializations[referenceLane]; + if (Value assembled = assemblyCache.lookup(&specialization)) { + results.push_back(assembled); + continue; + } + SmallVector available; + available.reserve(specialization.requirements.size()); + for (unsigned requirementIndex = 0; + requirementIndex < specialization.requirements.size(); + ++requirementIndex) { + FragmentTransfer &transfer = + specialization.transfers[requirementIndex]; + if (!transfer.local) { + Value received = receiveCache.lookup(&transfer); + if (!received) + return failure(); + available.push_back(received); + continue; + } + SmallVector requirements; + SmallVector transfers; + for (DeferredSpecialization &laneSpecialization : + exchange->specializations) { + requirements.push_back( + &laneSpecialization.requirements[requirementIndex]); + transfers.push_back( + &laneSpecialization.transfers[requirementIndex]); + } + auto local = materializeBatchLocalFragment( + requirements, lane, exchange->deferred, rewriter, transfers); + if (failed(local)) + return failure(); + available.push_back(*local); + } + auto result = materializeBatchResult( + *exchange, available, group, lane, rewriter); + if (failed(result)) + return failure(); + results.push_back(*result); + } + return results; +} + +static Value buildClassMembershipCondition( + ArrayRef members, unsigned laneCount, Value lane, + Operation *anchor, DenseConstantCache &denseConstants, + IRRewriter &rewriter, Location loc) { + if (members.size() == 1) { + Value selected = getOrCreateIndexConstant( + rewriter, anchor, members.front()); + return arith::CmpIOp::create( + rewriter, loc, arith::CmpIPredicate::eq, lane, selected); + } + SmallVector membership(laneCount); + for (unsigned member : members) + membership[member] = 1; + Value selected = buildLaneLookupCached( + membership, lane, anchor, denseConstants, rewriter, loc); + Value zero = getOrCreateIndexConstant(rewriter, anchor, 0); + return arith::CmpIOp::create( + rewriter, loc, arith::CmpIPredicate::ne, selected, zero); +} + +static FailureOr> tryEmitSimpleParentBoundary( + ScheduledBoundaryPlan &boundary, + ArrayRef classes, Value lane, + DenseConstantCache &denseConstants, IRRewriter &rewriter) { + if (classes.size() <= 2 || boundary.targetExchanges.size() != 1) + return std::optional(); + NormalizedDeferredExchange *exchange = boundary.targetExchanges.front(); + if (llvm::any_of(exchange->specializations, + [](const DeferredSpecialization &specialization) { + return specialization.transfers.size() != 1; + })) + return std::optional(); + + SmallVector remoteLanes, localLanes; + SmallVector preSendLanes, postSendLanes, localSendLanes; + const BoundaryAction *referenceReceive = nullptr; + for (unsigned classLane = 0; + classLane < boundary.actionsByLane.size(); ++classLane) { + ArrayRef actions = boundary.actionsByLane[classLane]; + if (actions.size() > 2) + return std::optional(); + const BoundaryAction *send = nullptr; + const BoundaryAction *receive = nullptr; + for (const BoundaryAction &action : actions) { + if (action.ref.exchange != exchange) + return std::optional(); + const BoundaryAction **slot = action.kind == BoundaryActionKind::Send + ? &send : &receive; + if (*slot) + return std::optional(); + *slot = &action; + } + FragmentTransfer &targetTransfer = + exchange->specializations[classLane].transfers.front(); + if ((receive == nullptr) != targetTransfer.local) + return std::optional(); + if (receive) { + const FragmentRequirement &requirement = + receive->ref.specialization->requirements[ + receive->ref.transfer->requirementIndex]; + if (referenceReceive) { + const FragmentRequirement &reference = + referenceReceive->ref.specialization->requirements[ + referenceReceive->ref.transfer->requirementIndex]; + if (receive->ref.transfer->requirementIndex + != referenceReceive->ref.transfer->requirementIndex + || requirement.leafIndex != reference.leafIndex + || requirement.selectedPosition != reference.selectedPosition + || requirement.publicationFragmentType + != reference.publicationFragmentType) + return std::optional(); + } else { + referenceReceive = receive; + } + } + if (!receive) { + localLanes.push_back(classLane); + if (send) + localSendLanes.push_back(classLane); + continue; + } + remoteLanes.push_back(classLane); + if (send) + (actions.front().kind == BoundaryActionKind::Send + ? preSendLanes : postSendLanes).push_back(classLane); + } + if (remoteLanes.empty()) + return std::optional(); + + auto filteredBoundary = [&](ArrayRef lanes, + std::optional kind, + bool withTarget) { + ScheduledBoundaryPlan filtered; + filtered.scheduled = boundary.scheduled; + filtered.insertionStep = boundary.insertionStep; + filtered.actionsByLane.resize(boundary.actionsByLane.size()); + for (unsigned classLane : lanes) + for (const BoundaryAction &action : + boundary.actionsByLane[classLane]) + if (!kind || action.kind == *kind) + filtered.actionsByLane[classLane].push_back(action); + if (withTarget) + filtered.targetExchanges.push_back(exchange); + return filtered; + }; + auto emitClass = [&](ScheduledBoundaryPlan &filtered, + ArrayRef lanes) -> FailureOr> { + BatchBoundaryClass group; + llvm::append_range(group.lanes, lanes); + return emitBatchBoundaryClass( + filtered, group, lane, denseConstants, rewriter); + }; + auto emitOptionalSends = [&](ArrayRef executingLanes, + ArrayRef sendLanes) -> LogicalResult { + if (sendLanes.empty()) + return success(); + ScheduledBoundaryPlan sends = filteredBoundary( + sendLanes, BoundaryActionKind::Send, false); + auto emit = [&]() -> LogicalResult { + auto results = emitClass(sends, sendLanes); + return failed(results) ? failure() : success(); + }; + if (sendLanes.size() == executingLanes.size()) + return emit(); + Value condition = buildClassMembershipCondition( + sendLanes, boundary.actionsByLane.size(), lane, + boundary.scheduled->op, denseConstants, rewriter, + boundary.scheduled->op->getLoc()); + auto selection = scf::IfOp::create( + rewriter, boundary.scheduled->op->getLoc(), TypeRange {}, condition, + false); + OpBuilder::InsertionGuard guard(rewriter); + Region ®ion = selection.getThenRegion(); + auto yield = dyn_cast_or_null( + region.front().empty() ? nullptr : ®ion.front().back()); + if (yield) + rewriter.setInsertionPoint(yield); + else + rewriter.setInsertionPointToStart(®ion.front()); + if (failed(emit())) + return failure(); + if (!yield) + scf::YieldOp::create(rewriter, boundary.scheduled->op->getLoc()); + return success(); + }; + auto emitRemote = [&]() -> FailureOr { + if (failed(emitOptionalSends(remoteLanes, preSendLanes))) + return failure(); + ScheduledBoundaryPlan receives = filteredBoundary( + remoteLanes, BoundaryActionKind::Receive, true); + auto results = emitClass(receives, remoteLanes); + if (failed(results) || results->size() != 1) + return failure(); + if (failed(emitOptionalSends(remoteLanes, postSendLanes))) + return failure(); + return results->front(); + }; + auto emitLocal = [&]() -> FailureOr { + if (failed(emitOptionalSends(localLanes, localSendLanes))) + return failure(); + ScheduledBoundaryPlan local = filteredBoundary( + localLanes, BoundaryActionKind::Receive, true); + auto results = emitClass(local, localLanes); + if (failed(results) || results->size() != 1) + return failure(); + return results->front(); + }; + if (localLanes.empty()) { + auto result = emitRemote(); + if (failed(result)) + return failure(); + return std::optional(*result); + } + + bool remoteIsThen = remoteLanes.size() <= localLanes.size(); + ArrayRef thenLanes = remoteIsThen + ? ArrayRef(remoteLanes) : ArrayRef(localLanes); + Value condition = buildClassMembershipCondition( + thenLanes, boundary.actionsByLane.size(), lane, + boundary.scheduled->op, denseConstants, rewriter, + boundary.scheduled->op->getLoc()); + auto selection = scf::IfOp::create( + rewriter, boundary.scheduled->op->getLoc(), + TypeRange {exchange->deferred.getOutput().getType()}, condition, true); + auto buildBranch = [&](Region ®ion, auto &&emit) -> LogicalResult { + OpBuilder::InsertionGuard guard(rewriter); + auto yield = dyn_cast_or_null( + region.front().empty() ? nullptr : ®ion.front().back()); + if (yield) + rewriter.setInsertionPoint(yield); + else + rewriter.setInsertionPointToStart(®ion.front()); + auto result = emit(); + if (failed(result)) + return failure(); + if (yield) + yield->setOperands(ValueRange {*result}); + else + scf::YieldOp::create( + rewriter, boundary.scheduled->op->getLoc(), ValueRange {*result}); + return success(); + }; + if (failed(buildBranch(selection.getThenRegion(), [&] { + return remoteIsThen ? emitRemote() : emitLocal(); + })) + || failed(buildBranch(selection.getElseRegion(), [&] { + return remoteIsThen ? emitLocal() : emitRemote(); + }))) + return failure(); + return std::optional(selection.getResult(0)); +} + +static LogicalResult realizeBatchBoundaryPart( + ScheduledBoundaryPlan &boundary, IRRewriter &rewriter, + DenseConstantCache &denseConstants, + DeferredEraseSet &erase) { + auto scheduled = cast(boundary.scheduled->op); + Value lane = *scheduled.getLaneArgument(); + SmallVector classes = classifyBatchBoundary(boundary); + auto simple = tryEmitSimpleParentBoundary( + boundary, classes, lane, denseConstants, rewriter); + if (failed(simple)) + return failure(); + if (*simple) { + replaceDeferred( + boundary.targetExchanges.front()->deferred, **simple, erase); + return success(); + } + SmallVector resultTypes; + for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) + resultTypes.push_back(exchange->deferred.getOutput().getType()); + SmallVector replacements; + if (classes.size() == 1) { + auto results = emitBatchBoundaryClass( + boundary, classes.front(), lane, denseConstants, rewriter); + if (failed(results)) + return failure(); + replacements = std::move(*results); + } else if (classes.size() == 2) { + unsigned thenIndex = classes[0].lanes.size() <= classes[1].lanes.size() + ? 0 : 1; + unsigned elseIndex = 1 - thenIndex; + Value condition; + if (classes[thenIndex].lanes.size() == 1) { + Value selectedLane = getOrCreateIndexConstant( + rewriter, scheduled, classes[thenIndex].lanes.front()); + condition = arith::CmpIOp::create( + rewriter, scheduled.getLoc(), arith::CmpIPredicate::eq, lane, + selectedLane); + } else { + SmallVector membership(boundary.actionsByLane.size()); + for (unsigned member : classes[thenIndex].lanes) + membership[member] = 1; + Value selected = buildLaneLookupCached( + membership, lane, scheduled, denseConstants, rewriter, + scheduled.getLoc()); + Value zero = getOrCreateIndexConstant(rewriter, scheduled, 0); + condition = arith::CmpIOp::create( + rewriter, scheduled.getLoc(), arith::CmpIPredicate::ne, selected, + zero); + } + auto selection = scf::IfOp::create( + rewriter, scheduled.getLoc(), resultTypes, condition, true); + auto buildBranch = [&](Region ®ion, const BatchBoundaryClass &group) + -> LogicalResult { + OpBuilder::InsertionGuard guard(rewriter); + auto existingYield = dyn_cast_or_null( + region.front().empty() ? nullptr : ®ion.front().back()); + if (existingYield) + rewriter.setInsertionPoint(existingYield); + else + rewriter.setInsertionPointToStart(®ion.front()); + auto results = emitBatchBoundaryClass( + boundary, group, lane, denseConstants, rewriter); + if (failed(results)) + return failure(); + if (existingYield) + existingYield->setOperands(*results); + else + scf::YieldOp::create(rewriter, scheduled.getLoc(), *results); + return success(); + }; + if (failed(buildBranch(selection.getThenRegion(), classes[thenIndex])) + || failed(buildBranch(selection.getElseRegion(), classes[elseIndex]))) + return failure(); + replacements.assign(selection.getResults().begin(), + selection.getResults().end()); + } else { + llvm::stable_sort(classes, [](const BatchBoundaryClass &lhs, + const BatchBoundaryClass &rhs) { + return lhs.lanes.front() < rhs.lanes.front(); + }); + SmallVector classByLane(boundary.actionsByLane.size()); + for (auto [classIndex, group] : llvm::enumerate(classes)) + for (unsigned member : group.lanes) + classByLane[member] = classIndex; + Value classId = buildLaneLookupCached( + classByLane, lane, scheduled, denseConstants, rewriter, + scheduled.getLoc()); + SmallVector cases; + for (int64_t classIndex = 0; + classIndex < static_cast(classes.size()) - 1; + ++classIndex) + cases.push_back(classIndex); + auto selection = scf::IndexSwitchOp::create( + rewriter, scheduled.getLoc(), resultTypes, classId, cases, + cases.size()); + auto buildRegion = [&](Region ®ion, const BatchBoundaryClass &group) + -> LogicalResult { + OpBuilder::InsertionGuard guard(rewriter); + Block *block = rewriter.createBlock(®ion); + rewriter.setInsertionPointToEnd(block); + auto results = emitBatchBoundaryClass( + boundary, group, lane, denseConstants, rewriter); + if (failed(results)) + return failure(); + scf::YieldOp::create(rewriter, scheduled.getLoc(), *results); + return success(); + }; + for (auto [region, group] : llvm::zip( + selection.getCaseRegions(), ArrayRef(classes).drop_back())) + if (failed(buildRegion(region, group))) + return failure(); + if (failed(buildRegion(selection.getDefaultRegion(), classes.back()))) + return failure(); + replacements.assign(selection.getResults().begin(), + selection.getResults().end()); + } + + if (replacements.size() != boundary.targetExchanges.size()) + return failure(); + for (auto [exchange, replacement] : + llvm::zip_equal(boundary.targetExchanges, replacements)) + replaceDeferred(exchange->deferred, replacement, erase); + return success(); +} + +static LogicalResult realizeBatchBoundary( + ScheduledBoundaryPlan &boundary, IRRewriter &rewriter, + DenseConstantCache &denseConstants, + DeferredEraseSet &erase) { + setInsertionAtBoundary( + rewriter, *boundary.scheduled, boundary.insertionStep); + if (classifyBatchBoundary(boundary).size() <= 2) + return realizeBatchBoundaryPart( + boundary, rewriter, denseConstants, erase); + + SmallVector> orderedParents; + llvm::SmallDenseSet seenParents; + for (ArrayRef actions : boundary.actionsByLane) + for (const BoundaryAction &action : actions) + if (seenParents.insert(action.ref.exchange->exchangeId).second) + orderedParents.push_back( + {action.globalOrder, action.ref.exchange->exchangeId}); + llvm::sort(orderedParents); + + SmallVector parts; + DenseMap partByParent; + for (auto [globalOrder, parent] : orderedParents) { + (void)globalOrder; + ScheduledBoundaryPlan part; + part.scheduled = boundary.scheduled; + part.insertionStep = boundary.insertionStep; + part.actionsByLane.resize(boundary.actionsByLane.size()); + partByParent[parent] = parts.size(); + parts.push_back(std::move(part)); + } + for (auto [lane, actions] : llvm::enumerate(boundary.actionsByLane)) + for (const BoundaryAction &action : actions) + parts[partByParent.lookup(action.ref.exchange->exchangeId)] + .actionsByLane[lane] + .push_back(action); + for (NormalizedDeferredExchange *exchange : boundary.targetExchanges) { + auto it = partByParent.find(exchange->exchangeId); + if (it == partByParent.end()) { + ScheduledBoundaryPlan part; + part.scheduled = boundary.scheduled; + part.insertionStep = boundary.insertionStep; + part.actionsByLane.resize(boundary.actionsByLane.size()); + partByParent[exchange->exchangeId] = parts.size(); + parts.push_back(std::move(part)); + it = partByParent.find(exchange->exchangeId); + } + parts[it->second].targetExchanges.push_back(exchange); + } + + for (ScheduledBoundaryPlan &part : parts) + if (failed(realizeBatchBoundaryPart( + part, rewriter, denseConstants, erase))) + return failure(); return success(); } static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { IRRewriter rewriter(funcOp.getContext()); + auto boundaries = buildScheduledBoundaryPlans(plan); + if (failed(boundaries)) + return funcOp.emitOpError("phase 2 failed to build scheduled boundary plans"); for (ScheduledInfo& info : plan.scheduled) { if (info.isBatch()) { if (failed(linearizeBatch(info, rewriter))) @@ -1041,44 +3397,34 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { } } - SmallVector erase; - // `plan.external` is the scheduler's global order; never recover it from a map. - for (const PlannedCommunicationTransfer &planned : plan.external) { - bool emitted = false; - for (NormalizedDeferredExchange &exchange : plan.exchanges) - for (DeferredSpecialization &specialization : exchange.specializations) - for (FragmentTransfer &transfer : specialization.transfers) - if (!transfer.local && transfer.channelId == planned.exchangeId) { - FragmentRequirement &requirement = - specialization.requirements[transfer.requirementIndex]; - if (requirement.producer->scheduled == exchange.target) { - emitted = true; - continue; - } - if (failed(emitFragmentSend(exchange, specialization, transfer, rewriter))) - return exchange.deferred.emitOpError("phase 2 failed to emit a fragment send"); - emitted = true; - } - if (!emitted) return funcOp.emitOpError("phase 2 cannot locate planned fragment transfer"); + DeferredEraseSet erase; + ReceiveCache receiveCache; + AssemblyCache assemblyCache; + IndexConstantCache indexConstants; + DenseConstantCache denseConstants; + for (Operation &op : *getConstantInsertionBlock(funcOp)) { + auto constant = dyn_cast(&op); + if (!constant) + continue; + if (std::optional value = + matchConstantIndexValue(constant.getResult())) + indexConstants.try_emplace(*value, constant.getResult()); + if (auto typed = dyn_cast(constant.getValue())) + denseConstants.try_emplace(typed, constant.getResult()); + } + for (ScheduledBoundaryPlan &boundary : *boundaries) { + if (boundary.scheduled->isBatch()) { + if (failed(realizeBatchBoundary( + boundary, rewriter, denseConstants, erase))) + return boundary.scheduled->op->emitOpError( + "phase 2 failed to realize a batch communication boundary"); + } else if (failed(realizeScalarBoundary( + boundary, receiveCache, assemblyCache, indexConstants, + denseConstants, rewriter, erase))) { + return boundary.scheduled->op->emitOpError( + "phase 2 failed to realize a scalar communication boundary"); + } } - DenseMap orderByParent; - for (auto [order, transfer] : llvm::enumerate(plan.external)) - orderByParent.try_emplace(transfer.parentExchangeId, order); - for (NormalizedDeferredExchange &exchange : plan.exchanges) - if (exchange.externalTransferCount == 0 && - failed(realizeFragmentExchange(exchange, rewriter, erase))) - return exchange.deferred.emitOpError("phase 2 failed to realize a local fragment exchange"); - SmallVector external; - for (NormalizedDeferredExchange &exchange : plan.exchanges) - if (exchange.externalTransferCount != 0) - external.push_back(&exchange); - llvm::stable_sort(external, [&](NormalizedDeferredExchange *lhs, - NormalizedDeferredExchange *rhs) { - return orderByParent.lookup(lhs->exchangeId) < orderByParent.lookup(rhs->exchangeId); - }); - for (NormalizedDeferredExchange *exchange : external) - if (failed(realizeFragmentExchange(*exchange, rewriter, erase))) - return exchange->deferred.emitOpError("phase 2 failed to realize an external fragment exchange"); for (Operation* op : erase) if (op && op->getBlock() && isa(op)) { if (!op->use_empty()) @@ -1149,6 +3495,18 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { }); for (SpatBlueprintOp blueprint : deadBlueprints) rewriter.eraseOp(blueprint); + DominanceInfo dominance(funcOp); + WalkResult dominanceResult = funcOp.walk([&](Operation *op) { + for (auto [index, operand] : llvm::enumerate(op->getOperands())) + if (!dominance.dominates(operand, op)) { + op->emitOpError() << "phase 2 produced non-dominating operand " + << index << ": " << operand; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (dominanceResult.wasInterrupted()) + return failure(); return success(); } @@ -1172,7 +3530,7 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp) { }); if (hasDeferred) return failure(); - return verifyRealizedCommunicationDeadlockFree(funcOp); + return success(); } } // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp index 2eee59c..1122882 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp @@ -1,9 +1,11 @@ #include "DeferredProjectionAnalysis.hpp" -#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/Matchers.h" -#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" +#include "llvm/ADT/SmallPtrSet.h" +#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" #include @@ -11,43 +13,83 @@ namespace onnx_mlir::spatial { using namespace mlir; namespace { +static FailureOr getSignedInt64(IntegerAttr value) { + return value.getValue().isSignedIntN(64) ? FailureOr(value.getValue().getSExtValue()) + : FailureOr(failure()); +} + +static FailureOr evaluate(Value value, const StaticIndexEnvironment &environment, + llvm::SmallDenseSet &visiting); + +static FailureOr evaluateDenseExtract(tensor::ExtractOp extract, + const StaticIndexEnvironment &environment, + llvm::SmallDenseSet &visiting) { + auto constant = extract.getTensor().getDefiningOp(); + auto elements = constant ? dyn_cast(constant.getValue()) : DenseIntElementsAttr(); + auto type = elements ? dyn_cast(elements.getType()) : RankedTensorType(); + if (!elements || !type || !type.hasStaticShape() + || extract.getIndices().size() != static_cast(type.getRank())) + return failure(); + int64_t linear = 0; + for (auto [index, dim] : llvm::zip(extract.getIndices(), type.getShape())) { + auto folded = evaluate(index, environment, visiting); + if (failed(folded) || *folded < 0 || *folded >= dim + || llvm::MulOverflow(linear, dim, linear) || llvm::AddOverflow(linear, *folded, linear)) + return failure(); + } + APInt value = elements.getValues()[linear]; + return value.isSignedIntN(64) ? FailureOr(value.getSExtValue()) + : FailureOr(failure()); +} + static FailureOr evaluate(Value value, const StaticIndexEnvironment &environment, llvm::SmallDenseSet &visiting) { if (auto it = environment.bindings.find(value); it != environment.bindings.end()) return it->second; + Attribute constant; + if (matchPattern(value, m_Constant(&constant))) + if (auto integer = dyn_cast_or_null(constant)) + return getSignedInt64(integer); + if (isa(value) || !value.getDefiningOp()) + return failure(); if (!visiting.insert(value).second) return failure(); - if (auto constant = value.getDefiningOp()) - if (auto integer = dyn_cast(constant.getValue())) - { visiting.erase(value); return integer.getInt(); } - if (auto cast = value.getDefiningOp()) { - auto result = evaluate(cast.getIn(), environment, visiting); visiting.erase(value); return result; - } - auto binary = [&](auto op, auto fn) -> FailureOr { - auto lhs = evaluate(op.getLhs(), environment, visiting); - auto rhs = evaluate(op.getRhs(), environment, visiting); - if (failed(lhs) || failed(rhs)) return failure(); - return fn(*lhs, *rhs); - }; - if (auto op = value.getDefiningOp()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr { int64_t r; if (llvm::AddOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; } - if (auto op = value.getDefiningOp()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr { int64_t r; if (llvm::SubOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; } - if (auto op = value.getDefiningOp()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr { int64_t r; if (llvm::MulOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; } - if (auto apply = value.getDefiningOp()) { auto result = evaluateAffineApply(apply, [&](Value operand) { return evaluate(operand, environment, visiting); }); visiting.erase(value); return result; } 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 failure(); - int64_t linear = 0; - for (auto [index, dim] : llvm::zip(extract.getIndices(), type.getShape())) { - auto i = evaluate(index, environment, visiting); - if (failed(i) || *i < 0 || *i >= dim) return failure(); - linear = linear * dim + *i; + auto result = evaluateDenseExtract(extract, environment, visiting); + visiting.erase(value); + return result; + } + + Operation *definingOp = value.getDefiningOp(); + if (definingOp->getNumRegions() != 0 || !isPureIndexComputationOp(definingOp)) { + visiting.erase(value); + return failure(); + } + SmallVector operandConstants; + operandConstants.reserve(definingOp->getNumOperands()); + Builder builder(definingOp->getContext()); + for (Value operand : definingOp->getOperands()) { + auto folded = evaluate(operand, environment, visiting); + if (failed(folded)) { + visiting.erase(value); + return failure(); } - visiting.erase(value); return elements.getValues()[linear].getSExtValue(); + operandConstants.push_back(builder.getIntegerAttr(operand.getType(), *folded)); + } + SmallVector results; + if (failed(definingOp->fold(operandConstants, results)) || results.size() != 1) { + visiting.erase(value); + return failure(); + } + FailureOr result = failure(); + if (auto integer = dyn_cast(results.front())) { + if (auto attr = dyn_cast(integer)) + result = getSignedInt64(attr); + } else if (auto foldedValue = dyn_cast(results.front())) { + result = evaluate(foldedValue, environment, visiting); } visiting.erase(value); - return failure(); + return result; } static FailureOr> sourceArgument(Value value, SpatDeferredCommunicationOp deferred, @@ -57,50 +99,29 @@ static FailureOr> sourceArgument(Value value, SpatDeferr argument && argument.getOwner() == &deferred.getBody().front() && argument.getArgNumber() < deferred.getSources().size()) return std::optional(argument.getArgNumber()); - // Phase 1's selector ends in collapse(extract_slice(stacked, table[lane])). - auto collapse = value.getDefiningOp(); - if (!collapse) return std::optional(); - value = collapse.getSrc(); - auto slice = value.getDefiningOp(); - if (!slice) return std::optional(); - auto sourceType = dyn_cast(slice.getSourceType()); - if (!sourceType || slice.getMixedOffsets().size() != static_cast(sourceType.getRank())) + auto result = dyn_cast(value); + auto selection = result ? dyn_cast(result.getOwner()) : scf::IndexSwitchOp(); + if (!selection || result.getResultNumber() != 0 || selection.getNumResults() != 1) return std::optional(); - for (unsigned dim = 1; dim < sourceType.getRank(); ++dim) { - auto offset = evaluateDeferredIndex(slice.getMixedOffsets()[dim], environment); - auto size = evaluateDeferredIndex(slice.getMixedSizes()[dim], environment); - auto stride = evaluateDeferredIndex(slice.getMixedStrides()[dim], environment); - if (failed(offset) || failed(size) || failed(stride) || *offset != 0 || *size != sourceType.getDimSize(dim) || *stride != 1) - return std::optional(); - } - auto leadingSize = evaluateDeferredIndex(slice.getMixedSizes().front(), environment); - auto leadingStride = evaluateDeferredIndex(slice.getMixedStrides().front(), environment); - if (failed(leadingSize) || failed(leadingStride) || *leadingSize != 1 || *leadingStride != 1) - return std::optional(); - auto selected = evaluateDeferredIndex(slice.getMixedOffsets().front(), environment); - if (failed(selected)) return failure(); - Value stacked = slice.getSource(); - while (auto cast = stacked.getDefiningOp()) stacked = cast.getSource(); - while (auto insert = stacked.getDefiningOp()) stacked = insert.getDest(); - // The stack is a chain. Find the insertion at the selected leading offset. - for (Value cursor = slice.getSource(); auto insert = cursor.getDefiningOp(); cursor = insert.getDest()) { - auto offset = evaluateDeferredIndex(insert.getMixedOffsets().front(), environment); - if (succeeded(offset) && *offset == *selected) { - Value source = insert.getSource(); - if (auto expand = source.getDefiningOp()) source = expand.getSrc(); - if (auto arg = dyn_cast(source); arg && arg.getOwner() == &deferred.getBody().front()) - return std::optional(arg.getArgNumber()); + auto selector = evaluateDeferredIndex(selection.getArg(), environment); + if (failed(selector)) + return failure(); + Region *selectedRegion = &selection.getDefaultRegion(); + for (auto [caseValue, region] : llvm::zip(selection.getCases(), selection.getCaseRegions())) + if (caseValue == *selector) { + selectedRegion = ®ion; + break; } - } - return std::optional(); -} - -static bool isResidual(Operation *op) { - return isa(op); + if (!selectedRegion->hasOneBlock()) + return failure(); + Block &block = selectedRegion->front(); + auto yield = dyn_cast(block.getTerminator()); + if (!yield || yield.getResults().size() != 1) + return failure(); + for (Operation &op : block.without_terminator()) + if (!isa(op)) + return failure(); + return sourceArgument(yield.getResults().front(), deferred, environment); } static SpatGraphComputeBatch graphBatchOwner(Value value) { @@ -119,15 +140,394 @@ static Value getEnclosingScheduledLane(SpatDeferredCommunicationOp deferred, return block && !block->empty() ? block->getArgument(0) : Value(); } +static bool isInsideDeferredLoop(Operation *op, + SpatDeferredCommunicationOp deferred) { + for (Operation *parent = op->getParentOp(); parent && parent != deferred; + parent = parent->getParentOp()) + if (isa(parent)) + return true; + return false; +} + +static bool isAllowedStaticIndexExpression( + Value value, Value scheduledLane, + llvm::SmallDenseSet &visiting) { + if (value == scheduledLane) + return true; + Attribute constant; + if (matchPattern(value, m_Constant(&constant))) + return true; + if (isa(value) || !value.getDefiningOp() + || !visiting.insert(value).second) + return false; + Operation *op = value.getDefiningOp(); + bool allowed = op->getNumRegions() == 0 + && (isPureIndexComputationOp(op) || isCompileTimeOp(op)) + && llvm::all_of(op->getOperands(), [&](Value operand) { + return isAllowedStaticIndexExpression(operand, scheduledLane, + visiting); + }); + visiting.erase(value); + return allowed; +} + +static bool isAllowedStaticIndexExpression(Value value, Value scheduledLane) { + llvm::SmallDenseSet visiting; + return isAllowedStaticIndexExpression(value, scheduledLane, visiting); +} + +static bool originatesFromDeferredSource( + Value value, SpatDeferredCommunicationOp deferred, + llvm::SmallDenseSet &visited) { + if (!visited.insert(value).second) + return false; + if (auto argument = dyn_cast(value)) + return argument.getOwner() == &deferred.getBody().front() + && argument.getArgNumber() < deferred.getSources().size(); + Operation *op = value.getDefiningOp(); + if (!op) + return false; + if (isa(op) && op->getBlock() == &deferred.getBody().front()) + return true; + return llvm::any_of(op->getOperands(), [&](Value operand) { + return originatesFromDeferredSource(operand, deferred, visited); + }); +} + +static bool originatesFromDeferredSource( + Value value, SpatDeferredCommunicationOp deferred) { + llvm::SmallDenseSet visited; + return originatesFromDeferredSource(value, deferred, visited); +} + +static LogicalResult verifyCanonicalSourceSelector( + scf::IndexSwitchOp selection, SpatDeferredCommunicationOp deferred, + Value scheduledLane, int64_t laneCount) { + if (selection->getBlock() != &deferred.getBody().front() + || selection.getNumResults() != 1 + || !selection.getArg().getType().isIndex() + || !isAllowedStaticIndexExpression(selection.getArg(), scheduledLane)) + return selection.emitOpError("is not a canonical deferred source selector"); + if (laneCount < 2 + || selection.getCases().size() != static_cast(laneCount - 1) + || selection.getCaseRegions().size() != selection.getCases().size()) + return selection.emitOpError( + "must cover every non-default scheduled lane"); + for (auto [index, caseValue] : llvm::enumerate(selection.getCases())) + if (caseValue != static_cast(index)) + return selection.emitOpError( + "must use consecutive scheduled-lane cases starting at zero"); + + auto verifyRegion = [&](Region ®ion) -> LogicalResult { + if (!region.hasOneBlock()) + return selection.emitOpError("source-selector region must have one block"); + Block &block = region.front(); + auto yield = dyn_cast(block.getTerminator()); + if (!yield || yield.getResults().size() != 1 + || yield.getResults().front().getType() != selection.getResult(0).getType()) + return selection.emitOpError( + "source-selector region must yield one exact result type"); + for (Operation &op : block.without_terminator()) + if (!isa(op)) + return selection.emitOpError( + "source-selector regions may contain only tensor.cast before scf.yield"); + Value source = yield.getResults().front(); + while (auto cast = source.getDefiningOp()) { + if (cast->getBlock() != &block) + return selection.emitOpError( + "source-selector casts must be local to their region"); + source = cast.getSource(); + } + auto argument = dyn_cast(source); + if (!argument || argument.getOwner() != &deferred.getBody().front() + || argument.getArgNumber() >= deferred.getSources().size()) + return selection.emitOpError( + "source-selector branch must resolve to a deferred source argument"); + return success(); + }; + for (Region ®ion : selection.getCaseRegions()) + if (failed(verifyRegion(region))) + return failure(); + return verifyRegion(selection.getDefaultRegion()); +} + +static bool valueDependsOnAny( + Value value, const llvm::SmallDenseSet &dependencies, + llvm::SmallDenseSet &visited) { + if (dependencies.contains(value)) + return true; + if (!visited.insert(value).second) + return false; + Operation *op = value.getDefiningOp(); + return op && llvm::any_of(op->getOperands(), [&](Value operand) { + return valueDependsOnAny(operand, dependencies, visited); + }); +} + +static bool valueDependsOnAny( + Value value, const llvm::SmallDenseSet &dependencies) { + llvm::SmallDenseSet visited; + return valueDependsOnAny(value, dependencies, visited); +} + +static LogicalResult specializeDeferredLoopStaticValues( + scf::ForOp loop, SpatDeferredCommunicationOp deferred, + const StaticIndexEnvironment &environment, + SpecializedDeferredProgram &program, + llvm::SmallDenseSet dynamicIndices = {}) { + auto record = [&](Value value, StringRef diagnostic) -> LogicalResult { + auto folded = evaluateDeferredIndex(value, environment); + if (failed(folded)) + return deferred.emitOpError(diagnostic); + program.staticValues.try_emplace(value, *folded); + return success(); + }; + auto step = evaluateDeferredIndex(loop.getStep(), environment); + if (failed(record(loop.getLowerBound(), + "deferred shaping loop lower bound did not specialize")) + || failed(record(loop.getUpperBound(), + "deferred shaping loop upper bound did not specialize")) + || failed(step) || *step <= 0) + return deferred.emitOpError( + "deferred shaping loop requires specialized bounds and a positive step"); + program.staticValues.try_emplace(loop.getStep(), *step); + + dynamicIndices.insert(loop.getInductionVar()); + for (Operation &nested : loop.getBody()->without_terminator()) { + if (auto nestedLoop = dyn_cast(nested)) { + if (failed(specializeDeferredLoopStaticValues( + nestedLoop, deferred, environment, program, dynamicIndices))) + return failure(); + continue; + } + for (Value operand : nested.getOperands()) { + if ((!operand.getType().isIndex() + && !isa(operand.getType())) + || valueDependsOnAny(operand, dynamicIndices)) + continue; + if (failed(record(operand, + "deferred shaping loop captured an index that did not specialize"))) + return failure(); + } + } + return success(); +} + +static FailureOr> +analyzeDeferredInsertAssembly( + const SpecializedDeferredProgram &program, + const StaticIndexEnvironment &environment) { + auto finalInsert = program.yieldedValue.getDefiningOp(); + if (!finalInsert || program.leaves.empty()) + return std::optional(); + + DenseMap leafByRoot; + for (auto [leafIndex, leaf] : llvm::enumerate(program.leaves)) { + if (leaf.physicalSlots.size() != 1 + || !leafByRoot.try_emplace(leaf.replacementRoot, leafIndex).second) + return std::optional(); + } + + SmallPtrSet consumed; + SmallVector reverseEntries; + llvm::SmallDenseSet insertedLeaves; + Value current = program.yieldedValue; + while (auto insert = current.getDefiningOp()) { + Value source = insert.getSource(); + SmallVector rankShaping; + while (Operation *shape = source.getDefiningOp()) { + if (auto collapse = dyn_cast(shape)) + source = collapse.getSrc(); + else if (auto expand = dyn_cast(shape)) + source = expand.getSrc(); + else + break; + rankShaping.push_back(shape); + } + auto leafIt = leafByRoot.find(source); + if (leafIt == leafByRoot.end() + || !insertedLeaves.insert(leafIt->second).second) + return std::optional(); + + DeferredInsertAssemblyEntry entry; + entry.leafIndex = leafIt->second; + entry.requirementIndex = leafIt->second; + auto foldGeometry = [&](ArrayRef values, + SmallVectorImpl &folded) { + for (OpFoldResult value : values) { + auto result = evaluateDeferredIndex(value, environment); + if (failed(result)) + return failure(); + folded.push_back(*result); + } + return success(); + }; + if (failed(foldGeometry(insert.getMixedOffsets(), + entry.targetGeometry.offsets)) + || failed(foldGeometry(insert.getMixedSizes(), + entry.targetGeometry.sizes)) + || failed(foldGeometry(insert.getMixedStrides(), + entry.targetGeometry.strides))) + return std::optional(); + reverseEntries.push_back(std::move(entry)); + consumed.insert(insert); + consumed.insert(rankShaping.begin(), rankShaping.end()); + current = insert.getDest(); + } + + auto initial = current.getDefiningOp(); + auto resultType = dyn_cast(program.yieldedValue.getType()); + if (!initial || !initial.getDynamicSizes().empty() || !resultType + || !resultType.hasStaticShape() || initial.getType() != resultType + || insertedLeaves.size() != program.leaves.size()) + return std::optional(); + consumed.insert(initial); + + for (const DeferredInsertAssemblyEntry &entry : reverseEntries) { + const StaticSliceGeometry &geometry = entry.targetGeometry; + if (geometry.offsets.size() != static_cast(resultType.getRank()) + || geometry.sizes.size() != geometry.offsets.size() + || geometry.strides.size() != geometry.offsets.size()) + return std::optional(); + for (auto [offset, size, stride, dim] : + llvm::zip_equal(geometry.offsets, geometry.sizes, + geometry.strides, resultType.getShape())) { + int64_t span = 0; + int64_t last = 0; + if (offset < 0 || size <= 0 || stride <= 0 + || llvm::MulOverflow(size - 1, stride, span) + || llvm::AddOverflow(offset, span, last) || last >= dim) + return std::optional(); + } + } + if (llvm::any_of(program.residualOps, + [&](Operation *op) { return !consumed.contains(op); }) + || consumed.size() != program.residualOps.size()) + return std::optional(); + + DeferredInsertAssembly assembly; + assembly.initialValue = initial; + assembly.resultType = resultType; + assembly.entries.assign(reverseEntries.rbegin(), reverseEntries.rend()); + return std::optional(std::move(assembly)); +} + } // namespace +LogicalResult verifyDeferredProgramContract( + SpatDeferredCommunicationOp deferred) { + if (!deferred.getBody().hasOneBlock()) + return deferred.emitOpError("deferred program must have exactly one body block"); + Block &body = deferred.getBody().front(); + auto terminator = dyn_cast(body.getTerminator()); + if (!terminator || terminator.getOutputs().size() != 1) + return deferred.emitOpError( + "deferred program must have exactly one yielded value"); + + Value scheduledLane; + int64_t laneCount = 1; + if (auto scheduled = deferred->getParentOfType()) { + scheduledLane = getEnclosingScheduledLane(deferred, scheduled); + laneCount = scheduled.getLaneCount(); + if (!scheduledLane || laneCount <= 0) + return deferred.emitOpError( + "deferred program cannot locate a valid enclosing scheduled lane"); + } + + bool invalid = false; + WalkResult result = deferred.getBody().walk([&](Operation *op) { + if (invalid) + return WalkResult::interrupt(); + auto reject = [&](StringRef message) { + op->emitOpError(message); + invalid = true; + return WalkResult::interrupt(); + }; + + if (auto yield = dyn_cast(op)) { + if (op != body.getTerminator()) + return reject("must be the unique top-level deferred terminator"); + return WalkResult::advance(); + } + if (auto yield = dyn_cast(op)) { + Operation *parent = op->getParentOp(); + if (!parent || !isa(parent) + || op != op->getBlock()->getTerminator()) + return reject("is not a valid deferred control-flow terminator"); + return WalkResult::advance(); + } + if (auto selection = dyn_cast(op)) { + if (failed(verifyCanonicalSourceSelector( + selection, deferred, scheduledLane, laneCount))) { + invalid = true; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (auto loop = dyn_cast(op)) { + auto isStaticRankedTensor = [](Type type) { + auto tensor = dyn_cast(type); + return tensor && tensor.hasStaticShape(); + }; + if (!llvm::all_of(loop.getInitArgs(), [&](Value value) { + return isStaticRankedTensor(value.getType()); + }) + || !llvm::all_of(loop.getResultTypes(), isStaticRankedTensor)) + return reject( + "requires static ranked tensor iter arguments and results"); + auto yield = dyn_cast(loop.getBody()->getTerminator()); + if (!yield || yield.getResults().getTypes() != loop.getResultTypes()) + return reject("yield types must exactly match loop result types"); + for (Value bound : {loop.getLowerBound(), loop.getUpperBound(), + loop.getStep()}) + if (!isAllowedStaticIndexExpression(bound, scheduledLane)) + return reject( + "bounds must use only constants, the scheduled lane, and pure index computation"); + for (int64_t lane = 0; lane < laneCount; ++lane) { + StaticIndexEnvironment environment; + if (scheduledLane) + environment.bindings[scheduledLane] = lane; + auto lower = evaluateDeferredIndex(loop.getLowerBound(), environment); + auto upper = evaluateDeferredIndex(loop.getUpperBound(), environment); + auto step = evaluateDeferredIndex(loop.getStep(), environment); + if (failed(lower) || failed(upper) || failed(step) || *step <= 0) + return reject( + "bounds must specialize for every scheduled lane with a positive step"); + } + return WalkResult::advance(); + } + if (op->getNumRegions() != 0) + return reject("unsupported region operation in deferred program"); + if (!isShapingOnlyOp(op) && !isCompileTimeOp(op) + && !isPureIndexComputationOp(op)) + return reject( + "deferred program permits only shaping, compile-time, or pure index operations"); + + for (Value operand : op->getOperands()) { + if (auto argument = dyn_cast(operand)) { + Operation *owner = argument.getOwner()->getParentOp(); + bool local = owner == deferred + || (owner && deferred->isAncestor(owner)); + if (!local && operand != scheduledLane) + return reject("captures an unsupported external block argument"); + } + if (isInsideDeferredLoop(op, deferred) + && originatesFromDeferredSource(operand, deferred)) + return reject( + "deferred source projection must remain outside residual loops"); + } + return WalkResult::advance(); + }); + return success(!invalid && !result.wasInterrupted()); +} + FailureOr evaluateDeferredIndex(Value value, const StaticIndexEnvironment &environment) { llvm::SmallDenseSet visiting; return evaluate(value, environment, visiting); } FailureOr evaluateDeferredIndex(OpFoldResult value, const StaticIndexEnvironment &environment) { if (auto attr = dyn_cast(value)) - if (auto integer = dyn_cast(attr)) return integer.getInt(); + if (auto integer = dyn_cast(attr)) return getSignedInt64(integer); if (auto dynamic = dyn_cast(value)) return evaluateDeferredIndex(dynamic, environment); return failure(); } @@ -152,6 +552,8 @@ FailureOr requireResolvedDeferredSource(Value value, Spa FailureOr analyzeDeferredProgram(SpatDeferredCommunicationOp deferred, std::optional targetScheduledLane) { + // The Phase 1 verifier must establish the deferred-program contract once; + // this analysis only specializes lane-dependent static semantics. Block &body = deferred.getBody().front(); auto yield = dyn_cast(body.getTerminator()); if (!yield || yield.getOutputs().size() != 1) @@ -229,13 +631,37 @@ FailureOr analyzeDeferredProgram(SpatDeferredCommuni program.leaves.push_back(std::move(leaf)); return success(); } + if (value.getType().isIndex() || isa(value.getType())) { + auto folded = evaluateDeferredIndex(value, environment); + if (failed(folded)) + return deferred.emitOpError("deferred index expression is not statically evaluable after specialization"); + program.staticValues.try_emplace(value, *folded); + return success(); + } Operation *op = value.getDefiningOp(); - if (!op || op->getBlock() != &body || !isResidual(op)) - return deferred.emitOpError("deferred residual contains an unsupported operation"); + if (!op || op->getBlock() != &body) + return deferred.emitOpError("deferred residual escapes its verified body"); + if (auto loop = dyn_cast(op)) { + if (failed(specializeDeferredLoopStaticValues( + loop, deferred, environment, program))) + return failure(); + for (Value init : loop.getInitArgs()) + if (failed(visit(init))) + return failure(); + program.residualOps.push_back(op); + return success(); + } for (Value operand : op->getOperands()) if (failed(visit(operand))) return failure(); program.residualOps.push_back(op); return success(); }; if (failed(visit(program.yieldedValue))) return failure(); + StaticIndexEnvironment assemblyEnvironment = environment; + for (auto [value, staticValue] : program.staticValues) + assemblyEnvironment.bindings.try_emplace(value, staticValue); + auto assembly = analyzeDeferredInsertAssembly(program, assemblyEnvironment); + if (failed(assembly)) + return failure(); + program.insertAssembly = std::move(*assembly); return std::move(program); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp index 2c9e4fd..0fc418b 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp @@ -35,6 +35,18 @@ struct DeferredProjectionLeaf { mlir::RankedTensorType reconstructedType; }; +struct DeferredInsertAssemblyEntry { + unsigned requirementIndex = 0; + unsigned leafIndex = 0; + StaticSliceGeometry targetGeometry; +}; + +struct DeferredInsertAssembly { + mlir::tensor::EmptyOp initialValue; + mlir::RankedTensorType resultType; + llvm::SmallVector entries; +}; + struct SpecializedDeferredProgram { SpatDeferredCommunicationOp deferred; std::optional targetScheduledLane; @@ -42,6 +54,8 @@ struct SpecializedDeferredProgram { mlir::Value yieldedValue; llvm::SmallVector leaves; llvm::SmallVector residualOps; + llvm::DenseMap staticValues; + std::optional insertAssembly; }; struct ResolvedDeferredSource { @@ -55,6 +69,8 @@ mlir::FailureOr> tryResolveDeferredSource( mlir::FailureOr requireResolvedDeferredSource( mlir::Value value, SpatDeferredCommunicationOp deferred, const StaticIndexEnvironment &environment); +mlir::LogicalResult verifyDeferredProgramContract( + SpatDeferredCommunicationOp deferred); mlir::FailureOr analyzeDeferredProgram( SpatDeferredCommunicationOp deferred, std::optional targetScheduledLane); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp index ee6090d..53f2488 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp @@ -1,6 +1,7 @@ #include "ScheduledComputeMaterialization.hpp" #include "ScheduledComputeReport.hpp" #include "ScheduledComputeVerification.hpp" +#include "SpatialDataflowCsvExporter.hpp" #include "DeferredCommunicationRealization.hpp" #include "DeferredCommunicationDeadlock.hpp" @@ -82,6 +83,13 @@ struct MergeComputeNodesPass final : PassWrapper buildSourceLaneStartForScheduledLane(OpBuilder &builder, Value scheduledLane, const SourceLaneSelector &selector, Operation *constantAnchor) { - (void)constantAnchor; if (selector.kind == SourceLaneSelector::Kind::Affine) { if (selector.affine.baseLaneStart == 0 && selector.affine.laneCount == 1) return scheduledLane; @@ -214,9 +213,7 @@ static FailureOr buildSourceLaneStartForScheduledLane(OpBuilder &builder, expr = d0 * selector.affine.laneCount; if (selector.affine.baseLaneStart != 0) expr = expr + selector.affine.baseLaneStart; - return affine::AffineApplyOp::create( - builder, loc, AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr), ValueRange {scheduledLane}) - .getResult(); + return createOrFoldAffineApply(builder, loc, expr, ValueRange {scheduledLane}, constantAnchor); } if (!selector.table) @@ -569,18 +566,13 @@ static SmallVector buildScheduledOutputInsertOffsets(OpBuilder &bu Location loc, Value scheduledLane, int64_t lanesPerScheduledLane, - RankedTensorType localFragmentType) { + RankedTensorType localFragmentType, + Operation *constantAnchor) { SmallVector offsets; Value scheduledOutputLane = scheduledLane; if (lanesPerScheduledLane != 1) { - scheduledOutputLane = - affine::AffineApplyOp::create(builder, - loc, - AffineMap::get(/*dimCount=*/1, - /*symbolCount=*/0, - builder.getAffineDimExpr(0) * lanesPerScheduledLane), - ValueRange {scheduledLane}) - .getResult(); + scheduledOutputLane = affineMulConst( + builder, loc, scheduledLane, lanesPerScheduledLane, constantAnchor); } offsets.push_back(scheduledOutputLane); offsets.append(localFragmentType.getRank() - 1, OpFoldResult(builder.getIndexAttr(0))); @@ -707,14 +699,12 @@ static LogicalResult materializeMultiCpuPeftClass( [&](OpBuilder &builder, Location bodyLoc, Value innerLane, ValueRange iterArgs, SmallVectorImpl &yielded) -> LogicalResult { IRMapping mapper; - Value sourceLane = - affine::AffineApplyOp::create(builder, - bodyLoc, - AffineMap::get(/*dimCount=*/2, - /*symbolCount=*/0, - builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1)), - ValueRange {*sourceLaneStart, innerLane}) - .getResult(); + Value sourceLane = createOrFoldAffineApply( + builder, + bodyLoc, + builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1), + ValueRange {*sourceLaneStart, innerLane}, + scheduled.getOperation()); mapper.map(*batch.getLaneArgument(), sourceLane); for (auto [index, weight] : llvm::enumerate(batch.getWeights())) mapper.map(*batch.getWeightArgument(index), @@ -789,8 +779,13 @@ static LogicalResult materializeMultiCpuPeftClass( finalLocalFragments.assign(loop->results.begin(), loop->results.end()); } - auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc()); - rewriter.setInsertionPointToStart(&inParallel.getRegion().front()); + struct Publication { + Value fragment; + SmallVector offsets; + SmallVector sizes; + SmallVector strides; + }; + SmallVector publications; for (auto [resultIndex, localFragment] : llvm::enumerate(finalLocalFragments)) { auto localFragmentType = cast(localFragment.getType()); int64_t lanesPerScheduledLane = isa(representative.op) ? 1 : representative.laneCount; @@ -799,22 +794,29 @@ static LogicalResult materializeMultiCpuPeftClass( scheduled.getLoc(), scheduledLane, lanesPerScheduledLane, - localFragmentType); + localFragmentType, + scheduled.getOperation()); SmallVector sizes; SmallVector strides; for (int64_t dim : localFragmentType.getShape()) { sizes.push_back(rewriter.getIndexAttr(dim)); strides.push_back(rewriter.getIndexAttr(1)); } + publications.push_back( + {localFragment, std::move(offsets), std::move(sizes), + std::move(strides)}); + } + auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc()); + rewriter.setInsertionPointToStart(&inParallel.getRegion().front()); + for (auto [resultIndex, publication] : llvm::enumerate(publications)) tensor::ParallelInsertSliceOp::create( rewriter, scheduled.getLoc(), - localFragment, + publication.fragment, block->getArgument(getScheduledBatchResultArgBase(scheduled) + stepPlan.resultOffset + resultIndex), - offsets, - sizes, - strides); - } + publication.offsets, + publication.sizes, + publication.strides); } return success(); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp index f9f48a4..aac65f0 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp @@ -22,6 +22,7 @@ #include "Scheduling/ComputeInstanceUtils.hpp" #include "Scheduling/MergeSchedulingAnalysis.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" namespace onnx_mlir { namespace spatial { @@ -267,11 +268,9 @@ inline SmallVector collectSourceLaneStarts(const ComputeStepTuple &ste } inline Value createI64LookupTableConstant(OpBuilder &builder, Operation *constantAnchor, ArrayRef values) { - OpBuilder::InsertionGuard guard(builder); - builder.setInsertionPoint(constantAnchor); RankedTensorType tableType = RankedTensorType::get({static_cast(values.size())}, builder.getI64Type()); DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values); - return arith::ConstantOp::create(builder, constantAnchor->getLoc(), tableType, tableAttr).getResult(); + return getOrCreateConstant(builder, constantAnchor, tableAttr, tableType); } inline FailureOr createEmptyTensorForType(OpBuilder &builder, Location loc, Type type) { diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp index fe5f20f..868b7bd 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp @@ -82,24 +82,36 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) { pim::CappedDiagnosticReporter diagnostics; GraphBatchPublicationCache publicationCache; funcOp.walk([&](SpatDeferredCommunicationOp transfer) { + bool ownershipValid = true; for (Value source : transfer.getSources()) { auto result = dyn_cast(source); if (!result || !isa(result.getOwner())) { + ownershipValid = false; diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { illegalOp->emitOpError("phase-check deferred communication source operand must be an original graph SSA result"); }); } } if (!transfer->getParentOfType() && - !transfer->getParentOfType()) + !transfer->getParentOfType()) { + ownershipValid = false; diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { illegalOp->emitOpError("phase-check deferred communication must be inside a scheduled compute"); }); + } + if (!ownershipValid) + return; + if (failed(verifyDeferredProgramContract(transfer))) { + diagnostics.report(transfer.getOperation(), [&](Operation *) {}); + return; + } if (auto scheduled = transfer->getParentOfType()) { for (unsigned lane = 0; lane < static_cast(scheduled.getLaneCount()); ++lane) { auto program = analyzeDeferredProgram(transfer, lane); - if (failed(program)) + if (failed(program)) { + diagnostics.report(transfer.getOperation(), [&](Operation *) {}); continue; + } for (const DeferredProjectionLeaf &leaf : program->leaves) { if (leaf.kind == DeferredLeafKind::ScalarSource) continue; @@ -110,8 +122,9 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) { diagnostics.report(transfer.getOperation(), [&](Operation *) {}); } } - } else - (void)analyzeDeferredProgram(transfer, std::nullopt); + } else if (failed(analyzeDeferredProgram(transfer, std::nullopt))) { + diagnostics.report(transfer.getOperation(), [&](Operation *) {}); + } }); diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed"); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp new file mode 100644 index 0000000..3a02086 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp @@ -0,0 +1,898 @@ +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/AsmState.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Value.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include +#include +#include + +#include "SpatialDataflowCsvExporter.hpp" +#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +using namespace mlir; + +namespace onnx_mlir { +namespace spatial { + +namespace { + +struct TopLevelOpInfo { + Operation* op = nullptr; + size_t opId = 0; + bool isScheduled = false; + std::optional scalarCore; +}; + +struct ExpandedNodeInfo { + std::string id; + std::optional core; + std::optional lane; +}; + +struct ChannelSendRecord { + std::string sourceId; + std::optional sourceLane; +}; + +enum class LogicalNodeSelector { + Scalar, + Lane, + RangeRepresentative, +}; + +struct ResolvedProducer { + Operation* op = nullptr; + size_t resultIndex = 0; + LogicalNodeSelector selector = LogicalNodeSelector::Scalar; + uint32_t lane = 0; + uint32_t laneStart = 0; + uint32_t laneCount = 1; +}; + +struct EdgeSource { + std::string id; + std::optional sourceLane; +}; + +using ScheduledNodeByGraphLane = DenseMap, ExpandedNodeInfo>; + +void emitEdgeRow(std::fstream& edgesFile, + StringRef sourceId, + StringRef targetId, + std::optional byteSize, + Type propagatedType, + StringRef stage, + std::optional sourceLane, + std::optional targetLane, + std::optional channelId); + +std::string csvEscape(StringRef field) { + bool needsQuotes = field.contains(',') || field.contains('"') || field.contains('\n') || field.contains('\r'); + if (!needsQuotes) + return field.str(); + + std::string escaped; + escaped.reserve(field.size() + 2); + escaped.push_back('"'); + for (char ch : field) + if (ch == '"') + escaped += "\"\""; + else + escaped.push_back(ch); + escaped.push_back('"'); + return escaped; +} + +void writeCsvRow(std::fstream& file, ArrayRef fields) { + for (size_t i = 0; i < fields.size(); ++i) { + if (i != 0) + file << ","; + file << csvEscape(fields[i]); + } + file << "\n"; +} + +template +std::string maybeNumber(std::optional value) { + if (!value) + return ""; + return std::to_string(*value); +} + +std::string stringifyType(Type type) { + std::string storage; + llvm::raw_string_ostream os(storage); + type.print(os); + return os.str(); +} + +std::string stringifyValueAsOperand(Value value, AsmState& asmState) { + std::string storage; + llvm::raw_string_ostream os(storage); + value.printAsOperand(os, asmState); + return os.str(); +} + +std::string stringifyResultSsaNames(Operation* op, AsmState* asmState) { + if (!asmState || op->getNumResults() == 0) + return ""; + + std::string storage; + llvm::raw_string_ostream os(storage); + llvm::interleave( + op->getResults(), [&](Value result) { os << stringifyValueAsOperand(result, *asmState); }, [&]() { os << ";"; }); + return os.str(); +} + +std::optional getTypeSizeBytes(Type type) { + if (auto shapedType = dyn_cast(type)) { + if (!shapedType.hasStaticShape() || !hasByteSizedElementType(shapedType.getElementType())) + return std::nullopt; + return static_cast(getShapedTypeSizeInBytes(shapedType)); + } + + if (isa(type)) + return static_cast(getElementTypeSizeInBytes(type)); + if (auto intType = dyn_cast(type)) { + if (intType.getWidth() <= 0 || intType.getWidth() % 8 != 0) + return std::nullopt; + return static_cast(getElementTypeSizeInBytes(type)); + } + if (auto floatType = dyn_cast(type)) { + if (floatType.getWidth() <= 0 || floatType.getWidth() % 8 != 0) + return std::nullopt; + return static_cast(getElementTypeSizeInBytes(type)); + } + return std::nullopt; +} + +std::string getScalarId(bool isScheduled, size_t opId) { return (isScheduled ? "sc:" : "gc:") + std::to_string(opId); } + +std::string getBatchLaneId(bool isScheduled, size_t opId, uint32_t lane) { + return (isScheduled ? "scb:" : "gcb:") + std::to_string(opId) + ":" + std::to_string(lane); +} + +template +bool isTopLevelRelevantCompute(Operation& op) { + return isa(&op); +} + +template +FailureOr buildTopLevelOpInfo(Operation& op, bool isScheduled, size_t opId) { + TopLevelOpInfo info; + info.op = &op; + info.opId = opId; + info.isScheduled = isScheduled; + + if constexpr (std::is_same_v) { + if (auto compute = dyn_cast(&op)) { + auto coreId = getOptionalScheduledCoreId(compute, "spatial dataflow export core id"); + if (failed(coreId)) + return failure(); + if (*coreId) + info.scalarCore = **coreId; + } + } + + return info; +} + +template +FailureOr> getBatchLaneCoreIds(BatchOpTy batch) { + if constexpr (std::is_same_v) { + auto coreIds = getOptionalScheduledBatchCoreIds(batch, "spatial dataflow export core ids"); + if (failed(coreIds)) + return failure(); + if (!*coreIds) + return SmallVector {}; + return SmallVector((**coreIds).begin(), (**coreIds).end()); + } + return SmallVector {}; +} + +std::string getExpandedNodeId(const DenseMap, ExpandedNodeInfo>& expandedNodes, + Operation* op, + uint32_t lane) { + auto it = expandedNodes.find({op, lane}); + if (it == expandedNodes.end()) + return ""; + return it->second.id; +} + +void addScalarNodeRow(std::fstream& nodesFile, + DenseMap, ExpandedNodeInfo>& expandedNodes, + const TopLevelOpInfo& info, + AsmState* asmState = nullptr) { + std::string id = getScalarId(info.isScheduled, info.opId); + SmallVector row {id, std::to_string(info.opId), "", maybeNumber(info.scalarCore)}; + if (asmState) + row.push_back(stringifyResultSsaNames(info.op, asmState)); + writeCsvRow(nodesFile, row); + expandedNodes[{info.op, 0}] = {id, info.scalarCore, std::nullopt}; +} + +template +void addBatchNodeRows(std::fstream& nodesFile, + DenseMap, ExpandedNodeInfo>& expandedNodes, + const TopLevelOpInfo& info, + BatchOpTy batch, + ArrayRef> laneCoreIds, + AsmState* asmState = nullptr) { + for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { + std::string id = getBatchLaneId(info.isScheduled, info.opId, lane); + SmallVector row { + id, std::to_string(info.opId), std::to_string(lane), maybeNumber(laneCoreIds[lane])}; + if (asmState) + row.push_back(stringifyResultSsaNames(info.op, asmState)); + writeCsvRow(nodesFile, row); + expandedNodes[{info.op, lane}] = {id, laneCoreIds[lane], lane}; + } +} + +std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t lane); + +std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t lane) { + if (value == laneArg) + return static_cast(lane); + + if (std::optional constant = matchConstantIndexValue(value)) + return *constant; + + if (auto constant = value.getDefiningOp()) { + if (auto intAttr = dyn_cast(constant.getValue())) + return intAttr.getInt(); + } + + if (auto extract = value.getDefiningOp()) { + auto constant = extract.getTensor().getDefiningOp(); + auto elements = constant ? dyn_cast(constant.getValue()) : nullptr; + auto shapedType = elements ? dyn_cast(elements.getType()) : nullptr; + if (!elements || !shapedType || shapedType.getRank() != 1 || extract.getIndices().size() != 1) + return std::nullopt; + + std::optional index = evaluateIndexLike(extract.getIndices().front(), laneArg, lane); + if (!index || *index < 0 || *index >= static_cast(elements.getNumElements())) + return std::nullopt; + + if (auto denseInts = dyn_cast(elements)) + return (*(denseInts.value_begin() + *index)).getSExtValue(); + return std::nullopt; + } + + if (auto affineApply = value.getDefiningOp()) + if (FailureOr folded = evaluateAffineApply(affineApply, + [&](Value operand) -> FailureOr { + if (std::optional resolved = + evaluateIndexLike(operand, laneArg, lane)) + return *resolved; + return failure(); + }); + succeeded(folded)) { + return *folded; + } + + return std::nullopt; +} + +SmallVector collectPossibleIntValues(Value value, Value laneArg, uint32_t lane) { + if (std::optional exact = evaluateIndexLike(value, laneArg, lane)) + return {*exact}; + + auto extract = value.getDefiningOp(); + auto constant = extract ? extract.getTensor().getDefiningOp() : nullptr; + auto elements = constant ? dyn_cast(constant.getValue()) : nullptr; + if (!elements) + return {}; + + SmallVector values; + if (auto denseInts = dyn_cast(elements)) { + values.reserve(elements.getNumElements()); + for (APInt element : denseInts.getValues()) + if (!llvm::is_contained(values, element.getSExtValue())) + values.push_back(element.getSExtValue()); + } + return values; +} + +template +std::optional getBatchLaneInput(BatchOpTy batch, uint32_t lane, unsigned inputIndex) { + if (batch.getNumResults() != 0) + return batch.getInputs()[inputIndex]; + + size_t laneCount = static_cast(batch.getLaneCount()); + if (laneCount == 0 || batch.getInputs().size() % laneCount != 0) + return std::nullopt; + + size_t inputsPerLane = batch.getInputs().size() / laneCount; + size_t flatIndex = static_cast(lane) * inputsPerLane + inputIndex; + if (flatIndex >= batch.getInputs().size()) + return std::nullopt; + return batch.getInputs()[flatIndex]; +} + +template +unsigned getBatchLaneInputCount(BatchOpTy batch) { + if (batch.getNumResults() != 0) + return batch.getInputs().size(); + + size_t laneCount = static_cast(batch.getLaneCount()); + if (laneCount == 0 || batch.getInputs().size() % laneCount != 0) + return 0; + return static_cast(batch.getInputs().size() / laneCount); +} + +template +std::optional resolveProducerForValue(Value value, std::optional consumerLane) { + Operation* op = value.getDefiningOp(); + if (!op) + return std::nullopt; + + while (auto extract = dyn_cast(op)) { + Value source = extract.getSource(); + Operation* sourceOp = source.getDefiningOp(); + auto sourceBatch = dyn_cast_or_null(sourceOp); + if (sourceBatch && sourceBatch.getNumResults() != 0) { + auto staticOffsets = extract.getStaticOffsets(); + if (!staticOffsets.empty() && staticOffsets.front() != ShapedType::kDynamic) { + uint32_t lane = static_cast(staticOffsets.front()); + return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, lane, lane, 1}; + } + if (consumerLane) + return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1}; + return ResolvedProducer { + sourceOp, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast(sourceBatch.getLaneCount())}; + } + value = source; + op = sourceOp; + if (!op) + return std::nullopt; + } + + if (auto compute = dyn_cast(op)) + return ResolvedProducer {compute.getOperation(), + static_cast(cast(value).getResultNumber()), + LogicalNodeSelector::Scalar, + 0, + 0, + 1}; + + if (auto batch = dyn_cast(op)) { + if (batch.getNumResults() != 0) { + if (consumerLane) + return ResolvedProducer {op, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1}; + return ResolvedProducer { + op, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast(batch.getLaneCount())}; + } + + uint32_t lane = static_cast(cast(value).getResultNumber()); + return ResolvedProducer {op, static_cast(lane), LogicalNodeSelector::Lane, lane, lane, 1}; + } + + return std::nullopt; +} + +SmallVector +resolveProducerSourcesForCsv(const ResolvedProducer& producer, + const DenseMap, ExpandedNodeInfo>& expandedNodes) { + SmallVector sources; + + if (producer.selector == LogicalNodeSelector::Scalar) { + std::string id = getExpandedNodeId(expandedNodes, producer.op, 0); + if (!id.empty()) + sources.push_back({id, std::nullopt}); + return sources; + } + + if (producer.selector == LogicalNodeSelector::Lane) { + std::string id = getExpandedNodeId(expandedNodes, producer.op, producer.lane); + if (!id.empty()) + sources.push_back({id, producer.lane}); + return sources; + } + + for (uint32_t lane = producer.laneStart; lane < producer.laneStart + producer.laneCount; ++lane) { + std::string id = getExpandedNodeId(expandedNodes, producer.op, lane); + if (!id.empty()) + sources.push_back({id, lane}); + } + return sources; +} + +FailureOr> getIntegerValues(Operation* op, StringRef name) { + Attribute attr = op->getAttr(name); + if (auto array = dyn_cast_or_null(attr)) + return SmallVector(array.asArrayRef()); + if (auto elements = dyn_cast_or_null(attr)) + return SmallVector(elements.getValues()); + return op->emitOpError() << "expected " << name << " integer array for Spatial dataflow report"; +} + +FailureOr +buildScheduledNodesByGraphLane(const DenseMap& topLevelInfo, + const DenseMap, ExpandedNodeInfo>& expandedNodes, + const DenseMap& graphOpsById) { + ScheduledNodeByGraphLane nodesByGraphLane; + for (const auto& entry : topLevelInfo) { + Operation* scheduledOp = entry.first; + auto sourceIds = getIntegerValues(scheduledOp, "scheduled.step_source_ids"); + auto sourceStarts = getIntegerValues(scheduledOp, "scheduled.source_lane_starts"); + auto sourceCounts = getIntegerValues(scheduledOp, "scheduled.source_lane_counts"); + if (failed(sourceIds) || failed(sourceStarts) || failed(sourceCounts)) + return failure(); + + uint32_t scheduledLaneCount = 1; + if (auto batch = dyn_cast(scheduledOp)) + scheduledLaneCount = static_cast(batch.getLaneCount()); + size_t expectedEntries = sourceIds->size() * scheduledLaneCount; + if (sourceStarts->size() != expectedEntries || sourceCounts->size() != expectedEntries) + return scheduledOp->emitOpError("inconsistent scheduling provenance arrays for Spatial dataflow report"); + + for (auto [step, graphId] : llvm::enumerate(*sourceIds)) { + auto graphIt = graphOpsById.find(graphId); + if (graphIt == graphOpsById.end()) + return scheduledOp->emitOpError() << "references unknown scheduled graph id " << graphId; + bool graphIsBatch = isa(graphIt->second); + for (uint32_t scheduledLane = 0; scheduledLane < scheduledLaneCount; ++scheduledLane) { + auto nodeIt = expandedNodes.find({scheduledOp, scheduledLane}); + if (nodeIt == expandedNodes.end()) + continue; + size_t index = step * scheduledLaneCount + scheduledLane; + int64_t start = graphIsBatch ? (*sourceStarts)[index] : 0; + int64_t count = graphIsBatch ? (*sourceCounts)[index] : 1; + if (start < 0 || count < 0) + return scheduledOp->emitOpError("negative scheduling provenance range for Spatial dataflow report"); + for (int64_t lane = start; lane < start + count; ++lane) + nodesByGraphLane[{graphId, static_cast(lane)}] = nodeIt->second; + } + } + } + return nodesByGraphLane; +} + +SmallVector resolveScheduledProducerNodes(const ResolvedProducer& producer, + const ScheduledNodeByGraphLane& nodesByGraphLane) { + SmallVector nodes; + auto graphId = producer.op->getAttrOfType("scheduled.graph_id"); + if (!graphId) + return nodes; + + uint32_t laneStart = producer.selector == LogicalNodeSelector::Scalar ? 0 : producer.laneStart; + uint32_t laneCount = producer.selector == LogicalNodeSelector::RangeRepresentative ? producer.laneCount : 1; + for (uint32_t lane = laneStart; lane < laneStart + laneCount; ++lane) + if (auto it = nodesByGraphLane.find({graphId.getInt(), lane}); it != nodesByGraphLane.end()) + nodes.push_back(it->second); + return nodes; +} + +LogicalResult +emitScheduledPlanningEdges(std::fstream& edgesFile, + func::FuncOp func, + const DenseMap& topLevelInfo, + const DenseMap, ExpandedNodeInfo>& expandedNodes, + StringRef stage) { + DenseMap graphOpsById; + for (Operation& op : func.getBody().front()) + if (auto graphId = op.getAttrOfType("scheduled.graph_id")) + graphOpsById[graphId.getInt()] = &op; + + auto nodesByGraphLane = buildScheduledNodesByGraphLane(topLevelInfo, expandedNodes, graphOpsById); + if (failed(nodesByGraphLane)) + return failure(); + + auto emitMappedEdge = + [&](const ResolvedProducer& producer, int64_t targetGraphId, uint32_t targetGraphLane, Type type) { + auto targetIt = nodesByGraphLane->find({targetGraphId, targetGraphLane}); + if (targetIt == nodesByGraphLane->end()) + return; + for (const ExpandedNodeInfo& source : resolveScheduledProducerNodes(producer, *nodesByGraphLane)) { + if (source.id == targetIt->second.id) + continue; + emitEdgeRow(edgesFile, + source.id, + targetIt->second.id, + getTypeSizeBytes(type), + type, + stage, + source.lane, + targetIt->second.lane, + std::nullopt); + } + }; + + for (Operation& op : func.getBody().front()) { + auto graphId = op.getAttrOfType("scheduled.graph_id"); + if (!graphId) + continue; + if (auto compute = dyn_cast(&op)) { + for (Value input : compute.getInputs()) + if (auto producer = resolveProducerForValue(input, std::nullopt)) + emitMappedEdge(*producer, graphId.getInt(), 0, input.getType()); + continue; + } + auto batch = dyn_cast(&op); + if (!batch) + continue; + unsigned inputCount = getBatchLaneInputCount(batch); + for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) + for (unsigned inputIndex = 0; inputIndex < inputCount; ++inputIndex) + if (std::optional input = getBatchLaneInput(batch, lane, inputIndex)) + if (auto producer = resolveProducerForValue(*input, lane)) + emitMappedEdge(*producer, graphId.getInt(), lane, input->getType()); + } + return success(); +} + +void emitEdgeRow(std::fstream& edgesFile, + StringRef sourceId, + StringRef targetId, + std::optional byteSize, + Type propagatedType, + StringRef stage, + std::optional sourceLane, + std::optional targetLane, + std::optional channelId) { + writeCsvRow(edgesFile, + {sourceId.str(), + targetId.str(), + maybeNumber(byteSize), + stringifyType(propagatedType), + stage.str(), + maybeNumber(sourceLane), + maybeNumber(targetLane), + maybeNumber(channelId)}); +} + +template +LogicalResult emitDataEdges(std::fstream& edgesFile, + const DenseMap& topLevelInfo, + const DenseMap, ExpandedNodeInfo>& expandedNodes, + StringRef stage) { + for (const auto& entry : topLevelInfo) { + Operation* op = entry.first; + const TopLevelOpInfo& info = entry.second; + + if (auto compute = dyn_cast(op)) { + for (Value input : compute.getInputs()) { + if (isa_and_nonnull(input.getDefiningOp())) + continue; + + auto producer = resolveProducerForValue(input, std::nullopt); + if (!producer) + continue; + + SmallVector sources = resolveProducerSourcesForCsv(*producer, expandedNodes); + std::optional byteSize = getTypeSizeBytes(input.getType()); + std::string targetId = getScalarId(info.isScheduled, info.opId); + for (const EdgeSource& source : sources) + emitEdgeRow(edgesFile, + source.id, + targetId, + byteSize, + input.getType(), + stage, + source.sourceLane, + std::nullopt, + std::nullopt); + } + continue; + } + + auto batch = dyn_cast(op); + if (!batch) + continue; + + unsigned inputCount = getBatchLaneInputCount(batch); + for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { + std::string targetId = getBatchLaneId(info.isScheduled, info.opId, lane); + for (unsigned inputIndex = 0; inputIndex < inputCount; ++inputIndex) { + std::optional input = getBatchLaneInput(batch, lane, inputIndex); + if (!input || isa_and_nonnull((*input).getDefiningOp())) + continue; + + auto producer = resolveProducerForValue(*input, lane); + if (!producer) + continue; + + SmallVector sources = resolveProducerSourcesForCsv(*producer, expandedNodes); + std::optional byteSize = getTypeSizeBytes((*input).getType()); + for (const EdgeSource& source : sources) + emitEdgeRow( + edgesFile, source.id, targetId, byteSize, (*input).getType(), stage, source.sourceLane, lane, std::nullopt); + } + } + } + + return success(); +} + +template +void collectChannelSends(DenseMap>& sendsByChannelId, + const DenseMap, ExpandedNodeInfo>& expandedNodes, + BatchOpTy batch) { + std::optional laneArg = batch.getLaneArgument(); + if (!laneArg) + return; + + for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { + std::string sourceId = getExpandedNodeId(expandedNodes, batch.getOperation(), lane); + if (sourceId.empty()) + continue; + batch.getBody().walk([&](SpatChannelSendOp send) { + std::optional channelId = evaluateIndexLike(send.getChannelId(), *laneArg, lane); + if (!channelId) + return; + sendsByChannelId[*channelId].push_back({sourceId, lane}); + }); + } +} + +void collectChannelSends(DenseMap>& sendsByChannelId, + const DenseMap, ExpandedNodeInfo>& expandedNodes, + SpatScheduledCompute compute) { + std::string sourceId = getExpandedNodeId(expandedNodes, compute.getOperation(), 0); + if (sourceId.empty()) + return; + compute.getBody().walk([&](SpatChannelSendOp send) { + std::optional channelId = evaluateIndexLike(send.getChannelId(), Value(), 0); + if (!channelId) + return; + sendsByChannelId[*channelId].push_back({sourceId, std::nullopt}); + }); +} + +DenseMap> +buildNodesByCore(const DenseMap, ExpandedNodeInfo>& expandedNodes) { + DenseMap> nodesByCore; + for (const auto& entry : expandedNodes) { + const ExpandedNodeInfo& node = entry.second; + if (!node.core) + continue; + nodesByCore[*node.core].push_back({node.id, node.lane}); + } + return nodesByCore; +} + +template +LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile, + const DenseMap& topLevelInfo, + ResolveChannelSourcesFn&& resolveChannelSources, + StringRef stage) { + for (const auto& entry : topLevelInfo) { + Operation* op = entry.first; + const TopLevelOpInfo& info = entry.second; + + if (auto compute = dyn_cast(op)) { + compute.getBody().walk([&](SpatChannelReceiveOp receive) { + SmallVector sources = resolveChannelSources(receive, 0); + if (sources.empty()) + return; + std::optional channelId = evaluateIndexLike(receive.getChannelId(), Value(), 0); + std::string targetId = getScalarId(info.isScheduled, info.opId); + std::optional byteSize = getTypeSizeBytes(receive.getType()); + for (const ChannelSendRecord& source : sources) + emitEdgeRow(edgesFile, + source.sourceId, + targetId, + byteSize, + receive.getType(), + stage, + source.sourceLane, + std::nullopt, + channelId); + }); + continue; + } + + auto batch = dyn_cast(op); + if (!batch) + continue; + auto laneArg = batch.getLaneArgument(); + if (!laneArg) + continue; + for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { + std::string targetId = getBatchLaneId(info.isScheduled, info.opId, lane); + batch.getBody().walk([&](SpatChannelReceiveOp receive) { + SmallVector sources = resolveChannelSources(receive, lane); + if (sources.empty()) + return; + std::optional channelId = evaluateIndexLike(receive.getChannelId(), *laneArg, lane); + std::optional byteSize = getTypeSizeBytes(receive.getType()); + for (const ChannelSendRecord& source : sources) + emitEdgeRow(edgesFile, + source.sourceId, + targetId, + byteSize, + receive.getType(), + stage, + source.sourceLane, + lane, + channelId); + }); + } + } + + return success(); +} + +LogicalResult exportGraph(func::FuncOp func, StringRef reportName) { + std::fstream nodesFile = openDialectDumpFileWithExtension((reportName + ".nodes").str(), "/reports", "csv"); + std::fstream edgesFile = openDialectDumpFileWithExtension((reportName + ".edges").str(), "/reports", "csv"); + if (!nodesFile.is_open() || !edgesFile.is_open()) + return success(); + + writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core", "ssa_name"}); + writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"}); + + Operation* asmRoot = func.getOperation(); + if (auto moduleOp = func->getParentOfType()) + asmRoot = moduleOp.getOperation(); + OpPrintingFlags flags; + flags.elideLargeElementsAttrs().enableDebugInfo(true, false); + AsmState asmState(asmRoot, flags); + + DenseMap topLevelInfo; + DenseMap, ExpandedNodeInfo> expandedNodes; + + size_t opId = 0; + for (Operation& op : func.getBody().front()) { + if (!isTopLevelRelevantCompute(op)) + continue; + FailureOr info = buildTopLevelOpInfo(op, false, opId++); + if (failed(info)) + return failure(); + topLevelInfo[&op] = *info; + + if (auto compute = dyn_cast(&op)) { + addScalarNodeRow(nodesFile, expandedNodes, *info, &asmState); + continue; + } + + auto batch = cast(&op); + SmallVector, 8> laneCoreIds(batch.getLaneCount()); + addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState); + } + + return emitDataEdges(edgesFile, topLevelInfo, expandedNodes, "spatial1"); +} + +LogicalResult exportScheduled(func::FuncOp func, StringRef reportName, StringRef stage) { + std::fstream nodesFile = openDialectDumpFileWithExtension((reportName + ".nodes").str(), "/reports", "csv"); + std::fstream edgesFile = openDialectDumpFileWithExtension((reportName + ".edges").str(), "/reports", "csv"); + if (!nodesFile.is_open() || !edgesFile.is_open()) + return success(); + + writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core", "ssa_name"}); + writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"}); + + Operation* asmRoot = func.getOperation(); + if (auto moduleOp = func->getParentOfType()) + asmRoot = moduleOp.getOperation(); + OpPrintingFlags flags; + flags.elideLargeElementsAttrs().enableDebugInfo(true, false); + AsmState asmState(asmRoot, flags); + + DenseMap topLevelInfo; + DenseMap, ExpandedNodeInfo> expandedNodes; + + size_t opId = 0; + for (Operation& op : func.getBody().front()) { + if (!isTopLevelRelevantCompute(op)) + continue; + FailureOr info = + buildTopLevelOpInfo(op, true, opId++); + if (failed(info)) + return failure(); + topLevelInfo[&op] = *info; + + if (isa(&op)) { + addScalarNodeRow(nodesFile, expandedNodes, *info, &asmState); + continue; + } + + auto batch = cast(&op); + auto coreIds = getBatchLaneCoreIds(batch); + if (failed(coreIds)) + return failure(); + SmallVector, 8> laneCoreIds(batch.getLaneCount()); + for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) + if (lane < coreIds->size()) + laneCoreIds[lane] = (*coreIds)[lane]; + addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState); + } + + if (stage == "spatial2") + return emitScheduledPlanningEdges(edgesFile, func, topLevelInfo, expandedNodes, stage); + if (failed( + emitDataEdges(edgesFile, topLevelInfo, expandedNodes, stage))) + return failure(); + + DenseMap> sendsByChannelId; + for (const auto& entry : topLevelInfo) { + Operation* op = entry.first; + if (auto compute = dyn_cast(op)) + collectChannelSends(sendsByChannelId, expandedNodes, compute); + else if (auto batch = dyn_cast(op)) + collectChannelSends(sendsByChannelId, expandedNodes, batch); + } + + DenseMap> nodesByCore = buildNodesByCore(expandedNodes); + auto resolveChannelSources = [&](SpatChannelReceiveOp receive, uint32_t lane) { + SmallVector sources; + + Value laneArg; + if (auto owner = receive->getParentOfType()) + if (auto maybeLaneArg = owner.getLaneArgument()) + laneArg = *maybeLaneArg; + + if (std::optional channelId = evaluateIndexLike(receive.getChannelId(), laneArg, lane)) { + if (auto it = sendsByChannelId.find(*channelId); it != sendsByChannelId.end()) + return it->second; + } + + for (int64_t sourceCore : collectPossibleIntValues(receive.getSourceCoreId(), laneArg, lane)) { + auto it = nodesByCore.find(static_cast(sourceCore)); + if (it == nodesByCore.end()) + continue; + llvm::append_range(sources, it->second); + } + return sources; + }; + + return emitExplicitChannelEdges( + edgesFile, topLevelInfo, resolveChannelSources, stage); +} + +} // namespace + +SpatialDataflowExportStage getSpatialDataflowExportStage() { + switch (pimExportSpatialDataflow.getValue()) { + case SpatialDataflowExportNone: return SpatialDataflowExportStage::None; + case SpatialDataflowExportSpatial1: return SpatialDataflowExportStage::Spatial1; + case SpatialDataflowExportSpatial2: return SpatialDataflowExportStage::Spatial2; + case SpatialDataflowExportSpatial3: return SpatialDataflowExportStage::Spatial3; + case SpatialDataflowExportAll: return SpatialDataflowExportStage::All; + } + llvm_unreachable("unknown spatial dataflow export mode"); +} + +bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage) { + switch (mode) { + case SpatialDataflowExportStage::None: return false; + case SpatialDataflowExportStage::Spatial1: return stage == SpatialDataflowExportStage::Spatial1; + case SpatialDataflowExportStage::Spatial2: return stage == SpatialDataflowExportStage::Spatial2; + case SpatialDataflowExportStage::Spatial3: return stage == SpatialDataflowExportStage::Spatial3; + case SpatialDataflowExportStage::All: return stage != SpatialDataflowExportStage::None; + } + return false; +} + +LogicalResult exportSpatialDataflowCsvGraph(func::FuncOp func, StringRef reportName) { + return exportGraph(func, reportName); +} + +LogicalResult exportSpatialDataflowCsvScheduled(func::FuncOp func, StringRef reportName, StringRef stage) { + return exportScheduled(func, reportName, stage); +} + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp new file mode 100644 index 0000000..0d47156 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp @@ -0,0 +1,28 @@ +#pragma once + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Support/LogicalResult.h" + +#include "llvm/ADT/StringRef.h" + +namespace onnx_mlir { +namespace spatial { + +enum class SpatialDataflowExportStage { + None, + Spatial1, + Spatial2, + Spatial3, + All, +}; + +SpatialDataflowExportStage getSpatialDataflowExportStage(); + +mlir::LogicalResult exportSpatialDataflowCsvGraph(mlir::func::FuncOp func, llvm::StringRef reportName); +mlir::LogicalResult +exportSpatialDataflowCsvScheduled(mlir::func::FuncOp func, llvm::StringRef reportName, llvm::StringRef stage); + +bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage); + +} // namespace spatial +} // namespace onnx_mlir