fast pim bufferization using tensors

minor fixes

.gitignore

@@ -1,5 +1,17 @@
+.zed
 .idea
 **/.vscode
+
 .claude
+.codex
 AGENTS.md
+
+CMakeUserPresets.json
+
 build
+build_release
+cmake-build-debug
+cmake-build-release
+compile.sh
+
+**/__*

README.md

@@ -135,7 +135,7 @@ validate.py \
     --raptor-path ../cmake-build-release/Release/bin/onnx-mlir \
     --onnx-include-dir ../onnx-mlir/include \
     --operations-dir ./networks/yolo11n/depth_04 \
-    --crossbar-size 2048
+    --crossbar-size 2048 --crossbar-count 256 --core-count 1000
 ```
 
 Available networks under `validation/networks/`: `vgg16`, `yolo11n`.

src/PIM/Common/CMakeLists.txt

@@ -1,5 +1,12 @@
 add_pim_library(OMPimCommon
-  PimCommon.cpp
+  IR/AddressAnalysis.cpp
+  IR/CoreBlockUtils.cpp
+  IR/EntryPointUtils.cpp
+  IR/ShapeUtils.cpp
+  IR/WeightUtils.cpp
+  Support/DebugDump.cpp
+  Support/Diagnostics.cpp
+  Support/FileSystemUtils.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -0,0 +1,259 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
+
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+namespace onnx_mlir {
+
+mlir::memref::GlobalOp lookupGlobalForGetGlobal(mlir::ModuleOp moduleOp, mlir::memref::GetGlobalOp getGlobalOp) {
+  if (!moduleOp || !getGlobalOp)
+    return {};
+  return moduleOp.lookupSymbol<mlir::memref::GlobalOp>(getGlobalOp.getName());
+}
+
+namespace {
+
+mlir::Value resolveAlias(mlir::Value value, const StaticValueKnowledge* knowledge) {
+  if (!knowledge)
+    return value;
+
+  auto iter = knowledge->aliases.find(value);
+  while (iter != knowledge->aliases.end()) {
+    value = iter->second;
+    iter = knowledge->aliases.find(value);
+  }
+  return value;
+}
+
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
+  value = resolveAlias(value, knowledge);
+
+  if (mlir::isa<mlir::BlockArgument>(value))
+    return value;
+
+  mlir::Operation* definingOp = value.getDefiningOp();
+  if (!definingOp)
+    return value;
+
+  if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
+    if (auto result = mlir::dyn_cast<mlir::OpResult>(value))
+      if (mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(result))
+        return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge);
+  }
+
+  if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge);
+  if (auto collapseOp = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge);
+  if (auto expandOp = mlir::dyn_cast<mlir::memref::ExpandShapeOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge);
+
+  return value;
+}
+
+llvm::FailureOr<int64_t> resolveOpFoldResult(mlir::OpFoldResult ofr, const StaticValueKnowledge* knowledge);
+
+llvm::FailureOr<int64_t> resolveIndexValueImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
+  value = resolveAlias(value, knowledge);
+
+  if (knowledge) {
+    auto iter = knowledge->indexValues.find(value);
+    if (iter != knowledge->indexValues.end())
+      return iter->second;
+  }
+
+  auto constantOp = value.getDefiningOp<mlir::arith::ConstantOp>();
+  if (constantOp) {
+    if (auto integerAttr = mlir::dyn_cast<mlir::IntegerAttr>(constantOp.getValue()))
+      return integerAttr.getInt();
+  }
+
+  mlir::Operation* definingOp = value.getDefiningOp();
+  if (!definingOp)
+    return mlir::failure();
+
+  if (auto indexCastOp = mlir::dyn_cast<mlir::arith::IndexCastOp>(definingOp))
+    return resolveIndexValueImpl(indexCastOp.getIn(), knowledge);
+
+  if (auto addOp = mlir::dyn_cast<mlir::arith::AddIOp>(definingOp)) {
+    auto lhs = resolveIndexValueImpl(addOp.getLhs(), knowledge);
+    auto rhs = resolveIndexValueImpl(addOp.getRhs(), knowledge);
+    if (failed(lhs) || failed(rhs))
+      return mlir::failure();
+    return *lhs + *rhs;
+  }
+
+  if (auto subOp = mlir::dyn_cast<mlir::arith::SubIOp>(definingOp)) {
+    auto lhs = resolveIndexValueImpl(subOp.getLhs(), knowledge);
+    auto rhs = resolveIndexValueImpl(subOp.getRhs(), knowledge);
+    if (failed(lhs) || failed(rhs))
+      return mlir::failure();
+    return *lhs - *rhs;
+  }
+
+  if (auto mulOp = mlir::dyn_cast<mlir::arith::MulIOp>(definingOp)) {
+    auto lhs = resolveIndexValueImpl(mulOp.getLhs(), knowledge);
+    auto rhs = resolveIndexValueImpl(mulOp.getRhs(), knowledge);
+    if (failed(lhs) || failed(rhs))
+      return mlir::failure();
+    return *lhs * *rhs;
+  }
+
+  if (auto divOp = mlir::dyn_cast<mlir::arith::DivUIOp>(definingOp)) {
+    auto lhs = resolveIndexValueImpl(divOp.getLhs(), knowledge);
+    auto rhs = resolveIndexValueImpl(divOp.getRhs(), knowledge);
+    if (failed(lhs) || failed(rhs) || *rhs == 0)
+      return mlir::failure();
+    return static_cast<int64_t>(static_cast<uint64_t>(*lhs) / static_cast<uint64_t>(*rhs));
+  }
+
+  if (auto remOp = mlir::dyn_cast<mlir::arith::RemUIOp>(definingOp)) {
+    auto lhs = resolveIndexValueImpl(remOp.getLhs(), knowledge);
+    auto rhs = resolveIndexValueImpl(remOp.getRhs(), knowledge);
+    if (failed(lhs) || failed(rhs) || *rhs == 0)
+      return mlir::failure();
+    return static_cast<int64_t>(static_cast<uint64_t>(*lhs) % static_cast<uint64_t>(*rhs));
+  }
+
+  return mlir::failure();
+}
+
+llvm::FailureOr<int64_t> resolveOpFoldResult(mlir::OpFoldResult ofr, const StaticValueKnowledge* knowledge) {
+  if (auto attr = mlir::dyn_cast<mlir::Attribute>(ofr)) {
+    auto integerAttr = mlir::dyn_cast<mlir::IntegerAttr>(attr);
+    if (!integerAttr)
+      return mlir::failure();
+    return integerAttr.getInt();
+  }
+
+  return resolveIndexValueImpl(mlir::cast<mlir::Value>(ofr), knowledge);
+}
+
+llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Value value,
+                                                                        const StaticValueKnowledge* knowledge) {
+  int64_t byteOffset = 0;
+  value = resolveAlias(value, knowledge);
+
+  while (true) {
+    if (mlir::isa<mlir::BlockArgument>(value))
+      return ResolvedContiguousAddress {value, byteOffset};
+
+    mlir::Operation* definingOp = value.getDefiningOp();
+    if (!definingOp)
+      return mlir::failure();
+
+    if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
+      mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast<mlir::OpResult>(value));
+      if (!tiedOperand)
+        return mlir::failure();
+      value = resolveAlias(tiedOperand->get(), knowledge);
+      continue;
+    }
+
+    if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(definingOp)) {
+      auto result = mlir::dyn_cast<mlir::OpResult>(value);
+      if (!result)
+        return mlir::failure();
+
+      auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
+      mlir::Value yieldedValue = resolveLoopCarriedAliasImpl(yieldOp.getOperand(result.getResultNumber()), knowledge);
+      if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(yieldedValue)) {
+        if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0
+            && static_cast<unsigned>(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size()) {
+          value = resolveAlias(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge);
+          continue;
+        }
+      }
+
+      value = yieldedValue;
+      continue;
+    }
+
+    if (auto subviewOp = mlir::dyn_cast<mlir::memref::SubViewOp>(definingOp)) {
+      auto sourceType = mlir::dyn_cast<mlir::MemRefType>(subviewOp.getSource().getType());
+      auto subviewType = mlir::dyn_cast<mlir::MemRefType>(subviewOp.getType());
+      if (!sourceType || !subviewType || !sourceType.hasStaticShape() || !subviewType.hasStaticShape())
+        return mlir::failure();
+
+      llvm::SmallVector<int64_t> offsets;
+      llvm::SmallVector<int64_t> sizes;
+      llvm::SmallVector<int64_t> strides;
+      offsets.reserve(subviewOp.getMixedOffsets().size());
+      sizes.reserve(subviewOp.getMixedSizes().size());
+      strides.reserve(subviewOp.getMixedStrides().size());
+
+      for (mlir::OpFoldResult offset : subviewOp.getMixedOffsets()) {
+        auto resolvedOffset = resolveOpFoldResult(offset, knowledge);
+        if (failed(resolvedOffset))
+          return mlir::failure();
+        offsets.push_back(*resolvedOffset);
+      }
+
+      for (mlir::OpFoldResult size : subviewOp.getMixedSizes()) {
+        auto resolvedSize = resolveOpFoldResult(size, knowledge);
+        if (failed(resolvedSize))
+          return mlir::failure();
+        sizes.push_back(*resolvedSize);
+      }
+
+      for (mlir::OpFoldResult stride : subviewOp.getMixedStrides()) {
+        auto resolvedStride = resolveOpFoldResult(stride, knowledge);
+        if (failed(resolvedStride))
+          return mlir::failure();
+        strides.push_back(*resolvedStride);
+      }
+
+      if (!isMemoryContiguous(sourceType.getShape(), offsets, sizes, strides))
+        return mlir::failure();
+
+      auto sourceStrides = computeRowMajorStrides(sourceType.getShape());
+      byteOffset += linearizeIndex(offsets, sourceStrides) * subviewType.getElementTypeBitWidth() / 8;
+      value = resolveAlias(subviewOp.getSource(), knowledge);
+      continue;
+    }
+
+    if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(definingOp)) {
+      value = resolveAlias(castOp.getSource(), knowledge);
+      continue;
+    }
+    if (auto collapseOp = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(definingOp)) {
+      value = resolveAlias(collapseOp.getSrc(), knowledge);
+      continue;
+    }
+    if (auto expandOp = mlir::dyn_cast<mlir::memref::ExpandShapeOp>(definingOp)) {
+      value = resolveAlias(expandOp.getSrc(), knowledge);
+      continue;
+    }
+
+    if (mlir::isa<onnx_mlir::pim::PimEmptyManyOp, mlir::memref::AllocOp, mlir::memref::GetGlobalOp>(definingOp))
+      return ResolvedContiguousAddress {value, byteOffset};
+
+    return mlir::failure();
+  }
+}
+
+} // namespace
+
+llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value) { return resolveIndexValueImpl(value, nullptr); }
+
+llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge) {
+  return resolveIndexValueImpl(value, &knowledge);
+}
+
+llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value) {
+  return resolveContiguousAddressImpl(value, nullptr);
+}
+
+llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value,
+                                                                    const StaticValueKnowledge& knowledge) {
+  return resolveContiguousAddressImpl(value, &knowledge);
+}
+
+mlir::Value resolveLoopCarriedAlias(mlir::Value value, const StaticValueKnowledge& knowledge) {
+  return resolveLoopCarriedAliasImpl(value, &knowledge);
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/AddressAnalysis.hpp

@@ -0,0 +1,43 @@
+#pragma once
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Value.h"
+
+#include "llvm/ADT/DenseMap.h"
+
+namespace onnx_mlir {
+
+/// Describes a value as a base addressable object plus a statically known
+/// byte offset after peeling aliases, casts, and contiguous subviews.
+struct ResolvedContiguousAddress {
+  mlir::Value base;
+  int64_t byteOffset = 0;
+};
+
+/// Records compile-time facts used when interpreting address arithmetic and
+/// loop-carried aliases inside PIM regions.
+struct StaticValueKnowledge {
+  llvm::DenseMap<mlir::Value, int64_t> indexValues;
+  llvm::DenseMap<mlir::Value, mlir::Value> aliases;
+
+  StaticValueKnowledge() {}
+};
+
+mlir::memref::GlobalOp lookupGlobalForGetGlobal(mlir::ModuleOp moduleOp, mlir::memref::GetGlobalOp getGlobalOp);
+
+/// Resolves a value to contiguous backing storage when that storage can be
+/// proven statically from aliases, DPS ties, casts, and subviews.
+llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value);
+llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value,
+                                                                    const StaticValueKnowledge& knowledge);
+
+/// Statically evaluates index-like SSA values, including simple integer
+/// arithmetic and loop facts recorded in `knowledge`.
+llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value);
+llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge);
+
+/// Follows alias, view, and DPS chains to recover the backing value of a
+/// loop-carried memref/result.
+mlir::Value resolveLoopCarriedAlias(mlir::Value value, const StaticValueKnowledge& knowledge);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/CompactAsmUtils.hpp

@@ -0,0 +1,755 @@
+#ifndef ONNX_MLIR_PIM_COMPACT_ASM_UTILS_HPP
+#define ONNX_MLIR_PIM_COMPACT_ASM_UTILS_HPP
+
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/LogicalResult.h"
+
+namespace onnx_mlir {
+namespace compact_asm {
+
+using namespace mlir;
+
+enum class ListDelimiter {
+  Square,
+  Paren
+};
+
+inline ParseResult parseOpenDelimiter(OpAsmParser& parser, ListDelimiter delimiter) {
+  if (delimiter == ListDelimiter::Square)
+    return parser.parseLSquare();
+  return parser.parseLParen();
+}
+
+inline ParseResult parseOptionalCloseDelimiter(OpAsmParser& parser, ListDelimiter delimiter) {
+  if (delimiter == ListDelimiter::Square)
+    return parser.parseOptionalRSquare();
+  return parser.parseOptionalRParen();
+}
+
+inline void printOpenDelimiter(OpAsmPrinter& printer, ListDelimiter delimiter) {
+  printer << (delimiter == ListDelimiter::Square ? "[" : "(");
+}
+
+inline void printCloseDelimiter(OpAsmPrinter& printer, ListDelimiter delimiter) {
+  printer << (delimiter == ListDelimiter::Square ? "]" : ")");
+}
+
+template <typename EntryT, typename ParseEntryFn>
+inline ParseResult parseCompressedRepeatedList(OpAsmParser& parser,
+                                               ListDelimiter delimiter,
+                                               SmallVectorImpl<EntryT>& entries,
+                                               ParseEntryFn parseEntry) {
+  if (parseOpenDelimiter(parser, delimiter))
+    return failure();
+  if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+    return success();
+
+  while (true) {
+    EntryT entry;
+    if (parseEntry(entry))
+      return failure();
+
+    int64_t repeatCount = 1;
+    if (succeeded(parser.parseOptionalKeyword("x"))) {
+      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+    }
+    for (int64_t index = 0; index < repeatCount; ++index)
+      entries.push_back(entry);
+
+    if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+      return success();
+    if (parser.parseComma())
+      return failure();
+  }
+}
+
+template <typename IntT>
+inline ParseResult parseCompressedIntegerEntries(OpAsmParser& parser,
+                                                 ListDelimiter delimiter,
+                                                 SmallVectorImpl<IntT>& values) {
+  if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+    return success();
+
+  while (true) {
+    if (succeeded(parser.parseOptionalLParen())) {
+      SmallVector<IntT> subgroup;
+      if (parseCompressedIntegerEntries(parser, ListDelimiter::Paren, subgroup))
+        return failure();
+
+      int64_t repeatCount = 1;
+      if (succeeded(parser.parseOptionalKeyword("x"))) {
+        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+      }
+      for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
+        llvm::append_range(values, subgroup);
+    }
+    else {
+      int64_t first = 0;
+      if (parser.parseInteger(first))
+        return failure();
+
+      if (succeeded(parser.parseOptionalKeyword("to"))) {
+        int64_t last = 0;
+        if (parser.parseInteger(last) || last < first)
+          return parser.emitError(parser.getCurrentLocation(), "invalid ascending range");
+
+        int64_t step = 1;
+        if (succeeded(parser.parseOptionalKeyword("by"))) {
+          if (parser.parseInteger(step) || step <= 0)
+            return parser.emitError(parser.getCurrentLocation(), "step after 'by' must be positive");
+        }
+        int64_t repeatCount = 1;
+        if (succeeded(parser.parseOptionalKeyword("x"))) {
+          if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+            return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+        }
+        if ((last - first) % step != 0) {
+          return parser.emitError(
+            parser.getCurrentLocation(), "range end must be reachable from start using the given step");
+        }
+
+        for (int64_t value = first; value <= last; value += step)
+          for (int64_t index = 0; index < repeatCount; ++index)
+            values.push_back(static_cast<IntT>(value));
+      }
+      else {
+        int64_t repeatCount = 1;
+        if (succeeded(parser.parseOptionalKeyword("x"))) {
+          if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+            return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+        }
+        for (int64_t index = 0; index < repeatCount; ++index)
+          values.push_back(static_cast<IntT>(first));
+      }
+    }
+
+    if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+      return success();
+    if (parser.parseComma())
+      return failure();
+  }
+}
+
+template <typename IntT>
+inline ParseResult parseCompressedIntegerSequence(OpAsmParser& parser,
+                                                  ListDelimiter delimiter,
+                                                  SmallVectorImpl<IntT>& values) {
+  if (parseOpenDelimiter(parser, delimiter))
+    return failure();
+  return parseCompressedIntegerEntries(parser, delimiter, values);
+}
+
+template <typename RangeT, typename PrintEntryFn>
+inline void printCompressedEqualRuns(OpAsmPrinter& printer, RangeT entries, PrintEntryFn printEntry) {
+  for (size_t index = 0; index < entries.size();) {
+    size_t runEnd = index + 1;
+    while (runEnd < entries.size() && entries[runEnd] == entries[index])
+      ++runEnd;
+
+    if (index != 0)
+      printer << ", ";
+    printEntry(entries[index]);
+    size_t runLength = runEnd - index;
+    if (runLength > 1)
+      printer << " x" << runLength;
+    index = runEnd;
+  }
+}
+
+template <typename IntT>
+inline void printCompressedIntegerSequence(OpAsmPrinter& printer,
+                                           ArrayRef<IntT> values,
+                                           ListDelimiter delimiter) {
+  struct FlatCompression {
+    enum class Kind {
+      Single,
+      EqualRun,
+      Progression
+    };
+
+    Kind kind = Kind::Single;
+    size_t covered = 1;
+    size_t repeatCount = 1;
+    size_t progressionValueCount = 1;
+    int64_t step = 1;
+    IntT firstValue {};
+    IntT lastValue {};
+  };
+
+  auto computeFlatCompression = [&](size_t start) {
+    FlatCompression compression;
+    compression.firstValue = values[start];
+    compression.lastValue = values[start];
+
+    auto findEqualRunEnd = [&](size_t runStart) {
+      size_t runEnd = runStart + 1;
+      while (runEnd < values.size() && values[runEnd] == values[runStart])
+        ++runEnd;
+      return runEnd;
+    };
+
+    size_t firstRunEnd = findEqualRunEnd(start);
+    compression.repeatCount = firstRunEnd - start;
+    size_t progressionEnd = firstRunEnd;
+    int64_t step = 0;
+    IntT lastValue = values[start];
+
+    if (firstRunEnd < values.size()) {
+      size_t secondRunEnd = findEqualRunEnd(firstRunEnd);
+      step = static_cast<int64_t>(values[firstRunEnd]) - static_cast<int64_t>(values[start]);
+      if (step > 0 && secondRunEnd - firstRunEnd == compression.repeatCount) {
+        progressionEnd = secondRunEnd;
+        lastValue = values[firstRunEnd];
+        size_t currentRunStart = secondRunEnd;
+        while (currentRunStart < values.size()) {
+          size_t currentRunEnd = findEqualRunEnd(currentRunStart);
+          if (currentRunEnd - currentRunStart != compression.repeatCount)
+            break;
+          if (static_cast<int64_t>(values[currentRunStart]) != static_cast<int64_t>(lastValue) + step)
+            break;
+          lastValue = values[currentRunStart];
+          progressionEnd = currentRunEnd;
+          currentRunStart = currentRunEnd;
+        }
+      }
+      else {
+        step = 0;
+      }
+    }
+
+    compression.covered = 1;
+    if (progressionEnd > firstRunEnd) {
+      size_t progressionValueCount = (progressionEnd - start) / compression.repeatCount;
+      if (progressionValueCount >= 3) {
+        compression.kind = FlatCompression::Kind::Progression;
+        compression.covered = progressionEnd - start;
+        compression.progressionValueCount = progressionValueCount;
+        compression.step = step;
+        compression.lastValue = lastValue;
+        return compression;
+      }
+    }
+
+    if (compression.repeatCount > 1) {
+      compression.kind = FlatCompression::Kind::EqualRun;
+      compression.covered = compression.repeatCount;
+      return compression;
+    }
+
+    return compression;
+  };
+
+  auto findRepeatedSublist = [&](size_t start) {
+    size_t bestLength = 0;
+    size_t bestRepeatCount = 1;
+    size_t remaining = values.size() - start;
+
+    for (size_t length = 2; length * 2 <= remaining; ++length) {
+      size_t repeatCount = 1;
+      ArrayRef<IntT> candidate = values.slice(start, length);
+      while (start + (repeatCount + 1) * length <= values.size()
+             && llvm::equal(candidate, values.slice(start + repeatCount * length, length))) {
+        ++repeatCount;
+      }
+
+      if (repeatCount <= 1)
+        continue;
+
+      size_t covered = length * repeatCount;
+      size_t bestCovered = bestLength * bestRepeatCount;
+      if (covered > bestCovered || (covered == bestCovered && length < bestLength)) {
+        bestLength = length;
+        bestRepeatCount = repeatCount;
+      }
+    }
+
+    return std::pair(bestLength, bestRepeatCount);
+  };
+
+  printOpenDelimiter(printer, delimiter);
+  for (size_t index = 0; index < values.size();) {
+    if (index != 0)
+      printer << ", ";
+
+    FlatCompression flat = computeFlatCompression(index);
+    auto [sublistLength, sublistRepeatCount] = findRepeatedSublist(index);
+    size_t repeatedSublistCoverage = sublistLength * sublistRepeatCount;
+    if (sublistRepeatCount > 1 && sublistLength > 1 && repeatedSublistCoverage > flat.covered) {
+      printCompressedIntegerSequence(printer, values.slice(index, sublistLength), ListDelimiter::Paren);
+      printer << " x" << sublistRepeatCount;
+      index += repeatedSublistCoverage;
+      continue;
+    }
+
+    switch (flat.kind) {
+    case FlatCompression::Kind::Progression:
+      printer << flat.firstValue << " to " << flat.lastValue;
+      if (flat.step != 1)
+        printer << " by " << flat.step;
+      if (flat.repeatCount > 1)
+        printer << " x" << flat.repeatCount;
+      index += flat.covered;
+      break;
+    case FlatCompression::Kind::EqualRun:
+      printer << flat.firstValue << " x" << flat.repeatCount;
+      index += flat.covered;
+      break;
+    case FlatCompression::Kind::Single:
+      printer << flat.firstValue;
+      index += flat.covered;
+      break;
+    }
+  }
+  printCloseDelimiter(printer, delimiter);
+}
+
+template <typename IntT>
+inline ParseResult parseCompressedIntegerList(OpAsmParser& parser, SmallVectorImpl<IntT>& values) {
+  return parseCompressedIntegerSequence(parser, ListDelimiter::Square, values);
+}
+
+template <typename IntT>
+inline void printCompressedIntegerList(OpAsmPrinter& printer, ArrayRef<IntT> values) {
+  printCompressedIntegerSequence(printer, values, ListDelimiter::Square);
+}
+
+inline void printCompressedValueSequence(OpAsmPrinter& printer, ValueRange values) {
+  for (size_t index = 0; index < values.size();) {
+    size_t equalRunEnd = index + 1;
+    while (equalRunEnd < values.size() && values[equalRunEnd] == values[index])
+      ++equalRunEnd;
+
+    if (index != 0)
+      printer << ", ";
+    if (equalRunEnd - index > 1) {
+      printer.printOperand(values[index]);
+      printer << " x" << (equalRunEnd - index);
+      index = equalRunEnd;
+      continue;
+    }
+
+    size_t rangeEnd = index + 1;
+    if (auto firstResult = dyn_cast<OpResult>(values[index])) {
+      while (rangeEnd < values.size()) {
+        auto nextResult = dyn_cast<OpResult>(values[rangeEnd]);
+        if (!nextResult || nextResult.getOwner() != firstResult.getOwner()
+            || nextResult.getResultNumber() != firstResult.getResultNumber() + (rangeEnd - index)) {
+          break;
+        }
+        ++rangeEnd;
+      }
+    }
+    else if (auto firstArg = dyn_cast<BlockArgument>(values[index])) {
+      while (rangeEnd < values.size()) {
+        auto nextArg = dyn_cast<BlockArgument>(values[rangeEnd]);
+        if (!nextArg || nextArg.getOwner() != firstArg.getOwner()
+            || nextArg.getArgNumber() != firstArg.getArgNumber() + (rangeEnd - index)) {
+          break;
+        }
+        ++rangeEnd;
+      }
+    }
+
+    printer.printOperand(values[index]);
+    if (rangeEnd - index >= 3) {
+      printer << " to ";
+      printer.printOperand(values[rangeEnd - 1]);
+    }
+    else if (rangeEnd - index == 2) {
+      printer << ", ";
+      printer.printOperand(values[index + 1]);
+    }
+    index = rangeEnd;
+  }
+}
+
+inline void printCompressedTypeSequence(OpAsmPrinter& printer, TypeRange types) {
+  printCompressedEqualRuns(printer, types, [&](Type type) { printer.printType(type); });
+}
+
+inline void printCompressedValueList(OpAsmPrinter& printer, ValueRange values, ListDelimiter delimiter) {
+  printOpenDelimiter(printer, delimiter);
+  printCompressedValueSequence(printer, values);
+  printCloseDelimiter(printer, delimiter);
+}
+
+inline void printCompressedTypeList(OpAsmPrinter& printer, TypeRange types, ListDelimiter delimiter) {
+  printOpenDelimiter(printer, delimiter);
+  printCompressedTypeSequence(printer, types);
+  printCloseDelimiter(printer, delimiter);
+}
+
+inline ParseResult parseCompressedTypeSequence(OpAsmParser& parser,
+                                               SmallVectorImpl<Type>& types,
+                                               bool allowEmpty) {
+  Type firstType;
+  OptionalParseResult firstTypeResult = parser.parseOptionalType(firstType);
+  if (!firstTypeResult.has_value()) {
+    if (allowEmpty)
+      return success();
+    return parser.emitError(parser.getCurrentLocation(), "expected type");
+  }
+  if (failed(*firstTypeResult))
+    return failure();
+
+  auto appendType = [&](Type type) -> ParseResult {
+    int64_t repeatCount = 1;
+    if (succeeded(parser.parseOptionalKeyword("x"))) {
+      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+    }
+    for (int64_t index = 0; index < repeatCount; ++index)
+      types.push_back(type);
+    return success();
+  };
+
+  if (appendType(firstType))
+    return failure();
+  while (succeeded(parser.parseOptionalComma())) {
+    Type nextType;
+    if (parser.parseType(nextType) || appendType(nextType))
+      return failure();
+  }
+  return success();
+}
+
+inline ParseResult parseCompressedOperandEntryWithFirst(
+  OpAsmParser& parser,
+  OpAsmParser::UnresolvedOperand firstOperand,
+  SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  if (succeeded(parser.parseOptionalKeyword("to"))) {
+    OpAsmParser::UnresolvedOperand lastOperand;
+    if (parser.parseOperand(lastOperand))
+      return failure();
+    if (firstOperand.name != lastOperand.name || firstOperand.number > lastOperand.number)
+      return parser.emitError(parser.getCurrentLocation(), "invalid operand range");
+    for (unsigned number = firstOperand.number; number <= lastOperand.number; ++number)
+      operands.push_back({firstOperand.location, firstOperand.name, number});
+    return success();
+  }
+
+  int64_t repeatCount = 1;
+  if (succeeded(parser.parseOptionalKeyword("x"))) {
+    if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+      return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+  }
+  for (int64_t index = 0; index < repeatCount; ++index)
+    operands.push_back(firstOperand);
+  return success();
+}
+
+inline ParseResult parseOneCompressedOperandEntry(
+  OpAsmParser& parser,
+  SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  OpAsmParser::UnresolvedOperand firstOperand;
+  if (parser.parseOperand(firstOperand))
+    return failure();
+  return parseCompressedOperandEntryWithFirst(parser, firstOperand, operands);
+}
+
+inline ParseResult parseCompressedOperandList(OpAsmParser& parser,
+                                              ListDelimiter delimiter,
+                                              SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  if (parseOpenDelimiter(parser, delimiter))
+    return failure();
+  if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+    return success();
+
+  while (true) {
+    if (parseOneCompressedOperandEntry(parser, operands))
+      return failure();
+    if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+      return success();
+    if (parser.parseComma())
+      return failure();
+  }
+}
+
+inline ParseResult parseCompressedOperandSequence(
+  OpAsmParser& parser,
+  SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  if (parseOneCompressedOperandEntry(parser, operands))
+    return failure();
+  while (succeeded(parser.parseOptionalComma()))
+    if (parseOneCompressedOperandEntry(parser, operands))
+      return failure();
+  return success();
+}
+
+inline ParseResult parseCompressedTypeList(OpAsmParser& parser,
+                                           ListDelimiter delimiter,
+                                           SmallVectorImpl<Type>& types) {
+  if (parseOpenDelimiter(parser, delimiter))
+    return failure();
+  if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+    return success();
+
+  if (parseCompressedTypeSequence(parser, types, /*allowEmpty=*/false))
+    return failure();
+  return parseOptionalCloseDelimiter(parser, delimiter);
+}
+
+inline bool hasRepeatedTuple(ValueRange values, size_t tupleSize) {
+  if (tupleSize == 0 || values.empty() || values.size() % tupleSize != 0)
+    return false;
+
+  SmallVector<Value> valueVec(values.begin(), values.end());
+  ArrayRef<Value> tuple(valueVec.data(), tupleSize);
+  for (size_t index = tupleSize; index < values.size(); index += tupleSize)
+    if (!llvm::equal(tuple, ArrayRef<Value>(valueVec).slice(index, tupleSize)))
+      return false;
+  return true;
+}
+
+inline bool hasRepeatedTuple(TypeRange types, size_t tupleSize) {
+  if (tupleSize == 0 || types.empty() || types.size() % tupleSize != 0)
+    return false;
+
+  SmallVector<Type> typeVec(types.begin(), types.end());
+  ArrayRef<Type> tuple(typeVec.data(), tupleSize);
+  for (size_t index = tupleSize; index < types.size(); index += tupleSize)
+    if (!llvm::equal(tuple, ArrayRef<Type>(typeVec).slice(index, tupleSize)))
+      return false;
+  return true;
+}
+
+inline void printValueTupleRun(OpAsmPrinter& printer,
+                               ValueRange values,
+                               size_t tupleSize,
+                               ListDelimiter delimiter) {
+  printOpenDelimiter(printer, delimiter);
+  printOpenDelimiter(printer, ListDelimiter::Paren);
+  for (size_t index = 0; index < tupleSize; ++index) {
+    if (index != 0)
+      printer << ", ";
+    printer.printOperand(values[index]);
+  }
+  printCloseDelimiter(printer, ListDelimiter::Paren);
+  printer << " x" << (values.size() / tupleSize);
+  printCloseDelimiter(printer, delimiter);
+}
+
+inline void printTypeTupleRun(OpAsmPrinter& printer,
+                              TypeRange types,
+                              size_t tupleSize,
+                              ListDelimiter delimiter) {
+  printOpenDelimiter(printer, delimiter);
+  printOpenDelimiter(printer, ListDelimiter::Paren);
+  for (size_t index = 0; index < tupleSize; ++index) {
+    if (index != 0)
+      printer << ", ";
+    printer.printType(types[index]);
+  }
+  printCloseDelimiter(printer, ListDelimiter::Paren);
+  printer << " x" << (types.size() / tupleSize);
+  printCloseDelimiter(printer, delimiter);
+}
+
+inline ParseResult parseCompressedOrTupleOperandList(
+  OpAsmParser& parser,
+  ListDelimiter delimiter,
+  SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  if (parseOpenDelimiter(parser, delimiter))
+    return failure();
+  if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+    return success();
+
+  if (succeeded(parser.parseOptionalLParen())) {
+    SmallVector<OpAsmParser::UnresolvedOperand> tupleOperands;
+    if (parseCompressedOperandSequence(parser, tupleOperands) || parser.parseRParen())
+      return failure();
+
+    int64_t repeatCount = 1;
+    if (succeeded(parser.parseOptionalKeyword("x"))) {
+      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+    }
+    for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
+      llvm::append_range(operands, tupleOperands);
+
+    while (succeeded(parser.parseOptionalComma())) {
+      if (parser.parseLParen())
+        return failure();
+      tupleOperands.clear();
+      if (parseCompressedOperandSequence(parser, tupleOperands) || parser.parseRParen())
+        return failure();
+
+      repeatCount = 1;
+      if (succeeded(parser.parseOptionalKeyword("x"))) {
+        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+      }
+      for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
+        llvm::append_range(operands, tupleOperands);
+    }
+    return parseOptionalCloseDelimiter(parser, delimiter);
+  }
+
+  while (true) {
+    if (parseOneCompressedOperandEntry(parser, operands))
+      return failure();
+    if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+      return success();
+    if (parser.parseComma())
+      return failure();
+  }
+}
+
+inline ParseResult parseCompressedOrTupleTypeList(OpAsmParser& parser,
+                                                  ListDelimiter delimiter,
+                                                  SmallVectorImpl<Type>& types) {
+  if (parseOpenDelimiter(parser, delimiter))
+    return failure();
+  if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+    return success();
+
+  if (succeeded(parser.parseOptionalLParen())) {
+    SmallVector<Type> tupleTypes;
+    if (parseCompressedTypeSequence(parser, tupleTypes, /*allowEmpty=*/false) || parser.parseRParen())
+      return failure();
+
+    int64_t repeatCount = 1;
+    if (succeeded(parser.parseOptionalKeyword("x"))) {
+      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+    }
+    for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
+      llvm::append_range(types, tupleTypes);
+
+    while (succeeded(parser.parseOptionalComma())) {
+      if (parser.parseLParen())
+        return failure();
+      tupleTypes.clear();
+      if (parseCompressedTypeSequence(parser, tupleTypes, /*allowEmpty=*/false) || parser.parseRParen())
+        return failure();
+
+      repeatCount = 1;
+      if (succeeded(parser.parseOptionalKeyword("x"))) {
+        if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+          return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+      }
+      for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
+        llvm::append_range(types, tupleTypes);
+    }
+    return parseOptionalCloseDelimiter(parser, delimiter);
+  }
+
+  while (true) {
+    Type type;
+    if (parser.parseType(type))
+      return failure();
+
+    int64_t repeatCount = 1;
+    if (succeeded(parser.parseOptionalKeyword("x"))) {
+      if (parser.parseInteger(repeatCount) || repeatCount <= 0)
+        return parser.emitError(parser.getCurrentLocation(), "repeat count after 'x' must be positive");
+    }
+    for (int64_t repeat = 0; repeat < repeatCount; ++repeat)
+      types.push_back(type);
+
+    if (succeeded(parseOptionalCloseDelimiter(parser, delimiter)))
+      return success();
+    if (parser.parseComma())
+      return failure();
+  }
+}
+
+inline void printArgumentBindings(OpAsmPrinter& printer, Block& block, ValueRange operands) {
+  if (block.getNumArguments() == 0) {
+    printer << "() = ()";
+    return;
+  }
+
+  if (block.getNumArguments() == 1) {
+    printer.printOperand(block.getArgument(0));
+    printer << " = ";
+    printCompressedValueList(printer, operands, ListDelimiter::Paren);
+    return;
+  }
+
+  printCompressedValueList(printer, ValueRange(block.getArguments()), ListDelimiter::Paren);
+  printer << " = ";
+  printCompressedValueList(printer, operands, ListDelimiter::Paren);
+}
+
+inline ParseResult parseCompressedArgumentEntryWithFirst(OpAsmParser& parser,
+                                                         OpAsmParser::Argument firstArgument,
+                                                         SmallVectorImpl<OpAsmParser::Argument>& arguments) {
+  if (succeeded(parser.parseOptionalKeyword("to"))) {
+    OpAsmParser::Argument lastArgument;
+    if (parser.parseArgument(lastArgument))
+      return failure();
+    if (firstArgument.ssaName.name != lastArgument.ssaName.name
+        || firstArgument.ssaName.number > lastArgument.ssaName.number) {
+      return parser.emitError(parser.getCurrentLocation(), "invalid argument range");
+    }
+    for (unsigned number = firstArgument.ssaName.number; number <= lastArgument.ssaName.number; ++number) {
+      OpAsmParser::Argument argument;
+      argument.ssaName = {firstArgument.ssaName.location, firstArgument.ssaName.name, number};
+      arguments.push_back(argument);
+    }
+    return success();
+  }
+
+  arguments.push_back(firstArgument);
+  return success();
+}
+
+inline ParseResult parseOneCompressedArgumentEntry(OpAsmParser& parser,
+                                                   SmallVectorImpl<OpAsmParser::Argument>& arguments) {
+  OpAsmParser::Argument firstArgument;
+  if (parser.parseArgument(firstArgument))
+    return failure();
+  return parseCompressedArgumentEntryWithFirst(parser, firstArgument, arguments);
+}
+
+inline void applyArgumentTypes(ArrayRef<Type> inputTypes, SmallVectorImpl<OpAsmParser::Argument>& arguments) {
+  for (auto [argument, inputType] : llvm::zip_equal(arguments, inputTypes))
+    argument.type = inputType;
+}
+
+inline ParseResult parseArgumentBindings(OpAsmParser& parser,
+                                         SmallVectorImpl<OpAsmParser::Argument>& arguments,
+                                         SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  if (succeeded(parser.parseOptionalLParen())) {
+    if (succeeded(parser.parseOptionalRParen())) {
+      if (parser.parseEqual() || parseCompressedOperandList(parser, ListDelimiter::Paren, operands))
+        return failure();
+      return success();
+    }
+
+    OpAsmParser::Argument firstArgument;
+    if (parser.parseArgument(firstArgument) || parseCompressedArgumentEntryWithFirst(parser, firstArgument, arguments))
+      return failure();
+    while (succeeded(parser.parseOptionalComma()))
+      if (parseOneCompressedArgumentEntry(parser, arguments))
+        return failure();
+    if (parser.parseRParen() || parser.parseEqual()
+        || parseCompressedOperandList(parser, ListDelimiter::Paren, operands)) {
+      return failure();
+    }
+    return success();
+  }
+
+  OpAsmParser::Argument argument;
+  if (parser.parseArgument(argument) || parser.parseEqual()
+      || parseCompressedOperandList(parser, ListDelimiter::Paren, operands)) {
+    return failure();
+  }
+  arguments.push_back(argument);
+  return success();
+}
+
+} // namespace compact_asm
+} // namespace onnx_mlir
+
+#endif

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -0,0 +1,67 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+
+#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+namespace onnx_mlir {
+
+bool isCoreStaticAddressOp(mlir::Operation* op) {
+  return mlir::isa<mlir::arith::ConstantOp,
+                   mlir::arith::AddIOp,
+                   mlir::arith::SubIOp,
+                   mlir::arith::MulIOp,
+                   mlir::arith::DivUIOp,
+                   mlir::arith::RemUIOp,
+                   mlir::arith::IndexCastOp,
+                   mlir::memref::AllocOp,
+                   mlir::memref::SubViewOp,
+                   mlir::memref::CastOp,
+                   mlir::memref::CollapseShapeOp,
+                   mlir::memref::ExpandShapeOp>(op);
+}
+
+mlir::LogicalResult
+walkPimCoreBlock(mlir::Block& block,
+                 const StaticValueKnowledge& knowledge,
+                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
+  bool hasFailure = false;
+  for (mlir::Operation& op : block) {
+    if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
+      continue;
+
+    if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(op)) {
+      mlir::Block& loopBody = forOp.getRegion().front();
+      auto lowerBound = resolveIndexValue(forOp.getLowerBound(), knowledge);
+      auto upperBound = resolveIndexValue(forOp.getUpperBound(), knowledge);
+      auto step = resolveIndexValue(forOp.getStep(), knowledge);
+      if (failed(lowerBound) || failed(upperBound) || failed(step) || *step <= 0) {
+        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
+        hasFailure = true;
+        continue;
+      }
+
+      llvm::SmallVector<mlir::Value> iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end());
+      for (int64_t inductionValue = *lowerBound; inductionValue < *upperBound; inductionValue += *step) {
+        StaticValueKnowledge loopKnowledge = knowledge;
+        loopKnowledge.indexValues[forOp.getInductionVar()] = inductionValue;
+        for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), iterValues))
+          loopKnowledge.aliases[iterArg] = iterValue;
+
+        if (failed(walkPimCoreBlock(loopBody, loopKnowledge, callback)))
+          hasFailure = true;
+
+        auto yieldOp = mlir::cast<mlir::scf::YieldOp>(loopBody.getTerminator());
+        for (auto [index, yieldedValue] : llvm::enumerate(yieldOp.getOperands()))
+          iterValues[index] = resolveLoopCarriedAlias(yieldedValue, loopKnowledge);
+      }
+      continue;
+    }
+
+    if (failed(callback(op, knowledge)))
+      hasFailure = true;
+  }
+  return mlir::success(!hasFailure);
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/CoreBlockUtils.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "mlir/IR/Block.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/STLFunctionalExtras.h"
+
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+
+namespace onnx_mlir {
+
+/// Returns true for ops in a `pim.core` body that only participate in static
+/// address or index computation and therefore do not emit PIM instructions.
+bool isCoreStaticAddressOp(mlir::Operation* op);
+
+/// Walks a `pim.core` body, statically unrolling nested `scf.for` loops when
+/// their bounds are known and invoking `callback` only on instruction-emitting
+/// operations.
+mlir::LogicalResult
+walkPimCoreBlock(mlir::Block& block,
+                 const StaticValueKnowledge& knowledge,
+                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/EntryPointUtils.cpp

@@ -0,0 +1,45 @@
+#include "src/Accelerators/PIM/Common/IR/EntryPointUtils.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+namespace onnx_mlir {
+
+llvm::FailureOr<mlir::func::FuncOp> getPimEntryFunc(mlir::ModuleOp moduleOp) {
+  if (!moduleOp)
+    return mlir::failure();
+
+  llvm::SmallVector<mlir::ONNXEntryPointOp> entryPoints(moduleOp.getOps<mlir::ONNXEntryPointOp>());
+  if (entryPoints.size() > 1) {
+    moduleOp.emitError("PIM pipeline requires a single ONNX entry point, but found ") << entryPoints.size();
+    return mlir::failure();
+  }
+  if (!entryPoints.empty()) {
+    auto entryPointAttr =
+      entryPoints.front()->getAttrOfType<mlir::SymbolRefAttr>(mlir::ONNXEntryPointOp::getEntryPointFuncAttrName());
+    if (!entryPointAttr) {
+      entryPoints.front().emitOpError("is missing the entry point function attribute");
+      return mlir::failure();
+    }
+    auto entryFunc = moduleOp.lookupSymbol<mlir::func::FuncOp>(entryPointAttr.getLeafReference().getValue());
+    if (!entryFunc) {
+      entryPoints.front().emitOpError("references an unknown entry function ")
+        << entryPointAttr.getLeafReference().getValue();
+      return mlir::failure();
+    }
+    return entryFunc;
+  }
+
+  if (auto mainGraphFunc = moduleOp.lookupSymbol<mlir::func::FuncOp>("main_graph"))
+    return mainGraphFunc;
+
+  llvm::SmallVector<mlir::func::FuncOp> nonExternalFuncs;
+  for (auto funcOp : moduleOp.getOps<mlir::func::FuncOp>())
+    if (!funcOp.isExternal())
+      nonExternalFuncs.push_back(funcOp);
+  if (nonExternalFuncs.size() == 1)
+    return nonExternalFuncs.front();
+
+  moduleOp.emitError("could not resolve a unique PIM entry function");
+  return mlir::failure();
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/EntryPointUtils.hpp

@@ -0,0 +1,13 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+
+namespace onnx_mlir {
+
+/// Resolves the function the PIM pipeline should treat as its entry point.
+/// Prefers ONNX entry-point metadata, then `main_graph`, then the only
+/// non-external function if the module is otherwise unambiguous.
+llvm::FailureOr<mlir::func::FuncOp> getPimEntryFunc(mlir::ModuleOp moduleOp);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/ShapeUtils.cpp

@@ -0,0 +1,89 @@
+#include "llvm/ADT/STLExtras.h"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+
+namespace onnx_mlir {
+
+llvm::SmallVector<int64_t> computeRowMajorStrides(llvm::ArrayRef<int64_t> shape) {
+  llvm::SmallVector<int64_t> strides(shape.size(), 1);
+  for (int64_t dim = static_cast<int64_t>(shape.size()) - 2; dim >= 0; --dim)
+    strides[dim] = strides[dim + 1] * shape[dim + 1];
+  return strides;
+}
+
+llvm::SmallVector<int64_t>
+delinearizeIndex(int64_t linearIndex, llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides) {
+  llvm::SmallVector<int64_t> indices(shape.size(), 0);
+  for (auto [dim, stride] : llvm::enumerate(strides)) {
+    indices[dim] = linearIndex / stride;
+    linearIndex %= stride;
+  }
+  return indices;
+}
+
+int64_t linearizeIndex(llvm::ArrayRef<int64_t> indices, llvm::ArrayRef<int64_t> strides) {
+  int64_t linearIndex = 0;
+  for (auto [index, stride] : llvm::zip_equal(indices, strides))
+    linearIndex += index * stride;
+  return linearIndex;
+}
+
+int64_t getNumElements(llvm::ArrayRef<int64_t> shape) {
+  int64_t numElements = 1;
+  for (int64_t dim : shape)
+    numElements *= dim;
+  return numElements;
+}
+
+bool isMemoryContiguous(llvm::ArrayRef<int64_t> srcShape,
+                        llvm::ArrayRef<int64_t> offsets,
+                        llvm::ArrayRef<int64_t> sizes,
+                        llvm::ArrayRef<int64_t> strides) {
+  if (std::any_of(strides.begin(), strides.end(), [](int64_t stride) -> bool { return stride != 1; }))
+    return false;
+
+  auto offsetsAndSizesAndShape = llvm::zip_equal(llvm::make_range(offsets.rbegin(), offsets.rend()),
+                                                 llvm::make_range(sizes.rbegin(), sizes.rend()),
+                                                 llvm::make_range(srcShape.rbegin(), srcShape.rend()));
+
+  auto firstNonZeroOffset = std::find_if(
+    offsetsAndSizesAndShape.begin(), offsetsAndSizesAndShape.end(), [&](auto offsetAndSizeAndShape) -> bool {
+      auto [offset, _size, _dimension] = offsetAndSizeAndShape;
+      return offset != 0;
+    });
+
+  if (firstNonZeroOffset != offsetsAndSizesAndShape.end()) {
+    auto [offset, size, dimension] = *firstNonZeroOffset;
+    if (size > dimension - offset)
+      return false;
+    ++firstNonZeroOffset;
+
+    if (std::any_of(firstNonZeroOffset, offsetsAndSizesAndShape.end(), [](auto offsetAndSizeAndShape) -> bool {
+          auto [_offset, size, _dimension] = offsetAndSizeAndShape;
+          return size != 1;
+        }))
+      return false;
+  }
+
+  auto sizesAndShape = llvm::zip_equal(llvm::make_range(sizes.rbegin(), sizes.rend()),
+                                       llvm::make_range(srcShape.rbegin(), srcShape.rend()));
+
+  auto firstDifferentSize = std::find_if(sizesAndShape.begin(), sizesAndShape.end(), [&](auto sizeAndShape) -> bool {
+    auto [size, dimension] = sizeAndShape;
+    return size != dimension;
+  });
+
+  if (firstDifferentSize != sizesAndShape.end()) {
+    ++firstDifferentSize;
+
+    if (std::any_of(firstDifferentSize, sizesAndShape.end(), [](auto sizeAndShape) -> bool {
+          auto [size, _dimension] = sizeAndShape;
+          return size != 1;
+        }))
+      return false;
+  }
+
+  return true;
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/ShapeUtils.hpp

@@ -0,0 +1,22 @@
+#pragma once
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace onnx_mlir {
+
+llvm::SmallVector<int64_t> computeRowMajorStrides(llvm::ArrayRef<int64_t> shape);
+
+llvm::SmallVector<int64_t>
+delinearizeIndex(int64_t linearIndex, llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides);
+
+int64_t linearizeIndex(llvm::ArrayRef<int64_t> indices, llvm::ArrayRef<int64_t> strides);
+
+int64_t getNumElements(llvm::ArrayRef<int64_t> shape);
+
+bool isMemoryContiguous(llvm::ArrayRef<int64_t> srcShape,
+                        llvm::ArrayRef<int64_t> offsets,
+                        llvm::ArrayRef<int64_t> sizes,
+                        llvm::ArrayRef<int64_t> strides);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/WeightUtils.cpp

@@ -0,0 +1,101 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
+
+#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+namespace onnx_mlir {
+
+bool hasWeightAlways(mlir::Operation* op) { return op && op->getAttr(PimWeightAlwaysAttrName) != nullptr; }
+
+void markWeightAlways(mlir::Operation* op) {
+  assert(op && "expected valid op");
+  op->setAttr(PimWeightAlwaysAttrName, mlir::UnitAttr::get(op->getContext()));
+}
+
+namespace {
+
+template <typename MVMOpTy, typename VMMOpTy, typename ParentOpTy>
+bool hasMvmVmmWeightUse(ParentOpTy parentOp, unsigned weightIndex) {
+  bool found = false;
+  parentOp.walk([&](mlir::Operation* op) {
+    if (auto mvmOp = mlir::dyn_cast<MVMOpTy>(op))
+      found |= mvmOp.getWeightIndex() == weightIndex;
+    else if (auto vmmOp = mlir::dyn_cast<VMMOpTy>(op))
+      found |= vmmOp.getWeightIndex() == weightIndex;
+  });
+  return found;
+}
+
+template <typename MVMOpTy, typename VMMOpTy, typename ParentOpTy>
+void walkMvmVmmWeightUses(ParentOpTy parentOp, llvm::function_ref<void(mlir::OpOperand&)> callback) {
+  auto weights = parentOp.getWeights();
+  llvm::SmallSet<unsigned, 8> visited;
+  auto walkWeightIndex = [&](unsigned weightIndex) {
+    if (weightIndex < weights.size() && visited.insert(weightIndex).second)
+      callback(parentOp->getOpOperand(weightIndex));
+  };
+
+  parentOp.walk([&](MVMOpTy op) { walkWeightIndex(op.getWeightIndex()); });
+  parentOp.walk([&](VMMOpTy op) { walkWeightIndex(op.getWeightIndex()); });
+}
+
+} // namespace
+
+bool isSpatialMvmVmmWeightUse(mlir::OpOperand& use) {
+  mlir::Operation* user = use.getOwner();
+  unsigned operandIndex = use.getOperandNumber();
+
+  auto computeOp = mlir::dyn_cast<spatial::SpatCompute>(user);
+  if (!computeOp || operandIndex >= computeOp.getWeights().size())
+    return false;
+
+  return hasMvmVmmWeightUse<spatial::SpatMVMOp, spatial::SpatVMMOp>(computeOp, operandIndex);
+}
+
+bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value) {
+  llvm::SmallPtrSet<mlir::Value, 8> visited;
+  auto walkUses = [&](mlir::Value currentValue, auto& self) -> bool {
+    if (!visited.insert(currentValue).second)
+      return true;
+    if (currentValue.use_empty())
+      return false;
+
+    return llvm::all_of(currentValue.getUses(), [&](mlir::OpOperand& use) {
+      if (isSpatialMvmVmmWeightUse(use))
+        return true;
+
+      mlir::Operation* user = use.getOwner();
+      if (auto extractSliceOp = mlir::dyn_cast<mlir::tensor::ExtractSliceOp>(user))
+        return extractSliceOp.getSource() == currentValue && self(extractSliceOp.getResult(), self);
+      if (auto expandShapeOp = mlir::dyn_cast<mlir::tensor::ExpandShapeOp>(user))
+        return expandShapeOp.getSrc() == currentValue && self(expandShapeOp.getResult(), self);
+      if (auto collapseShapeOp = mlir::dyn_cast<mlir::tensor::CollapseShapeOp>(user))
+        return collapseShapeOp.getSrc() == currentValue && self(collapseShapeOp.getResult(), self);
+      if (auto transposeOp = mlir::dyn_cast<mlir::ONNXTransposeOp>(user))
+        return transposeOp.getData() == currentValue && self(transposeOp.getResult(), self);
+
+      return false;
+    });
+  };
+
+  return walkUses(value, walkUses);
+}
+
+void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir::OpOperand&)> callback) {
+  assert(root && "expected valid root op");
+  root->walk([&](pim::PimCoreOp coreOp) { walkMvmVmmWeightUses<pim::PimMVMOp, pim::PimVMMOp>(coreOp, callback); });
+  root->walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    auto weights = coreBatchOp.getWeights();
+    for (auto weight : weights)
+      for (mlir::OpOperand& use : weight.getUses())
+        if (use.getOwner() == coreBatchOp.getOperation())
+          callback(use);
+  });
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/WeightUtils.hpp

@@ -0,0 +1,29 @@
+#pragma once
+
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/StringRef.h"
+
+inline constexpr llvm::StringRef PimWeightAlwaysAttrName = "weightAlways";
+
+namespace onnx_mlir {
+
+bool hasWeightAlways(mlir::Operation* op);
+
+/// Tags an op as producing a value that should stay materialized as a reusable
+/// weight across later PIM lowering/codegen stages.
+void markWeightAlways(mlir::Operation* op);
+
+bool isSpatialMvmVmmWeightUse(mlir::OpOperand& use);
+
+/// Returns true when a value flows only into Spatial weighted MVM/VMM operands,
+/// allowing later passes to preserve it as a dedicated weight-like object.
+bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value);
+
+/// Visits weight operands consumed by Pim core ops/core batches so downstream
+/// passes can identify globals that must remain weight-backed.
+void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir::OpOperand&)> callback);
+
+} // namespace onnx_mlir

src/PIM/Common/PimCommon.cpp

@@ -1,625 +0,0 @@
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/BuiltinTypeInterfaces.h"
-#include "mlir/Interfaces/DestinationStyleOpInterface.h"
-
-#include "llvm/ADT/SmallSet.h"
-#include "llvm/ADT/SmallPtrSet.h"
-#include "llvm/Support/raw_os_ostream.h"
-
-#include <filesystem>
-#include <fstream>
-
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Compiler/CompilerOptions.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-
-std::string getOutputDir() {
-  if (outputBaseName.empty() || outputBaseName == "-")
-    return {};
-
-  size_t lastSlash = outputBaseName.find_last_of('/');
-  if (lastSlash == std::string::npos)
-    return ".";
-  return outputBaseName.substr(0, lastSlash);
-}
-
-void createDirectory(const std::string& directory) {
-  std::error_code errorCode;
-  std::filesystem::create_directories(directory, errorCode);
-  assert(!errorCode && ("Failed to create directory: " + errorCode.message()).data());
-}
-
-void dumpModule(ModuleOp moduleOp, const std::string& name) {
-  std::string outputDir = getOutputDir();
-  if (outputDir.empty())
-    return;
-
-  std::string dialectsDir = outputDir + "/dialects";
-  createDirectory(dialectsDir);
-
-  std::fstream file(dialectsDir + "/" + name + ".mlir", std::ios::out);
-  llvm::raw_os_ostream os(file);
-  os << *moduleOp;
-  os.flush();
-  file.close();
-}
-
-FailureOr<func::FuncOp> getPimEntryFunc(ModuleOp moduleOp) {
-  if (!moduleOp)
-    return failure();
-
-  SmallVector<ONNXEntryPointOp> entryPoints(moduleOp.getOps<ONNXEntryPointOp>());
-  if (entryPoints.size() > 1) {
-    moduleOp.emitError("PIM pipeline requires a single ONNX entry point, but found ") << entryPoints.size();
-    return failure();
-  }
-  if (!entryPoints.empty()) {
-    auto entryPointAttr =
-      entryPoints.front()->getAttrOfType<SymbolRefAttr>(ONNXEntryPointOp::getEntryPointFuncAttrName());
-    if (!entryPointAttr) {
-      entryPoints.front().emitOpError("is missing the entry point function attribute");
-      return failure();
-    }
-    auto entryFunc = moduleOp.lookupSymbol<func::FuncOp>(entryPointAttr.getLeafReference().getValue());
-    if (!entryFunc) {
-      entryPoints.front().emitOpError("references an unknown entry function ")
-        << entryPointAttr.getLeafReference().getValue();
-      return failure();
-    }
-    return entryFunc;
-  }
-
-  if (auto mainGraphFunc = moduleOp.lookupSymbol<func::FuncOp>("main_graph"))
-    return mainGraphFunc;
-
-  SmallVector<func::FuncOp> nonExternalFuncs;
-  for (auto funcOp : moduleOp.getOps<func::FuncOp>())
-    if (!funcOp.isExternal())
-      nonExternalFuncs.push_back(funcOp);
-  if (nonExternalFuncs.size() == 1)
-    return nonExternalFuncs.front();
-
-  moduleOp.emitError("could not resolve a unique PIM entry function");
-  return failure();
-}
-
-bool hasWeightAlways(Operation* op) { return op && op->getAttr(PimWeightAlwaysAttrName) != nullptr; }
-
-void markWeightAlways(Operation* op) {
-  assert(op && "expected valid op");
-  op->setAttr(PimWeightAlwaysAttrName, UnitAttr::get(op->getContext()));
-}
-
-namespace {
-
-template <typename MVMOpTy, typename VMMOpTy, typename ParentOpTy>
-bool hasMvmVmmWeightUse(ParentOpTy parentOp, unsigned weightIndex) {
-  bool found = false;
-  parentOp.walk([&](Operation* op) {
-    if (auto mvmOp = dyn_cast<MVMOpTy>(op))
-      found |= mvmOp.getWeightIndex() == weightIndex;
-    else if (auto vmmOp = dyn_cast<VMMOpTy>(op))
-      found |= vmmOp.getWeightIndex() == weightIndex;
-  });
-  return found;
-}
-
-template <typename MVMOpTy, typename VMMOpTy, typename ParentOpTy>
-void walkMvmVmmWeightUses(ParentOpTy parentOp, function_ref<void(OpOperand&)> callback) {
-  auto weights = parentOp.getWeights();
-  llvm::SmallSet<unsigned, 8> visited;
-  auto walkWeightIndex = [&](unsigned weightIndex) {
-    if (weightIndex < weights.size() && visited.insert(weightIndex).second)
-      callback(parentOp->getOpOperand(weightIndex));
-  };
-
-  parentOp.walk([&](MVMOpTy op) { walkWeightIndex(op.getWeightIndex()); });
-  parentOp.walk([&](VMMOpTy op) { walkWeightIndex(op.getWeightIndex()); });
-}
-
-} // namespace
-
-bool isSpatialMvmVmmWeightUse(OpOperand& use) {
-  Operation* user = use.getOwner();
-  unsigned operandIndex = use.getOperandNumber();
-
-  auto computeOp = dyn_cast<spatial::SpatCompute>(user);
-  if (!computeOp || operandIndex >= computeOp.getWeights().size())
-    return false;
-
-  return hasMvmVmmWeightUse<spatial::SpatWeightedMVMOp, spatial::SpatWeightedVMMOp>(computeOp, operandIndex);
-}
-
-bool hasOnlySpatialMvmVmmWeightUses(Value value) {
-  SmallPtrSet<Value, 8> visited;
-  auto walkUses = [&](Value currentValue, auto& self) -> bool {
-    if (!visited.insert(currentValue).second)
-      return true;
-    if (currentValue.use_empty())
-      return false;
-
-    return llvm::all_of(currentValue.getUses(), [&](OpOperand& use) {
-      if (isSpatialMvmVmmWeightUse(use))
-        return true;
-
-      Operation* user = use.getOwner();
-      if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(user))
-        return extractSliceOp.getSource() == currentValue && self(extractSliceOp.getResult(), self);
-      if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(user))
-        return expandShapeOp.getSrc() == currentValue && self(expandShapeOp.getResult(), self);
-      if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(user))
-        return collapseShapeOp.getSrc() == currentValue && self(collapseShapeOp.getResult(), self);
-      if (auto transposeOp = dyn_cast<ONNXTransposeOp>(user))
-        return transposeOp.getData() == currentValue && self(transposeOp.getResult(), self);
-
-      return false;
-    });
-  };
-
-  return walkUses(value, walkUses);
-}
-
-void walkPimMvmVmmWeightUses(Operation* root, function_ref<void(OpOperand&)> callback) {
-  assert(root && "expected valid root op");
-  root->walk([&](pim::PimCoreOp coreOp) {
-    walkMvmVmmWeightUses<pim::PimMVMOp, pim::PimVMMOp>(coreOp, callback);
-  });
-}
-
-memref::GlobalOp lookupGlobalForGetGlobal(ModuleOp moduleOp, memref::GetGlobalOp getGlobalOp) {
-  if (!moduleOp || !getGlobalOp)
-    return {};
-  return moduleOp.lookupSymbol<memref::GlobalOp>(getGlobalOp.getName());
-}
-
-FailureOr<Operation*> getOtherEndOfChannel(Operation* op, bool opIsReceive, RewriterBase& rewriter) {
-
-  auto channelNewOp = op->getOperand(0).getDefiningOp<spatial::SpatChannelNewOp>();
-  if (!channelNewOp) {
-    op->emitError("User of Channel must have the first operand created by ChannelNewOp.");
-    return failure();
-  }
-  // channelNewOp should have two users: `op` and a
-  // `ChannelSendOp`/`ChannelReceiveOp`
-  auto channelUsers = channelNewOp->getUsers();
-  auto usersIterator = channelUsers.begin();
-  auto firstUser = *usersIterator;
-  usersIterator++;
-  if (usersIterator == channelUsers.end()) {
-    op->emitError("Operand generated by ChannelNewOp must have two users, "
-                  "only one found.");
-    channelNewOp->dump();
-    op->dump();
-    channelNewOp->getParentOp()->dump();
-    return failure();
-  }
-  auto secondUser = *usersIterator;
-  usersIterator++;
-  if (usersIterator != channelUsers.end()) {
-    op->emitError("Operand generated by ChannelNewOp must have two users, "
-                  "more than two found.");
-    return failure();
-  }
-  Operation* notOpUser;
-  if (firstUser == op) {
-    notOpUser = secondUser;
-  }
-  else if (secondUser == op) {
-    notOpUser = firstUser;
-  }
-  else {
-    op->emitError("Operand generated by ChannelNewOp must have two users, "
-                  "and one of them must be me, but"
-                  "none of them is actually me.");
-    return failure();
-  }
-
-  if (opIsReceive) {
-    if (!isa<spatial::SpatChannelSendOp>(notOpUser)) {
-      op->emitError("Operand generated by ChannelNewOp has two user, one is "
-                    "me, the other is not a ChannelSendOp.");
-      return failure();
-    }
-    return notOpUser;
-  }
-  else {
-    if (!isa<spatial::SpatChannelReceiveOp>(notOpUser)) {
-      op->emitError("Operand generated by ChannelNewOp has two user, one is "
-                    "me, the other is not a ChannelReceiveOp.");
-      return failure();
-    }
-    return notOpUser;
-  }
-}
-
-SmallVector<int64_t> computeRowMajorStrides(ArrayRef<int64_t> shape) {
-  SmallVector<int64_t> strides(shape.size(), 1);
-  for (int64_t dim = static_cast<int64_t>(shape.size()) - 2; dim >= 0; --dim)
-    strides[dim] = strides[dim + 1] * shape[dim + 1];
-  return strides;
-}
-
-SmallVector<int64_t> delinearizeIndex(int64_t linearIndex, ArrayRef<int64_t> shape, ArrayRef<int64_t> strides) {
-  SmallVector<int64_t> indices(shape.size(), 0);
-  for (auto [dim, stride] : llvm::enumerate(strides)) {
-    indices[dim] = linearIndex / stride;
-    linearIndex %= stride;
-  }
-  return indices;
-}
-
-int64_t linearizeIndex(ArrayRef<int64_t> indices, ArrayRef<int64_t> strides) {
-  int64_t linearIndex = 0;
-  for (auto [index, stride] : llvm::zip_equal(indices, strides))
-    linearIndex += index * stride;
-  return linearIndex;
-}
-
-int64_t getNumElements(ArrayRef<int64_t> shape) {
-  int64_t numElements = 1;
-  for (int64_t dim : shape)
-    numElements *= dim;
-  return numElements;
-}
-
-bool isMemoryContiguous(ArrayRef<int64_t> srcShape,
-                        ArrayRef<int64_t> offsets,
-                        ArrayRef<int64_t> sizes,
-                        ArrayRef<int64_t> strides) {
-  if (std::any_of(strides.begin(), strides.end(), [](int64_t stride) -> bool { return stride != 1; }))
-    return false;
-
-  auto offsetsAndSizesAndShape = llvm::zip_equal(llvm::make_range(offsets.rbegin(), offsets.rend()),
-                                                 llvm::make_range(sizes.rbegin(), sizes.rend()),
-                                                 llvm::make_range(srcShape.rbegin(), srcShape.rend()));
-
-  auto firstNonZeroOffset = std::find_if(
-    offsetsAndSizesAndShape.begin(), offsetsAndSizesAndShape.end(), [&](auto offsetAndSizeAndShape) -> bool {
-      auto [offset, _size, _dimension] = offsetAndSizeAndShape;
-      return offset != 0;
-    });
-
-  if (firstNonZeroOffset != offsetsAndSizesAndShape.end()) {
-    auto [offset, size, dimension] = *firstNonZeroOffset;
-    if (size > dimension - offset)
-      return false;
-    ++firstNonZeroOffset;
-
-    if (std::any_of(firstNonZeroOffset, offsetsAndSizesAndShape.end(), [](auto offsetAndSizeAndShape) -> bool {
-          auto [_offset, size, _dimension] = offsetAndSizeAndShape;
-          return size != 1;
-        }))
-      return false;
-  }
-
-  auto sizesAndShape = llvm::zip_equal(llvm::make_range(sizes.rbegin(), sizes.rend()),
-                                       llvm::make_range(srcShape.rbegin(), srcShape.rend()));
-
-  auto firstDifferentSize = std::find_if(sizesAndShape.begin(), sizesAndShape.end(), [&](auto sizeAndShape) -> bool {
-    auto [size, dimension] = sizeAndShape;
-    return size != dimension;
-  });
-
-  if (firstDifferentSize != sizesAndShape.end()) {
-    ++firstDifferentSize;
-
-    if (std::any_of(firstDifferentSize, sizesAndShape.end(), [](auto sizeAndShape) -> bool {
-          auto [size, _dimension] = sizeAndShape;
-          return size != 1;
-        }))
-      return false;
-  }
-
-  return true;
-}
-
-static Value resolveAlias(Value value, const StaticValueKnowledge* knowledge) {
-  if (!knowledge)
-    return value;
-
-  auto iter = knowledge->aliases.find(value);
-  while (iter != knowledge->aliases.end()) {
-    value = iter->second;
-    iter = knowledge->aliases.find(value);
-  }
-  return value;
-}
-
-// Walks through view-like ops and DPS tied operands to find the "underlying" memref value
-// behind an scf.for iter-arg. Used both when resolving a contiguous address inside a loop
-// and when propagating yielded values across iterations during static unrolling.
-static Value resolveLoopCarriedAliasImpl(Value value, const StaticValueKnowledge* knowledge) {
-  value = resolveAlias(value, knowledge);
-
-  if (auto blockArgument = dyn_cast<BlockArgument>(value))
-    return value;
-
-  Operation* definingOp = value.getDefiningOp();
-  if (!definingOp)
-    return value;
-
-  if (auto dpsDefiningOp = dyn_cast<DestinationStyleOpInterface>(definingOp)) {
-    if (auto result = dyn_cast<OpResult>(value))
-      if (OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(result))
-        return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge);
-  }
-
-  if (auto castOp = dyn_cast<memref::CastOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge);
-  if (auto collapseOp = dyn_cast<memref::CollapseShapeOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge);
-  if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge);
-
-  return value;
-}
-
-static FailureOr<int64_t> resolveOpFoldResult(OpFoldResult ofr, const StaticValueKnowledge* knowledge);
-
-static FailureOr<int64_t> resolveIndexValueImpl(Value value, const StaticValueKnowledge* knowledge) {
-  value = resolveAlias(value, knowledge);
-
-  if (knowledge) {
-    auto iter = knowledge->indexValues.find(value);
-    if (iter != knowledge->indexValues.end())
-      return iter->second;
-  }
-
-  auto constantOp = value.getDefiningOp<arith::ConstantOp>();
-  if (constantOp) {
-    if (auto integerAttr = dyn_cast<IntegerAttr>(constantOp.getValue()))
-      return integerAttr.getInt();
-  }
-
-  Operation* definingOp = value.getDefiningOp();
-  if (!definingOp)
-    return failure();
-
-  if (auto indexCastOp = dyn_cast<arith::IndexCastOp>(definingOp))
-    return resolveIndexValueImpl(indexCastOp.getIn(), knowledge);
-
-  if (auto addOp = dyn_cast<arith::AddIOp>(definingOp)) {
-    auto lhs = resolveIndexValueImpl(addOp.getLhs(), knowledge);
-    auto rhs = resolveIndexValueImpl(addOp.getRhs(), knowledge);
-    if (failed(lhs) || failed(rhs))
-      return failure();
-    return *lhs + *rhs;
-  }
-
-  if (auto subOp = dyn_cast<arith::SubIOp>(definingOp)) {
-    auto lhs = resolveIndexValueImpl(subOp.getLhs(), knowledge);
-    auto rhs = resolveIndexValueImpl(subOp.getRhs(), knowledge);
-    if (failed(lhs) || failed(rhs))
-      return failure();
-    return *lhs - *rhs;
-  }
-
-  if (auto mulOp = dyn_cast<arith::MulIOp>(definingOp)) {
-    auto lhs = resolveIndexValueImpl(mulOp.getLhs(), knowledge);
-    auto rhs = resolveIndexValueImpl(mulOp.getRhs(), knowledge);
-    if (failed(lhs) || failed(rhs))
-      return failure();
-    return *lhs * *rhs;
-  }
-
-  if (auto divOp = dyn_cast<arith::DivUIOp>(definingOp)) {
-    auto lhs = resolveIndexValueImpl(divOp.getLhs(), knowledge);
-    auto rhs = resolveIndexValueImpl(divOp.getRhs(), knowledge);
-    if (failed(lhs) || failed(rhs) || *rhs == 0)
-      return failure();
-    return static_cast<int64_t>(static_cast<uint64_t>(*lhs) / static_cast<uint64_t>(*rhs));
-  }
-
-  if (auto remOp = dyn_cast<arith::RemUIOp>(definingOp)) {
-    auto lhs = resolveIndexValueImpl(remOp.getLhs(), knowledge);
-    auto rhs = resolveIndexValueImpl(remOp.getRhs(), knowledge);
-    if (failed(lhs) || failed(rhs) || *rhs == 0)
-      return failure();
-    return static_cast<int64_t>(static_cast<uint64_t>(*lhs) % static_cast<uint64_t>(*rhs));
-  }
-
-  return failure();
-}
-
-static FailureOr<int64_t> resolveOpFoldResult(OpFoldResult ofr, const StaticValueKnowledge* knowledge) {
-  if (auto attr = dyn_cast<Attribute>(ofr)) {
-    auto integerAttr = dyn_cast<IntegerAttr>(attr);
-    if (!integerAttr)
-      return failure();
-    return integerAttr.getInt();
-  }
-
-  return resolveIndexValueImpl(cast<Value>(ofr), knowledge);
-}
-
-static FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(Value value,
-                                                                         const StaticValueKnowledge* knowledge) {
-  int64_t byteOffset = 0;
-  value = resolveAlias(value, knowledge);
-
-  while (true) {
-    if (isa<BlockArgument>(value))
-      return ResolvedContiguousAddress {value, byteOffset};
-
-    Operation* definingOp = value.getDefiningOp();
-    if (!definingOp)
-      return failure();
-
-    if (auto dpsDefiningOp = dyn_cast<DestinationStyleOpInterface>(definingOp)) {
-      OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(dyn_cast<OpResult>(value));
-      if (!tiedOperand)
-        return failure();
-      value = resolveAlias(tiedOperand->get(), knowledge);
-      continue;
-    }
-
-    if (auto forOp = dyn_cast<scf::ForOp>(definingOp)) {
-      auto result = dyn_cast<OpResult>(value);
-      if (!result)
-        return failure();
-
-      // Trace the loop carry back to its underlying memref, then if that memref is the
-      // loop's own iter-arg we know the base comes from the corresponding init arg
-      // (every iteration yields the same backing memory in the DPS sense).
-      auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
-      Value yieldedValue = resolveLoopCarriedAliasImpl(yieldOp.getOperand(result.getResultNumber()), knowledge);
-      if (auto blockArgument = dyn_cast<BlockArgument>(yieldedValue)) {
-        if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0
-            && static_cast<unsigned>(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size()) {
-          value = resolveAlias(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge);
-          continue;
-        }
-      }
-
-      value = yieldedValue;
-      continue;
-    }
-
-    if (auto subviewOp = dyn_cast<memref::SubViewOp>(definingOp)) {
-      auto sourceType = dyn_cast<MemRefType>(subviewOp.getSource().getType());
-      auto subviewType = dyn_cast<MemRefType>(subviewOp.getType());
-      if (!sourceType || !subviewType || !sourceType.hasStaticShape() || !subviewType.hasStaticShape())
-        return failure();
-
-      SmallVector<int64_t> offsets;
-      SmallVector<int64_t> sizes;
-      SmallVector<int64_t> strides;
-      offsets.reserve(subviewOp.getMixedOffsets().size());
-      sizes.reserve(subviewOp.getMixedSizes().size());
-      strides.reserve(subviewOp.getMixedStrides().size());
-
-      for (OpFoldResult offset : subviewOp.getMixedOffsets()) {
-        auto resolvedOffset = resolveOpFoldResult(offset, knowledge);
-        if (failed(resolvedOffset))
-          return failure();
-        offsets.push_back(*resolvedOffset);
-      }
-
-      for (OpFoldResult size : subviewOp.getMixedSizes()) {
-        auto resolvedSize = resolveOpFoldResult(size, knowledge);
-        if (failed(resolvedSize))
-          return failure();
-        sizes.push_back(*resolvedSize);
-      }
-
-      for (OpFoldResult stride : subviewOp.getMixedStrides()) {
-        auto resolvedStride = resolveOpFoldResult(stride, knowledge);
-        if (failed(resolvedStride))
-          return failure();
-        strides.push_back(*resolvedStride);
-      }
-
-      if (!isMemoryContiguous(sourceType.getShape(), offsets, sizes, strides))
-        return failure();
-
-      auto sourceStrides = computeRowMajorStrides(sourceType.getShape());
-      byteOffset += linearizeIndex(offsets, sourceStrides) * subviewType.getElementTypeBitWidth() / 8;
-      value = resolveAlias(subviewOp.getSource(), knowledge);
-      continue;
-    }
-
-    if (auto castOp = dyn_cast<memref::CastOp>(definingOp)) {
-      value = resolveAlias(castOp.getSource(), knowledge);
-      continue;
-    }
-    if (auto collapseOp = dyn_cast<memref::CollapseShapeOp>(definingOp)) {
-      value = resolveAlias(collapseOp.getSrc(), knowledge);
-      continue;
-    }
-    if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(definingOp)) {
-      value = resolveAlias(expandOp.getSrc(), knowledge);
-      continue;
-    }
-
-    if (isa<memref::AllocOp, memref::GetGlobalOp>(definingOp))
-      return ResolvedContiguousAddress {value, byteOffset};
-
-    return failure();
-  }
-}
-
-FailureOr<int64_t> resolveIndexValue(Value value) { return resolveIndexValueImpl(value, nullptr); }
-
-FailureOr<int64_t> resolveIndexValue(Value value, const StaticValueKnowledge& knowledge) {
-  return resolveIndexValueImpl(value, &knowledge);
-}
-
-FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(Value value) {
-  return resolveContiguousAddressImpl(value, nullptr);
-}
-
-FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(Value value, const StaticValueKnowledge& knowledge) {
-  return resolveContiguousAddressImpl(value, &knowledge);
-}
-
-Value resolveLoopCarriedAlias(Value value, const StaticValueKnowledge& knowledge) {
-  return resolveLoopCarriedAliasImpl(value, &knowledge);
-}
-
-bool isCoreStaticAddressOp(Operation* op) {
-  return isa<arith::ConstantOp,
-             arith::AddIOp,
-             arith::SubIOp,
-             arith::MulIOp,
-             arith::DivUIOp,
-             arith::RemUIOp,
-             arith::IndexCastOp,
-             memref::AllocOp,
-             memref::SubViewOp,
-             memref::CastOp,
-             memref::CollapseShapeOp,
-             memref::ExpandShapeOp>(op);
-}
-
-LogicalResult walkPimCoreBlock(Block& block,
-                               const StaticValueKnowledge& knowledge,
-                               llvm::function_ref<LogicalResult(Operation&, const StaticValueKnowledge&)> callback) {
-  bool hasFailure = false;
-  for (Operation& op : block) {
-    if (isa<pim::PimHaltOp, scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
-      continue;
-
-    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
-      Block& loopBody = forOp.getRegion().front();
-      auto lowerBound = resolveIndexValue(forOp.getLowerBound(), knowledge);
-      auto upperBound = resolveIndexValue(forOp.getUpperBound(), knowledge);
-      auto step = resolveIndexValue(forOp.getStep(), knowledge);
-      if (failed(lowerBound) || failed(upperBound) || failed(step) || *step <= 0) {
-        forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
-        hasFailure = true;
-        continue;
-      }
-
-      SmallVector<Value> iterValues(forOp.getInitArgs().begin(), forOp.getInitArgs().end());
-      for (int64_t inductionValue = *lowerBound; inductionValue < *upperBound; inductionValue += *step) {
-        StaticValueKnowledge loopKnowledge = knowledge;
-        loopKnowledge.indexValues[forOp.getInductionVar()] = inductionValue;
-        for (auto [iterArg, iterValue] : llvm::zip_equal(forOp.getRegionIterArgs(), iterValues))
-          loopKnowledge.aliases[iterArg] = iterValue;
-
-        if (failed(walkPimCoreBlock(loopBody, loopKnowledge, callback)))
-          hasFailure = true;
-
-        auto yieldOp = cast<scf::YieldOp>(loopBody.getTerminator());
-        for (auto [index, yieldedValue] : llvm::enumerate(yieldOp.getOperands()))
-          iterValues[index] = resolveLoopCarriedAlias(yieldedValue, loopKnowledge);
-      }
-      continue;
-    }
-
-    if (failed(callback(op, knowledge)))
-      hasFailure = true;
-  }
-  return success(!hasFailure);
-}
-
-} // namespace onnx_mlir

src/PIM/Common/PimCommon.hpp

@@ -11,84 +11,18 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/EntryPointUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
 #include "src/Compiler/CompilerOptions.hpp"
 
-inline constexpr llvm::StringRef PimWeightAlwaysAttrName = "weightAlways";
-
 namespace onnx_mlir {
 
-struct ResolvedContiguousAddress {
-  mlir::Value base;
-  int64_t byteOffset = 0;
-};
-
-struct StaticValueKnowledge {
-  llvm::DenseMap<mlir::Value, int64_t> indexValues;
-  llvm::DenseMap<mlir::Value, mlir::Value> aliases;
-
-  StaticValueKnowledge() {}
-};
-
-std::string getOutputDir();
-
-void createDirectory(const std::string& directory);
-
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name);
-
-llvm::FailureOr<mlir::func::FuncOp> getPimEntryFunc(mlir::ModuleOp moduleOp);
-
-bool hasWeightAlways(mlir::Operation* op);
-
-void markWeightAlways(mlir::Operation* op);
-
-bool isSpatialMvmVmmWeightUse(mlir::OpOperand& use);
-bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value);
-
-void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir::OpOperand&)> callback);
-
-mlir::memref::GlobalOp lookupGlobalForGetGlobal(mlir::ModuleOp moduleOp, mlir::memref::GetGlobalOp getGlobalOp);
-
-llvm::FailureOr<mlir::Operation*>
-getOtherEndOfChannel(mlir::Operation* op, bool opIsReceive, mlir::RewriterBase& rewriter);
-
-llvm::SmallVector<int64_t> computeRowMajorStrides(llvm::ArrayRef<int64_t> shape);
-
-llvm::SmallVector<int64_t>
-delinearizeIndex(int64_t linearIndex, llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> strides);
-
-int64_t linearizeIndex(llvm::ArrayRef<int64_t> indices, llvm::ArrayRef<int64_t> strides);
-
-int64_t getNumElements(llvm::ArrayRef<int64_t> shape);
-
-bool isMemoryContiguous(llvm::ArrayRef<int64_t> srcShape,
-                        llvm::ArrayRef<int64_t> offsets,
-                        llvm::ArrayRef<int64_t> sizes,
-                        llvm::ArrayRef<int64_t> strides);
-
-llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value);
-llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value,
-                                                                    const StaticValueKnowledge& knowledge);
-
-llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value);
-llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge);
-
-/// Follows alias and view/DPS chains using `knowledge` to find the value an scf.for
-/// iter-arg is ultimately backed by. Used when interpreting scf.for loop carries.
-mlir::Value resolveLoopCarriedAlias(mlir::Value value, const StaticValueKnowledge& knowledge);
-
-/// Returns true for ops inside a pim.core body that do not emit any PIM instruction and
-/// only contribute to static addressing or index computations (arith integer math,
-/// memref view ops, memref.alloc, arith.constant).
-bool isCoreStaticAddressOp(mlir::Operation* op);
-
-/// Walks `block` (the body of a pim.core region or an scf.for nested in it), statically
-/// unrolling any scf.for with resolvable bounds using `knowledge`. For each remaining op
-/// that is not skipped (pim.halt, scf.yield, or isCoreStaticAddressOp), `callback` is
-/// invoked with the op and the in-scope knowledge. The walker keeps going after a callback
-/// failure so callers can collect multiple diagnostics, but propagates the overall result.
-mlir::LogicalResult
-walkPimCoreBlock(mlir::Block& block,
-                 const StaticValueKnowledge& knowledge,
-                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback);
+inline constexpr llvm::StringLiteral kCoreIdAttrName = "coreId";
+inline constexpr llvm::StringLiteral kCoreIdsAttrName = "coreIds";
 
 } // namespace onnx_mlir

src/PIM/Common/Support/DebugDump.cpp

@@ -0,0 +1,27 @@
+#include "llvm/Support/raw_os_ostream.h"
+
+#include <fstream>
+
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
+
+namespace onnx_mlir {
+
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
+  std::string outputDir = getOutputDir();
+  if (outputDir.empty())
+    return;
+
+  std::string dialectsDir = outputDir + "/dialects";
+  createDirectory(dialectsDir);
+
+  std::fstream file(dialectsDir + "/" + name + ".mlir", std::ios::out);
+  llvm::raw_os_ostream os(file);
+  mlir::OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs();
+  moduleOp.print(os, flags);
+  os.flush();
+  file.close();
+}
+
+} // namespace onnx_mlir

src/PIM/Common/Support/DebugDump.hpp

@@ -0,0 +1,13 @@
+#pragma once
+
+#include "mlir/IR/BuiltinOps.h"
+
+#include <string>
+
+namespace onnx_mlir {
+
+/// Emits a MLIR snapshot under the current compiler output
+/// directory for pass-level debugging.
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name);
+
+} // namespace onnx_mlir

src/PIM/Common/Support/Diagnostics.cpp

@@ -0,0 +1,41 @@
+#include "llvm/ADT/STLExtras.h"
+
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
+
+namespace onnx_mlir::pim {
+
+mlir::InFlightDiagnostic emitUnsupportedStaticShapeDiagnostic(mlir::Operation* op, llvm::StringRef valueDescription) {
+  return op->emitOpError() << "requires statically shaped " << valueDescription;
+}
+
+mlir::InFlightDiagnostic emitUnsupportedRankDiagnostic(mlir::Operation* op,
+                                                       llvm::StringRef valueDescription,
+                                                       int64_t actualRank,
+                                                       llvm::ArrayRef<int64_t> supportedRanks) {
+  auto diag = op->emitOpError() << "has unsupported rank " << actualRank << " for " << valueDescription;
+  if (supportedRanks.empty())
+    return diag;
+
+  diag << "; supported rank";
+  if (supportedRanks.size() != 1)
+    diag << 's';
+  diag << ' ';
+
+  llvm::interleaveComma(supportedRanks, diag, [&](int64_t rank) { diag << rank; });
+  return diag;
+}
+
+mlir::InFlightDiagnostic
+emitMissingSymbolDiagnostic(mlir::Operation* op, llvm::StringRef symbolKind, llvm::StringRef symbolName) {
+  return op->emitOpError() << "references missing " << symbolKind << " `" << symbolName << "`";
+}
+
+mlir::LogicalResult emitFileSystemError(mlir::Location loc,
+                                        llvm::StringRef action,
+                                        llvm::StringRef path,
+                                        const std::error_code& errorCode) {
+  mlir::emitError(loc) << "failed to " << action << " `" << path << "`: " << errorCode.message();
+  return mlir::failure();
+}
+
+} // namespace onnx_mlir::pim

src/PIM/Common/Support/Diagnostics.hpp

@@ -0,0 +1,38 @@
+#pragma once
+
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+
+#include <system_error>
+
+namespace onnx_mlir::pim {
+
+/// Emits a consistent diagnostic for target paths that require static shapes.
+mlir::InFlightDiagnostic emitUnsupportedStaticShapeDiagnostic(mlir::Operation* op, llvm::StringRef valueDescription);
+
+/// Emits a consistent diagnostic for unsupported ranks while listing the ranks
+/// accepted by the current lowering/codegen path.
+mlir::InFlightDiagnostic emitUnsupportedRankDiagnostic(mlir::Operation* op,
+                                                       llvm::StringRef valueDescription,
+                                                       int64_t actualRank,
+                                                       llvm::ArrayRef<int64_t> supportedRanks);
+
+/// Emits a consistent diagnostic for missing symbol/global references.
+mlir::InFlightDiagnostic
+emitMissingSymbolDiagnostic(mlir::Operation* op, llvm::StringRef symbolKind, llvm::StringRef symbolName);
+
+/// Converts a filesystem error into an MLIR failure diagnostic anchored at
+/// the relevant IR location.
+mlir::LogicalResult
+emitFileSystemError(mlir::Location loc, llvm::StringRef action, llvm::StringRef path, const std::error_code& errorCode);
+
+template <typename T>
+mlir::LogicalResult failureOrToLogicalResult(const llvm::FailureOr<T>& value) {
+  return mlir::success(succeeded(value));
+}
+
+} // namespace onnx_mlir::pim

src/PIM/Common/Support/FileSystemUtils.cpp

@@ -0,0 +1,24 @@
+#include <filesystem>
+
+#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
+#include "src/Compiler/CompilerOptions.hpp"
+
+namespace onnx_mlir {
+
+std::string getOutputDir() {
+  if (outputBaseName.empty() || outputBaseName == "-")
+    return {};
+
+  size_t lastSlash = outputBaseName.find_last_of('/');
+  if (lastSlash == std::string::npos)
+    return ".";
+  return outputBaseName.substr(0, lastSlash);
+}
+
+void createDirectory(const std::string& directory) {
+  std::error_code errorCode;
+  std::filesystem::create_directories(directory, errorCode);
+  assert(!errorCode && ("Failed to create directory: " + errorCode.message()).data());
+}
+
+} // namespace onnx_mlir

src/PIM/Common/Support/FileSystemUtils.hpp

@@ -0,0 +1,13 @@
+#pragma once
+
+#include <string>
+
+namespace onnx_mlir {
+
+/// Returns the directory that should hold PIM artifacts/debug dumps for the
+/// current compiler invocation.
+std::string getOutputDir();
+
+void createDirectory(const std::string& directory);
+
+} // namespace onnx_mlir

src/PIM/Compiler/PimCodeGen.cpp

@@ -1,23 +1,34 @@
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/Verifier.h"
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Format.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <absl/types/compare.h>
 #include <algorithm>
 #include <cassert>
 #include <cmath>
+#include <cstdint>
+#include <fstream>
 #include <string>
 #include <utility>
 
 #include "Common/PimCommon.hpp"
-#include "Conversion/ONNXToSpatial/Common.hpp"
+#include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
@@ -26,6 +37,11 @@ using namespace llvm;
 using namespace mlir;
 using namespace onnx_mlir;
 
+static size_t getValueSizeInBytes(mlir::Value value) {
+  auto type = cast<ShapedType>(value.getType());
+  return type.getNumElements() * type.getElementTypeBitWidth() / 8;
+}
+
 MemEntry* PimMemory::gatherMemEntry(mlir::Value value) {
   auto type = cast<ShapedType>(value.getType());
   assert("Only static shape is supported" && type.hasStaticShape());
@@ -53,9 +69,22 @@ void PimMemory::allocateMemoryForValue(mlir::Value value, MemEntry& memEntry) {
 void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
   SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
   SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
+  SmallVector<mlir::Value> args;
+
+  for (mlir::Value arg : funcOp.getArguments()) {
+    gatherMemEntry(arg);
+    args.push_back(arg);
+  }
+
   funcOp.walk([&](memref::GetGlobalOp getGlobalOp) {
     if (!hasWeightAlways(getGlobalOp)) {
       auto globalMemrefOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
+      if (globalMemrefOp.getName().starts_with("arg")) {
+        StringRef indexStr = globalMemrefOp.getName().substr(4);
+        int index = 0;
+        llvm::to_integer(indexStr, index, 10);
+        globalAliases.push_back({getGlobalOp.getResult(), args[index]});
+      }
       auto [iter, inserted] = globalConstants.try_emplace(globalMemrefOp, getGlobalOp.getResult());
       if (inserted)
         gatherMemEntry(getGlobalOp.getResult());
@@ -64,13 +93,15 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
     }
   });
 
-  for (mlir::Value arg : funcOp.getArguments())
-    gatherMemEntry(arg);
-
   funcOp.walk([&](memref::AllocOp allocOp) {
     if (!allocOp->getParentOfType<pim::PimCoreOp>())
       gatherMemEntry(allocOp.getResult());
   });
+  funcOp.walk([&](pim::PimEmptyManyOp emptyManyOp) {
+    if (!emptyManyOp->getParentOfType<pim::PimCoreOp>() && !emptyManyOp->getParentOfType<pim::PimCoreBatchOp>())
+      for (mlir::Value output : emptyManyOp.getOutputs())
+        gatherMemEntry(output);
+  });
 
   allocateGatheredMemory();
 
@@ -80,10 +111,78 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
 
 void PimMemory::allocateCore(Operation* op) {
   op->walk([&](memref::AllocOp allocOp) { gatherMemEntry(allocOp); });
+  op->walk([&](pim::PimEmptyManyOp emptyManyOp) {
+    for (mlir::Value output : emptyManyOp.getOutputs())
+      gatherMemEntry(output);
+  });
 
   allocateGatheredMemory();
 }
 
+std::string formatMemory(uint64_t bytes) {
+  const char* units[] = {"B", "KB", "MB", "GB", "TB", "PB", "EB"};
+  int i = 0;
+  double size = static_cast<double>(bytes);
+  while (size >= 1024 && i < 6) {
+    size /= 1024;
+    i++;
+  }
+  // Formats to 2 decimal places
+  std::string out;
+  llvm::raw_string_ostream rss(out);
+  rss << llvm::format("%.2f ", size) << units[i];
+  return rss.str();
+}
+
+void PimMemory::report(llvm::raw_ostream& file) {
+  // Key: {OpName, Size}, Value: Vector of Addresses
+  // This groups all "memref.alloc" of "1KB" together
+  std::vector orderedList(globalMemEntriesMap.begin(), globalMemEntriesMap.end());
+  std::sort(
+    orderedList.begin(), orderedList.end(), [](auto lft, auto rgt) { return lft.second.address < rgt.second.address; });
+  auto newEnd = std::unique(orderedList.begin(), orderedList.end(), [](auto lft, auto rgt) {
+    return lft.second.address == rgt.second.address;
+  });
+  orderedList.erase(newEnd, orderedList.end());
+  std::sort(
+    orderedList.begin(), orderedList.end(), [](auto lft, auto rgt) { return lft.second.size < rgt.second.size; });
+  std::map<std::pair<std::string, uint64_t>, std::vector<uint64_t>> groupedStats;
+
+  for (auto& [value, memEntry] : orderedList) {
+    std::string opName = "Unknown/BlockArg";
+    if (auto op = value.getDefiningOp())
+      opName = op->getName().getStringRef().str();
+
+    groupedStats[{opName, memEntry.size}].push_back(memEntry.address);
+  }
+
+  file << "--- Memory Usage Report ---\n";
+
+  uint64_t totalMemory = 0;
+  for (auto const& [key, addresses] : groupedStats) {
+    const std::string& opName = key.first;
+    uint64_t size = key.second;
+
+    file.indent(4) << "Type: " << opName << " [" << formatMemory(size) << "]\n";
+    file.indent(6) << "Count:   " << addresses.size() << "\n";
+    file.indent(6) << "Total Memory:   " << formatMemory(size * addresses.size()) << "\n";
+    totalMemory += size * addresses.size();
+
+    // Optional: Print address range or first/last address to keep it concise
+    if (!addresses.empty()) {
+      auto [min, max] = std::minmax_element(addresses.begin(), addresses.end());
+      file.indent(6) << "Range:   " << llvm::format_hex(*min, 10) << " -> " << llvm::format_hex(*max, 10) << "\n";
+    }
+    file << "\n";
+    file << "Total Core Memory:  " << formatMemory(totalMemory) << "\n";
+  }
+}
+
+void PimMemory::remove(mlir::Value val) {
+  if (auto removeIter = globalMemEntriesMap.find(val); removeIter != globalMemEntriesMap.end())
+    globalMemEntriesMap.erase(removeIter);
+}
+
 MemEntry PimMemory::getMemEntry(mlir::Value value) const {
   auto iter = globalMemEntriesMap.find(value);
   assert("Missing memEntry for value" && iter != globalMemEntriesMap.end());
@@ -124,6 +223,28 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value, const StaticValu
   return iter->second.address + resolvedAddress->byteOffset;
 }
 
+void PimAcceleratorMemory::reportHost() {
+  llvm::raw_os_ostream os(fileReport);
+  os << "Host Memory\n";
+  hostMem.report(os);
+  os.flush();
+}
+
+void PimAcceleratorMemory::reportCore(size_t coreId) {
+  llvm::raw_os_ostream os(fileReport);
+  os << "Core " << coreId << " Memory\n";
+  deviceMem.at(coreId).report(os);
+  os.flush();
+}
+
+void PimAcceleratorMemory::clean(mlir::Operation* op) {
+  for (auto value : op->getResults()) {
+    hostMem.remove(value);
+    for (auto& device : deviceMem)
+      device.second.remove(value);
+  }
+}
+
 json::Object PimCodeGen::createEmptyOffset() {
   json::Object offset;
   offset["offset_select"] = 0;
@@ -131,6 +252,12 @@ json::Object PimCodeGen::createEmptyOffset() {
   return offset;
 }
 
+size_t PimCodeGen::remapCoreId(size_t coreId) const {
+  auto it = emittedCoreIds.find(coreId);
+  assert(it != emittedCoreIds.end() && "Missing emitted core id remapping");
+  return it->second;
+}
+
 static json::Object createRs1OnlyOffset() {
   json::Object offset;
   offset["offset_select"] = 1;
@@ -190,7 +317,7 @@ void PimCodeGen::emitCommunicationOp(StringRef opName, size_t bufferAddr, size_t
   json::Object json;
   json["op"] = opName;
   json["rd"] = 0;
-  json["core"] = coreId;
+  json["core"] = remapCoreId(coreId);
   json["size"] = size;
   json["offset"] = createEmptyOffset();
   emitInstruction(std::move(json));
@@ -242,10 +369,88 @@ void PimCodeGen::codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValue
     "recv", addressOf(receiveOp.getOutputBuffer(), knowledge), receiveOp.getSourceCoreId(), receiveOp.getSize());
 }
 
+void PimCodeGen::codeGenReceiveManyOp(pim::PimReceiveManyOp receiveManyOp,
+                                      const StaticValueKnowledge& knowledge) const {
+  for (auto [outputBuffer, sourceCoreId] :
+       llvm::zip(receiveManyOp.getOutputBuffers(), receiveManyOp.getSourceCoreIds()))
+    emitCommunicationOp("recv", addressOf(outputBuffer, knowledge), sourceCoreId, getValueSizeInBytes(outputBuffer));
+}
+
+void PimCodeGen::codeGenReceiveTensorOp(pim::PimReceiveTensorOp receiveTensorOp,
+                                        const StaticValueKnowledge& knowledge) const {
+  size_t outputAddr = addressOf(receiveTensorOp.getOutputBuffer(), knowledge);
+  size_t chunkSize = getValueSizeInBytes(receiveTensorOp.getOutputBuffer()) / receiveTensorOp.getSourceCoreIds().size();
+  for (auto [chunkIndex, sourceCoreId] : llvm::enumerate(receiveTensorOp.getSourceCoreIds()))
+    emitCommunicationOp("recv", outputAddr + chunkIndex * chunkSize, sourceCoreId, chunkSize);
+}
+
 void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const {
   emitCommunicationOp("send", addressOf(sendOp.getInput(), knowledge), sendOp.getTargetCoreId(), sendOp.getSize());
 }
 
+void PimCodeGen::codeGenSendManyOp(pim::PimSendManyOp sendManyOp, const StaticValueKnowledge& knowledge) const {
+  for (auto [input, targetCoreId] : llvm::zip(sendManyOp.getInputs(), sendManyOp.getTargetCoreIds()))
+    emitCommunicationOp("send", addressOf(input, knowledge), targetCoreId, getValueSizeInBytes(input));
+}
+
+void PimCodeGen::codeGenSendTensorOp(pim::PimSendTensorOp sendTensorOp, const StaticValueKnowledge& knowledge) const {
+  size_t inputAddr = addressOf(sendTensorOp.getInput(), knowledge);
+  size_t chunkSize = getValueSizeInBytes(sendTensorOp.getInput()) / sendTensorOp.getTargetCoreIds().size();
+  for (auto [chunkIndex, targetCoreId] : llvm::enumerate(sendTensorOp.getTargetCoreIds()))
+    emitCommunicationOp("send", inputAddr + chunkIndex * chunkSize, targetCoreId, chunkSize);
+}
+
+void PimCodeGen::codeGenExtractRowsOp(pim::PimExtractRowsOp extractRowsOp,
+                                      const StaticValueKnowledge& knowledge) const {
+  auto inputType = cast<ShapedType>(extractRowsOp.getInput().getType());
+  assert(inputType.hasStaticShape() && inputType.getRank() == 2 && "extract_rows codegen requires static rank-2 input");
+
+  size_t elementSize = inputType.getElementTypeBitWidth() / 8;
+  size_t rowSizeInBytes = static_cast<size_t>(inputType.getDimSize(1)) * elementSize;
+  size_t inputAddr = addressOf(extractRowsOp.getInput(), knowledge);
+
+  for (auto [rowIndex, outputBuffer] : llvm::enumerate(extractRowsOp.getOutputBuffers()))
+    emitMemCopyOp(
+      "lmv", addressOf(outputBuffer, knowledge), 0, inputAddr, rowIndex * rowSizeInBytes, rowSizeInBytes, "len");
+}
+
+void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const {
+  auto outputType = cast<ShapedType>(concatOp.getOutputBuffer().getType());
+  assert(outputType.hasStaticShape() && "concat codegen requires static output shape");
+
+  int64_t axis = concatOp.getAxis();
+  ArrayRef<int64_t> outputShape = outputType.getShape();
+  size_t elementSize = outputType.getElementTypeBitWidth() / 8;
+  size_t outputAddr = addressOf(concatOp.getOutputBuffer(), knowledge);
+
+  size_t outerCount = 1;
+  for (int64_t dim = 0; dim < axis; ++dim)
+    outerCount *= static_cast<size_t>(outputShape[dim]);
+
+  size_t innerCount = 1;
+  for (size_t dim = static_cast<size_t>(axis) + 1; dim < outputShape.size(); ++dim)
+    innerCount *= static_cast<size_t>(outputShape[dim]);
+
+  size_t outputConcatDim = static_cast<size_t>(outputShape[axis]);
+  size_t concatOffset = 0;
+  for (mlir::Value input : concatOp.getInputs()) {
+    auto inputType = cast<ShapedType>(input.getType());
+    assert(inputType.hasStaticShape() && "concat codegen requires static input shapes");
+
+    size_t inputConcatDim = static_cast<size_t>(inputType.getDimSize(axis));
+    size_t blockSizeInBytes = inputConcatDim * innerCount * elementSize;
+    size_t inputAddr = addressOf(input, knowledge);
+
+    for (size_t outerIndex = 0; outerIndex < outerCount; ++outerIndex) {
+      size_t dstOffset = (outerIndex * outputConcatDim + concatOffset) * innerCount * elementSize;
+      size_t srcOffset = outerIndex * inputConcatDim * innerCount * elementSize;
+      emitMemCopyOp("lmv", outputAddr, dstOffset, inputAddr, srcOffset, blockSizeInBytes, "len");
+    }
+
+    concatOffset += inputConcatDim;
+  }
+}
+
 template <typename MVMTy>
 void PimCodeGen::codeGenMVMLikeOp(size_t mvmId,
                                   MVMTy mvmLikeOp,
@@ -256,11 +461,6 @@ void PimCodeGen::codeGenMVMLikeOp(size_t mvmId,
   // TODO: save weights somewhere (if transposeMatrix=true, transpose the weight matrix)
 }
 
-static size_t getValueSizeInBytes(mlir::Value value) {
-  auto type = cast<ShapedType>(value.getType());
-  return type.getNumElements() * type.getElementTypeBitWidth() / 8;
-}
-
 void PimCodeGen::codeGenVVAddOp(pim::PimVVAddOp vvaddOp, const StaticValueKnowledge& knowledge) const {
   auto outputBufferAddr = addressOf(vvaddOp.getOutputBuffer(), knowledge);
   auto lhsAddr = addressOf(vvaddOp.getLhs(), knowledge);
@@ -412,6 +612,8 @@ void PimCodeGen::codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticVa
   emitInstruction(std::move(json));
 }
 
+void PimCodeGen::codeGetGlobalOp(memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const {}
+
 void PimCodeGen::codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const {
   auto srcAddr = addressOf(transposeOp.getInput(), knowledge);
   auto dstAddr = addressOf(transposeOp.getOutputBuffer(), knowledge);
@@ -474,19 +676,176 @@ std::string getMemorySizeAsString(size_t size) {
   return std::to_string(size) + " Bytes";
 }
 
-static SmallVector<unsigned, 8> getUsedWeightIndices(pim::PimCoreOp coreOp) {
+static SmallVector<unsigned, 8> getUsedWeightIndices(Block& block) {
   SmallVector<unsigned, 8> indices;
   auto addIndex = [&](unsigned weightIndex) {
     if (!llvm::is_contained(indices, weightIndex))
       indices.push_back(weightIndex);
   };
 
-  coreOp.walk([&](pim::PimMVMOp mvmOp) { addIndex(mvmOp.getWeightIndex()); });
-  coreOp.walk([&](pim::PimVMMOp vmmOp) { addIndex(vmmOp.getWeightIndex()); });
+  block.walk([&](pim::PimMVMOp mvmOp) { addIndex(mvmOp.getWeightIndex()); });
+  block.walk([&](pim::PimVMMOp vmmOp) { addIndex(vmmOp.getWeightIndex()); });
   llvm::sort(indices);
   return indices;
 }
 
+static SmallVector<unsigned, 8> getUsedWeightIndices(pim::PimCoreOp coreOp) {
+  return getUsedWeightIndices(coreOp.getBody().front());
+}
+
+static SmallVector<int32_t> getBatchCoreIds(pim::PimCoreBatchOp coreBatchOp) {
+  auto coreIdsAttr = coreBatchOp->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName);
+  assert(coreIdsAttr && "pim.core_batch requires coreIds array attribute");
+  return SmallVector<int32_t>(coreIdsAttr.asArrayRef().begin(), coreIdsAttr.asArrayRef().end());
+}
+
+static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
+  SmallVector<Operation*> coreLikeOps;
+  for (Operation& op : funcOp.getBody().front())
+    if (dyn_cast<pim::PimCoreOp>(&op) || dyn_cast<pim::PimCoreBatchOp>(&op))
+      coreLikeOps.push_back(&op);
+  return coreLikeOps;
+}
+
+static pim::PimCoreOp materializeScalarCoreFromBatchLane(pim::PimCoreBatchOp coreBatchOp, unsigned lane) {
+  OpBuilder builder(coreBatchOp);
+  builder.setInsertionPointAfter(coreBatchOp);
+
+  size_t laneCount = static_cast<size_t>(coreBatchOp.getLaneCount());
+  size_t weightsPerLane = coreBatchOp.getWeights().size() / laneCount;
+  SmallVector<mlir::Value> laneWeights;
+  laneWeights.reserve(weightsPerLane);
+  for (size_t weightIndex = 0; weightIndex < weightsPerLane; ++weightIndex)
+    laneWeights.push_back(coreBatchOp.getWeights()[lane * weightsPerLane + weightIndex]);
+
+  auto coreIds = getBatchCoreIds(coreBatchOp);
+  auto scalarCore = pim::PimCoreOp::create(
+    builder, coreBatchOp.getLoc(), ValueRange(laneWeights), builder.getI32IntegerAttr(coreIds[lane]));
+  Block* block = builder.createBlock(&scalarCore.getBody(), scalarCore.getBody().end());
+  IRMapping mapper;
+  if (coreBatchOp.getBody().front().getNumArguments() == 1)
+    mapper.map(coreBatchOp.getBody().front().getArgument(0), coreBatchOp.getInputs()[lane]);
+
+  builder.setInsertionPointToEnd(block);
+  for (Operation& op : coreBatchOp.getBody().front()) {
+    if (isa<pim::PimHaltOp>(op)) {
+      pim::PimHaltOp::create(builder, op.getLoc());
+      continue;
+    }
+
+    if (auto sendBatchOp = dyn_cast<pim::PimSendBatchOp>(op)) {
+      pim::PimSendOp::create(builder,
+                             sendBatchOp.getLoc(),
+                             mapper.lookup(sendBatchOp.getInput()),
+                             sendBatchOp.getSizeAttr(),
+                             builder.getI32IntegerAttr(sendBatchOp.getTargetCoreIds()[lane]));
+      continue;
+    }
+
+    if (auto sendManyBatchOp = dyn_cast<pim::PimSendManyBatchOp>(op)) {
+      SmallVector<int32_t> laneTargetCoreIds;
+      laneTargetCoreIds.reserve(sendManyBatchOp.getInputs().size());
+      for (auto valueIndex : llvm::seq<size_t>(0, sendManyBatchOp.getInputs().size()))
+        laneTargetCoreIds.push_back(
+          sendManyBatchOp.getTargetCoreIds()[valueIndex * laneCount + static_cast<size_t>(lane)]);
+
+      SmallVector<mlir::Value> mappedInputs;
+      mappedInputs.reserve(sendManyBatchOp.getInputs().size());
+      for (mlir::Value input : sendManyBatchOp.getInputs())
+        mappedInputs.push_back(mapper.lookup(input));
+
+      pim::PimSendManyOp::create(
+        builder, sendManyBatchOp.getLoc(), builder.getDenseI32ArrayAttr(laneTargetCoreIds), ValueRange(mappedInputs));
+      continue;
+    }
+
+    if (auto receiveBatchOp = dyn_cast<pim::PimReceiveBatchOp>(op)) {
+      auto scalarReceive =
+        pim::PimReceiveOp::create(builder,
+                                  receiveBatchOp.getLoc(),
+                                  receiveBatchOp.getOutput().getType(),
+                                  mapper.lookup(receiveBatchOp.getOutputBuffer()),
+                                  receiveBatchOp.getSizeAttr(),
+                                  builder.getI32IntegerAttr(receiveBatchOp.getSourceCoreIds()[lane]));
+      mapper.map(receiveBatchOp.getOutput(), scalarReceive.getOutput());
+      continue;
+    }
+
+    if (auto receiveManyBatchOp = dyn_cast<pim::PimReceiveManyBatchOp>(op)) {
+      SmallVector<int32_t> laneSourceCoreIds;
+      laneSourceCoreIds.reserve(receiveManyBatchOp.getOutputs().size());
+      for (auto valueIndex : llvm::seq<size_t>(0, receiveManyBatchOp.getOutputs().size()))
+        laneSourceCoreIds.push_back(
+          receiveManyBatchOp.getSourceCoreIds()[valueIndex * laneCount + static_cast<size_t>(lane)]);
+
+      SmallVector<mlir::Value> mappedOutputBuffers;
+      mappedOutputBuffers.reserve(receiveManyBatchOp.getOutputBuffers().size());
+      for (mlir::Value outputBuffer : receiveManyBatchOp.getOutputBuffers())
+        mappedOutputBuffers.push_back(mapper.lookup(outputBuffer));
+
+      auto scalarReceiveMany = pim::PimReceiveManyOp::create(builder,
+                                                             receiveManyBatchOp.getLoc(),
+                                                             receiveManyBatchOp->getResultTypes(),
+                                                             ValueRange(mappedOutputBuffers),
+                                                             builder.getDenseI32ArrayAttr(laneSourceCoreIds));
+      for (auto [originalOutput, scalarOutput] :
+           llvm::zip(receiveManyBatchOp.getOutputs(), scalarReceiveMany.getOutputs()))
+        mapper.map(originalOutput, scalarOutput);
+      continue;
+    }
+
+    if (auto memcpBatchOp = dyn_cast<pim::PimMemCopyHostToDevBatchOp>(op)) {
+      mlir::Value hostSource = mapper.lookupOrNull(memcpBatchOp.getHostSource());
+      if (!hostSource)
+        hostSource = memcpBatchOp.getHostSource();
+
+      auto scalarCopy = pim::PimMemCopyHostToDevOp::create(builder,
+                                                           memcpBatchOp.getLoc(),
+                                                           memcpBatchOp.getOutput().getType(),
+                                                           mapper.lookup(memcpBatchOp.getDeviceTarget()),
+                                                           hostSource,
+                                                           memcpBatchOp.getDeviceTargetOffsetAttr(),
+                                                           memcpBatchOp.getHostSourceOffsetAttr(),
+                                                           memcpBatchOp.getSizeAttr());
+      mapper.map(memcpBatchOp.getOutput(), scalarCopy.getOutput());
+      continue;
+    }
+
+    Operation* cloned = builder.clone(op, mapper);
+    for (auto [originalResult, clonedResult] : llvm::zip(op.getResults(), cloned->getResults()))
+      mapper.map(originalResult, clonedResult);
+  }
+
+  if (block->empty() || !isa<pim::PimHaltOp>(block->back()))
+    pim::PimHaltOp::create(builder, coreBatchOp.getLoc());
+  return scalarCore;
+}
+
+static void aliasMaterializedHostGlobals(ModuleOp moduleOp,
+                                         func::FuncOp funcOp,
+                                         pim::PimCoreOp coreOp,
+                                         PimAcceleratorMemory& memory) {
+  coreOp.walk([&](memref::GetGlobalOp getGlobalOp) {
+    if (hasWeightAlways(getGlobalOp) || memory.memEntriesMap.contains(getGlobalOp.getResult()))
+      return;
+
+    auto targetGlobal = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
+    if (!targetGlobal)
+      return;
+
+    mlir::Value aliasedValue;
+    funcOp.walk([&](memref::GetGlobalOp candidate) {
+      if (aliasedValue || candidate == getGlobalOp || !memory.memEntriesMap.contains(candidate.getResult()))
+        return;
+      if (lookupGlobalForGetGlobal(moduleOp, candidate) == targetGlobal)
+        aliasedValue = candidate.getResult();
+    });
+
+    if (aliasedValue)
+      memory.memEntriesMap[getGlobalOp.getResult()] = memory.memEntriesMap[aliasedValue];
+  });
+}
+
 /// Write global constant data into a binary memory image at their allocated addresses.
 static OnnxMlirCompilerErrorCodes
 writeMemoryBinary(ModuleOp moduleOp, func::FuncOp funcOp, PimAcceleratorMemory& memory, StringRef outputDirPath) {
@@ -553,8 +912,20 @@ static int64_t codeGenCoreOps(Block& block, PimCodeGen& coreCodeGen) {
         coreCodeGen.codeGenLmvOp(lmvOp, knowledge);
       else if (auto receiveOp = dyn_cast<pim::PimReceiveOp>(op))
         coreCodeGen.codeGenReceiveOp(receiveOp, knowledge);
+      else if (auto receiveManyOp = dyn_cast<pim::PimReceiveManyOp>(op))
+        coreCodeGen.codeGenReceiveManyOp(receiveManyOp, knowledge);
+      else if (auto receiveTensorOp = dyn_cast<pim::PimReceiveTensorOp>(op))
+        coreCodeGen.codeGenReceiveTensorOp(receiveTensorOp, knowledge);
       else if (auto sendOp = dyn_cast<pim::PimSendOp>(op))
         coreCodeGen.codeGenSendOp(sendOp, knowledge);
+      else if (auto sendManyOp = dyn_cast<pim::PimSendManyOp>(op))
+        coreCodeGen.codeGenSendManyOp(sendManyOp, knowledge);
+      else if (auto sendTensorOp = dyn_cast<pim::PimSendTensorOp>(op))
+        coreCodeGen.codeGenSendTensorOp(sendTensorOp, knowledge);
+      else if (auto extractRowsOp = dyn_cast<pim::PimExtractRowsOp>(op))
+        coreCodeGen.codeGenExtractRowsOp(extractRowsOp, knowledge);
+      else if (auto concatOp = dyn_cast<pim::PimConcatOp>(op))
+        coreCodeGen.codeGenConcatOp(concatOp, knowledge);
       else if (auto vmmOp = dyn_cast<pim::PimVMMOp>(op))
         coreCodeGen.codeGenMVMLikeOp<pim::PimVMMOp>(vmmOp.getWeightIndex(), vmmOp, true, knowledge);
       else if (auto mvmOp = dyn_cast<pim::PimMVMOp>(op))
@@ -581,6 +952,10 @@ static int64_t codeGenCoreOps(Block& block, PimCodeGen& coreCodeGen) {
         coreCodeGen.codeGenVSigmOp(vsigmOp, knowledge);
       else if (auto vsoftmaxOp = dyn_cast<pim::PimVSoftmaxOp>(op))
         coreCodeGen.codeGenVSoftmaxOp(vsoftmaxOp, knowledge);
+      else if (auto getGlobalOp = dyn_cast<memref::GetGlobalOp>(op))
+        coreCodeGen.codeGetGlobalOp(getGlobalOp, knowledge);
+      else if (isa<pim::PimEmptyManyOp>(op))
+        return success();
       else {
         op.emitError("Unsupported codegen for this operation");
         op.dump();
@@ -670,7 +1045,7 @@ static OnnxMlirCompilerErrorCodes writeCrossbarWeights(ModuleOp moduleOp,
   return CompilerSuccess;
 }
 
-llvm::DenseMap<pim::PimCoreOp, llvm::DenseMap<mlir::Value, std::string>>
+llvm::DenseMap<size_t, llvm::DenseMap<mlir::Value, std::string>>
 createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
   ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
   auto coreWeightsDirPath = outputDirPath + "/weights";
@@ -679,10 +1054,24 @@ createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
   size_t indexFileName = 0;
 
   int64_t xbarSize = crossbarSize.getValue();
-  llvm::DenseMap<pim::PimCoreOp, llvm::DenseMap<mlir::Value, std::string>> mapCoreWeightToFileName;
+  llvm::DenseMap<size_t, llvm::DenseMap<mlir::Value, std::string>> mapCoreWeightToFileName;
   llvm::DenseMap<memref::GlobalOp, std::string> mapGlobalOpToFileName;
 
-  for (pim::PimCoreOp coreOp : funcOp.getOps<pim::PimCoreOp>()) {
+  SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
+
+  for (Operation* op : coreLikeOps) {
+    SmallVector<pim::PimCoreOp> scalarCores;
+    if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
+      scalarCores.push_back(coreOp);
+    }
+    else {
+      auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
+      for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
+        scalarCores.push_back(materializeScalarCoreFromBatchLane(coreBatchOp, lane));
+    }
+
+    for (pim::PimCoreOp coreOp : scalarCores) {
+      size_t coreId = static_cast<size_t>(coreOp.getCoreId());
       for (unsigned index : getUsedWeightIndices(coreOp)) {
         if (index >= coreOp.getWeights().size()) {
           coreOp.emitWarning("Weight index " + std::to_string(index) + " is out of range");
@@ -717,7 +1106,7 @@ createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
         if (mapGlobalOpToFileName.contains(globalOp)) {
           auto& fileName = mapGlobalOpToFileName[globalOp];
           std::pair<mlir::Value, std::string> weightToFile = {weight, fileName};
-        mapCoreWeightToFileName[coreOp].insert(weightToFile);
+          mapCoreWeightToFileName[coreId].insert(weightToFile);
           continue;
         }
 
@@ -756,22 +1145,28 @@ createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
 
         weightFileStream.close();
         mapGlobalOpToFileName.insert({globalOp, newFileName});
-      mapCoreWeightToFileName[coreOp].insert({weight, newFileName});
+        mapCoreWeightToFileName[coreId].insert({weight, newFileName});
       }
     }
+
+    for (pim::PimCoreOp coreOp : scalarCores)
+      if (coreOp.getOperation() != op)
+        coreOp.erase();
+  }
   return mapCoreWeightToFileName;
 }
 
 /// Write the top-level PIM configuration JSON (core count, crossbar config, I/O addresses).
 static OnnxMlirCompilerErrorCodes writeConfigJson(func::FuncOp funcOp,
                                                   PimAcceleratorMemory& memory,
-                                                  size_t coreCount,
+                                                  size_t maxCoreId,
                                                   json::Object xbarsPerArrayGroup,
                                                   StringRef outputDirPath) {
   json::Object configJson;
 
-  // +1 because pimsim-nn also considers the host as a core
-  configJson["core_cnt"] = coreCount + 1;
+  // pimsim-nn indexes cores directly by their numeric core ID, with the host
+  // occupying core 0.
+  configJson["core_cnt"] = maxCoreId + 1;
 
   // TODO: Should this be based on the floating point type used in the model?
   // The 2 following values determine the bitwidth of the vectors' elements: bitwidth = adc_count * cell_precision
@@ -826,6 +1221,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
 
   PimAcceleratorMemory memory;
   memory.hostMem.allocateHost(moduleOp, funcOp);
+  memory.reportHost();
 
   if (auto err = writeMemoryBinary(moduleOp, funcOp, memory, outputDirPath))
     return err;
@@ -845,14 +1241,47 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
   // For each core, specify the number of crossbar per array group.
   // This implementation always assigns one crossbar per group.
   json::Object xbarsPerArrayGroup;
-  size_t coreCount = 0;
+  size_t maxCoreId = 0;
 
   // Create Weight Folder
   auto mapCoreWeightToFileName = createAndPopulateWeightFolder(funcOp, outputDirPath);
 
-  for (auto coreOp : funcOp.getOps<pim::PimCoreOp>()) {
-    auto coreId = coreOp.getCoreId();
-    coreCount++;
+  SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
+  llvm::DenseMap<size_t, size_t> emittedCoreIds;
+  size_t nextEmittedCoreId = 1;
+
+  for (Operation* op : coreLikeOps) {
+    if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
+      size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
+      if (!emittedCoreIds.contains(originalCoreId))
+        emittedCoreIds[originalCoreId] = nextEmittedCoreId++;
+      continue;
+    }
+
+    auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
+    auto batchCoreIds = getBatchCoreIds(coreBatchOp);
+    for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane) {
+      size_t originalCoreId = static_cast<size_t>(batchCoreIds[lane]);
+      if (!emittedCoreIds.contains(originalCoreId))
+        emittedCoreIds[originalCoreId] = nextEmittedCoreId++;
+    }
+  }
+
+  for (Operation* op : coreLikeOps) {
+    SmallVector<pim::PimCoreOp> scalarCores;
+    if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
+      scalarCores.push_back(coreOp);
+    }
+    else {
+      auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
+      for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
+        scalarCores.push_back(materializeScalarCoreFromBatchLane(coreBatchOp, lane));
+    }
+
+    for (pim::PimCoreOp coreOp : scalarCores) {
+      size_t originalCoreId = static_cast<size_t>(coreOp.getCoreId());
+      size_t coreId = emittedCoreIds.lookup(originalCoreId);
+      maxCoreId = std::max(maxCoreId, coreId);
 
       std::error_code errorCode;
       auto outputCorePath = outputDirPath + "/core_" + std::to_string(coreId) + ".json";
@@ -863,27 +1292,27 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
       }
       coreFileStream << '[';
 
-    PimCodeGen coreCodeGen(memory, coreFileStream);
+      PimCodeGen coreCodeGen(memory, coreFileStream, emittedCoreIds);
+      aliasMaterializedHostGlobals(moduleOp, funcOp, coreOp, memory);
       memory.getOrCreateDeviceMem(coreId).allocateCore(coreOp);
+      memory.reportCore(coreId);
 
       int64_t processedOperations = codeGenCoreOps(coreOp.getBody().front(), coreCodeGen);
       if (processedOperations < 0)
         return CompilerFailure;
       assert(processedOperations > 0);
 
-    // Remove trailing comma, close JSON array
       coreFileStream.seek(coreFileStream.tell() - 1);
       coreFileStream << ']';
       coreFileStream.close();
 
-    // Write crossbar weights for this core
       auto coreWeightsDirPath = outputDirPath + "/core_" + std::to_string(coreId);
       if (auto error = sys::fs::create_directory(coreWeightsDirPath)) {
         errs() << "Error creating core directory: " << coreWeightsDirPath << ": " << error.message() << '\n';
         return InvalidOutputFileAccess;
       }
 
-    auto& mapWeightToFile = mapCoreWeightToFileName[coreOp];
+      auto& mapWeightToFile = mapCoreWeightToFileName[originalCoreId];
       json::Array xbarsPerGroup;
       for (unsigned index : getUsedWeightIndices(coreOp)) {
         if (index >= coreOp.getWeights().size()) {
@@ -897,8 +1326,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
         if (auto error = sys::fs::create_link(outputDirPath + "/weights/" + fileName,
                                               coreWeightsDirPath + "/crossbar_" + std::to_string(index) + ".bin")) {
           errs() << "Error creating link file: " << (outputDirPath + "/weights/" + fileName) << " to "
-               << (coreWeightsDirPath + "/crossbar_" + std::to_string(index) + ".bin") << "\nError:" << error.message()
-               << '\n';
+                 << (coreWeightsDirPath + "/crossbar_" + std::to_string(index) + ".bin")
+                 << "\nError:" << error.message() << '\n';
           return InvalidOutputFileAccess;
         }
       }
@@ -906,5 +1335,12 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
       xbarsPerArrayGroup["core" + std::to_string(coreId)] = std::move(xbarsPerGroup);
     }
 
-  return writeConfigJson(funcOp, memory, coreCount, std::move(xbarsPerArrayGroup), outputDirPath);
+    for (pim::PimCoreOp coreOp : scalarCores)
+      if (coreOp.getOperation() != op) {
+        coreOp.walk([&memory](Operation* op) { memory.clean(op); });
+        coreOp.erase();
+      }
+  }
+
+  return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
 }

src/PIM/Compiler/PimCodeGen.hpp

@@ -1,7 +1,13 @@
 #pragma once
 
+#include "mlir/IR/Operation.h"
+
 #include "llvm-project/clang/include/clang/Basic/LLVM.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/JSON.h"
+#include "llvm/Support/raw_os_ostream.h"
+
+#include <fstream>
 
 #include "onnx-mlir/Compiler/OMCompilerTypes.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
@@ -33,6 +39,8 @@ public:
 
   void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
   void allocateCore(mlir::Operation* op);
+  void report(llvm::raw_ostream& os);
+  void remove(mlir::Value val);
 
   size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
   MemEntry getMemEntry(mlir::Value value) const;
@@ -45,23 +53,39 @@ public:
 
 private:
   llvm::SmallDenseMap<size_t, PimMemory> deviceMem;
+  std::fstream fileReport;
 
 public:
   PimAcceleratorMemory()
-  : hostMem(memEntriesMap) {}
+  : hostMem(memEntriesMap) {
+
+    std::string outputDir = getOutputDir();
+    if (outputDir.empty())
+      return;
+
+    std::string dialectsDir = outputDir + "/reports/";
+    createDirectory(dialectsDir);
+    std::fstream file(dialectsDir + "/memory_report.txt", std::ios::out);
+    fileReport = std::move(file);
+  }
 
   PimMemory& getOrCreateDeviceMem(size_t id);
 
   size_t getValueAddress(mlir::Value value, const StaticValueKnowledge& knowledge = {}) const;
+  void reportHost();
+  void reportCore(size_t coreId);
+  void clean(mlir::Operation* op);
 };
 
 class PimCodeGen {
   PimAcceleratorMemory& memory;
   llvm::raw_fd_ostream& coreFileStream;
+  const llvm::DenseMap<size_t, size_t>& emittedCoreIds;
 
   size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
     return memory.getValueAddress(value, knowledge);
   }
+  size_t remapCoreId(size_t coreId) const;
 
   static llvm::json::Object createEmptyOffset();
   void emitInstruction(llvm::json::Object instruction) const;
@@ -83,15 +107,23 @@ class PimCodeGen {
   void emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset) const;
 
 public:
-  PimCodeGen(PimAcceleratorMemory& memory, llvm::raw_fd_ostream& coreJson)
-  : memory(memory), coreFileStream(coreJson) {}
+  PimCodeGen(PimAcceleratorMemory& memory,
+             llvm::raw_fd_ostream& coreJson,
+             const llvm::DenseMap<size_t, size_t>& emittedCoreIds)
+  : memory(memory), coreFileStream(coreJson), emittedCoreIds(emittedCoreIds) {}
 
   void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const;
   void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const;
   void codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledge& knowledge) const;
 
   void codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenReceiveManyOp(pim::PimReceiveManyOp receiveManyOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenReceiveTensorOp(pim::PimReceiveTensorOp receiveTensorOp, const StaticValueKnowledge& knowledge) const;
   void codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenSendManyOp(pim::PimSendManyOp sendManyOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenSendTensorOp(pim::PimSendTensorOp sendTensorOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenExtractRowsOp(pim::PimExtractRowsOp extractRowsOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const;
 
   template <typename MVMTy>
   void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge);
@@ -106,6 +138,7 @@ public:
   void codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowledge& knowledge) const;
   void codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowledge& knowledge) const;
   void codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticValueKnowledge& knowledge) const;
+  void codeGetGlobalOp(mlir::memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const;
   void codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const;
 };
 

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -18,7 +18,9 @@ add_pim_library(OMONNXToSpatial
   Patterns/Tensor/Reshape.cpp
   Patterns/Tensor/Split.cpp
   ONNXToSpatialPass.cpp
-  Common.cpp
+  Common/ComputeRegionBuilder.cpp
+  Common/ShapeTilingUtils.cpp
+  Common/WeightMaterialization.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

src/PIM/Conversion/ONNXToSpatial/Common.hpp

@@ -1,279 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/Block.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/ValueRange.h"
-#include "mlir/Transforms/DialectConversion.h"
-
-#include <cassert>
-#include <type_traits>
-#include <utility>
-
-#include "llvm/ADT/SmallPtrSet.h"
-
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
-namespace onnx_mlir {
-
-template <class ShapedType>
-inline auto getImageWidth(const ShapedType& shapedType) {
-  return shapedType.getDimSize(2);
-}
-
-template <class ShapedType>
-inline auto getImageHeight(const ShapedType& shapedType) {
-  return shapedType.getDimSize(3);
-}
-
-template <class ShapedType>
-inline auto getImageChannel(const ShapedType& shapedType) {
-  return shapedType.getDimSize(1);
-}
-
-template <class ShapedType>
-inline auto getImageN(const ShapedType& shapedType) {
-  return shapedType.getDimSize(0);
-}
-
-template <class ShapedType>
-inline auto getKernelWidth(const ShapedType& shapedType) {
-  return shapedType.getDimSize(2);
-}
-
-template <class ShapedType>
-inline auto getKernelHeight(const ShapedType& shapedType) {
-  return shapedType.getDimSize(3);
-}
-
-template <class ShapedType>
-inline auto getFilterCount(const ShapedType& shapedType) {
-  return shapedType.getDimSize(0);
-}
-
-using HSliceId = size_t;
-using CoreId = size_t;
-
-template <class A, class B, class C = std::common_type_t<A, B>>
-constexpr C ceilIntegerDivide(A a, B b) {
-  static_assert(std::is_integral_v<A>, "A must be an integer type");
-  static_assert(std::is_integral_v<B>, "B must be an integer type");
-  C ac = static_cast<C>(a);
-  C bc = static_cast<C>(b);
-  return 1 + (ac - 1) / bc;
-}
-
-template <class A, class B, class C = std::common_type_t<A, B>>
-constexpr std::pair<C, C> ceilIntegerDivideWithRemainder(A a, B b) {
-  static_assert(std::is_integral_v<A>, "A must be an integer type");
-  static_assert(std::is_integral_v<B>, "B must be an integer type");
-  C ac = static_cast<C>(a);
-  C bc = static_cast<C>(b);
-  return {ceilIntegerDivide(ac, bc), ac % bc};
-}
-
-template <class T>
-bool isVectorShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2 && (shape[0] == 1 || shape[1] == 1);
-}
-
-template <class T>
-bool isMatrixShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2;
-}
-
-template <class T>
-bool isHVectorShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2 && shape[0] == 1;
-}
-
-template <class T>
-bool isVVectorShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2 && shape[1] == 1;
-}
-
-template <class T>
-T getVectorLength(mlir::ArrayRef<T> shape) {
-  assert(isVectorShape(shape));
-  return shape[0] != 1 ? shape[0] : shape[1];
-}
-
-inline auto getTensorShape(mlir::Value tensor) {
-  return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
-}
-
-inline bool isWeightLikeComputeOperand(mlir::Value value) {
-  auto rankedType = mlir::dyn_cast<mlir::RankedTensorType>(value.getType());
-  if (!rankedType || !isMatrixShape(rankedType.getShape()))
-    return false;
-
-  llvm::SmallPtrSet<mlir::Operation*, 8> visited;
-
-  while (auto* definingOp = value.getDefiningOp()) {
-    if (!visited.insert(definingOp).second)
-      return false;
-    if (hasWeightAlways(definingOp))
-      return true;
-
-    if (auto extractSliceOp = mlir::dyn_cast<mlir::tensor::ExtractSliceOp>(definingOp)) {
-      value = extractSliceOp.getSource();
-      continue;
-    }
-    if (auto expandShapeOp = mlir::dyn_cast<mlir::tensor::ExpandShapeOp>(definingOp)) {
-      value = expandShapeOp.getSrc();
-      continue;
-    }
-    if (auto collapseShapeOp = mlir::dyn_cast<mlir::tensor::CollapseShapeOp>(definingOp)) {
-      value = collapseShapeOp.getSrc();
-      continue;
-    }
-    if (auto transposeOp = mlir::dyn_cast<mlir::ONNXTransposeOp>(definingOp)) {
-      value = transposeOp.getData();
-      continue;
-    }
-
-    return false;
-  }
-
-  return false;
-}
-
-namespace detail {
-
-inline mlir::ValueRange getBlockArgs(mlir::Block* block) { return mlir::ValueRange(block->getArguments()); }
-
-template <typename Fn, size_t... Is>
-decltype(auto) invokeWithBlockArgs(Fn&& fn, mlir::Block* block, std::index_sequence<Is...>) {
-  return std::forward<Fn>(fn)(block->getArgument(Is)...);
-}
-
-template <typename Fn, size_t... Is>
-decltype(auto) invokeWithValues(Fn&& fn, mlir::ArrayRef<mlir::Value> values, std::index_sequence<Is...>) {
-  return std::forward<Fn>(fn)(values[Is]...);
-}
-
-template <size_t>
-using ValueArg = mlir::Value;
-
-template <typename Fn, typename Seq>
-struct InvokeWithBlockArgsResult;
-
-template <typename Fn, size_t... Is>
-struct InvokeWithBlockArgsResult<Fn, std::index_sequence<Is...>> {
-  using type = std::invoke_result_t<Fn, ValueArg<Is>...>;
-};
-
-template <typename Fn, typename Seq>
-using InvokeWithBlockArgsResultT = typename InvokeWithBlockArgsResult<Fn, Seq>::type;
-
-template <typename Fn>
-using InvokeWithValueRangeResultT = std::invoke_result_t<Fn, mlir::ValueRange>;
-
-} // namespace detail
-
-template <size_t NumInputs, typename RewriterT, typename BodyFn>
-auto createSpatCompute(RewriterT& rewriter,
-                       mlir::Location loc,
-                       mlir::TypeRange resultTypes,
-                       mlir::ValueRange weights,
-                       mlir::ValueRange inputs,
-                       BodyFn&& body) {
-  assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
-  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
-
-  using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
-  if constexpr (std::is_same_v<BodyResult, mlir::LogicalResult>) {
-    auto bodyResult =
-      detail::invokeWithBlockArgs(std::forward<BodyFn>(body), block, std::make_index_sequence<NumInputs> {});
-    if (mlir::failed(bodyResult)) {
-      rewriter.setInsertionPointAfter(computeOp);
-      rewriter.eraseOp(computeOp);
-      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
-    }
-    rewriter.setInsertionPointAfter(computeOp);
-    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
-  }
-  else {
-    static_assert(std::is_same_v<BodyResult, void>, "createSpatCompute body must return void or mlir::LogicalResult");
-    detail::invokeWithBlockArgs(std::forward<BodyFn>(body), block, std::make_index_sequence<NumInputs> {});
-
-    rewriter.setInsertionPointAfter(computeOp);
-    return computeOp;
-  }
-}
-
-template <typename RewriterT, typename BodyFn>
-auto createSpatCompute(RewriterT& rewriter,
-                       mlir::Location loc,
-                       mlir::TypeRange resultTypes,
-                       mlir::ValueRange weights,
-                       mlir::ValueRange inputs,
-                       BodyFn&& body) {
-  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
-
-  using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
-  if constexpr (std::is_same_v<BodyResult, mlir::LogicalResult>) {
-    auto bodyResult = std::forward<BodyFn>(body)(detail::getBlockArgs(block));
-    if (mlir::failed(bodyResult)) {
-      rewriter.setInsertionPointAfter(computeOp);
-      rewriter.eraseOp(computeOp);
-      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
-    }
-    rewriter.setInsertionPointAfter(computeOp);
-    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
-  }
-  else {
-    static_assert(std::is_same_v<BodyResult, void>, "createSpatCompute body must return void or mlir::LogicalResult");
-    std::forward<BodyFn>(body)(detail::getBlockArgs(block));
-
-    rewriter.setInsertionPointAfter(computeOp);
-    return computeOp;
-  }
-}
-
-llvm::SmallVector<mlir::Value> sliceTensor(const mlir::Value& tensorToSlice,
-                                           size_t axis,
-                                           int64_t sliceSize,
-                                           mlir::ConversionPatternRewriter& rewriter,
-                                           mlir::Location loc);
-
-llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
-                                           int64_t sliceSize,
-                                           mlir::ConversionPatternRewriter& rewriter,
-                                           mlir::Location loc);
-
-llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
-  const mlir::Value& vectorToSlice, mlir::ConversionPatternRewriter& rewriter, mlir::Location loc);
-
-llvm::DenseMap<HSliceId, llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>>>
-tileMatrix(mlir::Value& matrixToTile,
-           int64_t hSliceSize,
-           int64_t vSliceSize,
-           mlir::ConversionPatternRewriter& rewriter,
-           mlir::Location& loc);
-
-mlir::tensor::SplatOp broadcastToVector(mlir::Value scalarToBroadcast,
-                                        int64_t length,
-                                        mlir::ConversionPatternRewriter& rewriter,
-                                        mlir::Location loc);
-
-mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::ConversionPatternRewriter& rewriter);
-
-}; // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/Common.hpp

@@ -0,0 +1,8 @@
+#pragma once
+
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+#include "ComputeRegionBuilder.hpp"
+#include "ShapeTilingUtils.hpp"
+#include "WeightMaterialization.hpp"

src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.cpp

@@ -0,0 +1,39 @@
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+Value sumTensors(ArrayRef<Value> tensors, ConversionPatternRewriter& rewriter) {
+  if (tensors.size() == 1)
+    return tensors[0];
+
+  SmallVector<Value> tensors1 = {tensors.begin(), tensors.end()};
+  SmallVector<Value> tensors2;
+  tensors2.reserve(tensors.size() / 2);
+
+  auto* currTensors = &tensors1;
+  auto* nextTensors = &tensors2;
+  while (currTensors->size() > 1) {
+    for (size_t i = 0; i < currTensors->size() - 1; i += 2) {
+      Value a = (*currTensors)[i];
+      Value b = (*currTensors)[i + 1];
+      rewriter.setInsertionPointAfterValue(b);
+      auto addedValue = spatial::SpatVAddOp::create(rewriter, a.getLoc(), a.getType(), a, b);
+      nextTensors->push_back(addedValue);
+    }
+    if (currTensors->size() % 2 == 1)
+      nextTensors->push_back(currTensors->back());
+    std::swap(currTensors, nextTensors);
+    nextTensors->clear();
+  }
+  assert(currTensors->size() == 1 && "Expected a single input at this point.");
+  return (*currTensors)[0];
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp

@@ -0,0 +1,153 @@
+#pragma once
+
+#include "mlir/IR/Block.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/ValueRange.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include <cassert>
+#include <cstddef>
+#include <type_traits>
+#include <utility>
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+
+namespace detail {
+
+inline mlir::ValueRange getBlockArgs(mlir::Block* block) { return mlir::ValueRange(block->getArguments()); }
+
+template <typename Fn, size_t... Is>
+decltype(auto) invokeWithBlockArgs(Fn&& fn, mlir::Block* block, std::index_sequence<Is...>) {
+  return std::forward<Fn>(fn)(block->getArgument(Is)...);
+}
+
+template <typename Fn, size_t... Is>
+decltype(auto) invokeWithValues(Fn&& fn, mlir::ArrayRef<mlir::Value> values, std::index_sequence<Is...>) {
+  return std::forward<Fn>(fn)(values[Is]...);
+}
+
+template <size_t>
+using ValueArg = mlir::Value;
+
+template <typename Fn, typename Seq>
+struct InvokeWithBlockArgsResult;
+
+template <typename Fn, size_t... Is>
+struct InvokeWithBlockArgsResult<Fn, std::index_sequence<Is...>> {
+  using type = std::invoke_result_t<Fn, ValueArg<Is>...>;
+};
+
+template <typename Fn, typename Seq>
+using InvokeWithBlockArgsResultT = typename InvokeWithBlockArgsResult<Fn, Seq>::type;
+
+template <typename Fn>
+using InvokeWithValueRangeResultT = std::invoke_result_t<Fn, mlir::ValueRange>;
+
+} // namespace detail
+
+template <typename RewriterT>
+inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int64_t axis, mlir::ValueRange inputs) {
+  assert(!inputs.empty() && "spat.concat requires at least one input");
+  if (inputs.size() == 1)
+    return inputs.front();
+
+  auto firstType = mlir::cast<mlir::RankedTensorType>(inputs.front().getType());
+  auto outputShape = llvm::to_vector(firstType.getShape());
+  int64_t concatDimSize = 0;
+  bool concatDimDynamic = false;
+
+  for (mlir::Value input : inputs) {
+    auto inputType = mlir::cast<mlir::RankedTensorType>(input.getType());
+    assert(inputType.getRank() == firstType.getRank() && "spat.concat expects same-rank inputs");
+    if (mlir::ShapedType::isDynamic(inputType.getDimSize(axis)))
+      concatDimDynamic = true;
+    else
+      concatDimSize += inputType.getDimSize(axis);
+  }
+
+  outputShape[axis] = concatDimDynamic ? mlir::ShapedType::kDynamic : concatDimSize;
+  auto outputType = mlir::RankedTensorType::get(outputShape, firstType.getElementType(), firstType.getEncoding());
+  return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
+}
+
+/// Builds a `spat.compute` with a fixed number of SSA inputs and erases it if
+/// the body callback reports failure.
+template <size_t NumInputs, typename RewriterT, typename BodyFn>
+auto createSpatCompute(RewriterT& rewriter,
+                       mlir::Location loc,
+                       mlir::TypeRange resultTypes,
+                       mlir::ValueRange weights,
+                       mlir::ValueRange inputs,
+                       BodyFn&& body) {
+  assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
+  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
+
+  auto* block = new mlir::Block();
+  for (mlir::Value input : inputs)
+    block->addArgument(input.getType(), loc);
+
+  computeOp.getBody().push_back(block);
+  rewriter.setInsertionPointToStart(block);
+
+  using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
+  if constexpr (std::is_same_v<BodyResult, void>) {
+    detail::invokeWithBlockArgs(std::forward<BodyFn>(body), block, std::make_index_sequence<NumInputs> {});
+
+    rewriter.setInsertionPointAfter(computeOp);
+    return computeOp;
+  }
+  else {
+    auto bodyResult =
+      detail::invokeWithBlockArgs(std::forward<BodyFn>(body), block, std::make_index_sequence<NumInputs> {});
+    if (mlir::failed(bodyResult)) {
+      rewriter.setInsertionPointAfter(computeOp);
+      rewriter.eraseOp(computeOp);
+      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
+    }
+    rewriter.setInsertionPointAfter(computeOp);
+    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
+  }
+}
+
+/// Builds a `spat.compute` whose body consumes the block arguments as a single
+/// `ValueRange`, which is convenient for variadic reductions/concats.
+template <typename RewriterT, typename BodyFn>
+auto createSpatCompute(RewriterT& rewriter,
+                       mlir::Location loc,
+                       mlir::TypeRange resultTypes,
+                       mlir::ValueRange weights,
+                       mlir::ValueRange inputs,
+                       BodyFn&& body) {
+  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
+
+  auto* block = new mlir::Block();
+  for (mlir::Value input : inputs)
+    block->addArgument(input.getType(), loc);
+
+  computeOp.getBody().push_back(block);
+  rewriter.setInsertionPointToStart(block);
+
+  using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
+  if constexpr (std::is_same_v<BodyResult, void>) {
+    std::forward<BodyFn>(body)(detail::getBlockArgs(block));
+
+    rewriter.setInsertionPointAfter(computeOp);
+    return computeOp;
+  }
+  else {
+    auto bodyResult = std::forward<BodyFn>(body)(detail::getBlockArgs(block));
+    if (mlir::failed(bodyResult)) {
+      rewriter.setInsertionPointAfter(computeOp);
+      rewriter.eraseOp(computeOp);
+      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
+    }
+    rewriter.setInsertionPointAfter(computeOp);
+    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
+  }
+}
+
+mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::ConversionPatternRewriter& rewriter);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.cpp

@@ -1,24 +1,10 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tosa/IR/TosaOps.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Location.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/Value.h"
-#include "mlir/Transforms/DialectConversion.h"
 
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/Twine.h"
-#include "llvm/Support/Casting.h"
 
-#include <cassert>
-#include <optional>
-#include <utility>
-
-#include "Common.hpp"
+#include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
 
 using namespace mlir;
 
@@ -107,31 +93,4 @@ broadcastToVector(Value scalarToBroadcast, int64_t length, ConversionPatternRewr
   return tensor::SplatOp::create(rewriter, loc, type, elementValue);
 }
 
-Value sumTensors(ArrayRef<Value> tensors, ConversionPatternRewriter& rewriter) {
-  if (tensors.size() == 1)
-    return tensors[0];
-
-  SmallVector<Value> tensors1 = {tensors.begin(), tensors.end()};
-  SmallVector<Value> tensors2;
-  tensors2.reserve(tensors.size() / 2);
-
-  auto* currTensors = &tensors1;
-  auto* nextTensors = &tensors2;
-  while (currTensors->size() > 1) {
-    for (size_t i = 0; i < currTensors->size() - 1; i += 2) {
-      Value a = (*currTensors)[i];
-      Value b = (*currTensors)[i + 1];
-      rewriter.setInsertionPointAfterValue(b);
-      auto addedValue = spatial::SpatVAddOp::create(rewriter, a.getLoc(), a.getType(), a, b);
-      nextTensors->push_back(addedValue);
-    }
-    if (currTensors->size() % 2 == 1)
-      nextTensors->push_back(currTensors->back());
-    std::swap(currTensors, nextTensors);
-    nextTensors->clear();
-  }
-  assert(currTensors->size() == 1 && "Expected a single input at this point.");
-  return (*currTensors)[0];
-}
-
-}; // namespace onnx_mlir
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp

@@ -0,0 +1,143 @@
+#pragma once
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include <cassert>
+#include <cstddef>
+#include <type_traits>
+#include <utility>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace onnx_mlir {
+
+template <class ShapedType>
+inline auto getImageWidth(const ShapedType& shapedType) {
+  return shapedType.getDimSize(2);
+}
+
+template <class ShapedType>
+inline auto getImageHeight(const ShapedType& shapedType) {
+  return shapedType.getDimSize(3);
+}
+
+template <class ShapedType>
+inline auto getImageChannel(const ShapedType& shapedType) {
+  return shapedType.getDimSize(1);
+}
+
+template <class ShapedType>
+inline auto getImageN(const ShapedType& shapedType) {
+  return shapedType.getDimSize(0);
+}
+
+template <class ShapedType>
+inline auto getKernelWidth(const ShapedType& shapedType) {
+  return shapedType.getDimSize(2);
+}
+
+template <class ShapedType>
+inline auto getKernelHeight(const ShapedType& shapedType) {
+  return shapedType.getDimSize(3);
+}
+
+template <class ShapedType>
+inline auto getFilterCount(const ShapedType& shapedType) {
+  return shapedType.getDimSize(0);
+}
+
+using HSliceId = size_t;
+using CoreId = size_t;
+
+template <class A, class B, class C = std::common_type_t<A, B>>
+constexpr C ceilIntegerDivide(A a, B b) {
+  static_assert(std::is_integral_v<A>, "A must be an integer type");
+  static_assert(std::is_integral_v<B>, "B must be an integer type");
+  C ac = static_cast<C>(a);
+  C bc = static_cast<C>(b);
+  return 1 + (ac - 1) / bc;
+}
+
+template <class A, class B, class C = std::common_type_t<A, B>>
+constexpr std::pair<C, C> ceilIntegerDivideWithRemainder(A a, B b) {
+  static_assert(std::is_integral_v<A>, "A must be an integer type");
+  static_assert(std::is_integral_v<B>, "B must be an integer type");
+  C ac = static_cast<C>(a);
+  C bc = static_cast<C>(b);
+  return {ceilIntegerDivide(ac, bc), ac % bc};
+}
+
+template <class T>
+bool isVectorShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2 && (shape[0] == 1 || shape[1] == 1);
+}
+
+template <class T>
+bool isMatrixShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2;
+}
+
+template <class T>
+bool isHVectorShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2 && shape[0] == 1;
+}
+
+template <class T>
+bool isVVectorShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2 && shape[1] == 1;
+}
+
+template <class T>
+T getVectorLength(mlir::ArrayRef<T> shape) {
+  assert(isVectorShape(shape));
+  return shape[0] != 1 ? shape[0] : shape[1];
+}
+
+inline auto getTensorShape(mlir::Value tensor) {
+  return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
+}
+
+inline bool haveSameStaticShape(mlir::Value lhs, mlir::Value rhs) {
+  auto lhsType = mlir::dyn_cast<mlir::RankedTensorType>(lhs.getType());
+  auto rhsType = mlir::dyn_cast<mlir::RankedTensorType>(rhs.getType());
+  return lhsType && rhsType && lhsType.hasStaticShape() && rhsType.hasStaticShape() && lhsType.getShape() == rhsType.getShape();
+}
+
+/// Slices a statically shaped tensor along one axis into contiguous pieces of
+/// at most `sliceSize` elements.
+llvm::SmallVector<mlir::Value> sliceTensor(const mlir::Value& tensorToSlice,
+                                           size_t axis,
+                                           int64_t sliceSize,
+                                           mlir::ConversionPatternRewriter& rewriter,
+                                           mlir::Location loc);
+
+llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
+                                           int64_t sliceSize,
+                                           mlir::ConversionPatternRewriter& rewriter,
+                                           mlir::Location loc);
+
+/// Partitions one logical vector into per-core crossbar-sized slices using the
+/// current PIM target geometry.
+llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
+  const mlir::Value& vectorToSlice, mlir::ConversionPatternRewriter& rewriter, mlir::Location loc);
+
+/// Tiles a matrix first across output columns and then across input rows so it
+/// can be assigned to crossbars grouped by core.
+llvm::DenseMap<HSliceId, llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>>>
+tileMatrix(mlir::Value& matrixToTile,
+           int64_t hSliceSize,
+           int64_t vSliceSize,
+           mlir::ConversionPatternRewriter& rewriter,
+           mlir::Location& loc);
+
+mlir::tensor::SplatOp broadcastToVector(mlir::Value scalarToBroadcast,
+                                        int64_t length,
+                                        mlir::ConversionPatternRewriter& rewriter,
+                                        mlir::Location loc);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.cpp

@@ -0,0 +1,114 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include "WeightMaterialization.hpp"
+#include "ShapeTilingUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+bool isWeightLikeComputeOperand(Value value) {
+  auto rankedType = dyn_cast<RankedTensorType>(value.getType());
+  if (!rankedType || !isMatrixShape(rankedType.getShape()))
+    return false;
+
+  llvm::SmallPtrSet<Operation*, 8> visited;
+
+  while (auto* definingOp = value.getDefiningOp()) {
+    if (!visited.insert(definingOp).second)
+      return false;
+    if (hasWeightAlways(definingOp))
+      return true;
+
+    if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {
+      value = extractSliceOp.getSource();
+      continue;
+    }
+    if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(definingOp)) {
+      value = expandShapeOp.getSrc();
+      continue;
+    }
+    if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(definingOp)) {
+      value = collapseShapeOp.getSrc();
+      continue;
+    }
+    if (auto transposeOp = dyn_cast<ONNXTransposeOp>(definingOp)) {
+      value = transposeOp.getData();
+      continue;
+    }
+
+    return false;
+  }
+
+  return false;
+}
+
+FailureOr<Value> materializeWeightLikeValueInBlock(Value value, IRRewriter& rewriter, IRMapping& mapper) {
+  if (auto mapped = mapper.lookupOrNull(value))
+    return cast<Value>(mapped);
+
+  Operation* definingOp = value.getDefiningOp();
+  if (!definingOp)
+    return failure();
+
+  if (isa<arith::ConstantOp, ONNXConstantOp>(definingOp)) {
+    auto tensorType = dyn_cast<RankedTensorType>(value.getType());
+    if (!tensorType || !tensorType.hasStaticShape())
+      return failure();
+
+    SmallVector<OpFoldResult> offsets(tensorType.getRank(), rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> sizes;
+    SmallVector<OpFoldResult> strides(tensorType.getRank(), rewriter.getIndexAttr(1));
+    sizes.reserve(tensorType.getRank());
+    for (int64_t dim : tensorType.getShape())
+      sizes.push_back(rewriter.getIndexAttr(dim));
+
+    auto referencedValue =
+      tensor::ExtractSliceOp::create(rewriter, value.getLoc(), tensorType, value, offsets, sizes, strides);
+    mapper.map(value, referencedValue.getResult());
+    return referencedValue.getResult();
+  }
+
+  if (!isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, ONNXTransposeOp>(definingOp))
+    return failure();
+
+  IRMapping localMapper;
+  for (Value operand : definingOp->getOperands()) {
+    if (auto mapped = mapper.lookupOrNull(operand)) {
+      localMapper.map(operand, cast<Value>(mapped));
+      continue;
+    }
+
+    if (isWeightLikeComputeOperand(operand)) {
+      auto clonedOperand = materializeWeightLikeValueInBlock(operand, rewriter, mapper);
+      if (failed(clonedOperand))
+        return failure();
+      localMapper.map(operand, *clonedOperand);
+      continue;
+    }
+
+    localMapper.map(operand, operand);
+  }
+
+  Operation* clonedOp = rewriter.clone(*definingOp, localMapper);
+  for (auto [oldResult, newResult] : llvm::zip(definingOp->getResults(), clonedOp->getResults()))
+    mapper.map(oldResult, newResult);
+
+  auto mapped = mapper.lookupOrNull(value);
+  if (!mapped)
+    return failure();
+  return cast<Value>(mapped);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+
+namespace onnx_mlir {
+
+/// Returns true when a matrix-valued compute operand is ultimately backed by a
+/// weight-marked constant/view chain and can be promoted into weights.
+bool isWeightLikeComputeOperand(mlir::Value value);
+
+/// Rebuilds the view/transpose chain of a promoted weight operand inside a new
+/// compute body while reusing already-materialized intermediate values.
+llvm::FailureOr<mlir::Value>
+materializeWeightLikeValueInBlock(mlir::Value value, mlir::IRRewriter& rewriter, mlir::IRMapping& mapper);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -17,7 +18,7 @@
 #include <iterator>
 #include <utility>
 
-#include "Common.hpp"
+#include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
@@ -30,8 +31,6 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-bool haveSameStaticShape(Value lhs, Value rhs);
-
 namespace {
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
@@ -48,12 +47,51 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
 
 private:
   void annotateWeightsConstants(func::FuncOp funcOp) const;
-  void encapsulateGlobalInstruction(func::FuncOp funcOp);
+  LogicalResult encapsulateGlobalInstruction(func::FuncOp funcOp);
   LogicalResult promoteConstantInputsToWeights(func::FuncOp funcOp);
+  void populateEmptyFunction(func::FuncOp funcOp);
 };
 
 } // namespace
 
+static void foldSingleLaneComputeBatches(func::FuncOp funcOp) {
+  IRRewriter rewriter(funcOp.getContext());
+  SmallVector<spatial::SpatComputeBatch> batchOps;
+  funcOp.walk([&](spatial::SpatComputeBatch batchOp) { batchOps.push_back(batchOp); });
+
+  for (auto batchOp : batchOps) {
+    if (batchOp.getLaneCount() != 1)
+      continue;
+
+    auto loc = batchOp.getLoc();
+    rewriter.setInsertionPoint(batchOp);
+    auto computeOp =
+      spatial::SpatCompute::create(rewriter, loc, batchOp.getResultTypes(), batchOp.getWeights(), batchOp.getInputs());
+    computeOp.getProperties().setOperandSegmentSizes(
+      {static_cast<int>(batchOp.getWeights().size()), static_cast<int>(batchOp.getInputs().size())});
+
+    Block& templateBlock = batchOp.getBody().front();
+    SmallVector<Type> blockArgTypes;
+    SmallVector<Location> blockArgLocs;
+    for (BlockArgument arg : templateBlock.getArguments()) {
+      blockArgTypes.push_back(arg.getType());
+      blockArgLocs.push_back(loc);
+    }
+    auto* newBlock =
+      rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
+
+    IRMapping mapper;
+    for (auto [oldArg, newArg] : llvm::zip(templateBlock.getArguments(), newBlock->getArguments()))
+      mapper.map(oldArg, newArg);
+    rewriter.setInsertionPointToEnd(newBlock);
+    for (Operation& op : templateBlock)
+      rewriter.clone(op, mapper);
+
+    batchOp.replaceAllUsesWith(computeOp.getResults());
+    rewriter.eraseOp(batchOp);
+  }
+}
+
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   MLIRContext* ctx = &getContext();
@@ -124,6 +162,8 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  foldSingleLaneComputeBatches(*entryFunc);
+
   // Count the number of compute ops and check they do not exceed the core count
   if (coresCount != -1) {
     int computeOpsCount = 0;
@@ -144,7 +184,13 @@ void ONNXToSpatialPass::runOnOperation() {
     llvm::dbgs() << "Failed to run canonicalization cleanup, continuing...\n";
 
   annotateWeightsConstants(*entryFunc);
-  encapsulateGlobalInstruction(*entryFunc);
+
+  populateEmptyFunction(*entryFunc);
+
+  if (failed(encapsulateGlobalInstruction(*entryFunc))) {
+    signalPassFailure();
+    return;
+  }
 
   if (failed(promoteConstantInputsToWeights(*entryFunc))) {
     signalPassFailure();
@@ -160,7 +206,6 @@ bool encapsulator(IRRewriter& rewriter, Location loc, Operation* inst, std::func
   if (T toRemoveOp = llvm::dyn_cast_if_present<T>(inst)) {
     Value source = funcSource(toRemoveOp);
     rewriter.setInsertionPointAfter(toRemoveOp);
-    if (isa_and_present<spatial::SpatCompute>(source.getDefiningOp())) {
     auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), source);
     auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
     newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
@@ -173,6 +218,24 @@ bool encapsulator(IRRewriter& rewriter, Location loc, Operation* inst, std::func
     inst->erase();
     return true;
   }
+  return false;
+}
+
+bool encapsulateSlice(IRRewriter& rewriter, Location loc, Operation* inst) {
+  if (tensor::ExtractSliceOp toRemoveOp = llvm::dyn_cast_if_present<tensor::ExtractSliceOp>(inst)) {
+    auto source = toRemoveOp.getSource();
+    rewriter.setInsertionPointAfter(toRemoveOp);
+    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), source);
+    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
+    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
+    rewriter.setInsertionPointToEnd(BB);
+    IRMapping mapper;
+    mapper.map(source, BB->getArgument(0));
+    auto newInst = rewriter.clone(*inst, mapper);
+    spatial::SpatYieldOp::create(rewriter, loc, newInst->getResults());
+    inst->replaceAllUsesWith(newCompute->getResults());
+    inst->erase();
+    return true;
   }
   return false;
 }
@@ -194,6 +257,29 @@ bool encapsulateConcat(IRRewriter& rewriter, Location loc, Operation* inst) {
       newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sources.size()});
       rewriter.setInsertionPointToEnd(BB);
       IRMapping mapper;
+      for (auto [source, bbArg] : llvm::zip(sources, BB->getArguments()))
+        mapper.map(source, bbArg);
+      auto newConcat = spatial::SpatConcatOp::create(rewriter,
+                                                     loc,
+                                                     toRemoveOp.getType(),
+                                                     rewriter.getI64IntegerAttr(toRemoveOp.getDim()),
+                                                     ValueRange(BB->getArguments()));
+      spatial::SpatYieldOp::create(rewriter, loc, newConcat.getOutput());
+      inst->replaceAllUsesWith(newCompute->getResults());
+      inst->erase();
+      return true;
+    }
+    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), sources);
+    SmallVector<Type> sourceTypes;
+    SmallVector<Location> sourceLoc;
+    for (auto source : sources) {
+      sourceTypes.push_back(source.getType());
+      sourceLoc.push_back(loc);
+    }
+    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
+    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sources.size()});
+    rewriter.setInsertionPointToEnd(BB);
+    IRMapping mapper;
     for (auto [source, bbArg] : llvm::zip(sources, BB->getArguments()))
       mapper.map(source, bbArg);
     auto newConcat = rewriter.clone(*inst, mapper);
@@ -202,77 +288,109 @@ bool encapsulateConcat(IRRewriter& rewriter, Location loc, Operation* inst) {
     inst->erase();
     return true;
   }
-  }
   return false;
 }
 
-static FailureOr<Value> materializeWeightLikeValueInBlock(Value value, IRRewriter& rewriter, IRMapping& mapper) {
-  if (auto mapped = mapper.lookupOrNull(value))
-    return cast<Value>(mapped);
+static FailureOr<bool> sourceOperandHasWeightAlways(Operation* op) {
+  if (op == nullptr)
+    return false;
 
-  Operation* definingOp = value.getDefiningOp();
-  if (!definingOp)
-    return failure();
+  Operation* source = nullptr;
+  do {
 
-  if (isa<arith::ConstantOp, ONNXConstantOp>(definingOp)) {
-    auto tensorType = dyn_cast<RankedTensorType>(value.getType());
-    if (!tensorType || !tensorType.hasStaticShape())
-      return failure();
-
-    SmallVector<OpFoldResult> offsets(tensorType.getRank(), rewriter.getIndexAttr(0));
-    SmallVector<OpFoldResult> sizes;
-    SmallVector<OpFoldResult> strides(tensorType.getRank(), rewriter.getIndexAttr(1));
-    sizes.reserve(tensorType.getRank());
-    for (int64_t dim : tensorType.getShape())
-      sizes.push_back(rewriter.getIndexAttr(dim));
-
-    auto referencedValue =
-      tensor::ExtractSliceOp::create(rewriter, value.getLoc(), tensorType, value, offsets, sizes, strides);
-    mapper.map(value, referencedValue.getResult());
-    return referencedValue.getResult();
+    if (isa<spatial::SpatCompute, spatial::SpatComputeBatch>(*op)) {
+      return false;
     }
-
-  if (!isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, ONNXTransposeOp>(definingOp))
-    return failure();
-
-  IRMapping localMapper;
-  for (Value operand : definingOp->getOperands()) {
-    if (auto mapped = mapper.lookupOrNull(operand)) {
-      localMapper.map(operand, cast<Value>(mapped));
-      continue;
+    else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(*op)) {
+      auto tmpSource = extractSliceOp.getSource();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
     }
-
-    if (isWeightLikeComputeOperand(operand)) {
-      auto clonedOperand = materializeWeightLikeValueInBlock(operand, rewriter, mapper);
-      if (failed(clonedOperand))
-        return failure();
-      localMapper.map(operand, *clonedOperand);
-      continue;
+    else if (auto extractRowsOp = dyn_cast<spatial::SpatExtractRowsOp>(*op)) {
+      auto tmpSource = extractRowsOp.getInput();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
     }
-
-    localMapper.map(operand, operand);
+    else if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(*op)) {
+      auto tmpSource = expandShapeOp.getSrc();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
     }
-
-  Operation* clonedOp = rewriter.clone(*definingOp, localMapper);
-  for (auto [oldResult, newResult] : llvm::zip(definingOp->getResults(), clonedOp->getResults()))
-    mapper.map(oldResult, newResult);
-
-  auto mapped = mapper.lookupOrNull(value);
-  if (!mapped)
+    else if (auto transposeOp = dyn_cast<ONNXTransposeOp>(*op)) {
+      auto tmpSource = transposeOp.getData();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
+    }
+    else if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(*op)) {
+      auto tmpSource = collapseShapeOp.getSrc();
+      auto definingOp = tmpSource.getDefiningOp();
+      if (definingOp)
+        op = definingOp;
+      else
+        return false;
+    }
+    else if (auto constantOp = dyn_cast<arith::ConstantOp>(*op)) {
+      source = constantOp;
+    }
+    else if (auto concatOp = dyn_cast<tensor::ConcatOp>(*op)) {
+      bool res = false;
+      for (auto operand : concatOp.getOperands()) {
+        res |= hasWeightAlways(operand.getDefiningOp());
+        if (res)
+          return res;
+      }
+      return res;
+    }
+    else if (auto concatOp = dyn_cast<spatial::SpatConcatOp>(*op)) {
+      bool res = false;
+      for (auto operand : concatOp.getOperands()) {
+        res |= hasWeightAlways(operand.getDefiningOp());
+        if (res)
+          return res;
+      }
+      return res;
+    }
+    else {
+      op->emitOpError("unsupported global instruction while promoting weight-backed operands into Spatial computes");
       return failure();
-  return cast<Value>(mapped);
+    }
+  }
+  while (source == nullptr);
+
+  return hasWeightAlways(source);
 }
 
 // TODO what we want to keep in global?
-void ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
+LogicalResult ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
   Location loc = funcOp.getLoc();
   IRRewriter rewriter(&getContext());
   bool keep = true;
   while (keep) {
     keep = false;
     for (auto& instruction : llvm::make_early_inc_range(funcOp.getOps())) {
-      keep |= encapsulator<tensor::ExtractSliceOp>(
-        rewriter, loc, &instruction, [](tensor::ExtractSliceOp extract) { return extract.getSource(); });
+      if (isa<spatial::SpatCompute, spatial::SpatComputeBatch, spatial::SpatExtractRowsOp>(instruction)
+          || isa<func::ReturnOp>(instruction))
+        continue;
+
+      auto weightBacked = sourceOperandHasWeightAlways(&instruction);
+      if (failed(weightBacked))
+        return failure();
+      if (*weightBacked)
+        continue;
+
+      keep |= encapsulateSlice(rewriter, loc, &instruction);
 
       keep |= encapsulator<tensor::ExpandShapeOp>(
         rewriter, loc, &instruction, [](tensor::ExpandShapeOp expand) { return expand.getSrc(); });
@@ -286,6 +404,7 @@ void ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
       keep |= encapsulateConcat(rewriter, loc, &instruction);
     }
   }
+  return success();
 }
 
 void ONNXToSpatialPass::annotateWeightsConstants(func::FuncOp funcOp) const {
@@ -374,6 +493,47 @@ LogicalResult ONNXToSpatialPass::promoteConstantInputsToWeights(func::FuncOp fun
   return success();
 }
 
+void ONNXToSpatialPass::populateEmptyFunction(func::FuncOp funcOp) {
+  IRRewriter rewriter(&getContext());
+  IRMapping mapper;
+  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
+  if (!computes.empty())
+    return;
+  auto returnOp = llvm::cast<func::ReturnOp>(funcOp.getRegion().front().getTerminator());
+  rewriter.setInsertionPoint(returnOp);
+
+  SmallVector<Type> sourceTypes;
+  SmallVector<Location> sourceLoc;
+  for (auto source : funcOp.getArguments()) {
+    sourceTypes.push_back(source.getType());
+    sourceLoc.push_back(source.getLoc());
+  }
+
+  auto newCompute = spatial::SpatCompute::create(
+    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
+  auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
+  for (auto [bbArg, computeArg] : llvm::zip(BB->getArguments(), newCompute.getOperands()))
+    mapper.map(computeArg, bbArg);
+  newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sourceTypes.size()});
+  rewriter.setInsertionPointToEnd(BB);
+  for (Operation& inst : funcOp.getOps())
+    if (!isa<spatial::SpatCompute, func::ReturnOp>(&inst))
+      rewriter.clone(inst, mapper);
+
+  auto yield = spatial::SpatYieldOp::create(rewriter, funcOp.getLoc(), returnOp.getOperands());
+  for (size_t i = 0; i < yield.getNumOperands(); ++i)
+    yield.setOperand(i, mapper.lookupOrDefault(yield.getOperand(i)));
+
+  for (Operation& inst : llvm::make_early_inc_range(funcOp.getOps()))
+    if (!isa<spatial::SpatCompute, func::ReturnOp>(&inst)){
+      inst.dropAllUses();
+      rewriter.eraseOp(&inst);
+  }
+
+  for (auto [index, computeResult] : llvm::enumerate(newCompute.getResults()))
+    returnOp.setOperand(index, computeResult);
+}
+
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
 } // namespace onnx_mlir

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

@@ -7,11 +7,10 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include <algorithm>
-#include <cassert>
 
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -147,11 +146,11 @@ static Value buildPackedBias(bool hasBias,
   return arith::ConstantOp::create(rewriter, loc, packedBiasType, packedBiasAttr).getResult();
 }
 
-static SmallVector<Value> createIm2colRowComputes(Value x,
+static Value createIm2colRowComputes(Value x,
                                      RankedTensorType xType,
                                      RankedTensorType im2colType,
                                      RankedTensorType im2colRowType,
-                                                  RankedTensorType gemmInputRowType,
+                                     RankedTensorType gemmInputRowsType,
                                      int64_t batchSize,
                                      int64_t numChannelsIn,
                                      int64_t xHeight,
@@ -176,8 +175,8 @@ static SmallVector<Value> createIm2colRowComputes(Value x,
   auto elemType = xType.getElementType();
   constexpr size_t numInputs = 1;
   const int64_t packedNumRows = ceilIntegerDivide(numPatches, packFactor);
-  SmallVector<Type> resultTypes(packedNumRows, gemmInputRowType);
-  auto im2colComputeOp = createSpatCompute<numInputs>(rewriter, loc, resultTypes, {}, x, [&](Value xArg) {
+  auto im2colComputeOp =
+    createSpatCompute<numInputs>(rewriter, loc, TypeRange {gemmInputRowsType}, {}, x, [&](Value xArg) {
       Value paddedInput = xArg;
 
       // Pad input with zeros if needed:
@@ -285,23 +284,10 @@ static SmallVector<Value> createIm2colRowComputes(Value x,
         });
       }
 
-    SmallVector<Value> rowResults;
-    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> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-      rowResults.push_back(
-        tensor::ExtractSliceOp::create(rewriter, loc, gemmInputRowType, gemmInputRows, offsets, sizes, strides));
-    }
-    spatial::SpatYieldOp::create(rewriter, loc, rowResults);
+      spatial::SpatYieldOp::create(rewriter, loc, gemmInputRows);
     });
 
-  SmallVector<Value> rows;
-  rows.reserve(im2colComputeOp.getNumResults());
-  for (Value result : im2colComputeOp.getResults())
-    rows.push_back(result);
-  return rows;
+  return im2colComputeOp.getResult(0);
 }
 
 static Value createCollectedConvOutput(ValueRange gemmRows,
@@ -319,15 +305,12 @@ static Value createCollectedConvOutput(ValueRange gemmRows,
   auto collectComputeOp = createSpatCompute(rewriter, loc, convType, {}, gemmRows, [&](ValueRange gemmRowArgs) {
     Value gemmOut;
     if (packFactor == 1) {
-      gemmOut = gemmRowArgs.size() == 1 ? gemmRowArgs.front()
-                                        : tensor::ConcatOp::create(rewriter, loc, /*axis=*/0, gemmRowArgs).getResult();
+      gemmOut = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs);
     }
     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 = createSpatConcat(rewriter, loc, /*axis=*/0, gemmRowArgs);
       Value expandedOutput = tensor::ExpandShapeOp::create(rewriter,
                                                            loc,
                                                            expandedType,
@@ -386,11 +369,34 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
   auto wType = cast<RankedTensorType>(w.getType());
   auto outType = cast<RankedTensorType>(convOp.getY().getType());
 
-  assert("Only support static shapes" && xType.hasStaticShape() && wType.hasStaticShape() && outType.hasStaticShape());
-  assert("Only support 2D convolution" && xType.getRank() == 4);
-
-  // We need to understand what is group
-  assert("Only support group=1" && convOp.getGroup() == 1);
+  if (!xType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv input");
+    return failure();
+  }
+  if (!wType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv weight");
+    return failure();
+  }
+  if (!outType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv result");
+    return failure();
+  }
+  if (xType.getRank() != 4) {
+    pim::emitUnsupportedRankDiagnostic(convOp, "conv input", xType.getRank(), {4});
+    return failure();
+  }
+  if (wType.getRank() != 4) {
+    pim::emitUnsupportedRankDiagnostic(convOp, "conv weight", wType.getRank(), {4});
+    return failure();
+  }
+  if (outType.getRank() != 4) {
+    pim::emitUnsupportedRankDiagnostic(convOp, "conv result", outType.getRank(), {4});
+    return failure();
+  }
+  if (convOp.getGroup() != 1) {
+    convOp.emitOpError("only group=1 convolution is supported for Spatial lowering");
+    return failure();
+  }
 
   const int64_t batchSize = xType.getDimSize(0);
   const int64_t numChannelsIn = xType.getDimSize(1);
@@ -407,6 +413,19 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
   const auto dilationsAttr = convOp.getDilations();
   const auto padsAttr = convOp.getPads();
 
+  if (stridesAttr && stridesAttr->size() != 2) {
+    convOp.emitOpError("requires exactly two stride values for Spatial lowering");
+    return failure();
+  }
+  if (dilationsAttr && dilationsAttr->size() != 2) {
+    convOp.emitOpError("requires exactly two dilation values for Spatial lowering");
+    return failure();
+  }
+  if (padsAttr && padsAttr->size() != 4) {
+    convOp.emitOpError("requires exactly four pad values for 2D Spatial lowering");
+    return failure();
+  }
+
   const int64_t strideHeight = stridesAttr ? getI64FromArrayAttr(*stridesAttr, 0) : 1;
   const int64_t strideWidth = stridesAttr ? getI64FromArrayAttr(*stridesAttr, 1) : 1;
   const int64_t dilationHeight = dilationsAttr ? getI64FromArrayAttr(*dilationsAttr, 0) : 1;
@@ -447,6 +466,10 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
         padWidthBegin = totalPadW - padWidthEnd;
       }
     }
+    else if (autoPad != "NOTSET" && autoPad != "VALID") {
+      convOp.emitOpError() << "unsupported auto_pad value `" << autoPad << "` for Spatial lowering";
+      return failure();
+    }
     // "NOTSET" or "VALID" -> all pads stay 0
   }
 
@@ -509,14 +532,15 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
   //   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);
-  auto gemmOutputRowType =
-    RankedTensorType::get({1, effectiveMaxParallelPixels * numChannelsOut}, outType.getElementType());
-  SmallVector<Value> gemmInputRows = createIm2colRowComputes(x,
+  const int64_t packedNumRows = ceilIntegerDivide(numPatches, effectiveMaxParallelPixels);
+  auto gemmInputRowsType = RankedTensorType::get({packedNumRows, effectiveMaxParallelPixels * patchSize}, elemType);
+  auto gemmOutputRowsType =
+    RankedTensorType::get({packedNumRows, effectiveMaxParallelPixels * numChannelsOut}, outType.getElementType());
+  Value gemmInputRows = createIm2colRowComputes(x,
                                                 xType,
                                                 im2colType,
                                                 rowType,
-                                                             gemmInputRowType,
+                                                gemmInputRowsType,
                                                 batchSize,
                                                 numChannelsIn,
                                                 xHeight,
@@ -553,13 +577,10 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
   Value gemmC = buildPackedBias(
     hasB, gemmBias, biasMatrix, biasDenseAttr, outType, numChannelsOut, effectiveMaxParallelPixels, rewriter, loc);
 
-  SmallVector<Value> gemmRows;
-  gemmRows.reserve(gemmInputRows.size());
-  for (Value gemmInputRow : gemmInputRows) {
-    Value gemmRow = ONNXGemmOp::create(rewriter,
+  Value gemmRows = ONNXGemmOp::create(rewriter,
                                       loc,
-                                       gemmOutputRowType,
-                                       gemmInputRow,
+                                      gemmOutputRowsType,
+                                      gemmInputRows,
                                       gemmB,
                                       gemmC,
                                       rewriter.getF32FloatAttr(1.0f),
@@ -567,11 +588,9 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
                                       rewriter.getBoolAttr(false),
                                       rewriter.getBoolAttr(false))
                      .getY();
-    gemmRows.push_back(gemmRow);
-  }
 
   rewriter.replaceOp(convOp,
-                     createCollectedConvOutput(gemmRows,
+                     createCollectedConvOutput(ValueRange {gemmRows},
                                                convOp.getType(),
                                                gemmOutType,
                                                nhwcType,

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

@@ -5,7 +5,8 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -15,13 +16,6 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-static SmallVector<int64_t> computeRowMajorStrides(ArrayRef<int64_t> shape) {
-  SmallVector<int64_t> strides(shape.size(), 1);
-  for (int64_t i = static_cast<int64_t>(shape.size()) - 2; i >= 0; --i)
-    strides[i] = strides[i + 1] * shape[i + 1];
-  return strides;
-}
-
 static DenseElementsAttr getDenseConstantAttr(Value value) {
   if (auto constantOp = value.getDefiningOp<arith::ConstantOp>())
     return dyn_cast<DenseElementsAttr>(constantOp.getValue());

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

@@ -1,16 +1,16 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 #include "llvm/ADT/SmallVector.h"
 
-#include <cassert>
-
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -65,6 +65,66 @@ struct GemvToSpatialCompute : OpConversionPattern<ONNXGemmOp> {
                                 ConversionPatternRewriter& rewriter) const override;
 };
 
+struct GemmToSpatialComputeBatch : OpConversionPattern<ONNXGemmOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
+                                ONNXGemmOpAdaptor gemmOpAdaptor,
+                                ConversionPatternRewriter& rewriter) const override;
+};
+
+static SmallVector<Value> materializeBatchRowSlices(Value matrix,
+                                                    RankedTensorType matrixType,
+                                                    ConversionPatternRewriter& rewriter,
+                                                    Location loc) {
+  const int64_t numRows = matrixType.getDimSize(0);
+  auto rowType = RankedTensorType::get({1, matrixType.getDimSize(1)}, matrixType.getElementType());
+  SmallVector<Type> resultTypes(static_cast<size_t>(numRows), rowType);
+
+  auto buildRowSlices = [&](Value matrixArg) {
+    auto extractRowsOp = spatial::SpatExtractRowsOp::create(rewriter, loc, TypeRange(resultTypes), matrixArg);
+    return SmallVector<Value>(extractRowsOp->result_begin(), extractRowsOp->result_end());
+  };
+
+  auto cloneBatchInputChainIntoSliceCompute =
+    [&](Value rootInput, SmallVector<Operation*> chainOps, Value rootValue) -> SmallVector<Value> {
+    auto sliceCompute =
+      createSpatCompute<1>(rewriter, loc, TypeRange(resultTypes), {}, ValueRange {rootInput}, [&](Value input) {
+        Value transformedMatrix = input;
+        if (!chainOps.empty()) {
+          IRMapping mapper;
+          mapper.map(rootValue, input);
+          for (Operation* chainOp : chainOps)
+            rewriter.clone(*chainOp, mapper);
+          transformedMatrix = cast<Value>(mapper.lookup(matrix));
+        }
+        spatial::SpatYieldOp::create(rewriter, loc, buildRowSlices(transformedMatrix));
+      });
+    SmallVector<Value> rowSlices(sliceCompute->result_begin(), sliceCompute->result_end());
+    return rowSlices;
+  };
+
+  SmallVector<Operation*> chainOps;
+  Value rootValue = matrix;
+  while (Operation* definingOp = rootValue.getDefiningOp()) {
+    if (auto rootCompute = dyn_cast<spatial::SpatCompute>(definingOp)) {
+      SmallVector<Operation*> reversedChainOps(chainOps.rbegin(), chainOps.rend());
+      return cloneBatchInputChainIntoSliceCompute(
+        rootCompute.getResult(cast<OpResult>(rootValue).getResultNumber()), reversedChainOps, rootValue);
+    }
+
+    if (definingOp->getNumOperands() != 1)
+      break;
+    if (!isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, ONNXTransposeOp>(definingOp))
+      break;
+
+    chainOps.push_back(definingOp);
+    rootValue = definingOp->getOperand(0);
+  }
+
+  return buildRowSlices(matrix);
+}
+
 } // namespace
 
 LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
@@ -75,13 +135,23 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
   Value b = gemmOpAdaptor.getB();
   Value c = gemmOpAdaptor.getC();
 
-  assert("A should have been transposed already" && !gemmOpAdaptor.getTransA());
+  if (gemmOpAdaptor.getTransA()) {
+    gemmOp.emitOpError("requires transA=false before Gemm row decomposition");
+    return failure();
+  }
 
   bool hasC = !isa<ONNXNoneOp>(c.getDefiningOp());
 
   auto aType = cast<RankedTensorType>(a.getType());
   auto outType = cast<RankedTensorType>(gemmOp.getY().getType());
-  assert("Only support static shapes" && aType.hasStaticShape() && outType.hasStaticShape());
+  if (!aType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input A");
+    return failure();
+  }
+  if (!outType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm result");
+    return failure();
+  }
 
   const int64_t numOutRows = aType.getDimSize(0);
 
@@ -114,7 +184,14 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
       });
       cType = expandedType;
     }
-    assert("Only support rank 2 tensor for C" && cType.getRank() == 2);
+    if (!cType.hasStaticShape()) {
+      pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+      return failure();
+    }
+    if (cType.getRank() != 2) {
+      pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm bias", cType.getRank(), {1, 2});
+      return failure();
+    }
     cHasNumOutRows = cType.getDimSize(0) == numOutRows;
   }
 
@@ -138,8 +215,10 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
         auto cSliceType = RankedTensorType::get({1, cType.getDimSize(1)}, cType.getElementType());
         cSlice = tensor::ExtractSliceOp::create(rewriter, loc, cSliceType, c, offsets, sizes, strides).getResult();
       }
-      else
-        assert("C should be a vector" && isVectorShape(getTensorShape(c)));
+      else if (!isVectorShape(getTensorShape(c))) {
+        gemmOp.emitOpError("requires Gemm bias C to be vector-like when shared across decomposed rows");
+        return failure();
+      }
     }
 
     auto gemvOp = ONNXGemmOp::create(rewriter,
@@ -156,8 +235,7 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
   }
 
   auto concatComputeOp = createSpatCompute(rewriter, loc, gemmOp.getType(), {}, gemvOps, [&](ValueRange gemvOpsArgs) {
-    auto concatOp = tensor::ConcatOp::create(rewriter, loc, /*axis=*/0, gemvOpsArgs);
-    spatial::SpatYieldOp::create(rewriter, loc, concatOp.getResult());
+    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/0, gemvOpsArgs));
   });
 
   rewriter.replaceOp(gemmOp, concatComputeOp);
@@ -198,11 +276,28 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
       });
       cType = expandedType;
     }
-    assert("Only support rank 2 tensor for C" && cType.getRank() == 2);
+    if (!cType.hasStaticShape()) {
+      pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+      return failure();
+    }
+    if (cType.getRank() != 2) {
+      pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm bias", cType.getRank(), {1, 2});
+      return failure();
+    }
   }
 
-  assert("Only support static shapes" && aType.hasStaticShape() && bType.hasStaticShape()
-         && (!hasC || cType.hasStaticShape()) && outType.hasStaticShape());
+  if (!aType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input A");
+    return failure();
+  }
+  if (!bType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input B");
+    return failure();
+  }
+  if (!outType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm result");
+    return failure();
+  }
 
   if (!isVectorShape(aType.getShape()) || (hasC && !isVectorShape(cType.getShape())))
     // Not a gemv
@@ -281,19 +376,25 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
         weights.push_back(bTiles[outSliceId][coreId][aSliceId]);
 
       auto computeOp = createSpatCompute(
-        rewriter, gemmLoc, currOutHSliceType, weights, aHSlices[coreId], [&](ValueRange aHSlicesArgs) {
+        rewriter, gemmLoc, currOutHSliceType, weights, aHSlices[coreId], [&](ValueRange aHSlicesArgs) -> LogicalResult {
           SmallVector<Value> vmmOutputs;
           vmmOutputs.reserve(aHSlicesArgs.size());
           for (auto [aHSliceId, computeArg] : llvm::enumerate(aHSlicesArgs))
             vmmOutputs.push_back(
-              spatial::SpatWeightedVMMOp::create(rewriter, gemmLoc, currOutHSliceType, aHSliceId, computeArg));
-          assert(!vmmOutputs.empty() && "vmmOutputs must be non-empty");
+              spatial::SpatVMMOp::create(rewriter, gemmLoc, currOutHSliceType, aHSliceId, computeArg));
+          if (vmmOutputs.empty()) {
+            gemmOp.emitOpError("requires at least one non-empty slice when lowering tiled Gemm to Spatial VMMs");
+            return failure();
+          }
 
           Value partialVmmSum = sumTensors(vmmOutputs, rewriter);
           spatial::SpatYieldOp::create(rewriter, gemmLoc, partialVmmSum);
+          return success();
         });
+      if (failed(computeOp))
+        return failure();
 
-      partialResults.push_back(computeOp.getResult(0));
+      partialResults.push_back(computeOp->getResult(0));
     }
 
     if (hasC) {
@@ -313,8 +414,129 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
 
   auto concatComputeOp =
     createSpatCompute(rewriter, gemmLoc, gemmOp.getType(), {}, outHSlices, [&](ValueRange blockArgs) {
-      auto concatOp = tensor::ConcatOp::create(rewriter, gemmLoc, /*axis=*/1, blockArgs);
-      spatial::SpatYieldOp::create(rewriter, gemmLoc, concatOp.getResult());
+      spatial::SpatYieldOp::create(rewriter, gemmLoc, createSpatConcat(rewriter, gemmLoc, /*axis=*/1, blockArgs));
+    });
+
+  rewriter.replaceOp(gemmOp, concatComputeOp);
+  return success();
+}
+
+LogicalResult GemmToSpatialComputeBatch::matchAndRewrite(ONNXGemmOp gemmOp,
+                                                         ONNXGemmOpAdaptor gemmOpAdaptor,
+                                                         ConversionPatternRewriter& rewriter) const {
+  Location loc = gemmOp.getLoc();
+  Value a = gemmOpAdaptor.getA();
+  Value b = gemmOpAdaptor.getB();
+  Value c = gemmOpAdaptor.getC();
+
+  if (gemmOpAdaptor.getTransA()) {
+    gemmOp.emitOpError("requires transA=false before batch Gemm lowering");
+    return failure();
+  }
+
+  bool hasC = !isa<ONNXNoneOp>(c.getDefiningOp());
+
+  auto aType = cast<RankedTensorType>(a.getType());
+  auto bType = cast<RankedTensorType>(b.getType());
+  auto outType = cast<RankedTensorType>(gemmOp.getY().getType());
+  if (!aType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input A");
+    return failure();
+  }
+  if (!bType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input B");
+    return failure();
+  }
+  if (!outType.hasStaticShape()) {
+    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm result");
+    return failure();
+  }
+
+  const int64_t numOutRows = aType.getDimSize(0);
+  if (numOutRows <= 1)
+    return failure();
+
+  // Only handle the single-tile case: K <= crossbarSize and N <= crossbarSize
+  if (aType.getDimSize(1) > static_cast<int64_t>(crossbarSize.getValue())
+      || outType.getDimSize(1) > static_cast<int64_t>(crossbarSize.getValue()))
+    return failure();
+
+  auto scaledB = materializeScaledConstantTensor(b, gemmOpAdaptor.getAlpha().convertToFloat(), rewriter, loc);
+  if (failed(scaledB))
+    return failure();
+  b = *scaledB;
+  bType = cast<RankedTensorType>(b.getType());
+
+  if (gemmOpAdaptor.getTransB()) {
+    auto bShape = bType.getShape();
+    auto transposedType = bType.cloneWith(ArrayRef({bShape[1], bShape[0]}), bType.getElementType());
+    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0}));
+    bType = cast<RankedTensorType>(b.getType());
+  }
+  (void) bType;
+
+  Value sharedBias;
+  if (hasC) {
+    auto scaledC = materializeScaledConstantTensor(c, gemmOpAdaptor.getBeta().convertToFloat(), rewriter, loc);
+    if (failed(scaledC))
+      return failure();
+    c = *scaledC;
+    auto cType = cast<RankedTensorType>(c.getType());
+    if (cType.getRank() == 1) {
+      auto expandedType = RankedTensorType::get({1, cType.getDimSize(0)}, cType.getElementType());
+      c = tensor::ExpandShapeOp::create(rewriter,
+                                        loc,
+                                        expandedType,
+                                        c,
+                                        SmallVector<ReassociationIndices> {
+                                          {0, 1}
+      });
+      cType = cast<RankedTensorType>(c.getType());
+    }
+    if (!cType.hasStaticShape()) {
+      pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+      return failure();
+    }
+    if (cType.getRank() != 2) {
+      pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm bias", cType.getRank(), {1, 2});
+      return failure();
+    }
+    // Row-specific bias can't share a single template body; fall through to GemmToManyGemv
+    if (cType.getDimSize(0) == numOutRows && numOutRows > 1)
+      return failure();
+    if (cType.getDimSize(0) == 1 && cType.getDimSize(1) == 1)
+      c = broadcastToVector(c, outType.getDimSize(1), rewriter, loc);
+    sharedBias = c;
+  }
+
+  SmallVector<Value> aSlices = materializeBatchRowSlices(a, aType, rewriter, loc);
+  auto aSliceType = cast<RankedTensorType>(aSlices.front().getType());
+
+  auto outRowType = RankedTensorType::get({1, outType.getDimSize(1)}, outType.getElementType());
+  SmallVector<Type> resultTypes(static_cast<size_t>(numOutRows), outRowType);
+  SmallVector<Value> weights(static_cast<size_t>(numOutRows), b);
+
+  auto batchOp = spatial::SpatComputeBatch::create(rewriter,
+                                                   loc,
+                                                   TypeRange(resultTypes),
+                                                   rewriter.getI32IntegerAttr(static_cast<int32_t>(numOutRows)),
+                                                   ValueRange(weights),
+                                                   ValueRange(aSlices));
+
+  Block* body = rewriter.createBlock(
+    &batchOp.getBody(), batchOp.getBody().end(), TypeRange {aSliceType}, SmallVector<Location>(1, loc));
+  rewriter.setInsertionPointToEnd(body);
+
+  Value vmmResult = spatial::SpatVMMOp::create(rewriter, loc, outRowType, 0, body->getArgument(0)).getResult();
+  Value laneResult = vmmResult;
+  if (sharedBias)
+    laneResult = spatial::SpatVAddOp::create(rewriter, loc, outRowType, vmmResult, sharedBias).getResult();
+  spatial::SpatYieldOp::create(rewriter, loc, laneResult);
+
+  rewriter.setInsertionPointAfter(batchOp);
+  SmallVector<Value> laneResults(batchOp->result_begin(), batchOp->result_end());
+  auto concatComputeOp = createSpatCompute(rewriter, loc, gemmOp.getType(), {}, laneResults, [&](ValueRange args) {
+    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/0, args));
   });
 
   rewriter.replaceOp(gemmOp, concatComputeOp);
@@ -322,6 +544,7 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
 }
 
 void populateGemmPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.insert<GemmToSpatialComputeBatch>(ctx, PatternBenefit(2));
   patterns.insert<GemmToManyGemv>(ctx);
   patterns.insert<GemvToSpatialCompute>(ctx);
 }

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

@@ -5,7 +5,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -232,9 +232,7 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
       }));
     }
 
-    Value result = batchResults.size() == 1
-                   ? batchResults.front()
-                   : tensor::ConcatOp::create(rewriter, loc, /*axis=*/0, batchResults).getResult();
+    Value result = createSpatConcat(rewriter, loc, /*axis=*/0, batchResults);
     rewriter.replaceOp(matmulOp, result);
     return success();
   }

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

@@ -5,7 +5,7 @@
 
 #include <algorithm>
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -100,8 +100,7 @@ static Value buildReduceMeanKeepdims(Value input,
   for (Value slice : slices)
     reducedSlices.push_back(buildReduceMeanKeepdims(slice, reducedAxes, axis + 1, leafType, rewriter, loc));
 
-  return reducedSlices.size() == 1 ? reducedSlices.front()
-                                   : tensor::ConcatOp::create(rewriter, loc, axis, reducedSlices).getResult();
+  return createSpatConcat(rewriter, loc, axis, reducedSlices);
 }
 
 static Value squeezeReducedAxes(Value keepdimsValue,

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp

@@ -6,13 +6,12 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include <algorithm>
-#include <cassert>
 #include <optional>
 #include <type_traits>
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -31,11 +30,14 @@ static int64_t getOptionalI64(std::optional<ArrayAttrT> arrayAttr, size_t index,
   return arrayAttr ? getI64(*arrayAttr, index) : defaultValue;
 }
 
-static Value concatAlongAxis(ConversionPatternRewriter& rewriter, Location loc, int64_t axis, ArrayRef<Value> values) {
-  assert(!values.empty() && "Expected at least one value to concatenate.");
-  if (values.size() == 1)
-    return values.front();
-  return tensor::ConcatOp::create(rewriter, loc, axis, values);
+template <typename PoolOp>
+static FailureOr<Value>
+concatAlongAxis(ConversionPatternRewriter& rewriter, Location loc, PoolOp poolOp, int64_t axis, ArrayRef<Value> values) {
+  if (values.empty()) {
+    poolOp.emitOpError("failed to build pooled output because an intermediate concatenation input list was empty");
+    return failure();
+  }
+  return createSpatConcat(rewriter, loc, axis, values);
 }
 
 static Value materializeContiguousTile(ConversionPatternRewriter& rewriter, Location loc, Value tile) {
@@ -53,8 +55,12 @@ static Value materializeContiguousTile(ConversionPatternRewriter& rewriter, Loca
 }
 
 template <typename ReduceOp>
-static Value reduceWindowValues(ConversionPatternRewriter& rewriter, Location loc, ArrayRef<Value> windowValues) {
-  assert(!windowValues.empty() && "Expected at least one pool window value.");
+static FailureOr<Value>
+reduceWindowValues(ConversionPatternRewriter& rewriter, Location loc, Operation* op, ArrayRef<Value> windowValues) {
+  if (windowValues.empty()) {
+    op->emitOpError("pool window resolved to zero valid elements");
+    return failure();
+  }
 
   Value reduced = windowValues.front();
   for (Value value : windowValues.drop_front())
@@ -62,9 +68,12 @@ static Value reduceWindowValues(ConversionPatternRewriter& rewriter, Location lo
   return reduced;
 }
 
-static Value
-scaleAverageWindow(ConversionPatternRewriter& rewriter, Location loc, Value reducedWindow, int64_t divisor) {
-  assert(divisor > 0 && "AveragePool divisor must be positive.");
+static FailureOr<Value>
+scaleAverageWindow(ConversionPatternRewriter& rewriter, Location loc, Operation* op, Value reducedWindow, int64_t divisor) {
+  if (divisor <= 0) {
+    op->emitOpError("AveragePool divisor must be positive");
+    return failure();
+  }
   if (divisor == 1)
     return reducedWindow;
 
@@ -72,7 +81,7 @@ scaleAverageWindow(ConversionPatternRewriter& rewriter, Location loc, Value redu
   double scale = 1.0 / static_cast<double>(divisor);
   auto scaleAttr = DenseElementsAttr::get(tileType, rewriter.getFloatAttr(tileType.getElementType(), scale));
   Value scaleTensor = arith::ConstantOp::create(rewriter, loc, tileType, scaleAttr);
-  return spatial::SpatVMulOp::create(rewriter, loc, tileType, reducedWindow, scaleTensor);
+  return spatial::SpatVMulOp::create(rewriter, loc, tileType, reducedWindow, scaleTensor).getResult();
 }
 
 template <typename PoolOp>
@@ -211,28 +220,45 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
                 if (windowValues.empty())
                   return rewriter.notifyMatchFailure(poolOp, "pool window resolved to zero valid elements.");
 
-                Value reducedWindow = reduceWindowValues<ReduceOp>(rewriter, loc, windowValues);
+                auto reducedWindow = reduceWindowValues<ReduceOp>(rewriter, loc, poolOp, windowValues);
+                if (failed(reducedWindow))
+                  return failure();
+                Value reducedWindowValue = *reducedWindow;
                 if constexpr (std::is_same_v<PoolOp, ONNXAveragePoolOp>) {
                   const bool countIncludePad = poolOp.getCountIncludePad() == 1;
                   const int64_t divisor =
                     countIncludePad ? kernelHeight * kernelWidth : static_cast<int64_t>(windowValues.size());
-                  reducedWindow = scaleAverageWindow(rewriter, loc, reducedWindow, divisor);
+                  auto scaledWindow = scaleAverageWindow(rewriter, loc, poolOp, reducedWindowValue, divisor);
+                  if (failed(scaledWindow))
+                    return failure();
+                  reducedWindowValue = *scaledWindow;
                 }
 
-                outputChannelTiles.push_back(reducedWindow);
+                outputChannelTiles.push_back(reducedWindowValue);
               }
 
-              rowPixels.push_back(concatAlongAxis(rewriter, loc, /*axis=*/1, outputChannelTiles));
+              auto rowPixel = concatAlongAxis(rewriter, loc, poolOp, /*axis=*/1, outputChannelTiles);
+              if (failed(rowPixel))
+                return failure();
+              rowPixels.push_back(*rowPixel);
             }
 
-            rows.push_back(concatAlongAxis(rewriter, loc, /*axis=*/3, rowPixels));
+            auto row = concatAlongAxis(rewriter, loc, poolOp, /*axis=*/3, rowPixels);
+            if (failed(row))
+              return failure();
+            rows.push_back(*row);
           }
 
-          batchResults.push_back(concatAlongAxis(rewriter, loc, /*axis=*/2, rows));
+          auto batchResult = concatAlongAxis(rewriter, loc, poolOp, /*axis=*/2, rows);
+          if (failed(batchResult))
+            return failure();
+          batchResults.push_back(*batchResult);
         }
 
-        Value pooledOutput = concatAlongAxis(rewriter, loc, /*axis=*/0, batchResults);
-        spatial::SpatYieldOp::create(rewriter, loc, pooledOutput);
+        auto pooledOutput = concatAlongAxis(rewriter, loc, poolOp, /*axis=*/0, batchResults);
+        if (failed(pooledOutput))
+          return failure();
+        spatial::SpatYieldOp::create(rewriter, loc, *pooledOutput);
         return success();
       });
     if (failed(computeOp))

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Relu.cpp

@@ -1,6 +1,6 @@
 #include "mlir/Transforms/DialectConversion.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Sigmoid.cpp

@@ -1,6 +1,6 @@
 #include "mlir/Transforms/DialectConversion.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Softmax.cpp

@@ -1,7 +1,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -47,8 +47,7 @@ buildSoftmax(Value input, int64_t softmaxAxis, int64_t axis, ConversionPatternRe
   for (Value slice : slices)
     rebuiltSlices.push_back(buildSoftmax(slice, softmaxAxis, axis + 1, rewriter, loc));
 
-  return rebuiltSlices.size() == 1 ? rebuiltSlices.front()
-                                   : tensor::ConcatOp::create(rewriter, loc, axis, rebuiltSlices).getResult();
+  return createSpatConcat(rewriter, loc, axis, rebuiltSlices);
 }
 
 struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp

@@ -1,7 +1,8 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/PatternMatch.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
 using namespace mlir;
@@ -17,7 +18,7 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
     auto inputs = adaptor.getInputs();
     int64_t axis = adaptor.getAxis();
 
-    rewriter.replaceOpWithNewOp<tensor::ConcatOp>(maxpoolOp, axis, inputs);
+    rewriter.replaceOp(maxpoolOp, createSpatConcat(rewriter, maxpoolOp.getLoc(), axis, inputs));
 
     return success();
   }

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Gather.cpp

@@ -5,7 +5,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -49,7 +49,7 @@ static Value concatGatherSlices(Value data,
   }
   if (slices.empty())
     return {};
-  return slices.size() == 1 ? slices.front() : tensor::ConcatOp::create(rewriter, loc, axis, slices).getResult();
+  return createSpatConcat(rewriter, loc, axis, slices);
 }
 
 static Value addLeadingGatherDim(Value value, int64_t axis, ConversionPatternRewriter& rewriter, Location loc) {
@@ -130,9 +130,7 @@ struct Gather : OpConversionPattern<ONNXGatherOp> {
                                    return failure();
                                  rows.push_back(addLeadingGatherDim(gatheredRow, axis, rewriter, loc));
                                }
-                               result = rows.size() == 1
-                                        ? rows.front()
-                                        : tensor::ConcatOp::create(rewriter, loc, /*axis=*/axis, rows).getResult();
+                               result = createSpatConcat(rewriter, loc, /*axis=*/axis, rows);
                              }
                              else {
                                return failure();

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp

@@ -5,7 +5,7 @@
 
 #include <algorithm>
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -50,7 +50,7 @@ static Value buildNearestResize(Value input,
     slices.push_back(buildNearestResize(slice, inputShape, outputShape, axis + 1, rewriter, loc));
   }
 
-  return slices.size() == 1 ? slices.front() : tensor::ConcatOp::create(rewriter, loc, axis, slices).getResult();
+  return createSpatConcat(rewriter, loc, axis, slices);
 }
 
 struct Resize : OpConversionPattern<ONNXResizeOp> {

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Split.cpp

@@ -1,7 +1,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -23,7 +23,10 @@ static Value extractSliceAt(
     sizes.push_back(rewriter.getIndexAttr(dim));
   offsets[axis] = rewriter.getIndexAttr(offset);
   sizes[axis] = rewriter.getIndexAttr(size);
-  return tensor::ExtractSliceOp::create(rewriter, loc, input, offsets, sizes, strides);
+  SmallVector<int64_t> resultShape(inputType.getShape());
+  resultShape[axis] = size;
+  auto resultType = RankedTensorType::get(resultShape, inputType.getElementType());
+  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, input, offsets, sizes, strides);
 }
 
 struct Split : OpConversionPattern<ONNXSplitOp> {
@@ -49,12 +52,7 @@ struct Split : OpConversionPattern<ONNXSplitOp> {
       if (!resultType || !resultType.hasStaticShape())
         return failure();
       int64_t sliceSize = resultType.getShape()[axis];
-      auto computeOp =
-        createSpatCompute<1>(rewriter, splitOp.getLoc(), TypeRange {resultType}, {}, adaptor.getInput(), [&](Value x) {
-          Value output = extractSliceAt(x, axis, offset, sliceSize, rewriter, splitOp.getLoc());
-          spatial::SpatYieldOp::create(rewriter, splitOp.getLoc(), output);
-        });
-      outputs.push_back(computeOp.getResult(0));
+      outputs.push_back(extractSliceAt(adaptor.getInput(), axis, offset, sliceSize, rewriter, splitOp.getLoc()));
       offset += sliceSize;
     }
 

src/PIM/Conversion/SpatialToPim/CMakeLists.txt

@@ -5,6 +5,7 @@ add_public_tablegen_target(SpatialToPimIncGen)
 add_pim_library(OMSpatialToPim
   SpatialToPimPass.cpp
   Common.cpp
+  Patterns.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -7,23 +7,12 @@
 #include <cstddef>
 
 #include "Common.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace llvm;
 using namespace mlir;
 
 namespace onnx_mlir {
 
-namespace {
-
-IntegerAttr getRequiredI32Attr(Builder& builder, Operation* op, llvm::StringRef attrName) {
-  auto attr = op->getAttrOfType<IntegerAttr>(attrName);
-  assert(attr && "required precomputed channel attr is missing");
-  return IntegerAttr::get(builder.getI32Type(), attr.getInt());
-}
-
-} // namespace
-
 size_t getSliceActualOffset(tensor::ExtractSliceOp& sliceOp, ShapedType& inputShape) {
   /*
     EXAMPLE RUN:
@@ -74,37 +63,6 @@ IntegerAttr getTensorSizeInBytesAttr(Builder& builder, mlir::Value value) {
   return builder.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(value.getType()))));
 }
 
-IntegerAttr getSpatialChannelSourceCoreIdAttr(Builder& builder, mlir::Value channel) {
-  auto channelNewOp = channel.getDefiningOp<spatial::SpatChannelNewOp>();
-  assert(channelNewOp && "spatial channel value must come from spat.channel_new");
-  return getRequiredI32Attr(builder, channelNewOp, kChannelSourceCoreIdAttrName);
-}
-
-IntegerAttr getSpatialChannelTargetCoreIdAttr(Builder& builder, mlir::Value channel) {
-  auto channelNewOp = channel.getDefiningOp<spatial::SpatChannelNewOp>();
-  assert(channelNewOp && "spatial channel value must come from spat.channel_new");
-  return getRequiredI32Attr(builder, channelNewOp, kChannelTargetCoreIdAttrName);
-}
-
-bool hasSpatialChannelSourceCoreIdAttr(mlir::Value channel) {
-  auto channelNewOp = channel.getDefiningOp<spatial::SpatChannelNewOp>();
-  return channelNewOp && channelNewOp->hasAttr(kChannelSourceCoreIdAttrName);
-}
-
-bool hasSpatialChannelTargetCoreIdAttr(mlir::Value channel) {
-  auto channelNewOp = channel.getDefiningOp<spatial::SpatChannelNewOp>();
-  return channelNewOp && channelNewOp->hasAttr(kChannelTargetCoreIdAttrName);
-}
-
-mlir::Value
-createPimReceiveFromSpatialChannel(PatternRewriter& rewriter, Location loc, mlir::Value output, mlir::Value channel) {
-  mlir::Value outputBuffer = getBestOutputTensorFromOperandsOrAllocate(rewriter, output.getDefiningOp());
-  auto sizeAttr = getTensorSizeInBytesAttr(rewriter, output);
-  auto sourceCoreIdAttr = getSpatialChannelSourceCoreIdAttr(rewriter, channel);
-  return pim::PimReceiveOp::create(rewriter, loc, outputBuffer.getType(), outputBuffer, sizeAttr, sourceCoreIdAttr)
-    .getOutput();
-}
-
 Operation* getEarliestUserWithinBlock(mlir::Value value) {
   auto users = value.getUsers();
 
@@ -137,7 +95,7 @@ bool hasLaterUserInBlock(mlir::Value value, Operation* operation) {
   return false;
 }
 
-mlir::Value getBestOutputTensorFromOperandsOrAllocate(PatternRewriter& rewriter, Operation* operation) {
+mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Operation* operation) {
   assert("Only support operations with a single result" && operation->getNumResults() == 1);
   mlir::Value result = operation->getResult(0);
   auto resultType = result.getType();

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -2,16 +2,10 @@
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
-#include "llvm/ADT/StringRef.h"
-
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 namespace onnx_mlir {
 
-inline constexpr llvm::StringLiteral kChannelSourceCoreIdAttrName = "precomp_source_core_id";
-inline constexpr llvm::StringLiteral kChannelTargetCoreIdAttrName = "precomp_target_core_id";
-
 /**
  * \brief Get the offset of the ExtractSliceOp based on its static offsets and
  * its static tensor input.
@@ -30,17 +24,6 @@ size_t getShapedTypeSizeInBytes(mlir::ShapedType shapedType);
 
 mlir::IntegerAttr getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Value value);
 
-mlir::IntegerAttr getSpatialChannelSourceCoreIdAttr(mlir::Builder& builder, mlir::Value channel);
-
-mlir::IntegerAttr getSpatialChannelTargetCoreIdAttr(mlir::Builder& builder, mlir::Value channel);
-
-bool hasSpatialChannelSourceCoreIdAttr(mlir::Value channel);
-
-bool hasSpatialChannelTargetCoreIdAttr(mlir::Value channel);
-
-mlir::Value createPimReceiveFromSpatialChannel(
-  mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value output, mlir::Value channel);
-
 template <class T>
 size_t rangeLength(const mlir::iterator_range<T> range) {
   return std::distance(range.begin(), range.end());
@@ -58,7 +41,7 @@ mlir::Operation* getEarliestUserWithinBlock(mlir::Value value);
 
 mlir::SmallVector<mlir::Value> getOpOperandsSortedByUses(mlir::Operation* operation);
 
-mlir::Value getBestOutputTensorFromOperandsOrAllocate(mlir::PatternRewriter& rewriter, mlir::Operation* operation);
+mlir::Value getBestOutputTensorFromOperandsOrAllocate(mlir::RewriterBase& rewriter, mlir::Operation* operation);
 
 inline mlir::tensor::EmptyOp
 createEmptyTensorFromShaped(mlir::IRRewriter& rewriter, mlir::Location loc, mlir::ShapedType shapedType) {

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -0,0 +1,424 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/LogicalResult.h"
+
+#include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+
+static std::optional<unsigned> getDirectComputeInputIndex(Operation* owner, unsigned operandNumber) {
+  if (auto compute = dyn_cast<spatial::SpatCompute>(owner)) {
+    unsigned inputCount = compute.getInputs().size();
+    if (inputCount == 0)
+      return std::nullopt;
+
+    unsigned inputBegin = compute->getNumOperands() - inputCount;
+    if (operandNumber < inputBegin)
+      return std::nullopt;
+    return operandNumber - inputBegin;
+  }
+
+  if (auto computeBatch = dyn_cast<spatial::SpatComputeBatch>(owner)) {
+    unsigned inputCount = computeBatch.getInputs().size();
+    if (inputCount == 0)
+      return std::nullopt;
+
+    unsigned inputBegin = computeBatch->getNumOperands() - inputCount;
+    if (operandNumber < inputBegin)
+      return std::nullopt;
+    return operandNumber - inputBegin;
+  }
+
+  return std::nullopt;
+}
+
+struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::ExtractSliceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::tensor::ExtractSliceOp extractSliceOp, PatternRewriter& rewriter) const override {
+    if (!isa<func::FuncOp>(extractSliceOp->getParentOp()))
+      return failure();
+
+    for (auto& uses : extractSliceOp->getUses()) {
+      if (isa<spatial::SpatCompute>(uses.getOwner())) {
+        if (!getDirectComputeInputIndex(uses.getOwner(), uses.getOperandNumber()))
+          return failure();
+      }
+      else if (isa_and_present<func::FuncOp>(uses.getOwner()->getParentOp())) {
+        return failure();
+      }
+    }
+
+    llvm::DenseMap<Operation*, Value> mapSpatToExtract;
+
+    for (auto& uses : llvm::make_early_inc_range(extractSliceOp->getUses())) {
+
+      if (auto spatCompute = dyn_cast<spatial::SpatCompute>(uses.getOwner())) {
+        auto inputIndex = getDirectComputeInputIndex(spatCompute, uses.getOperandNumber());
+        if (!inputIndex)
+          return failure();
+        auto BBArgIndex = *inputIndex;
+        auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+
+        if (BBArgValue.use_empty())
+          continue;
+
+        rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+        if (!mapSpatToExtract.contains(spatCompute.getOperation())) {
+          auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
+          mapSpatToExtract.insert({spatCompute.getOperation(), newExtractSlice->getResult(0)});
+        }
+
+        rewriter.startOpModification(spatCompute.getOperation());
+        BBArgValue.replaceAllUsesWith(mapSpatToExtract[spatCompute.getOperation()]);
+        spatCompute.getInputsMutable().erase(BBArgIndex);
+        spatCompute.getBody().front().eraseArgument(BBArgIndex);
+        rewriter.finalizeOpModification(spatCompute.getOperation());
+      }
+      else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(uses.getOwner())) {
+        auto inputIndex = getDirectComputeInputIndex(spatComputeBatch, uses.getOperandNumber());
+        if (!inputIndex)
+          return failure();
+        auto BBArgIndex = *inputIndex;
+        auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+
+        if (BBArgValue.use_empty())
+          continue;
+
+        rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+        if (!mapSpatToExtract.contains(spatComputeBatch.getOperation())) {
+          auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
+          mapSpatToExtract.insert({spatComputeBatch.getOperation(), newExtractSlice->getResult(0)});
+        }
+
+        rewriter.startOpModification(spatComputeBatch.getOperation());
+        BBArgValue.replaceAllUsesWith(mapSpatToExtract[spatComputeBatch.getOperation()]);
+        spatComputeBatch.getInputsMutable().erase(BBArgIndex);
+        spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
+        rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+      }
+      else {
+        {
+          if (auto spatCompute = uses.getOwner()->getParentOfType<spatial::SpatCompute>()) {
+            rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+            if (!mapSpatToExtract.contains(spatCompute.getOperation())) {
+              auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
+              mapSpatToExtract.insert({spatCompute.getOperation(), newExtractSlice->getResult(0)});
+            }
+
+            rewriter.startOpModification(spatCompute.getOperation());
+            uses.set(mapSpatToExtract[spatCompute.getOperation()]);
+            rewriter.finalizeOpModification(spatCompute.getOperation());
+          }
+          else if (auto spatComputeBatch = uses.getOwner()->getParentOfType<spatial::SpatComputeBatch>()) {
+            rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+            if (!mapSpatToExtract.contains(spatComputeBatch.getOperation())) {
+              auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
+              mapSpatToExtract.insert({spatComputeBatch.getOperation(), newExtractSlice->getResult(0)});
+            }
+
+            rewriter.startOpModification(spatComputeBatch.getOperation());
+            uses.set(mapSpatToExtract[spatComputeBatch.getOperation()]);
+            rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+          }
+        }
+      }
+    }
+
+    rewriter.eraseOp(extractSliceOp);
+    return success();
+  }
+};
+
+struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::ConstantOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::arith::ConstantOp constantOp, PatternRewriter& rewriter) const override {
+    static int i = 0;
+    Location loc = constantOp.getLoc();
+
+    if (hasWeightAlways(constantOp))
+      return failure();
+
+    if (!isa<func::FuncOp>(constantOp->getParentOp()))
+      return failure();
+
+    if (llvm::all_of(constantOp->getUsers(), [](Operation* op) {
+          if (isa<spatial::SpatCompute>(op))
+            return false;
+          if (isa<func::FuncOp>(op->getParentOp()))
+            return true;
+          return false;
+        }))
+      return failure();
+
+    rewriter.setInsertionPoint(constantOp->getParentOfType<func::FuncOp>());
+
+    auto constRankedTensorType = llvm::dyn_cast<mlir::RankedTensorType>(constantOp.getType());
+
+    if (constRankedTensorType) {
+      mlir::MemRefType memRefType =
+        mlir::MemRefType::get(constRankedTensorType.getShape(), constRankedTensorType.getElementType());
+      std::string argName = "const_" + std::to_string(i++);
+      memref::GlobalOp::create(rewriter,
+                               loc,
+                               rewriter.getStringAttr(argName),
+                               rewriter.getStringAttr("private"),
+                               TypeAttr::get(memRefType),
+                               constantOp.getValueAttr(),
+                               rewriter.getUnitAttr(),
+                               {});
+
+      llvm::DenseMap<Operation*, Value> mapSpatComputeToConst;
+
+      for (auto& constUses : llvm::make_early_inc_range(constantOp->getUses())) {
+        auto constUsers = constUses.getOwner();
+
+        if (auto spatCompute = llvm::dyn_cast<spatial::SpatCompute>(constUsers)) {
+          auto inputIndex = getDirectComputeInputIndex(spatCompute, constUses.getOperandNumber());
+          if (!inputIndex)
+            return failure();
+          auto BBArgIndex = *inputIndex;
+          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+          if (!mapSpatComputeToConst.contains(spatCompute.getOperation())) {
+            auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+            auto toTensor = bufferization::ToTensorOp::create(
+              rewriter, loc, constRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+            mapSpatComputeToConst.insert({spatCompute.getOperation(), toTensor.getResult()});
+          }
+
+          rewriter.startOpModification(spatCompute.getOperation());
+          BBArgValue.replaceAllUsesWith(mapSpatComputeToConst[spatCompute.getOperation()]);
+          spatCompute.getInputsMutable().erase(BBArgIndex);
+          spatCompute.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatCompute.getOperation());
+        }
+        else if (auto spatComputeBatch = llvm::dyn_cast<spatial::SpatComputeBatch>(constUsers)) {
+          auto inputIndex = getDirectComputeInputIndex(spatComputeBatch, constUses.getOperandNumber());
+          if (!inputIndex)
+            return failure();
+          auto BBArgIndex = *inputIndex;
+          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+          if (!mapSpatComputeToConst.contains(spatComputeBatch.getOperation())) {
+            auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+            auto toTensor = bufferization::ToTensorOp::create(
+              rewriter, loc, constRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+            mapSpatComputeToConst.insert({spatComputeBatch.getOperation(), toTensor.getResult()});
+          }
+
+          rewriter.startOpModification(spatComputeBatch.getOperation());
+          BBArgValue.replaceAllUsesWith(mapSpatComputeToConst[spatComputeBatch.getOperation()]);
+          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
+          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+        }
+        else {
+          {
+
+            if (auto spatCompute = constUses.getOwner()->getParentOfType<spatial::SpatCompute>()) {
+              rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+              if (!mapSpatComputeToConst.contains(spatCompute.getOperation())) {
+                auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+                auto toTensor = bufferization::ToTensorOp::create(
+                  rewriter, loc, constRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+                mapSpatComputeToConst.insert({spatCompute.getOperation(), toTensor.getResult()});
+              }
+
+              rewriter.startOpModification(spatCompute.getOperation());
+              constUses.set(mapSpatComputeToConst[spatCompute.getOperation()]);
+              rewriter.finalizeOpModification(spatCompute.getOperation());
+            }
+            else if (auto spatComputeBatch = constUses.getOwner()->getParentOfType<spatial::SpatComputeBatch>()) {
+              rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+              if (!mapSpatComputeToConst.contains(spatComputeBatch.getOperation())) {
+                auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+                auto toTensor = bufferization::ToTensorOp::create(
+                  rewriter, loc, constRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+                mapSpatComputeToConst.insert({spatComputeBatch.getOperation(), toTensor.getResult()});
+              }
+
+              rewriter.startOpModification(spatComputeBatch.getOperation());
+              constUses.set(mapSpatComputeToConst[spatComputeBatch.getOperation()]);
+              rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+            }
+          }
+        }
+      }
+    }
+    else if (constantOp.getType().isIntOrIndexOrFloat()) {
+      llvm::DenseMap<Operation*, Value> mapSpatComputeToConst;
+
+      for (auto& constUses : llvm::make_early_inc_range(constantOp->getUses())) {
+        auto constUsers = constUses.getOwner();
+
+        if (auto spatCompute = llvm::dyn_cast<spatial::SpatCompute>(constUsers)) {
+          auto inputIndex = getDirectComputeInputIndex(spatCompute, constUses.getOperandNumber());
+          if (!inputIndex)
+            return failure();
+          auto BBArgIndex = *inputIndex;
+          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+          auto newConst = rewriter.clone(*constantOp);
+
+          rewriter.startOpModification(spatCompute.getOperation());
+          BBArgValue.replaceAllUsesWith(newConst->getResult(0));
+          spatCompute.getInputsMutable().erase(BBArgIndex);
+          spatCompute.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatCompute.getOperation());
+        }
+        else if (auto spatComputeBatch = llvm::dyn_cast<spatial::SpatComputeBatch>(constUsers)) {
+          auto inputIndex = getDirectComputeInputIndex(spatComputeBatch, constUses.getOperandNumber());
+          if (!inputIndex)
+            return failure();
+          auto BBArgIndex = *inputIndex;
+          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+          auto newConst = rewriter.clone(*constantOp);
+
+          rewriter.startOpModification(spatComputeBatch.getOperation());
+          BBArgValue.replaceAllUsesWith(newConst->getResult(0));
+          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
+          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+        }
+        else if (auto parent = constUsers->getParentOfType<spatial::SpatCompute>()) {
+          if (!mapSpatComputeToConst.contains(parent)) {
+            rewriter.setInsertionPoint(&parent.getBody().front().front());
+            auto newConst = rewriter.clone(*constantOp);
+            mapSpatComputeToConst.insert({parent.getOperation(), newConst->getResult(0)});
+          }
+          constUses.set(mapSpatComputeToConst[parent.getOperation()]);
+        }
+        else {
+          auto batchParent = constUsers->getParentOfType<spatial::SpatComputeBatch>();
+          assert(batchParent && "Global Constant used direcly not within a compute");
+          if (!mapSpatComputeToConst.contains(batchParent.getOperation())) {
+            rewriter.setInsertionPoint(&batchParent.getBody().front().front());
+            auto newConst = rewriter.clone(*constantOp);
+            mapSpatComputeToConst.insert({batchParent.getOperation(), newConst->getResult(0)});
+          }
+          constUses.set(mapSpatComputeToConst[batchParent.getOperation()]);
+        }
+      }
+    }
+    rewriter.eraseOp(constantOp);
+    return success();
+  }
+};
+
+struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::func::FuncOp funcOp, PatternRewriter& rewriter) const override {
+
+    if (funcOp.getArguments().empty())
+      return failure();
+
+    if (llvm::all_of(funcOp.getArguments(),
+                     [](mlir::BlockArgument blockArgument) { return blockArgument.use_empty(); }))
+      return failure();
+
+    Location loc = funcOp.getLoc();
+
+    for (auto [index, arg] : llvm::enumerate(funcOp.getArguments())) {
+      if (arg.getUses().empty())
+        continue;
+
+      rewriter.setInsertionPoint(funcOp.getOperation());
+
+      assert(isa<mlir::RankedTensorType>(arg.getType()));
+
+      auto argRankedTensorType = llvm::dyn_cast<mlir::RankedTensorType>(arg.getType());
+      mlir::MemRefType memRefType =
+        mlir::MemRefType::get(argRankedTensorType.getShape(), argRankedTensorType.getElementType());
+
+      std::string argName = "arg_" + std::to_string(index);
+
+      memref::GlobalOp::create(rewriter,
+                               loc,
+                               rewriter.getStringAttr(argName),
+                               rewriter.getStringAttr("private"),
+                               TypeAttr::get(memRefType),
+                               {},
+                               {},
+                               {});
+
+      for (auto& argUses : llvm::make_early_inc_range(arg.getUses())) {
+        auto argUser = argUses.getOwner();
+        if (auto spatCompute = dyn_cast<spatial::SpatCompute>(argUser)) {
+          auto inputIndex = getDirectComputeInputIndex(spatCompute, argUses.getOperandNumber());
+          if (!inputIndex)
+            return failure();
+          auto BBArgIndex = *inputIndex;
+          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
+          auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+          auto toTensor = bufferization::ToTensorOp::create(
+            rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+
+          rewriter.startOpModification(spatCompute.getOperation());
+          BBArgValue.replaceAllUsesWith(toTensor);
+          spatCompute.getInputsMutable().erase(BBArgIndex);
+          spatCompute.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatCompute.getOperation());
+        }
+        else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(argUser)) {
+          auto inputIndex = getDirectComputeInputIndex(spatComputeBatch, argUses.getOperandNumber());
+          if (!inputIndex)
+            return failure();
+          auto BBArgIndex = *inputIndex;
+          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+          rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
+          auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+          auto toTensor = bufferization::ToTensorOp::create(
+            rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+
+          rewriter.startOpModification(spatComputeBatch.getOperation());
+          BBArgValue.replaceAllUsesWith(toTensor);
+          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
+          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
+          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+        }
+        else {
+          rewriter.setInsertionPoint(argUser);
+          auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
+          auto toTensor = bufferization::ToTensorOp::create(
+            rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+          rewriter.startOpModification(argUser);
+          argUses.set(toTensor);
+          rewriter.finalizeOpModification(argUser);
+        }
+      }
+    }
+
+    return success();
+  }
+};
+
+} // namespace
+void populateGlobalTensorToMemrefPatterns(RewritePatternSet& patterns) {
+  patterns.add<MoveExtractSliceIntoCompute, FuncOpArgToGlobalMemoryPattern, ArithConstToGlobalMemoryPattern>(
+    patterns.getContext());
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Patterns.hpp

@@ -0,0 +1,10 @@
+#pragma once
+
+#include "mlir/IR/PatternMatch.h"
+
+
+namespace onnx_mlir {
+
+void populateGlobalTensorToMemrefPatterns(mlir::RewritePatternSet& patterns); 
+
+}

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

@@ -9,17 +9,6 @@ include "src/Accelerators/PIM/Dialect/Spatial/Spatial.td"
 include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
 #endif // OP_BASE
 
-def HasSpatialChannelSourceCoreIdAttr: Constraint<
-  CPred<"onnx_mlir::hasSpatialChannelSourceCoreIdAttr($0)">,
-  "spatial channel has precomputed source core id">;
-
-def HasSpatialChannelTargetCoreIdAttr: Constraint<
-  CPred<"onnx_mlir::hasSpatialChannelTargetCoreIdAttr($0)">,
-  "spatial channel has precomputed target core id">;
-
-def createPimReceiveFromSpatialChannelValue: NativeCodeCall<
-  "onnx_mlir::createPimReceiveFromSpatialChannel($_builder, $_loc, $0, $1)">;
-
 def onnxToPimTranspose : Pat<
   (ONNXTransposeOp:$srcOpRes $data, $perms),
   (PimTransposeOp $data, $perms,
@@ -27,13 +16,13 @@ def onnxToPimTranspose : Pat<
 >;
 
 def spatToPimVMM : Pat<
-  (SpatWeightedVMMOp:$srcOpRes $weightIndex, $vector),
+  (SpatVMMOp:$srcOpRes $weightIndex, $vector),
   (PimVMMOp $weightIndex, $vector,
     (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
 >;
 
 def spatToPimMVM : Pat<
-  (SpatWeightedMVMOp:$srcOpRes $weightIndex, $vector),
+  (SpatMVMOp:$srcOpRes $weightIndex, $vector),
   (PimMVMOp $weightIndex, $vector,
     (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
 >;
@@ -80,18 +69,4 @@ def spatToPimVSoftmax : Pat<
     (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
 >;
 
-def spatChannelSendToPimSend : Pat<
-  (SpatChannelSendOp $channel, $input),
-  (PimSendOp $input,
-    (NativeCodeCall<"onnx_mlir::getTensorSizeInBytesAttr($_builder, $0)"> $input),
-    (NativeCodeCall<"onnx_mlir::getSpatialChannelTargetCoreIdAttr($_builder, $0)"> $channel)),
-  [(HasSpatialChannelTargetCoreIdAttr $channel)]
->;
-
-def spatChannelReceiveToPimReceive : Pat<
-  (SpatChannelReceiveOp:$srcOpRes $channel),
-  (createPimReceiveFromSpatialChannelValue $srcOpRes, $channel),
-  [(HasSpatialChannelSourceCoreIdAttr $channel)]
->;
-
 #endif // SPATIAL_TO_PIM

src/PIM/Dialect/Pim/CMakeLists.txt

@@ -6,6 +6,8 @@ add_subdirectory(Transforms/Bufferization)
 add_pim_library(PimOps
   PimOps.hpp
   PimOps.cpp
+  PimOpsAsm.cpp
+  PimOpsVerify.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

src/PIM/Dialect/Pim/Pim.td

@@ -24,7 +24,7 @@ def PimTensor :
 // Execution
 //===----------------------------------------------------------------------===//
 
-def PimCoreOp : PimOp<"core", [SingleBlock]> {
+def PimCoreOp : PimOp<"core", [SingleBlock, IsolatedFromAbove]> {
   let summary = "Execute a block on a PIM core";
 
   let regions = (region SizedRegion<1>:$body);
@@ -39,6 +39,20 @@ def PimCoreOp : PimOp<"core", [SingleBlock]> {
   }];
 }
 
+def PimCoreBatchOp : PimOp<"core_batch", [SingleBlock, IsolatedFromAbove, AttrSizedOperandSegments]> {
+  let summary = "Execute equivalent batched core bodies";
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let arguments = (ins
+    I32Attr:$laneCount,
+    Variadic<PimTensor>:$weights,
+    Variadic<PimTensor>:$inputs
+  );
+
+  let hasCustomAssemblyFormat = 1;
+}
+
 def PimHaltOp : PimOp<"halt", [Terminator]> {
   let summary = "Halt execution of the core";
 
@@ -47,6 +61,48 @@ def PimHaltOp : PimOp<"halt", [Terminator]> {
   }];
 }
 
+def PimYieldOp : PimOp<"yield", [Terminator]> {
+  let summary = "Yield results from a Pim region";
+
+  let arguments = (ins
+    Variadic<PimTensor>:$outputs
+  );
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimMapOp : PimOp<"map", [SingleBlock]> {
+  let summary = "Apply the same lane-local region to many independent tensors";
+
+  let arguments = (ins
+    Variadic<PimTensor>:$inputs
+  );
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Tensor Utilities
+//===----------------------------------------------------------------------===//
+
+def PimEmptyManyOp : PimOp<"empty_many", []> {
+  let summary = "Create many identical empty tensors";
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Communication
 //===----------------------------------------------------------------------===//
@@ -65,6 +121,54 @@ def PimSendOp : PimOp<"send", []> {
   }];
 }
 
+def PimSendManyOp : PimOp<"send_many", []> {
+  let summary = "Send multiple tensors to target cores";
+
+  let arguments = (ins
+    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<PimTensor>:$inputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimSendTensorOp : PimOp<"send_tensor", []> {
+  let summary = "Send equal contiguous chunks of one tensor to target cores";
+
+  let arguments = (ins
+    PimTensor:$input,
+    DenseI32ArrayAttr:$targetCoreIds
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimSendBatchOp : PimOp<"send_batch", []> {
+  let summary = "Send a per-lane tensor to target cores from a batched core";
+
+  let arguments = (ins
+    PimTensor:$input,
+    I32Attr:$size,
+    DenseI32ArrayAttr:$targetCoreIds
+  );
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimSendManyBatchOp : PimOp<"send_many_batch", []> {
+  let summary = "Send multiple per-lane tensors to target cores from a batched core";
+
+  let arguments = (ins
+    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<PimTensor>:$inputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   let summary = "Receive a tensor from another core";
 
@@ -89,6 +193,94 @@ def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimReceiveManyOp : PimOp<"receive_many", [DestinationStyleOpInterface]> {
+  let summary = "Receive multiple tensors from source cores";
+
+  let arguments = (ins
+    Variadic<PimTensor>:$outputBuffers,
+    DenseI32ArrayAttr:$sourceCoreIds
+  );
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBuffersMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimReceiveTensorOp : PimOp<"receive_tensor", [DestinationStyleOpInterface]> {
+  let summary = "Receive equal contiguous chunks from source cores into one tensor";
+
+  let arguments = (ins
+    PimTensor:$outputBuffer,
+    DenseI32ArrayAttr:$sourceCoreIds
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBufferMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimReceiveBatchOp : PimOp<"receive_batch", [DestinationStyleOpInterface]> {
+  let summary = "Receive per-lane tensors from source cores into a batched core";
+
+  let arguments = (ins
+    PimTensor:$outputBuffer,
+    I32Attr:$size,
+    DenseI32ArrayAttr:$sourceCoreIds
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBufferMutable();
+    }
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimReceiveManyBatchOp : PimOp<"receive_many_batch", [DestinationStyleOpInterface]> {
+  let summary = "Receive multiple per-lane tensors from source cores into a batched core";
+
+  let arguments = (ins
+    Variadic<PimTensor>:$outputBuffers,
+    DenseI32ArrayAttr:$sourceCoreIds
+  );
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBuffersMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from host memory into device memory";
 
@@ -115,6 +307,32 @@ def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimMemCopyHostToDevBatchOp : PimOp<"memcp_hd_batch", [DestinationStyleOpInterface]> {
+  let summary = "Copy a per-lane tensor from host memory into device memory inside a batched core";
+
+  let arguments = (ins
+    PimTensor:$deviceTarget,
+    PimTensor:$hostSource,
+    I32Attr:$deviceTargetOffset,
+    I32Attr:$hostSourceOffset,
+    I32Attr:$size
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getDeviceTargetMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $deviceTarget `,` $hostSource `)` attr-dict `:` `(` type($deviceTarget) `,` type($hostSource) `)` `->` type($output)
+  }];
+}
+
 def PimMemCopyDevToHostOp : PimOp<"memcp_dh", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from device memory into host memory";
 
@@ -167,6 +385,55 @@ def PimMemCopyOp : PimOp<"memcp", [DestinationStyleOpInterface]> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Tensor utilities
+//===----------------------------------------------------------------------===//
+
+def PimExtractRowsOp : PimOp<"extract_rows", [DestinationStyleOpInterface]> {
+  let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";
+
+  let arguments = (ins
+    PimTensor:$input,
+    Variadic<PimTensor>:$outputBuffers
+  );
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBuffersMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimConcatOp : PimOp<"concat", [DestinationStyleOpInterface]> {
+  let summary = "Concatenate tensors";
+
+  let arguments = (ins
+    I64Attr:$axis,
+    Variadic<PimTensor>:$inputs,
+    PimTensor:$outputBuffer
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBufferMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Math
 //===----------------------------------------------------------------------===//

src/PIM/Dialect/Pim/PimOps.cpp

@@ -1,19 +1,5 @@
-#include "mlir/Dialect/Shape/IR/Shape.h"
-#include "mlir/IR/Block.h"
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/IntegerSet.h"
-#include "mlir/IR/Matchers.h"
-#include "mlir/IR/OpImplementation.h"
-#include "mlir/IR/PatternMatch.h"
-
-#include "llvm/ADT/SetVector.h"
-#include "llvm/ADT/SmallBitVector.h"
-
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
-using namespace mlir;
-
 namespace onnx_mlir {
 namespace pim {
 

src/PIM/Dialect/Pim/PimOpsAsm.cpp

@@ -0,0 +1,550 @@
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/Value.h"
+
+#include "llvm/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Common/IR/CompactAsmUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace pim {
+
+namespace {
+using namespace onnx_mlir::compact_asm;
+
+static DenseI32ArrayAttr getDenseI32ArrayAttr(OpAsmParser& parser, ArrayRef<int32_t> values) {
+  return parser.getBuilder().getDenseI32ArrayAttr(values);
+}
+
+static void printCoreIdList(OpAsmPrinter& printer, StringRef keyword, ArrayRef<int32_t> coreIds) {
+  printer << " " << keyword << " ";
+  printCompressedIntegerList(printer, coreIds);
+}
+
+static ParseResult parseOptionalCoreIdList(OpAsmParser& parser, StringRef keyword, SmallVectorImpl<int32_t>& coreIds) {
+  if (failed(parser.parseOptionalKeyword(keyword)))
+    return success();
+  return parseCompressedIntegerList(parser, coreIds);
+}
+
+} // namespace
+
+void PimCoreBatchOp::print(OpAsmPrinter& printer) {
+  printer << " lanes " << getLaneCount() << " ";
+  size_t weightsPerLane = getLaneCount() > 0 ? getWeights().size() / static_cast<size_t>(getLaneCount()) : 0;
+  if (getLaneCount() > 1 && hasRepeatedTuple(getWeights(), weightsPerLane))
+    printValueTupleRun(printer, getWeights(), weightsPerLane, ListDelimiter::Paren);
+  else
+    printCompressedValueList(printer, getWeights(), ListDelimiter::Paren);
+  printer << " ";
+  printCompressedValueList(printer, getInputs(), ListDelimiter::Square);
+
+  if (auto coreIdsAttr = (*this)->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName))
+    printCoreIdList(printer, "coreIds", coreIdsAttr.asArrayRef());
+
+  printer.printOptionalAttrDict(
+    (*this)->getAttrs(),
+    {getLaneCountAttrName().getValue(), getOperandSegmentSizesAttrName().getValue(), onnx_mlir::kCoreIdsAttrName});
+  printer << " ";
+  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
+  printer << " : ";
+  if (getLaneCount() > 1 && hasRepeatedTuple(TypeRange(getWeights()), weightsPerLane))
+    printTypeTupleRun(printer, TypeRange(getWeights()), weightsPerLane, ListDelimiter::Paren);
+  else
+    printCompressedTypeList(printer, TypeRange(getWeights()), ListDelimiter::Paren);
+  printer << " ";
+  printCompressedTypeList(printer, TypeRange(getInputs()), ListDelimiter::Square);
+  printer << " -> ()";
+}
+
+ParseResult PimCoreBatchOp::parse(OpAsmParser& parser, OperationState& result) {
+  int32_t laneCount = 0;
+  SmallVector<OpAsmParser::UnresolvedOperand> weights;
+  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
+  SmallVector<Type> weightTypes;
+  SmallVector<Type> inputTypes;
+  SmallVector<int32_t> coreIds;
+
+  if (parser.parseKeyword("lanes") || parser.parseInteger(laneCount)
+      || parseCompressedOrTupleOperandList(parser, ListDelimiter::Paren, weights)
+      || parseCompressedOperandList(parser, ListDelimiter::Square, inputs))
+    return failure();
+
+  bool hasCoreIds = succeeded(parser.parseOptionalKeyword("coreIds"));
+  if (hasCoreIds && parseCompressedIntegerList(parser, coreIds))
+    return failure();
+
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return failure();
+
+  Region* body = result.addRegion();
+  if (parser.parseRegion(*body))
+    return failure();
+
+  if (parser.parseColon() || parseCompressedOrTupleTypeList(parser, ListDelimiter::Paren, weightTypes)
+      || parseCompressedTypeList(parser, ListDelimiter::Square, inputTypes) || parser.parseArrow()
+      || parser.parseLParen() || parser.parseRParen())
+    return failure();
+
+  if (weights.size() != weightTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of weights and weight types must match");
+  if (inputs.size() != inputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
+  if (hasCoreIds && result.attributes.get(onnx_mlir::kCoreIdsAttrName))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "coreIds cannot be specified both positionally and in attr-dict");
+
+  auto& builder = parser.getBuilder();
+  result.addAttribute("laneCount", builder.getI32IntegerAttr(laneCount));
+  result.addAttribute(
+    "operandSegmentSizes",
+    builder.getDenseI32ArrayAttr({static_cast<int32_t>(weights.size()), static_cast<int32_t>(inputs.size())}));
+  if (hasCoreIds)
+    result.addAttribute(onnx_mlir::kCoreIdsAttrName, getDenseI32ArrayAttr(parser, coreIds));
+
+  if (parser.resolveOperands(weights, weightTypes, parser.getCurrentLocation(), result.operands)
+      || parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands)) {
+    return failure();
+  }
+  return success();
+}
+
+void PimYieldOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getOutputs());
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printCompressedTypeSequence(printer, getOutputs().getTypes());
+}
+
+ParseResult PimYieldOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> outputs;
+  SmallVector<Type> outputTypes;
+
+  OpAsmParser::UnresolvedOperand firstOutput;
+  OptionalParseResult firstOutputResult = parser.parseOptionalOperand(firstOutput);
+  if (firstOutputResult.has_value()) {
+    if (failed(*firstOutputResult))
+      return failure();
+    if (parseCompressedOperandEntryWithFirst(parser, firstOutput, outputs))
+      return failure();
+    while (succeeded(parser.parseOptionalComma()))
+      if (parseOneCompressedOperandEntry(parser, outputs))
+        return failure();
+  }
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
+    return failure();
+
+  if (outputs.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of outputs and output types must match");
+
+  return parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands);
+}
+
+void PimMapOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printArgumentBindings(printer, getBody().front(), getInputs());
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printer.printType(getInputs().front().getType());
+  printer << " -> ";
+  printer.printType(getOutputs().front().getType());
+  printer << " ";
+  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
+}
+
+ParseResult PimMapOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::Argument> regionArgs;
+  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
+  Type inputType;
+  Type outputType;
+
+  if (parseArgumentBindings(parser, regionArgs, inputs))
+    return failure();
+  if (inputs.empty())
+    return parser.emitError(parser.getCurrentLocation(), "map requires at least one input");
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() || parser.parseType(inputType)
+      || parser.parseArrow() || parser.parseType(outputType))
+    return failure();
+
+  SmallVector<Type> inputTypes(inputs.size(), inputType);
+  SmallVector<Type> outputTypes(inputs.size(), outputType);
+  if (regionArgs.size() != inputs.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of argument bindings and input operands must match");
+  if (parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputTypes);
+
+  applyArgumentTypes(inputTypes, regionArgs);
+  Region* body = result.addRegion();
+  return parser.parseRegion(*body, regionArgs);
+}
+
+void PimEmptyManyOp::print(OpAsmPrinter& printer) {
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printer.printType(getOutputs().front().getType());
+  printer << " x" << getOutputs().size();
+}
+
+ParseResult PimEmptyManyOp::parse(OpAsmParser& parser, OperationState& result) {
+  Type outputType;
+  int64_t resultCount = 0;
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() || parser.parseType(outputType)
+      || parser.parseKeyword("x") || parser.parseInteger(resultCount))
+    return failure();
+
+  if (resultCount <= 0)
+    return parser.emitError(parser.getCurrentLocation(), "result count after 'x' must be positive");
+
+  SmallVector<Type> resultTypes(resultCount, outputType);
+  result.addTypes(resultTypes);
+  return success();
+}
+
+void PimSendBatchOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printer.printOperand(getInput());
+  printCoreIdList(printer, "to", getTargetCoreIds());
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getTargetCoreIdsAttrName().getValue()});
+  printer << " : ";
+  printer.printType(getInput().getType());
+}
+
+ParseResult PimSendBatchOp::parse(OpAsmParser& parser, OperationState& result) {
+  OpAsmParser::UnresolvedOperand input;
+  Type inputType;
+  SmallVector<int32_t> targetCoreIds;
+
+  if (parser.parseOperand(input) || parseOptionalCoreIdList(parser, "to", targetCoreIds)
+      || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() || parser.parseType(inputType))
+    return failure();
+
+  if (!targetCoreIds.empty() && result.attributes.get("targetCoreIds"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "targetCoreIds cannot be specified both positionally and in attr-dict");
+  if (!targetCoreIds.empty())
+    result.addAttribute("targetCoreIds", getDenseI32ArrayAttr(parser, targetCoreIds));
+
+  return parser.resolveOperand(input, inputType, result.operands);
+}
+
+void PimSendManyOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getInputs());
+  printCoreIdList(printer, "to", getTargetCoreIds());
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getTargetCoreIdsAttrName().getValue()});
+  printer << " : ";
+  printCompressedTypeSequence(printer, TypeRange(getInputs()));
+}
+
+ParseResult PimSendManyOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
+  SmallVector<Type> inputTypes;
+  SmallVector<int32_t> targetCoreIds;
+
+  if (parseCompressedOperandSequence(parser, inputs) || parseOptionalCoreIdList(parser, "to", targetCoreIds)
+      || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, inputTypes, /*allowEmpty=*/false))
+    return failure();
+
+  if (inputs.size() != inputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
+  if (!targetCoreIds.empty() && result.attributes.get("targetCoreIds"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "targetCoreIds cannot be specified both positionally and in attr-dict");
+  if (!targetCoreIds.empty())
+    result.addAttribute("targetCoreIds", getDenseI32ArrayAttr(parser, targetCoreIds));
+
+  return parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands);
+}
+
+void PimSendTensorOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printer.printOperand(getInput());
+  printCoreIdList(printer, "to", getTargetCoreIds());
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getTargetCoreIdsAttrName().getValue()});
+  printer << " : ";
+  printer.printType(getInput().getType());
+}
+
+ParseResult PimSendTensorOp::parse(OpAsmParser& parser, OperationState& result) {
+  OpAsmParser::UnresolvedOperand input;
+  Type inputType;
+  SmallVector<int32_t> targetCoreIds;
+
+  if (parser.parseOperand(input) || parseOptionalCoreIdList(parser, "to", targetCoreIds)
+      || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() || parser.parseType(inputType))
+    return failure();
+
+  if (!targetCoreIds.empty() && result.attributes.get("targetCoreIds"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "targetCoreIds cannot be specified both positionally and in attr-dict");
+  if (!targetCoreIds.empty())
+    result.addAttribute("targetCoreIds", getDenseI32ArrayAttr(parser, targetCoreIds));
+
+  return parser.resolveOperand(input, inputType, result.operands);
+}
+
+void PimSendManyBatchOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getInputs());
+  printCoreIdList(printer, "to", getTargetCoreIds());
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getTargetCoreIdsAttrName().getValue()});
+  printer << " : ";
+  printCompressedTypeSequence(printer, TypeRange(getInputs()));
+}
+
+ParseResult PimSendManyBatchOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
+  SmallVector<Type> inputTypes;
+  SmallVector<int32_t> targetCoreIds;
+
+  if (parseCompressedOperandSequence(parser, inputs) || parseOptionalCoreIdList(parser, "to", targetCoreIds)
+      || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, inputTypes, /*allowEmpty=*/false))
+    return failure();
+
+  if (inputs.size() != inputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
+  if (!targetCoreIds.empty() && result.attributes.get("targetCoreIds"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "targetCoreIds cannot be specified both positionally and in attr-dict");
+  if (!targetCoreIds.empty())
+    result.addAttribute("targetCoreIds", getDenseI32ArrayAttr(parser, targetCoreIds));
+
+  return parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands);
+}
+
+void PimReceiveManyOp::print(OpAsmPrinter& printer) {
+  printCoreIdList(printer, "from", getSourceCoreIds());
+  printer << " into ";
+  printOpenDelimiter(printer, ListDelimiter::Paren);
+  printCompressedValueSequence(printer, getOutputBuffers());
+  printCloseDelimiter(printer, ListDelimiter::Paren);
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getSourceCoreIdsAttrName().getValue()});
+  printer << " : ";
+  printCompressedTypeSequence(printer, getOutputs().getTypes());
+}
+
+ParseResult PimReceiveManyOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> outputBuffers;
+  SmallVector<Type> outputTypes;
+  SmallVector<int32_t> sourceCoreIds;
+
+  if (parseOptionalCoreIdList(parser, "from", sourceCoreIds) || parser.parseKeyword("into") || parser.parseLParen()
+      || parseCompressedOperandSequence(parser, outputBuffers) || parser.parseRParen()
+      || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/false))
+    return failure();
+
+  if (outputBuffers.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of output buffers and output types must match");
+  if (!sourceCoreIds.empty() && result.attributes.get("sourceCoreIds"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "sourceCoreIds cannot be specified both positionally and in attr-dict");
+  if (!sourceCoreIds.empty())
+    result.addAttribute("sourceCoreIds", getDenseI32ArrayAttr(parser, sourceCoreIds));
+
+  if (parser.resolveOperands(outputBuffers, outputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputTypes);
+  return success();
+}
+
+void PimReceiveTensorOp::print(OpAsmPrinter& printer) {
+  printCoreIdList(printer, "from", getSourceCoreIds());
+  printer << " into ";
+  printOpenDelimiter(printer, ListDelimiter::Paren);
+  printer.printOperand(getOutputBuffer());
+  printCloseDelimiter(printer, ListDelimiter::Paren);
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getSourceCoreIdsAttrName().getValue()});
+  printer << " : ";
+  printer.printType(getOutputBuffer().getType());
+  printer << " -> ";
+  printer.printType(getOutput().getType());
+}
+
+ParseResult PimReceiveTensorOp::parse(OpAsmParser& parser, OperationState& result) {
+  OpAsmParser::UnresolvedOperand outputBuffer;
+  Type outputBufferType;
+  Type outputType;
+  SmallVector<int32_t> sourceCoreIds;
+
+  if (parseOptionalCoreIdList(parser, "from", sourceCoreIds) || parser.parseKeyword("into") || parser.parseLParen()
+      || parser.parseOperand(outputBuffer) || parser.parseRParen() || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon() || parser.parseType(outputBufferType) || parser.parseArrow()
+      || parser.parseType(outputType))
+    return failure();
+
+  if (!sourceCoreIds.empty() && result.attributes.get("sourceCoreIds"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "sourceCoreIds cannot be specified both positionally and in attr-dict");
+  if (!sourceCoreIds.empty())
+    result.addAttribute("sourceCoreIds", getDenseI32ArrayAttr(parser, sourceCoreIds));
+
+  if (parser.resolveOperand(outputBuffer, outputBufferType, result.operands))
+    return failure();
+  result.addTypes(outputType);
+  return success();
+}
+
+void PimReceiveBatchOp::print(OpAsmPrinter& printer) {
+  printCoreIdList(printer, "from", getSourceCoreIds());
+  printer << " into ";
+  printOpenDelimiter(printer, ListDelimiter::Paren);
+  printer.printOperand(getOutputBuffer());
+  printCloseDelimiter(printer, ListDelimiter::Paren);
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getSourceCoreIdsAttrName().getValue()});
+  printer << " : ";
+  printer.printType(getOutputBuffer().getType());
+  printer << " -> ";
+  printer.printType(getOutput().getType());
+}
+
+ParseResult PimReceiveBatchOp::parse(OpAsmParser& parser, OperationState& result) {
+  OpAsmParser::UnresolvedOperand outputBuffer;
+  Type outputBufferType;
+  Type outputType;
+  SmallVector<int32_t> sourceCoreIds;
+
+  if (parseOptionalCoreIdList(parser, "from", sourceCoreIds) || parser.parseKeyword("into") || parser.parseLParen()
+      || parser.parseOperand(outputBuffer) || parser.parseRParen() || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon() || parser.parseType(outputBufferType) || parser.parseArrow()
+      || parser.parseType(outputType))
+    return failure();
+
+  if (!sourceCoreIds.empty() && result.attributes.get("sourceCoreIds"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "sourceCoreIds cannot be specified both positionally and in attr-dict");
+  if (!sourceCoreIds.empty())
+    result.addAttribute("sourceCoreIds", getDenseI32ArrayAttr(parser, sourceCoreIds));
+
+  if (parser.resolveOperand(outputBuffer, outputBufferType, result.operands))
+    return failure();
+  result.addTypes(outputType);
+  return success();
+}
+
+void PimReceiveManyBatchOp::print(OpAsmPrinter& printer) {
+  printCoreIdList(printer, "from", getSourceCoreIds());
+  printer << " into ";
+  printOpenDelimiter(printer, ListDelimiter::Paren);
+  printCompressedValueSequence(printer, getOutputBuffers());
+  printCloseDelimiter(printer, ListDelimiter::Paren);
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getSourceCoreIdsAttrName().getValue()});
+  printer << " : ";
+  printCompressedTypeSequence(printer, getOutputs().getTypes());
+}
+
+ParseResult PimReceiveManyBatchOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> outputBuffers;
+  SmallVector<Type> outputTypes;
+  SmallVector<int32_t> sourceCoreIds;
+
+  if (parseOptionalCoreIdList(parser, "from", sourceCoreIds) || parser.parseKeyword("into") || parser.parseLParen()
+      || parseCompressedOperandSequence(parser, outputBuffers) || parser.parseRParen()
+      || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/false))
+    return failure();
+
+  if (outputBuffers.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of output buffers and output types must match");
+  if (!sourceCoreIds.empty() && result.attributes.get("sourceCoreIds"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "sourceCoreIds cannot be specified both positionally and in attr-dict");
+  if (!sourceCoreIds.empty())
+    result.addAttribute("sourceCoreIds", getDenseI32ArrayAttr(parser, sourceCoreIds));
+
+  if (parser.resolveOperands(outputBuffers, outputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputTypes);
+  return success();
+}
+
+void PimExtractRowsOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printer.printOperand(getInput());
+  printer << " into ";
+  printOpenDelimiter(printer, ListDelimiter::Paren);
+  printCompressedValueSequence(printer, getOutputBuffers());
+  printCloseDelimiter(printer, ListDelimiter::Paren);
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printer.printType(getInput().getType());
+  printer << " -> ";
+  printCompressedTypeSequence(printer, getOutputs().getTypes());
+}
+
+ParseResult PimExtractRowsOp::parse(OpAsmParser& parser, OperationState& result) {
+  OpAsmParser::UnresolvedOperand input;
+  SmallVector<OpAsmParser::UnresolvedOperand> outputBuffers;
+  Type inputType;
+  SmallVector<Type> outputTypes;
+
+  if (parser.parseOperand(input) || parser.parseKeyword("into") || parser.parseLParen()
+      || parseCompressedOperandSequence(parser, outputBuffers) || parser.parseRParen()
+      || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() || parser.parseType(inputType)
+      || parser.parseArrow() || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/false))
+    return failure();
+
+  if (outputBuffers.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of output buffers and output types must match");
+  if (parser.resolveOperand(input, inputType, result.operands)
+      || parser.resolveOperands(outputBuffers, outputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputTypes);
+  return success();
+}
+
+void PimConcatOp::print(OpAsmPrinter& printer) {
+  printer << " axis " << getAxis() << " ";
+  printCompressedValueSequence(printer, getInputs());
+  printer << " into ";
+  printer.printOperand(getOutputBuffer());
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getAxisAttrName().getValue()});
+  printer << " : ";
+  printOpenDelimiter(printer, ListDelimiter::Paren);
+  printCompressedTypeSequence(printer, TypeRange(getInputs()));
+  printCloseDelimiter(printer, ListDelimiter::Paren);
+  printer << " -> ";
+  printer.printType(getOutput().getType());
+}
+
+ParseResult PimConcatOp::parse(OpAsmParser& parser, OperationState& result) {
+  int64_t axis = 0;
+  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
+  OpAsmParser::UnresolvedOperand outputBuffer;
+  SmallVector<Type> inputTypes;
+  Type outputType;
+
+  if (parser.parseKeyword("axis") || parser.parseInteger(axis) || parseCompressedOperandSequence(parser, inputs)
+      || parser.parseKeyword("into") || parser.parseOperand(outputBuffer)
+      || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() || parser.parseLParen()
+      || parseCompressedTypeSequence(parser, inputTypes, /*allowEmpty=*/false) || parser.parseRParen()
+      || parser.parseArrow() || parser.parseType(outputType))
+    return failure();
+
+  if (inputs.size() != inputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
+  if (result.attributes.get("axis"))
+    return parser.emitError(parser.getCurrentLocation(), "axis cannot be specified both positionally and in attr-dict");
+
+  result.addAttribute("axis", parser.getBuilder().getI64IntegerAttr(axis));
+  if (parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands)
+      || parser.resolveOperand(outputBuffer, outputType, result.operands))
+    return failure();
+  result.addTypes(outputType);
+  return success();
+}
+
+} // namespace pim
+} // namespace onnx_mlir

src/PIM/Dialect/Pim/PimOpsVerify.cpp

@@ -0,0 +1,300 @@
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/TypeUtilities.h"
+
+#include "llvm/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace pim {
+
+namespace {
+
+static LogicalResult verifyManyCommunicationSizes(Operation* op, ArrayRef<int32_t> coreIds, size_t valueCount) {
+  if (coreIds.size() != valueCount)
+    return op->emitError("core id metadata length must match the number of values");
+  return success();
+}
+
+static bool haveSameShapedContainerKind(Type lhs, Type rhs) {
+  return (isa<RankedTensorType>(lhs) && isa<RankedTensorType>(rhs)) || (isa<MemRefType>(lhs) && isa<MemRefType>(rhs));
+}
+
+static LogicalResult verifyCompatibleShapedTypes(Operation* op, Type lhs, Type rhs, StringRef message) {
+  auto lhsShaped = dyn_cast<ShapedType>(lhs);
+  auto rhsShaped = dyn_cast<ShapedType>(rhs);
+  if (!lhsShaped || !rhsShaped || !haveSameShapedContainerKind(lhs, rhs))
+    return op->emitError(message);
+  if (lhsShaped.getElementType() != rhsShaped.getElementType() || lhsShaped.getShape() != rhsShaped.getShape())
+    return op->emitError(message);
+  return success();
+}
+
+static LogicalResult verifyManyCommunicationTypes(Operation* op, TypeRange types, StringRef kind) {
+  if (types.empty())
+    return op->emitError() << kind << " must carry at least one value";
+
+  Type firstType = types.front();
+  auto firstShapedType = dyn_cast<ShapedType>(firstType);
+  bool firstIsTensor = isa<RankedTensorType>(firstType);
+  bool firstIsMemRef = isa<MemRefType>(firstType);
+  for (Type type : types.drop_front())
+    if (type != firstType) {
+      auto shapedType = dyn_cast<ShapedType>(type);
+      if (!firstShapedType || !shapedType)
+        return op->emitError() << kind << " values must all have the same type";
+      if (firstIsTensor != isa<RankedTensorType>(type) || firstIsMemRef != isa<MemRefType>(type))
+        return op->emitError() << kind << " values must all use the same shaped container kind";
+      if (firstShapedType.getElementType() != shapedType.getElementType()
+          || firstShapedType.getShape() != shapedType.getShape())
+        return op->emitError() << kind << " values must all have the same shape and element type";
+    }
+  return success();
+}
+
+static LogicalResult verifyTensorCommunication(Operation* op, Type type, ArrayRef<int32_t> coreIds, StringRef kind) {
+  if (coreIds.empty())
+    return op->emitError() << kind << " must carry at least one chunk";
+
+  auto shapedType = dyn_cast<ShapedType>(type);
+  if (!shapedType || !shapedType.hasStaticShape())
+    return op->emitError() << kind << " requires a static shaped tensor or memref";
+
+  int64_t elementBits = shapedType.getElementTypeBitWidth();
+  if (elementBits <= 0 || elementBits % 8 != 0)
+    return op->emitError() << kind << " requires byte-sized elements";
+
+  int64_t totalBytes = shapedType.getNumElements() * elementBits / 8;
+  if (totalBytes % static_cast<int64_t>(coreIds.size()) != 0)
+    return op->emitError() << kind << " tensor byte size must be divisible by the number of core ids";
+
+  return success();
+}
+
+static FailureOr<int32_t> getParentBatchLaneCount(Operation* op) {
+  auto coreBatchOp = op->getParentOfType<PimCoreBatchOp>();
+  if (!coreBatchOp)
+    return failure();
+  return coreBatchOp.getLaneCount();
+}
+
+static LogicalResult verifyManyBatchCommunicationSizes(Operation* op, ArrayRef<int32_t> coreIds, size_t valueCount) {
+  auto laneCount = getParentBatchLaneCount(op);
+  if (failed(laneCount))
+    return op->emitError("must be nested inside pim.core_batch");
+  if (coreIds.size() != valueCount * static_cast<size_t>(*laneCount))
+    return op->emitError("core id metadata length must match the number of values times parent laneCount");
+  return success();
+}
+
+} // namespace
+
+LogicalResult PimEmptyManyOp::verify() {
+  if (getOutputs().empty())
+    return emitError("must produce at least one output");
+
+  Type firstType = getOutputs().front().getType();
+  auto firstShapedType = dyn_cast<ShapedType>(firstType);
+  if (!firstShapedType || !firstShapedType.hasRank())
+    return emitError("outputs must all be ranked shaped types");
+
+  for (Value output : getOutputs().drop_front())
+    if (output.getType() != firstType)
+      return emitError("outputs must all have the same type");
+
+  return success();
+}
+
+LogicalResult PimMapOp::verify() {
+  if (getInputs().empty())
+    return emitError("requires at least one input");
+  if (getOutputs().size() != getInputs().size())
+    return emitError("number of outputs must match number of inputs");
+
+  Type inputType = getInputs().front().getType();
+  for (Value input : getInputs().drop_front())
+    if (input.getType() != inputType)
+      return emitError("all inputs must have the same type");
+
+  Type outputType = getOutputs().front().getType();
+  for (Value output : getOutputs().drop_front())
+    if (output.getType() != outputType)
+      return emitError("all outputs must have the same type");
+
+  Block& block = getBody().front();
+  if (block.getNumArguments() != 1)
+    return emitError("body must have exactly one block argument");
+  if (failed(verifyCompatibleShapedTypes(
+        getOperation(), block.getArgument(0).getType(), inputType, "body block argument type must match input type")))
+    return emitError("body block argument type must match input type");
+
+  auto yieldOp = dyn_cast_or_null<PimYieldOp>(block.getTerminator());
+  if (!yieldOp)
+    return emitError("body must terminate with pim.yield");
+  if (yieldOp.getNumOperands() != 1)
+    return emitError("body yield must produce exactly one value");
+  if (failed(verifyCompatibleShapedTypes(
+        getOperation(), yieldOp.getOperand(0).getType(), outputType, "body yield type must match output type")))
+    return emitError("body yield type must match output type");
+
+  return success();
+}
+
+LogicalResult PimSendManyOp::verify() {
+  if (failed(verifyManyCommunicationSizes(getOperation(), getTargetCoreIds(), getInputs().size())))
+    return failure();
+  return verifyManyCommunicationTypes(getOperation(), getInputs().getTypes(), "send_many");
+}
+
+LogicalResult PimSendTensorOp::verify() {
+  return verifyTensorCommunication(getOperation(), getInput().getType(), getTargetCoreIds(), "send_tensor");
+}
+
+LogicalResult PimSendManyBatchOp::verify() {
+  if (failed(verifyManyBatchCommunicationSizes(getOperation(), getTargetCoreIds(), getInputs().size())))
+    return failure();
+  return verifyManyCommunicationTypes(getOperation(), getInputs().getTypes(), "send_many_batch");
+}
+
+LogicalResult PimReceiveManyOp::verify() {
+  if (getOutputBuffers().size() != getOutputs().size())
+    return emitError("number of output buffers must match the number of outputs");
+  if (failed(verifyManyCommunicationSizes(getOperation(), getSourceCoreIds(), getOutputs().size())))
+    return failure();
+
+  if (failed(verifyManyCommunicationTypes(getOperation(), getOutputBuffers().getTypes(), "receive_many")))
+    return failure();
+  if (failed(verifyManyCommunicationTypes(getOperation(), getOperation()->getResultTypes(), "receive_many")))
+    return failure();
+
+  for (auto [outputBuffer, output] : llvm::zip(getOutputBuffers(), getOutputs()))
+    if (outputBuffer.getType() != output.getType())
+      return emitError("output buffers and outputs must have matching types");
+
+  return success();
+}
+
+LogicalResult PimReceiveTensorOp::verify() {
+  if (failed(verifyCompatibleShapedTypes(
+        getOperation(), getOutputBuffer().getType(), getOutput().getType(), "output buffer and output must match")))
+    return failure();
+
+  return verifyTensorCommunication(getOperation(), getOutput().getType(), getSourceCoreIds(), "receive_tensor");
+}
+
+LogicalResult PimReceiveManyBatchOp::verify() {
+  if (getOutputBuffers().size() != getOutputs().size())
+    return emitError("number of output buffers must match the number of outputs");
+  if (failed(verifyManyBatchCommunicationSizes(getOperation(), getSourceCoreIds(), getOutputs().size())))
+    return failure();
+
+  if (failed(verifyManyCommunicationTypes(getOperation(), getOutputBuffers().getTypes(), "receive_many_batch")))
+    return failure();
+  if (failed(verifyManyCommunicationTypes(getOperation(), getOperation()->getResultTypes(), "receive_many_batch")))
+    return failure();
+
+  for (auto [outputBuffer, output] : llvm::zip(getOutputBuffers(), getOutputs()))
+    if (outputBuffer.getType() != output.getType())
+      return emitError("output buffers and outputs must have matching types");
+
+  return success();
+}
+
+LogicalResult PimExtractRowsOp::verify() {
+  if (getOutputBuffers().size() != getOutputs().size())
+    return emitError("number of output buffers must match the number of outputs");
+
+  auto inputType = dyn_cast<ShapedType>(getInput().getType());
+  if (!inputType || !inputType.hasRank() || inputType.getRank() != 2)
+    return emitError("input must be a rank-2 shaped type");
+
+  int64_t numRows = inputType.getShape()[0];
+  int64_t numCols = inputType.getShape()[1];
+  Type elementType = inputType.getElementType();
+
+  if (numRows >= 0 && static_cast<int64_t>(getOutputs().size()) != numRows)
+    return emitError("number of outputs must match the number of input rows");
+
+  for (auto [outputBuffer, output] : llvm::zip(getOutputBuffers(), getOutputs())) {
+    if (failed(verifyCompatibleShapedTypes(
+          getOperation(), outputBuffer.getType(), output.getType(), "output buffers and outputs must match")))
+      return failure();
+
+    auto outputType = dyn_cast<ShapedType>(output.getType());
+    if (!outputType || !outputType.hasRank() || outputType.getRank() != 2)
+      return emitError("outputs must all be rank-2 shaped types");
+    if (!haveSameShapedContainerKind(getInput().getType(), output.getType()))
+      return emitError("outputs must use the same shaped container kind as the input");
+    if (outputType.getElementType() != elementType)
+      return emitError("output element types must match input element type");
+    auto outputShape = outputType.getShape();
+    if (outputShape[0] != 1)
+      return emitError("each output must have exactly one row");
+    if (numCols >= 0 && outputShape[1] != numCols)
+      return emitError("output column count must match input column count");
+  }
+
+  return success();
+}
+
+LogicalResult PimConcatOp::verify() {
+  if (getInputs().empty())
+    return emitError("requires at least one input");
+
+  if (failed(verifyCompatibleShapedTypes(
+        getOperation(), getOutputBuffer().getType(), getOutput().getType(), "output buffer and output must match")))
+    return failure();
+
+  auto outputType = dyn_cast<ShapedType>(getOutput().getType());
+  if (!outputType || !outputType.hasRank())
+    return emitError("output must be a ranked shaped type");
+
+  int64_t axis = getAxis();
+  int64_t rank = outputType.getRank();
+  if (axis < 0 || axis >= rank)
+    return emitError("axis must be within the output rank");
+
+  int64_t concatenatedDimSize = 0;
+  bool concatenatedDimDynamic = false;
+  Type outputElementType = outputType.getElementType();
+
+  for (Value input : getInputs()) {
+    auto inputType = dyn_cast<ShapedType>(input.getType());
+    if (!inputType || !inputType.hasRank())
+      return emitError("inputs must be ranked shaped types");
+    if (!haveSameShapedContainerKind(input.getType(), getOutput().getType()))
+      return emitError("inputs and output must use the same shaped container kind");
+    if (inputType.getRank() != rank)
+      return emitError("all inputs must have the same rank as the output");
+    if (inputType.getElementType() != outputElementType)
+      return emitError("all inputs must have the same element type as the output");
+
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      if (dim == axis)
+        continue;
+      int64_t inputDim = inputType.getDimSize(dim);
+      int64_t outputDim = outputType.getDimSize(dim);
+      if (!ShapedType::isDynamic(inputDim) && !ShapedType::isDynamic(outputDim) && inputDim != outputDim)
+        return emitError("non-concatenated dimensions must match the output shape");
+    }
+
+    int64_t inputConcatDim = inputType.getDimSize(axis);
+    if (ShapedType::isDynamic(inputConcatDim)) {
+      concatenatedDimDynamic = true;
+      continue;
+    }
+    concatenatedDimSize += inputConcatDim;
+  }
+
+  int64_t outputConcatDim = outputType.getDimSize(axis);
+  if (!concatenatedDimDynamic && !ShapedType::isDynamic(outputConcatDim) && concatenatedDimSize != outputConcatDim)
+    return emitError("output concatenated dimension must equal the sum of input sizes");
+
+  return success();
+}
+
+} // namespace pim
+} // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.cpp

@@ -1,6 +1,7 @@
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 #include "OpBufferizationInterfaces.hpp"
@@ -33,6 +34,13 @@ static Value materializeContiguousMemRef(Value memrefValue, Location loc, Rewrit
     .getOutput();
 }
 
+static FailureOr<Value>
+getBufferOrValue(RewriterBase& rewriter, Value value, const BufferizationOptions& options, BufferizationState& state) {
+  if (isa<BufferLikeType>(value.getType()))
+    return value;
+  return getBuffer(rewriter, value, options, state);
+}
+
 struct MemCopyHostToDevOpInterface
 : DstBufferizableOpInterfaceExternalModel<MemCopyHostToDevOpInterface, PimMemCopyHostToDevOp> {
   LogicalResult bufferize(Operation* op,
@@ -43,12 +51,12 @@ struct MemCopyHostToDevOpInterface
     auto deviceTarget = memCopyHostToDevOp.getDeviceTarget();
     auto hostSource = memCopyHostToDevOp.getHostSource();
 
-    auto deviceTargetOpt = getBuffer(rewriter, deviceTarget, options, state);
+    auto deviceTargetOpt = getBufferOrValue(rewriter, deviceTarget, options, state);
     if (failed(deviceTargetOpt))
       return failure();
     auto deviceTargetMemRef = *deviceTargetOpt;
 
-    auto hostSourceOpt = getBuffer(rewriter, hostSource, options, state);
+    auto hostSourceOpt = getBufferOrValue(rewriter, hostSource, options, state);
     if (failed(hostSourceOpt))
       return failure();
     auto hostSourceMemRef = *hostSourceOpt;
@@ -65,6 +73,32 @@ struct MemCopyHostToDevOpInterface
   }
 };
 
+struct MemCopyHostToDevBatchOpInterface
+: DstBufferizableOpInterfaceExternalModel<MemCopyHostToDevBatchOpInterface, PimMemCopyHostToDevBatchOp> {
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto memCopyHostToDevOp = cast<PimMemCopyHostToDevBatchOp>(op);
+    auto deviceTargetOpt = getBufferOrValue(rewriter, memCopyHostToDevOp.getDeviceTarget(), options, state);
+    if (failed(deviceTargetOpt))
+      return failure();
+    auto hostSourceOpt = getBufferOrValue(rewriter, memCopyHostToDevOp.getHostSource(), options, state);
+    if (failed(hostSourceOpt))
+      return failure();
+
+    replaceOpWithNewBufferizedOp<PimMemCopyHostToDevBatchOp>(rewriter,
+                                                             memCopyHostToDevOp,
+                                                             deviceTargetOpt->getType(),
+                                                             *deviceTargetOpt,
+                                                             *hostSourceOpt,
+                                                             memCopyHostToDevOp.getDeviceTargetOffsetAttr(),
+                                                             memCopyHostToDevOp.getHostSourceOffsetAttr(),
+                                                             memCopyHostToDevOp.getSizeAttr());
+    return success();
+  }
+};
+
 struct MemCopyDevToHostOpInterface
 : DstBufferizableOpInterfaceExternalModel<MemCopyDevToHostOpInterface, PimMemCopyDevToHostOp> {
   LogicalResult bufferize(Operation* op,
@@ -74,13 +108,13 @@ struct MemCopyDevToHostOpInterface
     auto memCopyDevToHostOp = cast<PimMemCopyDevToHostOp>(op);
 
     auto hostTarget = memCopyDevToHostOp.getHostTarget();
-    auto hostTargetOpt = getBuffer(rewriter, hostTarget, options, state);
+    auto hostTargetOpt = getBufferOrValue(rewriter, hostTarget, options, state);
     if (failed(hostTargetOpt))
       return failure();
     auto hostTargetMemRef = *hostTargetOpt;
 
     auto deviceSource = memCopyDevToHostOp.getDeviceSource();
-    auto deviceSourceOpt = getBuffer(rewriter, deviceSource, options, state);
+    auto deviceSourceOpt = getBufferOrValue(rewriter, deviceSource, options, state);
     if (failed(deviceSourceOpt))
       return failure();
     auto deviceSourceMemRef = *deviceSourceOpt;
@@ -108,7 +142,7 @@ struct ReceiveOpInterface : DstBufferizableOpInterfaceExternalModel<ReceiveOpInt
                           BufferizationState& state) const {
     auto receiveOp = cast<PimReceiveOp>(op);
 
-    auto outputBufferOpt = getBuffer(rewriter, receiveOp.getOutputBuffer(), options, state);
+    auto outputBufferOpt = getBufferOrValue(rewriter, receiveOp.getOutputBuffer(), options, state);
     if (failed(outputBufferOpt))
       return failure();
 
@@ -122,6 +156,415 @@ struct ReceiveOpInterface : DstBufferizableOpInterfaceExternalModel<ReceiveOpInt
   }
 };
 
+struct ReceiveBatchOpInterface : DstBufferizableOpInterfaceExternalModel<ReceiveBatchOpInterface, PimReceiveBatchOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto receiveOp = cast<PimReceiveBatchOp>(op);
+    auto outputBufferOpt = getBufferOrValue(rewriter, receiveOp.getOutputBuffer(), options, state);
+    if (failed(outputBufferOpt))
+      return failure();
+
+    replaceOpWithNewBufferizedOp<PimReceiveBatchOp>(rewriter,
+                                                    op,
+                                                    outputBufferOpt->getType(),
+                                                    *outputBufferOpt,
+                                                    receiveOp.getSizeAttr(),
+                                                    receiveOp.getSourceCoreIdsAttr());
+    return success();
+  }
+};
+
+struct ReceiveManyOpInterface : DstBufferizableOpInterfaceExternalModel<ReceiveManyOpInterface, PimReceiveManyOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto receiveOp = cast<PimReceiveManyOp>(op);
+    SmallVector<Value> outputBuffers;
+    SmallVector<Type> resultTypes;
+    outputBuffers.reserve(receiveOp.getOutputBuffers().size());
+    resultTypes.reserve(receiveOp.getOutputBuffers().size());
+
+    for (Value outputBuffer : receiveOp.getOutputBuffers()) {
+      auto outputBufferOpt = getBufferOrValue(rewriter, outputBuffer, options, state);
+      if (failed(outputBufferOpt))
+        return failure();
+      outputBuffers.push_back(*outputBufferOpt);
+      resultTypes.push_back(outputBufferOpt->getType());
+    }
+
+    auto newOp = PimReceiveManyOp::create(rewriter,
+                                          receiveOp.getLoc(),
+                                          TypeRange(resultTypes),
+                                          ValueRange(outputBuffers),
+                                          receiveOp.getSourceCoreIdsAttr());
+    rewriter.replaceOp(receiveOp, newOp.getOutputs());
+    return success();
+  }
+};
+
+struct ReceiveTensorOpInterface
+: DstBufferizableOpInterfaceExternalModel<ReceiveTensorOpInterface, PimReceiveTensorOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto receiveOp = cast<PimReceiveTensorOp>(op);
+    auto outputBufferOpt = getBufferOrValue(rewriter, receiveOp.getOutputBuffer(), options, state);
+    if (failed(outputBufferOpt))
+      return failure();
+
+    replaceOpWithNewBufferizedOp<PimReceiveTensorOp>(
+      rewriter, op, outputBufferOpt->getType(), *outputBufferOpt, receiveOp.getSourceCoreIdsAttr());
+    return success();
+  }
+};
+
+struct ReceiveManyBatchOpInterface
+: DstBufferizableOpInterfaceExternalModel<ReceiveManyBatchOpInterface, PimReceiveManyBatchOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto receiveOp = cast<PimReceiveManyBatchOp>(op);
+    SmallVector<Value> outputBuffers;
+    SmallVector<Type> resultTypes;
+    outputBuffers.reserve(receiveOp.getOutputBuffers().size());
+    resultTypes.reserve(receiveOp.getOutputBuffers().size());
+
+    for (Value outputBuffer : receiveOp.getOutputBuffers()) {
+      auto outputBufferOpt = getBufferOrValue(rewriter, outputBuffer, options, state);
+      if (failed(outputBufferOpt))
+        return failure();
+      outputBuffers.push_back(*outputBufferOpt);
+      resultTypes.push_back(outputBufferOpt->getType());
+    }
+
+    auto newOp = PimReceiveManyBatchOp::create(rewriter,
+                                               receiveOp.getLoc(),
+                                               TypeRange(resultTypes),
+                                               ValueRange(outputBuffers),
+                                               receiveOp.getSourceCoreIdsAttr());
+    rewriter.replaceOp(receiveOp, newOp.getOutputs());
+    return success();
+  }
+};
+
+struct ExtractRowsOpInterface : DstBufferizableOpInterfaceExternalModel<ExtractRowsOpInterface, PimExtractRowsOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto extractRowsOp = cast<PimExtractRowsOp>(op);
+    auto inputOpt = getBufferOrValue(rewriter, extractRowsOp.getInput(), options, state);
+    if (failed(inputOpt))
+      return failure();
+
+    SmallVector<Value> outputBuffers;
+    SmallVector<Type> resultTypes;
+    outputBuffers.reserve(extractRowsOp.getOutputBuffers().size());
+    resultTypes.reserve(extractRowsOp.getOutputBuffers().size());
+
+    for (Value outputBuffer : extractRowsOp.getOutputBuffers()) {
+      auto outputBufferOpt = getBufferOrValue(rewriter, outputBuffer, options, state);
+      if (failed(outputBufferOpt))
+        return failure();
+      outputBuffers.push_back(*outputBufferOpt);
+      resultTypes.push_back(outputBufferOpt->getType());
+    }
+
+    auto newOp = PimExtractRowsOp::create(rewriter,
+                                          extractRowsOp.getLoc(),
+                                          TypeRange(resultTypes),
+                                          materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter),
+                                          ValueRange(outputBuffers));
+    rewriter.replaceOp(extractRowsOp, newOp.getOutputs());
+    return success();
+  }
+};
+
+struct ConcatOpInterface : DstBufferizableOpInterfaceExternalModel<ConcatOpInterface, PimConcatOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto concatOp = cast<PimConcatOp>(op);
+    SmallVector<Value> inputs;
+    inputs.reserve(concatOp.getInputs().size());
+    for (Value input : concatOp.getInputs()) {
+      auto inputOpt = getBufferOrValue(rewriter, input, options, state);
+      if (failed(inputOpt))
+        return failure();
+      inputs.push_back(materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter));
+    }
+
+    auto outputBufferOpt = getBufferOrValue(rewriter, concatOp.getOutputBuffer(), options, state);
+    if (failed(outputBufferOpt))
+      return failure();
+
+    replaceOpWithNewBufferizedOp<PimConcatOp>(
+      rewriter, op, outputBufferOpt->getType(), concatOp.getAxisAttr(), ValueRange(inputs), *outputBufferOpt);
+    return success();
+  }
+};
+
+struct EmptyManyOpInterface : BufferizableOpInterface::ExternalModel<EmptyManyOpInterface, PimEmptyManyOp> {
+  bool bufferizesToAllocation(Operation* op, Value value) const { return true; }
+
+  bool resultBufferizesToMemoryWrite(Operation* op, OpResult opResult, const AnalysisState& state) const {
+    return false;
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto emptyManyOp = cast<PimEmptyManyOp>(op);
+
+    SmallVector<Type> resultTypes;
+    resultTypes.reserve(emptyManyOp.getOutputs().size());
+    for (Value output : emptyManyOp.getOutputs()) {
+      auto shapedType = cast<ShapedType>(output.getType());
+      resultTypes.push_back(MemRefType::get(shapedType.getShape(), shapedType.getElementType()));
+    }
+
+    replaceOpWithNewBufferizedOp<PimEmptyManyOp>(rewriter, emptyManyOp, TypeRange(resultTypes));
+    return success();
+  }
+};
+
+struct SendTensorOpInterface : BufferizableOpInterface::ExternalModel<SendTensorOpInterface, PimSendTensorOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { return true; }
+
+  bool bufferizesToMemoryWrite(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { return false; }
+
+  AliasingValueList getAliasingValues(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return {};
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto sendOp = cast<PimSendTensorOp>(op);
+    auto inputOpt = getBufferOrValue(rewriter, sendOp.getInput(), options, state);
+    if (failed(inputOpt))
+      return failure();
+
+    replaceOpWithNewBufferizedOp<PimSendTensorOp>(
+      rewriter, op, materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter), sendOp.getTargetCoreIdsAttr());
+    return success();
+  }
+};
+
+struct MapOpInterface : BufferizableOpInterface::ExternalModel<MapOpInterface, PimMapOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { return true; }
+
+  bool bufferizesToMemoryWrite(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { return false; }
+
+  AliasingValueList getAliasingValues(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return {};
+  }
+
+  AliasingOpOperandList getAliasingOpOperands(Operation* op, Value value, const AnalysisState& state) const {
+    auto mapOp = cast<PimMapOp>(op);
+    auto bbArg = dyn_cast<BlockArgument>(value);
+    if (!bbArg || bbArg.getOwner() != &mapOp.getBody().front() || bbArg.getArgNumber() != 0
+        || mapOp.getInputs().empty())
+      return {};
+
+    return {
+      {&mapOp->getOpOperand(0), BufferRelation::Equivalent}
+    };
+  }
+
+  bool isWritable(Operation* op, Value value, const AnalysisState& state) const { return false; }
+
+  FailureOr<BufferLikeType> getBufferType(Operation* op,
+                                          Value value,
+                                          const BufferizationOptions& options,
+                                          const BufferizationState& state,
+                                          SmallVector<Value>& invocationStack) const {
+    auto mapOp = cast<PimMapOp>(op);
+    auto bbArg = dyn_cast<BlockArgument>(value);
+    if (!bbArg || bbArg.getOwner() != &mapOp.getBody().front() || bbArg.getArgNumber() != 0
+        || mapOp.getInputs().empty())
+      return failure();
+
+    auto inputType = dyn_cast<BufferLikeType>(mapOp.getInputs().front().getType());
+    if (inputType)
+      return inputType;
+
+    auto shapedType = cast<ShapedType>(mapOp.getInputs().front().getType());
+    return BufferLikeType(MemRefType::get(shapedType.getShape(), shapedType.getElementType()));
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto mapOp = cast<PimMapOp>(op);
+
+    SmallVector<Value> inputs;
+    SmallVector<Type> resultTypes;
+    inputs.reserve(mapOp.getInputs().size());
+    resultTypes.reserve(mapOp.getOutputs().size());
+
+    for (Value input : mapOp.getInputs()) {
+      if (isa<TensorType>(input.getType())) {
+        auto inputOpt = getBufferOrValue(rewriter, input, options, state);
+        if (failed(inputOpt))
+          return failure();
+        inputs.push_back(*inputOpt);
+      }
+      else {
+        inputs.push_back(input);
+      }
+    }
+
+    for (Value output : mapOp.getOutputs()) {
+      auto shapedType = cast<ShapedType>(output.getType());
+      resultTypes.push_back(MemRefType::get(shapedType.getShape(), shapedType.getElementType()));
+    }
+
+    rewriter.setInsertionPoint(mapOp);
+    auto newOp = PimMapOp::create(rewriter, mapOp.getLoc(), TypeRange(resultTypes), ValueRange(inputs));
+    rewriter.inlineRegionBefore(mapOp.getBody(), newOp.getBody(), newOp.getBody().begin());
+    for (Block& block : newOp.getBody())
+      if (failed(bufferization::bufferizeBlockSignature(&block, rewriter, options, state)))
+        return failure();
+
+    rewriter.replaceOp(mapOp, newOp.getOutputs());
+    return success();
+  }
+};
+
+struct CoreBatchOpInterface : BufferizableOpInterface::ExternalModel<CoreBatchOpInterface, PimCoreBatchOp> {
+  bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { return true; }
+
+  bool bufferizesToMemoryWrite(Operation* op, OpOperand& opOperand, const AnalysisState& state) const { return false; }
+
+  AliasingValueList getAliasingValues(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
+    return {};
+  }
+
+  AliasingOpOperandList getAliasingOpOperands(Operation* op, Value value, const AnalysisState& state) const {
+    auto coreBatchOp = cast<PimCoreBatchOp>(op);
+    auto bbArg = dyn_cast<BlockArgument>(value);
+    if (!bbArg || bbArg.getOwner() != &coreBatchOp.getBody().front())
+      return {};
+
+    unsigned inputOperandIndex = coreBatchOp.getWeights().size() + bbArg.getArgNumber();
+    return {
+      {&coreBatchOp->getOpOperand(inputOperandIndex), BufferRelation::Equivalent}
+    };
+  }
+
+  bool isWritable(Operation* op, Value value, const AnalysisState& state) const { return false; }
+
+  FailureOr<BufferLikeType> getBufferType(Operation* op,
+                                          Value value,
+                                          const BufferizationOptions& options,
+                                          const BufferizationState& state,
+                                          SmallVector<Value>& invocationStack) const {
+    auto coreBatchOp = cast<PimCoreBatchOp>(op);
+    auto bbArg = dyn_cast<BlockArgument>(value);
+    if (!bbArg || bbArg.getOwner() != &coreBatchOp.getBody().front())
+      return failure();
+
+    Value tiedInput = coreBatchOp.getInputs()[bbArg.getArgNumber()];
+    if (auto memRefType = dyn_cast<BufferLikeType>(tiedInput.getType()))
+      return memRefType;
+
+    return bufferization::getBufferType(tiedInput, options, state, invocationStack);
+  }
+
+  LogicalResult bufferize(Operation* op,
+                          RewriterBase& rewriter,
+                          const BufferizationOptions& options,
+                          BufferizationState& state) const {
+    auto coreBatchOp = cast<PimCoreBatchOp>(op);
+
+    bool alreadyBufferized =
+      llvm::all_of(coreBatchOp.getWeights(), [](Value weight) { return isa<BufferLikeType>(weight.getType()); })
+      && llvm::all_of(coreBatchOp.getInputs(), [](Value input) { return isa<BufferLikeType>(input.getType()); })
+      && llvm::all_of(coreBatchOp.getBody().front().getArguments(),
+                      [](BlockArgument arg) { return isa<BufferLikeType>(arg.getType()); });
+    if (alreadyBufferized)
+      return success();
+
+    SmallVector<Value> weights;
+    SmallVector<Value> inputs;
+    weights.reserve(coreBatchOp.getWeights().size());
+    inputs.reserve(coreBatchOp.getInputs().size());
+
+    for (Value weight : coreBatchOp.getWeights()) {
+      if (isa<TensorType>(weight.getType())) {
+        auto weightOpt = getBufferOrValue(rewriter, weight, options, state);
+        if (failed(weightOpt))
+          return failure();
+        weights.push_back(*weightOpt);
+      }
+      else {
+        weights.push_back(weight);
+      }
+    }
+
+    for (Value input : coreBatchOp.getInputs()) {
+      if (isa<TensorType>(input.getType())) {
+        auto inputOpt = getBufferOrValue(rewriter, input, options, state);
+        if (failed(inputOpt))
+          return failure();
+        inputs.push_back(*inputOpt);
+      }
+      else {
+        inputs.push_back(input);
+      }
+    }
+
+    rewriter.setInsertionPoint(coreBatchOp);
+    auto newOp = PimCoreBatchOp::create(
+      rewriter, coreBatchOp.getLoc(), coreBatchOp.getLaneCountAttr(), ValueRange(weights), ValueRange(inputs));
+    newOp.getProperties().setOperandSegmentSizes({static_cast<int>(weights.size()), static_cast<int>(inputs.size())});
+    if (auto coreIdsAttr = coreBatchOp->getAttr(onnx_mlir::kCoreIdsAttrName))
+      newOp->setAttr(onnx_mlir::kCoreIdsAttrName, coreIdsAttr);
+
+    rewriter.inlineRegionBefore(coreBatchOp.getBody(), newOp.getBody(), newOp.getBody().begin());
+    for (Block& block : newOp.getBody())
+      if (failed(bufferization::bufferizeBlockSignature(&block, rewriter, options, state)))
+        return failure();
+
+    rewriter.eraseOp(coreBatchOp);
+    return success();
+  }
+};
+
 struct TransposeOpInterface : DstBufferizableOpInterfaceExternalModel<TransposeOpInterface, PimTransposeOp> {
   bool bufferizesToMemoryRead(Operation* op, OpOperand& opOperand, const AnalysisState& state) const {
     return !cast<DestinationStyleOpInterface>(op).isDpsInit(&opOperand);
@@ -133,11 +576,11 @@ struct TransposeOpInterface : DstBufferizableOpInterfaceExternalModel<TransposeO
                           BufferizationState& state) const {
     auto transposeOp = cast<PimTransposeOp>(op);
 
-    auto inputOpt = getBuffer(rewriter, transposeOp.getInput(), options, state);
+    auto inputOpt = getBufferOrValue(rewriter, transposeOp.getInput(), options, state);
     if (failed(inputOpt))
       return failure();
 
-    auto outputBufferOpt = getBuffer(rewriter, transposeOp.getOutputBuffer(), options, state);
+    auto outputBufferOpt = getBufferOrValue(rewriter, transposeOp.getOutputBuffer(), options, state);
     if (failed(outputBufferOpt))
       return failure();
 
@@ -170,11 +613,11 @@ struct VMMOpInterface : DstBufferizableOpInterfaceExternalModel<VMMOpInterface,
                           BufferizationState& state) const {
     auto vmmOp = cast<PimVMMOp>(op);
 
-    auto inputOpt = getBuffer(rewriter, vmmOp.getInput(), options, state);
+    auto inputOpt = getBufferOrValue(rewriter, vmmOp.getInput(), options, state);
     if (failed(inputOpt))
       return failure();
 
-    auto outputBufferOpt = getBuffer(rewriter, vmmOp.getOutputBuffer(), options, state);
+    auto outputBufferOpt = getBufferOrValue(rewriter, vmmOp.getOutputBuffer(), options, state);
     if (failed(outputBufferOpt))
       return failure();
 
@@ -197,11 +640,11 @@ struct MVMOpInterface : DstBufferizableOpInterfaceExternalModel<MVMOpInterface,
                           BufferizationState& state) const {
     auto mvmOp = cast<PimMVMOp>(op);
 
-    auto inputOpt = getBuffer(rewriter, mvmOp.getInput(), options, state);
+    auto inputOpt = getBufferOrValue(rewriter, mvmOp.getInput(), options, state);
     if (failed(inputOpt))
       return failure();
 
-    auto outputBufferOpt = getBuffer(rewriter, mvmOp.getOutputBuffer(), options, state);
+    auto outputBufferOpt = getBufferOrValue(rewriter, mvmOp.getOutputBuffer(), options, state);
     if (failed(outputBufferOpt))
       return failure();
 
@@ -231,15 +674,15 @@ struct BinaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<BinaryDstO
                           BufferizationState& state) const {
     auto binaryOp = cast<OpTy>(op);
 
-    auto lhsOpt = getBuffer(rewriter, binaryOp.getLhs(), options, state);
+    auto lhsOpt = getBufferOrValue(rewriter, binaryOp.getLhs(), options, state);
     if (failed(lhsOpt))
       return failure();
 
-    auto rhsOpt = getBuffer(rewriter, binaryOp.getRhs(), options, state);
+    auto rhsOpt = getBufferOrValue(rewriter, binaryOp.getRhs(), options, state);
     if (failed(rhsOpt))
       return failure();
 
-    auto outputBufferOpt = getBuffer(rewriter, binaryOp.getOutputBuffer(), options, state);
+    auto outputBufferOpt = getBufferOrValue(rewriter, binaryOp.getOutputBuffer(), options, state);
     if (failed(outputBufferOpt))
       return failure();
 
@@ -270,11 +713,11 @@ struct UnaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<UnaryDstOpI
                           BufferizationState& state) const {
     auto unaryOp = cast<OpTy>(op);
 
-    auto inputOpt = getBuffer(rewriter, unaryOp.getInput(), options, state);
+    auto inputOpt = getBufferOrValue(rewriter, unaryOp.getInput(), options, state);
     if (failed(inputOpt))
       return failure();
 
-    auto outputBufferOpt = getBuffer(rewriter, unaryOp.getOutputBuffer(), options, state);
+    auto outputBufferOpt = getBufferOrValue(rewriter, unaryOp.getOutputBuffer(), options, state);
     if (failed(outputBufferOpt))
       return failure();
 
@@ -287,8 +730,19 @@ struct UnaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<UnaryDstOpI
 
 void registerOpBufferizationInterfaces(DialectRegistry& registry) {
   registry.addExtension(+[](MLIRContext* ctx, PimDialect* dialect) {
+    PimEmptyManyOp::attachInterface<EmptyManyOpInterface>(*ctx);
+    PimMapOp::attachInterface<MapOpInterface>(*ctx);
+    PimCoreBatchOp::attachInterface<CoreBatchOpInterface>(*ctx);
     PimReceiveOp::attachInterface<ReceiveOpInterface>(*ctx);
+    PimReceiveManyOp::attachInterface<ReceiveManyOpInterface>(*ctx);
+    PimReceiveTensorOp::attachInterface<ReceiveTensorOpInterface>(*ctx);
+    PimReceiveBatchOp::attachInterface<ReceiveBatchOpInterface>(*ctx);
+    PimReceiveManyBatchOp::attachInterface<ReceiveManyBatchOpInterface>(*ctx);
+    PimSendTensorOp::attachInterface<SendTensorOpInterface>(*ctx);
+    PimExtractRowsOp::attachInterface<ExtractRowsOpInterface>(*ctx);
+    PimConcatOp::attachInterface<ConcatOpInterface>(*ctx);
     PimMemCopyHostToDevOp::attachInterface<MemCopyHostToDevOpInterface>(*ctx);
+    PimMemCopyHostToDevBatchOp::attachInterface<MemCopyHostToDevBatchOpInterface>(*ctx);
     PimMemCopyDevToHostOp::attachInterface<MemCopyDevToHostOpInterface>(*ctx);
     PimTransposeOp::attachInterface<TransposeOpInterface>(*ctx);
     PimVMMOp::attachInterface<VMMOpInterface>(*ctx);

src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp

@@ -3,12 +3,19 @@
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Threading.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
+
 #include "Common/PimCommon.hpp"
 #include "Compiler/PimCodeGen.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 
@@ -40,15 +47,54 @@ private:
 
 void PimBufferizationPass::runOnOperation() {
   auto moduleOp = getOperation();
+  // Refactor this into a function
+  {
+    auto funcOp = *getPimEntryFunc(moduleOp);
+
+    SmallVector<Operation*> coreOps;
+    funcOp->walk<WalkOrder::PreOrder>([&](Operation* op) {
+      if (isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
+        coreOps.push_back(op);
+    });
+    MLIRContext* ctx = moduleOp.getContext();
+    // failableParallelForEach will run the lambda in parallel and stop if any thread fails
+    LogicalResult result = mlir::failableParallelForEach(ctx, coreOps, [&](Operation* coreOp) {
+      // Again, allocate state LOCALLY per thread/function
+      bufferization::OneShotBufferizationOptions options;
+      options.allowUnknownOps = true;
+      if (isa<pim::PimCoreBatchOp>(coreOp))
+        options.opFilter.denyOperation([coreOp](Operation* op) { return op == coreOp; });
+      bufferization::BufferizationState state;
+      if (failed(bufferization::runOneShotBufferize(coreOp, options, state))) {
+        coreOp->emitError("Failed to bufferize PIM and Spatial ops");
+        return failure();
+      }
+      return success();
+    });
+
+    if (failed(result)) {
+      moduleOp.emitError("Failed to bufferize-parallel PIM and Spatial ops");
+      signalPassFailure();
+    }
+
+    funcOp->walk([&](bufferization::ToTensorOp toTensorOp) {
+      if (llvm::isa_and_present<pim::PimCoreOp, pim::PimCoreBatchOp>(toTensorOp->getParentOp()))
+        toTensorOp->setAttr("restrict", UnitAttr::get(ctx));
+    });
 
     // One-Shot-Bufferization
     bufferization::OneShotBufferizationOptions options;
     options.allowUnknownOps = true;
+    options.opFilter.denyOperation([](Operation* op) {
+      return op->getParentOfType<pim::PimCoreOp>() || op->getParentOfType<pim::PimCoreBatchOp>();
+    });
     bufferization::BufferizationState state;
+
     if (failed(bufferization::runOneShotBufferize(moduleOp, options, state))) {
       moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
       signalPassFailure();
     }
+  }
 
   MLIRContext* ctx = moduleOp.getContext();
   ConversionTarget target(*ctx);
@@ -57,7 +103,18 @@ void PimBufferizationPass::runOnOperation() {
   RewritePatternSet patterns(ctx);
   populateWithGenerated(patterns);
 
-  if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
+  // Only convert memref.copy → pim.memcp inside pim.core / pim.core_batch bodies.
+  // Host-level copies (e.g. from split/slice ops) must remain as memref.copy for CPU lowering.
+  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+  bool hasFailed = false;
+  moduleOp.walk<WalkOrder::PreOrder>([&](Operation* op) {
+    if (!isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
+      return WalkResult::advance();
+    if (failed(applyPartialConversion(op, target, frozenPatterns)))
+      hasFailed = true;
+    return WalkResult::skip();
+  });
+  if (hasFailed) {
     signalPassFailure();
     return;
   }
@@ -93,8 +150,8 @@ void PimBufferizationPass::runOnOperation() {
 }
 
 void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const {
-  funcOp.walk([&](PimCoreOp coreOp) {
-    walkPimMvmVmmWeightUses(coreOp, [&](OpOperand& weightUse) {
+  auto markWeights = [&](Operation* op) {
+    walkPimMvmVmmWeightUses(op, [&](OpOperand& weightUse) {
       Value weight = weightUse.get();
       auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
       if (!getGlobalOp)
@@ -104,7 +161,10 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO
       markWeightAlways(getGlobalOp);
       markWeightAlways(globalMemrefOp);
     });
-  });
+  };
+
+  funcOp.walk([&](PimCoreOp coreOp) { markWeights(coreOp); });
+  funcOp.walk([&](PimCoreBatchOp coreBatchOp) { markWeights(coreBatchOp); });
 }
 
 std::unique_ptr<Pass> createPimBufferizationPass() { return std::make_unique<PimBufferizationPass>(); }

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -2,8 +2,13 @@ add_onnx_mlir_dialect(Spatial spat)
 add_onnx_mlir_dialect_doc(spat Spatial.td)
 
 add_pim_library(SpatialOps
+  Channels.cpp
   SpatialOps.cpp
+  SpatialOpsAsm.cpp
+  SpatialOpsVerify.cpp
+  SpatialOpsCanonicalization.cpp
   Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
+  Transforms/MergeComputeNodes/RegularOpCompaction.cpp
   Transforms/MergeComputeNodes/DCPGraph/Graph.cpp
   Transforms/MergeComputeNodes/DCPGraph/GraphDebug.cpp
   Transforms/MergeComputeNodes/DCPGraph/GraphSupport.cpp

src/PIM/Dialect/Spatial/Channels.cpp

@@ -0,0 +1,120 @@
+#include "src/Accelerators/PIM/Dialect/Spatial/Channels.hpp"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Diagnostics.h"
+
+using namespace mlir;
+
+namespace onnx_mlir::spatial {
+
+namespace {
+
+static Channels::ChannelId getChannelId(SpatChannelSendOp sendOp) { return sendOp.getChannelId(); }
+
+static Channels::ChannelId getChannelId(SpatChannelReceiveOp receiveOp) { return receiveOp.getChannelId(); }
+
+static LogicalResult verifyEndpointPair(ChannelEndpoints endpoints) {
+  if (!endpoints.send || !endpoints.receive)
+    return failure();
+
+  if (endpoints.send.getSourceCoreId() != endpoints.receive.getSourceCoreId()) {
+    endpoints.send.emitOpError("sourceCoreId does not match paired spat.channel_receive");
+    return failure();
+  }
+  if (endpoints.send.getTargetCoreId() != endpoints.receive.getTargetCoreId()) {
+    endpoints.send.emitOpError("targetCoreId does not match paired spat.channel_receive");
+    return failure();
+  }
+  if (endpoints.send.getInput().getType() != endpoints.receive.getOutput().getType()) {
+    endpoints.send.emitOpError("input type does not match paired spat.channel_receive result type");
+    return failure();
+  }
+
+  return success();
+}
+
+} // namespace
+
+Channels::Channels(func::FuncOp funcOp) {
+  if (!funcOp)
+    return;
+
+  funcOp.walk([&](SpatChannelSendOp sendOp) { insertSend(sendOp); });
+  funcOp.walk([&](SpatChannelReceiveOp receiveOp) { insertReceive(receiveOp); });
+}
+
+Channels::ChannelId Channels::allocate() { return nextChannelId++; }
+
+void Channels::insertSend(SpatChannelSendOp sendOp) {
+  ChannelId channelId = getChannelId(sendOp);
+  nextChannelId = std::max(nextChannelId, channelId + 1);
+  endpoints[channelId].send = sendOp;
+}
+
+void Channels::insertReceive(SpatChannelReceiveOp receiveOp) {
+  ChannelId channelId = getChannelId(receiveOp);
+  nextChannelId = std::max(nextChannelId, channelId + 1);
+  endpoints[channelId].receive = receiveOp;
+}
+
+void Channels::eraseSend(SpatChannelSendOp sendOp) {
+  ChannelId channelId = getChannelId(sendOp);
+  auto it = endpoints.find(channelId);
+  if (it == endpoints.end())
+    return;
+  it->second.send = {};
+  if (!it->second.receive)
+    endpoints.erase(it);
+}
+
+void Channels::eraseReceive(SpatChannelReceiveOp receiveOp) {
+  ChannelId channelId = getChannelId(receiveOp);
+  auto it = endpoints.find(channelId);
+  if (it == endpoints.end())
+    return;
+  it->second.receive = {};
+  if (!it->second.send)
+    endpoints.erase(it);
+}
+
+FailureOr<ChannelEndpoints> Channels::lookup(ChannelId id) const {
+  auto it = endpoints.find(id);
+  if (it == endpoints.end())
+    return failure();
+  return it->second;
+}
+
+FailureOr<SpatChannelReceiveOp> Channels::getReceiveFor(SpatChannelSendOp sendOp) const {
+  auto endpointsOr = lookup(getChannelId(sendOp));
+  if (failed(endpointsOr) || !endpointsOr->receive)
+    return failure();
+  return endpointsOr->receive;
+}
+
+FailureOr<SpatChannelSendOp> Channels::getSendFor(SpatChannelReceiveOp receiveOp) const {
+  auto endpointsOr = lookup(getChannelId(receiveOp));
+  if (failed(endpointsOr) || !endpointsOr->send)
+    return failure();
+  return endpointsOr->send;
+}
+
+LogicalResult Channels::verify() const {
+  for (const auto& [channelId, pair] : endpoints) {
+    if (!pair.send || !pair.receive) {
+      if (pair.send) {
+        auto sendOp = pair.send;
+        sendOp.emitOpError("channel_id ") << channelId << " is missing a paired spat.channel_receive";
+      }
+      else if (pair.receive) {
+        auto receiveOp = pair.receive;
+        receiveOp.emitOpError("channel_id ") << channelId << " is missing a paired spat.channel_send";
+      }
+      return failure();
+    }
+    if (failed(verifyEndpointPair(pair)))
+      return failure();
+  }
+  return success();
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Channels.hpp

@@ -0,0 +1,43 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct ChannelEndpoints {
+  SpatChannelSendOp send;
+  SpatChannelReceiveOp receive;
+};
+
+class Channels {
+public:
+  using ChannelId = int64_t;
+
+  explicit Channels(mlir::func::FuncOp funcOp);
+
+  ChannelId allocate();
+
+  void insertSend(SpatChannelSendOp sendOp);
+  void insertReceive(SpatChannelReceiveOp receiveOp);
+  void eraseSend(SpatChannelSendOp sendOp);
+  void eraseReceive(SpatChannelReceiveOp receiveOp);
+
+  llvm::FailureOr<ChannelEndpoints> lookup(ChannelId id) const;
+  llvm::FailureOr<SpatChannelReceiveOp> getReceiveFor(SpatChannelSendOp sendOp) const;
+  llvm::FailureOr<SpatChannelSendOp> getSendFor(SpatChannelReceiveOp receiveOp) const;
+
+  mlir::LogicalResult verify() const;
+
+private:
+  ChannelId nextChannelId = 0;
+  llvm::DenseMap<ChannelId, ChannelEndpoints> endpoints;
+};
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Spatial.td

@@ -9,7 +9,6 @@ def SpatialDialect : Dialect {
   let name = "spat";
   let summary = "Dialect designed for deep learning computation in a spatial architecture";
   let cppNamespace = "::onnx_mlir::spatial";
-  let useDefaultTypePrinterParser = 1;
 }
 
 class SpatOp<string mnemonic, list<Trait> traits = []> :
@@ -19,15 +18,6 @@ class SpatOp<string mnemonic, list<Trait> traits = []> :
 def SpatTensor :
   AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;
 
-class SpatType<string name, string typeMnemonic, list<Trait> traits = []>
-    : TypeDef<SpatialDialect, name, traits> {
-  let mnemonic = typeMnemonic;
-}
-
-def SpatChannelType : SpatType<"SpatChannel", "ch"> {
-  let summary = "Virtual channel type";
-}
-
 //===----------------------------------------------------------------------===//
 // Execution
 //===----------------------------------------------------------------------===//
@@ -48,10 +38,27 @@ def SpatCompute : SpatOp<"compute", [SingleBlock, AttrSizedOperandSegments]> {
 
   let hasVerifier = 1;
   let hasFolder = 1;
+  let hasCustomAssemblyFormat = 1;
+}
 
-  let assemblyFormat = [{
-    `[` $weights `]` `(` $inputs `)` attr-dict `:` `[` type($weights) `]` `(` type($inputs) `)` `->` type($outputs) $body
-  }];
+def SpatComputeBatch : SpatOp<"compute_batch",
+    [SingleBlock, AttrSizedOperandSegments]> {
+  let summary = "Compressed batch of independent equivalent compute lanes";
+
+  let arguments = (ins
+    I32Attr:$laneCount,
+    Variadic<SpatTensor>:$weights,
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
 def SpatYieldOp : SpatOp<"yield", [Terminator]> {
@@ -61,51 +68,83 @@ def SpatYieldOp : SpatOp<"yield", [Terminator]> {
     Variadic<SpatTensor>:$outputs
   );
 
-  let assemblyFormat = [{
-    $outputs attr-dict `:` type($outputs)
-  }];
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatExtractRowsOp : SpatOp<"extract_rows", []> {
+  let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";
+
+  let arguments = (ins
+    SpatTensor:$input
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatConcatOp : SpatOp<"concat", []> {
+  let summary = "Concatenate tensors with compact Spatial operand syntax";
+
+  let arguments = (ins
+    I64Attr:$axis,
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatMapOp : SpatOp<"map", [SingleBlock]> {
+  let summary = "Apply the same lane-local region to many independent tensors";
+
+  let arguments = (ins
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Communication
 //===----------------------------------------------------------------------===//
 
-def SpatChannelNewOp : SpatOp<"channel_new", []> {
-  let summary = "Create a new virtual channel";
-
-  let results = (outs
-    SpatChannelType:$channel
-  );
-
-  let builders = [
-    OpBuilder<(ins ), [{
-      $_state.addTypes(SpatChannelType());
-    }]>
-  ];
-
-  let assemblyFormat = [{
-    attr-dict
-  }];
-}
-
 def SpatChannelSendOp : SpatOp<"channel_send", []> {
-  let summary = "Send a tensor through a channel";
+  let summary = "Send a tensor through a logical channel";
 
   let arguments = (ins
-    SpatChannelType:$channel,
+    I64Attr:$channelId,
+    I32Attr:$sourceCoreId,
+    I32Attr:$targetCoreId,
     SpatTensor:$input
   );
 
   let assemblyFormat = [{
-    $input `to` $channel attr-dict `:` `(` type($input) `->` type($channel) `)`
+    $input attr-dict `:` type($input)
   }];
 }
 
 def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
-  let summary = "Receive a tensor from a channel";
+  let summary = "Receive a tensor from a logical channel";
 
   let arguments = (ins
-    SpatChannelType:$channel
+    I64Attr:$channelId,
+    I32Attr:$sourceCoreId,
+    I32Attr:$targetCoreId
   );
 
   let results = (outs
@@ -113,44 +152,108 @@ def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
   );
 
   let assemblyFormat = [{
-    $channel attr-dict `:` `(` type($channel) `->` type($output) `)`
+    attr-dict `:` type($output)
   }];
 }
 
-def SpatChannelBroadcastSendOp : SpatOp<"channel_broadcast_send", []> {
-  let summary = "Broadcast a tensor through a shared channel buffer";
+def SpatChannelSendManyOp : SpatOp<"channel_send_many", []> {
+  let summary = "Send multiple tensors through logical channels";
 
   let arguments = (ins
-    SpatChannelType:$channel,
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatChannelReceiveManyOp : SpatOp<"channel_receive_many", []> {
+  let summary = "Receive multiple tensors from logical channels";
+
+  let arguments = (ins
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatChannelSendBatchOp : SpatOp<"channel_send_batch", []> {
+  let summary = "Send per-lane tensors through logical channels in a batch body";
+
+  let arguments = (ins
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds,
     SpatTensor:$input
   );
 
-  let assemblyFormat = [{
-    $input `to` $channel attr-dict `:` `(` type($input) `->` type($channel) `)`
-  }];
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
-def SpatChannelBroadcastReceiveOp : SpatOp<"channel_broadcast_receive", []> {
-  let summary = "Receive a tensor from a shared channel buffer";
+def SpatChannelSendManyBatchOp : SpatOp<"channel_send_many_batch", []> {
+  let summary = "Send multiple per-lane tensors through logical channels in a batch body";
 
   let arguments = (ins
-    SpatChannelType:$channel
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatChannelReceiveBatchOp : SpatOp<"channel_receive_batch", []> {
+  let summary = "Receive a per-lane tensor through logical channels in a batch body";
+
+  let arguments = (ins
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds
   );
 
   let results = (outs
     SpatTensor:$output
   );
 
-  let assemblyFormat = [{
-    $channel attr-dict `:` `(` type($channel) `->` type($output) `)`
-  }];
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatChannelReceiveManyBatchOp : SpatOp<"channel_receive_many_batch", []> {
+  let summary = "Receive multiple per-lane tensors through logical channels in a batch body";
+
+  let arguments = (ins
+    DenseI64ArrayAttr:$channelIds,
+    DenseI32ArrayAttr:$sourceCoreIds,
+    DenseI32ArrayAttr:$targetCoreIds
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Math
 //===----------------------------------------------------------------------===//
 
-def SpatWeightedVMMOp : SpatOp<"Wvmm", []> {
+def SpatVMMOp : SpatOp<"wvmm", []> {
   let summary = "Vector-matrix multiplication within a weighted compute operation";
 
   let arguments = (ins
@@ -169,7 +272,7 @@ def SpatWeightedVMMOp : SpatOp<"Wvmm", []> {
   }];
 }
 
-def SpatWeightedMVMOp : SpatOp<"Wmvm", []> {
+def SpatMVMOp : SpatOp<"Wmvm", []> {
   let summary = "Matrix-vector multiplication within a weighted compute operation";
 
   let arguments = (ins

src/PIM/Dialect/Spatial/SpatialOps.cpp

@@ -1,27 +1,3 @@
-#include "mlir/Dialect/Shape/IR/Shape.h"
-#include "mlir/IR/Block.h"
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/BuiltinTypeInterfaces.h"
-#include "mlir/IR/Diagnostics.h"
-#include "mlir/IR/DialectImplementation.h"
-#include "mlir/IR/IntegerSet.h"
-#include "mlir/IR/Matchers.h"
-#include "mlir/IR/OpDefinition.h"
-#include "mlir/IR/OpImplementation.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/TypeUtilities.h"
-#include "mlir/IR/Value.h"
-#include "mlir/Support/LLVM.h"
-
-#include "llvm/ADT/TypeSwitch.h"
-#include "llvm/Support/LogicalResult.h"
-
-#include <cstdint>
-
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -41,235 +17,6 @@ void SpatialDialect::initialize() {
 
     >();
 }
-
-inline LogicalResult mvmOpVerifySize2(SpatWeightedMVMOp* emitter,
-                                      ArrayRef<int64_t>& matrixShape,
-                                      ArrayRef<int64_t>& vectorShape,
-                                      ArrayRef<int64_t>& outputShape) {
-
-  // Verify that the matrix, vector and output shapes have rank 2
-  if (matrixShape.size() != 2 || vectorShape.size() != 2 || outputShape.size() != 2)
-    return emitter->emitError("matrix, vector and output must have rank 2");
-
-  // Verify that the matrix shape is (N, M)
-  int64_t N = matrixShape[0];
-  int64_t M = matrixShape[1];
-  if (N <= 0 || M <= 0)
-    return emitter->emitError("matrix shape must be (N, M) with N > 0 and M > 0");
-
-  // Verify that the vector shape is (M, 1)
-  int64_t vectorM = vectorShape[0];
-  int64_t vector1 = vectorShape[1];
-  if (vectorM != M || vector1 != 1)
-    return emitter->emitError("vector shape must be (M, 1)");
-
-  // Verify that the output shape is (N, 1)
-  int64_t outputN = outputShape[0];
-  int64_t output1 = outputShape[1];
-  if (outputN != N || output1 != 1)
-    return emitter->emitError("output shape must be (N, 1)");
-
-  return success();
-}
-
-inline LogicalResult mvmOpVerifySize4(SpatWeightedMVMOp* emitter,
-                                      ArrayRef<int64_t>& matrixShape,
-                                      ArrayRef<int64_t>& vectorShape,
-                                      ArrayRef<int64_t>& outputShape) {
-
-  // Verify that the matrix, vector and output shapes have rank 4
-  if (matrixShape.size() != 4 || vectorShape.size() != 4 || outputShape.size() != 4)
-    return emitter->emitError("matrix, vector and output must have rank 4");
-
-  // Verify that the matrix shape is (N, M, 1, 1)
-  int64_t N = matrixShape[0];
-  int64_t M = matrixShape[1];
-  int64_t matrix1First = matrixShape[2];
-  int64_t matrix1Second = matrixShape[3];
-  if (N <= 0 || M <= 0 || matrix1First != 1 || matrix1Second != 1)
-    return emitter->emitError("matrix shape must be (N, M, 1, 1) with N > 0 and M > 0");
-
-  // Verify that the vector shape is (1, M, 1, 1)
-  int64_t vector1First = vectorShape[0];
-  int64_t vectorM = vectorShape[1];
-  int64_t vector1Second = vectorShape[2];
-  int64_t vector1Third = vectorShape[3];
-  if (vector1First != 1 || vectorM != M || vector1Second != 1 || vector1Third != 1) {
-    if (vector1First == 1 && vector1Second == 1 && vector1Third == 1 && ignoreConcatError == true) {
-      // This is ok, it was caused by the simplification of the concat error
-    }
-    else {
-      return emitter->emitError("vector shape must be (1, M, 1, 1)");
-    }
-  }
-
-  // Verify that the output shape is (1, N, 1, 1)
-  int64_t output1First = outputShape[0];
-  int64_t outputN = outputShape[1];
-  int64_t output1Second = outputShape[2];
-  int64_t output1Third = outputShape[3];
-  if (output1First != 1 || outputN != N || output1Second != 1 || output1Third != 1)
-    return emitter->emitError("output shape must be (1, N, 1, 1)");
-
-  return success();
-}
-
-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();
-
-  if (auto coreOp = dyn_cast<pim::PimCoreOp>(weigthedOp->getParentOp()))
-    return cast<ShapedType>(coreOp.getWeights()[weightIndex].getType()).getShape();
-
-  return failure();
-}
-
-LogicalResult SpatWeightedMVMOp::verify() {
-  auto matrixShapeOpt = getWeightShapeForWeightedOp(this->getOperation(), this->getWeightIndex());
-  if (failed(matrixShapeOpt))
-    return emitError("SpatWeightedMVMOp was not within a SpatCompute or Core op");
-  auto matrixShape = *matrixShapeOpt;
-  auto vectorShape = getInput().getType().getShape();
-  auto outputShape = getOutput().getType().getShape();
-
-  /* Two possible accepted shapes:
-    1. matrix: (N, M); vector: (M, 1); output: (N, 1)
-    2. matrix: (N, M, 1, 1); vector: (1, M, 1, 1); output: (1, N, 1, 1)
-  */
-
-  if (matrixShape.size() == 2)
-    return mvmOpVerifySize2(this, matrixShape, vectorShape, outputShape);
-  else if (matrixShape.size() == 4)
-    return mvmOpVerifySize4(this, matrixShape, vectorShape, outputShape);
-  else
-    return emitError("matrix rank must be 2 or 4");
-}
-
-LogicalResult SpatWeightedVMMOp::verify() {
-  auto matrixShapeOpt = getWeightShapeForWeightedOp(this->getOperation(), this->getWeightIndex());
-  if (failed(matrixShapeOpt))
-    return emitError("SpatWeightedVMMOp was not within a SpatCompute or Core op");
-  auto matrixShape = *matrixShapeOpt;
-  auto vectorShape = getInput().getType().getShape();
-  auto outputShape = getOutput().getType().getShape();
-
-  /* Accepted shape:
-   1. vector: (1, N); matrix: (N, M); output: (1, M)
-  */
-  if (matrixShape.size() != 2 || vectorShape.size() != 2 || outputShape.size() != 2)
-    return emitError("matrix, vector and output must have rank 2");
-
-  int64_t N = matrixShape[0];
-  int64_t M = matrixShape[1];
-  if (N <= 0 || M <= 0)
-    return emitError("matrix shape must be (N, M) with N > 0 and M > 0");
-
-  int64_t vector1 = vectorShape[0];
-  int64_t vectorN = vectorShape[1];
-  if (vectorN != N || vector1 != 1)
-    return emitError("vector shape must be (N, 1)");
-
-  int64_t output1 = outputShape[0];
-  int64_t outputM = outputShape[1];
-  if (outputM != M || output1 != 1)
-    return emitError("output shape must be (M, 1)");
-
-  return success();
-}
-
-LogicalResult SpatVAddOp::verify() {
-  // At least two operands
-  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
-    return failure();
-
-  return OpTrait::impl::verifySameOperandsAndResultType(*this);
-}
-
-LogicalResult SpatVMaxOp::verify() {
-  // At least two operands
-  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
-    return failure();
-
-  return OpTrait::impl::verifySameOperandsAndResultType(*this);
-}
-
-LogicalResult SpatCompute::verify() {
-  // Check that the terminator yields the same number and types as the compute results.
-  auto& block = getBody().front();
-  if (block.mightHaveTerminator()) {
-    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
-    if (!yieldOp)
-      return emitError("ComputeOp must have a single yield operation");
-
-    auto resultTypes = getResultTypes();
-    auto yieldTypes = yieldOp->getOperandTypes();
-    if (resultTypes.size() != yieldTypes.size()) {
-      return emitError("ComputeOp must have same number of results as yieldOp "
-                       "operands");
-    }
-
-    for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldTypes))) {
-      auto resultType = std::get<0>(it);
-      auto yieldType = std::get<1>(it);
-
-      // Same type and compatible shape
-      if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType))) {
-        return emitError("ComputeOp output must be of the same type as yieldOp "
-                         "operand");
-      }
-
-      // Same encoding
-      if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
-        if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
-          if (resultRankedType.getEncoding() != yieldRankedType.getEncoding()) {
-            return emitError("ComputeOp output must have the same encoding as "
-                             "yieldOp operand");
-          }
-        }
-        else {
-          return emitError("ComputeOp output has an encoding while yieldOp "
-                           "operand does not have one");
-        }
-      }
-      else {
-        // If result does not have an encoding, yield shouldn't either
-        if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
-          return emitError("ComputeOp output must not have an encoding if "
-                           "yieldOp operand has one");
-        }
-      }
-    }
-  }
-
-  // Check that each block argument is used
-  for (auto arg : block.getArguments())
-    if (arg.use_empty())
-      return emitError("ComputeOp block argument is not used");
-
-  return success();
-}
-
-LogicalResult SpatCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
-  Block& block = getBody().front();
-  if (!llvm::hasSingleElement(block))
-    return failure();
-
-  auto yieldOp = dyn_cast<SpatYieldOp>(block.front());
-  if (!yieldOp)
-    return failure();
-
-  for (Value yieldedValue : yieldOp.getOperands()) {
-    if (auto blockArg = dyn_cast<BlockArgument>(yieldedValue)) {
-      if (blockArg.getOwner() == &block) {
-        results.push_back(getOperand(blockArg.getArgNumber()));
-        continue;
-      }
-    }
-    results.push_back(yieldedValue);
-  }
-  return success();
-}
-
 } // namespace spatial
 } // namespace onnx_mlir
 

src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp

@@ -0,0 +1,35 @@
+#include "mlir/IR/Block.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace spatial {
+
+LogicalResult SpatCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
+  Block& block = getBody().front();
+  if (!llvm::hasSingleElement(block))
+    return failure();
+
+  auto yieldOp = dyn_cast<SpatYieldOp>(block.front());
+  if (!yieldOp)
+    return failure();
+
+  for (Value yieldedValue : yieldOp.getOperands()) {
+    if (auto blockArg = dyn_cast<BlockArgument>(yieldedValue)) {
+      if (blockArg.getOwner() == &block) {
+        results.push_back(getOperand(blockArg.getArgNumber()));
+        continue;
+      }
+    }
+    results.push_back(yieldedValue);
+  }
+  return success();
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -0,0 +1,502 @@
+#include "mlir/IR/Block.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Support/LLVM.h"
+
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace spatial {
+
+namespace {
+
+inline LogicalResult mvmOpVerifySize2(SpatMVMOp* emitter,
+                                      ArrayRef<int64_t>& matrixShape,
+                                      ArrayRef<int64_t>& vectorShape,
+                                      ArrayRef<int64_t>& outputShape) {
+  if (matrixShape.size() != 2 || vectorShape.size() != 2 || outputShape.size() != 2)
+    return emitter->emitError("matrix, vector and output must have rank 2");
+
+  int64_t N = matrixShape[0];
+  int64_t M = matrixShape[1];
+  if (N <= 0 || M <= 0)
+    return emitter->emitError("matrix shape must be (N, M) with N > 0 and M > 0");
+
+  int64_t vectorM = vectorShape[0];
+  int64_t vector1 = vectorShape[1];
+  if (vectorM != M || vector1 != 1)
+    return emitter->emitError("vector shape must be (M, 1)");
+
+  int64_t outputN = outputShape[0];
+  int64_t output1 = outputShape[1];
+  if (outputN != N || output1 != 1)
+    return emitter->emitError("output shape must be (N, 1)");
+
+  return success();
+}
+
+inline LogicalResult mvmOpVerifySize4(SpatMVMOp* emitter,
+                                      ArrayRef<int64_t>& matrixShape,
+                                      ArrayRef<int64_t>& vectorShape,
+                                      ArrayRef<int64_t>& outputShape) {
+  if (matrixShape.size() != 4 || vectorShape.size() != 4 || outputShape.size() != 4)
+    return emitter->emitError("matrix, vector and output must have rank 4");
+
+  int64_t N = matrixShape[0];
+  int64_t M = matrixShape[1];
+  int64_t matrix1First = matrixShape[2];
+  int64_t matrix1Second = matrixShape[3];
+  if (N <= 0 || M <= 0 || matrix1First != 1 || matrix1Second != 1)
+    return emitter->emitError("matrix shape must be (N, M, 1, 1) with N > 0 and M > 0");
+
+  int64_t vector1First = vectorShape[0];
+  int64_t vectorM = vectorShape[1];
+  int64_t vector1Second = vectorShape[2];
+  int64_t vector1Third = vectorShape[3];
+  if (vector1First != 1 || vectorM != M || vector1Second != 1 || vector1Third != 1) {
+    if (vector1First == 1 && vector1Second == 1 && vector1Third == 1 && ignoreConcatError == true) {
+      // This is ok, it was caused by the simplification of the concat error.
+    }
+    else {
+      return emitter->emitError("vector shape must be (1, M, 1, 1)");
+    }
+  }
+
+  int64_t output1First = outputShape[0];
+  int64_t outputN = outputShape[1];
+  int64_t output1Second = outputShape[2];
+  int64_t output1Third = outputShape[3];
+  if (output1First != 1 || outputN != N || output1Second != 1 || output1Third != 1)
+    return emitter->emitError("output shape must be (1, N, 1, 1)");
+
+  return success();
+}
+
+static FailureOr<ArrayRef<int64_t>> getWeightShapeForWeightedOp(Operation* weightedOp, size_t weightIndex) {
+  if (auto computeOp = weightedOp->getParentOfType<SpatCompute>())
+    return cast<ShapedType>(computeOp.getWeights()[weightIndex].getType()).getShape();
+
+  if (auto coreOp = weightedOp->getParentOfType<pim::PimCoreOp>())
+    return cast<ShapedType>(coreOp.getWeights()[weightIndex].getType()).getShape();
+
+  if (auto batchOp = weightedOp->getParentOfType<SpatComputeBatch>()) {
+    if (batchOp.getWeights().empty() || weightIndex >= batchOp.getWeights().size())
+      return failure();
+    return cast<ShapedType>(batchOp.getWeights()[weightIndex].getType()).getShape();
+  }
+
+  return failure();
+}
+
+static FailureOr<int32_t> getParentBatchLaneCount(Operation* op) {
+  auto batchOp = op->getParentOfType<SpatComputeBatch>();
+  if (!batchOp)
+    return failure();
+  return batchOp.getLaneCount();
+}
+
+static LogicalResult verifyManyChannelSizes(Operation* op,
+                                            ArrayRef<int64_t> channelIds,
+                                            ArrayRef<int32_t> sourceCoreIds,
+                                            ArrayRef<int32_t> targetCoreIds,
+                                            size_t valueCount) {
+  if (channelIds.size() != sourceCoreIds.size() || channelIds.size() != targetCoreIds.size())
+    return op->emitError("channelIds, sourceCoreIds, and targetCoreIds must have the same length");
+  if (channelIds.size() != valueCount)
+    return op->emitError("channel metadata length must match the number of values");
+  return success();
+}
+
+static LogicalResult verifyManyChannelTypes(Operation* op, TypeRange types, StringRef kind) {
+  if (types.empty())
+    return op->emitError() << kind << " must carry at least one value";
+
+  Type firstType = types.front();
+  for (Type type : types.drop_front())
+    if (type != firstType)
+      return op->emitError() << kind << " values must all have the same type";
+  return success();
+}
+
+static LogicalResult verifyBatchChannelSizes(Operation* op,
+                                             ArrayRef<int64_t> channelIds,
+                                             ArrayRef<int32_t> sourceCoreIds,
+                                             ArrayRef<int32_t> targetCoreIds) {
+  if (channelIds.size() != sourceCoreIds.size() || channelIds.size() != targetCoreIds.size())
+    return op->emitError("channelIds, sourceCoreIds, and targetCoreIds must have the same length");
+
+  auto laneCount = getParentBatchLaneCount(op);
+  if (failed(laneCount))
+    return op->emitError("must be nested inside spat.compute_batch");
+  if (channelIds.size() != static_cast<size_t>(*laneCount))
+    return op->emitError("channel metadata length must match parent laneCount");
+
+  return success();
+}
+
+static LogicalResult verifyManyBatchChannelSizes(Operation* op,
+                                                 ArrayRef<int64_t> channelIds,
+                                                 ArrayRef<int32_t> sourceCoreIds,
+                                                 ArrayRef<int32_t> targetCoreIds,
+                                                 size_t valueCount) {
+  if (channelIds.size() != sourceCoreIds.size() || channelIds.size() != targetCoreIds.size())
+    return op->emitError("channelIds, sourceCoreIds, and targetCoreIds must have the same length");
+
+  auto laneCount = getParentBatchLaneCount(op);
+  if (failed(laneCount))
+    return op->emitError("must be nested inside spat.compute_batch");
+  if (channelIds.size() != valueCount * static_cast<size_t>(*laneCount))
+    return op->emitError("channel metadata length must match the number of values times parent laneCount");
+
+  return success();
+}
+
+static LogicalResult verifyBatchBody(Operation* op, Block& block, TypeRange outputTypes, size_t weightsPerLane) {
+  auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
+  if (!yieldOp)
+    return op->emitError("body must terminate with spat.yield");
+  if (outputTypes.empty()) {
+    if (yieldOp.getNumOperands() != 0)
+      return op->emitError("body yield must be empty when compute_batch has no results");
+  }
+  else {
+    if (yieldOp.getNumOperands() != 1)
+      return op->emitError("body yield must produce exactly one value");
+    if (yieldOp.getOperand(0).getType() != outputTypes[0])
+      return op->emitError("body yield type must match output type");
+  }
+
+  for (auto& bodyOp : block) {
+    if (auto wvmm = dyn_cast<SpatVMMOp>(&bodyOp))
+      if (wvmm.getWeightIndex() < 0 || static_cast<size_t>(wvmm.getWeightIndex()) >= weightsPerLane)
+        return op->emitError("compute_batch body Wvmm weightIndex is out of range for one lane");
+    if (auto wmvm = dyn_cast<SpatMVMOp>(&bodyOp))
+      if (wmvm.getWeightIndex() < 0 || static_cast<size_t>(wmvm.getWeightIndex()) >= weightsPerLane)
+        return op->emitError("compute_batch body Wmvm weightIndex is out of range for one lane");
+  }
+  return success();
+}
+
+} // namespace
+
+LogicalResult SpatMVMOp::verify() {
+  auto matrixShapeOpt = getWeightShapeForWeightedOp(this->getOperation(), this->getWeightIndex());
+  if (failed(matrixShapeOpt))
+    return emitError("SpatMVMOp was not within a SpatCompute or Core op");
+  auto matrixShape = *matrixShapeOpt;
+  auto vectorShape = getInput().getType().getShape();
+  auto outputShape = getOutput().getType().getShape();
+
+  if (matrixShape.size() == 2)
+    return mvmOpVerifySize2(this, matrixShape, vectorShape, outputShape);
+  if (matrixShape.size() == 4)
+    return mvmOpVerifySize4(this, matrixShape, vectorShape, outputShape);
+  return emitError("matrix rank must be 2 or 4");
+}
+
+LogicalResult SpatVMMOp::verify() {
+  auto matrixShapeOpt = getWeightShapeForWeightedOp(this->getOperation(), this->getWeightIndex());
+  if (failed(matrixShapeOpt))
+    return emitError("SpatVMMOp was not within a SpatCompute or Core op");
+  auto matrixShape = *matrixShapeOpt;
+  auto vectorShape = getInput().getType().getShape();
+  auto outputShape = getOutput().getType().getShape();
+
+  if (matrixShape.size() != 2 || vectorShape.size() != 2 || outputShape.size() != 2)
+    return emitError("matrix, vector and output must have rank 2");
+
+  int64_t N = matrixShape[0];
+  int64_t M = matrixShape[1];
+  if (N <= 0 || M <= 0)
+    return emitError("matrix shape must be (N, M) with N > 0 and M > 0");
+
+  int64_t vector1 = vectorShape[0];
+  int64_t vectorN = vectorShape[1];
+  if (vectorN != N || vector1 != 1)
+    return emitError("vector shape must be (1, N)");
+
+  int64_t output1 = outputShape[0];
+  int64_t outputM = outputShape[1];
+  if (outputM != M || output1 != 1)
+    return emitError("output shape must be (1, M)");
+
+  return success();
+}
+
+LogicalResult SpatVAddOp::verify() {
+  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
+    return failure();
+  return OpTrait::impl::verifySameOperandsAndResultType(*this);
+}
+
+LogicalResult SpatVMaxOp::verify() {
+  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
+    return failure();
+  return OpTrait::impl::verifySameOperandsAndResultType(*this);
+}
+
+LogicalResult SpatExtractRowsOp::verify() {
+  auto inputType = dyn_cast<ShapedType>(getInput().getType());
+  if (!inputType || !inputType.hasRank() || inputType.getRank() != 2)
+    return emitError("input must be a rank-2 shaped type");
+
+  int64_t numRows = inputType.getShape()[0];
+  int64_t numCols = inputType.getShape()[1];
+  Type elementType = inputType.getElementType();
+
+  if (numRows >= 0 && static_cast<int64_t>(getNumResults()) != numRows)
+    return emitError("number of outputs must match the number of input rows");
+
+  for (Type output : getResultTypes()) {
+    auto outputType = dyn_cast<ShapedType>(output);
+    if (!outputType || !outputType.hasRank() || outputType.getRank() != 2)
+      return emitError("outputs must all be rank-2 shaped types");
+    if (outputType.getElementType() != elementType)
+      return emitError("output element types must match input element type");
+    auto outputShape = outputType.getShape();
+    if (outputShape[0] != 1)
+      return emitError("each output must have exactly one row");
+    if (numCols >= 0 && outputShape[1] != numCols)
+      return emitError("output column count must match input column count");
+  }
+
+  return success();
+}
+
+LogicalResult SpatConcatOp::verify() {
+  if (getInputs().empty())
+    return emitError("requires at least one input");
+
+  auto outputType = dyn_cast<ShapedType>(getOutput().getType());
+  if (!outputType || !outputType.hasRank())
+    return emitError("output must be a ranked shaped type");
+
+  int64_t axis = getAxis();
+  int64_t rank = outputType.getRank();
+  if (axis < 0 || axis >= rank)
+    return emitError("axis must be within the output rank");
+
+  int64_t concatenatedDimSize = 0;
+  bool concatenatedDimDynamic = false;
+  Type outputElementType = outputType.getElementType();
+
+  for (Value input : getInputs()) {
+    auto inputType = dyn_cast<ShapedType>(input.getType());
+    if (!inputType || !inputType.hasRank())
+      return emitError("inputs must be ranked shaped types");
+    if (inputType.getRank() != rank)
+      return emitError("all inputs must have the same rank as the output");
+    if (inputType.getElementType() != outputElementType)
+      return emitError("all inputs must have the same element type as the output");
+
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      if (dim == axis)
+        continue;
+      int64_t inputDim = inputType.getDimSize(dim);
+      int64_t outputDim = outputType.getDimSize(dim);
+      if (!ShapedType::isDynamic(inputDim) && !ShapedType::isDynamic(outputDim) && inputDim != outputDim)
+        return emitError("non-concatenated dimensions must match the output shape");
+    }
+
+    int64_t inputConcatDim = inputType.getDimSize(axis);
+    if (ShapedType::isDynamic(inputConcatDim)) {
+      concatenatedDimDynamic = true;
+      continue;
+    }
+    concatenatedDimSize += inputConcatDim;
+  }
+
+  int64_t outputConcatDim = outputType.getDimSize(axis);
+  if (!concatenatedDimDynamic && !ShapedType::isDynamic(outputConcatDim) && concatenatedDimSize != outputConcatDim)
+    return emitError("output concatenated dimension must equal the sum of input sizes");
+
+  return success();
+}
+
+LogicalResult SpatMapOp::verify() {
+  if (getInputs().empty())
+    return emitError("requires at least one input");
+  if (getOutputs().size() != getInputs().size())
+    return emitError("number of outputs must match number of inputs");
+
+  Type inputType = getInputs().front().getType();
+  for (Value input : getInputs().drop_front())
+    if (input.getType() != inputType)
+      return emitError("all inputs must have the same type");
+
+  Type outputType = getOutputs().front().getType();
+  for (Value output : getOutputs().drop_front())
+    if (output.getType() != outputType)
+      return emitError("all outputs must have the same type");
+
+  Block& block = getBody().front();
+  if (block.getNumArguments() != 1)
+    return emitError("body must have exactly one block argument");
+  if (block.getArgument(0).getType() != inputType)
+    return emitError("body block argument type must match input type");
+
+  auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
+  if (!yieldOp)
+    return emitError("body must terminate with spat.yield");
+  if (yieldOp.getNumOperands() != 1)
+    return emitError("body yield must produce exactly one value");
+  if (yieldOp.getOperand(0).getType() != outputType)
+    return emitError("body yield type must match output type");
+
+  return success();
+}
+
+LogicalResult SpatCompute::verify() {
+  auto& block = getBody().front();
+  if (block.mightHaveTerminator()) {
+    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
+    if (!yieldOp)
+      return emitError("ComputeOp must have a single yield operation");
+
+    auto resultTypes = getResultTypes();
+    auto yieldTypes = yieldOp->getOperandTypes();
+    if (resultTypes.size() != yieldTypes.size())
+      return emitError("ComputeOp must have same number of results as yieldOp operands");
+
+    for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldTypes))) {
+      auto resultType = std::get<0>(it);
+      auto yieldType = std::get<1>(it);
+
+      if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
+        return emitError("ComputeOp output must be of the same type as yieldOp operand");
+
+      if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
+        if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
+          if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
+            return emitError("ComputeOp output must have the same encoding as yieldOp operand");
+        }
+        else {
+          return emitError("ComputeOp output has an encoding while yieldOp operand does not have one");
+        }
+      }
+      else if (dyn_cast<RankedTensorType>(yieldType)) {
+        return emitError("ComputeOp output must not have an encoding if yieldOp operand has one");
+      }
+    }
+  }
+
+  for (auto arg : block.getArguments())
+    if (arg.use_empty())
+      return emitError("ComputeOp block argument is not used");
+
+  return success();
+}
+
+LogicalResult SpatChannelSendManyOp::verify() {
+  if (failed(verifyManyChannelSizes(
+        getOperation(), getChannelIds(), getSourceCoreIds(), getTargetCoreIds(), getInputs().size())))
+    return failure();
+  return verifyManyChannelTypes(getOperation(), getInputs().getTypes(), "channel_send_many");
+}
+
+LogicalResult SpatChannelReceiveManyOp::verify() {
+  if (failed(verifyManyChannelSizes(
+        getOperation(), getChannelIds(), getSourceCoreIds(), getTargetCoreIds(), getOutputs().size())))
+    return failure();
+  return verifyManyChannelTypes(getOperation(), getOperation()->getResultTypes(), "channel_receive_many");
+}
+
+LogicalResult SpatChannelSendBatchOp::verify() {
+  return verifyBatchChannelSizes(getOperation(), getChannelIds(), getSourceCoreIds(), getTargetCoreIds());
+}
+
+LogicalResult SpatChannelSendManyBatchOp::verify() {
+  if (failed(verifyManyBatchChannelSizes(
+        getOperation(), getChannelIds(), getSourceCoreIds(), getTargetCoreIds(), getInputs().size())))
+    return failure();
+  return verifyManyChannelTypes(getOperation(), getInputs().getTypes(), "channel_send_many_batch");
+}
+
+LogicalResult SpatChannelReceiveBatchOp::verify() {
+  return verifyBatchChannelSizes(getOperation(), getChannelIds(), getSourceCoreIds(), getTargetCoreIds());
+}
+
+LogicalResult SpatChannelReceiveManyBatchOp::verify() {
+  if (failed(verifyManyBatchChannelSizes(
+        getOperation(), getChannelIds(), getSourceCoreIds(), getTargetCoreIds(), getOutputs().size())))
+    return failure();
+  return verifyManyChannelTypes(getOperation(), getOperation()->getResultTypes(), "channel_receive_many_batch");
+}
+
+LogicalResult SpatComputeBatch::verify() {
+  int32_t count = getLaneCount();
+  if (count <= 0)
+    return emitError("laneCount must be positive");
+
+  auto laneCountSz = static_cast<size_t>(count);
+  if (getWeights().size() % laneCountSz != 0)
+    return emitError("number of weights must be a multiple of laneCount");
+
+  if (!getInputs().empty() && getInputs().size() != laneCountSz)
+    return emitError("number of inputs must be either 0 or laneCount");
+  if (!getOutputs().empty() && getOutputs().size() != laneCountSz)
+    return emitError("number of outputs must be either 0 or laneCount");
+
+  size_t weightsPerLane = getWeights().size() / laneCountSz;
+  for (size_t weightIndex = 0; weightIndex < weightsPerLane; ++weightIndex) {
+    Type weightType = getWeights()[weightIndex].getType();
+    for (size_t lane = 1; lane < laneCountSz; ++lane)
+      if (getWeights()[lane * weightsPerLane + weightIndex].getType() != weightType)
+        return emitError("corresponding weights across lanes must have the same type");
+  }
+
+  if (!getInputs().empty()) {
+    Type inputType = getInputs()[0].getType();
+    for (Value in : getInputs().drop_front())
+      if (in.getType() != inputType)
+        return emitError("all inputs must have the same type");
+  }
+
+  if (!getOutputs().empty()) {
+    Type outputType = getOutputs()[0].getType();
+    for (Value out : getOutputs().drop_front())
+      if (out.getType() != outputType)
+        return emitError("all outputs must have the same type");
+  }
+
+  if (auto coreIdAttr = (*this)->getAttr(onnx_mlir::kCoreIdsAttrName)) {
+    auto coreIdsAttr = dyn_cast<DenseI32ArrayAttr>(coreIdAttr);
+    if (!coreIdsAttr)
+      return emitError("compute_batch coreIds attribute must be a dense i32 array");
+    if (coreIdsAttr.size() != laneCountSz)
+      return emitError("compute_batch coreIds array length must match laneCount");
+    if (llvm::any_of(coreIdsAttr.asArrayRef(), [](int32_t coreId) { return coreId <= 0; }))
+      return emitError("compute_batch coreIds values must be positive");
+    llvm::SmallDenseSet<int32_t, 8> seenCoreIds;
+    for (int32_t coreId : coreIdsAttr.asArrayRef())
+      if (!seenCoreIds.insert(coreId).second)
+        return emitError("compute_batch coreIds values must be distinct");
+  }
+
+  Block& block = getBody().front();
+  if (getInputs().empty()) {
+    if (block.getNumArguments() != 0)
+      return emitError("compute_batch body must have no block arguments when there are no inputs");
+  }
+  else {
+    if (block.getNumArguments() != 1)
+      return emitError("compute_batch body must have exactly one block argument");
+    if (block.getArgument(0).getType() != getInputs()[0].getType())
+      return emitError("body block argument type must match input type");
+  }
+
+  return verifyBatchBody(getOperation(), block, getResultTypes(), weightsPerLane);
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -28,6 +28,8 @@ namespace spatial {
 using namespace mlir;
 
 namespace {
+using SpatCompute = onnx_mlir::spatial::SpatCompute;
+using SpatComputeBatch = onnx_mlir::spatial::SpatComputeBatch;
 
 struct VirtualNode {
   SmallVector<size_t, 4> originalComputeIndices;
@@ -54,6 +56,43 @@ struct WindowScheduleResult {
   size_t maxMergeGroupSize = 0;
 };
 
+size_t getSchedulingCpuBudget() {
+  if (coresCount.getValue() > 0)
+    return static_cast<size_t>(coresCount.getValue());
+  return std::numeric_limits<size_t>::max();
+}
+
+size_t getBatchChunkTargetCount(int32_t laneCount) {
+  assert(laneCount > 0 && "laneCount must be positive");
+  return std::min(static_cast<size_t>(laneCount), std::max<size_t>(1, getSchedulingCpuBudget()));
+}
+
+ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex) {
+  size_t totalLanes = static_cast<size_t>(batch.getLaneCount());
+  size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
+  size_t baseChunkSize = totalLanes / chunkCount;
+  size_t largeChunkCount = totalLanes % chunkCount;
+
+  size_t laneStart = chunkIndex * baseChunkSize + std::min(chunkIndex, largeChunkCount);
+  size_t laneCount = baseChunkSize + (chunkIndex < largeChunkCount ? 1 : 0);
+  return {batch.getOperation(), static_cast<uint32_t>(laneStart), static_cast<uint32_t>(laneCount)};
+}
+
+ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) {
+  size_t totalLanes = static_cast<size_t>(batch.getLaneCount());
+  size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
+  size_t baseChunkSize = totalLanes / chunkCount;
+  size_t largeChunkCount = totalLanes % chunkCount;
+  size_t largeChunkSpan = largeChunkCount * (baseChunkSize + 1);
+
+  size_t chunkIndex = 0;
+  if (static_cast<size_t>(lane) < largeChunkSpan)
+    chunkIndex = static_cast<size_t>(lane) / (baseChunkSize + 1);
+  else
+    chunkIndex = largeChunkCount + (static_cast<size_t>(lane) - largeChunkSpan) / baseChunkSize;
+  return getBatchChunkForIndex(batch, chunkIndex);
+}
+
 std::vector<IndexedEdge> aggregateEdges(ArrayRef<IndexedEdge> edges) {
   llvm::DenseMap<std::pair<size_t, size_t>, Weight> edgeWeights;
   for (auto [start, end, weight] : edges) {
@@ -81,14 +120,122 @@ std::vector<IndexedEdge> aggregateEdges(ArrayRef<IndexedEdge> edges) {
   return aggregatedEdges;
 }
 
-VirtualGraph buildInitialVirtualGraph(ArrayRef<SpatCompute> spatComputes, ArrayRef<IndexedEdge> edges) {
+Weight getComputeBodyWeight(Region& body) {
+  constexpr Weight kOperationWeight = 100;
+  Weight numOperations = 0;
+  for (auto& block : body)
+    for ([[maybe_unused]] auto& op : block)
+      numOperations = checkedAdd(numOperations, static_cast<Weight>(1));
+  return checkedMultiply(numOperations, kOperationWeight);
+}
+
+CrossbarUsage getComputeBodyCrossbarUsage(Region& body) {
+  CrossbarUsage crossbarUsage = 0;
+  for (auto& block : body)
+    for (auto& op : block)
+      if (isa<SpatVMMOp>(op))
+        crossbarUsage = checkedAdd(crossbarUsage, static_cast<CrossbarUsage>(1));
+  return crossbarUsage;
+}
+
+Weight getComputeInstanceWeight(const ComputeInstance& instance) {
+  if (auto spatCompute = dyn_cast<SpatCompute>(instance.op))
+    return getSpatComputeWeight(spatCompute);
+  auto batch = cast<SpatComputeBatch>(instance.op);
+  return checkedMultiply(getComputeBodyWeight(batch.getBody()), static_cast<Weight>(instance.laneCount));
+}
+
+CrossbarUsage getComputeInstanceCrossbarUsage(const ComputeInstance& instance) {
+  if (auto spatCompute = dyn_cast<SpatCompute>(instance.op))
+    return getSpatComputeCrossbarUsage(spatCompute);
+  auto batch = cast<SpatComputeBatch>(instance.op);
+  return checkedMultiply(getComputeBodyCrossbarUsage(batch.getBody()), static_cast<CrossbarUsage>(instance.laneCount));
+}
+
+SmallVector<Value, 4> getComputeInstanceInputs(const ComputeInstance& instance) {
+  if (auto spatCompute = dyn_cast<SpatCompute>(instance.op))
+    return SmallVector<Value, 4>(spatCompute.getInputs().begin(), spatCompute.getInputs().end());
+  auto batch = cast<SpatComputeBatch>(instance.op);
+  SmallVector<Value, 4> inputs;
+  inputs.reserve(instance.laneCount);
+  for (uint32_t lane = instance.laneStart; lane < instance.laneStart + instance.laneCount; ++lane)
+    inputs.push_back(batch.getInputs()[lane]);
+  return inputs;
+}
+
+std::optional<ComputeInstance> getOriginalComputeInstance(Value value) {
+  Operation* op = value.getDefiningOp();
+  if (!op)
+    return std::nullopt;
+
+  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
+    value = extract.getSource();
+    op = value.getDefiningOp();
+    if (!op)
+      return std::nullopt;
+  }
+
+  if (auto spatCompute = dyn_cast<SpatCompute>(op))
+    return ComputeInstance {spatCompute.getOperation(), 0, 1};
+  if (auto batch = dyn_cast<SpatComputeBatch>(op))
+    return getBatchChunkForLane(batch, static_cast<uint32_t>(cast<OpResult>(value).getResultNumber()));
+  return std::nullopt;
+}
+
+SmallVector<ComputeInstance> collectComputeInstances(Operation* entryOp) {
+  SmallVector<ComputeInstance> instances;
+  auto isUsedAsWeightOnly = [](Operation* producerOp) {
+    if (producerOp->getNumResults() == 0)
+      return false;
+    for (Value result : producerOp->getResults()) {
+      if (result.use_empty())
+        return false;
+      for (Operation* user : result.getUsers()) {
+        if (auto compute = dyn_cast<SpatCompute>(user)) {
+          if (!llvm::is_contained(compute.getWeights(), result))
+            return false;
+          continue;
+        }
+        if (auto batch = dyn_cast<SpatComputeBatch>(user)) {
+          if (!llvm::is_contained(batch.getWeights(), result))
+            return false;
+          continue;
+        }
+        return false;
+      }
+    }
+    return true;
+  };
+  for (Region& region : entryOp->getRegions()) {
+    for (Block& block : region) {
+      for (Operation& op : block) {
+        if (auto spatCompute = dyn_cast<SpatCompute>(&op)) {
+          if (isUsedAsWeightOnly(spatCompute.getOperation()))
+            continue;
+          instances.push_back({spatCompute.getOperation(), 0, 1});
+          continue;
+        }
+        if (auto batch = dyn_cast<SpatComputeBatch>(&op)) {
+          if (isUsedAsWeightOnly(batch.getOperation()))
+            continue;
+          size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
+          for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex)
+            instances.push_back(getBatchChunkForIndex(batch, chunkIndex));
+        }
+      }
+    }
+  }
+  return instances;
+}
+
+VirtualGraph buildInitialVirtualGraph(ArrayRef<ComputeInstance> computeInstances, ArrayRef<IndexedEdge> edges) {
   VirtualGraph graph;
-  graph.nodes.reserve(spatComputes.size());
-  for (auto [index, spatCompute] : llvm::enumerate(spatComputes)) {
+  graph.nodes.reserve(computeInstances.size());
+  for (auto [index, computeInstance] : llvm::enumerate(computeInstances)) {
     VirtualNode node;
     node.originalComputeIndices.push_back(index);
-    node.weight = getSpatComputeWeight(spatCompute);
-    node.crossbarUsage = getSpatComputeCrossbarUsage(spatCompute);
+    node.weight = getComputeInstanceWeight(computeInstance);
+    node.crossbarUsage = getComputeInstanceCrossbarUsage(computeInstance);
     graph.nodes.push_back(std::move(node));
   }
   graph.edges = aggregateEdges(edges);
@@ -116,22 +263,34 @@ TimingInfo computeTiming(const VirtualGraph& graph) {
     incomingEdgeCount[endIndex]++;
   }
 
-  std::vector<size_t> readyNodes;
-  readyNodes.reserve(nodeCount);
+  auto getVirtualNodeOrderKey = [&](size_t nodeIndex) {
+    const VirtualNode& node = graph.nodes[nodeIndex];
+    if (!node.originalComputeIndices.empty())
+      return node.originalComputeIndices.front();
+    return nodeIndex;
+  };
+  auto readyNodeGreater = [&](size_t lhs, size_t rhs) {
+    size_t lhsKey = getVirtualNodeOrderKey(lhs);
+    size_t rhsKey = getVirtualNodeOrderKey(rhs);
+    if (lhsKey != rhsKey)
+      return lhsKey > rhsKey;
+    return lhs > rhs;
+  };
+  std::priority_queue<size_t, std::vector<size_t>, decltype(readyNodeGreater)> readyNodes(readyNodeGreater);
   for (size_t i = 0; i < nodeCount; ++i)
     if (incomingEdgeCount[i] == 0)
-      readyNodes.push_back(i);
+      readyNodes.push(i);
 
-  size_t readyIndex = 0;
-  while (readyIndex != readyNodes.size()) {
-    size_t current = readyNodes[readyIndex++];
+  while (!readyNodes.empty()) {
+    size_t current = readyNodes.top();
+    readyNodes.pop();
     timing.topologicalOrder.push_back(current);
     for (auto [child, weight] : children[current]) {
       (void) weight;
       assert(incomingEdgeCount[child] > 0 && "incoming edge count underflow");
       incomingEdgeCount[child]--;
       if (incomingEdgeCount[child] == 0)
-        readyNodes.push_back(child);
+        readyNodes.push(child);
     }
   }
 
@@ -287,17 +446,21 @@ std::vector<IndexedEdge> buildWindowEdges(const VirtualGraph& graph, const std::
 WindowScheduleResult scheduleWindow(const VirtualGraph& graph, ArrayRef<size_t> selectedNodes, MLIRContext* context) {
   std::vector<Weight> windowWeights;
   std::vector<CrossbarUsage> windowCrossbarUsage;
+  std::vector<int64_t> windowNodeOrderKeys;
   std::vector<int64_t> nodeToWindowIndex(graph.nodes.size(), -1);
   windowWeights.reserve(selectedNodes.size());
   windowCrossbarUsage.reserve(selectedNodes.size());
+  windowNodeOrderKeys.reserve(selectedNodes.size());
 
   for (auto [windowIndex, nodeIndex] : llvm::enumerate(selectedNodes)) {
     nodeToWindowIndex[nodeIndex] = static_cast<int64_t>(windowIndex);
     windowWeights.push_back(graph.nodes[nodeIndex].weight);
     windowCrossbarUsage.push_back(graph.nodes[nodeIndex].crossbarUsage);
+    windowNodeOrderKeys.push_back(static_cast<int64_t>(nodeIndex));
   }
 
-  GraphDCP windowGraph(windowWeights, buildWindowEdges(graph, nodeToWindowIndex), windowCrossbarUsage);
+  GraphDCP windowGraph(
+    windowWeights, buildWindowEdges(graph, nodeToWindowIndex), windowNodeOrderKeys, windowCrossbarUsage);
   if (coresCount.getValue() > 0)
     windowGraph.setMaxCpuCount(static_cast<int>(coresCount.getValue()));
   windowGraph.setContext(context);
@@ -414,13 +577,7 @@ bool coarsenGraph(const VirtualGraph& graph,
   return true;
 }
 
-constexpr CPU kDefaultMaxCpuCount = 1000;
-
-CPU getVirtualGraphMaxCpuCount() {
-  if (coresCount.getValue() > 0)
-    return static_cast<CPU>(coresCount.getValue());
-  return kDefaultMaxCpuCount;
-}
+CPU getVirtualGraphMaxCpuCount() { return static_cast<CPU>(getSchedulingCpuBudget()); }
 
 size_t getDcpCoarseningWindowSize(size_t nodeCount) {
   size_t windowSize = std::min(dcpCriticalWindowSize.getValue(), nodeCount);
@@ -430,7 +587,7 @@ size_t getDcpCoarseningWindowSize(size_t nodeCount) {
   return windowSize;
 }
 
-DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph, ArrayRef<SpatCompute> spatComputes) {
+DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph, ArrayRef<ComputeInstance> computeInstances) {
   DCPAnalysisResult result;
 
   TimingInfo timing = computeTiming(graph);
@@ -443,19 +600,22 @@ DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph, ArrayRe
     std::iota(virtualNodeOrder.begin(), virtualNodeOrder.end(), 0);
   }
 
-  std::vector<size_t> originalComputeToCpu(spatComputes.size(), 0);
+  std::vector<size_t> originalComputeToCpu(computeInstances.size(), 0);
   for (auto [cpu, virtualNodeIndex] : llvm::enumerate(virtualNodeOrder)) {
     const VirtualNode& virtualNode = graph.nodes[virtualNodeIndex];
     for (size_t originalIndex : virtualNode.originalComputeIndices)
       originalComputeToCpu[originalIndex] = cpu;
   }
 
-  result.dominanceOrderCompute.reserve(spatComputes.size());
-  for (auto [originalIndex, spatCompute] : llvm::enumerate(spatComputes)) {
+  result.dominanceOrderCompute.reserve(computeInstances.size());
+  llvm::DenseMap<size_t, size_t> nextCpuSlot;
+  for (auto [originalIndex, computeInstance] : llvm::enumerate(computeInstances)) {
     size_t cpu = originalComputeToCpu[originalIndex];
-    result.dominanceOrderCompute.push_back(spatCompute);
-    result.computeToCpuMap[spatCompute] = cpu;
-    result.cpuToLastComputeMap[cpu] = spatCompute;
+    result.dominanceOrderCompute.push_back(computeInstance);
+    result.computeToCpuMap[computeInstance] = cpu;
+    result.computeToCpuSlotMap[computeInstance] = nextCpuSlot[cpu]++;
+    result.computeToAestMap[computeInstance] = originalIndex;
+    result.cpuToLastComputeMap[cpu] = computeInstance;
   }
   for (const auto& [cpu, lastCompute] : result.cpuToLastComputeMap)
     result.isLastComputeOfCpu.insert(lastCompute);
@@ -463,13 +623,48 @@ DCPAnalysisResult buildResultFromVirtualGraph(const VirtualGraph& graph, ArrayRe
   return result;
 }
 
-DCPAnalysisResult runLegacyDcp(ArrayRef<SpatCompute> spatComputes, ArrayRef<IndexedEdge> edges, MLIRContext* context) {
-  GraphDCP graphDCP(spatComputes, edges);
+DCPAnalysisResult buildResultFromScheduledGraph(GraphDCP& graphDCP, ArrayRef<ComputeInstance> computeInstances) {
+  DCPAnalysisResult result;
+  result.dominanceOrderCompute.assign(computeInstances.begin(), computeInstances.end());
+
+  for (CPU cpu = 0; cpu < graphDCP.cpuCount(); ++cpu) {
+    auto scheduledTasks = graphDCP.getScheduledTasks(cpu);
+    if (scheduledTasks.empty())
+      continue;
+
+    for (auto [slot, task] : llvm::enumerate(scheduledTasks)) {
+      ComputeInstance instance = computeInstances[task.nodeIndex];
+      result.computeToCpuMap[instance] = cpu;
+      result.computeToCpuSlotMap[instance] = slot;
+      result.computeToAestMap[instance] = static_cast<uint64_t>(task.aest);
+    }
+    result.cpuToLastComputeMap[cpu] = computeInstances[scheduledTasks.back().nodeIndex];
+    result.isLastComputeOfCpu.insert(computeInstances[scheduledTasks.back().nodeIndex]);
+  }
+
+  return result;
+}
+
+DCPAnalysisResult
+runLegacyDcp(ArrayRef<ComputeInstance> computeInstances, ArrayRef<IndexedEdge> edges, MLIRContext* context) {
+  SmallVector<Weight> nodeWeights;
+  SmallVector<CrossbarUsage> nodeCrossbarUsage;
+  SmallVector<int64_t> nodeOrderKeys;
+  nodeWeights.reserve(computeInstances.size());
+  nodeCrossbarUsage.reserve(computeInstances.size());
+  nodeOrderKeys.reserve(computeInstances.size());
+  for (auto [index, instance] : llvm::enumerate(computeInstances)) {
+    nodeWeights.push_back(getComputeInstanceWeight(instance));
+    nodeCrossbarUsage.push_back(getComputeInstanceCrossbarUsage(instance));
+    nodeOrderKeys.push_back(static_cast<int64_t>(index));
+  }
+
+  GraphDCP graphDCP(nodeWeights, edges, nodeOrderKeys, nodeCrossbarUsage);
   if (coresCount.getValue() > 0)
     graphDCP.setMaxCpuCount(static_cast<int>(coresCount.getValue()));
   graphDCP.setContext(context);
   graphDCP.runDcp();
-  return graphDCP.getResult();
+  return buildResultFromScheduledGraph(graphDCP, computeInstances);
 }
 
 } // namespace
@@ -488,27 +683,41 @@ SpatCompute getOriginalSpatCompute(Operation* op) {
 }
 
 DCPAnalysisResult DCPAnalysis::run() {
-  SmallVector<SpatCompute, 10> spatComputes;
+  SmallVector<ComputeInstance> computeInstances = collectComputeInstances(entryOp);
   SmallVector<IndexedEdge, 10> edges;
-  for (auto& region : entryOp->getRegions())
-    for (SpatCompute spatCompute : region.getOps<SpatCompute>())
-      spatComputes.push_back(spatCompute);
 
-  for (auto [indexEndEdge, spatCompute] : llvm::enumerate(spatComputes)) {
-    for (Value input : spatCompute.getInputs()) {
-      if (auto producerCompute = getOriginalSpatCompute(input.getDefiningOp())) {
-        auto producerIt = llvm::find(spatComputes, producerCompute);
-        assert(producerIt != spatComputes.end());
-        auto indexStartEdge = std::distance(spatComputes.begin(), producerIt);
-        edges.push_back({indexStartEdge, indexEndEdge, getSizeInBytes(cast<ShapedType>(input.getType()))});
+  llvm::DenseMap<ComputeInstance, size_t> instanceToIndex;
+  instanceToIndex.reserve(computeInstances.size());
+  for (auto [index, instance] : llvm::enumerate(computeInstances))
+    instanceToIndex[instance] = index;
+
+  for (auto [indexEndEdge, computeInstance] : llvm::enumerate(computeInstances)) {
+    for (Value input : getComputeInstanceInputs(computeInstance)) {
+      if (auto producerInstance = getOriginalComputeInstance(input)) {
+        auto producerIt = instanceToIndex.find(*producerInstance);
+        assert(producerIt != instanceToIndex.end());
+        auto indexStartEdge = producerIt->second;
+        edges.push_back({static_cast<int64_t>(indexStartEdge),
+                         static_cast<int64_t>(indexEndEdge),
+                         static_cast<int64_t>(getSizeInBytes(cast<ShapedType>(input.getType())))});
       }
     }
   }
 
+  if (coresCount.getValue() > 0) {
+    size_t schedulingCpuBudget = getSchedulingCpuBudget();
+    bool needsExactScheduledBatches = llvm::any_of(computeInstances, [&](const ComputeInstance& instance) {
+      auto batch = dyn_cast<SpatComputeBatch>(instance.op);
+      return batch && static_cast<size_t>(batch.getLaneCount()) > schedulingCpuBudget;
+    });
+    if (needsExactScheduledBatches)
+      return runLegacyDcp(computeInstances, edges, entryOp->getContext());
+  }
+
   if (dcpCriticalWindowSize.getValue() == 0)
-    return runLegacyDcp(spatComputes, edges, entryOp->getContext());
+    return runLegacyDcp(computeInstances, edges, entryOp->getContext());
 
-  VirtualGraph virtualGraph = buildInitialVirtualGraph(spatComputes, edges);
+  VirtualGraph virtualGraph = buildInitialVirtualGraph(computeInstances, edges);
   size_t iteration = 0;
   auto tryCoarsenSelectedNodes = [&](ArrayRef<size_t> selectedNodes) {
     size_t oldNodeCount = virtualGraph.nodes.size();
@@ -545,6 +754,13 @@ DCPAnalysisResult DCPAnalysis::run() {
   };
 
   while (virtualGraph.nodes.size() > 1) {
+    if (virtualGraph.nodes.size() <= getSchedulingCpuBudget()) {
+      if (virtualGraph.nodes.size() >= 200)
+        llvm::errs() << llvm::formatv(
+          "[DCP-COARSEN] iter={0} old={1} stop=cpu-budget\n", iteration, virtualGraph.nodes.size());
+      break;
+    }
+
     iteration++;
     TimingInfo timing = computeTiming(virtualGraph);
     if (!timing.valid) {
@@ -576,7 +792,7 @@ DCPAnalysisResult DCPAnalysis::run() {
     break;
   }
 
-  return buildResultFromVirtualGraph(virtualGraph, spatComputes);
+  return buildResultFromVirtualGraph(virtualGraph, computeInstances);
 }
 
 } // namespace spatial

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

@@ -5,15 +5,30 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 
+#include <cstdint>
 #include <vector>
 
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
+// A scheduling identity that covers both spat.compute and scheduled shards of
+// spat.compute_batch.
+struct ComputeInstance {
+  mlir::Operation* op = nullptr;
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 1;
+
+  bool operator==(const ComputeInstance& other) const {
+    return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount;
+  }
+};
+
 struct DCPAnalysisResult {
-  std::vector<onnx_mlir::spatial::SpatCompute> dominanceOrderCompute;
-  llvm::DenseMap<onnx_mlir::spatial::SpatCompute, size_t> computeToCpuMap;
-  llvm::DenseSet<onnx_mlir::spatial::SpatCompute> isLastComputeOfCpu;
-  llvm::DenseMap<size_t, onnx_mlir::spatial::SpatCompute> cpuToLastComputeMap;
+  std::vector<ComputeInstance> dominanceOrderCompute;
+  llvm::DenseMap<ComputeInstance, size_t> computeToCpuMap;
+  llvm::DenseMap<ComputeInstance, size_t> computeToCpuSlotMap;
+  llvm::DenseMap<ComputeInstance, uint64_t> computeToAestMap;
+  llvm::DenseSet<ComputeInstance> isLastComputeOfCpu;
+  llvm::DenseMap<size_t, ComputeInstance> cpuToLastComputeMap;
 };
 
 namespace onnx_mlir {
@@ -34,3 +49,21 @@ public:
 
 } // namespace spatial
 } // namespace onnx_mlir
+
+namespace llvm {
+template <>
+struct DenseMapInfo<ComputeInstance> {
+  static ComputeInstance getEmptyKey() {
+    return {DenseMapInfo<mlir::Operation*>::getEmptyKey(), UINT32_MAX, UINT32_MAX};
+  }
+  static ComputeInstance getTombstoneKey() {
+    return {DenseMapInfo<mlir::Operation*>::getTombstoneKey(), UINT32_MAX, UINT32_MAX};
+  }
+  static unsigned getHashValue(const ComputeInstance& v) {
+    return llvm::hash_combine(v.op, v.laneStart, v.laneCount);
+  }
+  static bool isEqual(const ComputeInstance& a, const ComputeInstance& b) {
+    return a == b;
+  }
+};
+} // namespace llvm

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

@@ -1491,18 +1491,21 @@ void GraphDCP::runDcp() {
   struct ReadyEntry {
     Time slack;
     Time aest;
+    int64_t orderKey;
     TaskDCP* task;
     bool operator>(const ReadyEntry& other) const {
       if (slack != other.slack)
         return slack > other.slack;
+      if (aest != other.aest)
         return aest > other.aest;
+      return orderKey > other.orderKey;
     }
   };
   std::priority_queue<ReadyEntry, std::vector<ReadyEntry>, std::greater<ReadyEntry>> readyQueue;
   size_t readyCount = 0;
 
   auto pushReady = [&](TaskDCP* node) {
-    readyQueue.push({slackOrZero(node->getAest(), node->getAlst()), node->getAest(), node});
+    readyQueue.push({slackOrZero(node->getAest(), node->getAlst()), node->getAest(), node->Id(), node});
   };
 
   for (auto& node : nodes) {
@@ -1528,7 +1531,7 @@ void GraphDCP::runDcp() {
         candidate = entry.task;
         break;
       }
-      readyQueue.push({curSlack, curAest, entry.task});
+      readyQueue.push({curSlack, curAest, entry.orderKey, entry.task});
     }
     assert(candidate != nullptr && "readyCount > 0 but heap exhausted");
     --readyCount;
@@ -1579,8 +1582,11 @@ DCPAnalysisResult GraphDCP::getResult() {
 
   auto dominanceOrder = dcp_graph::collectDominanceOrder(getRoots(), nodes.size());
   ret.dominanceOrderCompute.reserve(dominanceOrder.size());
-  for (auto elem : dominanceOrder)
-    ret.dominanceOrderCompute.push_back(elem->getSpatCompute());
+  for (auto elem : dominanceOrder) {
+    auto spatCompute = elem->getSpatCompute();
+    if (spatCompute)
+      ret.dominanceOrderCompute.push_back({spatCompute.getOperation(), 0});
+  }
 
   for (CPU cpu = 0; cpu < getLastCpu(); ++cpu) {
     const CpuTaskList* tasks = findCpuTasks(cpu);
@@ -1588,10 +1594,14 @@ DCPAnalysisResult GraphDCP::getResult() {
       continue;
     size_t i = 0;
     for (auto node : *tasks) {
-      ret.computeToCpuMap[node->getSpatCompute()] = cpu;
+      auto spatCompute = node->getSpatCompute();
+      if (!spatCompute)
+        continue;
+      ComputeInstance instance {spatCompute.getOperation(), 0};
+      ret.computeToCpuMap[instance] = cpu;
       if (i++ == tasks->size() - 1) {
-        ret.isLastComputeOfCpu.insert(node->getSpatCompute());
-        ret.cpuToLastComputeMap[cpu] = node->getSpatCompute();
+        ret.isLastComputeOfCpu.insert(instance);
+        ret.cpuToLastComputeMap[cpu] = instance;
       }
     }
   }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Graph.hpp

@@ -138,13 +138,18 @@ public:
 
   GraphDCP(llvm::ArrayRef<Weight> nodeWeights,
            llvm::ArrayRef<IndexedEdge> edges,
+           llvm::ArrayRef<int64_t> nodeOrderKeys = {},
            llvm::ArrayRef<CrossbarUsage> nodeCrossbarUsage = {})
   : nodes(), cpuTasks(), cpuCrossbarUsage() {
     assert((nodeCrossbarUsage.empty() || nodeCrossbarUsage.size() == nodeWeights.size())
            && "synthetic crossbar usage must match synthetic node weights");
+    assert((nodeOrderKeys.empty() || nodeOrderKeys.size() == nodeWeights.size())
+           && "synthetic node order keys must match synthetic node weights");
     nodes.reserve(nodeWeights.size());
     for (auto [index, weight] : llvm::enumerate(nodeWeights))
-      nodes.emplace_back(index, weight, nodeCrossbarUsage.empty() ? 0 : nodeCrossbarUsage[index]);
+      nodes.emplace_back(nodeOrderKeys.empty() ? static_cast<int64_t>(index) : nodeOrderKeys[index],
+                         weight,
+                         nodeCrossbarUsage.empty() ? 0 : nodeCrossbarUsage[index]);
     for (auto [start, end, weight] : edges)
       makeEdge(start, end, weight);
   }

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DCPGraph/Utils.hpp

@@ -105,7 +105,7 @@ inline CrossbarUsage getSpatComputeCrossbarUsage(onnx_mlir::spatial::SpatCompute
   CrossbarUsage crossbarUsage = 0;
   for (auto& region : spatCompute.getBody())
     for (auto& inst : region)
-      if (llvm::isa<onnx_mlir::spatial::SpatWeightedVMMOp>(inst))
+      if (llvm::isa<onnx_mlir::spatial::SpatVMMOp>(inst))
         crossbarUsage = checkedAdd(crossbarUsage, static_cast<CrossbarUsage>(1));
   return crossbarUsage;
 }

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

@@ -0,0 +1,577 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <tuple>
+
+#include "RegularOpCompaction.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+enum class RegularStepKind {
+  Wvmm,
+  VAddLhs,
+  VAddRhs,
+};
+
+struct RegularStep {
+  RegularStepKind kind;
+  int32_t weightIndex = 0;
+  Value invariantOperand;
+  Type resultType;
+};
+
+struct RegularChunk {
+  Operation* startOp = nullptr;
+  SmallVector<Operation*> ops;
+  SmallVector<RegularStep> steps;
+  Value input;
+  Value output;
+};
+
+static bool areEquivalentRegularSteps(const RegularStep& lhs, const RegularStep& rhs) {
+  return lhs.kind == rhs.kind && lhs.weightIndex == rhs.weightIndex && lhs.invariantOperand == rhs.invariantOperand
+      && lhs.resultType == rhs.resultType;
+}
+
+static bool areEquivalentRegularChunks(const RegularChunk& lhs, const RegularChunk& rhs) {
+  if (lhs.input.getType() != rhs.input.getType() || lhs.output.getType() != rhs.output.getType()
+      || lhs.steps.size() != rhs.steps.size()) {
+    return false;
+  }
+
+  return llvm::all_of(llvm::zip_equal(lhs.steps, rhs.steps),
+                      [](auto pair) { return areEquivalentRegularSteps(std::get<0>(pair), std::get<1>(pair)); });
+}
+
+static FailureOr<RegularChunk> analyzeRegularChunk(spatial::SpatVMMOp startOp) {
+  RegularChunk chunk;
+  chunk.startOp = startOp.getOperation();
+  chunk.input = startOp.getInput();
+  chunk.output = startOp.getOutput();
+  chunk.ops.push_back(startOp.getOperation());
+  chunk.steps.push_back(
+    {RegularStepKind::Wvmm, static_cast<int32_t>(startOp.getWeightIndex()), Value(), startOp.getOutput().getType()});
+
+  Value currentValue = startOp.getOutput();
+  while (currentValue.hasOneUse()) {
+    Operation* user = *currentValue.getUsers().begin();
+    if (user->getBlock() != startOp->getBlock())
+      break;
+
+    auto vaddOp = dyn_cast<spatial::SpatVAddOp>(user);
+    if (!vaddOp)
+      break;
+
+    if (vaddOp.getLhs() == currentValue)
+      chunk.steps.push_back({RegularStepKind::VAddLhs, 0, vaddOp.getRhs(), vaddOp.getOutput().getType()});
+    else if (vaddOp.getRhs() == currentValue)
+      chunk.steps.push_back({RegularStepKind::VAddRhs, 0, vaddOp.getLhs(), vaddOp.getOutput().getType()});
+    else
+      break;
+
+    chunk.ops.push_back(vaddOp);
+    chunk.output = vaddOp.getOutput();
+    currentValue = vaddOp.getOutput();
+  }
+
+  return chunk;
+}
+
+static void buildRegularMapBody(spatial::SpatMapOp mapOp, const RegularChunk& anchorChunk, IRRewriter& rewriter) {
+  auto* block = rewriter.createBlock(
+    &mapOp.getBody(), mapOp.getBody().end(), TypeRange {anchorChunk.input.getType()}, {anchorChunk.startOp->getLoc()});
+  rewriter.setInsertionPointToEnd(block);
+
+  IRMapping mapping;
+  mapping.map(anchorChunk.input, block->getArgument(0));
+
+  for (Operation* op : anchorChunk.ops) {
+    Operation* cloned = rewriter.clone(*op, mapping);
+    for (auto [oldResult, newResult] : llvm::zip(op->getResults(), cloned->getResults()))
+      mapping.map(oldResult, newResult);
+  }
+
+  spatial::SpatYieldOp::create(
+    rewriter, anchorChunk.startOp->getLoc(), ValueRange {mapping.lookup(anchorChunk.output)});
+}
+
+static void compactRegularChunkRun(IRRewriter& rewriter, ArrayRef<RegularChunk> run) {
+  assert(!run.empty() && "expected a non-empty regular chunk run");
+  const RegularChunk& anchorChunk = run.front();
+
+  SmallVector<Value> inputs;
+  SmallVector<Type> outputTypes;
+  inputs.reserve(run.size());
+  outputTypes.reserve(run.size());
+  for (const RegularChunk& chunk : run) {
+    inputs.push_back(chunk.input);
+    outputTypes.push_back(chunk.output.getType());
+  }
+
+  rewriter.setInsertionPoint(anchorChunk.startOp);
+  auto mapOp =
+    spatial::SpatMapOp::create(rewriter, anchorChunk.startOp->getLoc(), TypeRange(outputTypes), ValueRange(inputs));
+  buildRegularMapBody(mapOp, anchorChunk, rewriter);
+
+  for (auto [index, chunk] : llvm::enumerate(run)) {
+    Value output = chunk.output;
+    output.replaceAllUsesWith(mapOp.getResult(index));
+  }
+
+  SmallVector<Operation*> opsToErase;
+  for (const RegularChunk& chunk : run)
+    llvm::append_range(opsToErase, chunk.ops);
+  for (Operation* op : llvm::reverse(opsToErase))
+    rewriter.eraseOp(op);
+}
+
+} // namespace
+
+void compactScalarChannelRuns(func::FuncOp funcOp, int64_t& nextChannelId) {
+  IRRewriter rewriter(funcOp.getContext());
+
+  for (auto compute : funcOp.getOps<spatial::SpatCompute>()) {
+    Block& block = compute.getBody().front();
+    for (auto it = block.begin(); it != block.end();) {
+      auto receiveOp = dyn_cast<spatial::SpatChannelReceiveOp>(&*it);
+      if (receiveOp) {
+        SmallVector<spatial::SpatChannelReceiveOp> run;
+        Type outputType = receiveOp.getOutput().getType();
+        auto runIt = it;
+        while (runIt != block.end()) {
+          auto current = dyn_cast<spatial::SpatChannelReceiveOp>(&*runIt);
+          if (!current || current.getOutput().getType() != outputType)
+            break;
+          run.push_back(current);
+          ++runIt;
+        }
+
+        if (run.size() > 1) {
+          struct ReceiveEntry {
+            spatial::SpatChannelReceiveOp op;
+            size_t originalIndex = 0;
+            uint32_t sourceCoreId = 0;
+            uint32_t targetCoreId = 0;
+            uint64_t channelId = 0;
+          };
+          SmallVector<ReceiveEntry> sortedEntries;
+          sortedEntries.reserve(run.size());
+          for (auto [originalIndex, op] : llvm::enumerate(run))
+            sortedEntries.push_back({op, originalIndex, op.getSourceCoreId(), op.getTargetCoreId(), op.getChannelId()});
+          llvm::stable_sort(sortedEntries, [](const ReceiveEntry& lhs, const ReceiveEntry& rhs) {
+            return std::tuple(lhs.sourceCoreId, lhs.targetCoreId, lhs.channelId)
+                 < std::tuple(rhs.sourceCoreId, rhs.targetCoreId, rhs.channelId);
+          });
+
+          SmallVector<int64_t> channelIds;
+          SmallVector<int32_t> sourceCoreIds;
+          SmallVector<int32_t> targetCoreIds;
+          SmallVector<Type> outputTypes;
+          channelIds.reserve(sortedEntries.size());
+          sourceCoreIds.reserve(sortedEntries.size());
+          targetCoreIds.reserve(sortedEntries.size());
+          outputTypes.reserve(sortedEntries.size());
+          for (ReceiveEntry& entry : sortedEntries) {
+            (void) entry;
+            channelIds.push_back(nextChannelId++);
+            sourceCoreIds.push_back(static_cast<int32_t>(entry.sourceCoreId));
+            targetCoreIds.push_back(static_cast<int32_t>(entry.targetCoreId));
+            outputTypes.push_back(entry.op.getOutput().getType());
+          }
+
+          rewriter.setInsertionPoint(run.front());
+          auto compactReceive = spatial::SpatChannelReceiveManyOp::create(rewriter,
+                                                                          run.front().getLoc(),
+                                                                          TypeRange(outputTypes),
+                                                                          rewriter.getDenseI64ArrayAttr(channelIds),
+                                                                          rewriter.getDenseI32ArrayAttr(sourceCoreIds),
+                                                                          rewriter.getDenseI32ArrayAttr(targetCoreIds));
+          for (auto [sortedIndex, entry] : llvm::enumerate(sortedEntries))
+            entry.op.getOutput().replaceAllUsesWith(compactReceive.getResult(sortedIndex));
+          for (auto op : run)
+            rewriter.eraseOp(op);
+
+          it = compactReceive->getIterator();
+          ++it;
+          continue;
+        }
+      }
+
+      auto sendOp = dyn_cast<spatial::SpatChannelSendOp>(&*it);
+      if (sendOp) {
+        SmallVector<spatial::SpatChannelSendOp> run;
+        Type inputType = sendOp.getInput().getType();
+        auto runIt = it;
+        while (runIt != block.end()) {
+          auto current = dyn_cast<spatial::SpatChannelSendOp>(&*runIt);
+          if (!current || current.getInput().getType() != inputType)
+            break;
+          run.push_back(current);
+          ++runIt;
+        }
+
+        if (run.size() > 1) {
+          struct SendEntry {
+            spatial::SpatChannelSendOp op;
+            uint32_t sourceCoreId = 0;
+            uint32_t targetCoreId = 0;
+            uint64_t channelId = 0;
+          };
+          SmallVector<SendEntry> sortedEntries;
+          sortedEntries.reserve(run.size());
+          for (auto op : run)
+            sortedEntries.push_back({op, op.getSourceCoreId(), op.getTargetCoreId(), op.getChannelId()});
+          llvm::stable_sort(sortedEntries, [](const SendEntry& lhs, const SendEntry& rhs) {
+            return std::tuple(lhs.sourceCoreId, lhs.targetCoreId, lhs.channelId)
+                 < std::tuple(rhs.sourceCoreId, rhs.targetCoreId, rhs.channelId);
+          });
+
+          SmallVector<int64_t> channelIds;
+          SmallVector<int32_t> sourceCoreIds;
+          SmallVector<int32_t> targetCoreIds;
+          SmallVector<Value> inputs;
+          channelIds.reserve(sortedEntries.size());
+          sourceCoreIds.reserve(sortedEntries.size());
+          targetCoreIds.reserve(sortedEntries.size());
+          inputs.reserve(sortedEntries.size());
+          for (SendEntry& entry : sortedEntries) {
+            (void) entry;
+            channelIds.push_back(nextChannelId++);
+            sourceCoreIds.push_back(static_cast<int32_t>(entry.sourceCoreId));
+            targetCoreIds.push_back(static_cast<int32_t>(entry.targetCoreId));
+            inputs.push_back(entry.op.getInput());
+          }
+
+          rewriter.setInsertionPoint(run.front());
+          spatial::SpatChannelSendManyOp::create(rewriter,
+                                                 run.front().getLoc(),
+                                                 rewriter.getDenseI64ArrayAttr(channelIds),
+                                                 rewriter.getDenseI32ArrayAttr(sourceCoreIds),
+                                                 rewriter.getDenseI32ArrayAttr(targetCoreIds),
+                                                 ValueRange(inputs));
+          for (auto op : run)
+            rewriter.eraseOp(op);
+
+          it = runIt;
+          continue;
+        }
+      }
+
+      ++it;
+    }
+  }
+}
+
+void compactBatchChannelRuns(func::FuncOp funcOp) {
+  IRRewriter rewriter(funcOp.getContext());
+
+  for (auto batch : funcOp.getOps<spatial::SpatComputeBatch>()) {
+    Block& block = batch.getBody().front();
+    for (auto it = block.begin(); it != block.end();) {
+      auto receiveOp = dyn_cast<spatial::SpatChannelReceiveBatchOp>(&*it);
+      if (receiveOp) {
+        SmallVector<spatial::SpatChannelReceiveBatchOp> run;
+        Type outputType = receiveOp.getOutput().getType();
+        auto runIt = it;
+        while (runIt != block.end()) {
+          auto current = dyn_cast<spatial::SpatChannelReceiveBatchOp>(&*runIt);
+          if (!current || current.getOutput().getType() != outputType)
+            break;
+          run.push_back(current);
+          ++runIt;
+        }
+
+        if (run.size() > 1) {
+          SmallVector<int64_t> channelIds;
+          SmallVector<int32_t> sourceCoreIds;
+          SmallVector<int32_t> targetCoreIds;
+          SmallVector<Type> outputTypes;
+          outputTypes.reserve(run.size());
+          for (auto op : run) {
+            llvm::append_range(channelIds, op.getChannelIds());
+            llvm::append_range(sourceCoreIds, op.getSourceCoreIds());
+            llvm::append_range(targetCoreIds, op.getTargetCoreIds());
+            outputTypes.push_back(op.getOutput().getType());
+          }
+
+          rewriter.setInsertionPoint(run.front());
+          auto compactReceive =
+            spatial::SpatChannelReceiveManyBatchOp::create(rewriter,
+                                                           run.front().getLoc(),
+                                                           TypeRange(outputTypes),
+                                                           rewriter.getDenseI64ArrayAttr(channelIds),
+                                                           rewriter.getDenseI32ArrayAttr(sourceCoreIds),
+                                                           rewriter.getDenseI32ArrayAttr(targetCoreIds));
+          for (auto [index, op] : llvm::enumerate(run))
+            op.getOutput().replaceAllUsesWith(compactReceive.getResult(index));
+          for (auto op : run)
+            rewriter.eraseOp(op);
+
+          it = compactReceive->getIterator();
+          ++it;
+          continue;
+        }
+      }
+
+      auto sendOp = dyn_cast<spatial::SpatChannelSendBatchOp>(&*it);
+      if (sendOp) {
+        SmallVector<spatial::SpatChannelSendBatchOp> run;
+        Type inputType = sendOp.getInput().getType();
+        auto runIt = it;
+        while (runIt != block.end()) {
+          auto current = dyn_cast<spatial::SpatChannelSendBatchOp>(&*runIt);
+          if (!current || current.getInput().getType() != inputType)
+            break;
+          run.push_back(current);
+          ++runIt;
+        }
+
+        if (run.size() > 1) {
+          SmallVector<int64_t> channelIds;
+          SmallVector<int32_t> sourceCoreIds;
+          SmallVector<int32_t> targetCoreIds;
+          SmallVector<Value> inputs;
+          inputs.reserve(run.size());
+          for (auto op : run) {
+            llvm::append_range(channelIds, op.getChannelIds());
+            llvm::append_range(sourceCoreIds, op.getSourceCoreIds());
+            llvm::append_range(targetCoreIds, op.getTargetCoreIds());
+            inputs.push_back(op.getInput());
+          }
+
+          rewriter.setInsertionPoint(run.front());
+          spatial::SpatChannelSendManyBatchOp::create(rewriter,
+                                                      run.front().getLoc(),
+                                                      rewriter.getDenseI64ArrayAttr(channelIds),
+                                                      rewriter.getDenseI32ArrayAttr(sourceCoreIds),
+                                                      rewriter.getDenseI32ArrayAttr(targetCoreIds),
+                                                      ValueRange(inputs));
+          for (auto op : run)
+            rewriter.eraseOp(op);
+
+          it = runIt;
+          continue;
+        }
+      }
+
+      ++it;
+    }
+  }
+}
+
+void compactRegularOpRuns(func::FuncOp funcOp) {
+  IRRewriter rewriter(funcOp.getContext());
+
+  auto compactInBlock = [&](Block& block) {
+    for (auto it = block.begin(); it != block.end();) {
+      auto startOp = dyn_cast<spatial::SpatVMMOp>(&*it);
+      if (!startOp) {
+        ++it;
+        continue;
+      }
+
+      auto anchorChunk = analyzeRegularChunk(startOp);
+      if (failed(anchorChunk)) {
+        ++it;
+        continue;
+      }
+
+      SmallVector<RegularChunk> run {*anchorChunk};
+      auto runIt = std::next(it, static_cast<std::ptrdiff_t>(anchorChunk->ops.size()));
+      while (runIt != block.end()) {
+        auto candidateStart = dyn_cast<spatial::SpatVMMOp>(&*runIt);
+        if (!candidateStart)
+          break;
+
+        auto candidateChunk = analyzeRegularChunk(candidateStart);
+        if (failed(candidateChunk) || !areEquivalentRegularChunks(*anchorChunk, *candidateChunk))
+          break;
+
+        run.push_back(*candidateChunk);
+        runIt = std::next(runIt, static_cast<std::ptrdiff_t>(candidateChunk->ops.size()));
+      }
+
+      if (run.size() <= 1) {
+        ++it;
+        continue;
+      }
+
+      compactRegularChunkRun(rewriter, run);
+      it = runIt;
+    }
+  };
+
+  for (auto compute : funcOp.getOps<spatial::SpatCompute>())
+    compactInBlock(compute.getBody().front());
+  for (auto batch : funcOp.getOps<spatial::SpatComputeBatch>())
+    compactInBlock(batch.getBody().front());
+}
+
+void compactRowWiseWvmmRuns(func::FuncOp funcOp) {
+  IRRewriter rewriter(funcOp.getContext());
+
+  for (auto compute : funcOp.getOps<spatial::SpatCompute>()) {
+    Block& block = compute.getBody().front();
+    for (auto it = block.begin(); it != block.end();) {
+      auto wvmmOp = dyn_cast<spatial::SpatVMMOp>(&*it);
+      if (!wvmmOp) {
+        ++it;
+        continue;
+      }
+
+      auto extractRowsOp = wvmmOp.getInput().getDefiningOp<spatial::SpatExtractRowsOp>();
+      auto rowResult = dyn_cast<OpResult>(wvmmOp.getInput());
+      auto outputType = dyn_cast<RankedTensorType>(wvmmOp.getOutput().getType());
+      if (!extractRowsOp || !rowResult || rowResult.getOwner() != extractRowsOp || !outputType
+          || !outputType.hasStaticShape() || outputType.getRank() != 2 || outputType.getShape()[0] != 1) {
+        ++it;
+        continue;
+      }
+
+      SmallVector<spatial::SpatVMMOp> run;
+      auto runIt = it;
+      int64_t expectedRow = static_cast<int64_t>(rowResult.getResultNumber());
+      while (runIt != block.end()) {
+        auto current = dyn_cast<spatial::SpatVMMOp>(&*runIt);
+        if (!current || current.getWeightIndex() != wvmmOp.getWeightIndex()
+            || current.getInput().getDefiningOp<spatial::SpatExtractRowsOp>() != extractRowsOp
+            || current.getInput().getType() != wvmmOp.getInput().getType()
+            || current.getOutput().getType() != wvmmOp.getOutput().getType()) {
+          break;
+        }
+
+        auto currentRow = dyn_cast<OpResult>(current.getInput());
+        if (!currentRow || currentRow.getResultNumber() != static_cast<unsigned>(expectedRow))
+          break;
+
+        run.push_back(current);
+        ++expectedRow;
+        ++runIt;
+      }
+
+      if (run.size() <= 1) {
+        ++it;
+        continue;
+      }
+
+      if (!run.front().getOutput().hasOneUse()) {
+        ++it;
+        continue;
+      }
+      auto concatUse = run.front().getOutput().getUses().begin();
+      auto concatOp = dyn_cast<spatial::SpatConcatOp>(concatUse->getOwner());
+      if (!concatOp) {
+        ++it;
+        continue;
+      }
+
+      unsigned concatStartIndex = concatUse->getOperandNumber();
+      bool validConcatRun = true;
+      for (auto [index, op] : llvm::enumerate(run)) {
+        if (!op.getOutput().hasOneUse()) {
+          validConcatRun = false;
+          break;
+        }
+        OpOperand& use = *op.getOutput().getUses().begin();
+        if (use.getOwner() != concatOp || use.getOperandNumber() != concatStartIndex + index) {
+          validConcatRun = false;
+          break;
+        }
+      }
+      if (!validConcatRun) {
+        ++it;
+        continue;
+      }
+
+      auto inputType = dyn_cast<RankedTensorType>(wvmmOp.getInput().getType());
+      auto sourceType = dyn_cast<RankedTensorType>(extractRowsOp.getInput().getType());
+      if (!inputType || !sourceType || !inputType.hasStaticShape() || !sourceType.hasStaticShape()) {
+        ++it;
+        continue;
+      }
+
+      int64_t inputCols = inputType.getShape()[1];
+      int64_t outputCols = outputType.getShape()[1];
+      if (ShapedType::isDynamic(inputCols) || ShapedType::isDynamic(outputCols)) {
+        ++it;
+        continue;
+      }
+
+      int64_t firstRow = static_cast<int64_t>(rowResult.getResultNumber());
+      int64_t runLength = static_cast<int64_t>(run.size());
+      auto packedType = RankedTensorType::get({runLength, outputCols}, outputType.getElementType());
+
+      rewriter.setInsertionPoint(run.front());
+      auto zero = arith::ConstantIndexOp::create(rewriter, run.front().getLoc(), 0);
+      auto upper = arith::ConstantIndexOp::create(rewriter, run.front().getLoc(), runLength);
+      auto step = arith::ConstantIndexOp::create(rewriter, run.front().getLoc(), 1);
+      auto packedInit =
+        tensor::EmptyOp::create(rewriter, run.front().getLoc(), packedType.getShape(), packedType.getElementType());
+      auto loop =
+        scf::ForOp::create(rewriter, run.front().getLoc(), zero, upper, step, ValueRange {packedInit.getResult()});
+
+      {
+        OpBuilder::InsertionGuard guard(rewriter);
+        Block* loopBlock = loop.getBody();
+        rewriter.setInsertionPointToStart(loopBlock);
+        Value iv = loopBlock->getArgument(0);
+        Value acc = loopBlock->getArgument(1);
+
+        Value sourceRow = iv;
+        if (firstRow != 0) {
+          auto firstRowValue = arith::ConstantIndexOp::create(rewriter, run.front().getLoc(), firstRow);
+          sourceRow = arith::AddIOp::create(rewriter, run.front().getLoc(), iv, firstRowValue);
+        }
+
+        SmallVector<OpFoldResult> extractOffsets = {sourceRow, rewriter.getIndexAttr(0)};
+        SmallVector<OpFoldResult> extractSizes = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(inputCols)};
+        SmallVector<OpFoldResult> extractStrides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+        auto extractedRow = tensor::ExtractSliceOp::create(rewriter,
+                                                           run.front().getLoc(),
+                                                           inputType,
+                                                           extractRowsOp.getInput(),
+                                                           extractOffsets,
+                                                           extractSizes,
+                                                           extractStrides);
+        auto loopWvmm = spatial::SpatVMMOp::create(
+          rewriter, run.front().getLoc(), outputType, wvmmOp.getWeightIndex(), extractedRow.getResult());
+
+        SmallVector<OpFoldResult> insertOffsets = {iv, rewriter.getIndexAttr(0)};
+        SmallVector<OpFoldResult> insertSizes = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(outputCols)};
+        SmallVector<OpFoldResult> insertStrides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+        auto inserted = tensor::InsertSliceOp::create(
+          rewriter, run.front().getLoc(), loopWvmm.getResult(), acc, insertOffsets, insertSizes, insertStrides);
+        scf::YieldOp::create(rewriter, run.front().getLoc(), inserted.getResult());
+      }
+
+      SmallVector<Value> newConcatInputs;
+      newConcatInputs.reserve(concatOp.getInputs().size() - run.size() + 1);
+      for (auto [operandIndex, operand] : llvm::enumerate(concatOp.getInputs())) {
+        if (operandIndex == concatStartIndex)
+          newConcatInputs.push_back(loop.getResult(0));
+        if (operandIndex < concatStartIndex || operandIndex >= concatStartIndex + run.size())
+          newConcatInputs.push_back(operand);
+      }
+      rewriter.modifyOpInPlace(concatOp, [&] { concatOp->setOperands(newConcatInputs); });
+      for (auto op : run)
+        rewriter.eraseOp(op);
+
+      it = loop->getIterator();
+      ++it;
+    }
+  }
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,14 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+#include <cstdint>
+
+namespace onnx_mlir {
+
+void compactScalarChannelRuns(mlir::func::FuncOp funcOp, int64_t& nextChannelId);
+void compactBatchChannelRuns(mlir::func::FuncOp funcOp);
+void compactRegularOpRuns(mlir::func::FuncOp funcOp);
+void compactRowWiseWvmmRuns(mlir::func::FuncOp funcOp);
+
+} // namespace onnx_mlir

src/PIM/Pass/PimCodegen/HostConstantFolding/Patterns/Constant.cpp

@@ -116,10 +116,9 @@ struct FoldConstantCoreMapPattern final : OpRewritePattern<linalg::MapOp> {
     auto globalOp = createFoldedGlobal(moduleOp, mapOp.getLoc(), initType, splatAttr, "pim_core_fill");
 
     OpBuilder::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPoint(coreOp);
-    auto getGlobalOp = memref::GetGlobalOp::create(rewriter, mapOp.getLoc(), initType, globalOp.getName());
 
     rewriter.setInsertionPoint(mapOp);
+    auto getGlobalOp = memref::GetGlobalOp::create(rewriter, mapOp.getLoc(), initType, globalOp.getName());
     auto sizeInBytes = initType.getNumElements() * initType.getElementTypeBitWidth() / 8;
     pim::PimMemCopyOp::create(rewriter,
                               mapOp.getLoc(),
@@ -258,9 +257,18 @@ struct FoldConstantTransposePattern final : OpRewritePattern<pim::PimTransposeOp
     if (!resultType || !resultType.hasStaticShape())
       return failure();
 
+    // Look through an optional pim.memcp_hd to find the source get_global.
+    // This occurs when the constant was staged into device memory before transposing.
+    pim::PimMemCopyHostToDevOp memcpHd;
     auto sourceGetGlobal = transposeOp.getInput().getDefiningOp<memref::GetGlobalOp>();
+    if (!sourceGetGlobal) {
+      memcpHd = transposeOp.getInput().getDefiningOp<pim::PimMemCopyHostToDevOp>();
+      if (!memcpHd)
+        return failure();
+      sourceGetGlobal = memcpHd.getHostSource().getDefiningOp<memref::GetGlobalOp>();
       if (!sourceGetGlobal)
         return failure();
+    }
 
     auto moduleOp = transposeOp->getParentOfType<ModuleOp>();
     if (!moduleOp)
@@ -298,13 +306,26 @@ struct FoldConstantTransposePattern final : OpRewritePattern<pim::PimTransposeOp
 
     bool isAlwaysWeight =
       !transposeOp->getUsers().empty()
-      && llvm::all_of(transposeOp->getUsers(), [](Operation* user) { return isa<pim::PimCoreOp>(user); });
+      && llvm::all_of(transposeOp->getUsers(), [](Operation* user) {
+           return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user);
+         });
     if (isAlwaysWeight) {
       markWeightAlways(newGlobal);
       markWeightAlways(newGetGlobal);
     }
 
+    auto outputAllocOp = transposeOp.getOutputBuffer().getDefiningOp<memref::AllocOp>();
     rewriter.replaceOp(transposeOp, newGetGlobal.getResult());
+
+    if (memcpHd && memcpHd.use_empty()) {
+      auto deviceAllocOp = memcpHd.getDeviceTarget().getDefiningOp<memref::AllocOp>();
+      rewriter.eraseOp(memcpHd);
+      if (deviceAllocOp && deviceAllocOp->use_empty())
+        rewriter.eraseOp(deviceAllocOp);
+    }
+    if (outputAllocOp && outputAllocOp->use_empty())
+      rewriter.eraseOp(outputAllocOp);
+
     return success();
   }
 };
@@ -341,18 +362,25 @@ struct FoldConstantAllocPattern final : OpRewritePattern<memref::AllocOp> {
         continue;
       }
 
-      if (!isa<pim::PimCoreOp>(user))
+      if (!isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user))
         return failure();
     }
 
     if (!llvm::all_of(castsToReplace, [](memref::CastOp castOp) {
-          return llvm::all_of(castOp->getUsers(), [](Operation* user) { return isa<pim::PimCoreOp>(user); });
+          return llvm::all_of(castOp->getUsers(), [](Operation* user) {
+            return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user);
+          });
         })) {
       allLiveUsersAreCoreOps = false;
     }
 
     if (!llvm::all_of(allocOp->getUsers(), [](Operation* user) {
-          return isa<linalg::MapOp, memref::SubViewOp, memref::DeallocOp, memref::CastOp, pim::PimCoreOp>(user);
+          return isa<linalg::MapOp,
+                     memref::SubViewOp,
+                     memref::DeallocOp,
+                     memref::CastOp,
+                     pim::PimCoreOp,
+                     pim::PimCoreBatchOp>(user);
         })) {
       return failure();
     }
@@ -389,6 +417,83 @@ struct FoldConstantAllocPattern final : OpRewritePattern<memref::AllocOp> {
   }
 };
 
+struct FoldConstantHostCopyPattern final : OpRewritePattern<memref::CopyOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::CopyOp copyOp, PatternRewriter& rewriter) const override {
+    if (copyOp->getParentOfType<pim::PimCoreOp>())
+      return failure();
+
+    auto allocOp = copyOp.getTarget().getDefiningOp<memref::AllocOp>();
+    if (!allocOp)
+      return failure();
+    auto allocType = dyn_cast<MemRefType>(allocOp.getType());
+    if (!allocType || !allocType.hasStaticShape())
+      return failure();
+
+    auto srcSubview = getStaticSubviewInfo(copyOp.getSource());
+    Value globalSource = succeeded(srcSubview) ? srcSubview->source : stripMemRefCasts(copyOp.getSource());
+
+    auto moduleOp = copyOp->getParentOfType<ModuleOp>();
+    if (!moduleOp)
+      return failure();
+
+    auto denseAttr = getDenseGlobalValue(moduleOp, globalSource);
+    if (failed(denseAttr))
+      return failure();
+
+    DenseElementsAttr foldedAttr;
+    if (succeeded(srcSubview)) {
+      if (llvm::any_of(srcSubview->strides, [](int64_t stride) { return stride != 1; }))
+        return failure();
+      auto staticOffsets = getStaticSubviewOffsets(*srcSubview);
+      if (failed(staticOffsets))
+        return failure();
+
+      auto maybeFoldedAttr = foldDenseSubview(*denseAttr, *staticOffsets, allocType.getShape());
+      if (failed(maybeFoldedAttr))
+        return failure();
+      foldedAttr = *maybeFoldedAttr;
+    }
+    else {
+      auto resultTensorType = RankedTensorType::get(allocType.getShape(), allocType.getElementType());
+      if (resultTensorType != denseAttr->getType())
+        return failure();
+      foldedAttr = *denseAttr;
+    }
+
+    bool allLiveUsersAreCores = true;
+    for (Operation* user : allocOp->getUsers()) {
+      if (user == copyOp)
+        continue;
+      if (isa<memref::DeallocOp>(user))
+        continue;
+      if (isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user))
+        continue;
+      if (isa<memref::SubViewOp>(user)) {
+        allLiveUsersAreCores = false;
+        continue;
+      }
+      return failure();
+    }
+
+    auto newGlobal = createFoldedGlobal(moduleOp, allocOp.getLoc(), allocType, foldedAttr, "pim_folded_host_copy");
+    if (allLiveUsersAreCores)
+      markWeightAlways(newGlobal);
+
+    rewriter.setInsertionPoint(allocOp);
+    auto newGetGlobal = memref::GetGlobalOp::create(rewriter, allocOp.getLoc(), allocType, newGlobal.getName());
+    if (allLiveUsersAreCores)
+      markWeightAlways(newGetGlobal);
+
+    rewriter.replaceAllUsesWith(allocOp.getResult(), newGetGlobal.getResult());
+    rewriter.eraseOp(copyOp);
+    if (allocOp.use_empty())
+      rewriter.eraseOp(allocOp);
+    return success();
+  }
+};
+
 struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -443,7 +548,7 @@ struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
         continue;
       if (isa<memref::DeallocOp>(user))
         continue;
-      if (isa<pim::PimCoreOp>(user))
+      if (isa<pim::PimCoreOp, pim::PimCoreBatchOp>(user))
         continue;
       if (isa<memref::SubViewOp>(user)) {
         allLiveUsersAreCores = false;
@@ -473,7 +578,11 @@ struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
 
 void populateConstantFoldingConstantPatterns(RewritePatternSet& patterns) {
   patterns
-    .add<FoldConstantTransposePattern, FoldConstantAllocPattern, FoldConstantCoreMapPattern, FoldConstantMemCpPattern>(
+    .add<FoldConstantTransposePattern,
+         FoldConstantAllocPattern,
+         FoldConstantCoreMapPattern,
+         FoldConstantHostCopyPattern,
+         FoldConstantMemCpPattern>(
       patterns.getContext());
 }
 

src/PIM/Pass/PimCodegen/MaterializeHostConstantsPass.cpp

@@ -1,9 +1,12 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/MathExtras.h"
 
@@ -31,6 +34,35 @@ static int64_t getValueSizeInBytes(Value value) {
   return type.getNumElements() * type.getElementTypeBitWidth() / 8;
 }
 
+static void expandPimMapOps(func::FuncOp funcOp, IRRewriter& rewriter) {
+  SmallVector<pim::PimMapOp> mapOps;
+  funcOp.walk([&](pim::PimMapOp mapOp) { mapOps.push_back(mapOp); });
+
+  for (auto mapOp : mapOps) {
+    Block& body = mapOp.getBody().front();
+    auto yieldOp = cast<pim::PimYieldOp>(body.getTerminator());
+
+    SmallVector<Value> replacements;
+    replacements.reserve(mapOp.getInputs().size());
+    rewriter.setInsertionPoint(mapOp);
+    for (Value input : mapOp.getInputs()) {
+      IRMapping mapping;
+      mapping.map(body.getArgument(0), input);
+
+      for (Operation& op : body.without_terminator()) {
+        Operation* cloned = rewriter.clone(op, mapping);
+        for (auto [originalResult, clonedResult] : llvm::zip(op.getResults(), cloned->getResults()))
+          mapping.map(originalResult, clonedResult);
+        rewriter.setInsertionPointAfter(cloned);
+      }
+
+      replacements.push_back(mapping.lookupOrDefault(yieldOp.getOperand(0)));
+    }
+
+    rewriter.replaceOp(mapOp, replacements);
+  }
+}
+
 struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MaterializeHostConstantsPass)
 
@@ -41,13 +73,15 @@ struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass,
 
   void runOnOperation() override {
     ModuleOp moduleOp = getOperation();
-    OpBuilder rewriter(moduleOp.getContext());
+    IRRewriter rewriter(moduleOp.getContext());
     bool hasFailure = false;
 
     for (func::FuncOp funcOp : moduleOp.getOps<func::FuncOp>()) {
       if (funcOp.isExternal())
         continue;
 
+      expandPimMapOps(funcOp, rewriter);
+
       for (pim::PimCoreOp coreOp : funcOp.getOps<pim::PimCoreOp>()) {
         DenseMap<Value, DenseMap<int64_t, DenseMap<Type, Value>>> materializedValues;
 
@@ -113,6 +147,45 @@ struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass,
           }
         }
       }
+
+      SmallVector<Operation*> hostCompactOps;
+      for (Operation& op : funcOp.getBody().front())
+        if (isa<pim::PimExtractRowsOp, pim::PimConcatOp>(op))
+          hostCompactOps.push_back(&op);
+
+      for (Operation* op : hostCompactOps) {
+        rewriter.setInsertionPoint(op);
+
+        if (auto extractRowsOp = dyn_cast<pim::PimExtractRowsOp>(op)) {
+          auto inputType = dyn_cast<ShapedType>(extractRowsOp.getInput().getType());
+          if (!inputType || !inputType.hasStaticShape() || inputType.getRank() != 2) {
+            extractRowsOp.emitOpError("host-side extract_rows lowering requires a static rank-2 input");
+            hasFailure = true;
+            continue;
+          }
+
+          int64_t numCols = inputType.getDimSize(1);
+          SmallVector<Value> replacementRows;
+          replacementRows.reserve(extractRowsOp.getOutputs().size());
+          for (auto rowIndex : llvm::seq<size_t>(0, extractRowsOp.getOutputs().size())) {
+            SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(static_cast<int64_t>(rowIndex)),
+                                                 rewriter.getIndexAttr(0)};
+            SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(numCols)};
+            SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            replacementRows.push_back(memref::SubViewOp::create(
+                                        rewriter, extractRowsOp.getLoc(), extractRowsOp.getInput(), offsets, sizes, strides)
+                                        .getResult());
+          }
+
+          extractRowsOp->replaceAllUsesWith(ValueRange(replacementRows));
+          extractRowsOp->erase();
+          continue;
+        }
+
+        auto concatOp = cast<pim::PimConcatOp>(op);
+        concatOp.emitOpError("host-side concat must be folded away or lowered into pim.core before materialization");
+        hasFailure = true;
+      }
     }
 
     if (hasFailure) {

src/PIM/Pass/PimCodegen/VerificationPass.cpp

@@ -18,13 +18,33 @@ namespace {
 
 static bool isAddressOnlyHostOp(Operation* op) {
   return isa<arith::ConstantOp,
+             pim::PimEmptyManyOp,
              memref::AllocOp,
              memref::GetGlobalOp,
              memref::SubViewOp,
              memref::CastOp,
              memref::CollapseShapeOp,
              memref::ExpandShapeOp,
-             spatial::SpatChannelNewOp>(op);
+             memref::CopyOp>(op);
+}
+
+// Looser than isCodegenAddressableValue: follows view ops without requiring contiguity.
+// Used for memref.copy operands which may be non-contiguous subviews.
+static bool isBaseAddressableValue(Value value) {
+  while (true) {
+    if (isa<BlockArgument>(value))
+      return true;
+    Operation* defOp = value.getDefiningOp();
+    if (!defOp)
+      return false;
+    if (isa<pim::PimEmptyManyOp, memref::AllocOp, memref::GetGlobalOp>(defOp))
+      return true;
+    if (auto subview = dyn_cast<memref::SubViewOp>(defOp)) { value = subview.getSource(); continue; }
+    if (auto cast = dyn_cast<memref::CastOp>(defOp)) { value = cast.getSource(); continue; }
+    if (auto collapse = dyn_cast<memref::CollapseShapeOp>(defOp)) { value = collapse.getSrc(); continue; }
+    if (auto expand = dyn_cast<memref::ExpandShapeOp>(defOp)) { value = expand.getSrc(); continue; }
+    return false;
+  }
 }
 
 static bool isCodegenAddressableValue(Value value) {
@@ -32,12 +52,14 @@ static bool isCodegenAddressableValue(Value value) {
   if (failed(resolvedAddress))
     return false;
   return isa<BlockArgument>(resolvedAddress->base)
-      || isa<memref::AllocOp, memref::GetGlobalOp>(resolvedAddress->base.getDefiningOp());
+      || isa<pim::PimEmptyManyOp, memref::AllocOp, memref::GetGlobalOp>(resolvedAddress->base.getDefiningOp());
 }
 
 static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
   if (isa<pim::PimMemCopyHostToDevOp>(op))
     return operandIndex == 1;
+  if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
+    return operandIndex == 1;
   if (isa<pim::PimMemCopyDevToHostOp>(op))
     return operandIndex == 0;
   return false;
@@ -69,6 +91,12 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
           continue;
         }
 
+        if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(&op)) {
+          if (failed(verifyCoreWeights(moduleOp, coreBatchOp)) || failed(verifyCoreOperands(coreBatchOp)))
+            hasFailure = true;
+          continue;
+        }
+
         if (auto returnOp = dyn_cast<func::ReturnOp>(&op)) {
           if (failed(verifyReturnOp(returnOp)))
             hasFailure = true;
@@ -92,10 +120,11 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
   }
 
 private:
-  static LogicalResult verifyCoreWeights(ModuleOp moduleOp, pim::PimCoreOp coreOp) {
+  template <typename CoreOpTy>
+  static LogicalResult verifyCoreWeights(ModuleOp moduleOp, CoreOpTy coreOp) {
     bool hasFailure = false;
     for (auto [weightIndex, weight] : llvm::enumerate(coreOp.getWeights())) {
-      auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
+      auto getGlobalOp = weight.template getDefiningOp<memref::GetGlobalOp>();
       if (!getGlobalOp) {
         coreOp.emitOpError() << "weight #" << weightIndex
                              << " must be materialized as memref.get_global before JSON codegen";
@@ -131,7 +160,8 @@ private:
     return success(!hasFailure);
   }
 
-  static LogicalResult verifyCoreOperands(pim::PimCoreOp coreOp) {
+  template <typename CoreOpTy>
+  static LogicalResult verifyCoreOperands(CoreOpTy coreOp) {
     return walkPimCoreBlock(
       coreOp.getBody().front(), StaticValueKnowledge {}, [](Operation& op, const StaticValueKnowledge& knowledge) {
         bool hasFailure = false;
@@ -155,7 +185,7 @@ private:
             continue;
           }
 
-          if (!isa<memref::AllocOp>(resolvedAddress->base.getDefiningOp())) {
+          if (!isa<pim::PimEmptyManyOp, memref::AllocOp>(resolvedAddress->base.getDefiningOp())) {
             op.emitOpError() << "operand #" << operandIndex
                              << " must be backed by device-local memory; materialize host values with pim.memcp_hd";
             hasFailure = true;
@@ -174,6 +204,13 @@ private:
       return verifyAddressOnlySource(op, collapseOp.getSrc());
     if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(op))
       return verifyAddressOnlySource(op, expandOp.getSrc());
+    if (auto copyOp = dyn_cast<memref::CopyOp>(op)) {
+      if (!isBaseAddressableValue(copyOp.getSource()) || !isBaseAddressableValue(copyOp.getTarget())) {
+        op->emitOpError("depends on a value that is not backed by addressable storage");
+        return failure();
+      }
+      return success();
+    }
     return success();
   }
 

test/PIM/DCPTest.cpp

@@ -477,7 +477,7 @@ int testDCPGraphCrossbarExhaustion() {
 
   const std::vector<Weight> nodeWeights = {10, 10, 10};
   const std::vector<CrossbarUsage> nodeCrossbarUsage = {1, 1, 1};
-  GraphDCP graph(nodeWeights, {}, nodeCrossbarUsage);
+  GraphDCP graph(nodeWeights, {}, {},nodeCrossbarUsage);
   graph.setMaxCpuCount(3);
   graph.runDcp();
 

validation/raptor.py

@@ -1,4 +1,5 @@
 import re
+import shlex
 import subprocess
 from pathlib import Path
 from colorama import Fore, Style
@@ -37,8 +38,12 @@ def _parse_pim_pass_timings(output_text):
     return pass_timings
 
 
+def _format_command(cmd):
+    return shlex.join(str(arg) for arg in cmd)
+
+
 def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
-                        crossbar_size, crossbar_count, cwd=None, reporter=None):
+                        crossbar_size, crossbar_count, core_count=None, cwd=None, reporter=None):
     # Define the arguments, with the possibility to set crossbar size and count
     args = [
         network_path,
@@ -51,10 +56,18 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
         f"--crossbar-count={crossbar_count}",
         "--enable-timing",
     ]
+    if core_count is not None:
+        args.append(f"--core-count={core_count}")
+
+    cmd = [str(raptor_onnx_path)] + [str(arg) for arg in args]
+    if reporter is not None:
+        reporter.log(f"  Raptor command: {_format_command(cmd)}")
+    else:
+        print(f"Raptor command: {_format_command(cmd)}")
 
     try:
         output_text = run_command_with_reporter(
-            [str(raptor_onnx_path)] + [str(arg) for arg in args],
+            cmd,
             cwd=cwd,
             reporter=reporter,
             capture_output=True,

validation/validate.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import argparse
-import shlex
 import signal
 import subprocess
 import sys
@@ -11,12 +10,6 @@ from validate_one import ProgressReporter, clean_workspace_artifacts, validate_n
 from raptor import PIM_PASS_LABELS
 
 
-def format_command(cmd):
-    if isinstance(cmd, (list, tuple)):
-        return shlex.join(str(arg) for arg in cmd)
-    return str(cmd)
-
-
 def format_return_status(returncode):
     if returncode < 0:
         signal_num = -returncode
@@ -34,8 +27,6 @@ def print_validation_error(reporter, rel, exc):
           file=sys.stderr, flush=True)
     if isinstance(exc, subprocess.CalledProcessError):
         print(format_return_status(exc.returncode), file=sys.stderr, flush=True)
-        print("Retry command:", file=sys.stderr, flush=True)
-        print(format_command(exc.cmd), file=sys.stderr, flush=True)
     else:
         print(f"{type(exc).__name__}: {exc}", file=sys.stderr, flush=True)
     print("=" * 72, file=sys.stderr, flush=True)
@@ -65,6 +56,8 @@ def main():
     ap.add_argument("--threshold", type=float, default=1e-3, help="Max allowed diff per output element.")
     ap.add_argument("--crossbar-size", type=int, default=64)
     ap.add_argument("--crossbar-count", type=int, default=8)
+    ap.add_argument("--core-count", type=int, default=None,
+                    help="Core count to pass to Raptor. If omitted, Raptor uses its default.")
     ap.add_argument("--clean", action="store_true",
                     help="Remove generated validation artifacts under each model workspace and exit.")
     a = ap.parse_args()
@@ -114,7 +107,7 @@ def main():
         try:
             result = validate_network(
                 onnx_path, a.raptor_path, a.onnx_include_dir, simulator_dir,
-                crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count,
+                crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count, core_count=a.core_count,
                 threshold=a.threshold,
                 reporter=reporter,
                 model_index=index,

validation/validate_one.py

@@ -1,4 +1,3 @@
-import argparse
 import json
 import numpy as np
 import subprocess
@@ -37,7 +36,7 @@ class ValidationResult:
 
 
 class ProgressReporter:
-    def __init__(self, total_models, stages_per_model=STAGE_COUNT):
+    def __init__(self, total_models, stages_per_model=STAGE_COUNT, enabled=None):
         self.total_models = total_models
         self.stages_per_model = stages_per_model
         self.total_steps = max(1, total_models * stages_per_model)
@@ -45,7 +44,7 @@ class ProgressReporter:
         self.passed_models = 0
         self.failed_models = 0
         self.current_label = ""
-        self.enabled = True
+        self.enabled = sys.stdout.isatty() if enabled is None else enabled
         self.columns = shutil.get_terminal_size((100, 20)).columns
         self.suspended = False
 
@@ -258,7 +257,7 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1
 
 
 def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
-                     simulator_dir, crossbar_size=64, crossbar_count=8, threshold=1e-3,
+                     simulator_dir, crossbar_size=64, crossbar_count=8, core_count=None, threshold=1e-3,
                      reporter=None, model_index=1, model_total=1):
     network_onnx_path = Path(network_onnx_path).resolve()
     raptor_path = Path(raptor_path).resolve()
@@ -313,7 +312,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
         pim_pass_timings = compile_with_raptor(
             network_mlir_path, raptor_path, raptor_dir / network_onnx_path.stem,
-            crossbar_size, crossbar_count,
+            crossbar_size, crossbar_count, core_count=core_count,
             cwd=raptor_dir, reporter=reporter)
         print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
         reporter.advance()
@@ -350,18 +349,3 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             reporter.log("=" * 72)
         if owns_reporter:
             reporter.finish()
-
-
-if __name__ == '__main__':
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--network-onnx", required=True)
-    ap.add_argument("--raptor-path", required=True)
-    ap.add_argument("--onnx-include-dir", required=True)
-    a = ap.parse_args()
-
-    simulator_dir = Path(__file__).parent.resolve() / ".." / "backend-simulators" / "pim" / "pim-simulator"
-
-    passed = validate_network(
-        a.network_onnx, a.raptor_path, a.onnx_include_dir, simulator_dir
-    )
-    raise SystemExit(0 if passed.passed else 1)