diff --git a/README.md b/README.md index 5519851..f243e35 100644 --- a/README.md +++ b/README.md @@ -118,8 +118,6 @@ options; `onnx-mlir --help` lists the inherited ONNX-MLIR options. elements per convolution before streaming. Default is `1048576`. - `--pim-conv-stream-chunk-positions=` - maximum output positions per streamed convolution chunk. Default is `1024`. -- `--pim-report-conv-lowering=` - emit the bounded convolution - lowering report. Default is `true`. - `--use-experimental-conv-impl` - use the alternate convolution lowering. - `--pim-detect-communication-deadlock` - statically simulate expanded send/receive ordering and reject blocking deadlocks. Default is off. diff --git a/src/PIM/Common/IR/AddressAnalysis.cpp b/src/PIM/Common/IR/AddressAnalysis.cpp index 1fb13bc..0e3eb92 100644 --- a/src/PIM/Common/IR/AddressAnalysis.cpp +++ b/src/PIM/Common/IR/AddressAnalysis.cpp @@ -1,10 +1,13 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/Interfaces/DestinationStyleOpInterface.h" +#include "llvm/ADT/SmallPtrSet.h" + #include #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp" @@ -36,6 +39,10 @@ mlir::Value resolveAlias(mlir::Value value, const StaticValueKnowledge* knowledg llvm::FailureOr compileIndexValueImpl(mlir::Value value); llvm::FailureOr compileContiguousAddressExprImpl(mlir::Value value); +using AliasResolutionSet = llvm::SmallPtrSet; +mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, + const StaticValueKnowledge* knowledge, + AliasResolutionSet& visited); mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge); template @@ -45,18 +52,23 @@ CompiledIndexExpr makeCompiledIndexExpr(Args&&... args) { static mlir::Value resolveForYieldedAliasToInit(mlir::scf::ForOp forOp, mlir::Value yieldedValue, - const StaticValueKnowledge* knowledge) { - yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge); + const StaticValueKnowledge* knowledge, + AliasResolutionSet& visited) { + yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge, visited); if (auto blockArgument = mlir::dyn_cast(yieldedValue)) { if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0 && static_cast(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size()) - return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge); + return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge, visited); } return yieldedValue; } -mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) { +mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, + const StaticValueKnowledge* knowledge, + AliasResolutionSet& visited) { value = resolveAlias(value, knowledge); + if (!value || !visited.insert(value).second) + return value; if (auto blockArgument = mlir::dyn_cast(value)) { auto forOp = mlir::dyn_cast_or_null(blockArgument.getOwner()->getParentOp()); @@ -64,9 +76,12 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow const unsigned iterArgIndex = blockArgument.getArgNumber() - 1; auto yieldOp = mlir::dyn_cast(forOp.getBody()->getTerminator()); if (iterArgIndex < forOp.getInitArgs().size() && yieldOp - && iterArgIndex < yieldOp.getNumOperands() - && resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge) == blockArgument) - return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge); + && iterArgIndex < yieldOp.getNumOperands()) { + mlir::Value yieldedValue = resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge); + if (yieldedValue == blockArgument + || (yieldedValue && resolveLoopCarriedAliasImpl(yieldedValue, knowledge, visited) == blockArgument)) + return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge, visited); + } } return value; } @@ -75,10 +90,15 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow if (!definingOp) return value; + if (auto toBufferOp = mlir::dyn_cast(definingOp)) + return resolveLoopCarriedAliasImpl(toBufferOp.getTensor(), knowledge, visited); + if (auto toTensorOp = mlir::dyn_cast(definingOp)) + return resolveLoopCarriedAliasImpl(toTensorOp.getBuffer(), knowledge, visited); + if (auto dpsDefiningOp = mlir::dyn_cast(definingOp)) { if (auto result = mlir::dyn_cast(value)) if (mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(result)) - return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge); + return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge, visited); } if (auto forOp = mlir::dyn_cast(definingOp)) { @@ -86,20 +106,26 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow if (result) { auto yieldOp = mlir::dyn_cast(forOp.getBody()->getTerminator()); if (yieldOp && result.getResultNumber() < yieldOp.getNumOperands()) - return resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge); + return resolveForYieldedAliasToInit( + forOp, yieldOp.getOperand(result.getResultNumber()), knowledge, visited); } } if (auto castOp = mlir::dyn_cast(definingOp)) - return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge); + return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge, visited); if (auto collapseOp = mlir::dyn_cast(definingOp)) - return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge); + return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge, visited); if (auto expandOp = mlir::dyn_cast(definingOp)) - return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge); + return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge, visited); return value; } +mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) { + AliasResolutionSet visited; + return resolveLoopCarriedAliasImpl(value, knowledge, visited); +} + llvm::FailureOr resolveOpFoldResult(mlir::OpFoldResult ofr, const StaticValueKnowledge* knowledge); llvm::FailureOr resolveIndexValueImpl(mlir::Value value, const StaticValueKnowledge* knowledge); @@ -524,6 +550,15 @@ llvm::FailureOr resolveContiguousAddressImpl(mlir::Va if (!definingOp) return mlir::failure(); + if (auto toBufferOp = mlir::dyn_cast(definingOp)) { + value = resolveAlias(toBufferOp.getTensor(), knowledge); + continue; + } + if (auto toTensorOp = mlir::dyn_cast(definingOp)) { + value = resolveAlias(toTensorOp.getBuffer(), knowledge); + continue; + } + if (auto dpsDefiningOp = mlir::dyn_cast(definingOp)) { mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast(value)); if (!tiedOperand) @@ -538,7 +573,9 @@ llvm::FailureOr resolveContiguousAddressImpl(mlir::Va return mlir::failure(); auto yieldOp = mlir::cast(forOp.getBody()->getTerminator()); - value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge); + AliasResolutionSet visited; + value = resolveForYieldedAliasToInit( + forOp, yieldOp.getOperand(result.getResultNumber()), knowledge, visited); continue; } @@ -643,6 +680,15 @@ llvm::FailureOr compileContiguousAddressExprImpl(mlir::Valu if (!definingOp) return mlir::failure(); + if (auto toBufferOp = mlir::dyn_cast(definingOp)) { + value = toBufferOp.getTensor(); + continue; + } + if (auto toTensorOp = mlir::dyn_cast(definingOp)) { + value = toTensorOp.getBuffer(); + continue; + } + if (auto dpsDefiningOp = mlir::dyn_cast(definingOp)) { mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast(value)); if (!tiedOperand) @@ -657,7 +703,9 @@ llvm::FailureOr compileContiguousAddressExprImpl(mlir::Valu return mlir::failure(); auto yieldOp = mlir::cast(forOp.getBody()->getTerminator()); - value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), nullptr); + AliasResolutionSet visited; + value = resolveForYieldedAliasToInit( + forOp, yieldOp.getOperand(result.getResultNumber()), nullptr, visited); continue; } diff --git a/src/PIM/Compiler/PimCompilerOptions.cpp b/src/PIM/Compiler/PimCompilerOptions.cpp index dfa5e6c..cd22550 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -87,11 +87,6 @@ llvm::cl::opt pimConvStreamChunkPositions( llvm::cl::init(1024), llvm::cl::cat(OnnxMlirOptions)); -llvm::cl::opt pimReportConvLowering("pim-report-conv-lowering", - llvm::cl::desc("Emit a bounded Conv lowering report"), - llvm::cl::init(true), - llvm::cl::cat(OnnxMlirOptions)); - llvm::cl::opt pimEmitJson("pim-emit-json", llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"), llvm::cl::init(false), diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index 00c2472..84cf80f 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -57,7 +57,6 @@ extern llvm::cl::opt pimExportSpatialDataflow; extern llvm::cl::opt pimOnlyCodegen; extern llvm::cl::opt useExperimentalConvImpl; extern llvm::cl::opt pimEmitJson; -extern llvm::cl::opt pimReportConvLowering; extern llvm::cl::opt pimDetectCommunicationDeadlock; extern llvm::cl::opt pimMaterializeScalarFanoutGlobalOrder; extern llvm::cl::opt pimTraceCommunicationMaterialization; diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp index 9f430b0..93db327 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp @@ -47,13 +47,17 @@ FailureOr createFragmentAssemblyBlueprint(Value physicalBatch, llvm::append_range(offsets, entry.destinationOffsets); llvm::append_range(sizes, entry.sizes); } - return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {}, + auto blueprint = spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {}, rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout), rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes), rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"), rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots), rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides), - rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput(); + rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")); + if (indexMap == spatial::kContiguousRowMajorFragments + && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint)) + blueprint.setIndexMapAttr(rewriter.getStringAttr("fragment_assembly")); + return blueprint.getOutput(); } Value sumTensors(ArrayRef tensors, PatternRewriter& rewriter) { diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp index 398a520..0132655 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp @@ -394,6 +394,39 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter, rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides}); } +template +mlir::FailureOr mapGraphBatchFragments(mlir::Value input, + mlir::RankedTensorType outputType, + mlir::PatternRewriter& rewriter, + mlir::Location loc, + BodyFn&& build) { + auto inputType = mlir::dyn_cast(input.getType()); + if (!inputType || !inputType.hasStaticShape() || !outputType.hasStaticShape() + || inputType.getRank() != outputType.getRank() || inputType.getRank() < 2 + || inputType.getDimSize(0) != outputType.getDimSize(0)) + return mlir::failure(); + auto inputFragmentType = mlir::RankedTensorType::get( + inputType.getShape().drop_front(), inputType.getElementType(), inputType.getEncoding()); + auto outputFragmentType = mlir::RankedTensorType::get( + outputType.getShape().drop_front(), outputType.getElementType(), outputType.getEncoding()); + auto batch = createSpatComputeBatch( + rewriter, loc, mlir::TypeRange {outputType}, inputType.getDimSize(0), {}, mlir::ValueRange {input}, + [&](detail::SpatComputeBatchBodyArgs args) -> mlir::LogicalResult { + auto fragment = extractGraphBatchPhysicalFragment( + rewriter, loc, args.inputs.front(), args.lane, inputFragmentType); + if (mlir::failed(fragment)) + return mlir::failure(); + mlir::FailureOr result = build(*fragment, outputFragmentType); + if (mlir::failed(result) || result->getType() != outputFragmentType) + return mlir::failure(); + publishGraphBatchPhysicalFragment(rewriter, loc, *result, args.outputs.front(), args.lane); + return mlir::success(); + }); + if (mlir::failed(batch)) + return mlir::failure(); + return batch->getResult(0); +} + template mlir::Value materializeOrComputeUnary(mlir::Value input, mlir::RankedTensorType resultType, diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp index 413911c..a39de5d 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp @@ -186,25 +186,9 @@ static FailureOr applyRowStripActivation(const RowStripPhysicalValue& val Location loc, BuildActivation buildActivation) { auto storageType = cast(value.storage.getType()); - const int64_t laneCount = storageType.getDimSize(0); - auto batchOp = createSpatComputeBatch(rewriter, - loc, - TypeRange {storageType}, - laneCount, - {}, - ValueRange {value.storage}, - [&](detail::SpatComputeBatchBodyArgs args) { - FailureOr fragment = extractGraphBatchPhysicalFragment( - rewriter, loc, args.inputs.front(), args.lane, value.fragmentType); - if (failed(fragment)) return failure(); - Value result = buildActivation(*fragment); - publishGraphBatchPhysicalFragment( - rewriter, loc, result, args.outputs.front(), args.lane); - return success(); - }); - if (failed(batchOp)) - return failure(); - return batchOp->getResult(0); + return mapGraphBatchFragments(value.storage, storageType, rewriter, loc, [&](Value fragment, RankedTensorType) { + return FailureOr(buildActivation(fragment)); + }); } FailureOr applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) { diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp index cdd6ff0..74bb30b 100644 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp @@ -399,6 +399,42 @@ struct LowerSpatialPlansPass final : PassWrapper(&op)) { + FailureOr input = + getRowStripValue(rowStripValues, planOp.getInput()); + rewriter.setInsertionPoint(planOp); + auto lowered = lowerSelectedResizeNearestPlan( + planOp, succeeded(input) ? std::optional(input->storage) : std::nullopt, + rewriter); + if (failed(lowered)) { + planOp.emitOpError("failed to lower selected nearest Resize plan"); + signalPassFailure(); + return; + } + if (failed(input)) { + rewriter.replaceOp(planOp, *lowered); + continue; + } + 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 Resize plan requires a row-strip blueprint result"); + signalPassFailure(); + return; + } + auto blueprint = cast(*outputBlueprint); + auto output = buildRowStripValue(blueprint, *lowered); + if (failed(output)) { + signalPassFailure(); + return; + } + rowStripValues[blueprint.getResult()] = *output; + eraseAfterLowering.insert(planOp); + eraseAfterLowering.insert(blueprint); + continue; + } if (auto planOp = dyn_cast(&op)) { auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { auto blueprint = dyn_cast(user); @@ -701,6 +737,7 @@ struct LowerSpatialPlansPass final : PassWrapper(op) diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp index 1e25d72..a7b4d72 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp @@ -52,13 +52,16 @@ static void populateEmptyFunction(func::FuncOp funcOp) { SmallVector concatPlans(funcOp.getOps()); SmallVector reluPlans(funcOp.getOps()); SmallVector siluPlans(funcOp.getOps()); + SmallVector resizePlans( + funcOp.getOps()); SmallVector maxPoolPlans(funcOp.getOps()); SmallVector globalAveragePoolPlans( funcOp.getOps()); SmallVector blueprints(funcOp.getOps()); SmallVector materializers(funcOp.getOps()); if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !addPlans.empty() - || !concatPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty() + || !concatPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !resizePlans.empty() + || !maxPoolPlans.empty() || !blueprints.empty() || !globalAveragePoolPlans.empty() || !materializers.empty()) { return; } @@ -123,6 +126,14 @@ void ONNXToSpatialPass::runOnOperation() { return; } + RewritePatternSet matmulPatterns(ctx); + populateMatMulFusionPatterns(matmulPatterns, ctx); + if (failed(applyPatternsGreedily(moduleOp, std::move(matmulPatterns)))) { + moduleOp.emitError("failed to lower MatMul before producer conversion"); + signalPassFailure(); + return; + } + RewritePatternSet fusionPatterns(ctx); populateElementwiseFusionPatterns(fusionPatterns, ctx); if (failed(applyPatternsGreedily(moduleOp, std::move(fusionPatterns)))) { diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp index 9e92dcf..258fc89 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp @@ -150,6 +150,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter spatial::SpatConcatPlanOp, spatial::SpatReluPlanOp, spatial::SpatSiluPlanOp, + spatial::SpatResizeNearestPlanOp, spatial::SpatMaxPool2DPlanOp, spatial::SpatGlobalAveragePoolPlanOp, spatial::SpatBlueprintOp, diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp b/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp index e3cec73..bb9d069 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp @@ -20,6 +20,7 @@ void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRCo void populateElementwiseFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateGemmPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); +void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateReduceMeanPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp index 2450228..4eb5450 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -1,28 +1,19 @@ #include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypes.h" -#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallVector.h" -#include "llvm/Support/raw_ostream.h" #include -#include -#include -#include #include -#include -#include #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp" #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" @@ -46,199 +37,12 @@ struct ConvToGemm : OpConversionPattern { ConversionPatternRewriter& rewriter) const override; }; -struct ConvLoweringDecision { - PimConvLoweringType strategy; - std::string reason; - bool isAuto = false; - std::string fallbackReason; - std::string rejectedAutoStrategy; -}; - struct PreparedConvInput { Value value; RankedTensorType type; }; -struct ConvStrategyEstimate { - uint64_t estimatedMvmCount = 0; - uint64_t estimatedReductionVAddCount = 0; - uint64_t estimatedOutputFragments = 1; - bool perOutputPositionReduction = false; - bool requiresFuncReturnMaterialization = false; - bool giantCollectorConcatExpected = false; - bool fullInputBroadcastExpected = false; - uint64_t concatOperandCount = 0; - std::string materializationKind = "none"; - std::string collectorCore = "-"; -}; - -struct ConvReportEntry { - uint64_t id; - std::string where; - std::string strategy; - std::string mode; - std::string inputShape; - std::string weightShape; - std::string outputShape; - int64_t groups; - int64_t k; - int64_t c; - int64_t p; - int64_t xbarSize; - int64_t pack; - uint64_t im2colElements; - uint64_t im2colBudget; - std::string chunkText; - int64_t batchSize; - int64_t numberOfBatches; - std::string spatialComputeBatch; - std::string batchedInstructionEmission; - std::string reason; - std::string fallbackReason; - std::string rejectedAutoStrategy; - uint64_t estimatedMvmCount; - uint64_t estimatedReductionVAddCount; - uint64_t estimatedOutputFragments; - std::string materializationRequiredAtReturn; - std::string materializationKind; - uint64_t concatOperandCount; - std::string collectorCore; - std::string giantCollectorConcatExpected; - std::string fullInputBroadcastExpected; -}; - -enum class DistributedTensorOpKind { - Relu, - Sigmoid, - Add, - Sub, - Mul, - Div, - Conv, -}; - -enum class DistributedConvBarrierKind { - Return, - UnsupportedConsumer, - Fanout, - DeadValue, - GroupedConv, - Depthwise, -}; - -enum class DistributedTensorConstantKind { - None, - Splat, - PerChannel, -}; - -enum class DistributedTensorLayoutKind { - NchwRowStrip, -}; - -struct DistributedFragmentInfo { - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - int64_t producerLane = 0; -}; - -struct DistributedTensorInfo { - Value storage; - RankedTensorType logicalType; - DistributedTensorLayoutKind layoutKind = DistributedTensorLayoutKind::NchwRowStrip; - SmallVector fragments; - int64_t laneCount = 0; - int64_t fragmentHeight = 1; - int64_t channels = 0; - int64_t height = 0; - int64_t width = 0; - - bool isRowStripNchw() const { return layoutKind == DistributedTensorLayoutKind::NchwRowStrip; } -}; - -struct DistributedTensorRegistry { - llvm::DenseMap infos; - - void bind(Value value, const DistributedTensorInfo& info) { infos[value] = info; } - - const DistributedTensorInfo* lookup(Value value) const { - auto it = infos.find(value); - if (it == infos.end()) - return nullptr; - return &it->second; - } -}; - -struct DistributedTensorStep { - Operation* op = nullptr; - DistributedTensorOpKind kind; - DenseElementsAttr constantAttr; - DistributedTensorConstantKind constantKind = DistributedTensorConstantKind::None; - bool fragmentOnLhs = true; - std::optional convState; -}; - -struct DistributedConvAnalysis { - SmallVector steps; - Operation* replacementOp = nullptr; - DistributedConvBarrierKind barrierKind = DistributedConvBarrierKind::UnsupportedConsumer; - std::string barrierDetail; - - bool hasLocalConsumers() const { return !steps.empty(); } - bool hasDistributedConvConsumer() const { - return llvm::any_of(steps, [](const DistributedTensorStep& step) { return step.kind == DistributedTensorOpKind::Conv; }); - } -}; - -struct DistributedChainReportEntry { - uint64_t chainId = 0; - uint64_t chainLength = 0; - std::string producerKind; - std::string distributedOps; - std::string materializationPoints; - std::string firstMaterializationReason; - uint64_t maxLiveFragments = 0; - uint64_t maxFragmentFanout = 0; - uint64_t patchBuilderCoreCount = 0; - std::string centralJunctionDetected = "no"; - std::string convInputMaterializationKind = "dense_materialization_fallback"; - uint64_t localPatchFragments = 0; - uint64_t remotePatchFragments = 0; - uint64_t haloTransferCount = 0; - uint64_t groupedTransferCount = 0; - std::string fallbackReason; -}; - -struct DistributedConvReportTotals { - uint64_t totalConvs = 0; - uint64_t distributedTensorsCreated = 0; - uint64_t distributedValuesPropagated = 0; - uint64_t distributedConsumersHandled = 0; - uint64_t distributedConvInputsSeen = 0; - uint64_t distributedConvInputsConsumed = 0; - uint64_t materializationBarriersInserted = 0; - std::map fallbackReasons; - std::map barrierReasons; - SmallVector chains; -}; - static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter); - -static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b); - -static StringRef stringifyDistributedConvBarrierKind(DistributedConvBarrierKind kind) { - switch (kind) { - case DistributedConvBarrierKind::Return: return "func.return"; - case DistributedConvBarrierKind::UnsupportedConsumer: return "unsupported_consumer"; - case DistributedConvBarrierKind::Fanout: return "fanout"; - case DistributedConvBarrierKind::DeadValue: return "dead_value"; - case DistributedConvBarrierKind::GroupedConv: return "grouped_conv"; - case DistributedConvBarrierKind::Depthwise: return "depthwise_conv"; - } - llvm_unreachable("unknown distributed conv barrier kind"); -} - static StringRef stringifyConvLoweringStrategy(PimConvLoweringType strategy) { switch (strategy) { case PimConvLoweringAuto: return "auto"; @@ -254,825 +58,29 @@ static StringRef stringifyConvLoweringStrategy(PimConvLoweringType strategy) { llvm_unreachable("unknown conv lowering strategy"); } -static bool requiresFuncReturnMaterialization(const DistributedConvAnalysis& analysis) { - return !analysis.hasLocalConsumers() && analysis.barrierKind == DistributedConvBarrierKind::Return; -} - -static ConvStrategyEstimate estimateConvStrategy(const ConvGeometry& geo, - PimConvLoweringType strategy, - const DistributedConvAnalysis& analysis) { - ConvStrategyEstimate estimate; - estimate.requiresFuncReturnMaterialization = requiresFuncReturnMaterialization(analysis); - estimate.materializationKind = estimate.requiresFuncReturnMaterialization ? "func.return" : "none"; - - switch (strategy) { - case PimConvLoweringLegacy: - case PimConvLoweringPackedIm2Col: - estimate.estimatedMvmCount = static_cast(std::max(1, geo.p)); - break; - case PimConvLoweringStreamedPatch: - case PimConvLoweringStreamedPacked: - case PimConvLoweringOutputChannelTiled: { - uint64_t chunkPositions = chooseStreamChunkPositions(geo, /*packFactor=*/1); - estimate.estimatedMvmCount = static_cast(std::max(1, geo.p)); - estimate.estimatedOutputFragments = - std::max(1, static_cast(ceilIntegerDivide(geo.p, static_cast(chunkPositions)))); - break; - } - case PimConvLoweringInputKTiled: { - const int64_t numKSlices = ceilIntegerDivide(geo.k, geo.xbarSize); - const uint64_t maxLanesPerBatch = - std::max(1, - static_cast(crossbarCountInCore.getValue()) - / static_cast(std::max(1, numKSlices * 4))); - const uint64_t rowChunkWidth = std::max( - 1, - std::min({chooseStreamChunkPositions(geo, /*packFactor=*/1), - maxLanesPerBatch, - static_cast(std::max(1, geo.outWidth))})); - estimate.estimatedMvmCount = - static_cast(std::max(1, geo.p)) * static_cast(std::max(1, numKSlices)); - estimate.estimatedReductionVAddCount = - static_cast(std::max(1, geo.p)) - * static_cast(std::max(0, numKSlices - 1) + (geo.hasBias ? 1 : 0)); - estimate.estimatedOutputFragments = static_cast(std::max(1, geo.batchSize)) - * static_cast(std::max(1, geo.outHeight)) - * static_cast( - ceilIntegerDivide(geo.outWidth, static_cast(rowChunkWidth))); - estimate.perOutputPositionReduction = numKSlices > 1; - estimate.fullInputBroadcastExpected = estimate.estimatedOutputFragments > 1; - if (estimate.requiresFuncReturnMaterialization && estimate.estimatedOutputFragments >= 128) { - estimate.giantCollectorConcatExpected = true; - estimate.materializationKind = "single_collector_concat"; - estimate.concatOperandCount = estimate.estimatedOutputFragments; - estimate.collectorCore = "scheduled_post_merge"; - } - break; - } - case PimConvLoweringTiled2D: - estimate.estimatedMvmCount = static_cast(std::max(1, geo.p)); - break; - case PimConvLoweringDepthwise: - case PimConvLoweringAuto: - break; - } - - if (estimate.requiresFuncReturnMaterialization && estimate.materializationKind == "none") - estimate.materializationKind = "func.return"; - return estimate; -} - -static std::string formatShape(ArrayRef dims) { - std::string text; - llvm::raw_string_ostream os(text); - os << "["; - for (size_t i = 0; i < dims.size(); ++i) { - if (i != 0) - os << "x"; - os << dims[i]; - } - os << "]"; - return text; -} - -static std::string collapseWhitespace(StringRef text) { - std::string out; - out.reserve(text.size()); - bool lastWasSpace = false; - for (char c : text) { - bool isSpace = std::isspace(static_cast(c)); - if (isSpace) { - if (!lastWasSpace && !out.empty()) - out.push_back(' '); - lastWasSpace = true; - continue; - } - out.push_back(c); - lastWasSpace = false; - } - return out; -} - -static std::string abbreviate(StringRef text, size_t maxLen) { - if (text.size() <= maxLen) - return text.str(); - return (text.take_front(maxLen - 3) + "...").str(); -} - -static std::string abbreviateFromEnd(StringRef text, size_t maxLen) { - if (text.size() <= maxLen) - return text.str(); - return ("..." + text.take_back(maxLen - 3)).str(); -} - -static std::string summarizeLocation(Location loc, size_t maxLen = 44) { - std::string text; - llvm::raw_string_ostream os(text); - loc.print(os); - os.flush(); - std::string collapsed = collapseWhitespace(text); - if (collapsed.size() <= maxLen) - return collapsed; - if (collapsed.find('/') != std::string::npos || collapsed.find('#') != std::string::npos) - return abbreviateFromEnd(collapsed, maxLen); - return abbreviate(collapsed, maxLen); -} - -static std::string alignCell(StringRef text, size_t width, bool rightAlign = false) { - std::string cell = text.str(); - if (cell.size() < width) { - size_t padding = width - cell.size(); - if (rightAlign) - cell.insert(cell.begin(), padding, ' '); - else - cell.append(padding, ' '); - } - return cell; -} - -static bool hasSameStaticTensorType(Value value, Type expectedType) { - auto valueType = dyn_cast(value.getType()); - auto expectedTensorType = dyn_cast(expectedType); - return valueType && expectedTensorType && valueType.hasStaticShape() && valueType == expectedTensorType; -} - -static bool isSplatConstantValue(Value value, DenseElementsAttr& denseAttr) { - denseAttr = getHostConstDenseElementsAttr(value); - return static_cast(denseAttr) && denseAttr.isSplat(); -} - -static bool isPerChannelConstantValue(Value value, RankedTensorType currentType, DenseElementsAttr& denseAttr) { - denseAttr = getHostConstDenseElementsAttr(value); - if (!denseAttr || denseAttr.isSplat()) - return false; - - auto constantType = dyn_cast(denseAttr.getType()); - if (!constantType || !constantType.hasStaticShape()) - return false; - - const int64_t channels = currentType.getDimSize(1); - if (constantType.getRank() == 1) - return constantType.getDimSize(0) == channels; - if (constantType.getRank() == 2) - return constantType.getDimSize(0) == 1 && constantType.getDimSize(1) == channels; - if (constantType.getRank() != 4) - return false; - return constantType.getDimSize(0) == 1 && constantType.getDimSize(1) == channels - && constantType.getDimSize(2) == 1 && constantType.getDimSize(3) == 1; -} - -static std::optional -classifyDistributedBinaryConsumer(Operation* user, - Value currentValue, - Value lhs, - Value rhs, - DistributedTensorOpKind kind, - bool allowFragmentOnRhs, - std::string& failureDetail) { - if (user->getNumResults() != 1 || !hasSameStaticTensorType(user->getResult(0), currentValue.getType())) { - failureDetail = "result type mismatch"; - return std::nullopt; - } - - auto currentType = cast(currentValue.getType()); - DenseElementsAttr constantAttr; - if (lhs == currentValue) { - DistributedTensorConstantKind constantKind = DistributedTensorConstantKind::None; - if (isSplatConstantValue(rhs, constantAttr)) - constantKind = DistributedTensorConstantKind::Splat; - else if (isPerChannelConstantValue(rhs, currentType, constantAttr)) - constantKind = DistributedTensorConstantKind::PerChannel; - else - failureDetail = "unsupported rhs broadcast"; - if (constantKind == DistributedTensorConstantKind::None) - return std::nullopt; - return DistributedTensorStep {user, kind, constantAttr, constantKind, /*fragmentOnLhs=*/true, std::nullopt}; - } - - if (rhs == currentValue && allowFragmentOnRhs) { - DistributedTensorConstantKind constantKind = DistributedTensorConstantKind::None; - if (isSplatConstantValue(lhs, constantAttr)) - constantKind = DistributedTensorConstantKind::Splat; - else if (isPerChannelConstantValue(lhs, currentType, constantAttr)) - constantKind = DistributedTensorConstantKind::PerChannel; - else - failureDetail = "unsupported lhs broadcast"; - if (constantKind == DistributedTensorConstantKind::None) - return std::nullopt; - return DistributedTensorStep {user, kind, constantAttr, constantKind, /*fragmentOnLhs=*/false, std::nullopt}; - } - - failureDetail = "conv result is not the supported binary operand"; - return std::nullopt; -} - -static std::string stringifyDistributedTensorOpKind(DistributedTensorOpKind kind) { - switch (kind) { - case DistributedTensorOpKind::Relu: return "Relu"; - case DistributedTensorOpKind::Sigmoid: return "Sigmoid"; - case DistributedTensorOpKind::Add: return "Add"; - case DistributedTensorOpKind::Sub: return "Sub"; - case DistributedTensorOpKind::Mul: return "Mul"; - case DistributedTensorOpKind::Div: return "Div"; - case DistributedTensorOpKind::Conv: return "Conv"; - } - llvm_unreachable("unknown distributed tensor op kind"); -} - -[[maybe_unused]] static DistributedConvAnalysis analyzeDistributedConvConsumers(ONNXConvOp convOp) { - DistributedConvAnalysis analysis; - analysis.replacementOp = convOp; - - Value currentValue = convOp.getResult(); - while (true) { - if (currentValue.use_empty()) { - analysis.barrierKind = DistributedConvBarrierKind::DeadValue; - analysis.barrierDetail = "result has no users"; - return analysis; - } - - if (!currentValue.hasOneUse()) { - analysis.barrierKind = DistributedConvBarrierKind::Fanout; - analysis.barrierDetail = "value has multiple users"; - return analysis; - } - - Operation* user = *currentValue.getUsers().begin(); - if (isa(user)) { - analysis.barrierKind = DistributedConvBarrierKind::Return; - analysis.barrierDetail = "materialize at func.return"; - return analysis; - } - - std::optional step; - std::string failureDetail; - if (auto reluOp = dyn_cast(user)) { - if (hasSameStaticTensorType(reluOp.getResult(), currentValue.getType())) - step = DistributedTensorStep { - user, DistributedTensorOpKind::Relu, {}, DistributedTensorConstantKind::None, true, std::nullopt}; - else - failureDetail = "relu result type mismatch"; - } - else if (auto sigmoidOp = dyn_cast(user)) { - if (hasSameStaticTensorType(sigmoidOp.getResult(), currentValue.getType())) - step = DistributedTensorStep { - user, DistributedTensorOpKind::Sigmoid, {}, DistributedTensorConstantKind::None, true, std::nullopt}; - else - failureDetail = "sigmoid result type mismatch"; - } - else if (auto addOp = dyn_cast(user)) { - step = classifyDistributedBinaryConsumer( - user, currentValue, addOp.getA(), addOp.getB(), DistributedTensorOpKind::Add, /*allowFragmentOnRhs=*/true, - failureDetail); - } - else if (auto subOp = dyn_cast(user)) { - step = classifyDistributedBinaryConsumer( - user, currentValue, subOp.getA(), subOp.getB(), DistributedTensorOpKind::Sub, /*allowFragmentOnRhs=*/true, - failureDetail); - } - else if (auto mulOp = dyn_cast(user)) { - step = classifyDistributedBinaryConsumer( - user, currentValue, mulOp.getA(), mulOp.getB(), DistributedTensorOpKind::Mul, /*allowFragmentOnRhs=*/true, - failureDetail); - } - else if (auto divOp = dyn_cast(user)) { - step = classifyDistributedBinaryConsumer( - user, currentValue, divOp.getA(), divOp.getB(), DistributedTensorOpKind::Div, /*allowFragmentOnRhs=*/false, - failureDetail); - if (step) { - auto denseAttr = dyn_cast(step->constantAttr); - if (!denseAttr) { - failureDetail = "div requires floating-point splat constant"; - step.reset(); - } - } - } - else if (auto nextConv = dyn_cast(user)) { - failureDetail = "onnx.Conv distributed consumer blocked by dim0-only whole-batch materialization in MergeComputeNodes"; - } - else { - failureDetail = (user->getName().getStringRef() + " is not distributed-aware yet").str(); - } - - if (!step) { - analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer; - analysis.barrierDetail = failureDetail; - return analysis; - } - - analysis.replacementOp = user; - analysis.steps.push_back(*step); - currentValue = user->getResult(0); - } -} - -static void rewriteDistributedConvReport(const DistributedConvReportTotals& totals) { - std::fstream reportFile = openReportFile("conv_distributed_consumption_report"); - if (!reportFile.is_open()) - return; - - reportFile << "# PIM Conv Distributed Consumption Report\n\n"; - reportFile << "Totals:\n"; - reportFile << "- convs_seen: " << totals.totalConvs << "\n"; - reportFile << "- distributed_tensors_created: " << totals.distributedTensorsCreated << "\n"; - reportFile << "- distributed_values_propagated: " << totals.distributedValuesPropagated << "\n"; - reportFile << "- distributed_consumers_handled_locally: " << totals.distributedConsumersHandled << "\n"; - reportFile << "- distributed_conv_inputs_seen: " << totals.distributedConvInputsSeen << "\n"; - reportFile << "- distributed_conv_inputs_consumed: " << totals.distributedConvInputsConsumed << "\n"; - reportFile << "- materialization_barriers_inserted: " << totals.materializationBarriersInserted << "\n\n"; - - if (!totals.barrierReasons.empty()) { - reportFile << "Materialization barriers:\n"; - for (const auto& [reason, count] : totals.barrierReasons) - reportFile << "- " << reason << ": " << count << "\n"; - reportFile << "\n"; - } - - if (!totals.fallbackReasons.empty()) { - reportFile << "Fallback / no-distribution reasons:\n"; - for (const auto& [reason, count] : totals.fallbackReasons) - reportFile << "- " << reason << ": " << count << "\n"; - reportFile << "\n"; - } - - if (!totals.chains.empty()) { - reportFile << "Chains:\n"; - for (const DistributedChainReportEntry& chain : totals.chains) { - reportFile << "- chain_id: " << chain.chainId << "\n"; - reportFile << " chain_length: " << chain.chainLength << "\n"; - reportFile << " producer_kind: " << chain.producerKind << "\n"; - reportFile << " distributed_ops: " << chain.distributedOps << "\n"; - reportFile << " materialization_points: " << chain.materializationPoints << "\n"; - reportFile << " first_materialization_reason: " << chain.firstMaterializationReason << "\n"; - reportFile << " max_live_fragments: " << chain.maxLiveFragments << "\n"; - reportFile << " max_fragment_fanout: " << chain.maxFragmentFanout << "\n"; - reportFile << " patch_builder_core_count: " << chain.patchBuilderCoreCount << "\n"; - reportFile << " central_junction_detected: " << chain.centralJunctionDetected << "\n"; - reportFile << " conv_input_materialization_kind: " << chain.convInputMaterializationKind << "\n"; - reportFile << " local_patch_fragments: " << chain.localPatchFragments << "\n"; - reportFile << " remote_patch_fragments: " << chain.remotePatchFragments << "\n"; - reportFile << " halo_transfer_count: " << chain.haloTransferCount << "\n"; - reportFile << " grouped_transfer_count: " << chain.groupedTransferCount << "\n"; - if (!chain.fallbackReason.empty()) - reportFile << " fallback_reason: " << chain.fallbackReason << "\n"; - } - } -} - -[[maybe_unused]] static void recordDistributedConvOutcome(const DistributedConvAnalysis& analysis) { - static std::mutex reportMutex; - static DistributedConvReportTotals totals; - - std::string barrierKey = stringifyDistributedConvBarrierKind(analysis.barrierKind).str(); - if (!analysis.barrierDetail.empty()) - barrierKey += ": " + analysis.barrierDetail; - - std::lock_guard guard(reportMutex); - totals.totalConvs++; - if (analysis.hasLocalConsumers()) { - totals.distributedTensorsCreated++; - totals.distributedValuesPropagated += analysis.steps.size(); - totals.distributedConsumersHandled += llvm::count_if(analysis.steps, [](const DistributedTensorStep& step) { - return step.kind != DistributedTensorOpKind::Conv; - }); - totals.distributedConvInputsSeen += llvm::count_if(analysis.steps, [](const DistributedTensorStep& step) { - return step.kind == DistributedTensorOpKind::Conv; - }); - totals.distributedConvInputsConsumed += llvm::count_if(analysis.steps, [](const DistributedTensorStep& step) { - return step.kind == DistributedTensorOpKind::Conv; - }); - totals.materializationBarriersInserted++; - totals.barrierReasons[barrierKey]++; - } - else { - totals.fallbackReasons[barrierKey]++; - } - DistributedChainReportEntry chain; - chain.chainId = totals.totalConvs; - chain.chainLength = analysis.steps.size() + 1; - chain.producerKind = "Conv"; - chain.materializationPoints = stringifyDistributedConvBarrierKind(analysis.barrierKind).str(); - chain.firstMaterializationReason = analysis.barrierDetail; - chain.maxFragmentFanout = analysis.barrierKind == DistributedConvBarrierKind::Fanout ? 2 : 1; - std::string ops; - for (size_t index = 0; index < analysis.steps.size(); ++index) { - if (!ops.empty()) - ops += ", "; - ops += stringifyDistributedTensorOpKind(analysis.steps[index].kind); - } - chain.distributedOps = ops; - if (analysis.hasDistributedConvConsumer()) { - chain.convInputMaterializationKind = "distributed_with_halo_exchange"; - chain.patchBuilderCoreCount = 1; - } - if (totals.chains.size() == 16) - totals.chains.erase(totals.chains.begin()); - totals.chains.push_back(std::move(chain)); - rewriteDistributedConvReport(totals); -} - -static std::string makeDivider(ArrayRef widths) { - std::string divider = "+"; - for (size_t width : widths) { - divider.append(width + 2, '-'); - divider.push_back('+'); - } - return divider; -} - -static void printConvReportLegend(std::fstream& reportFile) { - reportFile << "# PIM Conv Lowering Report\n\n"; - reportFile << "Legend:\n"; - reportFile << "- `id`: sequential Conv report entry index within this compiler invocation.\n"; - reportFile << "- `where`: summarized MLIR location of the Conv op.\n"; - reportFile << "- `mode`: whether the selected strategy came from `auto` policy or a forced compiler option.\n"; - reportFile << "- `strategy`: selected Conv lowering algorithm.\n"; - reportFile << "- `input`, `weight`, `output`: tensor shapes of the Conv operands/result.\n"; - reportFile << "- `groups`: ONNX Conv group count.\n"; - reportFile << "- `K`: logical reduction size, `CinPerGroup * Kh * Kw`.\n"; - reportFile << "- `C`: logical output-channel width handled by the selected strategy.\n"; - reportFile << "- `P`: total output positions, `N * Hout * Wout`.\n"; - reportFile << "- `X`: crossbar size.\n"; - reportFile << "- `pack`: packed spatial positions per MVM group, `floor(X / max(K, C))`.\n"; - reportFile << "- `im2col`: total explicit im2col element count, `P * K`.\n"; - reportFile << "- `im2col_budget`: maximum im2col element budget allowed by the compiler option.\n"; - reportFile << "- `stream_chunk`: output positions materialized per streamed chunk, or `-` when not applicable.\n"; - reportFile << "- `batch_size`: logical batch size passed into the compute lowering for this Conv form.\n"; - reportFile << "- `batches`: number of repeated compute batches emitted for the Conv.\n"; - reportFile << "- `spatial_compute_batch`: whether ONNX-to-Spatial used `spat.compute_batch` for this Conv.\n"; - reportFile << "- `batched_instruction_emission`: whether the lowering is expected to reach batched PIM emission.\n"; - reportFile << "- `reason`: strategy-selection reason.\n"; - reportFile << "- Per-conv details below the table include profitability estimates and materialization diagnostics.\n"; - reportFile << "- placeholders like `[7]`: value too long for the table cell; see the appendix at the end.\n\n"; -} - -struct ConvReportOverflow { - uint64_t placeholderId; - SmallVector rowIds; - std::string column; - std::string value; -}; - -static StringRef describeConvLoweringStrategy(PimConvLoweringType strategy) { - switch (strategy) { - case PimConvLoweringAuto: return "Automatic policy selection."; - case PimConvLoweringLegacy: return "Legacy Conv lowering path kept for compatibility."; - case PimConvLoweringDepthwise: return "Specialized depthwise lowering that avoids generic cross-channel GEMM mixing."; - case PimConvLoweringPackedIm2Col: return "Explicit im2col plus packed GEMM lowering for Conv shapes that fit well in one crossbar."; - case PimConvLoweringStreamedPatch: return "Chunked per-patch streaming Conv lowering without global im2col materialization."; - case PimConvLoweringStreamedPacked: return "Chunked streamed Conv lowering that still packs multiple output positions per MVM group."; - case PimConvLoweringOutputChannelTiled: return "Conv lowering that splits output channels across tiles when C exceeds one crossbar width."; - case PimConvLoweringInputKTiled: return "Conv lowering that splits the reduction dimension K across tiles and accumulates partial sums."; - case PimConvLoweringTiled2D: return "Conv lowering that tiles both K and output channels because neither dimension fits one crossbar."; - } - llvm_unreachable("unknown conv lowering strategy"); -} - -static std::string fitConvReportCell(StringRef text, - size_t width, - uint64_t rowId, - StringRef column, - std::vector& overflows, - uint64_t& nextPlaceholderId, - bool rightAlign = false) { - if (text.size() <= width) - return alignCell(text, width, rightAlign); - - for (ConvReportOverflow& overflow : overflows) { - if (overflow.column == column && overflow.value == text) { - if (llvm::find(overflow.rowIds, rowId) == overflow.rowIds.end()) - overflow.rowIds.push_back(rowId); - std::string placeholder = "[" + std::to_string(overflow.placeholderId) + "]"; - return alignCell(placeholder, width, rightAlign); - } - } - - std::string placeholder = "[" + std::to_string(nextPlaceholderId++) + "]"; - overflows.push_back({nextPlaceholderId - 1, {rowId}, column.str(), text.str()}); - return alignCell(placeholder, width, rightAlign); -} - -static void writeConvReportTable(std::fstream& reportFile, ArrayRef entries) { - static constexpr size_t kIdWidth = 4; - static constexpr size_t kWhereWidth = 24; - static constexpr size_t kModeWidth = 6; - static constexpr size_t kStrategyWidth = 20; - static constexpr size_t kShapeWidth = 14; - static constexpr size_t kGroupsWidth = 3; - static constexpr size_t kSmallWidth = 5; - static constexpr size_t kPWidth = 10; - static constexpr size_t kIm2colWidth = 10; - static constexpr size_t kChunkWidth = 8; - static constexpr size_t kFlagWidth = 3; - static constexpr size_t kReasonWidth = 24; - - const SmallVector widths = { - kIdWidth, kWhereWidth, kModeWidth, kStrategyWidth, kShapeWidth, kShapeWidth, kShapeWidth, kGroupsWidth, - kSmallWidth, kSmallWidth, kPWidth, kSmallWidth, kSmallWidth, kIm2colWidth, kIm2colWidth, - kChunkWidth, kSmallWidth, kSmallWidth, kFlagWidth, kFlagWidth, kReasonWidth, - }; - const std::string divider = makeDivider(widths); - std::vector overflows; - uint64_t nextPlaceholderId = 1; - - auto printRow = [&](ArrayRef cells) { - reportFile << "|"; - for (size_t i = 0; i < cells.size(); ++i) - reportFile << " " << cells[i] << " |"; - reportFile << "\n"; - }; - - reportFile << divider << "\n"; - printRow({ - alignCell("id", kIdWidth, true), - alignCell("where", kWhereWidth), - alignCell("mode", kModeWidth), - alignCell("strategy", kStrategyWidth), - alignCell("input", kShapeWidth), - alignCell("weight", kShapeWidth), - alignCell("output", kShapeWidth), - alignCell("grp", kGroupsWidth, true), - alignCell("K", kSmallWidth, true), - alignCell("C", kSmallWidth, true), - alignCell("P", kPWidth, true), - alignCell("X", kSmallWidth, true), - alignCell("pack", kSmallWidth, true), - alignCell("im2col", kIm2colWidth, true), - alignCell("budget", kIm2colWidth, true), - alignCell("chunk", kChunkWidth, true), - alignCell("batch", kSmallWidth, true), - alignCell("nbat", kSmallWidth, true), - alignCell("scb", kFlagWidth), - alignCell("bie", kFlagWidth), - alignCell("reason", kReasonWidth), - }); - reportFile << divider << "\n"; - - for (const ConvReportEntry& entry : entries) { - printRow({ - alignCell(std::to_string(entry.id), kIdWidth, true), - fitConvReportCell(entry.where, kWhereWidth, entry.id, "where", overflows, nextPlaceholderId), - alignCell(entry.mode, kModeWidth), - fitConvReportCell(entry.strategy, kStrategyWidth, entry.id, "strategy", overflows, nextPlaceholderId), - fitConvReportCell(entry.inputShape, kShapeWidth, entry.id, "input", overflows, nextPlaceholderId), - fitConvReportCell(entry.weightShape, kShapeWidth, entry.id, "weight", overflows, nextPlaceholderId), - fitConvReportCell(entry.outputShape, kShapeWidth, entry.id, "output", overflows, nextPlaceholderId), - alignCell(std::to_string(entry.groups), kGroupsWidth, true), - alignCell(std::to_string(entry.k), kSmallWidth, true), - alignCell(std::to_string(entry.c), kSmallWidth, true), - alignCell(std::to_string(entry.p), kPWidth, true), - alignCell(std::to_string(entry.xbarSize), kSmallWidth, true), - alignCell(std::to_string(entry.pack), kSmallWidth, true), - alignCell(std::to_string(entry.im2colElements), kIm2colWidth, true), - alignCell(std::to_string(entry.im2colBudget), kIm2colWidth, true), - fitConvReportCell(entry.chunkText, kChunkWidth, entry.id, "stream_chunk", overflows, nextPlaceholderId, true), - alignCell(std::to_string(entry.batchSize), kSmallWidth, true), - alignCell(std::to_string(entry.numberOfBatches), kSmallWidth, true), - alignCell(entry.spatialComputeBatch, kFlagWidth), - alignCell(entry.batchedInstructionEmission, kFlagWidth), - fitConvReportCell(entry.reason, kReasonWidth, entry.id, "reason", overflows, nextPlaceholderId), - }); - } - reportFile << divider << "\n"; - - if (overflows.empty()) - reportFile << "\n"; - else { - reportFile << "\nAppendix:\n"; - for (const ConvReportOverflow& overflow : overflows) { - reportFile << " [" << overflow.placeholderId << "] rows "; - for (size_t i = 0; i < overflow.rowIds.size(); ++i) { - if (i != 0) - reportFile << ", "; - reportFile << overflow.rowIds[i]; - } - reportFile << ", " << overflow.column << ": " << overflow.value << "\n"; - } - reportFile << "\n"; - } - - reportFile << "Per-Conv Details:\n"; - for (const ConvReportEntry& entry : entries) { - reportFile << "- Conv " << entry.id << ": mode=" << entry.mode << ", strategy=" << entry.strategy - << ", reason=" << entry.reason << "\n"; - reportFile << " K=" << entry.k << ", Cout=" << entry.c << ", output_positions=" << entry.p - << ", estimated_mvm_count=" << entry.estimatedMvmCount - << ", estimated_reduction_vadd_count=" << entry.estimatedReductionVAddCount - << ", estimated_output_fragments=" << entry.estimatedOutputFragments << "\n"; - reportFile << " materialization_required_at_func_return=" << entry.materializationRequiredAtReturn - << ", materialization_kind=" << entry.materializationKind - << ", giant_collector_concat_expected=" << entry.giantCollectorConcatExpected - << ", concat_operand_count=" << entry.concatOperandCount - << ", collector_core=" << entry.collectorCore - << ", full_input_broadcast_expected=" << entry.fullInputBroadcastExpected << "\n"; - if (!entry.rejectedAutoStrategy.empty()) - reportFile << " rejected_auto_strategy=" << entry.rejectedAutoStrategy << "\n"; - if (!entry.fallbackReason.empty()) - reportFile << " fallback_reason=" << entry.fallbackReason << "\n"; - } - reportFile << "\n"; - - llvm::SmallVector usedStrategies; - for (const ConvReportEntry& entry : entries) { - PimConvLoweringType strategy = PimConvLoweringAuto; - for (PimConvLoweringType candidate : { - PimConvLoweringAuto, - PimConvLoweringLegacy, - PimConvLoweringDepthwise, - PimConvLoweringPackedIm2Col, - PimConvLoweringStreamedPatch, - PimConvLoweringStreamedPacked, - PimConvLoweringOutputChannelTiled, - PimConvLoweringInputKTiled, - PimConvLoweringTiled2D, - }) { - if (entry.strategy == stringifyConvLoweringStrategy(candidate)) { - strategy = candidate; - break; - } - } - if (llvm::find(usedStrategies, strategy) == usedStrategies.end()) - usedStrategies.push_back(strategy); - } - - reportFile << "Strategies used in this report:\n"; - for (PimConvLoweringType strategy : usedStrategies) - reportFile << "- `" << stringifyConvLoweringStrategy(strategy).str() << "`: " - << describeConvLoweringStrategy(strategy).str() << "\n"; -} - -static void rewriteConvLoweringReport(ArrayRef entries) { - std::fstream reportFile = openReportFile("conv_lowering_report"); - if (!reportFile.is_open()) - return; - printConvReportLegend(reportFile); - writeConvReportTable(reportFile, entries); -} - -[[maybe_unused]] static FailureOr resolveRequestedConvLoweringStrategy(ONNXConvOp convOp) { - if (!useExperimentalConvImpl) - return pimConvLowering.getValue(); - - if (pimConvLowering != PimConvLoweringAuto && pimConvLowering != PimConvLoweringPackedIm2Col) { - convOp.emitOpError() << "--use-experimental-conv-impl conflicts with --pim-conv-lowering=" - << stringifyConvLoweringStrategy(pimConvLowering); - return failure(); - } - return PimConvLoweringPackedIm2Col; -} - -static ConvLoweringDecision chooseConvLoweringStrategy(const ConvGeometry& geo, - PimConvLoweringType requested, - const DistributedConvAnalysis& analysis) { +static PimConvLoweringType chooseConvLoweringStrategy(const ConvGeometry& geo, + PimConvLoweringType requested) { if (requested != PimConvLoweringAuto) - return {requested, "forced by compiler option", /*isAuto=*/false, "", ""}; + return requested; // Transform-based convolution is intentionally not selected for this ISA: // it would require explicit transform sequences and staging traffic on top of // the same crossbar MVM primitive, which is not attractive here. if (geo.isDepthwise) - return {PimConvLoweringDepthwise, "depthwise convolution", /*isAuto=*/true, "", ""}; + return PimConvLoweringDepthwise; if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 && geo.im2colElements <= pimConvIm2colMaxElements) - return {PimConvLoweringPackedIm2Col, - "fits crossbar, packing useful, and global im2col fits budget", - /*isAuto=*/true, - "", - ""}; + return PimConvLoweringPackedIm2Col; if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 && geo.im2colElements > pimConvIm2colMaxElements) - return {PimConvLoweringStreamedPacked, - "fits crossbar and packing useful, but global im2col exceeds budget", - /*isAuto=*/true, - "", - ""}; + return PimConvLoweringStreamedPacked; if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize) - return {PimConvLoweringStreamedPatch, "fits crossbar but packing is not useful", /*isAuto=*/true, "", ""}; + return PimConvLoweringStreamedPatch; if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize) - return {PimConvLoweringOutputChannelTiled, - "output channels exceed one crossbar width", - /*isAuto=*/true, - "", - ""}; - if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize) { - ConvStrategyEstimate estimate = estimateConvStrategy(geo, PimConvLoweringInputKTiled, analysis); - std::string fallbackReason = "auto rejects input-k-tiled because the reduction-heavy path is force-only for now"; - if (estimate.requiresFuncReturnMaterialization && estimate.perOutputPositionReduction - && estimate.giantCollectorConcatExpected) { - fallbackReason += "; func.return would materialize " + std::to_string(estimate.concatOperandCount) - + " output fragments through a single collector concat after per-position reductions"; - } - if (estimate.fullInputBroadcastExpected) - fallbackReason += "; the current lowering also broadcasts the padded input to many workers"; - return {PimConvLoweringLegacy, - "fall back to legacy explicit-im2col for the current auto policy", - /*isAuto=*/true, - fallbackReason, - stringifyConvLoweringStrategy(PimConvLoweringInputKTiled).str()}; - } - return {PimConvLoweringTiled2D, "both reduction K and output channels exceed one crossbar", /*isAuto=*/true, "", ""}; + return PimConvLoweringOutputChannelTiled; + if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize) + return PimConvLoweringLegacy; + return PimConvLoweringTiled2D; } -[[maybe_unused]] static LogicalResult verifyForcedConvLoweringStrategy(ONNXConvOp convOp, - const ConvGeometry& geo, - PimConvLoweringType strategy) { - switch (strategy) { - case PimConvLoweringAuto: - case PimConvLoweringLegacy: - return success(); - case PimConvLoweringDepthwise: - if (geo.isDepthwise) - return success(); - return convOp.emitOpError("forced depthwise Conv lowering requires a depthwise convolution"); - case PimConvLoweringPackedIm2Col: - if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 && geo.im2colElements <= pimConvIm2colMaxElements) - return success(); - return convOp.emitOpError("forced packed-im2col Conv lowering requires K/C to fit, pack >= 2, and im2col within budget"); - case PimConvLoweringStreamedPatch: - if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize) - return success(); - return convOp.emitOpError("forced streamed-patch Conv lowering requires K and C to each fit one crossbar"); - case PimConvLoweringStreamedPacked: - if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2) - return success(); - return convOp.emitOpError("forced streamed-packed Conv lowering requires K/C to fit and pack >= 2"); - case PimConvLoweringOutputChannelTiled: - if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize) - return success(); - return convOp.emitOpError("forced output-channel-tiled Conv lowering requires K <= X and C > X"); - case PimConvLoweringInputKTiled: - if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize) - return success(); - return convOp.emitOpError("forced input-k-tiled Conv lowering requires K > X and C <= X"); - case PimConvLoweringTiled2D: - if (geo.k > geo.xbarSize && geo.c > geo.xbarSize) - return success(); - return convOp.emitOpError("forced tiled-2d Conv lowering requires K > X and C > X"); - } - llvm_unreachable("unknown conv lowering strategy"); -} - -static void reportConvLoweringDecision(ONNXConvOp convOp, - const ConvGeometry& geo, - const ConvLoweringDecision& decision, - const ConvStrategyEstimate& estimate, - int64_t batchSize, - int64_t numberOfBatches, - bool usesComputeBatch, - bool usesBatchedInstructionEmission, - std::optional streamChunkPositions = std::nullopt) { - if (!pimReportConvLowering) - return; - - const std::string location = summarizeLocation(convOp.getLoc()); - const std::string strategy = stringifyConvLoweringStrategy(decision.strategy).str(); - const std::string mode = decision.isAuto ? "auto" : "forced"; - const std::string inputShape = formatShape(cast(convOp.getX().getType()).getShape()); - const std::string weightShape = formatShape(cast(convOp.getW().getType()).getShape()); - const std::string outputShape = formatShape(cast(convOp.getY().getType()).getShape()); - const std::string chunkText = streamChunkPositions ? std::to_string(*streamChunkPositions) : "-"; - const std::string scbText = usesComputeBatch ? "yes" : "no"; - const std::string bieText = usesBatchedInstructionEmission ? "yes" : "no"; - - static uint64_t reportIndex = 0; - const uint64_t currentIndex = ++reportIndex; - static std::mutex reportMutex; - static std::vector reportEntries; - std::lock_guard lock(reportMutex); - reportEntries.push_back({ - currentIndex, - location, - strategy, - mode, - inputShape, - weightShape, - outputShape, - geo.group, - geo.k, - geo.c, - geo.p, - geo.xbarSize, - geo.pack, - geo.im2colElements, - pimConvIm2colMaxElements, - chunkText, - batchSize, - numberOfBatches, - scbText, - bieText, - decision.reason, - decision.fallbackReason, - decision.rejectedAutoStrategy, - estimate.estimatedMvmCount, - estimate.estimatedReductionVAddCount, - estimate.estimatedOutputFragments, - estimate.requiresFuncReturnMaterialization ? "yes" : "no", - estimate.materializationKind, - estimate.concatOperandCount, - estimate.collectorCore, - estimate.giantCollectorConcatExpected ? "yes" : "no", - estimate.fullInputBroadcastExpected ? "yes" : "no", - }); - rewriteConvLoweringReport(reportEntries); -} static Value expandBiasIfNeeded(Value bias, PatternRewriter& rewriter, Location loc) { auto biasType = cast(bias.getType()); @@ -1184,7 +192,6 @@ static Value createCollectedConvOutput(ValueRange gemmRows, int64_t numPatches, int64_t numChannelsOut, int64_t packFactor, - ArrayRef distributedConsumers, PatternRewriter& rewriter, Location loc); static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b); @@ -1241,10 +248,12 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr, RankedTensorType wType, const Tiling& tiling, PatternRewriter& rewriter, - Location loc) { + Location loc, + int64_t paddedInputRows = -1) { const int64_t paddedOutputChannels = static_cast(crossbarSize.getValue()); + const int64_t packedInputRows = paddedInputRows > 0 ? paddedInputRows : tiling.tileInputRows; auto packedWeightType = RankedTensorType::get( - {tiling.numChannelTiles, tiling.tileInputRows, paddedOutputChannels}, wType.getElementType()); + {tiling.numChannelTiles, packedInputRows, paddedOutputChannels}, wType.getElementType()); SmallVector packedValues(packedWeightType.getNumElements(), cast(rewriter.getZeroAttr(wType.getElementType()))); SmallVector sourceValues(wDenseAttr.getValues()); @@ -1263,7 +272,7 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr, ((globalOutChannel * wType.getDimSize(1) * wType.getDimSize(2)) + kernelH) * wType.getDimSize(3) + kernelW; const int64_t targetCol = localChannel * tiling.outputMultiplier + multiplierIndex; const int64_t targetFlatIndex = - ((tileIndex * tiling.tileInputRows) + targetRow) * paddedOutputChannels + targetCol; + ((tileIndex * packedInputRows) + targetRow) * paddedOutputChannels + targetCol; packedValues[targetFlatIndex] = sourceValues[sourceFlatIndex]; } } @@ -1703,7 +712,6 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& tiling->totalPatches, state.outType.getDimSize(1), /*packFactor=*/1, - {}, rewriter, loc); } @@ -1952,7 +960,6 @@ static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sour } static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, - ArrayRef distributedConsumers, PatternRewriter& rewriter, Location loc) { PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc); @@ -2139,7 +1146,7 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, elementType); return createCollectedConvOutput( chunkRows, state.outType, cast(chunkRows.front().getType()), nhwcType, state.outType, totalPatches, - state.numChannelsOut, /*packFactor=*/1, distributedConsumers, rewriter, loc); + state.numChannelsOut, /*packFactor=*/1, rewriter, loc); } static Value buildPackedWeights(DenseElementsAttr wDenseAttr, @@ -2434,7 +1441,6 @@ static Value createStreamedConvRows(const ConvLoweringState& state, } static Value rewritePackedIm2ColConv(const ConvLoweringState& state, - ArrayRef distributedConsumers, PatternRewriter& rewriter, Location loc) { auto wDenseAttr = getHostConstDenseElementsAttr(state.w); @@ -2480,13 +1486,11 @@ static Value rewritePackedIm2ColConv(const ConvLoweringState& state, plan.chunkNumPatches, state.numChannelsOut, plan.effectiveMaxParallelPixels, - distributedConsumers, rewriter, loc); } static Value rewriteStreamedConv(const ConvLoweringState& state, - ArrayRef distributedConsumers, PatternRewriter& rewriter, Location loc, int64_t forcedPackFactor) { @@ -2516,38 +1520,11 @@ static Value rewriteStreamedConv(const ConvLoweringState& state, state.outType.getElementType()); return createCollectedConvOutput( ValueRange {collectedRows}, state.outType, gemmOutType, nhwcType, state.outType, gemmOutType.getDimSize(0), - state.numChannelsOut, /*packFactor=*/1, distributedConsumers, rewriter, loc); + state.numChannelsOut, /*packFactor=*/1, rewriter, loc); } } // namespace standard -static SmallVector buildRowStripFragments(RankedTensorType tensorType) { - SmallVector fragments; - 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 { - {offsets.begin() + row * rank, offsets.begin() + (row + 1) * rank}, - {sizes.begin() + row * rank, sizes.begin() + (row + 1) * rank}, - {1, 1, 1, 1}, - row, - }); - } - return fragments; -} - -static DistributedTensorInfo makeDistributedTensorInfo(Value storage, RankedTensorType logicalType) { - DistributedTensorInfo info; - info.storage = storage; - info.logicalType = logicalType; - info.fragments = buildRowStripFragments(logicalType); - info.laneCount = logicalType.getDimSize(2); - info.channels = logicalType.getDimSize(1); - info.height = logicalType.getDimSize(2); - info.width = logicalType.getDimSize(3); - return info; -} static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter) { auto zeroAttr = DenseElementsAttr::get(gemmResultType, rewriter.getZeroAttr(gemmResultType.getElementType())); @@ -2817,55 +1794,50 @@ static Value extractDenseConvWindowRow(Value denseInput, rewriter, loc, fragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})); } -static FailureOr createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) { - auto elementType = state.xType.getElementType(); - auto floatType = dyn_cast(elementType); - if (!floatType) - return failure(); - +static Value createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) { + auto elementType = cast(state.xType.getElementType()); Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - auto tableType = RankedTensorType::get({state.outHeight * state.wHeight, 1, state.xWidth, state.numChannelsIn}, - elementType, - state.xType.getEncoding()); - Attribute zero = rewriter.getZeroAttr(elementType); - Attribute one = rewriter.getFloatAttr(floatType, 1.0); + auto tableType = RankedTensorType::get( + {2, 1, state.xWidth, state.numChannelsIn}, elementType, state.xType.getEncoding()); + SmallVector values(tableType.getNumElements(), rewriter.getZeroAttr(elementType)); + std::fill(values.begin() + tableType.getNumElements() / 2, values.end(), rewriter.getFloatAttr(elementType, 1.0)); + return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType); +} + +static Value createRowStripWindowMaskIndexTable(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 outputRow = 0; outputRow < state.outHeight; ++outputRow) for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) { int64_t sourceRow = outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin; - Attribute value = (sourceRow < 0 || sourceRow >= state.xHeight) ? zero : one; - for (int64_t width = 0; width < state.xWidth; ++width) - for (int64_t channel = 0; channel < state.numChannelsIn; ++channel) - values.push_back(value); + values.push_back(rewriter.getIndexAttr(sourceRow >= 0 && sourceRow < state.xHeight)); } - } - return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType); } static Value extractProjectedRowStripWindowMask(Value maskTable, + Value maskIndexTable, const ConvLoweringState& state, Value outputHeight, Value kernelRow, PatternRewriter& rewriter, Location loc) { Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc); + Value maskIndex = tensor::ExtractOp::create(rewriter, loc, maskIndexTable, ValueRange {tableIndex}).getResult(); auto fragmentType = getRowStripFragmentType(state.xType); - SmallVector offsets { - tableIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector sizes {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.xWidth), - rewriter.getIndexAttr(state.numChannelsIn)}; - return tensor::ExtractSliceOp::create(rewriter, - loc, - fragmentType, - maskTable, - offsets, - sizes, - getUnitStrides(rewriter, 4)); + return tensor::ExtractSliceOp::create( + rewriter, + loc, + fragmentType, + maskTable, + SmallVector { + maskIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(state.xWidth), rewriter.getIndexAttr(state.numChannelsIn)}, + getUnitStrides(rewriter, 4)); } static FailureOr createConvInputWindow(Value input, @@ -2876,20 +1848,23 @@ static FailureOr createConvInputWindow(Value input, auto fragmentType = getRowStripFragmentType(state.xType); auto inputType = dyn_cast(input.getType()); const bool denseInput = inputType == state.xType; - if (!denseInput && failed(describeRowStripPhysicalValue(input, state.xType))) + FailureOr physicalInput = describeRowStripPhysicalValue(input, state.xType); + if (!denseInput && failed(physicalInput)) return failure(); + if (!denseInput && physicalInput->tilesPerRow == 1 && state.wHeight == 1 && state.wWidth == 1 + && state.strideHeight == 1 && state.strideWidth == 1 && state.padHeightBegin == 0 + && state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0) + return extractGraphBatchPhysicalFragment( + rewriter, loc, input, outputHeight, physicalInput->fragmentType); auto paddedWindowType = RankedTensorType::get( {1, state.wHeight, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn}, state.xType.getElementType(), state.xType.getEncoding()); - FailureOr physicalInput = - denseInput ? FailureOr(failure()) : describeRowStripPhysicalValue(input, state.xType); Value sourceIndexTable = denseInput ? createRowStripWindowSourceRowTable(state, rewriter) : createRowStripWindowSourceSlotTable(state, physicalInput->tilesPerRow, rewriter); - FailureOr maskTable = createRowStripWindowMaskTable(state, rewriter); - if (failed(maskTable)) - return failure(); + Value maskTable = createRowStripWindowMaskTable(state, rewriter); + Value maskIndexTable = createRowStripWindowMaskIndexTable(state, rewriter); Value initWindow = createZeroTensorConstant(paddedWindowType, rewriter); Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); @@ -2914,7 +1889,8 @@ static FailureOr createConvInputWindow(Value input, return failure(); Value semanticRow = *sourceRow; if (state.padHeightBegin != 0 || state.padHeightEnd != 0) { - Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, rowLoc); + Value mask = extractProjectedRowStripWindowMask( + maskTable, maskIndexTable, state, outputHeight, kernelRow, rewriter, rowLoc); semanticRow = spatial::SpatVMulOp::create(rewriter, rowLoc, fragmentType, semanticRow, mask).getResult(); } Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, rowLoc); @@ -3632,126 +2608,269 @@ static FailureOr createPointwiseOutputFromRowStripFragments(Value rowStri return batch->getResult(0); } +static bool canConsumeDepthwiseRowStrip(const ConvLoweringState& state) { + if (state.batchSize != 1 || state.group != state.numChannelsIn + || state.dilationHeight != 1 || state.dilationWidth != 1 + || !isa(state.xType.getElementType()) + || !getHostConstDenseElementsAttr(state.w) + || (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType))) + return false; + auto tiling = depthwise::computeTiling(state.batchSize, + state.numChannelsIn, + state.numChannelsOut, + state.wHeight, + state.wWidth, + state.outHeight, + state.outWidth); + return tiling && tiling->numChannelTiles <= static_cast(crossbarCountInCore.getValue()); +} + +static Value insertDepthwiseInputSegment(Value inputWindow, + Value scratch, + Value tileIndex, + Value kernelRow, + Value sourceWidth, + Value scratchOffset, + int64_t inputChannel, + const depthwise::Tiling& tiling, + PatternRewriter& rewriter, + Location loc) { + auto inputWindowType = cast(inputWindow.getType()); + auto inputPixelType = RankedTensorType::get( + {1, 1, 1, tiling.channelsPerTile}, inputWindowType.getElementType(), inputWindowType.getEncoding()); + Value inputPixel = tensor::ExtractSliceOp::create( + rewriter, + loc, + inputPixelType, + inputWindow, + SmallVector {rewriter.getIndexAttr(0), kernelRow, sourceWidth, + rewriter.getIndexAttr(inputChannel)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling.channelsPerTile)}, + getUnitStrides(rewriter, 4)); + return tensor::InsertSliceOp::create( + rewriter, + loc, + inputPixel, + scratch, + SmallVector {tileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), scratchOffset}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(tiling.channelsPerTile)}, + getUnitStrides(rewriter, 4)); +} + +static FailureOr assembleDepthwiseInputScratch(Value inputWindow, + Value scratch, + Value inputWidth, + const ConvLoweringState& state, + const depthwise::Tiling& tiling, + PatternRewriter& rewriter, + Location loc) { + Operation* anchor = rewriter.getInsertionBlock()->getParentOp(); + Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0); + Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1); + Value cKernelElements = getOrCreateIndexConstant(rewriter, anchor, tiling.kernelElements); + SmallVector tileIndices; + tileIndices.reserve(tiling.numChannelTiles); + for (int64_t tile = 0; tile < tiling.numChannelTiles; ++tile) + tileIndices.push_back(getOrCreateIndexConstant(rewriter, anchor, tile)); + auto kernelLoop = buildNormalizedScfFor( + rewriter, + loc, + c0, + cKernelElements, + c1, + ValueRange {scratch}, + [&](OpBuilder&, Location kernelLoc, Value kernelIndex, ValueRange iterArgs, + SmallVectorImpl& yielded) { + Value kernelRow = affineFloorDivConst(rewriter, kernelLoc, kernelIndex, state.wWidth, anchor); + Value kernelColumn = affineModConst(rewriter, kernelLoc, kernelIndex, state.wWidth, anchor); + Value sourceWidth = createOrFoldAffineApply( + rewriter, + kernelLoc, + getAffineDimExpr(0, rewriter.getContext()) + getAffineDimExpr(1, rewriter.getContext()), + ValueRange {inputWidth, kernelColumn}, + anchor); + Value scratchOffset = affineMulConst( + rewriter, kernelLoc, kernelIndex, tiling.channelsPerTile, anchor); + Value nextScratch = iterArgs.front(); + for (int64_t tile = 0; tile < tiling.numChannelTiles; ++tile) + nextScratch = insertDepthwiseInputSegment(inputWindow, + nextScratch, + tileIndices[tile], + kernelRow, + sourceWidth, + scratchOffset, + tile * tiling.channelsPerTile, + tiling, + rewriter, + kernelLoc); + yielded.push_back(nextScratch); + return success(); + }); + if (failed(kernelLoop)) + return failure(); + return kernelLoop->results.front(); +} + +static FailureOr createDepthwiseOutputFromRowStripFragments(Value rowStripStorage, + const ConvLoweringState& state, + PatternRewriter& rewriter, + Location loc) { + if (!canConsumeDepthwiseRowStrip(state) + || failed(describeRowStripPhysicalValue(rowStripStorage, state.xType))) + return failure(); + auto tiling = depthwise::computeTiling(state.batchSize, + state.numChannelsIn, + state.numChannelsOut, + state.wHeight, + state.wWidth, + state.outHeight, + state.outWidth); + auto weight = getHostConstDenseElementsAttr(state.w); + if (!tiling || !weight) + return failure(); + + Value packedWeights = depthwise::buildPackedWeights( + weight, state.wType, *tiling, rewriter, loc, static_cast(crossbarSize.getValue())); + Value bias = state.hasBias ? expandBiasIfNeeded(state.b, rewriter, loc) : Value(); + auto paddedOutputType = RankedTensorType::get( + {1, static_cast(crossbarSize.getValue())}, state.outType.getElementType()); + auto outputTileType = RankedTensorType::get( + {1, tiling->tileOutputChannels}, state.outType.getElementType()); + auto outputPixelType = RankedTensorType::get( + {1, 1, 1, tiling->tileOutputChannels}, state.outType.getElementType()); + auto fragmentType = getRowStripFragmentType(state.outType); + auto storageType = getRowStripStorageType(state.outType); + + auto batch = createSpatComputeBatch( + rewriter, + loc, + TypeRange {storageType}, + state.outHeight, + ValueRange {packedWeights}, + state.hasBias ? ValueRange {rowStripStorage, bias} : ValueRange {rowStripStorage}, + [&](detail::SpatComputeBatchBodyArgs args) { + Operation* anchor = rewriter.getInsertionBlock()->getParentOp(); + FailureOr inputWindow = + createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc); + if (failed(inputWindow)) + return failure(); + Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0); + Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1); + Value cOutWidth = getOrCreateIndexConstant(rewriter, anchor, state.outWidth); + const int64_t xbarDim = static_cast(crossbarSize.getValue()); + auto paddedInputScratchType = RankedTensorType::get( + {tiling->numChannelTiles, 1, 1, xbarDim}, state.xType.getElementType(), state.xType.getEncoding()); + auto tileScratchType = RankedTensorType::get( + {1, 1, 1, xbarDim}, state.xType.getElementType(), state.xType.getEncoding()); + auto vmmInputType = RankedTensorType::get( + {1, xbarDim}, state.xType.getElementType(), state.xType.getEncoding()); + Value zeroScratch = createZeroTensorConstant(paddedInputScratchType, rewriter); + Value fragment = tensor::EmptyOp::create( + rewriter, loc, fragmentType.getShape(), fragmentType.getElementType()); + SmallVector weightTiles; + SmallVector biasTiles; + SmallVector tileIndices; + SmallVector weightTileSizes { + rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)}; + for (int64_t tile = 0; tile < tiling->numChannelTiles; ++tile) { + Value tileIndex = getOrCreateIndexConstant(rewriter, anchor, tile); + tileIndices.push_back(tileIndex); + weightTiles.push_back(extractMixedSliceOrIdentity( + rewriter, + loc, + args.weights.front(), + RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType()), + {SmallVector {tileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + weightTileSizes, + getUnitStrides(rewriter, 3)})); + if (state.hasBias) + biasTiles.push_back(depthwise::createBiasTile(args.inputs[1], tileIndex, *tiling, rewriter, loc)); + } + auto widthLoop = buildNormalizedScfFor( + rewriter, loc, c0, cOutWidth, c1, ValueRange {fragment, zeroScratch}, + [&](OpBuilder&, Location widthLoc, Value width, ValueRange iterArgs, + SmallVectorImpl& yielded) { + Value next = iterArgs.front(); + Value scratch = iterArgs[1]; + // Valid prefix entries are overwritten per tile; the padded tail stays zero. + Value inputWidth = affineMulConst( + rewriter, widthLoc, width, state.strideWidth, anchor); + FailureOr nextScratch = assembleDepthwiseInputScratch( + *inputWindow, scratch, inputWidth, state, *tiling, rewriter, widthLoc); + if (failed(nextScratch)) + return failure(); + scratch = *nextScratch; + for (int64_t tile = 0; tile < tiling->numChannelTiles; ++tile) { + Value tileScratch = tensor::ExtractSliceOp::create( + rewriter, + widthLoc, + tileScratchType, + scratch, + SmallVector {tileIndices[tile], rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), + rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)}, + getUnitStrides(rewriter, 4)); + Value vmmInput = tensor::CollapseShapeOp::create( + rewriter, + widthLoc, + vmmInputType, + tileScratch, + SmallVector {{0, 1, 2}, {3}}); + Value output = spatial::SpatVMMOp::create( + rewriter, widthLoc, paddedOutputType, weightTiles[tile], vmmInput); + Value validOutput = tensor::ExtractSliceOp::create( + rewriter, + widthLoc, + outputTileType, + output, + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), + rewriter.getIndexAttr(tiling->tileOutputChannels)}, + getUnitStrides(rewriter, 2)); + if (state.hasBias) + validOutput = spatial::SpatVAddOp::create( + rewriter, widthLoc, outputTileType, validOutput, biasTiles[tile]); + Value pixel = tensor::ExpandShapeOp::create( + rewriter, widthLoc, outputPixelType, validOutput, + SmallVector {{0, 1, 2}, {3}}); + next = tensor::InsertSliceOp::create( + rewriter, + widthLoc, + pixel, + next, + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), width, + rewriter.getIndexAttr(tile * tiling->tileOutputChannels)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling->tileOutputChannels)}, + getUnitStrides(rewriter, 4)); + } + yielded.push_back(next); + yielded.push_back(scratch); + return success(); + }); + if (failed(widthLoop)) + return failure(); + publishGraphBatchPhysicalFragment(rewriter, loc, widthLoop->results.front(), args.outputs.front(), args.lane); + return success(); + }); + return failed(batch) ? FailureOr(failure()) : FailureOr(batch->getResult(0)); +} + static FailureOr createConvOutputFromRowStripInput(const ConvLoweringState& state, - [[maybe_unused]] const ConvLoweringDecision& decision, Value rowStripInput, + PimConvLoweringType strategy, PatternRewriter& rewriter, Location loc) { + if (strategy == PimConvLoweringDepthwise) + return createDepthwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc); if (state.xHeight == 1 && state.xWidth == 1 && state.wHeight == 1 && state.wWidth == 1) return createPointwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc); return createConvOutputFromPixelMajorRowStripFragments(rowStripInput, state, rewriter, loc); } -static Value createFragmentConstant(const DistributedTensorStep& step, - RankedTensorType fragmentType, - PatternRewriter& 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, - rewriter.getInsertionBlock()->getParentOp(), - DenseElementsAttr::get(fragmentType, splatValue), - fragmentType); -} - -static Value createFragmentReciprocalConstant(const DistributedTensorStep& step, - RankedTensorType fragmentType, - PatternRewriter& rewriter) { - SmallVector values; - if (step.constantKind == DistributedTensorConstantKind::PerChannel) { - auto denseType = cast(step.constantAttr.getType()); - SmallVector channelValues; - for (const APFloat& value : step.constantAttr.getValues()) - channelValues.push_back(value); - values.reserve(fragmentType.getNumElements()); - for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n) - for (int64_t h = 0; h < fragmentType.getDimSize(1); ++h) - for (int64_t w = 0; w < fragmentType.getDimSize(2); ++w) - for (int64_t channel = 0; channel < fragmentType.getDimSize(3); ++channel) { - APFloat reciprocal = channelValues[channel]; - APFloat one(reciprocal.getSemantics(), 1); - [[maybe_unused]] APFloat::opStatus status = one.divide(reciprocal, APFloat::rmNearestTiesToEven); - assert(!(status & APFloat::opInvalidOp) && "distributed conv div requires finite non-zero constant"); - values.push_back(one); - } - (void)denseType; - } - else { - APFloat reciprocal = cast(step.constantAttr).getSplatValue(); - APFloat one(reciprocal.getSemantics(), 1); - [[maybe_unused]] APFloat::opStatus status = one.divide(reciprocal, APFloat::rmNearestTiesToEven); - assert(!(status & APFloat::opInvalidOp) && "distributed conv div requires finite non-zero constant"); - values.assign(fragmentType.getNumElements(), one); - } - return getOrCreateConstant(rewriter, - rewriter.getInsertionBlock()->getParentOp(), - DenseFPElementsAttr::get(fragmentType, values), - fragmentType); -} - -[[maybe_unused]] static FailureOr applyDistributedPreservingStep(const DistributedTensorInfo& inputInfo, - const DistributedTensorStep& step, - PatternRewriter& rewriter, - Location loc) { - auto logicalType = inputInfo.logicalType; - auto fragmentType = getRowStripFragmentType(logicalType); - auto storageType = getRowStripStorageType(logicalType); - auto batchOp = createSpatComputeBatch(rewriter, - loc, - TypeRange {storageType}, - inputInfo.laneCount, - {}, - ValueRange {inputInfo.storage}, - [&](detail::SpatComputeBatchBodyArgs args) { - 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(); - break; - case DistributedTensorOpKind::Sigmoid: - fragment = spatial::SpatSigmoidOp::create(rewriter, loc, fragmentType, fragment).getResult(); - break; - case DistributedTensorOpKind::Add: { - Value constant = createFragmentConstant(step, fragmentType, rewriter); - fragment = - spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult(); - break; - } - case DistributedTensorOpKind::Sub: { - Value constant = createFragmentConstant(step, fragmentType, rewriter); - Value lhs = step.fragmentOnLhs ? fragment : constant; - Value rhs = step.fragmentOnLhs ? constant : fragment; - fragment = spatial::SpatVSubOp::create(rewriter, loc, fragmentType, lhs, rhs).getResult(); - break; - } - case DistributedTensorOpKind::Mul: { - Value constant = createFragmentConstant(step, fragmentType, rewriter); - fragment = - spatial::SpatVMulOp::create(rewriter, loc, fragmentType, fragment, constant).getResult(); - break; - } - case DistributedTensorOpKind::Div: { - Value constant = createFragmentReciprocalConstant(step, fragmentType, rewriter); - fragment = - spatial::SpatVMulOp::create(rewriter, loc, fragmentType, fragment, constant).getResult(); - break; - } - case DistributedTensorOpKind::Conv: - return failure(); - } - insertRowStripFragment( - fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc); - return success(); - }); - if (failed(batchOp)) - return failure(); - return makeDistributedTensorInfo(batchOp->getResult(0), logicalType); -} static Value createCollectedConvOutput(ValueRange gemmRows, Type convType, @@ -3761,77 +2880,15 @@ static Value createCollectedConvOutput(ValueRange gemmRows, int64_t numPatches, int64_t numChannelsOut, int64_t packFactor, - ArrayRef distributedConsumers, PatternRewriter& rewriter, Location loc) { - auto materializeSplatTensor = [&](DenseElementsAttr denseAttr, RankedTensorType targetType) { - Attribute splatValue = denseAttr.getSplatValue(); - auto targetAttr = DenseElementsAttr::get(targetType, splatValue); - return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), targetAttr, targetType); - }; - - auto materializeReciprocalSplatTensor = [&](DenseFPElementsAttr denseAttr, RankedTensorType targetType) { - APFloat reciprocal = denseAttr.getSplatValue(); - APFloat one(reciprocal.getSemantics(), 1); - [[maybe_unused]] APFloat::opStatus status = one.divide(reciprocal, APFloat::rmNearestTiesToEven); - assert(!(status & APFloat::opInvalidOp) && "distributed conv div consumer requires finite non-zero scalar"); - return getOrCreateConstant( - rewriter, rewriter.getInsertionBlock()->getParentOp(), DenseFPElementsAttr::get(targetType, one), targetType); - }; - - auto applyDistributedConsumers = [&](Value fragment) { - Value current = fragment; - for (const DistributedTensorStep& step : distributedConsumers) { - auto fragmentType = cast(current.getType()); - switch (step.kind) { - case DistributedTensorOpKind::Relu: - current = spatial::SpatReluOp::create(rewriter, loc, fragmentType, current).getResult(); - break; - case DistributedTensorOpKind::Sigmoid: - current = spatial::SpatSigmoidOp::create(rewriter, loc, fragmentType, current).getResult(); - break; - case DistributedTensorOpKind::Add: { - Value splat = materializeSplatTensor(step.constantAttr, fragmentType); - current = spatial::SpatVAddOp::create(rewriter, loc, fragmentType, current, splat).getResult(); - break; - } - case DistributedTensorOpKind::Sub: { - Value splat = materializeSplatTensor(step.constantAttr, fragmentType); - Value lhs = step.fragmentOnLhs ? current : splat; - Value rhs = step.fragmentOnLhs ? splat : current; - current = spatial::SpatVSubOp::create(rewriter, loc, fragmentType, lhs, rhs).getResult(); - break; - } - case DistributedTensorOpKind::Mul: { - Value splat = materializeSplatTensor(step.constantAttr, fragmentType); - current = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, current, splat).getResult(); - break; - } - case DistributedTensorOpKind::Div: { - auto reciprocalAttr = cast(step.constantAttr); - Value reciprocal = materializeReciprocalSplatTensor(reciprocalAttr, fragmentType); - current = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, current, reciprocal).getResult(); - break; - } - case DistributedTensorOpKind::Conv: - llvm_unreachable("conv-consuming distributed chains should not materialize through createCollectedConvOutput"); - } - } - return current; - }; - auto collectComputeOp = createSpatCompute(rewriter, loc, convType, {}, gemmRows, [&](ValueRange gemmRowArgs) { - SmallVector transformedRows; - transformedRows.reserve(gemmRowArgs.size()); - for (Value row : gemmRowArgs) - transformedRows.push_back(applyDistributedConsumers(row)); - Value gemmOut; if (packFactor == 1) { - gemmOut = createSpatConcat(rewriter, loc, /*axis=*/0, transformedRows); + gemmOut = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs); } else { - Value packedOutput = createSpatConcat(rewriter, loc, /*axis=*/0, transformedRows); + Value packedOutput = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs); gemmOut = standard::unpackRowsFromParallelGemm( packedOutput, cast(packedOutput.getType()), numPatches, numChannelsOut, packFactor, rewriter, loc); } @@ -4108,6 +3165,22 @@ static LogicalResult verifyForcedConvLoweringStrategy(Operation* op, llvm_unreachable("unknown conv lowering strategy"); } +static FailureOr selectConvLoweringStrategy(Operation* op, + const ConvLoweringState& state) { + FailureOr requested = resolveRequestedConvLoweringStrategy(op); + if (failed(requested)) + return failure(); + + ConvGeometry geometry = buildConvGeometry(state); + PimConvLoweringType strategy = chooseConvLoweringStrategy(geometry, *requested); + if (strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(state) + && *requested == PimConvLoweringAuto) + strategy = PimConvLoweringLegacy; + if (failed(verifyForcedConvLoweringStrategy(op, geometry, strategy))) + return failure(); + return strategy; +} + static FailureOr lowerDenseSelectedConvPlan(Operation* op, const ConvLoweringState& state, PimConvLoweringType strategy, @@ -4123,15 +3196,13 @@ static ConvLoweringState makeGroupedConvLoweringState(const ConvLoweringState& p static FailureOr buildConvValueForStrategy(Operation* op, Location loc, const ConvLoweringState& state, - const ConvLoweringDecision& decision, - const DistributedConvAnalysis& analysis, - ArrayRef distributedConsumers, + PimConvLoweringType strategy, PatternRewriter& rewriter); static FailureOr buildGroupedConvValue(Operation* op, Location loc, const ConvLoweringState& state, - const ConvLoweringDecision& decision, + PimConvLoweringType strategy, PatternRewriter& rewriter); static FailureOr lowerGroupedSelectedConvPlan(Operation* op, @@ -4139,8 +3210,7 @@ static FailureOr lowerGroupedSelectedConvPlan(Operation* op, PimConvLoweringType strategy, PatternRewriter& rewriter, Location loc) { - ConvLoweringDecision decision {strategy, "", false, "", ""}; - return buildGroupedConvValue(op, loc, state, decision, rewriter); + return buildGroupedConvValue(op, loc, state, strategy, rewriter); } static FailureOr lowerDenseSelectedConvPlan(Operation* op, @@ -4148,40 +3218,33 @@ static FailureOr lowerDenseSelectedConvPlan(Operation* op, PimConvLoweringType strategy, PatternRewriter& rewriter, Location loc) { - DistributedConvAnalysis analysis; - analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer; - analysis.barrierDetail = "selected dense layout"; - ConvLoweringDecision decision {strategy, "", false, "", ""}; - return buildConvValueForStrategy(op, loc, state, decision, analysis, {}, rewriter); + return buildConvValueForStrategy(op, loc, state, strategy, rewriter); } static FailureOr buildConvValueForStrategy(Operation* op, Location loc, const ConvLoweringState& state, - const ConvLoweringDecision& decision, - const DistributedConvAnalysis& analysis, - ArrayRef distributedConsumers, + PimConvLoweringType strategy, PatternRewriter& rewriter) { - (void)analysis; const ConvGeometry geo = buildConvGeometry(state); - switch (decision.strategy) { + switch (strategy) { case PimConvLoweringDepthwise: { return depthwise::rewriteConv(op, state, rewriter, loc); } case PimConvLoweringLegacy: case PimConvLoweringPackedIm2Col: { - return standard::rewritePackedIm2ColConv(state, distributedConsumers, rewriter, loc); + return standard::rewritePackedIm2ColConv(state, rewriter, loc); } case PimConvLoweringStreamedPatch: case PimConvLoweringOutputChannelTiled: case PimConvLoweringTiled2D: { - return standard::rewriteStreamedConv(state, distributedConsumers, rewriter, loc, /*forcedPackFactor=*/1); + return standard::rewriteStreamedConv(state, rewriter, loc, /*forcedPackFactor=*/1); } case PimConvLoweringInputKTiled: { - return standard::rewriteInputKTiledConv(state, distributedConsumers, rewriter, loc); + return standard::rewriteInputKTiledConv(state, rewriter, loc); } case PimConvLoweringStreamedPacked: { - return standard::rewriteStreamedConv(state, distributedConsumers, rewriter, loc, geo.pack); + return standard::rewriteStreamedConv(state, rewriter, loc, geo.pack); } case PimConvLoweringAuto: break; @@ -4190,130 +3253,6 @@ static FailureOr buildConvValueForStrategy(Operation* op, return failure(); } -static LogicalResult -createConvValueForStrategy(ONNXConvOp convOp, - const ConvLoweringState& state, - const ConvLoweringDecision& decision, - const DistributedConvAnalysis& analysis, - ArrayRef distributedConsumers, - PatternRewriter& rewriter, - FailureOr& result) { - result = buildConvValueForStrategy(convOp, convOp.getLoc(), state, decision, analysis, distributedConsumers, rewriter); - if (failed(result)) - return failure(); - - const ConvGeometry geo = buildConvGeometry(state); - const ConvStrategyEstimate estimate = estimateConvStrategy(geo, decision.strategy, analysis); - switch (decision.strategy) { - case PimConvLoweringDepthwise: - reportConvLoweringDecision( - convOp, geo, decision, estimate, /*batchSize=*/geo.p, /*numberOfBatches=*/1, /*usesComputeBatch=*/true, - /*usesBatchedInstructionEmission=*/true, std::nullopt); - return success(); - case PimConvLoweringLegacy: - case PimConvLoweringPackedIm2Col: - reportConvLoweringDecision( - convOp, geo, decision, estimate, /*batchSize=*/geo.pack, /*numberOfBatches=*/1, /*usesComputeBatch=*/true, - /*usesBatchedInstructionEmission=*/true, std::nullopt); - return success(); - case PimConvLoweringStreamedPatch: - case PimConvLoweringOutputChannelTiled: - case PimConvLoweringTiled2D: { - uint64_t chunkPositions = chooseStreamChunkPositions(geo, /*packFactor=*/1); - const int64_t batches = ceilIntegerDivide(geo.p, static_cast(chunkPositions)); - reportConvLoweringDecision(convOp, - geo, - decision, - estimate, - /*batchSize=*/1, - batches, - /*usesComputeBatch=*/true, - /*usesBatchedInstructionEmission=*/true, - chunkPositions); - return success(); - } - case PimConvLoweringInputKTiled: { - const int64_t numKSlices = ceilIntegerDivide(geo.k, geo.xbarSize); - const uint64_t maxLanesPerBatch = - std::max(1, - static_cast(crossbarCountInCore.getValue()) - / static_cast(std::max(1, numKSlices * 4))); - const uint64_t rowChunkWidth = std::max( - 1, - std::min({chooseStreamChunkPositions(geo, /*packFactor=*/1), - maxLanesPerBatch, - static_cast(state.outWidth)})); - const int64_t batches = - state.batchSize * state.outHeight * ceilIntegerDivide(state.outWidth, static_cast(rowChunkWidth)); - reportConvLoweringDecision(convOp, - geo, - decision, - estimate, - /*batchSize=*/1, - batches, - /*usesComputeBatch=*/false, - /*usesBatchedInstructionEmission=*/false, - rowChunkWidth); - return success(); - } - case PimConvLoweringStreamedPacked: { - uint64_t chunkPositions = chooseStreamChunkPositions(geo, geo.pack); - const int64_t batches = ceilIntegerDivide(geo.p, static_cast(chunkPositions)); - reportConvLoweringDecision(convOp, - geo, - decision, - estimate, - /*batchSize=*/geo.pack, - batches, - /*usesComputeBatch=*/true, - /*usesBatchedInstructionEmission=*/true, - chunkPositions); - return success(); - } - case PimConvLoweringAuto: - break; - } - return convOp.emitOpError("unexpected auto strategy at Conv lowering dispatch"); -} - -static LogicalResult -rewriteSelectedConv(ONNXConvOp convOp, - const ConvLoweringState& state, - const ConvLoweringDecision& decision, - const DistributedConvAnalysis& analysis, - PatternRewriter& rewriter) { - FailureOr result = failure(); - if (failed(createConvValueForStrategy(convOp, state, decision, analysis, analysis.steps, rewriter, result))) - return failure(); - - if (!analysis.hasLocalConsumers()) { - rewriter.replaceOp(convOp, *result); - return success(); - } - - assert(analysis.replacementOp && "conv rewrite expects a replacement op"); - rewriter.replaceOp(analysis.replacementOp, *result); - for (auto it = analysis.steps.rbegin(); it != analysis.steps.rend(); ++it) - if (it->op != analysis.replacementOp) - rewriter.eraseOp(it->op); - rewriter.eraseOp(convOp); - return success(); -} - -[[maybe_unused]] static LogicalResult -rewriteUngroupedConv(ONNXConvOp convOp, - const ConvLoweringState& state, - const ConvLoweringDecision& decision, - const DistributedConvAnalysis& analysis, - PatternRewriter& rewriter) { - return rewriteSelectedConv(convOp, state, decision, analysis, rewriter); -} - -static LogicalResult -rewriteGroupedConv(ONNXConvOp convOp, - const ConvLoweringState& state, - const ConvLoweringDecision& decision, - PatternRewriter& rewriter); static ConvLoweringState makeGroupedConvLoweringState(const ConvLoweringState& parent, Value groupX, @@ -4349,7 +3288,7 @@ static ConvLoweringState makeGroupedConvLoweringState( static FailureOr buildGroupedConvValue(Operation* op, Location loc, const ConvLoweringState& state, - const ConvLoweringDecision& decision, + PimConvLoweringType strategy, PatternRewriter& rewriter) { SmallVector xSlices = sliceTensor(state.x, /*axis=*/1, state.numChannelsInPerGroup, rewriter, loc); SmallVector wSlices = sliceTensor(state.w, /*axis=*/0, state.numChannelsOutPerGroup, rewriter, loc); @@ -4384,11 +3323,7 @@ static FailureOr buildGroupedConvValue(Operation* op, Value groupW = wSlices[groupId]; Value groupB = state.hasBias ? bSlices[groupId] : Value(); ConvLoweringState groupState = makeGroupedConvLoweringState(state, groupX, groupW, groupB, groupOutType); - DistributedConvAnalysis groupAnalysis; - groupAnalysis.barrierKind = DistributedConvBarrierKind::GroupedConv; - groupAnalysis.barrierDetail = "grouped convolution still materializes densely"; - FailureOr groupResult = - buildConvValueForStrategy(op, loc, groupState, decision, groupAnalysis, {}, rewriter); + FailureOr groupResult = buildConvValueForStrategy(op, loc, groupState, strategy, rewriter); if (failed(groupResult)) return failure(); groupResults.push_back(*groupResult); @@ -4403,17 +3338,6 @@ static FailureOr buildGroupedConvValue(Operation* op, return concatCompute.getResult(0); } -[[maybe_unused]] static LogicalResult -rewriteGroupedConv(ONNXConvOp convOp, - const ConvLoweringState& state, - const ConvLoweringDecision& decision, - PatternRewriter& rewriter) { - FailureOr result = buildGroupedConvValue(convOp.getOperation(), convOp.getLoc(), state, decision, rewriter); - if (failed(result)) - return failure(); - rewriter.replaceOp(convOp, *result); - return success(); -} } // namespace @@ -4459,29 +3383,14 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) { if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType)) return failure(); - FailureOr requestedStrategy = resolveRequestedConvLoweringStrategy(planOp.getOperation()); - if (failed(requestedStrategy)) - return failure(); - - DistributedConvAnalysis analysis; - analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer; - analysis.barrierDetail = "selected row-strip layout"; ConvGeometry geometry = buildConvGeometry(*state); if (!rowStripOutputChannelTileFitsOneCore(geometry)) return failure(); - ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis); - if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state) - && *requestedStrategy == PimConvLoweringAuto) { - decision = {PimConvLoweringLegacy, - "depthwise auto fallback when structured depthwise lowering is not representable", - /*isAuto=*/true, - "", - ""}; - } - if (failed(verifyForcedConvLoweringStrategy(planOp.getOperation(), geometry, decision.strategy))) + FailureOr strategy = selectConvLoweringStrategy(planOp.getOperation(), *state); + if (failed(strategy)) return failure(); - switch (decision.strategy) { + switch (*strategy) { case PimConvLoweringLegacy: case PimConvLoweringDepthwise: case PimConvLoweringPackedIm2Col: @@ -4502,6 +3411,11 @@ LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) { if (failed(state)) return failure(); + FailureOr strategy = selectConvLoweringStrategy(planOp.getOperation(), *state); + if (failed(strategy)) + return failure(); + if (*strategy == PimConvLoweringDepthwise) + return canConsumeDepthwiseRowStrip(*state) ? success() : failure(); StringRef failureReason; return canConsumePixelMajorRowStripFragments(*state, failureReason) ? success() : failure(); } @@ -4515,31 +3429,16 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, if (failed(state)) return failure(); - FailureOr requestedStrategy = resolveRequestedConvLoweringStrategy(planOp.getOperation()); - if (failed(requestedStrategy)) - return failure(); - - DistributedConvAnalysis analysis; - analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer; - analysis.barrierDetail = emitRowStripLayout ? "selected row-strip layout" : "selected dense layout"; - ConvGeometry geometry = buildConvGeometry(*state); - ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis); - if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state) - && *requestedStrategy == PimConvLoweringAuto) { - decision = {PimConvLoweringLegacy, - "depthwise auto fallback when structured depthwise lowering is not representable", - /*isAuto=*/true, - "", - ""}; - } - if (failed(verifyForcedConvLoweringStrategy(planOp.getOperation(), geometry, decision.strategy))) + FailureOr strategy = selectConvLoweringStrategy(planOp.getOperation(), *state); + if (failed(strategy)) return failure(); if (emitRowStripLayout) { if (rowStripInput) { if (failed(canConsumeAndProduceRowStrip(planOp))) return planOp.emitOpError("selected row-strip input/output layout is not supported for this Conv plan"), failure(); - return createConvOutputFromRowStripInput(*state, decision, *rowStripInput, rewriter, planOp.getLoc()); + return createConvOutputFromRowStripInput( + *state, *rowStripInput, *strategy, rewriter, planOp.getLoc()); } if (failed(canLowerConvPlanToRowStrip(planOp))) return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure(); @@ -4549,11 +3448,11 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, return *rowStripStorage; } - if (decision.strategy == PimConvLoweringDepthwise) - return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, decision.strategy, rewriter, planOp.getLoc()); + if (*strategy == PimConvLoweringDepthwise) + return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, *strategy, rewriter, planOp.getLoc()); if (state->group != 1) - return lowerGroupedSelectedConvPlan(planOp.getOperation(), *state, decision.strategy, rewriter, planOp.getLoc()); - return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, decision.strategy, rewriter, planOp.getLoc()); + return lowerGroupedSelectedConvPlan(planOp.getOperation(), *state, *strategy, rewriter, planOp.getLoc()); + return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, *strategy, rewriter, planOp.getLoc()); } } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp index a71808d..871640d 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp @@ -48,6 +48,56 @@ static DenseElementsAttr getDenseConstantAttr(Value value) { return nullptr; } +struct BlueprintSplatMulToSpatial : OpConversionPattern { + explicit BlueprintSplatMulToSpatial(MLIRContext* ctx) : OpConversionPattern(ctx, 2) {} + + LogicalResult + matchAndRewrite(ONNXMulOp op, ONNXMulOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override { + auto blueprint = adaptor.getA().getDefiningOp(); + Value scalar = adaptor.getB(); + if (!blueprint) { + blueprint = adaptor.getB().getDefiningOp(); + scalar = adaptor.getA(); + } + auto scalarAttr = getDenseConstantAttr(scalar); + auto resultType = dyn_cast(op.getResult().getType()); + auto storageType = blueprint ? dyn_cast(blueprint.getInput().getType()) : RankedTensorType(); + if (!blueprint || !blueprint.getFragments().empty() || !scalarAttr || !scalarAttr.isSplat() || !resultType + || resultType != blueprint.getOutput().getType() || !storageType) + return failure(); + + auto mapped = mapGraphBatchFragments( + blueprint.getInput(), storageType, rewriter, op.getLoc(), [&](Value fragment, RankedTensorType fragmentType) { + auto splat = DenseElementsAttr::get(fragmentType, scalarAttr.getSplatValue()); + Value constant = arith::ConstantOp::create(rewriter, op.getLoc(), fragmentType, splat); + return FailureOr( + spatial::SpatVMulOp::create(rewriter, op.getLoc(), fragmentType, fragment, constant).getResult()); + }); + if (failed(mapped)) + return failure(); + + auto result = spatial::SpatBlueprintOp::create(rewriter, + op.getLoc(), + resultType, + *mapped, + ValueRange {}, + blueprint.getLogicalLayoutAttr(), + blueprint.getPhysicalLayoutAttr(), + blueprint.getFragmentOffsetsAttr(), + blueprint.getFragmentSizesAttr(), + blueprint.getIndexMapAttr(), + blueprint.getModeAttr(), + blueprint.getFragmentOperandIndicesAttr(), + blueprint.getFragmentSourceSlotsAttr(), + blueprint.getFragmentSourceOffsetsAttr(), + blueprint.getFragmentStridesAttr(), + blueprint.getConflictPolicyAttr(), + blueprint.getCoveragePolicyAttr()); + rewriter.replaceOp(op, result.getOutput()); + return success(); + } +}; + static FailureOr materializeBroadcastedConstantTensor(Value value, RankedTensorType resultType, ConversionPatternRewriter& rewriter, @@ -246,6 +296,7 @@ void populateElementwiseFusionPatterns(RewritePatternSet& patterns, MLIRContext* } void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { + patterns.add(ctx); patterns.add(ctx); patterns.add>(ctx); patterns.add>(ctx); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp index fd26d78..3332073 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp @@ -106,6 +106,92 @@ static Value mapOutputBatchIndexToSourceBatchIndex(Value outputBatchIndex, return sourceBatchIndex; } +static FailureOr collapseFragmentAssemblyBatchDims(Value value, + RankedTensorType resultType, + PatternRewriter& rewriter, + Location loc) { + auto blueprint = value.getDefiningOp(); + auto inputType = dyn_cast(value.getType()); + auto storageType = blueprint ? dyn_cast(blueprint.getInput().getType()) : RankedTensorType(); + auto operandIndices = blueprint ? blueprint.getFragmentOperandIndices() : std::nullopt; + auto sourceOffsets = blueprint ? blueprint.getFragmentSourceOffsets() : std::nullopt; + auto fragmentStrides = blueprint ? blueprint.getFragmentStrides() : std::nullopt; + if (!blueprint || !inputType || !storageType || !inputType.hasStaticShape() || !storageType.hasStaticShape() + || inputType.getRank() <= 3 || resultType.getRank() != 3 || !blueprint.getFragments().empty() + || blueprint.getMode() != "fragment_assembly" || !operandIndices || !sourceOffsets || !fragmentStrides + || storageType.getRank() != inputType.getRank() + 1) + return failure(); + if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments + && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint)) + return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure(); + + const int64_t batchRank = inputType.getRank() - 2; + SmallVector reassociation {ReassociationIndices {}, + ReassociationIndices {batchRank}, + ReassociationIndices {batchRank + 1}}; + for (int64_t dim = 0; dim < batchRank; ++dim) + reassociation.front().push_back(dim); + SmallVector outputFragmentShape {1, storageType.getDimSize(batchRank + 1), storageType.getDimSize(batchRank + 2)}; + auto outputFragmentType = RankedTensorType::get(outputFragmentShape, storageType.getElementType()); + auto outputStorageType = spatial::getGraphBatchPhysicalResultType(storageType.getDimSize(0), outputFragmentType); + SmallVector storageReassociation { + ReassociationIndices {0}, ReassociationIndices {}, ReassociationIndices {batchRank + 1}, + ReassociationIndices {batchRank + 2}}; + for (int64_t dim = 0; dim < batchRank; ++dim) + storageReassociation[1].push_back(dim + 1); + Value collapsedStorage = tensor::CollapseShapeOp::create( + rewriter, loc, outputStorageType, blueprint.getInput(), storageReassociation); + + const int64_t inputRank = inputType.getRank(); + const int64_t fragmentCount = operandIndices->size(); + ArrayRef inputOffsets = blueprint.getFragmentOffsets(); + ArrayRef inputSizes = blueprint.getFragmentSizes(); + SmallVector batchShape(inputType.getShape().drop_back(2)); + SmallVector batchStrides = computeRowMajorStrides(batchShape); + SmallVector offsets, sizes, strides; + offsets.reserve(fragmentCount * 3); + sizes.reserve(fragmentCount * 3); + strides.reserve(fragmentCount * 3); + for (int64_t fragment = 0; fragment < fragmentCount; ++fragment) { + int64_t flatBatch = 0; + for (int64_t dim = 0; dim < batchRank; ++dim) { + const int64_t index = fragment * inputRank + dim; + if (inputSizes[index] != 1 || (*fragmentStrides)[index] != 1) + return failure(); + flatBatch += inputOffsets[index] * batchStrides[dim]; + } + offsets.push_back(flatBatch); + sizes.push_back(1); + strides.push_back(1); + for (int64_t dim = batchRank; dim < inputRank; ++dim) { + const int64_t index = fragment * inputRank + dim; + offsets.push_back(inputOffsets[index]); + sizes.push_back(inputSizes[index]); + strides.push_back((*fragmentStrides)[index]); + } + } + auto collapsedBlueprint = spatial::SpatBlueprintOp::create(rewriter, + loc, + resultType, + collapsedStorage, + ValueRange {}, + blueprint.getLogicalLayoutAttr(), + rewriter.getStringAttr("fragmented"), + rewriter.getDenseI64ArrayAttr(offsets), + rewriter.getDenseI64ArrayAttr(sizes), + rewriter.getStringAttr("collapsed_fragments"), + blueprint.getModeAttr(), + blueprint.getFragmentOperandIndicesAttr(), + blueprint.getFragmentSourceSlotsAttr(), + blueprint.getFragmentSourceOffsetsAttr(), + rewriter.getDenseI64ArrayAttr(strides), + blueprint.getConflictPolicyAttr(), + blueprint.getCoveragePolicyAttr()); + if (spatial::isCanonicalContiguousRowMajorFragmentAssembly(collapsedBlueprint)) + collapsedBlueprint.setIndexMapAttr(rewriter.getStringAttr(spatial::kContiguousRowMajorFragments)); + return collapsedBlueprint.getOutput(); +} + static Value collapseBatchDims(Value value, int64_t batchSize, int64_t rows, int64_t cols, PatternRewriter& rewriter, Location loc) { auto type = cast(value.getType()); @@ -113,6 +199,8 @@ collapseBatchDims(Value value, int64_t batchSize, int64_t rows, int64_t cols, Pa return value; auto collapsedType = RankedTensorType::get({batchSize, rows, cols}, type.getElementType(), type.getEncoding()); + if (auto collapsed = collapseFragmentAssemblyBatchDims(value, collapsedType, rewriter, loc); succeeded(collapsed)) + return *collapsed; SmallVector reassociation = {ReassociationIndices {}, ReassociationIndices {static_cast(type.getRank() - 2)}, ReassociationIndices {static_cast(type.getRank() - 1)}}; @@ -241,7 +329,34 @@ static Value extractBatchMatrix(Value value, return materializeOrComputeUnary(value, matrixType, rewriter, loc, buildMatrix); } +static Value getLastTwoTransposeInput(Value value) { + auto type = cast(value.getType()); + if (auto transpose = value.getDefiningOp()) { + auto permutation = getTransposePermutationChecked(transpose.getPermAttr(), type.getRank()); + if (succeeded(permutation) && llvm::all_of(llvm::seq(0, type.getRank()), [&](int64_t dim) { + return (*permutation)[dim] == (dim < type.getRank() - 2 ? dim : 2 * type.getRank() - 3 - dim); + })) + return transpose.getData(); + } + return {}; +} + +static std::pair splitSplatMultiply(Value value) { + if (!value) + return {}; + auto multiply = value.getDefiningOp(); + if (!multiply) + return {}; + for (auto [data, scale] : {std::pair {multiply.getA(), multiply.getB()}, + std::pair {multiply.getB(), multiply.getA()}}) + if (auto constant = getHostConstDenseElementsAttr(scale); constant && constant.isSplat()) + return {data, scale}; + return {}; +} + static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Location loc) { + if (Value input = getLastTwoTransposeInput(value)) + return input; auto type = cast(value.getType()); auto shape = type.getShape(); auto createONNXTranspose = [&](RankedTensorType resultType, ArrayRef permutation) { @@ -407,123 +522,152 @@ static Value extractDynamicBatchedRowVector(Value matrix, {offsets, sizes, getUnitStrides(rewriter, 3)}); } +static int64_t chooseDynamicMatMulRowsPerLane(int64_t rows, int64_t reductionSize, int64_t columns) { + const int64_t crossbarElements = static_cast(crossbarSize.getValue() * crossbarSize.getValue()); + const int64_t target = std::min(rows, ceilIntegerDivide(reductionSize * columns, crossbarElements)); + int64_t rowsPerLane = 1; + for (int64_t candidate = 2; candidate <= target; ++candidate) + if (rows % candidate == 0) + rowsPerLane = candidate; + return rowsPerLane; +} + static FailureOr createBatchedVvdmulBatch(Value a, ArrayRef aBatchShape, Value b, ArrayRef bBatchShape, ArrayRef outputBatchShape, - RankedTensorType aType, - RankedTensorType bType, - RankedTensorType columnPiecesType, + int64_t reductionSize, + int64_t rowsPerLane, + RankedTensorType rowPiecesType, RankedTensorType outType, + RankedTensorType publicationFragmentType, PatternRewriter& rewriter, Location loc) { const int64_t numBatches = outType.getDimSize(0); const int64_t numOutRows = outType.getDimSize(1); const int64_t numOutCols = outType.getDimSize(2); - const int64_t reductionSize = aType.getDimSize(2); - const int64_t laneCount = numBatches * numOutCols; - auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType()); + const int64_t rowGroups = numOutRows / rowsPerLane; + const int64_t laneCount = numBatches * rowGroups; + auto vectorType = RankedTensorType::get({1, reductionSize}, outType.getElementType()); auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType()); - auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType()); + auto rowType = RankedTensorType::get({1, numOutCols}, outType.getElementType()); + auto rowsType = RankedTensorType::get({rowsPerLane, numOutCols}, outType.getElementType()); auto batchOp = createSpatComputeBatch( rewriter, loc, - TypeRange {columnPiecesType}, + TypeRange {rowPiecesType}, laneCount, ValueRange {}, ValueRange {a, b}, - [&](detail::SpatComputeBatchBodyArgs args) { + [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - Value batch = affineFloorDivConst(rewriter, loc, args.lane, numOutCols, anchorOp); - Value column = affineModConst(rewriter, loc, args.lane, numOutCols, anchorOp); - Value bVector = extractDynamicBatchedRowVector( - args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc); - Value columnInit = tensor::EmptyOp::create(rewriter, loc, columnType.getShape(), columnType.getElementType()); + Value batch = affineFloorDivConst(rewriter, loc, args.lane, rowGroups, anchorOp); + Value rowGroup = affineModConst(rewriter, loc, args.lane, rowGroups, anchorOp); + Value rowBase = affineMulConst(rewriter, loc, rowGroup, rowsPerLane, anchorOp); Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0); Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1); - Value cNumOutRows = - getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows); - auto loop = buildNormalizedScfFor( + Value cRows = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), rowsPerLane); + Value cNumOutCols = + getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutCols); + Value rowsInit = tensor::EmptyOp::create(rewriter, loc, rowsType.getShape(), rowsType.getElementType()); + auto rowsLoop = buildNormalizedScfFor( rewriter, loc, c0, - cNumOutRows, + cRows, c1, - ValueRange {columnInit}, - [&](OpBuilder&, Location nestedLoc, Value row, ValueRange iterArgs, SmallVectorImpl& yielded) { + ValueRange {rowsInit}, + [&](OpBuilder&, Location nestedLoc, Value rowOffset, ValueRange iterArgs, SmallVectorImpl& yielded) { + Value row = arith::AddIOp::create(rewriter, nestedLoc, rowBase, rowOffset); Value aVector = extractDynamicBatchedRowVector( args.inputs[0], aBatchShape, outputBatchShape, batch, row, vectorType, rewriter, nestedLoc); - Value scalar = spatial::SpatVVDMulOp::create(rewriter, nestedLoc, scalarType, aVector, bVector).getResult(); - Value next = tensor::InsertSliceOp::create(rewriter, - nestedLoc, - scalar, - iterArgs.front(), - SmallVector {row, rewriter.getIndexAttr(0)}, - SmallVector {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(1)}, - getUnitStrides(rewriter, 2)); - yielded.push_back(next); + Value rowInit = tensor::EmptyOp::create(rewriter, nestedLoc, rowType.getShape(), rowType.getElementType()); + auto columnsLoop = buildNormalizedScfFor( + rewriter, + nestedLoc, + c0, + cNumOutCols, + c1, + ValueRange {rowInit}, + [&](OpBuilder&, Location columnLoc, Value column, ValueRange columnArgs, SmallVectorImpl& rowYielded) { + Value bVector = extractDynamicBatchedRowVector( + args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, columnLoc); + Value scalar = spatial::SpatVVDMulOp::create(rewriter, columnLoc, scalarType, aVector, bVector).getResult(); + rowYielded.push_back(tensor::InsertSliceOp::create( + rewriter, + columnLoc, + scalar, + columnArgs.front(), + SmallVector {rewriter.getIndexAttr(0), column}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}, + getUnitStrides(rewriter, 2))); + return success(); + }); + assert(succeeded(columnsLoop) && "dynamic MatMul column loop construction must succeed"); + yielded.push_back(tensor::InsertSliceOp::create( + rewriter, + nestedLoc, + columnsLoop->results.front(), + iterArgs.front(), + SmallVector {rowOffset, rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(numOutCols)}, + getUnitStrides(rewriter, 2))); return success(); }); - assert(succeeded(loop) && "dynamic MatMul row loop construction must succeed"); - publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane); + assert(succeeded(rowsLoop) && "dynamic MatMul row-group loop construction must succeed"); + Value fragment = rowsLoop->results.front(); + while (fragment.getType() != publicationFragmentType) { + auto expanded = addLeadingUnitTensorDimension(rewriter, loc, fragment); + if (failed(expanded)) + return failure(); + fragment = *expanded; + } + publishGraphBatchPhysicalFragment(rewriter, loc, fragment, args.outputs.front(), args.lane); + return success(); }); if (failed(batchOp)) return failure(); return *batchOp; } -static FailureOr createBatchedDynamicOutputCompute(Value scalarPieces, - RankedTensorType scalarPiecesType, - RankedTensorType outType, - PatternRewriter& rewriter, - Location loc) { - const int64_t laneCount = scalarPiecesType.getDimSize(0); - const int64_t numOutCols = outType.getDimSize(2); - auto columnType = RankedTensorType::get({outType.getDimSize(1), 1}, outType.getElementType()); - - auto computeOp = createSpatCompute<1>( - rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult { - Value outputInit = - tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType()).getResult(); - Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0); - Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1); - Value cLaneCount = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), laneCount); - auto loop = buildNormalizedScfFor( - rewriter, - loc, - c0, - cLaneCount, - c1, - ValueRange {outputInit}, - [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl& yielded) { - Value outputAcc = iterArgs.front(); - Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - Value batch = affineFloorDivConst(rewriter, nestedLoc, lane, numOutCols, anchorOp); - Value column = affineModConst(rewriter, nestedLoc, lane, numOutCols, anchorOp); - FailureOr columnPiece = - extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, columnType); - if (failed(columnPiece)) - return failure(); - SmallVector outputOffsets {batch, rewriter.getIndexAttr(0), column}; - SmallVector outputSizes = { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(outType.getDimSize(1)), rewriter.getIndexAttr(1)}; - Value next = - tensor::InsertSliceOp::create( - rewriter, nestedLoc, *columnPiece, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3)) - .getResult(); - yielded.push_back(next); - return success(); - }); - if (failed(loop)) - return failure(); - spatial::SpatYieldOp::create(rewriter, loc, loop->results.front()); - return success(); - }); - if (failed(computeOp)) - return failure(); - return computeOp->getResult(0); +static FailureOr createBatchedRowOutputBlueprint(Value rowPieces, + RankedTensorType outType, + ArrayRef batchShape, + int64_t rowsPerFragment, + PatternRewriter& rewriter, + Location loc) { + SmallVector entries; + const int64_t rowAxis = outType.getRank() - 2; + const int64_t rows = outType.getDimSize(rowAxis); + const int64_t columns = outType.getDimSize(rowAxis + 1); + const int64_t batches = batchShape.empty() ? 1 : getStaticShapeElementCount(batchShape); + SmallVector batchStrides = computeRowMajorStrides(batchShape); + const int64_t rowGroups = rows / rowsPerFragment; + entries.reserve(batches * rowGroups); + for (int64_t batch = 0; batch < batches; ++batch) + for (int64_t row = 0; row < rows; row += rowsPerFragment) { + SmallVector offsets; + for (auto [dim, size] : llvm::enumerate(batchShape)) + offsets.push_back((batch / batchStrides[dim]) % size); + offsets.push_back(row); + offsets.push_back(0); + SmallVector sizes(outType.getRank(), 1); + sizes[rowAxis] = rowsPerFragment; + sizes.back() = columns; + entries.push_back({batch * rowGroups + row / rowsPerFragment, + 0, + std::move(offsets), + std::move(sizes)}); + } + return createFragmentAssemblyBlueprint( + rowPieces, + outType, + entries, + "dense_nchw", + rowsPerFragment == 1 ? spatial::kContiguousRowMajorFragments : "row_group_fragments", + rewriter, + loc); } static Value extractBatchedReductionPiece(Value partialPiecesArg, @@ -822,6 +966,8 @@ static Value finalizeNormalizedMatMulResult(Value value, const NormalizedMatMulInfo& info, PatternRewriter& rewriter, Location loc) { + if (value.getType() == info.outType) + return value; // The direct lowered result is always [flatBatch, normalizedM, normalizedN]. // Restore ONNX MatMul result rank by expanding right-aligned batch dimensions // and removing the synthetic unit matrix axes introduced for vector operands. @@ -923,17 +1069,34 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern { Location loc = matmulOp.getLoc(); bool useTransposedForm = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB()); + Value rhsRows = getLastTwoTransposeInput(matmulOp.getB()); + ONNXTransposeOp foldedTranspose = matmulOp.getB().getDefiningOp(); + ONNXMulOp foldedMultiply = rhsRows ? rhsRows.getDefiningOp() : ONNXMulOp {}; + auto [unscaledRhsRows, outputScale] = splitSplatMultiply(rhsRows); + if (unscaledRhsRows) + rhsRows = unscaledRhsRows; + const bool rhsStoredAsRows = rhsRows && !useTransposedForm; Value lhs = normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType, shapeInfo->lhsWasVector, rewriter, loc); - Value rhs = - normalizeMatMulOperand(matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc); + Value rhs = normalizeMatMulOperand( + rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc); lhs = collapseBatchDims(lhs, shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc); - rhs = collapseBatchDims(rhs, shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc); + rhs = collapseBatchDims(rhs, + shapeInfo->rhsBatch, + rhsStoredAsRows ? shapeInfo->n : shapeInfo->k, + rhsStoredAsRows ? shapeInfo->k : shapeInfo->n, + rewriter, + loc); MatMulLoweringPlan plan = buildLoweringPlan(lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc); plan.lhs = ensureBatchedTensor(plan.lhs, plan.lhsBatch, plan.m, plan.k, rewriter, loc); - plan.rhs = ensureBatchedTensor(plan.rhs, plan.rhsBatch, plan.k, plan.n, rewriter, loc); + plan.rhs = ensureBatchedTensor(plan.rhs, + plan.rhsBatch, + rhsStoredAsRows ? plan.n : plan.k, + rhsStoredAsRows ? plan.k : plan.n, + rewriter, + loc); plan.lhsType = cast(plan.lhs.getType()); plan.rhsType = cast(plan.rhs.getType()); auto directOutType = RankedTensorType::get( @@ -997,26 +1160,34 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern { return success(); } } - const int64_t laneCount = plan.batch * plan.n; - auto columnType = RankedTensorType::get({plan.m, 1}, shapeInfo->outType.getElementType()); - auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType( - laneCount, columnType); - Value transposedRhs = transposeLastTwoDims(plan.rhs, rewriter, loc); + RankedTensorType blueprintType = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector + ? shapeInfo->outType : directOutType; + SmallVector blueprintBatchShape = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector + ? shapeInfo->outputBatchShape : SmallVector {plan.batch}; + const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.m, plan.k, plan.n); + const int64_t laneCount = plan.batch * plan.m / rowsPerLane; + SmallVector fragmentShape(blueprintType.getRank(), 1); + fragmentShape[fragmentShape.size() - 2] = rowsPerLane; + fragmentShape.back() = plan.n; + auto fragmentType = RankedTensorType::get(fragmentShape, shapeInfo->outType.getElementType()); + auto rowPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount, fragmentType); + Value transposedRhs = rhsStoredAsRows ? plan.rhs : transposeLastTwoDims(plan.rhs, rewriter, loc); auto batchOp = createBatchedVvdmulBatch(plan.lhs, plan.lhsBatchShape, transposedRhs, plan.rhsBatchShape, plan.outputBatchShape, - plan.lhsType, - plan.rhsType, - scalarPiecesType, + plan.k, + rowsPerLane, + rowPiecesType, directOutType, + fragmentType, rewriter, loc); if (failed(batchOp)) return failure(); - auto result = - createBatchedDynamicOutputCompute(batchOp->getResult(0), scalarPiecesType, directOutType, rewriter, loc); + auto result = createBatchedRowOutputBlueprint( + batchOp->getResult(0), blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc); if (failed(result)) return failure(); Value finalResult = *result; @@ -1029,13 +1200,34 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern { .getResult(); } finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc); + if (outputScale) + finalResult = ONNXMulOp::create( + rewriter, loc, shapeInfo->outType, finalResult, outputScale).getResult(); rewriter.replaceOp(matmulOp, finalResult); + if (foldedTranspose && foldedTranspose->use_empty()) + rewriter.eraseOp(foldedTranspose); + if (foldedMultiply && foldedMultiply->use_empty()) + rewriter.eraseOp(foldedMultiply); return success(); } }; +struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes { + using MatMulBatchedToSpatialComputes::MatMulBatchedToSpatialComputes; + + LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override { + if (!getLastTwoTransposeInput(matmulOp.getB())) + return failure(); + return MatMulBatchedToSpatialComputes::matchAndRewrite(matmulOp, rewriter); + } +}; + } // namespace +void populateMatMulFusionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { + patterns.add(ctx); +} + void populateMatMulRewritePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.insert(ctx); } diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp index 769a943..764ea71 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp @@ -5,8 +5,11 @@ #include "llvm/ADT/STLExtras.h" +#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" #include "src/Accelerators/PIM/Common/IR/LoopUtils.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/PlanLowering.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -17,126 +20,144 @@ namespace onnx_mlir { namespace { static Value buildNearestAsymmetricIndex( - Value outputIndex, int64_t inputDim, int64_t outputDim, ConversionPatternRewriter& rewriter, Location loc) { + Value outputIndex, int64_t inputDim, int64_t outputDim, PatternRewriter& rewriter, Location loc) { + if (inputDim == outputDim) + return outputIndex; Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); + if (outputDim % inputDim == 0) + return affineFloorDivConst(rewriter, loc, outputIndex, outputDim / inputDim, anchorOp); + if (inputDim % outputDim == 0) + return affineMulConst(rewriter, loc, outputIndex, inputDim / outputDim, anchorOp); Value cInputDim = getOrCreateIndexConstant(rewriter, anchorOp, inputDim); Value cOutputDim = getOrCreateIndexConstant(rewriter, anchorOp, outputDim); - Value cInputDimLast = getOrCreateIndexConstant(rewriter, anchorOp, inputDim - 1); Value scaledIndex = arith::MulIOp::create(rewriter, loc, outputIndex, cInputDim); - Value inputIndex = arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim); - return arith::MinUIOp::create(rewriter, loc, inputIndex, cInputDimLast); + return arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim); } -static FailureOr buildNearestResizeLoop(Value input, - RankedTensorType inputType, - RankedTensorType resultType, - ConversionPatternRewriter& rewriter, - Location loc) { - auto elemType = resultType.getElementType(); - SmallVector unitShape(resultType.getRank(), 1); - auto unitTensorType = RankedTensorType::get(unitShape, elemType); - - SmallVector unitSizes(resultType.getRank(), rewriter.getIndexAttr(1)); - SmallVector unitStrides(resultType.getRank(), rewriter.getIndexAttr(1)); - - Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); - Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); - Value cOutputN = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(0)); - Value cOutputC = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(1)); - Value cOutputH = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(2)); - Value cOutputW = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(3)); - - Value outputInit = tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), elemType); - - auto batchLoop = buildNormalizedScfFor( - rewriter, - loc, - c0, - cOutputN, - c1, - ValueRange {outputInit}, - [&](OpBuilder&, Location nestedLoc, Value outputN, ValueRange batchIterArgs, SmallVectorImpl& batchYielded) { - Value outputBatchAcc = batchIterArgs.front(); - Value inputN = - buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), resultType.getDimSize(0), rewriter, nestedLoc); - - auto channelLoop = buildNormalizedScfFor( - rewriter, - nestedLoc, - c0, - cOutputC, - c1, - ValueRange {outputBatchAcc}, - [&](OpBuilder&, - Location channelLoc, - Value outputC, - ValueRange channelIterArgs, - SmallVectorImpl& channelYielded) { - Value outputChannelAcc = channelIterArgs.front(); - Value inputC = buildNearestAsymmetricIndex( - outputC, inputType.getDimSize(1), resultType.getDimSize(1), rewriter, channelLoc); - - auto heightLoop = buildNormalizedScfFor( - rewriter, - channelLoc, - c0, - cOutputH, - c1, - ValueRange {outputChannelAcc}, - [&](OpBuilder&, - Location heightLoc, - Value outputH, - ValueRange heightIterArgs, - SmallVectorImpl& heightYielded) { - Value outputHeightAcc = heightIterArgs.front(); - Value inputH = buildNearestAsymmetricIndex( - outputH, inputType.getDimSize(2), resultType.getDimSize(2), rewriter, heightLoc); - - auto widthLoop = buildNormalizedScfFor( - rewriter, - heightLoc, - c0, - cOutputW, - c1, - ValueRange {outputHeightAcc}, - [&](OpBuilder&, - Location widthLoc, - Value outputW, - ValueRange widthIterArgs, - SmallVectorImpl& widthYielded) { - Value outputWidthAcc = widthIterArgs.front(); - Value inputW = buildNearestAsymmetricIndex( - outputW, inputType.getDimSize(3), resultType.getDimSize(3), rewriter, widthLoc); - - SmallVector inputOffsets = {inputN, inputC, inputH, inputW}; - Value inputSlice = tensor::ExtractSliceOp::create( - rewriter, widthLoc, unitTensorType, input, inputOffsets, unitSizes, unitStrides); - - SmallVector outputOffsets = {outputN, outputC, outputH, outputW}; - Value updatedOutput = tensor::InsertSliceOp::create( - rewriter, widthLoc, inputSlice, outputWidthAcc, outputOffsets, unitSizes, unitStrides); - widthYielded.push_back(updatedOutput); - return success(); - }); - if (failed(widthLoop)) - return failure(); - heightYielded.push_back(widthLoop->results.front()); - return success(); - }); - if (failed(heightLoop)) - return failure(); - channelYielded.push_back(heightLoop->results.front()); +static FailureOr buildDenseNearestResize(Value input, + RankedTensorType inputType, + RankedTensorType resultType, + PatternRewriter& rewriter, + Location loc) { + ArrayRef shape = resultType.getShape(); + int64_t rowCount = shape[0] * shape[1] * shape[2]; + auto scalarType = RankedTensorType::get({1, 1, 1, 1}, resultType.getElementType()); + auto rowType = RankedTensorType::get({1, 1, 1, shape[3]}, resultType.getElementType()); + auto rowsType = RankedTensorType::get({rowCount, 1, 1, 1, shape[3]}, resultType.getElementType()); + auto batch = createSpatComputeBatch( + rewriter, loc, TypeRange {rowsType}, rowCount, {}, ValueRange {input}, + [&](detail::SpatComputeBatchBodyArgs args) { + Operation* anchor = rewriter.getInsertionBlock()->getParentOp(); + Value outputN = affineFloorDivConst(rewriter, loc, args.lane, shape[1] * shape[2], anchor); + Value channelRow = affineModConst(rewriter, loc, args.lane, shape[1] * shape[2], anchor); + Value outputC = affineFloorDivConst(rewriter, loc, channelRow, shape[2], anchor); + Value outputH = affineModConst(rewriter, loc, channelRow, shape[2], anchor); + Value inputN = buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), shape[0], rewriter, loc); + Value inputC = buildNearestAsymmetricIndex(outputC, inputType.getDimSize(1), shape[1], rewriter, loc); + Value inputH = buildNearestAsymmetricIndex(outputH, inputType.getDimSize(2), shape[2], rewriter, loc); + Value row = tensor::EmptyOp::create(rewriter, loc, rowType.getShape(), rowType.getElementType()); + Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0); + Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1); + Value width = getOrCreateIndexConstant(rewriter, anchor, shape[3]); + auto loop = buildNormalizedScfFor( + rewriter, loc, c0, width, c1, ValueRange {row}, + [&](OpBuilder&, Location nestedLoc, Value outputW, ValueRange iterArgs, SmallVectorImpl& yielded) { + Value inputW = buildNearestAsymmetricIndex( + outputW, inputType.getDimSize(3), shape[3], rewriter, nestedLoc); + SmallVector unitSizes(4, rewriter.getIndexAttr(1)); + SmallVector unitStrides(4, rewriter.getIndexAttr(1)); + Value scalar = tensor::ExtractSliceOp::create( + rewriter, nestedLoc, scalarType, args.inputs.front(), + SmallVector {inputN, inputC, inputH, inputW}, unitSizes, unitStrides); + yielded.push_back(tensor::InsertSliceOp::create( + rewriter, nestedLoc, scalar, iterArgs.front(), + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), + rewriter.getIndexAttr(0), outputW}, + unitSizes, unitStrides)); return success(); }); - if (failed(channelLoop)) + assert(succeeded(loop) && "nearest Resize row loop construction must succeed"); + publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane); + }); + if (failed(batch)) + return failure(); + + SmallVector entries; + entries.reserve(rowCount); + for (int64_t n = 0; n < shape[0]; ++n) + for (int64_t c = 0; c < shape[1]; ++c) + for (int64_t h = 0; h < shape[2]; ++h) + entries.push_back({(n * shape[1] + c) * shape[2] + h, 0, {n, c, h, 0}, {1, 1, 1, shape[3]}}); + return createFragmentAssemblyBlueprint( + batch->getResult(0), resultType, entries, "dense_nchw", spatial::kContiguousRowMajorFragments, rewriter, loc); +} + +static FailureOr buildRowStripNearestResize( + Value storage, RankedTensorType inputType, RankedTensorType resultType, + PatternRewriter& rewriter, Location loc) { + auto input = describeRowStripPhysicalValue(storage, inputType); + if (failed(input)) + return failure(); + int64_t tilesPerRow = input->tilesPerRow; + int64_t outputHeight = resultType.getDimSize(2); + int64_t outputWidth = resultType.getDimSize(3); + int64_t tileChannels = input->fragmentType.getDimSize(3); + int64_t laneCount = outputHeight * tilesPerRow; + auto outputFragmentType = RankedTensorType::get( + {1, 1, outputWidth, tileChannels}, resultType.getElementType()); + auto outputStorageType = spatial::getGraphBatchPhysicalResultType( + laneCount, outputFragmentType); + auto pixelType = RankedTensorType::get( + {1, 1, 1, tileChannels}, resultType.getElementType()); + auto batch = createSpatComputeBatch( + rewriter, loc, TypeRange {outputStorageType}, laneCount, {}, ValueRange {storage}, + [&](detail::SpatComputeBatchBodyArgs args) { + Operation* anchor = rewriter.getInsertionBlock()->getParentOp(); + Value outputRow = affineFloorDivConst(rewriter, loc, args.lane, tilesPerRow, anchor); + Value tile = affineModConst(rewriter, loc, args.lane, tilesPerRow, anchor); + Value inputRow = buildNearestAsymmetricIndex( + outputRow, inputType.getDimSize(2), outputHeight, rewriter, loc); + Value inputSlot = arith::AddIOp::create( + rewriter, loc, affineMulConst(rewriter, loc, inputRow, tilesPerRow, anchor), tile); + auto source = extractGraphBatchPhysicalFragment( + rewriter, loc, args.inputs.front(), inputSlot, input->fragmentType); + if (failed(source)) return failure(); - batchYielded.push_back(channelLoop->results.front()); + Value initial = tensor::EmptyOp::create( + rewriter, loc, outputFragmentType.getShape(), resultType.getElementType()); + Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0); + Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1); + Value width = getOrCreateIndexConstant(rewriter, anchor, outputWidth); + auto loop = buildNormalizedScfFor( + rewriter, loc, c0, width, c1, ValueRange {initial}, + [&](OpBuilder&, Location nestedLoc, Value outputColumn, ValueRange iterArgs, + SmallVectorImpl& yielded) { + Value inputColumn = buildNearestAsymmetricIndex( + outputColumn, inputType.getDimSize(3), outputWidth, rewriter, nestedLoc); + Value pixel = tensor::ExtractSliceOp::create( + rewriter, nestedLoc, pixelType, *source, + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), + inputColumn, rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)}, + getUnitStrides(rewriter, 4)); + yielded.push_back(tensor::InsertSliceOp::create( + rewriter, nestedLoc, pixel, iterArgs.front(), + SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), + outputColumn, rewriter.getIndexAttr(0)}, + SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)}, + getUnitStrides(rewriter, 4))); + return success(); + }); + if (failed(loop)) + return failure(); + publishGraphBatchPhysicalFragment( + rewriter, loc, loop->results.front(), args.outputs.front(), args.lane); return success(); }); - if (failed(batchLoop)) - return failure(); - return batchLoop->results.front(); + return failed(batch) ? FailureOr(failure()) + : FailureOr(batch->getResult(0)); } struct Resize : OpConversionPattern { @@ -161,23 +182,38 @@ struct Resize : OpConversionPattern { || llvm::any_of(resultType.getShape(), [](int64_t dim) { return dim <= 0; })) return rewriter.notifyMatchFailure(resizeOp, "resize lowering requires positive static dimensions."); - auto computeOp = createSpatCompute<1>( - rewriter, resizeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getX(), [&](Value x) -> LogicalResult { - auto result = buildNearestResizeLoop(x, inputType, resultType, rewriter, resizeOp.getLoc()); - if (failed(result)) - return failure(); - spatial::SpatYieldOp::create(rewriter, resizeOp.getLoc(), *result); - return success(); - }); - if (failed(computeOp)) - return failure(); - rewriter.replaceOp(resizeOp, computeOp->getResults()); + auto plan = spatial::SpatResizeNearestPlanOp::create( + rewriter, resizeOp.getLoc(), resultType, adaptor.getX(), rewriter.getStringAttr("nchw")); + rewriter.replaceOp(resizeOp, plan.getResult()); return success(); } }; } // namespace +LogicalResult canLowerResizeNearestPlanToRowStrip( + spatial::SpatResizeNearestPlanOp planOp) { + auto inputType = dyn_cast(planOp.getInput().getType()); + auto outputType = dyn_cast(planOp.getOutput().getType()); + return success(inputType && outputType && inputType.hasStaticShape() + && outputType.hasStaticShape() && inputType.getRank() == 4 + && outputType.getRank() == 4 && inputType.getDimSize(0) == 1 + && outputType.getDimSize(0) == 1 + && inputType.getDimSize(1) == outputType.getDimSize(1)); +} + +FailureOr lowerSelectedResizeNearestPlan( + spatial::SpatResizeNearestPlanOp planOp, std::optional rowStripInput, + PatternRewriter& rewriter) { + auto inputType = cast(planOp.getInput().getType()); + auto outputType = cast(planOp.getOutput().getType()); + if (rowStripInput) + return buildRowStripNearestResize( + *rowStripInput, inputType, outputType, rewriter, planOp.getLoc()); + return buildDenseNearestResize( + planOp.getInput(), inputType, outputType, rewriter, planOp.getLoc()); +} + void populateResizePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add(ctx); } } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp index 4a37d9e..2088717 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp @@ -61,6 +61,74 @@ static FailureOr materializeTransposedConstant(Value input, resultType); } +static FailureOr transposeFragmentAssemblyBlueprint(spatial::SpatBlueprintOp blueprint, + RankedTensorType resultType, + ArrayRef permutation, + ConversionPatternRewriter& rewriter, + Location loc) { + auto storageType = dyn_cast(blueprint.getInput().getType()); + auto sourceOffsets = blueprint.getFragmentSourceOffsets(); + auto fragmentStrides = blueprint.getFragmentStrides(); + if (!storageType || !storageType.hasStaticShape() || !resultType.hasStaticShape() + || !blueprint.getFragments().empty() || blueprint.getMode() != "fragment_assembly" + || !blueprint.getFragmentOperandIndices() || !sourceOffsets || !fragmentStrides + || llvm::any_of(*sourceOffsets, [](int64_t offset) { return offset != 0; }) + || storageType.getRank() != resultType.getRank() + 1) + return failure(); + if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments + && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint)) + return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure(); + + SmallVector outputStorageShape {storageType.getDimSize(0)}; + for (int64_t sourceDim : permutation) + outputStorageShape.push_back(storageType.getDimSize(sourceDim + 1)); + auto outputStorageType = RankedTensorType::get(outputStorageShape, storageType.getElementType()); + auto mapped = mapGraphBatchFragments( + blueprint.getInput(), outputStorageType, rewriter, loc, [&](Value fragment, RankedTensorType fragmentType) { + Value init = createTransposeInit(fragment, fragmentType, permutation, rewriter, loc); + return FailureOr( + linalg::TransposeOp::create(rewriter, loc, fragment, init, permutation).getResult()[0]); + }); + if (failed(mapped)) + return failure(); + + const int64_t rank = resultType.getRank(); + const int64_t fragmentCount = blueprint.getFragmentOperandIndices()->size(); + SmallVector offsets, sizes, strides; + offsets.reserve(fragmentCount * rank); + sizes.reserve(fragmentCount * rank); + strides.reserve(fragmentCount * rank); + ArrayRef inputOffsets = blueprint.getFragmentOffsets(); + ArrayRef inputSizes = blueprint.getFragmentSizes(); + for (int64_t fragment = 0; fragment < fragmentCount; ++fragment) + for (int64_t sourceDim : permutation) { + const int64_t index = fragment * rank + sourceDim; + offsets.push_back(inputOffsets[index]); + sizes.push_back(inputSizes[index]); + strides.push_back((*fragmentStrides)[index]); + } + auto transposedBlueprint = spatial::SpatBlueprintOp::create(rewriter, + loc, + resultType, + *mapped, + ValueRange {}, + blueprint.getLogicalLayoutAttr(), + rewriter.getStringAttr("fragmented"), + rewriter.getDenseI64ArrayAttr(offsets), + rewriter.getDenseI64ArrayAttr(sizes), + rewriter.getStringAttr("permuted_fragments"), + blueprint.getModeAttr(), + blueprint.getFragmentOperandIndicesAttr(), + blueprint.getFragmentSourceSlotsAttr(), + blueprint.getFragmentSourceOffsetsAttr(), + rewriter.getDenseI64ArrayAttr(strides), + blueprint.getConflictPolicyAttr(), + blueprint.getCoveragePolicyAttr()); + if (spatial::isCanonicalContiguousRowMajorFragmentAssembly(transposedBlueprint)) + transposedBlueprint.setIndexMapAttr(rewriter.getStringAttr(spatial::kContiguousRowMajorFragments)); + return transposedBlueprint.getOutput(); +} + struct TransposeToLinalgTranspose : OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -75,6 +143,14 @@ struct TransposeToLinalgTranspose : OpConversionPattern { auto permutation = getTransposePermutationChecked(transposeOp.getPermAttr(), inputType.getRank()); if (failed(permutation)) return failure(); + if (auto blueprint = adaptor.getData().getDefiningOp()) { + auto transposed = + transposeFragmentAssemblyBlueprint(blueprint, resultType, *permutation, rewriter, transposeOp.getLoc()); + if (succeeded(transposed)) { + rewriter.replaceOp(transposeOp, *transposed); + return success(); + } + } if (isCompileTimeComputable(adaptor.getData())) { auto constantTranspose = materializeTransposedConstant(adaptor.getData(), resultType, *permutation, rewriter, transposeOp.getLoc()); diff --git a/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp b/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp index be2308e..7a952a5 100644 --- a/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp @@ -20,6 +20,14 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp); mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp); +mlir::LogicalResult canLowerResizeNearestPlanToRowStrip( + spatial::SpatResizeNearestPlanOp planOp); + +mlir::FailureOr lowerSelectedResizeNearestPlan( + spatial::SpatResizeNearestPlanOp planOp, + std::optional rowStripInput, + mlir::PatternRewriter& rewriter); + mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp); mlir::FailureOr diff --git a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp index df69231..0b39303 100644 --- a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp @@ -36,6 +36,8 @@ static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap(user)) return getSelectedLayout(layouts, siluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; + if (auto resizePlan = dyn_cast(user)) + return getSelectedLayout(layouts, resizePlan.getResult()) == SelectedLayout::PixelMajorRowStrip; if (auto biasAddPlan = dyn_cast(user)) return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; if (auto addPlan = dyn_cast(user)) @@ -66,6 +68,8 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap(user)) return true; + if (auto resizePlan = dyn_cast(user)) + return succeeded(canLowerResizeNearestPlanToRowStrip(resizePlan)); if (auto biasAddPlan = dyn_cast(user)) { auto resultType = dyn_cast(biasAddPlan.getOutput().getType()); return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType); @@ -117,6 +121,14 @@ static SelectedLayout chooseActivationLayout(Value input, return SelectedLayout::PixelMajorRowStrip; } +static SelectedLayout chooseResizeLayout( + spatial::SpatResizeNearestPlanOp resizePlan, + llvm::DenseMap& layouts) { + return getSelectedLayout(layouts, resizePlan.getInput()) == SelectedLayout::PixelMajorRowStrip + && succeeded(canLowerResizeNearestPlanToRowStrip(resizePlan)) + ? SelectedLayout::PixelMajorRowStrip : SelectedLayout::DenseNchw; +} + static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan, llvm::DenseMap& layouts) { if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip) @@ -255,6 +267,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper(&op)) { + SelectedLayout selected = chooseResizeLayout(resizePlan, layouts); + if (layouts[resizePlan.getResult()] != selected) { + layouts[resizePlan.getResult()] = selected; + changed = true; + } + continue; + } if (auto biasAddPlan = dyn_cast(&op)) { SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts); if (layouts[biasAddPlan.getResult()] != selected) { @@ -312,6 +332,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper(&op)) producedValue = siluPlan.getResult(); + else if (auto resizePlan = dyn_cast(&op)) + producedValue = resizePlan.getResult(); else if (auto maxPoolPlan = dyn_cast(&op)) producedValue = maxPoolPlan.getResult(); else if (auto averagePoolPlan = dyn_cast(&op)) diff --git a/src/PIM/Conversion/SpatialToPim/Common.cpp b/src/PIM/Conversion/SpatialToPim/Common.cpp index 8e3d840..e4eddac 100644 --- a/src/PIM/Conversion/SpatialToPim/Common.cpp +++ b/src/PIM/Conversion/SpatialToPim/Common.cpp @@ -129,6 +129,32 @@ LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprin return success(); } +FailureOr reshapeContiguousRowMajorFragments(RewriterBase& rewriter, + Location loc, + mlir::Value source, + RankedTensorType resultType) { + auto sourceType = dyn_cast(source.getType()); + if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || resultType.getRank() < 2 + || sourceType.getRank() != resultType.getRank() + 1 || sourceType.getElementType() != resultType.getElementType() + || sourceType.getNumElements() != resultType.getNumElements() + || sourceType.getDimSize(0) != getStaticShapeElementCount(resultType.getShape().drop_back()) + || sourceType.getDimSize(sourceType.getRank() - 1) != resultType.getDimSize(resultType.getRank() - 1) + || llvm::any_of(sourceType.getShape().slice(1, sourceType.getRank() - 2), [](int64_t dim) { return dim != 1; })) + return failure(); + + SmallVector collapse {{}, {sourceType.getRank() - 1}}; + for (int64_t dim = 0; dim < sourceType.getRank() - 1; ++dim) + collapse.front().push_back(dim); + auto flatType = RankedTensorType::get( + {sourceType.getDimSize(0), sourceType.getDimSize(sourceType.getRank() - 1)}, resultType.getElementType()); + mlir::Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, source, collapse); + + SmallVector expand {{}, {resultType.getRank() - 1}}; + for (int64_t dim = 0; dim < resultType.getRank() - 1; ++dim) + expand.front().push_back(dim); + return tensor::ExpandShapeOp::create(rewriter, loc, resultType, flat, expand).getResult(); +} + static SmallVector expandFlatElementIndex(int64_t flatIndex, ArrayRef shape) { SmallVector indices(shape.size(), 0); for (int64_t dim = static_cast(shape.size()) - 1; dim >= 0; --dim) { diff --git a/src/PIM/Conversion/SpatialToPim/Common.hpp b/src/PIM/Conversion/SpatialToPim/Common.hpp index 481976c..61e74a4 100644 --- a/src/PIM/Conversion/SpatialToPim/Common.hpp +++ b/src/PIM/Conversion/SpatialToPim/Common.hpp @@ -51,6 +51,11 @@ mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlu llvm::ArrayRef flatSizes, llvm::ArrayRef flatStrides); +mlir::FailureOr reshapeContiguousRowMajorFragments(mlir::RewriterBase& rewriter, + mlir::Location loc, + mlir::Value source, + mlir::RankedTensorType resultType); + mlir::FailureOr> getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor, mlir::ShapedType sourceType, diff --git a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp b/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp index 8b1ad80..f2cb16d 100644 --- a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp @@ -71,6 +71,16 @@ static FailureOr lowerFragmentAssemblyBlueprint(IRRewriter& rewriter, flatStrides))) return failure(); + if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments) { + if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint)) + return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure(); + Value source = mapping.lookupOrDefault(blueprint.getInput()); + auto reshaped = reshapeContiguousRowMajorFragments( + rewriter, blueprint.getLoc(), source, cast(resultType)); + if (failed(reshaped)) + return blueprint.emitOpError("contiguous row-major fragment storage does not match its logical result"), failure(); + return *reshaped; + } SmallVector hostStrides = computeRowMajorStrides(resultType.getShape()); SmallVector copies; for (int64_t fragmentIndex = 0; fragmentIndex < static_cast(operandIndices.size()); ++fragmentIndex) { diff --git a/src/PIM/Conversion/SpatialToPim/Patterns.cpp b/src/PIM/Conversion/SpatialToPim/Patterns.cpp index 68ede56..c450e0b 100644 --- a/src/PIM/Conversion/SpatialToPim/Patterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/Patterns.cpp @@ -49,6 +49,16 @@ struct LowerFragmentAssemblyBlueprintPattern op, rank, fragmentOperands.size(), operandIndices, sourceOffsets, flatOffsets, flatSizes, flatStrides))) return failure(); + if (op.getIndexMap() == spatial::kContiguousRowMajorFragments) { + if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(op)) + return op.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"); + auto reshaped = reshapeContiguousRowMajorFragments( + rewriter, op.getLoc(), adaptor.getInput(), cast(resultType)); + if (failed(reshaped)) + return op.emitOpError("contiguous row-major fragment storage does not match its logical result"); + rewriter.replaceOp(op, *reshaped); + return success(); + } Value currentOutput = tensor::EmptyOp::create(rewriter, op.getLoc(), resultType.getShape(), resultType.getElementType()).getResult(); for (int64_t fragmentIndex = 0; fragmentIndex < static_cast(operandIndices.size()); ++fragmentIndex) { diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp b/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp index ad557d3..cb4e22b 100644 --- a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp +++ b/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp @@ -532,54 +532,74 @@ struct FoldConstantMemCpPattern final : OpRewritePattern { } }; -static bool isOne(Attribute value) { - if (auto floatValue = dyn_cast(value)) - return floatValue.getValue().isExactlyValue(1.0); - if (auto integerValue = dyn_cast(value)) - return integerValue.getValue() == 1; - return false; +enum class MultiplicationConstant { Other, Zero, One }; + +static MultiplicationConstant classifyMultiplicationConstant(Attribute value) { + if (auto floatValue = dyn_cast(value)) { + const APFloat& number = floatValue.getValue(); + if (number.isZero() && !number.isNegative()) + return MultiplicationConstant::Zero; + if (number.isExactlyValue(1.0)) + return MultiplicationConstant::One; + } + if (auto integerValue = dyn_cast(value)) { + if (integerValue.getValue().isZero()) + return MultiplicationConstant::Zero; + if (integerValue.getValue() == 1) + return MultiplicationConstant::One; + } + return MultiplicationConstant::Other; } -static bool isAllOneHostCopy(pim::PimMemCopyHostToDevOp copyOp, ModuleOp moduleOp, MemRefType copiedType) { +static MultiplicationConstant classifyUniformHostCopy( + pim::PimMemCopyHostToDevOp copyOp, ModuleOp moduleOp, MemRefType copiedType) { auto targetOffset = resolveIndexValue(copyOp.getDeviceTargetOffset()); auto sourceOffset = resolveIndexValue(copyOp.getHostSourceOffset()); if (failed(targetOffset) || failed(sourceOffset) || *targetOffset != 0) - return false; + return MultiplicationConstant::Other; Type elementType = copiedType.getElementType(); if (!elementType.isIntOrFloat()) - return false; + return MultiplicationConstant::Other; unsigned bitWidth = elementType.getIntOrFloatBitWidth(); if (bitWidth == 0 || bitWidth % 8 != 0) - return false; + return MultiplicationConstant::Other; int64_t elementBytes = bitWidth / 8; int64_t copiedElements = copiedType.getNumElements(); if (*sourceOffset % elementBytes != 0 || copyOp.getSize() != copiedElements * elementBytes) - return false; + return MultiplicationConstant::Other; auto source = getDenseGlobalValue(moduleOp, copyOp.getHostSource()); if (failed(source) || source->getElementType() != elementType) - return false; + return MultiplicationConstant::Other; int64_t firstElement = *sourceOffset / elementBytes; int64_t endElement = firstElement + copiedElements; if (firstElement < 0 || endElement > source->getNumElements()) - return false; + return MultiplicationConstant::Other; if (source->isSplat()) - return isOne(source->getSplatValue()); + return classifyMultiplicationConstant(source->getSplatValue()); + MultiplicationConstant classification = MultiplicationConstant::Other; int64_t index = 0; for (Attribute value : source->getValues()) { - if (index >= firstElement && index < endElement && !isOne(value)) - return false; + if (index >= firstElement && index < endElement) { + MultiplicationConstant current = classifyMultiplicationConstant(value); + if (current == MultiplicationConstant::Other) + return current; + if (classification == MultiplicationConstant::Other) + classification = current; + else if (classification != current) + return MultiplicationConstant::Other; + } if (++index >= endElement) break; } - return true; + return classification; } -struct FoldMultiplyByOnePattern final : OpRewritePattern { +struct FoldMultiplyByConstantPattern final : OpRewritePattern { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(pim::PimVVMulOp mulOp, PatternRewriter& rewriter) const override { @@ -605,14 +625,19 @@ struct FoldMultiplyByOnePattern final : OpRewritePattern { copyOp = candidate; } auto maskType = dyn_cast(mask.getType()); - if (!copyOp || !copyOp.use_empty() || !maskType || !isAllOneHostCopy(copyOp, moduleOp, maskType)) + if (!copyOp || !copyOp.use_empty() || !maskType) + continue; + MultiplicationConstant constant = classifyUniformHostCopy(copyOp, moduleOp, maskType); + if (constant == MultiplicationConstant::Other) continue; auto outputAlloc = mulOp.getOutputBuffer().getDefiningOp(); - rewriter.replaceOp(mulOp, input); - rewriter.eraseOp(copyOp); - if (maskAlloc.use_empty()) - rewriter.eraseOp(maskAlloc); + rewriter.replaceOp(mulOp, constant == MultiplicationConstant::One ? input : mask); + if (constant == MultiplicationConstant::One) { + rewriter.eraseOp(copyOp); + if (maskAlloc.use_empty()) + rewriter.eraseOp(maskAlloc); + } if (outputAlloc && outputAlloc.use_empty()) rewriter.eraseOp(outputAlloc); return success(); @@ -629,7 +654,7 @@ void populateConstantFoldingConstantPatterns(RewritePatternSet& patterns) { FoldConstantCoreMapPattern, FoldConstantHostCopyPattern, FoldConstantMemCpPattern, - FoldMultiplyByOnePattern>(patterns.getContext()); + FoldMultiplyByConstantPattern>(patterns.getContext()); } } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Spatial.td b/src/PIM/Dialect/Spatial/Spatial.td index 00af2bc..22f099f 100644 --- a/src/PIM/Dialect/Spatial/Spatial.td +++ b/src/PIM/Dialect/Spatial/Spatial.td @@ -303,6 +303,21 @@ def SpatSiluPlanOp : SpatOp<"silu_plan", []> { let hasVerifier = 1; } +def SpatResizeNearestPlanOp : SpatOp<"resize_nearest_plan", []> { + let summary = "Layout-aware nearest asymmetric Resize planning op"; + + let arguments = (ins + SpatTensor:$input, + StrAttr:$logicalLayout + ); + + let results = (outs + SpatTensor:$output + ); + + let hasVerifier = 1; +} + def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> { let summary = "Layout-aware 2D NCHW MaxPool planning op"; diff --git a/src/PIM/Dialect/Spatial/SpatialOps.cpp b/src/PIM/Dialect/Spatial/SpatialOps.cpp index 6acd8b5..166b8a8 100644 --- a/src/PIM/Dialect/Spatial/SpatialOps.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOps.cpp @@ -11,6 +11,70 @@ using namespace mlir; namespace onnx_mlir { namespace spatial { +bool hasCanonicalContiguousRowMajorFragments(RankedTensorType logicalType, + ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides) { + if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() <= 0 + || logicalType.getDimSize(logicalType.getRank() - 1) <= 0) + return false; + const int64_t rank = logicalType.getRank(); + const int64_t rowCount = logicalType.getNumElements() / logicalType.getDimSize(rank - 1); + if (offsets.size() != static_cast(rowCount * rank) || sizes.size() != offsets.size() + || strides.size() != offsets.size()) + return false; + for (int64_t row = 0; row < rowCount; ++row) { + int64_t remaining = row; + for (int64_t dim = rank - 2; dim >= 0; --dim) { + const int64_t index = row * rank + dim; + if (offsets[index] != remaining % logicalType.getDimSize(dim) || sizes[index] != 1 || strides[index] != 1) + return false; + remaining /= logicalType.getDimSize(dim); + } + const int64_t last = row * rank + rank - 1; + if (offsets[last] != 0 || sizes[last] != logicalType.getDimSize(rank - 1) || strides[last] != 1) + return false; + } + return true; +} + +bool isCanonicalContiguousRowMajorFragmentAssembly(SpatBlueprintOp blueprint) { + auto logicalType = dyn_cast(blueprint.getOutput().getType()); + auto physicalType = dyn_cast(blueprint.getInput().getType()); + auto operandIndices = blueprint.getFragmentOperandIndices(); + auto sourceSlots = blueprint.getFragmentSourceSlots(); + auto sourceOffsets = blueprint.getFragmentSourceOffsets(); + auto fragmentStrides = blueprint.getFragmentStrides(); + if (!logicalType || !physicalType || !logicalType.hasStaticShape() || !physicalType.hasStaticShape() + || logicalType.getRank() < 2 || !blueprint.getFragments().empty() + || blueprint.getMode() != "fragment_assembly" || !operandIndices || !sourceSlots || !sourceOffsets + || !fragmentStrides) + return false; + + ArrayRef offsets = blueprint.getFragmentOffsets(); + ArrayRef sizes = blueprint.getFragmentSizes(); + if (!hasCanonicalContiguousRowMajorFragments(logicalType, offsets, sizes, *fragmentStrides) + || operandIndices->empty() || operandIndices->size() != sourceSlots->size() + || operandIndices->size() != sourceOffsets->size() + || operandIndices->size() * static_cast(logicalType.getRank()) != offsets.size() + || physicalType.getRank() != logicalType.getRank() + 1 + || physicalType.getDimSize(0) != static_cast(operandIndices->size()) + || physicalType.getDimSize(0) + != logicalType.getNumElements() / logicalType.getDimSize(logicalType.getRank() - 1) + || physicalType.getElementType() != logicalType.getElementType() + || physicalType.getNumElements() != logicalType.getNumElements() + || physicalType.getDimSize(physicalType.getRank() - 1) != logicalType.getDimSize(logicalType.getRank() - 1) + || llvm::any_of(physicalType.getShape().slice(1, physicalType.getRank() - 2), + [](int64_t dim) { return dim != 1; })) + return false; + + for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) + if (operandIndex != 0 || (*sourceSlots)[fragmentIndex] != static_cast(fragmentIndex) + || (*sourceOffsets)[fragmentIndex] != 0) + return false; + return true; +} + RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, RankedTensorType fragmentType) { SmallVector shape {laneCount}; llvm::append_range(shape, fragmentType.getShape()); diff --git a/src/PIM/Dialect/Spatial/SpatialOps.hpp b/src/PIM/Dialect/Spatial/SpatialOps.hpp index 222766e..0953af3 100644 --- a/src/PIM/Dialect/Spatial/SpatialOps.hpp +++ b/src/PIM/Dialect/Spatial/SpatialOps.hpp @@ -30,6 +30,15 @@ namespace onnx_mlir { namespace spatial { +inline constexpr llvm::StringLiteral kContiguousRowMajorFragments = "contiguous_row_major_fragments"; + +bool hasCanonicalContiguousRowMajorFragments(mlir::RankedTensorType logicalType, + llvm::ArrayRef offsets, + llvm::ArrayRef sizes, + llvm::ArrayRef strides); + +bool isCanonicalContiguousRowMajorFragmentAssembly(SpatBlueprintOp blueprint); + mlir::RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, mlir::RankedTensorType fragmentType); mlir::FailureOr getGraphBatchFragmentType(mlir::RankedTensorType physicalType, int64_t expectedLaneCount); diff --git a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp index 09f7eed..7001d62 100644 --- a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp +++ b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp @@ -480,6 +480,23 @@ LogicalResult SpatSiluPlanOp::verify() { return success(); } +LogicalResult SpatResizeNearestPlanOp::verify() { + if (failed(verifyPlanTensorTypes( + getOperation(), getInput(), getOutput(), "spat.resize_nearest_plan"))) + return failure(); + auto inputType = dyn_cast(getInput().getType()); + auto outputType = dyn_cast(getOutput().getType()); + if (!inputType.hasStaticShape() || !outputType.hasStaticShape() + || inputType.getRank() != 4 || outputType.getRank() != 4) + return emitError("requires static rank-4 input and output tensors"); + if (getLogicalLayout() != "nchw") + return emitError("requires logical layout \"nchw\""); + if (llvm::any_of(inputType.getShape(), [](int64_t dim) { return dim <= 0; }) + || llvm::any_of(outputType.getShape(), [](int64_t dim) { return dim <= 0; })) + return emitError("requires positive dimensions"); + return success(); +} + LogicalResult SpatMaxPool2DPlanOp::verify() { if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.max_pool2d_plan"))) return failure(); @@ -587,8 +604,12 @@ LogicalResult SpatBlueprintOp::verify() { if (offsets.size() != sizes.size()) return emitError("fragment offset and size arrays must have the same length"); int64_t rank = logicalType.getRank(); - if (offsets.empty()) + bool isContiguousRowMajor = getIndexMap() == kContiguousRowMajorFragments; + if (offsets.empty()) { + if (isContiguousRowMajor) + return emitError("contiguous row-major fragment destination geometry is not canonical"); return success(); + } if (rank <= 0 || offsets.size() % rank != 0) return emitError("fragment metadata must be a whole number of rank-sized fragments"); @@ -611,6 +632,8 @@ LogicalResult SpatBlueprintOp::verify() { }; if (!isFragmentAssembly) { + if (isContiguousRowMajor) + return emitError("contiguous row-major fragments require fragment assembly metadata"); if (failed(verifyBoundsOnly({}))) return failure(); if (!getFragments().empty()) @@ -662,6 +685,13 @@ LogicalResult SpatBlueprintOp::verify() { if (failed(verifyBoundsOnly(strides))) return failure(); + if (isContiguousRowMajor) { + if (!hasCanonicalContiguousRowMajorFragments(logicalType, offsets, sizes, strides)) + return emitError("contiguous row-major fragment destination geometry is not canonical"); + if (!isCanonicalContiguousRowMajorFragmentAssembly(*this)) + return emitError("contiguous row-major fragment physical source order or storage is not canonical"); + } + SmallVector, SmallVector>, 8> slices; slices.reserve(static_cast(fragmentCount)); SmallVector fragmentCountsByOperand(static_cast(operandCount), 0); @@ -713,20 +743,22 @@ LogicalResult SpatBlueprintOp::verify() { if (sourceSliceOffsets[dim] + fragmentSizes[dim] > fragmentType.getDimSize(dim)) return emitError("fragment assembly source offset must describe a valid unit-stride slice"); - for (const auto& [existingOffsets, existingSizes] : slices) { - bool overlaps = true; - for (int64_t dim = 0; dim < rank; ++dim) { - int64_t begin = fragmentOffsets[dim]; - int64_t end = begin + fragmentSizes[dim]; - int64_t existingBegin = existingOffsets[dim]; - int64_t existingEnd = existingBegin + existingSizes[dim]; - if (end <= existingBegin || existingEnd <= begin) { - overlaps = false; - break; + if (!isContiguousRowMajor) { + for (const auto& [existingOffsets, existingSizes] : slices) { + bool overlaps = true; + for (int64_t dim = 0; dim < rank; ++dim) { + int64_t begin = fragmentOffsets[dim]; + int64_t end = begin + fragmentSizes[dim]; + int64_t existingBegin = existingOffsets[dim]; + int64_t existingEnd = existingBegin + existingSizes[dim]; + if (end <= existingBegin || existingEnd <= begin) { + overlaps = false; + break; + } } + if (overlaps) + return emitError("fragment assembly blueprint requires disjoint static slices"); } - if (overlaps) - return emitError("fragment assembly blueprint requires disjoint static slices"); } slices.push_back({std::move(fragmentOffsets), std::move(fragmentSizes)}); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp index e6846b2..0fc3985 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp @@ -9,6 +9,7 @@ #include "llvm/ADT/SmallPtrSet.h" #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp" #include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp" #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp" @@ -37,6 +38,63 @@ static SmallVector getBlueprintFragments(SpatBlueprintOp blueprint) { return fragments; } +static FailureOr buildContiguousRowMajorReconstruction( + OpBuilder &builder, Location loc, SpatBlueprintOp blueprint, + Value source) { + auto resultType = dyn_cast(blueprint.getOutput().getType()); + auto sourceType = dyn_cast(source.getType()); + if (!resultType || !sourceType || !resultType.hasStaticShape() + || !sourceType.hasStaticShape() || resultType.getRank() <= 0 + || sourceType.getRank() != resultType.getRank() + 1) + return failure(); + int64_t rank = resultType.getRank(); + int64_t width = resultType.getDimSize(rank - 1); + int64_t rowCount = resultType.getNumElements() / width; + if (sourceType.getDimSize(0) != rowCount || sourceType.getDimSize(rank) != width) + return failure(); + + auto flatType = RankedTensorType::get({rowCount, width}, resultType.getElementType()); + auto rowType = RankedTensorType::get({1, width}, resultType.getElementType()); + auto physicalRowType = RankedTensorType::get(sourceType.getShape().drop_front(), resultType.getElementType()); + Value init = tensor::EmptyOp::create(builder, loc, flatType.getShape(), flatType.getElementType()); + Value c0 = arith::ConstantIndexOp::create(builder, loc, 0); + Value c1 = arith::ConstantIndexOp::create(builder, loc, 1); + Value rows = arith::ConstantIndexOp::create(builder, loc, rowCount); + auto loop = buildNormalizedScfFor( + builder, loc, c0, rows, c1, ValueRange {init}, + [&](OpBuilder &nested, Location nestedLoc, Value row, ValueRange iterArgs, + SmallVectorImpl &yielded) { + SmallVector offsets {row}; + SmallVector sizes {nested.getIndexAttr(1)}; + SmallVector strides {nested.getIndexAttr(1)}; + for (int64_t dim : physicalRowType.getShape()) { + offsets.push_back(nested.getIndexAttr(0)); + sizes.push_back(nested.getIndexAttr(dim)); + strides.push_back(nested.getIndexAttr(1)); + } + Value physicalRow = tensor::ExtractSliceOp::create( + nested, nestedLoc, physicalRowType, source, offsets, sizes, strides); + SmallVector collapse {{}}; + for (int64_t dim = 0; dim < rank - 1; ++dim) + collapse.front().push_back(dim); + collapse.push_back({rank - 1}); + Value flatRow = tensor::CollapseShapeOp::create( + nested, nestedLoc, rowType, physicalRow, collapse); + yielded.push_back(tensor::InsertSliceOp::create( + nested, nestedLoc, flatRow, iterArgs.front(), + SmallVector {row, nested.getIndexAttr(0)}, + SmallVector {nested.getIndexAttr(1), nested.getIndexAttr(width)}, + SmallVector {nested.getIndexAttr(1), nested.getIndexAttr(1)})); + return success(); + }); + if (failed(loop)) + return failure(); + SmallVector expand {{}, {rank - 1}}; + for (int64_t dim = 0; dim < rank - 1; ++dim) + expand.front().push_back(dim); + return tensor::ExpandShapeOp::create(builder, loc, resultType, loop->results.front(), expand).getResult(); +} + static FailureOr buildBlueprintReconstruction( OpBuilder &builder, Location loc, SpatBlueprintOp blueprint, ValueRange sourceBlockArgs) { @@ -57,6 +115,13 @@ static FailureOr buildBlueprintReconstruction( sourceOffsets->size() != operandIndices->size()) return blueprint.emitOpError("phase 1 fragment assembly metadata has inconsistent sizes"), failure(); + if (blueprint.getIndexMap() == kContiguousRowMajorFragments) { + if (!isCanonicalContiguousRowMajorFragmentAssembly(blueprint)) + return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure(); + if (sourceBlockArgs.size() != 1) + return blueprint.emitOpError("contiguous row-major fragment reconstruction requires one physical source"), failure(); + return buildContiguousRowMajorReconstruction(builder, loc, blueprint, sourceBlockArgs.front()); + } Value result = tensor::EmptyOp::create(builder, loc, resultType.getShape(), resultType.getElementType()); for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) { @@ -178,6 +243,15 @@ static Operation *getTopLevelDeferredOperation( return op && isTopLevelDeferredOperation(op, body, plan) ? op : nullptr; } +static bool isDefinedInside(Operation *owner, Value value) { + if (Operation *definition = value.getDefiningOp()) + return owner->isProperAncestor(definition); + auto argument = dyn_cast(value); + Region *region = argument ? argument.getOwner()->getParent() : nullptr; + return region == &owner->getRegion(0) + || (region && owner->getRegion(0).isAncestor(region)); +} + static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan, llvm::SmallPtrSetImpl &seen) { if (value == plan.graphInput || value == plan.graphLane || value == plan.scheduledLane) @@ -195,18 +269,10 @@ static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan, loop.getRegion().walk([&](Operation *nested) { if (isa(nested) && nested != loop) eligible = false; - for (Value operand : nested->getOperands()) { - Operation *definition = operand.getDefiningOp(); - auto argument = dyn_cast(operand); - Region *argumentRegion = argument - ? argument.getOwner()->getParent() : nullptr; - bool definedInside = definition - ? loop->isProperAncestor(definition) - : argumentRegion == &loop.getRegion() - || (argumentRegion && loop.getRegion().isAncestor(argumentRegion)); - if (!definedInside && !isEligible(operand, body, plan, seen)) + for (Value operand : nested->getOperands()) + if (!isDefinedInside(loop, operand) + && !isEligible(operand, body, plan, seen)) eligible = false; - } }); if (!eligible) return false; @@ -267,19 +333,9 @@ static FailureOr clonePayloadRoot(Value root, Block &body, const Deferred if (auto loop = dyn_cast(op)) { SmallVector captures; loop.getRegion().walk([&](Operation *nested) { - for (Value operand : nested->getOperands()) { - Operation *definition = operand.getDefiningOp(); - auto argument = dyn_cast(operand); - Region *argumentRegion = argument - ? argument.getOwner()->getParent() : nullptr; - bool definedInside = definition - ? loop->isProperAncestor(definition) - : argumentRegion == &loop.getRegion() - || (argumentRegion - && loop.getRegion().isAncestor(argumentRegion)); - if (!definedInside && !mapping.contains(operand)) + for (Value operand : nested->getOperands()) + if (!isDefinedInside(loop, operand) && !mapping.contains(operand)) captures.push_back(operand); - } }); for (Value capture : captures) if (!mapping.contains(capture) && failed(clone(capture))) @@ -304,7 +360,10 @@ static bool dependsOnGraphLane(Value value, Value graphLane, Block &body, if (auto loop = dyn_cast(op)) { bool depends = false; loop.getRegion().walk([&](Operation *nested) { - depends |= llvm::is_contained(nested->getOperands(), graphLane); + for (Value operand : nested->getOperands()) + if (!isDefinedInside(loop, operand) + && dependsOnGraphLane(operand, graphLane, body, plan, seen)) + depends = true; }); if (depends) return true; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp index 8af6bf9..cdf9f11 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp @@ -14,6 +14,26 @@ namespace onnx_mlir::spatial { using namespace mlir; namespace { +static LogicalResult verifyNoEscapingRegionValues(Operation* owner, StringRef phase) { + Operation* escapingDefinition = nullptr; + Operation* escapingUser = nullptr; + owner->walk([&](Operation* nested) { + for (Value result : nested->getResults()) + for (Operation* user : result.getUsers()) + if (!owner->isProperAncestor(user)) { + escapingDefinition = nested; + escapingUser = user; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (!escapingDefinition) + return success(); + return owner->emitOpError() << phase << " left a value defined by " << escapingDefinition->getName() + << " at " << escapingDefinition->getLoc() << " captured by " + << escapingUser->getName() << " at " << escapingUser->getLoc(); +} + static LogicalResult placeLogicalProcessorsOnPhysicalCores(DeferredTransferPlan& plan, const SchedulingTarget& target) { std::vector logicalTrafficFlits(target.processorCount * target.processorCount, 0); for (const std::unique_ptr& exchange : plan.exchanges) @@ -138,6 +158,8 @@ static LogicalResult eraseOldGraph(func::FuncOp funcOp, IRRewriter& rewriter) { } } } + if (failed(verifyNoEscapingRegionValues(op, "phase 2"))) + return failure(); rewriter.eraseOp(op); } return success(); @@ -217,6 +239,8 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp, op->getResult(0).replaceAllUsesWith(replacement); if (!op->use_empty()) return op->emitOpError("phase 2 cannot erase deferred communication with live uses"); + if (failed(verifyNoEscapingRegionValues(op, "phase 2 deferred communication"))) + return failure(); rewriter.eraseOp(op); } if (failed(eraseDeferredSourceSelectors(funcOp, rewriter)) || failed(eraseOldGraph(funcOp, rewriter)) diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp index 42ae3a1..20551ec 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp @@ -180,13 +180,18 @@ static bool originatesFromDeferredSource( return originatesFromDeferredSource(value, deferred, visited); } -static bool isInsideDeferredLoop(Operation *op, - SpatDeferredCommunicationOp deferred) { +static scf::ForOp getEnclosingDeferredLoop( + Operation *op, SpatDeferredCommunicationOp deferred) { for (Operation *parent = op->getParentOp(); parent && parent != deferred; parent = parent->getParentOp()) - if (isa(parent)) - return true; - return false; + if (auto loop = dyn_cast(parent)) + return loop; + return {}; +} + +static bool isInsideDeferredLoop( + Operation *op, SpatDeferredCommunicationOp deferred) { + return static_cast(getEnclosingDeferredLoop(op, deferred)); } static FailureOr getLoopIterationCount( @@ -297,7 +302,7 @@ static LogicalResult validateDeferredProgram( && llvm::any_of(op->getOperands(), [&](Value operand) { return originatesFromDeferredSource(operand, deferred); })) { - auto loop = op->getParentOfType(); + auto loop = getEnclosingDeferredLoop(op, deferred); scf::ForOp outerLoop; for (Operation *parent = loop ? loop->getParentOp() : nullptr; parent && parent != deferred; parent = parent->getParentOp()) @@ -578,7 +583,7 @@ FailureOr analyzeDeferredProgramTemplate( SmallVector( ArrayRef(slice.getMixedStrides()).drop_front())}; leaf.reconstructedType = cast(value.getType()); - leaf.enclosingLoop = slice->getParentOfType(); + leaf.enclosingLoop = getEnclosingDeferredLoop(slice, deferred); if (graphProjection && slice.getSourceType().getRank() == leaf.reconstructedType.getRank() + 1 @@ -609,8 +614,11 @@ FailureOr analyzeDeferredProgramTemplate( program.leaves.push_back(std::move(leaf)); return success(); } - if (value.getType().isIndex() || isa(value.getType())) - return success(); + if (value.getType().isIndex() || isa(value.getType())) { + Operation *definition = value.getDefiningOp(); + if (!definition || !deferred->isProperAncestor(definition)) + return success(); + } if (auto argument = dyn_cast(value)) { auto loop = dyn_cast_or_null( argument.getOwner()->getParentOp()); @@ -619,7 +627,7 @@ FailureOr analyzeDeferredProgramTemplate( } Operation *op = value.getDefiningOp(); if (!op || (op->getBlock() != &body - && !op->getParentOfType())) + && !getEnclosingDeferredLoop(op, deferred))) return deferred.emitOpError( "deferred residual escapes its verified body: ") << value; if (auto loop = dyn_cast(op)) { diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp index 6cdb525..dbd01a1 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp @@ -289,7 +289,9 @@ static Value cloneResidual( mapping.map(oldValue, newValue); } for (Operation *op : exchange.program.residualOps) { - if (op->hasTrait()) + if (op->hasTrait() + || llvm::all_of(op->getResults(), + [&](Value result) { return mapping.contains(result); })) continue; if (auto oldLoop = dyn_cast(op)) { SmallVector initArgs; diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp index f6e554b..fcb6729 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp @@ -454,6 +454,9 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa OpBuilder builder(blueprint); blueprint->setAttr("fragmentOperandIndices", builder.getDenseI64ArrayAttr(newOperands)); blueprint->setAttr("fragmentSourceSlots", builder.getDenseI64ArrayAttr(newSlots)); + if (blueprint.getIndexMap() == kContiguousRowMajorFragments + && !isCanonicalContiguousRowMajorFragmentAssembly(blueprint)) + blueprint.setIndexMapAttr(builder.getStringAttr("fragment_assembly")); return success(); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp index 66d35ea..0f0a5dd 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp @@ -400,6 +400,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu std::vector processorResidentWeights(processorCount); std::vector schedules(nodeCount); std::vector> tasksByProcessor(processorCount); + std::vector> timelineByProcessor(processorCount); size_t scheduledCount = 0; while (!readyQueue.empty()) { @@ -441,7 +442,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu Time currentEnd = 0; bool foundGap = false; - for (size_t schedTaskIndex : tasksByProcessor[processor]) { + for (size_t schedTaskIndex : timelineByProcessor[processor]) { const ScheduledTask& schedTask = schedules[schedTaskIndex]; Time gapStart = std::max(currentEnd, dataReady); @@ -532,10 +533,13 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu insertResidentWeights(capacityReservations[bestProcessor], graph.nodes[task].residentWeights); insertResidentWeights(processorResidentWeights[bestProcessor], graph.nodes[task].residentWeights); - // 3. CRITICAL FIX: Topological Append - // Because the readyQueue pops in strict topological order, simply pushing to the - // back guarantees the Monoliths will be physically generated cycle-free. - // The hardware will still benefit from the processor assignment chosen by PEFT. + auto& timeline = timelineByProcessor[bestProcessor]; + timeline.insert(llvm::upper_bound(timeline, task, [&](size_t lhs, size_t rhs) { + return schedules[lhs].startTime < schedules[rhs].startTime; + }), task); + + // Materialization requires topological order; gap placement requires the + // separate chronological timeline above. tasksByProcessor[bestProcessor].push_back(task); for (const auto& [child, weight] : graph.successors[task]) { diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp index 301c3b3..1ab3838 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.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/AsmState.h" #include "mlir/IR/BuiltinAttributes.h" @@ -16,6 +17,7 @@ #include #include +#include #include #include #include @@ -54,6 +56,12 @@ struct ChannelSendRecord { std::optional sourceLane; }; +struct ChannelEvaluationContext { + Value laneArg; + uint32_t lane = 0; + DenseMap bindings; +}; + enum class LogicalNodeSelector { Scalar, Lane, @@ -249,10 +257,20 @@ void addBatchNodeRows(std::fstream& nodesFile, } } -std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t lane); +std::optional evaluateIndexLike(Value value, + Value laneArg, + uint32_t lane, + const DenseMap* bindings); -std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t lane) { - if (value == laneArg) +std::optional evaluateIndexLike(Value value, + Value laneArg, + uint32_t lane, + const DenseMap* bindings) { + if (bindings) + if (auto it = bindings->find(value); it != bindings->end()) + return it->second; + + if (laneArg && value == laneArg) return static_cast(lane); if (std::optional constant = matchConstantIndexValue(value)) @@ -270,7 +288,8 @@ std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t la if (!elements || !shapedType || shapedType.getRank() != 1 || extract.getIndices().size() != 1) return std::nullopt; - std::optional index = evaluateIndexLike(extract.getIndices().front(), laneArg, lane); + std::optional index = + evaluateIndexLike(extract.getIndices().front(), laneArg, lane, bindings); if (!index || *index < 0 || *index >= static_cast(elements.getNumElements())) return std::nullopt; @@ -279,11 +298,104 @@ std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t la return std::nullopt; } + if (auto indexCast = value.getDefiningOp()) + return evaluateIndexLike(indexCast.getIn(), laneArg, lane, bindings); + + if (auto add = value.getDefiningOp()) { + auto lhs = evaluateIndexLike(add.getLhs(), laneArg, lane, bindings); + auto rhs = evaluateIndexLike(add.getRhs(), laneArg, lane, bindings); + if (lhs && rhs) + return *lhs + *rhs; + return std::nullopt; + } + + if (auto sub = value.getDefiningOp()) { + auto lhs = evaluateIndexLike(sub.getLhs(), laneArg, lane, bindings); + auto rhs = evaluateIndexLike(sub.getRhs(), laneArg, lane, bindings); + if (lhs && rhs) + return *lhs - *rhs; + return std::nullopt; + } + + if (auto mul = value.getDefiningOp()) { + auto lhs = evaluateIndexLike(mul.getLhs(), laneArg, lane, bindings); + auto rhs = evaluateIndexLike(mul.getRhs(), laneArg, lane, bindings); + if (lhs && rhs) + return *lhs * *rhs; + return std::nullopt; + } + + if (auto div = value.getDefiningOp()) { + auto lhs = evaluateIndexLike(div.getLhs(), laneArg, lane, bindings); + auto rhs = evaluateIndexLike(div.getRhs(), laneArg, lane, bindings); + if (!lhs || !rhs || *rhs == 0 + || (*lhs == std::numeric_limits::min() && *rhs == -1)) + return std::nullopt; + return *lhs / *rhs; + } + + if (auto div = value.getDefiningOp()) { + auto lhs = evaluateIndexLike(div.getLhs(), laneArg, lane, bindings); + auto rhs = evaluateIndexLike(div.getRhs(), laneArg, lane, bindings); + if (!lhs || !rhs || *rhs == 0) + return std::nullopt; + return static_cast(static_cast(*lhs) / static_cast(*rhs)); + } + + if (auto rem = value.getDefiningOp()) { + auto lhs = evaluateIndexLike(rem.getLhs(), laneArg, lane, bindings); + auto rhs = evaluateIndexLike(rem.getRhs(), laneArg, lane, bindings); + if (!lhs || !rhs || *rhs == 0) + return std::nullopt; + if (*lhs == std::numeric_limits::min() && *rhs == -1) + return 0; + return *lhs % *rhs; + } + + if (auto rem = value.getDefiningOp()) { + auto lhs = evaluateIndexLike(rem.getLhs(), laneArg, lane, bindings); + auto rhs = evaluateIndexLike(rem.getRhs(), laneArg, lane, bindings); + if (!lhs || !rhs || *rhs == 0) + return std::nullopt; + return static_cast(static_cast(*lhs) % static_cast(*rhs)); + } + + if (auto cmp = value.getDefiningOp()) { + auto lhs = evaluateIndexLike(cmp.getLhs(), laneArg, lane, bindings); + auto rhs = evaluateIndexLike(cmp.getRhs(), laneArg, lane, bindings); + if (!lhs || !rhs) + return std::nullopt; + bool result = false; + switch (cmp.getPredicate()) { + case arith::CmpIPredicate::eq: result = *lhs == *rhs; break; + case arith::CmpIPredicate::ne: result = *lhs != *rhs; break; + case arith::CmpIPredicate::slt: result = *lhs < *rhs; break; + case arith::CmpIPredicate::sle: result = *lhs <= *rhs; break; + case arith::CmpIPredicate::sgt: result = *lhs > *rhs; break; + case arith::CmpIPredicate::sge: result = *lhs >= *rhs; break; + case arith::CmpIPredicate::ult: result = static_cast(*lhs) < static_cast(*rhs); break; + case arith::CmpIPredicate::ule: result = static_cast(*lhs) <= static_cast(*rhs); break; + case arith::CmpIPredicate::ugt: result = static_cast(*lhs) > static_cast(*rhs); break; + case arith::CmpIPredicate::uge: result = static_cast(*lhs) >= static_cast(*rhs); break; + } + return result ? 1 : 0; + } + + if (auto select = value.getDefiningOp()) { + auto condition = evaluateIndexLike(select.getCondition(), laneArg, lane, bindings); + if (!condition) + return std::nullopt; + return evaluateIndexLike(*condition ? select.getTrueValue() : select.getFalseValue(), + laneArg, + lane, + bindings); + } + if (auto affineApply = value.getDefiningOp()) if (FailureOr folded = evaluateAffineApply(affineApply, [&](Value operand) -> FailureOr { if (std::optional resolved = - evaluateIndexLike(operand, laneArg, lane)) + evaluateIndexLike(operand, laneArg, lane, bindings)) return *resolved; return failure(); }); @@ -294,24 +406,70 @@ std::optional evaluateIndexLike(Value value, Value laneArg, uint32_t la return std::nullopt; } -SmallVector collectPossibleIntValues(Value value, Value laneArg, uint32_t lane) { - if (std::optional exact = evaluateIndexLike(value, laneArg, lane)) - return {*exact}; +bool containsChannelOperation(Operation* root) { + bool found = false; + root->walk([&](Operation* op) { + found |= isa(op); + }); + return found; +} - auto extract = value.getDefiningOp(); - auto constant = extract ? extract.getTensor().getDefiningOp() : nullptr; - auto elements = constant ? dyn_cast(constant.getValue()) : nullptr; - if (!elements) - return {}; +template +LogicalResult walkChannelRegion(Region& region, const ChannelEvaluationContext& context, Emit& emit) { + if (region.empty()) + return success(); - SmallVector values; - if (auto denseInts = dyn_cast(elements)) { - values.reserve(elements.getNumElements()); - for (APInt element : denseInts.getValues()) - if (!llvm::is_contained(values, element.getSExtValue())) - values.push_back(element.getSExtValue()); + for (Operation& op : region.front()) { + if (auto ifOp = dyn_cast(&op)) { + auto condition = evaluateIndexLike( + ifOp.getCondition(), context.laneArg, context.lane, &context.bindings); + if (!condition) + return ifOp.emitOpError("has an unresolved condition in Spatial dataflow export"); + Region& selected = *condition ? ifOp.getThenRegion() : ifOp.getElseRegion(); + if (failed(walkChannelRegion(selected, context, emit))) + return failure(); + continue; + } + + if (auto forOp = dyn_cast(&op)) { + if (!containsChannelOperation(forOp)) + continue; + + auto lower = evaluateIndexLike(forOp.getLowerBound(), context.laneArg, context.lane, &context.bindings); + auto upper = evaluateIndexLike(forOp.getUpperBound(), context.laneArg, context.lane, &context.bindings); + auto step = evaluateIndexLike(forOp.getStep(), context.laneArg, context.lane, &context.bindings); + if (!lower || !upper || !step || *step == 0) + return forOp.emitOpError("has unresolved or invalid bounds in Spatial dataflow export"); + + constexpr uint64_t kMaxExportedLoopIterations = 1 << 20; + uint64_t iterationCount = 0; + int64_t induction = *lower; + while ((*step > 0 && induction < *upper) || (*step < 0 && induction > *upper)) { + if (++iterationCount > kMaxExportedLoopIterations) + return forOp.emitOpError("exceeds the bounded iteration limit in Spatial dataflow export"); + ChannelEvaluationContext iterationContext = context; + iterationContext.bindings[forOp.getInductionVar()] = induction; + if (failed(walkChannelRegion(forOp.getRegion(), iterationContext, emit))) + return failure(); + if ((*step > 0 && induction > std::numeric_limits::max() - *step) + || (*step < 0 && induction < std::numeric_limits::min() - *step)) + return forOp.emitOpError("overflows while enumerating Spatial dataflow export iterations"); + induction += *step; + } + continue; + } + + if (!isa(&op)) + continue; + auto channel = dyn_cast(&op); + Value channelValue = channel ? channel.getChannelId() : cast(&op).getChannelId(); + auto channelId = evaluateIndexLike(channelValue, context.laneArg, context.lane, &context.bindings); + if (!channelId) + return op.emitError("has an unresolved channel identity in Spatial dataflow export"); + if (failed(emit(op, *channelId, context))) + return failure(); } - return values; + return success(); } template @@ -604,50 +762,44 @@ LogicalResult emitDataEdges(std::fstream& edgesFile, } template -void collectChannelSends(DenseMap>& sendsByChannelId, - const DenseMap, ExpandedNodeInfo>& expandedNodes, - BatchOpTy batch) { +LogicalResult collectChannelSends(DenseMap>& sendsByChannelId, + const DenseMap, ExpandedNodeInfo>& expandedNodes, + BatchOpTy batch) { std::optional laneArg = batch.getLaneArgument(); if (!laneArg) - return; + return success(); for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { std::string sourceId = getExpandedNodeId(expandedNodes, batch.getOperation(), lane); if (sourceId.empty()) continue; - batch.getBody().walk([&](SpatChannelSendOp send) { - std::optional channelId = evaluateIndexLike(send.getChannelId(), *laneArg, lane); - if (!channelId) - return; - sendsByChannelId[*channelId].push_back({sourceId, lane}); - }); + ChannelEvaluationContext context; + context.laneArg = *laneArg; + context.lane = lane; + auto emit = [&](Operation& op, int64_t channelId, const ChannelEvaluationContext&) { + if (auto send = dyn_cast(&op)) + sendsByChannelId[channelId].push_back({sourceId, lane}); + return success(); + }; + if (failed(walkChannelRegion(batch.getBody(), context, emit))) + return failure(); } + return success(); } -void collectChannelSends(DenseMap>& sendsByChannelId, - const DenseMap, ExpandedNodeInfo>& expandedNodes, - SpatScheduledCompute compute) { +LogicalResult collectChannelSends(DenseMap>& sendsByChannelId, + const DenseMap, ExpandedNodeInfo>& expandedNodes, + SpatScheduledCompute compute) { std::string sourceId = getExpandedNodeId(expandedNodes, compute.getOperation(), 0); if (sourceId.empty()) - return; - compute.getBody().walk([&](SpatChannelSendOp send) { - std::optional channelId = evaluateIndexLike(send.getChannelId(), Value(), 0); - if (!channelId) - return; - sendsByChannelId[*channelId].push_back({sourceId, std::nullopt}); - }); -} - -DenseMap> -buildNodesByCore(const DenseMap, ExpandedNodeInfo>& expandedNodes) { - DenseMap> nodesByCore; - for (const auto& entry : expandedNodes) { - const ExpandedNodeInfo& node = entry.second; - if (!node.core) - continue; - nodesByCore[*node.core].push_back({node.id, node.lane}); - } - return nodesByCore; + return success(); + ChannelEvaluationContext context; + auto emit = [&](Operation& op, int64_t channelId, const ChannelEvaluationContext&) { + if (isa(&op)) + sendsByChannelId[channelId].push_back({sourceId, std::nullopt}); + return success(); + }; + return walkChannelRegion(compute.getBody(), context, emit); } template @@ -660,14 +812,18 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile, const TopLevelOpInfo& info = entry.second; if (auto compute = dyn_cast(op)) { - compute.getBody().walk([&](SpatChannelReceiveOp receive) { - SmallVector sources = resolveChannelSources(receive, 0); - if (sources.empty()) - return; - std::optional channelId = evaluateIndexLike(receive.getChannelId(), Value(), 0); + ChannelEvaluationContext context; + auto emit = [&](Operation& channelOp, int64_t channelId, const ChannelEvaluationContext&) { + auto receive = dyn_cast(&channelOp); + if (!receive) + return success(); + FailureOr> sources = + resolveChannelSources(receive, channelId, 0); + if (failed(sources)) + return failure(); std::string targetId = getScalarId(info.isScheduled, info.opId); std::optional byteSize = getTypeSizeBytes(receive.getType()); - for (const ChannelSendRecord& source : sources) + for (const ChannelSendRecord& source : *sources) emitEdgeRow(edgesFile, source.sourceId, targetId, @@ -677,7 +833,10 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile, source.sourceLane, std::nullopt, channelId); - }); + return success(); + }; + if (failed(walkChannelRegion(compute.getBody(), context, emit))) + return failure(); continue; } @@ -688,14 +847,20 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile, if (!laneArg) continue; for (uint32_t lane = 0; lane < static_cast(batch.getLaneCount()); ++lane) { - std::string targetId = getBatchLaneId(info.isScheduled, info.opId, lane); - batch.getBody().walk([&](SpatChannelReceiveOp receive) { - SmallVector sources = resolveChannelSources(receive, lane); - if (sources.empty()) - return; - std::optional channelId = evaluateIndexLike(receive.getChannelId(), *laneArg, lane); + ChannelEvaluationContext context; + context.laneArg = *laneArg; + context.lane = lane; + auto emit = [&](Operation& channelOp, int64_t channelId, const ChannelEvaluationContext& eventContext) { + auto receive = dyn_cast(&channelOp); + if (!receive) + return success(); + FailureOr> sources = + resolveChannelSources(receive, channelId, eventContext.lane); + if (failed(sources)) + return failure(); + std::string targetId = getBatchLaneId(info.isScheduled, info.opId, eventContext.lane); std::optional byteSize = getTypeSizeBytes(receive.getType()); - for (const ChannelSendRecord& source : sources) + for (const ChannelSendRecord& source : *sources) emitEdgeRow(edgesFile, source.sourceId, targetId, @@ -703,9 +868,12 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile, receive.getType(), stage, source.sourceLane, - lane, + eventContext.lane, channelId); - }); + return success(); + }; + if (failed(walkChannelRegion(batch.getBody(), context, emit))) + return failure(); } } @@ -810,33 +978,25 @@ LogicalResult exportScheduled(func::FuncOp func, DenseMap> sendsByChannelId; for (const auto& entry : topLevelInfo) { Operation* op = entry.first; + LogicalResult collected = success(); if (auto compute = dyn_cast(op)) - collectChannelSends(sendsByChannelId, expandedNodes, compute); + collected = collectChannelSends(sendsByChannelId, expandedNodes, compute); else if (auto batch = dyn_cast(op)) - collectChannelSends(sendsByChannelId, expandedNodes, batch); + collected = collectChannelSends(sendsByChannelId, expandedNodes, batch); + if (failed(collected)) + return failure(); } - DenseMap> nodesByCore = buildNodesByCore(expandedNodes); - auto resolveChannelSources = [&](SpatChannelReceiveOp receive, uint32_t lane) { + DenseMap consumedSendsByChannelId; + auto resolveChannelSources = [&](SpatChannelReceiveOp receive, int64_t channelId, uint32_t) { SmallVector sources; - - Value laneArg; - if (auto owner = receive->getParentOfType()) - if (auto maybeLaneArg = owner.getLaneArgument()) - laneArg = *maybeLaneArg; - - if (std::optional channelId = evaluateIndexLike(receive.getChannelId(), laneArg, lane)) { - if (auto it = sendsByChannelId.find(*channelId); it != sendsByChannelId.end()) - return it->second; - } - - for (int64_t sourceCore : collectPossibleIntValues(receive.getSourceCoreId(), laneArg, lane)) { - auto it = nodesByCore.find(static_cast(sourceCore)); - if (it == nodesByCore.end()) - continue; - llvm::append_range(sources, it->second); - } - return sources; + auto sends = sendsByChannelId.find(channelId); + size_t& consumed = consumedSendsByChannelId[channelId]; + if (sends == sendsByChannelId.end() || consumed >= sends->second.size()) + return receive.emitOpError("has no matching realized channel send in Spatial dataflow export"), + FailureOr>(failure()); + sources.push_back(sends->second[consumed++]); + return FailureOr>(std::move(sources)); }; return emitExplicitChannelEdges( diff --git a/validation/.gitignore b/validation/.gitignore index 3737e4c..edfd889 100644 --- a/validation/.gitignore +++ b/validation/.gitignore @@ -3,6 +3,7 @@ operations/**/outputs operations/**/raptor operations/**/runner operations/**/simulation +operations/**/*.csv networks/**/inputs networks/**/outputs networks/**/raptor @@ -14,3 +15,5 @@ networks/**/*.png networks/**/*.jpg networks/**/*.csv !networks/pimcomp_models/results.csv +!networks/pimcomp_models/validation_results.csv +!operations/validation_results.csv diff --git a/validation/networks/pimcomp_models/validation_results.csv b/validation/networks/pimcomp_models/validation_results.csv new file mode 100644 index 0000000..4549cd5 --- /dev/null +++ b/validation/networks/pimcomp_models/validation_results.csv @@ -0,0 +1,6 @@ +Operation,Result,Compile,Host mem,Cores mem,Cores,Xbars,Latency,Power,Energy +vgg8-mnist-reconstructed,PASS,1.009 s,1.37 MiB,3.14 MiB,141,761,1.465778 ms,325.627854 mW,477298145.040001 pJ +resnet18-v1-7,PASS,11.548 s,9.89 MiB,40.24 MiB,168,7676,28.099952 ms,312.513408 mW,8781611766.119984 pJ +resnet34-v1-7,PASS,28.495 s,9.90 MiB,48.89 MiB,168,15292,45.781486 ms,326.833870 mW,14962940227.679951 pJ +googlenet-12-latency,PASS,6.573 s,10.74 MiB,22.41 MiB,168,7176,13.371204 ms,457.538139 mW,6117835798.919991 pJ +yolo11n-latency,FAIL,58.572 s,82.55 MiB,185.68 MiB,168,6484,885.264931 ms,189.218985 mW,167508931321.001465 pJ diff --git a/validation/operations/README.md b/validation/operations/README.md index 861a54e..cee977c 100644 --- a/validation/operations/README.md +++ b/validation/operations/README.md @@ -43,7 +43,7 @@ and writes the same rows to `validation_results.csv`. ## Complete inventory -The suite contains 165 models. Tensor shapes, attributes, and constants are +The suite contains 168 models. Tensor shapes, attributes, and constants are defined in `gen_tests.py` and in the checked-in ONNX models. ### Add (5) @@ -64,7 +64,7 @@ defined in `gen_tests.py` and in the checked-in ONNX models. | `negative_axis` | Concatenates tensors using a negative axis. | | `three_inputs_channel_axis` | Concatenates three runtime NCHW tensors along the channel axis. | -### Conv (32) +### Conv (34) | Case | Description | |---|---| @@ -99,6 +99,8 @@ defined in `gen_tests.py` and in the checked-in ONNX models. | `with_bias_3x3` | Multi-channel 3x3 Conv with bias. | | `with_constant` | Hand-authored SAME_UPPER Conv with constant weight and bias. | | `without_kernel_shape_attr` | Conv whose kernel shape is inferred from its weight tensor. | +| `yolo11n_depthwise_head` | YOLO11n pointwise-to-depthwise head boundary at `80x80`, preserving row fragments. | +| `yolo11n_heavy` | Two largest standard YOLO11n Conv-SiLU blocks by MAC count at `64x80x80`. | | `yolo11n_stem` | First two YOLO11n `Conv-SiLU` blocks at `640x640`, including the distributed activation boundary. | ### Div (6) @@ -158,7 +160,7 @@ defined in `gen_tests.py` and in the checked-in ONNX models. | `with_homogeneous_constant` | Adds a constant bias matching the output shape. | | `with_scalar_constant` | Adds a scalar broadcast bias. | -### MatMul (11) +### MatMul (12) | Case | Description | |---|---| @@ -173,6 +175,7 @@ defined in `gen_tests.py` and in the checked-in ONNX models. | `left_constant` | Direct 2D MatMul with constant left-hand matrix. | | `matrix_vector` | Matrix-vector multiplication producing a 1D output. | | `vector_matrix` | Vector-matrix multiplication producing a 1D output. | +| `yolo_attention` | YOLO11n rank-4 dynamic MatMul-scale-transpose-MatMul attention chain. | ### Mul (5) diff --git a/validation/operations/conv/yolo11n_depthwise_head/conv_yolo11n_depthwise_head.onnx b/validation/operations/conv/yolo11n_depthwise_head/conv_yolo11n_depthwise_head.onnx new file mode 100644 index 0000000..5921c1a Binary files /dev/null and b/validation/operations/conv/yolo11n_depthwise_head/conv_yolo11n_depthwise_head.onnx differ diff --git a/validation/operations/conv/yolo11n_heavy/conv_yolo11n_heavy.onnx b/validation/operations/conv/yolo11n_heavy/conv_yolo11n_heavy.onnx new file mode 100644 index 0000000..f3ef7ed Binary files /dev/null and b/validation/operations/conv/yolo11n_heavy/conv_yolo11n_heavy.onnx differ diff --git a/validation/operations/gen_tests.py b/validation/operations/gen_tests.py index 9ca077f..88930b3 100644 --- a/validation/operations/gen_tests.py +++ b/validation/operations/gen_tests.py @@ -242,6 +242,50 @@ def conv_yolo11n_stem(): save_model(model, "conv/yolo11n_stem", "conv_yolo11n_stem.onnx") +def conv_yolo11n_heavy(): + """Two largest YOLO11n standard Conv-SiLU blocks by MAC count.""" + X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 64, 80, 80]) + Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 64, 80, 80]) + rng = np.random.default_rng(111) + W0 = numpy_helper.from_array(rng.uniform(-1, 1, (64, 64, 3, 3)).astype(np.float32), name="W0") + B0 = numpy_helper.from_array(rng.uniform(-1, 1, (64,)).astype(np.float32), name="B0") + W1 = numpy_helper.from_array(rng.uniform(-1, 1, (64, 64, 3, 3)).astype(np.float32), name="W1") + B1 = numpy_helper.from_array(rng.uniform(-1, 1, (64,)).astype(np.float32), name="B1") + nodes = [ + helper.make_node("Conv", ["X", "W0", "B0"], ["C0"], + kernel_shape=[3, 3], strides=[1, 1], pads=[1, 1, 1, 1]), + helper.make_node("Sigmoid", ["C0"], ["S0"]), + helper.make_node("Mul", ["C0", "S0"], ["A0"]), + helper.make_node("Conv", ["A0", "W1", "B1"], ["C1"], + kernel_shape=[3, 3], strides=[1, 1], pads=[1, 1, 1, 1]), + helper.make_node("Sigmoid", ["C1"], ["S1"]), + helper.make_node("Mul", ["C1", "S1"], ["Y"]), + ] + graph = helper.make_graph(nodes, "conv_yolo11n_heavy", [X], [Y], initializer=[W0, B0, W1, B1]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + save_model(model, "conv/yolo11n_heavy", "conv_yolo11n_heavy.onnx") + + +def conv_yolo11n_depthwise_head(): + """YOLO11n pointwise-to-depthwise head boundary at its largest feature map.""" + X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 64, 80, 80]) + Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 64, 80, 80]) + rng = np.random.default_rng(110) + W0 = numpy_helper.from_array(rng.uniform(-1, 1, (64, 64, 1, 1)).astype(np.float32), name="W0") + B0 = numpy_helper.from_array(rng.uniform(-1, 1, (64,)).astype(np.float32), name="B0") + W1 = numpy_helper.from_array(rng.uniform(-1, 1, (64, 1, 3, 3)).astype(np.float32), name="W1") + B1 = numpy_helper.from_array(rng.uniform(-1, 1, (64,)).astype(np.float32), name="B1") + nodes = [ + helper.make_node("Conv", ["X", "W0", "B0"], ["P"], + kernel_shape=[1, 1], strides=[1, 1], pads=[0, 0, 0, 0]), + helper.make_node("Conv", ["P", "W1", "B1"], ["Y"], + kernel_shape=[3, 3], strides=[1, 1], pads=[1, 1, 1, 1], group=64), + ] + graph = helper.make_graph(nodes, "conv_yolo11n_depthwise_head", [X], [Y], initializer=[W0, B0, W1, B1]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + save_model(model, "conv/yolo11n_depthwise_head", "conv_yolo11n_depthwise_head.onnx") + + def conv_pointwise_tiled_chain(): """Chained pointwise Convs with a tiled intermediate.""" X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1024, 1, 1]) @@ -760,6 +804,22 @@ def matmul_batched_3d_dynamic(): save_model(model, "matmul/batched_3d_dynamic", "matmul_batched_3d_dynamic.onnx") +def matmul_yolo_attention(): + """YOLO11n attention chain with rank-4 dynamic matrices.""" + Q = helper.make_tensor_value_info("Q", TensorProto.FLOAT, [1, 2, 400, 32]) + K = helper.make_tensor_value_info("K", TensorProto.FLOAT, [1, 2, 32, 400]) + V = helper.make_tensor_value_info("V", TensorProto.FLOAT, [1, 2, 64, 400]) + Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 2, 64, 400]) + scale = numpy_helper.from_array(np.asarray([0.1767767], dtype=np.float32), name="scale") + nodes = [helper.make_node("MatMul", ["Q", "K"], ["scores"]), + helper.make_node("Mul", ["scores", "scale"], ["scaled"]), + helper.make_node("Transpose", ["scaled"], ["weights"], perm=[0, 1, 3, 2]), + helper.make_node("MatMul", ["V", "weights"], ["Y"])] + graph = helper.make_graph(nodes, "matmul_yolo_attention", [Q, K, V], [Y], initializer=[scale]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + save_model(model, "matmul/yolo_attention", "matmul_yolo_attention.onnx") + + def matmul_batched_left_constant(): """Batched 3D MatMul with constant LHS and runtime RHS.""" rng = np.random.default_rng(70) @@ -2057,6 +2117,8 @@ if __name__ == "__main__": conv_huge_pointwise_1024() conv_huge_pointwise_1024_dynamic() conv_yolo11n_stem() + conv_yolo11n_heavy() + conv_yolo11n_depthwise_head() conv_pointwise_tiled_chain() conv_large_output_channels_1x1() conv_large_input_channels_1x1() @@ -2078,6 +2140,7 @@ if __name__ == "__main__": matmul_dynamic() matmul_batched_3d() matmul_batched_3d_dynamic() + matmul_yolo_attention() matmul_batched_left_constant() matmul_batched_rhs_broadcast() matmul_batched_lhs_broadcast() diff --git a/validation/operations/matmul/yolo_attention/matmul_yolo_attention.onnx b/validation/operations/matmul/yolo_attention/matmul_yolo_attention.onnx new file mode 100644 index 0000000..85b0463 Binary files /dev/null and b/validation/operations/matmul/yolo_attention/matmul_yolo_attention.onnx differ diff --git a/validation/operations/validation_results.csv b/validation/operations/validation_results.csv index e00c427..c14c1c3 100644 --- a/validation/operations/validation_results.csv +++ b/validation/operations/validation_results.csv @@ -1,166 +1,169 @@ Operation,Result,Compile,Host mem,Cores mem,Cores,Xbars,Latency,Power,Energy -add/after_gemm,PASS,0.141 s,0.01 MiB,0.01 MiB,5,4,0.007784 ms,104.703618 mW,815012.960000 pJ -add/basic,PASS,0.114 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -add/broadcast_row,PASS,0.110 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -add/channel_broadcast_1024,PASS,0.105 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ -add/leading_dimension_broadcast,PASS,0.137 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -concat/channel_axis,PASS,0.127 s,0.00 MiB,0.00 MiB,1,0,0.000457 ms,78.157549 mW,35718.000000 pJ -concat/negative_axis,PASS,0.102 s,0.00 MiB,0.00 MiB,1,0,0.001043 ms,78.092042 mW,81450.000000 pJ -concat/three_inputs_channel_axis,PASS,0.127 s,0.00 MiB,0.00 MiB,1,0,0.000644 ms,78.149068 mW,50328.000000 pJ -conv/batch_2,PASS,0.102 s,0.00 MiB,0.00 MiB,2,2,0.013694 ms,82.623885 mW,1131451.480000 pJ -conv/batch_4_pointwise,PASS,0.110 s,0.00 MiB,0.01 MiB,5,4,0.003932 ms,116.078576 mW,456420.960000 pJ -conv/depthwise_1024_channels,PASS,0.167 s,0.19 MiB,0.38 MiB,129,128,0.220751 ms,178.454307 mW,39393966.720000 pJ -conv/depthwise_grouped,PASS,0.109 s,0.01 MiB,0.00 MiB,5,4,0.006024 ms,108.326521 mW,652558.960000 pJ -conv/dilated_3x3,PASS,0.143 s,0.00 MiB,0.00 MiB,3,3,0.004045 ms,110.234541 mW,445898.720000 pJ -conv/dynamic,PASS,0.107 s,0.00 MiB,0.00 MiB,5,0,0.001835 ms,92.281199 mW,169336.000000 pJ -conv/explicit_padding,PASS,0.139 s,0.00 MiB,0.00 MiB,4,4,0.004327 ms,115.794768 mW,501043.960000 pJ -conv/grouped_many_groups,PASS,0.616 s,0.05 MiB,0.09 MiB,65,64,0.181845 ms,142.210104 mW,25860196.360000 pJ -conv/grouped_two_groups,PASS,0.145 s,0.00 MiB,0.00 MiB,3,2,0.005360 ms,101.459418 mW,543822.480000 pJ -conv/huge_pointwise_1024,PASS,0.749 s,0.01 MiB,0.01 MiB,1,64,0.028261 ms,133.488743 mW,3772525.360000 pJ -conv/huge_pointwise_1024_dynamic,PASS,0.097 s,8.04 MiB,12.61 MiB,168,0,2.627964 ms,169.518697 mW,445489032.000000 pJ -conv/kernel_3x3,PASS,0.155 s,0.00 MiB,0.00 MiB,3,3,0.003091 ms,115.862414 mW,358130.720000 pJ -conv/kernel_equals_input_spatial,PASS,0.156 s,0.00 MiB,0.00 MiB,1,2,0.008443 ms,83.863849 mW,708062.480000 pJ -conv/large_input_channels_1x1,PASS,0.160 s,0.01 MiB,0.01 MiB,1,8,0.017167 ms,89.416900 mW,1535019.920000 pJ -conv/large_output_channels_1x1,PASS,0.141 s,0.00 MiB,0.01 MiB,1,8,0.004964 ms,117.628106 mW,583905.920000 pJ -conv/large_spatial,PASS,0.141 s,0.00 MiB,0.01 MiB,6,6,0.004096 ms,129.015000 mW,528445.440000 pJ -conv/multi_channel,PASS,0.143 s,0.00 MiB,0.00 MiB,3,3,0.005148 ms,106.453520 mW,548022.720000 pJ -conv/non_square_kernel_1x3,PASS,0.127 s,0.00 MiB,0.00 MiB,5,5,0.004029 ms,123.600943 mW,497988.200000 pJ -conv/non_square_kernel_3x1,PASS,0.085 s,0.00 MiB,0.00 MiB,3,3,0.005526 ms,105.464843 mW,582798.720000 pJ -conv/non_uniform_stride,PASS,0.120 s,0.00 MiB,0.00 MiB,4,4,0.005808 ms,110.081433 mW,639352.960000 pJ -conv/pointwise_1x1,PASS,0.139 s,0.00 MiB,0.00 MiB,4,4,0.004539 ms,114.835858 mW,521239.960000 pJ -conv/pointwise_tiled_chain,PASS,0.943 s,0.01 MiB,0.02 MiB,2,80,0.084437 ms,102.289307 mW,8637002.200000 pJ -conv/real_asymmetric_padding,PASS,0.115 s,0.00 MiB,0.00 MiB,4,4,0.005232 ms,111.870214 mW,585304.960000 pJ -conv/relu_conv_store,PASS,0.107 s,0.05 MiB,0.08 MiB,32,32,0.062978 ms,246.649827 mW,15533512.800000 pJ -conv/same_lower_3x3,PASS,0.094 s,0.00 MiB,0.00 MiB,5,5,0.004700 ms,119.232170 mW,560391.200000 pJ -conv/same_padding_3x3,PASS,0.126 s,0.00 MiB,0.00 MiB,5,5,0.004700 ms,119.232170 mW,560391.200000 pJ -conv/simple,PASS,0.125 s,0.00 MiB,0.00 MiB,2,2,0.003148 ms,94.665972 mW,298008.480000 pJ -conv/stride_2,PASS,0.146 s,0.00 MiB,0.00 MiB,2,2,0.002827 ms,96.393873 mW,272505.480000 pJ -conv/with_bias_3x3,PASS,0.121 s,0.00 MiB,0.00 MiB,3,3,0.004898 ms,107.176546 mW,524950.720000 pJ -conv/with_constant,PASS,0.106 s,0.00 MiB,0.00 MiB,3,3,0.004273 ms,109.362677 mW,467306.720000 pJ -conv/without_kernel_shape_attr,PASS,0.105 s,0.00 MiB,0.00 MiB,3,3,0.003091 ms,115.862414 mW,358130.720000 pJ -conv/yolo11n_stem,PASS,3.091 s,29.20 MiB,31.38 MiB,168,488,9.799213 ms,361.075203 mW,3538252825.000010 pJ -div/after_gemm,PASS,0.130 s,0.01 MiB,0.01 MiB,5,4,0.007784 ms,104.703618 mW,815012.960000 pJ -div/basic,PASS,0.116 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -div/channel_broadcast_1024,PASS,0.121 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ -div/leading_dimension_broadcast,PASS,0.095 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -div/runtime_scalar_rhs,PASS,0.126 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ -div/scalar_constant,PASS,0.069 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -gather/3d_input_axis1,PASS,0.121 s,0.00 MiB,0.00 MiB,1,0,0.000589 ms,78.081494 mW,45990.000000 pJ -gather/axis0_matrix_indices,PASS,0.088 s,0.00 MiB,0.00 MiB,1,0,0.000697 ms,78.068867 mW,54414.000000 pJ -gather/axis1,PASS,0.123 s,0.00 MiB,0.00 MiB,1,0,0.000801 ms,78.059925 mW,62526.000000 pJ -gather/negative_axis,PASS,0.087 s,0.00 MiB,0.00 MiB,1,0,0.001437 ms,78.033403 mW,112134.000000 pJ -gather/negative_indices,PASS,0.115 s,0.00 MiB,0.00 MiB,1,0,0.000376 ms,78.127660 mW,29376.000000 pJ -gemm/alpha_beta,PASS,0.121 s,0.01 MiB,0.01 MiB,5,4,0.007456 ms,105.272125 mW,784908.960000 pJ -gemm/bias_rank2_broadcast,PASS,0.127 s,0.00 MiB,0.01 MiB,5,4,0.007072 ms,105.979208 mW,749484.960000 pJ -gemm/dynamic,PASS,0.084 s,0.00 MiB,0.00 MiB,5,0,0.002421 ms,91.480793 mW,221475.000000 pJ -gemm/dynamic_alpha,PASS,0.104 s,0.00 MiB,0.00 MiB,5,0,0.003262 ms,91.415696 mW,298198.000000 pJ -gemm/dynamic_beta,PASS,0.110 s,0.00 MiB,0.00 MiB,5,0,0.004365 ms,91.316151 mW,398595.000000 pJ -gemm/dynamic_bias,PASS,0.092 s,0.00 MiB,0.00 MiB,5,0,0.002665 ms,91.445779 mW,243703.000000 pJ -gemm/dynamic_bias_alpha_beta,PASS,0.089 s,0.00 MiB,0.00 MiB,5,0,0.005629 ms,91.279268 mW,513811.000000 pJ -gemm/dynamic_transB,PASS,0.077 s,0.00 MiB,0.00 MiB,5,0,0.001301 ms,91.378171 mW,118883.000000 pJ -gemm/huge_1024,PASS,0.219 s,0.01 MiB,0.10 MiB,73,64,0.017522 ms,215.037402 mW,3767885.360000 pJ -gemm/large,PASS,0.097 s,0.02 MiB,0.03 MiB,17,16,0.011229 ms,140.152181 mW,1573768.840000 pJ -gemm/large_k_small_n,PASS,0.106 s,0.01 MiB,0.01 MiB,9,8,0.004748 ms,133.481449 mW,633769.920000 pJ -gemm/non_square,PASS,0.102 s,0.00 MiB,0.01 MiB,5,4,0.003527 ms,118.958310 mW,419565.960000 pJ -gemm/scalar_bias,PASS,0.111 s,0.00 MiB,0.01 MiB,5,4,0.007072 ms,105.979208 mW,749484.960000 pJ -gemm/simple,PASS,0.135 s,0.03 MiB,0.08 MiB,42,40,0.021640 ms,151.774196 mW,3284393.600000 pJ -gemm/small,PASS,0.074 s,0.00 MiB,0.00 MiB,2,2,0.004420 ms,90.144000 mW,398436.480000 pJ -gemm/small_k_large_n,PASS,0.129 s,0.01 MiB,0.02 MiB,17,8,0.007962 ms,131.005014 mW,1043061.920000 pJ -gemm/transA,PASS,0.089 s,0.00 MiB,0.01 MiB,5,4,0.005762 ms,109.140743 mW,628868.960000 pJ -gemm/transA_transB,PASS,0.115 s,0.00 MiB,0.01 MiB,5,4,0.005762 ms,109.140743 mW,628868.960000 pJ -gemm/transB,PASS,0.068 s,0.00 MiB,0.01 MiB,5,4,0.003527 ms,118.958310 mW,419565.960000 pJ -gemm/transB_with_bias,PASS,0.069 s,0.01 MiB,0.01 MiB,5,4,0.005046 ms,110.546762 mW,557818.960000 pJ -gemm/with_bias,PASS,0.113 s,0.01 MiB,0.01 MiB,5,4,0.005562 ms,108.767882 mW,604966.960000 pJ -gemv/constant,PASS,0.112 s,0.00 MiB,0.00 MiB,0,0,0.000000 ms,2.000000 mW,0.000000 pJ -gemv/simple,PASS,0.136 s,0.00 MiB,0.01 MiB,6,4,0.005160 ms,111.150380 mW,573535.960000 pJ -gemv/with_heterogeneous_constant,PASS,0.138 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ -gemv/with_homogeneous_constant,PASS,0.140 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ -gemv/with_scalar_constant,PASS,0.124 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ -matmul/basic,PASS,0.081 s,0.00 MiB,0.00 MiB,2,2,0.004420 ms,90.144000 mW,398436.480000 pJ -matmul/batched_3d,PASS,0.133 s,0.00 MiB,0.01 MiB,5,4,0.005958 ms,108.588949 mW,646972.960000 pJ -matmul/batched_3d_dynamic,PASS,0.105 s,0.00 MiB,0.00 MiB,9,0,0.003525 ms,92.330213 mW,325464.000000 pJ -matmul/batched_left_constant,PASS,0.136 s,0.00 MiB,0.02 MiB,9,8,0.008822 ms,114.385164 mW,1009105.920000 pJ -matmul/batched_lhs_broadcast,PASS,0.133 s,0.00 MiB,0.01 MiB,5,4,0.005681 ms,109.389361 mW,621440.960000 pJ -matmul/batched_rhs_broadcast,PASS,0.134 s,0.00 MiB,0.01 MiB,5,4,0.005958 ms,108.588949 mW,646972.960000 pJ -matmul/dynamic,PASS,0.093 s,0.00 MiB,0.00 MiB,5,0,0.001621 ms,91.421962 mW,148195.000000 pJ -matmul/huge_1024,PASS,0.277 s,0.01 MiB,0.10 MiB,73,64,0.017522 ms,215.037402 mW,3767885.360000 pJ -matmul/left_constant,PASS,0.123 s,0.00 MiB,0.01 MiB,5,4,0.005853 ms,108.861944 mW,637168.960000 pJ -matmul/matrix_vector,PASS,0.150 s,0.52 MiB,0.78 MiB,168,173,0.384660 ms,202.131271 mW,77751814.880000 pJ -matmul/vector_matrix,PASS,0.186 s,0.01 MiB,0.01 MiB,9,8,0.007409 ms,118.680243 mW,879301.920000 pJ -mul/after_conv,PASS,0.127 s,0.00 MiB,0.00 MiB,4,3,0.005453 ms,107.639046 mW,586955.720000 pJ -mul/basic,PASS,0.070 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -mul/channel_broadcast_1024,PASS,0.065 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ -mul/leading_dimension_broadcast,PASS,0.108 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -mul/scalar_constant,PASS,0.082 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -pool/avg_basic,PASS,0.070 s,0.00 MiB,0.00 MiB,1,0,0.011939 ms,78.022112 mW,931506.000000 pJ -pool/avg_ceil_mode,PASS,0.107 s,0.00 MiB,0.00 MiB,1,0,0.004359 ms,78.033035 mW,340146.000000 pJ -pool/avg_explicit_padding,PASS,0.111 s,0.00 MiB,0.00 MiB,1,0,0.008822 ms,78.027205 mW,688356.000000 pJ -pool/avg_include_pad,PASS,0.075 s,0.00 MiB,0.00 MiB,1,0,0.008506 ms,78.016929 mW,663612.000000 pJ -pool/avg_large_channels,PASS,0.064 s,0.04 MiB,0.02 MiB,1,0,0.178249 ms,78.280327 mW,13953390.000000 pJ -pool/avg_non_uniform_stride,PASS,0.068 s,0.00 MiB,0.00 MiB,1,0,0.014513 ms,78.016537 mW,1132254.000000 pJ -pool/avg_real_asymmetric_padding,PASS,0.064 s,0.00 MiB,0.00 MiB,1,0,0.025206 ms,78.024756 mW,1966692.000000 pJ -pool/max_after_conv,PASS,0.070 s,0.00 MiB,0.00 MiB,6,4,0.006452 ms,96.374606 mW,621808.960000 pJ +add/after_gemm,PASS,0.072 s,0.01 MiB,0.01 MiB,5,4,0.007784 ms,104.703618 mW,815012.960000 pJ +add/basic,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +add/broadcast_row,PASS,0.048 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +add/channel_broadcast_1024,PASS,0.051 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ +add/leading_dimension_broadcast,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +concat/channel_axis,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,0.000457 ms,78.157549 mW,35718.000000 pJ +concat/negative_axis,PASS,0.061 s,0.00 MiB,0.00 MiB,1,0,0.001043 ms,78.092042 mW,81450.000000 pJ +concat/three_inputs_channel_axis,PASS,0.054 s,0.00 MiB,0.00 MiB,1,0,0.000644 ms,78.149068 mW,50328.000000 pJ +conv/batch_2,PASS,0.059 s,0.00 MiB,0.00 MiB,2,2,0.013694 ms,82.623885 mW,1131451.480000 pJ +conv/batch_4_pointwise,PASS,0.056 s,0.00 MiB,0.01 MiB,5,4,0.003932 ms,116.078576 mW,456420.960000 pJ +conv/depthwise_1024_channels,PASS,0.081 s,0.19 MiB,0.38 MiB,129,128,0.220751 ms,178.454307 mW,39393966.720000 pJ +conv/depthwise_grouped,PASS,0.059 s,0.01 MiB,0.00 MiB,5,4,0.006024 ms,108.326521 mW,652558.960000 pJ +conv/dilated_3x3,PASS,0.059 s,0.00 MiB,0.00 MiB,3,3,0.004045 ms,110.234541 mW,445898.720000 pJ +conv/dynamic,PASS,0.059 s,0.00 MiB,0.00 MiB,5,0,0.001835 ms,92.281199 mW,169336.000000 pJ +conv/explicit_padding,PASS,0.058 s,0.00 MiB,0.00 MiB,4,4,0.004327 ms,115.794768 mW,501043.960000 pJ +conv/grouped_many_groups,PASS,0.505 s,0.05 MiB,0.09 MiB,65,64,0.181845 ms,142.210104 mW,25860196.360000 pJ +conv/grouped_two_groups,PASS,0.061 s,0.00 MiB,0.00 MiB,3,2,0.005360 ms,101.459418 mW,543822.480000 pJ +conv/huge_pointwise_1024,PASS,0.643 s,0.01 MiB,0.01 MiB,1,64,0.028261 ms,133.488743 mW,3772525.360000 pJ +conv/huge_pointwise_1024_dynamic,PASS,0.081 s,8.04 MiB,12.61 MiB,168,0,2.627964 ms,169.518697 mW,445489032.000000 pJ +conv/kernel_3x3,PASS,0.056 s,0.00 MiB,0.00 MiB,3,3,0.003091 ms,115.862414 mW,358130.720000 pJ +conv/kernel_equals_input_spatial,PASS,0.060 s,0.00 MiB,0.00 MiB,1,2,0.008443 ms,83.863849 mW,708062.480000 pJ +conv/large_input_channels_1x1,PASS,0.088 s,0.01 MiB,0.01 MiB,1,8,0.017167 ms,89.416900 mW,1535019.920000 pJ +conv/large_output_channels_1x1,PASS,0.087 s,0.00 MiB,0.01 MiB,1,8,0.004964 ms,117.628106 mW,583905.920000 pJ +conv/large_spatial,PASS,0.055 s,0.00 MiB,0.01 MiB,6,6,0.004096 ms,129.015000 mW,528445.440000 pJ +conv/multi_channel,PASS,0.056 s,0.00 MiB,0.00 MiB,3,3,0.005148 ms,106.453520 mW,548022.720000 pJ +conv/non_square_kernel_1x3,PASS,0.056 s,0.00 MiB,0.00 MiB,5,5,0.004029 ms,123.600943 mW,497988.200000 pJ +conv/non_square_kernel_3x1,PASS,0.058 s,0.00 MiB,0.00 MiB,3,3,0.005526 ms,105.464843 mW,582798.720000 pJ +conv/non_uniform_stride,PASS,0.062 s,0.00 MiB,0.00 MiB,4,4,0.005808 ms,110.081433 mW,639352.960000 pJ +conv/pointwise_1x1,PASS,0.057 s,0.00 MiB,0.00 MiB,4,4,0.004539 ms,114.835858 mW,521239.960000 pJ +conv/pointwise_tiled_chain,PASS,0.771 s,0.01 MiB,0.02 MiB,2,80,0.084437 ms,102.289307 mW,8637002.200000 pJ +conv/real_asymmetric_padding,PASS,0.055 s,0.00 MiB,0.00 MiB,4,4,0.005232 ms,111.870214 mW,585304.960000 pJ +conv/relu_conv_store,PASS,0.076 s,0.02 MiB,0.10 MiB,32,32,0.064390 ms,243.916397 mW,15705776.800000 pJ +conv/same_lower_3x3,PASS,0.057 s,0.00 MiB,0.00 MiB,5,5,0.004700 ms,119.232170 mW,560391.200000 pJ +conv/same_padding_3x3,PASS,0.060 s,0.00 MiB,0.00 MiB,5,5,0.004700 ms,119.232170 mW,560391.200000 pJ +conv/simple,PASS,0.055 s,0.00 MiB,0.00 MiB,2,2,0.003148 ms,94.665972 mW,298008.480000 pJ +conv/stride_2,PASS,0.053 s,0.00 MiB,0.00 MiB,2,2,0.002827 ms,96.393873 mW,272505.480000 pJ +conv/with_bias_3x3,PASS,0.055 s,0.00 MiB,0.00 MiB,3,3,0.004898 ms,107.176546 mW,524950.720000 pJ +conv/with_constant,PASS,0.073 s,0.00 MiB,0.00 MiB,3,3,0.004273 ms,109.362677 mW,467306.720000 pJ +conv/without_kernel_shape_attr,PASS,0.057 s,0.00 MiB,0.00 MiB,3,3,0.003091 ms,115.862414 mW,358130.720000 pJ +conv/yolo11n_depthwise_head,PASS,0.509 s,4.82 MiB,15.92 MiB,160,720,4.105780 ms,492.826315 mW,2023436428.000010 pJ +conv/yolo11n_heavy,PASS,0.488 s,4.82 MiB,20.66 MiB,160,800,6.275385 ms,418.028199 mW,2623287892.000010 pJ +conv/yolo11n_stem,PASS,1.730 s,12.86 MiB,31.38 MiB,168,488,9.799204 ms,361.075380 mW,3538251304.000010 pJ +div/after_gemm,PASS,0.059 s,0.01 MiB,0.01 MiB,5,4,0.007784 ms,104.703618 mW,815012.960000 pJ +div/basic,PASS,0.048 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +div/channel_broadcast_1024,PASS,0.051 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ +div/leading_dimension_broadcast,PASS,0.063 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +div/runtime_scalar_rhs,PASS,0.057 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ +div/scalar_constant,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +gather/3d_input_axis1,PASS,0.050 s,0.00 MiB,0.00 MiB,1,0,0.000589 ms,78.081494 mW,45990.000000 pJ +gather/axis0_matrix_indices,PASS,0.050 s,0.00 MiB,0.00 MiB,1,0,0.000697 ms,78.068867 mW,54414.000000 pJ +gather/axis1,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.000801 ms,78.059925 mW,62526.000000 pJ +gather/negative_axis,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.001437 ms,78.033403 mW,112134.000000 pJ +gather/negative_indices,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,0.000376 ms,78.127660 mW,29376.000000 pJ +gemm/alpha_beta,PASS,0.056 s,0.01 MiB,0.01 MiB,5,4,0.007456 ms,105.272125 mW,784908.960000 pJ +gemm/bias_rank2_broadcast,PASS,0.056 s,0.00 MiB,0.01 MiB,5,4,0.007072 ms,105.979208 mW,749484.960000 pJ +gemm/dynamic,PASS,0.054 s,0.00 MiB,0.00 MiB,5,0,0.002421 ms,91.480793 mW,221475.000000 pJ +gemm/dynamic_alpha,PASS,0.056 s,0.00 MiB,0.00 MiB,5,0,0.003262 ms,91.415696 mW,298198.000000 pJ +gemm/dynamic_beta,PASS,0.059 s,0.00 MiB,0.00 MiB,5,0,0.004365 ms,91.316151 mW,398595.000000 pJ +gemm/dynamic_bias,PASS,0.056 s,0.00 MiB,0.00 MiB,5,0,0.002665 ms,91.445779 mW,243703.000000 pJ +gemm/dynamic_bias_alpha_beta,PASS,0.053 s,0.00 MiB,0.00 MiB,5,0,0.005629 ms,91.279268 mW,513811.000000 pJ +gemm/dynamic_transB,PASS,0.054 s,0.00 MiB,0.00 MiB,5,0,0.001301 ms,91.378171 mW,118883.000000 pJ +gemm/huge_1024,PASS,0.157 s,0.01 MiB,0.10 MiB,73,64,0.017522 ms,215.037402 mW,3767885.360000 pJ +gemm/large,PASS,0.067 s,0.02 MiB,0.03 MiB,17,16,0.011229 ms,140.152181 mW,1573768.840000 pJ +gemm/large_k_small_n,PASS,0.097 s,0.01 MiB,0.01 MiB,9,8,0.004748 ms,133.481449 mW,633769.920000 pJ +gemm/non_square,PASS,0.060 s,0.00 MiB,0.01 MiB,5,4,0.003527 ms,118.958310 mW,419565.960000 pJ +gemm/scalar_bias,PASS,0.054 s,0.00 MiB,0.01 MiB,5,4,0.007072 ms,105.979208 mW,749484.960000 pJ +gemm/simple,PASS,0.071 s,0.03 MiB,0.08 MiB,42,40,0.021640 ms,151.774196 mW,3284393.600000 pJ +gemm/small,PASS,0.052 s,0.00 MiB,0.00 MiB,2,2,0.004420 ms,90.144000 mW,398436.480000 pJ +gemm/small_k_large_n,PASS,0.089 s,0.01 MiB,0.02 MiB,17,8,0.007962 ms,131.005014 mW,1043061.920000 pJ +gemm/transA,PASS,0.057 s,0.00 MiB,0.01 MiB,5,4,0.005762 ms,109.140743 mW,628868.960000 pJ +gemm/transA_transB,PASS,0.053 s,0.00 MiB,0.01 MiB,5,4,0.005762 ms,109.140743 mW,628868.960000 pJ +gemm/transB,PASS,0.062 s,0.00 MiB,0.01 MiB,5,4,0.003527 ms,118.958310 mW,419565.960000 pJ +gemm/transB_with_bias,PASS,0.057 s,0.01 MiB,0.01 MiB,5,4,0.005046 ms,110.546762 mW,557818.960000 pJ +gemm/with_bias,PASS,0.056 s,0.01 MiB,0.01 MiB,5,4,0.005562 ms,108.767882 mW,604966.960000 pJ +gemv/constant,PASS,0.051 s,0.00 MiB,0.00 MiB,0,0,0.000000 ms,2.000000 mW,0.000000 pJ +gemv/simple,PASS,0.066 s,0.00 MiB,0.01 MiB,6,4,0.005160 ms,111.150380 mW,573535.960000 pJ +gemv/with_heterogeneous_constant,PASS,0.066 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ +gemv/with_homogeneous_constant,PASS,0.063 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ +gemv/with_scalar_constant,PASS,0.064 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ +matmul/basic,PASS,0.064 s,0.00 MiB,0.00 MiB,2,2,0.004420 ms,90.144000 mW,398436.480000 pJ +matmul/batched_3d,PASS,0.058 s,0.00 MiB,0.01 MiB,5,4,0.005958 ms,108.588949 mW,646972.960000 pJ +matmul/batched_3d_dynamic,PASS,0.053 s,0.00 MiB,0.00 MiB,4,0,0.001822 ms,92.192645 mW,167975.000000 pJ +matmul/batched_left_constant,PASS,0.058 s,0.00 MiB,0.02 MiB,9,8,0.008822 ms,114.385164 mW,1009105.920000 pJ +matmul/batched_lhs_broadcast,PASS,0.056 s,0.00 MiB,0.01 MiB,5,4,0.005681 ms,109.389361 mW,621440.960000 pJ +matmul/batched_rhs_broadcast,PASS,0.059 s,0.00 MiB,0.01 MiB,5,4,0.005958 ms,108.588949 mW,646972.960000 pJ +matmul/dynamic,PASS,0.054 s,0.00 MiB,0.00 MiB,5,0,0.001621 ms,91.421962 mW,148195.000000 pJ +matmul/huge_1024,PASS,0.142 s,0.01 MiB,0.10 MiB,73,64,0.017522 ms,215.037402 mW,3767885.360000 pJ +matmul/left_constant,PASS,0.064 s,0.00 MiB,0.01 MiB,5,4,0.005853 ms,108.861944 mW,637168.960000 pJ +matmul/matrix_vector,PASS,0.097 s,0.52 MiB,0.78 MiB,168,173,0.384660 ms,202.131271 mW,77751814.880000 pJ +matmul/vector_matrix,PASS,0.086 s,0.01 MiB,0.01 MiB,9,8,0.007409 ms,118.680243 mW,879301.920000 pJ +matmul/yolo_attention,PASS,0.417 s,1.02 MiB,43.44 MiB,168,0,8.151445 ms,170.003707 mW,1385775865.000000 pJ +mul/after_conv,PASS,0.060 s,0.00 MiB,0.00 MiB,4,3,0.005453 ms,107.639046 mW,586955.720000 pJ +mul/basic,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +mul/channel_broadcast_1024,PASS,0.060 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ +mul/leading_dimension_broadcast,PASS,0.058 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +mul/scalar_constant,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +pool/avg_basic,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,0.011939 ms,78.022112 mW,931506.000000 pJ +pool/avg_ceil_mode,PASS,0.054 s,0.00 MiB,0.00 MiB,1,0,0.004359 ms,78.033035 mW,340146.000000 pJ +pool/avg_explicit_padding,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,0.008822 ms,78.027205 mW,688356.000000 pJ +pool/avg_include_pad,PASS,0.055 s,0.00 MiB,0.00 MiB,1,0,0.008506 ms,78.016929 mW,663612.000000 pJ +pool/avg_large_channels,PASS,0.067 s,0.04 MiB,0.02 MiB,1,0,0.178249 ms,78.280327 mW,13953390.000000 pJ +pool/avg_non_uniform_stride,PASS,0.060 s,0.00 MiB,0.00 MiB,1,0,0.014513 ms,78.016537 mW,1132254.000000 pJ +pool/avg_real_asymmetric_padding,PASS,0.071 s,0.00 MiB,0.00 MiB,1,0,0.025206 ms,78.024756 mW,1966692.000000 pJ +pool/max_after_conv,PASS,0.066 s,0.00 MiB,0.00 MiB,6,4,0.006452 ms,96.374606 mW,621808.960000 pJ pool/max_basic,PASS,0.063 s,0.00 MiB,0.00 MiB,3,0,0.001634 ms,92.132191 mW,150544.000000 pJ -pool/max_ceil_mode,PASS,0.065 s,0.00 MiB,0.00 MiB,2,0,0.001297 ms,79.111025 mW,102607.000000 pJ -pool/max_global_style_kernel_equals_input,PASS,0.089 s,0.00 MiB,0.00 MiB,1,0,0.004366 ms,78.010994 mW,340596.000000 pJ -pool/max_non_square_kernel,PASS,0.103 s,0.00 MiB,0.00 MiB,4,0,0.003409 ms,93.253447 mW,317901.000000 pJ -pool/max_real_asymmetric_padding,PASS,0.088 s,0.00 MiB,0.00 MiB,4,0,0.003078 ms,93.124756 mW,286638.000000 pJ -pool/max_same_upper,PASS,0.065 s,0.00 MiB,0.00 MiB,3,0,0.003024 ms,92.095238 mW,278496.000000 pJ -pool/max_stride2_multichannel,PASS,0.086 s,0.00 MiB,0.00 MiB,3,0,0.004012 ms,92.269192 mW,370184.000000 pJ -reduce_mean/4d_spatial,PASS,0.061 s,0.00 MiB,0.00 MiB,3,0,0.000321 ms,92.448598 mW,29676.000000 pJ -reduce_mean/4d_spatial_keepdims_0,PASS,0.067 s,0.00 MiB,0.00 MiB,4,0,0.000655 ms,94.352672 mW,61801.000000 pJ -reduce_mean/after_conv,PASS,0.115 s,0.00 MiB,0.00 MiB,5,3,0.005342 ms,106.951089 mW,571332.720000 pJ -reduce_mean/all_axes_keepdims_0,PASS,0.103 s,0.00 MiB,0.00 MiB,2,0,0.000391 ms,79.237852 mW,30982.000000 pJ -reduce_mean/all_axes_keepdims_1,PASS,0.086 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ -reduce_mean/basic,PASS,0.099 s,0.00 MiB,0.00 MiB,4,0,0.000373 ms,93.514745 mW,34881.000000 pJ -reduce_mean/channel_axis_nchw,PASS,0.078 s,0.03 MiB,0.02 MiB,4,0,0.164926 ms,93.596631 mW,15436518.000000 pJ -reduce_mean/keepdims_0,PASS,0.073 s,0.00 MiB,0.00 MiB,5,0,0.000748 ms,91.401070 mW,68368.000000 pJ -reduce_mean/large_dimension_1024,PASS,0.115 s,0.01 MiB,0.00 MiB,1,0,0.002785 ms,78.017235 mW,217278.000000 pJ -reduce_mean/legacy_axes_1_2_keepdims_1,PASS,0.074 s,0.00 MiB,0.00 MiB,2,0,0.000271 ms,79.354244 mW,21505.000000 pJ -reduce_mean/legacy_axis1_keepdims_0,PASS,0.111 s,0.00 MiB,0.00 MiB,9,0,0.001986 ms,92.501511 mW,183708.000000 pJ -reduce_mean/legacy_axis1_keepdims_1,PASS,0.086 s,0.00 MiB,0.00 MiB,8,0,0.001373 ms,94.559359 mW,129830.000000 pJ -reduce_mean/legacy_empty_axes_noop,PASS,0.115 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ -reduce_mean/legacy_nchw_spatial,PASS,0.090 s,0.00 MiB,0.00 MiB,3,0,0.000321 ms,92.448598 mW,29676.000000 pJ -reduce_mean/legacy_negative_axis,PASS,0.110 s,0.00 MiB,0.00 MiB,6,0,0.000553 ms,93.520796 mW,51717.000000 pJ -reduce_mean/legacy_reduce_all_keepdims_1,PASS,0.107 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ -reduce_mean/negative_axis,PASS,0.120 s,0.00 MiB,0.00 MiB,6,0,0.000553 ms,93.520796 mW,51717.000000 pJ -relu/4d,PASS,0.116 s,0.00 MiB,0.00 MiB,1,0,0.000521 ms,78.184261 mW,40734.000000 pJ -relu/after_conv,PASS,0.078 s,0.00 MiB,0.00 MiB,3,3,0.004956 ms,106.998935 mW,530286.720000 pJ -relu/after_gemm,PASS,0.077 s,0.01 MiB,0.01 MiB,5,4,0.007513 ms,105.158653 mW,790056.960000 pJ -relu/basic,PASS,0.066 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ -reshape/4d_to_2d_flatten,PASS,0.098 s,0.00 MiB,0.00 MiB,1,0,0.000258 ms,78.279070 mW,20196.000000 pJ -reshape/infer_dim_minus_one,PASS,0.105 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ -reshape/same_rank,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ -reshape/zero_copies_input_dim,PASS,0.090 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ -resize/height_only,PASS,0.104 s,0.00 MiB,0.00 MiB,1,0,0.000795 ms,78.060377 mW,62058.000000 pJ -resize/nearest_2x,PASS,0.104 s,0.00 MiB,0.00 MiB,1,0,0.001422 ms,78.033755 mW,110964.000000 pJ -resize/nearest_downsample,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,0.000481 ms,78.099792 mW,37566.000000 pJ -resize/non_uniform,PASS,0.063 s,0.00 MiB,0.00 MiB,1,0,0.002108 ms,78.034156 mW,164496.000000 pJ -resize/width_only,PASS,0.068 s,0.00 MiB,0.00 MiB,1,0,0.000792 ms,78.060606 mW,61824.000000 pJ -resize/with_sizes,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,0.000951 ms,78.050473 mW,74226.000000 pJ -sigmoid/4d,PASS,0.104 s,0.00 MiB,0.00 MiB,1,0,0.000521 ms,78.184261 mW,40734.000000 pJ -sigmoid/after_gemm,PASS,0.118 s,0.01 MiB,0.01 MiB,5,4,0.007513 ms,105.158653 mW,790056.960000 pJ -sigmoid/basic,PASS,0.083 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ -slice/2d_basic,PASS,0.067 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ -slice/after_conv,PASS,0.111 s,0.00 MiB,0.01 MiB,7,6,0.011296 ms,118.190765 mW,1335082.880000 pJ -slice/default_axes,PASS,0.065 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ -slice/large_channel_1024,PASS,0.060 s,0.01 MiB,0.00 MiB,1,0,0.002832 ms,78.144068 mW,221304.000000 pJ -slice/nchw_spatial_crop,PASS,0.067 s,0.00 MiB,0.00 MiB,1,0,0.001302 ms,78.239631 mW,101868.000000 pJ -slice/negative_axis,PASS,0.098 s,0.00 MiB,0.00 MiB,1,0,0.000562 ms,78.298932 mW,44004.000000 pJ -slice/negative_indices,PASS,0.106 s,0.00 MiB,0.00 MiB,1,0,0.000322 ms,78.298137 mW,25212.000000 pJ -slice/step2,PASS,0.068 s,0.00 MiB,0.00 MiB,1,0,0.002042 ms,78.293830 mW,159876.000000 pJ -softmax/3d_last_axis,PASS,0.072 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED -softmax/basic,PASS,0.088 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED -softmax/channel_axis,PASS,0.110 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED -softmax/large_dimension_1024,PASS,0.104 s,0.01 MiB,0.01 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED -softmax/negative_axis,PASS,0.118 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED -split/basic,PASS,0.077 s,0.00 MiB,0.00 MiB,1,0,0.000403 ms,78.297767 mW,31554.000000 pJ -split/equal_three_way,PASS,0.107 s,0.00 MiB,0.00 MiB,1,0,0.000564 ms,78.297872 mW,44160.000000 pJ -split/negative_axis,PASS,0.111 s,0.00 MiB,0.00 MiB,1,0,0.001083 ms,78.288089 mW,84786.000000 pJ -split/uneven_channel_axis_4d,PASS,0.061 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ +pool/max_ceil_mode,PASS,0.078 s,0.00 MiB,0.00 MiB,2,0,0.001297 ms,79.111025 mW,102607.000000 pJ +pool/max_global_style_kernel_equals_input,PASS,0.067 s,0.00 MiB,0.00 MiB,1,0,0.004366 ms,78.010994 mW,340596.000000 pJ +pool/max_non_square_kernel,PASS,0.073 s,0.00 MiB,0.00 MiB,4,0,0.003409 ms,93.253447 mW,317901.000000 pJ +pool/max_real_asymmetric_padding,PASS,0.064 s,0.00 MiB,0.00 MiB,4,0,0.003078 ms,93.124756 mW,286638.000000 pJ +pool/max_same_upper,PASS,0.063 s,0.00 MiB,0.00 MiB,3,0,0.003024 ms,92.095238 mW,278496.000000 pJ +pool/max_stride2_multichannel,PASS,0.066 s,0.00 MiB,0.00 MiB,3,0,0.004012 ms,92.269192 mW,370184.000000 pJ +reduce_mean/4d_spatial,PASS,0.078 s,0.00 MiB,0.00 MiB,3,0,0.000321 ms,92.448598 mW,29676.000000 pJ +reduce_mean/4d_spatial_keepdims_0,PASS,0.066 s,0.00 MiB,0.00 MiB,4,0,0.000655 ms,94.352672 mW,61801.000000 pJ +reduce_mean/after_conv,PASS,0.068 s,0.00 MiB,0.00 MiB,5,3,0.005342 ms,106.951089 mW,571332.720000 pJ +reduce_mean/all_axes_keepdims_0,PASS,0.067 s,0.00 MiB,0.00 MiB,2,0,0.000391 ms,79.237852 mW,30982.000000 pJ +reduce_mean/all_axes_keepdims_1,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ +reduce_mean/basic,PASS,0.052 s,0.00 MiB,0.00 MiB,4,0,0.000373 ms,93.514745 mW,34881.000000 pJ +reduce_mean/channel_axis_nchw,PASS,0.054 s,0.03 MiB,0.02 MiB,4,0,0.164926 ms,93.596631 mW,15436518.000000 pJ +reduce_mean/keepdims_0,PASS,0.054 s,0.00 MiB,0.00 MiB,5,0,0.000748 ms,91.401070 mW,68368.000000 pJ +reduce_mean/large_dimension_1024,PASS,0.056 s,0.01 MiB,0.00 MiB,1,0,0.002785 ms,78.017235 mW,217278.000000 pJ +reduce_mean/legacy_axes_1_2_keepdims_1,PASS,0.064 s,0.00 MiB,0.00 MiB,2,0,0.000271 ms,79.354244 mW,21505.000000 pJ +reduce_mean/legacy_axis1_keepdims_0,PASS,0.071 s,0.00 MiB,0.00 MiB,9,0,0.001986 ms,92.501511 mW,183708.000000 pJ +reduce_mean/legacy_axis1_keepdims_1,PASS,0.065 s,0.00 MiB,0.00 MiB,8,0,0.001373 ms,94.559359 mW,129830.000000 pJ +reduce_mean/legacy_empty_axes_noop,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ +reduce_mean/legacy_nchw_spatial,PASS,0.050 s,0.00 MiB,0.00 MiB,3,0,0.000321 ms,92.448598 mW,29676.000000 pJ +reduce_mean/legacy_negative_axis,PASS,0.064 s,0.00 MiB,0.00 MiB,6,0,0.000553 ms,93.520796 mW,51717.000000 pJ +reduce_mean/legacy_reduce_all_keepdims_1,PASS,0.054 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ +reduce_mean/negative_axis,PASS,0.059 s,0.00 MiB,0.00 MiB,6,0,0.000553 ms,93.520796 mW,51717.000000 pJ +relu/4d,PASS,0.068 s,0.00 MiB,0.00 MiB,1,0,0.000521 ms,78.184261 mW,40734.000000 pJ +relu/after_conv,PASS,0.057 s,0.00 MiB,0.00 MiB,3,3,0.004956 ms,106.998935 mW,530286.720000 pJ +relu/after_gemm,PASS,0.064 s,0.01 MiB,0.01 MiB,5,4,0.007513 ms,105.158653 mW,790056.960000 pJ +relu/basic,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ +reshape/4d_to_2d_flatten,PASS,0.068 s,0.00 MiB,0.00 MiB,1,0,0.000258 ms,78.279070 mW,20196.000000 pJ +reshape/infer_dim_minus_one,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ +reshape/same_rank,PASS,0.065 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ +reshape/zero_copies_input_dim,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ +resize/height_only,PASS,0.063 s,0.00 MiB,0.00 MiB,4,0,0.000693 ms,93.554113 mW,64833.000000 pJ +resize/nearest_2x,PASS,0.064 s,0.00 MiB,0.00 MiB,4,0,0.001173 ms,93.572890 mW,109761.000000 pJ +resize/nearest_downsample,PASS,0.062 s,0.00 MiB,0.00 MiB,2,0,0.000427 ms,79.449649 mW,33925.000000 pJ +resize/non_uniform,PASS,0.068 s,0.00 MiB,0.00 MiB,6,0,0.001753 ms,93.575014 mW,164037.000000 pJ +resize/width_only,PASS,0.070 s,0.00 MiB,0.00 MiB,2,0,0.000667 ms,79.503748 mW,53029.000000 pJ +resize/with_sizes,PASS,0.052 s,0.00 MiB,0.00 MiB,3,0,0.000797 ms,92.542033 mW,73756.000000 pJ +sigmoid/4d,PASS,0.067 s,0.00 MiB,0.00 MiB,1,0,0.000521 ms,78.184261 mW,40734.000000 pJ +sigmoid/after_gemm,PASS,0.070 s,0.01 MiB,0.01 MiB,5,4,0.007513 ms,105.158653 mW,790056.960000 pJ +sigmoid/basic,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ +slice/2d_basic,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ +slice/after_conv,PASS,0.062 s,0.00 MiB,0.01 MiB,7,6,0.011296 ms,118.190765 mW,1335082.880000 pJ +slice/default_axes,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ +slice/large_channel_1024,PASS,0.049 s,0.01 MiB,0.00 MiB,1,0,0.002832 ms,78.144068 mW,221304.000000 pJ +slice/nchw_spatial_crop,PASS,0.063 s,0.00 MiB,0.00 MiB,1,0,0.001302 ms,78.239631 mW,101868.000000 pJ +slice/negative_axis,PASS,0.058 s,0.00 MiB,0.00 MiB,1,0,0.000562 ms,78.298932 mW,44004.000000 pJ +slice/negative_indices,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,0.000322 ms,78.298137 mW,25212.000000 pJ +slice/step2,PASS,0.054 s,0.00 MiB,0.00 MiB,1,0,0.002042 ms,78.293830 mW,159876.000000 pJ +softmax/3d_last_axis,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED +softmax/basic,PASS,0.069 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED +softmax/channel_axis,PASS,0.061 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED +softmax/large_dimension_1024,PASS,0.045 s,0.01 MiB,0.01 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED +softmax/negative_axis,PASS,0.064 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED +split/basic,PASS,0.066 s,0.00 MiB,0.00 MiB,1,0,0.000403 ms,78.297767 mW,31554.000000 pJ +split/equal_three_way,PASS,0.065 s,0.00 MiB,0.00 MiB,1,0,0.000564 ms,78.297872 mW,44160.000000 pJ +split/negative_axis,PASS,0.058 s,0.00 MiB,0.00 MiB,1,0,0.001083 ms,78.288089 mW,84786.000000 pJ +split/uneven_channel_axis_4d,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ sub/after_gemm,PASS,0.068 s,0.01 MiB,0.01 MiB,5,4,0.007784 ms,104.703618 mW,815012.960000 pJ -sub/basic,PASS,0.063 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -sub/broadcast_row,PASS,0.075 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ -sub/channel_broadcast_1024,PASS,0.052 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ +sub/basic,PASS,0.060 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +sub/broadcast_row,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +sub/channel_broadcast_1024,PASS,0.057 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ sub/constant_lhs_broadcast,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,0.000322 ms,78.223602 mW,25188.000000 pJ -sub/leading_dimension_broadcast,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ +sub/leading_dimension_broadcast,PASS,0.058 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ diff --git a/validation/tools/analyze_yolo11n_attention.py b/validation/tools/analyze_yolo11n_attention.py new file mode 100644 index 0000000..8453e4f --- /dev/null +++ b/validation/tools/analyze_yolo11n_attention.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +"""Save exact attention-tap arrays and quantify the MatMul error sources.""" + +import argparse +import hashlib +import json +import sys +from pathlib import Path + +import numpy as np +import onnx +from onnx import numpy_helper + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT / "validation")) + +from raptor_validation.onnx_utils import onnx_io # noqa: E402 +from raptor_validation.validate_one import ( # noqa: E402 + parse_pim_simulator_outputs, + sanitize_output_name, +) + + +TAP_NAMES = { + "v": "/model.10/m/m.0/attn/Split_output_2", + "raw": "/model.10/m/m.0/attn/MatMul_output_0", + "scaled": "/model.10/m/m.0/attn/Mul_output_0", + "rhs": "/model.10/m/m.0/attn/Transpose_1_output_0", + "c": "/model.10/m/m.0/attn/MatMul_1_output_0", +} +ABSOLUTE_TOLERANCE = 1e-3 +RELATIVE_TOLERANCE = 1e-5 + + +def sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(1 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def metric(actual, expected): + difference = np.abs(actual.astype(np.float64) - expected.astype(np.float64)) + allowed = ABSOLUTE_TOLERANCE + RELATIVE_TOLERANCE * np.abs(expected.astype(np.float64)) + return { + "max_abs": float(np.max(difference)), + "mean_abs": float(np.mean(difference)), + "rms": float(np.sqrt(np.mean(np.square(difference)))), + "elements_over_validator_limit": int(np.count_nonzero(difference > allowed)), + } + + +def f32_matmul(lhs, rhs): + return np.matmul(lhs.astype(np.float32), rhs.astype(np.float32)).astype(np.float32) + + +def f64_matmul(lhs, rhs): + return np.matmul(lhs.astype(np.float64), rhs.astype(np.float64)).astype(np.float64) + + +def load_constant(model, output_name): + for initializer in model.graph.initializer: + if initializer.name == output_name: + return float(numpy_helper.to_array(initializer).reshape(-1)[0]) + for node in model.graph.node: + if output_name not in node.output: + continue + for attribute in node.attribute: + if attribute.name == "value" and attribute.HasField("t"): + return float(numpy_helper.to_array(attribute.t).reshape(-1)[0]) + raise ValueError(f"could not find ONNX Constant producing {output_name}") + + +def load_arrays(workspace, model_path): + model = onnx.load(model_path) + descriptors = onnx_io(model_path) + output_descriptors = {name: (index, dtype, shape) for index, name, dtype, shape in descriptors[1]} + missing = sorted(set(TAP_NAMES.values()) - set(output_descriptors)) + if missing: + raise ValueError("tap model is missing outputs: " + ", ".join(missing)) + + sim_arrays = parse_pim_simulator_outputs( + workspace / "simulation" / "out.bin", descriptors[1] + ) + reference = {} + simulated = {} + input_files = {} + for key, name in TAP_NAMES.items(): + index, _dtype, shape = output_descriptors[name] + csv_path = workspace / "outputs" / f"output{index}_{sanitize_output_name(name)}.csv" + reference[key] = np.loadtxt(csv_path, delimiter=",", dtype=np.float32).reshape(shape) + simulated[key] = np.asarray(sim_arrays[index], dtype=np.float32).reshape(shape) + input_files[key] = csv_path + return reference, simulated, input_files + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--workspace", type=Path, required=True, + help="validator workspace containing inputs, outputs, and simulation") + parser.add_argument("--model", type=Path, required=True, help="five-output ONNX tap model") + parser.add_argument("--output-dir", type=Path, required=True, + help="directory for arrays.npz, metadata.json, and decomposition.json") + args = parser.parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + + model = onnx.load(args.model) + reference, simulated, source_files = load_arrays(args.workspace, args.model) + scale = np.float32(load_constant(model, "/model.10/m/m.0/attn/Constant_1_output_0")) + ref_v, ref_raw, ref_scaled, ref_rhs, ref_c = (reference[key] for key in ("v", "raw", "scaled", "rhs", "c")) + sim_v, sim_raw, sim_scaled, sim_rhs, sim_c = (simulated[key] for key in ("v", "raw", "scaled", "rhs", "c")) + + ref_score_transpose = np.swapaxes(ref_scaled, -1, -2) + sim_score_transpose = np.swapaxes(sim_scaled, -1, -2) + ref_ss_f32 = f32_matmul(ref_v, ref_rhs) + sim_ss_f32 = f32_matmul(sim_v, sim_rhs) + ref_ss_f64 = f64_matmul(ref_v, ref_rhs) + sim_ss_f64 = f64_matmul(sim_v, sim_rhs) + ref_split_f32 = (f32_matmul(ref_v, np.swapaxes(ref_raw, -1, -2)) * scale).astype(np.float32) + sim_split_f32 = (f32_matmul(sim_v, np.swapaxes(sim_raw, -1, -2)) * scale).astype(np.float32) + ref_split_f64 = f64_matmul(ref_v, np.swapaxes(ref_raw, -1, -2)) * np.float64(scale) + sim_split_f64 = f64_matmul(sim_v, np.swapaxes(sim_raw, -1, -2)) * np.float64(scale) + + arrays = { + **{f"ref_{key}": value for key, value in reference.items()}, + **{f"sim_{key}": value for key, value in simulated.items()}, + "ref_ss_f32": ref_ss_f32, + "sim_ss_f32": sim_ss_f32, + "ref_ss_f64": ref_ss_f64, + "sim_ss_f64": sim_ss_f64, + "ref_split_f32": ref_split_f32, + "sim_split_f32": sim_split_f32, + "ref_split_f64": ref_split_f64, + "sim_split_f64": sim_split_f64, + } + arrays_path = args.output_dir / "arrays.npz" + np.savez_compressed(arrays_path, **arrays) + + metrics = { + "validator_policy": { + "absolute_tolerance": ABSOLUTE_TOLERANCE, + "relative_tolerance": RELATIVE_TOLERANCE, + }, + "scale": float(scale), + "shape": list(ref_c.shape), + "tap_differences": {key: metric(simulated[key], reference[key]) for key in TAP_NAMES}, + "rhs_transpose_consistency": metric(ref_rhs, ref_score_transpose), + "sim_rhs_transpose_consistency": metric(sim_rhs, sim_score_transpose), + "c_sim_vs_ss_f32": metric(sim_c, ref_ss_f32), + "c_ref_vs_ss_f32": metric(ref_c, ref_ss_f32), + "c_sim_vs_simulated_inputs_ss_f32": metric(sim_c, sim_ss_f32), + "v_drift_only": metric(f32_matmul(sim_v, ref_rhs), ref_ss_f32), + "rhs_drift_only": metric(f32_matmul(ref_v, sim_rhs), ref_ss_f32), + "joint_input_drift": metric(sim_ss_f32, ref_ss_f32), + "scale_reassociation_reference": metric(ref_split_f32, ref_ss_f32), + "scale_reassociation_simulated": metric(sim_split_f32, sim_ss_f32), + "reference_accumulation_f32_vs_f64": metric(ref_ss_f32, ref_ss_f64), + "simulated_accumulation_f32_vs_f64": metric(sim_ss_f32, sim_ss_f64), + "split_accumulation_reference_f32_vs_f64": metric(ref_split_f32, ref_split_f64), + "split_accumulation_simulated_f32_vs_f64": metric(sim_split_f32, sim_split_f64), + } + decomposition_path = args.output_dir / "decomposition.json" + decomposition_path.write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8") + + metadata = { + "model": str(args.model), + "model_sha256": sha256(args.model), + "workspace": str(args.workspace), + "arrays_sha256": sha256(arrays_path), + "source_sha256": {key: sha256(path) for key, path in source_files.items()}, + "simulator_output_sha256": sha256(args.workspace / "simulation" / "out.bin"), + "input_sha256": sha256(args.workspace / "inputs" / "in0.csv"), + "outputs": TAP_NAMES, + "arrays": {key: {"dtype": str(value.dtype), "shape": list(value.shape)} for key, value in arrays.items()}, + } + (args.output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8") + print(json.dumps(metrics, indent=2)) + + +if __name__ == "__main__": + main()