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

README.md

@@ -53,19 +53,20 @@ 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
-   `SpatialTargetResources`. Layout-aware plan ops advertise typed alternatives
+   `SpatialTargetInfo`. 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. Passes and their invariant/layout analyses live under
+   separate layers.
-   `Passes/Transforms` and `Passes/Analyses`.
 
 2. **Merge, schedule, and realize Spatial communication**
-   (`src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes`).
+   (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
-   `TrivialGraphComputeMerge` performs local graph merging. One
+   `TrivialGraphComputeMerge` performs local graph merging, then
-   `ScheduleAndRealizeSpatial` pass then owns scheduling, intermediate
+   `ScheduleSpatialGraph` materializes scheduled computes and explicit deferred
-   verification, communication realization, and final verification. Supporting
+   communication. `VerifyScheduledSpatial` checks that intermediate contract;
+   `RealizeSpatialCommunication` resolves transfers and forwarding; and
+   `VerifyRealizedSpatial` checks the final scheduled graph. Supporting
    scheduling code lives under `MergeComputeNodes/Scheduling`.
 
 3. **Spatial -> Pim** (`src/PIM/Conversion/SpatialToPim`).
@@ -73,7 +74,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/Passes/Transforms/Bufferization`).
+4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
    `PimBufferizationPreparation` establishes writable destinations without
    duplicating the one-shot copy analysis, `PimOneShotBufferization` runs
    MLIR's one-shot analysis,
@@ -82,11 +83,10 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    absence, contiguity, and copy address spaces.
 
 5. **PIM local-memory planning**
-   (`src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning`).
+   (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`).
    Computes whole-core lifetimes, reuses addresses for non-overlapping
-   allocations, and records the explicit plan in PIM IR. Reusable lifetime
+   allocations, and records the explicit plan in PIM IR.
-   analysis lives under `src/PIM/Dialect/Pim/Passes/Analyses`.
+6. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and
-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/Passes` - pass registration and auxiliary passes.
+- `src/PIM/Pass` - pass registration and auxiliary passes.
 - `src/PIM/PimAccelerator.{cpp,hpp}` - ONNX-MLIR accelerator entry point.
 
 ## PIM compiler options
@@ -132,8 +132,14 @@ 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,13 +326,9 @@ 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,11 +601,7 @@ fn json_to_wait(
     inst_data_builder: &mut InstructionDataBuilder,
     json: &Value,
 ) -> Result<()> {
-    inst_data_builder.set_offset_select_value(
+    todo!("Not present in the compiler");
-        json_i64!(json, "event_register") as i32,
-        json_i64!(json, "wait_value") as i32,
-    );
-    inst_builder.make_inst(wait, inst_data_builder.build());
     Ok(())
 }
 
@@ -614,10 +610,7 @@ fn json_to_sync(
     inst_data_builder: &mut InstructionDataBuilder,
     json: &Value,
 ) -> Result<()> {
-    inst_data_builder
+    todo!("Not present in the compiler");
-        .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,8 +93,6 @@ struct DeadlockInfo {
     states: String,
 }
 
-type SyncEvents = Vec<[i32; 32]>;
-
 fn print_status(core_instructions: &[CoreInstructions]) {
     let mut tot_instructions = 0;
     let mut progress = 0;
@@ -137,7 +135,6 @@ 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();
 
@@ -172,9 +169,7 @@ impl<'a> Executable<'a> {
                     now = SystemTime::now();
                 }
             }
-            if handle_wait_sync(cores_instructions, &mut sync_events, core_result) {
+            handle_wait_sync(cpu, cores_instructions, core_result);
-                cpu_progressed = 0;
-            }
             match handle_send_recv(cpu, cores_instructions, send_recv, core_result) {
                 (true, other_cpu_index) => {
                     cpu_progressed = 0;
@@ -354,31 +349,12 @@ fn detect_deadlock(cores_instructions: &[CoreInstructions]) -> Option<DeadlockIn
     None
 }
 
-fn handle_wait_sync(
+fn handle_wait_sync<'a, 'b, 'c>(
-    core_instructions: &mut [CoreInstructions],
+    cpu: &'b mut CPU<'a>,
-    events: &mut SyncEvents,
+    core_instructions: &'c mut [CoreInstructions],
     core_result: InstructionStatus,
-) -> bool {
+) where
-    match core_result {
+    'a: 'b,
-        InstructionStatus::Sync(data) => {
+    'a: 'c,
-            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, if transfered { receiver } else { 0 })
+            (transfered, receiver)
         }
         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, if transfered { sender } else { 0 })
+            (transfered, sender)
         }
         _ => (false, 0),
     }

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

@@ -295,68 +295,3 @@ 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(Passes)
+add_subdirectory(Pass)
 add_subdirectory(Compiler)
 add_subdirectory(Conversion)
 

src/PIM/Common/PimCommon.hpp

@@ -32,9 +32,6 @@ 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, "",     "event_register", "wait_value", ""}, // wait
+  {false, false, false, "",     "",     "",      ""    }, // wait
-  {false, false, false, "core", "event_register", "", ""}, // sync
+  {false, false, false, "",     "",     "",      ""    }, // sync
 }};
 static_assert(kInstructionJsonFormats.size() == kOpcodeCount);
 

src/PIM/Compiler/PimCodeGen.cpp

@@ -692,34 +692,6 @@ 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");
@@ -1019,8 +991,6 @@ 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,8 +217,6 @@ 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,8 +2,6 @@
 
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
-#include <limits>
-
 #define DEBUG_TYPE "PimCompilerOptions"
 
 namespace onnx_mlir {
@@ -72,6 +70,11 @@ 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"),
@@ -84,11 +87,6 @@ llvm::cl::opt<uint64_t> pimConvStreamChunkPositions(
   llvm::cl::init(1024),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> pimReportConvLowering("pim-report-conv-lowering",
-                                          llvm::cl::desc("Emit a bounded Conv lowering report"),
-                                          llvm::cl::init(true),
-                                          llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
                                 llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"),
                                 llvm::cl::init(false),
@@ -100,9 +98,15 @@ llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> pimVerifyBufferizationCopyFreedom(
+llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder(
-  "pim-verify-bufferization-copy-freedom",
+  "pim-materialize-scalar-fanout-global-order",
-  llvm::cl::desc("Run the expensive official PIM tensor-copy freedom proof before bufferization"),
+  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::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
@@ -112,12 +116,6 @@ 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));
@@ -128,6 +126,11 @@ 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() {
@@ -137,14 +140,4 @@ 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> pimVerifyBufferizationCopyFreedom;
+extern llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder;
+extern llvm::cl::opt<bool> pimTraceCommunicationMaterialization;
 
 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,6 +70,13 @@ extern llvm::cl::opt<uint64_t> pimConvStreamChunkPositions;
 
 bool hasExplicitPimCoreCount();
 void verifyExplicitPimCoreCount();
-void verifyPimPipelineStages();
+
+// 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;
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerUtils.cpp

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

src/PIM/Compiler/PimCoreProgram.cpp

@@ -17,8 +17,6 @@ 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,8 +17,6 @@ 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
-  Passes/Analyses/ONNXToSpatialVerifier.cpp
+  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
-  Passes/Transforms/ONNXToSpatialPass.cpp
+  ONNXToSpatialPass.cpp
-  Passes/Analyses/SpatialLayoutCapabilities.cpp
+  SpatialLayoutCapabilities.cpp
-  Passes/Transforms/SpatialLayoutPlanningPass.cpp
+  SpatialLayoutPlanningPass.cpp
-  Passes/Transforms/SpatialPlanLoweringPatterns.cpp
+  LowerSpatialPlansPass.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

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

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

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

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

@@ -12,30 +12,56 @@ 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::SpatialTargetResources& target,
+    const spatial::SpatialTargetInfo& 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));
-  const int64_t rowsPerLane = std::max<int64_t>(
+  plan.rowTiles = std::max<int64_t>(1, ceilDivide(problem.m, plan.tileM));
-      1, fragmentRows != 0 ? fragmentRows : target.matrixShape.rows);
+  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);
 
   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, rowsPerLane);
+    plan.laneCount = problem.batch * ceilDivide(problem.m, plan.fragmentRows);
   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/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
 
 namespace onnx_mlir {
 
@@ -13,16 +13,27 @@ 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::SpatialTargetResources& target,
+    const spatial::SpatialTargetInfo& target,
     ContractionPlanKind kind,
     int64_t laneCount = 0,
     int64_t fragmentRows = 0);

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

@@ -8,6 +8,8 @@
 
 namespace onnx_mlir {
 
+enum class ContractionOrigin { Gemm, MatMul };
+
 struct ContractionProblem {
   llvm::SmallVector<int64_t> lhsBatchShape;
   llvm::SmallVector<int64_t> rhsBatchShape;
@@ -18,9 +20,16 @@ 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 SpatFlattenPlanOp;
+class SpatGraphCompute;
-struct SpatialTargetResources;
+struct SpatialTargetInfo;
 } // 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::SpatFlattenPlanOp flattenOp,
+    spatial::SpatGraphCompute flattenOp,
-    const spatial::SpatialTargetResources& target);
+    const spatial::SpatialTargetInfo& target);
 
 mlir::LogicalResult lowerFlattenFromRowStrip(
     const RowStripPhysicalValue& input,
-    spatial::SpatFlattenPlanOp flattenOp,
+    spatial::SpatGraphCompute flattenOp,
-    const spatial::SpatialTargetResources& target,
+    const spatial::SpatialTargetInfo& 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::SpatialTargetResources& target) {
+                              const spatial::SpatialTargetInfo& 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/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.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::SpatialTargetResources& target);
+  const spatial::SpatialTargetInfo& target);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

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

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp

@@ -1,32 +0,0 @@
-#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/ONNXToSpatialPass.cpp

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

src/PIM/Conversion/ONNXToSpatial/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/Passes/Analyses/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -108,9 +108,7 @@ void verifyScheduledInputs(ComputeOpTy compute,
   for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
     size_t currentInputIndex = inputIndex;
     Operation* definingOp = input.getDefiningOp();
-    if (allowChannelReceiveInputs
+    if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
-        && isa_and_nonnull<spatial::SpatChannelReceiveOp,
-                           spatial::SpatHostWaitLoadOp>(definingOp))
       continue;
     if (isScheduledPhase1Value(input))
       continue;
@@ -149,7 +147,6 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatBiasAddPlanOp,
             spatial::SpatAddPlanOp,
             spatial::SpatConcatPlanOp,
-            spatial::SpatFlattenPlanOp,
             spatial::SpatReluPlanOp,
             spatial::SpatSiluPlanOp,
             spatial::SpatResizeNearestPlanOp,
@@ -165,8 +162,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
       });
       continue;
     }
-    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp,
+    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp>(&op)) {
-            spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
         illegalOp->emitOpError() << kPhaseMarker
                                  << " explicit channel communication is not expected before merge materialization";
@@ -185,8 +181,7 @@ 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,
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
-            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/Transforms/LowerSpatialPlansPass.cpp

@@ -1,136 +0,0 @@
-#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/SpatialLayoutPlanningPass.cpp

@@ -1,318 +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/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.hpp

@@ -1,20 +0,0 @@
-#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,12 +9,11 @@ void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { popula
 
 void populateConversionPatterns(RewritePatternSet& patterns,
                                 MLIRContext* ctx,
-                                const spatial::SpatialTargetResources& target,
+                                const spatial::SpatialTargetInfo& target) {
-                                const ONNXToSpatialPlanningOptions& options) {
   populateElementwisePatterns(patterns, ctx);
   populateMatMulRewritePatterns(patterns, ctx, target);
   populateGemmPatterns(patterns, ctx, target);
-  populateConvPatterns(patterns, ctx, target, options);
+  populateConvPatterns(patterns, ctx, target);
   populatePoolPatterns(patterns, ctx, target);
   populateReduceMeanPatterns(patterns, ctx);
   populateReluPatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -4,20 +4,18 @@
 #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 SpatialTargetResources;
+struct SpatialTargetInfo;
 }
 
 void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateConversionPatterns(mlir::RewritePatternSet& patterns,
                                 mlir::MLIRContext* ctx,
-                                const spatial::SpatialTargetResources& target,
+                                const spatial::SpatialTargetInfo& target);
-                                const ONNXToSpatialPlanningOptions& options);
 void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
@@ -25,22 +23,21 @@ void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::ML
 
 void populateConvPatterns(mlir::RewritePatternSet& patterns,
                           mlir::MLIRContext* ctx,
-                          const spatial::SpatialTargetResources& target,
+                          const spatial::SpatialTargetInfo& 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::SpatialTargetResources& target);
+                          const spatial::SpatialTargetInfo& target);
 void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns,
                                    mlir::MLIRContext* ctx,
-                                   const spatial::SpatialTargetResources& target);
+                                   const spatial::SpatialTargetInfo& target);
 void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns,
                                   mlir::MLIRContext* ctx,
-                                  const spatial::SpatialTargetResources& target);
+                                  const spatial::SpatialTargetInfo& target);
 void populatePoolPatterns(mlir::RewritePatternSet& patterns,
                           mlir::MLIRContext* ctx,
-                          const spatial::SpatialTargetResources& target);
+                          const spatial::SpatialTargetInfo& 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,7 +1,6 @@
 #include "ConvGeometry.hpp"
 
 #include <algorithm>
-#include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 
@@ -9,22 +8,12 @@ namespace onnx_mlir {
 
 namespace {
 
-static const ONNXToSpatialPlanningOptions& defaultPlanningOptions() {
+static int64_t ceilDivide(int64_t value, int64_t divisor) {
-  static const ONNXToSpatialPlanningOptions options {
+  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
-    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;
 }
@@ -42,8 +31,20 @@ void classifyConvProblem(ConvProblem& problem) {
 }
 
 ConvGeometry buildConvGeometry(const ConvProblem& problem,
-                               const spatial::SpatialTargetResources& target) {
+                               const spatial::SpatialTargetInfo& 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,
@@ -51,6 +52,11 @@ 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));
@@ -58,76 +64,216 @@ ConvGeometry buildConvGeometry(const ConvProblem& problem,
 }
 
 static ConvMaterializationKind getMaterializationKind(
-    spatial::ConvLoweringStrategy strategy) {
+    const ConvProblem& problem, spatial::ConvLoweringStrategy strategy) {
+  if (strategy == spatial::ConvLoweringStrategy::Depthwise)
+    return ConvMaterializationKind::StructuredDepthwise;
+  if (problem.isPointwise)
+    return ConvMaterializationKind::PointwiseContraction;
   switch (strategy) {
   case spatial::ConvLoweringStrategy::Depthwise:
     return ConvMaterializationKind::StructuredDepthwise;
-  case spatial::ConvLoweringStrategy::Legacy:
   case spatial::ConvLoweringStrategy::PackedIm2Col:
+  case spatial::ConvLoweringStrategy::Legacy:
     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 ONNXToSpatialPlanningOptions& options) {
+                             const spatial::SpatialTargetInfo& target) {
   return fitsSingleCrossbar(geo) && geo.pack >= 2
-    && geo.im2colElements <= options.convIm2colMaxElements;
+    && 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);
 }
 
 mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
                                        spatial::ConvLoweringStrategy strategy,
-                                       const spatial::SpatialTargetResources& target,
+                                       const spatial::SpatialTargetInfo& 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 plan();
+    return buildLegacyCandidate(problem, target);
   case spatial::ConvLoweringStrategy::Depthwise:
-    return ifApplicable(problem.isDepthwise);
+    return buildDepthwiseCandidate(problem, target);
   case spatial::ConvLoweringStrategy::PackedIm2Col:
-    return ifApplicable(fitsPackedIm2Col(geo, options));
+    return buildPackedIm2ColCandidate(problem, target);
   case spatial::ConvLoweringStrategy::StreamedPatch:
-    return ifApplicable(fitsSingleCrossbar(geo));
+    return buildStreamedPatchCandidate(problem, target);
   case spatial::ConvLoweringStrategy::StreamedPacked:
-    return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2);
+    return buildStreamedPackedCandidate(problem, target);
   case spatial::ConvLoweringStrategy::OutputChannelTiled:
-    return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize);
+    return buildOutputChannelTiledCandidate(problem, target);
   case spatial::ConvLoweringStrategy::InputKTiled:
-    return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize);
+    return buildInputKTiledCandidate(problem, target);
   case spatial::ConvLoweringStrategy::Tiled2D:
-    return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize);
+    return buildTiled2DCandidate(problem, target);
   }
   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 ONNXToSpatialPlanningOptions& options) {
+                                    const spatial::SpatialTargetInfo& target) {
   const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
-  uint64_t chunkPositions = std::max<uint64_t>(1, options.convIm2colMaxElements / patchElements);
+  uint64_t chunkPositions = std::max<uint64_t>(1, target.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, options.convStreamChunkPositions));
+  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, target.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/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 #include <cstdint>
@@ -45,20 +45,29 @@ struct ConvProblem {
   bool isPointwise = false;
 };
 
-struct ConvLoweringState {
+struct ConvLoweringState : ConvProblem {
-  ConvProblem problem;
   mlir::Operation* diagnosticAnchor = nullptr;
   mlir::Value x;
   mlir::Value w;
   mlir::Value b;
-  const spatial::SpatialTargetResources* target = nullptr;
+  const spatial::SpatialTargetInfo* target = nullptr;
-  const ONNXToSpatialPlanningOptions* options = nullptr;
 
-  const spatial::SpatialTargetResources& targetInfo() const { return *target; }
+  const spatial::SpatialTargetInfo& 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;
@@ -66,6 +75,8 @@ struct ConvGeometry {
   int64_t matrixUnitsPerProcessor;
   int64_t pack;
   uint64_t im2colElements;
+  bool hasBias;
+  bool isDepthwise;
 };
 
 struct RowInterval {
@@ -83,14 +94,35 @@ struct ConvRowDemand {
 
 enum class ConvMaterializationKind : uint8_t {
   StructuredDepthwise,
+  PointwiseContraction,
   PackedIm2Col,
   StreamedPatch,
   StreamedPacked,
+  OutputChannelTiled,
   InputKTiled,
+  Tiled2D,
 };
 
 struct ConvPlan {
-  ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col;
+  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;
 };
 
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
@@ -98,16 +130,18 @@ bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOu
 void classifyConvProblem(ConvProblem& problem);
 
 ConvGeometry buildConvGeometry(const ConvProblem& problem,
-                               const spatial::SpatialTargetResources& target);
+                               const spatial::SpatialTargetInfo& target);
 
 mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
                                        spatial::ConvLoweringStrategy strategy,
-                                       const spatial::SpatialTargetResources& target,
+                                       const spatial::SpatialTargetInfo& target);
-                                       const ONNXToSpatialPlanningOptions& options);
+
+llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target);
 
 uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
                                     int64_t packFactor,
-                                    const ONNXToSpatialPlanningOptions& options);
+                                    const spatial::SpatialTargetInfo& target);
 
 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 = op.getA().getDefiningOp<spatial::SpatBlueprintOp>();
+    auto blueprint = adaptor.getA().getDefiningOp<spatial::SpatBlueprintOp>();
     Value scalar = adaptor.getB();
     if (!blueprint) {
-      blueprint = op.getB().getDefiningOp<spatial::SpatBlueprintOp>();
+      blueprint = adaptor.getB().getDefiningOp<spatial::SpatBlueprintOp>();
       scalar = adaptor.getA();
     }
     auto scalarAttr = getDenseConstantAttr(scalar);

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

@@ -23,6 +23,7 @@
 #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"
@@ -328,8 +329,6 @@ 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);
 }
@@ -661,14 +660,14 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
 }
 
 struct GemmToSpatialComputes : OpConversionPattern<ONNXGemmOp> {
-  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
       : OpConversionPattern<ONNXGemmOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
                                 ONNXGemmOpAdaptor gemmOpAdaptor,
                                 ConversionPatternRewriter& rewriter) const override;
 
-  const spatial::SpatialTargetResources& target;
+  const spatial::SpatialTargetInfo& target;
 };
 
 } // namespace
@@ -683,7 +682,7 @@ FailureOr<Value> lowerGemmToSpatial(
     bool transB,
     float alpha,
     float beta,
-    const spatial::SpatialTargetResources& target,
+    const spatial::SpatialTargetInfo& target,
     PatternRewriter& rewriter,
     Location loc) {
   auto aType = dyn_cast<RankedTensorType>(a.getType());
@@ -732,9 +731,14 @@ 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)) {
@@ -785,19 +789,13 @@ FailureOr<Value> lowerGemmToSpatial(
     auto bShape = bType.getShape();
     auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
     if (isCompileTimeComputable(b)) {
-      auto denseAttr = getHostConstDenseElementsAttr(b);
+      auto transposedConstant = materializeTransposedContractionConstant(
-      auto inputType = denseAttr ? dyn_cast<RankedTensorType>(denseAttr.getType()) : nullptr;
+          b, transposedType, {1, 0}, rewriter, loc);
-      auto transposedAttr = inputType && inputType.hasStaticShape() && transposedType.hasStaticShape()
+      if (failed(transposedConstant)) {
-        ? 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 = getOrCreateConstant(rewriter,
+      b = *transposedConstant;
-                              rewriter.getInsertionBlock()->getParentOp(),
-                              *transposedAttr,
-                              transposedType);
     } else {
       b = createLinalgTranspose(b, transposedType, {1, 0}, rewriter, loc);
     }
@@ -833,7 +831,7 @@ FailureOr<Value> lowerGemmToSpatial(
   }
   b = *paddedB;
   auto paddedAType = RankedTensorType::get({problem.m, paddedReductionSize}, aType.getElementType());
-  a = createPaddedInputCompute(a, paddedAType, rewriter, loc);
+  a = materializePaddedContractionInput(a, paddedAType, rewriter, loc);
   aType = paddedAType;
 
   Value bias;
@@ -898,7 +896,7 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
 
 void populateGemmPatterns(RewritePatternSet& patterns,
                           MLIRContext* ctx,
-                          const spatial::SpatialTargetResources& target) {
+                          const spatial::SpatialTargetInfo& target) {
   patterns.insert<GemmToSpatialComputes>(ctx, target);
 }
 

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

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

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

@@ -12,10 +12,9 @@
 #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"
@@ -363,16 +362,16 @@ static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Locati
     return input;
   auto type = cast<RankedTensorType>(value.getType());
   auto shape = type.getShape();
-  auto createTranspose = [&](RankedTensorType resultType, ArrayRef<int64_t> permutation) {
+  auto createONNXTranspose = [&](RankedTensorType resultType, ArrayRef<int64_t> permutation) {
-    return createLinalgTranspose(value, resultType, permutation, rewriter, loc);
+    return ONNXTransposeOp::create(rewriter, loc, resultType, value, rewriter.getI64ArrayAttr(permutation)).getResult();
   };
   if (type.getRank() == 2) {
     auto resultType = RankedTensorType::get({shape[1], shape[0]}, type.getElementType(), type.getEncoding());
-    return createTranspose(resultType, {1, 0});
+    return createONNXTranspose(resultType, {1, 0});
   }
 
   auto resultType = RankedTensorType::get({shape[0], shape[2], shape[1]}, type.getElementType(), type.getEncoding());
-  return createTranspose(resultType, {0, 2, 1});
+  return createONNXTranspose(resultType, {0, 2, 1});
 }
 
 static FailureOr<Value> materializePaddedBatchedWeight(Value value,
@@ -836,7 +835,7 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
   return computeOp->getResult(0);
 }
 
-struct NormalizedMatMulInfo {
+struct NormalizedMatMulInfo : ContractionProblem {
   NormalizedMatMulInfo(RankedTensorType lhsType,
                        RankedTensorType rhsType,
                        RankedTensorType outType,
@@ -845,7 +844,7 @@ struct NormalizedMatMulInfo {
                        ContractionProblem problem,
                        bool lhsWasVector,
                        bool rhsWasVector)
-      : problem(std::move(problem)),
+      : ContractionProblem(std::move(problem)),
         lhsType(lhsType),
         rhsType(rhsType),
         outType(outType),
@@ -854,7 +853,6 @@ struct NormalizedMatMulInfo {
         lhsWasVector(lhsWasVector),
         rhsWasVector(rhsWasVector) {}
 
-  ContractionProblem problem;
   RankedTensorType lhsType;
   RankedTensorType rhsType;
   RankedTensorType outType;
@@ -864,15 +862,14 @@ struct NormalizedMatMulInfo {
   bool rhsWasVector;
 };
 
-struct MatMulLoweringPlan {
+struct MatMulLoweringPlan : ContractionProblem {
   MatMulLoweringPlan(Value lhs, Value rhs, const NormalizedMatMulInfo& info)
-      : problem(info.problem),
+      : ContractionProblem(info),
         lhs(lhs),
         rhs(rhs),
         lhsType(cast<RankedTensorType>(lhs.getType())),
         rhsType(cast<RankedTensorType>(rhs.getType())) {}
 
-  ContractionProblem problem;
   Value lhs;
   Value rhs;
   RankedTensorType lhsType;
@@ -955,9 +952,14 @@ static FailureOr<NormalizedMatMulInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp)
                         m,
                         k,
                         n,
+                        ContractionOrigin::MatMul,
                         lhsType.getElementType(),
                         rhsType.getElementType(),
-                        outType.getElementType()},
+                        outType.getElementType(),
+                        false,
+                        false,
+                        lhsWasVector,
+                        rhsWasVector},
     lhsWasVector,
     rhsWasVector);
 }
@@ -977,10 +979,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.problem.lhsBatchShape, plan.problem.rhsBatchShape);
+  std::swap(plan.lhsBatchShape, plan.rhsBatchShape);
-  std::swap(plan.problem.lhsBatch, plan.problem.rhsBatch);
+  std::swap(plan.lhsBatch, plan.rhsBatch);
-  plan.problem.m = info.problem.n;
+  plan.m = info.n;
-  plan.problem.n = info.problem.m;
+  plan.n = info.m;
   plan.transposedResult = true;
   return plan;
 }
@@ -1004,17 +1006,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.problem.outputBatchShape.size() > 1) {
+  if (info.outputBatchShape.size() > 1) {
-    SmallVector<int64_t> expandedShape(info.problem.outputBatchShape.begin(), info.problem.outputBatchShape.end());
+    SmallVector<int64_t> expandedShape(info.outputBatchShape.begin(), info.outputBatchShape.end());
-    expandedShape.push_back(info.problem.m);
+    expandedShape.push_back(info.m);
-    expandedShape.push_back(info.problem.n);
+    expandedShape.push_back(info.n);
     auto expandedType = RankedTensorType::get(expandedShape, info.outType.getElementType(), info.outType.getEncoding());
-    result = expandBatchDims(result, expandedType, info.problem.outputBatchShape.size(), rewriter, loc);
+    result = expandBatchDims(result, expandedType, info.outputBatchShape.size(), rewriter, loc);
     currentType = expandedType;
   }
 
   SmallVector<bool> removedAxes(currentType.getRank(), false);
-  if (info.problem.outputBatchShape.empty())
+  if (info.outputBatchShape.empty())
     removedAxes[0] = true;
   if (info.lhsWasVector)
     removedAxes[currentType.getRank() - 2] = true;
@@ -1024,77 +1026,79 @@ static Value finalizeNormalizedMatMulResult(Value value,
 }
 
 struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
-  explicit MatMulToGemm(MLIRContext* ctx,
+  using OpRewritePattern::OpRewritePattern;
-                        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->problem.outputBatchShape.empty())
+    if (!shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
     bool useTransposedForm = isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
 
-    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->problem.lhsBatch, shapeInfo->problem.m, shapeInfo->problem.k, rewriter, loc);
+    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
-    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->problem.rhsBatch, shapeInfo->problem.k, shapeInfo->problem.n, rewriter, loc);
+    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
-    int64_t lhsBatchForGemm = shapeInfo->problem.lhsBatch;
+    int64_t lhsBatchForGemm = shapeInfo->lhsBatch;
-    int64_t rhsBatchForGemm = shapeInfo->problem.rhsBatch;
+    int64_t rhsBatchForGemm = shapeInfo->rhsBatch;
-    int64_t gemmM = shapeInfo->problem.m;
+    int64_t gemmM = shapeInfo->m;
-    int64_t gemmK = shapeInfo->problem.k;
+    int64_t gemmK = shapeInfo->k;
-    int64_t gemmN = shapeInfo->problem.n;
+    int64_t gemmN = shapeInfo->n;
     if (useTransposedForm) {
       lhs = transposeLastTwoDims(matmulOp.getB(), rewriter, loc);
-      lhsBatchForGemm = shapeInfo->problem.rhsBatch;
+      lhsBatchForGemm = shapeInfo->rhsBatch;
       rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc);
-      rhsBatchForGemm = shapeInfo->problem.lhsBatch;
+      rhsBatchForGemm = shapeInfo->lhsBatch;
-      gemmM = shapeInfo->problem.n;
+      gemmM = shapeInfo->n;
-      gemmN = shapeInfo->problem.m;
+      gemmN = shapeInfo->m;
     }
 
-    auto gemmType = RankedTensorType::get(
+    auto gemmType = RankedTensorType::get({gemmM, gemmN}, shapeInfo->outType.getElementType());
-        {gemmM, gemmN}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+    Value none = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());
     Value lhsMatrix = extractBatchMatrix(lhs, /*batchIndex=*/0, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
     Value rhsMatrix = extractBatchMatrix(rhs, /*batchIndex=*/0, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
-    FailureOr<Value> gemmResult = lowerGemmToSpatial(
+    Value gemmResult = ONNXGemmOp::create(rewriter,
-        matmulOp.getOperation(), lhsMatrix, rhsMatrix, Value(), gemmType,
+                                          loc,
-        /*transA=*/false, /*transB=*/false, /*alpha=*/1.0f, /*beta=*/1.0f,
+                                          gemmType,
-        target, rewriter, loc);
+                                          lhsMatrix,
-    if (failed(gemmResult))
+                                          rhsMatrix,
-      return failure();
+                                          none,
-    Value result = *gemmResult;
+                                          rewriter.getF32FloatAttr(1.0f),
+                                          rewriter.getF32FloatAttr(1.0f),
+                                          rewriter.getBoolAttr(false),
+                                          rewriter.getBoolAttr(false))
+                         .getY();
     if (useTransposedForm)
-      result = transposeLastTwoDims(result, rewriter, loc);
+      gemmResult =
+        ONNXTransposeOp::create(rewriter, loc, shapeInfo->outType, gemmResult, rewriter.getI64ArrayAttr({1, 0}))
+          .getResult();
 
-    if (shapeInfo->problem.outputBatchShape.empty()) {
+    if (shapeInfo->outputBatchShape.empty()) {
-      rewriter.replaceOp(matmulOp, result);
+      rewriter.replaceOp(matmulOp, gemmResult);
       return success();
     }
 
     auto directOutType =
-      RankedTensorType::get({1, shapeInfo->problem.m, shapeInfo->problem.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+      RankedTensorType::get({1, shapeInfo->m, shapeInfo->n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
-    Value batchedResult = ensureBatchedTensor(result, /*batchSize=*/1, shapeInfo->problem.m, shapeInfo->problem.n, rewriter, loc);
+    Value batchedResult = ensureBatchedTensor(gemmResult, /*batchSize=*/1, shapeInfo->m, shapeInfo->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::SpatialTargetResources& target)
+                                           const spatial::SpatialTargetInfo& 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->problem.outputBatchShape.empty())
+    if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
@@ -1110,63 +1114,62 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     const bool rhsStoredAsRows = rhsRows && !useTransposedForm;
 
     Value lhs =
-      normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType,
+      normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType, shapeInfo->lhsWasVector, rewriter, loc);
-                             shapeInfo->lhsWasVector, rewriter, loc);
     Value rhs = normalizeMatMulOperand(
-      rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType,
+      rhsStoredAsRows ? rhsRows : matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc);
-      shapeInfo->rhsWasVector, rewriter, loc);
+    lhs = collapseBatchDims(lhs, shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
-    lhs = collapseBatchDims(lhs, shapeInfo->problem.lhsBatch, shapeInfo->problem.m, shapeInfo->problem.k, rewriter, loc);
     rhs = collapseBatchDims(rhs,
-                            shapeInfo->problem.rhsBatch,
+                            shapeInfo->rhsBatch,
-                            rhsStoredAsRows ? shapeInfo->problem.n : shapeInfo->problem.k,
+                            rhsStoredAsRows ? shapeInfo->n : shapeInfo->k,
-                            rhsStoredAsRows ? shapeInfo->problem.k : shapeInfo->problem.n,
+                            rhsStoredAsRows ? shapeInfo->k : shapeInfo->n,
                             rewriter,
                             loc);
     MatMulLoweringPlan plan = buildLoweringPlan(
       lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc);
 
-    plan.lhs = ensureBatchedTensor(plan.lhs, plan.problem.lhsBatch, plan.problem.m, plan.problem.k, rewriter, loc);
+    plan.lhs = ensureBatchedTensor(plan.lhs, plan.lhsBatch, plan.m, plan.k, rewriter, loc);
     plan.rhs = ensureBatchedTensor(plan.rhs,
-                                   plan.problem.rhsBatch,
+                                   plan.rhsBatch,
-                                   rhsStoredAsRows ? plan.problem.n : plan.problem.k,
+                                   rhsStoredAsRows ? plan.n : plan.k,
-                                   rhsStoredAsRows ? plan.problem.k : plan.problem.n,
+                                   rhsStoredAsRows ? plan.k : plan.n,
                                    rewriter,
                                    loc);
     plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
     plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
     auto directOutType = RankedTensorType::get(
-      {plan.problem.batch, plan.problem.m, plan.problem.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+      {plan.batch, plan.m, plan.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
 
     if (isCompileTimeComputable(plan.rhs)) {
       ContractionPlan contractionPlan = makeContractionPlan(
-        plan.problem, target, ContractionPlanKind::StaticTiled);
+        plan, 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.problem.lhsBatch, plan.problem.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
+        {plan.lhsBatch, plan.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
       auto paddedRhsType = RankedTensorType::get(
-        {plan.problem.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding());
+        {plan.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding());
       auto paddedOutType =
-        RankedTensorType::get({plan.problem.batch, plan.problem.m, paddedOutCols}, shapeInfo->outType.getElementType());
+        RankedTensorType::get({plan.batch, plan.m, paddedOutCols}, shapeInfo->outType.getElementType());
 
       auto paddedRhs =
-        materializePaddedBatchedWeight(plan.rhs, plan.problem.rhsBatchShape, plan.problem.outputBatchShape, paddedRhsType, rewriter);
+        materializePaddedBatchedWeight(plan.rhs, plan.rhsBatchShape, plan.outputBatchShape, paddedRhsType, rewriter);
       if (succeeded(paddedRhs)) {
-        Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
+        Value paddedLhs = materializePaddedContractionInput(
+          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.problem.lhsBatchShape,
+                                             plan.lhsBatchShape,
                                              paddedRhsType,
-                                             plan.problem.rhsBatchShape,
+                                             plan.rhsBatchShape,
-                                             plan.problem.outputBatchShape,
+                                             plan.outputBatchShape,
                                              partialPiecesType,
-                                             plan.problem.m,
+                                             plan.m,
                                              numKSlices,
                                              numOutHSlices,
                                              xbarSize,
@@ -1178,7 +1181,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                                     partialPiecesType,
                                                     directOutType,
                                                     paddedOutType,
-                                                    plan.problem.batch,
+                                                    plan.batch,
                                                     numKSlices,
                                                     xbarSize,
                                                     rewriter,
@@ -1187,7 +1190,12 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
           return failure();
         Value finalResult = *result;
         if (plan.transposedResult) {
-          finalResult = transposeLastTwoDims(finalResult, rewriter, loc);
+          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 = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
         rewriter.replaceOp(matmulOp, finalResult);
@@ -1197,24 +1205,24 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     RankedTensorType blueprintType = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
                                    ? shapeInfo->outType : directOutType;
     SmallVector<int64_t> blueprintBatchShape = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
-                                             ? shapeInfo->problem.outputBatchShape : SmallVector<int64_t> {plan.problem.batch};
+                                             ? shapeInfo->outputBatchShape : SmallVector<int64_t> {plan.batch};
-    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.problem.m, plan.problem.k, plan.problem.n, xbarSize);
+    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.m, plan.k, plan.n, xbarSize);
     ContractionPlan contractionPlan = makeContractionPlan(
-      plan.problem, target, ContractionPlanKind::GroupedRowDynamicVVD,
+      plan, target, ContractionPlanKind::GroupedRowDynamicVVD,
-      /*laneCount=*/plan.problem.batch * plan.problem.m / rowsPerLane, rowsPerLane);
+      /*laneCount=*/plan.batch * plan.m / rowsPerLane, rowsPerLane);
     const int64_t laneCount = contractionPlan.laneCount;
     SmallVector<int64_t> fragmentShape(blueprintType.getRank(), 1);
     fragmentShape[fragmentShape.size() - 2] = rowsPerLane;
-    fragmentShape.back() = plan.problem.n;
+    fragmentShape.back() = plan.n;
     auto fragmentType = RankedTensorType::get(fragmentShape, shapeInfo->outType.getElementType());
     auto rowPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount, fragmentType);
     Value transposedRhs = rhsStoredAsRows ? plan.rhs : transposeLastTwoDims(plan.rhs, rewriter, loc);
     auto batchOp = createBatchedVvdmulBatch(plan.lhs,
-                                            plan.problem.lhsBatchShape,
+                                            plan.lhsBatchShape,
                                             transposedRhs,
-                                            plan.problem.rhsBatchShape,
+                                            plan.rhsBatchShape,
-                                            plan.problem.outputBatchShape,
+                                            plan.outputBatchShape,
-                                            plan.problem.k,
+                                            plan.k,
                                             rowsPerLane,
                                             rowPiecesType,
                                             directOutType,
@@ -1223,33 +1231,23 @@ 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(
-      rowPieces, blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc);
+      batchOp->getResult(0), blueprintType, blueprintBatchShape, rowsPerLane, rewriter, loc);
     if (failed(result))
       return failure();
     Value finalResult = *result;
     if (plan.transposedResult) {
-      finalResult = transposeLastTwoDims(finalResult, rewriter, loc);
+      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 = 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);
@@ -1258,7 +1256,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     return success();
   }
 
-  const spatial::SpatialTargetResources& target;
+  const spatial::SpatialTargetInfo& target;
 };
 
 struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
@@ -1275,14 +1273,14 @@ struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
 
 void populateMatMulFusionPatterns(RewritePatternSet& patterns,
                                   MLIRContext* ctx,
-                                  const spatial::SpatialTargetResources& target) {
+                                  const spatial::SpatialTargetInfo& target) {
   patterns.add<TransposedRhsMatMulToSpatial>(ctx, target);
 }
 
 void populateMatMulRewritePatterns(RewritePatternSet& patterns,
                                    MLIRContext* ctx,
-                                   const spatial::SpatialTargetResources& target) {
+                                   const spatial::SpatialTargetInfo& target) {
-  patterns.insert<MatMulToGemm>(ctx, target);
+  patterns.insert<MatMulToGemm>(ctx);
   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/Passes/Transforms/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/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::SpatialTargetResources& target)
+  PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
       : OpConversionPattern<PoolOp>(ctx), target(target) {}
 
-  const spatial::SpatialTargetResources& target;
+  const spatial::SpatialTargetInfo& 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::SpatialTargetResources&) {
+                                            const spatial::SpatialTargetInfo&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -447,10 +447,9 @@ LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
 }
 
 FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
-                                         Value input,
+                                         const spatial::SpatialTargetInfo& target,
-                                         const spatial::SpatialTargetResources& target,
                                          PatternRewriter& rewriter) {
-  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
       || inputType.getRank() != 4 || outputType.getRank() != 4)
@@ -477,7 +476,7 @@ FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
 
   auto compute = createSpatCompute<1>(
-    rewriter, planOp.getLoc(), outputType, {}, input,
+    rewriter, planOp.getLoc(), outputType, {}, planOp.getInput(),
     [&](Value input) -> LogicalResult {
       Value paddedInput = createPaddedPoolInput(
         rewriter, planOp.getLoc(), input, inputType,
@@ -616,15 +615,14 @@ static Value extractPoolIndex(PatternRewriter& rewriter,
 }
 
 FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
-                                            Value input,
                                             std::optional<Value> rowStripInput,
-                                            const spatial::SpatialTargetResources& target,
+                                            const spatial::SpatialTargetInfo& target,
                                             PatternRewriter& rewriter) {
   if (failed(canLowerMaxPoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
-  auto inputType = cast<RankedTensorType>(input.getType());
+  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   const int64_t channels = inputType.getDimSize(1);
   const int64_t inputHeight = inputType.getDimSize(2);
@@ -633,9 +631,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 actualInput = rowStripInput.value_or(input);
+  Value input = rowStripInput.value_or(planOp.getInput());
-  auto actualInputType = dyn_cast<RankedTensorType>(actualInput.getType());
+  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(actualInput, inputType);
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
   const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
@@ -683,7 +681,7 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
     TypeRange {outputStorageType},
     outputHeight * tilesPerRow,
     {},
-    ValueRange {actualInput},
+    ValueRange {input},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       SmallVector<Value> inputRows;
       inputRows.reserve(kernelHeight);
@@ -808,7 +806,7 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
 }
 
 LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
-  spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetResources&) {
+  spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetInfo&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -822,10 +820,9 @@ LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
 
 FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
     spatial::SpatGlobalAveragePoolPlanOp planOp,
-    Value input,
+    const spatial::SpatialTargetInfo& target,
-    const spatial::SpatialTargetResources& target,
     PatternRewriter& rewriter) {
-  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape()
       || !outputType.hasStaticShape() || inputType.getRank() != 4
@@ -847,7 +844,7 @@ FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
   const double scaleValue = 1.0 / static_cast<double>(height * width);
 
   auto compute = createSpatCompute<1>(
-    rewriter, planOp.getLoc(), outputType, {}, input,
+    rewriter, planOp.getLoc(), outputType, {}, planOp.getInput(),
     [&](Value input) -> LogicalResult {
       Value output = tensor::EmptyOp::create(
         rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
@@ -898,23 +895,22 @@ FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
 }
 
 FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                                    Value input,
                                                     std::optional<Value> rowStripInput,
-                                                    const spatial::SpatialTargetResources& target,
+                                                    const spatial::SpatialTargetInfo& target,
                                                     PatternRewriter& rewriter) {
   if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
-  auto inputType = cast<RankedTensorType>(input.getType());
+  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   auto elementType = dyn_cast<FloatType>(inputType.getElementType());
   if (!elementType)
     return failure();
 
-  Value actualInput = rowStripInput.value_or(input);
+  Value input = rowStripInput.value_or(planOp.getInput());
-  auto actualInputType = dyn_cast<RankedTensorType>(actualInput.getType());
+  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(actualInput, inputType);
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
   const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
@@ -944,7 +940,7 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
     TypeRange {outputStorageType},
     tilesPerRow,
     ValueRange {zero, scale},
-    ValueRange {actualInput},
+    ValueRange {input},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       Value reduced = args.weights[0];
       for (int64_t row = 0; row < height; ++row) {
@@ -1015,7 +1011,7 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
 
 void populatePoolPatterns(RewritePatternSet& patterns,
                           MLIRContext* ctx,
-                          const spatial::SpatialTargetResources& target) {
+                          const spatial::SpatialTargetInfo& 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/Passes/Transforms/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -47,11 +47,11 @@ static SmallVector<ReassociationIndices> getExpandFrom1DReassociation(int64_t ra
   return reassociation;
 }
 
-static Value buildFlattenBody(Value input,
+static Value buildFlatten(Value input,
                           RankedTensorType sourceType,
                           RankedTensorType resultType,
                           int64_t axis,
-                              PatternRewriter& rewriter,
+                          ConversionPatternRewriter& rewriter,
                           Location loc) {
   if (sourceType == resultType)
     return input;
@@ -75,25 +75,6 @@ static Value buildFlattenBody(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;
 
@@ -116,17 +97,21 @@ struct Flatten : OpConversionPattern<ONNXFlattenOp> {
     if (resultType.getShape()[0] != outerDim || resultType.getShape()[1] != innerDim)
       return failure();
 
-    auto plan = spatial::SpatFlattenPlanOp::create(
+    auto replaceWithFlatten = [&](auto build) -> LogicalResult {
-        rewriter, flattenOp.getLoc(), resultType, adaptor.getInput(),
+      Value flattened = materializeOrComputeUnary(adaptor.getInput(), resultType, rewriter, flattenOp.getLoc(), build);
-        rewriter.getI64IntegerAttr(*axis),
+      rewriter.replaceOp(flattenOp, flattened);
-        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;
@@ -134,11 +119,12 @@ struct RowStripFlattenAnalysis {
 };
 
 static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
-    spatial::SpatFlattenPlanOp flattenOp, const spatial::SpatialTargetResources& target) {
+  spatial::SpatGraphCompute flattenOp, const spatial::SpatialTargetInfo& target) {
-  if (flattenOp.getAxis() != 1)
+  if (flattenOp.getWeights().size() != 0 || flattenOp.getInputs().size() != 1
+      || flattenOp.getOutputs().size() != 1)
     return failure();
-  auto sourceType = dyn_cast<RankedTensorType>(flattenOp.getInput().getType());
+  auto sourceType = dyn_cast<RankedTensorType>(flattenOp.getInputs().front().getType());
-  auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOutput().getType());
+  auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOutputs().front().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())
@@ -148,21 +134,17 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
   if (channels > xbarDim && channels % xbarDim != 0)
     return failure();
 
-  Value consumerInput = flattenOp.getOutput();
+  auto yieldOp = dyn_cast<spatial::SpatYieldOp>(flattenOp.getBody().front().getTerminator());
-  Operation* consumerOp = nullptr;
+  if (!yieldOp || yieldOp.getOutputs().size() != 1)
-  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 consumer = dyn_cast<spatial::SpatGraphComputeBatch>(consumerOp);
+  auto collapse = yieldOp.getOutputs().front().getDefiningOp<tensor::CollapseShapeOp>();
-  if (!consumer || consumer.getInputs().size() != 1 || consumer.getInputs().front() != consumerInput
+  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)
       || consumer.getWeights().size() != 1)
     return failure();
   auto weightType = dyn_cast<RankedTensorType>(consumer.getWeights().front().getType());
@@ -173,32 +155,21 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
   if (llvm::none_of(consumer.getBody().getOps<spatial::SpatVMMOp>(),
                     [](spatial::SpatVMMOp) { return true; }))
     return failure();
-  return RowStripFlattenAnalysis {consumer, sourceType, resultType, weightType, weight};
+  return RowStripFlattenAnalysis {consumer, collapse, sourceType, resultType, weightType, weight};
 }
 
 } // namespace
 
 void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Flatten>(ctx); }
 
-FailureOr<Value> lowerDenseFlattenPlan(spatial::SpatFlattenPlanOp planOp,
+LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp,
-                                        Value input,
+                                          const spatial::SpatialTargetInfo& target) {
-                                        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::SpatFlattenPlanOp flattenOp,
+                                       spatial::SpatGraphCompute flattenOp,
-                                       const spatial::SpatialTargetResources& target,
+                                       const spatial::SpatialTargetInfo& target,
                                        PatternRewriter& rewriter) {
   FailureOr<RowStripFlattenAnalysis> analysis = analyzeRowStripFlatten(flattenOp, target);
   if (failed(analysis))
@@ -235,20 +206,19 @@ LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,
                                               analysis->weightType);
   analysis->consumer->setOperand(0, reorderedWeight);
 
-  auto compute = createSpatCompute<1>(
+  BlockArgument flattenInput = *flattenOp.getInputArgument(0);
-      rewriter, flattenOp.getLoc(), TypeRange {analysis->resultType}, {},
+  flattenOp.getInputsMutable().assign(input.storage);
-      ValueRange {input.storage}, [&](Value storage) {
+  flattenInput.setType(storageType);
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(analysis->collapse);
   auto flatType = RankedTensorType::get(
     {storageType.getNumElements()}, storageType.getElementType(), storageType.getEncoding());
   Value flat = tensor::CollapseShapeOp::create(
-            rewriter, flattenOp.getLoc(), flatType, storage,
+    rewriter, flattenOp.getLoc(), flatType, flattenInput, getCollapseTo1DReassociation(storageType.getRank()));
-            getCollapseTo1DReassociation(storageType.getRank()));
   Value logicalInput = tensor::ExpandShapeOp::create(
-            rewriter, flattenOp.getLoc(), analysis->resultType, flat,
+    rewriter, flattenOp.getLoc(), analysis->resultType, flat, getExpandFrom1DReassociation(2));
-            getExpandFrom1DReassociation(2));
+  rewriter.replaceOp(analysis->collapse, logicalInput);
-        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/Passes/Transforms/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -193,7 +193,7 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
 
 LogicalResult canLowerResizeNearestPlanToRowStrip(
     spatial::SpatResizeNearestPlanOp planOp,
-    const spatial::SpatialTargetResources&) {
+    const spatial::SpatialTargetInfo&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   return success(inputType && outputType && inputType.hasStaticShape()
@@ -204,17 +204,16 @@ LogicalResult canLowerResizeNearestPlanToRowStrip(
 }
 
 FailureOr<Value> lowerSelectedResizeNearestPlan(
-    spatial::SpatResizeNearestPlanOp planOp, Value input,
+    spatial::SpatResizeNearestPlanOp planOp, std::optional<Value> rowStripInput,
-    std::optional<Value> rowStripInput,
+    const spatial::SpatialTargetInfo&,
-    const spatial::SpatialTargetResources&,
     PatternRewriter& rewriter) {
-  auto inputType = cast<RankedTensorType>(input.getType());
+  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   if (rowStripInput)
     return buildRowStripNearestResize(
       *rowStripInput, inputType, outputType, rewriter, planOp.getLoc());
   return buildDenseNearestResize(
-    input, inputType, outputType, rewriter, planOp.getLoc());
+    planOp.getInput(), inputType, outputType, rewriter, planOp.getLoc());
 }
 
 void populateResizePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Resize>(ctx); }

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -10,83 +10,55 @@
 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::SpatialTargetResources& target,
+                        const spatial::SpatialTargetInfo& target,
-                        const ONNXToSpatialPlanningOptions& options,
                         mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
-                                               const spatial::SpatialTargetResources& target,
+                                               const spatial::SpatialTargetInfo& target);
-                                               const ONNXToSpatialPlanningOptions* options = nullptr);
 mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
-                                                 const spatial::SpatialTargetResources& target,
+                                                 const spatial::SpatialTargetInfo& target);
-                                                 const ONNXToSpatialPlanningOptions* options = nullptr);
 
 mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
-  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetResources& target);
+  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetInfo& target);
 
 mlir::FailureOr<mlir::Value> lowerSelectedResizeNearestPlan(
   spatial::SpatResizeNearestPlanOp planOp,
-  mlir::Value input,
   std::optional<mlir::Value> rowStripInput,
-  const spatial::SpatialTargetResources& target,
+  const spatial::SpatialTargetInfo& target,
   mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
-                                                 const spatial::SpatialTargetResources& target);
+                                                 const spatial::SpatialTargetInfo& target);
 
 mlir::FailureOr<mlir::Value>
 lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
-                        mlir::Value input,
+                        const spatial::SpatialTargetInfo& target,
-                        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::SpatialTargetResources& target,
+                           const spatial::SpatialTargetInfo& target,
                            mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult
 canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                        const spatial::SpatialTargetResources& target);
+                                        const spatial::SpatialTargetInfo& target);
 
 mlir::FailureOr<mlir::Value>
 lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
-                                mlir::Value input,
+                                const spatial::SpatialTargetInfo& target,
-                                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::SpatialTargetResources& target,
+                                   const spatial::SpatialTargetInfo& target,
                                    mlir::PatternRewriter& rewriter);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutCapabilities.cpp

@@ -1,6 +1,5 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.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"
 
 using namespace mlir;
@@ -29,7 +28,7 @@ static bool hasRowStripInput(ArrayRef<PhysicalLayout> operandLayouts, unsigned i
 }
 
 SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
-    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)) {
     if (succeeded(canConsumeAndProduceRowStrip(*this, target)))
@@ -44,26 +43,8 @@ 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 SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0))
     alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
@@ -71,7 +52,7 @@ SmallVector<LayoutAlternative> SpatReluPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
-    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)) {
     LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
@@ -82,7 +63,7 @@ SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
-    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (hasRowStripInput(operandLayouts, 0)
       && succeeded(canLowerResizeNearestPlanToRowStrip(*this, target)))
@@ -91,7 +72,7 @@ SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
-    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) {
     LayoutAlternative alternative = denseAlternative(getOperation());
@@ -105,7 +86,7 @@ SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternatives(
-    const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) {
     LayoutAlternative alternative = denseAlternative(getOperation());
@@ -119,7 +100,7 @@ SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternative
 }
 
 SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
-    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   auto resultType = dyn_cast<RankedTensorType>(getOutput().getType());
   if (resultType && hasRowStripInput(operandLayouts, 0)
@@ -131,7 +112,7 @@ SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
-    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
   SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
   if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0)
       && hasRowStripInput(operandLayouts, 1))
@@ -140,7 +121,7 @@ SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
 }
 
 SmallVector<LayoutAlternative> SpatConcatPlanOp::getLayoutAlternatives(
-    const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
+    const SpatialTargetInfo&, 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/SpatialLayoutPlanningPass.cpp

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

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,10 +1,7 @@
 #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"
 
@@ -31,49 +28,6 @@ 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,7 +10,6 @@
 #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"
@@ -24,12 +23,6 @@ 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,39 +345,20 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     auto blockArg = computeOp.getInputArgument(inputIndex);
     if (!blockArg)
       return computeOp.emitOpError("expected compute input block arguments during lowering");
-    auto channelReceive = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(
+    auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(input.getDefiningOp());
-      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(
+      auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
-        rewriter, receiveOp->getLoc(), outputType);
       auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
       if (failed(sizeAttr))
         return failure();
-      Value zero = arith::ConstantIndexOp::create(
+      Value received =
-        rewriter, receiveOp->getLoc(), 0);
+        PimReceiveOp::create(
-      Value received;
+          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer,
-      if (hostWaitLoad) {
+          arith::ConstantIndexOp::create(rewriter, receiveOp.getLoc(), 0),
-        auto hostBuffer = getPipelineHostBuffer(rewriter, hostWaitLoad);
+          *sizeAttr, receiveOp.getSourceCoreId())
-        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;
@@ -402,8 +383,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     if (rangeLength(resultUses) == 1) {
       OpOperand& resultUse = *resultUses.begin();
       Operation* resultUser = resultUse.getOwner();
-      if (isa<spatial::SpatChannelSendOp,
+      if (isa<spatial::SpatChannelSendOp>(resultUser))
-              spatial::SpatHostStoreSyncOp>(resultUser))
         continue;
     }
 

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

@@ -57,29 +57,10 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   }
 };
 
-struct HostStoreSyncLowering : OpRewritePattern<spatial::SpatHostStoreSyncOp> {
+struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
   using OpRewritePattern::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatHostStoreSyncOp op, PatternRewriter& rewriter) const override {
+  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp 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();
@@ -105,11 +86,12 @@ static LogicalResult lowerReceive(
     if (failed(sizeAttr))
       return failure();
     Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
-    auto received = createReceive(outputBuffer, zero, *sizeAttr);
+    auto receive = pim::PimReceiveOp::create(
-    if (failed(received))
+      rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, *sizeAttr, op.getSourceCoreId());
-      return failure();
+    copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
+    Value received = receive.getOutput();
     if (!destinationInsert) {
-      rewriter.replaceOp(op, *received);
+      rewriter.replaceOp(op, received);
       return success();
     }
 
@@ -117,42 +99,10 @@ static LogicalResult lowerReceive(
     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();
-    });
   }
 };
 
@@ -198,9 +148,7 @@ struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
 } // namespace
 
 void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
-  patterns.add<ChannelSendLowering, ChannelReceiveLowering,
+  patterns.add<ChannelSendLowering, ChannelReceiveLowering, ExtractRowsLowering, ConcatLowering>(patterns.getContext());
-               HostStoreSyncLowering, HostWaitLoadLowering,
-               ExtractRowsLowering, ConcatLowering>(patterns.getContext());
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -859,10 +859,6 @@ 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/Passes/Analyses/ONNXToSpatialVerifier.hpp"
+#include "Conversion/ONNXToSpatial/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 "Passes/PIMPasses.h"
+#include "Pass/PIMPasses.h"
 #include "SpatialToPimPass.hpp"
 
 using namespace mlir;
@@ -66,20 +66,17 @@ createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTe
   return padOp.getResult();
 }
 
-static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter,
+static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
-                                                      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] <= crossbarSize && "vector width must fit in one crossbar");
+  assert(shape[1] <= static_cast<int64_t>(crossbarSize) && "vector width must fit in one crossbar");
 
-  if (shape[1] == crossbarSize)
+  if (shape[1] == static_cast<int64_t>(crossbarSize))
     return vector;
 
   auto paddedType = RankedTensorType::get(
-    {shape[0], crossbarSize}, vectorType.getElementType(), vectorType.getEncoding());
+    {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
   return createZeroPaddedTensor(rewriter, loc, vector, paddedType);
 }
 
@@ -87,11 +84,6 @@ 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);
@@ -126,8 +118,6 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                     spatial::SpatConcatOp,
                     spatial::SpatChannelReceiveOp,
                     spatial::SpatChannelSendOp,
-                    spatial::SpatHostStoreSyncOp,
-                    spatial::SpatHostWaitLoadOp,
                     spatial::SpatExtractRowsOp>();
 
   RewritePatternSet initialPatterns(ctx);
@@ -142,12 +132,6 @@ 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))) {
@@ -190,17 +174,6 @@ 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);
@@ -221,8 +194,6 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                             spatial::SpatConcatOp,
                             spatial::SpatChannelReceiveOp,
                             spatial::SpatChannelSendOp,
-                            spatial::SpatHostStoreSyncOp,
-                            spatial::SpatHostWaitLoadOp,
                             spatial::SpatExtractRowsOp>();
 
   SmallVector<pim::PimCoreOp> coreOps;
@@ -272,8 +243,6 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   communicationTarget.addIllegalOp<spatial::SpatConcatOp,
                                    spatial::SpatChannelReceiveOp,
                                    spatial::SpatChannelSendOp,
-                                   spatial::SpatHostStoreSyncOp,
-                                   spatial::SpatHostWaitLoadOp,
                                    spatial::SpatExtractRowsOp>();
 
   RewritePatternSet communicationPatterns(ctx);
@@ -296,16 +265,15 @@ 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>(targetResources.matrixShape.columns);
+    const int64_t xbarDim = static_cast<int64_t>(crossbarSize);
     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>(targetResources.matrixUnitsPerProcessor)
+    assert(paddedOutputWidth / xbarDim <= static_cast<int64_t>(crossbarCountInCore)
            && "output width must fit in one core");
 
     rewriter.setInsertionPoint(vmmOp);
-    auto paddedInput = padHVectorInputToCrossbarSize(
+    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
-      rewriter, vmmOp.getLoc(), vmmOp.getInput(), xbarDim);
     if (failed(paddedInput)) {
       hasFailure = true;
       return WalkResult::interrupt();
@@ -407,9 +375,4 @@ 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,7 +18,6 @@
 #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 {
@@ -29,10 +28,7 @@ struct SpatialToPimPass : mlir::PassWrapper<SpatialToPimPass, mlir::OperationPas
   llvm::StringRef getDescription() const override { return "Lower Spatial ops to PIM-ready format"; }
 
   SpatialToPimPass() = default;
-  explicit SpatialToPimPass(const spatial::SpatialTargetResources& target)
+  SpatialToPimPass(const SpatialToPimPass& pass) {}
-      : targetResources(target), hasTarget(true) {}
-  SpatialToPimPass(const SpatialToPimPass& pass)
-      : targetResources(pass.targetResources), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() final;
 
@@ -41,8 +37,6 @@ 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/Analysis/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/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;

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(Passes/Analyses)
+add_subdirectory(Analysis)
-add_subdirectory(Passes/Transforms/Bufferization)
+add_subdirectory(Transforms/Bufferization)
-add_subdirectory(Passes/Transforms/HostConstantFolding)
+add_subdirectory(Transforms/HostConstantFolding)
-add_subdirectory(Passes/Transforms/InstructionSelection)
+add_subdirectory(Transforms/InstructionSelection)
-add_subdirectory(Passes/Transforms/LocalMemoryPlanning)
+add_subdirectory(Transforms/LocalMemoryPlanning)
-add_subdirectory(Passes/Transforms/Verification)
+add_subdirectory(Transforms/Verification)
 
 add_pim_library(PimOps
   PimOps.hpp

src/PIM/Dialect/Pim/Pim.td

@@ -118,32 +118,6 @@ 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,6 +10,7 @@
 
 #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;
@@ -156,10 +157,16 @@ 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 <= 0)
+  if (vector1 != 1 || vectorWidth != xbarDim)
-    return emitError("vector shape must be (1, positive-width)");
+    return emitError("vector shape must be (1, crossbar-size)");
 
   int64_t output1 = outputShape[0];
   int64_t outputWidth = outputShape[1];

src/PIM/Dialect/Pim/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/Passes/Transforms/Bufferization/BufferizationUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/Common.hpp"
 
 using namespace mlir;
 using namespace bufferization;

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

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

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

@@ -430,7 +430,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
 
   auto targetBytes = getShapedByteSize(targetType);
   auto sourceBytes = getShapedByteSize(sourceType);
-  if (succeeded(targetBytes) && succeeded(sourceBytes) && size <= *targetBytes && size <= *sourceBytes) {
+  if (targetType.getElementType() == sourceType.getElementType() && 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/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/Passes/Transforms/Bufferization/BufferizationUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp"
 
 using namespace mlir;
 using namespace bufferization;

src/PIM/Dialect/Pim/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/Passes/Transforms/Bufferization/Common.hpp"
+#include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
-#include "Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 
 using namespace mlir;
@@ -382,51 +382,8 @@ static bufferization::OneShotBufferizationOptions makePimBufferizationOptions()
   return options;
 }
 
-static LogicalResult verifyPimCoresNeedNoTensorCopies(
+static LogicalResult preparePimBufferization(func::FuncOp funcOp) {
-    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();
 }
 
@@ -669,9 +626,6 @@ 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.";
@@ -684,12 +638,9 @@ struct PimBufferizationPreparationPass
       signalPassFailure();
       return;
     }
-    if (failed(preparePimBufferization(moduleOp, *funcOp, verifyCopyFreedom)))
+    if (failed(preparePimBufferization(*funcOp)))
       signalPassFailure();
   }
-
-private:
-  bool verifyCopyFreedom;
 };
 
 struct PimOneShotBufferizationPass
@@ -784,10 +735,6 @@ 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/Transforms/HostConstantFolding/HostConstantFoldingPass.cpp

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

src/PIM/Dialect/Pim/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/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
-#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace llvm;
 using namespace mlir;

src/PIM/Dialect/Pim/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/Passes/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning/LocalMemoryPlanning.hpp"
-#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace llvm;
 using namespace mlir;

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

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

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

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

src/PIM/Dialect/Pim/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/Dialect/Pim/Passes/Analyses/LocalMemoryLifetimeAnalysis.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Analysis/LocalMemoryLifetimeAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 using namespace mlir;
 
@@ -241,8 +241,6 @@ static bool isSupportedCoreInstructionOp(Operation* op) {
              pim::PimVMVOp,
              pim::PimReceiveOp,
              pim::PimSendOp,
-             pim::PimSyncOp,
-             pim::PimWaitOp,
              pim::PimConcatOp,
              pim::PimVMMOp,
              pim::PimVVAddOp,
@@ -750,43 +748,12 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
   }
 
   VerificationPass() {}
-  VerificationPass(const spatial::SpatialTargetResources& target,
+  VerificationPass(const VerificationPass& pass) {}
-                   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;
@@ -844,7 +811,7 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
     }
 
     bool hasFailure = false;
-    if (detectCommunicationDeadlock && failed(verifyNoStaticCommunicationDeadlock(moduleOp, diagnostics)))
+    if (pimDetectCommunicationDeadlock && failed(verifyNoStaticCommunicationDeadlock(moduleOp, diagnostics)))
       hasFailure = true;
 
     if (diagnostics.hasFailure()) {
@@ -858,10 +825,6 @@ 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) {
@@ -1087,10 +1050,4 @@ 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/Spatial/CMakeLists.txt

@@ -6,6 +6,7 @@ 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)
@@ -16,27 +17,26 @@ add_pim_library(SpatialOps
   SpatialOpsVerify.cpp
   SpatialOpsCanonicalization.cpp
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
-  Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
+  Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
-  Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
+  Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
+  Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp
+  Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp
+  Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
+  Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredResultRealization.cpp
+  Transforms/MergeComputeNodes/DeferredResultRealization.cpp
-  Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
-  Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp
+  Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp
-  Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
-  Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
+  Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
-  Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
-  Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
-  Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
+  Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
-  Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
+  Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
-  Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
+  Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
-  Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp
+  Transforms/TrivialGraphComputeMergePass.cpp
-  Passes/Transforms/TrivialGraphComputeMergePass.cpp
 
   EXCLUDE_FROM_OM_LIBS
 
@@ -52,5 +52,6 @@ add_pim_library(SpatialOps
   MLIRBufferizationTransforms
   OMMlirDialects
   OMONNXOps
+  OMPimCompilerOptions
   PimOps
 )

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

@@ -1,215 +0,0 @@
-#include "mlir/Pass/Pass.h"
-
-#include "DeferredCommunicationRealization.hpp"
-#include "ScheduledComputeReport.hpp"
-#include "ScheduledComputeVerification.hpp"
-#include "Scheduling/PipelineScheduling.hpp"
-#include "SpatialDataflowCsvExporter.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Passes/PIMPasses.h"
-
-#include <limits>
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace spatial {
-namespace {
-
-static bool hasValidTarget(const SchedulingTarget& target) {
-  return target.processorCount != 0 && target.residentWeightCapacity != 0
-      && target.transferWidthBytes != 0
-      && target.interProcessorLatencyNs.size() == target.processorCount * target.processorCount
-      && (target.processorCount == 1 || target.averageInterProcessorLatencyNs != 0);
-}
-
-static FailureOr<func::FuncOp> requireEntry(ModuleOp moduleOp) {
-  auto entry = getPimEntryFunc(moduleOp);
-  if (failed(entry)) {
-    moduleOp.emitError("failed to locate the PIM entry function during Spatial scheduling and realization");
-    return failure();
-  }
-  return *entry;
-}
-
-static SchedulingTarget getPipelineSchedulingTarget(
-    const SchedulingTarget& physicalTarget, size_t pipelineStages) {
-  if (pipelineStages == 1)
-    return physicalTarget;
-
-  SchedulingTarget schedulingTarget = physicalTarget;
-  schedulingTarget.processorCount = physicalTarget.processorCount / pipelineStages;
-  schedulingTarget.residentWeightCapacity = checkedMultiply(
-    physicalTarget.residentWeightCapacity, pipelineStages);
-  schedulingTarget.interProcessorLatencyNs.assign(
-    schedulingTarget.processorCount * schedulingTarget.processorCount, 0);
-  Cost latencySum = 0;
-  size_t pairCount = 0;
-  for (size_t source = 0; source < schedulingTarget.processorCount; ++source)
-    for (size_t destination = 0;
-         destination < schedulingTarget.processorCount; ++destination) {
-      Cost latency = physicalTarget.getInterProcessorLatencyNs(
-        source, destination);
-      schedulingTarget.interProcessorLatencyNs[
-        source * schedulingTarget.processorCount + destination] = latency;
-      if (source != destination) {
-        latencySum = checkedAdd(latencySum, latency);
-        ++pairCount;
-      }
-    }
-  schedulingTarget.averageInterProcessorLatencyNs = pairCount == 0
-    ? 0
-    : (latencySum + pairCount - 1) / pairCount;
-  return schedulingTarget;
-}
-
-struct ScheduleAndRealizeSpatialPass final
-    : PassWrapper<ScheduleAndRealizeSpatialPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleAndRealizeSpatialPass)
-
-  ScheduleAndRealizeSpatialPass() = default;
-  ScheduleAndRealizeSpatialPass(const SchedulingTarget& target,
-                                SpatialDataflowExportStage exportStage,
-                                size_t pipelineStages)
-      : target(target), exportStage(exportStage),
-        pipelineStages(pipelineStages), hasTarget(true) {}
-
-  StringRef getArgument() const override { return "schedule-and-realize-spatial"; }
-  StringRef getDescription() const override {
-    return "Schedule Spatial graph computes, verify them, and realize deferred communication.";
-  }
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-    if (!hasTarget || !hasValidTarget(target)) {
-      moduleOp.emitError("ScheduleAndRealizeSpatial requires an explicit valid target");
-      signalPassFailure();
-      return;
-    }
-    if (pipelineStages == 0 || target.processorCount % pipelineStages != 0
-        || target.residentWeightCapacity
-             > std::numeric_limits<size_t>::max() / pipelineStages) {
-      moduleOp.emitError("ScheduleAndRealizeSpatial requires valid pipeline stages and resource counts");
-      signalPassFailure();
-      return;
-    }
-    auto entry = requireEntry(moduleOp);
-    if (failed(entry)) {
-      signalPassFailure();
-      return;
-    }
-    func::FuncOp entryFunc = *entry;
-
-    SchedulingTarget schedulingTarget = getPipelineSchedulingTarget(
-      target, pipelineStages);
-    ComputeGraph scheduledGraph;
-    MergeScheduleResult schedule;
-    for (;;) {
-      MergeSchedulingAnalysis analysis(
-        entryFunc, schedulingTarget,
-        pipelineStages > 1 ? target.processorCount : 0);
-      scheduledGraph = analysis.getGraph();
-      schedule = std::move(analysis.getResult());
-      std::string pipelineError;
-      if (succeeded(applyPipelineScheduling(
-            scheduledGraph, schedule, pipelineStages, target, pipelineError)))
-        break;
-      std::string splitError;
-      if (pipelineStages == 1
-          || failed(splitPipelineWorkload(
-                scheduledGraph, schedule, pipelineStages, target, splitError))) {
-        if (!splitError.empty())
-          pipelineError = splitError;
-        moduleOp.emitError() << pipelineError;
-        signalPassFailure();
-        return;
-      }
-    }
-    PatternRewriter rewriter(moduleOp.getContext());
-    FailureOr<ScheduledComputeMaterializationResult> materialization =
-        materializeScheduledCompute(entryFunc, schedule, rewriter);
-    if (failed(materialization)) {
-      signalPassFailure();
-      return;
-    }
-    ScheduledComputeMaterializationResult materializationResult = std::move(*materialization);
-
-    if (failed(verifyMaterializedScheduleMapping(
-            entryFunc, schedule, materializationResult.peftClassPlans,
-            materializationResult.graphComputeToBlockMap,
-            materializationResult.materializedSchedules))
-        || failed(verifyDeferredTransferPhase1Invariants(entryFunc))
-        || failed(verifyScheduledMaterializationRecords(materializationResult.materializedSchedules))
-        || failed(verifyScheduledSpatialInvariants(entryFunc))) {
-      moduleOp.emitError("scheduled Spatial phase verification failed");
-      signalPassFailure();
-      return;
-    }
-
-    dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true);
-    if (shouldExportSpatialDataflowStage(exportStage, SpatialDataflowExportStage::Spatial3)
-        && failed(exportSpatialDataflowCsvScheduled(
-            entryFunc, materializationResult.materializedSchedules,
-            "spatial3_scheduled_no_comm", "spatial3"))) {
-      signalPassFailure();
-      return;
-    }
-    dumpScheduledComputeReport(
-        moduleOp, entryFunc, schedule, materializationResult.peftClassPlans,
-        materializationResult.materializedSchedules);
-
-    if (failed(realizeDeferredCommunication(
-          entryFunc, materializationResult, target, pipelineStages))) {
-      moduleOp.emitError("Spatial communication realization failed");
-      signalPassFailure();
-      return;
-    }
-
-    bool deferredRemains = false;
-    entryFunc.walk([&](SpatDeferredCommunicationOp deferred) {
-      if (deferredRemains)
-        return;
-      deferred.emitOpError("realized Spatial graph still contains deferred communication");
-      deferredRemains = true;
-    });
-    if (deferredRemains
-        || failed(verifyScheduledResultsLive(materializationResult.materializedSchedules))
-        || failed(verifyScheduledSpatialInvariants(entryFunc))) {
-      moduleOp.emitError("realized Spatial communication verification failed");
-      signalPassFailure();
-      return;
-    }
-
-    dumpModule(moduleOp, "spatial4_scheduled", /*assumeVerified=*/true);
-    if (shouldExportSpatialDataflowStage(exportStage, SpatialDataflowExportStage::Spatial4)
-        && failed(exportSpatialDataflowCsvScheduled(
-            entryFunc, materializationResult.materializedSchedules,
-            "spatial4_scheduled", "spatial4")))
-      signalPassFailure();
-  }
-
-private:
-  SchedulingTarget target;
-  SpatialDataflowExportStage exportStage = SpatialDataflowExportStage::None;
-  size_t pipelineStages = 1;
-  bool hasTarget = false;
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createScheduleAndRealizeSpatialPass() {
-  return std::make_unique<ScheduleAndRealizeSpatialPass>();
-}
-
-std::unique_ptr<Pass> createScheduleAndRealizeSpatialPass(
-    const SchedulingTarget& target,
-    SpatialDataflowExportStage exportStage,
-    size_t pipelineStages) {
-  return std::make_unique<ScheduleAndRealizeSpatialPass>(
-    target, exportStage, pipelineStages);
-}
-
-} // namespace spatial
-} // namespace onnx_mlir

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

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

src/PIM/Dialect/Spatial/Spatial.td

@@ -10,6 +10,7 @@ include "mlir/IR/RegionKindInterface.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/ParallelCombiningOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "src/Accelerators/PIM/Dialect/Spatial/SpatialLayoutInterface.td"
 
 def SpatialDialect : Dialect {
   let name = "spat";
@@ -24,24 +25,6 @@ def SpatialDialect : Dialect {
   }];
 }
 
-def SpatialLayoutCapabilityInterface : OpInterface<"SpatialLayoutCapabilityInterface"> {
-  let description = [{
-    Contract implemented by logical Spatial planning operations that expose
-    their legal physical layout alternatives to the Spatial planner.
-  }];
-
-  let methods = [
-    InterfaceMethod<
-      "Return legal physical layout alternatives for this operation and its current operand layouts.",
-      "::llvm::SmallVector<::onnx_mlir::spatial::LayoutAlternative>",
-      "getLayoutAlternatives",
-      (ins "const ::onnx_mlir::spatial::SpatialTargetResources &":$target,
-           "::llvm::ArrayRef<::onnx_mlir::spatial::PhysicalLayout>":$operandLayouts)>
-  ];
-
-  let cppNamespace = "::onnx_mlir::spatial";
-}
-
 def SpatLogicalLayoutNCHW : I32EnumAttrCase<"NCHW", 0, "nchw">;
 def SpatLogicalLayout : I32EnumAttr<"LogicalLayout", "Logical tensor layout", [
   SpatLogicalLayoutNCHW
@@ -347,22 +330,6 @@ def SpatConv2DPlanOp : SpatLayoutPlanOp<"conv2d_plan"> {
   let hasVerifier = 1;
 }
 
-def SpatFlattenPlanOp : SpatLayoutPlanOp<"flatten_plan"> {
-  let summary = "Layout-aware static Flatten planning op";
-
-  let arguments = (ins
-    SpatTensor:$input,
-    I64Attr:$axis,
-    SpatLogicalLayoutAttr:$logicalLayout
-  );
-
-  let results = (outs
-    SpatTensor:$output
-  );
-
-  let hasVerifier = 1;
-}
-
 def SpatReluPlanOp : SpatLayoutPlanOp<"relu_plan"> {
   let summary = "Layout-aware ReLU planning op";
 
@@ -550,8 +517,7 @@ def SpatChannelSendOp : SpatOp<"channel_send", []> {
   );
 
   let assemblyFormat = [{
-    $input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
+    $input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($input)
-      attr-dict `:` type($input)
   }];
 }
 
@@ -569,45 +535,7 @@ def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
   );
 
   let assemblyFormat = [{
-    `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
+    `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($output)
-      attr-dict `:` type($output)
-  }];
-}
-
-def SpatHostStoreSyncOp : SpatOp<"host_store_sync", []> {
-  let summary = "Store a tensor to host memory and signal its consumer";
-
-  let arguments = (ins
-    Index:$sourceCoreId,
-    Index:$targetCoreId,
-    Index:$hostOffset,
-    Index:$eventRegister,
-    SpatTensor:$input
-  );
-
-  let assemblyFormat = [{
-    $input `from` $sourceCoreId `to` $targetCoreId
-      `host_offset` $hostOffset `event` $eventRegister attr-dict `:` type($input)
-  }];
-}
-
-def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> {
-  let summary = "Wait for a producer and load its tensor from host memory";
-
-  let arguments = (ins
-    Index:$sourceCoreId,
-    Index:$targetCoreId,
-    Index:$hostOffset,
-    Index:$eventRegister
-  );
-
-  let results = (outs
-    SpatTensor:$output
-  );
-
-  let assemblyFormat = [{
-    `from` $sourceCoreId `to` $targetCoreId
-      `host_offset` $hostOffset `event` $eventRegister attr-dict `:` type($output)
   }];
 }
 

src/PIM/Dialect/Spatial/SpatialLayoutInterface.td

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

src/PIM/Dialect/Spatial/SpatialOps.hpp

@@ -20,7 +20,7 @@
 #include <string>
 #include <tuple>
 
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
 
 /// Include the auto-generated header files containing the declarations
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialDialect.hpp.inc"

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -10,7 +10,7 @@
 #include "src/Accelerators/PIM/Common/IR/CompactAsmUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp"
 
 using namespace mlir;
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -463,33 +463,6 @@ LogicalResult SpatConv2DPlanOp::verify() {
   return success();
 }
 
-LogicalResult SpatFlattenPlanOp::verify() {
-  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.flatten_plan")))
-    return failure();
-  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
-  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
-  if (!inputType.hasStaticShape() || !outputType.hasStaticShape()
-      || outputType.getRank() != 2)
-    return emitError("requires static input and rank-2 output tensors");
-  if (getLogicalLayout() != LogicalLayout::NCHW)
-    return emitError("requires logical layout \"nchw\"");
-  if (llvm::any_of(inputType.getShape(), [](int64_t dim) { return dim <= 0; })
-      || llvm::any_of(outputType.getShape(), [](int64_t dim) { return dim <= 0; }))
-    return emitError("requires positive dimensions");
-  if (getAxis() > static_cast<uint64_t>(inputType.getRank()))
-    return emitError("requires axis in the input rank range");
-
-  int64_t outer = 1;
-  for (int64_t dim : inputType.getShape().take_front(getAxis()))
-    outer *= dim;
-  int64_t inner = 1;
-  for (int64_t dim : inputType.getShape().drop_front(getAxis()))
-    inner *= dim;
-  if (outputType.getDimSize(0) != outer || outputType.getDimSize(1) != inner)
-    return emitError("output shape does not match the requested flatten axis");
-  return success();
-}
-
 LogicalResult SpatReluPlanOp::verify() {
   if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.relu_plan")))
     return failure();

src/PIM/Dialect/Spatial/SpatialTargetInfo.hpp

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

src/PIM/Dialect/Spatial/SpatialTargetResources.hpp

@@ -1,27 +0,0 @@
-#pragma once
-
-#include "mlir/Support/LogicalResult.h"
-
-#include <cstddef>
-
-namespace onnx_mlir::spatial {
-
-struct MatrixUnitShape {
-  size_t rows = 0;
-  size_t columns = 0;
-};
-
-struct SpatialTargetResources {
-  MatrixUnitShape matrixShape;
-  size_t matrixUnitsPerProcessor = 0;
-  size_t processorCount = 0;
-  size_t vectorWidth = 0;
-
-  mlir::LogicalResult verify() const {
-    return mlir::success(matrixShape.rows != 0 && matrixShape.columns != 0
-                         && matrixUnitsPerProcessor != 0 && processorCount != 0
-                         && vectorWidth != 0);
-  }
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -219,12 +219,10 @@ static void appendReceive(BoundaryProgram &boundary,
         run->entryOffsets[run->entryOffsets.size() - 2]].family->requirement;
       CollectionTarget previousTarget {run->collection, run->positions.back()};
       bool sameEntry = previous == requirement;
-      bool sameRoute = run->slices.back().family->hostRouted
+      if (sameEntry
-                    == 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/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp

@@ -8,7 +8,6 @@
 #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;
@@ -19,8 +18,6 @@ struct LogicalTransferMetadataView {
   StaticIntSequenceChain parentCounts;
   StaticIntSequenceChain sourceCores;
   StaticIntSequenceChain targetCores;
-  StaticIntSequenceChain hostOffsets;
-  StaticIntSequenceChain eventRegisters;
   StaticIntSequenceChain targetLanes;
   StaticIntSequenceChain localOffsets;
   SmallVector<StaticIntSequenceChain> projectionOffsets;
@@ -31,8 +28,7 @@ struct LogicalTransferMetadataView {
 };
 using MetadataMember = StaticIntSequenceChain LogicalTransferMetadataView::*;
 static constexpr std::array<MetadataMember, 3> transferMetadataMembers{
-    &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores,
+    &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores, &LogicalTransferMetadataView::targetCores};
-    &LogicalTransferMetadataView::targetCores};
 struct TransferGrids {
   std::array<StaticIntGrid, 3> values;
   StaticIntGrid &channels() { return values[0]; }
@@ -45,8 +41,7 @@ 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),
+  return TransferGrids{{std::move(*channels), std::move(*sourceCores), std::move(*targetCores)}};
-                        std::move(*targetCores)}};
 }
 using GridGeometry = DeferredGridSliceGeometry;
 using StaticGeometryMember = SmallVector<StaticIntSequence> DeferredStaticSliceGeometry::*;
@@ -87,11 +82,6 @@ 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);
@@ -182,7 +172,6 @@ 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());
@@ -196,20 +185,6 @@ 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;
@@ -232,6 +207,7 @@ 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();
@@ -241,25 +217,10 @@ 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();
-    Value sourceCore = transferGrids->sourceCores().emitLookup(
+    auto send = SpatChannelSendOp::create(
-      action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+        context.rewriter, loc, transferGrids->channels().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
-    Value targetCore = transferGrids->targetCores().emitLookup(
+        transferGrids->sourceCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
-      action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+        transferGrids->targetCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), *payload);
-    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();
   };
@@ -294,45 +255,14 @@ 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);
-  Value sourceCore = grids->sourceCores().emitLookup(
+  auto receive = SpatChannelReceiveOp::create(context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
-    row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
+    grids->channels().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
-  Value targetCore = grids->targetCores().emitLookup(
+    grids->sourceCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
-    row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
+    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 output;
+  return receive.getOutput();
 }
 
 static FailureOr<SmallVector<LogicalTransferMetadataView, 0>>
@@ -385,9 +315,6 @@ 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) {
@@ -397,16 +324,6 @@ 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));
@@ -451,16 +368,6 @@ 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(
@@ -479,34 +386,15 @@ 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;
-    Value sourceCore = transferGrids->sourceCores().emitLookup(
+    auto receive =
-      entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
+        SpatChannelReceiveOp::create(context.rewriter, loc, fragmentType,
-    Value targetCore = transferGrids->targetCores().emitLookup(
+                                     transferGrids->channels().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
-      entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
+                                     transferGrids->sourceCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
-    Operation *receive;
+                                     transferGrids->targetCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
-    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(output, position, entry, runtimeLane, current);
+    return insert(receive.getOutput(), position, entry, runtimeLane, current);
   };
   if (actionCount == 1 && llvm::all_of(counts, [](int64_t count) { return count == 1; }))
     return emitEntry(context.constants.getIndex(0), initial);