Raptor sync wait

src/PIM/Common/PimCommon.hpp

@@ -32,6 +32,9 @@ inline constexpr llvm::StringLiteral kCoreIdAttrName = "coreId";
 inline constexpr llvm::StringLiteral kCoreIdsAttrName = "coreIds";
 inline constexpr llvm::StringLiteral kLocalMemoryAddressAttrName = "pim.local_memory_address";
 inline constexpr llvm::StringLiteral kLocalMemorySizeAttrName = "pim.local_memory_size";
+inline constexpr llvm::StringLiteral kPipelineHostBufferBytesAttrName = "pim.pipeline_host_buffer_bytes";
+inline constexpr llvm::StringLiteral kPipelineHostBufferName = "pim_pipeline_channels";
+inline constexpr size_t kPimEventRegisterCount = 32;
 inline constexpr std::array<llvm::StringLiteral, 4> kRemovedLocalMemoryPlanAttrNames = {
   "pim.local_memory_slot",
   "pim.local_memory_slot_size",

src/PIM/Compiler/PimBinaryFormat.hpp

@@ -162,8 +162,8 @@ inline constexpr std::array<InstructionJsonFormat, kOpcodeCount> kInstructionJso
   {true,  true,  true,  "",     "",     "",      "len" }, // lmv
   {true,  false, true,  "core", "",     "",      "size"}, // send
   {true,  false, true,  "core", "",     "",      "size"}, // recv
-  {false, false, false, "",     "",     "",      ""    }, // wait
+  {false, false, false, "",     "event_register", "wait_value", ""}, // wait
-  {false, false, false, "",     "",     "",      ""    }, // sync
+  {false, false, false, "core", "event_register", "", ""}, // sync
 }};
 static_assert(kInstructionJsonFormats.size() == kOpcodeCount);
 

src/PIM/Compiler/PimCodeGen.cpp

@@ -692,6 +692,34 @@ void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge
     pim_binary::Opcode::send, addressOf(sendOp.getInput(), knowledge), *targetCoreId, sendOp.getSize());
 }
 
+void PimCodeGen::codeGenWaitOp(
+    pim::PimWaitOp waitOp, const StaticValueKnowledge& knowledge) const {
+  auto eventRegister = indexOf(waitOp.getEventRegister(), knowledge);
+  assert(succeeded(eventRegister)
+         && "pim.wait event register must be statically resolvable during codegen");
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::wait;
+  instruction.generic1 = pim::checkedI32OrCrash(
+    *eventRegister, "wait event register");
+  instruction.generic2 = waitOp.getWaitValue();
+  emitInstruction(instruction);
+}
+
+void PimCodeGen::codeGenSyncOp(
+    pim::PimSyncOp syncOp, const StaticValueKnowledge& knowledge) const {
+  auto targetCoreId = indexOf(syncOp.getTargetCoreId(), knowledge);
+  auto eventRegister = indexOf(syncOp.getEventRegister(), knowledge);
+  assert(succeeded(targetCoreId) && succeeded(eventRegister)
+         && "pim.sync operands must be statically resolvable during codegen");
+  pim_binary::InstructionRecord instruction;
+  instruction.opcode = pim_binary::Opcode::sync;
+  instruction.r2OrImm = pim::checkedI32OrCrash(
+    *targetCoreId, "sync target core id");
+  instruction.generic1 = pim::checkedI32OrCrash(
+    *eventRegister, "sync event register");
+  emitInstruction(instruction);
+}
+
 void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const {
   auto outputType = cast<ShapedType>(concatOp.getOutputBuffer().getType());
   assert(outputType.hasStaticShape() && "concat codegen requires static output shape");
@@ -991,6 +1019,8 @@ static LogicalResult executeCompiledCorePlan(
     case CompiledCoreOpKind::VMV:     coreCodeGen.codeGenVMVOp(cast<pim::PimVMVOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Receive: coreCodeGen.codeGenReceiveOp(cast<pim::PimReceiveOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Send:    coreCodeGen.codeGenSendOp(cast<pim::PimSendOp>(node.op), knowledge); break;
+    case CompiledCoreOpKind::Wait:    coreCodeGen.codeGenWaitOp(cast<pim::PimWaitOp>(node.op), knowledge); break;
+    case CompiledCoreOpKind::Sync:    coreCodeGen.codeGenSyncOp(cast<pim::PimSyncOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Concat:  coreCodeGen.codeGenConcatOp(cast<pim::PimConcatOp>(node.op), knowledge); break;
     case CompiledCoreOpKind::Vmm:
       if (auto weightSlot = resolveWeightSlot(cast<pim::PimVMMOp>(node.op), knowledge); succeeded(weightSlot))

src/PIM/Compiler/PimCodeGen.hpp

@@ -217,6 +217,8 @@ public:
 
   void codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const;
   void codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenWaitOp(pim::PimWaitOp waitOp, const StaticValueKnowledge& knowledge) const;
+  void codeGenSyncOp(pim::PimSyncOp syncOp, const StaticValueKnowledge& knowledge) const;
   void codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const;
 
   template <typename MVMTy>

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -2,6 +2,8 @@
 
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
+#include <limits>
+
 #define DEBUG_TYPE "PimCompilerOptions"
 
 namespace onnx_mlir {
@@ -110,6 +112,12 @@ llvm::cl::opt<size_t>
 llvm::cl::opt<size_t>
   crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(64));
 
+llvm::cl::opt<size_t> pipelineStages(
+  "pipeline",
+  llvm::cl::desc("Number of throughput pipeline stages (1 preserves latency scheduling)"),
+  llvm::cl::init(1),
+  llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<long> coresCount("core-count",
                                llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."),
                                llvm::cl::init(-1));
@@ -129,4 +137,14 @@ void verifyExplicitPimCoreCount() {
     llvm::report_fatal_error("PIM compilation requires --core-count to be a positive integer");
 }
 
+void verifyPimPipelineStages() {
+  if (pipelineStages.getValue() == 0)
+    llvm::report_fatal_error("PIM compilation requires --pipeline to be positive");
+  if (static_cast<size_t>(coresCount.getValue()) % pipelineStages.getValue() != 0)
+    llvm::report_fatal_error("PIM compilation requires --core-count to be divisible by --pipeline");
+  if (crossbarCountInCore.getValue()
+      > std::numeric_limits<size_t>::max() / pipelineStages.getValue())
+    llvm::report_fatal_error("PIM compilation --crossbar-count * --pipeline overflows");
+}
+
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -62,6 +62,7 @@ extern llvm::cl::opt<bool> pimVerifyBufferizationCopyFreedom;
 
 extern llvm::cl::opt<size_t> crossbarSize;
 extern llvm::cl::opt<size_t> crossbarCountInCore;
+extern llvm::cl::opt<size_t> pipelineStages;
 extern llvm::cl::opt<long> coresCount;
 extern llvm::cl::opt<std::string> pimTargetConfig;
 extern llvm::cl::opt<uint64_t> pimConvIm2colMaxElements;
@@ -69,5 +70,6 @@ extern llvm::cl::opt<uint64_t> pimConvStreamChunkPositions;
 
 bool hasExplicitPimCoreCount();
 void verifyExplicitPimCoreCount();
+void verifyPimPipelineStages();
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -330,6 +330,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
                   EmissionTargetType& emissionTarget,
                   std::string outputNameNoExt) {
   verifyExplicitPimCoreCount();
+  verifyPimPipelineStages();
   spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
   spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget);
 
@@ -354,7 +355,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(createTrivialGraphComputeMergePass(
       schedulingTarget.residentWeightCapacity, exportStage));
     pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
-      schedulingTarget, exportStage));
+      schedulingTarget, exportStage, pipelineStages.getValue()));
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
 

src/PIM/Compiler/PimCoreProgram.cpp

@@ -17,6 +17,8 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   if (isa<pim::PimVMVOp>(op)) return CompiledCoreOpKind::VMV;
   if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
   if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
+  if (isa<pim::PimWaitOp>(op)) return CompiledCoreOpKind::Wait;
+  if (isa<pim::PimSyncOp>(op)) return CompiledCoreOpKind::Sync;
   if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
   if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
   if (isa<pim::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;

src/PIM/Compiler/PimCoreProgram.hpp

@@ -17,6 +17,8 @@ enum class CompiledCoreOpKind : uint8_t {
   VMV,
   Receive,
   Send,
+  Wait,
+  Sync,
   Concat,
   Vmm,
   VVAdd,

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

@@ -0,0 +1,39 @@
+#include "ContractionMaterialization.hpp"
+
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "MatrixProductLowering.hpp"
+
+namespace onnx_mlir {
+
+mlir::Value materializePaddedContractionInput(
+    mlir::Value input,
+    mlir::RankedTensorType paddedType,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc) {
+  return createPaddedInputCompute(input, paddedType, rewriter, loc);
+}
+
+mlir::FailureOr<mlir::Value> materializeTransposedContractionConstant(
+    mlir::Value input,
+    mlir::RankedTensorType resultType,
+    llvm::ArrayRef<int64_t> permutation,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc) {
+  auto denseAttr = getHostConstDenseElementsAttr(input);
+  auto inputType = denseAttr ? mlir::dyn_cast<mlir::RankedTensorType>(denseAttr.getType()) : nullptr;
+  if (!inputType || !inputType.hasStaticShape() || !resultType || !resultType.hasStaticShape()
+      || inputType.getRank() != resultType.getRank())
+    return mlir::failure();
+
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (mlir::failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return mlir::failure();
+
+  return getOrCreateConstant(rewriter,
+                             rewriter.getInsertionBlock()->getParentOp(),
+                             *transposedAttr,
+                             resultType);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,23 @@
+#pragma once
+
+#include "llvm/ADT/ArrayRef.h"
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+mlir::Value materializePaddedContractionInput(
+    mlir::Value input,
+    mlir::RankedTensorType paddedType,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> materializeTransposedContractionConstant(
+    mlir::Value input,
+    mlir::RankedTensorType resultType,
+    llvm::ArrayRef<int64_t> permutation,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -0,0 +1,902 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "mlir/Transforms/Passes.h"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<RowStripPhysicalValue> getRowStripValue(Value value) {
+  return getRowStripPhysicalValue(value);
+}
+
+static FailureOr<Value> publishRowStripValue(Operation* planOp,
+                                             Value storage,
+                                             PatternRewriter& rewriter) {
+  auto logicalType = dyn_cast<RankedTensorType>(planOp->getResult(0).getType());
+  if (!logicalType)
+    return planOp->emitOpError("requires ranked logical output type"), failure();
+  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
+  if (failed(value))
+    return planOp->emitOpError("lowering produced invalid row-strip physical storage"), failure();
+  FailureOr<Value> blueprint = createRowStripStorageBlueprint(
+      storage, logicalType, rewriter, planOp->getLoc());
+  if (failed(blueprint))
+    return planOp->emitOpError("failed to create row-strip storage Blueprint"), failure();
+  rewriter.replaceOp(planOp, *blueprint);
+  return *blueprint;
+}
+
+static bool isRowStripSelected(Operation* op) {
+  auto selected = spatial::getSelectedPhysicalLayout(op);
+  return selected && *selected == spatial::PhysicalLayout::NHWCRowStrip;
+}
+
+static bool isDenseSelected(Operation* op) {
+  auto selected = spatial::getSelectedPhysicalLayout(op);
+  return selected && *selected == spatial::PhysicalLayout::DenseNCHW;
+}
+
+static spatial::PhysicalLayout getKnownPhysicalLayout(Value value) {
+  if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
+    return materialize.getTargetPhysicalLayout();
+  if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
+    return blueprint.getPhysicalLayout();
+  if (Operation* producer = value.getDefiningOp()) {
+    if (auto selected = spatial::getSelectedPhysicalLayout(producer))
+      return *selected;
+  }
+  return spatial::PhysicalLayout::DenseNCHW;
+}
+
+static LogicalResult verifySelectedLayouts(
+    func::FuncOp funcOp, const spatial::SpatialTargetInfo& target) {
+  LogicalResult result = success();
+  funcOp.walk([&](Operation* op) {
+    auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
+    if (!capability)
+      return;
+    auto selected = spatial::getSelectedPhysicalLayout(op);
+    if (!selected) {
+      op->emitOpError("requires a selected physical layout from SpatialLayoutPlanning");
+      result = failure();
+      return;
+    }
+    if (*selected != spatial::PhysicalLayout::DenseNCHW
+        && *selected != spatial::PhysicalLayout::NHWCRowStrip) {
+      op->emitOpError("has an unsupported selected physical layout");
+      result = failure();
+      return;
+    }
+    SmallVector<spatial::PhysicalLayout> operandLayouts;
+    operandLayouts.reserve(op->getNumOperands());
+    for (Value operand : op->getOperands())
+      operandLayouts.push_back(getKnownPhysicalLayout(operand));
+    auto alternatives = capability.getLayoutAlternatives(target, operandLayouts);
+    if (llvm::none_of(alternatives, [&](const spatial::LayoutAlternative& alternative) {
+          return alternative.resultLayout == *selected
+              && alternative.operandLayouts == operandLayouts;
+        })) {
+      op->emitOpError("selected physical layout is not lowerable for its explicit operand layouts");
+      result = failure();
+    }
+  });
+  return result;
+}
+
+static FailureOr<Value>
+lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripRelu(input, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value>
+lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripSilu(input, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
+                                             spatial::SpatBiasAddPlanOp planOp,
+                                             PatternRewriter& rewriter) {
+  return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
+                                         const RowStripPhysicalValue& rhs,
+                                         spatial::SpatAddPlanOp planOp,
+                                         PatternRewriter& rewriter) {
+  return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
+                                            spatial::SpatConcatPlanOp planOp,
+                                            PatternRewriter& rewriter) {
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!outputType)
+    return failure();
+  return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value>
+materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
+  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
+    return failure();
+  return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc);
+}
+
+static FailureOr<Value> materializeDenseToRowStrip(
+    Value input, RankedTensorType logicalType, Location loc, PatternRewriter& rewriter) {
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4
+      || logicalType.getDimSize(0) != 1)
+    return failure();
+  auto nhwcType = RankedTensorType::get(
+    {1, logicalType.getDimSize(2), logicalType.getDimSize(3), logicalType.getDimSize(1)},
+    logicalType.getElementType(), logicalType.getEncoding());
+  auto rowsType = RankedTensorType::get(
+    {logicalType.getDimSize(2) * logicalType.getDimSize(3), logicalType.getDimSize(1)},
+    logicalType.getElementType(), logicalType.getEncoding());
+  auto rowsCompute = createSpatCompute<1>(
+    rewriter, loc, rowsType, {}, input, [&](Value denseInput) {
+      Value nhwc = createLinalgTranspose(
+        denseInput, nhwcType, {0, 2, 3, 1}, rewriter, loc);
+      Value rows = tensor::CollapseShapeOp::create(
+        rewriter, loc, rowsType, nhwc,
+        SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
+      spatial::SpatYieldOp::create(rewriter, loc, rows);
+    });
+  Value rows = rowsCompute->getResult(0);
+  FailureOr<Value> storage = createRowStripStorageFromRows(rows, logicalType, rewriter, loc);
+  if (failed(storage))
+    return failure();
+  return createRowStripStorageBlueprint(*storage, logicalType, rewriter, loc);
+}
+
+static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
+                                               PatternRewriter& rewriter, Location loc) {
+  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
+  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
+      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
+      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
+    return failure();
+  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
+  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
+  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
+      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
+      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
+    return failure();
+  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
+    if (outputDim > inputDim)
+      return failure();
+
+  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
+      if (failed(fragment))
+        return failure();
+      MixedSliceGeometry biasSlice;
+      for (int64_t dim : inputFragmentType->getShape()) {
+        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
+        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
+      }
+      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
+      if (!biasFragment)
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
+      MixedSliceGeometry outputSlice;
+      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
+      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
+      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
+      for (int64_t dim : outputFragmentType->getShape())
+        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
+      if (!output)
+        return failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
+struct LowerDenseReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
+    if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
+      return failure();
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
+        auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
+    if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
+      return failure();
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
+        Value sigmoid = spatial::SpatSigmoidOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
+        Value silu = spatial::SpatVMulOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
+  explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(planOp, std::nullopt, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense nearest Resize plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetInfo& target;
+};
+
+struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+    if (!resultType)
+      return planOp.emitOpError("requires ranked output type");
+
+    FailureOr<Value> denseBias = materializeDenseBiasAddTensor(
+      planOp.getBias(), resultType, rewriter, planOp.getLoc());
+    if (failed(denseBias))
+      return planOp.emitOpError("failed to materialize dense Conv-style bias");
+    if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+      FailureOr<Value> lowered = lowerDenseBatchBiasAdd(
+        planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+      if (succeeded(lowered)) {
+        rewriter.replaceOp(planOp, *lowered);
+        return success();
+      }
+    }
+    auto computeOp = createSpatCompute<2>(
+      rewriter,
+      planOp.getLoc(),
+      planOp.getOutput().getType(),
+      {},
+      ValueRange {planOp.getInput(), *denseBias},
+      [&](Value x, Value y) {
+        auto added = spatial::SpatVAddOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto compute = createSpatCompute<2>(
+      rewriter,
+      planOp.getLoc(),
+      planOp.getOutput().getType(),
+      {},
+      ValueRange {planOp.getLhs(), planOp.getRhs()},
+      [&](Value lhsValue, Value rhsValue) {
+        Value added = spatial::SpatVAddOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added);
+      });
+    rewriter.replaceOp(planOp, compute.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto compute = createSpatCompute(
+      rewriter,
+      planOp.getLoc(),
+      TypeRange {planOp.getOutput().getType()},
+      {},
+      planOp.getInputs(),
+      [&](ValueRange values) {
+        Value concatenated = spatial::SpatConcatOp::create(
+          rewriter,
+          planOp.getLoc(),
+          planOp.getOutput().getType(),
+          rewriter.getI64IntegerAttr(planOp.getAxis()),
+          values);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated);
+      });
+    rewriter.replaceOp(planOp, compute.getResults());
+    return success();
+  }
+};
+
+static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
+                                  PatternRewriter& rewriter) {
+  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(planOp.getLhs());
+  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(planOp.getRhs());
+  if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) {
+    if (getKnownPhysicalLayout(planOp.getLhs()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation()) && failed(rhs)) {
+    if (getKnownPhysicalLayout(planOp.getRhs()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation())) {
+    rewriter.setInsertionPoint(planOp);
+    FailureOr<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+  return planOp.emitOpError("dense Add plan was not lowered by the selected-plan patterns");
+}
+
+static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
+                                     PatternRewriter& rewriter) {
+  SmallVector<RowStripPhysicalValue> inputs;
+  for (Value input : planOp.getInputs()) {
+    FailureOr<RowStripPhysicalValue> physical = getRowStripValue(input);
+    if (failed(physical)) {
+      inputs.clear();
+      break;
+    }
+    inputs.push_back(*physical);
+  }
+  if (isRowStripSelected(planOp.getOperation()) && inputs.size() != planOp.getInputs().size()) {
+    if (llvm::any_of(planOp.getInputs(), [](Value input) {
+          return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip;
+        }))
+      return failure();
+    return planOp.emitOpError("selected row-strip Concat plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation())) {
+    rewriter.setInsertionPoint(planOp);
+    FailureOr<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+  return planOp.emitOpError("dense Concat plan was not lowered by the selected-plan patterns");
+}
+
+struct LowerSelectedConvPlan final : OpRewritePattern<spatial::SpatConv2DPlanOp> {
+  explicit LowerSelectedConvPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<spatial::SpatConv2DPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatConv2DPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (isDenseSelected(planOp.getOperation())) {
+      FailureOr<Value> lowered = lowerSelectedConv2DPlan(
+        planOp, std::nullopt, /*emitRowStripLayout=*/false, target, rewriter);
+      if (failed(lowered))
+        return planOp.emitOpError("failed to lower selected dense Spatial Conv plan");
+      rewriter.replaceOp(planOp, *lowered);
+      return success();
+    }
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+
+    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(planOp.getInput());
+    if (failed(rowStripInput)
+        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    std::optional<Value> physicalInput;
+    if (succeeded(rowStripInput))
+      physicalInput = rowStripInput->storage;
+    FailureOr<Value> lowered = lowerSelectedConv2DPlan(
+      planOp, physicalInput, /*emitRowStripLayout=*/true, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetInfo& target;
+};
+
+struct LowerRowStripReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial ReLU plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
+  explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip Resize plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(planOp, input->storage, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Resize plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetInfo& target;
+};
+
+struct LowerDenseMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
+  explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerDenseMaxPool2DPlan(planOp, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense Spatial MaxPool plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetInfo& target;
+};
+
+struct LowerRowStripMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
+  explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    if (failed(input)
+        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    std::optional<Value> physicalInput;
+    if (succeeded(input))
+      physicalInput = input->storage;
+    FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(planOp, physicalInput, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetInfo& target;
+};
+
+struct LowerRowStripGlobalAveragePoolPlan
+    final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
+  explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    if (failed(input)
+        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    std::optional<Value> physicalInput;
+    if (succeeded(input))
+      physicalInput = input->storage;
+    FailureOr<Value> lowered = lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetInfo& target;
+};
+
+struct LowerDenseGlobalAveragePoolPlan
+    final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
+  explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx,
+                                           const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerDenseGlobalAveragePoolPlan(planOp, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense Spatial global AveragePool plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetInfo& target;
+};
+
+struct LowerRowStripBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip bias_add plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    return lowerAddPlan(planOp, rewriter);
+  }
+};
+
+struct LowerRowStripConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
+                                PatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    return lowerConcatPlan(planOp, rewriter);
+  }
+};
+
+struct LowerMaterializeLayout final
+    : OpRewritePattern<spatial::SpatMaterializeLayoutOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp,
+                                PatternRewriter& rewriter) const override {
+    auto source = materializeOp.getSourcePhysicalLayout();
+    auto target = materializeOp.getTargetPhysicalLayout();
+    if (source == spatial::PhysicalLayout::DenseNCHW
+        && target == spatial::PhysicalLayout::DenseNCHW) {
+      rewriter.replaceOp(materializeOp, materializeOp.getInput());
+      return success();
+    }
+    if (source == spatial::PhysicalLayout::DenseNCHW
+        && target == spatial::PhysicalLayout::NHWCRowStrip) {
+      auto logicalType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
+      if (!logicalType)
+        return materializeOp.emitOpError("requires a ranked dense input"), failure();
+      FailureOr<Value> rowStrip = materializeDenseToRowStrip(
+        materializeOp.getInput(), logicalType, materializeOp.getLoc(), rewriter);
+      if (failed(rowStrip))
+        return materializeOp.emitOpError(
+          "failed to materialize dense NCHW storage to row-strip layout"), failure();
+      rewriter.replaceOp(materializeOp, *rowStrip);
+      return success();
+    }
+    if (source != spatial::PhysicalLayout::NHWCRowStrip
+        || target != spatial::PhysicalLayout::DenseNCHW)
+      return materializeOp.emitOpError(
+          "unsupported Spatial layout materialization direction"), failure();
+    auto inputType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
+    if (!inputType)
+      return materializeOp.emitOpError("requires a ranked row-strip input"), failure();
+    FailureOr<RowStripPhysicalValue> rowStripValue =
+        getRowStripValue(materializeOp.getInput());
+    if (failed(rowStripValue))
+      return materializeOp.emitOpError(
+          "requires an explicitly defining row-strip physical value"), failure();
+    FailureOr<Value> dense = materializeRowStripToDense(
+        *rowStripValue, materializeOp.getLoc(), rewriter);
+    if (failed(dense))
+      return materializeOp.emitOpError(
+          "failed to materialize row-strip storage to dense NCHW"), failure();
+    rewriter.replaceOp(materializeOp, *dense);
+    return success();
+  }
+};
+
+struct LowerRowStripFlatten final
+    : OpRewritePattern<spatial::SpatGraphCompute> {
+  explicit LowerRowStripFlatten(MLIRContext* context,
+                                const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<spatial::SpatGraphCompute>(context), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatGraphCompute flattenOp,
+                                PatternRewriter& rewriter) const override {
+    if (flattenOp.getInputs().size() != 1)
+      return failure();
+    FailureOr<RowStripPhysicalValue> input =
+        getRowStripValue(flattenOp.getInputs().front());
+    if (failed(input) || failed(canLowerFlattenFromRowStrip(flattenOp, target)))
+      return failure();
+    if (failed(lowerFlattenFromRowStrip(*input, flattenOp, target, rewriter)))
+      return flattenOp.emitOpError(
+          "failed to preserve row-strip layout through Flatten"), failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetInfo& target;
+};
+
+struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
+
+  StringRef getArgument() const override { return "lower-spatial-plans"; }
+  StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; }
+
+  LowerSpatialPlansPass() = default;
+  explicit LowerSpatialPlansPass(const spatial::SpatialTargetInfo& target)
+      : target(target), hasTarget(true) {}
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget) {
+      moduleOp.emitError("Spatial plan lowering requires an injected SpatialTargetInfo");
+      signalPassFailure();
+      return;
+    }
+    MLIRContext* ctx = moduleOp.getContext();
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
+      signalPassFailure();
+      return;
+    }
+    func::FuncOp funcOp = *entryFunc;
+    PatternRewriter rewriter(ctx);
+    auto verifyLogicalPhase = [&](StringRef stage) -> bool {
+      if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc)))
+        return true;
+      moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
+      signalPassFailure();
+      return false;
+    };
+
+    if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
+      return;
+    if (failed(verifySelectedLayouts(funcOp, target))) {
+      moduleOp.emitError("selected Spatial layout verification failed");
+      signalPassFailure();
+      return;
+    }
+
+    RewritePatternSet selectedPlanPatterns(ctx);
+    selectedPlanPatterns.add<LowerDenseReluPlan,
+                             LowerRowStripReluPlan,
+                             LowerDenseSiluPlan,
+                             LowerRowStripSiluPlan,
+                             LowerDenseBiasAddPlan,
+                             LowerRowStripBiasAddPlan,
+                             LowerDenseAddPlan,
+                             LowerRowStripAddPlan,
+                             LowerDenseConcatPlan,
+                             LowerRowStripConcatPlan>(ctx);
+    selectedPlanPatterns.add<LowerSelectedConvPlan,
+                             LowerDenseResizePlan,
+                             LowerRowStripResizePlan,
+                             LowerDenseMaxPoolPlan,
+                             LowerRowStripMaxPoolPlan,
+                             LowerDenseGlobalAveragePoolPlan,
+                             LowerRowStripGlobalAveragePoolPlan>(ctx, target);
+    if (failed(applyPatternsGreedily(funcOp, std::move(selectedPlanPatterns)))) {
+      moduleOp.emitError("failed to lower selected Spatial plans");
+      signalPassFailure();
+      return;
+    }
+
+    RewritePatternSet layoutPatterns(ctx);
+    layoutPatterns.add<LowerMaterializeLayout>(ctx);
+    layoutPatterns.add<LowerRowStripFlatten>(ctx, target);
+    ConversionTarget layoutTarget(*ctx);
+    layoutTarget.addLegalDialect<spatial::SpatialDialect,
+                                 tensor::TensorDialect,
+                                 linalg::LinalgDialect,
+                                 affine::AffineDialect,
+                                 arith::ArithDialect,
+                                 scf::SCFDialect,
+                                 func::FuncDialect>();
+    layoutTarget.addIllegalDialect<ONNXDialect>();
+    layoutTarget.addIllegalOp<spatial::SpatMaterializeLayoutOp>();
+    layoutTarget.addDynamicallyLegalOp<spatial::SpatGraphCompute>(
+        [&](spatial::SpatGraphCompute computeOp) {
+          if (computeOp.getInputs().size() != 1)
+            return true;
+          FailureOr<RowStripPhysicalValue> input =
+              getRowStripValue(computeOp.getInputs().front());
+          return failed(input) || failed(canLowerFlattenFromRowStrip(computeOp, target));
+        });
+    FrozenRewritePatternSet frozenLayoutPatterns(std::move(layoutPatterns));
+    if (failed(applyFullConversion(funcOp, layoutTarget,
+                                   frozenLayoutPatterns))) {
+      moduleOp.emitError("failed to lower explicit Spatial layout materialization");
+      signalPassFailure();
+      return;
+    }
+
+    if (!verifyLogicalPhase("after selected-plan conversion"))
+      return;
+    SmallVector<spatial::SpatBlueprintOp> deadPhysicalViews;
+    funcOp.walk([&](spatial::SpatBlueprintOp blueprint) {
+      if (spatial::isPhysicalView(blueprint.getMode()) && blueprint.use_empty())
+        deadPhysicalViews.push_back(blueprint);
+    });
+    for (spatial::SpatBlueprintOp blueprint : deadPhysicalViews)
+      rewriter.eraseOp(blueprint);
+    bool hasIllegalOps = false;
+    moduleOp.walk([&](Operation* op) {
+      if (isa<ONNXEntryPointOp>(op))
+        return;
+      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+        if (spatial::isFragmentAssembly(blueprint.getMode()))
+          return;
+        op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
+        hasIllegalOps = true;
+      }
+      else if (isa<spatial::SpatConv2DPlanOp,
+                   spatial::SpatBiasAddPlanOp,
+                   spatial::SpatAddPlanOp,
+                   spatial::SpatReluPlanOp,
+                   spatial::SpatSiluPlanOp,
+                   spatial::SpatResizeNearestPlanOp,
+                   spatial::SpatMaxPool2DPlanOp,
+                   spatial::SpatGlobalAveragePoolPlanOp,
+                   spatial::SpatMaterializeLayoutOp>(op)
+               || op->getDialect()->getNamespace() == "onnx") {
+        op->emitOpError("operation must not remain after LowerSpatialPlans");
+        hasIllegalOps = true;
+      }
+    });
+
+    PassManager canonicalizationPM(ctx);
+    canonicalizationPM.addPass(createCanonicalizerPass());
+    if (failed(canonicalizationPM.run(moduleOp)))
+      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");
+
+    if (hasIllegalOps) {
+      signalPassFailure();
+    } else {
+      dumpModule(moduleOp, "spatial1_graph");
+      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
+      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
+          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
+        signalPassFailure();
+        return;
+      }
+    }
+
+    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
+      return;
+  }
+
+  spatial::SpatialTargetInfo target;
+  bool hasTarget = false;
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }
+
+std::unique_ptr<Pass> createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target) {
+  return std::make_unique<LowerSpatialPlansPass>(target);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.cpp

@@ -108,7 +108,9 @@ void verifyScheduledInputs(ComputeOpTy compute,
   for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
     size_t currentInputIndex = inputIndex;
     Operation* definingOp = input.getDefiningOp();
-    if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
+    if (allowChannelReceiveInputs
+        && isa_and_nonnull<spatial::SpatChannelReceiveOp,
+                           spatial::SpatHostWaitLoadOp>(definingOp))
       continue;
     if (isScheduledPhase1Value(input))
       continue;
@@ -163,7 +165,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
       });
       continue;
     }
-    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp>(&op)) {
+    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp,
+            spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
         illegalOp->emitOpError() << kPhaseMarker
                                  << " explicit channel communication is not expected before merge materialization";
@@ -182,7 +185,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
 
 void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
   for (Operation& op : funcOp.getOps()) {
-    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp,
+            spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
         illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
       });

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -0,0 +1,64 @@
+#pragma once
+
+#include <optional>
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+
+struct RowStripPhysicalValue;
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
+                        std::optional<mlir::Value> rowStripInput,
+                        bool emitRowStripLayout,
+                        const spatial::SpatialTargetInfo& target,
+                        mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
+                                               const spatial::SpatialTargetInfo& target);
+mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
+                                                 const spatial::SpatialTargetInfo& target);
+
+mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
+  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetInfo& target);
+
+mlir::FailureOr<mlir::Value> lowerSelectedResizeNearestPlan(
+  spatial::SpatResizeNearestPlanOp planOp,
+  std::optional<mlir::Value> rowStripInput,
+  const spatial::SpatialTargetInfo& target,
+  mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
+                                                 const spatial::SpatialTargetInfo& target);
+
+mlir::FailureOr<mlir::Value>
+lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                        const spatial::SpatialTargetInfo& target,
+                       mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                           std::optional<mlir::Value> rowStripInput,
+                           const spatial::SpatialTargetInfo& target,
+                           mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult
+canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                        const spatial::SpatialTargetInfo& target);
+
+mlir::FailureOr<mlir::Value>
+lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                const spatial::SpatialTargetInfo& target,
+                                mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                   std::optional<mlir::Value> rowStripInput,
+                                   const spatial::SpatialTargetInfo& target,
+                                   mlir::PatternRewriter& rewriter);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutCapabilities.cpp

@@ -0,0 +1,133 @@
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir::spatial {
+
+static LayoutAlternative denseAlternative(Operation *op) {
+  LayoutAlternative alternative;
+  alternative.operandLayouts.assign(op->getNumOperands(), PhysicalLayout::DenseNCHW);
+  alternative.resultLayout = PhysicalLayout::DenseNCHW;
+  return alternative;
+}
+
+static LayoutAlternative rowStripAlternative(Operation *op,
+                                             ArrayRef<PhysicalLayout> operandLayouts) {
+  LayoutAlternative alternative;
+  alternative.operandLayouts.assign(operandLayouts.begin(), operandLayouts.end());
+  alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+  alternative.intrinsicCost = -2;
+  return alternative;
+}
+
+static bool hasRowStripInput(ArrayRef<PhysicalLayout> operandLayouts, unsigned index) {
+  return index < operandLayouts.size()
+      && operandLayouts[index] == PhysicalLayout::NHWCRowStrip;
+}
+
+SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)) {
+    if (succeeded(canConsumeAndProduceRowStrip(*this, target)))
+      alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  }
+  else if (succeeded(canLowerConvPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatReluPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)) {
+    LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.intrinsicCost = -3;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)
+      && succeeded(canLowerResizeNearestPlanToRowStrip(*this, target)))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    if (hasRowStripInput(operandLayouts, 0))
+      alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    if (hasRowStripInput(operandLayouts, 0))
+      alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  auto resultType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (resultType && hasRowStripInput(operandLayouts, 0)
+      && isSupportedBiasAddValue(getBias(), resultType))
+    alternatives.push_back(rowStripAlternative(getOperation(),
+                                                {PhysicalLayout::NHWCRowStrip,
+                                                 PhysicalLayout::DenseNCHW}));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0)
+      && hasRowStripInput(operandLayouts, 1))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatConcatPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (!operandLayouts.empty() && llvm::all_of(operandLayouts, [](PhysicalLayout layout) {
+        return layout == PhysicalLayout::NHWCRowStrip;
+      }))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -0,0 +1,265 @@
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+#include "llvm/ADT/DenseMap.h"
+
+#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+
+#include <algorithm>
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+using LayoutMap = llvm::DenseMap<Value, spatial::PhysicalLayout>;
+
+static spatial::PhysicalLayout getSelectedLayout(const LayoutMap& layouts, Value value) {
+  if (auto it = layouts.find(value); it != layouts.end())
+    return it->second;
+  if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
+    return materialize.getTargetPhysicalLayout();
+  if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
+    return blueprint.getPhysicalLayout();
+  return spatial::PhysicalLayout::DenseNCHW;
+}
+
+static SmallVector<spatial::PhysicalLayout> getOperandLayouts(
+    Operation* op, const LayoutMap& layouts) {
+  SmallVector<spatial::PhysicalLayout> operandLayouts;
+  operandLayouts.reserve(op->getNumOperands());
+  for (Value operand : op->getOperands())
+    operandLayouts.push_back(getSelectedLayout(layouts, operand));
+  return operandLayouts;
+}
+
+static FailureOr<SmallVector<spatial::LayoutAlternative>> getAlternatives(
+    Operation* op, const LayoutMap& layouts, const spatial::SpatialTargetInfo& target) {
+  auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
+  if (!capability)
+    return failure();
+  SmallVector<spatial::LayoutAlternative> alternatives =
+      capability.getLayoutAlternatives(target, getOperandLayouts(op, layouts));
+  if (alternatives.empty())
+    return op->emitOpError("does not advertise a legal Spatial layout alternative"), failure();
+  for (const spatial::LayoutAlternative& alternative : alternatives)
+    if (alternative.operandLayouts.size() != op->getNumOperands())
+      return op->emitOpError("advertises a layout alternative with the wrong operand count"), failure();
+  return alternatives;
+}
+
+static unsigned findCurrentAlternative(
+    Operation* op, ArrayRef<spatial::LayoutAlternative> alternatives,
+    spatial::PhysicalLayout selectedResult) {
+  for (auto [index, alternative] : llvm::enumerate(alternatives))
+    if (alternative.resultLayout == selectedResult)
+      return index;
+  return 0;
+}
+
+static int64_t alternativeCost(Operation* op,
+                               const spatial::LayoutAlternative& alternative,
+                               const LayoutMap& layouts,
+                               const LayoutMap& selectedResults,
+                               const spatial::SpatialTargetInfo& target) {
+  int64_t cost = alternative.intrinsicCost;
+  SmallVector<spatial::PhysicalLayout> operandLayouts = getOperandLayouts(op, layouts);
+  for (auto [actual, required] : llvm::zip(operandLayouts, alternative.operandLayouts))
+    cost += actual != required;
+
+  Value result = op->getResult(0);
+  for (OpOperand& use : result.getUses()) {
+    auto user = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(use.getOwner());
+    if (!user) {
+      if (alternative.resultLayout != spatial::PhysicalLayout::DenseNCHW) {
+        auto flatten = dyn_cast<spatial::SpatGraphCompute>(use.getOwner());
+        if (!flatten || failed(canLowerFlattenFromRowStrip(flatten, target)))
+          ++cost;
+      }
+      continue;
+    }
+    auto userAlternatives = getAlternatives(use.getOwner(), selectedResults, target);
+    if (failed(userAlternatives))
+      continue;
+    spatial::PhysicalLayout userResult =
+        selectedResults.lookup(use.getOwner()->getResult(0));
+    unsigned userIndex = findCurrentAlternative(use.getOwner(), *userAlternatives, userResult);
+    if (use.getOperandNumber() < (*userAlternatives)[userIndex].operandLayouts.size()
+        && (*userAlternatives)[userIndex].operandLayouts[use.getOperandNumber()]
+             != alternative.resultLayout)
+      ++cost;
+  }
+  return cost;
+}
+
+static LogicalResult materializeMismatchedUses(
+    IRRewriter& rewriter, Value value, const LayoutMap& layouts,
+    const spatial::SpatialTargetInfo& target) {
+  spatial::PhysicalLayout sourceLayout = getSelectedLayout(layouts, value);
+  SmallVector<std::pair<OpOperand*, spatial::PhysicalLayout>> mismatches;
+  for (OpOperand& use : value.getUses()) {
+    Operation* userOp = use.getOwner();
+    spatial::PhysicalLayout required = spatial::PhysicalLayout::DenseNCHW;
+    if (auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(userOp)) {
+      auto alternatives = getAlternatives(userOp, layouts, target);
+      if (failed(alternatives))
+        return failure();
+      spatial::PhysicalLayout selected =
+          getSelectedLayout(layouts, userOp->getResult(0));
+      unsigned selectedIndex = findCurrentAlternative(userOp, *alternatives, selected);
+      required = (*alternatives)[selectedIndex].operandLayouts[use.getOperandNumber()];
+    }
+    else if (auto flatten = dyn_cast<spatial::SpatGraphCompute>(userOp);
+             flatten && sourceLayout == spatial::PhysicalLayout::NHWCRowStrip
+                  && succeeded(canLowerFlattenFromRowStrip(flatten, target))) {
+      continue;
+    }
+    if (required != sourceLayout)
+      mismatches.push_back({&use, required});
+  }
+
+  for (auto [use, required] : mismatches) {
+    Operation* userOp = use->getOwner();
+    rewriter.setInsertionPoint(userOp);
+    auto materialized = spatial::SpatMaterializeLayoutOp::create(
+        rewriter, userOp->getLoc(), use->get().getType(), use->get(),
+        spatial::LogicalLayoutAttr::get(
+          rewriter.getContext(), spatial::LogicalLayout::NCHW),
+        spatial::PhysicalLayoutAttr::get(rewriter.getContext(), sourceLayout),
+        spatial::PhysicalLayoutAttr::get(rewriter.getContext(),
+                                          required));
+    use->set(materialized.getResult());
+  }
+  return success();
+}
+
+static LogicalResult verifySelectedLayouts(
+    ArrayRef<Operation*> planOps, const LayoutMap& layouts,
+    const spatial::SpatialTargetInfo& target) {
+  for (Operation* op : planOps) {
+    auto selected = spatial::getSelectedPhysicalLayout(op);
+    if (!selected)
+      return op->emitOpError("requires a selected physical layout"), failure();
+    auto alternatives = getAlternatives(op, layouts, target);
+    if (failed(alternatives))
+      return failure();
+    if (llvm::none_of(*alternatives, [&](const spatial::LayoutAlternative& alternative) {
+          return alternative.resultLayout == *selected;
+        }))
+      return op->emitOpError("selected physical layout is not advertised by its layout contract"), failure();
+  }
+  return success();
+}
+
+struct SpatialLayoutPlanningPass final
+    : PassWrapper<SpatialLayoutPlanningPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass)
+
+  StringRef getArgument() const override { return "spatial-layout-planning"; }
+  StringRef getDescription() const override {
+    return "Select Spatial layout alternatives and insert explicit reconciliation barriers.";
+  }
+
+  SpatialLayoutPlanningPass() = default;
+  explicit SpatialLayoutPlanningPass(const spatial::SpatialTargetInfo& target)
+      : target(target), hasTarget(true) {}
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget) {
+      moduleOp.emitError("Spatial layout planning requires an injected SpatialTargetInfo");
+      signalPassFailure();
+      return;
+    }
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during Spatial layout planning");
+      signalPassFailure();
+      return;
+    }
+
+    func::FuncOp funcOp = *entryFunc;
+    SmallVector<Operation*> planOps;
+    for (Operation& op : funcOp.getBody().front())
+      if (isa<spatial::SpatialLayoutCapabilityInterface>(&op))
+        planOps.push_back(&op);
+
+    LayoutMap layouts;
+    for (Operation* op : planOps)
+      layouts[op->getResult(0)] = spatial::PhysicalLayout::DenseNCHW;
+
+    const size_t maxRounds = 2 * planOps.size() + 1;
+    bool converged = false;
+    for (size_t round = 0; round < maxRounds && !converged; ++round) {
+      converged = true;
+      SmallVector<Operation*> order(planOps);
+      if (round % 2)
+        std::reverse(order.begin(), order.end());
+      for (Operation* op : order) {
+        auto alternatives = getAlternatives(op, layouts, target);
+        if (failed(alternatives)) {
+          signalPassFailure();
+          return;
+        }
+        spatial::PhysicalLayout current = layouts.lookup(op->getResult(0));
+        unsigned currentIndex = findCurrentAlternative(op, *alternatives, current);
+        int64_t bestCost = alternativeCost(
+            op, (*alternatives)[currentIndex], layouts, layouts, target);
+        unsigned bestIndex = currentIndex;
+        for (auto [index, alternative] : llvm::enumerate(*alternatives)) {
+          int64_t cost = alternativeCost(op, alternative, layouts, layouts, target);
+          if (cost < bestCost) {
+            bestCost = cost;
+            bestIndex = index;
+          }
+        }
+        spatial::PhysicalLayout selected = (*alternatives)[bestIndex].resultLayout;
+        if (selected != current) {
+          layouts[op->getResult(0)] = selected;
+          converged = false;
+        }
+      }
+    }
+    if (!converged) {
+      moduleOp.emitError("Spatial layout selection did not converge within its bounded iteration budget");
+      signalPassFailure();
+      return;
+    }
+    IRRewriter rewriter(&getContext());
+    for (Operation* op : planOps) {
+      op->setAttr(spatial::kSelectedLayoutAttrName,
+                  spatial::PhysicalLayoutAttr::get(
+                      rewriter.getContext(), layouts.lookup(op->getResult(0))));
+      if (failed(materializeMismatchedUses(rewriter, op->getResult(0), layouts, target))) {
+        signalPassFailure();
+        return;
+      }
+    }
+    if (failed(verifySelectedLayouts(planOps, layouts, target))
+        || failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("Spatial layout planning verification failed");
+      signalPassFailure();
+    }
+  }
+
+  spatial::SpatialTargetInfo target;
+  bool hasTarget = false;
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass() {
+  return std::make_unique<SpatialLayoutPlanningPass>();
+}
+
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass(
+    const spatial::SpatialTargetInfo& target) {
+  return std::make_unique<SpatialLayoutPlanningPass>(target);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,7 +1,10 @@
 #include "mlir/IR/ValueRange.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
 
 #include "llvm/ADT/STLExtras.h"
 
@@ -28,6 +31,49 @@ FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anc
   return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
 }
 
+LogicalResult materializePipelineHostBuffer(
+    func::FuncOp funcOp, RewriterBase &rewriter) {
+  auto bytes = funcOp->getAttrOfType<IntegerAttr>(
+    kPipelineHostBufferBytesAttrName);
+  if (!bytes)
+    return success();
+  if (bytes.getInt() <= 0)
+    return funcOp.emitOpError(
+      "pipeline host transfer buffer must be positive");
+  ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
+  if (moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName))
+    return funcOp.emitOpError(
+      "pipeline host transfer buffer symbol already exists");
+  auto type = MemRefType::get(
+    {bytes.getInt()}, rewriter.getI8Type());
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(moduleOp.getBody());
+  memref::GlobalOp::create(
+    rewriter, funcOp.getLoc(),
+    rewriter.getStringAttr(kPipelineHostBufferName),
+    rewriter.getStringAttr("private"), TypeAttr::get(type), Attribute(),
+    UnitAttr(), IntegerAttr());
+  return success();
+}
+
+FailureOr<mlir::Value> getPipelineHostBuffer(
+    OpBuilder &builder, Operation *anchor) {
+  auto funcOp = anchor->getParentOfType<func::FuncOp>();
+  auto moduleOp = anchor->getParentOfType<ModuleOp>();
+  auto bytes = funcOp
+    ? funcOp->getAttrOfType<IntegerAttr>(kPipelineHostBufferBytesAttrName)
+    : IntegerAttr();
+  auto global = moduleOp
+    ? moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName)
+    : memref::GlobalOp();
+  if (!bytes || !global)
+    return anchor->emitOpError(
+      "requires the pipeline host transfer buffer"), failure();
+  auto type = MemRefType::get({bytes.getInt()}, builder.getI8Type());
+  return memref::GetGlobalOp::create(
+    builder, anchor->getLoc(), type, kPipelineHostBufferName).getResult();
+}
+
 Operation* getEarliestUserWithinBlock(mlir::Value value) {
   auto users = value.getUsers();
 

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -10,6 +10,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Support/LogicalResult.h"
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
@@ -23,6 +24,12 @@ namespace onnx_mlir {
 mlir::FailureOr<mlir::IntegerAttr>
 getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);
 
+mlir::LogicalResult materializePipelineHostBuffer(
+  mlir::func::FuncOp funcOp, mlir::RewriterBase &rewriter);
+
+mlir::FailureOr<mlir::Value> getPipelineHostBuffer(
+  mlir::OpBuilder &builder, mlir::Operation *anchor);
+
 template <class T>
 size_t rangeLength(const mlir::iterator_range<T> range) {
   return std::distance(range.begin(), range.end());

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -345,20 +345,39 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     auto blockArg = computeOp.getInputArgument(inputIndex);
     if (!blockArg)
       return computeOp.emitOpError("expected compute input block arguments during lowering");
-    auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(input.getDefiningOp());
+    auto channelReceive = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(
+      input.getDefiningOp());
+    auto hostWaitLoad = dyn_cast_or_null<spatial::SpatHostWaitLoadOp>(
+      input.getDefiningOp());
+    Operation *receiveOp = channelReceive
+      ? channelReceive.getOperation() : hostWaitLoad.getOperation();
     if (receiveOp && !blockArg->use_empty()) {
       rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg));
       auto outputType = cast<ShapedType>(blockArg->getType());
-      auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
+      auto outputBuffer = createEmptyTensorFromShaped(
+        rewriter, receiveOp->getLoc(), outputType);
       auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
       if (failed(sizeAttr))
         return failure();
-      Value received =
+      Value zero = arith::ConstantIndexOp::create(
-        PimReceiveOp::create(
+        rewriter, receiveOp->getLoc(), 0);
-          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer,
+      Value received;
-          arith::ConstantIndexOp::create(rewriter, receiveOp.getLoc(), 0),
+      if (hostWaitLoad) {
-          *sizeAttr, receiveOp.getSourceCoreId())
+        auto hostBuffer = getPipelineHostBuffer(rewriter, hostWaitLoad);
+        if (failed(hostBuffer))
+          return failure();
+        PimWaitOp::create(
+          rewriter, receiveOp->getLoc(), hostWaitLoad.getEventRegister(),
+          rewriter.getI32IntegerAttr(1));
+        received = PimMemCopyHostToDevOp::create(
+          rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero,
+          hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr)
           .getOutput();
+      } else {
+        received = PimReceiveOp::create(
+          rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,
+          zero, *sizeAttr, channelReceive.getSourceCoreId()).getOutput();
+      }
       blockArg->replaceAllUsesWith(received);
       markOpToRemove(receiveOp);
       continue;
@@ -383,7 +402,8 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     if (rangeLength(resultUses) == 1) {
       OpOperand& resultUse = *resultUses.begin();
       Operation* resultUser = resultUse.getOwner();
-      if (isa<spatial::SpatChannelSendOp>(resultUser))
+      if (isa<spatial::SpatChannelSendOp,
+              spatial::SpatHostStoreSyncOp>(resultUser))
         continue;
     }
 

src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp

@@ -57,10 +57,29 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   }
 };
 
-struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
+struct HostStoreSyncLowering : OpRewritePattern<spatial::SpatHostStoreSyncOp> {
   using OpRewritePattern::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
+  LogicalResult matchAndRewrite(spatial::SpatHostStoreSyncOp op, PatternRewriter& rewriter) const override {
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
+    auto hostBuffer = getPipelineHostBuffer(rewriter, op);
+    if (failed(sizeAttr) || failed(hostBuffer))
+      return failure();
+    Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
+    pim::PimMemCopyDevToHostOp::create(
+      rewriter, op.getLoc(), hostBuffer->getType(), op.getHostOffset(), zero,
+      *hostBuffer, op.getInput(), *sizeAttr);
+    auto sync = pim::PimSyncOp::create(
+      rewriter, op.getLoc(), op.getTargetCoreId(), op.getEventRegister());
+    copyRaptorDebugAttrs(op.getOperation(), sync.getOperation());
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+template <typename ReceiveOp, typename CreateReceive>
+static LogicalResult lowerReceive(
+    ReceiveOp op, PatternRewriter& rewriter, CreateReceive createReceive) {
     if (op->use_empty()) {
       rewriter.eraseOp(op);
       return success();
@@ -86,12 +105,11 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     if (failed(sizeAttr))
       return failure();
     Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
-    auto receive = pim::PimReceiveOp::create(
+    auto received = createReceive(outputBuffer, zero, *sizeAttr);
-      rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, *sizeAttr, op.getSourceCoreId());
+    if (failed(received))
-    copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
+      return failure();
-    Value received = receive.getOutput();
     if (!destinationInsert) {
-      rewriter.replaceOp(op, received);
+      rewriter.replaceOp(op, *received);
       return success();
     }
 
@@ -99,10 +117,42 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert);
     auto copy = pim::PimMemCopyOp::create(
       rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero,
-      destinationInsert.getDest(), received, *sizeAttr);
+      destinationInsert.getDest(), *received, *sizeAttr);
     rewriter.replaceOp(destinationInsert, copy.getOutput());
     rewriter.eraseOp(op);
     return success();
+}
+
+struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
+    return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
+      auto receive = pim::PimReceiveOp::create(
+        rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero,
+        sizeAttr, op.getSourceCoreId());
+      copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
+      return receive.getOutput();
+    });
+  }
+};
+
+struct HostWaitLoadLowering : OpRewritePattern<spatial::SpatHostWaitLoadOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatHostWaitLoadOp op, PatternRewriter& rewriter) const override {
+    return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
+      auto hostBuffer = getPipelineHostBuffer(rewriter, op);
+      if (failed(hostBuffer))
+        return failure();
+      auto wait = pim::PimWaitOp::create(
+        rewriter, op.getLoc(), op.getEventRegister(),
+        rewriter.getI32IntegerAttr(1));
+      copyRaptorDebugAttrs(op.getOperation(), wait.getOperation());
+      return pim::PimMemCopyHostToDevOp::create(
+        rewriter, op.getLoc(), outputBuffer.getType(), zero,
+        op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput();
+    });
   }
 };
 
@@ -148,7 +198,9 @@ struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
 } // namespace
 
 void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
-  patterns.add<ChannelSendLowering, ChannelReceiveLowering, ExtractRowsLowering, ConcatLowering>(patterns.getContext());
+  patterns.add<ChannelSendLowering, ChannelReceiveLowering,
+               HostStoreSyncLowering, HostWaitLoadLowering,
+               ExtractRowsLowering, ConcatLowering>(patterns.getContext());
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -859,6 +859,10 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
       markOpToRemove(receiveOp);
       return;
     }
+    if (auto receiveOp = dyn_cast<spatial::SpatHostWaitLoadOp>(op)) {
+      markOpToRemove(receiveOp);
+      return;
+    }
   };
 
   SmallVector<Value> originalOperands(returnOp.getOperands().begin(), returnOp.getOperands().end());

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -126,6 +126,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                     spatial::SpatConcatOp,
                     spatial::SpatChannelReceiveOp,
                     spatial::SpatChannelSendOp,
+                    spatial::SpatHostStoreSyncOp,
+                    spatial::SpatHostWaitLoadOp,
                     spatial::SpatExtractRowsOp>();
 
   RewritePatternSet initialPatterns(ctx);
@@ -140,6 +142,12 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
   walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));
 
+  if (funcOp->hasAttr(kPipelineHostBufferBytesAttrName)
+      && failed(materializePipelineHostBuffer(funcOp, rewriter))) {
+    signalPassFailure();
+    return;
+  }
+
   auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
   addReturnOutputBuffers(returnOp, rewriter);
   if (failed(allocateAndInitializeCoreLocalVariables(funcOp, rewriter))) {
@@ -182,6 +190,17 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
       continue;
     }
   }
+  SmallVector<spatial::SpatHostWaitLoadOp> hostWaitLoadOps;
+  for (auto op : funcOp.getOps<spatial::SpatHostWaitLoadOp>())
+    hostWaitLoadOps.push_back(op);
+  for (auto op : hostWaitLoadOps) {
+    bool onlyPendingRemovalUsers = llvm::all_of(
+      op->getUsers(), [&](Operation* user) {
+        return llvm::is_contained(operationsToRemove, user);
+      });
+    if (onlyPendingRemovalUsers)
+      markOpToRemove(op);
+  }
 
   RewritePatternSet coreBodyPatterns(ctx);
   populateCoreBodyPatterns(coreBodyPatterns);
@@ -202,6 +221,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                             spatial::SpatConcatOp,
                             spatial::SpatChannelReceiveOp,
                             spatial::SpatChannelSendOp,
+                            spatial::SpatHostStoreSyncOp,
+                            spatial::SpatHostWaitLoadOp,
                             spatial::SpatExtractRowsOp>();
 
   SmallVector<pim::PimCoreOp> coreOps;
@@ -251,6 +272,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   communicationTarget.addIllegalOp<spatial::SpatConcatOp,
                                    spatial::SpatChannelReceiveOp,
                                    spatial::SpatChannelSendOp,
+                                   spatial::SpatHostStoreSyncOp,
+                                   spatial::SpatHostWaitLoadOp,
                                    spatial::SpatExtractRowsOp>();
 
   RewritePatternSet communicationPatterns(ctx);

src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.cpp

@@ -430,8 +430,7 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
 
   auto targetBytes = getShapedByteSize(targetType);
   auto sourceBytes = getShapedByteSize(sourceType);
-  if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes)
+  if (succeeded(targetBytes) && succeeded(sourceBytes) && size <= *targetBytes && size <= *sourceBytes) {
-      && size <= *targetBytes && size <= *sourceBytes) {
     auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape());
     auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape());
     if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank)

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

@@ -241,6 +241,8 @@ static bool isSupportedCoreInstructionOp(Operation* op) {
              pim::PimVMVOp,
              pim::PimReceiveOp,
              pim::PimSendOp,
+             pim::PimSyncOp,
+             pim::PimWaitOp,
              pim::PimConcatOp,
              pim::PimVMMOp,
              pim::PimVVAddOp,

src/PIM/Dialect/Pim/Pim.td

@@ -118,6 +118,32 @@ def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimSyncOp : PimOp<"sync", []> {
+  let summary = "Signal an event register on another core";
+
+  let arguments = (ins
+    Index:$targetCoreId,
+    Index:$eventRegister
+  );
+
+  let assemblyFormat = [{
+    $targetCoreId `event` $eventRegister attr-dict
+  }];
+}
+
+def PimWaitOp : PimOp<"wait", []> {
+  let summary = "Wait for an event register value";
+
+  let arguments = (ins
+    Index:$eventRegister,
+    I32Attr:$waitValue
+  );
+
+  let assemblyFormat = [{
+    $eventRegister `value` $waitValue attr-dict
+  }];
+}
+
 def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from host memory into device memory";
 

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -35,6 +35,7 @@ add_pim_library(SpatialOps
   Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
   Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
   Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp
   Passes/Transforms/TrivialGraphComputeMergePass.cpp
 
   EXCLUDE_FROM_OM_LIBS

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp

@@ -219,10 +219,12 @@ static void appendReceive(BoundaryProgram &boundary,
         run->entryOffsets[run->entryOffsets.size() - 2]].family->requirement;
       CollectionTarget previousTarget {run->collection, run->positions.back()};
       bool sameEntry = previous == requirement;
-      if (sameEntry
+      bool sameRoute = run->slices.back().family->hostRouted
+                    == slice.family->hostRouted;
+      if (sameRoute && (sameEntry
           || (sameCollectionEmissionContract(previousTarget, target)
               && previous->publicationFragmentType
-                   == requirement->publicationFragmentType)) {
+                   == requirement->publicationFragmentType))) {
         run->slices.push_back(slice);
         if (sameEntry) {
           run->entryOffsets.back() = run->slices.size();

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

@@ -8,6 +8,7 @@
 #include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp"
 #include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include <array>
 namespace onnx_mlir::spatial {
 using namespace mlir;
@@ -18,6 +19,8 @@ struct LogicalTransferMetadataView {
   StaticIntSequenceChain parentCounts;
   StaticIntSequenceChain sourceCores;
   StaticIntSequenceChain targetCores;
+  StaticIntSequenceChain hostOffsets;
+  StaticIntSequenceChain eventRegisters;
   StaticIntSequenceChain targetLanes;
   StaticIntSequenceChain localOffsets;
   SmallVector<StaticIntSequenceChain> projectionOffsets;
@@ -28,7 +31,8 @@ struct LogicalTransferMetadataView {
 };
 using MetadataMember = StaticIntSequenceChain LogicalTransferMetadataView::*;
 static constexpr std::array<MetadataMember, 3> transferMetadataMembers{
-    &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores, &LogicalTransferMetadataView::targetCores};
+    &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores,
+    &LogicalTransferMetadataView::targetCores};
 struct TransferGrids {
   std::array<StaticIntGrid, 3> values;
   StaticIntGrid &channels() { return values[0]; }
@@ -41,7 +45,8 @@ template <typename Build> static FailureOr<TransferGrids> buildTransferGrids(Bui
   auto targetCores = build(transferMetadataMembers[2]);
   if (failed(channels) || failed(sourceCores) || failed(targetCores))
     return failure();
-  return TransferGrids{{std::move(*channels), std::move(*sourceCores), std::move(*targetCores)}};
+  return TransferGrids{{std::move(*channels), std::move(*sourceCores),
+                        std::move(*targetCores)}};
 }
 using GridGeometry = DeferredGridSliceGeometry;
 using StaticGeometryMember = SmallVector<StaticIntSequence> DeferredStaticSliceGeometry::*;
@@ -82,6 +87,11 @@ static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferM
   metadata.parentCounts.append(StaticIntSequence::uniform(family.requirement->exchange->externalTransferCount, count));
   metadata.sourceCores.append(family.sourceCores, familyIndex, count);
   metadata.targetCores.append(family.targetCores, familyIndex, count);
+  if (family.hostRouted) {
+    metadata.hostOffsets.append(family.hostOffsets, familyIndex, count);
+    metadata.eventRegisters.append(
+      family.eventRegisters, familyIndex, count);
+  }
   metadata.targetLanes.append(StaticIntSequence::affine(targetLane, 1, count));
   if (family.requirement->producerLocalOffsets)
     metadata.localOffsets.append(*family.requirement->producerLocalOffsets, targetLane - requirementLanes.begin, count);
@@ -172,6 +182,7 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
     appendMetadata(slice, metadataByLane[sourceLane]);
   }
   LogicalTransferMetadataView logical = buildMetadataView(run.slices);
+  ExternalTransferFamily &firstFamily = *run.slices.front().family;
   size_t actionCount = 0;
   for (const LogicalTransferMetadataView &laneMetadata : metadataByLane)
     actionCount = std::max(actionCount, laneMetadata.size());
@@ -185,6 +196,20 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
   FailureOr<StaticIntGrid> localOffsets = buildGrid(&LogicalTransferMetadataView::localOffsets, logical.localOffsets.valueAt(0));
   if (failed(transferGrids) || failed(localOffsets))
     return failure();
+  std::optional<StaticIntGrid> hostOffsets;
+  std::optional<StaticIntGrid> eventRegisters;
+  if (firstFamily.hostRouted) {
+    auto offsets = buildGrid(
+      &LogicalTransferMetadataView::hostOffsets,
+      logical.hostOffsets.valueAt(0));
+    auto events = buildGrid(
+      &LogicalTransferMetadataView::eventRegisters,
+      logical.eventRegisters.valueAt(0));
+    if (failed(offsets) || failed(events))
+      return failure();
+    hostOffsets = std::move(*offsets);
+    eventRegisters = std::move(*events);
+  }
   GridGeometry projectionGrids;
   for (auto [geometryIndex, sourceMember] : llvm::enumerate(metadataGeometryMembers)) {
     const auto &logicalValues = logical.*sourceMember;
@@ -207,7 +232,6 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
     const LogicalTransferMetadataView &source = metadataByLane[sourceLane];
     counts[sourceLane] = source.size();
   }
-  ExternalTransferFamily &firstFamily = *run.slices.front().family;
   RequirementFamily &requirement = *firstFamily.requirement;
   Operation *anchor = requirement.exchange->deferred;
   Location loc = requirement.exchange->deferred.getLoc();
@@ -217,10 +241,25 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
     auto payload = materializeSendPayload(requirement, localOffset, projectionGrids[0].empty() ? nullptr : &projection, context, loc);
     if (failed(payload))
       return failure();
-    auto send = SpatChannelSendOp::create(
+    Value sourceCore = transferGrids->sourceCores().emitLookup(
-        context.rewriter, loc, transferGrids->channels().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+      action, runtimeLane, anchor, context.constants, context.rewriter, loc);
-        transferGrids->sourceCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+    Value targetCore = transferGrids->targetCores().emitLookup(
-        transferGrids->targetCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), *payload);
+      action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    Operation *send;
+    if (firstFamily.hostRouted)
+      send = SpatHostStoreSyncOp::create(
+        context.rewriter, loc, sourceCore, targetCore,
+        hostOffsets->emitLookup(
+          action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        eventRegisters->emitLookup(
+          action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        *payload);
+    else
+      send = SpatChannelSendOp::create(
+        context.rewriter, loc,
+        transferGrids->channels().emitLookup(
+          action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        sourceCore, targetCore, *payload);
     setLogicalTransferMetadata(send, logical);
     return success();
   };
@@ -255,14 +294,45 @@ static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices
   };
   auto grids = buildTransferGrids([&](MetadataMember member) { return buildGrid(metadata.*member); });
   if (failed(grids)) return failure();
+  std::optional<StaticIntGrid> hostOffsets;
+  std::optional<StaticIntGrid> eventRegisters;
+  if (slices.front().family->hostRouted) {
+    auto offsets = buildGrid(metadata.hostOffsets);
+    auto events = buildGrid(metadata.eventRegisters);
+    if (failed(offsets) || failed(events))
+      return failure();
+    hostOffsets = std::move(*offsets);
+    eventRegisters = std::move(*events);
+  }
   Value position = lane ? lane : context.constants.getIndex(0);
   Value row = context.constants.getIndex(0);
-  auto receive = SpatChannelReceiveOp::create(context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
+  Value sourceCore = grids->sourceCores().emitLookup(
-    grids->channels().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+    row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
-    grids->sourceCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+  Value targetCore = grids->targetCores().emitLookup(
-    grids->targetCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
+    row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
+  Operation *receive;
+  Value output;
+  if (slices.front().family->hostRouted) {
+    auto op = SpatHostWaitLoadOp::create(
+      context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
+      sourceCore, targetCore,
+      hostOffsets->emitLookup(
+        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+      eventRegisters->emitLookup(
+        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
+    receive = op;
+    output = op.getOutput();
+  } else {
+    auto op = SpatChannelReceiveOp::create(
+      context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
+      grids->channels().emitLookup(
+        row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+      sourceCore, targetCore);
+    receive = op;
+    output = op.getOutput();
+  }
   setLogicalTransferMetadata(receive, metadata);
-  return receive.getOutput();
+  return output;
 }
 
 static FailureOr<SmallVector<LogicalTransferMetadataView, 0>>
@@ -315,6 +385,9 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
   SmallVector<int64_t> counts(laneCount);
   std::optional<TransferGrids> transferGrids;
   std::optional<StaticIntGrid> positions;
+  std::optional<StaticIntGrid> hostOffsets;
+  std::optional<StaticIntGrid> eventRegisters;
+  bool hostRouted = run.slices.front().family->hostRouted;
   auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount);
   if (succeeded(metadataByEntry)) {
     auto buildRows = [&](auto member) {
@@ -324,6 +397,16 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
       return StaticIntGrid::fromRows(rows);
     };
     auto grids = buildTransferGrids(buildRows);
+    if (hostRouted) {
+      auto offsets = buildRows(
+        &LogicalTransferMetadataView::hostOffsets);
+      auto events = buildRows(
+        &LogicalTransferMetadataView::eventRegisters);
+      if (failed(offsets) || failed(events))
+        return failure();
+      hostOffsets = std::move(*offsets);
+      eventRegisters = std::move(*events);
+    }
     SmallVector<StaticIntSequence> positionRows;
     for (unsigned position : run.positions)
       positionRows.push_back(StaticIntSequence::uniform(position, laneCount));
@@ -368,6 +451,16 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
       return StaticIntGrid::fromColumns(actionCount, columns, defaultValue);
     };
     auto grids = buildTransferGrids(buildGrid);
+    if (hostRouted) {
+      auto offsets = buildGrid(
+        &LogicalTransferMetadataView::hostOffsets);
+      auto events = buildGrid(
+        &LogicalTransferMetadataView::eventRegisters);
+      if (failed(offsets) || failed(events))
+        return failure();
+      hostOffsets = std::move(*offsets);
+      eventRegisters = std::move(*events);
+    }
     SmallVector<StaticIntSequence> positionColumns;
     for (const StaticIntSequenceChain &values : positionsByLane)
       positionColumns.push_back(
@@ -386,15 +479,34 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
   Value runtimeLane = lane ? lane : context.constants.getIndex(0);
   auto emitEntry = [&](Value entry, Value current) -> FailureOr<Value> {
     Type fragmentType = run.slices.front().family->requirement->publicationFragmentType;
-    auto receive =
+    Value sourceCore = transferGrids->sourceCores().emitLookup(
-        SpatChannelReceiveOp::create(context.rewriter, loc, fragmentType,
+      entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
-                                     transferGrids->channels().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+    Value targetCore = transferGrids->targetCores().emitLookup(
-                                     transferGrids->sourceCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+      entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
-                                     transferGrids->targetCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
+    Operation *receive;
+    Value output;
+    if (hostRouted) {
+      auto op = SpatHostWaitLoadOp::create(
+        context.rewriter, loc, fragmentType, sourceCore, targetCore,
+        hostOffsets->emitLookup(
+          entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        eventRegisters->emitLookup(
+          entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
+      receive = op;
+      output = op.getOutput();
+    } else {
+      auto op = SpatChannelReceiveOp::create(
+        context.rewriter, loc, fragmentType,
+        transferGrids->channels().emitLookup(
+          entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        sourceCore, targetCore);
+      receive = op;
+      output = op.getOutput();
+    }
     setLogicalTransferMetadata(receive, logical);
     Value position = positions->emitLookup(
       entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
-    return insert(receive.getOutput(), position, entry, runtimeLane, current);
+    return insert(output, position, entry, runtimeLane, current);
   };
   if (actionCount == 1 && llvm::all_of(counts, [](int64_t count) { return count == 1; }))
     return emitEntry(context.constants.getIndex(0), initial);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp

@@ -28,6 +28,11 @@ static std::optional<Event> getPlannedHead(
   while (cursor.slice < plan.slices.size()) {
     const ScheduledTransferSlice &slice = plan.slices[cursor.slice];
     ExternalTransferFamily &family = *slice.family;
+    if (family.hostRouted) {
+      ++cursor.slice;
+      cursor.offset = 0;
+      continue;
+    }
     size_t begin = slice.familyOffset + cursor.offset;
     size_t length = slice.transferCount - cursor.offset;
     auto source = family.sourceStreams.find(stream, begin, length);
@@ -243,6 +248,8 @@ LogicalResult verifyPlannedCommunicationDeadlockFree(
   DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily *family = slice.family;
+    if (family->hostRouted)
+      continue;
     if (!familyIndex.try_emplace(family, familyIndex.size()).second)
       continue;
     size_t count = family->channelIds.size();
@@ -258,18 +265,6 @@ LogicalResult verifyPlannedCommunicationDeadlockFree(
     familyChannels.emplace_back(
       first, first + static_cast<int64_t>(count));
   }
-  llvm::sort(familyChannels);
-  int64_t nextChannel = 0;
-  for (auto [firstChannel, endChannel] : familyChannels) {
-    if (firstChannel != nextChannel)
-      return anchor->emitError(
-        "planned communication channels are not exactly contiguous");
-    nextChannel = endChannel;
-  }
-  if (static_cast<uint64_t>(nextChannel) != plan.logicalTransferCount)
-    return anchor->emitError(
-      "planned communication channel count is inconsistent");
-
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily &family = *slice.family;
     for (size_t offset = 0; offset < slice.transferCount; ++offset) {
@@ -296,6 +291,8 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
   DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily *family = slice.family;
+    if (family->hostRouted)
+      continue;
     if (!familyIndex.try_emplace(family, familyIndex.size()).second)
       continue;
     for (size_t index = 0; index < family->channelIds.size(); ++index)
@@ -305,6 +302,8 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
   DenseMap<int64_t, StaticIntSequenceChain> expected;
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily &family = *slice.family;
+    if (family.hostRouted)
+      continue;
     appendEventsByCore(expected, family.channelIds, family.sourceCores,
                        slice.familyOffset, slice.transferCount, true);
     appendEventsByCore(expected, family.channelIds, family.targetCores,

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

@@ -198,6 +198,7 @@ struct ScheduledInfo {
   llvm::SmallVector<mlir::Block*> blocks;
   llvm::SmallVector<mlir::Operation*> stepAnchors;
   llvm::SmallVector<int64_t> cores;
+  llvm::SmallVector<unsigned> pipelineStages;
   unsigned stepCount = 0;
   llvm::SmallVector<ProducedValue*> produced;
   llvm::SmallVector<unsigned> streamIds;
@@ -233,6 +234,9 @@ struct ExternalTransferFamily {
   StaticIntSequence sourceCores = StaticIntSequence::uniform(0, 1);
   StaticIntSequence targetCores = StaticIntSequence::uniform(0, 1);
   StaticIntSequence channelIds = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence hostOffsets = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1);
+  bool hostRouted = false;
 };
 
 struct DeferredExchangePlan {

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

@@ -209,15 +209,26 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) {
 
 LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
                                            const ScheduledComputeMaterializationResult& materialization,
-                                           const SchedulingTarget& target) {
+                                           const SchedulingTarget& target,
+                                           size_t pipelineStages) {
   IRRewriter rewriter(funcOp.getContext());
   eraseUnusedIdentityDeferredCommunications(funcOp, rewriter);
 
-  auto transfers = buildDeferredTransferPlan(funcOp, materialization);
+  auto transfers = buildDeferredTransferPlan(
+    funcOp, materialization, pipelineStages, target.processorCount);
   if (failed(transfers))
     return funcOp.emitOpError("phase 2 failed to build symbolic transfer families");
   if (failed(placeLogicalProcessorsOnPhysicalCores(*transfers, target)))
     return failure();
+  if (transfers->pipelineHostBufferBytes != 0) {
+    auto bytes = pim::checkedCast<int64_t>(
+      transfers->pipelineHostBufferBytes, funcOp,
+      "pipeline host transfer storage");
+    if (failed(bytes))
+      return failure();
+    funcOp->setAttr(kPipelineHostBufferBytesAttrName,
+                    rewriter.getI64IntegerAttr(*bytes));
+  }
   auto schedule = scheduleDeferredCommunication(funcOp, *transfers);
   if (failed(schedule) || failed(verifyPlannedCommunicationDeadlockFree(funcOp, transfers->stepCounts, *schedule)))
     return funcOp.emitOpError("phase 2 failed to schedule symbolic communication");

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp

@@ -9,6 +9,7 @@ struct SchedulingTarget;
 
 mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp,
                                                  const ScheduledComputeMaterializationResult& materialization,
-                                                 const SchedulingTarget& target);
+                                                 const SchedulingTarget& target,
+                                                 size_t pipelineStages = 1);
 
 } // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp

@@ -11,7 +11,7 @@ using namespace mlir;
 namespace {
 
 using TransferEmissionSignature =
-  std::tuple<ScheduledInfo*, Value, Type, bool, bool, bool>;
+  std::tuple<ScheduledInfo*, Value, Type, bool, bool, bool, bool>;
 
 static TransferEmissionSignature getTransferEmissionSignature(
     const ExternalTransferFamily& family) {
@@ -21,7 +21,8 @@ static TransferEmissionSignature getTransferEmissionSignature(
           family.requirement->publicationFragmentType,
           family.requirement->graphLanes.has_value(),
           family.requirement->producerProjection.has_value(),
-          producer->scheduled->isBatch()};
+          producer->scheduled->isBatch(),
+          family.hostRouted};
 }
 
 struct StreamThreshold {

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

@@ -5,6 +5,7 @@
 #include "DeferredProjectionAnalysis.hpp"
 #include "DeferredTransferPlanning.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
 namespace onnx_mlir::spatial {
 using namespace mlir;
@@ -28,7 +29,12 @@ static FailureOr<unsigned> getStepIndex(
 
 static LogicalResult collectScheduledOperations(
     const ScheduledComputeMaterializationResult &materialization,
-    DeferredTransferPlan &plan) {
+    DeferredTransferPlan &plan,
+    size_t pipelineStageCount,
+    size_t processorCount) {
+  if (pipelineStageCount == 0 || processorCount % pipelineStageCount != 0)
+    return failure();
+  size_t stageSize = processorCount / pipelineStageCount;
   unsigned nextStream = 0;
   for (const ScheduledMaterializationRecord &record :
        materialization.materializedSchedules) {
@@ -46,8 +52,13 @@ static LogicalResult collectScheduledOperations(
     if (llvm::any_of(info.stepAnchors,
                      [](Operation *anchor) { return !anchor; }))
       return op.emitOpError("phase 2 scheduled step anchor is missing");
-    for (size_t core : record.cpus)
+    for (size_t core : record.cpus) {
+      if (core >= processorCount)
+        return op.emitOpError("phase 2 scheduled core is outside the target");
       info.cores.push_back(core);
+      if (pipelineStageCount > 1)
+        info.pipelineStages.push_back(core / stageSize);
+    }
     for (size_t lane = 0; lane < info.cores.size(); ++lane)
       info.streamIds.push_back(nextStream++);
     plan.scheduled.push_back(std::move(info));
@@ -308,17 +319,22 @@ static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
   return success();
 }
 
-static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t& nextChannel) {
+static LogicalResult buildAvailabilityFamilies(
+    DeferredTransferPlan &plan,
+    DeferredExchangePlan& exchange,
+    uint64_t& nextChannel,
+    DenseMap<int64_t, DenseMap<int64_t, unsigned>>& eventRegistersByTarget) {
+  enum class Availability { Local, Direct, Host };
   for (RequirementFamily& requirement : exchange.requirements) {
     for (LaneInterval interval : requirement.targetLanes.intervals()) {
       unsigned runBegin = interval.begin;
-      bool runLocal = false;
+      Availability runAvailability = Availability::Local;
       bool haveRun = false;
-      auto flush = [&](unsigned end) {
+      auto flush = [&](unsigned end) -> LogicalResult {
         if (!haveRun || runBegin == end)
-          return;
+          return success();
         LaneSet lanes = LaneSet::range(runBegin, end);
-        if (runLocal) {
+        if (runAvailability == Availability::Local) {
           exchange.local.push_back({&requirement, lanes});
         }
         else {
@@ -339,25 +355,72 @@ static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t&
           family.sourceCores = StaticIntSequence::uniform(requirement.producer->core, count);
           family.targetCores = StaticIntSequence::fromValues(targetCores);
           family.channelIds = StaticIntSequence::affine(nextChannel, 1, count);
+          family.hostRouted = runAvailability == Availability::Host;
+          if (family.hostRouted) {
+            SmallVector<int64_t> eventRegisters;
+            for (int64_t targetCore : targetCores) {
+              auto &registers = eventRegistersByTarget[targetCore];
+              auto it = registers.try_emplace(
+                requirement.producer->core, registers.size()).first;
+              if (it->second >= kPimEventRegisterCount)
+                return exchange.deferred.emitOpError(
+                  "pipeline host transfer requires more event registers than the target core provides");
+              eventRegisters.push_back(it->second);
+            }
+            family.eventRegisters = StaticIntSequence::fromValues(
+              eventRegisters);
+            auto fragmentType = dyn_cast<ShapedType>(
+              requirement.publicationFragmentType);
+            auto fragmentBytes = fragmentType
+              ? pim::getCheckedShapedTypeSizeInBytes(
+                  fragmentType, exchange.deferred,
+                  "pipeline host transfer fragment")
+              : FailureOr<uint64_t>(failure());
+            if (failed(fragmentBytes))
+              return failure();
+            auto bytes = pim::checkedMul<size_t>(
+              count, static_cast<size_t>(*fragmentBytes), exchange.deferred,
+              "pipeline host transfer storage");
+            if (failed(bytes))
+              return failure();
+            family.hostOffsets = StaticIntSequence::affine(
+              plan.pipelineHostBufferBytes, *fragmentBytes, count);
+            auto endOffset = pim::checkedAdd<size_t>(
+              plan.pipelineHostBufferBytes, *bytes, exchange.deferred,
+              "pipeline host transfer storage");
+            if (failed(endOffset))
+              return failure();
+            plan.pipelineHostBufferBytes = *endOffset;
+          }
           nextChannel += count;
           exchange.externalTransferCount += count;
           exchange.external.push_back(std::move(family));
         }
+        return success();
       };
       for (unsigned lane = interval.begin; lane < interval.end; ++lane) {
         unsigned sourceStream = requirement.producer->scheduled->streamIds[requirement.producer->scheduledLane];
         bool local =
           sourceStream == exchange.target->streamIds[lane] && requirement.producer->step < exchange.consumerStep;
-        if (haveRun && local != runLocal) {
+        bool crossStage = !exchange.target->pipelineStages.empty()
-          flush(lane);
+          && requirement.producer->scheduled->pipelineStages[
+               requirement.producer->scheduledLane]
+             != exchange.target->pipelineStages[lane];
+        Availability availability = local ? Availability::Local
+          : crossStage ? Availability::Host : Availability::Direct;
+        if (haveRun && availability != runAvailability) {
+          if (failed(flush(lane)))
+            return failure();
           runBegin = lane;
         }
-        runLocal = local;
+        runAvailability = availability;
         haveRun = true;
       }
-      flush(interval.end);
+      if (failed(flush(interval.end)))
+        return failure();
     }
   }
+  return success();
 }
 
 static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& plan) {
@@ -368,6 +431,7 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p
   funcOp.walk([&](SpatDeferredCommunicationOp op) { deferredOps.push_back(op); });
   GraphBatchPublicationCache publicationCache;
   uint64_t nextChannel = 0;
+  DenseMap<int64_t, DenseMap<int64_t, unsigned>> eventRegistersByTarget;
   for (SpatDeferredCommunicationOp deferred : deferredOps) {
     Operation* targetOp = deferred->getParentOfType<SpatScheduledCompute>();
     if (!targetOp)
@@ -387,7 +451,9 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p
     exchange->program = std::move(*program);
     if (failed(buildRequirementFamilies(plan, *exchange, publicationCache)))
       return failure();
-    buildAvailabilityFamilies(*exchange, nextChannel);
+    if (failed(buildAvailabilityFamilies(
+          plan, *exchange, nextChannel, eventRegistersByTarget)))
+      return failure();
     plan.exchanges.push_back(std::move(exchange));
   }
   return success();
@@ -464,9 +530,12 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa
 
 FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(
     func::FuncOp funcOp,
-    const ScheduledComputeMaterializationResult &materialization) {
+    const ScheduledComputeMaterializationResult &materialization,
+    size_t pipelineStages,
+    size_t processorCount) {
   DeferredTransferPlan plan;
-  if (failed(collectScheduledOperations(materialization, plan))
+  if (failed(collectScheduledOperations(
+        materialization, plan, pipelineStages, processorCount))
       || failed(collectProducedValues(materialization, plan))
       || failed(buildExchanges(funcOp, plan)))
     return failure();

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp

@@ -13,11 +13,14 @@ struct DeferredTransferPlan {
   llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph;
   llvm::SmallVector<std::unique_ptr<DeferredExchangePlan>> exchanges;
   llvm::SmallVector<unsigned> stepCounts;
+  size_t pipelineHostBufferBytes = 0;
 };
 
 mlir::FailureOr<DeferredTransferPlan>
 buildDeferredTransferPlan(mlir::func::FuncOp funcOp,
-                          const ScheduledComputeMaterializationResult &materialization);
+                          const ScheduledComputeMaterializationResult &materialization,
+                          size_t pipelineStages,
+                          size_t processorCount);
 
 mlir::LogicalResult retargetDeferredPublications(mlir::func::FuncOp funcOp, DeferredTransferPlan& plan);
 

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp

@@ -3,12 +3,15 @@
 #include "DeferredCommunicationRealization.hpp"
 #include "ScheduledComputeReport.hpp"
 #include "ScheduledComputeVerification.hpp"
+#include "Scheduling/PipelineScheduling.hpp"
 #include "SpatialDataflowCsvExporter.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Passes/PIMPasses.h"
 
+#include <limits>
+
 using namespace mlir;
 
 namespace onnx_mlir {
@@ -31,14 +34,47 @@ static FailureOr<func::FuncOp> requireEntry(ModuleOp moduleOp) {
   return *entry;
 }
 
+static SchedulingTarget getPipelineSchedulingTarget(
+    const SchedulingTarget& physicalTarget, size_t pipelineStages) {
+  if (pipelineStages == 1)
+    return physicalTarget;
+
+  SchedulingTarget schedulingTarget = physicalTarget;
+  schedulingTarget.processorCount = physicalTarget.processorCount / pipelineStages;
+  schedulingTarget.residentWeightCapacity = checkedMultiply(
+    physicalTarget.residentWeightCapacity, pipelineStages);
+  schedulingTarget.interProcessorLatencyNs.assign(
+    schedulingTarget.processorCount * schedulingTarget.processorCount, 0);
+  Cost latencySum = 0;
+  size_t pairCount = 0;
+  for (size_t source = 0; source < schedulingTarget.processorCount; ++source)
+    for (size_t destination = 0;
+         destination < schedulingTarget.processorCount; ++destination) {
+      Cost latency = physicalTarget.getInterProcessorLatencyNs(
+        source, destination);
+      schedulingTarget.interProcessorLatencyNs[
+        source * schedulingTarget.processorCount + destination] = latency;
+      if (source != destination) {
+        latencySum = checkedAdd(latencySum, latency);
+        ++pairCount;
+      }
+    }
+  schedulingTarget.averageInterProcessorLatencyNs = pairCount == 0
+    ? 0
+    : (latencySum + pairCount - 1) / pairCount;
+  return schedulingTarget;
+}
+
 struct ScheduleAndRealizeSpatialPass final
     : PassWrapper<ScheduleAndRealizeSpatialPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleAndRealizeSpatialPass)
 
   ScheduleAndRealizeSpatialPass() = default;
   ScheduleAndRealizeSpatialPass(const SchedulingTarget& target,
-                                SpatialDataflowExportStage exportStage)
+                                SpatialDataflowExportStage exportStage,
-      : target(target), exportStage(exportStage), hasTarget(true) {}
+                                size_t pipelineStages)
+      : target(target), exportStage(exportStage),
+        pipelineStages(pipelineStages), hasTarget(true) {}
 
   StringRef getArgument() const override { return "schedule-and-realize-spatial"; }
   StringRef getDescription() const override {
@@ -52,6 +88,13 @@ struct ScheduleAndRealizeSpatialPass final
       signalPassFailure();
       return;
     }
+    if (pipelineStages == 0 || target.processorCount % pipelineStages != 0
+        || target.residentWeightCapacity
+             > std::numeric_limits<size_t>::max() / pipelineStages) {
+      moduleOp.emitError("ScheduleAndRealizeSpatial requires valid pipeline stages and resource counts");
+      signalPassFailure();
+      return;
+    }
     auto entry = requireEntry(moduleOp);
     if (failed(entry)) {
       signalPassFailure();
@@ -59,8 +102,31 @@ struct ScheduleAndRealizeSpatialPass final
     }
     func::FuncOp entryFunc = *entry;
 
-    MergeSchedulingAnalysis analysis(entryFunc, target);
+    SchedulingTarget schedulingTarget = getPipelineSchedulingTarget(
-    MergeScheduleResult schedule = std::move(analysis.getResult());
+      target, pipelineStages);
+    ComputeGraph scheduledGraph;
+    MergeScheduleResult schedule;
+    for (;;) {
+      MergeSchedulingAnalysis analysis(
+        entryFunc, schedulingTarget,
+        pipelineStages > 1 ? target.processorCount : 0);
+      scheduledGraph = analysis.getGraph();
+      schedule = std::move(analysis.getResult());
+      std::string pipelineError;
+      if (succeeded(applyPipelineScheduling(
+            scheduledGraph, schedule, pipelineStages, target, pipelineError)))
+        break;
+      std::string splitError;
+      if (pipelineStages == 1
+          || failed(splitPipelineWorkload(
+                scheduledGraph, schedule, pipelineStages, target, splitError))) {
+        if (!splitError.empty())
+          pipelineError = splitError;
+        moduleOp.emitError() << pipelineError;
+        signalPassFailure();
+        return;
+      }
+    }
     PatternRewriter rewriter(moduleOp.getContext());
     FailureOr<ScheduledComputeMaterializationResult> materialization =
         materializeScheduledCompute(entryFunc, schedule, rewriter);
@@ -94,7 +160,8 @@ struct ScheduleAndRealizeSpatialPass final
         moduleOp, entryFunc, schedule, materializationResult.peftClassPlans,
         materializationResult.materializedSchedules);
 
-    if (failed(realizeDeferredCommunication(entryFunc, materializationResult, target))) {
+    if (failed(realizeDeferredCommunication(
+          entryFunc, materializationResult, target, pipelineStages))) {
       moduleOp.emitError("Spatial communication realization failed");
       signalPassFailure();
       return;
@@ -126,6 +193,7 @@ struct ScheduleAndRealizeSpatialPass final
 private:
   SchedulingTarget target;
   SpatialDataflowExportStage exportStage = SpatialDataflowExportStage::None;
+  size_t pipelineStages = 1;
   bool hasTarget = false;
 };
 
@@ -136,8 +204,11 @@ std::unique_ptr<Pass> createScheduleAndRealizeSpatialPass() {
 }
 
 std::unique_ptr<Pass> createScheduleAndRealizeSpatialPass(
-    const SchedulingTarget& target, SpatialDataflowExportStage exportStage) {
+    const SchedulingTarget& target,
-  return std::make_unique<ScheduleAndRealizeSpatialPass>(target, exportStage);
+    SpatialDataflowExportStage exportStage,
+    size_t pipelineStages) {
+  return std::make_unique<ScheduleAndRealizeSpatialPass>(
+    target, exportStage, pipelineStages);
 }
 
 } // namespace spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "ScheduledComputeMaterialization.hpp"
+#include "Scheduling/MergeSchedulingAnalysis.hpp"
+
+#include <memory>
+#include <optional>
+
+namespace onnx_mlir::spatial {
+
+struct ScheduledSpatialState {
+  std::optional<MergeScheduleResult> logicalSchedule;
+  std::optional<ScheduledComputeMaterializationResult> materialization;
+};
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -772,6 +772,11 @@ std::vector<ComputeGraphEdge> aggregateEdges(llvm::ArrayRef<ComputeGraphEdge> ed
 
 } // namespace
 
+TransferCost getTransferCostFromBytes(Cost bytes,
+                                      const SchedulingTarget& target) {
+  return SchedulerCostModel {target}.getTransferCostFromBytes(bytes);
+}
+
 uint64_t countComputeBodyInstructions(Region& body) {
   uint64_t numOperations = 0;
   body.walk([&](Operation* op) { numOperations = checkedAdd(numOperations, static_cast<uint64_t>(1)); });
@@ -875,9 +880,13 @@ ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& insta
   return tiled;
 }
 
-ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& target) {
+ComputeGraph buildComputeGraph(Operation* entryOp,
+                               const SchedulingTarget& target,
+                               size_t computePartitionCount) {
   ComputeGraph graph;
   SchedulerCostModel costModel {target};
+  if (computePartitionCount == 0)
+    computePartitionCount = target.processorCount;
 
   for (Region& region : entryOp->getRegions()) {
     for (Block& block : region) {
@@ -898,10 +907,10 @@ ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& targe
           if (isUsedAsWeightOnly(batch.getOperation()))
             continue;
           size_t chunkCount =
-            getBatchChunkTargetCount(batch, target.processorCount);
+            getBatchChunkTargetCount(batch, computePartitionCount);
           for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) {
             ComputeInstance instance = getBatchChunkForIndex(
-              batch, chunkIndex, target.processorCount);
+              batch, chunkIndex, computePartitionCount);
             size_t index = graph.nodes.size();
             graph.nodes.push_back({instance,
                                    getComputeInstanceCost(instance, target),
@@ -920,7 +929,7 @@ ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& targe
     for (Value input : inputs) {
       for (const ProducerValueRef& producerRef :
            collectProducerValueRefs(input, node.instance,
-                                    target.processorCount)) {
+                                    computePartitionCount)) {
         auto producerIt = graph.instanceToIndex.find(producerRef.instance);
         if (producerIt == graph.instanceToIndex.end())
           continue;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp

@@ -61,9 +61,13 @@ struct ComputeGraph {
   llvm::DenseMap<ComputeInstance, size_t> instanceToIndex;
 };
 
-ComputeGraph buildComputeGraph(mlir::Operation* entryOp, const SchedulingTarget& target);
+ComputeGraph buildComputeGraph(mlir::Operation* entryOp,
+                               const SchedulingTarget& target,
+                               size_t computePartitionCount = 0);
 bool verifyAcyclic(const ComputeGraph& graph);
 
+TransferCost getTransferCostFromBytes(Cost bytes,
+                                      const SchedulingTarget& target);
 uint64_t countComputeBodyInstructions(mlir::Region& body);
 uint64_t countComputeBodyOperationInstances(mlir::Region& body);
 Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp

@@ -89,13 +89,14 @@ void verifySchedule(const ComputeGraph& graph,
 } // namespace
 
 MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op,
-                                                 const SchedulingTarget& schedulingTarget)
+                                                 const SchedulingTarget& schedulingTarget,
-: entryOp(op), target(schedulingTarget) {
+                                                 size_t partitionCount)
+: entryOp(op), target(schedulingTarget), computePartitionCount(partitionCount) {
   result = run();
 }
 
 MergeScheduleResult MergeSchedulingAnalysis::run() {
-  ComputeGraph graph = buildComputeGraph(entryOp, target);
+  graph = buildComputeGraph(entryOp, target, computePartitionCount);
   if (!verifyAcyclic(graph))
     llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
 

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp

@@ -3,6 +3,7 @@
 #include "mlir/IR/Operation.h"
 
 #include "MergeSchedule.hpp"
+#include "ComputeGraph.hpp"
 #include "SchedulingTarget.hpp"
 
 namespace onnx_mlir {
@@ -10,12 +11,17 @@ namespace spatial {
 
 class MergeSchedulingAnalysis {
 public:
-  MergeSchedulingAnalysis(mlir::Operation* op, const SchedulingTarget& target);
+  MergeSchedulingAnalysis(mlir::Operation* op,
+                          const SchedulingTarget& target,
+                          size_t computePartitionCount = 0);
   MergeScheduleResult& getResult() { return result; }
+  const ComputeGraph& getGraph() const { return graph; }
 
 private:
   mlir::Operation* entryOp = nullptr;
   const SchedulingTarget& target;
+  size_t computePartitionCount = 0;
+  ComputeGraph graph;
   MergeScheduleResult result;
 
   MergeScheduleResult run();

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "mlir/Support/LogicalResult.h"
+
+#include <cstddef>
+#include <string>
+
+#include "ComputeGraph.hpp"
+#include "MergeSchedule.hpp"
+#include "SchedulingTarget.hpp"
+
+namespace onnx_mlir::spatial {
+
+mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
+                                            MergeScheduleResult& schedule,
+                                            size_t pipelineStages,
+                                            const SchedulingTarget& physicalTarget,
+                                            std::string& error);
+
+mlir::LogicalResult splitPipelineWorkload(const ComputeGraph& graph,
+                                          const MergeScheduleResult& schedule,
+                                          size_t pipelineStages,
+                                          const SchedulingTarget& physicalTarget,
+                                          std::string& error);
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Spatial.td

@@ -550,7 +550,8 @@ def SpatChannelSendOp : SpatOp<"channel_send", []> {
   );
 
   let assemblyFormat = [{
-    $input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($input)
+    $input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
+      attr-dict `:` type($input)
   }];
 }
 
@@ -568,7 +569,45 @@ def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
   );
 
   let assemblyFormat = [{
-    `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($output)
+    `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
+      attr-dict `:` type($output)
+  }];
+}
+
+def SpatHostStoreSyncOp : SpatOp<"host_store_sync", []> {
+  let summary = "Store a tensor to host memory and signal its consumer";
+
+  let arguments = (ins
+    Index:$sourceCoreId,
+    Index:$targetCoreId,
+    Index:$hostOffset,
+    Index:$eventRegister,
+    SpatTensor:$input
+  );
+
+  let assemblyFormat = [{
+    $input `from` $sourceCoreId `to` $targetCoreId
+      `host_offset` $hostOffset `event` $eventRegister attr-dict `:` type($input)
+  }];
+}
+
+def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> {
+  let summary = "Wait for a producer and load its tensor from host memory";
+
+  let arguments = (ins
+    Index:$sourceCoreId,
+    Index:$targetCoreId,
+    Index:$hostOffset,
+    Index:$eventRegister
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let assemblyFormat = [{
+    `from` $sourceCoreId `to` $targetCoreId
+      `host_offset` $hostOffset `event` $eventRegister attr-dict `:` type($output)
   }];
 }
 

src/PIM/Dialect/Spatial/SpatialLayoutInterface.td

@@ -0,0 +1,24 @@
+#ifndef SPATIAL_LAYOUT_INTERFACE_TD
+#define SPATIAL_LAYOUT_INTERFACE_TD
+
+include "mlir/IR/OpBase.td"
+
+def SpatialLayoutCapabilityInterface : OpInterface<"SpatialLayoutCapabilityInterface"> {
+  let description = [{
+    Contract implemented by logical Spatial planning operations that expose
+    their legal physical layout alternatives to the Spatial planner.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      "Return legal physical layout alternatives for this operation and its current operand layouts.",
+      "::llvm::SmallVector<::onnx_mlir::spatial::LayoutAlternative>",
+      "getLayoutAlternatives",
+      (ins "const ::onnx_mlir::spatial::SpatialTargetInfo &":$target,
+           "::llvm::ArrayRef<::onnx_mlir::spatial::PhysicalLayout>":$operandLayouts)>
+  ];
+
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+#endif

src/PIM/Dialect/Spatial/SpatialTargetInfo.hpp

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+namespace onnx_mlir::spatial {
+
+struct MatrixUnitShape {
+  size_t rows = 128;
+  size_t columns = 128;
+};
+
+enum class ConvLoweringStrategy : uint8_t {
+  Auto,
+  Legacy,
+  Depthwise,
+  PackedIm2Col,
+  StreamedPatch,
+  StreamedPacked,
+  OutputChannelTiled,
+  InputKTiled,
+  Tiled2D,
+};
+
+struct SpatialTargetInfo {
+  MatrixUnitShape matrixShape;
+  size_t matrixUnitsPerProcessor = 64;
+  size_t processorCount = 1;
+  size_t vectorWidth = 16;
+
+  uint64_t convIm2colMaxElements = 1ull << 20;
+  uint64_t convStreamChunkPositions = 1024;
+  ConvLoweringStrategy convLoweringStrategy = ConvLoweringStrategy::Auto;
+  bool useExperimentalConvImplementation = false;
+};
+
+} // namespace onnx_mlir::spatial

src/PIM/Pass/PIMPasses.h

@@ -0,0 +1,63 @@
+#pragma once
+
+#include "mlir/Pass/Pass.h"
+
+#include <cstddef>
+#include <memory>
+#include <string>
+
+namespace onnx_mlir {
+namespace spatial {
+struct SchedulingTarget;
+struct ScheduledSpatialState;
+struct SpatialTargetInfo;
+
+std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass();
+std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass(const SchedulingTarget& target);
+std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass(
+  const SchedulingTarget& target,
+  std::shared_ptr<ScheduledSpatialState> state);
+std::unique_ptr<mlir::Pass> createVerifyScheduledSpatialPass();
+std::unique_ptr<mlir::Pass> createVerifyScheduledSpatialPass(
+  std::shared_ptr<ScheduledSpatialState> state);
+std::unique_ptr<mlir::Pass> createRealizeSpatialCommunicationPass();
+std::unique_ptr<mlir::Pass> createRealizeSpatialCommunicationPass(
+  const SchedulingTarget& target,
+  std::shared_ptr<ScheduledSpatialState> state);
+std::unique_ptr<mlir::Pass> createVerifyRealizedSpatialPass();
+std::unique_ptr<mlir::Pass> createVerifyRealizedSpatialPass(
+  std::shared_ptr<ScheduledSpatialState> state);
+}
+
+std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
+std::unique_ptr<mlir::Pass> createONNXToSpatialPass(const spatial::SpatialTargetInfo& target);
+std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass();
+std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass(const spatial::SpatialTargetInfo& target);
+std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass();
+std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target);
+
+std::unique_ptr<mlir::Pass> createSpatialToPimPass();
+
+std::unique_ptr<mlir::Pass> createPimBufferizationPreparationPass();
+std::unique_ptr<mlir::Pass> createPimOneShotBufferizationPass();
+std::unique_ptr<mlir::Pass> createPimMemoryNormalizationPass();
+std::unique_ptr<mlir::Pass> createPimBufferizationVerificationPass();
+
+
+std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass();
+std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass(
+  size_t residentWeightCapacity);
+
+std::unique_ptr<mlir::Pass> createPimHostConstantFoldingPass();
+
+std::unique_ptr<mlir::Pass> createPimInstructionSelectionPass();
+
+std::unique_ptr<mlir::Pass> createPimLocalMemoryPlanningPass();
+
+std::unique_ptr<mlir::Pass> createPimVerificationPass();
+
+std::unique_ptr<mlir::Pass> createEmitPimCodePass();
+
+std::unique_ptr<mlir::Pass> createMessagePass(std::string message);
+
+} // namespace onnx_mlir

src/PIM/Passes/PIMPasses.h

@@ -16,7 +16,8 @@ enum class SpatialDataflowExportStage;
 std::unique_ptr<mlir::Pass> createScheduleAndRealizeSpatialPass();
 std::unique_ptr<mlir::Pass> createScheduleAndRealizeSpatialPass(
   const SchedulingTarget& target,
-  SpatialDataflowExportStage exportStage);
+  SpatialDataflowExportStage exportStage,
+  size_t pipelineStages = 1);
 }
 
 std::unique_ptr<mlir::Pass> createONNXToSpatialPass();