temp commit: i will soft-reset and recommit after next changes

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -1,10 +1,13 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 
+#include "llvm/ADT/SmallPtrSet.h"
+
 #include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
@@ -36,6 +39,10 @@ mlir::Value resolveAlias(mlir::Value value, const StaticValueKnowledge* knowledg
 
 llvm::FailureOr<CompiledIndexExpr> compileIndexValueImpl(mlir::Value value);
 llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Value value);
+using AliasResolutionSet = llvm::SmallPtrSet<mlir::Value, 8>;
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value,
+                                        const StaticValueKnowledge* knowledge,
+                                        AliasResolutionSet& visited);
 mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
 template <typename... Args>
@@ -45,18 +52,23 @@ CompiledIndexExpr makeCompiledIndexExpr(Args&&... args) {
 
 static mlir::Value resolveForYieldedAliasToInit(mlir::scf::ForOp forOp,
                                                 mlir::Value yieldedValue,
-                                                const StaticValueKnowledge* knowledge) {
-  yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge);
+                                                const StaticValueKnowledge* knowledge,
+                                                AliasResolutionSet& visited) {
+  yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge, visited);
   if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(yieldedValue)) {
     if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0
         && static_cast<unsigned>(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size())
-      return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge);
+      return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge, visited);
   }
   return yieldedValue;
 }
 
-mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value,
+                                        const StaticValueKnowledge* knowledge,
+                                        AliasResolutionSet& visited) {
   value = resolveAlias(value, knowledge);
+  if (!value || !visited.insert(value).second)
+    return value;
 
   if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(value)) {
     auto forOp = mlir::dyn_cast_or_null<mlir::scf::ForOp>(blockArgument.getOwner()->getParentOp());
@@ -64,9 +76,12 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
       const unsigned iterArgIndex = blockArgument.getArgNumber() - 1;
       auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
       if (iterArgIndex < forOp.getInitArgs().size() && yieldOp
-          && iterArgIndex < yieldOp.getNumOperands()
-          && resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge) == blockArgument)
-        return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge);
+          && iterArgIndex < yieldOp.getNumOperands()) {
+        mlir::Value yieldedValue = resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge);
+        if (yieldedValue == blockArgument
+            || (yieldedValue && resolveLoopCarriedAliasImpl(yieldedValue, knowledge, visited) == blockArgument))
+          return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge, visited);
+      }
     }
     return value;
   }
@@ -75,10 +90,15 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
   if (!definingOp)
     return value;
 
+  if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(toBufferOp.getTensor(), knowledge, visited);
+  if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(toTensorOp.getBuffer(), knowledge, visited);
+
   if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
     if (auto result = mlir::dyn_cast<mlir::OpResult>(value))
       if (mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(result))
-        return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge);
+        return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge, visited);
   }
 
   if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(definingOp)) {
@@ -86,20 +106,26 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
     if (result) {
       auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
       if (yieldOp && result.getResultNumber() < yieldOp.getNumOperands())
-        return resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
+        return resolveForYieldedAliasToInit(
+          forOp, yieldOp.getOperand(result.getResultNumber()), knowledge, visited);
     }
   }
 
   if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge);
+    return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge, visited);
   if (auto collapseOp = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge);
+    return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge, visited);
   if (auto expandOp = mlir::dyn_cast<mlir::memref::ExpandShapeOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge);
+    return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge, visited);
 
   return value;
 }
 
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
+  AliasResolutionSet visited;
+  return resolveLoopCarriedAliasImpl(value, knowledge, visited);
+}
+
 llvm::FailureOr<int64_t> resolveOpFoldResult(mlir::OpFoldResult ofr, const StaticValueKnowledge* knowledge);
 llvm::FailureOr<int64_t> resolveIndexValueImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
@@ -524,6 +550,15 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
     if (!definingOp)
       return mlir::failure();
 
+    if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp)) {
+      value = resolveAlias(toBufferOp.getTensor(), knowledge);
+      continue;
+    }
+    if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp)) {
+      value = resolveAlias(toTensorOp.getBuffer(), knowledge);
+      continue;
+    }
+
     if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
       mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast<mlir::OpResult>(value));
       if (!tiedOperand)
@@ -538,7 +573,9 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
+      AliasResolutionSet visited;
+      value = resolveForYieldedAliasToInit(
+        forOp, yieldOp.getOperand(result.getResultNumber()), knowledge, visited);
       continue;
     }
 
@@ -643,6 +680,15 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
     if (!definingOp)
       return mlir::failure();
 
+    if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp)) {
+      value = toBufferOp.getTensor();
+      continue;
+    }
+    if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp)) {
+      value = toTensorOp.getBuffer();
+      continue;
+    }
+
     if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
       mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast<mlir::OpResult>(value));
       if (!tiedOperand)
@@ -657,7 +703,9 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), nullptr);
+      AliasResolutionSet visited;
+      value = resolveForYieldedAliasToInit(
+        forOp, yieldOp.getOperand(result.getResultNumber()), nullptr, visited);
       continue;
     }
 

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -87,11 +87,6 @@ llvm::cl::opt<uint64_t> pimConvStreamChunkPositions(
   llvm::cl::init(1024),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> pimReportConvLowering("pim-report-conv-lowering",
-                                          llvm::cl::desc("Emit a bounded Conv lowering report"),
-                                          llvm::cl::init(true),
-                                          llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
                                 llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"),
                                 llvm::cl::init(false),

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -57,7 +57,6 @@ extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;
 extern llvm::cl::opt<bool> pimOnlyCodegen;
 extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<bool> pimEmitJson;
-extern llvm::cl::opt<bool> pimReportConvLowering;
 extern llvm::cl::opt<bool> pimDetectCommunicationDeadlock;
 extern llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder;
 extern llvm::cl::opt<bool> pimTraceCommunicationMaterialization;

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

@@ -47,13 +47,17 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
     llvm::append_range(offsets, entry.destinationOffsets);
     llvm::append_range(sizes, entry.sizes);
   }
-  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
+  auto blueprint = spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
     rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout),
     rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
     rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"),
     rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
     rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
-    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete"));
+  if (indexMap == spatial::kContiguousRowMajorFragments
+      && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    blueprint.setIndexMapAttr(rewriter.getStringAttr("fragment_assembly"));
+  return blueprint.getOutput();
 }
 
 Value sumTensors(ArrayRef<Value> tensors, PatternRewriter& rewriter) {

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

@@ -394,6 +394,39 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
     rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
 }
 
+template <typename BodyFn>
+mlir::FailureOr<mlir::Value> mapGraphBatchFragments(mlir::Value input,
+                                                    mlir::RankedTensorType outputType,
+                                                    mlir::PatternRewriter& rewriter,
+                                                    mlir::Location loc,
+                                                    BodyFn&& build) {
+  auto inputType = mlir::dyn_cast<mlir::RankedTensorType>(input.getType());
+  if (!inputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
+      || inputType.getRank() != outputType.getRank() || inputType.getRank() < 2
+      || inputType.getDimSize(0) != outputType.getDimSize(0))
+    return mlir::failure();
+  auto inputFragmentType = mlir::RankedTensorType::get(
+    inputType.getShape().drop_front(), inputType.getElementType(), inputType.getEncoding());
+  auto outputFragmentType = mlir::RankedTensorType::get(
+    outputType.getShape().drop_front(), outputType.getElementType(), outputType.getEncoding());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, mlir::TypeRange {outputType}, inputType.getDimSize(0), {}, mlir::ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) -> mlir::LogicalResult {
+      auto fragment = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), args.lane, inputFragmentType);
+      if (mlir::failed(fragment))
+        return mlir::failure();
+      mlir::FailureOr<mlir::Value> result = build(*fragment, outputFragmentType);
+      if (mlir::failed(result) || result->getType() != outputFragmentType)
+        return mlir::failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, *result, args.outputs.front(), args.lane);
+      return mlir::success();
+    });
+  if (mlir::failed(batch))
+    return mlir::failure();
+  return batch->getResult(0);
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -186,25 +186,9 @@ static FailureOr<Value> applyRowStripActivation(const RowStripPhysicalValue& val
                                                 Location loc,
                                                 BuildActivation buildActivation) {
   auto storageType = cast<RankedTensorType>(value.storage.getType());
-  const int64_t laneCount = storageType.getDimSize(0);
-  auto batchOp = createSpatComputeBatch(rewriter,
-                                        loc,
-                                        TypeRange {storageType},
-                                        laneCount,
-                                        {},
-                                        ValueRange {value.storage},
-                                        [&](detail::SpatComputeBatchBodyArgs args) {
-                                          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
-                                            rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
-                                          if (failed(fragment)) return failure();
-                                          Value result = buildActivation(*fragment);
-                                          publishGraphBatchPhysicalFragment(
-                                            rewriter, loc, result, args.outputs.front(), args.lane);
-                                          return success();
-                                        });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  return mapGraphBatchFragments(value.storage, storageType, rewriter, loc, [&](Value fragment, RankedTensorType) {
+    return FailureOr<Value>(buildActivation(fragment));
+  });
 }
 
 FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -399,6 +399,42 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         rewriter.replaceOp(planOp, computeOp.getResults());
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatResizeNearestPlanOp>(&op)) {
+        FailureOr<RowStripPhysicalValue> input =
+          getRowStripValue(rowStripValues, planOp.getInput());
+        rewriter.setInsertionPoint(planOp);
+        auto lowered = lowerSelectedResizeNearestPlan(
+          planOp, succeeded(input) ? std::optional<Value>(input->storage) : std::nullopt,
+          rewriter);
+        if (failed(lowered)) {
+          planOp.emitOpError("failed to lower selected nearest Resize plan");
+          signalPassFailure();
+          return;
+        }
+        if (failed(input)) {
+          rewriter.replaceOp(planOp, *lowered);
+          continue;
+        }
+        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 Resize plan requires a row-strip blueprint result");
+          signalPassFailure();
+          return;
+        }
+        auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+        auto output = buildRowStripValue(blueprint, *lowered);
+        if (failed(output)) {
+          signalPassFailure();
+          return;
+        }
+        rowStripValues[blueprint.getResult()] = *output;
+        eraseAfterLowering.insert(planOp);
+        eraseAfterLowering.insert(blueprint);
+        continue;
+      }
       if (auto planOp = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
         auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
           auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
@@ -701,6 +737,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
                    spatial::SpatAddPlanOp,
                    spatial::SpatReluPlanOp,
                    spatial::SpatSiluPlanOp,
+                   spatial::SpatResizeNearestPlanOp,
                    spatial::SpatMaxPool2DPlanOp,
                    spatial::SpatGlobalAveragePoolPlanOp,
                    spatial::SpatMaterializeLayoutOp>(op)

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -52,13 +52,16 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   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::SpatResizeNearestPlanOp> resizePlans(
+    funcOp.getOps<spatial::SpatResizeNearestPlanOp>());
   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() || !siluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty()
+      || !concatPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !resizePlans.empty()
+      || !maxPoolPlans.empty() || !blueprints.empty()
       || !globalAveragePoolPlans.empty() || !materializers.empty()) {
     return;
   }
@@ -123,6 +126,14 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  RewritePatternSet matmulPatterns(ctx);
+  populateMatMulFusionPatterns(matmulPatterns, ctx);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(matmulPatterns)))) {
+    moduleOp.emitError("failed to lower MatMul before producer conversion");
+    signalPassFailure();
+    return;
+  }
+
   RewritePatternSet fusionPatterns(ctx);
   populateElementwiseFusionPatterns(fusionPatterns, ctx);
   if (failed(applyPatternsGreedily(moduleOp, std::move(fusionPatterns)))) {

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

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

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -20,6 +20,7 @@ void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRCo
 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 populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReduceMeanPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

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

@@ -48,6 +48,56 @@ static DenseElementsAttr getDenseConstantAttr(Value value) {
   return nullptr;
 }
 
+struct BlueprintSplatMulToSpatial : OpConversionPattern<ONNXMulOp> {
+  explicit BlueprintSplatMulToSpatial(MLIRContext* ctx) : OpConversionPattern(ctx, 2) {}
+
+  LogicalResult
+  matchAndRewrite(ONNXMulOp op, ONNXMulOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto blueprint = adaptor.getA().getDefiningOp<spatial::SpatBlueprintOp>();
+    Value scalar = adaptor.getB();
+    if (!blueprint) {
+      blueprint = adaptor.getB().getDefiningOp<spatial::SpatBlueprintOp>();
+      scalar = adaptor.getA();
+    }
+    auto scalarAttr = getDenseConstantAttr(scalar);
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    auto storageType = blueprint ? dyn_cast<RankedTensorType>(blueprint.getInput().getType()) : RankedTensorType();
+    if (!blueprint || !blueprint.getFragments().empty() || !scalarAttr || !scalarAttr.isSplat() || !resultType
+        || resultType != blueprint.getOutput().getType() || !storageType)
+      return failure();
+
+    auto mapped = mapGraphBatchFragments(
+      blueprint.getInput(), storageType, rewriter, op.getLoc(), [&](Value fragment, RankedTensorType fragmentType) {
+        auto splat = DenseElementsAttr::get(fragmentType, scalarAttr.getSplatValue<Attribute>());
+        Value constant = arith::ConstantOp::create(rewriter, op.getLoc(), fragmentType, splat);
+        return FailureOr<Value>(
+          spatial::SpatVMulOp::create(rewriter, op.getLoc(), fragmentType, fragment, constant).getResult());
+      });
+    if (failed(mapped))
+      return failure();
+
+    auto result = spatial::SpatBlueprintOp::create(rewriter,
+                                                   op.getLoc(),
+                                                   resultType,
+                                                   *mapped,
+                                                   ValueRange {},
+                                                   blueprint.getLogicalLayoutAttr(),
+                                                   blueprint.getPhysicalLayoutAttr(),
+                                                   blueprint.getFragmentOffsetsAttr(),
+                                                   blueprint.getFragmentSizesAttr(),
+                                                   blueprint.getIndexMapAttr(),
+                                                   blueprint.getModeAttr(),
+                                                   blueprint.getFragmentOperandIndicesAttr(),
+                                                   blueprint.getFragmentSourceSlotsAttr(),
+                                                   blueprint.getFragmentSourceOffsetsAttr(),
+                                                   blueprint.getFragmentStridesAttr(),
+                                                   blueprint.getConflictPolicyAttr(),
+                                                   blueprint.getCoveragePolicyAttr());
+    rewriter.replaceOp(op, result.getOutput());
+    return success();
+  }
+};
+
 static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
                                                              RankedTensorType resultType,
                                                              ConversionPatternRewriter& rewriter,
@@ -246,6 +296,7 @@ void populateElementwiseFusionPatterns(RewritePatternSet& patterns, MLIRContext*
 }
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<BlueprintSplatMulToSpatial>(ctx);
   patterns.add<AddToSpatialCompute>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);

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

@@ -106,6 +106,92 @@ static Value mapOutputBatchIndexToSourceBatchIndex(Value outputBatchIndex,
   return sourceBatchIndex;
 }
 
+static FailureOr<Value> collapseFragmentAssemblyBatchDims(Value value,
+                                                          RankedTensorType resultType,
+                                                          PatternRewriter& rewriter,
+                                                          Location loc) {
+  auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>();
+  auto inputType = dyn_cast<RankedTensorType>(value.getType());
+  auto storageType = blueprint ? dyn_cast<RankedTensorType>(blueprint.getInput().getType()) : RankedTensorType();
+  auto operandIndices = blueprint ? blueprint.getFragmentOperandIndices() : std::nullopt;
+  auto sourceOffsets = blueprint ? blueprint.getFragmentSourceOffsets() : std::nullopt;
+  auto fragmentStrides = blueprint ? blueprint.getFragmentStrides() : std::nullopt;
+  if (!blueprint || !inputType || !storageType || !inputType.hasStaticShape() || !storageType.hasStaticShape()
+      || inputType.getRank() <= 3 || resultType.getRank() != 3 || !blueprint.getFragments().empty()
+      || blueprint.getMode() != "fragment_assembly" || !operandIndices || !sourceOffsets || !fragmentStrides
+      || storageType.getRank() != inputType.getRank() + 1)
+    return failure();
+  if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments
+      && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+
+  const int64_t batchRank = inputType.getRank() - 2;
+  SmallVector<ReassociationIndices> reassociation {ReassociationIndices {},
+                                                   ReassociationIndices {batchRank},
+                                                   ReassociationIndices {batchRank + 1}};
+  for (int64_t dim = 0; dim < batchRank; ++dim)
+    reassociation.front().push_back(dim);
+  SmallVector<int64_t> outputFragmentShape {1, storageType.getDimSize(batchRank + 1), storageType.getDimSize(batchRank + 2)};
+  auto outputFragmentType = RankedTensorType::get(outputFragmentShape, storageType.getElementType());
+  auto outputStorageType = spatial::getGraphBatchPhysicalResultType(storageType.getDimSize(0), outputFragmentType);
+  SmallVector<ReassociationIndices> storageReassociation {
+    ReassociationIndices {0}, ReassociationIndices {}, ReassociationIndices {batchRank + 1},
+    ReassociationIndices {batchRank + 2}};
+  for (int64_t dim = 0; dim < batchRank; ++dim)
+    storageReassociation[1].push_back(dim + 1);
+  Value collapsedStorage = tensor::CollapseShapeOp::create(
+    rewriter, loc, outputStorageType, blueprint.getInput(), storageReassociation);
+
+  const int64_t inputRank = inputType.getRank();
+  const int64_t fragmentCount = operandIndices->size();
+  ArrayRef<int64_t> inputOffsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> inputSizes = blueprint.getFragmentSizes();
+  SmallVector<int64_t> batchShape(inputType.getShape().drop_back(2));
+  SmallVector<int64_t> batchStrides = computeRowMajorStrides(batchShape);
+  SmallVector<int64_t> offsets, sizes, strides;
+  offsets.reserve(fragmentCount * 3);
+  sizes.reserve(fragmentCount * 3);
+  strides.reserve(fragmentCount * 3);
+  for (int64_t fragment = 0; fragment < fragmentCount; ++fragment) {
+    int64_t flatBatch = 0;
+    for (int64_t dim = 0; dim < batchRank; ++dim) {
+      const int64_t index = fragment * inputRank + dim;
+      if (inputSizes[index] != 1 || (*fragmentStrides)[index] != 1)
+        return failure();
+      flatBatch += inputOffsets[index] * batchStrides[dim];
+    }
+    offsets.push_back(flatBatch);
+    sizes.push_back(1);
+    strides.push_back(1);
+    for (int64_t dim = batchRank; dim < inputRank; ++dim) {
+      const int64_t index = fragment * inputRank + dim;
+      offsets.push_back(inputOffsets[index]);
+      sizes.push_back(inputSizes[index]);
+      strides.push_back((*fragmentStrides)[index]);
+    }
+  }
+  auto collapsedBlueprint = spatial::SpatBlueprintOp::create(rewriter,
+                                                              loc,
+                                                              resultType,
+                                                              collapsedStorage,
+                                                              ValueRange {},
+                                                              blueprint.getLogicalLayoutAttr(),
+                                                              rewriter.getStringAttr("fragmented"),
+                                                              rewriter.getDenseI64ArrayAttr(offsets),
+                                                              rewriter.getDenseI64ArrayAttr(sizes),
+                                                              rewriter.getStringAttr("collapsed_fragments"),
+                                                              blueprint.getModeAttr(),
+                                                              blueprint.getFragmentOperandIndicesAttr(),
+                                                              blueprint.getFragmentSourceSlotsAttr(),
+                                                              blueprint.getFragmentSourceOffsetsAttr(),
+                                                              rewriter.getDenseI64ArrayAttr(strides),
+                                                              blueprint.getConflictPolicyAttr(),
+                                                              blueprint.getCoveragePolicyAttr());
+  if (spatial::isCanonicalContiguousRowMajorFragmentAssembly(collapsedBlueprint))
+    collapsedBlueprint.setIndexMapAttr(rewriter.getStringAttr(spatial::kContiguousRowMajorFragments));
+  return collapsedBlueprint.getOutput();
+}
+
 static Value
 collapseBatchDims(Value value, int64_t batchSize, int64_t rows, int64_t cols, PatternRewriter& rewriter, Location loc) {
   auto type = cast<RankedTensorType>(value.getType());
@@ -113,6 +199,8 @@ collapseBatchDims(Value value, int64_t batchSize, int64_t rows, int64_t cols, Pa
     return value;
 
   auto collapsedType = RankedTensorType::get({batchSize, rows, cols}, type.getElementType(), type.getEncoding());
+  if (auto collapsed = collapseFragmentAssemblyBatchDims(value, collapsedType, rewriter, loc); succeeded(collapsed))
+    return *collapsed;
   SmallVector<ReassociationIndices> reassociation = {ReassociationIndices {},
                                                      ReassociationIndices {static_cast<int64_t>(type.getRank() - 2)},
                                                      ReassociationIndices {static_cast<int64_t>(type.getRank() - 1)}};
@@ -241,7 +329,34 @@ static Value extractBatchMatrix(Value value,
   return materializeOrComputeUnary(value, matrixType, rewriter, loc, buildMatrix);
 }
 
+static Value getLastTwoTransposeInput(Value value) {
+  auto type = cast<RankedTensorType>(value.getType());
+  if (auto transpose = value.getDefiningOp<ONNXTransposeOp>()) {
+    auto permutation = getTransposePermutationChecked(transpose.getPermAttr(), type.getRank());
+    if (succeeded(permutation) && llvm::all_of(llvm::seq<int64_t>(0, type.getRank()), [&](int64_t dim) {
+          return (*permutation)[dim] == (dim < type.getRank() - 2 ? dim : 2 * type.getRank() - 3 - dim);
+        }))
+      return transpose.getData();
+  }
+  return {};
+}
+
+static std::pair<Value, Value> splitSplatMultiply(Value value) {
+  if (!value)
+    return {};
+  auto multiply = value.getDefiningOp<ONNXMulOp>();
+  if (!multiply)
+    return {};
+  for (auto [data, scale] : {std::pair {multiply.getA(), multiply.getB()},
+                             std::pair {multiply.getB(), multiply.getA()}})
+    if (auto constant = getHostConstDenseElementsAttr(scale); constant && constant.isSplat())
+      return {data, scale};
+  return {};
+}
+
 static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Location loc) {
+  if (Value input = getLastTwoTransposeInput(value))
+    return input;
   auto type = cast<RankedTensorType>(value.getType());
   auto shape = type.getShape();
   auto createONNXTranspose = [&](RankedTensorType resultType, ArrayRef<int64_t> permutation) {
@@ -407,123 +522,152 @@ static Value extractDynamicBatchedRowVector(Value matrix,
     {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
+static int64_t chooseDynamicMatMulRowsPerLane(int64_t rows, int64_t reductionSize, int64_t columns) {
+  const int64_t crossbarElements = static_cast<int64_t>(crossbarSize.getValue() * crossbarSize.getValue());
+  const int64_t target = std::min(rows, ceilIntegerDivide(reductionSize * columns, crossbarElements));
+  int64_t rowsPerLane = 1;
+  for (int64_t candidate = 2; candidate <= target; ++candidate)
+    if (rows % candidate == 0)
+      rowsPerLane = candidate;
+  return rowsPerLane;
+}
+
 static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
                                                                      ArrayRef<int64_t> aBatchShape,
                                                                      Value b,
                                                                      ArrayRef<int64_t> bBatchShape,
                                                                      ArrayRef<int64_t> outputBatchShape,
-                                                                     RankedTensorType aType,
-                                                                     RankedTensorType bType,
-                                                                     RankedTensorType columnPiecesType,
+                                                                     int64_t reductionSize,
+                                                                     int64_t rowsPerLane,
+                                                                     RankedTensorType rowPiecesType,
                                                                      RankedTensorType outType,
+                                                                     RankedTensorType publicationFragmentType,
                                                                      PatternRewriter& rewriter,
                                                                      Location loc) {
   const int64_t numBatches = outType.getDimSize(0);
   const int64_t numOutRows = outType.getDimSize(1);
   const int64_t numOutCols = outType.getDimSize(2);
-  const int64_t reductionSize = aType.getDimSize(2);
-  const int64_t laneCount = numBatches * numOutCols;
-  auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
+  const int64_t rowGroups = numOutRows / rowsPerLane;
+  const int64_t laneCount = numBatches * rowGroups;
+  auto vectorType = RankedTensorType::get({1, reductionSize}, outType.getElementType());
   auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-  auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
+  auto rowType = RankedTensorType::get({1, numOutCols}, outType.getElementType());
+  auto rowsType = RankedTensorType::get({rowsPerLane, numOutCols}, outType.getElementType());
   auto batchOp = createSpatComputeBatch(
     rewriter,
     loc,
-    TypeRange {columnPiecesType},
+    TypeRange {rowPiecesType},
     laneCount,
     ValueRange {},
     ValueRange {a, b},
-    [&](detail::SpatComputeBatchBodyArgs args) {
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-      Value batch = affineFloorDivConst(rewriter, loc, args.lane, numOutCols, anchorOp);
-      Value column = affineModConst(rewriter, loc, args.lane, numOutCols, anchorOp);
-      Value bVector = extractDynamicBatchedRowVector(
-        args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc);
-      Value columnInit = tensor::EmptyOp::create(rewriter, loc, columnType.getShape(), columnType.getElementType());
+      Value batch = affineFloorDivConst(rewriter, loc, args.lane, rowGroups, anchorOp);
+      Value rowGroup = affineModConst(rewriter, loc, args.lane, rowGroups, anchorOp);
+      Value rowBase = affineMulConst(rewriter, loc, rowGroup, rowsPerLane, anchorOp);
       Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
       Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
-      Value cNumOutRows =
-        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutRows);
-      auto loop = buildNormalizedScfFor(
+      Value cRows = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), rowsPerLane);
+      Value cNumOutCols =
+        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutCols);
+      Value rowsInit = tensor::EmptyOp::create(rewriter, loc, rowsType.getShape(), rowsType.getElementType());
+      auto rowsLoop = buildNormalizedScfFor(
         rewriter,
         loc,
         c0,
-        cNumOutRows,
+        cRows,
         c1,
-        ValueRange {columnInit},
-        [&](OpBuilder&, Location nestedLoc, Value row, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+        ValueRange {rowsInit},
+        [&](OpBuilder&, Location nestedLoc, Value rowOffset, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value row = arith::AddIOp::create(rewriter, nestedLoc, rowBase, rowOffset);
           Value aVector = extractDynamicBatchedRowVector(
             args.inputs[0], aBatchShape, outputBatchShape, batch, row, vectorType, rewriter, nestedLoc);
-          Value scalar = spatial::SpatVVDMulOp::create(rewriter, nestedLoc, scalarType, aVector, bVector).getResult();
-          Value next = tensor::InsertSliceOp::create(rewriter,
-                                                     nestedLoc,
-                                                     scalar,
-                                                     iterArgs.front(),
-                                                     SmallVector<OpFoldResult> {row, rewriter.getIndexAttr(0)},
-                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
-                                                                               rewriter.getIndexAttr(1)},
-                                                     getUnitStrides(rewriter, 2));
-          yielded.push_back(next);
+          Value rowInit = tensor::EmptyOp::create(rewriter, nestedLoc, rowType.getShape(), rowType.getElementType());
+          auto columnsLoop = buildNormalizedScfFor(
+            rewriter,
+            nestedLoc,
+            c0,
+            cNumOutCols,
+            c1,
+            ValueRange {rowInit},
+            [&](OpBuilder&, Location columnLoc, Value column, ValueRange columnArgs, SmallVectorImpl<Value>& rowYielded) {
+              Value bVector = extractDynamicBatchedRowVector(
+                args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, columnLoc);
+              Value scalar = spatial::SpatVVDMulOp::create(rewriter, columnLoc, scalarType, aVector, bVector).getResult();
+              rowYielded.push_back(tensor::InsertSliceOp::create(
+                rewriter,
+                columnLoc,
+                scalar,
+                columnArgs.front(),
+                SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), column},
+                SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+                getUnitStrides(rewriter, 2)));
+              return success();
+            });
+          assert(succeeded(columnsLoop) && "dynamic MatMul column loop construction must succeed");
+          yielded.push_back(tensor::InsertSliceOp::create(
+            rewriter,
+            nestedLoc,
+            columnsLoop->results.front(),
+            iterArgs.front(),
+            SmallVector<OpFoldResult> {rowOffset, rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(numOutCols)},
+            getUnitStrides(rewriter, 2)));
           return success();
         });
-      assert(succeeded(loop) && "dynamic MatMul row loop construction must succeed");
-      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
+      assert(succeeded(rowsLoop) && "dynamic MatMul row-group loop construction must succeed");
+      Value fragment = rowsLoop->results.front();
+      while (fragment.getType() != publicationFragmentType) {
+        auto expanded = addLeadingUnitTensorDimension(rewriter, loc, fragment);
+        if (failed(expanded))
+          return failure();
+        fragment = *expanded;
+      }
+      publishGraphBatchPhysicalFragment(rewriter, loc, fragment, args.outputs.front(), args.lane);
+      return success();
     });
   if (failed(batchOp))
     return failure();
   return *batchOp;
 }
 
-static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
-                                                          RankedTensorType scalarPiecesType,
-                                                          RankedTensorType outType,
-                                                          PatternRewriter& rewriter,
-                                                          Location loc) {
-  const int64_t laneCount = scalarPiecesType.getDimSize(0);
-  const int64_t numOutCols = outType.getDimSize(2);
-  auto columnType = RankedTensorType::get({outType.getDimSize(1), 1}, outType.getElementType());
-
-  auto computeOp = createSpatCompute<1>(
-    rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult {
-      Value outputInit =
-        tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType()).getResult();
-      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
-      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
-      Value cLaneCount = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), laneCount);
-      auto loop = buildNormalizedScfFor(
-        rewriter,
-        loc,
-        c0,
-        cLaneCount,
-        c1,
-        ValueRange {outputInit},
-        [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
-          Value outputAcc = iterArgs.front();
-          Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-          Value batch = affineFloorDivConst(rewriter, nestedLoc, lane, numOutCols, anchorOp);
-          Value column = affineModConst(rewriter, nestedLoc, lane, numOutCols, anchorOp);
-          FailureOr<Value> columnPiece =
-            extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, columnType);
-          if (failed(columnPiece))
-            return failure();
-          SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), column};
-          SmallVector<OpFoldResult> outputSizes = {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(outType.getDimSize(1)), rewriter.getIndexAttr(1)};
-          Value next =
-            tensor::InsertSliceOp::create(
-              rewriter, nestedLoc, *columnPiece, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
-              .getResult();
-          yielded.push_back(next);
-          return success();
-        });
-      if (failed(loop))
-        return failure();
-      spatial::SpatYieldOp::create(rewriter, loc, loop->results.front());
-      return success();
-    });
-  if (failed(computeOp))
-    return failure();
-  return computeOp->getResult(0);
+static FailureOr<Value> createBatchedRowOutputBlueprint(Value rowPieces,
+                                                        RankedTensorType outType,
+                                                        ArrayRef<int64_t> batchShape,
+                                                        int64_t rowsPerFragment,
+                                                        PatternRewriter& rewriter,
+                                                        Location loc) {
+  SmallVector<FragmentAssemblyEntry> entries;
+  const int64_t rowAxis = outType.getRank() - 2;
+  const int64_t rows = outType.getDimSize(rowAxis);
+  const int64_t columns = outType.getDimSize(rowAxis + 1);
+  const int64_t batches = batchShape.empty() ? 1 : getStaticShapeElementCount(batchShape);
+  SmallVector<int64_t> batchStrides = computeRowMajorStrides(batchShape);
+  const int64_t rowGroups = rows / rowsPerFragment;
+  entries.reserve(batches * rowGroups);
+  for (int64_t batch = 0; batch < batches; ++batch)
+    for (int64_t row = 0; row < rows; row += rowsPerFragment) {
+      SmallVector<int64_t, 4> offsets;
+      for (auto [dim, size] : llvm::enumerate(batchShape))
+        offsets.push_back((batch / batchStrides[dim]) % size);
+      offsets.push_back(row);
+      offsets.push_back(0);
+      SmallVector<int64_t, 4> sizes(outType.getRank(), 1);
+      sizes[rowAxis] = rowsPerFragment;
+      sizes.back() = columns;
+      entries.push_back({batch * rowGroups + row / rowsPerFragment,
+                         0,
+                         std::move(offsets),
+                         std::move(sizes)});
+  }
+  return createFragmentAssemblyBlueprint(
+    rowPieces,
+    outType,
+    entries,
+    "dense_nchw",
+    rowsPerFragment == 1 ? spatial::kContiguousRowMajorFragments : "row_group_fragments",
+    rewriter,
+    loc);
 }
 
 static Value extractBatchedReductionPiece(Value partialPiecesArg,
@@ -822,6 +966,8 @@ static Value finalizeNormalizedMatMulResult(Value value,
                                             const NormalizedMatMulInfo& info,
                                             PatternRewriter& rewriter,
                                             Location loc) {
+  if (value.getType() == info.outType)
+    return value;
   // The direct lowered result is always [flatBatch, normalizedM, normalizedN].
   // Restore ONNX MatMul result rank by expanding right-aligned batch dimensions
   // and removing the synthetic unit matrix axes introduced for vector operands.
@@ -923,17 +1069,34 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     Location loc = matmulOp.getLoc();
     bool useTransposedForm = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
                           && isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
+    Value rhsRows = getLastTwoTransposeInput(matmulOp.getB());
+    ONNXTransposeOp foldedTranspose = matmulOp.getB().getDefiningOp<ONNXTransposeOp>();
+    ONNXMulOp foldedMultiply = rhsRows ? rhsRows.getDefiningOp<ONNXMulOp>() : ONNXMulOp {};
+    auto [unscaledRhsRows, outputScale] = splitSplatMultiply(rhsRows);
+    if (unscaledRhsRows)
+      rhsRows = unscaledRhsRows;
+    const bool rhsStoredAsRows = rhsRows && !useTransposedForm;
 
     Value lhs =
       normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType, shapeInfo->lhsWasVector, rewriter, loc);
-    Value rhs =
-      normalizeMatMulOperand(matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc);
+    Value rhs = normalizeMatMulOperand(
+      rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc);
     lhs = collapseBatchDims(lhs, shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
-    rhs = collapseBatchDims(rhs, shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
+    rhs = collapseBatchDims(rhs,
+                            shapeInfo->rhsBatch,
+                            rhsStoredAsRows ? shapeInfo->n : shapeInfo->k,
+                            rhsStoredAsRows ? shapeInfo->k : shapeInfo->n,
+                            rewriter,
+                            loc);
     MatMulLoweringPlan plan = buildLoweringPlan(lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc);
 
     plan.lhs = ensureBatchedTensor(plan.lhs, plan.lhsBatch, plan.m, plan.k, rewriter, loc);
-    plan.rhs = ensureBatchedTensor(plan.rhs, plan.rhsBatch, plan.k, plan.n, rewriter, loc);
+    plan.rhs = ensureBatchedTensor(plan.rhs,
+                                   plan.rhsBatch,
+                                   rhsStoredAsRows ? plan.n : plan.k,
+                                   rhsStoredAsRows ? plan.k : plan.n,
+                                   rewriter,
+                                   loc);
     plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
     plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
     auto directOutType = RankedTensorType::get(
@@ -997,26 +1160,34 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
         return success();
       }
     }
-    const int64_t laneCount = plan.batch * plan.n;
-    auto columnType = RankedTensorType::get({plan.m, 1}, shapeInfo->outType.getElementType());
-    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(
-      laneCount, columnType);
-    Value transposedRhs = transposeLastTwoDims(plan.rhs, rewriter, loc);
+    RankedTensorType blueprintType = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
+                                   ? shapeInfo->outType : directOutType;
+    SmallVector<int64_t> blueprintBatchShape = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
+                                             ? shapeInfo->outputBatchShape : SmallVector<int64_t> {plan.batch};
+    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.m, plan.k, plan.n);
+    const int64_t laneCount = plan.batch * plan.m / rowsPerLane;
+    SmallVector<int64_t> fragmentShape(blueprintType.getRank(), 1);
+    fragmentShape[fragmentShape.size() - 2] = rowsPerLane;
+    fragmentShape.back() = plan.n;
+    auto fragmentType = RankedTensorType::get(fragmentShape, shapeInfo->outType.getElementType());
+    auto rowPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount, fragmentType);
+    Value transposedRhs = rhsStoredAsRows ? plan.rhs : transposeLastTwoDims(plan.rhs, rewriter, loc);
     auto batchOp = createBatchedVvdmulBatch(plan.lhs,
                                             plan.lhsBatchShape,
                                             transposedRhs,
                                             plan.rhsBatchShape,
                                             plan.outputBatchShape,
-                                            plan.lhsType,
-                                            plan.rhsType,
-                                            scalarPiecesType,
+                                            plan.k,
+                                            rowsPerLane,
+                                            rowPiecesType,
                                             directOutType,
+                                            fragmentType,
                                             rewriter,
                                             loc);
     if (failed(batchOp))
       return failure();
-    auto result =
-      createBatchedDynamicOutputCompute(batchOp->getResult(0), scalarPiecesType, directOutType, rewriter, loc);
+    auto result = createBatchedRowOutputBlueprint(
+      batchOp->getResult(0), blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc);
     if (failed(result))
       return failure();
     Value finalResult = *result;
@@ -1029,13 +1200,34 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
           .getResult();
     }
     finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
+    if (outputScale)
+      finalResult = ONNXMulOp::create(
+        rewriter, loc, shapeInfo->outType, finalResult, outputScale).getResult();
     rewriter.replaceOp(matmulOp, finalResult);
+    if (foldedTranspose && foldedTranspose->use_empty())
+      rewriter.eraseOp(foldedTranspose);
+    if (foldedMultiply && foldedMultiply->use_empty())
+      rewriter.eraseOp(foldedMultiply);
     return success();
   }
 };
 
+struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
+  using MatMulBatchedToSpatialComputes::MatMulBatchedToSpatialComputes;
+
+  LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
+    if (!getLastTwoTransposeInput(matmulOp.getB()))
+      return failure();
+    return MatMulBatchedToSpatialComputes::matchAndRewrite(matmulOp, rewriter);
+  }
+};
+
 } // namespace
 
+void populateMatMulFusionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<TransposedRhsMatMulToSpatial>(ctx);
+}
+
 void populateMatMulRewritePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   patterns.insert<MatMulToGemm, MatMulBatchedToSpatialComputes>(ctx);
 }

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

@@ -5,8 +5,11 @@
 
 #include "llvm/ADT/STLExtras.h"
 
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -17,126 +20,144 @@ namespace onnx_mlir {
 namespace {
 
 static Value buildNearestAsymmetricIndex(
-  Value outputIndex, int64_t inputDim, int64_t outputDim, ConversionPatternRewriter& rewriter, Location loc) {
+  Value outputIndex, int64_t inputDim, int64_t outputDim, PatternRewriter& rewriter, Location loc) {
+  if (inputDim == outputDim)
+    return outputIndex;
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  if (outputDim % inputDim == 0)
+    return affineFloorDivConst(rewriter, loc, outputIndex, outputDim / inputDim, anchorOp);
+  if (inputDim % outputDim == 0)
+    return affineMulConst(rewriter, loc, outputIndex, inputDim / outputDim, anchorOp);
   Value cInputDim = getOrCreateIndexConstant(rewriter, anchorOp, inputDim);
   Value cOutputDim = getOrCreateIndexConstant(rewriter, anchorOp, outputDim);
-  Value cInputDimLast = getOrCreateIndexConstant(rewriter, anchorOp, inputDim - 1);
   Value scaledIndex = arith::MulIOp::create(rewriter, loc, outputIndex, cInputDim);
-  Value inputIndex = arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim);
-  return arith::MinUIOp::create(rewriter, loc, inputIndex, cInputDimLast);
+  return arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim);
 }
 
-static FailureOr<Value> buildNearestResizeLoop(Value input,
-                                               RankedTensorType inputType,
-                                               RankedTensorType resultType,
-                                               ConversionPatternRewriter& rewriter,
-                                               Location loc) {
-  auto elemType = resultType.getElementType();
-  SmallVector<int64_t> unitShape(resultType.getRank(), 1);
-  auto unitTensorType = RankedTensorType::get(unitShape, elemType);
-
-  SmallVector<OpFoldResult> unitSizes(resultType.getRank(), rewriter.getIndexAttr(1));
-  SmallVector<OpFoldResult> unitStrides(resultType.getRank(), rewriter.getIndexAttr(1));
-
-  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
-  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-  Value cOutputN = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(0));
-  Value cOutputC = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(1));
-  Value cOutputH = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(2));
-  Value cOutputW = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(3));
-
-  Value outputInit = tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), elemType);
-
-  auto batchLoop = buildNormalizedScfFor(
-    rewriter,
-    loc,
-    c0,
-    cOutputN,
-    c1,
-    ValueRange {outputInit},
-    [&](OpBuilder&, Location nestedLoc, Value outputN, ValueRange batchIterArgs, SmallVectorImpl<Value>& batchYielded) {
-      Value outputBatchAcc = batchIterArgs.front();
-      Value inputN =
-        buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), resultType.getDimSize(0), rewriter, nestedLoc);
-
-      auto channelLoop = buildNormalizedScfFor(
-        rewriter,
-        nestedLoc,
-        c0,
-        cOutputC,
-        c1,
-        ValueRange {outputBatchAcc},
-        [&](OpBuilder&,
-            Location channelLoc,
-            Value outputC,
-            ValueRange channelIterArgs,
-            SmallVectorImpl<Value>& channelYielded) {
-          Value outputChannelAcc = channelIterArgs.front();
-          Value inputC = buildNearestAsymmetricIndex(
-            outputC, inputType.getDimSize(1), resultType.getDimSize(1), rewriter, channelLoc);
-
-          auto heightLoop = buildNormalizedScfFor(
-            rewriter,
-            channelLoc,
-            c0,
-            cOutputH,
-            c1,
-            ValueRange {outputChannelAcc},
-            [&](OpBuilder&,
-                Location heightLoc,
-                Value outputH,
-                ValueRange heightIterArgs,
-                SmallVectorImpl<Value>& heightYielded) {
-              Value outputHeightAcc = heightIterArgs.front();
-              Value inputH = buildNearestAsymmetricIndex(
-                outputH, inputType.getDimSize(2), resultType.getDimSize(2), rewriter, heightLoc);
-
-              auto widthLoop = buildNormalizedScfFor(
-                rewriter,
-                heightLoc,
-                c0,
-                cOutputW,
-                c1,
-                ValueRange {outputHeightAcc},
-                [&](OpBuilder&,
-                    Location widthLoc,
-                    Value outputW,
-                    ValueRange widthIterArgs,
-                    SmallVectorImpl<Value>& widthYielded) {
-                  Value outputWidthAcc = widthIterArgs.front();
-                  Value inputW = buildNearestAsymmetricIndex(
-                    outputW, inputType.getDimSize(3), resultType.getDimSize(3), rewriter, widthLoc);
-
-                  SmallVector<OpFoldResult> inputOffsets = {inputN, inputC, inputH, inputW};
-                  Value inputSlice = tensor::ExtractSliceOp::create(
-                    rewriter, widthLoc, unitTensorType, input, inputOffsets, unitSizes, unitStrides);
-
-                  SmallVector<OpFoldResult> outputOffsets = {outputN, outputC, outputH, outputW};
-                  Value updatedOutput = tensor::InsertSliceOp::create(
-                    rewriter, widthLoc, inputSlice, outputWidthAcc, outputOffsets, unitSizes, unitStrides);
-                  widthYielded.push_back(updatedOutput);
-                  return success();
-                });
-              if (failed(widthLoop))
-                return failure();
-              heightYielded.push_back(widthLoop->results.front());
-              return success();
-            });
-          if (failed(heightLoop))
-            return failure();
-          channelYielded.push_back(heightLoop->results.front());
+static FailureOr<Value> buildDenseNearestResize(Value input,
+                                                RankedTensorType inputType,
+                                                RankedTensorType resultType,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  ArrayRef<int64_t> shape = resultType.getShape();
+  int64_t rowCount = shape[0] * shape[1] * shape[2];
+  auto scalarType = RankedTensorType::get({1, 1, 1, 1}, resultType.getElementType());
+  auto rowType = RankedTensorType::get({1, 1, 1, shape[3]}, resultType.getElementType());
+  auto rowsType = RankedTensorType::get({rowCount, 1, 1, 1, shape[3]}, resultType.getElementType());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, TypeRange {rowsType}, rowCount, {}, ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value outputN = affineFloorDivConst(rewriter, loc, args.lane, shape[1] * shape[2], anchor);
+      Value channelRow = affineModConst(rewriter, loc, args.lane, shape[1] * shape[2], anchor);
+      Value outputC = affineFloorDivConst(rewriter, loc, channelRow, shape[2], anchor);
+      Value outputH = affineModConst(rewriter, loc, channelRow, shape[2], anchor);
+      Value inputN = buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), shape[0], rewriter, loc);
+      Value inputC = buildNearestAsymmetricIndex(outputC, inputType.getDimSize(1), shape[1], rewriter, loc);
+      Value inputH = buildNearestAsymmetricIndex(outputH, inputType.getDimSize(2), shape[2], rewriter, loc);
+      Value row = tensor::EmptyOp::create(rewriter, loc, rowType.getShape(), rowType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value width = getOrCreateIndexConstant(rewriter, anchor, shape[3]);
+      auto loop = buildNormalizedScfFor(
+        rewriter, loc, c0, width, c1, ValueRange {row},
+        [&](OpBuilder&, Location nestedLoc, Value outputW, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value inputW = buildNearestAsymmetricIndex(
+            outputW, inputType.getDimSize(3), shape[3], rewriter, nestedLoc);
+          SmallVector<OpFoldResult> unitSizes(4, rewriter.getIndexAttr(1));
+          SmallVector<OpFoldResult> unitStrides(4, rewriter.getIndexAttr(1));
+          Value scalar = tensor::ExtractSliceOp::create(
+            rewriter, nestedLoc, scalarType, args.inputs.front(),
+            SmallVector<OpFoldResult> {inputN, inputC, inputH, inputW}, unitSizes, unitStrides);
+          yielded.push_back(tensor::InsertSliceOp::create(
+            rewriter, nestedLoc, scalar, iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                       rewriter.getIndexAttr(0), outputW},
+            unitSizes, unitStrides));
           return success();
         });
-      if (failed(channelLoop))
+      assert(succeeded(loop) && "nearest Resize row loop construction must succeed");
+      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
+    });
+  if (failed(batch))
+    return failure();
+
+  SmallVector<FragmentAssemblyEntry> entries;
+  entries.reserve(rowCount);
+  for (int64_t n = 0; n < shape[0]; ++n)
+    for (int64_t c = 0; c < shape[1]; ++c)
+      for (int64_t h = 0; h < shape[2]; ++h)
+        entries.push_back({(n * shape[1] + c) * shape[2] + h, 0, {n, c, h, 0}, {1, 1, 1, shape[3]}});
+  return createFragmentAssemblyBlueprint(
+    batch->getResult(0), resultType, entries, "dense_nchw", spatial::kContiguousRowMajorFragments, rewriter, loc);
+}
+
+static FailureOr<Value> buildRowStripNearestResize(
+    Value storage, RankedTensorType inputType, RankedTensorType resultType,
+    PatternRewriter& rewriter, Location loc) {
+  auto input = describeRowStripPhysicalValue(storage, inputType);
+  if (failed(input))
+    return failure();
+  int64_t tilesPerRow = input->tilesPerRow;
+  int64_t outputHeight = resultType.getDimSize(2);
+  int64_t outputWidth = resultType.getDimSize(3);
+  int64_t tileChannels = input->fragmentType.getDimSize(3);
+  int64_t laneCount = outputHeight * tilesPerRow;
+  auto outputFragmentType = RankedTensorType::get(
+    {1, 1, outputWidth, tileChannels}, resultType.getElementType());
+  auto outputStorageType = spatial::getGraphBatchPhysicalResultType(
+    laneCount, outputFragmentType);
+  auto pixelType = RankedTensorType::get(
+    {1, 1, 1, tileChannels}, resultType.getElementType());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, TypeRange {outputStorageType}, laneCount, {}, ValueRange {storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value outputRow = affineFloorDivConst(rewriter, loc, args.lane, tilesPerRow, anchor);
+      Value tile = affineModConst(rewriter, loc, args.lane, tilesPerRow, anchor);
+      Value inputRow = buildNearestAsymmetricIndex(
+        outputRow, inputType.getDimSize(2), outputHeight, rewriter, loc);
+      Value inputSlot = arith::AddIOp::create(
+        rewriter, loc, affineMulConst(rewriter, loc, inputRow, tilesPerRow, anchor), tile);
+      auto source = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), inputSlot, input->fragmentType);
+      if (failed(source))
         return failure();
-      batchYielded.push_back(channelLoop->results.front());
+      Value initial = tensor::EmptyOp::create(
+        rewriter, loc, outputFragmentType.getShape(), resultType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value width = getOrCreateIndexConstant(rewriter, anchor, outputWidth);
+      auto loop = buildNormalizedScfFor(
+        rewriter, loc, c0, width, c1, ValueRange {initial},
+        [&](OpBuilder&, Location nestedLoc, Value outputColumn, ValueRange iterArgs,
+            SmallVectorImpl<Value>& yielded) {
+          Value inputColumn = buildNearestAsymmetricIndex(
+            outputColumn, inputType.getDimSize(3), outputWidth, rewriter, nestedLoc);
+          Value pixel = tensor::ExtractSliceOp::create(
+            rewriter, nestedLoc, pixelType, *source,
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                      inputColumn, rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)},
+            getUnitStrides(rewriter, 4));
+          yielded.push_back(tensor::InsertSliceOp::create(
+            rewriter, nestedLoc, pixel, iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                      outputColumn, rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)},
+            getUnitStrides(rewriter, 4)));
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      publishGraphBatchPhysicalFragment(
+        rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
       return success();
     });
-  if (failed(batchLoop))
-    return failure();
-  return batchLoop->results.front();
+  return failed(batch) ? FailureOr<Value>(failure())
+                       : FailureOr<Value>(batch->getResult(0));
 }
 
 struct Resize : OpConversionPattern<ONNXResizeOp> {
@@ -161,23 +182,38 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
         || llvm::any_of(resultType.getShape(), [](int64_t dim) { return dim <= 0; }))
       return rewriter.notifyMatchFailure(resizeOp, "resize lowering requires positive static dimensions.");
 
-    auto computeOp = createSpatCompute<1>(
-      rewriter, resizeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getX(), [&](Value x) -> LogicalResult {
-        auto result = buildNearestResizeLoop(x, inputType, resultType, rewriter, resizeOp.getLoc());
-        if (failed(result))
-          return failure();
-        spatial::SpatYieldOp::create(rewriter, resizeOp.getLoc(), *result);
-        return success();
-      });
-    if (failed(computeOp))
-      return failure();
-    rewriter.replaceOp(resizeOp, computeOp->getResults());
+    auto plan = spatial::SpatResizeNearestPlanOp::create(
+      rewriter, resizeOp.getLoc(), resultType, adaptor.getX(), rewriter.getStringAttr("nchw"));
+    rewriter.replaceOp(resizeOp, plan.getResult());
     return success();
   }
 };
 
 } // namespace
 
+LogicalResult canLowerResizeNearestPlanToRowStrip(
+    spatial::SpatResizeNearestPlanOp planOp) {
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  return success(inputType && outputType && inputType.hasStaticShape()
+    && outputType.hasStaticShape() && inputType.getRank() == 4
+    && outputType.getRank() == 4 && inputType.getDimSize(0) == 1
+    && outputType.getDimSize(0) == 1
+    && inputType.getDimSize(1) == outputType.getDimSize(1));
+}
+
+FailureOr<Value> lowerSelectedResizeNearestPlan(
+    spatial::SpatResizeNearestPlanOp planOp, std::optional<Value> rowStripInput,
+    PatternRewriter& rewriter) {
+  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
+  if (rowStripInput)
+    return buildRowStripNearestResize(
+      *rowStripInput, inputType, outputType, rewriter, planOp.getLoc());
+  return buildDenseNearestResize(
+    planOp.getInput(), inputType, outputType, rewriter, planOp.getLoc());
+}
+
 void populateResizePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Resize>(ctx); }
 
 } // namespace onnx_mlir

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

@@ -61,6 +61,74 @@ static FailureOr<Value> materializeTransposedConstant(Value input,
                              resultType);
 }
 
+static FailureOr<Value> transposeFragmentAssemblyBlueprint(spatial::SpatBlueprintOp blueprint,
+                                                           RankedTensorType resultType,
+                                                           ArrayRef<int64_t> permutation,
+                                                           ConversionPatternRewriter& rewriter,
+                                                           Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(blueprint.getInput().getType());
+  auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+  auto fragmentStrides = blueprint.getFragmentStrides();
+  if (!storageType || !storageType.hasStaticShape() || !resultType.hasStaticShape()
+      || !blueprint.getFragments().empty() || blueprint.getMode() != "fragment_assembly"
+      || !blueprint.getFragmentOperandIndices() || !sourceOffsets || !fragmentStrides
+      || llvm::any_of(*sourceOffsets, [](int64_t offset) { return offset != 0; })
+      || storageType.getRank() != resultType.getRank() + 1)
+    return failure();
+  if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments
+      && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+
+  SmallVector<int64_t> outputStorageShape {storageType.getDimSize(0)};
+  for (int64_t sourceDim : permutation)
+    outputStorageShape.push_back(storageType.getDimSize(sourceDim + 1));
+  auto outputStorageType = RankedTensorType::get(outputStorageShape, storageType.getElementType());
+  auto mapped = mapGraphBatchFragments(
+    blueprint.getInput(), outputStorageType, rewriter, loc, [&](Value fragment, RankedTensorType fragmentType) {
+      Value init = createTransposeInit(fragment, fragmentType, permutation, rewriter, loc);
+      return FailureOr<Value>(
+        linalg::TransposeOp::create(rewriter, loc, fragment, init, permutation).getResult()[0]);
+    });
+  if (failed(mapped))
+    return failure();
+
+  const int64_t rank = resultType.getRank();
+  const int64_t fragmentCount = blueprint.getFragmentOperandIndices()->size();
+  SmallVector<int64_t> offsets, sizes, strides;
+  offsets.reserve(fragmentCount * rank);
+  sizes.reserve(fragmentCount * rank);
+  strides.reserve(fragmentCount * rank);
+  ArrayRef<int64_t> inputOffsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> inputSizes = blueprint.getFragmentSizes();
+  for (int64_t fragment = 0; fragment < fragmentCount; ++fragment)
+    for (int64_t sourceDim : permutation) {
+      const int64_t index = fragment * rank + sourceDim;
+      offsets.push_back(inputOffsets[index]);
+      sizes.push_back(inputSizes[index]);
+      strides.push_back((*fragmentStrides)[index]);
+    }
+  auto transposedBlueprint = spatial::SpatBlueprintOp::create(rewriter,
+                                                               loc,
+                                                               resultType,
+                                                               *mapped,
+                                                               ValueRange {},
+                                                               blueprint.getLogicalLayoutAttr(),
+                                                               rewriter.getStringAttr("fragmented"),
+                                                               rewriter.getDenseI64ArrayAttr(offsets),
+                                                               rewriter.getDenseI64ArrayAttr(sizes),
+                                                               rewriter.getStringAttr("permuted_fragments"),
+                                                               blueprint.getModeAttr(),
+                                                               blueprint.getFragmentOperandIndicesAttr(),
+                                                               blueprint.getFragmentSourceSlotsAttr(),
+                                                               blueprint.getFragmentSourceOffsetsAttr(),
+                                                               rewriter.getDenseI64ArrayAttr(strides),
+                                                               blueprint.getConflictPolicyAttr(),
+                                                               blueprint.getCoveragePolicyAttr());
+  if (spatial::isCanonicalContiguousRowMajorFragmentAssembly(transposedBlueprint))
+    transposedBlueprint.setIndexMapAttr(rewriter.getStringAttr(spatial::kContiguousRowMajorFragments));
+  return transposedBlueprint.getOutput();
+}
+
 struct TransposeToLinalgTranspose : OpConversionPattern<ONNXTransposeOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -75,6 +143,14 @@ struct TransposeToLinalgTranspose : OpConversionPattern<ONNXTransposeOp> {
     auto permutation = getTransposePermutationChecked(transposeOp.getPermAttr(), inputType.getRank());
     if (failed(permutation))
       return failure();
+    if (auto blueprint = adaptor.getData().getDefiningOp<spatial::SpatBlueprintOp>()) {
+      auto transposed =
+        transposeFragmentAssemblyBlueprint(blueprint, resultType, *permutation, rewriter, transposeOp.getLoc());
+      if (succeeded(transposed)) {
+        rewriter.replaceOp(transposeOp, *transposed);
+        return success();
+      }
+    }
     if (isCompileTimeComputable(adaptor.getData())) {
       auto constantTranspose =
         materializeTransposedConstant(adaptor.getData(), resultType, *permutation, rewriter, transposeOp.getLoc());

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -20,6 +20,14 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
 mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp);
 mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp);
 
+mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
+  spatial::SpatResizeNearestPlanOp planOp);
+
+mlir::FailureOr<mlir::Value> lowerSelectedResizeNearestPlan(
+  spatial::SpatResizeNearestPlanOp planOp,
+  std::optional<mlir::Value> rowStripInput,
+  mlir::PatternRewriter& rewriter);
+
 mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp);
 
 mlir::FailureOr<mlir::Value>

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -36,6 +36,8 @@ static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, Selected
     return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
   if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(user))
     return getSelectedLayout(layouts, siluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  if (auto resizePlan = dyn_cast<spatial::SpatResizeNearestPlanOp>(user))
+    return getSelectedLayout(layouts, resizePlan.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))
@@ -66,6 +68,8 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, Selecte
 static bool canConsumeRowStripAsUser(Operation* user) {
   if (isa<spatial::SpatReluPlanOp, spatial::SpatSiluPlanOp>(user))
     return true;
+  if (auto resizePlan = dyn_cast<spatial::SpatResizeNearestPlanOp>(user))
+    return succeeded(canLowerResizeNearestPlanToRowStrip(resizePlan));
   if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
     auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
     return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
@@ -117,6 +121,14 @@ static SelectedLayout chooseActivationLayout(Value input,
   return SelectedLayout::PixelMajorRowStrip;
 }
 
+static SelectedLayout chooseResizeLayout(
+    spatial::SpatResizeNearestPlanOp resizePlan,
+    llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  return getSelectedLayout(layouts, resizePlan.getInput()) == SelectedLayout::PixelMajorRowStrip
+      && succeeded(canLowerResizeNearestPlanToRowStrip(resizePlan))
+    ? SelectedLayout::PixelMajorRowStrip : SelectedLayout::DenseNchw;
+}
+
 static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
                                           llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
@@ -255,6 +267,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto resizePlan = dyn_cast<spatial::SpatResizeNearestPlanOp>(&op)) {
+          SelectedLayout selected = chooseResizeLayout(resizePlan, layouts);
+          if (layouts[resizePlan.getResult()] != selected) {
+            layouts[resizePlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
         if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
           SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
           if (layouts[biasAddPlan.getResult()] != selected) {
@@ -312,6 +332,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
         producedValue = reluPlan.getResult();
       else if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op))
         producedValue = siluPlan.getResult();
+      else if (auto resizePlan = dyn_cast<spatial::SpatResizeNearestPlanOp>(&op))
+        producedValue = resizePlan.getResult();
       else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
         producedValue = maxPoolPlan.getResult();
       else if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op))

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -129,6 +129,32 @@ LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprin
   return success();
 }
 
+FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(RewriterBase& rewriter,
+                                                          Location loc,
+                                                          mlir::Value source,
+                                                          RankedTensorType resultType) {
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || resultType.getRank() < 2
+      || sourceType.getRank() != resultType.getRank() + 1 || sourceType.getElementType() != resultType.getElementType()
+      || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getDimSize(0) != getStaticShapeElementCount(resultType.getShape().drop_back())
+      || sourceType.getDimSize(sourceType.getRank() - 1) != resultType.getDimSize(resultType.getRank() - 1)
+      || llvm::any_of(sourceType.getShape().slice(1, sourceType.getRank() - 2), [](int64_t dim) { return dim != 1; }))
+    return failure();
+
+  SmallVector<ReassociationIndices> collapse {{}, {sourceType.getRank() - 1}};
+  for (int64_t dim = 0; dim < sourceType.getRank() - 1; ++dim)
+    collapse.front().push_back(dim);
+  auto flatType = RankedTensorType::get(
+    {sourceType.getDimSize(0), sourceType.getDimSize(sourceType.getRank() - 1)}, resultType.getElementType());
+  mlir::Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, source, collapse);
+
+  SmallVector<ReassociationIndices> expand {{}, {resultType.getRank() - 1}};
+  for (int64_t dim = 0; dim < resultType.getRank() - 1; ++dim)
+    expand.front().push_back(dim);
+  return tensor::ExpandShapeOp::create(rewriter, loc, resultType, flat, expand).getResult();
+}
+
 static SmallVector<int64_t, 4> expandFlatElementIndex(int64_t flatIndex, ArrayRef<int64_t> shape) {
   SmallVector<int64_t, 4> indices(shape.size(), 0);
   for (int64_t dim = static_cast<int64_t>(shape.size()) - 1; dim >= 0; --dim) {

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -51,6 +51,11 @@ mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlu
                                                      llvm::ArrayRef<int64_t> flatSizes,
                                                      llvm::ArrayRef<int64_t> flatStrides);
 
+mlir::FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(mlir::RewriterBase& rewriter,
+                                                                mlir::Location loc,
+                                                                mlir::Value source,
+                                                                mlir::RankedTensorType resultType);
+
 mlir::FailureOr<mlir::SmallVector<int64_t, 4>>
 getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor,
                                       mlir::ShapedType sourceType,

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -71,6 +71,16 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
                                               flatStrides)))
     return failure();
 
+  if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments) {
+    if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+      return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+    Value source = mapping.lookupOrDefault(blueprint.getInput());
+    auto reshaped = reshapeContiguousRowMajorFragments(
+      rewriter, blueprint.getLoc(), source, cast<RankedTensorType>(resultType));
+    if (failed(reshaped))
+      return blueprint.emitOpError("contiguous row-major fragment storage does not match its logical result"), failure();
+    return *reshaped;
+  }
   SmallVector<int64_t> hostStrides = computeRowMajorStrides(resultType.getShape());
   SmallVector<FragmentAssemblyCopy, 8> copies;
   for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -49,6 +49,16 @@ struct LowerFragmentAssemblyBlueprintPattern
           op, rank, fragmentOperands.size(), operandIndices, sourceOffsets, flatOffsets, flatSizes, flatStrides)))
       return failure();
 
+    if (op.getIndexMap() == spatial::kContiguousRowMajorFragments) {
+      if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(op))
+        return op.emitOpError("contiguous row-major fragment physical source order or storage is not canonical");
+      auto reshaped = reshapeContiguousRowMajorFragments(
+        rewriter, op.getLoc(), adaptor.getInput(), cast<RankedTensorType>(resultType));
+      if (failed(reshaped))
+        return op.emitOpError("contiguous row-major fragment storage does not match its logical result");
+      rewriter.replaceOp(op, *reshaped);
+      return success();
+    }
     Value currentOutput =
       tensor::EmptyOp::create(rewriter, op.getLoc(), resultType.getShape(), resultType.getElementType()).getResult();
     for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {

src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp

@@ -532,54 +532,74 @@ struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
   }
 };
 
-static bool isOne(Attribute value) {
-  if (auto floatValue = dyn_cast<FloatAttr>(value))
-    return floatValue.getValue().isExactlyValue(1.0);
-  if (auto integerValue = dyn_cast<IntegerAttr>(value))
-    return integerValue.getValue() == 1;
-  return false;
+enum class MultiplicationConstant { Other, Zero, One };
+
+static MultiplicationConstant classifyMultiplicationConstant(Attribute value) {
+  if (auto floatValue = dyn_cast<FloatAttr>(value)) {
+    const APFloat& number = floatValue.getValue();
+    if (number.isZero() && !number.isNegative())
+      return MultiplicationConstant::Zero;
+    if (number.isExactlyValue(1.0))
+      return MultiplicationConstant::One;
+  }
+  if (auto integerValue = dyn_cast<IntegerAttr>(value)) {
+    if (integerValue.getValue().isZero())
+      return MultiplicationConstant::Zero;
+    if (integerValue.getValue() == 1)
+      return MultiplicationConstant::One;
+  }
+  return MultiplicationConstant::Other;
 }
 
-static bool isAllOneHostCopy(pim::PimMemCopyHostToDevOp copyOp, ModuleOp moduleOp, MemRefType copiedType) {
+static MultiplicationConstant classifyUniformHostCopy(
+  pim::PimMemCopyHostToDevOp copyOp, ModuleOp moduleOp, MemRefType copiedType) {
   auto targetOffset = resolveIndexValue(copyOp.getDeviceTargetOffset());
   auto sourceOffset = resolveIndexValue(copyOp.getHostSourceOffset());
   if (failed(targetOffset) || failed(sourceOffset) || *targetOffset != 0)
-    return false;
+    return MultiplicationConstant::Other;
 
   Type elementType = copiedType.getElementType();
   if (!elementType.isIntOrFloat())
-    return false;
+    return MultiplicationConstant::Other;
   unsigned bitWidth = elementType.getIntOrFloatBitWidth();
   if (bitWidth == 0 || bitWidth % 8 != 0)
-    return false;
+    return MultiplicationConstant::Other;
 
   int64_t elementBytes = bitWidth / 8;
   int64_t copiedElements = copiedType.getNumElements();
   if (*sourceOffset % elementBytes != 0 || copyOp.getSize() != copiedElements * elementBytes)
-    return false;
+    return MultiplicationConstant::Other;
 
   auto source = getDenseGlobalValue(moduleOp, copyOp.getHostSource());
   if (failed(source) || source->getElementType() != elementType)
-    return false;
+    return MultiplicationConstant::Other;
 
   int64_t firstElement = *sourceOffset / elementBytes;
   int64_t endElement = firstElement + copiedElements;
   if (firstElement < 0 || endElement > source->getNumElements())
-    return false;
+    return MultiplicationConstant::Other;
   if (source->isSplat())
-    return isOne(source->getSplatValue<Attribute>());
+    return classifyMultiplicationConstant(source->getSplatValue<Attribute>());
 
+  MultiplicationConstant classification = MultiplicationConstant::Other;
   int64_t index = 0;
   for (Attribute value : source->getValues<Attribute>()) {
-    if (index >= firstElement && index < endElement && !isOne(value))
-      return false;
+    if (index >= firstElement && index < endElement) {
+      MultiplicationConstant current = classifyMultiplicationConstant(value);
+      if (current == MultiplicationConstant::Other)
+        return current;
+      if (classification == MultiplicationConstant::Other)
+        classification = current;
+      else if (classification != current)
+        return MultiplicationConstant::Other;
+    }
     if (++index >= endElement)
       break;
   }
-  return true;
+  return classification;
 }
 
-struct FoldMultiplyByOnePattern final : OpRewritePattern<pim::PimVVMulOp> {
+struct FoldMultiplyByConstantPattern final : OpRewritePattern<pim::PimVVMulOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(pim::PimVVMulOp mulOp, PatternRewriter& rewriter) const override {
@@ -605,14 +625,19 @@ struct FoldMultiplyByOnePattern final : OpRewritePattern<pim::PimVVMulOp> {
         copyOp = candidate;
       }
       auto maskType = dyn_cast<MemRefType>(mask.getType());
-      if (!copyOp || !copyOp.use_empty() || !maskType || !isAllOneHostCopy(copyOp, moduleOp, maskType))
+      if (!copyOp || !copyOp.use_empty() || !maskType)
+        continue;
+      MultiplicationConstant constant = classifyUniformHostCopy(copyOp, moduleOp, maskType);
+      if (constant == MultiplicationConstant::Other)
         continue;
 
       auto outputAlloc = mulOp.getOutputBuffer().getDefiningOp<memref::AllocOp>();
-      rewriter.replaceOp(mulOp, input);
-      rewriter.eraseOp(copyOp);
-      if (maskAlloc.use_empty())
-        rewriter.eraseOp(maskAlloc);
+      rewriter.replaceOp(mulOp, constant == MultiplicationConstant::One ? input : mask);
+      if (constant == MultiplicationConstant::One) {
+        rewriter.eraseOp(copyOp);
+        if (maskAlloc.use_empty())
+          rewriter.eraseOp(maskAlloc);
+      }
       if (outputAlloc && outputAlloc.use_empty())
         rewriter.eraseOp(outputAlloc);
       return success();
@@ -629,7 +654,7 @@ void populateConstantFoldingConstantPatterns(RewritePatternSet& patterns) {
                FoldConstantCoreMapPattern,
                FoldConstantHostCopyPattern,
                FoldConstantMemCpPattern,
-               FoldMultiplyByOnePattern>(patterns.getContext());
+               FoldMultiplyByConstantPattern>(patterns.getContext());
 }
 
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Spatial.td

@@ -303,6 +303,21 @@ def SpatSiluPlanOp : SpatOp<"silu_plan", []> {
   let hasVerifier = 1;
 }
 
+def SpatResizeNearestPlanOp : SpatOp<"resize_nearest_plan", []> {
+  let summary = "Layout-aware nearest asymmetric Resize 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/SpatialOps.cpp

@@ -11,6 +11,70 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace spatial {
 
+bool hasCanonicalContiguousRowMajorFragments(RankedTensorType logicalType,
+                                             ArrayRef<int64_t> offsets,
+                                             ArrayRef<int64_t> sizes,
+                                             ArrayRef<int64_t> strides) {
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() <= 0
+      || logicalType.getDimSize(logicalType.getRank() - 1) <= 0)
+    return false;
+  const int64_t rank = logicalType.getRank();
+  const int64_t rowCount = logicalType.getNumElements() / logicalType.getDimSize(rank - 1);
+  if (offsets.size() != static_cast<size_t>(rowCount * rank) || sizes.size() != offsets.size()
+      || strides.size() != offsets.size())
+    return false;
+  for (int64_t row = 0; row < rowCount; ++row) {
+    int64_t remaining = row;
+    for (int64_t dim = rank - 2; dim >= 0; --dim) {
+      const int64_t index = row * rank + dim;
+      if (offsets[index] != remaining % logicalType.getDimSize(dim) || sizes[index] != 1 || strides[index] != 1)
+        return false;
+      remaining /= logicalType.getDimSize(dim);
+    }
+    const int64_t last = row * rank + rank - 1;
+    if (offsets[last] != 0 || sizes[last] != logicalType.getDimSize(rank - 1) || strides[last] != 1)
+      return false;
+  }
+  return true;
+}
+
+bool isCanonicalContiguousRowMajorFragmentAssembly(SpatBlueprintOp blueprint) {
+  auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  auto physicalType = dyn_cast<RankedTensorType>(blueprint.getInput().getType());
+  auto operandIndices = blueprint.getFragmentOperandIndices();
+  auto sourceSlots = blueprint.getFragmentSourceSlots();
+  auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+  auto fragmentStrides = blueprint.getFragmentStrides();
+  if (!logicalType || !physicalType || !logicalType.hasStaticShape() || !physicalType.hasStaticShape()
+      || logicalType.getRank() < 2 || !blueprint.getFragments().empty()
+      || blueprint.getMode() != "fragment_assembly" || !operandIndices || !sourceSlots || !sourceOffsets
+      || !fragmentStrides)
+    return false;
+
+  ArrayRef<int64_t> offsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> sizes = blueprint.getFragmentSizes();
+  if (!hasCanonicalContiguousRowMajorFragments(logicalType, offsets, sizes, *fragmentStrides)
+      || operandIndices->empty() || operandIndices->size() != sourceSlots->size()
+      || operandIndices->size() != sourceOffsets->size()
+      || operandIndices->size() * static_cast<size_t>(logicalType.getRank()) != offsets.size()
+      || physicalType.getRank() != logicalType.getRank() + 1
+      || physicalType.getDimSize(0) != static_cast<int64_t>(operandIndices->size())
+      || physicalType.getDimSize(0)
+           != logicalType.getNumElements() / logicalType.getDimSize(logicalType.getRank() - 1)
+      || physicalType.getElementType() != logicalType.getElementType()
+      || physicalType.getNumElements() != logicalType.getNumElements()
+      || physicalType.getDimSize(physicalType.getRank() - 1) != logicalType.getDimSize(logicalType.getRank() - 1)
+      || llvm::any_of(physicalType.getShape().slice(1, physicalType.getRank() - 2),
+                     [](int64_t dim) { return dim != 1; }))
+    return false;
+
+  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices))
+    if (operandIndex != 0 || (*sourceSlots)[fragmentIndex] != static_cast<int64_t>(fragmentIndex)
+        || (*sourceOffsets)[fragmentIndex] != 0)
+      return false;
+  return true;
+}
+
 RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, RankedTensorType fragmentType) {
   SmallVector<int64_t> shape {laneCount};
   llvm::append_range(shape, fragmentType.getShape());

src/PIM/Dialect/Spatial/SpatialOps.hpp

@@ -30,6 +30,15 @@
 namespace onnx_mlir {
 namespace spatial {
 
+inline constexpr llvm::StringLiteral kContiguousRowMajorFragments = "contiguous_row_major_fragments";
+
+bool hasCanonicalContiguousRowMajorFragments(mlir::RankedTensorType logicalType,
+                                             llvm::ArrayRef<int64_t> offsets,
+                                             llvm::ArrayRef<int64_t> sizes,
+                                             llvm::ArrayRef<int64_t> strides);
+
+bool isCanonicalContiguousRowMajorFragmentAssembly(SpatBlueprintOp blueprint);
+
 mlir::RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, mlir::RankedTensorType fragmentType);
 mlir::FailureOr<mlir::RankedTensorType>
 getGraphBatchFragmentType(mlir::RankedTensorType physicalType, int64_t expectedLaneCount);

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -480,6 +480,23 @@ LogicalResult SpatSiluPlanOp::verify() {
   return success();
 }
 
+LogicalResult SpatResizeNearestPlanOp::verify() {
+  if (failed(verifyPlanTensorTypes(
+        getOperation(), getInput(), getOutput(), "spat.resize_nearest_plan")))
+    return failure();
+  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!inputType.hasStaticShape() || !outputType.hasStaticShape()
+      || inputType.getRank() != 4 || outputType.getRank() != 4)
+    return emitError("requires static rank-4 input and output tensors");
+  if (getLogicalLayout() != "nchw")
+    return emitError("requires logical layout \"nchw\"");
+  if (llvm::any_of(inputType.getShape(), [](int64_t dim) { return dim <= 0; })
+      || llvm::any_of(outputType.getShape(), [](int64_t dim) { return dim <= 0; }))
+    return emitError("requires positive dimensions");
+  return success();
+}
+
 LogicalResult SpatMaxPool2DPlanOp::verify() {
   if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.max_pool2d_plan")))
     return failure();
@@ -587,8 +604,12 @@ LogicalResult SpatBlueprintOp::verify() {
   if (offsets.size() != sizes.size())
     return emitError("fragment offset and size arrays must have the same length");
   int64_t rank = logicalType.getRank();
-  if (offsets.empty())
+  bool isContiguousRowMajor = getIndexMap() == kContiguousRowMajorFragments;
+  if (offsets.empty()) {
+    if (isContiguousRowMajor)
+      return emitError("contiguous row-major fragment destination geometry is not canonical");
     return success();
+  }
   if (rank <= 0 || offsets.size() % rank != 0)
     return emitError("fragment metadata must be a whole number of rank-sized fragments");
 
@@ -611,6 +632,8 @@ LogicalResult SpatBlueprintOp::verify() {
   };
 
   if (!isFragmentAssembly) {
+    if (isContiguousRowMajor)
+      return emitError("contiguous row-major fragments require fragment assembly metadata");
     if (failed(verifyBoundsOnly({})))
       return failure();
     if (!getFragments().empty())
@@ -662,6 +685,13 @@ LogicalResult SpatBlueprintOp::verify() {
   if (failed(verifyBoundsOnly(strides)))
     return failure();
 
+  if (isContiguousRowMajor) {
+    if (!hasCanonicalContiguousRowMajorFragments(logicalType, offsets, sizes, strides))
+      return emitError("contiguous row-major fragment destination geometry is not canonical");
+    if (!isCanonicalContiguousRowMajorFragmentAssembly(*this))
+      return emitError("contiguous row-major fragment physical source order or storage is not canonical");
+  }
+
   SmallVector<std::pair<SmallVector<int64_t, 4>, SmallVector<int64_t, 4>>, 8> slices;
   slices.reserve(static_cast<size_t>(fragmentCount));
   SmallVector<int64_t, 8> fragmentCountsByOperand(static_cast<size_t>(operandCount), 0);
@@ -713,20 +743,22 @@ LogicalResult SpatBlueprintOp::verify() {
       if (sourceSliceOffsets[dim] + fragmentSizes[dim] > fragmentType.getDimSize(dim))
         return emitError("fragment assembly source offset must describe a valid unit-stride slice");
 
-    for (const auto& [existingOffsets, existingSizes] : slices) {
-      bool overlaps = true;
-      for (int64_t dim = 0; dim < rank; ++dim) {
-        int64_t begin = fragmentOffsets[dim];
-        int64_t end = begin + fragmentSizes[dim];
-        int64_t existingBegin = existingOffsets[dim];
-        int64_t existingEnd = existingBegin + existingSizes[dim];
-        if (end <= existingBegin || existingEnd <= begin) {
-          overlaps = false;
-          break;
+    if (!isContiguousRowMajor) {
+      for (const auto& [existingOffsets, existingSizes] : slices) {
+        bool overlaps = true;
+        for (int64_t dim = 0; dim < rank; ++dim) {
+          int64_t begin = fragmentOffsets[dim];
+          int64_t end = begin + fragmentSizes[dim];
+          int64_t existingBegin = existingOffsets[dim];
+          int64_t existingEnd = existingBegin + existingSizes[dim];
+          if (end <= existingBegin || existingEnd <= begin) {
+            overlaps = false;
+            break;
+          }
         }
+        if (overlaps)
+          return emitError("fragment assembly blueprint requires disjoint static slices");
       }
-      if (overlaps)
-        return emitError("fragment assembly blueprint requires disjoint static slices");
     }
     slices.push_back({std::move(fragmentOffsets), std::move(fragmentSizes)});
   }

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

@@ -9,6 +9,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
@@ -37,6 +38,63 @@ static SmallVector<Value> getBlueprintFragments(SpatBlueprintOp blueprint) {
   return fragments;
 }
 
+static FailureOr<Value> buildContiguousRowMajorReconstruction(
+    OpBuilder &builder, Location loc, SpatBlueprintOp blueprint,
+    Value source) {
+  auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  if (!resultType || !sourceType || !resultType.hasStaticShape()
+      || !sourceType.hasStaticShape() || resultType.getRank() <= 0
+      || sourceType.getRank() != resultType.getRank() + 1)
+    return failure();
+  int64_t rank = resultType.getRank();
+  int64_t width = resultType.getDimSize(rank - 1);
+  int64_t rowCount = resultType.getNumElements() / width;
+  if (sourceType.getDimSize(0) != rowCount || sourceType.getDimSize(rank) != width)
+    return failure();
+
+  auto flatType = RankedTensorType::get({rowCount, width}, resultType.getElementType());
+  auto rowType = RankedTensorType::get({1, width}, resultType.getElementType());
+  auto physicalRowType = RankedTensorType::get(sourceType.getShape().drop_front(), resultType.getElementType());
+  Value init = tensor::EmptyOp::create(builder, loc, flatType.getShape(), flatType.getElementType());
+  Value c0 = arith::ConstantIndexOp::create(builder, loc, 0);
+  Value c1 = arith::ConstantIndexOp::create(builder, loc, 1);
+  Value rows = arith::ConstantIndexOp::create(builder, loc, rowCount);
+  auto loop = buildNormalizedScfFor(
+    builder, loc, c0, rows, c1, ValueRange {init},
+    [&](OpBuilder &nested, Location nestedLoc, Value row, ValueRange iterArgs,
+        SmallVectorImpl<Value> &yielded) {
+      SmallVector<OpFoldResult> offsets {row};
+      SmallVector<OpFoldResult> sizes {nested.getIndexAttr(1)};
+      SmallVector<OpFoldResult> strides {nested.getIndexAttr(1)};
+      for (int64_t dim : physicalRowType.getShape()) {
+        offsets.push_back(nested.getIndexAttr(0));
+        sizes.push_back(nested.getIndexAttr(dim));
+        strides.push_back(nested.getIndexAttr(1));
+      }
+      Value physicalRow = tensor::ExtractSliceOp::create(
+        nested, nestedLoc, physicalRowType, source, offsets, sizes, strides);
+      SmallVector<ReassociationIndices> collapse {{}};
+      for (int64_t dim = 0; dim < rank - 1; ++dim)
+        collapse.front().push_back(dim);
+      collapse.push_back({rank - 1});
+      Value flatRow = tensor::CollapseShapeOp::create(
+        nested, nestedLoc, rowType, physicalRow, collapse);
+      yielded.push_back(tensor::InsertSliceOp::create(
+        nested, nestedLoc, flatRow, iterArgs.front(),
+        SmallVector<OpFoldResult> {row, nested.getIndexAttr(0)},
+        SmallVector<OpFoldResult> {nested.getIndexAttr(1), nested.getIndexAttr(width)},
+        SmallVector<OpFoldResult> {nested.getIndexAttr(1), nested.getIndexAttr(1)}));
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  SmallVector<ReassociationIndices> expand {{}, {rank - 1}};
+  for (int64_t dim = 0; dim < rank - 1; ++dim)
+    expand.front().push_back(dim);
+  return tensor::ExpandShapeOp::create(builder, loc, resultType, loop->results.front(), expand).getResult();
+}
+
 static FailureOr<Value> buildBlueprintReconstruction(
     OpBuilder &builder, Location loc, SpatBlueprintOp blueprint,
     ValueRange sourceBlockArgs) {
@@ -57,6 +115,13 @@ static FailureOr<Value> buildBlueprintReconstruction(
       sourceOffsets->size() != operandIndices->size())
     return blueprint.emitOpError("phase 1 fragment assembly metadata has inconsistent sizes"), failure();
 
+  if (blueprint.getIndexMap() == kContiguousRowMajorFragments) {
+    if (!isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+      return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+    if (sourceBlockArgs.size() != 1)
+      return blueprint.emitOpError("contiguous row-major fragment reconstruction requires one physical source"), failure();
+    return buildContiguousRowMajorReconstruction(builder, loc, blueprint, sourceBlockArgs.front());
+  }
   Value result = tensor::EmptyOp::create(builder, loc, resultType.getShape(),
                                          resultType.getElementType());
   for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
@@ -178,6 +243,15 @@ static Operation *getTopLevelDeferredOperation(
   return op && isTopLevelDeferredOperation(op, body, plan) ? op : nullptr;
 }
 
+static bool isDefinedInside(Operation *owner, Value value) {
+  if (Operation *definition = value.getDefiningOp())
+    return owner->isProperAncestor(definition);
+  auto argument = dyn_cast<BlockArgument>(value);
+  Region *region = argument ? argument.getOwner()->getParent() : nullptr;
+  return region == &owner->getRegion(0)
+      || (region && owner->getRegion(0).isAncestor(region));
+}
+
 static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
                        llvm::SmallPtrSetImpl<Operation *> &seen) {
   if (value == plan.graphInput || value == plan.graphLane || value == plan.scheduledLane)
@@ -195,18 +269,10 @@ static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
     loop.getRegion().walk([&](Operation *nested) {
       if (isa<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))
+      for (Value operand : nested->getOperands())
+        if (!isDefinedInside(loop, operand)
+            && !isEligible(operand, body, plan, seen))
           eligible = false;
-      }
     });
     if (!eligible)
       return false;
@@ -267,19 +333,9 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
     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))
+        for (Value operand : nested->getOperands())
+          if (!isDefinedInside(loop, operand) && !mapping.contains(operand))
             captures.push_back(operand);
-        }
       });
       for (Value capture : captures)
         if (!mapping.contains(capture) && failed(clone(capture)))
@@ -304,7 +360,10 @@ static bool dependsOnGraphLane(Value value, Value graphLane, Block &body,
   if (auto loop = dyn_cast<scf::ForOp>(op)) {
     bool depends = false;
     loop.getRegion().walk([&](Operation *nested) {
-      depends |= llvm::is_contained(nested->getOperands(), graphLane);
+      for (Value operand : nested->getOperands())
+        if (!isDefinedInside(loop, operand)
+            && dependsOnGraphLane(operand, graphLane, body, plan, seen))
+          depends = true;
     });
     if (depends)
       return true;

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

@@ -14,6 +14,26 @@ namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
 
+static LogicalResult verifyNoEscapingRegionValues(Operation* owner, StringRef phase) {
+  Operation* escapingDefinition = nullptr;
+  Operation* escapingUser = nullptr;
+  owner->walk([&](Operation* nested) {
+    for (Value result : nested->getResults())
+      for (Operation* user : result.getUsers())
+        if (!owner->isProperAncestor(user)) {
+          escapingDefinition = nested;
+          escapingUser = user;
+          return WalkResult::interrupt();
+        }
+    return WalkResult::advance();
+  });
+  if (!escapingDefinition)
+    return success();
+  return owner->emitOpError() << phase << " left a value defined by " << escapingDefinition->getName()
+                              << " at " << escapingDefinition->getLoc() << " captured by "
+                              << escapingUser->getName() << " at " << escapingUser->getLoc();
+}
+
 static LogicalResult placeLogicalProcessorsOnPhysicalCores(DeferredTransferPlan& plan, const SchedulingTarget& target) {
   std::vector<Cost> logicalTrafficFlits(target.processorCount * target.processorCount, 0);
   for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
@@ -138,6 +158,8 @@ static LogicalResult eraseOldGraph(func::FuncOp funcOp, IRRewriter& rewriter) {
         }
       }
     }
+    if (failed(verifyNoEscapingRegionValues(op, "phase 2")))
+      return failure();
     rewriter.eraseOp(op);
   }
   return success();
@@ -217,6 +239,8 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
     op->getResult(0).replaceAllUsesWith(replacement);
     if (!op->use_empty())
       return op->emitOpError("phase 2 cannot erase deferred communication with live uses");
+    if (failed(verifyNoEscapingRegionValues(op, "phase 2 deferred communication")))
+      return failure();
     rewriter.eraseOp(op);
   }
   if (failed(eraseDeferredSourceSelectors(funcOp, rewriter)) || failed(eraseOldGraph(funcOp, rewriter))

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

@@ -180,13 +180,18 @@ static bool originatesFromDeferredSource(
   return originatesFromDeferredSource(value, deferred, visited);
 }
 
-static bool isInsideDeferredLoop(Operation *op,
-                                 SpatDeferredCommunicationOp deferred) {
+static scf::ForOp getEnclosingDeferredLoop(
+    Operation *op, SpatDeferredCommunicationOp deferred) {
   for (Operation *parent = op->getParentOp(); parent && parent != deferred;
        parent = parent->getParentOp())
-    if (isa<scf::ForOp>(parent))
-      return true;
-  return false;
+    if (auto loop = dyn_cast<scf::ForOp>(parent))
+      return loop;
+  return {};
+}
+
+static bool isInsideDeferredLoop(
+    Operation *op, SpatDeferredCommunicationOp deferred) {
+  return static_cast<bool>(getEnclosingDeferredLoop(op, deferred));
 }
 
 static FailureOr<unsigned> getLoopIterationCount(
@@ -297,7 +302,7 @@ static LogicalResult validateDeferredProgram(
         && llvm::any_of(op->getOperands(), [&](Value operand) {
              return originatesFromDeferredSource(operand, deferred);
            })) {
-      auto loop = op->getParentOfType<scf::ForOp>();
+      auto loop = getEnclosingDeferredLoop(op, deferred);
       scf::ForOp outerLoop;
       for (Operation *parent = loop ? loop->getParentOp() : nullptr;
            parent && parent != deferred; parent = parent->getParentOp())
@@ -578,7 +583,7 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
             SmallVector<OpFoldResult>(
               ArrayRef(slice.getMixedStrides()).drop_front())};
         leaf.reconstructedType = cast<RankedTensorType>(value.getType());
-        leaf.enclosingLoop = slice->getParentOfType<scf::ForOp>();
+        leaf.enclosingLoop = getEnclosingDeferredLoop(slice, deferred);
         if (graphProjection
             && slice.getSourceType().getRank()
                  == leaf.reconstructedType.getRank() + 1
@@ -609,8 +614,11 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
       program.leaves.push_back(std::move(leaf));
       return success();
     }
-    if (value.getType().isIndex() || isa<IntegerType>(value.getType()))
-      return success();
+    if (value.getType().isIndex() || isa<IntegerType>(value.getType())) {
+      Operation *definition = value.getDefiningOp();
+      if (!definition || !deferred->isProperAncestor(definition))
+        return success();
+    }
     if (auto argument = dyn_cast<BlockArgument>(value)) {
       auto loop = dyn_cast_or_null<scf::ForOp>(
         argument.getOwner()->getParentOp());
@@ -619,7 +627,7 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
     }
     Operation *op = value.getDefiningOp();
     if (!op || (op->getBlock() != &body
-                && !op->getParentOfType<scf::ForOp>()))
+                && !getEnclosingDeferredLoop(op, deferred)))
       return deferred.emitOpError(
         "deferred residual escapes its verified body: ") << value;
     if (auto loop = dyn_cast<scf::ForOp>(op)) {

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

@@ -289,7 +289,9 @@ static Value cloneResidual(
       mapping.map(oldValue, newValue);
   }
   for (Operation *op : exchange.program.residualOps) {
-    if (op->hasTrait<OpTrait::ConstantLike>())
+    if (op->hasTrait<OpTrait::ConstantLike>()
+        || llvm::all_of(op->getResults(),
+             [&](Value result) { return mapping.contains(result); }))
       continue;
     if (auto oldLoop = dyn_cast<scf::ForOp>(op)) {
       SmallVector<Value> initArgs;

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

@@ -454,6 +454,9 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa
   OpBuilder builder(blueprint);
   blueprint->setAttr("fragmentOperandIndices", builder.getDenseI64ArrayAttr(newOperands));
   blueprint->setAttr("fragmentSourceSlots", builder.getDenseI64ArrayAttr(newSlots));
+  if (blueprint.getIndexMap() == kContiguousRowMajorFragments
+      && !isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    blueprint.setIndexMapAttr(builder.getStringAttr("fragment_assembly"));
   return success();
 }
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -400,6 +400,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
   std::vector<ResidentWeightSet> processorResidentWeights(processorCount);
   std::vector<ScheduledTask> schedules(nodeCount);
   std::vector<std::vector<size_t>> tasksByProcessor(processorCount);
+  std::vector<std::vector<size_t>> timelineByProcessor(processorCount);
 
   size_t scheduledCount = 0;
   while (!readyQueue.empty()) {
@@ -441,7 +442,7 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
       Time currentEnd = 0;
       bool foundGap = false;
 
-      for (size_t schedTaskIndex : tasksByProcessor[processor]) {
+      for (size_t schedTaskIndex : timelineByProcessor[processor]) {
         const ScheduledTask& schedTask = schedules[schedTaskIndex];
         Time gapStart = std::max(currentEnd, dataReady);
 
@@ -532,10 +533,13 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
     insertResidentWeights(capacityReservations[bestProcessor], graph.nodes[task].residentWeights);
     insertResidentWeights(processorResidentWeights[bestProcessor], graph.nodes[task].residentWeights);
 
-    // 3. CRITICAL FIX: Topological Append
-    // Because the readyQueue pops in strict topological order, simply pushing to the
-    // back guarantees the Monoliths will be physically generated cycle-free.
-    // The hardware will still benefit from the processor assignment chosen by PEFT.
+    auto& timeline = timelineByProcessor[bestProcessor];
+    timeline.insert(llvm::upper_bound(timeline, task, [&](size_t lhs, size_t rhs) {
+      return schedules[lhs].startTime < schedules[rhs].startTime;
+    }), task);
+
+    // Materialization requires topological order; gap placement requires the
+    // separate chronological timeline above.
     tasksByProcessor[bestProcessor].push_back(task);
 
     for (const auto& [child, weight] : graph.successors[task]) {

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

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinAttributes.h"
@@ -16,6 +17,7 @@
 
 #include <cstdint>
 #include <fstream>
+#include <limits>
 #include <optional>
 #include <string>
 #include <utility>
@@ -54,6 +56,12 @@ struct ChannelSendRecord {
   std::optional<uint32_t> sourceLane;
 };
 
+struct ChannelEvaluationContext {
+  Value laneArg;
+  uint32_t lane = 0;
+  DenseMap<Value, int64_t> bindings;
+};
+
 enum class LogicalNodeSelector {
   Scalar,
   Lane,
@@ -249,10 +257,20 @@ void addBatchNodeRows(std::fstream& nodesFile,
   }
 }
 
-std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t lane);
+std::optional<int64_t> evaluateIndexLike(Value value,
+                                         Value laneArg,
+                                         uint32_t lane,
+                                         const DenseMap<Value, int64_t>* bindings);
 
-std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t lane) {
-  if (value == laneArg)
+std::optional<int64_t> evaluateIndexLike(Value value,
+                                         Value laneArg,
+                                         uint32_t lane,
+                                         const DenseMap<Value, int64_t>* bindings) {
+  if (bindings)
+    if (auto it = bindings->find(value); it != bindings->end())
+      return it->second;
+
+  if (laneArg && value == laneArg)
     return static_cast<int64_t>(lane);
 
   if (std::optional<int64_t> constant = matchConstantIndexValue(value))
@@ -270,7 +288,8 @@ std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t la
     if (!elements || !shapedType || shapedType.getRank() != 1 || extract.getIndices().size() != 1)
       return std::nullopt;
 
-    std::optional<int64_t> index = evaluateIndexLike(extract.getIndices().front(), laneArg, lane);
+    std::optional<int64_t> index =
+      evaluateIndexLike(extract.getIndices().front(), laneArg, lane, bindings);
     if (!index || *index < 0 || *index >= static_cast<int64_t>(elements.getNumElements()))
       return std::nullopt;
 
@@ -279,11 +298,104 @@ std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t la
     return std::nullopt;
   }
 
+  if (auto indexCast = value.getDefiningOp<arith::IndexCastOp>())
+    return evaluateIndexLike(indexCast.getIn(), laneArg, lane, bindings);
+
+  if (auto add = value.getDefiningOp<arith::AddIOp>()) {
+    auto lhs = evaluateIndexLike(add.getLhs(), laneArg, lane, bindings);
+    auto rhs = evaluateIndexLike(add.getRhs(), laneArg, lane, bindings);
+    if (lhs && rhs)
+      return *lhs + *rhs;
+    return std::nullopt;
+  }
+
+  if (auto sub = value.getDefiningOp<arith::SubIOp>()) {
+    auto lhs = evaluateIndexLike(sub.getLhs(), laneArg, lane, bindings);
+    auto rhs = evaluateIndexLike(sub.getRhs(), laneArg, lane, bindings);
+    if (lhs && rhs)
+      return *lhs - *rhs;
+    return std::nullopt;
+  }
+
+  if (auto mul = value.getDefiningOp<arith::MulIOp>()) {
+    auto lhs = evaluateIndexLike(mul.getLhs(), laneArg, lane, bindings);
+    auto rhs = evaluateIndexLike(mul.getRhs(), laneArg, lane, bindings);
+    if (lhs && rhs)
+      return *lhs * *rhs;
+    return std::nullopt;
+  }
+
+  if (auto div = value.getDefiningOp<arith::DivSIOp>()) {
+    auto lhs = evaluateIndexLike(div.getLhs(), laneArg, lane, bindings);
+    auto rhs = evaluateIndexLike(div.getRhs(), laneArg, lane, bindings);
+    if (!lhs || !rhs || *rhs == 0
+        || (*lhs == std::numeric_limits<int64_t>::min() && *rhs == -1))
+      return std::nullopt;
+    return *lhs / *rhs;
+  }
+
+  if (auto div = value.getDefiningOp<arith::DivUIOp>()) {
+    auto lhs = evaluateIndexLike(div.getLhs(), laneArg, lane, bindings);
+    auto rhs = evaluateIndexLike(div.getRhs(), laneArg, lane, bindings);
+    if (!lhs || !rhs || *rhs == 0)
+      return std::nullopt;
+    return static_cast<int64_t>(static_cast<uint64_t>(*lhs) / static_cast<uint64_t>(*rhs));
+  }
+
+  if (auto rem = value.getDefiningOp<arith::RemSIOp>()) {
+    auto lhs = evaluateIndexLike(rem.getLhs(), laneArg, lane, bindings);
+    auto rhs = evaluateIndexLike(rem.getRhs(), laneArg, lane, bindings);
+    if (!lhs || !rhs || *rhs == 0)
+      return std::nullopt;
+    if (*lhs == std::numeric_limits<int64_t>::min() && *rhs == -1)
+      return 0;
+    return *lhs % *rhs;
+  }
+
+  if (auto rem = value.getDefiningOp<arith::RemUIOp>()) {
+    auto lhs = evaluateIndexLike(rem.getLhs(), laneArg, lane, bindings);
+    auto rhs = evaluateIndexLike(rem.getRhs(), laneArg, lane, bindings);
+    if (!lhs || !rhs || *rhs == 0)
+      return std::nullopt;
+    return static_cast<int64_t>(static_cast<uint64_t>(*lhs) % static_cast<uint64_t>(*rhs));
+  }
+
+  if (auto cmp = value.getDefiningOp<arith::CmpIOp>()) {
+    auto lhs = evaluateIndexLike(cmp.getLhs(), laneArg, lane, bindings);
+    auto rhs = evaluateIndexLike(cmp.getRhs(), laneArg, lane, bindings);
+    if (!lhs || !rhs)
+      return std::nullopt;
+    bool result = false;
+    switch (cmp.getPredicate()) {
+    case arith::CmpIPredicate::eq:  result = *lhs == *rhs; break;
+    case arith::CmpIPredicate::ne:  result = *lhs != *rhs; break;
+    case arith::CmpIPredicate::slt: result = *lhs < *rhs; break;
+    case arith::CmpIPredicate::sle: result = *lhs <= *rhs; break;
+    case arith::CmpIPredicate::sgt: result = *lhs > *rhs; break;
+    case arith::CmpIPredicate::sge: result = *lhs >= *rhs; break;
+    case arith::CmpIPredicate::ult: result = static_cast<uint64_t>(*lhs) < static_cast<uint64_t>(*rhs); break;
+    case arith::CmpIPredicate::ule: result = static_cast<uint64_t>(*lhs) <= static_cast<uint64_t>(*rhs); break;
+    case arith::CmpIPredicate::ugt: result = static_cast<uint64_t>(*lhs) > static_cast<uint64_t>(*rhs); break;
+    case arith::CmpIPredicate::uge: result = static_cast<uint64_t>(*lhs) >= static_cast<uint64_t>(*rhs); break;
+    }
+    return result ? 1 : 0;
+  }
+
+  if (auto select = value.getDefiningOp<arith::SelectOp>()) {
+    auto condition = evaluateIndexLike(select.getCondition(), laneArg, lane, bindings);
+    if (!condition)
+      return std::nullopt;
+    return evaluateIndexLike(*condition ? select.getTrueValue() : select.getFalseValue(),
+                             laneArg,
+                             lane,
+                             bindings);
+  }
+
   if (auto affineApply = value.getDefiningOp<affine::AffineApplyOp>())
     if (FailureOr<int64_t> folded = evaluateAffineApply(affineApply,
                                                         [&](Value operand) -> FailureOr<int64_t> {
                                                           if (std::optional<int64_t> resolved =
-                                                                evaluateIndexLike(operand, laneArg, lane))
+                                                                evaluateIndexLike(operand, laneArg, lane, bindings))
                                                             return *resolved;
                                                           return failure();
                                                         });
@@ -294,24 +406,70 @@ std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t la
   return std::nullopt;
 }
 
-SmallVector<int64_t, 8> collectPossibleIntValues(Value value, Value laneArg, uint32_t lane) {
-  if (std::optional<int64_t> exact = evaluateIndexLike(value, laneArg, lane))
-    return {*exact};
+bool containsChannelOperation(Operation* root) {
+  bool found = false;
+  root->walk([&](Operation* op) {
+    found |= isa<SpatChannelSendOp, SpatChannelReceiveOp>(op);
+  });
+  return found;
+}
 
-  auto extract = value.getDefiningOp<tensor::ExtractOp>();
-  auto constant = extract ? extract.getTensor().getDefiningOp<arith::ConstantOp>() : nullptr;
-  auto elements = constant ? dyn_cast<ElementsAttr>(constant.getValue()) : nullptr;
-  if (!elements)
-    return {};
+template <typename Emit>
+LogicalResult walkChannelRegion(Region& region, const ChannelEvaluationContext& context, Emit& emit) {
+  if (region.empty())
+    return success();
 
-  SmallVector<int64_t, 8> values;
-  if (auto denseInts = dyn_cast<DenseIntElementsAttr>(elements)) {
-    values.reserve(elements.getNumElements());
-    for (APInt element : denseInts.getValues<APInt>())
-      if (!llvm::is_contained(values, element.getSExtValue()))
-        values.push_back(element.getSExtValue());
+  for (Operation& op : region.front()) {
+    if (auto ifOp = dyn_cast<scf::IfOp>(&op)) {
+      auto condition = evaluateIndexLike(
+        ifOp.getCondition(), context.laneArg, context.lane, &context.bindings);
+      if (!condition)
+        return ifOp.emitOpError("has an unresolved condition in Spatial dataflow export");
+      Region& selected = *condition ? ifOp.getThenRegion() : ifOp.getElseRegion();
+      if (failed(walkChannelRegion(selected, context, emit)))
+        return failure();
+      continue;
+    }
+
+    if (auto forOp = dyn_cast<scf::ForOp>(&op)) {
+      if (!containsChannelOperation(forOp))
+        continue;
+
+      auto lower = evaluateIndexLike(forOp.getLowerBound(), context.laneArg, context.lane, &context.bindings);
+      auto upper = evaluateIndexLike(forOp.getUpperBound(), context.laneArg, context.lane, &context.bindings);
+      auto step = evaluateIndexLike(forOp.getStep(), context.laneArg, context.lane, &context.bindings);
+      if (!lower || !upper || !step || *step == 0)
+        return forOp.emitOpError("has unresolved or invalid bounds in Spatial dataflow export");
+
+      constexpr uint64_t kMaxExportedLoopIterations = 1 << 20;
+      uint64_t iterationCount = 0;
+      int64_t induction = *lower;
+      while ((*step > 0 && induction < *upper) || (*step < 0 && induction > *upper)) {
+        if (++iterationCount > kMaxExportedLoopIterations)
+          return forOp.emitOpError("exceeds the bounded iteration limit in Spatial dataflow export");
+        ChannelEvaluationContext iterationContext = context;
+        iterationContext.bindings[forOp.getInductionVar()] = induction;
+        if (failed(walkChannelRegion(forOp.getRegion(), iterationContext, emit)))
+          return failure();
+        if ((*step > 0 && induction > std::numeric_limits<int64_t>::max() - *step)
+            || (*step < 0 && induction < std::numeric_limits<int64_t>::min() - *step))
+          return forOp.emitOpError("overflows while enumerating Spatial dataflow export iterations");
+        induction += *step;
+      }
+      continue;
+    }
+
+    if (!isa<SpatChannelSendOp, SpatChannelReceiveOp>(&op))
+      continue;
+    auto channel = dyn_cast<SpatChannelSendOp>(&op);
+    Value channelValue = channel ? channel.getChannelId() : cast<SpatChannelReceiveOp>(&op).getChannelId();
+    auto channelId = evaluateIndexLike(channelValue, context.laneArg, context.lane, &context.bindings);
+    if (!channelId)
+      return op.emitError("has an unresolved channel identity in Spatial dataflow export");
+    if (failed(emit(op, *channelId, context)))
+      return failure();
   }
-  return values;
+  return success();
 }
 
 template <typename BatchOpTy>
@@ -604,50 +762,44 @@ LogicalResult emitDataEdges(std::fstream& edgesFile,
 }
 
 template <typename BatchOpTy>
-void collectChannelSends(DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>>& sendsByChannelId,
-                         const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                         BatchOpTy batch) {
+LogicalResult collectChannelSends(DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>>& sendsByChannelId,
+                                  const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                                  BatchOpTy batch) {
   std::optional<BlockArgument> laneArg = batch.getLaneArgument();
   if (!laneArg)
-    return;
+    return success();
 
   for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
     std::string sourceId = getExpandedNodeId(expandedNodes, batch.getOperation(), lane);
     if (sourceId.empty())
       continue;
-    batch.getBody().walk([&](SpatChannelSendOp send) {
-      std::optional<int64_t> channelId = evaluateIndexLike(send.getChannelId(), *laneArg, lane);
-      if (!channelId)
-        return;
-      sendsByChannelId[*channelId].push_back({sourceId, lane});
-    });
+    ChannelEvaluationContext context;
+    context.laneArg = *laneArg;
+    context.lane = lane;
+    auto emit = [&](Operation& op, int64_t channelId, const ChannelEvaluationContext&) {
+      if (auto send = dyn_cast<SpatChannelSendOp>(&op))
+        sendsByChannelId[channelId].push_back({sourceId, lane});
+      return success();
+    };
+    if (failed(walkChannelRegion(batch.getBody(), context, emit)))
+      return failure();
   }
+  return success();
 }
 
-void collectChannelSends(DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>>& sendsByChannelId,
-                         const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                         SpatScheduledCompute compute) {
+LogicalResult collectChannelSends(DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>>& sendsByChannelId,
+                                  const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                                  SpatScheduledCompute compute) {
   std::string sourceId = getExpandedNodeId(expandedNodes, compute.getOperation(), 0);
   if (sourceId.empty())
-    return;
-  compute.getBody().walk([&](SpatChannelSendOp send) {
-    std::optional<int64_t> channelId = evaluateIndexLike(send.getChannelId(), Value(), 0);
-    if (!channelId)
-      return;
-    sendsByChannelId[*channelId].push_back({sourceId, std::nullopt});
-  });
-}
-
-DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>>
-buildNodesByCore(const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes) {
-  DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>> nodesByCore;
-  for (const auto& entry : expandedNodes) {
-    const ExpandedNodeInfo& node = entry.second;
-    if (!node.core)
-      continue;
-    nodesByCore[*node.core].push_back({node.id, node.lane});
-  }
-  return nodesByCore;
+    return success();
+  ChannelEvaluationContext context;
+  auto emit = [&](Operation& op, int64_t channelId, const ChannelEvaluationContext&) {
+    if (isa<SpatChannelSendOp>(&op))
+      sendsByChannelId[channelId].push_back({sourceId, std::nullopt});
+    return success();
+  };
+  return walkChannelRegion(compute.getBody(), context, emit);
 }
 
 template <typename ComputeOpTy, typename BatchOpTy, typename ResolveChannelSourcesFn>
@@ -660,14 +812,18 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile,
     const TopLevelOpInfo& info = entry.second;
 
     if (auto compute = dyn_cast<ComputeOpTy>(op)) {
-      compute.getBody().walk([&](SpatChannelReceiveOp receive) {
-        SmallVector<ChannelSendRecord, 4> sources = resolveChannelSources(receive, 0);
-        if (sources.empty())
-          return;
-        std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), Value(), 0);
+      ChannelEvaluationContext context;
+      auto emit = [&](Operation& channelOp, int64_t channelId, const ChannelEvaluationContext&) {
+        auto receive = dyn_cast<SpatChannelReceiveOp>(&channelOp);
+        if (!receive)
+          return success();
+        FailureOr<SmallVector<ChannelSendRecord, 4>> sources =
+          resolveChannelSources(receive, channelId, 0);
+        if (failed(sources))
+          return failure();
         std::string targetId = getScalarId(info.isScheduled, info.opId);
         std::optional<uint64_t> byteSize = getTypeSizeBytes(receive.getType());
-        for (const ChannelSendRecord& source : sources)
+        for (const ChannelSendRecord& source : *sources)
           emitEdgeRow(edgesFile,
                       source.sourceId,
                       targetId,
@@ -677,7 +833,10 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile,
                       source.sourceLane,
                       std::nullopt,
                       channelId);
-      });
+        return success();
+      };
+      if (failed(walkChannelRegion(compute.getBody(), context, emit)))
+        return failure();
       continue;
     }
 
@@ -688,14 +847,20 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile,
     if (!laneArg)
       continue;
     for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
-      std::string targetId = getBatchLaneId(info.isScheduled, info.opId, lane);
-      batch.getBody().walk([&](SpatChannelReceiveOp receive) {
-        SmallVector<ChannelSendRecord, 4> sources = resolveChannelSources(receive, lane);
-        if (sources.empty())
-          return;
-        std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), *laneArg, lane);
+      ChannelEvaluationContext context;
+      context.laneArg = *laneArg;
+      context.lane = lane;
+      auto emit = [&](Operation& channelOp, int64_t channelId, const ChannelEvaluationContext& eventContext) {
+        auto receive = dyn_cast<SpatChannelReceiveOp>(&channelOp);
+        if (!receive)
+          return success();
+        FailureOr<SmallVector<ChannelSendRecord, 4>> sources =
+          resolveChannelSources(receive, channelId, eventContext.lane);
+        if (failed(sources))
+          return failure();
+        std::string targetId = getBatchLaneId(info.isScheduled, info.opId, eventContext.lane);
         std::optional<uint64_t> byteSize = getTypeSizeBytes(receive.getType());
-        for (const ChannelSendRecord& source : sources)
+        for (const ChannelSendRecord& source : *sources)
           emitEdgeRow(edgesFile,
                       source.sourceId,
                       targetId,
@@ -703,9 +868,12 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile,
                       receive.getType(),
                       stage,
                       source.sourceLane,
-                      lane,
+                      eventContext.lane,
                       channelId);
-      });
+        return success();
+      };
+      if (failed(walkChannelRegion(batch.getBody(), context, emit)))
+        return failure();
     }
   }
 
@@ -810,33 +978,25 @@ LogicalResult exportScheduled(func::FuncOp func,
   DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>> sendsByChannelId;
   for (const auto& entry : topLevelInfo) {
     Operation* op = entry.first;
+    LogicalResult collected = success();
     if (auto compute = dyn_cast<SpatScheduledCompute>(op))
-      collectChannelSends(sendsByChannelId, expandedNodes, compute);
+      collected = collectChannelSends(sendsByChannelId, expandedNodes, compute);
     else if (auto batch = dyn_cast<SpatScheduledComputeBatch>(op))
-      collectChannelSends(sendsByChannelId, expandedNodes, batch);
+      collected = collectChannelSends(sendsByChannelId, expandedNodes, batch);
+    if (failed(collected))
+      return failure();
   }
 
-  DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>> nodesByCore = buildNodesByCore(expandedNodes);
-  auto resolveChannelSources = [&](SpatChannelReceiveOp receive, uint32_t lane) {
+  DenseMap<int64_t, size_t> consumedSendsByChannelId;
+  auto resolveChannelSources = [&](SpatChannelReceiveOp receive, int64_t channelId, uint32_t) {
     SmallVector<ChannelSendRecord, 4> sources;
-
-    Value laneArg;
-    if (auto owner = receive->getParentOfType<SpatScheduledComputeBatch>())
-      if (auto maybeLaneArg = owner.getLaneArgument())
-        laneArg = *maybeLaneArg;
-
-    if (std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), laneArg, lane)) {
-      if (auto it = sendsByChannelId.find(*channelId); it != sendsByChannelId.end())
-        return it->second;
-    }
-
-    for (int64_t sourceCore : collectPossibleIntValues(receive.getSourceCoreId(), laneArg, lane)) {
-      auto it = nodesByCore.find(static_cast<int32_t>(sourceCore));
-      if (it == nodesByCore.end())
-        continue;
-      llvm::append_range(sources, it->second);
-    }
-    return sources;
+    auto sends = sendsByChannelId.find(channelId);
+    size_t& consumed = consumedSendsByChannelId[channelId];
+    if (sends == sendsByChannelId.end() || consumed >= sends->second.size())
+      return receive.emitOpError("has no matching realized channel send in Spatial dataflow export"),
+             FailureOr<SmallVector<ChannelSendRecord, 4>>(failure());
+    sources.push_back(sends->second[consumed++]);
+    return FailureOr<SmallVector<ChannelSendRecord, 4>>(std::move(sources));
   };
 
   return emitExplicitChannelEdges<SpatScheduledCompute, SpatScheduledComputeBatch>(