Cose belle

2026-07-14 16:57:58 +02:00
parent d1a29ace3c
commit 51fdb830e5
38 changed files with 5365 additions and 914 deletions
@@ -105,6 +105,9 @@ Pass these to `onnx-mlir` when compiling for PIM:
  the codegen tail.
 - `--pim-emit-json` - also emit `core_*.json` instruction files alongside
  `core_*.pim`.
+- `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|all>` - control Spatial
+  dataflow CSV reports. The default `all` emits graph, scheduled, and realized
+  snapshots under `reports/`.
 - `--use-experimental-conv-impl` - use the alternate convolution lowering.
 - `--ignore-concat-error` - soft-fail a ConcatOp corner case.

@@ -167,7 +170,8 @@ Each validation run writes artifacts in the model workspace, for example under
 - `simulation/out.bin` - raw simulator output used for comparison.

 The compiler currently dumps dialect snapshots such as `spatial0.mlir`,
-`spatial1_dcp_merged.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
+`spatial1_graph.mlir`, `spatial2_scheduled_no_comm.mlir`,
+`spatial3_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
 `pim2_coalesced.mlir`, and `pim3_folded.mlir` when an output directory is
 available.

@@ -8,6 +8,7 @@ add_pim_library(OMPimCommon
  IR/IndexingUtils.cpp
  IR/LoopUtils.cpp
  IR/ShapeUtils.cpp
+  IR/ShapingUtils.cpp
  IR/SubviewUtils.cpp
  IR/TensorSliceUtils.cpp
  IR/WeightUtils.cpp
@@ -31,7 +31,7 @@ static FailureOr<int64_t> ceilDivSigned(int64_t lhs, int64_t rhs) {
 }

 Value createOrFoldAffineApply(
-  RewriterBase& rewriter, Location loc, AffineMap map, ValueRange operands, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, AffineMap map, ValueRange operands, Operation* constantAnchor) {
  assert(constantAnchor && "expected a valid constant anchor");
  assert(map.getNumResults() == 1 && "affine.apply expects a single-result affine map");

@@ -40,91 +40,91 @@ Value createOrFoldAffineApply(
  for (Value operand : operands) {
    std::optional<int64_t> constantValue = matchConstantIndexValue(operand);
    if (!constantValue)
-      return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
-    operandConstants.push_back(rewriter.getIndexAttr(*constantValue));
+      return affine::AffineApplyOp::create(builder, loc, map, operands).getResult();
+    operandConstants.push_back(builder.getIndexAttr(*constantValue));
  }

  SmallVector<Attribute> foldedResults;
  if (succeeded(map.constantFold(operandConstants, foldedResults)) && foldedResults.size() == 1)
    if (auto constantResult = dyn_cast<IntegerAttr>(foldedResults.front()))
-      return getOrCreateIndexConstant(rewriter, constantAnchor, constantResult.getInt());
+      return getOrCreateIndexConstant(builder, constantAnchor, constantResult.getInt());

-  return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
+  return affine::AffineApplyOp::create(builder, loc, map, operands).getResult();
 }

 Value createOrFoldAffineApply(
-  RewriterBase& rewriter, Location loc, AffineExpr expr, ValueRange dims, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, AffineExpr expr, ValueRange dims, Operation* constantAnchor) {
  AffineMap map = AffineMap::get(/*dimCount=*/dims.size(), /*symbolCount=*/0, expr);
-  return createOrFoldAffineApply(rewriter, loc, map, dims, constantAnchor);
+  return createOrFoldAffineApply(builder, loc, map, dims, constantAnchor);
 }

-Value affineMulConst(RewriterBase& rewriter, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) {
+Value affineMulConst(OpBuilder& builder, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) {
  assert(constantAnchor && "expected a valid constant anchor");
  if (multiplier == 0)
-    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);
  if (multiplier == 1)
    return value;

-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0 * multiplier, ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 * multiplier, ValueRange {value}, constantAnchor);
 }

-Value affineAddConst(RewriterBase& rewriter, Location loc, Value value, int64_t offset, Operation* constantAnchor) {
+Value affineAddConst(OpBuilder& builder, Location loc, Value value, int64_t offset, Operation* constantAnchor) {
  assert(constantAnchor && "expected a valid constant anchor");
  if (offset == 0)
    return value;

-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0 + offset, ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 + offset, ValueRange {value}, constantAnchor);
 }

-Value affineModConst(RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
+Value affineModConst(OpBuilder& builder, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
  assert(constantAnchor && "expected a valid constant anchor");
  assert(divisor > 0 && "expected a positive affine.mod divisor");
  if (divisor == 1)
-    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);

-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0 % divisor, ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 % divisor, ValueRange {value}, constantAnchor);
 }

 Value affineFloorDivConst(
-  RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
  assert(constantAnchor && "expected a valid constant anchor");
  assert(divisor > 0 && "expected a positive affine.floor_div divisor");
  if (divisor == 1)
    return value;

-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor);
 }

 Value affineAddModConst(
-  RewriterBase& rewriter, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
  assert(constantAnchor && "expected a valid constant anchor");
  assert(divisor > 0 && "expected a positive affine.mod divisor");
  if (divisor == 1)
-    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);

-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
  AffineExpr expr = d0;
  if (offset != 0)
    expr = expr + offset;
-  return createOrFoldAffineApply(rewriter, loc, expr % divisor, ValueRange {value}, constantAnchor);
+  return createOrFoldAffineApply(builder, loc, expr % divisor, ValueRange {value}, constantAnchor);
 }

 Value affineAddFloorDivConst(
-  RewriterBase& rewriter, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
  assert(constantAnchor && "expected a valid constant anchor");
  assert(divisor > 0 && "expected a positive affine.floor_div divisor");
  if (divisor == 1)
-    return offset == 0 ? value : affineAddConst(rewriter, loc, value, offset, constantAnchor);
+    return offset == 0 ? value : affineAddConst(builder, loc, value, offset, constantAnchor);

-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
  AffineExpr expr = d0;
  if (offset != 0)
    expr = expr + offset;
-  return createOrFoldAffineApply(rewriter, loc, expr.floorDiv(divisor), ValueRange {value}, constantAnchor);
+  return createOrFoldAffineApply(builder, loc, expr.floorDiv(divisor), ValueRange {value}, constantAnchor);
 }

 FailureOr<int64_t> evaluateAffineExpr(AffineExpr expr, ArrayRef<int64_t> dims, ArrayRef<int64_t> symbols) {
@@ -11,50 +11,50 @@ namespace onnx_mlir {

 using IndexValueResolver = llvm::function_ref<llvm::FailureOr<int64_t>(mlir::Value)>;

-mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter,
+mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder,
                                    mlir::Location loc,
                                    mlir::AffineMap map,
                                    mlir::ValueRange operands,
                                    mlir::Operation* constantAnchor);

-mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter,
+mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder,
                                    mlir::Location loc,
                                    mlir::AffineExpr expr,
                                    mlir::ValueRange dims,
                                    mlir::Operation* constantAnchor);

-mlir::Value affineMulConst(mlir::RewriterBase& rewriter,
+mlir::Value affineMulConst(mlir::OpBuilder& builder,
                           mlir::Location loc,
                           mlir::Value value,
                           int64_t multiplier,
                           mlir::Operation* constantAnchor);

-mlir::Value affineAddConst(mlir::RewriterBase& rewriter,
+mlir::Value affineAddConst(mlir::OpBuilder& builder,
                           mlir::Location loc,
                           mlir::Value value,
                           int64_t offset,
                           mlir::Operation* constantAnchor);

-mlir::Value affineModConst(mlir::RewriterBase& rewriter,
+mlir::Value affineModConst(mlir::OpBuilder& builder,
                           mlir::Location loc,
                           mlir::Value value,
                           int64_t divisor,
                           mlir::Operation* constantAnchor);

-mlir::Value affineFloorDivConst(mlir::RewriterBase& rewriter,
+mlir::Value affineFloorDivConst(mlir::OpBuilder& builder,
                                mlir::Location loc,
                                mlir::Value value,
                                int64_t divisor,
                                mlir::Operation* constantAnchor);

-mlir::Value affineAddModConst(mlir::RewriterBase& rewriter,
+mlir::Value affineAddModConst(mlir::OpBuilder& builder,
                              mlir::Location loc,
                              mlir::Value value,
                              int64_t offset,
                              int64_t divisor,
                              mlir::Operation* constantAnchor);

-mlir::Value affineAddFloorDivConst(mlir::RewriterBase& rewriter,
+mlir::Value affineAddFloorDivConst(mlir::OpBuilder& builder,
                                   mlir::Location loc,
                                   mlir::Value value,
                                   int64_t offset,
@@ -49,7 +49,7 @@ Value getOrCreateConstant(OperationFolder& folder, Operation* anchorOp, Attribut
  return folder.getOrCreateConstant(hostBlock, arithDialect, value, type);
 }

-Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute value, Type type) {
+Value getOrCreateConstant(OpBuilder& builder, Operation* anchorOp, Attribute value, Type type) {
  assert(anchorOp && "expected a valid anchor operation");
  Block* hostBlock = getConstantInsertionBlock(anchorOp);
  for (Operation& op : *hostBlock) {
@@ -59,9 +59,16 @@ Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute
    return constantOp.getResult();
  }

-  OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPointToStart(hostBlock);
-  return arith::ConstantOp::create(rewriter, anchorOp->getLoc(), type, cast<TypedAttr>(value)).getResult();
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(hostBlock);
+  return arith::ConstantOp::create(builder, anchorOp->getLoc(), type, cast<TypedAttr>(value)).getResult();
+}
+
+Value createConstantAtHostBlockStart(OpBuilder& builder, Operation* anchorOp, TypedAttr value) {
+  assert(anchorOp && "expected a valid anchor operation");
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(getConstantInsertionBlock(anchorOp));
+  return arith::ConstantOp::create(builder, anchorOp->getLoc(), value).getResult();
 }

 Value getOrCreateConstantLike(OperationFolder& folder, arith::ConstantOp constantOp) {
@@ -73,9 +80,8 @@ Value getOrCreateIndexConstant(OperationFolder& folder, Operation* anchorOp, int
  return getOrCreateConstant(folder, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
 }

-Value getOrCreateIndexConstant(RewriterBase& rewriter, Operation* anchorOp, int64_t value) {
-  Builder builder(anchorOp->getContext());
-  return getOrCreateConstant(rewriter, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
+Value getOrCreateIndexConstant(OpBuilder& builder, Operation* anchorOp, int64_t value) {
+  return getOrCreateConstant(builder, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
 }

 void hoistAndUniquifyIndexConstants(func::FuncOp funcOp, RewriterBase& rewriter) {
@@ -16,13 +16,16 @@ mlir::Value
 getOrCreateConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);

 mlir::Value
-getOrCreateConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
+getOrCreateConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
+
+mlir::Value
+createConstantAtHostBlockStart(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::TypedAttr value);

 mlir::Value getOrCreateConstantLike(mlir::OperationFolder& folder, mlir::arith::ConstantOp constantOp);

 mlir::Value getOrCreateIndexConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, int64_t value);

-mlir::Value getOrCreateIndexConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, int64_t value);
+mlir::Value getOrCreateIndexConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, int64_t value);

 void hoistAndUniquifyIndexConstants(mlir::func::FuncOp funcOp, mlir::RewriterBase& rewriter);

@@ -36,9 +36,10 @@ bool isCoreStaticAddressOp(mlir::Operation* op) {

 mlir::LogicalResult
 walkPimCoreBlock(mlir::Block& block,
-                 const StaticValueKnowledge& knowledge,
+                 const StaticValueKnowledge& initialKnowledge,
                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
  bool hasFailure = false;
+  StaticValueKnowledge knowledge = initialKnowledge;
  for (mlir::Operation& op : block) {
    if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
      continue;
@@ -89,6 +90,27 @@ walkPimCoreBlock(mlir::Block& block,
      continue;
    }

+    if (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        hasFailure = true;
+        continue;
+      }
+      mlir::Region* selected = &switchOp.getDefaultRegion();
+      for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions()))
+        if (caseValue == *selector) {
+          selected = &caseRegion;
+          break;
+        }
+      if (failed(walkPimCoreBlock(selected->front(), knowledge, callback)))
+        hasFailure = true;
+      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
    if (failed(callback(op, knowledge)))
      hasFailure = true;
  }
@@ -97,9 +119,10 @@ walkPimCoreBlock(mlir::Block& block,

 mlir::LogicalResult walkPimCoreBlockStructurally(
  mlir::Block& block,
-  const StaticValueKnowledge& knowledge,
+  const StaticValueKnowledge& initialKnowledge,
  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
  bool hasFailure = false;
+  StaticValueKnowledge knowledge = initialKnowledge;
  for (mlir::Operation& op : block) {
    if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
      continue;
@@ -159,6 +182,28 @@ mlir::LogicalResult walkPimCoreBlockStructurally(
      continue;
    }

+    if (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM verification");
+        hasFailure = true;
+        continue;
+      }
+      mlir::Region* selected = &switchOp.getDefaultRegion();
+      for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions()))
+        if (caseValue == *selector) {
+          selected = &caseRegion;
+          break;
+        }
+      for (mlir::Region& region : switchOp->getRegions())
+        if (failed(walkPimCoreBlockStructurally(region.front(), knowledge, callback)))
+          hasFailure = true;
+      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
    if (failed(callback(op, knowledge)))
      hasFailure = true;
  }
@@ -0,0 +1,39 @@
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+bool isShapingOnlyOp(Operation *op) {
+  return isa<tensor::CastOp,
+             tensor::CollapseShapeOp,
+             tensor::ExpandShapeOp,
+             tensor::ExtractSliceOp,
+             tensor::InsertSliceOp,
+             tensor::ConcatOp,
+             tensor::EmptyOp,
+             tensor::ExtractOp,
+             tensor::InsertOp,
+             tensor::SplatOp,
+             linalg::TransposeOp,
+             ONNXTransposeOp,
+             spatial::SpatConcatOp,
+             spatial::SpatExtractRowsOp>(op);
+}
+
+bool isPureIndexComputationOp(Operation *op) {
+  if (op->getNumRegions() != 0 || op->getNumResults() == 0 || op->hasTrait<OpTrait::IsTerminator>()
+      || !isMemoryEffectFree(op))
+    return false;
+  auto isIndexOrInteger = [](Type type) { return type.isIndex() || isa<IntegerType>(type); };
+  return llvm::all_of(op->getOperandTypes(), isIndexOrInteger)
+      && llvm::all_of(op->getResultTypes(), isIndexOrInteger);
+}
+
+} // namespace onnx_mlir
@@ -0,0 +1,13 @@
+#pragma once
+
+namespace mlir {
+class Operation;
+}
+
+namespace onnx_mlir {
+
+bool isShapingOnlyOp(mlir::Operation *op);
+
+bool isPureIndexComputationOp(mlir::Operation *op);
+
+} // namespace onnx_mlir
@@ -68,4 +68,39 @@ Value insertStaticSlice(
    .getResult();
 }

+FailureOr<Value> addLeadingUnitTensorDimension(OpBuilder& builder, Location loc, Value value) {
+  auto type = dyn_cast<RankedTensorType>(value.getType());
+  if (!type || !type.hasStaticShape())
+    return failure();
+  SmallVector<int64_t> shape {1};
+  llvm::append_range(shape, type.getShape());
+  auto resultType = RankedTensorType::get(shape, type.getElementType(), type.getEncoding());
+  SmallVector<ReassociationIndices> reassociation;
+  if (type.getRank() != 0) {
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < type.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+  }
+  return tensor::ExpandShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
+}
+
+FailureOr<Value> removeLeadingUnitTensorDimension(
+  OpBuilder& builder, Location loc, Value value, RankedTensorType resultType) {
+  if (value.getType() == resultType)
+    return value;
+  auto type = dyn_cast<RankedTensorType>(value.getType());
+  if (!type || !resultType || !type.hasStaticShape() || !resultType.hasStaticShape()
+      || type.getRank() != resultType.getRank() + 1 || type.getDimSize(0) != 1
+      || type.getElementType() != resultType.getElementType()
+      || !llvm::equal(type.getShape().drop_front(), resultType.getShape()))
+    return failure();
+  SmallVector<ReassociationIndices> reassociation;
+  if (resultType.getRank() != 0) {
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < resultType.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+  }
+  return tensor::CollapseShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
+}
+
 } // namespace onnx_mlir
@@ -25,4 +25,10 @@ mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
                              mlir::Value dest,
                              llvm::ArrayRef<mlir::OpFoldResult> offsets);

+mlir::FailureOr<mlir::Value>
+addLeadingUnitTensorDimension(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value);
+
+mlir::FailureOr<mlir::Value> removeLeadingUnitTensorDimension(
+  mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, mlir::RankedTensorType resultType);
+
 } // namespace onnx_mlir
@@ -1119,7 +1119,8 @@ struct CompiledCoreNode {
  enum class Kind : uint8_t {
    Op,
    Loop,
-    If
+    If,
+    IndexSwitch
  };

  Kind kind = Kind::Op;
@@ -1132,6 +1133,9 @@ struct CompiledCoreNode {
  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> loopBody;
  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> thenBody;
  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> elseBody;
+  llvm::SmallVector<int64_t> caseValues;
+  llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
 };

 static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
@@ -1231,6 +1235,31 @@ compileCoreEmissionPlan(Block& block, Operation* weightOwner, llvm::SmallVectorI
      continue;
    }

+    if (auto switchOp = dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = compileIndexExpr(switchOp.getArg());
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        return failure();
+      }
+      CompiledCoreNode switchNode;
+      switchNode.kind = CompiledCoreNode::Kind::IndexSwitch;
+      switchNode.op = switchOp.getOperation();
+      switchNode.condition = *selector;
+      llvm::append_range(switchNode.caseValues, switchOp.getCases());
+      for (mlir::Region& region : switchOp.getCaseRegions()) {
+        auto body = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+        if (failed(compileCoreEmissionPlan(region.front(), weightOwner, *body)))
+          return failure();
+        switchNode.caseBodies.push_back(std::move(body));
+      }
+      switchNode.defaultBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(
+            switchOp.getDefaultRegion().front(), weightOwner, *switchNode.defaultBody)))
+        return failure();
+      plan.push_back(std::move(switchNode));
+      continue;
+    }
+
    auto opKind = classifyCompiledCoreOpKind(op);
    if (failed(opKind)) {
      InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'";
@@ -1313,6 +1342,31 @@ static LogicalResult executeCompiledCorePlan(
      continue;
    }

+    if (node.kind == CompiledCoreNode::Kind::IndexSwitch) {
+      auto selector = node.condition.evaluate(knowledge);
+      auto switchOp = cast<mlir::scf::IndexSwitchOp>(node.op);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        return failure();
+      }
+      const llvm::SmallVectorImpl<CompiledCoreNode>* selectedBody = node.defaultBody.get();
+      mlir::Region* selectedRegion = &switchOp.getDefaultRegion();
+      for (auto [index, caseValue] : llvm::enumerate(node.caseValues))
+        if (caseValue == *selector) {
+          selectedBody = node.caseBodies[index].get();
+          selectedRegion = &switchOp.getCaseRegions()[index];
+          break;
+        }
+      if (failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge,
+                                         resolveWeightSlot, processedOperations,
+                                         batchLane, batchLaneCount)))
+        return failure();
+      auto yield = cast<mlir::scf::YieldOp>(selectedRegion->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
    switch (node.opKind) {
    case CompiledCoreOpKind::Load:
      coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);
@@ -1413,6 +1467,36 @@ static int64_t codeGenCoreOps(
  return failed(result) ? -1 : static_cast<int64_t>(processedOperations);
 }

+static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
+  std::string outputCorePath =
+    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
+  std::error_code errorCode;
+  raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
+  if (errorCode) {
+    errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
+    return InvalidOutputFileAccess;
+  }
+
+  pim_binary::writeHeader(coreBinaryStream);
+  pim_binary::patchInstructionCount(coreBinaryStream, 0);
+  coreBinaryStream.close();
+
+  if (!pimEmitJson.getValue())
+    return CompilerSuccess;
+
+  std::string outputCoreJsonPath =
+    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
+  errorCode = std::error_code();
+  raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
+  if (errorCode) {
+    errs() << "Error while opening core json file `" << outputCoreJsonPath << "`: " << errorCode.message() << '\n';
+    return InvalidOutputFileAccess;
+  }
+  coreJsonStream << "[]";
+  coreJsonStream.close();
+  return CompilerSuccess;
+}
+
 OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::string& outputDirPath) {
  if (!outputDirPath.empty()) {
    if (auto error = sys::fs::create_directory(outputDirPath)) {
@@ -1657,6 +1741,13 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
    if (jobResults[jobIndex].status != CompilerSuccess)
      return jobResults[jobIndex].status;

+  if (jobs.empty()) {
+    if (auto err = emitEmptyCoreArtifacts(outputDirPath, 0))
+      return err;
+    xbarsPerArrayGroup["core0"] = json::Array {};
+    memory.recordCoreReport(0, MemoryReportRow {});
+  }
+
  llvm::SmallVector<WeightFileRequest, 8> weightRequests;
  weightRequests.reserve(jobs.size());
  for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
@@ -59,13 +59,15 @@ llvm::cl::opt<PimConvLoweringType> pimConvLowering(

 llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow(
  "pim-export-spatial-dataflow",
-  llvm::cl::desc("Emit Gephi-importable CSV dataflow reports around MergeComputeNodes materialization"),
+  llvm::cl::desc("Emit Gephi-importable CSV dataflow reports for Spatial pipeline snapshots"),
  llvm::cl::values(clEnumValN(SpatialDataflowExportNone, "none", "Do not emit Spatial dataflow CSV reports")),
-  llvm::cl::values(clEnumValN(SpatialDataflowExportPre, "pre", "Emit pre-materialization Spatial dataflow CSV reports")),
  llvm::cl::values(
-    clEnumValN(SpatialDataflowExportPost, "post", "Emit post-materialization Spatial dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit spatial1 graph dataflow CSV reports")),
  llvm::cl::values(
-    clEnumValN(SpatialDataflowExportBoth, "both", "Emit both pre- and post-materialization Spatial dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 scheduled dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 realized dataflow CSV reports")),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportAll, "all", "Emit all Spatial dataflow CSV reports")),
  llvm::cl::init(SpatialDataflowExportNone),
  llvm::cl::cat(OnnxMlirOptions));

@@ -44,9 +44,10 @@ typedef enum {

 typedef enum {
  SpatialDataflowExportNone = 0,
-  SpatialDataflowExportPre = 1,
-  SpatialDataflowExportPost = 2,
-  SpatialDataflowExportBoth = 3,
+  SpatialDataflowExportSpatial1 = 1,
+  SpatialDataflowExportSpatial2 = 2,
+  SpatialDataflowExportSpatial3 = 3,
+  SpatialDataflowExportAll = 4,
 } PimSpatialDataflowExportType;

 extern llvm::cl::OptionCategory OnnxMlirOptions;
@@ -291,7 +291,26 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord

      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
        auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
-        if (!forOp) {
+        auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp());
+        auto indexSwitch = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp());
+        if (ifOp) {
+          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+            if (operand != value)
+              continue;
+            pendingValues.push_back(ifOp.getResult(index));
+            appendAliasDescription(interval.aliasesFollowed, ifOp.getResult(index));
+          }
+        }
+        else if (indexSwitch) {
+          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+            if (operand != value)
+              continue;
+            pendingValues.push_back(indexSwitch.getResult(index));
+            appendAliasDescription(interval.aliasesFollowed,
+                                   indexSwitch.getResult(index));
+          }
+        }
+        else if (!forOp) {
          addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
        }
        else {
@@ -9,10 +9,12 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"

+#include <cstring>
 #include <utility>

 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -21,24 +23,7 @@ using namespace mlir;

 namespace onnx_mlir {

-namespace {
-
-static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
-  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
-}
-
-static bool hasConstantIndices(tensor::ExtractOp extractOp) {
-  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
-}
-
-static bool isStaticTensorResult(Operation* op) {
-  return llvm::all_of(op->getResultTypes(), [](Type type) {
-    auto shapedType = dyn_cast<ShapedType>(type);
-    return shapedType && shapedType.hasStaticShape();
-  });
-}
-
-static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
+FailureOr<DenseElementsAttr> transposeDenseElementsAttr(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
  auto tensorType = dyn_cast<RankedTensorType>(denseAttr.getType());
  if (!tensorType)
    return failure();
@@ -59,7 +44,45 @@ static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr den

  auto transposedType = RankedTensorType::get(transposedShape, tensorType.getElementType(), tensorType.getEncoding());
  if (denseAttr.isSplat())
-    return DenseElementsAttr::get(transposedType, denseAttr.getSplatValue<Attribute>());
+    return DenseElementsAttr::getFromRawBuffer(transposedType, denseAttr.getRawData());
+
+  const unsigned elementBitWidth = tensorType.getElementTypeBitWidth();
+  const ArrayRef<char> inputData = denseAttr.getRawData();
+  if (elementBitWidth % 8 == 0) {
+    const size_t elementBytes = elementBitWidth / 8;
+    const size_t expectedBytes = denseAttr.getNumElements() * elementBytes;
+    if (inputData.size() == expectedBytes) {
+      SmallVector<char> transposedData(expectedBytes);
+      if (rank == 2 && perms[0] == 1 && perms[1] == 0) {
+        const int64_t rows = tensorType.getDimSize(0);
+        const int64_t columns = tensorType.getDimSize(1);
+        for (int64_t row = 0; row < rows; ++row)
+          for (int64_t column = 0; column < columns; ++column)
+            std::memcpy(transposedData.data() + (column * rows + row) * elementBytes,
+                        inputData.data() + (row * columns + column) * elementBytes,
+                        elementBytes);
+        return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData);
+      }
+
+      SmallVector<int64_t> originalStrides = computeRowMajorStrides(tensorType.getShape());
+      SmallVector<int64_t> transposedStrides = computeRowMajorStrides(transposedShape);
+      SmallVector<int64_t> originalIndices(rank);
+      for (int64_t linearIndex = 0; linearIndex < tensorType.getNumElements(); ++linearIndex) {
+        int64_t remaining = linearIndex;
+        for (int64_t dim = 0; dim < rank; ++dim) {
+          originalIndices[dim] = originalStrides.empty() ? 0 : remaining / originalStrides[dim];
+          remaining = originalStrides.empty() ? 0 : remaining % originalStrides[dim];
+        }
+        int64_t transposedLinearIndex = 0;
+        for (int64_t dim = 0; dim < rank; ++dim)
+          transposedLinearIndex += originalIndices[perms[dim]] * transposedStrides[dim];
+        std::memcpy(transposedData.data() + transposedLinearIndex * elementBytes,
+                    inputData.data() + linearIndex * elementBytes,
+                    elementBytes);
+      }
+      return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData);
+    }
+  }

  SmallVector<Attribute> originalValues(denseAttr.getValues<Attribute>());
  SmallVector<Attribute> transposedValues(originalValues.size());
@@ -84,16 +107,30 @@ static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr den
  return DenseElementsAttr::get(transposedType, transposedValues);
 }

+namespace {
+
+static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
+  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
+}
+
+static bool hasConstantIndices(tensor::ExtractOp extractOp) {
+  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
+}
+
+static bool isStaticTensorResult(Operation* op) {
+  return llvm::all_of(op->getResultTypes(), [](Type type) {
+    auto shapedType = dyn_cast<ShapedType>(type);
+    return shapedType && shapedType.hasStaticShape();
+  });
+}
+
 static FailureOr<DenseElementsAttr> reshapeDenseElements(DenseElementsAttr denseAttr, RankedTensorType resultType) {
  auto sourceType = dyn_cast<RankedTensorType>(denseAttr.getType());
-  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements())
+  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getElementType() != resultType.getElementType())
    return failure();

-  if (denseAttr.isSplat())
-    return DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>());
-
-  SmallVector<Attribute> values(denseAttr.getValues<Attribute>());
-  return DenseElementsAttr::get(resultType, values);
+  return DenseElementsAttr::getFromRawBuffer(resultType, denseAttr.getRawData());
 }

 static FailureOr<DenseElementsAttr> extractSliceDenseElements(DenseElementsAttr denseAttr,
@@ -161,7 +198,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
    perm.reserve(transposeOp.getPermAttr().size());
    for (IntegerAttr attr : transposeOp.getPermAttr().getAsRange<IntegerAttr>())
      perm.push_back(attr.getInt());
-    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm);
    return succeeded(transposedAttr) ? *transposedAttr : nullptr;
  }

@@ -171,7 +208,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
      return nullptr;

    SmallVector<int64_t> perm(transposeOp.getPermutation().begin(), transposeOp.getPermutation().end());
-    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm);
    return succeeded(transposedAttr) ? *transposedAttr : nullptr;
  }

@@ -219,6 +256,9 @@ getCompileTimeSourceImpl(Operation* op, llvm::SmallPtrSetImpl<Operation*>& visit

  chainLength += 1;

+  if (!isShapingOnlyOp(op))
+    return std::nullopt;
+
  if (auto extractOp = dyn_cast<tensor::ExtractOp>(op))
    return hasConstantIndices(extractOp)
           ? getCompileTimeSourceImpl(extractOp.getTensor().getDefiningOp(), visited, chainLength)
@@ -4,6 +4,8 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"

+#include "llvm/ADT/ArrayRef.h"
+
 namespace onnx_mlir {

 struct CompileTimeSource {
@@ -19,4 +21,7 @@ bool isCompileTimeOp(mlir::Operation* op);

 mlir::DenseElementsAttr getHostConstDenseElementsAttr(mlir::Value value);

+mlir::FailureOr<mlir::DenseElementsAttr> transposeDenseElementsAttr(
+  mlir::DenseElementsAttr denseAttr, llvm::ArrayRef<int64_t> permutation);
+
 } // namespace onnx_mlir
@@ -20,6 +20,7 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Dialect/ONNX/ONNXOps.hpp"

@@ -330,17 +331,25 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
    RewritePatternSet helperPatterns(ctx);
    populateGemmPatterns(helperPatterns, ctx);
    populateTransposePatterns(helperPatterns, ctx);
-    if (failed(applyPartialConversion(moduleOp, helperTarget, std::move(helperPatterns)))) {
-      moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
-      signalPassFailure();
-      return;
+    FrozenRewritePatternSet frozenHelperPatterns(
+      std::move(helperPatterns));
+    SmallVector<Operation*> topLevelHelperOps;
+    funcOp.walk([&](Operation* op) {
+      if (isa<spatial::SpatGraphCompute,
+              spatial::SpatGraphComputeBatch>(op))
+        return WalkResult::skip();
+      if (isa<ONNXGemmOp, ONNXTransposeOp>(op))
+        topLevelHelperOps.push_back(op);
+      return WalkResult::advance();
+    });
+    for (Operation *helper : topLevelHelperOps) {
+      if (failed(applyPartialConversion(
+            helper, helperTarget, frozenHelperPatterns))) {
+        moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
+        signalPassFailure();
+        return;
+      }
    }
-    FrozenRewritePatternSet nestedHelperPatterns([&] {
-      RewritePatternSet patterns(ctx);
-      populateGemmPatterns(patterns, ctx);
-      populateTransposePatterns(patterns, ctx);
-      return patterns;
-    }());
    ConversionTarget nestedHelperTarget(*ctx);
    nestedHelperTarget.addLegalDialect<spatial::SpatialDialect,
                                       tensor::TensorDialect,
@@ -356,7 +365,8 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
        computeLikeOps.push_back(op);
    });
    for (Operation* op : computeLikeOps) {
-      if (failed(applyFullConversion(op, nestedHelperTarget, nestedHelperPatterns))) {
+      if (failed(applyFullConversion(
+            op, nestedHelperTarget, frozenHelperPatterns))) {
        op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering");
        signalPassFailure();
        return;
@@ -392,6 +402,12 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
      signalPassFailure();
    } else {
      dumpModule(moduleOp, "spatial1_graph");
+      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
+      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
+          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
+        signalPassFailure();
+        return;
+      }
    }

    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
@@ -5,7 +5,6 @@
 #include "llvm/ADT/SmallVector.h"

 #include <algorithm>
-#include <numeric>
 #include <optional>
 #include <type_traits>

@@ -122,14 +121,6 @@ static RankedTensorType getKeepdimsType(RankedTensorType inputType, Type element
  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
 }

-static RankedTensorType getCompactKeptType(RankedTensorType inputType, Type elementType, ArrayRef<bool> reducedAxes) {
-  SmallVector<int64_t> shape;
-  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
-    if (!isReduced)
-      shape.push_back(dim);
-  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
-}
-
 static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> shape;
  shape.reserve(inputType.getRank());
@@ -228,59 +219,80 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
  return (*batchOp).getResult(0);
 }

-static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
-                                              RankedTensorType keepdimsType,
-                                              RankedTensorType compactKeptType,
-                                              ArrayRef<bool> reducedAxes,
-                                              ConversionPatternRewriter& rewriter,
-                                              Location loc) {
-  auto batchType = cast<RankedTensorType>(batchValue.getType());
-  if (batchType == keepdimsType)
-    return batchValue;
+static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
+    Value batchValue, RankedTensorType keepdimsType,
+    ArrayRef<bool> reducedAxes, ConversionPatternRewriter& rewriter,
+    Location loc) {
+  auto batchType = dyn_cast<RankedTensorType>(batchValue.getType());
+  int64_t rank = keepdimsType.getRank();
+  if (!batchType || !batchType.hasStaticShape()
+      || !keepdimsType.hasStaticShape()
+      || static_cast<int64_t>(reducedAxes.size()) != rank
+      || batchType.getRank() != rank + 1
+      || batchType.getElementType() != keepdimsType.getElementType())
+    return failure();

-  SmallVector<ReassociationIndices> collapseToFlat {{}};
-  for (int64_t axis = 0; axis < batchType.getRank(); ++axis)
-    collapseToFlat.front().push_back(axis);
-
-  SmallVector<ReassociationIndices> expandFlatToCompact(1);
-  for (int64_t axis = 0; axis < compactKeptType.getRank(); ++axis)
-    expandFlatToCompact.front().push_back(axis);
-
-  SmallVector<ReassociationIndices> expandCompactToKeepdims;
-  ReassociationIndices pendingLeadingReducedAxes;
+  int64_t laneCount = 1;
+  SmallVector<int64_t> keptAxes;
+  SmallVector<int64_t> keptAxisStrides;
  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
-    if (isReduced) {
-      if (expandCompactToKeepdims.empty())
-        pendingLeadingReducedAxes.push_back(axis);
-      else
-        expandCompactToKeepdims.back().push_back(axis);
-      continue;
-    }
-
-    expandCompactToKeepdims.emplace_back();
-    auto& group = expandCompactToKeepdims.back();
-    group.append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
-    pendingLeadingReducedAxes.clear();
-    group.push_back(axis);
+    int64_t dim = keepdimsType.getDimSize(axis);
+    if (dim <= 0 || (isReduced && dim != 1))
+      return failure();
+    if (!isReduced)
+      keptAxes.push_back(axis);
  }
-  if (!pendingLeadingReducedAxes.empty())
-    expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
+  keptAxisStrides.resize(keptAxes.size(), 1);
+  for (int64_t index = static_cast<int64_t>(keptAxes.size()) - 1;
+       index >= 0; --index) {
+    keptAxisStrides[index] = laneCount;
+    int64_t dim = keepdimsType.getDimSize(keptAxes[index]);
+    if (laneCount > std::numeric_limits<int64_t>::max() / dim)
+      return failure();
+    laneCount *= dim;
+  }
+  if (batchType.getDimSize(0) != laneCount
+      || llvm::any_of(batchType.getShape().drop_front(),
+                      [](int64_t dim) { return dim != 1; }))
+    return failure();

-  if (batchType.getNumElements() != batchType.getDimSize(0))
-    return {};
-  auto reshapeCompute =
-    createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
-      auto flatType = RankedTensorType::get({batchType.getNumElements()}, batchType.getElementType(), batchType.getEncoding());
-      Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
-      Value compact = flat;
-      if (compactKeptType != flatType)
-        compact = tensor::ExpandShapeOp::create(rewriter, loc, compactKeptType, flat, expandFlatToCompact);
-      Value keepdims = compact;
-      if (keepdimsType != compactKeptType)
-        keepdims = tensor::ExpandShapeOp::create(rewriter, loc, keepdimsType, compact, expandCompactToKeepdims);
-      spatial::SpatYieldOp::create(rewriter, loc, keepdims);
-    });
-  return reshapeCompute.getResult(0);
+  SmallVector<int64_t> operandIndices(laneCount, 0);
+  SmallVector<int64_t> sourceSlots;
+  SmallVector<int64_t> sourceOffsets(laneCount, 0);
+  SmallVector<int64_t> fragmentOffsets;
+  sourceSlots.reserve(laneCount);
+  fragmentOffsets.reserve(laneCount * rank);
+  for (int64_t lane = 0; lane < laneCount; ++lane) {
+    sourceSlots.push_back(lane);
+    size_t keptAxisIndex = 0;
+    for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
+      if (isReduced) {
+        fragmentOffsets.push_back(0);
+        continue;
+      }
+      int64_t dim = keepdimsType.getDimSize(axis);
+      fragmentOffsets.push_back(
+        (lane / keptAxisStrides[keptAxisIndex]) % dim);
+      ++keptAxisIndex;
+    }
+  }
+  SmallVector<int64_t> fragmentSizes(fragmentOffsets.size(), 1);
+  SmallVector<int64_t> fragmentStrides(fragmentOffsets.size(), 1);
+  return spatial::SpatBlueprintOp::create(
+           rewriter, loc, keepdimsType, batchValue, ValueRange {},
+           rewriter.getStringAttr("nchw"),
+           rewriter.getStringAttr("fragmented"),
+           rewriter.getDenseI64ArrayAttr(fragmentOffsets),
+           rewriter.getDenseI64ArrayAttr(fragmentSizes),
+           rewriter.getStringAttr("reduce_mean_keepdims_fragments"),
+           rewriter.getStringAttr("fragment_assembly"),
+           rewriter.getDenseI64ArrayAttr(operandIndices),
+           rewriter.getDenseI64ArrayAttr(sourceSlots),
+           rewriter.getDenseI64ArrayAttr(sourceOffsets),
+           rewriter.getDenseI64ArrayAttr(fragmentStrides),
+           rewriter.getStringAttr("disjoint"),
+           rewriter.getStringAttr("complete"))
+    .getOutput();
 }

 static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> reducedAxes) {
@@ -357,26 +369,36 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {

    Location loc = reduceMeanOp.getLoc();
    RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType());
-    RankedTensorType compactKeptType = getCompactKeptType(inputType, resultType.getElementType(), reducedAxes);
    RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
    int64_t laneCount = 1;
-    for (int64_t dim : compactKeptType.getShape())
+    for (auto [dim, isReduced] : llvm::zip_equal(keepdimsType.getShape(), reducedAxes)) {
+      if (isReduced)
+        continue;
+      if (dim <= 0 || laneCount > std::numeric_limits<int32_t>::max() / dim)
+        return rewriter.notifyMatchFailure(
+          reduceMeanOp, "ReduceMean physical lane count is not representable");
      laneCount *= dim;
+    }
    RankedTensorType batchType = getLanePackedKeepdimsType(laneCount, leafType);

    auto lanePackedKeepdims =
      buildReduceMeanKeepdimsBatch(adaptor.getData(), reducedAxes, batchType, leafType, rewriter, loc);
    if (failed(lanePackedKeepdims))
      return failure();
-    Value reducedKeepdims =
-      buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc);
+    auto reducedKeepdims = buildReduceMeanKeepdimsBlueprint(
+      *lanePackedKeepdims, keepdimsType, reducedAxes, rewriter, loc);
+    if (failed(reducedKeepdims))
+      return rewriter.notifyMatchFailure(
+        reduceMeanOp,
+        "cannot build physical-fragment ReduceMean keepdims reconstruction");

    if (semantics->keepdims != 0) {
-      rewriter.replaceOp(reduceMeanOp, reducedKeepdims);
+      rewriter.replaceOp(reduceMeanOp, *reducedKeepdims);
      return success();
    }

-    Value reduced = squeezeReducedAxes(reducedKeepdims, resultType, reducedAxes, rewriter, loc);
+    Value reduced = squeezeReducedAxes(
+      *reducedKeepdims, resultType, reducedAxes, rewriter, loc);
    rewriter.replaceOp(reduceMeanOp, reduced);
    return success();
  }
@@ -5,7 +5,7 @@

 #include "llvm/ADT/SmallVector.h"

-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -52,35 +52,12 @@ static FailureOr<Value> materializeTransposedConstant(Value input,
    return failure();
  }

-  if (denseAttr.isSplat())
-    return getOrCreateConstant(rewriter,
-                               rewriter.getInsertionBlock()->getParentOp(),
-                               DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>()),
-                               resultType);
-
-  SmallVector<Attribute> inputValues(denseAttr.getValues<Attribute>());
-  SmallVector<Attribute> resultValues(inputValues.size());
-  SmallVector<int64_t> inputStrides = computeRowMajorStrides(inputType.getShape());
-  SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultType.getShape());
-  SmallVector<int64_t> inputIndices(inputType.getRank(), 0);
-
-  for (auto [linearIndex, value] : llvm::enumerate(inputValues)) {
-    int64_t remaining = static_cast<int64_t>(linearIndex);
-    for (int64_t dim = 0; dim < inputType.getRank(); ++dim) {
-      inputIndices[dim] = inputStrides.empty() ? 0 : remaining / inputStrides[dim];
-      remaining = inputStrides.empty() ? 0 : remaining % inputStrides[dim];
-    }
-
-    int64_t resultLinearIndex = 0;
-    for (int64_t dim = 0; dim < resultType.getRank(); ++dim)
-      resultLinearIndex += inputIndices[permutation[dim]] * resultStrides[dim];
-
-    resultValues[resultLinearIndex] = value;
-  }
-
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return failure();
  return getOrCreateConstant(rewriter,
                             rewriter.getInsertionBlock()->getParentOp(),
-                             DenseElementsAttr::get(resultType, resultValues),
+                             *transposedAttr,
                             resultType);
 }

@@ -400,6 +400,11 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
    if (isa<spatial::SpatYieldOp>(op))
      continue;

+    // Cloning a region-bearing operation may leave the rewriter inside that
+    // region. Every old-block operation is lowered at the core-batch body
+    // boundary.
+    rewriter.setInsertionPointToEnd(newBlock);
+
    if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
      std::optional<StringRef> modeAttr = blueprint.getMode();
      if (modeAttr && *modeAttr == "fragment_assembly") {
@@ -8,6 +8,8 @@
 #include <limits>

 #include "Common.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
@@ -192,21 +194,23 @@ forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
 }

 static mlir::Value
-createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index, int64_t stepBytes) {
+createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index,
+                    int64_t stepBytes, Operation *constantAnchor) {
  if (stepBytes == 0)
    return start;
-  mlir::Value step = arith::ConstantIndexOp::create(builder, loc, stepBytes);
-  mlir::Value scaled = arith::MulIOp::create(builder, loc, index, step).getResult();
-  return arith::AddIOp::create(builder, loc, start, scaled).getResult();
+  return createOrFoldAffineApply(
+    builder, loc, builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1) * stepBytes,
+    ValueRange {start, index}, constantAnchor);
 }

 static mlir::Value createIndexedOffset(OpBuilder& builder,
                                       Location loc,
                                       mlir::Value indexArg,
-                                       ArrayRef<int64_t> values) {
+                                       ArrayRef<int64_t> values,
+                                       Operation *constantAnchor) {
  assert(!values.empty() && "expected lane-indexed values");
  if (llvm::all_of(values.drop_front(), [&](int64_t value) { return value == values.front(); }))
-    return arith::ConstantIndexOp::create(builder, loc, values.front());
+    return getOrCreateIndexConstant(builder, constantAnchor, values.front());

  if (values.size() >= 2) {
    int64_t step = values[1] - values[0];
@@ -214,21 +218,18 @@ static mlir::Value createIndexedOffset(OpBuilder& builder,
      return values[index] == values.front() + static_cast<int64_t>(index) * step;
    });
    if (arithmetic) {
-      mlir::Value base = arith::ConstantIndexOp::create(builder, loc, values.front());
-      mlir::Value stepValue = arith::ConstantIndexOp::create(builder, loc, step);
-      mlir::Value scaledIndex = arith::MulIOp::create(builder, loc, indexArg, stepValue).getResult();
-      return arith::AddIOp::create(builder, loc, base, scaledIndex).getResult();
+      return createOrFoldAffineApply(
+        builder, loc, builder.getAffineDimExpr(0) * step + values.front(),
+        ValueRange {indexArg}, constantAnchor);
    }
  }

-  mlir::Value selected = arith::ConstantIndexOp::create(builder, loc, values.front());
-  for (auto [lane, value] : llvm::enumerate(values.drop_front())) {
-    mlir::Value indexValue = arith::ConstantIndexOp::create(builder, loc, static_cast<int64_t>(lane + 1));
-    mlir::Value cmp = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::eq, indexArg, indexValue);
-    mlir::Value candidate = arith::ConstantIndexOp::create(builder, loc, value);
-    selected = arith::SelectOp::create(builder, loc, cmp, candidate, selected);
-  }
-  return selected;
+  RankedTensorType tableType = RankedTensorType::get(
+    {static_cast<int64_t>(values.size())}, builder.getI64Type());
+  DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
+  mlir::Value table = getOrCreateConstant(builder, constantAnchor, tableAttr, tableType);
+  mlir::Value selected = tensor::ExtractOp::create(builder, loc, table, ValueRange {indexArg});
+  return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), selected).getResult();
 }

 struct FragmentAssemblyCopyRunFamily {
@@ -433,11 +434,11 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
  mlir::Value hostStart;
  mlir::Value sourceStart;
  if (laneArg) {
-    hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane);
-    sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane);
+    hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane, anchor);
+    sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane, anchor);
  } else {
-    hostStart = arith::ConstantIndexOp::create(builder, loc, run.hostStartBytesByLane.front());
-    sourceStart = arith::ConstantIndexOp::create(builder, loc, run.sourceStartBytesByLane.front());
+    hostStart = getOrCreateIndexConstant(builder, anchor, run.hostStartBytesByLane.front());
+    sourceStart = getOrCreateIndexConstant(builder, anchor, run.sourceStartBytesByLane.front());
  }

  if (hostRunStartDelta)
@@ -459,9 +460,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
      .getOutput();
  }

-  mlir::Value lowerBound = arith::ConstantIndexOp::create(builder, loc, 0);
-  mlir::Value upperBound = arith::ConstantIndexOp::create(builder, loc, run.count);
-  mlir::Value step = arith::ConstantIndexOp::create(builder, loc, 1);
+  mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
+  mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, run.count);
+  mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
  FailureOr<NormalizedLoopResult> loop = buildNormalizedScfFor(
    builder,
    loc,
@@ -474,9 +475,10 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
        mlir::Value flatIndex,
        ValueRange iterArgs,
        SmallVectorImpl<mlir::Value>& yielded) {
-      mlir::Value hostOffset = createSteppedOffset(loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes);
+      mlir::Value hostOffset = createSteppedOffset(
+        loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes, anchor);
      mlir::Value sourceOffset =
-        createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes);
+        createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes, anchor);
      mlir::Value copied =
        pim::PimMemCopyDevToHostOp::create(loopBuilder,
                                           bodyLoc,
@@ -506,9 +508,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& build
    return emitFragmentAssemblyCopyRun(
      builder, loc, family.prototype, hostTarget, anchor, laneArg, baseHostOffset);

-  mlir::Value lowerBound = arith::ConstantIndexOp::create(builder, loc, 0);
-  mlir::Value upperBound = arith::ConstantIndexOp::create(builder, loc, family.sourceRunStartDeltas.size());
-  mlir::Value step = arith::ConstantIndexOp::create(builder, loc, 1);
+  mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
+  mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, family.sourceRunStartDeltas.size());
+  mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
  FailureOr<NormalizedLoopResult> outerLoop = buildNormalizedScfFor(
    builder,
    loc,
@@ -522,9 +524,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& build
        ValueRange iterArgs,
        SmallVectorImpl<mlir::Value>& yielded) {
      mlir::Value sourceRunStartDelta =
-        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas);
+        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas, anchor);
      mlir::Value hostRunStartDelta =
-        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas);
+        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas, anchor);
      FailureOr<mlir::Value> copied = emitFragmentAssemblyCopyRun(loopBuilder,
                                                                  bodyLoc,
                                                                  family.prototype,
@@ -10,7 +10,9 @@
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -180,16 +182,79 @@ static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute com
  return success();
 }

+static bool isHostMaterializableHelperOp(Operation* op) {
+  if (isa<spatial::SpatYieldOp>(op))
+    return true;
+  if (isa<arith::ConstantOp>(op) || op->hasTrait<OpTrait::ConstantLike>())
+    return true;
+  if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+    std::optional<StringRef> mode = blueprint.getMode();
+    return mode && *mode == "fragment_assembly";
+  }
+  return isShapingOnlyOp(op) || isPureIndexComputationOp(op);
+}
+
+static FailureOr<DenseMap<Value, Attribute>>
+analyzeHostMaterializableHelper(spatial::SpatScheduledCompute computeOp) {
+  DenseMap<Value, Attribute> folded;
+  for (auto [weightIndex, weight] : llvm::enumerate(computeOp.getWeights())) {
+    auto argument = computeOp.getWeightArgument(weightIndex);
+    if (!argument)
+      return failure();
+    Attribute constant;
+    if (matchPattern(weight, m_Constant(&constant)))
+      folded[*argument] = constant;
+  }
+  Block& block = computeOp.getBody().front();
+  for (Operation& op : block) {
+    if (!isHostMaterializableHelperOp(&op))
+      return failure();
+    if (isa<spatial::SpatYieldOp, spatial::SpatBlueprintOp>(op)
+        || (isShapingOnlyOp(&op) && !isPureIndexComputationOp(&op)))
+      continue;
+    if (isa<arith::ConstantOp>(op) || op.hasTrait<OpTrait::ConstantLike>()) {
+      for (Value result : op.getResults()) {
+        Attribute constant;
+        if (!matchPattern(result, m_Constant(&constant)))
+          return failure();
+        folded[result] = constant;
+      }
+      continue;
+    }
+    if (!isPureIndexComputationOp(&op) || op.getNumRegions() != 0)
+      return failure();
+    SmallVector<Attribute> operands;
+    for (Value operand : op.getOperands()) {
+      auto it = folded.find(operand);
+      if (it == folded.end())
+        return failure();
+      operands.push_back(it->second);
+    }
+    SmallVector<OpFoldResult> results;
+    if (failed(op.fold(operands, results))
+        || results.size() != op.getNumResults())
+      return failure();
+    for (auto [result, foldResult] : llvm::zip(op.getResults(), results)) {
+      auto attribute = dyn_cast<Attribute>(foldResult);
+      if (!attribute)
+        return failure();
+      folded[result] = attribute;
+    }
+  }
+  return folded;
+}
+
 static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatScheduledCompute computeOp,
                                                           IRRewriter& rewriter,
                                                           OperationFolder& constantFolder) {
  if (!computeOp.getInputs().empty() || computeOp.getNumResults() != 1)
    return false;
+  if (computeOp.getResult(0).use_empty())
+    return false;
  if (!llvm::all_of(computeOp.getResult(0).getUsers(), [](Operation* user) {
        return isa<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(user);
      }))
    return false;
-
  Block& block = computeOp.getBody().front();
  if (block.getNumArguments() != computeOp.getWeights().size())
    return false;
@@ -197,6 +262,9 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
  auto yieldOp = dyn_cast<spatial::SpatYieldOp>(block.getTerminator());
  if (!yieldOp || yieldOp.getNumOperands() != 1)
    return false;
+  auto folded = analyzeHostMaterializableHelper(computeOp);
+  if (failed(folded))
+    return false;

  rewriter.setInsertionPoint(computeOp);
  IRMapping mapping;
@@ -218,6 +286,20 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
      }
    }

+    if (isa<arith::ConstantOp>(op) || op.hasTrait<OpTrait::ConstantLike>()
+        || isPureIndexComputationOp(&op)) {
+      for (Value result : op.getResults()) {
+        auto it = folded->find(result);
+        if (it == folded->end())
+          return false;
+        mapping.map(
+          result,
+          getOrCreateConstant(constantFolder, computeOp, it->second,
+                              result.getType()));
+      }
+      continue;
+    }
+
    cloneMappedHelperOperands(&op, mapping, rewriter, constantFolder);
    Operation* clonedOp = rewriter.clone(op, mapping);
    for (auto [originalResult, newResult] : llvm::zip(op.getResults(), clonedOp->getResults()))
@@ -1,10 +1,23 @@
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

 using namespace mlir;

+static SmallVector<Region *> getSelectionRegions(OpResult result) {
+  SmallVector<Region *> regions;
+  if (auto selection = dyn_cast<scf::IndexSwitchOp>(result.getOwner()))
+    for (Region &region : selection->getRegions())
+      regions.push_back(&region);
+  else if (auto selection = dyn_cast<scf::IfOp>(result.getOwner())) {
+    regions.push_back(&selection.getThenRegion());
+    regions.push_back(&selection.getElseRegion());
+  }
+  return regions;
+}
+
 static bool isCoreBatchInputArgument(Value value) {
  auto blockArg = dyn_cast<BlockArgument>(value);
  if (!blockArg)
@@ -92,20 +105,46 @@ FailureOr<Value> onnx_mlir::pim::getPimAddressBase(Value value, const StaticValu
 }

 bool onnx_mlir::pim::isHostBackedPimAddress(Value value, const StaticValueKnowledge& knowledge) {
-  auto base = getPimStorageBase(value, knowledge);
-  if (failed(base))
-    return false;
-
-  if (isCoreBatchInputArgument(*base))
-    return true;
-
-  return isa_and_nonnull<memref::GetGlobalOp>(base->getDefiningOp());
+  llvm::SmallPtrSet<Value, 8> visited;
+  std::function<bool(Value)> isHost = [&](Value current) {
+    auto base = getPimStorageBase(current, knowledge);
+    if (failed(base) || !visited.insert(*base).second)
+      return false;
+    bool resultIsHost = isCoreBatchInputArgument(*base)
+                        || isa_and_nonnull<memref::GetGlobalOp>(base->getDefiningOp());
+    auto result = dyn_cast<OpResult>(*base);
+    SmallVector<Region *> regions = result ? getSelectionRegions(result)
+                                          : SmallVector<Region *>();
+    if (!resultIsHost && !regions.empty())
+      resultIsHost = llvm::all_of(regions, [&](Region *region) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        return yield && result.getResultNumber() < yield.getNumOperands()
+               && isHost(yield.getOperand(result.getResultNumber()));
+      });
+    visited.erase(*base);
+    return resultIsHost;
+  };
+  return isHost(value);
 }

 bool onnx_mlir::pim::isDeviceLocalPimAddress(Value value, const StaticValueKnowledge& knowledge) {
-  auto base = getPimStorageBase(value, knowledge);
-  if (failed(base))
-    return false;
-
-  return isa_and_nonnull<memref::AllocOp>(base->getDefiningOp());
+  llvm::SmallPtrSet<Value, 8> visited;
+  std::function<bool(Value)> isDevice = [&](Value current) {
+    auto base = getPimStorageBase(current, knowledge);
+    if (failed(base) || !visited.insert(*base).second)
+      return false;
+    bool resultIsDevice = isa_and_nonnull<memref::AllocOp>(base->getDefiningOp());
+    auto result = dyn_cast<OpResult>(*base);
+    SmallVector<Region *> regions = result ? getSelectionRegions(result)
+                                          : SmallVector<Region *>();
+    if (!resultIsDevice && !regions.empty())
+      resultIsDevice = llvm::all_of(regions, [&](Region *region) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        return yield && result.getResultNumber() < yield.getNumOperands()
+               && isDevice(yield.getOperand(result.getResultNumber()));
+      });
+    visited.erase(*base);
+    return resultIsDevice;
+  };
+  return isDevice(value);
 }
@@ -2,6 +2,8 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"

+#include "llvm/Support/MathExtras.h"
+
 #include "ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
@@ -33,6 +35,7 @@ struct CopyEndpointPlan {

 struct CopyLoopPlan {
  SmallVector<int64_t> outerShape;
+  int64_t outerElements = 0;
  int64_t chunkBytes = 0;
  ByteOffsetExpr targetBaseOffset;
  ByteOffsetExpr sourceBaseOffset;
@@ -74,6 +77,24 @@ static void appendTerm(ByteOffsetExpr& expr, Value value, int64_t scale) {
    expr.terms.push_back(ByteOffsetTerm {value, scale});
 }

+static FailureOr<int64_t> checkedPositiveMul(int64_t lhs, int64_t rhs) {
+  int64_t result = 0;
+  if (lhs < 0 || rhs < 0 || llvm::MulOverflow(lhs, rhs, result))
+    return failure();
+  return result;
+}
+
+static FailureOr<int64_t> checkedPositiveProduct(ArrayRef<int64_t> values) {
+  int64_t result = 1;
+  for (int64_t value : values) {
+    auto product = checkedPositiveMul(result, value);
+    if (failed(product))
+      return failure();
+    result = *product;
+  }
+  return result;
+}
+
 static FailureOr<SmallVector<int64_t>> getStaticMemRefStrides(MemRefType type) {
  SmallVector<int64_t> strides;
  int64_t offset = 0;
@@ -84,6 +105,165 @@ static FailureOr<SmallVector<int64_t>> getStaticMemRefStrides(MemRefType type) {
  return strides;
 }

+static FailureOr<SmallVector<int64_t>> getProvenMemRefStrides(Value value) {
+  llvm::SmallPtrSet<Value, 8> visiting;
+  std::function<FailureOr<SmallVector<int64_t>>(Value)> prove =
+    [&](Value current) -> FailureOr<SmallVector<int64_t>> {
+      auto type = dyn_cast<MemRefType>(current.getType());
+      if (!type || !visiting.insert(current).second)
+        return failure();
+      if (auto strides = getStaticMemRefStrides(type); succeeded(strides)) {
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto castOp = current.getDefiningOp<memref::CastOp>()) {
+        auto strides = prove(castOp.getSource());
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto subview = current.getDefiningOp<memref::SubViewOp>()) {
+        auto sourceStrides = prove(subview.getSource());
+        if (failed(sourceStrides) || subview.getSourceType().getRank() != subview.getType().getRank()) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides;
+        for (auto [sourceStride, viewStride] :
+             llvm::zip_equal(*sourceStrides, subview.getStaticStrides())) {
+          if (ShapedType::isDynamic(viewStride) || viewStride < 0) {
+            visiting.erase(current);
+            return failure();
+          }
+          auto stride = checkedPositiveMul(sourceStride, viewStride);
+          if (failed(stride)) {
+            visiting.erase(current);
+            return failure();
+          }
+          strides.push_back(*stride);
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto expand = current.getDefiningOp<memref::ExpandShapeOp>()) {
+        auto sourceStrides = prove(expand.getSrc());
+        auto resultType = dyn_cast<MemRefType>(expand.getResult().getType());
+        auto sourceType = dyn_cast<MemRefType>(expand.getSrc().getType());
+        if (failed(sourceStrides) || !sourceType || !resultType
+            || !resultType.hasStaticShape()
+            || sourceStrides->size() != static_cast<size_t>(sourceType.getRank())
+            || llvm::any_of(resultType.getShape(), [](int64_t dim) {
+                 return dim <= 0;
+               })) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides(resultType.getRank());
+        SmallVector<bool> assigned(resultType.getRank(), false);
+        for (auto [sourceDim, group] :
+             llvm::enumerate(expand.getReassociationIndices())) {
+          if (sourceDim >= sourceStrides->size() || group.empty()) {
+            visiting.erase(current);
+            return failure();
+          }
+          int64_t stride = (*sourceStrides)[sourceDim];
+          for (int64_t resultDim : llvm::reverse(group)) {
+            if (resultDim < 0 || resultDim >= resultType.getRank()
+                || assigned[resultDim]) {
+              visiting.erase(current);
+              return failure();
+            }
+            strides[resultDim] = stride;
+            assigned[resultDim] = true;
+            auto nextStride = checkedPositiveMul(
+              stride, resultType.getDimSize(resultDim));
+            if (failed(nextStride)) {
+              visiting.erase(current);
+              return failure();
+            }
+            stride = *nextStride;
+          }
+        }
+        if (llvm::is_contained(assigned, false)) {
+          visiting.erase(current);
+          return failure();
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto collapse = current.getDefiningOp<memref::CollapseShapeOp>()) {
+        auto sourceStrides = prove(collapse.getSrc());
+        auto sourceType = dyn_cast<MemRefType>(collapse.getSrc().getType());
+        if (failed(sourceStrides) || !sourceType
+            || !sourceType.hasStaticShape()
+            || sourceStrides->size() != static_cast<size_t>(sourceType.getRank())) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides;
+        for (ArrayRef<int64_t> group : collapse.getReassociationIndices()) {
+          if (group.empty()) {
+            visiting.erase(current);
+            return failure();
+          }
+          for (int64_t dim : group)
+            if (dim < 0 || dim >= sourceType.getRank()
+                || sourceType.getDimSize(dim) <= 0
+                || (*sourceStrides)[dim] < 0) {
+              visiting.erase(current);
+              return failure();
+            }
+          for (auto pair : llvm::zip(group.drop_back(), group.drop_front())) {
+            int64_t outer = std::get<0>(pair);
+            int64_t inner = std::get<1>(pair);
+            auto expectedOuterStride = checkedPositiveMul(
+              (*sourceStrides)[inner], sourceType.getDimSize(inner));
+            if (failed(expectedOuterStride)
+                || (*sourceStrides)[outer] != *expectedOuterStride) {
+              visiting.erase(current);
+              return failure();
+            }
+          }
+          strides.push_back((*sourceStrides)[group.back()]);
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      auto result = dyn_cast<OpResult>(current);
+      SmallVector<Region *> regions;
+      if (result) {
+        if (auto selection = dyn_cast<scf::IndexSwitchOp>(result.getOwner()))
+          for (Region &region : selection->getRegions())
+            regions.push_back(&region);
+        else if (auto selection = dyn_cast<scf::IfOp>(result.getOwner())) {
+          regions.push_back(&selection.getThenRegion());
+          regions.push_back(&selection.getElseRegion());
+        }
+      }
+      if (regions.empty()) {
+        visiting.erase(current);
+        return failure();
+      }
+      std::optional<SmallVector<int64_t>> common;
+      for (Region *region : regions) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        if (!yield || result.getResultNumber() >= yield.getNumOperands()) {
+          visiting.erase(current);
+          return failure();
+        }
+        auto strides = prove(yield.getOperand(result.getResultNumber()));
+        if (failed(strides) || (common && *common != *strides)) {
+          visiting.erase(current);
+          return failure();
+        }
+        common = std::move(*strides);
+      }
+      visiting.erase(current);
+      return common ? FailureOr<SmallVector<int64_t>>(std::move(*common))
+                    : FailureOr<SmallVector<int64_t>>(failure());
+    };
+  return prove(value);
+}
+
 static FailureOr<int64_t> getShapedByteSize(MemRefType type) {
  if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType()))
    return failure();
@@ -119,12 +299,15 @@ inferLogicalCopyShape(MemRefType targetType, MemRefType sourceType, int64_t size
  return failure();
 }

-static FailureOr<int64_t> getContiguousSuffixRank(MemRefType type, ArrayRef<int64_t> copyShape) {
-  if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())
+static FailureOr<int64_t> getContiguousSuffixRank(Value value, ArrayRef<int64_t> copyShape) {
+  auto type = dyn_cast<MemRefType>(value.getType());
+  if (!type || !type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())
      || type.getRank() != static_cast<int64_t>(copyShape.size()))
    return failure();
+  if (llvm::any_of(copyShape, [](int64_t dim) { return dim <= 0; }))
+    return failure();

-  auto strides = getStaticMemRefStrides(type);
+  auto strides = getProvenMemRefStrides(value);
  if (failed(strides))
    return failure();

@@ -134,7 +317,10 @@ static FailureOr<int64_t> getContiguousSuffixRank(MemRefType type, ArrayRef<int6
    if ((*strides)[dim] != expectedStride)
      break;
    ++contiguousSuffixRank;
-    expectedStride *= copyShape[dim];
+    auto nextStride = checkedPositiveMul(expectedStride, copyShape[dim]);
+    if (failed(nextStride))
+      return failure();
+    expectedStride = *nextStride;
  }
  return contiguousSuffixRank;
 }
@@ -174,18 +360,25 @@ static FailureOr<CopyEndpointPlan> analyzeCopyEndpoint(Value value, Value initia
    if (!sourceType || !sourceType.hasStaticShape() || !hasByteSizedElementType(sourceType.getElementType()))
      return failure();

-    auto sourceStrides = getStaticMemRefStrides(sourceType);
+    auto sourceStrides = getProvenMemRefStrides(subviewOp.getSource());
    if (failed(sourceStrides))
      return failure();

    int64_t elementByteWidth = static_cast<int64_t>(getElementTypeSizeInBytes(sourceType.getElementType()));
    for (auto [offset, stride] : llvm::zip_equal(subviewOp.getMixedOffsets(), *sourceStrides)) {
-      int64_t byteScale = stride * elementByteWidth;
+      auto byteScale = checkedPositiveMul(stride, elementByteWidth);
+      if (failed(byteScale))
+        return failure();
      if (auto attr = dyn_cast<Attribute>(offset)) {
-        endpoint.offset.constant += cast<IntegerAttr>(attr).getInt() * byteScale;
+        auto constantOffset = checkedPositiveMul(
+          cast<IntegerAttr>(attr).getInt(), *byteScale);
+        if (failed(constantOffset)
+            || llvm::AddOverflow(endpoint.offset.constant, *constantOffset,
+                                 endpoint.offset.constant))
+          return failure();
        continue;
      }
-      appendTerm(endpoint.offset, cast<Value>(offset), byteScale);
+      appendTerm(endpoint.offset, cast<Value>(offset), *byteScale);
    }

    endpoint.base = subviewOp.getSource();
@@ -213,8 +406,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
  auto sourceBytes = getShapedByteSize(sourceType);
  if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes)
      && *targetBytes == size && *sourceBytes == size) {
-    auto targetSuffixRank = getContiguousSuffixRank(targetType, targetType.getShape());
-    auto sourceSuffixRank = getContiguousSuffixRank(sourceType, sourceType.getShape());
+    auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape());
+    auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape());
    if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank)
        && *targetSuffixRank == targetType.getRank() && *sourceSuffixRank == sourceType.getRank()) {
      CopyRewritePlan plan;
@@ -230,8 +423,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
  if (failed(logicalCopyShape))
    return failure();

-  auto targetSuffixRank = getContiguousSuffixRank(targetType, *logicalCopyShape);
-  auto sourceSuffixRank = getContiguousSuffixRank(sourceType, *logicalCopyShape);
+  auto targetSuffixRank = getContiguousSuffixRank(target, *logicalCopyShape);
+  auto sourceSuffixRank = getContiguousSuffixRank(source, *logicalCopyShape);
  if (failed(targetSuffixRank) || failed(sourceSuffixRank))
    return failure();

@@ -246,8 +439,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
    return plan;
  }

-  auto targetStrides = getStaticMemRefStrides(targetType);
-  auto sourceStrides = getStaticMemRefStrides(sourceType);
+  auto targetStrides = getProvenMemRefStrides(target);
+  auto sourceStrides = getProvenMemRefStrides(source);
  if (failed(targetStrides) || failed(sourceStrides))
    return failure();

@@ -257,11 +450,27 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
  plan.loop.sourceBaseOffset = plan.source.offset;
  plan.loop.outerShape.assign(logicalCopyShape->begin(), logicalCopyShape->end() - contiguousSuffixRank);
  SmallVector<int64_t> chunkShape(logicalCopyShape->end() - contiguousSuffixRank, logicalCopyShape->end());
-  plan.loop.chunkBytes = getNumElements(chunkShape) * elementByteWidth;
-  for (int64_t stride : ArrayRef<int64_t>(*targetStrides).take_front(plan.loop.outerShape.size()))
-    plan.loop.targetOuterByteStrides.push_back(stride * elementByteWidth);
-  for (int64_t stride : ArrayRef<int64_t>(*sourceStrides).take_front(plan.loop.outerShape.size()))
-    plan.loop.sourceOuterByteStrides.push_back(stride * elementByteWidth);
+  auto outerElements = checkedPositiveProduct(plan.loop.outerShape);
+  auto chunkElements = checkedPositiveProduct(chunkShape);
+  auto chunkBytes = failed(chunkElements)
+                      ? FailureOr<int64_t>(failure())
+                      : checkedPositiveMul(*chunkElements, elementByteWidth);
+  if (failed(outerElements) || failed(chunkBytes))
+    return failure();
+  plan.loop.outerElements = *outerElements;
+  plan.loop.chunkBytes = *chunkBytes;
+  for (int64_t stride : ArrayRef<int64_t>(*targetStrides).take_front(plan.loop.outerShape.size())) {
+    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
+    if (failed(byteStride))
+      return failure();
+    plan.loop.targetOuterByteStrides.push_back(*byteStride);
+  }
+  for (int64_t stride : ArrayRef<int64_t>(*sourceStrides).take_front(plan.loop.outerShape.size())) {
+    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
+    if (failed(byteStride))
+      return failure();
+    plan.loop.sourceOuterByteStrides.push_back(*byteStride);
+  }
  if (plan.loop.chunkBytes <= 0)
    return failure();
  return plan;
@@ -361,7 +570,7 @@ static LogicalResult rewriteCopyLikeOp(CopyOp copyOp,
  }

  Value c0 = createIndexConstant(rewriter, anchorOp, 0);
-  Value cUpper = createIndexConstant(rewriter, anchorOp, getNumElements(plan->loop.outerShape));
+  Value cUpper = createIndexConstant(rewriter, anchorOp, plan->loop.outerElements);
  Value cStep = createIndexConstant(rewriter, anchorOp, 1);
  auto loop = buildNormalizedScfFor(
    rewriter,
@@ -17,6 +17,7 @@ add_pim_library(SpatialOps
  Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
+  Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
  Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
  Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -8,6 +8,7 @@
 #include "llvm/Support/LogicalResult.h"

 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

 using namespace mlir;
@@ -130,7 +131,7 @@ struct CanonicalizeSingleLaneComputeBatchPattern : OpRewritePattern<ComputeBatch
    rewriter.setInsertionPointToStart(newBlock);

    IRMapping mapper;
-    Value zero = arith::ConstantIndexOp::create(rewriter, compute.getLoc(), 0);
+    Value zero = getOrCreateIndexConstant(rewriter, compute.getOperation(), 0);
    mapper.map(*oldLaneArg, zero);
    for (auto [index, weight] : llvm::enumerate(compute.getWeights())) {
      auto oldArg = compute.getWeightArgument(index);
@@ -4,7 +4,6 @@

 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/DenseSet.h"

 namespace onnx_mlir::spatial {

@@ -23,54 +22,73 @@ static LogicalResult simulate(Operation *anchor,
                              ArrayRef<SmallVector<Event>> streams,
                              StringRef phase) {
  SmallVector<size_t> cursor(streams.size());
-  while (true) {
-    bool allFinished = true;
-    bool progressed = false;
-    for (unsigned stream = 0; stream < streams.size(); ++stream) {
-      if (cursor[stream] == streams[stream].size())
-        continue;
-      allFinished = false;
-      if (streams[stream][cursor[stream]].kind == EventKind::Compute) {
-        ++cursor[stream];
-        progressed = true;
-      }
-    }
-    if (allFinished)
-      return success();
+  DenseMap<uint64_t, unsigned> headSends;
+  DenseMap<uint64_t, unsigned> headReceives;
+  SmallVector<unsigned> readyComputes;
+  SmallVector<uint64_t> readyExchanges;
+  size_t computeCursor = 0;
+  size_t exchangeCursor = 0;
+  unsigned finishedStreams = 0;

-    for (unsigned source = 0; source < streams.size(); ++source) {
-      if (cursor[source] == streams[source].size())
-        continue;
-      const Event &send = streams[source][cursor[source]];
-      if (send.kind != EventKind::Send)
-        continue;
-      for (unsigned target = 0; target < streams.size(); ++target) {
-        if (cursor[target] == streams[target].size())
-          continue;
-        const Event &receive = streams[target][cursor[target]];
-        if (receive.kind == EventKind::Receive && receive.exchangeId == send.exchangeId) {
-          ++cursor[source];
-          ++cursor[target];
-          progressed = true;
-          break;
-        }
-      }
+  auto registerHead = [&](unsigned stream) {
+    if (cursor[stream] == streams[stream].size()) {
+      ++finishedStreams;
+      return;
    }
-    if (!progressed) {
-      InFlightDiagnostic diagnostic = anchor->emitError()
-                                      << phase << " communication rendezvous simulation made no progress";
-      unsigned reported = 0;
-      for (unsigned stream = 0; stream < streams.size() && reported < 8; ++stream) {
-        if (cursor[stream] == streams[stream].size())
-          continue;
-        const Event &event = streams[stream][cursor[stream]];
-        diagnostic << (reported == 0 ? "; blocked " : ", ") << "stream " << stream
-                   << " at exchange " << event.exchangeId;
-        ++reported;
-      }
-      return failure();
+    const Event &event = streams[stream][cursor[stream]];
+    if (event.kind == EventKind::Compute) {
+      readyComputes.push_back(stream);
+      return;
    }
+    DenseMap<uint64_t, unsigned> &heads =
+      event.kind == EventKind::Send ? headSends : headReceives;
+    DenseMap<uint64_t, unsigned> &peers =
+      event.kind == EventKind::Send ? headReceives : headSends;
+    heads[event.exchangeId] = stream;
+    if (peers.contains(event.exchangeId))
+      readyExchanges.push_back(event.exchangeId);
+  };
+  for (unsigned stream = 0; stream < streams.size(); ++stream)
+    registerHead(stream);
+
+  while (computeCursor != readyComputes.size()
+         || exchangeCursor != readyExchanges.size()) {
+    if (computeCursor != readyComputes.size()) {
+      unsigned stream = readyComputes[computeCursor++];
+      ++cursor[stream];
+      registerHead(stream);
+      continue;
+    }
+
+    uint64_t exchange = readyExchanges[exchangeCursor++];
+    auto send = headSends.find(exchange);
+    auto receive = headReceives.find(exchange);
+    if (send == headSends.end() || receive == headReceives.end())
+      continue;
+    unsigned source = send->second;
+    unsigned target = receive->second;
+    headSends.erase(send);
+    headReceives.erase(receive);
+    ++cursor[source];
+    ++cursor[target];
+    registerHead(source);
+    registerHead(target);
  }
+
+  if (finishedStreams == streams.size())
+    return success();
+  InFlightDiagnostic diagnostic = anchor->emitError()
+                                  << phase << " communication rendezvous simulation made no progress";
+  unsigned reported = 0;
+  for (unsigned stream = 0; stream < streams.size() && reported < 8; ++stream) {
+    if (cursor[stream] == streams[stream].size())
+      continue;
+    const Event &event = streams[stream][cursor[stream]];
+    diagnostic << (reported == 0 ? "; blocked " : ", ") << "stream " << stream
+               << " at exchange " << event.exchangeId;
+    ++reported;
+  }
+  return failure();
 }

 static std::optional<int64_t> getI64Attr(Operation *op, StringRef name) {
@@ -79,6 +97,98 @@ static std::optional<int64_t> getI64Attr(Operation *op, StringRef name) {
  return std::nullopt;
 }

+static LogicalResult getI64ArrayAttr(
+    Operation *op, StringRef name,
+    std::optional<SmallVector<int64_t>> &values) {
+  Attribute attr = op->getAttr(name);
+  if (!attr)
+    return success();
+  if (auto array = dyn_cast<DenseI64ArrayAttr>(attr)) {
+    values.emplace(array.asArrayRef());
+    return success();
+  }
+  auto elements = dyn_cast<DenseIntElementsAttr>(attr);
+  auto type = elements ? dyn_cast<RankedTensorType>(elements.getType())
+                       : RankedTensorType();
+  if (!elements || !type || type.getRank() != 1
+      || !type.getElementType().isInteger(64))
+    return op->emitOpError() << "has invalid " << name << " metadata";
+  values.emplace();
+  values->reserve(elements.getNumElements());
+  for (const APInt &value : elements.getValues<APInt>())
+    values->push_back(value.getSExtValue());
+  return success();
+}
+
+struct RealizedLogicalTransfer {
+  int64_t channelId = -1;
+  int64_t parentExchangeId = -1;
+  int64_t parentTransferCount = 0;
+  int64_t sourceCore = -1;
+  int64_t targetCore = -1;
+};
+
+static LogicalResult forEachRealizedLogicalTransfer(
+    Operation *op,
+    function_ref<LogicalResult(const RealizedLogicalTransfer &)> callback) {
+  auto scalarChannel = getI64Attr(op, "raptor.channel_id");
+  std::optional<SmallVector<int64_t>> batchChannels;
+  if (failed(getI64ArrayAttr(
+        op, "raptor.batch_channel_ids", batchChannels)))
+    return failure();
+  if (scalarChannel && batchChannels)
+    return op->emitOpError(
+      "mixes scalar and compact logical transfer metadata");
+
+  if (scalarChannel) {
+    auto exchange = getI64Attr(op, "raptor.exchange_id");
+    auto parent = getI64Attr(op, "raptor.parent_exchange_id");
+    auto count = getI64Attr(op, "raptor.parent_transfer_count");
+    auto source = getI64Attr(op, "raptor.source_core");
+    auto target = getI64Attr(op, "raptor.target_core");
+    if (!exchange || !parent || !count || !source || !target)
+      return op->emitOpError(
+        "is missing scalar logical transfer metadata");
+    RealizedLogicalTransfer transfer {
+      *scalarChannel, *parent, *count, *source, *target};
+    if (*exchange != transfer.channelId || transfer.channelId < 0
+        || transfer.parentExchangeId < 0 || transfer.parentTransferCount <= 0
+        || transfer.sourceCore < 0 || transfer.targetCore < 0)
+      return op->emitOpError("has invalid scalar logical transfer metadata");
+    return callback(transfer);
+  }
+
+  std::optional<SmallVector<int64_t>> sources, targets, parents, counts;
+  if (failed(getI64ArrayAttr(op, "raptor.batch_source_cores", sources))
+      || failed(getI64ArrayAttr(op, "raptor.batch_target_cores", targets))
+      || failed(getI64ArrayAttr(
+        op, "raptor.batch_parent_exchange_ids", parents))
+      || failed(getI64ArrayAttr(
+        op, "raptor.batch_parent_transfer_counts", counts)))
+    return failure();
+  if (!batchChannels || !sources || !targets || !parents || !counts)
+    return op->emitOpError(
+      "is missing compact logical transfer metadata");
+  size_t size = batchChannels->size();
+  if (size == 0 || sources->size() != size || targets->size() != size
+      || parents->size() != size || counts->size() != size)
+    return op->emitOpError(
+      "has non-parallel compact logical transfer metadata");
+  for (auto values : llvm::zip_equal(
+         *batchChannels, *parents, *counts, *sources, *targets)) {
+    RealizedLogicalTransfer transfer {
+      std::get<0>(values), std::get<1>(values), std::get<2>(values),
+      std::get<3>(values), std::get<4>(values)};
+    if (transfer.channelId < 0 || transfer.parentExchangeId < 0
+        || transfer.parentTransferCount <= 0 || transfer.sourceCore < 0
+        || transfer.targetCore < 0)
+      return op->emitOpError("has invalid compact logical transfer metadata");
+    if (failed(callback(transfer)))
+      return failure();
+  }
+  return success();
+}
+
 } // namespace

 LogicalResult verifyPlannedCommunicationDeadlockFree(
@@ -118,7 +228,11 @@ LogicalResult verifyPlannedCommunicationDeadlockFree(
 }

 LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) {
-  DenseMap<int64_t, SmallVector<Operation *, 2>> operationsByExchange;
+  struct LogicalOperation {
+    Operation *op = nullptr;
+    RealizedLogicalTransfer transfer;
+  };
+  DenseMap<int64_t, SmallVector<LogicalOperation, 2>> operationsByExchange;
  struct ParentExchange {
    std::optional<int64_t> expectedTransfers;
    DenseSet<int64_t> channels;
@@ -130,39 +244,25 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) {
  funcOp.walk([&](Operation *op) {
    if (!isa<SpatChannelSendOp, SpatChannelReceiveOp>(op))
      return;
-    std::optional<int64_t> exchangeId = getI64Attr(op, "raptor.exchange_id");
-    if (exchangeId)
-      operationsByExchange[*exchangeId].push_back(op);
-    if (auto channels = op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids"))
-      for (int64_t channel : channels.asArrayRef())
-        operationsByExchange[channel].push_back(op);
-    if (std::optional<int64_t> parent = getI64Attr(op, "raptor.parent_exchange_id")) {
-      ParentExchange &group = parentExchanges[*parent];
-      std::optional<int64_t> expected =
-        getI64Attr(op, "raptor.parent_transfer_count");
-      if (!expected || *expected <= 0
-          || (group.expectedTransfers && group.expectedTransfers != expected)) {
-        op->emitOpError(
-          "realized parent exchange has missing or inconsistent transfer count metadata");
-        invalid = true;
-      } else {
-        group.expectedTransfers = expected;
-      }
-      if (exchangeId)
-        group.channels.insert(*exchangeId);
-      if (auto channels =
-            op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids"))
-        group.channels.insert(channels.asArrayRef().begin(),
-                              channels.asArrayRef().end());
-    }
-    for (StringRef attrName : {"raptor.source_core", "raptor.target_core"})
-      if (std::optional<int64_t> core = getI64Attr(op, attrName); core && !llvm::is_contained(cores, *core))
-        cores.push_back(*core);
-    for (StringRef attrName : {"raptor.batch_source_cores", "raptor.batch_target_cores"})
-      if (auto batchCores = op->getAttrOfType<DenseI64ArrayAttr>(attrName))
-        for (int64_t core : batchCores.asArrayRef())
-          if (!llvm::is_contained(cores, core))
-            cores.push_back(core);
+    if (failed(forEachRealizedLogicalTransfer(
+          op, [&](const RealizedLogicalTransfer &transfer) -> LogicalResult {
+            operationsByExchange[transfer.channelId].push_back(
+              {op, transfer});
+            ParentExchange &parent =
+              parentExchanges[transfer.parentExchangeId];
+            if (parent.expectedTransfers
+                && *parent.expectedTransfers
+                     != transfer.parentTransferCount)
+              return op->emitOpError(
+                "declares an inconsistent parent transfer count");
+            parent.expectedTransfers = transfer.parentTransferCount;
+            parent.channels.insert(transfer.channelId);
+            for (int64_t core : {transfer.sourceCore, transfer.targetCore})
+              if (!llvm::is_contained(cores, core))
+                cores.push_back(core);
+            return success();
+          })))
+      invalid = true;
  });
  llvm::sort(cores);
  for (auto [index, core] : llvm::enumerate(cores))
@@ -172,38 +272,18 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) {
  funcOp.walk([&](Operation *op) {
    if (!isa<SpatChannelSendOp, SpatChannelReceiveOp>(op))
      return;
-    auto exchangeId = getI64Attr(op, "raptor.exchange_id");
-    if (!exchangeId) {
-      auto channels = op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids");
-      auto sourceCores = op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_source_cores");
-      auto targetCores = op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_target_cores");
-      if (!channels || !sourceCores || !targetCores
-          || channels.size() != sourceCores.size() || channels.size() != targetCores.size()) {
-        op->emitOpError("realized compact batch communication is missing channel/core metadata");
-        invalid = true;
-        return;
-      }
-      if (isa<SpatChannelSendOp>(op))
-        for (auto [channel, sourceCore] : llvm::zip(channels.asArrayRef(), sourceCores.asArrayRef()))
-          streams[streamByCore.lookup(sourceCore)].push_back(
-            {EventKind::Send, static_cast<uint64_t>(channel)});
-      else
-        for (auto [channel, targetCore] : llvm::zip(channels.asArrayRef(), targetCores.asArrayRef()))
-          streams[streamByCore.lookup(targetCore)].push_back(
-            {EventKind::Receive, static_cast<uint64_t>(channel)});
-      return;
-    }
-    auto sourceCore = getI64Attr(op, "raptor.source_core");
-    auto targetCore = getI64Attr(op, "raptor.target_core");
-    if (!sourceCore || !targetCore) {
-      op->emitOpError("realized communication is missing core metadata");
+    if (failed(forEachRealizedLogicalTransfer(
+          op, [&](const RealizedLogicalTransfer &transfer) {
+            unsigned stream = streamByCore.lookup(
+              isa<SpatChannelSendOp>(op) ? transfer.sourceCore
+                                         : transfer.targetCore);
+            streams[stream].push_back(
+              {isa<SpatChannelSendOp>(op) ? EventKind::Send
+                                          : EventKind::Receive,
+               static_cast<uint64_t>(transfer.channelId)});
+            return success();
+          })))
      invalid = true;
-      return;
-    }
-    if (isa<SpatChannelSendOp>(op))
-      streams[streamByCore.lookup(*sourceCore)].push_back({EventKind::Send, static_cast<uint64_t>(*exchangeId)});
-    else if (isa<SpatChannelReceiveOp>(op))
-      streams[streamByCore.lookup(*targetCore)].push_back({EventKind::Receive, static_cast<uint64_t>(*exchangeId)});
  });
  if (invalid)
    return failure();
@@ -216,56 +296,38 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) {
             << " does not contain its declared lane transfer set";

  for (const auto &entry : operationsByExchange) {
-    if (entry.second.size() != 2 || !isa<SpatChannelSendOp>(entry.second[0])
-        == !isa<SpatChannelSendOp>(entry.second[1]))
-      return funcOp.emitOpError() << "exchange " << entry.first << " does not have exactly one send and one receive";
-    auto send = dyn_cast<SpatChannelSendOp>(entry.second[0]);
-    auto receive = dyn_cast<SpatChannelReceiveOp>(entry.second[1]);
-    if (!send) {
-      send = cast<SpatChannelSendOp>(entry.second[1]);
-      receive = cast<SpatChannelReceiveOp>(entry.second[0]);
+    if (entry.second.size() != 2
+        || isa<SpatChannelSendOp>(entry.second[0].op)
+             == isa<SpatChannelSendOp>(entry.second[1].op)) {
+      return funcOp.emitOpError()
+             << "exchange " << entry.first
+             << " does not have exactly one send and one receive (sends="
+             << llvm::count_if(entry.second, [](const LogicalOperation &item) {
+                  return isa<SpatChannelSendOp>(item.op);
+                })
+             << ", receives="
+             << llvm::count_if(entry.second, [](const LogicalOperation &item) {
+                  return isa<SpatChannelReceiveOp>(item.op);
+                })
+             << ")";
    }
+    const LogicalOperation &first = entry.second[0];
+    const LogicalOperation &second = entry.second[1];
+    const LogicalOperation &sendRecord =
+      isa<SpatChannelSendOp>(first.op) ? first : second;
+    const LogicalOperation &receiveRecord =
+      isa<SpatChannelReceiveOp>(first.op) ? first : second;
+    auto send = cast<SpatChannelSendOp>(sendRecord.op);
+    auto receive = cast<SpatChannelReceiveOp>(receiveRecord.op);
    if (send.getInput().getType() != receive.getOutput().getType())
      return send.emitOpError("send and receive payload types do not match");
-    int64_t sendSource = 0;
-    int64_t sendTarget = 0;
-    if (auto channels = send->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids")) {
-      auto channelIt = llvm::find(channels.asArrayRef(), entry.first);
-      auto sources = send->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_source_cores");
-      auto targets = send->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_target_cores");
-      if (channelIt == channels.asArrayRef().end() || !sources || !targets)
-        return send.emitOpError("batch send channel metadata is incomplete");
-      size_t index = std::distance(channels.asArrayRef().begin(), channelIt);
-      sendSource = sources.asArrayRef()[index];
-      sendTarget = targets.asArrayRef()[index];
-    } else {
-      auto source = getI64Attr(send, "raptor.source_core");
-      auto target = getI64Attr(send, "raptor.target_core");
-      if (!source || !target)
-        return send.emitOpError("send core metadata is incomplete");
-      sendSource = *source;
-      sendTarget = *target;
-    }
-    int64_t receiveSource = 0;
-    int64_t receiveTarget = 0;
-    if (auto channels = receive->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids")) {
-      auto channelIt = llvm::find(channels.asArrayRef(), entry.first);
-      auto sources = receive->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_source_cores");
-      auto targets = receive->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_target_cores");
-      if (channelIt == channels.asArrayRef().end() || !sources || !targets)
-        return receive.emitOpError("batch receive channel metadata is incomplete");
-      size_t index = std::distance(channels.asArrayRef().begin(), channelIt);
-      receiveSource = sources.asArrayRef()[index];
-      receiveTarget = targets.asArrayRef()[index];
-    } else {
-      auto source = getI64Attr(receive, "raptor.source_core");
-      auto target = getI64Attr(receive, "raptor.target_core");
-      if (!source || !target)
-        return receive.emitOpError("receive core metadata is incomplete");
-      receiveSource = *source;
-      receiveTarget = *target;
-    }
-    if (receiveSource != sendSource || receiveTarget != sendTarget)
+    if (receiveRecord.transfer.sourceCore != sendRecord.transfer.sourceCore
+        || receiveRecord.transfer.targetCore
+             != sendRecord.transfer.targetCore
+        || receiveRecord.transfer.parentExchangeId
+             != sendRecord.transfer.parentExchangeId
+        || receiveRecord.transfer.parentTransferCount
+             != sendRecord.transfer.parentTransferCount)
      return receive.emitOpError("receive core metadata does not match its send");
  }
  return simulate(funcOp, streams, "realized");
@@ -8,6 +8,10 @@

 #include "llvm/ADT/SmallPtrSet.h"

+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+
 namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
@@ -112,11 +116,34 @@ static FailureOr<Value> buildBlueprintReconstruction(
  return result;
 }

-static bool isSupportedDeferredShapingOp(Operation *op) {
-  return isa<tensor::ExtractSliceOp, tensor::InsertSliceOp, tensor::CollapseShapeOp,
-             tensor::ExpandShapeOp, tensor::CastOp, tensor::EmptyOp, tensor::ExtractOp,
-             arith::ConstantOp, arith::IndexCastOp, arith::AddIOp, arith::SubIOp,
-             arith::MulIOp, affine::AffineApplyOp>(op);
+static FailureOr<Value> buildIndexSwitchSelection(OpBuilder &builder, Location loc,
+                                                  Value selector, ValueRange candidates,
+                                                  Operation *diagnosticOwner) {
+  if (candidates.empty())
+    return diagnosticOwner->emitOpError("direct selection requires at least one candidate"), failure();
+  Type type = candidates.front().getType();
+  if (llvm::any_of(candidates, [&](Value candidate) { return candidate.getType() != type; }))
+    return diagnosticOwner->emitOpError("direct selection requires identical candidate types"), failure();
+  if (candidates.size() == 1)
+    return candidates.front();
+
+  SmallVector<int64_t> cases;
+  for (int64_t index = 0; index < static_cast<int64_t>(candidates.size()) - 1; ++index)
+    cases.push_back(index);
+  auto selection = scf::IndexSwitchOp::create(
+    builder, loc, TypeRange {type}, selector, cases, cases.size());
+  auto buildYield = [&](Region &region, Value candidate) {
+    OpBuilder::InsertionGuard guard(builder);
+    Block *block = builder.createBlock(&region);
+    builder.setInsertionPointToEnd(block);
+    scf::YieldOp::create(builder, loc, candidate);
+  };
+  for (auto [region, candidate] : llvm::zip(selection.getCaseRegions(), candidates.drop_back()))
+    buildYield(region, candidate);
+  // The scheduled-lane verifier guarantees an in-range selector, so default is
+  // the final lane without an otherwise-unreachable extra branch.
+  buildYield(selection.getDefaultRegion(), candidates.back());
+  return selection.getResult(0);
 }

 static FailureOr<Value> buildSelectedDeferredSource(OpBuilder &builder, Location loc,
@@ -128,46 +155,28 @@ static FailureOr<Value> buildSelectedDeferredSource(OpBuilder &builder, Location
    return sourceBlockArgs.front();
  if (!scheduledLane || sourceOperandForScheduledLane.empty())
    return transfer.emitOpError("multiple deferred sources require the enclosing scheduled lane"), failure();
-  Value table = createI64LookupTableConstant(builder, transfer.getOperation(), sourceOperandForScheduledLane);
-  Value i64 = tensor::ExtractOp::create(builder, loc, table, ValueRange {scheduledLane}).getResult();
-  Value index = arith::IndexCastOp::create(builder, loc, builder.getIndexType(), i64).getResult();
-  auto type = dyn_cast<RankedTensorType>(sourceBlockArgs.front().getType());
-  if (!type || !type.hasStaticShape())
-    return transfer.emitOpError("multiple deferred sources require static ranked tensors"), failure();
-  for (Value source : sourceBlockArgs)
-    if (source.getType() != type)
-      return transfer.emitOpError("multiple deferred sources require identical tensor types"), failure();
-  SmallVector<int64_t> shape {static_cast<int64_t>(sourceBlockArgs.size())};
-  llvm::append_range(shape, type.getShape());
-  auto stacked = createEmptyTensorForType(builder, loc, RankedTensorType::get(shape, type.getElementType()));
-  if (failed(stacked))
-    return failure();
-  Value value = *stacked;
-  SmallVector<OpFoldResult> sizes, strides(shape.size(), builder.getIndexAttr(1));
-  sizes.push_back(builder.getIndexAttr(1));
-  for (int64_t dim : type.getShape()) sizes.push_back(builder.getIndexAttr(dim));
-  for (auto [i, source] : llvm::enumerate(sourceBlockArgs)) {
-    SmallVector<int64_t> expandedShape {1}; llvm::append_range(expandedShape, type.getShape());
-    SmallVector<ReassociationIndices> reassociation {{0, 1}};
-    for (int64_t dim = 1; dim < type.getRank(); ++dim) reassociation.push_back({dim + 1});
-    Value expanded = tensor::ExpandShapeOp::create(builder, loc,
-      RankedTensorType::get(expandedShape, type.getElementType()), source, reassociation).getResult();
-    SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0));
-    offsets[0] = builder.getIndexAttr(i);
-    value = tensor::InsertSliceOp::create(builder, loc, expanded, value, offsets, sizes, strides).getResult();
+  auto scheduled = transfer->getParentOfType<SpatScheduledComputeBatch>();
+  if (!scheduled || sourceOperandForScheduledLane.size() != static_cast<size_t>(scheduled.getLaneCount()))
+    return transfer.emitOpError("deferred source mapping must cover every scheduled lane"), failure();
+  SmallVector<Value> candidates;
+  candidates.reserve(sourceOperandForScheduledLane.size());
+  for (int64_t sourceIndex : sourceOperandForScheduledLane) {
+    if (sourceIndex < 0 || sourceIndex >= static_cast<int64_t>(sourceBlockArgs.size()))
+      return transfer.emitOpError("deferred source mapping operand is out of range"), failure();
+    candidates.push_back(sourceBlockArgs[sourceIndex]);
  }
-  SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0)); offsets[0] = index;
-  SmallVector<int64_t> sliceShape {1}; llvm::append_range(sliceShape, type.getShape());
-  auto slice = tensor::ExtractSliceOp::create(builder, loc,
-    RankedTensorType::get(sliceShape, type.getElementType()), value, offsets, sizes, strides);
-  // extract has a leading unit dimension; remove it without changing the payload.
-  SmallVector<ReassociationIndices> reassociation {{0, 1}};
-  for (int64_t dim = 1; dim < type.getRank(); ++dim) reassociation.push_back({dim + 1});
-  return tensor::CollapseShapeOp::create(builder, loc, type, slice.getResult(), reassociation).getResult();
+  return buildIndexSwitchSelection(builder, loc, scheduledLane, candidates, transfer.getOperation());
 }

-static bool isTopLevelShaping(Operation *op, Block &body) {
-  return op->getBlock() == &body && isSupportedDeferredShapingOp(op);
+static bool isDeferredPayloadCandidateOp(Operation *op) {
+  return isShapingOnlyOp(op) || isCompileTimeOp(op) || isPureIndexComputationOp(op);
+}
+
+static bool isTopLevelDeferredOperation(Operation *op, Block &body,
+                                        const DeferredInputPlan &plan) {
+  (void)plan;
+  return op->getBlock() == &body
+      && (isDeferredPayloadCandidateOp(op) || isa<scf::ForOp>(op));
 }

 static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
@@ -180,7 +189,7 @@ static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
  Operation *op = value.getDefiningOp();
  if (op && op->hasTrait<OpTrait::ConstantLike>())
    return true;
-  if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second)
+  if (!op || !isTopLevelDeferredOperation(op, body, plan) || !seen.insert(op).second)
    return op && seen.contains(op);
  return llvm::all_of(op->getOperands(), [&](Value operand) { return isEligible(operand, body, plan, seen); });
 }
@@ -196,7 +205,7 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
    if (isa<BlockArgument>(value))
      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
    Operation *op = value.getDefiningOp();
-    if (!op || !isSupportedDeferredShapingOp(op))
+    if (!op || (!isDeferredPayloadCandidateOp(op) && !op->hasTrait<OpTrait::ConstantLike>()))
      return transfer.emitOpError("phase 1 cannot clone the scheduled graph-lane expression"), failure();
    for (Value operand : op->getOperands()) if (failed(cloneScheduledLane(operand))) return failure();
    Operation *copy = builder.clone(*op, mapping);
@@ -214,7 +223,7 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
    if (isa<BlockArgument>(value))
      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
    Operation *op = value.getDefiningOp();
-    if (!op || (!isTopLevelShaping(op, body) && !op->hasTrait<OpTrait::ConstantLike>()))
+    if (!op || (!isTopLevelDeferredOperation(op, body, plan) && !op->hasTrait<OpTrait::ConstantLike>()))
      return transfer.emitOpError("phase 1 payload shaping contains an unsupported operation"), failure();
    for (Value operand : op->getOperands()) if (failed(clone(operand))) return failure();
    Operation *copy = builder.clone(*op, mapping);
@@ -225,65 +234,32 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
 }

 static bool dependsOnGraphLane(Value value, Value graphLane, Block &body,
+                               const DeferredInputPlan &plan,
                               llvm::SmallPtrSetImpl<Operation *> &seen) {
  if (value == graphLane)
    return true;
  Operation *op = value.getDefiningOp();
-  if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second)
+  if (!op || !isTopLevelDeferredOperation(op, body, plan) || !seen.insert(op).second)
    return false;
-  return llvm::any_of(op->getOperands(), [&](Value operand) {
-    return dependsOnGraphLane(operand, graphLane, body, seen);
-  });
-}
-
-static FailureOr<Value> buildPayloadAggregate(OpBuilder &builder, Location loc,
-                                              ArrayRef<Value> payloads) {
-  auto payloadType = dyn_cast<RankedTensorType>(payloads.front().getType());
-  if (!payloadType || !payloadType.hasStaticShape())
-    return failure();
-  SmallVector<int64_t> shape {static_cast<int64_t>(payloads.size())};
-  llvm::append_range(shape, payloadType.getShape());
-  auto aggregateType = RankedTensorType::get(shape, payloadType.getElementType());
-  auto empty = createEmptyTensorForType(builder, loc, aggregateType);
-  if (failed(empty)) return failure();
-  Value aggregate = *empty;
-  SmallVector<OpFoldResult> sizes, strides(shape.size(), builder.getIndexAttr(1));
-  sizes.push_back(builder.getIndexAttr(1));
-  for (int64_t dim : payloadType.getShape()) sizes.push_back(builder.getIndexAttr(dim));
-  SmallVector<ReassociationIndices> reassociation {{0, 1}};
-  for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1});
-  SmallVector<int64_t> expandedShape {1}; llvm::append_range(expandedShape, payloadType.getShape());
-  auto expandedType = RankedTensorType::get(expandedShape, payloadType.getElementType());
-  for (auto [index, payload] : llvm::enumerate(payloads)) {
-    Value expanded = tensor::ExpandShapeOp::create(builder, loc, expandedType, payload, reassociation).getResult();
-    SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0));
-    offsets[0] = builder.getIndexAttr(index);
-    aggregate = tensor::InsertSliceOp::create(builder, loc, expanded, aggregate, offsets, sizes, strides).getResult();
+  if (auto loop = dyn_cast<scf::ForOp>(op)) {
+    bool depends = false;
+    loop.getRegion().walk([&](Operation *nested) {
+      depends |= llvm::is_contained(nested->getOperands(), graphLane);
+    });
+    if (depends)
+      return true;
  }
-  return aggregate;
-}
-
-static FailureOr<Value> selectPayloadAggregate(OpBuilder &builder, Location loc, Value aggregate,
-                                               Value localLane) {
-  auto aggregateType = cast<RankedTensorType>(aggregate.getType());
-  SmallVector<int64_t> payloadShape(aggregateType.getShape().begin() + 1, aggregateType.getShape().end());
-  auto payloadType = RankedTensorType::get(payloadShape, aggregateType.getElementType());
-  SmallVector<OpFoldResult> offsets(aggregateType.getRank(), builder.getIndexAttr(0)); offsets[0] = localLane;
-  SmallVector<OpFoldResult> sizes, strides(aggregateType.getRank(), builder.getIndexAttr(1));
-  sizes.push_back(builder.getIndexAttr(1));
-  for (int64_t dim : payloadShape) sizes.push_back(builder.getIndexAttr(dim));
-  SmallVector<int64_t> unitShape {1}; llvm::append_range(unitShape, payloadShape);
-  Value unit = tensor::ExtractSliceOp::create(builder, loc,
-    RankedTensorType::get(unitShape, aggregateType.getElementType()), aggregate, offsets, sizes, strides).getResult();
-  SmallVector<ReassociationIndices> reassociation {{0, 1}};
-  for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1});
-  return tensor::CollapseShapeOp::create(builder, loc, payloadType, unit, reassociation).getResult();
+  return llvm::any_of(op->getOperands(), [&](Value operand) {
+    return dependsOnGraphLane(operand, graphLane, body, plan, seen);
+  });
 }

 static void collectClosure(Value value, Block &body, const DeferredInputPlan &plan,
                           llvm::SmallPtrSetImpl<Operation *> &ops) {
  Operation *op = value.getDefiningOp();
-  if (!op || !isTopLevelShaping(op, body) || !ops.insert(op).second) return;
+  if (!op || !isTopLevelDeferredOperation(op, body, plan) || !ops.insert(op).second) return;
+  if (auto loop = dyn_cast<scf::ForOp>(op))
+    loop.getRegion().walk([&](Operation *nested) { ops.insert(nested); });
  for (Value operand : op->getOperands())
    if (operand != plan.graphInput && operand != plan.graphLane) collectClosure(operand, body, plan, ops);
 }
@@ -371,12 +347,12 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc
      if (!seen.insert(value).second) continue;
      for (OpOperand &use : value.getUses()) {
        Operation *user = use.getOwner();
-        if (!isTopLevelShaping(user, body)) { needsIdentity = true; continue; }
+        if (!isTopLevelDeferredOperation(user, body, plan)) { needsIdentity = true; continue; }
        llvm::SmallPtrSet<Operation *, 16> eligibility;
        if (!isEligible(user->getResult(0), body, plan, eligibility)) { needsIdentity = true; continue; }
        for (Value result : user->getResults()) {
-          bool hasShapingUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return isTopLevelShaping(next.getOwner(), body); });
-          bool hasOtherUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return !isTopLevelShaping(next.getOwner(), body); });
+          bool hasShapingUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return isTopLevelDeferredOperation(next.getOwner(), body, plan); });
+          bool hasOtherUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return !isTopLevelDeferredOperation(next.getOwner(), body, plan); });
          if (hasOtherUse) roots.push_back(result);
          if (hasShapingUse) worklist.push_back(result);
        }
@@ -388,7 +364,7 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc
    for (Value root : roots) {
      llvm::SmallPtrSet<Operation *, 16> laneDependencies;
      bool scalarize = plan.scalarizedGraphLaneBase
-                    && dependsOnGraphLane(root, plan.graphLane, body, laneDependencies);
+                    && dependsOnGraphLane(root, plan.graphLane, body, plan, laneDependencies);
      OpBuilder::InsertPoint restore = builder.saveInsertionPoint();
      Operation *loop = nullptr;
      if (scalarize) {
@@ -419,9 +395,8 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc
      if (failed(selected)) return failure();
      Value boundGraphLane;
      if (scalarize) {
-        Value offsetValue = arith::ConstantIndexOp::create(builder, loc, offset);
-        boundGraphLane = offset ? arith::AddIOp::create(builder, loc, plan.scalarizedGraphLaneBase, offsetValue).getResult()
-                                : plan.scalarizedGraphLaneBase;
+        boundGraphLane = affineAddConst(
+          builder, loc, plan.scalarizedGraphLaneBase, offset, transfer.getOperation());
      }
      auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected, boundGraphLane);
      if (failed(payload)) return failure();
@@ -431,30 +406,26 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc
      }
      if (scalarize) {
        builder.restoreInsertionPoint(restore);
-        if (payloads.size() == 1) {
-          mapper.map(root, payloads.front());
-        } else {
-          if (loop) builder.setInsertionPoint(loop);
-          else builder.setInsertionPointToEnd(plan.scalarizedHoistBlock);
-          auto aggregate = buildPayloadAggregate(builder, loc, payloads);
-          if (failed(aggregate)) return failure();
-          builder.restoreInsertionPoint(restore);
-          auto selected = selectPayloadAggregate(builder, loc, *aggregate, plan.scalarizedLocalLane);
-          if (failed(selected)) return failure();
-          mapper.map(root, *selected);
-        }
+        auto selected = buildIndexSwitchSelection(
+          builder, loc, plan.scalarizedLocalLane, payloads, root.getDefiningOp());
+        if (failed(selected)) return failure();
+        mapper.map(root, *selected);
      } else {
        mapper.map(root, payloads.front());
      }
      collectClosure(root, body, plan, absorbed);
    }
  }
+  SmallVector<Operation *> notFullyAbsorbed;
  for (Operation *op : absorbed) {
    bool allResultsMapped = llvm::all_of(op->getResults(), [&](Value result) {
      return mapper.contains(result) || llvm::all_of(result.getUses(), [&](OpOperand &use) { return absorbed.contains(use.getOwner()); });
    });
-    if (!allResultsMapped) absorbed.erase(op);
+    if (!allResultsMapped)
+      notFullyAbsorbed.push_back(op);
  }
+  for (Operation *op : notFullyAbsorbed)
+    absorbed.erase(op);
  return success();
 }

@@ -1,9 +1,11 @@
 #include "DeferredProjectionAnalysis.hpp"

-#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/Matchers.h"
-#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"

 #include <limits>

@@ -11,43 +13,83 @@ namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {

+static FailureOr<int64_t> getSignedInt64(IntegerAttr value) {
+  return value.getValue().isSignedIntN(64) ? FailureOr<int64_t>(value.getValue().getSExtValue())
+                                           : FailureOr<int64_t>(failure());
+}
+
+static FailureOr<int64_t> evaluate(Value value, const StaticIndexEnvironment &environment,
+                                   llvm::SmallDenseSet<Value, 16> &visiting);
+
+static FailureOr<int64_t> evaluateDenseExtract(tensor::ExtractOp extract,
+                                               const StaticIndexEnvironment &environment,
+                                               llvm::SmallDenseSet<Value, 16> &visiting) {
+  auto constant = extract.getTensor().getDefiningOp<arith::ConstantOp>();
+  auto elements = constant ? dyn_cast<DenseIntElementsAttr>(constant.getValue()) : DenseIntElementsAttr();
+  auto type = elements ? dyn_cast<RankedTensorType>(elements.getType()) : RankedTensorType();
+  if (!elements || !type || !type.hasStaticShape()
+      || extract.getIndices().size() != static_cast<size_t>(type.getRank()))
+    return failure();
+  int64_t linear = 0;
+  for (auto [index, dim] : llvm::zip(extract.getIndices(), type.getShape())) {
+    auto folded = evaluate(index, environment, visiting);
+    if (failed(folded) || *folded < 0 || *folded >= dim
+        || llvm::MulOverflow(linear, dim, linear) || llvm::AddOverflow(linear, *folded, linear))
+      return failure();
+  }
+  APInt value = elements.getValues<APInt>()[linear];
+  return value.isSignedIntN(64) ? FailureOr<int64_t>(value.getSExtValue())
+                                : FailureOr<int64_t>(failure());
+}
+
 static FailureOr<int64_t> evaluate(Value value, const StaticIndexEnvironment &environment,
                                   llvm::SmallDenseSet<Value, 16> &visiting) {
  if (auto it = environment.bindings.find(value); it != environment.bindings.end())
    return it->second;
+  Attribute constant;
+  if (matchPattern(value, m_Constant(&constant)))
+    if (auto integer = dyn_cast_or_null<IntegerAttr>(constant))
+      return getSignedInt64(integer);
+  if (isa<BlockArgument>(value) || !value.getDefiningOp())
+    return failure();
  if (!visiting.insert(value).second)
    return failure();
-  if (auto constant = value.getDefiningOp<arith::ConstantOp>())
-    if (auto integer = dyn_cast<IntegerAttr>(constant.getValue()))
-      { visiting.erase(value); return integer.getInt(); }
-  if (auto cast = value.getDefiningOp<arith::IndexCastOp>()) {
-    auto result = evaluate(cast.getIn(), environment, visiting); visiting.erase(value); return result;
-  }
-  auto binary = [&](auto op, auto fn) -> FailureOr<int64_t> {
-    auto lhs = evaluate(op.getLhs(), environment, visiting);
-    auto rhs = evaluate(op.getRhs(), environment, visiting);
-    if (failed(lhs) || failed(rhs)) return failure();
-    return fn(*lhs, *rhs);
-  };
-  if (auto op = value.getDefiningOp<arith::AddIOp>()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr<int64_t> { int64_t r; if (llvm::AddOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; }
-  if (auto op = value.getDefiningOp<arith::SubIOp>()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr<int64_t> { int64_t r; if (llvm::SubOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; }
-  if (auto op = value.getDefiningOp<arith::MulIOp>()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr<int64_t> { int64_t r; if (llvm::MulOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; }
-  if (auto apply = value.getDefiningOp<affine::AffineApplyOp>()) { auto result = evaluateAffineApply(apply, [&](Value operand) { return evaluate(operand, environment, visiting); }); visiting.erase(value); return result; }
  if (auto extract = value.getDefiningOp<tensor::ExtractOp>()) {
-    auto constant = extract.getTensor().getDefiningOp<arith::ConstantOp>();
-    auto elements = constant ? dyn_cast<DenseIntElementsAttr>(constant.getValue()) : DenseIntElementsAttr();
-    auto type = elements ? dyn_cast<RankedTensorType>(elements.getType()) : RankedTensorType();
-    if (!elements || !type || extract.getIndices().size() != static_cast<size_t>(type.getRank())) return failure();
-    int64_t linear = 0;
-    for (auto [index, dim] : llvm::zip(extract.getIndices(), type.getShape())) {
-      auto i = evaluate(index, environment, visiting);
-      if (failed(i) || *i < 0 || *i >= dim) return failure();
-      linear = linear * dim + *i;
+    auto result = evaluateDenseExtract(extract, environment, visiting);
+    visiting.erase(value);
+    return result;
+  }
+
+  Operation *definingOp = value.getDefiningOp();
+  if (definingOp->getNumRegions() != 0 || !isPureIndexComputationOp(definingOp)) {
+    visiting.erase(value);
+    return failure();
+  }
+  SmallVector<Attribute> operandConstants;
+  operandConstants.reserve(definingOp->getNumOperands());
+  Builder builder(definingOp->getContext());
+  for (Value operand : definingOp->getOperands()) {
+    auto folded = evaluate(operand, environment, visiting);
+    if (failed(folded)) {
+      visiting.erase(value);
+      return failure();
    }
-    visiting.erase(value); return elements.getValues<APInt>()[linear].getSExtValue();
+    operandConstants.push_back(builder.getIntegerAttr(operand.getType(), *folded));
+  }
+  SmallVector<OpFoldResult> results;
+  if (failed(definingOp->fold(operandConstants, results)) || results.size() != 1) {
+    visiting.erase(value);
+    return failure();
+  }
+  FailureOr<int64_t> result = failure();
+  if (auto integer = dyn_cast<Attribute>(results.front())) {
+    if (auto attr = dyn_cast<IntegerAttr>(integer))
+      result = getSignedInt64(attr);
+  } else if (auto foldedValue = dyn_cast<Value>(results.front())) {
+    result = evaluate(foldedValue, environment, visiting);
  }
  visiting.erase(value);
-  return failure();
+  return result;
 }

 static FailureOr<std::optional<unsigned>> sourceArgument(Value value, SpatDeferredCommunicationOp deferred,
@@ -57,50 +99,29 @@ static FailureOr<std::optional<unsigned>> sourceArgument(Value value, SpatDeferr
      argument && argument.getOwner() == &deferred.getBody().front()
      && argument.getArgNumber() < deferred.getSources().size())
    return std::optional<unsigned>(argument.getArgNumber());
-  // Phase 1's selector ends in collapse(extract_slice(stacked, table[lane])).
-  auto collapse = value.getDefiningOp<tensor::CollapseShapeOp>();
-  if (!collapse) return std::optional<unsigned>();
-  value = collapse.getSrc();
-  auto slice = value.getDefiningOp<tensor::ExtractSliceOp>();
-  if (!slice) return std::optional<unsigned>();
-  auto sourceType = dyn_cast<RankedTensorType>(slice.getSourceType());
-  if (!sourceType || slice.getMixedOffsets().size() != static_cast<size_t>(sourceType.getRank()))
+  auto result = dyn_cast<OpResult>(value);
+  auto selection = result ? dyn_cast<scf::IndexSwitchOp>(result.getOwner()) : scf::IndexSwitchOp();
+  if (!selection || result.getResultNumber() != 0 || selection.getNumResults() != 1)
    return std::optional<unsigned>();
-  for (unsigned dim = 1; dim < sourceType.getRank(); ++dim) {
-    auto offset = evaluateDeferredIndex(slice.getMixedOffsets()[dim], environment);
-    auto size = evaluateDeferredIndex(slice.getMixedSizes()[dim], environment);
-    auto stride = evaluateDeferredIndex(slice.getMixedStrides()[dim], environment);
-    if (failed(offset) || failed(size) || failed(stride) || *offset != 0 || *size != sourceType.getDimSize(dim) || *stride != 1)
-      return std::optional<unsigned>();
-  }
-  auto leadingSize = evaluateDeferredIndex(slice.getMixedSizes().front(), environment);
-  auto leadingStride = evaluateDeferredIndex(slice.getMixedStrides().front(), environment);
-  if (failed(leadingSize) || failed(leadingStride) || *leadingSize != 1 || *leadingStride != 1)
-    return std::optional<unsigned>();
-  auto selected = evaluateDeferredIndex(slice.getMixedOffsets().front(), environment);
-  if (failed(selected)) return failure();
-  Value stacked = slice.getSource();
-  while (auto cast = stacked.getDefiningOp<tensor::CastOp>()) stacked = cast.getSource();
-  while (auto insert = stacked.getDefiningOp<tensor::InsertSliceOp>()) stacked = insert.getDest();
-  // The stack is a chain. Find the insertion at the selected leading offset.
-  for (Value cursor = slice.getSource(); auto insert = cursor.getDefiningOp<tensor::InsertSliceOp>(); cursor = insert.getDest()) {
-    auto offset = evaluateDeferredIndex(insert.getMixedOffsets().front(), environment);
-    if (succeeded(offset) && *offset == *selected) {
-      Value source = insert.getSource();
-      if (auto expand = source.getDefiningOp<tensor::ExpandShapeOp>()) source = expand.getSrc();
-      if (auto arg = dyn_cast<BlockArgument>(source); arg && arg.getOwner() == &deferred.getBody().front())
-        return std::optional<unsigned>(arg.getArgNumber());
+  auto selector = evaluateDeferredIndex(selection.getArg(), environment);
+  if (failed(selector))
+    return failure();
+  Region *selectedRegion = &selection.getDefaultRegion();
+  for (auto [caseValue, region] : llvm::zip(selection.getCases(), selection.getCaseRegions()))
+    if (caseValue == *selector) {
+      selectedRegion = &region;
+      break;
    }
-  }
-  return std::optional<unsigned>();
-}
-
-static bool isResidual(Operation *op) {
-  return isa<tensor::CollapseShapeOp, tensor::ExpandShapeOp, tensor::CastOp,
-             tensor::ExtractSliceOp, tensor::InsertSliceOp, tensor::ConcatOp,
-             tensor::EmptyOp, tensor::ExtractOp, arith::ConstantOp,
-             arith::IndexCastOp, arith::AddIOp, arith::SubIOp, arith::MulIOp,
-             affine::AffineApplyOp>(op);
+  if (!selectedRegion->hasOneBlock())
+    return failure();
+  Block &block = selectedRegion->front();
+  auto yield = dyn_cast<scf::YieldOp>(block.getTerminator());
+  if (!yield || yield.getResults().size() != 1)
+    return failure();
+  for (Operation &op : block.without_terminator())
+    if (!isa<tensor::CastOp>(op))
+      return failure();
+  return sourceArgument(yield.getResults().front(), deferred, environment);
 }

 static SpatGraphComputeBatch graphBatchOwner(Value value) {
@@ -119,15 +140,394 @@ static Value getEnclosingScheduledLane(SpatDeferredCommunicationOp deferred,
  return block && !block->empty() ? block->getArgument(0) : Value();
 }

+static bool isInsideDeferredLoop(Operation *op,
+                                 SpatDeferredCommunicationOp deferred) {
+  for (Operation *parent = op->getParentOp(); parent && parent != deferred;
+       parent = parent->getParentOp())
+    if (isa<scf::ForOp>(parent))
+      return true;
+  return false;
+}
+
+static bool isAllowedStaticIndexExpression(
+    Value value, Value scheduledLane,
+    llvm::SmallDenseSet<Value, 16> &visiting) {
+  if (value == scheduledLane)
+    return true;
+  Attribute constant;
+  if (matchPattern(value, m_Constant(&constant)))
+    return true;
+  if (isa<BlockArgument>(value) || !value.getDefiningOp()
+      || !visiting.insert(value).second)
+    return false;
+  Operation *op = value.getDefiningOp();
+  bool allowed = op->getNumRegions() == 0
+      && (isPureIndexComputationOp(op) || isCompileTimeOp(op))
+      && llvm::all_of(op->getOperands(), [&](Value operand) {
+           return isAllowedStaticIndexExpression(operand, scheduledLane,
+                                                 visiting);
+         });
+  visiting.erase(value);
+  return allowed;
+}
+
+static bool isAllowedStaticIndexExpression(Value value, Value scheduledLane) {
+  llvm::SmallDenseSet<Value, 16> visiting;
+  return isAllowedStaticIndexExpression(value, scheduledLane, visiting);
+}
+
+static bool originatesFromDeferredSource(
+    Value value, SpatDeferredCommunicationOp deferred,
+    llvm::SmallDenseSet<Value, 16> &visited) {
+  if (!visited.insert(value).second)
+    return false;
+  if (auto argument = dyn_cast<BlockArgument>(value))
+    return argument.getOwner() == &deferred.getBody().front()
+        && argument.getArgNumber() < deferred.getSources().size();
+  Operation *op = value.getDefiningOp();
+  if (!op)
+    return false;
+  if (isa<scf::IndexSwitchOp>(op) && op->getBlock() == &deferred.getBody().front())
+    return true;
+  return llvm::any_of(op->getOperands(), [&](Value operand) {
+    return originatesFromDeferredSource(operand, deferred, visited);
+  });
+}
+
+static bool originatesFromDeferredSource(
+    Value value, SpatDeferredCommunicationOp deferred) {
+  llvm::SmallDenseSet<Value, 16> visited;
+  return originatesFromDeferredSource(value, deferred, visited);
+}
+
+static LogicalResult verifyCanonicalSourceSelector(
+    scf::IndexSwitchOp selection, SpatDeferredCommunicationOp deferred,
+    Value scheduledLane, int64_t laneCount) {
+  if (selection->getBlock() != &deferred.getBody().front()
+      || selection.getNumResults() != 1
+      || !selection.getArg().getType().isIndex()
+      || !isAllowedStaticIndexExpression(selection.getArg(), scheduledLane))
+    return selection.emitOpError("is not a canonical deferred source selector");
+  if (laneCount < 2
+      || selection.getCases().size() != static_cast<size_t>(laneCount - 1)
+      || selection.getCaseRegions().size() != selection.getCases().size())
+    return selection.emitOpError(
+      "must cover every non-default scheduled lane");
+  for (auto [index, caseValue] : llvm::enumerate(selection.getCases()))
+    if (caseValue != static_cast<int64_t>(index))
+      return selection.emitOpError(
+        "must use consecutive scheduled-lane cases starting at zero");
+
+  auto verifyRegion = [&](Region &region) -> LogicalResult {
+    if (!region.hasOneBlock())
+      return selection.emitOpError("source-selector region must have one block");
+    Block &block = region.front();
+    auto yield = dyn_cast<scf::YieldOp>(block.getTerminator());
+    if (!yield || yield.getResults().size() != 1
+        || yield.getResults().front().getType() != selection.getResult(0).getType())
+      return selection.emitOpError(
+        "source-selector region must yield one exact result type");
+    for (Operation &op : block.without_terminator())
+      if (!isa<tensor::CastOp>(op))
+        return selection.emitOpError(
+          "source-selector regions may contain only tensor.cast before scf.yield");
+    Value source = yield.getResults().front();
+    while (auto cast = source.getDefiningOp<tensor::CastOp>()) {
+      if (cast->getBlock() != &block)
+        return selection.emitOpError(
+          "source-selector casts must be local to their region");
+      source = cast.getSource();
+    }
+    auto argument = dyn_cast<BlockArgument>(source);
+    if (!argument || argument.getOwner() != &deferred.getBody().front()
+        || argument.getArgNumber() >= deferred.getSources().size())
+      return selection.emitOpError(
+        "source-selector branch must resolve to a deferred source argument");
+    return success();
+  };
+  for (Region &region : selection.getCaseRegions())
+    if (failed(verifyRegion(region)))
+      return failure();
+  return verifyRegion(selection.getDefaultRegion());
+}
+
+static bool valueDependsOnAny(
+    Value value, const llvm::SmallDenseSet<Value, 16> &dependencies,
+    llvm::SmallDenseSet<Value, 16> &visited) {
+  if (dependencies.contains(value))
+    return true;
+  if (!visited.insert(value).second)
+    return false;
+  Operation *op = value.getDefiningOp();
+  return op && llvm::any_of(op->getOperands(), [&](Value operand) {
+    return valueDependsOnAny(operand, dependencies, visited);
+  });
+}
+
+static bool valueDependsOnAny(
+    Value value, const llvm::SmallDenseSet<Value, 16> &dependencies) {
+  llvm::SmallDenseSet<Value, 16> visited;
+  return valueDependsOnAny(value, dependencies, visited);
+}
+
+static LogicalResult specializeDeferredLoopStaticValues(
+    scf::ForOp loop, SpatDeferredCommunicationOp deferred,
+    const StaticIndexEnvironment &environment,
+    SpecializedDeferredProgram &program,
+    llvm::SmallDenseSet<Value, 16> dynamicIndices = {}) {
+  auto record = [&](Value value, StringRef diagnostic) -> LogicalResult {
+    auto folded = evaluateDeferredIndex(value, environment);
+    if (failed(folded))
+      return deferred.emitOpError(diagnostic);
+    program.staticValues.try_emplace(value, *folded);
+    return success();
+  };
+  auto step = evaluateDeferredIndex(loop.getStep(), environment);
+  if (failed(record(loop.getLowerBound(),
+                    "deferred shaping loop lower bound did not specialize"))
+      || failed(record(loop.getUpperBound(),
+                       "deferred shaping loop upper bound did not specialize"))
+      || failed(step) || *step <= 0)
+    return deferred.emitOpError(
+      "deferred shaping loop requires specialized bounds and a positive step");
+  program.staticValues.try_emplace(loop.getStep(), *step);
+
+  dynamicIndices.insert(loop.getInductionVar());
+  for (Operation &nested : loop.getBody()->without_terminator()) {
+    if (auto nestedLoop = dyn_cast<scf::ForOp>(nested)) {
+      if (failed(specializeDeferredLoopStaticValues(
+            nestedLoop, deferred, environment, program, dynamicIndices)))
+        return failure();
+      continue;
+    }
+    for (Value operand : nested.getOperands()) {
+      if ((!operand.getType().isIndex()
+           && !isa<IntegerType>(operand.getType()))
+          || valueDependsOnAny(operand, dynamicIndices))
+        continue;
+      if (failed(record(operand,
+                        "deferred shaping loop captured an index that did not specialize")))
+        return failure();
+    }
+  }
+  return success();
+}
+
+static FailureOr<std::optional<DeferredInsertAssembly>>
+analyzeDeferredInsertAssembly(
+    const SpecializedDeferredProgram &program,
+    const StaticIndexEnvironment &environment) {
+  auto finalInsert = program.yieldedValue.getDefiningOp<tensor::InsertSliceOp>();
+  if (!finalInsert || program.leaves.empty())
+    return std::optional<DeferredInsertAssembly>();
+
+  DenseMap<Value, unsigned> leafByRoot;
+  for (auto [leafIndex, leaf] : llvm::enumerate(program.leaves)) {
+    if (leaf.physicalSlots.size() != 1
+        || !leafByRoot.try_emplace(leaf.replacementRoot, leafIndex).second)
+      return std::optional<DeferredInsertAssembly>();
+  }
+
+  SmallPtrSet<Operation *, 32> consumed;
+  SmallVector<DeferredInsertAssemblyEntry> reverseEntries;
+  llvm::SmallDenseSet<unsigned, 16> insertedLeaves;
+  Value current = program.yieldedValue;
+  while (auto insert = current.getDefiningOp<tensor::InsertSliceOp>()) {
+    Value source = insert.getSource();
+    SmallVector<Operation *> rankShaping;
+    while (Operation *shape = source.getDefiningOp()) {
+      if (auto collapse = dyn_cast<tensor::CollapseShapeOp>(shape))
+        source = collapse.getSrc();
+      else if (auto expand = dyn_cast<tensor::ExpandShapeOp>(shape))
+        source = expand.getSrc();
+      else
+        break;
+      rankShaping.push_back(shape);
+    }
+    auto leafIt = leafByRoot.find(source);
+    if (leafIt == leafByRoot.end()
+        || !insertedLeaves.insert(leafIt->second).second)
+      return std::optional<DeferredInsertAssembly>();
+
+    DeferredInsertAssemblyEntry entry;
+    entry.leafIndex = leafIt->second;
+    entry.requirementIndex = leafIt->second;
+    auto foldGeometry = [&](ArrayRef<OpFoldResult> values,
+                            SmallVectorImpl<int64_t> &folded) {
+      for (OpFoldResult value : values) {
+        auto result = evaluateDeferredIndex(value, environment);
+        if (failed(result))
+          return failure();
+        folded.push_back(*result);
+      }
+      return success();
+    };
+    if (failed(foldGeometry(insert.getMixedOffsets(),
+                            entry.targetGeometry.offsets))
+        || failed(foldGeometry(insert.getMixedSizes(),
+                               entry.targetGeometry.sizes))
+        || failed(foldGeometry(insert.getMixedStrides(),
+                               entry.targetGeometry.strides)))
+      return std::optional<DeferredInsertAssembly>();
+    reverseEntries.push_back(std::move(entry));
+    consumed.insert(insert);
+    consumed.insert(rankShaping.begin(), rankShaping.end());
+    current = insert.getDest();
+  }
+
+  auto initial = current.getDefiningOp<tensor::EmptyOp>();
+  auto resultType = dyn_cast<RankedTensorType>(program.yieldedValue.getType());
+  if (!initial || !initial.getDynamicSizes().empty() || !resultType
+      || !resultType.hasStaticShape() || initial.getType() != resultType
+      || insertedLeaves.size() != program.leaves.size())
+    return std::optional<DeferredInsertAssembly>();
+  consumed.insert(initial);
+
+  for (const DeferredInsertAssemblyEntry &entry : reverseEntries) {
+    const StaticSliceGeometry &geometry = entry.targetGeometry;
+    if (geometry.offsets.size() != static_cast<size_t>(resultType.getRank())
+        || geometry.sizes.size() != geometry.offsets.size()
+        || geometry.strides.size() != geometry.offsets.size())
+      return std::optional<DeferredInsertAssembly>();
+    for (auto [offset, size, stride, dim] :
+         llvm::zip_equal(geometry.offsets, geometry.sizes,
+                         geometry.strides, resultType.getShape())) {
+      int64_t span = 0;
+      int64_t last = 0;
+      if (offset < 0 || size <= 0 || stride <= 0
+          || llvm::MulOverflow(size - 1, stride, span)
+          || llvm::AddOverflow(offset, span, last) || last >= dim)
+        return std::optional<DeferredInsertAssembly>();
+    }
+  }
+  if (llvm::any_of(program.residualOps,
+                   [&](Operation *op) { return !consumed.contains(op); })
+      || consumed.size() != program.residualOps.size())
+    return std::optional<DeferredInsertAssembly>();
+
+  DeferredInsertAssembly assembly;
+  assembly.initialValue = initial;
+  assembly.resultType = resultType;
+  assembly.entries.assign(reverseEntries.rbegin(), reverseEntries.rend());
+  return std::optional<DeferredInsertAssembly>(std::move(assembly));
+}
+
 } // namespace

+LogicalResult verifyDeferredProgramContract(
+    SpatDeferredCommunicationOp deferred) {
+  if (!deferred.getBody().hasOneBlock())
+    return deferred.emitOpError("deferred program must have exactly one body block");
+  Block &body = deferred.getBody().front();
+  auto terminator = dyn_cast<SpatYieldOp>(body.getTerminator());
+  if (!terminator || terminator.getOutputs().size() != 1)
+    return deferred.emitOpError(
+      "deferred program must have exactly one yielded value");
+
+  Value scheduledLane;
+  int64_t laneCount = 1;
+  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>()) {
+    scheduledLane = getEnclosingScheduledLane(deferred, scheduled);
+    laneCount = scheduled.getLaneCount();
+    if (!scheduledLane || laneCount <= 0)
+      return deferred.emitOpError(
+        "deferred program cannot locate a valid enclosing scheduled lane");
+  }
+
+  bool invalid = false;
+  WalkResult result = deferred.getBody().walk([&](Operation *op) {
+    if (invalid)
+      return WalkResult::interrupt();
+    auto reject = [&](StringRef message) {
+      op->emitOpError(message);
+      invalid = true;
+      return WalkResult::interrupt();
+    };
+
+    if (auto yield = dyn_cast<SpatYieldOp>(op)) {
+      if (op != body.getTerminator())
+        return reject("must be the unique top-level deferred terminator");
+      return WalkResult::advance();
+    }
+    if (auto yield = dyn_cast<scf::YieldOp>(op)) {
+      Operation *parent = op->getParentOp();
+      if (!parent || !isa<scf::ForOp, scf::IndexSwitchOp>(parent)
+          || op != op->getBlock()->getTerminator())
+        return reject("is not a valid deferred control-flow terminator");
+      return WalkResult::advance();
+    }
+    if (auto selection = dyn_cast<scf::IndexSwitchOp>(op)) {
+      if (failed(verifyCanonicalSourceSelector(
+            selection, deferred, scheduledLane, laneCount))) {
+        invalid = true;
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    }
+    if (auto loop = dyn_cast<scf::ForOp>(op)) {
+      auto isStaticRankedTensor = [](Type type) {
+        auto tensor = dyn_cast<RankedTensorType>(type);
+        return tensor && tensor.hasStaticShape();
+      };
+      if (!llvm::all_of(loop.getInitArgs(), [&](Value value) {
+            return isStaticRankedTensor(value.getType());
+          })
+          || !llvm::all_of(loop.getResultTypes(), isStaticRankedTensor))
+        return reject(
+          "requires static ranked tensor iter arguments and results");
+      auto yield = dyn_cast<scf::YieldOp>(loop.getBody()->getTerminator());
+      if (!yield || yield.getResults().getTypes() != loop.getResultTypes())
+        return reject("yield types must exactly match loop result types");
+      for (Value bound : {loop.getLowerBound(), loop.getUpperBound(),
+                          loop.getStep()})
+        if (!isAllowedStaticIndexExpression(bound, scheduledLane))
+          return reject(
+            "bounds must use only constants, the scheduled lane, and pure index computation");
+      for (int64_t lane = 0; lane < laneCount; ++lane) {
+        StaticIndexEnvironment environment;
+        if (scheduledLane)
+          environment.bindings[scheduledLane] = lane;
+        auto lower = evaluateDeferredIndex(loop.getLowerBound(), environment);
+        auto upper = evaluateDeferredIndex(loop.getUpperBound(), environment);
+        auto step = evaluateDeferredIndex(loop.getStep(), environment);
+        if (failed(lower) || failed(upper) || failed(step) || *step <= 0)
+          return reject(
+            "bounds must specialize for every scheduled lane with a positive step");
+      }
+      return WalkResult::advance();
+    }
+    if (op->getNumRegions() != 0)
+      return reject("unsupported region operation in deferred program");
+    if (!isShapingOnlyOp(op) && !isCompileTimeOp(op)
+        && !isPureIndexComputationOp(op))
+      return reject(
+        "deferred program permits only shaping, compile-time, or pure index operations");
+
+    for (Value operand : op->getOperands()) {
+      if (auto argument = dyn_cast<BlockArgument>(operand)) {
+        Operation *owner = argument.getOwner()->getParentOp();
+        bool local = owner == deferred
+            || (owner && deferred->isAncestor(owner));
+        if (!local && operand != scheduledLane)
+          return reject("captures an unsupported external block argument");
+      }
+      if (isInsideDeferredLoop(op, deferred)
+          && originatesFromDeferredSource(operand, deferred))
+        return reject(
+          "deferred source projection must remain outside residual loops");
+    }
+    return WalkResult::advance();
+  });
+  return success(!invalid && !result.wasInterrupted());
+}
+
 FailureOr<int64_t> evaluateDeferredIndex(Value value, const StaticIndexEnvironment &environment) {
  llvm::SmallDenseSet<Value, 16> visiting;
  return evaluate(value, environment, visiting);
 }
 FailureOr<int64_t> evaluateDeferredIndex(OpFoldResult value, const StaticIndexEnvironment &environment) {
  if (auto attr = dyn_cast<Attribute>(value))
-    if (auto integer = dyn_cast<IntegerAttr>(attr)) return integer.getInt();
+    if (auto integer = dyn_cast<IntegerAttr>(attr)) return getSignedInt64(integer);
  if (auto dynamic = dyn_cast<Value>(value)) return evaluateDeferredIndex(dynamic, environment);
  return failure();
 }
@@ -152,6 +552,8 @@ FailureOr<ResolvedDeferredSource> requireResolvedDeferredSource(Value value, Spa

 FailureOr<SpecializedDeferredProgram> analyzeDeferredProgram(SpatDeferredCommunicationOp deferred,
                                                              std::optional<unsigned> targetScheduledLane) {
+  // The Phase 1 verifier must establish the deferred-program contract once;
+  // this analysis only specializes lane-dependent static semantics.
  Block &body = deferred.getBody().front();
  auto yield = dyn_cast<SpatYieldOp>(body.getTerminator());
  if (!yield || yield.getOutputs().size() != 1)
@@ -229,13 +631,37 @@ FailureOr<SpecializedDeferredProgram> analyzeDeferredProgram(SpatDeferredCommuni
      program.leaves.push_back(std::move(leaf));
      return success();
    }
+    if (value.getType().isIndex() || isa<IntegerType>(value.getType())) {
+      auto folded = evaluateDeferredIndex(value, environment);
+      if (failed(folded))
+        return deferred.emitOpError("deferred index expression is not statically evaluable after specialization");
+      program.staticValues.try_emplace(value, *folded);
+      return success();
+    }
    Operation *op = value.getDefiningOp();
-    if (!op || op->getBlock() != &body || !isResidual(op))
-      return deferred.emitOpError("deferred residual contains an unsupported operation");
+    if (!op || op->getBlock() != &body)
+      return deferred.emitOpError("deferred residual escapes its verified body");
+    if (auto loop = dyn_cast<scf::ForOp>(op)) {
+      if (failed(specializeDeferredLoopStaticValues(
+            loop, deferred, environment, program)))
+        return failure();
+      for (Value init : loop.getInitArgs())
+        if (failed(visit(init)))
+          return failure();
+      program.residualOps.push_back(op);
+      return success();
+    }
    for (Value operand : op->getOperands()) if (failed(visit(operand))) return failure();
    program.residualOps.push_back(op); return success();
  };
  if (failed(visit(program.yieldedValue))) return failure();
+  StaticIndexEnvironment assemblyEnvironment = environment;
+  for (auto [value, staticValue] : program.staticValues)
+    assemblyEnvironment.bindings.try_emplace(value, staticValue);
+  auto assembly = analyzeDeferredInsertAssembly(program, assemblyEnvironment);
+  if (failed(assembly))
+    return failure();
+  program.insertAssembly = std::move(*assembly);
  return std::move(program);
 }

@@ -35,6 +35,18 @@ struct DeferredProjectionLeaf {
  mlir::RankedTensorType reconstructedType;
 };

+struct DeferredInsertAssemblyEntry {
+  unsigned requirementIndex = 0;
+  unsigned leafIndex = 0;
+  StaticSliceGeometry targetGeometry;
+};
+
+struct DeferredInsertAssembly {
+  mlir::tensor::EmptyOp initialValue;
+  mlir::RankedTensorType resultType;
+  llvm::SmallVector<DeferredInsertAssemblyEntry, 0> entries;
+};
+
 struct SpecializedDeferredProgram {
  SpatDeferredCommunicationOp deferred;
  std::optional<unsigned> targetScheduledLane;
@@ -42,6 +54,8 @@ struct SpecializedDeferredProgram {
  mlir::Value yieldedValue;
  llvm::SmallVector<DeferredProjectionLeaf, 0> leaves;
  llvm::SmallVector<mlir::Operation *> residualOps;
+  llvm::DenseMap<mlir::Value, int64_t> staticValues;
+  std::optional<DeferredInsertAssembly> insertAssembly;
 };

 struct ResolvedDeferredSource {
@@ -55,6 +69,8 @@ mlir::FailureOr<std::optional<ResolvedDeferredSource>> tryResolveDeferredSource(
 mlir::FailureOr<ResolvedDeferredSource> requireResolvedDeferredSource(
  mlir::Value value, SpatDeferredCommunicationOp deferred,
  const StaticIndexEnvironment &environment);
+mlir::LogicalResult verifyDeferredProgramContract(
+  SpatDeferredCommunicationOp deferred);
 mlir::FailureOr<SpecializedDeferredProgram> analyzeDeferredProgram(
  SpatDeferredCommunicationOp deferred,
  std::optional<unsigned> targetScheduledLane);
@@ -1,6 +1,7 @@
 #include "ScheduledComputeMaterialization.hpp"
 #include "ScheduledComputeReport.hpp"
 #include "ScheduledComputeVerification.hpp"
+#include "SpatialDataflowCsvExporter.hpp"
 #include "DeferredCommunicationRealization.hpp"
 #include "DeferredCommunicationDeadlock.hpp"

@@ -82,6 +83,13 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
      return;
    }

+    SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial2)
+        && failed(exportSpatialDataflowCsvScheduled(funcOp, "spatial2_scheduled_no_comm", "spatial2"))) {
+      signalPassFailure();
+      return;
+    }
+
    dumpScheduledComputeReportAndModule(moduleOp,
                                        funcOp,
                                        schedule,
@@ -92,6 +100,7 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
      signalPassFailure();
      return;
    }
+    dumpModule(moduleOp, "spatial3_scheduled", /*assumeVerified=*/true);
    if (failed(verifyRealizedCommunicationDeadlockFree(funcOp))) {
      moduleOp.emitError("MergeComputeNodes final communication verification failed");
      signalPassFailure();
@@ -100,6 +109,11 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
      moduleOp.emitError("scheduled Spatial phase 2 verification failed");
      signalPassFailure();
+      return;
+    }
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
+        && failed(exportSpatialDataflowCsvScheduled(funcOp, "spatial3_scheduled", "spatial3"))) {
+      signalPassFailure();
    }
  }
 };
@@ -204,7 +204,6 @@ static FailureOr<Value> buildSourceLaneStartForScheduledLane(OpBuilder &builder,
                                                             Value scheduledLane,
                                                             const SourceLaneSelector &selector,
                                                             Operation *constantAnchor) {
-  (void)constantAnchor;
  if (selector.kind == SourceLaneSelector::Kind::Affine) {
    if (selector.affine.baseLaneStart == 0 && selector.affine.laneCount == 1)
      return scheduledLane;
@@ -214,9 +213,7 @@ static FailureOr<Value> buildSourceLaneStartForScheduledLane(OpBuilder &builder,
      expr = d0 * selector.affine.laneCount;
    if (selector.affine.baseLaneStart != 0)
      expr = expr + selector.affine.baseLaneStart;
-    return affine::AffineApplyOp::create(
-             builder, loc, AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr), ValueRange {scheduledLane})
-      .getResult();
+    return createOrFoldAffineApply(builder, loc, expr, ValueRange {scheduledLane}, constantAnchor);
  }

  if (!selector.table)
@@ -569,18 +566,13 @@ static SmallVector<OpFoldResult> buildScheduledOutputInsertOffsets(OpBuilder &bu
                                                                   Location loc,
                                                                   Value scheduledLane,
                                                                   int64_t lanesPerScheduledLane,
-                                                                   RankedTensorType localFragmentType) {
+                                                                   RankedTensorType localFragmentType,
+                                                                   Operation *constantAnchor) {
  SmallVector<OpFoldResult> offsets;
  Value scheduledOutputLane = scheduledLane;
  if (lanesPerScheduledLane != 1) {
-    scheduledOutputLane =
-      affine::AffineApplyOp::create(builder,
-                                    loc,
-                                    AffineMap::get(/*dimCount=*/1,
-                                                   /*symbolCount=*/0,
-                                                   builder.getAffineDimExpr(0) * lanesPerScheduledLane),
-                                    ValueRange {scheduledLane})
-        .getResult();
+    scheduledOutputLane = affineMulConst(
+      builder, loc, scheduledLane, lanesPerScheduledLane, constantAnchor);
  }
  offsets.push_back(scheduledOutputLane);
  offsets.append(localFragmentType.getRank() - 1, OpFoldResult(builder.getIndexAttr(0)));
@@ -707,14 +699,12 @@ static LogicalResult materializeMultiCpuPeftClass(
        [&](OpBuilder &builder, Location bodyLoc, Value innerLane, ValueRange iterArgs, SmallVectorImpl<Value> &yielded) -> LogicalResult {

          IRMapping mapper;
-          Value sourceLane =
-            affine::AffineApplyOp::create(builder,
-                                          bodyLoc,
-                                          AffineMap::get(/*dimCount=*/2,
-                                                         /*symbolCount=*/0,
-                                                         builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1)),
-                                          ValueRange {*sourceLaneStart, innerLane})
-              .getResult();
+          Value sourceLane = createOrFoldAffineApply(
+            builder,
+            bodyLoc,
+            builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1),
+            ValueRange {*sourceLaneStart, innerLane},
+            scheduled.getOperation());
          mapper.map(*batch.getLaneArgument(), sourceLane);
          for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
            mapper.map(*batch.getWeightArgument(index),
@@ -789,8 +779,13 @@ static LogicalResult materializeMultiCpuPeftClass(
      finalLocalFragments.assign(loop->results.begin(), loop->results.end());
    }

-    auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
-    rewriter.setInsertionPointToStart(&inParallel.getRegion().front());
+    struct Publication {
+      Value fragment;
+      SmallVector<OpFoldResult> offsets;
+      SmallVector<OpFoldResult> sizes;
+      SmallVector<OpFoldResult> strides;
+    };
+    SmallVector<Publication> publications;
    for (auto [resultIndex, localFragment] : llvm::enumerate(finalLocalFragments)) {
      auto localFragmentType = cast<RankedTensorType>(localFragment.getType());
      int64_t lanesPerScheduledLane = isa<SpatCompute>(representative.op) ? 1 : representative.laneCount;
@@ -799,22 +794,29 @@ static LogicalResult materializeMultiCpuPeftClass(
        scheduled.getLoc(),
        scheduledLane,
        lanesPerScheduledLane,
-        localFragmentType);
+        localFragmentType,
+        scheduled.getOperation());
      SmallVector<OpFoldResult> sizes;
      SmallVector<OpFoldResult> strides;
      for (int64_t dim : localFragmentType.getShape()) {
        sizes.push_back(rewriter.getIndexAttr(dim));
        strides.push_back(rewriter.getIndexAttr(1));
      }
+      publications.push_back(
+        {localFragment, std::move(offsets), std::move(sizes),
+         std::move(strides)});
+    }
+    auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
+    rewriter.setInsertionPointToStart(&inParallel.getRegion().front());
+    for (auto [resultIndex, publication] : llvm::enumerate(publications))
      tensor::ParallelInsertSliceOp::create(
        rewriter,
        scheduled.getLoc(),
-        localFragment,
+        publication.fragment,
        block->getArgument(getScheduledBatchResultArgBase(scheduled) + stepPlan.resultOffset + resultIndex),
-        offsets,
-        sizes,
-        strides);
-    }
+        publication.offsets,
+        publication.sizes,
+        publication.strides);
  }
  return success();
 }
@@ -22,6 +22,7 @@
 #include "Scheduling/ComputeInstanceUtils.hpp"
 #include "Scheduling/MergeSchedulingAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"

 namespace onnx_mlir {
 namespace spatial {
@@ -267,11 +268,9 @@ inline SmallVector<uint32_t> collectSourceLaneStarts(const ComputeStepTuple &ste
 }

 inline Value createI64LookupTableConstant(OpBuilder &builder, Operation *constantAnchor, ArrayRef<int64_t> values) {
-  OpBuilder::InsertionGuard guard(builder);
-  builder.setInsertionPoint(constantAnchor);
  RankedTensorType tableType = RankedTensorType::get({static_cast<int64_t>(values.size())}, builder.getI64Type());
  DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
-  return arith::ConstantOp::create(builder, constantAnchor->getLoc(), tableType, tableAttr).getResult();
+  return getOrCreateConstant(builder, constantAnchor, tableAttr, tableType);
 }

 inline FailureOr<Value> createEmptyTensorForType(OpBuilder &builder, Location loc, Type type) {
@@ -82,24 +82,36 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
  pim::CappedDiagnosticReporter diagnostics;
  GraphBatchPublicationCache publicationCache;
  funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
+    bool ownershipValid = true;
    for (Value source : transfer.getSources()) {
      auto result = dyn_cast<OpResult>(source);
      if (!result || !isa<SpatGraphCompute, SpatGraphComputeBatch>(result.getOwner())) {
+        ownershipValid = false;
        diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
          illegalOp->emitOpError("phase-check deferred communication source operand must be an original graph SSA result");
        });
      }
    }
    if (!transfer->getParentOfType<SpatScheduledCompute>() &&
-        !transfer->getParentOfType<SpatScheduledComputeBatch>())
+        !transfer->getParentOfType<SpatScheduledComputeBatch>()) {
+      ownershipValid = false;
      diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
        illegalOp->emitOpError("phase-check deferred communication must be inside a scheduled compute");
        });
+    }
+    if (!ownershipValid)
+      return;
+    if (failed(verifyDeferredProgramContract(transfer))) {
+      diagnostics.report(transfer.getOperation(), [&](Operation *) {});
+      return;
+    }
    if (auto scheduled = transfer->getParentOfType<SpatScheduledComputeBatch>()) {
      for (unsigned lane = 0; lane < static_cast<unsigned>(scheduled.getLaneCount()); ++lane) {
        auto program = analyzeDeferredProgram(transfer, lane);
-        if (failed(program))
+        if (failed(program)) {
+          diagnostics.report(transfer.getOperation(), [&](Operation *) {});
          continue;
+        }
        for (const DeferredProjectionLeaf &leaf : program->leaves) {
          if (leaf.kind == DeferredLeafKind::ScalarSource)
            continue;
@@ -110,8 +122,9 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
            diagnostics.report(transfer.getOperation(), [&](Operation *) {});
        }
      }
-    } else
-      (void)analyzeDeferredProgram(transfer, std::nullopt);
+    } else if (failed(analyzeDeferredProgram(transfer, std::nullopt))) {
+      diagnostics.report(transfer.getOperation(), [&](Operation *) {});
+    }
  });

  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed");
@@ -0,0 +1,898 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AsmState.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cstdint>
+#include <fstream>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace spatial {
+
+namespace {
+
+struct TopLevelOpInfo {
+  Operation* op = nullptr;
+  size_t opId = 0;
+  bool isScheduled = false;
+  std::optional<int32_t> scalarCore;
+};
+
+struct ExpandedNodeInfo {
+  std::string id;
+  std::optional<int32_t> core;
+  std::optional<uint32_t> lane;
+};
+
+struct ChannelSendRecord {
+  std::string sourceId;
+  std::optional<uint32_t> sourceLane;
+};
+
+enum class LogicalNodeSelector {
+  Scalar,
+  Lane,
+  RangeRepresentative,
+};
+
+struct ResolvedProducer {
+  Operation* op = nullptr;
+  size_t resultIndex = 0;
+  LogicalNodeSelector selector = LogicalNodeSelector::Scalar;
+  uint32_t lane = 0;
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 1;
+};
+
+struct EdgeSource {
+  std::string id;
+  std::optional<uint32_t> sourceLane;
+};
+
+using ScheduledNodeByGraphLane = DenseMap<std::pair<int64_t, uint32_t>, ExpandedNodeInfo>;
+
+void emitEdgeRow(std::fstream& edgesFile,
+                 StringRef sourceId,
+                 StringRef targetId,
+                 std::optional<uint64_t> byteSize,
+                 Type propagatedType,
+                 StringRef stage,
+                 std::optional<uint32_t> sourceLane,
+                 std::optional<uint32_t> targetLane,
+                 std::optional<int64_t> channelId);
+
+std::string csvEscape(StringRef field) {
+  bool needsQuotes = field.contains(',') || field.contains('"') || field.contains('\n') || field.contains('\r');
+  if (!needsQuotes)
+    return field.str();
+
+  std::string escaped;
+  escaped.reserve(field.size() + 2);
+  escaped.push_back('"');
+  for (char ch : field)
+    if (ch == '"')
+      escaped += "\"\"";
+    else
+      escaped.push_back(ch);
+  escaped.push_back('"');
+  return escaped;
+}
+
+void writeCsvRow(std::fstream& file, ArrayRef<std::string> fields) {
+  for (size_t i = 0; i < fields.size(); ++i) {
+    if (i != 0)
+      file << ",";
+    file << csvEscape(fields[i]);
+  }
+  file << "\n";
+}
+
+template <typename NumberT>
+std::string maybeNumber(std::optional<NumberT> value) {
+  if (!value)
+    return "";
+  return std::to_string(*value);
+}
+
+std::string stringifyType(Type type) {
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  type.print(os);
+  return os.str();
+}
+
+std::string stringifyValueAsOperand(Value value, AsmState& asmState) {
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  value.printAsOperand(os, asmState);
+  return os.str();
+}
+
+std::string stringifyResultSsaNames(Operation* op, AsmState* asmState) {
+  if (!asmState || op->getNumResults() == 0)
+    return "";
+
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  llvm::interleave(
+    op->getResults(), [&](Value result) { os << stringifyValueAsOperand(result, *asmState); }, [&]() { os << ";"; });
+  return os.str();
+}
+
+std::optional<uint64_t> getTypeSizeBytes(Type type) {
+  if (auto shapedType = dyn_cast<ShapedType>(type)) {
+    if (!shapedType.hasStaticShape() || !hasByteSizedElementType(shapedType.getElementType()))
+      return std::nullopt;
+    return static_cast<uint64_t>(getShapedTypeSizeInBytes(shapedType));
+  }
+
+  if (isa<IndexType>(type))
+    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
+  if (auto intType = dyn_cast<IntegerType>(type)) {
+    if (intType.getWidth() <= 0 || intType.getWidth() % 8 != 0)
+      return std::nullopt;
+    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
+  }
+  if (auto floatType = dyn_cast<FloatType>(type)) {
+    if (floatType.getWidth() <= 0 || floatType.getWidth() % 8 != 0)
+      return std::nullopt;
+    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
+  }
+  return std::nullopt;
+}
+
+std::string getScalarId(bool isScheduled, size_t opId) { return (isScheduled ? "sc:" : "gc:") + std::to_string(opId); }
+
+std::string getBatchLaneId(bool isScheduled, size_t opId, uint32_t lane) {
+  return (isScheduled ? "scb:" : "gcb:") + std::to_string(opId) + ":" + std::to_string(lane);
+}
+
+template <typename ComputeOpTy, typename BatchOpTy>
+bool isTopLevelRelevantCompute(Operation& op) {
+  return isa<ComputeOpTy, BatchOpTy>(&op);
+}
+
+template <typename ComputeOpTy, typename BatchOpTy>
+FailureOr<TopLevelOpInfo> buildTopLevelOpInfo(Operation& op, bool isScheduled, size_t opId) {
+  TopLevelOpInfo info;
+  info.op = &op;
+  info.opId = opId;
+  info.isScheduled = isScheduled;
+
+  if constexpr (std::is_same_v<ComputeOpTy, SpatScheduledCompute>) {
+    if (auto compute = dyn_cast<ComputeOpTy>(&op)) {
+      auto coreId = getOptionalScheduledCoreId(compute, "spatial dataflow export core id");
+      if (failed(coreId))
+        return failure();
+      if (*coreId)
+        info.scalarCore = **coreId;
+    }
+  }
+
+  return info;
+}
+
+template <typename BatchOpTy>
+FailureOr<SmallVector<int32_t, 8>> getBatchLaneCoreIds(BatchOpTy batch) {
+  if constexpr (std::is_same_v<BatchOpTy, SpatScheduledComputeBatch>) {
+    auto coreIds = getOptionalScheduledBatchCoreIds(batch, "spatial dataflow export core ids");
+    if (failed(coreIds))
+      return failure();
+    if (!*coreIds)
+      return SmallVector<int32_t, 8> {};
+    return SmallVector<int32_t, 8>((**coreIds).begin(), (**coreIds).end());
+  }
+  return SmallVector<int32_t, 8> {};
+}
+
+std::string getExpandedNodeId(const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                              Operation* op,
+                              uint32_t lane) {
+  auto it = expandedNodes.find({op, lane});
+  if (it == expandedNodes.end())
+    return "";
+  return it->second.id;
+}
+
+void addScalarNodeRow(std::fstream& nodesFile,
+                      DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                      const TopLevelOpInfo& info,
+                      AsmState* asmState = nullptr) {
+  std::string id = getScalarId(info.isScheduled, info.opId);
+  SmallVector<std::string, 5> row {id, std::to_string(info.opId), "", maybeNumber<int32_t>(info.scalarCore)};
+  if (asmState)
+    row.push_back(stringifyResultSsaNames(info.op, asmState));
+  writeCsvRow(nodesFile, row);
+  expandedNodes[{info.op, 0}] = {id, info.scalarCore, std::nullopt};
+}
+
+template <typename BatchOpTy>
+void addBatchNodeRows(std::fstream& nodesFile,
+                      DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                      const TopLevelOpInfo& info,
+                      BatchOpTy batch,
+                      ArrayRef<std::optional<int32_t>> laneCoreIds,
+                      AsmState* asmState = nullptr) {
+  for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
+    std::string id = getBatchLaneId(info.isScheduled, info.opId, lane);
+    SmallVector<std::string, 5> row {
+      id, std::to_string(info.opId), std::to_string(lane), maybeNumber<int32_t>(laneCoreIds[lane])};
+    if (asmState)
+      row.push_back(stringifyResultSsaNames(info.op, asmState));
+    writeCsvRow(nodesFile, row);
+    expandedNodes[{info.op, lane}] = {id, laneCoreIds[lane], lane};
+  }
+}
+
+std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t lane);
+
+std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t lane) {
+  if (value == laneArg)
+    return static_cast<int64_t>(lane);
+
+  if (std::optional<int64_t> constant = matchConstantIndexValue(value))
+    return *constant;
+
+  if (auto constant = value.getDefiningOp<arith::ConstantOp>()) {
+    if (auto intAttr = dyn_cast<IntegerAttr>(constant.getValue()))
+      return intAttr.getInt();
+  }
+
+  if (auto extract = value.getDefiningOp<tensor::ExtractOp>()) {
+    auto constant = extract.getTensor().getDefiningOp<arith::ConstantOp>();
+    auto elements = constant ? dyn_cast<ElementsAttr>(constant.getValue()) : nullptr;
+    auto shapedType = elements ? dyn_cast<ShapedType>(elements.getType()) : nullptr;
+    if (!elements || !shapedType || shapedType.getRank() != 1 || extract.getIndices().size() != 1)
+      return std::nullopt;
+
+    std::optional<int64_t> index = evaluateIndexLike(extract.getIndices().front(), laneArg, lane);
+    if (!index || *index < 0 || *index >= static_cast<int64_t>(elements.getNumElements()))
+      return std::nullopt;
+
+    if (auto denseInts = dyn_cast<DenseIntElementsAttr>(elements))
+      return (*(denseInts.value_begin<APInt>() + *index)).getSExtValue();
+    return std::nullopt;
+  }
+
+  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))
+                                                            return *resolved;
+                                                          return failure();
+                                                        });
+        succeeded(folded)) {
+      return *folded;
+    }
+
+  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};
+
+  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 {};
+
+  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());
+  }
+  return values;
+}
+
+template <typename BatchOpTy>
+std::optional<Value> getBatchLaneInput(BatchOpTy batch, uint32_t lane, unsigned inputIndex) {
+  if (batch.getNumResults() != 0)
+    return batch.getInputs()[inputIndex];
+
+  size_t laneCount = static_cast<size_t>(batch.getLaneCount());
+  if (laneCount == 0 || batch.getInputs().size() % laneCount != 0)
+    return std::nullopt;
+
+  size_t inputsPerLane = batch.getInputs().size() / laneCount;
+  size_t flatIndex = static_cast<size_t>(lane) * inputsPerLane + inputIndex;
+  if (flatIndex >= batch.getInputs().size())
+    return std::nullopt;
+  return batch.getInputs()[flatIndex];
+}
+
+template <typename BatchOpTy>
+unsigned getBatchLaneInputCount(BatchOpTy batch) {
+  if (batch.getNumResults() != 0)
+    return batch.getInputs().size();
+
+  size_t laneCount = static_cast<size_t>(batch.getLaneCount());
+  if (laneCount == 0 || batch.getInputs().size() % laneCount != 0)
+    return 0;
+  return static_cast<unsigned>(batch.getInputs().size() / laneCount);
+}
+
+template <typename ComputeOpTy, typename BatchOpTy>
+std::optional<ResolvedProducer> resolveProducerForValue(Value value, std::optional<uint32_t> consumerLane) {
+  Operation* op = value.getDefiningOp();
+  if (!op)
+    return std::nullopt;
+
+  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
+    Value source = extract.getSource();
+    Operation* sourceOp = source.getDefiningOp();
+    auto sourceBatch = dyn_cast_or_null<BatchOpTy>(sourceOp);
+    if (sourceBatch && sourceBatch.getNumResults() != 0) {
+      auto staticOffsets = extract.getStaticOffsets();
+      if (!staticOffsets.empty() && staticOffsets.front() != ShapedType::kDynamic) {
+        uint32_t lane = static_cast<uint32_t>(staticOffsets.front());
+        return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, lane, lane, 1};
+      }
+      if (consumerLane)
+        return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1};
+      return ResolvedProducer {
+        sourceOp, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast<uint32_t>(sourceBatch.getLaneCount())};
+    }
+    value = source;
+    op = sourceOp;
+    if (!op)
+      return std::nullopt;
+  }
+
+  if (auto compute = dyn_cast<ComputeOpTy>(op))
+    return ResolvedProducer {compute.getOperation(),
+                             static_cast<size_t>(cast<OpResult>(value).getResultNumber()),
+                             LogicalNodeSelector::Scalar,
+                             0,
+                             0,
+                             1};
+
+  if (auto batch = dyn_cast<BatchOpTy>(op)) {
+    if (batch.getNumResults() != 0) {
+      if (consumerLane)
+        return ResolvedProducer {op, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1};
+      return ResolvedProducer {
+        op, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast<uint32_t>(batch.getLaneCount())};
+    }
+
+    uint32_t lane = static_cast<uint32_t>(cast<OpResult>(value).getResultNumber());
+    return ResolvedProducer {op, static_cast<size_t>(lane), LogicalNodeSelector::Lane, lane, lane, 1};
+  }
+
+  return std::nullopt;
+}
+
+SmallVector<EdgeSource, 8>
+resolveProducerSourcesForCsv(const ResolvedProducer& producer,
+                             const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes) {
+  SmallVector<EdgeSource, 8> sources;
+
+  if (producer.selector == LogicalNodeSelector::Scalar) {
+    std::string id = getExpandedNodeId(expandedNodes, producer.op, 0);
+    if (!id.empty())
+      sources.push_back({id, std::nullopt});
+    return sources;
+  }
+
+  if (producer.selector == LogicalNodeSelector::Lane) {
+    std::string id = getExpandedNodeId(expandedNodes, producer.op, producer.lane);
+    if (!id.empty())
+      sources.push_back({id, producer.lane});
+    return sources;
+  }
+
+  for (uint32_t lane = producer.laneStart; lane < producer.laneStart + producer.laneCount; ++lane) {
+    std::string id = getExpandedNodeId(expandedNodes, producer.op, lane);
+    if (!id.empty())
+      sources.push_back({id, lane});
+  }
+  return sources;
+}
+
+FailureOr<SmallVector<int64_t>> getIntegerValues(Operation* op, StringRef name) {
+  Attribute attr = op->getAttr(name);
+  if (auto array = dyn_cast_or_null<DenseI64ArrayAttr>(attr))
+    return SmallVector<int64_t>(array.asArrayRef());
+  if (auto elements = dyn_cast_or_null<DenseIntElementsAttr>(attr))
+    return SmallVector<int64_t>(elements.getValues<int64_t>());
+  return op->emitOpError() << "expected " << name << " integer array for Spatial dataflow report";
+}
+
+FailureOr<ScheduledNodeByGraphLane>
+buildScheduledNodesByGraphLane(const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
+                               const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                               const DenseMap<int64_t, Operation*>& graphOpsById) {
+  ScheduledNodeByGraphLane nodesByGraphLane;
+  for (const auto& entry : topLevelInfo) {
+    Operation* scheduledOp = entry.first;
+    auto sourceIds = getIntegerValues(scheduledOp, "scheduled.step_source_ids");
+    auto sourceStarts = getIntegerValues(scheduledOp, "scheduled.source_lane_starts");
+    auto sourceCounts = getIntegerValues(scheduledOp, "scheduled.source_lane_counts");
+    if (failed(sourceIds) || failed(sourceStarts) || failed(sourceCounts))
+      return failure();
+
+    uint32_t scheduledLaneCount = 1;
+    if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduledOp))
+      scheduledLaneCount = static_cast<uint32_t>(batch.getLaneCount());
+    size_t expectedEntries = sourceIds->size() * scheduledLaneCount;
+    if (sourceStarts->size() != expectedEntries || sourceCounts->size() != expectedEntries)
+      return scheduledOp->emitOpError("inconsistent scheduling provenance arrays for Spatial dataflow report");
+
+    for (auto [step, graphId] : llvm::enumerate(*sourceIds)) {
+      auto graphIt = graphOpsById.find(graphId);
+      if (graphIt == graphOpsById.end())
+        return scheduledOp->emitOpError() << "references unknown scheduled graph id " << graphId;
+      bool graphIsBatch = isa<SpatGraphComputeBatch>(graphIt->second);
+      for (uint32_t scheduledLane = 0; scheduledLane < scheduledLaneCount; ++scheduledLane) {
+        auto nodeIt = expandedNodes.find({scheduledOp, scheduledLane});
+        if (nodeIt == expandedNodes.end())
+          continue;
+        size_t index = step * scheduledLaneCount + scheduledLane;
+        int64_t start = graphIsBatch ? (*sourceStarts)[index] : 0;
+        int64_t count = graphIsBatch ? (*sourceCounts)[index] : 1;
+        if (start < 0 || count < 0)
+          return scheduledOp->emitOpError("negative scheduling provenance range for Spatial dataflow report");
+        for (int64_t lane = start; lane < start + count; ++lane)
+          nodesByGraphLane[{graphId, static_cast<uint32_t>(lane)}] = nodeIt->second;
+      }
+    }
+  }
+  return nodesByGraphLane;
+}
+
+SmallVector<ExpandedNodeInfo, 8> resolveScheduledProducerNodes(const ResolvedProducer& producer,
+                                                               const ScheduledNodeByGraphLane& nodesByGraphLane) {
+  SmallVector<ExpandedNodeInfo, 8> nodes;
+  auto graphId = producer.op->getAttrOfType<IntegerAttr>("scheduled.graph_id");
+  if (!graphId)
+    return nodes;
+
+  uint32_t laneStart = producer.selector == LogicalNodeSelector::Scalar ? 0 : producer.laneStart;
+  uint32_t laneCount = producer.selector == LogicalNodeSelector::RangeRepresentative ? producer.laneCount : 1;
+  for (uint32_t lane = laneStart; lane < laneStart + laneCount; ++lane)
+    if (auto it = nodesByGraphLane.find({graphId.getInt(), lane}); it != nodesByGraphLane.end())
+      nodes.push_back(it->second);
+  return nodes;
+}
+
+LogicalResult
+emitScheduledPlanningEdges(std::fstream& edgesFile,
+                           func::FuncOp func,
+                           const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
+                           const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                           StringRef stage) {
+  DenseMap<int64_t, Operation*> graphOpsById;
+  for (Operation& op : func.getBody().front())
+    if (auto graphId = op.getAttrOfType<IntegerAttr>("scheduled.graph_id"))
+      graphOpsById[graphId.getInt()] = &op;
+
+  auto nodesByGraphLane = buildScheduledNodesByGraphLane(topLevelInfo, expandedNodes, graphOpsById);
+  if (failed(nodesByGraphLane))
+    return failure();
+
+  auto emitMappedEdge =
+    [&](const ResolvedProducer& producer, int64_t targetGraphId, uint32_t targetGraphLane, Type type) {
+      auto targetIt = nodesByGraphLane->find({targetGraphId, targetGraphLane});
+      if (targetIt == nodesByGraphLane->end())
+        return;
+      for (const ExpandedNodeInfo& source : resolveScheduledProducerNodes(producer, *nodesByGraphLane)) {
+        if (source.id == targetIt->second.id)
+          continue;
+        emitEdgeRow(edgesFile,
+                    source.id,
+                    targetIt->second.id,
+                    getTypeSizeBytes(type),
+                    type,
+                    stage,
+                    source.lane,
+                    targetIt->second.lane,
+                    std::nullopt);
+      }
+    };
+
+  for (Operation& op : func.getBody().front()) {
+    auto graphId = op.getAttrOfType<IntegerAttr>("scheduled.graph_id");
+    if (!graphId)
+      continue;
+    if (auto compute = dyn_cast<SpatGraphCompute>(&op)) {
+      for (Value input : compute.getInputs())
+        if (auto producer = resolveProducerForValue<SpatGraphCompute, SpatGraphComputeBatch>(input, std::nullopt))
+          emitMappedEdge(*producer, graphId.getInt(), 0, input.getType());
+      continue;
+    }
+    auto batch = dyn_cast<SpatGraphComputeBatch>(&op);
+    if (!batch)
+      continue;
+    unsigned inputCount = getBatchLaneInputCount(batch);
+    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane)
+      for (unsigned inputIndex = 0; inputIndex < inputCount; ++inputIndex)
+        if (std::optional<Value> input = getBatchLaneInput(batch, lane, inputIndex))
+          if (auto producer = resolveProducerForValue<SpatGraphCompute, SpatGraphComputeBatch>(*input, lane))
+            emitMappedEdge(*producer, graphId.getInt(), lane, input->getType());
+  }
+  return success();
+}
+
+void emitEdgeRow(std::fstream& edgesFile,
+                 StringRef sourceId,
+                 StringRef targetId,
+                 std::optional<uint64_t> byteSize,
+                 Type propagatedType,
+                 StringRef stage,
+                 std::optional<uint32_t> sourceLane,
+                 std::optional<uint32_t> targetLane,
+                 std::optional<int64_t> channelId) {
+  writeCsvRow(edgesFile,
+              {sourceId.str(),
+               targetId.str(),
+               maybeNumber<uint64_t>(byteSize),
+               stringifyType(propagatedType),
+               stage.str(),
+               maybeNumber<uint32_t>(sourceLane),
+               maybeNumber<uint32_t>(targetLane),
+               maybeNumber<int64_t>(channelId)});
+}
+
+template <typename ComputeOpTy, typename BatchOpTy>
+LogicalResult emitDataEdges(std::fstream& edgesFile,
+                            const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
+                            const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                            StringRef stage) {
+  for (const auto& entry : topLevelInfo) {
+    Operation* op = entry.first;
+    const TopLevelOpInfo& info = entry.second;
+
+    if (auto compute = dyn_cast<ComputeOpTy>(op)) {
+      for (Value input : compute.getInputs()) {
+        if (isa_and_nonnull<SpatChannelReceiveOp>(input.getDefiningOp()))
+          continue;
+
+        auto producer = resolveProducerForValue<ComputeOpTy, BatchOpTy>(input, std::nullopt);
+        if (!producer)
+          continue;
+
+        SmallVector<EdgeSource, 8> sources = resolveProducerSourcesForCsv(*producer, expandedNodes);
+        std::optional<uint64_t> byteSize = getTypeSizeBytes(input.getType());
+        std::string targetId = getScalarId(info.isScheduled, info.opId);
+        for (const EdgeSource& source : sources)
+          emitEdgeRow(edgesFile,
+                      source.id,
+                      targetId,
+                      byteSize,
+                      input.getType(),
+                      stage,
+                      source.sourceLane,
+                      std::nullopt,
+                      std::nullopt);
+      }
+      continue;
+    }
+
+    auto batch = dyn_cast<BatchOpTy>(op);
+    if (!batch)
+      continue;
+
+    unsigned inputCount = getBatchLaneInputCount(batch);
+    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
+      std::string targetId = getBatchLaneId(info.isScheduled, info.opId, lane);
+      for (unsigned inputIndex = 0; inputIndex < inputCount; ++inputIndex) {
+        std::optional<Value> input = getBatchLaneInput(batch, lane, inputIndex);
+        if (!input || isa_and_nonnull<SpatChannelReceiveOp>((*input).getDefiningOp()))
+          continue;
+
+        auto producer = resolveProducerForValue<ComputeOpTy, BatchOpTy>(*input, lane);
+        if (!producer)
+          continue;
+
+        SmallVector<EdgeSource, 8> sources = resolveProducerSourcesForCsv(*producer, expandedNodes);
+        std::optional<uint64_t> byteSize = getTypeSizeBytes((*input).getType());
+        for (const EdgeSource& source : sources)
+          emitEdgeRow(
+            edgesFile, source.id, targetId, byteSize, (*input).getType(), stage, source.sourceLane, lane, std::nullopt);
+      }
+    }
+  }
+
+  return success();
+}
+
+template <typename BatchOpTy>
+void 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;
+
+  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});
+    });
+  }
+}
+
+void 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;
+}
+
+template <typename ComputeOpTy, typename BatchOpTy, typename ResolveChannelSourcesFn>
+LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile,
+                                       const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
+                                       ResolveChannelSourcesFn&& resolveChannelSources,
+                                       StringRef stage) {
+  for (const auto& entry : topLevelInfo) {
+    Operation* op = entry.first;
+    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);
+        std::string targetId = getScalarId(info.isScheduled, info.opId);
+        std::optional<uint64_t> byteSize = getTypeSizeBytes(receive.getType());
+        for (const ChannelSendRecord& source : sources)
+          emitEdgeRow(edgesFile,
+                      source.sourceId,
+                      targetId,
+                      byteSize,
+                      receive.getType(),
+                      stage,
+                      source.sourceLane,
+                      std::nullopt,
+                      channelId);
+      });
+      continue;
+    }
+
+    auto batch = dyn_cast<BatchOpTy>(op);
+    if (!batch)
+      continue;
+    auto laneArg = batch.getLaneArgument();
+    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);
+        std::optional<uint64_t> byteSize = getTypeSizeBytes(receive.getType());
+        for (const ChannelSendRecord& source : sources)
+          emitEdgeRow(edgesFile,
+                      source.sourceId,
+                      targetId,
+                      byteSize,
+                      receive.getType(),
+                      stage,
+                      source.sourceLane,
+                      lane,
+                      channelId);
+      });
+    }
+  }
+
+  return success();
+}
+
+LogicalResult exportGraph(func::FuncOp func, StringRef reportName) {
+  std::fstream nodesFile = openDialectDumpFileWithExtension((reportName + ".nodes").str(), "/reports", "csv");
+  std::fstream edgesFile = openDialectDumpFileWithExtension((reportName + ".edges").str(), "/reports", "csv");
+  if (!nodesFile.is_open() || !edgesFile.is_open())
+    return success();
+
+  writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core", "ssa_name"});
+  writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"});
+
+  Operation* asmRoot = func.getOperation();
+  if (auto moduleOp = func->getParentOfType<ModuleOp>())
+    asmRoot = moduleOp.getOperation();
+  OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  AsmState asmState(asmRoot, flags);
+
+  DenseMap<Operation*, TopLevelOpInfo> topLevelInfo;
+  DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo> expandedNodes;
+
+  size_t opId = 0;
+  for (Operation& op : func.getBody().front()) {
+    if (!isTopLevelRelevantCompute<SpatGraphCompute, SpatGraphComputeBatch>(op))
+      continue;
+    FailureOr<TopLevelOpInfo> info = buildTopLevelOpInfo<SpatGraphCompute, SpatGraphComputeBatch>(op, false, opId++);
+    if (failed(info))
+      return failure();
+    topLevelInfo[&op] = *info;
+
+    if (auto compute = dyn_cast<SpatGraphCompute>(&op)) {
+      addScalarNodeRow(nodesFile, expandedNodes, *info, &asmState);
+      continue;
+    }
+
+    auto batch = cast<SpatGraphComputeBatch>(&op);
+    SmallVector<std::optional<int32_t>, 8> laneCoreIds(batch.getLaneCount());
+    addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState);
+  }
+
+  return emitDataEdges<SpatGraphCompute, SpatGraphComputeBatch>(edgesFile, topLevelInfo, expandedNodes, "spatial1");
+}
+
+LogicalResult exportScheduled(func::FuncOp func, StringRef reportName, StringRef stage) {
+  std::fstream nodesFile = openDialectDumpFileWithExtension((reportName + ".nodes").str(), "/reports", "csv");
+  std::fstream edgesFile = openDialectDumpFileWithExtension((reportName + ".edges").str(), "/reports", "csv");
+  if (!nodesFile.is_open() || !edgesFile.is_open())
+    return success();
+
+  writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core", "ssa_name"});
+  writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"});
+
+  Operation* asmRoot = func.getOperation();
+  if (auto moduleOp = func->getParentOfType<ModuleOp>())
+    asmRoot = moduleOp.getOperation();
+  OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  AsmState asmState(asmRoot, flags);
+
+  DenseMap<Operation*, TopLevelOpInfo> topLevelInfo;
+  DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo> expandedNodes;
+
+  size_t opId = 0;
+  for (Operation& op : func.getBody().front()) {
+    if (!isTopLevelRelevantCompute<SpatScheduledCompute, SpatScheduledComputeBatch>(op))
+      continue;
+    FailureOr<TopLevelOpInfo> info =
+      buildTopLevelOpInfo<SpatScheduledCompute, SpatScheduledComputeBatch>(op, true, opId++);
+    if (failed(info))
+      return failure();
+    topLevelInfo[&op] = *info;
+
+    if (isa<SpatScheduledCompute>(&op)) {
+      addScalarNodeRow(nodesFile, expandedNodes, *info, &asmState);
+      continue;
+    }
+
+    auto batch = cast<SpatScheduledComputeBatch>(&op);
+    auto coreIds = getBatchLaneCoreIds(batch);
+    if (failed(coreIds))
+      return failure();
+    SmallVector<std::optional<int32_t>, 8> laneCoreIds(batch.getLaneCount());
+    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane)
+      if (lane < coreIds->size())
+        laneCoreIds[lane] = (*coreIds)[lane];
+    addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState);
+  }
+
+  if (stage == "spatial2")
+    return emitScheduledPlanningEdges(edgesFile, func, topLevelInfo, expandedNodes, stage);
+  if (failed(
+        emitDataEdges<SpatScheduledCompute, SpatScheduledComputeBatch>(edgesFile, topLevelInfo, expandedNodes, stage)))
+    return failure();
+
+  DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>> sendsByChannelId;
+  for (const auto& entry : topLevelInfo) {
+    Operation* op = entry.first;
+    if (auto compute = dyn_cast<SpatScheduledCompute>(op))
+      collectChannelSends(sendsByChannelId, expandedNodes, compute);
+    else if (auto batch = dyn_cast<SpatScheduledComputeBatch>(op))
+      collectChannelSends(sendsByChannelId, expandedNodes, batch);
+  }
+
+  DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>> nodesByCore = buildNodesByCore(expandedNodes);
+  auto resolveChannelSources = [&](SpatChannelReceiveOp receive, uint32_t lane) {
+    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;
+  };
+
+  return emitExplicitChannelEdges<SpatScheduledCompute, SpatScheduledComputeBatch>(
+    edgesFile, topLevelInfo, resolveChannelSources, stage);
+}
+
+} // namespace
+
+SpatialDataflowExportStage getSpatialDataflowExportStage() {
+  switch (pimExportSpatialDataflow.getValue()) {
+  case SpatialDataflowExportNone:     return SpatialDataflowExportStage::None;
+  case SpatialDataflowExportSpatial1: return SpatialDataflowExportStage::Spatial1;
+  case SpatialDataflowExportSpatial2: return SpatialDataflowExportStage::Spatial2;
+  case SpatialDataflowExportSpatial3: return SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportAll:      return SpatialDataflowExportStage::All;
+  }
+  llvm_unreachable("unknown spatial dataflow export mode");
+}
+
+bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage) {
+  switch (mode) {
+  case SpatialDataflowExportStage::None:     return false;
+  case SpatialDataflowExportStage::Spatial1: return stage == SpatialDataflowExportStage::Spatial1;
+  case SpatialDataflowExportStage::Spatial2: return stage == SpatialDataflowExportStage::Spatial2;
+  case SpatialDataflowExportStage::Spatial3: return stage == SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportStage::All:      return stage != SpatialDataflowExportStage::None;
+  }
+  return false;
+}
+
+LogicalResult exportSpatialDataflowCsvGraph(func::FuncOp func, StringRef reportName) {
+  return exportGraph(func, reportName);
+}
+
+LogicalResult exportSpatialDataflowCsvScheduled(func::FuncOp func, StringRef reportName, StringRef stage) {
+  return exportScheduled(func, reportName, stage);
+}
+
+} // namespace spatial
+} // namespace onnx_mlir
@@ -0,0 +1,28 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/StringRef.h"
+
+namespace onnx_mlir {
+namespace spatial {
+
+enum class SpatialDataflowExportStage {
+  None,
+  Spatial1,
+  Spatial2,
+  Spatial3,
+  All,
+};
+
+SpatialDataflowExportStage getSpatialDataflowExportStage();
+
+mlir::LogicalResult exportSpatialDataflowCsvGraph(mlir::func::FuncOp func, llvm::StringRef reportName);
+mlir::LogicalResult
+exportSpatialDataflowCsvScheduled(mlir::func::FuncOp func, llvm::StringRef reportName, llvm::StringRef stage);
+
+bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage);
+
+} // namespace spatial
+} // namespace onnx_mlir