multiple-output spat computes

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp

@@ -289,7 +289,8 @@ static SmallVector<Value> createIm2colRowComputes(Value x,
     rowResults.reserve(packedNumRows);
     for (int64_t rowIdx = 0; rowIdx < packedNumRows; rowIdx++) {
       SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(rowIdx), rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(packFactor * patchSize)};
+      SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
+                                         rewriter.getIndexAttr(packFactor * patchSize)};
       SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
       rowResults.push_back(
         tensor::ExtractSliceOp::create(rewriter, loc, gemmInputRowType, gemmInputRows, offsets, sizes, strides));
@@ -325,9 +326,10 @@ static Value createCollectedConvOutput(ValueRange gemmRows,
     else {
       auto expandedType = RankedTensorType::get({packedNumRows, packFactor, numChannelsOut}, outType.getElementType());
       auto paddedType = RankedTensorType::get({paddedNumPatches, numChannelsOut}, outType.getElementType());
-      Value packedOutput = gemmRowArgs.size() == 1
-                           ? gemmRowArgs.front()
-                           : tensor::ConcatOp::create(rewriter, loc, /*axis=*/0, gemmRowArgs).getResult();
+      Value packedOutput =
+        gemmRowArgs.size() == 1
+          ? gemmRowArgs.front()
+          : tensor::ConcatOp::create(rewriter, loc, /*axis=*/0, gemmRowArgs).getResult();
       Value expandedOutput = tensor::ExpandShapeOp::create(rewriter,
                                                            loc,
                                                            expandedType,
@@ -503,41 +505,38 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
   //   B (weights): [patchSize, cOut]        -- W^T, stored in crossbar columns
   // and optionally repack several old rows into one GEMM row to use the available crossbar size better.
   //
-  // We want to process N pixels at the same time. Instead of doing N separate operations
-  // of (1 x patchSize) x (patchSize x cOut), we construct a block-diagonal weight matrix
-  // containing N copies of W^T and concatenate N im2col rows into one longer row:
-  //   A_packed: [ceil(numPatches / N), N * patchSize]
-  //   B_packed: [N * patchSize, N * cOut]
-  //   Y_packed: [ceil(numPatches / N), N * cOut]
-  auto gemmInputRowType = RankedTensorType::get({1, effectiveMaxParallelPixels * patchSize}, elemType);
+  // The im2col compute yields each GEMM input row as a separate result so every GEMM consumes only
+  // the row it needs instead of receiving a full packed tensor and slicing it locally.
+  auto gemmInputRowType =
+    RankedTensorType::get({1, effectiveMaxParallelPixels * patchSize}, elemType);
   auto gemmOutputRowType =
     RankedTensorType::get({1, effectiveMaxParallelPixels * numChannelsOut}, outType.getElementType());
   SmallVector<Value> gemmInputRows = createIm2colRowComputes(x,
-                                                             xType,
-                                                             im2colType,
-                                                             rowType,
-                                                             gemmInputRowType,
-                                                             batchSize,
-                                                             numChannelsIn,
-                                                             xHeight,
-                                                             xWidth,
-                                                             wHeight,
-                                                             wWidth,
-                                                             padHeightBegin,
-                                                             padHeightEnd,
-                                                             padWidthBegin,
-                                                             padWidthEnd,
-                                                             strideHeight,
-                                                             strideWidth,
-                                                             dilationHeight,
-                                                             dilationWidth,
-                                                             outWidth,
-                                                             patchSize,
-                                                             numPatches,
-                                                             numPatchesPerBatch,
-                                                             effectiveMaxParallelPixels,
-                                                             rewriter,
-                                                             loc);
+                                                            xType,
+                                                            im2colType,
+                                                            rowType,
+                                                            gemmInputRowType,
+                                                            batchSize,
+                                                            numChannelsIn,
+                                                            xHeight,
+                                                            xWidth,
+                                                            wHeight,
+                                                            wWidth,
+                                                            padHeightBegin,
+                                                            padHeightEnd,
+                                                            padWidthBegin,
+                                                            padWidthEnd,
+                                                            strideHeight,
+                                                            strideWidth,
+                                                            dilationHeight,
+                                                            dilationWidth,
+                                                            outWidth,
+                                                            patchSize,
+                                                            numPatches,
+                                                            numPatchesPerBatch,
+                                                            effectiveMaxParallelPixels,
+                                                            rewriter,
+                                                            loc);
 
   Value gemmB = buildPackedWeight(wDenseAttr,
                                   wTrans,

src/PIM/Dialect/Spatial/SpatialOps.cpp

@@ -1,4 +1,5 @@
 #include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -13,7 +14,10 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
 
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/LogicalResult.h"
 
@@ -115,10 +119,13 @@ inline LogicalResult mvmOpVerifySize4(SpatWeightedMVMOp* emitter,
 }
 
 llvm::FailureOr<ArrayRef<int64_t>> getWeightShapeForWeightedOp(Operation* weigthedOp, size_t weightIndex) {
-  if (auto computeOp = dyn_cast<SpatCompute>(weigthedOp->getParentOp()))
-    return cast<ShapedType>(computeOp.getWeights()[weightIndex].getType()).getShape();
+  auto wcomputeOp = dyn_cast<SpatCompute>(weigthedOp->getParentOp());
+  if (wcomputeOp)
+    return cast<ShapedType>(wcomputeOp.getWeights()[weightIndex].getType()).getShape();
 
-  if (auto coreOp = dyn_cast<pim::PimCoreOp>(weigthedOp->getParentOp()))
+  auto coreOp = dyn_cast<pim::PimCoreOp>(weigthedOp->getParentOp());
+
+  if (coreOp)
     return cast<ShapedType>(coreOp.getWeights()[weightIndex].getType()).getShape();
 
   return failure();

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/DCPAnalysis.cpp

@@ -28,7 +28,7 @@ using namespace mlir;
 namespace {
 
 struct VirtualNode {
-  SmallVector<size_t, 4> originalComputeIndices;
+  llvm::SmallVector<size_t, 4> originalComputeIndices;
   Weight weight = 0;
   CrossbarUsage crossbarUsage = 0;
 };
@@ -50,7 +50,7 @@ struct WindowScheduleResult {
   bool usedAllAvailableCpus = false;
 };
 
-std::vector<IndexedEdge> aggregateEdges(ArrayRef<IndexedEdge> edges) {
+std::vector<IndexedEdge> aggregateEdges(llvm::ArrayRef<IndexedEdge> edges) {
   std::map<std::pair<size_t, size_t>, Weight> edgeWeights;
   for (auto [start, end, weight] : edges) {
     size_t startIndex = static_cast<size_t>(start);
@@ -74,7 +74,8 @@ std::vector<IndexedEdge> aggregateEdges(ArrayRef<IndexedEdge> edges) {
   return aggregatedEdges;
 }
 
-VirtualGraph buildInitialVirtualGraph(ArrayRef<SpatCompute> spatComputes, ArrayRef<IndexedEdge> edges) {
+VirtualGraph buildInitialVirtualGraph(llvm::ArrayRef<SpatCompute> spatComputes,
+                                      llvm::ArrayRef<IndexedEdge> edges) {
   VirtualGraph graph;
   graph.nodes.reserve(spatComputes.size());
   for (auto [index, spatCompute] : llvm::enumerate(spatComputes)) {
@@ -173,7 +174,7 @@ std::vector<size_t> selectCriticalWindow(const TimingInfo& timing, size_t window
   return selected;
 }
 
-std::vector<size_t> getOriginalSignature(const VirtualGraph& graph, ArrayRef<size_t> selectedNodes) {
+std::vector<size_t> getOriginalSignature(const VirtualGraph& graph, llvm::ArrayRef<size_t> selectedNodes) {
   std::vector<size_t> signature;
   for (size_t nodeIndex : selectedNodes) {
     const VirtualNode& node = graph.nodes[nodeIndex];
@@ -196,7 +197,8 @@ std::vector<IndexedEdge> buildWindowEdges(const VirtualGraph& graph, const std::
   return aggregateEdges(windowEdges);
 }
 
-WindowScheduleResult scheduleWindow(const VirtualGraph& graph, ArrayRef<size_t> selectedNodes, MLIRContext* context) {
+WindowScheduleResult
+scheduleWindow(const VirtualGraph& graph, llvm::ArrayRef<size_t> selectedNodes, MLIRContext* context) {
   std::vector<Weight> windowWeights;
   std::vector<CrossbarUsage> windowCrossbarUsage;
   std::vector<int64_t> nodeToWindowIndex(graph.nodes.size(), -1);
@@ -232,7 +234,9 @@ WindowScheduleResult scheduleWindow(const VirtualGraph& graph, ArrayRef<size_t>
   return result;
 }
 
-bool coarsenGraph(const VirtualGraph& graph, ArrayRef<std::vector<size_t>> mergeGroups, VirtualGraph& coarsenedGraph) {
+bool coarsenGraph(const VirtualGraph& graph,
+                  llvm::ArrayRef<std::vector<size_t>> mergeGroups,
+                  VirtualGraph& coarsenedGraph) {
   std::vector<int64_t> nodeToMergeGroup(graph.nodes.size(), -1);
   for (auto [groupIndex, mergeGroup] : llvm::enumerate(mergeGroups)) {
     if (mergeGroup.size() < 2)
@@ -299,7 +303,7 @@ bool coarsenGraph(const VirtualGraph& graph, ArrayRef<std::vector<size_t>> merge
 }
 
 bool coarsenGraphWithFallback(const VirtualGraph& graph,
-                              ArrayRef<std::vector<size_t>> mergeGroups,
+                              llvm::ArrayRef<std::vector<size_t>> mergeGroups,
                               VirtualGraph& coarsenedGraph) {
   if (coarsenGraph(graph, mergeGroups, coarsenedGraph))
     return true;
@@ -326,7 +330,7 @@ bool coarsenGraphWithFallback(const VirtualGraph& graph,
   return !acceptedMergeGroups.empty();
 }
 
-std::vector<size_t> computeOriginalTopologicalOrder(size_t computeCount, ArrayRef<IndexedEdge> edges) {
+std::vector<size_t> computeOriginalTopologicalOrder(size_t computeCount, llvm::ArrayRef<IndexedEdge> edges) {
   VirtualGraph graph;
   graph.nodes.resize(computeCount);
   graph.edges = aggregateEdges(edges);
@@ -340,8 +344,8 @@ std::vector<size_t> computeOriginalTopologicalOrder(size_t computeCount, ArrayRe
 }
 
 DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph,
-                                              ArrayRef<SpatCompute> spatComputes,
-                                              ArrayRef<IndexedEdge> originalEdges) {
+                                              llvm::ArrayRef<SpatCompute> spatComputes,
+                                              llvm::ArrayRef<IndexedEdge> originalEdges) {
   DCPAnalysisResult result;
   std::vector<size_t> originalToVirtualNode(spatComputes.size(), 0);
   for (auto [virtualNodeIndex, virtualNode] : llvm::enumerate(graph.nodes))
@@ -363,7 +367,9 @@ DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph,
   return result;
 }
 
-DCPAnalysisResult runLegacyDcp(ArrayRef<SpatCompute> spatComputes, ArrayRef<IndexedEdge> edges, MLIRContext* context) {
+DCPAnalysisResult runLegacyDcp(llvm::ArrayRef<SpatCompute> spatComputes,
+                               llvm::ArrayRef<IndexedEdge> edges,
+                               MLIRContext* context) {
   GraphDCP graphDCP(spatComputes, edges);
   if (coresCount.getValue() > 0)
     graphDCP.setMaxCpuCount(static_cast<int>(coresCount.getValue()));
@@ -377,12 +383,12 @@ DCPAnalysisResult runLegacyDcp(ArrayRef<SpatCompute> spatComputes, ArrayRef<Inde
 SpatCompute getOriginalSpatCompute(Operation* op) {
   if (!op)
     return {};
-  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
+  while (auto extract = llvm::dyn_cast<tensor::ExtractSliceOp>(op)) {
     op = extract.getSource().getDefiningOp();
     if (!op)
       return {};
   }
-  if (auto res = dyn_cast<SpatCompute>(op))
+  if (auto res = llvm::dyn_cast<SpatCompute>(op))
     return res;
   return {};
 }