diff --git a/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp b/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp index 037f0a3..f30f73a 100644 --- a/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp @@ -141,7 +141,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe std::optional mode = blueprint.getMode(); std::optional> operandIndicesAttr = blueprint.getFragmentOperandIndices(); std::optional> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets(); - if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr) + std::optional> sourceSlotsAttr = blueprint.getFragmentSourceSlots(); + if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr) return failure(); if (!blueprint.getOutput().hasOneUse() || !isa(*blueprint.getOutput().getUsers().begin())) return failure(); @@ -153,6 +154,9 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe ArrayRef operandIndices = *operandIndicesAttr; ArrayRef sourceOffsets = *sourceOffsetsAttr; + ArrayRef sourceSlots = *sourceSlotsAttr; + if (sourceSlots.size() != operandIndices.size()) + return failure(); ArrayRef flatOffsets = blueprint.getFragmentOffsets(); ArrayRef flatSizes = blueprint.getFragmentSizes(); ArrayRef flatStrides = *stridesAttr; @@ -174,7 +178,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe if (operandIndices[fragmentIndex] != static_cast(use.getOperandNumber())) continue; - int64_t sourceElementOffset = sourceOffsets[fragmentIndex]; + int64_t sourceElementOffset = + sourceSlots[fragmentIndex] * payloadElementCount + sourceOffsets[fragmentIndex]; int64_t lane = sourceElementOffset / payloadElementCount; if (lane < 0 || lane >= static_cast(laneCount)) return failure(); diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp b/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp index 995b5a7..2df4402 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp +++ b/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp @@ -204,15 +204,32 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO if (!targetType || !sourceType || size <= 0) return failure(); - auto logicalCopyShape = inferLogicalCopyShape(targetType, sourceType, size); - if (failed(logicalCopyShape)) - return failure(); - auto targetPlan = analyzeCopyEndpoint(target, targetOffset, targetType); auto sourcePlan = analyzeCopyEndpoint(source, sourceOffset, sourceType); if (failed(targetPlan) || failed(sourcePlan)) return failure(); + auto targetBytes = getShapedByteSize(targetType); + auto sourceBytes = getShapedByteSize(sourceType); + if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes) + && *targetBytes == size && *sourceBytes == size) { + auto targetSuffixRank = getContiguousSuffixRank(targetType, targetType.getShape()); + auto sourceSuffixRank = getContiguousSuffixRank(sourceType, sourceType.getShape()); + if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank) + && *targetSuffixRank == targetType.getRank() && *sourceSuffixRank == sourceType.getRank()) { + CopyRewritePlan plan; + plan.kind = CopyRewritePlan::Kind::Direct; + plan.target = *targetPlan; + plan.source = *sourcePlan; + plan.directBytes = size; + return plan; + } + } + + auto logicalCopyShape = inferLogicalCopyShape(targetType, sourceType, size); + if (failed(logicalCopyShape)) + return failure(); + auto targetSuffixRank = getContiguousSuffixRank(targetType, *logicalCopyShape); auto sourceSuffixRank = getContiguousSuffixRank(sourceType, *logicalCopyShape); if (failed(targetSuffixRank) || failed(sourceSuffixRank)) diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index 1dd4f88..33f02a9 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -10,6 +10,7 @@ add_pim_library(SpatialOps Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp + Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp index d3f726c..6b08094 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp @@ -2,6 +2,7 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/IRMapping.h" @@ -24,6 +25,93 @@ static FailureOr getOriginalProducerValue(const ProducerValueRef &produce return outputs[producer.resultIndex]; } +static SmallVector getBlueprintFragments(SpatBlueprintOp blueprint) { + SmallVector fragments {blueprint.getInput()}; + llvm::append_range(fragments, blueprint.getFragments()); + return fragments; +} + +static FailureOr buildBlueprintReconstruction( + OpBuilder &builder, Location loc, SpatBlueprintOp blueprint, + ValueRange sourceBlockArgs) { + auto resultType = dyn_cast(blueprint.getOutput().getType()); + auto operandIndices = blueprint.getFragmentOperandIndices(); + auto sourceSlots = blueprint.getFragmentSourceSlots(); + auto sourceOffsets = blueprint.getFragmentSourceOffsets(); + auto strides = blueprint.getFragmentStrides(); + if (!resultType || !resultType.hasStaticShape() || !operandIndices || + !sourceSlots || !sourceOffsets || !strides) + return blueprint.emitOpError("phase 1 requires complete static fragment assembly metadata"), failure(); + int64_t rank = resultType.getRank(); + ArrayRef offsets = blueprint.getFragmentOffsets(); + ArrayRef sizes = blueprint.getFragmentSizes(); + if (offsets.size() != sizes.size() || offsets.size() != strides->size() || + offsets.size() != operandIndices->size() * rank || + sourceSlots->size() != operandIndices->size() || + sourceOffsets->size() != operandIndices->size()) + return blueprint.emitOpError("phase 1 fragment assembly metadata has inconsistent sizes"), failure(); + + Value result = tensor::EmptyOp::create(builder, loc, resultType.getShape(), + resultType.getElementType()); + for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) { + if (operandIndex < 0 || operandIndex >= static_cast(sourceBlockArgs.size())) + return blueprint.emitOpError("phase 1 fragment assembly operand index is out of range"), failure(); + auto physicalType = dyn_cast(sourceBlockArgs[operandIndex].getType()); + if (!physicalType || !physicalType.hasStaticShape() || physicalType.getRank() != rank + 1) + return blueprint.emitOpError("phase 1 fragment assembly source is not a physical fragment batch"), failure(); + SmallVector fragmentShape(physicalType.getShape().drop_front()); + int64_t linearOffset = (*sourceOffsets)[fragmentIndex]; + SmallVector sourceCoordinates(rank); + for (int64_t dim = rank - 1; dim >= 0; --dim) { + sourceCoordinates[dim] = linearOffset % fragmentShape[dim]; + linearOffset /= fragmentShape[dim]; + } + if (linearOffset != 0) + return blueprint.emitOpError("phase 1 fragment source offset is out of range"), failure(); + + SmallVector sliceOffsets, sliceSizes, sliceStrides; + sliceOffsets.push_back(builder.getIndexAttr((*sourceSlots)[fragmentIndex])); + sliceSizes.push_back(builder.getIndexAttr(1)); + sliceStrides.push_back(builder.getIndexAttr(1)); + SmallVector selectedShape {1}; + for (int64_t dim = 0; dim < rank; ++dim) { + int64_t index = fragmentIndex * rank + dim; + int64_t size = sizes[index]; + if ((*strides)[index] != 1 || sourceCoordinates[dim] < 0 || size <= 0 || + sourceCoordinates[dim] + size > fragmentShape[dim]) + return blueprint.emitOpError("phase 1 fragment geometry is unsupported"), failure(); + sliceOffsets.push_back(builder.getIndexAttr(sourceCoordinates[dim])); + sliceSizes.push_back(builder.getIndexAttr(size)); + sliceStrides.push_back(builder.getIndexAttr(1)); + selectedShape.push_back(size); + } + auto selectedType = RankedTensorType::get(selectedShape, resultType.getElementType()); + Value selected = tensor::ExtractSliceOp::create( + builder, loc, selectedType, sourceBlockArgs[operandIndex], sliceOffsets, + sliceSizes, sliceStrides); + SmallVector fragmentResultShape(selectedShape.begin() + 1, + selectedShape.end()); + auto fragmentType = RankedTensorType::get(fragmentResultShape, + resultType.getElementType()); + SmallVector reassociation {{0, 1}}; + for (int64_t dim = 1; dim < rank; ++dim) + reassociation.push_back({dim + 1}); + Value fragment = tensor::CollapseShapeOp::create( + builder, loc, fragmentType, selected, reassociation); + SmallVector targetOffsets, targetSizes, targetStrides; + for (int64_t dim = 0; dim < rank; ++dim) { + int64_t index = fragmentIndex * rank + dim; + targetOffsets.push_back(builder.getIndexAttr(offsets[index])); + targetSizes.push_back(builder.getIndexAttr(sizes[index])); + targetStrides.push_back(builder.getIndexAttr((*strides)[index])); + } + result = tensor::InsertSliceOp::create(builder, loc, fragment, result, + targetOffsets, targetSizes, + targetStrides); + } + return result; +} + static bool isSupportedDeferredShapingOp(Operation *op) { return isa clonePayloadRoot(Value root, Block &body, const DeferredInputPlan &plan, OpBuilder &builder, SpatDeferredCommunicationOp transfer, - Value selectedSource) { + Value selectedSource, Value boundGraphLane) { IRMapping mapping; mapping.map(plan.graphInput, selectedSource); std::function(Value)> cloneScheduledLane = [&](Value value) -> FailureOr { @@ -118,7 +206,7 @@ static FailureOr clonePayloadRoot(Value root, Block &body, const Deferred std::function(Value)> clone = [&](Value value) -> FailureOr { if (mapping.contains(value)) return mapping.lookup(value); if (value == plan.graphLane) { - auto mappedLane = cloneScheduledLane(plan.scheduledGraphLane); + auto mappedLane = cloneScheduledLane(boundGraphLane ? boundGraphLane : plan.scheduledGraphLane); if (failed(mappedLane)) return failure(); mapping.map(value, *mappedLane); return *mappedLane; @@ -136,6 +224,62 @@ static FailureOr clonePayloadRoot(Value root, Block &body, const Deferred return clone(root); } +static bool dependsOnGraphLane(Value value, Value graphLane, Block &body, + llvm::SmallPtrSetImpl &seen) { + if (value == graphLane) + return true; + Operation *op = value.getDefiningOp(); + if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second) + return false; + return llvm::any_of(op->getOperands(), [&](Value operand) { + return dependsOnGraphLane(operand, graphLane, body, seen); + }); +} + +static FailureOr buildPayloadAggregate(OpBuilder &builder, Location loc, + ArrayRef payloads) { + auto payloadType = dyn_cast(payloads.front().getType()); + if (!payloadType || !payloadType.hasStaticShape()) + return failure(); + SmallVector shape {static_cast(payloads.size())}; + llvm::append_range(shape, payloadType.getShape()); + auto aggregateType = RankedTensorType::get(shape, payloadType.getElementType()); + auto empty = createEmptyTensorForType(builder, loc, aggregateType); + if (failed(empty)) return failure(); + Value aggregate = *empty; + SmallVector sizes, strides(shape.size(), builder.getIndexAttr(1)); + sizes.push_back(builder.getIndexAttr(1)); + for (int64_t dim : payloadType.getShape()) sizes.push_back(builder.getIndexAttr(dim)); + SmallVector reassociation {{0, 1}}; + for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1}); + SmallVector expandedShape {1}; llvm::append_range(expandedShape, payloadType.getShape()); + auto expandedType = RankedTensorType::get(expandedShape, payloadType.getElementType()); + for (auto [index, payload] : llvm::enumerate(payloads)) { + Value expanded = tensor::ExpandShapeOp::create(builder, loc, expandedType, payload, reassociation).getResult(); + SmallVector offsets(shape.size(), builder.getIndexAttr(0)); + offsets[0] = builder.getIndexAttr(index); + aggregate = tensor::InsertSliceOp::create(builder, loc, expanded, aggregate, offsets, sizes, strides).getResult(); + } + return aggregate; +} + +static FailureOr selectPayloadAggregate(OpBuilder &builder, Location loc, Value aggregate, + Value localLane) { + auto aggregateType = cast(aggregate.getType()); + SmallVector payloadShape(aggregateType.getShape().begin() + 1, aggregateType.getShape().end()); + auto payloadType = RankedTensorType::get(payloadShape, aggregateType.getElementType()); + SmallVector offsets(aggregateType.getRank(), builder.getIndexAttr(0)); offsets[0] = localLane; + SmallVector sizes, strides(aggregateType.getRank(), builder.getIndexAttr(1)); + sizes.push_back(builder.getIndexAttr(1)); + for (int64_t dim : payloadShape) sizes.push_back(builder.getIndexAttr(dim)); + SmallVector unitShape {1}; llvm::append_range(unitShape, payloadShape); + Value unit = tensor::ExtractSliceOp::create(builder, loc, + RankedTensorType::get(unitShape, aggregateType.getElementType()), aggregate, offsets, sizes, strides).getResult(); + SmallVector reassociation {{0, 1}}; + for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1}); + return tensor::CollapseShapeOp::create(builder, loc, payloadType, unit, reassociation).getResult(); +} + static void collectClosure(Value value, Block &body, const DeferredInputPlan &plan, llvm::SmallPtrSetImpl &ops) { Operation *op = value.getDefiningOp(); @@ -146,12 +290,27 @@ static void collectClosure(Value value, Block &body, const DeferredInputPlan &pl } // namespace +bool isDeferredFragmentAssemblyInput( + Value input, const ComputeInstance &consumerInstance) { + auto blueprint = input.getDefiningOp(); + if (!blueprint || blueprint.getMode() != "fragment_assembly") + return false; + return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) { + return getProducerValueRef(fragment, &consumerInstance).has_value(); + }); +} + LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput, const ComputeInstance &consumerInstance, const MergeScheduleResult &, ValueRange scheduledInputs, Block &block, unsigned firstInputArgument, ArrayRef carriedKeys, Value graphLane, Value scheduledGraphLane, DeferredInputPlan &plan) { - plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}}; + plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}, {}, {}, {}, 1, nullptr}; + if (isDeferredFragmentAssemblyInput(input, consumerInstance)) { + plan.blueprint = input.getDefiningOp(); + plan.originalSources = getBlueprintFragments(plan.blueprint); + return success(); + } auto producer = getProducerValueRef(input, &consumerInstance); if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); } ProducerValueKey key {producer->instance, producer->resultIndex}; @@ -171,8 +330,13 @@ LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input, const MergeScheduleResult &, ValueRange scheduledInputs, Block &block, unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane, Value scheduledLane, DeferredInputPlan &plan) { - plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane}; const ComputeInstance &representative = tuple.instances.front(); + plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane, {}, {}, {}, 1, nullptr}; + if (isDeferredFragmentAssemblyInput(input, representative)) { + plan.blueprint = input.getDefiningOp(); + plan.originalSources = getBlueprintFragments(plan.blueprint); + return success(); + } auto producer = getProducerValueRef(input, &representative); if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); } auto inputs = getComputeInstanceInputs(representative); @@ -222,19 +386,67 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc llvm::sort(roots, [](Value a, Value b) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); }); roots.erase(std::unique(roots.begin(), roots.end()), roots.end()); for (Value root : roots) { + llvm::SmallPtrSet laneDependencies; + bool scalarize = plan.scalarizedGraphLaneBase + && dependsOnGraphLane(root, plan.graphLane, body, laneDependencies); + OpBuilder::InsertPoint restore = builder.saveInsertionPoint(); + Operation *loop = nullptr; + if (scalarize) { + loop = builder.getInsertionBlock()->getParentOp(); + if (loop && !isa(loop)) + loop = loop->getParentOfType(); + if (loop) + builder.setInsertionPoint(loop); + else if (plan.scalarizedHoistBlock) + builder.setInsertionPointToEnd(plan.scalarizedHoistBlock); + else + return emitError(loc) << "phase 1 scalarized deferred payload is missing a hoist point"; + } + SmallVector payloads; + unsigned count = scalarize ? plan.scalarizedLaneCount : 1; + for (unsigned offset = 0; offset < count; ++offset) { auto transfer = SpatDeferredCommunicationOp::create(builder, loc, root.getType(), plan.originalSources); Block *deferred = builder.createBlock(&transfer.getBody(), transfer.getBody().end(), TypeRange {transfer.getSources().getTypes()}, SmallVector(transfer.getSources().size(), loc)); builder.setInsertionPointToStart(deferred); - auto selected = buildSelectedDeferredSource(builder, loc, transfer, plan.scheduledLane, - deferred->getArguments(), plan.sourceOperandForScheduledLane); + auto selected = plan.blueprint + ? buildBlueprintReconstruction(builder, loc, plan.blueprint, + deferred->getArguments()) + : buildSelectedDeferredSource(builder, loc, transfer, + plan.scheduledLane, + deferred->getArguments(), + plan.sourceOperandForScheduledLane); if (failed(selected)) return failure(); - auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected); + Value boundGraphLane; + if (scalarize) { + Value offsetValue = arith::ConstantIndexOp::create(builder, loc, offset); + boundGraphLane = offset ? arith::AddIOp::create(builder, loc, plan.scalarizedGraphLaneBase, offsetValue).getResult() + : plan.scalarizedGraphLaneBase; + } + auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected, boundGraphLane); if (failed(payload)) return failure(); SpatYieldOp::create(builder, loc, *payload); - mapper.map(root, transfer.getOutput()); - collectClosure(root, body, plan, absorbed); + payloads.push_back(transfer.getOutput()); builder.setInsertionPointAfter(transfer); + } + if (scalarize) { + builder.restoreInsertionPoint(restore); + if (payloads.size() == 1) { + mapper.map(root, payloads.front()); + } else { + if (loop) builder.setInsertionPoint(loop); + else builder.setInsertionPointToEnd(plan.scalarizedHoistBlock); + auto aggregate = buildPayloadAggregate(builder, loc, payloads); + if (failed(aggregate)) return failure(); + builder.restoreInsertionPoint(restore); + auto selected = selectPayloadAggregate(builder, loc, *aggregate, plan.scalarizedLocalLane); + if (failed(selected)) return failure(); + mapper.map(root, *selected); + } + } else { + mapper.map(root, payloads.front()); + } + collectClosure(root, body, plan, absorbed); } } for (Operation *op : absorbed) { diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp index 4165b9d..0e42563 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp @@ -15,8 +15,16 @@ struct DeferredInputPlan { Value graphLane; Value scheduledGraphLane; Value scheduledLane; + SpatBlueprintOp blueprint; + Value scalarizedLocalLane; + Value scalarizedGraphLaneBase; + int64_t scalarizedLaneCount = 1; + Block *scalarizedHoistBlock = nullptr; }; +bool isDeferredFragmentAssemblyInput(Value input, + const ComputeInstance &consumerInstance); + LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input, BlockArgument graphInput, const ComputeInstance &consumerInstance, diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp index 6db9d78..b65ddc3 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp @@ -13,6 +13,7 @@ #include #include "DeferredCommunicationDeadlock.hpp" +#include "DeferredProjectionAnalysis.hpp" #include "DeferredCommunicationRealization.hpp" #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" @@ -27,18 +28,6 @@ namespace { struct ScheduledInfo; -struct DeferredBodyAnalysis { - struct SourceUse { - unsigned sourceIndex = 0; - tensor::ExtractSliceOp leadingSlice; - std::optional staticSourceLane; - }; - - Value yieldedValue; - SmallVector sourceSideOps; - SmallVector sourceUses; -}; - struct ProducedValue { ScheduledInfo* scheduled = nullptr; unsigned step = 0; @@ -47,6 +36,7 @@ struct ProducedValue { int64_t core = -1; int64_t laneStart = 0; int64_t laneCount = 1; + unsigned scheduledLane = 0; Value payload; Value published; }; @@ -65,46 +55,52 @@ struct ScheduledInfo { bool isBatch() const { return isa(op); } }; -struct LaneTransfer { - ProducedValue* source = nullptr; +struct FragmentRequirement { + unsigned leafIndex = 0; + unsigned selectedPosition = 0; + unsigned sourceOperandIndex = 0; + std::optional physicalSlot; + std::optional graphLane; + ProducedValue *producer = nullptr; + Type publicationFragmentType; + std::optional targetScheduledLane; +}; + +struct FragmentTransfer { + unsigned requirementIndex = 0; unsigned sourceStream = 0; unsigned targetStream = 0; int64_t sourceCore = -1; int64_t targetCore = -1; - unsigned targetLane = 0; - unsigned sourceIndex = 0; uint64_t channelId = 0; + unsigned globalOrder = 0; unsigned sourceInsertionStep = 0; unsigned targetInsertionStep = 0; - Type transportType; - Operation* bodyRoot = nullptr; bool local = false; }; -struct NormalizedScalarExchange { - uint64_t id = 0; - SpatDeferredCommunicationOp deferred; - ScheduledInfo* target = nullptr; - unsigned consumerStep = 0; - DeferredBodyAnalysis body; - LaneTransfer transfer; +struct DeferredSpecialization { + std::optional targetScheduledLane; + SpecializedDeferredProgram program; + SmallVector requirements; + SmallVector transfers; }; -struct NormalizedBatchedExchange { - uint64_t id = 0; +struct NormalizedDeferredExchange { + uint64_t exchangeId = 0; SpatDeferredCommunicationOp deferred; - ScheduledInfo* target = nullptr; + ScheduledInfo *target = nullptr; unsigned consumerStep = 0; - DeferredBodyAnalysis body; - SmallVector lanes; + SmallVector specializations; + unsigned externalTransferCount = 0; }; struct RealizationPlan { SmallVector scheduled; SmallVector> producedStorage; - DenseMap> producedByGraph; - SmallVector scalar; - SmallVector batched; + DenseMap> producedByGraph; + GraphBatchPublicationCache publicationCache; + SmallVector exchanges; SmallVector external; SmallVector stepCounts; }; @@ -260,6 +256,7 @@ static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan info.cores.front(), laneStarts[step], std::max(laneCounts[step], 1), + 0, *payload, info.op->getResult(globalResult), }); @@ -281,6 +278,7 @@ static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan info.cores[lane], laneStarts[laneIndex], std::max(laneCounts[laneIndex], 1), + lane, *payload, info.op->getResult(globalResult), }); @@ -294,213 +292,8 @@ static LogicalResult collectScheduled(func::FuncOp funcOp, RealizationPlan& plan return success(); } -static bool isSupportedDeferredBodyOp(Operation* op) { - return isa(op); -} - -static LogicalResult collectDeferredDependencies(SpatDeferredCommunicationOp deferred, - Value value, - llvm::SmallDenseSet& sourceIndices, - llvm::SmallPtrSetImpl& requiredOps, - Value allowedExternalLane = {}) { - Block& body = deferred.getBody().front(); - if (auto argument = dyn_cast(value)) { - if (argument.getOwner() != &body && value == allowedExternalLane) - return success(); - if (argument.getArgNumber() >= deferred.getSources().size()) - return deferred.emitOpError() << "deferred body yielded value depends on unknown block argument " - << argument.getArgNumber(); - sourceIndices.insert(argument.getArgNumber()); - return success(); - } - Operation* definingOp = value.getDefiningOp(); - if (!definingOp || definingOp->getBlock() != &body) { - if (value == allowedExternalLane) - return success(); - if (definingOp && definingOp->hasTrait()) - return success(); - return deferred.emitOpError("deferred body payload derivation captures a value unavailable in the source stream"); - } - if (!isSupportedDeferredBodyOp(definingOp)) - return deferred.emitOpError() << "deferred body payload derivation contains unsupported op '" - << definingOp->getName() << "'"; - if (!requiredOps.insert(definingOp).second) - return success(); - for (Value operand : definingOp->getOperands()) - if (failed(collectDeferredDependencies(deferred, operand, sourceIndices, requiredOps, allowedExternalLane))) - return failure(); - return success(); -} - -static bool dependsOn(Value value, Value needle, llvm::SmallPtrSetImpl& visited) { - if (value == needle) - return true; - Operation* op = value.getDefiningOp(); - if (!op || !visited.insert(op).second) - return false; - return llvm::any_of(op->getOperands(), [&](Value operand) { return dependsOn(operand, needle, visited); }); -} - -static std::optional evaluateConstantInteger(Value value, Value lane = {}, std::optional laneValue = std::nullopt) { - if (!value) - return std::nullopt; - if (value == lane && laneValue) - return laneValue; - if (auto constant = value.getDefiningOp()) - if (auto integer = dyn_cast(constant.getValue())) - return integer.getInt(); - if (auto cast = value.getDefiningOp()) - return evaluateConstantInteger(cast.getIn(), lane, laneValue); - if (auto add = value.getDefiningOp()) { - auto lhs = evaluateConstantInteger(add.getLhs(), lane, laneValue); - auto rhs = evaluateConstantInteger(add.getRhs(), lane, laneValue); - if (lhs && rhs) - return *lhs + *rhs; - } - if (auto sub = value.getDefiningOp()) { - auto lhs = evaluateConstantInteger(sub.getLhs(), lane, laneValue); - auto rhs = evaluateConstantInteger(sub.getRhs(), lane, laneValue); - if (lhs && rhs) - return *lhs - *rhs; - } - if (auto mul = value.getDefiningOp()) { - auto lhs = evaluateConstantInteger(mul.getLhs(), lane, laneValue); - auto rhs = evaluateConstantInteger(mul.getRhs(), lane, laneValue); - if (lhs && rhs) - return *lhs * *rhs; - } - if (auto apply = value.getDefiningOp()) { - auto evaluated = evaluateAffineApply(apply, [lane, laneValue](Value operand) -> FailureOr { - auto constant = evaluateConstantInteger(operand, lane, laneValue); - return constant ? FailureOr(*constant) : FailureOr(failure()); - }); - if (succeeded(evaluated)) - return *evaluated; - } - if (auto extract = value.getDefiningOp()) { - auto constant = extract.getTensor().getDefiningOp(); - auto elements = constant ? dyn_cast(constant.getValue()) : DenseIntElementsAttr(); - auto type = elements ? dyn_cast(elements.getType()) : RankedTensorType(); - if (!elements || !type || extract.getIndices().size() != static_cast(type.getRank())) - return std::nullopt; - int64_t linearIndex = 0; - for (auto [index, dimension] : llvm::zip(extract.getIndices(), type.getShape())) { - auto constantIndex = evaluateConstantInteger(index, lane, laneValue); - if (!constantIndex || *constantIndex < 0 || *constantIndex >= dimension) - return std::nullopt; - linearIndex = linearIndex * dimension + *constantIndex; - } - auto values = elements.getValues(); - return values[linearIndex].getSExtValue(); - } - return std::nullopt; -} - -static FailureOr analyzeDeferredBody(SpatDeferredCommunicationOp deferred, - std::optional targetLane = std::nullopt) { - Block& body = deferred.getBody().front(); - auto yield = dyn_cast(body.getTerminator()); - if (!yield || yield.getOutputs().size() != 1) - return deferred.emitOpError("phase 2 requires a single deferred yielded value"); - - DeferredBodyAnalysis analysis; - analysis.yieldedValue = yield.getOutputs().front(); - llvm::SmallDenseSet sourceIndices; - llvm::SmallPtrSet requiredOps; - Value scheduledLane; - if (auto scheduled = deferred->getParentOfType()) - scheduledLane = scheduled.getLaneArgument().value(); - if (failed(collectDeferredDependencies(deferred, analysis.yieldedValue, sourceIndices, requiredOps, scheduledLane))) - return failure(); - for (Operation& op : body.without_terminator()) - if (requiredOps.contains(&op)) - analysis.sourceSideOps.push_back(&op); - - for (unsigned sourceIndex : sourceIndices) { - DeferredBodyAnalysis::SourceUse sourceUse; - sourceUse.sourceIndex = sourceIndex; - Value sourceArgument = body.getArgument(sourceIndex); - for (Operation* op : analysis.sourceSideOps) { - auto slice = dyn_cast(op); - if (!slice) - continue; - llvm::SmallPtrSet visited; - if (!dependsOn(slice.getSource(), sourceArgument, visited)) - continue; - auto sourceType = dyn_cast(slice.getSourceType()); - if (!sourceType || sourceType.getRank() == 0 || slice.getMixedSizes().empty()) - continue; - auto leadingSize = getConstantIntValue(slice.getMixedSizes().front()); - if (!leadingSize || *leadingSize != 1) - continue; - sourceUse.leadingSlice = slice; - if (auto offset = dyn_cast(slice.getMixedOffsets().front())) - sourceUse.staticSourceLane = cast(offset).getInt(); - else { - Value dynamicOffset = cast(slice.getMixedOffsets().front()); - sourceUse.staticSourceLane = evaluateConstantInteger(dynamicOffset, scheduledLane, - targetLane ? std::optional(*targetLane) : std::nullopt); - } - break; - } - analysis.sourceUses.push_back(sourceUse); - } - return analysis; -} - -static DeferredBodyAnalysis::SourceUse* findSourceUse(DeferredBodyAnalysis& analysis, unsigned sourceIndex) { - auto it = llvm::find_if(analysis.sourceUses, [&](const auto& use) { return use.sourceIndex == sourceIndex; }); - return it == analysis.sourceUses.end() ? nullptr : &*it; -} - -static bool isIdentityDeferredBody(SpatDeferredCommunicationOp deferred, - const DeferredBodyAnalysis& analysis, - unsigned sourceIndex) { - Block& block = deferred.getBody().front(); - return analysis.sourceSideOps.empty() && analysis.yieldedValue == block.getArgument(sourceIndex); -} - -static LogicalResult configurePayloadAndReplacement(SpatDeferredCommunicationOp deferred, - DeferredBodyAnalysis& analysis, - ProducedValue& source, - unsigned sourceIndex, - int64_t sourceLane, - LaneTransfer& lane) { - lane.sourceIndex = sourceIndex; - lane.transportType = analysis.yieldedValue.getType(); - if (isIdentityDeferredBody(deferred, analysis, sourceIndex) && source.payload.getType() == analysis.yieldedValue.getType()) - return success(); - - // A matching source argument lets Phase 2 clone only the body-owned shaping. - if (deferred.getBody().front().getArgument(sourceIndex).getType() == source.payload.getType()) - return success(); - - if (auto* sourceUse = findSourceUse(analysis, sourceIndex); - sourceUse && sourceUse->leadingSlice - && sourceUse->leadingSlice.getResult().getType() == source.payload.getType()) { - if (sourceUse->staticSourceLane && *sourceUse->staticSourceLane != sourceLane) - return deferred.emitOpError("deferred-body source lane disagrees with scheduled source metadata"); - lane.bodyRoot = sourceUse->leadingSlice; - return success(); - } - return deferred.emitOpError( - "phase 1 must place all graph-result shaping inside the deferred_communication body"); -} - static FailureOr -findProducer(RealizationPlan& plan, SpatDeferredCommunicationOp deferred, int64_t graphId, unsigned resultIndex, +findProducer(RealizationPlan& plan, Operation *diagnosticOwner, int64_t graphId, unsigned resultIndex, std::optional graphLane) { ProducedValue* match = nullptr; for (ProducedValue* produced : plan.producedByGraph.lookup(graphId)) { @@ -509,52 +302,153 @@ findProducer(RealizationPlan& plan, SpatDeferredCommunicationOp deferred, int64_ if (graphLane && (*graphLane < produced->laneStart || *graphLane >= produced->laneStart + produced->laneCount)) continue; if (match) - return deferred.emitOpError("phase 2 cannot uniquely resolve graph publication ownership"), failure(); + return diagnosticOwner->emitOpError("phase 2 cannot uniquely resolve graph publication ownership"), failure(); match = produced; } if (!match) - return deferred.emitOpError("phase 2 cannot map a deferred source to a scheduled producer"), failure(); + return diagnosticOwner->emitOpError("phase 2 cannot map a graph lane to a scheduled producer"), failure(); return match; } -static LogicalResult buildLaneTransfer(RealizationPlan& plan, - SpatDeferredCommunicationOp deferred, - DeferredBodyAnalysis& body, - ScheduledInfo& target, - unsigned consumerStep, - unsigned targetLane, - unsigned sourceIndex, - LaneTransfer& lane) { - Value sourceValue = deferred.getSources()[sourceIndex]; - auto sourceResult = dyn_cast(sourceValue); - if (!sourceResult) - return deferred.emitOpError("phase 2 requires deferred sources to be original graph results"); - Operation* graphProducer = sourceResult.getOwner(); - if (!isa_and_nonnull(graphProducer)) - return deferred.emitOpError("phase 2 requires deferred sources to be graph compute results"); - auto graphId = graphProducer ? graphProducer->getAttrOfType("scheduled.graph_id") : IntegerAttr(); - if (!graphId) - return deferred.emitOpError("phase 2 cannot identify the serialized graph producer"); - auto sourceUse = findSourceUse(body, sourceIndex); - std::optional sourceLane = sourceUse ? sourceUse->staticSourceLane : std::nullopt; - if (isa(graphProducer) && !sourceLane) - return deferred.emitOpError("phase 2 requires a statically resolved physical source slot for graph batch payloads"); - auto source = findProducer(plan, deferred, graphId.getInt(), sourceResult.getResultNumber(), sourceLane); - if (failed(source)) - return failure(); - if (failed(configurePayloadAndReplacement(deferred, body, **source, sourceIndex, sourceLane.value_or(0), lane))) - return failure(); +static LogicalResult retargetBlueprintToScheduledPublications( + RealizationPlan &plan, SpatBlueprintOp blueprint) { + if (blueprint.getMode() != "fragment_assembly") + return success(); + auto operandIndices = blueprint.getFragmentOperandIndices(); + auto sourceSlots = blueprint.getFragmentSourceSlots(); + if (!operandIndices || !sourceSlots) + return blueprint.emitOpError("phase 2 requires explicit Blueprint fragment ownership"); + SmallVector oldOperands {blueprint.getInput()}; + llvm::append_range(oldOperands, blueprint.getFragments()); + SmallVector publications; + SmallVector scheduledOperandIndices; + SmallVector scheduledSourceSlots; + for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) { + if (operandIndex < 0 || operandIndex >= static_cast(oldOperands.size())) + return blueprint.emitOpError("phase 2 Blueprint operand index is out of range"); + auto sourceResult = dyn_cast(oldOperands[operandIndex]); + auto graphBatch = sourceResult + ? dyn_cast(sourceResult.getOwner()) + : SpatGraphComputeBatch(); + if (!graphBatch) { + Value existing = oldOperands[operandIndex]; + auto it = llvm::find(publications, existing); + if (it == publications.end()) { + publications.push_back(existing); + it = std::prev(publications.end()); + } + scheduledOperandIndices.push_back(std::distance(publications.begin(), it)); + scheduledSourceSlots.push_back((*sourceSlots)[fragmentIndex]); + continue; + } + auto graphId = graphBatch->getAttrOfType("scheduled.graph_id"); + auto publication = getGraphBatchPublicationMap( + graphBatch, sourceResult.getResultNumber(), plan.publicationCache); + int64_t physicalSlot = (*sourceSlots)[fragmentIndex]; + if (!graphId || failed(publication) || physicalSlot < 0 || + physicalSlot >= static_cast((*publication)->physicalSlotToGraphLane.size())) + return blueprint.emitOpError("phase 2 cannot resolve Blueprint physical fragment ownership"); + int64_t graphLane = (*publication)->physicalSlotToGraphLane[physicalSlot]; + auto producer = findProducer(plan, blueprint.getOperation(), graphId.getInt(), + sourceResult.getResultNumber(), graphLane); + if (failed(producer)) + return failure(); + auto it = llvm::find(publications, (*producer)->published); + if (it == publications.end()) { + publications.push_back((*producer)->published); + it = std::prev(publications.end()); + } + scheduledOperandIndices.push_back(std::distance(publications.begin(), it)); + scheduledSourceSlots.push_back((*producer)->scheduled->isBatch() + ? (*producer)->scheduledLane + : graphLane - (*producer)->laneStart); + Operation *scheduled = (*producer)->published.getDefiningOp(); + if (blueprint->getBlock() == scheduled->getBlock() && + blueprint->isBeforeInBlock(scheduled)) + blueprint->moveAfter(scheduled); + } + blueprint->setOperands(publications); + blueprint->setAttr("fragmentOperandIndices", + OpBuilder(blueprint).getDenseI64ArrayAttr(scheduledOperandIndices)); + blueprint->setAttr("fragmentSourceSlots", + OpBuilder(blueprint).getDenseI64ArrayAttr(scheduledSourceSlots)); + return success(); +} - lane.source = *source; - lane.sourceCore = (*source)->core; - lane.targetCore = target.cores[target.isBatch() ? targetLane : 0]; - lane.targetLane = targetLane; - lane.sourceStream = - (*source)->scheduled->streamIds[(*source)->scheduled->isBatch() - ? llvm::find((*source)->scheduled->cores, (*source)->core) - (*source)->scheduled->cores.begin() - : 0]; - lane.targetStream = target.streamIds[target.isBatch() ? targetLane : 0]; - lane.local = (*source)->scheduled == &target && lane.sourceStream == lane.targetStream && (*source)->step < consumerStep; +static LogicalResult configureFragmentTransfer(NormalizedDeferredExchange &exchange, + DeferredSpecialization &specialization, + unsigned requirementIndex, + FragmentTransfer &transfer) { + FragmentRequirement &requirement = specialization.requirements[requirementIndex]; + ProducedValue *producer = requirement.producer; + transfer.requirementIndex = requirementIndex; + transfer.sourceStream = producer->scheduled->streamIds[producer->scheduledLane]; + transfer.sourceCore = producer->core; + unsigned targetLane = specialization.targetScheduledLane.value_or(0); + transfer.targetStream = exchange.target->streamIds[targetLane]; + transfer.targetCore = exchange.target->cores[targetLane]; + transfer.local = transfer.sourceStream == transfer.targetStream && + producer->step < exchange.consumerStep; + return success(); +} + +static LogicalResult buildDeferredSpecialization(RealizationPlan &plan, + NormalizedDeferredExchange &exchange, + std::optional targetScheduledLane) { + auto program = analyzeDeferredProgram(exchange.deferred, targetScheduledLane); + if (failed(program)) + return exchange.deferred.emitOpError("phase 2 deferred projection analysis failed"); + DeferredSpecialization specialization; + specialization.targetScheduledLane = targetScheduledLane; + specialization.program = std::move(*program); + for (auto [leafIndex, leaf] : llvm::enumerate(specialization.program.leaves)) { + Value sourceValue = exchange.deferred.getSources()[leaf.sourceOperandIndex]; + auto sourceResult = dyn_cast(sourceValue); + if (!sourceResult) + return exchange.deferred.emitOpError("phase 2 requires graph-result deferred sources"); + Operation *owner = sourceResult.getOwner(); + auto graphId = owner->getAttrOfType("scheduled.graph_id"); + if (!graphId) + return exchange.deferred.emitOpError("phase 2 cannot identify graph producer"); + unsigned requirementLeafIndex = leafIndex; + unsigned sourceOperandIndex = leaf.sourceOperandIndex; + auto append = [&](std::optional physicalSlot, std::optional graphLane, + Type fragmentType, unsigned selectedPosition) -> LogicalResult { + auto producer = findProducer(plan, exchange.deferred.getOperation(), graphId.getInt(), + sourceResult.getResultNumber(), graphLane); + if (failed(producer)) return failure(); + FragmentRequirement requirement {requirementLeafIndex, selectedPosition, + sourceOperandIndex, physicalSlot, graphLane, *producer, + fragmentType, targetScheduledLane}; + specialization.requirements.push_back(std::move(requirement)); + FragmentTransfer transfer; + if (failed(configureFragmentTransfer(exchange, specialization, + specialization.requirements.size() - 1, transfer))) + return failure(); + specialization.transfers.push_back(std::move(transfer)); + return success(); + }; + if (leaf.kind == DeferredLeafKind::ScalarSource) { + if (!isa(owner)) + return exchange.deferred.emitOpError("scalar deferred leaf requires graph_compute"); + if (failed(append(std::nullopt, std::nullopt, sourceResult.getType(), 0))) return failure(); + continue; + } + auto graphBatch = dyn_cast(owner); + if (!graphBatch) + return exchange.deferred.emitOpError("projection deferred leaf requires graph_compute_batch"); + auto publication = getGraphBatchPublicationMap(graphBatch, sourceResult.getResultNumber(), + plan.publicationCache); + if (failed(publication)) + return exchange.deferred.emitOpError("phase 2 graph-batch publication analysis failed"); + for (auto [position, physicalSlot] : llvm::enumerate(leaf.physicalSlots)) { + if (physicalSlot < 0 || physicalSlot >= static_cast((*publication)->physicalSlotToGraphLane.size())) + return exchange.deferred.emitOpError("projection physical slot is outside publication map"); + int64_t graphLane = (*publication)->physicalSlotToGraphLane[physicalSlot]; + if (failed(append(physicalSlot, graphLane, (*publication)->publicationFragmentType, position))) return failure(); + } + } + exchange.specializations.push_back(std::move(specialization)); return success(); } @@ -575,59 +469,21 @@ static LogicalResult normalizeDeferred(func::FuncOp funcOp, RealizationPlan& pla auto step = target && scheduledBlock ? getStepIndex(*target, scheduledBlock) : FailureOr(failure()); if (!target || failed(step)) return deferred.emitOpError("phase 2 cannot locate the deferred consumer scheduled step"); - auto body = analyzeDeferredBody(deferred); - if (failed(body)) - return failure(); - - bool batched = target->isBatch(); - if (!batched) { - NormalizedScalarExchange exchange; - exchange.id = nextExchangeId++; - exchange.deferred = deferred; - exchange.target = target; - exchange.consumerStep = *step; - exchange.body = *body; - if (failed(buildLaneTransfer(plan, - deferred, - exchange.body, - *target, - *step, - 0, - 0, - exchange.transfer))) - return failure(); - plan.scalar.push_back(std::move(exchange)); - continue; - } - - size_t lanes = target->cores.size(); - - NormalizedBatchedExchange exchange; - exchange.id = nextExchangeId++; + NormalizedDeferredExchange exchange; + exchange.exchangeId = nextExchangeId++; exchange.deferred = deferred; exchange.target = target; exchange.consumerStep = *step; - exchange.body = *body; - for (unsigned laneIndex = 0; laneIndex < lanes; ++laneIndex) { - auto laneBody = analyzeDeferredBody(deferred, laneIndex); - if (failed(laneBody)) - return failure(); - if (laneBody->sourceUses.size() != 1) - return deferred.emitOpError("phase 2 scheduled-batch specialization must resolve exactly one source"); - unsigned sourceIndex = laneBody->sourceUses.front().sourceIndex; - LaneTransfer lane; - if (failed(buildLaneTransfer(plan, - deferred, - *laneBody, - *target, - *step, - laneIndex, - sourceIndex, - lane))) - return failure(); - exchange.lanes.push_back(std::move(lane)); + if (!target->isBatch()) { + if (failed(buildDeferredSpecialization(plan, exchange, std::nullopt))) + return deferred.emitOpError("phase 2 failed to build scalar deferred specialization"); + } else { + for (unsigned lane = 0; lane < target->cores.size(); ++lane) + if (failed(buildDeferredSpecialization(plan, exchange, lane))) + return deferred.emitOpError() + << "phase 2 failed to build deferred specialization for scheduled lane " << lane; } - plan.batched.push_back(std::move(exchange)); + plan.exchanges.push_back(std::move(exchange)); } return success(); } @@ -657,14 +513,13 @@ static LogicalResult scheduleCommunication(func::FuncOp funcOp, RealizationPlan& } SmallVector readyGroups; - for (auto& entry : groups) { + for (auto &entry : groups) { if (completedTransfers[entry.second.front()]) continue; - bool ready = llvm::all_of(entry.second, [&](unsigned index) { - const PlannedCommunicationTransfer& transfer = plan.external[index]; - return completedSteps[transfer.sourceStream] > transfer.producerStep; - }); - if (ready) + if (llvm::all_of(entry.second, [&](unsigned index) { + const PlannedCommunicationTransfer &transfer = plan.external[index]; + return completedSteps[transfer.sourceStream] > transfer.producerStep; + })) readyGroups.push_back(entry.first); } if (!readyGroups.empty()) { @@ -675,7 +530,7 @@ static LogicalResult scheduleCommunication(func::FuncOp funcOp, RealizationPlan& unsigned source = std::numeric_limits::max(); unsigned target = std::numeric_limits::max(); for (unsigned index : indices) { - const PlannedCommunicationTransfer& transfer = plan.external[index]; + const PlannedCommunicationTransfer &transfer = plan.external[index]; consumer = std::min(consumer, transfer.consumerStep); source = std::min(source, transfer.sourceStream); target = std::min(target, transfer.targetStream); @@ -684,13 +539,13 @@ static LogicalResult scheduleCommunication(func::FuncOp funcOp, RealizationPlan& }; return priority(lhs) < priority(rhs); }); - ArrayRef selected = groups[readyGroups.front()]; - SmallVector ordered(selected.begin(), selected.end()); + SmallVector ordered(groups[readyGroups.front()].begin(), + groups[readyGroups.front()].end()); llvm::stable_sort(ordered, [&](unsigned lhs, unsigned rhs) { - const auto& left = plan.external[lhs]; - const auto& right = plan.external[rhs]; - return std::tie(left.sourceStream, left.targetStream, left.exchangeId) - < std::tie(right.sourceStream, right.targetStream, right.exchangeId); + const auto &left = plan.external[lhs]; + const auto &right = plan.external[rhs]; + return std::tie(left.sourceStream, left.targetStream, left.exchangeId) < + std::tie(right.sourceStream, right.targetStream, right.exchangeId); }); for (unsigned index : ordered) { completedTransfers[index] = true; @@ -701,82 +556,64 @@ static LogicalResult scheduleCommunication(func::FuncOp funcOp, RealizationPlan& } progressed = true; } - if (!progressed) - return funcOp.emitOpError("global communication scheduler made no progress before IR mutation"); + if (!progressed) { + InFlightDiagnostic diagnostic = funcOp.emitOpError( + "global communication scheduler made no progress before IR mutation"); + unsigned reported = 0; + for (auto [index, transfer] : llvm::enumerate(plan.external)) { + if (completedTransfers[index] || reported == 8) + continue; + diagnostic << (reported++ == 0 ? "; blocked " : ", ") + << "channel " << transfer.exchangeId << " source stream " + << transfer.sourceStream << " at step " + << completedSteps[transfer.sourceStream] << " after producer " + << transfer.producerStep << " target stream " + << transfer.targetStream << " at step " + << completedSteps[transfer.targetStream] << " before consumer " + << transfer.consumerStep; + } + return failure(); + } } plan.external = std::move(scheduled); return success(); } static LogicalResult buildAndVerifyPlan(func::FuncOp funcOp, RealizationPlan& plan) { - for (Operation& op : funcOp.getOps()) { - if (!isa(op)) - continue; - auto graphId = op.getAttrOfType("scheduled.graph_id"); - if (!graphId) - continue; - for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) { - bool hasExternalUse = llvm::any_of(result.getUses(), [](OpOperand& use) { - return !isa(use.getOwner()); - }); - if (!hasExternalUse) - continue; - SmallVector exact; - for (ProducedValue* produced : plan.producedByGraph.lookup(graphId.getInt())) - if (produced->resultIndex == resultIndex && produced->published.getType() == result.getType() - && !llvm::is_contained(exact, produced->published)) - exact.push_back(produced->published); - if (exact.size() != 1) - return op.emitOpError("phase 2 cannot prove unique final publication ownership for a partitioned graph result"); - } - } - uint64_t nextChannel = 0; - for (NormalizedScalarExchange& exchange : plan.scalar) { - if (exchange.transfer.local) - continue; - exchange.transfer.channelId = nextChannel++; - plan.external.push_back({ - exchange.transfer.channelId, - exchange.id, - exchange.transfer.sourceStream, - exchange.transfer.targetStream, - exchange.transfer.source->step, - exchange.consumerStep, - }); - } - for (NormalizedBatchedExchange& exchange : plan.batched) - for (LaneTransfer& lane : exchange.lanes) { - if (lane.local) - continue; - lane.channelId = nextChannel++; - plan.external.push_back({ - lane.channelId, - exchange.id, - lane.sourceStream, - lane.targetStream, - lane.source->step, - exchange.consumerStep, - }); - } + for (NormalizedDeferredExchange &exchange : plan.exchanges) + for (DeferredSpecialization &specialization : exchange.specializations) + for (FragmentTransfer &transfer : specialization.transfers) { + if (transfer.local) + continue; + transfer.channelId = nextChannel++; + ++exchange.externalTransferCount; + const FragmentRequirement &requirement = specialization.requirements[transfer.requirementIndex]; + plan.external.push_back({transfer.channelId, exchange.exchangeId, + transfer.sourceStream, transfer.targetStream, + requirement.producer->step, exchange.consumerStep}); + } if (failed(scheduleCommunication(funcOp, plan))) return failure(); DenseMap plannedByChannel; - for (PlannedCommunicationTransfer& transfer : plan.external) + DenseMap orderByChannel; + for (auto [order, transfer] : llvm::enumerate(plan.external)) { plannedByChannel[transfer.exchangeId] = &transfer; - auto applyInsertionSteps = [&](LaneTransfer& lane) { - if (lane.local) + orderByChannel[transfer.exchangeId] = order; + } + auto applyInsertionSteps = [&](FragmentTransfer& transfer) { + if (transfer.local) return; - PlannedCommunicationTransfer* transfer = plannedByChannel.lookup(lane.channelId); - assert(transfer && "every external lane must have one planned transfer"); - lane.sourceInsertionStep = transfer->sourceInsertionStep; - lane.targetInsertionStep = transfer->targetInsertionStep; + PlannedCommunicationTransfer* planned = plannedByChannel.lookup(transfer.channelId); + assert(planned && "every external fragment must have one planned transfer"); + transfer.globalOrder = orderByChannel.lookup(transfer.channelId); + transfer.sourceInsertionStep = planned->sourceInsertionStep; + transfer.targetInsertionStep = planned->targetInsertionStep; }; - for (NormalizedScalarExchange& exchange : plan.scalar) - applyInsertionSteps(exchange.transfer); - for (NormalizedBatchedExchange& exchange : plan.batched) - for (LaneTransfer& lane : exchange.lanes) - applyInsertionSteps(lane); + for (NormalizedDeferredExchange &exchange : plan.exchanges) + for (DeferredSpecialization &specialization : exchange.specializations) + for (FragmentTransfer &transfer : specialization.transfers) + applyInsertionSteps(transfer); return verifyPlannedCommunicationDeadlockFree(funcOp, plan.stepCounts.size(), plan.stepCounts, plan.external); } @@ -904,48 +741,6 @@ setTransferAttrs(Operation* op, uint64_t exchangeId, uint64_t channelId, int64_t op->setAttr("raptor.target_core", builder.getI64IntegerAttr(targetCore)); } -static FailureOr materializeSourcePayload(DeferredBodyAnalysis& body, LaneTransfer& lane, IRRewriter& rewriter) { - if (body.sourceSideOps.empty()) - return lane.source->payload; - - Block& deferredBlock = - cast(body.yieldedValue.getParentBlock()->getParentOp()).getBody().front(); - IRMapping mapping; - BlockArgument sourceArgument = deferredBlock.getArgument(lane.sourceIndex); - if (sourceArgument.getType() == lane.source->payload.getType()) - mapping.map(sourceArgument, lane.source->payload); - if (lane.bodyRoot) - for (Value result : lane.bodyRoot->getResults()) - if (result.getType() == lane.source->payload.getType()) - mapping.map(result, lane.source->payload); - - for (Operation* op : body.sourceSideOps) { - if (op == lane.bodyRoot) - continue; - bool needed = llvm::any_of(op->getResults(), [&](Value result) { - if (result == body.yieldedValue) - return true; - return llvm::any_of(body.sourceSideOps, [&](Operation* candidate) { - return candidate != lane.bodyRoot && llvm::is_contained(candidate->getOperands(), result); - }); - }); - if (!needed) - continue; - for (Value operand : op->getOperands()) { - if (mapping.contains(operand) || operand.getParentBlock() != &deferredBlock) - continue; - return op->emitOpError("source-side deferred-body clone has an unmapped source operand"); - } - Operation* clone = rewriter.clone(*op, mapping); - for (auto [result, clonedResult] : llvm::zip(op->getResults(), clone->getResults())) - mapping.map(result, clonedResult); - } - Value payload = mapping.lookupOrDefault(body.yieldedValue); - if (!payload || payload.getType() != lane.transportType) - return failure(); - return payload; -} - static void replaceDeferred(SpatDeferredCommunicationOp deferred, Value replacement, SmallVectorImpl& erase) { @@ -954,13 +749,6 @@ static void replaceDeferred(SpatDeferredCommunicationOp deferred, erase.push_back(deferred); } -static Value createI64TableValue(OpBuilder& builder, Location loc, ArrayRef values, Value index) { - RankedTensorType type = RankedTensorType::get({static_cast(values.size())}, builder.getI64Type()); - Value table = arith::ConstantOp::create(builder, loc, type, DenseElementsAttr::get(type, values)); - Value value = tensor::ExtractOp::create(builder, loc, table, ValueRange {index}); - return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), value); -} - static void setInsertionAtBoundary(IRRewriter& rewriter, ScheduledInfo& scheduled, unsigned step) { if (step < scheduled.stepAnchors.size() && scheduled.stepAnchors[step]->getBlock()) { rewriter.setInsertionPoint(scheduled.stepAnchors[step]); @@ -969,222 +757,275 @@ static void setInsertionAtBoundary(IRRewriter& rewriter, ScheduledInfo& schedule rewriter.setInsertionPoint(scheduled.blocks.front()->getTerminator()); } -static unsigned countExternalLanes(const NormalizedBatchedExchange& exchange) { - return llvm::count_if(exchange.lanes, [](const LaneTransfer& lane) { return !lane.local; }); -} - -static LogicalResult emitGuardedBatchSends(NormalizedBatchedExchange& exchange, - ScheduledInfo& source, - ArrayRef lanes, - IRRewriter& rewriter) { - Value laneIndex = cast(source.op).getLaneArgument().value(); - unsigned transferCount = countExternalLanes(exchange); - for (LaneTransfer* lane : lanes) { - auto coreIt = llvm::find(source.cores, lane->sourceCore); - if (coreIt == source.cores.end()) - return exchange.deferred.emitOpError("phase 2 cannot locate a partial batch send source lane"); - setInsertionAtBoundary(rewriter, *lane->source->scheduled, lane->sourceInsertionStep); - Location loc = exchange.deferred.getLoc(); - Value sourceLane = arith::ConstantIndexOp::create(rewriter, loc, std::distance(source.cores.begin(), coreIt)); - Value selected = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, laneIndex, sourceLane); - auto ifOp = scf::IfOp::create(rewriter, loc, selected, false); - rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); - auto payload = materializeSourcePayload(exchange.body, *lane, rewriter); - if (failed(payload)) - return exchange.deferred.emitOpError("failed to materialize partial batch source payload"); - Value channel = arith::ConstantIndexOp::create(rewriter, loc, lane->channelId); - Value sourceCore = arith::ConstantIndexOp::create(rewriter, loc, lane->sourceCore); - Value targetCore = arith::ConstantIndexOp::create(rewriter, loc, lane->targetCore); - auto send = SpatChannelSendOp::create(rewriter, loc, channel, sourceCore, targetCore, *payload); - setTransferAttrs(send, lane->channelId, lane->channelId, lane->sourceCore, lane->targetCore); - send->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); - send->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(transferCount)); +static FailureOr materializeProducedFragment(const FragmentRequirement &requirement, + IRRewriter &rewriter, Location loc) { + Value payload = requirement.producer->payload; + if (!requirement.graphLane) { + if (payload.getType() != requirement.publicationFragmentType) + return failure(); + return payload; } - return success(); -} - -static LogicalResult emitCompactBatchSends(NormalizedBatchedExchange& exchange, - IRRewriter& rewriter) { - llvm::MapVector> bySource; - for (LaneTransfer& lane : exchange.lanes) - if (!lane.local) - bySource[lane.source->scheduled].push_back(&lane); - - for (auto& entry : bySource) { - ScheduledInfo* source = entry.first; - ArrayRef lanes = entry.second; - if (!source->isBatch()) { - for (LaneTransfer* lane : lanes) { - setInsertionAtBoundary(rewriter, *lane->source->scheduled, lane->sourceInsertionStep); - auto payload = materializeSourcePayload(exchange.body, *lane, rewriter); - if (failed(payload)) - return exchange.deferred.emitOpError("failed to materialize scalar source payload"); - Location loc = exchange.deferred.getLoc(); - Value channel = arith::ConstantIndexOp::create(rewriter, loc, lane->channelId); - Value sourceCore = arith::ConstantIndexOp::create(rewriter, loc, lane->sourceCore); - Value targetCore = arith::ConstantIndexOp::create(rewriter, loc, lane->targetCore); - auto send = SpatChannelSendOp::create(rewriter, loc, channel, sourceCore, targetCore, *payload); - setTransferAttrs(send, lane->channelId, lane->channelId, lane->sourceCore, lane->targetCore); - send->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); - send->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(countExternalLanes(exchange))); - } - continue; - } - if (lanes.size() != source->cores.size()) { - if (failed(emitGuardedBatchSends(exchange, *source, lanes, rewriter))) - return failure(); - continue; - } - SmallVector channels(source->cores.size(), -1); - SmallVector sourceCores(source->cores.size(), -1); - SmallVector targetCores(source->cores.size(), -1); - Value sourcePayload = lanes.front()->source->payload; - for (LaneTransfer* lane : lanes) { - auto coreIt = llvm::find(source->cores, lane->sourceCore); - if (coreIt == source->cores.end() || lane->source->payload != sourcePayload - || lane->sourceIndex != lanes.front()->sourceIndex || lane->bodyRoot != lanes.front()->bodyRoot - || lane->transportType != lanes.front()->transportType) - return exchange.deferred.emitOpError("phase 2 cannot compact heterogeneous batch sends"); - unsigned sourceLane = std::distance(source->cores.begin(), coreIt); - channels[sourceLane] = lane->channelId; - sourceCores[sourceLane] = lane->sourceCore; - targetCores[sourceLane] = lane->targetCore; - } - if (llvm::is_contained(channels, -1)) - return exchange.deferred.emitOpError("phase 2 compact batch send has an unmapped physical lane"); - - setInsertionAtBoundary(rewriter, *lanes.front()->source->scheduled, lanes.front()->sourceInsertionStep); - auto payload = materializeSourcePayload(exchange.body, *lanes.front(), rewriter); - if (failed(payload)) - return exchange.deferred.emitOpError("failed to materialize compact batch source payload"); - Location loc = exchange.deferred.getLoc(); - Value laneIndex = cast(source->op).getLaneArgument().value(); - Value channel = createI64TableValue(rewriter, loc, channels, laneIndex); - Value sourceCore = createI64TableValue(rewriter, loc, sourceCores, laneIndex); - Value targetCore = createI64TableValue(rewriter, loc, targetCores, laneIndex); - auto send = SpatChannelSendOp::create(rewriter, loc, channel, sourceCore, targetCore, *payload); - send->setAttr("raptor.batch_channel_ids", rewriter.getDenseI64ArrayAttr(channels)); - send->setAttr("raptor.batch_source_cores", rewriter.getDenseI64ArrayAttr(sourceCores)); - send->setAttr("raptor.batch_target_cores", rewriter.getDenseI64ArrayAttr(targetCores)); - send->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); - send->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(countExternalLanes(exchange))); - } - return success(); -} - -static LogicalResult -emitCompactBatchReceive(NormalizedBatchedExchange& exchange, IRRewriter& rewriter, SmallVectorImpl& erase) { - ScheduledInfo& target = *exchange.target; - if (!target.isBatch() || exchange.lanes.size() != target.cores.size() - || llvm::any_of(exchange.lanes, [](const LaneTransfer& lane) { return lane.local; })) - return exchange.deferred.emitOpError( - "phase 2 compact batch receive requires one external transfer for every target lane"); - - SmallVector channels(target.cores.size(), -1); - SmallVector sourceCores(target.cores.size(), -1); - SmallVector targetCores(target.cores.size(), -1); - Type payloadType = exchange.lanes.front().transportType; - for (LaneTransfer& lane : exchange.lanes) { - if (lane.targetLane >= target.cores.size() || lane.transportType != payloadType) - return exchange.deferred.emitOpError("phase 2 cannot compact heterogeneous batch receives"); - channels[lane.targetLane] = lane.channelId; - sourceCores[lane.targetLane] = lane.sourceCore; - targetCores[lane.targetLane] = lane.targetCore; - } - if (llvm::is_contained(channels, -1)) - return exchange.deferred.emitOpError("phase 2 compact batch receive has an unmapped target lane"); - - setInsertionAtBoundary(rewriter, target, exchange.lanes.front().targetInsertionStep); - Location loc = exchange.deferred.getLoc(); - Value laneIndex = cast(target.op).getLaneArgument().value(); - Value channel = createI64TableValue(rewriter, loc, channels, laneIndex); - Value sourceCore = createI64TableValue(rewriter, loc, sourceCores, laneIndex); - Value targetCore = createI64TableValue(rewriter, loc, targetCores, laneIndex); - auto receive = SpatChannelReceiveOp::create(rewriter, loc, payloadType, channel, sourceCore, targetCore); - receive->setAttr("raptor.batch_channel_ids", rewriter.getDenseI64ArrayAttr(channels)); - receive->setAttr("raptor.batch_source_cores", rewriter.getDenseI64ArrayAttr(sourceCores)); - receive->setAttr("raptor.batch_target_cores", rewriter.getDenseI64ArrayAttr(targetCores)); - receive->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); - receive->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(countExternalLanes(exchange))); - replaceDeferred(exchange.deferred, receive.getOutput(), erase); - return success(); -} - -static LogicalResult realizeScalarExchange(NormalizedScalarExchange& exchange, - IRRewriter& rewriter, - SmallVectorImpl& erase) { - LaneTransfer& transfer = exchange.transfer; - if (transfer.local) - rewriter.setInsertionPointAfterValue(transfer.source->payload); - else - setInsertionAtBoundary(rewriter, *transfer.source->scheduled, transfer.sourceInsertionStep); - auto payload = materializeSourcePayload(exchange.body, transfer, rewriter); - if (failed(payload)) - return exchange.deferred.emitOpError("failed to materialize scalar source payload"); - if (transfer.local) { - replaceDeferred(exchange.deferred, *payload, erase); - return success(); - } - - Location loc = exchange.deferred.getLoc(); - Value channel = arith::ConstantIndexOp::create(rewriter, loc, transfer.channelId); - Value sourceCore = arith::ConstantIndexOp::create(rewriter, loc, transfer.sourceCore); - Value targetCore = arith::ConstantIndexOp::create(rewriter, loc, transfer.targetCore); - auto createSend = [&] { - auto send = SpatChannelSendOp::create(rewriter, loc, channel, sourceCore, targetCore, *payload); - setTransferAttrs(send, transfer.channelId, transfer.channelId, transfer.sourceCore, transfer.targetCore); - send->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); - send->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(1)); - }; - if (transfer.source->scheduled->isBatch()) { - auto batch = cast(transfer.source->scheduled->op); - auto it = llvm::find(transfer.source->scheduled->cores, transfer.sourceCore); - if (it == transfer.source->scheduled->cores.end()) - return exchange.deferred.emitOpError("phase 2 cannot locate scalar exchange source batch lane"); - Value lane = arith::ConstantIndexOp::create(rewriter, loc, std::distance(transfer.source->scheduled->cores.begin(), it)); - Value selected = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, batch.getLaneArgument().value(), lane); - auto ifOp = scf::IfOp::create(rewriter, loc, selected, false); - rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); - createSend(); - } else { - createSend(); - } - - setInsertionAtBoundary(rewriter, *exchange.target, transfer.targetInsertionStep); - channel = arith::ConstantIndexOp::create(rewriter, loc, transfer.channelId); - sourceCore = arith::ConstantIndexOp::create(rewriter, loc, transfer.sourceCore); - targetCore = arith::ConstantIndexOp::create(rewriter, loc, transfer.targetCore); - auto receive = SpatChannelReceiveOp::create(rewriter, loc, transfer.transportType, channel, sourceCore, targetCore); - setTransferAttrs(receive, transfer.channelId, transfer.channelId, transfer.sourceCore, transfer.targetCore); - receive->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.id)); - receive->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(1)); - replaceDeferred(exchange.deferred, receive.getOutput(), erase); - return success(); -} - -static LogicalResult realizeBatchedExchange(NormalizedBatchedExchange& exchange, - IRRewriter& rewriter, - SmallVectorImpl& erase) { - bool allLocal = llvm::all_of(exchange.lanes, [](const LaneTransfer& lane) { return lane.local; }); - if (allLocal) { - LaneTransfer& first = exchange.lanes.front(); - for (LaneTransfer& lane : exchange.lanes) - if (lane.source->payload != first.source->payload || lane.sourceIndex != first.sourceIndex - || lane.bodyRoot != first.bodyRoot) - return exchange.deferred.emitOpError("phase 2 cannot compact heterogeneous local batch payloads"); - rewriter.setInsertionPointAfterValue(first.source->payload); - auto payload = materializeSourcePayload(exchange.body, first, rewriter); - if (failed(payload)) - return exchange.deferred.emitOpError("failed to materialize local compact batch payload"); - replaceDeferred(exchange.deferred, *payload, erase); - return success(); - } - if (llvm::any_of(exchange.lanes, [](const LaneTransfer& lane) { return lane.local; })) - return exchange.deferred.emitOpError( - "phase 2 cannot compact a batch exchange mixing local and external target lanes"); - if (failed(emitCompactBatchSends(exchange, rewriter)) - || failed(emitCompactBatchReceive(exchange, rewriter, erase))) + if (payload.getType() == requirement.publicationFragmentType) + return payload; + auto payloadType = dyn_cast(payload.getType()); + auto fragmentType = dyn_cast(requirement.publicationFragmentType); + if (!payloadType || !fragmentType || payloadType.getRank() != fragmentType.getRank() + 1 || + payloadType.getDimSize(0) != requirement.producer->laneCount || + !llvm::equal(payloadType.getShape().drop_front(), fragmentType.getShape())) return failure(); + int64_t localOffset = *requirement.graphLane - requirement.producer->laneStart; + if (localOffset < 0 || localOffset >= requirement.producer->laneCount) + return failure(); + SmallVector offsets(payloadType.getRank(), rewriter.getIndexAttr(0)); + SmallVector sizes, strides(payloadType.getRank(), rewriter.getIndexAttr(1)); + offsets[0] = rewriter.getIndexAttr(localOffset); + sizes.push_back(rewriter.getIndexAttr(1)); + for (int64_t dim : fragmentType.getShape()) sizes.push_back(rewriter.getIndexAttr(dim)); + SmallVector unitShape {1}; llvm::append_range(unitShape, fragmentType.getShape()); + auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType()); + Value unit = tensor::ExtractSliceOp::create(rewriter, loc, unitType, payload, offsets, sizes, strides); + SmallVector reassociation {{0, 1}}; + for (int64_t dim = 1; dim < fragmentType.getRank(); ++dim) reassociation.push_back({dim + 1}); + return tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, unit, reassociation).getResult(); +} + +static FailureOr applyInnerFragmentGeometry(Value fragment, const DeferredProjectionLeaf &leaf, + IRRewriter &rewriter, Location loc) { + if (leaf.kind != DeferredLeafKind::GraphBatchProjection) + return fragment; + auto type = dyn_cast(fragment.getType()); + if (!type || leaf.innerGeometry.offsets.size() != static_cast(type.getRank())) return failure(); + bool identity = true; + for (unsigned i = 0; i < type.getRank(); ++i) + identity &= leaf.innerGeometry.offsets[i] == 0 && leaf.innerGeometry.sizes[i] == type.getDimSize(i) && + leaf.innerGeometry.strides[i] == 1; + if (identity) return fragment; + if (leaf.reconstructedType.getRank() != type.getRank() + 1) return failure(); + SmallVector shape(leaf.reconstructedType.getShape().drop_front()); + auto innerType = RankedTensorType::get(shape, type.getElementType()); + SmallVector offsets, sizes, strides; + for (unsigned i = 0; i < type.getRank(); ++i) { + offsets.push_back(rewriter.getIndexAttr(leaf.innerGeometry.offsets[i])); + sizes.push_back(rewriter.getIndexAttr(leaf.innerGeometry.sizes[i])); + strides.push_back(rewriter.getIndexAttr(leaf.innerGeometry.strides[i])); + } + return tensor::ExtractSliceOp::create(rewriter, loc, innerType, fragment, offsets, sizes, strides).getResult(); +} + +static FailureOr reconstructProjectionLeaf(const DeferredProjectionLeaf &leaf, unsigned leafIndex, + const DeferredSpecialization &specialization, + ArrayRef available, IRRewriter &rewriter, Location loc) { + SmallVector selected(leaf.kind == DeferredLeafKind::ScalarSource ? 1 : leaf.physicalSlots.size()); + for (auto [index, requirement] : llvm::enumerate(specialization.requirements)) + if (requirement.leafIndex == leafIndex && requirement.selectedPosition < selected.size()) { + if (selected[requirement.selectedPosition]) return failure(); + selected[requirement.selectedPosition] = available[index]; + } + if (llvm::any_of(selected, [](Value value) { return !value; })) return failure(); + if (leaf.kind == DeferredLeafKind::ScalarSource) return selected.front(); + if (selected.size() == 1 && selected.front().getType() == leaf.reconstructedType) return selected.front(); + Value result = tensor::EmptyOp::create(rewriter, loc, leaf.reconstructedType.getShape(), + leaf.reconstructedType.getElementType()); + auto outputType = leaf.reconstructedType; + for (auto [position, fragment] : llvm::enumerate(selected)) { + auto shaped = applyInnerFragmentGeometry(fragment, leaf, rewriter, loc); + if (failed(shaped)) return failure(); + auto fragmentType = cast((*shaped).getType()); + SmallVector unitShape {1}; llvm::append_range(unitShape, fragmentType.getShape()); + auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType()); + SmallVector reassociation {{0, 1}}; + for (int64_t dim = 1; dim < fragmentType.getRank(); ++dim) reassociation.push_back({dim + 1}); + Value unit = tensor::ExpandShapeOp::create(rewriter, loc, unitType, *shaped, reassociation); + SmallVector offsets(outputType.getRank(), rewriter.getIndexAttr(0)); + offsets[0] = rewriter.getIndexAttr(position); + SmallVector sizes, strides(outputType.getRank(), rewriter.getIndexAttr(1)); + sizes.push_back(rewriter.getIndexAttr(1)); + for (int64_t dim : fragmentType.getShape()) sizes.push_back(rewriter.getIndexAttr(dim)); + result = tensor::InsertSliceOp::create(rewriter, loc, unit, result, offsets, sizes, strides); + } + return result; +} + +static FailureOr materializeDeferredSpecializationResult(DeferredSpecialization &specialization, + ArrayRef available, + IRRewriter &rewriter, Location loc) { + IRMapping mapping; + for (auto [index, leaf] : llvm::enumerate(specialization.program.leaves)) { + auto reconstructed = reconstructProjectionLeaf(leaf, index, specialization, available, rewriter, loc); + if (failed(reconstructed)) return failure(); + mapping.map(leaf.replacementRoot, *reconstructed); + } + if (specialization.targetScheduledLane) { + mapping.map(specialization.program.scheduledLane, + arith::ConstantIndexOp::create(rewriter, loc, + *specialization.targetScheduledLane)); + } + for (Operation *op : specialization.program.residualOps) { + Operation *copy = rewriter.clone(*op, mapping); + for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair)); + } + Value result = mapping.lookupOrDefault(specialization.program.yieldedValue); + return result && result.getType() == specialization.program.deferred.getOutput().getType() + ? FailureOr(result) : FailureOr(failure()); +} + +static LogicalResult emitFragmentSendHere( + NormalizedDeferredExchange &exchange, + DeferredSpecialization &specialization, FragmentTransfer &transfer, + IRRewriter &rewriter) { + Location loc = exchange.deferred.getLoc(); + auto payload = materializeProducedFragment( + specialization.requirements[transfer.requirementIndex], rewriter, loc); + if (failed(payload)) return failure(); + auto send = SpatChannelSendOp::create( + rewriter, loc, + arith::ConstantIndexOp::create(rewriter, loc, transfer.channelId), + arith::ConstantIndexOp::create(rewriter, loc, transfer.sourceCore), + arith::ConstantIndexOp::create(rewriter, loc, transfer.targetCore), + *payload); + setTransferAttrs(send, transfer.channelId, transfer.channelId, + transfer.sourceCore, transfer.targetCore); + send->setAttr("raptor.parent_exchange_id", + rewriter.getI64IntegerAttr(exchange.exchangeId)); + send->setAttr("raptor.parent_transfer_count", + rewriter.getI64IntegerAttr(exchange.externalTransferCount)); + return success(); +} + +static LogicalResult emitFragmentSend(NormalizedDeferredExchange &exchange, DeferredSpecialization &specialization, + FragmentTransfer &transfer, IRRewriter &rewriter) { + if (transfer.local) return success(); + setInsertionAtBoundary(rewriter, *specialization.requirements[transfer.requirementIndex].producer->scheduled, + transfer.sourceInsertionStep); + Location loc = exchange.deferred.getLoc(); + ProducedValue *producer = specialization.requirements[transfer.requirementIndex].producer; + if (!producer->scheduled->isBatch()) + return emitFragmentSendHere(exchange, specialization, transfer, rewriter); + auto batch = cast(producer->scheduled->op); + Value lane = arith::ConstantIndexOp::create(rewriter, loc, producer->scheduledLane); + Value selected = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, *batch.getLaneArgument(), lane); + auto ifOp = scf::IfOp::create(rewriter, loc, selected, false); + rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); + return emitFragmentSendHere(exchange, specialization, transfer, rewriter); +} + +static FailureOr receiveFragment(NormalizedDeferredExchange &exchange, + DeferredSpecialization &specialization, + FragmentTransfer &transfer, IRRewriter &rewriter) { + if (transfer.local) + return materializeProducedFragment(specialization.requirements[transfer.requirementIndex], rewriter, + exchange.deferred.getLoc()); + Location loc = exchange.deferred.getLoc(); + auto receive = SpatChannelReceiveOp::create(rewriter, loc, + specialization.requirements[transfer.requirementIndex].publicationFragmentType, + arith::ConstantIndexOp::create(rewriter, loc, transfer.channelId), + arith::ConstantIndexOp::create(rewriter, loc, transfer.sourceCore), + arith::ConstantIndexOp::create(rewriter, loc, transfer.targetCore)); + setTransferAttrs(receive, transfer.channelId, transfer.channelId, transfer.sourceCore, transfer.targetCore); + receive->setAttr("raptor.parent_exchange_id", rewriter.getI64IntegerAttr(exchange.exchangeId)); + receive->setAttr("raptor.parent_transfer_count", rewriter.getI64IntegerAttr(exchange.externalTransferCount)); + return receive.getOutput(); +} + +static FailureOr realizeSpecialization(NormalizedDeferredExchange &exchange, + DeferredSpecialization &specialization, IRRewriter &rewriter) { + SmallVector available(specialization.requirements.size()); + for (FragmentTransfer &transfer : specialization.transfers) + if (transfer.local) { + auto value = receiveFragment(exchange, specialization, transfer, rewriter); + if (failed(value)) return failure(); + available[transfer.requirementIndex] = *value; + } + struct OrderedTransfer { + DeferredSpecialization *specialization; + FragmentTransfer *transfer; + }; + SmallVector ordered; + for (DeferredSpecialization &owner : exchange.specializations) + for (FragmentTransfer &transfer : owner.transfers) + if (!transfer.local) + ordered.push_back({&owner, &transfer}); + llvm::sort(ordered, [](const OrderedTransfer &lhs, const OrderedTransfer &rhs) { + return lhs.transfer->globalOrder < rhs.transfer->globalOrder; + }); + unsigned targetScheduledLane = specialization.targetScheduledLane.value_or(0); + for (OrderedTransfer event : ordered) { + FragmentRequirement &requirement = + event.specialization->requirements[event.transfer->requirementIndex]; + if (requirement.producer->scheduled == exchange.target && + requirement.producer->scheduledLane == targetScheduledLane && + failed(emitFragmentSendHere(exchange, *event.specialization, + *event.transfer, rewriter))) + return failure(); + if (event.specialization == &specialization) { + auto value = receiveFragment(exchange, specialization, *event.transfer, + rewriter); + if (failed(value)) return failure(); + available[event.transfer->requirementIndex] = *value; + } + } + return materializeDeferredSpecializationResult(specialization, available, rewriter, exchange.deferred.getLoc()); +} + +static LogicalResult realizeFragmentExchange(NormalizedDeferredExchange &exchange, IRRewriter &rewriter, + SmallVectorImpl &erase) { + if (exchange.specializations.size() != (exchange.target->isBatch() ? exchange.target->cores.size() : 1)) + return exchange.deferred.emitOpError("phase 2 specialization count does not match target lanes"); + if (!exchange.target->isBatch()) { + unsigned targetStep = exchange.consumerStep; + for (const FragmentTransfer &transfer : exchange.specializations.front().transfers) + if (!transfer.local) { + targetStep = transfer.targetInsertionStep; + break; + } + setInsertionAtBoundary(rewriter, *exchange.target, targetStep); + auto result = realizeSpecialization(exchange, exchange.specializations.front(), rewriter); + if (failed(result)) return failure(); + replaceDeferred(exchange.deferred, *result, erase); + return success(); + } + // Each branch specializes the deferred program to the explicit scheduled lane. + auto batch = cast(exchange.target->op); + auto resultType = dyn_cast(exchange.deferred.getOutput().getType()); + if (!resultType || !resultType.hasStaticShape()) + return exchange.deferred.emitOpError( + "phase 2 scheduled-batch lane dispatch requires a static ranked result"); + Location loc = exchange.deferred.getLoc(); + setInsertionAtBoundary(rewriter, *exchange.target, exchange.consumerStep); + Value dispatchBuffer = tensor::EmptyOp::create( + rewriter, loc, resultType.getShape(), resultType.getElementType()); + auto storeResult = [&](Value value) -> Value { + SmallVector offsets(resultType.getRank(), + rewriter.getIndexAttr(0)); + SmallVector sizes; + SmallVector strides(resultType.getRank(), + rewriter.getIndexAttr(1)); + for (int64_t dimension : resultType.getShape()) + sizes.push_back(rewriter.getIndexAttr(dimension)); + return tensor::InsertSliceOp::create(rewriter, loc, value, + dispatchBuffer, offsets, sizes, + strides); + }; + Value result = dispatchBuffer; + for (unsigned lane = 0; lane < exchange.specializations.size(); ++lane) { + Value required = arith::ConstantIndexOp::create(rewriter, loc, lane); + Value selected = arith::CmpIOp::create( + rewriter, loc, arith::CmpIPredicate::eq, *batch.getLaneArgument(), + required); + auto ifOp = scf::IfOp::create(rewriter, loc, TypeRange {resultType}, + selected, true); + rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); + auto value = realizeSpecialization(exchange, exchange.specializations[lane], + rewriter); + if (failed(value)) + return failure(); + scf::YieldOp::create(rewriter, loc, storeResult(*value)); + rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front()); + scf::YieldOp::create(rewriter, loc, result); + rewriter.setInsertionPointAfter(ifOp); + result = ifOp.getResult(0); + } + replaceDeferred(exchange.deferred, result, erase); return success(); } @@ -1201,36 +1042,43 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { } SmallVector erase; - for (NormalizedScalarExchange& exchange : plan.scalar) - if (exchange.transfer.local && failed(realizeScalarExchange(exchange, rewriter, erase))) - return failure(); - for (NormalizedBatchedExchange& exchange : plan.batched) - if (llvm::all_of(exchange.lanes, [](const LaneTransfer& lane) { return lane.local; }) - && failed(realizeBatchedExchange(exchange, rewriter, erase))) - return failure(); - + // `plan.external` is the scheduler's global order; never recover it from a map. + for (const PlannedCommunicationTransfer &planned : plan.external) { + bool emitted = false; + for (NormalizedDeferredExchange &exchange : plan.exchanges) + for (DeferredSpecialization &specialization : exchange.specializations) + for (FragmentTransfer &transfer : specialization.transfers) + if (!transfer.local && transfer.channelId == planned.exchangeId) { + FragmentRequirement &requirement = + specialization.requirements[transfer.requirementIndex]; + if (requirement.producer->scheduled == exchange.target) { + emitted = true; + continue; + } + if (failed(emitFragmentSend(exchange, specialization, transfer, rewriter))) + return exchange.deferred.emitOpError("phase 2 failed to emit a fragment send"); + emitted = true; + } + if (!emitted) return funcOp.emitOpError("phase 2 cannot locate planned fragment transfer"); + } DenseMap orderByParent; for (auto [order, transfer] : llvm::enumerate(plan.external)) orderByParent.try_emplace(transfer.parentExchangeId, order); - struct ExternalExchange { - unsigned order; - NormalizedScalarExchange* scalar = nullptr; - NormalizedBatchedExchange* batched = nullptr; - }; - SmallVector external; - for (NormalizedScalarExchange& exchange : plan.scalar) - if (!exchange.transfer.local) - external.push_back({orderByParent.lookup(exchange.id), &exchange, nullptr}); - for (NormalizedBatchedExchange& exchange : plan.batched) - if (llvm::any_of(exchange.lanes, [](const LaneTransfer& lane) { return !lane.local; })) - external.push_back({orderByParent.lookup(exchange.id), nullptr, &exchange}); - llvm::sort(external, [](const ExternalExchange& lhs, const ExternalExchange& rhs) { return lhs.order < rhs.order; }); - for (ExternalExchange& exchange : external) { - LogicalResult result = exchange.scalar ? realizeScalarExchange(*exchange.scalar, rewriter, erase) - : realizeBatchedExchange(*exchange.batched, rewriter, erase); - if (failed(result)) - return failure(); - } + for (NormalizedDeferredExchange &exchange : plan.exchanges) + if (exchange.externalTransferCount == 0 && + failed(realizeFragmentExchange(exchange, rewriter, erase))) + return exchange.deferred.emitOpError("phase 2 failed to realize a local fragment exchange"); + SmallVector external; + for (NormalizedDeferredExchange &exchange : plan.exchanges) + if (exchange.externalTransferCount != 0) + external.push_back(&exchange); + llvm::stable_sort(external, [&](NormalizedDeferredExchange *lhs, + NormalizedDeferredExchange *rhs) { + return orderByParent.lookup(lhs->exchangeId) < orderByParent.lookup(rhs->exchangeId); + }); + for (NormalizedDeferredExchange *exchange : external) + if (failed(realizeFragmentExchange(*exchange, rewriter, erase))) + return exchange->deferred.emitOpError("phase 2 failed to realize an external fragment exchange"); for (Operation* op : erase) if (op && op->getBlock() && isa(op)) { if (!op->use_empty()) @@ -1238,6 +1086,12 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { rewriter.eraseOp(op); } + SmallVector blueprints; + funcOp.walk([&](SpatBlueprintOp blueprint) { blueprints.push_back(blueprint); }); + for (SpatBlueprintOp blueprint : blueprints) + if (failed(retargetBlueprintToScheduledPublications(plan, blueprint))) + return failure(); + for (Operation& op : funcOp.getOps()) { if (!isa(op)) continue; @@ -1257,8 +1111,14 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { && !llvm::is_contained(exact, produced->published)) exact.push_back(produced->published); if (exact.size() == 1) { - for (OpOperand* use : externalUses) + for (OpOperand* use : externalUses) { + Operation *consumer = use->getOwner(); + Operation *producer = exact.front().getDefiningOp(); + if (consumer->getBlock() == producer->getBlock() && + consumer->isBeforeInBlock(producer)) + consumer->moveAfter(producer); use->set(exact.front()); + } continue; } return op.emitOpError("phase 2 final publication ownership changed after planning"); @@ -1282,6 +1142,13 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { } rewriter.eraseOp(op); } + SmallVector deadBlueprints; + funcOp.walk([&](SpatBlueprintOp blueprint) { + if (blueprint.getOutput().use_empty()) + deadBlueprints.push_back(blueprint); + }); + for (SpatBlueprintOp blueprint : deadBlueprints) + rewriter.eraseOp(blueprint); return success(); } @@ -1289,9 +1156,12 @@ static LogicalResult realizePlan(func::FuncOp funcOp, RealizationPlan& plan) { LogicalResult realizeDeferredCommunication(func::FuncOp funcOp) { RealizationPlan plan; - if (failed(collectScheduled(funcOp, plan)) || failed(normalizeDeferred(funcOp, plan)) - || failed(buildAndVerifyPlan(funcOp, plan))) - return failure(); + if (failed(collectScheduled(funcOp, plan))) + return funcOp.emitOpError("phase 2 failed to collect scheduled producers"); + if (failed(normalizeDeferred(funcOp, plan))) + return funcOp.emitOpError("phase 2 failed to normalize deferred projections"); + if (failed(buildAndVerifyPlan(funcOp, plan))) + return funcOp.emitOpError("phase 2 failed to plan fragment communication"); if (failed(realizePlan(funcOp, plan))) return failure(); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp new file mode 100644 index 0000000..2eee59c --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp @@ -0,0 +1,287 @@ +#include "DeferredProjectionAnalysis.hpp" + +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/Matchers.h" +#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" + +#include + +namespace onnx_mlir::spatial { +using namespace mlir; +namespace { + +static FailureOr evaluate(Value value, const StaticIndexEnvironment &environment, + llvm::SmallDenseSet &visiting) { + if (auto it = environment.bindings.find(value); it != environment.bindings.end()) + return it->second; + if (!visiting.insert(value).second) + return failure(); + if (auto constant = value.getDefiningOp()) + if (auto integer = dyn_cast(constant.getValue())) + { visiting.erase(value); return integer.getInt(); } + if (auto cast = value.getDefiningOp()) { + auto result = evaluate(cast.getIn(), environment, visiting); visiting.erase(value); return result; + } + auto binary = [&](auto op, auto fn) -> FailureOr { + auto lhs = evaluate(op.getLhs(), environment, visiting); + auto rhs = evaluate(op.getRhs(), environment, visiting); + if (failed(lhs) || failed(rhs)) return failure(); + return fn(*lhs, *rhs); + }; + if (auto op = value.getDefiningOp()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr { int64_t r; if (llvm::AddOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; } + if (auto op = value.getDefiningOp()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr { int64_t r; if (llvm::SubOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; } + if (auto op = value.getDefiningOp()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr { int64_t r; if (llvm::MulOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; } + if (auto apply = value.getDefiningOp()) { auto result = evaluateAffineApply(apply, [&](Value operand) { return evaluate(operand, environment, visiting); }); visiting.erase(value); return result; } + if (auto extract = value.getDefiningOp()) { + auto constant = extract.getTensor().getDefiningOp(); + auto elements = constant ? dyn_cast(constant.getValue()) : DenseIntElementsAttr(); + auto type = elements ? dyn_cast(elements.getType()) : RankedTensorType(); + if (!elements || !type || extract.getIndices().size() != static_cast(type.getRank())) return failure(); + int64_t linear = 0; + for (auto [index, dim] : llvm::zip(extract.getIndices(), type.getShape())) { + auto i = evaluate(index, environment, visiting); + if (failed(i) || *i < 0 || *i >= dim) return failure(); + linear = linear * dim + *i; + } + visiting.erase(value); return elements.getValues()[linear].getSExtValue(); + } + visiting.erase(value); + return failure(); +} + +static FailureOr> sourceArgument(Value value, SpatDeferredCommunicationOp deferred, + const StaticIndexEnvironment &environment) { + while (auto cast = value.getDefiningOp()) value = cast.getSource(); + if (auto argument = dyn_cast(value); + argument && argument.getOwner() == &deferred.getBody().front() + && argument.getArgNumber() < deferred.getSources().size()) + return std::optional(argument.getArgNumber()); + // Phase 1's selector ends in collapse(extract_slice(stacked, table[lane])). + auto collapse = value.getDefiningOp(); + if (!collapse) return std::optional(); + value = collapse.getSrc(); + auto slice = value.getDefiningOp(); + if (!slice) return std::optional(); + auto sourceType = dyn_cast(slice.getSourceType()); + if (!sourceType || slice.getMixedOffsets().size() != static_cast(sourceType.getRank())) + return std::optional(); + for (unsigned dim = 1; dim < sourceType.getRank(); ++dim) { + auto offset = evaluateDeferredIndex(slice.getMixedOffsets()[dim], environment); + auto size = evaluateDeferredIndex(slice.getMixedSizes()[dim], environment); + auto stride = evaluateDeferredIndex(slice.getMixedStrides()[dim], environment); + if (failed(offset) || failed(size) || failed(stride) || *offset != 0 || *size != sourceType.getDimSize(dim) || *stride != 1) + return std::optional(); + } + auto leadingSize = evaluateDeferredIndex(slice.getMixedSizes().front(), environment); + auto leadingStride = evaluateDeferredIndex(slice.getMixedStrides().front(), environment); + if (failed(leadingSize) || failed(leadingStride) || *leadingSize != 1 || *leadingStride != 1) + return std::optional(); + auto selected = evaluateDeferredIndex(slice.getMixedOffsets().front(), environment); + if (failed(selected)) return failure(); + Value stacked = slice.getSource(); + while (auto cast = stacked.getDefiningOp()) stacked = cast.getSource(); + while (auto insert = stacked.getDefiningOp()) stacked = insert.getDest(); + // The stack is a chain. Find the insertion at the selected leading offset. + for (Value cursor = slice.getSource(); auto insert = cursor.getDefiningOp(); cursor = insert.getDest()) { + auto offset = evaluateDeferredIndex(insert.getMixedOffsets().front(), environment); + if (succeeded(offset) && *offset == *selected) { + Value source = insert.getSource(); + if (auto expand = source.getDefiningOp()) source = expand.getSrc(); + if (auto arg = dyn_cast(source); arg && arg.getOwner() == &deferred.getBody().front()) + return std::optional(arg.getArgNumber()); + } + } + return std::optional(); +} + +static bool isResidual(Operation *op) { + return isa(op); +} + +static SpatGraphComputeBatch graphBatchOwner(Value value) { + if (auto result = dyn_cast(value)) + return dyn_cast(result.getOwner()); + return {}; +} + +static Value getEnclosingScheduledLane(SpatDeferredCommunicationOp deferred, + SpatScheduledComputeBatch scheduled) { + Block *block = deferred->getBlock(); + while (block && block->getParentOp() != scheduled) { + Operation *parent = block->getParentOp(); + block = parent ? parent->getBlock() : nullptr; + } + return block && !block->empty() ? block->getArgument(0) : Value(); +} + +} // namespace + +FailureOr evaluateDeferredIndex(Value value, const StaticIndexEnvironment &environment) { + llvm::SmallDenseSet visiting; + return evaluate(value, environment, visiting); +} +FailureOr evaluateDeferredIndex(OpFoldResult value, const StaticIndexEnvironment &environment) { + if (auto attr = dyn_cast(value)) + if (auto integer = dyn_cast(attr)) return integer.getInt(); + if (auto dynamic = dyn_cast(value)) return evaluateDeferredIndex(dynamic, environment); + return failure(); +} + +FailureOr> tryResolveDeferredSource(Value value, SpatDeferredCommunicationOp deferred, + const StaticIndexEnvironment &environment) { + auto index = sourceArgument(value, deferred, environment); + if (failed(index)) + return failure(); + if (!*index) + return std::optional(); + return std::optional(ResolvedDeferredSource {**index, deferred.getSources()[**index]}); +} + +FailureOr requireResolvedDeferredSource(Value value, SpatDeferredCommunicationOp deferred, + const StaticIndexEnvironment &environment) { + auto source = tryResolveDeferredSource(value, deferred, environment); + if (failed(source) || !*source) + return deferred.emitOpError("cannot statically resolve deferred source selection"), failure(); + return **source; +} + +FailureOr analyzeDeferredProgram(SpatDeferredCommunicationOp deferred, + std::optional targetScheduledLane) { + Block &body = deferred.getBody().front(); + auto yield = dyn_cast(body.getTerminator()); + if (!yield || yield.getOutputs().size() != 1) + return deferred.emitOpError("requires exactly one deferred yielded value"), failure(); + StaticIndexEnvironment environment; + if (auto scheduled = deferred->getParentOfType()) { + if (!targetScheduledLane) return deferred.emitOpError("scheduled-batch deferred program requires lane specialization"), failure(); + Value scheduledLane = getEnclosingScheduledLane(deferred, scheduled); + if (!scheduledLane) + return deferred.emitOpError("cannot locate the enclosing scheduled lane"), failure(); + environment.bindings[scheduledLane] = *targetScheduledLane; + } else if (targetScheduledLane) return deferred.emitOpError("scalar deferred program cannot have a target lane"), failure(); + SpecializedDeferredProgram program; + program.deferred = deferred; + program.targetScheduledLane = targetScheduledLane; + if (auto scheduled = deferred->getParentOfType()) + program.scheduledLane = getEnclosingScheduledLane(deferred, scheduled); + program.yieldedValue = yield.getOutputs().front(); + llvm::SmallDenseSet visited; + std::function visit = [&](Value value) -> LogicalResult { + if (!visited.insert(value).second) return success(); + // A graph-batch projection is semantically different from the canonical + // source-selector scaffold even though both contain extract/collapse ops. + if (auto slice = value.getDefiningOp()) { + auto source = tryResolveDeferredSource(slice.getSource(), deferred, environment); + if (failed(source)) return failure(); + if (*source && graphBatchOwner((*source)->selectedValue)) { + auto type = dyn_cast((*source)->selectedValue.getType()); + auto result = dyn_cast(slice.getResult().getType()); + if (!type || !result || type.getRank() == 0) return deferred.emitOpError("graph projection requires ranked tensors"); + auto offset = evaluateDeferredIndex(slice.getMixedOffsets().front(), environment); + if (failed(offset)) return deferred.emitOpError("graph projection leading offset is not statically resolvable"); + auto size = evaluateDeferredIndex(slice.getMixedSizes().front(), environment); + if (failed(size)) return deferred.emitOpError("graph projection leading size is not statically resolvable"); + auto stride = evaluateDeferredIndex(slice.getMixedStrides().front(), environment); + if (failed(stride)) return deferred.emitOpError("graph projection leading stride is not statically resolvable"); + if (*offset < 0) return deferred.emitOpError("graph projection leading offset is negative"); + if (*size <= 0) return deferred.emitOpError("graph projection leading size must be positive"); + if (*stride <= 0) return deferred.emitOpError("graph projection leading stride must be positive"); + DeferredProjectionLeaf leaf; + leaf.kind = DeferredLeafKind::GraphBatchProjection; leaf.sourceOperandIndex = (*source)->sourceOperandIndex; + leaf.replacementRoot = value; leaf.leadingProjection = slice; leaf.reconstructedType = result; + for (int64_t i = 0; i < *size; ++i) { + int64_t slot; + if (llvm::MulOverflow(i, *stride, slot) || llvm::AddOverflow(*offset, slot, slot) + || slot >= type.getDimSize(0)) + return deferred.emitOpError("graph projection selects a physical slot outside the batch"); + leaf.physicalSlots.push_back(slot); + } + for (unsigned i = 1; i < type.getRank(); ++i) { + auto innerOffset = evaluateDeferredIndex(slice.getMixedOffsets()[i], environment); + auto innerSize = evaluateDeferredIndex(slice.getMixedSizes()[i], environment); + auto innerStride = evaluateDeferredIndex(slice.getMixedStrides()[i], environment); + if (failed(innerOffset) || failed(innerSize) || failed(innerStride)) + return deferred.emitOpError("graph projection has unresolved inner geometry"); + leaf.innerGeometry.offsets.push_back(*innerOffset); leaf.innerGeometry.sizes.push_back(*innerSize); leaf.innerGeometry.strides.push_back(*innerStride); + } + program.leaves.push_back(std::move(leaf)); return success(); + } + } + + auto source = tryResolveDeferredSource(value, deferred, environment); + if (failed(source)) return failure(); + if (*source) { + auto type = dyn_cast((*source)->selectedValue.getType()); + if (!type) return deferred.emitOpError("deferred source is not a ranked tensor"); + DeferredProjectionLeaf leaf; + leaf.sourceOperandIndex = (*source)->sourceOperandIndex; + leaf.replacementRoot = value; + leaf.reconstructedType = type; + if (graphBatchOwner((*source)->selectedValue)) { + leaf.kind = DeferredLeafKind::GraphBatchIdentity; + for (int64_t slot = 0; slot < type.getDimSize(0); ++slot) leaf.physicalSlots.push_back(slot); + } else leaf.kind = DeferredLeafKind::ScalarSource; + program.leaves.push_back(std::move(leaf)); + return success(); + } + Operation *op = value.getDefiningOp(); + if (!op || op->getBlock() != &body || !isResidual(op)) + return deferred.emitOpError("deferred residual contains an unsupported operation"); + for (Value operand : op->getOperands()) if (failed(visit(operand))) return failure(); + program.residualOps.push_back(op); return success(); + }; + if (failed(visit(program.yieldedValue))) return failure(); + return std::move(program); +} + +FailureOr getGraphBatchPublicationMap( + SpatGraphComputeBatch graphBatch, unsigned resultIndex, GraphBatchPublicationCache &cache) { + GraphBatchPublicationKey key {graphBatch.getOperation(), resultIndex}; + if (auto it = cache.find(key); it != cache.end()) return &it->second; + auto resultType = dyn_cast(graphBatch.getResult(resultIndex).getType()); + auto output = graphBatch.getOutputArgument(resultIndex); + auto lane = graphBatch.getLaneArgument(); + if (!resultType || !output || !lane || resultType.getRank() == 0) + return graphBatch.emitOpError("graph batch publication is malformed"), failure(); + tensor::ParallelInsertSliceOp publication; + auto parallel = dyn_cast(graphBatch.getBody().front().getTerminator()); + if (!parallel) return graphBatch.emitOpError("graph batch lacks publication region"), failure(); + for (Operation &op : parallel.getRegion().front()) if (auto insert = dyn_cast(op); insert && insert.getDest() == *output) { + if (publication) return graphBatch.emitOpError("graph result has multiple publications"), failure(); + publication = insert; + } + if (!publication) return graphBatch.emitOpError("graph result lacks publication"), failure(); + auto fragment = dyn_cast(publication.getSource().getType()); + if (!fragment || resultType.getRank() != fragment.getRank() + 1) return graphBatch.emitOpError("graph publication fragment type is invalid"), failure(); + GraphBatchPublicationMap map; + map.physicalResultType = resultType; + map.publicationFragmentType = fragment; + map.graphLaneToPhysicalSlot.resize(graphBatch.getLaneCount(), -1); + map.physicalSlotToGraphLane.resize(resultType.getDimSize(0), -1); + for (int64_t index = 0; index < graphBatch.getLaneCount(); ++index) { + StaticIndexEnvironment environment; environment.bindings[*lane] = index; + auto slot = evaluateDeferredIndex(publication.getMixedOffsets().front(), environment); + auto size = evaluateDeferredIndex(publication.getMixedSizes().front(), environment); + auto stride = evaluateDeferredIndex(publication.getMixedStrides().front(), environment); + if (failed(slot) || failed(size) || failed(stride) || *size != 1 || *stride != 1 || *slot < 0 || *slot >= resultType.getDimSize(0)) + return graphBatch.emitOpError("graph publication leading geometry is invalid"), failure(); + for (unsigned dim = 1; dim < resultType.getRank(); ++dim) { + auto offset = evaluateDeferredIndex(publication.getMixedOffsets()[dim], environment); + auto extent = evaluateDeferredIndex(publication.getMixedSizes()[dim], environment); + auto step = evaluateDeferredIndex(publication.getMixedStrides()[dim], environment); + if (failed(offset) || failed(extent) || failed(step) || *offset != 0 || *extent != fragment.getDimSize(dim - 1) || *step != 1) + return graphBatch.emitOpError("graph publication inner geometry is invalid"), failure(); + } + if (map.physicalSlotToGraphLane[*slot] != -1) return graphBatch.emitOpError("graph publication has duplicate physical slot"), failure(); + map.graphLaneToPhysicalSlot[index] = *slot; map.physicalSlotToGraphLane[*slot] = index; + } + if (llvm::is_contained(map.physicalSlotToGraphLane, -1)) return graphBatch.emitOpError("graph publication has missing physical slot"), failure(); + return &cache.try_emplace(key, std::move(map)).first->second; +} + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp new file mode 100644 index 0000000..2c9e4fd --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp @@ -0,0 +1,102 @@ +#pragma once + +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Operation.h" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +#include "llvm/ADT/DenseMap.h" + +namespace onnx_mlir::spatial { + +struct StaticIndexEnvironment { + llvm::DenseMap bindings; +}; + +mlir::FailureOr evaluateDeferredIndex( + mlir::Value value, const StaticIndexEnvironment &environment); +mlir::FailureOr evaluateDeferredIndex( + mlir::OpFoldResult value, const StaticIndexEnvironment &environment); + +enum class DeferredLeafKind { ScalarSource, GraphBatchProjection, GraphBatchIdentity }; + +struct StaticSliceGeometry { + llvm::SmallVector offsets; + llvm::SmallVector sizes; + llvm::SmallVector strides; +}; + +struct DeferredProjectionLeaf { + DeferredLeafKind kind = DeferredLeafKind::ScalarSource; + unsigned sourceOperandIndex = 0; + mlir::Value replacementRoot; + mlir::tensor::ExtractSliceOp leadingProjection; + llvm::SmallVector physicalSlots; + StaticSliceGeometry innerGeometry; + mlir::RankedTensorType reconstructedType; +}; + +struct SpecializedDeferredProgram { + SpatDeferredCommunicationOp deferred; + std::optional targetScheduledLane; + mlir::Value scheduledLane; + mlir::Value yieldedValue; + llvm::SmallVector leaves; + llvm::SmallVector residualOps; +}; + +struct ResolvedDeferredSource { + unsigned sourceOperandIndex = 0; + mlir::Value selectedValue; +}; + +mlir::FailureOr> tryResolveDeferredSource( + mlir::Value value, SpatDeferredCommunicationOp deferred, + const StaticIndexEnvironment &environment); +mlir::FailureOr requireResolvedDeferredSource( + mlir::Value value, SpatDeferredCommunicationOp deferred, + const StaticIndexEnvironment &environment); +mlir::FailureOr analyzeDeferredProgram( + SpatDeferredCommunicationOp deferred, + std::optional targetScheduledLane); + +struct GraphBatchPublicationMap { + mlir::RankedTensorType physicalResultType; + mlir::RankedTensorType publicationFragmentType; + llvm::SmallVector graphLaneToPhysicalSlot; + llvm::SmallVector physicalSlotToGraphLane; +}; + +struct GraphBatchPublicationKey { + mlir::Operation *graphBatch = nullptr; + unsigned resultIndex = 0; + bool operator==(const GraphBatchPublicationKey &other) const { + return graphBatch == other.graphBatch && resultIndex == other.resultIndex; + } +}; + +using GraphBatchPublicationCache = + llvm::DenseMap; + +mlir::FailureOr getGraphBatchPublicationMap( + SpatGraphComputeBatch graphBatch, unsigned resultIndex, + GraphBatchPublicationCache &cache); + +} // namespace onnx_mlir::spatial + +namespace llvm { +template <> struct DenseMapInfo { + static onnx_mlir::spatial::GraphBatchPublicationKey getEmptyKey() { + return {DenseMapInfo::getEmptyKey(), 0}; + } + static onnx_mlir::spatial::GraphBatchPublicationKey getTombstoneKey() { + return {DenseMapInfo::getTombstoneKey(), 0}; + } + static unsigned getHashValue(const onnx_mlir::spatial::GraphBatchPublicationKey &key) { + return hash_combine(key.graphBatch, key.resultIndex); + } + static bool isEqual(const onnx_mlir::spatial::GraphBatchPublicationKey &lhs, + const onnx_mlir::spatial::GraphBatchPublicationKey &rhs) { + return lhs == rhs; + } +}; +} // namespace llvm diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp index c725c05..ee6090d 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp @@ -9,6 +9,7 @@ #include "Scheduling/MergeSchedulingAnalysis.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" #include "src/Accelerators/PIM/Pass/PIMPasses.h" using namespace mlir; @@ -44,6 +45,10 @@ struct MergeComputeNodesPass final : PassWrapperpeftClassPlans, diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp index fc856a9..55a019f 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp @@ -299,7 +299,8 @@ static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClass for (Value weight : getComputeInstanceWeights(instance)) appendUnique(peftClassPlan.weights, weight); for (Value input : getComputeInstanceInputs(instance)) - if (!getProducerValueRef(input, &instance)) + if (!getProducerValueRef(input, &instance) && + !isDeferredFragmentAssemblyInput(input, instance)) appendUnique(peftClassPlan.inputs, input); } return success(); @@ -333,7 +334,8 @@ static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClass for (Value weight : getComputeInstanceWeights(instance)) appendUnique(peftClassPlan.weights, weight); for (Value input : getComputeInstanceInputs(instance)) - if (!getProducerValueRef(input, &instance)) + if (!getProducerValueRef(input, &instance) && + !isDeferredFragmentAssemblyInput(input, instance)) appendUnique(peftClassPlan.inputs, input); } } @@ -387,6 +389,11 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew IRMapping mapper; mapper.map(*batch.getLaneArgument(), originalLane); + Value localLane = arith::SubIOp::create(builder, + bodyLoc, + originalLane, + getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart)) + .getResult(); for (auto [index, weight] : llvm::enumerate(batch.getWeights())) mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight)); SmallVector inputPlans; @@ -403,9 +410,13 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew scheduledWeights.size(), ArrayRef {}, *batch.getLaneArgument(), - lower, + originalLane, plan))) return failure(); + plan.scalarizedLocalLane = localLane; + plan.scalarizedGraphLaneBase = lower; + plan.scalarizedLaneCount = instance.laneCount; + plan.scalarizedHoistBlock = █ inputPlans.push_back(std::move(plan)); } for (auto [index, outputArg] : llvm::enumerate(batch.getOutputs())) @@ -422,11 +433,6 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew auto inParallel = dyn_cast(source.getTerminator()); if (!inParallel) return batch.emitOpError("expected spat.in_parallel in resultful spat.graph_compute_batch"), failure(); - Value localLane = arith::SubIOp::create(builder, - bodyLoc, - originalLane, - getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart)) - .getResult(); DenseMap outputIndexByArg; for (size_t index = 0; index < batch.getNumResults(); ++index) outputIndexByArg[*batch.getOutputArgument(index)] = index; @@ -624,7 +630,8 @@ static LogicalResult materializeMultiCpuPeftClass( if (auto compute = dyn_cast(representative.op)) { IRMapping mapper; for (auto [index, weight] : llvm::enumerate(compute.getWeights())) - mapper.map(*compute.getWeightArgument(index), block->getArgument(1 + index)); + mapper.map(*compute.getWeightArgument(index), + getBlockOperand(*block, scheduled.getWeights(), weight, 1)); unsigned firstInputArg = 1 + scheduled.getWeights().size(); SmallVector inputPlans; for (auto [index, input] : llvm::enumerate(compute.getInputs())) { @@ -686,6 +693,10 @@ static LogicalResult materializeMultiCpuPeftClass( Value lower = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 0); Value upper = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), representative.laneCount); Value step = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 1); + FailureOr sourceLaneStart = + buildSourceLaneStartForScheduledLane(rewriter, batch.getLoc(), scheduledLane, sourceLaneSelector, scheduled.getOperation()); + if (failed(sourceLaneStart)) + return failure(); auto loop = buildNormalizedScfFor( rewriter, batch.getLoc(), @@ -696,10 +707,6 @@ static LogicalResult materializeMultiCpuPeftClass( [&](OpBuilder &builder, Location bodyLoc, Value innerLane, ValueRange iterArgs, SmallVectorImpl &yielded) -> LogicalResult { IRMapping mapper; - FailureOr sourceLaneStart = - buildSourceLaneStartForScheduledLane(builder, bodyLoc, scheduledLane, sourceLaneSelector, scheduled.getOperation()); - if (failed(sourceLaneStart)) - return failure(); Value sourceLane = affine::AffineApplyOp::create(builder, bodyLoc, @@ -710,7 +717,8 @@ static LogicalResult materializeMultiCpuPeftClass( .getResult(); mapper.map(*batch.getLaneArgument(), sourceLane); for (auto [index, weight] : llvm::enumerate(batch.getWeights())) - mapper.map(*batch.getWeightArgument(index), block->getArgument(1 + index)); + mapper.map(*batch.getWeightArgument(index), + getBlockOperand(*block, scheduled.getWeights(), weight, 1)); unsigned firstInputArg = 1 + scheduled.getWeights().size(); SmallVector inputPlans; for (auto [index, input] : llvm::enumerate(batch.getInputs())) { @@ -730,6 +738,10 @@ static LogicalResult materializeMultiCpuPeftClass( scheduledLane, plan))) return failure(); + plan.scalarizedLocalLane = innerLane; + plan.scalarizedGraphLaneBase = *sourceLaneStart; + plan.scalarizedLaneCount = representative.laneCount; + plan.scalarizedHoistBlock = block; inputPlans.push_back(std::move(plan)); } for (unsigned index = 0; index < batch.getNumResults(); ++index) @@ -838,7 +850,10 @@ materializeScheduledCompute(func::FuncOp funcOp, llvm::sort(peftClassPlan.cpus); for (size_t cpu : peftClassPlan.cpus) llvm::sort(peftClassPlan.instancesByCpu[cpu], [&](const ComputeInstance &lhs, const ComputeInstance &rhs) { - return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs); + return std::tie(graphIds.find(lhs.op)->second, + schedule.computeToCpuSlotMap.find(lhs)->second) < + std::tie(graphIds.find(rhs.op)->second, + schedule.computeToCpuSlotMap.find(rhs)->second); }); if (failed(verifyPeftClassPlan(funcOp.getOperation(), peftClassPlan, schedule))) return failure(); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp index f7ba008..1ca2b30 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp @@ -64,8 +64,6 @@ struct PeftMaterializationReportSummary { size_t scheduledCompute = 0; size_t scheduledComputeBatch = 0; size_t deferredCommunication = 0; - size_t deferredCommunicationScalarMetadata = 0; - size_t deferredCommunicationScheduledLaneMetadata = 0; size_t deferredCommunicationMultiSourcePayloads = 0; }; @@ -95,10 +93,6 @@ static PeftMaterializationReportSummary buildPeftMaterializationReportSummary( } funcOp.walk([&](SpatDeferredCommunicationOp transfer) { summary.deferredCommunication++; - if (auto batchedAttr = transfer->getAttrOfType("batched"); batchedAttr && batchedAttr.getValue()) - summary.deferredCommunicationScheduledLaneMetadata++; - else - summary.deferredCommunicationScalarMetadata++; if (transfer.getSources().size() > 1) summary.deferredCommunicationMultiSourcePayloads++; }); @@ -200,8 +194,6 @@ static void dumpPeftMaterializationReport(ModuleOp moduleOp, os << " scheduled_compute_batch: " << summary.scheduledComputeBatch << "\n"; os << "Deferred communications:\n"; os << " total: " << summary.deferredCommunication << "\n"; - os << " scalar metadata: " << summary.deferredCommunicationScalarMetadata << "\n"; - os << " scheduled-lane metadata: " << summary.deferredCommunicationScheduledLaneMetadata << "\n"; os << " multi-source payloads: " << summary.deferredCommunicationMultiSourcePayloads << "\n\n"; os << "PEFT Classes\n"; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp index 097711b..fe5f20f 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp @@ -1,4 +1,5 @@ #include "ScheduledComputeVerification.hpp" +#include "DeferredProjectionAnalysis.hpp" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" @@ -79,6 +80,7 @@ LogicalResult verifyMaterializedScheduleMapping( LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) { pim::CappedDiagnosticReporter diagnostics; + GraphBatchPublicationCache publicationCache; funcOp.walk([&](SpatDeferredCommunicationOp transfer) { for (Value source : transfer.getSources()) { auto result = dyn_cast(source); @@ -92,7 +94,24 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) { !transfer->getParentOfType()) diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { illegalOp->emitOpError("phase-check deferred communication must be inside a scheduled compute"); - }); + }); + if (auto scheduled = transfer->getParentOfType()) { + for (unsigned lane = 0; lane < static_cast(scheduled.getLaneCount()); ++lane) { + auto program = analyzeDeferredProgram(transfer, lane); + if (failed(program)) + continue; + for (const DeferredProjectionLeaf &leaf : program->leaves) { + if (leaf.kind == DeferredLeafKind::ScalarSource) + continue; + auto source = dyn_cast(transfer.getSources()[leaf.sourceOperandIndex]); + auto graph = source ? dyn_cast(source.getOwner()) : SpatGraphComputeBatch(); + if (!graph) continue; + if (failed(getGraphBatchPublicationMap(graph, source.getResultNumber(), publicationCache))) + diagnostics.report(transfer.getOperation(), [&](Operation *) {}); + } + } + } else + (void)analyzeDeferredProgram(transfer, std::nullopt); }); diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed");