From 83a54e28e4a3955b9f3f0a89802f11e08fa25375 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Mon, 6 Jul 2026 10:12:20 +0200 Subject: [PATCH] meno diamantini --- src/PIM/CMakeLists.txt | 1 - src/PIM/Common/IR/CoreBlockUtils.cpp | 31 + src/PIM/Compiler/PimCodeGen.cpp | 94 +- src/PIM/Conversion/CMakeLists.txt | 3 +- .../Conversion/ONNXToSpatial/CMakeLists.txt | 2 + .../ONNXToSpatial/Common/BiasAddUtils.cpp | 112 + .../ONNXToSpatial/Common/BiasAddUtils.hpp | 30 + .../Common/RowStripLayoutUtils.cpp | 239 ++ .../Common/RowStripLayoutUtils.hpp | 69 + .../ONNXToSpatial/LowerSpatialPlansPass.cpp | 175 +- .../ONNXToSpatial/ONNXToSpatialPass.cpp | 5 +- .../ONNXToSpatial/ONNXToSpatialVerifier.cpp | 1 + .../ONNXToSpatial/Patterns/Math/Conv.cpp | 642 ++-- .../Patterns/Math/Elementwise.cpp | 59 +- .../SpatialLayoutPlanningPass.cpp | 64 +- .../SpatialToGraphviz/CMakeLists.txt | 17 - .../SpatialToGraphviz/SpatialToGraphviz.cpp | 259 -- .../Bufferization/PimBufferizationPass.cpp | 149 +- src/PIM/Dialect/Spatial/Spatial.td | 16 + src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp | 51 +- .../MaterializeMergeSchedule.cpp | 2940 ++++++++++++++--- .../MaterializedClassState.hpp | 54 + .../MergeComputeNodes/ProjectedFragments.hpp | 18 + src/PIM/Pass/PIMPasses.h | 2 - src/PIM/PimAccelerator.cpp | 1 - validation/validate.py | 0 26 files changed, 3756 insertions(+), 1278 deletions(-) create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp delete mode 100644 src/PIM/Conversion/SpatialToGraphviz/CMakeLists.txt delete mode 100644 src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp mode change 100644 => 100755 validation/validate.py diff --git a/src/PIM/CMakeLists.txt b/src/PIM/CMakeLists.txt index 31588b2..86c13c3 100644 --- a/src/PIM/CMakeLists.txt +++ b/src/PIM/CMakeLists.txt @@ -117,7 +117,6 @@ add_pim_library(OMPIMAccel SpatialOps PimOps OMONNXToSpatial - OMSpatialToGraphviz OMSpatialToPim OMPimCommon OMPimBufferization diff --git a/src/PIM/Common/IR/CoreBlockUtils.cpp b/src/PIM/Common/IR/CoreBlockUtils.cpp index 03bdc8f..42a6699 100644 --- a/src/PIM/Common/IR/CoreBlockUtils.cpp +++ b/src/PIM/Common/IR/CoreBlockUtils.cpp @@ -74,6 +74,21 @@ walkPimCoreBlock(mlir::Block& block, continue; } + if (auto ifOp = mlir::dyn_cast(op)) { + auto condition = resolveIndexValue(ifOp.getCondition(), knowledge); + if (failed(condition)) { + ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen"); + hasFailure = true; + continue; + } + + mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion(); + if (!selectedRegion.empty()) + if (failed(walkPimCoreBlock(selectedRegion.front(), knowledge, callback))) + hasFailure = true; + continue; + } + if (failed(callback(op, knowledge))) hasFailure = true; } @@ -128,6 +143,22 @@ mlir::LogicalResult walkPimCoreBlockStructurally( continue; } + if (auto ifOp = mlir::dyn_cast(op)) { + if (failed(resolveIndexValue(ifOp.getCondition(), knowledge))) { + ifOp.emitOpError("requires statically evaluable scf.if condition for PIM verification"); + hasFailure = true; + continue; + } + + if (!ifOp.getThenRegion().empty()) + if (failed(walkPimCoreBlockStructurally(ifOp.getThenRegion().front(), knowledge, callback))) + hasFailure = true; + if (!ifOp.getElseRegion().empty()) + if (failed(walkPimCoreBlockStructurally(ifOp.getElseRegion().front(), knowledge, callback))) + hasFailure = true; + continue; + } + if (failed(callback(op, knowledge))) hasFailure = true; } diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index 30b0ac4..b949eb7 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -414,31 +414,35 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value, const StaticValueKnowledge& knowledge, std::optional lane) const { value = resolveCachedAlias(value, knowledge); - auto compiledIt = compiledAddressExprs.find(value); - if (compiledIt == compiledAddressExprs.end()) { - auto compiledExpr = compileContiguousAddressExpr(value); - if (failed(compiledExpr)) { - errs() << "Failed to compile contiguous address for value: "; + + FailureOr resolvedAddress = resolveContiguousAddress(value, knowledge); + if (failed(resolvedAddress)) { + auto compiledIt = compiledAddressExprs.find(value); + if (compiledIt == compiledAddressExprs.end()) { + auto compiledExpr = compileContiguousAddressExpr(value); + if (failed(compiledExpr)) { + errs() << "Failed to compile contiguous address for value: "; + value.print(errs()); + errs() << " : " << value.getType(); + errs() << "\n"; + llvm_unreachable("Failed to compile contiguous address"); + } + compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first; + } + + resolvedAddress = compiledIt->second.evaluate(knowledge, lane); + if (failed(resolvedAddress)) { + errs() << "Failed to evaluate contiguous address for value: "; value.print(errs()); errs() << " : " << value.getType(); errs() << "\n"; - llvm_unreachable("Failed to compile contiguous address"); + if (auto* definingOp = value.getDefiningOp()) { + errs() << "Defining op:\n"; + definingOp->print(errs()); + errs() << "\n"; + } + llvm_unreachable("Failed to resolve contiguous address"); } - compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first; - } - - auto resolvedAddress = compiledIt->second.evaluate(knowledge, lane); - if (failed(resolvedAddress)) { - errs() << "Failed to evaluate contiguous address for value: "; - value.print(errs()); - errs() << " : " << value.getType(); - errs() << "\n"; - if (auto* definingOp = value.getDefiningOp()) { - errs() << "Defining op:\n"; - definingOp->print(errs()); - errs() << "\n"; - } - llvm_unreachable("Failed to resolve contiguous address"); } MemoryValueKey key = getMemoryValueKey(resolvedAddress->base, lane); @@ -1114,7 +1118,8 @@ enum class CompiledCoreOpKind : uint8_t { struct CompiledCoreNode { enum class Kind : uint8_t { Op, - Loop + Loop, + If }; Kind kind = Kind::Op; @@ -1123,7 +1128,10 @@ struct CompiledCoreNode { CompiledIndexExpr lowerBound; CompiledIndexExpr upperBound; CompiledIndexExpr step; + CompiledIndexExpr condition; std::unique_ptr> loopBody; + std::unique_ptr> thenBody; + std::unique_ptr> elseBody; }; static FailureOr classifyCompiledCoreOpKind(Operation& op) { @@ -1201,6 +1209,28 @@ compileCoreEmissionPlan(Block& block, Operation* weightOwner, llvm::SmallVectorI continue; } + if (auto ifOp = dyn_cast(op)) { + auto condition = compileIndexExpr(ifOp.getCondition()); + if (failed(condition)) { + ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen"); + return failure(); + } + + CompiledCoreNode ifNode; + ifNode.kind = CompiledCoreNode::Kind::If; + ifNode.op = ifOp.getOperation(); + ifNode.condition = *condition; + ifNode.thenBody = std::make_unique>(); + if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), weightOwner, *ifNode.thenBody))) + return failure(); + ifNode.elseBody = std::make_unique>(); + if (!ifOp.getElseRegion().empty()) + if (failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), weightOwner, *ifNode.elseBody))) + return failure(); + plan.push_back(std::move(ifNode)); + continue; + } + auto opKind = classifyCompiledCoreOpKind(op); if (failed(opKind)) { InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'"; @@ -1263,6 +1293,26 @@ static LogicalResult executeCompiledCorePlan( continue; } + if (node.kind == CompiledCoreNode::Kind::If) { + auto condition = node.condition.evaluate(knowledge); + auto ifOp = cast(node.op); + if (failed(condition)) { + ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen"); + return failure(); + } + + const auto& selectedBody = *condition != 0 ? node.thenBody : node.elseBody; + if (selectedBody && failed(executeCompiledCorePlan(*selectedBody, + coreCodeGen, + knowledge, + resolveWeightSlot, + processedOperations, + batchLane, + batchLaneCount))) + return failure(); + continue; + } + switch (node.opKind) { case CompiledCoreOpKind::Load: coreCodeGen.codeGenLoadOp(cast(node.op), knowledge); diff --git a/src/PIM/Conversion/CMakeLists.txt b/src/PIM/Conversion/CMakeLists.txt index dcbc3d9..0e91043 100644 --- a/src/PIM/Conversion/CMakeLists.txt +++ b/src/PIM/Conversion/CMakeLists.txt @@ -1,3 +1,2 @@ add_subdirectory(ONNXToSpatial) -add_subdirectory(SpatialToGraphviz) -add_subdirectory(SpatialToPim) \ No newline at end of file +add_subdirectory(SpatialToPim) diff --git a/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt b/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt index ae4bc06..60e9e44 100644 --- a/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt +++ b/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt @@ -31,8 +31,10 @@ add_pim_library(OMONNXToSpatial SpatialLayoutPlanningPass.cpp LowerSpatialPlansPass.cpp Common/AttributeUtils.cpp + Common/BiasAddUtils.cpp Common/ComputeRegionBuilder.cpp Common/MatrixProductLowering.cpp + Common/RowStripLayoutUtils.cpp Common/ShapeTilingUtils.cpp Common/WeightMaterialization.cpp diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.cpp new file mode 100644 index 0000000..868dc2f --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.cpp @@ -0,0 +1,112 @@ +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" + +#include "llvm/ADT/SmallVector.h" + +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" + +using namespace mlir; + +namespace onnx_mlir { + +LogicalResult isSupportedBiasAddShape(RankedTensorType biasType, RankedTensorType resultType) { + if (!biasType || !resultType || !biasType.hasStaticShape() || !resultType.hasStaticShape()) + return failure(); + if (resultType.getRank() != 4) + return failure(); + if (biasType.getElementType() != resultType.getElementType()) + return failure(); + + const int64_t channels = resultType.getDimSize(1); + ArrayRef shape = biasType.getShape(); + if (shape.empty()) + return success(); + if (shape.size() == 1) + return success(shape[0] == channels); + if (shape.size() == 2) + return success(shape[0] == 1 && shape[1] == channels); + if (shape.size() == 4) + return success(shape[0] == 1 && shape[1] == channels && shape[2] == 1 && shape[3] == 1); + return failure(); +} + +FailureOr> getBiasChannelValues(DenseElementsAttr denseAttr, RankedTensorType resultType) { + auto biasType = dyn_cast(denseAttr.getType()); + if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType))) + return failure(); + + const int64_t channels = resultType.getDimSize(1); + if (denseAttr.isSplat()) { + return SmallVector(channels, denseAttr.getSplatValue()); + } + + SmallVector flattened(denseAttr.getValues()); + if (biasType.getRank() == 1) + return flattened; + if (biasType.getRank() == 2) + return flattened; + + SmallVector channelValues; + channelValues.reserve(channels); + const int64_t channelStride = biasType.getDimSize(2) * biasType.getDimSize(3); + for (int64_t channel = 0; channel < channels; ++channel) + channelValues.push_back(flattened[channel * channelStride]); + return channelValues; +} + +bool isSupportedBiasAddValue(Value bias, RankedTensorType resultType, DenseElementsAttr* denseAttr) { + auto attr = getHostConstDenseElementsAttr(bias); + if (!attr) + return false; + auto biasType = dyn_cast(attr.getType()); + if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType))) + return false; + if (failed(getBiasChannelValues(attr, resultType))) + return false; + if (denseAttr) + *denseAttr = attr; + return true; +} + +FailureOr classifyBiasAddPlanCandidate(Value lhs, Value rhs, RankedTensorType resultType) { + auto lhsType = dyn_cast(lhs.getType()); + auto rhsType = dyn_cast(rhs.getType()); + if (!lhsType || !rhsType) + return failure(); + if (lhsType == resultType && isSupportedBiasAddValue(rhs, resultType)) + return BiasAddPlanCandidate {lhs, rhs}; + if (rhsType == resultType && isSupportedBiasAddValue(lhs, resultType)) + return BiasAddPlanCandidate {rhs, lhs}; + return failure(); +} + +FailureOr +materializeDenseBiasAddTensor(Value bias, RankedTensorType resultType, RewriterBase& rewriter, Location loc) { + DenseElementsAttr denseAttr; + if (!isSupportedBiasAddValue(bias, resultType, &denseAttr)) + return failure(); + + FailureOr> channelValues = getBiasChannelValues(denseAttr, resultType); + if (failed(channelValues)) + return failure(); + + SmallVector resultValues; + resultValues.reserve(resultType.getNumElements()); + const int64_t batches = resultType.getDimSize(0); + const int64_t channels = resultType.getDimSize(1); + const int64_t height = resultType.getDimSize(2); + const int64_t width = resultType.getDimSize(3); + for (int64_t n = 0; n < batches; ++n) + for (int64_t c = 0; c < channels; ++c) + for (int64_t h = 0; h < height; ++h) + for (int64_t w = 0; w < width; ++w) + resultValues.push_back((*channelValues)[c]); + + auto resultAttr = DenseElementsAttr::get(resultType, resultValues); + return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp new file mode 100644 index 0000000..6ce84e5 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LogicalResult.h" + +namespace onnx_mlir { + +struct BiasAddPlanCandidate { + mlir::Value data; + mlir::Value bias; +}; + +mlir::LogicalResult isSupportedBiasAddShape(mlir::RankedTensorType biasType, mlir::RankedTensorType resultType); +bool isSupportedBiasAddValue(mlir::Value bias, + mlir::RankedTensorType resultType, + mlir::DenseElementsAttr* denseAttr = nullptr); +mlir::FailureOr> +getBiasChannelValues(mlir::DenseElementsAttr denseAttr, mlir::RankedTensorType resultType); +mlir::FailureOr classifyBiasAddPlanCandidate(mlir::Value lhs, + mlir::Value rhs, + mlir::RankedTensorType resultType); +mlir::FailureOr materializeDenseBiasAddTensor(mlir::Value bias, + mlir::RankedTensorType resultType, + mlir::RewriterBase& rewriter, + mlir::Location loc); + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp new file mode 100644 index 0000000..0553d32 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp @@ -0,0 +1,239 @@ +#include "llvm/ADT/SmallVector.h" + +#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Dialect/ONNX/ONNXOps.hpp" + +using namespace mlir; + +namespace onnx_mlir { + +RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) { + return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)}, + logicalType.getElementType(), + logicalType.getEncoding()); +} + +RankedTensorType getRowStripStorageType(RankedTensorType logicalType) { + return RankedTensorType::get({logicalType.getDimSize(2), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)}, + logicalType.getElementType(), + logicalType.getEncoding()); +} + +std::pair, SmallVector> buildRowStripMetadata(RankedTensorType type) { + SmallVector offsets; + SmallVector sizes; + const int64_t channels = type.getDimSize(1); + const int64_t height = type.getDimSize(2); + const int64_t width = type.getDimSize(3); + offsets.reserve(height * 4); + sizes.reserve(height * 4); + for (int64_t row = 0; row < height; ++row) { + offsets.append({0, 0, row, 0}); + sizes.append({1, channels, 1, width}); + } + return {offsets, sizes}; +} + +SmallVector buildRowStripFragmentOffsets(PatternRewriter& rewriter, OpFoldResult row) { + return {row, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; +} + +SmallVector buildRowStripFragmentSizes(PatternRewriter& rewriter, RankedTensorType logicalType) { + return {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(logicalType.getDimSize(1)), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(logicalType.getDimSize(3))}; +} + +Value extractRowStripFragment(Value storage, + RankedTensorType logicalType, + OpFoldResult row, + PatternRewriter& rewriter, + Location loc) { + return tensor::ExtractSliceOp::create(rewriter, + loc, + getRowStripFragmentType(logicalType), + storage, + buildRowStripFragmentOffsets(rewriter, row), + buildRowStripFragmentSizes(rewriter, logicalType), + getUnitStrides(rewriter, 4)); +} + +void insertRowStripFragment(Value fragment, + Value output, + RankedTensorType logicalType, + OpFoldResult row, + PatternRewriter& rewriter, + Location loc) { + createParallelInsertSliceIntoBatchOutput(rewriter, + loc, + fragment, + output, + buildRowStripFragmentOffsets(rewriter, row), + buildRowStripFragmentSizes(rewriter, logicalType), + getUnitStrides(rewriter, 4)); +} + +FailureOr createPerChannelConstantFragment(DenseElementsAttr denseAttr, + RankedTensorType fragmentType, + PatternRewriter& rewriter) { + FailureOr> channelValues = getBiasChannelValues(denseAttr, fragmentType); + if (failed(channelValues)) + return failure(); + + SmallVector values; + values.reserve(fragmentType.getNumElements()); + for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n) + for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel) + for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h) + for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w) + values.push_back((*channelValues)[channel]); + + auto attr = DenseElementsAttr::get(fragmentType, values); + return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType); +} + +FailureOr createRowStripStorageFromRows(Value rows, + RankedTensorType logicalType, + PatternRewriter& rewriter, + Location loc) { + auto rowsType = dyn_cast(rows.getType()); + if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2) + return failure(); + if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4) + return failure(); + if (logicalType.getDimSize(0) != 1) + return failure(); + if (rowsType.getElementType() != logicalType.getElementType()) + return failure(); + + const int64_t channels = logicalType.getDimSize(1); + const int64_t height = logicalType.getDimSize(2); + const int64_t width = logicalType.getDimSize(3); + if (rowsType.getDimSize(0) != height * width) + return failure(); + if (rowsType.getDimSize(1) != channels) + return failure(); + + auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding()); + auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding()); + auto fragmentType = getRowStripFragmentType(logicalType); + auto storageType = getRowStripStorageType(logicalType); + auto batchOp = createSpatComputeBatch( + rewriter, loc, TypeRange {storageType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) { + Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp); + SmallVector rowOffsets {rowStart, rewriter.getIndexAttr(0)}; + SmallVector rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)}; + Value rowSlice = tensor::ExtractSliceOp::create( + rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2)); + Value channelWidth = ONNXTransposeOp::create( + rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult(); + Value fragment = tensor::ExpandShapeOp::create( + rewriter, loc, fragmentType, channelWidth, SmallVector {{0, 1}, {2, 3}}); + insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc); + return success(); + }); + if (failed(batchOp)) + return failure(); + return batchOp->getResult(0); +} + +FailureOr +materializeRowStripStorageToDense(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) { + auto storageType = dyn_cast(storage.getType()); + if (!storageType || storageType != getRowStripStorageType(logicalType)) + return failure(); + + auto batchOp = createSpatComputeBatch( + rewriter, loc, TypeRange {logicalType}, logicalType.getDimSize(2), {}, ValueRange {storage}, + [&](detail::SpatComputeBatchBodyArgs args) { + Value fragment = extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc); + createParallelInsertSliceIntoBatchOutput(rewriter, + loc, + fragment, + args.outputs.front(), + SmallVector {rewriter.getIndexAttr(0), + rewriter.getIndexAttr(0), + args.lane, + rewriter.getIndexAttr(0)}, + buildRowStripFragmentSizes(rewriter, logicalType), + getUnitStrides(rewriter, 4)); + return success(); + }); + if (failed(batchOp)) + return failure(); + return batchOp->getResult(0); +} + +FailureOr +applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) { + auto fragmentType = getRowStripFragmentType(logicalType); + auto storageType = getRowStripStorageType(logicalType); + auto batchOp = createSpatComputeBatch(rewriter, + loc, + TypeRange {storageType}, + logicalType.getDimSize(2), + {}, + ValueRange {storage}, + [&](detail::SpatComputeBatchBodyArgs args) { + Value fragment = + extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc); + fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult(); + insertRowStripFragment( + fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc); + return success(); + }); + if (failed(batchOp)) + return failure(); + return batchOp->getResult(0); +} + +FailureOr +applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) { + DenseElementsAttr denseAttr; + if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr)) + return failure(); + auto fragmentType = getRowStripFragmentType(logicalType); + auto storageType = getRowStripStorageType(logicalType); + auto batchOp = createSpatComputeBatch(rewriter, + loc, + TypeRange {storageType}, + logicalType.getDimSize(2), + {}, + ValueRange {storage}, + [&](detail::SpatComputeBatchBodyArgs args) { + Value fragment = + extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc); + Value constant; + if (denseAttr.isSplat()) { + constant = getOrCreateConstant( + rewriter, + rewriter.getInsertionBlock()->getParentOp(), + DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue()), + fragmentType); + } + else { + FailureOr perChannel = + createPerChannelConstantFragment(denseAttr, fragmentType, rewriter); + if (failed(perChannel)) + return failure(); + constant = *perChannel; + } + fragment = + spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult(); + insertRowStripFragment( + fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc); + return success(); + }); + if (failed(batchOp)) + return failure(); + return batchOp->getResult(0); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp new file mode 100644 index 0000000..3a436ba --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp @@ -0,0 +1,69 @@ +#pragma once + +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/PatternMatch.h" + +namespace onnx_mlir { + +inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments"; + +struct RowStripPhysicalValue { + mlir::Value storage; + mlir::RankedTensorType logicalType; + llvm::SmallVector fragmentOffsets; + llvm::SmallVector fragmentSizes; +}; + +std::pair, llvm::SmallVector> +buildRowStripMetadata(mlir::RankedTensorType type); + +mlir::RankedTensorType getRowStripFragmentType(mlir::RankedTensorType logicalType); + +mlir::RankedTensorType getRowStripStorageType(mlir::RankedTensorType logicalType); + +llvm::SmallVector buildRowStripFragmentOffsets(mlir::PatternRewriter& rewriter, + mlir::OpFoldResult row); + +llvm::SmallVector buildRowStripFragmentSizes(mlir::PatternRewriter& rewriter, + mlir::RankedTensorType logicalType); + +mlir::Value extractRowStripFragment(mlir::Value storage, + mlir::RankedTensorType logicalType, + mlir::OpFoldResult row, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +void insertRowStripFragment(mlir::Value fragment, + mlir::Value output, + mlir::RankedTensorType logicalType, + mlir::OpFoldResult row, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +mlir::FailureOr createPerChannelConstantFragment(mlir::DenseElementsAttr denseAttr, + mlir::RankedTensorType fragmentType, + mlir::PatternRewriter& rewriter); + +mlir::FailureOr createRowStripStorageFromRows(mlir::Value rows, + mlir::RankedTensorType logicalType, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +mlir::FailureOr materializeRowStripStorageToDense(mlir::Value storage, + mlir::RankedTensorType logicalType, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +mlir::FailureOr applyRowStripRelu(mlir::Value storage, + mlir::RankedTensorType logicalType, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +mlir::FailureOr applyRowStripBiasAdd(mlir::Value storage, + mlir::RankedTensorType logicalType, + mlir::Value bias, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp index e674fa2..6ef425f 100644 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -13,7 +13,9 @@ #include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp" @@ -29,14 +31,6 @@ namespace { static constexpr StringLiteral kDenseLayout = "dense_nchw"; static constexpr StringLiteral kRowStripLayout = "nchw_row_strip"; -struct RowStripPhysicalValue { - Value physicalValue; - RankedTensorType logicalType; - SmallVector fragmentOffsets; - SmallVector fragmentSizes; - std::string indexMap; -}; - static FailureOr getRowStripValue(llvm::DenseMap& rowStripValues, Value value) { auto it = rowStripValues.find(value); @@ -46,112 +40,42 @@ static FailureOr getRowStripValue(llvm::DenseMap buildRowStripValue(spatial::SpatBlueprintOp blueprint, - Value physicalValue) { + Value storage) { auto logicalType = dyn_cast(blueprint.getOutput().getType()); if (!logicalType) return blueprint.emitOpError("requires ranked logical output type"), failure(); RowStripPhysicalValue value; - value.physicalValue = physicalValue; + value.storage = storage; value.logicalType = logicalType; value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end()); value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end()); - value.indexMap = blueprint.getIndexMap().str(); + if (blueprint.getIndexMap() != kRowStripIndexMap) + return blueprint.emitOpError("requires the canonical row-strip index map"), failure(); + auto storageType = dyn_cast(storage.getType()); + if (!storageType || storageType != getRowStripStorageType(logicalType)) + return blueprint.emitOpError("requires physical row-strip fragment storage"), failure(); return value; } static FailureOr lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) { - auto packedType = cast(input.physicalValue.getType()); - auto computeOp = - createSpatCompute<1>(rewriter, planOp.getLoc(), TypeRange {packedType}, {}, input.physicalValue, [&](Value x) { - auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), packedType, x); - spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult()); - }); - return computeOp.getResult(0); + return applyRowStripRelu(input.storage, input.logicalType, rewriter, planOp.getLoc()); +} + +static FailureOr lowerRowStripBiasAdd(const RowStripPhysicalValue& input, + spatial::SpatBiasAddPlanOp planOp, + PatternRewriter& rewriter) { + return applyRowStripBiasAdd(input.storage, input.logicalType, planOp.getBias(), rewriter, planOp.getLoc()); } static FailureOr materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) { - auto packedType = dyn_cast(rowStripValue.physicalValue.getType()); - if (!packedType || packedType.getRank() != 3 || !packedType.hasStaticShape()) - return failure(); if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape()) return failure(); - if (rowStripValue.indexMap != "packed_hwc_rows_to_nchw") + auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType); + if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes)) return failure(); - - const int64_t rank = rowStripValue.logicalType.getRank(); - const int64_t fragmentCount = rowStripValue.fragmentOffsets.size() / rank; - const int64_t packedWidth = packedType.getDimSize(1); - const int64_t packedChannels = packedType.getDimSize(2); - if (fragmentCount != packedType.getDimSize(0)) - return failure(); - for (int64_t fragmentIndex = 0; fragmentIndex < fragmentCount; ++fragmentIndex) { - if (rowStripValue.fragmentOffsets[fragmentIndex * rank + 0] != 0 - || rowStripValue.fragmentOffsets[fragmentIndex * rank + 1] != 0 - || rowStripValue.fragmentOffsets[fragmentIndex * rank + 2] != fragmentIndex - || rowStripValue.fragmentOffsets[fragmentIndex * rank + 3] != 0) - return failure(); - if (rowStripValue.fragmentSizes[fragmentIndex * rank + 0] != 1 - || rowStripValue.fragmentSizes[fragmentIndex * rank + 1] != packedChannels - || rowStripValue.fragmentSizes[fragmentIndex * rank + 2] != 1 - || rowStripValue.fragmentSizes[fragmentIndex * rank + 3] != packedWidth) - return failure(); - } - - auto packedSliceType = - RankedTensorType::get({1, packedWidth, packedChannels}, packedType.getElementType(), packedType.getEncoding()); - auto expandedType = - RankedTensorType::get({1, 1, packedWidth, packedChannels}, packedType.getElementType(), packedType.getEncoding()); - auto logicalFragmentType = - RankedTensorType::get({1, packedChannels, 1, packedWidth}, packedType.getElementType(), packedType.getEncoding()); - auto batchOp = createSpatComputeBatch( - rewriter, - loc, - TypeRange {rowStripValue.logicalType}, - fragmentCount, - {}, - ValueRange {rowStripValue.physicalValue}, - [&](detail::SpatComputeBatchBodyArgs args) { - SmallVector packedOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector packedSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(packedWidth), rewriter.getIndexAttr(packedChannels)}; - Value packedSlice = tensor::ExtractSliceOp::create( - rewriter, loc, packedSliceType, args.inputs.front(), packedOffsets, packedSizes, getUnitStrides(rewriter, 3)); - - Value expanded = tensor::ExpandShapeOp::create(rewriter, - loc, - expandedType, - packedSlice, - SmallVector { - {0, 1}, - {2}, - {3} - }); - Value transposeInit = - tensor::EmptyOp::create(rewriter, loc, logicalFragmentType.getShape(), logicalFragmentType.getElementType()); - Value logicalFragment = - linalg::TransposeOp::create(rewriter, loc, expanded, transposeInit, SmallVector {0, 3, 1, 2}) - .getResult()[0]; - - SmallVector logicalOffsets { - rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), args.lane, rewriter.getIndexAttr(0)}; - SmallVector logicalSizes {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(packedChannels), - rewriter.getIndexAttr(1), - rewriter.getIndexAttr(packedWidth)}; - createParallelInsertSliceIntoBatchOutput(rewriter, - loc, - logicalFragment, - args.outputs.front(), - logicalOffsets, - logicalSizes, - getUnitStrides(rewriter, 4)); - return success(); - }); - if (failed(batchOp)) - return failure(); - return batchOp->getResult(0); + return materializeRowStripStorageToDense(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc); } struct LowerSpatialPlansPass final : PassWrapper> { @@ -194,7 +118,7 @@ struct LowerSpatialPlansPass final : PassWrapper lowered = lowerSelectedConv2DPlan( planOp, - succeeded(rowStripInput) ? std::optional {rowStripInput->physicalValue} : std::nullopt, + succeeded(rowStripInput) ? std::optional {rowStripInput->storage} : std::nullopt, /*emitRowStripLayout=*/true, rewriter); if (failed(lowered)) { @@ -266,6 +190,64 @@ struct LowerSpatialPlansPass final : PassWrapper(&op)) { + if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) { + auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { + auto blueprint = dyn_cast(user); + return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout; + }); + if (outputBlueprint == planOp.getResult().getUsers().end()) { + planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result"); + signalPassFailure(); + return; + } + + FailureOr input = getRowStripValue(rowStripValues, planOp.getInput()); + rewriter.setInsertionPoint(planOp); + FailureOr lowered = lowerRowStripBiasAdd(*input, planOp, rewriter); + if (failed(lowered)) { + planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan"); + signalPassFailure(); + return; + } + auto blueprint = cast(*outputBlueprint); + FailureOr output = buildRowStripValue(blueprint, *lowered); + if (failed(output)) { + signalPassFailure(); + return; + } + rowStripValues[blueprint.getResult()] = *output; + eraseAfterLowering.insert(planOp); + eraseAfterLowering.insert(blueprint); + continue; + } + + auto resultType = dyn_cast(planOp.getOutput().getType()); + if (!resultType) { + planOp.emitOpError("requires ranked output type"); + signalPassFailure(); + return; + } + rewriter.setInsertionPoint(planOp); + FailureOr denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc()); + if (failed(denseBias)) { + planOp.emitOpError("failed to materialize dense Conv-style bias"); + signalPassFailure(); + return; + } + auto computeOp = createSpatCompute<2>(rewriter, + planOp.getLoc(), + planOp.getOutput().getType(), + {}, + ValueRange {planOp.getInput(), *denseBias}, + [&](Value x, Value y) { + auto added = spatial::SpatVAddOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult()); + }); + rewriter.replaceOp(planOp, computeOp.getResults()); + continue; + } if (auto materializeOp = dyn_cast(&op)) { if (materializeOp.getSourcePhysicalLayout() == kDenseLayout && materializeOp.getTargetPhysicalLayout() == kDenseLayout) { @@ -385,6 +367,7 @@ struct LowerSpatialPlansPass final : PassWrapper(op)) return; if (isa(op) diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp index acba5f0..118fc25 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp @@ -45,11 +45,12 @@ static void populateEmptyFunction(func::FuncOp funcOp) { SmallVector computes(funcOp.getOps()); SmallVector computeBatches(funcOp.getOps()); SmallVector convPlans(funcOp.getOps()); + SmallVector biasAddPlans(funcOp.getOps()); SmallVector reluPlans(funcOp.getOps()); SmallVector blueprints(funcOp.getOps()); SmallVector materializers(funcOp.getOps()); - if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !reluPlans.empty() || !blueprints.empty() - || !materializers.empty()) { + if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !reluPlans.empty() + || !blueprints.empty() || !materializers.empty()) { return; } diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp index bbb6604..f6cd5aa 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp @@ -132,6 +132,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch, spatial::SpatConv2DPlanOp, + spatial::SpatBiasAddPlanOp, spatial::SpatReluPlanOp, spatial::SpatBlueprintOp, spatial::SpatMaterializeLayoutOp>(&op)) { diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp index e41583a..3f142b9 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -23,7 +23,9 @@ #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp" @@ -221,10 +223,6 @@ struct DistributedConvReportTotals { }; static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter); -static FailureOr createRowStripPackedRows(Value rows, - const ConvLoweringState& state, - PatternRewriter& rewriter, - Location loc); static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b); @@ -467,59 +465,6 @@ classifyDistributedBinaryConsumer(Operation* user, return std::nullopt; } -static bool covers(RowInterval acquired, RowInterval needed) { - return acquired.begin <= needed.begin && acquired.end >= needed.end; -} - -static bool canConsumeRowStripHwcInput(const ConvLoweringState& state, StringRef& failureReason) { - if (state.batchSize != 1) { - failureReason = "unsupported_batch"; - return false; - } - if (state.group != 1) { - failureReason = "unsupported_groups"; - return false; - } - if (isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup)) { - failureReason = "unsupported_depthwise"; - return false; - } - if (state.strideHeight != 1 || state.strideWidth != 1) { - failureReason = "unsupported_stride"; - return false; - } - if (state.dilationHeight != 1 || state.dilationWidth != 1) { - failureReason = "unsupported_dilation"; - return false; - } - if (state.padHeightBegin != state.padHeightEnd || state.padWidthBegin != state.padWidthEnd) { - failureReason = "unsupported_padding"; - return false; - } - if (state.padHeightBegin != 1 || state.padWidthBegin != 1) { - failureReason = "unsupported_padding"; - return false; - } - if (state.wHeight != 3 || state.wWidth != 3) { - failureReason = "unsupported_kernel"; - return false; - } - if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) { - failureReason = "unsupported_output_shape"; - return false; - } - if (!getHostConstDenseElementsAttr(state.w)) { - failureReason = "non_constant_weights"; - return false; - } - if (state.hasBias && !getHostConstDenseElementsAttr(state.b)) { - failureReason = "non_constant_bias"; - return false; - } - failureReason = ""; - return true; -} - static std::string stringifyDistributedTensorOpKind(DistributedTensorOpKind kind) { switch (kind) { case DistributedTensorOpKind::Relu: return "Relu"; @@ -2493,21 +2438,15 @@ static Value rewriteStreamedConv(const ConvLoweringState& state, } // namespace standard -static RankedTensorType getRowStripFragmentType(RankedTensorType tensorType, int64_t width) { - return RankedTensorType::get( - {tensorType.getDimSize(0), tensorType.getDimSize(1), 1, width}, tensorType.getElementType(), tensorType.getEncoding()); -} - static SmallVector buildRowStripFragments(RankedTensorType tensorType) { SmallVector fragments; - const int64_t height = tensorType.getDimSize(2); - const int64_t width = tensorType.getDimSize(3); - const int64_t channels = tensorType.getDimSize(1); - fragments.reserve(height); - for (int64_t row = 0; row < height; ++row) { + auto [offsets, sizes] = buildRowStripMetadata(tensorType); + const int64_t rank = tensorType.getRank(); + fragments.reserve(offsets.size() / rank); + for (int64_t row = 0; row < static_cast(offsets.size() / rank); ++row) { fragments.push_back(DistributedFragmentInfo { - {0, 0, row, 0}, - {1, channels, 1, width}, + {offsets.begin() + row * rank, offsets.begin() + (row + 1) * rank}, + {sizes.begin() + row * rank, sizes.begin() + (row + 1) * rank}, {1, 1, 1, 1}, row, }); @@ -2527,101 +2466,266 @@ static DistributedTensorInfo makeDistributedTensorInfo(Value storage, RankedTens return info; } -static Value createPerChannelConstantFragment(DenseElementsAttr denseAttr, - RankedTensorType fragmentType, - PatternRewriter& rewriter) { - auto denseType = cast(denseAttr.getType()); - SmallVector channelValues; - channelValues.reserve(fragmentType.getDimSize(1)); - SmallVector flattened(denseAttr.getValues()); - if (denseType.getRank() == 1) { - channelValues = flattened; - } - else if (denseType.getRank() == 2) { - channelValues = flattened; - } - else { - for (int64_t channel = 0; channel < denseType.getDimSize(1); ++channel) - channelValues.push_back(flattened[channel]); - } - - SmallVector values; - values.reserve(fragmentType.getNumElements()); - for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n) - for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel) - for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h) - for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w) - values.push_back(channelValues[channel]); - - auto attr = DenseElementsAttr::get(fragmentType, values); - return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType); -} - static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter) { auto zeroAttr = DenseElementsAttr::get(gemmResultType, rewriter.getZeroAttr(gemmResultType.getElementType())); return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, gemmResultType); } -static bool canDirectLowerRowStripConv(const ConvLoweringState& state, StringRef& failureReason) { - if (!canConsumeRowStripHwcInput(state, failureReason)) - return false; - - ConvGeometry geometry = buildConvGeometry(state); - if (state.numChannelsOut > geometry.xbarSize) { - failureReason = "unsupported_output_channels"; +static bool canConsumeNchwRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) { + if (state.batchSize != 1) { + failureReason = "batch_not_one"; + return false; + } + if (state.group != 1) { + failureReason = "grouped_conv"; + return false; + } + if (!state.xType.hasStaticShape() || !state.wType.hasStaticShape() || !state.outType.hasStaticShape()) { + failureReason = "dynamic_shape"; + return false; + } + if (!isa(state.xType.getElementType())) { + failureReason = "non_float_input"; + return false; + } + if (state.strideHeight != 1 || state.strideWidth != 1) { + failureReason = "stride_not_one"; + return false; + } + if (state.dilationHeight != 1 || state.dilationWidth != 1) { + failureReason = "dilation_not_one"; + return false; + } + if (state.wHeight != 3 || state.wWidth != 3) { + failureReason = "kernel_not_3x3"; + return false; + } + if (state.padHeightBegin != 1 || state.padHeightEnd != 1 || state.padWidthBegin != 1 || state.padWidthEnd != 1) { + failureReason = "padding_not_1"; + return false; + } + if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) { + failureReason = "not_same_spatial_shape"; + return false; + } + if (!getHostConstDenseElementsAttr(state.w)) { + failureReason = "non_constant_weight"; + return false; + } + if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) { + failureReason = "unsupported_bias"; + return false; + } + ConvGeometry geometry = buildConvGeometry(state); + if (geometry.c > geometry.xbarSize) { + failureReason = "output_channels_exceed_crossbar"; return false; } - - failureReason = ""; return true; } -static FailureOr createRowStripPackedRows(Value rows, - const ConvLoweringState& state, - PatternRewriter& rewriter, - Location loc) { - auto rowsType = dyn_cast(rows.getType()); - if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2) - return failure(); - - if (state.batchSize != 1) - return failure(); - if (state.outType.getRank() != 4 || !state.outType.hasStaticShape()) - return failure(); - - const int64_t outHeight = state.outType.getDimSize(2); - const int64_t outWidth = state.outType.getDimSize(3); - const int64_t outChannels = state.outType.getDimSize(1); - if (rowsType.getDimSize(0) != outHeight * outWidth || rowsType.getDimSize(1) != outChannels) - return failure(); - - auto packedType = RankedTensorType::get({outHeight, outWidth, outChannels}, rowsType.getElementType(), rowsType.getEncoding()); - auto packedRows = - createSpatCompute<1>(rewriter, loc, TypeRange {packedType}, {}, rows, [&](Value rowValues) { - Value packed = tensor::ExpandShapeOp::create( - rewriter, loc, packedType, rowValues, SmallVector {{0, 1}, {2}}); - spatial::SpatYieldOp::create(rewriter, loc, packed); - }); - return packedRows.getResult(0); +static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& rewriter) { + auto zeroAttr = DenseElementsAttr::get(type, rewriter.getZeroAttr(type.getElementType())); + return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, type); } -static FailureOr createConvOutputFromRowStripHwc(Value inputHwc, - const ConvLoweringState& state, - PatternRewriter& rewriter, - Location loc) { - auto inputType = dyn_cast(inputHwc.getType()); - if (!inputType || !inputType.hasStaticShape() || inputType.getRank() != 3) +static FailureOr createPaddedBiasRowConstant(const ConvLoweringState& state, + int64_t paddedChannels, + PatternRewriter& rewriter) { + DenseElementsAttr denseAttr; + if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr)) return failure(); - if (inputType.getDimSize(0) != state.xHeight || inputType.getDimSize(1) != state.xWidth - || inputType.getDimSize(2) != state.numChannelsIn) + FailureOr> channelValues = getBiasChannelValues(denseAttr, state.outType); + if (failed(channelValues)) + return failure(); + + auto biasType = RankedTensorType::get({1, paddedChannels}, state.outType.getElementType()); + SmallVector values(biasType.getNumElements(), cast(rewriter.getZeroAttr(biasType.getElementType()))); + for (int64_t channel = 0; channel < state.numChannelsOut; ++channel) + values[channel] = (*channelValues)[channel]; + auto biasAttr = DenseElementsAttr::get(biasType, values); + return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), biasAttr, biasType); +} + +static Value createHorizontallyPaddedRowStripFragment(Value fragment, + const ConvLoweringState& state, + PatternRewriter& rewriter, + Location loc) { + auto paddedType = RankedTensorType::get({1, state.numChannelsIn, 1, state.xWidth + 2}, + state.xType.getElementType(), + state.xType.getEncoding()); + return createZeroPaddedTensor(fragment, paddedType, {0, 0, 0, 1}, {0, 0, 0, 1}, rewriter, loc); +} + +static Value createRowStripWindowSourceRowTable(const ConvLoweringState& state, PatternRewriter& rewriter) { + Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + auto tableType = RankedTensorType::get({state.outHeight * state.wHeight}, rewriter.getIndexType()); + SmallVector values; + values.reserve(tableType.getNumElements()); + for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) { + for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) { + int64_t sourceRow = outputRow + kernelRow - state.padHeightBegin; + sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1); + values.push_back(rewriter.getIndexAttr(sourceRow)); + } + } + + return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType); +} + +static Value createRowStripWindowTableIndex(Value outputHeight, + Value kernelRow, + const ConvLoweringState& state, + PatternRewriter& rewriter, + Location loc) { + Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + MLIRContext* ctx = rewriter.getContext(); + AffineExpr outputRowExpr = getAffineDimExpr(0, ctx); + AffineExpr kernelRowExpr = getAffineDimExpr(1, ctx); + return createOrFoldAffineApply( + rewriter, loc, outputRowExpr * state.wHeight + kernelRowExpr, ValueRange {outputHeight, kernelRow}, anchorOp); +} + +static Value extractProjectedRowStripWindowRow(Value rowStripStorage, + Value sourceRowTable, + const ConvLoweringState& state, + Value outputHeight, + Value kernelRow, + PatternRewriter& rewriter, + Location loc) { + Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc); + Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult(); + return extractRowStripFragment(rowStripStorage, state.xType, sourceRow, rewriter, loc); +} + +static FailureOr createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) { + auto elementType = state.xType.getElementType(); + auto floatType = dyn_cast(elementType); + if (!floatType) + return failure(); + + Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + auto tableType = RankedTensorType::get({state.outHeight * state.wHeight, state.numChannelsIn, 1, state.xWidth}, + elementType, + state.xType.getEncoding()); + Attribute zero = rewriter.getZeroAttr(elementType); + Attribute one = rewriter.getFloatAttr(floatType, 1.0); + SmallVector values; + values.reserve(tableType.getNumElements()); + for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) { + for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) { + int64_t sourceRow = outputRow + kernelRow - state.padHeightBegin; + Attribute value = (sourceRow < 0 || sourceRow >= state.xHeight) ? zero : one; + for (int64_t channel = 0; channel < state.numChannelsIn; ++channel) + for (int64_t width = 0; width < state.xWidth; ++width) + values.push_back(value); + } + } + + return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType); +} + +static Value extractProjectedRowStripWindowMask(Value maskTable, + const ConvLoweringState& state, + Value outputHeight, + Value kernelRow, + PatternRewriter& rewriter, + Location loc) { + Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc); + auto fragmentType = getRowStripFragmentType(state.xType); + SmallVector offsets { + tableIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; + SmallVector sizes {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(state.numChannelsIn), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(state.xWidth)}; + return tensor::ExtractSliceOp::create(rewriter, + loc, + fragmentType, + maskTable, + offsets, + sizes, + getUnitStrides(rewriter, 4)); +} + +static FailureOr createNchwRowStripConvWindow(Value rowStripStorage, + const ConvLoweringState& state, + Value outputHeight, + PatternRewriter& rewriter, + Location loc) { + auto fragmentType = getRowStripFragmentType(state.xType); + auto paddedWindowType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.xWidth + 2}, + state.xType.getElementType(), + state.xType.getEncoding()); + Value sourceRowTable = createRowStripWindowSourceRowTable(state, rewriter); + FailureOr maskTable = createRowStripWindowMaskTable(state, rewriter); + if (failed(maskTable)) + return failure(); + Value initWindow = createZeroTensorConstant(paddedWindowType, rewriter); + + Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + Value window = initWindow; + for (int64_t kernelRowIndex = 0; kernelRowIndex < state.wHeight; ++kernelRowIndex) { + Value kernelRow = getOrCreateIndexConstant(rewriter, anchorOp, kernelRowIndex); + Value sourceRow = + extractProjectedRowStripWindowRow(rowStripStorage, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc); + Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, loc); + Value semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, sourceRow, mask).getResult(); + Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, loc); + window = tensor::InsertSliceOp::create(rewriter, + loc, + paddedRow, + window, + SmallVector {rewriter.getIndexAttr(0), + rewriter.getIndexAttr(0), + rewriter.getIndexAttr(kernelRowIndex), + rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(state.numChannelsIn), + rewriter.getIndexAttr(1), + rewriter.getIndexAttr(state.xWidth + 2)}, + getUnitStrides(rewriter, 4)); + } + return window; +} + +static FailureOr createNchwRowStripConvPatchRow(Value paddedWindow, + const ConvLoweringState& state, + Value outputWidth, + PatternRewriter& rewriter, + Location loc) { + Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; + auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, + state.xType.getElementType(), + state.xType.getEncoding()); + auto rowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType(), state.xType.getEncoding()); + Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); + Value patch = createConvInputPatch(paddedWindow, + patchType, + c0, + c0, + c0, + outputWidth, + state.dilationHeight, + state.dilationWidth, + rewriter, + loc); + return tensor::CollapseShapeOp::create( + rewriter, loc, rowType, patch, SmallVector {{0}, {1, 2, 3}}) + .getResult(); +} + +static FailureOr createConvOutputFromNchwRowStripFragments(Value rowStripStorage, + const ConvLoweringState& state, + PatternRewriter& rewriter, + Location loc) { + auto inputType = dyn_cast(rowStripStorage.getType()); + if (!inputType || inputType != getRowStripStorageType(state.xType)) return failure(); StringRef failureReason; - if (!canDirectLowerRowStripConv(state, failureReason)) - return failure(); - - ConvRowDemand demand = buildConvRowDemand(RowInterval {0, state.outHeight}, state); - if (!covers(demand.acquiredInputRows, demand.neededInputRows)) + if (!canConsumeNchwRowStripFragments(state, failureReason)) return failure(); ConvGeometry geometry = buildConvGeometry(state); @@ -2629,121 +2733,61 @@ static FailureOr createConvOutputFromRowStripHwc(Value inputHwc, const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim); const int64_t paddedK = numKSlices * xbarDim; - auto elementType = inputType.getElementType(); - auto paddedInputType = RankedTensorType::get({state.xHeight + state.padHeightBegin + state.padHeightEnd, - state.xWidth + state.padWidthBegin + state.padWidthEnd, - state.numChannelsIn}, - elementType, - inputType.getEncoding()); - auto paddedPatchType = - RankedTensorType::get({state.wHeight, state.wWidth, 1}, elementType, inputType.getEncoding()); - auto flatPatchType = RankedTensorType::get({state.wHeight * state.wWidth}, elementType, inputType.getEncoding()); - auto rowChunkType = RankedTensorType::get({1, state.wHeight * state.wWidth}, elementType, inputType.getEncoding()); + auto elementType = state.outType.getElementType(); auto rowType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType()); - auto packedOutputType = - RankedTensorType::get({state.outHeight, state.outWidth, state.numChannelsOut}, state.outType.getElementType()); - auto packedOutputSliceType = - RankedTensorType::get({1, 1, state.numChannelsOut}, state.outType.getElementType()); + auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType); auto paddedRowType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType()); auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType, inputType.getEncoding()); auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType()); + auto outputStorageType = getRowStripStorageType(state.outType); auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); if (!weightDenseAttr) return failure(); Value paddedWeights = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter); - - Value paddedBias; - if (state.hasBias) { - Value biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc); - auto biasMatrixType = cast(biasMatrix.getType()); - auto paddedBiasType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType()); - if (auto biasDenseAttr = getHostConstDenseElementsAttr(state.b)) - paddedBias = standard::createPaddedConstantMatrix(biasDenseAttr, biasMatrixType, paddedBiasType, rewriter); - else - paddedBias = materializeOrComputeUnary( - biasMatrix, paddedBiasType, rewriter, loc, [&](Value biasValue) { - return standard::createPaddedConvMatrix(biasValue, biasMatrixType, paddedBiasType, rewriter, loc); - }); - } - - auto paddedInputOp = - createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, inputHwc, [&](Value hwcInputArg) { - Value paddedInput = createZeroPaddedTensor(hwcInputArg, - paddedInputType, - {state.padHeightBegin, state.padWidthBegin, 0}, - {state.padHeightEnd, state.padWidthEnd, 0}, - rewriter, - loc); - spatial::SpatYieldOp::create(rewriter, loc, paddedInput); - }); - - SmallVector batchInputs {paddedInputOp.getResult(0)}; + FailureOr paddedBias = failure(); if (state.hasBias) - batchInputs.push_back(paddedBias); + paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter); + if (state.hasBias && failed(paddedBias)) + return failure(); + auto batchOp = createSpatComputeBatch( rewriter, loc, - TypeRange {packedOutputType}, + TypeRange {outputStorageType}, state.outHeight, ValueRange {paddedWeights}, - batchInputs, + state.hasBias ? ValueRange {rowStripStorage, *paddedBias} : ValueRange {rowStripStorage}, [&](detail::SpatComputeBatchBodyArgs args) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices); Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth); - Value cNumChannels = getOrCreateIndexConstant(rewriter, anchorOp, state.numChannelsIn); - Value localHeightOffset = args.lane; - Value packedRowInit = - tensor::EmptyOp::create(rewriter, loc, ArrayRef {1, state.outWidth, state.numChannelsOut}, elementType); + Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim); + auto fragmentType = getRowStripFragmentType(state.outType); + FailureOr inputWindow = createNchwRowStripConvWindow(args.inputs.front(), state, args.lane, rewriter, loc); + if (failed(inputWindow)) + return failure(); + Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType); auto widthLoop = buildNormalizedScfFor( rewriter, loc, c0, cOutWidth, c1, - ValueRange {packedRowInit}, + ValueRange {fragmentInit}, [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl& widthYielded) { - Value localWidthOffset = widthIndex; - Value rowInit = tensor::EmptyOp::create(rewriter, widthLoc, ArrayRef {1, patchSize}, elementType); - auto rowLoop = buildNormalizedScfFor( - rewriter, - widthLoc, - c0, - cNumChannels, - c1, - ValueRange {rowInit}, - [&](OpBuilder&, Location rowLoc, Value channel, ValueRange rowIterArgs, SmallVectorImpl& rowYielded) { - SmallVector patchOffsets {localHeightOffset, localWidthOffset, channel}; - SmallVector patchSizes { - rewriter.getIndexAttr(state.wHeight), rewriter.getIndexAttr(state.wWidth), rewriter.getIndexAttr(1)}; - Value channelPatch = tensor::ExtractSliceOp::create( - rewriter, rowLoc, paddedPatchType, args.inputs.front(), patchOffsets, patchSizes, getUnitStrides(rewriter, 3)); - Value flatPatch = tensor::CollapseShapeOp::create( - rewriter, rowLoc, flatPatchType, channelPatch, SmallVector {{0, 1, 2}}); - Value rowChunk = tensor::ExpandShapeOp::create( - rewriter, rowLoc, rowChunkType, flatPatch, SmallVector {{0, 1}}); - Value flatOffset = affineMulConst( - rewriter, rowLoc, channel, state.wHeight * state.wWidth, anchorOp); - SmallVector rowOffsets {rewriter.getIndexAttr(0), flatOffset}; - SmallVector rowSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.wHeight * state.wWidth)}; - Value nextRow = tensor::InsertSliceOp::create( - rewriter, rowLoc, rowChunk, rowIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2)); - rowYielded.push_back(nextRow); - return success(); - }); - if (failed(rowLoop)) + FailureOr patchRow = + createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc); + if (failed(patchRow)) return failure(); - Value paddedRow = rowLoop->results.front(); + Value paddedRow = *patchRow; if (patchSize != paddedK) paddedRow = createZeroPaddedTensor( paddedRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc); - auto zeroAttr = DenseElementsAttr::get(paddedRowType, rewriter.getZeroAttr(state.outType.getElementType())); - Value zeroRow = getOrCreateConstant(rewriter, anchorOp, zeroAttr, paddedRowType); + Value zeroRow = createZeroTensorConstant(paddedRowType, rewriter); auto kLoop = buildNormalizedScfFor( rewriter, widthLoc, @@ -2752,7 +2796,7 @@ static FailureOr createConvOutputFromRowStripHwc(Value inputHwc, c1, ValueRange {zeroRow}, [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl& reduceYielded) { - Value kOffset = affineMulConst(rewriter, reduceLoc, kSlice, xbarDim, anchorOp); + Value kOffset = arith::MulIOp::create(rewriter, reduceLoc, kSlice, cXbar); SmallVector aOffsets {rewriter.getIndexAttr(0), kOffset}; SmallVector aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)}; Value aTile = tensor::ExtractSliceOp::create( @@ -2776,9 +2820,7 @@ static FailureOr createConvOutputFromRowStripHwc(Value inputHwc, Value rowResult = kLoop->results.front(); if (state.hasBias) - rowResult = - spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, rowResult, args.inputs[1]).getResult(); - + rowResult = spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, rowResult, args.inputs[1]).getResult(); Value outputRow = rowResult; if (state.numChannelsOut != xbarDim) { SmallVector outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; @@ -2790,25 +2832,23 @@ static FailureOr createConvOutputFromRowStripHwc(Value inputHwc, Value outputFragment = tensor::ExpandShapeOp::create(rewriter, widthLoc, - packedOutputSliceType, + outputPixelType, outputRow, - SmallVector {{0}, {1, 2}}); - SmallVector rowOffsets {rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)}; + SmallVector {{0}, {1, 2, 3}}); + SmallVector rowOffsets { + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex}; SmallVector rowSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)}; - Value nextPackedRow = tensor::InsertSliceOp::create( - rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 3)); - widthYielded.push_back(nextPackedRow); + rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1)}; + Value nextFragment = tensor::InsertSliceOp::create( + rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4)); + widthYielded.push_back(nextFragment); return success(); }); if (failed(widthLoop)) return failure(); - SmallVector batchOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector batchSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.outWidth), rewriter.getIndexAttr(state.numChannelsOut)}; - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, widthLoop->results.front(), args.outputs.front(), batchOffsets, batchSizes, getUnitStrides(rewriter, 3)); + insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc); return success(); }); if (failed(batchOp)) @@ -2816,19 +2856,22 @@ static FailureOr createConvOutputFromRowStripHwc(Value inputHwc, return batchOp->getResult(0); } -static FailureOr createConvRowsFromRowStripInput(const ConvLoweringState& state, - [[maybe_unused]] const ConvLoweringDecision& decision, - Value rowStripInput, - PatternRewriter& rewriter, - Location loc) { - return createConvOutputFromRowStripHwc(rowStripInput, state, rewriter, loc); +static FailureOr createConvOutputFromRowStripInput(const ConvLoweringState& state, + [[maybe_unused]] const ConvLoweringDecision& decision, + Value rowStripInput, + PatternRewriter& rewriter, + Location loc) { + return createConvOutputFromNchwRowStripFragments(rowStripInput, state, rewriter, loc); } static Value createFragmentConstant(const DistributedTensorStep& step, RankedTensorType fragmentType, PatternRewriter& rewriter) { - if (step.constantKind == DistributedTensorConstantKind::PerChannel) - return createPerChannelConstantFragment(step.constantAttr, fragmentType, rewriter); + if (step.constantKind == DistributedTensorConstantKind::PerChannel) { + FailureOr constant = createPerChannelConstantFragment(step.constantAttr, fragmentType, rewriter); + assert(succeeded(constant) && "distributed per-channel constants are classified before lowering"); + return *constant; + } Attribute splatValue = step.constantAttr.getSplatValue(); return getOrCreateConstant(rewriter, @@ -2938,67 +2981,22 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step, } } -[[maybe_unused]] static FailureOr createDistributedTensorFromRows(Value rows, - RankedTensorType logicalType, - PatternRewriter& rewriter, - Location loc) { - const int64_t width = logicalType.getDimSize(3); - const int64_t height = logicalType.getDimSize(2); - auto rowsType = cast(rows.getType()); - auto rowSliceType = - RankedTensorType::get({width, logicalType.getDimSize(1)}, logicalType.getElementType(), rowsType.getEncoding()); - auto channelWidthType = - RankedTensorType::get({logicalType.getDimSize(1), width}, logicalType.getElementType(), rowsType.getEncoding()); - auto fragmentType = getRowStripFragmentType(logicalType, width); - auto batchOp = createSpatComputeBatch( - rewriter, loc, TypeRange {logicalType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) { - Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp); - SmallVector rowOffsets {rowStart, rewriter.getIndexAttr(0)}; - SmallVector rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(logicalType.getDimSize(1))}; - Value rowSlice = tensor::ExtractSliceOp::create( - rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2)); - Value channelWidth = ONNXTransposeOp::create( - rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult(); - Value fragment = tensor::ExpandShapeOp::create(rewriter, - loc, - fragmentType, - channelWidth, - SmallVector {{0, 1}, {2, 3}}); - SmallVector outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), args.lane, - rewriter.getIndexAttr(0)}; - SmallVector outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(logicalType.getDimSize(1)), - rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)}; - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, fragment, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 4)); - return success(); - }); - if (failed(batchOp)) - return failure(); - return makeDistributedTensorInfo(batchOp->getResult(0), logicalType); -} - [[maybe_unused]] static FailureOr applyDistributedPreservingStep(const DistributedTensorInfo& inputInfo, const DistributedTensorStep& step, PatternRewriter& rewriter, Location loc) { auto logicalType = inputInfo.logicalType; - const int64_t width = logicalType.getDimSize(3); - auto fragmentType = getRowStripFragmentType(logicalType, width); + auto fragmentType = getRowStripFragmentType(logicalType); + auto storageType = getRowStripStorageType(logicalType); auto batchOp = createSpatComputeBatch(rewriter, loc, - TypeRange {logicalType}, + TypeRange {storageType}, inputInfo.laneCount, {}, ValueRange {inputInfo.storage}, [&](detail::SpatComputeBatchBodyArgs args) { - SmallVector offsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), - args.lane, rewriter.getIndexAttr(0)}; - SmallVector sizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(logicalType.getDimSize(1)), - rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)}; - Value fragment = tensor::ExtractSliceOp::create( - rewriter, loc, fragmentType, args.inputs.front(), offsets, sizes, getUnitStrides(rewriter, 4)); + Value fragment = + extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc); switch (step.kind) { case DistributedTensorOpKind::Relu: fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult(); @@ -3034,8 +3032,8 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step, case DistributedTensorOpKind::Conv: return failure(); } - createParallelInsertSliceIntoBatchOutput( - rewriter, loc, fragment, args.outputs.front(), offsets, sizes, getUnitStrides(rewriter, 4)); + insertRowStripFragment( + fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc); return success(); }); if (failed(batchOp)) @@ -3743,6 +3741,8 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) { return failure(); if (state->outType.getRank() != 4 || !state->outType.hasStaticShape()) return failure(); + if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType)) + return failure(); FailureOr requestedStrategy = resolveRequestedConvLoweringStrategy(planOp.getOperation()); if (failed(requestedStrategy)) @@ -3786,7 +3786,7 @@ LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) { return failure(); StringRef failureReason; - return canDirectLowerRowStripConv(*state, failureReason) ? success() : failure(); + return canConsumeNchwRowStripFragments(*state, failureReason) ? success() : failure(); } FailureOr @@ -3822,17 +3822,33 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, if (rowStripInput) { if (failed(canConsumeAndProduceRowStrip(planOp))) return planOp.emitOpError("selected row-strip input/output layout is not supported for this Conv plan"), failure(); - return createConvRowsFromRowStripInput(*state, decision, *rowStripInput, rewriter, planOp.getLoc()); + return createConvOutputFromRowStripInput(*state, decision, *rowStripInput, rewriter, planOp.getLoc()); } if (failed(canLowerConvPlanToRowStrip(planOp))) return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure(); - FailureOr rows = createConvRowsForStrategy(*state, decision, rewriter, planOp.getLoc()); + ConvLoweringState rowState = *state; + const bool applyBiasAfterStorage = rowState.hasBias; + Value originalBias = rowState.b; + if (applyBiasAfterStorage) { + if (!isSupportedBiasAddValue(originalBias, rowState.outType)) + return planOp.emitOpError("selected row-strip Conv bias must be host-constant scalar/per-channel NCHW"), + failure(); + rowState.b = Value(); + rowState.hasBias = false; + } + + FailureOr rows = createConvRowsForStrategy(rowState, decision, rewriter, planOp.getLoc()); if (failed(rows)) return failure(); - FailureOr packedRows = createRowStripPackedRows(*rows, *state, rewriter, planOp.getLoc()); - if (failed(packedRows)) - return planOp.emitOpError("failed to pack Conv rows into the selected row-strip physical layout"), failure(); - return *packedRows; + FailureOr rowStripStorage = createRowStripStorageFromRows(*rows, state->outType, rewriter, planOp.getLoc()); + if (failed(rowStripStorage)) + return planOp.emitOpError("failed to build row-strip fragment storage for the selected Conv plan"), failure(); + if (applyBiasAfterStorage) { + rowStripStorage = applyRowStripBiasAdd(*rowStripStorage, rowState.outType, originalBias, rewriter, planOp.getLoc()); + if (failed(rowStripStorage)) + return planOp.emitOpError("failed to apply row-strip Conv bias per fragment"), failure(); + } + return *rowStripStorage; } if (decision.strategy == PimConvLoweringDepthwise) diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp index 506456f..acb3c54 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp @@ -5,7 +5,7 @@ #include "llvm/ADT/SmallVector.h" -#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" @@ -47,38 +47,28 @@ static FailureOr materializeBroadcastedConstantTensor(Value value, return failure(); const int64_t rankOffset = static_cast(resultShape.size() - sourceShape.size()); - for (int64_t i = 0; i < static_cast(resultShape.size()); ++i) { - const int64_t sourceIndex = i - rankOffset; - const int64_t sourceDim = sourceIndex < 0 ? 1 : sourceShape[sourceIndex]; - const int64_t resultDim = resultShape[i]; - if (sourceDim != 1 && sourceDim != resultDim) - return failure(); - } - - SmallVector sourceValues(denseAttr.getValues()); SmallVector sourceStrides = computeRowMajorStrides(sourceShape); SmallVector resultStrides = computeRowMajorStrides(resultShape); - + SmallVector sourceValues(denseAttr.getValues()); SmallVector resultValues; resultValues.reserve(resultType.getNumElements()); for (int64_t flatIndex = 0; flatIndex < resultType.getNumElements(); ++flatIndex) { int64_t remaining = flatIndex; int64_t sourceFlatIndex = 0; - for (int64_t i = 0; i < static_cast(resultShape.size()); ++i) { const int64_t resultIndex = resultStrides.empty() ? 0 : remaining / resultStrides[i]; remaining = resultStrides.empty() ? 0 : remaining % resultStrides[i]; - const int64_t sourceIndex = i - rankOffset; if (sourceIndex < 0) continue; - const int64_t sourceDim = sourceShape[sourceIndex]; + const int64_t resultDim = resultShape[i]; + if (sourceDim != 1 && sourceDim != resultDim) + return failure(); const int64_t mappedIndex = sourceDim == 1 ? 0 : resultIndex; sourceFlatIndex += mappedIndex * sourceStrides[sourceIndex]; } - resultValues.push_back(sourceValues[sourceFlatIndex]); } @@ -106,7 +96,7 @@ static FailureOr materializeReciprocalTensor(Value value, if (failed(broadcastedValue)) return failure(); - auto denseAttr = dyn_cast(getDenseConstantAttr(*broadcastedValue)); + auto denseAttr = dyn_cast(getHostConstDenseElementsAttr(*broadcastedValue)); if (!denseAttr) return failure(); @@ -185,10 +175,45 @@ struct DivToSpatialCompute : OpConversionPattern { } }; +struct AddToSpatialCompute : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(ONNXAddOp op, ONNXAddOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override { + auto resultType = dyn_cast(op.getResult().getType()); + if (!resultType || !resultType.hasStaticShape()) + return failure(); + + FailureOr candidate = + classifyBiasAddPlanCandidate(adaptor.getA(), adaptor.getB(), resultType); + if (succeeded(candidate)) { + auto plan = spatial::SpatBiasAddPlanOp::create( + rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, rewriter.getStringAttr("nchw")); + rewriter.replaceOp(op, plan.getResult()); + return success(); + } + + auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc()); + if (failed(lhs)) + return failure(); + auto rhs = prepareElementwiseOperand(adaptor.getB(), resultType, rewriter, op.getLoc()); + if (failed(rhs)) + return failure(); + + auto computeOp = + createSpatCompute<2>(rewriter, op.getLoc(), resultType, {}, ValueRange {*lhs, *rhs}, [&](Value x, Value y) { + auto loweredOp = spatial::SpatVAddOp::create(rewriter, op.getLoc(), resultType, x, y); + spatial::SpatYieldOp::create(rewriter, op.getLoc(), loweredOp.getResult()); + }); + rewriter.replaceOp(op, computeOp); + return success(); + } +}; + } // namespace void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { - patterns.add>(ctx); + patterns.add(ctx); patterns.add>(ctx); patterns.add>(ctx); patterns.add(ctx); diff --git a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp index 815ef49..6e20e14 100644 --- a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp @@ -6,10 +6,11 @@ #include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Accelerators/PIM/Pass/PIMPasses.h" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" using namespace mlir; @@ -19,7 +20,6 @@ namespace { static constexpr StringLiteral kLogicalLayout = "nchw"; static constexpr StringLiteral kDenseLayout = "dense_nchw"; static constexpr StringLiteral kRowStripLayout = "nchw_row_strip"; -static constexpr StringLiteral kRowStripIndexMap = "packed_hwc_rows_to_nchw"; enum class SelectedLayout { DenseNchw, @@ -34,6 +34,8 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap& l static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap& layouts) { if (auto reluPlan = dyn_cast(user)) return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip; + if (auto biasAddPlan = dyn_cast(user)) + return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip; if (auto convPlan = dyn_cast(user)) return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip; return false; @@ -49,21 +51,26 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap, SmallVector> buildRowStripMetadata(RankedTensorType type) { - SmallVector offsets; - SmallVector sizes; - const int64_t channels = type.getDimSize(1); - const int64_t height = type.getDimSize(2); - const int64_t width = type.getDimSize(3); - offsets.reserve(height * 4); - sizes.reserve(height * 4); - for (int64_t row = 0; row < height; ++row) { - offsets.append({0, 0, row, 0}); - sizes.append({1, channels, 1, width}); +static bool canConsumeRowStripAsUser(Operation* user) { + if (isa(user)) + return true; + if (auto biasAddPlan = dyn_cast(user)) { + auto resultType = dyn_cast(biasAddPlan.getOutput().getType()); + return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType); } - return {offsets, sizes}; + if (auto convPlan = dyn_cast(user)) + return succeeded(canConsumeAndProduceRowStrip(convPlan)); + return false; } +static bool hasRowStripConsumer(Value value) { + for (Operation* user : value.getUsers()) + if (canConsumeRowStripAsUser(user)) + return true; + return false; +} + + static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan, llvm::DenseMap& layouts) { SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput()); @@ -76,6 +83,9 @@ static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan, llvm::DenseMap& layouts) { if (!canSelectConvRowStrip(convPlan, layouts)) return SelectedLayout::DenseNchw; + if (getSelectedLayout(layouts, convPlan.getInput()) != SelectedLayout::NchwRowStrip + && !hasRowStripConsumer(convPlan.getResult())) + return SelectedLayout::DenseNchw; if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts)) return SelectedLayout::DenseNchw; return SelectedLayout::NchwRowStrip; @@ -85,11 +95,27 @@ static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan, llvm::DenseMap& layouts) { if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip) return SelectedLayout::DenseNchw; + if (!hasRowStripConsumer(reluPlan.getResult())) + return SelectedLayout::DenseNchw; if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts)) return SelectedLayout::DenseNchw; return SelectedLayout::NchwRowStrip; } +static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan, + llvm::DenseMap& layouts) { + if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip) + return SelectedLayout::DenseNchw; + auto resultType = dyn_cast(biasAddPlan.getOutput().getType()); + if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType)) + return SelectedLayout::DenseNchw; + if (!hasRowStripConsumer(biasAddPlan.getResult())) + return SelectedLayout::DenseNchw; + if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts)) + return SelectedLayout::DenseNchw; + return SelectedLayout::NchwRowStrip; +} + static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) { auto outputType = cast(value.getType()); auto [offsets, sizes] = buildRowStripMetadata(outputType); @@ -173,6 +199,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper(&op)) { + SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts); + if (layouts[biasAddPlan.getResult()] != selected) { + layouts[biasAddPlan.getResult()] = selected; + changed = true; + } + continue; + } } } @@ -180,6 +214,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper(&op)) producedValue = convPlan.getResult(); + else if (auto biasAddPlan = dyn_cast(&op)) + producedValue = biasAddPlan.getResult(); else if (auto reluPlan = dyn_cast(&op)) producedValue = reluPlan.getResult(); else diff --git a/src/PIM/Conversion/SpatialToGraphviz/CMakeLists.txt b/src/PIM/Conversion/SpatialToGraphviz/CMakeLists.txt deleted file mode 100644 index dea0370..0000000 --- a/src/PIM/Conversion/SpatialToGraphviz/CMakeLists.txt +++ /dev/null @@ -1,17 +0,0 @@ -add_onnx_mlir_rewriter(SpatialToGraphviz) - -add_pim_library(OMSpatialToGraphviz - SpatialToGraphviz.cpp - - EXCLUDE_FROM_OM_LIBS - - LINK_LIBS PUBLIC - MLIRTosaDialect - OMCompilerOptions - OMPimCommon - OMONNXOps - SpatialOps - - ACCEL_INCLUDE_DIRS PRIVATE - ${PIM_GENERATED_INCLUDE_DIRS} -) diff --git a/src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp b/src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp deleted file mode 100644 index 61f90c1..0000000 --- a/src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp +++ /dev/null @@ -1,259 +0,0 @@ -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Dialect/Tosa/IR/TosaOps.h" -#include "mlir/IR/Block.h" -#include "mlir/IR/Diagnostics.h" -#include "mlir/IR/Value.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -#include "llvm/Support/Casting.h" -#include "llvm/Support/Format.h" - -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" -#include "src/Dialect/ONNX/ONNXOps.hpp" - -#define FORMAT_OPERATION(op) 'x' << llvm::format_hex_no_prefix(reinterpret_cast(op), 0) -#define FORMAT_ARGUMENT(computeOpPointer, argumentNum) llvm::format("Arg_%p_%u", computeOpPointer, argumentNum) - -using namespace mlir; - -namespace onnx_mlir { - -namespace { - -struct SpatialToGraphvizPass : public PassWrapper> { - - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialToGraphvizPass) - - StringRef getArgument() const override { return "convert-spatial-to-graphviz"; } - - StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; } - - SpatialToGraphvizPass(raw_ostream& os = llvm::errs()) - : os(os) {} - SpatialToGraphvizPass(const SpatialToGraphvizPass& pass) - : SpatialToGraphvizPass(pass.os) {} - void runOnOperation() final; - -private: - raw_ostream& os; - - /** - * Draws the subgraph for a given spatial::SpatCompute, including: - * 1. Input nodes (block arguments) - * 2. Operations - * 3. Edges between yield (output) and its users - * - * @param op The spatial::SpatCompute to draw the subgraph for. - * @param computeNum The number of the compute operation. - */ - void drawComputeOpSubgraph(spatial::SpatCompute op, size_t computeNum) { - os << "\tsubgraph cluster" << computeNum << " {\n\t\tlabel=\"Compute" << computeNum << "\";\n" - << "\t\tstyle=filled;\n" - << "\t\tcolor=lightblue;\n"; - - Block& block = op.getBody().front(); - - // Inputs - size_t inputNum = 0; - for (BlockArgument& input : block.getArguments()) { - - auto fromOp = FORMAT_ARGUMENT(op.getOperation(), inputNum); - - os << "\t\t" << fromOp << " [label=\"Arg" << inputNum << "\",shape=box];\n"; - for (auto userOp : input.getUsers()) - os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n"; - inputNum++; - } - - // Iterate operations - for (auto& childOp : block.getOperations()) { - os << "\t\t" << FORMAT_OPERATION(&childOp) << " [label=\"" << childOp.getName() << "\"];\n"; - - drawEdgesFromOpToItsUsers(&childOp); - } - - os << "\t}\n"; - - // Draw edges from the yield to the users of this computeOp - Operation* yieldOp = block.getTerminator(); - if (!isa(yieldOp)) { - yieldOp->emitError("Terminator of block must be YieldOp ???"); - signalPassFailure(); - return; - } - - for (auto computeOpResult : op->getResults()) { - for (auto& computeOpUse : computeOpResult.getUses()) { - auto toOp = FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber()); - os << "\t" << FORMAT_OPERATION(yieldOp) << " -> " << toOp << ";\n"; - } - } - } - - /** - * @brief Draws the subgraph for a concatOp. - * - * This function draws a subgraph for a concatOp. The subgraph consists of a - * node for each input of the concatOp, as well as an output node. Edges are - * created from the output node to each user of the concatOp. - * - * @param concatOp The concatOp for which the subgraph is drawn. - * @param concatOpNum The number of the concatOp. - */ - void drawConcatOpSubgraph(Operation* concatOp, size_t concatOpNum) { - os << "\tsubgraph clusterconcat" << concatOpNum << " {\n\t\tlabel=\"ConcatOp" << concatOpNum << "\";\n" - << "\t\tstyle=filled;\n" - << "\t\tcolor=orange;\n"; - - // Inputs - size_t inputNum = 0; - for (Value input : concatOp->getOperands()) { - auto fromOp = FORMAT_ARGUMENT(concatOp, inputNum); - - os << "\t\t" << fromOp << " [label=\"Input" << inputNum << "\"];\n"; - for (auto userOp : input.getUsers()) - os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n"; - inputNum++; - } - - // Output - os << "\t\t" << FORMAT_OPERATION(concatOp) << " [label=Out];\n"; - - os << "\t}\n"; - - // Edges from output to users - - for (auto& computeOpUse : concatOp->getResult(0).getUses()) { - os << "\t" << FORMAT_OPERATION(concatOp) << " -> " - << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber()) << ";\n"; - } - } - - /** - * Draws the ExtractSliceOp in the graph visualization. - * - * This function takes a tensor::ExtractSliceOp and adds the corresponding - * node and edges to the graph visualization. It creates a node with the - * label as the static offsets attribute of the sliceOp, and connects it to - * the compute operations that use the result of the sliceOp. - * - * @param sliceOp The tensor::ExtractSliceOp to be drawn in the graph - * visualization. - */ - void drawExtractSliceOp(tensor::ExtractSliceOp sliceOp) { - auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0); - os << "\t" << nodeId << " [label=\"Slice: "; - sliceOp.getStaticOffsetsAttr().print(os); - os << "\",color=lawngreen];\n"; - - for (auto& computeOpUse : sliceOp.getResult().getUses()) { - os << "\t" << nodeId << " -> " << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber()) - << ";\n"; - } - } - - void drawBiasTileOp(tensor::ExtractSliceOp sliceOp) { - auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0); - os << "\t" << nodeId << " [label=\"Bias: "; - sliceOp.getStaticOffsetsAttr().print(os); - os << "\",color=lightpink];\n"; - - for (auto user : sliceOp.getResult().getUsers()) - os << "\t" << nodeId << " -> " << FORMAT_OPERATION(user) << ";\n"; - } - - /** - * Draws edges from the given operation to its users. - * - * @param fromOp The operation from which the edges are drawn. - */ - void drawEdgesFromOpToItsUsers(mlir::Operation* fromOp) { - for (auto result : fromOp->getResults()) - for (auto userOp : result.getUsers()) - os << "\t\t" << FORMAT_OPERATION(fromOp) << " -> " << FORMAT_OPERATION(userOp) << ";\n"; - } - - /** - * Draws input node and edges for the given `funcOp`. - * - * @param funcOp The `funcOp` for which to draw input nodes and edges. - */ - void drawInputNodesAndEdges(func::FuncOp& funcOp) { - os << "\tinput [label=\"Module Input\",color=green];\n"; - - size_t funcOpArgNum = 0; - for (BlockArgument& arg : funcOp.getArguments()) { - - for (auto& useOp : arg.getUses()) { - os << "\tinput -> " << FORMAT_ARGUMENT(useOp.getOwner(), useOp.getOperandNumber()) << "[label=" << funcOpArgNum - << "];\n"; - } - funcOpArgNum++; - } - } -}; - -void SpatialToGraphvizPass::runOnOperation() { - ModuleOp module = getOperation(); - - auto entryFunc = getPimEntryFunc(module); - if (failed(entryFunc)) { - module.emitError("failed to locate the PIM entry function for Spatial graph visualization"); - signalPassFailure(); - return; - } - func::FuncOp func = *entryFunc; - - os << "digraph G {\n" - << "\tnode [style=filled,color=white];\n"; - - size_t computeNum = 0; - size_t concatNum = 0; - - // Iterate over the ComputeOps within FuncOp: - // 1. Print their subgraph - // 2. Print the edges from its inputs to its outputs - for (Operation& op : func.getOps()) { - if (auto computeOp = dyn_cast(op)) { - drawComputeOpSubgraph(computeOp, computeNum++); - } - else if (auto concatOp = dyn_cast(op)) { - drawConcatOpSubgraph(concatOp, concatNum++); - } - else if (auto extractSliceOp = dyn_cast(op)) { - auto producerOp = extractSliceOp->getOperand(0).getDefiningOp(); - if (producerOp) { - // Skip extractSliceOp if producer is constant weights (ONNXConstantOp) - if (llvm::isa(producerOp)) - continue; - // If produced by tosa::ReshapeOp (i.e. it is a bias tile) connect - // directly to its user, which is not a ComputeOp argument. - if (llvm::isa(producerOp)) { - drawBiasTileOp(extractSliceOp); - continue; - } - } - - drawExtractSliceOp(extractSliceOp); - } - } - - // Draw input node, and edges to it users - drawInputNodesAndEdges(func); - - // Draw output node (use the return Operation - argument number=0 - as nodeId) - auto returnOp = func.getBody().front().getTerminator(); - os << '\t' << FORMAT_ARGUMENT(returnOp, 0) << " [label=\"Module Output\",color=green];\n"; - - os << "}\n"; -} - -} // namespace - -std::unique_ptr createSpatialToGraphvizPass() { return std::make_unique(); } - -} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp b/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp index b64a8f7..619f65c 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp +++ b/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp @@ -302,76 +302,87 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const { bool hasFailure = false; - moduleOp.walk([&](Operation* op) { - auto verifyOperand = [&](Value operand, unsigned operandIndex) { - if (!isa(operand.getType())) - return; - if (succeeded(resolveContiguousAddress(operand)) || succeeded(compileContiguousAddressExpr(operand))) - return; - op->emitOpError() << "operand #" << operandIndex - << " is not backed by contiguous addressable storage after PIM bufferization"; - hasFailure = true; - }; - if (auto memCopyOp = dyn_cast(op)) { - if (!pim::isNormalizedCopyOp(memCopyOp)) { - memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization"); - hasFailure = true; - } - verifyOperand(memCopyOp.getTarget(), 0); - verifyOperand(memCopyOp.getSource(), 1); - return; - } - if (auto loadOp = dyn_cast(op)) { - if (!pim::isNormalizedCopyOp(loadOp)) { - loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization"); - hasFailure = true; - } - verifyOperand(loadOp.getDeviceTarget(), 2); - verifyOperand(loadOp.getHostSource(), 3); - return; - } - if (auto storeOp = dyn_cast(op)) { - if (!pim::isNormalizedCopyOp(storeOp)) { - storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization"); - hasFailure = true; - } - verifyOperand(storeOp.getHostTarget(), 2); - verifyOperand(storeOp.getDeviceSource(), 3); - return; - } - if (auto sendOp = dyn_cast(op)) { - verifyOperand(sendOp.getInput(), 0); - return; - } - if (auto receiveOp = dyn_cast(op)) { - verifyOperand(receiveOp.getOutputBuffer(), 0); - return; - } - if (auto concatOp = dyn_cast(op)) { - verifyOperand(concatOp.getOutputBuffer(), 0); - for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs())) - verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1); - return; - } - if (isa(op)) { - for (auto operandAndIndex : llvm::enumerate(op->getOperands())) { - if (auto vmmOp = dyn_cast(op); vmmOp && operandAndIndex.index() == 0) - continue; - verifyOperand(operandAndIndex.value(), operandAndIndex.index()); - } - } + auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) { + (void) walkPimCoreBlockStructurally( + coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) { + auto verifyOperand = [&](Value operand, unsigned operandIndex) { + if (!isa(operand.getType())) + return; + if (succeeded(resolveContiguousAddress(operand, knowledge)) || succeeded(compileContiguousAddressExpr(operand))) + return; + op.emitOpError() << "operand #" << operandIndex + << " is not backed by contiguous addressable storage after PIM bufferization"; + hasFailure = true; + }; + + if (auto memCopyOp = dyn_cast(&op)) { + if (!pim::isNormalizedCopyOp(memCopyOp)) { + memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization"); + hasFailure = true; + } + verifyOperand(memCopyOp.getTarget(), 0); + verifyOperand(memCopyOp.getSource(), 1); + return success(); + } + if (auto loadOp = dyn_cast(&op)) { + if (!pim::isNormalizedCopyOp(loadOp)) { + loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization"); + hasFailure = true; + } + verifyOperand(loadOp.getDeviceTarget(), 2); + verifyOperand(loadOp.getHostSource(), 3); + return success(); + } + if (auto storeOp = dyn_cast(&op)) { + if (!pim::isNormalizedCopyOp(storeOp)) { + storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization"); + hasFailure = true; + } + verifyOperand(storeOp.getHostTarget(), 2); + verifyOperand(storeOp.getDeviceSource(), 3); + return success(); + } + if (auto sendOp = dyn_cast(&op)) { + verifyOperand(sendOp.getInput(), 0); + return success(); + } + if (auto receiveOp = dyn_cast(&op)) { + verifyOperand(receiveOp.getOutputBuffer(), 0); + return success(); + } + if (auto concatOp = dyn_cast(&op)) { + verifyOperand(concatOp.getOutputBuffer(), 0); + for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs())) + verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1); + return success(); + } + if (isa(&op)) { + for (auto operandAndIndex : llvm::enumerate(op.getOperands())) { + if (auto vmmOp = dyn_cast(&op); vmmOp && operandAndIndex.index() == 0) + continue; + verifyOperand(operandAndIndex.value(), operandAndIndex.index()); + } + } + return success(); + }); + }; + + moduleOp.walk([&](pim::PimCoreOp coreOp) { verifyWithKnowledge(coreOp, seedCoreKnowledge(coreOp)); }); + moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) { + StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, 0); + verifyWithKnowledge(coreBatchOp, knowledge); }); if (hasFailure) { diff --git a/src/PIM/Dialect/Spatial/Spatial.td b/src/PIM/Dialect/Spatial/Spatial.td index 603e626..3d1bb03 100644 --- a/src/PIM/Dialect/Spatial/Spatial.td +++ b/src/PIM/Dialect/Spatial/Spatial.td @@ -232,6 +232,22 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> { let hasVerifier = 1; } +def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> { + let summary = "Layout-aware Conv-style bias add planning op"; + + let arguments = (ins + SpatTensor:$input, + SpatTensor:$bias, + StrAttr:$logicalLayout + ); + + let results = (outs + SpatTensor:$output + ); + + let hasVerifier = 1; +} + def SpatBlueprintOp : SpatOp<"blueprint", []> { let summary = "Blueprint for assembling logical tensors from published fragments"; diff --git a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp index 67561ff..ab9cc49 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp @@ -1,5 +1,6 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/Block.h" @@ -59,6 +60,21 @@ static LogicalResult verifyStaticWeights(ComputeOpTy computeOp, StringRef kind) return success(); } +static bool isStaticScfForInductionVar(Value value) { + auto blockArg = dyn_cast(value); + if (!blockArg) + return false; + + auto loop = dyn_cast_or_null(blockArg.getOwner()->getParentOp()); + if (!loop || loop.getInductionVar() != value) + return false; + + std::optional lowerBound = matchConstantIndexValue(loop.getLowerBound()); + std::optional upperBound = matchConstantIndexValue(loop.getUpperBound()); + std::optional step = matchConstantIndexValue(loop.getStep()); + return lowerBound && upperBound && step && *step > 0 && *upperBound >= *lowerBound; +} + static bool isStaticIndexExpr(Value value) { if (matchConstantIndexValue(value)) return true; @@ -80,7 +96,7 @@ static bool isStaticIndexExpr(Value value) { } static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) { - if (value == laneArg || isStaticIndexExpr(value)) + if (value == laneArg || isStaticIndexExpr(value) || isStaticScfForInductionVar(value)) return true; auto affineApply = value.getDefiningOp(); @@ -436,6 +452,39 @@ LogicalResult SpatReluPlanOp::verify() { return success(); } +LogicalResult SpatBiasAddPlanOp::verify() { + if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.bias_add_plan"))) + return failure(); + if (!isKnownLogicalLayout(getLogicalLayout())) + return emitError("requires a known logical layout"); + + auto inputType = dyn_cast(getInput().getType()); + auto biasType = dyn_cast(getBias().getType()); + auto outputType = dyn_cast(getOutput().getType()); + if (!inputType || !biasType || !outputType) + return emitError("requires ranked tensor input, bias, and output"); + if (!inputType.hasStaticShape() || !biasType.hasStaticShape() || !outputType.hasStaticShape()) + return emitError("requires static tensor input, bias, and output"); + if (inputType != outputType) + return emitError("requires matching input and output tensor types"); + if (outputType.getRank() != 4) + return emitError("requires rank-4 input/output tensors"); + if (getLogicalLayout() != "nchw") + return emitError("requires logical layout \"nchw\""); + if (biasType.getElementType() != outputType.getElementType()) + return emitError("requires bias element type to match the output element type"); + + ArrayRef biasShape = biasType.getShape(); + const int64_t channels = outputType.getDimSize(1); + const bool supported = biasShape.empty() || (biasShape.size() == 1 && biasShape[0] == channels) + || (biasShape.size() == 2 && biasShape[0] == 1 && biasShape[1] == channels) + || (biasShape.size() == 4 && biasShape[0] == 1 && biasShape[1] == channels + && biasShape[2] == 1 && biasShape[3] == 1); + if (!supported) + return emitError("requires scalar or per-channel bias broadcastable over NCHW"); + return success(); +} + LogicalResult SpatBlueprintOp::verify() { auto modeAttr = getModeAttr(); bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly"; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp index 0c03134..f710b66 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp @@ -15,11 +15,13 @@ #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringMap.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/raw_ostream.h" #include #include +#include #include #include #include @@ -44,6 +46,13 @@ using namespace mlir; namespace onnx_mlir { namespace spatial { namespace { + +struct ProjectedProducerFragmentDemand { + ProducerKey producer; + SmallVector fragmentOffsets; + unsigned ordinal = 0; +}; + FailureOr recordProjectedScalarHostFragmentsFromPackedValue(MaterializerState& state, MaterializedClass& sourceClass, ArrayRef keys, @@ -72,6 +81,8 @@ FailureOr materializeTensorValueForMaterializedClassUse(MaterializerState IRMapping* mapper = nullptr); FailureOr materializeWholeBatchInput( MaterializerState& state, MaterializedClass& targetClass, ProducerKey key, Type resultType, Location loc); +ScalarPeerEdgeKey makeScalarPeerEdgeKey(int64_t sourceCore, int64_t targetCore, Type payloadType); +bool hasPlannedScalarPeerReceive(const MaterializerState& state, const ScalarPeerEdgeKey& key); FailureOr localizeMaterializedClassOperand(MaterializerState& state, MaterializedClass& targetClass, Value value, @@ -91,6 +102,22 @@ bool isProjectedInputSliceCompatibleWithProducerFragments(SpatComputeBatch consu const AffineProjectedInputSliceMatch& match, ProducerKey producer, uint32_t consumerLane); +SmallVector +collectProjectedFragmentDemandsForBatchInput(MaterializerState& state, + SpatComputeBatch consumerBatch, + unsigned inputIndex, + Value input, + ComputeInstance logicalConsumer); +SmallVector collectProjectedProducerKeysForBatchInput(MaterializerState& state, + SpatComputeBatch consumerBatch, + unsigned inputIndex, + Value input, + ComputeInstance logicalConsumer); +LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, + MaterializedClass& sourceClass, + ArrayRef keys, + Value payload, + Location loc); std::optional getProjectedInputSliceMatch(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex); std::optional @@ -434,22 +461,7 @@ bool canUseProjectedLaneInput(MaterializerState& state, unsigned inputIndex, Value input, ComputeInstance logicalConsumer) { - auto producerResult = dyn_cast(input); - if (!producerResult) - return false; - - auto producerBatch = dyn_cast_or_null(producerResult.getOwner()); - if (!producerBatch || producerBatch.getNumResults() == 0) - return false; - - std::optional match = getProjectedInputSliceMatch(state, consumerBatch, inputIndex); - if (!match) - return false; - - ProducerKey laneProducer = - getBatchLaneProducerKey(producerBatch, logicalConsumer.laneStart, 1, producerResult.getResultNumber()); - return isProjectedInputSliceCompatibleWithProducerFragments( - consumerBatch, *match, laneProducer, logicalConsumer.laneStart); + return !collectProjectedFragmentDemandsForBatchInput(state, consumerBatch, inputIndex, input, logicalConsumer).empty(); } FailureOr classifyComputeBatchInputDemand(MaterializerState& state, @@ -497,6 +509,11 @@ SmallVector collectProducerKeysForBatchInputDestinations(Materia ComputeInstance logicalConsumer) { if (std::optional wholeBatchProducer = getWholeBatchProducerKeyForDirectBatchResult(input)) { if (!canUseProjectedLaneInput(state, consumerBatch, inputIndex, input, logicalConsumer)) { + SmallVector projectedKeys = + collectProjectedProducerKeysForBatchInput(state, consumerBatch, inputIndex, input, logicalConsumer); + if (!projectedKeys.empty()) + return projectedKeys; + auto producerBatch = cast(wholeBatchProducer->instance.op); SmallVector keys; for (uint32_t lane = 0; lane < static_cast(producerBatch.getLaneCount()); ++lane) @@ -580,6 +597,165 @@ LogicalResult buildMaterializationWorkStreams(MaterializerState& state) { return success(); } + +using MaterializedBatchingSignature = SmallVector; + +static int64_t getProjectedInputExchangeDirection(int64_t sourceCoreId, int64_t targetCoreId) { + if (sourceCoreId == targetCoreId) + return 0; + return sourceCoreId < targetCoreId ? 1 : 2; +} + +static bool isLowerToHigherProjectedInputFragment(const ProjectedInputTransferFragment& fragment) { + return fragment.sourceCoreId < fragment.targetCoreId; +} + +static bool isHigherToLowerProjectedInputFragment(const ProjectedInputTransferFragment& fragment) { + return fragment.sourceCoreId > fragment.targetCoreId; +} + +static LogicalResult appendProjectedBatchingSignatureForCpu(MaterializerState& state, + CpuId cpu, + MaterializedBatchingSignature& signature) { + auto streamIt = state.logicalInstancesByCpu.find(cpu); + if (streamIt == state.logicalInstancesByCpu.end()) + return success(); + + for (const ComputeInstance& logicalConsumer : streamIt->second) { + auto batchConsumer = dyn_cast(logicalConsumer.op); + if (!batchConsumer) + continue; + + SmallVector consumerInputs = getComputeInstanceInputs(logicalConsumer); + for (auto [inputIndex, input] : llvm::enumerate(consumerInputs)) { + SmallVector demands = collectProjectedFragmentDemandsForBatchInput( + state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); + if (demands.empty()) + continue; + + signature.push_back(-1); + signature.push_back(static_cast(inputIndex)); + for (const ProjectedProducerFragmentDemand& demand : demands) { + ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, demand.producer.instance); + auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); + if (producerCpuIt == state.schedule.computeToCpuMap.end()) + return logicalConsumer.op->emitError( + "projected input batching signature found a fragment produced by an unscheduled compute"); + signature.push_back(getProjectedInputExchangeDirection(producerCpuIt->second, cpu)); + } + } + } + + return success(); +} + +static DenseMap buildCpuToMaterializedGroup(ArrayRef> groups) { + DenseMap cpuToGroup; + for (auto [groupIndex, group] : llvm::enumerate(groups)) + for (CpuId cpu : group) + cpuToGroup[cpu] = groupIndex; + return cpuToGroup; +} + +static LogicalResult appendOrdinaryInputBatchingSignatureForCpu(MaterializerState& state, + CpuId cpu, + const DenseMap& cpuToGroup, + MaterializedBatchingSignature& signature) { + auto streamIt = state.logicalInstancesByCpu.find(cpu); + if (streamIt == state.logicalInstancesByCpu.end()) + return success(); + + for (const ComputeInstance& logicalConsumer : streamIt->second) { + SmallVector consumerInputs = getComputeInstanceInputs(logicalConsumer); + for (auto [inputIndex, input] : llvm::enumerate(consumerInputs)) { + SmallVector producerKeys; + if (auto batchConsumer = dyn_cast(logicalConsumer.op)) { + SmallVector projectedDemands = collectProjectedFragmentDemandsForBatchInput( + state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); + if (!projectedDemands.empty()) + continue; + producerKeys = collectProducerKeysForBatchInputDestinations( + state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); + } + else { + producerKeys = collectProducerKeysForDestinations(input, logicalConsumer); + } + + if (producerKeys.empty()) + continue; + + signature.push_back(-2); + signature.push_back(static_cast(inputIndex)); + for (ProducerKey producerKey : producerKeys) { + ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producerKey.instance); + auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); + if (producerCpuIt == state.schedule.computeToCpuMap.end()) + return logicalConsumer.op->emitError( + "ordinary input batching signature found a fragment produced by an unscheduled compute"); + auto producerGroupIt = cpuToGroup.find(producerCpuIt->second); + if (producerGroupIt == cpuToGroup.end()) + return logicalConsumer.op->emitError( + "ordinary input batching signature found a producer outside materialized CPU groups"); + signature.push_back(static_cast(producerGroupIt->second)); + } + } + } + + return success(); +} + +static LogicalResult buildMaterializedBatchingSignature(MaterializerState& state, + CpuId cpu, + const DenseMap& cpuToGroup, + MaterializedBatchingSignature& signature) { + if (failed(appendProjectedBatchingSignatureForCpu(state, cpu, signature))) + return failure(); + return appendOrdinaryInputBatchingSignatureForCpu(state, cpu, cpuToGroup, signature); +} + +static std::string stringifyMaterializedBatchingSignature(ArrayRef signature) { + std::string key; + llvm::raw_string_ostream os(key); + for (int64_t value : signature) + os << value << ','; + os.flush(); + return key; +} + +static LogicalResult refineCpuGroupsByMaterializedBatchingSignature(MaterializerState& state, + SmallVectorImpl>& groups) { + bool changed = true; + while (changed) { + changed = false; + DenseMap cpuToGroup = buildCpuToMaterializedGroup(groups); + SmallVector, 8> refinedGroups; + refinedGroups.reserve(groups.size()); + + for (ArrayRef group : groups) { + llvm::StringMap groupBySignature; + size_t groupCountBefore = refinedGroups.size(); + for (CpuId cpu : group) { + MaterializedBatchingSignature signature; + if (failed(buildMaterializedBatchingSignature(state, cpu, cpuToGroup, signature))) + return failure(); + + std::string key = stringifyMaterializedBatchingSignature(signature); + auto [it, inserted] = groupBySignature.try_emplace(key, refinedGroups.size()); + if (inserted) + refinedGroups.emplace_back(); + refinedGroups[it->second].push_back(cpu); + } + changed |= refinedGroups.size() != groupCountBefore + 1; + } + + groups.clear(); + for (SmallVector& group : refinedGroups) + groups.push_back(std::move(group)); + } + + return success(); +} + LogicalResult buildMaterializationClassesFromScheduleEquivalence(MaterializerState& state) { DenseSet usedCpus; for (const auto& entry : state.schedule.cpuToLastComputeMap) @@ -609,11 +785,21 @@ LogicalResult buildMaterializationClassesFromScheduleEquivalence(MaterializerSta roots.push_back(entry.first); llvm::sort(roots); - state.classes.reserve(roots.size()); + SmallVector, 8> classCpuGroups; + classCpuGroups.reserve(roots.size()); for (CpuId root : roots) { + SmallVector rootCpus = groupsByRoot.lookup(root); + llvm::sort(rootCpus); + classCpuGroups.push_back(std::move(rootCpus)); + } + if (failed(refineCpuGroupsByMaterializedBatchingSignature(state, classCpuGroups))) + return failure(); + + state.classes.reserve(classCpuGroups.size()); + for (SmallVector& cpus : classCpuGroups) { MaterializedClass materializedClass; materializedClass.id = state.classes.size(); - materializedClass.cpus = groupsByRoot.lookup(root); + materializedClass.cpus = std::move(cpus); llvm::sort(materializedClass.cpus); materializedClass.isBatch = materializedClass.cpus.size() > 1; for (auto [lane, cpu] : llvm::enumerate(materializedClass.cpus)) { @@ -2010,6 +2196,32 @@ FailureOr materializeIndexedBatchRunPackedValue(MaterializerState& state, using IndexedFragmentBuilder = llvm::function_ref(Value flatIndex)>; using IndexedInsertOffsetBuilder = llvm::function_ref(Value flatIndex)>; +struct ReceiveMessagePartition { + SmallVector criticalStaticIndices; + SmallVector remainingIndices; +}; + +std::optional getConstantIndexValue(Value value); + +ReceiveMessagePartition partitionReceivesByPendingSendUnlocks(const MaterializerState& state, + const MaterializedClass& targetClass, + Type payloadType, + const MessageVector& messages, + bool preserveMessageOrder = false); +ScalarPeerReceiveKey makeScalarPeerReceiveKey(int64_t sourceCore, + int64_t targetCore, + Type payloadType, + std::optional channelId = std::nullopt); +FailureOr createScalarPeerReceiveAndFlush(MaterializerState& state, + MaterializedClass& targetClass, + Type payloadType, + Value channelId, + Value sourceCoreId, + Value targetCoreId, + std::optional staticKey, + Location loc); +MessageVector filterMessageVector(const MessageVector& messages, ArrayRef indices); + bool isDeferredLocalPackedScalarRun(const PackedScalarRunValue& run) { return run.kind == PackedScalarRunKind::DeferredLocalCompute; } @@ -2125,22 +2337,60 @@ FailureOr materializePackedScalarRunValue(MaterializerState& state, Value init = tensor::EmptyOp::create(state.rewriter, loc, fullPackedType->getShape(), fullPackedType->getElementType()) .getResult(); - auto packed = emitIndexedFragmentInsertLoop( - state, - targetClass, - init, - static_cast(run.messages.size()), - [&](Value index) -> FailureOr { - Value channelId = createIndexedChannelId(state, targetClass.op, run.messages, index, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, run.messages, index, loc); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, run.messages, index, loc); - return SpatChannelReceiveOp::create(state.rewriter, loc, run.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); - }, - [&](Value index) -> FailureOr { - return scaleIndexByDim0SizeInClass(state, targetClass, index, run.fragmentType.getDimSize(0), loc); - }, - loc); + Value packedInit = init; + ReceiveMessagePartition partition = + partitionReceivesByPendingSendUnlocks(state, targetClass, run.fragmentType, run.messages, /*preserveMessageOrder=*/true); + for (size_t index : partition.criticalStaticIndices) { + state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); + FailureOr received = createScalarPeerReceiveAndFlush( + state, + targetClass, + run.fragmentType, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, run.messages.channelIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, run.messages.sourceCoreIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, run.messages.targetCoreIds[index]), + makeScalarPeerReceiveKey( + run.messages.sourceCoreIds[index], run.messages.targetCoreIds[index], run.fragmentType, run.messages.channelIds[index]), + loc); + if (failed(received)) + return failure(); + FailureOr offset = scaleIndexByDim0SizeInClass( + state, + targetClass, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(index)), + run.fragmentType.getDimSize(0), + loc); + if (failed(offset)) + return failure(); + FailureOr updated = createDim0InsertSliceInClass(state, targetClass, loc, *received, packedInit, *offset); + if (failed(updated)) + return failure(); + packedInit = *updated; + } + + auto remainingMessages = filterMessageVector(run.messages, partition.remainingIndices); + SmallVector remainingPackedRowOffsets; + remainingPackedRowOffsets.reserve(partition.remainingIndices.size()); + for (size_t originalIndex : partition.remainingIndices) + remainingPackedRowOffsets.push_back(static_cast(originalIndex) * run.fragmentType.getDimSize(0)); + auto packed = remainingMessages.empty() + ? FailureOr(packedInit) + : emitIndexedFragmentInsertLoop( + state, + targetClass, + packedInit, + static_cast(remainingMessages.size()), + [&](Value index) -> FailureOr { + Value channelId = createIndexedChannelId(state, targetClass.op, remainingMessages, index, loc); + Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingMessages, index, loc); + Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingMessages, index, loc); + return SpatChannelReceiveOp::create(state.rewriter, loc, run.fragmentType, channelId, sourceCoreId, targetCoreId) + .getOutput(); + }, + [&](Value index) -> FailureOr { + return createIndexedIndexValue(state, targetClass.op, remainingPackedRowOffsets, index, loc); + }, + loc); if (failed(packed)) return failure(); run.packed = *packed; @@ -2602,11 +2852,22 @@ LogicalResult collectProducerDestinations(MaterializerState& state) { SmallVector consumerInputs = getComputeInstanceInputs(scheduledConsumer); for (auto [inputIndex, input] : llvm::enumerate(consumerInputs)) { SmallVector producerKeys; - if (auto batchConsumer = dyn_cast(logicalConsumer.op)) - producerKeys = collectProducerKeysForBatchInputDestinations( + bool projectedPlanOwnsFanout = false; + if (auto batchConsumer = dyn_cast(logicalConsumer.op)) { + SmallVector projectedKeys = collectProjectedProducerKeysForBatchInput( state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); - else + projectedPlanOwnsFanout = !state.classes[targetClass].isBatch && !projectedKeys.empty(); + producerKeys = projectedPlanOwnsFanout + ? std::move(projectedKeys) + : collectProducerKeysForBatchInputDestinations( + state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); + } + else { producerKeys = collectProducerKeysForDestinations(input, logicalConsumer); + } + + if (projectedPlanOwnsFanout) + continue; for (ProducerKey producerKey : producerKeys) { ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producerKey.instance); @@ -2712,6 +2973,45 @@ static SmallVector collectEnclosingStaticProjectedLo return loops; } +static bool isProjectedOffsetValue(Value value, + Value laneArg, + ArrayRef loops, + bool& usesDynamicBinding); + +static DenseIntElementsAttr getDenseIndexVector(Value value) { + auto constant = value.getDefiningOp(); + if (!constant) + return nullptr; + auto denseAttr = dyn_cast(constant.getValue()); + if (!denseAttr || denseAttr.getType().getRank() != 1) + return nullptr; + return denseAttr; +} + +static FailureOr getDenseIndexVectorElement(DenseIntElementsAttr denseAttr, int64_t index) { + if (index < 0 || index >= denseAttr.getNumElements()) + return failure(); + auto values = denseAttr.getValues(); + return values[index].getSExtValue(); +} + +static bool isProjectedIndexLookup(Value value, + Value laneArg, + ArrayRef loops, + bool& usesDynamicBinding) { + auto extract = value.getDefiningOp(); + if (!extract || extract.getIndices().size() != 1) + return false; + if (!getDenseIndexVector(extract.getTensor())) + return false; + + bool indexUsesDynamicBinding = false; + if (!isProjectedOffsetValue(extract.getIndices().front(), laneArg, loops, indexUsesDynamicBinding)) + return false; + usesDynamicBinding = usesDynamicBinding || indexUsesDynamicBinding; + return true; +} + static bool isProjectedOffsetValue(Value value, Value laneArg, ArrayRef loops, bool& usesDynamicBinding) { if (value == laneArg) { @@ -2729,6 +3029,9 @@ isProjectedOffsetValue(Value value, Value laneArg, ArrayRef(); if (!affineApply || affineApply.getAffineMap().getNumResults() != 1) return false; @@ -2769,6 +3072,19 @@ static FailureOr evaluateProjectedOffsetValue(OpFoldResult value, return loop.lowerBound + loopIterationIndices[index] * loop.step; } + if (auto extract = current.getDefiningOp()) { + if (extract.getIndices().size() != 1) + return failure(); + DenseIntElementsAttr denseAttr = getDenseIndexVector(extract.getTensor()); + if (!denseAttr) + return failure(); + FailureOr index = + evaluateProjectedOffsetValue(extract.getIndices().front(), laneArg, lane, loops, loopIterationIndices); + if (failed(index)) + return failure(); + return getDenseIndexVectorElement(denseAttr, *index); + } + if (auto affineApply = current.getDefiningOp()) { return evaluateAffineApply(affineApply, [&](Value operand) { return evaluateProjectedOffsetValue(operand, laneArg, lane, loops, loopIterationIndices); @@ -2803,24 +3119,14 @@ static std::optional getConstantIndex(OpFoldResult value) { return std::nullopt; } -static std::optional matchAffineProjectedInputSlice(SpatComputeBatch batch, - unsigned inputIndex) { +static std::optional matchAffineProjectedInputSlice( + BlockArgument inputArg, BlockArgument laneArg, tensor::ExtractSliceOp extract) { const auto fail = [&](StringRef) -> std::optional { return std::nullopt; }; - std::optional inputArg = batch.getInputArgument(inputIndex); - std::optional laneArg = batch.getLaneArgument(); - if (!inputArg || !laneArg) - return fail("missing-input-or-lane-arg"); + if (extract.getSource() != inputArg) + return fail("extract-source-is-not-input-arg"); - if (!inputArg->hasOneUse()) - return fail("input-arg-not-one-use"); - - Operation* user = *inputArg->getUsers().begin(); - auto extract = dyn_cast(user); - if (!extract || extract.getSource() != *inputArg) - return fail("input-user-is-not-direct-extract-slice"); - - auto inputType = dyn_cast(inputArg->getType()); + auto inputType = dyn_cast(inputArg.getType()); auto fragmentType = dyn_cast(extract.getResult().getType()); if (!inputType || !fragmentType || !inputType.hasStaticShape() || !fragmentType.hasStaticShape()) return fail("non-static-ranked-input-or-fragment"); @@ -2845,7 +3151,7 @@ static std::optional matchAffineProjectedInputSl for (auto [dim, offset] : llvm::enumerate(offsets)) { bool usesDynamicBinding = false; if (auto value = dyn_cast(offset)) { - if (!isProjectedOffsetValue(value, *laneArg, loops, usesDynamicBinding)) + if (!isProjectedOffsetValue(value, laneArg, loops, usesDynamicBinding)) return std::nullopt; } else if (!isa(offset)) @@ -2875,6 +3181,57 @@ static std::optional matchAffineProjectedInputSl return match; } +static std::optional matchAffineProjectedInputSlice(SpatComputeBatch batch, + unsigned inputIndex) { + const auto fail = [&](StringRef) -> std::optional { return std::nullopt; }; + + std::optional inputArg = batch.getInputArgument(inputIndex); + std::optional laneArg = batch.getLaneArgument(); + if (!inputArg || !laneArg) + return fail("missing-input-or-lane-arg"); + + std::optional projectedMatch; + for (OpOperand& use : inputArg->getUses()) { + auto extract = dyn_cast(use.getOwner()); + if (!extract || extract.getSource() != *inputArg) + return fail("input-user-is-not-direct-extract-slice"); + + std::optional current = + matchAffineProjectedInputSlice(*inputArg, *laneArg, extract); + if (!current) + return fail("input-extract-is-not-projectable"); + if (projectedMatch) + return fail("multiple-projected-input-extracts"); + projectedMatch = std::move(current); + } + + if (!projectedMatch) + return fail("input-arg-has-no-uses"); + return projectedMatch; +} + +static SmallVector collectAffineProjectedInputSlices(SpatComputeBatch batch, + unsigned inputIndex) { + SmallVector matches; + std::optional inputArg = batch.getInputArgument(inputIndex); + std::optional laneArg = batch.getLaneArgument(); + if (!inputArg || !laneArg) + return matches; + + for (OpOperand& use : inputArg->getUses()) { + auto extract = dyn_cast(use.getOwner()); + if (!extract || extract.getSource() != *inputArg) + return {}; + + std::optional match = + matchAffineProjectedInputSlice(*inputArg, *laneArg, extract); + if (!match) + return {}; + matches.push_back(std::move(*match)); + } + return matches; +} + std::optional getProjectedInputSliceMatch(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex) { ProjectedBatchInputKey key {batch.getOperation(), inputIndex}; @@ -2951,6 +3308,18 @@ FailureOr evaluateProjectionIndexLike(Value value, Value laneArg, uint3 if (std::optional constant = matchConstantIndexValue(value)) return *constant; + if (auto extract = value.getDefiningOp()) { + if (extract.getIndices().size() != 1) + return failure(); + DenseIntElementsAttr denseAttr = getDenseIndexVector(extract.getTensor()); + if (!denseAttr) + return failure(); + FailureOr index = evaluateProjectionIndexLike(extract.getIndices().front(), laneArg, lane); + if (failed(index)) + return failure(); + return getDenseIndexVectorElement(denseAttr, *index); + } + auto affineApply = value.getDefiningOp(); if (!affineApply || affineApply.getAffineMap().getNumResults() != 1) return failure(); @@ -3092,6 +3461,124 @@ bool isProjectedInputSliceCompatibleWithProducerFragments(SpatComputeBatch consu return recurse(recurse, 0); } + +SmallVector +collectProjectedFragmentDemandsForMatch(SpatComputeBatch consumerBatch, + Value input, + ComputeInstance logicalConsumer, + const AffineProjectedInputSliceMatch& match) { + SmallVector demands; + + auto producerResult = dyn_cast(input); + if (!producerResult) + return demands; + + auto producerBatch = dyn_cast_or_null(producerResult.getOwner()); + if (!producerBatch || producerBatch.getNumResults() == 0) + return demands; + + FailureOr producerProjection = + getBatchResultProjectionInsert(producerBatch, producerResult.getResultNumber()); + if (failed(producerProjection)) + return demands; + + std::optional producerLaneArg = producerBatch.getLaneArgument(); + std::optional consumerLaneArg = consumerBatch.getLaneArgument(); + if (!producerLaneArg || !consumerLaneArg) + return demands; + + unsigned ordinal = 0; + + const auto appendProducerForCurrentProjectedFragment = [&](const AffineProjectedInputSliceMatch& match, + ArrayRef loopIterationIndices) -> LogicalResult { + SmallVector consumerSizes(match.fragmentShape.begin(), match.fragmentShape.end()); + SmallVector consumerOffsets; + consumerOffsets.reserve(match.offsets.size()); + for (OpFoldResult offset : match.offsets) { + FailureOr evaluated = evaluateProjectedOffsetValue( + offset, *consumerLaneArg, logicalConsumer.laneStart, match.loops, loopIterationIndices); + if (failed(evaluated)) + return failure(); + consumerOffsets.push_back(*evaluated); + } + + for (uint32_t producerLane = 0; producerLane < static_cast(producerBatch.getLaneCount()); ++producerLane) { + FailureOr> producerOffsets = + evaluateStaticProjectionIndices(producerProjection->getMixedOffsets(), *producerLaneArg, producerLane); + FailureOr> producerSizes = + evaluateStaticProjectionIndices(producerProjection->getMixedSizes(), *producerLaneArg, producerLane); + FailureOr> producerStrides = + evaluateStaticProjectionIndices(producerProjection->getMixedStrides(), *producerLaneArg, producerLane); + if (failed(producerOffsets) || failed(producerSizes) || failed(producerStrides)) + return failure(); + if (!areAllUnitStrides(*producerStrides)) + return failure(); + if (!isStaticSliceContainedIn(consumerOffsets, consumerSizes, *producerOffsets, *producerSizes)) + continue; + + demands.push_back(ProjectedProducerFragmentDemand { + getBatchLaneProducerKey(producerBatch, producerLane, 1, producerResult.getResultNumber()), + std::move(consumerOffsets), + ordinal++}); + return success(); + } + + return failure(); + }; + + SmallVector loopIterationIndices(match.loops.size(), 0); + const auto recurse = [&](auto&& self, size_t loopIndex) -> LogicalResult { + if (loopIndex == match.loops.size()) + return appendProducerForCurrentProjectedFragment(match, loopIterationIndices); + + for (int64_t iteration = 0; iteration < match.loops[loopIndex].tripCount; ++iteration) { + loopIterationIndices[loopIndex] = iteration; + if (failed(self(self, loopIndex + 1))) + return failure(); + } + return success(); + }; + + LogicalResult result = match.loops.empty() ? appendProducerForCurrentProjectedFragment(match, loopIterationIndices) + : recurse(recurse, 0); + + if (failed(result)) + demands.clear(); + return demands; +} + +SmallVector +collectProjectedFragmentDemandsForBatchInput(MaterializerState& state, + SpatComputeBatch consumerBatch, + unsigned inputIndex, + Value input, + ComputeInstance logicalConsumer) { + (void)state; + SmallVector demands; + for (const AffineProjectedInputSliceMatch& match : collectAffineProjectedInputSlices(consumerBatch, inputIndex)) { + SmallVector matchDemands = + collectProjectedFragmentDemandsForMatch(consumerBatch, input, logicalConsumer, match); + if (matchDemands.empty()) + return {}; + llvm::append_range(demands, matchDemands); + } + return demands; +} + +SmallVector collectProjectedProducerKeysForBatchInput(MaterializerState& state, + SpatComputeBatch consumerBatch, + unsigned inputIndex, + Value input, + ComputeInstance logicalConsumer) { + SmallVector keys; + for (const ProjectedProducerFragmentDemand& demand : collectProjectedFragmentDemandsForBatchInput( + state, consumerBatch, inputIndex, input, logicalConsumer)) { + if (!llvm::is_contained(keys, demand.producer)) + keys.push_back(demand.producer); + } + return keys; +} + LogicalResult collectProjectedTransfers(MaterializerState& state) { struct PendingProjectedTransferDescriptor { ProjectedBatchInputKey inputKey; @@ -3187,7 +3674,90 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { } for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs())) { - SmallVector producers = collectProducerKeysForDestinations(input, logicalConsumer); + SmallVector fragmentDemands = + collectProjectedFragmentDemandsForBatchInput( + state, batch, static_cast(inputIndex), input, logicalConsumer); + SmallVector matches = + collectAffineProjectedInputSlices(batch, static_cast(inputIndex)); + + if (!matches.empty() && !fragmentDemands.empty()) { + for (AffineProjectedInputSliceMatch& match : matches) { + SmallVector matchDemands = + collectProjectedFragmentDemandsForMatch(batch, input, logicalConsumer, match); + if (matchDemands.empty()) + return targetClass.op->emitError("failed to collect projected input transfer fragments"); + + ProjectedInputTransferPlan& plan = + state.projectedInputTransferPlans[match.extract.getOperation()][targetClassId]; + ProjectedBatchInputKey currentInputKey {batch.getOperation(), static_cast(inputIndex)}; + if (plan.fragments.empty()) { + plan.inputKey = currentInputKey; + plan.extractOp = match.extract.getOperation(); + plan.layout.fragmentType = match.fragmentType; + plan.layout.fragmentShape = match.fragmentShape; + plan.layout.loopLowerBounds.reserve(match.loops.size()); + plan.layout.loopSteps.reserve(match.loops.size()); + plan.layout.loopTripCounts.reserve(match.loops.size()); + for (const StaticProjectedLoopInfo& loop : match.loops) { + plan.layout.loopLowerBounds.push_back(loop.lowerBound); + plan.layout.loopSteps.push_back(loop.step); + plan.layout.loopTripCounts.push_back(loop.tripCount); + } + plan.layout.fragmentsPerLogicalSlot = getProjectedFragmentsPerLogicalSlot(plan.layout.loopTripCounts); + plan.layout.payloadFragmentCount = plan.layout.fragmentsPerLogicalSlot; + } + if (!(plan.inputKey == currentInputKey) || plan.extractOp != match.extract.getOperation() + || plan.layout.fragmentType != match.fragmentType || plan.layout.fragmentShape != match.fragmentShape + || plan.layout.loopTripCounts.size() != match.loops.size()) + return targetClass.op->emitError("inconsistent projected input transfer plan"); + + auto targetCpu = getCheckedCoreId(targetClass.op, cpu, "projected input target core id"); + if (failed(targetCpu)) + return failure(); + + for (const ProjectedProducerFragmentDemand& demand : matchDemands) { + ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, demand.producer.instance); + auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); + if (producerCpuIt == state.schedule.computeToCpuMap.end()) + return logicalConsumer.op->emitError( + "projected input transfer found a fragment produced by an unscheduled compute"); + + auto sourceCpu = getCheckedCoreId( + targetClass.op, producerCpuIt->second, "projected input source core id"); + if (failed(sourceCpu)) + return failure(); + + ProjectedInputTransferFragment fragment; + fragment.producer = demand.producer; + fragment.fragmentOffsets = demand.fragmentOffsets; + fragment.targetLane = targetLane; + fragment.ordinal = targetClass.isBatch + ? targetLane * plan.layout.fragmentsPerLogicalSlot + demand.ordinal + : demand.ordinal; + fragment.channelId = state.nextChannelId++; + fragment.sourceCoreId = *sourceCpu; + fragment.targetCoreId = *targetCpu; + if (fragment.sourceCoreId > fragment.targetCoreId) { + state.projectedInputPhaseBarrierClasses.insert(targetClassId); + ++state.pendingProjectedHighToLowReceives[targetClassId]; + } + plan.fragments.push_back(std::move(fragment)); + } + } + continue; + } + + std::optional match = + getProjectedInputSliceMatch(state, batch, static_cast(inputIndex)); + SmallVector producers; + if (!fragmentDemands.empty()) { + for (const ProjectedProducerFragmentDemand& demand : fragmentDemands) + if (!llvm::is_contained(producers, demand.producer)) + producers.push_back(demand.producer); + } + else { + producers = collectProducerKeysForDestinations(input, logicalConsumer); + } if (producers.size() != 1) continue; ProducerKey producer = producers.front(); @@ -3201,8 +3771,6 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { if (sourceClassId == targetClassId) continue; - std::optional match = - getProjectedInputSliceMatch(state, batch, static_cast(inputIndex)); if (!match) continue; if (!isProjectedInputSliceCompatibleWithProducerFragments( @@ -3291,15 +3859,6 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { if (payloadFragmentCount == 0) continue; - // Batch-target projected replacements currently select fragments with the - // local materialization-run slot index. That is only unambiguous when each - // target lane receives one projected fragment. Multi-fragment payloads - // need an explicit producer-key to payload-slot mapping; otherwise two - // independently materialized runs can both select fragment 0 from the same - // packed receive and duplicate rows. - if (payloadFragmentCount != 1) - continue; - bool uniform = true; for (ArrayRef> laneFragments : pendingDescriptor.fragmentOffsetsByLane) { if (laneFragments.size() != payloadFragmentCount) { @@ -3396,23 +3955,6 @@ collectScalarTargetProjectedDescriptor(MaterializerState& state, return combined; } -bool haveSameDestinationClasses(MaterializerState& state, ArrayRef keys) { - if (keys.empty()) - return true; - - auto firstIt = state.producerDestClasses.find(keys.front()); - ArrayRef first = firstIt == state.producerDestClasses.end() ? ArrayRef() : firstIt->second; - for (ProducerKey key : keys.drop_front()) { - auto it = state.producerDestClasses.find(key); - ArrayRef current = it == state.producerDestClasses.end() ? ArrayRef() : it->second; - if (first.size() != current.size()) - return false; - for (auto [lhs, rhs] : llvm::zip(first, current)) - if (lhs != rhs) - return false; - } - return true; -} ArrayRef getDestinationClasses(MaterializerState& state, ProducerKey key) { auto it = state.producerDestClasses.find(key); @@ -3421,10 +3963,126 @@ ArrayRef getDestinationClasses(MaterializerState& state, ProducerKey ke return it->second; } +SmallVector filterProducerKeysForDestination(MaterializerState& state, + ArrayRef keys, + ClassId destinationClass) { + SmallVector filtered; + for (ProducerKey key : keys) + if (llvm::is_contained(getDestinationClasses(state, key), destinationClass)) + filtered.push_back(key); + return filtered; +} + // ----------------------------------------------------------------------------- // Communication materialization helpers. // ----------------------------------------------------------------------------- +void appendScalarSendAtCurrentInsertionPoint(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + int64_t channelId, + int32_t sourceCoreId, + int32_t targetCoreId, + Location loc) { + Value channelIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, channelId); + Value sourceCoreIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, sourceCoreId); + Value targetCoreIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, targetCoreId); + SpatChannelSendOp::create(state.rewriter, loc, channelIdValue, sourceCoreIdValue, targetCoreIdValue, payload); +} + +bool isMaterializedScalarCore(const MaterializerState& state, int64_t coreId) { + return llvm::any_of(state.classes, [&](const MaterializedClass& klass) { + return !klass.isBatch && klass.cpus.size() == 1 && static_cast(klass.cpus.front()) == coreId; + }); +} + +bool shouldEmitScalarPeerSendStatically(const MaterializerState& state, + const MaterializedClass& sourceClass, + int64_t sourceCoreId, + int64_t targetCoreId) { + if (sourceClass.isBatch || sourceCoreId >= targetCoreId) + return false; + return isMaterializedScalarCore(state, targetCoreId); +} + +LogicalResult appendScalarSendLoopAtCurrentInsertionPoint(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + const MessageVector& messages, + Location loc) { + MessageVector loopMessages; + MessageVector delayedPeerMessages; + loopMessages.channelIds.reserve(messages.size()); + loopMessages.sourceCoreIds.reserve(messages.size()); + loopMessages.targetCoreIds.reserve(messages.size()); + delayedPeerMessages.channelIds.reserve(messages.size()); + delayedPeerMessages.sourceCoreIds.reserve(messages.size()); + delayedPeerMessages.targetCoreIds.reserve(messages.size()); + + for (size_t index = 0; index < messages.size(); ++index) { + int64_t sourceCoreId = messages.sourceCoreIds[index]; + int64_t targetCoreId = messages.targetCoreIds[index]; + if (shouldEmitScalarPeerSendStatically(state, sourceClass, sourceCoreId, targetCoreId)) { + delayedPeerMessages.append(messages.channelIds[index], sourceCoreId, targetCoreId); + continue; + } + loopMessages.append(messages.channelIds[index], sourceCoreId, targetCoreId); + } + + auto emitStaticPeerSends = [&]() { + for (size_t index = 0; index < delayedPeerMessages.size(); ++index) + appendScalarSendAtCurrentInsertionPoint(state, + sourceClass, + payload, + delayedPeerMessages.channelIds[index], + static_cast(delayedPeerMessages.sourceCoreIds[index]), + static_cast(delayedPeerMessages.targetCoreIds[index]), + loc); + }; + + if (loopMessages.empty()) { + emitStaticPeerSends(); + return success(); + } + + if (loopMessages.size() == 1) { + appendScalarSendAtCurrentInsertionPoint(state, + sourceClass, + payload, + loopMessages.channelIds.front(), + static_cast(loopMessages.sourceCoreIds.front()), + static_cast(loopMessages.targetCoreIds.front()), + loc); + emitStaticPeerSends(); + return success(); + } + + Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0); + Value upperBound = + getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast(loopMessages.size())); + Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1); + + auto sendLoop = buildNormalizedScfFor( + state.rewriter, + loc, + lowerBound, + upperBound, + step, + ValueRange {}, + [&](OpBuilder&, Location, Value index, ValueRange, SmallVectorImpl&) { + Value channelId = createIndexedChannelId(state, sourceClass.op, loopMessages, index, loc); + Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, loopMessages, index, loc); + Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, loopMessages, index, loc); + SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); + return success(); + }); + if (failed(sendLoop)) + return failure(); + + emitStaticPeerSends(); + return success(); +} + void appendScalarSend(MaterializerState& state, MaterializedClass& sourceClass, Value payload, @@ -3435,10 +4093,7 @@ void appendScalarSend(MaterializerState& state, assert(!sourceClass.isBatch && "scalar send helper expects a scalar source class"); state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - Value channelIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, channelId); - Value sourceCoreIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, sourceCoreId); - Value targetCoreIdValue = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, targetCoreId); - SpatChannelSendOp::create(state.rewriter, loc, channelIdValue, sourceCoreIdValue, targetCoreIdValue, payload); + appendScalarSendAtCurrentInsertionPoint(state, sourceClass, payload, channelId, sourceCoreId, targetCoreId, loc); } LogicalResult appendScalarSendLoop(MaterializerState& state, @@ -3451,29 +4106,7 @@ LogicalResult appendScalarSendLoop(MaterializerState& state, assert(succeeded(messages.verify(sourceClass.op)) && "message metadata is inconsistent"); state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0); - Value upperBound = - getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast(messages.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1); - - auto sendLoop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {}, - [&](OpBuilder&, Location, Value index, ValueRange, SmallVectorImpl&) { - Value channelId = createIndexedChannelId(state, sourceClass.op, messages, index, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, messages, index, loc); - Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, messages, index, loc); - SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); - return success(); - }); - if (failed(sendLoop)) - return failure(); - return success(); + return appendScalarSendLoopAtCurrentInsertionPoint(state, sourceClass, payload, messages, loc); } FailureOr buildProjectedPackedPayload(MaterializerState& state, @@ -3579,19 +4212,23 @@ FailureOr buildProjectedPayloadForMessage(MaterializerState& state, return buildProjectedPackedPayload(state, targetClass, *localizedPayload, descriptor, messageIndex, loc); } -LogicalResult appendProjectedScalarSendLoop(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const ProjectedTransferDescriptor& descriptor, - const MessageVector& messages, - Location loc) { - assert(!sourceClass.isBatch && "projected scalar send expects scalar source class"); - assert(succeeded(messages.verify(sourceClass.op)) && "message metadata is inconsistent"); - if (failed(verifyProjectedSendDescriptor(sourceClass.op, descriptor, messages))) - return failure(); +FailureOr getProjectedCommunicationType(Operation* anchor, const ProjectedTransferDescriptor& descriptor) { + if (descriptor.layout.payloadFragmentCount == 1) + return Type(descriptor.layout.fragmentType); - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); + FailureOr packedType = + getPackedBatchTensorType(descriptor.layout.fragmentType, descriptor.layout.payloadFragmentCount); + if (failed(packedType)) + return anchor->emitError("cannot compute projected communication payload type"), failure(); + return Type(*packedType); +} +LogicalResult appendProjectedScalarSendLoopAtCurrentInsertionPoint(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + const ProjectedTransferDescriptor& descriptor, + const MessageVector& messages, + Location loc) { if (messages.size() == 1) { Value channelId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, messages.channelIds.front()); Value sourceCoreId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, messages.sourceCoreIds.front()); @@ -3633,6 +4270,21 @@ LogicalResult appendProjectedScalarSendLoop(MaterializerState& state, return success(); } +LogicalResult appendProjectedScalarSendLoop(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + const ProjectedTransferDescriptor& descriptor, + const MessageVector& messages, + Location loc) { + assert(!sourceClass.isBatch && "projected scalar send expects scalar source class"); + assert(succeeded(messages.verify(sourceClass.op)) && "message metadata is inconsistent"); + if (failed(verifyProjectedSendDescriptor(sourceClass.op, descriptor, messages))) + return failure(); + + state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); + return appendProjectedScalarSendLoopAtCurrentInsertionPoint(state, sourceClass, payload, descriptor, messages, loc); +} + LogicalResult appendSend(MaterializerState& state, MaterializedClass& sourceClass, Value payload, @@ -3654,6 +4306,7 @@ LogicalResult appendSend(MaterializerState& state, return success(); } + state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); if (messages.size() == 1) { appendScalarSend(state, sourceClass, @@ -3797,6 +4450,81 @@ struct ProjectedScalarSendGroup { ProjectedTransferDescriptor descriptor; }; +ProjectedScalarSendGroup orderedProjectedScalarSendGroup(const ProjectedScalarSendGroup& group); +LogicalResult finalizeSplitProjectedScalarSendGroup(Operation* anchor, ProjectedScalarSendGroup& group); + +struct PeerExchangeOrderKey { + int64_t maxCore = 0; + int64_t sourceCore = 0; + int64_t targetCore = 0; + int64_t channelId = 0; +}; + +PeerExchangeOrderKey makePeerExchangeOrderKey(int64_t sourceCore, int64_t targetCore, int64_t channelId) { + return PeerExchangeOrderKey { + .maxCore = std::max(sourceCore, targetCore), + .sourceCore = sourceCore, + .targetCore = targetCore, + .channelId = channelId, + }; +} + +PeerExchangeOrderKey makePeerExchangeOrderKey(const MessageVector& messages, size_t index) { + return makePeerExchangeOrderKey(messages.sourceCoreIds[index], messages.targetCoreIds[index], messages.channelIds[index]); +} + +bool peerExchangeOrderLess(PeerExchangeOrderKey lhs, PeerExchangeOrderKey rhs) { + if (lhs.maxCore != rhs.maxCore) + return lhs.maxCore > rhs.maxCore; + if (lhs.sourceCore != rhs.sourceCore) + return lhs.sourceCore > rhs.sourceCore; + if (lhs.targetCore != rhs.targetCore) + return lhs.targetCore > rhs.targetCore; + return lhs.channelId < rhs.channelId; +} + +void sortMessageIndicesByPeerExchangeOrder(const MessageVector& messages, SmallVectorImpl& indices) { + llvm::stable_sort(indices, [&](size_t lhs, size_t rhs) { + return peerExchangeOrderLess(makePeerExchangeOrderKey(messages, lhs), makePeerExchangeOrderKey(messages, rhs)); + }); +} + +MessageVector orderedMessageVector(const MessageVector& messages) { + SmallVector indices; + indices.reserve(messages.size()); + for (size_t index = 0; index < messages.size(); ++index) + indices.push_back(index); + sortMessageIndicesByPeerExchangeOrder(messages, indices); + + MessageVector ordered; + for (size_t index : indices) + ordered.append(messages.channelIds[index], messages.sourceCoreIds[index], messages.targetCoreIds[index]); + return ordered; +} + +std::optional getBestPeerExchangeOrderKey(const MessageVector& messages) { + if (messages.empty()) + return std::nullopt; + + PeerExchangeOrderKey best = makePeerExchangeOrderKey(messages, 0); + for (size_t index = 1; index < messages.size(); ++index) { + PeerExchangeOrderKey candidate = makePeerExchangeOrderKey(messages, index); + if (peerExchangeOrderLess(candidate, best)) + best = candidate; + } + return best; +} + +void sortBatchRunSendPlansByPeerExchangeOrder(SmallVectorImpl& plans) { + llvm::stable_sort(plans, [](const BatchRunSendPlan& lhs, const BatchRunSendPlan& rhs) { + std::optional lhsKey = getBestPeerExchangeOrderKey(lhs.messages); + std::optional rhsKey = getBestPeerExchangeOrderKey(rhs.messages); + if (!lhsKey || !rhsKey) + return lhsKey.has_value() && !rhsKey.has_value(); + return peerExchangeOrderLess(*lhsKey, *rhsKey); + }); +} + struct ScalarSourceFanoutPlan { SmallVector receivePlans; std::optional ordinaryMessages; @@ -3910,8 +4638,11 @@ FailureOr buildScalarSourceFanoutPlan(MaterializerState& if (*descriptor) { const ProjectedTransferDescriptor& projectedDescriptor = **descriptor; + FailureOr projectedCommunicationType = getProjectedCommunicationType(sourceClass.op, projectedDescriptor); + if (failed(projectedCommunicationType)) + return failure(); - receivePlan.receiveType = projectedDescriptor.payloadType; + receivePlan.receiveType = *projectedCommunicationType; receivePlan.projectedExtractOp = projectedDescriptor.extractOp; receivePlan.projectedLayout = projectedDescriptor.layout; @@ -3952,18 +4683,1215 @@ FailureOr buildScalarSourceFanoutPlan(MaterializerState& return fanoutPlan; } +LogicalResult collectPlannedScalarPeerReceivesForSource(MaterializerState& state, + MaterializedClass& sourceClass, + ArrayRef keys, + Value payload) { + SmallVector destinationClasses = collectDestinationClassesForKeys(state, keys); + int64_t savedNextChannelId = state.nextChannelId; + auto fanoutPlan = buildScalarSourceFanoutPlan(state, sourceClass, keys, destinationClasses, payload); + state.nextChannelId = savedNextChannelId; + if (failed(fanoutPlan)) + return failure(); + + for (const ScalarSourceReceivePlan& plan : fanoutPlan->receivePlans) { + for (size_t index = 0; index < plan.messages.size(); ++index) { + ScalarPeerEdgeKey receiveKey { + .sourceCore = plan.messages.sourceCoreIds[index], + .targetCore = plan.messages.targetCoreIds[index], + .payloadType = plan.receiveType, + }; + bool alreadyPlanned = llvm::any_of(state.plannedScalarPeerReceives, [&](const ScalarPeerEdgeKey& record) { + return record.sourceCore == receiveKey.sourceCore && record.targetCore == receiveKey.targetCore + && record.payloadType == receiveKey.payloadType; + }); + if (!alreadyPlanned) + state.plannedScalarPeerReceives.push_back(receiveKey); + } + } + + return success(); +} + LogicalResult emitScalarSourceFanoutSends(MaterializerState& state, MaterializedClass& sourceClass, Value payload, const ScalarSourceFanoutPlan& plan, Location loc) { - if (plan.ordinaryMessages && failed(appendSend(state, sourceClass, payload, *plan.ordinaryMessages, loc))) + if (plan.ordinaryMessages) { + MessageVector orderedMessages = orderedMessageVector(*plan.ordinaryMessages); + if (failed(appendSend(state, sourceClass, payload, orderedMessages, loc))) + return failure(); + } + + for (const ProjectedScalarSendGroup& group : plan.projectedSendGroups) { + ProjectedScalarSendGroup orderedGroup = orderedProjectedScalarSendGroup(group); + if (failed(finalizeSplitProjectedScalarSendGroup(sourceClass.op, orderedGroup))) + return failure(); + if (failed(appendProjectedScalarSendLoop( + state, sourceClass, payload, orderedGroup.descriptor, orderedGroup.messages, loc))) + return failure(); + } + + return success(); +} + +ProjectedScalarSendGroup makeEmptyProjectedScalarSendGroup(const ProjectedScalarSendGroup& group) { + ProjectedScalarSendGroup splitGroup; + splitGroup.descriptor = group.descriptor; + splitGroup.descriptor.fragmentOffsets.clear(); + splitGroup.descriptor.fragmentOffsetsByDim.clear(); + return splitGroup; +} + +void appendProjectedScalarSendMessage(ProjectedScalarSendGroup& destination, + const ProjectedScalarSendGroup& source, + size_t index) { + const unsigned payloadFragmentCount = source.descriptor.layout.payloadFragmentCount; + destination.messages.append(source.messages.channelIds[index], + source.messages.sourceCoreIds[index], + source.messages.targetCoreIds[index]); + ArrayRef> offsets( + source.descriptor.fragmentOffsets.data() + index * payloadFragmentCount, payloadFragmentCount); + llvm::append_range(destination.descriptor.fragmentOffsets, offsets); +} + +ProjectedScalarSendGroup orderedProjectedScalarSendGroup(const ProjectedScalarSendGroup& group) { + SmallVector indices; + indices.reserve(group.messages.size()); + for (size_t index = 0; index < group.messages.size(); ++index) + indices.push_back(index); + sortMessageIndicesByPeerExchangeOrder(group.messages, indices); + + ProjectedScalarSendGroup ordered = makeEmptyProjectedScalarSendGroup(group); + for (size_t index : indices) + appendProjectedScalarSendMessage(ordered, group, index); + return ordered; +} + +LogicalResult finalizeSplitProjectedScalarSendGroup(Operation* anchor, ProjectedScalarSendGroup& group) { + if (group.messages.empty()) + return success(); + if (failed(finalizeProjectedTransferDescriptor(anchor, group.descriptor))) + return failure(); + if (failed(verifyProjectedSendDescriptor(anchor, group.descriptor, group.messages))) + return failure(); + return success(); +} + +ScalarPeerEdgeKey makeScalarPeerEdgeKey(int64_t sourceCore, int64_t targetCore, Type payloadType) { + return ScalarPeerEdgeKey {.sourceCore = sourceCore, .targetCore = targetCore, .payloadType = payloadType}; +} + +ScalarPeerReceiveKey makeScalarPeerReceiveKey(int64_t sourceCore, + int64_t targetCore, + Type payloadType, + std::optional channelId) { + return ScalarPeerReceiveKey { + .sourceCore = sourceCore, + .targetCore = targetCore, + .payloadType = payloadType, + .channelId = channelId, + }; +} + +ScalarPeerEdgeKey makeRequiredPeerReceiveForLowerToHigherSend(int64_t sourceCore, int64_t targetCore, Type payloadType) { + return makeScalarPeerEdgeKey(targetCore, sourceCore, payloadType); +} + +bool hasMaterializedScalarPeerReceive(const MaterializerState& state, const ScalarPeerEdgeKey& key) { + return llvm::any_of(state.materializedScalarPeerReceives, [&](const ScalarPeerEdgeKey& record) { + return record.sourceCore == key.sourceCore && record.targetCore == key.targetCore && record.payloadType == key.payloadType; + }); +} + +bool hasPlannedScalarPeerReceive(const MaterializerState& state, const ScalarPeerEdgeKey& key) { + return llvm::any_of(state.plannedScalarPeerReceives, [&](const ScalarPeerEdgeKey& record) { + return record.sourceCore == key.sourceCore && record.targetCore == key.targetCore && record.payloadType == key.payloadType; + }); +} + +bool receiveKeyMatchesPendingWait(const ScalarPeerReceiveKey& receiveKey, const ScalarPeerEdgeKey& pendingWait) { + return receiveKey.sourceCore == pendingWait.sourceCore && receiveKey.targetCore == pendingWait.targetCore + && receiveKey.payloadType == pendingWait.payloadType; +} + +bool hasPendingScalarSendWaitingForReceive(const MaterializerState& state, const ScalarPeerReceiveKey& receiveKey) { + return llvm::any_of(state.pendingScalarSends, [&](const PendingScalarSend& pending) { + return receiveKeyMatchesPendingWait(receiveKey, pending.waitForReceive); + }) || llvm::any_of(state.pendingProjectedScalarSends, [&](const PendingProjectedScalarSend& pending) { + return receiveKeyMatchesPendingWait(receiveKey, pending.waitForReceive); + }); +} + +ReceiveMessagePartition partitionReceivesByPendingSendUnlocks(const MaterializerState& state, + const MaterializedClass& targetClass, + Type payloadType, + const MessageVector& messages, + bool preserveMessageOrder) { + (void) targetClass; + ReceiveMessagePartition partition; + for (size_t index = 0; index < messages.size(); ++index) { + ScalarPeerReceiveKey receiveKey = makeScalarPeerReceiveKey(messages.sourceCoreIds[index], + messages.targetCoreIds[index], + payloadType, + messages.channelIds[index]); + bool isCritical = hasPendingScalarSendWaitingForReceive(state, receiveKey); + (isCritical ? partition.criticalStaticIndices : partition.remainingIndices).push_back(index); + } + + if (!preserveMessageOrder) { + sortMessageIndicesByPeerExchangeOrder(messages, partition.criticalStaticIndices); + sortMessageIndicesByPeerExchangeOrder(messages, partition.remainingIndices); + } + return partition; +} + +FailureOr findTopLevelDefAnchorInClass(MaterializedClass& klass, Value value) { + if (auto blockArg = dyn_cast(value)) { + if (blockArg.getOwner() == klass.body) + return static_cast(nullptr); + klass.op->emitError("pending scalar send payload block argument is not local to class"); + return failure(); + } + + Operation* definingOp = value.getDefiningOp(); + if (!definingOp) + return klass.op->emitError("pending scalar send payload has no defining op"), failure(); + + Operation* current = definingOp; + while (current) { + if (current->getBlock() == klass.body) + return current; + current = current->getParentOp(); + } + + klass.op->emitError("pending scalar send payload is defined outside the source materialized class"); + return failure(); +} + +Operation* latestAnchorInBlock(Operation* a, Operation* b) { + if (!a) + return b; + if (!b) + return a; + if (a->getBlock() != b->getBlock()) + return nullptr; + return a->isBeforeInBlock(b) ? b : a; +} + +FailureOr resolvePendingSendInsertionAnchor(MaterializedClass& sourceClass, + Operation* receiveAnchor, + Operation* payloadAnchor) { + if (!receiveAnchor) + return sourceClass.op->emitError("missing receive anchor while flushing pending scalar send"), failure(); + if (!payloadAnchor) + return receiveAnchor; + if (receiveAnchor->getBlock() == payloadAnchor->getBlock()) + return latestAnchorInBlock(receiveAnchor, payloadAnchor); + + Operation* current = receiveAnchor; + while (current) { + Operation* parent = current->getParentOp(); + if (!parent) + break; + if (payloadAnchor->getBlock() == parent->getBlock()) { + if (payloadAnchor->isBeforeInBlock(current)) + return receiveAnchor; + return sourceClass.op->emitError("pending scalar send payload does not dominate nested receive flush point"), + failure(); + } + current = parent; + } + + return sourceClass.op->emitError("pending scalar send payload anchor is incompatible with nested receive flush point"), + failure(); +} + +LogicalResult enqueuePendingScalarSend(MaterializerState& state, + ClassId sourceClass, + int64_t sourceCore, + int64_t targetCore, + Type payloadType, + Value payload, + const MessageVector& messages, + Location loc) { + MaterializedClass& sourceMaterializedClass = state.classes[sourceClass]; + if (failed(messages.verify(sourceMaterializedClass.op))) + return failure(); + FailureOr payloadAnchor = findTopLevelDefAnchorInClass(sourceMaterializedClass, payload); + if (failed(payloadAnchor)) + return failure(); + state.pendingScalarSends.push_back(PendingScalarSend { + .sourceClass = sourceClass, + .sourceCore = sourceCore, + .targetCore = targetCore, + .payloadType = payloadType, + .waitForReceive = makeRequiredPeerReceiveForLowerToHigherSend(sourceCore, targetCore, payloadType), + .payloadAnchor = *payloadAnchor, + .payload = payload, + .messages = messages, + .loc = loc, + }); + return success(); +} + +LogicalResult enqueuePendingProjectedScalarSend(MaterializerState& state, + ClassId sourceClass, + int64_t sourceCore, + int64_t targetCore, + Type payloadType, + Value payload, + const ProjectedScalarSendGroup& group, + Location loc) { + MaterializedClass& sourceMaterializedClass = state.classes[sourceClass]; + if (failed(verifyProjectedSendDescriptor(sourceMaterializedClass.op, group.descriptor, group.messages))) + return failure(); + FailureOr payloadAnchor = findTopLevelDefAnchorInClass(sourceMaterializedClass, payload); + if (failed(payloadAnchor)) + return failure(); + state.pendingProjectedScalarSends.push_back(PendingProjectedScalarSend { + .sourceClass = sourceClass, + .sourceCore = sourceCore, + .targetCore = targetCore, + .payloadType = payloadType, + .waitForReceive = makeRequiredPeerReceiveForLowerToHigherSend(sourceCore, targetCore, payloadType), + .payloadAnchor = *payloadAnchor, + .payload = payload, + .messages = group.messages, + .descriptor = group.descriptor, + .loc = loc, + }); + return success(); +} + +LogicalResult flushPendingScalarSendsForReceive( + MaterializerState& state, const ScalarPeerEdgeKey& receiveKey, Operation* anchor) { + for (size_t index = 0; index < state.pendingProjectedScalarSends.size();) { + const PendingProjectedScalarSend& pending = state.pendingProjectedScalarSends[index]; + if (pending.waitForReceive.sourceCore != receiveKey.sourceCore + || pending.waitForReceive.targetCore != receiveKey.targetCore + || pending.waitForReceive.payloadType != receiveKey.payloadType) { + ++index; + continue; + } + + MaterializedClass& sourceClass = state.classes[pending.sourceClass]; + FailureOr sendAnchor = resolvePendingSendInsertionAnchor(sourceClass, anchor, pending.payloadAnchor); + if (failed(sendAnchor)) + return failure(); + OpBuilder::InsertionGuard guard(state.rewriter); + state.rewriter.setInsertionPointAfter(*sendAnchor); + if (failed(appendProjectedScalarSendLoopAtCurrentInsertionPoint( + state, sourceClass, pending.payload, pending.descriptor, pending.messages, pending.loc))) + return failure(); + state.pendingProjectedScalarSends.erase(state.pendingProjectedScalarSends.begin() + index); + } + + for (size_t index = 0; index < state.pendingScalarSends.size();) { + const PendingScalarSend& pending = state.pendingScalarSends[index]; + if (pending.waitForReceive.sourceCore != receiveKey.sourceCore + || pending.waitForReceive.targetCore != receiveKey.targetCore + || pending.waitForReceive.payloadType != receiveKey.payloadType) { + ++index; + continue; + } + + MaterializedClass& sourceClass = state.classes[pending.sourceClass]; + FailureOr sendAnchor = resolvePendingSendInsertionAnchor(sourceClass, anchor, pending.payloadAnchor); + if (failed(sendAnchor)) + return failure(); + OpBuilder::InsertionGuard guard(state.rewriter); + state.rewriter.setInsertionPointAfter(*sendAnchor); + if (pending.messages.size() == 1) + appendScalarSendAtCurrentInsertionPoint(state, + sourceClass, + pending.payload, + pending.messages.channelIds.front(), + pending.messages.sourceCoreIds.front(), + pending.messages.targetCoreIds.front(), + pending.loc); + else if (failed(appendScalarSendLoopAtCurrentInsertionPoint( + state, sourceClass, pending.payload, pending.messages, pending.loc))) + return failure(); + state.pendingScalarSends.erase(state.pendingScalarSends.begin() + index); + } + + return success(); +} + +FailureOr createScalarPeerReceiveAndFlush(MaterializerState& state, + MaterializedClass& targetClass, + Type payloadType, + Value channelId, + Value sourceCoreId, + Value targetCoreId, + std::optional staticKey, + Location loc) { + OpBuilder::InsertionGuard guard(state.rewriter); + + Value received = + SpatChannelReceiveOp::create(state.rewriter, loc, payloadType, channelId, sourceCoreId, targetCoreId).getOutput(); + if (targetClass.isBatch || !staticKey || staticKey->sourceCore <= staticKey->targetCore) + return received; + + ScalarPeerEdgeKey receiveKey = + makeScalarPeerEdgeKey(staticKey->sourceCore, staticKey->targetCore, staticKey->payloadType); + if (!hasMaterializedScalarPeerReceive(state, receiveKey)) + state.materializedScalarPeerReceives.push_back(receiveKey); + if (failed(flushPendingScalarSendsForReceive(state, receiveKey, received.getDefiningOp()))) + return failure(); + return received; +} + +FailureOr appendReceiveAndFlushPendingScalarSends( + MaterializerState& state, MaterializedClass& targetClass, Type type, const MessageVector& messages, Location loc) { + assert(succeeded(messages.verify(targetClass.op)) && "message metadata is inconsistent"); + assert(!messages.empty() && "expected at least one receive"); + + state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); + if (targetClass.isBatch) + return appendReceive(state, targetClass, type, messages, loc); + if (messages.size() != 1) + return targetClass.op->emitError("scalar receive flush expects exactly one message") + << " messageCount=" << messages.size(), + failure(); + + return createScalarPeerReceiveAndFlush(state, + targetClass, + type, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages.channelIds.front()), + getOrCreateIndexConstant( + state.constantFolder, targetClass.op, messages.sourceCoreIds.front()), + getOrCreateIndexConstant( + state.constantFolder, targetClass.op, messages.targetCoreIds.front()), + makeScalarPeerReceiveKey(messages.sourceCoreIds.front(), + messages.targetCoreIds.front(), + type, + messages.channelIds.front()), + loc); +} + + + +LogicalResult appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + ArrayRef fragments, + Location loc) { + assert(sourceClass.isBatch && "conditional projected sends require a batch source class"); + if (fragments.empty()) + return success(); + + auto batch = cast(sourceClass.op); + auto laneArg = batch.getLaneArgument(); + if (!laneArg) + return batch.emitOpError("expected lane argument while emitting projected input sends"); + + SmallVector sourceLanes; + MessageVector messages; + sourceLanes.reserve(fragments.size()); + messages.channelIds.reserve(fragments.size()); + messages.sourceCoreIds.reserve(fragments.size()); + messages.targetCoreIds.reserve(fragments.size()); + + for (ProjectedInputTransferFragment* fragment : fragments) { + auto laneIt = sourceClass.cpuToLane.find(fragment->sourceCoreId); + if (laneIt == sourceClass.cpuToLane.end()) + return sourceClass.op->emitError("projected input send could not map source core to batch lane") + << " sourceCore=" << fragment->sourceCoreId; + sourceLanes.push_back(laneIt->second); + messages.append(fragment->channelId, fragment->sourceCoreId, fragment->targetCoreId); + } + + Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0); + Value upperBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast(fragments.size())); + Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1); + auto sendLoop = buildNormalizedScfFor( + state.rewriter, + loc, + lowerBound, + upperBound, + step, + ValueRange {}, + [&](OpBuilder&, Location loopLoc, Value index, ValueRange, SmallVectorImpl&) { + Value sourceLane = createIndexedIndexValue(state, sourceClass.op, sourceLanes, index, loopLoc); + Value isSender = arith::CmpIOp::create(state.rewriter, loopLoc, arith::CmpIPredicate::eq, *laneArg, sourceLane); + auto ifOp = scf::IfOp::create(state.rewriter, loopLoc, TypeRange {}, isSender, false); + OpBuilder::InsertionGuard guard(state.rewriter); + Block& thenBlock = ifOp.getThenRegion().front(); + if (!thenBlock.empty()) + if (auto yieldOp = dyn_cast(thenBlock.back())) + yieldOp.erase(); + state.rewriter.setInsertionPointToEnd(&thenBlock); + Value channelId = createIndexedChannelId(state, sourceClass.op, messages, index, loopLoc); + Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, messages, index, loopLoc); + Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, messages, index, loopLoc); + SpatChannelSendOp::create(state.rewriter, loopLoc, channelId, sourceCoreId, targetCoreId, payload); + scf::YieldOp::create(state.rewriter, loopLoc); + return success(); + }); + return failed(sendLoop) ? failure() : success(); +} + +LogicalResult appendConditionalBatchProjectedInputSends(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + ArrayRef fragments, + Location loc) { + state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); + return appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint(state, sourceClass, payload, fragments, loc); +} + +LogicalResult appendScalarProjectedInputSendNow(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + ProjectedInputTransferFragment& fragment, + Location loc) { + assert(!sourceClass.isBatch && "scalar projected input send requires a scalar source class"); + if (sourceClass.cpus.empty()) + return sourceClass.op->emitError("scalar projected input send requires one source core"); + int64_t sourceCore = sourceClass.cpus.front(); + if (sourceCore != fragment.sourceCoreId) + return sourceClass.op->emitError("projected input send source core does not match materialized scalar class") + << " expected=" << sourceCore << " actual=" << fragment.sourceCoreId; + + appendScalarSend(state, sourceClass, payload, fragment.channelId, fragment.sourceCoreId, fragment.targetCoreId, loc); + return success(); +} + +LogicalResult appendScalarProjectedInputSendAtCurrentInsertionPoint(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + ProjectedInputTransferFragment& fragment, + Location loc) { + assert(!sourceClass.isBatch && "scalar projected input send requires a scalar source class"); + if (sourceClass.cpus.empty()) + return sourceClass.op->emitError("scalar projected input send requires one source core"); + int64_t sourceCore = sourceClass.cpus.front(); + if (sourceCore != fragment.sourceCoreId) + return sourceClass.op->emitError("projected input send source core does not match materialized scalar class") + << " expected=" << sourceCore << " actual=" << fragment.sourceCoreId; + + appendScalarSendAtCurrentInsertionPoint( + state, sourceClass, payload, fragment.channelId, fragment.sourceCoreId, fragment.targetCoreId, loc); + return success(); +} + +LogicalResult emitProjectedInputSendsAtCurrentInsertionPoint(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + ArrayRef fragments, + Location loc) { + if (fragments.empty()) + return success(); + if (sourceClass.isBatch) + return appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint( + state, sourceClass, payload, fragments, loc); + + for (ProjectedInputTransferFragment* fragment : fragments) + if (failed(appendScalarProjectedInputSendAtCurrentInsertionPoint(state, sourceClass, payload, *fragment, loc))) + return failure(); + return success(); +} + +LogicalResult emitProjectedInputSendsNow(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + ArrayRef fragments, + Location loc) { + if (fragments.empty()) + return success(); + if (sourceClass.isBatch) + return appendConditionalBatchProjectedInputSends(state, sourceClass, payload, fragments, loc); + + for (ProjectedInputTransferFragment* fragment : fragments) + if (failed(appendScalarProjectedInputSendNow(state, sourceClass, payload, *fragment, loc))) + return failure(); + return success(); +} + +void enqueuePendingProjectedInputSends(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + ArrayRef fragments, + Location loc) { + if (fragments.empty()) + return; + state.pendingProjectedInputSends.push_back(PendingProjectedInputSend { + .sourceClass = sourceClass.id, + .payload = payload, + .fragments = SmallVector(fragments.begin(), fragments.end()), + .loc = loc, + }); +} + +LogicalResult flushPendingProjectedInputSendsForClass(MaterializerState& state, MaterializedClass& sourceClass) { + for (size_t index = 0; index < state.pendingProjectedInputSends.size();) { + PendingProjectedInputSend& pending = state.pendingProjectedInputSends[index]; + if (pending.sourceClass != sourceClass.id) { + ++index; + continue; + } + + if (failed(emitProjectedInputSendsAtCurrentInsertionPoint( + state, sourceClass, pending.payload, pending.fragments, pending.loc))) + return failure(); + state.pendingProjectedInputSends.erase(state.pendingProjectedInputSends.begin() + index); + } + return success(); +} + +LogicalResult flushProjectedInputPhaseBSends(MaterializerState& state, MaterializedClass& sourceClass) { + auto countIt = state.pendingProjectedHighToLowReceives.find(sourceClass.id); + if (countIt != state.pendingProjectedHighToLowReceives.end() && countIt->second != 0) + return sourceClass.op->emitError("projected input low-to-high phase started before high-to-low receives completed") + << " classId=" << sourceClass.id << " remainingHighToLowReceives=" << countIt->second, + failure(); + + return flushPendingProjectedInputSendsForClass(state, sourceClass); +} + +LogicalResult releaseProjectedInputPhaseBarrierReceives(MaterializerState& state, + MaterializedClass& targetClass, + unsigned receiveCount) { + if (receiveCount == 0) + return success(); + + auto countIt = state.pendingProjectedHighToLowReceives.find(targetClass.id); + if (countIt == state.pendingProjectedHighToLowReceives.end() || countIt->second < receiveCount) + return targetClass.op->emitError("projected input high-to-low receive count underflow") + << " classId=" << targetClass.id << " receiveCount=" << receiveCount, + failure(); + + countIt->second -= receiveCount; + if (countIt->second != 0) + return success(); + + state.pendingProjectedHighToLowReceives.erase(countIt); + state.projectedInputPhaseBarrierClasses.erase(targetClass.id); + return flushProjectedInputPhaseBSends(state, targetClass); +} + +LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, + MaterializedClass& sourceClass, + ArrayRef keys, + Value payload, + Location loc) { + if (keys.empty()) + return success(); + + DenseSet keySet; + for (ProducerKey key : keys) + keySet.insert(key); + + SmallVector readyFragments; + SmallVector deferredFragments; + for (auto& extractEntry : state.projectedInputTransferPlans) { + for (auto& classEntry : extractEntry.second) { + ProjectedInputTransferPlan& plan = classEntry.second; + for (ProjectedInputTransferFragment& fragment : plan.fragments) { + if (fragment.sendEmitted || !keySet.contains(fragment.producer)) + continue; + + if (fragment.sourceCoreId == fragment.targetCoreId) { + fragment.sendEmitted = true; + continue; + } + + if (!isHigherToLowerProjectedInputFragment(fragment) + && state.projectedInputPhaseBarrierClasses.contains(sourceClass.id)) { + deferredFragments.push_back(&fragment); + fragment.sendEmitted = true; + continue; + } + + readyFragments.push_back(&fragment); + fragment.sendEmitted = true; + } + } + } + + if (failed(emitProjectedInputSendsNow(state, sourceClass, payload, readyFragments, loc))) + return failure(); + enqueuePendingProjectedInputSends(state, sourceClass, payload, deferredFragments, loc); + return success(); +} + + +static FailureOr> +collectOrderedBatchProjectedInputFragments(MaterializedClass& targetClass, + const ProjectedInputTransferPlan& plan, + unsigned fragmentsPerLane) { + const size_t expectedFragmentCount = targetClass.cpus.size() * static_cast(fragmentsPerLane); + if (plan.fragments.size() != expectedFragmentCount) { + targetClass.op->emitError("projected input batch transfer did not collect one complete payload per lane") + << " expected=" << expectedFragmentCount << " actual=" << plan.fragments.size(); + return failure(); + } + + SmallVector fragments(expectedFragmentCount, nullptr); + for (const ProjectedInputTransferFragment& fragment : plan.fragments) { + if (fragment.targetLane >= targetClass.cpus.size()) { + targetClass.op->emitError("projected input batch transfer has an out-of-range target lane") + << " targetLane=" << fragment.targetLane << " laneCount=" << targetClass.cpus.size(); + return failure(); + } + if (fragment.ordinal >= expectedFragmentCount) { + targetClass.op->emitError("projected input batch transfer has an out-of-range ordinal") + << " ordinal=" << fragment.ordinal << " fragmentCount=" << expectedFragmentCount; + return failure(); + } + unsigned expectedOrdinal = fragment.targetLane * fragmentsPerLane + (fragment.ordinal % fragmentsPerLane); + if (fragment.ordinal != expectedOrdinal) { + targetClass.op->emitError("projected input batch transfer ordinal does not match target lane order") + << " ordinal=" << fragment.ordinal << " expected=" << expectedOrdinal; + return failure(); + } + if (fragments[fragment.ordinal]) { + targetClass.op->emitError("projected input batch transfer has duplicate fragment ordinal") + << " ordinal=" << fragment.ordinal; + return failure(); + } + fragments[fragment.ordinal] = &fragment; + } + + for (auto [index, fragment] : llvm::enumerate(fragments)) { + if (fragment) + continue; + targetClass.op->emitError("projected input batch transfer is missing fragment ordinal") << " ordinal=" << index; + return failure(); + } + return fragments; +} + +static SmallVector +getBatchProjectedInputPayloadSlot(ArrayRef fragments, + size_t laneCount, + unsigned fragmentsPerLane, + unsigned slot) { + SmallVector slotFragments; + slotFragments.reserve(laneCount); + for (size_t lane = 0; lane < laneCount; ++lane) + slotFragments.push_back(fragments[lane * fragmentsPerLane + slot]); + return slotFragments; +} + +static bool isLocalProjectedInputFragment(const ProjectedInputTransferFragment& fragment) { + return fragment.sourceCoreId == fragment.targetCoreId; +} + +enum class ProjectedInputExchangeDirection { + Local, + LowerToHigher, + HigherToLower +}; + +static ProjectedInputExchangeDirection getProjectedInputExchangeDirectionKind( + const ProjectedInputTransferFragment& fragment) { + if (isLocalProjectedInputFragment(fragment)) + return ProjectedInputExchangeDirection::Local; + return isLowerToHigherProjectedInputFragment(fragment) ? ProjectedInputExchangeDirection::LowerToHigher + : ProjectedInputExchangeDirection::HigherToLower; +} + +static FailureOr getUniformProjectedInputSlotDirection( + MaterializedClass& targetClass, ArrayRef fragments) { + if (fragments.empty()) + return targetClass.op->emitError("projected input batch transfer has an empty slot"), failure(); + + ProjectedInputExchangeDirection direction = getProjectedInputExchangeDirectionKind(*fragments.front()); + for (const ProjectedInputTransferFragment* fragment : fragments.drop_front()) { + if (getProjectedInputExchangeDirectionKind(*fragment) == direction) + continue; + return targetClass.op->emitError( + "projected input batch transfer slot mixes communication directions; split lanes before materialization") + << " fragmentCount=" << fragments.size(), + failure(); + } + return direction; +} + +static FailureOr resolveUniformLocalBatchProjectedInputSlot(MaterializerState& state, + MaterializedClass& targetClass, + ArrayRef fragments) { + std::optional slotValue; + for (const ProjectedInputTransferFragment* fragment : fragments) { + std::optional local = state.availableValues.lookup(state, fragment->producer, targetClass.id); + if (!local) + return targetClass.op->emitError("projected input batch transfer could not resolve local fragment") + << " sourceCore=" << fragment->sourceCoreId << " targetCore=" << fragment->targetCoreId + << " targetLane=" << fragment->targetLane << " ordinal=" << fragment->ordinal, + failure(); + if (slotValue && *slotValue != *local) + return targetClass.op->emitError( + "projected input batch transfer local slot is not represented by one batch SSA value") + << " ordinal=" << fragment->ordinal, + failure(); + slotValue = *local; + } + + if (!slotValue) + return targetClass.op->emitError("projected input batch transfer has an empty local slot"), failure(); + return *slotValue; +} + +static FailureOr materializeRemoteBatchProjectedInputSlotReceive( + MaterializerState& state, + MaterializedClass& targetClass, + Type fragmentType, + ArrayRef fragments, + Value laneArg, + Location loc) { + MessageVector messages; + messages.channelIds.reserve(fragments.size()); + messages.sourceCoreIds.reserve(fragments.size()); + messages.targetCoreIds.reserve(fragments.size()); + for (const ProjectedInputTransferFragment* fragment : fragments) + messages.append(fragment->channelId, fragment->sourceCoreId, fragment->targetCoreId); + + return SpatChannelReceiveOp::create(state.rewriter, + loc, + fragmentType, + createIndexedChannelId(state, targetClass.op, messages, laneArg, loc), + createIndexedSourceCoreId(state, targetClass.op, messages, laneArg, loc), + createIndexedTargetCoreId(state, targetClass.op, messages, laneArg, loc)) + .getOutput(); +} + +static FailureOr materializeBatchProjectedInputPayloadSlot( + MaterializerState& state, + MaterializedClass& targetClass, + const ProjectedInputTransferPlan& plan, + ArrayRef fragments, + Value laneArg, + Location loc) { + FailureOr direction = getUniformProjectedInputSlotDirection(targetClass, fragments); + if (failed(direction)) return failure(); - for (const ProjectedScalarSendGroup& group : plan.projectedSendGroups) - if (failed(appendProjectedScalarSendLoop(state, sourceClass, payload, group.descriptor, group.messages, loc))) + if (*direction == ProjectedInputExchangeDirection::Local) + return resolveUniformLocalBatchProjectedInputSlot(state, targetClass, fragments); + + return materializeRemoteBatchProjectedInputSlotReceive( + state, targetClass, plan.layout.fragmentType, fragments, laneArg, loc); +} + +static LogicalResult materializeBatchProjectedInputPayloadSlotsInOrder( + MaterializerState& state, + MaterializedClass& targetClass, + const ProjectedInputTransferPlan& plan, + ArrayRef fragments, + unsigned fragmentsPerLane, + Value laneArg, + Location loc, + Value& packed) { + SmallVector localSlots; + SmallVector lowerToHigherSlots; + SmallVector higherToLowerSlots; + + for (unsigned slot = 0; slot < fragmentsPerLane; ++slot) { + SmallVector slotFragments = + getBatchProjectedInputPayloadSlot(fragments, targetClass.cpus.size(), fragmentsPerLane, slot); + FailureOr direction = + getUniformProjectedInputSlotDirection(targetClass, slotFragments); + if (failed(direction)) return failure(); + switch (*direction) { + case ProjectedInputExchangeDirection::Local: + localSlots.push_back(slot); + break; + case ProjectedInputExchangeDirection::LowerToHigher: + lowerToHigherSlots.push_back(slot); + break; + case ProjectedInputExchangeDirection::HigherToLower: + higherToLowerSlots.push_back(slot); + break; + } + } + + const auto materializeSlot = [&](unsigned slot) -> LogicalResult { + SmallVector slotFragments = + getBatchProjectedInputPayloadSlot(fragments, targetClass.cpus.size(), fragmentsPerLane, slot); + FailureOr fragmentValue = + materializeBatchProjectedInputPayloadSlot(state, targetClass, plan, slotFragments, laneArg, loc); + if (failed(fragmentValue)) + return failure(); + + Value payloadOrdinal = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(slot)); + FailureOr offset = scaleIndexByDim0SizeInClass( + state, targetClass, payloadOrdinal, plan.layout.fragmentType.getDimSize(0), loc); + if (failed(offset)) + return failure(); + FailureOr updated = createDim0InsertSliceInClass(state, targetClass, loc, *fragmentValue, packed, *offset); + if (failed(updated)) + return failure(); + packed = *updated; + return success(); + }; + + for (unsigned slot : localSlots) + if (failed(materializeSlot(slot))) + return failure(); + unsigned highToLowReceiveCount = 0; + for (unsigned slot : higherToLowerSlots) { + if (failed(materializeSlot(slot))) + return failure(); + highToLowReceiveCount += static_cast(targetClass.cpus.size()); + } + if (failed(releaseProjectedInputPhaseBarrierReceives(state, targetClass, highToLowReceiveCount))) + return failure(); + for (unsigned slot : lowerToHigherSlots) + if (failed(materializeSlot(slot))) + return failure(); + return success(); +} + + +FailureOr materializeBatchProjectedInputTransferPlanReceives(MaterializerState& state, + MaterializedClass& targetClass, + const ProjectedInputTransferPlan& plan, + Location loc, + Operation* insertionAnchor) { + auto batch = dyn_cast(targetClass.op); + if (!batch) + return targetClass.op->emitError("projected input transfer batch target must be a scheduled compute_batch"), + failure(); + std::optional laneArg = batch.getLaneArgument(); + if (!laneArg) + return targetClass.op->emitError("projected input transfer batch target is missing its lane argument"), + failure(); + + const unsigned fragmentsPerLane = plan.layout.fragmentsPerLogicalSlot; + if (fragmentsPerLane == 0 || plan.layout.payloadFragmentCount != fragmentsPerLane) + return targetClass.op->emitError("projected input batch transfer requires one logical payload per target lane"), + failure(); + + FailureOr> maybeFragments = + collectOrderedBatchProjectedInputFragments(targetClass, plan, fragmentsPerLane); + if (failed(maybeFragments)) + return failure(); + ArrayRef fragments = *maybeFragments; + + FailureOr payloadType = + getProjectedPayloadType(targetClass.op, plan.layout.fragmentType, fragmentsPerLane); + if (failed(payloadType)) + return failure(); + + OpBuilder::InsertionGuard guard(state.rewriter); + if (insertionAnchor) { + if (insertionAnchor->getBlock() != targetClass.body) + return targetClass.op->emitError("projected input transfer insertion anchor is outside the target class body"), + failure(); + state.rewriter.setInsertionPoint(insertionAnchor); + } + + Value packed = tensor::EmptyOp::create(state.rewriter, loc, payloadType->getShape(), payloadType->getElementType()).getResult(); + if (failed(materializeBatchProjectedInputPayloadSlotsInOrder( + state, targetClass, plan, fragments, fragmentsPerLane, *laneArg, loc, packed))) + return failure(); + + return packed; +} + +FailureOr materializeProjectedInputTransferPlanReceives(MaterializerState& state, + MaterializedClass& targetClass, + const ProjectedInputTransferPlan& plan, + Location loc, + Operation* insertionAnchor = nullptr) { + if (failed(verifyProjectedFragmentLayout(targetClass.op, plan.layout))) + return failure(); + if (targetClass.isBatch) + return materializeBatchProjectedInputTransferPlanReceives(state, targetClass, plan, loc, insertionAnchor); + + OpBuilder::InsertionGuard guard(state.rewriter); + if (insertionAnchor) { + if (insertionAnchor->getBlock() != targetClass.body) + return targetClass.op->emitError("projected input transfer insertion anchor is outside the target class body"), + failure(); + state.rewriter.setInsertionPoint(insertionAnchor); + } + + SmallVector fragments; + fragments.reserve(plan.fragments.size()); + for (const ProjectedInputTransferFragment& fragment : plan.fragments) + fragments.push_back(&fragment); + llvm::sort(fragments, [](const ProjectedInputTransferFragment* lhs, const ProjectedInputTransferFragment* rhs) { + return lhs->ordinal < rhs->ordinal; + }); + + if (fragments.size() != plan.layout.payloadFragmentCount) + return targetClass.op->emitError("projected input transfer did not collect one complete logical payload") + << " expected=" << plan.layout.payloadFragmentCount << " actual=" << fragments.size(); + for (auto [index, fragment] : llvm::enumerate(fragments)) + if (fragment->ordinal != index) + return targetClass.op->emitError("projected input transfer fragments are not ordinal-contiguous"); + + const auto materializeOneFragment = [&](const ProjectedInputTransferFragment& fragment) -> FailureOr { + if (fragment.sourceCoreId == fragment.targetCoreId) { + std::optional local = state.availableValues.lookup(state, fragment.producer, targetClass.id); + if (!local) + return targetClass.op->emitError("projected input transfer could not resolve local fragment"), failure(); + return *local; + } + + return createScalarPeerReceiveAndFlush(state, + targetClass, + plan.layout.fragmentType, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.channelId), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.sourceCoreId), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.targetCoreId), + makeScalarPeerReceiveKey(fragment.sourceCoreId, + fragment.targetCoreId, + plan.layout.fragmentType, + fragment.channelId), + loc); + }; + + if (fragments.size() == 1) { + const ProjectedInputTransferFragment& fragment = *fragments.front(); + if (isLowerToHigherProjectedInputFragment(fragment)) + if (failed(flushProjectedInputPhaseBSends(state, targetClass))) + return failure(); + FailureOr value = materializeOneFragment(fragment); + if (failed(value)) + return failure(); + if (isHigherToLowerProjectedInputFragment(fragment)) + if (failed(releaseProjectedInputPhaseBarrierReceives(state, targetClass, 1))) + return failure(); + return *value; + } + + FailureOr payloadType = + getProjectedPayloadType(targetClass.op, plan.layout.fragmentType, plan.layout.payloadFragmentCount); + if (failed(payloadType)) + return failure(); + Value packed = tensor::EmptyOp::create(state.rewriter, loc, payloadType->getShape(), payloadType->getElementType()).getResult(); + + SmallVector local; + SmallVector lowerToHigher; + SmallVector higherToLower; + for (auto [index, fragment] : llvm::enumerate(fragments)) { + if (isHigherToLowerProjectedInputFragment(*fragment)) + higherToLower.push_back(index); + else if (isLowerToHigherProjectedInputFragment(*fragment)) + lowerToHigher.push_back(index); + else + local.push_back(index); + } + + const auto materializeAndInsert = [&](size_t index) -> LogicalResult { + const ProjectedInputTransferFragment& fragment = *fragments[index]; + FailureOr received = materializeOneFragment(fragment); + if (failed(received)) + return failure(); + + Value offset = getOrCreateIndexConstant( + state.constantFolder, targetClass.op, static_cast(index) * plan.layout.fragmentType.getDimSize(0)); + FailureOr next = createDim0InsertSliceInClass(state, targetClass, loc, *received, packed, offset); + if (failed(next)) { + targetClass.op->emitError("failed to pack projected input transfer fragment") + << " targetClass=" << targetClass.id << " ordinal=" << fragment.ordinal + << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId + << " fragmentType=" << plan.layout.fragmentType << " payloadType=" << *payloadType; + return failure(); + } + packed = *next; + return success(); + }; + + for (size_t index : local) + if (failed(materializeAndInsert(index))) + return failure(); + + for (size_t index : higherToLower) + if (failed(materializeAndInsert(index))) + return failure(); + + if (failed(releaseProjectedInputPhaseBarrierReceives( + state, targetClass, static_cast(higherToLower.size())))) + return failure(); + + for (size_t index : lowerToHigher) + if (failed(materializeAndInsert(index))) + return failure(); + + return packed; +} + +LogicalResult verifyNoPendingProjectedInputSends(MaterializerState& state) { + if (!state.pendingProjectedHighToLowReceives.empty()) { + auto countIt = state.pendingProjectedHighToLowReceives.begin(); + MaterializedClass& targetClass = state.classes[countIt->first]; + return targetClass.op->emitError("projected input high-to-low receive phase was not fully materialized") + << " classId=" << countIt->first << " remainingReceives=" << countIt->second, + failure(); + } + + if (state.pendingProjectedInputSends.empty()) + return success(); + + PendingProjectedInputSend& pending = state.pendingProjectedInputSends.front(); + MaterializedClass& sourceClass = state.classes[pending.sourceClass]; + InFlightDiagnostic diag = sourceClass.op->emitError("projected input send phase barrier was never released") + << " sourceClass=" << pending.sourceClass + << " fragmentCount=" << pending.fragments.size(); + if (!pending.fragments.empty()) { + ProjectedInputTransferFragment* fragment = pending.fragments.front(); + diag << " firstSourceCore=" << fragment->sourceCoreId << " firstTargetCore=" << fragment->targetCoreId; + } + return failure(); +} + +MessageVector filterMessageVector(const MessageVector& messages, ArrayRef indices) { + MessageVector filtered; + for (size_t index : indices) + filtered.append(messages.channelIds[index], messages.sourceCoreIds[index], messages.targetCoreIds[index]); + return filtered; +} + +SmallVector filterInt64Vector(ArrayRef values, ArrayRef indices) { + SmallVector filtered; + filtered.reserve(indices.size()); + for (size_t index : indices) + filtered.push_back(values[index]); + return filtered; +} + + +LogicalResult verifyNoPendingScalarPeerSends(MaterializerState& state) { + struct PendingSummary { + ClassId sourceClass = 0; + int64_t sourceCore = 0; + int64_t targetCore = 0; + Type payloadType; + size_t ordinaryCount = 0; + size_t projectedCount = 0; + }; + + SmallVector summaries; + auto accumulate = [&](ClassId sourceClass, int64_t sourceCore, int64_t targetCore, Type payloadType, bool projected) { + auto it = llvm::find_if(summaries, [&](const PendingSummary& summary) { + return summary.sourceClass == sourceClass && summary.sourceCore == sourceCore + && summary.targetCore == targetCore && summary.payloadType == payloadType; + }); + if (it == summaries.end()) { + summaries.push_back(PendingSummary { + .sourceClass = sourceClass, + .sourceCore = sourceCore, + .targetCore = targetCore, + .payloadType = payloadType, + }); + it = std::prev(summaries.end()); + } + if (projected) + ++it->projectedCount; + else + ++it->ordinaryCount; + }; + + for (const PendingScalarSend& pending : state.pendingScalarSends) + accumulate(pending.sourceClass, pending.sourceCore, pending.targetCore, pending.payloadType, /*projected=*/false); + for (const PendingProjectedScalarSend& pending : state.pendingProjectedScalarSends) + accumulate(pending.sourceClass, pending.sourceCore, pending.targetCore, pending.payloadType, /*projected=*/true); + + if (summaries.empty()) + return success(); + + auto diag = state.func.emitError("scalar peer communication materialization left pending sends"); + for (const PendingSummary& summary : summaries) + diag << " [sourceClass=" << summary.sourceClass << " sourceCore=" << summary.sourceCore + << " targetCore=" << summary.targetCore << " payloadType=" << summary.payloadType + << " ordinaryPending=" << summary.ordinaryCount << " projectedPending=" << summary.projectedCount << "]"; + return failure(); +} + +LogicalResult splitScalarPeerFanoutForImmediateEmission(MaterializerState& state, + MaterializedClass& sourceClass, + Value payload, + const ScalarSourceFanoutPlan& fanoutPlan, + Location loc, + ScalarSourceFanoutPlan& immediatePlan) { + int64_t sourceCpu = static_cast(sourceClass.cpus.front()); + + immediatePlan.receivePlans = fanoutPlan.receivePlans; + if (fanoutPlan.ordinaryMessages) { + for (size_t index = 0; index < fanoutPlan.ordinaryMessages->size(); ++index) { + MessageVector message; + message.append(fanoutPlan.ordinaryMessages->channelIds[index], + fanoutPlan.ordinaryMessages->sourceCoreIds[index], + fanoutPlan.ordinaryMessages->targetCoreIds[index]); + int64_t targetCore = fanoutPlan.ordinaryMessages->targetCoreIds[index]; + ScalarPeerEdgeKey requiredReceive = makeRequiredPeerReceiveForLowerToHigherSend(sourceCpu, targetCore, payload.getType()); + bool deferSend = sourceCpu < targetCore && hasPlannedScalarPeerReceive(state, requiredReceive) + && !hasMaterializedScalarPeerReceive(state, requiredReceive); + if (deferSend) { + if (failed(enqueuePendingScalarSend( + state, sourceClass.id, sourceCpu, targetCore, payload.getType(), payload, message, loc))) + return failure(); + continue; + } + if (!immediatePlan.ordinaryMessages) + immediatePlan.ordinaryMessages = MessageVector {}; + immediatePlan.ordinaryMessages->append(message.channelIds, message.sourceCoreIds, message.targetCoreIds); + } + } + + for (const ProjectedScalarSendGroup& group : fanoutPlan.projectedSendGroups) { + assert(succeeded(verifyProjectedSendDescriptor(sourceClass.op, group.descriptor, group.messages)) + && "projected scalar send group metadata is inconsistent"); + FailureOr projectedCommunicationType = getProjectedCommunicationType(sourceClass.op, group.descriptor); + if (failed(projectedCommunicationType)) + return failure(); + + std::optional immediateGroup; + SmallVector deferredGroups; + for (size_t index = 0; index < group.messages.size(); ++index) { + int64_t targetCore = group.messages.targetCoreIds[index]; + ScalarPeerEdgeKey requiredReceive = + makeRequiredPeerReceiveForLowerToHigherSend(sourceCpu, targetCore, *projectedCommunicationType); + bool deferSend = sourceCpu < targetCore && hasPlannedScalarPeerReceive(state, requiredReceive) + && !hasMaterializedScalarPeerReceive(state, requiredReceive); + if (!deferSend) { + if (!immediateGroup) + immediateGroup = makeEmptyProjectedScalarSendGroup(group); + appendProjectedScalarSendMessage(*immediateGroup, group, index); + continue; + } + + auto deferredIt = llvm::find_if(deferredGroups, [&](const ProjectedScalarSendGroup& deferredGroup) { + return !deferredGroup.messages.empty() && deferredGroup.messages.targetCoreIds.front() == targetCore; + }); + if (deferredIt == deferredGroups.end()) { + deferredGroups.push_back(makeEmptyProjectedScalarSendGroup(group)); + deferredIt = std::prev(deferredGroups.end()); + } + appendProjectedScalarSendMessage(*deferredIt, group, index); + } + + if (immediateGroup) { + if (failed(finalizeSplitProjectedScalarSendGroup(sourceClass.op, *immediateGroup))) + return failure(); + immediatePlan.projectedSendGroups.push_back(std::move(*immediateGroup)); + } + + for (ProjectedScalarSendGroup& deferredGroup : deferredGroups) { + if (failed(finalizeSplitProjectedScalarSendGroup(sourceClass.op, deferredGroup))) + return failure(); + int64_t targetCore = deferredGroup.messages.targetCoreIds.front(); + if (failed(enqueuePendingProjectedScalarSend(state, + sourceClass.id, + sourceCpu, + targetCore, + *projectedCommunicationType, + payload, + deferredGroup, + loc))) + return failure(); + } + } + return success(); } @@ -3978,22 +5906,30 @@ LogicalResult emitScalarSourceCommunication( auto fanoutPlan = buildScalarSourceFanoutPlan(state, sourceClass, keys, destinationClasses, payload); if (failed(fanoutPlan)) return failure(); - if (failed(emitScalarSourceFanoutSends(state, sourceClass, payload, *fanoutPlan, loc))) + ScalarSourceFanoutPlan immediatePlan; + if (failed(splitScalarPeerFanoutForImmediateEmission(state, sourceClass, payload, *fanoutPlan, loc, immediatePlan))) + return failure(); + + if (failed(emitScalarSourceFanoutSends(state, sourceClass, payload, immediatePlan, loc))) return failure(); for (const ScalarSourceReceivePlan& plan : fanoutPlan->receivePlans) { MaterializedClass& targetClass = state.classes[plan.targetClass]; - Value received = appendReceive(state, targetClass, plan.receiveType, plan.messages, loc); + FailureOr received = appendReceiveAndFlushPendingScalarSends(state, targetClass, plan.receiveType, plan.messages, loc); + if (failed(received)) + return failure(); if (plan.projectedExtractOp) { - state.projectedExtractReplacements[plan.projectedExtractOp][plan.targetClass] = - ProjectedExtractReplacement {received, plan.projectedLayout}; + state.projectedExtractReplacements[plan.projectedExtractOp][plan.targetClass] = ProjectedExtractReplacement { + *received, + plan.projectedLayout + }; continue; } for (ProducerKey key : keys) - state.availableValues.record(key, targetClass.id, received); + state.availableValues.record(key, targetClass.id, *received); } return success(); @@ -4048,12 +5984,17 @@ LogicalResult emitClassToClassCommunication(MaterializerState& state, if (sourceClass.cpus.size() != targetClass.cpus.size()) return sourceClass.op->emitError( - "cannot materialize batch communication between equivalence classes of different sizes"); + "cannot materialize batch communication between equivalence classes of different sizes") + << " sourceClass=" << sourceClass.id << " sourceSize=" << sourceClass.cpus.size() + << " targetClass=" << targetClass.id << " targetSize=" << targetClass.cpus.size() + << " keyCount=" << keys.size(); MessageVector messages; messages.channelIds.reserve(sourceClass.cpus.size()); messages.sourceCoreIds.reserve(sourceClass.cpus.size()); messages.targetCoreIds.reserve(targetClass.cpus.size()); + SmallVector lowerToHigherOrLocal; + SmallVector higherToLower; for (auto [lane, sourceCpu] : llvm::enumerate(sourceClass.cpus)) { auto checkedSourceCpu = getCheckedCoreId(sourceClass.op, sourceCpu, "batch source core id"); @@ -4063,14 +6004,38 @@ LogicalResult emitClassToClassCommunication(MaterializerState& state, if (failed(checkedTargetCpu)) return failure(); messages.append(state.nextChannelId++, *checkedSourceCpu, *checkedTargetCpu); + if (*checkedSourceCpu <= *checkedTargetCpu) + lowerToHigherOrLocal.push_back(lane); + else + higherToLower.push_back(lane); } - if (failed(appendSend(state, sourceClass, payload, messages, loc))) - return failure(); - Value received = appendReceive(state, targetClass, payload.getType(), messages, loc); + auto recordReceived = [&](ArrayRef indices, Value received) { + for (size_t index : indices) + state.availableValues.record(keys[index], targetClass.id, received); + }; - for (ProducerKey key : keys) - state.availableValues.record(key, targetClass.id, received); + if (!lowerToHigherOrLocal.empty()) { + MessageVector sendFirst = filterMessageVector(messages, lowerToHigherOrLocal); + if (failed(appendSend(state, sourceClass, payload, sendFirst, loc))) + return failure(); + FailureOr received = + appendReceiveAndFlushPendingScalarSends(state, targetClass, payload.getType(), sendFirst, loc); + if (failed(received)) + return failure(); + recordReceived(lowerToHigherOrLocal, *received); + } + + if (!higherToLower.empty()) { + MessageVector receiveFirst = filterMessageVector(messages, higherToLower); + FailureOr received = + appendReceiveAndFlushPendingScalarSends(state, targetClass, payload.getType(), receiveFirst, loc); + if (failed(received)) + return failure(); + if (failed(appendSend(state, sourceClass, payload, receiveFirst, loc))) + return failure(); + recordReceived(higherToLower, *received); + } return success(); } @@ -4162,8 +6127,11 @@ emitHostCommunication(MaterializerState& state, MaterializedClass& sourceClass, if (failed(appendSend(state, sourceClass, payload, messages, payload.getLoc()))) return failure(); - Value ownerPayload = appendReceive(state, ownerClass, payload.getType(), messages, payload.getLoc()); - return setHostOutputValue(state, ownerClass, originalOutput, ownerPayload); + FailureOr ownerPayload = + appendReceiveAndFlushPendingScalarSends(state, ownerClass, payload.getType(), messages, payload.getLoc()); + if (failed(ownerPayload)) + return failure(); + return setHostOutputValue(state, ownerClass, originalOutput, *ownerPayload); } if (sourceClass.isBatch) @@ -4183,8 +6151,11 @@ emitHostCommunication(MaterializerState& state, MaterializedClass& sourceClass, if (failed(appendSend(state, sourceClass, payload, messages, payload.getLoc()))) return failure(); - Value ownerPayload = appendReceive(state, ownerClass, payload.getType(), messages, payload.getLoc()); - return setHostOutputValue(state, ownerClass, originalOutput, ownerPayload); + FailureOr ownerPayload = + appendReceiveAndFlushPendingScalarSends(state, ownerClass, payload.getType(), messages, payload.getLoc()); + if (failed(ownerPayload)) + return failure(); + return setHostOutputValue(state, ownerClass, originalOutput, *ownerPayload); } LogicalResult emitOutputFanout(MaterializerState& state, @@ -4199,6 +6170,8 @@ LogicalResult emitOutputFanout(MaterializerState& state, if (!sourceClass.isBatch) { if (failed(emitScalarSourceCommunication(state, sourceClass, keys, payload, loc))) return failure(); + if (failed(emitPlannedProjectedInputSendsForKeys(state, sourceClass, keys, payload, loc))) + return failure(); return emitHostCommunication(state, sourceClass, payload, originalOutput); } @@ -4212,13 +6185,20 @@ LogicalResult emitOutputFanout(MaterializerState& state, recordedProjectedHostFragments = *recorded; } - if (!haveSameDestinationClasses(state, keys)) - return sourceClass.op->emitError( - "cannot materialize batched output whose lanes have different destination equivalence classes"); + if (failed(emitPlannedProjectedInputSendsForKeys(state, sourceClass, keys, payload, loc))) + return failure(); - for (ClassId destinationClass : getDestinationClasses(state, keys.front())) - if (failed(emitClassToClassCommunication(state, sourceClass, state.classes[destinationClass], keys, payload, loc))) + for (ClassId destinationClass : collectDestinationClassesForKeys(state, keys)) { + SmallVector destinationKeys = filterProducerKeysForDestination(state, keys, destinationClass); + if (destinationKeys.empty()) + continue; + if (sourceClass.isBatch && destinationKeys.size() != keys.size()) + return sourceClass.op->emitError( + "non-uniform batch fanout requires an explicit projected input transfer plan"); + if (failed(emitClassToClassCommunication( + state, sourceClass, state.classes[destinationClass], destinationKeys, payload, loc))) return failure(); + } if (hasLiveExternalUseCached(state, originalOutput) && !recordedProjectedHostFragments) return sourceClass.op->emitError("batch host publication requires explicit fragment assembly metadata"); @@ -4234,6 +6214,33 @@ LogicalResult emitOutputFanout(MaterializerState& state, return success(); } +LogicalResult collectPlannedScalarPeerReceives(MaterializerState& state) { + for (const ComputeInstance& instance : state.schedule.dominanceOrderCompute) { + auto cpuIt = state.schedule.computeToCpuMap.find(instance); + if (cpuIt == state.schedule.computeToCpuMap.end()) + return instance.op->emitError("missing CPU assignment while precomputing scalar peer receives"); + + auto classIt = state.cpuToClass.find(cpuIt->second); + if (classIt == state.cpuToClass.end()) + return instance.op->emitError("missing materialized class while precomputing scalar peer receives"); + + MaterializedClass& sourceClass = state.classes[classIt->second]; + if (sourceClass.isBatch) + continue; + + ArrayRef outputs = getComputeInstanceOutputValuesCached(state, instance); + for (Value output : outputs) { + SmallVector keys = collectProducerKeysForDestinations(output, instance); + if (keys.empty()) + continue; + if (failed(collectPlannedScalarPeerReceivesForSource(state, sourceClass, keys, output))) + return failure(); + } + } + + return success(); +} + struct DirectWholeBatchFragment { ProducerKey key; Value fragment; @@ -4299,6 +6306,23 @@ struct WholeBatchAssemblyPlan { SmallVector directFragments; }; +ProjectedWholeBatchFragmentGroup filterProjectedDeferredReceiveGroup(const ProjectedWholeBatchFragmentGroup& group, + ArrayRef indices) { + ProjectedWholeBatchFragmentGroup filtered; + filtered.kind = group.kind; + filtered.fragmentType = group.fragmentType; + filtered.messages = filterMessageVector(group.messages, indices); + filtered.offsetsByDim.resize(group.offsetsByDim.size()); + filtered.sizesByDim.resize(group.sizesByDim.size()); + filtered.stridesByDim.resize(group.stridesByDim.size()); + for (size_t dim = 0; dim < group.offsetsByDim.size(); ++dim) { + filtered.offsetsByDim[dim] = filterInt64Vector(group.offsetsByDim[dim], indices); + filtered.sizesByDim[dim] = filterInt64Vector(group.sizesByDim[dim], indices); + filtered.stridesByDim[dim] = filterInt64Vector(group.stridesByDim[dim], indices); + } + return filtered; +} + bool wholeBatchLaneCovered(const WholeBatchAssemblyPlan& plan, uint32_t lane) { return lane < plan.coveredLanes.size() && plan.coveredLanes[lane] != 0; } @@ -4884,23 +6908,56 @@ FailureOr emitWholeBatchFragmentGroup(MaterializerState& state, Location loc) { switch (group.kind) { case WholeBatchFragmentSourceKind::DeferredReceive: { - FailureOr updated = emitIndexedFragmentInsertLoop( - state, - targetClass, - destination, - static_cast(group.outputOffsets.size()), - [&](Value flatIndex) -> FailureOr { - Value channelId = createIndexedChannelId(state, targetClass.op, group.messages, flatIndex, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, group.messages, flatIndex, loc); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, group.messages, flatIndex, loc); - return SpatChannelReceiveOp::create( - state.rewriter, loc, group.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); - }, - [&](Value flatIndex) -> FailureOr { - return createIndexedIndexValue(state, targetClass.op, group.outputOffsets, flatIndex, loc); - }, - loc); + Value updatedDestination = destination; + ReceiveMessagePartition partition = + partitionReceivesByPendingSendUnlocks(state, targetClass, group.fragmentType, group.messages); + for (size_t index : partition.criticalStaticIndices) { + state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); + FailureOr received = createScalarPeerReceiveAndFlush( + state, + targetClass, + group.fragmentType, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.channelIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.sourceCoreIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.targetCoreIds[index]), + makeScalarPeerReceiveKey( + group.messages.sourceCoreIds[index], group.messages.targetCoreIds[index], group.fragmentType, group.messages.channelIds[index]), + loc); + if (failed(received)) + return failure(); + FailureOr updated = createDim0InsertSliceInClass(state, + targetClass, + loc, + *received, + updatedDestination, + getOrCreateIndexConstant( + state.constantFolder, targetClass.op, group.outputOffsets[index])); + if (failed(updated)) + return failure(); + updatedDestination = *updated; + } + + auto remainingMessages = filterMessageVector(group.messages, partition.remainingIndices); + auto remainingOffsets = filterInt64Vector(group.outputOffsets, partition.remainingIndices); + FailureOr updated = remainingMessages.empty() + ? FailureOr(updatedDestination) + : emitIndexedFragmentInsertLoop( + state, + targetClass, + updatedDestination, + static_cast(remainingOffsets.size()), + [&](Value flatIndex) -> FailureOr { + Value channelId = createIndexedChannelId(state, targetClass.op, remainingMessages, flatIndex, loc); + Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); + Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); + return SpatChannelReceiveOp::create( + state.rewriter, loc, group.fragmentType, channelId, sourceCoreId, targetCoreId) + .getOutput(); + }, + [&](Value flatIndex) -> FailureOr { + return createIndexedIndexValue(state, targetClass.op, remainingOffsets, flatIndex, loc); + }, + loc); if (failed(updated)) return failure(); @@ -5320,20 +7377,60 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( FailureOr updated = failure(); switch (group.kind) { case ProjectedWholeBatchFragmentSourceKind::DeferredReceive: - updated = emitProjectedWholeBatchFragmentInsertLoop( - state, - targetClass, - result, - group, - [&](Value flatIndex) -> FailureOr { - Value channelId = createIndexedChannelId(state, targetClass.op, group.messages, flatIndex, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, group.messages, flatIndex, loc); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, group.messages, flatIndex, loc); - return SpatChannelReceiveOp::create( - state.rewriter, loc, group.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); - }, - loc); + updated = result; + for (size_t index : partitionReceivesByPendingSendUnlocks(state, targetClass, group.fragmentType, group.messages) + .criticalStaticIndices) { + state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); + FailureOr received = createScalarPeerReceiveAndFlush( + state, + targetClass, + group.fragmentType, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.channelIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.sourceCoreIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.targetCoreIds[index]), + makeScalarPeerReceiveKey( + group.messages.sourceCoreIds[index], group.messages.targetCoreIds[index], group.fragmentType, group.messages.channelIds[index]), + loc); + if (failed(received)) + return failure(); + + SmallVector offsets; + SmallVector sizes; + SmallVector strides; + offsets.reserve(group.offsetsByDim.size()); + sizes.reserve(group.sizesByDim.size()); + strides.reserve(group.stridesByDim.size()); + for (size_t dim = 0; dim < group.offsetsByDim.size(); ++dim) { + offsets.push_back(state.rewriter.getIndexAttr(group.offsetsByDim[dim][index])); + sizes.push_back(state.rewriter.getIndexAttr(group.sizesByDim[dim][index])); + strides.push_back(state.rewriter.getIndexAttr(group.stridesByDim[dim][index])); + } + state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); + updated = tensor::InsertSliceOp::create(state.rewriter, loc, *received, *updated, offsets, sizes, strides) + .getResult(); + } + { + ReceiveMessagePartition partition = + partitionReceivesByPendingSendUnlocks(state, targetClass, group.fragmentType, group.messages); + if (!partition.remainingIndices.empty()) { + ProjectedWholeBatchFragmentGroup remainingGroup = + filterProjectedDeferredReceiveGroup(group, partition.remainingIndices); + updated = emitProjectedWholeBatchFragmentInsertLoop( + state, + targetClass, + *updated, + remainingGroup, + [&](Value flatIndex) -> FailureOr { + Value channelId = createIndexedChannelId(state, targetClass.op, remainingGroup.messages, flatIndex, loc); + Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingGroup.messages, flatIndex, loc); + Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingGroup.messages, flatIndex, loc); + return SpatChannelReceiveOp::create( + state.rewriter, loc, remainingGroup.fragmentType, channelId, sourceCoreId, targetCoreId) + .getOutput(); + }, + loc); + } + } break; case ProjectedWholeBatchFragmentSourceKind::PackedValue: updated = emitProjectedWholeBatchFragmentInsertLoop( @@ -5681,8 +7778,11 @@ FailureOr resolveInputValue(MaterializerState& state, else { size_t laneCount = targetClass.cpus.size(); MessageVector messages = indexedRun->messages.slice(slotIndex * laneCount, laneCount); - received = - appendReceive(state, targetClass, indexedRun->fragmentType, messages, consumerInstance.op->getLoc()); + FailureOr appended = appendReceiveAndFlushPendingScalarSends( + state, targetClass, indexedRun->fragmentType, messages, consumerInstance.op->getLoc()); + if (failed(appended)) + return failure(); + received = *appended; } state.availableValues.record(*producer, targetClass.id, received); return rejectNonLocalResolvedValue(received); @@ -5739,6 +7839,15 @@ bool hasPlannedProjectedInputTransfer(MaterializerState& state, ComputeInstance logicalConsumer, ClassId classId) { ProjectedBatchInputKey inputKey {batch.getOperation(), inputIndex}; + if (std::optional match = getProjectedInputSliceMatch(state, batch, inputIndex)) { + auto planIt = state.projectedInputTransferPlans.find(match->extract.getOperation()); + if (planIt != state.projectedInputTransferPlans.end()) { + auto classIt = planIt->second.find(classId); + if (classIt != planIt->second.end() && classIt->second.inputKey == inputKey) + return true; + } + } + SmallVector producers = collectProducerKeysForDestinations(input, logicalConsumer); for (ProducerKey producer : producers) { auto producerIt = state.projectedTransfers.find(producer); @@ -5898,25 +8007,7 @@ LogicalResult mapInputs(MaterializerState& state, return batch.emitOpError("expected compute_batch input block argument while materializing inputs"); std::optional producer = getInputRequestProducerKey(input, instance); if (demand->kind == BatchInputDemandKind::ProjectedFragment) { - std::optional match = - getProjectedInputSliceMatch(state, batch, static_cast(index)); - auto mappedType = dyn_cast((*mapped).getType()); - if (!match || !mappedType || mappedType != match->fragmentType) { - InFlightDiagnostic diagnostic = batch.emitOpError( - "projected compute_batch input resolved to an incompatible fragment") - << " #" << index << " resolvedType=" << (*mapped).getType(); - if (match) - diagnostic << " expectedFragmentType=" << match->fragmentType; - if (producer) { - diagnostic << " from '" << producer->instance.op->getName() - << "' laneStart=" << producer->instance.laneStart - << " laneCount=" << producer->instance.laneCount - << " resultIndex=" << producer->resultIndex; - } - return failure(); - } - - mapper.map(*inputArg, *mapped); + (void)mapped; continue; } FailureOr remapped = mapResolvedInput(*mapped, input, batch, producer); @@ -6018,16 +8109,84 @@ std::optional lookupProjectedExtractReplacement(Mat return classIt->second; } +FailureOr> getOrMaterializeProjectedExtractReplacement( + MaterializerState& state, + MaterializedClass& targetClass, + tensor::ExtractSliceOp extract, + Operation* insertionAnchor = nullptr) { + if (std::optional replacement = + lookupProjectedExtractReplacement(state, targetClass, extract)) + return replacement; + + auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); + if (planIt == state.projectedInputTransferPlans.end()) + return std::optional {}; + + auto classIt = planIt->second.find(targetClass.id); + if (classIt == planIt->second.end()) + return std::optional {}; + + FailureOr payload = materializeProjectedInputTransferPlanReceives( + state, targetClass, classIt->second, extract.getLoc(), insertionAnchor); + if (failed(payload)) { + targetClass.op->emitError("failed to materialize projected input transfer replacement") + << " targetClass=" << targetClass.id << " fragmentCount=" << classIt->second.fragments.size() + << " payloadFragmentCount=" << classIt->second.layout.payloadFragmentCount; + return failure(); + } + + ProjectedExtractReplacement replacement {*payload, classIt->second.layout}; + state.projectedExtractReplacements[extract.getOperation()][targetClass.id] = replacement; + return std::optional {replacement}; +} + +Operation* getTopLevelMaterializedClassBodyOp(Operation* nestedOp, MaterializedClass& targetClass) { + if (!nestedOp) + return nullptr; + + Operation* current = nestedOp; + while (Operation* parent = current->getParentOp()) { + if (parent == targetClass.op) + return current; + current = parent; + } + + return nullptr; +} + +void moveNewTopLevelOpsBeforeAnchor(MaterializedClass& targetClass, Operation* anchor) { + if (!anchor || anchor->getBlock() != targetClass.body) + return; + + SmallVector insertedOps; + auto it = std::next(anchor->getIterator()); + auto end = targetClass.body->getTerminator()->getIterator(); + for (; it != end; ++it) + insertedOps.push_back(&*it); + + for (Operation* op : insertedOps) + op->moveBefore(anchor); +} + FailureOr materializeProjectedWholeBatchExtractReplacement(MaterializerState& state, MaterializedClass& targetClass, tensor::ExtractSliceOp extract, ProducerKey producer, IRMapping* mapper) { + OpBuilder::InsertPoint replacementPoint = state.rewriter.saveInsertionPoint(); + Operation* topLevelAnchor = getTopLevelMaterializedClassBodyOp(extract.getOperation(), targetClass); + FailureOr fullSource = materializeWholeBatchInput(state, targetClass, producer, extract.getSource().getType(), extract.getLoc()); if (failed(fullSource)) return failure(); + moveNewTopLevelOpsBeforeAnchor(targetClass, topLevelAnchor); + if (topLevelAnchor) + state.rewriter.setInsertionPoint(extract); + else + state.rewriter.restoreInsertionPoint(replacementPoint); + auto remapFoldResult = [&](OpFoldResult value) -> FailureOr { if (auto mappedValue = dyn_cast_if_present(value)) { FailureOr localized = rematerializeIndexValueInClass( @@ -6077,12 +8236,17 @@ LogicalResult applyProjectedExtractReplacementsInClonedOp(MaterializerState& sta CloneIndexingContext indexing, IRMapping& mapper) { if (auto originalExtract = dyn_cast(&originalOp)) { - if (std::optional replacement = - lookupProjectedExtractReplacement(state, targetClass, originalExtract)) { - auto clonedExtract = dyn_cast(&clonedOp); + auto clonedExtract = dyn_cast(&clonedOp); + Operation* insertionAnchor = clonedExtract ? getTopLevelMaterializedClassBodyOp(clonedExtract, targetClass) : nullptr; + FailureOr> maybeReplacement = + getOrMaterializeProjectedExtractReplacement(state, targetClass, originalExtract, insertionAnchor); + if (failed(maybeReplacement)) + return failure(); + if (std::optional replacement = *maybeReplacement) { if (!clonedExtract) return targetClass.op->emitError("projected replacement lost extract structure during cloning"); + OpBuilder::InsertionGuard guard(state.rewriter); state.rewriter.setInsertionPoint(clonedExtract); FailureOr projected = materializeProjectedExtractReplacement( state, targetClass, clonedExtract, *replacement, indexing.projectionSlotIndex, &mapper); @@ -6099,6 +8263,7 @@ LogicalResult applyProjectedExtractReplacementsInClonedOp(MaterializerState& sta if (!clonedExtract) return targetClass.op->emitError("projected whole-batch replacement lost extract structure during cloning"); + OpBuilder::InsertionGuard guard(state.rewriter); state.rewriter.setInsertionPoint(clonedExtract); FailureOr projected = materializeProjectedWholeBatchExtractReplacement(state, targetClass, clonedExtract, *producer, &mapper); @@ -6243,8 +8408,11 @@ LogicalResult cloneComputeTemplateBody(MaterializerState& state, Block& sourceBlock = getComputeInstanceTemplateBlock(instance); for (Operation& op : sourceBlock.without_terminator()) { if (auto extract = dyn_cast(&op)) { - if (std::optional replacement = - lookupProjectedExtractReplacement(state, targetClass, extract)) { + FailureOr> maybeReplacement = + getOrMaterializeProjectedExtractReplacement(state, targetClass, extract); + if (failed(maybeReplacement)) + return failure(); + if (std::optional replacement = *maybeReplacement) { FailureOr projected = materializeProjectedExtractReplacement( state, targetClass, extract, *replacement, indexing.projectionSlotIndex, &mapper); if (failed(projected)) @@ -6330,6 +8498,7 @@ LogicalResult cloneComputeTemplateBody(MaterializerState& state, return failure(); for (auto [oldResult, newResult] : llvm::zip(op.getResults(), cloned->getResults())) mapper.map(oldResult, newResult); + state.rewriter.setInsertionPointAfter(cloned); } return success(); @@ -6597,12 +8766,20 @@ FailureOr> cloneInstanceBody(MaterializerState& state, OpBuilder::InsertPoint cloneInsertionPoint = state.rewriter.saveInsertionPoint(); mapWeights(state, targetClass, instance, mapper); - if (failed(mapInputs(state, targetClass, instance, mapper, indexing))) + if (failed(mapInputs(state, targetClass, instance, mapper, indexing))) { + sourceOp->emitError("failed to map inputs while cloning compute/compute_batch body") + << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount + << " targetClass=" << targetClass.id; return failure(); + } state.rewriter.restoreInsertionPoint(cloneInsertionPoint); - if (failed(cloneComputeTemplateBody(state, targetClass, instance, mapper, indexing))) + if (failed(cloneComputeTemplateBody(state, targetClass, instance, mapper, indexing))) { + sourceOp->emitError("failed to clone compute/compute_batch body") + << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount + << " targetClass=" << targetClass.id; return failure(); + } if (auto compute = dyn_cast(sourceOp)) { Block& sourceBlock = getComputeInstanceTemplateBlock(instance); @@ -6857,21 +9034,29 @@ LogicalResult emitPackedRunFanout(MaterializerState& state, auto fanoutPlan = buildScalarSourceFanoutPlan(state, sourceClass, keys, destinationClasses, packed); if (failed(fanoutPlan)) return failure(); - if (failed(emitScalarSourceFanoutSends(state, sourceClass, packed, *fanoutPlan, loc))) + ScalarSourceFanoutPlan immediatePlan; + if (failed(splitScalarPeerFanoutForImmediateEmission(state, sourceClass, packed, *fanoutPlan, loc, immediatePlan))) + return failure(); + + if (failed(emitScalarSourceFanoutSends(state, sourceClass, packed, immediatePlan, loc))) return failure(); for (const ScalarSourceReceivePlan& plan : fanoutPlan->receivePlans) { MaterializedClass& targetClass = state.classes[plan.targetClass]; - Value received = appendReceive(state, targetClass, plan.receiveType, plan.messages, loc); + FailureOr received = appendReceiveAndFlushPendingScalarSends(state, targetClass, plan.receiveType, plan.messages, loc); + if (failed(received)) + return failure(); if (plan.projectedExtractOp) { - state.projectedExtractReplacements[plan.projectedExtractOp][plan.targetClass] = - ProjectedExtractReplacement {received, plan.projectedLayout}; + state.projectedExtractReplacements[plan.projectedExtractOp][plan.targetClass] = ProjectedExtractReplacement { + *received, + plan.projectedLayout + }; continue; } - if (failed(registerPackedRunValue(state, targetClass, keys, received, fragmentType, loc))) + if (failed(registerPackedRunValue(state, targetClass, keys, *received, fragmentType, loc))) return failure(); } @@ -7301,8 +9486,17 @@ LogicalResult materializeScalarBatchRun(MaterializerState& state, } FailureOr> packedOutputs = materializeBatchOutputGroupLoop(state, targetClass, run, group); - if (failed(packedOutputs)) + if (failed(packedOutputs)) { + InFlightDiagnostic diagnostic = sourceBatch.emitOpError("failed to materialize scalar batch output group") + << " targetClass=" << targetClass.id << " runSize=" << run.size(); + if (!run.empty() && !run.front().peers.empty()) + diagnostic << " firstLaneStart=" << run.front().peers.front().laneStart + << " firstLaneCount=" << run.front().peers.front().laneCount; + diagnostic << " resultIndices="; + for (size_t resultIndex : group.resultIndices) + diagnostic << " " << resultIndex; return failure(); + } for (auto [groupOutputIndex, resultIndex] : llvm::enumerate(group.resultIndices)) { Value packed = (*packedOutputs)[groupOutputIndex]; @@ -7324,18 +9518,32 @@ LogicalResult materializeScalarBatchRun(MaterializerState& state, rankedFragmentType, representativeOriginalOutput, loc); - if (failed(recordedProjectedHostFragments)) + if (failed(recordedProjectedHostFragments)) { + sourceBatch.emitOpError("failed to record projected scalar host fragments") + << " targetClass=" << targetClass.id << " resultIndex=" << resultIndex; return failure(); + } if (run.size() == 1) { if (*recordedProjectedHostFragments) { - if (failed(emitScalarSourceCommunication(state, targetClass, keys, packed, loc))) + if (failed(emitScalarSourceCommunication(state, targetClass, keys, packed, loc))) { + sourceBatch.emitOpError("failed to emit scalar source communication for projected host fragment") + << " targetClass=" << targetClass.id << " resultIndex=" << resultIndex; return failure(); + } + if (failed(emitPlannedProjectedInputSendsForKeys(state, targetClass, keys, packed, loc))) { + sourceBatch.emitOpError("failed to emit projected input sends for scalar projected host fragment") + << " targetClass=" << targetClass.id << " resultIndex=" << resultIndex; + return failure(); + } continue; } - if (failed(emitOutputFanout(state, targetClass, keys, packed, representativeOriginalOutput, loc))) + if (failed(emitOutputFanout(state, targetClass, keys, packed, representativeOriginalOutput, loc))) { + sourceBatch.emitOpError("failed to emit singleton scalar batch output fanout") + << " targetClass=" << targetClass.id << " resultIndex=" << resultIndex; return failure(); + } continue; } @@ -7679,6 +9887,7 @@ LogicalResult emitPackedBatchRunSends(MaterializerState& state, return failure(); if (sendPlans.empty()) return success(); + sortBatchRunSendPlansByPeerExchangeOrder(sendPlans); DenseMap packedOutputIndexByResult; for (auto [packedIndex, output] : llvm::enumerate(plan.outputs)) @@ -7854,238 +10063,6 @@ LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeIns return success(); } -std::optional getStaticCommunicationCore(Value value) { - if (std::optional coreId = mlir::getConstantIntValue(value)) - return coreId; - return std::nullopt; -} - -bool isTopLevelCommunicationOp(Operation* op) { - return isa(op); -} - -bool valueMayEvaluateToCore(Value value, int64_t coreId) { - if (std::optional constant = getConstantIndexValue(value)) - return *constant == coreId; - - auto affineApply = value.getDefiningOp(); - if (!affineApply) - return false; - - AffineMap map = affineApply.getAffineMap(); - if (map.getNumResults() != 1 || map.getNumDims() != 1 || map.getNumSymbols() != 0 - || affineApply.getMapOperands().size() != 1) - return false; - - auto iv = dyn_cast(affineApply.getMapOperands().front()); - if (!iv) - return false; - - auto loop = dyn_cast_or_null(iv.getOwner()->getParentOp()); - if (!loop || loop.getInductionVar() != iv) - return false; - - std::optional lower = getConstantIndexValue(loop.getLowerBound()); - std::optional upper = getConstantIndexValue(loop.getUpperBound()); - std::optional step = getConstantIndexValue(loop.getStep()); - if (!lower || !upper || !step || *step <= 0) - return false; - - for (int64_t iteration = *lower; iteration < *upper; iteration += *step) { - FailureOr evaluated = evaluateSingleResultAffineMap(map, ArrayRef{iteration}); - if (succeeded(evaluated) && *evaluated == coreId) - return true; - } - - return false; -} - -bool operationContainsCommunication(Operation& op) { - bool found = false; - WalkResult walkResult = op.walk([&](Operation* nestedOp) -> WalkResult { - if (!isa(nestedOp)) - return WalkResult::advance(); - found = true; - return WalkResult::interrupt(); - }); - (void) walkResult; - return found; -} - -bool operationContainsSend(Operation& op) { - bool found = false; - WalkResult walkResult = op.walk([&](SpatChannelSendOp) -> WalkResult { - found = true; - return WalkResult::interrupt(); - }); - (void) walkResult; - return found; -} - -bool operationContainsReceiveFromPeer(Operation& op, int64_t sourceCoreId, int64_t targetCoreId, Type payloadType) { - bool found = false; - WalkResult walkResult = op.walk([&](SpatChannelReceiveOp receive) -> WalkResult { - if (receive.getType() != payloadType) - return WalkResult::advance(); - if (!valueMayEvaluateToCore(receive.getSourceCoreId(), targetCoreId)) - return WalkResult::advance(); - if (!valueMayEvaluateToCore(receive.getTargetCoreId(), sourceCoreId)) - return WalkResult::advance(); - - found = true; - return WalkResult::interrupt(); - }); - (void) walkResult; - return found; -} - -bool valueDependsOn(Value value, Value dependency) { - if (value == dependency) - return true; - - 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* definingOp = current.getDefiningOp(); - if (!definingOp) - continue; - llvm::append_range(worklist, definingOp->getOperands()); - } - - return false; -} - -bool opDependsOnValue(Operation& op, Value dependency) { - bool found = false; - WalkResult walkResult = op.walk([&](Operation* nestedOp) -> WalkResult { - for (Value operand : nestedOp->getOperands()) { - if (!valueDependsOn(operand, dependency)) - continue; - found = true; - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - (void) walkResult; - return found; -} - -LogicalResult orderLowerCoreScalarSendsAfterMatchingReceives(MaterializerState& state) { - for (MaterializedClass& materializedClass : state.classes) { - if (materializedClass.isBatch) - continue; - - Block& body = *materializedClass.body; - SmallVector topLevelOps; - for (Operation& op : body.without_terminator()) - topLevelOps.push_back(&op); - - for (Operation* op : topLevelOps) { - auto send = dyn_cast(op); - if (!send || send->getBlock() != &body) - continue; - - std::optional sourceCoreId = getStaticCommunicationCore(send.getSourceCoreId()); - std::optional targetCoreId = getStaticCommunicationCore(send.getTargetCoreId()); - if (!sourceCoreId || !targetCoreId || *sourceCoreId >= *targetCoreId) - continue; - - Operation* anchor = nullptr; - for (Operation* cursor = send->getNextNode(); cursor && cursor != body.getTerminator(); cursor = cursor->getNextNode()) { - if (isTopLevelCommunicationOp(cursor)) { - auto receive = dyn_cast(cursor); - if (!receive) - break; - - std::optional receiveSourceCoreId = getStaticCommunicationCore(receive.getSourceCoreId()); - std::optional receiveTargetCoreId = getStaticCommunicationCore(receive.getTargetCoreId()); - if (!receiveSourceCoreId || !receiveTargetCoreId) - break; - if (*receiveSourceCoreId != *targetCoreId || *receiveTargetCoreId != *sourceCoreId) - break; - if (receive.getType() != send.getInput().getType()) - break; - - anchor = receive; - break; - } - - if (!operationContainsCommunication(*cursor)) - continue; - if (operationContainsSend(*cursor)) - break; - if (!operationContainsReceiveFromPeer(*cursor, *sourceCoreId, *targetCoreId, send.getInput().getType())) - break; - - anchor = cursor; - break; - } - - if (!anchor) - continue; - send->moveAfter(anchor); - } - - for (Operation* op : topLevelOps) { - auto receive = dyn_cast(op); - if (!receive || receive->getBlock() != &body) - continue; - - std::optional sourceCoreId = getStaticCommunicationCore(receive.getSourceCoreId()); - std::optional targetCoreId = getStaticCommunicationCore(receive.getTargetCoreId()); - if (!sourceCoreId || !targetCoreId || *targetCoreId >= *sourceCoreId) - continue; - - Operation* anchor = nullptr; - for (Operation* cursor = receive->getNextNode(); cursor && cursor != body.getTerminator(); cursor = cursor->getNextNode()) { - if (!isTopLevelCommunicationOp(cursor)) { - if (opDependsOnValue(*cursor, receive.getOutput())) - break; - continue; - } - - auto send = dyn_cast(cursor); - if (!send || send.getInput().getType() != receive.getType()) - break; - - std::optional sendSourceCoreId = getStaticCommunicationCore(send.getSourceCoreId()); - std::optional sendTargetCoreId = getStaticCommunicationCore(send.getTargetCoreId()); - if (!sendSourceCoreId || !sendTargetCoreId) - break; - if (*sendSourceCoreId != *targetCoreId || *sendTargetCoreId != *sourceCoreId) - break; - - bool hasInterveningUse = false; - for (Operation* between = receive->getNextNode(); between; between = between->getNextNode()) { - if (opDependsOnValue(*between, receive.getOutput())) { - hasInterveningUse = true; - break; - } - if (between == cursor) - break; - } - if (hasInterveningUse) - continue; - - anchor = cursor; - break; - } - - if (!anchor) - continue; - receive->moveAfter(anchor); - } - } - - return success(); -} - void replaceHostUses(MaterializerState& state) { for (const auto& [oldValue, replacement] : state.hostReplacements) replaceLiveExternalUses(oldValue, replacement, state.oldComputeOps); @@ -8109,40 +10086,79 @@ MergeScheduleMaterializer::run(func::FuncOp func, const MergeScheduleResult& sch if (schedule.dominanceOrderCompute.empty()) return success(); + auto runStage = [&](StringRef stage, auto&& callback) -> LogicalResult { + if (succeeded(callback())) + return success(); + return func.emitError("merge schedule materialization failed during ") << stage; + }; + MaterializerState state(func, schedule, nextChannelId); - if (failed(buildMaterializationWorkStreams(state))) + if (failed(runStage("work-stream construction", [&]() { return buildMaterializationWorkStreams(state); }))) return failure(); - if (failed(buildMaterializationClassesFromScheduleEquivalence(state))) + if (failed(runStage("equivalence-class construction", [&]() { + return buildMaterializationClassesFromScheduleEquivalence(state); + }))) return failure(); - if (failed(verifyScheduleEquivalenceMatchesLogicalStreams(state))) + if (failed(runStage("equivalence/logical-stream verification", [&]() { + return verifyScheduleEquivalenceMatchesLogicalStreams(state); + }))) return failure(); if (state.classes.empty()) return success(); - if (failed(collectHostOutputs(state))) + if (failed(runStage("host-output collection", [&]() { return collectHostOutputs(state); }))) return failure(); - if (failed(createEmptyMaterializedOps(state))) + if (failed(runStage("empty materialized-op creation", [&]() { return createEmptyMaterializedOps(state); }))) return failure(); - if (failed(collectProducerDestinations(state))) + if (failed(runStage("producer-destination collection", [&]() { return collectProducerDestinations(state); }))) return failure(); - if (failed(collectProjectedTransfers(state))) + if (failed(runStage("projected-transfer collection", [&]() { return collectProjectedTransfers(state); }))) + return failure(); + if (failed(runStage("planned scalar-peer receive collection", [&]() { + return collectPlannedScalarPeerReceives(state); + }))) return failure(); - for (const ComputeInstance& instance : schedule.dominanceOrderCompute) - if (failed(materializeInstanceSlot(state, instance))) - return failure(); - - if (failed(finalizeProjectedHostOutputFragments(state))) - return failure(); - if (failed(orderLowerCoreScalarSendsAfterMatchingReceives(state))) + if (failed(runStage("instance-slot materialization", [&]() -> LogicalResult { + for (const ComputeInstance& instance : schedule.dominanceOrderCompute) { + if (succeeded(materializeInstanceSlot(state, instance))) + continue; + InFlightDiagnostic diag = instance.op->emitError( + "merge schedule materialization failed while materializing scheduled compute instance"); + diag << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount; + return failure(); + } + return success(); + }))) return failure(); - for (MaterializedClass& materializedClass : state.classes) - if (failed(localizeAllScheduledBodyCaptures(state, materializedClass))) - return failure(); + if (failed(runStage("projected input send phase verification", [&]() { + return verifyNoPendingProjectedInputSends(state); + }))) + return failure(); + + if (failed(runStage("projected host-output finalization", [&]() { + return finalizeProjectedHostOutputFragments(state); + }))) + return failure(); + + if (failed(runStage("scheduled-body capture localization", [&]() -> LogicalResult { + for (MaterializedClass& materializedClass : state.classes) { + if (succeeded(localizeAllScheduledBodyCaptures(state, materializedClass))) + continue; + return materializedClass.op->emitError( + "merge schedule materialization failed while localizing scheduled-body captures") + << " classId=" << materializedClass.id; + } + return success(); + }))) + return failure(); + + if (failed(runStage("pending scalar-peer send verification", [&]() { return verifyNoPendingScalarPeerSends(state); }))) + return failure(); replaceHostUses(state); - if (failed(eraseOldComputeOps(state))) + if (failed(runStage("old compute-op erasure", [&]() { return eraseOldComputeOps(state); }))) return failure(); LogicalResult _ = runRegionDCE(state.rewriter, state.func.getBody()); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp index 3c243d7..5f41098 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp @@ -121,6 +121,51 @@ struct CompactRunPlan { llvm::SmallVector outputs; }; +struct ScalarPeerEdgeKey { + int64_t sourceCore = 0; + int64_t targetCore = 0; + mlir::Type payloadType; +}; + +struct ScalarPeerReceiveKey { + int64_t sourceCore = 0; + int64_t targetCore = 0; + mlir::Type payloadType; + std::optional channelId; +}; + +struct PendingScalarSend { + ClassId sourceClass = 0; + int64_t sourceCore = 0; + int64_t targetCore = 0; + mlir::Type payloadType; + ScalarPeerEdgeKey waitForReceive; + mlir::Operation* payloadAnchor = nullptr; + mlir::Value payload; + MessageVector messages; + mlir::Location loc; +}; + +struct PendingProjectedScalarSend { + ClassId sourceClass = 0; + int64_t sourceCore = 0; + int64_t targetCore = 0; + mlir::Type payloadType; + ScalarPeerEdgeKey waitForReceive; + mlir::Operation* payloadAnchor = nullptr; + mlir::Value payload; + MessageVector messages; + ProjectedTransferDescriptor descriptor; + mlir::Location loc; +}; + +struct PendingProjectedInputSend { + ClassId sourceClass = 0; + mlir::Value payload; + llvm::SmallVector fragments; + mlir::Location loc; +}; + enum class BatchInputDemandKind { LaneFragment, ProjectedFragment, @@ -235,10 +280,19 @@ struct MaterializerState { projectedTransfers; llvm::DenseMap> projectedExtractReplacements; + llvm::DenseMap> + projectedInputTransferPlans; AvailableValueStore availableValues; llvm::DenseMap hostReplacements; llvm::DenseMap hostOutputOwners; llvm::SmallVector pendingProjectedHostOutputFragments; + llvm::SmallVector plannedScalarPeerReceives; + llvm::SmallVector materializedScalarPeerReceives; + llvm::SmallVector pendingScalarSends; + llvm::SmallVector pendingProjectedScalarSends; + llvm::SmallVector pendingProjectedInputSends; + llvm::DenseSet projectedInputPhaseBarrierClasses; + llvm::DenseMap pendingProjectedHighToLowReceives; llvm::DenseSet oldComputeOps; MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId) diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp index 0a655b9..6ce50df 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp @@ -47,6 +47,24 @@ struct ProjectedExtractReplacement { ProjectedFragmentLayout layout; }; +struct ProjectedInputTransferFragment { + ProducerKey producer; + llvm::SmallVector fragmentOffsets; + unsigned targetLane = 0; + unsigned ordinal = 0; + int64_t channelId = 0; + int32_t sourceCoreId = 0; + int32_t targetCoreId = 0; + bool sendEmitted = false; +}; + +struct ProjectedInputTransferPlan { + ProjectedBatchInputKey inputKey; + mlir::Operation* extractOp = nullptr; + ProjectedFragmentLayout layout; + llvm::SmallVector fragments; +}; + struct PendingProjectedHostOutputFragment { mlir::Value originalOutput; ClassId sourceClass = 0; diff --git a/src/PIM/Pass/PIMPasses.h b/src/PIM/Pass/PIMPasses.h index 5e03534..efda377 100644 --- a/src/PIM/Pass/PIMPasses.h +++ b/src/PIM/Pass/PIMPasses.h @@ -11,8 +11,6 @@ std::unique_ptr createONNXToSpatialPass(); std::unique_ptr createSpatialLayoutPlanningPass(); std::unique_ptr createLowerSpatialPlansPass(); -std::unique_ptr createSpatialToGraphvizPass(); - std::unique_ptr createSpatialToPimPass(); std::unique_ptr createPimBufferizationPass(); diff --git a/src/PIM/PimAccelerator.cpp b/src/PIM/PimAccelerator.cpp index eaa44b2..6bde8ce 100644 --- a/src/PIM/PimAccelerator.cpp +++ b/src/PIM/PimAccelerator.cpp @@ -74,7 +74,6 @@ void PimAccelerator::registerPasses(int optLevel) const { registerPass(createONNXToSpatialPass); registerPass(createSpatialLayoutPlanningPass); registerPass(createLowerSpatialPlansPass); - registerPass(createSpatialToGraphvizPass); registerPass(createSpatialToPimPass); registerPass(createPimBufferizationPass); registerPass(createPimMemoryCoalescingPass); diff --git a/validation/validate.py b/validation/validate.py old mode 100644 new mode 100755