diff --git a/src/PIM/Common/Support/DebugDump.cpp b/src/PIM/Common/Support/DebugDump.cpp index 632a3a9..05b5f05 100644 --- a/src/PIM/Common/Support/DebugDump.cpp +++ b/src/PIM/Common/Support/DebugDump.cpp @@ -24,7 +24,7 @@ void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVer llvm::raw_os_ostream os(file); mlir::OpPrintingFlags flags; - flags.elideLargeElementsAttrs().enableDebugInfo(true, false); + flags.elideLargeElementsAttrs().enableDebugInfo(false, false); if (assumeVerified) flags.assumeVerified(); moduleOp.print(os, flags); diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp index 78f7328..da4f5b0 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp @@ -60,6 +60,56 @@ struct SpatComputeBatchBodyArgs { mlir::ValueRange outputs; }; +inline mlir::SmallVector getGraphComputeBlockArgTypes(mlir::ValueRange weights, mlir::ValueRange inputs) { + mlir::SmallVector blockArgTypes; + blockArgTypes.reserve(weights.size() + inputs.size()); + for (mlir::Value weight : weights) + blockArgTypes.push_back(weight.getType()); + for (mlir::Value input : inputs) + blockArgTypes.push_back(input.getType()); + return blockArgTypes; +} + +inline mlir::SmallVector getGraphComputeBlockArgLocs(mlir::Location defaultLoc, + mlir::ValueRange weights, + mlir::ValueRange inputs) { + mlir::SmallVector blockArgLocs; + blockArgLocs.reserve(weights.size() + inputs.size()); + for (mlir::Value weight : weights) + blockArgLocs.push_back(weight.getLoc()); + for (mlir::Value input : inputs) + blockArgLocs.push_back(input.getLoc()); + return blockArgLocs; +} + +inline mlir::SmallVector getGraphComputeBatchBlockArgTypes(mlir::OpBuilder& builder, + mlir::TypeRange resultTypes, + mlir::ValueRange weights, + mlir::ValueRange inputs) { + mlir::SmallVector blockArgTypes {builder.getIndexType()}; + blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size()); + for (mlir::Value weight : weights) + blockArgTypes.push_back(weight.getType()); + for (mlir::Value input : inputs) + blockArgTypes.push_back(input.getType()); + llvm::append_range(blockArgTypes, resultTypes); + return blockArgTypes; +} + +inline mlir::SmallVector getGraphComputeBatchBlockArgLocs(mlir::Location defaultLoc, + mlir::TypeRange resultTypes, + mlir::ValueRange weights, + mlir::ValueRange inputs) { + mlir::SmallVector blockArgLocs {defaultLoc}; + blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size()); + for (mlir::Value weight : weights) + blockArgLocs.push_back(weight.getLoc()); + for (mlir::Value input : inputs) + blockArgLocs.push_back(input.getLoc()); + blockArgLocs.append(resultTypes.size(), defaultLoc); + return blockArgLocs; +} + } // namespace detail template @@ -87,6 +137,31 @@ inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput(); } +template +spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter, + mlir::Location loc, + mlir::TypeRange resultTypes, + mlir::ValueRange weights, + mlir::ValueRange inputs, + mlir::TypeRange blockArgTypes, + llvm::ArrayRef blockArgLocs) { + auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs); + rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs); + rewriter.setInsertionPointToStart(&computeOp.getBody().front()); + return computeOp; +} + +template +spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter, + mlir::Location loc, + mlir::TypeRange resultTypes, + mlir::ValueRange weights, + mlir::ValueRange inputs) { + auto blockArgTypes = detail::getGraphComputeBlockArgTypes(weights, inputs); + auto blockArgLocs = detail::getGraphComputeBlockArgLocs(loc, weights, inputs); + return createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, blockArgTypes, blockArgLocs); +} + /// Builds a `spat.graph_compute` with a fixed number of SSA inputs and erases it if /// the body callback reports failure. template @@ -97,16 +172,8 @@ auto createSpatGraphCompute(RewriterT& rewriter, mlir::ValueRange inputs, BodyFn&& body) { assert(inputs.size() == NumInputs && "NumInputs must match the number of input values"); - auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs); - - auto* block = new mlir::Block(); - for (mlir::Value weight : weights) - block->addArgument(weight.getType(), loc); - for (mlir::Value input : inputs) - block->addArgument(input.getType(), loc); - - computeOp.getBody().push_back(block); - rewriter.setInsertionPointToStart(block); + auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs); + auto* block = &computeOp.getBody().front(); using BodyResult = detail::InvokeWithBlockArgsResultT, std::make_index_sequence>; if constexpr (std::is_same_v) { @@ -140,16 +207,8 @@ auto createSpatGraphCompute(RewriterT& rewriter, mlir::ValueRange weights, mlir::ValueRange inputs, BodyFn&& body) { - auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs); - - auto* block = new mlir::Block(); - for (mlir::Value weight : weights) - block->addArgument(weight.getType(), loc); - for (mlir::Value input : inputs) - block->addArgument(input.getType(), loc); - - computeOp.getBody().push_back(block); - rewriter.setInsertionPointToStart(block); + auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs); + auto* block = &computeOp.getBody().front(); using BodyResult = detail::InvokeWithValueRangeResultT>; if constexpr (std::is_same_v) { @@ -170,14 +229,15 @@ auto createSpatGraphCompute(RewriterT& rewriter, } } -template -auto createSpatGraphComputeBatch(RewriterT& rewriter, - mlir::Location loc, - mlir::TypeRange resultTypes, - int64_t laneCount, - mlir::ValueRange weights, - mlir::ValueRange inputs, - BodyFn&& body) { +template +auto createEmptySpatGraphComputeBatch(RewriterT& rewriter, + mlir::Location loc, + mlir::TypeRange resultTypes, + int64_t laneCount, + mlir::ValueRange weights, + mlir::ValueRange inputs, + mlir::TypeRange blockArgTypes, + llvm::ArrayRef blockArgLocs) { if (laneCount <= 0 || laneCount > std::numeric_limits::max()) return mlir::FailureOr(mlir::failure()); @@ -186,27 +246,36 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter, return mlir::FailureOr(mlir::failure()); auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs); + rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs); + rewriter.setInsertionPointToStart(&batchOp.getBody().front()); + return mlir::FailureOr(batchOp); +} - mlir::SmallVector blockArgTypes {rewriter.getIndexType()}; - mlir::SmallVector blockArgLocs {loc}; - blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size()); - blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size()); - for (mlir::Value weight : weights) { - blockArgTypes.push_back(weight.getType()); - blockArgLocs.push_back(weight.getLoc()); - } - for (mlir::Value input : inputs) { - blockArgTypes.push_back(input.getType()); - blockArgLocs.push_back(input.getLoc()); - } - for (mlir::Type resultType : resultTypes) { - blockArgTypes.push_back(resultType); - blockArgLocs.push_back(loc); - } +template +auto createEmptySpatGraphComputeBatch(RewriterT& rewriter, + mlir::Location loc, + mlir::TypeRange resultTypes, + int64_t laneCount, + mlir::ValueRange weights, + mlir::ValueRange inputs) { + auto blockArgTypes = detail::getGraphComputeBatchBlockArgTypes(rewriter, resultTypes, weights, inputs); + auto blockArgLocs = detail::getGraphComputeBatchBlockArgLocs(loc, resultTypes, weights, inputs); + return createEmptySpatGraphComputeBatch( + rewriter, loc, resultTypes, laneCount, weights, inputs, blockArgTypes, blockArgLocs); +} - auto* block = - rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs); - rewriter.setInsertionPointToStart(block); +template +auto createSpatGraphComputeBatch(RewriterT& rewriter, + mlir::Location loc, + mlir::TypeRange resultTypes, + int64_t laneCount, + mlir::ValueRange weights, + mlir::ValueRange inputs, + BodyFn&& body) { + auto batchOp = createEmptySpatGraphComputeBatch(rewriter, loc, resultTypes, laneCount, weights, inputs); + if (failed(batchOp)) + return mlir::FailureOr(mlir::failure()); + auto* block = &(*batchOp).getBody().front(); detail::SpatComputeBatchBodyArgs args { block->getArgument(0), @@ -217,18 +286,18 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter, using BodyResult = std::invoke_result_t; if constexpr (std::is_same_v) { std::forward(body)(args); - rewriter.setInsertionPointAfter(batchOp); - return mlir::FailureOr(batchOp); + rewriter.setInsertionPointAfter(*batchOp); + return batchOp; } else { auto bodyResult = std::forward(body)(args); if (mlir::failed(bodyResult)) { - rewriter.setInsertionPointAfter(batchOp); - rewriter.eraseOp(batchOp); + rewriter.setInsertionPointAfter(*batchOp); + rewriter.eraseOp(*batchOp); return mlir::FailureOr(mlir::failure()); } - rewriter.setInsertionPointAfter(batchOp); - return mlir::FailureOr(batchOp); + rewriter.setInsertionPointAfter(*batchOp); + return batchOp; } } diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp index e411a9d..2eb1cdb 100644 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -11,6 +11,7 @@ #include "llvm/ADT/SmallPtrSet.h" #include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "mlir/Transforms/Passes.h" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" @@ -375,6 +376,12 @@ struct LowerSpatialPlansPass final : PassWrappergetArguments(), newCompute.getOperands())) mapper.map(computeArg, blockArg); newCompute.getProperties().setOperandSegmentSizes({0, static_cast(sourceTypes.size())}); @@ -175,11 +176,6 @@ void ONNXToSpatialPass::runOnOperation() { arith::ArithDialect, scf::SCFDialect>(); - PassManager cleanupPM(ctx); - cleanupPM.addPass(createCanonicalizerPass()); - if (failed(cleanupPM.run(moduleOp))) - moduleOp.emitWarning("failed to run ONNX-to-Spatial canonicalization cleanup; continuing"); - annotateWeightsConstants(*entryFunc); if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) { @@ -215,13 +211,18 @@ void ONNXToSpatialPass::runOnOperation() { populateEmptyFunction(*entryFunc); - if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) { + PassManager canonicalizationPM(ctx); + canonicalizationPM.addPass(createCanonicalizerPass()); + if (failed(canonicalizationPM.run(moduleOp))) + moduleOp.emitWarning("failed to run ONNXToSpatial canonicalization; continuing"); + + dumpModule(moduleOp, "spatial0"); + + if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) { moduleOp.emitError("logical Spatial graph verification failed after ONNX-to-Spatial"); signalPassFailure(); return; } - dumpModule(moduleOp, "spatial0"); - if (failed(verifyONNXToSpatial(*entryFunc))) { moduleOp.emitError("ONNX-to-Spatial host legality verification failed"); signalPassFailure(); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp index b120b9e..56e2090 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp @@ -917,9 +917,7 @@ struct MatMulToGemm : OpRewritePattern { if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector) return failure(); - const bool hasNonSingletonOutputBatch = - !shapeInfo->outputBatchShape.empty() && getStaticShapeElementCount(shapeInfo->outputBatchShape) != 1; - if (hasNonSingletonOutputBatch) + if (!shapeInfo->outputBatchShape.empty()) return failure(); Location loc = matmulOp.getLoc(); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Post.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Post.cpp index a66b456..49630ba 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Post.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Post.cpp @@ -10,6 +10,7 @@ #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp" #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" @@ -128,8 +129,6 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePatternnewWeights, promoted->newInputs); SmallVector newBlockArgTypes; SmallVector newBlockArgLocs; for (Value weight : promoted->newWeights) { @@ -138,10 +137,14 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePatternnewInputTypes); llvm::append_range(newBlockArgLocs, promoted->newInputLocs); - auto* newBlock = rewriter.createBlock( - &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs); - newCompute.getProperties().setOperandSegmentSizes( - {static_cast(promoted->newWeights.size()), static_cast(promoted->newInputs.size())}); + auto newCompute = createEmptySpatGraphCompute(rewriter, + compute.getLoc(), + compute.getResultTypes(), + promoted->newWeights, + promoted->newInputs, + TypeRange(newBlockArgTypes), + newBlockArgLocs); + auto* newBlock = &newCompute.getBody().front(); rewriter.setInsertionPointToStart(newBlock); IRRewriter bodyRewriter(rewriter.getContext()); @@ -193,12 +196,6 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern(compute.getLaneCount()), "promoted compute_batch lane count"); - if (failed(laneCountAttr)) - return failure(); - auto newCompute = spatial::SpatGraphComputeBatch::create( - rewriter, compute.getLoc(), compute.getResultTypes(), *laneCountAttr, promoted->newWeights, promoted->newInputs); auto laneArg = compute.getLaneArgument(); if (!laneArg) return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument"); @@ -223,23 +220,30 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePatterngetLoc()); } - auto* newBlock = rewriter.createBlock( - &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs); - newCompute.getProperties().setOperandSegmentSizes( - {static_cast(promoted->newWeights.size()), static_cast(promoted->newInputs.size())}); + auto newCompute = createEmptySpatGraphComputeBatch(rewriter, + compute.getLoc(), + compute.getResultTypes(), + compute.getLaneCount(), + promoted->newWeights, + promoted->newInputs, + TypeRange(newBlockArgTypes), + newBlockArgLocs); + if (failed(newCompute)) + return failure(); + auto* newBlock = &(*newCompute).getBody().front(); rewriter.setInsertionPointToStart(newBlock); IRRewriter bodyRewriter(rewriter.getContext()); bodyRewriter.setInsertionPointToStart(newBlock); IRMapping mapper; - auto newLaneArg = newCompute.getLaneArgument(); + auto newLaneArg = (*newCompute).getLaneArgument(); if (!newLaneArg) return rewriter.notifyMatchFailure(compute, "missing rewritten compute_batch lane block argument"); mapper.map(*laneArg, *newLaneArg); for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) { auto oldWeightArg = compute.getWeightArgument(weightIndex); - auto newWeightArg = newCompute.getWeightArgument(weightIndex); + auto newWeightArg = (*newCompute).getWeightArgument(weightIndex); if (!oldWeightArg || !newWeightArg) return rewriter.notifyMatchFailure(compute, "missing compute_batch weight block argument during rewrite"); mapper.map(*oldWeightArg, *newWeightArg); @@ -249,7 +253,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern(0, compute.getNumResults())) { @@ -263,7 +267,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern : SpatOp { + let hasCanonicalizer = 1; let extraClassDeclaration = [{ std::optional<::mlir::BlockArgument> getLaneArgument(); std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx); @@ -114,6 +115,7 @@ def SpatGraphComputeBatch : SpatComputeBatchLikeBase<"graph_compute_batch"> { } def SpatScheduledComputeBatch : SpatComputeBatchLikeBase<"scheduled_compute_batch"> { + let hasCanonicalizer = 1; let extraClassDeclaration = [{ std::optional<::mlir::BlockArgument> getLaneArgument(); std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx); diff --git a/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp b/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp index b7ab8d3..2df4361 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp @@ -50,64 +50,6 @@ static void printBlockArgumentList(OpAsmPrinter& printer, ArrayRef()) - continue; - printer.getStream() << "\n"; - printer.getStream().indent(indent); - printer.printCustomOrGenericOp(&nestedOp); - } - } - - printer.getStream() << "\n"; - printer << "}"; -} - static ParseResult parseBlockArgumentList(OpAsmParser& parser, SmallVectorImpl& arguments) { if (parser.parseLParen()) return failure(); @@ -218,9 +160,7 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) { printer << " -> "; printCompressedTypeSequence(printer, op.getResultTypes()); printer << " "; - printRegionWithoutBlockArgLocs( - printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true, - /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true); + printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock()); } template @@ -350,9 +290,7 @@ void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) { printer << " -> "; printCompressedTypeSequence(printer, op.getResultTypes()); printer << " "; - printRegionWithoutBlockArgLocs( - printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true, - /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true); + printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock()); } template @@ -521,7 +459,7 @@ void SpatDeferredCommunicationOp::print(OpAsmPrinter& printer) { printer << " : "; printer.printFunctionalType(getSources().getTypes(), getOperation()->getResultTypes()); printer << " "; - printRegionWithoutBlockArgLocs(printer, getBody(), /*printEntryBlockArgs=*/false); + printer.printRegion(getBody(), /*printEntryBlockArgs=*/false); } ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) { @@ -763,8 +701,7 @@ ParseResult SpatScheduledComputeBatch::parse(OpAsmParser& parser, OperationState void SpatInParallelOp::print(OpAsmPrinter& printer) { printer << " "; - printRegionWithoutBlockArgLocs( - printer, getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false); + printer.printRegion(getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false); printer.printOptionalAttrDict((*this)->getAttrs()); } diff --git a/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp b/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp index 278dffb..7b3bb5b 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp @@ -1,8 +1,13 @@ +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Block.h" +#include "mlir/IR/IRMapping.h" +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" #include "llvm/Support/LogicalResult.h" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" using namespace mlir; @@ -40,5 +45,177 @@ LogicalResult SpatScheduledCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVecto return foldComputeLike(*this, results); } +template +static ScalarComputeOpTy createEmptyScalarCompute(PatternRewriter& rewriter, + Location loc, + TypeRange resultTypes, + ValueRange weights, + ValueRange inputs) { + auto computeOp = ScalarComputeOpTy::create(rewriter, loc, resultTypes, weights, inputs); + SmallVector blockArgTypes; + SmallVector blockArgLocs; + blockArgTypes.reserve(weights.size() + inputs.size()); + blockArgLocs.reserve(weights.size() + inputs.size()); + for (Value weight : weights) { + blockArgTypes.push_back(weight.getType()); + blockArgLocs.push_back(weight.getLoc()); + } + for (Value input : inputs) { + blockArgTypes.push_back(input.getType()); + blockArgLocs.push_back(input.getLoc()); + } + rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs); + rewriter.setInsertionPointToStart(&computeOp.getBody().front()); + return computeOp; +} + +static SmallVector remapMixedOffsets(ArrayRef mixedOffsets, IRMapping& mapper) { + SmallVector remapped; + remapped.reserve(mixedOffsets.size()); + for (OpFoldResult ofr : mixedOffsets) { + if (auto value = dyn_cast(ofr)) + remapped.push_back(cast(mapper.lookupOrDefault(value))); + else + remapped.push_back(cast(ofr)); + } + return remapped; +} + +static SmallVector createEmptyResults(PatternRewriter& rewriter, Location loc, TypeRange resultTypes) { + SmallVector resultValues; + resultValues.reserve(resultTypes.size()); + for (Type resultType : resultTypes) { + auto tensorType = dyn_cast(resultType); + if (!tensorType || !tensorType.hasStaticShape()) + return {}; + resultValues.push_back(tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), tensorType.getElementType())); + } + return resultValues; +} + +template +static void copyCanonicalizedBatchAttrs(ScalarComputeOpTy compute, ComputeBatchOpTy batch, PatternRewriter& rewriter) { + for (NamedAttribute attr : batch->getAttrs()) { + if (attr.getName() == batch.getOperandSegmentSizesAttrName() || attr.getName() == batch.getLaneCountAttrName() + || attr.getName() == onnx_mlir::kCoreIdsAttrName) + continue; + compute->setAttr(attr.getName(), attr.getValue()); + } + if constexpr (std::is_same_v) { + if (auto coreIds = batch->template getAttrOfType(onnx_mlir::kCoreIdsAttrName)) { + assert(coreIds.size() == 1 && "single-lane scheduled compute_batch canonicalization expects exactly one core id"); + compute->setAttr(onnx_mlir::kCoreIdAttrName, rewriter.getI32IntegerAttr(coreIds.asArrayRef().front())); + } + } +} + +template +struct CanonicalizeSingleLaneComputeBatchPattern : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(ComputeBatchOpTy compute, PatternRewriter& rewriter) const override { + if (compute.getLaneCount() != 1) + return rewriter.notifyMatchFailure(compute, "lane count is not 1"); + + Block& oldBlock = compute.getBody().front(); + auto oldLaneArg = compute.getLaneArgument(); + if (!oldLaneArg) + return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument"); + + rewriter.setInsertionPointAfter(compute); + auto newCompute = + createEmptyScalarCompute(rewriter, compute.getLoc(), compute.getResultTypes(), compute.getWeights(), compute.getInputs()); + copyCanonicalizedBatchAttrs(newCompute, compute, rewriter); + auto* newBlock = &newCompute.getBody().front(); + rewriter.setInsertionPointToStart(newBlock); + + IRMapping mapper; + Value zero = arith::ConstantIndexOp::create(rewriter, compute.getLoc(), 0); + mapper.map(*oldLaneArg, zero); + for (auto [index, weight] : llvm::enumerate(compute.getWeights())) { + auto oldArg = compute.getWeightArgument(index); + auto newArg = newCompute.getWeightArgument(index); + if (!oldArg || !newArg) + return rewriter.notifyMatchFailure(compute, "missing rewritten compute weight block argument"); + mapper.map(*oldArg, *newArg); + } + for (auto [index, input] : llvm::enumerate(compute.getInputs())) { + auto oldArg = compute.getInputArgument(index); + auto newArg = newCompute.getInputArgument(index); + if (!oldArg || !newArg) + return rewriter.notifyMatchFailure(compute, "missing rewritten compute input block argument"); + mapper.map(*oldArg, *newArg); + } + + SmallVector resultValues = createEmptyResults(rewriter, compute.getLoc(), compute.getResultTypes()); + if (resultValues.size() != compute.getNumResults()) + return rewriter.notifyMatchFailure(compute, "single-lane compute_batch canonicalization requires static ranked results"); + for (auto [index, resultValue] : llvm::enumerate(resultValues)) { + auto oldOutputArg = compute.getOutputArgument(index); + if (!oldOutputArg) + return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument"); + mapper.map(*oldOutputArg, resultValue); + } + + auto oldInParallel = dyn_cast(oldBlock.getTerminator()); + auto oldYield = dyn_cast(oldBlock.getTerminator()); + for (Operation& op : oldBlock.without_terminator()) + rewriter.clone(op, mapper); + + if (oldYield) { + SpatYieldOp::create(rewriter, oldYield.getLoc(), ValueRange {}); + rewriter.replaceOp(compute, newCompute.getResults()); + return success(); + } + if (!oldInParallel) + return rewriter.notifyMatchFailure(compute, "expected spat.in_parallel or empty spat.yield terminator"); + + DenseMap outputIndexByArg; + for (size_t index = 0; index < compute.getNumResults(); ++index) { + auto oldOutputArg = compute.getOutputArgument(index); + if (!oldOutputArg) + return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument"); + outputIndexByArg[*oldOutputArg] = index; + } + + for (Operation& op : oldInParallel.getRegion().front()) { + auto insertSlice = dyn_cast(&op); + if (!insertSlice) + return rewriter.notifyMatchFailure(compute, "expected only tensor.parallel_insert_slice in spat.in_parallel"); + auto oldDest = dyn_cast(insertSlice.getDest()); + if (!oldDest) + return rewriter.notifyMatchFailure(compute, "expected tensor.parallel_insert_slice destination to be a block argument"); + auto resultIndexIt = outputIndexByArg.find(oldDest); + if (resultIndexIt == outputIndexByArg.end()) + return rewriter.notifyMatchFailure(compute, "unexpected tensor.parallel_insert_slice destination"); + size_t resultIndex = resultIndexIt->second; + Value remappedSource = mapper.lookupOrDefault(insertSlice.getSource()); + auto remappedOffsets = remapMixedOffsets(insertSlice.getMixedOffsets(), mapper); + auto remappedSizes = remapMixedOffsets(insertSlice.getMixedSizes(), mapper); + auto remappedStrides = remapMixedOffsets(insertSlice.getMixedStrides(), mapper); + resultValues[resultIndex] = tensor::InsertSliceOp::create(rewriter, + insertSlice.getLoc(), + remappedSource, + resultValues[resultIndex], + remappedOffsets, + remappedSizes, + remappedStrides) + .getResult(); + } + + SpatYieldOp::create(rewriter, oldInParallel.getLoc(), resultValues); + rewriter.replaceOp(compute, newCompute.getResults()); + return success(); + } +}; + +void SpatGraphComputeBatch::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { + results.add>(context); +} + +void SpatScheduledComputeBatch::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) { + results.add>(context); +} + } // namespace spatial } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp index 4ae9432..e84f163 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp @@ -847,30 +847,33 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) return batch.emitOpError("compute_batch coreIds values must be unique"); } - Block& block = batch.getBody().front(); - if (block.getNumArguments() == 0) - return batch.emitOpError("compute_batch body must have exactly one lane block argument"); - unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults(); - if (block.getNumArguments() != expectedArgCount) - return batch.emitOpError("compute_batch body block arguments must match lane, weight, input, and output operands/results"); - auto laneArg = batch.getLaneArgument(); - if (!laneArg || !laneArg->getType().isIndex()) - return batch.emitOpError("compute_batch first block argument must have index type"); + if (batch.getBody().empty()) + return batch.emitOpError("compute_batch body must have at least one block"); - for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights())) { - auto blockArg = batch.getWeightArgument(weightIndex); - if (!blockArg || blockArg->getType() != weight.getType()) - return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly"); - } - for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs())) { - auto blockArg = batch.getInputArgument(inputIndex); - if (!blockArg || blockArg->getType() != input.getType()) - return batch.emitOpError("compute_batch input block argument types must match input operand types exactly"); - } - for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes())) { - auto blockArg = batch.getOutputArgument(resultIndex); - if (!blockArg || blockArg->getType() != resultType) - return batch.emitOpError("compute_batch output block argument types must match result types exactly"); + unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults(); + bool verifyLaneSliceOffsets = !isa(batch.getOperation()); + for (Block& block : batch.getBody()) { + if (block.getNumArguments() == 0) + return batch.emitOpError("compute_batch body must have exactly one lane block argument"); + if (block.getNumArguments() != expectedArgCount) + return batch.emitOpError( + "compute_batch body block arguments must match lane, weight, input, and output operands/results"); + if (!block.getArgument(0).getType().isIndex()) + return batch.emitOpError("compute_batch first block argument must have index type"); + + for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights())) + if (block.getArgument(1 + weightIndex).getType() != weight.getType()) + return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly"); + for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs())) + if (block.getArgument(1 + batch.getWeights().size() + inputIndex).getType() != input.getType()) + return batch.emitOpError("compute_batch input block argument types must match input operand types exactly"); + for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes())) + if (block.getArgument(1 + batch.getWeights().size() + batch.getInputs().size() + resultIndex).getType() + != resultType) + return batch.emitOpError("compute_batch output block argument types must match result types exactly"); + + if (failed(verifyBatchBody(batch, block, verifyLaneSliceOffsets))) + return failure(); } if (failed(verifyComputeResultsUses(batch.getOperation()))) @@ -879,7 +882,7 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) return failure(); if (failed(verifyOnlyConstantExternalValues(batch.getOperation(), batch.getBody(), opName))) return failure(); - return verifyBatchBody(batch, block, !isa(batch.getOperation())); + return success(); } LogicalResult SpatGraphComputeBatch::verify() { return verifyComputeBatchLikeOp(*this, "spat.graph_compute_batch"); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp new file mode 100644 index 0000000..7f71439 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp @@ -0,0 +1,426 @@ +#include "DeferredCommunicationPlanning.hpp" + +#include "mlir/Dialect/Tensor/IR/Tensor.h" + +namespace onnx_mlir { +namespace spatial { + +using namespace mlir; + +namespace { + +struct LaneProducerInfo { + ProducerValueRef producer; + Value originalSource; + DeferredEndpoint producerEndpoint; + DeferredEndpoint consumerEndpoint; +}; + +static Value getBlockOperand(Block &block, ValueRange operands, Value value, unsigned firstArgument = 0) { + auto it = llvm::find(operands, value); + assert(it != operands.end() && "missing scheduled operand"); + return block.getArgument(firstArgument + std::distance(operands.begin(), it)); +} + +static FailureOr getOriginalProducerValue(const ProducerValueRef &producer) { + auto outputs = getComputeInstanceOutputValues(producer.instance); + if (producer.resultIndex >= outputs.size()) + return failure(); + return outputs[producer.resultIndex]; +} + +static DeferredEndpoint getDeferredEndpoint(const ComputeInstance &instance, + const MergeScheduleResult &schedule) { + size_t cpu = getScheduledCpuForComputeInstance(instance, schedule); + return DeferredEndpoint { + instance, + cpu, + getCanonicalPeftClassId(cpu, schedule), + instance.laneStart, + instance.laneCount + }; +} + +static void setDeferredTransferMetadata(SpatDeferredCommunicationOp transfer, + OpBuilder &builder, + const DeferredTransferMetadata &transferMetadata) { + SmallVector sourceCpus; + SmallVector sourceClasses; + SmallVector sourceLaneRanges; + SmallVector targetCpus; + SmallVector targetClasses; + SmallVector targetLaneRanges; + for (const DeferredEndpoint &endpoint : transferMetadata.producers) { + sourceCpus.push_back(static_cast(endpoint.cpu)); + sourceClasses.push_back(static_cast(endpoint.canonicalClassId)); + sourceLaneRanges.push_back(static_cast(endpoint.laneStart)); + sourceLaneRanges.push_back(static_cast(endpoint.laneCount)); + } + for (const DeferredEndpoint &endpoint : transferMetadata.consumers) { + targetCpus.push_back(static_cast(endpoint.cpu)); + targetClasses.push_back(static_cast(endpoint.canonicalClassId)); + targetLaneRanges.push_back(static_cast(endpoint.laneStart)); + targetLaneRanges.push_back(static_cast(endpoint.laneCount)); + } + transfer->setAttr("source_cpus", builder.getDenseI64ArrayAttr(sourceCpus)); + transfer->setAttr("source_classes", builder.getDenseI64ArrayAttr(sourceClasses)); + transfer->setAttr("source_lane_ranges", builder.getDenseI64ArrayAttr(sourceLaneRanges)); + transfer->setAttr("target_cpus", builder.getDenseI64ArrayAttr(targetCpus)); + transfer->setAttr("target_classes", builder.getDenseI64ArrayAttr(targetClasses)); + transfer->setAttr("target_lane_ranges", builder.getDenseI64ArrayAttr(targetLaneRanges)); + // `batched=true` means this deferred communication belongs to a + // spat.scheduled_compute_batch block and the array attrs are indexed by + // scheduled lane. The legacy scalar attrs on the op keep representative values + // for compatibility; Phase 2 must consume the arrays when batched=true. + transfer->setAttr("batched", builder.getBoolAttr(transferMetadata.isBatched)); + if (transferMetadata.isBatched) { + transfer->setAttr("source_operand_for_scheduled_lane", + builder.getDenseI64ArrayAttr(transferMetadata.sourceOperandForScheduledLane)); + transfer->setAttr("multi_source_payload", builder.getBoolAttr(transferMetadata.multiSourcePayload)); + } +} + +static FailureOr projectDeferredPayload(Value selectedSource, + Value consumerInput, + OpBuilder &builder, + Location loc) { + if (selectedSource.getType() == consumerInput.getType()) + return selectedSource; + + auto slice = consumerInput.getDefiningOp(); + if (!slice || slice.getSource().getType() != selectedSource.getType()) + return failure(); + + IRMapping mapper; + mapper.map(slice.getSource(), selectedSource); + return cast(builder.clone(*slice, mapper)).getResult(); +} + +static FailureOr buildSelectedDeferredSource(OpBuilder &builder, + Location loc, + SpatDeferredCommunicationOp transfer, + Value scheduledLane, + ValueRange sourceBlockArgs, + DenseI64ArrayAttr sourceOperandForScheduledLaneAttr) { + SmallVector sourceOperandForScheduledLane(sourceOperandForScheduledLaneAttr.asArrayRef().begin(), + sourceOperandForScheduledLaneAttr.asArrayRef().end()); + Value sourceOperandIndexTable = + createI64LookupTableConstant(builder, transfer.getOperation(), sourceOperandForScheduledLane); + Value sourceIndexI64 = + tensor::ExtractOp::create(builder, loc, sourceOperandIndexTable, ValueRange {scheduledLane}).getResult(); + Value sourceIndex = arith::IndexCastOp::create(builder, loc, builder.getIndexType(), sourceIndexI64).getResult(); + + Type elementType = sourceBlockArgs.front().getType(); + auto sourceTensorType = dyn_cast(elementType); + if (!sourceTensorType || !sourceTensorType.hasStaticShape()) + return transfer.emitOpError( + "scheduled batch deferred communication with multiple original graph SSA sources requires compatible static ranked tensor sources"), + failure(); + + for (Value source : sourceBlockArgs) { + auto rankedSourceType = dyn_cast(source.getType()); + if (!rankedSourceType || !rankedSourceType.hasStaticShape() || rankedSourceType != sourceTensorType) + return transfer.emitOpError( + "scheduled batch deferred communication with multiple original graph SSA sources requires compatible static ranked tensor sources"), + failure(); + } + if (sourceTensorType.getRank() == 0) + return transfer.emitOpError( + "scheduled batch deferred communication with multiple original graph SSA sources cannot yet build compact tensor source selection"), + failure(); + + SmallVector stackedShape {static_cast(sourceBlockArgs.size())}; + llvm::append_range(stackedShape, sourceTensorType.getShape()); + RankedTensorType stackedType = RankedTensorType::get(stackedShape, sourceTensorType.getElementType()); + auto stackedEmpty = createEmptyTensorForType(builder, loc, stackedType); + if (failed(stackedEmpty)) + return transfer.emitOpError( + "scheduled batch deferred communication with multiple original graph SSA sources cannot yet build compact tensor source selection"), + failure(); + Value stackedSources = *stackedEmpty; + + SmallVector insertSizes; + SmallVector unitStrides(stackedType.getRank(), builder.getIndexAttr(1)); + insertSizes.push_back(builder.getIndexAttr(1)); + for (int64_t dim : sourceTensorType.getShape()) + insertSizes.push_back(builder.getIndexAttr(dim)); + + for (auto [sourceIndexValue, source] : llvm::enumerate(sourceBlockArgs)) { + SmallVector expandedShape {1}; + llvm::append_range(expandedShape, sourceTensorType.getShape()); + SmallVector reassociation; + reassociation.push_back({0, 1}); + for (int64_t dim = 1; dim < sourceTensorType.getRank(); ++dim) + reassociation.push_back({dim + 1}); + Value expandedSource = tensor::ExpandShapeOp::create(builder, + loc, + RankedTensorType::get(expandedShape, sourceTensorType.getElementType()), + source, + reassociation) + .getResult(); + SmallVector insertOffsets(stackedType.getRank(), builder.getIndexAttr(0)); + insertOffsets[0] = builder.getIndexAttr(sourceIndexValue); + stackedSources = tensor::InsertSliceOp::create( + builder, loc, expandedSource, stackedSources, insertOffsets, insertSizes, unitStrides) + .getResult(); + } + + SmallVector extractOffsets(stackedType.getRank(), builder.getIndexAttr(0)); + SmallVector extractSizes; + SmallVector extractStrides(stackedType.getRank(), builder.getIndexAttr(1)); + extractOffsets[0] = sourceIndex; + extractSizes.push_back(builder.getIndexAttr(1)); + for (int64_t dim : sourceTensorType.getShape()) + extractSizes.push_back(builder.getIndexAttr(dim)); + SmallVector selectedSliceShape {1}; + llvm::append_range(selectedSliceShape, sourceTensorType.getShape()); + RankedTensorType selectedSliceType = RankedTensorType::get(selectedSliceShape, sourceTensorType.getElementType()); + Value selectedSlice = tensor::ExtractSliceOp::create( + builder, + loc, + selectedSliceType, + stackedSources, + extractOffsets, + extractSizes, + extractStrides) + .getResult(); + + SmallVector collapseReassociation; + collapseReassociation.push_back({0, 1}); + for (int64_t dim = 1; dim < sourceTensorType.getRank(); ++dim) + collapseReassociation.push_back({dim + 1}); + return tensor::CollapseShapeOp::create(builder, loc, sourceTensorType, selectedSlice, collapseReassociation) + .getResult(); +} + +static LogicalResult finalizeDeferredCommunication(OpBuilder &builder, + Location loc, + Value consumerInput, + SpatDeferredCommunicationOp transfer, + Value &result) { + Block *block = builder.createBlock(&transfer.getBody(), + transfer.getBody().end(), + TypeRange {transfer.getSources().getTypes()}, + SmallVector(transfer.getSources().size(), loc)); + builder.setInsertionPointToStart(block); + + Value selectedSource = block->getArgument(0); + if (auto parentScheduled = dyn_cast(transfer->getParentOp())) { + auto sourceOperandForScheduledLaneAttr = + transfer->getAttrOfType("source_operand_for_scheduled_lane"); + auto multiSourcePayloadAttr = transfer->getAttrOfType("multi_source_payload"); + if (sourceOperandForScheduledLaneAttr && sourceOperandForScheduledLaneAttr.size() == parentScheduled.getLaneCount()) { + if (multiSourcePayloadAttr && multiSourcePayloadAttr.getValue()) { + FailureOr multiSourceSelection = buildSelectedDeferredSource( + builder, + loc, + transfer, + transfer->getBlock()->getArgument(0), + block->getArguments(),finalizeDeferredCommunication + sourceOperandForScheduledLaneAttr); + if (failed(multiSourceSelection)) + return failure(); + selectedSource = *multiSourceSelection; + } else { + selectedSource = block->getArgument(sourceOperandForScheduledLaneAttr.asArrayRef().front()); + } + } + } + + FailureOr yielded = projectDeferredPayload(selectedSource, consumerInput, builder, loc); + if (failed(yielded) || (*yielded).getType() != consumerInput.getType()) + return transfer.emitOpError("cannot derive deferred communication payload from original graph producer value"); + + SpatYieldOp::create(builder, loc, *yielded); + result = transfer.getOutput(); + builder.setInsertionPointAfter(transfer); + return success(); +} + +static LogicalResult createSingleLaneDeferredCommunication(OpBuilder &builder, + Location loc, + Value consumerInput, + Value originalSource, + const ProducerValueRef &producer, + const ComputeInstance &consumer, + const MergeScheduleResult &schedule, + uint64_t exchangeId, + Value &result) { + DeferredTransferMetadata transferMetadata; + transferMetadata.producers.push_back(getDeferredEndpoint(producer.instance, schedule)); + transferMetadata.consumers.push_back(getDeferredEndpoint(consumer, schedule)); + auto transfer = SpatDeferredCommunicationOp::create( + builder, + loc, + consumerInput.getType(), + ValueRange {originalSource}, + builder.getStringAttr(("x" + std::to_string(exchangeId)).c_str()), + builder.getStringAttr(getInstanceName(producer.instance)), + builder.getStringAttr(getInstanceName(consumer)), + builder.getI64IntegerAttr(static_cast(transferMetadata.producers.front().canonicalClassId)), + builder.getI64IntegerAttr(static_cast(transferMetadata.consumers.front().canonicalClassId)), + builder.getI64IntegerAttr(static_cast(transferMetadata.producers.front().cpu)), + builder.getI64IntegerAttr(static_cast(transferMetadata.consumers.front().cpu)), + builder.getI64IntegerAttr(producer.instance.laneStart), + builder.getI64IntegerAttr(consumer.laneStart), + builder.getStringAttr(consumerInput.getDefiningOp() ? "projected" : "ordinary"), + builder.getI64IntegerAttr(static_cast(producer.resultIndex)), + builder.getDenseI64ArrayAttr({}), + builder.getI64IntegerAttr(-1)); + setDeferredTransferMetadata(transfer, builder, transferMetadata); + return finalizeDeferredCommunication(builder, loc, consumerInput, transfer, result); +} + +static LogicalResult createScheduledLaneDeferredCommunication(OpBuilder &builder, + Location loc, + Value consumerInput, + ValueRange originalSources, + const ProducerValueRef &producer, + const DeferredTransferMetadata &transferMetadata, + uint64_t exchangeId, + Value &result) { + assert(transferMetadata.isBatched && "expected scheduled-lane deferred communication metadata"); + assert(!transferMetadata.producers.empty() && !transferMetadata.consumers.empty() && "expected non-empty endpoints"); + auto transfer = SpatDeferredCommunicationOp::create( + builder, + loc, + consumerInput.getType(), + originalSources, + builder.getStringAttr(("x" + std::to_string(exchangeId)).c_str()), + builder.getStringAttr(getInstanceName(producer.instance)), + builder.getStringAttr(getInstanceName(transferMetadata.consumers.front().instance)), + builder.getI64IntegerAttr(static_cast(transferMetadata.producers.front().canonicalClassId)), + builder.getI64IntegerAttr(static_cast(transferMetadata.consumers.front().canonicalClassId)), + builder.getI64IntegerAttr(static_cast(transferMetadata.producers.front().cpu)), + builder.getI64IntegerAttr(static_cast(transferMetadata.consumers.front().cpu)), + builder.getI64IntegerAttr(producer.instance.laneStart), + builder.getI64IntegerAttr(transferMetadata.consumers.front().laneStart), + builder.getStringAttr(consumerInput.getDefiningOp() ? "projected" : "ordinary"), + builder.getI64IntegerAttr(static_cast(producer.resultIndex)), + builder.getDenseI64ArrayAttr({}), + builder.getI64IntegerAttr(-1)); + setDeferredTransferMetadata(transfer, builder, transferMetadata); + return finalizeDeferredCommunication(builder, loc, consumerInput, transfer, result); +} + +} // namespace + +LogicalResult mapSingleCpuInput(OpBuilder &builder, + Location loc, + Value input, + const ComputeInstance &consumerInstance, + const MergeScheduleResult &schedule, + ValueRange scheduledInputs, + Block &block, + unsigned firstInputArgument, + ArrayRef carriedKeys, + uint64_t &exchangeId, + Value &mapped) { + std::optional producer = getProducerValueRef(input, &consumerInstance); + if (!producer) { + mapped = getBlockOperand(block, scheduledInputs, input, firstInputArgument); + return success(); + } + + ProducerValueKey producerKey {producer->instance, producer->resultIndex}; + auto carriedIt = llvm::find(carriedKeys, producerKey); + if (carriedIt != carriedKeys.end()) { + mapped = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carriedIt)); + return success(); + } + + FailureOr originalSource = getOriginalProducerValue(*producer); + if (failed(originalSource)) + return emitError(loc) << "cannot resolve original graph producer value for " << getInstanceName(producer->instance); + return createSingleLaneDeferredCommunication( + builder, loc, input, *originalSource, *producer, consumerInstance, schedule, exchangeId++, mapped); +} + +static FailureOr getComputeInstanceInputIndex(const ComputeInstance &instance, Value input) { + auto inputs = getComputeInstanceInputs(instance); + auto it = llvm::find(inputs, input); + if (it == inputs.end()) + return failure(); + return static_cast(std::distance(inputs.begin(), it)); +} + +LogicalResult mapMultiCpuTupleInput(OpBuilder &builder, + Location loc, + Value input, + const ComputeStepTuple &stepTuple, + const PeftClassPlan &peftClassPlan, + const MergeScheduleResult &schedule, + ValueRange scheduledInputs, + Block &block, + unsigned firstInputArgument, + uint64_t &exchangeId, + Value &mapped) { + assert(!stepTuple.instances.empty() && "expected non-empty step tuple"); + const ComputeInstance &representative = stepTuple.instances.front(); + std::optional representativeProducer = getProducerValueRef(input, &representative); + if (!representativeProducer) { + FailureOr inputIndex = getComputeInstanceInputIndex(representative, input); + if (failed(inputIndex)) + return emitError(loc) << "cannot resolve scheduled batch step host input index"; + for (const ComputeInstance &instance : stepTuple.instances) { + auto instanceInputs = getComputeInstanceInputs(instance); + if (*inputIndex >= instanceInputs.size()) + return emitError(loc) << "scheduled batch step host input index out of range"; + if (instanceInputs[*inputIndex] != input || getProducerValueRef(instanceInputs[*inputIndex], &instance)) + return emitError(loc) << "scheduled batch step requires identical host inputs across lanes"; + } + mapped = getBlockOperand(block, scheduledInputs, input, firstInputArgument); + return success(); + } + + FailureOr inputIndex = getComputeInstanceInputIndex(representative, input); + if (failed(inputIndex)) + return emitError(loc) << "cannot resolve scheduled batch step input index"; + + (void)peftClassPlan; + DeferredTransferMetadata transferMetadata; + transferMetadata.isBatched = true; + SmallVector laneProducerInfos; + SmallVector uniqueOriginalSources; + SmallVector sourceOperandForScheduledLane; + for (const ComputeInstance &instance : stepTuple.instances) { + auto instanceInputs = getComputeInstanceInputs(instance); + if (*inputIndex >= instanceInputs.size()) + return emitError(loc) << "scheduled batch step input index out of range"; + Value laneInput = instanceInputs[*inputIndex]; + std::optional laneProducer = getProducerValueRef(laneInput, &instance); + if (!laneProducer) + return emitError(loc) << "scheduled batch step mixes host and producer-derived inputs"; + FailureOr laneSource = getOriginalProducerValue(*laneProducer); + if (failed(laneSource)) + return emitError(loc) << "cannot resolve original graph producer value for " << getInstanceName(laneProducer->instance); + + auto sourceIt = llvm::find(uniqueOriginalSources, *laneSource); + if (sourceIt == uniqueOriginalSources.end()) { + sourceOperandForScheduledLane.push_back(static_cast(uniqueOriginalSources.size())); + uniqueOriginalSources.push_back(*laneSource); + } else { + sourceOperandForScheduledLane.push_back(std::distance(uniqueOriginalSources.begin(), sourceIt)); + } + + laneProducerInfos.push_back(LaneProducerInfo { + *laneProducer, + *laneSource, + getDeferredEndpoint(laneProducer->instance, schedule), + getDeferredEndpoint(instance, schedule), + }); + } + + for (const LaneProducerInfo &laneProducerInfo : laneProducerInfos) { + transferMetadata.producers.push_back(laneProducerInfo.producerEndpoint); + transferMetadata.consumers.push_back(laneProducerInfo.consumerEndpoint); + } + transferMetadata.sourceOperandForScheduledLane = sourceOperandForScheduledLane; + transferMetadata.multiSourcePayload = uniqueOriginalSources.size() > 1; + + return createScheduledLaneDeferredCommunication( + builder, loc, input, uniqueOriginalSources, *representativeProducer, transferMetadata, exchangeId++, mapped); +} + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp new file mode 100644 index 0000000..e4598cf --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp @@ -0,0 +1,33 @@ +#pragma once + +#include "ScheduledComputePlan.hpp" + +namespace onnx_mlir { +namespace spatial { + +LogicalResult mapSingleCpuInput(OpBuilder &builder, + Location loc, + Value input, + const ComputeInstance &consumerInstance, + const MergeScheduleResult &schedule, + ValueRange scheduledInputs, + Block &block, + unsigned firstInputArgument, + ArrayRef carriedKeys, + uint64_t &exchangeId, + Value &mapped); + +LogicalResult mapMultiCpuTupleInput(OpBuilder &builder, + Location loc, + Value input, + const ComputeStepTuple &stepTuple, + const PeftClassPlan &peftClassPlan, + const MergeScheduleResult &schedule, + ValueRange scheduledInputs, + Block &block, + unsigned firstInputArgument, + uint64_t &exchangeId, + Value &mapped); + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp new file mode 100644 index 0000000..4e012b2 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp @@ -0,0 +1,94 @@ +#include "ScheduledComputeMaterialization.hpp" +#include "ScheduledComputeReport.hpp" +#include "ScheduledComputeVerification.hpp" + +#include "mlir/Pass/Pass.h" + +#include "Scheduling/MergeSchedulingAnalysis.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Pass/PIMPasses.h" + +using namespace mlir; + +namespace onnx_mlir { +namespace spatial { +namespace { + +struct MergeComputeNodesPass final : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass) + + StringRef getArgument() const override { return "pim-merge-compute-nodes"; } + StringRef getDescription() const override { + return "Materialize scheduled Spatial compute with deferred communication placeholders."; + } + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + auto entryFunc = getPimEntryFunc(moduleOp); + if (failed(entryFunc)) { + moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes"); + signalPassFailure(); + return; + } + + func::FuncOp funcOp = *entryFunc; + MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult(); + PatternRewriter rewriter(moduleOp.getContext()); + FailureOr materialization = + materializeScheduledCompute(funcOp, schedule, rewriter); + if (failed(materialization)) { + signalPassFailure(); + return; + } + + if (failed(verifyMaterializedScheduleMapping(funcOp, + schedule, + materialization->peftClassPlans, + materialization->graphComputeToBlockMap, + materialization->materializedSchedules))) { + moduleOp.emitError("scheduled Spatial materialization mapping verification failed"); + signalPassFailure(); + return; + } + if (failed(verifyDeferredTransferPhase1Invariants(funcOp))) { + moduleOp.emitError("scheduled Spatial deferred communication verification failed"); + signalPassFailure(); + return; + } + if (failed(verifyScheduledMaterializationRecords(materialization->materializedSchedules))) { + moduleOp.emitError("scheduled Spatial materialization record verification failed"); + signalPassFailure(); + return; + } + if (failed(verifyPeftMaterializationReportSummary(funcOp, + schedule, + materialization->peftClassPlans, + materialization->materializedSchedules))) { + moduleOp.emitError("scheduled Spatial report verification failed"); + signalPassFailure(); + return; + } + if (failed(verifyScheduledSpatialInvariants(funcOp))) { + moduleOp.emitError("scheduled Spatial phase 1 verification failed"); + signalPassFailure(); + return; + } + + dumpScheduledComputeReportAndModule(moduleOp, + funcOp, + schedule, + materialization->peftClassPlans, + materialization->materializedSchedules); + moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; " + "Phase 2 communication realization is not implemented"); + signalPassFailure(); + } +}; + +} // namespace +} // namespace spatial + +std::unique_ptr createMergeComputeNodesPass() { return std::make_unique(); } + +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp new file mode 100644 index 0000000..515686c --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp @@ -0,0 +1,954 @@ +#include "ScheduledComputeMaterialization.hpp" +#include "DeferredCommunicationPlanning.hpp" + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" + +#include + +namespace onnx_mlir { +namespace spatial { + +using namespace mlir; +namespace { + +struct BatchFragmentSpec { + RankedTensorType resultType; + RankedTensorType sourceSliceType; + unsigned laneDim = 0; + SmallVector offsets; + SmallVector sizes; + SmallVector strides; +}; + +static SmallVector remapMixedOffsets(ArrayRef mixedOffsets, IRMapping &mapper) { + SmallVector remapped; + remapped.reserve(mixedOffsets.size()); + for (OpFoldResult ofr : mixedOffsets) { + if (auto value = dyn_cast(ofr)) + remapped.push_back(cast(mapper.lookupOrDefault(value))); + else + remapped.push_back(cast(ofr)); + } + return remapped; +} + +static void appendUnique(SmallVectorImpl &values, Value value) { + if (!llvm::is_contained(values, value)) + values.push_back(value); +} + +static Value getBlockOperand(Block &block, ValueRange operands, Value value, unsigned firstArgument = 0) { + auto it = llvm::find(operands, value); + assert(it != operands.end() && "missing scheduled operand"); + return block.getArgument(firstArgument + std::distance(operands.begin(), it)); +} + +static Value getScheduledComputeOutputArgument(Block &block, ValueRange scheduledWeights, ValueRange scheduledInputs, + ArrayRef carriedKeys, ProducerValueKey key) { + unsigned base = scheduledWeights.size() + scheduledInputs.size(); + auto it = llvm::find(carriedKeys, key); + assert(it != carriedKeys.end() && "missing carried output"); + return block.getArgument(base + std::distance(carriedKeys.begin(), it)); +} + +static unsigned getScheduledComputeResultArgBase(SpatScheduledCompute scheduled) { + return scheduled.getWeights().size() + scheduled.getInputs().size(); +} + +static void appendComputeBlockArguments(SmallVectorImpl &argTypes, + SmallVectorImpl &argLocs, + ValueRange weights, + ValueRange inputs, + ArrayRef carriedKeys, + Location loc) { + for (Value weight : weights) + argTypes.push_back(weight.getType()); + for (Value input : inputs) + argTypes.push_back(input.getType()); + for (ProducerValueKey key : carriedKeys) { + auto outputs = getComputeInstanceOutputValues(key.instance); + assert(key.resultIndex < outputs.size() && "missing carried result"); + argTypes.push_back(outputs[key.resultIndex].getType()); + } + argLocs.append(argTypes.size(), loc); +} + +static Block *createScheduledComputeBlock(PatternRewriter &rewriter, + SpatScheduledCompute scheduled, + ArrayRef carriedKeys, + Location loc) { + SmallVector argTypes; + SmallVector argLocs; + appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(), scheduled.getInputs(), carriedKeys, loc); + return rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), TypeRange(argTypes), argLocs); +} + +static void appendBlockYieldBaseAndCarriedOperands(Block &block, + unsigned baseArgCount, + size_t carriedCount, + SmallVectorImpl &operands) { + for (unsigned index = 0; index < baseArgCount; ++index) + operands.push_back(block.getArgument(index)); + for (size_t index = 0; index < carriedCount; ++index) + operands.push_back(block.getArgument(baseArgCount + index)); +} + +static void createBlockYield(PatternRewriter &rewriter, Location loc, ValueRange outputs, Block *next = nullptr) { + OperationState state(loc, SpatBlockYieldOp::getOperationName()); + state.addOperands(outputs); + if (next) + state.addSuccessors(next); + rewriter.create(state); +} + +static FailureOr getBatchFragmentSpec(SpatComputeBatch batch, + unsigned resultIndex, + uint32_t fragmentLaneCount) { + auto inParallel = dyn_cast(batch.getBody().front().getTerminator()); + if (!inParallel) + return batch.emitOpError("scheduled materialization only supports resultful spat.graph_compute_batch"); + + auto outputArg = batch.getOutputArgument(resultIndex); + if (!outputArg) + return batch.emitOpError("scheduled materialization could not locate batch output block argument"); + + for (Operation &op : inParallel.getRegion().front()) { + auto insert = dyn_cast(&op); + if (!insert) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + if (insert.getDest() != *outputArg) + continue; + + RankedTensorType destType = insert.getDestType(); + RankedTensorType sourceType = insert.getSourceType(); + if (!destType || !sourceType || !destType.hasStaticShape() || !sourceType.hasStaticShape()) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + if (destType.getRank() == 0 || sourceType.getRank() == 0 || destType.getRank() != sourceType.getRank()) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + if (!insert.hasUnitStride()) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + + auto offsets = insert.getMixedOffsets(); + auto sizes = insert.getMixedSizes(); + auto strides = insert.getMixedStrides(); + if (offsets.size() != static_cast(destType.getRank()) || sizes.size() != static_cast(destType.getRank()) + || strides.size() != static_cast(destType.getRank())) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + + std::optional laneDim; + for (unsigned dim = 0; dim < offsets.size(); ++dim) { + bool dimLooksLaneLike = sourceType.getShape()[dim] != destType.getShape()[dim]; + if (auto value = dyn_cast(offsets[dim])) + dimLooksLaneLike = dimLooksLaneLike || valueTransitivelyDependsOn(value, *batch.getLaneArgument()); + if (dimLooksLaneLike) { + if (laneDim && *laneDim != dim) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + laneDim = dim; + } + } + if (!laneDim) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + + SmallVector shape(destType.getShape().begin(), destType.getShape().end()); + shape[*laneDim] = fragmentLaneCount; + return BatchFragmentSpec { + RankedTensorType::get(shape, destType.getElementType()), + sourceType, + *laneDim, + SmallVector(offsets.begin(), offsets.end()), + SmallVector(sizes.begin(), sizes.end()), + SmallVector(strides.begin(), strides.end()) + }; + } + + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); +} + + +static SourceLaneSelector buildSourceLaneSelector(PatternRewriter &rewriter, + const ComputeStepTuple &stepTuple, + Operation *constantAnchor, + std::map, Value> &laneStartTableCache) { + if (std::optional affineMapping = getSourceLaneAffineMapping(stepTuple)) { + SourceLaneSelector selector; + selector.kind = SourceLaneSelector::Kind::Affine; + selector.affine = *affineMapping; + return selector; + } + + SmallVector tableValues = collectSourceLaneStarts(stepTuple); + std::vector cacheKey(tableValues.begin(), tableValues.end()); + auto cacheIt = laneStartTableCache.find(cacheKey); + if (cacheIt != laneStartTableCache.end()) { + SourceLaneSelector selector; + selector.kind = SourceLaneSelector::Kind::Table; + selector.table = cacheIt->second; + selector.tableValues = tableValues; + return selector; + } + + SmallVector tableValuesI64; + tableValuesI64.reserve(tableValues.size()); + for (uint32_t value : tableValues) + tableValuesI64.push_back(value); + Value table = createI64LookupTableConstant(rewriter, constantAnchor, tableValuesI64); + laneStartTableCache.emplace(std::move(cacheKey), table); + SourceLaneSelector selector; + selector.kind = SourceLaneSelector::Kind::Table; + selector.table = table; + selector.tableValues = tableValues; + return selector; +} + +static FailureOr buildSourceLaneStartForScheduledLane(OpBuilder &builder, + Location loc, + Value scheduledLane, + const SourceLaneSelector &selector, + Operation *constantAnchor) { + (void)constantAnchor; + if (selector.kind == SourceLaneSelector::Kind::Affine) { + if (selector.affine.baseLaneStart == 0 && selector.affine.laneCount == 1) + return scheduledLane; + AffineExpr d0 = builder.getAffineDimExpr(0); + AffineExpr expr = d0; + if (selector.affine.laneCount != 1) + expr = d0 * selector.affine.laneCount; + if (selector.affine.baseLaneStart != 0) + expr = expr + selector.affine.baseLaneStart; + return affine::AffineApplyOp::create( + builder, loc, AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr), ValueRange {scheduledLane}) + .getResult(); + } + + if (!selector.table) + return failure(); + Value sourceLaneStartI64 = + tensor::ExtractOp::create(builder, loc, selector.table, ValueRange {scheduledLane}).getResult(); + return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), sourceLaneStartI64).getResult(); +} + +static LogicalResult verifyPeftClassPlan(Operation *diagnosticAnchor, + const PeftClassPlan &peftClassPlan, + const MergeScheduleResult &schedule) { + if (peftClassPlan.cpus.empty()) + return diagnosticAnchor->emitOpError("PEFT materialization class has no CPUs"); + + SmallVector *> schedules; + for (size_t cpu : peftClassPlan.cpus) { + auto it = peftClassPlan.instancesByCpu.find(cpu); + if (it == peftClassPlan.instancesByCpu.end()) + return diagnosticAnchor->emitOpError("PEFT materialization class is missing a per-CPU schedule"); + schedules.push_back(&it->second); + for (const ComputeInstance &instance : it->second) + if (!schedule.computeToCpuSlotMap.count(instance)) + return diagnosticAnchor->emitOpError("PEFT materialization class references a compute instance without a scheduler position"); + } + + if (peftClassPlan.cpus.size() == 1) + return success(); + + auto emitNonIso = [&](size_t stepPosition) -> LogicalResult { + std::string cpus; + llvm::raw_string_ostream os(cpus); + llvm::interleaveComma(peftClassPlan.cpus, os, [&](size_t cpu) { os << cpu; }); + diagnosticAnchor->emitOpError("PEFT equivalence class has non-isomorphic per-CPU schedules") + << " class " << peftClassPlan.canonicalClassId << " cpus [" << os.str() << "] step " << stepPosition; + return failure(); + }; + + size_t tupleCount = schedules.front()->size(); + for (const SmallVector *cpuSchedule : schedules) + if (cpuSchedule->size() != tupleCount) + return emitNonIso(0); + + for (size_t stepPosition = 0; stepPosition < tupleCount; ++stepPosition) { + const ComputeInstance &reference = (*schedules.front())[stepPosition]; + bool refIsScalar = isa(reference.op); + for (size_t cpuIndex = 1; cpuIndex < schedules.size(); ++cpuIndex) { + const ComputeInstance &instance = (*schedules[cpuIndex])[stepPosition]; + if (instance.op != reference.op || instance.laneCount != reference.laneCount) + return emitNonIso(stepPosition); + if (isa(instance.op) != refIsScalar) + return emitNonIso(stepPosition); + } + } + + return success(); +} + +static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClassPlan, + const MergeScheduleResult &schedule) { + peftClassPlan.weights.clear(); + peftClassPlan.inputs.clear(); + peftClassPlan.resultTypes.clear(); + + if (peftClassPlan.cpus.size() == 1) { + size_t cpu = peftClassPlan.cpus.front(); + for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(cpu)) { + if (auto compute = dyn_cast(instance.op)) { + llvm::append_range(peftClassPlan.resultTypes, compute.getResultTypes()); + } else { + auto batch = cast(instance.op); + for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) { + auto spec = getBatchFragmentSpec(batch, resultIndex, instance.laneCount); + if (failed(spec)) + return failure(); + peftClassPlan.resultTypes.push_back(spec->resultType); + } + } + + for (Value weight : getComputeInstanceWeights(instance)) + appendUnique(peftClassPlan.weights, weight); + for (Value input : getComputeInstanceInputs(instance)) + if (!getProducerValueRef(input, &instance)) + appendUnique(peftClassPlan.inputs, input); + } + return success(); + } + + for (const ScheduledStepPlan &stepPlan : buildScheduledStepPlans(peftClassPlan)) { + const ComputeStepTuple &stepTuple = stepPlan.stepTuple; + const ComputeInstance &representative = stepTuple.instances.front(); + if (auto compute = dyn_cast(representative.op)) { + for (Type type : compute.getResultTypes()) { + auto tensorType = dyn_cast(type); + if (!tensorType || !tensorType.hasStaticShape()) + return compute.emitOpError("scheduled materialization only supports static ranked tensor scalar results"); + SmallVector shape; + shape.push_back(static_cast(peftClassPlan.cpus.size())); + llvm::append_range(shape, tensorType.getShape()); + peftClassPlan.resultTypes.push_back(RankedTensorType::get(shape, tensorType.getElementType())); + } + } else { + auto batch = cast(representative.op); + uint32_t totalLanes = static_cast(peftClassPlan.cpus.size()) * representative.laneCount; + for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) { + auto spec = getBatchFragmentSpec(batch, resultIndex, totalLanes); + if (failed(spec)) + return failure(); + peftClassPlan.resultTypes.push_back(spec->resultType); + } + } + + for (const ComputeInstance &instance : stepTuple.instances) { + for (Value weight : getComputeInstanceWeights(instance)) + appendUnique(peftClassPlan.weights, weight); + for (Value input : getComputeInstanceInputs(instance)) + if (!getProducerValueRef(input, &instance)) + appendUnique(peftClassPlan.inputs, input); + } + } + + return success(); +} + +static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper, SmallVectorImpl &yieldedValues) { + for (Operation &op : source.without_terminator()) + builder.clone(op, mapper); + auto yield = cast(source.getTerminator()); + for (Value output : yield.getOutputs()) + yieldedValues.push_back(mapper.lookup(output)); +} + +static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rewriter, + SpatComputeBatch batch, + const ComputeInstance &instance, + ValueRange scheduledWeights, + ValueRange scheduledInputs, + Block &block, + const MergeScheduleResult &schedule, + uint64_t &exchangeId, + SmallVectorImpl &yieldedValues) { + SmallVector initResults; + SmallVector fragmentSpecs; + for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) { + auto spec = getBatchFragmentSpec(batch, resultIndex, instance.laneCount); + if (failed(spec)) + return failure(); + fragmentSpecs.push_back(*spec); + auto empty = createEmptyTensorForType(rewriter, batch.getLoc(), spec->resultType); + if (failed(empty)) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + initResults.push_back(*empty); + } + + Value lower = getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart); + Value upper = getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart + instance.laneCount); + Value step = getOrCreateIndexConstant(rewriter, batch.getOperation(), 1); + auto loop = buildNormalizedScfFor( + rewriter, + batch.getLoc(), + lower, + upper, + step, + initResults, + [&](OpBuilder &builder, Location bodyLoc, Value originalLane, ValueRange iterArgs, SmallVectorImpl &yielded) -> LogicalResult { + + IRMapping mapper; + mapper.map(*batch.getLaneArgument(), originalLane); + for (auto [index, weight] : llvm::enumerate(batch.getWeights())) + mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight)); + for (auto [index, input] : llvm::enumerate(batch.getInputs())) { + Value mapped; + if (failed(mapSingleCpuInput(builder, + input.getLoc(), + input, + instance, + schedule, + scheduledInputs, + block, + scheduledWeights.size(), + ArrayRef {}, + exchangeId, + mapped))) + return failure(); + mapper.map(*batch.getInputArgument(index), mapped); + } + for (auto [index, outputArg] : llvm::enumerate(batch.getOutputs())) + (void)outputArg, mapper.map(*batch.getOutputArgument(index), iterArgs[index]); + + Block &source = batch.getBody().front(); + for (Operation &op : source.without_terminator()) + builder.clone(op, mapper); + + 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; + + SmallVector current(iterArgs.begin(), iterArgs.end()); + for (Operation &op : inParallel.getRegion().front()) { + auto insert = dyn_cast(&op); + if (!insert) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + auto oldDest = dyn_cast(insert.getDest()); + if (!oldDest || !outputIndexByArg.count(oldDest)) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"), failure(); + size_t resultIndex = outputIndexByArg.lookup(oldDest); + SmallVector offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper); + offsets[fragmentSpecs[resultIndex].laneDim] = localLane; + current[resultIndex] = tensor::InsertSliceOp::create(builder, + insert.getLoc(), + mapper.lookup(insert.getSource()), + current[resultIndex], + offsets, + remapMixedOffsets(insert.getMixedSizes(), mapper), + remapMixedOffsets(insert.getMixedStrides(), mapper)) + .getResult(); + } + llvm::append_range(yielded, current); + return success(); + }); + if (failed(loop)) + return failure(); + llvm::append_range(yieldedValues, loop->results); + return success(); +} + +static LogicalResult materializeSingleCpuPeftClass( + PatternRewriter &rewriter, + SpatScheduledCompute scheduled, + const PeftClassPlan &peftClassPlan, + const MergeScheduleResult &schedule, + DenseMap &graphComputeToBlockMap, + ScheduledMaterializationRecord &record, + uint64_t &exchangeId) { + size_t cpu = peftClassPlan.cpus.front(); + auto instancesIt = peftClassPlan.instancesByCpu.find(cpu); + assert(instancesIt != peftClassPlan.instancesByCpu.end() && "missing single-cpu schedule"); + const SmallVector &instances = instancesIt->second; + + SmallVector carriedKeys; + Block *block = nullptr; + for (auto [ordinal, instance] : llvm::enumerate(instances)) { + if (!block) + block = createScheduledComputeBlock(rewriter, scheduled, carriedKeys, instance.op->getLoc()); + + GraphComputeBlockKey key = getGraphComputeBlockKey(instance); + graphComputeToBlockMap[key] = block; + record.computeKeys.push_back(key); + record.blocks.push_back(block); + + rewriter.setInsertionPointToStart(block); + SmallVector yieldedValues; + if (auto compute = dyn_cast(instance.op)) { + IRMapping mapper; + for (auto [index, weight] : llvm::enumerate(compute.getWeights())) + mapper.map(*compute.getWeightArgument(index), getBlockOperand(*block, scheduled.getWeights(), weight)); + for (auto [index, input] : llvm::enumerate(compute.getInputs())) { + Value mapped; + if (failed(mapSingleCpuInput(rewriter, + input.getLoc(), + input, + instance, + schedule, + scheduled.getInputs(), + *block, + scheduled.getWeights().size(), + carriedKeys, + exchangeId, + mapped))) + return failure(); + mapper.map(*compute.getInputArgument(index), mapped); + } + cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues); + } else { + auto batch = cast(instance.op); + if (failed(materializeResultfulBatchChunkAsScalar(rewriter, + batch, + instance, + scheduled.getWeights(), + scheduled.getInputs(), + *block, + schedule, + exchangeId, + yieldedValues))) + return failure(); + } + + SmallVector currentKeys; + for (size_t index = 0; index < yieldedValues.size(); ++index) + currentKeys.push_back({instance, index}); + unsigned baseArgCount = getScheduledComputeResultArgBase(scheduled); + SmallVector blockYieldOperands; + bool hasNextBlock = ordinal + 1 < instances.size(); + if (hasNextBlock) { + SmallVector nextCarriedKeys(carriedKeys); + llvm::append_range(nextCarriedKeys, currentKeys); + Block *nextBlock = createScheduledComputeBlock(rewriter, scheduled, nextCarriedKeys, instance.op->getLoc()); + appendBlockYieldBaseAndCarriedOperands(*block, baseArgCount, carriedKeys.size(), blockYieldOperands); + llvm::append_range(blockYieldOperands, yieldedValues); + rewriter.setInsertionPointToEnd(block); + createBlockYield(rewriter, instance.op->getLoc(), blockYieldOperands, nextBlock); + carriedKeys = std::move(nextCarriedKeys); + block = nextBlock; + } else { + for (ProducerValueKey carried : carriedKeys) + blockYieldOperands.push_back(getScheduledComputeOutputArgument(*block, + scheduled.getWeights(), + scheduled.getInputs(), + carriedKeys, + carried)); + llvm::append_range(blockYieldOperands, yieldedValues); + rewriter.setInsertionPointToEnd(block); + createBlockYield(rewriter, instance.op->getLoc(), blockYieldOperands); + } + } + return success(); +} + +// Builds offsets for inserting one per-CPU local fragment into the +// scheduled_compute_batch output. The lane offset is in scheduled-output +// lane space, not local fragment lane space. +static SmallVector buildScheduledOutputInsertOffsets(OpBuilder &builder, + Location loc, + Value scheduledLane, + int64_t lanesPerScheduledLane, + RankedTensorType localFragmentType, + unsigned laneDim) { + SmallVector offsets; + Value scheduledOutputLane = scheduledLane; + if (lanesPerScheduledLane != 1) { + scheduledOutputLane = + affine::AffineApplyOp::create(builder, + loc, + AffineMap::get(/*dimCount=*/1, + /*symbolCount=*/0, + builder.getAffineDimExpr(0) * lanesPerScheduledLane), + ValueRange {scheduledLane}) + .getResult(); + } + for (unsigned dim = 0; dim < static_cast(localFragmentType.getRank()); ++dim) + offsets.push_back(dim == laneDim ? OpFoldResult(scheduledOutputLane) : OpFoldResult(builder.getIndexAttr(0))); + return offsets; +} + +static LogicalResult materializeMultiCpuPeftClass( + PatternRewriter &rewriter, + SpatScheduledComputeBatch scheduled, + const PeftClassPlan &peftClassPlan, + const MergeScheduleResult &schedule, + DenseMap &graphComputeToBlockMap, + ScheduledMaterializationRecord &record, + uint64_t &exchangeId) { + std::map, Value> laneStartTableCache; + ArrayRef stepPlans = record.stepPlans; + for (const ScheduledStepPlan &stepPlan : stepPlans) { + const ComputeStepTuple &stepTuple = stepPlan.stepTuple; + SourceLaneSelector sourceLaneSelector = + buildSourceLaneSelector(rewriter, stepTuple, scheduled.getOperation(), laneStartTableCache); + SmallVector blockArgTypes {rewriter.getIndexType()}; + SmallVector blockArgLocs {scheduled.getLoc()}; + for (Value weight : scheduled.getWeights()) { + blockArgTypes.push_back(weight.getType()); + blockArgLocs.push_back(weight.getLoc()); + } + for (Value input : scheduled.getInputs()) { + blockArgTypes.push_back(input.getType()); + blockArgLocs.push_back(input.getLoc()); + } + for (Type resultType : scheduled.getResultTypes()) { + blockArgTypes.push_back(resultType); + blockArgLocs.push_back(scheduled.getLoc()); + } + Block *block = rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), blockArgTypes, blockArgLocs); + for (const ComputeInstance &instance : stepTuple.instances) { + GraphComputeBlockKey key = getGraphComputeBlockKey(instance); + graphComputeToBlockMap[key] = block; + record.computeKeys.push_back(key); + record.blocks.push_back(block); + } + + rewriter.setInsertionPointToStart(block); + Value scheduledLane = block->getArgument(0); + const ComputeInstance &representative = stepTuple.instances.front(); + SmallVector finalLocalFragments; + SmallVector finalLaneDims; + + if (auto compute = dyn_cast(representative.op)) { + IRMapping mapper; + for (auto [index, weight] : llvm::enumerate(compute.getWeights())) + mapper.map(*compute.getWeightArgument(index), block->getArgument(1 + index)); + unsigned firstInputArg = 1 + scheduled.getWeights().size(); + for (auto [index, input] : llvm::enumerate(compute.getInputs())) { + Value mapped; + if (failed(mapMultiCpuTupleInput(rewriter, + input.getLoc(), + input, + stepTuple, + peftClassPlan, + schedule, + scheduled.getInputs(), + *block, + firstInputArg, + exchangeId, + mapped))) + return failure(); + mapper.map(*compute.getInputArgument(index), mapped); + } + SmallVector yieldedValues; + cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues); + for (Value yielded : yieldedValues) { + auto tensorType = dyn_cast(yielded.getType()); + if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0) + return compute.emitOpError("scheduled materialization only supports static ranked tensor scalar step results"); + SmallVector reassociation; + reassociation.push_back({0, 1}); + for (int64_t dim = 1; dim < tensorType.getRank(); ++dim) + reassociation.push_back({static_cast(dim + 1)}); + SmallVector expandedShape {1}; + llvm::append_range(expandedShape, tensorType.getShape()); + finalLocalFragments.push_back(tensor::ExpandShapeOp::create(rewriter, + scheduled.getLoc(), + RankedTensorType::get(expandedShape, tensorType.getElementType()), + yielded, + reassociation) + .getResult()); + finalLaneDims.push_back(0); + } + } else { + auto batch = cast(representative.op); + SmallVector localFragments; + SmallVector fragmentSpecs; + for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) { + auto spec = getBatchFragmentSpec(batch, resultIndex, representative.laneCount); + if (failed(spec)) + return failure(); + fragmentSpecs.push_back(*spec); + auto empty = createEmptyTensorForType(rewriter, batch.getLoc(), spec->resultType); + if (failed(empty)) + return failure(); + localFragments.push_back(*empty); + } + + Value lower = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 0); + Value upper = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), representative.laneCount); + Value step = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 1); + auto loop = buildNormalizedScfFor( + rewriter, + batch.getLoc(), + lower, + upper, + step, + localFragments, + [&](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, + AffineMap::get(/*dimCount=*/2, + /*symbolCount=*/0, + builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1)), + ValueRange {*sourceLaneStart, innerLane}) + .getResult(); + mapper.map(*batch.getLaneArgument(), sourceLane); + for (auto [index, weight] : llvm::enumerate(batch.getWeights())) + mapper.map(*batch.getWeightArgument(index), block->getArgument(1 + index)); + unsigned firstInputArg = 1 + scheduled.getWeights().size(); + for (auto [index, input] : llvm::enumerate(batch.getInputs())) { + Value mapped; + if (failed(mapMultiCpuTupleInput(builder, + input.getLoc(), + input, + stepTuple, + peftClassPlan, + schedule, + scheduled.getInputs(), + *block, + firstInputArg, + exchangeId, + mapped))) + return failure(); + mapper.map(*batch.getInputArgument(index), mapped); + } + for (unsigned index = 0; index < batch.getNumResults(); ++index) + mapper.map(*batch.getOutputArgument(index), iterArgs[index]); + for (Operation &op : batch.getBody().front().without_terminator()) + builder.clone(op, mapper); + + auto inParallel = dyn_cast(batch.getBody().front().getTerminator()); + if (!inParallel) + return batch.emitOpError("expected spat.in_parallel in resultful spat.graph_compute_batch"), failure(); + + DenseMap outputIndexByArg; + for (size_t index = 0; index < batch.getNumResults(); ++index) + outputIndexByArg[*batch.getOutputArgument(index)] = index; + + SmallVector current(iterArgs.begin(), iterArgs.end()); + for (Operation &op : inParallel.getRegion().front()) { + auto insert = dyn_cast(&op); + if (!insert) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"); + auto oldDest = dyn_cast(insert.getDest()); + if (!oldDest || !outputIndexByArg.count(oldDest)) + return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"), failure(); + size_t resultIndex = outputIndexByArg.lookup(oldDest); + SmallVector offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper); + offsets[fragmentSpecs[resultIndex].laneDim] = innerLane; + current[resultIndex] = tensor::InsertSliceOp::create(builder, + insert.getLoc(), + mapper.lookup(insert.getSource()), + current[resultIndex], + offsets, + remapMixedOffsets(insert.getMixedSizes(), mapper), + remapMixedOffsets(insert.getMixedStrides(), mapper)) + .getResult(); + } + llvm::append_range(yielded, current); + return success(); + }); + if (failed(loop)) + return failure(); + finalLocalFragments.assign(loop->results.begin(), loop->results.end()); + for (const BatchFragmentSpec &spec : fragmentSpecs) + finalLaneDims.push_back(spec.laneDim); + } + + auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc()); + rewriter.setInsertionPointToStart(&inParallel.getRegion().front()); + for (auto [resultIndex, localFragment] : llvm::enumerate(finalLocalFragments)) { + auto localFragmentType = cast(localFragment.getType()); + int64_t lanesPerScheduledLane = isa(representative.op) ? 1 : representative.laneCount; + SmallVector offsets = buildScheduledOutputInsertOffsets( + rewriter, + scheduled.getLoc(), + scheduledLane, + lanesPerScheduledLane, + localFragmentType, + finalLaneDims[resultIndex]); + SmallVector sizes; + SmallVector strides; + for (int64_t dim : localFragmentType.getShape()) { + sizes.push_back(rewriter.getIndexAttr(dim)); + strides.push_back(rewriter.getIndexAttr(1)); + } + tensor::ParallelInsertSliceOp::create( + rewriter, + scheduled.getLoc(), + localFragment, + block->getArgument(getScheduledBatchResultArgBase(scheduled) + stepPlan.resultOffset + resultIndex), + offsets, + sizes, + strides); + } + } + return success(); +} + + +} // namespace + + +FailureOr +materializeScheduledCompute(func::FuncOp funcOp, + const MergeScheduleResult &schedule, + PatternRewriter &rewriter) { + llvm::MapVector peftClassPlans; + for (const ComputeInstance &instance : schedule.dominanceOrderCompute) { + size_t cpu = schedule.computeToCpuMap.lookup(instance); + size_t canonicalPeftClassId = getCanonicalPeftClassId(cpu, schedule); + auto &peftClassPlan = peftClassPlans[canonicalPeftClassId]; + peftClassPlan.canonicalClassId = canonicalPeftClassId; + if (!llvm::is_contained(peftClassPlan.cpus, cpu)) + peftClassPlan.cpus.push_back(cpu); + peftClassPlan.instancesByCpu[cpu].push_back(instance); + } + for (auto &entry : peftClassPlans) { + PeftClassPlan &peftClassPlan = entry.second; + 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); + }); + if (failed(verifyPeftClassPlan(funcOp.getOperation(), peftClassPlan, schedule))) + return failure(); + if (failed(collectPeftClassOperandsAndResults(peftClassPlan, schedule))) + return failure(); + } + + Operation *insertionPoint = funcOp.getBody().front().getTerminator(); + DenseMap graphComputeToBlockMap; + DenseMap scheduledComputes; + DenseMap scheduledComputeBatches; + DenseMap classToRecordIndex; + std::vector materializedSchedules; + uint64_t exchangeId = 0; + + for (auto &entry : peftClassPlans) { + PeftClassPlan &peftClassPlan = entry.second; + rewriter.setInsertionPoint(insertionPoint); + + ScheduledMaterializationRecord record; + record.canonicalPeftClassId = peftClassPlan.canonicalClassId; + record.cpus = peftClassPlan.cpus; + record.stepPlans = buildScheduledStepPlans(peftClassPlan); + + if (peftClassPlan.cpus.size() == 1) { + auto scheduled = SpatScheduledCompute::create( + rewriter, + peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).front().op->getLoc(), + TypeRange(peftClassPlan.resultTypes), + peftClassPlan.weights, + peftClassPlan.inputs); + scheduled->setAttr(kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast(peftClassPlan.cpus.front()))); + scheduled->setAttr("scheduled.peft_cpus", rewriter.getDenseI64ArrayAttr(toI64Array(peftClassPlan.cpus))); + SmallVector stepSources; + SmallVector sourceLaneSelectors; + SmallVector stepResultOffsets; + SmallVector stepResultCounts; + SmallVector sourceLaneStarts; + SmallVector sourceLaneCounts; + size_t resultOffset = 0; + for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front())) { + stepSources.push_back(rewriter.getStringAttr(getInstanceName(instance))); + sourceLaneSelectors.push_back(rewriter.getStringAttr(isa(instance.op) ? "scalar" : "affine")); + size_t resultCount = getComputeInstanceResultValueCount(instance); + stepResultOffsets.push_back(static_cast(resultOffset)); + stepResultCounts.push_back(static_cast(resultCount)); + resultOffset += resultCount; + if (isa(instance.op)) { + sourceLaneStarts.push_back(0); + sourceLaneCounts.push_back(0); + } else { + sourceLaneStarts.push_back(instance.laneStart); + sourceLaneCounts.push_back(instance.laneCount); + } + } + scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources)); + scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(stepResultOffsets)); + scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(stepResultCounts)); + scheduled->setAttr("scheduled.source_lane_starts", rewriter.getDenseI64ArrayAttr(sourceLaneStarts)); + scheduled->setAttr("scheduled.source_lane_counts", rewriter.getDenseI64ArrayAttr(sourceLaneCounts)); + scheduled->setAttr("scheduled.source_lane_selector", rewriter.getArrayAttr(sourceLaneSelectors)); + record.scheduledOp = scheduled.getOperation(); + scheduledComputes[peftClassPlan.canonicalClassId] = scheduled; + } else { + auto scheduled = SpatScheduledComputeBatch::create(rewriter, + peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).front().op->getLoc(), + TypeRange(peftClassPlan.resultTypes), + rewriter.getI32IntegerAttr(static_cast(peftClassPlan.cpus.size())), + peftClassPlan.weights, + peftClassPlan.inputs); + scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(toI32Array(peftClassPlan.cpus))); + scheduled->setAttr("scheduled.peft_cpus", rewriter.getDenseI64ArrayAttr(toI64Array(peftClassPlan.cpus))); + SmallVector stepSources; + SmallVector sourceLaneSelectors; + SmallVector resultOffsets; + SmallVector resultCounts; + SmallVector sourceLaneStarts; + SmallVector sourceLaneCounts; + for (const ScheduledStepPlan &stepPlan : record.stepPlans) { + stepSources.push_back(rewriter.getStringAttr(getInstanceName(stepPlan.stepTuple.instances.front()))); + sourceLaneSelectors.push_back(rewriter.getStringAttr(usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table")); + resultOffsets.push_back(static_cast(stepPlan.resultOffset)); + resultCounts.push_back(static_cast(stepPlan.resultCount)); + for (const ComputeInstance &instance : stepPlan.stepTuple.instances) { + sourceLaneStarts.push_back(instance.laneStart); + sourceLaneCounts.push_back(instance.laneCount); + } + } + RankedTensorType sourceLaneTableType = RankedTensorType::get( + {static_cast(record.stepPlans.size()), static_cast(peftClassPlan.cpus.size())}, + rewriter.getI64Type()); + scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources)); + scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(resultOffsets)); + scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(resultCounts)); + scheduled->setAttr("scheduled.source_lane_starts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef(sourceLaneStarts))); + scheduled->setAttr("scheduled.source_lane_counts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef(sourceLaneCounts))); + scheduled->setAttr("scheduled.source_lane_selector", rewriter.getArrayAttr(sourceLaneSelectors)); + record.scheduledOp = scheduled.getOperation(); + scheduledComputeBatches[peftClassPlan.canonicalClassId] = scheduled; + } + + classToRecordIndex[peftClassPlan.canonicalClassId] = materializedSchedules.size(); + materializedSchedules.push_back(std::move(record)); + } + + for (auto &entry : peftClassPlans) { + PeftClassPlan &peftClassPlan = entry.second; + ScheduledMaterializationRecord &record = + materializedSchedules[classToRecordIndex.lookup(peftClassPlan.canonicalClassId)]; + if (peftClassPlan.cpus.size() == 1) { + if (failed(materializeSingleCpuPeftClass(rewriter, + scheduledComputes.lookup(peftClassPlan.canonicalClassId), + peftClassPlan, + schedule, + graphComputeToBlockMap, + record, + exchangeId))) + return failure(); + } else { + if (failed(materializeMultiCpuPeftClass(rewriter, + scheduledComputeBatches.lookup(peftClassPlan.canonicalClassId), + peftClassPlan, + schedule, + graphComputeToBlockMap, + record, + exchangeId))) + return failure(); + } + } + + return ScheduledComputeMaterializationResult {std::move(peftClassPlans), std::move(materializedSchedules), std::move(graphComputeToBlockMap)}; +} + + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp new file mode 100644 index 0000000..b594a88 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "ScheduledComputePlan.hpp" + +namespace onnx_mlir { +namespace spatial { + +struct ScheduledComputeMaterializationResult { + llvm::MapVector peftClassPlans; + std::vector materializedSchedules; + DenseMap graphComputeToBlockMap; +}; + +FailureOr +materializeScheduledCompute(func::FuncOp funcOp, + const MergeScheduleResult &schedule, + PatternRewriter &rewriter); + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp new file mode 100644 index 0000000..cb19067 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp @@ -0,0 +1,337 @@ +#pragma once + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/AsmState.h" +#include "mlir/IR/IRMapping.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include +#include + +#include "Scheduling/ComputeInstanceUtils.hpp" +#include "Scheduling/MergeSchedulingAnalysis.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +namespace onnx_mlir { +namespace spatial { + +using namespace mlir; + +struct ProducerValueKey { + ComputeInstance instance; + size_t resultIndex = 0; + + bool operator==(const ProducerValueKey &other) const { + return instance == other.instance && resultIndex == other.resultIndex; + } +}; + +struct GraphComputeBlockKey { + Operation *op = nullptr; + uint32_t laneStart = 0; + uint32_t laneCount = 1; + + bool operator==(const GraphComputeBlockKey &other) const { + return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount; + } +}; + +struct PeftClassPlan { + size_t canonicalClassId = 0; + SmallVector cpus; + llvm::MapVector> instancesByCpu; + + SmallVector weights; + SmallVector inputs; + SmallVector resultTypes; +}; + +struct ComputeStepTuple { + SmallVector instances; +}; + +struct ScheduledStepPlan { + ComputeStepTuple stepTuple; + size_t stepIndex = 0; + size_t resultOffset = 0; + size_t resultCount = 0; +}; + +struct SourceLaneAffineMapping { + uint32_t baseLaneStart = 0; + uint32_t laneCount = 1; +}; + +struct SourceLaneSelector { + enum class Kind { Affine, Table } kind = Kind::Affine; + SourceLaneAffineMapping affine; + Value table; + SmallVector tableValues; +}; + +struct DeferredEndpoint { + ComputeInstance instance; + size_t cpu = 0; + size_t canonicalClassId = 0; + uint32_t laneStart = 0; + uint32_t laneCount = 1; +}; + +struct DeferredTransferMetadata { + SmallVector producers; + SmallVector consumers; + bool isBatched = false; + SmallVector sourceOperandForScheduledLane; + bool multiSourcePayload = false; +}; + +struct ScheduledMaterializationRecord { + Operation *scheduledOp = nullptr; + size_t canonicalPeftClassId = 0; + SmallVector cpus; + SmallVector stepPlans; + SmallVector computeKeys; + SmallVector blocks; +}; + +struct ScheduledComputePrintContext { + mlir::AsmState asmState; + explicit ScheduledComputePrintContext(ModuleOp module, const OpPrintingFlags &flags = OpPrintingFlags()) + : asmState(module.getOperation(), flags) {} +}; + +inline GraphComputeBlockKey getGraphComputeBlockKey(const ComputeInstance &instance) { + return {instance.op, instance.laneStart, instance.laneCount}; +} + +inline SmallVector getPeftClassCpus(size_t cpu, const MergeScheduleResult &schedule) { + llvm::SmallDenseSet seen; + SmallVector cpus; + auto append = [&](size_t value) { + if (seen.insert(value).second) + cpus.push_back(value); + }; + append(cpu); + if (auto it = schedule.equivalentClass.find(cpu); it != schedule.equivalentClass.end()) + for (size_t equivalentCpu : it->second) + append(equivalentCpu); + llvm::sort(cpus); + return cpus; +} + +inline size_t getCanonicalPeftClassId(size_t cpu, const MergeScheduleResult &schedule) { + SmallVector cpus = getPeftClassCpus(cpu, schedule); + return cpus.empty() ? cpu : cpus.front(); +} + +inline size_t getScheduledCpuForComputeInstance(const ComputeInstance &instance, const MergeScheduleResult &schedule) { + if (auto it = schedule.computeToCpuMap.find(instance); it != schedule.computeToCpuMap.end()) + return it->second; + + auto batch = dyn_cast(instance.op); + assert(batch && instance.laneCount != 0 && "missing scheduled CPU for non-batch compute instance"); + assert(instance.laneStart < static_cast(batch.getLaneCount()) && "batch lane start out of range"); + ComputeInstance chunk = getBatchChunkForLane(batch, instance.laneStart); + auto it = schedule.computeToCpuMap.find(chunk); + assert(it != schedule.computeToCpuMap.end() && "missing scheduled CPU for batch chunk"); + return it->second; +} + +inline std::string getInstanceName(const ComputeInstance &instance) { + return llvm::formatv("{0}[lanes={1}:{2}]", + instance.op->getName().getStringRef(), + instance.laneStart, + instance.laneStart + instance.laneCount) + .str(); +} + +inline SmallVector toI64Array(ArrayRef values) { + SmallVector converted; + converted.reserve(values.size()); + for (size_t value : values) + converted.push_back(static_cast(value)); + return converted; +} + +inline SmallVector toI32Array(ArrayRef values) { + SmallVector converted; + converted.reserve(values.size()); + for (size_t value : values) + converted.push_back(static_cast(value)); + return converted; +} + +inline unsigned getScheduledBatchResultArgBase(SpatScheduledComputeBatch scheduled) { + unsigned weightArgBase = 1; + unsigned inputArgBase = weightArgBase + scheduled.getWeights().size(); + return inputArgBase + scheduled.getInputs().size(); +} + +inline SmallVector collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) { + SmallVector keys; + for (Operation &op : funcOp.getOps()) { + if (auto compute = dyn_cast(&op)) + keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1})); + else if (auto batch = dyn_cast(&op)) + for (ComputeInstance chunk : getBatchChunksForRange(batch, 0, static_cast(batch.getLaneCount()))) + keys.push_back(getGraphComputeBlockKey(chunk)); + } + return keys; +} + +inline size_t countPeftEquivalenceClasses(const MergeScheduleResult &schedule) { + llvm::SmallDenseSet classes; + for (const ComputeInstance &instance : schedule.dominanceOrderCompute) { + auto cpuIt = schedule.computeToCpuMap.find(instance); + if (cpuIt == schedule.computeToCpuMap.end()) + continue; + classes.insert(getCanonicalPeftClassId(cpuIt->second, schedule)); + } + return classes.size(); +} + +inline SmallVector buildComputeStepTuples(const PeftClassPlan &peftClassPlan) { + SmallVector stepTuples; + if (peftClassPlan.cpus.empty()) + return stepTuples; + size_t stepCount = peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).size(); + stepTuples.resize(stepCount); + for (size_t stepIndex = 0; stepIndex < stepCount; ++stepIndex) + for (size_t cpu : peftClassPlan.cpus) + stepTuples[stepIndex].instances.push_back(peftClassPlan.instancesByCpu.lookup(cpu)[stepIndex]); + return stepTuples; +} + +inline size_t getComputeInstanceResultValueCount(const ComputeInstance &instance) { + return TypeSwitch(instance.op) + .Case([](SpatCompute compute) { return compute.getNumResults(); }) + .Case([](SpatComputeBatch batch) { return batch.getNumResults(); }) + .Default([](Operation *) -> size_t { + llvm_unreachable("expected graph compute or graph compute batch"); + }); +} + +inline SmallVector buildScheduledStepPlans(const PeftClassPlan &peftClassPlan) { + SmallVector stepPlans; + size_t resultOffset = 0; + for (auto [stepIndex, stepTuple] : llvm::enumerate(buildComputeStepTuples(peftClassPlan))) { + assert(!stepTuple.instances.empty() && "expected non-empty step tuple"); + size_t resultCount = getComputeInstanceResultValueCount(stepTuple.instances.front()); + stepPlans.push_back(ScheduledStepPlan {stepTuple, stepIndex, resultOffset, resultCount}); + resultOffset += resultCount; + } + return stepPlans; +} + +inline bool valueTransitivelyDependsOn(Value value, Value dependency) { + SmallVector worklist {value}; + DenseSet visited; + while (!worklist.empty()) { + Value current = worklist.pop_back_val(); + if (!visited.insert(current).second) + continue; + if (current == dependency) + return true; + Operation *def = current.getDefiningOp(); + if (!def) + continue; + for (Value operand : def->getOperands()) + worklist.push_back(operand); + } + return false; +} + +inline std::optional getSourceLaneAffineMapping(const ComputeStepTuple &stepTuple) { + if (stepTuple.instances.empty()) + return std::nullopt; + const ComputeInstance &reference = stepTuple.instances.front(); + for (const ComputeInstance &instance : stepTuple.instances) { + if (instance.op != reference.op || instance.laneCount != reference.laneCount) + return std::nullopt; + } + + for (size_t index = 0; index < stepTuple.instances.size(); ++index) { + uint32_t expectedLaneStart = reference.laneStart + static_cast(index) * reference.laneCount; + if (stepTuple.instances[index].laneStart != expectedLaneStart) + return std::nullopt; + } + + return SourceLaneAffineMapping {reference.laneStart, reference.laneCount}; +} + +inline bool usesAffineSourceLaneMapping(const ComputeStepTuple &stepTuple) { + return getSourceLaneAffineMapping(stepTuple).has_value(); +} + +inline SmallVector collectSourceLaneStarts(const ComputeStepTuple &stepTuple) { + SmallVector sourceLaneStarts; + sourceLaneStarts.reserve(stepTuple.instances.size()); + for (const ComputeInstance &instance : stepTuple.instances) + sourceLaneStarts.push_back(instance.laneStart); + return sourceLaneStarts; +} + +inline Value createI64LookupTableConstant(OpBuilder &builder, Operation *constantAnchor, ArrayRef values) { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(constantAnchor); + RankedTensorType tableType = RankedTensorType::get({static_cast(values.size())}, builder.getI64Type()); + DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values); + return arith::ConstantOp::create(builder, constantAnchor->getLoc(), tableType, tableAttr).getResult(); +} + +inline FailureOr createEmptyTensorForType(OpBuilder &builder, Location loc, Type type) { + auto tensorType = dyn_cast(type); + if (!tensorType || !tensorType.hasStaticShape()) + return failure(); + return tensor::EmptyOp::create(builder, loc, tensorType.getShape(), tensorType.getElementType()).getResult(); +} + +} // namespace spatial +} // namespace onnx_mlir + +namespace llvm { +template <> +struct DenseMapInfo { + static onnx_mlir::spatial::ProducerValueKey getEmptyKey() { + return {DenseMapInfo::getEmptyKey(), std::numeric_limits::max()}; + } + static onnx_mlir::spatial::ProducerValueKey getTombstoneKey() { + return {DenseMapInfo::getTombstoneKey(), std::numeric_limits::max()}; + } + static unsigned getHashValue(const onnx_mlir::spatial::ProducerValueKey &key) { + return hash_combine(DenseMapInfo::getHashValue(key.instance), key.resultIndex); + } + static bool isEqual(const onnx_mlir::spatial::ProducerValueKey &lhs, + const onnx_mlir::spatial::ProducerValueKey &rhs) { + return lhs == rhs; + } +}; + +template <> +struct DenseMapInfo { + static onnx_mlir::spatial::GraphComputeBlockKey getEmptyKey() { + return {DenseMapInfo::getEmptyKey(), UINT32_MAX, UINT32_MAX}; + } + static onnx_mlir::spatial::GraphComputeBlockKey getTombstoneKey() { + return {DenseMapInfo::getTombstoneKey(), UINT32_MAX, UINT32_MAX - 1}; + } + static unsigned getHashValue(const onnx_mlir::spatial::GraphComputeBlockKey &key) { + return hash_combine(key.op, key.laneStart, key.laneCount); + } + static bool isEqual(const onnx_mlir::spatial::GraphComputeBlockKey &lhs, + const onnx_mlir::spatial::GraphComputeBlockKey &rhs) { + return lhs == rhs; + } +}; +} // namespace llvm diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp new file mode 100644 index 0000000..f7ba008 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp @@ -0,0 +1,312 @@ +#include "ScheduledComputeReport.hpp" + +#include "llvm/Support/raw_os_ostream.h" + +#include + +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" + +namespace onnx_mlir { +namespace spatial { + +using namespace mlir; +namespace { + +static std::string formatValueLabel(Value value, AsmState &asmState) { + std::string storage; + llvm::raw_string_ostream os(storage); + value.printAsOperand(os, asmState); + return storage; +} + +static std::string formatOperationLabel(Operation *op, AsmState &asmState) { + if (op->getNumResults() == 0) + return op->getName().getStringRef().str(); + std::string storage; + llvm::raw_string_ostream os(storage); + llvm::interleaveComma(op->getResults(), os, [&](Value result) { os << formatValueLabel(result, asmState); }); + return os.str(); +} + +static std::string formatGraphComputeBlockKey(const GraphComputeBlockKey &key, AsmState &asmState) { + return llvm::formatv("{0} {1}", formatOperationLabel(key.op, asmState), key.op->getName().getStringRef()).str(); +} + +static std::string formatComputeInstanceForReport(const ComputeInstance &instance, AsmState &asmState) { + std::string opLabel = formatGraphComputeBlockKey(getGraphComputeBlockKey(instance), asmState); + if (isa(instance.op)) + return opLabel; + return llvm::formatv("{0} sourceLanes [{1}:{2}]", + opLabel, + instance.laneStart, + instance.laneStart + instance.laneCount) + .str(); +} + +template +static void printIndexedList(raw_ostream &os, ArrayRef values) { + os << "["; + llvm::interleaveComma(llvm::enumerate(values), os, [&](auto indexedValue) { + os << indexedValue.index() << ":" << indexedValue.value(); + }); + os << "]"; +} + +struct PeftMaterializationReportSummary { + size_t scalarGraphCompute = 0; + size_t graphComputeBatchOps = 0; + size_t scalarGraphComputeInstances = 0; + size_t graphComputeBatchInstances = 0; + size_t peftClasses = 0; + size_t singleCpuClasses = 0; + size_t multiCpuClasses = 0; + size_t scheduledCompute = 0; + size_t scheduledComputeBatch = 0; + size_t deferredCommunication = 0; + size_t deferredCommunicationScalarMetadata = 0; + size_t deferredCommunicationScheduledLaneMetadata = 0; + size_t deferredCommunicationMultiSourcePayloads = 0; +}; + +static PeftMaterializationReportSummary buildPeftMaterializationReportSummary( + func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + ArrayRef materializedSchedules) { + PeftMaterializationReportSummary summary; + for (Operation &op : funcOp.getOps()) { + if (isa(op)) + summary.scalarGraphCompute++; + else if (isa(op)) { + summary.graphComputeBatchOps++; + } + } + for (const ComputeInstance &instance : schedule.dominanceOrderCompute) + (isa(instance.op) ? summary.scalarGraphComputeInstances : summary.graphComputeBatchInstances)++; + summary.peftClasses = peftClassPlans.size(); + for (const auto &entry : peftClassPlans) + (entry.second.cpus.size() == 1 ? summary.singleCpuClasses : summary.multiCpuClasses)++; + for (const ScheduledMaterializationRecord &record : materializedSchedules) { + if (isa(record.scheduledOp)) + summary.scheduledCompute++; + else + summary.scheduledComputeBatch++; + } + 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++; + }); + return summary; +} + +} // namespace + +LogicalResult verifyPeftMaterializationReportSummary(func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + ArrayRef materializedSchedules) { + PeftMaterializationReportSummary summary = + buildPeftMaterializationReportSummary(funcOp, schedule, peftClassPlans, materializedSchedules); + pim::CappedDiagnosticReporter diagnostics; + if (summary.peftClasses != peftClassPlans.size()) + diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check report PEFT total " << summary.peftClasses + << " does not match classes.size() " << peftClassPlans.size(); + }); + if (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances != schedule.dominanceOrderCompute.size()) + diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check report total compute instances " + << (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances) + << " does not match schedule size " << schedule.dominanceOrderCompute.size(); + }); + if (summary.scheduledCompute + summary.scheduledComputeBatch != materializedSchedules.size()) + diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check report scheduled total " + << (summary.scheduledCompute + summary.scheduledComputeBatch) + << " does not match materialized scheduled ops " << materializedSchedules.size(); + }); + diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial report verification failed"); + return success(!diagnostics.hasFailure()); +} + +namespace { + +static std::string formatStepResultRange(size_t resultOffset, size_t resultCount) { + if (resultCount == 1) + return llvm::formatv("result[{0}]", resultOffset).str(); + return llvm::formatv("result[{0}:{1}]", resultOffset, resultOffset + resultCount).str(); +} + +static void printMultiSourceDeferredInputs(raw_ostream &os, Block &block) { + unsigned deferredInputIndex = 0; + for (Operation &op : block.getOperations()) { + auto transfer = dyn_cast(&op); + if (!transfer) + continue; + auto multiSourcePayload = transfer->getAttrOfType("multi_source_payload"); + auto sourceOperandForScheduledLane = + transfer->getAttrOfType("source_operand_for_scheduled_lane"); + if (multiSourcePayload && multiSourcePayload.getValue() && sourceOperandForScheduledLane) { + SmallVector sourceOperandIndexes; + for (int64_t sourceOperandIndex : sourceOperandForScheduledLane.asArrayRef()) + sourceOperandIndexes.push_back(static_cast(sourceOperandIndex)); + os << " deferred input " << deferredInputIndex << ": multi-source uniqueSources=" + << transfer.getSources().size() << " sourceOperandForScheduledLane="; + printIndexedList(os, ArrayRef(sourceOperandIndexes)); + os << "\n"; + } + deferredInputIndex++; + } +} + +static void dumpPeftMaterializationReport(ModuleOp moduleOp, + func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + ArrayRef materializedSchedules, + ScheduledComputePrintContext &printContext) { + std::fstream file = openDialectDumpFileWithExtension("spatial2_scheduled_no_comm", "/reports", "txt"); + if (!file.is_open()) + return; + + llvm::raw_os_ostream os(file); + AsmState &asmState = printContext.asmState; + PeftMaterializationReportSummary summary = + buildPeftMaterializationReportSummary(funcOp, schedule, peftClassPlans, materializedSchedules); + + os << "Summary\n"; + os << "=======\n"; + os << "Graph computes:\n"; + os << " total: " << (summary.scalarGraphCompute + summary.graphComputeBatchOps) << "\n"; + os << " scalar graph_compute: " << summary.scalarGraphCompute << "\n"; + os << " graph_compute_batch: " << summary.graphComputeBatchOps << "\n"; + os << "Compute instances:\n"; + os << " total: " << (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances) << "\n"; + os << " scalar graph_compute instances: " << summary.scalarGraphComputeInstances << "\n"; + os << " graph_compute_batch instances: " << summary.graphComputeBatchInstances << "\n"; + os << "PEFT classes:\n"; + os << " total: " << summary.peftClasses << "\n"; + os << " single-cpu: " << summary.singleCpuClasses << "\n"; + os << " multi-cpu: " << summary.multiCpuClasses << "\n"; + os << "Scheduled ops:\n"; + os << " total: " << (summary.scheduledCompute + summary.scheduledComputeBatch) << "\n"; + os << " scheduled_compute: " << summary.scheduledCompute << "\n"; + 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"; + os << "============\n"; + for (const auto &entry : peftClassPlans) { + const PeftClassPlan &peftClassPlan = entry.second; + os << "C" << peftClassPlan.canonicalClassId << " " + << (peftClassPlan.cpus.size() == 1 ? "single-cpu" : "multi-cpu") << " PEFT class\n"; + if (peftClassPlan.cpus.size() == 1) { + size_t cpu = peftClassPlan.cpus.front(); + os << " cpu: " << cpu << "\n"; + os << " steps: " << peftClassPlan.instancesByCpu.lookup(cpu).size() << "\n"; + for (auto [stepIndex, instance] : llvm::enumerate(peftClassPlan.instancesByCpu.lookup(cpu))) + os << " step " << stepIndex << ": " << formatComputeInstanceForReport(instance, asmState) << "\n"; + } else { + os << " scheduled lanes: " << peftClassPlan.cpus.size() << "\n"; + os << " steps: " << peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).size() << "\n"; + os << " cpus by scheduled lane:\n"; + os << " "; + printIndexedList(os, ArrayRef(peftClassPlan.cpus)); + os << "\n"; + os << " step sources:\n"; + for (auto [stepIndex, stepTuple] : llvm::enumerate(buildComputeStepTuples(peftClassPlan))) + os << " step " << stepIndex << ": " + << formatGraphComputeBlockKey(getGraphComputeBlockKey(stepTuple.instances.front()), asmState) << "\n"; + } + os << "\n"; + } + + os << "Materialized Scheduled Ops\n"; + os << "=========================\n"; + for (const ScheduledMaterializationRecord &record : materializedSchedules) { + os << "C" << record.canonicalPeftClassId << " -> " << formatOperationLabel(record.scheduledOp, asmState) << " " + << record.scheduledOp->getName().getStringRef() << "\n"; + os << " kind: " + << (isa(record.scheduledOp) ? "single-cpu scheduled_compute" + : "multi-cpu scheduled_compute_batch") + << "\n"; + if (isa(record.scheduledOp)) + os << " cpu: " << record.cpus.front() << "\n"; + else + os << " scheduled lanes: " << record.cpus.size() << "\n"; + os << " results: " << record.scheduledOp->getNumResults() << "\n"; + os << " steps: " + << (isa(record.scheduledOp) + ? peftClassPlans.lookup(record.canonicalPeftClassId).instancesByCpu.lookup(record.cpus.front()).size() + : record.stepPlans.size()) + << "\n"; + if (isa(record.scheduledOp)) { + os << " cpus by scheduled lane:\n"; + os << " "; + printIndexedList(os, ArrayRef(record.cpus)); + os << "\n\n"; + } + if (isa(record.scheduledOp)) { + const PeftClassPlan &peftClassPlan = peftClassPlans.lookup(record.canonicalPeftClassId); + size_t cpu = peftClassPlan.cpus.front(); + size_t resultOffset = 0; + for (auto [stepIndex, instance] : llvm::enumerate(peftClassPlan.instancesByCpu.lookup(cpu))) { + size_t resultCount = getComputeInstanceResultValueCount(instance); + os << " step " << stepIndex << " " << formatStepResultRange(resultOffset, resultCount) << " " + << formatComputeInstanceForReport(instance, asmState) << "\n"; + resultOffset += resultCount; + } + } else { + auto scheduledBatch = cast(record.scheduledOp); + for (auto [stepIndex, stepPlan] : llvm::enumerate(record.stepPlans)) { + const ComputeInstance &representative = stepPlan.stepTuple.instances.front(); + SmallVector sourceLaneStarts = collectSourceLaneStarts(stepPlan.stepTuple); + os << " step " << stepIndex << " " << formatStepResultRange(stepPlan.resultOffset, stepPlan.resultCount) << " " + << formatGraphComputeBlockKey(getGraphComputeBlockKey(representative), asmState) + << " lanesPerScheduledLane=" << representative.laneCount << " sourceLaneSelector=" + << (usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table") << "\n"; + os << " source lanes by scheduled lane:\n"; + os << " "; + printIndexedList(os, ArrayRef(sourceLaneStarts)); + os << "\n"; + Block &stepBlock = *std::next(scheduledBatch.getBody().begin(), stepIndex); + printMultiSourceDeferredInputs(os, stepBlock); + } + } + os << "\n"; + } +} + + +} // namespace + +void dumpScheduledComputeReportAndModule(ModuleOp moduleOp, + func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + ArrayRef materializedSchedules) { + OpPrintingFlags flags; + flags.elideLargeElementsAttrs().enableDebugInfo(false, false).assumeVerified(); + ScheduledComputePrintContext printContext(moduleOp, flags); + dumpPeftMaterializationReport(moduleOp, funcOp, schedule, peftClassPlans, materializedSchedules, printContext); + + std::fstream file = openDialectDumpFileWithExtension("spatial2_scheduled_no_comm", "/dialects", "mlir"); + if (!file.is_open()) + return; + llvm::raw_os_ostream os(file); + moduleOp.getOperation()->print(os, printContext.asmState); + os.flush(); +} + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp new file mode 100644 index 0000000..84d0fb2 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp @@ -0,0 +1,21 @@ +#pragma once + +#include "ScheduledComputeMaterialization.hpp" + +namespace onnx_mlir { +namespace spatial { + +LogicalResult verifyPeftMaterializationReportSummary( + func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + ArrayRef materializedSchedules); + +void dumpScheduledComputeReportAndModule(ModuleOp moduleOp, + func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + ArrayRef materializedSchedules); + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp new file mode 100644 index 0000000..455deb5 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp @@ -0,0 +1,316 @@ +#include "ScheduledComputeVerification.hpp" + +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" + +namespace onnx_mlir { +namespace spatial { + +using namespace mlir; + +LogicalResult verifyMaterializedScheduleMapping( + func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + const DenseMap &graphComputeToBlockMap, + ArrayRef materializedSchedules) { + pim::CappedDiagnosticReporter diagnostics; + size_t expectedClassCount = countPeftEquivalenceClasses(schedule); + if (expectedClassCount != materializedSchedules.size()) { + diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check expected " << expectedClassCount + << " PEFT equivalence classes but materialized " << materializedSchedules.size() + << " scheduled computes"; + }); + } + + llvm::SmallDenseSet seenClasses; + for (const ScheduledMaterializationRecord &record : materializedSchedules) { + if (!seenClasses.insert(record.canonicalPeftClassId).second) { + diagnostics.report(record.scheduledOp, [&](Operation *illegalOp) { + illegalOp->emitOpError("phase-check multiple scheduled ops own the same PEFT equivalence class"); + }); + } + if (!peftClassPlans.count(record.canonicalPeftClassId)) { + diagnostics.report(record.scheduledOp, [&](Operation *illegalOp) { + illegalOp->emitOpError("phase-check scheduled op refers to a missing PEFT materialization class"); + }); + } + } + + for (const auto &entry : peftClassPlans) { + if (!seenClasses.count(entry.first)) { + diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check PEFT equivalence class " << entry.first + << " was not materialized by any scheduled op"; + }); + } + } + + for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) { + if (graphComputeToBlockMap.count(key)) + continue; + diagnostics.report(key.op, [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check graph compute is missing a scheduled MLIR block mapping for lanes [" + << key.laneStart << ":" << (key.laneStart + key.laneCount) << "]"; + }); + } + + for (const auto &entry : graphComputeToBlockMap) { + Block *block = entry.second; + if (!block || !isa(block->getParentOp())) { + diagnostics.report(entry.first.op, [&](Operation *illegalOp) { + illegalOp->emitOpError("phase-check graph compute block mapping does not target a scheduled compute block"); + }); + } + } + + if (graphComputeToBlockMap.size() != collectExpectedGraphComputeBlockKeys(funcOp).size()) { + diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check expected " + << collectExpectedGraphComputeBlockKeys(funcOp).size() + << " graph compute block mappings but saw " << graphComputeToBlockMap.size(); + }); + } + + diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial materialization verification failed"); + return success(!diagnostics.hasFailure()); +} + +LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) { + pim::CappedDiagnosticReporter diagnostics; + funcOp.walk([&](SpatDeferredCommunicationOp transfer) { + for (Value source : transfer.getSources()) { + if (isa_and_nonnull(source.getDefiningOp())) { + diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError("phase-check deferred communication source operand must remain an original graph SSA producer value"); + }); + } + } + + if (auto scheduled = dyn_cast(transfer->getParentOp())) { + auto batchedAttr = transfer->getAttrOfType("batched"); + auto sourceOperandForScheduledLane = + transfer->getAttrOfType("source_operand_for_scheduled_lane"); + auto multiSourcePayload = transfer->getAttrOfType("multi_source_payload"); + auto sourceCpus = transfer->getAttrOfType("source_cpus"); + auto sourceClasses = transfer->getAttrOfType("source_classes"); + auto sourceLaneRanges = transfer->getAttrOfType("source_lane_ranges"); + auto targetCpus = transfer->getAttrOfType("target_cpus"); + auto targetClasses = transfer->getAttrOfType("target_classes"); + auto targetLaneRanges = transfer->getAttrOfType("target_lane_ranges"); + int64_t laneCount = scheduled.getLaneCount(); + auto hasExpectedSize = [&](DenseI64ArrayAttr attr, int64_t expected, StringRef name) { + if (!attr || attr.size() != expected) + diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check scheduled batch deferred communication requires " << name + << " length " << expected; + }); + }; + if (!batchedAttr || !batchedAttr.getValue()) + diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError("phase-check scheduled batch deferred communication requires batched metadata"); + }); + hasExpectedSize(sourceCpus, laneCount, "source_cpus"); + hasExpectedSize(sourceClasses, laneCount, "source_classes"); + hasExpectedSize(targetCpus, laneCount, "target_cpus"); + hasExpectedSize(targetClasses, laneCount, "target_classes"); + hasExpectedSize(sourceLaneRanges, laneCount * 2, "source_lane_ranges"); + hasExpectedSize(targetLaneRanges, laneCount * 2, "target_lane_ranges"); + hasExpectedSize(sourceOperandForScheduledLane, laneCount, "source_operand_for_scheduled_lane"); + if (!multiSourcePayload || multiSourcePayload.getValue() != (transfer.getSources().size() > 1)) + diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError("phase-check scheduled batch deferred communication has inconsistent multi_source_payload metadata"); + }); + if (sourceOperandForScheduledLane) { + for (int64_t sourceOperandIndex : sourceOperandForScheduledLane.asArrayRef()) { + if (sourceOperandIndex < 0 || sourceOperandIndex >= static_cast(transfer.getSources().size())) { + diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError( + "phase-check scheduled batch deferred communication source_operand_for_scheduled_lane index is out of range"); + }); + break; + } + } + } + } + }); + + diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed"); + return success(!diagnostics.hasFailure()); +} + +LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled, + ArrayRef stepPlans) { + pim::CappedDiagnosticReporter diagnostics; + unsigned resultArgBase = getScheduledBatchResultArgBase(scheduled); + if (scheduled.getBody().getBlocks().size() != stepPlans.size()) { + diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "scheduled batch step routing expected " << stepPlans.size() + << " blocks but saw " << scheduled.getBody().getBlocks().size(); + }); + diagnostics.emitSuppressedSummary(scheduled.getOperation(), + "scheduled batch step routing verification failed"); + return failure(); + } + + SmallVector globalResultWrites(scheduled.getNumResults(), 0); + size_t stepIndex = 0; + for (Block &block : scheduled.getBody().getBlocks()) { + const ScheduledStepPlan &stepPlan = stepPlans[stepIndex++]; + SmallVector localWrites(stepPlan.resultCount, false); + auto inParallel = dyn_cast(block.getTerminator()); + if (!inParallel) { + diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex + << " is missing spat.in_parallel"; + }); + continue; + } + + for (Operation &op : inParallel.getRegion().front()) { + auto insert = dyn_cast(&op); + if (!insert) + continue; + auto dest = dyn_cast(insert.getDest()); + if (!dest || dest.getOwner() != &block) { + diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex + << " writes to a non-block result destination"; + }); + continue; + } + unsigned resultIndex = dest.getArgNumber() - resultArgBase; + if (dest.getArgNumber() < resultArgBase || resultIndex >= scheduled.getNumResults()) { + diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex + << " writes to invalid result block argument " << dest.getArgNumber(); + }); + continue; + } + if (resultIndex < stepPlan.resultOffset + || resultIndex >= stepPlan.resultOffset + stepPlan.resultCount) { + diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex + << " expected result range [" << stepPlan.resultOffset << ":" + << (stepPlan.resultOffset + stepPlan.resultCount) + << ") but wrote result " << resultIndex; + }); + continue; + } + localWrites[resultIndex - stepPlan.resultOffset] = true; + globalResultWrites[resultIndex]++; + } + + for (size_t index = 0; index < localWrites.size(); ++index) + if (!localWrites[index]) + diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex + << " did not write expected result " << (stepPlan.resultOffset + index); + }); + } + + for (size_t resultIndex = 0; resultIndex < globalResultWrites.size(); ++resultIndex) + if (globalResultWrites[resultIndex] != 1) + diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "scheduled batch result " << resultIndex << " expected one producing step but saw " + << globalResultWrites[resultIndex]; + }); + + diagnostics.emitSuppressedSummary(scheduled.getOperation(), "scheduled batch step routing verification failed"); + return success(!diagnostics.hasFailure()); +} + +LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch scheduled) { + pim::CappedDiagnosticReporter diagnostics; + unsigned resultArgBase = getScheduledBatchResultArgBase(scheduled); + for (auto enumeratedBlock : llvm::enumerate(scheduled.getBody().getBlocks())) { + size_t stepIndex = enumeratedBlock.index(); + Block &block = enumeratedBlock.value(); + Value scheduledLane = block.getArgument(0); + auto inParallel = dyn_cast(block.getTerminator()); + if (!inParallel) { + diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() << "phase-check scheduled batch step " << stepIndex + << " is missing spat.in_parallel"; + }); + continue; + } + auto isFinalScheduledOutputInsert = [&](Operation *op) { + auto insert = dyn_cast(op); + if (!insert || op->getParentOp() != inParallel.getOperation()) + return false; + auto dest = dyn_cast(insert.getDest()); + return dest && dest.getOwner() == &block && dest.getArgNumber() >= resultArgBase; + }; + + block.walk([&](Operation *op) { + if (op == block.getTerminator()) + return; + if (isFinalScheduledOutputInsert(op)) { + if (scheduled.getLaneCount() > 1) { + auto insert = cast(op); + bool dependsOnScheduledLane = false; + for (OpFoldResult offset : insert.getMixedOffsets()) { + if (auto value = dyn_cast(offset); value && valueTransitivelyDependsOn(value, scheduledLane)) { + dependsOnScheduledLane = true; + break; + } + } + if (!dependsOnScheduledLane) + diagnostics.report(insert.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError( + "phase-check scheduled batch final output insert must be indexed by scheduled lane"); + }); + } + return; + } + + auto insertSlice = dyn_cast(op); + if (!insertSlice) + return; + auto dest = dyn_cast(insertSlice.getDest()); + if (dest && dest.getOwner() == &block && dest.getArgNumber() >= resultArgBase) + return; + + auto destType = dyn_cast(insertSlice.getDestType()); + if (!destType || !destType.hasStaticShape() || destType.getRank() == 0) + return; + + for (OpFoldResult offset : insertSlice.getMixedOffsets()) { + auto value = dyn_cast(offset); + if (!value) + continue; + if (!valueTransitivelyDependsOn(value, scheduledLane)) + continue; + diagnostics.report(insertSlice.getOperation(), [&](Operation *illegalOp) { + illegalOp->emitOpError() + << "phase-check scheduled batch local fragment insert offset must use the source-instance inner lane, not the scheduled lane" + << " step " << stepIndex; + }); + break; + } + }); + } + + diagnostics.emitSuppressedSummary(scheduled.getOperation(), + "scheduled batch local fragment offset verification failed"); + return success(!diagnostics.hasFailure()); +} + + +LogicalResult verifyScheduledMaterializationRecords(ArrayRef materializedSchedules) { + for (const ScheduledMaterializationRecord &record : materializedSchedules) { + auto scheduled = dyn_cast(record.scheduledOp); + if (!scheduled) + continue; + if (failed(verifyMultiCpuStepResultRouting(scheduled, record.stepPlans))) + return failure(); + if (failed(verifyMultiCpuLocalFragmentOffsets(scheduled))) + return failure(); + } + return success(); +} + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp new file mode 100644 index 0000000..17f00b1 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include "ScheduledComputeMaterialization.hpp" + +namespace onnx_mlir { +namespace spatial { + +LogicalResult verifyMaterializedScheduleMapping( + func::FuncOp funcOp, + const MergeScheduleResult &schedule, + const llvm::MapVector &peftClassPlans, + const DenseMap &graphComputeToBlockMap, + ArrayRef materializedSchedules); + +LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp); +LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled, + ArrayRef stepPlans); +LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch scheduled); +LogicalResult verifyScheduledMaterializationRecords(ArrayRef materializedSchedules); + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp deleted file mode 100644 index 15bbf57..0000000 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp +++ /dev/null @@ -1,832 +0,0 @@ -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/AsmState.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/Pass/Pass.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/MapVector.h" -#include "llvm/ADT/SmallString.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/FormatVariadic.h" -#include "llvm/Support/raw_os_ostream.h" -#include "llvm/Support/raw_ostream.h" - -#include -#include -#include -#include - -#include "ComputeInstanceUtils.hpp" -#include "MergeSchedulingAnalysis.hpp" -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" -#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" - -using namespace mlir; - -namespace onnx_mlir { -namespace spatial { -namespace { - -struct ProducerValueKey { - ComputeInstance instance; - size_t resultIndex = 0; - - bool operator==(const ProducerValueKey& other) const { - return instance == other.instance && resultIndex == other.resultIndex; - } -}; - -struct GraphComputeBlockKey { - Operation* op = nullptr; - uint32_t laneStart = 0; - uint32_t laneCount = 1; - - bool operator==(const GraphComputeBlockKey& other) const { - return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount; - } -}; - -struct ScheduledMaterializationRecord { - Operation* scheduledOp = nullptr; - SmallVector canonicalCpuClasses; - SmallVector cpus; - SmallVector computeKeys; - SmallVector blocks; -}; - -} // namespace -} // namespace spatial -} // namespace onnx_mlir - -namespace llvm { -template <> -struct DenseMapInfo { - static onnx_mlir::spatial::ProducerValueKey getEmptyKey() { - return {DenseMapInfo::getEmptyKey(), std::numeric_limits::max()}; - } - static onnx_mlir::spatial::ProducerValueKey getTombstoneKey() { - return {DenseMapInfo::getTombstoneKey(), std::numeric_limits::max()}; - } - static unsigned getHashValue(const onnx_mlir::spatial::ProducerValueKey& key) { - return hash_combine(DenseMapInfo::getHashValue(key.instance), key.resultIndex); - } - static bool isEqual(const onnx_mlir::spatial::ProducerValueKey& lhs, - const onnx_mlir::spatial::ProducerValueKey& rhs) { - return lhs == rhs; - } -}; - -template <> -struct DenseMapInfo { - static onnx_mlir::spatial::GraphComputeBlockKey getEmptyKey() { - return {DenseMapInfo::getEmptyKey(), UINT32_MAX, UINT32_MAX}; - } - static onnx_mlir::spatial::GraphComputeBlockKey getTombstoneKey() { - return {DenseMapInfo::getTombstoneKey(), UINT32_MAX, UINT32_MAX - 1}; - } - static unsigned getHashValue(const onnx_mlir::spatial::GraphComputeBlockKey& key) { - return hash_combine(key.op, key.laneStart, key.laneCount); - } - static bool isEqual(const onnx_mlir::spatial::GraphComputeBlockKey& lhs, - const onnx_mlir::spatial::GraphComputeBlockKey& rhs) { - return lhs == rhs; - } -}; -} // namespace llvm - -namespace onnx_mlir { -namespace spatial { -namespace { - -struct ScalarClass { - size_t cpu = 0; - SmallVector cpus; - SmallVector instances; - SmallVector weights; - SmallVector inputs; - SmallVector resultTypes; -}; - -struct CompactCpuIndexMap { - DenseMap rawToCompact; - - size_t lookup(size_t rawCpu) const { - auto it = rawToCompact.find(rawCpu); - assert(it != rawToCompact.end() && "missing compact cpu index"); - return it->second; - } -}; - -static ComputeInstance getWholeBatchInstance(SpatComputeBatch batch) { - return {batch.getOperation(), 0, static_cast(batch.getLaneCount())}; -} - -static GraphComputeBlockKey getGraphComputeBlockKey(const ComputeInstance& instance) { - return {instance.op, instance.laneStart, instance.laneCount}; -} - -static CompactCpuIndexMap buildCompactCpuIndexMap(const MergeScheduleResult& schedule) { - SmallVector rawCpus; - auto append = [&](size_t rawCpu) { - if (!llvm::is_contained(rawCpus, rawCpu)) - rawCpus.push_back(rawCpu); - }; - for (const auto& entry : schedule.computeToCpuMap) - append(entry.second); - for (const auto& entry : schedule.equivalentClass) { - append(entry.first); - for (size_t rawCpu : entry.second) - append(rawCpu); - } - for (const auto& entry : schedule.cpuToLastComputeMap) - append(entry.first); - llvm::sort(rawCpus); - - CompactCpuIndexMap compactMap; - for (auto [index, rawCpu] : llvm::enumerate(rawCpus)) - compactMap.rawToCompact[rawCpu] = index; - return compactMap; -} - -static SmallVector getEquivalentCpus(size_t cpu, const MergeScheduleResult& schedule) { - llvm::SmallDenseSet seen; - SmallVector cpus; - auto append = [&](size_t value) { - if (seen.insert(value).second) - cpus.push_back(value); - }; - append(cpu); - if (auto it = schedule.equivalentClass.find(cpu); it != schedule.equivalentClass.end()) - for (size_t equivalentCpu : it->second) - append(equivalentCpu); - llvm::sort(cpus); - return cpus; -} - -static size_t getCanonicalCpuClass(size_t cpu, const MergeScheduleResult& schedule) { - SmallVector cpus = getEquivalentCpus(cpu, schedule); - return cpus.empty() ? cpu : cpus.front(); -} - -static Value getProducerOutput(const ProducerValueRef& ref) { - auto outputs = getComputeInstanceOutputValues(ref.instance); - assert(ref.resultIndex < outputs.size() && "producer result index out of range"); - return outputs[ref.resultIndex]; -} - -static SmallVector getCanonicalCpuClassesForBatch(SpatComputeBatch batch, const MergeScheduleResult& schedule) { - llvm::SmallDenseSet seen; - SmallVector classes; - size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount()); - for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) { - ComputeInstance chunk = getBatchChunkForIndex(batch, chunkIndex); - auto cpuIt = schedule.computeToCpuMap.find(chunk); - if (cpuIt == schedule.computeToCpuMap.end()) - continue; - size_t cpuClass = getCanonicalCpuClass(cpuIt->second, schedule); - if (seen.insert(cpuClass).second) - classes.push_back(cpuClass); - } - llvm::sort(classes); - return classes; -} - -static SmallVector remapCpuList(ArrayRef rawCpus, const CompactCpuIndexMap& compactCpuMap) { - SmallVector compactCpus; - compactCpus.reserve(rawCpus.size()); - for (size_t rawCpu : rawCpus) - compactCpus.push_back(compactCpuMap.lookup(rawCpu)); - llvm::sort(compactCpus); - return compactCpus; -} - -static std::string getInstanceName(const ComputeInstance& instance) { - return llvm::formatv("{0}[lanes={1}:{2}]", - instance.op->getName().getStringRef(), - instance.laneStart, - instance.laneStart + instance.laneCount) - .str(); -} - -static void appendUnique(SmallVectorImpl& values, Value value) { - if (!llvm::is_contained(values, value)) - values.push_back(value); -} - -static Value getBlockOperand(Block& block, ValueRange operands, Value value, unsigned firstArgument = 0) { - auto it = llvm::find(operands, value); - assert(it != operands.end() && "missing scheduled operand"); - return block.getArgument(firstArgument + std::distance(operands.begin(), it)); -} - -static void appendComputeBlockArguments(SmallVectorImpl& argTypes, - SmallVectorImpl& argLocs, - ValueRange weights, - ValueRange inputs, - ArrayRef carriedKeys, - Location loc) { - for (Value weight : weights) - argTypes.push_back(weight.getType()); - for (Value input : inputs) - argTypes.push_back(input.getType()); - for (ProducerValueKey key : carriedKeys) { - auto outputs = getComputeInstanceOutputValues(key.instance); - assert(key.resultIndex < outputs.size() && "missing carried result"); - argTypes.push_back(outputs[key.resultIndex].getType()); - } - argLocs.append(argTypes.size(), loc); -} - -static Block* createScheduledComputeBlock(PatternRewriter& rewriter, - SpatScheduledCompute scheduled, - ArrayRef carriedKeys, - Location loc) { - SmallVector argTypes; - SmallVector argLocs; - appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(), scheduled.getInputs(), carriedKeys, loc); - return rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), TypeRange(argTypes), argLocs); -} - -static void appendBlockYieldBaseAndCarriedOperands(Block& block, - unsigned baseArgCount, - size_t carriedCount, - SmallVectorImpl& operands) { - for (unsigned index = 0; index < baseArgCount; ++index) - operands.push_back(block.getArgument(index)); - for (size_t index = 0; index < carriedCount; ++index) - operands.push_back(block.getArgument(baseArgCount + index)); -} - -static void createBlockYield(PatternRewriter& rewriter, Location loc, ValueRange outputs, Block* next = nullptr) { - OperationState state(loc, SpatBlockYieldOp::getOperationName()); - state.addOperands(outputs); - if (next) - state.addSuccessors(next); - rewriter.create(state); -} - -static std::string formatValueLabel(Value value, AsmState& asmState) { - std::string storage; - llvm::raw_string_ostream os(storage); - value.printAsOperand(os, asmState); - return storage; -} - -static std::string formatBlockLabel(Block* block, AsmState& asmState) { - std::string storage; - llvm::raw_string_ostream os(storage); - block->printAsOperand(os, asmState); - return storage; -} - -static std::string formatOperationLabel(Operation* op, AsmState& asmState) { - if (op->getNumResults() == 0) - return op->getName().getStringRef().str(); - std::string storage; - llvm::raw_string_ostream os(storage); - llvm::interleaveComma(op->getResults(), os, [&](Value result) { os << formatValueLabel(result, asmState); }); - return os.str(); -} - -static std::string formatGraphComputeBlockKey(const GraphComputeBlockKey& key, AsmState& asmState) { - return llvm::formatv("{0} {1} lanes [{2}:{3}]", - formatOperationLabel(key.op, asmState), - key.op->getName().getStringRef(), - key.laneStart, - key.laneStart + key.laneCount) - .str(); -} - -static size_t countPeftEquivalenceClasses(const MergeScheduleResult& schedule) { - llvm::SmallDenseSet classes; - for (const ComputeInstance& instance : schedule.dominanceOrderCompute) { - auto cpuIt = schedule.computeToCpuMap.find(instance); - if (cpuIt == schedule.computeToCpuMap.end()) - continue; - classes.insert(getCanonicalCpuClass(cpuIt->second, schedule)); - } - return classes.size(); -} - -static SmallVector collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) { - SmallVector keys; - for (Operation& op : funcOp.getOps()) { - if (auto compute = dyn_cast(&op)) - keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1})); - else if (auto batch = dyn_cast(&op)) - keys.push_back(getGraphComputeBlockKey(getWholeBatchInstance(batch))); - } - return keys; -} - -static LogicalResult verifyMaterializedScheduleMapping( - func::FuncOp funcOp, - const MergeScheduleResult& schedule, - const DenseMap& graphComputeToBlockMap, - ArrayRef materializedSchedules) { - pim::CappedDiagnosticReporter diagnostics; - size_t expectedClassCount = countPeftEquivalenceClasses(schedule); - if (expectedClassCount != materializedSchedules.size()) { - diagnostics.report(funcOp.getOperation(), [&](Operation* illegalOp) { - illegalOp->emitOpError() << "phase-check expected " << expectedClassCount - << " PEFT equivalence classes but materialized " << materializedSchedules.size() - << " scheduled computes"; - }); - } - - for (const ScheduledMaterializationRecord& record : materializedSchedules) { - if (record.canonicalCpuClasses.empty()) { - diagnostics.report(record.scheduledOp, [&](Operation* illegalOp) { - illegalOp->emitOpError("phase-check scheduled compute is missing an owning PEFT equivalence class"); - }); - continue; - } - if (record.canonicalCpuClasses.size() > 1) { - diagnostics.report(record.scheduledOp, [&](Operation* illegalOp) { - std::string classes; - llvm::raw_string_ostream os(classes); - llvm::interleaveComma(record.canonicalCpuClasses, os, [&](size_t cpuClass) { os << cpuClass; }); - illegalOp->emitOpError("phase-check scheduled compute spans multiple PEFT equivalence classes") - << " (classes " << os.str() << ")"; - }); - } - } - - for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) { - if (graphComputeToBlockMap.count(key)) - continue; - diagnostics.report(key.op, [&](Operation* illegalOp) { - illegalOp->emitOpError() << "phase-check graph compute is missing a scheduled MLIR block mapping for lanes [" - << key.laneStart << ":" << (key.laneStart + key.laneCount) << "]"; - }); - } - - for (const auto& entry : graphComputeToBlockMap) { - Block* block = entry.second; - if (!block || !isa(block->getParentOp())) { - diagnostics.report(entry.first.op, [&](Operation* illegalOp) { - illegalOp->emitOpError("phase-check graph compute block mapping does not target a scheduled compute block"); - }); - } - } - - diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial materialization verification failed"); - return success(!diagnostics.hasFailure()); -} - -static void dumpPeftMaterializationReport(func::FuncOp funcOp, - const MergeScheduleResult& schedule, - const CompactCpuIndexMap& compactCpuMap, - ArrayRef materializedSchedules) { - std::fstream file = openDialectDumpFileWithExtension("spatial1_schedule_materialization", "/reports", "txt"); - if (!file.is_open()) - return; - - llvm::raw_os_ostream os(file); - AsmState asmState(funcOp); - - llvm::MapVector> classToCpus; - DenseMap> cpuToInstances; - for (const ComputeInstance& instance : schedule.dominanceOrderCompute) { - auto cpuIt = schedule.computeToCpuMap.find(instance); - if (cpuIt == schedule.computeToCpuMap.end()) - continue; - size_t cpu = compactCpuMap.lookup(cpuIt->second); - size_t cpuClass = compactCpuMap.lookup(getCanonicalCpuClass(cpuIt->second, schedule)); - if (!llvm::is_contained(classToCpus[cpuClass], cpu)) - classToCpus[cpuClass].push_back(cpu); - cpuToInstances[cpu].push_back(instance); - } - - os << "PEFT Schedule\n"; - os << "=============\n"; - os << "equivalence classes: " << classToCpus.size() << "\n\n"; - for (auto& entry : classToCpus) { - size_t cpuClass = entry.first; - llvm::sort(entry.second); - os << "class " << cpuClass << " cpus ["; - llvm::interleaveComma(entry.second, os, [&](size_t cpu) { os << cpu; }); - os << "]\n"; - for (size_t cpu : entry.second) { - os << " cpu " << cpu << ":\n"; - auto& instances = cpuToInstances[cpu]; - llvm::sort(instances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) { - return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs); - }); - for (const ComputeInstance& instance : instances) { - os << " slot " << schedule.computeToCpuSlotMap.lookup(instance) << ": " - << formatGraphComputeBlockKey(getGraphComputeBlockKey(instance), asmState) << "\n"; - } - } - os << "\n"; - } - - os << "Materialized Scheduled Computes\n"; - os << "===============================\n"; - os << "scheduled ops: " << materializedSchedules.size() << "\n\n"; - for (const ScheduledMaterializationRecord& record : materializedSchedules) { - os << formatOperationLabel(record.scheduledOp, asmState) << " " << record.scheduledOp->getName().getStringRef(); - os << " classes ["; - llvm::interleaveComma(record.canonicalCpuClasses, os, [&](size_t cpuClass) { os << cpuClass; }); - os << "] cpus ["; - llvm::interleaveComma(record.cpus, os, [&](size_t cpu) { os << cpu; }); - os << "]\n"; - for (auto [key, block] : llvm::zip(record.computeKeys, record.blocks)) - os << " " << formatGraphComputeBlockKey(key, asmState) << " -> " << formatBlockLabel(block, asmState) << "\n"; - os << "\n"; - } -} - -static LogicalResult createDeferredCommunication(PatternRewriter& rewriter, - Location loc, - Value consumerInput, - Value source, - const ProducerValueRef& producer, - const ComputeInstance& consumer, - const MergeScheduleResult& schedule, - const CompactCpuIndexMap& compactCpuMap, - uint64_t exchangeId, - Value& result) { - auto sourceCpu = compactCpuMap.lookup(schedule.computeToCpuMap.lookup(producer.instance)); - auto targetCpu = compactCpuMap.lookup(schedule.computeToCpuMap.lookup(consumer)); - auto transfer = SpatDeferredCommunicationOp::create( - rewriter, - loc, - consumerInput.getType(), - ValueRange {source}, - rewriter.getStringAttr(("x" + std::to_string(exchangeId)).c_str()), - rewriter.getStringAttr(getInstanceName(producer.instance)), - rewriter.getStringAttr(getInstanceName(consumer)), - rewriter.getI64IntegerAttr(static_cast(sourceCpu)), - rewriter.getI64IntegerAttr(static_cast(targetCpu)), - rewriter.getI64IntegerAttr(static_cast(sourceCpu)), - rewriter.getI64IntegerAttr(static_cast(targetCpu)), - rewriter.getI64IntegerAttr(producer.instance.laneStart), - rewriter.getI64IntegerAttr(consumer.laneStart), - rewriter.getStringAttr(consumerInput.getDefiningOp() ? "projected" : "ordinary"), - rewriter.getI64IntegerAttr(static_cast(producer.resultIndex)), - rewriter.getDenseI64ArrayAttr({}), - rewriter.getI64IntegerAttr(-1)); - - Block* block = rewriter.createBlock(&transfer.getBody(), transfer.getBody().end(), source.getType(), loc); - rewriter.setInsertionPointToStart(block); - Value yielded = block->getArgument(0); - if (yielded.getType() != consumerInput.getType()) { - auto slice = consumerInput.getDefiningOp(); - if (!slice || slice.getSource().getType() != source.getType()) - return transfer.emitOpError("cannot derive deferred communication payload type from recorded source"); - IRMapping mapper; - mapper.map(slice.getSource(), yielded); - auto cloned = cast(rewriter.clone(*slice, mapper)); - yielded = cloned.getResult(); - } - SpatYieldOp::create(rewriter, loc, yielded); - result = transfer.getOutput(); - rewriter.setInsertionPointAfter(transfer); - return success(); -} - -static LogicalResult mapComputeInputs(PatternRewriter& rewriter, - Block& block, - SpatCompute compute, - const ComputeInstance& instance, - const ScalarClass& klass, - const MergeScheduleResult& schedule, - const CompactCpuIndexMap& compactCpuMap, - ValueRange scheduledWeights, - ValueRange scheduledInputs, - IRMapping& mapper, - uint64_t& exchangeId) { - for (auto [index, weight] : llvm::enumerate(compute.getWeights())) - mapper.map(*compute.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight)); - - for (auto [index, input] : llvm::enumerate(compute.getInputs())) { - Value mapped; - std::optional producer = getProducerValueRef(input, &instance); - if (producer) { - Value source = getProducerOutput(*producer); - if (failed(createDeferredCommunication( - rewriter, input.getLoc(), input, source, *producer, instance, schedule, compactCpuMap, exchangeId++, mapped))) - return failure(); - } - else { - mapped = getBlockOperand(block, scheduledInputs, input, scheduledWeights.size()); - } - mapper.map(*compute.getInputArgument(index), mapped); - } - return success(); -} - -static LogicalResult mapBatchInputs(PatternRewriter& rewriter, - Block& block, - SpatComputeBatch batch, - const MergeScheduleResult& schedule, - const CompactCpuIndexMap& compactCpuMap, - ValueRange scheduledWeights, - ValueRange scheduledInputs, - IRMapping& mapper, - uint64_t& exchangeId) { - mapper.map(*batch.getLaneArgument(), block.getArgument(0)); - unsigned firstWeightArg = 1; - unsigned firstInputArg = firstWeightArg + scheduledWeights.size(); - for (auto [index, weight] : llvm::enumerate(batch.getWeights())) - mapper.map(*batch.getWeightArgument(index), block.getArgument(firstWeightArg + index)); - - ComputeInstance consumer {batch.getOperation(), 0, static_cast(batch.getLaneCount())}; - for (auto [index, input] : llvm::enumerate(batch.getInputs())) { - Value mapped; - std::optional producer = getProducerValueRef(input, &consumer); - if (producer) { - Value source = getProducerOutput(*producer); - if (failed(createDeferredCommunication( - rewriter, input.getLoc(), input, source, *producer, consumer, schedule, compactCpuMap, exchangeId++, mapped))) - return failure(); - } - else { - auto it = llvm::find(scheduledInputs, input); - if (it == scheduledInputs.end()) - return batch.emitOpError("host batch input is missing from scheduled batch operands"); - mapped = block.getArgument(firstInputArg + std::distance(scheduledInputs.begin(), it)); - } - mapper.map(*batch.getInputArgument(index), mapped); - } - - unsigned firstOutputArg = firstInputArg + scheduledInputs.size(); - for (unsigned index = 0; index < batch.getNumResults(); ++index) - mapper.map(*batch.getOutputArgument(index), block.getArgument(firstOutputArg + index)); - return success(); -} - -static void cloneComputeBody(PatternRewriter& rewriter, - Block& source, - IRMapping& mapper, - SmallVectorImpl& yieldedValues) { - for (Operation& op : source.without_terminator()) - rewriter.clone(op, mapper); - auto yield = cast(source.getTerminator()); - for (Value output : yield.getOutputs()) - yieldedValues.push_back(mapper.lookup(output)); -} - -struct MergeComputeNodesPass final : PassWrapper> { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass) - - StringRef getArgument() const override { return "pim-merge-compute-nodes"; } - StringRef getDescription() const override { - return "Materialize scheduled Spatial compute with deferred communication placeholders."; - } - - void runOnOperation() override { - ModuleOp moduleOp = getOperation(); - auto entryFunc = getPimEntryFunc(moduleOp); - if (failed(entryFunc)) { - moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes"); - signalPassFailure(); - return; - } - - func::FuncOp funcOp = *entryFunc; - MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult(); - CompactCpuIndexMap compactCpuMap = buildCompactCpuIndexMap(schedule); - llvm::MapVector scalarClasses; - SmallVector batchOps; - for (const ComputeInstance& instance : schedule.dominanceOrderCompute) { - if (isa(instance.op)) { - size_t cpu = schedule.computeToCpuMap.lookup(instance); - size_t cpuClass = getCanonicalCpuClass(cpu, schedule); - auto& klass = scalarClasses[cpuClass]; - klass.cpu = cpuClass; - if (klass.cpus.empty()) - klass.cpus = getEquivalentCpus(cpu, schedule); - klass.instances.push_back(instance); - } - else { - auto batch = cast(instance.op); - if (!llvm::is_contained(batchOps, batch)) - batchOps.push_back(batch); - } - } - for (auto& entry : scalarClasses) { - llvm::sort(entry.second.instances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) { - return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs); - }); - } - - PatternRewriter rewriter(moduleOp.getContext()); - Operation* insertionPoint = funcOp.getBody().front().getTerminator(); - DenseMap scalarOps; - DenseMap batchReplacements; - DenseMap scheduledOpToRecordIndex; - DenseMap graphComputeToBlockMap; - std::vector materializedSchedules; - - for (auto& entry : scalarClasses) { - ScalarClass& klass = entry.second; - for (const ComputeInstance& instance : klass.instances) { - auto compute = cast(instance.op); - llvm::append_range(klass.resultTypes, compute.getResultTypes()); - for (Value weight : compute.getWeights()) - appendUnique(klass.weights, weight); - for (Value input : compute.getInputs()) { - std::optional producer = getProducerValueRef(input, &instance); - if (!producer) - appendUnique(klass.inputs, input); - } - } - rewriter.setInsertionPoint(insertionPoint); - auto scheduled = SpatScheduledCompute::create( - rewriter, klass.instances.front().op->getLoc(), TypeRange(klass.resultTypes), klass.weights, klass.inputs); - scheduled->setAttr( - kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast(compactCpuMap.lookup(klass.cpu)))); - SmallVector equivalentCpus; - for (size_t cpu : klass.cpus) - equivalentCpus.push_back(static_cast(compactCpuMap.lookup(cpu))); - scheduled->setAttr("scheduled.equivalent_cpus", rewriter.getDenseI64ArrayAttr(equivalentCpus)); - SmallVector sources; - SmallVector lanes; - for (const ComputeInstance& instance : klass.instances) { - sources.push_back(rewriter.getStringAttr(getInstanceName(instance))); - lanes.push_back(instance.laneStart); - lanes.push_back(instance.laneCount); - } - scheduled->setAttr("scheduled.sources", rewriter.getArrayAttr(sources)); - scheduled->setAttr("scheduled.lane_ranges", rewriter.getDenseI64ArrayAttr(lanes)); - scalarOps[&klass] = scheduled; - ScheduledMaterializationRecord record; - record.scheduledOp = scheduled.getOperation(); - record.canonicalCpuClasses.push_back(compactCpuMap.lookup(klass.cpu)); - record.cpus = remapCpuList(klass.cpus, compactCpuMap); - scheduledOpToRecordIndex[scheduled.getOperation()] = materializedSchedules.size(); - materializedSchedules.push_back(std::move(record)); - } - - for (SpatComputeBatch batch : batchOps) { - SmallVector inputs; - ComputeInstance wholeBatch = getWholeBatchInstance(batch); - for (Value input : batch.getInputs()) - if (!getProducerValueRef(input, &wholeBatch)) - appendUnique(inputs, input); - - rewriter.setInsertionPoint(insertionPoint); - auto scheduled = SpatScheduledComputeBatch::create(rewriter, - batch.getLoc(), - batch.getResultTypes(), - rewriter.getI32IntegerAttr(batch.getLaneCount()), - batch.getWeights(), - inputs); - SmallVector coreIds; - uint32_t laneCount = batch.getLaneCount(); - coreIds.reserve(laneCount); - for (uint32_t lane = 0; lane < laneCount; ++lane) { - ComputeInstance chunk = getBatchChunkForLane(batch, lane); - coreIds.push_back(static_cast(compactCpuMap.lookup(schedule.computeToCpuMap.lookup(chunk)))); - } - scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(coreIds)); - scheduled->setAttr("scheduled.sources", rewriter.getArrayAttr({rewriter.getStringAttr(getInstanceName(wholeBatch))})); - batchReplacements[batch.getOperation()] = scheduled; - ScheduledMaterializationRecord record; - record.scheduledOp = scheduled.getOperation(); - record.canonicalCpuClasses = remapCpuList(getCanonicalCpuClassesForBatch(batch, schedule), compactCpuMap); - for (int32_t coreId : coreIds) - if (!llvm::is_contained(record.cpus, static_cast(coreId))) - record.cpus.push_back(static_cast(coreId)); - scheduledOpToRecordIndex[scheduled.getOperation()] = materializedSchedules.size(); - materializedSchedules.push_back(std::move(record)); - } - - uint64_t exchangeId = 0; - for (auto& entry : scalarClasses) { - ScalarClass& klass = entry.second; - SpatScheduledCompute scheduled = scalarOps.lookup(&klass); - SmallVector carriedKeys; - Block* block = nullptr; - for (auto [ordinal, instance] : llvm::enumerate(klass.instances)) { - auto compute = cast(instance.op); - if (!block) - block = createScheduledComputeBlock(rewriter, scheduled, carriedKeys, compute.getLoc()); - GraphComputeBlockKey key = getGraphComputeBlockKey(instance); - graphComputeToBlockMap[key] = block; - auto& record = materializedSchedules[scheduledOpToRecordIndex.lookup(scheduled.getOperation())]; - record.computeKeys.push_back(key); - record.blocks.push_back(block); - rewriter.setInsertionPointToStart(block); - IRMapping mapper; - if (failed(mapComputeInputs(rewriter, - *block, - compute, - instance, - klass, - schedule, - compactCpuMap, - scheduled.getWeights(), - scheduled.getInputs(), - mapper, - exchangeId))) { - signalPassFailure(); - return; - } - SmallVector yieldedValues; - cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues); - SmallVector currentKeys; - for (auto [index, value] : llvm::enumerate(yieldedValues)) { - ProducerValueKey key {instance, index}; - currentKeys.push_back(key); - } - - unsigned baseArgCount = scheduled.getWeights().size() + scheduled.getInputs().size(); - SmallVector blockYieldOperands; - bool hasNextBlock = ordinal + 1 < klass.instances.size(); - if (hasNextBlock) { - SmallVector nextCarriedKeys(carriedKeys); - llvm::append_range(nextCarriedKeys, currentKeys); - Block* nextBlock = createScheduledComputeBlock(rewriter, scheduled, nextCarriedKeys, compute.getLoc()); - appendBlockYieldBaseAndCarriedOperands(*block, baseArgCount, carriedKeys.size(), blockYieldOperands); - llvm::append_range(blockYieldOperands, yieldedValues); - rewriter.setInsertionPointToEnd(block); - createBlockYield(rewriter, compute.getLoc(), blockYieldOperands, nextBlock); - carriedKeys = std::move(nextCarriedKeys); - block = nextBlock; - } - else { - for (size_t index = 0; index < carriedKeys.size(); ++index) - blockYieldOperands.push_back(block->getArgument(baseArgCount + index)); - llvm::append_range(blockYieldOperands, yieldedValues); - rewriter.setInsertionPointToEnd(block); - createBlockYield(rewriter, compute.getLoc(), blockYieldOperands); - } - } - } - - for (SpatComputeBatch batch : batchOps) { - SpatScheduledComputeBatch scheduled = batchReplacements.lookup(batch.getOperation()); - SmallVector blockArgTypes {rewriter.getIndexType()}; - SmallVector blockArgLocs {batch.getLoc()}; - for (Value weight : scheduled.getWeights()) { - blockArgTypes.push_back(weight.getType()); - blockArgLocs.push_back(weight.getLoc()); - } - for (Value input : scheduled.getInputs()) { - blockArgTypes.push_back(input.getType()); - blockArgLocs.push_back(input.getLoc()); - } - for (Type resultType : scheduled.getResultTypes()) { - blockArgTypes.push_back(resultType); - blockArgLocs.push_back(batch.getLoc()); - } - Block* block = rewriter.createBlock( - &scheduled.getBody(), scheduled.getBody().end(), TypeRange(blockArgTypes), blockArgLocs); - GraphComputeBlockKey key = getGraphComputeBlockKey(getWholeBatchInstance(batch)); - graphComputeToBlockMap[key] = block; - auto& record = materializedSchedules[scheduledOpToRecordIndex.lookup(scheduled.getOperation())]; - record.computeKeys.push_back(key); - record.blocks.push_back(block); - rewriter.setInsertionPointToStart(block); - IRMapping mapper; - if (failed(mapBatchInputs(rewriter, - *block, - batch, - schedule, - compactCpuMap, - scheduled.getWeights(), - scheduled.getInputs(), - mapper, - exchangeId))) { - signalPassFailure(); - return; - } - for (Operation& op : batch.getBody().front()) - rewriter.clone(op, mapper); - } - - dumpPeftMaterializationReport(funcOp, schedule, compactCpuMap, materializedSchedules); - if (failed(verifyMaterializedScheduleMapping(funcOp, schedule, graphComputeToBlockMap, materializedSchedules))) { - moduleOp.emitError("scheduled Spatial materialization mapping verification failed"); - signalPassFailure(); - return; - } - if (failed(verifyScheduledSpatialInvariants(funcOp))) { - moduleOp.emitError("scheduled Spatial phase 1 verification failed"); - signalPassFailure(); - return; - } - dumpModule(moduleOp, "spatial1_scheduled_no_comm", /*assumeVerified=*/true); - moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; " - "Phase 2 communication realization is not implemented"); - signalPassFailure(); - } -}; - -} // namespace - -} // namespace spatial - -std::unique_ptr createMergeComputeNodesPass() { return std::make_unique(); } - -} // namespace onnx_mlir