From 10b6ee6c32c311e6dc053de2a43ee1e9d5a95312 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Tue, 4 Aug 2026 11:28:05 +0200 Subject: [PATCH] big refactor --- README.md | 46 +- src/PIM/CMakeLists.txt | 2 +- src/PIM/Common/IR/AddressAnalysis.cpp | 76 +- src/PIM/Compiler/PimCompilerOptions.cpp | 22 +- src/PIM/Compiler/PimCompilerOptions.hpp | 12 +- src/PIM/Compiler/PimCompilerUtils.cpp | 83 +- .../Conversion/ONNXToSpatial/CMakeLists.txt | 13 +- .../Common/ComputeRegionBuilder.cpp | 16 +- .../Common/ComputeRegionBuilder.hpp | 33 + .../Common/ContractionPlanning.cpp | 42 + .../Common/ContractionPlanning.hpp | 30 + .../Common/ContractionProblem.hpp | 26 + .../Common/MatrixProductLowering.cpp | 55 + .../Common/MatrixProductLowering.hpp | 8 + .../Common/RowStripLayoutUtils.cpp | 69 +- .../Common/RowStripLayoutUtils.hpp | 23 + .../ONNXToSpatial/Common/ShapeTilingUtils.cpp | 11 +- .../ONNXToSpatial/Common/ShapeTilingUtils.hpp | 6 +- .../ONNXToSpatial/LowerSpatialPlansPass.cpp | 739 ---- .../ONNXToSpatial/ONNXToSpatialOptions.hpp | 32 + .../Analyses}/ONNXToSpatialVerifier.cpp | 7 +- .../Analyses}/ONNXToSpatialVerifier.hpp | 0 .../Analyses/SpatialLayoutCapabilities.cpp | 152 + .../Transforms/LowerSpatialPlansPass.cpp | 136 + .../Transforms}/ONNXToSpatialPass.cpp | 43 +- .../Passes/Transforms/PlanLowering.hpp | 92 + .../Transforms/SpatialLayoutPlanningPass.cpp | 318 ++ .../SpatialPlanLoweringPatterns.cpp | 845 +++++ .../SpatialPlanLoweringPatterns.hpp | 20 + src/PIM/Conversion/ONNXToSpatial/Patterns.cpp | 13 +- src/PIM/Conversion/ONNXToSpatial/Patterns.hpp | 30 +- .../ONNXToSpatial/Patterns/Math/Conv.cpp | 3180 ++++++----------- .../Patterns/Math/ConvGeometry.cpp | 150 +- .../Patterns/Math/ConvGeometry.hpp | 74 +- .../Patterns/Math/Elementwise.cpp | 59 +- .../ONNXToSpatial/Patterns/Math/Gemm.cpp | 267 +- .../ONNXToSpatial/Patterns/Math/Gemm.hpp | 27 + .../ONNXToSpatial/Patterns/Math/MatMul.cpp | 730 ++-- .../Patterns/Math/ReduceMean.cpp | 6 +- .../ONNXToSpatial/Patterns/NN/Pool.cpp | 278 +- .../ONNXToSpatial/Patterns/NN/Relu.cpp | 2 +- .../ONNXToSpatial/Patterns/Tensor/Concat.cpp | 3 +- .../ONNXToSpatial/Patterns/Tensor/Flatten.cpp | 136 +- .../ONNXToSpatial/Patterns/Tensor/Resize.cpp | 279 +- .../Patterns/Tensor/Transpose.cpp | 76 + .../Conversion/ONNXToSpatial/PlanLowering.hpp | 44 - .../SpatialLayoutPlanningPass.cpp | 341 -- .../BatchCoreLoweringPatterns.cpp | 9 +- src/PIM/Conversion/SpatialToPim/Common.cpp | 26 + src/PIM/Conversion/SpatialToPim/Common.hpp | 5 + .../SpatialToPim/CoreLoweringPatterns.cpp | 19 +- src/PIM/Conversion/SpatialToPim/Patterns.cpp | 13 +- .../SpatialToPim/ReturnPathNormalization.cpp | 6 +- .../SpatialToPim/SpatialToPimPass.cpp | 32 +- .../SpatialToPim/SpatialToPimPass.hpp | 8 +- src/PIM/Dialect/Pim/CMakeLists.txt | 12 +- .../Analyses}/CMakeLists.txt | 0 .../Analyses}/LocalMemoryLifetimeAnalysis.cpp | 2 +- .../Analyses}/LocalMemoryLifetimeAnalysis.hpp | 0 .../Bufferization/BufferizationUtils.cpp | 4 +- .../Bufferization/BufferizationUtils.hpp | 0 .../Transforms/Bufferization/CMakeLists.txt | 0 .../Transforms/Bufferization/Common.cpp | 2 +- .../Transforms/Bufferization/Common.hpp | 0 .../Bufferization/ContiguityPatterns.cpp | 0 .../Bufferization/ContiguityPatterns.hpp | 0 .../OpBufferizationInterfaces.cpp | 2 +- .../OpBufferizationInterfaces.hpp | 0 .../Bufferization/PimBufferizationPass.cpp | 359 +- .../HostConstantFolding/CMakeLists.txt | 0 .../Transforms/HostConstantFolding/Common.cpp | 0 .../Transforms/HostConstantFolding/Common.hpp | 0 .../HostConstantFoldingPass.cpp | 2 +- .../HostConstantFolding/Patterns.hpp | 0 .../HostConstantFolding/Patterns/Constant.cpp | 73 +- .../HostConstantFolding/Patterns/Subview.cpp | 0 .../InstructionSelection/CMakeLists.txt | 0 .../InstructionSelectionPass.cpp | 4 +- .../LocalMemoryPlanning/CMakeLists.txt | 0 .../LocalMemoryPlanning.cpp | 4 +- .../LocalMemoryPlanning.hpp | 2 +- .../Transforms/Verification/CMakeLists.txt | 1 - .../Verification/VerificationPass.cpp | 51 +- src/PIM/Dialect/Pim/PimOpsVerify.cpp | 11 +- src/PIM/Dialect/Spatial/CMakeLists.txt | 52 +- .../DeferredBoundaryPlanning.cpp | 0 .../DeferredBoundaryPlanning.hpp | 0 .../DeferredBoundaryRealization.cpp | 0 .../DeferredBoundaryRealization.hpp | 0 .../DeferredCommunicationDeadlock.cpp | 0 .../DeferredCommunicationDeadlock.hpp | 0 .../DeferredCommunicationModel.hpp | 0 .../DeferredCommunicationPlanning.cpp | 109 +- .../DeferredCommunicationPlanning.hpp | 0 .../DeferredCommunicationRealization.cpp | 24 + .../DeferredCommunicationRealization.hpp | 0 .../DeferredCommunicationScheduling.cpp | 0 .../DeferredCommunicationScheduling.hpp | 0 .../DeferredProjectionAnalysis.cpp | 28 +- .../DeferredProjectionAnalysis.hpp | 0 .../DeferredResultRealization.cpp | 4 +- .../DeferredResultRealization.hpp | 0 .../DeferredTransferPlanning.cpp | 5 +- .../DeferredTransferPlanning.hpp | 0 .../ScheduledComputeMaterialization.cpp | 0 .../ScheduledComputeMaterialization.hpp | 0 .../ScheduledComputePlan.hpp | 0 .../ScheduledComputePlanning.cpp | 2 +- .../ScheduledComputeReport.cpp | 0 .../ScheduledComputeReport.hpp | 0 .../ScheduledComputeVerification.cpp | 0 .../ScheduledComputeVerification.hpp | 0 .../ScheduledSpatialPasses.cpp | 144 + .../Scheduling/ComputeGraph.cpp | 0 .../Scheduling/ComputeGraph.hpp | 0 .../Scheduling/ComputeInstance.hpp | 0 .../Scheduling/ComputeInstanceUtils.cpp | 0 .../Scheduling/ComputeInstanceUtils.hpp | 0 .../Scheduling/MergeSchedule.hpp | 0 .../Scheduling/MergeSchedulingAnalysis.cpp | 0 .../Scheduling/MergeSchedulingAnalysis.hpp | 0 .../Scheduling/PeftScheduler.cpp | 342 +- .../Scheduling/PeftScheduler.hpp | 0 .../Scheduling/SchedulingTarget.hpp | 0 .../MergeComputeNodes/Scheduling/Utils.hpp | 0 .../SpatialDataflowCsvExporter.cpp | 353 +- .../SpatialDataflowCsvExporter.hpp | 2 - .../TrivialGraphComputeMergePass.cpp | 26 +- src/PIM/Dialect/Spatial/Spatial.td | 149 +- src/PIM/Dialect/Spatial/SpatialOps.cpp | 98 + src/PIM/Dialect/Spatial/SpatialOps.hpp | 73 + src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp | 23 +- src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp | 123 +- .../Spatial/SpatialTargetResources.hpp | 27 + .../MergeComputeNodesPass.cpp | 128 - src/PIM/Pass/PIMPasses.h | 41 - src/PIM/{Pass => Passes}/CMakeLists.txt | 0 src/PIM/{Pass => Passes}/MessagePass.cpp | 0 src/PIM/Passes/PIMPasses.h | 65 + .../PimCodegen/EmitPimCodePass.cpp | 0 src/PIM/PimAccelerator.cpp | 21 +- test/PIM/PimMemoryLivenessPlannerTest.cpp | 2 +- test/PIM/SpatialSchedulingTargetTest.cpp | 2 +- validation/.gitignore | 6 + validation/operations/README.md | 9 +- .../conv_yolo11n_depthwise_head.onnx | Bin 0 -> 19547 bytes .../yolo11n_heavy/conv_yolo11n_heavy.onnx | Bin 0 -> 295828 bytes validation/operations/gen_tests.py | 63 + .../yolo_attention/matmul_yolo_attention.onnx | Bin 0 -> 301 bytes validation/operations/validation_results.csv | 333 +- 150 files changed, 6737 insertions(+), 4816 deletions(-) create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp delete mode 100644 src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp rename src/PIM/Conversion/ONNXToSpatial/{ => Passes/Analyses}/ONNXToSpatialVerifier.cpp (97%) rename src/PIM/Conversion/ONNXToSpatial/{ => Passes/Analyses}/ONNXToSpatialVerifier.hpp (100%) create mode 100644 src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/SpatialLayoutCapabilities.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/LowerSpatialPlansPass.cpp rename src/PIM/Conversion/ONNXToSpatial/{ => Passes/Transforms}/ONNXToSpatialPass.cpp (85%) create mode 100644 src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialLayoutPlanningPass.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.cpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp create mode 100644 src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp delete mode 100644 src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp delete mode 100644 src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp rename src/PIM/Dialect/Pim/{Analysis => Passes/Analyses}/CMakeLists.txt (100%) rename src/PIM/Dialect/Pim/{Analysis => Passes/Analyses}/LocalMemoryLifetimeAnalysis.cpp (98%) rename src/PIM/Dialect/Pim/{Analysis => Passes/Analyses}/LocalMemoryLifetimeAnalysis.hpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/BufferizationUtils.cpp (91%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/BufferizationUtils.hpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/CMakeLists.txt (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/Common.cpp (98%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/Common.hpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/ContiguityPatterns.cpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/ContiguityPatterns.hpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/OpBufferizationInterfaces.cpp (99%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/OpBufferizationInterfaces.hpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Bufferization/PimBufferizationPass.cpp (79%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/HostConstantFolding/CMakeLists.txt (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/HostConstantFolding/Common.cpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/HostConstantFolding/Common.hpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp (95%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/HostConstantFolding/Patterns.hpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/HostConstantFolding/Patterns/Constant.cpp (90%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/HostConstantFolding/Patterns/Subview.cpp (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/InstructionSelection/CMakeLists.txt (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/InstructionSelection/InstructionSelectionPass.cpp (98%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/LocalMemoryPlanning/CMakeLists.txt (100%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp (97%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp (87%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Verification/CMakeLists.txt (89%) rename src/PIM/Dialect/Pim/{ => Passes}/Transforms/Verification/VerificationPass.cpp (94%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp (85%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp (91%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp (97%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredResultRealization.cpp (99%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredResultRealization.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp (98%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp (99%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp (100%) create mode 100644 src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/ComputeInstance.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp (71%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/Scheduling/Utils.hpp (100%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp (72%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp (94%) rename src/PIM/Dialect/Spatial/{ => Passes}/Transforms/TrivialGraphComputeMergePass.cpp (96%) create mode 100644 src/PIM/Dialect/Spatial/SpatialTargetResources.hpp delete mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp delete mode 100644 src/PIM/Pass/PIMPasses.h rename src/PIM/{Pass => Passes}/CMakeLists.txt (100%) rename src/PIM/{Pass => Passes}/MessagePass.cpp (100%) create mode 100644 src/PIM/Passes/PIMPasses.h rename src/PIM/{Pass => Passes}/PimCodegen/EmitPimCodePass.cpp (100%) create mode 100644 validation/operations/conv/yolo11n_depthwise_head/conv_yolo11n_depthwise_head.onnx create mode 100644 validation/operations/conv/yolo11n_heavy/conv_yolo11n_heavy.onnx create mode 100644 validation/operations/matmul/yolo_attention/matmul_yolo_attention.onnx diff --git a/README.md b/README.md index 5519851..f546d39 100644 --- a/README.md +++ b/README.md @@ -52,27 +52,41 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts `Patterns/{Math,NN,Tensor}` and currently cover Conv, Gemm, MatMul, elementwise Add/Mul/Div, ReduceMean, pooling, Relu, Sigmoid, Softmax, Concat, Gather, Reshape, Resize, and Split. + The compiler-layer target adapter supplies the target-neutral + `SpatialTargetResources`. Layout-aware plan ops advertise typed alternatives + through the Spatial layout interface; the layout planner records the + selected layout and explicit materialization edges. `LowerSpatialPlans` + then pattern-lowers those selected plans. Contraction and Conv lowering + keep semantic problems, target-dependent plans, and IR materializers in + separate layers. Passes and their invariant/layout analyses live under + `Passes/Transforms` and `Passes/Analyses`. -2. **Merge compute nodes** - (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`). - Builds a compute graph, schedules it with the PEFT scheduler, and materializes - the merge schedule into Spatial IR. Supporting scheduling code lives under - `MergeComputeNodes/Scheduling`. +2. **Merge, schedule, and realize Spatial communication** + (`src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes`). + `TrivialGraphComputeMerge` performs local graph merging. One + `ScheduleAndRealizeSpatial` pass then owns scheduling, intermediate + verification, communication realization, and final verification. Supporting + scheduling code lives under `MergeComputeNodes/Scheduling`. 3. **Spatial -> Pim** (`src/PIM/Conversion/SpatialToPim`). Lowers Spatial operations to the `pim` dialect (`src/PIM/Dialect/Pim`), including `pim.core`, `pim.core_batch`, communication, tensor packing, global tensor materialization, and return-path normalization. -4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`). - Converts tensor-semantics PIM IR into memref-semantics PIM IR using MLIR's - bufferization interfaces. +4. **Bufferization** (`src/PIM/Dialect/Pim/Passes/Transforms/Bufferization`). + `PimBufferizationPreparation` establishes writable destinations without + duplicating the one-shot copy analysis, `PimOneShotBufferization` runs + MLIR's one-shot analysis, + `PimMemoryNormalization` forwards/removes redundant copies and normalizes + addressable accesses, and `PimBufferizationVerification` checks tensor + absence, contiguity, and copy address spaces. 5. **PIM local-memory planning** - (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`). + (`src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning`). Computes whole-core lifetimes, reuses addresses for non-overlapping - allocations, and records the explicit plan in PIM IR. -6. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and + allocations, and records the explicit plan in PIM IR. Reusable lifetime + analysis lives under `src/PIM/Dialect/Pim/Passes/Analyses`. +6. **PIM verification and code generation** (`src/PIM/Passes/PimCodegen` and `src/PIM/Compiler`). Verifies the memory plan and other PIM invariants, then emits `.pim` core files, weights, and `memory.bin` / `config.json` without rerunning liveness. @@ -85,7 +99,7 @@ Supporting pieces: points. - `src/PIM/Conversion/SpatialToGraphviz` - optional Spatial graphviz conversion pass. -- `src/PIM/Pass` - pass registration and auxiliary passes. +- `src/PIM/Passes` - pass registration and auxiliary passes. - `src/PIM/PimAccelerator.{cpp,hpp}` - ONNX-MLIR accelerator entry point. ## PIM compiler options @@ -118,16 +132,8 @@ 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. -- `--pim-materialize-scalar-fanout-global-order` - use the experimental, - expensive globally ordered scalar-fanout materializer. Default is off. -- `--pim-trace-communication-materialization` - emit verbose communication - materialization diagnostics and provenance attributes. Default is off. -- `--ignore-concat-error` - soft-fail a ConcatOp corner case. ## Standard PIM hardware profile diff --git a/src/PIM/CMakeLists.txt b/src/PIM/CMakeLists.txt index cd94643..b575b53 100644 --- a/src/PIM/CMakeLists.txt +++ b/src/PIM/CMakeLists.txt @@ -94,7 +94,7 @@ endfunction() add_subdirectory(Dialect) add_subdirectory(Common) -add_subdirectory(Pass) +add_subdirectory(Passes) add_subdirectory(Compiler) add_subdirectory(Conversion) 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..62704b2 100644 --- a/src/PIM/Compiler/PimCompilerOptions.cpp +++ b/src/PIM/Compiler/PimCompilerOptions.cpp @@ -70,11 +70,6 @@ llvm::cl::opt llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions)); -llvm::cl::opt useExperimentalConvImpl("use-experimental-conv-impl", - llvm::cl::desc("Use experimental implementation for convolution"), - llvm::cl::init(false), - llvm::cl::cat(OnnxMlirOptions)); - llvm::cl::opt pimConvIm2colMaxElements( "pim-conv-im2col-max-elements", llvm::cl::desc("Maximum number of im2col elements to materialize globally for one Conv before streaming/chunking"), @@ -103,15 +98,9 @@ llvm::cl::opt pimDetectCommunicationDeadlock( llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions)); -llvm::cl::opt pimMaterializeScalarFanoutGlobalOrder( - "pim-materialize-scalar-fanout-global-order", - llvm::cl::desc("Experimental expensive materializer mode: emit scalar-source fanout as globally ordered communication events instead of all-send fanout loops"), - llvm::cl::init(false), - llvm::cl::cat(OnnxMlirOptions)); - -llvm::cl::opt pimTraceCommunicationMaterialization( - "pim-trace-communication-materialization", - llvm::cl::desc("Emit verbose materializer-time diagnostics and provenance attributes for every Spatial communication op"), +llvm::cl::opt pimVerifyBufferizationCopyFreedom( + "pim-verify-bufferization-copy-freedom", + llvm::cl::desc("Run the expensive official PIM tensor-copy freedom proof before bufferization"), llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions)); @@ -131,11 +120,6 @@ llvm::cl::opt pimTargetConfig( llvm::cl::init(""), llvm::cl::cat(OnnxMlirOptions)); -llvm::cl::opt - ignoreConcatError("ignore-concat-error", - llvm::cl::desc("Ignore ConcatOp corner case: do not assert and do a simplification"), - llvm::cl::init(false)); - bool hasExplicitPimCoreCount() { return coresCount.getNumOccurrences() != 0; } void verifyExplicitPimCoreCount() { diff --git a/src/PIM/Compiler/PimCompilerOptions.hpp b/src/PIM/Compiler/PimCompilerOptions.hpp index 00c2472..baff5c7 100644 --- a/src/PIM/Compiler/PimCompilerOptions.hpp +++ b/src/PIM/Compiler/PimCompilerOptions.hpp @@ -55,12 +55,10 @@ extern llvm::cl::opt pimConvLowering; 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; +extern llvm::cl::opt pimVerifyBufferizationCopyFreedom; extern llvm::cl::opt crossbarSize; extern llvm::cl::opt crossbarCountInCore; @@ -72,12 +70,4 @@ extern llvm::cl::opt pimConvStreamChunkPositions; bool hasExplicitPimCoreCount(); void verifyExplicitPimCoreCount(); -// This option, by default set to false, will ignore an error when resolving a -// specific tiles of the operands of a concat. This specific case is when the -// wanted tile is generated by two separate operands of the concat. If this is -// set to false, this corner case will assert an error. If this is set to true, -// a simplification is performed and only the tile from the first operand is -// taken. -extern llvm::cl::opt ignoreConcatError; - } // namespace onnx_mlir diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index 6e64791..0ed3c4f 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -14,9 +14,12 @@ #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp" +#include "src/Accelerators/PIM/Passes/PIMPasses.h" #include "src/Compiler/CompilerPasses.hpp" #define DEBUG_TYPE "PimCompilerUtils" @@ -80,6 +83,54 @@ spatial::SchedulingTarget getDefaultPimSchedulingTarget() { return target; } +spatial::ConvLoweringStrategy getSpatialConvLoweringStrategy(PimConvLoweringType strategy) { + switch (strategy) { + case PimConvLoweringAuto: return spatial::ConvLoweringStrategy::Auto; + case PimConvLoweringLegacy: return spatial::ConvLoweringStrategy::Legacy; + case PimConvLoweringDepthwise: return spatial::ConvLoweringStrategy::Depthwise; + case PimConvLoweringPackedIm2Col: return spatial::ConvLoweringStrategy::PackedIm2Col; + case PimConvLoweringStreamedPatch: return spatial::ConvLoweringStrategy::StreamedPatch; + case PimConvLoweringStreamedPacked: return spatial::ConvLoweringStrategy::StreamedPacked; + case PimConvLoweringOutputChannelTiled: return spatial::ConvLoweringStrategy::OutputChannelTiled; + case PimConvLoweringInputKTiled: return spatial::ConvLoweringStrategy::InputKTiled; + case PimConvLoweringTiled2D: return spatial::ConvLoweringStrategy::Tiled2D; + } + llvm_unreachable("unknown PIM Conv lowering strategy"); +} + +spatial::SpatialDataflowExportStage getPimSpatialDataflowExportStage( + PimSpatialDataflowExportType stage) { + switch (stage) { + case SpatialDataflowExportNone: return spatial::SpatialDataflowExportStage::None; + case SpatialDataflowExportSpatial1: return spatial::SpatialDataflowExportStage::Spatial1; + case SpatialDataflowExportSpatial2: return spatial::SpatialDataflowExportStage::Spatial2; + case SpatialDataflowExportSpatial3: return spatial::SpatialDataflowExportStage::Spatial3; + case SpatialDataflowExportSpatial4: return spatial::SpatialDataflowExportStage::Spatial4; + case SpatialDataflowExportAll: return spatial::SpatialDataflowExportStage::All; + } + llvm_unreachable("unknown PIM Spatial dataflow export stage"); +} + +spatial::SpatialTargetResources getPimSpatialTargetResources(const spatial::SchedulingTarget& target) { + spatial::SpatialTargetResources resources; + resources.matrixShape = {target.matrixRows, target.matrixColumns}; + resources.matrixUnitsPerProcessor = target.residentWeightCapacity; + resources.processorCount = target.processorCount; + resources.vectorWidth = target.vectorWidth; + if (failed(resources.verify())) + llvm::report_fatal_error("PIM target resources are incomplete"); + return resources; +} + +ONNXToSpatialPlanningOptions getPimONNXToSpatialPlanningOptions() { + ONNXToSpatialPlanningOptions options; + options.convIm2colMaxElements = pimConvIm2colMaxElements.getValue(); + options.convStreamChunkPositions = pimConvStreamChunkPositions.getValue(); + options.forcedConvStrategy = getSpatialConvLoweringStrategy(pimConvLowering.getValue()); + options.reportConvLowering = pimReportConvLowering.getValue(); + return options; +} + const llvm::json::Object& requireObject(const llvm::json::Object& object, llvm::StringRef key, llvm::StringRef path) { @@ -279,11 +330,13 @@ void addPassesPim(OwningOpRef& module, EmissionTargetType& emissionTarget, std::string outputNameNoExt) { verifyExplicitPimCoreCount(); + spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget(); + spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget); if (pimOnlyCodegen) { pm.addPass(createPimInstructionSelectionPass()); pm.addPass(createPimLocalMemoryPlanningPass()); - pm.addPass(createPimVerificationPass()); + pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue())); pm.addPass(createEmitPimCodePass()); return; } @@ -292,23 +345,29 @@ void addPassesPim(OwningOpRef& module, addONNXToMLIRPasses(pm, /*target CPU*/ false); if (pimEmissionTarget >= EmitSpatial) { - spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget(); - pm.addPass(createONNXToSpatialPass()); - pm.addPass(createSpatialLayoutPlanningPass()); - pm.addPass(createLowerSpatialPlansPass()); + ONNXToSpatialPlanningOptions planningOptions = getPimONNXToSpatialPlanningOptions(); + spatial::SpatialDataflowExportStage exportStage = + getPimSpatialDataflowExportStage(pimExportSpatialDataflow.getValue()); + pm.addPass(createONNXToSpatialPass(targetResources, planningOptions)); + pm.addPass(createSpatialLayoutPlanningPass(targetResources)); + pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage)); pm.addPass(createTrivialGraphComputeMergePass( - schedulingTarget.residentWeightCapacity)); - pm.addPass(createMergeComputeNodesPass(schedulingTarget)); + schedulingTarget.residentWeightCapacity, exportStage)); + pm.addPass(spatial::createScheduleAndRealizeSpatialPass( + schedulingTarget, exportStage)); pm.addPass(createMessagePass("Onnx lowered to Spatial")); } if (pimEmissionTarget >= EmitPim) { - pm.addPass(createSpatialToPimPass()); + pm.addPass(createSpatialToPimPass(targetResources)); pm.addPass(createMessagePass("Spatial lowered to Pim")); } if (pimEmissionTarget >= EmitPimBufferized) { - pm.addPass(createPimBufferizationPass()); + pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue())); + pm.addPass(createPimOneShotBufferizationPass()); + pm.addPass(createPimMemoryNormalizationPass()); + pm.addPass(createPimBufferizationVerificationPass()); pm.addPass(createMessagePass("Pim bufferized")); } @@ -320,7 +379,7 @@ void addPassesPim(OwningOpRef& module, pm.addPass(createMessagePass("Pim instructions selected")); pm.addPass(createPimLocalMemoryPlanningPass()); pm.addPass(createMessagePass("Pim local memory planned")); - pm.addPass(createPimVerificationPass()); + pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue())); pm.addPass(createMessagePass("Pim verified")); pm.addPass(createEmitPimCodePass()); pm.addPass(createMessagePass("Pim code emitted")); diff --git a/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt b/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt index 6d9395c..e218c5d 100644 --- a/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt +++ b/src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt @@ -5,7 +5,7 @@ add_public_tablegen_target(ONNXToSpatialIncGen) add_pim_library(OMONNXToSpatial Patterns.cpp CompileTime.cpp - ONNXToSpatialVerifier.cpp + Passes/Analyses/ONNXToSpatialVerifier.cpp Patterns/Pre.cpp Patterns/Post.cpp Patterns/Math/Conv.cpp @@ -26,12 +26,15 @@ add_pim_library(OMONNXToSpatial Patterns/Tensor/Slice.cpp Patterns/Tensor/Split.cpp Patterns/Tensor/Transpose.cpp - ONNXToSpatialPass.cpp - SpatialLayoutPlanningPass.cpp - LowerSpatialPlansPass.cpp + Passes/Transforms/ONNXToSpatialPass.cpp + Passes/Analyses/SpatialLayoutCapabilities.cpp + Passes/Transforms/SpatialLayoutPlanningPass.cpp + Passes/Transforms/SpatialPlanLoweringPatterns.cpp + Passes/Transforms/LowerSpatialPlansPass.cpp Common/AttributeUtils.cpp Common/BiasAddUtils.cpp Common/ComputeRegionBuilder.cpp + Common/ContractionPlanning.cpp Common/MatrixProductLowering.cpp Common/RowStripLayoutUtils.cpp Common/ShapeTilingUtils.cpp @@ -46,8 +49,6 @@ add_pim_library(OMONNXToSpatial MLIRLinalgDialect MLIRSCFDialect MLIRTosaDialect - OMCompilerOptions - OMPimCompilerOptions OMONNXOps SpatialOps OMPimCommon diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp index 9f430b0..5320441 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp @@ -25,6 +25,9 @@ FailureOr createFragmentAssemblyBlueprint(Value physicalBatch, const int64_t laneCount = physicalType.getDimSize(0); if (laneCount <= 0) return emitError(loc, "fragment assembly requires at least one physical source slot"), failure(); + auto physicalLayoutValue = spatial::symbolizePhysicalLayout(physicalLayout); + if (!physicalLayoutValue) + return emitError(loc, "unknown physical layout for fragment assembly"), failure(); const int64_t fragmentElements = physicalType.getNumElements() / laneCount; SmallVector operandIndices(entries.size(), 0), sourceSlots, sourceOffsets, offsets, sizes, strides(entries.size() * rank, 1); @@ -47,13 +50,18 @@ 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 {}, - rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout), + auto blueprint = spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {}, + spatial::getNCHWLayout(rewriter.getContext()), + spatial::PhysicalLayoutAttr::get(rewriter.getContext(), *physicalLayoutValue), rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes), - rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"), + rewriter.getStringAttr(indexMap), spatial::getFragmentAssemblyMode(rewriter.getContext()), 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/ContractionPlanning.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.cpp new file mode 100644 index 0000000..73acfef --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.cpp @@ -0,0 +1,42 @@ +#include "ContractionPlanning.hpp" + +#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" + +#include + +namespace onnx_mlir { + +namespace { + +static int64_t ceilDivide(int64_t value, int64_t divisor) { + return divisor == 0 ? 0 : (value + divisor - 1) / divisor; +} + +} // namespace + +ContractionPlan makeContractionPlan( + const ContractionProblem& problem, + const spatial::SpatialTargetResources& target, + ContractionPlanKind kind, + int64_t laneCount, + int64_t fragmentRows) { + ContractionPlan plan; + plan.tileK = std::max(1, target.matrixShape.rows); + plan.tileN = std::max(1, target.matrixShape.columns); + plan.reductionSlices = std::max(1, ceilDivide(problem.k, plan.tileK)); + plan.outputTiles = std::max(1, ceilDivide(problem.n, plan.tileN)); + const int64_t rowsPerLane = std::max( + 1, fragmentRows != 0 ? fragmentRows : target.matrixShape.rows); + + if (laneCount != 0) + plan.laneCount = laneCount; + else if (kind == ContractionPlanKind::StaticTiled) + plan.laneCount = problem.batch * problem.m * plan.reductionSlices * plan.outputTiles; + else if (kind == ContractionPlanKind::GroupedRowDynamicVVD) + plan.laneCount = problem.batch * ceilDivide(problem.m, rowsPerLane); + else + plan.laneCount = problem.batch * problem.m * problem.n; + return plan; +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp new file mode 100644 index 0000000..7caf908 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include "ContractionProblem.hpp" + +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" + +namespace onnx_mlir { + +enum class ContractionPlanKind { + StaticTiled, + BatchedDynamicVVD, + GroupedRowDynamicVVD, +}; + +struct ContractionPlan { + int64_t tileK = 1; + int64_t tileN = 1; + int64_t reductionSlices = 1; + int64_t outputTiles = 1; + int64_t laneCount = 0; +}; + +ContractionPlan makeContractionPlan( + const ContractionProblem& problem, + const spatial::SpatialTargetResources& target, + ContractionPlanKind kind, + int64_t laneCount = 0, + int64_t fragmentRows = 0); + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp new file mode 100644 index 0000000..db11675 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "mlir/IR/BuiltinTypes.h" + +#include "llvm/ADT/SmallVector.h" + +#include + +namespace onnx_mlir { + +struct ContractionProblem { + llvm::SmallVector lhsBatchShape; + llvm::SmallVector rhsBatchShape; + llvm::SmallVector outputBatchShape; + int64_t lhsBatch = 1; + int64_t rhsBatch = 1; + int64_t batch = 1; + int64_t m = 0; + int64_t k = 0; + int64_t n = 0; + mlir::Type lhsElementType; + mlir::Type rhsElementType; + mlir::Type resultElementType; +}; + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.cpp index 4b0b00d..668bf7c 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.cpp @@ -1,15 +1,70 @@ #include "MatrixProductLowering.hpp" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp" +#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" using namespace mlir; namespace onnx_mlir { +static bool isInsideSpatialCompute(Operation* op) { + for (Operation* parent = op; parent; parent = parent->getParentOp()) + if (spatial::isAnySpatialComputeLike(parent)) + return true; + return false; +} + +static Value buildLinalgTranspose(Value value, + RankedTensorType resultType, + ArrayRef permutation, + PatternRewriter& rewriter, + Location loc) { + Value init = tensor::EmptyOp::create( + rewriter, loc, resultType.getShape(), resultType.getElementType()); + return linalg::TransposeOp::create( + rewriter, loc, value, init, permutation).getResult()[0]; +} + +static Value materializeConstantTranspose(Value value, + RankedTensorType resultType, + ArrayRef permutation, + PatternRewriter& rewriter) { + auto denseAttr = getHostConstDenseElementsAttr(value); + if (!denseAttr) + return {}; + auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation); + if (failed(transposedAttr) || transposedAttr->getType() != resultType) + return {}; + return getOrCreateConstant( + rewriter, rewriter.getInsertionBlock()->getParentOp(), *transposedAttr, resultType); +} + +Value createLinalgTranspose(Value value, + RankedTensorType resultType, + ArrayRef permutation, + PatternRewriter& rewriter, + Location loc) { + if (Value constant = materializeConstantTranspose(value, resultType, permutation, rewriter)) + return constant; + + if (isInsideSpatialCompute(rewriter.getInsertionBlock()->getParentOp())) + return buildLinalgTranspose(value, resultType, permutation, rewriter, loc); + + auto compute = createSpatCompute<1>( + rewriter, loc, TypeRange {resultType}, {}, ValueRange {value}, + [&](Value input) { + spatial::SpatYieldOp::create( + rewriter, loc, buildLinalgTranspose(input, resultType, permutation, rewriter, loc)); + }); + return compute.getResult(0); +} + Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) { auto sourceType = cast(value.getType()); SmallVector lowPads(sourceType.getRank(), rewriter.getIndexAttr(0)); diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp index 45eaff2..1125437 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp @@ -5,8 +5,16 @@ #include "mlir/IR/Value.h" #include "mlir/Transforms/DialectConversion.h" +#include "llvm/ADT/ArrayRef.h" + namespace onnx_mlir { +mlir::Value createLinalgTranspose(mlir::Value value, + mlir::RankedTensorType resultType, + llvm::ArrayRef permutation, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + mlir::Value createZeroPaddedTensor(mlir::Value value, mlir::RankedTensorType resultType, mlir::PatternRewriter& rewriter, diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp index 413911c..180f972 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp @@ -5,9 +5,9 @@ #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" -#include "src/Dialect/ONNX/ONNXOps.hpp" #include @@ -33,6 +33,16 @@ FailureOr describeRowStripPhysicalValue(Value storage, Ra tilesPerRow}; } +FailureOr getRowStripPhysicalValue(Value value) { + auto blueprint = value.getDefiningOp(); + auto logicalType = dyn_cast(value.getType()); + if (!blueprint || !logicalType || blueprint.getOutput() != value + || blueprint.getPhysicalLayout() != spatial::PhysicalLayout::NHWCRowStrip + || !spatial::isPhysicalView(blueprint.getMode())) + return failure(); + return describeRowStripPhysicalValue(blueprint.getInput(), logicalType); +} + RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) { return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3), logicalType.getDimSize(1)}, @@ -144,6 +154,35 @@ FailureOr createRowStripStorageFromRows(Value rows, return batchOp->getResult(0); } +FailureOr createRowStripStorageBlueprint(Value storage, + RankedTensorType logicalType, + PatternRewriter& rewriter, + Location loc) { + FailureOr value = describeRowStripPhysicalValue(storage, logicalType); + if (failed(value)) + return failure(); + + auto blueprint = spatial::SpatBlueprintOp::create( + rewriter, + loc, + logicalType, + storage, + ValueRange {}, + spatial::getNCHWLayout(rewriter.getContext()), + spatial::getNHWCRowStripLayout(rewriter.getContext()), + rewriter.getDenseI64ArrayAttr({}), + rewriter.getDenseI64ArrayAttr({}), + rewriter.getStringAttr(kRowStripIndexMap), + spatial::getPhysicalViewMode(rewriter.getContext()), + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr); + return blueprint.getOutput(); +} + FailureOr createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) { @@ -160,8 +199,8 @@ FailureOr createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va rewriter, loc, args.inputs.front(), args.lane, value.fragmentType); if (failed(fragment)) return failure(); - Value nchw = ONNXTransposeOp::create( - rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2})); + Value nchw = createLinalgTranspose( + *fragment, nchwFragmentType, {0, 3, 1, 2}, rewriter, loc); publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane); return success(); }); @@ -176,7 +215,7 @@ FailureOr createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1, value.logicalType.getDimSize(3)}}); } - return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip", + return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "dense_nchw", kRowStripIndexMap, rewriter, loc); } @@ -186,25 +225,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/Common/RowStripLayoutUtils.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp index a7f75e9..4f6df35 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp @@ -6,6 +6,12 @@ namespace onnx_mlir { +namespace spatial { +class SpatBlueprintOp; +class SpatFlattenPlanOp; +struct SpatialTargetResources; +} // namespace spatial + inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments"; struct RowStripPhysicalValue { @@ -18,6 +24,8 @@ struct RowStripPhysicalValue { mlir::FailureOr describeRowStripPhysicalValue(mlir::Value storage, mlir::RankedTensorType logicalType); +mlir::FailureOr getRowStripPhysicalValue(mlir::Value value); + std::pair, llvm::SmallVector> buildRowStripMetadata(mlir::RankedTensorType type); @@ -53,6 +61,11 @@ mlir::FailureOr createRowStripStorageFromRows(mlir::Value rows, mlir::PatternRewriter& rewriter, mlir::Location loc); +mlir::FailureOr createRowStripStorageBlueprint(mlir::Value storage, + mlir::RankedTensorType logicalType, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + mlir::FailureOr createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value, mlir::PatternRewriter& rewriter, mlir::Location loc); @@ -80,4 +93,14 @@ mlir::FailureOr applyRowStripConcat(llvm::ArrayRef> -sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, PatternRewriter& rewriter, Location loc) { - SmallVector slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc); +sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, + PatternRewriter& rewriter, + Location loc, + const spatial::SpatialTargetResources& target) { + SmallVector slices = sliceVector( + vectorToSlice, static_cast(target.matrixShape.rows), rewriter, loc); DenseMap> slicesPerCore; for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) { - size_t coreId = sliceId / crossbarCountInCore; + size_t coreId = sliceId / target.matrixUnitsPerProcessor; slicesPerCore[coreId].push_back(slices[sliceId]); } return slicesPerCore; diff --git a/src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp b/src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp index 4fb9021..aba0c7d 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp @@ -7,6 +7,7 @@ #include "llvm/ADT/SmallVector.h" #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" namespace onnx_mlir { @@ -26,6 +27,9 @@ llvm::SmallVector sliceVector(const mlir::Value& vectorToSlice, /// Partitions one logical vector into per-core crossbar-sized slices using the /// current PIM target geometry. llvm::DenseMap> sliceVectorPerCrossbarPerCore( - const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc); + const mlir::Value& vectorToSlice, + mlir::PatternRewriter& rewriter, + mlir::Location loc, + const spatial::SpatialTargetResources& target); } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp deleted file mode 100644 index cdd6ff0..0000000 --- a/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp +++ /dev/null @@ -1,739 +0,0 @@ -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Transforms/DialectConversion.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/SmallPtrSet.h" - -#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" -#include "mlir/Transforms/Passes.h" -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" -#include "src/Dialect/ONNX/ONNXOps.hpp" - -using namespace mlir; - -namespace onnx_mlir { -namespace { - -static constexpr StringLiteral kDenseLayout = "dense_nchw"; -static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip"; - -static FailureOr getRowStripValue(llvm::DenseMap& rowStripValues, - Value value) { - auto it = rowStripValues.find(value); - if (it == rowStripValues.end()) - return failure(); - return it->second; -} - -static FailureOr buildRowStripValue(spatial::SpatBlueprintOp blueprint, - Value storage) { - auto logicalType = dyn_cast(blueprint.getOutput().getType()); - if (!logicalType) - return blueprint.emitOpError("requires ranked logical output type"), failure(); - if (blueprint.getIndexMap() != kRowStripIndexMap) - return blueprint.emitOpError("requires the canonical row-strip index map"), failure(); - FailureOr value = describeRowStripPhysicalValue(storage, logicalType); - if (failed(value)) - return blueprint.emitOpError("requires physical row-strip fragment storage"), failure(); - return *value; -} - -static FailureOr -lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) { - return applyRowStripRelu(input, rewriter, planOp.getLoc()); -} - -static FailureOr -lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) { - return applyRowStripSilu(input, rewriter, planOp.getLoc()); -} - -static FailureOr lowerRowStripBiasAdd(const RowStripPhysicalValue& input, - spatial::SpatBiasAddPlanOp planOp, - PatternRewriter& rewriter) { - return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc()); -} - -static FailureOr lowerRowStripAdd(const RowStripPhysicalValue& lhs, - const RowStripPhysicalValue& rhs, - spatial::SpatAddPlanOp planOp, - PatternRewriter& rewriter) { - return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc()); -} - -static FailureOr lowerRowStripConcat(ArrayRef inputs, - spatial::SpatConcatPlanOp planOp, - PatternRewriter& rewriter) { - auto outputType = dyn_cast(planOp.getOutput().getType()); - if (!outputType) - return failure(); - return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc()); -} - -static FailureOr -materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) { - if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape()) - return failure(); - return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc); -} - -static FailureOr lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType, - PatternRewriter& rewriter, Location loc) { - auto producer = input.getDefiningOp(); - auto inputType = dyn_cast(input.getType()); - auto biasType = dyn_cast(bias.getType()); - if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape() - || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount() - || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount()) - return failure(); - auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount()); - auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount()); - if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank() - || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front() - || inputFragmentType->getRank() != outputFragmentType->getRank() + 1) - return failure(); - for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape())) - if (outputDim > inputDim) - return failure(); - - auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias}, - [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult { - FailureOr fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType); - if (failed(fragment)) - return failure(); - MixedSliceGeometry biasSlice; - for (int64_t dim : inputFragmentType->getShape()) { - biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0)); - biasSlice.sizes.push_back(rewriter.getIndexAttr(dim)); - biasSlice.strides.push_back(rewriter.getIndexAttr(1)); - } - Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice); - if (!biasFragment) - return failure(); - Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment); - MixedSliceGeometry outputSlice; - outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0)); - outputSlice.sizes.push_back(rewriter.getIndexAttr(1)); - outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1)); - for (int64_t dim : outputFragmentType->getShape()) - outputSlice.sizes.push_back(rewriter.getIndexAttr(dim)); - Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice); - if (!output) - return failure(); - publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane); - return success(); - }); - if (failed(batch)) - return failure(); - return batch->getResult(0); -} - -static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp, - llvm::DenseMap& rowStripValues, - llvm::SmallPtrSetImpl& eraseAfterLowering, - PatternRewriter& rewriter) { - FailureOr lhs = getRowStripValue(rowStripValues, planOp.getLhs()); - FailureOr rhs = getRowStripValue(rowStripValues, planOp.getRhs()); - if (succeeded(lhs) && succeeded(rhs)) { - 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()) - return planOp.emitOpError("row-strip add plan requires a row-strip blueprint result"); - - rewriter.setInsertionPoint(planOp); - FailureOr lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter); - if (failed(lowered)) - return planOp.emitOpError("failed to lower selected row-strip Spatial add plan"); - auto blueprint = cast(*outputBlueprint); - FailureOr output = buildRowStripValue(blueprint, *lowered); - if (failed(output)) - return failure(); - rowStripValues[blueprint.getResult()] = *output; - eraseAfterLowering.insert(planOp); - eraseAfterLowering.insert(blueprint); - return success(); - } - - rewriter.setInsertionPoint(planOp); - auto compute = createSpatCompute<2>(rewriter, - planOp.getLoc(), - planOp.getOutput().getType(), - {}, - ValueRange {planOp.getLhs(), planOp.getRhs()}, - [&](Value lhsValue, Value rhsValue) { - Value added = spatial::SpatVAddOp::create( - rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue); - spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added); - }); - rewriter.replaceOp(planOp, compute.getResults()); - return success(); -} - -static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp, - llvm::DenseMap& rowStripValues, - llvm::SmallPtrSetImpl& eraseAfterLowering, - PatternRewriter& rewriter) { - SmallVector inputs; - for (Value input : planOp.getInputs()) { - FailureOr physical = getRowStripValue(rowStripValues, input); - if (failed(physical)) { - inputs.clear(); - break; - } - inputs.push_back(*physical); - } - if (inputs.size() == planOp.getInputs().size()) { - 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()) - return planOp.emitOpError("row-strip concat plan requires a row-strip blueprint result"); - rewriter.setInsertionPoint(planOp); - FailureOr lowered = lowerRowStripConcat(inputs, planOp, rewriter); - if (failed(lowered)) - return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan"); - auto blueprint = cast(*outputBlueprint); - FailureOr output = buildRowStripValue(blueprint, *lowered); - if (failed(output)) - return failure(); - rowStripValues[blueprint.getResult()] = *output; - eraseAfterLowering.insert(planOp); - eraseAfterLowering.insert(blueprint); - return success(); - } - - rewriter.setInsertionPoint(planOp); - auto compute = createSpatCompute( - rewriter, - planOp.getLoc(), - TypeRange {planOp.getOutput().getType()}, - {}, - planOp.getInputs(), - [&](ValueRange values) { - Value concatenated = spatial::SpatConcatOp::create( - rewriter, planOp.getLoc(), planOp.getOutput().getType(), rewriter.getI64IntegerAttr(planOp.getAxis()), values); - spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated); - }); - rewriter.replaceOp(planOp, compute.getResults()); - return success(); -} - -struct LowerSpatialPlansPass final : PassWrapper> { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass) - - StringRef getArgument() const override { return "lower-spatial-plans"; } - StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; } - - void runOnOperation() override { - ModuleOp moduleOp = getOperation(); - MLIRContext* ctx = moduleOp.getContext(); - auto entryFunc = getPimEntryFunc(moduleOp); - if (failed(entryFunc)) { - moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans"); - signalPassFailure(); - return; - } - func::FuncOp funcOp = *entryFunc; - PatternRewriter rewriter(ctx); - llvm::DenseMap rowStripValues; - llvm::SmallPtrSet eraseAfterLowering; - auto verifyLogicalPhase = [&](StringRef stage) -> bool { - if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc))) - return true; - moduleOp.emitError() << "logical Spatial graph verification failed " << stage; - signalPassFailure(); - return false; - }; - - if (!verifyLogicalPhase("at the start of LowerSpatialPlans")) - return; - for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) { - if (auto planOp = dyn_cast(&op)) { - FailureOr rowStripInput = getRowStripValue(rowStripValues, planOp.getInput()); - auto rowStripBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { - auto blueprint = dyn_cast(user); - return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout; - }); - if (rowStripBlueprint != planOp.getResult().getUsers().end()) { - rewriter.setInsertionPoint(planOp); - std::optional physicalInput; - if (succeeded(rowStripInput)) - physicalInput = rowStripInput->storage; - FailureOr lowered = lowerSelectedConv2DPlan( - planOp, - physicalInput, - /*emitRowStripLayout=*/true, - rewriter); - if (failed(lowered)) { - auto diagnostic = planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan with input "); - diagnostic << planOp.getInput().getType() << " and output " << planOp.getResult().getType(); - if (physicalInput) - diagnostic << " from physical storage " << physicalInput->getType(); - signalPassFailure(); - return; - } - auto blueprint = cast(*rowStripBlueprint); - FailureOr rowStripValue = buildRowStripValue(blueprint, *lowered); - if (failed(rowStripValue)) { - signalPassFailure(); - return; - } - rowStripValues[blueprint.getResult()] = *rowStripValue; - eraseAfterLowering.insert(planOp); - eraseAfterLowering.insert(blueprint); - continue; - } - rewriter.setInsertionPoint(planOp); - FailureOr lowered = - lowerSelectedConv2DPlan(planOp, std::nullopt, /*emitRowStripLayout=*/false, rewriter); - if (failed(lowered)) { - planOp.emitOpError("failed to lower selected Spatial Conv plan"); - signalPassFailure(); - return; - } - rewriter.replaceOp(planOp, *lowered); - continue; - } - - if (auto planOp = dyn_cast(&op)) { - if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) { - auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { - auto blueprint = dyn_cast(user); - return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout; - }); - if (outputBlueprint == planOp.getResult().getUsers().end()) { - planOp.emitOpError("row-strip Relu plan requires a row-strip blueprint result"); - signalPassFailure(); - return; - } - - FailureOr input = getRowStripValue(rowStripValues, planOp.getInput()); - rewriter.setInsertionPoint(planOp); - FailureOr lowered = lowerRowStripRelu(*input, planOp, rewriter); - if (failed(lowered)) { - planOp.emitOpError("failed to lower selected row-strip Spatial Relu plan"); - signalPassFailure(); - return; - } - auto blueprint = cast(*outputBlueprint); - FailureOr output = buildRowStripValue(blueprint, *lowered); - if (failed(output)) { - signalPassFailure(); - return; - } - rowStripValues[blueprint.getResult()] = *output; - eraseAfterLowering.insert(planOp); - eraseAfterLowering.insert(blueprint); - continue; - } - - rewriter.setInsertionPoint(planOp); - auto computeOp = createSpatCompute<1>( - rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) { - auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x); - spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult()); - }); - rewriter.replaceOp(planOp, computeOp.getResults()); - continue; - } - - if (auto planOp = dyn_cast(&op)) { - if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) { - auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { - auto blueprint = dyn_cast(user); - return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout; - }); - if (outputBlueprint == planOp.getResult().getUsers().end()) { - planOp.emitOpError("row-strip SiLU plan requires a row-strip blueprint result"); - signalPassFailure(); - return; - } - - FailureOr input = getRowStripValue(rowStripValues, planOp.getInput()); - rewriter.setInsertionPoint(planOp); - FailureOr lowered = lowerRowStripSilu(*input, planOp, rewriter); - if (failed(lowered)) { - planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan"); - signalPassFailure(); - return; - } - auto blueprint = cast(*outputBlueprint); - FailureOr output = buildRowStripValue(blueprint, *lowered); - if (failed(output)) { - signalPassFailure(); - return; - } - rowStripValues[blueprint.getResult()] = *output; - eraseAfterLowering.insert(planOp); - eraseAfterLowering.insert(blueprint); - continue; - } - - rewriter.setInsertionPoint(planOp); - auto computeOp = createSpatCompute<1>( - rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) { - Value sigmoid = spatial::SpatSigmoidOp::create( - rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult(); - Value silu = spatial::SpatVMulOp::create( - rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult(); - spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu); - }); - rewriter.replaceOp(planOp, computeOp.getResults()); - continue; - } - if (auto planOp = dyn_cast(&op)) { - 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("selected MaxPool plan requires a row-strip blueprint result"); - signalPassFailure(); - return; - } - - FailureOr input = getRowStripValue(rowStripValues, planOp.getInput()); - rewriter.setInsertionPoint(planOp); - std::optional physicalInput; - if (succeeded(input)) - physicalInput = input->storage; - FailureOr lowered = lowerSelectedMaxPool2DPlan( - planOp, physicalInput, rewriter); - if (failed(lowered)) { - planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan"); - signalPassFailure(); - return; - } - auto blueprint = cast(*outputBlueprint); - FailureOr output = buildRowStripValue(blueprint, *lowered); - if (failed(output)) { - signalPassFailure(); - return; - } - rowStripValues[blueprint.getResult()] = *output; - eraseAfterLowering.insert(planOp); - eraseAfterLowering.insert(blueprint); - continue; - } - if (auto planOp = dyn_cast(&op)) { - 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("selected global AveragePool plan requires a row-strip blueprint result"); - signalPassFailure(); - return; - } - - FailureOr input = getRowStripValue(rowStripValues, planOp.getInput()); - rewriter.setInsertionPoint(planOp); - std::optional physicalInput; - if (succeeded(input)) - physicalInput = input->storage; - FailureOr lowered = - lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, rewriter); - if (failed(lowered)) { - planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan"); - signalPassFailure(); - return; - } - auto blueprint = cast(*outputBlueprint); - FailureOr output = buildRowStripValue(blueprint, *lowered); - if (failed(output)) { - signalPassFailure(); - return; - } - rowStripValues[blueprint.getResult()] = *output; - eraseAfterLowering.insert(planOp); - eraseAfterLowering.insert(blueprint); - continue; - } - if (auto planOp = dyn_cast(&op)) { - if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) { - auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) { - auto blueprint = dyn_cast(user); - return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout; - }); - if (outputBlueprint == planOp.getResult().getUsers().end()) { - planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result"); - signalPassFailure(); - return; - } - - FailureOr input = getRowStripValue(rowStripValues, planOp.getInput()); - rewriter.setInsertionPoint(planOp); - FailureOr lowered = lowerRowStripBiasAdd(*input, planOp, rewriter); - if (failed(lowered)) { - planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan"); - signalPassFailure(); - return; - } - auto blueprint = cast(*outputBlueprint); - FailureOr output = buildRowStripValue(blueprint, *lowered); - if (failed(output)) { - signalPassFailure(); - return; - } - rowStripValues[blueprint.getResult()] = *output; - eraseAfterLowering.insert(planOp); - eraseAfterLowering.insert(blueprint); - continue; - } - - auto resultType = dyn_cast(planOp.getOutput().getType()); - if (!resultType) { - planOp.emitOpError("requires ranked output type"); - signalPassFailure(); - return; - } - rewriter.setInsertionPoint(planOp); - FailureOr denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc()); - if (failed(denseBias)) { - planOp.emitOpError("failed to materialize dense Conv-style bias"); - signalPassFailure(); - return; - } - if (planOp.getInput().getDefiningOp()) { - FailureOr lowered = lowerDenseBatchBiasAdd(planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc()); - if (succeeded(lowered)) { - rewriter.replaceOp(planOp, *lowered); - continue; - } - } - auto computeOp = createSpatCompute<2>(rewriter, - planOp.getLoc(), - planOp.getOutput().getType(), - {}, - ValueRange {planOp.getInput(), *denseBias}, - [&](Value x, Value y) { - auto added = spatial::SpatVAddOp::create( - rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y); - spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult()); - }); - rewriter.replaceOp(planOp, computeOp.getResults()); - continue; - } - if (auto planOp = dyn_cast(&op)) { - if (failed(lowerAddPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) { - signalPassFailure(); - return; - } - continue; - } - if (auto planOp = dyn_cast(&op)) { - if (failed(lowerConcatPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) { - signalPassFailure(); - return; - } - continue; - } - if (auto flattenOp = dyn_cast(&op)) { - if (flattenOp.getInputs().size() == 1) { - FailureOr input = - getRowStripValue(rowStripValues, flattenOp.getInputs().front()); - if (succeeded(input) && succeeded(canLowerFlattenFromRowStrip(flattenOp))) { - rewriter.setInsertionPoint(flattenOp); - if (failed(lowerFlattenFromRowStrip(*input, flattenOp, rewriter))) { - flattenOp.emitOpError("failed to preserve row-strip layout through Flatten"); - signalPassFailure(); - return; - } - continue; - } - } - } - if (auto materializeOp = dyn_cast(&op)) { - if (materializeOp.getSourcePhysicalLayout() == kDenseLayout - && materializeOp.getTargetPhysicalLayout() == kDenseLayout) { - rewriter.replaceOp(materializeOp, materializeOp.getInput()); - continue; - } - if (materializeOp.getSourcePhysicalLayout() != kRowStripLayout - || materializeOp.getTargetPhysicalLayout() != kDenseLayout) { - materializeOp.emitOpError("non-dense materialize_layout lowering is not supported yet"); - signalPassFailure(); - return; - } - FailureOr rowStripValue = getRowStripValue(rowStripValues, materializeOp.getInput()); - if (failed(rowStripValue)) { - materializeOp.emitOpError("expected a row-strip blueprint input during row-strip materialization"); - signalPassFailure(); - return; - } - rewriter.setInsertionPoint(materializeOp); - FailureOr dense = materializeRowStripToDense(*rowStripValue, materializeOp.getLoc(), rewriter); - if (failed(dense)) { - materializeOp.emitOpError("failed to materialize selected row-strip layout back to dense NCHW"); - signalPassFailure(); - return; - } - rewriter.replaceOp(materializeOp, *dense); - continue; - } - if (auto blueprintOp = dyn_cast(&op)) { - if (std::optional mode = blueprintOp.getMode(); mode && *mode == "fragment_assembly") - continue; - if (blueprintOp.getPhysicalLayout() == kDenseLayout) { - rewriter.replaceOp(blueprintOp, blueprintOp.getInput()); - continue; - } - if (blueprintOp.getPhysicalLayout() != kRowStripLayout) { - blueprintOp.emitOpError("non-dense blueprint lowering is not supported yet"); - signalPassFailure(); - return; - } - if (!eraseAfterLowering.contains(blueprintOp)) { - blueprintOp.emitOpError("unhandled row-strip blueprint remained during LowerSpatialPlans"); - signalPassFailure(); - return; - } - } - } - bool erasedAny = true; - while (erasedAny) { - erasedAny = false; - for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) { - if (!eraseAfterLowering.contains(&op)) - continue; - if (!op.use_empty()) - continue; - eraseAfterLowering.erase(&op); - rewriter.eraseOp(&op); - erasedAny = true; - } - } - if (!eraseAfterLowering.empty()) { - for (Operation& op : funcOp.getBody().front()) - if (eraseAfterLowering.contains(&op)) - op.emitOpError("selected row-strip planning op could not be fully eliminated during LowerSpatialPlans"); - signalPassFailure(); - return; - } - ConversionTarget helperTarget(*ctx); - helperTarget.addLegalDialect(); - helperTarget.addLegalOp(); - helperTarget.addIllegalOp(); - helperTarget.markOpRecursivelyLegal(); - - RewritePatternSet helperPatterns(ctx); - populateGemmPatterns(helperPatterns, ctx); - populateTransposePatterns(helperPatterns, ctx); - FrozenRewritePatternSet frozenHelperPatterns( - std::move(helperPatterns)); - SmallVector topLevelHelperOps; - funcOp.walk([&](Operation* op) { - if (isa(op)) - return WalkResult::skip(); - if (isa(op)) - topLevelHelperOps.push_back(op); - return WalkResult::advance(); - }); - for (Operation *helper : topLevelHelperOps) { - if (failed(applyPartialConversion( - helper, helperTarget, frozenHelperPatterns))) { - moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering"); - signalPassFailure(); - return; - } - } - ConversionTarget nestedHelperTarget(*ctx); - nestedHelperTarget.addLegalDialect(); - nestedHelperTarget.addIllegalOp(); - SmallVector computeLikeOps; - funcOp.walk([&](Operation* op) { - if (isa(op)) - computeLikeOps.push_back(op); - }); - for (Operation* op : computeLikeOps) { - if (failed(applyFullConversion( - op, nestedHelperTarget, frozenHelperPatterns))) { - op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering"); - signalPassFailure(); - return; - } - } - if (!verifyLogicalPhase("after nested helper conversions")) - return; - bool hasIllegalOps = false; - moduleOp.walk([&](Operation* op) { - if (isa(op)) - return; - if (auto blueprint = dyn_cast(op)) { - if (std::optional mode = blueprint.getMode(); mode && *mode == "fragment_assembly") - return; - op->emitOpError("planning blueprint must not remain after LowerSpatialPlans"); - hasIllegalOps = true; - } - else if (isa(op) - || op->getDialect()->getNamespace() == "onnx") { - op->emitOpError("operation must not remain after LowerSpatialPlans"); - hasIllegalOps = true; - } - }); - - PassManager canonicalizationPM(ctx); - canonicalizationPM.addPass(createCanonicalizerPass()); - if (failed(canonicalizationPM.run(moduleOp))) - moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing"); - - if (hasIllegalOps) { - signalPassFailure(); - } else { - dumpModule(moduleOp, "spatial1_graph"); - spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage(); - if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1) - && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) { - signalPassFailure(); - return; - } - } - - if (!verifyLogicalPhase("at the end of LowerSpatialPlans")) - return; - } -}; - -} // namespace - -std::unique_ptr createLowerSpatialPlansPass() { return std::make_unique(); } - -} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp new file mode 100644 index 0000000..3710579 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" + +#include + +namespace onnx_mlir::spatial { + +enum class ConvLoweringStrategy : uint8_t { + Auto, + Legacy, + Depthwise, + PackedIm2Col, + StreamedPatch, + StreamedPacked, + OutputChannelTiled, + InputKTiled, + Tiled2D, +}; + +} // namespace onnx_mlir::spatial + +namespace onnx_mlir { + +struct ONNXToSpatialPlanningOptions { + uint64_t convIm2colMaxElements = 0; + uint64_t convStreamChunkPositions = 0; + spatial::ConvLoweringStrategy forcedConvStrategy = spatial::ConvLoweringStrategy::Auto; + bool reportConvLowering = true; +}; + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp similarity index 97% rename from src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp rename to src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp index 9e92dcf..186a945 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp @@ -6,7 +6,7 @@ #include "Common/IR/WeightUtils.hpp" #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" using namespace mlir; @@ -130,8 +130,7 @@ template void verifyNoNestedFragmentAssemblyBlueprints(ComputeOpTy compute, pim::CappedDiagnosticReporter& diagnostics) { compute.getBody().walk([&](spatial::SpatBlueprintOp blueprint) { - std::optional mode = blueprint.getMode(); - if (!mode || *mode != "fragment_assembly") + if (!spatial::isFragmentAssembly(blueprint.getMode())) return; diagnostics.report(blueprint.getOperation(), [&](Operation* illegalOp) { illegalOp->emitOpError("fragment assembly blueprint must be host-level after merge materialization"); @@ -148,8 +147,10 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter spatial::SpatBiasAddPlanOp, spatial::SpatAddPlanOp, spatial::SpatConcatPlanOp, + spatial::SpatFlattenPlanOp, spatial::SpatReluPlanOp, spatial::SpatSiluPlanOp, + spatial::SpatResizeNearestPlanOp, spatial::SpatMaxPool2DPlanOp, spatial::SpatGlobalAveragePoolPlanOp, spatial::SpatBlueprintOp, diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp similarity index 100% rename from src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp rename to src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp diff --git a/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/SpatialLayoutCapabilities.cpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/SpatialLayoutCapabilities.cpp new file mode 100644 index 0000000..69a6390 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/SpatialLayoutCapabilities.cpp @@ -0,0 +1,152 @@ +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +using namespace mlir; + +namespace onnx_mlir::spatial { + +static LayoutAlternative denseAlternative(Operation *op) { + LayoutAlternative alternative; + alternative.operandLayouts.assign(op->getNumOperands(), PhysicalLayout::DenseNCHW); + alternative.resultLayout = PhysicalLayout::DenseNCHW; + return alternative; +} + +static LayoutAlternative rowStripAlternative(Operation *op, + ArrayRef operandLayouts) { + LayoutAlternative alternative; + alternative.operandLayouts.assign(operandLayouts.begin(), operandLayouts.end()); + alternative.resultLayout = PhysicalLayout::NHWCRowStrip; + alternative.intrinsicCost = -2; + return alternative; +} + +static bool hasRowStripInput(ArrayRef operandLayouts, unsigned index) { + return index < operandLayouts.size() + && operandLayouts[index] == PhysicalLayout::NHWCRowStrip; +} + +SmallVector SpatConv2DPlanOp::getLayoutAlternatives( + const SpatialTargetResources& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (hasRowStripInput(operandLayouts, 0)) { + if (succeeded(canConsumeAndProduceRowStrip(*this, target))) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + } + else if (succeeded(canLowerConvPlanToRowStrip(*this, target))) { + LayoutAlternative alternative = denseAlternative(getOperation()); + alternative.resultLayout = PhysicalLayout::NHWCRowStrip; + alternative.intrinsicCost = -2; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatFlattenPlanOp::getLayoutAlternatives( + const SpatialTargetResources& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (!operandLayouts.empty() + && operandLayouts[0] == PhysicalLayout::Fragmented) { + LayoutAlternative alternative = denseAlternative(getOperation()); + alternative.operandLayouts[0] = PhysicalLayout::Fragmented; + alternatives.push_back(std::move(alternative)); + } + if (hasRowStripInput(operandLayouts, 0) + && succeeded(canLowerFlattenFromRowStrip(*this, target))) { + LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts); + alternative.resultLayout = PhysicalLayout::DenseNCHW; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatReluPlanOp::getLayoutAlternatives( + const SpatialTargetResources&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (hasRowStripInput(operandLayouts, 0)) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + return alternatives; +} + +SmallVector SpatSiluPlanOp::getLayoutAlternatives( + const SpatialTargetResources&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (hasRowStripInput(operandLayouts, 0)) { + LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts); + alternative.intrinsicCost = -3; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatResizeNearestPlanOp::getLayoutAlternatives( + const SpatialTargetResources& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (hasRowStripInput(operandLayouts, 0) + && succeeded(canLowerResizeNearestPlanToRowStrip(*this, target))) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + return alternatives; +} + +SmallVector SpatMaxPool2DPlanOp::getLayoutAlternatives( + const SpatialTargetResources& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) { + LayoutAlternative alternative = denseAlternative(getOperation()); + if (hasRowStripInput(operandLayouts, 0)) + alternative = rowStripAlternative(getOperation(), operandLayouts); + alternative.resultLayout = PhysicalLayout::NHWCRowStrip; + alternative.intrinsicCost = -2; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatGlobalAveragePoolPlanOp::getLayoutAlternatives( + const SpatialTargetResources& target, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) { + LayoutAlternative alternative = denseAlternative(getOperation()); + if (hasRowStripInput(operandLayouts, 0)) + alternative = rowStripAlternative(getOperation(), operandLayouts); + alternative.resultLayout = PhysicalLayout::NHWCRowStrip; + alternative.intrinsicCost = -2; + alternatives.push_back(std::move(alternative)); + } + return alternatives; +} + +SmallVector SpatBiasAddPlanOp::getLayoutAlternatives( + const SpatialTargetResources&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + auto resultType = dyn_cast(getOutput().getType()); + if (resultType && hasRowStripInput(operandLayouts, 0) + && isSupportedBiasAddValue(getBias(), resultType)) + alternatives.push_back(rowStripAlternative(getOperation(), + {PhysicalLayout::NHWCRowStrip, + PhysicalLayout::DenseNCHW})); + return alternatives; +} + +SmallVector SpatAddPlanOp::getLayoutAlternatives( + const SpatialTargetResources&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0) + && hasRowStripInput(operandLayouts, 1)) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + return alternatives; +} + +SmallVector SpatConcatPlanOp::getLayoutAlternatives( + const SpatialTargetResources&, ArrayRef operandLayouts) { + SmallVector alternatives {denseAlternative(getOperation())}; + if (!operandLayouts.empty() && llvm::all_of(operandLayouts, [](PhysicalLayout layout) { + return layout == PhysicalLayout::NHWCRowStrip; + })) + alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts)); + return alternatives; +} + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/LowerSpatialPlansPass.cpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/LowerSpatialPlansPass.cpp new file mode 100644 index 0000000..89ed83d --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/LowerSpatialPlansPass.cpp @@ -0,0 +1,136 @@ +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp" +#include "src/Accelerators/PIM/Passes/PIMPasses.h" +#include "src/Dialect/ONNX/ONNXOps.hpp" + +using namespace mlir; + +namespace onnx_mlir { +namespace { + +struct LowerSpatialPlansPass final + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass) + + StringRef getArgument() const override { return "lower-spatial-plans"; } + StringRef getDescription() const override { + return "Lower selected Spatial planning ops to low-level Spatial IR."; + } + + LowerSpatialPlansPass() = default; + LowerSpatialPlansPass(const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options, + spatial::SpatialDataflowExportStage exportStage) + : target(target), planningOptions(options), exportStage(exportStage), hasTarget(true) {} + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + if (!hasTarget) { + moduleOp.emitError("Spatial plan lowering requires an injected SpatialTargetResources"); + signalPassFailure(); + return; + } + auto entryFunc = getPimEntryFunc(moduleOp); + if (failed(entryFunc)) { + moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans"); + signalPassFailure(); + return; + } + func::FuncOp funcOp = *entryFunc; + auto verifyLogicalPhase = [&](StringRef stage) -> bool { + if (succeeded(verifyLogicalSpatialGraphInvariants(funcOp))) + return true; + moduleOp.emitError() << "logical Spatial graph verification failed " << stage; + signalPassFailure(); + return false; + }; + + if (!verifyLogicalPhase("at the start of LowerSpatialPlans")) + return; + if (failed(verifySelectedSpatialLayouts(funcOp, target))) { + moduleOp.emitError("selected Spatial layout verification failed"); + signalPassFailure(); + return; + } + + MLIRContext* ctx = moduleOp.getContext(); + RewritePatternSet patterns(ctx); + populateSpatialPlanLoweringPatterns(patterns, ctx, target, planningOptions); + + ConversionTarget conversionTarget(*ctx); + conversionTarget.addLegalDialect(); + conversionTarget.addIllegalDialect(); + conversionTarget.addLegalOp(); + conversionTarget.addIllegalOp(); + conversionTarget.addDynamicallyLegalOp( + [](spatial::SpatBlueprintOp blueprint) { + return spatial::isFragmentAssembly(blueprint.getMode()); + }); + + if (failed(applyFullConversion(funcOp, conversionTarget, + std::move(patterns)))) { + moduleOp.emitError("failed to lower Spatial plans and layout materialization"); + signalPassFailure(); + return; + } + + dumpModule(moduleOp, "spatial1_graph"); + if (spatial::shouldExportSpatialDataflowStage( + exportStage, spatial::SpatialDataflowExportStage::Spatial1) + && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) { + signalPassFailure(); + return; + } + verifyLogicalPhase("at the end of LowerSpatialPlans"); + } + + spatial::SpatialTargetResources target; + ONNXToSpatialPlanningOptions planningOptions; + spatial::SpatialDataflowExportStage exportStage = spatial::SpatialDataflowExportStage::None; + bool hasTarget = false; +}; + +} // namespace + +std::unique_ptr createLowerSpatialPlansPass() { + return std::make_unique(); +} + +std::unique_ptr createLowerSpatialPlansPass( + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options, + spatial::SpatialDataflowExportStage exportStage) { + return std::make_unique(target, options, exportStage); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/ONNXToSpatialPass.cpp similarity index 85% rename from src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp rename to src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/ONNXToSpatialPass.cpp index 1e25d72..41eb330 100644 --- a/src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/ONNXToSpatialPass.cpp @@ -12,15 +12,15 @@ #include "llvm/ADT/SmallVector.h" -#include "Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" #include "Common/PimCommon.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" -#include "ONNXToSpatialVerifier.hpp" using namespace mlir; @@ -34,9 +34,17 @@ struct ONNXToSpatialPass : PassWrapper biasAddPlans(funcOp.getOps()); SmallVector addPlans(funcOp.getOps()); SmallVector concatPlans(funcOp.getOps()); + SmallVector flattenPlans(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() || !flattenPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !resizePlans.empty() + || !maxPoolPlans.empty() || !blueprints.empty() || !globalAveragePoolPlans.empty() || !materializers.empty()) { return; } @@ -103,6 +115,11 @@ static void populateEmptyFunction(func::FuncOp funcOp) { void ONNXToSpatialPass::runOnOperation() { ModuleOp moduleOp = getOperation(); + if (!hasTarget) { + moduleOp.emitError("ONNX-to-Spatial lowering requires an injected SpatialTargetResources"); + signalPassFailure(); + return; + } MLIRContext* ctx = &getContext(); ConversionTarget preTarget(*ctx); @@ -123,6 +140,14 @@ void ONNXToSpatialPass::runOnOperation() { return; } + RewritePatternSet matmulPatterns(ctx); + populateMatMulFusionPatterns(matmulPatterns, ctx, target); + 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)))) { @@ -171,7 +196,7 @@ void ONNXToSpatialPass::runOnOperation() { target.addIllegalOp(); RewritePatternSet conversionPatterns(ctx); - populateConversionPatterns(conversionPatterns, ctx); + populateConversionPatterns(conversionPatterns, ctx, this->target, planningOptions); if (failed(applyPartialConversion(moduleOp, target, std::move(conversionPatterns)))) { moduleOp.emitError("failed to convert required ONNX ops to Spatial ops"); signalPassFailure(); @@ -247,4 +272,10 @@ void ONNXToSpatialPass::runOnOperation() { std::unique_ptr createONNXToSpatialPass() { return std::make_unique(); } +std::unique_ptr createONNXToSpatialPass( + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) { + return std::make_unique(target, options); +} + } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp new file mode 100644 index 0000000..c4677aa --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp @@ -0,0 +1,92 @@ +#pragma once + +#include + +#include "mlir/IR/PatternMatch.h" +#include "mlir/Support/LogicalResult.h" + +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + +namespace onnx_mlir { + +struct RowStripPhysicalValue; +struct ONNXToSpatialPlanningOptions; + +inline spatial::PhysicalLayout getSpatialPlanOperandLayout(mlir::Value value) { + if (auto materialize = value.getDefiningOp()) + return materialize.getTargetPhysicalLayout(); + if (auto blueprint = value.getDefiningOp()) + return blueprint.getPhysicalLayout(); + if (mlir::Operation* producer = value.getDefiningOp()) + if (auto selected = spatial::getSelectedPhysicalLayout(producer)) + return *selected; + return spatial::PhysicalLayout::DenseNCHW; +} + +mlir::FailureOr +lowerDenseFlattenPlan(spatial::SpatFlattenPlanOp planOp, + mlir::Value input, + mlir::PatternRewriter& rewriter); + +mlir::FailureOr +lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, + mlir::Value input, + mlir::Value weight, + mlir::Value bias, + std::optional rowStripInput, + bool emitRowStripLayout, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options, + mlir::PatternRewriter& rewriter); + +mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions* options = nullptr); +mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions* options = nullptr); + +mlir::LogicalResult canLowerResizeNearestPlanToRowStrip( + spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetResources& target); + +mlir::FailureOr lowerSelectedResizeNearestPlan( + spatial::SpatResizeNearestPlanOp planOp, + mlir::Value input, + std::optional rowStripInput, + const spatial::SpatialTargetResources& target, + mlir::PatternRewriter& rewriter); + +mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp, + const spatial::SpatialTargetResources& target); + +mlir::FailureOr +lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, + mlir::Value input, + const spatial::SpatialTargetResources& target, + mlir::PatternRewriter& rewriter); + +mlir::FailureOr +lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, + mlir::Value input, + std::optional rowStripInput, + const spatial::SpatialTargetResources& target, + mlir::PatternRewriter& rewriter); + +mlir::LogicalResult +canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp, + const spatial::SpatialTargetResources& target); + +mlir::FailureOr +lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp, + mlir::Value input, + const spatial::SpatialTargetResources& target, + mlir::PatternRewriter& rewriter); + +mlir::FailureOr +lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp, + mlir::Value input, + std::optional rowStripInput, + const spatial::SpatialTargetResources& target, + mlir::PatternRewriter& rewriter); + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialLayoutPlanningPass.cpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialLayoutPlanningPass.cpp new file mode 100644 index 0000000..7099344 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialLayoutPlanningPass.cpp @@ -0,0 +1,318 @@ +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/DenseMap.h" + +#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Passes/PIMPasses.h" + +#include +#include + +using namespace mlir; + +namespace onnx_mlir { +namespace { + +struct SpatialLayoutSelection { + llvm::DenseMap selectedAlternative; + llvm::DenseMap resultLayouts; +}; + +static spatial::PhysicalLayout getKnownLayout( + const SpatialLayoutSelection& selection, Value value) { + if (auto it = selection.resultLayouts.find(value); it != selection.resultLayouts.end()) + return it->second; + return getSpatialPlanOperandLayout(value); +} + +static SmallVector getOperandLayouts( + Operation* op, const SpatialLayoutSelection& selection) { + SmallVector operandLayouts; + operandLayouts.reserve(op->getNumOperands()); + for (Value operand : op->getOperands()) + operandLayouts.push_back(getKnownLayout(selection, operand)); + return operandLayouts; +} + +class SpatialLayoutAnalysis { +public: + SpatialLayoutAnalysis(func::FuncOp funcOp, + const spatial::SpatialTargetResources& target) + : funcOp(funcOp), target(target) {} + + FailureOr run() { + SpatialLayoutSelection selection; + SmallVector planOps; + for (Operation& op : funcOp.getBody().front()) { + if (!isa(&op)) + continue; + planOps.push_back(&op); + selection.resultLayouts[op.getResult(0)] = spatial::PhysicalLayout::DenseNCHW; + selection.selectedAlternative[&op] = 0; + } + + const size_t maxRounds = 2 * planOps.size() + 1; + for (size_t round = 0; round < maxRounds; ++round) { + bool changed = false; + SmallVector order(planOps); + if (round % 2) + std::reverse(order.begin(), order.end()); + + for (Operation* op : order) { + FailureOr> alternatives = + getAlternatives(op, selection); + if (failed(alternatives)) + return failure(); + + unsigned currentIndex = selection.selectedAlternative.lookup(op); + if (currentIndex >= alternatives->size()) + currentIndex = 0; + if (selection.selectedAlternative.lookup(op) != currentIndex) { + selection.selectedAlternative[op] = currentIndex; + changed = true; + } + Value result = op->getResult(0); + if (selection.resultLayouts.lookup(result) != + (*alternatives)[currentIndex].resultLayout) { + selection.resultLayouts[result] = (*alternatives)[currentIndex].resultLayout; + changed = true; + } + int64_t bestCost = alternativeCost(op, (*alternatives)[currentIndex], selection); + unsigned bestIndex = currentIndex; + for (auto [index, alternative] : llvm::enumerate(*alternatives)) { + int64_t cost = alternativeCost(op, alternative, selection); + if (cost < bestCost) { + bestCost = cost; + bestIndex = index; + } + } + if (bestIndex == currentIndex) + continue; + selection.selectedAlternative[op] = bestIndex; + selection.resultLayouts[result] = (*alternatives)[bestIndex].resultLayout; + changed = true; + } + if (!changed) + return selection; + } + funcOp.emitError("Spatial layout selection did not converge within its bounded iteration budget"); + return failure(); + } + + FailureOr> getAlternatives( + Operation* op, const SpatialLayoutSelection& selection) { + auto capability = dyn_cast(op); + if (!capability) + return failure(); + SmallVector alternatives = + capability.getLayoutAlternatives(target, getOperandLayouts(op, selection)); + if (alternatives.empty()) + return op->emitOpError("does not advertise a legal Spatial layout alternative"), failure(); + for (const spatial::LayoutAlternative& alternative : alternatives) { + if (alternative.operandLayouts.size() != op->getNumOperands()) + return op->emitOpError("advertises a layout alternative with the wrong operand count"), failure(); + } + if (llvm::any_of(op->getResult(0).getUses(), [](OpOperand& use) { + return isa(use.getOwner()); + }) + && llvm::none_of(alternatives, [](const spatial::LayoutAlternative& alternative) { + return alternative.resultLayout == spatial::PhysicalLayout::DenseNCHW; + })) + return op->emitOpError("does not provide the required DenseNCHW function-result layout"), failure(); + return alternatives; + } + +private: + int64_t alternativeCost(Operation* op, + const spatial::LayoutAlternative& alternative, + const SpatialLayoutSelection& selection) { + if (llvm::any_of(op->getResult(0).getUses(), [](OpOperand& use) { + return isa(use.getOwner()); + }) + && alternative.resultLayout != spatial::PhysicalLayout::DenseNCHW) + return std::numeric_limits::max() / 4; + + int64_t cost = alternative.intrinsicCost; + SmallVector operandLayouts = getOperandLayouts(op, selection); + for (auto [actual, required] : llvm::zip(operandLayouts, alternative.operandLayouts)) + cost += actual != required; + + Value result = op->getResult(0); + for (OpOperand& use : result.getUses()) { + auto user = dyn_cast(use.getOwner()); + if (!user) + continue; + SmallVector userOperandLayouts = + getOperandLayouts(use.getOwner(), selection); + for (auto [index, operand] : llvm::enumerate(use.getOwner()->getOperands())) + if (operand == result) + userOperandLayouts[index] = alternative.resultLayout; + SmallVector userAlternatives = + user.getLayoutAlternatives(target, userOperandLayouts); + if (llvm::none_of(userAlternatives, + [&](const spatial::LayoutAlternative& userAlternative) { + return userAlternative.operandLayouts.size() + == use.getOwner()->getNumOperands() + && userAlternative.operandLayouts[use.getOperandNumber()] + == alternative.resultLayout; + })) + ++cost; + } + return cost; + } + + func::FuncOp funcOp; + const spatial::SpatialTargetResources& target; +}; + +static LogicalResult materializeMismatchedUses( + IRRewriter& rewriter, const SpatialLayoutSelection& selection, + Operation* op, SpatialLayoutAnalysis& analysis) { + Value value = op->getResult(0); + spatial::PhysicalLayout sourceLayout = getKnownLayout(selection, value); + SmallVector> mismatches; + for (OpOperand& use : value.getUses()) { + Operation* userOp = use.getOwner(); + auto capability = dyn_cast(userOp); + if (!capability) { + if (isa(userOp) || sourceLayout == spatial::PhysicalLayout::DenseNCHW) + continue; + mismatches.push_back({&use, spatial::PhysicalLayout::DenseNCHW}); + continue; + } + FailureOr> alternatives = + analysis.getAlternatives(userOp, selection); + if (failed(alternatives)) + return failure(); + unsigned selectedIndex = selection.selectedAlternative.lookup(userOp); + if (selectedIndex >= alternatives->size()) + return userOp->emitOpError() + << "has no selected Spatial layout alternative (index " << selectedIndex + << ", alternatives " << alternatives->size() << ")", + failure(); + spatial::PhysicalLayout required = + (*alternatives)[selectedIndex].operandLayouts[use.getOperandNumber()]; + if (required != sourceLayout) + mismatches.push_back({&use, required}); + } + + for (auto [use, required] : mismatches) { + Operation* userOp = use->getOwner(); + rewriter.setInsertionPoint(userOp); + auto materialized = spatial::SpatMaterializeLayoutOp::create( + rewriter, userOp->getLoc(), use->get().getType(), use->get(), + spatial::LogicalLayoutAttr::get( + rewriter.getContext(), spatial::LogicalLayout::NCHW), + spatial::PhysicalLayoutAttr::get(rewriter.getContext(), sourceLayout), + spatial::PhysicalLayoutAttr::get(rewriter.getContext(), required)); + use->set(materialized.getResult()); + } + return success(); +} + +static LogicalResult verifySelectedLayouts( + const SpatialLayoutSelection& selection, + ArrayRef planOps, + SpatialLayoutAnalysis& analysis) { + for (Operation* op : planOps) { + auto selected = spatial::getSelectedPhysicalLayout(op); + if (!selected) + return op->emitOpError("requires a selected physical layout"), failure(); + FailureOr> alternatives = + analysis.getAlternatives(op, selection); + if (failed(alternatives)) + return failure(); + unsigned selectedIndex = selection.selectedAlternative.lookup(op); + if (selectedIndex >= alternatives->size()) + return op->emitOpError() + << "has no selected Spatial layout alternative (index " << selectedIndex + << ", alternatives " << alternatives->size() << ")", + failure(); + const spatial::LayoutAlternative& alternative = (*alternatives)[selectedIndex]; + if (alternative.resultLayout != *selected + || getOperandLayouts(op, selection) != alternative.operandLayouts) + return op->emitOpError("selected physical layout does not satisfy its exact layout contract"), failure(); + } + return success(); +} + +struct SpatialLayoutPlanningPass final + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass) + + StringRef getArgument() const override { return "spatial-layout-planning"; } + StringRef getDescription() const override { + return "Select Spatial layout alternatives and insert explicit reconciliation barriers."; + } + + SpatialLayoutPlanningPass() = default; + explicit SpatialLayoutPlanningPass(const spatial::SpatialTargetResources& target) + : target(target), hasTarget(true) {} + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + if (!hasTarget) { + moduleOp.emitError("Spatial layout planning requires an injected SpatialTargetResources"); + signalPassFailure(); + return; + } + auto entryFunc = getPimEntryFunc(moduleOp); + if (failed(entryFunc)) { + moduleOp.emitError("failed to locate the PIM entry function during Spatial layout planning"); + signalPassFailure(); + return; + } + + func::FuncOp funcOp = *entryFunc; + SpatialLayoutAnalysis analysis(funcOp, target); + FailureOr selection = analysis.run(); + if (failed(selection)) { + signalPassFailure(); + return; + } + + SmallVector planOps; + for (Operation& op : funcOp.getBody().front()) + if (isa(&op)) + planOps.push_back(&op); + + IRRewriter rewriter(&getContext()); + for (Operation* op : planOps) { + op->setAttr(spatial::kSelectedLayoutAttrName, + spatial::PhysicalLayoutAttr::get( + rewriter.getContext(), selection->resultLayouts.lookup(op->getResult(0)))); + if (failed(materializeMismatchedUses( + rewriter, *selection, op, analysis))) { + signalPassFailure(); + return; + } + } + if (failed(verifySelectedLayouts(*selection, planOps, analysis)) + || failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) { + moduleOp.emitError("Spatial layout planning verification failed"); + signalPassFailure(); + return; + } + } + + spatial::SpatialTargetResources target; + bool hasTarget = false; +}; + +} // namespace + +std::unique_ptr createSpatialLayoutPlanningPass() { + return std::make_unique(); +} + +std::unique_ptr createSpatialLayoutPlanningPass( + const spatial::SpatialTargetResources& target) { + return std::make_unique(target); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.cpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.cpp new file mode 100644 index 0000000..bbf14da --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.cpp @@ -0,0 +1,845 @@ +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp" +#include "src/Accelerators/PIM/Passes/PIMPasses.h" + +using namespace mlir; + +namespace onnx_mlir { +namespace { + +static FailureOr getRowStripValue(Value value) { + return getRowStripPhysicalValue(value); +} + +static FailureOr publishRowStripValue(Operation* planOp, + Value storage, + PatternRewriter& rewriter) { + auto logicalType = dyn_cast(planOp->getResult(0).getType()); + if (!logicalType) + return planOp->emitOpError("requires ranked logical output type"), failure(); + FailureOr value = describeRowStripPhysicalValue(storage, logicalType); + if (failed(value)) + return planOp->emitOpError("lowering produced invalid row-strip physical storage"), failure(); + FailureOr blueprint = createRowStripStorageBlueprint( + storage, logicalType, rewriter, planOp->getLoc()); + if (failed(blueprint)) + return planOp->emitOpError("failed to create row-strip storage Blueprint"), failure(); + rewriter.replaceOp(planOp, *blueprint); + return *blueprint; +} + +static bool isRowStripSelected(Operation* op) { + auto selected = spatial::getSelectedPhysicalLayout(op); + return selected && *selected == spatial::PhysicalLayout::NHWCRowStrip; +} + +static bool isDenseSelected(Operation* op) { + auto selected = spatial::getSelectedPhysicalLayout(op); + return selected && *selected == spatial::PhysicalLayout::DenseNCHW; +} + +static spatial::PhysicalLayout getKnownPhysicalLayout(Value value) { + return getSpatialPlanOperandLayout(value); +} + +static LogicalResult verifySelectedLayouts( + func::FuncOp funcOp, const spatial::SpatialTargetResources& target) { + LogicalResult result = success(); + funcOp.walk([&](Operation* op) { + auto capability = dyn_cast(op); + if (!capability) + return; + auto selected = spatial::getSelectedPhysicalLayout(op); + if (!selected) { + op->emitOpError("requires a selected physical layout from SpatialLayoutPlanning"); + result = failure(); + return; + } + if (*selected != spatial::PhysicalLayout::DenseNCHW + && *selected != spatial::PhysicalLayout::NHWCRowStrip) { + op->emitOpError("has an unsupported selected physical layout"); + result = failure(); + return; + } + SmallVector operandLayouts; + operandLayouts.reserve(op->getNumOperands()); + for (Value operand : op->getOperands()) + operandLayouts.push_back(getKnownPhysicalLayout(operand)); + auto alternatives = capability.getLayoutAlternatives(target, operandLayouts); + if (llvm::none_of(alternatives, [&](const spatial::LayoutAlternative& alternative) { + return alternative.resultLayout == *selected + && alternative.operandLayouts == operandLayouts; + })) { + op->emitOpError("selected physical layout is not lowerable for its explicit operand layouts"); + result = failure(); + } + }); + return result; +} + +static FailureOr +lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) { + return applyRowStripRelu(input, rewriter, planOp.getLoc()); +} + +static FailureOr +lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) { + return applyRowStripSilu(input, rewriter, planOp.getLoc()); +} + +static FailureOr lowerRowStripAdd(const RowStripPhysicalValue& lhs, + const RowStripPhysicalValue& rhs, + spatial::SpatAddPlanOp planOp, + PatternRewriter& rewriter) { + return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc()); +} + +static FailureOr lowerRowStripConcat(ArrayRef inputs, + spatial::SpatConcatPlanOp planOp, + PatternRewriter& rewriter) { + auto outputType = dyn_cast(planOp.getOutput().getType()); + if (!outputType) + return failure(); + return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc()); +} + +static FailureOr +materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) { + if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape()) + return failure(); + return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc); +} + +static FailureOr materializeDenseToRowStrip( + Value input, RankedTensorType logicalType, Location loc, PatternRewriter& rewriter) { + if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4 + || logicalType.getDimSize(0) != 1) + return failure(); + auto nhwcType = RankedTensorType::get( + {1, logicalType.getDimSize(2), logicalType.getDimSize(3), logicalType.getDimSize(1)}, + logicalType.getElementType(), logicalType.getEncoding()); + auto rowsType = RankedTensorType::get( + {logicalType.getDimSize(2) * logicalType.getDimSize(3), logicalType.getDimSize(1)}, + logicalType.getElementType(), logicalType.getEncoding()); + auto rowsCompute = createSpatCompute<1>( + rewriter, loc, rowsType, {}, input, [&](Value denseInput) { + Value nhwc = createLinalgTranspose( + denseInput, nhwcType, {0, 2, 3, 1}, rewriter, loc); + Value rows = tensor::CollapseShapeOp::create( + rewriter, loc, rowsType, nhwc, + SmallVector {{0, 1, 2}, {3}}); + spatial::SpatYieldOp::create(rewriter, loc, rows); + }); + Value rows = rowsCompute->getResult(0); + FailureOr storage = createRowStripStorageFromRows(rows, logicalType, rewriter, loc); + if (failed(storage)) + return failure(); + return createRowStripStorageBlueprint(*storage, logicalType, rewriter, loc); +} + +static FailureOr lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType, + PatternRewriter& rewriter, Location loc) { + auto producer = input.getDefiningOp(); + auto inputType = dyn_cast(input.getType()); + auto biasType = dyn_cast(bias.getType()); + if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape() + || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount() + || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount()) + return failure(); + auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount()); + auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount()); + if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank() + || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front() + || inputFragmentType->getRank() != outputFragmentType->getRank() + 1) + return failure(); + for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape())) + if (outputDim > inputDim) + return failure(); + + auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias}, + [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult { + FailureOr fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType); + if (failed(fragment)) + return failure(); + MixedSliceGeometry biasSlice; + for (int64_t dim : inputFragmentType->getShape()) { + biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0)); + biasSlice.sizes.push_back(rewriter.getIndexAttr(dim)); + biasSlice.strides.push_back(rewriter.getIndexAttr(1)); + } + Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice); + if (!biasFragment) + return failure(); + Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment); + MixedSliceGeometry outputSlice; + outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0)); + outputSlice.sizes.push_back(rewriter.getIndexAttr(1)); + outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1)); + for (int64_t dim : outputFragmentType->getShape()) + outputSlice.sizes.push_back(rewriter.getIndexAttr(dim)); + Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice); + if (!output) + return failure(); + publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane); + return success(); + }); + if (failed(batch)) + return failure(); + return batch->getResult(0); +} + +struct LowerDenseReluPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp, + spatial::SpatReluPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation()); + if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW) + return failure(); + + auto computeOp = createSpatCompute<1>( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, adaptor.getInput(), [&](Value x) { + auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult()); + }); + rewriter.replaceOp(planOp, computeOp.getResults()); + return success(); + } +}; + +struct LowerDenseSiluPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp, + spatial::SpatSiluPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation()); + if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW) + return failure(); + + auto computeOp = createSpatCompute<1>( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, adaptor.getInput(), [&](Value x) { + Value sigmoid = spatial::SpatSigmoidOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult(); + Value silu = spatial::SpatVMulOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult(); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu); + }); + rewriter.replaceOp(planOp, computeOp.getResults()); + return success(); + } +}; + +struct LowerDenseResizePlan final : OpConversionPattern { + explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target) + : OpConversionPattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp, + spatial::SpatResizeNearestPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + FailureOr lowered = lowerSelectedResizeNearestPlan( + planOp, adaptor.getInput(), std::nullopt, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense nearest Resize plan"); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + + const spatial::SpatialTargetResources& target; +}; + +struct LowerDenseBiasAddPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp, + spatial::SpatBiasAddPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + auto resultType = dyn_cast(planOp.getOutput().getType()); + if (!resultType) + return planOp.emitOpError("requires ranked output type"); + + FailureOr denseBias = materializeDenseBiasAddTensor( + adaptor.getBias(), resultType, rewriter, planOp.getLoc()); + if (failed(denseBias)) + return planOp.emitOpError("failed to materialize dense Conv-style bias"); + if (adaptor.getInput().getDefiningOp()) { + FailureOr lowered = lowerDenseBatchBiasAdd( + adaptor.getInput(), *denseBias, resultType, rewriter, planOp.getLoc()); + if (succeeded(lowered)) { + rewriter.replaceOp(planOp, *lowered); + return success(); + } + } + auto computeOp = createSpatCompute<2>( + rewriter, + planOp.getLoc(), + planOp.getOutput().getType(), + {}, + ValueRange {adaptor.getInput(), *denseBias}, + [&](Value x, Value y) { + auto added = spatial::SpatVAddOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult()); + }); + rewriter.replaceOp(planOp, computeOp.getResults()); + return success(); + } +}; + +struct LowerDenseAddPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp, + spatial::SpatAddPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + auto compute = createSpatCompute<2>( + rewriter, + planOp.getLoc(), + planOp.getOutput().getType(), + {}, + ValueRange {adaptor.getLhs(), adaptor.getRhs()}, + [&](Value lhsValue, Value rhsValue) { + Value added = spatial::SpatVAddOp::create( + rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added); + }); + rewriter.replaceOp(planOp, compute.getResults()); + return success(); + } +}; + +struct LowerDenseConcatPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp, + spatial::SpatConcatPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + auto compute = createSpatCompute( + rewriter, + planOp.getLoc(), + TypeRange {planOp.getOutput().getType()}, + {}, + adaptor.getInputs(), + [&](ValueRange values) { + Value concatenated = spatial::SpatConcatOp::create( + rewriter, + planOp.getLoc(), + planOp.getOutput().getType(), + rewriter.getI64IntegerAttr(planOp.getAxis()), + values); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated); + }); + rewriter.replaceOp(planOp, compute.getResults()); + return success(); + } +}; + +static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp, + Value lhsValue, Value rhsValue, + PatternRewriter& rewriter) { + FailureOr lhs = getRowStripValue(lhsValue); + FailureOr rhs = getRowStripValue(rhsValue); + if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) { + if (getKnownPhysicalLayout(lhsValue) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs"); + } + if (isRowStripSelected(planOp.getOperation()) && failed(rhs)) { + if (getKnownPhysicalLayout(rhsValue) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs"); + } + if (isRowStripSelected(planOp.getOperation())) { + rewriter.setInsertionPoint(planOp); + FailureOr lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial add plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + return planOp.emitOpError("dense Add plan was not lowered by the selected-plan patterns"); +} + +static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp, + ValueRange inputValues, + PatternRewriter& rewriter) { + SmallVector inputs; + for (Value input : inputValues) { + FailureOr physical = getRowStripValue(input); + if (failed(physical)) { + inputs.clear(); + break; + } + inputs.push_back(*physical); + } + if (isRowStripSelected(planOp.getOperation()) && inputs.size() != inputValues.size()) { + if (llvm::any_of(inputValues, [](Value input) { + return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip; + })) + return failure(); + return planOp.emitOpError("selected row-strip Concat plan requires row-strip inputs"); + } + if (isRowStripSelected(planOp.getOperation())) { + rewriter.setInsertionPoint(planOp); + FailureOr lowered = lowerRowStripConcat(inputs, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + return planOp.emitOpError("dense Concat plan was not lowered by the selected-plan patterns"); +} + +struct LowerSelectedConvPlan final : OpConversionPattern { + explicit LowerSelectedConvPlan(MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) + : OpConversionPattern(ctx), target(target), options(options) {} + + LogicalResult matchAndRewrite(spatial::SpatConv2DPlanOp planOp, + spatial::SpatConv2DPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (isDenseSelected(planOp.getOperation())) { + FailureOr lowered = lowerSelectedConv2DPlan( + planOp, adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(), + std::nullopt, /*emitRowStripLayout=*/false, target, options, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense Spatial Conv plan"); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + + FailureOr rowStripInput = getRowStripValue(adaptor.getInput()); + if (failed(rowStripInput) + && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + std::optional physicalInput; + if (succeeded(rowStripInput)) + physicalInput = rowStripInput->storage; + FailureOr lowered = lowerSelectedConv2DPlan( + planOp, adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(), + physicalInput, /*emitRowStripLayout=*/true, target, options, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + + const spatial::SpatialTargetResources& target; + const ONNXToSpatialPlanningOptions& options; +}; + +struct LowerRowStripReluPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp, + spatial::SpatReluPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(adaptor.getInput()); + if (failed(input)) { + if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input"); + } + FailureOr lowered = lowerRowStripRelu(*input, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial ReLU plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } +}; + +struct LowerRowStripSiluPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp, + spatial::SpatSiluPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(adaptor.getInput()); + if (failed(input)) { + if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input"); + } + FailureOr lowered = lowerRowStripSilu(*input, planOp, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } +}; + +struct LowerRowStripResizePlan final : OpConversionPattern { + explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target) + : OpConversionPattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp, + spatial::SpatResizeNearestPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(adaptor.getInput()); + if (failed(input)) { + if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip Resize plan requires a row-strip input"); + } + FailureOr lowered = lowerSelectedResizeNearestPlan( + planOp, adaptor.getInput(), input->storage, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Resize plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + + const spatial::SpatialTargetResources& target; +}; + +struct LowerDenseMaxPoolPlan final : OpConversionPattern { + explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target) + : OpConversionPattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp, + spatial::SpatMaxPool2DPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + FailureOr lowered = lowerDenseMaxPool2DPlan( + planOp, adaptor.getInput(), target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense Spatial MaxPool plan"); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + + const spatial::SpatialTargetResources& target; +}; + +struct LowerRowStripMaxPoolPlan final : OpConversionPattern { + explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target) + : OpConversionPattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp, + spatial::SpatMaxPool2DPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(adaptor.getInput()); + if (failed(input) + && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + std::optional physicalInput; + if (succeeded(input)) + physicalInput = input->storage; + FailureOr lowered = lowerSelectedMaxPool2DPlan( + planOp, adaptor.getInput(), physicalInput, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + + const spatial::SpatialTargetResources& target; +}; + +struct LowerRowStripGlobalAveragePoolPlan + final : OpConversionPattern { + explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target) + : OpConversionPattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp, + spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(adaptor.getInput()); + if (failed(input) + && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + std::optional physicalInput; + if (succeeded(input)) + physicalInput = input->storage; + FailureOr lowered = lowerSelectedGlobalAveragePoolPlan( + planOp, adaptor.getInput(), physicalInput, target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } + + const spatial::SpatialTargetResources& target; +}; + +struct LowerDenseGlobalAveragePoolPlan + final : OpConversionPattern { + explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx, + const spatial::SpatialTargetResources& target) + : OpConversionPattern(ctx), target(target) {} + + LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp, + spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + FailureOr lowered = lowerDenseGlobalAveragePoolPlan( + planOp, adaptor.getInput(), target, rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense Spatial global AveragePool plan"); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + + const spatial::SpatialTargetResources& target; +}; + +struct LowerRowStripBiasAddPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp, + spatial::SpatBiasAddPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + FailureOr input = getRowStripValue(adaptor.getInput()); + if (failed(input)) { + if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip) + return failure(); + return planOp.emitOpError("selected row-strip bias_add plan requires a row-strip input"); + } + FailureOr lowered = applyRowStripBiasAdd( + *input, adaptor.getBias(), rewriter, planOp.getLoc()); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan"); + if (failed(publishRowStripValue(planOp, *lowered, rewriter))) + return failure(); + return success(); + } +}; + +struct LowerRowStripAddPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp, + spatial::SpatAddPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + return lowerAddPlan(planOp, adaptor.getLhs(), adaptor.getRhs(), rewriter); + } +}; + +struct LowerRowStripConcatPlan final : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp, + spatial::SpatConcatPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isRowStripSelected(planOp.getOperation())) + return failure(); + return lowerConcatPlan(planOp, adaptor.getInputs(), rewriter); + } +}; + +struct LowerMaterializeLayout final + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp, + spatial::SpatMaterializeLayoutOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + auto source = materializeOp.getSourcePhysicalLayout(); + auto target = materializeOp.getTargetPhysicalLayout(); + if (source == spatial::PhysicalLayout::DenseNCHW + && target == spatial::PhysicalLayout::DenseNCHW) { + rewriter.replaceOp(materializeOp, adaptor.getInput()); + return success(); + } + if (source == spatial::PhysicalLayout::DenseNCHW + && target == spatial::PhysicalLayout::NHWCRowStrip) { + auto logicalType = dyn_cast(adaptor.getInput().getType()); + if (!logicalType) + return materializeOp.emitOpError("requires a ranked dense input"), failure(); + FailureOr rowStrip = materializeDenseToRowStrip( + adaptor.getInput(), logicalType, materializeOp.getLoc(), rewriter); + if (failed(rowStrip)) + return materializeOp.emitOpError( + "failed to materialize dense NCHW storage to row-strip layout"), failure(); + rewriter.replaceOp(materializeOp, *rowStrip); + return success(); + } + if (source != spatial::PhysicalLayout::NHWCRowStrip + || target != spatial::PhysicalLayout::DenseNCHW) + return materializeOp.emitOpError( + "unsupported Spatial layout materialization direction"), failure(); + auto inputType = dyn_cast(adaptor.getInput().getType()); + if (!inputType) + return materializeOp.emitOpError("requires a ranked row-strip input"), failure(); + FailureOr rowStripValue = + getRowStripValue(adaptor.getInput()); + if (failed(rowStripValue)) + return failure(); + FailureOr dense = materializeRowStripToDense( + *rowStripValue, materializeOp.getLoc(), rewriter); + if (failed(dense)) + return materializeOp.emitOpError( + "failed to materialize row-strip storage to dense NCHW"), failure(); + rewriter.replaceOp(materializeOp, *dense); + return success(); + } +}; + +struct LowerSelectedFlattenPlan final + : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(spatial::SpatFlattenPlanOp planOp, + spatial::SpatFlattenPlanOpAdaptor adaptor, + ConversionPatternRewriter& rewriter) const override { + if (!isDenseSelected(planOp.getOperation())) + return failure(); + FailureOr rowStripInput = getRowStripValue(adaptor.getInput()); + if (succeeded(rowStripInput)) { + if (failed(canLowerFlattenFromRowStrip(planOp, target)) + || failed(lowerFlattenFromRowStrip(*rowStripInput, planOp, target, rewriter))) + return planOp.emitOpError("failed to lower selected Spatial Flatten plan"), failure(); + return success(); + } + FailureOr lowered = lowerDenseFlattenPlan(planOp, adaptor.getInput(), rewriter); + if (failed(lowered)) + return planOp.emitOpError("failed to lower selected dense Spatial Flatten plan"), failure(); + rewriter.replaceOp(planOp, *lowered); + return success(); + } + + explicit LowerSelectedFlattenPlan(MLIRContext* context, + const spatial::SpatialTargetResources& target) + : OpConversionPattern(context), target(target) {} + + const spatial::SpatialTargetResources& target; +}; + +struct EraseDeadPhysicalViewBlueprint final + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(spatial::SpatBlueprintOp blueprint, + PatternRewriter& rewriter) const override { + if (!spatial::isPhysicalView(blueprint.getMode()) || !blueprint.use_empty()) + return failure(); + rewriter.eraseOp(blueprint); + return success(); + } +}; + +static void populateConvPlanLoweringPatterns( + RewritePatternSet& patterns, MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) { + patterns.add(ctx, target, options); +} + +static void populateElementwisePlanLoweringPatterns( + RewritePatternSet& patterns, MLIRContext* ctx) { + patterns.add(ctx); +} + +static void populatePoolPlanLoweringPatterns( + RewritePatternSet& patterns, MLIRContext* ctx, + const spatial::SpatialTargetResources& target) { + patterns.add(ctx, target); +} + +static void populateResizePlanLoweringPatterns( + RewritePatternSet& patterns, MLIRContext* ctx, + const spatial::SpatialTargetResources& target) { + patterns.add(ctx, target); +} + +static void populateConcatPlanLoweringPatterns( + RewritePatternSet& patterns, MLIRContext* ctx) { + patterns.add(ctx); +} + +static void populateFlattenPlanLoweringPatterns( + RewritePatternSet& patterns, MLIRContext* ctx, + const spatial::SpatialTargetResources& target) { + patterns.add(ctx, target); +} + +static void populateLayoutMaterializationPatterns( + RewritePatternSet& patterns, MLIRContext* ctx) { + patterns.add(ctx); +} + +} // namespace + +void populateSpatialPlanLoweringPatterns( + RewritePatternSet& patterns, MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) { + populateConvPlanLoweringPatterns(patterns, ctx, target, options); + populateElementwisePlanLoweringPatterns(patterns, ctx); + populatePoolPlanLoweringPatterns(patterns, ctx, target); + populateResizePlanLoweringPatterns(patterns, ctx, target); + populateConcatPlanLoweringPatterns(patterns, ctx); + populateFlattenPlanLoweringPatterns(patterns, ctx, target); + populateLayoutMaterializationPatterns(patterns, ctx); +} + +LogicalResult verifySelectedSpatialLayouts( + func::FuncOp funcOp, const spatial::SpatialTargetResources& target) { + return verifySelectedLayouts(funcOp, target); +} + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp new file mode 100644 index 0000000..91d0009 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Support/LogicalResult.h" + +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" + +namespace onnx_mlir { + +void populateSpatialPlanLoweringPatterns( + mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options); + +mlir::LogicalResult verifySelectedSpatialLayouts( + mlir::func::FuncOp funcOp, const spatial::SpatialTargetResources& target); + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns.cpp index b2106fd..b6aa6e8 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns.cpp @@ -7,12 +7,15 @@ namespace onnx_mlir { void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { populateGeneratedPrePatterns(patterns, ctx); } -void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { +void populateConversionPatterns(RewritePatternSet& patterns, + MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) { populateElementwisePatterns(patterns, ctx); - populateMatMulRewritePatterns(patterns, ctx); - populateGemmPatterns(patterns, ctx); - populateConvPatterns(patterns, ctx); - populatePoolPatterns(patterns, ctx); + populateMatMulRewritePatterns(patterns, ctx, target); + populateGemmPatterns(patterns, ctx, target); + populateConvPatterns(patterns, ctx, target, options); + populatePoolPatterns(patterns, ctx, target); populateReduceMeanPatterns(patterns, ctx); populateReluPatterns(patterns, ctx); populateSigmoidPatterns(patterns, ctx); diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp b/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp index e3cec73..bc65a60 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns.hpp @@ -4,23 +4,43 @@ #include "mlir/IR/MLIRContext.h" #include "mlir/Transforms/DialectConversion.h" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" namespace onnx_mlir { +namespace spatial { +struct SpatialTargetResources; +} + void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); -void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); +void populateConversionPatterns(mlir::RewritePatternSet& patterns, + mlir::MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options); void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); -void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); +void populateConvPatterns(mlir::RewritePatternSet& patterns, + mlir::MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options); void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); 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 populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); +void populateGemmPatterns(mlir::RewritePatternSet& patterns, + mlir::MLIRContext* ctx, + const spatial::SpatialTargetResources& target); +void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns, + mlir::MLIRContext* ctx, + const spatial::SpatialTargetResources& target); +void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns, + mlir::MLIRContext* ctx, + const spatial::SpatialTargetResources& target); +void populatePoolPatterns(mlir::RewritePatternSet& patterns, + mlir::MLIRContext* ctx, + const spatial::SpatialTargetResources& target); void populateReduceMeanPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx); void populateSigmoidPatterns(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..39677aa 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp @@ -1,5 +1,4 @@ #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" @@ -7,28 +6,27 @@ #include "mlir/IR/BuiltinTypes.h" #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallPtrSet.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" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -39,19 +37,17 @@ namespace onnx_mlir { namespace { struct ConvToGemm : OpConversionPattern { - using OpConversionPattern::OpConversionPattern; + explicit ConvToGemm(MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) + : OpConversionPattern(ctx), target(target), options(options) {} LogicalResult matchAndRewrite(ONNXConvOp convOp, ONNXConvOpAdaptor convOpAdaptor, ConversionPatternRewriter& rewriter) const override; -}; -struct ConvLoweringDecision { - PimConvLoweringType strategy; - std::string reason; - bool isAuto = false; - std::string fallbackReason; - std::string rejectedAutoStrategy; + const spatial::SpatialTargetResources& target; + const ONNXToSpatialPlanningOptions& options; }; struct PreparedConvInput { @@ -59,1019 +55,207 @@ struct PreparedConvInput { 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) { +static StringRef stringifyConvLoweringStrategy(spatial::ConvLoweringStrategy strategy) { switch (strategy) { - case PimConvLoweringAuto: return "auto"; - case PimConvLoweringLegacy: return "legacy"; - case PimConvLoweringDepthwise: return "depthwise"; - case PimConvLoweringPackedIm2Col: return "packed-im2col"; - case PimConvLoweringStreamedPatch: return "streamed-patch"; - case PimConvLoweringStreamedPacked: return "streamed-packed"; - case PimConvLoweringOutputChannelTiled: return "output-channel-tiled"; - case PimConvLoweringInputKTiled: return "input-k-tiled"; - case PimConvLoweringTiled2D: return "tiled-2d"; + case spatial::ConvLoweringStrategy::Auto: return "auto"; + case spatial::ConvLoweringStrategy::Legacy: return "legacy"; + case spatial::ConvLoweringStrategy::Depthwise: return "depthwise"; + case spatial::ConvLoweringStrategy::PackedIm2Col: return "packed-im2col"; + case spatial::ConvLoweringStrategy::StreamedPatch: return "streamed-patch"; + case spatial::ConvLoweringStrategy::StreamedPacked: return "streamed-packed"; + case spatial::ConvLoweringStrategy::OutputChannelTiled: return "output-channel-tiled"; + case spatial::ConvLoweringStrategy::InputKTiled: return "input-k-tiled"; + case spatial::ConvLoweringStrategy::Tiled2D: return "tiled-2d"; } llvm_unreachable("unknown conv lowering strategy"); } -static bool requiresFuncReturnMaterialization(const DistributedConvAnalysis& analysis) { - return !analysis.hasLocalConsumers() && analysis.barrierKind == DistributedConvBarrierKind::Return; +enum class ConvLoweringReportPhase { Planning, Realization }; + +struct ConvLoweringReportEntry { + size_t convId; + std::string phase; + std::string location; + std::string strategy; + std::string implementation; +}; + +struct ConvLoweringReportState { + std::mutex mutex; + llvm::SmallPtrSet planned; + llvm::SmallPtrSet realized; + llvm::DenseMap convIds; + llvm::SmallVector entries; + size_t nextConvId = 1; +}; + +static StringRef stringifyConvLoweringReportPhase(ConvLoweringReportPhase phase) { + return phase == ConvLoweringReportPhase::Planning ? "planning" : "realization"; } -static ConvStrategyEstimate estimateConvStrategy(const ConvGeometry& geo, - PimConvLoweringType strategy, - const DistributedConvAnalysis& analysis) { - ConvStrategyEstimate estimate; - estimate.requiresFuncReturnMaterialization = requiresFuncReturnMaterialization(analysis); - estimate.materializationKind = estimate.requiresFuncReturnMaterialization ? "func.return" : "none"; +static std::string convReportLocation(Operation* op) { + std::string location; + llvm::raw_string_ostream stream(location); + op->getLoc().print(stream); + if (location.size() > 120) + location.replace(117, std::string::npos, "..."); + return location; +} - 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; +static StringRef stringifyConvMaterializationKind(ConvMaterializationKind kind) { + switch (kind) { + case ConvMaterializationKind::StructuredDepthwise: + return "structured-depthwise"; + case ConvMaterializationKind::PackedIm2Col: + return "packed-im2col"; + case ConvMaterializationKind::StreamedPatch: + return "streamed-patch"; + case ConvMaterializationKind::StreamedPacked: + return "streamed-packed"; + case ConvMaterializationKind::InputKTiled: + return "input-k-tiled"; } - 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; + llvm_unreachable("unknown Conv materialization kind"); +} + +static StringRef convLoweringImplementation(ConvMaterializationKind kind) { + switch (kind) { + case ConvMaterializationKind::StructuredDepthwise: + return "DW"; + case ConvMaterializationKind::PackedIm2Col: + return "PIC"; + case ConvMaterializationKind::StreamedPatch: + return "STR"; + case ConvMaterializationKind::InputKTiled: + return "IKT"; + case ConvMaterializationKind::StreamedPacked: + return "STP"; } - 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, ' '); + llvm_unreachable("unknown conv lowering implementation"); +} + +static StringRef convRowStripInputImplementation(const ConvLoweringState& state, + ConvMaterializationKind kind) { + if (kind == ConvMaterializationKind::StructuredDepthwise) + return "RSDW"; + if (state.problem.xHeight == 1 && state.problem.xWidth == 1 && state.problem.wHeight == 1 && state.problem.wWidth == 1) + return "RSP"; + return "RSM"; +} + +static constexpr size_t kConvReportIdWidth = 4; +static constexpr size_t kConvReportLocationWidth = 24; +static constexpr size_t kConvReportStrategyWidth = 20; +static constexpr size_t kConvReportCodeWidth = 8; + +static std::string convReportCell(StringRef value, size_t width) { + std::string cell = value.str(); + if (cell.size() > width) { + cell = width <= 3 ? std::string(width, '.') : cell.substr(0, width - 3) + "..."; } + cell.append(width - cell.size(), ' '); 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 void writeConvReportTableHeader(std::fstream& reportFile, StringRef fourthColumn) { + reportFile << "+------+--------------------------+----------------------+----------+\n"; + reportFile << "| " << convReportCell("Conv", kConvReportIdWidth) << " | " + << convReportCell("Location", kConvReportLocationWidth) << " | " + << convReportCell("Strategy", kConvReportStrategyWidth) << " | " + << convReportCell(fourthColumn, kConvReportCodeWidth) << " |\n"; + reportFile << "+------+--------------------------+----------------------+----------+\n"; } -static bool isSplatConstantValue(Value value, DenseElementsAttr& denseAttr) { - denseAttr = getHostConstDenseElementsAttr(value); - return static_cast(denseAttr) && denseAttr.isSplat(); +static void writeConvReportRow(std::fstream& reportFile, + const ConvLoweringReportEntry& entry) { + reportFile << "| " << convReportCell(std::to_string(entry.convId), kConvReportIdWidth) << " | " + << convReportCell(entry.location, kConvReportLocationWidth) << " | " + << convReportCell(entry.strategy, kConvReportStrategyWidth) << " | " + << convReportCell(entry.implementation, kConvReportCodeWidth) << " |\n"; } -static bool isPerChannelConstantValue(Value value, RankedTensorType currentType, DenseElementsAttr& denseAttr) { - denseAttr = getHostConstDenseElementsAttr(value); - if (!denseAttr || denseAttr.isSplat()) +static void writeConvReportLegend(std::fstream& reportFile) { + reportFile << "Legend: Conv is shared by both sections; codes expand to:\n"; + reportFile << " SEL selectConvLoweringPlan\n"; + reportFile << " DW depthwise::rewriteConv\n"; + reportFile << " PIC standard::rewritePackedIm2ColConv\n"; + reportFile << " STR standard::rewriteStreamedConv(pack=1)\n"; + reportFile << " IKT standard::rewriteInputKTiledConv\n"; + reportFile << " STP standard::rewriteStreamedConv(pack=geo.pack)\n"; + reportFile << " AUTO unresolved strategy\n"; + reportFile << " RSD createRowStripConvOutputFromDenseInput -> createRowStripConvOutput\n"; + reportFile << " RSDW createConvOutputFromRowStripInput -> createDepthwiseOutputFromRowStripFragments\n"; + reportFile << " RSP createConvOutputFromRowStripInput -> createPointwiseOutputFromRowStripFragments\n"; + reportFile << " RSM createConvOutputFromRowStripInput -> createConvOutputFromPixelMajorRowStripFragments\n\n"; +} + +static bool writeConvLoweringReport(const ConvLoweringReportEntry& entry, + ConvLoweringReportState& state) { + state.entries.push_back(entry); + std::fstream reportFile = openReportFile("conv_lowering_report"); + if (!reportFile.is_open()) { + state.entries.pop_back(); 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; + reportFile << "# PIM Conv Lowering Report (bounded to 512 rows)\n\n"; + reportFile << "## Plan selection\n"; + writeConvReportTableHeader(reportFile, "Selector"); + bool realizationSectionStarted = false; + for (const ConvLoweringReportEntry& reportEntry : state.entries) { + if (reportEntry.phase == "realization" && !realizationSectionStarted) { + reportFile << "\n## Realization\n"; + writeConvReportTableHeader(reportFile, "Code"); + realizationSectionStarted = true; } - - 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"; + writeConvReportRow(reportFile, reportEntry); } 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); + writeConvReportLegend(reportFile); + if (!reportFile.good()) { + state.entries.pop_back(); + return false; } - - reportFile << "Strategies used in this report:\n"; - for (PimConvLoweringType strategy : usedStrategies) - reportFile << "- `" << stringifyConvLoweringStrategy(strategy).str() << "`: " - << describeConvLoweringStrategy(strategy).str() << "\n"; + return true; } -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) { - if (requested != PimConvLoweringAuto) - return {requested, "forced by compiler option", /*isAuto=*/false, "", ""}; - - // 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, "", ""}; - 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, - "", - ""}; - 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, - "", - ""}; - if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize) - return {PimConvLoweringStreamedPatch, "fits crossbar but packing is not useful", /*isAuto=*/true, "", ""}; - 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, "", ""}; -} - -[[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) +static void recordConvLoweringReport(Operation* op, + ConvLoweringReportPhase phase, + StringRef strategy, + StringRef implementation, + bool enabled) { + if (!enabled) 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 ConvLoweringReportState state; + std::lock_guard lock(state.mutex); + if (state.entries.size() >= 512) + return; + if (phase == ConvLoweringReportPhase::Planning) { + if (state.planned.contains(op)) + return; + } + else { + if (state.realized.contains(op)) + return; + } + size_t convId = state.convIds.lookup(op); + if (!convId) { + convId = state.nextConvId++; + state.convIds[op] = convId; + } + ConvLoweringReportEntry entry {convId, + stringifyConvLoweringReportPhase(phase).str(), + convReportLocation(op), + strategy.str(), + implementation.str()}; + if (!writeConvLoweringReport(entry, state)) + return; + if (phase == ConvLoweringReportPhase::Planning) + state.planned.insert(op); + else + state.realized.insert(op); } static Value expandBiasIfNeeded(Value bias, PatternRewriter& rewriter, Location loc) { @@ -1184,11 +368,8 @@ 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); - namespace depthwise { struct Tiling { @@ -1208,10 +389,10 @@ static std::optional computeTiling(int64_t batchSize, int64_t wHeight, int64_t wWidth, int64_t outHeight, - int64_t outWidth) { + int64_t outWidth, + int64_t xbarDim) { const int64_t kernelElements = wHeight * wWidth; const int64_t outputMultiplier = numChannelsOut / numChannelsIn; - const int64_t xbarDim = static_cast(crossbarSize.getValue()); if (kernelElements <= 0 || outputMultiplier <= 0 || kernelElements > xbarDim || outputMultiplier > xbarDim) return std::nullopt; @@ -1241,10 +422,13 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr, RankedTensorType wType, const Tiling& tiling, PatternRewriter& rewriter, - Location loc) { - const int64_t paddedOutputChannels = static_cast(crossbarSize.getValue()); + Location loc, + int64_t xbarDim, + int64_t paddedInputRows = -1) { + const int64_t paddedOutputChannels = xbarDim; + 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 +447,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]; } } @@ -1387,7 +571,7 @@ static Value createWeightTile(Value packedWeights, PatternRewriter& rewriter, Location loc) { SmallVector offsets {channelTileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - const int64_t paddedOutputChannels = static_cast(crossbarSize.getValue()); + const int64_t paddedOutputChannels = packedWeightType.getDimSize(2); SmallVector sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling.tileInputRows), rewriter.getIndexAttr(paddedOutputChannels)}; @@ -1513,27 +697,28 @@ static bool canUseStructuredRewrite(const ConvLoweringState& state) { if (!getHostConstDenseElementsAttr(state.w)) return false; - auto tiling = computeTiling(state.batchSize, - state.numChannelsIn, - state.numChannelsOut, - state.wHeight, - state.wWidth, - state.outHeight, - state.outWidth); + auto tiling = computeTiling(state.problem.batchSize, + state.problem.numChannelsIn, + state.problem.numChannelsOut, + state.problem.wHeight, + state.problem.wWidth, + state.problem.outHeight, + state.problem.outWidth, + state.targetInfo().matrixShape.rows); if (!tiling) return false; - if (!state.hasBias) + if (!state.problem.hasBias) return true; auto biasType = dyn_cast(state.b.getType()); if (!biasType) return false; if (biasType.getRank() == 1) - return biasType.getDimSize(0) == state.numChannelsOut; + return biasType.getDimSize(0) == state.problem.numChannelsOut; if (biasType.getRank() != 2) return false; - return biasType.getDimSize(0) == 1 && biasType.getDimSize(1) == state.numChannelsOut; + return biasType.getDimSize(0) == 1 && biasType.getDimSize(1) == state.problem.numChannelsOut; } static FailureOr @@ -1544,24 +729,25 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& return failure(); } - auto tiling = computeTiling(state.xType.getDimSize(0), - state.xType.getDimSize(1), - state.outType.getDimSize(1), - state.wType.getDimSize(2), - state.wType.getDimSize(3), - state.outType.getDimSize(2), - state.outType.getDimSize(3)); + auto tiling = computeTiling(state.problem.xType.getDimSize(0), + state.problem.xType.getDimSize(1), + state.problem.outType.getDimSize(1), + state.problem.wType.getDimSize(2), + state.problem.wType.getDimSize(3), + state.problem.outType.getDimSize(2), + state.problem.outType.getDimSize(3), + state.targetInfo().matrixShape.rows); if (!tiling) { convOp->emitOpError("failed to derive a structured depthwise tiling that fits Spatial weighted VMM lowering"); return failure(); } Value paddedInput = createPaddedInput(state.x, - state.xType, - state.padHeightBegin, - state.padHeightEnd, - state.padWidthBegin, - state.padWidthEnd, + state.problem.xType, + state.problem.padHeightBegin, + state.problem.padHeightEnd, + state.problem.padWidthBegin, + state.problem.padWidthEnd, rewriter, loc); auto paddedInputType = cast(paddedInput.getType()); @@ -1570,17 +756,18 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& paddedInputType.getDimSize(3), paddedInputType.getDimSize(1)}, paddedInputType.getElementType()); - Value channelLastInput = ONNXTransposeOp::create( - rewriter, loc, channelLastInputType, paddedInput, rewriter.getI64ArrayAttr({0, 2, 3, 1})); - Value packedWeights = buildPackedWeights(wDenseAttr, state.wType, *tiling, rewriter, loc); + Value channelLastInput = createLinalgTranspose( + paddedInput, channelLastInputType, {0, 2, 3, 1}, rewriter, loc); + Value packedWeights = buildPackedWeights( + wDenseAttr, state.problem.wType, *tiling, rewriter, loc, state.targetInfo().matrixShape.rows); Value expandedBias; SmallVector batchInputs {channelLastInput}; - if (state.hasBias) { + if (state.problem.hasBias) { expandedBias = expandBiasIfNeeded(state.b, rewriter, loc); auto biasType = dyn_cast(expandedBias.getType()); if (!biasType || biasType.getRank() != 2 || biasType.getDimSize(0) != 1 - || biasType.getDimSize(1) != state.outType.getDimSize(1)) { + || biasType.getDimSize(1) != state.problem.outType.getDimSize(1)) { convOp->emitOpError("requires bias sliceable as tensor<1xCout> for structured depthwise Spatial lowering"); return failure(); } @@ -1588,14 +775,14 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& } auto gemmOutType = - RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType()); - auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType()); + RankedTensorType::get({tiling->totalPatches, state.problem.outType.getDimSize(1)}, state.problem.outType.getElementType()); + auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.problem.outType.getElementType()); auto paddedRowTileType = RankedTensorType::get( - {1, static_cast(crossbarSize.getValue())}, state.outType.getElementType()); + {1, static_cast(state.targetInfo().matrixShape.rows)}, state.problem.outType.getElementType()); auto piecesType = spatial::getGraphBatchPhysicalResultType( tiling->totalPatches * tiling->numChannelTiles, rowTileType); auto inputTileType = - RankedTensorType::get({1, state.wType.getDimSize(2), state.wType.getDimSize(3), tiling->channelsPerTile}, + RankedTensorType::get({1, state.problem.wType.getDimSize(2), state.problem.wType.getDimSize(3), tiling->channelsPerTile}, paddedInputType.getElementType()); SmallVector batchWeights; if (tiling->numChannelTiles == 1) { @@ -1646,11 +833,11 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& channelTileIndex, inputTileType, *tiling, - state.strideHeight, - state.strideWidth, - state.dilationHeight, - state.dilationWidth, - state.outType.getDimSize(3), + state.problem.strideHeight, + state.problem.strideWidth, + state.problem.dilationHeight, + state.problem.dilationWidth, + state.problem.outType.getDimSize(3), rewriter, loc); Value weightTile = tiling->numChannelTiles == 1 @@ -1688,22 +875,21 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& return failure(); auto nhwcType = RankedTensorType::get( - {state.xType.getDimSize(0), state.outType.getDimSize(2), state.outType.getDimSize(3), state.outType.getDimSize(1)}, - state.outType.getElementType()); + {state.problem.xType.getDimSize(0), state.problem.outType.getDimSize(2), state.problem.outType.getDimSize(3), state.problem.outType.getDimSize(1)}, + state.problem.outType.getElementType()); auto reconstructedRows = reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc); if (failed(reconstructedRows)) return failure(); Value collectedRows = *reconstructedRows; return createCollectedConvOutput(ValueRange {collectedRows}, - state.outType, + state.problem.outType, gemmOutType, nhwcType, - state.outType, + state.problem.outType, tiling->totalPatches, - state.outType.getDimSize(1), + state.problem.outType.getDimSize(1), /*packFactor=*/1, - {}, rewriter, loc); } @@ -1741,20 +927,20 @@ buildConvGemmPlan(const ConvLoweringState& state, static PreparedConvInput prepareInputForIm2Col(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) { - if (state.padHeightBegin == 0 && state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0) - return {state.x, state.xType}; + if (state.problem.padHeightBegin == 0 && state.problem.padHeightEnd == 0 && state.problem.padWidthBegin == 0 && state.problem.padWidthEnd == 0) + return {state.x, state.problem.xType}; - auto paddedType = RankedTensorType::get({state.batchSize, - state.numChannelsIn, - state.xHeight + state.padHeightBegin + state.padHeightEnd, - state.xWidth + state.padWidthBegin + state.padWidthEnd}, - state.xType.getElementType()); + auto paddedType = RankedTensorType::get({state.problem.batchSize, + state.problem.numChannelsIn, + state.problem.xHeight + state.problem.padHeightBegin + state.problem.padHeightEnd, + state.problem.xWidth + state.problem.padWidthBegin + state.problem.padWidthEnd}, + state.problem.xType.getElementType()); auto paddedInputOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedType}, {}, state.x, [&](Value inputArg) { Value paddedInput = createZeroPaddedTensor(inputArg, paddedType, - {0, 0, state.padHeightBegin, state.padWidthBegin}, - {0, 0, state.padHeightEnd, state.padWidthEnd}, + {0, 0, state.problem.padHeightBegin, state.problem.padWidthBegin}, + {0, 0, state.problem.padHeightEnd, state.problem.padWidthEnd}, rewriter, loc); spatial::SpatYieldOp::create(rewriter, loc, paddedInput); @@ -1818,8 +1004,7 @@ static Value createWeightMatrix( }); if (!transpose) return flattened; - return ONNXTransposeOp::create(rewriter, loc, plan.wTransType, flattened, rewriter.getI64ArrayAttr({1, 0})) - .getResult(); + return createLinalgTranspose(flattened, plan.wTransType, {1, 0}, rewriter, loc); }; if (isCompileTimeComputable(weights)) @@ -1871,17 +1056,17 @@ static Value createPaddedInputKTiledWeightConstant(DenseElementsAttr sourceAttr, int64_t paddedK, int64_t paddedC, PatternRewriter& rewriter) { - auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.wType.getElementType()); + auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.problem.wType.getElementType()); SmallVector sourceValues(sourceAttr.getValues()); SmallVector paddedValues( paddedType.getNumElements(), cast(rewriter.getZeroAttr(paddedType.getElementType()))); - for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) { - for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) { - for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) { - for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) { + for (int64_t outChannel = 0; outChannel < state.problem.numChannelsOut; ++outChannel) { + for (int64_t inChannel = 0; inChannel < state.problem.numChannelsIn; ++inChannel) { + for (int64_t kernelH = 0; kernelH < state.problem.wHeight; ++kernelH) { + for (int64_t kernelW = 0; kernelW < state.problem.wWidth; ++kernelW) { const int64_t sourceFlatIndex = - (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW; - const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW; + (((outChannel * state.problem.numChannelsIn) + inChannel) * state.problem.wHeight + kernelH) * state.problem.wWidth + kernelW; + const int64_t patchIndex = ((inChannel * state.problem.wHeight) + kernelH) * state.problem.wWidth + kernelW; paddedValues[patchIndex * paddedC + outChannel] = sourceValues[sourceFlatIndex]; } } @@ -1897,22 +1082,22 @@ static Value createPaddedPixelMajorWeightConstant(DenseElementsAttr sourceAttr, int64_t paddedC, int64_t packFactor, PatternRewriter& rewriter) { - auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.wType.getElementType()); + auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.problem.wType.getElementType()); SmallVector sourceValues(sourceAttr.getValues()); SmallVector paddedValues( paddedType.getNumElements(), cast(rewriter.getZeroAttr(paddedType.getElementType()))); - const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; + const int64_t patchSize = state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; for (int64_t copy = 0; copy < packFactor; ++copy) - for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) - for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) - for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) - for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) { + for (int64_t outChannel = 0; outChannel < state.problem.numChannelsOut; ++outChannel) + for (int64_t kernelH = 0; kernelH < state.problem.wHeight; ++kernelH) + for (int64_t kernelW = 0; kernelW < state.problem.wWidth; ++kernelW) + for (int64_t inChannel = 0; inChannel < state.problem.numChannelsIn; ++inChannel) { const int64_t sourceFlatIndex = - (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW; + (((outChannel * state.problem.numChannelsIn) + inChannel) * state.problem.wHeight + kernelH) * state.problem.wWidth + kernelW; const int64_t patchIndex = - ((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel; + ((kernelH * state.problem.wWidth) + kernelW) * state.problem.numChannelsIn + inChannel; const int64_t packedRow = copy * patchSize + patchIndex; - const int64_t packedColumn = copy * state.numChannelsOut + outChannel; + const int64_t packedColumn = copy * state.problem.numChannelsOut + outChannel; paddedValues[packedRow * paddedC + packedColumn] = sourceValues[sourceFlatIndex]; } return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), @@ -1924,22 +1109,22 @@ static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sour int64_t paddedK, int64_t xbarDim, PatternRewriter& rewriter) { - const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim); + const int64_t outputTileCount = ceilIntegerDivide(state.problem.numChannelsOut, xbarDim); auto paddedType = - RankedTensorType::get({outputTileCount, paddedK, xbarDim}, state.wType.getElementType()); + RankedTensorType::get({outputTileCount, paddedK, xbarDim}, state.problem.wType.getElementType()); SmallVector sourceValues(sourceAttr.getValues()); SmallVector paddedValues( paddedType.getNumElements(), cast(rewriter.getZeroAttr(paddedType.getElementType()))); - for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) { + for (int64_t outChannel = 0; outChannel < state.problem.numChannelsOut; ++outChannel) { const int64_t outputTile = outChannel / xbarDim; const int64_t tileChannel = outChannel % xbarDim; - for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) { - for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) { - for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) { + for (int64_t inChannel = 0; inChannel < state.problem.numChannelsIn; ++inChannel) { + for (int64_t kernelH = 0; kernelH < state.problem.wHeight; ++kernelH) { + for (int64_t kernelW = 0; kernelW < state.problem.wWidth; ++kernelW) { const int64_t sourceFlatIndex = - (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW; + (((outChannel * state.problem.numChannelsIn) + inChannel) * state.problem.wHeight + kernelH) * state.problem.wWidth + kernelW; const int64_t patchIndex = - ((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel; + ((kernelH * state.problem.wWidth) + kernelW) * state.problem.numChannelsIn + inChannel; const int64_t destinationFlatIndex = ((outputTile * paddedK) + patchIndex) * xbarDim + tileChannel; paddedValues[destinationFlatIndex] = sourceValues[sourceFlatIndex]; @@ -1952,24 +1137,24 @@ static Value createPaddedOutputChannelTiledWeightConstant(DenseElementsAttr sour } static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, - ArrayRef distributedConsumers, PatternRewriter& rewriter, Location loc) { PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc); - ConvGeometry geo = buildConvGeometry(state); + ConvGeometry geo = buildConvGeometry(state.problem, state.targetInfo()); const int64_t xbarDim = geo.xbarSize; const int64_t numKSlices = ceilIntegerDivide(geo.k, xbarDim); const int64_t paddedK = numKSlices * xbarDim; const uint64_t maxLanesPerBatch = std::max(1, - static_cast(crossbarCountInCore.getValue()) + static_cast(state.targetInfo().matrixUnitsPerProcessor) / static_cast(std::max(1, numKSlices * 4))); const uint64_t rowChunkWidth = std::max( 1, - std::min({chooseStreamChunkPositions(geo, /*packFactor=*/1), + std::min({chooseStreamChunkPositions(geo, /*packFactor=*/1, + state.planningOptions()), maxLanesPerBatch, - static_cast(state.outWidth)})); - const auto elementType = state.outType.getElementType(); + static_cast(state.problem.outWidth)})); + const auto elementType = state.problem.outType.getElementType(); auto wDenseAttr = getHostConstDenseElementsAttr(state.w); if (!wDenseAttr) return failure(); @@ -1978,7 +1163,7 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, Value paddedBias; RankedTensorType paddedBiasType; - if (state.hasBias) { + if (state.problem.hasBias) { Value biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc); auto biasMatrixType = cast(biasMatrix.getType()); paddedBiasType = RankedTensorType::get({1, xbarDim}, elementType); @@ -1992,24 +1177,24 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, } SmallVector chunkRows; - const int64_t totalPatches = state.batchSize * state.outHeight * state.outWidth; + const int64_t totalPatches = state.problem.batchSize * state.problem.outHeight * state.problem.outWidth; chunkRows.reserve( - state.batchSize * state.outHeight * ceilIntegerDivide(state.outWidth, static_cast(rowChunkWidth))); - for (int64_t batchIndex = 0; batchIndex < state.batchSize; ++batchIndex) { - for (int64_t outHeightIndex = 0; outHeightIndex < state.outHeight; ++outHeightIndex) { - for (int64_t outWidthChunkStart = 0; outWidthChunkStart < state.outWidth; + state.problem.batchSize * state.problem.outHeight * ceilIntegerDivide(state.problem.outWidth, static_cast(rowChunkWidth))); + for (int64_t batchIndex = 0; batchIndex < state.problem.batchSize; ++batchIndex) { + for (int64_t outHeightIndex = 0; outHeightIndex < state.problem.outHeight; ++outHeightIndex) { + for (int64_t outWidthChunkStart = 0; outWidthChunkStart < state.problem.outWidth; outWidthChunkStart += static_cast(rowChunkWidth)) { const int64_t chunkNumPatches = - std::min(static_cast(rowChunkWidth), state.outWidth - outWidthChunkStart); - auto chunkRowsType = RankedTensorType::get({chunkNumPatches, state.numChannelsOut}, elementType); + std::min(static_cast(rowChunkWidth), state.problem.outWidth - outWidthChunkStart); + auto chunkRowsType = RankedTensorType::get({chunkNumPatches, state.problem.numChannelsOut}, elementType); auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType); auto paddedChunkRowType = RankedTensorType::get({1, paddedK}, elementType); - auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elementType); + auto patchType = RankedTensorType::get({1, state.problem.numChannelsIn, state.problem.wHeight, state.problem.wWidth}, elementType); auto collapsedPatchType = RankedTensorType::get({1, geo.k}, elementType); - auto weightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType()); - auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType); + auto weightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.problem.wType.getElementType()); + auto rowType = RankedTensorType::get({1, state.problem.numChannelsOut}, elementType); SmallVector inputsStorage {preparedInput.value}; - if (state.hasBias) + if (state.problem.hasBias) inputsStorage.push_back(paddedBias); ValueRange inputs(inputsStorage); auto chunkCompute = spatial::SpatCompute::create(rewriter, loc, TypeRange {chunkRowsType}, ValueRange {paddedWeight}, inputs); @@ -2023,7 +1208,7 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, auto buildChunk = [&]() -> LogicalResult { Value weightArg = block->getArgument(0); Value inputArg = block->getArgument(1); - Value biasArg = state.hasBias ? block->getArgument(2) : Value(); + Value biasArg = state.problem.hasBias ? block->getArgument(2) : Value(); Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); Value cBatchIndex = getOrCreateIndexConstant(rewriter, anchorOp, batchIndex); Value cZero = getOrCreateIndexConstant(rewriter, anchorOp, 0); @@ -2031,7 +1216,7 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, Value cOne = getOrCreateIndexConstant(rewriter, anchorOp, 1); Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim); Value cInputHeightOffset = - getOrCreateIndexConstant(rewriter, anchorOp, outHeightIndex * state.strideHeight); + getOrCreateIndexConstant(rewriter, anchorOp, outHeightIndex * state.problem.strideHeight); Value chunkRowsValue = tensor::EmptyOp::create(rewriter, loc, chunkRowsType.getShape(), elementType); auto widthLoop = buildNormalizedScfFor( @@ -2045,7 +1230,7 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, Value laneWithChunkOffset = affineAddConst(rewriter, nestedLoc, widthIndex, outWidthChunkStart, anchorOp); Value inputWidthOffset = createOrFoldAffineApply(rewriter, nestedLoc, - getAffineDimExpr(0, rewriter.getContext()) * state.strideWidth, + getAffineDimExpr(0, rewriter.getContext()) * state.problem.strideWidth, ValueRange {laneWithChunkOffset}, anchorOp); Value patch = createConvInputPatch(inputArg, @@ -2054,8 +1239,8 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, cZero, cInputHeightOffset, inputWidthOffset, - state.dilationHeight, - state.dilationWidth, + state.problem.dilationHeight, + state.problem.dilationWidth, rewriter, nestedLoc); Value patchRow = tensor::CollapseShapeOp::create(rewriter, @@ -2099,21 +1284,21 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, return failure(); Value reduced = kLoop->results.front(); - if (state.hasBias) + if (state.problem.hasBias) reduced = spatial::SpatVAddOp::create(rewriter, nestedLoc, paddedRowType, reduced, biasArg).getResult(); Value row = reduced; - if (state.numChannelsOut != xbarDim) { + if (state.problem.numChannelsOut != xbarDim) { SmallVector rowOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; SmallVector rowSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)}; + rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.problem.numChannelsOut)}; row = tensor::ExtractSliceOp::create( rewriter, nestedLoc, rowType, reduced, rowOffsets, rowSizes, getUnitStrides(rewriter, 2)); } SmallVector outputOffsets {widthIndex, rewriter.getIndexAttr(0)}; SmallVector outputSizes { - rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)}; + rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.problem.numChannelsOut)}; Value updatedRows = tensor::InsertSliceOp::create( rewriter, nestedLoc, row, iterArgs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 2)); yielded.push_back(updatedRows); @@ -2135,11 +1320,11 @@ static FailureOr rewriteInputKTiledConv(const ConvLoweringState& state, } } - auto nhwcType = RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut}, + auto nhwcType = RankedTensorType::get({state.problem.batchSize, state.problem.outHeight, state.problem.outWidth, state.problem.numChannelsOut}, elementType); return createCollectedConvOutput( - chunkRows, state.outType, cast(chunkRows.front().getType()), nhwcType, state.outType, totalPatches, - state.numChannelsOut, /*packFactor=*/1, distributedConsumers, rewriter, loc); + chunkRows, state.problem.outType, cast(chunkRows.front().getType()), nhwcType, state.problem.outType, totalPatches, + state.problem.numChannelsOut, /*packFactor=*/1, rewriter, loc); } static Value buildPackedWeights(DenseElementsAttr wDenseAttr, @@ -2152,22 +1337,22 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr, return wTrans; auto packedWeightType = RankedTensorType::get( - {plan.effectiveMaxParallelPixels * plan.patchSize, plan.effectiveMaxParallelPixels * state.numChannelsOut}, - state.wType.getElementType()); + {plan.effectiveMaxParallelPixels * plan.patchSize, plan.effectiveMaxParallelPixels * state.problem.numChannelsOut}, + state.problem.wType.getElementType()); SmallVector sourceValues(wDenseAttr.getValues()); SmallVector packedValues(packedWeightType.getNumElements(), - cast(rewriter.getZeroAttr(state.wType.getElementType()))); + cast(rewriter.getZeroAttr(state.problem.wType.getElementType()))); for (int64_t copyId = 0; copyId < plan.effectiveMaxParallelPixels; ++copyId) { - for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) { - for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) { - for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) { - for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) { + for (int64_t outChannel = 0; outChannel < state.problem.numChannelsOut; ++outChannel) { + for (int64_t inChannel = 0; inChannel < state.problem.numChannelsIn; ++inChannel) { + for (int64_t kernelH = 0; kernelH < state.problem.wHeight; ++kernelH) { + for (int64_t kernelW = 0; kernelW < state.problem.wWidth; ++kernelW) { const int64_t sourceFlatIndex = - (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW; - const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW; + (((outChannel * state.problem.numChannelsIn) + inChannel) * state.problem.wHeight + kernelH) * state.problem.wWidth + kernelW; + const int64_t patchIndex = ((inChannel * state.problem.wHeight) + kernelH) * state.problem.wWidth + kernelW; const int64_t targetRow = copyId * plan.patchSize + patchIndex; - const int64_t targetCol = copyId * state.numChannelsOut + outChannel; + const int64_t targetCol = copyId * state.problem.numChannelsOut + outChannel; packedValues[targetRow * packedWeightType.getDimSize(1) + targetCol] = sourceValues[sourceFlatIndex]; } } @@ -2186,7 +1371,7 @@ static Value buildPackedBias(Value gemmBias, const ConvGemmPlan& plan, PatternRewriter& rewriter, Location loc) { - if (!state.hasBias) + if (!state.problem.hasBias) return gemmBias; if (plan.effectiveMaxParallelPixels == 1) @@ -2194,12 +1379,12 @@ static Value buildPackedBias(Value gemmBias, SmallVector sourceValues(biasDenseAttr.getValues()); SmallVector packedValues; - packedValues.reserve(plan.effectiveMaxParallelPixels * state.numChannelsOut); + packedValues.reserve(plan.effectiveMaxParallelPixels * state.problem.numChannelsOut); for (int64_t copyId = 0; copyId < plan.effectiveMaxParallelPixels; ++copyId) packedValues.append(sourceValues.begin(), sourceValues.end()); auto packedBiasType = - RankedTensorType::get({1, plan.effectiveMaxParallelPixels * state.numChannelsOut}, state.outType.getElementType()); + RankedTensorType::get({1, plan.effectiveMaxParallelPixels * state.problem.numChannelsOut}, state.problem.outType.getElementType()); auto packedBiasAttr = DenseElementsAttr::get(packedBiasType, packedValues); return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), packedBiasAttr, packedBiasType); } @@ -2212,30 +1397,30 @@ buildConvGemmPlan(const ConvLoweringState& state, int64_t chunkNumPatches, std::optional forcedPackFactor) { ConvGemmPlan plan; - plan.patchSize = state.numChannelsIn * state.wHeight * state.wWidth; - plan.numPatchesPerBatch = state.outHeight * state.outWidth; - plan.globalNumPatches = state.batchSize * plan.numPatchesPerBatch; + plan.patchSize = state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; + plan.numPatchesPerBatch = state.problem.outHeight * state.problem.outWidth; + plan.globalNumPatches = state.problem.batchSize * plan.numPatchesPerBatch; plan.chunkStart = chunkStart; plan.chunkNumPatches = chunkNumPatches; - const int64_t wMaxDim = std::max(plan.patchSize, state.numChannelsOut); + const int64_t wMaxDim = std::max(plan.patchSize, state.problem.numChannelsOut); plan.maxParallelPixels = forcedPackFactor ? *forcedPackFactor - : std::max(1, static_cast(crossbarSize.getValue()) / wMaxDim); + : std::max(1, static_cast(state.targetInfo().matrixShape.rows) / wMaxDim); plan.effectiveMaxParallelPixels = (canPackWeightsAsConstants && canPackBiasAsConstants) ? plan.maxParallelPixels : 1; plan.packedNumRows = ceilIntegerDivide(plan.chunkNumPatches, plan.effectiveMaxParallelPixels); - auto elemType = state.xType.getElementType(); - auto outElemType = state.outType.getElementType(); + auto elemType = state.problem.xType.getElementType(); + auto outElemType = state.problem.outType.getElementType(); plan.gemmInputRowsType = RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * plan.patchSize}, elemType); - plan.wFlatType = RankedTensorType::get({state.numChannelsOut, plan.patchSize}, state.wType.getElementType()); - plan.wTransType = RankedTensorType::get({plan.patchSize, state.numChannelsOut}, state.wType.getElementType()); - plan.gemmOutType = RankedTensorType::get({plan.chunkNumPatches, state.numChannelsOut}, outElemType); + plan.wFlatType = RankedTensorType::get({state.problem.numChannelsOut, plan.patchSize}, state.problem.wType.getElementType()); + plan.wTransType = RankedTensorType::get({plan.patchSize, state.problem.numChannelsOut}, state.problem.wType.getElementType()); + plan.gemmOutType = RankedTensorType::get({plan.chunkNumPatches, state.problem.numChannelsOut}, outElemType); plan.gemmOutputRowsType = - RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * state.numChannelsOut}, outElemType); + RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * state.problem.numChannelsOut}, outElemType); plan.nhwcType = - RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut}, outElemType); + RankedTensorType::get({state.problem.batchSize, state.problem.outHeight, state.problem.outWidth, state.problem.numChannelsOut}, outElemType); return plan; } @@ -2244,7 +1429,8 @@ static Value createIm2colRows(const ConvLoweringState& state, const ConvGemmPlan& plan, PatternRewriter& rewriter, Location loc) { - if (plan.gemmInputRowsType.getDimSize(1) > crossbarSize.getValue()) { + if (plan.gemmInputRowsType.getDimSize(1) + > static_cast(state.targetInfo().matrixShape.rows)) { assert(plan.effectiveMaxParallelPixels == 1 && "multi-crossbar im2col rows cannot pack pixels"); auto compute = createSpatCompute<1>( rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value input) { @@ -2255,7 +1441,7 @@ static Value createIm2colRows(const ConvLoweringState& state, Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1); Value upper = getOrCreateIndexConstant(rewriter, anchor, plan.chunkNumPatches); auto patchType = RankedTensorType::get( - {1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType); + {1, state.problem.numChannelsIn, state.problem.wHeight, state.problem.wWidth}, elemType); auto rowType = RankedTensorType::get({plan.patchSize}, elemType); auto loop = buildNormalizedScfFor( rewriter, loc, c0, upper, c1, ValueRange {empty}, @@ -2266,14 +1452,14 @@ static Value createIm2colRows(const ConvLoweringState& state, Value batchPatchIndex = affineAddModConst( rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor); Value outHeight = affineFloorDivConst( - rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor); + rewriter, nestedLoc, batchPatchIndex, state.problem.outWidth, anchor); Value outWidth = affineModConst( - rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor); + rewriter, nestedLoc, batchPatchIndex, state.problem.outWidth, anchor); Value patch = createConvInputPatch( input, patchType, batchIndex, c0, - affineMulConst(rewriter, nestedLoc, outHeight, state.strideHeight, anchor), - affineMulConst(rewriter, nestedLoc, outWidth, state.strideWidth, anchor), - state.dilationHeight, state.dilationWidth, rewriter, nestedLoc); + affineMulConst(rewriter, nestedLoc, outHeight, state.problem.strideHeight, anchor), + affineMulConst(rewriter, nestedLoc, outWidth, state.problem.strideWidth, anchor), + state.problem.dilationHeight, state.problem.dilationWidth, rewriter, nestedLoc); Value row = tensor::CollapseShapeOp::create( rewriter, nestedLoc, rowType, patch, SmallVector {{0, 1, 2, 3}}); @@ -2297,7 +1483,7 @@ static Value createIm2colRows(const ConvLoweringState& state, auto elemType = preparedInput.type.getElementType(); auto packedRowType = RankedTensorType::get( {plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding()); - auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType); + auto patchType = RankedTensorType::get({1, state.problem.numChannelsIn, state.problem.wHeight, state.problem.wWidth}, elemType); auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType); bool hasPartialLane = plan.chunkNumPatches % plan.effectiveMaxParallelPixels != 0; SmallVector im2colInputs {preparedInput.value}; @@ -2341,25 +1527,25 @@ static Value createIm2colRows(const ConvLoweringState& state, ValueRange {rowInit}, [&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex); - Value batchIndex = state.batchSize == 1 + Value batchIndex = state.problem.batchSize == 1 ? c0 : affineAddFloorDivConst( rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp); Value batchPatchIndex = affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp); - Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp); - Value outWidthIndex = affineModConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp); - Value inputHeightOffset = affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp); - Value inputWidthOffset = affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp); + Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.problem.outWidth, anchorOp); + Value outWidthIndex = affineModConst(rewriter, nestedLoc, batchPatchIndex, state.problem.outWidth, anchorOp); + Value inputHeightOffset = affineMulConst(rewriter, nestedLoc, outHeightIndex, state.problem.strideHeight, anchorOp); + Value inputWidthOffset = affineMulConst(rewriter, nestedLoc, outWidthIndex, state.problem.strideWidth, anchorOp); Value patch = createConvInputPatch(args.inputs.front(), patchType, batchIndex, c0, inputHeightOffset, inputWidthOffset, - state.dilationHeight, - state.dilationWidth, + state.problem.dilationHeight, + state.problem.dilationWidth, rewriter, nestedLoc); Value patchRow = tensor::CollapseShapeOp::create(rewriter, @@ -2412,43 +1598,46 @@ static Value maybeUnpackChunkRows(Value gemmRows, return unpackCompute.getResult(0); } -static Value createStreamedConvRows(const ConvLoweringState& state, - const PreparedConvInput& preparedInput, - Value weightMatrix, - Value biasMatrix, - DenseElementsAttr wDenseAttr, - DenseElementsAttr biasDenseAttr, - int64_t forcedPackFactor, - PatternRewriter& rewriter, - Location loc) { - const int64_t totalPatches = state.batchSize * state.outHeight * state.outWidth; +static FailureOr createStreamedConvRows(const ConvLoweringState& state, + const PreparedConvInput& preparedInput, + Value weightMatrix, + Value biasMatrix, + DenseElementsAttr wDenseAttr, + DenseElementsAttr biasDenseAttr, + int64_t forcedPackFactor, + PatternRewriter& rewriter, + Location loc) { + const int64_t totalPatches = state.problem.batchSize * state.problem.outHeight * state.problem.outWidth; ConvGemmPlan plan = buildConvGemmPlan(state, static_cast(wDenseAttr), - !state.hasBias || static_cast(biasDenseAttr), 0, totalPatches, forcedPackFactor); + !state.problem.hasBias || static_cast(biasDenseAttr), 0, totalPatches, forcedPackFactor); Value inputRows = createIm2colRows(state, preparedInput, plan, rewriter, loc); Value packedWeights = buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc); - Value gemmBias = state.hasBias ? state.b : createZeroGemmBias(plan.gemmOutputRowsType, rewriter); + Value gemmBias = state.problem.hasBias ? state.b : createZeroGemmBias(plan.gemmOutputRowsType, rewriter); Value packedBias = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc); - Value gemmRows = ONNXGemmOp::create(rewriter, loc, plan.gemmOutputRowsType, inputRows, - packedWeights, packedBias, APFloat(1.0f), APFloat(1.0f), 0, !wDenseAttr).getY(); - return maybeUnpackChunkRows(gemmRows, plan, rewriter, loc); + FailureOr gemmRows = lowerGemmToSpatial( + state.diagnosticAnchor, inputRows, packedWeights, packedBias, + plan.gemmOutputRowsType, /*transA=*/false, /*transB=*/!wDenseAttr, + /*alpha=*/1.0f, /*beta=*/1.0f, state.targetInfo(), rewriter, loc); + if (failed(gemmRows)) + return failure(); + return maybeUnpackChunkRows(*gemmRows, plan, rewriter, loc); } -static Value rewritePackedIm2ColConv(const ConvLoweringState& state, - ArrayRef distributedConsumers, - PatternRewriter& rewriter, - Location loc) { +static FailureOr rewritePackedIm2ColConv(const ConvLoweringState& state, + PatternRewriter& rewriter, + Location loc) { auto wDenseAttr = getHostConstDenseElementsAttr(state.w); PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc); Value biasMatrix; DenseElementsAttr biasDenseAttr; - if (state.hasBias) { + if (state.problem.hasBias) { biasDenseAttr = getHostConstDenseElementsAttr(state.b); biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc); } ConvGemmPlan plan = - buildConvGemmPlan(state, static_cast(wDenseAttr), !state.hasBias || static_cast(biasDenseAttr), 0, - state.batchSize * state.outHeight * state.outWidth); + buildConvGemmPlan(state, static_cast(wDenseAttr), !state.problem.hasBias || static_cast(biasDenseAttr), 0, + state.problem.batchSize * state.problem.outHeight * state.problem.outWidth); // Static weights use the crossbar [patchSize, Cout] layout. Runtime weights // stay in ONNX's contiguous [Cout, patchSize] layout and Gemm consumes them // through transB without materializing a transpose. @@ -2456,98 +1645,66 @@ static Value rewritePackedIm2ColConv(const ConvLoweringState& state, Value gemmInputRows = createIm2colRows(state, preparedInput, plan, rewriter, loc); Value gemmB = buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc); Value gemmBias = createZeroGemmBias(plan.gemmOutputRowsType, rewriter); - if (state.hasBias) + if (state.problem.hasBias) gemmBias = state.b; Value gemmC = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc); - Value gemmRows = ONNXGemmOp::create(rewriter, - loc, - plan.gemmOutputRowsType, - gemmInputRows, - gemmB, - gemmC, - APFloat(1.0f), - APFloat(1.0f), - /*transA=*/0, - /*transB=*/!wDenseAttr) - .getY(); + FailureOr gemmRows = lowerGemmToSpatial( + state.diagnosticAnchor, gemmInputRows, gemmB, gemmC, + plan.gemmOutputRowsType, /*transA=*/false, /*transB=*/!wDenseAttr, + /*alpha=*/1.0f, /*beta=*/1.0f, state.targetInfo(), rewriter, loc); + if (failed(gemmRows)) + return failure(); - return createCollectedConvOutput(ValueRange {gemmRows}, - state.outType, + return createCollectedConvOutput(ValueRange {*gemmRows}, + state.problem.outType, plan.gemmOutType, plan.nhwcType, - state.outType, + state.problem.outType, plan.chunkNumPatches, - state.numChannelsOut, + state.problem.numChannelsOut, plan.effectiveMaxParallelPixels, - distributedConsumers, rewriter, loc); } -static Value rewriteStreamedConv(const ConvLoweringState& state, - ArrayRef distributedConsumers, - PatternRewriter& rewriter, - Location loc, - int64_t forcedPackFactor) { +static FailureOr rewriteStreamedConv(const ConvLoweringState& state, + PatternRewriter& rewriter, + Location loc, + int64_t forcedPackFactor) { auto wDenseAttr = getHostConstDenseElementsAttr(state.w); PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc); Value biasMatrix; DenseElementsAttr biasDenseAttr; - if (state.hasBias) { + if (state.problem.hasBias) { biasDenseAttr = getHostConstDenseElementsAttr(state.b); biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc); } ConvGemmPlan seedPlan = buildConvGemmPlan( - state, static_cast(wDenseAttr), !state.hasBias || static_cast(biasDenseAttr), 0, 1, forcedPackFactor); + state, static_cast(wDenseAttr), !state.problem.hasBias || static_cast(biasDenseAttr), 0, 1, forcedPackFactor); Value weightMatrix = createWeightMatrix(state.w, seedPlan, static_cast(wDenseAttr), rewriter, loc); - Value collectedRows = createStreamedConvRows(state, - preparedInput, - weightMatrix, - biasMatrix, - wDenseAttr, - biasDenseAttr, - forcedPackFactor, - rewriter, - loc); - auto gemmOutType = cast(collectedRows.getType()); - auto nhwcType = RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut}, - state.outType.getElementType()); + FailureOr collectedRows = createStreamedConvRows(state, + preparedInput, + weightMatrix, + biasMatrix, + wDenseAttr, + biasDenseAttr, + forcedPackFactor, + rewriter, + loc); + if (failed(collectedRows)) + return failure(); + auto gemmOutType = cast(collectedRows->getType()); + auto nhwcType = RankedTensorType::get({state.problem.batchSize, state.problem.outHeight, state.problem.outWidth, state.problem.numChannelsOut}, + state.problem.outType.getElementType()); return createCollectedConvOutput( - ValueRange {collectedRows}, state.outType, gemmOutType, nhwcType, state.outType, gemmOutType.getDimSize(0), - state.numChannelsOut, /*packFactor=*/1, distributedConsumers, rewriter, loc); + ValueRange {*collectedRows}, state.problem.outType, gemmOutType, nhwcType, state.problem.outType, gemmOutType.getDimSize(0), + state.problem.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())); @@ -2557,46 +1714,46 @@ static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) { return ceilIntegerDivide(geometry.k, geometry.xbarSize) * ceilIntegerDivide(geometry.c, geometry.xbarSize) - <= static_cast(crossbarCountInCore.getValue()); + <= geometry.matrixUnitsPerProcessor; } static bool rowStripOutputChannelTileFitsOneCore(const ConvGeometry& geometry) { return ceilIntegerDivide(geometry.k, geometry.xbarSize) - <= static_cast(crossbarCountInCore.getValue()); + <= geometry.matrixUnitsPerProcessor; } static int64_t chooseRowStripPixelPackFactor(const ConvLoweringState& state, int64_t xbarDim) { - const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; + const int64_t patchSize = state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; const int64_t baseWeightGroups = ceilIntegerDivide(patchSize, xbarDim) - * ceilIntegerDivide(state.numChannelsOut, xbarDim); - int64_t factor = std::min(state.outWidth, xbarDim / state.numChannelsOut); + * ceilIntegerDivide(state.problem.numChannelsOut, xbarDim); + int64_t factor = std::min(state.problem.outWidth, xbarDim / state.problem.numChannelsOut); while (factor > 1 - && (state.outWidth % factor != 0 + && (state.problem.outWidth % factor != 0 || ceilIntegerDivide(factor * patchSize, xbarDim) - * ceilIntegerDivide(factor * state.numChannelsOut, xbarDim) + * ceilIntegerDivide(factor * state.problem.numChannelsOut, xbarDim) > baseWeightGroups)) --factor; return std::max(factor, 1); } static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) { - if (state.batchSize != 1) { + if (state.problem.batchSize != 1) { failureReason = "batch_not_one"; return false; } - if (state.group != 1) { + if (state.problem.group != 1) { failureReason = "grouped_conv"; return false; } - if (!state.xType.hasStaticShape() || !state.wType.hasStaticShape() || !state.outType.hasStaticShape()) { + if (!state.problem.xType.hasStaticShape() || !state.problem.wType.hasStaticShape() || !state.problem.outType.hasStaticShape()) { failureReason = "dynamic_shape"; return false; } - if (!isa(state.xType.getElementType())) { + if (!isa(state.problem.xType.getElementType())) { failureReason = "non_float_input"; return false; } - if (state.dilationHeight != 1 || state.dilationWidth != 1) { + if (state.problem.dilationHeight != 1 || state.problem.dilationWidth != 1) { failureReason = "dilation_not_one"; return false; } @@ -2604,11 +1761,11 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state failureReason = "non_constant_weight"; return false; } - if (!rowStripOutputChannelTileFitsOneCore(buildConvGeometry(state))) { + if (!rowStripOutputChannelTileFitsOneCore(buildConvGeometry(state.problem, state.targetInfo()))) { failureReason = "output_channel_tile_does_not_fit_one_core"; return false; } - if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) { + if (state.problem.hasBias && !isSupportedBiasAddValue(state.b, state.problem.outType)) { failureReason = "unsupported_bias"; return false; } @@ -2624,17 +1781,17 @@ static FailureOr createBiasRowConstant(const ConvLoweringState& state, int64_t packFactor, PatternRewriter& rewriter) { DenseElementsAttr denseAttr; - if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr)) + if (!isSupportedBiasAddValue(state.b, state.problem.outType, &denseAttr)) return failure(); - FailureOr> channelValues = getBiasChannelValues(denseAttr, state.outType); + FailureOr> channelValues = getBiasChannelValues(denseAttr, state.problem.outType); if (failed(channelValues)) return failure(); SmallVector packedValues; - packedValues.reserve(packFactor * state.numChannelsOut); + packedValues.reserve(packFactor * state.problem.numChannelsOut); for (int64_t copy = 0; copy < packFactor; ++copy) packedValues.append(channelValues->begin(), channelValues->end()); - auto biasType = RankedTensorType::get({1, packFactor * state.numChannelsOut}, state.outType.getElementType()); + auto biasType = RankedTensorType::get({1, packFactor * state.problem.numChannelsOut}, state.problem.outType.getElementType()); return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), DenseElementsAttr::get(biasType, packedValues), @@ -2645,16 +1802,16 @@ static FailureOr createPaddedBiasTileConstant(const ConvLoweringState& st int64_t tileChannels, PatternRewriter& rewriter) { DenseElementsAttr denseAttr; - if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr)) + if (!isSupportedBiasAddValue(state.b, state.problem.outType, &denseAttr)) return failure(); - FailureOr> channelValues = getBiasChannelValues(denseAttr, state.outType); + FailureOr> channelValues = getBiasChannelValues(denseAttr, state.problem.outType); if (failed(channelValues)) return failure(); - const int64_t tileCount = ceilIntegerDivide(state.numChannelsOut, tileChannels); - auto tileType = RankedTensorType::get({tileCount, 1, tileChannels}, state.outType.getElementType()); + const int64_t tileCount = ceilIntegerDivide(state.problem.numChannelsOut, tileChannels); + auto tileType = RankedTensorType::get({tileCount, 1, tileChannels}, state.problem.outType.getElementType()); SmallVector values( tileType.getNumElements(), cast(rewriter.getZeroAttr(tileType.getElementType()))); - for (int64_t channel = 0; channel < state.numChannelsOut; ++channel) + for (int64_t channel = 0; channel < state.problem.numChannelsOut; ++channel) values[channel] = (*channelValues)[channel]; return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), DenseElementsAttr::get(tileType, values), tileType); @@ -2665,27 +1822,27 @@ static Value createHorizontallyPaddedRowStripFragment(Value fragment, PatternRewriter& rewriter, Location loc) { auto paddedType = RankedTensorType::get( - {1, 1, state.xWidth + state.padWidthBegin + state.padWidthEnd, state.numChannelsIn}, - state.xType.getElementType(), - state.xType.getEncoding()); + {1, 1, state.problem.xWidth + state.problem.padWidthBegin + state.problem.padWidthEnd, state.problem.numChannelsIn}, + state.problem.xType.getElementType(), + state.problem.xType.getEncoding()); return createZeroPaddedTensor(fragment, paddedType, - {0, 0, state.padWidthBegin, 0}, - {0, 0, state.padWidthEnd, 0}, + {0, 0, state.problem.padWidthBegin, 0}, + {0, 0, state.problem.padWidthEnd, 0}, rewriter, loc); } static Value createRowStripWindowSourceRowTable(const ConvLoweringState& state, PatternRewriter& rewriter) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - auto tableType = RankedTensorType::get({state.outHeight * state.wHeight}, rewriter.getIndexType()); + auto tableType = RankedTensorType::get({state.problem.outHeight * state.problem.wHeight}, rewriter.getIndexType()); SmallVector values; values.reserve(tableType.getNumElements()); - for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) { - for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) { + for (int64_t outputRow = 0; outputRow < state.problem.outHeight; ++outputRow) { + for (int64_t kernelRow = 0; kernelRow < state.problem.wHeight; ++kernelRow) { int64_t sourceRow = - outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin; - sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1); + outputRow * state.problem.strideHeight + kernelRow * state.problem.dilationHeight - state.problem.padHeightBegin; + sourceRow = std::clamp(sourceRow, int64_t {0}, state.problem.xHeight - 1); values.push_back(rewriter.getIndexAttr(sourceRow)); } } @@ -2698,14 +1855,14 @@ static Value createRowStripWindowSourceSlotTable(const ConvLoweringState& state, PatternRewriter& rewriter) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); auto tableType = - RankedTensorType::get({state.outHeight * state.wHeight * tilesPerRow}, rewriter.getIndexType()); + RankedTensorType::get({state.problem.outHeight * state.problem.wHeight * tilesPerRow}, rewriter.getIndexType()); SmallVector values; values.reserve(tableType.getNumElements()); - for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) - for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) { + for (int64_t outputRow = 0; outputRow < state.problem.outHeight; ++outputRow) + for (int64_t kernelRow = 0; kernelRow < state.problem.wHeight; ++kernelRow) { int64_t sourceRow = - outputRow * state.strideHeight + kernelRow * state.dilationHeight - state.padHeightBegin; - sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1); + outputRow * state.problem.strideHeight + kernelRow * state.problem.dilationHeight - state.problem.padHeightBegin; + sourceRow = std::clamp(sourceRow, int64_t {0}, state.problem.xHeight - 1); for (int64_t tile = 0; tile < tilesPerRow; ++tile) values.push_back(rewriter.getIndexAttr(sourceRow * tilesPerRow + tile)); } @@ -2722,7 +1879,7 @@ static Value createRowStripWindowTableIndex(Value outputHeight, AffineExpr outputRowExpr = getAffineDimExpr(0, ctx); AffineExpr kernelRowExpr = getAffineDimExpr(1, ctx); return createOrFoldAffineApply( - rewriter, loc, outputRowExpr * state.wHeight + kernelRowExpr, ValueRange {outputHeight, kernelRow}, anchorOp); + rewriter, loc, outputRowExpr * state.problem.wHeight + kernelRowExpr, ValueRange {outputHeight, kernelRow}, anchorOp); } static FailureOr extractProjectedRowStripWindowRow(Value rowStripStorage, @@ -2732,7 +1889,7 @@ static FailureOr extractProjectedRowStripWindowRow(Value rowStripStorage, Value kernelRow, PatternRewriter& rewriter, Location loc) { - FailureOr physical = describeRowStripPhysicalValue(rowStripStorage, state.xType); + FailureOr physical = describeRowStripPhysicalValue(rowStripStorage, state.problem.xType); if (failed(physical)) return failure(); Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc); @@ -2746,7 +1903,7 @@ static FailureOr extractProjectedRowStripWindowRow(Value rowStripStorage, rewriter, loc, rowStripStorage, sourceSlot, physical->fragmentType); } - auto fullFragmentType = getRowStripFragmentType(state.xType); + auto fullFragmentType = getRowStripFragmentType(state.problem.xType); Value fullFragment = tensor::EmptyOp::create( rewriter, loc, fullFragmentType.getShape(), fullFragmentType.getElementType()); const int64_t tileChannels = physical->fragmentType.getDimSize(3); @@ -2759,9 +1916,9 @@ static FailureOr extractProjectedRowStripWindowRow(Value rowStripStorage, if (failed(fragment)) return failure(); const int64_t channelOffset = tile * tileChannels; - const int64_t validChannels = std::min(tileChannels, state.numChannelsIn - channelOffset); + const int64_t validChannels = std::min(tileChannels, state.problem.numChannelsIn - channelOffset); auto validType = RankedTensorType::get( - {1, 1, state.xWidth, validChannels}, state.xType.getElementType(), state.xType.getEncoding()); + {1, 1, state.problem.xWidth, validChannels}, state.problem.xType.getElementType(), state.problem.xType.getEncoding()); Value validFragment = *fragment; if (validChannels != tileChannels) validFragment = tensor::ExtractSliceOp::create( @@ -2772,7 +1929,7 @@ static FailureOr extractProjectedRowStripWindowRow(Value rowStripStorage, SmallVector(4, rewriter.getIndexAttr(0)), SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.xWidth), + rewriter.getIndexAttr(state.problem.xWidth), rewriter.getIndexAttr(validChannels)}, getUnitStrides(rewriter, 4)); fullFragment = tensor::InsertSliceOp::create( @@ -2786,7 +1943,7 @@ static FailureOr extractProjectedRowStripWindowRow(Value rowStripStorage, rewriter.getIndexAttr(channelOffset)}, SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.xWidth), + rewriter.getIndexAttr(state.problem.xWidth), rewriter.getIndexAttr(validChannels)}, getUnitStrides(rewriter, 4)); } @@ -2803,69 +1960,63 @@ static Value extractDenseConvWindowRow(Value denseInput, Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc); Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult(); auto nchwType = RankedTensorType::get( - {1, state.numChannelsIn, 1, state.xWidth}, state.xType.getElementType(), state.xType.getEncoding()); - auto fragmentType = getRowStripFragmentType(state.xType); + {1, state.problem.numChannelsIn, 1, state.problem.xWidth}, state.problem.xType.getElementType(), state.problem.xType.getEncoding()); + auto fragmentType = getRowStripFragmentType(state.problem.xType); SmallVector offsets { rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceRow, rewriter.getIndexAttr(0)}; SmallVector sizes {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.numChannelsIn), + rewriter.getIndexAttr(state.problem.numChannelsIn), rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.xWidth)}; + rewriter.getIndexAttr(state.problem.xWidth)}; Value nchw = tensor::ExtractSliceOp::create( rewriter, loc, nchwType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4)); - return ONNXTransposeOp::create( - rewriter, loc, fragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})); + return createLinalgTranspose(nchw, fragmentType, {0, 2, 3, 1}, rewriter, loc); } -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.problem.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.problem.xWidth, state.problem.numChannelsIn}, elementType, state.problem.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.problem.outHeight * state.problem.wHeight}, rewriter.getIndexType()); SmallVector values; values.reserve(tableType.getNumElements()); - for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) { - for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) { + for (int64_t outputRow = 0; outputRow < state.problem.outHeight; ++outputRow) + for (int64_t kernelRow = 0; kernelRow < state.problem.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); + outputRow * state.problem.strideHeight + kernelRow * state.problem.dilationHeight - state.problem.padHeightBegin; + values.push_back(rewriter.getIndexAttr(sourceRow >= 0 && sourceRow < state.problem.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); - 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)); + Value maskIndex = tensor::ExtractOp::create(rewriter, loc, maskIndexTable, ValueRange {tableIndex}).getResult(); + auto fragmentType = getRowStripFragmentType(state.problem.xType); + 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.problem.xWidth), rewriter.getIndexAttr(state.problem.numChannelsIn)}, + getUnitStrides(rewriter, 4)); } static FailureOr createConvInputWindow(Value input, @@ -2873,29 +2024,32 @@ static FailureOr createConvInputWindow(Value input, Value outputHeight, PatternRewriter& rewriter, Location loc) { - auto fragmentType = getRowStripFragmentType(state.xType); + auto fragmentType = getRowStripFragmentType(state.problem.xType); auto inputType = dyn_cast(input.getType()); - const bool denseInput = inputType == state.xType; - if (!denseInput && failed(describeRowStripPhysicalValue(input, state.xType))) + const bool denseInput = inputType == state.problem.xType; + FailureOr physicalInput = describeRowStripPhysicalValue(input, state.problem.xType); + if (!denseInput && failed(physicalInput)) return failure(); + if (!denseInput && physicalInput->tilesPerRow == 1 && state.problem.wHeight == 1 && state.problem.wWidth == 1 + && state.problem.strideHeight == 1 && state.problem.strideWidth == 1 && state.problem.padHeightBegin == 0 + && state.problem.padHeightEnd == 0 && state.problem.padWidthBegin == 0 && state.problem.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); + {1, state.problem.wHeight, state.problem.xWidth + state.problem.padWidthBegin + state.problem.padWidthEnd, state.problem.numChannelsIn}, + state.problem.xType.getElementType(), + state.problem.xType.getEncoding()); 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(); Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); - Value cKernelRows = getOrCreateIndexConstant(rewriter, anchorOp, state.wHeight); + Value cKernelRows = getOrCreateIndexConstant(rewriter, anchorOp, state.problem.wHeight); auto loop = buildNormalizedScfFor( rewriter, loc, @@ -2913,8 +2067,9 @@ static FailureOr createConvInputWindow(Value input, if (failed(sourceRow)) return failure(); Value semanticRow = *sourceRow; - if (state.padHeightBegin != 0 || state.padHeightEnd != 0) { - Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, rowLoc); + if (state.problem.padHeightBegin != 0 || state.problem.padHeightEnd != 0) { + 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); @@ -2926,9 +2081,9 @@ static FailureOr createConvInputWindow(Value input, SmallVector {rewriter.getIndexAttr(0), kernelRow, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.xWidth + state.padWidthBegin - + state.padWidthEnd), - rewriter.getIndexAttr(state.numChannelsIn)}, + rewriter.getIndexAttr(state.problem.xWidth + state.problem.padWidthBegin + + state.problem.padWidthEnd), + rewriter.getIndexAttr(state.problem.numChannelsIn)}, getUnitStrides(rewriter, 4))); return success(); }); @@ -2942,33 +2097,33 @@ static FailureOr createPixelMajorConvPatchRow(Value paddedWindow, PatternRewriter& rewriter, Location loc) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; - auto patchType = RankedTensorType::get({1, state.wHeight, state.wWidth, state.numChannelsIn}, - state.xType.getElementType(), - state.xType.getEncoding()); - auto rowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType(), state.xType.getEncoding()); - Value inputWidthOffset = affineMulConst(rewriter, loc, outputWidth, state.strideWidth, anchorOp); + const int64_t patchSize = state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; + auto patchType = RankedTensorType::get({1, state.problem.wHeight, state.problem.wWidth, state.problem.numChannelsIn}, + state.problem.xType.getElementType(), + state.problem.xType.getEncoding()); + auto rowType = RankedTensorType::get({1, patchSize}, state.problem.xType.getElementType(), state.problem.xType.getEncoding()); + Value inputWidthOffset = affineMulConst(rewriter, loc, outputWidth, state.problem.strideWidth, anchorOp); SmallVector offsets { rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), inputWidthOffset, rewriter.getIndexAttr(0)}; SmallVector sizes {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.wHeight), - rewriter.getIndexAttr(state.wWidth), - rewriter.getIndexAttr(state.numChannelsIn)}; + rewriter.getIndexAttr(state.problem.wHeight), + rewriter.getIndexAttr(state.problem.wWidth), + rewriter.getIndexAttr(state.problem.numChannelsIn)}; Value patch; - if (state.dilationWidth == 1) + if (state.problem.dilationWidth == 1) patch = tensor::ExtractSliceOp::create( rewriter, loc, patchType, paddedWindow, offsets, sizes, getUnitStrides(rewriter, 4)); else { - auto columnType = RankedTensorType::get({1, state.wHeight, 1, state.numChannelsIn}, - state.xType.getElementType(), state.xType.getEncoding()); + auto columnType = RankedTensorType::get({1, state.problem.wHeight, 1, state.problem.numChannelsIn}, + state.problem.xType.getElementType(), state.problem.xType.getEncoding()); patch = tensor::EmptyOp::create(rewriter, loc, patchType.getShape(), patchType.getElementType()); - for (int64_t kernelColumn = 0; kernelColumn < state.wWidth; ++kernelColumn) { + for (int64_t kernelColumn = 0; kernelColumn < state.problem.wWidth; ++kernelColumn) { Value sourceWidth = - affineAddConst(rewriter, loc, inputWidthOffset, kernelColumn * state.dilationWidth, anchorOp); + affineAddConst(rewriter, loc, inputWidthOffset, kernelColumn * state.problem.dilationWidth, anchorOp); SmallVector columnSizes {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.wHeight), + rewriter.getIndexAttr(state.problem.wHeight), rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.numChannelsIn)}; + rewriter.getIndexAttr(state.problem.numChannelsIn)}; Value column = tensor::ExtractSliceOp::create( rewriter, loc, columnType, paddedWindow, SmallVector {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceWidth, @@ -2996,9 +2151,9 @@ static FailureOr createPackedPixelMajorConvPatchRow(Value paddedWindow, return createPixelMajorConvPatchRow(paddedWindow, state, outputGroup, rewriter, loc); Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); - const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; + const int64_t patchSize = state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; auto packedType = RankedTensorType::get( - {1, packFactor * patchSize}, state.xType.getElementType(), state.xType.getEncoding()); + {1, packFactor * patchSize}, state.problem.xType.getElementType(), state.problem.xType.getEncoding()); Value packed = tensor::EmptyOp::create(rewriter, loc, packedType.getShape(), packedType.getElementType()); Value outputStart = affineMulConst(rewriter, loc, outputGroup, packFactor, anchorOp); Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); @@ -3046,25 +2201,25 @@ static FailureOr> createConvInputTiles(Value paddedWindow, int64_t xbarDim, PatternRewriter& rewriter, Location loc) { - auto elementType = state.xType.getElementType(); + auto elementType = state.problem.xType.getElementType(); auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType); SmallVector inputTiles; inputTiles.reserve(numKSlices); - if (packFactor == 1 && state.numChannelsIn % xbarDim == 0) { + if (packFactor == 1 && state.problem.numChannelsIn % xbarDim == 0) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); auto inputTileType = RankedTensorType::get( - {1, 1, 1, xbarDim}, elementType, state.xType.getEncoding()); + {1, 1, 1, xbarDim}, elementType, state.problem.xType.getEncoding()); for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) { const int64_t linearOffset = kSlice * xbarDim; - const int64_t kernelPixel = linearOffset / state.numChannelsIn; - const int64_t kernelRow = kernelPixel / state.wWidth; - const int64_t kernelColumn = kernelPixel % state.wWidth; - const int64_t channelOffset = linearOffset % state.numChannelsIn; + const int64_t kernelPixel = linearOffset / state.problem.numChannelsIn; + const int64_t kernelRow = kernelPixel / state.problem.wWidth; + const int64_t kernelColumn = kernelPixel % state.problem.wWidth; + const int64_t channelOffset = linearOffset % state.problem.numChannelsIn; Value inputWidthOffset = - affineMulConst(rewriter, loc, outputWidth, state.strideWidth, anchorOp); + affineMulConst(rewriter, loc, outputWidth, state.problem.strideWidth, anchorOp); inputWidthOffset = affineAddConst( - rewriter, loc, inputWidthOffset, kernelColumn * state.dilationWidth, anchorOp); + rewriter, loc, inputWidthOffset, kernelColumn * state.problem.dilationWidth, anchorOp); Value inputTile = tensor::ExtractSliceOp::create( rewriter, loc, @@ -3235,14 +2390,14 @@ static FailureOr createRowStripConvOutput(const ConvLoweringState& state, int64_t xbarDim, PatternRewriter& rewriter, Location loc) { - const int64_t laneCount = state.outHeight; - const int64_t patchSize = packFactor * state.numChannelsIn * state.wHeight * state.wWidth; + const int64_t laneCount = state.problem.outHeight; + const int64_t patchSize = packFactor * state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; const bool hasPartialInputTile = patchSize % xbarDim != 0; - auto elementType = state.outType.getElementType(); + auto elementType = state.problem.outType.getElementType(); auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType); - auto outputPixelType = RankedTensorType::get({1, 1, packFactor, state.numChannelsOut}, elementType); - auto fragmentType = getRowStripFragmentType(state.outType); - auto storageType = getRowStripStorageType(state.outType); + auto outputPixelType = RankedTensorType::get({1, 1, packFactor, state.problem.numChannelsOut}, elementType); + auto fragmentType = getRowStripFragmentType(state.problem.outType); + auto storageType = getRowStripStorageType(state.problem.outType); auto batch = createSpatComputeBatch( rewriter, @@ -3255,7 +2410,7 @@ static FailureOr createRowStripConvOutput(const ConvLoweringState& state, Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); - Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth / packFactor); + Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.problem.outWidth / packFactor); FailureOr inputWindow = createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc); if (failed(inputWindow)) @@ -3291,7 +2446,7 @@ static FailureOr createRowStripConvOutput(const ConvLoweringState& state, return failure(); FailureOr output = createConvOutputRow(*inputTiles, paddedK, - packFactor * state.numChannelsOut, + packFactor * state.problem.numChannelsOut, args.weights.front(), bias ? args.inputs[1] : Value(), xbarDim, @@ -3313,7 +2468,7 @@ static FailureOr createRowStripConvOutput(const ConvLoweringState& state, SmallVector {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(packFactor), - rewriter.getIndexAttr(state.numChannelsOut)}, + rewriter.getIndexAttr(state.problem.numChannelsOut)}, getUnitStrides(rewriter, 4)); yielded.push_back(next); if (hasPartialInputTile) @@ -3339,21 +2494,21 @@ static FailureOr createOutputChannelTiledRowStripConvOutput(const ConvLow int64_t xbarDim, PatternRewriter& rewriter, Location loc) { - const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim); - const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth; + const int64_t outputTileCount = ceilIntegerDivide(state.problem.numChannelsOut, xbarDim); + const int64_t patchSize = state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; const bool hasPartialInputTile = patchSize % xbarDim != 0; - auto elementType = state.outType.getElementType(); + auto elementType = state.problem.outType.getElementType(); auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType); auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType); auto tilePixelType = RankedTensorType::get({1, 1, 1, xbarDim}, elementType); - auto tileFragmentType = RankedTensorType::get({1, 1, state.outWidth, xbarDim}, elementType); - auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.wType.getElementType()); - const int64_t laneCount = state.outHeight * outputTileCount; + auto tileFragmentType = RankedTensorType::get({1, 1, state.problem.outWidth, xbarDim}, elementType); + auto tileWeightsType = RankedTensorType::get({paddedK, xbarDim}, state.problem.wType.getElementType()); + const int64_t laneCount = state.problem.outHeight * outputTileCount; auto tileStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, tileFragmentType); FailureOr paddedBias = failure(); - if (state.hasBias) + if (state.problem.hasBias) paddedBias = createPaddedBiasTileConstant(state, xbarDim, rewriter); - if (state.hasBias && failed(paddedBias)) + if (state.problem.hasBias && failed(paddedBias)) return failure(); auto tileBatch = createSpatComputeBatch( @@ -3362,12 +2517,12 @@ static FailureOr createOutputChannelTiledRowStripConvOutput(const ConvLow TypeRange {tileStorageType}, laneCount, ValueRange {paddedWeights}, - state.hasBias ? ValueRange {input, *paddedBias} : ValueRange {input}, + state.problem.hasBias ? ValueRange {input, *paddedBias} : ValueRange {input}, [&](detail::SpatComputeBatchBodyArgs args) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1); - Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth); + Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.problem.outWidth); Value outputRow = affineFloorDivConst(rewriter, loc, args.lane, outputTileCount, anchorOp); Value outputTile = affineModConst(rewriter, loc, args.lane, outputTileCount, anchorOp); SmallVector weightOffsets { @@ -3377,9 +2532,9 @@ static FailureOr createOutputChannelTiledRowStripConvOutput(const ConvLow Value tileWeights = tensor::ExtractSliceOp::create( rewriter, loc, tileWeightsType, args.weights.front(), weightOffsets, weightSizes, getUnitStrides(rewriter, 3)); FailureOr biasTile = failure(); - if (state.hasBias) + if (state.problem.hasBias) biasTile = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], outputTile, paddedRowType); - if (state.hasBias && failed(biasTile)) + if (state.problem.hasBias && failed(biasTile)) return failure(); FailureOr inputWindow = createConvInputWindow(args.inputs.front(), state, outputRow, rewriter, loc); @@ -3418,7 +2573,7 @@ static FailureOr createOutputChannelTiledRowStripConvOutput(const ConvLow createConvOutputTile(*inputTiles, tileWeights, xbarDim, xbarDim, rewriter, widthLoc); if (failed(paddedOutputRow)) return failure(); - if (state.hasBias) + if (state.problem.hasBias) paddedOutputRow = spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, *paddedOutputRow, *biasTile).getResult(); Value outputPixel = tensor::ExpandShapeOp::create( @@ -3453,18 +2608,18 @@ static FailureOr createOutputChannelTiledRowStripConvOutput(const ConvLow static FailureOr createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) { - ConvGeometry geometry = buildConvGeometry(state); - if (state.group != 1 || state.batchSize != 1 || !rowStripOutputChannelTileFitsOneCore(geometry)) + ConvGeometry geometry = buildConvGeometry(state.problem, state.targetInfo()); + if (state.problem.group != 1 || state.problem.batchSize != 1 || !rowStripOutputChannelTileFitsOneCore(geometry)) return failure(); auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); if (!weightDenseAttr) return failure(); - if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) + if (state.problem.hasBias && !isSupportedBiasAddValue(state.b, state.problem.outType)) return failure(); const int64_t xbarDim = geometry.xbarSize; - const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth; + const int64_t basePatchSize = state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim); const int64_t basePaddedK = baseNumKSlices * xbarDim; @@ -3479,18 +2634,18 @@ static FailureOr const int64_t packedPatchSize = packFactor * basePatchSize; const int64_t numKSlices = ceilIntegerDivide(packedPatchSize, xbarDim); const int64_t paddedK = numKSlices * xbarDim; - const int64_t packedOutputChannels = packFactor * state.numChannelsOut; + const int64_t packedOutputChannels = packFactor * state.problem.numChannelsOut; const int64_t paddedOutputChannels = ceilIntegerDivide(packedOutputChannels, xbarDim) * xbarDim; Value paddedWeights = standard::createPaddedPixelMajorWeightConstant( weightDenseAttr, state, paddedK, paddedOutputChannels, packFactor, rewriter); FailureOr bias = failure(); - if (state.hasBias) + if (state.problem.hasBias) bias = createBiasRowConstant(state, packFactor, rewriter); - if (state.hasBias && failed(bias)) + if (state.problem.hasBias && failed(bias)) return failure(); return createRowStripConvOutput( - state, state.x, paddedWeights, state.hasBias ? *bias : Value(), + state, state.x, paddedWeights, state.problem.hasBias ? *bias : Value(), packFactor, paddedK, numKSlices, xbarDim, rewriter, loc); } @@ -3498,16 +2653,16 @@ static FailureOr createConvOutputFromPixelMajorRowStripFragments(Value ro const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) { - if (failed(describeRowStripPhysicalValue(rowStripStorage, state.xType))) + if (failed(describeRowStripPhysicalValue(rowStripStorage, state.problem.xType))) return failure(); StringRef failureReason; if (!canConsumePixelMajorRowStripFragments(state, failureReason)) return failure(); - ConvGeometry geometry = buildConvGeometry(state); + ConvGeometry geometry = buildConvGeometry(state.problem, state.targetInfo()); const int64_t xbarDim = geometry.xbarSize; - const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth; + const int64_t basePatchSize = state.problem.numChannelsIn * state.problem.wHeight * state.problem.wWidth; const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim); const int64_t basePaddedK = baseNumKSlices * xbarDim; auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); @@ -3524,18 +2679,18 @@ static FailureOr createConvOutputFromPixelMajorRowStripFragments(Value ro const int64_t packedPatchSize = packFactor * basePatchSize; const int64_t numKSlices = ceilIntegerDivide(packedPatchSize, xbarDim); const int64_t paddedK = numKSlices * xbarDim; - const int64_t packedOutputChannels = packFactor * state.numChannelsOut; + const int64_t packedOutputChannels = packFactor * state.problem.numChannelsOut; const int64_t paddedOutputChannels = ceilIntegerDivide(packedOutputChannels, xbarDim) * xbarDim; Value paddedWeights = standard::createPaddedPixelMajorWeightConstant( weightDenseAttr, state, paddedK, paddedOutputChannels, packFactor, rewriter); FailureOr bias = failure(); - if (state.hasBias) + if (state.problem.hasBias) bias = createBiasRowConstant(state, packFactor, rewriter); - if (state.hasBias && failed(bias)) + if (state.problem.hasBias && failed(bias)) return failure(); return createRowStripConvOutput( - state, rowStripStorage, paddedWeights, state.hasBias ? *bias : Value(), + state, rowStripStorage, paddedWeights, state.problem.hasBias ? *bias : Value(), packFactor, paddedK, numKSlices, xbarDim, rewriter, loc); } @@ -3543,34 +2698,34 @@ static FailureOr createPointwiseOutputFromRowStripFragments(Value rowStri const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) { - FailureOr input = describeRowStripPhysicalValue(rowStripStorage, state.xType); + FailureOr input = describeRowStripPhysicalValue(rowStripStorage, state.problem.xType); if (failed(input)) return failure(); - ConvGeometry geometry = buildConvGeometry(state); + ConvGeometry geometry = buildConvGeometry(state.problem, state.targetInfo()); const int64_t xbarDim = geometry.xbarSize; const int64_t inputFragmentChannels = input->fragmentType.getDimSize(3); - if (inputFragmentChannels % xbarDim != 0 || state.numChannelsIn % xbarDim != 0) + if (inputFragmentChannels % xbarDim != 0 || state.problem.numChannelsIn % xbarDim != 0) return failure(); auto weightDenseAttr = getHostConstDenseElementsAttr(state.w); if (!weightDenseAttr) return failure(); - const int64_t outputTileCount = ceilIntegerDivide(state.numChannelsOut, xbarDim); - const int64_t numKSlices = state.numChannelsIn / xbarDim; - auto elementType = state.outType.getElementType(); + const int64_t outputTileCount = ceilIntegerDivide(state.problem.numChannelsOut, xbarDim); + const int64_t numKSlices = state.problem.numChannelsIn / xbarDim; + auto elementType = state.problem.outType.getElementType(); auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType); auto inputRowType = RankedTensorType::get({1, inputFragmentChannels}, elementType); - auto weightTileType = RankedTensorType::get({state.numChannelsIn, xbarDim}, state.wType.getElementType()); - auto weightSliceType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType()); + auto weightTileType = RankedTensorType::get({state.problem.numChannelsIn, xbarDim}, state.problem.wType.getElementType()); + auto weightSliceType = RankedTensorType::get({xbarDim, xbarDim}, state.problem.wType.getElementType()); auto outputFragmentType = RankedTensorType::get({1, 1, 1, xbarDim}, elementType); auto outputStorageType = spatial::getGraphBatchPhysicalResultType(outputTileCount, outputFragmentType); Value paddedWeights = standard::createPaddedOutputChannelTiledWeightConstant( - weightDenseAttr, state, state.numChannelsIn, xbarDim, rewriter); + weightDenseAttr, state, state.problem.numChannelsIn, xbarDim, rewriter); FailureOr paddedBias = failure(); - if (state.hasBias) paddedBias = createPaddedBiasTileConstant(state, xbarDim, rewriter); - if (state.hasBias && failed(paddedBias)) return failure(); + if (state.problem.hasBias) paddedBias = createPaddedBiasTileConstant(state, xbarDim, rewriter); + if (state.problem.hasBias && failed(paddedBias)) return failure(); auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {outputStorageType}, outputTileCount, ValueRange {paddedWeights}, - state.hasBias ? ValueRange {rowStripStorage, *paddedBias} : ValueRange {rowStripStorage}, + state.problem.hasBias ? ValueRange {rowStripStorage, *paddedBias} : ValueRange {rowStripStorage}, [&](detail::SpatComputeBatchBodyArgs args) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0); @@ -3578,7 +2733,7 @@ static FailureOr createPointwiseOutputFromRowStripFragments(Value rowStri Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices); SmallVector weightOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; SmallVector weightSizes {rewriter.getIndexAttr(1), - rewriter.getIndexAttr(state.numChannelsIn), rewriter.getIndexAttr(xbarDim)}; + rewriter.getIndexAttr(state.problem.numChannelsIn), rewriter.getIndexAttr(xbarDim)}; Value weightTile = tensor::ExtractSliceOp::create( rewriter, loc, weightTileType, args.weights.front(), weightOffsets, weightSizes, getUnitStrides(rewriter, 3)); auto createPiece = [&](Value kSlice, Location pieceLoc) -> FailureOr { @@ -3617,7 +2772,7 @@ static FailureOr createPointwiseOutputFromRowStripFragments(Value rowStri if (failed(reduction)) return failure(); result = reduction->results.front(); } - if (state.hasBias) { + if (state.problem.hasBias) { FailureOr bias = extractGraphBatchPhysicalFragment( rewriter, loc, args.inputs[1], args.lane, paddedRowType); if (failed(bias)) return failure(); @@ -3632,126 +2787,279 @@ static FailureOr createPointwiseOutputFromRowStripFragments(Value rowStri return batch->getResult(0); } +static bool canConsumeDepthwiseRowStrip(const ConvLoweringState& state) { + if (state.problem.batchSize != 1 || state.problem.group != state.problem.numChannelsIn + || state.problem.dilationHeight != 1 || state.problem.dilationWidth != 1 + || !isa(state.problem.xType.getElementType()) + || !getHostConstDenseElementsAttr(state.w) + || (state.problem.hasBias && !isSupportedBiasAddValue(state.b, state.problem.outType))) + return false; + auto tiling = depthwise::computeTiling(state.problem.batchSize, + state.problem.numChannelsIn, + state.problem.numChannelsOut, + state.problem.wHeight, + state.problem.wWidth, + state.problem.outHeight, + state.problem.outWidth, + state.targetInfo().matrixShape.rows); + return tiling && tiling->numChannelTiles + <= static_cast(state.targetInfo().matrixUnitsPerProcessor); +} + +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.problem.wWidth, anchor); + Value kernelColumn = affineModConst(rewriter, kernelLoc, kernelIndex, state.problem.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.problem.xType))) + return failure(); + auto tiling = depthwise::computeTiling(state.problem.batchSize, + state.problem.numChannelsIn, + state.problem.numChannelsOut, + state.problem.wHeight, + state.problem.wWidth, + state.problem.outHeight, + state.problem.outWidth, + state.targetInfo().matrixShape.rows); + auto weight = getHostConstDenseElementsAttr(state.w); + if (!tiling || !weight) + return failure(); + + Value packedWeights = depthwise::buildPackedWeights( + weight, + state.problem.wType, + *tiling, + rewriter, + loc, + static_cast(state.targetInfo().matrixShape.rows), + static_cast(state.targetInfo().matrixShape.rows)); + Value bias = state.problem.hasBias ? expandBiasIfNeeded(state.b, rewriter, loc) : Value(); + auto paddedOutputType = RankedTensorType::get( + {1, static_cast(state.targetInfo().matrixShape.rows)}, state.problem.outType.getElementType()); + auto outputTileType = RankedTensorType::get( + {1, tiling->tileOutputChannels}, state.problem.outType.getElementType()); + auto outputPixelType = RankedTensorType::get( + {1, 1, 1, tiling->tileOutputChannels}, state.problem.outType.getElementType()); + auto fragmentType = getRowStripFragmentType(state.problem.outType); + auto storageType = getRowStripStorageType(state.problem.outType); + + auto batch = createSpatComputeBatch( + rewriter, + loc, + TypeRange {storageType}, + state.problem.outHeight, + ValueRange {packedWeights}, + state.problem.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.problem.outWidth); + const int64_t xbarDim = static_cast(state.targetInfo().matrixShape.rows); + auto paddedInputScratchType = RankedTensorType::get( + {tiling->numChannelTiles, 1, 1, xbarDim}, state.problem.xType.getElementType(), state.problem.xType.getEncoding()); + auto tileScratchType = RankedTensorType::get( + {1, 1, 1, xbarDim}, state.problem.xType.getElementType(), state.problem.xType.getEncoding()); + auto vmmInputType = RankedTensorType::get( + {1, xbarDim}, state.problem.xType.getElementType(), state.problem.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.problem.wType.getElementType()), + {SmallVector {tileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + weightTileSizes, + getUnitStrides(rewriter, 3)})); + if (state.problem.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.problem.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.problem.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, + ConvMaterializationKind kind, PatternRewriter& rewriter, Location loc) { - if (state.xHeight == 1 && state.xWidth == 1 && state.wHeight == 1 && state.wWidth == 1) + if (kind == ConvMaterializationKind::StructuredDepthwise) + return createDepthwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc); + if (state.problem.xHeight == 1 && state.problem.xWidth == 1 && state.problem.wHeight == 1 && state.problem.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 +3069,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); } @@ -3848,86 +3094,94 @@ static Value createCollectedConvOutput(ValueRange gemmRows, {0, 1, 2}, {3} }); - Value nchwOut = ONNXTransposeOp::create(rewriter, loc, outType, nhwcOut, rewriter.getI64ArrayAttr({0, 3, 1, 2})); + Value nchwOut = createLinalgTranspose(nhwcOut, outType, {0, 3, 1, 2}, rewriter, loc); spatial::SpatYieldOp::create(rewriter, loc, nchwOut); }); return collectComputeOp.getResult(0); } -static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b) { +static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, + Value x, + Value w, + Value b, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) { ConvLoweringState state; + state.diagnosticAnchor = convOp.getOperation(); state.x = x; state.w = w; state.b = b; - state.xType = cast(state.x.getType()); - state.wType = cast(state.w.getType()); - state.outType = cast(convOp.getY().getType()); + state.target = ⌖ + state.options = &options; + state.problem.xType = cast(state.x.getType()); + state.problem.wType = cast(state.w.getType()); + state.problem.outType = cast(convOp.getY().getType()); - if (!state.xType.hasStaticShape()) { + if (!state.problem.xType.hasStaticShape()) { pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv input"); return failure(); } - if (!state.wType.hasStaticShape()) { + if (!state.problem.wType.hasStaticShape()) { pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv weight"); return failure(); } - if (!state.outType.hasStaticShape()) { + if (!state.problem.outType.hasStaticShape()) { pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv result"); return failure(); } - if (state.xType.getRank() != 4) { - pim::emitUnsupportedRankDiagnostic(convOp, "conv input", state.xType.getRank(), {4}); + if (state.problem.xType.getRank() != 4) { + pim::emitUnsupportedRankDiagnostic(convOp, "conv input", state.problem.xType.getRank(), {4}); return failure(); } - if (state.wType.getRank() != 4) { - pim::emitUnsupportedRankDiagnostic(convOp, "conv weight", state.wType.getRank(), {4}); + if (state.problem.wType.getRank() != 4) { + pim::emitUnsupportedRankDiagnostic(convOp, "conv weight", state.problem.wType.getRank(), {4}); return failure(); } - if (state.outType.getRank() != 4) { - pim::emitUnsupportedRankDiagnostic(convOp, "conv result", state.outType.getRank(), {4}); + if (state.problem.outType.getRank() != 4) { + pim::emitUnsupportedRankDiagnostic(convOp, "conv result", state.problem.outType.getRank(), {4}); return failure(); } - state.group = convOp.getGroup(); - if (state.group < 1) { + state.problem.group = convOp.getGroup(); + if (state.problem.group < 1) { convOp.emitOpError("requires group >= 1 for Spatial lowering"); return failure(); } - state.batchSize = state.xType.getDimSize(0); - state.numChannelsIn = state.xType.getDimSize(1); - state.xHeight = state.xType.getDimSize(2); - state.xWidth = state.xType.getDimSize(3); - state.numChannelsOut = state.wType.getDimSize(0); - state.wHeight = state.wType.getDimSize(2); - state.wWidth = state.wType.getDimSize(3); - state.outHeight = state.outType.getDimSize(2); - state.outWidth = state.outType.getDimSize(3); - state.hasBias = + state.problem.batchSize = state.problem.xType.getDimSize(0); + state.problem.numChannelsIn = state.problem.xType.getDimSize(1); + state.problem.xHeight = state.problem.xType.getDimSize(2); + state.problem.xWidth = state.problem.xType.getDimSize(3); + state.problem.numChannelsOut = state.problem.wType.getDimSize(0); + state.problem.wHeight = state.problem.wType.getDimSize(2); + state.problem.wWidth = state.problem.wType.getDimSize(3); + state.problem.outHeight = state.problem.outType.getDimSize(2); + state.problem.outWidth = state.problem.outType.getDimSize(3); + state.problem.hasBias = state.b && !isa(state.b.getDefiningOp()) && !isZeroSplatHostConstant(state.b); - if (state.numChannelsIn % state.group != 0) { - convOp.emitOpError() << "requires input channels " << state.numChannelsIn << " to be divisible by group " - << state.group << " for Spatial lowering"; + if (state.problem.numChannelsIn % state.problem.group != 0) { + convOp.emitOpError() << "requires input channels " << state.problem.numChannelsIn << " to be divisible by group " + << state.problem.group << " for Spatial lowering"; return failure(); } - if (state.numChannelsOut % state.group != 0) { - convOp.emitOpError() << "requires output channels " << state.numChannelsOut << " to be divisible by group " - << state.group << " for Spatial lowering"; + if (state.problem.numChannelsOut % state.problem.group != 0) { + convOp.emitOpError() << "requires output channels " << state.problem.numChannelsOut << " to be divisible by group " + << state.problem.group << " for Spatial lowering"; return failure(); } - state.numChannelsInPerGroup = state.numChannelsIn / state.group; - state.numChannelsOutPerGroup = state.numChannelsOut / state.group; - if (state.wType.getDimSize(1) != state.numChannelsInPerGroup) { - convOp.emitOpError() << "requires grouped conv weight input channels " << state.wType.getDimSize(1) - << " to match input channels per group " << state.numChannelsInPerGroup + state.problem.numChannelsInPerGroup = state.problem.numChannelsIn / state.problem.group; + state.problem.numChannelsOutPerGroup = state.problem.numChannelsOut / state.problem.group; + if (state.problem.wType.getDimSize(1) != state.problem.numChannelsInPerGroup) { + convOp.emitOpError() << "requires grouped conv weight input channels " << state.problem.wType.getDimSize(1) + << " to match input channels per group " << state.problem.numChannelsInPerGroup << " for Spatial lowering"; return failure(); } - if (state.wType.getDimSize(0) != state.numChannelsOut) { - convOp.emitOpError() << "requires weight output channels " << state.wType.getDimSize(0) - << " to match result channels " << state.numChannelsOut << " for Spatial lowering"; + if (state.problem.wType.getDimSize(0) != state.problem.numChannelsOut) { + convOp.emitOpError() << "requires weight output channels " << state.problem.wType.getDimSize(0) + << " to match result channels " << state.problem.numChannelsOut << " for Spatial lowering"; return failure(); } @@ -3948,44 +3202,46 @@ static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, return failure(); } - state.strideHeight = getOptionalI64Attr(stridesAttr, 0, 1); - state.strideWidth = getOptionalI64Attr(stridesAttr, 1, 1); - state.dilationHeight = getOptionalI64Attr(dilationsAttr, 0, 1); - state.dilationWidth = getOptionalI64Attr(dilationsAttr, 1, 1); - state.padHeightBegin = 0; - state.padHeightEnd = 0; - state.padWidthBegin = 0; - state.padWidthEnd = 0; + state.problem.strideHeight = getOptionalI64Attr(stridesAttr, 0, 1); + state.problem.strideWidth = getOptionalI64Attr(stridesAttr, 1, 1); + state.problem.dilationHeight = getOptionalI64Attr(dilationsAttr, 0, 1); + state.problem.dilationWidth = getOptionalI64Attr(dilationsAttr, 1, 1); + state.problem.padHeightBegin = 0; + state.problem.padHeightEnd = 0; + state.problem.padWidthBegin = 0; + state.problem.padWidthEnd = 0; if (padsAttr) { - state.padHeightBegin = getI64Attr(*padsAttr, 0); - state.padWidthBegin = getI64Attr(*padsAttr, 1); - state.padHeightEnd = getI64Attr(*padsAttr, 2); - state.padWidthEnd = getI64Attr(*padsAttr, 3); + state.problem.padHeightBegin = getI64Attr(*padsAttr, 0); + state.problem.padWidthBegin = getI64Attr(*padsAttr, 1); + state.problem.padHeightEnd = getI64Attr(*padsAttr, 2); + state.problem.padWidthEnd = getI64Attr(*padsAttr, 3); + classifyConvProblem(state.problem); return state; } const auto autoPad = convOp.getAutoPad(); if (autoPad == "SAME_UPPER" || autoPad == "SAME_LOWER") { - const int64_t effectiveKernelH = (state.wHeight - 1) * state.dilationHeight + 1; - const int64_t effectiveKernelW = (state.wWidth - 1) * state.dilationWidth + 1; + const int64_t effectiveKernelH = (state.problem.wHeight - 1) * state.problem.dilationHeight + 1; + const int64_t effectiveKernelW = (state.problem.wWidth - 1) * state.problem.dilationWidth + 1; const int64_t totalPadH = - std::max(static_cast(0), (state.outHeight - 1) * state.strideHeight + effectiveKernelH - state.xHeight); + std::max(static_cast(0), (state.problem.outHeight - 1) * state.problem.strideHeight + effectiveKernelH - state.problem.xHeight); const int64_t totalPadW = - std::max(static_cast(0), (state.outWidth - 1) * state.strideWidth + effectiveKernelW - state.xWidth); + std::max(static_cast(0), (state.problem.outWidth - 1) * state.problem.strideWidth + effectiveKernelW - state.problem.xWidth); if (autoPad == "SAME_UPPER") { - state.padHeightBegin = totalPadH / 2; - state.padHeightEnd = totalPadH - state.padHeightBegin; - state.padWidthBegin = totalPadW / 2; - state.padWidthEnd = totalPadW - state.padWidthBegin; + state.problem.padHeightBegin = totalPadH / 2; + state.problem.padHeightEnd = totalPadH - state.problem.padHeightBegin; + state.problem.padWidthBegin = totalPadW / 2; + state.problem.padWidthEnd = totalPadW - state.problem.padWidthBegin; } else { - state.padHeightEnd = totalPadH / 2; - state.padHeightBegin = totalPadH - state.padHeightEnd; - state.padWidthEnd = totalPadW / 2; - state.padWidthBegin = totalPadW - state.padWidthEnd; + state.problem.padHeightEnd = totalPadH / 2; + state.problem.padHeightBegin = totalPadH - state.problem.padHeightEnd; + state.problem.padWidthEnd = totalPadW / 2; + state.problem.padWidthBegin = totalPadW - state.problem.padWidthEnd; } + classifyConvProblem(state.problem); return state; } @@ -3994,50 +3250,64 @@ static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, return failure(); } + classifyConvProblem(state.problem); return state; } -static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, ONNXConvOpAdaptor convOpAdaptor) { - return analyzeConvLoweringState(convOp, convOpAdaptor.getX(), convOpAdaptor.getW(), convOpAdaptor.getB()); +static FailureOr analyzeConvLoweringState(ONNXConvOp convOp, + ONNXConvOpAdaptor convOpAdaptor, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) { + return analyzeConvLoweringState( + convOp, convOpAdaptor.getX(), convOpAdaptor.getW(), convOpAdaptor.getB(), target, options); } -static FailureOr analyzeConvLoweringState(spatial::SpatConv2DPlanOp planOp) { +static FailureOr analyzeConvLoweringState( + spatial::SpatConv2DPlanOp planOp, + Value x, + Value w, + Value b, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions* options = nullptr) { ConvLoweringState state; - state.x = planOp.getInput(); - state.w = planOp.getWeight(); - state.b = planOp.getBias() ? planOp.getBias() : Value(); - state.xType = dyn_cast(state.x.getType()); - state.wType = dyn_cast(state.w.getType()); - state.outType = dyn_cast(planOp.getOutput().getType()); + state.diagnosticAnchor = planOp.getOperation(); + state.x = x; + state.w = w; + state.b = b; + state.target = ⌖ + state.options = options; + state.problem.xType = dyn_cast(state.x.getType()); + state.problem.wType = dyn_cast(state.w.getType()); + state.problem.outType = dyn_cast(planOp.getOutput().getType()); - if (!state.xType || !state.wType || !state.outType) + if (!state.problem.xType || !state.problem.wType || !state.problem.outType) return planOp.emitOpError("requires ranked tensor input, weight, and output"), failure(); - if (!state.xType.hasStaticShape() || !state.wType.hasStaticShape() || !state.outType.hasStaticShape()) + if (!state.problem.xType.hasStaticShape() || !state.problem.wType.hasStaticShape() || !state.problem.outType.hasStaticShape()) return planOp.emitOpError("requires static input, weight, and output shapes"), failure(); - if (state.xType.getRank() != 4 || state.wType.getRank() != 4 || state.outType.getRank() != 4) + if (state.problem.xType.getRank() != 4 || state.problem.wType.getRank() != 4 || state.problem.outType.getRank() != 4) return planOp.emitOpError("requires rank-4 input, weight, and output tensors"), failure(); - state.group = planOp.getGroup(); - if (state.group < 1) + state.problem.group = planOp.getGroup(); + if (state.problem.group < 1) return planOp.emitOpError("requires group >= 1"), failure(); - state.batchSize = state.xType.getDimSize(0); - state.numChannelsIn = state.xType.getDimSize(1); - state.xHeight = state.xType.getDimSize(2); - state.xWidth = state.xType.getDimSize(3); - state.numChannelsOut = state.wType.getDimSize(0); - state.wHeight = state.wType.getDimSize(2); - state.wWidth = state.wType.getDimSize(3); - state.outHeight = state.outType.getDimSize(2); - state.outWidth = state.outType.getDimSize(3); - state.hasBias = planOp.getBias() && !isZeroSplatHostConstant(planOp.getBias()); + state.problem.batchSize = state.problem.xType.getDimSize(0); + state.problem.numChannelsIn = state.problem.xType.getDimSize(1); + state.problem.xHeight = state.problem.xType.getDimSize(2); + state.problem.xWidth = state.problem.xType.getDimSize(3); + state.problem.numChannelsOut = state.problem.wType.getDimSize(0); + state.problem.wHeight = state.problem.wType.getDimSize(2); + state.problem.wWidth = state.problem.wType.getDimSize(3); + state.problem.outHeight = state.problem.outType.getDimSize(2); + state.problem.outWidth = state.problem.outType.getDimSize(3); + state.problem.hasBias = state.b && !isZeroSplatHostConstant(state.b); - if (state.numChannelsIn % state.group != 0 || state.numChannelsOut % state.group != 0) + if (state.problem.numChannelsIn % state.problem.group != 0 || state.problem.numChannelsOut % state.problem.group != 0) return planOp.emitOpError("requires input and output channels divisible by group"), failure(); - state.numChannelsInPerGroup = state.numChannelsIn / state.group; - state.numChannelsOutPerGroup = state.numChannelsOut / state.group; - if (state.wType.getDimSize(1) != state.numChannelsInPerGroup) + state.problem.numChannelsInPerGroup = state.problem.numChannelsIn / state.problem.group; + state.problem.numChannelsOutPerGroup = state.problem.numChannelsOut / state.problem.group; + if (state.problem.wType.getDimSize(1) != state.problem.numChannelsInPerGroup) return planOp.emitOpError("requires grouped conv weight channels to match input channels per group"), failure(); auto pads = planOp.getPads(); @@ -4046,71 +3316,94 @@ static FailureOr analyzeConvLoweringState(spatial::SpatConv2D if (pads.size() != 4 || strides.size() != 2 || dilations.size() != 2) return planOp.emitOpError("requires 4 pads, 2 strides, and 2 dilations"), failure(); - state.padHeightBegin = pads[0]; - state.padWidthBegin = pads[1]; - state.padHeightEnd = pads[2]; - state.padWidthEnd = pads[3]; - state.strideHeight = strides[0]; - state.strideWidth = strides[1]; - state.dilationHeight = dilations[0]; - state.dilationWidth = dilations[1]; + state.problem.padHeightBegin = pads[0]; + state.problem.padWidthBegin = pads[1]; + state.problem.padHeightEnd = pads[2]; + state.problem.padWidthEnd = pads[3]; + state.problem.strideHeight = strides[0]; + state.problem.strideWidth = strides[1]; + state.problem.dilationHeight = dilations[0]; + state.problem.dilationWidth = dilations[1]; + classifyConvProblem(state.problem); return state; } -static FailureOr resolveRequestedConvLoweringStrategy(Operation* op) { - if (!useExperimentalConvImpl) - return pimConvLowering.getValue(); - - if (pimConvLowering != PimConvLoweringAuto && pimConvLowering != PimConvLoweringPackedIm2Col) { - op->emitOpError() << "--use-experimental-conv-impl conflicts with --pim-conv-lowering=" - << stringifyConvLoweringStrategy(pimConvLowering); - return failure(); - } - return PimConvLoweringPackedIm2Col; +static FailureOr analyzeConvLoweringState( + spatial::SpatConv2DPlanOp planOp, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions* options = nullptr) { + return analyzeConvLoweringState( + planOp, planOp.getInput(), planOp.getWeight(), + planOp.getBias() ? planOp.getBias() : Value(), target, options); } -static LogicalResult verifyForcedConvLoweringStrategy(Operation* op, - const ConvGeometry& geo, - PimConvLoweringType strategy) { - switch (strategy) { - case PimConvLoweringAuto: - case PimConvLoweringLegacy: - return success(); - case PimConvLoweringDepthwise: - if (geo.isDepthwise) - return success(); - return op->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 op->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 op->emitOpError("forced streamed-patch Conv lowering requires K and C to each fit one crossbar"); - case PimConvLoweringStreamedPacked: +static FailureOr selectConvLoweringPlan( + Operation* op, const ConvLoweringState& state, bool reportPlanning) { + const spatial::ConvLoweringStrategy requested = state.planningOptions().forcedConvStrategy; + + auto tryStrategy = [&](spatial::ConvLoweringStrategy strategy) -> FailureOr { + return makeConvPlan(state.problem, strategy, state.targetInfo(), state.planningOptions()); + }; + auto reportAndReturn = [&](FailureOr candidate) -> FailureOr { + if (failed(candidate)) + return failure(); + if (reportPlanning) + recordConvLoweringReport(op, ConvLoweringReportPhase::Planning, + stringifyConvMaterializationKind(candidate->kind), "SEL", + state.planningOptions().reportConvLowering); + return *candidate; + }; + + if (requested == spatial::ConvLoweringStrategy::Auto) { + ConvGeometry geo = buildConvGeometry(state.problem, state.targetInfo()); + if (state.problem.isDepthwise) { + FailureOr depthwise = tryStrategy(spatial::ConvLoweringStrategy::Depthwise); + if (succeeded(depthwise) && depthwise::canUseStructuredRewrite(state)) + return reportAndReturn(std::move(depthwise)); + FailureOr legacy = tryStrategy(spatial::ConvLoweringStrategy::Legacy); + if (succeeded(legacy)) + return reportAndReturn(std::move(legacy)); + } + if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 + && geo.im2colElements <= state.planningOptions().convIm2colMaxElements) + if (auto candidate = tryStrategy(spatial::ConvLoweringStrategy::PackedIm2Col); succeeded(candidate)) + return reportAndReturn(std::move(candidate)); if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2) - return success(); - return op->emitOpError("forced streamed-packed Conv lowering requires K/C to fit and pack >= 2"); - case PimConvLoweringOutputChannelTiled: + if (auto candidate = tryStrategy(spatial::ConvLoweringStrategy::StreamedPacked); succeeded(candidate)) + return reportAndReturn(std::move(candidate)); + if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize) + if (auto candidate = tryStrategy(spatial::ConvLoweringStrategy::StreamedPatch); succeeded(candidate)) + return reportAndReturn(std::move(candidate)); if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize) - return success(); - return op->emitOpError("forced output-channel-tiled Conv lowering requires K <= X and C > X"); - case PimConvLoweringInputKTiled: + if (auto candidate = tryStrategy(spatial::ConvLoweringStrategy::OutputChannelTiled); succeeded(candidate)) + return reportAndReturn(std::move(candidate)); if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize) - return success(); - return op->emitOpError("forced input-k-tiled Conv lowering requires K > X and C <= X"); - case PimConvLoweringTiled2D: + if (auto candidate = tryStrategy(spatial::ConvLoweringStrategy::Legacy); succeeded(candidate)) + return reportAndReturn(std::move(candidate)); if (geo.k > geo.xbarSize && geo.c > geo.xbarSize) - return success(); - return op->emitOpError("forced tiled-2d Conv lowering requires K > X and C > X"); + if (auto candidate = tryStrategy(spatial::ConvLoweringStrategy::Tiled2D); succeeded(candidate)) + return reportAndReturn(std::move(candidate)); + op->emitOpError("has no applicable Conv lowering candidate for the injected Spatial target"); + return failure(); } - llvm_unreachable("unknown conv lowering strategy"); + + FailureOr candidate = tryStrategy(requested); + if (failed(candidate)) { + op->emitOpError() << "forced Conv lowering `" + << stringifyConvLoweringStrategy(requested) + << "` is not applicable to this Conv problem"; + return failure(); + } + if (reportPlanning) + recordConvLoweringReport(op, ConvLoweringReportPhase::Planning, + stringifyConvMaterializationKind(candidate->kind), "SEL", + state.planningOptions().reportConvLowering); + return *candidate; } static FailureOr lowerDenseSelectedConvPlan(Operation* op, const ConvLoweringState& state, - PimConvLoweringType strategy, + ConvMaterializationKind kind, PatternRewriter& rewriter, Location loc); @@ -4120,200 +3413,61 @@ static ConvLoweringState makeGroupedConvLoweringState(const ConvLoweringState& p Value groupB, RankedTensorType groupOutType); -static FailureOr buildConvValueForStrategy(Operation* op, - Location loc, - const ConvLoweringState& state, - const ConvLoweringDecision& decision, - const DistributedConvAnalysis& analysis, - ArrayRef distributedConsumers, - PatternRewriter& rewriter); +static FailureOr lowerConvByMaterialization(Operation* op, + Location loc, + const ConvLoweringState& state, + ConvMaterializationKind kind, + PatternRewriter& rewriter); static FailureOr buildGroupedConvValue(Operation* op, Location loc, const ConvLoweringState& state, - const ConvLoweringDecision& decision, + ConvMaterializationKind kind, PatternRewriter& rewriter); static FailureOr lowerGroupedSelectedConvPlan(Operation* op, const ConvLoweringState& state, - PimConvLoweringType strategy, + ConvMaterializationKind kind, PatternRewriter& rewriter, Location loc) { - ConvLoweringDecision decision {strategy, "", false, "", ""}; - return buildGroupedConvValue(op, loc, state, decision, rewriter); + return buildGroupedConvValue(op, loc, state, kind, rewriter); } static FailureOr lowerDenseSelectedConvPlan(Operation* op, const ConvLoweringState& state, - PimConvLoweringType strategy, + ConvMaterializationKind kind, 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 lowerConvByMaterialization(op, loc, state, kind, rewriter); } -static FailureOr buildConvValueForStrategy(Operation* op, - Location loc, - const ConvLoweringState& state, - const ConvLoweringDecision& decision, - const DistributedConvAnalysis& analysis, - ArrayRef distributedConsumers, - PatternRewriter& rewriter) { - (void)analysis; - const ConvGeometry geo = buildConvGeometry(state); - switch (decision.strategy) { - case PimConvLoweringDepthwise: { +static FailureOr lowerConvByMaterialization(Operation* op, + Location loc, + const ConvLoweringState& state, + ConvMaterializationKind kind, + PatternRewriter& rewriter) { + const ConvGeometry geo = buildConvGeometry(state.problem, state.targetInfo()); + switch (kind) { + case ConvMaterializationKind::StructuredDepthwise: { return depthwise::rewriteConv(op, state, rewriter, loc); } - case PimConvLoweringLegacy: - case PimConvLoweringPackedIm2Col: { - return standard::rewritePackedIm2ColConv(state, distributedConsumers, rewriter, loc); + case ConvMaterializationKind::PackedIm2Col: { + return standard::rewritePackedIm2ColConv(state, rewriter, loc); } - case PimConvLoweringStreamedPatch: - case PimConvLoweringOutputChannelTiled: - case PimConvLoweringTiled2D: { - return standard::rewriteStreamedConv(state, distributedConsumers, rewriter, loc, /*forcedPackFactor=*/1); + case ConvMaterializationKind::StreamedPatch: { + return standard::rewriteStreamedConv(state, rewriter, loc, /*forcedPackFactor=*/1); } - case PimConvLoweringInputKTiled: { - return standard::rewriteInputKTiledConv(state, distributedConsumers, rewriter, loc); + case ConvMaterializationKind::InputKTiled: { + return standard::rewriteInputKTiledConv(state, rewriter, loc); } - case PimConvLoweringStreamedPacked: { - return standard::rewriteStreamedConv(state, distributedConsumers, rewriter, loc, geo.pack); + case ConvMaterializationKind::StreamedPacked: { + return standard::rewriteStreamedConv(state, rewriter, loc, geo.pack); } - case PimConvLoweringAuto: - break; } - op->emitOpError("unexpected auto strategy at Conv lowering dispatch"); + op->emitOpError("unexpected Conv materialization kind"); 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, @@ -4327,34 +3481,35 @@ static ConvLoweringState makeGroupedConvLoweringState( state.x = groupX; state.w = groupW; state.b = groupB; - state.xType = cast(groupX.getType()); - state.wType = cast(groupW.getType()); - state.outType = groupOutType; - state.batchSize = state.xType.getDimSize(0); - state.numChannelsIn = state.xType.getDimSize(1); - state.xHeight = state.xType.getDimSize(2); - state.xWidth = state.xType.getDimSize(3); - state.numChannelsOut = state.wType.getDimSize(0); - state.wHeight = state.wType.getDimSize(2); - state.wWidth = state.wType.getDimSize(3); - state.outHeight = state.outType.getDimSize(2); - state.outWidth = state.outType.getDimSize(3); - state.group = 1; - state.numChannelsInPerGroup = state.numChannelsIn; - state.numChannelsOutPerGroup = state.numChannelsOut; - state.hasBias = static_cast(groupB); + state.problem.xType = cast(groupX.getType()); + state.problem.wType = cast(groupW.getType()); + state.problem.outType = groupOutType; + state.problem.batchSize = state.problem.xType.getDimSize(0); + state.problem.numChannelsIn = state.problem.xType.getDimSize(1); + state.problem.xHeight = state.problem.xType.getDimSize(2); + state.problem.xWidth = state.problem.xType.getDimSize(3); + state.problem.numChannelsOut = state.problem.wType.getDimSize(0); + state.problem.wHeight = state.problem.wType.getDimSize(2); + state.problem.wWidth = state.problem.wType.getDimSize(3); + state.problem.outHeight = state.problem.outType.getDimSize(2); + state.problem.outWidth = state.problem.outType.getDimSize(3); + state.problem.group = 1; + state.problem.numChannelsInPerGroup = state.problem.numChannelsIn; + state.problem.numChannelsOutPerGroup = state.problem.numChannelsOut; + state.problem.hasBias = static_cast(groupB); + classifyConvProblem(state.problem); return state; } static FailureOr buildGroupedConvValue(Operation* op, Location loc, const ConvLoweringState& state, - const ConvLoweringDecision& decision, + ConvMaterializationKind kind, PatternRewriter& rewriter) { - SmallVector xSlices = sliceTensor(state.x, /*axis=*/1, state.numChannelsInPerGroup, rewriter, loc); - SmallVector wSlices = sliceTensor(state.w, /*axis=*/0, state.numChannelsOutPerGroup, rewriter, loc); + SmallVector xSlices = sliceTensor(state.x, /*axis=*/1, state.problem.numChannelsInPerGroup, rewriter, loc); + SmallVector wSlices = sliceTensor(state.w, /*axis=*/0, state.problem.numChannelsOutPerGroup, rewriter, loc); SmallVector bSlices; - if (state.hasBias) { + if (state.problem.hasBias) { auto biasType = cast(state.b.getType()); int64_t biasAxis = -1; if (biasType.getRank() == 1) @@ -4366,29 +3521,25 @@ static FailureOr buildGroupedConvValue(Operation* op, << biasType.getRank(); return failure(); } - bSlices = sliceTensor(state.b, biasAxis, state.numChannelsOutPerGroup, rewriter, loc); + bSlices = sliceTensor(state.b, biasAxis, state.problem.numChannelsOutPerGroup, rewriter, loc); } - if (xSlices.size() != static_cast(state.group) || wSlices.size() != static_cast(state.group) - || (state.hasBias && bSlices.size() != static_cast(state.group))) { + if (xSlices.size() != static_cast(state.problem.group) || wSlices.size() != static_cast(state.problem.group) + || (state.problem.hasBias && bSlices.size() != static_cast(state.problem.group))) { op->emitOpError("failed to partition grouped convolution operands for Spatial lowering"); return failure(); } SmallVector groupResults; - groupResults.reserve(state.group); + groupResults.reserve(state.problem.group); auto groupOutType = RankedTensorType::get( - {state.batchSize, state.numChannelsOutPerGroup, state.outHeight, state.outWidth}, state.outType.getElementType()); - for (int64_t groupId = 0; groupId < state.group; groupId++) { + {state.problem.batchSize, state.problem.numChannelsOutPerGroup, state.problem.outHeight, state.problem.outWidth}, state.problem.outType.getElementType()); + for (int64_t groupId = 0; groupId < state.problem.group; groupId++) { Value groupX = xSlices[groupId]; Value groupW = wSlices[groupId]; - Value groupB = state.hasBias ? bSlices[groupId] : Value(); + Value groupB = state.problem.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 = lowerConvByMaterialization(op, loc, groupState, kind, rewriter); if (failed(groupResult)) return failure(); groupResults.push_back(*groupResult); @@ -4397,37 +3548,27 @@ static FailureOr buildGroupedConvValue(Operation* op, if (llvm::all_of(groupResults, isCompileTimeComputable)) return createSpatConcat(rewriter, loc, /*axis=*/1, groupResults); - auto concatCompute = createSpatCompute(rewriter, loc, TypeRange {state.outType}, {}, groupResults, [&](ValueRange args) { + auto concatCompute = createSpatCompute(rewriter, loc, TypeRange {state.problem.outType}, {}, groupResults, [&](ValueRange args) { spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/1, args)); }); 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 LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp, ONNXConvOpAdaptor convOpAdaptor, ConversionPatternRewriter& rewriter) const { - FailureOr state = analyzeConvLoweringState(convOp, convOpAdaptor); + FailureOr state = analyzeConvLoweringState( + convOp, convOpAdaptor, target, options); if (failed(state)) return failure(); SmallVector pads { - state->padHeightBegin, state->padWidthBegin, state->padHeightEnd, state->padWidthEnd}; - SmallVector strides {state->strideHeight, state->strideWidth}; - SmallVector dilations {state->dilationHeight, state->dilationWidth}; - Value bias = state->hasBias ? convOpAdaptor.getB() : Value(); + state->problem.padHeightBegin, state->problem.padWidthBegin, state->problem.padHeightEnd, state->problem.padWidthEnd}; + SmallVector strides {state->problem.strideHeight, state->problem.strideWidth}; + SmallVector dilations {state->problem.dilationHeight, state->problem.dilationWidth}; + Value bias = state->problem.hasBias ? convOpAdaptor.getB() : Value(); auto convPlan = spatial::SpatConv2DPlanOp::create(rewriter, convOp.getLoc(), convOp.getY().getType(), @@ -4437,123 +3578,122 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp, rewriter.getDenseI64ArrayAttr(pads), rewriter.getDenseI64ArrayAttr(strides), rewriter.getDenseI64ArrayAttr(dilations), - rewriter.getI64IntegerAttr(state->group), - rewriter.getStringAttr("nchw")); + rewriter.getI64IntegerAttr(state->problem.group), + spatial::getNCHWLayout(rewriter.getContext())); rewriter.replaceOp(convOp, convPlan.getResult()); return success(); } -void populateConvPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.insert(ctx); } +void populateConvPatterns(RewritePatternSet& patterns, + MLIRContext* ctx, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) { + patterns.insert(ctx, target, options); +} -LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) { - FailureOr state = analyzeConvLoweringState(planOp); +LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions* options) { + FailureOr state = analyzeConvLoweringState(planOp, target, options); if (failed(state)) return failure(); - if (state->group != 1 || state->batchSize != 1) + if (state->problem.group != 1 || state->problem.batchSize != 1) return failure(); - if (state->outType.getRank() != 4 || !state->outType.hasStaticShape()) + if (state->problem.outType.getRank() != 4 || !state->problem.outType.hasStaticShape()) return failure(); if (!getHostConstDenseElementsAttr(state->w)) return failure(); - if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType)) + if (state->problem.hasBias && !isSupportedBiasAddValue(state->b, state->problem.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); + ConvGeometry geometry = buildConvGeometry(state->problem, state->targetInfo()); 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 plan = + selectConvLoweringPlan(planOp.getOperation(), *state, /*reportPlanning=*/true); + if (failed(plan)) return failure(); - switch (decision.strategy) { - case PimConvLoweringLegacy: - case PimConvLoweringDepthwise: - case PimConvLoweringPackedIm2Col: - case PimConvLoweringStreamedPatch: - case PimConvLoweringOutputChannelTiled: - case PimConvLoweringTiled2D: - case PimConvLoweringStreamedPacked: + switch (plan->kind) { + case ConvMaterializationKind::StructuredDepthwise: + case ConvMaterializationKind::PackedIm2Col: + case ConvMaterializationKind::StreamedPatch: + case ConvMaterializationKind::StreamedPacked: return success(); - case PimConvLoweringAuto: - case PimConvLoweringInputKTiled: + case ConvMaterializationKind::InputKTiled: return failure(); } - llvm_unreachable("unknown conv lowering strategy"); + llvm_unreachable("unknown Conv materialization kind"); } -LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) { - FailureOr state = analyzeConvLoweringState(planOp); +LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions* options) { + FailureOr state = analyzeConvLoweringState(planOp, target, options); if (failed(state)) return failure(); + FailureOr plan = + selectConvLoweringPlan(planOp.getOperation(), *state, /*reportPlanning=*/true); + if (failed(plan)) + return failure(); + if (plan->kind == ConvMaterializationKind::StructuredDepthwise) + return canConsumeDepthwiseRowStrip(*state) ? success() : failure(); StringRef failureReason; return canConsumePixelMajorRowStripFragments(*state, failureReason) ? success() : failure(); } FailureOr lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, + Value input, + Value weight, + Value bias, std::optional rowStripInput, bool emitRowStripLayout, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options, PatternRewriter& rewriter) { - FailureOr state = analyzeConvLoweringState(planOp); + FailureOr state = analyzeConvLoweringState( + planOp, input, weight, bias, target, &options); if (failed(state)) return failure(); - FailureOr requestedStrategy = resolveRequestedConvLoweringStrategy(planOp.getOperation()); - if (failed(requestedStrategy)) + FailureOr plan = + selectConvLoweringPlan(planOp.getOperation(), *state, /*reportPlanning=*/false); + if (failed(plan)) 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))) - return failure(); + auto reportRealization = [&](StringRef implementation) { + recordConvLoweringReport( + planOp.getOperation(), ConvLoweringReportPhase::Realization, + stringifyConvMaterializationKind(plan->kind), + implementation, state->planningOptions().reportConvLowering); + }; if (emitRowStripLayout) { if (rowStripInput) { - if (failed(canConsumeAndProduceRowStrip(planOp))) + if (failed(canConsumeAndProduceRowStrip(planOp, target, &options))) 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()); + reportRealization(convRowStripInputImplementation(*state, plan->kind)); + return createConvOutputFromRowStripInput( + *state, *rowStripInput, plan->kind, rewriter, planOp.getLoc()); } - if (failed(canLowerConvPlanToRowStrip(planOp))) + if (failed(canLowerConvPlanToRowStrip(planOp, target, &options))) return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure(); + reportRealization("RSD"); FailureOr rowStripStorage = createRowStripConvOutputFromDenseInput(*state, rewriter, planOp.getLoc()); if (failed(rowStripStorage)) return planOp.emitOpError("failed to build row-strip fragment storage for the selected Conv plan"), failure(); return *rowStripStorage; } - if (decision.strategy == PimConvLoweringDepthwise) - return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, decision.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()); + reportRealization(convLoweringImplementation(plan->kind)); + if (plan->kind == ConvMaterializationKind::StructuredDepthwise) + return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, plan->kind, rewriter, planOp.getLoc()); + if (state->problem.group != 1) + return lowerGroupedSelectedConvPlan(planOp.getOperation(), *state, plan->kind, rewriter, planOp.getLoc()); + return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, plan->kind, rewriter, planOp.getLoc()); } } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.cpp index 4b1fa71..a674b81 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.cpp @@ -1,49 +1,133 @@ #include "ConvGeometry.hpp" #include +#include #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" namespace onnx_mlir { +namespace { + +static const ONNXToSpatialPlanningOptions& defaultPlanningOptions() { + static const ONNXToSpatialPlanningOptions options { + std::numeric_limits::max(), + std::numeric_limits::max(), + spatial::ConvLoweringStrategy::Auto, + false, + }; + return options; +} + +} // namespace + +const ONNXToSpatialPlanningOptions& ConvLoweringState::planningOptions() const { + return options ? *options : defaultPlanningOptions(); +} + bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup) { return group == numChannelsIn && numChannelsInPerGroup == 1 && numChannelsOut % group == 0; } -ConvGeometry buildConvGeometry(const ConvLoweringState& state) { +void classifyConvProblem(ConvProblem& problem) { + problem.isDepthwise = isDepthwiseConv( + problem.group, problem.numChannelsIn, problem.numChannelsOut, + problem.numChannelsInPerGroup); + problem.isGrouped = problem.group > 1; + problem.isPointwise = problem.wHeight == 1 && problem.wWidth == 1 + && problem.strideHeight == 1 && problem.strideWidth == 1 + && problem.dilationHeight == 1 && problem.dilationWidth == 1 + && problem.padHeightBegin == 0 && problem.padHeightEnd == 0 + && problem.padWidthBegin == 0 && problem.padWidthEnd == 0; +} + +ConvGeometry buildConvGeometry(const ConvProblem& problem, + const spatial::SpatialTargetResources& target) { ConvGeometry geo { - state.batchSize, - state.numChannelsIn, - state.xHeight, - state.xWidth, - state.numChannelsOut, - state.wHeight, - state.wWidth, - state.outHeight, - state.outWidth, - state.group, - state.numChannelsInPerGroup, - state.numChannelsOutPerGroup, - state.numChannelsInPerGroup * state.wHeight * state.wWidth, - state.numChannelsOutPerGroup, - state.batchSize * state.outHeight * state.outWidth, - static_cast(crossbarSize.getValue()), + problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth, + problem.numChannelsOutPerGroup, + problem.batchSize * problem.outHeight * problem.outWidth, + static_cast(target.matrixShape.rows), + static_cast(target.matrixUnitsPerProcessor), 1, 0, - state.hasBias, - isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup), }; geo.pack = std::max(1, geo.xbarSize / std::max(geo.k, geo.c)); geo.im2colElements = static_cast(std::max(0, geo.p)) * static_cast(std::max(0, geo.k)); return geo; } -uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor) { +static ConvMaterializationKind getMaterializationKind( + spatial::ConvLoweringStrategy strategy) { + switch (strategy) { + case spatial::ConvLoweringStrategy::Depthwise: + return ConvMaterializationKind::StructuredDepthwise; + case spatial::ConvLoweringStrategy::Legacy: + case spatial::ConvLoweringStrategy::PackedIm2Col: + return ConvMaterializationKind::PackedIm2Col; + case spatial::ConvLoweringStrategy::StreamedPatch: + case spatial::ConvLoweringStrategy::OutputChannelTiled: + case spatial::ConvLoweringStrategy::Tiled2D: + return ConvMaterializationKind::StreamedPatch; + case spatial::ConvLoweringStrategy::StreamedPacked: + return ConvMaterializationKind::StreamedPacked; + case spatial::ConvLoweringStrategy::InputKTiled: + return ConvMaterializationKind::InputKTiled; + case spatial::ConvLoweringStrategy::Auto: + break; + } + llvm_unreachable("auto is not a Conv materialization kind"); +} + +static bool fitsSingleCrossbar(const ConvGeometry& geo) { + return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize; +} + +static bool fitsPackedIm2Col(const ConvGeometry& geo, + const ONNXToSpatialPlanningOptions& options) { + return fitsSingleCrossbar(geo) && geo.pack >= 2 + && geo.im2colElements <= options.convIm2colMaxElements; +} + +mlir::FailureOr makeConvPlan(const ConvProblem& problem, + spatial::ConvLoweringStrategy strategy, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options) { + ConvGeometry geo = buildConvGeometry(problem, target); + auto plan = [&]() { return ConvPlan {getMaterializationKind(strategy)}; }; + auto ifApplicable = [&](bool applicable) -> mlir::FailureOr { + return applicable ? mlir::FailureOr(plan()) : mlir::FailureOr(mlir::failure()); + }; + switch (strategy) { + case spatial::ConvLoweringStrategy::Auto: + return mlir::failure(); + case spatial::ConvLoweringStrategy::Legacy: + return plan(); + case spatial::ConvLoweringStrategy::Depthwise: + return ifApplicable(problem.isDepthwise); + case spatial::ConvLoweringStrategy::PackedIm2Col: + return ifApplicable(fitsPackedIm2Col(geo, options)); + case spatial::ConvLoweringStrategy::StreamedPatch: + return ifApplicable(fitsSingleCrossbar(geo)); + case spatial::ConvLoweringStrategy::StreamedPacked: + return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2); + case spatial::ConvLoweringStrategy::OutputChannelTiled: + return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize); + case spatial::ConvLoweringStrategy::InputKTiled: + return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize); + case spatial::ConvLoweringStrategy::Tiled2D: + return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize); + } + llvm_unreachable("unknown Conv lowering strategy"); +} + +uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, + int64_t packFactor, + const ONNXToSpatialPlanningOptions& options) { const uint64_t patchElements = static_cast(std::max(1, geo.k)); - uint64_t chunkPositions = std::max(1, pimConvIm2colMaxElements / patchElements); + uint64_t chunkPositions = std::max(1, options.convIm2colMaxElements / patchElements); chunkPositions = std::min(chunkPositions, static_cast(std::max(1, geo.p))); - chunkPositions = std::min(chunkPositions, std::max(1, pimConvStreamChunkPositions)); + chunkPositions = std::min(chunkPositions, std::max(1, options.convStreamChunkPositions)); if (packFactor > 1 && chunkPositions > static_cast(packFactor)) { chunkPositions -= chunkPositions % static_cast(packFactor); @@ -52,24 +136,26 @@ uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor) return std::max(1, chunkPositions); } -RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state) { - const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin; +RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem) { + const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin; const int64_t rawEnd = - (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1; - return {std::max(0, rawBegin), std::min(state.xHeight, rawEnd)}; + (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin + + problem.dilationHeight * (problem.wHeight - 1) + 1; + return {std::max(0, rawBegin), std::min(problem.xHeight, rawEnd)}; } -ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state) { +ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem) { ConvRowDemand demand; demand.outputRows = outputRows; - demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, state); + demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, problem); demand.acquiredInputRows = demand.neededInputRows; - const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin; + const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin; const int64_t rawEnd = - (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1; + (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin + + problem.dilationHeight * (problem.wHeight - 1) + 1; demand.topHaloRows = std::max(0, -rawBegin); - demand.bottomHaloRows = std::max(0, rawEnd - state.xHeight); + demand.bottomHaloRows = std::max(0, rawEnd - problem.xHeight); demand.acquiredInputRows = demand.neededInputRows; return demand; } diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp index 60564c6..a712d6d 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp @@ -3,14 +3,19 @@ #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Value.h" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" + #include +namespace mlir { +class Operation; +} // namespace mlir + namespace onnx_mlir { -struct ConvLoweringState { - mlir::Value x; - mlir::Value w; - mlir::Value b; +struct ConvProblem { mlir::RankedTensorType xType; mlir::RankedTensorType wType; mlir::RankedTensorType outType; @@ -35,29 +40,32 @@ struct ConvLoweringState { int64_t dilationHeight; int64_t dilationWidth; bool hasBias; + bool isDepthwise = false; + bool isGrouped = false; + bool isPointwise = false; +}; + +struct ConvLoweringState { + ConvProblem problem; + mlir::Operation* diagnosticAnchor = nullptr; + mlir::Value x; + mlir::Value w; + mlir::Value b; + const spatial::SpatialTargetResources* target = nullptr; + const ONNXToSpatialPlanningOptions* options = nullptr; + + const spatial::SpatialTargetResources& targetInfo() const { return *target; } + const ONNXToSpatialPlanningOptions& planningOptions() const; }; struct ConvGeometry { - int64_t batchSize; - int64_t numChannelsIn; - int64_t xHeight; - int64_t xWidth; - int64_t numChannelsOut; - int64_t wHeight; - int64_t wWidth; - int64_t outHeight; - int64_t outWidth; - int64_t group; - int64_t numChannelsInPerGroup; - int64_t numChannelsOutPerGroup; int64_t k; int64_t c; int64_t p; int64_t xbarSize; + int64_t matrixUnitsPerProcessor; int64_t pack; uint64_t im2colElements; - bool hasBias; - bool isDepthwise; }; struct RowInterval { @@ -73,14 +81,36 @@ struct ConvRowDemand { int64_t bottomHaloRows = 0; }; +enum class ConvMaterializationKind : uint8_t { + StructuredDepthwise, + PackedIm2Col, + StreamedPatch, + StreamedPacked, + InputKTiled, +}; + +struct ConvPlan { + ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col; +}; + bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup); -ConvGeometry buildConvGeometry(const ConvLoweringState& state); +void classifyConvProblem(ConvProblem& problem); -uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor); +ConvGeometry buildConvGeometry(const ConvProblem& problem, + const spatial::SpatialTargetResources& target); -RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state); +mlir::FailureOr makeConvPlan(const ConvProblem& problem, + spatial::ConvLoweringStrategy strategy, + const spatial::SpatialTargetResources& target, + const ONNXToSpatialPlanningOptions& options); -ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state); +uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, + int64_t packFactor, + const ONNXToSpatialPlanningOptions& options); + +RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem); + +ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem); } // 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..dd666de 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp @@ -31,7 +31,7 @@ struct SiluToSpatialPlan : OpRewritePattern { return failure(); auto plan = spatial::SpatSiluPlanOp::create( - rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, rewriter.getStringAttr("nchw")); + rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, spatial::getNCHWLayout(rewriter.getContext())); rewriter.replaceOp(mulOp, plan.getResult()); rewriter.eraseOp(sigmoidOp); return success(); @@ -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 = op.getA().getDefiningOp(); + Value scalar = adaptor.getB(); + if (!blueprint) { + blueprint = op.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, @@ -210,14 +260,16 @@ struct AddToSpatialCompute : OpConversionPattern { classifyBiasAddPlanCandidate(adaptor.getA(), adaptor.getB(), resultType); if (succeeded(candidate)) { auto plan = spatial::SpatBiasAddPlanOp::create( - rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, rewriter.getStringAttr("nchw")); + rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, + spatial::getNCHWLayout(rewriter.getContext())); rewriter.replaceOp(op, plan.getResult()); return success(); } if (resultType.getRank() == 4 && adaptor.getA().getType() == resultType && adaptor.getB().getType() == resultType) { auto plan = spatial::SpatAddPlanOp::create( - rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(), rewriter.getStringAttr("nchw")); + rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(), + spatial::getNCHWLayout(rewriter.getContext())); rewriter.replaceOp(op, plan.getResult()); return success(); } @@ -246,6 +298,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/Gemm.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp index 082ef04..ac4ef02 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp @@ -1,5 +1,6 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.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/BuiltinTypes.h" @@ -21,6 +22,9 @@ #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -31,7 +35,7 @@ namespace onnx_mlir { namespace { static FailureOr -materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewriter& rewriter, Location loc) { +materializeScaledConstantTensor(Value value, float factor, PatternRewriter& rewriter, Location loc) { if (factor == 1.0f) return value; @@ -57,7 +61,12 @@ materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewr } static Value createGemmBatchKOffset( - Value lane, int64_t numOutRows, int64_t numKSlices, ConversionPatternRewriter& rewriter, Location loc) { + Value lane, + int64_t numOutRows, + int64_t numKSlices, + int64_t xbarSize, + PatternRewriter& rewriter, + Location loc) { if (numKSlices == 1) return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0); @@ -65,7 +74,7 @@ static Value createGemmBatchKOffset( AffineExpr d0 = getAffineDimExpr(0, context); return createOrFoldAffineApply(rewriter, loc, - (d0.floorDiv(numOutRows) % numKSlices) * crossbarSize.getValue(), + (d0.floorDiv(numOutRows) % numKSlices) * xbarSize, ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp()); } @@ -74,7 +83,8 @@ static Value createGemmBatchHOffset(Value lane, int64_t numOutRows, int64_t numKSlices, int64_t numOutHSlices, - ConversionPatternRewriter& rewriter, + int64_t xbarSize, + PatternRewriter& rewriter, Location loc) { if (numOutHSlices == 1) return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0); @@ -83,14 +93,14 @@ static Value createGemmBatchHOffset(Value lane, AffineExpr d0 = getAffineDimExpr(0, context); return createOrFoldAffineApply(rewriter, loc, - d0.floorDiv(numOutRows * numKSlices) * crossbarSize.getValue(), + d0.floorDiv(numOutRows * numKSlices) * xbarSize, ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp()); } static FailureOr materializePaddedConstantMatrix(Value value, RankedTensorType resultType, - ConversionPatternRewriter& rewriter, + PatternRewriter& rewriter, Location loc) { auto sourceType = cast(value.getType()); if (sourceType == resultType) @@ -121,7 +131,7 @@ static FailureOr materializePaddedConstantMatrix(Value value, static FailureOr materializePaddedBroadcastedConstantTensor(Value value, RankedTensorType resultType, int64_t unpaddedColumns, - ConversionPatternRewriter& rewriter, + PatternRewriter& rewriter, Location loc) { auto denseAttr = getHostConstDenseElementsAttr(value); if (!denseAttr) @@ -187,7 +197,7 @@ static FailureOr materializePaddedBroadcastedConstantTensor(Value value, static FailureOr prepareBias(Value c, RankedTensorType outType, RankedTensorType paddedOutType, - ConversionPatternRewriter& rewriter, + PatternRewriter& rewriter, Location loc) { auto cType = cast(c.getType()); if (!cType.hasStaticShape()) @@ -203,9 +213,15 @@ static FailureOr prepareBias(Value c, } static Value extractATile( - Value a, Value row, Value kOffset, RankedTensorType aTileType, ConversionPatternRewriter& rewriter, Location loc) { + Value a, + Value row, + Value kOffset, + RankedTensorType aTileType, + int64_t xbarSize, + PatternRewriter& rewriter, + Location loc) { SmallVector offsets {row, kOffset}; - SmallVector sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())}; + SmallVector sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)}; SmallVector strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; return tensor::ExtractSliceOp::create(rewriter, loc, aTileType, a, offsets, sizes, strides).getResult(); @@ -219,7 +235,8 @@ static FailureOr createVmmBatch(Value a, int64_t numOutRows, int64_t numKSlices, int64_t numOutHSlices, - ConversionPatternRewriter& rewriter, + int64_t xbarSize, + PatternRewriter& rewriter, Location loc) { const int64_t laneCount = partialPiecesType.getDimSize(0); auto batchOp = createSpatComputeBatch( @@ -232,21 +249,21 @@ static FailureOr createVmmBatch(Value a, [&](detail::SpatComputeBatchBodyArgs args) { Value row = onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutRows, rewriter.getInsertionBlock()->getParentOp()); - Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, rewriter, loc); - Value hOffset = createGemmBatchHOffset(args.lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc); + Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, xbarSize, rewriter, loc); + Value hOffset = createGemmBatchHOffset( + args.lane, numOutRows, numKSlices, numOutHSlices, xbarSize, rewriter, loc); auto aTileType = - RankedTensorType::get({1, static_cast(crossbarSize.getValue())}, aType.getElementType()); + RankedTensorType::get({1, xbarSize}, aType.getElementType()); auto bTileType = RankedTensorType::get( - {static_cast(crossbarSize.getValue()), static_cast(crossbarSize.getValue())}, + {xbarSize, xbarSize}, paddedBType.getElementType()); auto pieceType = - RankedTensorType::get({1, static_cast(crossbarSize.getValue())}, partialPiecesType.getElementType()); - Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, rewriter, loc); + RankedTensorType::get({1, xbarSize}, partialPiecesType.getElementType()); + Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, xbarSize, rewriter, loc); SmallVector bOffsets {kOffset, hOffset}; - SmallVector bSizes {rewriter.getIndexAttr(crossbarSize.getValue()), - rewriter.getIndexAttr(crossbarSize.getValue())}; + SmallVector bSizes {rewriter.getIndexAttr(xbarSize), rewriter.getIndexAttr(xbarSize)}; SmallVector unitStrides = getUnitStrides(rewriter, 2); Value bTile = extractStaticSliceOrIdentity( rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides); @@ -260,7 +277,7 @@ static FailureOr createVmmBatch(Value a, } static Value extractDynamicGemmBColumn( - Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) { + Value matrix, Value column, RankedTensorType vectorType, PatternRewriter& rewriter, Location loc) { SmallVector offsets {rewriter.getIndexAttr(0), column}; SmallVector sizes {rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)}; SmallVector strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; @@ -280,7 +297,7 @@ static Value extractDynamicGemmBColumn( } static Value extractDynamicGemmRowVector( - Value matrix, Value row, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) { + Value matrix, Value row, RankedTensorType vectorType, PatternRewriter& rewriter, Location loc) { SmallVector offsets {row, rewriter.getIndexAttr(0)}; SmallVector sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))}; SmallVector strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; @@ -311,13 +328,15 @@ static FailureOr verifyDynamicGemmBiasType(RankedTensorType cT } static bool hasGemmBias(Value c) { + if (!c) + return false; Operation* definingOp = c.getDefiningOp(); return (!definingOp || !isa(definingOp)) && !isZeroSplatHostConstant(c); } static Value createScalarTensorConstant(RankedTensorType scalarType, float value, - ConversionPatternRewriter& rewriter, + PatternRewriter& rewriter, Location loc) { auto elementType = scalarType.getElementType(); auto scalarAttr = rewriter.getFloatAttr(elementType, value); @@ -330,7 +349,7 @@ static Value createBroadcastedBiasScalar(Value bias, Value row, Value column, RankedTensorType scalarType, - ConversionPatternRewriter& rewriter, + PatternRewriter& rewriter, Location loc) { SmallVector unitStrides(biasType.getRank(), rewriter.getIndexAttr(1)); if (biasType.getRank() == 1) { @@ -365,7 +384,7 @@ static FailureOr createVvdmulBatch(Value a, RankedTensorType columnPiecesType, RankedTensorType outType, bool transposeB, - ConversionPatternRewriter& rewriter, + PatternRewriter& rewriter, Location loc) { const int64_t numOutRows = outType.getDimSize(0); const int64_t numOutCols = outType.getDimSize(1); @@ -425,7 +444,7 @@ static FailureOr createDynamicGemmOutputCompute(Value scal RankedTensorType outType, float alpha, float beta, - ConversionPatternRewriter& rewriter, + PatternRewriter& rewriter, Location loc) { const int64_t numOutRows = outType.getDimSize(0); const int64_t numOutCols = outType.getDimSize(1); @@ -510,7 +529,7 @@ static Value createPartialGroupOffset(Value hSlice, int64_t kSlice, int64_t numKSlices, int64_t numOutRows, - ConversionPatternRewriter& rewriter, + PatternRewriter& rewriter, Location loc) { MLIRContext* context = rewriter.getContext(); AffineExpr d0 = getAffineDimExpr(0, context); @@ -527,10 +546,12 @@ static Value extractReductionPiece(Value partialPiecesArg, RankedTensorType pieceType, int64_t numKSlices, int64_t numOutRows, - ConversionPatternRewriter& rewriter, + int64_t xbarSize, + PatternRewriter& rewriter, Location loc) { SmallVector unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; - SmallVector pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())}; + SmallVector pieceSizes { + rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)}; SmallVector pieceOffsets { createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0), @@ -545,13 +566,15 @@ static Value reducePartialPiecesForHSlice(Value partialPiecesArg, RankedTensorType pieceType, int64_t numKSlices, int64_t numOutRows, - ConversionPatternRewriter& rewriter, + int64_t xbarSize, + PatternRewriter& rewriter, Location loc) { SmallVector activePieces; activePieces.reserve(numKSlices); for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) activePieces.push_back( - extractReductionPiece(partialPiecesArg, hSlice, kSlice, pieceType, numKSlices, numOutRows, rewriter, loc)); + extractReductionPiece( + partialPiecesArg, hSlice, kSlice, pieceType, numKSlices, numOutRows, xbarSize, rewriter, loc)); while (activePieces.size() > 1) { SmallVector nextPieces; @@ -574,11 +597,12 @@ static FailureOr createReductionOutput(Value partialPieces, RankedTensorType outType, RankedTensorType paddedOutType, int64_t numKSlices, - ConversionPatternRewriter& rewriter, + int64_t xbarSize, + PatternRewriter& rewriter, Location loc) { const int64_t numOutRows = outType.getDimSize(0); - const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue()); - auto pieceType = RankedTensorType::get({numOutRows, static_cast(crossbarSize.getValue())}, + const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), xbarSize); + auto pieceType = RankedTensorType::get({numOutRows, xbarSize}, partialPiecesType.getElementType()); if (bias && cast(bias.getType()) != paddedOutType) @@ -590,20 +614,20 @@ static FailureOr createReductionOutput(Value partialPieces, SmallVector outputSlices; outputSlices.reserve(numOutHSlices); for (int64_t hSlice = 0; hSlice < numOutHSlices; ++hSlice) { - const int64_t columnOffset = hSlice * crossbarSize.getValue(); + const int64_t columnOffset = hSlice * xbarSize; const int64_t columns = - std::min(static_cast(crossbarSize.getValue()), outType.getDimSize(1) - columnOffset); + std::min(xbarSize, outType.getDimSize(1) - columnOffset); auto outputSliceType = RankedTensorType::get({numOutRows, columns}, outType.getElementType()); auto computeOp = createSpatCompute( rewriter, loc, TypeRange {outputSliceType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult { Value hSliceValue = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), hSlice); Value reduced = reducePartialPiecesForHSlice( - blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, rewriter, loc); + blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, xbarSize, rewriter, loc); if (bias) { SmallVector biasOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(columnOffset)}; SmallVector pieceSizes {rewriter.getIndexAttr(numOutRows), - rewriter.getIndexAttr(crossbarSize.getValue())}; + rewriter.getIndexAttr(xbarSize)}; SmallVector unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}; Value biasSlice = tensor::ExtractSliceOp::create(rewriter, loc, pieceType, blockArgs[1], biasOffsets, pieceSizes, unitStrides) @@ -637,79 +661,96 @@ static FailureOr createReductionOutput(Value partialPieces, } struct GemmToSpatialComputes : OpConversionPattern { - using OpConversionPattern::OpConversionPattern; + explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetResources& target) + : OpConversionPattern(ctx), target(target) {} LogicalResult matchAndRewrite(ONNXGemmOp gemmOp, ONNXGemmOpAdaptor gemmOpAdaptor, ConversionPatternRewriter& rewriter) const override; + + const spatial::SpatialTargetResources& target; }; } // namespace -LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp, - ONNXGemmOpAdaptor gemmOpAdaptor, - ConversionPatternRewriter& rewriter) const { - Location loc = gemmOp.getLoc(); - Value a = gemmOpAdaptor.getA(); - Value b = gemmOpAdaptor.getB(); - Value c = gemmOpAdaptor.getC(); - +FailureOr lowerGemmToSpatial( + Operation* diagnosticAnchor, + Value a, + Value b, + Value c, + RankedTensorType outType, + bool transA, + bool transB, + float alpha, + float beta, + const spatial::SpatialTargetResources& target, + PatternRewriter& rewriter, + Location loc) { auto aType = dyn_cast(a.getType()); auto bType = dyn_cast(b.getType()); - auto outType = dyn_cast(gemmOp.getY().getType()); - if (!aType || !bType || !outType) + if (!diagnosticAnchor || !aType || !bType || !outType) return failure(); if (!aType.hasStaticShape()) { - pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input A"); + pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm input A"); return failure(); } if (!bType.hasStaticShape()) { - pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input B"); + pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm input B"); return failure(); } if (!outType.hasStaticShape()) { - pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm result"); + pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm result"); return failure(); } if (aType.getRank() != 2) { - pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm input A", aType.getRank(), {2}); + pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm input A", aType.getRank(), {2}); return failure(); } if (bType.getRank() != 2) { - pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm input B", bType.getRank(), {2}); + pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm input B", bType.getRank(), {2}); return failure(); } if (outType.getRank() != 2) { - pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm result", outType.getRank(), {2}); + pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm result", outType.getRank(), {2}); return failure(); } - if (gemmOpAdaptor.getTransA()) { + if (transA) { auto aShape = aType.getShape(); auto transposedType = RankedTensorType::get({aShape[1], aShape[0]}, aType.getElementType(), aType.getEncoding()); - a = ONNXTransposeOp::create(rewriter, loc, transposedType, a, rewriter.getI64ArrayAttr({1, 0})).getResult(); + a = createLinalgTranspose(a, transposedType, {1, 0}, rewriter, loc); aType = transposedType; } - const int64_t numOutRows = outType.getDimSize(0); - const int64_t numOutCols = outType.getDimSize(1); - const int64_t reductionSize = aType.getDimSize(1); - const bool transposeB = gemmOpAdaptor.getTransB(); + ContractionProblem problem; + problem.lhsBatchShape = {}; + problem.rhsBatchShape = {}; + problem.outputBatchShape = {}; + problem.lhsBatch = 1; + problem.rhsBatch = 1; + problem.batch = 1; + problem.m = outType.getDimSize(0); + problem.k = aType.getDimSize(1); + problem.n = outType.getDimSize(1); + problem.lhsElementType = aType.getElementType(); + problem.rhsElementType = bType.getElementType(); + problem.resultElementType = outType.getElementType(); + const bool transposeB = transB; if (!isCompileTimeComputable(b)) { + ContractionPlan plan = makeContractionPlan( + problem, target, ContractionPlanKind::BatchedDynamicVVD); bool hasC = hasGemmBias(c); - float alpha = gemmOpAdaptor.getAlpha().convertToFloat(); - float beta = gemmOpAdaptor.getBeta().convertToFloat(); RankedTensorType biasType; if (hasC) { auto cType = dyn_cast(c.getType()); if (!cType || !cType.hasStaticShape()) { - pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias"); + pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm bias"); return failure(); } auto verifiedBiasType = verifyDynamicGemmBiasType(cType, outType); if (failed(verifiedBiasType)) { - gemmOp.emitOpError("requires Gemm bias C to be broadcastable to the output shape"); + diagnosticAnchor->emitOpError("requires Gemm bias C to be broadcastable to the output shape"); return failure(); } biasType = *verifiedBiasType; @@ -717,19 +758,19 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp, const int64_t bReductionSize = bType.getDimSize(transposeB ? 1 : 0); const int64_t bOutputColumns = bType.getDimSize(transposeB ? 0 : 1); - if (aType.getDimSize(0) != numOutRows || bReductionSize != reductionSize || bOutputColumns != numOutCols) { - gemmOp.emitOpError("has inconsistent A, B, and output shapes"); + if (aType.getDimSize(0) != problem.m || bReductionSize != problem.k || bOutputColumns != problem.n) { + diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes"); return failure(); } - const int64_t laneCount64 = numOutRows * numOutCols; + const int64_t laneCount64 = plan.laneCount; if (laneCount64 > std::numeric_limits::max()) { - gemmOp.emitOpError("requires Gemm dynamic batch lane count to fit in i32"); + diagnosticAnchor->emitOpError("requires Gemm dynamic batch lane count to fit in i32"); return failure(); } - auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType()); - auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(numOutCols, columnType); + auto columnType = RankedTensorType::get({problem.m, 1}, outType.getElementType()); + auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(problem.n, columnType); auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, transposeB, rewriter, loc); if (failed(batchOp)) return failure(); @@ -737,94 +778,128 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp, batchOp->getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc); if (failed(outputCompute)) return failure(); - rewriter.replaceOp(gemmOp, outputCompute->getResults()); - return success(); + return outputCompute->getResult(0); } if (transposeB) { auto bShape = bType.getShape(); auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding()); - b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult(); + if (isCompileTimeComputable(b)) { + auto denseAttr = getHostConstDenseElementsAttr(b); + auto inputType = denseAttr ? dyn_cast(denseAttr.getType()) : nullptr; + auto transposedAttr = inputType && inputType.hasStaticShape() && transposedType.hasStaticShape() + ? transposeDenseElementsAttr(denseAttr, {1, 0}) + : FailureOr(failure()); + if (failed(transposedAttr) || transposedAttr->getType() != transposedType) { + diagnosticAnchor->emitOpError("requires Gemm input B transpose to remain statically materializable"); + return failure(); + } + b = getOrCreateConstant(rewriter, + rewriter.getInsertionBlock()->getParentOp(), + *transposedAttr, + transposedType); + } else { + b = createLinalgTranspose(b, transposedType, {1, 0}, rewriter, loc); + } bType = transposedType; } - auto scaledB = materializeScaledConstantTensor(b, gemmOpAdaptor.getAlpha().convertToFloat(), rewriter, loc); + auto scaledB = materializeScaledConstantTensor(b, alpha, rewriter, loc); if (failed(scaledB)) { - gemmOp.emitOpError("requires constant Gemm input B when alpha is not 1.0"); + diagnosticAnchor->emitOpError("requires constant Gemm input B when alpha is not 1.0"); return failure(); } b = *scaledB; bType = cast(b.getType()); - if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize || bType.getDimSize(1) != numOutCols) { - gemmOp.emitOpError("has inconsistent A, B, and output shapes after transpose handling"); + if (aType.getDimSize(0) != problem.m || bType.getDimSize(0) != problem.k || bType.getDimSize(1) != problem.n) { + diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes after transpose handling"); return failure(); } - const int64_t numKSlices = ceilIntegerDivide(reductionSize, crossbarSize.getValue()); - const int64_t numOutHSlices = ceilIntegerDivide(numOutCols, crossbarSize.getValue()); - const int64_t paddedReductionSize = numKSlices * static_cast(crossbarSize.getValue()); - const int64_t paddedOutCols = numOutHSlices * static_cast(crossbarSize.getValue()); + ContractionPlan plan = makeContractionPlan( + problem, target, ContractionPlanKind::StaticTiled); + const int64_t xbarSize = plan.tileK; + const int64_t numKSlices = plan.reductionSlices; + const int64_t numOutHSlices = plan.outputTiles; + const int64_t paddedReductionSize = numKSlices * plan.tileK; + const int64_t paddedOutCols = numOutHSlices * plan.tileN; auto paddedBType = RankedTensorType::get({paddedReductionSize, paddedOutCols}, bType.getElementType()); auto paddedB = materializePaddedConstantMatrix(b, paddedBType, rewriter, loc); if (failed(paddedB)) { - gemmOp.emitOpError("requires constant Gemm input B so tiled weights can be padded statically"); + diagnosticAnchor->emitOpError("requires constant Gemm input B so tiled weights can be padded statically"); return failure(); } b = *paddedB; - auto paddedAType = RankedTensorType::get({numOutRows, paddedReductionSize}, aType.getElementType()); + auto paddedAType = RankedTensorType::get({problem.m, paddedReductionSize}, aType.getElementType()); a = createPaddedInputCompute(a, paddedAType, rewriter, loc); aType = paddedAType; Value bias; bool hasC = hasGemmBias(c); - auto paddedOutType = RankedTensorType::get({numOutRows, paddedOutCols}, outType.getElementType()); + auto paddedOutType = RankedTensorType::get({problem.m, paddedOutCols}, outType.getElementType()); if (hasC) { auto cType = dyn_cast(c.getType()); if (!cType || !cType.hasStaticShape()) { - pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias"); + pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm bias"); return failure(); } - auto scaledC = materializeScaledConstantTensor(c, gemmOpAdaptor.getBeta().convertToFloat(), rewriter, loc); + auto scaledC = materializeScaledConstantTensor(c, beta, rewriter, loc); if (failed(scaledC)) { - gemmOp.emitOpError("requires constant Gemm bias C when beta is not 1.0"); + diagnosticAnchor->emitOpError("requires constant Gemm bias C when beta is not 1.0"); return failure(); } c = *scaledC; auto preparedBias = prepareBias(c, outType, paddedOutType, rewriter, loc); if (failed(preparedBias)) { - gemmOp.emitOpError("requires Gemm bias C to be broadcastable to the output shape"); + diagnosticAnchor->emitOpError("requires Gemm bias C to be broadcastable to the output shape"); return failure(); } bias = *preparedBias; } - const int64_t laneCount64 = numOutHSlices * numKSlices * numOutRows; + const int64_t laneCount64 = plan.laneCount; if (laneCount64 > std::numeric_limits::max()) { - gemmOp.emitOpError("requires Gemm tiled batch lane count to fit in i32"); + diagnosticAnchor->emitOpError("requires Gemm tiled batch lane count to fit in i32"); return failure(); } auto partialPiecesType = spatial::getGraphBatchPhysicalResultType( - laneCount64, RankedTensorType::get({1, static_cast(crossbarSize.getValue())}, outType.getElementType())); + laneCount64, RankedTensorType::get({1, xbarSize}, outType.getElementType())); auto batchOp = - createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc); + createVmmBatch( + a, b, aType, paddedBType, partialPiecesType, problem.m, numKSlices, numOutHSlices, xbarSize, rewriter, loc); if (failed(batchOp)) return failure(); auto reductionOutput = createReductionOutput( - batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc); + batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, xbarSize, rewriter, loc); if (failed(reductionOutput)) return failure(); - rewriter.replaceOp(gemmOp, *reductionOutput); + return *reductionOutput; +} + +LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp, + ONNXGemmOpAdaptor gemmOpAdaptor, + ConversionPatternRewriter& rewriter) const { + FailureOr result = lowerGemmToSpatial( + gemmOp.getOperation(), gemmOpAdaptor.getA(), gemmOpAdaptor.getB(), gemmOpAdaptor.getC(), + cast(gemmOp.getY().getType()), gemmOpAdaptor.getTransA(), + gemmOpAdaptor.getTransB(), gemmOpAdaptor.getAlpha().convertToFloat(), + gemmOpAdaptor.getBeta().convertToFloat(), target, rewriter, gemmOp.getLoc()); + if (failed(result)) + return failure(); + rewriter.replaceOp(gemmOp, *result); return success(); } -void populateGemmPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { - patterns.insert(ctx); +void populateGemmPatterns(RewritePatternSet& patterns, + MLIRContext* ctx, + const spatial::SpatialTargetResources& target) { + patterns.insert(ctx, target); } } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp new file mode 100644 index 0000000..b4104a0 --- /dev/null +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp @@ -0,0 +1,27 @@ +#pragma once + +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/PatternMatch.h" + +namespace onnx_mlir { +namespace spatial { +struct SpatialTargetResources; +} + +mlir::FailureOr lowerGemmToSpatial( + mlir::Operation* diagnosticAnchor, + mlir::Value a, + mlir::Value b, + mlir::Value c, + mlir::RankedTensorType outputType, + bool transA, + bool transB, + float alpha, + float beta, + const spatial::SpatialTargetResources& target, + mlir::PatternRewriter& rewriter, + mlir::Location loc); + +} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp index fd26d78..d2df42e 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp @@ -11,7 +11,11 @@ #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -106,6 +110,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() + || !spatial::isFragmentAssembly(blueprint.getMode()) || !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(), + spatial::getFragmentedLayout(rewriter.getContext()), + 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 +203,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,19 +333,46 @@ 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) { - return ONNXTransposeOp::create(rewriter, loc, resultType, value, rewriter.getI64ArrayAttr(permutation)).getResult(); + auto createTranspose = [&](RankedTensorType resultType, ArrayRef permutation) { + return createLinalgTranspose(value, resultType, permutation, rewriter, loc); }; if (type.getRank() == 2) { auto resultType = RankedTensorType::get({shape[1], shape[0]}, type.getElementType(), type.getEncoding()); - return createONNXTranspose(resultType, {1, 0}); + return createTranspose(resultType, {1, 0}); } auto resultType = RankedTensorType::get({shape[0], shape[2], shape[1]}, type.getElementType(), type.getEncoding()); - return createONNXTranspose(resultType, {0, 2, 1}); + return createTranspose(resultType, {0, 2, 1}); } static FailureOr materializePaddedBatchedWeight(Value value, @@ -347,6 +466,7 @@ static FailureOr createBatchedVmmBatch(Value a, int64_t numOutRows, int64_t numKSlices, int64_t numOutHSlices, + int64_t xbarSize, PatternRewriter& rewriter, Location loc) { const int64_t laneCount = partialPiecesType.getDimSize(0); @@ -365,16 +485,16 @@ static FailureOr createBatchedVmmBatch(Value a, Value sliceLane = affineModConst(rewriter, loc, outerLane, numKSlices * numOutHSlices, anchorOp); Value kSlice = affineModConst(rewriter, loc, sliceLane, numKSlices, anchorOp); Value hSlice = affineFloorDivConst(rewriter, loc, sliceLane, numKSlices, anchorOp); - Value kOffset = affineMulConst(rewriter, loc, kSlice, crossbarSize.getValue(), anchorOp); - Value hOffset = affineMulConst(rewriter, loc, hSlice, crossbarSize.getValue(), anchorOp); + Value kOffset = affineMulConst(rewriter, loc, kSlice, xbarSize, anchorOp); + Value hOffset = affineMulConst(rewriter, loc, hSlice, xbarSize, anchorOp); auto aTileType = - RankedTensorType::get({1, static_cast(crossbarSize.getValue())}, aType.getElementType()); + RankedTensorType::get({1, xbarSize}, aType.getElementType()); auto bTileType = RankedTensorType::get( - {static_cast(crossbarSize.getValue()), static_cast(crossbarSize.getValue())}, + {xbarSize, xbarSize}, bType.getElementType()); auto pieceType = - RankedTensorType::get({1, static_cast(crossbarSize.getValue())}, partialPiecesType.getElementType()); + RankedTensorType::get({1, xbarSize}, partialPiecesType.getElementType()); Value aTile = extractBatchedATile( args.inputs.front(), aBatchShape, outputBatchShape, batch, row, kOffset, aTileType, rewriter, loc); @@ -407,123 +527,155 @@ static Value extractDynamicBatchedRowVector(Value matrix, {offsets, sizes, getUnitStrides(rewriter, 3)}); } +static int64_t chooseDynamicMatMulRowsPerLane(int64_t rows, + int64_t reductionSize, + int64_t columns, + int64_t xbarSize) { + const int64_t crossbarElements = xbarSize * xbarSize; + 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, @@ -534,6 +686,7 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg, int64_t numKSlices, int64_t numOutHSlices, int64_t numOutRows, + int64_t xbarSize, PatternRewriter& rewriter, Location loc) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); @@ -543,7 +696,8 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg, Value batchAndHSlice = arith::AddIOp::create(rewriter, loc, batchOffset, hOffset); Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset); SmallVector offsets {pieceOffset, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}; - SmallVector sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())}; + SmallVector sizes { + rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)}; return extractMixedSliceOrIdentity( rewriter, loc, partialPiecesArg, pieceType, {offsets, sizes, getUnitStrides(rewriter, 3)}); @@ -556,13 +710,24 @@ static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg, int64_t numKSlices, int64_t numOutHSlices, int64_t numOutRows, + int64_t xbarSize, PatternRewriter& rewriter, Location loc) { SmallVector activePieces; activePieces.reserve(numKSlices); for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice) activePieces.push_back(extractBatchedReductionPiece( - partialPiecesArg, batch, hSlice, kSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, loc)); + partialPiecesArg, + batch, + hSlice, + kSlice, + pieceType, + numKSlices, + numOutHSlices, + numOutRows, + xbarSize, + rewriter, + loc)); while (activePieces.size() > 1) { SmallVector nextPieces; @@ -585,13 +750,14 @@ static FailureOr createBatchedReductionCompute(Value partialPieces, RankedTensorType paddedOutType, int64_t numBatches, int64_t numKSlices, + int64_t xbarSize, PatternRewriter& rewriter, Location loc) { auto computeOp = createSpatCompute<1>( rewriter, loc, TypeRange {outType}, {}, ValueRange {partialPieces}, [&](Value partialPiecesArg) -> LogicalResult { const int64_t numOutRows = outType.getDimSize(1); - const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), crossbarSize.getValue()); - auto pieceType = RankedTensorType::get({numOutRows, static_cast(crossbarSize.getValue())}, + const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), xbarSize); + auto pieceType = RankedTensorType::get({numOutRows, xbarSize}, partialPiecesType.getElementType()); Value outputInit = @@ -621,13 +787,22 @@ static FailureOr createBatchedReductionCompute(Value partialPieces, [&](OpBuilder&, Location hLoc, Value hSlice, ValueRange hIterArgs, SmallVectorImpl& hYielded) { Value outputAcc = hIterArgs.front(); Value reduced = reduceBatchedPartialPiecesForHSlice( - partialPiecesArg, batch, hSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, hLoc); + partialPiecesArg, + batch, + hSlice, + pieceType, + numKSlices, + numOutHSlices, + numOutRows, + xbarSize, + rewriter, + hLoc); Value hOffset = affineMulConst( - rewriter, hLoc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp()); + rewriter, hLoc, hSlice, xbarSize, rewriter.getInsertionBlock()->getParentOp()); SmallVector outputOffsets {batch, rewriter.getIndexAttr(0), hOffset}; SmallVector outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(numOutRows), - rewriter.getIndexAttr(crossbarSize.getValue())}; + rewriter.getIndexAttr(xbarSize)}; Value next = tensor::InsertSliceOp::create( rewriter, hLoc, reduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3)) @@ -662,38 +837,46 @@ static FailureOr createBatchedReductionCompute(Value partialPieces, } struct NormalizedMatMulInfo { + NormalizedMatMulInfo(RankedTensorType lhsType, + RankedTensorType rhsType, + RankedTensorType outType, + RankedTensorType normalizedLhsType, + RankedTensorType normalizedRhsType, + ContractionProblem problem, + bool lhsWasVector, + bool rhsWasVector) + : problem(std::move(problem)), + lhsType(lhsType), + rhsType(rhsType), + outType(outType), + normalizedLhsType(normalizedLhsType), + normalizedRhsType(normalizedRhsType), + lhsWasVector(lhsWasVector), + rhsWasVector(rhsWasVector) {} + + ContractionProblem problem; RankedTensorType lhsType; RankedTensorType rhsType; RankedTensorType outType; RankedTensorType normalizedLhsType; RankedTensorType normalizedRhsType; - SmallVector lhsBatchShape; - SmallVector rhsBatchShape; - SmallVector outputBatchShape; bool lhsWasVector; bool rhsWasVector; - int64_t lhsBatch; - int64_t rhsBatch; - int64_t batch; - int64_t m; - int64_t k; - int64_t n; }; struct MatMulLoweringPlan { + MatMulLoweringPlan(Value lhs, Value rhs, const NormalizedMatMulInfo& info) + : problem(info.problem), + lhs(lhs), + rhs(rhs), + lhsType(cast(lhs.getType())), + rhsType(cast(rhs.getType())) {} + + ContractionProblem problem; Value lhs; Value rhs; RankedTensorType lhsType; RankedTensorType rhsType; - SmallVector lhsBatchShape; - SmallVector rhsBatchShape; - SmallVector outputBatchShape; - int64_t lhsBatch; - int64_t rhsBatch; - int64_t batch; - int64_t m; - int64_t k; - int64_t n; bool transposedResult; }; @@ -757,22 +940,26 @@ static FailureOr analyzeMatMulShape(ONNXMatMulOp matmulOp) return failure(); } - return NormalizedMatMulInfo {lhsType, - rhsType, - outType, - normalizedLhsType, - normalizedRhsType, - lhsBatchShape, - rhsBatchShape, - *outputBatchShape, - lhsWasVector, - rhsWasVector, - lhsBatch, - rhsBatch, - batch, - m, - k, - n}; + return NormalizedMatMulInfo( + lhsType, + rhsType, + outType, + normalizedLhsType, + normalizedRhsType, + ContractionProblem {lhsBatchShape, + rhsBatchShape, + *outputBatchShape, + lhsBatch, + rhsBatch, + batch, + m, + k, + n, + lhsType.getElementType(), + rhsType.getElementType(), + outType.getElementType()}, + lhsWasVector, + rhsWasVector); } static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs, @@ -781,20 +968,8 @@ static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs, bool useTransposedForm, PatternRewriter& rewriter, Location loc) { - MatMulLoweringPlan plan {normalizedLhs, - normalizedRhs, - cast(normalizedLhs.getType()), - cast(normalizedRhs.getType()), - info.lhsBatchShape, - info.rhsBatchShape, - info.outputBatchShape, - info.lhsBatch, - info.rhsBatch, - info.batch, - info.m, - info.k, - info.n, - false}; + MatMulLoweringPlan plan(normalizedLhs, normalizedRhs, info); + plan.transposedResult = false; if (!useTransposedForm) return plan; @@ -802,10 +977,10 @@ static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs, plan.rhs = transposeLastTwoDims(normalizedLhs, rewriter, loc); plan.lhsType = cast(plan.lhs.getType()); plan.rhsType = cast(plan.rhs.getType()); - std::swap(plan.lhsBatchShape, plan.rhsBatchShape); - std::swap(plan.lhsBatch, plan.rhsBatch); - plan.m = info.n; - plan.n = info.m; + std::swap(plan.problem.lhsBatchShape, plan.problem.rhsBatchShape); + std::swap(plan.problem.lhsBatch, plan.problem.rhsBatch); + plan.problem.m = info.problem.n; + plan.problem.n = info.problem.m; plan.transposedResult = true; return plan; } @@ -822,22 +997,24 @@ 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. Value result = value; RankedTensorType currentType = directOutType; - if (info.outputBatchShape.size() > 1) { - SmallVector expandedShape(info.outputBatchShape.begin(), info.outputBatchShape.end()); - expandedShape.push_back(info.m); - expandedShape.push_back(info.n); + if (info.problem.outputBatchShape.size() > 1) { + SmallVector expandedShape(info.problem.outputBatchShape.begin(), info.problem.outputBatchShape.end()); + expandedShape.push_back(info.problem.m); + expandedShape.push_back(info.problem.n); auto expandedType = RankedTensorType::get(expandedShape, info.outType.getElementType(), info.outType.getEncoding()); - result = expandBatchDims(result, expandedType, info.outputBatchShape.size(), rewriter, loc); + result = expandBatchDims(result, expandedType, info.problem.outputBatchShape.size(), rewriter, loc); currentType = expandedType; } SmallVector removedAxes(currentType.getRank(), false); - if (info.outputBatchShape.empty()) + if (info.problem.outputBatchShape.empty()) removedAxes[0] = true; if (info.lhsWasVector) removedAxes[currentType.getRank() - 2] = true; @@ -847,128 +1024,152 @@ static Value finalizeNormalizedMatMulResult(Value value, } struct MatMulToGemm : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; + explicit MatMulToGemm(MLIRContext* ctx, + const spatial::SpatialTargetResources& target) + : OpRewritePattern(ctx), target(target) {} LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override { auto shapeInfo = analyzeMatMulShape(matmulOp); if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector) return failure(); - if (!shapeInfo->outputBatchShape.empty()) + if (!shapeInfo->problem.outputBatchShape.empty()) return failure(); Location loc = matmulOp.getLoc(); bool useTransposedForm = isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB()); - Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc); - Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc); - int64_t lhsBatchForGemm = shapeInfo->lhsBatch; - int64_t rhsBatchForGemm = shapeInfo->rhsBatch; - int64_t gemmM = shapeInfo->m; - int64_t gemmK = shapeInfo->k; - int64_t gemmN = shapeInfo->n; + Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->problem.lhsBatch, shapeInfo->problem.m, shapeInfo->problem.k, rewriter, loc); + Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->problem.rhsBatch, shapeInfo->problem.k, shapeInfo->problem.n, rewriter, loc); + int64_t lhsBatchForGemm = shapeInfo->problem.lhsBatch; + int64_t rhsBatchForGemm = shapeInfo->problem.rhsBatch; + int64_t gemmM = shapeInfo->problem.m; + int64_t gemmK = shapeInfo->problem.k; + int64_t gemmN = shapeInfo->problem.n; if (useTransposedForm) { lhs = transposeLastTwoDims(matmulOp.getB(), rewriter, loc); - lhsBatchForGemm = shapeInfo->rhsBatch; + lhsBatchForGemm = shapeInfo->problem.rhsBatch; rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc); - rhsBatchForGemm = shapeInfo->lhsBatch; - gemmM = shapeInfo->n; - gemmN = shapeInfo->m; + rhsBatchForGemm = shapeInfo->problem.lhsBatch; + gemmM = shapeInfo->problem.n; + gemmN = shapeInfo->problem.m; } - auto gemmType = RankedTensorType::get({gemmM, gemmN}, shapeInfo->outType.getElementType()); - Value none = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType()); + auto gemmType = RankedTensorType::get( + {gemmM, gemmN}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding()); Value lhsMatrix = extractBatchMatrix(lhs, /*batchIndex=*/0, lhsBatchForGemm, gemmM, gemmK, rewriter, loc); Value rhsMatrix = extractBatchMatrix(rhs, /*batchIndex=*/0, rhsBatchForGemm, gemmK, gemmN, rewriter, loc); - Value gemmResult = ONNXGemmOp::create(rewriter, - loc, - gemmType, - lhsMatrix, - rhsMatrix, - none, - rewriter.getF32FloatAttr(1.0f), - rewriter.getF32FloatAttr(1.0f), - rewriter.getBoolAttr(false), - rewriter.getBoolAttr(false)) - .getY(); + FailureOr gemmResult = lowerGemmToSpatial( + matmulOp.getOperation(), lhsMatrix, rhsMatrix, Value(), gemmType, + /*transA=*/false, /*transB=*/false, /*alpha=*/1.0f, /*beta=*/1.0f, + target, rewriter, loc); + if (failed(gemmResult)) + return failure(); + Value result = *gemmResult; if (useTransposedForm) - gemmResult = - ONNXTransposeOp::create(rewriter, loc, shapeInfo->outType, gemmResult, rewriter.getI64ArrayAttr({1, 0})) - .getResult(); + result = transposeLastTwoDims(result, rewriter, loc); - if (shapeInfo->outputBatchShape.empty()) { - rewriter.replaceOp(matmulOp, gemmResult); + if (shapeInfo->problem.outputBatchShape.empty()) { + rewriter.replaceOp(matmulOp, result); return success(); } auto directOutType = - RankedTensorType::get({1, shapeInfo->m, shapeInfo->n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding()); - Value batchedResult = ensureBatchedTensor(gemmResult, /*batchSize=*/1, shapeInfo->m, shapeInfo->n, rewriter, loc); + RankedTensorType::get({1, shapeInfo->problem.m, shapeInfo->problem.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding()); + Value batchedResult = ensureBatchedTensor(result, /*batchSize=*/1, shapeInfo->problem.m, shapeInfo->problem.n, rewriter, loc); Value finalResult = finalizeNormalizedMatMulResult(batchedResult, directOutType, *shapeInfo, rewriter, loc); rewriter.replaceOp(matmulOp, finalResult); return success(); } + + const spatial::SpatialTargetResources& target; }; struct MatMulBatchedToSpatialComputes : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; + explicit MatMulBatchedToSpatialComputes(MLIRContext* ctx, + const spatial::SpatialTargetResources& target) + : OpRewritePattern(ctx), target(target) {} LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override { auto shapeInfo = analyzeMatMulShape(matmulOp); if (failed(shapeInfo)) return failure(); - if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->outputBatchShape.empty()) + if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->problem.outputBatchShape.empty()) return failure(); Location loc = matmulOp.getLoc(); + const int64_t xbarSize = static_cast(target.matrixShape.rows); 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); - lhs = collapseBatchDims(lhs, shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc); - rhs = collapseBatchDims(rhs, shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc); - MatMulLoweringPlan plan = buildLoweringPlan(lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc); + normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType, + shapeInfo->lhsWasVector, rewriter, loc); + Value rhs = normalizeMatMulOperand( + rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType, + shapeInfo->rhsWasVector, rewriter, loc); + lhs = collapseBatchDims(lhs, shapeInfo->problem.lhsBatch, shapeInfo->problem.m, shapeInfo->problem.k, rewriter, loc); + rhs = collapseBatchDims(rhs, + shapeInfo->problem.rhsBatch, + rhsStoredAsRows ? shapeInfo->problem.n : shapeInfo->problem.k, + rhsStoredAsRows ? shapeInfo->problem.k : shapeInfo->problem.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.lhs = ensureBatchedTensor(plan.lhs, plan.problem.lhsBatch, plan.problem.m, plan.problem.k, rewriter, loc); + plan.rhs = ensureBatchedTensor(plan.rhs, + plan.problem.rhsBatch, + rhsStoredAsRows ? plan.problem.n : plan.problem.k, + rhsStoredAsRows ? plan.problem.k : plan.problem.n, + rewriter, + loc); plan.lhsType = cast(plan.lhs.getType()); plan.rhsType = cast(plan.rhs.getType()); auto directOutType = RankedTensorType::get( - {plan.batch, plan.m, plan.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding()); + {plan.problem.batch, plan.problem.m, plan.problem.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding()); if (isCompileTimeComputable(plan.rhs)) { - const int64_t numKSlices = ceilIntegerDivide(plan.k, crossbarSize.getValue()); - const int64_t numOutHSlices = ceilIntegerDivide(plan.n, crossbarSize.getValue()); - const int64_t paddedReductionSize = numKSlices * static_cast(crossbarSize.getValue()); - const int64_t paddedOutCols = numOutHSlices * static_cast(crossbarSize.getValue()); + ContractionPlan contractionPlan = makeContractionPlan( + plan.problem, target, ContractionPlanKind::StaticTiled); + const int64_t numKSlices = contractionPlan.reductionSlices; + const int64_t numOutHSlices = contractionPlan.outputTiles; + const int64_t paddedReductionSize = numKSlices * xbarSize; + const int64_t paddedOutCols = numOutHSlices * xbarSize; auto paddedLhsType = RankedTensorType::get( - {plan.lhsBatch, plan.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding()); + {plan.problem.lhsBatch, plan.problem.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding()); auto paddedRhsType = RankedTensorType::get( - {plan.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding()); + {plan.problem.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding()); auto paddedOutType = - RankedTensorType::get({plan.batch, plan.m, paddedOutCols}, shapeInfo->outType.getElementType()); + RankedTensorType::get({plan.problem.batch, plan.problem.m, paddedOutCols}, shapeInfo->outType.getElementType()); auto paddedRhs = - materializePaddedBatchedWeight(plan.rhs, plan.rhsBatchShape, plan.outputBatchShape, paddedRhsType, rewriter); + materializePaddedBatchedWeight(plan.rhs, plan.problem.rhsBatchShape, plan.problem.outputBatchShape, paddedRhsType, rewriter); if (succeeded(paddedRhs)) { Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc); - const int64_t laneCount = plan.batch * plan.m * numKSlices * numOutHSlices; + const int64_t laneCount = contractionPlan.laneCount; auto partialPiecesType = spatial::getGraphBatchPhysicalResultType( - laneCount, RankedTensorType::get({1, static_cast(crossbarSize.getValue())}, shapeInfo->outType.getElementType())); + laneCount, RankedTensorType::get({1, xbarSize}, shapeInfo->outType.getElementType())); auto batchOp = createBatchedVmmBatch(paddedLhs, *paddedRhs, paddedLhsType, - plan.lhsBatchShape, + plan.problem.lhsBatchShape, paddedRhsType, - plan.rhsBatchShape, - plan.outputBatchShape, + plan.problem.rhsBatchShape, + plan.problem.outputBatchShape, partialPiecesType, - plan.m, + plan.problem.m, numKSlices, numOutHSlices, + xbarSize, rewriter, loc); if (failed(batchOp)) @@ -977,67 +1178,112 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern { partialPiecesType, directOutType, paddedOutType, - plan.batch, + plan.problem.batch, numKSlices, + xbarSize, rewriter, loc); if (failed(result)) return failure(); Value finalResult = *result; if (plan.transposedResult) { - auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n}, - shapeInfo->outType.getElementType(), - shapeInfo->outType.getEncoding()); - finalResult = - ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1})) - .getResult(); + finalResult = transposeLastTwoDims(finalResult, rewriter, loc); } finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc); rewriter.replaceOp(matmulOp, finalResult); 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->problem.outputBatchShape : SmallVector {plan.problem.batch}; + const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.problem.m, plan.problem.k, plan.problem.n, xbarSize); + ContractionPlan contractionPlan = makeContractionPlan( + plan.problem, target, ContractionPlanKind::GroupedRowDynamicVVD, + /*laneCount=*/plan.problem.batch * plan.problem.m / rowsPerLane, rowsPerLane); + const int64_t laneCount = contractionPlan.laneCount; + SmallVector fragmentShape(blueprintType.getRank(), 1); + fragmentShape[fragmentShape.size() - 2] = rowsPerLane; + fragmentShape.back() = plan.problem.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, + plan.problem.lhsBatchShape, transposedRhs, - plan.rhsBatchShape, - plan.outputBatchShape, - plan.lhsType, - plan.rhsType, - scalarPiecesType, + plan.problem.rhsBatchShape, + plan.problem.outputBatchShape, + plan.problem.k, + rowsPerLane, + rowPiecesType, directOutType, + fragmentType, rewriter, loc); if (failed(batchOp)) return failure(); - auto result = - createBatchedDynamicOutputCompute(batchOp->getResult(0), scalarPiecesType, directOutType, rewriter, loc); + Value rowPieces = batchOp->getResult(0); + if (outputScale) { + auto scaleAttr = getHostConstDenseElementsAttr(outputScale); + if (!scaleAttr || !scaleAttr.isSplat()) + return failure(); + auto scaled = mapGraphBatchFragments( + rowPieces, rowPiecesType, rewriter, loc, [&](Value fragment, RankedTensorType fragmentType) { + auto splat = DenseElementsAttr::get(fragmentType, scaleAttr.getSplatValue()); + Value scale = arith::ConstantOp::create(rewriter, loc, fragmentType, splat); + return FailureOr( + spatial::SpatVMulOp::create(rewriter, loc, fragmentType, fragment, scale).getResult()); + }); + if (failed(scaled)) + return failure(); + rowPieces = *scaled; + if (auto scaleBatch = rowPieces.getDefiningOp()) + scaleBatch->setAttr("spat.preserve_compute_boundary", rewriter.getUnitAttr()); + } + auto result = createBatchedRowOutputBlueprint( + rowPieces, blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc); if (failed(result)) return failure(); Value finalResult = *result; if (plan.transposedResult) { - auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n}, - shapeInfo->outType.getElementType(), - shapeInfo->outType.getEncoding()); - finalResult = - ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1})) - .getResult(); + finalResult = transposeLastTwoDims(finalResult, rewriter, loc); } finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc); rewriter.replaceOp(matmulOp, finalResult); + if (foldedTranspose && foldedTranspose->use_empty()) + rewriter.eraseOp(foldedTranspose); + if (foldedMultiply && foldedMultiply->use_empty()) + rewriter.eraseOp(foldedMultiply); return success(); } + + const spatial::SpatialTargetResources& target; +}; + +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 populateMatMulRewritePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { - patterns.insert(ctx); +void populateMatMulFusionPatterns(RewritePatternSet& patterns, + MLIRContext* ctx, + const spatial::SpatialTargetResources& target) { + patterns.add(ctx, target); +} + +void populateMatMulRewritePatterns(RewritePatternSet& patterns, + MLIRContext* ctx, + const spatial::SpatialTargetResources& target) { + patterns.insert(ctx, target); + patterns.insert(ctx, target); } } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp index 8904feb..359c1d3 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp @@ -280,12 +280,12 @@ static FailureOr buildReduceMeanKeepdimsBlueprint( SmallVector fragmentStrides(fragmentOffsets.size(), 1); return spatial::SpatBlueprintOp::create( rewriter, loc, keepdimsType, batchValue, ValueRange {}, - rewriter.getStringAttr("nchw"), - rewriter.getStringAttr("fragmented"), + spatial::getNCHWLayout(rewriter.getContext()), + spatial::getFragmentedLayout(rewriter.getContext()), rewriter.getDenseI64ArrayAttr(fragmentOffsets), rewriter.getDenseI64ArrayAttr(fragmentSizes), rewriter.getStringAttr("reduce_mean_keepdims_fragments"), - rewriter.getStringAttr("fragment_assembly"), + spatial::getFragmentAssemblyMode(rewriter.getContext()), rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots), rewriter.getDenseI64ArrayAttr(sourceOffsets), diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp index 26692de..171dc0f 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp @@ -14,10 +14,10 @@ #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -32,8 +32,10 @@ static Value materializeTileTensor(PatternRewriter& rewriter, Location loc, Valu return insertStaticSlice(rewriter, loc, tile, empty, getZeroOffsets(rewriter, tileType.getRank())); } -static Value -createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type elementType, bool useMinimumValue) { +static Value createPoolFillElement(OpBuilder& rewriter, + Location loc, + Type elementType, + bool useMinimumValue) { Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); if (!useMinimumValue) return getOrCreateConstant(rewriter, anchorOp, rewriter.getZeroAttr(elementType), elementType); @@ -51,7 +53,7 @@ createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type el llvm_unreachable("unsupported pool element type"); } -static Value createPoolFillTensor(ConversionPatternRewriter& rewriter, +static Value createPoolFillTensor(OpBuilder& rewriter, Location loc, RankedTensorType tensorType, bool useMinimumValue) { @@ -59,16 +61,15 @@ static Value createPoolFillTensor(ConversionPatternRewriter& rewriter, return tensor::SplatOp::create(rewriter, loc, tensorType, fillElement); } -template -static Value createPaddedPoolInput(ConversionPatternRewriter& rewriter, +static Value createPaddedPoolInput(OpBuilder& rewriter, Location loc, - PoolOp poolOp, Value input, RankedTensorType inputType, int64_t padTop, int64_t padLeft, int64_t padBottom, - int64_t padRight) { + int64_t padRight, + bool useMinimumValue) { if (padTop == 0 && padLeft == 0 && padBottom == 0 && padRight == 0) return input; @@ -90,8 +91,8 @@ static Value createPaddedPoolInput(ConversionPatternRewriter& rewriter, padBlock->addArgument(rewriter.getIndexType(), loc); padOp.getRegion().push_back(padBlock); rewriter.setInsertionPointToStart(padBlock); - Value padValue = - createPoolFillElement(rewriter, loc, inputType.getElementType(), std::is_same_v); + Value padValue = createPoolFillElement( + rewriter, loc, inputType.getElementType(), useMinimumValue); tensor::YieldOp::create(rewriter, loc, padValue); rewriter.setInsertionPointAfter(padOp); return padOp.getResult(); @@ -160,7 +161,10 @@ struct PoolToSpatialCompute; template struct PoolToSpatialComputeBase : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; + PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetResources& target) + : OpConversionPattern(ctx), target(target) {} + + const spatial::SpatialTargetResources& target; LogicalResult matchAndRewrite(PoolOp poolOp, PoolOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const final { Location loc = poolOp.getLoc(); @@ -241,7 +245,7 @@ struct PoolToSpatialComputeBase : public OpConversionPattern { rewriter.getDenseI64ArrayAttr({padTop, padLeft, padBottom, padRight}), rewriter.getDenseI64ArrayAttr({strideHeight, strideWidth}), rewriter.getDenseI64ArrayAttr({dilationHeight, dilationWidth}), - rewriter.getStringAttr("nchw")); + spatial::getNCHWLayout(rewriter.getContext())); rewriter.replaceOp(poolOp, plan.getResult()); return success(); } @@ -251,12 +255,12 @@ struct PoolToSpatialComputeBase : public OpConversionPattern { && dilationHeight == 1 && dilationWidth == 1 && padTop == 0 && padLeft == 0 && padBottom == 0 && padRight == 0) { auto plan = spatial::SpatGlobalAveragePoolPlanOp::create( - rewriter, loc, outType, x, rewriter.getStringAttr("nchw")); + rewriter, loc, outType, x, spatial::getNCHWLayout(rewriter.getContext())); rewriter.replaceOp(poolOp, plan.getResult()); return success(); } - const int64_t xbarSize = static_cast(crossbarSize.getValue()); + const int64_t xbarSize = static_cast(target.matrixShape.rows); const int64_t channelTileCount = (channels + xbarSize - 1) / xbarSize; const int64_t outputPatchCount = batchSize * outputHeight * outputWidth; const bool countIncludePad = [&]() { @@ -292,7 +296,9 @@ struct PoolToSpatialComputeBase : public OpConversionPattern { auto computeOp = createSpatCompute(rewriter, loc, outType, {}, ValueRange {x}, [&](Value xArg) -> LogicalResult { Value paddedInput = - createPaddedPoolInput(rewriter, loc, poolOp, xArg, xType, padTop, padLeft, padBottom, padRight); + createPaddedPoolInput(rewriter, loc, xArg, xType, padTop, padLeft, + padBottom, padRight, + std::is_same_v); Value pooledOutputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType()); Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp(); @@ -424,7 +430,8 @@ struct PoolToSpatialCompute } // namespace -LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp) { +LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp, + const spatial::SpatialTargetResources&) { auto inputType = dyn_cast(planOp.getInput().getType()); auto outputType = dyn_cast(planOp.getOutput().getType()); if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()) @@ -439,6 +446,119 @@ LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp) return success(); } +FailureOr lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, + Value input, + const spatial::SpatialTargetResources& target, + PatternRewriter& rewriter) { + auto inputType = dyn_cast(input.getType()); + auto outputType = dyn_cast(planOp.getOutput().getType()); + if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape() + || inputType.getRank() != 4 || outputType.getRank() != 4) + return planOp.emitOpError("dense MaxPool lowering requires static rank-4 tensors"), failure(); + + auto kernel = planOp.getKernelShape(); + auto pads = planOp.getPads(); + auto strides = planOp.getStrides(); + auto dilations = planOp.getDilations(); + if (kernel.size() != 2 || pads.size() != 4 || strides.size() != 2 || dilations.size() != 2 + || llvm::any_of(kernel, [](int64_t value) { return value <= 0; }) + || llvm::any_of(strides, [](int64_t value) { return value <= 0; }) + || llvm::any_of(dilations, [](int64_t value) { return value <= 0; }) + || llvm::any_of(pads, [](int64_t value) { return value < 0; })) + return planOp.emitOpError("dense MaxPool lowering requires valid kernel, padding, stride, and dilation attributes"), + failure(); + + const int64_t batchSize = inputType.getDimSize(0); + const int64_t channels = inputType.getDimSize(1); + const int64_t outputHeight = outputType.getDimSize(2); + const int64_t outputWidth = outputType.getDimSize(3); + const int64_t tileWidth = std::max(1, target.matrixShape.rows); + const int64_t channelTileCount = (channels + tileWidth - 1) / tileWidth; + const int64_t outputPatchCount = batchSize * outputHeight * outputWidth; + + auto compute = createSpatCompute<1>( + rewriter, planOp.getLoc(), outputType, {}, input, + [&](Value input) -> LogicalResult { + Value paddedInput = createPaddedPoolInput( + rewriter, planOp.getLoc(), input, inputType, + pads[0], pads[1], pads[2], pads[3], /*useMinimumValue=*/true); + Value outputInit = tensor::EmptyOp::create( + rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType()); + Operation* anchor = rewriter.getInsertionBlock()->getParentOp(); + Value zero = getOrCreateIndexConstant(rewriter, anchor, 0); + Value one = getOrCreateIndexConstant(rewriter, anchor, 1); + Value patchCount = getOrCreateIndexConstant(rewriter, anchor, outputPatchCount); + Value pixelsPerBatch = getOrCreateIndexConstant( + rewriter, anchor, outputHeight * outputWidth); + Value outputWidthValue = getOrCreateIndexConstant(rewriter, anchor, outputWidth); + Value strideHeight = getOrCreateIndexConstant(rewriter, anchor, strides[0]); + Value strideWidth = getOrCreateIndexConstant(rewriter, anchor, strides[1]); + + auto loop = buildNormalizedScfFor( + rewriter, planOp.getLoc(), zero, patchCount, one, ValueRange {outputInit}, + [&](OpBuilder&, Location loc, Value patch, ValueRange iterArgs, + SmallVectorImpl& yielded) { + Value batch = arith::DivUIOp::create(rewriter, loc, patch, pixelsPerBatch); + Value batchPatch = arith::RemUIOp::create(rewriter, loc, patch, pixelsPerBatch); + Value outputRow = arith::DivUIOp::create(rewriter, loc, batchPatch, outputWidthValue); + Value outputColumn = arith::RemUIOp::create(rewriter, loc, batchPatch, outputWidthValue); + Value windowRow = arith::MulIOp::create(rewriter, loc, outputRow, strideHeight); + Value windowColumn = arith::MulIOp::create(rewriter, loc, outputColumn, strideWidth); + Value updated = iterArgs.front(); + + for (int64_t tile = 0; tile < channelTileCount; ++tile) { + const int64_t tileChannels = std::min(tileWidth, channels - tile * tileWidth); + auto tileType = RankedTensorType::get( + {1, tileChannels, 1, 1}, outputType.getElementType()); + Value reduced = createPoolFillTensor( + rewriter, loc, tileType, /*useMinimumValue=*/true); + for (int64_t kernelRow = 0; kernelRow < kernel[0]; ++kernelRow) { + Value sourceRow = windowRow; + if (kernelRow * dilations[0] != 0) + sourceRow = arith::AddIOp::create( + rewriter, loc, sourceRow, + getOrCreateIndexConstant(rewriter, anchor, kernelRow * dilations[0])); + for (int64_t kernelColumn = 0; kernelColumn < kernel[1]; ++kernelColumn) { + Value sourceColumn = windowColumn; + if (kernelColumn * dilations[1] != 0) + sourceColumn = arith::AddIOp::create( + rewriter, loc, sourceColumn, + getOrCreateIndexConstant(rewriter, anchor, kernelColumn * dilations[1])); + Value point = tensor::ExtractSliceOp::create( + rewriter, loc, tileType, paddedInput, + SmallVector { + batch, rewriter.getIndexAttr(tile * tileWidth), sourceRow, sourceColumn}, + SmallVector { + rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}, + getUnitStrides(rewriter, 4)); + point = materializeTileTensor(rewriter, loc, point); + reduced = spatial::SpatVMaxOp::create( + rewriter, loc, tileType, reduced, point); + } + } + updated = tensor::InsertSliceOp::create( + rewriter, loc, reduced, updated, + SmallVector { + batch, rewriter.getIndexAttr(tile * tileWidth), outputRow, outputColumn}, + SmallVector { + rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}, + getUnitStrides(rewriter, 4)); + } + yielded.push_back(updated); + return success(); + }); + if (failed(loop)) + return failure(); + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), loop->results.front()); + return success(); + }); + if (failed(compute)) + return failure(); + return compute->getResult(0); +} + static Value createClampedPoolIndexTable(PatternRewriter& rewriter, Operation* anchorOp, int64_t outputSize, @@ -496,13 +616,15 @@ static Value extractPoolIndex(PatternRewriter& rewriter, } FailureOr lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, + Value input, std::optional rowStripInput, + const spatial::SpatialTargetResources& target, PatternRewriter& rewriter) { - if (failed(canLowerMaxPoolPlanToRowStrip(planOp))) + if (failed(canLowerMaxPoolPlanToRowStrip(planOp, target))) return failure(); Location loc = planOp.getLoc(); - auto inputType = cast(planOp.getInput().getType()); + auto inputType = cast(input.getType()); auto outputType = cast(planOp.getOutput().getType()); const int64_t channels = inputType.getDimSize(1); const int64_t inputHeight = inputType.getDimSize(2); @@ -511,9 +633,9 @@ FailureOr lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, const int64_t outputWidth = outputType.getDimSize(3); const int64_t kernelHeight = planOp.getKernelShape()[0]; const int64_t kernelWidth = planOp.getKernelShape()[1]; - Value input = rowStripInput.value_or(planOp.getInput()); - auto actualInputType = dyn_cast(input.getType()); - FailureOr physicalValue = describeRowStripPhysicalValue(input, inputType); + Value actualInput = rowStripInput.value_or(input); + auto actualInputType = dyn_cast(actualInput.getType()); + FailureOr physicalValue = describeRowStripPhysicalValue(actualInput, inputType); const bool physicalInput = succeeded(physicalValue); if (!physicalInput && actualInputType != inputType) return failure(); @@ -561,7 +683,7 @@ FailureOr lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, TypeRange {outputStorageType}, outputHeight * tilesPerRow, {}, - ValueRange {input}, + ValueRange {actualInput}, [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult { SmallVector inputRows; inputRows.reserve(kernelHeight); @@ -590,8 +712,8 @@ FailureOr lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, rewriter.getIndexAttr(1), rewriter.getIndexAttr(inputWidth)}, getUnitStrides(rewriter, 4)); - inputRows.push_back(ONNXTransposeOp::create( - rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1}))); + inputRows.push_back(createLinalgTranspose( + nchw, inputFragmentType, {0, 2, 3, 1}, rewriter, loc)); } } @@ -685,7 +807,8 @@ FailureOr lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, return batch->getResult(0); } -LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp) { +LogicalResult canLowerGlobalAveragePoolPlanToRowStrip( + spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetResources&) { auto inputType = dyn_cast(planOp.getInput().getType()); auto outputType = dyn_cast(planOp.getOutput().getType()); if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()) @@ -697,22 +820,101 @@ LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAverage return success(); } +FailureOr lowerDenseGlobalAveragePoolPlan( + spatial::SpatGlobalAveragePoolPlanOp planOp, + Value input, + const spatial::SpatialTargetResources& target, + PatternRewriter& rewriter) { + auto inputType = dyn_cast(input.getType()); + auto outputType = dyn_cast(planOp.getOutput().getType()); + if (!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) + || outputType.getDimSize(2) != 1 || outputType.getDimSize(3) != 1) + return planOp.emitOpError("dense global AveragePool lowering requires static rank-4 floating-point tensors"), + failure(); + auto elementType = dyn_cast(inputType.getElementType()); + if (!elementType) + return planOp.emitOpError("dense global AveragePool lowering requires floating-point tensors"), + failure(); + + const int64_t channels = inputType.getDimSize(1); + const int64_t height = inputType.getDimSize(2); + const int64_t width = inputType.getDimSize(3); + const int64_t tileWidth = std::max(1, target.matrixShape.rows); + const int64_t channelTileCount = (channels + tileWidth - 1) / tileWidth; + const double scaleValue = 1.0 / static_cast(height * width); + + auto compute = createSpatCompute<1>( + rewriter, planOp.getLoc(), outputType, {}, input, + [&](Value input) -> LogicalResult { + Value output = tensor::EmptyOp::create( + rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType()); + Operation* anchor = rewriter.getInsertionBlock()->getParentOp(); + for (int64_t tile = 0; tile < channelTileCount; ++tile) { + const int64_t tileChannels = std::min(tileWidth, channels - tile * tileWidth); + auto tileType = RankedTensorType::get( + {1, tileChannels, 1, 1}, outputType.getElementType()); + Value reduced = createPoolFillTensor( + rewriter, planOp.getLoc(), tileType, /*useMinimumValue=*/false); + for (int64_t row = 0; row < height; ++row) { + for (int64_t column = 0; column < width; ++column) { + Value point = tensor::ExtractSliceOp::create( + rewriter, planOp.getLoc(), tileType, input, + SmallVector { + rewriter.getIndexAttr(0), rewriter.getIndexAttr(tile * tileWidth), + rewriter.getIndexAttr(row), rewriter.getIndexAttr(column)}, + SmallVector { + rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}, + getUnitStrides(rewriter, 4)); + point = materializeTileTensor(rewriter, planOp.getLoc(), point); + reduced = spatial::SpatVAddOp::create( + rewriter, planOp.getLoc(), tileType, reduced, point); + } + } + auto scaleAttr = DenseElementsAttr::get( + tileType, rewriter.getFloatAttr(elementType, scaleValue)); + Value scale = getOrCreateConstant(rewriter, anchor, scaleAttr, tileType); + reduced = spatial::SpatVMulOp::create( + rewriter, planOp.getLoc(), tileType, reduced, scale); + output = tensor::InsertSliceOp::create( + rewriter, planOp.getLoc(), reduced, output, + SmallVector { + rewriter.getIndexAttr(0), rewriter.getIndexAttr(tile * tileWidth), + rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)}, + SmallVector { + rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels), + rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)}, + getUnitStrides(rewriter, 4)); + } + spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), output); + return success(); + }); + if (failed(compute)) + return failure(); + return compute->getResult(0); +} + FailureOr lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp, + Value input, std::optional rowStripInput, + const spatial::SpatialTargetResources& target, PatternRewriter& rewriter) { - if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp))) + if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp, target))) return failure(); Location loc = planOp.getLoc(); - auto inputType = cast(planOp.getInput().getType()); + auto inputType = cast(input.getType()); auto outputType = cast(planOp.getOutput().getType()); auto elementType = dyn_cast(inputType.getElementType()); if (!elementType) return failure(); - Value input = rowStripInput.value_or(planOp.getInput()); - auto actualInputType = dyn_cast(input.getType()); - FailureOr physicalValue = describeRowStripPhysicalValue(input, inputType); + Value actualInput = rowStripInput.value_or(input); + auto actualInputType = dyn_cast(actualInput.getType()); + FailureOr physicalValue = describeRowStripPhysicalValue(actualInput, inputType); const bool physicalInput = succeeded(physicalValue); if (!physicalInput && actualInputType != inputType) return failure(); @@ -742,7 +944,7 @@ FailureOr lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo TypeRange {outputStorageType}, tilesPerRow, ValueRange {zero, scale}, - ValueRange {input}, + ValueRange {actualInput}, [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult { Value reduced = args.weights[0]; for (int64_t row = 0; row < height; ++row) { @@ -777,8 +979,8 @@ FailureOr lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)}, getUnitStrides(rewriter, 4)); - fragment = ONNXTransposeOp::create( - rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})); + fragment = createLinalgTranspose( + nchw, inputFragmentType, {0, 2, 3, 1}, rewriter, loc); } for (int64_t column = 0; column < width; ++column) { Value point = tensor::ExtractSliceOp::create( @@ -811,9 +1013,11 @@ FailureOr lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo return batch->getResult(0); } -void populatePoolPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { - patterns.insert>(ctx); - patterns.insert>(ctx); +void populatePoolPatterns(RewritePatternSet& patterns, + MLIRContext* ctx, + const spatial::SpatialTargetResources& target) { + patterns.insert>(ctx, target); + patterns.insert>(ctx, target); } } // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Relu.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Relu.cpp index 008865e..6866b7e 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Relu.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Relu.cpp @@ -17,7 +17,7 @@ struct ReluToSpatialCompute : OpConversionPattern { Location loc = reluOp.getLoc(); Type resultType = reluOp.getResult().getType(); auto reluPlan = spatial::SpatReluPlanOp::create( - rewriter, loc, resultType, adaptor.getX(), rewriter.getStringAttr("nchw")); + rewriter, loc, resultType, adaptor.getX(), spatial::getNCHWLayout(rewriter.getContext())); rewriter.replaceOp(reluOp, reluPlan.getResult()); return success(); } diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp index a23076e..07a1a2d 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp @@ -32,7 +32,8 @@ struct Concat : public OpConversionPattern { return type && type.hasStaticShape() && type.getRank() == 4; })) { rewriter.replaceOpWithNewOp( - maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis), rewriter.getStringAttr("nchw")); + maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis), + spatial::getNCHWLayout(rewriter.getContext())); return success(); } diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Flatten.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Flatten.cpp index 9fdb01b..5eb5b00 100644 --- a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Flatten.cpp +++ b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Flatten.cpp @@ -4,11 +4,10 @@ #include "llvm/ADT/SmallVector.h" #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp" #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" #include "src/Dialect/ONNX/ONNXOps.hpp" @@ -48,12 +47,12 @@ static SmallVector getExpandFrom1DReassociation(int64_t ra return reassociation; } -static Value buildFlatten(Value input, - RankedTensorType sourceType, - RankedTensorType resultType, - int64_t axis, - ConversionPatternRewriter& rewriter, - Location loc) { +static Value buildFlattenBody(Value input, + RankedTensorType sourceType, + RankedTensorType resultType, + int64_t axis, + PatternRewriter& rewriter, + Location loc) { if (sourceType == resultType) return input; @@ -76,6 +75,25 @@ static Value buildFlatten(Value input, rewriter, loc, resultType, flattened, getExpandFrom1DReassociation(resultType.getRank())); } +static Value buildFlatten(Value input, + RankedTensorType sourceType, + RankedTensorType resultType, + int64_t axis, + PatternRewriter& rewriter, + Location loc) { + if (spatial::isAnySpatialComputeLike(rewriter.getInsertionBlock()->getParentOp())) + return buildFlattenBody(input, sourceType, resultType, axis, rewriter, loc); + + auto compute = createSpatCompute<1>( + rewriter, loc, TypeRange {resultType}, {}, ValueRange {input}, + [&](Value computeInput) { + spatial::SpatYieldOp::create( + rewriter, loc, + buildFlattenBody(computeInput, sourceType, resultType, axis, rewriter, loc)); + }); + return compute.getResult(0); +} + struct Flatten : OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -98,53 +116,53 @@ struct Flatten : OpConversionPattern { if (resultType.getShape()[0] != outerDim || resultType.getShape()[1] != innerDim) return failure(); - auto replaceWithFlatten = [&](auto build) -> LogicalResult { - Value flattened = materializeOrComputeUnary(adaptor.getInput(), resultType, rewriter, flattenOp.getLoc(), build); - rewriter.replaceOp(flattenOp, flattened); - return success(); - }; - - return replaceWithFlatten([&](Value input) { - return buildFlatten(input, sourceType, resultType, *axis, rewriter, flattenOp.getLoc()); - }); + auto plan = spatial::SpatFlattenPlanOp::create( + rewriter, flattenOp.getLoc(), resultType, adaptor.getInput(), + rewriter.getI64IntegerAttr(*axis), + spatial::getNCHWLayout(rewriter.getContext())); + rewriter.replaceOp(flattenOp, plan.getOutput()); + return success(); } }; struct RowStripFlattenAnalysis { spatial::SpatGraphComputeBatch consumer; - tensor::CollapseShapeOp collapse; RankedTensorType sourceType; RankedTensorType resultType; RankedTensorType weightType; DenseElementsAttr weight; }; -static FailureOr analyzeRowStripFlatten(spatial::SpatGraphCompute flattenOp) { - if (flattenOp.getWeights().size() != 0 || flattenOp.getInputs().size() != 1 - || flattenOp.getOutputs().size() != 1) +static FailureOr analyzeRowStripFlatten( + spatial::SpatFlattenPlanOp flattenOp, const spatial::SpatialTargetResources& target) { + if (flattenOp.getAxis() != 1) return failure(); - auto sourceType = dyn_cast(flattenOp.getInputs().front().getType()); - auto resultType = dyn_cast(flattenOp.getOutputs().front().getType()); + auto sourceType = dyn_cast(flattenOp.getInput().getType()); + auto resultType = dyn_cast(flattenOp.getOutput().getType()); if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || sourceType.getRank() != 4 || resultType.getRank() != 2 || sourceType.getDimSize(0) != 1 || resultType.getDimSize(0) != 1 || resultType.getDimSize(1) != sourceType.getNumElements()) return failure(); const int64_t channels = sourceType.getDimSize(1); - const int64_t xbarDim = static_cast(crossbarSize.getValue()); + const int64_t xbarDim = static_cast(target.matrixShape.rows); if (channels > xbarDim && channels % xbarDim != 0) return failure(); - auto yieldOp = dyn_cast(flattenOp.getBody().front().getTerminator()); - if (!yieldOp || yieldOp.getOutputs().size() != 1) + Value consumerInput = flattenOp.getOutput(); + Operation* consumerOp = nullptr; + while (consumerInput.hasOneUse()) { + Operation* user = *consumerInput.getUsers().begin(); + if (auto materialize = dyn_cast(user)) { + consumerInput = materialize.getOutput(); + continue; + } + consumerOp = user; + break; + } + if (!consumerOp) return failure(); - auto collapse = yieldOp.getOutputs().front().getDefiningOp(); - if (!collapse || collapse.getSrc() != *flattenOp.getInputArgument(0)) - return failure(); - - if (!flattenOp.getResult(0).hasOneUse()) - return failure(); - auto consumer = dyn_cast(*flattenOp.getResult(0).getUsers().begin()); - if (!consumer || consumer.getInputs().size() != 1 || consumer.getInputs().front() != flattenOp.getResult(0) + auto consumer = dyn_cast(consumerOp); + if (!consumer || consumer.getInputs().size() != 1 || consumer.getInputs().front() != consumerInput || consumer.getWeights().size() != 1) return failure(); auto weightType = dyn_cast(consumer.getWeights().front().getType()); @@ -155,21 +173,34 @@ static FailureOr analyzeRowStripFlatten(spatial::SpatGr if (llvm::none_of(consumer.getBody().getOps(), [](spatial::SpatVMMOp) { return true; })) return failure(); - return RowStripFlattenAnalysis {consumer, collapse, sourceType, resultType, weightType, weight}; + return RowStripFlattenAnalysis {consumer, sourceType, resultType, weightType, weight}; } } // namespace void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add(ctx); } -LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp) { - return succeeded(analyzeRowStripFlatten(flattenOp)) ? success() : failure(); +FailureOr lowerDenseFlattenPlan(spatial::SpatFlattenPlanOp planOp, + Value input, + PatternRewriter& rewriter) { + auto sourceType = dyn_cast(input.getType()); + auto resultType = dyn_cast(planOp.getOutput().getType()); + if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape()) + return failure(); + return buildFlatten(input, sourceType, resultType, planOp.getAxis(), rewriter, + planOp.getLoc()); +} + +LogicalResult canLowerFlattenFromRowStrip(spatial::SpatFlattenPlanOp flattenOp, + const spatial::SpatialTargetResources& target) { + return succeeded(analyzeRowStripFlatten(flattenOp, target)) ? success() : failure(); } LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input, - spatial::SpatGraphCompute flattenOp, + spatial::SpatFlattenPlanOp flattenOp, + const spatial::SpatialTargetResources& target, PatternRewriter& rewriter) { - FailureOr analysis = analyzeRowStripFlatten(flattenOp); + FailureOr analysis = analyzeRowStripFlatten(flattenOp, target); if (failed(analysis)) return failure(); auto storageType = dyn_cast(input.storage.getType()); @@ -204,19 +235,20 @@ LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input, analysis->weightType); analysis->consumer->setOperand(0, reorderedWeight); - BlockArgument flattenInput = *flattenOp.getInputArgument(0); - flattenOp.getInputsMutable().assign(input.storage); - flattenInput.setType(storageType); - - OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPoint(analysis->collapse); - auto flatType = RankedTensorType::get( - {storageType.getNumElements()}, storageType.getElementType(), storageType.getEncoding()); - Value flat = tensor::CollapseShapeOp::create( - rewriter, flattenOp.getLoc(), flatType, flattenInput, getCollapseTo1DReassociation(storageType.getRank())); - Value logicalInput = tensor::ExpandShapeOp::create( - rewriter, flattenOp.getLoc(), analysis->resultType, flat, getExpandFrom1DReassociation(2)); - rewriter.replaceOp(analysis->collapse, logicalInput); + auto compute = createSpatCompute<1>( + rewriter, flattenOp.getLoc(), TypeRange {analysis->resultType}, {}, + ValueRange {input.storage}, [&](Value storage) { + auto flatType = RankedTensorType::get( + {storageType.getNumElements()}, storageType.getElementType(), storageType.getEncoding()); + Value flat = tensor::CollapseShapeOp::create( + rewriter, flattenOp.getLoc(), flatType, storage, + getCollapseTo1DReassociation(storageType.getRank())); + Value logicalInput = tensor::ExpandShapeOp::create( + rewriter, flattenOp.getLoc(), analysis->resultType, flat, + getExpandFrom1DReassociation(2)); + spatial::SpatYieldOp::create(rewriter, flattenOp.getLoc(), logicalInput); + }); + rewriter.replaceOp(flattenOp, compute.getResult(0)); return success(); } diff --git a/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp b/src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp index 769a943..e1db67f 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/Passes/Transforms/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,41 @@ 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(), spatial::getNCHWLayout(rewriter.getContext())); + rewriter.replaceOp(resizeOp, plan.getResult()); return success(); } }; } // namespace +LogicalResult canLowerResizeNearestPlanToRowStrip( + spatial::SpatResizeNearestPlanOp planOp, + const spatial::SpatialTargetResources&) { + 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, Value input, + std::optional rowStripInput, + const spatial::SpatialTargetResources&, + PatternRewriter& rewriter) { + auto inputType = cast(input.getType()); + auto outputType = cast(planOp.getOutput().getType()); + if (rowStripInput) + return buildRowStripNearestResize( + *rowStripInput, inputType, outputType, rewriter, planOp.getLoc()); + return buildDenseNearestResize( + input, 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..b77a9dc 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() || !spatial::isFragmentAssembly(blueprint.getMode()) + || !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(), + spatial::getFragmentedLayout(rewriter.getContext()), + 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 deleted file mode 100644 index be2308e..0000000 --- a/src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp +++ /dev/null @@ -1,44 +0,0 @@ -#pragma once - -#include - -#include "mlir/IR/PatternMatch.h" -#include "mlir/Support/LogicalResult.h" - -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" - -namespace onnx_mlir { - -struct RowStripPhysicalValue; - -mlir::FailureOr -lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp, - std::optional rowStripInput, - bool emitRowStripLayout, - mlir::PatternRewriter& rewriter); - -mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp); -mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp); - -mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp); - -mlir::FailureOr -lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp, - std::optional rowStripInput, - mlir::PatternRewriter& rewriter); - -mlir::LogicalResult -canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp); - -mlir::FailureOr -lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp, - std::optional rowStripInput, - mlir::PatternRewriter& rewriter); - -mlir::LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp); - -mlir::LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input, - spatial::SpatGraphCompute flattenOp, - mlir::PatternRewriter& rewriter); - -} // namespace onnx_mlir diff --git a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp b/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp deleted file mode 100644 index df69231..0000000 --- a/src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp +++ /dev/null @@ -1,341 +0,0 @@ -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Pass/Pass.h" - -#include "llvm/ADT/DenseMap.h" - -#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" -#include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp" -#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" - -using namespace mlir; - -namespace onnx_mlir { -namespace { - -static constexpr StringLiteral kLogicalLayout = "nchw"; -static constexpr StringLiteral kDenseLayout = "dense_nchw"; -static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip"; - -enum class SelectedLayout { - DenseNchw, - PixelMajorRowStrip, -}; - -static SelectedLayout getSelectedLayout(llvm::DenseMap& layouts, Value value) { - auto it = layouts.find(value); - return it == layouts.end() ? SelectedLayout::DenseNchw : it->second; -} - -static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap& layouts) { - if (auto reluPlan = dyn_cast(user)) - return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; - if (auto siluPlan = dyn_cast(user)) - return getSelectedLayout(layouts, siluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; - if (auto biasAddPlan = dyn_cast(user)) - return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; - if (auto addPlan = dyn_cast(user)) - return getSelectedLayout(layouts, addPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; - if (auto concatPlan = dyn_cast(user)) - return getSelectedLayout(layouts, concatPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; - if (auto convPlan = dyn_cast(user)) - return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; - if (auto maxPoolPlan = dyn_cast(user)) - return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; - if (auto averagePoolPlan = dyn_cast(user)) - return getSelectedLayout(layouts, averagePoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip; - if (auto flattenCompute = dyn_cast(user)) - return succeeded(canLowerFlattenFromRowStrip(flattenCompute)); - return false; -} - -static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap& layouts) { - for (Operation* user : value.getUsers()) { - if (usesSelectedRowStrip(user, layouts)) - continue; - // Dense-only users must be materialized explicitly. - continue; - } - return true; -} - -static bool canConsumeRowStripAsUser(Operation* user) { - if (isa(user)) - return true; - if (auto biasAddPlan = dyn_cast(user)) { - auto resultType = dyn_cast(biasAddPlan.getOutput().getType()); - return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType); - } - if (isa(user)) - return true; - if (isa(user)) - return true; - if (auto convPlan = dyn_cast(user)) - return succeeded(canConsumeAndProduceRowStrip(convPlan)); - if (auto maxPoolPlan = dyn_cast(user)) - return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)); - if (auto averagePoolPlan = dyn_cast(user)) - return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan)); - return false; -} - -static bool hasRowStripConsumer(Value value) { - for (Operation* user : value.getUsers()) - if (canConsumeRowStripAsUser(user)) - return true; - return false; -} - -static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan, - llvm::DenseMap& layouts) { - SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput()); - if (inputLayout == SelectedLayout::PixelMajorRowStrip) - return succeeded(canConsumeAndProduceRowStrip(convPlan)); - return succeeded(canLowerConvPlanToRowStrip(convPlan)); -} - -static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan, - llvm::DenseMap& layouts) { - if (!canSelectConvRowStrip(convPlan, layouts)) - return SelectedLayout::DenseNchw; - if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts)) - return SelectedLayout::DenseNchw; - return SelectedLayout::PixelMajorRowStrip; -} - -static SelectedLayout chooseActivationLayout(Value input, - Value result, - llvm::DenseMap& layouts) { - if (getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip) - return SelectedLayout::DenseNchw; - if (!allUsersCanHandleRowStrip(result, layouts)) - return SelectedLayout::DenseNchw; - return SelectedLayout::PixelMajorRowStrip; -} - -static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan, - llvm::DenseMap& layouts) { - if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip) - return SelectedLayout::DenseNchw; - auto resultType = dyn_cast(biasAddPlan.getOutput().getType()); - if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType)) - return SelectedLayout::DenseNchw; - if (!hasRowStripConsumer(biasAddPlan.getResult())) - return SelectedLayout::DenseNchw; - if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts)) - return SelectedLayout::DenseNchw; - return SelectedLayout::PixelMajorRowStrip; -} - -static SelectedLayout chooseAddLayout(spatial::SpatAddPlanOp addPlan, llvm::DenseMap& layouts) { - if (getSelectedLayout(layouts, addPlan.getLhs()) != SelectedLayout::PixelMajorRowStrip - || getSelectedLayout(layouts, addPlan.getRhs()) != SelectedLayout::PixelMajorRowStrip) - return SelectedLayout::DenseNchw; - if (!allUsersCanHandleRowStrip(addPlan.getResult(), layouts)) - return SelectedLayout::DenseNchw; - return SelectedLayout::PixelMajorRowStrip; -} - -static SelectedLayout chooseConcatLayout(spatial::SpatConcatPlanOp concatPlan, - llvm::DenseMap& layouts) { - if (llvm::any_of(concatPlan.getInputs(), [&](Value input) { - return getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip; - })) - return SelectedLayout::DenseNchw; - if (!allUsersCanHandleRowStrip(concatPlan.getResult(), layouts)) - return SelectedLayout::DenseNchw; - return SelectedLayout::PixelMajorRowStrip; -} - -static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) { - return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::PixelMajorRowStrip - : SelectedLayout::DenseNchw; -} - -static SelectedLayout chooseGlobalAveragePoolLayout( - spatial::SpatGlobalAveragePoolPlanOp averagePoolPlan) { - return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan)) - ? SelectedLayout::PixelMajorRowStrip - : SelectedLayout::DenseNchw; -} - -static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) { - auto outputType = cast(value.getType()); - auto [offsets, sizes] = buildRowStripMetadata(outputType); - return spatial::SpatBlueprintOp::create(rewriter, - value.getLoc(), - outputType, - value, - ValueRange {}, - rewriter.getStringAttr(kLogicalLayout), - rewriter.getStringAttr(kRowStripLayout), - rewriter.getDenseI64ArrayAttr(offsets), - rewriter.getDenseI64ArrayAttr(sizes), - rewriter.getStringAttr(kRowStripIndexMap), - nullptr, - nullptr, - nullptr, - nullptr, - nullptr, - nullptr, - nullptr); -} - -static void materializeDenseUses(IRRewriter& rewriter, - Value layoutValue, - llvm::DenseMap& layouts) { - SmallVector denseUses; - for (OpOperand& use : layoutValue.getUses()) { - if (usesSelectedRowStrip(use.getOwner(), layouts)) - continue; - denseUses.push_back(&use); - } - - for (OpOperand* use : denseUses) { - Operation* owner = use->getOwner(); - rewriter.setInsertionPoint(owner); - auto materialized = spatial::SpatMaterializeLayoutOp::create(rewriter, - owner->getLoc(), - use->get().getType(), - use->get(), - rewriter.getStringAttr(kLogicalLayout), - rewriter.getStringAttr(kRowStripLayout), - rewriter.getStringAttr(kDenseLayout)); - use->set(materialized.getResult()); - } -} - -struct SpatialLayoutPlanningPass final : PassWrapper> { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass) - - StringRef getArgument() const override { return "spatial-layout-planning"; } - StringRef getDescription() const override { return "Select conservative Spatial layouts and insert reconciliation barriers."; } - - void runOnOperation() override { - auto entryFunc = getPimEntryFunc(getOperation()); - if (failed(entryFunc)) { - getOperation().emitError("failed to locate the PIM entry function during Spatial layout planning"); - signalPassFailure(); - return; - } - - func::FuncOp funcOp = *entryFunc; - IRRewriter rewriter(&getContext()); - llvm::DenseMap layouts; - - bool changed = true; - while (changed) { - changed = false; - for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) { - if (auto convPlan = dyn_cast(&op)) { - SelectedLayout selected = chooseConvLayout(convPlan, layouts); - if (layouts[convPlan.getResult()] != selected) { - layouts[convPlan.getResult()] = selected; - changed = true; - } - continue; - } - if (auto reluPlan = dyn_cast(&op)) { - SelectedLayout selected = chooseActivationLayout(reluPlan.getInput(), reluPlan.getResult(), layouts); - if (layouts[reluPlan.getResult()] != selected) { - layouts[reluPlan.getResult()] = selected; - changed = true; - } - continue; - } - if (auto siluPlan = dyn_cast(&op)) { - SelectedLayout selected = chooseActivationLayout(siluPlan.getInput(), siluPlan.getResult(), layouts); - if (layouts[siluPlan.getResult()] != selected) { - layouts[siluPlan.getResult()] = selected; - changed = true; - } - continue; - } - if (auto biasAddPlan = dyn_cast(&op)) { - SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts); - if (layouts[biasAddPlan.getResult()] != selected) { - layouts[biasAddPlan.getResult()] = selected; - changed = true; - } - continue; - } - if (auto addPlan = dyn_cast(&op)) { - SelectedLayout selected = chooseAddLayout(addPlan, layouts); - if (layouts[addPlan.getResult()] != selected) { - layouts[addPlan.getResult()] = selected; - changed = true; - } - continue; - } - if (auto concatPlan = dyn_cast(&op)) { - SelectedLayout selected = chooseConcatLayout(concatPlan, layouts); - if (layouts[concatPlan.getResult()] != selected) { - layouts[concatPlan.getResult()] = selected; - changed = true; - } - continue; - } - if (auto maxPoolPlan = dyn_cast(&op)) { - SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan); - if (layouts[maxPoolPlan.getResult()] != selected) { - layouts[maxPoolPlan.getResult()] = selected; - changed = true; - } - continue; - } - if (auto averagePoolPlan = dyn_cast(&op)) { - SelectedLayout selected = chooseGlobalAveragePoolLayout(averagePoolPlan); - if (layouts[averagePoolPlan.getResult()] != selected) { - layouts[averagePoolPlan.getResult()] = selected; - changed = true; - } - continue; - } - } - } - - for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) { - Value producedValue; - if (auto convPlan = dyn_cast(&op)) - producedValue = convPlan.getResult(); - else if (auto biasAddPlan = dyn_cast(&op)) - producedValue = biasAddPlan.getResult(); - else if (auto addPlan = dyn_cast(&op)) - producedValue = addPlan.getResult(); - else if (auto concatPlan = dyn_cast(&op)) - producedValue = concatPlan.getResult(); - else if (auto reluPlan = dyn_cast(&op)) - producedValue = reluPlan.getResult(); - else if (auto siluPlan = dyn_cast(&op)) - producedValue = siluPlan.getResult(); - else if (auto maxPoolPlan = dyn_cast(&op)) - producedValue = maxPoolPlan.getResult(); - else if (auto averagePoolPlan = dyn_cast(&op)) - producedValue = averagePoolPlan.getResult(); - else - continue; - - if (getSelectedLayout(layouts, producedValue) != SelectedLayout::PixelMajorRowStrip) - continue; - - rewriter.setInsertionPointAfter(&op); - auto blueprint = insertRowStripBlueprint(rewriter, producedValue); - rewriter.replaceAllUsesExcept(producedValue, blueprint.getResult(), blueprint); - materializeDenseUses(rewriter, blueprint.getResult(), layouts); - } - if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) { - getOperation().emitError("logical Spatial graph verification failed after SpatialLayoutPlanning"); - signalPassFailure(); - } - } -}; - -} // namespace - -std::unique_ptr createSpatialLayoutPlanningPass() { return std::make_unique(); } - -} // namespace onnx_mlir diff --git a/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp b/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp index 7a58074..83e9384 100644 --- a/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp @@ -149,11 +149,10 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe auto blueprint = dyn_cast(use.getOwner()); if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType()) return failure(); - std::optional mode = blueprint.getMode(); std::optional> operandIndicesAttr = blueprint.getFragmentOperandIndices(); std::optional> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets(); std::optional> sourceSlotsAttr = blueprint.getFragmentSourceSlots(); - if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr) + if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr) return failure(); if (!blueprint.getOutput().hasOneUse() || !isa(*blueprint.getOutput().getUsers().begin())) return failure(); @@ -418,8 +417,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul rewriter.setInsertionPointToEnd(newBlock); if (auto blueprint = dyn_cast(op)) { - std::optional modeAttr = blueprint.getMode(); - if (modeAttr && *modeAttr == "fragment_assembly") { + if (spatial::isFragmentAssembly(blueprint.getMode())) { for (Operation* user : blueprint.getOutput().getUsers()) { if (!isa(user)) return blueprint.emitOpError( @@ -483,8 +481,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul auto hostTargetType = cast(hostTarget.getType()); if (auto blueprint = insertSlice.getSource().getDefiningOp()) { - std::optional modeAttr = blueprint.getMode(); - if (modeAttr && *modeAttr == "fragment_assembly") { + if (spatial::isFragmentAssembly(blueprint.getMode())) { FailureOr> fragmentAssemblyCopies = collectFragmentAssemblyCopiesFromBlueprint(blueprint, mapper, /*lane=*/0, /*hostTargetIndex=*/0); if (failed(fragmentAssemblyCopies)) 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..318a1e2 100644 --- a/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp @@ -42,12 +42,11 @@ static FailureOr lowerFragmentAssemblyBlueprint(IRRewriter& rewriter, if (!resultType || !resultType.hasStaticShape()) return blueprint.emitOpError("fragment assembly lowering requires a static ranked tensor result"); - std::optional modeAttr = blueprint.getMode(); std::optional> operandIndicesAttr = blueprint.getFragmentOperandIndices(); std::optional> sourceSlotsAttr = blueprint.getFragmentSourceSlots(); std::optional> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets(); std::optional> fragmentStridesAttr = blueprint.getFragmentStrides(); - if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceSlotsAttr + if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceSlotsAttr || !sourceOffsetsAttr || !fragmentStridesAttr) return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata"); @@ -71,6 +70,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) { @@ -193,8 +202,7 @@ static bool isHostMaterializableHelperOp(Operation* op) { if (isa(op) || op->hasTrait()) return true; if (auto blueprint = dyn_cast(op)) { - std::optional mode = blueprint.getMode(); - return mode && *mode == "fragment_assembly"; + return spatial::isFragmentAssembly(blueprint.getMode()); } return isShapingOnlyOp(op) || isPureIndexComputationOp(op); } @@ -281,8 +289,7 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule } for (Operation& op : block.without_terminator()) { if (auto blueprint = dyn_cast(op)) { - std::optional modeAttr = blueprint.getMode(); - if (modeAttr && *modeAttr == "fragment_assembly") { + if (spatial::isFragmentAssembly(blueprint.getMode())) { auto lowered = lowerFragmentAssemblyBlueprint(rewriter, blueprint, mapping); if (failed(lowered)) return false; diff --git a/src/PIM/Conversion/SpatialToPim/Patterns.cpp b/src/PIM/Conversion/SpatialToPim/Patterns.cpp index 68ede56..ebb9f34 100644 --- a/src/PIM/Conversion/SpatialToPim/Patterns.cpp +++ b/src/PIM/Conversion/SpatialToPim/Patterns.cpp @@ -22,8 +22,7 @@ struct LowerFragmentAssemblyBlueprintPattern LogicalResult matchAndRewrite(spatial::SpatBlueprintOp op, OpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override { - std::optional modeAttr = op.getMode(); - if (!modeAttr || *modeAttr != "fragment_assembly") + if (!spatial::isFragmentAssembly(op.getMode())) return failure(); auto resultType = dyn_cast(op.getOutput().getType()); @@ -49,6 +48,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/Conversion/SpatialToPim/ReturnPathNormalization.cpp b/src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp index 0c8e9a6..b33c6d0 100644 --- a/src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp +++ b/src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp @@ -158,8 +158,7 @@ analyzeTopLevelFragmentAssemblyUses(Value value) { auto blueprint = dyn_cast(use.getOwner()); if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType()) return failure(); - std::optional mode = blueprint.getMode(); - if (!mode || *mode != "fragment_assembly") + if (!spatial::isFragmentAssembly(blueprint.getMode())) return failure(); if (!blueprint.getOutput().hasOneUse() || !isa(*blueprint.getOutput().getUsers().begin())) return failure(); @@ -819,8 +818,7 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret } if (auto blueprint = dyn_cast(op)) { - std::optional mode = blueprint.getMode(); - if (mode && *mode == "fragment_assembly") { + if (spatial::isFragmentAssembly(blueprint.getMode())) { markOpToRemove(blueprint.getOperation()); for (Value operand : blueprint->getOperands()) markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain); diff --git a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp index 60f1324..b25db93 100644 --- a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp +++ b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp @@ -29,13 +29,13 @@ #include "Common/IR/ConstantUtils.hpp" #include "Common/PimCommon.hpp" #include "Common/Support/CheckedArithmetic.hpp" -#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp" +#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp" #include "Conversion/ONNXToSpatial/Common/Common.hpp" #include "Conversion/SpatialToPim/Common.hpp" #include "Conversion/SpatialToPim/Patterns.hpp" #include "Dialect/Pim/PimOps.hpp" #include "Dialect/Spatial/SpatialOps.hpp" -#include "Pass/PIMPasses.h" +#include "Passes/PIMPasses.h" #include "SpatialToPimPass.hpp" using namespace mlir; @@ -66,17 +66,20 @@ createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTe return padOp.getResult(); } -static FailureOr padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) { +static FailureOr padHVectorInputToCrossbarSize(IRRewriter& rewriter, + Location loc, + Value vector, + int64_t crossbarSize) { auto vectorType = cast(vector.getType()); ArrayRef shape = vectorType.getShape(); assert(isHVectorShape(shape) && "expected a horizontal vector"); - assert(shape[1] <= static_cast(crossbarSize) && "vector width must fit in one crossbar"); + assert(shape[1] <= crossbarSize && "vector width must fit in one crossbar"); - if (shape[1] == static_cast(crossbarSize)) + if (shape[1] == crossbarSize) return vector; auto paddedType = RankedTensorType::get( - {shape[0], static_cast(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding()); + {shape[0], crossbarSize}, vectorType.getElementType(), vectorType.getEncoding()); return createZeroPaddedTensor(rewriter, loc, vector, paddedType); } @@ -84,6 +87,11 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() { outputTensors.clear(); operationsToRemove.clear(); ModuleOp moduleOp = getOperation(); + if (!hasTarget || failed(targetResources.verify())) { + moduleOp.emitError("Spatial-to-PIM lowering requires valid injected target resources"); + signalPassFailure(); + return; + } MLIRContext* ctx = moduleOp.getContext(); auto entryFunc = getPimEntryFunc(moduleOp); @@ -265,15 +273,16 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func: ArrayRef outputShape = outputType.getShape(); assert(isHVectorShape(outputShape) && "expected a horizontal vector output"); auto weightType = cast(vmmOp.getWeight().getType()); - const int64_t xbarDim = static_cast(crossbarSize); + const int64_t xbarDim = static_cast(targetResources.matrixShape.columns); const int64_t paddedOutputWidth = ceilIntegerDivide(outputShape[1], xbarDim) * xbarDim; assert(weightType.getRank() == 2 && weightType.getDimSize(1) == paddedOutputWidth && "expected VMM weight width to match the padded output width"); - assert(paddedOutputWidth / xbarDim <= static_cast(crossbarCountInCore) + assert(paddedOutputWidth / xbarDim <= static_cast(targetResources.matrixUnitsPerProcessor) && "output width must fit in one core"); rewriter.setInsertionPoint(vmmOp); - auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput()); + auto paddedInput = padHVectorInputToCrossbarSize( + rewriter, vmmOp.getLoc(), vmmOp.getInput(), xbarDim); if (failed(paddedInput)) { hasFailure = true; return WalkResult::interrupt(); @@ -375,4 +384,9 @@ void raptor::SpatialToPimPass::eraseOpsToRemove() { std::unique_ptr createSpatialToPimPass() { return std::make_unique(); } +std::unique_ptr createSpatialToPimPass( + const spatial::SpatialTargetResources& target) { + return std::make_unique(target); +} + } // namespace onnx_mlir diff --git a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp index 702bb7c..cfa1118 100644 --- a/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp +++ b/src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp @@ -18,6 +18,7 @@ #include "Conversion/SpatialToPim/Common.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" namespace onnx_mlir { namespace raptor { @@ -28,7 +29,10 @@ struct SpatialToPimPass : mlir::PassWrapper outputTensors; llvm::SmallVector operationsToRemove; + spatial::SpatialTargetResources targetResources; + bool hasTarget = false; mlir::LogicalResult allocateAndInitializeCoreLocalVariables(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter); mlir::LogicalResult diff --git a/src/PIM/Dialect/Pim/CMakeLists.txt b/src/PIM/Dialect/Pim/CMakeLists.txt index eca8dfb..f121461 100644 --- a/src/PIM/Dialect/Pim/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/CMakeLists.txt @@ -1,12 +1,12 @@ add_onnx_mlir_dialect(Pim pim) add_onnx_mlir_dialect_doc(pim Pim.td) -add_subdirectory(Analysis) -add_subdirectory(Transforms/Bufferization) -add_subdirectory(Transforms/HostConstantFolding) -add_subdirectory(Transforms/InstructionSelection) -add_subdirectory(Transforms/LocalMemoryPlanning) -add_subdirectory(Transforms/Verification) +add_subdirectory(Passes/Analyses) +add_subdirectory(Passes/Transforms/Bufferization) +add_subdirectory(Passes/Transforms/HostConstantFolding) +add_subdirectory(Passes/Transforms/InstructionSelection) +add_subdirectory(Passes/Transforms/LocalMemoryPlanning) +add_subdirectory(Passes/Transforms/Verification) add_pim_library(PimOps PimOps.hpp diff --git a/src/PIM/Dialect/Pim/Analysis/CMakeLists.txt b/src/PIM/Dialect/Pim/Passes/Analyses/CMakeLists.txt similarity index 100% rename from src/PIM/Dialect/Pim/Analysis/CMakeLists.txt rename to src/PIM/Dialect/Pim/Passes/Analyses/CMakeLists.txt diff --git a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp b/src/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.cpp similarity index 98% rename from src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp rename to src/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.cpp index 4bd7b94..8285227 100644 --- a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp +++ b/src/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.cpp @@ -8,7 +8,7 @@ #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" using namespace mlir; diff --git a/src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp b/src/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp similarity index 100% rename from src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp rename to src/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.cpp similarity index 91% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.cpp index 6d69664..f01514c 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.cpp @@ -3,8 +3,8 @@ #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/Common.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp" using namespace mlir; using namespace bufferization; diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.hpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.hpp diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/CMakeLists.txt b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/CMakeLists.txt similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/CMakeLists.txt rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/CMakeLists.txt diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/Common.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.cpp similarity index 98% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/Common.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.cpp index 6b752b5..af1e9b5 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/Common.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.cpp @@ -1,4 +1,4 @@ -#include "Dialect/Pim/Transforms/Bufferization/Common.hpp" +#include "Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp" #include "mlir/Dialect/SCF/IR/SCF.h" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/Common.hpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/Common.hpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.cpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.cpp diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/OpBufferizationInterfaces.cpp similarity index 99% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/OpBufferizationInterfaces.cpp index 69db269..1eb6aae 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/OpBufferizationInterfaces.cpp @@ -7,7 +7,7 @@ #include "OpBufferizationInterfaces.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.hpp" using namespace mlir; using namespace bufferization; diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.hpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/OpBufferizationInterfaces.hpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.hpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/OpBufferizationInterfaces.hpp diff --git a/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/PimBufferizationPass.cpp similarity index 79% rename from src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/PimBufferizationPass.cpp index 038d428..7988e3f 100644 --- a/src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/PimBufferizationPass.cpp @@ -20,11 +20,11 @@ #include "Common/Support/Diagnostics.hpp" #include "Compiler/PimCodeGen.hpp" #include "Dialect/Pim/PimOps.hpp" -#include "Dialect/Pim/Transforms/Bufferization/Common.hpp" -#include "Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp" +#include "Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp" +#include "Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp" #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" +#include "src/Accelerators/PIM/Passes/PIMPasses.h" #include "src/Compiler/CompilerOptions.hpp" using namespace mlir; @@ -33,6 +33,9 @@ using namespace pim; namespace onnx_mlir { +static void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp); +static FailureOr requirePimEntryFunc(ModuleOp moduleOp, StringRef phase); + namespace { struct MemRefCopyWorkItem { @@ -333,22 +336,6 @@ static LogicalResult verifyPimCopyEndpoints(Operation* copy, return success(valid); } -struct PimBufferizationPass : PassWrapper> { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass) - StringRef getArgument() const override { return "bufferize-pim"; } - StringRef getDescription() const override { return "Bufferize PIM and Spatial ops."; } - - PimBufferizationPass() = default; - PimBufferizationPass(const PimBufferizationPass& pass) {} - - void runOnOperation() final; - -private: - void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const; - LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) const; - LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) const; -}; - static void materializeWritableConstantDestinations(func::FuncOp funcOp) { SmallVector constantBackedRoots; llvm::SmallPtrSet seenRoots; @@ -387,14 +374,22 @@ static void materializeWritableConstantDestinations(func::FuncOp funcOp) { } } +static bufferization::OneShotBufferizationOptions makePimBufferizationOptions() { + bufferization::OneShotBufferizationOptions options; + options.allowUnknownOps = true; + options.bufferizeFunctionBoundaries = true; + options.setFunctionBoundaryTypeConversion(bufferization::LayoutMapOption::IdentityLayoutMap); + return options; +} + static LogicalResult verifyPimCoresNeedNoTensorCopies( - ModuleOp module, const bufferization::OneShotBufferizationOptions& baseOptions) { + ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions& baseOptions) { static constexpr StringLiteral kExistingAlloc = "raptor.existing_core_alloc"; - OwningOpRef clone = module.clone(); + OwningOpRef clone = moduleOp.clone(); clone->walk([&](bufferization::AllocTensorOp alloc) { if (alloc->getParentOfType() || alloc->getParentOfType()) - alloc->setAttr(kExistingAlloc, UnitAttr::get(module.getContext())); + alloc->setAttr(kExistingAlloc, UnitAttr::get(moduleOp.getContext())); }); auto options = baseOptions; @@ -407,7 +402,7 @@ static LogicalResult verifyPimCoresNeedNoTensorCopies( bufferization::BufferizationState state; if (failed(bufferization::insertTensorCopies(*clone, options, state))) { - module.emitError("official one-shot analysis failed while verifying PIM core copy freedom"); + moduleOp.emitError("official one-shot analysis failed while verifying PIM core copy freedom"); return failure(); } @@ -423,29 +418,22 @@ static LogicalResult verifyPimCoresNeedNoTensorCopies( op->emitOpError("official one-shot bufferization requires a tensor copy inside a PIM core"); }); }); - diagnostics.emitSuppressedSummary(module, "required PIM core tensor copies"); + diagnostics.emitSuppressedSummary(moduleOp, "required PIM core tensor copies"); return success(!diagnostics.hasFailure()); } -} // namespace - -void PimBufferizationPass::runOnOperation() { - auto moduleOp = getOperation(); - auto funcOp = *getPimEntryFunc(moduleOp); - - bufferization::OneShotBufferizationOptions options; - options.allowUnknownOps = true; - options.bufferizeFunctionBoundaries = true; - options.setFunctionBoundaryTypeConversion(bufferization::LayoutMapOption::IdentityLayoutMap); - +static LogicalResult preparePimBufferization( + ModuleOp moduleOp, func::FuncOp funcOp, bool verifyCopyFreedom) { materializeWritableConstantDestinations(funcOp); - if (failed(verifyPimCoresNeedNoTensorCopies(moduleOp, options))) { - signalPassFailure(); - return; - } + if (verifyCopyFreedom) + return verifyPimCoresNeedNoTensorCopies(moduleOp, makePimBufferizationOptions()); + return success(); +} +static LogicalResult runOneShotPimBufferization( + ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions& options) { auto hostOptions = options; - hostOptions.opFilter.denyOperation([](Operation *op) { + hostOptions.opFilter.denyOperation([](Operation* op) { return op->getParentOfType() || op->getParentOfType(); }); @@ -453,84 +441,14 @@ void PimBufferizationPass::runOnOperation() { if (failed(bufferization::insertTensorCopies(moduleOp, hostOptions, state)) || failed(bufferization::bufferizeModuleOp(moduleOp, options, state))) { moduleOp.emitError("Failed to bufferize PIM and Spatial ops"); - signalPassFailure(); - return; + return failure(); } - - forwardSingleConsumerReceiveCopies(funcOp); - forwardSingleConsumerContiguousInputCopies(funcOp); - forwardSingleConsumerPimOutputCopies(funcOp); - - MLIRContext* ctx = moduleOp.getContext(); - PatternRewriter rewriter(ctx); - - SmallVector copyWorklist; - llvm::SmallPtrSet seenCopyOps; - auto addCopyOp = [&](memref::CopyOp copyOp, const StaticValueKnowledge& knowledge) { - if (seenCopyOps.insert(copyOp.getOperation()).second) - copyWorklist.push_back({copyOp, knowledge}); - }; - - moduleOp.walk([&](pim::PimCoreOp coreOp) { - StaticValueKnowledge knowledge = seedCoreKnowledge(coreOp); - (void) walkPimCoreBlockStructurally( - coreOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) { - if (auto copyOp = dyn_cast(&op)) - addCopyOp(copyOp, opKnowledge); - return success(); - }); - }); - moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) { - for (unsigned lane = 0; lane < coreBatchOp.getLaneCount(); ++lane) { - StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane); - (void) walkPimCoreBlockStructurally( - coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) { - if (auto copyOp = dyn_cast(&op)) - addCopyOp(copyOp, opKnowledge); - return success(); - }); - } - }); - - bool hasFailed = false; - Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0); - for (const MemRefCopyWorkItem& workItem : copyWorklist) { - memref::CopyOp copyOp = workItem.copyOp; - rewriter.setInsertionPoint(copyOp); - if (failed(lowerMemRefCopyToPimCopy(copyOp, zeroOffset, rewriter, workItem.knowledge))) - hasFailed = true; - } - if (hasFailed) { - signalPassFailure(); - return; - } - - RewritePatternSet contiguityPatterns(ctx); - populatePimContiguityNormalizationPatterns(contiguityPatterns); - - GreedyRewriteConfig contiguityConfig; - contiguityConfig.enableFolding(false); - if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) { - moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization"); - signalPassFailure(); - return; - } - if (failed(verifyContiguousRuntimeOperands(moduleOp))) { - signalPassFailure(); - return; - } - if (failed(verifyPimCopyAddressSpaces(moduleOp))) { - signalPassFailure(); - return; - } - - annotateWeightsMemrefs(moduleOp, funcOp); - - // Dump to file for debug - dumpModule(moduleOp, "pim1_buff"); + return success(); } -void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const { +} // namespace + +static void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) { auto markWeights = [&](Operation* op) { walkPimMvmVmmWeightUses(op, [&](OpOperand& weightUse) { Value weight = weightUse.get(); @@ -548,7 +466,7 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO funcOp.walk([&](PimCoreBatchOp coreBatchOp) { markWeights(coreBatchOp); }); } -LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const { +static LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) { bool hasFailure = false; auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) { @@ -640,7 +558,7 @@ LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp mod return success(); } -LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp) const { +static LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) { size_t failureCount = 0; auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) { (void) walkPimCoreBlockStructurally( @@ -675,6 +593,211 @@ LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp return success(failureCount == 0); } -std::unique_ptr createPimBufferizationPass() { return std::make_unique(); } +static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp) { + forwardSingleConsumerReceiveCopies(funcOp); + forwardSingleConsumerContiguousInputCopies(funcOp); + forwardSingleConsumerPimOutputCopies(funcOp); + + MLIRContext* ctx = moduleOp.getContext(); + PatternRewriter rewriter(ctx); + + SmallVector copyWorklist; + llvm::SmallPtrSet seenCopyOps; + auto addCopyOp = [&](memref::CopyOp copyOp, const StaticValueKnowledge& knowledge) { + if (seenCopyOps.insert(copyOp.getOperation()).second) + copyWorklist.push_back({copyOp, knowledge}); + }; + + moduleOp.walk([&](pim::PimCoreOp coreOp) { + StaticValueKnowledge knowledge = seedCoreKnowledge(coreOp); + (void) walkPimCoreBlockStructurally( + coreOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) { + if (auto copyOp = dyn_cast(&op)) + addCopyOp(copyOp, opKnowledge); + return success(); + }); + }); + moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) { + for (unsigned lane = 0; lane < coreBatchOp.getLaneCount(); ++lane) { + StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane); + (void) walkPimCoreBlockStructurally( + coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) { + if (auto copyOp = dyn_cast(&op)) + addCopyOp(copyOp, opKnowledge); + return success(); + }); + } + }); + + bool hasFailed = false; + Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0); + for (const MemRefCopyWorkItem& workItem : copyWorklist) { + memref::CopyOp copyOp = workItem.copyOp; + rewriter.setInsertionPoint(copyOp); + if (failed(lowerMemRefCopyToPimCopy(copyOp, zeroOffset, rewriter, workItem.knowledge))) + hasFailed = true; + } + if (hasFailed) + return failure(); + + RewritePatternSet contiguityPatterns(ctx); + populatePimContiguityNormalizationPatterns(contiguityPatterns); + + GreedyRewriteConfig contiguityConfig; + contiguityConfig.enableFolding(false); + if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) { + moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization"); + return failure(); + } + annotateWeightsMemrefs(moduleOp, funcOp); + dumpModule(moduleOp, "pim1_buff"); + return success(); +} + +static FailureOr requirePimEntryFunc(ModuleOp moduleOp, StringRef phase) { + auto entryFunc = getPimEntryFunc(moduleOp); + if (failed(entryFunc)) { + moduleOp.emitError("failed to locate the PIM entry function during ") << phase; + return failure(); + } + return *entryFunc; +} + +namespace { + +struct PimBufferizationPreparationPass + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPreparationPass) + + explicit PimBufferizationPreparationPass(bool verifyCopyFreedom = false) + : verifyCopyFreedom(verifyCopyFreedom) {} + + StringRef getArgument() const override { return "pim-bufferization-preparation"; } + StringRef getDescription() const override { + return "Prepare writable tensor destinations for PIM one-shot bufferization."; + } + + void runOnOperation() final { + ModuleOp moduleOp = getOperation(); + auto funcOp = requirePimEntryFunc(moduleOp, "PIM bufferization preparation"); + if (failed(funcOp)) { + signalPassFailure(); + return; + } + if (failed(preparePimBufferization(moduleOp, *funcOp, verifyCopyFreedom))) + signalPassFailure(); + } + +private: + bool verifyCopyFreedom; +}; + +struct PimOneShotBufferizationPass + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimOneShotBufferizationPass) + + StringRef getArgument() const override { return "pim-one-shot-bufferization"; } + StringRef getDescription() const override { + return "Run one-shot bufferization for PIM and Spatial tensors."; + } + + void runOnOperation() final { + if (failed(runOneShotPimBufferization(getOperation(), makePimBufferizationOptions()))) + signalPassFailure(); + } +}; + +struct PimMemoryNormalizationPass + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimMemoryNormalizationPass) + + StringRef getArgument() const override { return "pim-memory-normalization"; } + StringRef getDescription() const override { + return "Normalize PIM memory copies and verify addressable operands."; + } + + void runOnOperation() final { + ModuleOp moduleOp = getOperation(); + auto funcOp = requirePimEntryFunc(moduleOp, "PIM memory normalization"); + if (failed(funcOp)) { + signalPassFailure(); + return; + } + if (failed(normalizePimMemory(moduleOp, *funcOp))) + signalPassFailure(); + } +}; + +static LogicalResult verifyNoTensorValues(ModuleOp moduleOp) { + size_t failureCount = 0; + moduleOp.walk([&](Operation* op) { + if (failureCount >= 8) + return; + if (op->getDialect()->getNamespace() == "tensor") { + op->emitOpError("tensor operation remains after PIM bufferization"); + ++failureCount; + return; + } + for (Value value : op->getOperands()) { + if (isa(value.getType())) { + op->emitOpError("tensor operand remains after PIM bufferization"); + ++failureCount; + return; + } + } + for (Value value : op->getResults()) { + if (isa(value.getType())) { + op->emitOpError("tensor result remains after PIM bufferization"); + ++failureCount; + return; + } + } + }); + if (failureCount != 0) + moduleOp.emitError() << "found " << failureCount + << " tensor value(s) after PIM bufferization" + << (failureCount == 8 ? " (first 8 reported)" : ""); + return success(failureCount == 0); +} + +struct PimBufferizationVerificationPass + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationVerificationPass) + + StringRef getArgument() const override { return "pim-bufferization-verification"; } + StringRef getDescription() const override { + return "Verify tensor elimination, contiguity, and PIM copy address spaces."; + } + + void runOnOperation() final { + ModuleOp moduleOp = getOperation(); + if (failed(verifyNoTensorValues(moduleOp)) + || failed(verifyContiguousRuntimeOperands(moduleOp)) + || failed(verifyPimCopyAddressSpaces(moduleOp))) + signalPassFailure(); + } +}; + +} // namespace + +std::unique_ptr createPimBufferizationPreparationPass() { + return std::make_unique(); +} + +std::unique_ptr createPimBufferizationPreparationPass(bool verifyCopyFreedom) { + return std::make_unique(verifyCopyFreedom); +} + +std::unique_ptr createPimOneShotBufferizationPass() { + return std::make_unique(); +} + +std::unique_ptr createPimMemoryNormalizationPass() { + return std::make_unique(); +} + +std::unique_ptr createPimBufferizationVerificationPass() { + return std::make_unique(); +} } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/CMakeLists.txt b/src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/CMakeLists.txt similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/HostConstantFolding/CMakeLists.txt rename to src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/CMakeLists.txt diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Common.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Common.cpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Common.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Common.cpp diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Common.hpp b/src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Common.hpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Common.hpp rename to src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Common.hpp diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp similarity index 95% rename from src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp index a7d4cc9..ef553ff 100644 --- a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp @@ -5,7 +5,7 @@ #include "Patterns.hpp" #include "src/Accelerators/PIM/Common/PimCommon.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp" using namespace mlir; diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns.hpp b/src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Patterns.hpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns.hpp rename to src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Patterns.hpp diff --git a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Patterns/Constant.cpp similarity index 90% rename from src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Patterns/Constant.cpp index ad557d3..cb4e22b 100644 --- a/src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp +++ b/src/PIM/Dialect/Pim/Passes/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/Pim/Transforms/HostConstantFolding/Patterns/Subview.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Patterns/Subview.cpp similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Subview.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Patterns/Subview.cpp diff --git a/src/PIM/Dialect/Pim/Transforms/InstructionSelection/CMakeLists.txt b/src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/CMakeLists.txt similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/InstructionSelection/CMakeLists.txt rename to src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/CMakeLists.txt diff --git a/src/PIM/Dialect/Pim/Transforms/InstructionSelection/InstructionSelectionPass.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp similarity index 98% rename from src/PIM/Dialect/Pim/Transforms/InstructionSelection/InstructionSelectionPass.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp index cf5fbbb..206fe81 100644 --- a/src/PIM/Dialect/Pim/Transforms/InstructionSelection/InstructionSelectionPass.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp @@ -8,8 +8,8 @@ #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp" #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp" +#include "src/Accelerators/PIM/Passes/PIMPasses.h" using namespace llvm; using namespace mlir; diff --git a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt b/src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/CMakeLists.txt similarity index 100% rename from src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/CMakeLists.txt rename to src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/CMakeLists.txt diff --git a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp similarity index 97% rename from src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp index 4a85565..02e1bfb 100644 --- a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp @@ -8,8 +8,8 @@ #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp" -#include "src/Accelerators/PIM/Pass/PIMPasses.h" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp" +#include "src/Accelerators/PIM/Passes/PIMPasses.h" using namespace llvm; using namespace mlir; diff --git a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp b/src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp similarity index 87% rename from src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp rename to src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp index 76d3e13..f5987ea 100644 --- a/src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp @@ -1,6 +1,6 @@ #pragma once -#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp" namespace onnx_mlir { diff --git a/src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt b/src/PIM/Dialect/Pim/Passes/Transforms/Verification/CMakeLists.txt similarity index 89% rename from src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt rename to src/PIM/Dialect/Pim/Passes/Transforms/Verification/CMakeLists.txt index de55435..6c851cd 100644 --- a/src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Verification/CMakeLists.txt @@ -5,7 +5,6 @@ add_pim_library(OMPimVerification LINK_LIBS PUBLIC OMPimCommon - OMPimCompilerOptions OMPimBufferization OMPimLocalMemoryLifetimeAnalysis PimOps diff --git a/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp similarity index 94% rename from src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp rename to src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp index d63ae7a..d6285be 100644 --- a/src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Verification/VerificationPass.cpp @@ -18,11 +18,11 @@ #include "src/Accelerators/PIM/Common/PimCommon.hpp" #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp" #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" -#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp" +#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp" #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp" using namespace mlir; @@ -748,12 +748,43 @@ struct VerificationPass : PassWrapper> } VerificationPass() {} - VerificationPass(const VerificationPass& pass) {} + VerificationPass(const spatial::SpatialTargetResources& target, + bool detectCommunicationDeadlock) + : targetResources(target), hasTarget(true), + detectCommunicationDeadlock(detectCommunicationDeadlock) {} + VerificationPass(const VerificationPass& pass) + : targetResources(pass.targetResources), hasTarget(pass.hasTarget), + detectCommunicationDeadlock(pass.detectCommunicationDeadlock) {} void runOnOperation() override { ModuleOp moduleOp = getOperation(); pim::CappedDiagnosticReporter diagnostics; + if (!hasTarget || failed(targetResources.verify())) { + moduleOp.emitError("PIM codegen verification requires valid injected target resources"); + signalPassFailure(); + return; + } + + const int64_t xbarDim = static_cast(targetResources.matrixShape.columns); + moduleOp.walk([&](pim::PimVMMOp vmmOp) { + auto weightType = dyn_cast(vmmOp.getWeight().getType()); + auto inputType = dyn_cast(vmmOp.getInput().getType()); + if (!weightType || !inputType || weightType.getRank() != 2 || inputType.getRank() != 2) + return; + int64_t rows = weightType.getDimSize(0); + int64_t columns = weightType.getDimSize(1); + if (rows > xbarDim + || columns > xbarDim * static_cast(targetResources.matrixUnitsPerProcessor) + || columns % xbarDim != 0 + || inputType.getDimSize(1) != xbarDim) { + diagnostics.report(vmmOp.getOperation(), [xbarDim](Operation* op) { + op->emitOpError() << "VMM dimensions do not fit the injected target (crossbar size " + << xbarDim << ")"; + }); + } + }); + moduleOp.walk([&](Operation* op) { if (op->getDialect()->getNamespace() != "spat") return; @@ -811,7 +842,7 @@ struct VerificationPass : PassWrapper> } bool hasFailure = false; - if (pimDetectCommunicationDeadlock && failed(verifyNoStaticCommunicationDeadlock(moduleOp, diagnostics))) + if (detectCommunicationDeadlock && failed(verifyNoStaticCommunicationDeadlock(moduleOp, diagnostics))) hasFailure = true; if (diagnostics.hasFailure()) { @@ -825,6 +856,10 @@ struct VerificationPass : PassWrapper> } private: + spatial::SpatialTargetResources targetResources; + bool hasTarget = false; + bool detectCommunicationDeadlock = false; + template static LogicalResult verifyCoreWeights(ModuleOp moduleOp, CoreOpTy coreOp, pim::CappedDiagnosticReporter& diagnostics) { @@ -1050,4 +1085,10 @@ private: std::unique_ptr createPimVerificationPass() { return std::make_unique(); } +std::unique_ptr createPimVerificationPass( + const spatial::SpatialTargetResources& target, + bool detectCommunicationDeadlock) { + return std::make_unique(target, detectCommunicationDeadlock); +} + } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Pim/PimOpsVerify.cpp b/src/PIM/Dialect/Pim/PimOpsVerify.cpp index a16942c..04f4d1e 100644 --- a/src/PIM/Dialect/Pim/PimOpsVerify.cpp +++ b/src/PIM/Dialect/Pim/PimOpsVerify.cpp @@ -10,7 +10,6 @@ #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp" #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp" -#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp" using namespace mlir; @@ -157,16 +156,10 @@ LogicalResult PimVMMOp::verify() { int64_t M = matrixShape[1]; if (N <= 0 || M <= 0) return emitError("matrix shape must be (N, M) with N > 0 and M > 0"); - const int64_t xbarDim = static_cast(crossbarSize); - if (N > xbarDim || M > xbarDim * static_cast(crossbarCountInCore)) - return emitError("matrix dimensions must fit in one array group"); - if (M % xbarDim != 0) - return emitError("matrix output width must be padded to a whole number of crossbars"); - int64_t vector1 = vectorShape[0]; int64_t vectorWidth = vectorShape[1]; - if (vector1 != 1 || vectorWidth != xbarDim) - return emitError("vector shape must be (1, crossbar-size)"); + if (vector1 != 1 || vectorWidth <= 0) + return emitError("vector shape must be (1, positive-width)"); int64_t output1 = outputShape[0]; int64_t outputWidth = outputShape[1]; diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index 983cfa9..fa34968 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -1,38 +1,49 @@ add_onnx_mlir_dialect(Spatial spat) add_onnx_mlir_dialect_doc(spat Spatial.td) +set(LLVM_TARGET_DEFINITIONS Spatial.td) +mlir_tablegen(SpatialEnums.hpp.inc -gen-enum-decls "-I${ONNX_MLIR_SRC_ROOT}") +mlir_tablegen(SpatialEnums.cpp.inc -gen-enum-defs "-I${ONNX_MLIR_SRC_ROOT}") +add_public_tablegen_target(OMSpatialEnumsIncGen) + +mlir_tablegen(SpatialLayoutInterface.hpp.inc -gen-op-interface-decls "-I${ONNX_MLIR_SRC_ROOT}") +mlir_tablegen(SpatialLayoutInterface.cpp.inc -gen-op-interface-defs "-I${ONNX_MLIR_SRC_ROOT}") +add_public_tablegen_target(OMSpatialLayoutInterfaceIncGen) + add_pim_library(SpatialOps SpatialOps.cpp SpatialOpsAsm.cpp SpatialOpsVerify.cpp SpatialOpsCanonicalization.cpp ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp - Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp - Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp - Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp - Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp - Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp - Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp - Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp - Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp - Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp - Transforms/MergeComputeNodes/DeferredResultRealization.cpp - Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp - Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp - Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp - Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp - Transforms/MergeComputeNodes/ScheduledComputeReport.cpp - Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp - Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp - Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp - Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp - Transforms/TrivialGraphComputeMergePass.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp + Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp + Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp + Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp + Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp + Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp + Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp + Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp + Passes/Transforms/MergeComputeNodes/DeferredResultRealization.cpp + Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp + Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp + Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp + Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp + Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp + Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp + Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp + Passes/Transforms/TrivialGraphComputeMergePass.cpp EXCLUDE_FROM_OM_LIBS DEPENDS OMONNXIncGen OMSpatialIncGen + OMSpatialEnumsIncGen + OMSpatialLayoutInterfaceIncGen LINK_LIBS PUBLIC MLIRIR @@ -40,6 +51,5 @@ add_pim_library(SpatialOps MLIRBufferizationTransforms OMMlirDialects OMONNXOps - OMPimCompilerOptions PimOps ) diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp similarity index 85% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp index e6846b2..712110c 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/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; @@ -328,7 +387,7 @@ static void collectClosure(Value value, Block &body, const DeferredInputPlan &pl bool isDeferredFragmentAssemblyInput(Value input, size_t processorCount) { auto blueprint = input.getDefiningOp(); - if (!blueprint || blueprint.getMode() != "fragment_assembly") + if (!blueprint || !isFragmentAssembly(blueprint.getMode())) return false; return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) { return getProducerValueRef(fragment, nullptr, processorCount).has_value(); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp similarity index 91% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp index 8af6bf9..cdf9f11 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp +++ b/src/PIM/Dialect/Spatial/Passes/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/DeferredCommunicationRealization.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp similarity index 97% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp index 42ae3a1..b816e3d 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp +++ b/src/PIM/Dialect/Spatial/Passes/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 || definition->getBlock() != &body) + 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/DeferredProjectionAnalysis.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredResultRealization.cpp similarity index 99% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredResultRealization.cpp index 6cdb525..dbd01a1 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp +++ b/src/PIM/Dialect/Spatial/Passes/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/DeferredResultRealization.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredResultRealization.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredResultRealization.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp similarity index 98% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp index f6e554b..9cd9444 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp @@ -395,7 +395,7 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p static LogicalResult retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBatchPublicationCache& publicationCache) { - if (blueprint.getMode() != "fragment_assembly") + if (!isFragmentAssembly(blueprint.getMode())) return success(); bool escapesScheduledGraph = llvm::any_of( blueprint.getOutput().getUses(), [](OpOperand &use) { @@ -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/DeferredTransferPlanning.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp similarity index 99% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp index c98f3ae..8cdf867 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp @@ -24,7 +24,7 @@ bool requiresScheduledPublication(Value value, DenseSet &visited) { SpatDeferredCommunicationOp>(user)) return false; auto blueprint = dyn_cast(user); - return !blueprint || blueprint.getMode() != "fragment_assembly" + return !blueprint || !isFragmentAssembly(blueprint.getMode()) || requiresScheduledPublication(blueprint.getOutput(), visited); }); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp new file mode 100644 index 0000000..578bc25 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp @@ -0,0 +1,144 @@ +#include "mlir/Pass/Pass.h" + +#include "DeferredCommunicationRealization.hpp" +#include "ScheduledComputeReport.hpp" +#include "ScheduledComputeVerification.hpp" +#include "SpatialDataflowCsvExporter.hpp" +#include "src/Accelerators/PIM/Common/PimCommon.hpp" +#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp" +#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp" +#include "src/Accelerators/PIM/Passes/PIMPasses.h" + +using namespace mlir; + +namespace onnx_mlir { +namespace spatial { +namespace { + +static bool hasValidTarget(const SchedulingTarget& target) { + return target.processorCount != 0 && target.residentWeightCapacity != 0 + && target.transferWidthBytes != 0 + && target.interProcessorLatencyNs.size() == target.processorCount * target.processorCount + && (target.processorCount == 1 || target.averageInterProcessorLatencyNs != 0); +} + +static FailureOr requireEntry(ModuleOp moduleOp) { + auto entry = getPimEntryFunc(moduleOp); + if (failed(entry)) { + moduleOp.emitError("failed to locate the PIM entry function during Spatial scheduling and realization"); + return failure(); + } + return *entry; +} + +struct ScheduleAndRealizeSpatialPass final + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleAndRealizeSpatialPass) + + ScheduleAndRealizeSpatialPass() = default; + ScheduleAndRealizeSpatialPass(const SchedulingTarget& target, + SpatialDataflowExportStage exportStage) + : target(target), exportStage(exportStage), hasTarget(true) {} + + StringRef getArgument() const override { return "schedule-and-realize-spatial"; } + StringRef getDescription() const override { + return "Schedule Spatial graph computes, verify them, and realize deferred communication."; + } + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + if (!hasTarget || !hasValidTarget(target)) { + moduleOp.emitError("ScheduleAndRealizeSpatial requires an explicit valid target"); + signalPassFailure(); + return; + } + auto entry = requireEntry(moduleOp); + if (failed(entry)) { + signalPassFailure(); + return; + } + func::FuncOp entryFunc = *entry; + + MergeSchedulingAnalysis analysis(entryFunc, target); + MergeScheduleResult schedule = std::move(analysis.getResult()); + PatternRewriter rewriter(moduleOp.getContext()); + FailureOr materialization = + materializeScheduledCompute(entryFunc, schedule, rewriter); + if (failed(materialization)) { + signalPassFailure(); + return; + } + ScheduledComputeMaterializationResult materializationResult = std::move(*materialization); + + if (failed(verifyMaterializedScheduleMapping( + entryFunc, schedule, materializationResult.peftClassPlans, + materializationResult.graphComputeToBlockMap, + materializationResult.materializedSchedules)) + || failed(verifyDeferredTransferPhase1Invariants(entryFunc)) + || failed(verifyScheduledMaterializationRecords(materializationResult.materializedSchedules)) + || failed(verifyScheduledSpatialInvariants(entryFunc))) { + moduleOp.emitError("scheduled Spatial phase verification failed"); + signalPassFailure(); + return; + } + + dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true); + if (shouldExportSpatialDataflowStage(exportStage, SpatialDataflowExportStage::Spatial3) + && failed(exportSpatialDataflowCsvScheduled( + entryFunc, materializationResult.materializedSchedules, + "spatial3_scheduled_no_comm", "spatial3"))) { + signalPassFailure(); + return; + } + dumpScheduledComputeReport( + moduleOp, entryFunc, schedule, materializationResult.peftClassPlans, + materializationResult.materializedSchedules); + + if (failed(realizeDeferredCommunication(entryFunc, materializationResult, target))) { + moduleOp.emitError("Spatial communication realization failed"); + signalPassFailure(); + return; + } + + bool deferredRemains = false; + entryFunc.walk([&](SpatDeferredCommunicationOp deferred) { + if (deferredRemains) + return; + deferred.emitOpError("realized Spatial graph still contains deferred communication"); + deferredRemains = true; + }); + if (deferredRemains + || failed(verifyScheduledResultsLive(materializationResult.materializedSchedules)) + || failed(verifyScheduledSpatialInvariants(entryFunc))) { + moduleOp.emitError("realized Spatial communication verification failed"); + signalPassFailure(); + return; + } + + dumpModule(moduleOp, "spatial4_scheduled", /*assumeVerified=*/true); + if (shouldExportSpatialDataflowStage(exportStage, SpatialDataflowExportStage::Spatial4) + && failed(exportSpatialDataflowCsvScheduled( + entryFunc, materializationResult.materializedSchedules, + "spatial4_scheduled", "spatial4"))) + signalPassFailure(); + } + +private: + SchedulingTarget target; + SpatialDataflowExportStage exportStage = SpatialDataflowExportStage::None; + bool hasTarget = false; +}; + +} // namespace + +std::unique_ptr createScheduleAndRealizeSpatialPass() { + return std::make_unique(); +} + +std::unique_ptr createScheduleAndRealizeSpatialPass( + const SchedulingTarget& target, SpatialDataflowExportStage exportStage) { + return std::make_unique(target, exportStage); +} + +} // namespace spatial +} // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstance.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstance.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstance.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstance.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp similarity index 100% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp similarity index 71% rename from src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp rename to src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp index 66d35ea..6ba90ca 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp @@ -12,7 +12,7 @@ #include #include "PeftScheduler.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp" namespace onnx_mlir { namespace spatial { @@ -23,6 +23,7 @@ namespace { // Pressure means distinct weights exceed half the fleet's one-copy capacity. constexpr size_t kHighResidentWeightPressureCapacityDivisor = 2; +constexpr Time kMakespanMaterialityDivisor = 100; struct ScheduledTask { size_t processor = std::numeric_limits::max(); @@ -181,7 +182,7 @@ FailureOr buildLanePublicationSignatures(SpatComputeB for (auto [useIndex, use] : llvm::enumerate(result.getUses())) { auto blueprint = dyn_cast(use.getOwner()); - if (!blueprint || blueprint.getMode() != "fragment_assembly") + if (!blueprint || !isFragmentAssembly(blueprint.getMode())) continue; auto operandIndices = blueprint.getFragmentOperandIndices(); auto sourceSlots = blueprint.getFragmentSourceSlots(); @@ -388,166 +389,201 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu return lhs > rhs; }; - std::vector remainingParents(nodeCount, 0); - std::priority_queue, decltype(readyCompare)> readyQueue(readyCompare); - for (size_t node = 0; node < nodeCount; ++node) { - remainingParents[node] = graph.predecessors[node].size(); - if (remainingParents[node] == 0) - readyQueue.push(node); - } - - std::vector scheduled(nodeCount, false); - std::vector processorResidentWeights(processorCount); - std::vector schedules(nodeCount); - std::vector> tasksByProcessor(processorCount); - - size_t scheduledCount = 0; - while (!readyQueue.empty()) { - size_t task = readyQueue.top(); - readyQueue.pop(); - if (scheduled[task]) - continue; - - size_t bestProcessor = std::numeric_limits::max(); - Time bestEst = 0; - Time bestEft = 0; - Time bestOeft = std::numeric_limits