better comparison scripts

update pimsim-nn submodule

README.md

@@ -53,20 +53,19 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    elementwise Add/Mul/Div, ReduceMean, pooling, Relu, Sigmoid, Softmax,
    Concat, Gather, Reshape, Resize, and Split.
    The compiler-layer target adapter supplies the target-neutral
-   `SpatialTargetInfo`. Layout-aware plan ops advertise typed alternatives
+   `SpatialTargetResources`. Layout-aware plan ops advertise typed alternatives
    through the Spatial layout interface; the layout planner records the
    selected layout and explicit materialization edges. `LowerSpatialPlans`
    then pattern-lowers those selected plans. Contraction and Conv lowering
    keep semantic problems, target-dependent plans, and IR materializers in
-   separate layers.
+   separate layers. Passes and their invariant/layout analyses live under
+   `Passes/Transforms` and `Passes/Analyses`.
 
 2. **Merge, schedule, and realize Spatial communication**
-   (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
-   `TrivialGraphComputeMerge` performs local graph merging, then
-   `ScheduleSpatialGraph` materializes scheduled computes and explicit deferred
-   communication. `VerifyScheduledSpatial` checks that intermediate contract;
-   `RealizeSpatialCommunication` resolves transfers and forwarding; and
-   `VerifyRealizedSpatial` checks the final scheduled graph. Supporting
+   (`src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes`).
+   `TrivialGraphComputeMerge` performs local graph merging. One
+   `ScheduleAndRealizeSpatial` pass then owns scheduling, intermediate
+   verification, communication realization, and final verification. Supporting
    scheduling code lives under `MergeComputeNodes/Scheduling`.
 
 3. **Spatial -> Pim** (`src/PIM/Conversion/SpatialToPim`).
@@ -74,7 +73,7 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    including `pim.core`, `pim.core_batch`, communication, tensor packing, global
    tensor materialization, and return-path normalization.
 
-4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
+4. **Bufferization** (`src/PIM/Dialect/Pim/Passes/Transforms/Bufferization`).
    `PimBufferizationPreparation` establishes writable destinations without
    duplicating the one-shot copy analysis, `PimOneShotBufferization` runs
    MLIR's one-shot analysis,
@@ -83,10 +82,11 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    absence, contiguity, and copy address spaces.
 
 5. **PIM local-memory planning**
-   (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`).
+   (`src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning`).
    Computes whole-core lifetimes, reuses addresses for non-overlapping
-   allocations, and records the explicit plan in PIM IR.
-6. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and
+   allocations, and records the explicit plan in PIM IR. Reusable lifetime
+   analysis lives under `src/PIM/Dialect/Pim/Passes/Analyses`.
+6. **PIM verification and code generation** (`src/PIM/Passes/PimCodegen` and
    `src/PIM/Compiler`).
    Verifies the memory plan and other PIM invariants, then emits `.pim` core
    files, weights, and `memory.bin` / `config.json` without rerunning liveness.
@@ -99,7 +99,7 @@ Supporting pieces:
   points.
 - `src/PIM/Conversion/SpatialToGraphviz` - optional Spatial graphviz conversion
   pass.
-- `src/PIM/Pass` - pass registration and auxiliary passes.
+- `src/PIM/Passes` - pass registration and auxiliary passes.
 - `src/PIM/PimAccelerator.{cpp,hpp}` - ONNX-MLIR accelerator entry point.
 
 ## PIM compiler options
@@ -132,14 +132,8 @@ options; `onnx-mlir --help` lists the inherited ONNX-MLIR options.
   elements per convolution before streaming. Default is `1048576`.
 - `--pim-conv-stream-chunk-positions=<N>` - maximum output positions per
   streamed convolution chunk. Default is `1024`.
-- `--use-experimental-conv-impl` - use the alternate convolution lowering.
 - `--pim-detect-communication-deadlock` - statically simulate expanded
   send/receive ordering and reject blocking deadlocks. Default is off.
-- `--pim-materialize-scalar-fanout-global-order` - use the experimental,
-  expensive globally ordered scalar-fanout materializer. Default is off.
-- `--pim-trace-communication-materialization` - emit verbose communication
-  materialization diagnostics and provenance attributes. Default is off.
-- `--ignore-concat-error` - soft-fail a ConcatOp corner case.
 
 ## Standard PIM hardware profile
 

backend-simulators/pim/pim-simulator/src/lib/binary_to_instruction/mod.rs

@@ -326,9 +326,13 @@ fn append_record(
             inst_builder.make_inst(recv, inst_data_builder.build());
         }
         31 => {
+            inst_data_builder.set_offset_select_value(generic1, generic2);
             inst_builder.make_inst(wait, inst_data_builder.build());
         }
         32 => {
+            inst_data_builder
+                .set_imm_core(r2_or_imm + 1)
+                .set_offset_select_value(generic1, 0);
             inst_builder.make_inst(sync, inst_data_builder.build());
         }
         _ => bail!("unsupported PIM binary opcode {opcode}"),

backend-simulators/pim/pim-simulator/src/lib/json_to_instruction/json_isa.rs

@@ -601,7 +601,11 @@ fn json_to_wait(
     inst_data_builder: &mut InstructionDataBuilder,
     json: &Value,
 ) -> Result<()> {
-    todo!("Not present in the compiler");
+    inst_data_builder.set_offset_select_value(
+        json_i64!(json, "event_register") as i32,
+        json_i64!(json, "wait_value") as i32,
+    );
+    inst_builder.make_inst(wait, inst_data_builder.build());
     Ok(())
 }
 
@@ -610,7 +614,10 @@ fn json_to_sync(
     inst_data_builder: &mut InstructionDataBuilder,
     json: &Value,
 ) -> Result<()> {
-    todo!("Not present in the compiler");
+    inst_data_builder
+        .set_imm_core(json_i64!(json, "core") as i32 + 1)
+        .set_offset_select_value(json_i64!(json, "event_register") as i32, 0);
+    inst_builder.make_inst(sync, inst_data_builder.build());
     Ok(())
 }
 

backend-simulators/pim/pim-simulator/src/lib/pimcore.rs

@@ -93,6 +93,8 @@ struct DeadlockInfo {
     states: String,
 }
 
+type SyncEvents = Vec<[i32; 32]>;
+
 fn print_status(core_instructions: &[CoreInstructions]) {
     let mut tot_instructions = 0;
     let mut progress = 0;
@@ -135,6 +137,7 @@ impl<'a> Executable<'a> {
         } = self;
         let mut cpu_progressed = 0;
         let max_core = cpu.num_core();
+        let mut sync_events: SyncEvents = vec![[0; 32]; max_core];
         let mut cpu_index = 0;
         let mut now = SystemTime::now();
 
@@ -169,7 +172,9 @@ impl<'a> Executable<'a> {
                     now = SystemTime::now();
                 }
             }
-            handle_wait_sync(cpu, cores_instructions, core_result);
+            if handle_wait_sync(cores_instructions, &mut sync_events, core_result) {
+                cpu_progressed = 0;
+            }
             match handle_send_recv(cpu, cores_instructions, send_recv, core_result) {
                 (true, other_cpu_index) => {
                     cpu_progressed = 0;
@@ -349,12 +354,31 @@ fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockIn
     None
 }
 
-fn handle_wait_sync<'a, 'b, 'c>(
-    cpu: &'b mut CPU<'a>,
-    core_instructions: &'c mut [CoreInstructions],
+fn handle_wait_sync(
+    core_instructions: &mut [CoreInstructions],
+    events: &mut SyncEvents,
     core_result: InstructionStatus,
-) where
-    'a: 'b,
-    'a: 'c,
-{
+) -> bool {
+    match core_result {
+        InstructionStatus::Sync(data) => {
+            let (source, target) = data.get_core_immcore();
+            let register = data.offset_select() as usize;
+            events[target as usize][register] += 1;
+            core_instructions[source as usize].program_counter += 1;
+            true
+        }
+        InstructionStatus::Waiting(data) => {
+            let core = data.core_indx() as usize;
+            let register = data.offset_select() as usize;
+            let value = data.offset_value();
+            if events[core][register] >= value {
+                events[core][register] -= value;
+                core_instructions[core].program_counter += 1;
+                true
+            } else {
+                false
+            }
+        }
+        _ => false,
+    }
 }

backend-simulators/pim/pim-simulator/src/lib/send_recv.rs

@@ -134,7 +134,7 @@ where
                 send_recv.sending[sender] = None;
                 send_recv.receiving[receiver] = None;
             }
-            (transfered, receiver)
+            (transfered, if transfered { receiver } else { 0 })
         }
         InstructionStatus::Reciving(instruction_data) => {
             let (core_idx, imm_core) = instruction_data.get_core_immcore();
@@ -163,7 +163,7 @@ where
                 send_recv.sending[sender] = None;
                 send_recv.receiving[receiver] = None;
             }
-            (transfered, sender)
+            (transfered, if transfered { sender } else { 0 })
         }
         _ => (false, 0),
     }

backend-simulators/pim/pim-simulator/tests/sync.rs

@@ -295,3 +295,68 @@ fn multiple_send_recv_test() {
         "send_recv failed to store"
     );
 }
+
+#[test]
+fn sync_wait_tokens_test() {
+    let cpu = common::empty_cpu(2);
+    let mut cores = CoreInstructionsBuilder::new(2);
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+
+    data.set_core_indx(1).fix_core_indx();
+    for _ in 0..2 {
+        instructions.make_inst(
+            sync,
+            data.set_imm_core(2).set_offset_select_value(0, 0).build(),
+        );
+    }
+    cores.set_core(1, instructions.build());
+
+    data.set_core_indx(2).fix_core_indx();
+    for _ in 0..2 {
+        instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    }
+    cores.set_core(2, instructions.build());
+
+    Executable::new(cpu, cores.build()).execute().unwrap();
+}
+
+#[test]
+fn blocked_transfers_do_not_starve_sync_producer() {
+    let cpu = common::empty_cpu(4);
+    let mut cores = CoreInstructionsBuilder::new(4);
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(1).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(recv, data.set_rd(1).set_imm_core(2).set_imm_len(1).build());
+    instructions.make_inst(send, data.set_r1(1).set_imm_core(3).set_imm_len(1).build());
+    cores.set_core(1, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(2).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(wait, data.set_offset_select_value(0, 1).build());
+    instructions.make_inst(send, data.set_r1(1).set_imm_core(1).set_imm_len(1).build());
+    cores.set_core(2, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(3).fix_core_indx();
+    instructions.make_inst(sldi, data.set_rdimm(1, 0).build());
+    instructions.make_inst(recv, data.set_rd(1).set_imm_core(1).set_imm_len(1).build());
+    cores.set_core(3, instructions.build());
+
+    let mut instructions = InstructionsBuilder::new();
+    let mut data = InstructionDataBuilder::new();
+    data.set_core_indx(4).fix_core_indx();
+    instructions.make_inst(
+        sync,
+        data.set_imm_core(2).set_offset_select_value(0, 0).build(),
+    );
+    cores.set_core(4, instructions.build());
+
+    Executable::new(cpu, cores.build()).execute().unwrap();
+}

src/PIM/CMakeLists.txt

@@ -94,7 +94,7 @@ endfunction()
 
 add_subdirectory(Dialect)
 add_subdirectory(Common)
-add_subdirectory(Pass)
+add_subdirectory(Passes)
 add_subdirectory(Compiler)
 add_subdirectory(Conversion)
 

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, "",     "",     "",      ""    }, // sync
+  {false, false, false, "",     "event_register", "wait_value", ""}, // wait
+  {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 {
@@ -70,11 +72,6 @@ llvm::cl::opt<bool>
                  llvm::cl::init(false),
                  llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> useExperimentalConvImpl("use-experimental-conv-impl",
-                                            llvm::cl::desc("Use experimental implementation for convolution"),
-                                            llvm::cl::init(false),
-                                            llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<uint64_t> pimConvIm2colMaxElements(
   "pim-conv-im2col-max-elements",
   llvm::cl::desc("Maximum number of im2col elements to materialize globally for one Conv before streaming/chunking"),
@@ -87,6 +84,11 @@ llvm::cl::opt<uint64_t> pimConvStreamChunkPositions(
   llvm::cl::init(1024),
   llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool> pimReportConvLowering("pim-report-conv-lowering",
+                                          llvm::cl::desc("Emit a bounded Conv lowering report"),
+                                          llvm::cl::init(true),
+                                          llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
                                 llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"),
                                 llvm::cl::init(false),
@@ -98,15 +100,9 @@ llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder(
-  "pim-materialize-scalar-fanout-global-order",
-  llvm::cl::desc("Experimental expensive materializer mode: emit scalar-source fanout as globally ordered communication events instead of all-send fanout loops"),
-  llvm::cl::init(false),
-  llvm::cl::cat(OnnxMlirOptions));
-
-llvm::cl::opt<bool> pimTraceCommunicationMaterialization(
-  "pim-trace-communication-materialization",
-  llvm::cl::desc("Emit verbose materializer-time diagnostics and provenance attributes for every Spatial communication op"),
+llvm::cl::opt<bool> pimVerifyBufferizationCopyFreedom(
+  "pim-verify-bufferization-copy-freedom",
+  llvm::cl::desc("Run the expensive official PIM tensor-copy freedom proof before bufferization"),
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
@@ -116,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));
@@ -126,11 +128,6 @@ llvm::cl::opt<std::string> pimTargetConfig(
   llvm::cl::init(""),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool>
-  ignoreConcatError("ignore-concat-error",
-                    llvm::cl::desc("Ignore ConcatOp corner case: do not assert and do a simplification"),
-                    llvm::cl::init(false));
-
 bool hasExplicitPimCoreCount() { return coresCount.getNumOccurrences() != 0; }
 
 void verifyExplicitPimCoreCount() {
@@ -140,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

@@ -55,14 +55,14 @@ extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
 extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;
 
 extern llvm::cl::opt<bool> pimOnlyCodegen;
-extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<bool> pimEmitJson;
+extern llvm::cl::opt<bool> pimReportConvLowering;
 extern llvm::cl::opt<bool> pimDetectCommunicationDeadlock;
-extern llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder;
-extern llvm::cl::opt<bool> pimTraceCommunicationMaterialization;
+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;
@@ -70,13 +70,6 @@ extern llvm::cl::opt<uint64_t> pimConvStreamChunkPositions;
 
 bool hasExplicitPimCoreCount();
 void verifyExplicitPimCoreCount();
-
-// This option, by default set to false, will ignore an error when resolving a
-// specific tiles of the operands of a concat. This specific case is when the
-// wanted tile is generated by two separate operands of the concat. If this is
-// set to false, this corner case will assert an error. If this is set to true,
-// a simplification is performed and only the tile from the first operand is
-// taken.
-extern llvm::cl::opt<bool> ignoreConcatError;
+void verifyPimPipelineStages();
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -14,11 +14,12 @@
 
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 #include "src/Compiler/CompilerPasses.hpp"
 
 #define DEBUG_TYPE "PimCompilerUtils"
@@ -97,17 +98,37 @@ spatial::ConvLoweringStrategy getSpatialConvLoweringStrategy(PimConvLoweringType
   llvm_unreachable("unknown PIM Conv lowering strategy");
 }
 
-spatial::SpatialTargetInfo getPimSpatialTargetInfo(const spatial::SchedulingTarget& target) {
-  spatial::SpatialTargetInfo info;
-  info.matrixShape = {target.matrixRows, target.matrixColumns};
-  info.matrixUnitsPerProcessor = target.residentWeightCapacity;
-  info.processorCount = target.processorCount;
-  info.vectorWidth = target.vectorWidth;
-  info.convIm2colMaxElements = pimConvIm2colMaxElements.getValue();
-  info.convStreamChunkPositions = pimConvStreamChunkPositions.getValue();
-  info.convLoweringStrategy = getSpatialConvLoweringStrategy(pimConvLowering.getValue());
-  info.useExperimentalConvImplementation = useExperimentalConvImpl.getValue();
-  return info;
+spatial::SpatialDataflowExportStage getPimSpatialDataflowExportStage(
+    PimSpatialDataflowExportType stage) {
+  switch (stage) {
+  case SpatialDataflowExportNone:     return spatial::SpatialDataflowExportStage::None;
+  case SpatialDataflowExportSpatial1: return spatial::SpatialDataflowExportStage::Spatial1;
+  case SpatialDataflowExportSpatial2: return spatial::SpatialDataflowExportStage::Spatial2;
+  case SpatialDataflowExportSpatial3: return spatial::SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportSpatial4: return spatial::SpatialDataflowExportStage::Spatial4;
+  case SpatialDataflowExportAll:      return spatial::SpatialDataflowExportStage::All;
+  }
+  llvm_unreachable("unknown PIM Spatial dataflow export stage");
+}
+
+spatial::SpatialTargetResources getPimSpatialTargetResources(const spatial::SchedulingTarget& target) {
+  spatial::SpatialTargetResources resources;
+  resources.matrixShape = {target.matrixRows, target.matrixColumns};
+  resources.matrixUnitsPerProcessor = target.residentWeightCapacity;
+  resources.processorCount = target.processorCount;
+  resources.vectorWidth = target.vectorWidth;
+  if (failed(resources.verify()))
+    llvm::report_fatal_error("PIM target resources are incomplete");
+  return resources;
+}
+
+ONNXToSpatialPlanningOptions getPimONNXToSpatialPlanningOptions() {
+  ONNXToSpatialPlanningOptions options;
+  options.convIm2colMaxElements = pimConvIm2colMaxElements.getValue();
+  options.convStreamChunkPositions = pimConvStreamChunkPositions.getValue();
+  options.forcedConvStrategy = getSpatialConvLoweringStrategy(pimConvLowering.getValue());
+  options.reportConvLowering = pimReportConvLowering.getValue();
+  return options;
 }
 
 const llvm::json::Object& requireObject(const llvm::json::Object& object,
@@ -309,11 +330,14 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
                   EmissionTargetType& emissionTarget,
                   std::string outputNameNoExt) {
   verifyExplicitPimCoreCount();
+  verifyPimPipelineStages();
+  spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
+  spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget);
 
   if (pimOnlyCodegen) {
     pm.addPass(createPimInstructionSelectionPass());
     pm.addPass(createPimLocalMemoryPlanningPass());
-    pm.addPass(createPimVerificationPass());
+    pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
     pm.addPass(createEmitPimCodePass());
     return;
   }
@@ -322,28 +346,26 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     addONNXToMLIRPasses(pm, /*target CPU*/ false);
 
   if (pimEmissionTarget >= EmitSpatial) {
-    spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
-    spatial::SpatialTargetInfo targetInfo = getPimSpatialTargetInfo(schedulingTarget);
-    pm.addPass(createONNXToSpatialPass(targetInfo));
-    pm.addPass(createSpatialLayoutPlanningPass(targetInfo));
-    pm.addPass(createLowerSpatialPlansPass(targetInfo));
+    ONNXToSpatialPlanningOptions planningOptions = getPimONNXToSpatialPlanningOptions();
+    spatial::SpatialDataflowExportStage exportStage =
+      getPimSpatialDataflowExportStage(pimExportSpatialDataflow.getValue());
+    pm.addPass(createONNXToSpatialPass(targetResources, planningOptions));
+    pm.addPass(createSpatialLayoutPlanningPass(targetResources));
+    pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage));
     pm.addPass(createTrivialGraphComputeMergePass(
-      schedulingTarget.residentWeightCapacity));
-    auto scheduledState = std::make_shared<spatial::ScheduledSpatialState>();
-    pm.addPass(spatial::createScheduleSpatialGraphPass(schedulingTarget, scheduledState));
-    pm.addPass(spatial::createVerifyScheduledSpatialPass(scheduledState));
-    pm.addPass(spatial::createRealizeSpatialCommunicationPass(schedulingTarget, scheduledState));
-    pm.addPass(spatial::createVerifyRealizedSpatialPass(scheduledState));
+      schedulingTarget.residentWeightCapacity, exportStage));
+    pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
+      schedulingTarget, exportStage, pipelineStages.getValue()));
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
 
   if (pimEmissionTarget >= EmitPim) {
-    pm.addPass(createSpatialToPimPass());
+    pm.addPass(createSpatialToPimPass(targetResources));
     pm.addPass(createMessagePass("Spatial lowered to Pim"));
   }
 
   if (pimEmissionTarget >= EmitPimBufferized) {
-    pm.addPass(createPimBufferizationPreparationPass());
+    pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue()));
     pm.addPass(createPimOneShotBufferizationPass());
     pm.addPass(createPimMemoryNormalizationPass());
     pm.addPass(createPimBufferizationVerificationPass());
@@ -358,7 +380,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(createMessagePass("Pim instructions selected"));
     pm.addPass(createPimLocalMemoryPlanningPass());
     pm.addPass(createMessagePass("Pim local memory planned"));
-    pm.addPass(createPimVerificationPass());
+    pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
     pm.addPass(createMessagePass("Pim verified"));
     pm.addPass(createEmitPimCodePass());
     pm.addPass(createMessagePass("Pim code emitted"));

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/CMakeLists.txt

@@ -5,7 +5,7 @@ add_public_tablegen_target(ONNXToSpatialIncGen)
 add_pim_library(OMONNXToSpatial
   Patterns.cpp
   CompileTime.cpp
-  ONNXToSpatialVerifier.cpp
+  Passes/Analyses/ONNXToSpatialVerifier.cpp
   Patterns/Pre.cpp
   Patterns/Post.cpp
   Patterns/Math/Conv.cpp
@@ -26,14 +26,14 @@ add_pim_library(OMONNXToSpatial
   Patterns/Tensor/Slice.cpp
   Patterns/Tensor/Split.cpp
   Patterns/Tensor/Transpose.cpp
-  ONNXToSpatialPass.cpp
-  SpatialLayoutCapabilities.cpp
-  SpatialLayoutPlanningPass.cpp
-  LowerSpatialPlansPass.cpp
+  Passes/Transforms/ONNXToSpatialPass.cpp
+  Passes/Analyses/SpatialLayoutCapabilities.cpp
+  Passes/Transforms/SpatialLayoutPlanningPass.cpp
+  Passes/Transforms/SpatialPlanLoweringPatterns.cpp
+  Passes/Transforms/LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
   Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.cpp
-  Common/ContractionMaterialization.cpp
   Common/ContractionPlanning.cpp
   Common/MatrixProductLowering.cpp
   Common/RowStripLayoutUtils.cpp

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

@@ -1,39 +0,0 @@
-#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

@@ -1,23 +0,0 @@
-#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/Common/ContractionPlanning.cpp

@@ -12,56 +12,30 @@ static int64_t ceilDivide(int64_t value, int64_t divisor) {
   return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
 }
 
-static llvm::SmallVector<int64_t> buildBatchMap(
-    llvm::ArrayRef<int64_t> sourceShape,
-    llvm::ArrayRef<int64_t> outputShape) {
-  llvm::SmallVector<int64_t> map(outputShape.size(), -1);
-  const int64_t offset = outputShape.size() - sourceShape.size();
-  for (int64_t source = 0; source < static_cast<int64_t>(sourceShape.size()); ++source) {
-    const int64_t output = source + offset;
-    if (sourceShape[source] != 1)
-      map[output] = source;
-  }
-  return map;
-}
-
 } // namespace
 
 ContractionPlan makeContractionPlan(
     const ContractionProblem& problem,
-    const spatial::SpatialTargetInfo& target,
+    const spatial::SpatialTargetResources& target,
     ContractionPlanKind kind,
     int64_t laneCount,
     int64_t fragmentRows) {
   ContractionPlan plan;
-  plan.problem = problem;
-  plan.kind = kind;
-  plan.tileM = std::max<int64_t>(1, target.matrixShape.rows);
   plan.tileK = std::max<int64_t>(1, target.matrixShape.rows);
   plan.tileN = std::max<int64_t>(1, target.matrixShape.columns);
   plan.reductionSlices = std::max<int64_t>(1, ceilDivide(problem.k, plan.tileK));
   plan.outputTiles = std::max<int64_t>(1, ceilDivide(problem.n, plan.tileN));
-  plan.rowTiles = std::max<int64_t>(1, ceilDivide(problem.m, plan.tileM));
-  plan.fragmentRows = std::max<int64_t>(
-      1, fragmentRows != 0 ? fragmentRows : plan.tileM);
-  plan.lhsBatchMap = buildBatchMap(problem.lhsBatchShape, problem.outputBatchShape);
-  plan.rhsBatchMap = buildBatchMap(problem.rhsBatchShape, problem.outputBatchShape);
+  const int64_t rowsPerLane = std::max<int64_t>(
+      1, fragmentRows != 0 ? fragmentRows : target.matrixShape.rows);
 
   if (laneCount != 0)
     plan.laneCount = laneCount;
   else if (kind == ContractionPlanKind::StaticTiled)
     plan.laneCount = problem.batch * problem.m * plan.reductionSlices * plan.outputTiles;
   else if (kind == ContractionPlanKind::GroupedRowDynamicVVD)
-    plan.laneCount = problem.batch * ceilDivide(problem.m, plan.fragmentRows);
+    plan.laneCount = problem.batch * ceilDivide(problem.m, rowsPerLane);
   else
     plan.laneCount = problem.batch * problem.m * problem.n;
-
-  plan.expectedMvmCount = kind == ContractionPlanKind::StaticTiled ? plan.laneCount : 0;
-  plan.expectedVvdCount = kind == ContractionPlanKind::StaticTiled ? 0 : plan.laneCount;
-  plan.expectedVectorCount = plan.laneCount * plan.reductionSlices;
-  if (problem.resultElementType && problem.n > 0)
-    plan.physicalFragmentType = mlir::RankedTensorType::get(
-        {plan.fragmentRows, problem.n}, problem.resultElementType);
   return plan;
 }
 

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

@@ -2,7 +2,7 @@
 
 #include "ContractionProblem.hpp"
 
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 
@@ -13,27 +13,16 @@ enum class ContractionPlanKind {
 };
 
 struct ContractionPlan {
-  ContractionProblem problem;
-  ContractionPlanKind kind = ContractionPlanKind::StaticTiled;
-  int64_t tileM = 1;
   int64_t tileK = 1;
   int64_t tileN = 1;
-  int64_t fragmentRows = 1;
   int64_t reductionSlices = 1;
   int64_t outputTiles = 1;
-  int64_t rowTiles = 1;
   int64_t laneCount = 0;
-  int64_t expectedMvmCount = 0;
-  int64_t expectedVvdCount = 0;
-  int64_t expectedVectorCount = 0;
-  llvm::SmallVector<int64_t> lhsBatchMap;
-  llvm::SmallVector<int64_t> rhsBatchMap;
-  mlir::RankedTensorType physicalFragmentType;
 };
 
 ContractionPlan makeContractionPlan(
     const ContractionProblem& problem,
-    const spatial::SpatialTargetInfo& target,
+    const spatial::SpatialTargetResources& target,
     ContractionPlanKind kind,
     int64_t laneCount = 0,
     int64_t fragmentRows = 0);

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

@@ -8,8 +8,6 @@
 
 namespace onnx_mlir {
 
-enum class ContractionOrigin { Gemm, MatMul };
-
 struct ContractionProblem {
   llvm::SmallVector<int64_t> lhsBatchShape;
   llvm::SmallVector<int64_t> rhsBatchShape;
@@ -20,16 +18,9 @@ struct ContractionProblem {
   int64_t m = 0;
   int64_t k = 0;
   int64_t n = 0;
-  ContractionOrigin origin = ContractionOrigin::MatMul;
   mlir::Type lhsElementType;
   mlir::Type rhsElementType;
   mlir::Type resultElementType;
-  bool lhsTransposed = false;
-  bool rhsTransposed = false;
-  bool lhsWasVector = false;
-  bool rhsWasVector = false;
-  float alpha = 1.0f;
-  float beta = 1.0f;
 };
 
 } // namespace onnx_mlir

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

@@ -8,8 +8,8 @@ namespace onnx_mlir {
 
 namespace spatial {
 class SpatBlueprintOp;
-class SpatGraphCompute;
-struct SpatialTargetInfo;
+class SpatFlattenPlanOp;
+struct SpatialTargetResources;
 } // namespace spatial
 
 inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
@@ -94,13 +94,13 @@ mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysical
                                                  mlir::Location loc);
 
 mlir::LogicalResult canLowerFlattenFromRowStrip(
-    spatial::SpatGraphCompute flattenOp,
-    const spatial::SpatialTargetInfo& target);
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target);
 
 mlir::LogicalResult lowerFlattenFromRowStrip(
     const RowStripPhysicalValue& input,
-    spatial::SpatGraphCompute flattenOp,
-    const spatial::SpatialTargetInfo& target,
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target,
     mlir::PatternRewriter& rewriter);
 
 } // namespace onnx_mlir

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

@@ -69,7 +69,7 @@ DenseMap<CoreId, SmallVector<Value>>
 sliceVectorPerCrossbarPerCore(const Value& vectorToSlice,
                               PatternRewriter& rewriter,
                               Location loc,
-                              const spatial::SpatialTargetInfo& target) {
+                              const spatial::SpatialTargetResources& target) {
   SmallVector<Value> slices = sliceVector(
     vectorToSlice, static_cast<int64_t>(target.matrixShape.rows), rewriter, loc);
   DenseMap<CoreId, SmallVector<Value>> slicesPerCore;

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

@@ -7,7 +7,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 
@@ -30,6 +30,6 @@ llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPer
   const mlir::Value& vectorToSlice,
   mlir::PatternRewriter& rewriter,
   mlir::Location loc,
-  const spatial::SpatialTargetInfo& target);
+  const spatial::SpatialTargetResources& target);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp

@@ -0,0 +1,32 @@
+#pragma once
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+
+#include <cstdint>
+
+namespace onnx_mlir::spatial {
+
+enum class ConvLoweringStrategy : uint8_t {
+  Auto,
+  Legacy,
+  Depthwise,
+  PackedIm2Col,
+  StreamedPatch,
+  StreamedPacked,
+  OutputChannelTiled,
+  InputKTiled,
+  Tiled2D,
+};
+
+} // namespace onnx_mlir::spatial
+
+namespace onnx_mlir {
+
+struct ONNXToSpatialPlanningOptions {
+  uint64_t convIm2colMaxElements = 0;
+  uint64_t convStreamChunkPositions = 0;
+  spatial::ConvLoweringStrategy forcedConvStrategy = spatial::ConvLoweringStrategy::Auto;
+  bool reportConvLowering = true;
+};
+
+} // namespace onnx_mlir

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

@@ -6,7 +6,7 @@
 #include "Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -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;
@@ -147,6 +149,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatBiasAddPlanOp,
             spatial::SpatAddPlanOp,
             spatial::SpatConcatPlanOp,
+            spatial::SpatFlattenPlanOp,
             spatial::SpatReluPlanOp,
             spatial::SpatSiluPlanOp,
             spatial::SpatResizeNearestPlanOp,
@@ -162,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";
@@ -181,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/Passes/Analyses/SpatialLayoutCapabilities.cpp

@@ -1,5 +1,6 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -28,7 +29,7 @@ static bool hasRowStripInput(ArrayRef<PhysicalLayout> operandLayouts, unsigned i
 }
 
 SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)) {
     if (succeeded(canConsumeAndProduceRowStrip(*this, target)))
@@ -43,8 +44,26 @@ SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
   return alternatives;
 }
 
+SmallVector<LayoutAlternative> SpatFlattenPlanOp::getLayoutAlternatives(
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (!operandLayouts.empty()
+      && operandLayouts[0] == PhysicalLayout::Fragmented) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    alternative.operandLayouts[0] = PhysicalLayout::Fragmented;
+    alternatives.push_back(std::move(alternative));
+  }
+  if (hasRowStripInput(operandLayouts, 0)
+      && succeeded(canLowerFlattenFromRowStrip(*this, target))) {
+    LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.resultLayout = PhysicalLayout::DenseNCHW;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
 SmallVector<LayoutAlternative> SpatReluPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0))
     alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
@@ -52,7 +71,7 @@ SmallVector<LayoutAlternative> SpatReluPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)) {
     LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
@@ -63,7 +82,7 @@ SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)
       && succeeded(canLowerResizeNearestPlanToRowStrip(*this, target)))
@@ -72,7 +91,7 @@ SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) {
     LayoutAlternative alternative = denseAlternative(getOperation());
@@ -86,7 +105,7 @@ SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) {
     LayoutAlternative alternative = denseAlternative(getOperation());
@@ -100,7 +119,7 @@ SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternative
 }
 
 SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   auto resultType = dyn_cast<RankedTensorType>(getOutput().getType());
   if (resultType && hasRowStripInput(operandLayouts, 0)
@@ -112,7 +131,7 @@ SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0)
       && hasRowStripInput(operandLayouts, 1))
@@ -121,7 +140,7 @@ SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatConcatPlanOp::getLayoutAlternatives(
-    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (!operandLayouts.empty() && llvm::all_of(operandLayouts, [](PhysicalLayout layout) {
         return layout == PhysicalLayout::NHWCRowStrip;

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/LowerSpatialPlansPass.cpp

@@ -0,0 +1,136 @@
+#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/DialectConversion.h"
+
+#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+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;
+  LowerSpatialPlansPass(const spatial::SpatialTargetResources& target,
+                        const ONNXToSpatialPlanningOptions& options,
+                        spatial::SpatialDataflowExportStage exportStage)
+      : target(target), planningOptions(options), exportStage(exportStage), hasTarget(true) {}
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget) {
+      moduleOp.emitError("Spatial plan lowering requires an injected SpatialTargetResources");
+      signalPassFailure();
+      return;
+    }
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
+      signalPassFailure();
+      return;
+    }
+    func::FuncOp funcOp = *entryFunc;
+    auto verifyLogicalPhase = [&](StringRef stage) -> bool {
+      if (succeeded(verifyLogicalSpatialGraphInvariants(funcOp)))
+        return true;
+      moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
+      signalPassFailure();
+      return false;
+    };
+
+    if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
+      return;
+    if (failed(verifySelectedSpatialLayouts(funcOp, target))) {
+      moduleOp.emitError("selected Spatial layout verification failed");
+      signalPassFailure();
+      return;
+    }
+
+    MLIRContext* ctx = moduleOp.getContext();
+    RewritePatternSet patterns(ctx);
+    populateSpatialPlanLoweringPatterns(patterns, ctx, target, planningOptions);
+
+    ConversionTarget conversionTarget(*ctx);
+    conversionTarget.addLegalDialect<spatial::SpatialDialect,
+                                     tensor::TensorDialect,
+                                     linalg::LinalgDialect,
+                                     affine::AffineDialect,
+                                     arith::ArithDialect,
+                                     scf::SCFDialect,
+                                     func::FuncDialect>();
+    conversionTarget.addIllegalDialect<ONNXDialect>();
+    conversionTarget.addLegalOp<ONNXEntryPointOp>();
+    conversionTarget.addIllegalOp<spatial::SpatConv2DPlanOp,
+                                  spatial::SpatFlattenPlanOp,
+                                  spatial::SpatReluPlanOp,
+                                  spatial::SpatSiluPlanOp,
+                                  spatial::SpatResizeNearestPlanOp,
+                                  spatial::SpatMaxPool2DPlanOp,
+                                  spatial::SpatGlobalAveragePoolPlanOp,
+                                  spatial::SpatBiasAddPlanOp,
+                                  spatial::SpatAddPlanOp,
+                                  spatial::SpatConcatPlanOp,
+                                  spatial::SpatMaterializeLayoutOp>();
+    conversionTarget.addDynamicallyLegalOp<spatial::SpatBlueprintOp>(
+        [](spatial::SpatBlueprintOp blueprint) {
+          return spatial::isFragmentAssembly(blueprint.getMode());
+        });
+
+    if (failed(applyFullConversion(funcOp, conversionTarget,
+                                   std::move(patterns)))) {
+      moduleOp.emitError("failed to lower Spatial plans and layout materialization");
+      signalPassFailure();
+      return;
+    }
+
+    dumpModule(moduleOp, "spatial1_graph");
+    if (spatial::shouldExportSpatialDataflowStage(
+            exportStage, spatial::SpatialDataflowExportStage::Spatial1)
+        && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
+      signalPassFailure();
+      return;
+    }
+    verifyLogicalPhase("at the end of LowerSpatialPlans");
+  }
+
+  spatial::SpatialTargetResources target;
+  ONNXToSpatialPlanningOptions planningOptions;
+  spatial::SpatialDataflowExportStage exportStage = spatial::SpatialDataflowExportStage::None;
+  bool hasTarget = false;
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createLowerSpatialPlansPass() {
+  return std::make_unique<LowerSpatialPlansPass>();
+}
+
+std::unique_ptr<Pass> createLowerSpatialPlansPass(
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options,
+    spatial::SpatialDataflowExportStage exportStage) {
+  return std::make_unique<LowerSpatialPlansPass>(target, options, exportStage);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/ONNXToSpatialPass.cpp

@@ -12,15 +12,15 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
-#include "ONNXToSpatialVerifier.hpp"
 
 using namespace mlir;
 
@@ -34,14 +34,16 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
   StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
 
   ONNXToSpatialPass() = default;
-  explicit ONNXToSpatialPass(const spatial::SpatialTargetInfo& target)
-      : target(target), hasTarget(true) {}
+  ONNXToSpatialPass(const spatial::SpatialTargetResources& target,
+                    const ONNXToSpatialPlanningOptions& options)
+      : target(target), planningOptions(options), hasTarget(true) {}
   ONNXToSpatialPass(const ONNXToSpatialPass& pass)
-      : target(pass.target), hasTarget(pass.hasTarget) {}
+      : target(pass.target), planningOptions(pass.planningOptions), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() override;
 
-  spatial::SpatialTargetInfo target;
+  spatial::SpatialTargetResources target;
+  ONNXToSpatialPlanningOptions planningOptions;
   bool hasTarget = false;
 };
 
@@ -56,6 +58,7 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
   SmallVector<spatial::SpatAddPlanOp> addPlans(funcOp.getOps<spatial::SpatAddPlanOp>());
   SmallVector<spatial::SpatConcatPlanOp> concatPlans(funcOp.getOps<spatial::SpatConcatPlanOp>());
+  SmallVector<spatial::SpatFlattenPlanOp> flattenPlans(funcOp.getOps<spatial::SpatFlattenPlanOp>());
   SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
   SmallVector<spatial::SpatSiluPlanOp> siluPlans(funcOp.getOps<spatial::SpatSiluPlanOp>());
   SmallVector<spatial::SpatResizeNearestPlanOp> resizePlans(
@@ -66,7 +69,7 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
   SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
   if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !addPlans.empty()
-      || !concatPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !resizePlans.empty()
+      || !concatPlans.empty() || !flattenPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !resizePlans.empty()
       || !maxPoolPlans.empty() || !blueprints.empty()
       || !globalAveragePoolPlans.empty() || !materializers.empty()) {
     return;
@@ -113,7 +116,7 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   if (!hasTarget) {
-    moduleOp.emitError("ONNX-to-Spatial lowering requires an injected SpatialTargetInfo");
+    moduleOp.emitError("ONNX-to-Spatial lowering requires an injected SpatialTargetResources");
     signalPassFailure();
     return;
   }
@@ -193,7 +196,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXSplitOp>();
 
   RewritePatternSet conversionPatterns(ctx);
-  populateConversionPatterns(conversionPatterns, ctx, this->target);
+  populateConversionPatterns(conversionPatterns, ctx, this->target, planningOptions);
   if (failed(applyPartialConversion(moduleOp, target, std::move(conversionPatterns)))) {
     moduleOp.emitError("failed to convert required ONNX ops to Spatial ops");
     signalPassFailure();
@@ -269,8 +272,10 @@ void ONNXToSpatialPass::runOnOperation() {
 
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
-std::unique_ptr<Pass> createONNXToSpatialPass(const spatial::SpatialTargetInfo& target) {
-  return std::make_unique<ONNXToSpatialPass>(target);
+std::unique_ptr<Pass> createONNXToSpatialPass(
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  return std::make_unique<ONNXToSpatialPass>(target, options);
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp

@@ -10,55 +10,83 @@
 namespace onnx_mlir {
 
 struct RowStripPhysicalValue;
+struct ONNXToSpatialPlanningOptions;
+
+inline spatial::PhysicalLayout getSpatialPlanOperandLayout(mlir::Value value) {
+  if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
+    return materialize.getTargetPhysicalLayout();
+  if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
+    return blueprint.getPhysicalLayout();
+  if (mlir::Operation* producer = value.getDefiningOp())
+    if (auto selected = spatial::getSelectedPhysicalLayout(producer))
+      return *selected;
+  return spatial::PhysicalLayout::DenseNCHW;
+}
+
+mlir::FailureOr<mlir::Value>
+lowerDenseFlattenPlan(spatial::SpatFlattenPlanOp planOp,
+                      mlir::Value input,
+                      mlir::PatternRewriter& rewriter);
 
 mlir::FailureOr<mlir::Value>
 lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
+                        mlir::Value input,
+                        mlir::Value weight,
+                        mlir::Value bias,
                         std::optional<mlir::Value> rowStripInput,
                         bool emitRowStripLayout,
-                        const spatial::SpatialTargetInfo& target,
+                        const spatial::SpatialTargetResources& target,
+                        const ONNXToSpatialPlanningOptions& options,
                         mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
-                                               const spatial::SpatialTargetInfo& target);
+                                               const spatial::SpatialTargetResources& target,
+                                               const ONNXToSpatialPlanningOptions* options = nullptr);
 mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
-                                                 const spatial::SpatialTargetInfo& target);
+                                                 const spatial::SpatialTargetResources& target,
+                                                 const ONNXToSpatialPlanningOptions* options = nullptr);
 
 mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
-  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetInfo& target);
+  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetResources& target);
 
 mlir::FailureOr<mlir::Value> lowerSelectedResizeNearestPlan(
   spatial::SpatResizeNearestPlanOp planOp,
+  mlir::Value input,
   std::optional<mlir::Value> rowStripInput,
-  const spatial::SpatialTargetInfo& target,
+  const spatial::SpatialTargetResources& target,
   mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
-                                                 const spatial::SpatialTargetInfo& target);
+                                                 const spatial::SpatialTargetResources& target);
 
 mlir::FailureOr<mlir::Value>
 lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
-                        const spatial::SpatialTargetInfo& target,
+                        mlir::Value input,
+                        const spatial::SpatialTargetResources& target,
                        mlir::PatternRewriter& rewriter);
 
 mlir::FailureOr<mlir::Value>
 lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                           mlir::Value input,
                            std::optional<mlir::Value> rowStripInput,
-                           const spatial::SpatialTargetInfo& target,
+                           const spatial::SpatialTargetResources& target,
                            mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult
 canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                        const spatial::SpatialTargetInfo& target);
+                                        const spatial::SpatialTargetResources& target);
 
 mlir::FailureOr<mlir::Value>
 lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                const spatial::SpatialTargetInfo& target,
+                                mlir::Value input,
+                                const spatial::SpatialTargetResources& target,
                                 mlir::PatternRewriter& rewriter);
 
 mlir::FailureOr<mlir::Value>
 lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                   mlir::Value input,
                                    std::optional<mlir::Value> rowStripInput,
-                                   const spatial::SpatialTargetInfo& target,
+                                   const spatial::SpatialTargetResources& target,
                                    mlir::PatternRewriter& rewriter);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialLayoutPlanningPass.cpp

@@ -0,0 +1,318 @@
+#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/Passes/Analyses/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+
+#include <algorithm>
+#include <limits>
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct SpatialLayoutSelection {
+  llvm::DenseMap<Operation*, unsigned> selectedAlternative;
+  llvm::DenseMap<Value, spatial::PhysicalLayout> resultLayouts;
+};
+
+static spatial::PhysicalLayout getKnownLayout(
+    const SpatialLayoutSelection& selection, Value value) {
+  if (auto it = selection.resultLayouts.find(value); it != selection.resultLayouts.end())
+    return it->second;
+  return getSpatialPlanOperandLayout(value);
+}
+
+static SmallVector<spatial::PhysicalLayout> getOperandLayouts(
+    Operation* op, const SpatialLayoutSelection& selection) {
+  SmallVector<spatial::PhysicalLayout> operandLayouts;
+  operandLayouts.reserve(op->getNumOperands());
+  for (Value operand : op->getOperands())
+    operandLayouts.push_back(getKnownLayout(selection, operand));
+  return operandLayouts;
+}
+
+class SpatialLayoutAnalysis {
+public:
+  SpatialLayoutAnalysis(func::FuncOp funcOp,
+                        const spatial::SpatialTargetResources& target)
+      : funcOp(funcOp), target(target) {}
+
+  FailureOr<SpatialLayoutSelection> run() {
+    SpatialLayoutSelection selection;
+    SmallVector<Operation*> planOps;
+    for (Operation& op : funcOp.getBody().front()) {
+      if (!isa<spatial::SpatialLayoutCapabilityInterface>(&op))
+        continue;
+      planOps.push_back(&op);
+      selection.resultLayouts[op.getResult(0)] = spatial::PhysicalLayout::DenseNCHW;
+      selection.selectedAlternative[&op] = 0;
+    }
+
+    const size_t maxRounds = 2 * planOps.size() + 1;
+    for (size_t round = 0; round < maxRounds; ++round) {
+      bool changed = false;
+      SmallVector<Operation*> order(planOps);
+      if (round % 2)
+        std::reverse(order.begin(), order.end());
+
+      for (Operation* op : order) {
+        FailureOr<SmallVector<spatial::LayoutAlternative>> alternatives =
+            getAlternatives(op, selection);
+        if (failed(alternatives))
+          return failure();
+
+        unsigned currentIndex = selection.selectedAlternative.lookup(op);
+        if (currentIndex >= alternatives->size())
+          currentIndex = 0;
+        if (selection.selectedAlternative.lookup(op) != currentIndex) {
+          selection.selectedAlternative[op] = currentIndex;
+          changed = true;
+        }
+        Value result = op->getResult(0);
+        if (selection.resultLayouts.lookup(result) !=
+            (*alternatives)[currentIndex].resultLayout) {
+          selection.resultLayouts[result] = (*alternatives)[currentIndex].resultLayout;
+          changed = true;
+        }
+        int64_t bestCost = alternativeCost(op, (*alternatives)[currentIndex], selection);
+        unsigned bestIndex = currentIndex;
+        for (auto [index, alternative] : llvm::enumerate(*alternatives)) {
+          int64_t cost = alternativeCost(op, alternative, selection);
+          if (cost < bestCost) {
+            bestCost = cost;
+            bestIndex = index;
+          }
+        }
+        if (bestIndex == currentIndex)
+          continue;
+        selection.selectedAlternative[op] = bestIndex;
+        selection.resultLayouts[result] = (*alternatives)[bestIndex].resultLayout;
+        changed = true;
+      }
+      if (!changed)
+        return selection;
+    }
+    funcOp.emitError("Spatial layout selection did not converge within its bounded iteration budget");
+    return failure();
+  }
+
+  FailureOr<SmallVector<spatial::LayoutAlternative>> getAlternatives(
+      Operation* op, const SpatialLayoutSelection& selection) {
+    auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
+    if (!capability)
+      return failure();
+    SmallVector<spatial::LayoutAlternative> alternatives =
+      capability.getLayoutAlternatives(target, getOperandLayouts(op, selection));
+    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();
+    }
+    if (llvm::any_of(op->getResult(0).getUses(), [](OpOperand& use) {
+          return isa<func::ReturnOp>(use.getOwner());
+        })
+        && llvm::none_of(alternatives, [](const spatial::LayoutAlternative& alternative) {
+             return alternative.resultLayout == spatial::PhysicalLayout::DenseNCHW;
+           }))
+      return op->emitOpError("does not provide the required DenseNCHW function-result layout"), failure();
+    return alternatives;
+  }
+
+private:
+  int64_t alternativeCost(Operation* op,
+                          const spatial::LayoutAlternative& alternative,
+                          const SpatialLayoutSelection& selection) {
+    if (llvm::any_of(op->getResult(0).getUses(), [](OpOperand& use) {
+          return isa<func::ReturnOp>(use.getOwner());
+        })
+        && alternative.resultLayout != spatial::PhysicalLayout::DenseNCHW)
+      return std::numeric_limits<int64_t>::max() / 4;
+
+    int64_t cost = alternative.intrinsicCost;
+    SmallVector<spatial::PhysicalLayout> operandLayouts = getOperandLayouts(op, selection);
+    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)
+        continue;
+      SmallVector<spatial::PhysicalLayout> userOperandLayouts =
+          getOperandLayouts(use.getOwner(), selection);
+      for (auto [index, operand] : llvm::enumerate(use.getOwner()->getOperands()))
+        if (operand == result)
+          userOperandLayouts[index] = alternative.resultLayout;
+      SmallVector<spatial::LayoutAlternative> userAlternatives =
+          user.getLayoutAlternatives(target, userOperandLayouts);
+      if (llvm::none_of(userAlternatives,
+            [&](const spatial::LayoutAlternative& userAlternative) {
+              return userAlternative.operandLayouts.size()
+                       == use.getOwner()->getNumOperands()
+                  && userAlternative.operandLayouts[use.getOperandNumber()]
+                       == alternative.resultLayout;
+            }))
+        ++cost;
+    }
+    return cost;
+  }
+
+  func::FuncOp funcOp;
+  const spatial::SpatialTargetResources& target;
+};
+
+static LogicalResult materializeMismatchedUses(
+    IRRewriter& rewriter, const SpatialLayoutSelection& selection,
+    Operation* op, SpatialLayoutAnalysis& analysis) {
+  Value value = op->getResult(0);
+  spatial::PhysicalLayout sourceLayout = getKnownLayout(selection, value);
+  SmallVector<std::pair<OpOperand*, spatial::PhysicalLayout>> mismatches;
+  for (OpOperand& use : value.getUses()) {
+    Operation* userOp = use.getOwner();
+    auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(userOp);
+    if (!capability) {
+      if (isa<func::ReturnOp>(userOp) || sourceLayout == spatial::PhysicalLayout::DenseNCHW)
+        continue;
+      mismatches.push_back({&use, spatial::PhysicalLayout::DenseNCHW});
+      continue;
+    }
+    FailureOr<SmallVector<spatial::LayoutAlternative>> alternatives =
+        analysis.getAlternatives(userOp, selection);
+    if (failed(alternatives))
+      return failure();
+    unsigned selectedIndex = selection.selectedAlternative.lookup(userOp);
+    if (selectedIndex >= alternatives->size())
+      return userOp->emitOpError()
+                 << "has no selected Spatial layout alternative (index " << selectedIndex
+                                   << ", alternatives " << alternatives->size() << ")",
+             failure();
+    spatial::PhysicalLayout required =
+        (*alternatives)[selectedIndex].operandLayouts[use.getOperandNumber()];
+    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(
+    const SpatialLayoutSelection& selection,
+    ArrayRef<Operation*> planOps,
+    SpatialLayoutAnalysis& analysis) {
+  for (Operation* op : planOps) {
+    auto selected = spatial::getSelectedPhysicalLayout(op);
+    if (!selected)
+      return op->emitOpError("requires a selected physical layout"), failure();
+    FailureOr<SmallVector<spatial::LayoutAlternative>> alternatives =
+        analysis.getAlternatives(op, selection);
+    if (failed(alternatives))
+      return failure();
+    unsigned selectedIndex = selection.selectedAlternative.lookup(op);
+    if (selectedIndex >= alternatives->size())
+      return op->emitOpError()
+                 << "has no selected Spatial layout alternative (index " << selectedIndex
+                               << ", alternatives " << alternatives->size() << ")",
+             failure();
+    const spatial::LayoutAlternative& alternative = (*alternatives)[selectedIndex];
+    if (alternative.resultLayout != *selected
+        || getOperandLayouts(op, selection) != alternative.operandLayouts)
+      return op->emitOpError("selected physical layout does not satisfy its exact 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::SpatialTargetResources& target)
+  : target(target), hasTarget(true) {}
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget) {
+      moduleOp.emitError("Spatial layout planning requires an injected SpatialTargetResources");
+      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;
+    SpatialLayoutAnalysis analysis(funcOp, target);
+    FailureOr<SpatialLayoutSelection> selection = analysis.run();
+    if (failed(selection)) {
+      signalPassFailure();
+      return;
+    }
+
+    SmallVector<Operation*> planOps;
+    for (Operation& op : funcOp.getBody().front())
+      if (isa<spatial::SpatialLayoutCapabilityInterface>(&op))
+        planOps.push_back(&op);
+
+    IRRewriter rewriter(&getContext());
+    for (Operation* op : planOps) {
+      op->setAttr(spatial::kSelectedLayoutAttrName,
+                  spatial::PhysicalLayoutAttr::get(
+              rewriter.getContext(), selection->resultLayouts.lookup(op->getResult(0))));
+      if (failed(materializeMismatchedUses(
+              rewriter, *selection, op, analysis))) {
+        signalPassFailure();
+        return;
+      }
+    }
+    if (failed(verifySelectedLayouts(*selection, planOps, analysis))
+        || failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("Spatial layout planning verification failed");
+      signalPassFailure();
+      return;
+    }
+  }
+
+  spatial::SpatialTargetResources target;
+  bool hasTarget = false;
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass() {
+  return std::make_unique<SpatialLayoutPlanningPass>();
+}
+
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass(
+    const spatial::SpatialTargetResources& target) {
+  return std::make_unique<SpatialLayoutPlanningPass>(target);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.cpp

@@ -4,23 +4,21 @@
 #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 "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #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/ONNXToSpatialOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 
 using namespace mlir;
 
@@ -59,19 +57,11 @@ static bool isDenseSelected(Operation* op) {
 }
 
 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;
+  return getSpatialPlanOperandLayout(value);
 }
 
 static LogicalResult verifySelectedLayouts(
-    func::FuncOp funcOp, const spatial::SpatialTargetInfo& target) {
+    func::FuncOp funcOp, const spatial::SpatialTargetResources& target) {
   LogicalResult result = success();
   funcOp.walk([&](Operation* op) {
     auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
@@ -115,12 +105,6 @@ lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp pl
   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,
@@ -222,17 +206,18 @@ static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTe
   return batch->getResult(0);
 }
 
-struct LowerDenseReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseReluPlan final : OpConversionPattern<spatial::SpatReluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatReluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& 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) {
+      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, adaptor.getInput(), [&](Value x) {
         auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
         spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
       });
@@ -241,17 +226,18 @@ struct LowerDenseReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
   }
 };
 
-struct LowerDenseSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseSiluPlan final : OpConversionPattern<spatial::SpatSiluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatSiluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& 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) {
+      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, adaptor.getInput(), [&](Value x) {
         Value sigmoid = spatial::SpatSigmoidOp::create(
           rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
         Value silu = spatial::SpatVMulOp::create(
@@ -263,29 +249,32 @@ struct LowerDenseSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
   }
 };
 
-struct LowerDenseResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
-  explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+struct LowerDenseResizePlan final : OpConversionPattern<spatial::SpatResizeNearestPlanOp> {
+  explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatResizeNearestPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
-    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(planOp, std::nullopt, target, rewriter);
+    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(
+        planOp, adaptor.getInput(), 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;
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseBiasAddPlan final : OpConversionPattern<spatial::SpatBiasAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatBiasAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
     auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
@@ -293,12 +282,12 @@ struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp
       return planOp.emitOpError("requires ranked output type");
 
     FailureOr<Value> denseBias = materializeDenseBiasAddTensor(
-      planOp.getBias(), resultType, rewriter, planOp.getLoc());
+      adaptor.getBias(), resultType, rewriter, planOp.getLoc());
     if (failed(denseBias))
       return planOp.emitOpError("failed to materialize dense Conv-style bias");
-    if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+    if (adaptor.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
       FailureOr<Value> lowered = lowerDenseBatchBiasAdd(
-        planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+        adaptor.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
       if (succeeded(lowered)) {
         rewriter.replaceOp(planOp, *lowered);
         return success();
@@ -309,7 +298,7 @@ struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp
       planOp.getLoc(),
       planOp.getOutput().getType(),
       {},
-      ValueRange {planOp.getInput(), *denseBias},
+      ValueRange {adaptor.getInput(), *denseBias},
       [&](Value x, Value y) {
         auto added = spatial::SpatVAddOp::create(
           rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
@@ -320,11 +309,12 @@ struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp
   }
 };
 
-struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseAddPlan final : OpConversionPattern<spatial::SpatAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
     auto compute = createSpatCompute<2>(
@@ -332,7 +322,7 @@ struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
       planOp.getLoc(),
       planOp.getOutput().getType(),
       {},
-      ValueRange {planOp.getLhs(), planOp.getRhs()},
+      ValueRange {adaptor.getLhs(), adaptor.getRhs()},
       [&](Value lhsValue, Value rhsValue) {
         Value added = spatial::SpatVAddOp::create(
           rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
@@ -343,11 +333,12 @@ struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
   }
 };
 
-struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerDenseConcatPlan final : OpConversionPattern<spatial::SpatConcatPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatConcatPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
     auto compute = createSpatCompute(
@@ -355,7 +346,7 @@ struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp>
       planOp.getLoc(),
       TypeRange {planOp.getOutput().getType()},
       {},
-      planOp.getInputs(),
+      adaptor.getInputs(),
       [&](ValueRange values) {
         Value concatenated = spatial::SpatConcatOp::create(
           rewriter,
@@ -371,16 +362,17 @@ struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp>
 };
 
 static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
+                                  Value lhsValue, Value rhsValue,
                                   PatternRewriter& rewriter) {
-  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(planOp.getLhs());
-  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(planOp.getRhs());
+  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(lhsValue);
+  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rhsValue);
   if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) {
-    if (getKnownPhysicalLayout(planOp.getLhs()) == spatial::PhysicalLayout::NHWCRowStrip)
+    if (getKnownPhysicalLayout(lhsValue) == 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)
+    if (getKnownPhysicalLayout(rhsValue) == spatial::PhysicalLayout::NHWCRowStrip)
       return failure();
     return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
   }
@@ -397,9 +389,10 @@ static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
 }
 
 static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
+                                     ValueRange inputValues,
                                      PatternRewriter& rewriter) {
   SmallVector<RowStripPhysicalValue> inputs;
-  for (Value input : planOp.getInputs()) {
+  for (Value input : inputValues) {
     FailureOr<RowStripPhysicalValue> physical = getRowStripValue(input);
     if (failed(physical)) {
       inputs.clear();
@@ -407,8 +400,8 @@ static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
     }
     inputs.push_back(*physical);
   }
-  if (isRowStripSelected(planOp.getOperation()) && inputs.size() != planOp.getInputs().size()) {
-    if (llvm::any_of(planOp.getInputs(), [](Value input) {
+  if (isRowStripSelected(planOp.getOperation()) && inputs.size() != inputValues.size()) {
+    if (llvm::any_of(inputValues, [](Value input) {
           return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip;
         }))
       return failure();
@@ -426,15 +419,19 @@ static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
   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) {}
+struct LowerSelectedConvPlan final : OpConversionPattern<spatial::SpatConv2DPlanOp> {
+  explicit LowerSelectedConvPlan(MLIRContext* ctx,
+                                 const spatial::SpatialTargetResources& target,
+                                 const ONNXToSpatialPlanningOptions& options)
+      : OpConversionPattern<spatial::SpatConv2DPlanOp>(ctx), target(target), options(options) {}
 
   LogicalResult matchAndRewrite(spatial::SpatConv2DPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatConv2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (isDenseSelected(planOp.getOperation())) {
       FailureOr<Value> lowered = lowerSelectedConv2DPlan(
-        planOp, std::nullopt, /*emitRowStripLayout=*/false, target, rewriter);
+          planOp, adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(),
+          std::nullopt, /*emitRowStripLayout=*/false, target, options, rewriter);
       if (failed(lowered))
         return planOp.emitOpError("failed to lower selected dense Spatial Conv plan");
       rewriter.replaceOp(planOp, *lowered);
@@ -443,15 +440,16 @@ struct LowerSelectedConvPlan final : OpRewritePattern<spatial::SpatConv2DPlanOp>
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
 
-    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(adaptor.getInput());
     if (failed(rowStripInput)
-        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        && getKnownPhysicalLayout(adaptor.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);
+        planOp, adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(),
+        physicalInput, /*emitRowStripLayout=*/true, target, options, rewriter);
     if (failed(lowered))
       return planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
     if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
@@ -459,19 +457,21 @@ struct LowerSelectedConvPlan final : OpRewritePattern<spatial::SpatConv2DPlanOp>
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
+  const ONNXToSpatialPlanningOptions& options;
 };
 
-struct LowerRowStripReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripReluPlan final : OpConversionPattern<spatial::SpatReluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatReluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)) {
-      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
         return failure();
       return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input");
     }
@@ -484,16 +484,17 @@ struct LowerRowStripReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
   }
 };
 
-struct LowerRowStripSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripSiluPlan final : OpConversionPattern<spatial::SpatSiluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatSiluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)) {
-      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
         return failure();
       return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input");
     }
@@ -506,21 +507,23 @@ struct LowerRowStripSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
   }
 };
 
-struct LowerRowStripResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
-  explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+struct LowerRowStripResizePlan final : OpConversionPattern<spatial::SpatResizeNearestPlanOp> {
+  explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatResizeNearestPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)) {
-      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      if (getKnownPhysicalLayout(adaptor.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);
+    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(
+        planOp, adaptor.getInput(), input->storage, target, rewriter);
     if (failed(lowered))
       return planOp.emitOpError("failed to lower selected row-strip Resize plan");
     if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
@@ -528,43 +531,47 @@ struct LowerRowStripResizePlan final : OpRewritePattern<spatial::SpatResizeNeare
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerDenseMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
-  explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
+struct LowerDenseMaxPoolPlan final : OpConversionPattern<spatial::SpatMaxPool2DPlanOp> {
+  explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatMaxPool2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
-    FailureOr<Value> lowered = lowerDenseMaxPool2DPlan(planOp, target, rewriter);
+    FailureOr<Value> lowered = lowerDenseMaxPool2DPlan(
+        planOp, adaptor.getInput(), 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;
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerRowStripMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
-  explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
+struct LowerRowStripMaxPoolPlan final : OpConversionPattern<spatial::SpatMaxPool2DPlanOp> {
+  explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatMaxPool2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)
-        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
       return failure();
     std::optional<Value> physicalInput;
     if (succeeded(input))
       physicalInput = input->storage;
-    FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(planOp, physicalInput, target, rewriter);
+    FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
+        planOp, adaptor.getInput(), physicalInput, target, rewriter);
     if (failed(lowered))
       return planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
     if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
@@ -572,26 +579,28 @@ struct LowerRowStripMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DP
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
 struct LowerRowStripGlobalAveragePoolPlan
-    final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
-  explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+    final : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp> {
+  explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)
-        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
       return failure();
     std::optional<Value> physicalInput;
     if (succeeded(input))
       physicalInput = input->storage;
-    FailureOr<Value> lowered = lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, target, rewriter);
+    FailureOr<Value> lowered = lowerSelectedGlobalAveragePoolPlan(
+        planOp, adaptor.getInput(), 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)))
@@ -599,43 +608,47 @@ struct LowerRowStripGlobalAveragePoolPlan
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
 struct LowerDenseGlobalAveragePoolPlan
-    final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
+    final : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp> {
   explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx,
-                                           const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+                                           const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isDenseSelected(planOp.getOperation()))
       return failure();
-    FailureOr<Value> lowered = lowerDenseGlobalAveragePoolPlan(planOp, target, rewriter);
+    FailureOr<Value> lowered = lowerDenseGlobalAveragePoolPlan(
+        planOp, adaptor.getInput(), 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;
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerRowStripBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripBiasAddPlan final : OpConversionPattern<spatial::SpatBiasAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatBiasAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
     if (failed(input)) {
-      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      if (getKnownPhysicalLayout(adaptor.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);
+    FailureOr<Value> lowered = applyRowStripBiasAdd(
+        *input, adaptor.getBias(), rewriter, planOp.getLoc());
     if (failed(lowered))
       return planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
     if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
@@ -644,48 +657,51 @@ struct LowerRowStripBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPla
   }
 };
 
-struct LowerRowStripAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripAddPlan final : OpConversionPattern<spatial::SpatAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    return lowerAddPlan(planOp, rewriter);
+    return lowerAddPlan(planOp, adaptor.getLhs(), adaptor.getRhs(), rewriter);
   }
 };
 
-struct LowerRowStripConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct LowerRowStripConcatPlan final : OpConversionPattern<spatial::SpatConcatPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatConcatPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
     if (!isRowStripSelected(planOp.getOperation()))
       return failure();
-    return lowerConcatPlan(planOp, rewriter);
+    return lowerConcatPlan(planOp, adaptor.getInputs(), rewriter);
   }
 };
 
 struct LowerMaterializeLayout final
-    : OpRewritePattern<spatial::SpatMaterializeLayoutOp> {
-  using OpRewritePattern::OpRewritePattern;
+    : OpConversionPattern<spatial::SpatMaterializeLayoutOp> {
+  using OpConversionPattern::OpConversionPattern;
 
   LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp,
-                                PatternRewriter& rewriter) const override {
+                                spatial::SpatMaterializeLayoutOpAdaptor adaptor,
+                                ConversionPatternRewriter& 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());
+      rewriter.replaceOp(materializeOp, adaptor.getInput());
       return success();
     }
     if (source == spatial::PhysicalLayout::DenseNCHW
         && target == spatial::PhysicalLayout::NHWCRowStrip) {
-      auto logicalType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
+      auto logicalType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
       if (!logicalType)
         return materializeOp.emitOpError("requires a ranked dense input"), failure();
       FailureOr<Value> rowStrip = materializeDenseToRowStrip(
-        materializeOp.getInput(), logicalType, materializeOp.getLoc(), rewriter);
+        adaptor.getInput(), logicalType, materializeOp.getLoc(), rewriter);
       if (failed(rowStrip))
         return materializeOp.emitOpError(
           "failed to materialize dense NCHW storage to row-strip layout"), failure();
@@ -696,14 +712,13 @@ struct LowerMaterializeLayout final
         || target != spatial::PhysicalLayout::DenseNCHW)
       return materializeOp.emitOpError(
           "unsupported Spatial layout materialization direction"), failure();
-    auto inputType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
+    auto inputType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
     if (!inputType)
       return materializeOp.emitOpError("requires a ranked row-strip input"), failure();
     FailureOr<RowStripPhysicalValue> rowStripValue =
-        getRowStripValue(materializeOp.getInput());
+        getRowStripValue(adaptor.getInput());
     if (failed(rowStripValue))
-      return materializeOp.emitOpError(
-          "requires an explicitly defining row-strip physical value"), failure();
+      return failure();
     FailureOr<Value> dense = materializeRowStripToDense(
         *rowStripValue, materializeOp.getLoc(), rewriter);
     if (failed(dense))
@@ -714,189 +729,117 @@ struct LowerMaterializeLayout final
   }
 };
 
-struct LowerRowStripFlatten final
-    : OpRewritePattern<spatial::SpatGraphCompute> {
-  explicit LowerRowStripFlatten(MLIRContext* context,
-                                const spatial::SpatialTargetInfo& target)
-      : OpRewritePattern<spatial::SpatGraphCompute>(context), target(target) {}
+struct LowerSelectedFlattenPlan final
+    : OpConversionPattern<spatial::SpatFlattenPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatGraphCompute flattenOp,
-                                PatternRewriter& rewriter) const override {
-    if (flattenOp.getInputs().size() != 1)
+  LogicalResult matchAndRewrite(spatial::SpatFlattenPlanOp planOp,
+                                spatial::SpatFlattenPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
       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();
+    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(adaptor.getInput());
+    if (succeeded(rowStripInput)) {
+      if (failed(canLowerFlattenFromRowStrip(planOp, target))
+          || failed(lowerFlattenFromRowStrip(*rowStripInput, planOp, target, rewriter)))
+        return planOp.emitOpError("failed to lower selected Spatial Flatten plan"), failure();
+      return success();
+    }
+    FailureOr<Value> lowered = lowerDenseFlattenPlan(planOp, adaptor.getInput(), rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense Spatial Flatten plan"), failure();
+    rewriter.replaceOp(planOp, *lowered);
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  explicit LowerSelectedFlattenPlan(MLIRContext* context,
+                                    const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatFlattenPlanOp>(context), target(target) {}
+
+  const spatial::SpatialTargetResources& target;
 };
 
-struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
+struct EraseDeadPhysicalViewBlueprint final
+    : OpRewritePattern<spatial::SpatBlueprintOp> {
+  using OpRewritePattern::OpRewritePattern;
 
-  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;
+  LogicalResult matchAndRewrite(spatial::SpatBlueprintOp blueprint,
+                                PatternRewriter& rewriter) const override {
+    if (!spatial::isPhysicalView(blueprint.getMode()) || !blueprint.use_empty())
+      return failure();
+    rewriter.eraseOp(blueprint);
+    return success();
   }
-    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;
+static void populateConvPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  patterns.add<LowerSelectedConvPlan>(ctx, target, options);
 }
 
-    RewritePatternSet selectedPlanPatterns(ctx);
-    selectedPlanPatterns.add<LowerDenseReluPlan,
+static void populateElementwisePlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerDenseReluPlan,
                LowerRowStripReluPlan,
                LowerDenseSiluPlan,
                LowerRowStripSiluPlan,
                LowerDenseBiasAddPlan,
                LowerRowStripBiasAddPlan,
                LowerDenseAddPlan,
-                             LowerRowStripAddPlan,
-                             LowerDenseConcatPlan,
-                             LowerRowStripConcatPlan>(ctx);
-    selectedPlanPatterns.add<LowerSelectedConvPlan,
-                             LowerDenseResizePlan,
-                             LowerRowStripResizePlan,
-                             LowerDenseMaxPoolPlan,
+               LowerRowStripAddPlan>(ctx);
+}
+
+static void populatePoolPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<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;
+static void populateResizePlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<LowerDenseResizePlan, LowerRowStripResizePlan>(ctx, target);
 }
 
-    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;
-      }
+static void populateConcatPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerDenseConcatPlan, LowerRowStripConcatPlan>(ctx);
 }
 
-    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
-      return;
+static void populateFlattenPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<LowerSelectedFlattenPlan>(ctx, target);
 }
 
-  spatial::SpatialTargetInfo target;
-  bool hasTarget = false;
-};
+static void populateLayoutMaterializationPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerMaterializeLayout, EraseDeadPhysicalViewBlueprint>(ctx);
+}
 
 } // namespace
 
-std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }
+void populateSpatialPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  populateConvPlanLoweringPatterns(patterns, ctx, target, options);
+  populateElementwisePlanLoweringPatterns(patterns, ctx);
+  populatePoolPlanLoweringPatterns(patterns, ctx, target);
+  populateResizePlanLoweringPatterns(patterns, ctx, target);
+  populateConcatPlanLoweringPatterns(patterns, ctx);
+  populateFlattenPlanLoweringPatterns(patterns, ctx, target);
+  populateLayoutMaterializationPatterns(patterns, ctx);
+}
 
-std::unique_ptr<Pass> createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target) {
-  return std::make_unique<LowerSpatialPlansPass>(target);
+LogicalResult verifySelectedSpatialLayouts(
+    func::FuncOp funcOp, const spatial::SpatialTargetResources& target) {
+  return verifySelectedLayouts(funcOp, target);
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Passes/Transforms/SpatialPlanLoweringPatterns.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+
+namespace onnx_mlir {
+
+void populateSpatialPlanLoweringPatterns(
+    mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options);
+
+mlir::LogicalResult verifySelectedSpatialLayouts(
+    mlir::func::FuncOp funcOp, const spatial::SpatialTargetResources& target);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -9,11 +9,12 @@ void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { popula
 
 void populateConversionPatterns(RewritePatternSet& patterns,
                                 MLIRContext* ctx,
-                                const spatial::SpatialTargetInfo& target) {
+                                const spatial::SpatialTargetResources& target,
+                                const ONNXToSpatialPlanningOptions& options) {
   populateElementwisePatterns(patterns, ctx);
   populateMatMulRewritePatterns(patterns, ctx, target);
   populateGemmPatterns(patterns, ctx, target);
-  populateConvPatterns(patterns, ctx, target);
+  populateConvPatterns(patterns, ctx, target, options);
   populatePoolPatterns(patterns, ctx, target);
   populateReduceMeanPatterns(patterns, ctx);
   populateReluPatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -4,18 +4,20 @@
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 namespace onnx_mlir {
 
 namespace spatial {
-struct SpatialTargetInfo;
+struct SpatialTargetResources;
 }
 
 void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateConversionPatterns(mlir::RewritePatternSet& patterns,
                                 mlir::MLIRContext* ctx,
-                                const spatial::SpatialTargetInfo& target);
+                                const spatial::SpatialTargetResources& target,
+                                const ONNXToSpatialPlanningOptions& options);
 void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
@@ -23,21 +25,22 @@ void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::ML
 
 void populateConvPatterns(mlir::RewritePatternSet& patterns,
                           mlir::MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target);
+                          const spatial::SpatialTargetResources& target,
+                          const ONNXToSpatialPlanningOptions& options);
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateElementwiseFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateGemmPatterns(mlir::RewritePatternSet& patterns,
                           mlir::MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target);
+                          const spatial::SpatialTargetResources& target);
 void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns,
                                    mlir::MLIRContext* ctx,
-                                   const spatial::SpatialTargetInfo& target);
+                                   const spatial::SpatialTargetResources& target);
 void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns,
                                   mlir::MLIRContext* ctx,
-                                  const spatial::SpatialTargetInfo& target);
+                                  const spatial::SpatialTargetResources& target);
 void populatePoolPatterns(mlir::RewritePatternSet& patterns,
                           mlir::MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target);
+                          const spatial::SpatialTargetResources& target);
 void populateReduceMeanPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSigmoidPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

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

@@ -1,6 +1,7 @@
 #include "ConvGeometry.hpp"
 
 #include <algorithm>
+#include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 
@@ -8,12 +9,22 @@ namespace onnx_mlir {
 
 namespace {
 
-static int64_t ceilDivide(int64_t value, int64_t divisor) {
-  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
+static const ONNXToSpatialPlanningOptions& defaultPlanningOptions() {
+  static const ONNXToSpatialPlanningOptions options {
+    std::numeric_limits<uint64_t>::max(),
+    std::numeric_limits<uint64_t>::max(),
+    spatial::ConvLoweringStrategy::Auto,
+    false,
+  };
+  return options;
 }
 
 } // namespace
 
+const ONNXToSpatialPlanningOptions& ConvLoweringState::planningOptions() const {
+  return options ? *options : defaultPlanningOptions();
+}
+
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup) {
   return group == numChannelsIn && numChannelsInPerGroup == 1 && numChannelsOut % group == 0;
 }
@@ -31,20 +42,8 @@ void classifyConvProblem(ConvProblem& problem) {
 }
 
 ConvGeometry buildConvGeometry(const ConvProblem& problem,
-                               const spatial::SpatialTargetInfo& target) {
+                               const spatial::SpatialTargetResources& target) {
   ConvGeometry geo {
-    problem.batchSize,
-    problem.numChannelsIn,
-    problem.xHeight,
-    problem.xWidth,
-    problem.numChannelsOut,
-    problem.wHeight,
-    problem.wWidth,
-    problem.outHeight,
-    problem.outWidth,
-    problem.group,
-    problem.numChannelsInPerGroup,
-    problem.numChannelsOutPerGroup,
     problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth,
     problem.numChannelsOutPerGroup,
     problem.batchSize * problem.outHeight * problem.outWidth,
@@ -52,11 +51,6 @@ ConvGeometry buildConvGeometry(const ConvProblem& problem,
     static_cast<int64_t>(target.matrixUnitsPerProcessor),
     1,
     0,
-    problem.hasBias,
-    isDepthwiseConv(problem.group,
-                    problem.numChannelsIn,
-                    problem.numChannelsOut,
-                    problem.numChannelsInPerGroup),
   };
   geo.pack = std::max<int64_t>(1, geo.xbarSize / std::max<int64_t>(geo.k, geo.c));
   geo.im2colElements = static_cast<uint64_t>(std::max<int64_t>(0, geo.p)) * static_cast<uint64_t>(std::max<int64_t>(0, geo.k));
@@ -64,216 +58,76 @@ ConvGeometry buildConvGeometry(const ConvProblem& problem,
 }
 
 static ConvMaterializationKind getMaterializationKind(
-    const ConvProblem& problem, spatial::ConvLoweringStrategy strategy) {
-  if (strategy == spatial::ConvLoweringStrategy::Depthwise)
-    return ConvMaterializationKind::StructuredDepthwise;
-  if (problem.isPointwise)
-    return ConvMaterializationKind::PointwiseContraction;
+    spatial::ConvLoweringStrategy strategy) {
   switch (strategy) {
   case spatial::ConvLoweringStrategy::Depthwise:
     return ConvMaterializationKind::StructuredDepthwise;
-  case spatial::ConvLoweringStrategy::PackedIm2Col:
   case spatial::ConvLoweringStrategy::Legacy:
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
     return ConvMaterializationKind::PackedIm2Col;
   case spatial::ConvLoweringStrategy::StreamedPatch:
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+  case spatial::ConvLoweringStrategy::Tiled2D:
     return ConvMaterializationKind::StreamedPatch;
   case spatial::ConvLoweringStrategy::StreamedPacked:
     return ConvMaterializationKind::StreamedPacked;
-  case spatial::ConvLoweringStrategy::OutputChannelTiled:
-    return ConvMaterializationKind::OutputChannelTiled;
   case spatial::ConvLoweringStrategy::InputKTiled:
     return ConvMaterializationKind::InputKTiled;
-  case spatial::ConvLoweringStrategy::Tiled2D:
-    return ConvMaterializationKind::Tiled2D;
   case spatial::ConvLoweringStrategy::Auto:
     break;
   }
   llvm_unreachable("auto is not a Conv materialization kind");
 }
 
-static ConvPlan makeCandidatePlan(const ConvProblem& problem,
-                                  spatial::ConvLoweringStrategy strategy,
-                                  const spatial::SpatialTargetInfo& target) {
-  ConvPlan plan;
-  plan.geometry = buildConvGeometry(problem, target);
-  plan.strategy = strategy;
-  plan.materializationKind = getMaterializationKind(problem, strategy);
-  plan.laneCount = plan.geometry.p;
-  plan.reductionCount = std::max<int64_t>(
-    1, (plan.geometry.k + plan.geometry.xbarSize - 1) / plan.geometry.xbarSize);
-  plan.mvmCount = plan.laneCount * plan.reductionCount;
-  plan.vectorCount = plan.mvmCount;
-  plan.weightElements = static_cast<uint64_t>(std::max<int64_t>(0, problem.numChannelsOut))
-    * static_cast<uint64_t>(std::max<int64_t>(0, plan.geometry.k));
-  plan.scratchElements = plan.geometry.im2colElements;
-  plan.materializationElements = strategy == spatial::ConvLoweringStrategy::Depthwise
-    ? 0
-    : std::min<uint64_t>(plan.geometry.im2colElements, target.convIm2colMaxElements);
-  plan.requiresInputMaterialization = strategy != spatial::ConvLoweringStrategy::Depthwise;
-  plan.producesRowStrip = strategy != spatial::ConvLoweringStrategy::InputKTiled
-    && ceilDivide(plan.geometry.k, plan.geometry.xbarSize) <= plan.geometry.matrixUnitsPerProcessor;
-  plan.consumesRowStrip = plan.producesRowStrip;
-  // Conv materializers emit local compute and leave inter-core communication
-  // to Spatial scheduling; zero is an explicit ownership statement here.
-  plan.communicationElements = 0;
-  plan.usesContraction = problem.isPointwise || strategy != spatial::ConvLoweringStrategy::Depthwise;
-  if (problem.isPointwise) {
-    ContractionProblem contraction;
-    contraction.origin = ContractionOrigin::Gemm;
-    contraction.batch = 1;
-    contraction.m = plan.geometry.p;
-    contraction.k = plan.geometry.c;
-    contraction.n = problem.numChannelsOutPerGroup;
-    contraction.lhsElementType = problem.xType.getElementType();
-    contraction.rhsElementType = problem.wType.getElementType();
-    contraction.resultElementType = problem.outType.getElementType();
-    plan.contraction = makeContractionPlan(
-      contraction, target, ContractionPlanKind::StaticTiled);
-    plan.hasContractionPlan = true;
-    plan.laneCount = plan.contraction.laneCount;
-    plan.mvmCount = plan.contraction.expectedMvmCount;
-    plan.vectorCount = plan.contraction.expectedVectorCount;
-    plan.reductionCount = plan.contraction.reductionSlices;
-  }
-  return plan;
-}
-
 static bool fitsSingleCrossbar(const ConvGeometry& geo) {
   return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize;
 }
 
 static bool fitsPackedIm2Col(const ConvGeometry& geo,
-                             const spatial::SpatialTargetInfo& target) {
+                             const ONNXToSpatialPlanningOptions& options) {
   return fitsSingleCrossbar(geo) && geo.pack >= 2
-    && geo.im2colElements <= target.convIm2colMaxElements;
-}
-
-static mlir::FailureOr<ConvPlan> buildDepthwiseCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  if (!problem.isDepthwise)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Depthwise, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildPackedIm2ColCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (!fitsPackedIm2Col(geo, target))
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::PackedIm2Col, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildStreamedPatchCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  if (!fitsSingleCrossbar(buildConvGeometry(problem, target)))
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::StreamedPatch, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildStreamedPackedCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (!fitsSingleCrossbar(geo) || geo.pack < 2)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::StreamedPacked, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildOutputChannelTiledCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (geo.k > geo.xbarSize || geo.c <= geo.xbarSize)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::OutputChannelTiled, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildInputKTiledCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (geo.k <= geo.xbarSize || geo.c > geo.xbarSize)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::InputKTiled, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildTiled2DCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  if (geo.k <= geo.xbarSize || geo.c <= geo.xbarSize)
-    return mlir::failure();
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Tiled2D, target);
-}
-
-static mlir::FailureOr<ConvPlan> buildLegacyCandidate(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  // Legacy is retained as the explicit compatibility/debug materializer and
-  // as the safe fallback when structured depthwise lowering is unavailable.
-  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Legacy, target);
+    && geo.im2colElements <= options.convIm2colMaxElements;
 }
 
 mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
                                        spatial::ConvLoweringStrategy strategy,
-                                       const spatial::SpatialTargetInfo& target) {
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  auto plan = [&]() { return ConvPlan {getMaterializationKind(strategy)}; };
+  auto ifApplicable = [&](bool applicable) -> mlir::FailureOr<ConvPlan> {
+    return applicable ? mlir::FailureOr<ConvPlan>(plan()) : mlir::FailureOr<ConvPlan>(mlir::failure());
+  };
   switch (strategy) {
   case spatial::ConvLoweringStrategy::Auto:
     return mlir::failure();
   case spatial::ConvLoweringStrategy::Legacy:
-    return buildLegacyCandidate(problem, target);
+    return plan();
   case spatial::ConvLoweringStrategy::Depthwise:
-    return buildDepthwiseCandidate(problem, target);
+    return ifApplicable(problem.isDepthwise);
   case spatial::ConvLoweringStrategy::PackedIm2Col:
-    return buildPackedIm2ColCandidate(problem, target);
+    return ifApplicable(fitsPackedIm2Col(geo, options));
   case spatial::ConvLoweringStrategy::StreamedPatch:
-    return buildStreamedPatchCandidate(problem, target);
+    return ifApplicable(fitsSingleCrossbar(geo));
   case spatial::ConvLoweringStrategy::StreamedPacked:
-    return buildStreamedPackedCandidate(problem, target);
+    return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2);
   case spatial::ConvLoweringStrategy::OutputChannelTiled:
-    return buildOutputChannelTiledCandidate(problem, target);
+    return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize);
   case spatial::ConvLoweringStrategy::InputKTiled:
-    return buildInputKTiledCandidate(problem, target);
+    return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize);
   case spatial::ConvLoweringStrategy::Tiled2D:
-    return buildTiled2DCandidate(problem, target);
+    return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize);
   }
   llvm_unreachable("unknown Conv lowering strategy");
 }
 
-llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
-  ConvGeometry geo = buildConvGeometry(problem, target);
-  llvm::SmallVector<ConvPlan, 8> candidates;
-  auto append = [&](spatial::ConvLoweringStrategy strategy) {
-    mlir::FailureOr<ConvPlan> candidate = makeConvPlan(problem, strategy, target);
-    if (succeeded(candidate))
-      candidates.push_back(*candidate);
-  };
-
-  if (problem.isDepthwise) {
-    append(spatial::ConvLoweringStrategy::Depthwise);
-    append(spatial::ConvLoweringStrategy::Legacy);
-    return candidates;
-  }
-  if (fitsPackedIm2Col(geo, target))
-    append(spatial::ConvLoweringStrategy::PackedIm2Col);
-  if (fitsSingleCrossbar(geo) && geo.pack >= 2)
-    append(spatial::ConvLoweringStrategy::StreamedPacked);
-  if (fitsSingleCrossbar(geo))
-    append(spatial::ConvLoweringStrategy::StreamedPatch);
-  if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize)
-    append(spatial::ConvLoweringStrategy::OutputChannelTiled);
-  if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize)
-    append(spatial::ConvLoweringStrategy::Legacy);
-  if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize)
-    append(spatial::ConvLoweringStrategy::InputKTiled);
-  if (geo.k > geo.xbarSize && geo.c > geo.xbarSize)
-    append(spatial::ConvLoweringStrategy::Tiled2D);
-  return candidates;
-}
-
 uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
                                     int64_t packFactor,
-                                    const spatial::SpatialTargetInfo& target) {
+                                    const ONNXToSpatialPlanningOptions& options) {
   const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
-  uint64_t chunkPositions = std::max<uint64_t>(1, target.convIm2colMaxElements / patchElements);
+  uint64_t chunkPositions = std::max<uint64_t>(1, options.convIm2colMaxElements / patchElements);
   chunkPositions = std::min<uint64_t>(chunkPositions, static_cast<uint64_t>(std::max<int64_t>(1, geo.p)));
-  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, target.convStreamChunkPositions));
+  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, options.convStreamChunkPositions));
 
   if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
     chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp

@@ -3,8 +3,8 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Value.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 #include <cstdint>
@@ -45,29 +45,20 @@ struct ConvProblem {
   bool isPointwise = false;
 };
 
-struct ConvLoweringState : ConvProblem {
+struct ConvLoweringState {
+  ConvProblem problem;
   mlir::Operation* diagnosticAnchor = nullptr;
   mlir::Value x;
   mlir::Value w;
   mlir::Value b;
-  const spatial::SpatialTargetInfo* target = nullptr;
+  const spatial::SpatialTargetResources* target = nullptr;
+  const ONNXToSpatialPlanningOptions* options = nullptr;
 
-  const spatial::SpatialTargetInfo& targetInfo() const { return *target; }
+  const spatial::SpatialTargetResources& targetInfo() const { return *target; }
+  const ONNXToSpatialPlanningOptions& planningOptions() const;
 };
 
 struct ConvGeometry {
-  int64_t batchSize;
-  int64_t numChannelsIn;
-  int64_t xHeight;
-  int64_t xWidth;
-  int64_t numChannelsOut;
-  int64_t wHeight;
-  int64_t wWidth;
-  int64_t outHeight;
-  int64_t outWidth;
-  int64_t group;
-  int64_t numChannelsInPerGroup;
-  int64_t numChannelsOutPerGroup;
   int64_t k;
   int64_t c;
   int64_t p;
@@ -75,8 +66,6 @@ struct ConvGeometry {
   int64_t matrixUnitsPerProcessor;
   int64_t pack;
   uint64_t im2colElements;
-  bool hasBias;
-  bool isDepthwise;
 };
 
 struct RowInterval {
@@ -94,35 +83,14 @@ struct ConvRowDemand {
 
 enum class ConvMaterializationKind : uint8_t {
   StructuredDepthwise,
-  PointwiseContraction,
   PackedIm2Col,
   StreamedPatch,
   StreamedPacked,
-  OutputChannelTiled,
   InputKTiled,
-  Tiled2D,
 };
 
 struct ConvPlan {
-  ConvGeometry geometry;
-  spatial::ConvLoweringStrategy strategy = spatial::ConvLoweringStrategy::Auto;
-  ConvMaterializationKind materializationKind = ConvMaterializationKind::PackedIm2Col;
-  int64_t laneCount = 0;
-  int64_t mvmCount = 0;
-  int64_t vectorCount = 0;
-  int64_t reductionCount = 0;
-  uint64_t weightElements = 0;
-  uint64_t scratchElements = 0;
-  uint64_t materializationElements = 0;
-  uint64_t communicationElements = 0;
-  spatial::PhysicalLayout resultLayout = spatial::PhysicalLayout::DenseNCHW;
-  bool consumesRowStrip = false;
-  bool producesRowStrip = false;
-  bool requiresInputMaterialization = false;
-  bool requiresOutputMaterialization = false;
-  bool usesContraction = false;
-  bool hasContractionPlan = false;
-  ContractionPlan contraction;
+  ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col;
 };
 
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
@@ -130,18 +98,16 @@ bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOu
 void classifyConvProblem(ConvProblem& problem);
 
 ConvGeometry buildConvGeometry(const ConvProblem& problem,
-                               const spatial::SpatialTargetInfo& target);
+                               const spatial::SpatialTargetResources& target);
 
 mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
                                        spatial::ConvLoweringStrategy strategy,
-                                       const spatial::SpatialTargetInfo& target);
-
-llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
-    const ConvProblem& problem, const spatial::SpatialTargetInfo& target);
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options);
 
 uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
                                     int64_t packFactor,
-                                    const spatial::SpatialTargetInfo& target);
+                                    const ONNXToSpatialPlanningOptions& options);
 
 RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem);
 

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

@@ -53,10 +53,10 @@ struct BlueprintSplatMulToSpatial : OpConversionPattern<ONNXMulOp> {
 
   LogicalResult
   matchAndRewrite(ONNXMulOp op, ONNXMulOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
-    auto blueprint = adaptor.getA().getDefiningOp<spatial::SpatBlueprintOp>();
+    auto blueprint = op.getA().getDefiningOp<spatial::SpatBlueprintOp>();
     Value scalar = adaptor.getB();
     if (!blueprint) {
-      blueprint = adaptor.getB().getDefiningOp<spatial::SpatBlueprintOp>();
+      blueprint = op.getB().getDefiningOp<spatial::SpatBlueprintOp>();
       scalar = adaptor.getA();
     }
     auto scalarAttr = getDenseConstantAttr(scalar);

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

@@ -23,7 +23,6 @@
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
@@ -329,6 +328,8 @@ static FailureOr<RankedTensorType> verifyDynamicGemmBiasType(RankedTensorType cT
 }
 
 static bool hasGemmBias(Value c) {
+  if (!c)
+    return false;
   Operation* definingOp = c.getDefiningOp();
   return (!definingOp || !isa<ONNXNoneOp>(definingOp)) && !isZeroSplatHostConstant(c);
 }
@@ -660,14 +661,14 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
 }
 
 struct GemmToSpatialComputes : OpConversionPattern<ONNXGemmOp> {
-  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
       : OpConversionPattern<ONNXGemmOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
                                 ONNXGemmOpAdaptor gemmOpAdaptor,
                                 ConversionPatternRewriter& rewriter) const override;
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
 } // namespace
@@ -682,7 +683,7 @@ FailureOr<Value> lowerGemmToSpatial(
     bool transB,
     float alpha,
     float beta,
-    const spatial::SpatialTargetInfo& target,
+    const spatial::SpatialTargetResources& target,
     PatternRewriter& rewriter,
     Location loc) {
   auto aType = dyn_cast<RankedTensorType>(a.getType());
@@ -731,14 +732,9 @@ FailureOr<Value> lowerGemmToSpatial(
   problem.m = outType.getDimSize(0);
   problem.k = aType.getDimSize(1);
   problem.n = outType.getDimSize(1);
-  problem.origin = ContractionOrigin::Gemm;
   problem.lhsElementType = aType.getElementType();
   problem.rhsElementType = bType.getElementType();
   problem.resultElementType = outType.getElementType();
-  problem.lhsTransposed = transA;
-  problem.rhsTransposed = transB;
-  problem.alpha = alpha;
-  problem.beta = beta;
   const bool transposeB = transB;
 
   if (!isCompileTimeComputable(b)) {
@@ -789,13 +785,19 @@ FailureOr<Value> lowerGemmToSpatial(
     auto bShape = bType.getShape();
     auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
     if (isCompileTimeComputable(b)) {
-      auto transposedConstant = materializeTransposedContractionConstant(
-          b, transposedType, {1, 0}, rewriter, loc);
-      if (failed(transposedConstant)) {
+      auto denseAttr = getHostConstDenseElementsAttr(b);
+      auto inputType = denseAttr ? dyn_cast<RankedTensorType>(denseAttr.getType()) : nullptr;
+      auto transposedAttr = inputType && inputType.hasStaticShape() && transposedType.hasStaticShape()
+        ? transposeDenseElementsAttr(denseAttr, {1, 0})
+        : FailureOr<DenseElementsAttr>(failure());
+      if (failed(transposedAttr) || transposedAttr->getType() != transposedType) {
         diagnosticAnchor->emitOpError("requires Gemm input B transpose to remain statically materializable");
         return failure();
       }
-      b = *transposedConstant;
+      b = getOrCreateConstant(rewriter,
+                              rewriter.getInsertionBlock()->getParentOp(),
+                              *transposedAttr,
+                              transposedType);
     } else {
       b = createLinalgTranspose(b, transposedType, {1, 0}, rewriter, loc);
     }
@@ -831,7 +833,7 @@ FailureOr<Value> lowerGemmToSpatial(
   }
   b = *paddedB;
   auto paddedAType = RankedTensorType::get({problem.m, paddedReductionSize}, aType.getElementType());
-  a = materializePaddedContractionInput(a, paddedAType, rewriter, loc);
+  a = createPaddedInputCompute(a, paddedAType, rewriter, loc);
   aType = paddedAType;
 
   Value bias;
@@ -896,7 +898,7 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
 
 void populateGemmPatterns(RewritePatternSet& patterns,
                           MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target) {
+                          const spatial::SpatialTargetResources& target) {
   patterns.insert<GemmToSpatialComputes>(ctx, target);
 }
 

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

@@ -7,7 +7,7 @@
 
 namespace onnx_mlir {
 namespace spatial {
-struct SpatialTargetInfo;
+struct SpatialTargetResources;
 }
 
 mlir::FailureOr<mlir::Value> lowerGemmToSpatial(
@@ -20,7 +20,7 @@ mlir::FailureOr<mlir::Value> lowerGemmToSpatial(
     bool transB,
     float alpha,
     float beta,
-    const spatial::SpatialTargetInfo& target,
+    const spatial::SpatialTargetResources& target,
     mlir::PatternRewriter& rewriter,
     mlir::Location loc);
 

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

@@ -12,9 +12,10 @@
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -362,16 +363,16 @@ static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Locati
     return input;
   auto type = cast<RankedTensorType>(value.getType());
   auto shape = type.getShape();
-  auto createONNXTranspose = [&](RankedTensorType resultType, ArrayRef<int64_t> permutation) {
-    return ONNXTransposeOp::create(rewriter, loc, resultType, value, rewriter.getI64ArrayAttr(permutation)).getResult();
+  auto createTranspose = [&](RankedTensorType resultType, ArrayRef<int64_t> permutation) {
+    return createLinalgTranspose(value, resultType, permutation, rewriter, loc);
   };
   if (type.getRank() == 2) {
     auto resultType = RankedTensorType::get({shape[1], shape[0]}, type.getElementType(), type.getEncoding());
-    return createONNXTranspose(resultType, {1, 0});
+    return createTranspose(resultType, {1, 0});
   }
 
   auto resultType = RankedTensorType::get({shape[0], shape[2], shape[1]}, type.getElementType(), type.getEncoding());
-  return createONNXTranspose(resultType, {0, 2, 1});
+  return createTranspose(resultType, {0, 2, 1});
 }
 
 static FailureOr<Value> materializePaddedBatchedWeight(Value value,
@@ -835,7 +836,7 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
   return computeOp->getResult(0);
 }
 
-struct NormalizedMatMulInfo : ContractionProblem {
+struct NormalizedMatMulInfo {
   NormalizedMatMulInfo(RankedTensorType lhsType,
                        RankedTensorType rhsType,
                        RankedTensorType outType,
@@ -844,7 +845,7 @@ struct NormalizedMatMulInfo : ContractionProblem {
                        ContractionProblem problem,
                        bool lhsWasVector,
                        bool rhsWasVector)
-      : ContractionProblem(std::move(problem)),
+      : problem(std::move(problem)),
         lhsType(lhsType),
         rhsType(rhsType),
         outType(outType),
@@ -853,6 +854,7 @@ struct NormalizedMatMulInfo : ContractionProblem {
         lhsWasVector(lhsWasVector),
         rhsWasVector(rhsWasVector) {}
 
+  ContractionProblem problem;
   RankedTensorType lhsType;
   RankedTensorType rhsType;
   RankedTensorType outType;
@@ -862,14 +864,15 @@ struct NormalizedMatMulInfo : ContractionProblem {
   bool rhsWasVector;
 };
 
-struct MatMulLoweringPlan : ContractionProblem {
+struct MatMulLoweringPlan {
   MatMulLoweringPlan(Value lhs, Value rhs, const NormalizedMatMulInfo& info)
-      : ContractionProblem(info),
+      : problem(info.problem),
         lhs(lhs),
         rhs(rhs),
         lhsType(cast<RankedTensorType>(lhs.getType())),
         rhsType(cast<RankedTensorType>(rhs.getType())) {}
 
+  ContractionProblem problem;
   Value lhs;
   Value rhs;
   RankedTensorType lhsType;
@@ -952,14 +955,9 @@ static FailureOr<NormalizedMatMulInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp)
                         m,
                         k,
                         n,
-                        ContractionOrigin::MatMul,
                         lhsType.getElementType(),
                         rhsType.getElementType(),
-                        outType.getElementType(),
-                        false,
-                        false,
-                        lhsWasVector,
-                        rhsWasVector},
+                        outType.getElementType()},
     lhsWasVector,
     rhsWasVector);
 }
@@ -979,10 +977,10 @@ static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs,
   plan.rhs = transposeLastTwoDims(normalizedLhs, rewriter, loc);
   plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
   plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
-  std::swap(plan.lhsBatchShape, plan.rhsBatchShape);
-  std::swap(plan.lhsBatch, plan.rhsBatch);
-  plan.m = info.n;
-  plan.n = info.m;
+  std::swap(plan.problem.lhsBatchShape, plan.problem.rhsBatchShape);
+  std::swap(plan.problem.lhsBatch, plan.problem.rhsBatch);
+  plan.problem.m = info.problem.n;
+  plan.problem.n = info.problem.m;
   plan.transposedResult = true;
   return plan;
 }
@@ -1006,17 +1004,17 @@ static Value finalizeNormalizedMatMulResult(Value value,
   // and removing the synthetic unit matrix axes introduced for vector operands.
   Value result = value;
   RankedTensorType currentType = directOutType;
-  if (info.outputBatchShape.size() > 1) {
-    SmallVector<int64_t> expandedShape(info.outputBatchShape.begin(), info.outputBatchShape.end());
-    expandedShape.push_back(info.m);
-    expandedShape.push_back(info.n);
+  if (info.problem.outputBatchShape.size() > 1) {
+    SmallVector<int64_t> expandedShape(info.problem.outputBatchShape.begin(), info.problem.outputBatchShape.end());
+    expandedShape.push_back(info.problem.m);
+    expandedShape.push_back(info.problem.n);
     auto expandedType = RankedTensorType::get(expandedShape, info.outType.getElementType(), info.outType.getEncoding());
-    result = expandBatchDims(result, expandedType, info.outputBatchShape.size(), rewriter, loc);
+    result = expandBatchDims(result, expandedType, info.problem.outputBatchShape.size(), rewriter, loc);
     currentType = expandedType;
   }
 
   SmallVector<bool> removedAxes(currentType.getRank(), false);
-  if (info.outputBatchShape.empty())
+  if (info.problem.outputBatchShape.empty())
     removedAxes[0] = true;
   if (info.lhsWasVector)
     removedAxes[currentType.getRank() - 2] = true;
@@ -1026,79 +1024,77 @@ static Value finalizeNormalizedMatMulResult(Value value,
 }
 
 struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
-  using OpRewritePattern::OpRewritePattern;
+  explicit MatMulToGemm(MLIRContext* ctx,
+                        const spatial::SpatialTargetResources& target)
+      : OpRewritePattern<ONNXMatMulOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
     if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector)
       return failure();
 
-    if (!shapeInfo->outputBatchShape.empty())
+    if (!shapeInfo->problem.outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
     bool useTransposedForm = isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
 
-    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
-    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
-    int64_t lhsBatchForGemm = shapeInfo->lhsBatch;
-    int64_t rhsBatchForGemm = shapeInfo->rhsBatch;
-    int64_t gemmM = shapeInfo->m;
-    int64_t gemmK = shapeInfo->k;
-    int64_t gemmN = shapeInfo->n;
+    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->problem.lhsBatch, shapeInfo->problem.m, shapeInfo->problem.k, rewriter, loc);
+    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->problem.rhsBatch, shapeInfo->problem.k, shapeInfo->problem.n, rewriter, loc);
+    int64_t lhsBatchForGemm = shapeInfo->problem.lhsBatch;
+    int64_t rhsBatchForGemm = shapeInfo->problem.rhsBatch;
+    int64_t gemmM = shapeInfo->problem.m;
+    int64_t gemmK = shapeInfo->problem.k;
+    int64_t gemmN = shapeInfo->problem.n;
     if (useTransposedForm) {
       lhs = transposeLastTwoDims(matmulOp.getB(), rewriter, loc);
-      lhsBatchForGemm = shapeInfo->rhsBatch;
+      lhsBatchForGemm = shapeInfo->problem.rhsBatch;
       rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc);
-      rhsBatchForGemm = shapeInfo->lhsBatch;
-      gemmM = shapeInfo->n;
-      gemmN = shapeInfo->m;
+      rhsBatchForGemm = shapeInfo->problem.lhsBatch;
+      gemmM = shapeInfo->problem.n;
+      gemmN = shapeInfo->problem.m;
     }
 
-    auto gemmType = RankedTensorType::get({gemmM, gemmN}, shapeInfo->outType.getElementType());
-    Value none = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());
+    auto gemmType = RankedTensorType::get(
+        {gemmM, gemmN}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
     Value lhsMatrix = extractBatchMatrix(lhs, /*batchIndex=*/0, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
     Value rhsMatrix = extractBatchMatrix(rhs, /*batchIndex=*/0, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
-    Value gemmResult = ONNXGemmOp::create(rewriter,
-                                          loc,
-                                          gemmType,
-                                          lhsMatrix,
-                                          rhsMatrix,
-                                          none,
-                                          rewriter.getF32FloatAttr(1.0f),
-                                          rewriter.getF32FloatAttr(1.0f),
-                                          rewriter.getBoolAttr(false),
-                                          rewriter.getBoolAttr(false))
-                         .getY();
+    FailureOr<Value> gemmResult = lowerGemmToSpatial(
+        matmulOp.getOperation(), lhsMatrix, rhsMatrix, Value(), gemmType,
+        /*transA=*/false, /*transB=*/false, /*alpha=*/1.0f, /*beta=*/1.0f,
+        target, rewriter, loc);
+    if (failed(gemmResult))
+      return failure();
+    Value result = *gemmResult;
     if (useTransposedForm)
-      gemmResult =
-        ONNXTransposeOp::create(rewriter, loc, shapeInfo->outType, gemmResult, rewriter.getI64ArrayAttr({1, 0}))
-          .getResult();
+      result = transposeLastTwoDims(result, rewriter, loc);
 
-    if (shapeInfo->outputBatchShape.empty()) {
-      rewriter.replaceOp(matmulOp, gemmResult);
+    if (shapeInfo->problem.outputBatchShape.empty()) {
+      rewriter.replaceOp(matmulOp, result);
       return success();
     }
 
     auto directOutType =
-      RankedTensorType::get({1, shapeInfo->m, shapeInfo->n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
-    Value batchedResult = ensureBatchedTensor(gemmResult, /*batchSize=*/1, shapeInfo->m, shapeInfo->n, rewriter, loc);
+      RankedTensorType::get({1, shapeInfo->problem.m, shapeInfo->problem.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+    Value batchedResult = ensureBatchedTensor(result, /*batchSize=*/1, shapeInfo->problem.m, shapeInfo->problem.n, rewriter, loc);
     Value finalResult = finalizeNormalizedMatMulResult(batchedResult, directOutType, *shapeInfo, rewriter, loc);
     rewriter.replaceOp(matmulOp, finalResult);
     return success();
   }
+
+  const spatial::SpatialTargetResources& target;
 };
 
 struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
   explicit MatMulBatchedToSpatialComputes(MLIRContext* ctx,
-                                           const spatial::SpatialTargetInfo& target)
+                                           const spatial::SpatialTargetResources& target)
       : OpRewritePattern<ONNXMatMulOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
     if (failed(shapeInfo))
       return failure();
-    if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->outputBatchShape.empty())
+    if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->problem.outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
@@ -1114,62 +1110,63 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     const bool rhsStoredAsRows = rhsRows && !useTransposedForm;
 
     Value lhs =
-      normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType, shapeInfo->lhsWasVector, rewriter, loc);
+      normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType,
+                             shapeInfo->lhsWasVector, rewriter, loc);
     Value rhs = normalizeMatMulOperand(
-      rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc);
-    lhs = collapseBatchDims(lhs, shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
+      rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType,
+      shapeInfo->rhsWasVector, rewriter, loc);
+    lhs = collapseBatchDims(lhs, shapeInfo->problem.lhsBatch, shapeInfo->problem.m, shapeInfo->problem.k, rewriter, loc);
     rhs = collapseBatchDims(rhs,
-                            shapeInfo->rhsBatch,
-                            rhsStoredAsRows ? shapeInfo->n : shapeInfo->k,
-                            rhsStoredAsRows ? shapeInfo->k : shapeInfo->n,
+                            shapeInfo->problem.rhsBatch,
+                            rhsStoredAsRows ? shapeInfo->problem.n : shapeInfo->problem.k,
+                            rhsStoredAsRows ? shapeInfo->problem.k : shapeInfo->problem.n,
                             rewriter,
                             loc);
     MatMulLoweringPlan plan = buildLoweringPlan(
       lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc);
 
-    plan.lhs = ensureBatchedTensor(plan.lhs, plan.lhsBatch, plan.m, plan.k, rewriter, loc);
+    plan.lhs = ensureBatchedTensor(plan.lhs, plan.problem.lhsBatch, plan.problem.m, plan.problem.k, rewriter, loc);
     plan.rhs = ensureBatchedTensor(plan.rhs,
-                                   plan.rhsBatch,
-                                   rhsStoredAsRows ? plan.n : plan.k,
-                                   rhsStoredAsRows ? plan.k : plan.n,
+                                   plan.problem.rhsBatch,
+                                   rhsStoredAsRows ? plan.problem.n : plan.problem.k,
+                                   rhsStoredAsRows ? plan.problem.k : plan.problem.n,
                                    rewriter,
                                    loc);
     plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
     plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
     auto directOutType = RankedTensorType::get(
-      {plan.batch, plan.m, plan.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+      {plan.problem.batch, plan.problem.m, plan.problem.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
 
     if (isCompileTimeComputable(plan.rhs)) {
       ContractionPlan contractionPlan = makeContractionPlan(
-        plan, target, ContractionPlanKind::StaticTiled);
+        plan.problem, target, ContractionPlanKind::StaticTiled);
       const int64_t numKSlices = contractionPlan.reductionSlices;
       const int64_t numOutHSlices = contractionPlan.outputTiles;
       const int64_t paddedReductionSize = numKSlices * xbarSize;
       const int64_t paddedOutCols = numOutHSlices * xbarSize;
       auto paddedLhsType = RankedTensorType::get(
-        {plan.lhsBatch, plan.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
+        {plan.problem.lhsBatch, plan.problem.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
       auto paddedRhsType = RankedTensorType::get(
-        {plan.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding());
+        {plan.problem.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding());
       auto paddedOutType =
-        RankedTensorType::get({plan.batch, plan.m, paddedOutCols}, shapeInfo->outType.getElementType());
+        RankedTensorType::get({plan.problem.batch, plan.problem.m, paddedOutCols}, shapeInfo->outType.getElementType());
 
       auto paddedRhs =
-        materializePaddedBatchedWeight(plan.rhs, plan.rhsBatchShape, plan.outputBatchShape, paddedRhsType, rewriter);
+        materializePaddedBatchedWeight(plan.rhs, plan.problem.rhsBatchShape, plan.problem.outputBatchShape, paddedRhsType, rewriter);
       if (succeeded(paddedRhs)) {
-        Value paddedLhs = materializePaddedContractionInput(
-          plan.lhs, paddedLhsType, rewriter, loc);
+        Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
         const int64_t laneCount = contractionPlan.laneCount;
         auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
           laneCount, RankedTensorType::get({1, xbarSize}, shapeInfo->outType.getElementType()));
         auto batchOp = createBatchedVmmBatch(paddedLhs,
                                              *paddedRhs,
                                              paddedLhsType,
-                                             plan.lhsBatchShape,
+                                             plan.problem.lhsBatchShape,
                                              paddedRhsType,
-                                             plan.rhsBatchShape,
-                                             plan.outputBatchShape,
+                                             plan.problem.rhsBatchShape,
+                                             plan.problem.outputBatchShape,
                                              partialPiecesType,
-                                             plan.m,
+                                             plan.problem.m,
                                              numKSlices,
                                              numOutHSlices,
                                              xbarSize,
@@ -1181,7 +1178,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                                     partialPiecesType,
                                                     directOutType,
                                                     paddedOutType,
-                                                    plan.batch,
+                                                    plan.problem.batch,
                                                     numKSlices,
                                                     xbarSize,
                                                     rewriter,
@@ -1190,12 +1187,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
           return failure();
         Value finalResult = *result;
         if (plan.transposedResult) {
-          auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n},
-                                                         shapeInfo->outType.getElementType(),
-                                                         shapeInfo->outType.getEncoding());
-          finalResult =
-            ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
-              .getResult();
+          finalResult = transposeLastTwoDims(finalResult, rewriter, loc);
         }
         finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
         rewriter.replaceOp(matmulOp, finalResult);
@@ -1205,24 +1197,24 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     RankedTensorType blueprintType = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
                                    ? shapeInfo->outType : directOutType;
     SmallVector<int64_t> blueprintBatchShape = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
-                                             ? shapeInfo->outputBatchShape : SmallVector<int64_t> {plan.batch};
-    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.m, plan.k, plan.n, xbarSize);
+                                             ? shapeInfo->problem.outputBatchShape : SmallVector<int64_t> {plan.problem.batch};
+    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.problem.m, plan.problem.k, plan.problem.n, xbarSize);
     ContractionPlan contractionPlan = makeContractionPlan(
-      plan, target, ContractionPlanKind::GroupedRowDynamicVVD,
-      /*laneCount=*/plan.batch * plan.m / rowsPerLane, rowsPerLane);
+      plan.problem, target, ContractionPlanKind::GroupedRowDynamicVVD,
+      /*laneCount=*/plan.problem.batch * plan.problem.m / rowsPerLane, rowsPerLane);
     const int64_t laneCount = contractionPlan.laneCount;
     SmallVector<int64_t> fragmentShape(blueprintType.getRank(), 1);
     fragmentShape[fragmentShape.size() - 2] = rowsPerLane;
-    fragmentShape.back() = plan.n;
+    fragmentShape.back() = plan.problem.n;
     auto fragmentType = RankedTensorType::get(fragmentShape, shapeInfo->outType.getElementType());
     auto rowPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount, fragmentType);
     Value transposedRhs = rhsStoredAsRows ? plan.rhs : transposeLastTwoDims(plan.rhs, rewriter, loc);
     auto batchOp = createBatchedVvdmulBatch(plan.lhs,
-                                            plan.lhsBatchShape,
+                                            plan.problem.lhsBatchShape,
                                             transposedRhs,
-                                            plan.rhsBatchShape,
-                                            plan.outputBatchShape,
-                                            plan.k,
+                                            plan.problem.rhsBatchShape,
+                                            plan.problem.outputBatchShape,
+                                            plan.problem.k,
                                             rowsPerLane,
                                             rowPiecesType,
                                             directOutType,
@@ -1231,23 +1223,33 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                             loc);
     if (failed(batchOp))
       return failure();
+    Value rowPieces = batchOp->getResult(0);
+    if (outputScale) {
+      auto scaleAttr = getHostConstDenseElementsAttr(outputScale);
+      if (!scaleAttr || !scaleAttr.isSplat())
+        return failure();
+      auto scaled = mapGraphBatchFragments(
+        rowPieces, rowPiecesType, rewriter, loc, [&](Value fragment, RankedTensorType fragmentType) {
+          auto splat = DenseElementsAttr::get(fragmentType, scaleAttr.getSplatValue<Attribute>());
+          Value scale = arith::ConstantOp::create(rewriter, loc, fragmentType, splat);
+          return FailureOr<Value>(
+            spatial::SpatVMulOp::create(rewriter, loc, fragmentType, fragment, scale).getResult());
+        });
+      if (failed(scaled))
+        return failure();
+      rowPieces = *scaled;
+      if (auto scaleBatch = rowPieces.getDefiningOp<spatial::SpatGraphComputeBatch>())
+        scaleBatch->setAttr("spat.preserve_compute_boundary", rewriter.getUnitAttr());
+    }
     auto result = createBatchedRowOutputBlueprint(
-      batchOp->getResult(0), blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc);
+      rowPieces, blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc);
     if (failed(result))
       return failure();
     Value finalResult = *result;
     if (plan.transposedResult) {
-      auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n},
-                                                     shapeInfo->outType.getElementType(),
-                                                     shapeInfo->outType.getEncoding());
-      finalResult =
-        ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
-          .getResult();
+      finalResult = transposeLastTwoDims(finalResult, rewriter, loc);
     }
     finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
-    if (outputScale)
-      finalResult = ONNXMulOp::create(
-        rewriter, loc, shapeInfo->outType, finalResult, outputScale).getResult();
     rewriter.replaceOp(matmulOp, finalResult);
     if (foldedTranspose && foldedTranspose->use_empty())
       rewriter.eraseOp(foldedTranspose);
@@ -1256,7 +1258,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     return success();
   }
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 };
 
 struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
@@ -1273,14 +1275,14 @@ struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
 
 void populateMatMulFusionPatterns(RewritePatternSet& patterns,
                                   MLIRContext* ctx,
-                                  const spatial::SpatialTargetInfo& target) {
+                                  const spatial::SpatialTargetResources& target) {
   patterns.add<TransposedRhsMatMulToSpatial>(ctx, target);
 }
 
 void populateMatMulRewritePatterns(RewritePatternSet& patterns,
                                    MLIRContext* ctx,
-                                   const spatial::SpatialTargetInfo& target) {
-  patterns.insert<MatMulToGemm>(ctx);
+                                   const spatial::SpatialTargetResources& target) {
+  patterns.insert<MatMulToGemm>(ctx, target);
   patterns.insert<MatMulBatchedToSpatialComputes>(ctx, target);
 }
 

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp

@@ -17,7 +17,7 @@
 #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/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -161,10 +161,10 @@ struct PoolToSpatialCompute;
 
 template <typename PoolOp, typename PoolOpAdaptor, typename ReduceOp>
 struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
-  PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+  PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
       : OpConversionPattern<PoolOp>(ctx), target(target) {}
 
-  const spatial::SpatialTargetInfo& target;
+  const spatial::SpatialTargetResources& target;
 
   LogicalResult matchAndRewrite(PoolOp poolOp, PoolOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const final {
     Location loc = poolOp.getLoc();
@@ -431,7 +431,7 @@ struct PoolToSpatialCompute<ONNXAveragePoolOp>
 } // namespace
 
 LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
-                                            const spatial::SpatialTargetInfo&) {
+                                            const spatial::SpatialTargetResources&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -447,9 +447,10 @@ LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
 }
 
 FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
-                                         const spatial::SpatialTargetInfo& target,
+                                         Value input,
+                                         const spatial::SpatialTargetResources& target,
                                          PatternRewriter& rewriter) {
-  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
       || inputType.getRank() != 4 || outputType.getRank() != 4)
@@ -476,7 +477,7 @@ FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
 
   auto compute = createSpatCompute<1>(
-    rewriter, planOp.getLoc(), outputType, {}, planOp.getInput(),
+    rewriter, planOp.getLoc(), outputType, {}, input,
     [&](Value input) -> LogicalResult {
       Value paddedInput = createPaddedPoolInput(
         rewriter, planOp.getLoc(), input, inputType,
@@ -615,14 +616,15 @@ static Value extractPoolIndex(PatternRewriter& rewriter,
 }
 
 FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                                            Value input,
                                             std::optional<Value> rowStripInput,
-                                            const spatial::SpatialTargetInfo& target,
+                                            const spatial::SpatialTargetResources& target,
                                             PatternRewriter& rewriter) {
   if (failed(canLowerMaxPoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
-  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = cast<RankedTensorType>(input.getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   const int64_t channels = inputType.getDimSize(1);
   const int64_t inputHeight = inputType.getDimSize(2);
@@ -631,9 +633,9 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   const int64_t outputWidth = outputType.getDimSize(3);
   const int64_t kernelHeight = planOp.getKernelShape()[0];
   const int64_t kernelWidth = planOp.getKernelShape()[1];
-  Value input = rowStripInput.value_or(planOp.getInput());
-  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
+  Value actualInput = rowStripInput.value_or(input);
+  auto actualInputType = dyn_cast<RankedTensorType>(actualInput.getType());
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(actualInput, inputType);
   const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
@@ -681,7 +683,7 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
     TypeRange {outputStorageType},
     outputHeight * tilesPerRow,
     {},
-    ValueRange {input},
+    ValueRange {actualInput},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       SmallVector<Value> inputRows;
       inputRows.reserve(kernelHeight);
@@ -806,7 +808,7 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
 }
 
 LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
-  spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetInfo&) {
+  spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetResources&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -820,9 +822,10 @@ LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
 
 FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
     spatial::SpatGlobalAveragePoolPlanOp planOp,
-    const spatial::SpatialTargetInfo& target,
+    Value input,
+    const spatial::SpatialTargetResources& target,
     PatternRewriter& rewriter) {
-  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape()
       || !outputType.hasStaticShape() || inputType.getRank() != 4
@@ -844,7 +847,7 @@ FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
   const double scaleValue = 1.0 / static_cast<double>(height * width);
 
   auto compute = createSpatCompute<1>(
-    rewriter, planOp.getLoc(), outputType, {}, planOp.getInput(),
+    rewriter, planOp.getLoc(), outputType, {}, input,
     [&](Value input) -> LogicalResult {
       Value output = tensor::EmptyOp::create(
         rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
@@ -895,22 +898,23 @@ FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
 }
 
 FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                                    Value input,
                                                     std::optional<Value> rowStripInput,
-                                                    const spatial::SpatialTargetInfo& target,
+                                                    const spatial::SpatialTargetResources& target,
                                                     PatternRewriter& rewriter) {
   if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
-  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = cast<RankedTensorType>(input.getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   auto elementType = dyn_cast<FloatType>(inputType.getElementType());
   if (!elementType)
     return failure();
 
-  Value input = rowStripInput.value_or(planOp.getInput());
-  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
+  Value actualInput = rowStripInput.value_or(input);
+  auto actualInputType = dyn_cast<RankedTensorType>(actualInput.getType());
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(actualInput, inputType);
   const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
@@ -940,7 +944,7 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
     TypeRange {outputStorageType},
     tilesPerRow,
     ValueRange {zero, scale},
-    ValueRange {input},
+    ValueRange {actualInput},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       Value reduced = args.weights[0];
       for (int64_t row = 0; row < height; ++row) {
@@ -1011,7 +1015,7 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
 
 void populatePoolPatterns(RewritePatternSet& patterns,
                           MLIRContext* ctx,
-                          const spatial::SpatialTargetInfo& target) {
+                          const spatial::SpatialTargetResources& target) {
   patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx, target);
   patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx, target);
 }

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

@@ -7,7 +7,7 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -47,11 +47,11 @@ static SmallVector<ReassociationIndices> getExpandFrom1DReassociation(int64_t ra
   return reassociation;
 }
 
-static Value buildFlatten(Value input,
+static Value buildFlattenBody(Value input,
                               RankedTensorType sourceType,
                               RankedTensorType resultType,
                               int64_t axis,
-                          ConversionPatternRewriter& rewriter,
+                              PatternRewriter& rewriter,
                               Location loc) {
   if (sourceType == resultType)
     return input;
@@ -75,6 +75,25 @@ static Value buildFlatten(Value input,
     rewriter, loc, resultType, flattened, getExpandFrom1DReassociation(resultType.getRank()));
 }
 
+static Value buildFlatten(Value input,
+                          RankedTensorType sourceType,
+                          RankedTensorType resultType,
+                          int64_t axis,
+                          PatternRewriter& rewriter,
+                          Location loc) {
+  if (spatial::isAnySpatialComputeLike(rewriter.getInsertionBlock()->getParentOp()))
+    return buildFlattenBody(input, sourceType, resultType, axis, rewriter, loc);
+
+  auto compute = createSpatCompute<1>(
+      rewriter, loc, TypeRange {resultType}, {}, ValueRange {input},
+      [&](Value computeInput) {
+        spatial::SpatYieldOp::create(
+            rewriter, loc,
+            buildFlattenBody(computeInput, sourceType, resultType, axis, rewriter, loc));
+      });
+  return compute.getResult(0);
+}
+
 struct Flatten : OpConversionPattern<ONNXFlattenOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -97,21 +116,17 @@ struct Flatten : OpConversionPattern<ONNXFlattenOp> {
     if (resultType.getShape()[0] != outerDim || resultType.getShape()[1] != innerDim)
       return failure();
 
-    auto replaceWithFlatten = [&](auto build) -> LogicalResult {
-      Value flattened = materializeOrComputeUnary(adaptor.getInput(), resultType, rewriter, flattenOp.getLoc(), build);
-      rewriter.replaceOp(flattenOp, flattened);
+    auto plan = spatial::SpatFlattenPlanOp::create(
+        rewriter, flattenOp.getLoc(), resultType, adaptor.getInput(),
+        rewriter.getI64IntegerAttr(*axis),
+        spatial::getNCHWLayout(rewriter.getContext()));
+    rewriter.replaceOp(flattenOp, plan.getOutput());
     return success();
-    };
-
-    return replaceWithFlatten([&](Value input) {
-      return buildFlatten(input, sourceType, resultType, *axis, rewriter, flattenOp.getLoc());
-    });
   }
 };
 
 struct RowStripFlattenAnalysis {
   spatial::SpatGraphComputeBatch consumer;
-  tensor::CollapseShapeOp collapse;
   RankedTensorType sourceType;
   RankedTensorType resultType;
   RankedTensorType weightType;
@@ -119,12 +134,11 @@ struct RowStripFlattenAnalysis {
 };
 
 static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
-  spatial::SpatGraphCompute flattenOp, const spatial::SpatialTargetInfo& target) {
-  if (flattenOp.getWeights().size() != 0 || flattenOp.getInputs().size() != 1
-      || flattenOp.getOutputs().size() != 1)
+    spatial::SpatFlattenPlanOp flattenOp, const spatial::SpatialTargetResources& target) {
+  if (flattenOp.getAxis() != 1)
     return failure();
-  auto sourceType = dyn_cast<RankedTensorType>(flattenOp.getInputs().front().getType());
-  auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOutputs().front().getType());
+  auto sourceType = dyn_cast<RankedTensorType>(flattenOp.getInput().getType());
+  auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOutput().getType());
   if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape()
       || sourceType.getRank() != 4 || resultType.getRank() != 2 || sourceType.getDimSize(0) != 1
       || resultType.getDimSize(0) != 1 || resultType.getDimSize(1) != sourceType.getNumElements())
@@ -134,17 +148,21 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
   if (channels > xbarDim && channels % xbarDim != 0)
     return failure();
 
-  auto yieldOp = dyn_cast<spatial::SpatYieldOp>(flattenOp.getBody().front().getTerminator());
-  if (!yieldOp || yieldOp.getOutputs().size() != 1)
+  Value consumerInput = flattenOp.getOutput();
+  Operation* consumerOp = nullptr;
+  while (consumerInput.hasOneUse()) {
+    Operation* user = *consumerInput.getUsers().begin();
+    if (auto materialize = dyn_cast<spatial::SpatMaterializeLayoutOp>(user)) {
+      consumerInput = materialize.getOutput();
+      continue;
+    }
+    consumerOp = user;
+    break;
+  }
+  if (!consumerOp)
     return failure();
-  auto collapse = yieldOp.getOutputs().front().getDefiningOp<tensor::CollapseShapeOp>();
-  if (!collapse || collapse.getSrc() != *flattenOp.getInputArgument(0))
-    return failure();
-
-  if (!flattenOp.getResult(0).hasOneUse())
-    return failure();
-  auto consumer = dyn_cast<spatial::SpatGraphComputeBatch>(*flattenOp.getResult(0).getUsers().begin());
-  if (!consumer || consumer.getInputs().size() != 1 || consumer.getInputs().front() != flattenOp.getResult(0)
+  auto consumer = dyn_cast<spatial::SpatGraphComputeBatch>(consumerOp);
+  if (!consumer || consumer.getInputs().size() != 1 || consumer.getInputs().front() != consumerInput
       || consumer.getWeights().size() != 1)
     return failure();
   auto weightType = dyn_cast<RankedTensorType>(consumer.getWeights().front().getType());
@@ -155,21 +173,32 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
   if (llvm::none_of(consumer.getBody().getOps<spatial::SpatVMMOp>(),
                     [](spatial::SpatVMMOp) { return true; }))
     return failure();
-  return RowStripFlattenAnalysis {consumer, collapse, sourceType, resultType, weightType, weight};
+  return RowStripFlattenAnalysis {consumer, sourceType, resultType, weightType, weight};
 }
 
 } // namespace
 
 void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Flatten>(ctx); }
 
-LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp,
-                                          const spatial::SpatialTargetInfo& target) {
+FailureOr<Value> lowerDenseFlattenPlan(spatial::SpatFlattenPlanOp planOp,
+                                        Value input,
+                                        PatternRewriter& rewriter) {
+  auto sourceType = dyn_cast<RankedTensorType>(input.getType());
+  auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
+    return failure();
+  return buildFlatten(input, sourceType, resultType, planOp.getAxis(), rewriter,
+                      planOp.getLoc());
+}
+
+LogicalResult canLowerFlattenFromRowStrip(spatial::SpatFlattenPlanOp flattenOp,
+                                          const spatial::SpatialTargetResources& target) {
   return succeeded(analyzeRowStripFlatten(flattenOp, target)) ? success() : failure();
 }
 
 LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,
-                                       spatial::SpatGraphCompute flattenOp,
-                                       const spatial::SpatialTargetInfo& target,
+                                       spatial::SpatFlattenPlanOp flattenOp,
+                                       const spatial::SpatialTargetResources& target,
                                        PatternRewriter& rewriter) {
   FailureOr<RowStripFlattenAnalysis> analysis = analyzeRowStripFlatten(flattenOp, target);
   if (failed(analysis))
@@ -206,19 +235,20 @@ LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,
                                               analysis->weightType);
   analysis->consumer->setOperand(0, reorderedWeight);
 
-  BlockArgument flattenInput = *flattenOp.getInputArgument(0);
-  flattenOp.getInputsMutable().assign(input.storage);
-  flattenInput.setType(storageType);
-
-  OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPoint(analysis->collapse);
+  auto compute = createSpatCompute<1>(
+      rewriter, flattenOp.getLoc(), TypeRange {analysis->resultType}, {},
+      ValueRange {input.storage}, [&](Value storage) {
         auto flatType = RankedTensorType::get(
             {storageType.getNumElements()}, storageType.getElementType(), storageType.getEncoding());
         Value flat = tensor::CollapseShapeOp::create(
-    rewriter, flattenOp.getLoc(), flatType, flattenInput, getCollapseTo1DReassociation(storageType.getRank()));
+            rewriter, flattenOp.getLoc(), flatType, storage,
+            getCollapseTo1DReassociation(storageType.getRank()));
         Value logicalInput = tensor::ExpandShapeOp::create(
-    rewriter, flattenOp.getLoc(), analysis->resultType, flat, getExpandFrom1DReassociation(2));
-  rewriter.replaceOp(analysis->collapse, logicalInput);
+            rewriter, flattenOp.getLoc(), analysis->resultType, flat,
+            getExpandFrom1DReassociation(2));
+        spatial::SpatYieldOp::create(rewriter, flattenOp.getLoc(), logicalInput);
+      });
+  rewriter.replaceOp(flattenOp, compute.getResult(0));
   return success();
 }
 

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

@@ -9,7 +9,7 @@
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -193,7 +193,7 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
 
 LogicalResult canLowerResizeNearestPlanToRowStrip(
     spatial::SpatResizeNearestPlanOp planOp,
-    const spatial::SpatialTargetInfo&) {
+    const spatial::SpatialTargetResources&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   return success(inputType && outputType && inputType.hasStaticShape()
@@ -204,16 +204,17 @@ LogicalResult canLowerResizeNearestPlanToRowStrip(
 }
 
 FailureOr<Value> lowerSelectedResizeNearestPlan(
-    spatial::SpatResizeNearestPlanOp planOp, std::optional<Value> rowStripInput,
-    const spatial::SpatialTargetInfo&,
+    spatial::SpatResizeNearestPlanOp planOp, Value input,
+    std::optional<Value> rowStripInput,
+    const spatial::SpatialTargetResources&,
     PatternRewriter& rewriter) {
-  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = cast<RankedTensorType>(input.getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   if (rowStripInput)
     return buildRowStripNearestResize(
       *rowStripInput, inputType, outputType, rewriter, planOp.getLoc());
   return buildDenseNearestResize(
-    planOp.getInput(), inputType, outputType, rewriter, planOp.getLoc());
+    input, inputType, outputType, rewriter, planOp.getLoc());
 }
 
 void populateResizePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Resize>(ctx); }

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -1,265 +0,0 @@
-#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 =
-        PimReceiveOp::create(
-          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer,
-          arith::ConstantIndexOp::create(rewriter, receiveOp.getLoc(), 0),
-          *sizeAttr, receiveOp.getSourceCoreId())
+      Value zero = arith::ConstantIndexOp::create(
+        rewriter, receiveOp->getLoc(), 0);
+      Value received;
+      if (hostWaitLoad) {
+        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(
-      rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, *sizeAttr, op.getSourceCoreId());
-    copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
-    Value received = receive.getOutput();
+    auto received = createReceive(outputBuffer, zero, *sizeAttr);
+    if (failed(received))
+      return failure();
     if (!destinationInsert) {
-      rewriter.replaceOp(op, received);
+      rewriter.replaceOp(op, *received);
       return success();
     }
 
@@ -99,11 +117,43 @@ 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();
+    });
+  }
 };
 
 struct ExtractRowsLowering : OpRewritePattern<spatial::SpatExtractRowsOp> {
@@ -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

@@ -29,13 +29,13 @@
 #include "Common/IR/ConstantUtils.hpp"
 #include "Common/PimCommon.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
-#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "Conversion/SpatialToPim/Patterns.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Spatial/SpatialOps.hpp"
-#include "Pass/PIMPasses.h"
+#include "Passes/PIMPasses.h"
 #include "SpatialToPimPass.hpp"
 
 using namespace mlir;
@@ -66,17 +66,20 @@ createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTe
   return padOp.getResult();
 }
 
-static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
+static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter,
+                                                      Location loc,
+                                                      Value vector,
+                                                      int64_t crossbarSize) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
   assert(isHVectorShape(shape) && "expected a horizontal vector");
-  assert(shape[1] <= static_cast<int64_t>(crossbarSize) && "vector width must fit in one crossbar");
+  assert(shape[1] <= crossbarSize && "vector width must fit in one crossbar");
 
-  if (shape[1] == static_cast<int64_t>(crossbarSize))
+  if (shape[1] == crossbarSize)
     return vector;
 
   auto paddedType = RankedTensorType::get(
-    {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
+    {shape[0], crossbarSize}, vectorType.getElementType(), vectorType.getEncoding());
   return createZeroPaddedTensor(rewriter, loc, vector, paddedType);
 }
 
@@ -84,6 +87,11 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   outputTensors.clear();
   operationsToRemove.clear();
   ModuleOp moduleOp = getOperation();
+  if (!hasTarget || failed(targetResources.verify())) {
+    moduleOp.emitError("Spatial-to-PIM lowering requires valid injected target resources");
+    signalPassFailure();
+    return;
+  }
   MLIRContext* ctx = moduleOp.getContext();
 
   auto entryFunc = getPimEntryFunc(moduleOp);
@@ -118,6 +126,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                     spatial::SpatConcatOp,
                     spatial::SpatChannelReceiveOp,
                     spatial::SpatChannelSendOp,
+                    spatial::SpatHostStoreSyncOp,
+                    spatial::SpatHostWaitLoadOp,
                     spatial::SpatExtractRowsOp>();
 
   RewritePatternSet initialPatterns(ctx);
@@ -132,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))) {
@@ -174,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);
@@ -194,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;
@@ -243,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);
@@ -265,15 +296,16 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
     ArrayRef<int64_t> outputShape = outputType.getShape();
     assert(isHVectorShape(outputShape) && "expected a horizontal vector output");
     auto weightType = cast<RankedTensorType>(vmmOp.getWeight().getType());
-    const int64_t xbarDim = static_cast<int64_t>(crossbarSize);
+    const int64_t xbarDim = static_cast<int64_t>(targetResources.matrixShape.columns);
     const int64_t paddedOutputWidth = ceilIntegerDivide(outputShape[1], xbarDim) * xbarDim;
     assert(weightType.getRank() == 2 && weightType.getDimSize(1) == paddedOutputWidth
            && "expected VMM weight width to match the padded output width");
-    assert(paddedOutputWidth / xbarDim <= static_cast<int64_t>(crossbarCountInCore)
+    assert(paddedOutputWidth / xbarDim <= static_cast<int64_t>(targetResources.matrixUnitsPerProcessor)
            && "output width must fit in one core");
 
     rewriter.setInsertionPoint(vmmOp);
-    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
+    auto paddedInput = padHVectorInputToCrossbarSize(
+      rewriter, vmmOp.getLoc(), vmmOp.getInput(), xbarDim);
     if (failed(paddedInput)) {
       hasFailure = true;
       return WalkResult::interrupt();
@@ -375,4 +407,9 @@ void raptor::SpatialToPimPass::eraseOpsToRemove() {
 
 std::unique_ptr<Pass> createSpatialToPimPass() { return std::make_unique<raptor::SpatialToPimPass>(); }
 
+std::unique_ptr<Pass> createSpatialToPimPass(
+    const spatial::SpatialTargetResources& target) {
+  return std::make_unique<raptor::SpatialToPimPass>(target);
+}
+
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -18,6 +18,7 @@
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 namespace raptor {
@@ -28,7 +29,10 @@ struct SpatialToPimPass : mlir::PassWrapper<SpatialToPimPass, mlir::OperationPas
   llvm::StringRef getDescription() const override { return "Lower Spatial ops to PIM-ready format"; }
 
   SpatialToPimPass() = default;
-  SpatialToPimPass(const SpatialToPimPass& pass) {}
+  explicit SpatialToPimPass(const spatial::SpatialTargetResources& target)
+      : targetResources(target), hasTarget(true) {}
+  SpatialToPimPass(const SpatialToPimPass& pass)
+      : targetResources(pass.targetResources), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() final;
 
@@ -37,6 +41,8 @@ private:
 
   llvm::SmallVector<OutputTensorFactory> outputTensors;
   llvm::SmallVector<mlir::Operation*> operationsToRemove;
+  spatial::SpatialTargetResources targetResources;
+  bool hasTarget = false;
 
   mlir::LogicalResult allocateAndInitializeCoreLocalVariables(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
   mlir::LogicalResult

src/PIM/Dialect/Pim/CMakeLists.txt

@@ -1,12 +1,12 @@
 add_onnx_mlir_dialect(Pim pim)
 add_onnx_mlir_dialect_doc(pim Pim.td)
 
-add_subdirectory(Analysis)
-add_subdirectory(Transforms/Bufferization)
-add_subdirectory(Transforms/HostConstantFolding)
-add_subdirectory(Transforms/InstructionSelection)
-add_subdirectory(Transforms/LocalMemoryPlanning)
-add_subdirectory(Transforms/Verification)
+add_subdirectory(Passes/Analyses)
+add_subdirectory(Passes/Transforms/Bufferization)
+add_subdirectory(Passes/Transforms/HostConstantFolding)
+add_subdirectory(Passes/Transforms/InstructionSelection)
+add_subdirectory(Passes/Transforms/LocalMemoryPlanning)
+add_subdirectory(Passes/Transforms/Verification)
 
 add_pim_library(PimOps
   PimOps.hpp

src/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.cpp

@@ -8,7 +8,7 @@
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;

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

@@ -3,8 +3,8 @@
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp"
 
 using namespace mlir;
 using namespace bufferization;

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

@@ -1,4 +1,4 @@
-#include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"

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)
-      && size <= *targetBytes && size <= *sourceBytes) {
+  if (succeeded(targetBytes) && succeeded(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/Bufferization/OpBufferizationInterfaces.cpp

@@ -7,7 +7,7 @@
 #include "OpBufferizationInterfaces.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/BufferizationUtils.hpp"
 
 using namespace mlir;
 using namespace bufferization;

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

@@ -20,11 +20,11 @@
 #include "Common/Support/Diagnostics.hpp"
 #include "Compiler/PimCodeGen.hpp"
 #include "Dialect/Pim/PimOps.hpp"
-#include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
-#include "Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp"
+#include "Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 
 using namespace mlir;
@@ -382,8 +382,51 @@ static bufferization::OneShotBufferizationOptions makePimBufferizationOptions()
   return options;
 }
 
-static LogicalResult preparePimBufferization(func::FuncOp funcOp) {
+static LogicalResult verifyPimCoresNeedNoTensorCopies(
+    ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions& baseOptions) {
+  static constexpr StringLiteral kExistingAlloc = "raptor.existing_core_alloc";
+  OwningOpRef<ModuleOp> clone = moduleOp.clone();
+  clone->walk([&](bufferization::AllocTensorOp alloc) {
+    if (alloc->getParentOfType<pim::PimCoreOp>()
+        || alloc->getParentOfType<pim::PimCoreBatchOp>())
+      alloc->setAttr(kExistingAlloc, UnitAttr::get(moduleOp.getContext()));
+  });
+
+  auto options = baseOptions;
+  options.bufferizeFunctionBoundaries = false;
+  options.opFilter.allowOperation([](Operation* op) {
+    return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op)
+        || op->getParentOfType<pim::PimCoreOp>()
+        || op->getParentOfType<pim::PimCoreBatchOp>();
+  });
+
+  bufferization::BufferizationState state;
+  if (failed(bufferization::insertTensorCopies(*clone, options, state))) {
+    moduleOp.emitError("official one-shot analysis failed while verifying PIM core copy freedom");
+    return failure();
+  }
+
+  CappedDiagnosticReporter diagnostics;
+  clone->walk([&](bufferization::AllocTensorOp alloc) {
+    if (alloc->hasAttr(kExistingAlloc)
+        || (!alloc->getParentOfType<pim::PimCoreOp>()
+            && !alloc->getParentOfType<pim::PimCoreBatchOp>()))
+      return;
+    Operation* requiredBy = alloc->getUsers().empty()
+      ? alloc.getOperation() : *alloc->getUsers().begin();
+    diagnostics.report(requiredBy, [](Operation* op) {
+      op->emitOpError("official one-shot bufferization requires a tensor copy inside a PIM core");
+    });
+  });
+  diagnostics.emitSuppressedSummary(moduleOp, "required PIM core tensor copies");
+  return success(!diagnostics.hasFailure());
+}
+
+static LogicalResult preparePimBufferization(
+    ModuleOp moduleOp, func::FuncOp funcOp, bool verifyCopyFreedom) {
   materializeWritableConstantDestinations(funcOp);
+  if (verifyCopyFreedom)
+    return verifyPimCoresNeedNoTensorCopies(moduleOp, makePimBufferizationOptions());
   return success();
 }
 
@@ -626,6 +669,9 @@ struct PimBufferizationPreparationPass
     : PassWrapper<PimBufferizationPreparationPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPreparationPass)
 
+  explicit PimBufferizationPreparationPass(bool verifyCopyFreedom = false)
+      : verifyCopyFreedom(verifyCopyFreedom) {}
+
   StringRef getArgument() const override { return "pim-bufferization-preparation"; }
   StringRef getDescription() const override {
     return "Prepare writable tensor destinations for PIM one-shot bufferization.";
@@ -638,9 +684,12 @@ struct PimBufferizationPreparationPass
       signalPassFailure();
       return;
     }
-    if (failed(preparePimBufferization(*funcOp)))
+    if (failed(preparePimBufferization(moduleOp, *funcOp, verifyCopyFreedom)))
       signalPassFailure();
   }
+
+private:
+  bool verifyCopyFreedom;
 };
 
 struct PimOneShotBufferizationPass
@@ -735,6 +784,10 @@ std::unique_ptr<Pass> createPimBufferizationPreparationPass() {
   return std::make_unique<PimBufferizationPreparationPass>();
 }
 
+std::unique_ptr<Pass> createPimBufferizationPreparationPass(bool verifyCopyFreedom) {
+  return std::make_unique<PimBufferizationPreparationPass>(verifyCopyFreedom);
+}
+
 std::unique_ptr<Pass> createPimOneShotBufferizationPass() {
   return std::make_unique<PimOneShotBufferizationPass>();
 }

src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp

@@ -5,7 +5,7 @@
 
 #include "Patterns.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
 
 using namespace mlir;
 

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

@@ -8,8 +8,8 @@
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 
 using namespace llvm;
 using namespace mlir;

src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.cpp

@@ -8,8 +8,8 @@
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 
 using namespace llvm;
 using namespace mlir;

src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp"
 
 namespace onnx_mlir {
 

src/PIM/Dialect/Pim/Passes/Transforms/Verification/CMakeLists.txt

@@ -5,7 +5,6 @@ add_pim_library(OMPimVerification
 
   LINK_LIBS PUBLIC
   OMPimCommon
-  OMPimCompilerOptions
   OMPimBufferization
   OMPimLocalMemoryLifetimeAnalysis
   PimOps

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

@@ -18,11 +18,11 @@
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 using namespace mlir;
 
@@ -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,
@@ -748,12 +750,43 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
   }
 
   VerificationPass() {}
-  VerificationPass(const VerificationPass& pass) {}
+  VerificationPass(const spatial::SpatialTargetResources& target,
+                   bool detectCommunicationDeadlock)
+      : targetResources(target), hasTarget(true),
+        detectCommunicationDeadlock(detectCommunicationDeadlock) {}
+  VerificationPass(const VerificationPass& pass)
+      : targetResources(pass.targetResources), hasTarget(pass.hasTarget),
+        detectCommunicationDeadlock(pass.detectCommunicationDeadlock) {}
 
   void runOnOperation() override {
     ModuleOp moduleOp = getOperation();
     pim::CappedDiagnosticReporter diagnostics;
 
+    if (!hasTarget || failed(targetResources.verify())) {
+      moduleOp.emitError("PIM codegen verification requires valid injected target resources");
+      signalPassFailure();
+      return;
+    }
+
+    const int64_t xbarDim = static_cast<int64_t>(targetResources.matrixShape.columns);
+    moduleOp.walk([&](pim::PimVMMOp vmmOp) {
+      auto weightType = dyn_cast<ShapedType>(vmmOp.getWeight().getType());
+      auto inputType = dyn_cast<ShapedType>(vmmOp.getInput().getType());
+      if (!weightType || !inputType || weightType.getRank() != 2 || inputType.getRank() != 2)
+        return;
+      int64_t rows = weightType.getDimSize(0);
+      int64_t columns = weightType.getDimSize(1);
+      if (rows > xbarDim
+          || columns > xbarDim * static_cast<int64_t>(targetResources.matrixUnitsPerProcessor)
+          || columns % xbarDim != 0
+          || inputType.getDimSize(1) != xbarDim) {
+        diagnostics.report(vmmOp.getOperation(), [xbarDim](Operation* op) {
+          op->emitOpError() << "VMM dimensions do not fit the injected target (crossbar size "
+                            << xbarDim << ")";
+        });
+      }
+    });
+
     moduleOp.walk([&](Operation* op) {
       if (op->getDialect()->getNamespace() != "spat")
         return;
@@ -811,7 +844,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
     }
 
     bool hasFailure = false;
-    if (pimDetectCommunicationDeadlock && failed(verifyNoStaticCommunicationDeadlock(moduleOp, diagnostics)))
+    if (detectCommunicationDeadlock && failed(verifyNoStaticCommunicationDeadlock(moduleOp, diagnostics)))
       hasFailure = true;
 
     if (diagnostics.hasFailure()) {
@@ -825,6 +858,10 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
   }
 
 private:
+  spatial::SpatialTargetResources targetResources;
+  bool hasTarget = false;
+  bool detectCommunicationDeadlock = false;
+
   template <typename CoreOpTy>
   static LogicalResult
   verifyCoreWeights(ModuleOp moduleOp, CoreOpTy coreOp, pim::CappedDiagnosticReporter& diagnostics) {
@@ -1050,4 +1087,10 @@ private:
 
 std::unique_ptr<Pass> createPimVerificationPass() { return std::make_unique<VerificationPass>(); }
 
+std::unique_ptr<Pass> createPimVerificationPass(
+    const spatial::SpatialTargetResources& target,
+    bool detectCommunicationDeadlock) {
+  return std::make_unique<VerificationPass>(target, detectCommunicationDeadlock);
+}
+
 } // namespace onnx_mlir

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/Pim/PimOpsVerify.cpp

@@ -10,7 +10,6 @@
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;
@@ -157,16 +156,10 @@ LogicalResult PimVMMOp::verify() {
   int64_t M = matrixShape[1];
   if (N <= 0 || M <= 0)
     return emitError("matrix shape must be (N, M) with N > 0 and M > 0");
-  const int64_t xbarDim = static_cast<int64_t>(crossbarSize);
-  if (N > xbarDim || M > xbarDim * static_cast<int64_t>(crossbarCountInCore))
-    return emitError("matrix dimensions must fit in one array group");
-  if (M % xbarDim != 0)
-    return emitError("matrix output width must be padded to a whole number of crossbars");
-
   int64_t vector1 = vectorShape[0];
   int64_t vectorWidth = vectorShape[1];
-  if (vector1 != 1 || vectorWidth != xbarDim)
-    return emitError("vector shape must be (1, crossbar-size)");
+  if (vector1 != 1 || vectorWidth <= 0)
+    return emitError("vector shape must be (1, positive-width)");
 
   int64_t output1 = outputShape[0];
   int64_t outputWidth = outputShape[1];

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -6,7 +6,6 @@ mlir_tablegen(SpatialEnums.hpp.inc -gen-enum-decls "-I${ONNX_MLIR_SRC_ROOT}")
 mlir_tablegen(SpatialEnums.cpp.inc -gen-enum-defs "-I${ONNX_MLIR_SRC_ROOT}")
 add_public_tablegen_target(OMSpatialEnumsIncGen)
 
-set(LLVM_TARGET_DEFINITIONS SpatialLayoutInterface.td)
 mlir_tablegen(SpatialLayoutInterface.hpp.inc -gen-op-interface-decls "-I${ONNX_MLIR_SRC_ROOT}")
 mlir_tablegen(SpatialLayoutInterface.cpp.inc -gen-op-interface-defs "-I${ONNX_MLIR_SRC_ROOT}")
 add_public_tablegen_target(OMSpatialLayoutInterfaceIncGen)
@@ -17,26 +16,27 @@ add_pim_library(SpatialOps
   SpatialOpsVerify.cpp
   SpatialOpsCanonicalization.cpp
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
-  Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
-  Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
-  Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
-  Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
-  Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp
-  Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp
-  Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp
-  Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
-  Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
-  Transforms/MergeComputeNodes/DeferredResultRealization.cpp
-  Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
-  Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp
-  Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
-  Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
-  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
-  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
-  Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
-  Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
-  Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
-  Transforms/TrivialGraphComputeMergePass.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredResultRealization.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
+  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
 
@@ -52,6 +52,5 @@ add_pim_library(SpatialOps
   MLIRBufferizationTransforms
   OMMlirDialects
   OMONNXOps
-  OMPimCompilerOptions
   PimOps
 )

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(
-        context.rewriter, loc, transferGrids->channels().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
-        transferGrids->sourceCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
-        transferGrids->targetCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), *payload);
+    Value sourceCore = transferGrids->sourceCores().emitLookup(
+      action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    Value targetCore = transferGrids->targetCores().emitLookup(
+      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,
-    grids->channels().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
-    grids->sourceCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
-    grids->targetCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
+  Value sourceCore = grids->sourceCores().emitLookup(
+    row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
+  Value targetCore = grids->targetCores().emitLookup(
+    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 =
-        SpatChannelReceiveOp::create(context.rewriter, loc, fragmentType,
-                                     transferGrids->channels().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
-                                     transferGrids->sourceCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
-                                     transferGrids->targetCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
+    Value sourceCore = transferGrids->sourceCores().emitLookup(
+      entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    Value targetCore = 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 {