even faster on pimcomp models

src/PIM/CMakeLists.txt

@@ -121,6 +121,7 @@ add_pim_library(OMPIMAccel
   OMPimCommon
   OMPimBufferization
   OMPimHostConstantFolding
+  OMPimInstructionSelection
   OMPimVerification
   MLIRTensorInferTypeOpInterfaceImpl
 )

src/PIM/Compiler/CMakeLists.txt

@@ -31,6 +31,7 @@ add_pim_library(OMPimCompilerUtils
   OMPimCommon
   OMPimBufferization
   OMPimHostConstantFolding
+  OMPimInstructionSelection
   OMPimLocalMemoryPlanning
   OMPimVerification
   OMPimPasses

src/PIM/Compiler/PimCodeGen.cpp

@@ -108,6 +108,23 @@ static Operation* getDiagnosticAnchor(mlir::Value value) {
   return nullptr;
 }
 
+static bool isZeroSplatGlobal(mlir::Value value) {
+  auto getGlobalOp = value.getDefiningOp<memref::GetGlobalOp>();
+  auto moduleOp = getGlobalOp ? getGlobalOp->getParentOfType<ModuleOp>() : ModuleOp();
+  auto globalOp = moduleOp ? lookupGlobalForGetGlobal(moduleOp, getGlobalOp) : memref::GlobalOp();
+  if (!globalOp || !globalOp.getConstant() || !globalOp.getInitialValue())
+    return false;
+  auto denseAttr = dyn_cast<DenseElementsAttr>(*globalOp.getInitialValue());
+  if (!denseAttr || !denseAttr.isSplat())
+    return false;
+  Attribute valueAttr = denseAttr.getSplatValue<Attribute>();
+  if (auto floatAttr = dyn_cast<FloatAttr>(valueAttr))
+    return floatAttr.getValue().isZero() && !floatAttr.getValue().isNegative();
+  if (auto integerAttr = dyn_cast<IntegerAttr>(valueAttr))
+    return integerAttr.getValue().isZero();
+  return false;
+}
+
 // PIM instruction immediates are serialized as signed int32_t fields today
 // (`sldi` goes through checkedI32OrCrash), so local addresses must stay within
 // the non-negative int32_t range.
@@ -541,14 +558,23 @@ void PimCodeGen::emitMemCopyOp(pim_binary::Opcode opcode,
                                size_t rs1Offset,
                                size_t size,
                                StringRef sizeFieldName) const {
-  setupRdRs1(rdAddr, rdOffset, rs1Addr, rs1Offset);
-
   pim_binary::InstructionRecord instruction;
   instruction.opcode = opcode;
   instruction.rd = 0;
   instruction.r1 = 1;
-  instruction.generic1 = 0;
-  instruction.generic2 = 0;
+  if (rdOffset == rs1Offset) {
+    setupRdRs1(rdAddr, 0, rs1Addr, 0);
+    instruction.generic1 = rdOffset == 0 ? 0 : 3;
+    instruction.generic2 = pim::checkedI32OrCrash(rdOffset, "shared address offset");
+  } else if (rdOffset != 0) {
+    setupRdRs1(rdAddr, 0, rs1Addr, rs1Offset);
+    instruction.generic1 = 1;
+    instruction.generic2 = pim::checkedI32OrCrash(rdOffset, "rd address offset");
+  } else {
+    setupRdRs1(rdAddr, 0, rs1Addr, 0);
+    instruction.generic1 = rs1Offset == 0 ? 0 : 2;
+    instruction.generic2 = pim::checkedI32OrCrash(rs1Offset, "rs1 address offset");
+  }
   instruction.generic3 = pim::checkedI32OrCrash(size, sizeFieldName);
   emitInstruction(instruction);
 }
@@ -584,6 +610,16 @@ void PimCodeGen::codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticVa
   auto hostSourceOffset = indexOf(loadOp.getHostSourceOffset(), knowledge);
   assert(succeeded(deviceTargetOffset) && succeeded(hostSourceOffset)
          && "pim.memcp_hd offsets must be statically resolvable during codegen");
+  if (isZeroSplatGlobal(loadOp.getHostSource())) {
+    setupRd(addressOf(loadOp.getDeviceTarget(), knowledge), *deviceTargetOffset);
+    pim_binary::InstructionRecord instruction;
+    instruction.opcode = pim_binary::Opcode::lldi;
+    instruction.rd = 0;
+    instruction.r2OrImm = 0;
+    instruction.generic3 = loadOp.getSize();
+    emitInstruction(instruction);
+    return;
+  }
   emitMemCopyOp(pim_binary::Opcode::ld,
                 addressOf(loadOp.getDeviceTarget(), knowledge),
                 *deviceTargetOffset,
@@ -619,6 +655,26 @@ void PimCodeGen::codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledg
                 "len");
 }
 
+void PimCodeGen::codeGenVMVOp(pim::PimVMVOp vmvOp, const StaticValueKnowledge& knowledge) const {
+  auto targetOffset = indexOf(vmvOp.getTargetOffset(), knowledge);
+  auto sourceOffset = indexOf(vmvOp.getSourceOffset(), knowledge);
+  auto sourceStride = indexOf(vmvOp.getSourceStride(), knowledge);
+  assert(succeeded(targetOffset) && succeeded(sourceOffset) && succeeded(sourceStride)
+         && "pim.vmv operands must be statically resolvable during codegen");
+  auto sourceType = cast<ShapedType>(vmvOp.getSource().getType());
+  int32_t bitwidth = getVectorElementBitwidthOrCrash(sourceType);
+  ensureVectorBitwidth(bitwidth, bitwidth);
+  setupRdRs1Rs2(addressOf(vmvOp.getTarget(), knowledge), *targetOffset,
+                addressOf(vmvOp.getSource(), knowledge), *sourceOffset, 0, *sourceStride);
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::vmv;
+  instruction.rd = 0;
+  instruction.r1 = 1;
+  instruction.r2OrImm = 2;
+  instruction.generic3 = vmvOp.getLength();
+  emitInstruction(instruction);
+}
+
 void PimCodeGen::codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const {
   auto outputOffset = indexOf(receiveOp.getOutputOffset(), knowledge);
   auto sourceCoreId = indexOf(receiveOp.getSourceCoreId(), knowledge);
@@ -723,89 +779,6 @@ void PimCodeGen::emitUnaryVectorOp(pim_binary::Opcode opcode,
   emitInstruction(instruction);
 }
 
-void PimCodeGen::codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const {
-  auto srcAddr = addressOf(plan.source, knowledge);
-  auto dstAddr = addressOf(plan.destination, knowledge);
-
-  size_t maxElementOffset = plan.totalBytes == 0 ? 0 : plan.totalBytes - plan.elementBytes;
-  int32_t maxSourceAddress = pim::checkedI32OrCrash(
-    pim::checkedAddOrCrash(srcAddr, maxElementOffset, "transpose source address"), "transpose source address");
-  int32_t maxDestinationAddress =
-    pim::checkedI32OrCrash(pim::checkedAddOrCrash(dstAddr, maxElementOffset, "transpose destination address"),
-                           "transpose destination address");
-  (void) maxSourceAddress;
-  (void) maxDestinationAddress;
-
-  pim_binary::InstructionRecord copyInstruction;
-  copyInstruction.opcode = pim_binary::Opcode::lmv;
-  copyInstruction.rd = 0;
-  copyInstruction.r1 = 1;
-
-  size_t maxRunElements = static_cast<size_t>(std::numeric_limits<int32_t>::max()) / plan.elementBytes;
-  auto emitRun = [&](size_t sourceStart, size_t destinationStart, size_t runLength) {
-    while (runLength != 0) {
-      size_t chunkElements = std::min(runLength, maxRunElements);
-      // totalBytes was checked when the plan was compiled, so these bounded
-      // products cannot overflow.
-      size_t sourceOffset = sourceStart * plan.elementBytes;
-      size_t destinationOffset = destinationStart * plan.elementBytes;
-      size_t byteSize = chunkElements * plan.elementBytes;
-      assert(sourceOffset <= plan.totalBytes - byteSize && destinationOffset <= plan.totalBytes - byteSize);
-      genSetRegisterImmediate(0, static_cast<int32_t>(dstAddr + destinationOffset));
-      genSetRegisterImmediate(1, static_cast<int32_t>(srcAddr + sourceOffset));
-      copyInstruction.generic3 = static_cast<int32_t>(byteSize);
-      emitInstruction(copyInstruction);
-      sourceStart += chunkElements;
-      destinationStart += chunkElements;
-      runLength -= chunkElements;
-    }
-  };
-
-  if (plan.storagePreserving) {
-    emitRun(0, 0, plan.totalElements);
-    return;
-  }
-
-  size_t rank = plan.sourceShape.size();
-  SmallVector<size_t> sourceIndices(rank, 0);
-  size_t destinationFlat = 0;
-  size_t runSourceStart = 0;
-  size_t runDestinationStart = 0;
-  size_t runLength = 0;
-
-  for (size_t sourceFlat = 0; sourceFlat < plan.totalElements; ++sourceFlat) {
-    if (runLength != 0 && destinationFlat != runDestinationStart + runLength) {
-      emitRun(runSourceStart, runDestinationStart, runLength);
-      runSourceStart = sourceFlat;
-      runDestinationStart = destinationFlat;
-      runLength = 0;
-    }
-    if (runLength == 0) {
-      runSourceStart = sourceFlat;
-      runDestinationStart = destinationFlat;
-    }
-    ++runLength;
-
-    if (runLength == maxRunElements) {
-      emitRun(runSourceStart, runDestinationStart, runLength);
-      runLength = 0;
-    }
-
-    if (sourceFlat + 1 == plan.totalElements)
-      break;
-    for (size_t sourceDim = rank; sourceDim-- > 0;) {
-      destinationFlat += plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]];
-      if (++sourceIndices[sourceDim] < plan.sourceShape[sourceDim])
-        break;
-      sourceIndices[sourceDim] = 0;
-      destinationFlat -= plan.destinationRewinds[sourceDim];
-    }
-  }
-
-  if (runLength != 0)
-    emitRun(runSourceStart, runDestinationStart, runLength);
-}
-
 static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
   SmallVector<Operation*> coreLikeOps;
   for (Operation& op : funcOp.getBody().front())
@@ -1015,6 +988,7 @@ static LogicalResult executeCompiledCorePlan(
       coreCodeGen.codeGenStoreOp(cast<pim::PimMemCopyDevToHostOp>(node.op), knowledge);
       break;
     case CompiledCoreOpKind::Lmv:     coreCodeGen.codeGenLmvOp(cast<pim::PimMemCopyOp>(node.op), knowledge); break;
+    case CompiledCoreOpKind::VMV:     coreCodeGen.codeGenVMVOp(cast<pim::PimVMVOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Receive: coreCodeGen.codeGenReceiveOp(cast<pim::PimReceiveOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Send:    coreCodeGen.codeGenSendOp(cast<pim::PimSendOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Concat:  coreCodeGen.codeGenConcatOp(cast<pim::PimConcatOp>(node.op), knowledge); break;
@@ -1024,7 +998,6 @@ static LogicalResult executeCompiledCorePlan(
       else
         return failure();
       break;
-    case CompiledCoreOpKind::Transpose: coreCodeGen.codeGenTransposeOp(*node.transposePlan, knowledge); break;
     case CompiledCoreOpKind::VVAdd:     emitBinary(cast<pim::PimVVAddOp>(node.op), pim_binary::Opcode::vvadd); break;
     case CompiledCoreOpKind::VVSub:     emitBinary(cast<pim::PimVVSubOp>(node.op), pim_binary::Opcode::vvsub); break;
     case CompiledCoreOpKind::VVMul:     emitBinary(cast<pim::PimVVMulOp>(node.op), pim_binary::Opcode::vvmul); break;
@@ -1224,7 +1197,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
 
     for (auto [slot, fileName] : llvm::enumerate(weightFiles)) {
       xbarsPerGroup.push_back(weights[slot].shape[1] / static_cast<int64_t>(crossbarSize));
-      std::string sourcePath = outputDirPath + "/weights/" + fileName;
+      std::string sourcePath = "../weights/" + fileName;
       std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin";
       sys::fs::remove(targetPath);
       if (auto error = sys::fs::create_link(sourcePath, targetPath)) {

src/PIM/Compiler/PimCodeGen.hpp

@@ -25,7 +25,6 @@
 namespace onnx_mlir {
 
 struct CompiledCoreProgram;
-struct CompiledTransposePlan;
 class PimInstructionWriter;
 
 struct MemEntry {
@@ -214,6 +213,7 @@ public:
   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 codeGenVMVOp(pim::PimVMVOp vmvOp, const StaticValueKnowledge& knowledge) const;
 
   void codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const;
   void codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const;
@@ -222,7 +222,6 @@ public:
   template <typename MVMTy>
   void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge);
 
-  void codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const;
 };
 
 OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -281,6 +281,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
   verifyExplicitPimCoreCount();
 
   if (pimOnlyCodegen) {
+    pm.addPass(createPimInstructionSelectionPass());
     pm.addPass(createPimLocalMemoryPlanningPass());
     pm.addPass(createPimVerificationPass());
     pm.addPass(createEmitPimCodePass());
@@ -315,6 +316,8 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(mlir::createLowerAffinePass());
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
+    pm.addPass(createPimInstructionSelectionPass());
+    pm.addPass(createMessagePass("Pim instructions selected"));
     pm.addPass(createPimLocalMemoryPlanningPass());
     pm.addPass(createMessagePass("Pim local memory planned"));
     pm.addPass(createPimVerificationPass());

src/PIM/Compiler/PimCoreProgram.cpp

@@ -1,12 +1,6 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/BuiltinTypes.h"
-
-#include "llvm/ADT/STLExtras.h"
-
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -20,11 +14,11 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   if (isa<pim::PimMemCopyHostToDevOp>(op)) return CompiledCoreOpKind::Load;
   if (isa<pim::PimMemCopyDevToHostOp>(op)) return CompiledCoreOpKind::Store;
   if (isa<pim::PimMemCopyOp>(op)) return CompiledCoreOpKind::Lmv;
+  if (isa<pim::PimVMVOp>(op)) return CompiledCoreOpKind::VMV;
   if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
   if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
   if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
   if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
-  if (isa<pim::PimTransposeOp>(op)) return CompiledCoreOpKind::Transpose;
   if (isa<pim::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;
   if (isa<pim::PimVVSubOp>(op)) return CompiledCoreOpKind::VVSub;
   if (isa<pim::PimVVMulOp>(op)) return CompiledCoreOpKind::VVMul;
@@ -38,77 +32,6 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   return failure();
 }
 
-static bool isStoragePreservingTranspose(ArrayRef<size_t> sourceShape, ArrayRef<int64_t> permutation) {
-  SmallVector<unsigned> sourceNonUnitDims;
-  SmallVector<unsigned> destinationSourceNonUnitDims;
-  for (auto [dim, size] : llvm::enumerate(sourceShape))
-    if (size != 1)
-      sourceNonUnitDims.push_back(dim);
-  for (int64_t sourceDim : permutation)
-    if (sourceShape[sourceDim] != 1)
-      destinationSourceNonUnitDims.push_back(static_cast<unsigned>(sourceDim));
-  return sourceNonUnitDims == destinationSourceNonUnitDims;
-}
-
-static FailureOr<CompiledTransposePlan> compileTransposePlan(pim::PimTransposeOp transposeOp) {
-  auto sourceType = cast<ShapedType>(transposeOp.getInput().getType());
-  ArrayRef<int64_t> sourceShape = sourceType.getShape();
-  size_t rank = sourceShape.size();
-  CompiledTransposePlan plan;
-  plan.source = transposeOp.getInput();
-  plan.destination = transposeOp.getOutputBuffer();
-  plan.elementBytes = getElementTypeSizeInBytes(sourceType.getElementType());
-  auto totalElements = pim::checkedSize(sourceType.getNumElements(), transposeOp, "transpose elements");
-  if (failed(totalElements)) return failure();
-  plan.totalElements = *totalElements;
-  auto totalBytes = pim::checkedMul(plan.totalElements, plan.elementBytes, transposeOp, "transpose byte size");
-  if (failed(totalBytes)) return failure();
-  plan.totalBytes = *totalBytes;
-
-  SmallVector<int64_t> permutation = map_to_vector(transposeOp.getPermutation().getAsRange<IntegerAttr>(),
-                                                   [](IntegerAttr attr) { return attr.getInt(); });
-  if (permutation.size() != rank) {
-    transposeOp.emitOpError("requires permutation rank to match source rank for PIM codegen");
-    return failure();
-  }
-
-  SmallVector<size_t> destinationShape(rank);
-  plan.destinationStrides.assign(rank, 1);
-  plan.destinationDimensionForSource.assign(rank, 0);
-  plan.destinationRewinds.assign(rank, 0);
-  SmallVector<bool> seenSourceDimensions(rank, false);
-  for (size_t dim = 0; dim < rank; ++dim) {
-    auto size = pim::checkedSize(sourceShape[dim], transposeOp, "transpose source dimension");
-    if (failed(size)) return failure();
-    plan.sourceShape.push_back(*size);
-  }
-  for (auto [destinationDim, sourceDim] : llvm::enumerate(permutation)) {
-    if (sourceDim < 0 || static_cast<size_t>(sourceDim) >= rank || seenSourceDimensions[sourceDim]) {
-      transposeOp.emitOpError("requires a valid permutation containing each source dimension exactly once");
-      return failure();
-    }
-    seenSourceDimensions[sourceDim] = true;
-    destinationShape[destinationDim] = plan.sourceShape[sourceDim];
-    plan.destinationDimensionForSource[sourceDim] = destinationDim;
-  }
-  for (size_t dim = rank; dim > 1; --dim) {
-    auto stride = pim::checkedMul(
-      plan.destinationStrides[dim - 1], destinationShape[dim - 1], transposeOp, "transpose destination stride");
-    if (failed(stride)) return failure();
-    plan.destinationStrides[dim - 2] = *stride;
-  }
-  for (size_t sourceDim = 0; sourceDim < rank; ++sourceDim) {
-    auto rewind = pim::checkedMul(plan.sourceShape[sourceDim],
-                                  plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]],
-                                  transposeOp,
-                                  "transpose destination rewind");
-    if (failed(rewind)) return failure();
-    plan.destinationRewinds[sourceDim] = *rewind;
-  }
-  plan.storagePreserving = isStoragePreservingTranspose(plan.sourceShape, permutation);
-  return plan;
-}
-
 static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<CompiledCoreNode>& plan) {
   for (Operation& op : block) {
     if (isa<pim::PimHaltOp, scf::YieldOp, memref::GetGlobalOp>(op) || isCoreStaticAddressOp(&op))
@@ -188,11 +111,6 @@ static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<Compi
     CompiledCoreNode node;
     node.op = &op;
     node.opKind = *opKind;
-    if (*opKind == CompiledCoreOpKind::Transpose) {
-      auto transposePlan = compileTransposePlan(cast<pim::PimTransposeOp>(op));
-      if (failed(transposePlan)) return failure();
-      node.transposePlan = *transposePlan;
-    }
     plan.push_back(std::move(node));
   }
   return success();

src/PIM/Compiler/PimCoreProgram.hpp

@@ -6,34 +6,19 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include <memory>
-#include <optional>
-
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
 
 namespace onnx_mlir {
 
-struct CompiledTransposePlan {
-  mlir::Value source;
-  mlir::Value destination;
-  size_t elementBytes = 0;
-  size_t totalElements = 0;
-  size_t totalBytes = 0;
-  llvm::SmallVector<size_t> sourceShape;
-  llvm::SmallVector<size_t> destinationStrides;
-  llvm::SmallVector<unsigned> destinationDimensionForSource;
-  llvm::SmallVector<size_t> destinationRewinds;
-  bool storagePreserving = false;
-};
-
 enum class CompiledCoreOpKind : uint8_t {
   Load,
   Store,
   Lmv,
+  VMV,
   Receive,
   Send,
   Concat,
   Vmm,
-  Transpose,
   VVAdd,
   VVSub,
   VVMul,
@@ -62,7 +47,6 @@ struct CompiledCoreNode {
   llvm::SmallVector<int64_t> caseValues;
   llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
   std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
-  std::optional<CompiledTransposePlan> transposePlan;
 };
 
 struct CompiledCoreProgram {

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

@@ -180,7 +180,11 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
                                          kRowStripIndexMap, rewriter, loc);
 }
 
-FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+template <typename BuildActivation>
+static FailureOr<Value> applyRowStripActivation(const RowStripPhysicalValue& value,
+                                                PatternRewriter& rewriter,
+                                                Location loc,
+                                                BuildActivation buildActivation) {
   auto storageType = cast<RankedTensorType>(value.storage.getType());
   const int64_t laneCount = storageType.getDimSize(0);
   auto batchOp = createSpatComputeBatch(rewriter,
@@ -193,10 +197,9 @@ FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRe
                                           FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
                                             rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
                                           if (failed(fragment)) return failure();
-                                          Value relu = spatial::SpatReluOp::create(
-                                            rewriter, loc, value.fragmentType, *fragment).getResult();
+                                          Value result = buildActivation(*fragment);
                                           publishGraphBatchPhysicalFragment(
-                                            rewriter, loc, relu, args.outputs.front(), args.lane);
+                                            rewriter, loc, result, args.outputs.front(), args.lane);
                                           return success();
                                         });
   if (failed(batchOp))
@@ -204,6 +207,19 @@ FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRe
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+  return applyRowStripActivation(value, rewriter, loc, [&](Value fragment) {
+    return spatial::SpatReluOp::create(rewriter, loc, value.fragmentType, fragment).getResult();
+  });
+}
+
+FailureOr<Value> applyRowStripSilu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+  return applyRowStripActivation(value, rewriter, loc, [&](Value fragment) {
+    Value sigmoid = spatial::SpatSigmoidOp::create(rewriter, loc, value.fragmentType, fragment).getResult();
+    return spatial::SpatVMulOp::create(rewriter, loc, value.fragmentType, fragment, sigmoid).getResult();
+  });
+}
+
 FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
                                      Value bias,
                                      PatternRewriter& rewriter,

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

@@ -61,6 +61,10 @@ mlir::FailureOr<mlir::Value> applyRowStripRelu(const RowStripPhysicalValue& valu
                                                mlir::PatternRewriter& rewriter,
                                                mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> applyRowStripSilu(const RowStripPhysicalValue& value,
+                                               mlir::PatternRewriter& rewriter,
+                                               mlir::Location loc);
+
 mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
                                                   mlir::Value bias,
                                                   mlir::PatternRewriter& rewriter,

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -58,6 +58,11 @@ lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp pl
   return applyRowStripRelu(input, rewriter, planOp.getLoc());
 }
 
+static FailureOr<Value>
+lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripSilu(input, rewriter, planOp.getLoc());
+}
+
 static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
                                              spatial::SpatBiasAddPlanOp planOp,
                                              PatternRewriter& rewriter) {
@@ -349,6 +354,51 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         rewriter.replaceOp(planOp, computeOp.getResults());
         continue;
       }
+
+      if (auto planOp = dyn_cast<spatial::SpatSiluPlanOp>(&op)) {
+        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
+          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+          });
+          if (outputBlueprint == planOp.getResult().getUsers().end()) {
+            planOp.emitOpError("row-strip SiLU plan requires a row-strip blueprint result");
+            signalPassFailure();
+            return;
+          }
+
+          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+          rewriter.setInsertionPoint(planOp);
+          FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
+          if (failed(lowered)) {
+            planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan");
+            signalPassFailure();
+            return;
+          }
+          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+          if (failed(output)) {
+            signalPassFailure();
+            return;
+          }
+          rowStripValues[blueprint.getResult()] = *output;
+          eraseAfterLowering.insert(planOp);
+          eraseAfterLowering.insert(blueprint);
+          continue;
+        }
+
+        rewriter.setInsertionPoint(planOp);
+        auto computeOp = createSpatCompute<1>(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
+            Value sigmoid = spatial::SpatSigmoidOp::create(
+              rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
+            Value silu = spatial::SpatVMulOp::create(
+              rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
+            spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
+          });
+        rewriter.replaceOp(planOp, computeOp.getResults());
+        continue;
+      }
       if (auto planOp = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
         auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
           auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
@@ -650,6 +700,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
                    spatial::SpatBiasAddPlanOp,
                    spatial::SpatAddPlanOp,
                    spatial::SpatReluPlanOp,
+                   spatial::SpatSiluPlanOp,
                    spatial::SpatMaxPool2DPlanOp,
                    spatial::SpatGlobalAveragePoolPlanOp,
                    spatial::SpatMaterializeLayoutOp>(op)

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -8,6 +8,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "llvm/ADT/SmallVector.h"
 
@@ -50,13 +51,14 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatAddPlanOp> addPlans(funcOp.getOps<spatial::SpatAddPlanOp>());
   SmallVector<spatial::SpatConcatPlanOp> concatPlans(funcOp.getOps<spatial::SpatConcatPlanOp>());
   SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
+  SmallVector<spatial::SpatSiluPlanOp> siluPlans(funcOp.getOps<spatial::SpatSiluPlanOp>());
   SmallVector<spatial::SpatMaxPool2DPlanOp> maxPoolPlans(funcOp.getOps<spatial::SpatMaxPool2DPlanOp>());
   SmallVector<spatial::SpatGlobalAveragePoolPlanOp> globalAveragePoolPlans(
     funcOp.getOps<spatial::SpatGlobalAveragePoolPlanOp>());
   SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
   SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
   if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !addPlans.empty()
-      || !concatPlans.empty() || !reluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty()
+      || !concatPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty()
       || !globalAveragePoolPlans.empty() || !materializers.empty()) {
     return;
   }
@@ -121,6 +123,14 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  RewritePatternSet fusionPatterns(ctx);
+  populateElementwiseFusionPatterns(fusionPatterns, ctx);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(fusionPatterns)))) {
+    moduleOp.emitError("failed to fuse layout-aware ONNX elementwise patterns");
+    signalPassFailure();
+    return;
+  }
+
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
     moduleOp.emitError("failed to locate the PIM entry function during ONNX-to-Spatial lowering");

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -149,6 +149,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatAddPlanOp,
             spatial::SpatConcatPlanOp,
             spatial::SpatReluPlanOp,
+            spatial::SpatSiluPlanOp,
             spatial::SpatMaxPool2DPlanOp,
             spatial::SpatGlobalAveragePoolPlanOp,
             spatial::SpatBlueprintOp,

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -17,6 +17,7 @@ void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::ML
 
 void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateElementwiseFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateGemmPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

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

@@ -1256,7 +1256,7 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
       for (int64_t kernelIndex = 0; kernelIndex < tiling.kernelElements; ++kernelIndex) {
         const int64_t kernelH = kernelIndex / wType.getDimSize(3);
         const int64_t kernelW = kernelIndex % wType.getDimSize(3);
-        const int64_t targetRow = localChannel * tiling.kernelElements + kernelIndex;
+        const int64_t targetRow = kernelIndex * tiling.channelsPerTile + localChannel;
         for (int64_t multiplierIndex = 0; multiplierIndex < tiling.outputMultiplier; ++multiplierIndex) {
           const int64_t globalOutChannel = globalChannel * tiling.outputMultiplier + multiplierIndex;
           const int64_t sourceFlatIndex =
@@ -1326,16 +1326,49 @@ static Value createInputTile(Value input,
   Value channelOffset = tiling.channelsPerTile == 1
                         ? channelTileIndex
                         : affineMulConst(rewriter, loc, channelTileIndex, tiling.channelsPerTile, anchorOp);
-  Value tile4D = createConvInputPatch(input,
-                                      inputTileType,
-                                      batchIndex,
-                                      channelOffset,
-                                      inputHeightOffset,
-                                      inputWidthOffset,
-                                      dilationHeight,
-                                      dilationWidth,
-                                      rewriter,
-                                      loc);
+  Value tile4D;
+  if (dilationHeight == 1 && dilationWidth == 1) {
+    SmallVector<OpFoldResult> offsets {batchIndex, inputHeightOffset, inputWidthOffset, channelOffset};
+    SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                     rewriter.getIndexAttr(inputTileType.getDimSize(1)),
+                                     rewriter.getIndexAttr(inputTileType.getDimSize(2)),
+                                     rewriter.getIndexAttr(tiling.channelsPerTile)};
+    tile4D = tensor::ExtractSliceOp::create(
+      rewriter, loc, inputTileType, input, offsets, sizes, getUnitStrides(rewriter, 4));
+  }
+  else {
+    auto pixelType = RankedTensorType::get(
+      {1, 1, 1, tiling.channelsPerTile}, inputTileType.getElementType(), inputTileType.getEncoding());
+    tile4D = tensor::EmptyOp::create(rewriter, loc, inputTileType.getShape(), inputTileType.getElementType());
+    for (int64_t kernelH = 0; kernelH < inputTileType.getDimSize(1); ++kernelH)
+      for (int64_t kernelW = 0; kernelW < inputTileType.getDimSize(2); ++kernelW) {
+        Value sourceHeight = affineAddConst(rewriter, loc, inputHeightOffset, kernelH * dilationHeight, anchorOp);
+        Value sourceWidth = affineAddConst(rewriter, loc, inputWidthOffset, kernelW * dilationWidth, anchorOp);
+        SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                        rewriter.getIndexAttr(1),
+                                        rewriter.getIndexAttr(1),
+                                        rewriter.getIndexAttr(tiling.channelsPerTile)};
+        Value pixel = tensor::ExtractSliceOp::create(
+          rewriter,
+          loc,
+          pixelType,
+          input,
+          SmallVector<OpFoldResult> {batchIndex, sourceHeight, sourceWidth, channelOffset},
+          sizes,
+          getUnitStrides(rewriter, 4));
+        tile4D = tensor::InsertSliceOp::create(
+          rewriter,
+          loc,
+          pixel,
+          tile4D,
+          SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                     rewriter.getIndexAttr(kernelH),
+                                     rewriter.getIndexAttr(kernelW),
+                                     rewriter.getIndexAttr(0)},
+          sizes,
+          getUnitStrides(rewriter, 4));
+      }
+  }
   auto collapsedType = RankedTensorType::get({1, tiling.tileInputRows}, inputTileType.getElementType());
   return tensor::CollapseShapeOp::create(rewriter,
                                          loc,
@@ -1531,10 +1564,18 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
                                         state.padWidthEnd,
                                         rewriter,
                                         loc);
+  auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
+  auto channelLastInputType = RankedTensorType::get({paddedInputType.getDimSize(0),
+                                                      paddedInputType.getDimSize(2),
+                                                      paddedInputType.getDimSize(3),
+                                                      paddedInputType.getDimSize(1)},
+                                                     paddedInputType.getElementType());
+  Value channelLastInput = ONNXTransposeOp::create(
+    rewriter, loc, channelLastInputType, paddedInput, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
   Value packedWeights = buildPackedWeights(wDenseAttr, state.wType, *tiling, rewriter, loc);
 
   Value expandedBias;
-  SmallVector<Value> batchInputs {paddedInput};
+  SmallVector<Value> batchInputs {channelLastInput};
   if (state.hasBias) {
     expandedBias = expandBiasIfNeeded(state.b, rewriter, loc);
     auto biasType = dyn_cast<RankedTensorType>(expandedBias.getType());
@@ -1553,9 +1594,8 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
     {1, static_cast<int64_t>(crossbarSize.getValue())}, state.outType.getElementType());
   auto piecesType = spatial::getGraphBatchPhysicalResultType(
     tiling->totalPatches * tiling->numChannelTiles, rowTileType);
-  auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
   auto inputTileType =
-    RankedTensorType::get({1, tiling->channelsPerTile, state.wType.getDimSize(2), state.wType.getDimSize(3)},
+    RankedTensorType::get({1, state.wType.getDimSize(2), state.wType.getDimSize(3), tiling->channelsPerTile},
                           paddedInputType.getElementType());
   SmallVector<Value> batchWeights;
   if (tiling->numChannelTiles == 1) {
@@ -1766,7 +1806,7 @@ static Value unpackRowsFromParallelGemm(Value packedRows,
 }
 
 static Value createWeightMatrix(
-  Value weights, const ConvGemmPlan& plan, PatternRewriter& rewriter, Location loc) {
+  Value weights, const ConvGemmPlan& plan, bool transpose, PatternRewriter& rewriter, Location loc) {
   auto buildWeightMatrix = [&](Value weight) -> Value {
     Value flattened = tensor::CollapseShapeOp::create(rewriter,
                                                       loc,
@@ -1776,6 +1816,8 @@ static Value createWeightMatrix(
                                                         {0},
                                                         {1, 2, 3}
     });
+    if (!transpose)
+      return flattened;
     return ONNXTransposeOp::create(rewriter, loc, plan.wTransType, flattened, rewriter.getI64ArrayAttr({1, 0}))
       .getResult();
   };
@@ -1783,8 +1825,9 @@ static Value createWeightMatrix(
   if (isCompileTimeComputable(weights))
     return buildWeightMatrix(weights);
 
+  RankedTensorType resultType = transpose ? plan.wTransType : plan.wFlatType;
   auto computeOp =
-    createSpatCompute<1>(rewriter, loc, TypeRange {plan.wTransType}, {}, ValueRange {weights}, [&](Value weight) {
+    createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, ValueRange {weights}, [&](Value weight) {
       spatial::SpatYieldOp::create(rewriter, loc, buildWeightMatrix(weight));
     });
   return computeOp.getResult(0);
@@ -1852,21 +1895,26 @@ static Value createPaddedPixelMajorWeightConstant(DenseElementsAttr sourceAttr,
                                                   const ConvLoweringState& state,
                                                   int64_t paddedK,
                                                   int64_t paddedC,
+                                                  int64_t packFactor,
                                                   PatternRewriter& rewriter) {
   auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.wType.getElementType());
   SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
   SmallVector<Attribute> paddedValues(
     paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
-  for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel)
-    for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH)
-      for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW)
-        for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
-          const int64_t sourceFlatIndex =
-            (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
-          const int64_t patchIndex =
-            ((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel;
-          paddedValues[patchIndex * paddedC + outChannel] = sourceValues[sourceFlatIndex];
-        }
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  for (int64_t copy = 0; copy < packFactor; ++copy)
+    for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel)
+      for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH)
+        for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW)
+          for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
+            const int64_t sourceFlatIndex =
+              (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
+            const int64_t patchIndex =
+              ((kernelH * state.wWidth) + kernelW) * state.numChannelsIn + inChannel;
+            const int64_t packedRow = copy * patchSize + patchIndex;
+            const int64_t packedColumn = copy * state.numChannelsOut + outChannel;
+            paddedValues[packedRow * paddedC + packedColumn] = sourceValues[sourceFlatIndex];
+          }
   return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
     DenseElementsAttr::get(paddedType, paddedValues), paddedType);
 }
@@ -2381,7 +2429,7 @@ static Value createStreamedConvRows(const ConvLoweringState& state,
   Value gemmBias = state.hasBias ? state.b : createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
   Value packedBias = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);
   Value gemmRows = ONNXGemmOp::create(rewriter, loc, plan.gemmOutputRowsType, inputRows,
-    packedWeights, packedBias, APFloat(1.0f), APFloat(1.0f), 0, 0).getY();
+    packedWeights, packedBias, APFloat(1.0f), APFloat(1.0f), 0, !wDenseAttr).getY();
   return maybeUnpackChunkRows(gemmRows, plan, rewriter, loc);
 }
 
@@ -2401,9 +2449,10 @@ static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
   ConvGemmPlan plan =
     buildConvGemmPlan(state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0,
                       state.batchSize * state.outHeight * state.outWidth);
-  // Prepare weight matrix W for crossbar storage:
-  //   W: [Cout, Cin, KH, KW] -> [Cout, patchSize] -> [patchSize, Cout]
-  Value weightMatrix = createWeightMatrix(state.w, plan, rewriter, loc);
+  // Static weights use the crossbar [patchSize, Cout] layout. Runtime weights
+  // stay in ONNX's contiguous [Cout, patchSize] layout and Gemm consumes them
+  // through transB without materializing a transpose.
+  Value weightMatrix = createWeightMatrix(state.w, plan, static_cast<bool>(wDenseAttr), rewriter, loc);
   Value gemmInputRows = createIm2colRows(state, preparedInput, plan, rewriter, loc);
   Value gemmB = buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc);
   Value gemmBias = createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
@@ -2420,7 +2469,7 @@ static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
                                       APFloat(1.0f),
                                       APFloat(1.0f),
                                       /*transA=*/0,
-                                      /*transB=*/0)
+                                      /*transB=*/!wDenseAttr)
                      .getY();
 
   return createCollectedConvOutput(ValueRange {gemmRows},
@@ -2452,7 +2501,7 @@ static Value rewriteStreamedConv(const ConvLoweringState& state,
 
   ConvGemmPlan seedPlan = buildConvGemmPlan(
     state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0, 1, forcedPackFactor);
-  Value weightMatrix = createWeightMatrix(state.w, seedPlan, rewriter, loc);
+  Value weightMatrix = createWeightMatrix(state.w, seedPlan, static_cast<bool>(wDenseAttr), rewriter, loc);
   Value collectedRows = createStreamedConvRows(state,
                                               preparedInput,
                                               weightMatrix,
@@ -2516,6 +2565,20 @@ static bool rowStripOutputChannelTileFitsOneCore(const ConvGeometry& geometry) {
          <= static_cast<int64_t>(crossbarCountInCore.getValue());
 }
 
+static int64_t chooseRowStripPixelPackFactor(const ConvLoweringState& state, int64_t xbarDim) {
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t baseWeightGroups = ceilIntegerDivide(patchSize, xbarDim)
+                                 * ceilIntegerDivide(state.numChannelsOut, xbarDim);
+  int64_t factor = std::min(state.outWidth, xbarDim / state.numChannelsOut);
+  while (factor > 1
+         && (state.outWidth % factor != 0
+             || ceilIntegerDivide(factor * patchSize, xbarDim)
+                  * ceilIntegerDivide(factor * state.numChannelsOut, xbarDim)
+                > baseWeightGroups))
+    --factor;
+  return std::max<int64_t>(factor, 1);
+}
+
 static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
   if (state.batchSize != 1) {
     failureReason = "batch_not_one";
@@ -2558,6 +2621,7 @@ static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& re
 }
 
 static FailureOr<Value> createBiasRowConstant(const ConvLoweringState& state,
+                                              int64_t packFactor,
                                               PatternRewriter& rewriter) {
   DenseElementsAttr denseAttr;
   if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr))
@@ -2566,10 +2630,14 @@ static FailureOr<Value> createBiasRowConstant(const ConvLoweringState& state,
   if (failed(channelValues))
     return failure();
 
-  auto biasType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
+  SmallVector<Attribute> packedValues;
+  packedValues.reserve(packFactor * state.numChannelsOut);
+  for (int64_t copy = 0; copy < packFactor; ++copy)
+    packedValues.append(channelValues->begin(), channelValues->end());
+  auto biasType = RankedTensorType::get({1, packFactor * state.numChannelsOut}, state.outType.getElementType());
   return getOrCreateConstant(rewriter,
                              rewriter.getInsertionBlock()->getParentOp(),
-                             DenseElementsAttr::get(biasType, *channelValues),
+                             DenseElementsAttr::get(biasType, packedValues),
                              biasType);
 }
 
@@ -2825,39 +2893,47 @@ static FailureOr<Value> createConvInputWindow(Value input,
   Value initWindow = createZeroTensorConstant(paddedWindowType, rewriter);
 
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  Value window = initWindow;
-  for (int64_t kernelRowIndex = 0; kernelRowIndex < state.wHeight; ++kernelRowIndex) {
-    Value kernelRow = getOrCreateIndexConstant(rewriter, anchorOp, kernelRowIndex);
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+  Value cKernelRows = getOrCreateIndexConstant(rewriter, anchorOp, state.wHeight);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c0,
+    cKernelRows,
+    c1,
+    ValueRange {initWindow},
+    [&](OpBuilder&, Location rowLoc, Value kernelRow, ValueRange iterArgs,
+        SmallVectorImpl<Value>& yielded) -> LogicalResult {
     FailureOr<Value> sourceRow =
       denseInput
         ? FailureOr<Value>(
-            extractDenseConvWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, loc))
-        : extractProjectedRowStripWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, loc);
+            extractDenseConvWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, rowLoc))
+        : extractProjectedRowStripWindowRow(input, sourceIndexTable, state, outputHeight, kernelRow, rewriter, rowLoc);
     if (failed(sourceRow))
       return failure();
     Value semanticRow = *sourceRow;
     if (state.padHeightBegin != 0 || state.padHeightEnd != 0) {
-      Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, loc);
-      semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, semanticRow, mask).getResult();
+      Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, rowLoc);
+      semanticRow = spatial::SpatVMulOp::create(rewriter, rowLoc, fragmentType, semanticRow, mask).getResult();
     }
-    Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, loc);
-    window = tensor::InsertSliceOp::create(rewriter,
-                                           loc,
-                                           paddedRow,
-                                           window,
-                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
-                                                                      rewriter.getIndexAttr(kernelRowIndex),
-                                                                      rewriter.getIndexAttr(0),
-                                                                      rewriter.getIndexAttr(0)},
-                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
-                                                                      rewriter.getIndexAttr(1),
-                                                                      rewriter.getIndexAttr(
-                                                                        state.xWidth + state.padWidthBegin
-                                                                        + state.padWidthEnd),
-                                                                      rewriter.getIndexAttr(state.numChannelsIn)},
-                                           getUnitStrides(rewriter, 4));
-  }
-  return window;
+    Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, rowLoc);
+    yielded.push_back(tensor::InsertSliceOp::create(
+      rewriter,
+      rowLoc,
+      paddedRow,
+      iterArgs.front(),
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), kernelRow,
+                                 rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+      SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
+                                 rewriter.getIndexAttr(state.xWidth + state.padWidthBegin
+                                                       + state.padWidthEnd),
+                                 rewriter.getIndexAttr(state.numChannelsIn)},
+      getUnitStrides(rewriter, 4)));
+    return success();
+  });
+  return failed(loop) ? FailureOr<Value>(failure())
+                      : FailureOr<Value>(loop->results.front());
 }
 
 static FailureOr<Value> createPixelMajorConvPatchRow(Value paddedWindow,
@@ -2878,20 +2954,92 @@ static FailureOr<Value> createPixelMajorConvPatchRow(Value paddedWindow,
                                    rewriter.getIndexAttr(state.wHeight),
                                    rewriter.getIndexAttr(state.wWidth),
                                    rewriter.getIndexAttr(state.numChannelsIn)};
-  SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1),
-                                     rewriter.getIndexAttr(1),
-                                     rewriter.getIndexAttr(state.dilationWidth),
-                                     rewriter.getIndexAttr(1)};
-  Value patch = tensor::ExtractSliceOp::create(
-    rewriter, loc, patchType, paddedWindow, offsets, sizes, strides);
+  Value patch;
+  if (state.dilationWidth == 1)
+    patch = tensor::ExtractSliceOp::create(
+      rewriter, loc, patchType, paddedWindow, offsets, sizes, getUnitStrides(rewriter, 4));
+  else {
+    auto columnType = RankedTensorType::get({1, state.wHeight, 1, state.numChannelsIn},
+                                            state.xType.getElementType(), state.xType.getEncoding());
+    patch = tensor::EmptyOp::create(rewriter, loc, patchType.getShape(), patchType.getElementType());
+    for (int64_t kernelColumn = 0; kernelColumn < state.wWidth; ++kernelColumn) {
+      Value sourceWidth =
+        affineAddConst(rewriter, loc, inputWidthOffset, kernelColumn * state.dilationWidth, anchorOp);
+      SmallVector<OpFoldResult> columnSizes {rewriter.getIndexAttr(1),
+                                             rewriter.getIndexAttr(state.wHeight),
+                                             rewriter.getIndexAttr(1),
+                                             rewriter.getIndexAttr(state.numChannelsIn)};
+      Value column = tensor::ExtractSliceOp::create(
+        rewriter, loc, columnType, paddedWindow,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), sourceWidth,
+                                   rewriter.getIndexAttr(0)},
+        columnSizes, getUnitStrides(rewriter, 4));
+      patch = tensor::InsertSliceOp::create(
+        rewriter, loc, column, patch,
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                   rewriter.getIndexAttr(kernelColumn), rewriter.getIndexAttr(0)},
+        columnSizes, getUnitStrides(rewriter, 4));
+    }
+  }
   return tensor::CollapseShapeOp::create(
            rewriter, loc, rowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}})
     .getResult();
 }
 
+static FailureOr<Value> createPackedPixelMajorConvPatchRow(Value paddedWindow,
+                                                           const ConvLoweringState& state,
+                                                           Value outputGroup,
+                                                           int64_t packFactor,
+                                                           PatternRewriter& rewriter,
+                                                           Location loc) {
+  if (packFactor == 1)
+    return createPixelMajorConvPatchRow(paddedWindow, state, outputGroup, rewriter, loc);
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  auto packedType = RankedTensorType::get(
+    {1, packFactor * patchSize}, state.xType.getElementType(), state.xType.getEncoding());
+  Value packed = tensor::EmptyOp::create(rewriter, loc, packedType.getShape(), packedType.getElementType());
+  Value outputStart = affineMulConst(rewriter, loc, outputGroup, packFactor, anchorOp);
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+  Value cPackFactor = getOrCreateIndexConstant(rewriter, anchorOp, packFactor);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c0,
+    cPackFactor,
+    c1,
+    ValueRange {packed},
+    [&](OpBuilder&, Location copyLoc, Value copy, ValueRange iterArgs,
+        SmallVectorImpl<Value>& yielded) -> LogicalResult {
+      Value outputWidth = createOrFoldAffineApply(
+        rewriter, copyLoc, rewriter.getAffineDimExpr(0) + rewriter.getAffineDimExpr(1),
+        ValueRange {outputStart, copy}, anchorOp);
+      FailureOr<Value> patch =
+        createPixelMajorConvPatchRow(paddedWindow, state, outputWidth, rewriter, copyLoc);
+      if (failed(patch))
+        return failure();
+      Value packedOffset = affineMulConst(rewriter, copyLoc, copy, patchSize, anchorOp);
+      yielded.push_back(tensor::InsertSliceOp::create(
+        rewriter,
+        copyLoc,
+        *patch,
+        iterArgs.front(),
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), packedOffset},
+        SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(patchSize)},
+        getUnitStrides(rewriter, 2)));
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  return loop->results.front();
+}
+
 static FailureOr<SmallVector<Value>> createConvInputTiles(Value paddedWindow,
                                                           const ConvLoweringState& state,
                                                           Value outputWidth,
+                                                          int64_t packFactor,
                                                           Value& partialInputScratch,
                                                           int64_t patchSize,
                                                           int64_t numKSlices,
@@ -2903,7 +3051,7 @@ static FailureOr<SmallVector<Value>> createConvInputTiles(Value paddedWindow,
   SmallVector<Value> inputTiles;
   inputTiles.reserve(numKSlices);
 
-  if (state.numChannelsIn % xbarDim == 0) {
+  if (packFactor == 1 && state.numChannelsIn % xbarDim == 0) {
     Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
     auto inputTileType = RankedTensorType::get(
       {1, 1, 1, xbarDim}, elementType, state.xType.getEncoding());
@@ -2938,8 +3086,8 @@ static FailureOr<SmallVector<Value>> createConvInputTiles(Value paddedWindow,
     return inputTiles;
   }
 
-  FailureOr<Value> patchRow =
-    createPixelMajorConvPatchRow(paddedWindow, state, outputWidth, rewriter, loc);
+  FailureOr<Value> patchRow = createPackedPixelMajorConvPatchRow(
+    paddedWindow, state, outputWidth, packFactor, rewriter, loc);
   if (failed(patchRow))
     return failure();
 
@@ -3081,17 +3229,18 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
                                                  Value input,
                                                  Value paddedWeights,
                                                  Value bias,
+                                                 int64_t packFactor,
                                                  int64_t paddedK,
                                                  int64_t numKSlices,
                                                  int64_t xbarDim,
                                                  PatternRewriter& rewriter,
                                                  Location loc) {
   const int64_t laneCount = state.outHeight;
-  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t patchSize = packFactor * state.numChannelsIn * state.wHeight * state.wWidth;
   const bool hasPartialInputTile = patchSize % xbarDim != 0;
   auto elementType = state.outType.getElementType();
   auto partialInputScratchType = RankedTensorType::get({1, xbarDim}, elementType);
-  auto outputPixelType = RankedTensorType::get({1, 1, 1, state.numChannelsOut}, elementType);
+  auto outputPixelType = RankedTensorType::get({1, 1, packFactor, state.numChannelsOut}, elementType);
   auto fragmentType = getRowStripFragmentType(state.outType);
   auto storageType = getRowStripStorageType(state.outType);
 
@@ -3106,7 +3255,7 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
       Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
+      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth / packFactor);
       FailureOr<Value> inputWindow =
         createConvInputWindow(args.inputs.front(), state, args.lane, rewriter, loc);
       if (failed(inputWindow))
@@ -3131,6 +3280,7 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
           FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
                                                                           state,
                                                                           localColumn,
+                                                                          packFactor,
                                                                           partialInputScratch,
                                                                           patchSize,
                                                                           numKSlices,
@@ -3141,7 +3291,7 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
             return failure();
           FailureOr<Value> output = createConvOutputRow(*inputTiles,
                                                         paddedK,
-                                                        state.numChannelsOut,
+                                                        packFactor * state.numChannelsOut,
                                                         args.weights.front(),
                                                         bias ? args.inputs[1] : Value(),
                                                         xbarDim,
@@ -3150,7 +3300,7 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
           if (failed(output))
             return failure();
           Value outputPixel = tensor::ExpandShapeOp::create(
-            rewriter, pixelLoc, outputPixelType, *output, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
+            rewriter, pixelLoc, outputPixelType, *output, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
           Value next = tensor::InsertSliceOp::create(
             rewriter,
             pixelLoc,
@@ -3158,11 +3308,11 @@ static FailureOr<Value> createRowStripConvOutput(const ConvLoweringState& state,
             iterArgs.front(),
             SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
                                        rewriter.getIndexAttr(0),
-                                       localColumn,
+                                       affineMulConst(rewriter, pixelLoc, localColumn, packFactor, anchorOp),
                                        rewriter.getIndexAttr(0)},
             SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
                                        rewriter.getIndexAttr(1),
-                                       rewriter.getIndexAttr(1),
+                                       rewriter.getIndexAttr(packFactor),
                                        rewriter.getIndexAttr(state.numChannelsOut)},
             getUnitStrides(rewriter, 4));
           yielded.push_back(next);
@@ -3255,6 +3405,7 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
           FailureOr<SmallVector<Value>> inputTiles = createConvInputTiles(*inputWindow,
                                                                           state,
                                                                           widthIndex,
+                                                                          /*packFactor=*/1,
                                                                           partialInputScratch,
                                                                           patchSize,
                                                                           numKSlices,
@@ -3313,29 +3464,34 @@ static FailureOr<Value>
     return failure();
 
   const int64_t xbarDim = geometry.xbarSize;
-  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
-  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
-  const int64_t paddedK = numKSlices * xbarDim;
+  const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim);
+  const int64_t basePaddedK = baseNumKSlices * xbarDim;
 
   if (!rowStripOutputTileFitsOneCore(geometry)) {
     Value tiledWeights =
-      standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+      standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, basePaddedK, xbarDim, rewriter);
     return createOutputChannelTiledRowStripConvOutput(
-      state, state.x, tiledWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+      state, state.x, tiledWeights, basePaddedK, baseNumKSlices, xbarDim, rewriter, loc);
   }
 
-  const int64_t paddedOutputChannels = ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
-  Value paddedWeights =
-    standard::createPaddedPixelMajorWeightConstant(weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
+  const int64_t packFactor = chooseRowStripPixelPackFactor(state, xbarDim);
+  const int64_t packedPatchSize = packFactor * basePatchSize;
+  const int64_t numKSlices = ceilIntegerDivide(packedPatchSize, xbarDim);
+  const int64_t paddedK = numKSlices * xbarDim;
+  const int64_t packedOutputChannels = packFactor * state.numChannelsOut;
+  const int64_t paddedOutputChannels = ceilIntegerDivide(packedOutputChannels, xbarDim) * xbarDim;
+  Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
+    weightDenseAttr, state, paddedK, paddedOutputChannels, packFactor, rewriter);
   FailureOr<Value> bias = failure();
   if (state.hasBias)
-    bias = createBiasRowConstant(state, rewriter);
+    bias = createBiasRowConstant(state, packFactor, rewriter);
   if (state.hasBias && failed(bias))
     return failure();
 
   return createRowStripConvOutput(
     state, state.x, paddedWeights, state.hasBias ? *bias : Value(),
-    paddedK, numKSlices, xbarDim, rewriter, loc);
+    packFactor, paddedK, numKSlices, xbarDim, rewriter, loc);
 }
 
 static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value rowStripStorage,
@@ -3351,32 +3507,36 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
 
   ConvGeometry geometry = buildConvGeometry(state);
   const int64_t xbarDim = geometry.xbarSize;
-  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
-  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
-  const int64_t paddedK = numKSlices * xbarDim;
+  const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim);
+  const int64_t basePaddedK = baseNumKSlices * xbarDim;
   auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
   if (!weightDenseAttr)
     return failure();
   if (!rowStripOutputTileFitsOneCore(geometry)) {
     Value tiledWeights =
-      standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+      standard::createPaddedOutputChannelTiledWeightConstant(weightDenseAttr, state, basePaddedK, xbarDim, rewriter);
     return createOutputChannelTiledRowStripConvOutput(
-      state, rowStripStorage, tiledWeights, paddedK, numKSlices, xbarDim, rewriter, loc);
+      state, rowStripStorage, tiledWeights, basePaddedK, baseNumKSlices, xbarDim, rewriter, loc);
   }
 
-  const int64_t paddedOutputChannels =
-    ceilIntegerDivide(state.numChannelsOut, xbarDim) * xbarDim;
+  const int64_t packFactor = chooseRowStripPixelPackFactor(state, xbarDim);
+  const int64_t packedPatchSize = packFactor * basePatchSize;
+  const int64_t numKSlices = ceilIntegerDivide(packedPatchSize, xbarDim);
+  const int64_t paddedK = numKSlices * xbarDim;
+  const int64_t packedOutputChannels = packFactor * state.numChannelsOut;
+  const int64_t paddedOutputChannels = ceilIntegerDivide(packedOutputChannels, xbarDim) * xbarDim;
   Value paddedWeights = standard::createPaddedPixelMajorWeightConstant(
-    weightDenseAttr, state, paddedK, paddedOutputChannels, rewriter);
+    weightDenseAttr, state, paddedK, paddedOutputChannels, packFactor, rewriter);
   FailureOr<Value> bias = failure();
   if (state.hasBias)
-    bias = createBiasRowConstant(state, rewriter);
+    bias = createBiasRowConstant(state, packFactor, rewriter);
   if (state.hasBias && failed(bias))
     return failure();
 
   return createRowStripConvOutput(
     state, rowStripStorage, paddedWeights, state.hasBias ? *bias : Value(),
-    paddedK, numKSlices, xbarDim, rewriter, loc);
+    packFactor, paddedK, numKSlices, xbarDim, rewriter, loc);
 }
 
 static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStripStorage,

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

@@ -16,6 +16,28 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
+struct SiluToSpatialPlan : OpRewritePattern<ONNXMulOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ONNXMulOp mulOp, PatternRewriter& rewriter) const override {
+    ONNXSigmoidOp sigmoidOp = mulOp.getA().getDefiningOp<ONNXSigmoidOp>();
+    Value input = mulOp.getB();
+    if (!sigmoidOp) {
+      sigmoidOp = mulOp.getB().getDefiningOp<ONNXSigmoidOp>();
+      input = mulOp.getA();
+    }
+    if (!sigmoidOp || sigmoidOp.getX() != input || !sigmoidOp->hasOneUse()
+        || sigmoidOp.getResult().getType() != input.getType() || mulOp.getResult().getType() != input.getType())
+      return failure();
+
+    auto plan = spatial::SpatSiluPlanOp::create(
+      rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, rewriter.getStringAttr("nchw"));
+    rewriter.replaceOp(mulOp, plan.getResult());
+    rewriter.eraseOp(sigmoidOp);
+    return success();
+  }
+};
+
 static DenseElementsAttr getDenseConstantAttr(Value value) {
   if (auto constantOp = value.getDefiningOp<arith::ConstantOp>())
     return dyn_cast<DenseElementsAttr>(constantOp.getValue());
@@ -219,6 +241,10 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
 
 } // namespace
 
+void populateElementwiseFusionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<SiluToSpatialPlan>(ctx);
+}
+
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   patterns.add<AddToSpatialCompute>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);

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

@@ -259,17 +259,6 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
   return *batchOp;
 }
 
-static Value
-createDynamicGemmBatchRow(Value lane, int64_t numOutCols, ConversionPatternRewriter& rewriter, Location loc) {
-  if (numOutCols == 1)
-    return lane;
-
-  MLIRContext* context = rewriter.getContext();
-  AffineExpr d0 = getAffineDimExpr(0, context);
-  return createOrFoldAffineApply(
-    rewriter, loc, d0.floorDiv(numOutCols), ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp());
-}
-
 static Value extractDynamicGemmBColumn(
   Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), column};
@@ -373,33 +362,56 @@ static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
                                                               Value b,
                                                               RankedTensorType aType,
                                                               RankedTensorType bType,
-                                                              RankedTensorType scalarPiecesType,
+                                                              RankedTensorType columnPiecesType,
                                                               RankedTensorType outType,
+                                                              bool transposeB,
                                                               ConversionPatternRewriter& rewriter,
                                                               Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   const int64_t reductionSize = aType.getDimSize(1);
-  const int64_t laneCount = numOutRows * numOutCols;
+  auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
+  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
+  auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
   auto batchOp = createSpatComputeBatch(
     rewriter,
     loc,
-    TypeRange {scalarPiecesType},
-    laneCount,
+    TypeRange {columnPiecesType},
+    numOutCols,
     ValueRange {},
     ValueRange {a, b},
     [&](detail::SpatComputeBatchBodyArgs args) {
-      Value row = createDynamicGemmBatchRow(args.lane, numOutCols, rewriter, loc);
-      Value column =
-        onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-
-      auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
-      auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-      Value aVector = extractDynamicGemmRowVector(args.inputs[0], row, vectorType, rewriter, loc);
-      Value bVector = extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc);
-      Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-
-      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
+      Value bVector = transposeB
+                        ? extractDynamicGemmRowVector(args.inputs[1], args.lane, vectorType, rewriter, loc)
+                        : extractDynamicGemmBColumn(args.inputs[1], args.lane, vectorType, rewriter, loc);
+      Value columnInit = tensor::EmptyOp::create(rewriter, loc, columnType.getShape(), columnType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
+      Value cNumOutRows =
+        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cNumOutRows,
+        c1,
+        ValueRange {columnInit},
+        [&](OpBuilder&, Location nestedLoc, Value row, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value aVector = extractDynamicGemmRowVector(args.inputs[0], row, vectorType, rewriter, nestedLoc);
+          Value scalar = spatial::SpatVVDMulOp::create(rewriter, nestedLoc, scalarType, aVector, bVector).getResult();
+          Value next = tensor::InsertSliceOp::create(rewriter,
+                                                     nestedLoc,
+                                                     scalar,
+                                                     iterArgs.front(),
+                                                     SmallVector<OpFoldResult> {row, rewriter.getIndexAttr(0)},
+                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                               rewriter.getIndexAttr(1)},
+                                                     getUnitStrides(rewriter, 2));
+          yielded.push_back(next);
+          return success();
+        });
+      assert(succeeded(loop) && "dynamic Gemm row loop construction must succeed");
+      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -415,7 +427,7 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
                                                                       float beta,
                                                                       ConversionPatternRewriter& rewriter,
                                                                       Location loc) {
-  const int64_t laneCount = scalarPiecesType.getDimSize(0);
+  const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   SmallVector<Value> inputs {scalarPieces};
   if (bias)
@@ -428,43 +440,62 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
     Value outputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType()).getResult();
     Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
     Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
-    Value cLaneCount = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), laneCount);
+    Value cNumOutCols = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutCols);
+    auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
     auto loop = buildNormalizedScfFor(
       rewriter,
       loc,
       c0,
-      cLaneCount,
+      cNumOutCols,
       c1,
       ValueRange {outputInit},
-      [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+      [&](OpBuilder&, Location nestedLoc, Value column, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
         Value outputAcc = iterArgs.front();
-        Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, nestedLoc);
-        Value column =
-          onnx_mlir::affineModConst(rewriter, nestedLoc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-        SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
-        if (failed(scalar))
+        FailureOr<Value> columnPiece =
+          extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, column, columnType);
+        if (failed(columnPiece))
           return failure();
-        if (alpha != 1.0f) {
-          Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc);
-          *scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, *scalar, alphaTensor).getResult();
-        }
-        if (biasArg) {
-          Value biasScalar =
-            createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, nestedLoc);
-          if (beta != 1.0f) {
-            Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, nestedLoc);
-            biasScalar =
-              spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).getResult();
-          }
-          *scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, *scalar, biasScalar).getResult();
-        }
-        SmallVector<OpFoldResult> outputOffsets {row, column};
-        Value outputNext =
-          tensor::InsertSliceOp::create(rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
-            .getResult();
-        yielded.push_back(outputNext);
+        Value cNumOutRows =
+          getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
+        auto rowLoop = buildNormalizedScfFor(
+          rewriter,
+          nestedLoc,
+          c0,
+          cNumOutRows,
+          c1,
+          ValueRange {outputAcc},
+          [&](OpBuilder&, Location rowLoc, Value row, ValueRange rowIterArgs, SmallVectorImpl<Value>& rowYielded) {
+            SmallVector<OpFoldResult> scalarOffsets {row, rewriter.getIndexAttr(0)};
+            SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            Value scalar = tensor::ExtractSliceOp::create(
+              rewriter, rowLoc, scalarType, *columnPiece, scalarOffsets, scalarSizes, unitStrides);
+            if (alpha != 1.0f) {
+              Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, rowLoc);
+              scalar = spatial::SpatVMulOp::create(rewriter, rowLoc, scalarType, scalar, alphaTensor).getResult();
+            }
+            if (biasArg) {
+              Value biasScalar = createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, rowLoc);
+              if (beta != 1.0f) {
+                Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, rowLoc);
+                biasScalar =
+                  spatial::SpatVMulOp::create(rewriter, rowLoc, scalarType, biasScalar, betaTensor).getResult();
+              }
+              scalar = spatial::SpatVAddOp::create(rewriter, rowLoc, scalarType, scalar, biasScalar).getResult();
+            }
+            Value next = tensor::InsertSliceOp::create(rewriter,
+                                                       rowLoc,
+                                                       scalar,
+                                                       rowIterArgs.front(),
+                                                       SmallVector<OpFoldResult> {row, column},
+                                                       scalarSizes,
+                                                       unitStrides);
+            rowYielded.push_back(next);
+            return success();
+          });
+        if (failed(rowLoop))
+          return failure();
+        yielded.push_back(rowLoop->results.front());
         return success();
       });
     if (failed(loop))
@@ -660,16 +691,10 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     aType = transposedType;
   }
 
-  if (gemmOpAdaptor.getTransB()) {
-    auto bShape = bType.getShape();
-    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
-    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult();
-    bType = transposedType;
-  }
-
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   const int64_t reductionSize = aType.getDimSize(1);
+  const bool transposeB = gemmOpAdaptor.getTransB();
 
   if (!isCompileTimeComputable(b)) {
     bool hasC = hasGemmBias(c);
@@ -690,8 +715,9 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
       biasType = *verifiedBiasType;
     }
 
-    if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize
-        || bType.getDimSize(1) != numOutCols) {
+    const int64_t bReductionSize = bType.getDimSize(transposeB ? 1 : 0);
+    const int64_t bOutputColumns = bType.getDimSize(transposeB ? 0 : 1);
+    if (aType.getDimSize(0) != numOutRows || bReductionSize != reductionSize || bOutputColumns != numOutCols) {
       gemmOp.emitOpError("has inconsistent A, B, and output shapes");
       return failure();
     }
@@ -702,8 +728,9 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
       return failure();
     }
 
-    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount64, RankedTensorType::get({1, 1}, outType.getElementType()));
-    auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, rewriter, loc);
+    auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(numOutCols, columnType);
+    auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, transposeB, rewriter, loc);
     if (failed(batchOp))
       return failure();
     auto outputCompute = createDynamicGemmOutputCompute(
@@ -714,6 +741,13 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     return success();
   }
 
+  if (transposeB) {
+    auto bShape = bType.getShape();
+    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
+    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    bType = transposedType;
+  }
+
   auto scaledB = materializeScaledConstantTensor(b, gemmOpAdaptor.getAlpha().convertToFloat(), rewriter, loc);
   if (failed(scaledB)) {
     gemmOp.emitOpError("requires constant Gemm input B when alpha is not 1.0");

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

@@ -389,41 +389,6 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
   return *batchOp;
 }
 
-static Value extractDynamicBatchedBColumn(Value matrix,
-                                          ArrayRef<int64_t> sourceBatchShape,
-                                          ArrayRef<int64_t> outputBatchShape,
-                                          Value outputBatchIndex,
-                                          Value column,
-                                          RankedTensorType vectorType,
-                                          PatternRewriter& rewriter,
-                                          Location loc) {
-  auto columnSliceType = RankedTensorType::get({1, vectorType.getDimSize(1), 1}, vectorType.getElementType());
-  Value sourceBatchIndex =
-    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
-  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), rewriter.getIndexAttr(0), column};
-  SmallVector<OpFoldResult> sizes {
-    rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  Value columnSlice = tensor::ExtractSliceOp::create(rewriter, loc, columnSliceType, matrix, offsets, sizes, strides);
-  auto collapsedType = RankedTensorType::get({vectorType.getDimSize(1)}, vectorType.getElementType());
-  Value collapsed = tensor::CollapseShapeOp::create(rewriter,
-                                                    loc,
-                                                    collapsedType,
-                                                    columnSlice,
-                                                    SmallVector<ReassociationIndices> {
-                                                      {0, 1, 2}
-  })
-                      .getResult();
-  return tensor::ExpandShapeOp::create(rewriter,
-                                       loc,
-                                       vectorType,
-                                       collapsed,
-                                       SmallVector<ReassociationIndices> {
-                                         {0, 1}
-  })
-    .getResult();
-}
-
 static Value extractDynamicBatchedRowVector(Value matrix,
                                             ArrayRef<int64_t> sourceBatchShape,
                                             ArrayRef<int64_t> outputBatchShape,
@@ -449,7 +414,7 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
                                                                      ArrayRef<int64_t> outputBatchShape,
                                                                      RankedTensorType aType,
                                                                      RankedTensorType bType,
-                                                                     RankedTensorType scalarPiecesType,
+                                                                     RankedTensorType columnPiecesType,
                                                                      RankedTensorType outType,
                                                                      PatternRewriter& rewriter,
                                                                      Location loc) {
@@ -457,29 +422,52 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
   const int64_t numOutRows = outType.getDimSize(1);
   const int64_t numOutCols = outType.getDimSize(2);
   const int64_t reductionSize = aType.getDimSize(2);
-  const int64_t laneCount = numBatches * numOutRows * numOutCols;
+  const int64_t laneCount = numBatches * numOutCols;
+  auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
+  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
+  auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
   auto batchOp = createSpatComputeBatch(
     rewriter,
     loc,
-    TypeRange {scalarPiecesType},
+    TypeRange {columnPiecesType},
     laneCount,
     ValueRange {},
     ValueRange {a, b},
     [&](detail::SpatComputeBatchBodyArgs args) {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-      Value batch = affineFloorDivConst(rewriter, loc, args.lane, numOutRows * numOutCols, anchorOp);
-      Value batchLane = affineModConst(rewriter, loc, args.lane, numOutRows * numOutCols, anchorOp);
-      Value row = affineFloorDivConst(rewriter, loc, batchLane, numOutCols, anchorOp);
-      Value column = affineModConst(rewriter, loc, batchLane, numOutCols, anchorOp);
-
-      auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
-      auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-      Value aVector = extractDynamicBatchedRowVector(
-        args.inputs[0], aBatchShape, outputBatchShape, batch, row, vectorType, rewriter, loc);
-      Value bVector = extractDynamicBatchedBColumn(
+      Value batch = affineFloorDivConst(rewriter, loc, args.lane, numOutCols, anchorOp);
+      Value column = affineModConst(rewriter, loc, args.lane, numOutCols, anchorOp);
+      Value bVector = extractDynamicBatchedRowVector(
         args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc);
-      Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
+      Value columnInit = tensor::EmptyOp::create(rewriter, loc, columnType.getShape(), columnType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
+      Value cNumOutRows =
+        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cNumOutRows,
+        c1,
+        ValueRange {columnInit},
+        [&](OpBuilder&, Location nestedLoc, Value row, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value aVector = extractDynamicBatchedRowVector(
+            args.inputs[0], aBatchShape, outputBatchShape, batch, row, vectorType, rewriter, nestedLoc);
+          Value scalar = spatial::SpatVVDMulOp::create(rewriter, nestedLoc, scalarType, aVector, bVector).getResult();
+          Value next = tensor::InsertSliceOp::create(rewriter,
+                                                     nestedLoc,
+                                                     scalar,
+                                                     iterArgs.front(),
+                                                     SmallVector<OpFoldResult> {row, rewriter.getIndexAttr(0)},
+                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                               rewriter.getIndexAttr(1)},
+                                                     getUnitStrides(rewriter, 2));
+          yielded.push_back(next);
+          return success();
+        });
+      assert(succeeded(loop) && "dynamic MatMul row loop construction must succeed");
+      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -492,9 +480,8 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
                                                           PatternRewriter& rewriter,
                                                           Location loc) {
   const int64_t laneCount = scalarPiecesType.getDimSize(0);
-  const int64_t numOutRows = outType.getDimSize(1);
   const int64_t numOutCols = outType.getDimSize(2);
-  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
+  auto columnType = RankedTensorType::get({outType.getDimSize(1), 1}, outType.getElementType());
 
   auto computeOp = createSpatCompute<1>(
     rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult {
@@ -513,19 +500,18 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
         [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
           Value outputAcc = iterArgs.front();
           Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-          Value batch = affineFloorDivConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
-          Value batchLane = affineModConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
-          Value row = affineFloorDivConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
-          Value column = affineModConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
-          FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
-          if (failed(scalar))
+          Value batch = affineFloorDivConst(rewriter, nestedLoc, lane, numOutCols, anchorOp);
+          Value column = affineModConst(rewriter, nestedLoc, lane, numOutCols, anchorOp);
+          FailureOr<Value> columnPiece =
+            extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, columnType);
+          if (failed(columnPiece))
             return failure();
-          SmallVector<OpFoldResult> outputOffsets {batch, row, column};
+          SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), column};
           SmallVector<OpFoldResult> outputSizes = {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(outType.getDimSize(1)), rewriter.getIndexAttr(1)};
           Value next =
             tensor::InsertSliceOp::create(
-              rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+              rewriter, nestedLoc, *columnPiece, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
               .getResult();
           yielded.push_back(next);
           return success();
@@ -1011,12 +997,14 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
         return success();
       }
     }
-    const int64_t laneCount = plan.batch * plan.m * plan.n;
+    const int64_t laneCount = plan.batch * plan.n;
+    auto columnType = RankedTensorType::get({plan.m, 1}, shapeInfo->outType.getElementType());
     auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(
-      laneCount, RankedTensorType::get({1, 1}, shapeInfo->outType.getElementType()));
+      laneCount, columnType);
+    Value transposedRhs = transposeLastTwoDims(plan.rhs, rewriter, loc);
     auto batchOp = createBatchedVvdmulBatch(plan.lhs,
                                             plan.lhsBatchShape,
-                                            plan.rhs,
+                                            transposedRhs,
                                             plan.rhsBatchShape,
                                             plan.outputBatchShape,
                                             plan.lhsType,

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -34,6 +34,8 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& l
 static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
     return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(user))
+    return getSelectedLayout(layouts, siluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
     return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(user))
@@ -62,7 +64,7 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, Selecte
 }
 
 static bool canConsumeRowStripAsUser(Operation* user) {
-  if (isa<spatial::SpatReluPlanOp>(user))
+  if (isa<spatial::SpatReluPlanOp, spatial::SpatSiluPlanOp>(user))
     return true;
   if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
     auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
@@ -105,13 +107,12 @@ static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
   return SelectedLayout::PixelMajorRowStrip;
 }
 
-static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
-                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
+static SelectedLayout chooseActivationLayout(Value input,
+                                             Value result,
+                                             llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip)
     return SelectedLayout::DenseNchw;
-  if (!hasRowStripConsumer(reluPlan.getResult()))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
+  if (!allUsersCanHandleRowStrip(result, layouts))
     return SelectedLayout::DenseNchw;
   return SelectedLayout::PixelMajorRowStrip;
 }
@@ -239,13 +240,21 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           continue;
         }
         if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
-          SelectedLayout selected = chooseReluLayout(reluPlan, layouts);
+          SelectedLayout selected = chooseActivationLayout(reluPlan.getInput(), reluPlan.getResult(), layouts);
           if (layouts[reluPlan.getResult()] != selected) {
             layouts[reluPlan.getResult()] = selected;
             changed = true;
           }
           continue;
         }
+        if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op)) {
+          SelectedLayout selected = chooseActivationLayout(siluPlan.getInput(), siluPlan.getResult(), layouts);
+          if (layouts[siluPlan.getResult()] != selected) {
+            layouts[siluPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
         if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
           SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
           if (layouts[biasAddPlan.getResult()] != selected) {
@@ -301,6 +310,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
         producedValue = concatPlan.getResult();
       else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
         producedValue = reluPlan.getResult();
+      else if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op))
+        producedValue = siluPlan.getResult();
       else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
         producedValue = maxPoolPlan.getResult();
       else if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op))

src/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.cpp

@@ -47,11 +47,11 @@ static bool isRuntimeMemoryTouch(Operation* op) {
   return isa<PimMemCopyHostToDevOp,
              PimMemCopyDevToHostOp,
              PimMemCopyOp,
+             PimVMVOp,
              PimReceiveOp,
              PimSendOp,
              PimConcatOp,
              PimVMMOp,
-             PimTransposeOp,
              PimVVAddOp,
              PimVVSubOp,
              PimVVMulOp,

src/PIM/Dialect/Pim/CMakeLists.txt

@@ -4,6 +4,7 @@ add_onnx_mlir_dialect_doc(pim Pim.td)
 add_subdirectory(Analysis)
 add_subdirectory(Transforms/Bufferization)
 add_subdirectory(Transforms/HostConstantFolding)
+add_subdirectory(Transforms/InstructionSelection)
 add_subdirectory(Transforms/LocalMemoryPlanning)
 add_subdirectory(Transforms/Verification)
 

src/PIM/Dialect/Pim/Pim.td

@@ -202,6 +202,36 @@ def PimMemCopyOp : PimOp<"memcp", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimVMVOp : PimOp<"vmv", [DestinationStyleOpInterface]> {
+  let summary = "Gather a strided vector within device memory";
+
+  let arguments = (ins
+    Index:$targetOffset,
+    Index:$sourceOffset,
+    Index:$sourceStride,
+    PimTensor:$target,
+    PimTensor:$source,
+    I32Attr:$length
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getTargetMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let assemblyFormat = [{
+    `[` $targetOffset `,` $sourceOffset `,` $sourceStride `]`
+    `(` $target `,` $source `)` attr-dict
+    `:` type($target) `,` type($source) `->` type($output)
+  }];
+}
+
 def PimConcatOp : PimOp<"concat", [DestinationStyleOpInterface]> {
   let summary = "Concatenate tensors";
 

src/PIM/Dialect/Pim/PimOpsVerify.cpp

@@ -118,6 +118,20 @@ LogicalResult PimCoreBatchOp::verify() {
   return verifyOnlyConstantExternalValues(getOperation(), getBody(), "pim.core_batch");
 }
 
+LogicalResult PimVMVOp::verify() {
+  if (getLength() <= 0)
+    return emitError("length must be positive");
+  if (failed(verifyCompatibleShapedTypes(
+        getOperation(), getTarget().getType(), getOutput().getType(), "target and output types must match")))
+    return failure();
+  auto targetType = cast<ShapedType>(getTarget().getType());
+  auto sourceType = dyn_cast<ShapedType>(getSource().getType());
+  if (!sourceType || !haveSameShapedContainerKind(getTarget().getType(), getSource().getType())
+      || targetType.getElementType() != sourceType.getElementType())
+    return emitError("target and source must use the same shaped container kind and element type");
+  return success();
+}
+
 LogicalResult PimVMMOp::verify() {
   if (failed(verifyCompatibleShapedTypes(
         getOperation(), getOutputBuffer().getType(), getOutput().getType(), "output buffer and output must match")))

src/PIM/Dialect/Pim/Transforms/InstructionSelection/CMakeLists.txt

@@ -0,0 +1,12 @@
+add_pim_library(OMPimInstructionSelection
+  InstructionSelectionPass.cpp
+
+  EXCLUDE_FROM_OM_LIBS
+
+  LINK_LIBS PUBLIC
+  MLIRArithDialect
+  MLIRSCFDialect
+  OMPimCommon
+  OMPimBufferization
+  PimOps
+)

src/PIM/Dialect/Pim/Transforms/InstructionSelection/InstructionSelectionPass.cpp

@@ -0,0 +1,167 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "llvm/ADT/STLExtras.h"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+
+using namespace llvm;
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static bool isStoragePreservingTranspose(ArrayRef<int64_t> sourceShape, ArrayRef<int64_t> permutation) {
+  SmallVector<int64_t> sourceNonUnitDims;
+  SmallVector<int64_t> destinationSourceNonUnitDims;
+  for (auto [dim, size] : llvm::enumerate(sourceShape))
+    if (size != 1)
+      sourceNonUnitDims.push_back(dim);
+  for (int64_t sourceDim : permutation)
+    if (sourceShape[sourceDim] != 1)
+      destinationSourceNonUnitDims.push_back(sourceDim);
+  return sourceNonUnitDims == destinationSourceNonUnitDims;
+}
+
+static Value indexConstant(PatternRewriter& rewriter, Location loc, int64_t value) {
+  return arith::ConstantIndexOp::create(rewriter, loc, value);
+}
+
+struct LowerTransposePattern final : OpRewritePattern<pim::PimTransposeOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(pim::PimTransposeOp op, PatternRewriter& rewriter) const override {
+    auto sourceType = dyn_cast<MemRefType>(op.getInput().getType());
+    auto targetType = dyn_cast<MemRefType>(op.getOutputBuffer().getType());
+    if (!sourceType || !targetType || !sourceType.hasStaticShape() || !targetType.hasStaticShape())
+      return op.emitOpError("requires static memref operands before PIM instruction selection");
+
+    ArrayRef<int64_t> sourceShape = sourceType.getShape();
+    size_t rank = sourceShape.size();
+    SmallVector<int64_t> permutation = map_to_vector(
+      op.getPermutation().getAsRange<IntegerAttr>(), [](IntegerAttr attr) { return attr.getInt(); });
+    if (permutation.size() != rank)
+      return op.emitOpError("requires permutation rank to match source rank");
+    SmallVector<bool> seen(rank, false);
+    for (int64_t sourceDim : permutation) {
+      if (sourceDim < 0 || static_cast<size_t>(sourceDim) >= rank || seen[sourceDim])
+        return op.emitOpError("requires a permutation containing each source dimension exactly once");
+      seen[sourceDim] = true;
+    }
+
+    auto totalElements = pim::checkedSize(sourceType.getNumElements(), op, "transpose elements");
+    size_t elementBytes = getElementTypeSizeInBytes(sourceType.getElementType());
+    auto totalBytes = failed(totalElements)
+      ? FailureOr<size_t>(failure())
+      : pim::checkedMul(*totalElements, elementBytes, op, "transpose byte size");
+    if (failed(totalBytes))
+      return failure();
+    auto byteSizeAttr = pim::getCheckedI32Attr(rewriter, op, *totalBytes, "transpose byte size");
+    if (failed(byteSizeAttr))
+      return failure();
+
+    Location loc = op.getLoc();
+    Value zero = indexConstant(rewriter, loc, 0);
+    if (isStoragePreservingTranspose(sourceShape, permutation)) {
+      auto copy = pim::PimMemCopyOp::create(rewriter, loc, targetType, zero, zero,
+                                            op.getOutputBuffer(), op.getInput(), *byteSizeAttr);
+      rewriter.replaceOp(op, copy.getOutput());
+      return success();
+    }
+
+    SmallVector<int64_t> sourceStrides(rank, 1);
+    for (size_t dim = rank; dim > 1; --dim)
+      sourceStrides[dim - 2] = sourceStrides[dim - 1] * sourceShape[dim - 1];
+
+    SmallVector<int64_t> destinationShape;
+    destinationShape.reserve(rank);
+    for (int64_t sourceDim : permutation)
+      destinationShape.push_back(sourceShape[sourceDim]);
+
+    size_t runStart = rank;
+    int64_t runLength = 1;
+    int64_t sourceStride = 1;
+    for (size_t destinationDim = rank; destinationDim-- > 0;) {
+      if (destinationShape[destinationDim] == 1)
+        continue;
+      int64_t candidateStride = sourceStrides[permutation[destinationDim]];
+      if (runStart != rank && candidateStride != runLength * sourceStride)
+        break;
+      runStart = destinationDim;
+      if (runLength == 1)
+        sourceStride = candidateStride;
+      runLength *= destinationShape[destinationDim];
+    }
+    assert(runStart != rank && "non-storage-preserving transpose must have a non-unit dimension");
+
+    int64_t outerCount = sourceType.getNumElements() / runLength;
+    auto lengthAttr = pim::getCheckedI32Attr(rewriter, op, runLength, "vmv length");
+    auto runBytesAttr = pim::getCheckedI32Attr(rewriter, op, runLength * elementBytes, "transpose run byte size");
+    if (failed(lengthAttr) || failed(runBytesAttr))
+      return failure();
+
+    Value upper = indexConstant(rewriter, loc, outerCount);
+    Value one = indexConstant(rewriter, loc, 1);
+    auto loop = scf::ForOp::create(rewriter, loc, zero, upper, one, ValueRange{op.getOutputBuffer()});
+    rewriter.setInsertionPointToStart(loop.getBody());
+    Value outer = loop.getInductionVar();
+    Value remaining = outer;
+    Value sourceElementOffset = zero;
+    for (size_t destinationDim = runStart; destinationDim-- > 0;) {
+      int64_t dimension = destinationShape[destinationDim];
+      if (dimension == 1)
+        continue;
+      Value dimensionValue = indexConstant(rewriter, loc, dimension);
+      Value coordinate = arith::RemUIOp::create(rewriter, loc, remaining, dimensionValue);
+      remaining = arith::DivUIOp::create(rewriter, loc, remaining, dimensionValue);
+      Value stride = indexConstant(rewriter, loc, sourceStrides[permutation[destinationDim]]);
+      Value contribution = arith::MulIOp::create(rewriter, loc, coordinate, stride);
+      sourceElementOffset = arith::AddIOp::create(rewriter, loc, sourceElementOffset, contribution);
+    }
+    Value bytes = indexConstant(rewriter, loc, elementBytes);
+    Value sourceOffset = arith::MulIOp::create(rewriter, loc, sourceElementOffset, bytes);
+    Value runByteCount = indexConstant(rewriter, loc, runLength * elementBytes);
+    Value targetOffset = arith::MulIOp::create(rewriter, loc, outer, runByteCount);
+    Value target = loop.getRegionIterArg(0);
+    Value output;
+    if (sourceStride == 1)
+      output = pim::PimMemCopyOp::create(
+        rewriter, loc, targetType, targetOffset, sourceOffset, target, op.getInput(), *runBytesAttr).getOutput();
+    else
+      output = pim::PimVMVOp::create(rewriter, loc, targetType, targetOffset, sourceOffset,
+                                    indexConstant(rewriter, loc, sourceStride), target, op.getInput(), *lengthAttr)
+                 .getOutput();
+    scf::YieldOp::create(rewriter, loc, output);
+    rewriter.replaceOp(op, loop.getResult(0));
+    return success();
+  }
+};
+
+struct InstructionSelectionPass : PassWrapper<InstructionSelectionPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InstructionSelectionPass)
+
+  StringRef getArgument() const override { return "pim-instruction-selection"; }
+  StringRef getDescription() const override { return "Select explicit PIM ISA operations"; }
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<LowerTransposePattern>(&getContext());
+    pim::populatePimContiguityNormalizationPatterns(patterns);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createPimInstructionSelectionPass() {
+  return std::make_unique<InstructionSelectionPass>();
+}
+
+} // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp

@@ -238,11 +238,11 @@ static bool isSupportedCoreInstructionOp(Operation* op) {
   return isa<pim::PimMemCopyHostToDevOp,
              pim::PimMemCopyDevToHostOp,
              pim::PimMemCopyOp,
+             pim::PimVMVOp,
              pim::PimReceiveOp,
              pim::PimSendOp,
              pim::PimConcatOp,
              pim::PimVMMOp,
-             pim::PimTransposeOp,
              pim::PimVVAddOp,
              pim::PimVVSubOp,
              pim::PimVVMulOp,

src/PIM/Dialect/Spatial/Spatial.td

@@ -288,6 +288,21 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> {
   let hasVerifier = 1;
 }
 
+def SpatSiluPlanOp : SpatOp<"silu_plan", []> {
+  let summary = "Layout-aware SiLU planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    StrAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
 def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> {
   let summary = "Layout-aware 2D NCHW MaxPool planning op";
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -472,6 +472,14 @@ LogicalResult SpatReluPlanOp::verify() {
   return success();
 }
 
+LogicalResult SpatSiluPlanOp::verify() {
+  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.silu_plan")))
+    return failure();
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
+  return success();
+}
+
 LogicalResult SpatMaxPool2DPlanOp::verify() {
   if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.max_pool2d_plan")))
     return failure();

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

@@ -265,6 +265,46 @@ static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices
   return receive.getOutput();
 }
 
+static FailureOr<SmallVector<LogicalTransferMetadataView, 0>>
+buildRectangularReceiveMetadata(const EmitReceiveAssemblyRun &run,
+                                unsigned laneCount) {
+  if (run.lanes.size() != laneCount
+      || !llvm::all_of(run.entryLanes,
+            [&](const LaneSet &lanes) { return lanes.size() == laneCount; }))
+    return failure();
+  SmallVector<LogicalTransferMetadataView, 0> result;
+  for (size_t entry = 0; entry < run.positions.size(); ++entry) {
+    ArrayRef<ScheduledTransferSlice> slices = ArrayRef(run.slices).slice(
+      run.entryOffsets[entry],
+      run.entryOffsets[entry + 1] - run.entryOffsets[entry]);
+    SmallVector<const ScheduledTransferSlice *> ordered;
+    for (const ScheduledTransferSlice &slice : slices)
+      ordered.push_back(&slice);
+    llvm::sort(ordered,
+      [](const ScheduledTransferSlice *left,
+         const ScheduledTransferSlice *right) {
+        LaneInterval leftFamily = left->family->targetLanes.intervals().front();
+        LaneInterval rightFamily = right->family->targetLanes.intervals().front();
+        return leftFamily.begin + left->familyOffset
+             < rightFamily.begin + right->familyOffset;
+      });
+    unsigned covered = 0;
+    LogicalTransferMetadataView metadata;
+    for (const ScheduledTransferSlice *slice : ordered) {
+      LaneInterval family = slice->family->targetLanes.intervals().front();
+      unsigned begin = family.begin + slice->familyOffset;
+      if (begin != covered)
+        return failure();
+      appendMetadata(*slice, metadata);
+      covered += slice->transferCount;
+    }
+    if (covered != laneCount)
+      return failure();
+    result.push_back(std::move(metadata));
+  }
+  return result;
+}
+
 template <typename Insert>
 static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, Value lane, unsigned laneCount, Value initial,
                                             DeferredEmissionContext &context, Insert insert) {
@@ -275,39 +315,11 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
   SmallVector<int64_t> counts(laneCount);
   std::optional<TransferGrids> transferGrids;
   std::optional<StaticIntGrid> positions;
-  SmallVector<LogicalTransferMetadataView, 0> metadataByEntry;
-  bool rectangular = run.lanes.size() == laneCount && llvm::all_of(run.entryLanes, [&](const LaneSet &lanes) { return lanes.size() == laneCount; });
-  for (size_t entry = 0; rectangular && entry < run.positions.size(); ++entry) {
-    ArrayRef<ScheduledTransferSlice> slices =
-        ArrayRef(run.slices).slice(run.entryOffsets[entry], run.entryOffsets[entry + 1] - run.entryOffsets[entry]);
-    SmallVector<const ScheduledTransferSlice *> ordered;
-    for (const ScheduledTransferSlice &slice : slices)
-      ordered.push_back(&slice);
-    llvm::sort(ordered, [](const ScheduledTransferSlice *left, const ScheduledTransferSlice *right) {
-      LaneInterval leftFamily = left->family->targetLanes.intervals().front();
-      LaneInterval rightFamily = right->family->targetLanes.intervals().front();
-      return leftFamily.begin + left->familyOffset < rightFamily.begin + right->familyOffset;
-    });
-    unsigned covered = 0;
-    LogicalTransferMetadataView metadata;
-    for (const ScheduledTransferSlice *slice : ordered) {
-      LaneInterval family = slice->family->targetLanes.intervals().front();
-      unsigned begin = family.begin + slice->familyOffset;
-      if (begin != covered) {
-        rectangular = false;
-        break;
-      }
-      appendMetadata(*slice, metadata);
-      covered += slice->transferCount;
-    }
-    rectangular &= covered == laneCount;
-    if (rectangular)
-      metadataByEntry.push_back(std::move(metadata));
-  }
-  if (rectangular) {
+  auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount);
+  if (succeeded(metadataByEntry)) {
     auto buildRows = [&](auto member) {
       SmallVector<StaticIntSequence> rows;
-      for (const LogicalTransferMetadataView &metadata : metadataByEntry)
+      for (const LogicalTransferMetadataView &metadata : *metadataByEntry)
         rows.push_back((metadata.*member).canonicalize());
       return StaticIntGrid::fromRows(rows);
     };
@@ -400,6 +412,62 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
   return loop->results.front();
 }
 
+static FailureOr<Value> transformProjectionFragment(
+    Value fragment, Value geometryRow, Value runtimeLane,
+    const DeferredProjectionLeafTemplate &leaf, const GridGeometry &geometry,
+    DeferredExchangePlan &exchange, DeferredEmissionContext &context) {
+  if (leaf.form != DeferredLeafForm::GraphBatchProjection)
+    return fragment;
+  Value shaped = extractMixedSliceOrIdentity(
+    context.rewriter, exchange.deferred.getLoc(), fragment,
+    getDeferredProjectedFragmentType(leaf),
+    lookupGeometry(geometry, geometryRow, runtimeLane, exchange.deferred,
+                   context, exchange.deferred.getLoc()));
+  return shaped ? FailureOr<Value>(shaped)
+                : FailureOr<Value>(failure());
+}
+
+static bool canStreamLoopedReceiveResult(
+    const EmitReceiveAssemblyRun &run, unsigned laneCount,
+    const DeferredResultPlan &plan, DeferredEmissionContext &context) {
+  if (!run.collection
+      || run.collection->key.kind != FragmentCollectionKind::Leaf
+      || run.collection->key.leafIndex != 0
+      || context.fragmentCollections.count(run.collection->key)
+      || run.lanes.size() != laneCount)
+    return false;
+  if (!canRealizeDeferredLoopResult(plan))
+    return false;
+  if (run.collection->positionCount
+      != plan.exchange->program.leaves.front().loopIterationCount)
+    return false;
+  return llvm::all_of(run.positions, [&](unsigned position) {
+    return position < run.collection->positionCount;
+  });
+}
+
+static FailureOr<Value> emitLoopedReceiveResult(
+    const EmitReceiveAssemblyRun &run, Value lane, unsigned laneCount,
+    const DeferredResultPlan &plan, DeferredEmissionContext &context) {
+  return realizeDeferredLoopResult(
+    plan, lane,
+    [&](Value initial,
+        DeferredLoopIterationEmitter emitIteration) -> FailureOr<Value> {
+      return emitReceiveAssembly(
+        run, lane, laneCount, initial, context,
+        [&](Value fragment, Value position, Value, Value runtimeLane,
+            Value current) -> FailureOr<Value> {
+          auto shaped = transformProjectionFragment(
+            fragment, position, runtimeLane, plan.exchange->program.leaves.front(),
+            plan.innerGeometry.front(), *plan.exchange, context);
+          return failed(shaped)
+                   ? FailureOr<Value>(failure())
+                   : emitIteration(*shaped, position, current);
+        });
+    },
+    context);
+}
+
 template <typename Emit>
 static LogicalResult emitCollectionUpdate(const LaneSet &lanes, Value lane, unsigned laneCount, FragmentCollectionKey key, Value current,
                                           Operation *anchor, DeferredEmissionContext &context, Emit emit, bool local = false) {
@@ -424,26 +492,59 @@ static LogicalResult emitCollectionUpdate(const LaneSet &lanes, Value lane, unsi
   return success();
 }
 
-static FailureOr<Value> insertProjectionFragment(Value fragment, Value specialization, Value position, Value geometryRow, Value runtimeLane,
+static FailureOr<Value> insertProjectionFragment(Value fragment,
+                                                 Value specialization,
+                                                 Value position,
+                                                 unsigned fragmentPositionCount,
+                                                 Value runtimeLane,
                                                  Value assembled, const DeferredProjectionLeafTemplate &leaf, const GridGeometry &geometry,
                                                  DeferredExchangePlan &exchange, bool grouped, DeferredEmissionContext &context) {
-  Value shaped = fragment;
-  if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
-    RankedTensorType projectedType = getDeferredProjectedFragmentType(leaf);
-    shaped = extractMixedSliceOrIdentity(context.rewriter, exchange.deferred.getLoc(), shaped,
-                                         projectedType,
-                                         lookupGeometry(geometry, geometryRow, runtimeLane, exchange.deferred, context, exchange.deferred.getLoc()));
-    if (!shaped) return failure();
+  Value iteration = position;
+  Value fragmentPosition = position;
+  if (leaf.enclosingLoop && fragmentPositionCount > 1) {
+    Value divisor = context.constants.getIndex(fragmentPositionCount);
+    iteration = arith::DivUIOp::create(
+      context.rewriter, exchange.deferred.getLoc(), position, divisor);
+    fragmentPosition = arith::RemUIOp::create(
+      context.rewriter, exchange.deferred.getLoc(), position, divisor);
   }
+  Value geometryRow = leaf.enclosingLoop
+                    ? iteration
+                    : context.constants.getIndex(0);
+  if (grouped) {
+    Value base = arith::MulIOp::create(
+      context.rewriter, exchange.deferred.getLoc(), specialization,
+      context.constants.getIndex(leaf.loopIterationCount));
+    geometryRow = arith::AddIOp::create(
+      context.rewriter, exchange.deferred.getLoc(), base, geometryRow);
+  }
+  auto transformed = transformProjectionFragment(
+    fragment, geometryRow, runtimeLane, leaf, geometry, exchange, context);
+  if (failed(transformed))
+    return failure();
+  Value shaped = *transformed;
   auto sourceType = dyn_cast<RankedTensorType>(shaped.getType());
   auto assembledType = dyn_cast<RankedTensorType>(assembled.getType());
   int64_t rankDifference = sourceType && assembledType ? assembledType.getRank() - sourceType.getRank() : 0;
-  if (rankDifference < 0 || rankDifference > 2 || (grouped && rankDifference == 0)) return failure();
-  if (rankDifference == 0 && sourceType != assembledType) return failure();
+  bool leading = sourceType && leaf.reconstructedType
+              && leaf.reconstructedType.getRank() == sourceType.getRank() + 1
+              && leaf.reconstructedType.getShape().drop_front()
+                   == sourceType.getShape();
+  int64_t expectedRankDifference = static_cast<int64_t>(grouped)
+                                 + static_cast<int64_t>(!!leaf.enclosingLoop)
+                                 + static_cast<int64_t>(leading);
+  if (!sourceType || !assembledType
+      || rankDifference != expectedRankDifference)
+    return failure();
   MixedSliceGeometry slice;
   slice.offsets.assign(assembledType.getRank(), context.rewriter.getIndexAttr(0));
-  if (rankDifference) slice.offsets.front() = grouped ? specialization : position;
-  if (rankDifference == 2) slice.offsets[1] = position;
+  unsigned prefix = 0;
+  if (grouped)
+    slice.offsets[prefix++] = specialization;
+  if (leaf.enclosingLoop)
+    slice.offsets[prefix++] = iteration;
+  if (leading)
+    slice.offsets[prefix++] = fragmentPosition;
   slice.sizes.assign(rankDifference, context.rewriter.getIndexAttr(1));
   for (int64_t dimension : sourceType.getShape()) slice.sizes.push_back(context.rewriter.getIndexAttr(dimension));
   slice.strides.assign(assembledType.getRank(), context.rewriter.getIndexAttr(1));
@@ -488,7 +589,10 @@ static LogicalResult emitLeafCollectionUpdate(const EmitReceiveAssemblyRun &run,
             specialization = arith::DivUIOp::create(context.rewriter, exchange.deferred.getLoc(), position, divisor);
             leafPosition = arith::RemUIOp::create(context.rewriter, exchange.deferred.getLoc(), position, divisor);
           }
-          return insertProjectionFragment(fragment, specialization, leafPosition, grouped ? specialization : context.constants.getIndex(0), runtimeLane,
+          unsigned fragmentPositionCount =
+            collection.positionCount / leaf.loopIterationCount;
+          return insertProjectionFragment(fragment, specialization, leafPosition,
+                                          fragmentPositionCount, runtimeLane,
                                           assembled, leaf, geometry, exchange, grouped, context);
         });
   };
@@ -642,10 +746,11 @@ static LogicalResult emitLoopedLocalCollectionUpdate(
         unsigned leafIndex = collection.key.leafIndex;
         const DeferredProjectionLeafTemplate &leaf =
           exchange.program.leaves[leafIndex];
+        unsigned fragmentPositionCount =
+          collection.positionCount / leaf.loopIterationCount;
         auto inserted = insertProjectionFragment(
           *source, specialization, leafPosition,
-          grouped ? specialization : context.constants.getIndex(0),
-          runtimeLane, iterArgs.front(), leaf,
+          fragmentPositionCount, runtimeLane, iterArgs.front(), leaf,
           resultPlan.innerGeometry[leafIndex], exchange, grouped, context);
         if (failed(inserted))
           return failure();
@@ -818,8 +923,10 @@ static LogicalResult emitLocalCollectionUpdate(const EmitLocalCollectionRun &upd
   auto emit = [&](Value assembled) -> FailureOr<Value> {
     auto fragment = materialize();
     if (failed(fragment)) return failure();
+    unsigned fragmentPositionCount =
+      collection.positionCount / leaf.loopIterationCount;
     return insertProjectionFragment(*fragment, context.constants.getIndex(specialization),
-                                    context.constants.getIndex(leafPosition), context.constants.getIndex(specialization),
+                                    context.constants.getIndex(leafPosition), fragmentPositionCount,
                                     lane ? lane : context.constants.getIndex(0), assembled, leaf, resultPlan.innerGeometry[leafIndex], exchange, grouped, context);
   };
   return emitCollectionUpdate(update.lanes, lane, laneCount, collection.key, current, exchange.deferred, context, emit, true);
@@ -828,7 +935,29 @@ static LogicalResult emitLocalCollectionUpdate(const EmitLocalCollectionRun &upd
 static FailureOr<SmallVector<Value>> emitInstructions(ArrayRef<BoundaryInstruction> instructions, Value lane, unsigned laneCount,
                                                       ArrayRef<DeferredResultPlan> results, DeferredEmissionContext &context) {
   SmallVector<Value> produced;
-  for (const BoundaryInstruction &instruction : instructions) {
+  for (size_t instructionIndex = 0;
+       instructionIndex < instructions.size(); ++instructionIndex) {
+    const BoundaryInstruction &instruction = instructions[instructionIndex];
+    if (auto assembly = std::get_if<EmitReceiveAssemblyRun>(&instruction)) {
+      if (instructionIndex + 1 < instructions.size())
+        if (auto result = std::get_if<ProduceDeferredResult>(
+              &instructions[instructionIndex + 1])) {
+          const DeferredResultPlan *plan = findResultPlan(results, result->exchange);
+          if (plan && assembly->collection
+              && assembly->collection->key.exchange == result->exchange
+              && canStreamLoopedReceiveResult(
+                   *assembly, laneCount, *plan, context)) {
+            auto value = emitLoopedReceiveResult(
+              *assembly, lane, laneCount, *plan, context);
+            if (failed(value))
+              return result->exchange->deferred.emitOpError(
+                "failed to stream loop-indexed received result"), failure();
+            produced.push_back(*value);
+            ++instructionIndex;
+            continue;
+          }
+        }
+    }
     if (auto send = std::get_if<EmitSendRun>(&instruction)) {
       if (failed(emitConditionalSendRun(*send, lane, laneCount, context)))
         return failure();

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

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/Operation.h"
 
 #include "llvm/ADT/ArrayRef.h"
@@ -146,6 +147,8 @@ struct DeferredProjectionLeafTemplate {
   DeferredSliceTemplate innerGeometry;
   mlir::RankedTensorType reconstructedType;
   bool leadingRankReduced = false;
+  mlir::scf::ForOp enclosingLoop;
+  unsigned loopIterationCount = 1;
 };
 
 struct DeferredInsertAssemblyEntryTemplate {

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

@@ -171,6 +171,13 @@ static bool isTopLevelDeferredOperation(Operation *op, Block &body,
       && (isDeferredPayloadCandidateOp(op) || isa<scf::ForOp>(op));
 }
 
+static Operation *getTopLevelDeferredOperation(
+    Operation *op, Block &body, const DeferredInputPlan &plan) {
+  while (op && op->getBlock() != &body)
+    op = op->getParentOp();
+  return op && isTopLevelDeferredOperation(op, body, plan) ? op : nullptr;
+}
+
 static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
                        llvm::SmallPtrSetImpl<Operation *> &seen) {
   if (value == plan.graphInput || value == plan.graphLane || value == plan.scheduledLane)
@@ -183,6 +190,27 @@ static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
     return true;
   if (!op || !isTopLevelDeferredOperation(op, body, plan) || !seen.insert(op).second)
     return op && seen.contains(op);
+  if (auto loop = dyn_cast<scf::ForOp>(op)) {
+    bool eligible = true;
+    loop.getRegion().walk([&](Operation *nested) {
+      if (isa<scf::ForOp>(nested) && nested != loop)
+        eligible = false;
+      for (Value operand : nested->getOperands()) {
+        Operation *definition = operand.getDefiningOp();
+        auto argument = dyn_cast<BlockArgument>(operand);
+        Region *argumentRegion = argument
+          ? argument.getOwner()->getParent() : nullptr;
+        bool definedInside = definition
+          ? loop->isProperAncestor(definition)
+          : argumentRegion == &loop.getRegion()
+            || (argumentRegion && loop.getRegion().isAncestor(argumentRegion));
+        if (!definedInside && !isEligible(operand, body, plan, seen))
+          eligible = false;
+      }
+    });
+    if (!eligible)
+      return false;
+  }
   return llvm::all_of(op->getOperands(), [&](Value operand) { return isEligible(operand, body, plan, seen); });
 }
 
@@ -198,8 +226,16 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
         argument && argument.getOwner() == &transfer.getBody().front()
         && argument.getArgNumber() >= transfer.getSources().size())
       return value;
-    if (isa<BlockArgument>(value))
-      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
+    if (auto argument = dyn_cast<BlockArgument>(value)) {
+      Operation *owner = argument.getOwner()->getParentOp();
+      if (isa_and_nonnull<scf::ForOp>(owner)
+          && owner->isProperAncestor(transfer))
+        return value;
+      return transfer.emitOpError(
+        "phase 1 scheduled graph-lane expression captures an unsupported block argument: ")
+        << value << " owned by " << owner->getName(),
+        failure();
+    }
     Operation *op = value.getDefiningOp();
     if (!op || (!isDeferredPayloadCandidateOp(op) && !op->hasTrait<OpTrait::ConstantLike>()))
       return transfer.emitOpError("phase 1 cannot clone the scheduled graph-lane expression"), failure();
@@ -216,11 +252,39 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
       mapping.map(value, *mappedLane);
       return *mappedLane;
     }
-    if (isa<BlockArgument>(value))
-      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
+    if (auto argument = dyn_cast<BlockArgument>(value)) {
+      Operation *owner = argument.getOwner()->getParentOp();
+      if (isa_and_nonnull<scf::ForOp>(owner)
+          && owner->isProperAncestor(transfer))
+        return value;
+      return transfer.emitOpError(
+        "phase 1 payload shaping captures an unsupported block argument: ")
+        << value << " owned by " << owner->getName(), failure();
+    }
     Operation *op = value.getDefiningOp();
     if (!op || (!isTopLevelDeferredOperation(op, body, plan) && !op->hasTrait<OpTrait::ConstantLike>()))
       return transfer.emitOpError("phase 1 payload shaping contains an unsupported operation"), failure();
+    if (auto loop = dyn_cast<scf::ForOp>(op)) {
+      SmallVector<Value> captures;
+      loop.getRegion().walk([&](Operation *nested) {
+        for (Value operand : nested->getOperands()) {
+          Operation *definition = operand.getDefiningOp();
+          auto argument = dyn_cast<BlockArgument>(operand);
+          Region *argumentRegion = argument
+            ? argument.getOwner()->getParent() : nullptr;
+          bool definedInside = definition
+            ? loop->isProperAncestor(definition)
+            : argumentRegion == &loop.getRegion()
+              || (argumentRegion
+                  && loop.getRegion().isAncestor(argumentRegion));
+          if (!definedInside && !mapping.contains(operand))
+            captures.push_back(operand);
+        }
+      });
+      for (Value capture : captures)
+        if (!mapping.contains(capture) && failed(clone(capture)))
+          return failure();
+    }
     for (Value operand : op->getOperands()) if (failed(clone(operand))) return failure();
     Operation *copy = builder.clone(*op, mapping);
     for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair));
@@ -352,10 +416,19 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc
       Value value = worklist.pop_back_val();
       if (!seen.insert(value).second) continue;
       for (OpOperand &use : value.getUses()) {
-        Operation *user = use.getOwner();
-        if (!isTopLevelDeferredOperation(user, body, plan)) { needsIdentity = true; continue; }
+        Operation *user = getTopLevelDeferredOperation(
+          use.getOwner(), body, plan);
+        if (!user) {
+          if (value == plan.graphInput) needsIdentity = true;
+          else roots.push_back(value);
+          continue;
+        }
         llvm::SmallPtrSet<Operation *, 16> eligibility;
-        if (!isEligible(user->getResult(0), body, plan, eligibility)) { needsIdentity = true; continue; }
+        if (!isEligible(user->getResult(0), body, plan, eligibility)) {
+          if (value == plan.graphInput) needsIdentity = true;
+          else roots.push_back(value);
+          continue;
+        }
         for (Value result : user->getResults()) {
           bool hasShapingUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return isTopLevelDeferredOperation(next.getOwner(), body, plan); });
           bool hasOtherUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return !isTopLevelDeferredOperation(next.getOwner(), body, plan); });

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

@@ -8,6 +8,8 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 
+#include <limits>
+
 #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 
@@ -187,6 +189,35 @@ static bool isInsideDeferredLoop(Operation *op,
   return false;
 }
 
+static FailureOr<unsigned> getLoopIterationCount(
+    scf::ForOp loop, Value scheduledLane, Value specialization,
+    int64_t laneCount, int64_t specializationCount) {
+  std::optional<unsigned> commonCount;
+  for (int64_t specializationIndex = 0;
+       specializationIndex < specializationCount; ++specializationIndex)
+    for (int64_t lane = 0; lane < laneCount; ++lane) {
+      StaticIndexEnvironment environment;
+      if (scheduledLane)
+        environment.bindings[scheduledLane] = lane;
+      if (specialization)
+        environment.bindings[specialization] = specializationIndex;
+      auto lower = evaluateDeferredIndex(loop.getLowerBound(), environment);
+      auto upper = evaluateDeferredIndex(loop.getUpperBound(), environment);
+      auto step = evaluateDeferredIndex(loop.getStep(), environment);
+      if (failed(lower) || failed(upper) || failed(step) || *step <= 0)
+        return failure();
+      int64_t distance = std::max<int64_t>(0, *upper - *lower);
+      uint64_t count = (static_cast<uint64_t>(distance)
+                        + static_cast<uint64_t>(*step) - 1)
+                     / static_cast<uint64_t>(*step);
+      if (count > std::numeric_limits<unsigned>::max()
+          || (commonCount && *commonCount != count))
+        return failure();
+      commonCount = count;
+    }
+  return commonCount.value_or(0);
+}
+
 static FailureOr<SmallVector<unsigned>>
 getPossibleDeferredSourceOperandIndices(
     Value sourceRoot, SpatDeferredCommunicationOp deferred) {
@@ -220,7 +251,7 @@ static LogicalResult validateDeferredProgram(
       op->emitOpError(message);
       return WalkResult::interrupt();
     };
-    if (isa<SpatYieldOp, scf::YieldOp>(op)
+    if (isa<SpatYieldOp, scf::YieldOp, tensor::YieldOp>(op)
         || (isa<linalg::YieldOp>(op)
             && isa<linalg::TransposeOp>(op->getParentOp())))
       return WalkResult::advance();
@@ -250,15 +281,39 @@ static LogicalResult validateDeferredProgram(
         return reject("has a non-positive or non-static deferred loop step");
       return WalkResult::advance();
     }
-    if (op->getNumRegions() != 0 && !isa<linalg::TransposeOp>(op))
+    bool structuredLoopShaping = isInsideDeferredLoop(op, deferred)
+                              && (isShapingOnlyOp(op)
+                                  || isa<tensor::PadOp>(op));
+    if (op->getNumRegions() != 0 && !isa<linalg::TransposeOp>(op)
+        && !structuredLoopShaping)
       return reject("contains an unsupported region operation");
+    bool structuredLoopResidual = isInsideDeferredLoop(op, deferred)
+                               && (op->getNumRegions() == 0
+                                   || structuredLoopShaping);
     if (!isShapingOnlyOp(op) && !isCompileTimeOp(op)
-        && !isPureIndexComputationOp(op))
+        && !isPureIndexComputationOp(op) && !structuredLoopResidual)
       return reject("contains an unsupported deferred operation");
-    for (Value operand : op->getOperands())
-      if (isInsideDeferredLoop(op, deferred)
-          && originatesFromDeferredSource(operand, deferred))
-        return reject("projects a deferred source inside a residual loop");
+    if (isInsideDeferredLoop(op, deferred)
+        && llvm::any_of(op->getOperands(), [&](Value operand) {
+             return originatesFromDeferredSource(operand, deferred);
+           })) {
+      auto loop = op->getParentOfType<scf::ForOp>();
+      scf::ForOp outerLoop;
+      for (Operation *parent = loop ? loop->getParentOp() : nullptr;
+           parent && parent != deferred; parent = parent->getParentOp())
+        if (auto candidate = dyn_cast<scf::ForOp>(parent))
+          outerLoop = candidate;
+      if (!loop || outerLoop) {
+        op->emitOpError("projects a deferred source inside an unsupported loop nest")
+          << " (inner " << (loop ? loop.getLoc() : op->getLoc())
+          << ", outer " << (outerLoop ? outerLoop.getLoc() : op->getLoc())
+          << ")";
+        return WalkResult::interrupt();
+      }
+      if (getConstantIntValue(loop.getLowerBound()) != 0
+          || getConstantIntValue(loop.getStep()) != 1)
+        return reject("projects a deferred source inside a non-normalized loop");
+    }
     return WalkResult::advance();
   });
   return failure(result.wasInterrupted());
@@ -376,9 +431,10 @@ FailureOr<int64_t> evaluateDeferredIndex(
 
 DeferredLaneValueEvaluator::DeferredLaneValueEvaluator(
     const DeferredProgramTemplate &program, unsigned laneCount,
-    unsigned specializationIndex)
+    unsigned specializationIndex, scf::ForOp loop, unsigned loopIteration)
   : program(program), laneCount(laneCount),
-    specializationIndex(specializationIndex) {}
+    specializationIndex(specializationIndex), loop(loop),
+    loopIteration(loopIteration) {}
 
 FailureOr<StaticIntSequence> DeferredLaneValueEvaluator::evaluate(Value value) {
   if (auto it = values.find(value); it != values.end())
@@ -397,6 +453,8 @@ FailureOr<StaticIntSequence> DeferredLaneValueEvaluator::evaluate(Value value) {
     if (program.specializationArgument)
       environment.bindings[program.specializationArgument] =
         specializationIndex;
+    if (loop)
+      environment.bindings[loop.getInductionVar()] = loopIteration;
     auto result = evaluateDeferredIndex(value, environment);
     if (failed(result))
       return failure();
@@ -435,6 +493,8 @@ DeferredLaneValueEvaluator::resolveSourceOperandIndices(Value sourceRoot) {
     if (program.specializationArgument)
       environment.bindings[program.specializationArgument] =
         specializationIndex;
+    if (loop)
+      environment.bindings[loop.getInductionVar()] = loopIteration;
     auto index = sourceArgument(sourceRoot, program.deferred, environment);
     if (failed(index) || !*index)
       return failure();
@@ -518,6 +578,7 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
             SmallVector<OpFoldResult>(
               ArrayRef(slice.getMixedStrides()).drop_front())};
         leaf.reconstructedType = cast<RankedTensorType>(value.getType());
+        leaf.enclosingLoop = slice->getParentOfType<scf::ForOp>();
         if (graphProjection
             && slice.getSourceType().getRank()
                  == leaf.reconstructedType.getRank() + 1
@@ -540,30 +601,59 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
       if (!type)
         return deferred.emitOpError(
           "deferred source is not a ranked tensor");
-      program.leaves.push_back({DeferredLeafForm::DirectSource, value, value,
-                                {}, {}, type});
+      DeferredProjectionLeafTemplate leaf;
+      leaf.form = DeferredLeafForm::DirectSource;
+      leaf.sourceRoot = value;
+      leaf.replacementRoot = value;
+      leaf.reconstructedType = type;
+      program.leaves.push_back(std::move(leaf));
       return success();
     }
     if (value.getType().isIndex() || isa<IntegerType>(value.getType()))
       return success();
+    if (auto argument = dyn_cast<BlockArgument>(value)) {
+      auto loop = dyn_cast_or_null<scf::ForOp>(
+        argument.getOwner()->getParentOp());
+      if (loop && argument.getArgNumber() > 0)
+        return visit(loop.getInitArgs()[argument.getArgNumber() - 1]);
+    }
     Operation *op = value.getDefiningOp();
-    if (!op || (op->getBlock() != &body && !isa<scf::ForOp>(op)))
+    if (!op || (op->getBlock() != &body
+                && !op->getParentOfType<scf::ForOp>()))
       return deferred.emitOpError(
-        "deferred residual escapes its verified body");
+        "deferred residual escapes its verified body: ") << value;
     if (auto loop = dyn_cast<scf::ForOp>(op)) {
       for (Value init : loop.getInitArgs())
         if (failed(visit(init)))
           return failure();
+      auto loopYield = cast<scf::YieldOp>(loop.getBody()->getTerminator());
+      for (Value yielded : loopYield.getResults())
+        if (failed(visit(yielded)))
+          return failure();
     } else {
       for (Value operand : op->getOperands())
         if (failed(visit(operand)))
           return failure();
     }
-    program.residualOps.push_back(op);
+    if (op->getBlock() == &body)
+      program.residualOps.push_back(op);
     return success();
   };
   if (failed(visit(program.yieldedValue)))
     return failure();
+  for (DeferredProjectionLeafTemplate &leaf : program.leaves) {
+    if (!leaf.enclosingLoop)
+      continue;
+    auto count = getLoopIterationCount(
+      leaf.enclosingLoop, program.scheduledLane,
+      program.specializationArgument, laneCount,
+      program.specializationCount);
+    if (failed(count) || *count == 0)
+      return deferred.emitOpError(
+        "deferred source projection loop must have a common non-zero trip count"),
+        failure();
+    leaf.loopIterationCount = *count;
+  }
   auto assembly = analyzeInsertAssembly(program);
   if (failed(assembly))
     return failure();

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

@@ -25,7 +25,9 @@ class DeferredLaneValueEvaluator {
 public:
   DeferredLaneValueEvaluator(const DeferredProgramTemplate &program,
                              unsigned laneCount,
-                             unsigned specializationIndex = 0);
+                             unsigned specializationIndex = 0,
+                             mlir::scf::ForOp loop = {},
+                             unsigned loopIteration = 0);
 
   mlir::FailureOr<StaticIntSequence> evaluate(mlir::Value value);
   mlir::FailureOr<StaticIntSequence> evaluate(mlir::OpFoldResult value);
@@ -36,6 +38,8 @@ private:
   const DeferredProgramTemplate &program;
   unsigned laneCount;
   unsigned specializationIndex;
+  mlir::scf::ForOp loop;
+  unsigned loopIteration;
   llvm::DenseMap<mlir::Value, StaticIntSequence> values;
   llvm::DenseMap<mlir::Value, StaticIntSequence> sourceOperands;
 };

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

@@ -73,15 +73,19 @@ static LogicalResult buildLeafCollections(DeferredExchangePlan &exchange,
     RankedTensorType normalized = fragmentType;
     if (leaf.form == DeferredLeafForm::GraphBatchProjection)
       normalized = getDeferredProjectedFragmentType(leaf);
-    bool direct = positionCount == 1 && normalized == leaf.reconstructedType;
+    if (positionCount % leaf.loopIterationCount != 0)
+      return failure();
+    unsigned fragmentPositionCount = positionCount / leaf.loopIterationCount;
+    bool direct = fragmentPositionCount == 1
+               && normalized == leaf.reconstructedType;
     bool leading = leaf.reconstructedType.getRank() == normalized.getRank() + 1
-                && leaf.reconstructedType.getDimSize(0) == positionCount
+                && leaf.reconstructedType.getDimSize(0) == fragmentPositionCount
                 && leaf.reconstructedType.getShape().drop_front()
                      == normalized.getShape();
     if (!direct && !leading)
       return exchange.deferred.emitOpError(
                "cannot form fragment collection for leaf ")
-             << leafIndex << ": " << positionCount << " fragments of "
+             << leafIndex << ": " << fragmentPositionCount << " fragments of "
              << normalized << " do not reconstruct " << leaf.reconstructedType;
     FragmentCollectionKind kind = specializationCount == 1
                                 ? FragmentCollectionKind::Leaf
@@ -89,6 +93,8 @@ static LogicalResult buildLeafCollections(DeferredExchangePlan &exchange,
     SmallVector<int64_t> shape;
     if (specializationCount > 1)
       shape.push_back(specializationCount);
+    if (leaf.enclosingLoop)
+      shape.push_back(leaf.loopIterationCount);
     llvm::append_range(shape, leaf.reconstructedType.getShape());
     FragmentCollectionPlan collection;
     collection.key = {&exchange, kind, static_cast<unsigned>(leafIndex)};
@@ -196,11 +202,14 @@ static constexpr std::array<TemplateGeometryMember, 3> templateGeometryMembers{
 template <typename GetValue>
 static FailureOr<StaticIntGrid> buildGrid(
     const DeferredProgramTemplate &program, unsigned laneCount,
-    unsigned rowCount, bool specializeRows, GetValue getValue) {
+    unsigned rowCount, bool specializeRows, scf::ForOp loop,
+    unsigned loopIterationCount, GetValue getValue) {
   SmallVector<StaticIntSequence> rows;
   for (unsigned row = 0; row < rowCount; ++row) {
+    unsigned specialization = specializeRows ? row / loopIterationCount : 0;
+    unsigned iteration = loop ? row % loopIterationCount : 0;
     DeferredLaneValueEvaluator evaluator(
-      program, laneCount, specializeRows ? row : 0);
+      program, laneCount, specialization, loop, iteration);
     auto sequence = evaluator.evaluate(getValue(row));
     if (failed(sequence)) return failure();
     rows.push_back(std::move(*sequence));
@@ -211,7 +220,8 @@ static FailureOr<StaticIntGrid> buildGrid(
 template <typename GetGeometry>
 static FailureOr<DeferredGridSliceGeometry> buildGeometryGrids(
     const DeferredProgramTemplate &program, unsigned laneCount,
-    unsigned rowCount, bool specializeRows, GetGeometry getGeometry) {
+    unsigned rowCount, bool specializeRows, scf::ForOp loop,
+    unsigned loopIterationCount, GetGeometry getGeometry) {
   DeferredGridSliceGeometry result;
   const DeferredSliceTemplate &first = getGeometry(0);
   for (auto [group, sourceMember] : llvm::enumerate(templateGeometryMembers)) {
@@ -219,6 +229,7 @@ static FailureOr<DeferredGridSliceGeometry> buildGeometryGrids(
     for (unsigned dimension = 0;
          dimension < (first.*member).size(); ++dimension) {
       auto grid = buildGrid(program, laneCount, rowCount, specializeRows,
+        loop, loopIterationCount,
         [&](unsigned row) { return (getGeometry(row).*member)[dimension]; });
       if (failed(grid)) return failure();
       result[group].push_back(std::move(*grid));
@@ -226,8 +237,92 @@ static FailureOr<DeferredGridSliceGeometry> buildGeometryGrids(
   }
   return result;
 }
-static Value cloneResidual(DeferredExchangePlan &exchange, IRMapping &mapping, DeferredEmissionContext &context) {
+
+static LogicalResult cloneLoopIteration(
+    DeferredExchangePlan &exchange, scf::ForOp oldLoop,
+    const IRMapping &mapping, Value induction, ValueRange iterArgs,
+    llvm::function_ref<FailureOr<Value>(unsigned)> getLeaf,
+    DeferredEmissionContext &context, SmallVectorImpl<Value> &yielded) {
+  IRMapping nested(mapping);
+  nested.map(oldLoop.getInductionVar(), induction);
+  for (auto [oldArg, newArg] :
+       llvm::zip(oldLoop.getRegionIterArgs(), iterArgs))
+    nested.map(oldArg, newArg);
+  for (Operation &nestedOp : oldLoop.getBody()->without_terminator()) {
+    bool replaced = false;
+    for (auto [leafIndex, leaf] : llvm::enumerate(exchange.program.leaves)) {
+      if (leaf.enclosingLoop != oldLoop
+          || leaf.replacementRoot.getDefiningOp() != &nestedOp)
+        continue;
+      auto selected = getLeaf(leafIndex);
+      if (failed(selected))
+        return failure();
+      nested.map(leaf.replacementRoot, *selected);
+      replaced = true;
+    }
+    if (replaced)
+      continue;
+    Operation *copy = context.rewriter.clone(nestedOp, nested);
+    for (auto [oldValue, newValue] :
+         llvm::zip(nestedOp.getResults(), copy->getResults()))
+      nested.map(oldValue, newValue);
+  }
+  auto oldYield = cast<scf::YieldOp>(oldLoop.getBody()->getTerminator());
+  for (Value value : oldYield.getResults())
+    yielded.push_back(nested.lookupOrDefault(value));
+  return success();
+}
+
+static Value cloneResidual(
+    DeferredExchangePlan &exchange, IRMapping &mapping,
+    ArrayRef<Value> leafCollections,
+    DeferredEmissionContext &context) {
+  for (Operation &op :
+       exchange.deferred.getBody().front().without_terminator()) {
+    if (!op.hasTrait<OpTrait::ConstantLike>()
+        || llvm::all_of(op.getResults(),
+             [&](Value result) { return mapping.contains(result); }))
+      continue;
+    Operation *copy = context.rewriter.clone(op, mapping);
+    for (auto [oldValue, newValue] :
+         llvm::zip(op.getResults(), copy->getResults()))
+      mapping.map(oldValue, newValue);
+  }
   for (Operation *op : exchange.program.residualOps) {
+    if (op->hasTrait<OpTrait::ConstantLike>())
+      continue;
+    if (auto oldLoop = dyn_cast<scf::ForOp>(op)) {
+      SmallVector<Value> initArgs;
+      for (Value init : oldLoop.getInitArgs())
+        initArgs.push_back(mapping.lookupOrDefault(init));
+      auto loop = buildNormalizedScfFor(
+        context.rewriter, oldLoop.getLoc(),
+        mapping.lookupOrDefault(oldLoop.getLowerBound()),
+        mapping.lookupOrDefault(oldLoop.getUpperBound()),
+        mapping.lookupOrDefault(oldLoop.getStep()), initArgs,
+        [&](OpBuilder &, Location loc, Value induction, ValueRange iterArgs,
+            SmallVectorImpl<Value> &yielded) -> LogicalResult {
+          return cloneLoopIteration(
+            exchange, oldLoop, mapping, induction, iterArgs,
+            [&](unsigned leafIndex) -> FailureOr<Value> {
+              const auto &leaf = exchange.program.leaves[leafIndex];
+              Value selected = extractMixedSliceOrIdentity(
+                context.rewriter, loc, leafCollections[leafIndex],
+                leaf.reconstructedType,
+                leadingSlice(context.rewriter, leaf.reconstructedType,
+                             induction));
+              return selected ? FailureOr<Value>(selected)
+                              : FailureOr<Value>(failure());
+            },
+            context, yielded);
+        });
+      if (failed(loop))
+        return {};
+      for (auto [oldValue, newValue] :
+           llvm::zip(oldLoop.getResults(), loop->results))
+        mapping.map(oldValue, newValue);
+      continue;
+    }
     Operation *copy = context.rewriter.clone(*op, mapping);
     for (auto [oldValue, newValue] : llvm::zip(op->getResults(), copy->getResults())) mapping.map(oldValue, newValue);
   }
@@ -252,16 +347,26 @@ static FailureOr<Value> realizeOne(const DeferredResultPlan &plan, Value lane, D
     if (!value.getType().isIndex()) selected = arith::IndexCastOp::create(context.rewriter, exchange.deferred.getLoc(), value.getType(), selected);
     mapping.map(value, selected);
   }
+  SmallVector<Value> leafCollections;
   for (unsigned index = 0; index < exchange.program.leaves.size(); ++index) {
     Value value = context.fragmentCollections.lookup(
       {&exchange, FragmentCollectionKind::Leaf, index});
-    if (!value || value.getType() != exchange.program.leaves[index].reconstructedType) {
+    const auto &leaf = exchange.program.leaves[index];
+    SmallVector<int64_t> expectedShape;
+    if (leaf.enclosingLoop)
+      expectedShape.push_back(leaf.loopIterationCount);
+    llvm::append_range(expectedShape, leaf.reconstructedType.getShape());
+    Type expectedType = RankedTensorType::get(
+      expectedShape, leaf.reconstructedType.getElementType());
+    if (!value || value.getType() != expectedType) {
       exchange.deferred.emitOpError("failed to reconstruct deferred result leaf ") << index;
       return failure();
     }
-    mapping.map(exchange.program.leaves[index].replacementRoot, value);
+    leafCollections.push_back(value);
+    if (!exchange.program.leaves[index].enclosingLoop)
+      mapping.map(exchange.program.leaves[index].replacementRoot, value);
   }
-  Value result = cloneResidual(exchange, mapping, context);
+  Value result = cloneResidual(exchange, mapping, leafCollections, context);
   Type expected = exchange.program.specializationCount > 1 ? Type(exchange.program.specializationFragmentType) : exchange.deferred.getOutput().getType();
   return result && result.getType() == expected ? FailureOr<Value>(result) : FailureOr<Value>(failure());
 }
@@ -274,8 +379,10 @@ FailureOr<DeferredResultPlan> buildDeferredResultPlan(DeferredExchangePlan &exch
       : buildLeafCollections(exchange, plan))) return failure();
   unsigned specializations = exchange.program.specializationCount;
   for (const auto &leaf : exchange.program.leaves) {
+    unsigned geometryRows = specializations * leaf.loopIterationCount;
     auto geometry = buildGeometryGrids(exchange.program,
-      exchange.targetLaneCount, specializations, true,
+      exchange.targetLaneCount, geometryRows, true, leaf.enclosingLoop,
+      leaf.loopIterationCount,
       [&](unsigned) -> const DeferredSliceTemplate & {
         return leaf.innerGeometry;
       });
@@ -285,7 +392,7 @@ FailureOr<DeferredResultPlan> buildDeferredResultPlan(DeferredExchangePlan &exch
   if (exchange.program.insertAssembly) {
     const auto &entries = exchange.program.insertAssembly->entries;
     auto geometry = buildGeometryGrids(exchange.program,
-      exchange.targetLaneCount, entries.size(), false,
+      exchange.targetLaneCount, entries.size(), false, {}, 1,
       [&](unsigned row) -> const DeferredSliceTemplate & {
         return entries[row].targetGeometry;
       });
@@ -301,7 +408,7 @@ FailureOr<DeferredResultPlan> buildDeferredResultPlan(DeferredExchangePlan &exch
     });
   for (Value value : residualValues) {
     auto grid = buildGrid(exchange.program, exchange.targetLaneCount,
-      specializations, true,
+      specializations, true, {}, 1,
       [&](unsigned) { return OpFoldResult(value); });
     if (succeeded(grid))
       plan.residualValues.try_emplace(value, std::move(*grid));
@@ -320,6 +427,8 @@ FailureOr<Value> realizeDeferredResult(const DeferredResultPlan &plan, Value lan
       FragmentCollectionKey key{&exchange, FragmentCollectionKind::GroupedLeaf, static_cast<unsigned>(index)};
       Value collection = context.fragmentCollections.lookup(key);
       SmallVector<int64_t> shape{exchange.program.specializationCount};
+      if (leaf.enclosingLoop)
+        shape.push_back(leaf.loopIterationCount);
       llvm::append_range(shape, leaf.reconstructedType.getShape());
       if (!collection || collection.getType() != RankedTensorType::get(shape, leaf.reconstructedType.getElementType())) return failure();
       leafStacks.push_back(collection);
@@ -337,12 +446,102 @@ FailureOr<Value> realizeDeferredResult(const DeferredResultPlan &plan, Value lan
                                           if (!value.getType().isIndex()) selected = arith::IndexCastOp::create(context.rewriter, loc, value.getType(), selected);
                                           mapping.map(value, selected);
                                       }
-                                      for (auto [index, leaf] : llvm::enumerate(exchange.program.leaves)) { Value selected = extractMixedSliceOrIdentity(context.rewriter, loc, leafStacks[index], leaf.reconstructedType, leadingSlice(context.rewriter, leaf.reconstructedType, specialization)); if (!selected) return failure(); mapping.map(leaf.replacementRoot, selected); }
-                                      Value fragment = cloneResidual(exchange, mapping, context);
+                                      SmallVector<Value> selectedLeaves;
+                                      for (auto [index, leaf] : llvm::enumerate(exchange.program.leaves)) {
+                                        RankedTensorType selectedType =
+                                          leaf.reconstructedType;
+                                        if (leaf.enclosingLoop) {
+                                          SmallVector<int64_t> shape {
+                                            static_cast<int64_t>(leaf.loopIterationCount)};
+                                          llvm::append_range(shape, leaf.reconstructedType.getShape());
+                                          selectedType = RankedTensorType::get(
+                                            shape, leaf.reconstructedType.getElementType());
+                                        }
+                                        Value selected = extractMixedSliceOrIdentity(context.rewriter, loc, leafStacks[index], selectedType, leadingSlice(context.rewriter, selectedType, specialization));
+                                        if (!selected) return failure();
+                                        selectedLeaves.push_back(selected);
+                                        if (!leaf.enclosingLoop)
+                                          mapping.map(leaf.replacementRoot, selected);
+                                      }
+                                      Value fragment = cloneResidual(
+                                        exchange, mapping, selectedLeaves,
+                                        context);
                                       if (!fragment || fragment.getType() != fragmentType) return failure();
                                       yielded.push_back(insertMixedSlice(context.rewriter, loc, fragment, iterArgs.front(), leadingSlice(context.rewriter, fragmentType, specialization)));
                                       return success();
                                     });
   return failed(loop) || loop->results.size() != 1 ? FailureOr<Value>(failure()) : FailureOr<Value>(loop->results.front());
 }
+
+bool canRealizeDeferredLoopResult(const DeferredResultPlan &plan) {
+  DeferredExchangePlan &exchange = *plan.exchange;
+  if (exchange.program.specializationCount != 1
+      || exchange.program.leaves.size() != 1)
+    return false;
+  const DeferredProjectionLeafTemplate &leaf = exchange.program.leaves.front();
+  scf::ForOp loop = leaf.enclosingLoop;
+  return loop
+      && llvm::count_if(exchange.program.residualOps, [](Operation *op) {
+           return !op->hasTrait<OpTrait::ConstantLike>();
+         }) == 1
+      && llvm::is_contained(exchange.program.residualOps, loop.getOperation())
+      && loop.getInitArgs().size() == 1 && loop.getNumResults() == 1
+      && exchange.program.yieldedValue == loop.getResult(0)
+      && (leaf.form != DeferredLeafForm::GraphBatchProjection
+          || getDeferredProjectedFragmentType(leaf) == leaf.reconstructedType);
+}
+
+FailureOr<Value> realizeDeferredLoopResult(
+    const DeferredResultPlan &plan, Value lane,
+    DeferredLoopEmitter emitLoop,
+    DeferredEmissionContext &context) {
+  if (!canRealizeDeferredLoopResult(plan))
+    return failure();
+  DeferredExchangePlan &exchange = *plan.exchange;
+  scf::ForOp oldLoop = exchange.program.leaves.front().enclosingLoop;
+  IRMapping mapping;
+  if (exchange.program.scheduledLane)
+    mapping.map(exchange.program.scheduledLane, lane);
+  for (auto &[value, grid] : plan.residualValues) {
+    Value selected = grid.emitLookup(
+      context.constants.getIndex(0),
+      lane ? lane : context.constants.getIndex(0), exchange.deferred,
+      context.constants, context.rewriter, exchange.deferred.getLoc());
+    if (!value.getType().isIndex())
+      selected = arith::IndexCastOp::create(
+        context.rewriter, exchange.deferred.getLoc(), value.getType(), selected);
+    mapping.map(value, selected);
+  }
+  for (Operation &op :
+       exchange.deferred.getBody().front().without_terminator()) {
+    if (!op.hasTrait<OpTrait::ConstantLike>())
+      continue;
+    Operation *copy = context.rewriter.clone(op, mapping);
+    for (auto [oldValue, newValue] :
+         llvm::zip(op.getResults(), copy->getResults()))
+      mapping.map(oldValue, newValue);
+  }
+  Value initial = mapping.lookupOrDefault(oldLoop.getInitArgs().front());
+  auto result = emitLoop(
+    initial,
+    [&](Value fragment, Value iteration, Value current) -> FailureOr<Value> {
+      if (fragment.getType()
+          != exchange.program.leaves.front().reconstructedType)
+        return failure();
+      SmallVector<Value> yielded;
+      if (failed(cloneLoopIteration(
+            exchange, oldLoop, mapping, iteration, ValueRange {current},
+            [&](unsigned leafIndex) -> FailureOr<Value> {
+              return leafIndex == 0 ? FailureOr<Value>(fragment)
+                                    : FailureOr<Value>(failure());
+            },
+            context, yielded))
+          || yielded.size() != 1)
+        return failure();
+      return yielded.front();
+    });
+  return succeeded(result)
+             && result->getType() == exchange.deferred.getOutput().getType()
+           ? result : FailureOr<Value>(failure());
+}
 } // namespace onnx_mlir::spatial

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

@@ -2,6 +2,7 @@
 
 #include "DeferredCommunicationModel.hpp"
 #include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp"
+#include "llvm/ADT/STLFunctionalExtras.h"
 #include <array>
 
 namespace onnx_mlir::spatial {
@@ -54,4 +55,16 @@ mlir::FailureOr<mlir::Value> realizeDeferredResult(const DeferredResultPlan& pla
                                                    mlir::Value lane,
                                                    DeferredEmissionContext& context);
 
+using DeferredLoopIterationEmitter = llvm::function_ref<
+    mlir::FailureOr<mlir::Value>(mlir::Value, mlir::Value, mlir::Value)>;
+using DeferredLoopEmitter = llvm::function_ref<mlir::FailureOr<mlir::Value>(
+    mlir::Value, DeferredLoopIterationEmitter)>;
+
+bool canRealizeDeferredLoopResult(const DeferredResultPlan &plan);
+
+mlir::FailureOr<mlir::Value> realizeDeferredLoopResult(
+    const DeferredResultPlan &plan, mlir::Value lane,
+    DeferredLoopEmitter emitLoop,
+    DeferredEmissionContext &context);
+
 } // namespace onnx_mlir::spatial

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

@@ -133,11 +133,14 @@ static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
   for (unsigned specialization = 0;
        specialization < exchange.program.specializationCount;
        ++specialization) {
-    DeferredLaneValueEvaluator evaluator(
-      exchange.program, exchange.targetLaneCount, specialization);
     for (auto leafItem : llvm::enumerate(exchange.program.leaves)) {
       unsigned leafIndex = leafItem.index();
       const DeferredProjectionLeafTemplate &leaf = leafItem.value();
+      for (unsigned iteration = 0;
+           iteration < leaf.loopIterationCount; ++iteration) {
+      DeferredLaneValueEvaluator evaluator(
+        exchange.program, exchange.targetLaneCount, specialization,
+        leaf.enclosingLoop, iteration);
       auto sourceIndices = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
       if (failed(sourceIndices))
         return failure();
@@ -273,7 +276,9 @@ static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
             ++end;
           RequirementFamily family;
           family.exchange = &exchange;
-          family.coordinate = {specialization, leafIndex, position};
+          family.coordinate = {
+            specialization, leafIndex,
+            iteration * positionCount + position};
           family.targetLanes = LaneSet::range(begin, end);
           family.producer = producers[begin];
           family.publicationFragmentType = fragmentTypes[begin];
@@ -297,6 +302,7 @@ static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
           begin = end;
         }
       }
+      }
     }
   }
   return success();

src/PIM/Pass/PIMPasses.h

@@ -28,6 +28,8 @@ std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass(
 
 std::unique_ptr<mlir::Pass> createPimHostConstantFoldingPass();
 
+std::unique_ptr<mlir::Pass> createPimInstructionSelectionPass();
+
 std::unique_ptr<mlir::Pass> createPimLocalMemoryPlanningPass();
 
 std::unique_ptr<mlir::Pass> createPimVerificationPass();

src/PIM/PimAccelerator.cpp

@@ -79,6 +79,7 @@ void PimAccelerator::registerPasses(int optLevel) const {
   mlir::registerPass([] { return createTrivialGraphComputeMergePass(); });
   mlir::registerPass([] { return createMergeComputeNodesPass(); });
   registerPass(createPimHostConstantFoldingPass);
+  registerPass(createPimInstructionSelectionPass);
   registerPass(createPimLocalMemoryPlanningPass);
   registerPass(createPimVerificationPass);
   registerPass(createEmitPimCodePass);