Cose belle

2026-07-14 16:57:58 +02:00
parent d1a29ace3c
commit 51fdb830e5
38 changed files with 5365 additions and 914 deletions
@@ -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