compact syntax for spatial tensor ops

better IR compaction after dcp merge
remove pim.mvm op
better memory report

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp

@@ -1436,6 +1436,21 @@ public:
     compactBatchChannelRuns(func);
     compactRegularOpRuns(func);
     compactRowWiseWvmmRuns(func);
+    compactScalarChannelRuns(func, nextChannelId);
+    compactBatchChannelRuns(func);
+
+    auto eraseUnusedOps = [&](auto tag) {
+      using OpTy = decltype(tag);
+      SmallVector<OpTy> ops;
+      func.walk([&](OpTy op) { ops.push_back(op); });
+      for (auto op : llvm::reverse(ops))
+        if (op->use_empty())
+          op.erase();
+    };
+    eraseUnusedOps(tensor::ExtractSliceOp {});
+    eraseUnusedOps(spatial::SpatConcatOp {});
+    eraseUnusedOps(spatial::SpatExtractRowsOp {});
+
     if (!sortTopologically(&func.getBody().front())) {
       func.emitOpError("failed to topologically order merged Spatial IR");
       signalPassFailure();

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/RegularOpCompaction.cpp

@@ -13,6 +13,7 @@
 #include <tuple>
 
 #include "RegularOpCompaction.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/TensorPackingPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -41,133 +42,75 @@ struct RegularChunk {
   Value output;
 };
 
-static RankedTensorType getPackedTensorType(RankedTensorType elementType, int64_t count) {
-  SmallVector<int64_t> packedShape(elementType.getShape().begin(), elementType.getShape().end());
-  packedShape[0] *= count;
-  return RankedTensorType::get(packedShape, elementType.getElementType());
-}
+static spatial::SpatConcatOp getContiguousConcatUse(ValueRange values, unsigned& startOperandIndex) {
+  if (values.empty() || !values.front().hasOneUse())
+    return {};
 
-static Value
-extractPackedChunk(Value packedValue, RankedTensorType chunkType, unsigned index, IRRewriter& rewriter, Location loc) {
-  SmallVector<OpFoldResult> offsets;
-  SmallVector<OpFoldResult> sizes;
-  SmallVector<OpFoldResult> strides;
-  offsets.reserve(chunkType.getRank());
-  sizes.reserve(chunkType.getRank());
-  strides.reserve(chunkType.getRank());
+  OpOperand& firstUse = *values.front().getUses().begin();
+  auto concatOp = dyn_cast<spatial::SpatConcatOp>(firstUse.getOwner());
+  if (!concatOp)
+    return {};
 
-  offsets.push_back(rewriter.getIndexAttr(static_cast<int64_t>(index) * chunkType.getDimSize(0)));
-  sizes.push_back(rewriter.getIndexAttr(chunkType.getDimSize(0)));
-  strides.push_back(rewriter.getIndexAttr(1));
-  for (int64_t dim = 1; dim < chunkType.getRank(); ++dim) {
-    offsets.push_back(rewriter.getIndexAttr(0));
-    sizes.push_back(rewriter.getIndexAttr(chunkType.getDimSize(dim)));
-    strides.push_back(rewriter.getIndexAttr(1));
+  startOperandIndex = firstUse.getOperandNumber();
+  for (auto [index, value] : llvm::enumerate(values)) {
+    if (!value.hasOneUse())
+      return {};
+    OpOperand& use = *value.getUses().begin();
+    if (use.getOwner() != concatOp || use.getOperandNumber() != startOperandIndex + index)
+      return {};
   }
 
-  return tensor::ExtractSliceOp::create(rewriter, loc, chunkType, packedValue, offsets, sizes, strides).getResult();
+  return concatOp;
 }
 
-static Value createPackedExtractRowsSlice(
-  spatial::SpatExtractRowsOp extractRowsOp, unsigned startIndex, unsigned count, IRRewriter& rewriter, Location loc) {
-  auto rowType = dyn_cast<RankedTensorType>(extractRowsOp.getOutputs()[startIndex].getType());
-  auto inputType = dyn_cast<RankedTensorType>(extractRowsOp.getInput().getType());
-  if (!rowType || !inputType || !rowType.hasStaticShape() || !inputType.hasStaticShape() || rowType.getRank() == 0)
-    return {};
-
-  int64_t rowsPerValue = rowType.getDimSize(0);
-  if (ShapedType::isDynamic(rowsPerValue))
-    return {};
-
-  auto packedType = getPackedTensorType(rowType, static_cast<int64_t>(count));
-  SmallVector<OpFoldResult> offsets;
-  SmallVector<OpFoldResult> sizes;
-  SmallVector<OpFoldResult> strides;
-  offsets.reserve(inputType.getRank());
-  sizes.reserve(inputType.getRank());
-  strides.reserve(inputType.getRank());
-
-  offsets.push_back(rewriter.getIndexAttr(static_cast<int64_t>(startIndex) * rowsPerValue));
-  sizes.push_back(rewriter.getIndexAttr(static_cast<int64_t>(count) * rowsPerValue));
-  strides.push_back(rewriter.getIndexAttr(1));
-  for (int64_t dim = 1; dim < inputType.getRank(); ++dim) {
-    offsets.push_back(rewriter.getIndexAttr(0));
-    sizes.push_back(rewriter.getIndexAttr(inputType.getDimSize(dim)));
-    strides.push_back(rewriter.getIndexAttr(1));
+static void replaceConcatRunWithPackedValue(spatial::SpatConcatOp concatOp,
+                                            unsigned startOperandIndex,
+                                            unsigned operandCount,
+                                            Value packedValue,
+                                            IRRewriter& rewriter) {
+  SmallVector<Value> newInputs;
+  newInputs.reserve(concatOp.getInputs().size() - operandCount + 1);
+  for (auto [operandIndex, operand] : llvm::enumerate(concatOp.getInputs())) {
+    if (operandIndex == startOperandIndex)
+      newInputs.push_back(packedValue);
+    if (operandIndex < startOperandIndex || operandIndex >= startOperandIndex + operandCount)
+      newInputs.push_back(operand);
   }
-
-  return tensor::ExtractSliceOp::create(rewriter, loc, packedType, extractRowsOp.getInput(), offsets, sizes, strides)
-    .getResult();
+  if (newInputs.size() == 1 && newInputs.front().getType() == concatOp.getOutput().getType()) {
+    rewriter.replaceOp(concatOp, newInputs.front());
+    return;
+  }
+  rewriter.modifyOpInPlace(concatOp, [&] { concatOp->setOperands(newInputs); });
 }
 
-static Value createPackedExtractSliceTensor(ValueRange values, IRRewriter& rewriter, Location loc) {
-  if (values.empty())
-    return {};
-  if (values.size() == 1)
-    return values.front();
-
-  auto firstSliceOp = values.front().getDefiningOp<tensor::ExtractSliceOp>();
-  if (!firstSliceOp)
+static RankedTensorType
+getPackedConcatSliceType(spatial::SpatConcatOp concatOp, unsigned startOperandIndex, unsigned operandCount) {
+  auto firstType = dyn_cast<RankedTensorType>(concatOp.getInputs()[startOperandIndex].getType());
+  if (!firstType || !firstType.hasStaticShape())
     return {};
 
-  auto firstType = dyn_cast<RankedTensorType>(firstSliceOp.getResult().getType());
-  auto sourceType = dyn_cast<RankedTensorType>(firstSliceOp.getSource().getType());
-  if (!firstType || !sourceType || !firstType.hasStaticShape() || !sourceType.hasStaticShape()
-      || firstType.getRank() == 0)
+  int64_t axis = concatOp.getAxis();
+  if (axis < 0 || axis >= firstType.getRank())
     return {};
 
-  auto hasStaticValues = [](ArrayRef<int64_t> values) {
-    return llvm::all_of(values, [](int64_t value) { return !ShapedType::isDynamic(value); });
-  };
-  if (!hasStaticValues(firstSliceOp.getStaticOffsets()) || !hasStaticValues(firstSliceOp.getStaticSizes())
-      || !hasStaticValues(firstSliceOp.getStaticStrides()))
-    return {};
-
-  ArrayRef<int64_t> firstOffsets = firstSliceOp.getStaticOffsets();
-  ArrayRef<int64_t> firstSizes = firstSliceOp.getStaticSizes();
-  ArrayRef<int64_t> firstStrides = firstSliceOp.getStaticStrides();
-  int64_t rowsPerValue = firstSizes[0];
-  if (ShapedType::isDynamic(rowsPerValue))
-    return {};
-
-  for (size_t index = 1; index < values.size(); ++index) {
-    auto sliceOp = values[index].getDefiningOp<tensor::ExtractSliceOp>();
-    if (!sliceOp || sliceOp.getSource() != firstSliceOp.getSource()
-        || sliceOp.getResult().getType() != firstSliceOp.getResult().getType()
-        || !hasStaticValues(sliceOp.getStaticOffsets()) || !hasStaticValues(sliceOp.getStaticSizes())
-        || !hasStaticValues(sliceOp.getStaticStrides()))
+  SmallVector<int64_t> shape(firstType.getShape().begin(), firstType.getShape().end());
+  shape[axis] = 0;
+  for (unsigned index = 0; index < operandCount; ++index) {
+    auto operandType = dyn_cast<RankedTensorType>(concatOp.getInputs()[startOperandIndex + index].getType());
+    if (!operandType || !operandType.hasStaticShape() || operandType.getRank() != firstType.getRank())
       return {};
 
-    if (sliceOp.getStaticSizes() != firstSizes || sliceOp.getStaticStrides() != firstStrides)
-      return {};
-
-    if (sliceOp.getStaticOffsets()[0] != firstOffsets[0] + static_cast<int64_t>(index) * rowsPerValue)
-      return {};
-
-    for (int64_t dim = 1; dim < firstType.getRank(); ++dim)
-      if (sliceOp.getStaticOffsets()[dim] != firstOffsets[dim])
+    for (int64_t dim = 0; dim < firstType.getRank(); ++dim) {
+      if (dim == axis)
+        continue;
+      if (operandType.getShape()[dim] != shape[dim])
         return {};
+    }
+
+    shape[axis] += operandType.getShape()[axis];
   }
 
-  auto packedType = getPackedTensorType(firstType, static_cast<int64_t>(values.size()));
-  SmallVector<OpFoldResult> offsets;
-  SmallVector<OpFoldResult> sizes;
-  SmallVector<OpFoldResult> strides;
-  offsets.reserve(firstType.getRank());
-  sizes.reserve(firstType.getRank());
-  strides.reserve(firstType.getRank());
-
-  offsets.push_back(rewriter.getIndexAttr(firstOffsets[0]));
-  sizes.push_back(rewriter.getIndexAttr(rowsPerValue * static_cast<int64_t>(values.size())));
-  strides.push_back(rewriter.getIndexAttr(firstStrides[0]));
-  for (int64_t dim = 1; dim < firstType.getRank(); ++dim) {
-    offsets.push_back(rewriter.getIndexAttr(firstOffsets[dim]));
-    sizes.push_back(rewriter.getIndexAttr(firstSizes[dim]));
-    strides.push_back(rewriter.getIndexAttr(firstStrides[dim]));
-  }
-
-  return tensor::ExtractSliceOp::create(rewriter, loc, packedType, firstSliceOp.getSource(), offsets, sizes, strides)
-    .getResult();
+  return RankedTensorType::get(shape, firstType.getElementType());
 }
 
 static bool getContiguousOpResults(ValueRange values, Operation*& owner, unsigned& startIndex) {
@@ -207,8 +150,7 @@ static Value createPackedTensorForValues(ValueRange values, IRRewriter& rewriter
     return {};
   if (!llvm::all_of(values.drop_front(), [&](Value value) { return value.getType() == firstType; }))
     return {};
-
-  return tensor::ConcatOp::create(rewriter, loc, /*dim=*/0, values).getResult();
+  return {};
 }
 
 static bool areEquivalentRegularSteps(const RegularStep& lhs, const RegularStep& rhs) {
@@ -346,11 +288,28 @@ static void compactRegularChunkRun(IRRewriter& rewriter, ArrayRef<RegularChunk>
     scf::YieldOp::create(rewriter, anchorChunk.startOp->getLoc(), inserted.getResult());
   }
 
-  for (auto [index, chunk] : llvm::enumerate(run)) {
-    Value replacement = extractPackedChunk(
-      loop.getResult(0), outputType, static_cast<unsigned>(index), rewriter, chunk.startOp->getLoc());
-    Value output = chunk.output;
-    output.replaceAllUsesWith(replacement);
+  SmallVector<Value> outputs;
+  outputs.reserve(run.size());
+  for (const RegularChunk& chunk : run)
+    outputs.push_back(chunk.output);
+
+  unsigned concatStartIndex = 0;
+  auto concatOp = getContiguousConcatUse(ValueRange(outputs), concatStartIndex);
+  auto concatPackedType = concatOp
+                          ? getPackedConcatSliceType(concatOp, concatStartIndex, static_cast<unsigned>(outputs.size()))
+                          : RankedTensorType {};
+
+  if (concatOp && concatPackedType == packedOutputType) {
+    replaceConcatRunWithPackedValue(
+      concatOp, concatStartIndex, static_cast<unsigned>(outputs.size()), loop.getResult(0), rewriter);
+  }
+  else {
+    for (auto [index, chunk] : llvm::enumerate(run)) {
+      Value replacement = extractPackedChunk(
+        loop.getResult(0), outputType, static_cast<unsigned>(index), rewriter, chunk.startOp->getLoc());
+      Value output = chunk.output;
+      output.replaceAllUsesWith(replacement);
+    }
   }
 
   SmallVector<Operation*> opsToErase;
@@ -412,7 +371,18 @@ void compactScalarChannelRuns(func::FuncOp funcOp, int64_t& nextChannelId) {
           }
 
           auto rowType = cast<RankedTensorType>(run.front().getOutput().getType());
-          auto packedType = getPackedTensorType(rowType, static_cast<int64_t>(sortedEntries.size()));
+          auto fallbackPackedType = getPackedTensorType(rowType, static_cast<int64_t>(sortedEntries.size()));
+          SmallVector<Value> sortedOutputs;
+          sortedOutputs.reserve(sortedEntries.size());
+          for (ReceiveEntry& entry : sortedEntries)
+            sortedOutputs.push_back(entry.op.getOutput());
+
+          unsigned concatStartIndex = 0;
+          auto concatOp = getContiguousConcatUse(ValueRange(sortedOutputs), concatStartIndex);
+          auto concatPackedType =
+            concatOp ? getPackedConcatSliceType(concatOp, concatStartIndex, static_cast<unsigned>(sortedOutputs.size()))
+                     : RankedTensorType {};
+          auto packedType = concatPackedType ? concatPackedType : fallbackPackedType;
           rewriter.setInsertionPoint(run.front());
           auto compactReceive =
             spatial::SpatChannelReceiveTensorOp::create(rewriter,
@@ -421,9 +391,18 @@ void compactScalarChannelRuns(func::FuncOp funcOp, int64_t& nextChannelId) {
                                                         rewriter.getDenseI64ArrayAttr(channelIds),
                                                         rewriter.getDenseI32ArrayAttr(sourceCoreIds),
                                                         rewriter.getDenseI32ArrayAttr(targetCoreIds));
-          for (auto [sortedIndex, entry] : llvm::enumerate(sortedEntries))
-            entry.op.getOutput().replaceAllUsesWith(extractPackedChunk(
-              compactReceive.getOutput(), rowType, static_cast<unsigned>(sortedIndex), rewriter, entry.op.getLoc()));
+          if (concatOp && concatPackedType) {
+            replaceConcatRunWithPackedValue(concatOp,
+                                            concatStartIndex,
+                                            static_cast<unsigned>(sortedOutputs.size()),
+                                            compactReceive.getOutput(),
+                                            rewriter);
+          }
+          else {
+            for (auto [sortedIndex, entry] : llvm::enumerate(sortedEntries))
+              entry.op.getOutput().replaceAllUsesWith(extractPackedChunk(
+                compactReceive.getOutput(), rowType, static_cast<unsigned>(sortedIndex), rewriter, entry.op.getLoc()));
+          }
           for (auto op : run)
             rewriter.eraseOp(op);
 
@@ -531,7 +510,18 @@ void compactBatchChannelRuns(func::FuncOp funcOp) {
           }
 
           auto rowType = cast<RankedTensorType>(run.front().getOutput().getType());
-          auto packedType = getPackedTensorType(rowType, static_cast<int64_t>(run.size()));
+          auto fallbackPackedType = getPackedTensorType(rowType, static_cast<int64_t>(run.size()));
+          SmallVector<Value> outputs;
+          outputs.reserve(run.size());
+          for (auto op : run)
+            outputs.push_back(op.getOutput());
+
+          unsigned concatStartIndex = 0;
+          auto concatOp = getContiguousConcatUse(ValueRange(outputs), concatStartIndex);
+          auto concatPackedType =
+            concatOp ? getPackedConcatSliceType(concatOp, concatStartIndex, static_cast<unsigned>(outputs.size()))
+                     : RankedTensorType {};
+          auto packedType = concatPackedType ? concatPackedType : fallbackPackedType;
           rewriter.setInsertionPoint(run.front());
           auto compactReceive =
             spatial::SpatChannelReceiveTensorBatchOp::create(rewriter,
@@ -540,9 +530,15 @@ void compactBatchChannelRuns(func::FuncOp funcOp) {
                                                              rewriter.getDenseI64ArrayAttr(channelIds),
                                                              rewriter.getDenseI32ArrayAttr(sourceCoreIds),
                                                              rewriter.getDenseI32ArrayAttr(targetCoreIds));
-          for (auto [index, op] : llvm::enumerate(run))
-            op.getOutput().replaceAllUsesWith(extractPackedChunk(
-              compactReceive.getOutput(), rowType, static_cast<unsigned>(index), rewriter, op.getLoc()));
+          if (concatOp && concatPackedType) {
+            replaceConcatRunWithPackedValue(
+              concatOp, concatStartIndex, static_cast<unsigned>(outputs.size()), compactReceive.getOutput(), rewriter);
+          }
+          else {
+            for (auto [index, op] : llvm::enumerate(run))
+              op.getOutput().replaceAllUsesWith(extractPackedChunk(
+                compactReceive.getOutput(), rowType, static_cast<unsigned>(index), rewriter, op.getLoc()));
+          }
           for (auto op : run)
             rewriter.eraseOp(op);