diff --git a/src/PIM/Common/Support/DebugDump.cpp b/src/PIM/Common/Support/DebugDump.cpp index 4b7026e..632a3a9 100644 --- a/src/PIM/Common/Support/DebugDump.cpp +++ b/src/PIM/Common/Support/DebugDump.cpp @@ -17,7 +17,7 @@ std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::Str return std::fstream(dialectsDir + "/" + name + "." + extension.str(), std::ios::out); } -void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) { +void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified) { std::fstream file = openDialectDumpFileWithExtension(name, "/dialects", "mlir"); if (!file.is_open()) return; @@ -25,6 +25,8 @@ void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) { llvm::raw_os_ostream os(file); mlir::OpPrintingFlags flags; flags.elideLargeElementsAttrs().enableDebugInfo(true, false); + if (assumeVerified) + flags.assumeVerified(); moduleOp.print(os, flags); os.flush(); file.close(); diff --git a/src/PIM/Common/Support/DebugDump.hpp b/src/PIM/Common/Support/DebugDump.hpp index b0e8a99..820fa67 100644 --- a/src/PIM/Common/Support/DebugDump.hpp +++ b/src/PIM/Common/Support/DebugDump.hpp @@ -10,7 +10,7 @@ namespace onnx_mlir { /// Emits a MLIR snapshot under the current compiler output /// directory for pass-level debugging. -void dumpModule(mlir::ModuleOp moduleOp, const std::string& name); +void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified = false); /// Opens a file under the same dialect dump directory used by dumpModule. std::fstream openDialectDumpFileWithExtension(const std::string& name,llvm::StringRef destination = "/dialects", llvm::StringRef extension = "mlir"); diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp index 6ef425f..e411a9d 100644 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -18,7 +18,6 @@ #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/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Accelerators/PIM/Pass/PIMPasses.h" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -380,12 +379,6 @@ struct LowerSpatialPlansPass final : PassWrapperhasTrait(); } +bool isRecordedDeferredCommunicationSource(Operation* op, Value value) { + auto transfer = dyn_cast(op); + return transfer && llvm::is_contained(transfer.getSources(), value); +} + template void verifyComputeBodyCaptures(ComputeOpTy compute, StringRef kind, pim::CappedDiagnosticReporter& diagnostics) { Region& body = compute.getBody(); body.walk([&](Operation* nestedOp) { for (OpOperand& operand : nestedOp->getOpOperands()) { Value value = operand.get(); - if (isLegalExternalCapture(value, body)) + if (isLegalExternalCapture(value, body) || isRecordedDeferredCommunicationSource(nestedOp, value)) continue; Operation* definingOp = value.getDefiningOp(); @@ -90,21 +95,29 @@ bool isLegalHostBackedValue(Value value) { return definingOp->getDialect()->getNamespace() != "spat"; } +bool isScheduledPhase1Value(Value value) { + Operation* definingOp = value.getDefiningOp(); + return isa_and_nonnull(definingOp); +} + template void verifyScheduledInputs(ComputeOpTy compute, bool allowChannelReceiveInputs, StringRef kind, pim::CappedDiagnosticReporter& diagnostics) { for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) { + size_t currentInputIndex = inputIndex; Operation* definingOp = input.getDefiningOp(); if (allowChannelReceiveInputs && isa_and_nonnull(definingOp)) continue; + if (isScheduledPhase1Value(input)) + continue; if (isLegalHostBackedValue(input)) continue; diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) { InFlightDiagnostic diag = illegalOp->emitOpError() - << kPhaseMarker << " " << kind << " input #" << inputIndex + << kPhaseMarker << " " << kind << " input #" << currentInputIndex << (allowChannelReceiveInputs ? " must come from the host or explicit spat.channel_receive" : " must come from the host"); if (definingOp) @@ -163,9 +176,9 @@ 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 << " graph Spatial compute op remained after merge materialization"; + illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1"; }); } } diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index 0ca7852..60e229c 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -7,14 +7,9 @@ add_pim_library(SpatialOps SpatialOpsVerify.cpp SpatialOpsCanonicalization.cpp ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp - Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp - Transforms/MergeComputeNodes/CommunicationPlan.cpp - Transforms/MergeComputeNodes/HostOutputFinalization.cpp - Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp - Transforms/MergeComputeNodes/ProjectedFragments.cpp - Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp + Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp diff --git a/src/PIM/Dialect/Spatial/Spatial.td b/src/PIM/Dialect/Spatial/Spatial.td index 3d1bb03..340348f 100644 --- a/src/PIM/Dialect/Spatial/Spatial.td +++ b/src/PIM/Dialect/Spatial/Spatial.td @@ -6,6 +6,7 @@ include "mlir/IR/OpAsmInterface.td" include "mlir/IR/BuiltinTypes.td" include "mlir/IR/AttrTypeBase.td" include "mlir/IR/RegionKindInterface.td" +include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/ParallelCombiningOpInterface.td" include "mlir/Interfaces/SideEffectInterfaces.td" @@ -27,7 +28,7 @@ def SpatTensor : //===----------------------------------------------------------------------===// class SpatComputeLikeBase : SpatOp]> { let summary = "Compute region with attached constant weights"; @@ -40,7 +41,7 @@ class SpatComputeLikeBase : SpatOp:$outputs ); - let regions = (region SizedRegion<1>:$body); + let regions = (region MinSizedRegion<1>:$body); let hasVerifier = 1; let hasFolder = 1; @@ -76,7 +77,7 @@ def SpatScheduledCompute : SpatComputeLikeBase<"scheduled_compute"> { } class SpatComputeBatchLikeBase : SpatOp]> { let summary = "Tensor-native batch of equivalent compute lanes with shared weights and packed inputs"; @@ -90,7 +91,7 @@ class SpatComputeBatchLikeBase : SpatOp:$outputs ); - let regions = (region SizedRegion<1>:$body); + let regions = (region MinSizedRegion<1>:$body); let hasVerifier = 1; let hasCustomAssemblyFormat = 1; @@ -161,6 +162,54 @@ def SpatYieldOp : SpatOp<"yield", [Terminator]> { let hasCustomAssemblyFormat = 1; } +def SpatBlockYieldOp : SpatOp<"block_yield", [ + Terminator, + DeclareOpInterfaceMethods + ]> { + let summary = "Terminate a scheduled structural compute block"; + + let arguments = (ins + Variadic:$outputs + ); + + let successors = (successor + VariadicSuccessor:$next + ); + + let hasVerifier = 1; + let hasCustomAssemblyFormat = 1; +} + +def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]> { + let summary = "Temporary scheduled communication placeholder"; + + let arguments = (ins + Variadic:$sources, + StrAttr:$exchangeId, + StrAttr:$logicalProducer, + StrAttr:$logicalConsumer, + I64Attr:$sourceClass, + I64Attr:$targetClass, + I64Attr:$sourceCore, + I64Attr:$targetCore, + I64Attr:$sourceLane, + I64Attr:$targetLane, + StrAttr:$transferKind, + I64Attr:$resultIndex, + DenseI64ArrayAttr:$projectedTransfer, + I64Attr:$hostOutputOwner + ); + + let results = (outs + SpatTensor:$output + ); + + let regions = (region SizedRegion<1>:$body); + + let hasVerifier = 1; + let hasCustomAssemblyFormat = 1; +} + def SpatExtractRowsOp : SpatOp<"extract_rows", []> { let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors"; diff --git a/src/PIM/Dialect/Spatial/SpatialOps.cpp b/src/PIM/Dialect/Spatial/SpatialOps.cpp index d083848..f27d900 100644 --- a/src/PIM/Dialect/Spatial/SpatialOps.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOps.cpp @@ -238,6 +238,15 @@ void SpatScheduledCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValu setComputeAsmBlockArgumentNames(*this, region, setNameFn); } +SuccessorOperands SpatBlockYieldOp::getSuccessorOperands(unsigned index) { + assert(index == 0 && "invalid successor index"); + return SuccessorOperands(getOutputsMutable()); +} + +Block* SpatBlockYieldOp::getSuccessorForOperands(ArrayRef) { + return getOperation()->getNumSuccessors() == 0 ? nullptr : getOperation()->getSuccessor(0); +} + std::optional SpatGraphComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); } std::optional SpatGraphComputeBatch::getWeightArgument(unsigned idx) { return getBlockArgument(getBody(), 1 + idx); diff --git a/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp b/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp index 8d1c39f..b7ab8d3 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp @@ -50,6 +50,64 @@ static void printBlockArgumentList(OpAsmPrinter& printer, ArrayRef()) + continue; + printer.getStream() << "\n"; + printer.getStream().indent(indent); + printer.printCustomOrGenericOp(&nestedOp); + } + } + + printer.getStream() << "\n"; + printer << "}"; +} + static ParseResult parseBlockArgumentList(OpAsmParser& parser, SmallVectorImpl& arguments) { if (parser.parseLParen()) return failure(); @@ -160,7 +218,9 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) { printer << " -> "; printCompressedTypeSequence(printer, op.getResultTypes()); printer << " "; - printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false); + printRegionWithoutBlockArgLocs( + printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true, + /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true); } template @@ -290,7 +350,9 @@ void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) { printer << " -> "; printCompressedTypeSequence(printer, op.getResultTypes()); printer << " "; - printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false); + printRegionWithoutBlockArgLocs( + printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true, + /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true); } template @@ -407,6 +469,89 @@ ParseResult SpatYieldOp::parse(OpAsmParser& parser, OperationState& result) { return parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands); } +void SpatBlockYieldOp::print(OpAsmPrinter& printer) { + printer << " "; + printCompressedValueSequence(printer, getOutputs()); + if (getOperation()->getNumSuccessors() != 0) { + printer << " next "; + printer.printSuccessor(getOperation()->getSuccessor(0)); + } + printer.printOptionalAttrDict((*this)->getAttrs()); + printer << " : "; + printCompressedTypeSequence(printer, getOutputs().getTypes()); +} + +ParseResult SpatBlockYieldOp::parse(OpAsmParser& parser, OperationState& result) { + SmallVector outputs; + SmallVector outputTypes; + Block* successor = nullptr; + + OpAsmParser::UnresolvedOperand firstOutput; + OptionalParseResult firstOutputResult = parser.parseOptionalOperand(firstOutput); + if (firstOutputResult.has_value()) { + if (failed(*firstOutputResult)) + return failure(); + if (parseCompressedOperandEntryWithFirst(parser, firstOutput, outputs)) + return failure(); + while (succeeded(parser.parseOptionalComma())) + if (parseOneCompressedOperandEntry(parser, outputs)) + return failure(); + } + + if (succeeded(parser.parseOptionalKeyword("next")) && parser.parseSuccessor(successor)) + return failure(); + + if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() + || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true)) + return failure(); + + if (outputs.size() != outputTypes.size()) + return parser.emitError(parser.getCurrentLocation(), "number of outputs and output types must match"); + if (parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands)) + return failure(); + if (successor) + result.addSuccessors(successor); + return success(); +} + +void SpatDeferredCommunicationOp::print(OpAsmPrinter& printer) { + printer << " "; + printCompressedValueSequence(printer, getSources()); + printer.printOptionalAttrDict((*this)->getAttrs()); + printer << " : "; + printer.printFunctionalType(getSources().getTypes(), getOperation()->getResultTypes()); + printer << " "; + printRegionWithoutBlockArgLocs(printer, getBody(), /*printEntryBlockArgs=*/false); +} + +ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) { + SmallVector sources; + Type functionTypeStorage; + + if (parseCompressedOperandSequence(parser, sources) || parser.parseOptionalAttrDict(result.attributes) + || parser.parseColon() || parser.parseType(functionTypeStorage)) + return failure(); + + auto functionType = dyn_cast(functionTypeStorage); + if (!functionType) + return parser.emitError(parser.getCurrentLocation(), "expected deferred communication function type"); + if (sources.size() != functionType.getNumInputs()) + return parser.emitError(parser.getCurrentLocation(), "number of sources and source types must match"); + + if (parser.resolveOperands(sources, functionType.getInputs(), parser.getCurrentLocation(), result.operands)) + return failure(); + result.addTypes(functionType.getResults()); + + Region* body = result.addRegion(); + SmallVector bodyArgs; + for (Type type : functionType.getInputs()) { + OpAsmParser::Argument argument; + argument.type = type; + bodyArgs.push_back(argument); + } + return parser.parseRegion(*body, bodyArgs); +} + void SpatExtractRowsOp::print(OpAsmPrinter& printer) { printer << " "; printer.printOperand(getInput()); @@ -618,7 +763,8 @@ ParseResult SpatScheduledComputeBatch::parse(OpAsmParser& parser, OperationState void SpatInParallelOp::print(OpAsmPrinter& printer) { printer << " "; - printer.printRegion(getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false); + printRegionWithoutBlockArgLocs( + printer, getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false); printer.printOptionalAttrDict((*this)->getAttrs()); } diff --git a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp index ab9cc49..4ae9432 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp @@ -192,12 +192,18 @@ static bool isConstantExternalValue(Value value) { return definingOp && definingOp->hasTrait(); } +static bool isRecordedDeferredCommunicationSource(Operation* op, Value value) { + auto transfer = dyn_cast(op); + return transfer && llvm::is_contained(transfer.getSources(), value); +} + static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region& region, StringRef kind) { bool hasFailure = false; region.walk([&](Operation* op) { for (OpOperand& operand : op->getOpOperands()) { Value value = operand.get(); - if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value)) + if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value) + || isRecordedDeferredCommunicationSource(op, value)) continue; InFlightDiagnostic diagnostic = @@ -219,8 +225,35 @@ static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region return success(!hasFailure); } +static LogicalResult verifyYieldTypes(Operation* op, Region& region, TypeRange resultTypes, StringRef kind) { + if (region.empty()) + return op->emitOpError() << kind << " requires a body block"; + Block& block = region.front(); + auto yield = dyn_cast_or_null(block.getTerminator()); + if (!yield) + return op->emitOpError() << kind << " body must terminate with spat.yield"; + if (yield.getOutputs().size() != resultTypes.size()) + return op->emitOpError() << kind << " yield operand count must match result count"; + for (auto [yieldType, resultType] : llvm::zip(yield.getOutputs().getTypes(), resultTypes)) + if (yieldType != resultType) + return op->emitOpError() << kind << " yield operand types must match result types"; + return success(); +} + +static LogicalResult verifyRegionArguments(Operation* op, Region& region, ValueRange operands, StringRef kind) { + if (region.empty()) + return op->emitOpError() << kind << " requires a body block"; + Block& block = region.front(); + if (block.getNumArguments() != operands.size()) + return op->emitOpError() << kind << " body argument count must match operand count"; + for (auto [arg, operand] : llvm::zip(block.getArguments(), operands)) + if (arg.getType() != operand.getType()) + return op->emitOpError() << kind << " body argument types must match operand types"; + return success(); +} + template -static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block) { +static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block, bool verifyLaneSliceOffsets = true) { if (batchOp.getNumResults() == 0) { auto yieldOp = dyn_cast_or_null(block.getTerminator()); if (!yieldOp) @@ -235,11 +268,12 @@ static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block) { auto laneArg = batchOp.getLaneArgument(); if (!laneArg) return batchOp.emitError("compute_batch body must have a lane block argument"); - for (auto& bodyOp : block) { - if (auto extractSlice = dyn_cast(&bodyOp)) - if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice"))) - return failure(); - } + if (verifyLaneSliceOffsets) + for (auto& bodyOp : block) { + if (auto extractSlice = dyn_cast(&bodyOp)) + if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice"))) + return failure(); + } return success(); } @@ -679,7 +713,9 @@ LogicalResult verifyComputeResultsUses(Operation* op) { if (!isAnySpatialComputeLike(op)) return op->emitError("verifyComputeResultUses: op is not a Spatial compute-like operation"); if (!llvm::all_of(op->getResults(), [](Value result) { - return llvm::all_of(result.getUsers(), [](Operation* op) { + return llvm::all_of(result.getUsers(), [result](Operation* op) { + if (isRecordedDeferredCommunicationSource(op, result)) + return true; return !isAnySpatialComputeLike(op->getParentOp()); }); })) { @@ -690,57 +726,67 @@ LogicalResult verifyComputeResultsUses(Operation* op) { template LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) { - auto& block = compute.getBody().front(); unsigned expectedArgCount = compute.getWeights().size() + compute.getInputs().size(); - if (block.getNumArguments() != expectedArgCount) - return compute.emitOpError("compute body must have weight and input block arguments"); + bool isScheduled = isa(compute.getOperation()); + if (compute.getBody().empty()) + return compute.emitOpError("compute body must have at least one block"); - for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) { - auto blockArg = compute.getWeightArgument(weightIndex); - if (!blockArg || blockArg->getType() != weight.getType()) - return compute.emitOpError("compute weight block argument types must match weight operand types exactly"); - } - for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) { - auto blockArg = compute.getInputArgument(inputIndex); - if (!blockArg || blockArg->getType() != input.getType()) - return compute.emitOpError("compute input block argument types must match input operand types exactly"); - } + SmallVector yieldedTypes; + for (Block& block : compute.getBody()) { + if ((!isScheduled && block.getNumArguments() != expectedArgCount) + || (isScheduled && block.getNumArguments() < expectedArgCount)) + return compute.emitOpError("compute body must have weight and input block arguments"); - if (block.mightHaveTerminator()) { - auto yieldOp = dyn_cast_or_null(block.getTerminator()); - if (!yieldOp) + for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) + if (block.getArgument(weightIndex).getType() != weight.getType()) + return compute.emitOpError("compute weight block argument types must match weight operand types exactly"); + for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) + if (block.getArgument(compute.getWeights().size() + inputIndex).getType() != input.getType()) + return compute.emitOpError("compute input block argument types must match input operand types exactly"); + + Operation* terminator = block.getTerminator(); + if (auto yieldOp = dyn_cast_or_null(terminator)) { + if (isScheduled) + return compute.emitOpError("scheduled compute blocks must terminate with spat.block_yield"); + llvm::append_range(yieldedTypes, yieldOp->getOperandTypes()); + continue; + } + auto blockYield = dyn_cast_or_null(terminator); + if (!blockYield || !isScheduled) return compute.emitOpError("ComputeOp must have a single yield operation"); + if (blockYield->getNumSuccessors() == 0) + llvm::append_range(yieldedTypes, blockYield->getOperandTypes()); + } - auto resultTypes = compute.getResultTypes(); - auto yieldTypes = yieldOp->getOperandTypes(); - if (resultTypes.size() != yieldTypes.size()) - return compute.emitOpError("ComputeOp must have same number of results as yieldOp operands"); + auto resultTypes = compute.getResultTypes(); + if (resultTypes.size() != yieldedTypes.size()) + return compute.emitOpError("ComputeOp must have same number of results as yielded operands"); - for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldTypes))) { - auto resultType = std::get<0>(it); - auto yieldType = std::get<1>(it); + for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldedTypes))) { + auto resultType = std::get<0>(it); + auto yieldType = std::get<1>(it); - if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType))) - return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand"); + if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType))) + return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand"); - if (auto resultRankedType = dyn_cast(resultType)) { - if (auto yieldRankedType = dyn_cast(yieldType)) { - if (resultRankedType.getEncoding() != yieldRankedType.getEncoding()) - return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand"); - } - else { + if (auto resultRankedType = dyn_cast(resultType)) { + if (auto yieldRankedType = dyn_cast(yieldType)) { + if (resultRankedType.getEncoding() != yieldRankedType.getEncoding()) return compute.emitOpError("ComputeOp output has an encoding while yieldOp operand does not have one"); - } } - else if (dyn_cast(yieldType)) { - return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one"); + else { + return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand"); } } + else if (dyn_cast(yieldType)) { + return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one"); + } } - for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex) - if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty()) - return compute.emitOpError("ComputeOp block argument is not used"); + if (compute.getBody().hasOneBlock()) + for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex) + if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty()) + return compute.emitOpError("ComputeOp block argument is not used"); if (failed(verifyStaticWeights(compute, opName))) return failure(); if (failed(verifyOnlyConstantExternalValues(compute.getOperation(), compute.getBody(), opName))) @@ -754,6 +800,31 @@ LogicalResult SpatGraphCompute::verify() { return verifyComputeLikeOp(*this, "sp LogicalResult SpatScheduledCompute::verify() { return verifyComputeLikeOp(*this, "spat.scheduled_compute"); } +LogicalResult SpatBlockYieldOp::verify() { + if (getOperation()->getNumSuccessors() > 1) + return emitOpError("may target at most one next scheduled block"); + Operation* parent = getOperation()->getParentOp(); + if (!isa_and_nonnull(parent)) + return emitOpError("expected spat.scheduled_compute parent"); + if (getOperation()->getNumSuccessors() == 1) { + Block* next = getOperation()->getSuccessor(0); + if (getOperation()->getNumOperands() != next->getNumArguments()) + return emitOpError("successor operand count must match next block argument count"); + for (auto [operand, argument] : llvm::zip(getOperation()->getOperands(), next->getArguments())) + if (operand.getType() != argument.getType()) + return emitOpError("successor operand types must match next block argument types"); + } + return success(); +} + +LogicalResult SpatDeferredCommunicationOp::verify() { + if (getSources().empty()) + return emitOpError("requires at least one source"); + if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication"))) + return failure(); + return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication"); +} + template LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) { int32_t count = batch.getLaneCount(); @@ -808,7 +879,7 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) return failure(); if (failed(verifyOnlyConstantExternalValues(batch.getOperation(), batch.getBody(), opName))) return failure(); - return verifyBatchBody(batch, block); + return verifyBatchBody(batch, block, !isa(batch.getOperation())); } LogicalResult SpatGraphComputeBatch::verify() { return verifyComputeBatchLikeOp(*this, "spat.graph_compute_batch"); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp deleted file mode 100644 index 9b46571..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp +++ /dev/null @@ -1,656 +0,0 @@ -#include "CommunicationPlan.hpp" - -#include "mlir/IR/Diagnostics.h" - -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/STLExtras.h" - -#include - -using namespace mlir; - -namespace onnx_mlir::spatial { -namespace { - -StringRef stringifyExchangeKind(CommunicationExchangeKind kind) { - switch (kind) { - case CommunicationExchangeKind::HostForward: return "HostForward"; - case CommunicationExchangeKind::OrdinaryValue: return "OrdinaryValue"; - case CommunicationExchangeKind::PackedScalarRun: return "PackedScalarRun"; - case CommunicationExchangeKind::IndexedBatchRun: return "IndexedBatchRun"; - case CommunicationExchangeKind::ProjectedInput: return "ProjectedInput"; - } - llvm_unreachable("unknown communication exchange kind"); -} - -unsigned getExchangeKindOrder(CommunicationExchangeKind kind) { - switch (kind) { - case CommunicationExchangeKind::HostForward: return 0; - case CommunicationExchangeKind::OrdinaryValue: return 1; - case CommunicationExchangeKind::PackedScalarRun: return 2; - case CommunicationExchangeKind::IndexedBatchRun: return 3; - case CommunicationExchangeKind::ProjectedInput: return 4; - } - llvm_unreachable("unknown communication exchange kind"); -} - -void emitPlanError(Operation* anchor, StringRef message) { - if (anchor) - anchor->emitError(message); -} - -struct PeerPair { - int64_t minCore = 0; - int64_t maxCore = 0; - - bool operator==(const PeerPair& other) const { - return minCore == other.minCore && maxCore == other.maxCore; - } -}; - -struct PeerPairInfo { - static PeerPair getEmptyKey() { - return {llvm::DenseMapInfo::getEmptyKey(), llvm::DenseMapInfo::getEmptyKey()}; - } - static PeerPair getTombstoneKey() { - return {llvm::DenseMapInfo::getTombstoneKey(), llvm::DenseMapInfo::getTombstoneKey()}; - } - static unsigned getHashValue(PeerPair key) { return llvm::hash_combine(key.minCore, key.maxCore); } - static bool isEqual(PeerPair lhs, PeerPair rhs) { return lhs == rhs; } -}; - -PeerPair getPeerPair(const CommunicationExchange& exchange) { - return {std::min(exchange.sourceCore, exchange.targetCore), std::max(exchange.sourceCore, exchange.targetCore)}; -} - -bool exchangeOrderLess(const CommunicationExchange& lhs, const CommunicationExchange& rhs) { - if (getExchangeKindOrder(lhs.kind) != getExchangeKindOrder(rhs.kind)) - return getExchangeKindOrder(lhs.kind) < getExchangeKindOrder(rhs.kind); - if (lhs.sourceClass != rhs.sourceClass) - return lhs.sourceClass < rhs.sourceClass; - if (lhs.targetClass != rhs.targetClass) - return lhs.targetClass < rhs.targetClass; - if (lhs.producerOrder != rhs.producerOrder) - return lhs.producerOrder < rhs.producerOrder; - if (lhs.producer.resultIndex != rhs.producer.resultIndex) - return lhs.producer.resultIndex < rhs.producer.resultIndex; - if (lhs.producer.instance.laneStart != rhs.producer.instance.laneStart) - return lhs.producer.instance.laneStart < rhs.producer.instance.laneStart; - if (lhs.producer.instance.laneCount != rhs.producer.instance.laneCount) - return lhs.producer.instance.laneCount < rhs.producer.instance.laneCount; - if (lhs.sourceLane != rhs.sourceLane) - return lhs.sourceLane < rhs.sourceLane; - if (lhs.targetLane != rhs.targetLane) - return lhs.targetLane < rhs.targetLane; - if (lhs.ordinal != rhs.ordinal) - return lhs.ordinal < rhs.ordinal; - return lhs.channelId < rhs.channelId; -} - -void sortExchangeIds(ArrayRef exchanges, - SmallVectorImpl& ids) { - llvm::stable_sort(ids, [&](CommunicationExchangeId lhs, CommunicationExchangeId rhs) { - return exchangeOrderLess(exchanges[lhs.value], exchanges[rhs.value]); - }); -} - -void addEvent(CommunicationPlan& plan, - DenseMap>& eventsByClass, - CommunicationExchangeId id, - CommunicationEventKind kind, - ClassId classId, - int64_t core, - unsigned& order) { - eventsByClass[classId].push_back(CommunicationEvent {id, kind, classId, core, order++}); - (void)plan; -} - -void buildEventStreams(CommunicationPlan& plan, - SmallVectorImpl& exchanges, - DenseMap>& eventsByClass) { - DenseMap, PeerPairInfo> exchangesByPair; - for (const CommunicationExchange& exchange : exchanges) - exchangesByPair[getPeerPair(exchange)].push_back(exchange.id); - - SmallVector pairs; - pairs.reserve(exchangesByPair.size()); - for (const auto& entry : exchangesByPair) - pairs.push_back(entry.first); - llvm::sort(pairs, [](PeerPair lhs, PeerPair rhs) { - if (lhs.minCore != rhs.minCore) - return lhs.minCore < rhs.minCore; - return lhs.maxCore < rhs.maxCore; - }); - - unsigned order = 0; - for (PeerPair pair : pairs) { - SmallVector hiToLo; - SmallVector loToHi; - for (CommunicationExchangeId id : exchangesByPair.lookup(pair)) { - const CommunicationExchange& exchange = exchanges[id.value]; - if (exchange.sourceCore == pair.maxCore && exchange.targetCore == pair.minCore) - hiToLo.push_back(id); - else - loToHi.push_back(id); - } - - sortExchangeIds(exchanges, hiToLo); - sortExchangeIds(exchanges, loToHi); - - for (CommunicationExchangeId id : hiToLo) { - const CommunicationExchange& exchange = exchanges[id.value]; - addEvent(plan, eventsByClass, id, CommunicationEventKind::Receive, exchange.targetClass, exchange.targetCore, order); - } - for (CommunicationExchangeId id : hiToLo) { - const CommunicationExchange& exchange = exchanges[id.value]; - addEvent(plan, eventsByClass, id, CommunicationEventKind::Send, exchange.sourceClass, exchange.sourceCore, order); - } - for (CommunicationExchangeId id : loToHi) { - const CommunicationExchange& exchange = exchanges[id.value]; - addEvent(plan, eventsByClass, id, CommunicationEventKind::Send, exchange.sourceClass, exchange.sourceCore, order); - } - for (CommunicationExchangeId id : loToHi) { - const CommunicationExchange& exchange = exchanges[id.value]; - addEvent(plan, eventsByClass, id, CommunicationEventKind::Receive, exchange.targetClass, exchange.targetCore, order); - } - } -} - -} // namespace - -FailureOr -CommunicationPlan::build(Operation* anchor, - ArrayRef classes, - const DenseMap& cpuToClass, - const ProducerSourceClassMap& producerSourceClasses, - const DenseMap& producerPayloadTypes, - const DenseMap, ProducerKeyInfo>& producerDestClasses, - const DenseMap, ProducerKeyInfo>& ordinaryProducerDestClasses, - ProjectedInputPlanMap& projectedPlans, - const DenseMap& hostOutputOwners, - const DenseMap& liveExternalUseCache, - int64_t& nextChannelId) { - CommunicationPlan plan; - - DenseMap classById; - for (CommunicationClassInfo klass : classes) - classById[klass.classId] = std::move(klass); - - auto appendExchange = [&](CommunicationExchangeKind kind, - ProducerKey producer, - ClassId sourceClass, - ClassId targetClass, - Type payloadType, - int64_t sourceCore, - int64_t targetCore, - unsigned sourceLane, - unsigned targetLane, - unsigned ordinal, - Value originalHostOutput = Value(), - const ProjectedInputTransferFragment* fragment = nullptr) -> FailureOr { - if (sourceCore == targetCore) - return failure(); - - auto sourceInfoIt = classById.find(sourceClass); - auto targetInfoIt = classById.find(targetClass); - if (sourceInfoIt == classById.end() || targetInfoIt == classById.end()) - return failure(); - - CommunicationExchangeId exchangeId {plan.exchanges.size()}; - plan.exchanges.push_back(CommunicationExchange { - exchangeId, - kind, - PlaceholderId {static_cast(plan.demands.size())}, - producer, - sourceClass, - targetClass, - sourceInfoIt->second.rank, - targetInfoIt->second.rank, - payloadType, - nextChannelId++, - sourceCore, - targetCore, - sourceLane, - targetLane, - ordinal, - exchangeId.value, - originalHostOutput, - fragment, - }); - return exchangeId; - }; - - auto getPayloadType = [&](ProducerKey producer) -> Type { - auto typeIt = producerPayloadTypes.find(producer); - if (typeIt != producerPayloadTypes.end()) - return typeIt->second; - if (!producer.instance.op || producer.resultIndex >= producer.instance.op->getNumResults()) - return Type(); - return producer.instance.op->getResult(producer.resultIndex).getType(); - }; - - for (const auto& entry : producerDestClasses) { - ProducerKey producer = entry.first; - auto sourceClassIt = producerSourceClasses.find(producer); - if (sourceClassIt == producerSourceClasses.end()) { - emitPlanError(anchor, "communication plan could not find producer source class"); - return failure(); - } - ClassId sourceClass = sourceClassIt->second; - auto sourceInfoIt = classById.find(sourceClass); - if (sourceInfoIt == classById.end()) { - emitPlanError(anchor, "communication plan references an unknown producer source class"); - return failure(); - } - const CommunicationClassInfo& sourceInfo = sourceInfoIt->second; - Type payloadType = getPayloadType(producer); - if (!payloadType) { - emitPlanError(anchor, "communication plan could not determine producer payload type"); - return failure(); - } - - for (ClassId targetClass : entry.second) { - if (targetClass == sourceClass) - continue; - - auto targetInfoIt = classById.find(targetClass); - if (targetInfoIt == classById.end()) { - emitPlanError(anchor, "communication plan references an unknown target class"); - return failure(); - } - const CommunicationClassInfo& targetInfo = targetInfoIt->second; - - CommunicationExchangeKind kind = CommunicationExchangeKind::OrdinaryValue; - if (sourceInfo.isBatch && !targetInfo.isBatch) - kind = CommunicationExchangeKind::PackedScalarRun; - else if (sourceInfo.isBatch && targetInfo.isBatch) - kind = CommunicationExchangeKind::IndexedBatchRun; - else if (!llvm::is_contained(ordinaryProducerDestClasses.lookup(producer), targetClass)) - continue; - - SmallVector* bucket = nullptr; - ProducerTargetKey lookupKey {producer, targetClass}; - if (kind == CommunicationExchangeKind::OrdinaryValue) - bucket = &plan.ordinaryExchangesByProducerTarget[lookupKey]; - else if (kind == CommunicationExchangeKind::PackedScalarRun) - bucket = &plan.packedScalarRunExchangesByProducerTarget[lookupKey]; - else if (kind == CommunicationExchangeKind::IndexedBatchRun) - bucket = &plan.indexedBatchRunExchangesByProducerTarget[lookupKey]; - - if (!bucket) - continue; - - if (!sourceInfo.isBatch) { - if (sourceInfo.cpus.empty()) { - emitPlanError(anchor, "communication plan source class has no CPU"); - return failure(); - } - int64_t sourceCore = static_cast(sourceInfo.cpus.front()); - unsigned ordinal = 0; - for (auto [targetLane, targetCpu] : llvm::enumerate(targetInfo.cpus)) { - auto exchangeId = appendExchange(kind, - producer, - sourceClass, - targetClass, - payloadType, - sourceCore, - static_cast(targetCpu), - 0, - static_cast(targetLane), - ordinal++); - if (failed(exchangeId)) { - emitPlanError(anchor, "communication plan failed to append scalar source exchange"); - return failure(); - } - bucket->push_back(*exchangeId); - if (!targetInfo.isBatch) - break; - } - continue; - } - - if (sourceInfo.cpus.empty()) { - emitPlanError(anchor, "communication plan batch source class has no CPUs"); - return failure(); - } - if (targetInfo.isBatch && sourceInfo.cpus.size() != targetInfo.cpus.size()) { - emitPlanError(anchor, "communication plan cannot map batch classes of different sizes"); - return failure(); - } - - for (auto [sourceLane, sourceCpu] : llvm::enumerate(sourceInfo.cpus)) { - CpuId targetCpu = targetInfo.isBatch ? targetInfo.cpus[sourceLane] : targetInfo.cpus.front(); - auto exchangeId = appendExchange(kind, - producer, - sourceClass, - targetClass, - payloadType, - static_cast(sourceCpu), - static_cast(targetCpu), - static_cast(sourceLane), - targetInfo.isBatch ? static_cast(sourceLane) : 0, - static_cast(bucket->size())); - if (failed(exchangeId)) { - emitPlanError(anchor, "communication plan failed to append batch source exchange"); - return failure(); - } - bucket->push_back(*exchangeId); - } - } - } - - for (const auto& entry : producerSourceClasses) { - ProducerKey producer = entry.first; - ClassId sourceClass = entry.second; - if (!producer.instance.op || producer.resultIndex >= producer.instance.op->getNumResults()) - continue; - - Value originalOutput = producer.instance.op->getResult(producer.resultIndex); - auto liveIt = liveExternalUseCache.find(originalOutput); - if (liveIt == liveExternalUseCache.end() || !liveIt->second) - continue; - - auto ownerIt = hostOutputOwners.find(originalOutput); - if (ownerIt == hostOutputOwners.end()) - continue; - ClassId ownerClass = ownerIt->second; - if (ownerClass == sourceClass) - continue; - - auto sourceInfoIt = classById.find(sourceClass); - auto ownerInfoIt = classById.find(ownerClass); - if (sourceInfoIt == classById.end() || ownerInfoIt == classById.end()) { - emitPlanError(anchor, "communication plan references an unknown host-forward class"); - return failure(); - } - const CommunicationClassInfo& sourceInfo = sourceInfoIt->second; - const CommunicationClassInfo& ownerInfo = ownerInfoIt->second; - if (sourceInfo.isBatch != ownerInfo.isBatch) { - emitPlanError(anchor, "communication plan does not support mixed scalar/batch host forwarding"); - return failure(); - } - if (sourceInfo.isBatch && sourceInfo.cpus.size() != ownerInfo.cpus.size()) { - emitPlanError(anchor, "communication plan cannot host-forward between batch classes of different sizes"); - return failure(); - } - - Type payloadType = getPayloadType(producer); - if (!payloadType) { - emitPlanError(anchor, "communication plan could not determine host-forward payload type"); - return failure(); - } - - HostForwardKey lookupKey {originalOutput, sourceClass, ownerClass}; - SmallVector& bucket = plan.hostForwardExchangesByOutput[lookupKey]; - if (!sourceInfo.isBatch) { - if (!bucket.empty()) - continue; - auto exchangeId = appendExchange(CommunicationExchangeKind::HostForward, - producer, - sourceClass, - ownerClass, - payloadType, - static_cast(sourceInfo.cpus.front()), - static_cast(ownerInfo.cpus.front()), - 0, - 0, - 0, - originalOutput); - if (failed(exchangeId)) { - emitPlanError(anchor, "communication plan failed to append scalar host-forward exchange"); - return failure(); - } - bucket.push_back(*exchangeId); - continue; - } - - if (bucket.size() >= sourceInfo.cpus.size()) - continue; - unsigned ordinal = static_cast(bucket.size()); - if (producer.instance.laneStart >= sourceInfo.cpus.size()) { - emitPlanError(anchor, "communication plan host-forward producer lane is out of range"); - return failure(); - } - unsigned lane = producer.instance.laneStart; - auto exchangeId = appendExchange(CommunicationExchangeKind::HostForward, - producer, - sourceClass, - ownerClass, - payloadType, - static_cast(sourceInfo.cpus[lane]), - static_cast(ownerInfo.cpus[lane]), - lane, - lane, - ordinal, - originalOutput); - if (failed(exchangeId)) { - emitPlanError(anchor, "communication plan failed to append batch host-forward exchange"); - return failure(); - } - bucket.push_back(*exchangeId); - } - - DenseSet seenFragments; - for (auto& extractEntry : projectedPlans) { - for (auto& classEntry : extractEntry.second) { - ClassId targetClass = classEntry.first; - ProjectedInputTransferPlan& inputPlan = classEntry.second; - - auto targetInfoIt = classById.find(targetClass); - if (targetInfoIt == classById.end()) { - emitPlanError(anchor, "communication plan references an unknown target class"); - return failure(); - } - const CommunicationClassInfo& targetInfo = targetInfoIt->second; - - for (ProjectedInputTransferFragment& fragment : inputPlan.fragments) { - if (!seenFragments.insert(&fragment).second) { - emitPlanError(anchor, "communication plan found the same projected fragment twice"); - return failure(); - } - - auto sourceClassIt = cpuToClass.find(fragment.sourceCoreId); - if (sourceClassIt == cpuToClass.end()) { - emitPlanError(anchor, "communication plan could not map projected source core to a class"); - return failure(); - } - - ClassId sourceClass = sourceClassIt->second; - auto sourceInfoIt = classById.find(sourceClass); - if (sourceInfoIt == classById.end()) { - emitPlanError(anchor, "communication plan references an unknown source class"); - return failure(); - } - const CommunicationClassInfo& sourceInfo = sourceInfoIt->second; - - plan.demands.push_back(InputDemand {{static_cast(plan.demands.size())}, - fragment.producer, - targetClass, - inputPlan.layout.fragmentType}); - plan.publications.push_back(ProducerPublication {fragment.producer, sourceClass, inputPlan.layout.fragmentType}); - - if (fragment.sourceCoreId == fragment.targetCoreId) - continue; - - if (sourceInfo.rank == targetInfo.rank && sourceClass != targetClass) { - emitPlanError(anchor, "communication plan found a same-rank inter-class projected exchange"); - return failure(); - } - - FailureOr exchangeId = appendExchange( - CommunicationExchangeKind::ProjectedInput, - fragment.producer, - sourceClass, - targetClass, - inputPlan.layout.fragmentType, - fragment.sourceCoreId, - fragment.targetCoreId, - 0, - fragment.targetLane, - fragment.ordinal, - Value(), - &fragment); - if (failed(exchangeId)) { - emitPlanError(anchor, "communication plan failed to append projected input exchange"); - return failure(); - } - fragment.channelId = plan.exchanges[exchangeId->value].channelId; - plan.projectedExchangesByFragment[&fragment].push_back(*exchangeId); - } - } - } - - buildEventStreams(plan, plan.exchanges, plan.eventsByClass); - if (failed(plan.verify(anchor))) - return failure(); - return plan; -} - -LogicalResult CommunicationPlan::verify(Operation* anchor) const { - DenseSet channelIds; - SmallVector malformed; - size_t malformedCount = 0; - - auto recordMalformed = [&](const CommunicationExchange& exchange) { - ++malformedCount; - if (malformed.size() < 8) - malformed.push_back(&exchange); - }; - - for (auto [index, exchange] : llvm::enumerate(exchanges)) { - bool bad = false; - if (exchange.id.value != index) - bad = true; - if (exchange.sourceCore == exchange.targetCore) - bad = true; - if (!exchange.payloadType) - bad = true; - if (!channelIds.insert(exchange.channelId).second) - bad = true; - if (bad) - recordMalformed(exchange); - } - - DenseMap sendCounts; - DenseMap receiveCounts; - for (const auto& classEvents : eventsByClass) { - std::optional previousOrder; - DenseSet classOrders; - for (const CommunicationEvent& event : classEvents.second) { - if (event.exchangeId.value >= exchanges.size()) { - if (anchor) - anchor->emitError("communication plan event references an invalid exchange id"); - return failure(); - } - - const CommunicationExchange& exchange = exchanges[event.exchangeId.value]; - if (previousOrder && event.order < *previousOrder) { - if (anchor) - anchor->emitError("communication plan event stream is not sorted"); - return failure(); - } - previousOrder = event.order; - if (!classOrders.insert(event.order).second) { - if (anchor) - anchor->emitError("communication plan event stream contains duplicate event order"); - return failure(); - } - - if (event.kind == CommunicationEventKind::Send) { - ++sendCounts[event.exchangeId.value]; - if (event.classId != exchange.sourceClass || event.core != exchange.sourceCore) - recordMalformed(exchange); - } - else { - ++receiveCounts[event.exchangeId.value]; - if (event.classId != exchange.targetClass || event.core != exchange.targetCore) - recordMalformed(exchange); - } - } - } - - for (const CommunicationExchange& exchange : exchanges) { - if (sendCounts[exchange.id.value] != 1 || receiveCounts[exchange.id.value] != 1) - recordMalformed(exchange); - } - - auto verifyLookup = [&](ArrayRef ids) -> LogicalResult { - for (CommunicationExchangeId id : ids) - if (id.value >= exchanges.size()) - return failure(); - return success(); - }; - - for (const auto& entry : projectedExchangesByFragment) - if (failed(verifyLookup(entry.second))) { - if (anchor) - anchor->emitError("communication plan projected lookup references an invalid exchange id"); - return failure(); - } - - if (malformedCount != 0) { - if (anchor) { - InFlightDiagnostic diag = anchor->emitError("communication plan contains malformed exchanges"); - for (const CommunicationExchange* exchange : malformed) { - diag << " [exchange=" << exchange->id.value << " kind=" << stringifyExchangeKind(exchange->kind) - << " sourceClass=" << exchange->sourceClass << " sourceCore=" << exchange->sourceCore - << " targetClass=" << exchange->targetClass << " targetCore=" << exchange->targetCore - << " channelId=" << exchange->channelId << " payloadType=" << exchange->payloadType << "]"; - } - if (malformedCount > malformed.size()) - diag << " additionalMalformed=" << (malformedCount - malformed.size()); - } - return failure(); - } - - return success(); -} - -ArrayRef CommunicationPlan::getEventsForClass(ClassId classId) const { - auto it = eventsByClass.find(classId); - if (it == eventsByClass.end()) - return {}; - return it->second; -} - -ArrayRef -CommunicationPlan::getOrdinaryValueExchanges(ProducerKey producer, ClassId targetClass) const { - auto it = ordinaryExchangesByProducerTarget.find({producer, targetClass}); - if (it == ordinaryExchangesByProducerTarget.end()) - return {}; - return it->second; -} - -ArrayRef -CommunicationPlan::getProjectedInputExchanges(const ProjectedInputTransferFragment& fragment) const { - auto it = projectedExchangesByFragment.find(&fragment); - if (it == projectedExchangesByFragment.end()) - return {}; - return it->second; -} - -ArrayRef -CommunicationPlan::getPackedScalarRunExchanges(ProducerKey producer, ClassId targetClass) const { - auto it = packedScalarRunExchangesByProducerTarget.find({producer, targetClass}); - if (it == packedScalarRunExchangesByProducerTarget.end()) - return {}; - return it->second; -} - -ArrayRef -CommunicationPlan::getIndexedBatchRunExchanges(ProducerKey producer, ClassId targetClass) const { - auto it = indexedBatchRunExchangesByProducerTarget.find({producer, targetClass}); - if (it == indexedBatchRunExchangesByProducerTarget.end()) - return {}; - return it->second; -} - -ArrayRef -CommunicationPlan::getHostForwardExchanges(Value originalOutput, ClassId sourceClass, ClassId ownerClass) const { - auto it = hostForwardExchangesByOutput.find({originalOutput, sourceClass, ownerClass}); - if (it == hostForwardExchangesByOutput.end()) - return {}; - return it->second; -} - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp deleted file mode 100644 index d109545..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp +++ /dev/null @@ -1,229 +0,0 @@ -#pragma once - -#include "mlir/IR/Operation.h" -#include "mlir/IR/Types.h" -#include "mlir/Support/LLVM.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/SmallVector.h" - -#include -#include -#include - -#include "MergeScheduleKeys.hpp" -#include "ProjectedFragments.hpp" - -namespace onnx_mlir::spatial { - -enum class TransferKind { - DirectValue, - LocalFragment, - SameClassForward, - RemoteChannel, - WholeTensorBarrier -}; - -enum class CommunicationExchangeKind { - OrdinaryValue, - ProjectedInput, - PackedScalarRun, - IndexedBatchRun, - HostForward -}; - -enum class CommunicationEventKind { - Send, - Receive -}; - -struct CommunicationExchangeId { - size_t value = 0; - - bool operator==(const CommunicationExchangeId& other) const { return value == other.value; } - operator size_t() const { return value; } -}; - -struct CommunicationExchangeIdInfo { - static CommunicationExchangeId getEmptyKey() { return {llvm::DenseMapInfo::getEmptyKey()}; } - static CommunicationExchangeId getTombstoneKey() { return {llvm::DenseMapInfo::getTombstoneKey()}; } - static unsigned getHashValue(CommunicationExchangeId id) { return llvm::hash_value(id.value); } - static bool isEqual(CommunicationExchangeId lhs, CommunicationExchangeId rhs) { return lhs == rhs; } -}; - -struct PlaceholderId { - unsigned value = 0; - - bool operator==(const PlaceholderId& other) const { return value == other.value; } -}; - -struct PlaceholderIdInfo { - static PlaceholderId getEmptyKey() { return {static_cast(-1)}; } - static PlaceholderId getTombstoneKey() { return {static_cast(-2)}; } - static unsigned getHashValue(PlaceholderId id) { return llvm::hash_value(id.value); } - static bool isEqual(PlaceholderId lhs, PlaceholderId rhs) { return lhs == rhs; } -}; - -struct InputDemand { - PlaceholderId placeholder; - ProducerKey producer; - ClassId targetClass = 0; - mlir::Type requiredType; -}; - -struct ProducerPublication { - ProducerKey producer; - ClassId sourceClass = 0; - mlir::Type payloadType; -}; - -struct CommunicationClassInfo { - ClassId classId = 0; - size_t rank = 0; - bool isBatch = false; - mlir::Operation* op = nullptr; - llvm::SmallVector cpus; -}; - -struct CommunicationExchange { - CommunicationExchangeId id; - CommunicationExchangeKind kind = CommunicationExchangeKind::OrdinaryValue; - - PlaceholderId placeholder; - ProducerKey producer; - ClassId sourceClass = 0; - ClassId targetClass = 0; - - size_t sourceRank = 0; - size_t targetRank = 0; - mlir::Type payloadType; - int64_t channelId = 0; - int64_t sourceCore = 0; - int64_t targetCore = 0; - - unsigned sourceLane = 0; - unsigned targetLane = 0; - unsigned ordinal = 0; - uint64_t producerOrder = 0; - - mlir::Value originalHostOutput; - const ProjectedInputTransferFragment* fragment = nullptr; -}; - -using ExchangeDescriptor = CommunicationExchange; - -struct CommunicationEvent { - CommunicationExchangeId exchangeId; - CommunicationEventKind kind = CommunicationEventKind::Send; - ClassId classId = 0; - int64_t core = 0; - unsigned order = 0; -}; - -class CommunicationPlan { -public: - using ProjectedInputPlanMap = - llvm::DenseMap>; - using ProducerSourceClassMap = llvm::DenseMap; - - static mlir::FailureOr - build(mlir::Operation* anchor, - llvm::ArrayRef classes, - const llvm::DenseMap& cpuToClass, - const ProducerSourceClassMap& producerSourceClasses, - const llvm::DenseMap& producerPayloadTypes, - const llvm::DenseMap, ProducerKeyInfo>& producerDestClasses, - const llvm::DenseMap, ProducerKeyInfo>& ordinaryProducerDestClasses, - ProjectedInputPlanMap& projectedPlans, - const llvm::DenseMap& hostOutputOwners, - const llvm::DenseMap& liveExternalUseCache, - int64_t& nextChannelId); - - mlir::LogicalResult verify(mlir::Operation* anchor) const; - - llvm::ArrayRef getEventsForClass(ClassId classId) const; - llvm::ArrayRef getExchanges() const { return exchanges; } - const CommunicationExchange& getExchange(CommunicationExchangeId id) const { return exchanges[id.value]; } - const CommunicationExchange& getExchange(size_t id) const { return exchanges[id]; } - - llvm::ArrayRef - getOrdinaryValueExchanges(ProducerKey producer, ClassId targetClass) const; - llvm::ArrayRef - getProjectedInputExchanges(const ProjectedInputTransferFragment& fragment) const; - llvm::ArrayRef - getPackedScalarRunExchanges(ProducerKey producer, ClassId targetClass) const; - llvm::ArrayRef - getIndexedBatchRunExchanges(ProducerKey producer, ClassId targetClass) const; - llvm::ArrayRef - getHostForwardExchanges(mlir::Value originalOutput, ClassId sourceClass, ClassId ownerClass) const; - -private: - struct ProducerTargetKey { - ProducerKey producer; - ClassId targetClass = 0; - - bool operator==(const ProducerTargetKey& other) const { - return producer == other.producer && targetClass == other.targetClass; - } - }; - - struct ProducerTargetKeyInfo { - static ProducerTargetKey getEmptyKey() { - return {ProducerKeyInfo::getEmptyKey(), llvm::DenseMapInfo::getEmptyKey()}; - } - static ProducerTargetKey getTombstoneKey() { - return {ProducerKeyInfo::getTombstoneKey(), llvm::DenseMapInfo::getTombstoneKey()}; - } - static unsigned getHashValue(const ProducerTargetKey& key) { - return llvm::hash_combine(ProducerKeyInfo::getHashValue(key.producer), key.targetClass); - } - static bool isEqual(const ProducerTargetKey& lhs, const ProducerTargetKey& rhs) { return lhs == rhs; } - }; - - struct HostForwardKey { - mlir::Value originalOutput; - ClassId sourceClass = 0; - ClassId ownerClass = 0; - - bool operator==(const HostForwardKey& other) const { - return originalOutput == other.originalOutput && sourceClass == other.sourceClass - && ownerClass == other.ownerClass; - } - }; - - struct HostForwardKeyInfo { - static HostForwardKey getEmptyKey() { - return {llvm::DenseMapInfo::getEmptyKey(), - llvm::DenseMapInfo::getEmptyKey(), - llvm::DenseMapInfo::getEmptyKey()}; - } - static HostForwardKey getTombstoneKey() { - return {llvm::DenseMapInfo::getTombstoneKey(), - llvm::DenseMapInfo::getTombstoneKey(), - llvm::DenseMapInfo::getTombstoneKey()}; - } - static unsigned getHashValue(const HostForwardKey& key) { - return llvm::hash_combine(key.originalOutput, key.sourceClass, key.ownerClass); - } - static bool isEqual(const HostForwardKey& lhs, const HostForwardKey& rhs) { return lhs == rhs; } - }; - - llvm::SmallVector demands; - llvm::SmallVector publications; - llvm::SmallVector exchanges; - llvm::DenseMap> eventsByClass; - - llvm::DenseMap> - projectedExchangesByFragment; - llvm::DenseMap, ProducerTargetKeyInfo> - ordinaryExchangesByProducerTarget; - llvm::DenseMap, ProducerTargetKeyInfo> - packedScalarRunExchangesByProducerTarget; - llvm::DenseMap, ProducerTargetKeyInfo> - indexedBatchRunExchangesByProducerTarget; - llvm::DenseMap, HostForwardKeyInfo> - hostForwardExchangesByOutput; -}; - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/HostOutputFinalization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/HostOutputFinalization.cpp deleted file mode 100644 index ba683c2..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/HostOutputFinalization.cpp +++ /dev/null @@ -1,134 +0,0 @@ -#include "HostOutputFinalization.hpp" - -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/PatternMatch.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/STLExtras.h" - -#include "MaterializedClassState.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" - -using namespace mlir; - -namespace onnx_mlir::spatial { - -LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) { - if (state.pendingProjectedHostOutputFragments.empty()) - return success(); - - DenseMap> byOutput; - for (PendingProjectedHostOutputFragment& fragment : state.pendingProjectedHostOutputFragments) - byOutput[fragment.originalOutput].push_back(&fragment); - - SmallVector outputs; - outputs.reserve(byOutput.size()); - - auto returnOp = dyn_cast(state.func.getBody().front().getTerminator()); - if (!returnOp) - return state.func.emitError("expected func.return terminator while finalizing projected host output fragments"); - - DenseSet seenOutputs; - for (Value returned : returnOp.getOperands()) { - if (!byOutput.contains(returned) || !seenOutputs.insert(returned).second) - continue; - outputs.push_back(returned); - } - if (outputs.size() != byOutput.size()) - return state.func.emitError("projected host output fragments must be keyed by returned logical host outputs"); - - for (Value originalOutput : outputs) { - if (isa_and_present(originalOutput.getDefiningOp())) { - return state.func.emitError( - "projected host output assembly must be keyed by the original logical host output, not by a materialized scheduled result"); - } - - auto resultType = dyn_cast(originalOutput.getType()); - if (!resultType || !resultType.hasStaticShape()) - return state.func.emitError("projected host output must have static ranked tensor type"); - - SmallVector& fragments = byOutput[originalOutput]; - llvm::sort(fragments, [](const PendingProjectedHostOutputFragment* lhs, - const PendingProjectedHostOutputFragment* rhs) { - if (lhs->sourceClass != rhs->sourceClass) - return lhs->sourceClass < rhs->sourceClass; - if (lhs->publicationResultIndex != rhs->publicationResultIndex) - return lhs->publicationResultIndex < rhs->publicationResultIndex; - if (lhs->sourceFragmentOrdinal != rhs->sourceFragmentOrdinal) - return lhs->sourceFragmentOrdinal < rhs->sourceFragmentOrdinal; - return std::lexicographical_compare(lhs->offsets.begin(), - lhs->offsets.end(), - rhs->offsets.begin(), - rhs->offsets.end()); - }); - - state.rewriter.setInsertionPoint(returnOp); - Location loc = fragments.front()->loc; - SmallVector blueprintOperands; - SmallVector fragmentOperandIndices; - SmallVector fragmentSourceOffsets; - SmallVector flatOffsets; - SmallVector flatSizes; - SmallVector flatStrides; - DenseMap operandIndicesByValue; - - for (PendingProjectedHostOutputFragment* fragmentRecord : fragments) { - if (fragmentRecord->sourceClass >= state.classes.size()) - return state.func.emitError("projected host output fragment references an invalid source class"); - - MaterializedClass& sourceClass = state.classes[fragmentRecord->sourceClass]; - if (fragmentRecord->publicationResultIndex >= sourceClass.op->getNumResults()) { - return sourceClass.op->emitError("projected host output fragment references an invalid publication result") - << " sourceClass=" << sourceClass.id - << " resultIndex=" << fragmentRecord->publicationResultIndex - << " resultCount=" << sourceClass.op->getNumResults(); - } - - Value operand = sourceClass.op->getResult(fragmentRecord->publicationResultIndex); - auto [operandIt, inserted] = - operandIndicesByValue.try_emplace(operand, static_cast(blueprintOperands.size())); - if (inserted) - blueprintOperands.push_back(operand); - fragmentOperandIndices.push_back(operandIt->second); - fragmentSourceOffsets.push_back(fragmentRecord->sourceElementOffset); - llvm::append_range(flatOffsets, fragmentRecord->offsets); - llvm::append_range(flatSizes, fragmentRecord->sizes); - llvm::append_range(flatStrides, fragmentRecord->strides); - - auto operandType = dyn_cast(operand.getType()); - if (!operandType || !operandType.hasStaticShape()) - return state.func.emitError("projected host output assembly requires static ranked tensor operands"); - } - - if (blueprintOperands.empty()) - return state.func.emitError("missing projected host output fragments"); - - Value input = blueprintOperands.front(); - ValueRange extraFragments = ValueRange(blueprintOperands).drop_front(); - auto blueprint = SpatBlueprintOp::create( - state.rewriter, - loc, - resultType, - input, - extraFragments, - state.rewriter.getStringAttr("nchw"), - state.rewriter.getStringAttr("fragmented"), - state.rewriter.getDenseI64ArrayAttr(flatOffsets), - state.rewriter.getDenseI64ArrayAttr(flatSizes), - state.rewriter.getStringAttr("identity"), - state.rewriter.getStringAttr("fragment_assembly"), - state.rewriter.getDenseI64ArrayAttr(fragmentOperandIndices), - state.rewriter.getDenseI64ArrayAttr(fragmentSourceOffsets), - state.rewriter.getDenseI64ArrayAttr(flatStrides), - state.rewriter.getStringAttr("disjoint"), - state.rewriter.getStringAttr("complete")); - - state.hostReplacements[originalOutput] = blueprint.getOutput(); - } - - return success(); -} - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/HostOutputFinalization.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/HostOutputFinalization.hpp deleted file mode 100644 index 29a6fa2..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/HostOutputFinalization.hpp +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once - -#include "mlir/Support/LogicalResult.h" - -namespace onnx_mlir::spatial { - -struct MaterializerState; - -mlir::LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state); - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp deleted file mode 100644 index 05dfec2..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp +++ /dev/null @@ -1,10553 +0,0 @@ -#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/IR/BuiltinOps.h" -#include "mlir/IR/Dominance.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/IR/Matchers.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Transforms/FoldUtils.h" -#include "mlir/Transforms/RegionUtils.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringMap.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/raw_ostream.h" - -#include -#include -#include -#include -#include -#include -#include - -#include "HostOutputFinalization.hpp" -#include "MaterializedClassState.hpp" -#include "MergeMessages.hpp" -#include "MergeScheduleKeys.hpp" -#include "ProjectedFragments.hpp" -#include "Scheduling/ComputeInstanceUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp" -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" - -using namespace mlir; - -namespace onnx_mlir { -namespace spatial { -namespace { - -struct ProjectedProducerFragmentDemand { - ProducerKey producer; - SmallVector fragmentOffsets; - unsigned ordinal = 0; -}; - -FailureOr recordProjectedScalarHostFragmentsFromPackedValue(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef keys, - Value packed, - Value originalOutput, - Location loc); - -FailureOr materializeProjectedExtractReplacement(MaterializerState& state, - MaterializedClass& targetClass, - tensor::ExtractSliceOp extract, - const ProjectedExtractReplacement& replacement, - std::optional projectionSlotIndex, - IRMapping* mapper = nullptr); -FailureOr rematerializeTensorValueInClass(MaterializerState& state, - MaterializedClass& targetClass, - Value value, - Operation* anchor, - StringRef context, - IRMapping* mapper = nullptr); -FailureOr materializeTensorValueForMaterializedClassUse(MaterializerState& state, - MaterializedClass& targetClass, - Value value, - Operation* anchor, - StringRef context, - std::optional producer = std::nullopt, - IRMapping* mapper = nullptr); -FailureOr materializeWholeBatchInput( - MaterializerState& state, MaterializedClass& targetClass, ProducerKey key, Type resultType, Location loc); -FailureOr localizeMaterializedClassOperand(MaterializerState& state, - MaterializedClass& targetClass, - Value value, - Operation* anchor, - StringRef tensorContext, - StringRef genericContext, - IRMapping* mapper = nullptr); -LogicalResult localizeCapturesInClonedOp(MaterializerState& state, - MaterializedClass& targetClass, - Operation& clonedOp, - IRMapping* mapper = nullptr); -LogicalResult localizeAllScheduledBodyCaptures(MaterializerState& state, MaterializedClass& targetClass); -void createDim0ParallelInsertSlice( - MaterializerState& state, Location loc, Value fragment, Value destination, OpFoldResult firstOffset); -Value scaleIndexByDim0Size(MaterializerState& state, Operation* anchor, Value index, int64_t dim0Size, Location loc); -bool isProjectedInputSliceCompatibleWithProducerFragments(SpatComputeBatch consumerBatch, - const AffineProjectedInputSliceMatch& match, - ProducerKey producer, - uint32_t consumerLane); -SmallVector -collectProjectedFragmentDemandsForBatchInput(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer); -SmallVector collectProjectedProducerKeysForBatchInput(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer); -LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef keys, - Value payload, - Location loc); -static FailureOr emitReadyCommunicationEventsForClass(MaterializerState& state, - MaterializedClass& materializedClass, - Location loc); -static Operation* getProjectedReceiveInsertionAnchor(MaterializedClass& targetClass); -static LogicalResult materializeProjectedInputReceivesForClass(MaterializerState& state, - MaterializedClass& targetClass, - Operation* requestedExtract, - Location loc, - Operation* insertionAnchor); -static void tagProjectedInputTransfer(Operation* op); -std::optional -getProjectedInputSliceMatch(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex); -std::optional -getProjectedWholeBatchReplacementProducer(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex); -std::optional getProjectedWholeBatchReplacementProducer(MaterializerState& state, - tensor::ExtractSliceOp extract); -bool hasProjectedInputReplacement(MaterializerState& state, - SpatComputeBatch batch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer, - ClassId classId); -LogicalResult recordLocalPackedSlotValue(MaterializerState& state, - MaterializedClass& materializedClass, - ArrayRef keys, - Value packed); -Operation* getTopLevelMaterializedClassBodyOp(Operation* nestedOp, MaterializedClass& targetClass); -LogicalResult registerPackedRunValue(MaterializerState& state, - MaterializedClass& materializedClass, - ArrayRef keys, - Value packed, - Type fragmentType, - Location loc); -FailureOr getPackedSlotFragmentType(Value packed, size_t keyCount); -SmallVector& getBatchOutputFragmentTypesCached(MaterializerState& state, SpatComputeBatch batch); -static bool isProducerValueMaterialized(MaterializerState& state, ProducerKey key); -static void recordDeferredProducerValue(MaterializerState& state, ProducerKey key); -FailureOr materializeProjectedWholeBatchExtractReplacement(MaterializerState& state, - MaterializedClass& targetClass, - tensor::ExtractSliceOp extract, - ProducerKey producer, - IRMapping* mapper = nullptr); -FailureOr> materializeLocalProjectedInputExtractReplacement(MaterializerState& state, - MaterializedClass& targetClass, - tensor::ExtractSliceOp extract, - IRMapping* mapper = nullptr); -enum class MaterializedLayoutKind { - DenseLogical, - RowPackedNCHWFromRows, - Unknown -}; - -struct MaterializedTensorView { - Value value; - RankedTensorType logicalType; - RankedTensorType physicalType; - MaterializedLayoutKind layout = MaterializedLayoutKind::Unknown; -}; - -FailureOr materializeLogicalTensorView(MaterializerState& state, - MaterializedClass& targetClass, - MaterializedTensorView view, - Operation* anchor, - StringRef context, - std::optional producer = std::nullopt, - IRMapping* mapper = nullptr); -bool isConstantLike(Value value) { - Operation* definingOp = value.getDefiningOp(); - return definingOp && definingOp->hasTrait(); -} - -bool isInsideOldCompute(Operation* op, const DenseSet& oldComputeOps) { - for (Operation* current = op; current; current = current->getParentOp()) - if (oldComputeOps.contains(current)) - return true; - return false; -} - -bool hasLiveExternalUse(Value value, const DenseSet& oldComputeOps); -ArrayRef getComputeInstanceOutputValuesCached(MaterializerState& state, ComputeInstance instance); - -bool hasLiveExternalUseCached(MaterializerState& state, Value value) { - auto cached = state.liveExternalUseCache.find(value); - if (cached != state.liveExternalUseCache.end()) - return cached->second; - bool live = hasLiveExternalUse(value, state.oldComputeOps); - state.liveExternalUseCache[value] = live; - return live; -} - -std::optional getConstantFirstSliceOffset(tensor::ExtractSliceOp extract) { - if (extract.getMixedOffsets().empty()) - return std::nullopt; - - OpFoldResult offset = extract.getMixedOffsets().front(); - if (auto attr = dyn_cast(offset)) { - auto intAttr = dyn_cast(attr); - if (!intAttr || intAttr.getInt() < 0) - return std::nullopt; - return static_cast(intAttr.getInt()); - } - - auto value = cast(offset); - if (auto constantIndex = value.getDefiningOp()) { - if (constantIndex.value() < 0) - return std::nullopt; - return static_cast(constantIndex.value()); - } - - APInt constantValue; - if (matchPattern(value, m_ConstantInt(&constantValue))) { - if (constantValue.isNegative()) - return std::nullopt; - return static_cast(constantValue.getZExtValue()); - } - - return std::nullopt; -} - -ProducerKey -getBatchLaneProducerKey(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount, size_t resultIndex) { - return { - {batch.getOperation(), laneStart, laneCount}, - resultIndex - }; -} - -ProducerKey getWholeBatchProducerKey(SpatComputeBatch batch, size_t resultIndex) { - return getBatchLaneProducerKey(batch, 0, static_cast(batch.getLaneCount()), resultIndex); -} - -bool isWholeBatchProducerKey(ProducerKey key) { - auto batch = dyn_cast_or_null(key.instance.op); - return batch && batch.getNumResults() != 0 && key.instance.laneStart == 0 - && key.instance.laneCount == static_cast(batch.getLaneCount()); -} - -std::optional getContiguousProducerRangeForKeys(ArrayRef keys) { - if (keys.empty()) - return std::nullopt; - - ProducerKey first = keys.front(); - auto batch = dyn_cast_or_null(first.instance.op); - if (!batch) - return std::nullopt; - - SmallVector sorted(keys.begin(), keys.end()); - llvm::sort(sorted, [](ProducerKey lhs, ProducerKey rhs) { - return std::tie(lhs.instance.laneStart, lhs.instance.laneCount, lhs.resultIndex) - < std::tie(rhs.instance.laneStart, rhs.instance.laneCount, rhs.resultIndex); - }); - - uint32_t laneStart = sorted.front().instance.laneStart; - uint32_t nextLane = laneStart; - for (ProducerKey key : sorted) { - if (key.instance.op != first.instance.op || key.resultIndex != first.resultIndex || key.instance.laneCount == 0) - return std::nullopt; - if (key.instance.laneStart != nextLane) - return std::nullopt; - nextLane += key.instance.laneCount; - } - - uint32_t laneCount = nextLane - laneStart; - if (laneStart + laneCount > static_cast(batch.getLaneCount())) - return std::nullopt; - - return getBatchLaneProducerKey(batch, laneStart, laneCount, first.resultIndex); -} - -WholeBatchAssemblyLookupKey makeWholeBatchAssemblyLookupKey(Operation* sourceOp, size_t resultIndex, ClassId classId) { - return {sourceOp, resultIndex, classId}; -} - -WholeBatchAssemblyLookupKey makeWholeBatchAssemblyLookupKey(ProducerKey key, ClassId classId) { - return makeWholeBatchAssemblyLookupKey(key.instance.op, key.resultIndex, classId); -} - -FailureOr getPackedBatchTensorType(Type laneType, size_t laneCount) { - auto tensorType = dyn_cast(laneType); - if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0) - return failure(); - - SmallVector shape(tensorType.getShape()); - shape[0] *= static_cast(laneCount); - return RankedTensorType::get(shape, tensorType.getElementType(), tensorType.getEncoding()); -} - -ComputeInstance getScheduledChunkForLogicalInstance(MaterializerState& state, ComputeInstance logicalInstance) { - auto it = state.logicalInstanceToScheduledChunk.find(logicalInstance); - if (it != state.logicalInstanceToScheduledChunk.end()) - return it->second; - return logicalInstance; -} - -FailureOr -getPublicationLaneForProducerKey(MaterializerState& state, const MaterializedClass& sourceClass, ProducerKey key) { - if (!sourceClass.isBatch) - return 0; - - ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, key.instance); - auto cpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); - if (cpuIt == state.schedule.computeToCpuMap.end()) { - sourceClass.op->emitError( - "projected packed host publication could not resolve the producer CPU for a publication lane") - << " laneStart=" << key.instance.laneStart << " laneCount=" << key.instance.laneCount - << " resultIndex=" << key.resultIndex; - return failure(); - } - - auto laneIt = sourceClass.cpuToLane.find(cpuIt->second); - if (laneIt == sourceClass.cpuToLane.end()) { - sourceClass.op->emitError("projected packed host publication could not map a producer key to a publication lane") - << " cpu=" << cpuIt->second << " laneStart=" << key.instance.laneStart << " laneCount=" << key.instance.laneCount - << " resultIndex=" << key.resultIndex; - return failure(); - } - - return laneIt->second; -} - -SmallVector -collectProducerKeysForDestinations(Value value, std::optional logicalConsumer = std::nullopt) { - // Destination collection works in the materializer's logical one-lane key domain. - // Whole-batch resultful producers are expanded into per-lane producer keys here. - SmallVector keys; - Operation* definingOp = value.getDefiningOp(); - if (!definingOp) - return keys; - - while (auto extract = dyn_cast(definingOp)) { - Value source = extract.getSource(); - auto batch = dyn_cast_or_null(source.getDefiningOp()); - if (batch && batch.getNumResults() != 0) { - auto result = dyn_cast(source); - if (!result) - return {}; - - if (std::optional lane = getConstantFirstSliceOffset(extract)) { - if (*lane >= static_cast(batch.getLaneCount())) - return {}; - keys.push_back(getBatchLaneProducerKey(batch, *lane, 1, result.getResultNumber())); - return keys; - } - - if (logicalConsumer && isa(logicalConsumer->op)) { - keys.push_back(getBatchLaneProducerKey(batch, logicalConsumer->laneStart, 1, result.getResultNumber())); - return keys; - } - - return {}; - } - - value = source; - definingOp = value.getDefiningOp(); - if (!definingOp) - return {}; - } - - if (auto compute = dyn_cast(definingOp)) { - auto result = dyn_cast(value); - if (!result) - return {}; - keys.push_back({ - {compute.getOperation(), 0, 1}, - result.getResultNumber() - }); - return keys; - } - - if (auto batch = dyn_cast(definingOp)) { - auto result = dyn_cast(value); - if (!result) - return {}; - - if (batch.getNumResults() != 0) { - if (logicalConsumer && isa(logicalConsumer->op)) { - keys.push_back(getBatchLaneProducerKey(batch, logicalConsumer->laneStart, 1, result.getResultNumber())); - return keys; - } - - for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) - keys.push_back(getBatchLaneProducerKey(batch, lane, 1, result.getResultNumber())); - return keys; - } - - ComputeInstance chunk = getBatchChunkForLane(batch, result.getResultNumber()); - keys.push_back({chunk, static_cast(result.getResultNumber() - chunk.laneStart)}); - return keys; - } - - return keys; -} - -std::optional getInputRequestProducerKey(Value value, - std::optional logicalConsumer = std::nullopt) { - // Input resolution may request a whole-batch key for scalar consumers that read - // a complete resultful compute_batch value. - Operation* definingOp = value.getDefiningOp(); - if (!definingOp) - return std::nullopt; - - while (auto extract = dyn_cast(definingOp)) { - Value source = extract.getSource(); - auto batch = dyn_cast_or_null(source.getDefiningOp()); - if (batch && batch.getNumResults() != 0) { - auto result = dyn_cast(source); - if (!result) - return std::nullopt; - - if (std::optional lane = getConstantFirstSliceOffset(extract)) - return getBatchLaneProducerKey(batch, *lane, 1, result.getResultNumber()); - - if (logicalConsumer && isa(logicalConsumer->op)) - return getBatchLaneProducerKey(batch, logicalConsumer->laneStart, 1, result.getResultNumber()); - - return std::nullopt; - } - - value = source; - definingOp = value.getDefiningOp(); - if (!definingOp) - return std::nullopt; - } - - if (auto compute = dyn_cast(definingOp)) { - auto result = dyn_cast(value); - if (!result) - return std::nullopt; - return ProducerKey { - {compute.getOperation(), 0, 1}, - result.getResultNumber() - }; - } - - if (auto batch = dyn_cast(definingOp)) { - auto result = dyn_cast(value); - if (!result) - return std::nullopt; - - if (batch.getNumResults() != 0) { - if (logicalConsumer && isa(logicalConsumer->op)) - return getBatchLaneProducerKey(batch, logicalConsumer->laneStart, 1, result.getResultNumber()); - return getWholeBatchProducerKey(batch, result.getResultNumber()); - } - - return ProducerKey {getBatchChunkForLane(batch, result.getResultNumber()), 0}; - } - - return std::nullopt; -} - -std::optional getWholeBatchProducerKeyForDirectBatchResult(Value value) { - auto result = dyn_cast(value); - if (!result) - return std::nullopt; - - auto batch = dyn_cast_or_null(result.getOwner()); - if (!batch || batch.getNumResults() == 0) - return std::nullopt; - - return getWholeBatchProducerKey(batch, result.getResultNumber()); -} - -bool canUseProjectedLaneInput(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer) { - return !collectProjectedFragmentDemandsForBatchInput(state, consumerBatch, inputIndex, input, logicalConsumer).empty(); -} - -static bool isDebugPaddedVggInput(Value input) { - auto type = dyn_cast(input.getType()); - return type && type.hasStaticShape() && type.getRank() == 4 && type.getDimSize(0) == 1 - && type.getDimSize(1) == 3 && type.getDimSize(2) == 226 && type.getDimSize(3) == 226; -} - -static bool hasProjectedPlanForBatchInputKey(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - ClassId classId) { - ProjectedBatchInputKey inputKey {consumerBatch.getOperation(), inputIndex}; - for (const auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(classId); - if (classIt != extractEntry.second.end() && classIt->second.inputKey == inputKey) - return true; - } - - for (const auto& producerEntry : state.projectedTransfers) { - auto classIt = producerEntry.second.find(classId); - if (classIt != producerEntry.second.end() && classIt->second.inputKey == inputKey) - return true; - } - - return false; -} - -static SmallVector collectProjectedInputKeyProducersForClass(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - ClassId classId) { - ProjectedBatchInputKey inputKey {consumerBatch.getOperation(), inputIndex}; - SmallVector producers; - - auto appendProducer = [&](ProducerKey producer) { - if (!llvm::is_contained(producers, producer)) - producers.push_back(producer); - }; - - for (const auto& producerEntry : state.projectedTransfers) { - auto classIt = producerEntry.second.find(classId); - if (classIt != producerEntry.second.end() && classIt->second.inputKey == inputKey) - appendProducer(producerEntry.first); - } - - for (const auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(classId); - if (classIt == extractEntry.second.end() || !(classIt->second.inputKey == inputKey)) - continue; - for (const ProjectedInputTransferFragment& fragment : classIt->second.fragments) - appendProducer(fragment.producer); - } - - return producers; -} - -static bool hasUnreplacedBatchInputUse(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - ClassId classId) { - auto inputArg = consumerBatch.getInputArgument(inputIndex); - if (!inputArg) - return true; - - for (OpOperand& use : inputArg->getUses()) { - auto extract = dyn_cast(use.getOwner()); - if (!extract) - return true; - - auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); - if (planIt != state.projectedInputTransferPlans.end() && planIt->second.find(classId) != planIt->second.end()) - continue; - if (hasProjectedPlanForBatchInputKey(state, consumerBatch, inputIndex, classId)) - continue; - if (getProjectedWholeBatchReplacementProducer(state, extract)) - continue; - if (inputIndex == 0) { - static unsigned debugCount = 0; - if (debugCount++ < 16) { - llvm::errs() << "debug unreplaced batch input use classId=" << classId - << " consumer='" << consumerBatch->getName() << "'" - << " useOwner='" << use.getOwner()->getName() << "'" - << " isExtract=" << static_cast(extract) - << " hasPlan=" << (planIt != state.projectedInputTransferPlans.end()) - << " planHasClass=" - << (planIt != state.projectedInputTransferPlans.end() - && planIt->second.find(classId) != planIt->second.end()) - << " hasWholeReplacement=" - << getProjectedWholeBatchReplacementProducer(state, extract).has_value() - << "\n"; - } - } - return true; - } - return false; -} - -static bool hasUnreplacedComputeInputUse(MaterializerState& state, - SpatCompute compute, - unsigned inputIndex, - ClassId classId) { - auto inputArg = compute.getInputArgument(inputIndex); - if (!inputArg) - return true; - - for (OpOperand& use : inputArg->getUses()) { - auto extract = dyn_cast(use.getOwner()); - if (!extract) - return true; - - auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); - if (planIt != state.projectedInputTransferPlans.end() && planIt->second.find(classId) != planIt->second.end()) - continue; - return true; - } - return false; -} - -static bool hasProjectedComputeInputReplacement(MaterializerState& state, - SpatCompute compute, - unsigned inputIndex, - ClassId classId) { - auto inputArg = compute.getInputArgument(inputIndex); - if (!inputArg) - return false; - - for (OpOperand& use : inputArg->getUses()) { - auto extract = dyn_cast(use.getOwner()); - if (!extract) - continue; - - auto replacementIt = state.projectedExtractReplacements.find(extract.getOperation()); - if (replacementIt != state.projectedExtractReplacements.end() - && replacementIt->second.find(classId) != replacementIt->second.end()) - return true; - - auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); - if (planIt != state.projectedInputTransferPlans.end() && planIt->second.find(classId) != planIt->second.end()) - return true; - } - return false; -} - -FailureOr classifyComputeBatchInputDemand(MaterializerState& state, - MaterializedClass& targetClass, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer) { - if (std::optional wholeBatchProducer = getWholeBatchProducerKeyForDirectBatchResult(input)) { - if (hasUnreplacedBatchInputUse(state, consumerBatch, inputIndex, targetClass.id)) - return BatchInputDemand {.kind = BatchInputDemandKind::WholeTensorBarrier, - .wholeTensorProducer = wholeBatchProducer}; - - if (canUseProjectedLaneInput(state, consumerBatch, inputIndex, input, logicalConsumer)) - return BatchInputDemand {.kind = BatchInputDemandKind::ProjectedFragment, .wholeTensorProducer = std::nullopt}; - - if (getProjectedWholeBatchReplacementProducer(state, consumerBatch, inputIndex)) - return BatchInputDemand {.kind = BatchInputDemandKind::ProjectedFragment, .wholeTensorProducer = std::nullopt}; - - auto inputArg = consumerBatch.getInputArgument(inputIndex); - if (!inputArg) - return consumerBatch.emitOpError("expected compute_batch input block argument while classifying input demand") - << " #" << inputIndex; - - bool hasUses = false; - for (OpOperand& use : inputArg->getUses()) { - hasUses = true; - if (!isa(use.getOwner())) - return BatchInputDemand {.kind = BatchInputDemandKind::WholeTensorBarrier, - .wholeTensorProducer = wholeBatchProducer}; - } - if (!hasUses) - return BatchInputDemand {.kind = BatchInputDemandKind::LaneFragment, .wholeTensorProducer = std::nullopt}; - - return targetClass.op->emitError("failed to classify compute_batch input demand") - << " reason=direct whole-batch input only has projected uses, but no projected fragment path was proven" - << " consumerOp='" << consumerBatch->getName() << "' inputIndex=" << inputIndex << " producerOp='" - << wholeBatchProducer->instance.op->getName() << "' resultIndex=" << wholeBatchProducer->resultIndex - << " sourceClass=" << targetClass.id << " valueType=" << input.getType(); - } - - return BatchInputDemand {.kind = BatchInputDemandKind::LaneFragment, .wholeTensorProducer = std::nullopt}; -} - -SmallVector collectProducerKeysForBatchInputDestinations(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer) { - if (std::optional wholeBatchProducer = getWholeBatchProducerKeyForDirectBatchResult(input)) { - if (!canUseProjectedLaneInput(state, consumerBatch, inputIndex, input, logicalConsumer)) { - SmallVector projectedKeys = - collectProjectedProducerKeysForBatchInput(state, consumerBatch, inputIndex, input, logicalConsumer); - if (!projectedKeys.empty()) - return projectedKeys; - - auto producerBatch = cast(wholeBatchProducer->instance.op); - SmallVector keys; - for (uint32_t lane = 0; lane < static_cast(producerBatch.getLaneCount()); ++lane) - keys.push_back(getBatchLaneProducerKey(producerBatch, lane, 1, wholeBatchProducer->resultIndex)); - return keys; - } - } - - return collectProducerKeysForDestinations(input, logicalConsumer); -} - -class CpuUnionFind { -public: - void insert(CpuId cpu) { parent.try_emplace(cpu, cpu); } - - CpuId find(CpuId cpu) { - insert(cpu); - CpuId p = parent.lookup(cpu); - if (p == cpu) - return cpu; - CpuId root = find(p); - parent[cpu] = root; - return root; - } - - void unite(CpuId lhs, CpuId rhs) { - CpuId lhsRoot = find(lhs); - CpuId rhsRoot = find(rhs); - if (lhsRoot == rhsRoot) - return; - if (rhsRoot < lhsRoot) - std::swap(lhsRoot, rhsRoot); - parent[rhsRoot] = lhsRoot; - } - -private: - DenseMap parent; -}; - -LogicalResult buildMaterializationWorkStreams(MaterializerState& state) { - DenseMap> scheduledInstancesByCpu; - for (const auto& [instance, cpu] : state.schedule.computeToCpuMap) { - state.oldComputeOps.insert(instance.op); - scheduledInstancesByCpu[cpu].push_back(instance); - state.logicalInstancesByCpu.try_emplace(cpu); - } - - for (auto& [cpu, scheduledInstances] : scheduledInstancesByCpu) { - llvm::sort(scheduledInstances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) { - auto lhsIt = state.schedule.computeToCpuSlotMap.find(lhs); - auto rhsIt = state.schedule.computeToCpuSlotMap.find(rhs); - assert(lhsIt != state.schedule.computeToCpuSlotMap.end() && "missing scheduler slot"); - assert(rhsIt != state.schedule.computeToCpuSlotMap.end() && "missing scheduler slot"); - return lhsIt->second < rhsIt->second; - }); - - SmallVector& logicalInstances = state.logicalInstancesByCpu[cpu]; - SlotId logicalSlot = 0; - for (const ComputeInstance& instance : scheduledInstances) { - LogicalSlotRange range {logicalSlot, 1}; - if (isa(instance.op)) - range.count = instance.laneCount; - - state.scheduledInstanceToLogicalSlots[instance] = range; - - if (isa(instance.op)) { - for (uint32_t localLane = 0; localLane < instance.laneCount; ++localLane, ++logicalSlot) { - uint32_t logicalLane = instance.laneStart + localLane; - ComputeInstance logicalInstance {instance.op, logicalLane, 1}; - logicalInstances.push_back(logicalInstance); - state.logicalInstanceToScheduledChunk[logicalInstance] = instance; - } - continue; - } - - logicalInstances.push_back(instance); - ++logicalSlot; - } - } - - return success(); -} - - -using MaterializedBatchingSignature = SmallVector; - -static int64_t getProjectedInputExchangeDirection(int64_t sourceCoreId, int64_t targetCoreId) { - if (sourceCoreId == targetCoreId) - return 0; - return sourceCoreId < targetCoreId ? 1 : 2; -} - -static LogicalResult appendProjectedBatchingSignatureForCpu(MaterializerState& state, - CpuId cpu, - MaterializedBatchingSignature& signature) { - auto streamIt = state.logicalInstancesByCpu.find(cpu); - if (streamIt == state.logicalInstancesByCpu.end()) - return success(); - - for (const ComputeInstance& logicalConsumer : streamIt->second) { - auto batchConsumer = dyn_cast(logicalConsumer.op); - if (!batchConsumer) - continue; - - SmallVector consumerInputs = getComputeInstanceInputs(logicalConsumer); - for (auto [inputIndex, input] : llvm::enumerate(consumerInputs)) { - SmallVector demands = collectProjectedFragmentDemandsForBatchInput( - state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); - if (demands.empty()) - continue; - - signature.push_back(-1); - signature.push_back(static_cast(inputIndex)); - for (const ProjectedProducerFragmentDemand& demand : demands) { - ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, demand.producer.instance); - auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); - if (producerCpuIt == state.schedule.computeToCpuMap.end()) - return logicalConsumer.op->emitError( - "projected input batching signature found a fragment produced by an unscheduled compute"); - signature.push_back(getProjectedInputExchangeDirection(producerCpuIt->second, cpu)); - } - } - } - - return success(); -} - -static DenseMap buildCpuToMaterializedGroup(ArrayRef> groups) { - DenseMap cpuToGroup; - for (auto [groupIndex, group] : llvm::enumerate(groups)) - for (CpuId cpu : group) - cpuToGroup[cpu] = groupIndex; - return cpuToGroup; -} - -static LogicalResult appendOrdinaryInputBatchingSignatureForCpu(MaterializerState& state, - CpuId cpu, - const DenseMap& cpuToGroup, - MaterializedBatchingSignature& signature) { - auto streamIt = state.logicalInstancesByCpu.find(cpu); - if (streamIt == state.logicalInstancesByCpu.end()) - return success(); - - for (const ComputeInstance& logicalConsumer : streamIt->second) { - SmallVector consumerInputs = getComputeInstanceInputs(logicalConsumer); - for (auto [inputIndex, input] : llvm::enumerate(consumerInputs)) { - SmallVector producerKeys; - if (auto batchConsumer = dyn_cast(logicalConsumer.op)) { - if (hasProjectedInputReplacement( - state, batchConsumer, static_cast(inputIndex), input, logicalConsumer, state.cpuToClass.lookup(cpu)) - && !hasUnreplacedBatchInputUse(state, batchConsumer, static_cast(inputIndex), - state.cpuToClass.lookup(cpu))) - continue; - producerKeys = collectProducerKeysForBatchInputDestinations( - state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); - } - else { - producerKeys = collectProducerKeysForDestinations(input, logicalConsumer); - } - - if (producerKeys.empty()) - continue; - - signature.push_back(-2); - signature.push_back(static_cast(inputIndex)); - for (ProducerKey producerKey : producerKeys) { - ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producerKey.instance); - auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); - if (producerCpuIt == state.schedule.computeToCpuMap.end()) - return logicalConsumer.op->emitError( - "ordinary input batching signature found a fragment produced by an unscheduled compute"); - auto producerGroupIt = cpuToGroup.find(producerCpuIt->second); - if (producerGroupIt == cpuToGroup.end()) - return logicalConsumer.op->emitError( - "ordinary input batching signature found a producer outside materialized CPU groups"); - signature.push_back(static_cast(producerGroupIt->second)); - } - } - } - - return success(); -} - -static LogicalResult buildMaterializedBatchingSignature(MaterializerState& state, - CpuId cpu, - const DenseMap& cpuToGroup, - MaterializedBatchingSignature& signature) { - if (failed(appendProjectedBatchingSignatureForCpu(state, cpu, signature))) - return failure(); - return appendOrdinaryInputBatchingSignatureForCpu(state, cpu, cpuToGroup, signature); -} - -static std::string stringifyMaterializedBatchingSignature(ArrayRef signature) { - std::string key; - llvm::raw_string_ostream os(key); - for (int64_t value : signature) - os << value << ','; - os.flush(); - return key; -} - -static LogicalResult refineCpuGroupsByMaterializedBatchingSignature(MaterializerState& state, - SmallVectorImpl>& groups) { - bool changed = true; - while (changed) { - changed = false; - DenseMap cpuToGroup = buildCpuToMaterializedGroup(groups); - SmallVector, 8> refinedGroups; - refinedGroups.reserve(groups.size()); - - for (ArrayRef group : groups) { - llvm::StringMap groupBySignature; - size_t groupCountBefore = refinedGroups.size(); - for (CpuId cpu : group) { - MaterializedBatchingSignature signature; - if (failed(buildMaterializedBatchingSignature(state, cpu, cpuToGroup, signature))) - return failure(); - - std::string key = stringifyMaterializedBatchingSignature(signature); - auto [it, inserted] = groupBySignature.try_emplace(key, refinedGroups.size()); - if (inserted) - refinedGroups.emplace_back(); - refinedGroups[it->second].push_back(cpu); - } - changed |= refinedGroups.size() != groupCountBefore + 1; - } - - groups.clear(); - for (SmallVector& group : refinedGroups) - groups.push_back(std::move(group)); - } - - return success(); -} - -LogicalResult buildMaterializationClassesFromScheduleEquivalence(MaterializerState& state) { - DenseSet usedCpus; - for (const auto& entry : state.schedule.cpuToLastComputeMap) - usedCpus.insert(entry.first); - for (const auto& entry : state.schedule.computeToCpuMap) - usedCpus.insert(entry.second); - - CpuUnionFind unionFind; - for (CpuId cpu : usedCpus) - unionFind.insert(cpu); - - for (const auto& [cpu, equivalentCpus] : state.schedule.equivalentClass) { - if (!usedCpus.contains(cpu)) - continue; - for (CpuId equivalentCpu : equivalentCpus) - if (usedCpus.contains(equivalentCpu)) - unionFind.unite(cpu, equivalentCpu); - } - - DenseMap> groupsByRoot; - for (CpuId cpu : usedCpus) - groupsByRoot[unionFind.find(cpu)].push_back(cpu); - - SmallVector roots; - roots.reserve(groupsByRoot.size()); - for (const auto& entry : groupsByRoot) - roots.push_back(entry.first); - llvm::sort(roots); - - SmallVector, 8> classCpuGroups; - classCpuGroups.reserve(roots.size()); - for (CpuId root : roots) { - SmallVector rootCpus = groupsByRoot.lookup(root); - llvm::sort(rootCpus); - classCpuGroups.push_back(std::move(rootCpus)); - } - if (failed(refineCpuGroupsByMaterializedBatchingSignature(state, classCpuGroups))) - return failure(); - - state.classes.reserve(classCpuGroups.size()); - for (SmallVector& cpus : classCpuGroups) { - MaterializedClass materializedClass; - materializedClass.id = state.classes.size(); - materializedClass.cpus = std::move(cpus); - llvm::sort(materializedClass.cpus); - materializedClass.isBatch = materializedClass.cpus.size() > 1; - for (auto [lane, cpu] : llvm::enumerate(materializedClass.cpus)) { - materializedClass.cpuToLane[cpu] = static_cast(lane); - state.cpuToClass[cpu] = materializedClass.id; - } - state.classes.push_back(std::move(materializedClass)); - } - - return success(); -} - -LogicalResult verifyScheduleEquivalenceMatchesLogicalStreams(MaterializerState& state) { - for (const MaterializedClass& materializedClass : state.classes) { - if (materializedClass.cpus.empty()) - continue; - - auto referenceIt = state.logicalInstancesByCpu.find(materializedClass.cpus.front()); - if (referenceIt == state.logicalInstancesByCpu.end()) - return state.func.emitError("missing logical stream for materialized class reference CPU"); - - ArrayRef referenceStream(referenceIt->second); - for (CpuId cpu : materializedClass.cpus) { - auto streamIt = state.logicalInstancesByCpu.find(cpu); - if (streamIt == state.logicalInstancesByCpu.end()) - return state.func.emitError("missing logical stream for materialized class CPU"); - - ArrayRef stream(streamIt->second); - if (stream.size() != referenceStream.size()) - return state.func.emitError("materialized class CPUs have mismatched logical stream lengths"); - - for (auto [slot, zipped] : llvm::enumerate(llvm::zip(referenceStream, stream))) { - const ComputeInstance& referenceInstance = std::get<0>(zipped); - const ComputeInstance& currentInstance = std::get<1>(zipped); - if (referenceInstance.op != currentInstance.op) - return state.func.emitError("materialized class logical slot source op mismatch"); - if (isa(referenceInstance.op) != isa(currentInstance.op)) - return state.func.emitError("materialized class logical slot batch/scalar mismatch"); - (void) slot; - } - } - } - - return success(); -} - -LogicalResult forEachLogicalConsumerInMaterializationOrder( - MaterializerState& state, - llvm::function_ref - callback) { - for (const ComputeInstance& scheduledInstance : state.schedule.dominanceOrderCompute) { - auto cpuIt = state.schedule.computeToCpuMap.find(scheduledInstance); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return scheduledInstance.op->emitError("missing CPU assignment for scheduled logical-slot iteration"); - - auto rangeIt = state.scheduledInstanceToLogicalSlots.find(scheduledInstance); - if (rangeIt == state.scheduledInstanceToLogicalSlots.end()) - return scheduledInstance.op->emitError("missing logical slot range for scheduled logical-slot iteration"); - - CpuId cpu = cpuIt->second; - ClassId classId = state.cpuToClass.lookup(cpu); - LogicalSlotRange range = rangeIt->second; - auto streamIt = state.logicalInstancesByCpu.find(cpu); - if (streamIt == state.logicalInstancesByCpu.end()) - return scheduledInstance.op->emitError("missing logical stream for CPU"); - for (SlotId logicalSlot = range.start; logicalSlot < range.start + range.count; ++logicalSlot) { - if (logicalSlot >= streamIt->second.size()) - return scheduledInstance.op->emitError("missing logical slot materialization instance"); - if (failed(callback(cpu, classId, scheduledInstance, streamIt->second[logicalSlot], logicalSlot))) - return failure(); - } - } - - return success(); -} - -bool isTerminalHostBatchOutput(Value output, const DenseSet& oldComputeOps); - -LogicalResult collectHostOutputs(MaterializerState& state) { - DenseSet seenOutputs; - SmallVector orderedOutputs; - DenseMap preferredOwners; - for (const ComputeInstance& instance : state.schedule.dominanceOrderCompute) { - auto cpuIt = state.schedule.computeToCpuMap.find(instance); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return instance.op->emitError("schedule materialization expected a CPU assignment for every compute instance"); - - ClassId classId = state.cpuToClass.lookup(cpuIt->second); - MaterializedClass& materializedClass = state.classes[classId]; - for (Value output : getComputeInstanceOutputValuesCached(state, instance)) { - if (!hasLiveExternalUseCached(state, output)) - continue; - - if (seenOutputs.insert(output).second) { - orderedOutputs.push_back(output); - preferredOwners[output] = classId; - continue; - } - - auto batch = dyn_cast_or_null(output.getDefiningOp()); - if (!batch || batch.getNumResults() == 0) - continue; - - ClassId currentOwner = preferredOwners.lookup(output); - bool terminalHost = isTerminalHostBatchOutput(output, state.oldComputeOps); - if (terminalHost) { - // Terminal batch outputs should stay owned by the producing batch class - // so publication remains explicit in IR via the batch host-output path. - if (!state.classes[currentOwner].isBatch && materializedClass.isBatch) - preferredOwners[output] = classId; - continue; - } - - // Keep batch-defined outputs on a batch owner whenever one exists so - // publication and reconstruction remain explicit in Spatial IR instead of - // falling back to late scalar host forwarding. - if (!state.classes[currentOwner].isBatch && materializedClass.isBatch) - preferredOwners[output] = classId; - } - } - - for (MaterializedClass& materializedClass : state.classes) { - materializedClass.hostOutputs.clear(); - materializedClass.hostOutputToResultIndex.clear(); - materializedClass.publicationOutputToResultIndex.clear(); - } - state.hostOutputOwners.clear(); - - for (Value output : orderedOutputs) { - ClassId ownerClassId = preferredOwners.lookup(output); - MaterializedClass& ownerClass = state.classes[ownerClassId]; - ownerClass.hostOutputToResultIndex[output] = ownerClass.hostOutputs.size(); - ownerClass.hostOutputs.push_back(output); - state.hostOutputOwners[output] = ownerClassId; - } - - return success(); -} - -LogicalResult createEmptyMaterializedOps(MaterializerState& state) { - Location loc = state.func.getLoc(); - Block& funcBlock = state.func.getBody().front(); - - Operation* firstOldCompute = nullptr; - for (Operation& op : funcBlock) { - if (state.oldComputeOps.contains(&op)) { - firstOldCompute = &op; - break; - } - } - - if (firstOldCompute) - state.rewriter.setInsertionPoint(firstOldCompute); - else - state.rewriter.setInsertionPointToStart(&funcBlock); - - for (MaterializedClass& materializedClass : state.classes) { - SmallVector resultTypes; - resultTypes.reserve(materializedClass.hostOutputs.size()); - for (Value output : materializedClass.hostOutputs) - resultTypes.push_back(output.getType()); - - if (!materializedClass.isBatch) { - auto compute = - SpatScheduledCompute::create(state.rewriter, loc, TypeRange(resultTypes), ValueRange {}, ValueRange {}); - compute.getProperties().setOperandSegmentSizes({0, 0}); - auto coreIdAttr = - pim::getCheckedI32Attr(state.rewriter, state.func, materializedClass.cpus.front(), "materialized core id"); - if (failed(coreIdAttr)) - return failure(); - compute->setAttr(onnx_mlir::kCoreIdAttrName, *coreIdAttr); - Block* body = state.rewriter.createBlock(&compute.getBody()); - state.rewriter.setInsertionPointToEnd(body); - SmallVector placeholderOutputs; - placeholderOutputs.reserve(resultTypes.size()); - for (Type resultType : resultTypes) { - auto tensorType = dyn_cast(resultType); - if (!tensorType || !tensorType.hasStaticShape()) { - compute.emitOpError("host-facing materialized compute results must be static ranked tensors"); - return failure(); - } - placeholderOutputs.push_back( - tensor::EmptyOp::create(state.rewriter, loc, tensorType.getShape(), tensorType.getElementType()).getResult()); - } - SpatYieldOp::create(state.rewriter, loc, ValueRange(placeholderOutputs)); - materializedClass.op = compute.getOperation(); - materializedClass.body = body; - state.rewriter.setInsertionPointAfter(compute.getOperation()); - continue; - } - - auto batchLaneCountAttr = pim::getCheckedI32Attr( - state.rewriter, state.func, materializedClass.cpus.size(), "materialized batch lane count"); - if (failed(batchLaneCountAttr)) - return failure(); - auto batch = SpatScheduledComputeBatch::create( - state.rewriter, loc, TypeRange(resultTypes), *batchLaneCountAttr, ValueRange {}, ValueRange {}); - batch.getProperties().setOperandSegmentSizes({0, 0}); - auto coreIds = getCheckedCoreIds(state.func, materializedClass.cpus, "materialized batch core id"); - if (failed(coreIds)) - return failure(); - batch->setAttr(onnx_mlir::kCoreIdsAttrName, state.rewriter.getDenseI32ArrayAttr(*coreIds)); - - SmallVector blockArgTypes {state.rewriter.getIndexType()}; - SmallVector blockArgLocs {loc}; - llvm::append_range(blockArgTypes, resultTypes); - blockArgLocs.append(resultTypes.size(), loc); - Block* body = - state.rewriter.createBlock(&batch.getBody(), batch.getBody().end(), TypeRange(blockArgTypes), blockArgLocs); - state.rewriter.setInsertionPointToEnd(body); - if (resultTypes.empty()) - SpatYieldOp::create(state.rewriter, loc, ValueRange {}); - else - SpatInParallelOp::create(state.rewriter, loc); - materializedClass.op = batch.getOperation(); - materializedClass.body = body; - state.rewriter.setInsertionPointAfter(batch.getOperation()); - } - - return success(); -} - -BlockArgument appendWeight(MaterializerState& state, MaterializedClass& materializedClass, Value weight) { - auto it = materializedClass.weightArgs.find(weight); - if (it != materializedClass.weightArgs.end()) - return it->second; - - unsigned weightIndex = materializedClass.weights.size(); - materializedClass.weights.push_back(weight); - - if (auto compute = dyn_cast(materializedClass.op)) { - auto arg = compute.insertWeight(weightIndex, weight, weight.getLoc()); - assert(arg && "expected compute body while inserting a weight"); - materializedClass.weightArgs[weight] = std::get<1>(*arg); - return std::get<1>(*arg); - } - - auto batch = cast(materializedClass.op); - auto arg = batch.insertWeight(weightIndex, weight, weight.getLoc()); - assert(arg && "expected compute_batch body while inserting a weight argument"); - materializedClass.weightArgs[weight] = std::get<1>(*arg); - return std::get<1>(*arg); -} - -BlockArgument appendInput(MaterializerState& state, MaterializedClass& materializedClass, Value input) { - auto it = materializedClass.inputArgs.find(input); - if (it != materializedClass.inputArgs.end()) - return it->second; - - materializedClass.inputs.push_back(input); - if (auto compute = dyn_cast(materializedClass.op)) { - auto arg = compute.insertInput(materializedClass.inputs.size() - 1, input, input.getLoc()); - assert(arg && "expected compute body while inserting an input"); - materializedClass.inputArgs[input] = std::get<1>(*arg); - return std::get<1>(*arg); - } - - auto compute = cast(materializedClass.op); - auto arg = compute.insertInput(materializedClass.inputs.size() - 1, input, input.getLoc()); - assert(arg && "expected compute_batch body while inserting an input argument"); - materializedClass.inputArgs[input] = std::get<1>(*arg); - return std::get<1>(*arg); -} - -bool isOldComputeRegionBlockArgument(MaterializerState& state, Value value) { - auto blockArg = dyn_cast(value); - if (!blockArg) - return false; - - Operation* owner = blockArg.getOwner()->getParentOp(); - return owner && state.oldComputeOps.contains(owner); -} - -FailureOr appendScalarPublicationResult(MaterializerState& state, - MaterializedClass& materializedClass, - Value payload, - Location loc) { - auto existing = materializedClass.publicationOutputToResultIndex.find(payload); - if (existing != materializedClass.publicationOutputToResultIndex.end()) - return existing->second; - - auto compute = dyn_cast(materializedClass.op); - if (!compute) - return materializedClass.op->emitError("scalar publication result requires spat.scheduled_compute owner"); - - auto payloadType = dyn_cast(payload.getType()); - if (!payloadType || !payloadType.hasStaticShape()) - return materializedClass.op->emitError("scalar publication result requires static ranked tensor payload"); - - FailureOr> inserted = - compute.insertOutput(state.rewriter, compute.getNumResults(), payloadType, loc); - if (failed(inserted)) - return materializedClass.op->emitError("failed to append scalar publication result"); - - auto [result, newCompute] = *inserted; - materializedClass.op = newCompute.getOperation(); - materializedClass.body = &newCompute.getBody().front(); - materializedClass.publicationOutputToResultIndex[payload] = result.getResultNumber(); - - auto yieldOp = dyn_cast(materializedClass.body->getTerminator()); - if (!yieldOp) - return materializedClass.op->emitError("expected spat.yield terminator while appending scalar publication result"); - state.rewriter.modifyOpInPlace(yieldOp, [&] { yieldOp->insertOperands(yieldOp.getNumOperands(), payload); }); - return result.getResultNumber(); -} - -FailureOr appendBatchPublicationResult(MaterializerState& state, - MaterializedClass& materializedClass, - Value payload, - Location loc) { - auto existing = materializedClass.publicationOutputToResultIndex.find(payload); - if (existing != materializedClass.publicationOutputToResultIndex.end()) - return existing->second; - - auto batch = dyn_cast(materializedClass.op); - if (!batch) - return materializedClass.op->emitError("batch publication result requires spat.scheduled_compute_batch owner"); - - auto payloadType = dyn_cast(payload.getType()); - if (!payloadType || !payloadType.hasStaticShape() || payloadType.getRank() == 0) - return materializedClass.op->emitError( - "batch publication result requires a static ranked tensor payload with rank > 0"); - - SmallVector publishedShape(payloadType.getShape()); - publishedShape[0] *= static_cast(materializedClass.cpus.size()); - auto publishedType = RankedTensorType::get(publishedShape, payloadType.getElementType(), payloadType.getEncoding()); - - FailureOr> inserted = - batch.insertOutput(state.rewriter, batch.getNumResults(), publishedType, loc); - if (failed(inserted)) - return materializedClass.op->emitError("failed to append batch publication result"); - - auto [result, outputArg, newBatch] = *inserted; - materializedClass.op = newBatch.getOperation(); - materializedClass.body = &newBatch.getBody().front(); - materializedClass.publicationOutputToResultIndex[payload] = result.getResultNumber(); - - auto inParallelOp = dyn_cast(materializedClass.body->getTerminator()); - auto laneArg = newBatch.getLaneArgument(); - if (!laneArg) - return materializedClass.op->emitError("batch publication result requires a lane argument"); - if (!inParallelOp) { - auto yieldOp = dyn_cast(materializedClass.body->getTerminator()); - if (!yieldOp || yieldOp.getNumOperands() != 0) - return materializedClass.op->emitError( - "batch publication result requires either spat.in_parallel or an empty spat.yield terminator"); - state.rewriter.setInsertionPoint(yieldOp); - inParallelOp = SpatInParallelOp::create(state.rewriter, loc); - state.rewriter.eraseOp(yieldOp); - } - - state.rewriter.setInsertionPoint(inParallelOp); - Value firstOffset = scaleIndexByDim0Size(state, materializedClass.op, *laneArg, payloadType.getDimSize(0), loc); - state.rewriter.setInsertionPointToStart(&inParallelOp.getRegion().front()); - createDim0ParallelInsertSlice(state, loc, payload, outputArg, firstOffset); - return result.getResultNumber(); -} - -// ----------------------------------------------------------------------------- -// Materialized-class value localization helpers. -// ----------------------------------------------------------------------------- - -Operation* getEnclosingSpatialComputeLikeOp(Value value) { - Block* block = nullptr; - if (auto blockArg = dyn_cast(value)) - block = blockArg.getOwner(); - else if (Operation* definingOp = value.getDefiningOp()) - block = definingOp->getBlock(); - - if (!block) - return nullptr; - - for (Operation* current = block->getParentOp(); current; current = current->getParentOp()) - if (isa(current)) - return current; - return nullptr; -} - -static bool isValueDefinedInMaterializedClass(Value value, const MaterializedClass& targetClass) { - if (isConstantLike(value)) - return true; - - if (auto blockArg = dyn_cast(value)) { - for (Operation* current = blockArg.getOwner()->getParentOp(); current; current = current->getParentOp()) - if (current == targetClass.op) - return true; - return false; - } - - if (Operation* definingOp = value.getDefiningOp()) - for (Operation* current = definingOp; current; current = current->getParentOp()) - if (current == targetClass.op) - return true; - - return false; -} - -bool isTensorValueLocalToMaterializedClass(Value value, const MaterializedClass& targetClass) { - if (!isa(value.getType())) - return true; - return isValueDefinedInMaterializedClass(value, targetClass); -} - -bool isTensorValueDefinedInDifferentMaterializedClass(Value value, const MaterializedClass& targetClass) { - if (!isa(value.getType()) || isTensorValueLocalToMaterializedClass(value, targetClass)) - return false; - - Operation* owner = getEnclosingSpatialComputeLikeOp(value); - return owner && owner != targetClass.op; -} - -std::optional getRegionIndexInParentOp(Region* region) { - Operation* parent = region ? region->getParentOp() : nullptr; - if (!parent) - return std::nullopt; - - for (auto [index, candidate] : llvm::enumerate(parent->getRegions())) - if (&candidate == region) - return static_cast(index); - return std::nullopt; -} - -std::optional getBlockIndexInRegion(Block* block) { - Region* region = block ? block->getParent() : nullptr; - if (!region) - return std::nullopt; - - for (auto [index, candidate] : llvm::enumerate(region->getBlocks())) - if (&candidate == block) - return static_cast(index); - return std::nullopt; -} - -Block* getBlockByIndex(Region& region, unsigned blockIndex) { - unsigned index = 0; - for (Block& block : region) { - if (index == blockIndex) - return █ - ++index; - } - return nullptr; -} - -static bool isValueLegalInMaterializedClassBody(Value value, const MaterializedClass& targetClass) { - return isValueDefinedInMaterializedClass(value, targetClass); -} - -std::string stringifyOperationForMaterializerDebug(Operation* op) { - if (!op) - return std::string(""); - std::string storage; - llvm::raw_string_ostream stream(storage); - op->print(stream); - return storage; -} - -std::string stringifyValueForMaterializerDebug(Value value) { - std::string storage; - llvm::raw_string_ostream stream(storage); - value.print(stream); - return storage; -} - -std::string truncateMaterializerDebugString(std::string text, size_t limit = 1200) { - for (char& ch : text) - if (ch == '\n' || ch == '\r' || ch == '\t') - ch = ' '; - - if (text.size() <= limit) - return text; - text.resize(limit); - text += "..."; - return text; -} - -std::string formatMaterializerOperandListInline(Operation* op, const MaterializedClass& targetClass) { - if (!op) - return std::string(""); - - std::string storage; - llvm::raw_string_ostream stream(storage); - for (OpOperand& operand : op->getOpOperands()) { - if (operand.getOperandNumber() != 0) - stream << " | "; - Value value = operand.get(); - stream << "operand#" << operand.getOperandNumber() << " type=" << value.getType() - << " local=" << (isValueLegalInMaterializedClassBody(value, targetClass) ? 1 : 0) - << " value=" << stringifyValueForMaterializerDebug(value); - if (auto blockArg = dyn_cast(value)) { - stream << " blockArg#" << blockArg.getArgNumber(); - if (Operation* owner = blockArg.getOwner()->getParentOp()) - stream << " ownerOp='" << owner->getName() << "'"; - } - else if (Operation* definingOp = value.getDefiningOp()) { - stream << " definingOp='" << definingOp->getName() << "'"; - } - } - return truncateMaterializerDebugString(stream.str()); -} - -std::string formatMaterializerParentChainInline(Operation* op) { - if (!op) - return std::string(""); - - std::string storage; - llvm::raw_string_ostream stream(storage); - unsigned depth = 0; - for (Operation* current = op; current; current = current->getParentOp()) { - if (depth != 0) - stream << " <- "; - stream << "[" << depth++ << "]" << current->getName(); - } - return truncateMaterializerDebugString(stream.str()); -} - -void attachMaterializerOperationPrintNote(InFlightDiagnostic& diagnostic, Operation* op, StringRef label) { - if (!op) - return; - diagnostic.attachNote(op->getLoc()) << label << ":\n" << stringifyOperationForMaterializerDebug(op); -} - -void attachMaterializerParentChainNote(InFlightDiagnostic& diagnostic, Operation* op, StringRef label) { - if (!op) - return; - - std::string storage; - llvm::raw_string_ostream stream(storage); - unsigned depth = 0; - for (Operation* current = op; current; current = current->getParentOp()) - stream << " [" << depth++ << "] " << current->getName() << "\n"; - - diagnostic.attachNote(op->getLoc()) << label << ":\n" << stream.str(); -} - -void attachMaterializerOperandListNote(InFlightDiagnostic& diagnostic, - Operation* op, - const MaterializedClass& targetClass, - StringRef label) { - if (!op) - return; - - std::string storage; - llvm::raw_string_ostream stream(storage); - for (OpOperand& operand : op->getOpOperands()) { - Value value = operand.get(); - stream << " operand#" << operand.getOperandNumber() << " type=" << value.getType() - << " local=" << (isValueLegalInMaterializedClassBody(value, targetClass) ? 1 : 0) - << " value=" << stringifyValueForMaterializerDebug(value); - if (auto blockArg = dyn_cast(value)) { - stream << " blockArg#" << blockArg.getArgNumber(); - if (Operation* owner = blockArg.getOwner()->getParentOp()) - stream << " ownerOp='" << owner->getName() << "'"; - } - else if (Operation* definingOp = value.getDefiningOp()) { - stream << " definingOp='" << definingOp->getName() << "'"; - } - stream << "\n"; - } - - diagnostic.attachNote(op->getLoc()) << label << ":\n" << stream.str(); -} - -void attachMaterializerValueOriginNote(InFlightDiagnostic& diagnostic, Value value, StringRef label) { - if (auto blockArg = dyn_cast(value)) { - if (Operation* owner = blockArg.getOwner()->getParentOp()) - diagnostic.attachNote(owner->getLoc()) << label << " is block argument #" << blockArg.getArgNumber() << " of '" - << owner->getName() << "' with type " << blockArg.getType(); - else - diagnostic.attachNote(UnknownLoc::get(value.getContext())) - << label << " is a top-level block argument #" << blockArg.getArgNumber() << " with type " - << blockArg.getType(); - return; - } - - if (Operation* definingOp = value.getDefiningOp()) { - diagnostic.attachNote(definingOp->getLoc()) - << label << " is defined by '" << definingOp->getName() << "' with result type " << value.getType(); - return; - } - - diagnostic.attachNote(UnknownLoc::get(value.getContext())) - << label << " has no defining operation and is not a block argument, type " << value.getType(); -} - -void attachMaterializedClassBodySummary(InFlightDiagnostic& diagnostic, const MaterializedClass& targetClass) { - Block& body = *targetClass.body; - diagnostic.attachNote(targetClass.op->getLoc()) - << "target class " << targetClass.id << " op '" << targetClass.op->getName() << "' body has " - << body.getNumArguments() << " block arguments and " << std::distance(body.begin(), body.end()) - << " top-level operations"; -} - -FailureOr rematerializeIndexValueInClass( - MaterializerState& state, MaterializedClass& targetClass, Value value, Location loc, IRMapping* mapper = nullptr); - -FailureOr rematerializeIndexOpFoldResultInClass(MaterializerState& state, - MaterializedClass& targetClass, - OpFoldResult value, - Location loc, - IRMapping* mapper = nullptr) { - if (auto attr = dyn_cast(value)) - return OpFoldResult(attr); - - FailureOr rematerialized = rematerializeIndexValueInClass(state, targetClass, cast(value), loc, mapper); - if (failed(rematerialized)) - return failure(); - return OpFoldResult(*rematerialized); -} - -FailureOr rematerializeIndexValueInClass( - MaterializerState& state, MaterializedClass& targetClass, Value value, Location loc, IRMapping* mapper) { - Value originalValue = value; - bool mapperHadOriginalValue = false; - Value mappedOriginalValue; - - if (mapper && mapper->contains(value)) { - mapperHadOriginalValue = true; - Value mapped = mapper->lookup(value); - mappedOriginalValue = mapped; - if (isValueLegalInMaterializedClassBody(mapped, targetClass) || isConstantLike(mapped)) - return mapped; - value = mapped; - } - - if (isValueLegalInMaterializedClassBody(value, targetClass)) - return value; - - if (!value.getType().isIndex()) - return targetClass.op->emitError("cannot rematerialize non-index external value in materialized class body") - << " type=" << value.getType(); - - if (auto constantIndex = value.getDefiningOp()) - return getOrCreateIndexConstant(state.constantFolder, targetClass.op, constantIndex.value()); - - APInt constantValue; - if (matchPattern(value, m_ConstantInt(&constantValue))) { - if (!constantValue.isSignedIntN(64)) - return targetClass.op->emitError("cannot rematerialize out-of-range index constant") - << " value=" << llvm::toString(constantValue, 10, /*Signed=*/true); - return getOrCreateIndexConstant(state.constantFolder, targetClass.op, constantValue.getSExtValue()); - } - - if (auto affineApply = value.getDefiningOp()) { - SmallVector remappedOperands; - remappedOperands.reserve(affineApply.getMapOperands().size()); - for (Value operand : affineApply.getMapOperands()) { - FailureOr remapped = rematerializeIndexValueInClass(state, targetClass, operand, loc, mapper); - if (failed(remapped)) - return failure(); - remappedOperands.push_back(*remapped); - } - return createOrFoldAffineApply(state.rewriter, loc, affineApply.getAffineMap(), remappedOperands, state.func); - } - - if (auto addOp = value.getDefiningOp()) { - FailureOr lhs = rematerializeIndexValueInClass(state, targetClass, addOp.getLhs(), loc, mapper); - FailureOr rhs = rematerializeIndexValueInClass(state, targetClass, addOp.getRhs(), loc, mapper); - if (failed(lhs) || failed(rhs)) - return failure(); - return arith::AddIOp::create(state.rewriter, loc, *lhs, *rhs).getResult(); - } - - if (auto subOp = value.getDefiningOp()) { - FailureOr lhs = rematerializeIndexValueInClass(state, targetClass, subOp.getLhs(), loc, mapper); - FailureOr rhs = rematerializeIndexValueInClass(state, targetClass, subOp.getRhs(), loc, mapper); - if (failed(lhs) || failed(rhs)) - return failure(); - return arith::SubIOp::create(state.rewriter, loc, *lhs, *rhs).getResult(); - } - - if (auto mulOp = value.getDefiningOp()) { - FailureOr lhs = rematerializeIndexValueInClass(state, targetClass, mulOp.getLhs(), loc, mapper); - FailureOr rhs = rematerializeIndexValueInClass(state, targetClass, mulOp.getRhs(), loc, mapper); - if (failed(lhs) || failed(rhs)) - return failure(); - return arith::MulIOp::create(state.rewriter, loc, *lhs, *rhs).getResult(); - } - - if (auto divOp = value.getDefiningOp()) { - FailureOr lhs = rematerializeIndexValueInClass(state, targetClass, divOp.getLhs(), loc, mapper); - FailureOr rhs = rematerializeIndexValueInClass(state, targetClass, divOp.getRhs(), loc, mapper); - if (failed(lhs) || failed(rhs)) - return failure(); - return arith::DivUIOp::create(state.rewriter, loc, *lhs, *rhs).getResult(); - } - - if (auto extractOp = value.getDefiningOp()) { - SmallVector remappedIndices; - remappedIndices.reserve(extractOp.getIndices().size()); - for (Value index : extractOp.getIndices()) { - FailureOr remapped = rematerializeIndexValueInClass(state, targetClass, index, loc, mapper); - if (failed(remapped)) - return failure(); - remappedIndices.push_back(*remapped); - } - - Value tensor = extractOp.getTensor(); - if (!isConstantLike(tensor) && !isValueLegalInMaterializedClassBody(tensor, targetClass)) - return targetClass.op->emitError("cannot rematerialize indexed table lookup from external non-constant tensor") - << " tensorType=" << tensor.getType(); - return tensor::ExtractOp::create(state.rewriter, loc, tensor, remappedIndices).getResult(); - } - - if (auto blockArg = dyn_cast(value)) { - InFlightDiagnostic diagnostic = - targetClass.op->emitError("cannot rematerialize external block argument in materialized class body"); - diagnostic << " currentArg#" << blockArg.getArgNumber() << " currentType=" << blockArg.getType() - << " targetClass=" << targetClass.id << " targetOp='" << targetClass.op->getName() << "'"; - if (Operation* owner = blockArg.getOwner()->getParentOp()) { - diagnostic << " ownerOp='" << owner->getName() << "'"; - diagnostic << " ownerIR=\"" << truncateMaterializerDebugString(stringifyOperationForMaterializerDebug(owner)) - << "\""; - diagnostic << " ownerChain=\"" << formatMaterializerParentChainInline(owner) << "\""; - } - diagnostic << " targetIR=\"" - << truncateMaterializerDebugString(stringifyOperationForMaterializerDebug(targetClass.op)) << "\""; - if (mapper) { - diagnostic << " mapperPresent=1 mapperHadOriginal=" << (mapperHadOriginalValue ? 1 : 0); - if (mapperHadOriginalValue) - diagnostic << " mappedType=" << mappedOriginalValue.getType(); - } - else { - diagnostic << " mapperPresent=0"; - } - attachMaterializerValueOriginNote(diagnostic, originalValue, "original value"); - if (value != originalValue) - attachMaterializerValueOriginNote(diagnostic, value, "mapped/current value"); - if (mapperHadOriginalValue && mappedOriginalValue != value) - attachMaterializerValueOriginNote(diagnostic, mappedOriginalValue, "mapper value"); - if (Operation* owner = blockArg.getOwner()->getParentOp()) { - attachMaterializerOperationPrintNote(diagnostic, owner, "external block argument owner op"); - attachMaterializerParentChainNote(diagnostic, owner, "external block argument owner parent chain"); - } - attachMaterializerOperationPrintNote(diagnostic, targetClass.op, "target materialized op"); - attachMaterializedClassBodySummary(diagnostic, targetClass); - return failure(); - } - - InFlightDiagnostic diagnostic = - targetClass.op->emitError("cannot rematerialize external index value in materialized class body"); - diagnostic << " type=" << value.getType() << " targetClass=" << targetClass.id << " targetOp='" - << targetClass.op->getName() << "'"; - attachMaterializerValueOriginNote(diagnostic, originalValue, "original value"); - if (value != originalValue) - attachMaterializerValueOriginNote(diagnostic, value, "mapped/current value"); - attachMaterializedClassBodySummary(diagnostic, targetClass); - return failure(); -} - -InFlightDiagnostic emitNonLocalMaterializedClassValueDiagnostic(Operation* anchor, - const MaterializedClass& targetClass, - StringRef context, - Value value, - std::optional producer = std::nullopt) { - InFlightDiagnostic diagnostic = anchor->emitError(context) << " into target class " << targetClass.id; - - if (producer) { - diagnostic << " from '" << producer->instance.op->getName() << "' resultIndex=" << producer->resultIndex - << " laneStart=" << producer->instance.laneStart << " laneCount=" << producer->instance.laneCount; - } - else if (auto result = dyn_cast(value)) { - diagnostic << " from '" << result.getOwner()->getName() << "' resultIndex=" << result.getResultNumber(); - } - else if (auto blockArg = dyn_cast(value)) { - diagnostic << " from block argument #" << blockArg.getArgNumber(); - if (Operation* owner = blockArg.getOwner()->getParentOp()) - diagnostic << " of '" << owner->getName() << "'"; - } - - if (Operation* definingOp = value.getDefiningOp()) - diagnostic.attachNote(definingOp->getLoc()) << "offending tensor producer is '" << definingOp->getName() << "'"; - return diagnostic; -} - -FailureOr rematerializeTensorValueInClass(MaterializerState& state, - MaterializedClass& targetClass, - Value value, - Operation* anchor, - StringRef context, - IRMapping* mapper) { - auto extractSlice = value.getDefiningOp(); - if (extractSlice) { - FailureOr localizedSource = materializeTensorValueForMaterializedClassUse( - state, targetClass, extractSlice.getSource(), anchor, context, std::nullopt, mapper); - if (failed(localizedSource)) - return failure(); - - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - offsets.reserve(extractSlice.getMixedOffsets().size()); - sizes.reserve(extractSlice.getMixedSizes().size()); - strides.reserve(extractSlice.getMixedStrides().size()); - - for (OpFoldResult offset : extractSlice.getMixedOffsets()) { - FailureOr localized = - rematerializeIndexOpFoldResultInClass(state, targetClass, offset, anchor->getLoc(), mapper); - if (failed(localized)) - return failure(); - offsets.push_back(*localized); - } - for (OpFoldResult size : extractSlice.getMixedSizes()) { - FailureOr localized = - rematerializeIndexOpFoldResultInClass(state, targetClass, size, anchor->getLoc(), mapper); - if (failed(localized)) - return failure(); - sizes.push_back(*localized); - } - for (OpFoldResult stride : extractSlice.getMixedStrides()) { - FailureOr localized = - rematerializeIndexOpFoldResultInClass(state, targetClass, stride, anchor->getLoc(), mapper); - if (failed(localized)) - return failure(); - strides.push_back(*localized); - } - - auto resultType = dyn_cast(extractSlice.getResult().getType()); - if (!resultType) - return anchor->emitError("expected ranked tensor extract_slice while rematerializing tensor capture"); - - return extractStaticSliceOrIdentity( - state.rewriter, anchor->getLoc(), *localizedSource, resultType, offsets, sizes, strides); - } - - if (auto collapseShape = value.getDefiningOp()) { - FailureOr localizedSource = materializeTensorValueForMaterializedClassUse( - state, targetClass, collapseShape.getSrc(), anchor, context, std::nullopt, mapper); - if (failed(localizedSource)) - return failure(); - return tensor::CollapseShapeOp::create( - state.rewriter, anchor->getLoc(), *localizedSource, collapseShape.getReassociationIndices()) - .getResult(); - } - - if (auto empty = value.getDefiningOp()) { - auto emptyType = dyn_cast(empty.getType()); - if (!emptyType) - return anchor->emitError("expected ranked tensor.empty while rematerializing tensor capture"); - - SmallVector dynamicSizes; - dynamicSizes.reserve(empty.getDynamicSizes().size()); - for (Value dynamicSize : empty.getDynamicSizes()) { - FailureOr localizedSize = - rematerializeIndexValueInClass(state, targetClass, dynamicSize, anchor->getLoc(), mapper); - if (failed(localizedSize)) - return failure(); - dynamicSizes.push_back(*localizedSize); - } - return tensor::EmptyOp::create( - state.rewriter, anchor->getLoc(), emptyType.getShape(), emptyType.getElementType(), dynamicSizes) - .getResult(); - } - - return failure(); -} - -FailureOr materializeTensorValueForMaterializedClassUse(MaterializerState& state, - MaterializedClass& targetClass, - Value value, - Operation* anchor, - StringRef context, - std::optional producer, - IRMapping* mapper) { - if (mapper && mapper->contains(value)) - value = mapper->lookup(value); - - if (!isa(value.getType()) || isConstantLike(value) - || isTensorValueLocalToMaterializedClass(value, targetClass)) - return value; - - if (value.getDefiningOp() || value.getDefiningOp() - || value.getDefiningOp()) { - FailureOr rematerialized = - rematerializeTensorValueInClass(state, targetClass, value, anchor, context, mapper); - if (failed(rematerialized)) - return failure(); - return *rematerialized; - } - - if (isTensorValueDefinedInDifferentMaterializedClass(value, targetClass)) { - emitNonLocalMaterializedClassValueDiagnostic(anchor, targetClass, context, value, producer); - return failure(); - } - - if (isOldComputeRegionBlockArgument(state, value)) { - InFlightDiagnostic diagnostic = - anchor->emitError("cannot append old graph_compute region block argument as scheduled input"); - attachMaterializerValueOriginNote(diagnostic, value, "escaped input"); - return failure(); - } - - return appendInput(state, targetClass, value); -} - -MaterializedLayoutKind classifyMaterializedTensorLayout(RankedTensorType logicalType, RankedTensorType physicalType) { - if (logicalType == physicalType) - return MaterializedLayoutKind::DenseLogical; - - if (!logicalType || !physicalType || !logicalType.hasStaticShape() || !physicalType.hasStaticShape()) - return MaterializedLayoutKind::Unknown; - - if (logicalType.getRank() != 4 || physicalType.getRank() != 4) - return MaterializedLayoutKind::Unknown; - if (logicalType.getElementType() != physicalType.getElementType() - || logicalType.getEncoding() != physicalType.getEncoding()) - return MaterializedLayoutKind::Unknown; - - if (logicalType.getDimSize(0) != 1 || physicalType.getDimSize(2) != 1) - return MaterializedLayoutKind::Unknown; - if (logicalType.getDimSize(1) != physicalType.getDimSize(1) - || logicalType.getDimSize(2) != physicalType.getDimSize(0) - || logicalType.getDimSize(3) != physicalType.getDimSize(3)) - return MaterializedLayoutKind::Unknown; - - return MaterializedLayoutKind::RowPackedNCHWFromRows; -} - -FailureOr materializeLogicalTensorView(MaterializerState& state, - MaterializedClass& targetClass, - MaterializedTensorView view, - Operation* anchor, - StringRef context, - std::optional producer, - IRMapping* mapper) { - FailureOr localizedValue = materializeTensorValueForMaterializedClassUse( - state, targetClass, view.value, anchor, context, producer, mapper); - if (failed(localizedValue)) - return failure(); - - if (view.layout == MaterializedLayoutKind::DenseLogical) - return *localizedValue; - - if (view.layout == MaterializedLayoutKind::RowPackedNCHWFromRows) { - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value empty = tensor::EmptyOp::create(state.rewriter, - anchor->getLoc(), - view.logicalType.getShape(), - view.logicalType.getElementType()) - .getResult(); - auto permutation = DenseI64ArrayAttr::get(anchor->getContext(), {2, 1, 0, 3}); - auto transpose = - linalg::TransposeOp::create(state.rewriter, anchor->getLoc(), *localizedValue, empty, permutation); - return transpose.getResult().front(); - } - - InFlightDiagnostic diagnostic = anchor->emitError("materialized tensor replacement changed physical layout") - << " logicalType=" << view.logicalType << " physicalType=" << view.physicalType; - if (producer) { - diagnostic << " producerOp='" << producer->instance.op->getName() << "' laneStart=" << producer->instance.laneStart - << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex; - } - return failure(); -} - -std::optional mapExternalRegionBlockArgumentToLocalClone(const MaterializedClass& targetClass, - Operation* anchor, - BlockArgument externalArg) { - Block* sourceBlock = externalArg.getOwner(); - Region* sourceRegion = sourceBlock ? sourceBlock->getParent() : nullptr; - Operation* sourceParent = sourceRegion ? sourceRegion->getParentOp() : nullptr; - if (!sourceParent || !anchor) - return std::nullopt; - - std::optional sourceRegionIndex = getRegionIndexInParentOp(sourceRegion); - std::optional sourceBlockIndex = getBlockIndexInRegion(sourceBlock); - if (!sourceRegionIndex || !sourceBlockIndex) - return std::nullopt; - - for (Operation* current = anchor->getParentOp(); current && current != targetClass.op; - current = current->getParentOp()) { - if (current->getName() != sourceParent->getName()) - continue; - if (current->getNumRegions() <= *sourceRegionIndex) - continue; - - Region& localRegion = current->getRegion(*sourceRegionIndex); - Block* localBlock = getBlockByIndex(localRegion, *sourceBlockIndex); - if (!localBlock || localBlock->getNumArguments() <= externalArg.getArgNumber()) - continue; - - BlockArgument localArg = localBlock->getArgument(externalArg.getArgNumber()); - if (localArg.getType() != externalArg.getType()) - continue; - if (!isValueLegalInMaterializedClassBody(localArg, targetClass)) - continue; - return localArg; - } - - return std::nullopt; -} - -FailureOr localizeMaterializedClassOperand(MaterializerState& state, - MaterializedClass& targetClass, - Value value, - Operation* anchor, - StringRef tensorContext, - StringRef genericContext, - IRMapping* mapper) { - if (mapper && mapper->contains(value)) - value = mapper->lookup(value); - - if (auto blockArg = dyn_cast(value)) - if (std::optional localArg = mapExternalRegionBlockArgumentToLocalClone(targetClass, anchor, blockArg)) - return *localArg; - - if (isa(value.getType())) - return materializeTensorValueForMaterializedClassUse( - state, targetClass, value, anchor, tensorContext, std::nullopt, mapper); - - if (isValueLegalInMaterializedClassBody(value, targetClass)) - return value; - - if (value.getType().isIndex()) - return rematerializeIndexValueInClass(state, targetClass, value, anchor->getLoc(), mapper); - - InFlightDiagnostic diagnostic = anchor->emitError(genericContext); - diagnostic << " type=" << value.getType(); - if (auto blockArg = dyn_cast(value)) { - diagnostic << " blockArg#" << blockArg.getArgNumber(); - if (Operation* owner = blockArg.getOwner()->getParentOp()) - diagnostic.attachNote(owner->getLoc()) << "block argument belongs to '" << owner->getName() << "'"; - } - else if (Operation* definingOp = value.getDefiningOp()) { - diagnostic.attachNote(definingOp->getLoc()) - << "unsupported external operand producer is '" << definingOp->getName() << "'"; - } - return failure(); -} - -// ----------------------------------------------------------------------------- -// Tensor packing helpers. -// ----------------------------------------------------------------------------- - -struct Dim0SliceParams { - SmallVector offsets; - SmallVector sizes; - SmallVector strides; -}; - -Dim0SliceParams -buildDim0SliceParams(OpBuilder& builder, RankedTensorType referenceType, OpFoldResult firstOffset, int64_t firstSize) { - Dim0SliceParams params; - params.offsets.reserve(referenceType.getRank()); - params.sizes.reserve(referenceType.getRank()); - params.strides.reserve(referenceType.getRank()); - - params.offsets.push_back(firstOffset); - params.sizes.push_back(builder.getIndexAttr(firstSize)); - params.strides.push_back(builder.getIndexAttr(1)); - - for (int64_t dim = 1; dim < referenceType.getRank(); ++dim) { - params.offsets.push_back(builder.getIndexAttr(0)); - params.sizes.push_back(builder.getIndexAttr(referenceType.getDimSize(dim))); - params.strides.push_back(builder.getIndexAttr(1)); - } - - return params; -} - -Value createDim0ExtractSlice( - MaterializerState& state, Location loc, Value source, OpFoldResult firstOffset, int64_t firstSize) { - auto sourceType = cast(source.getType()); - Dim0SliceParams params = buildDim0SliceParams(state.rewriter, sourceType, firstOffset, firstSize); - return tensor::ExtractSliceOp::create(state.rewriter, loc, source, params.offsets, params.sizes, params.strides) - .getResult(); -} - -FailureOr createDim0ExtractSliceInClass(MaterializerState& state, - MaterializedClass& targetClass, - Location loc, - Value source, - OpFoldResult firstOffset, - int64_t firstSize) { - FailureOr localizedSource = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - source, - targetClass.op, - "createDim0ExtractSliceInClass tried to reuse a tensor from another materialized class"); - if (failed(localizedSource)) - return failure(); - FailureOr localizedOffset = rematerializeIndexOpFoldResultInClass(state, targetClass, firstOffset, loc); - if (failed(localizedOffset)) - return failure(); - return createDim0ExtractSlice(state, loc, *localizedSource, *localizedOffset, firstSize); -} - -Value createIndexedIndexValue(MaterializerState& state, - Operation* anchor, - ArrayRef values, - Value index, - Location loc, - std::optional preferredPeriod = std::nullopt, - bool allowExhaustiveTiledSearch = true); - -Value createDim0InsertSlice( - MaterializerState& state, Location loc, Value fragment, Value destination, OpFoldResult firstOffset) { - auto fragmentType = cast(fragment.getType()); - Dim0SliceParams params = buildDim0SliceParams(state.rewriter, fragmentType, firstOffset, fragmentType.getDimSize(0)); - return tensor::InsertSliceOp::create( - state.rewriter, loc, fragment, destination, params.offsets, params.sizes, params.strides) - .getResult(); -} - -FailureOr createDim0InsertSliceInClass(MaterializerState& state, - MaterializedClass& targetClass, - Location loc, - Value fragment, - Value destination, - OpFoldResult firstOffset) { - FailureOr localizedFragment = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - fragment, - targetClass.op, - "createDim0InsertSliceInClass tried to reuse a fragment tensor from another materialized class"); - if (failed(localizedFragment)) - return failure(); - FailureOr localizedDestination = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - destination, - targetClass.op, - "createDim0InsertSliceInClass tried to reuse a destination tensor from another materialized class"); - if (failed(localizedDestination)) - return failure(); - FailureOr localizedOffset = rematerializeIndexOpFoldResultInClass(state, targetClass, firstOffset, loc); - if (failed(localizedOffset)) - return failure(); - return createDim0InsertSlice(state, loc, *localizedFragment, *localizedDestination, *localizedOffset); -} - -void createDim0ParallelInsertSlice( - MaterializerState& state, Location loc, Value fragment, Value destination, OpFoldResult firstOffset) { - auto fragmentType = cast(fragment.getType()); - Dim0SliceParams params = buildDim0SliceParams(state.rewriter, fragmentType, firstOffset, fragmentType.getDimSize(0)); - tensor::ParallelInsertSliceOp::create( - state.rewriter, loc, fragment, destination, params.offsets, params.sizes, params.strides); -} - -Value scaleIndexByDim0Size(MaterializerState& state, Operation* anchor, Value index, int64_t dim0Size, Location loc) { - if (dim0Size == 1) - return index; - - MLIRContext* context = state.func.getContext(); - AffineExpr d0 = getAffineDimExpr(0, context); - AffineMap map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, d0 * dim0Size); - return createOrFoldAffineApply(state.rewriter, loc, map, ValueRange {index}, anchor); -} - -FailureOr scaleIndexByDim0SizeInClass( - MaterializerState& state, MaterializedClass& targetClass, Value index, int64_t dim0Size, Location loc) { - FailureOr localizedIndex = rematerializeIndexValueInClass(state, targetClass, index, loc); - if (failed(localizedIndex)) - return failure(); - if (dim0Size == 1) - return *localizedIndex; - - return scaleIndexByDim0Size(state, targetClass.op, *localizedIndex, dim0Size, loc); -} - -bool sameProducerResult(ProducerKey lhs, ProducerKey rhs) { - return lhs.instance.op == rhs.instance.op && lhs.resultIndex == rhs.resultIndex; -} - -bool containsProducerKey(ProducerKey outer, ProducerKey inner) { - if (!sameProducerResult(outer, inner)) - return false; - if (!isa(outer.instance.op)) - return false; - if (outer.instance.laneCount == 0 || inner.instance.laneCount == 0) - return false; - - uint32_t outerStart = outer.instance.laneStart; - uint32_t outerEnd = outerStart + outer.instance.laneCount; - uint32_t innerStart = inner.instance.laneStart; - uint32_t innerEnd = innerStart + inner.instance.laneCount; - - return outerStart <= innerStart && innerEnd <= outerEnd; -} - -std::optional extractPackedProducerSlice(MaterializerState& state, - MaterializedClass& materializedClass, - ProducerKey packedKey, - Value packed, - ProducerKey requestedKey) { - if (!containsProducerKey(packedKey, requestedKey)) - return std::nullopt; - - auto packedType = dyn_cast(packed.getType()); - if (!packedType || !packedType.hasStaticShape() || packedType.getRank() == 0) - return std::nullopt; - - if (packedKey.instance.laneCount == 0) - return std::nullopt; - - int64_t packedRows = packedType.getDimSize(0); - if (packedRows % static_cast(packedKey.instance.laneCount) != 0) - return std::nullopt; - - int64_t rowsPerLane = packedRows / static_cast(packedKey.instance.laneCount); - int64_t rowOffset = - static_cast(requestedKey.instance.laneStart - packedKey.instance.laneStart) * rowsPerLane; - int64_t rowCount = static_cast(requestedKey.instance.laneCount) * rowsPerLane; - - state.rewriter.setInsertionPoint(materializedClass.body->getTerminator()); - - Value firstOffset = getOrCreateIndexConstant(state.constantFolder, materializedClass.op, rowOffset); - return createDim0ExtractSlice(state, materializedClass.op->getLoc(), packed, firstOffset, rowCount); -} - -} // namespace - -std::optional AvailableValueStore::lookupExact(ProducerKey key, ClassId classId) const { - auto producerIt = exactValues.find(key); - if (producerIt == exactValues.end()) - return std::nullopt; - - auto valueIt = producerIt->second.find(classId); - if (valueIt == producerIt->second.end()) - return std::nullopt; - - return valueIt->second; -} - -namespace { - -using IndexedFragmentBuilder = llvm::function_ref(Value flatIndex)>; -using IndexedInsertOffsetBuilder = llvm::function_ref(Value flatIndex)>; - -SmallVector flattenPackedScalarRunKeys(const PackedScalarRunValue& run); -FailureOr emitIndexedFragmentInsertLoop(MaterializerState& state, - MaterializedClass& targetClass, - Value destination, - int64_t itemCount, - IndexedFragmentBuilder buildFragment, - IndexedInsertOffsetBuilder buildOffset, - Location loc); -FailureOr> cloneBatchBodyForLane(MaterializerState& state, - MaterializedClass& targetClass, - const ComputeInstance& instance, - Value laneValue, - ArrayRef resultIndices, - CloneIndexingContext indexing); -Value createIndexedChannelId( - MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc); -Value createIndexedSourceCoreId( - MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc); -Value createIndexedTargetCoreId( - MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc); -Value appendReceive( - MaterializerState& state, MaterializedClass& targetClass, Type type, const MessageVector& messages, Location loc); -FailureOr buildMessageVectorForExchangeIds(MaterializerState& state, - ArrayRef ids, - CommunicationExchangeKind expectedKind, - ClassId expectedSourceClass, - ClassId expectedTargetClass, - Type expectedPayloadType, - Operation* anchor); -LogicalResult bindCommunicationPayload(MaterializerState& state, - CommunicationExchangeId id, - MaterializedClass& sourceClass, - Value payload, - Operation* anchor, - Location loc); -FailureOr emitPlannedReceive(MaterializerState& state, - CommunicationExchangeId id, - MaterializedClass& targetClass, - Location loc); - -Value getPackedSliceForRunIndex(MaterializerState& state, - Operation* anchor, - Value packed, - RankedTensorType fragmentType, - size_t index, - Location loc) { - int64_t rowOffset = static_cast(index) * fragmentType.getDimSize(0); - Value firstOffset = getOrCreateIndexConstant(state.constantFolder, anchor, rowOffset); - return createDim0ExtractSlice(state, loc, packed, firstOffset, fragmentType.getDimSize(0)); -} - -Value getPackedSliceForDynamicRunIndex( - MaterializerState& state, Operation* anchor, Value packed, RankedTensorType fragmentType, Value index, Location loc) { - Value firstOffset = scaleIndexByDim0Size(state, anchor, index, fragmentType.getDimSize(0), loc); - return createDim0ExtractSlice(state, loc, packed, firstOffset, fragmentType.getDimSize(0)); -} - -FailureOr materializeIndexedBatchRunPackedValue(MaterializerState& state, - MaterializedClass& targetClass, - IndexedBatchRunValue& run, - ProducerKey key, - Location loc) { - if (!run.packed) - return failure(); - - size_t flattenedIndexBase = 0; - for (const PackedScalarRunSlot& slot : run.slots) { - std::optional contiguousKey = getContiguousProducerRangeForKeys(slot.keys); - auto keyIt = llvm::find(slot.keys, key); - if ((!contiguousKey && keyIt == slot.keys.end()) || (contiguousKey && !containsProducerKey(*contiguousKey, key))) { - flattenedIndexBase += slot.keys.size(); - continue; - } - - FailureOr packed = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - run.packed, - targetClass.op, - "indexed batch run packed resolution tried to reuse a tensor from another materialized class"); - if (failed(packed)) - return failure(); - - if (!contiguousKey) { - size_t globalFragmentIndex = flattenedIndexBase + static_cast(std::distance(slot.keys.begin(), keyIt)); - return getPackedSliceForRunIndex(state, targetClass.op, *packed, run.fragmentType, globalFragmentIndex, loc); - } - - FailureOr slotPackedType = getPackedBatchTensorType(run.fragmentType, slot.keys.size()); - if (failed(slotPackedType)) - return failure(); - - int64_t rowOffset = static_cast(flattenedIndexBase) * run.fragmentType.getDimSize(0); - Value firstOffset = getOrCreateIndexConstant(state.constantFolder, targetClass.op, rowOffset); - Value slotPacked = - createDim0ExtractSlice(state, loc, *packed, firstOffset, (*slotPackedType).getDimSize(0)); - - if (*contiguousKey == key) - return slotPacked; - - std::optional sliced = extractPackedProducerSlice(state, targetClass, *contiguousKey, slotPacked, key); - if (!sliced) - return failure(); - return *sliced; - } - - return failure(); -} - -using IndexedFragmentBuilder = llvm::function_ref(Value flatIndex)>; -using IndexedInsertOffsetBuilder = llvm::function_ref(Value flatIndex)>; - -struct ReceiveMessagePartition { - SmallVector criticalStaticIndices; - SmallVector remainingIndices; -}; - -std::optional getConstantIndexValue(Value value); - -ReceiveMessagePartition partitionReceiveMessagesInPlanOrder(const MessageVector& messages); -MessageVector filterMessageVector(const MessageVector& messages, ArrayRef indices); - -bool isDeferredLocalPackedScalarRun(const PackedScalarRunValue& run) { - return run.kind == PackedScalarRunKind::DeferredLocalCompute; -} - -size_t getPackedScalarRunReceiveCount(const PackedScalarRunValue& run) { - size_t count = 0; - for (const PackedScalarRunSlot& slot : run.slots) - count += slot.keys.size(); - return count; -} - -LogicalResult validatePackedScalarRunMetadata(Operation* anchor, const PackedScalarRunValue& run) { - if (run.kind == PackedScalarRunKind::DeferredLocalCompute) - return success(); - - size_t receiveCount = getPackedScalarRunReceiveCount(run); - - if (receiveCount == 0) - return anchor->emitError("packed scalar run has no receives"); - - if (!run.exchangeIds.empty()) { - if (run.exchangeIds.size() != receiveCount) - return anchor->emitError("packed scalar run exchange id count is inconsistent"); - return success(); - } - - return anchor->emitError("packed scalar run deferred receive is missing planned exchange ids"); -} - -FailureOr materializePackedScalarRunValue(MaterializerState& state, - MaterializedClass& targetClass, - PackedScalarRunValue& run, - Location loc) { - if (run.packed) - return run.packed; - - if (run.kind == PackedScalarRunKind::Materialized) - return targetClass.op->emitError("materialized packed scalar run has no packed value"); - - if (isDeferredLocalPackedScalarRun(run)) { - SmallVector keys = flattenPackedScalarRunKeys(run); - if (keys.empty()) - return failure(); - FailureOr packedType = getPackedBatchTensorType(run.fragmentType, keys.size()); - if (failed(packedType)) - return targetClass.op->emitError("cannot materialize deferred local packed run for non-static ranked tensor"); - - SmallVector sourceLanes; - sourceLanes.reserve(keys.size()); - for (ProducerKey key : keys) { - if (key.instance.laneCount != 1) - return failure(); - sourceLanes.push_back(key.instance.laneStart); - } - - SmallVector resultIndices {run.resultIndex}; - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value init = - tensor::EmptyOp::create(state.rewriter, loc, packedType->getShape(), packedType->getElementType()).getResult(); - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(keys.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1); - - auto loop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {init}, - [&](OpBuilder&, Location, Value loopIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { - Value acc = iterArgs.front(); - Value sourceLane = createIndexedIndexValue(state, targetClass.op, sourceLanes, loopIndex, loc); - - FailureOr> produced = - cloneBatchBodyForLane(state, - targetClass, - keys.front().instance, - sourceLane, - resultIndices, - CloneIndexingContext {.runSlotIndex = std::nullopt, .projectionSlotIndex = loopIndex}); - if (failed(produced) || produced->size() != 1) - return failure(); - - FailureOr firstOffset = - scaleIndexByDim0SizeInClass(state, targetClass, loopIndex, run.fragmentType.getDimSize(0), loc); - if (failed(firstOffset)) - return failure(); - FailureOr next = - createDim0InsertSliceInClass(state, targetClass, loc, produced->front(), acc, *firstOffset); - if (failed(next)) - return failure(); - yielded.push_back(*next); - return success(); - }); - if (failed(loop)) - return failure(); - run.packed = loop->results.front(); - return run.packed; - } - - if (failed(validatePackedScalarRunMetadata(targetClass.op, run))) - return failure(); - - FailureOr builtMessages = - buildMessageVectorForExchangeIds(state, - run.exchangeIds, - CommunicationExchangeKind::PackedScalarRun, - state.communicationPlan.getExchange(run.exchangeIds.front()).sourceClass, - targetClass.id, - run.fragmentType, - targetClass.op); - if (failed(builtMessages)) - return failure(); - const MessageVector& receiveMessages = *builtMessages; - - FailureOr fullPackedType = - getPackedBatchTensorType(run.fragmentType, getPackedScalarRunReceiveCount(run)); - if (failed(fullPackedType)) - return targetClass.op->emitError("cannot create lazy packed scalar run receive type"); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value init = - tensor::EmptyOp::create(state.rewriter, loc, fullPackedType->getShape(), fullPackedType->getElementType()) - .getResult(); - Value packedInit = init; - ReceiveMessagePartition partition = - partitionReceiveMessagesInPlanOrder(receiveMessages); - for (size_t index : partition.criticalStaticIndices) { - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value received = SpatChannelReceiveOp::create( - state.rewriter, - loc, - run.fragmentType, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, receiveMessages.channelIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, receiveMessages.sourceCoreIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, receiveMessages.targetCoreIds[index])) - .getOutput(); - FailureOr offset = scaleIndexByDim0SizeInClass( - state, - targetClass, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(index)), - run.fragmentType.getDimSize(0), - loc); - if (failed(offset)) - return failure(); - FailureOr updated = createDim0InsertSliceInClass(state, targetClass, loc, received, packedInit, *offset); - if (failed(updated)) - return failure(); - packedInit = *updated; - } - - auto remainingMessages = filterMessageVector(receiveMessages, partition.remainingIndices); - SmallVector remainingPackedRowOffsets; - remainingPackedRowOffsets.reserve(partition.remainingIndices.size()); - for (size_t originalIndex : partition.remainingIndices) - remainingPackedRowOffsets.push_back(static_cast(originalIndex) * run.fragmentType.getDimSize(0)); - auto packed = remainingMessages.empty() - ? FailureOr(packedInit) - : emitIndexedFragmentInsertLoop( - state, - targetClass, - packedInit, - static_cast(remainingMessages.size()), - [&](Value index) -> FailureOr { - Value channelId = createIndexedChannelId(state, targetClass.op, remainingMessages, index, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingMessages, index, loc); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingMessages, index, loc); - return SpatChannelReceiveOp::create(state.rewriter, loc, run.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); - }, - [&](Value index) -> FailureOr { - return createIndexedIndexValue(state, targetClass.op, remainingPackedRowOffsets, index, loc); - }, - loc); - if (failed(packed)) - return failure(); - run.packed = *packed; - return run.packed; -} - -} // namespace - -std::optional AvailableValueStore::lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId) { - for (PackedScalarRunValue& run : packedScalarRuns) { - if (run.targetClass != classId || run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex) - continue; - - size_t flattenedIndexBase = 0; - for (auto [slotIndex, slot] : llvm::enumerate(run.slots)) { - std::optional contiguousKey = getContiguousProducerRangeForKeys(slot.keys); - auto keyIt = llvm::find(slot.keys, key); - if ((!contiguousKey && keyIt == slot.keys.end()) || (contiguousKey && !containsProducerKey(*contiguousKey, key))) - { - flattenedIndexBase += slot.keys.size(); - continue; - } - - MaterializedClass& materializedClass = state.classes[classId]; - state.rewriter.setInsertionPoint(materializedClass.body->getTerminator()); - - FailureOr packed = - materializePackedScalarRunValue(state, materializedClass, run, materializedClass.op->getLoc()); - if (failed(packed)) - return std::nullopt; - - if (!contiguousKey) { - size_t globalFragmentIndex = flattenedIndexBase + static_cast(std::distance(slot.keys.begin(), keyIt)); - Value sliced = getPackedSliceForRunIndex(state, - materializedClass.op, - *packed, - run.fragmentType, - globalFragmentIndex, - (*packed).getLoc()); - record(key, classId, sliced); - return sliced; - } - - FailureOr slotPackedType = getPackedBatchTensorType(run.fragmentType, slot.keys.size()); - if (failed(slotPackedType)) - return std::nullopt; - - int64_t rowOffset = static_cast(flattenedIndexBase) * run.fragmentType.getDimSize(0); - Value firstOffset = getOrCreateIndexConstant(state.constantFolder, materializedClass.op, rowOffset); - Value slotPacked = - createDim0ExtractSlice(state, (*packed).getLoc(), *packed, firstOffset, (*slotPackedType).getDimSize(0)); - - if (*contiguousKey == key) { - record(key, classId, slotPacked); - return slotPacked; - } - - std::optional sliced = - extractPackedProducerSlice(state, materializedClass, *contiguousKey, slotPacked, key); - if (!sliced) - return std::nullopt; - - record(key, classId, *sliced); - return *sliced; - } - } - - return std::nullopt; -} - -IndexedBatchRunValue* AvailableValueStore::lookupIndexedBatchRun(ProducerKey key, ClassId classId) { - for (IndexedBatchRunValue& run : indexedBatchRuns) { - if (run.targetClass != classId || run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex) - continue; - for (const PackedScalarRunSlot& slot : run.slots) { - if (!llvm::is_contained(slot.keys, key)) - continue; - return &run; - } - } - return nullptr; -} - -std::optional AvailableValueStore::lookup(MaterializerState& state, ProducerKey key, ClassId classId) { - - if (std::optional exact = lookupExact(key, classId)) - return exact; - - auto exchangeProducerIt = exactExchangeIds.find(key); - if (exchangeProducerIt != exactExchangeIds.end()) { - auto exchangeClassIt = exchangeProducerIt->second.find(classId); - if (exchangeClassIt != exchangeProducerIt->second.end() && !exchangeClassIt->second.empty()) { - ArrayRef ids = exchangeClassIt->second; - const CommunicationExchange& firstExchange = state.communicationPlan.getExchange(ids.front()); - bool ready = true; - for (CommunicationExchangeId id : ids) { - if (id.value >= state.communicationExchangeStates.size() - || !state.communicationExchangeStates[id.value].sendEmitted) { - ready = false; - break; - } - } - if (!ready) - return std::nullopt; - - MaterializedClass& materializedClass = state.classes[classId]; - FailureOr messages = buildMessageVectorForExchangeIds(state, - ids, - firstExchange.kind, - firstExchange.sourceClass, - firstExchange.targetClass, - Type(), - materializedClass.op); - if (failed(messages)) - return std::nullopt; - Value received = - appendReceive(state, materializedClass, firstExchange.payloadType, *messages, materializedClass.op->getLoc()); - if (!received) - return std::nullopt; - - for (CommunicationExchangeId id : ids) { - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; - emission.receiveEmitted = true; - emission.receiveValue = received; - if (Operation* receiveDef = received.getDefiningOp()) - emission.receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, materializedClass); - } - - record(key, classId, received); - exchangeProducerIt->second.erase(exchangeClassIt); - return received; - } - } - - if (std::optional packedRunValue = lookupPackedRun(state, key, classId)) - return packedRunValue; - - MaterializedClass& materializedClass = state.classes[classId]; - - for (const auto& [candidateKey, classValues] : exactValues) { - if (!sameProducerResult(candidateKey, key) || !containsProducerKey(candidateKey, key)) - continue; - - auto valueIt = classValues.find(classId); - if (valueIt == classValues.end()) - continue; - - std::optional slice = - extractPackedProducerSlice(state, materializedClass, candidateKey, valueIt->second, key); - if (!slice) - return std::nullopt; - - record(key, classId, *slice); - return *slice; - } - return std::nullopt; -} - -namespace { - -Value createIndexTensorConstant(MaterializerState& state, Operation* anchor, ArrayRef values) { - SmallVector elements; - elements.reserve(values.size()); - for (int64_t value : values) - elements.push_back(APInt(64, value)); - - auto type = RankedTensorType::get({static_cast(values.size())}, state.rewriter.getIndexType()); - auto attr = DenseIntElementsAttr::get(type, elements); - return getOrCreateConstant(state.constantFolder, anchor, attr, type); -} - -bool allEqual(ArrayRef values) { - assert(!values.empty() && "expected at least one value"); - for (int64_t value : values.drop_front()) - if (value != values.front()) - return false; - return true; -} - -struct IndexedIndexPattern { - int64_t base = 0; - int64_t step = 0; - int64_t period = 1; - int64_t innerStep = 0; - int64_t outerStep = 0; - bool isTiled = false; -}; - -bool matchAffineSequence(ArrayRef values, IndexedIndexPattern& pattern) { - assert(!values.empty() && "expected at least one value"); - - pattern.base = values.front(); - pattern.step = values.size() == 1 ? 0 : values[1] - values[0]; - pattern.isTiled = false; - - for (auto [index, value] : llvm::enumerate(values)) { - int64_t expected = pattern.base + pattern.step * static_cast(index); - if (value != expected) - return false; - } - - return true; -} - -bool matchTiledAffineSequence(ArrayRef values, IndexedIndexPattern& pattern, int64_t period) { - assert(!values.empty() && "expected at least one value"); - if (period < 2 || period > static_cast(values.size() / 2)) - return false; - - int64_t base = values.front(); - int64_t innerStep = values[1] - values[0]; - int64_t outerStep = values[period] - values[0]; - - for (auto [index, value] : llvm::enumerate(values)) { - int64_t i = static_cast(index); - int64_t expected = base + outerStep * (i / period) + innerStep * (i % period); - if (value != expected) - return false; - } - - pattern.base = base; - pattern.period = period; - pattern.innerStep = innerStep; - pattern.outerStep = outerStep; - pattern.isTiled = true; - return true; -} - -bool matchTiledAffineSequence(ArrayRef values, IndexedIndexPattern& pattern) { - assert(!values.empty() && "expected at least one value"); - - for (int64_t period = 2; period <= static_cast(values.size() / 2); ++period) - if (matchTiledAffineSequence(values, pattern, period)) - return true; - - return false; -} - -std::optional getIndexedIndexPattern(ArrayRef values, - std::optional preferredPeriod = std::nullopt, - bool allowExhaustiveTiledSearch = true) { - assert(!values.empty() && "expected at least one value"); - - IndexedIndexPattern pattern; - if (matchAffineSequence(values, pattern)) - return pattern; - if (preferredPeriod && matchTiledAffineSequence(values, pattern, *preferredPeriod)) - return pattern; - if (allowExhaustiveTiledSearch && values.size() <= 256 && matchTiledAffineSequence(values, pattern)) - return pattern; - - return std::nullopt; -} - -Value createAffineIndexValue(MaterializerState& state, const IndexedIndexPattern& pattern, Value index, Location loc) { - MLIRContext* context = state.func.getContext(); - AffineExpr d0 = getAffineDimExpr(0, context); - - AffineExpr expr; - if (!pattern.isTiled) { - expr = getAffineConstantExpr(pattern.base, context) + d0 * pattern.step; - } - else { - expr = getAffineConstantExpr(pattern.base, context) + d0.floorDiv(pattern.period) * pattern.outerStep - + (d0 % pattern.period) * pattern.innerStep; - } - - AffineMap map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr); - return createOrFoldAffineApply(state.rewriter, loc, map, ValueRange {index}, state.func); -} - -Value createIndexedIndexValue(MaterializerState& state, - Operation* anchor, - ArrayRef values, - Value index, - Location loc, - std::optional preferredPeriod, - bool allowExhaustiveTiledSearch) { - assert(!values.empty() && "expected at least one indexed value"); - - if (allEqual(values)) - return getOrCreateIndexConstant(state.constantFolder, anchor, values.front()); - - if (std::optional pattern = - getIndexedIndexPattern(values, preferredPeriod, allowExhaustiveTiledSearch)) - return createAffineIndexValue(state, *pattern, index, loc); - Value table = createIndexTensorConstant(state, anchor, values); - return tensor::ExtractOp::create(state.rewriter, loc, table, ValueRange {index}).getResult(); -} - -Value createIndexedIndexValue( - MaterializerState& state, Operation* anchor, ArrayRef values, Value index, Location loc) { - assert(!values.empty() && "expected at least one indexed value"); - - SmallVector widened; - widened.reserve(values.size()); - for (int32_t value : values) - widened.push_back(value); - - return createIndexedIndexValue(state, anchor, ArrayRef(widened), index, loc, std::nullopt, true); -} - -OpFoldResult createIndexedOrStaticIndex( - MaterializerState& state, Operation* anchor, ArrayRef values, Value index, Location loc) { - assert(!values.empty() && "expected at least one indexed value"); - if (allEqual(values)) - return state.rewriter.getIndexAttr(values.front()); - return createIndexedIndexValue(state, anchor, values, index, loc); -} - -Value createIndexedChannelId( - MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc) { - return createIndexedIndexValue(state, anchor, ArrayRef(messages.channelIds), index, loc); -} - -Value createIndexedChannelId(MaterializerState& state, - Operation* anchor, - const MessageVector& messages, - Value index, - Location loc, - std::optional preferredPeriod) { - return createIndexedIndexValue( - state, anchor, ArrayRef(messages.channelIds), index, loc, preferredPeriod, true); -} - -Value createIndexedSourceCoreId( - MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc) { - return createIndexedIndexValue(state, anchor, ArrayRef(messages.sourceCoreIds), index, loc); -} - -Value createIndexedSourceCoreId(MaterializerState& state, - Operation* anchor, - const MessageVector& messages, - Value index, - Location loc, - std::optional preferredPeriod) { - SmallVector widened(messages.sourceCoreIds.begin(), messages.sourceCoreIds.end()); - return createIndexedIndexValue(state, anchor, ArrayRef(widened), index, loc, preferredPeriod, true); -} - -Value createIndexedTargetCoreId( - MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc) { - return createIndexedIndexValue(state, anchor, ArrayRef(messages.targetCoreIds), index, loc); -} - -Value createIndexedTargetCoreId(MaterializerState& state, - Operation* anchor, - const MessageVector& messages, - Value index, - Location loc, - std::optional preferredPeriod) { - SmallVector widened(messages.targetCoreIds.begin(), messages.targetCoreIds.end()); - return createIndexedIndexValue(state, anchor, ArrayRef(widened), index, loc, preferredPeriod, true); -} - -Value createLaneIndexedIndexValue(MaterializerState& state, - MaterializedClass& materializedClass, - ArrayRef values, - Location loc) { - assert(materializedClass.isBatch && "lane-indexed value requires a materialized batch class"); - assert(values.size() == materializedClass.cpus.size() && "expected one value per materialized batch lane"); - - auto batch = cast(materializedClass.op); - auto laneArg = batch.getLaneArgument(); - assert(laneArg && "expected compute_batch lane argument"); - - return createIndexedIndexValue(state, materializedClass.op, values, *laneArg, loc); -} - -Value createLaneIndexedIndexValue(MaterializerState& state, - MaterializedClass& materializedClass, - ArrayRef values, - Location loc) { - assert(materializedClass.isBatch && "lane-indexed value requires a materialized batch class"); - assert(values.size() == materializedClass.cpus.size() && "expected one value per materialized batch lane"); - - SmallVector widened; - widened.reserve(values.size()); - for (int32_t value : values) - widened.push_back(value); - - return createLaneIndexedIndexValue(state, materializedClass, ArrayRef(widened), loc); -} - -FailureOr> -getPeerLogicalInstances(MaterializerState& state, const MaterializedClass& materializedClass, SlotId logicalSlot) { - SmallVector peers; - peers.reserve(materializedClass.cpus.size()); - for (CpuId cpu : materializedClass.cpus) { - auto streamIt = state.logicalInstancesByCpu.find(cpu); - if (streamIt == state.logicalInstancesByCpu.end() || logicalSlot >= streamIt->second.size()) - return failure(); - peers.push_back(streamIt->second[logicalSlot]); - } - return peers; -} - -Value createOriginalLaneValue(MaterializerState& state, - MaterializedClass& materializedClass, - ArrayRef peers, - Location loc) { - assert(!peers.empty() && "expected at least one peer instance"); - if (!materializedClass.isBatch) - return getOrCreateIndexConstant(state.constantFolder, materializedClass.op, peers.front().laneStart); - - auto batch = cast(materializedClass.op); - auto laneArg = batch.getLaneArgument(); - assert(laneArg && "expected materialized compute_batch lane argument"); - - SmallVector laneValues; - laneValues.reserve(peers.size()); - for (const ComputeInstance& peer : peers) - laneValues.push_back(peer.laneStart); - - return createIndexedIndexValue(state, materializedClass.op, ArrayRef(laneValues), *laneArg, loc); -} - -bool hasLiveExternalUse(Value value, const DenseSet& oldComputeOps) { - SmallVector worklist {value}; - DenseSet visited; - - while (!worklist.empty()) { - Value current = worklist.pop_back_val(); - if (!visited.insert(current).second) - continue; - - for (OpOperand& use : current.getUses()) { - Operation* owner = use.getOwner(); - if (isInsideOldCompute(owner, oldComputeOps)) - continue; - if (isa(owner)) { - for (Value result : owner->getResults()) - worklist.push_back(result); - continue; - } - return true; - } - } - - return false; -} - -bool hasRealComputeConsumer(Value value, const DenseSet& oldComputeOps) { - SmallVector worklist {value}; - DenseSet visited; - - while (!worklist.empty()) { - Value current = worklist.pop_back_val(); - if (!visited.insert(current).second) - continue; - - for (OpOperand& use : current.getUses()) { - Operation* owner = use.getOwner(); - if (isInsideOldCompute(owner, oldComputeOps)) - continue; - if (isa(owner)) { - for (Value result : owner->getResults()) - worklist.push_back(result); - continue; - } - if (isa(owner)) - continue; - return true; - } - } - - return false; -} - -FailureOr getBatchResultProjectionInsert(SpatComputeBatch batch, size_t resultIndex); - -bool isTerminalHostBatchOutput(Value output, const DenseSet& oldComputeOps) { - auto batch = dyn_cast_or_null(output.getDefiningOp()); - if (!batch || batch.getNumResults() == 0) - return false; - if (!hasLiveExternalUse(output, oldComputeOps)) - return false; - return !hasRealComputeConsumer(output, oldComputeOps); -} - -void appendDestinationClass(MaterializerState& state, ProducerKey key, ClassId classId, bool ordinary = true) { - SmallVector& destinations = state.producerDestClasses[key]; - if (!llvm::is_contained(destinations, classId)) - destinations.push_back(classId); - - if (!ordinary) - return; - SmallVector& ordinaryDestinations = state.ordinaryProducerDestClasses[key]; - if (!llvm::is_contained(ordinaryDestinations, classId)) - ordinaryDestinations.push_back(classId); -} - -void replaceLiveExternalUses(Value oldValue, Value replacement, const DenseSet& oldComputeOps) { - SmallVector uses; - for (OpOperand& use : oldValue.getUses()) - uses.push_back(&use); - - for (OpOperand* use : uses) { - Operation* owner = use->getOwner(); - if (isInsideOldCompute(owner, oldComputeOps)) - continue; - use->set(replacement); - } -} - -LogicalResult collectProducerDestinations(MaterializerState& state) { - return forEachLogicalConsumerInMaterializationOrder( - state, - [&](CpuId, ClassId targetClass, ComputeInstance scheduledConsumer, ComputeInstance logicalConsumer, SlotId) - -> LogicalResult { - SmallVector consumerInputs = getComputeInstanceInputs(scheduledConsumer); - for (auto [inputIndex, input] : llvm::enumerate(consumerInputs)) { - SmallVector producerKeys; - bool projectedPlanOwnsFanout = false; - if (auto batchConsumer = dyn_cast(logicalConsumer.op)) { - bool hasProjectedReplacement = - hasProjectedInputReplacement( - state, batchConsumer, static_cast(inputIndex), input, logicalConsumer, targetClass); - bool hasUnreplacedUse = - hasUnreplacedBatchInputUse(state, batchConsumer, static_cast(inputIndex), targetClass); - if (isDebugPaddedVggInput(input)) { - static unsigned debugCount = 0; - if (debugCount++ < 12) { - llvm::errs() << "debug padded destination targetClass=" << targetClass - << " inputIndex=" << inputIndex - << " hasProjectedReplacement=" << hasProjectedReplacement - << " hasUnreplacedUse=" << hasUnreplacedUse - << " canUseProjectedLane=" - << canUseProjectedLaneInput( - state, batchConsumer, static_cast(inputIndex), input, logicalConsumer) - << " hasWholeBatchReplacement=" - << getProjectedWholeBatchReplacementProducer(state, batchConsumer, static_cast(inputIndex)) - .has_value() - << "\n"; - } - } - projectedPlanOwnsFanout = - hasProjectedReplacement && !hasUnreplacedUse; - producerKeys = projectedPlanOwnsFanout - ? collectProjectedInputKeyProducersForClass( - state, batchConsumer, static_cast(inputIndex), targetClass) - : collectProducerKeysForBatchInputDestinations( - state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); - } - else { - if (auto computeConsumer = dyn_cast(logicalConsumer.op)) { - if (hasProjectedComputeInputReplacement( - state, computeConsumer, static_cast(inputIndex), targetClass) - && !hasUnreplacedComputeInputUse(state, computeConsumer, static_cast(inputIndex), targetClass)) - continue; - } - producerKeys = collectProducerKeysForDestinations(input, logicalConsumer); - } - - for (ProducerKey producerKey : producerKeys) { - ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producerKey.instance); - auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); - if (producerCpuIt == state.schedule.computeToCpuMap.end()) - return logicalConsumer.op->emitError( - "schedule materialization found an input produced by an unscheduled compute"); - - ClassId sourceClass = state.cpuToClass.lookup(producerCpuIt->second); - if (sourceClass == targetClass) { - if (projectedPlanOwnsFanout) - continue; - SameClassConsumerLookupKey lookupKey {producerKey.instance.op, producerKey.resultIndex, targetClass}; - SmallVector& bucket = state.sameClassConsumerIndex[lookupKey]; - if (!llvm::is_contained(bucket, producerKey)) - bucket.push_back(producerKey); - continue; - } - - appendDestinationClass(state, producerKey, targetClass, /*ordinary=*/!projectedPlanOwnsFanout); - } - } - - return success(); - }); -} - -bool isStaticSliceInBounds(ArrayRef offsets, RankedTensorType sourceType, RankedTensorType fragmentType) { - if (offsets.size() != static_cast(sourceType.getRank()) - || offsets.size() != static_cast(fragmentType.getRank())) - return false; - - for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) { - int64_t offset = offsets[dim]; - if (offset < 0) - return false; - - int64_t sourceDimSize = sourceType.getDimSize(dim); - int64_t fragmentDimSize = fragmentType.getDimSize(dim); - if (fragmentDimSize < 0 || sourceDimSize < 0 || fragmentDimSize > sourceDimSize) - return false; - if (offset > sourceDimSize - fragmentDimSize) - return false; - } - - return true; -} - -bool isStaticSliceContainedIn(ArrayRef innerOffsets, - ArrayRef innerSizes, - ArrayRef outerOffsets, - ArrayRef outerSizes) { - if (innerOffsets.size() != innerSizes.size() || outerOffsets.size() != outerSizes.size() - || innerOffsets.size() != outerOffsets.size()) - return false; - - for (size_t dim = 0; dim < innerOffsets.size(); ++dim) { - if (innerSizes[dim] < 0 || outerSizes[dim] < 0) - return false; - - int64_t innerBegin = innerOffsets[dim]; - int64_t innerEnd = innerBegin + innerSizes[dim]; - int64_t outerBegin = outerOffsets[dim]; - int64_t outerEnd = outerBegin + outerSizes[dim]; - if (innerBegin < outerBegin || innerEnd > outerEnd) - return false; - } - - return true; -} - -bool areAllUnitStrides(ArrayRef strides) { - return llvm::all_of(strides, [](int64_t stride) { return stride == 1; }); -} - -static std::optional getStaticForTripCount(scf::ForOp loop) { - std::optional lowerBound = matchConstantIndexValue(loop.getLowerBound()); - std::optional upperBound = matchConstantIndexValue(loop.getUpperBound()); - std::optional step = matchConstantIndexValue(loop.getStep()); - if (!lowerBound || !upperBound || !step || *step <= 0 || *upperBound < *lowerBound) - return std::nullopt; - - int64_t distance = *upperBound - *lowerBound; - return (distance + *step - 1) / *step; -} - -static SmallVector collectEnclosingStaticProjectedLoops(Operation* op) { - SmallVector loops; - SmallVector reversedLoops; - for (Operation* current = op->getParentOp(); current; current = current->getParentOp()) - if (auto loop = dyn_cast(current)) - reversedLoops.push_back(loop); - - for (scf::ForOp loop : llvm::reverse(reversedLoops)) { - std::optional lowerBound = matchConstantIndexValue(loop.getLowerBound()); - std::optional step = matchConstantIndexValue(loop.getStep()); - std::optional tripCount = getStaticForTripCount(loop); - if (!lowerBound || !step || !tripCount) - return {}; - loops.push_back(StaticProjectedLoopInfo {.iv = cast(loop.getInductionVar()), - .lowerBound = *lowerBound, - .step = *step, - .tripCount = *tripCount}); - } - return loops; -} - -static bool isProjectedOffsetValue(Value value, - Value laneArg, - ArrayRef loops, - bool& usesDynamicBinding); - -static DenseIntElementsAttr getDenseIndexVector(Value value) { - auto constant = value.getDefiningOp(); - if (!constant) - return nullptr; - auto denseAttr = dyn_cast(constant.getValue()); - if (!denseAttr || denseAttr.getType().getRank() != 1) - return nullptr; - return denseAttr; -} - -static FailureOr getDenseIndexVectorElement(DenseIntElementsAttr denseAttr, int64_t index) { - if (index < 0 || index >= denseAttr.getNumElements()) - return failure(); - auto values = denseAttr.getValues(); - return values[index].getSExtValue(); -} - -static bool isProjectedIndexLookup(Value value, - Value laneArg, - ArrayRef loops, - bool& usesDynamicBinding) { - auto extract = value.getDefiningOp(); - if (!extract || extract.getIndices().size() != 1) - return false; - if (!getDenseIndexVector(extract.getTensor())) - return false; - - bool indexUsesDynamicBinding = false; - if (!isProjectedOffsetValue(extract.getIndices().front(), laneArg, loops, indexUsesDynamicBinding)) - return false; - usesDynamicBinding = usesDynamicBinding || indexUsesDynamicBinding; - return true; -} - -static bool -isProjectedOffsetValue(Value value, Value laneArg, ArrayRef loops, bool& usesDynamicBinding) { - if (value == laneArg) { - usesDynamicBinding = true; - return true; - } - - for (const StaticProjectedLoopInfo& loop : loops) { - if (value == loop.iv) { - usesDynamicBinding = true; - return true; - } - } - - if (matchPattern(value, m_Constant())) - return true; - - if (isProjectedIndexLookup(value, laneArg, loops, usesDynamicBinding)) - return true; - - auto affineApply = value.getDefiningOp(); - if (!affineApply || affineApply.getAffineMap().getNumResults() != 1) - return false; - - bool nestedUsesDynamicBinding = false; - for (Value operand : affineApply.getMapOperands()) { - bool operandUsesDynamicBinding = false; - if (!isProjectedOffsetValue(operand, laneArg, loops, operandUsesDynamicBinding)) - return false; - nestedUsesDynamicBinding = nestedUsesDynamicBinding || operandUsesDynamicBinding; - } - - usesDynamicBinding = usesDynamicBinding || nestedUsesDynamicBinding; - return true; -} - -static std::optional getConstantIndex(OpFoldResult value); - -static FailureOr evaluateProjectedOffsetValue(OpFoldResult value, - Value laneArg, - uint32_t lane, - ArrayRef loops, - ArrayRef loopIterationIndices) { - if (std::optional constant = getConstantIndex(value)) - return *constant; - - Value current = dyn_cast(value); - if (!current) - return failure(); - if (current == laneArg) - return static_cast(lane); - - for (auto [index, loop] : llvm::enumerate(loops)) { - if (current != loop.iv) - continue; - if (index >= loopIterationIndices.size()) - return failure(); - return loop.lowerBound + loopIterationIndices[index] * loop.step; - } - - if (auto extract = current.getDefiningOp()) { - if (extract.getIndices().size() != 1) - return failure(); - DenseIntElementsAttr denseAttr = getDenseIndexVector(extract.getTensor()); - if (!denseAttr) - return failure(); - FailureOr index = - evaluateProjectedOffsetValue(extract.getIndices().front(), laneArg, lane, loops, loopIterationIndices); - if (failed(index)) - return failure(); - return getDenseIndexVectorElement(denseAttr, *index); - } - - if (auto affineApply = current.getDefiningOp()) { - return evaluateAffineApply(affineApply, [&](Value operand) { - return evaluateProjectedOffsetValue(operand, laneArg, lane, loops, loopIterationIndices); - }); - } - - return failure(); -} - -static std::optional getConstantIndex(OpFoldResult value) { - if (auto attr = dyn_cast(value)) { - auto intAttr = dyn_cast(attr); - if (!intAttr) - return std::nullopt; - return intAttr.getInt(); - } - - Value operand = dyn_cast(value); - if (!operand) - return std::nullopt; - - if (auto constantIndex = operand.getDefiningOp()) - return constantIndex.value(); - - APInt apInt; - if (matchPattern(operand, m_ConstantInt(&apInt))) { - if (apInt.isNegative()) - return std::nullopt; - return static_cast(apInt.getSExtValue()); - } - - return std::nullopt; -} - -static std::optional matchAffineProjectedInputSlice( - BlockArgument inputArg, BlockArgument laneArg, tensor::ExtractSliceOp extract) { - const auto fail = [&](StringRef) -> std::optional { return std::nullopt; }; - - if (extract.getSource() != inputArg) - return fail("extract-source-is-not-input-arg"); - - auto inputType = dyn_cast(inputArg.getType()); - auto fragmentType = dyn_cast(extract.getResult().getType()); - if (!inputType || !fragmentType || !inputType.hasStaticShape() || !fragmentType.hasStaticShape()) - return fail("non-static-ranked-input-or-fragment"); - - if (inputType.getRank() == 0 || inputType.getRank() != fragmentType.getRank()) - return fail("rank-mismatch-or-rank-zero"); - - SmallVector offsets = extract.getMixedOffsets(); - SmallVector sizes = extract.getMixedSizes(); - SmallVector strides = extract.getMixedStrides(); - - if (offsets.size() != static_cast(inputType.getRank()) - || sizes.size() != static_cast(inputType.getRank()) - || strides.size() != static_cast(inputType.getRank())) - return fail("slice-rank-mismatch"); - - SmallVector loops = collectEnclosingStaticProjectedLoops(extract.getOperation()); - if (extract->getParentOfType() && loops.empty()) - return fail("unsupported-enclosing-loop"); - - bool hasDynamicProjection = false; - for (auto [dim, offset] : llvm::enumerate(offsets)) { - bool usesDynamicBinding = false; - if (auto value = dyn_cast(offset)) { - if (!isProjectedOffsetValue(value, laneArg, loops, usesDynamicBinding)) - return std::nullopt; - } - else if (!isa(offset)) - return std::nullopt; - if (std::optional stride = getConstantIndex(strides[dim]); !stride || *stride != 1) - return std::nullopt; - std::optional size = getConstantIndex(sizes[dim]); - if (!size || *size != fragmentType.getDimSize(dim)) - return std::nullopt; - hasDynamicProjection = hasDynamicProjection || usesDynamicBinding; - } - - if (!hasDynamicProjection) - return fail("no-dynamic-projection"); - - for (int64_t dim = 0; dim < inputType.getRank(); ++dim) - if (fragmentType.getDimSize(dim) <= 0 || fragmentType.getDimSize(dim) > inputType.getDimSize(dim)) - return std::nullopt; - - AffineProjectedInputSliceMatch match; - match.extract = extract; - match.sourceType = inputType; - match.fragmentType = fragmentType; - match.offsets.assign(offsets.begin(), offsets.end()); - match.fragmentShape.assign(fragmentType.getShape().begin(), fragmentType.getShape().end()); - match.loops = std::move(loops); - return match; -} - -static std::optional matchAffineProjectedInputSlice(SpatComputeBatch batch, - unsigned inputIndex) { - const auto fail = [&](StringRef) -> std::optional { return std::nullopt; }; - - std::optional inputArg = batch.getInputArgument(inputIndex); - std::optional laneArg = batch.getLaneArgument(); - if (!inputArg || !laneArg) - return fail("missing-input-or-lane-arg"); - - std::optional projectedMatch; - for (OpOperand& use : inputArg->getUses()) { - auto extract = dyn_cast(use.getOwner()); - if (!extract || extract.getSource() != *inputArg) - return fail("input-user-is-not-direct-extract-slice"); - - std::optional current = - matchAffineProjectedInputSlice(*inputArg, *laneArg, extract); - if (!current) - return fail("input-extract-is-not-projectable"); - if (projectedMatch) - return fail("multiple-projected-input-extracts"); - projectedMatch = std::move(current); - } - - if (!projectedMatch) - return fail("input-arg-has-no-uses"); - return projectedMatch; -} - -static SmallVector collectAffineProjectedInputSlices(SpatComputeBatch batch, - unsigned inputIndex) { - SmallVector matches; - std::optional inputArg = batch.getInputArgument(inputIndex); - std::optional laneArg = batch.getLaneArgument(); - if (!inputArg || !laneArg) - return matches; - - for (OpOperand& use : inputArg->getUses()) { - auto extract = dyn_cast(use.getOwner()); - if (!extract || extract.getSource() != *inputArg) - return {}; - - std::optional match = - matchAffineProjectedInputSlice(*inputArg, *laneArg, extract); - if (!match) - return {}; - matches.push_back(std::move(*match)); - } - return matches; -} - -std::optional -getProjectedInputSliceMatch(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex) { - ProjectedBatchInputKey key {batch.getOperation(), inputIndex}; - auto cached = state.projectedInputMatches.find(key); - if (cached != state.projectedInputMatches.end()) - return cached->second; - if (state.nonProjectedInputs.contains(key)) - return std::nullopt; - - std::optional match = matchAffineProjectedInputSlice(batch, inputIndex); - if (!match) { - state.nonProjectedInputs.insert(key); - return std::nullopt; - } - - state.projectedInputMatches.insert({key, *match}); - return match; -} - -std::optional -getProjectedWholeBatchReplacementProducer(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex) { - std::optional match = getProjectedInputSliceMatch(state, batch, inputIndex); - if (!match) - return std::nullopt; - - Value input = batch.getInputs()[inputIndex]; - std::optional wholeBatchProducer = getWholeBatchProducerKeyForDirectBatchResult(input); - if (!wholeBatchProducer) - return std::nullopt; - - if (canUseProjectedLaneInput(state, batch, inputIndex, input, ComputeInstance {batch.getOperation(), 0, 1})) - return std::nullopt; - - auto producerBatch = dyn_cast_or_null(wholeBatchProducer->instance.op); - if (!producerBatch) - return std::nullopt; - - if (failed(getBatchResultProjectionInsert(producerBatch, wholeBatchProducer->resultIndex))) - return std::nullopt; - - return wholeBatchProducer; -} - -std::optional getProjectedWholeBatchReplacementProducer(MaterializerState& state, - tensor::ExtractSliceOp extract) { - auto sourceArg = dyn_cast(extract.getSource()); - if (!sourceArg) - return std::nullopt; - - auto batch = dyn_cast_or_null(sourceArg.getOwner()->getParentOp()); - if (!batch) - return std::nullopt; - - for (unsigned inputIndex = 0; inputIndex < batch.getInputs().size(); ++inputIndex) { - std::optional inputArg = batch.getInputArgument(inputIndex); - if (!inputArg || *inputArg != sourceArg) - continue; - - std::optional match = getProjectedInputSliceMatch(state, batch, inputIndex); - if (!match || match->extract != extract) - return std::nullopt; - return getProjectedWholeBatchReplacementProducer(state, batch, inputIndex); - } - - return std::nullopt; -} - -FailureOr evaluateProjectionIndexLike(OpFoldResult value, Value laneArg, uint32_t lane); - -FailureOr evaluateProjectionIndexLike(Value value, Value laneArg, uint32_t lane) { - if (value == laneArg) - return static_cast(lane); - - if (std::optional constant = matchConstantIndexValue(value)) - return *constant; - - if (auto extract = value.getDefiningOp()) { - if (extract.getIndices().size() != 1) - return failure(); - DenseIntElementsAttr denseAttr = getDenseIndexVector(extract.getTensor()); - if (!denseAttr) - return failure(); - FailureOr index = evaluateProjectionIndexLike(extract.getIndices().front(), laneArg, lane); - if (failed(index)) - return failure(); - return getDenseIndexVectorElement(denseAttr, *index); - } - - auto affineApply = value.getDefiningOp(); - if (!affineApply || affineApply.getAffineMap().getNumResults() != 1) - return failure(); - - SmallVector operands; - operands.reserve(affineApply.getMapOperands().size()); - for (Value operand : affineApply.getMapOperands()) { - FailureOr evaluated = evaluateProjectionIndexLike(operand, laneArg, lane); - if (failed(evaluated)) - return failure(); - operands.push_back(IntegerAttr::get(IndexType::get(value.getContext()), *evaluated)); - } - - SmallVector results; - if (failed(affineApply.getAffineMap().constantFold(operands, results)) || results.size() != 1) - return failure(); - - auto intAttr = dyn_cast(results.front()); - if (!intAttr) - return failure(); - return intAttr.getInt(); -} - -FailureOr evaluateProjectionIndexLike(OpFoldResult value, Value laneArg, uint32_t lane) { - if (auto attr = llvm::dyn_cast(value)) { - auto intAttr = dyn_cast(attr); - if (!intAttr) - return failure(); - return intAttr.getInt(); - } - return evaluateProjectionIndexLike(llvm::cast(value), laneArg, lane); -} - -FailureOr getBatchResultProjectionInsert(SpatComputeBatch batch, size_t resultIndex) { - auto inParallel = dyn_cast_or_null(batch.getBody().front().getTerminator()); - if (!inParallel) - return failure(); - - auto firstOutputArg = batch.getOutputArgument(0); - if (!firstOutputArg) - return failure(); - - for (Operation& op : inParallel.getRegion().front()) { - auto insert = dyn_cast(&op); - if (!insert) - continue; - - auto outputArg = dyn_cast(insert.getDest()); - if (!outputArg || outputArg.getOwner() != &batch.getBody().front()) - continue; - - unsigned candidateIndex = outputArg.getArgNumber() - firstOutputArg->getArgNumber(); - if (candidateIndex == resultIndex) - return insert; - } - - return failure(); -} - -FailureOr> -evaluateStaticProjectionIndices(ArrayRef values, Value laneArg, uint32_t lane) { - SmallVector evaluated; - evaluated.reserve(values.size()); - for (OpFoldResult value : values) { - FailureOr index = evaluateProjectionIndexLike(value, laneArg, lane); - if (failed(index)) - return failure(); - evaluated.push_back(*index); - } - return evaluated; -} - -bool isProjectedInputSliceCompatibleWithProducerFragments(SpatComputeBatch consumerBatch, - const AffineProjectedInputSliceMatch& match, - ProducerKey producer, - uint32_t consumerLane) { - auto producerBatch = dyn_cast_or_null(producer.instance.op); - if (!producerBatch) - return true; - - FailureOr producerProjection = - getBatchResultProjectionInsert(producerBatch, producer.resultIndex); - if (failed(producerProjection)) - return true; - - std::optional producerLaneArg = producerBatch.getLaneArgument(); - std::optional consumerLaneArg = consumerBatch.getLaneArgument(); - if (!producerLaneArg || !consumerLaneArg) - return false; - - SmallVector consumerSizes(match.fragmentShape.begin(), match.fragmentShape.end()); - SmallVector loopIterationIndices(match.loops.size(), 0); - - const auto consumerSliceFitsOneProducerFragment = [&]() -> bool { - SmallVector consumerOffsets; - consumerOffsets.reserve(match.offsets.size()); - for (OpFoldResult offset : match.offsets) { - FailureOr evaluated = - evaluateProjectedOffsetValue(offset, *consumerLaneArg, consumerLane, match.loops, loopIterationIndices); - if (failed(evaluated)) - return false; - consumerOffsets.push_back(*evaluated); - } - - uint32_t producerLaneEnd = producer.instance.laneStart + producer.instance.laneCount; - for (uint32_t producerLane = producer.instance.laneStart; producerLane < producerLaneEnd; ++producerLane) { - FailureOr> producerOffsets = - evaluateStaticProjectionIndices(producerProjection->getMixedOffsets(), *producerLaneArg, producerLane); - FailureOr> producerSizes = - evaluateStaticProjectionIndices(producerProjection->getMixedSizes(), *producerLaneArg, producerLane); - FailureOr> producerStrides = - evaluateStaticProjectionIndices(producerProjection->getMixedStrides(), *producerLaneArg, producerLane); - if (failed(producerOffsets) || failed(producerSizes) || failed(producerStrides)) - return false; - if (!areAllUnitStrides(*producerStrides)) - return false; - if (isStaticSliceContainedIn(consumerOffsets, consumerSizes, *producerOffsets, *producerSizes)) - return true; - } - - return false; - }; - - if (match.loops.empty()) - return consumerSliceFitsOneProducerFragment(); - - const auto recurse = [&](auto&& self, size_t loopIndex) -> bool { - if (loopIndex == match.loops.size()) - return consumerSliceFitsOneProducerFragment(); - - for (int64_t iteration = 0; iteration < match.loops[loopIndex].tripCount; ++iteration) { - loopIterationIndices[loopIndex] = iteration; - if (!self(self, loopIndex + 1)) - return false; - } - return true; - }; - - return recurse(recurse, 0); -} - - -SmallVector -collectProjectedFragmentDemandsForMatch(SpatComputeBatch consumerBatch, - Value input, - ComputeInstance logicalConsumer, - const AffineProjectedInputSliceMatch& match) { - SmallVector demands; - - auto producerResult = dyn_cast(input); - if (!producerResult) - return demands; - - auto producerBatch = dyn_cast_or_null(producerResult.getOwner()); - if (!producerBatch || producerBatch.getNumResults() == 0) - return demands; - - FailureOr producerProjection = - getBatchResultProjectionInsert(producerBatch, producerResult.getResultNumber()); - if (failed(producerProjection)) - return demands; - - std::optional producerLaneArg = producerBatch.getLaneArgument(); - std::optional consumerLaneArg = consumerBatch.getLaneArgument(); - if (!producerLaneArg || !consumerLaneArg) - return demands; - - unsigned ordinal = 0; - - const auto appendProducerForCurrentProjectedFragment = [&](const AffineProjectedInputSliceMatch& match, - ArrayRef loopIterationIndices) -> LogicalResult { - SmallVector consumerSizes(match.fragmentShape.begin(), match.fragmentShape.end()); - SmallVector consumerOffsets; - consumerOffsets.reserve(match.offsets.size()); - for (OpFoldResult offset : match.offsets) { - FailureOr evaluated = evaluateProjectedOffsetValue( - offset, *consumerLaneArg, logicalConsumer.laneStart, match.loops, loopIterationIndices); - if (failed(evaluated)) - return failure(); - consumerOffsets.push_back(*evaluated); - } - - for (uint32_t producerLane = 0; producerLane < static_cast(producerBatch.getLaneCount()); ++producerLane) { - FailureOr> producerOffsets = - evaluateStaticProjectionIndices(producerProjection->getMixedOffsets(), *producerLaneArg, producerLane); - FailureOr> producerSizes = - evaluateStaticProjectionIndices(producerProjection->getMixedSizes(), *producerLaneArg, producerLane); - FailureOr> producerStrides = - evaluateStaticProjectionIndices(producerProjection->getMixedStrides(), *producerLaneArg, producerLane); - if (failed(producerOffsets) || failed(producerSizes) || failed(producerStrides)) - return failure(); - if (!areAllUnitStrides(*producerStrides)) - return failure(); - if (!isStaticSliceContainedIn(consumerOffsets, consumerSizes, *producerOffsets, *producerSizes)) - continue; - - demands.push_back(ProjectedProducerFragmentDemand { - getBatchLaneProducerKey(producerBatch, producerLane, 1, producerResult.getResultNumber()), - std::move(consumerOffsets), - ordinal++}); - return success(); - } - - return failure(); - }; - - SmallVector loopIterationIndices(match.loops.size(), 0); - const auto recurse = [&](auto&& self, size_t loopIndex) -> LogicalResult { - if (loopIndex == match.loops.size()) - return appendProducerForCurrentProjectedFragment(match, loopIterationIndices); - - for (int64_t iteration = 0; iteration < match.loops[loopIndex].tripCount; ++iteration) { - loopIterationIndices[loopIndex] = iteration; - if (failed(self(self, loopIndex + 1))) - return failure(); - } - return success(); - }; - - LogicalResult result = match.loops.empty() ? appendProducerForCurrentProjectedFragment(match, loopIterationIndices) - : recurse(recurse, 0); - - if (failed(result)) - demands.clear(); - return demands; -} - -SmallVector -collectProjectedFragmentDemandsForBatchInput(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer) { - (void)state; - SmallVector demands; - for (const AffineProjectedInputSliceMatch& match : collectAffineProjectedInputSlices(consumerBatch, inputIndex)) { - SmallVector matchDemands = - collectProjectedFragmentDemandsForMatch(consumerBatch, input, logicalConsumer, match); - if (matchDemands.empty()) - return {}; - llvm::append_range(demands, matchDemands); - } - return demands; -} - -SmallVector collectProjectedProducerKeysForBatchInput(MaterializerState& state, - SpatComputeBatch consumerBatch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer) { - SmallVector keys; - for (const ProjectedProducerFragmentDemand& demand : collectProjectedFragmentDemandsForBatchInput( - state, consumerBatch, inputIndex, input, logicalConsumer)) { - if (!llvm::is_contained(keys, demand.producer)) - keys.push_back(demand.producer); - } - return keys; -} - -LogicalResult collectProjectedTransfers(MaterializerState& state) { - struct PendingProjectedTransferDescriptor { - ProjectedBatchInputKey inputKey; - Operation* extractOp = nullptr; - RankedTensorType sourceType; - RankedTensorType fragmentType; - SmallVector fragmentShape; - SmallVector, 16>, 8> fragmentOffsetsByLane; - SmallVector loopLowerBounds; - SmallVector loopSteps; - SmallVector loopTripCounts; - bool invalid = false; - }; - - DenseMap, ProducerKeyInfo> pending; - - const auto isIdentityProjectedTransfer = [&](const PendingProjectedTransferDescriptor& descriptor) { - if (!descriptor.sourceType || descriptor.sourceType != descriptor.fragmentType) - return false; - - if (descriptor.fragmentOffsetsByLane.size() != 1) - return false; - - ArrayRef> fragments = descriptor.fragmentOffsetsByLane.front(); - if (fragments.size() != 1) - return false; - - return llvm::all_of(fragments.front(), [](int64_t offset) { return offset == 0; }); - }; - - const auto appendEvaluatedFragments = [&](PendingProjectedTransferDescriptor& descriptor, - unsigned targetLane, - const AffineProjectedInputSliceMatch& match, - Value laneArg, - uint32_t lane) -> LogicalResult { - SmallVector loopIterationIndices; - loopIterationIndices.resize(match.loops.size(), 0); - - const auto appendOneFragment = [&]() -> LogicalResult { - SmallVector evaluatedOffsets; - evaluatedOffsets.reserve(match.offsets.size()); - for (OpFoldResult offset : match.offsets) { - FailureOr evaluated = - evaluateProjectedOffsetValue(offset, laneArg, lane, match.loops, loopIterationIndices); - if (failed(evaluated)) - return failure(); - evaluatedOffsets.push_back(*evaluated); - } - - if (!isStaticSliceInBounds(evaluatedOffsets, match.sourceType, match.fragmentType)) - return failure(); - - descriptor.fragmentOffsetsByLane[targetLane].push_back(std::move(evaluatedOffsets)); - return success(); - }; - - if (match.loops.empty()) - return appendOneFragment(); - - const auto recurse = [&](auto&& self, size_t loopIndex) -> LogicalResult { - if (loopIndex == match.loops.size()) - return appendOneFragment(); - - for (int64_t iteration = 0; iteration < match.loops[loopIndex].tripCount; ++iteration) { - loopIterationIndices[loopIndex] = iteration; - if (failed(self(self, loopIndex + 1))) - return failure(); - } - return success(); - }; - - return recurse(recurse, 0); - }; - - if (failed(forEachLogicalConsumerInMaterializationOrder( - state, - [&](CpuId cpu, - ClassId targetClassId, - ComputeInstance consumer, - ComputeInstance logicalConsumer, - SlotId logicalSlot) -> LogicalResult { - auto batch = dyn_cast(consumer.op); - if (!batch) - return success(); - - MaterializedClass& targetClass = state.classes[targetClassId]; - unsigned targetLane = 0; - if (targetClass.isBatch) { - auto targetLaneIt = targetClass.cpuToLane.find(cpu); - if (targetLaneIt == targetClass.cpuToLane.end()) - return consumer.op->emitError("projected transfer collection could not recover target lane"); - targetLane = targetLaneIt->second; - } - - for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs())) { - ProjectedBatchInputKey currentInputKey {batch.getOperation(), static_cast(inputIndex)}; - SmallVector fragmentDemands = - collectProjectedFragmentDemandsForBatchInput( - state, batch, static_cast(inputIndex), input, logicalConsumer); - SmallVector matches = - collectAffineProjectedInputSlices(batch, static_cast(inputIndex)); - - if (!matches.empty() && !fragmentDemands.empty()) { - for (AffineProjectedInputSliceMatch& match : matches) { - SmallVector matchDemands = - collectProjectedFragmentDemandsForMatch(batch, input, logicalConsumer, match); - if (matchDemands.empty()) - return targetClass.op->emitError("failed to collect projected input transfer fragments"); - - ProjectedInputTransferPlan& plan = - state.projectedInputTransferPlans[match.extract.getOperation()][targetClassId]; - if (plan.fragments.empty()) { - plan.inputKey = currentInputKey; - plan.extractOp = match.extract.getOperation(); - plan.layout.fragmentType = match.fragmentType; - plan.layout.fragmentShape = match.fragmentShape; - plan.layout.loopLowerBounds.reserve(match.loops.size()); - plan.layout.loopSteps.reserve(match.loops.size()); - plan.layout.loopTripCounts.reserve(match.loops.size()); - for (const StaticProjectedLoopInfo& loop : match.loops) { - plan.layout.loopLowerBounds.push_back(loop.lowerBound); - plan.layout.loopSteps.push_back(loop.step); - plan.layout.loopTripCounts.push_back(loop.tripCount); - } - plan.layout.fragmentsPerLogicalSlot = getProjectedFragmentsPerLogicalSlot(plan.layout.loopTripCounts); - plan.layout.payloadFragmentCount = plan.layout.fragmentsPerLogicalSlot; - } - if (!(plan.inputKey == currentInputKey) || plan.extractOp != match.extract.getOperation() - || plan.layout.fragmentType != match.fragmentType || plan.layout.fragmentShape != match.fragmentShape - || plan.layout.loopTripCounts.size() != match.loops.size()) - return targetClass.op->emitError("inconsistent projected input transfer plan"); - - auto targetCpu = getCheckedCoreId(targetClass.op, cpu, "projected input target core id"); - if (failed(targetCpu)) - return failure(); - - for (const ProjectedProducerFragmentDemand& demand : matchDemands) { - ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, demand.producer.instance); - auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); - if (producerCpuIt == state.schedule.computeToCpuMap.end()) - return logicalConsumer.op->emitError( - "projected input transfer found a fragment produced by an unscheduled compute"); - - auto sourceCpu = getCheckedCoreId( - targetClass.op, producerCpuIt->second, "projected input source core id"); - if (failed(sourceCpu)) - return failure(); - - ProjectedInputTransferFragment fragment; - fragment.producer = demand.producer; - fragment.fragmentOffsets = demand.fragmentOffsets; - fragment.targetLane = targetLane; - fragment.ordinal = targetClass.isBatch - ? targetLane * plan.layout.fragmentsPerLogicalSlot + demand.ordinal - : demand.ordinal; - fragment.sourceCoreId = *sourceCpu; - fragment.targetCoreId = *targetCpu; - plan.fragments.push_back(std::move(fragment)); - } - } - continue; - } - - std::optional match = - getProjectedInputSliceMatch(state, batch, static_cast(inputIndex)); - SmallVector producers; - if (!fragmentDemands.empty()) { - for (const ProjectedProducerFragmentDemand& demand : fragmentDemands) - if (!llvm::is_contained(producers, demand.producer)) - producers.push_back(demand.producer); - } - else { - producers = collectProducerKeysForDestinations(input, logicalConsumer); - } - if (producers.size() != 1) - continue; - ProducerKey producer = producers.front(); - - ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producer.instance); - auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); - if (producerCpuIt == state.schedule.computeToCpuMap.end()) - continue; - - ClassId sourceClassId = state.cpuToClass.lookup(producerCpuIt->second); - if (sourceClassId == targetClassId) - continue; - - if (!match) - continue; - if (!isProjectedInputSliceCompatibleWithProducerFragments( - batch, *match, producer, logicalConsumer.laneStart)) - continue; - - PendingProjectedTransferDescriptor& descriptor = pending[producer][targetClassId]; - if (descriptor.fragmentOffsetsByLane.empty()) { - descriptor.inputKey = {batch.getOperation(), static_cast(inputIndex)}; - descriptor.extractOp = match->extract.getOperation(); - descriptor.sourceType = match->sourceType; - descriptor.fragmentType = match->fragmentType; - descriptor.fragmentShape = match->fragmentShape; - descriptor.fragmentOffsetsByLane.resize(targetClass.isBatch ? targetClass.cpus.size() : 1); - descriptor.loopLowerBounds.reserve(match->loops.size()); - descriptor.loopSteps.reserve(match->loops.size()); - descriptor.loopTripCounts.reserve(match->loops.size()); - for (const StaticProjectedLoopInfo& loop : match->loops) { - descriptor.loopLowerBounds.push_back(loop.lowerBound); - descriptor.loopSteps.push_back(loop.step); - descriptor.loopTripCounts.push_back(loop.tripCount); - } - } - - if (!(descriptor.inputKey == currentInputKey) || descriptor.extractOp != match->extract.getOperation() - || descriptor.sourceType != match->sourceType || descriptor.fragmentType != match->fragmentType - || descriptor.fragmentShape != match->fragmentShape - || descriptor.loopLowerBounds.size() != match->loops.size()) { - descriptor.invalid = true; - continue; - } - for (auto [index, loop] : llvm::enumerate(match->loops)) { - if (descriptor.loopLowerBounds[index] != loop.lowerBound || descriptor.loopSteps[index] != loop.step - || descriptor.loopTripCounts[index] != loop.tripCount) { - descriptor.invalid = true; - break; - } - } - if (descriptor.invalid) - continue; - - if (targetLane >= descriptor.fragmentOffsetsByLane.size()) { - descriptor.invalid = true; - continue; - } - - if (failed(appendEvaluatedFragments( - descriptor, targetLane, *match, *batch.getLaneArgument(), logicalConsumer.laneStart))) { - descriptor.invalid = true; - continue; - } - - (void) logicalSlot; - } - - return success(); - }))) - return failure(); - - for (auto& producerEntry : pending) { - ProducerKey producer = producerEntry.first; - for (auto& classEntry : producerEntry.second) { - ClassId targetClassId = classEntry.first; - PendingProjectedTransferDescriptor& pendingDescriptor = classEntry.second; - - if (pendingDescriptor.invalid) - continue; - if (pendingDescriptor.fragmentOffsetsByLane.empty()) - continue; - if (isIdentityProjectedTransfer(pendingDescriptor)) - continue; - - MaterializedClass& targetClass = state.classes[targetClassId]; - ProjectedTransferDescriptor descriptor; - descriptor.inputKey = pendingDescriptor.inputKey; - descriptor.extractOp = pendingDescriptor.extractOp; - descriptor.layout.fragmentType = pendingDescriptor.fragmentType; - descriptor.layout.fragmentShape = pendingDescriptor.fragmentShape; - descriptor.layout.loopLowerBounds = pendingDescriptor.loopLowerBounds; - descriptor.layout.loopSteps = pendingDescriptor.loopSteps; - descriptor.layout.loopTripCounts = pendingDescriptor.loopTripCounts; - descriptor.layout.fragmentsPerLogicalSlot = getProjectedFragmentsPerLogicalSlot(descriptor.layout.loopTripCounts); - if (targetClass.isBatch) { - unsigned payloadFragmentCount = pendingDescriptor.fragmentOffsetsByLane.front().size(); - if (payloadFragmentCount == 0) - continue; - - bool uniform = true; - for (ArrayRef> laneFragments : pendingDescriptor.fragmentOffsetsByLane) { - if (laneFragments.size() != payloadFragmentCount) { - uniform = false; - break; - } - } - if (!uniform) - continue; - - descriptor.layout.payloadFragmentCount = payloadFragmentCount; - descriptor.fragmentOffsets.reserve(pendingDescriptor.fragmentOffsetsByLane.size() * payloadFragmentCount); - for (ArrayRef> laneFragments : pendingDescriptor.fragmentOffsetsByLane) - llvm::append_range(descriptor.fragmentOffsets, laneFragments); - } - else { - if (pendingDescriptor.fragmentOffsetsByLane.size() != 1) - return targetClass.op->emitError("scalar projected transfer descriptor expected one local offset stream"); - if (pendingDescriptor.fragmentOffsetsByLane.front().empty()) - continue; - - descriptor.layout.payloadFragmentCount = pendingDescriptor.fragmentOffsetsByLane.front().size(); - llvm::append_range(descriptor.fragmentOffsets, pendingDescriptor.fragmentOffsetsByLane.front()); - if (descriptor.fragmentOffsets.size() != descriptor.layout.payloadFragmentCount) - return targetClass.op->emitError("scalar projected transfer offset count does not match the local run"); - } - if (failed(finalizeProjectedTransferDescriptor(targetClass.op, descriptor))) - return failure(); - - state.projectedTransfers[producer][targetClassId] = std::move(descriptor); - } - } - - return success(); -} - -LogicalResult buildGlobalCommunicationPlan(MaterializerState& state) { - SmallVector classes; - classes.reserve(state.classes.size()); - for (const MaterializedClass& klass : state.classes) { - classes.push_back(CommunicationClassInfo { - .classId = klass.id, - .rank = klass.id, - .isBatch = klass.isBatch, - .op = klass.op, - .cpus = klass.cpus, - }); - } - - CommunicationPlan::ProducerSourceClassMap producerSourceClasses; - DenseMap producerPayloadTypes; - for (const auto& [producer, destinations] : state.producerDestClasses) { - (void)destinations; - ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producer.instance); - auto cpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return producer.instance.op->emitError("global communication planning found an unscheduled producer"); - auto classIt = state.cpuToClass.find(cpuIt->second); - if (classIt == state.cpuToClass.end()) - return producer.instance.op->emitError("global communication planning could not map producer CPU to class"); - producerSourceClasses[producer] = classIt->second; - - if (auto batch = dyn_cast_or_null(producer.instance.op)) { - SmallVector& fragmentTypes = getBatchOutputFragmentTypesCached(state, batch); - if (producer.resultIndex >= fragmentTypes.size() || !fragmentTypes[producer.resultIndex]) - return producer.instance.op->emitError("global communication planning could not recover batch fragment type"); - producerPayloadTypes[producer] = fragmentTypes[producer.resultIndex]; - continue; - } - - if (producer.resultIndex >= producer.instance.op->getNumResults()) - return producer.instance.op->emitError("global communication planning found an invalid producer result index"); - producerPayloadTypes[producer] = producer.instance.op->getResult(producer.resultIndex).getType(); - } - - FailureOr plan = CommunicationPlan::build(state.func, - classes, - state.cpuToClass, - producerSourceClasses, - producerPayloadTypes, - state.producerDestClasses, - state.ordinaryProducerDestClasses, - state.projectedInputTransferPlans, - state.hostOutputOwners, - state.liveExternalUseCache, - state.nextChannelId); - if (failed(plan)) - return failure(); - state.communicationPlan = std::move(*plan); - state.communicationExchangeStates.assign(state.communicationPlan.getExchanges().size(), - CommunicationExchangeEmissionState {}); - return success(); -} - -ArrayRef getDestinationClasses(MaterializerState& state, ProducerKey key) { - auto it = state.producerDestClasses.find(key); - if (it == state.producerDestClasses.end()) - return {}; - return it->second; -} - -static bool hasOrdinaryDestinationClass(MaterializerState& state, ProducerKey key, ClassId classId) { - auto it = state.ordinaryProducerDestClasses.find(key); - return it != state.ordinaryProducerDestClasses.end() && llvm::is_contained(it->second, classId); -} - -SmallVector filterProducerKeysForDestination(MaterializerState& state, - ArrayRef keys, - ClassId destinationClass) { - SmallVector filtered; - for (ProducerKey key : keys) - if (llvm::is_contained(getDestinationClasses(state, key), destinationClass)) - filtered.push_back(key); - return filtered; -} - -// ----------------------------------------------------------------------------- -// Communication materialization helpers. -// ----------------------------------------------------------------------------- - -void appendScalarSendAtCurrentInsertionPoint(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - int64_t channelId, - int32_t sourceCoreId, - int32_t targetCoreId, - Location loc) { - Value channelIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, channelId); - Value sourceCoreIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, sourceCoreId); - Value targetCoreIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, targetCoreId); - SpatChannelSendOp::create(state.rewriter, loc, channelIdValue, sourceCoreIdValue, targetCoreIdValue, payload); -} - -LogicalResult appendScalarSendLoopAtCurrentInsertionPoint(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const MessageVector& messages, - Location loc) { - if (messages.size() == 1) { - appendScalarSendAtCurrentInsertionPoint(state, - sourceClass, - payload, - messages.channelIds.front(), - static_cast(messages.sourceCoreIds.front()), - static_cast(messages.targetCoreIds.front()), - loc); - return success(); - } - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0); - Value upperBound = - getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast(messages.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1); - - auto sendLoop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {}, - [&](OpBuilder&, Location, Value index, ValueRange, SmallVectorImpl&) { - Value channelId = createIndexedChannelId(state, sourceClass.op, messages, index, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, messages, index, loc); - Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, messages, index, loc); - SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); - return success(); - }); - if (failed(sendLoop)) - return failure(); - return success(); -} - -void appendScalarSend(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - int64_t channelId, - int32_t sourceCoreId, - int32_t targetCoreId, - Location loc) { - assert(!sourceClass.isBatch && "scalar send helper expects a scalar source class"); - - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - appendScalarSendAtCurrentInsertionPoint(state, sourceClass, payload, channelId, sourceCoreId, targetCoreId, loc); -} - -LogicalResult appendScalarSendLoop(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const MessageVector& messages, - Location loc) { - assert(!sourceClass.isBatch && "scalar send loop expects a scalar source class"); - assert(messages.size() > 1 && "send loop is only useful for multiple sends"); - assert(succeeded(messages.verify(sourceClass.op)) && "message metadata is inconsistent"); - - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - return appendScalarSendLoopAtCurrentInsertionPoint(state, sourceClass, payload, messages, loc); -} - -LogicalResult appendSend(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const MessageVector& messages, - Location loc) { - assert(succeeded(messages.verify(sourceClass.op)) && "message metadata is inconsistent"); - assert(!messages.empty() && "expected at least one send"); - - if (sourceClass.isBatch) { - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - if (messages.size() != sourceClass.cpus.size()) - return sourceClass.op->emitError("batch send expects exactly one message per materialized lane") - << " messageCount=" << messages.size() << " laneCount=" << sourceClass.cpus.size(); - - Value channelId = createLaneIndexedIndexValue(state, sourceClass, messages.channelIds, loc); - Value sourceCoreId = createLaneIndexedIndexValue(state, sourceClass, messages.sourceCoreIds, loc); - Value targetCoreId = createLaneIndexedIndexValue(state, sourceClass, messages.targetCoreIds, loc); - SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); - return success(); - } - - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - if (messages.size() == 1) { - appendScalarSend(state, - sourceClass, - payload, - messages.channelIds.front(), - messages.sourceCoreIds.front(), - messages.targetCoreIds.front(), - loc); - return success(); - } - - return appendScalarSendLoop(state, sourceClass, payload, messages, loc); -} - -Value appendScalarReceive(MaterializerState& state, - MaterializedClass& targetClass, - Type type, - int64_t channelId, - int32_t sourceCoreId, - int32_t targetCoreId, - Location loc) { - assert(!targetClass.isBatch && "scalar receive helper expects a scalar target class"); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value channelIdValue = getOrCreateIndexConstant(state.constantFolder, targetClass.op, channelId); - Value sourceCoreIdValue = getOrCreateIndexConstant(state.constantFolder, targetClass.op, sourceCoreId); - Value targetCoreIdValue = getOrCreateIndexConstant(state.constantFolder, targetClass.op, targetCoreId); - return SpatChannelReceiveOp::create(state.rewriter, loc, type, channelIdValue, sourceCoreIdValue, targetCoreIdValue) - .getOutput(); -} - -Value appendReceive( - MaterializerState& state, MaterializedClass& targetClass, Type type, const MessageVector& messages, Location loc) { - assert(succeeded(messages.verify(targetClass.op)) && "message metadata is inconsistent"); - assert(!messages.empty() && "expected at least one receive"); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - - if (targetClass.isBatch) { - if (messages.size() != targetClass.cpus.size()) { - targetClass.op->emitOpError("batch receive expects exactly one message per materialized lane") - << " messageCount=" << messages.size() << " laneCount=" << targetClass.cpus.size(); - return Value(); - } - Value channelId = createLaneIndexedIndexValue(state, targetClass, messages.channelIds, loc); - Value sourceCoreId = createLaneIndexedIndexValue(state, targetClass, messages.sourceCoreIds, loc); - Value targetCoreId = createLaneIndexedIndexValue(state, targetClass, messages.targetCoreIds, loc); - return SpatChannelReceiveOp::create(state.rewriter, loc, type, channelId, sourceCoreId, targetCoreId).getOutput(); - } - - assert(messages.size() == 1 && "scalar target class can only receive one message at a time"); - return appendScalarReceive(state, - targetClass, - type, - messages.channelIds.front(), - messages.sourceCoreIds.front(), - messages.targetCoreIds.front(), - loc); -} - -FailureOr buildMessageVectorForExchangeIds(MaterializerState& state, - ArrayRef ids, - CommunicationExchangeKind expectedKind, - ClassId expectedSourceClass, - ClassId expectedTargetClass, - Type expectedPayloadType, - Operation* anchor) { - MessageVector messages; - messages.channelIds.reserve(ids.size()); - messages.sourceCoreIds.reserve(ids.size()); - messages.targetCoreIds.reserve(ids.size()); - for (CommunicationExchangeId id : ids) { - const CommunicationExchange& exchange = state.communicationPlan.getExchange(id); - if (exchange.kind != expectedKind || exchange.sourceClass != expectedSourceClass - || exchange.targetClass != expectedTargetClass) { - return anchor->emitError("planned communication exchange does not match requested endpoint"), failure(); - } - if (expectedPayloadType && exchange.payloadType && exchange.payloadType != expectedPayloadType) { - // Packed scalar fanout can bind a wider packed payload than the per-lane - // producer type known during planning; the actual receive type is supplied - // by the materialization site. - } - messages.append(exchange.channelId, - static_cast(exchange.sourceCore), - static_cast(exchange.targetCore)); - } - if (failed(messages.verify(anchor))) - return failure(); - return messages; -} - -LogicalResult bindCommunicationPayload(MaterializerState& state, - CommunicationExchangeId id, - MaterializedClass& sourceClass, - Value payload, - Operation* anchor, - Location loc) { - if (id.value >= state.communicationExchangeStates.size()) - return anchor->emitError("communication payload binding references an invalid exchange id"); - const CommunicationExchange& exchange = state.communicationPlan.getExchange(id); - if (exchange.sourceClass != sourceClass.id) - return anchor->emitError("communication payload binding source class mismatch"); - - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; - if (emission.producerBound && emission.producerPayload != payload) - return anchor->emitError("communication exchange was bound to two different payload values"); - - emission.producerBound = true; - emission.producerPayload = payload; - emission.loc = loc; - if (Operation* payloadDef = payload.getDefiningOp()) - emission.producerAnchor = getTopLevelMaterializedClassBodyOp(payloadDef, sourceClass); - - return success(); -} - -static LogicalResult emitPlannedSend(MaterializerState& state, - CommunicationExchangeId id, - MaterializedClass& sourceClass, - Location loc) { - const CommunicationExchange& exchange = state.communicationPlan.getExchange(id); - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; - if (exchange.sourceClass != sourceClass.id) - return sourceClass.op->emitError("planned send source class mismatch"); - if (emission.sendEmitted) - return sourceClass.op->emitError("planned send emitted twice"); - if (!emission.producerBound) - return success(); - - FailureOr messages = buildMessageVectorForExchangeIds(state, - ArrayRef(id), - exchange.kind, - exchange.sourceClass, - exchange.targetClass, - Type(), - sourceClass.op); - if (failed(messages)) - return failure(); - - if (sourceClass.isBatch) { - auto batch = dyn_cast(sourceClass.op); - if (!batch) - return sourceClass.op->emitError("planned batch send requires a scheduled compute_batch"); - std::optional laneArg = batch.getLaneArgument(); - if (!laneArg) - return sourceClass.op->emitError("planned batch send requires a lane argument"); - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - Value sourceLane = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, exchange.sourceLane); - Value isSender = arith::CmpIOp::create(state.rewriter, loc, arith::CmpIPredicate::eq, *laneArg, sourceLane); - auto ifOp = scf::IfOp::create(state.rewriter, loc, TypeRange {}, isSender, false); - OpBuilder::InsertionGuard guard(state.rewriter); - Block& thenBlock = ifOp.getThenRegion().front(); - if (!thenBlock.empty()) - if (auto yieldOp = dyn_cast(thenBlock.back())) - yieldOp.erase(); - state.rewriter.setInsertionPointToEnd(&thenBlock); - SpatChannelSendOp::create(state.rewriter, - loc, - getOrCreateIndexConstant(state.constantFolder, sourceClass.op, exchange.channelId), - getOrCreateIndexConstant(state.constantFolder, sourceClass.op, exchange.sourceCore), - getOrCreateIndexConstant(state.constantFolder, sourceClass.op, exchange.targetCore), - emission.producerPayload); - scf::YieldOp::create(state.rewriter, loc); - } - else if (failed(appendSend(state, sourceClass, emission.producerPayload, *messages, loc))) { - return failure(); - } - emission.sendEmitted = true; - if (Operation* payloadDef = emission.producerPayload.getDefiningOp()) - emission.sendAnchor = getTopLevelMaterializedClassBodyOp(payloadDef, sourceClass); - return success(); -} - -FailureOr emitPlannedReceive(MaterializerState& state, - CommunicationExchangeId id, - MaterializedClass& targetClass, - Location loc) { - if (id.value >= state.communicationExchangeStates.size()) - return targetClass.op->emitError("planned receive references an invalid exchange id"), failure(); - const CommunicationExchange& exchange = state.communicationPlan.getExchange(id); - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; - if (exchange.targetClass != targetClass.id) - return targetClass.op->emitError("planned receive target class mismatch"), failure(); - if (emission.receiveEmitted) - return emission.receiveValue; - if (!emission.sendEmitted) - return failure(); - - FailureOr messages = buildMessageVectorForExchangeIds(state, - ArrayRef(id), - exchange.kind, - exchange.sourceClass, - exchange.targetClass, - Type(), - targetClass.op); - if (failed(messages)) - return failure(); - - Value received = appendReceive(state, targetClass, exchange.payloadType, *messages, loc); - if (!received) - return failure(); - emission.receiveEmitted = true; - emission.receiveValue = received; - if (Operation* receiveDef = received.getDefiningOp()) - emission.receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, targetClass); - return received; -} - -SmallVector collectDestinationClassesForKeys(MaterializerState& state, ArrayRef keys) { - SmallVector destinations; - - for (ProducerKey key : keys) - for (ClassId destinationClass : getDestinationClasses(state, key)) - destinations.push_back(destinationClass); - - llvm::sort(destinations); - destinations.erase(std::unique(destinations.begin(), destinations.end()), destinations.end()); - return destinations; -} - -SmallVector collectProjectedProducerKeysForOutput(MaterializerState& state, Value originalOutput) { - SmallVector keys; - auto appendKey = [&](ProducerKey key) { - if (!key.instance.op || key.resultIndex >= key.instance.op->getNumResults()) - return; - if (key.instance.op->getResult(key.resultIndex) != originalOutput) - return; - if (!llvm::is_contained(keys, key)) - keys.push_back(key); - }; - for (const auto& producerEntry : state.projectedTransfers) - appendKey(producerEntry.first); - for (const auto& extractEntry : state.projectedInputTransferPlans) - for (const auto& classEntry : extractEntry.second) - for (const ProjectedInputTransferFragment& fragment : classEntry.second.fragments) - appendKey(fragment.producer); - return keys; -} - -ReceiveMessagePartition partitionReceiveMessagesInPlanOrder(const MessageVector& messages) { - ReceiveMessagePartition partition; - for (size_t index = 0; index < messages.size(); ++index) - partition.remainingIndices.push_back(index); - return partition; -} - -LogicalResult appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - ArrayRef fragments, - Location loc) { - assert(sourceClass.isBatch && "conditional projected sends require a batch source class"); - if (fragments.empty()) - return success(); - - auto batch = cast(sourceClass.op); - auto laneArg = batch.getLaneArgument(); - if (!laneArg) - return batch.emitOpError("expected lane argument while emitting projected input sends"); - - SmallVector sourceLanes; - MessageVector messages; - sourceLanes.reserve(fragments.size()); - messages.channelIds.reserve(fragments.size()); - messages.sourceCoreIds.reserve(fragments.size()); - messages.targetCoreIds.reserve(fragments.size()); - - for (ProjectedInputTransferFragment* fragment : fragments) { - auto laneIt = sourceClass.cpuToLane.find(fragment->sourceCoreId); - if (laneIt == sourceClass.cpuToLane.end()) - return sourceClass.op->emitError("projected input send could not map source core to batch lane") - << " sourceCore=" << fragment->sourceCoreId; - sourceLanes.push_back(laneIt->second); - messages.append(fragment->channelId, fragment->sourceCoreId, fragment->targetCoreId); - } - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0); - Value upperBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast(fragments.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1); - auto sendLoop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {}, - [&](OpBuilder&, Location loopLoc, Value index, ValueRange, SmallVectorImpl&) { - Value sourceLane = createIndexedIndexValue(state, sourceClass.op, sourceLanes, index, loopLoc); - Value isSender = arith::CmpIOp::create(state.rewriter, loopLoc, arith::CmpIPredicate::eq, *laneArg, sourceLane); - auto ifOp = scf::IfOp::create(state.rewriter, loopLoc, TypeRange {}, isSender, false); - OpBuilder::InsertionGuard guard(state.rewriter); - Block& thenBlock = ifOp.getThenRegion().front(); - if (!thenBlock.empty()) - if (auto yieldOp = dyn_cast(thenBlock.back())) - yieldOp.erase(); - state.rewriter.setInsertionPointToEnd(&thenBlock); - Value channelId = createIndexedChannelId(state, sourceClass.op, messages, index, loopLoc); - Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, messages, index, loopLoc); - Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, messages, index, loopLoc); - auto send = SpatChannelSendOp::create(state.rewriter, loopLoc, channelId, sourceCoreId, targetCoreId, payload); - tagProjectedInputTransfer(send.getOperation()); - scf::YieldOp::create(state.rewriter, loopLoc); - return success(); - }); - if (failed(sendLoop)) - return failure(); - if (sendLoop->loop) - tagProjectedInputTransfer(sendLoop->loop.getOperation()); - return success(); -} - -LogicalResult appendScalarProjectedInputSendAtCurrentInsertionPoint(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - ProjectedInputTransferFragment& fragment, - Location loc) { - assert(!sourceClass.isBatch && "scalar projected input send requires a scalar source class"); - if (sourceClass.cpus.empty()) - return sourceClass.op->emitError("scalar projected input send requires one source core"); - int64_t sourceCore = sourceClass.cpus.front(); - if (sourceCore != fragment.sourceCoreId) - return sourceClass.op->emitError("projected input send source core does not match materialized scalar class") - << " expected=" << sourceCore << " actual=" << fragment.sourceCoreId; - - Value channelId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, fragment.channelId); - Value sourceCoreId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, fragment.sourceCoreId); - Value targetCoreId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, fragment.targetCoreId); - auto send = SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); - tagProjectedInputTransfer(send.getOperation()); - return success(); -} - -LogicalResult emitProjectedInputSendsAtCurrentInsertionPoint(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - ArrayRef fragments, - Location loc) { - if (fragments.empty()) - return success(); - if (sourceClass.isBatch) - return appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint( - state, sourceClass, payload, fragments, loc); - - for (ProjectedInputTransferFragment* fragment : fragments) - if (failed(appendScalarProjectedInputSendAtCurrentInsertionPoint(state, sourceClass, payload, *fragment, loc))) - return failure(); - return success(); -} - -static std::optional getProjectedExchangeId(MaterializerState& state, - const ProjectedInputTransferFragment& fragment, - Operation* anchor) { - ArrayRef ids = state.communicationPlan.getProjectedInputExchanges(fragment); - if (ids.size() > 1 && anchor) - anchor->emitError("projected remote fragment has multiple communication exchanges") - << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId; - if (ids.empty() && anchor) - anchor->emitError("projected remote fragment is missing from communication plan") - << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId - << " channelId=" << fragment.channelId; - if (ids.size() != 1) - return std::nullopt; - return ids.front(); -} - -static bool isProjectedRemoteFragmentReadyForReceive(MaterializerState& state, - const ProjectedInputTransferFragment& fragment, - Operation* anchor) { - std::optional exchangeId = getProjectedExchangeId(state, fragment, anchor); - if (!exchangeId) - return false; - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId->value]; - if (!emission.sendEmitted) { - recordDeferredProducerValue(state, fragment.producer); - return false; - } - return true; -} - -static LogicalResult preflightProjectedProducerReceives(MaterializerState& state, - MaterializedClass& targetClass, - ProducerKey producer) { - for (auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(targetClass.id); - if (classIt == extractEntry.second.end()) - continue; - - for (const ProjectedInputTransferFragment& fragment : classIt->second.fragments) { - if (!(fragment.producer == producer) || fragment.sourceCoreId == fragment.targetCoreId) - continue; - - std::optional exchangeId = - getProjectedExchangeId(state, fragment, targetClass.op); - if (!exchangeId) - return failure(); - - const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId->value]; - (void)emission; - } - } - return success(); -} - -static LogicalResult preflightProjectedWholeBatchProducerReceives(MaterializerState& state, - MaterializedClass& targetClass, - ProducerKey producer) { - auto batch = dyn_cast_or_null(producer.instance.op); - if (!batch) - return success(); - - uint32_t laneEnd = producer.instance.laneStart + producer.instance.laneCount; - if (laneEnd > static_cast(batch.getLaneCount())) - return failure(); - - for (uint32_t lane = producer.instance.laneStart; lane < laneEnd; ++lane) - if (failed(preflightProjectedProducerReceives( - state, targetClass, getBatchLaneProducerKey(batch, lane, 1, producer.resultIndex)))) - return failure(); - return success(); -} - -static bool isRemoteProjectedInputPlanReadyForReceive(MaterializerState& state, - const ProjectedInputTransferPlan& plan) { - bool hasRemoteFragment = false; - for (const ProjectedInputTransferFragment& fragment : plan.fragments) { - if (fragment.sourceCoreId == fragment.targetCoreId) - return false; - - hasRemoteFragment = true; - ArrayRef exchangeIds = state.communicationPlan.getProjectedInputExchanges(fragment); - if (exchangeIds.size() > 1) - return false; - std::optional exchangeId = - exchangeIds.empty() ? std::nullopt : std::optional(exchangeIds.front()); - if (!exchangeId) - return false; - if (!state.communicationExchangeStates[exchangeId->value].sendEmitted) - return false; - } - return hasRemoteFragment; -} - -static bool hasBoundReciprocalDifferentPayloadSendPending(MaterializerState& state, - const ExchangeDescriptor& receiveExchange) { - for (const ExchangeDescriptor& candidate : state.communicationPlan.getExchanges()) { - if (candidate.kind != CommunicationExchangeKind::ProjectedInput) - continue; - if (candidate.sourceCore != receiveExchange.targetCore || candidate.targetCore != receiveExchange.sourceCore) - continue; - if (candidate.payloadType == receiveExchange.payloadType) - continue; - - const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[candidate.id.value]; - if (emission.producerBound && !emission.sendEmitted) - return true; - } - return false; -} - -static LogicalResult materializeReadyProjectedInputReceiveForExchange(MaterializerState& state, - const ExchangeDescriptor& exchange, - Location loc) { - MaterializedClass& targetClass = state.classes[exchange.targetClass]; - for (auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(targetClass.id); - if (classIt == extractEntry.second.end()) - continue; - - bool containsExchange = llvm::any_of(classIt->second.fragments, [&](const ProjectedInputTransferFragment& fragment) { - return &fragment == exchange.fragment; - }); - if (!containsExchange) - continue; - - auto replacementIt = state.projectedExtractReplacements.find(extractEntry.first); - if (replacementIt != state.projectedExtractReplacements.end() - && replacementIt->second.find(targetClass.id) != replacementIt->second.end()) - return success(); - - if (!isRemoteProjectedInputPlanReadyForReceive(state, classIt->second)) - return success(); - if (hasBoundReciprocalDifferentPayloadSendPending(state, exchange)) - return success(); - - Operation* insertionAnchor = getProjectedReceiveInsertionAnchor(targetClass); - return materializeProjectedInputReceivesForClass(state, targetClass, extractEntry.first, loc, insertionAnchor); - } - return success(); -} - -LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef keys, - Value payload, - Location loc) { - if (keys.empty()) - return success(); - - DenseSet keySet; - for (ProducerKey key : keys) - keySet.insert(key); - - for (auto& extractEntry : state.projectedInputTransferPlans) { - for (auto& classEntry : extractEntry.second) { - ProjectedInputTransferPlan& plan = classEntry.second; - for (ProjectedInputTransferFragment& fragment : plan.fragments) { - if (!keySet.contains(fragment.producer)) - continue; - - if (fragment.sourceCoreId == fragment.targetCoreId) { - fragment.sendEmitted = true; - continue; - } - - std::optional exchangeId = getProjectedExchangeId(state, fragment, sourceClass.op); - if (!exchangeId) - return failure(); - if (state.communicationPlan.getExchange(*exchangeId).sourceClass != sourceClass.id) - continue; - if (failed(bindCommunicationPayload(state, *exchangeId, sourceClass, payload, sourceClass.op, loc))) - return failure(); - } - } - } - - FailureOr emittedBound = emitReadyCommunicationEventsForClass(state, sourceClass, loc); - if (failed(emittedBound)) - return failure(); - return success(); -} - -static LogicalResult emitProjectedInputSendForExchange(MaterializerState& state, - const ExchangeDescriptor& exchange, - CommunicationExchangeEmissionState& emission, - Location loc) { - if (!exchange.fragment) - return state.func.emitError("projected communication exchange is missing its fragment"); - if (!emission.producerBound || emission.sendEmitted) - return success(); - if (!emission.loc) - emission.loc = loc; - - MaterializedClass& sourceClass = state.classes[exchange.sourceClass]; - OpBuilder::InsertionGuard guard(state.rewriter); - if (emission.producerAnchor && emission.producerAnchor->getBlock() == sourceClass.body) - state.rewriter.setInsertionPointAfter(emission.producerAnchor); - else - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - - ProjectedInputTransferFragment* fragment = const_cast(exchange.fragment); - SmallVector fragments {fragment}; - if (failed(emitProjectedInputSendsAtCurrentInsertionPoint( - state, sourceClass, emission.producerPayload, fragments, *emission.loc))) - return failure(); - emission.sendEmitted = true; - emission.sendAnchor = fragment->sendEmitted && emission.producerAnchor ? emission.producerAnchor : nullptr; - fragment->sendEmitted = true; - return success(); -} - -static FailureOr emitReadyCommunicationEventsForClass(MaterializerState& state, - MaterializedClass& materializedClass, - Location loc) { - bool emittedAny = false; - for (const CommunicationEvent& event : state.communicationPlan.getEventsForClass(materializedClass.id)) { - const ExchangeDescriptor& exchange = state.communicationPlan.getExchange(event.exchangeId); - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchange.id]; - switch (event.kind) { - case CommunicationEventKind::Send: - if (!emission.producerBound || emission.sendEmitted) - break; - if (exchange.kind == CommunicationExchangeKind::ProjectedInput) { - bool wasEmitted = emission.sendEmitted; - if (failed(emitProjectedInputSendForExchange(state, exchange, emission, loc))) - return failure(); - emittedAny |= !wasEmitted && emission.sendEmitted; - } - else { - if (failed(emitPlannedSend(state, event.exchangeId, materializedClass, loc))) - return failure(); - emittedAny = true; - } - break; - case CommunicationEventKind::Receive: - if (!emission.sendEmitted || emission.receiveEmitted) - break; - if (exchange.kind == CommunicationExchangeKind::ProjectedInput) { - bool wasEmitted = emission.receiveEmitted; - if (failed(materializeReadyProjectedInputReceiveForExchange(state, exchange, loc))) - return failure(); - emittedAny |= !wasEmitted && emission.receiveEmitted; - } - else if (exchange.kind == CommunicationExchangeKind::OrdinaryValue - || exchange.kind == CommunicationExchangeKind::HostForward) { - if (materializedClass.isBatch) - break; - FailureOr received = emitPlannedReceive(state, event.exchangeId, materializedClass, loc); - if (failed(received)) - return failure(); - if (exchange.kind == CommunicationExchangeKind::OrdinaryValue) - state.availableValues.record(exchange.producer, exchange.targetClass, *received); - emittedAny = true; - } - break; - } - } - return emittedAny; -} - - -static FailureOr> -collectOrderedBatchProjectedInputFragments(MaterializedClass& targetClass, - const ProjectedInputTransferPlan& plan, - unsigned fragmentsPerLane) { - const size_t expectedFragmentCount = targetClass.cpus.size() * static_cast(fragmentsPerLane); - if (plan.fragments.size() != expectedFragmentCount) { - targetClass.op->emitError("projected input batch transfer did not collect one complete payload per lane") - << " expected=" << expectedFragmentCount << " actual=" << plan.fragments.size(); - return failure(); - } - - SmallVector fragments(expectedFragmentCount, nullptr); - for (const ProjectedInputTransferFragment& fragment : plan.fragments) { - if (fragment.targetLane >= targetClass.cpus.size()) { - targetClass.op->emitError("projected input batch transfer has an out-of-range target lane") - << " targetLane=" << fragment.targetLane << " laneCount=" << targetClass.cpus.size(); - return failure(); - } - if (fragment.ordinal >= expectedFragmentCount) { - targetClass.op->emitError("projected input batch transfer has an out-of-range ordinal") - << " ordinal=" << fragment.ordinal << " fragmentCount=" << expectedFragmentCount; - return failure(); - } - unsigned expectedOrdinal = fragment.targetLane * fragmentsPerLane + (fragment.ordinal % fragmentsPerLane); - if (fragment.ordinal != expectedOrdinal) { - targetClass.op->emitError("projected input batch transfer ordinal does not match target lane order") - << " ordinal=" << fragment.ordinal << " expected=" << expectedOrdinal; - return failure(); - } - if (fragments[fragment.ordinal]) { - targetClass.op->emitError("projected input batch transfer has duplicate fragment ordinal") - << " ordinal=" << fragment.ordinal; - return failure(); - } - fragments[fragment.ordinal] = &fragment; - } - - for (auto [index, fragment] : llvm::enumerate(fragments)) { - if (fragment) - continue; - targetClass.op->emitError("projected input batch transfer is missing fragment ordinal") << " ordinal=" << index; - return failure(); - } - return fragments; -} - -static SmallVector -getBatchProjectedInputPayloadSlot(ArrayRef fragments, - size_t laneCount, - unsigned fragmentsPerLane, - unsigned slot) { - SmallVector slotFragments; - slotFragments.reserve(laneCount); - for (size_t lane = 0; lane < laneCount; ++lane) - slotFragments.push_back(fragments[lane * fragmentsPerLane + slot]); - return slotFragments; -} - -static bool isLocalProjectedInputFragment(const ProjectedInputTransferFragment& fragment) { - return fragment.sourceCoreId == fragment.targetCoreId; -} - -enum class ProjectedInputExchangeDirection { - Local, - Remote -}; - -static constexpr StringLiteral projectedInputTransferAttr = "raptor.projected_input_transfer"; - -static void tagProjectedInputTransfer(Operation* op) { - if (!op) - return; - op->setAttr(projectedInputTransferAttr, UnitAttr::get(op->getContext())); -} - -static Operation* getProjectedReceiveInsertionAnchor(MaterializedClass& targetClass) { - Operation* lastProjectedTransfer = nullptr; - for (Operation& op : targetClass.body->without_terminator()) - if (op.hasAttr(projectedInputTransferAttr)) - lastProjectedTransfer = &op; - - if (!lastProjectedTransfer) - return targetClass.body->getTerminator(); - - Operation* next = lastProjectedTransfer->getNextNode(); - return next ? next : targetClass.body->getTerminator(); -} - -static FailureOr -getProjectedInputExchangeDirectionKind(MaterializerState&, - MaterializedClass& targetClass, - const ProjectedInputTransferFragment& fragment) { - (void)targetClass; - if (fragment.sourceCoreId == fragment.targetCoreId) - return ProjectedInputExchangeDirection::Local; - return ProjectedInputExchangeDirection::Remote; -} - -static void tagProjectedInputReceive(Value value, ProjectedInputExchangeDirection direction) { - if (direction == ProjectedInputExchangeDirection::Local) - return; - - auto receive = value.getDefiningOp(); - if (!receive) - return; - - tagProjectedInputTransfer(receive.getOperation()); -} - -static FailureOr getUniformProjectedInputSlotDirection( - MaterializerState& state, MaterializedClass& targetClass, ArrayRef fragments) { - if (fragments.empty()) - return targetClass.op->emitError("projected input batch transfer has an empty slot"), failure(); - - FailureOr direction = - getProjectedInputExchangeDirectionKind(state, targetClass, *fragments.front()); - if (failed(direction)) - return failure(); - for (const ProjectedInputTransferFragment* fragment : fragments.drop_front()) { - FailureOr fragmentDirection = - getProjectedInputExchangeDirectionKind(state, targetClass, *fragment); - if (failed(fragmentDirection)) - return failure(); - if (*fragmentDirection == *direction) - continue; - } - return *direction; -} - -static FailureOr resolveUniformLocalBatchProjectedInputSlot(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef fragments) { - std::optional slotValue; - for (const ProjectedInputTransferFragment* fragment : fragments) { - std::optional local = state.availableValues.lookup(state, fragment->producer, targetClass.id); - if (!local && !isProducerValueMaterialized(state, fragment->producer)) { - recordDeferredProducerValue(state, fragment->producer); - return failure(); - } - if (!local) - return targetClass.op->emitError("projected input batch transfer could not resolve local fragment") - << " sourceCore=" << fragment->sourceCoreId << " targetCore=" << fragment->targetCoreId - << " targetLane=" << fragment->targetLane << " ordinal=" << fragment->ordinal, - failure(); - if (slotValue && *slotValue != *local) - return targetClass.op->emitError( - "projected input batch transfer local slot is not represented by one batch SSA value") - << " ordinal=" << fragment->ordinal, - failure(); - slotValue = *local; - } - - if (!slotValue) - return targetClass.op->emitError("projected input batch transfer has an empty local slot"), failure(); - return *slotValue; -} - -struct ProjectedInputPlanMaterialization { - struct BatchSlotGroup { - unsigned slot = 0; - ProjectedInputExchangeDirection direction = ProjectedInputExchangeDirection::Local; - SmallVector lanes; - }; - - Operation* extractOp = nullptr; - const ProjectedInputTransferPlan* plan = nullptr; - SmallVector fragments; - Value payload; - Value laneArg; - SmallVector batchSlotGroups; - SmallVector localScalarFragments; - SmallVector remoteScalarFragments; -}; - -static FailureOr -prepareProjectedInputPlanMaterialization(MaterializerState& state, - MaterializedClass& targetClass, - Operation* extractOp, - const ProjectedInputTransferPlan& plan, - Location loc) { - if (failed(verifyProjectedFragmentLayout(targetClass.op, plan.layout))) - return failure(); - - ProjectedInputPlanMaterialization work; - work.extractOp = extractOp; - work.plan = &plan; - - if (targetClass.isBatch) { - auto batch = dyn_cast(targetClass.op); - if (!batch) - return targetClass.op->emitError("projected input transfer batch target must be a scheduled compute_batch"), - failure(); - std::optional laneArg = batch.getLaneArgument(); - if (!laneArg) - return targetClass.op->emitError("projected input transfer batch target is missing its lane argument"), - failure(); - - const unsigned fragmentsPerLane = plan.layout.fragmentsPerLogicalSlot; - if (fragmentsPerLane == 0 || plan.layout.payloadFragmentCount != fragmentsPerLane) - return targetClass.op->emitError("projected input batch transfer requires one logical payload per target lane"), - failure(); - - FailureOr> maybeFragments = - collectOrderedBatchProjectedInputFragments(targetClass, plan, fragmentsPerLane); - if (failed(maybeFragments)) - return failure(); - work.fragments = std::move(*maybeFragments); - work.laneArg = *laneArg; - - FailureOr payloadType = - getProjectedPayloadType(targetClass.op, plan.layout.fragmentType, fragmentsPerLane); - if (failed(payloadType)) - return failure(); - work.payload = - tensor::EmptyOp::create(state.rewriter, loc, payloadType->getShape(), payloadType->getElementType()).getResult(); - - for (unsigned slot = 0; slot < fragmentsPerLane; ++slot) { - SmallVector slotFragments = - getBatchProjectedInputPayloadSlot(work.fragments, targetClass.cpus.size(), fragmentsPerLane, slot); - FailureOr direction = - getUniformProjectedInputSlotDirection(state, targetClass, slotFragments); - if (failed(direction)) - return failure(); - ProjectedInputPlanMaterialization::BatchSlotGroup group; - group.slot = slot; - group.direction = *direction; - for (unsigned lane = 0; lane < targetClass.cpus.size(); ++lane) - group.lanes.push_back(lane); - work.batchSlotGroups.push_back(std::move(group)); - } - return work; - } - - work.fragments.reserve(plan.fragments.size()); - for (const ProjectedInputTransferFragment& fragment : plan.fragments) - work.fragments.push_back(&fragment); - llvm::sort(work.fragments, [](const ProjectedInputTransferFragment* lhs, const ProjectedInputTransferFragment* rhs) { - return lhs->ordinal < rhs->ordinal; - }); - - if (work.fragments.size() != plan.layout.payloadFragmentCount) - return targetClass.op->emitError("projected input transfer did not collect one complete logical payload") - << " expected=" << plan.layout.payloadFragmentCount << " actual=" << work.fragments.size(), - failure(); - for (auto [index, fragment] : llvm::enumerate(work.fragments)) - if (fragment->ordinal != index) - return targetClass.op->emitError("projected input transfer fragments are not ordinal-contiguous"), failure(); - - if (plan.layout.payloadFragmentCount != 1) { - FailureOr payloadType = - getProjectedPayloadType(targetClass.op, plan.layout.fragmentType, plan.layout.payloadFragmentCount); - if (failed(payloadType)) - return failure(); - work.payload = - tensor::EmptyOp::create(state.rewriter, loc, payloadType->getShape(), payloadType->getElementType()).getResult(); - } - - for (auto [index, fragment] : llvm::enumerate(work.fragments)) { - FailureOr direction = - getProjectedInputExchangeDirectionKind(state, targetClass, *fragment); - if (failed(direction)) - return failure(); - if (*direction == ProjectedInputExchangeDirection::Local) - work.localScalarFragments.push_back(index); - else - work.remoteScalarFragments.push_back(index); - } - - return work; -} - -static FailureOr materializeScalarProjectedInputFragment(MaterializerState& state, - MaterializedClass& targetClass, - const ProjectedInputTransferPlan& plan, - const ProjectedInputTransferFragment& fragment, - Location loc) { - FailureOr direction = - getProjectedInputExchangeDirectionKind(state, targetClass, fragment); - if (failed(direction)) - return failure(); - - if (*direction == ProjectedInputExchangeDirection::Local) { - std::optional local = state.availableValues.lookup(state, fragment.producer, targetClass.id); - if (!local && !isProducerValueMaterialized(state, fragment.producer)) { - recordDeferredProducerValue(state, fragment.producer); - return failure(); - } - if (!local) - return targetClass.op->emitError("projected input transfer could not resolve local fragment"), failure(); - return *local; - } - - if (!isProjectedRemoteFragmentReadyForReceive(state, fragment, targetClass.op)) - return failure(); - - std::optional exchangeId = getProjectedExchangeId(state, fragment, targetClass.op); - if (!exchangeId) - return failure(); - FailureOr received = emitPlannedReceive(state, *exchangeId, targetClass, loc); - if (failed(received)) - return failure(); - tagProjectedInputReceive(*received, *direction); - return *received; -} - -static FailureOr materializeBatchProjectedInputPayloadSlotGroupValue( - MaterializerState& state, - MaterializedClass& targetClass, - const ProjectedInputTransferPlan& plan, - ArrayRef slotFragments, - ArrayRef lanes, - Value laneArg, - Location loc) { - SmallVector groupFragments; - groupFragments.reserve(lanes.size()); - SmallVector laneToGroupIndex(targetClass.cpus.size(), 0); - for (auto [groupIndex, lane] : llvm::enumerate(lanes)) { - groupFragments.push_back(slotFragments[lane]); - laneToGroupIndex[lane] = static_cast(groupIndex); - } - - FailureOr direction = - getUniformProjectedInputSlotDirection(state, targetClass, groupFragments); - if (failed(direction)) - return failure(); - - if (*direction == ProjectedInputExchangeDirection::Local) - return resolveUniformLocalBatchProjectedInputSlot(state, targetClass, groupFragments); - - for (const ProjectedInputTransferFragment* fragment : groupFragments) - if (!isProjectedRemoteFragmentReadyForReceive(state, *fragment, targetClass.op)) - return failure(); - - SmallVector exchangeIds; - exchangeIds.reserve(groupFragments.size()); - for (const ProjectedInputTransferFragment* fragment : groupFragments) { - std::optional exchangeId = getProjectedExchangeId(state, *fragment, targetClass.op); - if (!exchangeId) - return failure(); - exchangeIds.push_back(*exchangeId); - } - - Value groupIndex = createIndexedIndexValue(state, targetClass.op, laneToGroupIndex, laneArg, loc); - FailureOr messages = - buildMessageVectorForExchangeIds(state, - exchangeIds, - CommunicationExchangeKind::ProjectedInput, - state.communicationPlan.getExchange(exchangeIds.front()).sourceClass, - targetClass.id, - plan.layout.fragmentType, - targetClass.op); - if (failed(messages)) - return failure(); - Value received = SpatChannelReceiveOp::create(state.rewriter, - loc, - plan.layout.fragmentType, - createIndexedChannelId(state, targetClass.op, *messages, groupIndex, loc), - createIndexedSourceCoreId(state, targetClass.op, *messages, groupIndex, loc), - createIndexedTargetCoreId(state, targetClass.op, *messages, groupIndex, loc)) - .getOutput(); - tagProjectedInputReceive(received, *direction); - Operation* receiveAnchor = nullptr; - if (Operation* receiveDef = received.getDefiningOp()) - receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, targetClass); - for (CommunicationExchangeId exchangeId : exchangeIds) { - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId.value]; - emission.receiveEmitted = true; - emission.receiveValue = received; - emission.receiveAnchor = receiveAnchor; - } - return received; -} - -static FailureOr insertProjectedBatchSlotGroup(MaterializerState& state, - MaterializedClass& targetClass, - ProjectedInputPlanMaterialization& work, - ArrayRef slotFragments, - ArrayRef lanes, - unsigned slot, - Location loc) { - const ProjectedInputTransferPlan& plan = *work.plan; - FailureOr fragmentValue = materializeBatchProjectedInputPayloadSlotGroupValue( - state, targetClass, plan, slotFragments, lanes, work.laneArg, loc); - if (failed(fragmentValue)) - return failure(); - - Value payloadOrdinal = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(slot)); - FailureOr offset = - scaleIndexByDim0SizeInClass(state, targetClass, payloadOrdinal, plan.layout.fragmentType.getDimSize(0), loc); - if (failed(offset)) - return failure(); - return createDim0InsertSliceInClass(state, targetClass, loc, *fragmentValue, work.payload, *offset); -} - -static LogicalResult materializeProjectedInputBatchSlotGroup( - MaterializerState& state, - MaterializedClass& targetClass, - ProjectedInputPlanMaterialization& work, - const ProjectedInputPlanMaterialization::BatchSlotGroup& group, - Location loc) { - const ProjectedInputTransferPlan& plan = *work.plan; - SmallVector slotFragments = - getBatchProjectedInputPayloadSlot( - work.fragments, targetClass.cpus.size(), plan.layout.fragmentsPerLogicalSlot, group.slot); - - if (group.lanes.size() != targetClass.cpus.size()) - return targetClass.op->emitError("projected input batch transfer produced a partial lane group"), failure(); - - FailureOr updated = - insertProjectedBatchSlotGroup(state, targetClass, work, slotFragments, group.lanes, group.slot, loc); - if (failed(updated)) - return failure(); - work.payload = *updated; - return success(); -} - -static LogicalResult materializeProjectedInputScalarFragment(MaterializerState& state, - MaterializedClass& targetClass, - ProjectedInputPlanMaterialization& work, - size_t index, - Location loc) { - const ProjectedInputTransferPlan& plan = *work.plan; - const ProjectedInputTransferFragment& fragment = *work.fragments[index]; - FailureOr received = materializeScalarProjectedInputFragment(state, targetClass, plan, fragment, loc); - if (failed(received)) - return failure(); - - if (plan.layout.payloadFragmentCount == 1) { - work.payload = *received; - return success(); - } - - Value offset = getOrCreateIndexConstant( - state.constantFolder, targetClass.op, static_cast(index) * plan.layout.fragmentType.getDimSize(0)); - FailureOr next = createDim0InsertSliceInClass(state, targetClass, loc, *received, work.payload, offset); - if (failed(next)) { - targetClass.op->emitError("failed to pack projected input transfer fragment") - << " targetClass=" << targetClass.id << " ordinal=" << fragment.ordinal - << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId - << " fragmentType=" << plan.layout.fragmentType; - return failure(); - } - work.payload = *next; - return success(); -} - -static LogicalResult preflightProjectedInputPlanReceives(MaterializerState& state, - MaterializedClass& targetClass, - const ProjectedInputTransferPlan& plan) { - for (const ProjectedInputTransferFragment& fragment : plan.fragments) { - if (isLocalProjectedInputFragment(fragment)) - continue; - std::optional exchangeId = getProjectedExchangeId(state, fragment, targetClass.op); - if (!exchangeId) - return failure(); - } - return success(); -} - -static LogicalResult preflightProjectedInputsForInstance(MaterializerState& state, - MaterializedClass& targetClass, - const ComputeInstance& instance) { - SmallVector requestedExtracts; - instance.op->walk([&](tensor::ExtractSliceOp extract) { - auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); - if (planIt == state.projectedInputTransferPlans.end()) - return; - if (planIt->second.find(targetClass.id) != planIt->second.end()) - requestedExtracts.push_back(extract.getOperation()); - }); - - for (Operation* requestedExtract : requestedExtracts) { - for (auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(targetClass.id); - if (classIt == extractEntry.second.end()) - continue; - - auto replacementIt = state.projectedExtractReplacements.find(extractEntry.first); - if (replacementIt != state.projectedExtractReplacements.end() - && replacementIt->second.find(targetClass.id) != replacementIt->second.end()) - continue; - if (extractEntry.first != requestedExtract) - continue; - - if (failed(preflightProjectedInputPlanReceives(state, targetClass, classIt->second))) - return failure(); - } - } - return success(); -} - -static FailureOr> -collectProjectedInputPlanMaterializationsForClass(MaterializerState& state, - MaterializedClass& targetClass, - Operation* requestedExtract, - Location loc) { - SmallVector works; - for (auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(targetClass.id); - if (classIt == extractEntry.second.end()) - continue; - - auto replacementIt = state.projectedExtractReplacements.find(extractEntry.first); - if (replacementIt != state.projectedExtractReplacements.end() - && replacementIt->second.find(targetClass.id) != replacementIt->second.end()) - continue; - if (extractEntry.first != requestedExtract) - continue; - - FailureOr work = - prepareProjectedInputPlanMaterialization(state, targetClass, extractEntry.first, classIt->second, loc); - if (failed(work)) - return failure(); - works.push_back(std::move(*work)); - } - return works; -} - -static LogicalResult materializeProjectedInputReceivesForClass(MaterializerState& state, - MaterializedClass& targetClass, - Operation* requestedExtract, - Location loc, - Operation* insertionAnchor = nullptr) { - OpBuilder::InsertionGuard guard(state.rewriter); - if (insertionAnchor) { - if (insertionAnchor->getBlock() != targetClass.body) - return targetClass.op->emitError("projected input transfer insertion anchor is outside the target class body"), - failure(); - state.rewriter.setInsertionPoint(insertionAnchor); - } - - FailureOr> maybeWorks = - collectProjectedInputPlanMaterializationsForClass(state, targetClass, requestedExtract, loc); - if (failed(maybeWorks)) - return failure(); - SmallVector& works = *maybeWorks; - if (works.empty()) - return success(); - - for (ProjectedInputPlanMaterialization& work : works) { - for (const auto& group : work.batchSlotGroups) - if (group.direction == ProjectedInputExchangeDirection::Local - && failed(materializeProjectedInputBatchSlotGroup(state, targetClass, work, group, loc))) - return failure(); - for (size_t index : work.localScalarFragments) - if (failed(materializeProjectedInputScalarFragment(state, targetClass, work, index, loc))) - return failure(); - } - - for (ProjectedInputPlanMaterialization& work : works) { - for (const auto& group : work.batchSlotGroups) - if (group.direction == ProjectedInputExchangeDirection::Remote - && failed(materializeProjectedInputBatchSlotGroup(state, targetClass, work, group, loc))) - return failure(); - for (size_t index : work.remoteScalarFragments) - if (failed(materializeProjectedInputScalarFragment(state, targetClass, work, index, loc))) - return failure(); - } - - for (const ProjectedInputPlanMaterialization& work : works) { - if (!work.payload) - return targetClass.op->emitError("projected input transfer did not produce a replacement payload") - << " targetClass=" << targetClass.id, - failure(); - state.projectedExtractReplacements[work.extractOp][targetClass.id] = - ProjectedExtractReplacement {work.payload, work.plan->layout}; - } - return success(); -} - -LogicalResult verifyCommunicationPlanFullyMaterialized(MaterializerState& state) { - size_t malformedCount = 0; - SmallVector malformed; - for (const ExchangeDescriptor& exchange : state.communicationPlan.getExchanges()) { - const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchange.id.value]; - bool bad = !emission.producerBound || !emission.sendEmitted || !emission.receiveEmitted || !emission.receiveValue; - if (emission.receiveValue && emission.receiveValue.getType() != exchange.payloadType) - bad = true; - if (!bad) - continue; - ++malformedCount; - if (malformed.size() < 8) - malformed.push_back(&exchange); - } - - if (malformedCount == 0) - return success(); - - InFlightDiagnostic diag = state.func.emitError("communication plan was not fully materialized"); - for (const ExchangeDescriptor* exchange : malformed) { - const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchange->id.value]; - diag << " [exchange=" << exchange->id.value << " kind=" << static_cast(exchange->kind) - << " sourceClass=" << exchange->sourceClass << " sourceCore=" << exchange->sourceCore - << " targetClass=" << exchange->targetClass << " targetCore=" << exchange->targetCore - << " channelId=" << exchange->channelId << " producerBound=" << emission.producerBound - << " sendEmitted=" << emission.sendEmitted << " receiveEmitted=" << emission.receiveEmitted - << " receiveType=" << (emission.receiveValue ? emission.receiveValue.getType() : Type()) << "]"; - } - if (malformedCount > malformed.size()) - diag << " additionalIncomplete=" << (malformedCount - malformed.size()); - return failure(); -} - -MessageVector filterMessageVector(const MessageVector& messages, ArrayRef indices) { - MessageVector filtered; - for (size_t index : indices) - filtered.append(messages.channelIds[index], messages.sourceCoreIds[index], messages.targetCoreIds[index]); - return filtered; -} - -SmallVector filterInt64Vector(ArrayRef values, ArrayRef indices) { - SmallVector filtered; - filtered.reserve(indices.size()); - for (size_t index : indices) - filtered.push_back(values[index]); - return filtered; -} - -template -SmallVector filterVector(ArrayRef values, ArrayRef indices) { - SmallVector filtered; - filtered.reserve(indices.size()); - for (size_t index : indices) - filtered.push_back(values[index]); - return filtered; -} - - -LogicalResult emitScalarSourceCommunication( - MaterializerState& state, MaterializedClass& sourceClass, ArrayRef keys, Value payload, Location loc) { - assert(!sourceClass.isBatch && "scalar-source communication expects a scalar source class"); - - for (ProducerKey key : keys) - state.availableValues.record(key, sourceClass.id, payload); - - SmallVector destinationClasses = collectDestinationClassesForKeys(state, keys); - for (ClassId destinationClass : destinationClasses) { - if (destinationClass == sourceClass.id) - continue; - for (ProducerKey key : keys) { - ArrayRef ids = - state.communicationPlan.getOrdinaryValueExchanges(key, destinationClass); - if (ids.empty()) - continue; - SmallVector boundIds; - for (CommunicationExchangeId id : ids) { - if (state.communicationPlan.getExchange(id).sourceClass != sourceClass.id) - continue; - if (failed(bindCommunicationPayload(state, id, sourceClass, payload, sourceClass.op, loc))) - return failure(); - boundIds.push_back(id); - } - if (!boundIds.empty()) - state.availableValues.recordDeferredExact(key, destinationClass, boundIds); - } - } - - FailureOr emitted = emitReadyCommunicationEventsForClass(state, sourceClass, loc); - if (failed(emitted)) - return failure(); - return success(); -} - -LogicalResult emitClassToClassCommunication(MaterializerState& state, - MaterializedClass& sourceClass, - MaterializedClass& targetClass, - ArrayRef keys, - Value payload, - Location loc) { - if (sourceClass.id == targetClass.id) { - if (keys.size() == 1) - state.availableValues.record(keys.front(), targetClass.id, payload); - else if (failed(recordLocalPackedSlotValue(state, targetClass, keys, payload))) - return failure(); - return success(); - } - - if (!sourceClass.isBatch) - return sourceClass.op->emitError("scalar-source communication must be emitted through the scalar fanout planner"); - - if (!targetClass.isBatch) { - if (keys.size() != sourceClass.cpus.size()) - return sourceClass.op->emitError("batch-to-scalar communication expects one producer key per source lane"); - - SmallVector exchangeIds; - for (ProducerKey key : keys) { - ArrayRef ids = - state.communicationPlan.getPackedScalarRunExchanges(key, targetClass.id); - if (ids.empty()) - return sourceClass.op->emitError("missing planned batch-to-scalar exchange"); - for (CommunicationExchangeId id : ids) { - if (failed(bindCommunicationPayload(state, id, sourceClass, payload, sourceClass.op, loc))) - return failure(); - exchangeIds.push_back(id); - } - } - - if (keys.size() == 1) { - state.availableValues.recordDeferredExact(keys.front(), targetClass.id, exchangeIds); - return emitReadyCommunicationEventsForClass(state, sourceClass, loc); - } - - FailureOr fragmentType = getPackedSlotFragmentType(payload, keys.size()); - if (failed(fragmentType)) - return sourceClass.op->emitError("batch-to-scalar communication could not recover packed fragment type"); - - PackedScalarRunValue packedRun; - packedRun.targetClass = targetClass.id; - packedRun.sourceOp = keys.front().instance.op; - packedRun.resultIndex = keys.front().resultIndex; - packedRun.kind = PackedScalarRunKind::DeferredReceive; - packedRun.fragmentType = *fragmentType; - packedRun.exchangeIds = std::move(exchangeIds); - packedRun.slots.reserve(keys.size()); - for (ProducerKey key : keys) { - PackedScalarRunSlot slot; - slot.keys.push_back(key); - packedRun.slots.push_back(std::move(slot)); - } - if (failed(validatePackedScalarRunMetadata(targetClass.op, packedRun))) - return failure(); - state.availableValues.recordPackedRun(std::move(packedRun)); - return emitReadyCommunicationEventsForClass(state, sourceClass, loc); - } - - if (sourceClass.cpus.size() != targetClass.cpus.size()) - return sourceClass.op->emitError( - "cannot materialize batch communication between equivalence classes of different sizes") - << " sourceClass=" << sourceClass.id << " sourceSize=" << sourceClass.cpus.size() - << " targetClass=" << targetClass.id << " targetSize=" << targetClass.cpus.size() - << " keyCount=" << keys.size(); - - if (keys.size() != sourceClass.cpus.size()) - return sourceClass.op->emitError("batch-to-batch communication expects one producer key per source lane"); - - auto fragmentType = dyn_cast(payload.getType()); - if (!fragmentType || !fragmentType.hasStaticShape() || fragmentType.getRank() == 0) - return sourceClass.op->emitError("batch-to-batch communication requires a static ranked lane payload"); - - Operation* sourceOp = keys.front().instance.op; - size_t resultIndex = keys.front().resultIndex; - for (ProducerKey key : keys) - if (key.instance.op != sourceOp || key.resultIndex != resultIndex || key.instance.laneCount != 1) - return sourceClass.op->emitError("batch-to-batch communication expects one-lane keys from one producer result"); - - SmallVector exchangeIds; - for (ProducerKey key : keys) { - ArrayRef ids = - state.communicationPlan.getIndexedBatchRunExchanges(key, targetClass.id); - if (ids.empty()) - return sourceClass.op->emitError("missing planned batch-to-batch exchange"); - for (CommunicationExchangeId id : ids) { - if (failed(bindCommunicationPayload(state, id, sourceClass, payload, sourceClass.op, loc))) - return failure(); - exchangeIds.push_back(id); - } - } - - IndexedBatchRunValue indexedRun; - indexedRun.targetClass = targetClass.id; - indexedRun.sourceOp = sourceOp; - indexedRun.resultIndex = resultIndex; - indexedRun.fragmentType = fragmentType; - indexedRun.exchangeIds = std::move(exchangeIds); - PackedScalarRunSlot slot; - slot.keys.append(keys.begin(), keys.end()); - indexedRun.slots.push_back(std::move(slot)); - state.availableValues.recordIndexedBatchRun(std::move(indexedRun)); - - return emitReadyCommunicationEventsForClass(state, sourceClass, loc); -} - -LogicalResult -setHostOutputValue(MaterializerState& state, MaterializedClass& sourceClass, Value originalOutput, Value payload) { - auto resultIt = sourceClass.hostOutputToResultIndex.find(originalOutput); - if (resultIt == sourceClass.hostOutputToResultIndex.end()) - return sourceClass.op->emitError("missing host result slot for materialized output") - << " ownerKind=" << (sourceClass.isBatch ? "batch" : "scalar") - << " hostOutputs=" << sourceClass.hostOutputs.size() << " originalDef=" - << (originalOutput.getDefiningOp() ? originalOutput.getDefiningOp()->getName().getStringRef() - : StringRef("")); - - unsigned resultIndex = resultIt->second; - state.hostReplacements[originalOutput] = sourceClass.op->getResult(resultIndex); - - if (!sourceClass.isBatch) { - auto yieldOp = dyn_cast(sourceClass.body->getTerminator()); - if (!yieldOp) - return sourceClass.op->emitError("expected spat.yield terminator in materialized compute"); - if (resultIndex >= yieldOp.getNumOperands()) - return sourceClass.op->emitError("host result index out of range for materialized compute"); - if (payload.getType() != originalOutput.getType()) - return sourceClass.op->emitError("cannot set scalar host output from fragment payload") - << " payloadType=" << payload.getType() << " outputType=" << originalOutput.getType(); - - state.rewriter.modifyOpInPlace(yieldOp, [&] { yieldOp->setOperand(resultIndex, payload); }); - return success(); - } - - auto batch = cast(sourceClass.op); - auto inParallelOp = dyn_cast(sourceClass.body->getTerminator()); - if (!inParallelOp) - return sourceClass.op->emitError("expected spat.in_parallel terminator in materialized compute_batch"); - - auto payloadType = dyn_cast(payload.getType()); - if (!payloadType || !payloadType.hasStaticShape()) - return sourceClass.op->emitError("host-facing compute_batch payload must be a static ranked tensor"); - - auto laneArg = batch.getLaneArgument(); - if (!laneArg) - return batch.emitOpError("expected compute_batch lane block argument while materializing batch output"); - - auto outputArg = batch.getOutputArgument(resultIndex); - if (!outputArg) - return batch.emitOpError("expected compute_batch output block argument while materializing batch output"); - - state.rewriter.setInsertionPointToStart(&inParallelOp.getRegion().front()); - createDim0ParallelInsertSlice(state, payload.getLoc(), payload, *outputArg, *laneArg); - return success(); -} - -LogicalResult -emitHostCommunication(MaterializerState& state, MaterializedClass& sourceClass, Value payload, Value originalOutput) { - if (!hasLiveExternalUseCached(state, originalOutput)) - return success(); - - auto ownerIt = state.hostOutputOwners.find(originalOutput); - if (ownerIt == state.hostOutputOwners.end()) - return sourceClass.op->emitError("missing host owner for live external output"); - - MaterializedClass& ownerClass = state.classes[ownerIt->second]; - if (sourceClass.id == ownerClass.id) - return setHostOutputValue(state, ownerClass, originalOutput, payload); - - ArrayRef ids = - state.communicationPlan.getHostForwardExchanges(originalOutput, sourceClass.id, ownerClass.id); - if (ids.empty()) - return sourceClass.op->emitError("missing planned host-forward exchange"); - for (CommunicationExchangeId id : ids) - if (failed(bindCommunicationPayload(state, id, sourceClass, payload, sourceClass.op, payload.getLoc()))) - return failure(); - FailureOr emitted = emitReadyCommunicationEventsForClass(state, sourceClass, payload.getLoc()); - if (failed(emitted)) - return failure(); - - if (sourceClass.isBatch && ownerClass.isBatch) { - if (sourceClass.cpus.size() != ownerClass.cpus.size()) - return sourceClass.op->emitError("batch host publication requires batch source/owner classes of equal size"); - if (ids.size() != ownerClass.cpus.size()) - return sourceClass.op->emitError("planned batch host-forward exchange count does not match owner lanes"); - - for (CommunicationExchangeId id : ids) - if (!state.communicationExchangeStates[id.value].sendEmitted) - return failure(); - FailureOr messages = - buildMessageVectorForExchangeIds(state, - ids, - CommunicationExchangeKind::HostForward, - sourceClass.id, - ownerClass.id, - payload.getType(), - ownerClass.op); - if (failed(messages)) - return failure(); - Value ownerPayload = appendReceive(state, ownerClass, payload.getType(), *messages, payload.getLoc()); - if (!ownerPayload) - return failure(); - Operation* receiveAnchor = nullptr; - if (Operation* receiveDef = ownerPayload.getDefiningOp()) - receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, ownerClass); - for (CommunicationExchangeId id : ids) { - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; - emission.receiveEmitted = true; - emission.receiveValue = ownerPayload; - emission.receiveAnchor = receiveAnchor; - } - return setHostOutputValue(state, ownerClass, originalOutput, ownerPayload); - } - - if (sourceClass.isBatch) - return sourceClass.op->emitError("batch host publication must be routed through the owning/projection-aware path"); - if (ownerClass.isBatch) - return ownerClass.op->emitError("generic host publication does not support batch host owners"); - if (payload.getType() != originalOutput.getType()) - return sourceClass.op->emitError("cannot forward fragment payload to scalar host owner") - << " payloadType=" << payload.getType() << " outputType=" << originalOutput.getType(); - - if (ids.size() != 1) - return sourceClass.op->emitError("scalar host forwarding expected one planned exchange"); - FailureOr ownerPayload = emitPlannedReceive(state, ids.front(), ownerClass, payload.getLoc()); - if (failed(ownerPayload)) - return failure(); - return setHostOutputValue(state, ownerClass, originalOutput, *ownerPayload); -} - -LogicalResult emitOutputFanout(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef keys, - Value payload, - Value originalOutput, - Location loc) { - SmallVector communicationKeys(keys.begin(), keys.end()); - for (ProducerKey projectedKey : collectProjectedProducerKeysForOutput(state, originalOutput)) - if (!llvm::is_contained(communicationKeys, projectedKey)) - communicationKeys.push_back(projectedKey); - - if (keys.empty() && communicationKeys.empty()) - return success(); - - if (!sourceClass.isBatch) { - if (!keys.empty()) - if (failed(emitScalarSourceCommunication(state, sourceClass, keys, payload, loc))) - return failure(); - if (failed(emitPlannedProjectedInputSendsForKeys(state, sourceClass, communicationKeys, payload, loc))) - return failure(); - - return emitHostCommunication(state, sourceClass, payload, originalOutput); - } - - bool recordedProjectedHostFragments = false; - if (hasLiveExternalUseCached(state, originalOutput)) { - FailureOr recorded = - recordProjectedScalarHostFragmentsFromPackedValue(state, sourceClass, keys, payload, originalOutput, loc); - if (failed(recorded)) - return failure(); - recordedProjectedHostFragments = *recorded; - } - - if (failed(emitPlannedProjectedInputSendsForKeys(state, sourceClass, communicationKeys, payload, loc))) - return failure(); - - for (ClassId destinationClass : collectDestinationClassesForKeys(state, keys)) { - SmallVector destinationKeys = - filterProducerKeysForDestination(state, keys, destinationClass); - if (destinationKeys.empty()) - continue; - if (sourceClass.isBatch && destinationKeys.size() != keys.size()) - return sourceClass.op->emitError( - "non-uniform batch fanout requires an explicit projected input transfer plan"); - if (failed(emitClassToClassCommunication( - state, sourceClass, state.classes[destinationClass], destinationKeys, payload, loc))) - return failure(); - } - - if (hasLiveExternalUseCached(state, originalOutput) && !recordedProjectedHostFragments) - return sourceClass.op->emitError("batch host publication requires explicit fragment assembly metadata"); - - if (!recordedProjectedHostFragments && failed(emitHostCommunication(state, sourceClass, payload, originalOutput))) - return failure(); - - if (keys.size() == 1) - state.availableValues.record(keys.front(), sourceClass.id, payload); - else if (!keys.empty() && failed(recordLocalPackedSlotValue(state, sourceClass, keys, payload))) - return failure(); - - return success(); -} - -struct DirectWholeBatchFragment { - ProducerKey key; - Value fragment; -}; - -enum class WholeBatchFragmentSourceKind { - DeferredReceive, - DeferredLocalCompute, - PackedValue, - DirectValue -}; - -struct WholeBatchFragmentGroup { - WholeBatchFragmentSourceKind kind = WholeBatchFragmentSourceKind::DirectValue; - RankedTensorType fragmentType; - SmallVector outputOffsets; - SmallVector exchangeIds; - Operation* sourceOp = nullptr; - size_t resultIndex = 0; - SmallVector sourceLanes; - Value packed; - RankedTensorType slotPackedType; - SmallVector packedRowOffsets; - SmallVector, 16> directFragments; - SmallVector redundantReceives; -}; - -enum class ProjectedWholeBatchFragmentSourceKind { - DeferredReceive, - PackedValue, - DirectValue -}; - -struct ProjectedWholeBatchDirectFragment { - Value fragment; - SmallVector offsets; - SmallVector sizes; - SmallVector strides; -}; - -struct ProjectedWholeBatchFragmentGroup { - ProjectedWholeBatchFragmentSourceKind kind = ProjectedWholeBatchFragmentSourceKind::DirectValue; - RankedTensorType fragmentType; - SmallVector, 4> offsetsByDim; - SmallVector, 4> sizesByDim; - SmallVector, 4> stridesByDim; - SmallVector exchangeIds; - SmallVector redundantOps; - Value packed; - RankedTensorType packedSourceType; - SmallVector packedIndices; - SmallVector directFragments; -}; - -struct WholeBatchAssemblyPlan { - RankedTensorType resultType; - int64_t rowsPerLane = 0; - uint32_t batchLaneCount = 0; - uint32_t coveredLaneCount = 0; - - SmallVector coveredLanes; - SmallVector packedRuns; - SmallVector directFragments; -}; - -ProjectedWholeBatchFragmentGroup filterProjectedDeferredReceiveGroup(const ProjectedWholeBatchFragmentGroup& group, - ArrayRef indices) { - ProjectedWholeBatchFragmentGroup filtered; - filtered.kind = group.kind; - filtered.fragmentType = group.fragmentType; - for (size_t index : indices) - filtered.exchangeIds.push_back(group.exchangeIds[index]); - filtered.offsetsByDim.resize(group.offsetsByDim.size()); - filtered.sizesByDim.resize(group.sizesByDim.size()); - filtered.stridesByDim.resize(group.stridesByDim.size()); - for (size_t dim = 0; dim < group.offsetsByDim.size(); ++dim) { - filtered.offsetsByDim[dim] = filterInt64Vector(group.offsetsByDim[dim], indices); - filtered.sizesByDim[dim] = filterInt64Vector(group.sizesByDim[dim], indices); - filtered.stridesByDim[dim] = filterInt64Vector(group.stridesByDim[dim], indices); - } - return filtered; -} - -bool wholeBatchLaneCovered(const WholeBatchAssemblyPlan& plan, uint32_t lane) { - return lane < plan.coveredLanes.size() && plan.coveredLanes[lane] != 0; -} - -bool wholeBatchRangeOverlaps(const WholeBatchAssemblyPlan& plan, uint32_t laneStart, uint32_t laneCount) { - if (laneCount == 0) - return false; - if (laneStart >= plan.coveredLanes.size()) - return false; - - uint32_t laneEnd = std::min(laneStart + laneCount, plan.coveredLanes.size()); - for (uint32_t lane = laneStart; lane < laneEnd; ++lane) - if (plan.coveredLanes[lane] != 0) - return true; - return false; -} - -void recordWholeBatchCoverage(WholeBatchAssemblyPlan& plan, uint32_t laneStart, uint32_t laneCount) { - assert(laneCount != 0 && "cannot cover an empty whole-batch range"); - assert(laneStart + laneCount <= plan.coveredLanes.size() && "whole-batch coverage out of bounds"); - - for (uint32_t lane = laneStart; lane < laneStart + laneCount; ++lane) { - if (plan.coveredLanes[lane] != 0) - continue; - plan.coveredLanes[lane] = 1; - ++plan.coveredLaneCount; - } -} - -bool localLaneRangeOverlaps(ArrayRef covered, uint32_t laneStart, uint32_t laneCount) { - if (laneCount == 0) - return false; - if (laneStart >= covered.size()) - return false; - - uint32_t laneEnd = std::min(laneStart + laneCount, covered.size()); - for (uint32_t lane = laneStart; lane < laneEnd; ++lane) - if (covered[lane] != 0) - return true; - return false; -} - -void markLocalLaneRangeCovered(MutableArrayRef covered, uint32_t laneStart, uint32_t laneCount) { - assert(laneStart + laneCount <= covered.size() && "local coverage out of bounds"); - for (uint32_t lane = laneStart; lane < laneStart + laneCount; ++lane) - covered[lane] = 1; -} - -LogicalResult -validateWholeBatchFragmentType(RankedTensorType resultType, RankedTensorType fragmentType, int64_t expectedRows) { - if (!fragmentType.hasStaticShape()) - return failure(); - if (fragmentType.getRank() != resultType.getRank()) - return failure(); - if (fragmentType.getDimSize(0) != expectedRows) - return failure(); - - for (int64_t dim = 1; dim < resultType.getRank(); ++dim) - if (fragmentType.getDimSize(dim) != resultType.getDimSize(dim)) - return failure(); - - return success(); -} - -// ----------------------------------------------------------------------------- -// Packed run tensor assembly helpers. -// ----------------------------------------------------------------------------- - -FailureOr insertFragmentIntoWholeBatch(MaterializerState& state, - MaterializedClass& targetClass, - Value fragment, - Value destination, - OpFoldResult firstOffset, - Location loc) { - return createDim0InsertSliceInClass(state, targetClass, loc, fragment, destination, firstOffset); -} - -FailureOr extractPackedSlotForIndex(MaterializerState& state, - MaterializedClass& targetClass, - Value packed, - RankedTensorType slotPackedType, - Value packedRowOffset, - Location loc) { - return createDim0ExtractSliceInClass( - state, targetClass, loc, packed, packedRowOffset, slotPackedType.getDimSize(0)); -} - -SmallVector flattenPackedScalarRunKeys(const PackedScalarRunValue& run) { - SmallVector keys; - for (const PackedScalarRunSlot& slot : run.slots) - llvm::append_range(keys, slot.keys); - return keys; -} - -bool packedScalarRunSlotsMatch(const PackedScalarRunValue& lhs, const PackedScalarRunValue& rhs) { - if (lhs.slots.size() != rhs.slots.size()) - return false; - - for (auto [lhsSlot, rhsSlot] : llvm::zip(lhs.slots, rhs.slots)) { - if (lhsSlot.keys.size() != rhsSlot.keys.size()) - return false; - if (!llvm::equal(lhsSlot.keys, rhsSlot.keys)) - return false; - } - - return true; -} - -std::optional getConstantIndexValue(Value value) { - APInt constant; - if (matchPattern(value, m_ConstantInt(&constant))) - return constant.getSExtValue(); - return std::nullopt; -} - -std::optional findPlannedExchangeForReceive(MaterializerState& state, - SpatChannelReceiveOp receive, - CommunicationExchangeKind expectedKind) { - std::optional channelId = getConstantIndexValue(receive.getChannelId()); - std::optional sourceCoreId = getConstantIndexValue(receive.getSourceCoreId()); - std::optional targetCoreId = getConstantIndexValue(receive.getTargetCoreId()); - if (!channelId || !sourceCoreId || !targetCoreId) - return std::nullopt; - for (const CommunicationExchange& exchange : state.communicationPlan.getExchanges()) { - if (exchange.kind == expectedKind && exchange.channelId == *channelId && exchange.sourceCore == *sourceCoreId - && exchange.targetCore == *targetCoreId) - return exchange.id; - } - return std::nullopt; -} - -std::optional findPlannedExchangeForReceive(MaterializerState& state, - SpatChannelReceiveOp receive) { - std::optional channelId = getConstantIndexValue(receive.getChannelId()); - std::optional sourceCoreId = getConstantIndexValue(receive.getSourceCoreId()); - std::optional targetCoreId = getConstantIndexValue(receive.getTargetCoreId()); - if (!channelId || !sourceCoreId || !targetCoreId) - return std::nullopt; - for (const CommunicationExchange& exchange : state.communicationPlan.getExchanges()) { - if (exchange.channelId == *channelId && exchange.sourceCore == *sourceCoreId && exchange.targetCore == *targetCoreId) - return exchange.id; - } - return std::nullopt; -} - -PackedScalarRunValue* findDeferredReceiveAlternativeForPackedRun(MaterializerState& state, - const MaterializedClass& targetClass, - const PackedScalarRunValue& run) { - WholeBatchAssemblyLookupKey lookupKey = - makeWholeBatchAssemblyLookupKey(run.sourceOp, run.resultIndex, targetClass.id); - ArrayRef runIndices = state.availableValues.getPackedRunIndicesForWholeBatch(lookupKey); - - for (size_t runIndex : runIndices) { - PackedScalarRunValue& candidate = state.availableValues.getPackedRun(runIndex); - if (&candidate == &run || candidate.kind != PackedScalarRunKind::DeferredReceive) - continue; - if (candidate.fragmentType != run.fragmentType) - continue; - if (!packedScalarRunSlotsMatch(candidate, run)) - continue; - return &candidate; - } - - return nullptr; -} - -FailureOr emitIndexedFragmentInsertLoop(MaterializerState& state, - MaterializedClass& targetClass, - Value destination, - int64_t itemCount, - IndexedFragmentBuilder buildFragment, - IndexedInsertOffsetBuilder buildOffset, - Location loc) { - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, itemCount); - Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1); - Operation* insertionPoint = targetClass.body->getTerminator(); - - state.rewriter.setInsertionPoint(insertionPoint); - auto loop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {destination}, - [&](OpBuilder&, Location, Value flatIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { - FailureOr fragment = buildFragment(flatIndex); - if (failed(fragment)) - return failure(); - FailureOr offset = buildOffset(flatIndex); - if (failed(offset)) - return failure(); - FailureOr next = - insertFragmentIntoWholeBatch(state, targetClass, *fragment, iterArgs.front(), *offset, loc); - if (failed(next)) - return failure(); - yielded.push_back(*next); - return success(); - }); - if (failed(loop)) - return failure(); - return loop->results.front(); -} - -FailureOr> cloneBatchBodyForLane(MaterializerState& state, - MaterializedClass& targetClass, - const ComputeInstance& instance, - Value laneValue, - ArrayRef resultIndices, - CloneIndexingContext indexing = {}); - -Value createBatchRunFlatIndex(MaterializerState& state, MaterializedClass& targetClass, Value slotIndex, Location loc); -FailureOr materializeIndexedBatchRunReceive(MaterializerState& state, - MaterializedClass& targetClass, - IndexedBatchRunValue& run, - Value runSlotIndex, - Location loc); - -} // namespace - -FailureOr materializeDeferredLocalPackedScalarRunValue(MaterializerState& state, - MaterializedClass& targetClass, - PackedScalarRunValue& run, - Location loc) { - SmallVector keys = flattenPackedScalarRunKeys(run); - if (keys.empty()) - return failure(); - FailureOr packedType = getPackedBatchTensorType(run.fragmentType, keys.size()); - if (failed(packedType)) - return targetClass.op->emitError("cannot materialize deferred local packed run for non-static ranked tensor"); - - SmallVector sourceLanes; - sourceLanes.reserve(keys.size()); - for (ProducerKey key : keys) { - if (key.instance.laneCount != 1) - return failure(); - sourceLanes.push_back(key.instance.laneStart); - } - - SmallVector resultIndices {run.resultIndex}; - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value init = - tensor::EmptyOp::create(state.rewriter, loc, packedType->getShape(), packedType->getElementType()).getResult(); - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(keys.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1); - - auto loop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {init}, - [&](OpBuilder&, Location, Value loopIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { - Value acc = iterArgs.front(); - Value sourceLane = createIndexedIndexValue(state, targetClass.op, sourceLanes, loopIndex, loc); - - FailureOr> produced = - cloneBatchBodyForLane(state, - targetClass, - keys.front().instance, - sourceLane, - resultIndices, - CloneIndexingContext {.runSlotIndex = std::nullopt, .projectionSlotIndex = loopIndex}); - if (failed(produced) || produced->size() != 1) - return failure(); - - FailureOr firstOffset = - scaleIndexByDim0SizeInClass(state, targetClass, loopIndex, run.fragmentType.getDimSize(0), loc); - if (failed(firstOffset)) - return failure(); - FailureOr next = - createDim0InsertSliceInClass(state, targetClass, loc, produced->front(), acc, *firstOffset); - if (failed(next)) - return failure(); - yielded.push_back(*next); - return success(); - }); - if (failed(loop)) - return failure(); - run.packed = loop->results.front(); - return run.packed; -} - -namespace { - -LogicalResult collectPackedRunsForWholeBatchInput(MaterializerState& state, - MaterializedClass& targetClass, - ProducerKey key, - WholeBatchAssemblyPlan& plan) { - WholeBatchAssemblyLookupKey lookupKey = makeWholeBatchAssemblyLookupKey(key, targetClass.id); - ArrayRef runIndices = state.availableValues.getPackedRunIndicesForWholeBatch(lookupKey); - - for (size_t runIndex : runIndices) { - PackedScalarRunValue& run = state.availableValues.getPackedRun(runIndex); - - SmallVector runKeys; - SmallVector runCoveredLanes(plan.batchLaneCount, 0); - - for (const PackedScalarRunSlot& slot : run.slots) { - for (ProducerKey fragmentKey : slot.keys) { - if (fragmentKey.instance.op != key.instance.op || fragmentKey.resultIndex != key.resultIndex) - return failure(); - - if (fragmentKey.instance.laneCount == 0) - return failure(); - - if (wholeBatchRangeOverlaps(plan, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount)) - return failure(); - - if (localLaneRangeOverlaps(runCoveredLanes, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount)) - return failure(); - - markLocalLaneRangeCovered(runCoveredLanes, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount); - runKeys.push_back(fragmentKey); - } - } - - if (runKeys.empty()) - continue; - - plan.packedRuns.push_back(&run); - - for (ProducerKey runKey : runKeys) - recordWholeBatchCoverage(plan, runKey.instance.laneStart, runKey.instance.laneCount); - } - - return success(); -} - -LogicalResult collectDirectFragmentsForWholeBatchInput(MaterializerState& state, - MaterializedClass& targetClass, - SpatComputeBatch batch, - ProducerKey key, - WholeBatchAssemblyPlan& plan) { - struct CandidateFragment { - ProducerKey key; - Value value; - }; - - uint32_t batchLaneCount = static_cast(batch.getLaneCount()); - if (plan.coveredLaneCount == plan.batchLaneCount) - return success(); - - WholeBatchAssemblyLookupKey lookupKey = makeWholeBatchAssemblyLookupKey(key, targetClass.id); - ArrayRef indexedFragments = - state.availableValues.getExactFragmentsForWholeBatch(lookupKey); - - SmallVector candidates; - candidates.reserve(indexedFragments.size()); - for (const AvailableValueStore::ExactBatchFragmentRecord& record : indexedFragments) { - ProducerKey candidateKey = record.key; - if (candidateKey.instance.op != batch.getOperation() || candidateKey.resultIndex != key.resultIndex - || candidateKey.instance.laneCount == 0) - continue; - if (!isTensorValueLocalToMaterializedClass(record.value, targetClass)) - continue; - if (wholeBatchRangeOverlaps(plan, candidateKey.instance.laneStart, candidateKey.instance.laneCount)) - continue; - - auto fragmentType = dyn_cast(record.value.getType()); - if (!fragmentType) - continue; - - int64_t expectedRows = plan.rowsPerLane * static_cast(candidateKey.instance.laneCount); - if (failed(validateWholeBatchFragmentType(plan.resultType, fragmentType, expectedRows))) - continue; - - candidates.push_back({candidateKey, record.value}); - } - - llvm::sort(candidates, [](const CandidateFragment& lhs, const CandidateFragment& rhs) { - if (lhs.key.instance.laneStart != rhs.key.instance.laneStart) - return lhs.key.instance.laneStart < rhs.key.instance.laneStart; - return lhs.key.instance.laneCount > rhs.key.instance.laneCount; - }); - - size_t candidateCursor = 0; - uint32_t lane = 0; - while (lane < batchLaneCount) { - while (lane < batchLaneCount && wholeBatchLaneCovered(plan, lane)) - ++lane; - - if (lane >= batchLaneCount) - break; - - while (candidateCursor < candidates.size() && candidates[candidateCursor].key.instance.laneStart < lane) - ++candidateCursor; - - size_t candidateIndex = candidateCursor; - const CandidateFragment* best = nullptr; - while (candidateIndex < candidates.size() && candidates[candidateIndex].key.instance.laneStart == lane) { - const CandidateFragment& candidate = candidates[candidateIndex]; - if (!wholeBatchRangeOverlaps(plan, lane, candidate.key.instance.laneCount)) { - best = &candidate; - break; - } - ++candidateIndex; - } - - if (!best) - return failure(); - - plan.directFragments.push_back({best->key, best->value}); - recordWholeBatchCoverage(plan, lane, best->key.instance.laneCount); - lane += best->key.instance.laneCount; - } - - return success(); -} - -LogicalResult collectWholeBatchFragmentGroups(MaterializerState& state, - MaterializedClass& targetClass, - const WholeBatchAssemblyPlan& plan, - SmallVectorImpl& groups) { - for (PackedScalarRunValue* run : plan.packedRuns) { - if (!run || run->slots.empty()) - continue; - if (run->fragmentType.getDimSize(0) != plan.rowsPerLane) - return failure(); - - if (run->kind == PackedScalarRunKind::Materialized && run->packed - && !isTensorValueLocalToMaterializedClass(run->packed, targetClass)) { - if (PackedScalarRunValue* deferredRun = findDeferredReceiveAlternativeForPackedRun(state, targetClass, *run)) - run = deferredRun; - else { - SmallVector keys = flattenPackedScalarRunKeys(*run); - std::optional packedKey = getContiguousProducerRangeForKeys(keys); - emitNonLocalMaterializedClassValueDiagnostic(targetClass.op, - targetClass, - "whole-batch assembly tried to reuse non-local PackedValue", - run->packed, - packedKey); - return failure(); - } - } - - if (run->kind == PackedScalarRunKind::DeferredReceive) { - if (failed(validatePackedScalarRunMetadata(targetClass.op, *run))) - return failure(); - - auto groupIt = llvm::find_if(groups, [&](const WholeBatchFragmentGroup& group) { - return group.kind == WholeBatchFragmentSourceKind::DeferredReceive && group.fragmentType == run->fragmentType; - }); - if (groupIt == groups.end()) { - WholeBatchFragmentGroup group; - group.kind = WholeBatchFragmentSourceKind::DeferredReceive; - group.fragmentType = run->fragmentType; - groups.push_back(std::move(group)); - groupIt = std::prev(groups.end()); - } - - llvm::append_range(groupIt->exchangeIds, run->exchangeIds); - for (const PackedScalarRunSlot& slot : run->slots) - for (ProducerKey fragmentKey : slot.keys) - groupIt->outputOffsets.push_back(static_cast(fragmentKey.instance.laneStart) * plan.rowsPerLane); - continue; - } - - if (run->kind == PackedScalarRunKind::DeferredLocalCompute) { - SmallVector keys = flattenPackedScalarRunKeys(*run); - if (keys.empty()) - return failure(); - - auto groupIt = llvm::find_if(groups, [&](const WholeBatchFragmentGroup& group) { - return group.kind == WholeBatchFragmentSourceKind::DeferredLocalCompute - && group.fragmentType == run->fragmentType && group.sourceOp == run->sourceOp - && group.resultIndex == run->resultIndex; - }); - if (groupIt == groups.end()) { - WholeBatchFragmentGroup group; - group.kind = WholeBatchFragmentSourceKind::DeferredLocalCompute; - group.fragmentType = run->fragmentType; - group.sourceOp = run->sourceOp; - group.resultIndex = run->resultIndex; - groups.push_back(std::move(group)); - groupIt = std::prev(groups.end()); - } - - for (ProducerKey fragmentKey : keys) { - if (fragmentKey.instance.laneCount != 1) - return failure(); - groupIt->sourceLanes.push_back(fragmentKey.instance.laneStart); - groupIt->outputOffsets.push_back(static_cast(fragmentKey.instance.laneStart) * plan.rowsPerLane); - } - continue; - } - - auto sourceBatch = dyn_cast_or_null(run->sourceOp); - if (!sourceBatch || !run->packed) - return failure(); - - auto getOrCreatePackedValueGroup = [&](RankedTensorType slotPackedType) -> WholeBatchFragmentGroup& { - auto groupIt = llvm::find_if(groups, [&](const WholeBatchFragmentGroup& group) { - return group.kind == WholeBatchFragmentSourceKind::PackedValue && group.fragmentType == run->fragmentType - && group.packed == run->packed && group.slotPackedType == slotPackedType; - }); - if (groupIt == groups.end()) { - WholeBatchFragmentGroup group; - group.kind = WholeBatchFragmentSourceKind::PackedValue; - group.fragmentType = run->fragmentType; - group.packed = run->packed; - group.slotPackedType = slotPackedType; - groups.push_back(std::move(group)); - groupIt = std::prev(groups.end()); - } - return *groupIt; - }; - - size_t flattenedIndexBase = 0; - for (auto [slotIndex, slot] : llvm::enumerate(run->slots)) { - std::optional contiguousKey = getContiguousProducerRangeForKeys(slot.keys); - if (contiguousKey) { - FailureOr slotPackedType = getPackedBatchTensorType(run->fragmentType, slot.keys.size()); - if (failed(slotPackedType)) - return failure(); - WholeBatchFragmentGroup& group = getOrCreatePackedValueGroup(*slotPackedType); - group.packedRowOffsets.push_back(static_cast(flattenedIndexBase) * plan.rowsPerLane); - group.outputOffsets.push_back(static_cast(contiguousKey->instance.laneStart) * plan.rowsPerLane); - flattenedIndexBase += slot.keys.size(); - continue; - } - - WholeBatchFragmentGroup& group = getOrCreatePackedValueGroup(run->fragmentType); - for (auto [keyIndex, fragmentKey] : llvm::enumerate(slot.keys)) { - group.packedRowOffsets.push_back(static_cast(flattenedIndexBase + keyIndex) * plan.rowsPerLane); - group.outputOffsets.push_back(static_cast(fragmentKey.instance.laneStart) * plan.rowsPerLane); - } - flattenedIndexBase += slot.keys.size(); - } - } - - auto getOrCreateDeferredReceiveGroup = [&](RankedTensorType fragmentType) -> WholeBatchFragmentGroup& { - auto groupIt = llvm::find_if(groups, [&](const WholeBatchFragmentGroup& group) { - return group.kind == WholeBatchFragmentSourceKind::DeferredReceive && group.fragmentType == fragmentType; - }); - if (groupIt == groups.end()) { - WholeBatchFragmentGroup group; - group.kind = WholeBatchFragmentSourceKind::DeferredReceive; - group.fragmentType = fragmentType; - groups.push_back(std::move(group)); - groupIt = std::prev(groups.end()); - } - return *groupIt; - }; - - auto getOrCreateDirectValueGroup = [&](RankedTensorType fragmentType) -> WholeBatchFragmentGroup& { - auto groupIt = llvm::find_if(groups, [&](const WholeBatchFragmentGroup& group) { - return group.kind == WholeBatchFragmentSourceKind::DirectValue && group.fragmentType == fragmentType; - }); - if (groupIt == groups.end()) { - WholeBatchFragmentGroup group; - group.kind = WholeBatchFragmentSourceKind::DirectValue; - group.fragmentType = fragmentType; - groups.push_back(std::move(group)); - groupIt = std::prev(groups.end()); - } - return *groupIt; - }; - - for (const DirectWholeBatchFragment& fragment : plan.directFragments) { - if (!isTensorValueLocalToMaterializedClass(fragment.fragment, targetClass)) { - emitNonLocalMaterializedClassValueDiagnostic(targetClass.op, - targetClass, - "whole-batch assembly tried to reuse non-local DirectValue", - fragment.fragment, - fragment.key); - return failure(); - } - - auto fragmentType = dyn_cast(fragment.fragment.getType()); - if (!fragmentType) - return failure(); - - int64_t outputOffset = static_cast(fragment.key.instance.laneStart) * plan.rowsPerLane; - - if (auto receive = fragment.fragment.getDefiningOp()) { - if (fragment.fragment.use_empty()) { - WholeBatchFragmentGroup& group = getOrCreateDeferredReceiveGroup(fragmentType); - if (std::optional exchangeId = - findPlannedExchangeForReceive(state, receive, CommunicationExchangeKind::PackedScalarRun)) { - group.exchangeIds.push_back(*exchangeId); - group.outputOffsets.push_back(outputOffset); - group.redundantReceives.push_back(receive.getOperation()); - continue; - } - } - } - - WholeBatchFragmentGroup& group = getOrCreateDirectValueGroup(fragmentType); - group.directFragments.push_back({fragment.fragment, outputOffset}); - } - - return success(); -} - -FailureOr emitWholeBatchFragmentGroup(MaterializerState& state, - MaterializedClass& targetClass, - Value destination, - const WholeBatchFragmentGroup& group, - Location loc) { - switch (group.kind) { - case WholeBatchFragmentSourceKind::DeferredReceive: { - FailureOr messages = - buildMessageVectorForExchangeIds(state, - group.exchangeIds, - CommunicationExchangeKind::PackedScalarRun, - state.communicationPlan.getExchange(group.exchangeIds.front()).sourceClass, - targetClass.id, - group.fragmentType, - targetClass.op); - if (failed(messages)) - return failure(); - Value updatedDestination = destination; - ReceiveMessagePartition partition = partitionReceiveMessagesInPlanOrder(*messages); - for (size_t index : partition.criticalStaticIndices) { - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value received = SpatChannelReceiveOp::create( - state.rewriter, - loc, - group.fragmentType, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->channelIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->sourceCoreIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->targetCoreIds[index])) - .getOutput(); - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[group.exchangeIds[index].value]; - emission.receiveEmitted = true; - emission.receiveValue = received; - if (Operation* receiveDef = received.getDefiningOp()) - emission.receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, targetClass); - FailureOr updated = createDim0InsertSliceInClass(state, - targetClass, - loc, - received, - updatedDestination, - getOrCreateIndexConstant( - state.constantFolder, targetClass.op, group.outputOffsets[index])); - if (failed(updated)) - return failure(); - updatedDestination = *updated; - } - - auto remainingMessages = filterMessageVector(*messages, partition.remainingIndices); - auto remainingOffsets = filterInt64Vector(group.outputOffsets, partition.remainingIndices); - auto remainingExchangeIds = - filterVector(ArrayRef(group.exchangeIds), partition.remainingIndices); - FailureOr updated = remainingMessages.empty() - ? FailureOr(updatedDestination) - : emitIndexedFragmentInsertLoop( - state, - targetClass, - updatedDestination, - static_cast(remainingOffsets.size()), - [&](Value flatIndex) -> FailureOr { - Value channelId = createIndexedChannelId(state, targetClass.op, remainingMessages, flatIndex, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); - Value received = - SpatChannelReceiveOp::create(state.rewriter, loc, group.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); - return received; - }, - [&](Value flatIndex) -> FailureOr { - return createIndexedIndexValue(state, targetClass.op, remainingOffsets, flatIndex, loc); - }, - loc); - if (failed(updated)) - return failure(); - if (!remainingMessages.empty()) { - for (CommunicationExchangeId exchangeId : remainingExchangeIds) { - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId.value]; - emission.receiveEmitted = true; - emission.receiveValue = *updated; - } - } - - for (Operation* receive : group.redundantReceives) - if (receive && receive->use_empty()) - receive->erase(); - - return *updated; - } - case WholeBatchFragmentSourceKind::DeferredLocalCompute: { - SmallVector resultIndices {group.resultIndex}; - return emitIndexedFragmentInsertLoop( - state, - targetClass, - destination, - static_cast(group.outputOffsets.size()), - [&](Value flatIndex) -> FailureOr { - Value sourceLane = createIndexedIndexValue(state, targetClass.op, group.sourceLanes, flatIndex, loc); - FailureOr> produced = - cloneBatchBodyForLane(state, - targetClass, - ComputeInstance {group.sourceOp, 0, 1}, - sourceLane, - resultIndices, - CloneIndexingContext {.runSlotIndex = flatIndex, .projectionSlotIndex = flatIndex}); - if (failed(produced) || produced->size() != 1) - return failure(); - return produced->front(); - }, - [&](Value flatIndex) -> FailureOr { - return createIndexedIndexValue(state, targetClass.op, group.outputOffsets, flatIndex, loc); - }, - loc); - } - case WholeBatchFragmentSourceKind::PackedValue: - return emitIndexedFragmentInsertLoop( - state, - targetClass, - destination, - static_cast(group.packedRowOffsets.size()), - [&](Value flatIndex) -> FailureOr { - Value packedRowOffset = - createIndexedIndexValue(state, targetClass.op, group.packedRowOffsets, flatIndex, loc); - FailureOr packed = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - group.packed, - targetClass.op, - "whole-batch packed fragment assembly tried to reuse a tensor from another materialized class"); - if (failed(packed)) - return failure(); - return extractPackedSlotForIndex(state, targetClass, *packed, group.slotPackedType, packedRowOffset, loc); - }, - [&](Value flatIndex) -> FailureOr { - return createIndexedIndexValue(state, targetClass.op, group.outputOffsets, flatIndex, loc); - }, - loc); - case WholeBatchFragmentSourceKind::DirectValue: - for (const auto& [fragment, offset] : group.directFragments) { - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - FailureOr localFragment = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - fragment, - targetClass.op, - "whole-batch direct fragment assembly tried to reuse a tensor from another materialized class"); - if (failed(localFragment)) - return failure(); - FailureOr updated = - createDim0InsertSliceInClass(state, - targetClass, - loc, - *localFragment, - destination, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, offset)); - if (failed(updated)) - return failure(); - destination = *updated; - } - return destination; - } - - return failure(); -} - -FailureOr emitProjectedWholeBatchFragmentInsertLoop(MaterializerState& state, - MaterializedClass& targetClass, - Value destination, - const ProjectedWholeBatchFragmentGroup& group, - llvm::function_ref(Value)> buildFragment, - Location loc) { - assert(group.fragmentType && "expected projected fragment type"); - assert(!group.offsetsByDim.empty() && "expected projected insert coordinates"); - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.offsetsByDim.front().size()); - Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - auto loop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {destination}, - [&](OpBuilder&, Location, Value flatIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { - FailureOr fragment = buildFragment(flatIndex); - if (failed(fragment)) - return failure(); - - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - unsigned rank = group.offsetsByDim.size(); - offsets.reserve(rank); - sizes.reserve(rank); - strides.reserve(rank); - for (unsigned dim = 0; dim < rank; ++dim) { - offsets.push_back(createIndexedOrStaticIndex(state, targetClass.op, group.offsetsByDim[dim], flatIndex, loc)); - sizes.push_back(createIndexedOrStaticIndex(state, targetClass.op, group.sizesByDim[dim], flatIndex, loc)); - strides.push_back(createIndexedOrStaticIndex(state, targetClass.op, group.stridesByDim[dim], flatIndex, loc)); - } - - Value updated = - tensor::InsertSliceOp::create(state.rewriter, loc, *fragment, iterArgs.front(), offsets, sizes, strides) - .getResult(); - yielded.push_back(updated); - return success(); - }); - if (failed(loop)) - return failure(); - return loop->results.front(); -} - -static bool isProducerValueMaterialized(MaterializerState& state, ProducerKey key) { - ComputeInstance scheduledInstance = getScheduledChunkForLogicalInstance(state, key.instance); - auto cpuIt = state.schedule.computeToCpuMap.find(scheduledInstance); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return false; - auto classIt = state.cpuToClass.find(cpuIt->second); - if (classIt == state.cpuToClass.end()) - return false; - auto rangeIt = state.scheduledInstanceToLogicalSlots.find(scheduledInstance); - if (rangeIt == state.scheduledInstanceToLogicalSlots.end()) - return false; - - LogicalSlotRange range = rangeIt->second; - if (isa(key.instance.op)) { - if (key.instance.laneStart < scheduledInstance.laneStart) - return false; - uint32_t localLaneStart = key.instance.laneStart - scheduledInstance.laneStart; - if (localLaneStart + key.instance.laneCount > range.count) - return false; - for (SlotId offset = localLaneStart; offset < localLaneStart + key.instance.laneCount; ++offset) - if (!state.materializedLogicalSlots.contains({classIt->second, range.start + offset})) - return false; - return true; - } - - if (range.count != 1) - return false; - if (!state.materializedLogicalSlots.contains({classIt->second, range.start})) - return false; - return true; -} - -static void unmarkProducerValueMaterialized(MaterializerState& state, ProducerKey key) { - ComputeInstance scheduledInstance = getScheduledChunkForLogicalInstance(state, key.instance); - auto cpuIt = state.schedule.computeToCpuMap.find(scheduledInstance); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return; - auto classIt = state.cpuToClass.find(cpuIt->second); - if (classIt == state.cpuToClass.end()) - return; - auto rangeIt = state.scheduledInstanceToLogicalSlots.find(scheduledInstance); - if (rangeIt == state.scheduledInstanceToLogicalSlots.end()) - return; - - LogicalSlotRange range = rangeIt->second; - if (!isa(key.instance.op)) { - if (range.count == 1) - state.materializedLogicalSlots.erase({classIt->second, range.start}); - return; - } - - if (key.instance.laneStart < scheduledInstance.laneStart) - return; - uint32_t localLaneStart = key.instance.laneStart - scheduledInstance.laneStart; - if (localLaneStart + key.instance.laneCount > range.count) - return; - for (SlotId offset = localLaneStart; offset < localLaneStart + key.instance.laneCount; ++offset) - state.materializedLogicalSlots.erase({classIt->second, range.start + offset}); -} - -static void recordDeferredProducerValue(MaterializerState& state, ProducerKey key) { - if (!llvm::is_contained(state.deferredProducerValues, key)) - state.deferredProducerValues.push_back(key); -} - -static int64_t getProducerSourceClassForDiagnostic(MaterializerState& state, ProducerKey key) { - ComputeInstance scheduled = getScheduledChunkForLogicalInstance(state, key.instance); - auto cpuIt = state.schedule.computeToCpuMap.find(scheduled); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return -1; - auto classIt = state.cpuToClass.find(cpuIt->second); - if (classIt == state.cpuToClass.end()) - return -1; - return static_cast(classIt->second); -} - -static bool hasDeferredMaterializationDependency(const MaterializerState& state) { - return !state.deferredProducerValues.empty(); -} - -std::optional getStaticProjectedPackedFragmentIndex(tensor::ExtractSliceOp extract) { - auto sourceType = dyn_cast(extract.getSource().getType()); - auto resultType = dyn_cast(extract.getResult().getType()); - if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() - || sourceType.getRank() == 0 || sourceType.getRank() != resultType.getRank()) - return std::nullopt; - - std::optional firstOffset = getConstantIndex(extract.getMixedOffsets().front()); - if (!firstOffset) - return std::nullopt; - - for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) { - std::optional offset = getConstantIndex(extract.getMixedOffsets()[dim]); - std::optional size = getConstantIndex(extract.getMixedSizes()[dim]); - std::optional stride = getConstantIndex(extract.getMixedStrides()[dim]); - if (!offset || !size || !stride || *stride != 1 || *size != resultType.getDimSize(dim)) - return std::nullopt; - if (dim != 0 && *offset != 0) - return std::nullopt; - } - - return *firstOffset; -} - -void appendProjectedInsertCoordinates(ProjectedWholeBatchFragmentGroup& group, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides) { - if (group.offsetsByDim.empty()) { - size_t rank = offsets.size(); - group.offsetsByDim.resize(rank); - group.sizesByDim.resize(rank); - group.stridesByDim.resize(rank); - } - - for (size_t dim = 0; dim < offsets.size(); ++dim) { - group.offsetsByDim[dim].push_back(offsets[dim]); - group.sizesByDim[dim].push_back(sizes[dim]); - group.stridesByDim[dim].push_back(strides[dim]); - } -} - -FailureOr buildWholeBatchAssemblyPlan(MaterializerState& state, - MaterializedClass& targetClass, - ProducerKey key, - Type resultType) { - auto batch = dyn_cast_or_null(key.instance.op); - auto resultTensorType = dyn_cast(resultType); - if (!batch || !resultTensorType || !resultTensorType.hasStaticShape() || resultTensorType.getRank() == 0) - return failure(); - - uint32_t batchLaneCount = static_cast(batch.getLaneCount()); - if (batchLaneCount == 0 || resultTensorType.getDimSize(0) % static_cast(batchLaneCount) != 0) - return failure(); - - WholeBatchAssemblyPlan plan; - plan.resultType = resultTensorType; - plan.rowsPerLane = resultTensorType.getDimSize(0) / static_cast(batchLaneCount); - plan.batchLaneCount = batchLaneCount; - plan.coveredLanes.assign(batchLaneCount, 0); - - if (failed(collectPackedRunsForWholeBatchInput(state, targetClass, key, plan))) - return failure(); - - if (plan.coveredLaneCount == plan.batchLaneCount) - return plan; - - if (failed(collectDirectFragmentsForWholeBatchInput(state, targetClass, batch, key, plan))) - return failure(); - - return plan; -} - -FailureOr emitWholeBatchAssemblyPlan(MaterializerState& state, - MaterializedClass& targetClass, - ProducerKey key, - WholeBatchAssemblyPlan& plan, - Location loc) { - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value result = - tensor::EmptyOp::create(state.rewriter, loc, plan.resultType.getShape(), plan.resultType.getElementType()) - .getResult(); - - SmallVector groups; - if (failed(collectWholeBatchFragmentGroups(state, targetClass, plan, groups))) - return failure(); - - for (const WholeBatchFragmentGroup& group : groups) { - FailureOr updated = emitWholeBatchFragmentGroup(state, targetClass, result, group, loc); - if (failed(updated)) - return failure(); - result = *updated; - } - - state.availableValues.record(key, targetClass.id, result); - return result; -} - -// ----------------------------------------------------------------------------- -// Run materialization helpers. -// ----------------------------------------------------------------------------- - -FailureOr materializeProjectedWholeBatchInputFromFragments( - MaterializerState& state, MaterializedClass& targetClass, ProducerKey key, Type resultType, Location loc) { - auto batch = dyn_cast_or_null(key.instance.op); - auto resultTensorType = dyn_cast(resultType); - if (!batch || !resultTensorType || !resultTensorType.hasStaticShape()) - return failure(); - - FailureOr projection = getBatchResultProjectionInsert(batch, key.resultIndex); - if (failed(projection)) - return failure(); - - auto laneArg = batch.getLaneArgument(); - if (!laneArg) - return batch.emitOpError("missing compute_batch lane argument while materializing projected whole-batch input"); - - uint32_t laneEnd = key.instance.laneStart + key.instance.laneCount; - if (laneEnd > static_cast(batch.getLaneCount())) - return failure(); - if (failed(preflightProjectedWholeBatchProducerReceives(state, targetClass, key))) - return failure(); - - if (targetClass.isBatch) { - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value result = - tensor::EmptyOp::create(state.rewriter, loc, resultTensorType.getShape(), resultTensorType.getElementType()) - .getResult(); - - for (uint32_t lane = key.instance.laneStart; lane < laneEnd; ++lane) { - ProducerKey laneKey = getBatchLaneProducerKey(batch, lane, 1, key.resultIndex); - std::optional fragment = state.availableValues.lookup(state, laneKey, targetClass.id); - if (!fragment && !isProducerValueMaterialized(state, laneKey)) { - recordDeferredProducerValue(state, laneKey); - return failure(); - } - if (!fragment) { - if (hasOrdinaryDestinationClass(state, laneKey, targetClass.id)) { - unmarkProducerValueMaterialized(state, laneKey); - recordDeferredProducerValue(state, laneKey); - return failure(); - } - return targetClass.op->emitError("projected whole-batch assembly is missing a lane fragment") - << " targetClass=" << targetClass.id << " lane=" << lane - << " sourceOp='" << batch->getName() << "' resultIndex=" << key.resultIndex - << " sourceClass=" << getProducerSourceClassForDiagnostic(state, laneKey) - << " ordinaryDestination=" << hasOrdinaryDestinationClass(state, laneKey, targetClass.id), - failure(); - } - - FailureOr> offsets = - evaluateStaticProjectionIndices(projection->getMixedOffsets(), *laneArg, lane); - FailureOr> sizes = - evaluateStaticProjectionIndices(projection->getMixedSizes(), *laneArg, lane); - FailureOr> strides = - evaluateStaticProjectionIndices(projection->getMixedStrides(), *laneArg, lane); - if (failed(offsets) || failed(sizes) || failed(strides)) - return failure(); - - SmallVector offsetAttrs; - SmallVector sizeAttrs; - SmallVector strideAttrs; - offsetAttrs.reserve(offsets->size()); - sizeAttrs.reserve(sizes->size()); - strideAttrs.reserve(strides->size()); - for (auto [offset, size, stride] : llvm::zip(*offsets, *sizes, *strides)) { - offsetAttrs.push_back(state.rewriter.getIndexAttr(offset)); - sizeAttrs.push_back(state.rewriter.getIndexAttr(size)); - strideAttrs.push_back(state.rewriter.getIndexAttr(stride)); - } - - FailureOr localFragment = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - *fragment, - targetClass.op, - "projected whole-batch assembly tried to reuse a tensor from another materialized class", - laneKey); - if (failed(localFragment)) - return failure(); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - result = - tensor::InsertSliceOp::create(state.rewriter, loc, *localFragment, result, offsetAttrs, sizeAttrs, strideAttrs) - .getResult(); - } - - state.availableValues.record(key, targetClass.id, result); - return result; - } - - SmallVector groups; - auto getOrCreateReceiveGroup = [&](RankedTensorType fragmentType) -> ProjectedWholeBatchFragmentGroup& { - auto groupIt = llvm::find_if(groups, [&](const ProjectedWholeBatchFragmentGroup& group) { - return group.kind == ProjectedWholeBatchFragmentSourceKind::DeferredReceive && group.fragmentType == fragmentType; - }); - if (groupIt == groups.end()) { - ProjectedWholeBatchFragmentGroup group; - group.kind = ProjectedWholeBatchFragmentSourceKind::DeferredReceive; - group.fragmentType = fragmentType; - groups.push_back(std::move(group)); - groupIt = std::prev(groups.end()); - } - return *groupIt; - }; - auto getOrCreatePackedGroup = [&](Value packed, - RankedTensorType packedSourceType, - RankedTensorType fragmentType) -> ProjectedWholeBatchFragmentGroup& { - auto groupIt = llvm::find_if(groups, [&](const ProjectedWholeBatchFragmentGroup& group) { - return group.kind == ProjectedWholeBatchFragmentSourceKind::PackedValue && group.fragmentType == fragmentType - && group.packed == packed && group.packedSourceType == packedSourceType; - }); - if (groupIt == groups.end()) { - ProjectedWholeBatchFragmentGroup group; - group.kind = ProjectedWholeBatchFragmentSourceKind::PackedValue; - group.fragmentType = fragmentType; - group.packed = packed; - group.packedSourceType = packedSourceType; - groups.push_back(std::move(group)); - groupIt = std::prev(groups.end()); - } - return *groupIt; - }; - auto getOrCreateDirectGroup = [&](RankedTensorType fragmentType) -> ProjectedWholeBatchFragmentGroup& { - auto groupIt = llvm::find_if(groups, [&](const ProjectedWholeBatchFragmentGroup& group) { - return group.kind == ProjectedWholeBatchFragmentSourceKind::DirectValue && group.fragmentType == fragmentType; - }); - if (groupIt == groups.end()) { - ProjectedWholeBatchFragmentGroup group; - group.kind = ProjectedWholeBatchFragmentSourceKind::DirectValue; - group.fragmentType = fragmentType; - groups.push_back(std::move(group)); - groupIt = std::prev(groups.end()); - } - return *groupIt; - }; - - for (uint32_t lane = key.instance.laneStart; lane < laneEnd; ++lane) { - ProducerKey laneKey = getBatchLaneProducerKey(batch, lane, 1, key.resultIndex); - FailureOr> offsets = - evaluateStaticProjectionIndices(projection->getMixedOffsets(), *laneArg, lane); - FailureOr> sizes = - evaluateStaticProjectionIndices(projection->getMixedSizes(), *laneArg, lane); - FailureOr> strides = - evaluateStaticProjectionIndices(projection->getMixedStrides(), *laneArg, lane); - if (failed(offsets) || failed(sizes) || failed(strides)) - return failure(); - - bool grouped = false; - if (std::optional exact = state.availableValues.lookupExact(laneKey, targetClass.id)) { - if (auto receive = exact->getDefiningOp()) { - auto fragmentType = dyn_cast(receive.getOutput().getType()); - if (fragmentType && receive.getOutput().use_empty()) { - ProjectedWholeBatchFragmentGroup& group = getOrCreateReceiveGroup(fragmentType); - if (std::optional exchangeId = findPlannedExchangeForReceive(state, receive)) { - group.exchangeIds.push_back(*exchangeId); - appendProjectedInsertCoordinates(group, *offsets, *sizes, *strides); - group.redundantOps.push_back(receive.getOperation()); - grouped = true; - } - } - } - } - - if (grouped) - continue; - - std::optional fragment = state.availableValues.lookup(state, laneKey, targetClass.id); - if (!fragment && !isProducerValueMaterialized(state, laneKey)) { - recordDeferredProducerValue(state, laneKey); - return failure(); - } - if (!fragment) { - if (hasOrdinaryDestinationClass(state, laneKey, targetClass.id)) { - unmarkProducerValueMaterialized(state, laneKey); - recordDeferredProducerValue(state, laneKey); - return failure(); - } - return targetClass.op->emitError("projected whole-batch assembly is missing a lane fragment") - << " targetClass=" << targetClass.id << " lane=" << lane - << " sourceOp='" << batch->getName() << "' resultIndex=" << key.resultIndex - << " sourceClass=" << getProducerSourceClassForDiagnostic(state, laneKey) - << " ordinaryDestination=" << hasOrdinaryDestinationClass(state, laneKey, targetClass.id), - failure(); - } - - auto fragmentType = dyn_cast(fragment->getType()); - if (!fragmentType) - return failure(); - - if (auto extract = fragment->getDefiningOp()) { - if (std::optional packedIndex = getStaticProjectedPackedFragmentIndex(extract)) { - auto packedSourceType = dyn_cast(extract.getSource().getType()); - if (packedSourceType) { - ProjectedWholeBatchFragmentGroup& group = - getOrCreatePackedGroup(extract.getSource(), packedSourceType, fragmentType); - group.packedIndices.push_back(*packedIndex); - appendProjectedInsertCoordinates(group, *offsets, *sizes, *strides); - group.redundantOps.push_back(extract.getOperation()); - continue; - } - } - } - - ProjectedWholeBatchFragmentGroup& group = getOrCreateDirectGroup(fragmentType); - group.directFragments.push_back({*fragment, *offsets, *sizes, *strides}); - } - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value result = - tensor::EmptyOp::create(state.rewriter, loc, resultTensorType.getShape(), resultTensorType.getElementType()) - .getResult(); - - for (const ProjectedWholeBatchFragmentGroup& group : groups) { - FailureOr updated = failure(); - switch (group.kind) { - case ProjectedWholeBatchFragmentSourceKind::DeferredReceive: { - const CommunicationExchange& firstExchange = state.communicationPlan.getExchange(group.exchangeIds.front()); - FailureOr messages = - buildMessageVectorForExchangeIds(state, - group.exchangeIds, - firstExchange.kind, - firstExchange.sourceClass, - targetClass.id, - group.fragmentType, - targetClass.op); - if (failed(messages)) - return failure(); - updated = result; - ReceiveMessagePartition partition = partitionReceiveMessagesInPlanOrder(*messages); - for (size_t index : partition.criticalStaticIndices) { - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value received = SpatChannelReceiveOp::create( - state.rewriter, - loc, - group.fragmentType, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->channelIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->sourceCoreIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->targetCoreIds[index])) - .getOutput(); - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[group.exchangeIds[index].value]; - emission.receiveEmitted = true; - emission.receiveValue = received; - if (Operation* receiveDef = received.getDefiningOp()) - emission.receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, targetClass); - - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - offsets.reserve(group.offsetsByDim.size()); - sizes.reserve(group.sizesByDim.size()); - strides.reserve(group.stridesByDim.size()); - for (size_t dim = 0; dim < group.offsetsByDim.size(); ++dim) { - offsets.push_back(state.rewriter.getIndexAttr(group.offsetsByDim[dim][index])); - sizes.push_back(state.rewriter.getIndexAttr(group.sizesByDim[dim][index])); - strides.push_back(state.rewriter.getIndexAttr(group.stridesByDim[dim][index])); - } - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - updated = tensor::InsertSliceOp::create(state.rewriter, loc, received, *updated, offsets, sizes, strides) - .getResult(); - } - if (!partition.remainingIndices.empty()) { - ProjectedWholeBatchFragmentGroup remainingGroup = - filterProjectedDeferredReceiveGroup(group, partition.remainingIndices); - MessageVector remainingMessages = filterMessageVector(*messages, partition.remainingIndices); - updated = emitProjectedWholeBatchFragmentInsertLoop( - state, - targetClass, - *updated, - remainingGroup, - [&](Value flatIndex) -> FailureOr { - Value channelId = createIndexedChannelId(state, targetClass.op, remainingMessages, flatIndex, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); - return SpatChannelReceiveOp::create( - state.rewriter, loc, remainingGroup.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); - }, - loc); - if (failed(updated)) - return failure(); - for (CommunicationExchangeId exchangeId : remainingGroup.exchangeIds) { - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId.value]; - emission.receiveEmitted = true; - emission.receiveValue = *updated; - } - } - break; - } - case ProjectedWholeBatchFragmentSourceKind::PackedValue: - updated = emitProjectedWholeBatchFragmentInsertLoop( - state, - targetClass, - result, - group, - [&](Value flatIndex) -> FailureOr { - SmallVector extractOffsets; - SmallVector extractSizes; - SmallVector extractStrides; - extractOffsets.reserve(group.packedSourceType.getRank()); - extractSizes.reserve(group.packedSourceType.getRank()); - extractStrides.reserve(group.packedSourceType.getRank()); - extractOffsets.push_back( - createIndexedOrStaticIndex(state, targetClass.op, group.packedIndices, flatIndex, loc)); - extractSizes.push_back(state.rewriter.getIndexAttr(1)); - extractStrides.push_back(state.rewriter.getIndexAttr(1)); - for (int64_t dim = 1; dim < group.packedSourceType.getRank(); ++dim) { - extractOffsets.push_back(state.rewriter.getIndexAttr(0)); - extractSizes.push_back(state.rewriter.getIndexAttr(group.packedSourceType.getDimSize(dim))); - extractStrides.push_back(state.rewriter.getIndexAttr(1)); - } - - FailureOr packed = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - group.packed, - targetClass.op, - "projected whole-batch packed fragment assembly tried to reuse a tensor from another materialized class"); - if (failed(packed)) - return failure(); - - return tensor::ExtractSliceOp::create( - state.rewriter, loc, group.fragmentType, *packed, extractOffsets, extractSizes, extractStrides) - .getResult(); - }, - loc); - break; - case ProjectedWholeBatchFragmentSourceKind::DirectValue: { - updated = result; - for (const ProjectedWholeBatchDirectFragment& fragment : group.directFragments) { - FailureOr localFragment = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - fragment.fragment, - targetClass.op, - "projected whole-batch assembly tried to reuse a tensor from another materialized class"); - if (failed(localFragment)) - return failure(); - - SmallVector offsetAttrs; - SmallVector sizeAttrs; - SmallVector strideAttrs; - for (auto [offset, size, stride] : llvm::zip(fragment.offsets, fragment.sizes, fragment.strides)) { - offsetAttrs.push_back(state.rewriter.getIndexAttr(offset)); - sizeAttrs.push_back(state.rewriter.getIndexAttr(size)); - strideAttrs.push_back(state.rewriter.getIndexAttr(stride)); - } - updated = tensor::InsertSliceOp::create( - state.rewriter, loc, *localFragment, *updated, offsetAttrs, sizeAttrs, strideAttrs) - .getResult(); - } - break; - } - } - if (failed(updated)) - return failure(); - result = *updated; - } - - for (const ProjectedWholeBatchFragmentGroup& group : groups) - for (Operation* redundantOp : group.redundantOps) - if (redundantOp && redundantOp->use_empty()) - redundantOp->erase(); - - state.availableValues.record(key, targetClass.id, result); - return result; -} - -FailureOr materializeWholeBatchInput( - MaterializerState& state, MaterializedClass& targetClass, ProducerKey key, Type resultType, Location loc) { - FailureOr plan = buildWholeBatchAssemblyPlan(state, targetClass, key, resultType); - if (succeeded(plan)) - return emitWholeBatchAssemblyPlan(state, targetClass, key, *plan, loc); - - return materializeProjectedWholeBatchInputFromFragments(state, targetClass, key, resultType, loc); -} - -FailureOr recordProjectedScalarHostFragmentsFromPackedRun(MaterializerState& state, - MaterializedClass& sourceClass, - SpatComputeBatch sourceBatch, - size_t resultIndex, - ArrayRef run, - Value packed, - RankedTensorType fragmentType, - Value originalOutput, - Location loc) { - if (!hasLiveExternalUseCached(state, originalOutput)) - return false; - if (packed.getType() == originalOutput.getType() || fragmentType == originalOutput.getType()) - return false; - - auto resultType = dyn_cast(originalOutput.getType()); - if (!resultType || !resultType.hasStaticShape()) - return false; - - FailureOr projection = getBatchResultProjectionInsert(sourceBatch, resultIndex); - if (failed(projection)) - return false; - - std::optional laneArg = sourceBatch.getLaneArgument(); - if (!laneArg) { - sourceBatch.emitOpError("missing compute_batch lane argument while recording projected host fragments"); - return failure(); - } - - FailureOr publicationResultIndex = appendScalarPublicationResult(state, sourceClass, packed, loc); - if (failed(publicationResultIndex)) - return failure(); - - int64_t fragmentElementCount = fragmentType.getNumElements(); - - for (auto [runIndex, slot] : llvm::enumerate(run)) { - if (slot.peers.size() != 1) { - sourceClass.op->emitError("projected scalar host output publication expects scalar one-peer run slots"); - return failure(); - } - - const ComputeInstance& peer = slot.peers.front(); - if (peer.op != sourceBatch.getOperation()) { - sourceClass.op->emitError("projected scalar host output run changed source operation"); - return failure(); - } - if (peer.laneCount != 1) { - sourceClass.op->emitError("projected scalar host output publication expects one logical lane per packed slot") - << " laneStart=" << peer.laneStart << " laneCount=" << peer.laneCount; - return failure(); - } - - FailureOr> offsets = - evaluateStaticProjectionIndices(projection->getMixedOffsets(), *laneArg, peer.laneStart); - FailureOr> sizes = - evaluateStaticProjectionIndices(projection->getMixedSizes(), *laneArg, peer.laneStart); - FailureOr> strides = - evaluateStaticProjectionIndices(projection->getMixedStrides(), *laneArg, peer.laneStart); - if (failed(offsets) || failed(sizes) || failed(strides)) { - sourceClass.op->emitError("failed to evaluate projected host output slice for logical lane ") << peer.laneStart; - return failure(); - } - - state.pendingProjectedHostOutputFragments.push_back(PendingProjectedHostOutputFragment { - originalOutput, - sourceClass.id, - ProducerKey {peer, resultIndex}, - *publicationResultIndex, - static_cast(runIndex), - static_cast(runIndex) * fragmentElementCount, - SmallVector(*offsets), - SmallVector(*sizes), - SmallVector(*strides), - peer.laneStart, - loc - }); - } - - return true; -} - -FailureOr recordProjectedScalarHostFragmentsFromPackedValue(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef keys, - Value packed, - Value originalOutput, - Location loc) { - if (!sourceClass.isBatch || keys.empty()) - return false; - if (!hasLiveExternalUseCached(state, originalOutput)) - return false; - if (packed.getType() == originalOutput.getType()) - return false; - - auto resultType = dyn_cast(originalOutput.getType()); - auto packedType = dyn_cast(packed.getType()); - auto sourceBatch = dyn_cast_or_null(keys.front().instance.op); - if (!resultType || !resultType.hasStaticShape() || !packedType || !packedType.hasStaticShape() || !sourceBatch) - return false; - if (keys.front().resultIndex >= static_cast(sourceBatch.getNumResults())) - return false; - - FailureOr projection = - getBatchResultProjectionInsert(sourceBatch, keys.front().resultIndex); - if (failed(projection)) - return false; - - auto laneArg = sourceBatch.getLaneArgument(); - if (!laneArg) - return sourceBatch.emitOpError("missing compute_batch lane argument while recording projected host fragments"); - - FailureOr> firstSizes = - evaluateStaticProjectionIndices(projection->getMixedSizes(), *laneArg, keys.front().instance.laneStart); - if (failed(firstSizes)) - return sourceClass.op->emitError("failed to evaluate projected host output slice for logical lane ") - << keys.front().instance.laneStart; - - SmallVector fragmentShape(*firstSizes); - auto fragmentType = RankedTensorType::get(fragmentShape, packedType.getElementType(), packedType.getEncoding()); - if (fragmentType == originalOutput.getType()) - return false; - - FailureOr publicationResultIndex = appendBatchPublicationResult(state, sourceClass, packed, loc); - if (failed(publicationResultIndex)) - return failure(); - - if (packedType != fragmentType) { - size_t keysPerPublishedPayload = keys.size(); - if (sourceClass.isBatch) { - if (sourceClass.cpus.empty() || keys.size() % sourceClass.cpus.size() != 0) - return sourceClass.op->emitError("projected packed host publication requires a stable per-lane key partition") - << " packedType=" << packedType << " fragmentType=" << fragmentType << " keyCount=" << keys.size() - << " laneCount=" << sourceClass.cpus.size(); - keysPerPublishedPayload = keys.size() / sourceClass.cpus.size(); - } - - if (packedType.getRank() == 0 || packedType.getDimSize(0) % static_cast(keysPerPublishedPayload) != 0) - return sourceClass.op->emitError("projected packed host publication requires either direct fragment operands or " - "evenly dim-0 packed fragments") - << " packedType=" << packedType << " fragmentType=" << fragmentType << " keyCount=" << keys.size(); - - SmallVector packedFragmentShape(packedType.getShape()); - packedFragmentShape[0] /= static_cast(keysPerPublishedPayload); - if (packedFragmentShape != fragmentShape) - return sourceClass.op->emitError( - "projected packed host publication fragment shape does not match projected slice size") - << " packedType=" << packedType << " fragmentType=" << fragmentType << " keyCount=" << keys.size(); - } - - int64_t payloadElementCount = packedType.getNumElements(); - int64_t fragmentElementCount = fragmentType.getNumElements(); - int64_t fragmentsPerPublishedPayload = payloadElementCount / fragmentElementCount; - size_t keysPerPublishedPayload = keys.size(); - if (sourceClass.isBatch) - keysPerPublishedPayload /= sourceClass.cpus.size(); - if (fragmentsPerPublishedPayload <= 0 - || static_cast(keysPerPublishedPayload) % fragmentsPerPublishedPayload != 0) - return sourceClass.op->emitOpError( - "projected packed host publication requires a deterministic publication packing layout") - << " packedType=" << packedType << " fragmentType=" << fragmentType << " keyCount=" << keys.size(); - - DenseMap publishedFragmentOrdinals; - for (auto [fragmentIndex, key] : llvm::enumerate(keys)) { - if (key.instance.op != sourceBatch.getOperation() || key.resultIndex != keys.front().resultIndex - || key.instance.laneCount != 1) - return sourceClass.op->emitError( - "projected packed host publication requires one-lane keys from one producer result"); - - FailureOr> offsets = - evaluateStaticProjectionIndices(projection->getMixedOffsets(), *laneArg, key.instance.laneStart); - FailureOr> sizes = - evaluateStaticProjectionIndices(projection->getMixedSizes(), *laneArg, key.instance.laneStart); - FailureOr> strides = - evaluateStaticProjectionIndices(projection->getMixedStrides(), *laneArg, key.instance.laneStart); - if (failed(offsets) || failed(sizes) || failed(strides)) - return sourceClass.op->emitError("failed to evaluate projected host output slice for logical lane ") - << key.instance.laneStart; - if (SmallVector(*sizes) != fragmentShape) - return sourceClass.op->emitError( - "projected packed host publication requires one operand to map to a consistent fragment shape"); - - FailureOr publishedLaneIndex = getPublicationLaneForProducerKey(state, sourceClass, key); - if (failed(publishedLaneIndex)) - return failure(); - int64_t ordinalWithinPublishedPayload = - sourceClass.isBatch ? publishedFragmentOrdinals[*publishedLaneIndex]++ : static_cast(fragmentIndex); - int64_t localFragmentOffsetWithinPublishedPayload = - (ordinalWithinPublishedPayload % fragmentsPerPublishedPayload) * fragmentElementCount; - - state.pendingProjectedHostOutputFragments.push_back(PendingProjectedHostOutputFragment { - originalOutput, - sourceClass.id, - key, - *publicationResultIndex, - static_cast(fragmentIndex), - static_cast(*publishedLaneIndex) * payloadElementCount + localFragmentOffsetWithinPublishedPayload, - SmallVector(*offsets), - SmallVector(*sizes), - SmallVector(*strides), - key.instance.laneStart, - loc}); - } - - return true; -} - -FailureOr resolveInputValue(MaterializerState& state, - MaterializedClass& targetClass, - Value input, - const ComputeInstance& consumerInstance, - CloneIndexingContext indexing, - bool allowWholeBatchFallback = true) { - auto rejectNonLocalResolvedValue = [&](Value resolved) -> FailureOr { - if (!isTensorValueDefinedInDifferentMaterializedClass(resolved, targetClass)) - return resolved; - - std::optional producer = getInputRequestProducerKey(input, consumerInstance); - emitNonLocalMaterializedClassValueDiagnostic( - consumerInstance.op, - targetClass, - "input resolution tried to reuse a tensor from another materialized class", - resolved, - producer); - return failure(); - }; - - if (isConstantLike(input)) - return input; - - if (std::optional producer = getInputRequestProducerKey(input, consumerInstance)) { - if (indexing.runSlotIndex) { - if (IndexedBatchRunValue* indexedRun = state.availableValues.lookupIndexedBatchRun(*producer, targetClass.id)) { - FailureOr received = materializeIndexedBatchRunReceive( - state, targetClass, *indexedRun, *indexing.runSlotIndex, consumerInstance.op->getLoc()); - if (failed(received)) - return failure(); - return rejectNonLocalResolvedValue(*received); - } - } - - if (std::optional value = state.availableValues.lookup(state, *producer, targetClass.id)) - return rejectNonLocalResolvedValue(*value); - - if (IndexedBatchRunValue* indexedRun = state.availableValues.lookupIndexedBatchRun(*producer, targetClass.id)) { - for (auto [slotIndex, slot] : llvm::enumerate(indexedRun->slots)) { - if (!llvm::is_contained(slot.keys, *producer)) - continue; - - Value received = Value(); - if (indexedRun->packed) { - FailureOr packed = - materializeIndexedBatchRunPackedValue(state, targetClass, *indexedRun, *producer, consumerInstance.op->getLoc()); - if (failed(packed)) - return failure(); - received = *packed; - } - else { - size_t laneCount = targetClass.cpus.size(); - MessageVector messages; - if (indexedRun->exchangeIds.empty()) - return targetClass.op->emitError("indexed batch run lookup is missing planned exchange ids"), failure(); - ArrayRef ids(indexedRun->exchangeIds.data() + slotIndex * laneCount, laneCount); - FailureOr plannedMessages = - buildMessageVectorForExchangeIds(state, - ids, - CommunicationExchangeKind::IndexedBatchRun, - state.communicationPlan.getExchange(ids.front()).sourceClass, - targetClass.id, - indexedRun->fragmentType, - targetClass.op); - if (failed(plannedMessages)) - return failure(); - messages = std::move(*plannedMessages); - received = appendReceive(state, targetClass, indexedRun->fragmentType, messages, consumerInstance.op->getLoc()); - if (!received) - return failure(); - } - state.availableValues.record(*producer, targetClass.id, received); - return rejectNonLocalResolvedValue(received); - } - } - - if (isWholeBatchProducerKey(*producer)) { - if (!allowWholeBatchFallback) { - if (!isProducerValueMaterialized(state, *producer)) { - recordDeferredProducerValue(state, *producer); - return failure(); - } - consumerInstance.op->emitError("failed to resolve compute_batch input without a compact fragment path") - << " from '" << producer->instance.op->getName() << "' laneStart=" << producer->instance.laneStart - << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex; - return failure(); - } - FailureOr wholeBatch = - materializeWholeBatchInput(state, targetClass, *producer, input.getType(), consumerInstance.op->getLoc()); - if (failed(wholeBatch) && !hasDeferredMaterializationDependency(state)) - consumerInstance.op->emitError("failed to materialize whole-batch input") - << " from '" << producer->instance.op->getName() << "' laneStart=" << producer->instance.laneStart - << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex; - if (failed(wholeBatch)) - return failure(); - return rejectNonLocalResolvedValue(*wholeBatch); - } - - if (!isProducerValueMaterialized(state, *producer)) { - recordDeferredProducerValue(state, *producer); - return failure(); - } - - ArrayRef deferredExactIds = - state.availableValues.getDeferredExactExchangeIds(*producer, targetClass.id); - if (!deferredExactIds.empty()) { - bool hasUnboundExchange = false; - for (CommunicationExchangeId id : deferredExactIds) { - if (id.value >= state.communicationExchangeStates.size() - || !state.communicationExchangeStates[id.value].producerBound) { - hasUnboundExchange = true; - break; - } - } - if (hasUnboundExchange) - unmarkProducerValueMaterialized(state, *producer); - recordDeferredProducerValue(state, *producer); - return failure(); - } - - std::optional producerClass; - if (auto cpuIt = state.schedule.computeToCpuMap.find(producer->instance); - cpuIt != state.schedule.computeToCpuMap.end()) { - if (auto classIt = state.cpuToClass.find(cpuIt->second); classIt != state.cpuToClass.end()) - producerClass = classIt->second; - } - consumerInstance.op->emitError("failed to resolve producer value") - << " from op '" << producer->instance.op->getName() << "' laneStart=" << producer->instance.laneStart - << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex - << " targetClass=" << targetClass.id - << " producerClass=" << (producerClass ? static_cast(*producerClass) : -1); - return failure(); - } - - if (isTensorValueDefinedInDifferentMaterializedClass(input, targetClass)) { - emitNonLocalMaterializedClassValueDiagnostic( - consumerInstance.op, - targetClass, - "input resolution tried to append a tensor from another materialized class as a normal input", - input); - return failure(); - } - - if (isOldComputeRegionBlockArgument(state, input)) { - InFlightDiagnostic diagnostic = - consumerInstance.op->emitError("cannot append old graph_compute region block argument as scheduled input"); - attachMaterializerValueOriginNote(diagnostic, input, "escaped input"); - return failure(); - } - - return appendInput(state, targetClass, input); -} - -bool hasPlannedProjectedInputTransfer(MaterializerState& state, - SpatComputeBatch batch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer, - ClassId classId) { - ProjectedBatchInputKey inputKey {batch.getOperation(), inputIndex}; - if (std::optional match = getProjectedInputSliceMatch(state, batch, inputIndex)) { - auto planIt = state.projectedInputTransferPlans.find(match->extract.getOperation()); - if (planIt != state.projectedInputTransferPlans.end()) { - auto classIt = planIt->second.find(classId); - if (classIt != planIt->second.end() && classIt->second.inputKey == inputKey) - return true; - } - } - - SmallVector producers = collectProducerKeysForDestinations(input, logicalConsumer); - for (ProducerKey producer : producers) { - auto producerIt = state.projectedTransfers.find(producer); - if (producerIt == state.projectedTransfers.end()) - continue; - - auto descriptorIt = producerIt->second.find(classId); - if (descriptorIt == producerIt->second.end()) - continue; - - if (descriptorIt->second.inputKey == inputKey) - return true; - } - - return false; -} - -bool hasProjectedInputReplacement(MaterializerState& state, - SpatComputeBatch batch, - unsigned inputIndex, - Value input, - ComputeInstance logicalConsumer, - ClassId classId) { - std::optional match = getProjectedInputSliceMatch(state, batch, inputIndex); - if (!match) - return false; - - auto replacementIt = state.projectedExtractReplacements.find(match->extract.getOperation()); - if (replacementIt != state.projectedExtractReplacements.end() - && replacementIt->second.find(classId) != replacementIt->second.end()) - return true; - - if (hasPlannedProjectedInputTransfer(state, batch, inputIndex, input, logicalConsumer, classId)) - return true; - - return getProjectedWholeBatchReplacementProducer(state, batch, inputIndex).has_value(); -} - -void mapWeights(MaterializerState& state, - MaterializedClass& targetClass, - const ComputeInstance& instance, - IRMapping& mapper) { - Operation* op = instance.op; - if (auto compute = dyn_cast(op)) { - for (auto [index, weight] : llvm::enumerate(compute.getWeights())) { - auto weightArg = compute.getWeightArgument(index); - assert(weightArg && "expected compute weight block argument"); - mapper.map(*weightArg, appendWeight(state, targetClass, weight)); - } - return; - } - - auto batch = cast(op); - for (auto [index, weight] : llvm::enumerate(batch.getWeights())) { - auto weightArg = batch.getWeightArgument(index); - assert(weightArg && "expected compute_batch weight block argument"); - mapper.map(*weightArg, appendWeight(state, targetClass, weight)); - } -} - -LogicalResult mapInputs(MaterializerState& state, - MaterializedClass& targetClass, - const ComputeInstance& instance, - IRMapping& mapper, - CloneIndexingContext indexing) { - auto mapResolvedInput = [&](Value resolved, - Value logicalInput, - Operation* anchor, - std::optional producer) -> FailureOr { - auto logicalType = dyn_cast(logicalInput.getType()); - auto physicalType = dyn_cast(resolved.getType()); - if (!logicalType || !physicalType || !logicalType.hasStaticShape() || !physicalType.hasStaticShape()) - return materializeTensorValueForMaterializedClassUse( - state, - targetClass, - resolved, - anchor, - "input mapping tried to reuse a tensor from another materialized class", - producer); - - return materializeLogicalTensorView(state, - targetClass, - MaterializedTensorView { - .value = resolved, - .logicalType = logicalType, - .physicalType = physicalType, - .layout = classifyMaterializedTensorLayout(logicalType, physicalType)}, - anchor, - "input mapping tried to reuse a tensor from another materialized class", - producer); - }; - - Operation* op = instance.op; - if (auto compute = dyn_cast(op)) { - for (auto [index, input] : llvm::enumerate(compute.getInputs())) { - if (hasProjectedComputeInputReplacement(state, compute, static_cast(index), targetClass.id) - && !hasUnreplacedComputeInputUse(state, compute, static_cast(index), targetClass.id)) - continue; - - FailureOr mapped = resolveInputValue(state, targetClass, input, instance, indexing); - if (failed(mapped)) { - if (hasDeferredMaterializationDependency(state)) - return failure(); - std::optional producer = getInputRequestProducerKey(input, instance); - auto diagnostic = compute.emitOpError("failed to resolve materialized compute input") << " #" << index; - if (producer) { - diagnostic << " from '" << producer->instance.op->getName() << "' laneStart=" << producer->instance.laneStart - << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex; - } - return failure(); - } - auto inputArg = compute.getInputArgument(index); - if (!inputArg) - return compute.emitOpError("expected compute input block argument while materializing inputs"); - std::optional producer = getInputRequestProducerKey(input, instance); - FailureOr remapped = mapResolvedInput(*mapped, input, compute, producer); - if (failed(remapped)) { - if (isTensorValueDefinedInDifferentMaterializedClass(*mapped, targetClass)) - emitNonLocalMaterializedClassValueDiagnostic( - compute, - targetClass, - "mapInputs tried to append a tensor from another materialized class", - *mapped, - producer); - return failure(); - } - mapper.map(*inputArg, *remapped); - } - return success(); - } - - auto batch = cast(op); - for (auto [index, input] : llvm::enumerate(batch.getInputs())) { - bool hasProjectedReplacement = - hasProjectedInputReplacement(state, batch, static_cast(index), input, instance, targetClass.id); - bool hasUnreplacedUse = - hasUnreplacedBatchInputUse(state, batch, static_cast(index), targetClass.id); - if (isDebugPaddedVggInput(input)) { - static unsigned debugCount = 0; - if (debugCount++ < 12) { - llvm::errs() << "debug padded mapInputs targetClass=" << targetClass.id - << " inputIndex=" << index - << " hasProjectedReplacement=" << hasProjectedReplacement - << " hasUnreplacedUse=" << hasUnreplacedUse - << " canUseProjectedLane=" - << canUseProjectedLaneInput(state, batch, static_cast(index), input, instance) - << " hasWholeBatchReplacement=" - << getProjectedWholeBatchReplacementProducer(state, batch, static_cast(index)).has_value() - << "\n"; - } - } - if (hasProjectedReplacement && !hasUnreplacedUse) - continue; - - FailureOr demand = - classifyComputeBatchInputDemand(state, targetClass, batch, static_cast(index), input, instance); - if (failed(demand)) - return batch.emitOpError("failed to classify materialized compute_batch input") << " #" << index; - - FailureOr mapped = failure(); - if (demand->kind == BatchInputDemandKind::WholeTensorBarrier) { - assert(demand->wholeTensorProducer && "whole-tensor input demand must carry a producer"); - mapped = materializeWholeBatchInput( - state, targetClass, *demand->wholeTensorProducer, input.getType(), batch.getOperation()->getLoc()); - if (failed(mapped) && !hasDeferredMaterializationDependency(state)) - return batch.emitOpError("failed to materialize whole-batch compute_batch input") - << " #" << index << " from '" << demand->wholeTensorProducer->instance.op->getName() - << "' laneStart=" << demand->wholeTensorProducer->instance.laneStart - << " laneCount=" << demand->wholeTensorProducer->instance.laneCount - << " resultIndex=" << demand->wholeTensorProducer->resultIndex; - if (failed(mapped)) - return failure(); - } - else { - mapped = resolveInputValue(state, targetClass, input, instance, indexing, /*allowWholeBatchFallback=*/false); - if (failed(mapped) && !hasDeferredMaterializationDependency(state)) - return batch.emitOpError("failed to resolve materialized compute_batch input"); - if (failed(mapped)) - return failure(); - } - - auto inputArg = batch.getInputArgument(index); - if (!inputArg) - return batch.emitOpError("expected compute_batch input block argument while materializing inputs"); - std::optional producer = getInputRequestProducerKey(input, instance); - if (demand->kind == BatchInputDemandKind::ProjectedFragment) { - (void)mapped; - continue; - } - FailureOr remapped = mapResolvedInput(*mapped, input, batch, producer); - if (failed(remapped)) { - if (isTensorValueDefinedInDifferentMaterializedClass(*mapped, targetClass)) - emitNonLocalMaterializedClassValueDiagnostic(batch, - targetClass, - "mapInputs tried to append a tensor from another materialized class", - *mapped, - producer); - return failure(); - } - mapper.map(*inputArg, *remapped); - } - return success(); -} - -SmallVector collectMappedBatchOutputs(SpatComputeBatch batch, IRMapping& mapper) { - SmallVector outputs(batch.getNumResults(), Value {}); - auto inParallel = dyn_cast_or_null(batch.getBody().front().getTerminator()); - if (!inParallel) - return outputs; - - for (Operation& op : inParallel.getRegion().front()) { - auto insert = dyn_cast(&op); - if (!insert) - continue; - - auto outputArg = dyn_cast(insert.getDest()); - if (!outputArg || outputArg.getOwner() != &batch.getBody().front()) - continue; - - auto firstOutputArg = batch.getOutputArgument(0); - if (!firstOutputArg) - return outputs; - unsigned resultIndex = outputArg.getArgNumber() - firstOutputArg->getArgNumber(); - if (resultIndex >= outputs.size()) - continue; - outputs[resultIndex] = mapper.lookupOrDefault(insert.getSource()); - } - - return outputs; -} - -SmallVector collectBatchOutputFragmentTypes(SpatComputeBatch batch) { - SmallVector types(batch.getNumResults(), Type {}); - auto inParallel = dyn_cast_or_null(batch.getBody().front().getTerminator()); - if (!inParallel) - return types; - - auto firstOutputArg = batch.getOutputArgument(0); - if (!firstOutputArg) - return types; - - for (Operation& op : inParallel.getRegion().front()) { - auto insert = dyn_cast(&op); - if (!insert) - continue; - - auto outputArg = dyn_cast(insert.getDest()); - if (!outputArg || outputArg.getOwner() != &batch.getBody().front()) - continue; - - unsigned resultIndex = outputArg.getArgNumber() - firstOutputArg->getArgNumber(); - if (resultIndex >= types.size()) - continue; - - types[resultIndex] = insert.getSource().getType(); - } - - return types; -} - -SmallVector& getBatchOutputFragmentTypesCached(MaterializerState& state, SpatComputeBatch batch) { - auto [it, inserted] = state.batchOutputFragmentTypesCache.try_emplace(batch.getOperation(), SmallVector {}); - if (inserted) - it->second = collectBatchOutputFragmentTypes(batch); - return it->second; -} - -ArrayRef getComputeInstanceOutputValuesCached(MaterializerState& state, ComputeInstance instance) { - auto [it, inserted] = state.computeInstanceOutputsCache.try_emplace(instance, SmallVector {}); - if (inserted) - it->second = getComputeInstanceOutputValues(instance); - return it->second; -} - -std::optional lookupProjectedExtractReplacement(MaterializerState& state, - MaterializedClass& targetClass, - tensor::ExtractSliceOp extract) { - auto replacementIt = state.projectedExtractReplacements.find(extract.getOperation()); - if (replacementIt == state.projectedExtractReplacements.end()) - return std::nullopt; - - auto classIt = replacementIt->second.find(targetClass.id); - if (classIt == replacementIt->second.end()) - return std::nullopt; - - return classIt->second; -} - -static std::optional getBatchInputKeyForExtractSource(tensor::ExtractSliceOp extract) { - auto sourceArg = dyn_cast(extract.getSource()); - if (!sourceArg) - return std::nullopt; - - auto batch = dyn_cast_or_null(sourceArg.getOwner()->getParentOp()); - if (!batch) - return std::nullopt; - - for (unsigned inputIndex = 0; inputIndex < batch.getInputs().size(); ++inputIndex) { - std::optional inputArg = batch.getInputArgument(inputIndex); - if (inputArg && *inputArg == sourceArg) - return ProjectedBatchInputKey {batch.getOperation(), inputIndex}; - } - - return std::nullopt; -} - -static std::optional -lookupProjectedExtractReplacementForInputKey(MaterializerState& state, - MaterializedClass& targetClass, - ProjectedBatchInputKey inputKey) { - auto lookupReplacementForExtract = [&](Operation* extractOp) -> std::optional { - auto replacementIt = state.projectedExtractReplacements.find(extractOp); - if (replacementIt == state.projectedExtractReplacements.end()) - return std::nullopt; - auto classIt = replacementIt->second.find(targetClass.id); - if (classIt == replacementIt->second.end()) - return std::nullopt; - return classIt->second; - }; - - for (auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(targetClass.id); - if (classIt != extractEntry.second.end() && classIt->second.inputKey == inputKey) { - if (std::optional replacement = lookupReplacementForExtract(extractEntry.first)) - return replacement; - } - } - - for (auto& producerEntry : state.projectedTransfers) { - auto classIt = producerEntry.second.find(targetClass.id); - if (classIt != producerEntry.second.end() && classIt->second.inputKey == inputKey) { - if (std::optional replacement = - lookupReplacementForExtract(classIt->second.extractOp)) - return replacement; - } - } - - return std::nullopt; -} - -static Operation* getProjectedInputPlanExtractForInputKey(MaterializerState& state, - MaterializedClass& targetClass, - ProjectedBatchInputKey inputKey) { - for (auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(targetClass.id); - if (classIt != extractEntry.second.end() && classIt->second.inputKey == inputKey) - return extractEntry.first; - } - return nullptr; -} - -static std::optional getProjectedInputKeyProducerForClass(MaterializerState& state, - MaterializedClass& targetClass, - ProjectedBatchInputKey inputKey) { - for (auto& producerEntry : state.projectedTransfers) { - auto classIt = producerEntry.second.find(targetClass.id); - if (classIt != producerEntry.second.end() && classIt->second.inputKey == inputKey) - return producerEntry.first; - } - - for (auto& extractEntry : state.projectedInputTransferPlans) { - auto classIt = extractEntry.second.find(targetClass.id); - if (classIt == extractEntry.second.end() || !(classIt->second.inputKey == inputKey) - || classIt->second.fragments.empty()) - continue; - return classIt->second.fragments.front().producer; - } - - return std::nullopt; -} - -FailureOr> getOrMaterializeProjectedExtractReplacement( - MaterializerState& state, - MaterializedClass& targetClass, - tensor::ExtractSliceOp extract, - Operation* insertionAnchor = nullptr) { - if (std::optional replacement = - lookupProjectedExtractReplacement(state, targetClass, extract)) - return replacement; - - auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); - if (planIt != state.projectedInputTransferPlans.end()) { - auto classIt = planIt->second.find(targetClass.id); - if (classIt != planIt->second.end()) { - if (failed(materializeProjectedInputReceivesForClass( - state, targetClass, extract.getOperation(), extract.getLoc(), insertionAnchor))) { - if (!hasDeferredMaterializationDependency(state)) - targetClass.op->emitError("failed to materialize projected input transfer replacements for class") - << " targetClass=" << targetClass.id; - return failure(); - } - - if (std::optional replacement = - lookupProjectedExtractReplacement(state, targetClass, extract)) - return replacement; - - targetClass.op->emitError("projected input transfer class materialization did not produce requested replacement") - << " targetClass=" << targetClass.id << " fragmentCount=" << classIt->second.fragments.size() - << " payloadFragmentCount=" << classIt->second.layout.payloadFragmentCount; - return failure(); - } - } - - if (std::optional inputKey = getBatchInputKeyForExtractSource(extract)) { - if (std::optional replacement = - lookupProjectedExtractReplacementForInputKey(state, targetClass, *inputKey)) - return replacement; - if (Operation* planExtract = getProjectedInputPlanExtractForInputKey(state, targetClass, *inputKey)) { - if (failed(materializeProjectedInputReceivesForClass( - state, targetClass, planExtract, extract.getLoc(), insertionAnchor))) { - if (!hasDeferredMaterializationDependency(state)) - targetClass.op->emitError("failed to materialize projected input key replacement for class") - << " targetClass=" << targetClass.id; - return failure(); - } - if (std::optional replacement = - lookupProjectedExtractReplacementForInputKey(state, targetClass, *inputKey)) - return replacement; - } - - if (std::optional producer = getProjectedInputKeyProducerForClass(state, targetClass, *inputKey)) { - if (getProducerSourceClassForDiagnostic(state, *producer) == static_cast(targetClass.id)) - return std::optional {}; - if (isProducerValueMaterialized(state, *producer) - && hasOrdinaryDestinationClass(state, *producer, targetClass.id)) - unmarkProducerValueMaterialized(state, *producer); - recordDeferredProducerValue(state, *producer); - return failure(); - } - } - - return std::optional {}; -} - -Operation* getTopLevelMaterializedClassBodyOp(Operation* nestedOp, MaterializedClass& targetClass) { - if (!nestedOp) - return nullptr; - - Operation* current = nestedOp; - while (Operation* parent = current->getParentOp()) { - if (parent == targetClass.op) - return current; - current = parent; - } - - return nullptr; -} - -void moveNewTopLevelOpsBeforeAnchor(MaterializedClass& targetClass, Operation* anchor) { - if (!anchor || anchor->getBlock() != targetClass.body) - return; - - SmallVector insertedOps; - auto it = std::next(anchor->getIterator()); - auto end = targetClass.body->getTerminator()->getIterator(); - for (; it != end; ++it) - insertedOps.push_back(&*it); - - for (Operation* op : insertedOps) - op->moveBefore(anchor); -} - -void moveDefinitionBeforeAnchorWithSameBlockDeps(Operation* def, Operation* anchor, DenseSet& moved) { - if (!def || !anchor || def == anchor || def->getBlock() != anchor->getBlock()) - return; - if (!moved.insert(def).second) - return; - - for (Value operand : def->getOperands()) - moveDefinitionBeforeAnchorWithSameBlockDeps(operand.getDefiningOp(), anchor, moved); - - if (anchor->isBeforeInBlock(def)) - def->moveBefore(anchor); -} - -void moveDefinitionBeforeAnchorWithSameBlockDeps(Operation* def, Operation* anchor) { - DenseSet moved; - moveDefinitionBeforeAnchorWithSameBlockDeps(def, anchor, moved); -} - -FailureOr materializeProjectedWholeBatchExtractReplacement(MaterializerState& state, - MaterializedClass& targetClass, - tensor::ExtractSliceOp extract, - ProducerKey producer, - IRMapping* mapper) { - OpBuilder::InsertPoint replacementPoint = state.rewriter.saveInsertionPoint(); - Operation* topLevelAnchor = getTopLevelMaterializedClassBodyOp(extract.getOperation(), targetClass); - - FailureOr fullSource = - materializeWholeBatchInput(state, targetClass, producer, extract.getSource().getType(), extract.getLoc()); - if (failed(fullSource)) - return failure(); - - moveNewTopLevelOpsBeforeAnchor(targetClass, topLevelAnchor); - if (topLevelAnchor) - state.rewriter.setInsertionPoint(extract); - else - state.rewriter.restoreInsertionPoint(replacementPoint); - - auto remapFoldResult = [&](OpFoldResult value) -> FailureOr { - if (auto mappedValue = dyn_cast_if_present(value)) { - FailureOr localized = rematerializeIndexValueInClass( - state, targetClass, mapper ? mapper->lookupOrDefault(mappedValue) : mappedValue, extract.getLoc(), mapper); - if (failed(localized)) - return failure(); - return OpFoldResult(*localized); - } - return value; - }; - - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - offsets.reserve(extract.getMixedOffsets().size()); - sizes.reserve(extract.getMixedSizes().size()); - strides.reserve(extract.getMixedStrides().size()); - - for (OpFoldResult value : extract.getMixedOffsets()) { - FailureOr localized = remapFoldResult(value); - if (failed(localized)) - return failure(); - offsets.push_back(*localized); - } - for (OpFoldResult value : extract.getMixedSizes()) { - FailureOr localized = remapFoldResult(value); - if (failed(localized)) - return failure(); - sizes.push_back(*localized); - } - for (OpFoldResult value : extract.getMixedStrides()) { - FailureOr localized = remapFoldResult(value); - if (failed(localized)) - return failure(); - strides.push_back(*localized); - } - - auto resultType = cast(extract.getType()); - return extractStaticSliceOrIdentity( - state.rewriter, extract.getLoc(), *fullSource, resultType, offsets, sizes, strides); -} - -FailureOr> materializeLocalProjectedInputExtractReplacement(MaterializerState& state, - MaterializedClass& targetClass, - tensor::ExtractSliceOp extract, - IRMapping* mapper) { - std::optional inputKey = getBatchInputKeyForExtractSource(extract); - if (!inputKey) - return std::optional {}; - - std::optional producer = getProjectedInputKeyProducerForClass(state, targetClass, *inputKey); - if (!producer) - return std::optional {}; - - if (!isProducerValueMaterialized(state, *producer)) { - recordDeferredProducerValue(state, *producer); - return failure(); - } - - std::optional fullSource = state.availableValues.lookup(state, *producer, targetClass.id); - if (!fullSource) - return std::optional {}; - - FailureOr localizedSource = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - *fullSource, - extract.getOperation(), - "local projected input replacement tried to reuse a tensor from another materialized class", - *producer, - mapper); - if (failed(localizedSource)) - return failure(); - - auto remapFoldResult = [&](OpFoldResult value) -> FailureOr { - if (auto mappedValue = dyn_cast_if_present(value)) { - FailureOr localized = rematerializeIndexValueInClass( - state, targetClass, mapper ? mapper->lookupOrDefault(mappedValue) : mappedValue, extract.getLoc(), mapper); - if (failed(localized)) - return failure(); - return OpFoldResult(*localized); - } - return value; - }; - - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - offsets.reserve(extract.getMixedOffsets().size()); - sizes.reserve(extract.getMixedSizes().size()); - strides.reserve(extract.getMixedStrides().size()); - - for (OpFoldResult value : extract.getMixedOffsets()) { - FailureOr localized = remapFoldResult(value); - if (failed(localized)) - return failure(); - offsets.push_back(*localized); - } - for (OpFoldResult value : extract.getMixedSizes()) { - FailureOr localized = remapFoldResult(value); - if (failed(localized)) - return failure(); - sizes.push_back(*localized); - } - for (OpFoldResult value : extract.getMixedStrides()) { - FailureOr localized = remapFoldResult(value); - if (failed(localized)) - return failure(); - strides.push_back(*localized); - } - - auto resultType = dyn_cast(extract.getType()); - if (!resultType) - return failure(); - FailureOr projected = - extractStaticSliceOrIdentity(state.rewriter, extract.getLoc(), *localizedSource, resultType, offsets, sizes, strides); - if (failed(projected)) - return failure(); - return std::optional {*projected}; -} - -LogicalResult applyProjectedExtractReplacementsInClonedOp(MaterializerState& state, - MaterializedClass& targetClass, - Operation& originalOp, - Operation& clonedOp, - CloneIndexingContext indexing, - IRMapping& mapper) { - if (auto originalExtract = dyn_cast(&originalOp)) { - auto clonedExtract = dyn_cast(&clonedOp); - Operation* insertionAnchor = clonedExtract ? getTopLevelMaterializedClassBodyOp(clonedExtract, targetClass) : nullptr; - FailureOr> maybeReplacement = - getOrMaterializeProjectedExtractReplacement(state, targetClass, originalExtract, insertionAnchor); - if (failed(maybeReplacement)) - return failure(); - if (std::optional replacement = *maybeReplacement) { - if (!clonedExtract) - return targetClass.op->emitError("projected replacement lost extract structure during cloning"); - - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPoint(clonedExtract); - FailureOr projected = materializeProjectedExtractReplacement( - state, targetClass, clonedExtract, *replacement, indexing.projectionSlotIndex, &mapper); - if (failed(projected)) - return failure(); - - clonedExtract.getResult().replaceAllUsesWith(*projected); - state.rewriter.eraseOp(clonedExtract); - return success(); - } - - if (auto clonedExtract = dyn_cast(&clonedOp)) { - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPoint(clonedExtract); - FailureOr> localProjected = - materializeLocalProjectedInputExtractReplacement(state, targetClass, clonedExtract, &mapper); - if (failed(localProjected)) - return failure(); - if (std::optional projected = *localProjected) { - clonedExtract.getResult().replaceAllUsesWith(*projected); - state.rewriter.eraseOp(clonedExtract); - return success(); - } - } - - if (std::optional producer = getProjectedWholeBatchReplacementProducer(state, originalExtract)) { - auto clonedExtract = dyn_cast(&clonedOp); - if (!clonedExtract) - return targetClass.op->emitError("projected whole-batch replacement lost extract structure during cloning"); - - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPoint(clonedExtract); - FailureOr projected = - materializeProjectedWholeBatchExtractReplacement(state, targetClass, clonedExtract, *producer, &mapper); - if (failed(projected)) - return failure(); - - clonedExtract.getResult().replaceAllUsesWith(*projected); - state.rewriter.eraseOp(clonedExtract); - return success(); - } - } - - if (originalOp.getNumRegions() != clonedOp.getNumRegions()) - return targetClass.op->emitError("projected replacement traversal found non-isomorphic cloned regions"); - - for (auto [originalRegion, clonedRegion] : llvm::zip(originalOp.getRegions(), clonedOp.getRegions())) { - if (std::distance(originalRegion.begin(), originalRegion.end()) - != std::distance(clonedRegion.begin(), clonedRegion.end())) - return targetClass.op->emitError("projected replacement traversal found non-isomorphic cloned blocks"); - - for (auto [originalBlock, clonedBlock] : llvm::zip(originalRegion.getBlocks(), clonedRegion.getBlocks())) { - auto originalIt = originalBlock.begin(); - auto clonedIt = clonedBlock.begin(); - while (originalIt != originalBlock.end() && clonedIt != clonedBlock.end()) { - Operation& originalNestedOp = *originalIt++; - Operation* currentClonedOp = &*clonedIt++; - if (failed(applyProjectedExtractReplacementsInClonedOp( - state, targetClass, originalNestedOp, *currentClonedOp, indexing, mapper))) - return failure(); - } - if (originalIt != originalBlock.end() || clonedIt != clonedBlock.end()) - return targetClass.op->emitError("projected replacement traversal found mismatched cloned operations"); - } - } - - return success(); -} - -LogicalResult mapClonedRegionBlockArguments(Operation& originalOp, Operation& clonedOp, IRMapping& mapper) { - if (originalOp.getNumRegions() != clonedOp.getNumRegions()) - return clonedOp.emitError("cloned operation has a different number of regions than the source operation"); - - for (auto [originalRegion, clonedRegion] : llvm::zip(originalOp.getRegions(), clonedOp.getRegions())) { - if (std::distance(originalRegion.begin(), originalRegion.end()) - != std::distance(clonedRegion.begin(), clonedRegion.end())) - return clonedOp.emitError("cloned operation has a different number of blocks than the source operation"); - - for (auto [originalBlock, clonedBlock] : llvm::zip(originalRegion.getBlocks(), clonedRegion.getBlocks())) { - if (originalBlock.getNumArguments() != clonedBlock.getNumArguments()) - return clonedOp.emitError("cloned operation block has a different number of arguments than the source block"); - - for (auto [originalArg, clonedArg] : llvm::zip(originalBlock.getArguments(), clonedBlock.getArguments())) - if (!mapper.contains(originalArg)) - mapper.map(originalArg, clonedArg); - - if (std::distance(originalBlock.begin(), originalBlock.end()) - != std::distance(clonedBlock.begin(), clonedBlock.end())) - return clonedOp.emitError("cloned operation block has a different number of operations than the source block"); - - auto originalIt = originalBlock.begin(); - auto clonedIt = clonedBlock.begin(); - while (originalIt != originalBlock.end()) { - Operation& originalNestedOp = *originalIt++; - Operation& clonedNestedOp = *clonedIt++; - if (failed(mapClonedRegionBlockArguments(originalNestedOp, clonedNestedOp, mapper))) - return failure(); - } - } - } - - return success(); -} - -static std::optional getConstantIndex(OpFoldResult value); -bool isStaticSliceInBounds(ArrayRef offsets, RankedTensorType sourceType, RankedTensorType fragmentType); - -FailureOr> tryNormalizeLocalizedExtractSlice(tensor::ExtractSliceOp extract, IRMapping& mapper) { - Value localizedSource = mapper.lookupOrDefault(extract.getSource()); - if (localizedSource == extract.getSource() || localizedSource.getType() == extract.getSource().getType()) - return std::optional {}; - - auto localizedSourceType = dyn_cast(localizedSource.getType()); - auto resultType = dyn_cast(extract.getType()); - if (!localizedSourceType || !resultType || !localizedSourceType.hasStaticShape() || !resultType.hasStaticShape()) - return std::optional {}; - if (localizedSourceType != resultType) - return std::optional {}; - - if (extract.getMixedSizes().size() != static_cast(localizedSourceType.getRank()) - || extract.getMixedStrides().size() != static_cast(localizedSourceType.getRank())) - return std::optional {}; - - for (int64_t dim = 0; dim < localizedSourceType.getRank(); ++dim) { - std::optional size = getConstantIndex(extract.getMixedSizes()[dim]); - std::optional stride = getConstantIndex(extract.getMixedStrides()[dim]); - if (!size || !stride || *stride != 1 || *size != localizedSourceType.getDimSize(dim)) - return std::optional {}; - } - - return std::optional {localizedSource}; -} - -LogicalResult verifyMaterializedStaticExtractSlice(Operation* anchor, - Value source, - RankedTensorType resultType, - ArrayRef offsets, - ArrayRef sizes, - ArrayRef strides) { - auto sourceType = dyn_cast(source.getType()); - if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape()) - return success(); - if (offsets.size() != static_cast(sourceType.getRank()) - || sizes.size() != static_cast(sourceType.getRank()) - || strides.size() != static_cast(sourceType.getRank())) - return success(); - - SmallVector staticOffsets; - staticOffsets.reserve(offsets.size()); - for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) { - std::optional offset = getConstantIndex(offsets[dim]); - std::optional size = getConstantIndex(sizes[dim]); - std::optional stride = getConstantIndex(strides[dim]); - if (!offset || !size || !stride) - return success(); - if (*stride != 1 || *size != resultType.getDimSize(dim)) - return success(); - staticOffsets.push_back(*offset); - } - - if (isStaticSliceInBounds(staticOffsets, sourceType, resultType)) - return success(); - - return anchor->emitError("materializer produced statically out-of-bounds extract_slice") - << " sourceType=" << sourceType << " resultType=" << resultType; -} - -LogicalResult cloneComputeTemplateBody(MaterializerState& state, - MaterializedClass& targetClass, - const ComputeInstance& instance, - IRMapping& mapper, - CloneIndexingContext indexing) { - Block& sourceBlock = getComputeInstanceTemplateBlock(instance); - for (Operation& op : sourceBlock.without_terminator()) { - if (auto extract = dyn_cast(&op)) { - FailureOr> maybeReplacement = - getOrMaterializeProjectedExtractReplacement(state, targetClass, extract); - if (failed(maybeReplacement)) - return failure(); - if (std::optional replacement = *maybeReplacement) { - FailureOr projected = materializeProjectedExtractReplacement( - state, targetClass, extract, *replacement, indexing.projectionSlotIndex, &mapper); - if (failed(projected)) - return failure(); - - mapper.map(extract.getResult(), *projected); - continue; - } - - FailureOr> localProjected = - materializeLocalProjectedInputExtractReplacement(state, targetClass, extract, &mapper); - if (failed(localProjected)) - return failure(); - if (std::optional projected = *localProjected) { - mapper.map(extract.getResult(), *projected); - continue; - } - - if (std::optional producer = getProjectedWholeBatchReplacementProducer(state, extract)) { - FailureOr projected = - materializeProjectedWholeBatchExtractReplacement(state, targetClass, extract, *producer, &mapper); - if (failed(projected)) - return failure(); - - mapper.map(extract.getResult(), *projected); - continue; - } - } - - for (Value operand : op.getOperands()) { - if (mapper.contains(operand)) - continue; - - FailureOr localized = localizeMaterializedClassOperand( - state, - targetClass, - operand, - &op, - "cloneComputeTemplateBody tried to reuse a tensor from another materialized class", - "cloneComputeTemplateBody produced an unsupported external non-tensor operand", - &mapper); - if (failed(localized)) - return failure(); - if (*localized != operand) - mapper.map(operand, *localized); - } - - if (auto extract = dyn_cast(&op)) { - FailureOr> normalized = tryNormalizeLocalizedExtractSlice(extract, mapper); - if (failed(normalized)) - return failure(); - if (*normalized) { - mapper.map(extract.getResult(), **normalized); - continue; - } - - auto remapFoldResult = [&](OpFoldResult value) -> OpFoldResult { - if (auto mappedValue = dyn_cast_if_present(value)) - return mapper.lookupOrDefault(mappedValue); - return value; - }; - - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - offsets.reserve(extract.getMixedOffsets().size()); - sizes.reserve(extract.getMixedSizes().size()); - strides.reserve(extract.getMixedStrides().size()); - - llvm::append_range(offsets, llvm::map_range(extract.getMixedOffsets(), remapFoldResult)); - llvm::append_range(sizes, llvm::map_range(extract.getMixedSizes(), remapFoldResult)); - llvm::append_range(strides, llvm::map_range(extract.getMixedStrides(), remapFoldResult)); - - auto resultType = cast(extract.getType()); - Value localizedSource = mapper.lookupOrDefault(extract.getSource()); - if (failed(verifyMaterializedStaticExtractSlice( - extract.getOperation(), localizedSource, resultType, offsets, sizes, strides))) - return failure(); - Value localizedExtract = extractStaticSliceOrIdentity( - state.rewriter, extract.getLoc(), localizedSource, resultType, offsets, sizes, strides); - mapper.map(extract.getResult(), localizedExtract); - continue; - } - - Operation* cloned = state.rewriter.clone(op, mapper); - if (failed(mapClonedRegionBlockArguments(op, *cloned, mapper))) - return failure(); - if (op.getNumRegions() != 0 - && failed(applyProjectedExtractReplacementsInClonedOp(state, targetClass, op, *cloned, indexing, mapper))) - return failure(); - if (failed(localizeCapturesInClonedOp(state, targetClass, *cloned, &mapper))) - return failure(); - for (auto [oldResult, newResult] : llvm::zip(op.getResults(), cloned->getResults())) - mapper.map(oldResult, newResult); - state.rewriter.setInsertionPointAfter(cloned); - } - - return success(); -} - -FailureOr materializeProjectedExtractReplacement(MaterializerState& state, - MaterializedClass& targetClass, - tensor::ExtractSliceOp extract, - const ProjectedExtractReplacement& replacement, - std::optional projectionSlotIndex, - IRMapping* mapper) { - if (failed(verifyProjectedFragmentLayout(targetClass.op, replacement.layout))) - return failure(); - - FailureOr localizedPayload = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - replacement.payload, - targetClass.op, - "projected extract replacement tried to reuse a tensor from another materialized class", - std::nullopt, - mapper); - if (failed(localizedPayload)) - return failure(); - Value payload = *localizedPayload; - Operation* topLevelAnchor = getTopLevelMaterializedClassBodyOp(extract.getOperation(), targetClass); - Operation* payloadDef = payload.getDefiningOp(); - Operation* topLevelPayloadDef = payloadDef ? getTopLevelMaterializedClassBodyOp(payloadDef, targetClass) : nullptr; - if (topLevelAnchor && topLevelPayloadDef && topLevelAnchor != topLevelPayloadDef - && topLevelAnchor->getBlock() == topLevelPayloadDef->getBlock() - && topLevelAnchor->isBeforeInBlock(topLevelPayloadDef)) - moveDefinitionBeforeAnchorWithSameBlockDeps(topLevelPayloadDef, topLevelAnchor); - if (payloadDef && payloadDef->getBlock() == extract->getBlock() && extract->isBeforeInBlock(payloadDef)) - moveDefinitionBeforeAnchorWithSameBlockDeps(payloadDef, extract); - - if (replacement.layout.payloadFragmentCount == 1) - return payload; - - if (replacement.layout.payloadFragmentCount < replacement.layout.fragmentsPerLogicalSlot) - return targetClass.op->emitError("projected replacement payload is smaller than one logical slot"); - - Value intraSlotFragmentIndex = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - const auto linearizeProjectedLoopIndices = [&]() -> FailureOr { - if (replacement.layout.loopTripCounts.empty()) - return intraSlotFragmentIndex; - - SmallVector surroundingLoops; - for (Operation* current = extract->getParentOp(); current; current = current->getParentOp()) { - if (auto loop = dyn_cast(current)) - surroundingLoops.push_back(loop); - if (current == targetClass.op) - break; - } - std::reverse(surroundingLoops.begin(), surroundingLoops.end()); - - if (surroundingLoops.size() != replacement.layout.loopTripCounts.size()) - return targetClass.op->emitError("projected replacement loop structure does not match the collected descriptor"); - - Value linearizedIndex = intraSlotFragmentIndex; - for (auto [index, loop] : llvm::enumerate(surroundingLoops)) { - FailureOr localizedIv = - rematerializeIndexValueInClass(state, targetClass, loop.getInductionVar(), extract.getLoc(), mapper); - if (failed(localizedIv)) - return failure(); - Value iv = *localizedIv; - MLIRContext* context = state.func.getContext(); - AffineExpr d0 = getAffineDimExpr(0, context); - AffineMap normalizedMap = - AffineMap::get(/*dimCount=*/1, - /*symbolCount=*/0, - (d0 - replacement.layout.loopLowerBounds[index]).floorDiv(replacement.layout.loopSteps[index])); - Value normalized = - createOrFoldAffineApply(state.rewriter, extract.getLoc(), normalizedMap, ValueRange {iv}, targetClass.op); - - AffineExpr d1 = getAffineDimExpr(1, context); - AffineMap linearizedMap = AffineMap::get( - /*dimCount=*/2, /*symbolCount=*/0, d0 * replacement.layout.loopTripCounts[index] + d1); - linearizedIndex = createOrFoldAffineApply( - state.rewriter, extract.getLoc(), linearizedMap, ValueRange {linearizedIndex, normalized}, targetClass.op); - } - return linearizedIndex; - }; - - FailureOr linearizedIndex = linearizeProjectedLoopIndices(); - if (failed(linearizedIndex)) - return failure(); - intraSlotFragmentIndex = *linearizedIndex; - - const auto computeProjectedPayloadFragmentIndex = [&]() -> FailureOr { - if (replacement.layout.payloadFragmentCount == replacement.layout.fragmentsPerLogicalSlot) { - if (replacement.layout.loopTripCounts.empty() && replacement.layout.fragmentsPerLogicalSlot != 1) - return targetClass.op->emitError("projected replacement is missing loop metadata for packed logical slot"); - return intraSlotFragmentIndex; - } - - if (!projectionSlotIndex) - return targetClass.op->emitError("packed projected extract replacement requires a fragment slot index"); - - FailureOr localProjectionSlotIndex = - rematerializeIndexValueInClass(state, targetClass, *projectionSlotIndex, extract.getLoc(), mapper); - if (failed(localProjectionSlotIndex)) - return failure(); - - MLIRContext* context = state.func.getContext(); - AffineExpr d0 = getAffineDimExpr(0, context); - AffineExpr d1 = getAffineDimExpr(1, context); - AffineMap packedIndexMap = AffineMap::get( - /*dimCount=*/2, - /*symbolCount=*/0, - d0 * replacement.layout.fragmentsPerLogicalSlot + d1); - return createOrFoldAffineApply(state.rewriter, - extract.getLoc(), - packedIndexMap, - ValueRange {*localProjectionSlotIndex, intraSlotFragmentIndex}, - targetClass.op); - }; - - FailureOr packedFragmentIndex = computeProjectedPayloadFragmentIndex(); - if (failed(packedFragmentIndex)) - return failure(); - - FailureOr packedOffset = scaleIndexByDim0SizeInClass( - state, targetClass, *packedFragmentIndex, replacement.layout.fragmentType.getDimSize(0), extract.getLoc()); - if (failed(packedOffset)) - return failure(); - return createDim0ExtractSliceInClass( - state, targetClass, extract.getLoc(), payload, *packedOffset, replacement.layout.fragmentType.getDimSize(0)); -} - -FailureOr materializeIndexedBatchRunReceive(MaterializerState& state, - MaterializedClass& targetClass, - IndexedBatchRunValue& run, - Value runSlotIndex, - Location loc) { - if (!targetClass.isBatch) - return targetClass.op->emitError("indexed batch run receive requires a batch target class"); - if (run.packed) - return getPackedSliceForDynamicRunIndex(state, targetClass.op, run.packed, run.fragmentType, runSlotIndex, loc); - if (run.exchangeIds.empty()) - return targetClass.op->emitError("indexed batch run receive is missing planned exchange ids"), failure(); - FailureOr messages = - buildMessageVectorForExchangeIds(state, - run.exchangeIds, - CommunicationExchangeKind::IndexedBatchRun, - state.communicationPlan.getExchange(run.exchangeIds.front()).sourceClass, - targetClass.id, - run.fragmentType, - targetClass.op); - if (failed(messages)) - return failure(); - - Value flatIndex = createBatchRunFlatIndex(state, targetClass, runSlotIndex, loc); - std::optional preferredPeriod = static_cast(targetClass.cpus.size()); - Value channelId = createIndexedChannelId(state, targetClass.op, *messages, flatIndex, loc, preferredPeriod); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, *messages, flatIndex, loc, preferredPeriod); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, *messages, flatIndex, loc, preferredPeriod); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - return SpatChannelReceiveOp::create(state.rewriter, loc, run.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); -} - -LogicalResult localizeCapturesInOperationTree(MaterializerState& state, - MaterializedClass& targetClass, - Operation& root, - StringRef tensorContext, - StringRef genericContext, - IRMapping* mapper = nullptr) { - WalkResult walkResult = root.walk([&](Operation* nestedOp) -> WalkResult { - for (OpOperand& operand : nestedOp->getOpOperands()) { - Value current = operand.get(); - if (isValueLegalInMaterializedClassBody(current, targetClass)) - continue; - - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPoint(nestedOp); - FailureOr localized = - localizeMaterializedClassOperand(state, targetClass, current, nestedOp, tensorContext, genericContext, mapper); - if (failed(localized)) { - InFlightDiagnostic diagnostic = targetClass.op->emitError("failed to localize cloned scheduled-body operand"); - diagnostic << " targetClass=" << targetClass.id << " nestedOp='" << nestedOp->getName() << "' operand#" - << operand.getOperandNumber() << " operandType=" << current.getType() << " offendingIR=\"" - << truncateMaterializerDebugString(stringifyOperationForMaterializerDebug(nestedOp)) - << "\" offendingOperands=\"" << formatMaterializerOperandListInline(nestedOp, targetClass) - << "\" parentChain=\"" << formatMaterializerParentChainInline(nestedOp) << "\""; - diagnostic.attachNote(nestedOp->getLoc()) << "offending nested operation"; - attachMaterializerOperationPrintNote(diagnostic, nestedOp, "offending nested operation IR"); - attachMaterializerOperandListNote(diagnostic, nestedOp, targetClass, "offending nested operation operands"); - attachMaterializerParentChainNote(diagnostic, nestedOp, "offending nested operation parent chain"); - attachMaterializerValueOriginNote(diagnostic, current, "offending operand"); - attachMaterializerOperationPrintNote(diagnostic, targetClass.op, "target materialized op"); - attachMaterializedClassBodySummary(diagnostic, targetClass); - return WalkResult::interrupt(); - } - operand.set(*localized); - } - return WalkResult::advance(); - }); - - return walkResult.wasInterrupted() ? failure() : success(); -} - -LogicalResult localizeCapturesInClonedOp(MaterializerState& state, - MaterializedClass& targetClass, - Operation& clonedOp, - IRMapping* mapper) { - return localizeCapturesInOperationTree( - state, - targetClass, - clonedOp, - "cloneComputeTemplateBody tried to reuse a tensor from another materialized class", - "cloneComputeTemplateBody produced an unsupported external non-tensor operand", - mapper); -} - -LogicalResult localizeAllScheduledBodyCaptures(MaterializerState& state, MaterializedClass& targetClass) { - SmallVector bodyOps; - for (Operation& op : *targetClass.body) - op.walk([&](Operation* nestedOp) { bodyOps.push_back(nestedOp); }); - - for (Operation* nestedOp : bodyOps) { - if (nestedOp->getBlock() == nullptr) - continue; - if (auto extract = dyn_cast(nestedOp)) { - if (!isValueLegalInMaterializedClassBody(extract.getSource(), targetClass)) { - FailureOr> maybeReplacement = - getOrMaterializeProjectedExtractReplacement(state, targetClass, extract, extract.getOperation()); - if (failed(maybeReplacement)) - return failure(); - if (std::optional replacement = *maybeReplacement) { - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPoint(extract); - FailureOr projected = materializeProjectedExtractReplacement( - state, targetClass, extract, *replacement, std::nullopt, nullptr); - if (failed(projected)) - return failure(); - extract.getResult().replaceAllUsesWith(*projected); - state.rewriter.eraseOp(extract); - continue; - } - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPoint(extract); - FailureOr> localProjected = - materializeLocalProjectedInputExtractReplacement(state, targetClass, extract, nullptr); - if (failed(localProjected)) - return failure(); - if (std::optional projected = *localProjected) { - extract.getResult().replaceAllUsesWith(*projected); - state.rewriter.eraseOp(extract); - continue; - } - } - } - for (OpOperand& operand : nestedOp->getOpOperands()) { - Value current = operand.get(); - if (isValueLegalInMaterializedClassBody(current, targetClass)) - continue; - - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPoint(nestedOp); - FailureOr localized = localizeMaterializedClassOperand( - state, - targetClass, - current, - nestedOp, - "final scheduled body capture localization tried to reuse a tensor from another materialized class", - "final scheduled body capture localization found an unsupported external non-tensor operand"); - if (failed(localized)) { - InFlightDiagnostic diagnostic = targetClass.op->emitError("failed to localize final scheduled-body operand"); - diagnostic << " targetClass=" << targetClass.id << " nestedOp='" << nestedOp->getName() << "' operand#" - << operand.getOperandNumber() << " operandType=" << current.getType() << " offendingIR=\"" - << truncateMaterializerDebugString(stringifyOperationForMaterializerDebug(nestedOp)) - << "\" offendingOperands=\"" << formatMaterializerOperandListInline(nestedOp, targetClass) - << "\" parentChain=\"" << formatMaterializerParentChainInline(nestedOp) << "\""; - diagnostic.attachNote(nestedOp->getLoc()) << "offending nested operation"; - attachMaterializerValueOriginNote(diagnostic, current, "offending operand"); - attachMaterializedClassBodySummary(diagnostic, targetClass); - return failure(); - } - operand.set(*localized); - } - } - - return success(); -} - -FailureOr> cloneInstanceBody(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef peers, - CloneIndexingContext indexing) { - assert(!peers.empty() && "expected at least one peer instance"); - const ComputeInstance& instance = peers.front(); - Operation* sourceOp = instance.op; - Location loc = sourceOp->getLoc(); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - - IRMapping mapper; - if (auto batch = dyn_cast(sourceOp)) { - for (const ComputeInstance& peer : peers) { - if (peer.op != sourceOp) { - sourceOp->emitError("equivalence class slot contains different source compute_batch operations"); - return failure(); - } - } - auto laneArg = batch.getLaneArgument(); - if (!laneArg) { - sourceOp->emitError("expected source compute_batch lane block argument"); - return failure(); - } - mapper.map(*laneArg, createOriginalLaneValue(state, targetClass, peers, loc)); - } - - OpBuilder::InsertPoint cloneInsertionPoint = state.rewriter.saveInsertionPoint(); - - mapWeights(state, targetClass, instance, mapper); - if (failed(mapInputs(state, targetClass, instance, mapper, indexing))) { - if (hasDeferredMaterializationDependency(state)) - return failure(); - sourceOp->emitError("failed to map inputs while cloning compute/compute_batch body") - << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount - << " targetClass=" << targetClass.id; - return failure(); - } - - state.rewriter.restoreInsertionPoint(cloneInsertionPoint); - if (failed(cloneComputeTemplateBody(state, targetClass, instance, mapper, indexing))) { - if (hasDeferredMaterializationDependency(state)) - return failure(); - sourceOp->emitError("failed to clone compute/compute_batch body") - << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount - << " targetClass=" << targetClass.id; - return failure(); - } - - if (auto compute = dyn_cast(sourceOp)) { - Block& sourceBlock = getComputeInstanceTemplateBlock(instance); - auto yield = dyn_cast_or_null(sourceBlock.getTerminator()); - if (!yield) { - compute.emitOpError("expected spat.yield terminator while materializing compute"); - return failure(); - } - - SmallVector outputs; - outputs.reserve(yield.getNumOperands()); - for (Value yielded : yield.getOutputs()) - outputs.push_back(mapper.lookupOrDefault(yielded)); - return outputs; - } - - auto batch = cast(sourceOp); - if (batch.getNumResults() == 0) - return SmallVector {}; - - SmallVector outputs = collectMappedBatchOutputs(batch, mapper); - for (Value output : outputs) - if (!output) { - batch.emitOpError("failed to recover yielded per-lane value for compute_batch result"); - return failure(); - } - return outputs; -} - -bool sameDestinationClasses(ArrayRef lhs, ArrayRef rhs) { - if (lhs.size() != rhs.size()) - return false; - for (auto [lhsClass, rhsClass] : llvm::zip(lhs, rhs)) - if (lhsClass != rhsClass) - return false; - return true; -} - -SmallVector -collectDestinationClassesForRun(MaterializerState& state, ArrayRef run, size_t resultIndex) { - SmallVector destinations; - - for (const MaterializationRunSlot& slot : run) { - for (const ComputeInstance& peer : slot.peers) { - ProducerKey key {peer, resultIndex}; - for (ClassId destinationClass : getDestinationClasses(state, key)) - if (!llvm::is_contained(destinations, destinationClass)) - destinations.push_back(destinationClass); - } - } - - llvm::sort(destinations); - return destinations; -} - -SmallVector groupBatchRunOutputsByDestination(MaterializerState& state, - ArrayRef run) { - assert(!run.empty() && "expected non-empty materialization run"); - assert(!run.front().peers.empty() && "expected non-empty materialization run slot"); - - SmallVector groups; - ArrayRef outputs = getComputeInstanceOutputValuesCached(state, run.front().peers.front()); - - for (auto [resultIndex, output] : llvm::enumerate(outputs)) { - SmallVector destinations = collectDestinationClassesForRun(state, run, resultIndex); - - auto existingGroup = llvm::find_if(groups, [&](const OutputDestinationGroup& group) { - return sameDestinationClasses(group.destinationClasses, destinations); - }); - - if (existingGroup != groups.end()) { - existingGroup->resultIndices.push_back(resultIndex); - continue; - } - - OutputDestinationGroup group; - group.resultIndices.push_back(resultIndex); - group.destinationClasses = std::move(destinations); - groups.push_back(std::move(group)); - } - - return groups; -} - -FailureOr getPackedRunTensorType(Type elementType, size_t runSize) { - auto tensorType = dyn_cast(elementType); - if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0) - return failure(); - - SmallVector shape(tensorType.getShape()); - shape[0] *= static_cast(runSize); - return RankedTensorType::get(shape, tensorType.getElementType(), tensorType.getEncoding()); -} - -LogicalResult registerDeferredLocalPackedRunValue(MaterializerState& state, - MaterializedClass& materializedClass, - ArrayRef keys, - Type fragmentType, - Location loc) { - if (keys.empty()) - return success(); - - auto rankedFragmentType = dyn_cast(fragmentType); - if (!rankedFragmentType || !rankedFragmentType.hasStaticShape() || rankedFragmentType.getRank() == 0) - return materializedClass.op->emitError("deferred local packed run expects static ranked fragment type"); - - Operation* sourceOp = keys.front().instance.op; - size_t resultIndex = keys.front().resultIndex; - - for (ProducerKey key : keys) { - if (key.instance.op != sourceOp || key.resultIndex != resultIndex) - return materializedClass.op->emitError("deferred local packed run expects one producer result"); - - if (key.instance.laneCount != 1) - return materializedClass.op->emitError("deferred local packed run expects one lane per fragment"); - } - - PackedScalarRunValue packedRun; - packedRun.targetClass = materializedClass.id; - packedRun.sourceOp = sourceOp; - packedRun.resultIndex = resultIndex; - packedRun.kind = PackedScalarRunKind::DeferredLocalCompute; - packedRun.fragmentType = rankedFragmentType; - - packedRun.slots.reserve(keys.size()); - for (ProducerKey key : keys) { - PackedScalarRunSlot slot; - slot.keys.push_back(key); - packedRun.slots.push_back(std::move(slot)); - } - - state.availableValues.recordPackedRun(std::move(packedRun)); - return success(); -} - -LogicalResult registerPackedRunValue(MaterializerState& state, - MaterializedClass& materializedClass, - ArrayRef keys, - Value packed, - Type fragmentType, - Location loc) { - if (keys.empty()) - return success(); - - FailureOr expectedPackedType = getPackedRunTensorType(fragmentType, keys.size()); - if (failed(expectedPackedType)) - return materializedClass.op->emitError("packed run registration expects static ranked fragment type"); - - if (packed.getType() != *expectedPackedType) - return materializedClass.op->emitError("packed run value has unexpected tensor type"); - - auto rankedFragmentType = dyn_cast(fragmentType); - if (!rankedFragmentType || !rankedFragmentType.hasStaticShape() || rankedFragmentType.getRank() == 0) - return materializedClass.op->emitError("packed run registration expects static ranked fragment type"); - - Operation* sourceOp = keys.front().instance.op; - size_t resultIndex = keys.front().resultIndex; - - for (ProducerKey key : keys) { - if (key.instance.op != sourceOp || key.resultIndex != resultIndex) - return materializedClass.op->emitError("packed run registration expects one producer result"); - if (key.instance.laneCount != 1) - return materializedClass.op->emitError("packed run registration expects one lane per packed fragment"); - } - - if (std::optional contiguousKey = getContiguousProducerRangeForKeys(keys)) { - state.availableValues.record(*contiguousKey, materializedClass.id, packed); - return success(); - } - - PackedScalarRunValue packedRun; - packedRun.targetClass = materializedClass.id; - packedRun.sourceOp = sourceOp; - packedRun.resultIndex = resultIndex; - packedRun.packed = packed; - packedRun.kind = PackedScalarRunKind::Materialized; - packedRun.fragmentType = rankedFragmentType; - - packedRun.slots.reserve(keys.size()); - for (ProducerKey key : keys) { - PackedScalarRunSlot slot; - slot.keys.push_back(key); - packedRun.slots.push_back(std::move(slot)); - } - - state.availableValues.recordPackedRun(std::move(packedRun)); - return success(); -} - -FailureOr getPackedSlotFragmentType(Value packed, size_t keyCount) { - auto packedType = dyn_cast(packed.getType()); - if (!packedType || !packedType.hasStaticShape() || packedType.getRank() == 0 || keyCount == 0) - return failure(); - - int64_t packedRows = packedType.getDimSize(0); - if (packedRows % static_cast(keyCount) != 0) - return failure(); - - SmallVector fragmentShape(packedType.getShape()); - fragmentShape[0] = packedRows / static_cast(keyCount); - return RankedTensorType::get(fragmentShape, packedType.getElementType(), packedType.getEncoding()); -} - -LogicalResult recordLocalPackedSlotValue(MaterializerState& state, - MaterializedClass& materializedClass, - ArrayRef keys, - Value packed) { - if (keys.empty()) - return success(); - - Operation* sourceOp = keys.front().instance.op; - size_t resultIndex = keys.front().resultIndex; - for (ProducerKey key : keys) { - if (key.instance.op != sourceOp || key.resultIndex != resultIndex) - return materializedClass.op->emitError("local packed slot registration expects one producer result"); - if (key.instance.laneCount != 1) - return materializedClass.op->emitError("local packed slot registration expects one lane per packed fragment"); - } - - FailureOr fragmentType = getPackedSlotFragmentType(packed, keys.size()); - if (failed(fragmentType)) { - for (ProducerKey key : keys) - state.availableValues.record(key, materializedClass.id, packed); - return success(); - } - - PackedScalarRunValue packedRun; - packedRun.targetClass = materializedClass.id; - packedRun.sourceOp = sourceOp; - packedRun.resultIndex = resultIndex; - packedRun.packed = packed; - packedRun.kind = PackedScalarRunKind::Materialized; - packedRun.fragmentType = *fragmentType; - - PackedScalarRunSlot slot; - llvm::append_range(slot.keys, keys); - packedRun.slots.push_back(std::move(slot)); - - state.availableValues.recordPackedRun(std::move(packedRun)); - return success(); -} - -LogicalResult emitPackedRunFanout(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef destinationClasses, - ArrayRef keys, - Value packed, - Type fragmentType, - Location loc) { - assert(!sourceClass.isBatch && "packed run fanout expects a scalar source class"); - - auto rankedFragmentType = dyn_cast(fragmentType); - if (!rankedFragmentType || !rankedFragmentType.hasStaticShape() || rankedFragmentType.getRank() == 0) - return sourceClass.op->emitError("packed run fanout expects static ranked fragment type"); - - Operation* sourceOp = keys.front().instance.op; - size_t resultIndex = keys.front().resultIndex; - for (ClassId destinationClass : destinationClasses) { - if (destinationClass == sourceClass.id) - continue; - - SmallVector exchangeIds; - for (ProducerKey key : keys) { - ArrayRef ids = - state.communicationPlan.getOrdinaryValueExchanges(key, destinationClass); - if (ids.empty()) - continue; - exchangeIds.push_back(ids.front()); - break; - } - if (exchangeIds.empty()) - continue; - - for (CommunicationExchangeId id : exchangeIds) - if (failed(bindCommunicationPayload(state, id, sourceClass, packed, sourceClass.op, loc))) - return failure(); - - PackedScalarRunValue packedRun; - packedRun.targetClass = destinationClass; - packedRun.sourceOp = sourceOp; - packedRun.resultIndex = resultIndex; - packedRun.kind = PackedScalarRunKind::DeferredReceive; - packedRun.fragmentType = rankedFragmentType; - llvm::append_range(packedRun.exchangeIds, exchangeIds); - PackedScalarRunSlot slot; - llvm::append_range(slot.keys, keys); - packedRun.slots.push_back(std::move(slot)); - state.availableValues.recordPackedRun(std::move(packedRun)); - } - - FailureOr emitted = emitReadyCommunicationEventsForClass(state, sourceClass, loc); - return failed(emitted) ? failure() : success(); -} - -FailureOr> cloneBatchBodyForLane(MaterializerState& state, - MaterializedClass& targetClass, - const ComputeInstance& instance, - Value laneValue, - ArrayRef resultIndices, - CloneIndexingContext indexing) { - auto batch = dyn_cast(instance.op); - if (!batch) - return failure(); - - IRMapping mapper; - auto sourceLaneArg = batch.getLaneArgument(); - if (!sourceLaneArg) - return batch.emitOpError("expected source compute_batch lane block argument"); - - mapper.map(*sourceLaneArg, laneValue); - - OpBuilder::InsertPoint cloneInsertionPoint = state.rewriter.saveInsertionPoint(); - - mapWeights(state, targetClass, instance, mapper); - if (failed(mapInputs(state, targetClass, instance, mapper, indexing))) - return failure(); - - state.rewriter.restoreInsertionPoint(cloneInsertionPoint); - if (failed(cloneComputeTemplateBody(state, targetClass, instance, mapper, indexing))) - return failure(); - - SmallVector allOutputs = collectMappedBatchOutputs(batch, mapper); - if (allOutputs.empty() && !resultIndices.empty()) - return batch.emitOpError("failed to recover source compute_batch outputs"); - - SmallVector selectedOutputs; - selectedOutputs.reserve(resultIndices.size()); - for (size_t resultIndex : resultIndices) { - if (resultIndex >= allOutputs.size() || !allOutputs[resultIndex]) - return batch.emitOpError("failed to recover selected compute_batch output"); - selectedOutputs.push_back(allOutputs[resultIndex]); - } - - return selectedOutputs; -} - -FailureOr> materializeBatchOutputGroupLoop(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef run, - const OutputDestinationGroup& group) { - assert(!run.empty() && "expected non-empty batch run"); - assert(!run.front().peers.empty() && "expected non-empty materialization run slot"); - - Operation* sourceOp = run.front().peers.front().op; - Location loc = sourceOp->getLoc(); - - if (run.size() == 1) { - if (run.front().peers.size() != 1) - return sourceOp->emitError("scalar batch output loop expects exactly one peer in singleton slot"); - - const ComputeInstance& item = run.front().peers.front(); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value laneValue = getOrCreateIndexConstant(state.constantFolder, targetClass.op, item.laneStart); - return cloneBatchBodyForLane(state, targetClass, item, laneValue, group.resultIndices, {}); - } - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - - auto sourceBatch = cast(sourceOp); - SmallVector& fragmentTypes = getBatchOutputFragmentTypesCached(state, sourceBatch); - SmallVector initValues; - for (size_t resultIndex : group.resultIndices) { - if (resultIndex >= fragmentTypes.size() || !fragmentTypes[resultIndex]) - return sourceBatch.emitOpError("failed to recover per-lane output type for packed batch run"); - - Type fragmentType = fragmentTypes[resultIndex]; - FailureOr packedType = getPackedRunTensorType(fragmentType, run.size()); - if (failed(packedType)) - return sourceBatch.emitOpError("cannot materialize packed batch run for non-static ranked output"); - - initValues.push_back( - tensor::EmptyOp::create(state.rewriter, loc, packedType->getShape(), packedType->getElementType()).getResult()); - } - - SmallVector logicalLanes; - logicalLanes.reserve(run.size()); - for (const MaterializationRunSlot& slot : run) { - if (slot.peers.size() != 1) - return sourceOp->emitError("scalar batch output loop expects exactly one peer per materialization slot"); - - const ComputeInstance& item = slot.peers.front(); - if (item.op != sourceOp) - return sourceOp->emitError("materialization run contains different source operations"); - - logicalLanes.push_back(item.laneStart); - } - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(run.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - auto loop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange(initValues), - [&](OpBuilder&, Location, Value loopIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { - Value sourceLane = createIndexedIndexValue(state, targetClass.op, logicalLanes, loopIndex, loc); - - FailureOr> produced = - cloneBatchBodyForLane(state, - targetClass, - run.front().peers.front(), - sourceLane, - group.resultIndices, - CloneIndexingContext {.runSlotIndex = loopIndex, .projectionSlotIndex = loopIndex}); - if (failed(produced)) - return failure(); - - yielded.reserve(produced->size()); - for (auto [outputIndex, output] : llvm::enumerate(*produced)) { - auto fragmentType = cast(output.getType()); - Value acc = iterArgs[outputIndex]; - Value firstOffset = scaleIndexByDim0Size(state, targetClass.op, loopIndex, fragmentType.getDimSize(0), loc); - yielded.push_back(createDim0InsertSlice(state, loc, output, acc, firstOffset)); - } - return success(); - }); - if (failed(loop)) - return failure(); - - SmallVector results; - results.reserve(loop->results.size()); - for (Value result : loop->results) - results.push_back(result); - return results; -} - -SmallVector getMaterializationRunSlotOutputKeys(const MaterializationRunSlot& slot, - size_t resultIndex) { - SmallVector keys; - keys.reserve(slot.peers.size()); - for (const ComputeInstance& peer : slot.peers) - keys.push_back({peer, resultIndex}); - return keys; -} - -FailureOr> -getMaterializationRunSlotPeers(MaterializerState& state, MaterializedClass& targetClass, SlotId logicalSlot) { - if (targetClass.isBatch) - return getPeerLogicalInstances(state, targetClass, logicalSlot); - - auto streamIt = state.logicalInstancesByCpu.find(targetClass.cpus.front()); - if (streamIt == state.logicalInstancesByCpu.end() || logicalSlot >= streamIt->second.size()) - return failure(); - - return SmallVector {streamIt->second[logicalSlot]}; -} - -FailureOr collectBatchMaterializationRun(MaterializerState& state, - MaterializedClass& targetClass, - SlotId startSlot, - Operation* sourceOp) { - MaterializationRun run; - - for (SlotId slot = startSlot;; ++slot) { - ClassSlotKey classSlot {targetClass.id, slot}; - if (state.materializedLogicalSlots.contains(classSlot)) - break; - - FailureOr> peers = getMaterializationRunSlotPeers(state, targetClass, slot); - if (failed(peers) || peers->empty()) - break; - - bool validSlot = true; - for (const ComputeInstance& peer : *peers) { - if (peer.op != sourceOp || !isa(peer.op)) { - validSlot = false; - break; - } - } - - if (!validSlot) - break; - - MaterializationRunSlot runSlot; - runSlot.peers = std::move(*peers); - run.push_back(std::move(runSlot)); - } - - if (run.empty()) - return failure(); - - return run; -} - -SmallVector getMaterializationRunOutputKeys(ArrayRef run, size_t resultIndex) { - SmallVector keys; - for (const MaterializationRunSlot& slot : run) - llvm::append_range(keys, getMaterializationRunSlotOutputKeys(slot, resultIndex)); - return keys; -} - -ArrayRef getFirstMaterializationRunOriginalOutputs(MaterializerState& state, - ArrayRef run) { - assert(!run.empty() && "expected non-empty materialization run"); - assert(!run.front().peers.empty() && "expected non-empty materialization run slot"); - return getComputeInstanceOutputValuesCached(state, run.front().peers.front()); -} - -Operation* getMaterializationRunSourceOp(ArrayRef run) { - assert(!run.empty() && "expected non-empty materialization run"); - assert(!run.front().peers.empty() && "expected non-empty materialization run slot"); - return run.front().peers.front().op; -} - -Location getMaterializationRunLoc(ArrayRef run) { - return getMaterializationRunSourceOp(run)->getLoc(); -} - -bool hasMaterializationRunResultLiveExternalUse(MaterializerState& state, - ArrayRef run, - size_t resultIndex) { - for (const MaterializationRunSlot& slot : run) { - for (const ComputeInstance& peer : slot.peers) { - ArrayRef outputs = getComputeInstanceOutputValuesCached(state, peer); - if (resultIndex >= outputs.size()) - return true; - - if (hasLiveExternalUseCached(state, outputs[resultIndex])) - return true; - } - } - - return false; -} - -bool hasMaterializationRunGroupLiveExternalUse(MaterializerState& state, - ArrayRef run, - const OutputDestinationGroup& group) { - for (size_t resultIndex : group.resultIndices) - if (hasMaterializationRunResultLiveExternalUse(state, run, resultIndex)) - return true; - - return false; -} - -bool hasSameClassConsumer(MaterializerState& state, ProducerKey producerKey, ClassId classId); - -bool hasMaterializationRunGroupSameClassConsumer(MaterializerState& state, - ClassId classId, - ArrayRef run, - const OutputDestinationGroup& group) { - for (size_t resultIndex : group.resultIndices) { - for (const MaterializationRunSlot& slot : run) { - for (const ComputeInstance& peer : slot.peers) - if (hasSameClassConsumer(state, {peer, resultIndex}, classId)) - return true; - } - } - - return false; -} - -bool hasPlannedProjectedInputSendForKeys(MaterializerState& state, ArrayRef keys) { - DenseSet keySet; - for (ProducerKey key : keys) - keySet.insert(key); - - for (auto& extractEntry : state.projectedInputTransferPlans) - for (auto& classEntry : extractEntry.second) - for (const ProjectedInputTransferFragment& fragment : classEntry.second.fragments) - if (keySet.contains(fragment.producer)) - return true; - return false; -} - -bool hasMaterializationRunGroupProjectedConsumer(MaterializerState& state, - ArrayRef run, - const OutputDestinationGroup& group) { - for (size_t resultIndex : group.resultIndices) { - SmallVector keys = getMaterializationRunOutputKeys(run, resultIndex); - if (hasPlannedProjectedInputSendForKeys(state, keys)) - return true; - } - return false; -} - -bool hasMaterializationRunResultSameClassConsumer(MaterializerState& state, - ClassId classId, - ArrayRef run, - size_t resultIndex) { - for (const MaterializationRunSlot& slot : run) - for (const ComputeInstance& peer : slot.peers) - if (hasSameClassConsumer(state, {peer, resultIndex}, classId)) - return true; - return false; -} - -StringRef describeWholeTensorBarrierReason(WholeTensorBarrierReason reason) { - switch (reason) { - case WholeTensorBarrierReason::FunctionReturnWithoutBlueprint: - return "function return or external use without spat.blueprint assembly"; - case WholeTensorBarrierReason::DenseLogicalConsumer: return "consumer requires a dense logical tensor"; - } - llvm_unreachable("unknown whole-tensor barrier reason"); -} - -FailureOr classifyRunOutputDemand(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef run, - ArrayRef destinationClasses, - size_t resultIndex) { - auto sourceBatch = dyn_cast(getMaterializationRunSourceOp(run)); - if (!sourceBatch) - return targetClass.op->emitError("compact batch demand classification expects a compute_batch source"); - - SmallVector& fragmentTypes = getBatchOutputFragmentTypesCached(state, sourceBatch); - ArrayRef firstOriginalOutputs = getFirstMaterializationRunOriginalOutputs(state, run); - if (resultIndex >= firstOriginalOutputs.size() || resultIndex >= fragmentTypes.size()) - return targetClass.op->emitError("compact batch demand classification found an invalid output index") - << " resultIndex=" << resultIndex; - - auto fragmentType = dyn_cast(fragmentTypes[resultIndex]); - if (!fragmentType || !fragmentType.hasStaticShape() || fragmentType.getRank() == 0) - return targetClass.op->emitError("compact batch demand classification requires static ranked fragment metadata") - << " resultIndex=" << resultIndex << " fragmentType=" << fragmentTypes[resultIndex]; - - RunOutputDemand demand; - demand.resultIndex = resultIndex; - demand.originalOutput = firstOriginalOutputs[resultIndex]; - demand.fragmentType = fragmentType; - - for (ClassId destinationClass : destinationClasses) - demand.actions.push_back(TensorDemandAction {.kind = TensorDemandActionKind::DestinationFanout, - .destinationClass = destinationClass, - .barrierReason = std::nullopt}); - - if (hasMaterializationRunResultSameClassConsumer(state, targetClass.id, run, resultIndex)) - demand.actions.push_back(TensorDemandAction {.kind = TensorDemandActionKind::SameClassIndexedFragment, - .destinationClass = std::nullopt, - .barrierReason = std::nullopt}); - - SmallVector keys = getMaterializationRunOutputKeys(run, resultIndex); - if (hasPlannedProjectedInputSendForKeys(state, keys)) - demand.actions.push_back(TensorDemandAction {.kind = TensorDemandActionKind::ProjectedInputPublication, - .destinationClass = std::nullopt, - .barrierReason = std::nullopt}); - - if (!hasMaterializationRunResultLiveExternalUse(state, run, resultIndex)) - return demand; - - Value originalOutput = demand.originalOutput; - if (!isTerminalHostBatchOutput(originalOutput, state.oldComputeOps)) { - demand.actions.push_back( - TensorDemandAction {.kind = TensorDemandActionKind::WholeTensorBarrier, - .destinationClass = std::nullopt, - .barrierReason = WholeTensorBarrierReason::FunctionReturnWithoutBlueprint}); - return demand; - } - - auto outputType = dyn_cast(originalOutput.getType()); - if (!outputType || !outputType.hasStaticShape()) - return targetClass.op->emitError("failed to classify compact batch output demand") - << " reason=terminal blueprint publication requires static ranked output metadata" - << " producerOp='" << sourceBatch->getName() << "' resultIndex=" << resultIndex - << " sourceClass=" << targetClass.id << " valueType=" << originalOutput.getType(); - - if (failed(getBatchResultProjectionInsert(sourceBatch, resultIndex))) - return targetClass.op->emitError("failed to classify compact batch output demand") - << " reason=terminal blueprint publication is missing projection metadata" - << " producerOp='" << sourceBatch->getName() << "' resultIndex=" << resultIndex - << " sourceClass=" << targetClass.id << " valueType=" << originalOutput.getType(); - - demand.actions.push_back(TensorDemandAction {.kind = TensorDemandActionKind::TerminalBlueprintPublication, - .destinationClass = std::nullopt, - .barrierReason = std::nullopt}); - return demand; -} - -bool hasWholeTensorBarrier(const RunOutputDemand& demand) { - return llvm::any_of(demand.actions, [](const TensorDemandAction& action) { - return action.kind == TensorDemandActionKind::WholeTensorBarrier; - }); -} - -FailureOr> tryBuildCompactRunPlan(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef run, - ArrayRef groups) { - if (run.size() < 2 || run.front().peers.empty()) - return std::optional {}; - - CompactRunPlan plan; - for (const OutputDestinationGroup& group : groups) { - for (size_t resultIndex : group.resultIndices) { - FailureOr demand = - classifyRunOutputDemand(state, targetClass, run, group.destinationClasses, resultIndex); - if (failed(demand)) - return failure(); - if (hasWholeTensorBarrier(*demand)) - return std::optional {}; - if (!demand->actions.empty()) - plan.outputs.push_back(std::move(*demand)); - } - } - - if (plan.outputs.empty()) - return std::optional {}; - return std::optional {std::move(plan)}; -} - -void markMaterializationRunSlots(MaterializerState& state, - ClassId classId, - SlotId startSlot, - ArrayRef run) { - for (auto slotIndex : llvm::seq(0, run.size())) - state.materializedLogicalSlots.insert({classId, startSlot + static_cast(slotIndex)}); -} - -void unmarkMaterializationRunSlots(MaterializerState& state, - ClassId classId, - SlotId startSlot, - ArrayRef run) { - for (auto slotIndex : llvm::seq(0, run.size())) - state.materializedLogicalSlots.erase({classId, startSlot + static_cast(slotIndex)}); -} - -LogicalResult materializeScalarBatchRun(MaterializerState& state, - MaterializedClass& targetClass, - SlotId startSlot, - ArrayRef run) { - assert(!targetClass.isBatch && "scalar batch run materialization expects scalar target class"); - assert(!run.empty() && "expected non-empty batch run"); - - markMaterializationRunSlots(state, targetClass.id, startSlot, run); - - SmallVector groups = groupBatchRunOutputsByDestination(state, run); - ArrayRef firstOriginalOutputs = getFirstMaterializationRunOriginalOutputs(state, run); - - auto sourceBatch = cast(getMaterializationRunSourceOp(run)); - SmallVector& fragmentTypes = getBatchOutputFragmentTypesCached(state, sourceBatch); - Location loc = getMaterializationRunLoc(run); - - for (const OutputDestinationGroup& group : groups) { - if (run.size() > 1 && group.destinationClasses.empty() - && !hasMaterializationRunGroupLiveExternalUse(state, run, group) - && !hasMaterializationRunGroupSameClassConsumer(state, targetClass.id, run, group) - && !hasMaterializationRunGroupProjectedConsumer(state, run, group)) { - for (size_t resultIndex : group.resultIndices) { - if (resultIndex >= fragmentTypes.size() || !fragmentTypes[resultIndex]) - return sourceBatch.emitOpError("failed to recover per-lane output type for deferred local packed run"); - - SmallVector keys = getMaterializationRunOutputKeys(run, resultIndex); - if (failed(registerDeferredLocalPackedRunValue(state, targetClass, keys, fragmentTypes[resultIndex], loc))) - return failure(); - } - - continue; - } - - FailureOr> packedOutputs = materializeBatchOutputGroupLoop(state, targetClass, run, group); - if (failed(packedOutputs)) { - if (hasDeferredMaterializationDependency(state)) { - unmarkMaterializationRunSlots(state, targetClass.id, startSlot, run); - return failure(); - } - InFlightDiagnostic diagnostic = sourceBatch.emitOpError("failed to materialize scalar batch output group") - << " targetClass=" << targetClass.id << " runSize=" << run.size(); - if (!run.empty() && !run.front().peers.empty()) - diagnostic << " firstLaneStart=" << run.front().peers.front().laneStart - << " firstLaneCount=" << run.front().peers.front().laneCount; - diagnostic << " resultIndices="; - for (size_t resultIndex : group.resultIndices) - diagnostic << " " << resultIndex; - return failure(); - } - - for (auto [groupOutputIndex, resultIndex] : llvm::enumerate(group.resultIndices)) { - Value packed = (*packedOutputs)[groupOutputIndex]; - if (resultIndex >= fragmentTypes.size() || !fragmentTypes[resultIndex]) - return sourceBatch.emitOpError("failed to recover per-lane output type for packed batch run"); - - Type fragmentType = fragmentTypes[resultIndex]; - SmallVector keys = getMaterializationRunOutputKeys(run, resultIndex); - - auto rankedFragmentType = cast(fragmentType); - Value representativeOriginalOutput = firstOriginalOutputs[resultIndex]; - FailureOr recordedProjectedHostFragments = - recordProjectedScalarHostFragmentsFromPackedRun(state, - targetClass, - sourceBatch, - resultIndex, - run, - packed, - rankedFragmentType, - representativeOriginalOutput, - loc); - if (failed(recordedProjectedHostFragments)) { - sourceBatch.emitOpError("failed to record projected scalar host fragments") - << " targetClass=" << targetClass.id << " resultIndex=" << resultIndex; - return failure(); - } - - if (run.size() == 1) { - if (*recordedProjectedHostFragments) { - if (failed(emitScalarSourceCommunication(state, targetClass, keys, packed, loc))) { - sourceBatch.emitOpError("failed to emit scalar source communication for projected host fragment") - << " targetClass=" << targetClass.id << " resultIndex=" << resultIndex; - return failure(); - } - if (failed(emitPlannedProjectedInputSendsForKeys(state, targetClass, keys, packed, loc))) { - sourceBatch.emitOpError("failed to emit projected input sends for scalar projected host fragment") - << " targetClass=" << targetClass.id << " resultIndex=" << resultIndex; - return failure(); - } - continue; - } - - if (failed(emitOutputFanout(state, targetClass, keys, packed, representativeOriginalOutput, loc))) { - sourceBatch.emitOpError("failed to emit singleton scalar batch output fanout") - << " targetClass=" << targetClass.id << " resultIndex=" << resultIndex; - return failure(); - } - continue; - } - - if (failed(emitPackedRunFanout(state, targetClass, group.destinationClasses, keys, packed, fragmentType, loc))) - return failure(); - - if (failed(registerPackedRunValue(state, targetClass, keys, packed, fragmentType, loc))) - return failure(); - - if (failed(emitPlannedProjectedInputSendsForKeys(state, targetClass, keys, packed, loc))) - return failure(); - - if (*recordedProjectedHostFragments) - continue; - - for (auto [runIndex, slot] : llvm::enumerate(run)) { - assert(slot.peers.size() == 1 && "scalar materialization run slot must contain exactly one peer"); - - ArrayRef originalOutputs = getComputeInstanceOutputValuesCached(state, slot.peers.front()); - Value originalOutput = originalOutputs[resultIndex]; - - if (!hasLiveExternalUseCached(state, originalOutput)) - continue; - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - Value fragment = getPackedSliceForRunIndex(state, targetClass.op, packed, rankedFragmentType, runIndex, loc); - - if (failed(emitHostCommunication(state, targetClass, fragment, originalOutput))) - return failure(); - } - } - } - - return success(); -} - -bool hasSameClassConsumer(MaterializerState& state, ProducerKey producerKey, ClassId classId) { - SameClassConsumerLookupKey lookupKey {producerKey.instance.op, producerKey.resultIndex, classId}; - auto it = state.sameClassConsumerIndex.find(lookupKey); - if (it == state.sameClassConsumerIndex.end()) - return false; - - for (ProducerKey existing : it->second) - if (containsProducerKey(existing, producerKey) || containsProducerKey(producerKey, existing)) - return true; - return false; -} - -Value createBatchRunFlatIndex(MaterializerState& state, MaterializedClass& targetClass, Value slotIndex, Location loc) { - auto batch = cast(targetClass.op); - auto laneArg = batch.getLaneArgument(); - assert(laneArg && "expected materialized compute_batch lane argument"); - - MLIRContext* context = state.func.getContext(); - AffineExpr d0 = getAffineDimExpr(0, context); - AffineExpr d1 = getAffineDimExpr(1, context); - - int64_t laneCount = static_cast(targetClass.cpus.size()); - AffineMap map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0, d0 * laneCount + d1); - return createOrFoldAffineApply(state.rewriter, loc, map, ValueRange {slotIndex, *laneArg}, state.func); -} - -Value createBatchClassRunSourceLane(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef run, - Value slotIndex, - Location loc) { - SmallVector sourceLanes; - sourceLanes.reserve(run.size() * targetClass.cpus.size()); - - for (auto [runSlotIndex, slot] : llvm::enumerate(run)) { - (void) runSlotIndex; - assert(slot.peers.size() == targetClass.cpus.size() && "expected one peer per materialized batch lane"); - for (const ComputeInstance& peer : slot.peers) - sourceLanes.push_back(peer.laneStart); - } - - Value flatIndex = createBatchRunFlatIndex(state, targetClass, slotIndex, loc); - return createIndexedIndexValue(state, - targetClass.op, - sourceLanes, - flatIndex, - loc, - static_cast(targetClass.cpus.size()), - /*allowExhaustiveTiledSearch=*/false); -} - -LogicalResult buildBatchRunSendPlans(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef run, - const CompactRunPlan& compactPlan, - SmallVectorImpl& plans) { - assert(sourceClass.isBatch && "batch run send planning expects a materialized batch source"); - - for (const RunOutputDemand& output : compactPlan.outputs) { - for (const TensorDemandAction& action : output.actions) { - if (action.kind != TensorDemandActionKind::DestinationFanout) - continue; - assert(action.destinationClass && "destination fanout action must carry a destination class"); - ClassId destinationClass = *action.destinationClass; - MaterializedClass& targetClass = state.classes[destinationClass]; - - if (targetClass.isBatch && targetClass.cpus.size() != sourceClass.cpus.size()) - return sourceClass.op->emitError( - "cannot compact batch run communication between batch classes of different sizes"); - - BatchRunSendPlan plan; - plan.resultIndex = output.resultIndex; - plan.destinationClass = destinationClass; - - size_t exchangeCount = run.size() * sourceClass.cpus.size(); - plan.exchangeIds.reserve(exchangeCount); - - for (size_t slotIndex = 0; slotIndex < run.size(); ++slotIndex) { - SmallVector slotKeys = getMaterializationRunSlotOutputKeys(run[slotIndex], output.resultIndex); - if (slotKeys.size() != sourceClass.cpus.size()) - return sourceClass.op->emitError("compact batch run send plan expected one key per source lane"); - for (ProducerKey key : slotKeys) { - ArrayRef ids = targetClass.isBatch - ? state.communicationPlan.getIndexedBatchRunExchanges(key, destinationClass) - : state.communicationPlan.getPackedScalarRunExchanges(key, destinationClass); - if (ids.empty()) - return sourceClass.op->emitError("compact batch run send plan is missing planned exchanges"); - plan.exchangeIds.push_back(ids.front()); - } - } - - plans.push_back(std::move(plan)); - } - } - - return success(); -} - -void appendBatchRunSend(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const BatchRunSendPlan& plan, - Value flatIndex, - Location loc) { - assert(sourceClass.isBatch && "batch run send expects a materialized batch source"); - - MaterializedClass& targetClass = state.classes[plan.destinationClass]; - const CommunicationExchange& firstExchange = state.communicationPlan.getExchange(plan.exchangeIds.front()); - FailureOr messages = - buildMessageVectorForExchangeIds(state, - plan.exchangeIds, - firstExchange.kind, - sourceClass.id, - targetClass.id, - payload.getType(), - sourceClass.op); - if (failed(messages)) - return; - std::optional preferredPeriod = static_cast(sourceClass.cpus.size()); - Value channelId = createIndexedChannelId(state, sourceClass.op, *messages, flatIndex, loc, preferredPeriod); - Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, *messages, flatIndex, loc, preferredPeriod); - Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, *messages, flatIndex, loc, preferredPeriod); - - SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); - for (CommunicationExchangeId id : plan.exchangeIds) { - CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; - emission.producerBound = true; - emission.producerPayload = payload; - emission.sendEmitted = true; - } -} - -LogicalResult appendPackedScalarRunReceives(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef run, - const BatchRunSendPlan& plan, - Type fragmentType, - Location loc) { - MaterializedClass& targetClass = state.classes[plan.destinationClass]; - assert(!targetClass.isBatch && "packed scalar run receives expect a scalar target class"); - - size_t laneCount = sourceClass.cpus.size(); - size_t receiveCount = run.size() * laneCount; - - if (receiveCount != plan.exchangeIds.size()) - return targetClass.op->emitError("inconsistent flattened batch run receive plan"); - - auto rankedFragmentType = dyn_cast(fragmentType); - if (!rankedFragmentType || !rankedFragmentType.hasStaticShape() || rankedFragmentType.getRank() == 0) - return targetClass.op->emitError("packed scalar run receive expects static ranked fragment type"); - - PackedScalarRunValue packedRun; - packedRun.targetClass = targetClass.id; - packedRun.sourceOp = run.front().peers.front().op; - packedRun.resultIndex = plan.resultIndex; - packedRun.kind = PackedScalarRunKind::DeferredReceive; - packedRun.fragmentType = rankedFragmentType; - - packedRun.exchangeIds = plan.exchangeIds; - - packedRun.slots.reserve(run.size()); - for (const MaterializationRunSlot& slot : run) { - PackedScalarRunSlot packedSlot; - packedSlot.keys = getMaterializationRunSlotOutputKeys(slot, plan.resultIndex); - packedRun.slots.push_back(std::move(packedSlot)); - } - - if (failed(validatePackedScalarRunMetadata(targetClass.op, packedRun))) - return failure(); - - state.availableValues.recordPackedRun(std::move(packedRun)); - return success(); -} - -LogicalResult recordIndexedBatchRunReceives(MaterializerState& state, - ArrayRef run, - const BatchRunSendPlan& plan, - Type fragmentType) { - MaterializedClass& targetClass = state.classes[plan.destinationClass]; - auto rankedFragmentType = dyn_cast(fragmentType); - if (!rankedFragmentType || !rankedFragmentType.hasStaticShape() || rankedFragmentType.getRank() == 0) - return targetClass.op->emitError("indexed batch run receive expects static ranked fragment type"); - - IndexedBatchRunValue indexedRun; - indexedRun.targetClass = targetClass.id; - indexedRun.sourceOp = run.front().peers.front().op; - indexedRun.resultIndex = plan.resultIndex; - indexedRun.fragmentType = rankedFragmentType; - indexedRun.exchangeIds = plan.exchangeIds; - indexedRun.slots.reserve(run.size()); - for (const MaterializationRunSlot& slot : run) { - PackedScalarRunSlot indexedSlot; - indexedSlot.keys = getMaterializationRunSlotOutputKeys(slot, plan.resultIndex); - indexedRun.slots.push_back(std::move(indexedSlot)); - } - - state.availableValues.recordIndexedBatchRun(std::move(indexedRun)); - return success(); -} - -LogicalResult recordLocalIndexedBatchRunValue(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef run, - size_t resultIndex, - Value packed, - RankedTensorType fragmentType) { - IndexedBatchRunValue indexedRun; - indexedRun.targetClass = targetClass.id; - indexedRun.sourceOp = run.front().peers.front().op; - indexedRun.resultIndex = resultIndex; - indexedRun.packed = packed; - indexedRun.fragmentType = fragmentType; - indexedRun.slots.reserve(run.size()); - for (const MaterializationRunSlot& slot : run) { - PackedScalarRunSlot indexedSlot; - indexedSlot.keys = getMaterializationRunSlotOutputKeys(slot, resultIndex); - indexedRun.slots.push_back(std::move(indexedSlot)); - } - - state.availableValues.recordIndexedBatchRun(std::move(indexedRun)); - return success(); -} - -LogicalResult appendBatchRunReceives(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef run, - const BatchRunSendPlan& plan, - Type fragmentType, - Location loc) { - MaterializedClass& targetClass = state.classes[plan.destinationClass]; - - if (!targetClass.isBatch) - return appendPackedScalarRunReceives(state, sourceClass, run, plan, fragmentType, loc); - return recordIndexedBatchRunReceives(state, run, plan, fragmentType); -} - -FailureOr> materializeCompactBatchOutputGroupLoop(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef run, - const CompactRunPlan& plan) { - assert(targetClass.isBatch && "compact batch output loop expects a batch target class"); - assert(!run.empty() && "expected non-empty compact batch run"); - assert(!run.front().peers.empty() && "expected non-empty compact batch run slot"); - - auto sourceBatch = dyn_cast(getMaterializationRunSourceOp(run)); - if (!sourceBatch) - return failure(); - - Location loc = getMaterializationRunLoc(run); - - SmallVector initValues; - initValues.reserve(plan.outputs.size()); - for (const RunOutputDemand& output : plan.outputs) { - FailureOr packedType = getPackedRunTensorType(output.fragmentType, run.size()); - if (failed(packedType)) - return sourceBatch.emitOpError("cannot compact batch run for non-static ranked output") - << " resultIndex=" << output.resultIndex; - initValues.push_back( - tensor::EmptyOp::create(state.rewriter, loc, packedType->getShape(), packedType->getElementType()).getResult()); - } - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(run.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - auto loop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange(initValues), - [&](OpBuilder&, Location, Value slotIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { - Value sourceLane = createBatchClassRunSourceLane(state, targetClass, run, slotIndex, loc); - SmallVector resultIndices; - resultIndices.reserve(plan.outputs.size()); - for (const RunOutputDemand& output : plan.outputs) - resultIndices.push_back(output.resultIndex); - - FailureOr> produced = - cloneBatchBodyForLane(state, - targetClass, - getScheduledChunkForLogicalInstance(state, run.front().peers.front()), - sourceLane, - resultIndices, - CloneIndexingContext {.runSlotIndex = slotIndex, .projectionSlotIndex = slotIndex}); - if (failed(produced)) - return failure(); - - yielded.reserve(produced->size()); - for (auto [outputIndex, output] : llvm::enumerate(*produced)) { - auto fragmentType = dyn_cast(output.getType()); - if (!fragmentType || !fragmentType.hasStaticShape()) - return failure(); - Value firstOffset = scaleIndexByDim0Size(state, targetClass.op, slotIndex, fragmentType.getDimSize(0), loc); - yielded.push_back(createDim0InsertSlice(state, loc, output, iterArgs[outputIndex], firstOffset)); - } - return success(); - }); - if (failed(loop)) - return failure(); - - SmallVector results; - results.reserve(loop->results.size()); - for (Value result : loop->results) - results.push_back(result); - return results; -} - -LogicalResult emitPackedBatchRunSends(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef run, - const CompactRunPlan& plan, - ArrayRef packedOutputs, - Location loc) { - SmallVector sendPlans; - if (failed(buildBatchRunSendPlans(state, targetClass, run, plan, sendPlans))) - return failure(); - if (sendPlans.empty()) - return success(); - - DenseMap packedOutputIndexByResult; - for (auto [packedIndex, output] : llvm::enumerate(plan.outputs)) - packedOutputIndexByResult[output.resultIndex] = packedIndex; - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(run.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - auto loop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {}, - [&](OpBuilder&, Location, Value slotIndex, ValueRange, SmallVectorImpl&) { - Value flatIndex = createBatchRunFlatIndex(state, targetClass, slotIndex, loc); - for (const BatchRunSendPlan& sendPlan : sendPlans) { - auto packedIt = packedOutputIndexByResult.find(sendPlan.resultIndex); - if (packedIt == packedOutputIndexByResult.end()) - return failure(); - - const RunOutputDemand& output = plan.outputs[packedIt->second]; - Value fragment = getPackedSliceForDynamicRunIndex( - state, targetClass.op, packedOutputs[packedIt->second], output.fragmentType, slotIndex, loc); - appendBatchRunSend(state, targetClass, fragment, sendPlan, flatIndex, loc); - } - return success(); - }); - if (failed(loop)) - return failure(); - - for (const BatchRunSendPlan& sendPlan : sendPlans) { - auto packedIt = packedOutputIndexByResult.find(sendPlan.resultIndex); - if (packedIt == packedOutputIndexByResult.end()) - return targetClass.op->emitError("missing packed output for compact batch run send plan"); - if (failed( - appendBatchRunReceives(state, targetClass, run, sendPlan, plan.outputs[packedIt->second].fragmentType, loc))) - return failure(); - } - - return success(); -} - -LogicalResult materializeBatchClassRun(MaterializerState& state, - MaterializedClass& targetClass, - SlotId startSlot, - ArrayRef run, - const CompactRunPlan& plan) { - assert(targetClass.isBatch && "batch-target run materialization expects a materialized batch class"); - assert(!run.empty() && "expected non-empty batch-target run"); - - markMaterializationRunSlots(state, targetClass.id, startSlot, run); - - auto sourceBatch = cast(run.front().peers.front().op); - Location loc = sourceBatch.getLoc(); - - FailureOr> packedOutputs = - materializeCompactBatchOutputGroupLoop(state, targetClass, run, plan); - if (failed(packedOutputs)) - return failure(); - - for (auto [packedIndex, output] : llvm::enumerate(plan.outputs)) { - SmallVector keys = getMaterializationRunOutputKeys(run, output.resultIndex); - - for (const TensorDemandAction& action : output.actions) { - switch (action.kind) { - case TensorDemandActionKind::DestinationFanout: break; - case TensorDemandActionKind::SameClassIndexedFragment: - if (failed(recordLocalIndexedBatchRunValue( - state, targetClass, run, output.resultIndex, (*packedOutputs)[packedIndex], output.fragmentType))) - return failure(); - break; - case TensorDemandActionKind::ProjectedInputPublication: - if (failed(emitPlannedProjectedInputSendsForKeys(state, targetClass, keys, (*packedOutputs)[packedIndex], loc))) - return failure(); - break; - case TensorDemandActionKind::TerminalBlueprintPublication: { - FailureOr recordedProjectedHostFragments = recordProjectedScalarHostFragmentsFromPackedValue( - state, targetClass, keys, (*packedOutputs)[packedIndex], output.originalOutput, loc); - if (failed(recordedProjectedHostFragments)) - return failure(); - if (!*recordedProjectedHostFragments) - return sourceBatch.emitOpError( - "compact batch blueprint publication requires explicit fragment assembly metadata") - << " resultIndex=" << output.resultIndex; - break; - } - case TensorDemandActionKind::WholeTensorBarrier: - return sourceBatch.emitOpError("compact batch materialization reached a whole-tensor barrier unexpectedly") - << " resultIndex=" << output.resultIndex - << " reason=" << describeWholeTensorBarrierReason(*action.barrierReason); - } - } - } - - if (failed(emitPackedBatchRunSends(state, targetClass, run, plan, *packedOutputs, loc))) - return failure(); - - return success(); -} - -LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeInstance& instance) { - auto cpuIt = state.schedule.computeToCpuMap.find(instance); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return instance.op->emitError("schedule materialization expected a CPU assignment for every compute instance"); - auto logicalRangeIt = state.scheduledInstanceToLogicalSlots.find(instance); - if (logicalRangeIt == state.scheduledInstanceToLogicalSlots.end()) - return instance.op->emitError("schedule materialization expected logical slots for every compute instance"); - - ClassId classId = state.cpuToClass.lookup(cpuIt->second); - MaterializedClass& targetClass = state.classes[classId]; - - LogicalSlotRange logicalRange = logicalRangeIt->second; - SlotId startLogicalSlot = logicalRange.start; - while (startLogicalSlot < logicalRange.start + logicalRange.count - && state.materializedLogicalSlots.contains({classId, startLogicalSlot})) { - ++startLogicalSlot; - } - if (startLogicalSlot == logicalRange.start + logicalRange.count) - return success(); - - if (isa(instance.op)) { - FailureOr run = - collectBatchMaterializationRun(state, targetClass, startLogicalSlot, instance.op); - - if (succeeded(run)) { - if (!targetClass.isBatch) - return materializeScalarBatchRun(state, targetClass, startLogicalSlot, *run); - - SmallVector groups = groupBatchRunOutputsByDestination(state, *run); - FailureOr> plan = tryBuildCompactRunPlan(state, targetClass, *run, groups); - if (failed(plan)) - return failure(); - if (*plan) { - if (failed(materializeBatchClassRun(state, targetClass, startLogicalSlot, *run, **plan))) - return failure(); - return success(); - } - } - } - - FailureOr> peers = - getMaterializationRunSlotPeers(state, targetClass, startLogicalSlot); - if (failed(peers)) - return instance.op->emitError("failed to collect peer compute instances for equivalence class logical slot"); - - Value projectionSlotIndex = getOrCreateIndexConstant( - state.constantFolder, targetClass.op, static_cast(startLogicalSlot - logicalRange.start)); - FailureOr> materializedOutputs = - cloneInstanceBody(state, - targetClass, - *peers, - CloneIndexingContext {.runSlotIndex = std::nullopt, .projectionSlotIndex = projectionSlotIndex}); - if (failed(materializedOutputs)) - return failure(); - - ArrayRef originalOutputs = getComputeInstanceOutputValuesCached(state, instance); - if (materializedOutputs->size() != originalOutputs.size()) - return instance.op->emitError("materialized output count does not match original compute instance output count"); - - for (auto [resultIndex, zipped] : llvm::enumerate(llvm::zip(*materializedOutputs, originalOutputs))) { - Value materializedOutput = std::get<0>(zipped); - Value originalOutput = std::get<1>(zipped); - MaterializationRunSlot slot; - slot.peers = *peers; - SmallVector keys = getMaterializationRunSlotOutputKeys(slot, resultIndex); - if (failed(emitOutputFanout(state, targetClass, keys, materializedOutput, originalOutput, instance.op->getLoc()))) - return failure(); - } - - state.materializedLogicalSlots.insert({classId, startLogicalSlot}); - return success(); -} - -enum class MaterializationAttemptKind { - Success, - DeferredCommunication, - Failure -}; - -static MaterializationAttemptKind attemptMaterializeInstanceSlot(MaterializerState& state, - const ComputeInstance& instance) { - state.deferredProducerValues.clear(); - - auto cpuIt = state.schedule.computeToCpuMap.find(instance); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return MaterializationAttemptKind::Failure; - ClassId classId = state.cpuToClass.lookup(cpuIt->second); - MaterializedClass& targetClass = state.classes[classId]; - if (failed(preflightProjectedInputsForInstance(state, targetClass, instance))) - return MaterializationAttemptKind::Failure; - if (!state.deferredProducerValues.empty()) - return MaterializationAttemptKind::DeferredCommunication; - - SmallVector newlyMarkedSlots; - auto rangeIt = state.scheduledInstanceToLogicalSlots.find(instance); - if (rangeIt != state.scheduledInstanceToLogicalSlots.end()) { - LogicalSlotRange range = rangeIt->second; - for (SlotId slot = range.start; slot < range.start + range.count; ++slot) { - ClassSlotKey classSlot {classId, slot}; - if (!state.materializedLogicalSlots.contains(classSlot)) - newlyMarkedSlots.push_back(classSlot); - } - } - - if (succeeded(materializeInstanceSlot(state, instance))) - return MaterializationAttemptKind::Success; - if (!state.deferredProducerValues.empty()) { - for (ClassSlotKey classSlot : newlyMarkedSlots) - state.materializedLogicalSlots.erase(classSlot); - return MaterializationAttemptKind::DeferredCommunication; - } - return MaterializationAttemptKind::Failure; -} - -static bool isScheduledInstanceFullyMaterialized(MaterializerState& state, const ComputeInstance& instance) { - auto cpuIt = state.schedule.computeToCpuMap.find(instance); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return false; - auto classIt = state.cpuToClass.find(cpuIt->second); - if (classIt == state.cpuToClass.end()) - return false; - auto rangeIt = state.scheduledInstanceToLogicalSlots.find(instance); - if (rangeIt == state.scheduledInstanceToLogicalSlots.end()) - return false; - - LogicalSlotRange range = rangeIt->second; - for (SlotId slot = range.start; slot < range.start + range.count; ++slot) - if (!state.materializedLogicalSlots.contains({classIt->second, slot})) - return false; - return true; -} - -static bool containsRemainingInstance(ArrayRef remaining, ComputeInstance instance) { - return llvm::is_contained(remaining, instance); -} - -static uint64_t getMaterializerOpId(Operation* op) { - return static_cast(reinterpret_cast(op)); -} - -static void appendDeferredProducerDiagnostic(InFlightDiagnostic& diag, - MaterializerState& state, - ProducerKey producer, - ArrayRef remaining) { - ComputeInstance scheduledInstance = getScheduledChunkForLogicalInstance(state, producer.instance); - diag << " [producerOp='" << producer.instance.op->getName() << "' laneStart=" << producer.instance.laneStart - << " laneCount=" << producer.instance.laneCount << " resultIndex=" << producer.resultIndex - << " opId=" << getMaterializerOpId(producer.instance.op) - << " sourceClass=" << getProducerSourceClassForDiagnostic(state, producer) - << " producerRemaining=" << containsRemainingInstance(remaining, producer.instance) - << " scheduledLaneStart=" << scheduledInstance.laneStart - << " scheduledLaneCount=" << scheduledInstance.laneCount - << " scheduledRemaining=" << containsRemainingInstance(remaining, scheduledInstance) - << " materialized=" << isProducerValueMaterialized(state, producer) - << " destinations="; - for (ClassId destination : getDestinationClasses(state, producer)) - diag << " " << destination; - diag << "]"; -} - -static LogicalResult emitCommunicationDeferralDiagnostic( - MaterializerState& state, - ArrayRef remaining, - const DenseMap, DenseMapInfo>& - deferredProducersByInstance) { - const ComputeInstance& instance = remaining.front(); - InFlightDiagnostic diag = instance.op->emitError("communication plan made no materialization progress") - << " remainingInstances=" << remaining.size() - << " schedule=communication-plan-events"; - unsigned reported = 0; - for (ProducerKey producer : state.deferredProducerValues) { - if (reported++ == 8) - break; - appendDeferredProducerDiagnostic(diag, state, producer, remaining); - } - unsigned exchangeReported = 0; - for (const ExchangeDescriptor& exchange : state.communicationPlan.getExchanges()) { - const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchange.id.value]; - if ((!emission.producerBound || emission.sendEmitted) && (emission.sendEmitted || emission.receiveEmitted)) - continue; - if (exchangeReported++ == 8) - break; - diag << " [exchange=" << exchange.id.value << " kind=" << static_cast(exchange.kind) - << " sourceClass=" << exchange.sourceClass << " targetClass=" << exchange.targetClass - << " sourceCore=" << exchange.sourceCore << " targetCore=" << exchange.targetCore - << " channel=" << exchange.channelId << " producerBound=" << emission.producerBound - << " sendEmitted=" << emission.sendEmitted << " receiveEmitted=" << emission.receiveEmitted << "]"; - } - unsigned remainingReported = 0; - for (const ComputeInstance& pending : remaining) { - if (remainingReported++ == 8) - break; - diag << " [pendingOp='" << pending.op->getName() << "' laneStart=" << pending.laneStart - << " laneCount=" << pending.laneCount << " opId=" << getMaterializerOpId(pending.op) << "]"; - auto producerIt = deferredProducersByInstance.find(pending); - if (producerIt != deferredProducersByInstance.end() && !producerIt->second.empty()) { - diag << " [pendingBlockedBy"; - appendDeferredProducerDiagnostic(diag, state, producerIt->second.front(), remaining); - diag << "]"; - } - } - return failure(); -} - -static LogicalResult materializeInstanceSlotsWithCommunicationPlan(MaterializerState& state) { - SmallVector remaining(state.schedule.dominanceOrderCompute.begin(), - state.schedule.dominanceOrderCompute.end()); - - while (!remaining.empty()) { - bool progress = false; - DenseMap, DenseMapInfo> deferredProducersByInstance; - for (size_t index = 0; index < remaining.size();) { - const ComputeInstance& instance = remaining[index]; - MaterializationAttemptKind attempt = attemptMaterializeInstanceSlot(state, instance); - if (attempt == MaterializationAttemptKind::Success) { - progress = true; - if (isScheduledInstanceFullyMaterialized(state, instance)) { - remaining.erase(remaining.begin() + index); - continue; - } - ++index; - continue; - } - if (attempt == MaterializationAttemptKind::Failure) { - InFlightDiagnostic diag = instance.op->emitError( - "merge schedule materialization failed while materializing scheduled compute instance"); - diag << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount; - return failure(); - } - deferredProducersByInstance[instance] = state.deferredProducerValues; - for (ProducerKey producer : state.deferredProducerValues) { - ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producer.instance); - if (state.schedule.computeToCpuMap.find(scheduledProducer) != state.schedule.computeToCpuMap.end() - && !isProducerValueMaterialized(state, producer) - && !containsRemainingInstance(remaining, scheduledProducer)) - remaining.push_back(scheduledProducer); - } - ++index; - } - - bool emittedAnyCommunication = false; - for (MaterializedClass& materializedClass : state.classes) { - FailureOr emitted = - emitReadyCommunicationEventsForClass(state, materializedClass, state.func.getLoc()); - if (failed(emitted)) - return failure(); - emittedAnyCommunication |= *emitted; - } - progress |= emittedAnyCommunication; - - if (progress) { - continue; - } - - return emitCommunicationDeferralDiagnostic(state, remaining, deferredProducersByInstance); - } - return success(); -} - -void replaceHostUses(MaterializerState& state) { - for (const auto& [oldValue, replacement] : state.hostReplacements) - replaceLiveExternalUses(oldValue, replacement, state.oldComputeOps); -} - -LogicalResult eraseOldComputeOps(MaterializerState& state) { - DenseSet seen; - for (const ComputeInstance& instance : state.schedule.dominanceOrderCompute) { - if (!seen.insert(instance.op).second) - continue; - instance.op->dropAllUses(); - instance.op->erase(); - } - return success(); -} - -LogicalResult verifyMaterializerDominanceForDebug(func::FuncOp func) { - DominanceInfo dominance(func); - WalkResult result = func.walk([&](Operation* op) -> WalkResult { - for (OpOperand& operand : op->getOpOperands()) { - Value value = operand.get(); - if (dominance.dominates(value, op)) - continue; - InFlightDiagnostic diag = op->emitError("materializer produced non-dominating operand") - << " operand#" << operand.getOperandNumber() - << " operandType=" << value.getType() - << " op=\"" << truncateMaterializerDebugString(stringifyOperationForMaterializerDebug(op)) - << "\""; - attachMaterializerValueOriginNote(diag, value, "non-dominating operand"); - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - return result.wasInterrupted() ? failure() : success(); -} - -} // namespace - -LogicalResult -MergeScheduleMaterializer::run(func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId) { - if (schedule.dominanceOrderCompute.empty()) - return success(); - - auto runStage = [&](StringRef stage, auto&& callback) -> LogicalResult { - if (succeeded(callback())) - return success(); - return func.emitError("merge schedule materialization failed during ") << stage; - }; - - MaterializerState state(func, schedule, nextChannelId); - if (failed(runStage("work-stream construction", [&]() { return buildMaterializationWorkStreams(state); }))) - return failure(); - if (failed(runStage("equivalence-class construction", [&]() { - return buildMaterializationClassesFromScheduleEquivalence(state); - }))) - return failure(); - if (failed(runStage("equivalence/logical-stream verification", [&]() { - return verifyScheduleEquivalenceMatchesLogicalStreams(state); - }))) - return failure(); - if (state.classes.empty()) - return success(); - - if (failed(runStage("host-output collection", [&]() { return collectHostOutputs(state); }))) - return failure(); - if (failed(runStage("empty materialized-op creation", [&]() { return createEmptyMaterializedOps(state); }))) - return failure(); - if (failed(runStage("projected-transfer collection", [&]() { return collectProjectedTransfers(state); }))) - return failure(); - if (failed(runStage("producer-destination collection", [&]() { return collectProducerDestinations(state); }))) - return failure(); - if (failed(runStage("global communication planning", [&]() { - return buildGlobalCommunicationPlan(state); - }))) - return failure(); - if (failed(runStage("instance-slot materialization", [&]() -> LogicalResult { - return materializeInstanceSlotsWithCommunicationPlan(state); - }))) - return failure(); - - if (failed(runStage("communication plan materialization verification", [&]() { - return verifyCommunicationPlanFullyMaterialized(state); - }))) - return failure(); - - if (failed(runStage("projected host-output finalization", [&]() { - return finalizeProjectedHostOutputFragments(state); - }))) - return failure(); - - if (failed(runStage("scheduled-body capture localization", [&]() -> LogicalResult { - for (MaterializedClass& materializedClass : state.classes) { - if (succeeded(localizeAllScheduledBodyCaptures(state, materializedClass))) - continue; - return materializedClass.op->emitError( - "merge schedule materialization failed while localizing scheduled-body captures") - << " classId=" << materializedClass.id; - } - return success(); - }))) - return failure(); - - replaceHostUses(state); - if (failed(runStage("old compute-op erasure", [&]() { return eraseOldComputeOps(state); }))) - return failure(); - - LogicalResult _ = runRegionDCE(state.rewriter, state.func.getBody()); - (void) _; - if (failed(runStage("materializer dominance verification", [&]() { - return verifyMaterializerDominanceForDebug(state.func); - }))) - return failure(); - - return success(); -} - -} // namespace spatial -} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.hpp deleted file mode 100644 index 3663086..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Support/LogicalResult.h" - -#include "Scheduling/MergeSchedule.hpp" - -namespace onnx_mlir { -namespace spatial { - -class MergeScheduleMaterializer { -public: - mlir::LogicalResult run(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId); -}; - -} // namespace spatial -} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp deleted file mode 100644 index 7337d2c..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp +++ /dev/null @@ -1,296 +0,0 @@ -#pragma once - -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Transforms/FoldUtils.h" - -#include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/SmallVector.h" - -#include - -#include "MaterializeMergeSchedule.hpp" -#include "CommunicationPlan.hpp" -#include "MergeMessages.hpp" -#include "MergeScheduleKeys.hpp" -#include "ProjectedFragments.hpp" -#include "Scheduling/ComputeInstanceUtils.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" - -namespace onnx_mlir::spatial { - -struct MaterializedClass { - ClassId id = 0; - llvm::SmallVector cpus; - mlir::Operation* op = nullptr; - mlir::Block* body = nullptr; - bool isBatch = false; - - llvm::DenseMap cpuToLane; - llvm::SmallVector weights; - llvm::SmallVector inputs; - llvm::SmallVector hostOutputs; - llvm::DenseMap publicationOutputToResultIndex; - llvm::DenseMap weightArgs; - llvm::DenseMap inputArgs; - llvm::DenseMap hostOutputToResultIndex; -}; - -struct PackedScalarRunSlot { - llvm::SmallVector keys; -}; - -enum class PackedScalarRunKind { - Materialized, - DeferredReceive, - DeferredLocalCompute -}; - -struct PackedScalarRunValue { - ClassId targetClass = 0; - mlir::Operation* sourceOp = nullptr; - size_t resultIndex = 0; - PackedScalarRunKind kind = PackedScalarRunKind::Materialized; - - mlir::Value packed; - - mlir::RankedTensorType fragmentType; - llvm::SmallVector slots; - llvm::SmallVector exchangeIds; -}; - -struct IndexedBatchRunValue { - ClassId targetClass = 0; - mlir::Operation* sourceOp = nullptr; - size_t resultIndex = 0; - mlir::Value packed; - mlir::RankedTensorType fragmentType; - llvm::SmallVector slots; - llvm::SmallVector exchangeIds; -}; - -struct LogicalSlotRange { - SlotId start = 0; - SlotId count = 0; -}; - -struct MaterializationRunSlot { - llvm::SmallVector peers; -}; - -using MaterializationRun = llvm::SmallVector; - -struct OutputDestinationGroup { - llvm::SmallVector resultIndices; - llvm::SmallVector destinationClasses; -}; - -struct BatchRunSendPlan { - size_t resultIndex = 0; - ClassId destinationClass = 0; - llvm::SmallVector exchangeIds; -}; - -enum class TensorDemandActionKind { - DestinationFanout, - SameClassIndexedFragment, - ProjectedInputPublication, - TerminalBlueprintPublication, - WholeTensorBarrier -}; - -enum class WholeTensorBarrierReason { - FunctionReturnWithoutBlueprint, - DenseLogicalConsumer -}; - -struct TensorDemandAction { - TensorDemandActionKind kind = TensorDemandActionKind::DestinationFanout; - std::optional destinationClass; - std::optional barrierReason; -}; - -struct RunOutputDemand { - size_t resultIndex = 0; - mlir::Value originalOutput; - mlir::RankedTensorType fragmentType; - llvm::SmallVector actions; -}; - -struct CompactRunPlan { - llvm::SmallVector outputs; -}; - -enum class BatchInputDemandKind { - LaneFragment, - ProjectedFragment, - WholeTensorBarrier -}; - -struct BatchInputDemand { - BatchInputDemandKind kind = BatchInputDemandKind::LaneFragment; - std::optional wholeTensorProducer; -}; - -struct CommunicationExchangeEmissionState { - bool producerBound = false; - bool sendEmitted = false; - bool receiveEmitted = false; - - mlir::Value producerPayload; - mlir::Value receiveValue; - - mlir::Operation* producerAnchor = nullptr; - mlir::Operation* sendAnchor = nullptr; - mlir::Operation* receiveAnchor = nullptr; - std::optional loc; -}; - -struct CloneIndexingContext { - std::optional runSlotIndex; - std::optional projectionSlotIndex; -}; - -struct MaterializerState; - -class AvailableValueStore { -public: - struct ExactBatchFragmentRecord { - ProducerKey key; - mlir::Value value; - }; - - void record(ProducerKey key, ClassId classId, mlir::Value value) { - exactValues[key][classId] = value; - - auto batch = mlir::dyn_cast_or_null(key.instance.op); - if (!batch || key.instance.laneCount == 0) - return; - - WholeBatchAssemblyLookupKey lookupKey {batch.getOperation(), key.resultIndex, classId}; - llvm::SmallVector& bucket = exactBatchFragmentsByProducerResultClass[lookupKey]; - for (ExactBatchFragmentRecord& record : bucket) { - if (!(record.key == key)) - continue; - record.value = value; - return; - } - bucket.push_back({key, value}); - } - - void recordPackedRun(PackedScalarRunValue run) { - size_t runIndex = packedScalarRuns.size(); - packedScalarRuns.push_back(std::move(run)); - const PackedScalarRunValue& storedRun = packedScalarRuns[runIndex]; - WholeBatchAssemblyLookupKey lookupKey {storedRun.sourceOp, storedRun.resultIndex, storedRun.targetClass}; - packedRunsByProducerResultClass[lookupKey].push_back(runIndex); - } - - void recordIndexedBatchRun(IndexedBatchRunValue run) { indexedBatchRuns.push_back(std::move(run)); } - void recordDeferredExact(ProducerKey key, ClassId classId, llvm::ArrayRef exchangeIds) { - llvm::SmallVector& bucket = exactExchangeIds[key][classId]; - for (CommunicationExchangeId id : exchangeIds) - if (!llvm::is_contained(bucket, id)) - bucket.push_back(id); - } - llvm::ArrayRef getDeferredExactExchangeIds(ProducerKey key, ClassId classId) const { - auto producerIt = exactExchangeIds.find(key); - if (producerIt == exactExchangeIds.end()) - return {}; - auto classIt = producerIt->second.find(classId); - if (classIt == producerIt->second.end()) - return {}; - return classIt->second; - } - bool hasDeferredExact(ProducerKey key, ClassId classId) const { - return !getDeferredExactExchangeIds(key, classId).empty(); - } - - std::optional lookupExact(ProducerKey key, ClassId classId) const; - std::optional lookup(MaterializerState& state, ProducerKey key, ClassId classId); - IndexedBatchRunValue* lookupIndexedBatchRun(ProducerKey key, ClassId classId); - - llvm::ArrayRef getPackedRunIndicesForWholeBatch(WholeBatchAssemblyLookupKey key) const { - auto it = packedRunsByProducerResultClass.find(key); - if (it == packedRunsByProducerResultClass.end()) - return {}; - return it->second; - } - - llvm::ArrayRef getExactFragmentsForWholeBatch(WholeBatchAssemblyLookupKey key) const { - auto it = exactBatchFragmentsByProducerResultClass.find(key); - if (it == exactBatchFragmentsByProducerResultClass.end()) - return {}; - return it->second; - } - - PackedScalarRunValue& getPackedRun(size_t index) { return packedScalarRuns[index]; } - -private: - std::optional lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId); - - llvm::DenseMap, ProducerKeyInfo> exactValues; - llvm::DenseMap>, - ProducerKeyInfo> - exactExchangeIds; - llvm::SmallVector packedScalarRuns; - llvm::SmallVector indexedBatchRuns; - llvm::DenseMap, - WholeBatchAssemblyLookupKeyInfo> - exactBatchFragmentsByProducerResultClass; - llvm::DenseMap, WholeBatchAssemblyLookupKeyInfo> - packedRunsByProducerResultClass; -}; - -struct MaterializerState { - mlir::func::FuncOp func; - const MergeScheduleResult& schedule; - mlir::IRRewriter rewriter; - mlir::OperationFolder constantFolder; - int64_t& nextChannelId; - llvm::SmallVector classes; - llvm::DenseMap cpuToClass; - llvm::DenseMap> logicalInstancesByCpu; - llvm::DenseMap scheduledInstanceToLogicalSlots; - llvm::DenseMap logicalInstanceToScheduledChunk; - llvm::DenseSet materializedLogicalSlots; - - llvm::DenseMap, ProducerKeyInfo> producerDestClasses; - llvm::DenseMap, ProducerKeyInfo> ordinaryProducerDestClasses; - llvm::DenseMap, SameClassConsumerLookupKeyInfo> - sameClassConsumerIndex; - llvm::DenseMap - projectedInputMatches; - llvm::DenseSet nonProjectedInputs; - llvm::DenseMap liveExternalUseCache; - llvm::DenseMap> batchOutputFragmentTypesCache; - llvm::DenseMap, llvm::DenseMapInfo> - computeInstanceOutputsCache; - llvm::DenseMap, ProducerKeyInfo> - projectedTransfers; - llvm::DenseMap> - projectedExtractReplacements; - llvm::DenseMap> - projectedInputTransferPlans; - CommunicationPlan communicationPlan; - llvm::SmallVector communicationExchangeStates; - llvm::SmallVector deferredProducerValues; - AvailableValueStore availableValues; - llvm::DenseMap hostReplacements; - llvm::DenseMap hostOutputOwners; - llvm::SmallVector pendingProjectedHostOutputFragments; - llvm::DenseSet oldComputeOps; - - MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId) - : func(func), - schedule(schedule), - rewriter(func.getContext()), - constantFolder(func.getContext()), - nextChannelId(nextChannelId) {} -}; - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp deleted file mode 100644 index a9f3995..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp +++ /dev/null @@ -1,398 +0,0 @@ -#include "mlir/Analysis/TopologicalSortUtils.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/IR/Location.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/IR/Region.h" -#include "mlir/IR/Value.h" -#include "mlir/IR/ValueRange.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Support/LLVM.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallSet.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/Support/raw_os_ostream.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "MaterializeMergeSchedule.hpp" -#include "SpatialDataflowCsvExporter.hpp" -#include "Scheduling/ComputeGraph.hpp" -#include "Scheduling/ComputeInstanceUtils.hpp" -#include "Scheduling/MergeSchedulingAnalysis.hpp" -#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/CompactAsmUtils.hpp" -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" -#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" - -using namespace mlir; - -namespace onnx_mlir { -namespace { -using namespace onnx_mlir::compact_asm; -using SpatCompute = spatial::SpatGraphCompute; -using SpatComputeBatch = spatial::SpatGraphComputeBatch; - -bool isTrivialSerialMergeCandidate(SpatCompute compute) { - if (!compute->hasOneUse()) - return false; - - auto& use = *compute->getUses().begin(); - auto user = dyn_cast(use.getOwner()); - return user && user.getInputs().size() == 1 && use.getOperandNumber() >= user.getWeights().size(); -} - -SmallVector appendMissingWeightsAndBuildIndexMap(SmallVectorImpl& targetWeights, - ValueRange sourceWeights) { - DenseMap> targetWeightIndices; - for (auto [weightIndex, weight] : llvm::enumerate(targetWeights)) - targetWeightIndices[weight].push_back(weightIndex); - - DenseMap usedSourceWeightOccurrences; - SmallVector sourceToTargetIndex; - sourceToTargetIndex.reserve(sourceWeights.size()); - for (Value weight : sourceWeights) { - size_t occurrence = usedSourceWeightOccurrences[weight]++; - auto& matchingIndices = targetWeightIndices[weight]; - if (occurrence >= matchingIndices.size()) { - size_t newIndex = targetWeights.size(); - targetWeights.push_back(weight); - matchingIndices.push_back(newIndex); - sourceToTargetIndex.push_back(newIndex); - continue; - } - sourceToTargetIndex.push_back(matchingIndices[occurrence]); - } - return sourceToTargetIndex; -} - -void mergeTriviallyConnectedComputes(func::FuncOp funcOp) { - Location loc = funcOp.getLoc(); - IRRewriter rewriter(funcOp->getContext()); - SmallVector trivialComputes; - llvm::SmallSet toErase; - - for (auto compute : funcOp.getOps()) - if (isTrivialSerialMergeCandidate(compute)) - trivialComputes.push_back(compute); - - while (!trivialComputes.empty()) { - auto compute = trivialComputes.front(); - - if (compute.use_empty()) { - std::swap(trivialComputes.front(), trivialComputes.back()); - trivialComputes.pop_back(); - continue; - } - - auto& computeUse = *compute->getUses().begin(); - auto child = cast(computeUse.getOwner()); - auto usedResult = cast(computeUse.get()).getResultNumber(); - auto childInputIndex = computeUse.getOperandNumber() - child.getWeights().size(); - - rewriter.setInsertionPointAfter(compute.getOperation()); - SmallVector mergedWeights(compute.getWeights().begin(), compute.getWeights().end()); - SmallVector childWeightToNewIndex = appendMissingWeightsAndBuildIndexMap(mergedWeights, child.getWeights()); - SmallVector mergedInputs(compute.getInputs().begin(), compute.getInputs().end()); - auto newCompute = SpatCompute::create(rewriter, loc, child.getResultTypes(), mergedWeights, mergedInputs); - Block* newBody = rewriter.createBlock(&newCompute.getBodyRegion()); - for (Value weight : mergedWeights) - newBody->addArgument(weight.getType(), loc); - for (Value input : mergedInputs) - newBody->addArgument(input.getType(), loc); - - IRMapping mapper; - for (auto [weightIndex, _] : llvm::enumerate(compute.getWeights())) { - auto oldWeightArg = compute.getWeightArgument(weightIndex); - auto newWeightArg = newCompute.getWeightArgument(weightIndex); - assert(oldWeightArg && newWeightArg && "expected compute weight block arguments"); - mapper.map(*oldWeightArg, *newWeightArg); - } - for (auto [inputIndex, _] : llvm::enumerate(compute.getInputs())) { - auto oldInputArg = compute.getInputArgument(inputIndex); - auto newInputArg = newCompute.getInputArgument(inputIndex); - assert(oldInputArg && newInputArg && "expected compute input block arguments"); - mapper.map(*oldInputArg, *newInputArg); - } - for (auto [oldIndex, weight] : llvm::enumerate(child.getWeights())) { - auto oldWeightArg = child.getWeightArgument(oldIndex); - auto newWeightArg = newCompute.getWeightArgument(childWeightToNewIndex[oldIndex]); - assert(oldWeightArg && newWeightArg && "expected child compute weight block arguments"); - mapper.map(*oldWeightArg, *newWeightArg); - } - - rewriter.setInsertionPointToEnd(newBody); - auto computeYield = cast(compute.getBody().front().getTerminator()); - for (Operation& op : compute.getBody().front().without_terminator()) - rewriter.clone(op, mapper); - auto childInputArg = child.getInputArgument(childInputIndex); - assert(childInputArg && "expected child compute input block argument"); - mapper.map(*childInputArg, mapper.lookupOrDefault(computeYield.getOperand(usedResult))); - - rewriter.setInsertionPointToEnd(newBody); - for (auto& op : child.getBody().front()) - rewriter.clone(op, mapper); - - child.replaceAllUsesWith(newCompute); - toErase.insert(child); - - std::swap(trivialComputes.front(), trivialComputes.back()); - trivialComputes.pop_back(); - toErase.insert(compute); - - if (isTrivialSerialMergeCandidate(newCompute)) - trivialComputes.push_back(newCompute); - } - - for (auto compute : toErase) { - for (Value result : compute->getResults()) - result.dropAllUses(); - compute.erase(); - } -} - -void generateReport(func::FuncOp funcOp, const std::string& name, size_t usedCpuCount = 0) { - std::fstream file = openReportFile(name); - if (!file.is_open()) - return; - llvm::raw_os_ostream os(file); - - struct ReportRow { - uint64_t id = 0; - uint64_t logicalComputeCount = 0; - uint64_t crossbarCount = 0; - uint64_t instructionCount = 0; - bool isRebatched = false; - SmallVector coreIds; - }; - - //TODO Used for report refactor - struct CollectorConcatRow { - uint64_t computeId = 0; - int32_t coreId = -1; - uint64_t operandCount = 0; - }; - - uint64_t totalComputeOps = 0; - uint64_t totalLogicalComputes = 0; - uint64_t totalBatchComputeOps = 0; - uint64_t totalInstructionCount = 0; - uint64_t totalCrossbarCount = 0; - uint64_t nextBatchId = 0; - //TODO Used for report refactor - std::vector collectedData; - //TODO Used for report refactor - std::vector collectorConcatRows; - - auto getPerInstanceCrossbarCount = [&](Operation* op) -> uint64_t { - return static_cast(spatial::collectDistinctCrossbarWeights(op).size()); - }; - - for (Operation& op : funcOp.getBody().front()) { - if (auto spatCompute = dyn_cast(&op)) { - uint64_t numInst = spatial::countComputeBodyInstructions(spatCompute.getBody()); - uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(spatCompute.getOperation()); - SmallVector coreIds; - auto coreId = getOptionalScheduledCoreId(spatCompute, "merge compute core id"); - if (failed(coreId)) - return; - if (*coreId) - coreIds.push_back(**coreId); - uint64_t computeId = totalComputeOps++; - collectedData.push_back({computeId, 1, perInstanceCrossbarCount, numInst, false, coreIds}); - uint64_t maxConcatOperands = 0; - spatCompute.getBody().walk([&](spatial::SpatConcatOp concatOp) { - maxConcatOperands = std::max(maxConcatOperands, concatOp.getInputs().size()); - }); - //TODO 128 is a magic number - if (maxConcatOperands >= 128 && !coreIds.empty()) - collectorConcatRows.push_back({computeId, coreIds.front(), maxConcatOperands}); - totalLogicalComputes += 1; - totalInstructionCount += numInst; - totalCrossbarCount += perInstanceCrossbarCount; - continue; - } - if (auto batch = dyn_cast(&op)) { - uint64_t numInst = spatial::countComputeBodyInstructions(batch.getBody()); - uint64_t logicalCount = static_cast(batch.getLaneCount()); - uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(batch.getOperation()); - SmallVector coreIds; - auto optionalCoreIds = getOptionalScheduledBatchCoreIds(batch, "merge compute_batch core id"); - if (failed(optionalCoreIds)) - return; - if (*optionalCoreIds) - coreIds = std::move(**optionalCoreIds); - collectedData.push_back( - {nextBatchId++, logicalCount, perInstanceCrossbarCount * logicalCount, numInst, true, coreIds}); - totalComputeOps += 1; - totalLogicalComputes += logicalCount; - totalBatchComputeOps += 1; - totalInstructionCount += numInst * logicalCount; - totalCrossbarCount += perInstanceCrossbarCount * logicalCount; - } - } - - llvm::SmallVector totalFields = { - {"Used cores", std::to_string(usedCpuCount) }, - {"Number of top-level compute ops", std::to_string(totalComputeOps) }, - {"Number of logical computes", std::to_string(totalLogicalComputes) }, - {"Number of top-level batch compute ops", std::to_string(totalBatchComputeOps) }, - {"Number of instructions", std::to_string(totalInstructionCount)}, - {"Number of used crossbars", std::to_string(totalCrossbarCount) } - }; - printReportTotalsBlock(os, totalFields); - if (!collectedData.empty() || !collectorConcatRows.empty()) - os << "\n"; - - if (!collectorConcatRows.empty()) { - os << "Collector concat materialization:\n"; - for (const CollectorConcatRow& row : collectorConcatRows) - os << "\tmaterialization_kind = single_collector_concat, compute = " << row.computeId - << ", concat_operand_count = " << row.operandCount << ", collector_core = " << row.coreId << "\n"; - os << "\n"; - } - - sortReportEntriesByFirstCore(collectedData); - - for (uint64_t cI = 0; cI < totalComputeOps; ++cI) { - uint64_t lastIndex = cI; - ReportRow current = collectedData[cI]; - - for (uint64_t nI = cI + 1; nI < totalComputeOps; ++nI) { - ReportRow next = collectedData[nI]; - if (current.isRebatched == next.isRebatched && current.crossbarCount == next.crossbarCount - && current.instructionCount == next.instructionCount - && current.logicalComputeCount == next.logicalComputeCount) - lastIndex = nI; - else - break; - } - - if (current.isRebatched) { - os << "Batch "; - for (uint64_t index = cI; index <= lastIndex; ++index) { - if (index != cI) - os << ",\n "; - os << collectedData[index].id << " (cores "; - if (collectedData[index].coreIds.empty()) - os << "unknown"; - else - printCompressedIntegerEntries(os, ArrayRef(collectedData[index].coreIds)); - os << ")"; - } - } - else { - os << "Compute "; - SmallVector opIds; - opIds.reserve(lastIndex - cI + 1); - for (uint64_t index = cI; index <= lastIndex; ++index) - opIds.push_back(collectedData[index].id); - printCompressedIntegerEntries(os, ArrayRef(opIds)); - } - - os << ":\n"; - uint64_t perCoreLogicalComputeCount = current.isRebatched ? 1 : current.logicalComputeCount; - uint64_t perCoreInstructionCount = current.instructionCount; - uint64_t perCoreCrossbarCount = - current.logicalComputeCount == 0 ? 0 : current.crossbarCount / current.logicalComputeCount; - uint64_t totalEntryInstructionCount = current.instructionCount * current.logicalComputeCount; - - llvm::SmallVector perCoreFields = { - {"Number of logical computes", std::to_string(perCoreLogicalComputeCount)}, - {"Number of instructions", std::to_string(perCoreInstructionCount) }, - {"Number of used crossbars", std::to_string(perCoreCrossbarCount) } - }; - if (current.isRebatched) { - llvm::SmallVector totalEntryFields = { - {"Number of logical computes", std::to_string(current.logicalComputeCount)}, - {"Number of instructions", std::to_string(totalEntryInstructionCount) }, - {"Number of used crossbars", std::to_string(current.crossbarCount) } - }; - printReportPerCoreAndTotalFields(os, perCoreFields, totalEntryFields); - } - else { - printReportFlatFields(os, perCoreFields); - } - printReportEntrySeparator(os, lastIndex + 1 < totalComputeOps); - cI = lastIndex; - } - - os.flush(); - file.close(); -} - -struct MergeComputeNodesPass : PassWrapper> { - -private: - int64_t nextChannelId = 0; - -public: - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass) - - StringRef getArgument() const override { return "pim-merge-compute-nodes-pass"; } - StringRef getDescription() const override { - return "Merge Spatial-Compute-Nodes in order to reduce the total " - "execution time"; - } - - LogicalResult initialize(MLIRContext* context) override { return success(); } - - void runOnOperation() override { - func::FuncOp func = getOperation(); - if (failed(verifyLogicalSpatialGraphInvariants(func))) { - func.emitOpError("logical Spatial graph verification failed at the start of MergeComputeNodes"); - signalPassFailure(); - return; - } - mergeTriviallyConnectedComputes(func); - if (failed(verifyLogicalSpatialGraphInvariants(func))) { - func.emitOpError("logical Spatial graph verification failed after trivial merge simplification"); - signalPassFailure(); - return; - } - - const spatial::MergeScheduleResult* analysisResult = nullptr; - analysisResult = &getAnalysis().getResult(); - spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage(); - if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) { - signalPassFailure(); - return; - } - - if (!sortTopologically(&func.getBody().front())) { - func.emitOpError("failed to topologically order merged Spatial IR"); - signalPassFailure(); - return; - } - if (failed(verifyScheduledSpatialInvariants(func))) { - func.emitOpError("scheduled Spatial verification failed after merge materialization"); - signalPassFailure(); - return; - } - if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Post) - && failed(spatial::exportSpatialDataflowCsvPost(func))) { - signalPassFailure(); - return; - } - dumpModule(cast(func->getParentOp()), "spatial2_merged"); - generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size()); - } -}; - -} // namespace - -std::unique_ptr createMergeComputeNodesPass() { return std::make_unique(); } - -} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeMessages.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeMessages.hpp deleted file mode 100644 index b992723..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeMessages.hpp +++ /dev/null @@ -1,67 +0,0 @@ -#pragma once - -#include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringRef.h" - -#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" - -namespace onnx_mlir::spatial { - -using CpuId = size_t; - -inline mlir::FailureOr getCheckedCoreId(mlir::Operation* anchor, CpuId cpu, llvm::StringRef fieldName) { - return pim::checkedI32(static_cast(cpu), anchor, fieldName); -} - -inline mlir::FailureOr> -getCheckedCoreIds(mlir::Operation* anchor, llvm::ArrayRef cpus, llvm::StringRef fieldName) { - llvm::SmallVector coreIds; - coreIds.reserve(cpus.size()); - for (CpuId cpu : cpus) { - auto checkedCoreId = getCheckedCoreId(anchor, cpu, fieldName); - if (mlir::failed(checkedCoreId)) - return mlir::failure(); - coreIds.push_back(*checkedCoreId); - } - return coreIds; -} - -struct MessageVector { - llvm::SmallVector channelIds; - llvm::SmallVector sourceCoreIds; - llvm::SmallVector targetCoreIds; - - size_t size() const { return channelIds.size(); } - bool empty() const { return channelIds.empty(); } - - mlir::LogicalResult verify(mlir::Operation* anchor) const { - if (channelIds.size() != sourceCoreIds.size() || channelIds.size() != targetCoreIds.size()) - return anchor->emitError("message metadata is inconsistent"); - return mlir::success(); - } - - void append(int64_t channelId, int32_t sourceCoreId, int32_t targetCoreId) { - channelIds.push_back(channelId); - sourceCoreIds.push_back(sourceCoreId); - targetCoreIds.push_back(targetCoreId); - } - - void append(llvm::ArrayRef channels, llvm::ArrayRef sources, llvm::ArrayRef targets) { - assert(channels.size() == sources.size() && "channel/source count mismatch"); - assert(channels.size() == targets.size() && "channel/target count mismatch"); - llvm::append_range(channelIds, channels); - llvm::append_range(sourceCoreIds, sources); - llvm::append_range(targetCoreIds, targets); - } - - MessageVector slice(size_t offset, size_t count) const { - MessageVector result; - result.append(llvm::ArrayRef(channelIds).slice(offset, count), - llvm::ArrayRef(sourceCoreIds).slice(offset, count), - llvm::ArrayRef(targetCoreIds).slice(offset, count)); - return result; - } -}; - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeScheduleKeys.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeScheduleKeys.hpp deleted file mode 100644 index 2fd1f2d..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeScheduleKeys.hpp +++ /dev/null @@ -1,134 +0,0 @@ -#pragma once - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/SmallVector.h" - -#include -#include -#include -#include - -#include "Scheduling/ComputeInstanceUtils.hpp" - -namespace onnx_mlir::spatial { - -using ClassId = size_t; -using SlotId = size_t; - -struct ProducerKey { - ComputeInstance instance; - size_t resultIndex = 0; - - bool operator==(const ProducerKey& other) const { - return instance == other.instance && resultIndex == other.resultIndex; - } -}; - -struct ProducerKeyInfo { - static ProducerKey getEmptyKey() { - return {llvm::DenseMapInfo::getEmptyKey(), std::numeric_limits::max()}; - } - - static ProducerKey getTombstoneKey() { - return {llvm::DenseMapInfo::getTombstoneKey(), std::numeric_limits::max()}; - } - - static unsigned getHashValue(const ProducerKey& key) { - return llvm::hash_combine(llvm::DenseMapInfo::getHashValue(key.instance), key.resultIndex); - } - - static bool isEqual(const ProducerKey& lhs, const ProducerKey& rhs) { return lhs == rhs; } -}; - -struct SameClassConsumerLookupKey { - mlir::Operation* sourceOp = nullptr; - size_t resultIndex = 0; - ClassId classId = 0; - - bool operator==(const SameClassConsumerLookupKey& other) const { - return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId; - } -}; - -struct SameClassConsumerLookupKeyInfo { - static SameClassConsumerLookupKey getEmptyKey() { - return {llvm::DenseMapInfo::getEmptyKey(), std::numeric_limits::max(), - std::numeric_limits::max()}; - } - - static SameClassConsumerLookupKey getTombstoneKey() { - return {llvm::DenseMapInfo::getTombstoneKey(), std::numeric_limits::max(), - std::numeric_limits::max()}; - } - - static unsigned getHashValue(const SameClassConsumerLookupKey& key) { - return llvm::hash_combine(llvm::DenseMapInfo::getHashValue(key.sourceOp), - key.resultIndex, - key.classId); - } - - static bool isEqual(const SameClassConsumerLookupKey& lhs, const SameClassConsumerLookupKey& rhs) { - return lhs == rhs; - } -}; - -struct WholeBatchAssemblyLookupKey { - mlir::Operation* sourceOp = nullptr; - size_t resultIndex = 0; - ClassId classId = 0; - - bool operator==(const WholeBatchAssemblyLookupKey& other) const { - return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId; - } -}; - -struct WholeBatchAssemblyLookupKeyInfo { - static WholeBatchAssemblyLookupKey getEmptyKey() { - return {llvm::DenseMapInfo::getEmptyKey(), std::numeric_limits::max(), - std::numeric_limits::max()}; - } - - static WholeBatchAssemblyLookupKey getTombstoneKey() { - return {llvm::DenseMapInfo::getTombstoneKey(), std::numeric_limits::max(), - std::numeric_limits::max()}; - } - - static unsigned getHashValue(const WholeBatchAssemblyLookupKey& key) { - return llvm::hash_combine(llvm::DenseMapInfo::getHashValue(key.sourceOp), - key.resultIndex, - key.classId); - } - - static bool isEqual(const WholeBatchAssemblyLookupKey& lhs, const WholeBatchAssemblyLookupKey& rhs) { - return lhs == rhs; - } -}; - -using ClassSlotKey = std::pair; - -struct ProjectedBatchInputKey { - mlir::Operation* consumerOp = nullptr; - unsigned inputIndex = 0; - - bool operator==(const ProjectedBatchInputKey& other) const { - return consumerOp == other.consumerOp && inputIndex == other.inputIndex; - } -}; - -struct ProjectedBatchInputKeyInfo { - static ProjectedBatchInputKey getEmptyKey() { - return {llvm::DenseMapInfo::getEmptyKey(), std::numeric_limits::max()}; - } - - static ProjectedBatchInputKey getTombstoneKey() { - return {llvm::DenseMapInfo::getTombstoneKey(), std::numeric_limits::max()}; - } - - static unsigned getHashValue(const ProjectedBatchInputKey& key) { - return llvm::hash_combine(key.consumerOp, key.inputIndex); - } - - static bool isEqual(const ProjectedBatchInputKey& lhs, const ProjectedBatchInputKey& rhs) { return lhs == rhs; } -}; - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.cpp deleted file mode 100644 index 1aa5610..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.cpp +++ /dev/null @@ -1,104 +0,0 @@ -#include "ProjectedFragments.hpp" - -#include "mlir/IR/BuiltinTypes.h" - -namespace onnx_mlir::spatial { - -static mlir::FailureOr getPackedBatchTensorType(mlir::Type laneType, size_t laneCount) { - auto tensorType = mlir::dyn_cast(laneType); - if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0) - return mlir::failure(); - - llvm::SmallVector shape(tensorType.getShape()); - shape[0] *= static_cast(laneCount); - return mlir::RankedTensorType::get(shape, tensorType.getElementType(), tensorType.getEncoding()); -} - -unsigned getProjectedFragmentsPerLogicalSlot(llvm::ArrayRef loopTripCounts) { - unsigned fragmentsPerLogicalSlot = 1; - for (int64_t tripCount : loopTripCounts) { - assert(tripCount > 0 && "projected loop trip counts must be positive"); - fragmentsPerLogicalSlot *= static_cast(tripCount); - } - return fragmentsPerLogicalSlot; -} - -mlir::LogicalResult verifyProjectedFragmentLayout(mlir::Operation* anchor, const ProjectedFragmentLayout& layout) { - if (!layout.fragmentType || layout.fragmentShape.empty()) - return anchor->emitError("projected fragment layout is missing fragment type metadata"); - if (layout.fragmentShape.size() != static_cast(layout.fragmentType.getRank())) - return anchor->emitError("projected fragment layout rank does not match fragment type"); - if (layout.payloadFragmentCount == 0 || layout.fragmentsPerLogicalSlot == 0) - return anchor->emitError("projected fragment layout has an invalid fragment count"); - if (layout.payloadFragmentCount % layout.fragmentsPerLogicalSlot != 0) - return anchor->emitError("projected fragment layout payload fragment count is incompatible with logical slots"); - return mlir::success(); -} - -mlir::FailureOr -getProjectedPayloadType(mlir::Operation* anchor, mlir::RankedTensorType fragmentType, unsigned payloadFragmentCount) { - auto packedType = getPackedBatchTensorType(fragmentType, payloadFragmentCount); - if (mlir::failed(packedType)) { - anchor->emitError("cannot create projected payload type"); - return mlir::failure(); - } - return *packedType; -} - -llvm::SmallVector, 4> -buildProjectedFragmentOffsetsByDim(llvm::ArrayRef> fragmentOffsets, size_t rank) { - llvm::SmallVector, 4> fragmentOffsetsByDim(rank); - for (llvm::ArrayRef offsets : fragmentOffsets) { - assert(offsets.size() == rank && "projected offset rank mismatch"); - for (size_t dim = 0; dim < rank; ++dim) - fragmentOffsetsByDim[dim].push_back(offsets[dim]); - } - return fragmentOffsetsByDim; -} - -mlir::LogicalResult verifyProjectedTransferDescriptor(mlir::Operation* anchor, - const ProjectedTransferDescriptor& descriptor) { - if (mlir::failed(verifyProjectedFragmentLayout(anchor, descriptor.layout))) - return mlir::failure(); - if (!descriptor.payloadType) - return anchor->emitError("projected transfer descriptor is missing payload type"); - if (descriptor.fragmentOffsets.empty()) - return anchor->emitError("projected transfer descriptor expected at least one fragment offset"); - if (descriptor.fragmentOffsetsByDim.size() != descriptor.layout.fragmentShape.size()) - return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent"); - for (llvm::ArrayRef dimOffsets : descriptor.fragmentOffsetsByDim) - if (dimOffsets.size() != descriptor.fragmentOffsets.size()) - return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent"); - for (llvm::ArrayRef offsets : descriptor.fragmentOffsets) - if (offsets.size() != descriptor.layout.fragmentShape.size()) - return anchor->emitError("projected transfer offset rank does not match fragment rank"); - return mlir::success(); -} - -mlir::LogicalResult verifyProjectedSendDescriptor(mlir::Operation* anchor, - const ProjectedTransferDescriptor& descriptor, - const MessageVector& messages) { - if (mlir::failed(verifyProjectedTransferDescriptor(anchor, descriptor))) - return mlir::failure(); - if (messages.size() * descriptor.layout.payloadFragmentCount != descriptor.fragmentOffsets.size()) - return anchor->emitError("projected send descriptor metadata is inconsistent"); - return mlir::success(); -} - -mlir::LogicalResult finalizeProjectedTransferDescriptor(mlir::Operation* anchor, - ProjectedTransferDescriptor& descriptor) { - descriptor.fragmentOffsetsByDim = - buildProjectedFragmentOffsetsByDim(descriptor.fragmentOffsets, descriptor.layout.fragmentShape.size()); - - auto payloadType = - getProjectedPayloadType(anchor, descriptor.layout.fragmentType, descriptor.layout.payloadFragmentCount); - if (mlir::failed(payloadType)) - return mlir::failure(); - if (descriptor.payloadType && descriptor.payloadType != *payloadType) - return anchor->emitError("projected transfer descriptor payload type does not match projected layout"); - descriptor.payloadType = *payloadType; - - return verifyProjectedTransferDescriptor(anchor, descriptor); -} - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp deleted file mode 100644 index 6ce50df..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp +++ /dev/null @@ -1,105 +0,0 @@ -#pragma once - -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Value.h" -#include "mlir/IR/ValueRange.h" - -#include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/SmallVector.h" - -#include - -#include "MergeMessages.hpp" -#include "MergeScheduleKeys.hpp" - -namespace onnx_mlir::spatial { - -struct ProjectedFragmentLayout { - mlir::RankedTensorType fragmentType; - llvm::SmallVector fragmentShape; - unsigned fragmentsPerLogicalSlot = 1; - unsigned payloadFragmentCount = 1; - llvm::SmallVector loopLowerBounds; - llvm::SmallVector loopSteps; - llvm::SmallVector loopTripCounts; -}; - -struct StaticProjectedLoopInfo { - mlir::BlockArgument iv; - int64_t lowerBound = 0; - int64_t step = 1; - int64_t tripCount = 1; -}; - -struct ProjectedTransferDescriptor { - ProjectedBatchInputKey inputKey; - mlir::Operation* extractOp = nullptr; - ProjectedFragmentLayout layout; - mlir::RankedTensorType payloadType; - llvm::SmallVector, 16> fragmentOffsets; - llvm::SmallVector, 4> fragmentOffsetsByDim; -}; - -struct ProjectedExtractReplacement { - mlir::Value payload; - ProjectedFragmentLayout layout; -}; - -struct ProjectedInputTransferFragment { - ProducerKey producer; - llvm::SmallVector fragmentOffsets; - unsigned targetLane = 0; - unsigned ordinal = 0; - int64_t channelId = 0; - int32_t sourceCoreId = 0; - int32_t targetCoreId = 0; - bool sendEmitted = false; -}; - -struct ProjectedInputTransferPlan { - ProjectedBatchInputKey inputKey; - mlir::Operation* extractOp = nullptr; - ProjectedFragmentLayout layout; - llvm::SmallVector fragments; -}; - -struct PendingProjectedHostOutputFragment { - mlir::Value originalOutput; - ClassId sourceClass = 0; - ProducerKey producerKey; - unsigned publicationResultIndex = 0; - int64_t sourceFragmentOrdinal = 0; - int64_t sourceElementOffset = 0; - llvm::SmallVector offsets; - llvm::SmallVector sizes; - llvm::SmallVector strides; - uint32_t sourceLane = 0; - mlir::Location loc; -}; - -struct AffineProjectedInputSliceMatch { - mlir::tensor::ExtractSliceOp extract; - mlir::RankedTensorType sourceType; - mlir::RankedTensorType fragmentType; - llvm::SmallVector fragmentShape; - llvm::SmallVector offsets; - llvm::SmallVector loops; -}; - -unsigned getProjectedFragmentsPerLogicalSlot(llvm::ArrayRef loopTripCounts); -mlir::LogicalResult verifyProjectedFragmentLayout(mlir::Operation* anchor, const ProjectedFragmentLayout& layout); -mlir::FailureOr -getProjectedPayloadType(mlir::Operation* anchor, mlir::RankedTensorType fragmentType, unsigned payloadFragmentCount); -llvm::SmallVector, 4> -buildProjectedFragmentOffsetsByDim(llvm::ArrayRef> fragmentOffsets, size_t rank); -mlir::LogicalResult verifyProjectedTransferDescriptor(mlir::Operation* anchor, - const ProjectedTransferDescriptor& descriptor); -mlir::LogicalResult verifyProjectedSendDescriptor(mlir::Operation* anchor, - const ProjectedTransferDescriptor& descriptor, - const MessageVector& messages); -mlir::LogicalResult finalizeProjectedTransferDescriptor(mlir::Operation* anchor, - ProjectedTransferDescriptor& descriptor); - -} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp index 6663c7b..b526828 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp @@ -106,10 +106,11 @@ static std::optional getConstantExtractLane(tensor::ExtractSliceOp ext static std::optional getResultfulBatchProducerValueRef(SpatComputeBatch batch, const ComputeInstance* consumerInstance) { - if (!consumerInstance) - return std::nullopt; - if (!isa(consumerInstance->op)) - return std::nullopt; + if (!consumerInstance || !isa(consumerInstance->op)) + return ProducerValueRef { + {batch.getOperation(), 0, static_cast(batch.getLaneCount())}, + 0 + }; if (consumerInstance->laneStart + consumerInstance->laneCount > static_cast(batch.getLaneCount())) return std::nullopt; return ProducerValueRef { diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp new file mode 100644 index 0000000..15bbf57 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp @@ -0,0 +1,832 @@ +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/AsmState.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/SmallString.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/raw_os_ostream.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include +#include + +#include "ComputeInstanceUtils.hpp" +#include "MergeSchedulingAnalysis.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Pass/PIMPasses.h" + +using namespace mlir; + +namespace onnx_mlir { +namespace spatial { +namespace { + +struct ProducerValueKey { + ComputeInstance instance; + size_t resultIndex = 0; + + bool operator==(const ProducerValueKey& other) const { + return instance == other.instance && resultIndex == other.resultIndex; + } +}; + +struct GraphComputeBlockKey { + Operation* op = nullptr; + uint32_t laneStart = 0; + uint32_t laneCount = 1; + + bool operator==(const GraphComputeBlockKey& other) const { + return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount; + } +}; + +struct ScheduledMaterializationRecord { + Operation* scheduledOp = nullptr; + SmallVector canonicalCpuClasses; + SmallVector cpus; + SmallVector computeKeys; + SmallVector blocks; +}; + +} // namespace +} // namespace spatial +} // namespace onnx_mlir + +namespace llvm { +template <> +struct DenseMapInfo { + static onnx_mlir::spatial::ProducerValueKey getEmptyKey() { + return {DenseMapInfo::getEmptyKey(), std::numeric_limits::max()}; + } + static onnx_mlir::spatial::ProducerValueKey getTombstoneKey() { + return {DenseMapInfo::getTombstoneKey(), std::numeric_limits::max()}; + } + static unsigned getHashValue(const onnx_mlir::spatial::ProducerValueKey& key) { + return hash_combine(DenseMapInfo::getHashValue(key.instance), key.resultIndex); + } + static bool isEqual(const onnx_mlir::spatial::ProducerValueKey& lhs, + const onnx_mlir::spatial::ProducerValueKey& rhs) { + return lhs == rhs; + } +}; + +template <> +struct DenseMapInfo { + static onnx_mlir::spatial::GraphComputeBlockKey getEmptyKey() { + return {DenseMapInfo::getEmptyKey(), UINT32_MAX, UINT32_MAX}; + } + static onnx_mlir::spatial::GraphComputeBlockKey getTombstoneKey() { + return {DenseMapInfo::getTombstoneKey(), UINT32_MAX, UINT32_MAX - 1}; + } + static unsigned getHashValue(const onnx_mlir::spatial::GraphComputeBlockKey& key) { + return hash_combine(key.op, key.laneStart, key.laneCount); + } + static bool isEqual(const onnx_mlir::spatial::GraphComputeBlockKey& lhs, + const onnx_mlir::spatial::GraphComputeBlockKey& rhs) { + return lhs == rhs; + } +}; +} // namespace llvm + +namespace onnx_mlir { +namespace spatial { +namespace { + +struct ScalarClass { + size_t cpu = 0; + SmallVector cpus; + SmallVector instances; + SmallVector weights; + SmallVector inputs; + SmallVector resultTypes; +}; + +struct CompactCpuIndexMap { + DenseMap rawToCompact; + + size_t lookup(size_t rawCpu) const { + auto it = rawToCompact.find(rawCpu); + assert(it != rawToCompact.end() && "missing compact cpu index"); + return it->second; + } +}; + +static ComputeInstance getWholeBatchInstance(SpatComputeBatch batch) { + return {batch.getOperation(), 0, static_cast(batch.getLaneCount())}; +} + +static GraphComputeBlockKey getGraphComputeBlockKey(const ComputeInstance& instance) { + return {instance.op, instance.laneStart, instance.laneCount}; +} + +static CompactCpuIndexMap buildCompactCpuIndexMap(const MergeScheduleResult& schedule) { + SmallVector rawCpus; + auto append = [&](size_t rawCpu) { + if (!llvm::is_contained(rawCpus, rawCpu)) + rawCpus.push_back(rawCpu); + }; + for (const auto& entry : schedule.computeToCpuMap) + append(entry.second); + for (const auto& entry : schedule.equivalentClass) { + append(entry.first); + for (size_t rawCpu : entry.second) + append(rawCpu); + } + for (const auto& entry : schedule.cpuToLastComputeMap) + append(entry.first); + llvm::sort(rawCpus); + + CompactCpuIndexMap compactMap; + for (auto [index, rawCpu] : llvm::enumerate(rawCpus)) + compactMap.rawToCompact[rawCpu] = index; + return compactMap; +} + +static SmallVector getEquivalentCpus(size_t cpu, const MergeScheduleResult& schedule) { + llvm::SmallDenseSet seen; + SmallVector cpus; + auto append = [&](size_t value) { + if (seen.insert(value).second) + cpus.push_back(value); + }; + append(cpu); + if (auto it = schedule.equivalentClass.find(cpu); it != schedule.equivalentClass.end()) + for (size_t equivalentCpu : it->second) + append(equivalentCpu); + llvm::sort(cpus); + return cpus; +} + +static size_t getCanonicalCpuClass(size_t cpu, const MergeScheduleResult& schedule) { + SmallVector cpus = getEquivalentCpus(cpu, schedule); + return cpus.empty() ? cpu : cpus.front(); +} + +static Value getProducerOutput(const ProducerValueRef& ref) { + auto outputs = getComputeInstanceOutputValues(ref.instance); + assert(ref.resultIndex < outputs.size() && "producer result index out of range"); + return outputs[ref.resultIndex]; +} + +static SmallVector getCanonicalCpuClassesForBatch(SpatComputeBatch batch, const MergeScheduleResult& schedule) { + llvm::SmallDenseSet seen; + SmallVector classes; + size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount()); + for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) { + ComputeInstance chunk = getBatchChunkForIndex(batch, chunkIndex); + auto cpuIt = schedule.computeToCpuMap.find(chunk); + if (cpuIt == schedule.computeToCpuMap.end()) + continue; + size_t cpuClass = getCanonicalCpuClass(cpuIt->second, schedule); + if (seen.insert(cpuClass).second) + classes.push_back(cpuClass); + } + llvm::sort(classes); + return classes; +} + +static SmallVector remapCpuList(ArrayRef rawCpus, const CompactCpuIndexMap& compactCpuMap) { + SmallVector compactCpus; + compactCpus.reserve(rawCpus.size()); + for (size_t rawCpu : rawCpus) + compactCpus.push_back(compactCpuMap.lookup(rawCpu)); + llvm::sort(compactCpus); + return compactCpus; +} + +static std::string getInstanceName(const ComputeInstance& instance) { + return llvm::formatv("{0}[lanes={1}:{2}]", + instance.op->getName().getStringRef(), + instance.laneStart, + instance.laneStart + instance.laneCount) + .str(); +} + +static void appendUnique(SmallVectorImpl& values, Value value) { + if (!llvm::is_contained(values, value)) + values.push_back(value); +} + +static Value getBlockOperand(Block& block, ValueRange operands, Value value, unsigned firstArgument = 0) { + auto it = llvm::find(operands, value); + assert(it != operands.end() && "missing scheduled operand"); + return block.getArgument(firstArgument + std::distance(operands.begin(), it)); +} + +static void appendComputeBlockArguments(SmallVectorImpl& argTypes, + SmallVectorImpl& argLocs, + ValueRange weights, + ValueRange inputs, + ArrayRef carriedKeys, + Location loc) { + for (Value weight : weights) + argTypes.push_back(weight.getType()); + for (Value input : inputs) + argTypes.push_back(input.getType()); + for (ProducerValueKey key : carriedKeys) { + auto outputs = getComputeInstanceOutputValues(key.instance); + assert(key.resultIndex < outputs.size() && "missing carried result"); + argTypes.push_back(outputs[key.resultIndex].getType()); + } + argLocs.append(argTypes.size(), loc); +} + +static Block* createScheduledComputeBlock(PatternRewriter& rewriter, + SpatScheduledCompute scheduled, + ArrayRef carriedKeys, + Location loc) { + SmallVector argTypes; + SmallVector argLocs; + appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(), scheduled.getInputs(), carriedKeys, loc); + return rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), TypeRange(argTypes), argLocs); +} + +static void appendBlockYieldBaseAndCarriedOperands(Block& block, + unsigned baseArgCount, + size_t carriedCount, + SmallVectorImpl& operands) { + for (unsigned index = 0; index < baseArgCount; ++index) + operands.push_back(block.getArgument(index)); + for (size_t index = 0; index < carriedCount; ++index) + operands.push_back(block.getArgument(baseArgCount + index)); +} + +static void createBlockYield(PatternRewriter& rewriter, Location loc, ValueRange outputs, Block* next = nullptr) { + OperationState state(loc, SpatBlockYieldOp::getOperationName()); + state.addOperands(outputs); + if (next) + state.addSuccessors(next); + rewriter.create(state); +} + +static std::string formatValueLabel(Value value, AsmState& asmState) { + std::string storage; + llvm::raw_string_ostream os(storage); + value.printAsOperand(os, asmState); + return storage; +} + +static std::string formatBlockLabel(Block* block, AsmState& asmState) { + std::string storage; + llvm::raw_string_ostream os(storage); + block->printAsOperand(os, asmState); + return storage; +} + +static std::string formatOperationLabel(Operation* op, AsmState& asmState) { + if (op->getNumResults() == 0) + return op->getName().getStringRef().str(); + std::string storage; + llvm::raw_string_ostream os(storage); + llvm::interleaveComma(op->getResults(), os, [&](Value result) { os << formatValueLabel(result, asmState); }); + return os.str(); +} + +static std::string formatGraphComputeBlockKey(const GraphComputeBlockKey& key, AsmState& asmState) { + return llvm::formatv("{0} {1} lanes [{2}:{3}]", + formatOperationLabel(key.op, asmState), + key.op->getName().getStringRef(), + key.laneStart, + key.laneStart + key.laneCount) + .str(); +} + +static size_t countPeftEquivalenceClasses(const MergeScheduleResult& schedule) { + llvm::SmallDenseSet classes; + for (const ComputeInstance& instance : schedule.dominanceOrderCompute) { + auto cpuIt = schedule.computeToCpuMap.find(instance); + if (cpuIt == schedule.computeToCpuMap.end()) + continue; + classes.insert(getCanonicalCpuClass(cpuIt->second, schedule)); + } + return classes.size(); +} + +static SmallVector collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) { + SmallVector keys; + for (Operation& op : funcOp.getOps()) { + if (auto compute = dyn_cast(&op)) + keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1})); + else if (auto batch = dyn_cast(&op)) + keys.push_back(getGraphComputeBlockKey(getWholeBatchInstance(batch))); + } + return keys; +} + +static LogicalResult verifyMaterializedScheduleMapping( + func::FuncOp funcOp, + const MergeScheduleResult& schedule, + const DenseMap& graphComputeToBlockMap, + ArrayRef materializedSchedules) { + pim::CappedDiagnosticReporter diagnostics; + size_t expectedClassCount = countPeftEquivalenceClasses(schedule); + if (expectedClassCount != materializedSchedules.size()) { + diagnostics.report(funcOp.getOperation(), [&](Operation* illegalOp) { + illegalOp->emitOpError() << "phase-check expected " << expectedClassCount + << " PEFT equivalence classes but materialized " << materializedSchedules.size() + << " scheduled computes"; + }); + } + + for (const ScheduledMaterializationRecord& record : materializedSchedules) { + if (record.canonicalCpuClasses.empty()) { + diagnostics.report(record.scheduledOp, [&](Operation* illegalOp) { + illegalOp->emitOpError("phase-check scheduled compute is missing an owning PEFT equivalence class"); + }); + continue; + } + if (record.canonicalCpuClasses.size() > 1) { + diagnostics.report(record.scheduledOp, [&](Operation* illegalOp) { + std::string classes; + llvm::raw_string_ostream os(classes); + llvm::interleaveComma(record.canonicalCpuClasses, os, [&](size_t cpuClass) { os << cpuClass; }); + illegalOp->emitOpError("phase-check scheduled compute spans multiple PEFT equivalence classes") + << " (classes " << os.str() << ")"; + }); + } + } + + for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) { + if (graphComputeToBlockMap.count(key)) + continue; + diagnostics.report(key.op, [&](Operation* illegalOp) { + illegalOp->emitOpError() << "phase-check graph compute is missing a scheduled MLIR block mapping for lanes [" + << key.laneStart << ":" << (key.laneStart + key.laneCount) << "]"; + }); + } + + for (const auto& entry : graphComputeToBlockMap) { + Block* block = entry.second; + if (!block || !isa(block->getParentOp())) { + diagnostics.report(entry.first.op, [&](Operation* illegalOp) { + illegalOp->emitOpError("phase-check graph compute block mapping does not target a scheduled compute block"); + }); + } + } + + diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial materialization verification failed"); + return success(!diagnostics.hasFailure()); +} + +static void dumpPeftMaterializationReport(func::FuncOp funcOp, + const MergeScheduleResult& schedule, + const CompactCpuIndexMap& compactCpuMap, + ArrayRef materializedSchedules) { + std::fstream file = openDialectDumpFileWithExtension("spatial1_schedule_materialization", "/reports", "txt"); + if (!file.is_open()) + return; + + llvm::raw_os_ostream os(file); + AsmState asmState(funcOp); + + llvm::MapVector> classToCpus; + DenseMap> cpuToInstances; + for (const ComputeInstance& instance : schedule.dominanceOrderCompute) { + auto cpuIt = schedule.computeToCpuMap.find(instance); + if (cpuIt == schedule.computeToCpuMap.end()) + continue; + size_t cpu = compactCpuMap.lookup(cpuIt->second); + size_t cpuClass = compactCpuMap.lookup(getCanonicalCpuClass(cpuIt->second, schedule)); + if (!llvm::is_contained(classToCpus[cpuClass], cpu)) + classToCpus[cpuClass].push_back(cpu); + cpuToInstances[cpu].push_back(instance); + } + + os << "PEFT Schedule\n"; + os << "=============\n"; + os << "equivalence classes: " << classToCpus.size() << "\n\n"; + for (auto& entry : classToCpus) { + size_t cpuClass = entry.first; + llvm::sort(entry.second); + os << "class " << cpuClass << " cpus ["; + llvm::interleaveComma(entry.second, os, [&](size_t cpu) { os << cpu; }); + os << "]\n"; + for (size_t cpu : entry.second) { + os << " cpu " << cpu << ":\n"; + auto& instances = cpuToInstances[cpu]; + llvm::sort(instances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) { + return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs); + }); + for (const ComputeInstance& instance : instances) { + os << " slot " << schedule.computeToCpuSlotMap.lookup(instance) << ": " + << formatGraphComputeBlockKey(getGraphComputeBlockKey(instance), asmState) << "\n"; + } + } + os << "\n"; + } + + os << "Materialized Scheduled Computes\n"; + os << "===============================\n"; + os << "scheduled ops: " << materializedSchedules.size() << "\n\n"; + for (const ScheduledMaterializationRecord& record : materializedSchedules) { + os << formatOperationLabel(record.scheduledOp, asmState) << " " << record.scheduledOp->getName().getStringRef(); + os << " classes ["; + llvm::interleaveComma(record.canonicalCpuClasses, os, [&](size_t cpuClass) { os << cpuClass; }); + os << "] cpus ["; + llvm::interleaveComma(record.cpus, os, [&](size_t cpu) { os << cpu; }); + os << "]\n"; + for (auto [key, block] : llvm::zip(record.computeKeys, record.blocks)) + os << " " << formatGraphComputeBlockKey(key, asmState) << " -> " << formatBlockLabel(block, asmState) << "\n"; + os << "\n"; + } +} + +static LogicalResult createDeferredCommunication(PatternRewriter& rewriter, + Location loc, + Value consumerInput, + Value source, + const ProducerValueRef& producer, + const ComputeInstance& consumer, + const MergeScheduleResult& schedule, + const CompactCpuIndexMap& compactCpuMap, + uint64_t exchangeId, + Value& result) { + auto sourceCpu = compactCpuMap.lookup(schedule.computeToCpuMap.lookup(producer.instance)); + auto targetCpu = compactCpuMap.lookup(schedule.computeToCpuMap.lookup(consumer)); + auto transfer = SpatDeferredCommunicationOp::create( + rewriter, + loc, + consumerInput.getType(), + ValueRange {source}, + rewriter.getStringAttr(("x" + std::to_string(exchangeId)).c_str()), + rewriter.getStringAttr(getInstanceName(producer.instance)), + rewriter.getStringAttr(getInstanceName(consumer)), + rewriter.getI64IntegerAttr(static_cast(sourceCpu)), + rewriter.getI64IntegerAttr(static_cast(targetCpu)), + rewriter.getI64IntegerAttr(static_cast(sourceCpu)), + rewriter.getI64IntegerAttr(static_cast(targetCpu)), + rewriter.getI64IntegerAttr(producer.instance.laneStart), + rewriter.getI64IntegerAttr(consumer.laneStart), + rewriter.getStringAttr(consumerInput.getDefiningOp() ? "projected" : "ordinary"), + rewriter.getI64IntegerAttr(static_cast(producer.resultIndex)), + rewriter.getDenseI64ArrayAttr({}), + rewriter.getI64IntegerAttr(-1)); + + Block* block = rewriter.createBlock(&transfer.getBody(), transfer.getBody().end(), source.getType(), loc); + rewriter.setInsertionPointToStart(block); + Value yielded = block->getArgument(0); + if (yielded.getType() != consumerInput.getType()) { + auto slice = consumerInput.getDefiningOp(); + if (!slice || slice.getSource().getType() != source.getType()) + return transfer.emitOpError("cannot derive deferred communication payload type from recorded source"); + IRMapping mapper; + mapper.map(slice.getSource(), yielded); + auto cloned = cast(rewriter.clone(*slice, mapper)); + yielded = cloned.getResult(); + } + SpatYieldOp::create(rewriter, loc, yielded); + result = transfer.getOutput(); + rewriter.setInsertionPointAfter(transfer); + return success(); +} + +static LogicalResult mapComputeInputs(PatternRewriter& rewriter, + Block& block, + SpatCompute compute, + const ComputeInstance& instance, + const ScalarClass& klass, + const MergeScheduleResult& schedule, + const CompactCpuIndexMap& compactCpuMap, + ValueRange scheduledWeights, + ValueRange scheduledInputs, + IRMapping& mapper, + uint64_t& exchangeId) { + for (auto [index, weight] : llvm::enumerate(compute.getWeights())) + mapper.map(*compute.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight)); + + for (auto [index, input] : llvm::enumerate(compute.getInputs())) { + Value mapped; + std::optional producer = getProducerValueRef(input, &instance); + if (producer) { + Value source = getProducerOutput(*producer); + if (failed(createDeferredCommunication( + rewriter, input.getLoc(), input, source, *producer, instance, schedule, compactCpuMap, exchangeId++, mapped))) + return failure(); + } + else { + mapped = getBlockOperand(block, scheduledInputs, input, scheduledWeights.size()); + } + mapper.map(*compute.getInputArgument(index), mapped); + } + return success(); +} + +static LogicalResult mapBatchInputs(PatternRewriter& rewriter, + Block& block, + SpatComputeBatch batch, + const MergeScheduleResult& schedule, + const CompactCpuIndexMap& compactCpuMap, + ValueRange scheduledWeights, + ValueRange scheduledInputs, + IRMapping& mapper, + uint64_t& exchangeId) { + mapper.map(*batch.getLaneArgument(), block.getArgument(0)); + unsigned firstWeightArg = 1; + unsigned firstInputArg = firstWeightArg + scheduledWeights.size(); + for (auto [index, weight] : llvm::enumerate(batch.getWeights())) + mapper.map(*batch.getWeightArgument(index), block.getArgument(firstWeightArg + index)); + + ComputeInstance consumer {batch.getOperation(), 0, static_cast(batch.getLaneCount())}; + for (auto [index, input] : llvm::enumerate(batch.getInputs())) { + Value mapped; + std::optional producer = getProducerValueRef(input, &consumer); + if (producer) { + Value source = getProducerOutput(*producer); + if (failed(createDeferredCommunication( + rewriter, input.getLoc(), input, source, *producer, consumer, schedule, compactCpuMap, exchangeId++, mapped))) + return failure(); + } + else { + auto it = llvm::find(scheduledInputs, input); + if (it == scheduledInputs.end()) + return batch.emitOpError("host batch input is missing from scheduled batch operands"); + mapped = block.getArgument(firstInputArg + std::distance(scheduledInputs.begin(), it)); + } + mapper.map(*batch.getInputArgument(index), mapped); + } + + unsigned firstOutputArg = firstInputArg + scheduledInputs.size(); + for (unsigned index = 0; index < batch.getNumResults(); ++index) + mapper.map(*batch.getOutputArgument(index), block.getArgument(firstOutputArg + index)); + return success(); +} + +static void cloneComputeBody(PatternRewriter& rewriter, + Block& source, + IRMapping& mapper, + SmallVectorImpl& yieldedValues) { + for (Operation& op : source.without_terminator()) + rewriter.clone(op, mapper); + auto yield = cast(source.getTerminator()); + for (Value output : yield.getOutputs()) + yieldedValues.push_back(mapper.lookup(output)); +} + +struct MergeComputeNodesPass final : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass) + + StringRef getArgument() const override { return "pim-merge-compute-nodes"; } + StringRef getDescription() const override { + return "Materialize scheduled Spatial compute with deferred communication placeholders."; + } + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + auto entryFunc = getPimEntryFunc(moduleOp); + if (failed(entryFunc)) { + moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes"); + signalPassFailure(); + return; + } + + func::FuncOp funcOp = *entryFunc; + MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult(); + CompactCpuIndexMap compactCpuMap = buildCompactCpuIndexMap(schedule); + llvm::MapVector scalarClasses; + SmallVector batchOps; + for (const ComputeInstance& instance : schedule.dominanceOrderCompute) { + if (isa(instance.op)) { + size_t cpu = schedule.computeToCpuMap.lookup(instance); + size_t cpuClass = getCanonicalCpuClass(cpu, schedule); + auto& klass = scalarClasses[cpuClass]; + klass.cpu = cpuClass; + if (klass.cpus.empty()) + klass.cpus = getEquivalentCpus(cpu, schedule); + klass.instances.push_back(instance); + } + else { + auto batch = cast(instance.op); + if (!llvm::is_contained(batchOps, batch)) + batchOps.push_back(batch); + } + } + for (auto& entry : scalarClasses) { + llvm::sort(entry.second.instances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) { + return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs); + }); + } + + PatternRewriter rewriter(moduleOp.getContext()); + Operation* insertionPoint = funcOp.getBody().front().getTerminator(); + DenseMap scalarOps; + DenseMap batchReplacements; + DenseMap scheduledOpToRecordIndex; + DenseMap graphComputeToBlockMap; + std::vector materializedSchedules; + + for (auto& entry : scalarClasses) { + ScalarClass& klass = entry.second; + for (const ComputeInstance& instance : klass.instances) { + auto compute = cast(instance.op); + llvm::append_range(klass.resultTypes, compute.getResultTypes()); + for (Value weight : compute.getWeights()) + appendUnique(klass.weights, weight); + for (Value input : compute.getInputs()) { + std::optional producer = getProducerValueRef(input, &instance); + if (!producer) + appendUnique(klass.inputs, input); + } + } + rewriter.setInsertionPoint(insertionPoint); + auto scheduled = SpatScheduledCompute::create( + rewriter, klass.instances.front().op->getLoc(), TypeRange(klass.resultTypes), klass.weights, klass.inputs); + scheduled->setAttr( + kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast(compactCpuMap.lookup(klass.cpu)))); + SmallVector equivalentCpus; + for (size_t cpu : klass.cpus) + equivalentCpus.push_back(static_cast(compactCpuMap.lookup(cpu))); + scheduled->setAttr("scheduled.equivalent_cpus", rewriter.getDenseI64ArrayAttr(equivalentCpus)); + SmallVector sources; + SmallVector lanes; + for (const ComputeInstance& instance : klass.instances) { + sources.push_back(rewriter.getStringAttr(getInstanceName(instance))); + lanes.push_back(instance.laneStart); + lanes.push_back(instance.laneCount); + } + scheduled->setAttr("scheduled.sources", rewriter.getArrayAttr(sources)); + scheduled->setAttr("scheduled.lane_ranges", rewriter.getDenseI64ArrayAttr(lanes)); + scalarOps[&klass] = scheduled; + ScheduledMaterializationRecord record; + record.scheduledOp = scheduled.getOperation(); + record.canonicalCpuClasses.push_back(compactCpuMap.lookup(klass.cpu)); + record.cpus = remapCpuList(klass.cpus, compactCpuMap); + scheduledOpToRecordIndex[scheduled.getOperation()] = materializedSchedules.size(); + materializedSchedules.push_back(std::move(record)); + } + + for (SpatComputeBatch batch : batchOps) { + SmallVector inputs; + ComputeInstance wholeBatch = getWholeBatchInstance(batch); + for (Value input : batch.getInputs()) + if (!getProducerValueRef(input, &wholeBatch)) + appendUnique(inputs, input); + + rewriter.setInsertionPoint(insertionPoint); + auto scheduled = SpatScheduledComputeBatch::create(rewriter, + batch.getLoc(), + batch.getResultTypes(), + rewriter.getI32IntegerAttr(batch.getLaneCount()), + batch.getWeights(), + inputs); + SmallVector coreIds; + uint32_t laneCount = batch.getLaneCount(); + coreIds.reserve(laneCount); + for (uint32_t lane = 0; lane < laneCount; ++lane) { + ComputeInstance chunk = getBatchChunkForLane(batch, lane); + coreIds.push_back(static_cast(compactCpuMap.lookup(schedule.computeToCpuMap.lookup(chunk)))); + } + scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(coreIds)); + scheduled->setAttr("scheduled.sources", rewriter.getArrayAttr({rewriter.getStringAttr(getInstanceName(wholeBatch))})); + batchReplacements[batch.getOperation()] = scheduled; + ScheduledMaterializationRecord record; + record.scheduledOp = scheduled.getOperation(); + record.canonicalCpuClasses = remapCpuList(getCanonicalCpuClassesForBatch(batch, schedule), compactCpuMap); + for (int32_t coreId : coreIds) + if (!llvm::is_contained(record.cpus, static_cast(coreId))) + record.cpus.push_back(static_cast(coreId)); + scheduledOpToRecordIndex[scheduled.getOperation()] = materializedSchedules.size(); + materializedSchedules.push_back(std::move(record)); + } + + uint64_t exchangeId = 0; + for (auto& entry : scalarClasses) { + ScalarClass& klass = entry.second; + SpatScheduledCompute scheduled = scalarOps.lookup(&klass); + SmallVector carriedKeys; + Block* block = nullptr; + for (auto [ordinal, instance] : llvm::enumerate(klass.instances)) { + auto compute = cast(instance.op); + if (!block) + block = createScheduledComputeBlock(rewriter, scheduled, carriedKeys, compute.getLoc()); + GraphComputeBlockKey key = getGraphComputeBlockKey(instance); + graphComputeToBlockMap[key] = block; + auto& record = materializedSchedules[scheduledOpToRecordIndex.lookup(scheduled.getOperation())]; + record.computeKeys.push_back(key); + record.blocks.push_back(block); + rewriter.setInsertionPointToStart(block); + IRMapping mapper; + if (failed(mapComputeInputs(rewriter, + *block, + compute, + instance, + klass, + schedule, + compactCpuMap, + scheduled.getWeights(), + scheduled.getInputs(), + mapper, + exchangeId))) { + signalPassFailure(); + return; + } + SmallVector yieldedValues; + cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues); + SmallVector currentKeys; + for (auto [index, value] : llvm::enumerate(yieldedValues)) { + ProducerValueKey key {instance, index}; + currentKeys.push_back(key); + } + + unsigned baseArgCount = scheduled.getWeights().size() + scheduled.getInputs().size(); + SmallVector blockYieldOperands; + bool hasNextBlock = ordinal + 1 < klass.instances.size(); + if (hasNextBlock) { + SmallVector nextCarriedKeys(carriedKeys); + llvm::append_range(nextCarriedKeys, currentKeys); + Block* nextBlock = createScheduledComputeBlock(rewriter, scheduled, nextCarriedKeys, compute.getLoc()); + appendBlockYieldBaseAndCarriedOperands(*block, baseArgCount, carriedKeys.size(), blockYieldOperands); + llvm::append_range(blockYieldOperands, yieldedValues); + rewriter.setInsertionPointToEnd(block); + createBlockYield(rewriter, compute.getLoc(), blockYieldOperands, nextBlock); + carriedKeys = std::move(nextCarriedKeys); + block = nextBlock; + } + else { + for (size_t index = 0; index < carriedKeys.size(); ++index) + blockYieldOperands.push_back(block->getArgument(baseArgCount + index)); + llvm::append_range(blockYieldOperands, yieldedValues); + rewriter.setInsertionPointToEnd(block); + createBlockYield(rewriter, compute.getLoc(), blockYieldOperands); + } + } + } + + for (SpatComputeBatch batch : batchOps) { + SpatScheduledComputeBatch scheduled = batchReplacements.lookup(batch.getOperation()); + SmallVector blockArgTypes {rewriter.getIndexType()}; + SmallVector blockArgLocs {batch.getLoc()}; + for (Value weight : scheduled.getWeights()) { + blockArgTypes.push_back(weight.getType()); + blockArgLocs.push_back(weight.getLoc()); + } + for (Value input : scheduled.getInputs()) { + blockArgTypes.push_back(input.getType()); + blockArgLocs.push_back(input.getLoc()); + } + for (Type resultType : scheduled.getResultTypes()) { + blockArgTypes.push_back(resultType); + blockArgLocs.push_back(batch.getLoc()); + } + Block* block = rewriter.createBlock( + &scheduled.getBody(), scheduled.getBody().end(), TypeRange(blockArgTypes), blockArgLocs); + GraphComputeBlockKey key = getGraphComputeBlockKey(getWholeBatchInstance(batch)); + graphComputeToBlockMap[key] = block; + auto& record = materializedSchedules[scheduledOpToRecordIndex.lookup(scheduled.getOperation())]; + record.computeKeys.push_back(key); + record.blocks.push_back(block); + rewriter.setInsertionPointToStart(block); + IRMapping mapper; + if (failed(mapBatchInputs(rewriter, + *block, + batch, + schedule, + compactCpuMap, + scheduled.getWeights(), + scheduled.getInputs(), + mapper, + exchangeId))) { + signalPassFailure(); + return; + } + for (Operation& op : batch.getBody().front()) + rewriter.clone(op, mapper); + } + + dumpPeftMaterializationReport(funcOp, schedule, compactCpuMap, materializedSchedules); + if (failed(verifyMaterializedScheduleMapping(funcOp, schedule, graphComputeToBlockMap, materializedSchedules))) { + moduleOp.emitError("scheduled Spatial materialization mapping verification failed"); + signalPassFailure(); + return; + } + if (failed(verifyScheduledSpatialInvariants(funcOp))) { + moduleOp.emitError("scheduled Spatial phase 1 verification failed"); + signalPassFailure(); + return; + } + dumpModule(moduleOp, "spatial1_scheduled_no_comm", /*assumeVerified=*/true); + moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; " + "Phase 2 communication realization is not implemented"); + signalPassFailure(); + } +}; + +} // namespace + +} // namespace spatial + +std::unique_ptr createMergeComputeNodesPass() { return std::make_unique(); } + +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp deleted file mode 100644 index 7286cec..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp +++ /dev/null @@ -1,728 +0,0 @@ -#include "SpatialDataflowCsvExporter.hpp" - -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/AsmState.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Value.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/Support/Casting.h" -#include "llvm/Support/ErrorHandling.h" -#include "llvm/Support/raw_ostream.h" - -#include -#include -#include -#include -#include - -#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" -#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" -#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" - -using namespace mlir; - -namespace onnx_mlir { -namespace spatial { - -namespace { - -struct TopLevelOpInfo { - Operation* op = nullptr; - size_t opId = 0; - bool isPost = false; - std::optional scalarCore; -}; - -struct ExpandedNodeInfo { - std::string id; - std::optional core; - std::optional lane; -}; - -struct ChannelSendRecord { - std::string sourceId; - std::optional sourceLane; -}; - -enum class LogicalNodeSelector { - Scalar, - Lane, - RangeRepresentative, -}; - -struct ResolvedProducer { - Operation* op = nullptr; - size_t resultIndex = 0; - LogicalNodeSelector selector = LogicalNodeSelector::Scalar; - uint32_t lane = 0; - uint32_t laneStart = 0; - uint32_t laneCount = 1; -}; - -struct EdgeSource { - std::string id; - std::optional sourceLane; -}; - -std::string csvEscape(StringRef field) { - bool needsQuotes = field.contains(',') || field.contains('"') || field.contains('\n') || field.contains('\r'); - if (!needsQuotes) - return field.str(); - - std::string escaped; - escaped.reserve(field.size() + 2); - escaped.push_back('"'); - for (char ch : field) { - if (ch == '"') - escaped += "\"\""; - else - escaped.push_back(ch); - } - escaped.push_back('"'); - return escaped; -} - -void writeCsvRow(std::fstream& file, ArrayRef fields) { - for (size_t i = 0; i < fields.size(); ++i) { - if (i != 0) - file << ","; - file << csvEscape(fields[i]); - } - file << "\n"; -} - -template -std::string maybeNumber(std::optional value) { - if (!value) - return ""; - return std::to_string(*value); -} - -std::string stringifyType(Type type) { - std::string storage; - llvm::raw_string_ostream os(storage); - type.print(os); - return os.str(); -} - -std::string stringifyValueAsOperand(Value value, AsmState& asmState) { - std::string storage; - llvm::raw_string_ostream os(storage); - value.printAsOperand(os, asmState); - return os.str(); -} - -std::string stringifyResultSsaNames(Operation* op, AsmState* asmState) { - if (!asmState || op->getNumResults() == 0) - return ""; - - std::string storage; - llvm::raw_string_ostream os(storage); - llvm::interleave( - op->getResults(), - [&](Value result) { os << stringifyValueAsOperand(result, *asmState); }, - [&]() { os << ";"; }); - return os.str(); -} - -std::optional getTypeSizeBytes(Type type) { - if (auto shapedType = dyn_cast(type)) { - if (!shapedType.hasStaticShape() || !hasByteSizedElementType(shapedType.getElementType())) - return std::nullopt; - return static_cast(getShapedTypeSizeInBytes(shapedType)); - } - - if (isa(type)) - return static_cast(getElementTypeSizeInBytes(type)); - if (auto intType = dyn_cast(type)) { - if (intType.getWidth() <= 0 || intType.getWidth() % 8 != 0) - return std::nullopt; - return static_cast(getElementTypeSizeInBytes(type)); - } - if (auto floatType = dyn_cast(type)) { - if (floatType.getWidth() <= 0 || floatType.getWidth() % 8 != 0) - return std::nullopt; - return static_cast(getElementTypeSizeInBytes(type)); - } - return std::nullopt; -} - -std::string getScalarId(bool isPost, size_t opId) { - return (isPost ? "sc:" : "gc:") + std::to_string(opId); -} - -std::string getBatchLaneId(bool isPost, size_t opId, uint32_t lane) { - return (isPost ? "scb:" : "gcb:") + std::to_string(opId) + ":" + std::to_string(lane); -} - -template -bool isTopLevelRelevantCompute(Operation& op) { - return isa(&op); -} - -template -FailureOr buildTopLevelOpInfo(Operation& op, bool isPost, size_t opId) { - TopLevelOpInfo info; - info.op = &op; - info.opId = opId; - info.isPost = isPost; - - if constexpr (std::is_same_v) { - if (auto compute = dyn_cast(&op)) { - auto coreId = getOptionalScheduledCoreId(compute, "spatial dataflow export core id"); - if (failed(coreId)) - return failure(); - if (*coreId) - info.scalarCore = **coreId; - } - } - - return info; -} - -template -FailureOr> getBatchLaneCoreIds(BatchOpTy batch) { - if constexpr (std::is_same_v) { - auto coreIds = getOptionalScheduledBatchCoreIds(batch, "spatial dataflow export core ids"); - if (failed(coreIds)) - return failure(); - if (!*coreIds) - return SmallVector {}; - return SmallVector((**coreIds).begin(), (**coreIds).end()); - } - return SmallVector {}; -} - -std::string getExpandedNodeId(const DenseMap, ExpandedNodeInfo>& expandedNodes, - Operation* op, - uint32_t lane) { - auto it = expandedNodes.find({op, lane}); - if (it == expandedNodes.end()) - return ""; - return it->second.id; -} - -void addScalarNodeRow(std::fstream& nodesFile, - DenseMap, ExpandedNodeInfo>& expandedNodes, - const TopLevelOpInfo& info, - AsmState* asmState = nullptr) { - std::string id = getScalarId(info.isPost, info.opId); - SmallVector row {id, std::to_string(info.opId), "", maybeNumber(info.scalarCore)}; - if (asmState) - row.push_back(stringifyResultSsaNames(info.op, asmState)); - writeCsvRow(nodesFile, row); - expandedNodes[{info.op, 0}] = {id, info.scalarCore, std::nullopt}; -} - -template -void addBatchNodeRows(std::fstream& nodesFile, - DenseMap, ExpandedNodeInfo>& expandedNodes, - const TopLevelOpInfo& info, - BatchOpTy batch, - ArrayRef> laneCoreIds, - AsmState* asmState = nullptr) { - for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { - std::string id = getBatchLaneId(info.isPost, info.opId, lane); - SmallVector row {id, - std::to_string(info.opId), - std::to_string(lane), - maybeNumber(laneCoreIds[lane])}; - if (asmState) - row.push_back(stringifyResultSsaNames(info.op, asmState)); - writeCsvRow(nodesFile, row); - expandedNodes[{info.op, lane}] = {id, laneCoreIds[lane], lane}; - } -} - -std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t lane); - -std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t lane) { - if (value == laneArg) - return static_cast(lane); - - if (std::optional constant = matchConstantIndexValue(value)) - return *constant; - - if (auto constant = value.getDefiningOp()) { - if (auto intAttr = dyn_cast(constant.getValue())) - return intAttr.getInt(); - } - - if (auto extract = value.getDefiningOp()) { - auto constant = extract.getTensor().getDefiningOp(); - auto elements = constant ? dyn_cast(constant.getValue()) : nullptr; - auto shapedType = elements ? dyn_cast(elements.getType()) : nullptr; - if (!elements || !shapedType || shapedType.getRank() != 1 || extract.getIndices().size() != 1) - return std::nullopt; - - std::optional index = evaluateIndexLike(extract.getIndices().front(), laneArg, lane); - if (!index || *index < 0 || *index >= static_cast(elements.getNumElements())) - return std::nullopt; - - if (auto denseInts = dyn_cast(elements)) - return (*(denseInts.value_begin() + *index)).getSExtValue(); - return std::nullopt; - } - - if (auto affineApply = value.getDefiningOp()) - if (FailureOr folded = evaluateAffineApply( - affineApply, - [&](Value operand) -> FailureOr { - if (std::optional resolved = evaluateIndexLike(operand, laneArg, lane)) - return *resolved; - return failure(); - }); - succeeded(folded)) { - return *folded; - } - - return std::nullopt; -} - -SmallVector collectPossibleIntValues(Value value, Value laneArg, uint32_t lane) { - if (std::optional exact = evaluateIndexLike(value, laneArg, lane)) - return {*exact}; - - auto extract = value.getDefiningOp(); - auto constant = extract ? extract.getTensor().getDefiningOp() : nullptr; - auto elements = constant ? dyn_cast(constant.getValue()) : nullptr; - if (!elements) - return {}; - - SmallVector values; - if (auto denseInts = dyn_cast(elements)) { - values.reserve(elements.getNumElements()); - for (APInt element : denseInts.getValues()) - if (!llvm::is_contained(values, element.getSExtValue())) - values.push_back(element.getSExtValue()); - } - return values; -} - -template -std::optional getBatchLaneInput(BatchOpTy batch, uint32_t lane, unsigned inputIndex) { - if (batch.getNumResults() != 0) - return batch.getInputs()[inputIndex]; - - size_t laneCount = static_cast(batch.getLaneCount()); - if (laneCount == 0 || batch.getInputs().size() % laneCount != 0) - return std::nullopt; - - size_t inputsPerLane = batch.getInputs().size() / laneCount; - size_t flatIndex = static_cast(lane) * inputsPerLane + inputIndex; - if (flatIndex >= batch.getInputs().size()) - return std::nullopt; - return batch.getInputs()[flatIndex]; -} - -template -unsigned getBatchLaneInputCount(BatchOpTy batch) { - if (batch.getNumResults() != 0) - return batch.getInputs().size(); - - size_t laneCount = static_cast(batch.getLaneCount()); - if (laneCount == 0 || batch.getInputs().size() % laneCount != 0) - return 0; - return static_cast(batch.getInputs().size() / laneCount); -} - -template -std::optional resolveProducerForValue(Value value, std::optional consumerLane) { - Operation* op = value.getDefiningOp(); - if (!op) - return std::nullopt; - - while (auto extract = dyn_cast(op)) { - Value source = extract.getSource(); - Operation* sourceOp = source.getDefiningOp(); - auto sourceBatch = dyn_cast_or_null(sourceOp); - if (sourceBatch && sourceBatch.getNumResults() != 0) { - auto staticOffsets = extract.getStaticOffsets(); - if (!staticOffsets.empty() && staticOffsets.front() != ShapedType::kDynamic) { - uint32_t lane = static_cast(staticOffsets.front()); - return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, lane, lane, 1}; - } - if (consumerLane) - return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1}; - return ResolvedProducer { - sourceOp, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast(sourceBatch.getLaneCount()) - }; - } - value = source; - op = sourceOp; - if (!op) - return std::nullopt; - } - - if (auto compute = dyn_cast(op)) - return ResolvedProducer { - compute.getOperation(), static_cast(cast(value).getResultNumber()), LogicalNodeSelector::Scalar, 0, 0, 1 - }; - - if (auto batch = dyn_cast(op)) { - if (batch.getNumResults() != 0) { - if (consumerLane) - return ResolvedProducer {op, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1}; - return ResolvedProducer { - op, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast(batch.getLaneCount()) - }; - } - - uint32_t lane = static_cast(cast(value).getResultNumber()); - return ResolvedProducer {op, static_cast(lane), LogicalNodeSelector::Lane, lane, lane, 1}; - } - - return std::nullopt; -} - -SmallVector -resolveProducerSourcesForCsv(const ResolvedProducer& producer, - const DenseMap, ExpandedNodeInfo>& expandedNodes) { - SmallVector sources; - - if (producer.selector == LogicalNodeSelector::Scalar) { - std::string id = getExpandedNodeId(expandedNodes, producer.op, 0); - if (!id.empty()) - sources.push_back({id, std::nullopt}); - return sources; - } - - if (producer.selector == LogicalNodeSelector::Lane) { - std::string id = getExpandedNodeId(expandedNodes, producer.op, producer.lane); - if (!id.empty()) - sources.push_back({id, producer.lane}); - return sources; - } - - for (uint32_t lane = producer.laneStart; lane < producer.laneStart + producer.laneCount; ++lane) { - std::string id = getExpandedNodeId(expandedNodes, producer.op, lane); - if (!id.empty()) - sources.push_back({id, lane}); - } - return sources; -} - -void emitEdgeRow(std::fstream& edgesFile, - StringRef sourceId, - StringRef targetId, - std::optional byteSize, - Type propagatedType, - StringRef stage, - std::optional sourceLane, - std::optional targetLane, - std::optional channelId) { - writeCsvRow(edgesFile, - {sourceId.str(), - targetId.str(), - maybeNumber(byteSize), - stringifyType(propagatedType), - stage.str(), - maybeNumber(sourceLane), - maybeNumber(targetLane), - maybeNumber(channelId)}); -} - -template -LogicalResult emitDataEdges(std::fstream& edgesFile, - const DenseMap& topLevelInfo, - const DenseMap, ExpandedNodeInfo>& expandedNodes, - StringRef stage) { - for (const auto& entry : topLevelInfo) { - Operation* op = entry.first; - const TopLevelOpInfo& info = entry.second; - - if (auto compute = dyn_cast(op)) { - for (Value input : compute.getInputs()) { - if (isa_and_nonnull(input.getDefiningOp())) - continue; - - auto producer = resolveProducerForValue(input, std::nullopt); - if (!producer) - continue; - - SmallVector sources = resolveProducerSourcesForCsv(*producer, expandedNodes); - std::optional byteSize = getTypeSizeBytes(input.getType()); - std::string targetId = getScalarId(info.isPost, info.opId); - for (const EdgeSource& source : sources) - emitEdgeRow(edgesFile, source.id, targetId, byteSize, input.getType(), stage, source.sourceLane, std::nullopt, std::nullopt); - } - continue; - } - - auto batch = dyn_cast(op); - if (!batch) - continue; - - unsigned inputCount = getBatchLaneInputCount(batch); - for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { - std::string targetId = getBatchLaneId(info.isPost, info.opId, lane); - for (unsigned inputIndex = 0; inputIndex < inputCount; ++inputIndex) { - std::optional input = getBatchLaneInput(batch, lane, inputIndex); - if (!input || isa_and_nonnull((*input).getDefiningOp())) - continue; - - auto producer = resolveProducerForValue(*input, lane); - if (!producer) - continue; - - SmallVector sources = resolveProducerSourcesForCsv(*producer, expandedNodes); - std::optional byteSize = getTypeSizeBytes((*input).getType()); - for (const EdgeSource& source : sources) - emitEdgeRow(edgesFile, source.id, targetId, byteSize, (*input).getType(), stage, source.sourceLane, lane, std::nullopt); - } - } - } - - return success(); -} - -template -void collectChannelSends(DenseMap>& sendsByChannelId, - const DenseMap, ExpandedNodeInfo>& expandedNodes, - BatchOpTy batch) { - std::optional laneArg = batch.getLaneArgument(); - if (!laneArg) - return; - - for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { - std::string sourceId = getExpandedNodeId(expandedNodes, batch.getOperation(), lane); - if (sourceId.empty()) - continue; - batch.getBody().walk([&](SpatChannelSendOp send) { - std::optional channelId = evaluateIndexLike(send.getChannelId(), *laneArg, lane); - if (!channelId) - return; - sendsByChannelId[*channelId].push_back({sourceId, lane}); - }); - } -} - -void collectChannelSends(DenseMap>& sendsByChannelId, - const DenseMap, ExpandedNodeInfo>& expandedNodes, - SpatScheduledCompute compute) { - std::string sourceId = getExpandedNodeId(expandedNodes, compute.getOperation(), 0); - if (sourceId.empty()) - return; - compute.getBody().walk([&](SpatChannelSendOp send) { - std::optional channelId = evaluateIndexLike(send.getChannelId(), Value(), 0); - if (!channelId) - return; - sendsByChannelId[*channelId].push_back({sourceId, std::nullopt}); - }); -} - -DenseMap> -buildNodesByCore(const DenseMap, ExpandedNodeInfo>& expandedNodes) { - DenseMap> nodesByCore; - for (const auto& entry : expandedNodes) { - const ExpandedNodeInfo& node = entry.second; - if (!node.core) - continue; - nodesByCore[*node.core].push_back({node.id, node.lane}); - } - return nodesByCore; -} - -template -LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile, - const DenseMap& topLevelInfo, - ResolveChannelSourcesFn&& resolveChannelSources, - StringRef stage) { - for (const auto& entry : topLevelInfo) { - Operation* op = entry.first; - const TopLevelOpInfo& info = entry.second; - - if (auto compute = dyn_cast(op)) { - compute.getBody().walk([&](SpatChannelReceiveOp receive) { - SmallVector sources = resolveChannelSources(receive, 0); - if (sources.empty()) - return; - std::optional channelId = evaluateIndexLike(receive.getChannelId(), Value(), 0); - std::string targetId = getScalarId(info.isPost, info.opId); - std::optional byteSize = getTypeSizeBytes(receive.getType()); - for (const ChannelSendRecord& source : sources) - emitEdgeRow(edgesFile, source.sourceId, targetId, byteSize, receive.getType(), stage, source.sourceLane, std::nullopt, channelId); - }); - continue; - } - - auto batch = dyn_cast(op); - if (!batch) - continue; - auto laneArg = batch.getLaneArgument(); - if (!laneArg) - continue; - for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { - std::string targetId = getBatchLaneId(info.isPost, info.opId, lane); - batch.getBody().walk([&](SpatChannelReceiveOp receive) { - SmallVector sources = resolveChannelSources(receive, lane); - if (sources.empty()) - return; - std::optional channelId = evaluateIndexLike(receive.getChannelId(), *laneArg, lane); - std::optional byteSize = getTypeSizeBytes(receive.getType()); - for (const ChannelSendRecord& source : sources) - emitEdgeRow(edgesFile, source.sourceId, targetId, byteSize, receive.getType(), stage, source.sourceLane, lane, channelId); - }); - } - } - - return success(); -} - -LogicalResult exportStagePre(func::FuncOp func) { - std::fstream nodesFile = openDialectDumpFileWithExtension("spatial1_graph.nodes", "/reports", "csv"); - std::fstream edgesFile = openDialectDumpFileWithExtension("spatial1_graph.edges","/reports", "csv"); - if (!nodesFile.is_open() || !edgesFile.is_open()) - return success(); - - writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core", "ssa_name"}); - writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"}); - - Operation* asmRoot = func.getOperation(); - if (auto moduleOp = func->getParentOfType()) - asmRoot = moduleOp.getOperation(); - OpPrintingFlags flags; - flags.elideLargeElementsAttrs().enableDebugInfo(true, false); - AsmState asmState(asmRoot, flags); - - DenseMap topLevelInfo; - DenseMap, ExpandedNodeInfo> expandedNodes; - - size_t opId = 0; - for (Operation& op : func.getBody().front()) { - if (!isTopLevelRelevantCompute(op)) - continue; - FailureOr info = buildTopLevelOpInfo(op, false, opId++); - if (failed(info)) - return failure(); - topLevelInfo[&op] = *info; - - if (auto compute = dyn_cast(&op)) { - addScalarNodeRow(nodesFile, expandedNodes, *info, &asmState); - continue; - } - - auto batch = cast(&op); - SmallVector, 8> laneCoreIds(batch.getLaneCount()); - addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState); - } - - return emitDataEdges(edgesFile, topLevelInfo, expandedNodes, "pre"); -} - -LogicalResult exportStagePost(func::FuncOp func) { - std::fstream nodesFile = openDialectDumpFileWithExtension("spatial2_merged.nodes", "/reports", "csv"); - std::fstream edgesFile = openDialectDumpFileWithExtension("spatial2_merged.edges", "/reports", "csv"); - if (!nodesFile.is_open() || !edgesFile.is_open()) - return success(); - - writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core"}); - writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"}); - - DenseMap topLevelInfo; - DenseMap, ExpandedNodeInfo> expandedNodes; - - size_t opId = 0; - for (Operation& op : func.getBody().front()) { - if (!isTopLevelRelevantCompute(op)) - continue; - FailureOr info = buildTopLevelOpInfo(op, true, opId++); - if (failed(info)) - return failure(); - topLevelInfo[&op] = *info; - - if (isa(&op)) { - addScalarNodeRow(nodesFile, expandedNodes, *info); - continue; - } - - auto batch = cast(&op); - auto coreIds = getBatchLaneCoreIds(batch); - if (failed(coreIds)) - return failure(); - SmallVector, 8> laneCoreIds(batch.getLaneCount()); - for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) - if (lane < coreIds->size()) - laneCoreIds[lane] = (*coreIds)[lane]; - addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds); - } - - if (failed(emitDataEdges(edgesFile, topLevelInfo, expandedNodes, "post"))) - return failure(); - - DenseMap> sendsByChannelId; - for (const auto& entry : topLevelInfo) { - Operation* op = entry.first; - if (auto compute = dyn_cast(op)) - collectChannelSends(sendsByChannelId, expandedNodes, compute); - else if (auto batch = dyn_cast(op)) - collectChannelSends(sendsByChannelId, expandedNodes, batch); - } - - DenseMap> nodesByCore = buildNodesByCore(expandedNodes); - auto resolveChannelSources = [&](SpatChannelReceiveOp receive, uint32_t lane) { - SmallVector sources; - - Value laneArg; - if (auto owner = receive->getParentOfType()) - if (auto maybeLaneArg = owner.getLaneArgument()) - laneArg = *maybeLaneArg; - - if (std::optional channelId = evaluateIndexLike(receive.getChannelId(), laneArg, lane)) { - if (auto it = sendsByChannelId.find(*channelId); it != sendsByChannelId.end()) - return it->second; - } - - for (int64_t sourceCore : collectPossibleIntValues(receive.getSourceCoreId(), laneArg, lane)) { - auto it = nodesByCore.find(static_cast(sourceCore)); - if (it == nodesByCore.end()) - continue; - llvm::append_range(sources, it->second); - } - return sources; - }; - - return emitExplicitChannelEdges( - edgesFile, topLevelInfo, resolveChannelSources, "post"); -} - -} // namespace - -SpatialDataflowExportStage getSpatialDataflowExportStage() { - switch (pimExportSpatialDataflow.getValue()) { - case SpatialDataflowExportNone: return SpatialDataflowExportStage::None; - case SpatialDataflowExportPre: return SpatialDataflowExportStage::Pre; - case SpatialDataflowExportPost: return SpatialDataflowExportStage::Post; - case SpatialDataflowExportBoth: return SpatialDataflowExportStage::Both; - } - llvm_unreachable("unknown spatial dataflow export mode"); -} - -bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage) { - switch (mode) { - case SpatialDataflowExportStage::None: return false; - case SpatialDataflowExportStage::Pre: return stage == SpatialDataflowExportStage::Pre; - case SpatialDataflowExportStage::Post: return stage == SpatialDataflowExportStage::Post; - case SpatialDataflowExportStage::Both: - return stage == SpatialDataflowExportStage::Pre || stage == SpatialDataflowExportStage::Post; - } - return false; -} - -LogicalResult exportSpatialDataflowCsvPre(func::FuncOp func) { return exportStagePre(func); } - -LogicalResult exportSpatialDataflowCsvPost(func::FuncOp func) { return exportStagePost(func); } - -} // namespace spatial -} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp deleted file mode 100644 index aef6efc..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Support/LogicalResult.h" - - -namespace onnx_mlir { -namespace spatial { - -enum class SpatialDataflowExportStage { - None, - Pre, - Post, - Both, -}; - -SpatialDataflowExportStage getSpatialDataflowExportStage(); - -mlir::LogicalResult exportSpatialDataflowCsvPre(mlir::func::FuncOp func); -mlir::LogicalResult exportSpatialDataflowCsvPost(mlir::func::FuncOp func); - -bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage); - -} // namespace spatial -} // namespace onnx_mlir