From a39fdba36609d45bde274aa15cc0f277a519ed16 Mon Sep 17 00:00:00 2001 From: ilgeco Date: Thu, 6 Aug 2026 14:32:46 +0200 Subject: [PATCH] Raptor sync wait --- src/PIM/Common/PimCommon.hpp | 3 + src/PIM/Compiler/PimBinaryFormat.hpp | 4 +- src/PIM/Compiler/PimCodeGen.cpp | 30 + src/PIM/Compiler/PimCodeGen.hpp | 2 + src/PIM/Compiler/PimCompilerOptions.cpp | 18 + src/PIM/Compiler/PimCompilerOptions.hpp | 2 + src/PIM/Compiler/PimCompilerUtils.cpp | 3 +- src/PIM/Compiler/PimCoreProgram.cpp | 2 + src/PIM/Compiler/PimCoreProgram.hpp | 2 + .../Common/ContractionMaterialization.cpp | 39 + .../Common/ContractionMaterialization.hpp | 23 + .../ONNXToSpatial/LowerSpatialPlansPass.cpp | 902 ++++++++++++++ .../Passes/Analyses/ONNXToSpatialVerifier.cpp | 10 +- .../Conversion/ONNXToSpatial/PlanLowering.hpp | 64 + .../SpatialLayoutCapabilities.cpp | 133 ++ .../SpatialLayoutPlanningPass.cpp | 265 ++++ src/PIM/Conversion/SpatialToPim/Common.cpp | 46 + src/PIM/Conversion/SpatialToPim/Common.hpp | 7 + .../SpatialToPim/CoreLoweringPatterns.cpp | 36 +- .../SpatialToPim/Patterns/ChannelLowering.cpp | 70 +- .../SpatialToPim/ReturnPathNormalization.cpp | 4 + .../SpatialToPim/SpatialToPimPass.cpp | 23 + .../Bufferization/ContiguityPatterns.cpp | 3 +- .../Verification/VerificationPass.cpp | 2 + src/PIM/Dialect/Pim/Pim.td | 26 + src/PIM/Dialect/Spatial/CMakeLists.txt | 1 + .../DeferredBoundaryPlanning.cpp | 6 +- .../DeferredBoundaryRealization.cpp | 148 ++- .../DeferredCommunicationDeadlock.cpp | 23 +- .../DeferredCommunicationModel.hpp | 4 + .../DeferredCommunicationRealization.cpp | 15 +- .../DeferredCommunicationRealization.hpp | 3 +- .../DeferredCommunicationScheduling.cpp | 5 +- .../DeferredTransferPlanning.cpp | 97 +- .../DeferredTransferPlanning.hpp | 5 +- .../ScheduledSpatialPasses.cpp | 85 +- .../ScheduledSpatialPasses.hpp | 16 + .../Scheduling/ComputeGraph.cpp | 17 +- .../Scheduling/ComputeGraph.hpp | 6 +- .../Scheduling/MergeSchedulingAnalysis.cpp | 7 +- .../Scheduling/MergeSchedulingAnalysis.hpp | 8 +- .../Scheduling/PipelineScheduling.cpp | 1092 +++++++++++++++++ .../Scheduling/PipelineScheduling.hpp | 26 + src/PIM/Dialect/Spatial/Spatial.td | 43 +- .../Dialect/Spatial/SpatialLayoutInterface.td | 24 + src/PIM/Dialect/Spatial/SpatialTargetInfo.hpp | 37 + src/PIM/Pass/PIMPasses.h | 63 + src/PIM/Passes/PIMPasses.h | 3 +- 48 files changed, 3357 insertions(+), 96 deletions(-) create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.hpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/SpatialLayoutCapabilities.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp create mode 100644 src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp create mode 100644 src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp create mode 100644 src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp create mode 100644 src/PIM/Dialect/Spatial/SpatialLayoutInterface.td create mode 100644 src/PIM/Dialect/Spatial/SpatialTargetInfo.hpp create mode 100644 src/PIM/Pass/PIMPasses.h diff --git a/src/PIM/Common/PimCommon.hpp b/src/PIM/Common/PimCommon.hpp index 10b6bf6..effad03 100644 --- a/src/PIM/Common/PimCommon.hpp +++ b/src/PIM/Common/PimCommon.hpp @@ -32,6 +32,9 @@ inline constexpr llvm::StringLiteral kCoreIdAttrName = "coreId"; inline constexpr llvm::StringLiteral kCoreIdsAttrName = "coreIds"; inline constexpr llvm::StringLiteral kLocalMemoryAddressAttrName = "pim.local_memory_address"; inline constexpr llvm::StringLiteral kLocalMemorySizeAttrName = "pim.local_memory_size"; +inline constexpr llvm::StringLiteral kPipelineHostBufferBytesAttrName = "pim.pipeline_host_buffer_bytes"; +inline constexpr llvm::StringLiteral kPipelineHostBufferName = "pim_pipeline_channels"; +inline constexpr size_t kPimEventRegisterCount = 32; inline constexpr std::array kRemovedLocalMemoryPlanAttrNames = { "pim.local_memory_slot", "pim.local_memory_slot_size", diff --git a/src/PIM/Compiler/PimBinaryFormat.hpp b/src/PIM/Compiler/PimBinaryFormat.hpp index 889631e..26673b6 100644 --- a/src/PIM/Compiler/PimBinaryFormat.hpp +++ b/src/PIM/Compiler/PimBinaryFormat.hpp @@ -162,8 +162,8 @@ inline constexpr std::array kInstructionJso {true, true, true, "", "", "", "len" }, // lmv {true, false, true, "core", "", "", "size"}, // send {true, false, true, "core", "", "", "size"}, // recv - {false, false, false, "", "", "", "" }, // wait - {false, false, false, "", "", "", "" }, // sync + {false, false, false, "", "event_register", "wait_value", ""}, // wait + {false, false, false, "core", "event_register", "", ""}, // sync }}; static_assert(kInstructionJsonFormats.size() == kOpcodeCount); diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index e6980f1..873fe21 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -692,6 +692,34 @@ void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge pim_binary::Opcode::send, addressOf(sendOp.getInput(), knowledge), *targetCoreId, sendOp.getSize()); } +void PimCodeGen::codeGenWaitOp( + pim::PimWaitOp waitOp, const StaticValueKnowledge& knowledge) const { + auto eventRegister = indexOf(waitOp.getEventRegister(), knowledge); + assert(succeeded(eventRegister) + && "pim.wait event register must be statically resolvable during codegen"); + pim_binary::InstructionRecord instruction; + instruction.opcode = pim_binary::Opcode::wait; + instruction.generic1 = pim::checkedI32OrCrash( + *eventRegister, "wait event register"); + instruction.generic2 = waitOp.getWaitValue(); + emitInstruction(instruction); +} + +void PimCodeGen::codeGenSyncOp( + pim::PimSyncOp syncOp, const StaticValueKnowledge& knowledge) const { + auto targetCoreId = indexOf(syncOp.getTargetCoreId(), knowledge); + auto eventRegister = indexOf(syncOp.getEventRegister(), knowledge); + assert(succeeded(targetCoreId) && succeeded(eventRegister) + && "pim.sync operands must be statically resolvable during codegen"); + pim_binary::InstructionRecord instruction; + instruction.opcode = pim_binary::Opcode::sync; + instruction.r2OrImm = pim::checkedI32OrCrash( + *targetCoreId, "sync target core id"); + instruction.generic1 = pim::checkedI32OrCrash( + *eventRegister, "sync event register"); + emitInstruction(instruction); +} + void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const { auto outputType = cast(concatOp.getOutputBuffer().getType()); assert(outputType.hasStaticShape() && "concat codegen requires static output shape"); @@ -991,6 +1019,8 @@ static LogicalResult executeCompiledCorePlan( case CompiledCoreOpKind::VMV: coreCodeGen.codeGenVMVOp(cast(node.op), knowledge); break; case CompiledCoreOpKind::Receive: coreCodeGen.codeGenReceiveOp(cast(node.op), knowledge); break; case CompiledCoreOpKind::Send: coreCodeGen.codeGenSendOp(cast(node.op), knowledge); break; + case CompiledCoreOpKind::Wait: coreCodeGen.codeGenWaitOp(cast(node.op), knowledge); break; + case CompiledCoreOpKind::Sync: coreCodeGen.codeGenSyncOp(cast(node.op), knowledge); break; case CompiledCoreOpKind::Concat: coreCodeGen.codeGenConcatOp(cast(node.op), knowledge); break; case CompiledCoreOpKind::Vmm: if (auto weightSlot = resolveWeightSlot(cast(node.op), knowledge); succeeded(weightSlot)) diff --git a/src/PIM/Compiler/PimCodeGen.hpp b/src/PIM/Compiler/PimCodeGen.hpp index 7c045c2..aab3338 100644 --- a/src/PIM/Compiler/PimCodeGen.hpp +++ b/src/PIM/Compiler/PimCodeGen.hpp @@ -217,6 +217,8 @@ public: void codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const; void codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const; + void codeGenWaitOp(pim::PimWaitOp waitOp, const StaticValueKnowledge& knowledge) const; + void codeGenSyncOp(pim::PimSyncOp syncOp, const StaticValueKnowledge& knowledge) const; void codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const; template diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index 62704b2..0a10796 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -2,6 +2,8 @@ #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" +#include + #define DEBUG_TYPE "PimCompilerOptions" namespace onnx_mlir { @@ -110,6 +112,12 @@ llvm::cl::opt llvm::cl::opt crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(64)); +llvm::cl::opt pipelineStages( + "pipeline", + llvm::cl::desc("Number of throughput pipeline stages (1 preserves latency scheduling)"), + llvm::cl::init(1), + llvm::cl::cat(OnnxMlirOptions)); + llvm::cl::opt coresCount("core-count", llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."), llvm::cl::init(-1)); @@ -129,4 +137,14 @@ void verifyExplicitPimCoreCount() { llvm::report_fatal_error("PIM compilation requires --core-count to be a positive integer"); } +void verifyPimPipelineStages() { + if (pipelineStages.getValue() == 0) + llvm::report_fatal_error("PIM compilation requires --pipeline to be positive"); + if (static_cast(coresCount.getValue()) % pipelineStages.getValue() != 0) + llvm::report_fatal_error("PIM compilation requires --core-count to be divisible by --pipeline"); + if (crossbarCountInCore.getValue() + > std::numeric_limits::max() / pipelineStages.getValue()) + llvm::report_fatal_error("PIM compilation --crossbar-count * --pipeline overflows"); +} + } // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index baff5c7..e6cff19 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -62,6 +62,7 @@ extern llvm::cl::opt pimVerifyBufferizationCopyFreedom; extern llvm::cl::opt crossbarSize; extern llvm::cl::opt crossbarCountInCore; +extern llvm::cl::opt pipelineStages; extern llvm::cl::opt coresCount; extern llvm::cl::opt pimTargetConfig; extern llvm::cl::opt pimConvIm2colMaxElements; @@ -69,5 +70,6 @@ extern llvm::cl::opt pimConvStreamChunkPositions; bool hasExplicitPimCoreCount(); void verifyExplicitPimCoreCount(); +void verifyPimPipelineStages(); } // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index 0ed3c4f..d172285 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -330,6 +330,7 @@ void addPassesPim(OwningOpRef& module, EmissionTargetType& emissionTarget, std::string outputNameNoExt) { verifyExplicitPimCoreCount(); + verifyPimPipelineStages(); spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget(); spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget); @@ -354,7 +355,7 @@ void addPassesPim(OwningOpRef& module, pm.addPass(createTrivialGraphComputeMergePass( schedulingTarget.residentWeightCapacity, exportStage)); pm.addPass(spatial::createScheduleAndRealizeSpatialPass( - schedulingTarget, exportStage)); + schedulingTarget, exportStage, pipelineStages.getValue())); pm.addPass(createMessagePass("Onnx lowered to Spatial")); } diff --git a/src/PIM/Compiler/PimCoreProgram.cpp b/src/PIM/Compiler/PimCoreProgram.cpp index 98ce549..3dde637 100644 --- a/src/PIM/Compiler/PimCoreProgram.cpp +++ b/src/PIM/Compiler/PimCoreProgram.cpp @@ -17,6 +17,8 @@ static FailureOr classifyCompiledCoreOpKind(Operation& op) { if (isa(op)) return CompiledCoreOpKind::VMV; if (isa(op)) return CompiledCoreOpKind::Receive; if (isa(op)) return CompiledCoreOpKind::Send; + if (isa(op)) return CompiledCoreOpKind::Wait; + if (isa(op)) return CompiledCoreOpKind::Sync; if (isa(op)) return CompiledCoreOpKind::Concat; if (isa(op)) return CompiledCoreOpKind::Vmm; if (isa(op)) return CompiledCoreOpKind::VVAdd; diff --git a/src/PIM/Compiler/PimCoreProgram.hpp b/src/PIM/Compiler/PimCoreProgram.hpp index 8afd4b4..411a655 100644 --- a/src/PIM/Compiler/PimCoreProgram.hpp +++ b/src/PIM/Compiler/PimCoreProgram.hpp @@ -17,6 +17,8 @@ enum class CompiledCoreOpKind : uint8_t { VMV, Receive, Send, + Wait, + Sync, Concat, Vmm, VVAdd, diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.cpp new file mode 100644 index 0000000..dc11798 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.cpp @@ -0,0 +1,39 @@ +#include "ContractionMaterialization.hpp" + +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" +#include "MatrixProductLowering.hpp" + +namespace onnx_mlir { + +mlir::Value materializePaddedContractionInput( + mlir::Value input, + mlir::RankedTensorType paddedType, + mlir::PatternRewriter& rewriter, + mlir::Location loc) { + return createPaddedInputCompute(input, paddedType, rewriter, loc); +} + +mlir::FailureOr materializeTransposedContractionConstant( + mlir::Value input, + mlir::RankedTensorType resultType, + llvm::ArrayRef permutation, + mlir::PatternRewriter& rewriter, + mlir::Location loc) { + auto denseAttr = getHostConstDenseElementsAttr(input); + auto inputType = denseAttr ? mlir::dyn_cast(denseAttr.getType()) : nullptr; + if (!inputType || !inputType.hasStaticShape() || !resultType || !resultType.hasStaticShape() + || inputType.getRank() != resultType.getRank()) + return mlir::failure(); + + auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation); + if (mlir::failed(transposedAttr) || transposedAttr->getType() != resultType) + return mlir::failure(); + + return getOrCreateConstant(rewriter, + rewriter.getInsertionBlock()->getParentOp(), + *transposedAttr, + resultType); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.hpp new file mode 100644 index 0000000..5f92d3d --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include "llvm/ADT/ArrayRef.h" + +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/PatternMatch.h" + +namespace onnx_mlir { + +mlir::Value materializePaddedContractionInput( + mlir::Value input, + mlir::RankedTensorType paddedType, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +mlir::FailureOr materializeTransposedContractionConstant( + mlir::Value input, + mlir::RankedTensorType resultType, + llvm::ArrayRef permutation, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp new file mode 100644 index 0000000..25a35b2 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -0,0 +1,902 @@ +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "mlir/Transforms/Passes.h" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" +#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" + +using namespace mlir; + +namespace onnx_mlir { +namespace { + +static FailureOr getRowStripValue(Value value) { + return getRowStripPhysicalValue(value); +} + +static FailureOr publishRowStripValue(Operation* planOp, + Value storage, + PatternRewriter& rewriter) { + auto logicalType = dyn_cast(planOp->getResult(0).getType()); + if (!logicalType) + return planOp->emitOpError("requires ranked logical output type"), failure(); + FailureOr value = describeRowStripPhysicalValue(storage, logicalType); + if (failed(value)) + return planOp->emitOpError("lowering produced invalid row-strip physical storage"), failure(); + FailureOr blueprint = createRowStripStorageBlueprint( + storage, logicalType, rewriter, planOp->getLoc()); + if (failed(blueprint)) + return planOp->emitOpError("failed to create row-strip storage Blueprint"), failure(); + rewriter.replaceOp(planOp, *blueprint); + return *blueprint; +} + +static bool isRowStripSelected(Operation* op) { + auto selected = spatial::getSelectedPhysicalLayout(op); + return selected && *selected == spatial::PhysicalLayout::NHWCRowStrip; +} + +static bool isDenseSelected(Operation* op) { + auto selected = spatial::getSelectedPhysicalLayout(op); + return selected && *selected == spatial::PhysicalLayout::DenseNCHW; +} + +static spatial::PhysicalLayout getKnownPhysicalLayout(Value value) { + if (auto materialize = value.getDefiningOp()) + return materialize.getTargetPhysicalLayout(); + if (auto blueprint = value.getDefiningOp()) + return blueprint.getPhysicalLayout(); + if (Operation* producer = value.getDefiningOp()) { + if (auto selected = spatial::getSelectedPhysicalLayout(producer)) + return *selected; + } + return spatial::PhysicalLayout::DenseNCHW; +} + +static LogicalResult verifySelectedLayouts( + func::FuncOp funcOp, const spatial::SpatialTargetInfo& target) { + LogicalResult result = success(); + funcOp.walk([&](Operation* op) { + auto capability = dyn_cast(op); + if (!capability) + return; + auto selected = spatial::getSelectedPhysicalLayout(op); + if (!selected) { + op->emitOpError("requires a selected physical layout from SpatialLayoutPlanning"); + result = failure(); + return; + } + if (*selected != spatial::PhysicalLayout::DenseNCHW + && *selected != spatial::PhysicalLayout::NHWCRowStrip) { + op->emitOpError("has an unsupported selected physical layout"); + result = failure(); + return; + } + SmallVector operandLayouts; + operandLayouts.reserve(op->getNumOperands()); + for (Value operand : op->getOperands()) + operandLayouts.push_back(getKnownPhysicalLayout(operand)); + auto alternatives = capability.getLayoutAlternatives(target, operandLayouts); + if (llvm::none_of(alternatives, [&](const spatial::LayoutAlternative& alternative) { + return alternative.resultLayout == *selected + && alternative.operandLayouts == operandLayouts; + })) { + op->emitOpError("selected physical layout is not lowerable for its explicit operand layouts"); + result = failure(); + } + }); + return result; +} + +static FailureOr +lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) { + return applyRowStripRelu(input, rewriter, planOp.getLoc()); +} + +static FailureOr +lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) { + return applyRowStripSilu(input, rewriter, planOp.getLoc()); +} + +static FailureOr lowerRowStripBiasAdd(const RowStripPhysicalValue& input, + spatial::SpatBiasAddPlanOp planOp, + PatternRewriter& rewriter) { + return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc()); +} + +static FailureOr lowerRowStripAdd(const RowStripPhysicalValue& lhs, + const RowStripPhysicalValue& rhs, + spatial::SpatAddPlanOp planOp, + PatternRewriter& rewriter) { + return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc()); +} + +static FailureOr lowerRowStripConcat(ArrayRef inputs, + spatial::SpatConcatPlanOp planOp, + PatternRewriter& rewriter) { + auto outputType = dyn_cast(planOp.getOutput().getType()); + if (!outputType) + return failure(); + return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc()); +} + +static FailureOr +materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) { + if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape()) + return failure(); + return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc); +} + +static FailureOr materializeDenseToRowStrip( + Value input, RankedTensorType logicalType, Location loc, PatternRewriter& rewriter) { + if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4 + || logicalType.getDimSize(0) != 1) + return failure(); + auto nhwcType = RankedTensorType::get( + {1, logicalType.getDimSize(2), logicalType.getDimSize(3), logicalType.getDimSize(1)}, + logicalType.getElementType(), logicalType.getEncoding()); + auto rowsType = RankedTensorType::get( + {logicalType.getDimSize(2) * logicalType.getDimSize(3), logicalType.getDimSize(1)}, + logicalType.getElementType(), logicalType.getEncoding()); + auto rowsCompute = createSpatCompute<1>( + rewriter, loc, rowsType, {}, input, [&](Value denseInput) { + Value nhwc = createLinalgTranspose( + denseInput, nhwcType, {0, 2, 3, 1}, rewriter, loc); + Value rows = tensor::CollapseShapeOp::create( + rewriter, loc, rowsType, nhwc, + SmallVector {{0, 1, 2}, {3}}); + spatial::SpatYieldOp::create(rewriter, loc, rows); + }); + Value rows = rowsCompute->getResult(0); + FailureOr storage = createRowStripStorageFromRows(rows, logicalType, rewriter, loc); + if (failed(storage)) + return failure(); + return createRowStripStorageBlueprint(*storage, logicalType, rewriter, loc); +} + +static FailureOr lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType, + PatternRewriter& rewriter, Location loc) { + auto producer = input.getDefiningOp(); + auto inputType = dyn_cast(input.getType()); + auto biasType = dyn_cast(bias.getType()); + if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape() + || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount() + || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount()) + return failure(); + auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount()); + auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount()); + if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank() + || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front() + || inputFragmentType->getRank() != outputFragmentType->getRank() + 1) + return failure(); + for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape())) + if (outputDim > inputDim) + return failure(); + + auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias}, + [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult { + FailureOr fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType); + if (failed(fragment)) + return failure(); + MixedSliceGeometry biasSlice; + for (int64_t dim : inputFragmentType->getShape()) { + biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0)); + biasSlice.sizes.push_back(rewriter.getIndexAttr(dim)); + biasSlice.strides.push_back(rewriter.getIndexAttr(1)); + } + Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice); + if (!biasFragment) + return failure(); + Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment); + MixedSliceGeometry outputSlice; + outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0)); + outputSlice.sizes.push_back(rewriter.getIndexAttr(1)); + outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1)); + for (int64_t dim : outputFragmentType->getShape()) + outputSlice.sizes.push_back(rewriter.getIndexAttr(dim)); + Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice); + if (!output) + return failure(); + publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane); + return success(); + }); + if (failed(batch)) + return failure(); + return batch->getResult(0); +} + +struct LowerDenseReluPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp, + PatternRewriter& rewriter) const override { + auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation()); + if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW) + return failure(); + + auto computeOp = createSpatCompute<1>( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) { + auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult()); + }); + rewriter.replaceOp(planOp, computeOp.getResults()); + return success(); + } +}; + +struct LowerDenseSiluPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp, + PatternRewriter& rewriter) const override { + auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation()); + if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW) + return failure(); + + auto computeOp = createSpatCompute<1>( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) { + Value sigmoid = spatial::SpatSigmoidOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult(); + Value silu = spatial::SpatVMulOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult(); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu); + }); + rewriter.replaceOp(planOp, computeOp.getResults()); + return success(); + } +}; + +struct LowerDenseResizePlan final : OpRewritePattern { + explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target) + : OpRewritePattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + FailureOr lowered = lowerSelectedResizeNearestPlan(planOp, std::nullopt, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense nearest Resize plan"); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + + const spatial::SpatialTargetInfo& target; +}; + +struct LowerDenseBiasAddPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + auto resultType = dyn_cast(planOp.getOutput().getType()); + if (!resultType) + return planOp.emitOpError("requires ranked output type"); + + FailureOr denseBias = materializeDenseBiasAddTensor( + planOp.getBias(), resultType, rewriter, planOp.getLoc()); + if (failed(denseBias)) + return planOp.emitOpError("failed to materialize dense Conv-style bias"); + if (planOp.getInput().getDefiningOp()) { + FailureOr lowered = lowerDenseBatchBiasAdd( + planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc()); + if (succeeded(lowered)) { + rewriter.replaceOp(planOp, *lowered); + return success(); + } + } + auto computeOp = createSpatCompute<2>( + rewriter, + planOp.getLoc(), + planOp.getOutput().getType(), + {}, + ValueRange {planOp.getInput(), *denseBias}, + [&](Value x, Value y) { + auto added = spatial::SpatVAddOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult()); + }); + rewriter.replaceOp(planOp, computeOp.getResults()); + return success(); + } +}; + +struct LowerDenseAddPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + auto compute = createSpatCompute<2>( + rewriter, + planOp.getLoc(), + planOp.getOutput().getType(), + {}, + ValueRange {planOp.getLhs(), planOp.getRhs()}, + [&](Value lhsValue, Value rhsValue) { + Value added = spatial::SpatVAddOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added); + }); + rewriter.replaceOp(planOp, compute.getResults()); + return success(); + } +}; + +struct LowerDenseConcatPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + auto compute = createSpatCompute( + rewriter, + planOp.getLoc(), + TypeRange {planOp.getOutput().getType()}, + {}, + planOp.getInputs(), + [&](ValueRange values) { + Value concatenated = spatial::SpatConcatOp::create( + rewriter, + planOp.getLoc(), + planOp.getOutput().getType(), + rewriter.getI64IntegerAttr(planOp.getAxis()), + values); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated); + }); + rewriter.replaceOp(planOp, compute.getResults()); + return success(); + } +}; + +static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp, + PatternRewriter& rewriter) { + FailureOr lhs = getRowStripValue(planOp.getLhs()); + FailureOr rhs = getRowStripValue(planOp.getRhs()); + if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) { + if (getKnownPhysicalLayout(planOp.getLhs()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs"); + } + if (isRowStripSelected(planOp.getOperation()) && failed(rhs)) { + if (getKnownPhysicalLayout(planOp.getRhs()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs"); + } + if (isRowStripSelected(planOp.getOperation())) { + rewriter.setInsertionPoint(planOp); + FailureOr lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial add plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + return planOp.emitOpError("dense Add plan was not lowered by the selected-plan patterns"); +} + +static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp, + PatternRewriter& rewriter) { + SmallVector inputs; + for (Value input : planOp.getInputs()) { + FailureOr physical = getRowStripValue(input); + if (failed(physical)) { + inputs.clear(); + break; + } + inputs.push_back(*physical); + } + if (isRowStripSelected(planOp.getOperation()) && inputs.size() != planOp.getInputs().size()) { + if (llvm::any_of(planOp.getInputs(), [](Value input) { + return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip; + })) + return failure(); + return planOp.emitOpError("selected row-strip Concat plan requires row-strip inputs"); + } + if (isRowStripSelected(planOp.getOperation())) { + rewriter.setInsertionPoint(planOp); + FailureOr lowered = lowerRowStripConcat(inputs, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + return planOp.emitOpError("dense Concat plan was not lowered by the selected-plan patterns"); +} + +struct LowerSelectedConvPlan final : OpRewritePattern { + explicit LowerSelectedConvPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target) + : OpRewritePattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatConv2DPlanOp planOp, + PatternRewriter& rewriter) const override { + if (isDenseSelected(planOp.getOperation())) { + FailureOr lowered = lowerSelectedConv2DPlan( + planOp, std::nullopt, /*emitRowStripLayout=*/false, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense Spatial Conv plan"); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + + FailureOr rowStripInput = getRowStripValue(planOp.getInput()); + if (failed(rowStripInput) + && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + std::optional physicalInput; + if (succeeded(rowStripInput)) + physicalInput = rowStripInput->storage; + FailureOr lowered = lowerSelectedConv2DPlan( + planOp, physicalInput, /*emitRowStripLayout=*/true, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + + const spatial::SpatialTargetInfo& target; +}; + +struct LowerRowStripReluPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(planOp.getInput()); + if (failed(input)) { + if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input"); + } + FailureOr lowered = lowerRowStripRelu(*input, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial ReLU plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } +}; + +struct LowerRowStripSiluPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(planOp.getInput()); + if (failed(input)) { + if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input"); + } + FailureOr lowered = lowerRowStripSilu(*input, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } +}; + +struct LowerRowStripResizePlan final : OpRewritePattern { + explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target) + : OpRewritePattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(planOp.getInput()); + if (failed(input)) { + if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip Resize plan requires a row-strip input"); + } + FailureOr lowered = lowerSelectedResizeNearestPlan(planOp, input->storage, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Resize plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + + const spatial::SpatialTargetInfo& target; +}; + +struct LowerDenseMaxPoolPlan final : OpRewritePattern { + explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target) + : OpRewritePattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + FailureOr lowered = lowerDenseMaxPool2DPlan(planOp, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense Spatial MaxPool plan"); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + + const spatial::SpatialTargetInfo& target; +}; + +struct LowerRowStripMaxPoolPlan final : OpRewritePattern { + explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target) + : OpRewritePattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(planOp.getInput()); + if (failed(input) + && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + std::optional physicalInput; + if (succeeded(input)) + physicalInput = input->storage; + FailureOr lowered = lowerSelectedMaxPool2DPlan(planOp, physicalInput, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + + const spatial::SpatialTargetInfo& target; +}; + +struct LowerRowStripGlobalAveragePoolPlan + final : OpRewritePattern { + explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target) + : OpRewritePattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(planOp.getInput()); + if (failed(input) + && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + std::optional physicalInput; + if (succeeded(input)) + physicalInput = input->storage; + FailureOr lowered = lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + + const spatial::SpatialTargetInfo& target; +}; + +struct LowerDenseGlobalAveragePoolPlan + final : OpRewritePattern { + explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx, + const spatial::SpatialTargetInfo& target) + : OpRewritePattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + FailureOr lowered = lowerDenseGlobalAveragePoolPlan(planOp, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense Spatial global AveragePool plan"); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + + const spatial::SpatialTargetInfo& target; +}; + +struct LowerRowStripBiasAddPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(planOp.getInput()); + if (failed(input)) { + if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip bias_add plan requires a row-strip input"); + } + FailureOr lowered = lowerRowStripBiasAdd(*input, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } +}; + +struct LowerRowStripAddPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + return lowerAddPlan(planOp, rewriter); + } +}; + +struct LowerRowStripConcatPlan final : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp, + PatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + return lowerConcatPlan(planOp, rewriter); + } +}; + +struct LowerMaterializeLayout final + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp, + PatternRewriter& rewriter) const override { + auto source = materializeOp.getSourcePhysicalLayout(); + auto target = materializeOp.getTargetPhysicalLayout(); + if (source == spatial::PhysicalLayout::DenseNCHW + && target == spatial::PhysicalLayout::DenseNCHW) { + rewriter.replaceOp(materializeOp, materializeOp.getInput()); + return success(); + } + if (source == spatial::PhysicalLayout::DenseNCHW + && target == spatial::PhysicalLayout::NHWCRowStrip) { + auto logicalType = dyn_cast(materializeOp.getInput().getType()); + if (!logicalType) + return materializeOp.emitOpError("requires a ranked dense input"), failure(); + FailureOr rowStrip = materializeDenseToRowStrip( + materializeOp.getInput(), logicalType, materializeOp.getLoc(), rewriter); + if (failed(rowStrip)) + return materializeOp.emitOpError( + "failed to materialize dense NCHW storage to row-strip layout"), failure(); + rewriter.replaceOp(materializeOp, *rowStrip); + return success(); + } + if (source != spatial::PhysicalLayout::NHWCRowStrip + || target != spatial::PhysicalLayout::DenseNCHW) + return materializeOp.emitOpError( + "unsupported Spatial layout materialization direction"), failure(); + auto inputType = dyn_cast(materializeOp.getInput().getType()); + if (!inputType) + return materializeOp.emitOpError("requires a ranked row-strip input"), failure(); + FailureOr rowStripValue = + getRowStripValue(materializeOp.getInput()); + if (failed(rowStripValue)) + return materializeOp.emitOpError( + "requires an explicitly defining row-strip physical value"), failure(); + FailureOr dense = materializeRowStripToDense( + *rowStripValue, materializeOp.getLoc(), rewriter); + if (failed(dense)) + return materializeOp.emitOpError( + "failed to materialize row-strip storage to dense NCHW"), failure(); + rewriter.replaceOp(materializeOp, *dense); + return success(); + } +}; + +struct LowerRowStripFlatten final + : OpRewritePattern { + explicit LowerRowStripFlatten(MLIRContext* context, + const spatial::SpatialTargetInfo& target) + : OpRewritePattern(context), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatGraphCompute flattenOp, + PatternRewriter& rewriter) const override { + if (flattenOp.getInputs().size() != 1) + return failure(); + FailureOr input = + getRowStripValue(flattenOp.getInputs().front()); + if (failed(input) || failed(canLowerFlattenFromRowStrip(flattenOp, target))) + return failure(); + if (failed(lowerFlattenFromRowStrip(*input, flattenOp, target, rewriter))) + return flattenOp.emitOpError( + "failed to preserve row-strip layout through Flatten"), failure(); + return success(); + } + + const spatial::SpatialTargetInfo& target; +}; + +struct LowerSpatialPlansPass final : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass) + + StringRef getArgument() const override { return "lower-spatial-plans"; } + StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; } + + LowerSpatialPlansPass() = default; + explicit LowerSpatialPlansPass(const spatial::SpatialTargetInfo& target) + : target(target), hasTarget(true) {} + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + if (!hasTarget) { + moduleOp.emitError("Spatial plan lowering requires an injected SpatialTargetInfo"); + signalPassFailure(); + return; + } + MLIRContext* ctx = moduleOp.getContext(); + auto entryFunc = getPimEntryFunc(moduleOp); + if (failed(entryFunc)) { + moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans"); + signalPassFailure(); + return; + } + func::FuncOp funcOp = *entryFunc; + PatternRewriter rewriter(ctx); + auto verifyLogicalPhase = [&](StringRef stage) -> bool { + if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc))) + return true; + moduleOp.emitError() << "logical Spatial graph verification failed " << stage; + signalPassFailure(); + return false; + }; + + if (!verifyLogicalPhase("at the start of LowerSpatialPlans")) + return; + if (failed(verifySelectedLayouts(funcOp, target))) { + moduleOp.emitError("selected Spatial layout verification failed"); + signalPassFailure(); + return; + } + + RewritePatternSet selectedPlanPatterns(ctx); + selectedPlanPatterns.add(ctx); + selectedPlanPatterns.add(ctx, target); + if (failed(applyPatternsGreedily(funcOp, std::move(selectedPlanPatterns)))) { + moduleOp.emitError("failed to lower selected Spatial plans"); + signalPassFailure(); + return; + } + + RewritePatternSet layoutPatterns(ctx); + layoutPatterns.add(ctx); + layoutPatterns.add(ctx, target); + ConversionTarget layoutTarget(*ctx); + layoutTarget.addLegalDialect(); + layoutTarget.addIllegalDialect(); + layoutTarget.addIllegalOp(); + layoutTarget.addDynamicallyLegalOp( + [&](spatial::SpatGraphCompute computeOp) { + if (computeOp.getInputs().size() != 1) + return true; + FailureOr input = + getRowStripValue(computeOp.getInputs().front()); + return failed(input) || failed(canLowerFlattenFromRowStrip(computeOp, target)); + }); + FrozenRewritePatternSet frozenLayoutPatterns(std::move(layoutPatterns)); + if (failed(applyFullConversion(funcOp, layoutTarget, + frozenLayoutPatterns))) { + moduleOp.emitError("failed to lower explicit Spatial layout materialization"); + signalPassFailure(); + return; + } + + if (!verifyLogicalPhase("after selected-plan conversion")) + return; + SmallVector deadPhysicalViews; + funcOp.walk([&](spatial::SpatBlueprintOp blueprint) { + if (spatial::isPhysicalView(blueprint.getMode()) && blueprint.use_empty()) + deadPhysicalViews.push_back(blueprint); + }); + for (spatial::SpatBlueprintOp blueprint : deadPhysicalViews) + rewriter.eraseOp(blueprint); + bool hasIllegalOps = false; + moduleOp.walk([&](Operation* op) { + if (isa(op)) + return; + if (auto blueprint = dyn_cast(op)) { + if (spatial::isFragmentAssembly(blueprint.getMode())) + return; + op->emitOpError("planning blueprint must not remain after LowerSpatialPlans"); + hasIllegalOps = true; + } + else if (isa(op) + || op->getDialect()->getNamespace() == "onnx") { + op->emitOpError("operation must not remain after LowerSpatialPlans"); + hasIllegalOps = true; + } + }); + + PassManager canonicalizationPM(ctx); + canonicalizationPM.addPass(createCanonicalizerPass()); + if (failed(canonicalizationPM.run(moduleOp))) + moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing"); + + if (hasIllegalOps) { + signalPassFailure(); + } else { + dumpModule(moduleOp, "spatial1_graph"); + spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage(); + if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1) + && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) { + signalPassFailure(); + return; + } + } + + if (!verifyLogicalPhase("at the end of LowerSpatialPlans")) + return; + } + + spatial::SpatialTargetInfo target; + bool hasTarget = false; +}; + +} // namespace + +std::unique_ptr createLowerSpatialPlansPass() { return std::make_unique(); } + +std::unique_ptr createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target) { + return std::make_unique(target); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp index 186a945..625f1cf 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp @@ -108,7 +108,9 @@ void verifyScheduledInputs(ComputeOpTy compute, for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) { size_t currentInputIndex = inputIndex; Operation* definingOp = input.getDefiningOp(); - if (allowChannelReceiveInputs && isa_and_nonnull(definingOp)) + if (allowChannelReceiveInputs + && isa_and_nonnull(definingOp)) continue; if (isScheduledPhase1Value(input)) continue; @@ -163,7 +165,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter }); continue; } - if (isa(&op)) { + if (isa(&op)) { diagnostics.report(&op, [&](Operation* illegalOp) { illegalOp->emitOpError() << kPhaseMarker << " explicit channel communication is not expected before merge materialization"; @@ -182,7 +185,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) { for (Operation& op : funcOp.getOps()) { - if (isa(&op)) { + if (isa(&op)) { diagnostics.report(&op, [&](Operation* illegalOp) { illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1"; }); diff --git a/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp b/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp new file mode 100644 index 0000000..2060508 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp @@ -0,0 +1,64 @@ +#pragma once + +#include + +#include "mlir/IR/PatternMatch.h" +#include "mlir/Support/LogicalResult.h" + +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +namespace onnx_mlir { + +struct RowStripPhysicalValue; + +mlir::FailureOr +lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, + std::optional rowStripInput, + bool emitRowStripLayout, + const spatial::SpatialTargetInfo& target, + mlir::PatternRewriter& rewriter); + +mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp, + const spatial::SpatialTargetInfo& target); +mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp, + const spatial::SpatialTargetInfo& target); + +mlir::LogicalResult canLowerResizeNearestPlanToRowStrip( + spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetInfo& target); + +mlir::FailureOr lowerSelectedResizeNearestPlan( + spatial::SpatResizeNearestPlanOp planOp, + std::optional rowStripInput, + const spatial::SpatialTargetInfo& target, + mlir::PatternRewriter& rewriter); + +mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp, + const spatial::SpatialTargetInfo& target); + +mlir::FailureOr +lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, + const spatial::SpatialTargetInfo& target, + mlir::PatternRewriter& rewriter); + +mlir::FailureOr +lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, + std::optional rowStripInput, + const spatial::SpatialTargetInfo& target, + mlir::PatternRewriter& rewriter); + +mlir::LogicalResult +canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp, + const spatial::SpatialTargetInfo& target); + +mlir::FailureOr +lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp, + const spatial::SpatialTargetInfo& target, + mlir::PatternRewriter& rewriter); + +mlir::FailureOr +lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp, + std::optional rowStripInput, + const spatial::SpatialTargetInfo& target, + mlir::PatternRewriter& rewriter); + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutCapabilities.cpp b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutCapabilities.cpp new file mode 100644 index 0000000..e96a8a5 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutCapabilities.cpp @@ -0,0 +1,133 @@ +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +using namespace mlir; + +namespace onnx_mlir::spatial { + +static LayoutAlternative denseAlternative(Operation *op) { + LayoutAlternative alternative; + alternative.operandLayouts.assign(op->getNumOperands(), PhysicalLayout::DenseNCHW); + alternative.resultLayout = PhysicalLayout::DenseNCHW; + return alternative; +} + +static LayoutAlternative rowStripAlternative(Operation *op, + ArrayRef operandLayouts) { + LayoutAlternative alternative; + alternative.operandLayouts.assign(operandLayouts.begin(), operandLayouts.end()); + alternative.resultLayout = PhysicalLayout::NHWCRowStrip; + alternative.intrinsicCost = -2; + return alternative; +} + +static bool hasRowStripInput(ArrayRef operandLayouts, unsigned index) { + return index < operandLayouts.size() + && operandLayouts[index] == PhysicalLayout::NHWCRowStrip; +} + +SmallVector SpatConv2DPlanOp::getLayoutAlternatives( + const SpatialTargetInfo& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (hasRowStripInput(operandLayouts, 0)) { + if (succeeded(canConsumeAndProduceRowStrip(*this, target))) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + } + else if (succeeded(canLowerConvPlanToRowStrip(*this, target))) { + LayoutAlternative alternative = denseAlternative(getOperation()); + alternative.resultLayout = PhysicalLayout::NHWCRowStrip; + alternative.intrinsicCost = -2; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatReluPlanOp::getLayoutAlternatives( + const SpatialTargetInfo&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (hasRowStripInput(operandLayouts, 0)) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + return alternatives; +} + +SmallVector SpatSiluPlanOp::getLayoutAlternatives( + const SpatialTargetInfo&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (hasRowStripInput(operandLayouts, 0)) { + LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts); + alternative.intrinsicCost = -3; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatResizeNearestPlanOp::getLayoutAlternatives( + const SpatialTargetInfo& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (hasRowStripInput(operandLayouts, 0) + && succeeded(canLowerResizeNearestPlanToRowStrip(*this, target))) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + return alternatives; +} + +SmallVector SpatMaxPool2DPlanOp::getLayoutAlternatives( + const SpatialTargetInfo& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) { + LayoutAlternative alternative = denseAlternative(getOperation()); + if (hasRowStripInput(operandLayouts, 0)) + alternative = rowStripAlternative(getOperation(), operandLayouts); + alternative.resultLayout = PhysicalLayout::NHWCRowStrip; + alternative.intrinsicCost = -2; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatGlobalAveragePoolPlanOp::getLayoutAlternatives( + const SpatialTargetInfo& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) { + LayoutAlternative alternative = denseAlternative(getOperation()); + if (hasRowStripInput(operandLayouts, 0)) + alternative = rowStripAlternative(getOperation(), operandLayouts); + alternative.resultLayout = PhysicalLayout::NHWCRowStrip; + alternative.intrinsicCost = -2; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatBiasAddPlanOp::getLayoutAlternatives( + const SpatialTargetInfo&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + auto resultType = dyn_cast(getOutput().getType()); + if (resultType && hasRowStripInput(operandLayouts, 0) + && isSupportedBiasAddValue(getBias(), resultType)) + alternatives.push_back(rowStripAlternative(getOperation(), + {PhysicalLayout::NHWCRowStrip, + PhysicalLayout::DenseNCHW})); + return alternatives; +} + +SmallVector SpatAddPlanOp::getLayoutAlternatives( + const SpatialTargetInfo&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0) + && hasRowStripInput(operandLayouts, 1)) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + return alternatives; +} + +SmallVector SpatConcatPlanOp::getLayoutAlternatives( + const SpatialTargetInfo&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (!operandLayouts.empty() && llvm::all_of(operandLayouts, [](PhysicalLayout layout) { + return layout == PhysicalLayout::NHWCRowStrip; + })) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + return alternatives; +} + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp new file mode 100644 index 0000000..56eb1e6 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp @@ -0,0 +1,265 @@ +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/DenseMap.h" + +#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Pass/PIMPasses.h" + +#include + +using namespace mlir; + +namespace onnx_mlir { +namespace { + +using LayoutMap = llvm::DenseMap; + +static spatial::PhysicalLayout getSelectedLayout(const LayoutMap& layouts, Value value) { + if (auto it = layouts.find(value); it != layouts.end()) + return it->second; + if (auto materialize = value.getDefiningOp()) + return materialize.getTargetPhysicalLayout(); + if (auto blueprint = value.getDefiningOp()) + return blueprint.getPhysicalLayout(); + return spatial::PhysicalLayout::DenseNCHW; +} + +static SmallVector getOperandLayouts( + Operation* op, const LayoutMap& layouts) { + SmallVector operandLayouts; + operandLayouts.reserve(op->getNumOperands()); + for (Value operand : op->getOperands()) + operandLayouts.push_back(getSelectedLayout(layouts, operand)); + return operandLayouts; +} + +static FailureOr> getAlternatives( + Operation* op, const LayoutMap& layouts, const spatial::SpatialTargetInfo& target) { + auto capability = dyn_cast(op); + if (!capability) + return failure(); + SmallVector alternatives = + capability.getLayoutAlternatives(target, getOperandLayouts(op, layouts)); + if (alternatives.empty()) + return op->emitOpError("does not advertise a legal Spatial layout alternative"), failure(); + for (const spatial::LayoutAlternative& alternative : alternatives) + if (alternative.operandLayouts.size() != op->getNumOperands()) + return op->emitOpError("advertises a layout alternative with the wrong operand count"), failure(); + return alternatives; +} + +static unsigned findCurrentAlternative( + Operation* op, ArrayRef alternatives, + spatial::PhysicalLayout selectedResult) { + for (auto [index, alternative] : llvm::enumerate(alternatives)) + if (alternative.resultLayout == selectedResult) + return index; + return 0; +} + +static int64_t alternativeCost(Operation* op, + const spatial::LayoutAlternative& alternative, + const LayoutMap& layouts, + const LayoutMap& selectedResults, + const spatial::SpatialTargetInfo& target) { + int64_t cost = alternative.intrinsicCost; + SmallVector operandLayouts = getOperandLayouts(op, layouts); + for (auto [actual, required] : llvm::zip(operandLayouts, alternative.operandLayouts)) + cost += actual != required; + + Value result = op->getResult(0); + for (OpOperand& use : result.getUses()) { + auto user = dyn_cast(use.getOwner()); + if (!user) { + if (alternative.resultLayout != spatial::PhysicalLayout::DenseNCHW) { + auto flatten = dyn_cast(use.getOwner()); + if (!flatten || failed(canLowerFlattenFromRowStrip(flatten, target))) + ++cost; + } + continue; + } + auto userAlternatives = getAlternatives(use.getOwner(), selectedResults, target); + if (failed(userAlternatives)) + continue; + spatial::PhysicalLayout userResult = + selectedResults.lookup(use.getOwner()->getResult(0)); + unsigned userIndex = findCurrentAlternative(use.getOwner(), *userAlternatives, userResult); + if (use.getOperandNumber() < (*userAlternatives)[userIndex].operandLayouts.size() + && (*userAlternatives)[userIndex].operandLayouts[use.getOperandNumber()] + != alternative.resultLayout) + ++cost; + } + return cost; +} + +static LogicalResult materializeMismatchedUses( + IRRewriter& rewriter, Value value, const LayoutMap& layouts, + const spatial::SpatialTargetInfo& target) { + spatial::PhysicalLayout sourceLayout = getSelectedLayout(layouts, value); + SmallVector> mismatches; + for (OpOperand& use : value.getUses()) { + Operation* userOp = use.getOwner(); + spatial::PhysicalLayout required = spatial::PhysicalLayout::DenseNCHW; + if (auto capability = dyn_cast(userOp)) { + auto alternatives = getAlternatives(userOp, layouts, target); + if (failed(alternatives)) + return failure(); + spatial::PhysicalLayout selected = + getSelectedLayout(layouts, userOp->getResult(0)); + unsigned selectedIndex = findCurrentAlternative(userOp, *alternatives, selected); + required = (*alternatives)[selectedIndex].operandLayouts[use.getOperandNumber()]; + } + else if (auto flatten = dyn_cast(userOp); + flatten && sourceLayout == spatial::PhysicalLayout::NHWCRowStrip + && succeeded(canLowerFlattenFromRowStrip(flatten, target))) { + continue; + } + if (required != sourceLayout) + mismatches.push_back({&use, required}); + } + + for (auto [use, required] : mismatches) { + Operation* userOp = use->getOwner(); + rewriter.setInsertionPoint(userOp); + auto materialized = spatial::SpatMaterializeLayoutOp::create( + rewriter, userOp->getLoc(), use->get().getType(), use->get(), + spatial::LogicalLayoutAttr::get( + rewriter.getContext(), spatial::LogicalLayout::NCHW), + spatial::PhysicalLayoutAttr::get(rewriter.getContext(), sourceLayout), + spatial::PhysicalLayoutAttr::get(rewriter.getContext(), + required)); + use->set(materialized.getResult()); + } + return success(); +} + +static LogicalResult verifySelectedLayouts( + ArrayRef planOps, const LayoutMap& layouts, + const spatial::SpatialTargetInfo& target) { + for (Operation* op : planOps) { + auto selected = spatial::getSelectedPhysicalLayout(op); + if (!selected) + return op->emitOpError("requires a selected physical layout"), failure(); + auto alternatives = getAlternatives(op, layouts, target); + if (failed(alternatives)) + return failure(); + if (llvm::none_of(*alternatives, [&](const spatial::LayoutAlternative& alternative) { + return alternative.resultLayout == *selected; + })) + return op->emitOpError("selected physical layout is not advertised by its layout contract"), failure(); + } + return success(); +} + +struct SpatialLayoutPlanningPass final + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass) + + StringRef getArgument() const override { return "spatial-layout-planning"; } + StringRef getDescription() const override { + return "Select Spatial layout alternatives and insert explicit reconciliation barriers."; + } + + SpatialLayoutPlanningPass() = default; + explicit SpatialLayoutPlanningPass(const spatial::SpatialTargetInfo& target) + : target(target), hasTarget(true) {} + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + if (!hasTarget) { + moduleOp.emitError("Spatial layout planning requires an injected SpatialTargetInfo"); + signalPassFailure(); + return; + } + auto entryFunc = getPimEntryFunc(moduleOp); + if (failed(entryFunc)) { + moduleOp.emitError("failed to locate the PIM entry function during Spatial layout planning"); + signalPassFailure(); + return; + } + + func::FuncOp funcOp = *entryFunc; + SmallVector planOps; + for (Operation& op : funcOp.getBody().front()) + if (isa(&op)) + planOps.push_back(&op); + + LayoutMap layouts; + for (Operation* op : planOps) + layouts[op->getResult(0)] = spatial::PhysicalLayout::DenseNCHW; + + const size_t maxRounds = 2 * planOps.size() + 1; + bool converged = false; + for (size_t round = 0; round < maxRounds && !converged; ++round) { + converged = true; + SmallVector order(planOps); + if (round % 2) + std::reverse(order.begin(), order.end()); + for (Operation* op : order) { + auto alternatives = getAlternatives(op, layouts, target); + if (failed(alternatives)) { + signalPassFailure(); + return; + } + spatial::PhysicalLayout current = layouts.lookup(op->getResult(0)); + unsigned currentIndex = findCurrentAlternative(op, *alternatives, current); + int64_t bestCost = alternativeCost( + op, (*alternatives)[currentIndex], layouts, layouts, target); + unsigned bestIndex = currentIndex; + for (auto [index, alternative] : llvm::enumerate(*alternatives)) { + int64_t cost = alternativeCost(op, alternative, layouts, layouts, target); + if (cost < bestCost) { + bestCost = cost; + bestIndex = index; + } + } + spatial::PhysicalLayout selected = (*alternatives)[bestIndex].resultLayout; + if (selected != current) { + layouts[op->getResult(0)] = selected; + converged = false; + } + } + } + if (!converged) { + moduleOp.emitError("Spatial layout selection did not converge within its bounded iteration budget"); + signalPassFailure(); + return; + } + IRRewriter rewriter(&getContext()); + for (Operation* op : planOps) { + op->setAttr(spatial::kSelectedLayoutAttrName, + spatial::PhysicalLayoutAttr::get( + rewriter.getContext(), layouts.lookup(op->getResult(0)))); + if (failed(materializeMismatchedUses(rewriter, op->getResult(0), layouts, target))) { + signalPassFailure(); + return; + } + } + if (failed(verifySelectedLayouts(planOps, layouts, target)) + || failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) { + moduleOp.emitError("Spatial layout planning verification failed"); + signalPassFailure(); + } + } + + spatial::SpatialTargetInfo target; + bool hasTarget = false; +}; + +} // namespace + +std::unique_ptr createSpatialLayoutPlanningPass() { + return std::make_unique(); +} + +std::unique_ptr createSpatialLayoutPlanningPass( + const spatial::SpatialTargetInfo& target) { + return std::make_unique(target); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/SpatialToPim/Common.cpp b/src/PIM/Conversion/SpatialToPim/Common.cpp index e4eddac..6d5dfe5 100644 --- a/src/PIM/Conversion/SpatialToPim/Common.cpp +++ b/src/PIM/Conversion/SpatialToPim/Common.cpp @@ -1,7 +1,10 @@ #include "mlir/IR/ValueRange.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" #include "llvm/ADT/STLExtras.h" @@ -28,6 +31,49 @@ FailureOr getTensorSizeInBytesAttr(Builder& builder, Operation* anc return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size"); } +LogicalResult materializePipelineHostBuffer( + func::FuncOp funcOp, RewriterBase &rewriter) { + auto bytes = funcOp->getAttrOfType( + kPipelineHostBufferBytesAttrName); + if (!bytes) + return success(); + if (bytes.getInt() <= 0) + return funcOp.emitOpError( + "pipeline host transfer buffer must be positive"); + ModuleOp moduleOp = funcOp->getParentOfType(); + if (moduleOp.lookupSymbol(kPipelineHostBufferName)) + return funcOp.emitOpError( + "pipeline host transfer buffer symbol already exists"); + auto type = MemRefType::get( + {bytes.getInt()}, rewriter.getI8Type()); + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(moduleOp.getBody()); + memref::GlobalOp::create( + rewriter, funcOp.getLoc(), + rewriter.getStringAttr(kPipelineHostBufferName), + rewriter.getStringAttr("private"), TypeAttr::get(type), Attribute(), + UnitAttr(), IntegerAttr()); + return success(); +} + +FailureOr getPipelineHostBuffer( + OpBuilder &builder, Operation *anchor) { + auto funcOp = anchor->getParentOfType(); + auto moduleOp = anchor->getParentOfType(); + auto bytes = funcOp + ? funcOp->getAttrOfType(kPipelineHostBufferBytesAttrName) + : IntegerAttr(); + auto global = moduleOp + ? moduleOp.lookupSymbol(kPipelineHostBufferName) + : memref::GlobalOp(); + if (!bytes || !global) + return anchor->emitOpError( + "requires the pipeline host transfer buffer"), failure(); + auto type = MemRefType::get({bytes.getInt()}, builder.getI8Type()); + return memref::GetGlobalOp::create( + builder, anchor->getLoc(), type, kPipelineHostBufferName).getResult(); +} + Operation* getEarliestUserWithinBlock(mlir::Value value) { auto users = value.getUsers(); diff --git a/src/PIM/Conversion/SpatialToPim/Common.hpp b/src/PIM/Conversion/SpatialToPim/Common.hpp index 61e74a4..dfa73a4 100644 --- a/src/PIM/Conversion/SpatialToPim/Common.hpp +++ b/src/PIM/Conversion/SpatialToPim/Common.hpp @@ -10,6 +10,7 @@ #include "mlir/IR/Builders.h" #include "mlir/IR/Value.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Support/LogicalResult.h" #include "src/Accelerators/PIM/Common/PimCommon.hpp" @@ -23,6 +24,12 @@ namespace onnx_mlir { mlir::FailureOr getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value); +mlir::LogicalResult materializePipelineHostBuffer( + mlir::func::FuncOp funcOp, mlir::RewriterBase &rewriter); + +mlir::FailureOr getPipelineHostBuffer( + mlir::OpBuilder &builder, mlir::Operation *anchor); + template size_t rangeLength(const mlir::iterator_range range) { return std::distance(range.begin(), range.end()); diff --git a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp b/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp index 318a1e2..4dc6101 100644 --- a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp @@ -345,20 +345,39 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom auto blockArg = computeOp.getInputArgument(inputIndex); if (!blockArg) return computeOp.emitOpError("expected compute input block arguments during lowering"); - auto receiveOp = dyn_cast_or_null(input.getDefiningOp()); + auto channelReceive = dyn_cast_or_null( + input.getDefiningOp()); + auto hostWaitLoad = dyn_cast_or_null( + input.getDefiningOp()); + Operation *receiveOp = channelReceive + ? channelReceive.getOperation() : hostWaitLoad.getOperation(); if (receiveOp && !blockArg->use_empty()) { rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg)); auto outputType = cast(blockArg->getType()); - auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType); + auto outputBuffer = createEmptyTensorFromShaped( + rewriter, receiveOp->getLoc(), outputType); auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg); if (failed(sizeAttr)) return failure(); - Value received = - PimReceiveOp::create( - rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer, - arith::ConstantIndexOp::create(rewriter, receiveOp.getLoc(), 0), - *sizeAttr, receiveOp.getSourceCoreId()) + Value zero = arith::ConstantIndexOp::create( + rewriter, receiveOp->getLoc(), 0); + Value received; + if (hostWaitLoad) { + auto hostBuffer = getPipelineHostBuffer(rewriter, hostWaitLoad); + if (failed(hostBuffer)) + return failure(); + PimWaitOp::create( + rewriter, receiveOp->getLoc(), hostWaitLoad.getEventRegister(), + rewriter.getI32IntegerAttr(1)); + received = PimMemCopyHostToDevOp::create( + rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero, + hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr) .getOutput(); + } else { + received = PimReceiveOp::create( + rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer, + zero, *sizeAttr, channelReceive.getSourceCoreId()).getOutput(); + } blockArg->replaceAllUsesWith(received); markOpToRemove(receiveOp); continue; @@ -383,7 +402,8 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom if (rangeLength(resultUses) == 1) { OpOperand& resultUse = *resultUses.begin(); Operation* resultUser = resultUse.getOwner(); - if (isa(resultUser)) + if (isa(resultUser)) continue; } diff --git a/src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp b/src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp index b1f0e4c..08d2902 100644 --- a/src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp +++ b/src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp @@ -57,10 +57,29 @@ struct ChannelSendLowering : OpRewritePattern { } }; -struct ChannelReceiveLowering : OpRewritePattern { +struct HostStoreSyncLowering : OpRewritePattern { using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override { + LogicalResult matchAndRewrite(spatial::SpatHostStoreSyncOp op, PatternRewriter& rewriter) const override { + auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput()); + auto hostBuffer = getPipelineHostBuffer(rewriter, op); + if (failed(sizeAttr) || failed(hostBuffer)) + return failure(); + Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0); + pim::PimMemCopyDevToHostOp::create( + rewriter, op.getLoc(), hostBuffer->getType(), op.getHostOffset(), zero, + *hostBuffer, op.getInput(), *sizeAttr); + auto sync = pim::PimSyncOp::create( + rewriter, op.getLoc(), op.getTargetCoreId(), op.getEventRegister()); + copyRaptorDebugAttrs(op.getOperation(), sync.getOperation()); + rewriter.eraseOp(op); + return success(); + } +}; + +template +static LogicalResult lowerReceive( + ReceiveOp op, PatternRewriter& rewriter, CreateReceive createReceive) { if (op->use_empty()) { rewriter.eraseOp(op); return success(); @@ -86,12 +105,11 @@ struct ChannelReceiveLowering : OpRewritePattern if (failed(sizeAttr)) return failure(); Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0); - auto receive = pim::PimReceiveOp::create( - rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, *sizeAttr, op.getSourceCoreId()); - copyRaptorDebugAttrs(op.getOperation(), receive.getOperation()); - Value received = receive.getOutput(); + auto received = createReceive(outputBuffer, zero, *sizeAttr); + if (failed(received)) + return failure(); if (!destinationInsert) { - rewriter.replaceOp(op, received); + rewriter.replaceOp(op, *received); return success(); } @@ -99,10 +117,42 @@ struct ChannelReceiveLowering : OpRewritePattern Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert); auto copy = pim::PimMemCopyOp::create( rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero, - destinationInsert.getDest(), received, *sizeAttr); + destinationInsert.getDest(), *received, *sizeAttr); rewriter.replaceOp(destinationInsert, copy.getOutput()); rewriter.eraseOp(op); return success(); +} + +struct ChannelReceiveLowering : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override { + return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr { + auto receive = pim::PimReceiveOp::create( + rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, + sizeAttr, op.getSourceCoreId()); + copyRaptorDebugAttrs(op.getOperation(), receive.getOperation()); + return receive.getOutput(); + }); + } +}; + +struct HostWaitLoadLowering : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatHostWaitLoadOp op, PatternRewriter& rewriter) const override { + return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr { + auto hostBuffer = getPipelineHostBuffer(rewriter, op); + if (failed(hostBuffer)) + return failure(); + auto wait = pim::PimWaitOp::create( + rewriter, op.getLoc(), op.getEventRegister(), + rewriter.getI32IntegerAttr(1)); + copyRaptorDebugAttrs(op.getOperation(), wait.getOperation()); + return pim::PimMemCopyHostToDevOp::create( + rewriter, op.getLoc(), outputBuffer.getType(), zero, + op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput(); + }); } }; @@ -148,7 +198,9 @@ struct ConcatLowering : OpRewritePattern { } // namespace void populateChannelLoweringPatterns(RewritePatternSet& patterns) { - patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); } } // namespace onnx_mlir diff --git a/src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp b/src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp index b33c6d0..9f802f3 100644 --- a/src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp +++ b/src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp @@ -859,6 +859,10 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret markOpToRemove(receiveOp); return; } + if (auto receiveOp = dyn_cast(op)) { + markOpToRemove(receiveOp); + return; + } }; SmallVector originalOperands(returnOp.getOperands().begin(), returnOp.getOperands().end()); diff --git a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp index b25db93..f9b4ebd 100644 --- a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp +++ b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp @@ -126,6 +126,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { spatial::SpatConcatOp, spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp, + spatial::SpatHostStoreSyncOp, + spatial::SpatHostWaitLoadOp, spatial::SpatExtractRowsOp>(); RewritePatternSet initialPatterns(ctx); @@ -140,6 +142,12 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { populateGlobalTensorMaterializationPatterns(globalTensorPatterns); walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns)); + if (funcOp->hasAttr(kPipelineHostBufferBytesAttrName) + && failed(materializePipelineHostBuffer(funcOp, rewriter))) { + signalPassFailure(); + return; + } + auto returnOp = cast(funcOp.front().getTerminator()); addReturnOutputBuffers(returnOp, rewriter); if (failed(allocateAndInitializeCoreLocalVariables(funcOp, rewriter))) { @@ -182,6 +190,17 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { continue; } } + SmallVector hostWaitLoadOps; + for (auto op : funcOp.getOps()) + hostWaitLoadOps.push_back(op); + for (auto op : hostWaitLoadOps) { + bool onlyPendingRemovalUsers = llvm::all_of( + op->getUsers(), [&](Operation* user) { + return llvm::is_contained(operationsToRemove, user); + }); + if (onlyPendingRemovalUsers) + markOpToRemove(op); + } RewritePatternSet coreBodyPatterns(ctx); populateCoreBodyPatterns(coreBodyPatterns); @@ -202,6 +221,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { spatial::SpatConcatOp, spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp, + spatial::SpatHostStoreSyncOp, + spatial::SpatHostWaitLoadOp, spatial::SpatExtractRowsOp>(); SmallVector coreOps; @@ -251,6 +272,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { communicationTarget.addIllegalOp(); RewritePatternSet communicationPatterns(ctx); diff --git a/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.cpp index 5ae8539..2e70a85 100644 --- a/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.cpp @@ -430,8 +430,7 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO auto targetBytes = getShapedByteSize(targetType); auto sourceBytes = getShapedByteSize(sourceType); - if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes) - && size <= *targetBytes && size <= *sourceBytes) { + if (succeeded(targetBytes) && succeeded(sourceBytes) && size <= *targetBytes && size <= *sourceBytes) { auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape()); auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape()); if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank) diff --git a/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp index d6285be..4db708c 100644 --- a/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp @@ -241,6 +241,8 @@ static bool isSupportedCoreInstructionOp(Operation* op) { pim::PimVMVOp, pim::PimReceiveOp, pim::PimSendOp, + pim::PimSyncOp, + pim::PimWaitOp, pim::PimConcatOp, pim::PimVMMOp, pim::PimVVAddOp, diff --git a/src/PIM/Dialect/Pim/Pim.td b/src/PIM/Dialect/Pim/Pim.td index 5134499..80a67d9 100644 --- a/src/PIM/Dialect/Pim/Pim.td +++ b/src/PIM/Dialect/Pim/Pim.td @@ -118,6 +118,32 @@ def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> { }]; } +def PimSyncOp : PimOp<"sync", []> { + let summary = "Signal an event register on another core"; + + let arguments = (ins + Index:$targetCoreId, + Index:$eventRegister + ); + + let assemblyFormat = [{ + $targetCoreId `event` $eventRegister attr-dict + }]; +} + +def PimWaitOp : PimOp<"wait", []> { + let summary = "Wait for an event register value"; + + let arguments = (ins + Index:$eventRegister, + I32Attr:$waitValue + ); + + let assemblyFormat = [{ + $eventRegister `value` $waitValue attr-dict + }]; +} + def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> { let summary = "Copy a memory region from host memory into device memory"; diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index fa34968..d46483f 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -35,6 +35,7 @@ add_pim_library(SpatialOps Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp Passes/Transforms/TrivialGraphComputeMergePass.cpp EXCLUDE_FROM_OM_LIBS diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp index 965e032..e685537 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp @@ -219,10 +219,12 @@ static void appendReceive(BoundaryProgram &boundary, run->entryOffsets[run->entryOffsets.size() - 2]].family->requirement; CollectionTarget previousTarget {run->collection, run->positions.back()}; bool sameEntry = previous == requirement; - if (sameEntry + bool sameRoute = run->slices.back().family->hostRouted + == slice.family->hostRouted; + if (sameRoute && (sameEntry || (sameCollectionEmissionContract(previousTarget, target) && previous->publicationFragmentType - == requirement->publicationFragmentType)) { + == requirement->publicationFragmentType))) { run->slices.push_back(slice); if (sameEntry) { run->entryOffsets.back() = run->slices.size(); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp index eca577b..28c0e14 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp @@ -8,6 +8,7 @@ #include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp" #include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp" #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" #include namespace onnx_mlir::spatial { using namespace mlir; @@ -18,6 +19,8 @@ struct LogicalTransferMetadataView { StaticIntSequenceChain parentCounts; StaticIntSequenceChain sourceCores; StaticIntSequenceChain targetCores; + StaticIntSequenceChain hostOffsets; + StaticIntSequenceChain eventRegisters; StaticIntSequenceChain targetLanes; StaticIntSequenceChain localOffsets; SmallVector projectionOffsets; @@ -28,7 +31,8 @@ struct LogicalTransferMetadataView { }; using MetadataMember = StaticIntSequenceChain LogicalTransferMetadataView::*; static constexpr std::array transferMetadataMembers{ - &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores, &LogicalTransferMetadataView::targetCores}; + &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores, + &LogicalTransferMetadataView::targetCores}; struct TransferGrids { std::array values; StaticIntGrid &channels() { return values[0]; } @@ -41,7 +45,8 @@ template static FailureOr buildTransferGrids(Bui auto targetCores = build(transferMetadataMembers[2]); if (failed(channels) || failed(sourceCores) || failed(targetCores)) return failure(); - return TransferGrids{{std::move(*channels), std::move(*sourceCores), std::move(*targetCores)}}; + return TransferGrids{{std::move(*channels), std::move(*sourceCores), + std::move(*targetCores)}}; } using GridGeometry = DeferredGridSliceGeometry; using StaticGeometryMember = SmallVector DeferredStaticSliceGeometry::*; @@ -82,6 +87,11 @@ static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferM metadata.parentCounts.append(StaticIntSequence::uniform(family.requirement->exchange->externalTransferCount, count)); metadata.sourceCores.append(family.sourceCores, familyIndex, count); metadata.targetCores.append(family.targetCores, familyIndex, count); + if (family.hostRouted) { + metadata.hostOffsets.append(family.hostOffsets, familyIndex, count); + metadata.eventRegisters.append( + family.eventRegisters, familyIndex, count); + } metadata.targetLanes.append(StaticIntSequence::affine(targetLane, 1, count)); if (family.requirement->producerLocalOffsets) metadata.localOffsets.append(*family.requirement->producerLocalOffsets, targetLane - requirementLanes.begin, count); @@ -172,6 +182,7 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la appendMetadata(slice, metadataByLane[sourceLane]); } LogicalTransferMetadataView logical = buildMetadataView(run.slices); + ExternalTransferFamily &firstFamily = *run.slices.front().family; size_t actionCount = 0; for (const LogicalTransferMetadataView &laneMetadata : metadataByLane) actionCount = std::max(actionCount, laneMetadata.size()); @@ -185,6 +196,20 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la FailureOr localOffsets = buildGrid(&LogicalTransferMetadataView::localOffsets, logical.localOffsets.valueAt(0)); if (failed(transferGrids) || failed(localOffsets)) return failure(); + std::optional hostOffsets; + std::optional eventRegisters; + if (firstFamily.hostRouted) { + auto offsets = buildGrid( + &LogicalTransferMetadataView::hostOffsets, + logical.hostOffsets.valueAt(0)); + auto events = buildGrid( + &LogicalTransferMetadataView::eventRegisters, + logical.eventRegisters.valueAt(0)); + if (failed(offsets) || failed(events)) + return failure(); + hostOffsets = std::move(*offsets); + eventRegisters = std::move(*events); + } GridGeometry projectionGrids; for (auto [geometryIndex, sourceMember] : llvm::enumerate(metadataGeometryMembers)) { const auto &logicalValues = logical.*sourceMember; @@ -207,7 +232,6 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la const LogicalTransferMetadataView &source = metadataByLane[sourceLane]; counts[sourceLane] = source.size(); } - ExternalTransferFamily &firstFamily = *run.slices.front().family; RequirementFamily &requirement = *firstFamily.requirement; Operation *anchor = requirement.exchange->deferred; Location loc = requirement.exchange->deferred.getLoc(); @@ -217,10 +241,25 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la auto payload = materializeSendPayload(requirement, localOffset, projectionGrids[0].empty() ? nullptr : &projection, context, loc); if (failed(payload)) return failure(); - auto send = SpatChannelSendOp::create( - context.rewriter, loc, transferGrids->channels().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), - transferGrids->sourceCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), - transferGrids->targetCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), *payload); + Value sourceCore = transferGrids->sourceCores().emitLookup( + action, runtimeLane, anchor, context.constants, context.rewriter, loc); + Value targetCore = transferGrids->targetCores().emitLookup( + action, runtimeLane, anchor, context.constants, context.rewriter, loc); + Operation *send; + if (firstFamily.hostRouted) + send = SpatHostStoreSyncOp::create( + context.rewriter, loc, sourceCore, targetCore, + hostOffsets->emitLookup( + action, runtimeLane, anchor, context.constants, context.rewriter, loc), + eventRegisters->emitLookup( + action, runtimeLane, anchor, context.constants, context.rewriter, loc), + *payload); + else + send = SpatChannelSendOp::create( + context.rewriter, loc, + transferGrids->channels().emitLookup( + action, runtimeLane, anchor, context.constants, context.rewriter, loc), + sourceCore, targetCore, *payload); setLogicalTransferMetadata(send, logical); return success(); }; @@ -255,14 +294,45 @@ static FailureOr emitReceiveValue(ArrayRef slices }; auto grids = buildTransferGrids([&](MetadataMember member) { return buildGrid(metadata.*member); }); if (failed(grids)) return failure(); + std::optional hostOffsets; + std::optional eventRegisters; + if (slices.front().family->hostRouted) { + auto offsets = buildGrid(metadata.hostOffsets); + auto events = buildGrid(metadata.eventRegisters); + if (failed(offsets) || failed(events)) + return failure(); + hostOffsets = std::move(*offsets); + eventRegisters = std::move(*events); + } Value position = lane ? lane : context.constants.getIndex(0); Value row = context.constants.getIndex(0); - auto receive = SpatChannelReceiveOp::create(context.rewriter, anchor->getLoc(), requirement.publicationFragmentType, - grids->channels().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()), - grids->sourceCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()), - grids->targetCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc())); + Value sourceCore = grids->sourceCores().emitLookup( + row, position, anchor, context.constants, context.rewriter, anchor->getLoc()); + Value targetCore = grids->targetCores().emitLookup( + row, position, anchor, context.constants, context.rewriter, anchor->getLoc()); + Operation *receive; + Value output; + if (slices.front().family->hostRouted) { + auto op = SpatHostWaitLoadOp::create( + context.rewriter, anchor->getLoc(), requirement.publicationFragmentType, + sourceCore, targetCore, + hostOffsets->emitLookup( + row, position, anchor, context.constants, context.rewriter, anchor->getLoc()), + eventRegisters->emitLookup( + row, position, anchor, context.constants, context.rewriter, anchor->getLoc())); + receive = op; + output = op.getOutput(); + } else { + auto op = SpatChannelReceiveOp::create( + context.rewriter, anchor->getLoc(), requirement.publicationFragmentType, + grids->channels().emitLookup( + row, position, anchor, context.constants, context.rewriter, anchor->getLoc()), + sourceCore, targetCore); + receive = op; + output = op.getOutput(); + } setLogicalTransferMetadata(receive, metadata); - return receive.getOutput(); + return output; } static FailureOr> @@ -315,6 +385,9 @@ static FailureOr emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V SmallVector counts(laneCount); std::optional transferGrids; std::optional positions; + std::optional hostOffsets; + std::optional eventRegisters; + bool hostRouted = run.slices.front().family->hostRouted; auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount); if (succeeded(metadataByEntry)) { auto buildRows = [&](auto member) { @@ -324,6 +397,16 @@ static FailureOr emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V return StaticIntGrid::fromRows(rows); }; auto grids = buildTransferGrids(buildRows); + if (hostRouted) { + auto offsets = buildRows( + &LogicalTransferMetadataView::hostOffsets); + auto events = buildRows( + &LogicalTransferMetadataView::eventRegisters); + if (failed(offsets) || failed(events)) + return failure(); + hostOffsets = std::move(*offsets); + eventRegisters = std::move(*events); + } SmallVector positionRows; for (unsigned position : run.positions) positionRows.push_back(StaticIntSequence::uniform(position, laneCount)); @@ -368,6 +451,16 @@ static FailureOr emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V return StaticIntGrid::fromColumns(actionCount, columns, defaultValue); }; auto grids = buildTransferGrids(buildGrid); + if (hostRouted) { + auto offsets = buildGrid( + &LogicalTransferMetadataView::hostOffsets); + auto events = buildGrid( + &LogicalTransferMetadataView::eventRegisters); + if (failed(offsets) || failed(events)) + return failure(); + hostOffsets = std::move(*offsets); + eventRegisters = std::move(*events); + } SmallVector positionColumns; for (const StaticIntSequenceChain &values : positionsByLane) positionColumns.push_back( @@ -386,15 +479,34 @@ static FailureOr emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V Value runtimeLane = lane ? lane : context.constants.getIndex(0); auto emitEntry = [&](Value entry, Value current) -> FailureOr { Type fragmentType = run.slices.front().family->requirement->publicationFragmentType; - auto receive = - SpatChannelReceiveOp::create(context.rewriter, loc, fragmentType, - transferGrids->channels().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc), - transferGrids->sourceCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc), - transferGrids->targetCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc)); + Value sourceCore = transferGrids->sourceCores().emitLookup( + entry, runtimeLane, anchor, context.constants, context.rewriter, loc); + Value targetCore = transferGrids->targetCores().emitLookup( + entry, runtimeLane, anchor, context.constants, context.rewriter, loc); + Operation *receive; + Value output; + if (hostRouted) { + auto op = SpatHostWaitLoadOp::create( + context.rewriter, loc, fragmentType, sourceCore, targetCore, + hostOffsets->emitLookup( + entry, runtimeLane, anchor, context.constants, context.rewriter, loc), + eventRegisters->emitLookup( + entry, runtimeLane, anchor, context.constants, context.rewriter, loc)); + receive = op; + output = op.getOutput(); + } else { + auto op = SpatChannelReceiveOp::create( + context.rewriter, loc, fragmentType, + transferGrids->channels().emitLookup( + entry, runtimeLane, anchor, context.constants, context.rewriter, loc), + sourceCore, targetCore); + receive = op; + output = op.getOutput(); + } setLogicalTransferMetadata(receive, logical); Value position = positions->emitLookup( entry, runtimeLane, anchor, context.constants, context.rewriter, loc); - return insert(receive.getOutput(), position, entry, runtimeLane, current); + return insert(output, position, entry, runtimeLane, current); }; if (actionCount == 1 && llvm::all_of(counts, [](int64_t count) { return count == 1; })) return emitEntry(context.constants.getIndex(0), initial); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp index ccb8670..697d89b 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp @@ -28,6 +28,11 @@ static std::optional getPlannedHead( while (cursor.slice < plan.slices.size()) { const ScheduledTransferSlice &slice = plan.slices[cursor.slice]; ExternalTransferFamily &family = *slice.family; + if (family.hostRouted) { + ++cursor.slice; + cursor.offset = 0; + continue; + } size_t begin = slice.familyOffset + cursor.offset; size_t length = slice.transferCount - cursor.offset; auto source = family.sourceStreams.find(stream, begin, length); @@ -243,6 +248,8 @@ LogicalResult verifyPlannedCommunicationDeadlockFree( DenseMap familyIndex; for (const ScheduledTransferSlice &slice : plan.slices) { ExternalTransferFamily *family = slice.family; + if (family->hostRouted) + continue; if (!familyIndex.try_emplace(family, familyIndex.size()).second) continue; size_t count = family->channelIds.size(); @@ -258,18 +265,6 @@ LogicalResult verifyPlannedCommunicationDeadlockFree( familyChannels.emplace_back( first, first + static_cast(count)); } - llvm::sort(familyChannels); - int64_t nextChannel = 0; - for (auto [firstChannel, endChannel] : familyChannels) { - if (firstChannel != nextChannel) - return anchor->emitError( - "planned communication channels are not exactly contiguous"); - nextChannel = endChannel; - } - if (static_cast(nextChannel) != plan.logicalTransferCount) - return anchor->emitError( - "planned communication channel count is inconsistent"); - for (const ScheduledTransferSlice &slice : plan.slices) { ExternalTransferFamily &family = *slice.family; for (size_t offset = 0; offset < slice.transferCount; ++offset) { @@ -296,6 +291,8 @@ LogicalResult verifyRealizedCommunicationDeadlockFree( DenseMap familyIndex; for (const ScheduledTransferSlice &slice : plan.slices) { ExternalTransferFamily *family = slice.family; + if (family->hostRouted) + continue; if (!familyIndex.try_emplace(family, familyIndex.size()).second) continue; for (size_t index = 0; index < family->channelIds.size(); ++index) @@ -305,6 +302,8 @@ LogicalResult verifyRealizedCommunicationDeadlockFree( DenseMap expected; for (const ScheduledTransferSlice &slice : plan.slices) { ExternalTransferFamily &family = *slice.family; + if (family.hostRouted) + continue; appendEventsByCore(expected, family.channelIds, family.sourceCores, slice.familyOffset, slice.transferCount, true); appendEventsByCore(expected, family.channelIds, family.targetCores, diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp index 4c734a8..6ac6617 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp @@ -198,6 +198,7 @@ struct ScheduledInfo { llvm::SmallVector blocks; llvm::SmallVector stepAnchors; llvm::SmallVector cores; + llvm::SmallVector pipelineStages; unsigned stepCount = 0; llvm::SmallVector produced; llvm::SmallVector streamIds; @@ -233,6 +234,9 @@ struct ExternalTransferFamily { StaticIntSequence sourceCores = StaticIntSequence::uniform(0, 1); StaticIntSequence targetCores = StaticIntSequence::uniform(0, 1); StaticIntSequence channelIds = StaticIntSequence::uniform(0, 1); + StaticIntSequence hostOffsets = StaticIntSequence::uniform(0, 1); + StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1); + bool hostRouted = false; }; struct DeferredExchangePlan { diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp index cdf9f11..a64dbbe 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp @@ -209,15 +209,26 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) { LogicalResult realizeDeferredCommunication(func::FuncOp funcOp, const ScheduledComputeMaterializationResult& materialization, - const SchedulingTarget& target) { + const SchedulingTarget& target, + size_t pipelineStages) { IRRewriter rewriter(funcOp.getContext()); eraseUnusedIdentityDeferredCommunications(funcOp, rewriter); - auto transfers = buildDeferredTransferPlan(funcOp, materialization); + auto transfers = buildDeferredTransferPlan( + funcOp, materialization, pipelineStages, target.processorCount); if (failed(transfers)) return funcOp.emitOpError("phase 2 failed to build symbolic transfer families"); if (failed(placeLogicalProcessorsOnPhysicalCores(*transfers, target))) return failure(); + if (transfers->pipelineHostBufferBytes != 0) { + auto bytes = pim::checkedCast( + transfers->pipelineHostBufferBytes, funcOp, + "pipeline host transfer storage"); + if (failed(bytes)) + return failure(); + funcOp->setAttr(kPipelineHostBufferBytesAttrName, + rewriter.getI64IntegerAttr(*bytes)); + } auto schedule = scheduleDeferredCommunication(funcOp, *transfers); if (failed(schedule) || failed(verifyPlannedCommunicationDeadlockFree(funcOp, transfers->stepCounts, *schedule))) return funcOp.emitOpError("phase 2 failed to schedule symbolic communication"); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp index 76795e2..7a19be1 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp @@ -9,6 +9,7 @@ struct SchedulingTarget; mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp, const ScheduledComputeMaterializationResult& materialization, - const SchedulingTarget& target); + const SchedulingTarget& target, + size_t pipelineStages = 1); } // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp index 830d960..58f9a61 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp @@ -11,7 +11,7 @@ using namespace mlir; namespace { using TransferEmissionSignature = - std::tuple; + std::tuple; static TransferEmissionSignature getTransferEmissionSignature( const ExternalTransferFamily& family) { @@ -21,7 +21,8 @@ static TransferEmissionSignature getTransferEmissionSignature( family.requirement->publicationFragmentType, family.requirement->graphLanes.has_value(), family.requirement->producerProjection.has_value(), - producer->scheduled->isBatch()}; + producer->scheduled->isBatch(), + family.hostRouted}; } struct StreamThreshold { diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp index 9cd9444..a8ebdac 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp @@ -5,6 +5,7 @@ #include "DeferredProjectionAnalysis.hpp" #include "DeferredTransferPlanning.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" namespace onnx_mlir::spatial { using namespace mlir; @@ -28,7 +29,12 @@ static FailureOr getStepIndex( static LogicalResult collectScheduledOperations( const ScheduledComputeMaterializationResult &materialization, - DeferredTransferPlan &plan) { + DeferredTransferPlan &plan, + size_t pipelineStageCount, + size_t processorCount) { + if (pipelineStageCount == 0 || processorCount % pipelineStageCount != 0) + return failure(); + size_t stageSize = processorCount / pipelineStageCount; unsigned nextStream = 0; for (const ScheduledMaterializationRecord &record : materialization.materializedSchedules) { @@ -46,8 +52,13 @@ static LogicalResult collectScheduledOperations( if (llvm::any_of(info.stepAnchors, [](Operation *anchor) { return !anchor; })) return op.emitOpError("phase 2 scheduled step anchor is missing"); - for (size_t core : record.cpus) + for (size_t core : record.cpus) { + if (core >= processorCount) + return op.emitOpError("phase 2 scheduled core is outside the target"); info.cores.push_back(core); + if (pipelineStageCount > 1) + info.pipelineStages.push_back(core / stageSize); + } for (size_t lane = 0; lane < info.cores.size(); ++lane) info.streamIds.push_back(nextStream++); plan.scheduled.push_back(std::move(info)); @@ -308,17 +319,22 @@ static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan, return success(); } -static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t& nextChannel) { +static LogicalResult buildAvailabilityFamilies( + DeferredTransferPlan &plan, + DeferredExchangePlan& exchange, + uint64_t& nextChannel, + DenseMap>& eventRegistersByTarget) { + enum class Availability { Local, Direct, Host }; for (RequirementFamily& requirement : exchange.requirements) { for (LaneInterval interval : requirement.targetLanes.intervals()) { unsigned runBegin = interval.begin; - bool runLocal = false; + Availability runAvailability = Availability::Local; bool haveRun = false; - auto flush = [&](unsigned end) { + auto flush = [&](unsigned end) -> LogicalResult { if (!haveRun || runBegin == end) - return; + return success(); LaneSet lanes = LaneSet::range(runBegin, end); - if (runLocal) { + if (runAvailability == Availability::Local) { exchange.local.push_back({&requirement, lanes}); } else { @@ -339,25 +355,72 @@ static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t& family.sourceCores = StaticIntSequence::uniform(requirement.producer->core, count); family.targetCores = StaticIntSequence::fromValues(targetCores); family.channelIds = StaticIntSequence::affine(nextChannel, 1, count); + family.hostRouted = runAvailability == Availability::Host; + if (family.hostRouted) { + SmallVector eventRegisters; + for (int64_t targetCore : targetCores) { + auto ®isters = eventRegistersByTarget[targetCore]; + auto it = registers.try_emplace( + requirement.producer->core, registers.size()).first; + if (it->second >= kPimEventRegisterCount) + return exchange.deferred.emitOpError( + "pipeline host transfer requires more event registers than the target core provides"); + eventRegisters.push_back(it->second); + } + family.eventRegisters = StaticIntSequence::fromValues( + eventRegisters); + auto fragmentType = dyn_cast( + requirement.publicationFragmentType); + auto fragmentBytes = fragmentType + ? pim::getCheckedShapedTypeSizeInBytes( + fragmentType, exchange.deferred, + "pipeline host transfer fragment") + : FailureOr(failure()); + if (failed(fragmentBytes)) + return failure(); + auto bytes = pim::checkedMul( + count, static_cast(*fragmentBytes), exchange.deferred, + "pipeline host transfer storage"); + if (failed(bytes)) + return failure(); + family.hostOffsets = StaticIntSequence::affine( + plan.pipelineHostBufferBytes, *fragmentBytes, count); + auto endOffset = pim::checkedAdd( + plan.pipelineHostBufferBytes, *bytes, exchange.deferred, + "pipeline host transfer storage"); + if (failed(endOffset)) + return failure(); + plan.pipelineHostBufferBytes = *endOffset; + } nextChannel += count; exchange.externalTransferCount += count; exchange.external.push_back(std::move(family)); } + return success(); }; for (unsigned lane = interval.begin; lane < interval.end; ++lane) { unsigned sourceStream = requirement.producer->scheduled->streamIds[requirement.producer->scheduledLane]; bool local = sourceStream == exchange.target->streamIds[lane] && requirement.producer->step < exchange.consumerStep; - if (haveRun && local != runLocal) { - flush(lane); + bool crossStage = !exchange.target->pipelineStages.empty() + && requirement.producer->scheduled->pipelineStages[ + requirement.producer->scheduledLane] + != exchange.target->pipelineStages[lane]; + Availability availability = local ? Availability::Local + : crossStage ? Availability::Host : Availability::Direct; + if (haveRun && availability != runAvailability) { + if (failed(flush(lane))) + return failure(); runBegin = lane; } - runLocal = local; + runAvailability = availability; haveRun = true; } - flush(interval.end); + if (failed(flush(interval.end))) + return failure(); } } + return success(); } static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& plan) { @@ -368,6 +431,7 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p funcOp.walk([&](SpatDeferredCommunicationOp op) { deferredOps.push_back(op); }); GraphBatchPublicationCache publicationCache; uint64_t nextChannel = 0; + DenseMap> eventRegistersByTarget; for (SpatDeferredCommunicationOp deferred : deferredOps) { Operation* targetOp = deferred->getParentOfType(); if (!targetOp) @@ -387,7 +451,9 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p exchange->program = std::move(*program); if (failed(buildRequirementFamilies(plan, *exchange, publicationCache))) return failure(); - buildAvailabilityFamilies(*exchange, nextChannel); + if (failed(buildAvailabilityFamilies( + plan, *exchange, nextChannel, eventRegistersByTarget))) + return failure(); plan.exchanges.push_back(std::move(exchange)); } return success(); @@ -464,9 +530,12 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa FailureOr buildDeferredTransferPlan( func::FuncOp funcOp, - const ScheduledComputeMaterializationResult &materialization) { + const ScheduledComputeMaterializationResult &materialization, + size_t pipelineStages, + size_t processorCount) { DeferredTransferPlan plan; - if (failed(collectScheduledOperations(materialization, plan)) + if (failed(collectScheduledOperations( + materialization, plan, pipelineStages, processorCount)) || failed(collectProducedValues(materialization, plan)) || failed(buildExchanges(funcOp, plan))) return failure(); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp index e3c0038..0b26fd3 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp @@ -13,11 +13,14 @@ struct DeferredTransferPlan { llvm::DenseMap> producedByGraph; llvm::SmallVector> exchanges; llvm::SmallVector stepCounts; + size_t pipelineHostBufferBytes = 0; }; mlir::FailureOr buildDeferredTransferPlan(mlir::func::FuncOp funcOp, - const ScheduledComputeMaterializationResult &materialization); + const ScheduledComputeMaterializationResult &materialization, + size_t pipelineStages, + size_t processorCount); mlir::LogicalResult retargetDeferredPublications(mlir::func::FuncOp funcOp, DeferredTransferPlan& plan); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp index 578bc25..f6fd9fd 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp @@ -3,12 +3,15 @@ #include "DeferredCommunicationRealization.hpp" #include "ScheduledComputeReport.hpp" #include "ScheduledComputeVerification.hpp" +#include "Scheduling/PipelineScheduling.hpp" #include "SpatialDataflowCsvExporter.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp" #include "src/Accelerators/PIM/Passes/PIMPasses.h" +#include + using namespace mlir; namespace onnx_mlir { @@ -31,14 +34,47 @@ static FailureOr requireEntry(ModuleOp moduleOp) { return *entry; } +static SchedulingTarget getPipelineSchedulingTarget( + const SchedulingTarget& physicalTarget, size_t pipelineStages) { + if (pipelineStages == 1) + return physicalTarget; + + SchedulingTarget schedulingTarget = physicalTarget; + schedulingTarget.processorCount = physicalTarget.processorCount / pipelineStages; + schedulingTarget.residentWeightCapacity = checkedMultiply( + physicalTarget.residentWeightCapacity, pipelineStages); + schedulingTarget.interProcessorLatencyNs.assign( + schedulingTarget.processorCount * schedulingTarget.processorCount, 0); + Cost latencySum = 0; + size_t pairCount = 0; + for (size_t source = 0; source < schedulingTarget.processorCount; ++source) + for (size_t destination = 0; + destination < schedulingTarget.processorCount; ++destination) { + Cost latency = physicalTarget.getInterProcessorLatencyNs( + source, destination); + schedulingTarget.interProcessorLatencyNs[ + source * schedulingTarget.processorCount + destination] = latency; + if (source != destination) { + latencySum = checkedAdd(latencySum, latency); + ++pairCount; + } + } + schedulingTarget.averageInterProcessorLatencyNs = pairCount == 0 + ? 0 + : (latencySum + pairCount - 1) / pairCount; + return schedulingTarget; +} + struct ScheduleAndRealizeSpatialPass final : PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleAndRealizeSpatialPass) ScheduleAndRealizeSpatialPass() = default; ScheduleAndRealizeSpatialPass(const SchedulingTarget& target, - SpatialDataflowExportStage exportStage) - : target(target), exportStage(exportStage), hasTarget(true) {} + SpatialDataflowExportStage exportStage, + size_t pipelineStages) + : target(target), exportStage(exportStage), + pipelineStages(pipelineStages), hasTarget(true) {} StringRef getArgument() const override { return "schedule-and-realize-spatial"; } StringRef getDescription() const override { @@ -52,6 +88,13 @@ struct ScheduleAndRealizeSpatialPass final signalPassFailure(); return; } + if (pipelineStages == 0 || target.processorCount % pipelineStages != 0 + || target.residentWeightCapacity + > std::numeric_limits::max() / pipelineStages) { + moduleOp.emitError("ScheduleAndRealizeSpatial requires valid pipeline stages and resource counts"); + signalPassFailure(); + return; + } auto entry = requireEntry(moduleOp); if (failed(entry)) { signalPassFailure(); @@ -59,8 +102,31 @@ struct ScheduleAndRealizeSpatialPass final } func::FuncOp entryFunc = *entry; - MergeSchedulingAnalysis analysis(entryFunc, target); - MergeScheduleResult schedule = std::move(analysis.getResult()); + SchedulingTarget schedulingTarget = getPipelineSchedulingTarget( + target, pipelineStages); + ComputeGraph scheduledGraph; + MergeScheduleResult schedule; + for (;;) { + MergeSchedulingAnalysis analysis( + entryFunc, schedulingTarget, + pipelineStages > 1 ? target.processorCount : 0); + scheduledGraph = analysis.getGraph(); + schedule = std::move(analysis.getResult()); + std::string pipelineError; + if (succeeded(applyPipelineScheduling( + scheduledGraph, schedule, pipelineStages, target, pipelineError))) + break; + std::string splitError; + if (pipelineStages == 1 + || failed(splitPipelineWorkload( + scheduledGraph, schedule, pipelineStages, target, splitError))) { + if (!splitError.empty()) + pipelineError = splitError; + moduleOp.emitError() << pipelineError; + signalPassFailure(); + return; + } + } PatternRewriter rewriter(moduleOp.getContext()); FailureOr materialization = materializeScheduledCompute(entryFunc, schedule, rewriter); @@ -94,7 +160,8 @@ struct ScheduleAndRealizeSpatialPass final moduleOp, entryFunc, schedule, materializationResult.peftClassPlans, materializationResult.materializedSchedules); - if (failed(realizeDeferredCommunication(entryFunc, materializationResult, target))) { + if (failed(realizeDeferredCommunication( + entryFunc, materializationResult, target, pipelineStages))) { moduleOp.emitError("Spatial communication realization failed"); signalPassFailure(); return; @@ -126,6 +193,7 @@ struct ScheduleAndRealizeSpatialPass final private: SchedulingTarget target; SpatialDataflowExportStage exportStage = SpatialDataflowExportStage::None; + size_t pipelineStages = 1; bool hasTarget = false; }; @@ -136,8 +204,11 @@ std::unique_ptr createScheduleAndRealizeSpatialPass() { } std::unique_ptr createScheduleAndRealizeSpatialPass( - const SchedulingTarget& target, SpatialDataflowExportStage exportStage) { - return std::make_unique(target, exportStage); + const SchedulingTarget& target, + SpatialDataflowExportStage exportStage, + size_t pipelineStages) { + return std::make_unique( + target, exportStage, pipelineStages); } } // namespace spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp new file mode 100644 index 0000000..4c98411 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include "ScheduledComputeMaterialization.hpp" +#include "Scheduling/MergeSchedulingAnalysis.hpp" + +#include +#include + +namespace onnx_mlir::spatial { + +struct ScheduledSpatialState { + std::optional logicalSchedule; + std::optional materialization; +}; + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp index b11cda5..85c6f27 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp @@ -772,6 +772,11 @@ std::vector aggregateEdges(llvm::ArrayRef ed } // namespace +TransferCost getTransferCostFromBytes(Cost bytes, + const SchedulingTarget& target) { + return SchedulerCostModel {target}.getTransferCostFromBytes(bytes); +} + uint64_t countComputeBodyInstructions(Region& body) { uint64_t numOperations = 0; body.walk([&](Operation* op) { numOperations = checkedAdd(numOperations, static_cast(1)); }); @@ -875,9 +880,13 @@ ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& insta return tiled; } -ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& target) { +ComputeGraph buildComputeGraph(Operation* entryOp, + const SchedulingTarget& target, + size_t computePartitionCount) { ComputeGraph graph; SchedulerCostModel costModel {target}; + if (computePartitionCount == 0) + computePartitionCount = target.processorCount; for (Region& region : entryOp->getRegions()) { for (Block& block : region) { @@ -898,10 +907,10 @@ ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& targe if (isUsedAsWeightOnly(batch.getOperation())) continue; size_t chunkCount = - getBatchChunkTargetCount(batch, target.processorCount); + getBatchChunkTargetCount(batch, computePartitionCount); for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) { ComputeInstance instance = getBatchChunkForIndex( - batch, chunkIndex, target.processorCount); + batch, chunkIndex, computePartitionCount); size_t index = graph.nodes.size(); graph.nodes.push_back({instance, getComputeInstanceCost(instance, target), @@ -920,7 +929,7 @@ ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& targe for (Value input : inputs) { for (const ProducerValueRef& producerRef : collectProducerValueRefs(input, node.instance, - target.processorCount)) { + computePartitionCount)) { auto producerIt = graph.instanceToIndex.find(producerRef.instance); if (producerIt == graph.instanceToIndex.end()) continue; diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp index 4ef76d7..3d65845 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp @@ -61,9 +61,13 @@ struct ComputeGraph { llvm::DenseMap instanceToIndex; }; -ComputeGraph buildComputeGraph(mlir::Operation* entryOp, const SchedulingTarget& target); +ComputeGraph buildComputeGraph(mlir::Operation* entryOp, + const SchedulingTarget& target, + size_t computePartitionCount = 0); bool verifyAcyclic(const ComputeGraph& graph); +TransferCost getTransferCostFromBytes(Cost bytes, + const SchedulingTarget& target); uint64_t countComputeBodyInstructions(mlir::Region& body); uint64_t countComputeBodyOperationInstances(mlir::Region& body); Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp index e4b6fe5..f9433a1 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp @@ -89,13 +89,14 @@ void verifySchedule(const ComputeGraph& graph, } // namespace MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op, - const SchedulingTarget& schedulingTarget) -: entryOp(op), target(schedulingTarget) { + const SchedulingTarget& schedulingTarget, + size_t partitionCount) +: entryOp(op), target(schedulingTarget), computePartitionCount(partitionCount) { result = run(); } MergeScheduleResult MergeSchedulingAnalysis::run() { - ComputeGraph graph = buildComputeGraph(entryOp, target); + graph = buildComputeGraph(entryOp, target, computePartitionCount); if (!verifyAcyclic(graph)) llvm::report_fatal_error("merge scheduling: compute graph is cyclic"); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp index 5b5e6ed..f22d65d 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp @@ -3,6 +3,7 @@ #include "mlir/IR/Operation.h" #include "MergeSchedule.hpp" +#include "ComputeGraph.hpp" #include "SchedulingTarget.hpp" namespace onnx_mlir { @@ -10,12 +11,17 @@ namespace spatial { class MergeSchedulingAnalysis { public: - MergeSchedulingAnalysis(mlir::Operation* op, const SchedulingTarget& target); + MergeSchedulingAnalysis(mlir::Operation* op, + const SchedulingTarget& target, + size_t computePartitionCount = 0); MergeScheduleResult& getResult() { return result; } + const ComputeGraph& getGraph() const { return graph; } private: mlir::Operation* entryOp = nullptr; const SchedulingTarget& target; + size_t computePartitionCount = 0; + ComputeGraph graph; MergeScheduleResult result; MergeScheduleResult run(); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp new file mode 100644 index 0000000..be2e30c --- /dev/null +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp @@ -0,0 +1,1092 @@ +#include "llvm/ADT/STLExtras.h" + +#include +#include +#include +#include +#include +#include + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/PatternMatch.h" + +#include "ComputeInstanceUtils.hpp" +#include "PeftScheduler.hpp" +#include "PipelineScheduling.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +namespace onnx_mlir::spatial { +namespace { + +using TaskList = std::vector; +using TaskCosts = std::vector; + +using namespace mlir; + +struct SplitBody { + scf::ForOp loop; + SmallVector vmms; +}; + +struct BlueprintAssemblyInfo { + Cost schedulingCost = 0; + Cost work = 0; + TaskList predecessors; +}; + +struct PipelineTaskModel { + TaskCosts schedulingCosts; + TaskCosts assemblyCosts; + std::vector predecessors; +}; + +static bool fits(const ComputeGraph &graph, + const TaskCosts &taskCosts, + const TaskList &tasks, + Cost maximumCost, + size_t residentWeightCapacity, + size_t maximumPacks); +static Cost findMaximumPackCost(const ComputeGraph &graph, + const TaskCosts &taskCosts, + const TaskList &tasks, + size_t residentWeightCapacity, + size_t pipelineStages); + +static BlueprintAssemblyInfo getBlueprintAssemblyInfo( + const ComputeGraph &graph, const ComputeGraphNode &node, + size_t consumerCpu, const MergeScheduleResult &schedule, + const SchedulingTarget &target) { + BlueprintAssemblyInfo info; + if (!node.instance.op) + return info; + for (Value input : getComputeInstanceInputs(node.instance)) { + auto blueprint = input.getDefiningOp(); + if (!blueprint || !isFragmentAssembly(blueprint.getMode())) + continue; + + auto operandIndices = blueprint.getFragmentOperandIndices(); + auto sourceSlots = blueprint.getFragmentSourceSlots(); + auto resultType = dyn_cast(blueprint.getOutput().getType()); + if (!operandIndices || !sourceSlots || !resultType + || !resultType.hasStaticShape()) + continue; + int64_t rank = resultType.getRank(); + ArrayRef fragmentSizes = blueprint.getFragmentSizes(); + if (rank <= 0 + || fragmentSizes.size() != operandIndices->size() * rank + || sourceSlots->size() != operandIndices->size()) + continue; + + SmallVector sources {blueprint.getInput()}; + llvm::append_range(sources, blueprint.getFragments()); + for (auto [fragment, operandIndex] : llvm::enumerate(*operandIndices)) { + if (operandIndex < 0 + || operandIndex >= static_cast(sources.size())) + continue; + auto producer = dyn_cast_or_null( + sources[operandIndex].getDefiningOp()); + int64_t sourceSlot = (*sourceSlots)[fragment]; + if (!producer || sourceSlot < 0 + || sourceSlot >= producer.getLaneCount()) + continue; + ComputeInstance producerInstance = getBatchChunkForLane( + producer, static_cast(sourceSlot), target.processorCount); + auto producerTask = graph.instanceToIndex.find(producerInstance); + if (producerTask == graph.instanceToIndex.end()) + continue; + if (!llvm::is_contained(info.predecessors, producerTask->second)) + info.predecessors.push_back(producerTask->second); + + Cost elements = 1; + for (int64_t dimension = 0; dimension < rank; ++dimension) + elements = checkedMultiply( + elements, static_cast(fragmentSizes[fragment * rank + dimension])); + Cost bytes = (checkedMultiply(elements, target.computeBitwidth) + 7) / 8; + TransferCost transferCost = getTransferCostFromBytes(bytes, target); + size_t producerCpu = schedule.computeToCpuMap.lookup(producerInstance); + Cost transferCycles = (producerCpu == consumerCpu + ? transferCost.fixed + : getPeftTransferTime( + transferCost, producerCpu, consumerCpu, target)) + / target.processorPeriodNs; + info.schedulingCost = checkedAdd( + info.schedulingCost, checkedAdd(elements, transferCycles)); + info.work = checkedAdd(info.work, elements); + } + } + return info; +} + +static PipelineTaskModel getPipelineTaskModel( + const ComputeGraph &graph, const MergeScheduleResult &schedule, + const SchedulingTarget &target) { + PipelineTaskModel model; + model.schedulingCosts.reserve(graph.nodes.size()); + model.assemblyCosts.reserve(graph.nodes.size()); + model.predecessors.resize(graph.nodes.size()); + for (const ComputeGraphNode &node : graph.nodes) + if (node.instance.op) { + model.schedulingCosts.push_back(node.cost); + model.assemblyCosts.push_back(1); + } else { + model.schedulingCosts.push_back(node.cost); + model.assemblyCosts.push_back(node.cost); + } + + for (size_t task = 0; task < graph.nodes.size(); ++task) { + size_t cpu = schedule.computeToCpuMap.lookup(graph.nodes[task].instance); + BlueprintAssemblyInfo blueprint = getBlueprintAssemblyInfo( + graph, graph.nodes[task], cpu, schedule, target); + model.schedulingCosts[task] = checkedAdd( + model.schedulingCosts[task], blueprint.schedulingCost); + model.assemblyCosts[task] = checkedAdd( + model.assemblyCosts[task], blueprint.work); + model.predecessors[task].insert( + model.predecessors[task].end(), blueprint.predecessors.begin(), + blueprint.predecessors.end()); + for (const auto &[predecessor, transferCost] : graph.predecessors[task]) { + if (!llvm::is_contained(model.predecessors[task], predecessor)) + model.predecessors[task].push_back(predecessor); + size_t predecessorCpu = + schedule.computeToCpuMap.lookup(graph.nodes[predecessor].instance); + if (predecessorCpu == cpu) + continue; + model.schedulingCosts[task] = checkedAdd( + model.schedulingCosts[task], getPeftTransferTime( + transferCost, predecessorCpu, cpu, target) + / target.processorPeriodNs); + } + } + return model; +} + +static TaskCosts getPipelineTaskCosts(const ComputeGraph &graph, + const MergeScheduleResult &schedule, + const SchedulingTarget &target) { + return getPipelineTaskModel(graph, schedule, target).schedulingCosts; +} + +static void copyPipelineComputeAttrs(SpatGraphComputeBatch source, + SpatGraphComputeBatch destination) { + for (NamedAttribute attr : source->getAttrs()) { + StringRef name = attr.getName().getValue(); + if (name == source.getOperandSegmentSizesAttrName() + || name == source.getLaneCountAttrName()) + continue; + destination->setAttr(attr.getName(), attr.getValue()); + } + destination->setAttr("pipeline.split", UnitAttr::get(source.getContext())); +} + +static FailureOr createZeroTensor(OpBuilder &builder, Location loc, Type type) { + auto tensorType = dyn_cast(type); + if (!tensorType || !tensorType.hasStaticShape()) + return failure(); + auto zero = DenseElementsAttr::get( + tensorType, builder.getZeroAttr(tensorType.getElementType())); + return arith::ConstantOp::create(builder, loc, zero).getResult(); +} + +static bool isInside(Operation *operation, Operation *ancestor) { + for (Operation *current = operation; current; current = current->getParentOp()) + if (current == ancestor) + return true; + return false; +} + +static FailureOr findSplitBody(SpatGraphComputeBatch batch) { + if (batch.getNumResults() != 1 || batch.getBody().empty()) + return failure(); + + SplitBody result; + Block &block = batch.getBody().front(); + for (Operation &operation : block.without_terminator()) { + if (auto loop = dyn_cast(&operation)) { + SmallVector vmms; + loop.walk([&](SpatVMMOp vmm) { vmms.push_back(vmm); }); + if (!vmms.empty()) { + if (result.loop) + return failure(); + result.loop = loop; + result.vmms = std::move(vmms); + } + continue; + } + bool containsVmm = false; + operation.walk([&](SpatVMMOp) { containsVmm = true; }); + if (containsVmm) + return failure(); + } + if (!result.loop || result.vmms.size() < 2 || result.loop.getNumResults() != 1) + return failure(); + if (!isa(block.getTerminator())) + return failure(); + if (!isa(result.loop.getResult(0).getType())) + return failure(); + return result; +} + +static LogicalResult verifySplittableVmmUses(ArrayRef vmms, + std::string &error) { + for (SpatVMMOp vmm : vmms) + for (OpOperand &use : vmm.getResult().getUses()) + if (!isa(use.getOwner())) { + error = "pipeline task split requires VMM results to feed an additive reduction"; + return failure(); + } + return success(); +} + +static SmallVector collectAdditiveInputArguments( + SpatGraphComputeBatch batch, scf::ForOp loop) { + SmallVector result; + Block &block = batch.getBody().front(); + for (BlockArgument argument : block.getArguments()) { + if (argument == *batch.getLaneArgument()) + continue; + bool hasUse = false; + bool additive = true; + for (OpOperand &use : argument.getUses()) { + if (!isInside(use.getOwner(), loop)) { + additive = false; + break; + } + hasUse = true; + if (!isa(use.getOwner())) { + additive = false; + break; + } + } + if (hasUse && additive && isa(argument.getType())) + result.push_back(argument); + } + return result; +} + +static void simplifyZeroAdds(Operation *root, ArrayRef zeroValues) { + bool changed = true; + while (changed) { + changed = false; + SmallVector adds; + root->walk([&](SpatVAddOp add) { adds.push_back(add); }); + for (SpatVAddOp add : adds) { + Value replacement; + if (llvm::is_contained(zeroValues, add.getLhs())) + replacement = add.getRhs(); + else if (llvm::is_contained(zeroValues, add.getRhs())) + replacement = add.getLhs(); + if (!replacement) + continue; + add.getResult().replaceAllUsesWith(replacement); + add.erase(); + changed = true; + } + } +} + +static FailureOr cloneSplitLoop(OpBuilder &builder, + scf::ForOp sourceLoop, + IRMapping &mapper, + ArrayRef selectedVmmIndices, + ArrayRef sourceVmms) { + Operation *cloned = builder.clone(*sourceLoop.getOperation(), mapper); + SmallVector clonedVmms; + cloned->walk([&](SpatVMMOp vmm) { clonedVmms.push_back(vmm); }); + if (clonedVmms.size() != sourceVmms.size()) { + cloned->erase(); + return failure(); + } + + SmallVector zeroValues; + for (auto [index, vmm] : llvm::enumerate(clonedVmms)) { + if (llvm::is_contained(selectedVmmIndices, static_cast(index))) + continue; + OpBuilder zeroBuilder(vmm.getOperation()); + FailureOr zero = createZeroTensor(zeroBuilder, vmm.getLoc(), vmm.getResult().getType()); + if (failed(zero)) { + cloned->erase(); + return failure(); + } + zeroValues.push_back(*zero); + vmm.getResult().replaceAllUsesWith(*zero); + vmm.erase(); + } + simplifyZeroAdds(cloned, zeroValues); + return cloned; +} + +static SmallVector makeFragmentOffsets(OpBuilder &builder, + Value lane, + RankedTensorType fragmentType) { + SmallVector offsets {lane}; + offsets.append(fragmentType.getRank(), builder.getIndexAttr(0)); + return offsets; +} + +static SmallVector makeFragmentSizes(OpBuilder &builder, + RankedTensorType fragmentType) { + SmallVector sizes {builder.getIndexAttr(1)}; + for (int64_t dimension : fragmentType.getShape()) + sizes.push_back(builder.getIndexAttr(dimension)); + return sizes; +} + +static SmallVector makeFragmentStrides(OpBuilder &builder, + RankedTensorType fragmentType) { + return SmallVector( + fragmentType.getRank() + 1, builder.getIndexAttr(1)); +} + +static FailureOr createPartialBatch( + OpBuilder &builder, + SpatGraphComputeBatch source, + SplitBody &splitBody, + ArrayRef selectedVmmIndices, + ArrayRef additiveInputs) { + auto fragmentType = dyn_cast( + splitBody.loop.getOperation()->getResult(0).getType()); + if (!fragmentType || !fragmentType.hasStaticShape()) + return failure(); + RankedTensorType resultType = + getGraphBatchPhysicalResultType(source.getLaneCount(), fragmentType); + SmallVector partialInputs; + for (auto [index, input] : llvm::enumerate(source.getInputs())) + if (!llvm::is_contained(additiveInputs, *source.getInputArgument(index))) + partialInputs.push_back(input); + auto partial = createEmptySpatGraphComputeBatch( + builder, source.getLoc(), TypeRange {resultType}, source.getLaneCount(), + source.getWeights(), partialInputs); + if (failed(partial)) + return failure(); + copyPipelineComputeAttrs(source, *partial); + + IRMapping mapper; + mapper.map(*source.getLaneArgument(), *(*partial).getLaneArgument()); + for (size_t index = 0; index < source.getWeights().size(); ++index) + mapper.map(*source.getWeightArgument(index), *(*partial).getWeightArgument(index)); + unsigned partialInputIndex = 0; + for (size_t index = 0; index < source.getInputs().size(); ++index) { + BlockArgument sourceArgument = *source.getInputArgument(index); + if (llvm::is_contained(additiveInputs, sourceArgument)) + continue; + mapper.map(sourceArgument, *(*partial).getInputArgument(partialInputIndex++)); + } + for (size_t index = 0; index < source.getOutputs().size(); ++index) + mapper.map(*source.getOutputArgument(index), *(*partial).getOutputArgument(index)); + Block &sourceBlock = source.getBody().front(); + Block &partialBlock = (*partial).getBody().front(); + builder.setInsertionPointToEnd(&partialBlock); + for (Operation &operation : sourceBlock.without_terminator()) { + if (&operation == splitBody.loop) + break; + builder.clone(operation, mapper); + } + + for (BlockArgument input : additiveInputs) { + FailureOr zero = createZeroTensor( + builder, input.getLoc(), input.getType()); + if (failed(zero)) { + partial->erase(); + return failure(); + } + mapper.map(input, *zero); + } + + if (failed(cloneSplitLoop(builder, splitBody.loop, mapper, + selectedVmmIndices, splitBody.vmms))) { + partial->erase(); + return failure(); + } + auto partialLoop = cast( + &partialBlock.back()); + auto output = (*partial).getOutputArgument(0); + if (!output) { + partial->erase(); + return failure(); + } + + builder.setInsertionPointToEnd(&partialBlock); + auto inParallel = SpatInParallelOp::create(builder, source.getLoc()); + builder.setInsertionPointToEnd(&inParallel.getRegion().front()); + tensor::ParallelInsertSliceOp::create( + builder, source.getLoc(), partialLoop.getResult(0), *output, + makeFragmentOffsets(builder, *(*partial).getLaneArgument(), fragmentType), + makeFragmentSizes(builder, fragmentType), + makeFragmentStrides(builder, fragmentType)); + return *partial; +} + +static FailureOr extractBatchFragment(OpBuilder &builder, + Value physicalValue, + Value lane, + RankedTensorType fragmentType, + Location loc) { + SmallVector offsets = makeFragmentOffsets(builder, lane, fragmentType); + SmallVector sizes = makeFragmentSizes(builder, fragmentType); + SmallVector strides = makeFragmentStrides(builder, fragmentType); + return tensor::ExtractSliceOp::create( + builder, loc, fragmentType, physicalValue, offsets, sizes, strides).getResult(); +} + +static bool hasUnsupportedContinuationOperands(SpatGraphComputeBatch source, + scf::ForOp loop) { + Block &block = source.getBody().front(); + bool afterLoop = false; + for (Operation &operation : block.without_terminator()) { + if (&operation == loop) { + afterLoop = true; + continue; + } + if (!afterLoop) + continue; + for (Value operand : operation.getOperands()) { + if (auto argument = dyn_cast(operand)) + if (argument.getOwner() == &block) + return true; + Operation *definingOp = operand.getDefiningOp(); + if (definingOp && definingOp->isBeforeInBlock(loop)) + return true; + } + } + return false; +} + +static FailureOr createContinuationBatch( + OpBuilder &builder, + SpatGraphComputeBatch source, + SplitBody &splitBody, + ArrayRef partialResults) { + if (partialResults.size() < 2 + || hasUnsupportedContinuationOperands(source, splitBody.loop)) + return failure(); + + SmallVector inputs(partialResults.begin(), partialResults.end()); + auto continuation = createEmptySpatGraphComputeBatch( + builder, source.getLoc(), source.getResultTypes(), source.getLaneCount(), + ValueRange {}, inputs); + if (failed(continuation)) + return failure(); + copyPipelineComputeAttrs(source, *continuation); + + IRMapping mapper; + mapper.map(*source.getLaneArgument(), *(*continuation).getLaneArgument()); + for (size_t index = 0; index < source.getOutputs().size(); ++index) + mapper.map(*source.getOutputArgument(index), *(*continuation).getOutputArgument(index)); + Block &block = (*continuation).getBody().front(); + builder.setInsertionPointToEnd(&block); + auto fragmentType = dyn_cast( + splitBody.loop.getOperation()->getResult(0).getType()); + if (!fragmentType || !fragmentType.hasStaticShape()) { + continuation->erase(); + return failure(); + } + + SmallVector partialFragments; + for (auto [index, partial] : llvm::enumerate(partialResults)) { + auto input = (*continuation).getInputArgument(index); + if (!input) { + continuation->erase(); + return failure(); + } + FailureOr fragment = extractBatchFragment( + builder, *input, *(*continuation).getLaneArgument(), fragmentType, + source.getLoc()); + if (failed(fragment)) { + continuation->erase(); + return failure(); + } + partialFragments.push_back(*fragment); + } + + Value combined = partialFragments.front(); + for (Value fragment : llvm::drop_begin(partialFragments)) + combined = SpatVAddOp::create( + builder, source.getLoc(), fragmentType, combined, fragment).getResult(); + mapper.map(splitBody.loop.getOperation()->getResult(0), combined); + + bool afterLoop = false; + for (Operation &operation : source.getBody().front().without_terminator()) { + if (&operation == splitBody.loop) { + afterLoop = true; + continue; + } + if (afterLoop) + builder.clone(operation, mapper); + } + builder.clone(*source.getBody().front().getTerminator(), mapper); + return *continuation; +} + +static SmallVector, 4> +partitionVmmIndices(size_t vmmCount, size_t partCount) { + SmallVector, 4> partitions; + size_t next = 0; + for (size_t part = 0; part < partCount; ++part) { + size_t remainingParts = partCount - part; + size_t remainingVmms = vmmCount - next; + size_t count = (remainingVmms + remainingParts - 1) / remainingParts; + SmallVector selected; + for (size_t index = 0; index < count; ++index) + selected.push_back(static_cast(next + index)); + next += count; + partitions.push_back(std::move(selected)); + } + return partitions; +} + +static void eraseCreatedSplits(ArrayRef created) { + for (Operation *operation : llvm::reverse(created)) + operation->erase(); +} + +static LogicalResult splitBatchCompute(SpatGraphComputeBatch batch, + size_t pipelineStages, + const SchedulingTarget &physicalTarget, + std::string &error) { + if (batch->hasAttr("pipeline.split")) + return failure(); + FailureOr splitBody = findSplitBody(batch); + if (failed(splitBody)) + return failure(); + if (failed(verifySplittableVmmUses(splitBody->vmms, error))) + return failure(); + + size_t partCount = std::min(pipelineStages, splitBody->vmms.size()); + SmallVector, 4> partitions = + partitionVmmIndices(splitBody->vmms.size(), partCount); + SmallVector additiveInputs = + collectAdditiveInputArguments(batch, splitBody->loop); + SmallVector created; + SmallVector partialResults; + OpBuilder builder(batch.getOperation()); + for (auto [partIndex, partition] : llvm::enumerate(partitions)) { + builder.setInsertionPoint(batch.getOperation()); + ArrayRef zeroInputs = + partIndex == 0 ? ArrayRef() : additiveInputs; + FailureOr partial = createPartialBatch( + builder, batch, *splitBody, partition, zeroInputs); + if (failed(partial)) { + eraseCreatedSplits(created); + return failure(); + } + ComputeInstance instance { + partial->getOperation(), 0, + static_cast(partial->getLaneCount())}; + if (getComputeInstanceResidentWeights(instance, physicalTarget).size() + > physicalTarget.residentWeightCapacity) { + eraseCreatedSplits(created); + partial->erase(); + return failure(); + } + created.push_back(partial->getOperation()); + partialResults.push_back(partial->getResult(0)); + } + + builder.setInsertionPoint(batch.getOperation()); + FailureOr continuation = createContinuationBatch( + builder, batch, *splitBody, partialResults); + if (failed(continuation)) { + eraseCreatedSplits(created); + return failure(); + } + created.push_back(continuation->getOperation()); + batch->replaceAllUsesWith(continuation->getOperation()->getResults()); + batch->erase(); + return success(); +} + +static LogicalResult splitPipelineWorkloadImpl(const ComputeGraph &graph, + const MergeScheduleResult &schedule, + size_t pipelineStages, + const SchedulingTarget &physicalTarget, + std::string &error) { + size_t groupSize = schedule.processorCount; + std::vector tasksByCpu(groupSize); + for (size_t task = 0; task < graph.nodes.size(); ++task) { + auto cpu = schedule.computeToCpuMap.find(graph.nodes[task].instance); + if (cpu == schedule.computeToCpuMap.end() || cpu->second >= groupSize) { + error = "pipeline split received an incomplete PEFT schedule for task " + + std::to_string(task) + " (cpu " + + (cpu == schedule.computeToCpuMap.end() + ? std::string("missing") + : std::to_string(cpu->second)) + + ", schedule processors " + std::to_string(groupSize) + ")"; + return failure(); + } + tasksByCpu[cpu->second].push_back(task); + } + for (TaskList &tasks : tasksByCpu) + llvm::sort(tasks, [&](size_t lhs, size_t rhs) { + return schedule.computeToCpuSlotMap.lookup(graph.nodes[lhs].instance) + < schedule.computeToCpuSlotMap.lookup(graph.nodes[rhs].instance); + }); + TaskCosts taskCosts = getPipelineTaskCosts(graph, schedule, physicalTarget); + + for (const TaskList &tasks : tasksByCpu) { + if (tasks.empty()) + continue; + Cost maximumCost = findMaximumPackCost( + graph, taskCosts, tasks, physicalTarget.residentWeightCapacity, pipelineStages); + if (fits(graph, taskCosts, tasks, maximumCost, + physicalTarget.residentWeightCapacity, pipelineStages)) + continue; + + SmallVector candidates(tasks.begin(), tasks.end()); + llvm::stable_sort(candidates, [&](size_t lhs, size_t rhs) { + if (graph.nodes[lhs].residentWeights.size() + != graph.nodes[rhs].residentWeights.size()) + return graph.nodes[lhs].residentWeights.size() + > graph.nodes[rhs].residentWeights.size(); + return taskCosts[lhs] > taskCosts[rhs]; + }); + std::string candidateError; + for (size_t task : candidates) { + auto batch = dyn_cast( + graph.nodes[task].instance.op); + if (!batch || batch->hasAttr("pipeline.split")) + continue; + std::string currentError; + if (succeeded(splitBatchCompute( + batch, pipelineStages, physicalTarget, currentError))) + return success(); + if (!currentError.empty()) + candidateError = currentError; + } + error = candidateError.empty() + ? "pipeline scheduling could not find a splittable graph compute" + : candidateError; + return failure(); + } + return failure(); +} + +bool fits(const ComputeGraph& graph, + const TaskCosts& taskCosts, + const TaskList& tasks, + Cost maximumCost, + size_t residentWeightCapacity, + size_t maximumPacks) { + size_t packs = 1; + Cost cost = 0; + ResidentWeightSet weights; + bool packEmpty = true; + for (size_t task : tasks) { + const ComputeGraphNode& node = graph.nodes[task]; + Cost taskCost = taskCosts[task]; + if (node.residentWeights.size() > residentWeightCapacity) + return false; + bool startsNewPack = !packEmpty + && (cost > maximumCost - taskCost + || getResidentWeightUnionSize(weights, node.residentWeights) > residentWeightCapacity); + if (startsNewPack) { + if (++packs > maximumPacks) + return false; + cost = 0; + weights.clear(); + packEmpty = true; + } + cost = checkedAdd(cost, taskCost); + insertResidentWeights(weights, node.residentWeights); + packEmpty = false; + } + return true; +} + +Cost findMaximumPackCost(const ComputeGraph& graph, + const TaskCosts& taskCosts, + const TaskList& tasks, + size_t residentWeightCapacity, + size_t pipelineStages) { + Cost low = 0; + Cost high = 0; + for (size_t task : tasks) { + low = std::max(low, taskCosts[task]); + high = checkedAdd(high, taskCosts[task]); + } + while (low < high) { + Cost middle = low + (high - low) / 2; + if (fits(graph, taskCosts, tasks, middle, residentWeightCapacity, pipelineStages)) + high = middle; + else + low = middle + 1; + } + return low; +} + +static Cost getCoreCost(const TaskList &tasks, const TaskCosts &taskCosts) { + Cost cost = 0; + for (size_t task : tasks) + cost = checkedAdd(cost, taskCosts[task]); + return cost; +} + +static Cost getStageMaximumAssemblyCost( + const std::vector &tasksByCpu, + const TaskCosts &assemblyCosts, size_t groupSize, size_t stage) { + Cost maximum = 0; + for (size_t cpu = stage * groupSize; + cpu < (stage + 1) * groupSize; ++cpu) + maximum = std::max( + maximum, getCoreCost(tasksByCpu[cpu], assemblyCosts)); + return maximum; +} + +static bool fitsResidentWeights(const ComputeGraph &graph, + const TaskList &tasks, size_t candidate, + size_t residentWeightCapacity) { + ResidentWeightSet weights; + for (size_t task : tasks) + insertResidentWeights(weights, graph.nodes[task].residentWeights); + return getResidentWeightUnionSize( + weights, graph.nodes[candidate].residentWeights) + <= residentWeightCapacity; +} + +static void repackPipelineStage( + const ComputeGraph &graph, const TaskCosts &schedulingCosts, + const TaskCosts &assemblyCosts, + std::vector &tasksByCpu, size_t groupSize, size_t stage, + size_t residentWeightCapacity) { + TaskList tasks; + Cost originalMaximum = 0; + for (size_t cpu = stage * groupSize; + cpu < (stage + 1) * groupSize; ++cpu) { + llvm::append_range(tasks, tasksByCpu[cpu]); + originalMaximum = std::max( + originalMaximum, + getCoreCost(tasksByCpu[cpu], schedulingCosts)); + } + llvm::sort(tasks, [&](size_t lhs, size_t rhs) { + return graph.nodes[lhs].originalOrder < graph.nodes[rhs].originalOrder; + }); + + std::vector packed(groupSize); + std::vector weights(groupSize); + TaskCosts loads(groupSize); + TaskCosts assemblyLoads(groupSize); + for (size_t task : tasks) { + std::optional bestCore; + std::optional> bestScore; + for (size_t core = 0; core < groupSize; ++core) { + size_t unionSize = getResidentWeightUnionSize( + weights[core], graph.nodes[task].residentWeights); + if (unionSize > residentWeightCapacity) + continue; + size_t addedWeights = unionSize - weights[core].size(); + auto score = std::make_tuple( + checkedAdd(assemblyLoads[core], assemblyCosts[task]), + checkedAdd(loads[core], schedulingCosts[task]), addedWeights, core); + if (!bestScore || score < *bestScore) { + bestCore = core; + bestScore = score; + } + } + if (!bestCore) + return; + packed[*bestCore].push_back(task); + insertResidentWeights( + weights[*bestCore], graph.nodes[task].residentWeights); + loads[*bestCore] = checkedAdd(loads[*bestCore], schedulingCosts[task]); + assemblyLoads[*bestCore] = checkedAdd( + assemblyLoads[*bestCore], assemblyCosts[task]); + } + if (*std::max_element(loads.begin(), loads.end()) > originalMaximum) + return; + for (size_t core = 0; core < groupSize; ++core) + tasksByCpu[stage * groupSize + core] = std::move(packed[core]); +} + +static void rebalancePipelineStages( + const ComputeGraph &graph, const PipelineTaskModel &model, + std::vector &tasksByCpu, size_t groupSize, + size_t pipelineStages, size_t residentWeightCapacity) { + size_t minimumAssemblyFanIn = std::numeric_limits::max(); + for (size_t task = 0; task < graph.nodes.size(); ++task) + if (graph.nodes[task].instance.op && model.assemblyCosts[task] > 1 + && !model.predecessors[task].empty()) + minimumAssemblyFanIn = std::min( + minimumAssemblyFanIn, model.predecessors[task].size()); + bool hasAssembly = minimumAssemblyFanIn != std::numeric_limits::max(); + if (hasAssembly && groupSize < minimumAssemblyFanIn) + return; + const TaskCosts &balanceCosts = + hasAssembly ? model.assemblyCosts : model.schedulingCosts; + + Cost schedulingLimit = 0; + for (const TaskList &tasks : tasksByCpu) + schedulingLimit = std::max( + schedulingLimit, getCoreCost(tasks, model.schedulingCosts)); + for (size_t stage = 0; stage < pipelineStages; ++stage) + repackPipelineStage( + graph, model.schedulingCosts, balanceCosts, tasksByCpu, + groupSize, stage, + residentWeightCapacity); + + std::vector taskToCpu(graph.nodes.size()); + for (size_t cpu = 0; cpu < tasksByCpu.size(); ++cpu) + for (size_t task : tasksByCpu[cpu]) + taskToCpu[task] = cpu; + + bool changed; + do { + changed = false; + for (size_t sourceStage = pipelineStages; sourceStage-- > 1;) { + size_t targetStage = sourceStage - 1; + while (true) { + Cost sourceMaximum = getStageMaximumAssemblyCost( + tasksByCpu, balanceCosts, groupSize, sourceStage); + Cost targetMaximum = getStageMaximumAssemblyCost( + tasksByCpu, balanceCosts, groupSize, targetStage); + if (targetMaximum >= sourceMaximum) + break; + + struct Move { + size_t sourceCpu; + size_t targetCpu; + size_t task; + }; + std::optional best; + std::optional> bestScore; + for (size_t sourceCpu = sourceStage * groupSize; + sourceCpu < (sourceStage + 1) * groupSize; ++sourceCpu) { + if (tasksByCpu[sourceCpu].empty()) + continue; + size_t task = tasksByCpu[sourceCpu].front(); + bool dependenciesReady = llvm::all_of( + model.predecessors[task], [&](size_t predecessor) { + return taskToCpu[predecessor] / groupSize <= targetStage; + }); + if (!dependenciesReady) + continue; + + Cost sourceAfter = getCoreCost( + tasksByCpu[sourceCpu], balanceCosts) + - balanceCosts[task]; + for (size_t targetCpu = targetStage * groupSize; + targetCpu < (targetStage + 1) * groupSize; ++targetCpu) { + const TaskList &targetTasks = tasksByCpu[targetCpu]; + if (!fitsResidentWeights( + graph, targetTasks, task, residentWeightCapacity)) + continue; + Cost targetAfter = checkedAdd( + getCoreCost(targetTasks, balanceCosts), balanceCosts[task]); + Cost targetSchedulingAfter = checkedAdd( + getCoreCost(targetTasks, model.schedulingCosts), + model.schedulingCosts[task]); + if (targetAfter >= sourceMaximum + || targetSchedulingAfter > schedulingLimit) + continue; + auto score = std::make_tuple( + graph.nodes[task].originalOrder, + std::max(sourceAfter, targetAfter), targetAfter, targetCpu); + if (!bestScore || score < *bestScore) { + best = Move {sourceCpu, targetCpu, task}; + bestScore = score; + } + } + } + if (!best) + break; + tasksByCpu[best->sourceCpu].erase( + tasksByCpu[best->sourceCpu].begin()); + TaskList &targetTasks = tasksByCpu[best->targetCpu]; + auto insertion = llvm::find_if(targetTasks, [&](size_t task) { + return graph.nodes[task].originalOrder + > graph.nodes[best->task].originalOrder; + }); + targetTasks.insert(insertion, best->task); + taskToCpu[best->task] = best->targetCpu; + changed = true; + } + } + } while (changed); +} + +mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph, + MergeScheduleResult& schedule, + size_t pipelineStages, + const SchedulingTarget& physicalTarget, + std::string& error) { + const size_t groupSize = schedule.processorCount; + std::vector tasksByCpu(groupSize); + for (size_t task = 0; task < graph.nodes.size(); ++task) { + const ComputeInstance& instance = graph.nodes[task].instance; + auto cpu = schedule.computeToCpuMap.find(instance); + auto slot = schedule.computeToCpuSlotMap.find(instance); + if (cpu == schedule.computeToCpuMap.end() || slot == schedule.computeToCpuSlotMap.end() + || cpu->second >= groupSize) { + error = "pipeline scheduling received an incomplete PEFT schedule"; + return mlir::failure(); + } + tasksByCpu[cpu->second].push_back(task); + } + for (TaskList& tasks : tasksByCpu) + llvm::sort(tasks, [&](size_t lhs, size_t rhs) { + return schedule.computeToCpuSlotMap.lookup(graph.nodes[lhs].instance) + < schedule.computeToCpuSlotMap.lookup(graph.nodes[rhs].instance); + }); + PipelineTaskModel taskModel = getPipelineTaskModel( + graph, schedule, physicalTarget); + const TaskCosts &taskCosts = taskModel.schedulingCosts; + + std::vector tasksByPhysicalCpu(physicalTarget.processorCount); + for (size_t sourceCpu = 0; sourceCpu < groupSize; ++sourceCpu) { + const TaskList& tasks = tasksByCpu[sourceCpu]; + if (tasks.empty()) + continue; + for (size_t task : tasks) + if (graph.nodes[task].residentWeights.size() > physicalTarget.residentWeightCapacity) { + error = "pipeline scheduling cannot fit one compute instance in a physical core's crossbars"; + return mlir::failure(); + } + + Cost maximumCost = findMaximumPackCost( + graph, taskCosts, tasks, physicalTarget.residentWeightCapacity, pipelineStages); + if (!fits(graph, taskCosts, tasks, maximumCost, + physicalTarget.residentWeightCapacity, pipelineStages)) { + error = "pipeline scheduling cannot partition one PEFT core within the physical crossbar limit"; + return mlir::failure(); + } + + const size_t desiredPacks = std::min(pipelineStages, tasks.size()); + size_t stage = 0; + Cost packCost = 0; + ResidentWeightSet packWeights; + bool packEmpty = true; + for (size_t index = 0; index < tasks.size(); ++index) { + size_t task = tasks[index]; + const ComputeGraphNode& node = graph.nodes[task]; + Cost taskCost = taskCosts[task]; + bool exceedsLimit = + !packEmpty + && (packCost > maximumCost - taskCost + || getResidentWeightUnionSize(packWeights, node.residentWeights) > physicalTarget.residentWeightCapacity); + bool reserveOneTaskPerPack = !packEmpty && tasks.size() - index == desiredPacks - stage - 1; + if (exceedsLimit || reserveOneTaskPerPack) { + ++stage; + packCost = 0; + packWeights.clear(); + packEmpty = true; + } + if (stage >= pipelineStages) { + error = "pipeline scheduling produced too many packs"; + return mlir::failure(); + } + size_t physicalCpu = sourceCpu + stage * groupSize; + tasksByPhysicalCpu[physicalCpu].push_back(task); + packCost = checkedAdd(packCost, taskCost); + insertResidentWeights(packWeights, node.residentWeights); + packEmpty = false; + } + } + + rebalancePipelineStages( + graph, taskModel, tasksByPhysicalCpu, groupSize, pipelineStages, + physicalTarget.residentWeightCapacity); + + schedule.computeToCpuMap.clear(); + schedule.processorCount = physicalTarget.processorCount; + schedule.computeToCpuSlotMap.clear(); + schedule.computeToAestMap.clear(); + schedule.isLastComputeOfCpu.clear(); + schedule.cpuToLastComputeMap.clear(); + schedule.equivalentClass.clear(); + for (size_t cpu = 0; cpu < tasksByPhysicalCpu.size(); ++cpu) { + const TaskList& tasks = tasksByPhysicalCpu[cpu]; + for (auto [slot, task] : llvm::enumerate(tasks)) { + schedule.computeToCpuMap[graph.nodes[task].instance] = cpu; + schedule.computeToCpuSlotMap[graph.nodes[task].instance] = slot; + } + if (!tasks.empty()) { + const ComputeInstance& last = graph.nodes[tasks.back()].instance; + schedule.isLastComputeOfCpu.insert(last); + schedule.cpuToLastComputeMap[cpu] = last; + } + } + + std::vector indegree(graph.nodes.size()); + std::vector successors(graph.nodes.size()); + for (size_t task = 0; task < graph.nodes.size(); ++task) { + indegree[task] = graph.predecessors[task].size(); + for (const auto& [successor, cost] : graph.successors[task]) { + (void) cost; + successors[task].push_back(successor); + } + } + for (const TaskList& tasks : tasksByPhysicalCpu) + for (size_t index = 1; index < tasks.size(); ++index) { + successors[tasks[index - 1]].push_back(tasks[index]); + ++indegree[tasks[index]]; + } + + auto laterOriginalOrder = [&](size_t lhs, size_t rhs) { + return graph.nodes[lhs].originalOrder > graph.nodes[rhs].originalOrder; + }; + std::priority_queue, decltype(laterOriginalOrder)> ready(laterOriginalOrder); + for (size_t task = 0; task < graph.nodes.size(); ++task) + if (indegree[task] == 0) + ready.push(task); + + std::vector