restore unwanted changes

README.md

@@ -52,27 +52,41 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    `Patterns/{Math,NN,Tensor}` and currently cover Conv, Gemm, MatMul,
    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
+   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
+   `Passes/Transforms` and `Passes/Analyses`.
 
-2. **Merge compute nodes**
-   (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
-   Builds a compute graph, schedules it with the PEFT scheduler, and materializes
-   the merge schedule into Spatial IR. Supporting scheduling code lives under
-   `MergeComputeNodes/Scheduling`.
+2. **Merge, schedule, and realize Spatial communication**
+   (`src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes`).
+   `TrivialGraphComputeMerge` performs local graph merging. One
+   `ScheduleAndRealizeSpatial` pass then owns scheduling, intermediate
+   verification, communication realization, and final verification. Supporting
+   scheduling code lives under `MergeComputeNodes/Scheduling`.
 
 3. **Spatial -> Pim** (`src/PIM/Conversion/SpatialToPim`).
    Lowers Spatial operations to the `pim` dialect (`src/PIM/Dialect/Pim`),
    including `pim.core`, `pim.core_batch`, communication, tensor packing, global
    tensor materialization, and return-path normalization.
 
-4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
-   Converts tensor-semantics PIM IR into memref-semantics PIM IR using MLIR's
-   bufferization interfaces.
+4. **Bufferization** (`src/PIM/Dialect/Pim/Passes/Transforms/Bufferization`).
+   `PimBufferizationPreparation` establishes writable destinations without
+   duplicating the one-shot copy analysis, `PimOneShotBufferization` runs
+   MLIR's one-shot analysis,
+   `PimMemoryNormalization` forwards/removes redundant copies and normalizes
+   addressable accesses, and `PimBufferizationVerification` checks tensor
+   absence, contiguity, and copy address spaces.
 
 5. **PIM local-memory planning**
-   (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`).
+   (`src/PIM/Dialect/Pim/Passes/Transforms/LocalMemoryPlanning`).
    Computes whole-core lifetimes, reuses addresses for non-overlapping
-   allocations, and records the explicit plan in PIM IR.
-6. **PIM verification and code generation** (`src/PIM/Pass/PimCodegen` and
+   allocations, and records the explicit plan in PIM IR. Reusable lifetime
+   analysis lives under `src/PIM/Dialect/Pim/Passes/Analyses`.
+6. **PIM verification and code generation** (`src/PIM/Passes/PimCodegen` and
    `src/PIM/Compiler`).
    Verifies the memory plan and other PIM invariants, then emits `.pim` core
    files, weights, and `memory.bin` / `config.json` without rerunning liveness.
@@ -85,7 +99,7 @@ Supporting pieces:
   points.
 - `src/PIM/Conversion/SpatialToGraphviz` - optional Spatial graphviz conversion
   pass.
-- `src/PIM/Pass` - pass registration and auxiliary passes.
+- `src/PIM/Passes` - pass registration and auxiliary passes.
 - `src/PIM/PimAccelerator.{cpp,hpp}` - ONNX-MLIR accelerator entry point.
 
 ## PIM compiler options
@@ -118,16 +132,8 @@ options; `onnx-mlir --help` lists the inherited ONNX-MLIR options.
   elements per convolution before streaming. Default is `1048576`.
 - `--pim-conv-stream-chunk-positions=<N>` - maximum output positions per
   streamed convolution chunk. Default is `1024`.
-- `--pim-report-conv-lowering=<true|false>` - emit the bounded convolution
-  lowering report. Default is `true`.
-- `--use-experimental-conv-impl` - use the alternate convolution lowering.
 - `--pim-detect-communication-deadlock` - statically simulate expanded
   send/receive ordering and reject blocking deadlocks. Default is off.
-- `--pim-materialize-scalar-fanout-global-order` - use the experimental,
-  expensive globally ordered scalar-fanout materializer. Default is off.
-- `--pim-trace-communication-materialization` - emit verbose communication
-  materialization diagnostics and provenance attributes. Default is off.
-- `--ignore-concat-error` - soft-fail a ConcatOp corner case.
 
 ## Standard PIM hardware profile
 

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

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

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

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

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

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

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

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

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

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

src/PIM/CMakeLists.txt

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

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -1,10 +1,13 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 
+#include "llvm/ADT/SmallPtrSet.h"
+
 #include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
@@ -36,6 +39,10 @@ mlir::Value resolveAlias(mlir::Value value, const StaticValueKnowledge* knowledg
 
 llvm::FailureOr<CompiledIndexExpr> compileIndexValueImpl(mlir::Value value);
 llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Value value);
+using AliasResolutionSet = llvm::SmallPtrSet<mlir::Value, 8>;
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value,
+                                        const StaticValueKnowledge* knowledge,
+                                        AliasResolutionSet& visited);
 mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
 template <typename... Args>
@@ -45,18 +52,23 @@ CompiledIndexExpr makeCompiledIndexExpr(Args&&... args) {
 
 static mlir::Value resolveForYieldedAliasToInit(mlir::scf::ForOp forOp,
                                                 mlir::Value yieldedValue,
-                                                const StaticValueKnowledge* knowledge) {
-  yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge);
+                                                const StaticValueKnowledge* knowledge,
+                                                AliasResolutionSet& visited) {
+  yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge, visited);
   if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(yieldedValue)) {
     if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0
         && static_cast<unsigned>(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size())
-      return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge);
+      return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge, visited);
   }
   return yieldedValue;
 }
 
-mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value,
+                                        const StaticValueKnowledge* knowledge,
+                                        AliasResolutionSet& visited) {
   value = resolveAlias(value, knowledge);
+  if (!value || !visited.insert(value).second)
+    return value;
 
   if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(value)) {
     auto forOp = mlir::dyn_cast_or_null<mlir::scf::ForOp>(blockArgument.getOwner()->getParentOp());
@@ -64,9 +76,12 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
       const unsigned iterArgIndex = blockArgument.getArgNumber() - 1;
       auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
       if (iterArgIndex < forOp.getInitArgs().size() && yieldOp
-          && iterArgIndex < yieldOp.getNumOperands()
-          && resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge) == blockArgument)
-        return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge);
+          && iterArgIndex < yieldOp.getNumOperands()) {
+        mlir::Value yieldedValue = resolveAlias(yieldOp.getOperand(iterArgIndex), knowledge);
+        if (yieldedValue == blockArgument
+            || (yieldedValue && resolveLoopCarriedAliasImpl(yieldedValue, knowledge, visited) == blockArgument))
+          return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[iterArgIndex], knowledge, visited);
+      }
     }
     return value;
   }
@@ -75,10 +90,15 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
   if (!definingOp)
     return value;
 
+  if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(toBufferOp.getTensor(), knowledge, visited);
+  if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp))
+    return resolveLoopCarriedAliasImpl(toTensorOp.getBuffer(), knowledge, visited);
+
   if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
     if (auto result = mlir::dyn_cast<mlir::OpResult>(value))
       if (mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(result))
-        return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge);
+        return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge, visited);
   }
 
   if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(definingOp)) {
@@ -86,20 +106,26 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
     if (result) {
       auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
       if (yieldOp && result.getResultNumber() < yieldOp.getNumOperands())
-        return resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
+        return resolveForYieldedAliasToInit(
+          forOp, yieldOp.getOperand(result.getResultNumber()), knowledge, visited);
     }
   }
 
   if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge);
+    return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge, visited);
   if (auto collapseOp = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge);
+    return resolveLoopCarriedAliasImpl(collapseOp.getSrc(), knowledge, visited);
   if (auto expandOp = mlir::dyn_cast<mlir::memref::ExpandShapeOp>(definingOp))
-    return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge);
+    return resolveLoopCarriedAliasImpl(expandOp.getSrc(), knowledge, visited);
 
   return value;
 }
 
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
+  AliasResolutionSet visited;
+  return resolveLoopCarriedAliasImpl(value, knowledge, visited);
+}
+
 llvm::FailureOr<int64_t> resolveOpFoldResult(mlir::OpFoldResult ofr, const StaticValueKnowledge* knowledge);
 llvm::FailureOr<int64_t> resolveIndexValueImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
@@ -524,6 +550,15 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
     if (!definingOp)
       return mlir::failure();
 
+    if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp)) {
+      value = resolveAlias(toBufferOp.getTensor(), knowledge);
+      continue;
+    }
+    if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp)) {
+      value = resolveAlias(toTensorOp.getBuffer(), knowledge);
+      continue;
+    }
+
     if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
       mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast<mlir::OpResult>(value));
       if (!tiedOperand)
@@ -538,7 +573,9 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
+      AliasResolutionSet visited;
+      value = resolveForYieldedAliasToInit(
+        forOp, yieldOp.getOperand(result.getResultNumber()), knowledge, visited);
       continue;
     }
 
@@ -643,6 +680,15 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
     if (!definingOp)
       return mlir::failure();
 
+    if (auto toBufferOp = mlir::dyn_cast<mlir::bufferization::ToBufferOp>(definingOp)) {
+      value = toBufferOp.getTensor();
+      continue;
+    }
+    if (auto toTensorOp = mlir::dyn_cast<mlir::bufferization::ToTensorOp>(definingOp)) {
+      value = toTensorOp.getBuffer();
+      continue;
+    }
+
     if (auto dpsDefiningOp = mlir::dyn_cast<mlir::DestinationStyleOpInterface>(definingOp)) {
       mlir::OpOperand* tiedOperand = dpsDefiningOp.getTiedOpOperand(mlir::dyn_cast<mlir::OpResult>(value));
       if (!tiedOperand)
@@ -657,7 +703,9 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), nullptr);
+      AliasResolutionSet visited;
+      value = resolveForYieldedAliasToInit(
+        forOp, yieldOp.getOperand(result.getResultNumber()), nullptr, visited);
       continue;
     }
 

src/PIM/Common/PimCommon.hpp

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

src/PIM/Compiler/PimBinaryFormat.hpp

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

src/PIM/Compiler/PimCodeGen.cpp

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

src/PIM/Compiler/PimCodeGen.hpp

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

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -2,6 +2,8 @@
 
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
+#include <limits>
+
 #define DEBUG_TYPE "PimCompilerOptions"
 
 namespace onnx_mlir {
@@ -70,11 +72,6 @@ llvm::cl::opt<bool>
                  llvm::cl::init(false),
                  llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> useExperimentalConvImpl("use-experimental-conv-impl",
-                                            llvm::cl::desc("Use experimental implementation for convolution"),
-                                            llvm::cl::init(false),
-                                            llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<uint64_t> pimConvIm2colMaxElements(
   "pim-conv-im2col-max-elements",
   llvm::cl::desc("Maximum number of im2col elements to materialize globally for one Conv before streaming/chunking"),
@@ -103,15 +100,9 @@ llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder(
-  "pim-materialize-scalar-fanout-global-order",
-  llvm::cl::desc("Experimental expensive materializer mode: emit scalar-source fanout as globally ordered communication events instead of all-send fanout loops"),
-  llvm::cl::init(false),
-  llvm::cl::cat(OnnxMlirOptions));
-
-llvm::cl::opt<bool> pimTraceCommunicationMaterialization(
-  "pim-trace-communication-materialization",
-  llvm::cl::desc("Emit verbose materializer-time diagnostics and provenance attributes for every Spatial communication op"),
+llvm::cl::opt<bool> pimVerifyBufferizationCopyFreedom(
+  "pim-verify-bufferization-copy-freedom",
+  llvm::cl::desc("Run the expensive official PIM tensor-copy freedom proof before bufferization"),
   llvm::cl::init(false),
   llvm::cl::cat(OnnxMlirOptions));
 
@@ -121,6 +112,12 @@ llvm::cl::opt<size_t>
 llvm::cl::opt<size_t>
   crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(64));
 
+llvm::cl::opt<size_t> pipelineStages(
+  "pipeline",
+  llvm::cl::desc("Number of throughput pipeline stages (1 preserves latency scheduling)"),
+  llvm::cl::init(1),
+  llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<long> coresCount("core-count",
                                llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."),
                                llvm::cl::init(-1));
@@ -131,11 +128,6 @@ llvm::cl::opt<std::string> pimTargetConfig(
   llvm::cl::init(""),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<bool>
-  ignoreConcatError("ignore-concat-error",
-                    llvm::cl::desc("Ignore ConcatOp corner case: do not assert and do a simplification"),
-                    llvm::cl::init(false));
-
 bool hasExplicitPimCoreCount() { return coresCount.getNumOccurrences() != 0; }
 
 void verifyExplicitPimCoreCount() {
@@ -145,4 +137,14 @@ void verifyExplicitPimCoreCount() {
     llvm::report_fatal_error("PIM compilation requires --core-count to be a positive integer");
 }
 
+void verifyPimPipelineStages() {
+  if (pipelineStages.getValue() == 0)
+    llvm::report_fatal_error("PIM compilation requires --pipeline to be positive");
+  if (static_cast<size_t>(coresCount.getValue()) % pipelineStages.getValue() != 0)
+    llvm::report_fatal_error("PIM compilation requires --core-count to be divisible by --pipeline");
+  if (crossbarCountInCore.getValue()
+      > std::numeric_limits<size_t>::max() / pipelineStages.getValue())
+    llvm::report_fatal_error("PIM compilation --crossbar-count * --pipeline overflows");
+}
+
 } // namespace onnx_mlir

src/PIM/Compiler/PimCompilerOptions.hpp

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

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -14,9 +14,12 @@
 
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 #include "src/Compiler/CompilerPasses.hpp"
 
 #define DEBUG_TYPE "PimCompilerUtils"
@@ -80,6 +83,54 @@ spatial::SchedulingTarget getDefaultPimSchedulingTarget() {
   return target;
 }
 
+spatial::ConvLoweringStrategy getSpatialConvLoweringStrategy(PimConvLoweringType strategy) {
+  switch (strategy) {
+  case PimConvLoweringAuto: return spatial::ConvLoweringStrategy::Auto;
+  case PimConvLoweringLegacy: return spatial::ConvLoweringStrategy::Legacy;
+  case PimConvLoweringDepthwise: return spatial::ConvLoweringStrategy::Depthwise;
+  case PimConvLoweringPackedIm2Col: return spatial::ConvLoweringStrategy::PackedIm2Col;
+  case PimConvLoweringStreamedPatch: return spatial::ConvLoweringStrategy::StreamedPatch;
+  case PimConvLoweringStreamedPacked: return spatial::ConvLoweringStrategy::StreamedPacked;
+  case PimConvLoweringOutputChannelTiled: return spatial::ConvLoweringStrategy::OutputChannelTiled;
+  case PimConvLoweringInputKTiled: return spatial::ConvLoweringStrategy::InputKTiled;
+  case PimConvLoweringTiled2D: return spatial::ConvLoweringStrategy::Tiled2D;
+  }
+  llvm_unreachable("unknown PIM Conv lowering strategy");
+}
+
+spatial::SpatialDataflowExportStage getPimSpatialDataflowExportStage(
+    PimSpatialDataflowExportType stage) {
+  switch (stage) {
+  case SpatialDataflowExportNone:     return spatial::SpatialDataflowExportStage::None;
+  case SpatialDataflowExportSpatial1: return spatial::SpatialDataflowExportStage::Spatial1;
+  case SpatialDataflowExportSpatial2: return spatial::SpatialDataflowExportStage::Spatial2;
+  case SpatialDataflowExportSpatial3: return spatial::SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportSpatial4: return spatial::SpatialDataflowExportStage::Spatial4;
+  case SpatialDataflowExportAll:      return spatial::SpatialDataflowExportStage::All;
+  }
+  llvm_unreachable("unknown PIM Spatial dataflow export stage");
+}
+
+spatial::SpatialTargetResources getPimSpatialTargetResources(const spatial::SchedulingTarget& target) {
+  spatial::SpatialTargetResources resources;
+  resources.matrixShape = {target.matrixRows, target.matrixColumns};
+  resources.matrixUnitsPerProcessor = target.residentWeightCapacity;
+  resources.processorCount = target.processorCount;
+  resources.vectorWidth = target.vectorWidth;
+  if (failed(resources.verify()))
+    llvm::report_fatal_error("PIM target resources are incomplete");
+  return resources;
+}
+
+ONNXToSpatialPlanningOptions getPimONNXToSpatialPlanningOptions() {
+  ONNXToSpatialPlanningOptions options;
+  options.convIm2colMaxElements = pimConvIm2colMaxElements.getValue();
+  options.convStreamChunkPositions = pimConvStreamChunkPositions.getValue();
+  options.forcedConvStrategy = getSpatialConvLoweringStrategy(pimConvLowering.getValue());
+  options.reportConvLowering = pimReportConvLowering.getValue();
+  return options;
+}
+
 const llvm::json::Object& requireObject(const llvm::json::Object& object,
                                         llvm::StringRef key,
                                         llvm::StringRef path) {
@@ -279,11 +330,14 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
                   EmissionTargetType& emissionTarget,
                   std::string outputNameNoExt) {
   verifyExplicitPimCoreCount();
+  verifyPimPipelineStages();
+  spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
+  spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget);
 
   if (pimOnlyCodegen) {
     pm.addPass(createPimInstructionSelectionPass());
     pm.addPass(createPimLocalMemoryPlanningPass());
-    pm.addPass(createPimVerificationPass());
+    pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
     pm.addPass(createEmitPimCodePass());
     return;
   }
@@ -292,23 +346,29 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     addONNXToMLIRPasses(pm, /*target CPU*/ false);
 
   if (pimEmissionTarget >= EmitSpatial) {
-    spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
-    pm.addPass(createONNXToSpatialPass());
-    pm.addPass(createSpatialLayoutPlanningPass());
-    pm.addPass(createLowerSpatialPlansPass());
+    ONNXToSpatialPlanningOptions planningOptions = getPimONNXToSpatialPlanningOptions();
+    spatial::SpatialDataflowExportStage exportStage =
+      getPimSpatialDataflowExportStage(pimExportSpatialDataflow.getValue());
+    pm.addPass(createONNXToSpatialPass(targetResources, planningOptions));
+    pm.addPass(createSpatialLayoutPlanningPass(targetResources));
+    pm.addPass(createLowerSpatialPlansPass(targetResources, planningOptions, exportStage));
     pm.addPass(createTrivialGraphComputeMergePass(
-      schedulingTarget.residentWeightCapacity));
-    pm.addPass(createMergeComputeNodesPass(schedulingTarget));
+      schedulingTarget.residentWeightCapacity, exportStage));
+    pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
+      schedulingTarget, exportStage, pipelineStages.getValue()));
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
 
   if (pimEmissionTarget >= EmitPim) {
-    pm.addPass(createSpatialToPimPass());
+    pm.addPass(createSpatialToPimPass(targetResources));
     pm.addPass(createMessagePass("Spatial lowered to Pim"));
   }
 
   if (pimEmissionTarget >= EmitPimBufferized) {
-    pm.addPass(createPimBufferizationPass());
+    pm.addPass(createPimBufferizationPreparationPass(pimVerifyBufferizationCopyFreedom.getValue()));
+    pm.addPass(createPimOneShotBufferizationPass());
+    pm.addPass(createPimMemoryNormalizationPass());
+    pm.addPass(createPimBufferizationVerificationPass());
     pm.addPass(createMessagePass("Pim bufferized"));
   }
 
@@ -320,7 +380,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(createMessagePass("Pim instructions selected"));
     pm.addPass(createPimLocalMemoryPlanningPass());
     pm.addPass(createMessagePass("Pim local memory planned"));
-    pm.addPass(createPimVerificationPass());
+    pm.addPass(createPimVerificationPass(targetResources, pimDetectCommunicationDeadlock.getValue()));
     pm.addPass(createMessagePass("Pim verified"));
     pm.addPass(createEmitPimCodePass());
     pm.addPass(createMessagePass("Pim code emitted"));

src/PIM/Compiler/PimCoreProgram.cpp

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

src/PIM/Compiler/PimCoreProgram.hpp

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

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -5,7 +5,7 @@ add_public_tablegen_target(ONNXToSpatialIncGen)
 add_pim_library(OMONNXToSpatial
   Patterns.cpp
   CompileTime.cpp
-  ONNXToSpatialVerifier.cpp
+  Passes/Analyses/ONNXToSpatialVerifier.cpp
   Patterns/Pre.cpp
   Patterns/Post.cpp
   Patterns/Math/Conv.cpp
@@ -26,12 +26,15 @@ add_pim_library(OMONNXToSpatial
   Patterns/Tensor/Slice.cpp
   Patterns/Tensor/Split.cpp
   Patterns/Tensor/Transpose.cpp
-  ONNXToSpatialPass.cpp
-  SpatialLayoutPlanningPass.cpp
-  LowerSpatialPlansPass.cpp
+  Passes/Transforms/ONNXToSpatialPass.cpp
+  Passes/Analyses/SpatialLayoutCapabilities.cpp
+  Passes/Transforms/SpatialLayoutPlanningPass.cpp
+  Passes/Transforms/SpatialPlanLoweringPatterns.cpp
+  Passes/Transforms/LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
   Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.cpp
+  Common/ContractionPlanning.cpp
   Common/MatrixProductLowering.cpp
   Common/RowStripLayoutUtils.cpp
   Common/ShapeTilingUtils.cpp
@@ -46,8 +49,6 @@ add_pim_library(OMONNXToSpatial
   MLIRLinalgDialect
   MLIRSCFDialect
   MLIRTosaDialect
-  OMCompilerOptions
-  OMPimCompilerOptions
   OMONNXOps
   SpatialOps
   OMPimCommon

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

@@ -25,6 +25,9 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
   const int64_t laneCount = physicalType.getDimSize(0);
   if (laneCount <= 0)
     return emitError(loc, "fragment assembly requires at least one physical source slot"), failure();
+  auto physicalLayoutValue = spatial::symbolizePhysicalLayout(physicalLayout);
+  if (!physicalLayoutValue)
+    return emitError(loc, "unknown physical layout for fragment assembly"), failure();
   const int64_t fragmentElements = physicalType.getNumElements() / laneCount;
   SmallVector<int64_t> operandIndices(entries.size(), 0), sourceSlots, sourceOffsets, offsets, sizes,
     strides(entries.size() * rank, 1);
@@ -47,13 +50,18 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
     llvm::append_range(offsets, entry.destinationOffsets);
     llvm::append_range(sizes, entry.sizes);
   }
-  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
-    rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout),
+  auto blueprint = spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
+    spatial::getNCHWLayout(rewriter.getContext()),
+    spatial::PhysicalLayoutAttr::get(rewriter.getContext(), *physicalLayoutValue),
     rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
-    rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getStringAttr(indexMap), spatial::getFragmentAssemblyMode(rewriter.getContext()),
     rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
     rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
-    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete"));
+  if (indexMap == spatial::kContiguousRowMajorFragments
+      && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    blueprint.setIndexMapAttr(rewriter.getStringAttr("fragment_assembly"));
+  return blueprint.getOutput();
 }
 
 Value sumTensors(ArrayRef<Value> tensors, PatternRewriter& rewriter) {

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

@@ -394,6 +394,39 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
     rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
 }
 
+template <typename BodyFn>
+mlir::FailureOr<mlir::Value> mapGraphBatchFragments(mlir::Value input,
+                                                    mlir::RankedTensorType outputType,
+                                                    mlir::PatternRewriter& rewriter,
+                                                    mlir::Location loc,
+                                                    BodyFn&& build) {
+  auto inputType = mlir::dyn_cast<mlir::RankedTensorType>(input.getType());
+  if (!inputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
+      || inputType.getRank() != outputType.getRank() || inputType.getRank() < 2
+      || inputType.getDimSize(0) != outputType.getDimSize(0))
+    return mlir::failure();
+  auto inputFragmentType = mlir::RankedTensorType::get(
+    inputType.getShape().drop_front(), inputType.getElementType(), inputType.getEncoding());
+  auto outputFragmentType = mlir::RankedTensorType::get(
+    outputType.getShape().drop_front(), outputType.getElementType(), outputType.getEncoding());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, mlir::TypeRange {outputType}, inputType.getDimSize(0), {}, mlir::ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) -> mlir::LogicalResult {
+      auto fragment = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), args.lane, inputFragmentType);
+      if (mlir::failed(fragment))
+        return mlir::failure();
+      mlir::FailureOr<mlir::Value> result = build(*fragment, outputFragmentType);
+      if (mlir::failed(result) || result->getType() != outputFragmentType)
+        return mlir::failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, *result, args.outputs.front(), args.lane);
+      return mlir::success();
+    });
+  if (mlir::failed(batch))
+    return mlir::failure();
+  return batch->getResult(0);
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -0,0 +1,42 @@
+#include "ContractionPlanning.hpp"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+
+#include <algorithm>
+
+namespace onnx_mlir {
+
+namespace {
+
+static int64_t ceilDivide(int64_t value, int64_t divisor) {
+  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
+}
+
+} // namespace
+
+ContractionPlan makeContractionPlan(
+    const ContractionProblem& problem,
+    const spatial::SpatialTargetResources& target,
+    ContractionPlanKind kind,
+    int64_t laneCount,
+    int64_t fragmentRows) {
+  ContractionPlan plan;
+  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>(
+      1, fragmentRows != 0 ? fragmentRows : target.matrixShape.rows);
+
+  if (laneCount != 0)
+    plan.laneCount = laneCount;
+  else if (kind == ContractionPlanKind::StaticTiled)
+    plan.laneCount = problem.batch * problem.m * plan.reductionSlices * plan.outputTiles;
+  else if (kind == ContractionPlanKind::GroupedRowDynamicVVD)
+    plan.laneCount = problem.batch * ceilDivide(problem.m, rowsPerLane);
+  else
+    plan.laneCount = problem.batch * problem.m * problem.n;
+  return plan;
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "ContractionProblem.hpp"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+
+namespace onnx_mlir {
+
+enum class ContractionPlanKind {
+  StaticTiled,
+  BatchedDynamicVVD,
+  GroupedRowDynamicVVD,
+};
+
+struct ContractionPlan {
+  int64_t tileK = 1;
+  int64_t tileN = 1;
+  int64_t reductionSlices = 1;
+  int64_t outputTiles = 1;
+  int64_t laneCount = 0;
+};
+
+ContractionPlan makeContractionPlan(
+    const ContractionProblem& problem,
+    const spatial::SpatialTargetResources& target,
+    ContractionPlanKind kind,
+    int64_t laneCount = 0,
+    int64_t fragmentRows = 0);
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstdint>
+
+namespace onnx_mlir {
+
+struct ContractionProblem {
+  llvm::SmallVector<int64_t> lhsBatchShape;
+  llvm::SmallVector<int64_t> rhsBatchShape;
+  llvm::SmallVector<int64_t> outputBatchShape;
+  int64_t lhsBatch = 1;
+  int64_t rhsBatch = 1;
+  int64_t batch = 1;
+  int64_t m = 0;
+  int64_t k = 0;
+  int64_t n = 0;
+  mlir::Type lhsElementType;
+  mlir::Type rhsElementType;
+  mlir::Type resultElementType;
+};
+
+} // namespace onnx_mlir

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

@@ -1,15 +1,70 @@
 #include "MatrixProductLowering.hpp"
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 
+static bool isInsideSpatialCompute(Operation* op) {
+  for (Operation* parent = op; parent; parent = parent->getParentOp())
+    if (spatial::isAnySpatialComputeLike(parent))
+      return true;
+  return false;
+}
+
+static Value buildLinalgTranspose(Value value,
+                                  RankedTensorType resultType,
+                                  ArrayRef<int64_t> permutation,
+                                  PatternRewriter& rewriter,
+                                  Location loc) {
+  Value init = tensor::EmptyOp::create(
+      rewriter, loc, resultType.getShape(), resultType.getElementType());
+  return linalg::TransposeOp::create(
+    rewriter, loc, value, init, permutation).getResult()[0];
+}
+
+static Value materializeConstantTranspose(Value value,
+                                          RankedTensorType resultType,
+                                          ArrayRef<int64_t> permutation,
+                                          PatternRewriter& rewriter) {
+  auto denseAttr = getHostConstDenseElementsAttr(value);
+  if (!denseAttr)
+    return {};
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return {};
+  return getOrCreateConstant(
+      rewriter, rewriter.getInsertionBlock()->getParentOp(), *transposedAttr, resultType);
+}
+
+Value createLinalgTranspose(Value value,
+                            RankedTensorType resultType,
+                            ArrayRef<int64_t> permutation,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  if (Value constant = materializeConstantTranspose(value, resultType, permutation, rewriter))
+    return constant;
+
+  if (isInsideSpatialCompute(rewriter.getInsertionBlock()->getParentOp()))
+    return buildLinalgTranspose(value, resultType, permutation, rewriter, loc);
+
+  auto compute = createSpatCompute<1>(
+    rewriter, loc, TypeRange {resultType}, {}, ValueRange {value},
+    [&](Value input) {
+      spatial::SpatYieldOp::create(
+        rewriter, loc, buildLinalgTranspose(input, resultType, permutation, rewriter, loc));
+    });
+  return compute.getResult(0);
+}
+
 Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));

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

@@ -5,8 +5,16 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "llvm/ADT/ArrayRef.h"
+
 namespace onnx_mlir {
 
+mlir::Value createLinalgTranspose(mlir::Value value,
+                                  mlir::RankedTensorType resultType,
+                                  llvm::ArrayRef<int64_t> permutation,
+                                  mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc);
+
 mlir::Value createZeroPaddedTensor(mlir::Value value,
                                    mlir::RankedTensorType resultType,
                                    mlir::PatternRewriter& rewriter,

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

@@ -5,9 +5,9 @@
 #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/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
 
 #include <numeric>
 
@@ -33,6 +33,16 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
     tilesPerRow};
 }
 
+FailureOr<RowStripPhysicalValue> getRowStripPhysicalValue(Value value) {
+  auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>();
+  auto logicalType = dyn_cast<RankedTensorType>(value.getType());
+  if (!blueprint || !logicalType || blueprint.getOutput() != value
+      || blueprint.getPhysicalLayout() != spatial::PhysicalLayout::NHWCRowStrip
+      || !spatial::isPhysicalView(blueprint.getMode()))
+    return failure();
+  return describeRowStripPhysicalValue(blueprint.getInput(), logicalType);
+}
+
 RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
   return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
                                 logicalType.getDimSize(1)},
@@ -144,6 +154,35 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> createRowStripStorageBlueprint(Value storage,
+                                                RankedTensorType logicalType,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
+  if (failed(value))
+    return failure();
+
+  auto blueprint = spatial::SpatBlueprintOp::create(
+      rewriter,
+      loc,
+      logicalType,
+      storage,
+      ValueRange {},
+      spatial::getNCHWLayout(rewriter.getContext()),
+      spatial::getNHWCRowStripLayout(rewriter.getContext()),
+      rewriter.getDenseI64ArrayAttr({}),
+      rewriter.getDenseI64ArrayAttr({}),
+      rewriter.getStringAttr(kRowStripIndexMap),
+      spatial::getPhysicalViewMode(rewriter.getContext()),
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr);
+  return blueprint.getOutput();
+}
+
 FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                  PatternRewriter& rewriter,
                                                  Location loc) {
@@ -160,8 +199,8 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
       if (failed(fragment))
         return failure();
-      Value nchw = ONNXTransposeOp::create(
-        rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
+      Value nchw = createLinalgTranspose(
+        *fragment, nchwFragmentType, {0, 3, 1, 2}, rewriter, loc);
       publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
       return success();
     });
@@ -176,7 +215,7 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
          value.logicalType.getDimSize(3)}});
     }
-  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
+  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "dense_nchw",
                                          kRowStripIndexMap, rewriter, loc);
 }
 
@@ -186,25 +225,9 @@ static FailureOr<Value> applyRowStripActivation(const RowStripPhysicalValue& val
                                                 Location loc,
                                                 BuildActivation buildActivation) {
   auto storageType = cast<RankedTensorType>(value.storage.getType());
-  const int64_t laneCount = storageType.getDimSize(0);
-  auto batchOp = createSpatComputeBatch(rewriter,
-                                        loc,
-                                        TypeRange {storageType},
-                                        laneCount,
-                                        {},
-                                        ValueRange {value.storage},
-                                        [&](detail::SpatComputeBatchBodyArgs args) {
-                                          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
-                                            rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
-                                          if (failed(fragment)) return failure();
-                                          Value result = buildActivation(*fragment);
-                                          publishGraphBatchPhysicalFragment(
-                                            rewriter, loc, result, args.outputs.front(), args.lane);
-                                          return success();
-                                        });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  return mapGraphBatchFragments(value.storage, storageType, rewriter, loc, [&](Value fragment, RankedTensorType) {
+    return FailureOr<Value>(buildActivation(fragment));
+  });
 }
 
 FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {

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

@@ -6,6 +6,12 @@
 
 namespace onnx_mlir {
 
+namespace spatial {
+class SpatBlueprintOp;
+class SpatFlattenPlanOp;
+struct SpatialTargetResources;
+} // namespace spatial
+
 inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
 
 struct RowStripPhysicalValue {
@@ -18,6 +24,8 @@ struct RowStripPhysicalValue {
 mlir::FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(mlir::Value storage,
                                                                      mlir::RankedTensorType logicalType);
 
+mlir::FailureOr<RowStripPhysicalValue> getRowStripPhysicalValue(mlir::Value value);
+
 std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
 buildRowStripMetadata(mlir::RankedTensorType type);
 
@@ -53,6 +61,11 @@ mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
                                                            mlir::PatternRewriter& rewriter,
                                                            mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> createRowStripStorageBlueprint(mlir::Value storage,
+                                                            mlir::RankedTensorType logicalType,
+                                                            mlir::PatternRewriter& rewriter,
+                                                            mlir::Location loc);
+
 mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                                mlir::PatternRewriter& rewriter,
                                                                mlir::Location loc);
@@ -80,4 +93,14 @@ mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysical
                                                  mlir::PatternRewriter& rewriter,
                                                  mlir::Location loc);
 
+mlir::LogicalResult canLowerFlattenFromRowStrip(
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target);
+
+mlir::LogicalResult lowerFlattenFromRowStrip(
+    const RowStripPhysicalValue& input,
+    spatial::SpatFlattenPlanOp flattenOp,
+    const spatial::SpatialTargetResources& target,
+    mlir::PatternRewriter& rewriter);
+
 } // namespace onnx_mlir

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

@@ -5,7 +5,6 @@
 
 #include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 
 using namespace mlir;
@@ -67,11 +66,15 @@ sliceVector(const Value& vectorToSlice, int64_t sliceSize, PatternRewriter& rewr
 }
 
 DenseMap<CoreId, SmallVector<Value>>
-sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, PatternRewriter& rewriter, Location loc) {
-  SmallVector<Value> slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc);
+sliceVectorPerCrossbarPerCore(const Value& vectorToSlice,
+                              PatternRewriter& rewriter,
+                              Location loc,
+                              const spatial::SpatialTargetResources& target) {
+  SmallVector<Value> slices = sliceVector(
+    vectorToSlice, static_cast<int64_t>(target.matrixShape.rows), rewriter, loc);
   DenseMap<CoreId, SmallVector<Value>> slicesPerCore;
   for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) {
-    size_t coreId = sliceId / crossbarCountInCore;
+    size_t coreId = sliceId / target.matrixUnitsPerProcessor;
     slicesPerCore[coreId].push_back(slices[sliceId]);
   }
   return slicesPerCore;

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

@@ -7,6 +7,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 
@@ -26,6 +27,9 @@ llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
 /// Partitions one logical vector into per-core crossbar-sized slices using the
 /// current PIM target geometry.
 llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
-  const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc);
+  const mlir::Value& vectorToSlice,
+  mlir::PatternRewriter& rewriter,
+  mlir::Location loc,
+  const spatial::SpatialTargetResources& target);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -1,739 +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 "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallPtrSet.h"
-
-#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "mlir/Transforms/Passes.h"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace {
-
-static constexpr StringLiteral kDenseLayout = "dense_nchw";
-static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
-
-static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
-                                                         Value value) {
-  auto it = rowStripValues.find(value);
-  if (it == rowStripValues.end())
-    return failure();
-  return it->second;
-}
-
-static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprintOp blueprint,
-                                                           Value storage) {
-  auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
-  if (!logicalType)
-    return blueprint.emitOpError("requires ranked logical output type"), failure();
-  if (blueprint.getIndexMap() != kRowStripIndexMap)
-    return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
-  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
-  if (failed(value))
-    return blueprint.emitOpError("requires physical row-strip fragment storage"), failure();
-  return *value;
-}
-
-static FailureOr<Value>
-lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
-  return applyRowStripRelu(input, rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value>
-lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
-  return applyRowStripSilu(input, rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
-                                             spatial::SpatBiasAddPlanOp planOp,
-                                             PatternRewriter& rewriter) {
-  return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
-                                         const RowStripPhysicalValue& rhs,
-                                         spatial::SpatAddPlanOp planOp,
-                                         PatternRewriter& rewriter) {
-  return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
-                                            spatial::SpatConcatPlanOp planOp,
-                                            PatternRewriter& rewriter) {
-  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
-  if (!outputType)
-    return failure();
-  return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
-}
-
-static FailureOr<Value>
-materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
-  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
-    return failure();
-  return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc);
-}
-
-static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
-                                               PatternRewriter& rewriter, Location loc) {
-  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
-  auto inputType = dyn_cast<RankedTensorType>(input.getType());
-  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
-  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
-      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
-      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
-    return failure();
-  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
-  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
-  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
-      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
-      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
-    return failure();
-  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
-    if (outputDim > inputDim)
-      return failure();
-
-  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
-    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
-      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
-      if (failed(fragment))
-        return failure();
-      MixedSliceGeometry biasSlice;
-      for (int64_t dim : inputFragmentType->getShape()) {
-        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
-        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
-        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
-      }
-      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
-      if (!biasFragment)
-        return failure();
-      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
-      MixedSliceGeometry outputSlice;
-      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
-      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
-      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
-      for (int64_t dim : outputFragmentType->getShape())
-        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
-      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
-      if (!output)
-        return failure();
-      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
-      return success();
-    });
-  if (failed(batch))
-    return failure();
-  return batch->getResult(0);
-}
-
-static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
-                                  llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
-                                  llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
-                                  PatternRewriter& rewriter) {
-  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(rowStripValues, planOp.getLhs());
-  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rowStripValues, planOp.getRhs());
-  if (succeeded(lhs) && succeeded(rhs)) {
-    auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-      auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-      return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-    });
-    if (outputBlueprint == planOp.getResult().getUsers().end())
-      return planOp.emitOpError("row-strip add plan requires a row-strip blueprint result");
-
-    rewriter.setInsertionPoint(planOp);
-    FailureOr<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
-    if (failed(lowered))
-      return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
-    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-    FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-    if (failed(output))
-      return failure();
-    rowStripValues[blueprint.getResult()] = *output;
-    eraseAfterLowering.insert(planOp);
-    eraseAfterLowering.insert(blueprint);
-    return success();
-  }
-
-  rewriter.setInsertionPoint(planOp);
-  auto compute = createSpatCompute<2>(rewriter,
-                                      planOp.getLoc(),
-                                      planOp.getOutput().getType(),
-                                      {},
-                                      ValueRange {planOp.getLhs(), planOp.getRhs()},
-                                      [&](Value lhsValue, Value rhsValue) {
-                                        Value added = spatial::SpatVAddOp::create(
-                                          rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
-                                        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added);
-                                      });
-  rewriter.replaceOp(planOp, compute.getResults());
-  return success();
-}
-
-static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
-                                     llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
-                                     llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
-                                     PatternRewriter& rewriter) {
-  SmallVector<RowStripPhysicalValue> inputs;
-  for (Value input : planOp.getInputs()) {
-    FailureOr<RowStripPhysicalValue> physical = getRowStripValue(rowStripValues, input);
-    if (failed(physical)) {
-      inputs.clear();
-      break;
-    }
-    inputs.push_back(*physical);
-  }
-  if (inputs.size() == planOp.getInputs().size()) {
-    auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-      auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-      return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-    });
-    if (outputBlueprint == planOp.getResult().getUsers().end())
-      return planOp.emitOpError("row-strip concat plan requires a row-strip blueprint result");
-    rewriter.setInsertionPoint(planOp);
-    FailureOr<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
-    if (failed(lowered))
-      return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
-    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-    FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-    if (failed(output))
-      return failure();
-    rowStripValues[blueprint.getResult()] = *output;
-    eraseAfterLowering.insert(planOp);
-    eraseAfterLowering.insert(blueprint);
-    return success();
-  }
-
-  rewriter.setInsertionPoint(planOp);
-  auto compute = createSpatCompute(
-    rewriter,
-    planOp.getLoc(),
-    TypeRange {planOp.getOutput().getType()},
-    {},
-    planOp.getInputs(),
-    [&](ValueRange values) {
-      Value concatenated = spatial::SpatConcatOp::create(
-        rewriter, planOp.getLoc(), planOp.getOutput().getType(), rewriter.getI64IntegerAttr(planOp.getAxis()), values);
-      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated);
-    });
-  rewriter.replaceOp(planOp, compute.getResults());
-  return success();
-}
-
-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."; }
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-    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);
-    llvm::DenseMap<Value, RowStripPhysicalValue> rowStripValues;
-    llvm::SmallPtrSet<Operation*, 16> eraseAfterLowering;
-    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;
-    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
-      if (auto planOp = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
-        FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(rowStripValues, planOp.getInput());
-        auto rowStripBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-        });
-        if (rowStripBlueprint != planOp.getResult().getUsers().end()) {
-          rewriter.setInsertionPoint(planOp);
-          std::optional<Value> physicalInput;
-          if (succeeded(rowStripInput))
-            physicalInput = rowStripInput->storage;
-          FailureOr<Value> lowered = lowerSelectedConv2DPlan(
-            planOp,
-            physicalInput,
-            /*emitRowStripLayout=*/true,
-            rewriter);
-          if (failed(lowered)) {
-            auto diagnostic = planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan with input ");
-            diagnostic << planOp.getInput().getType() << " and output " << planOp.getResult().getType();
-            if (physicalInput)
-              diagnostic << " from physical storage " << physicalInput->getType();
-            signalPassFailure();
-            return;
-          }
-          auto blueprint = cast<spatial::SpatBlueprintOp>(*rowStripBlueprint);
-          FailureOr<RowStripPhysicalValue> rowStripValue = buildRowStripValue(blueprint, *lowered);
-          if (failed(rowStripValue)) {
-            signalPassFailure();
-            return;
-          }
-          rowStripValues[blueprint.getResult()] = *rowStripValue;
-          eraseAfterLowering.insert(planOp);
-          eraseAfterLowering.insert(blueprint);
-          continue;
-        }
-        rewriter.setInsertionPoint(planOp);
-        FailureOr<Value> lowered =
-          lowerSelectedConv2DPlan(planOp, std::nullopt, /*emitRowStripLayout=*/false, rewriter);
-        if (failed(lowered)) {
-          planOp.emitOpError("failed to lower selected Spatial Conv plan");
-          signalPassFailure();
-          return;
-        }
-        rewriter.replaceOp(planOp, *lowered);
-        continue;
-      }
-
-      if (auto planOp = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
-        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
-          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-          });
-          if (outputBlueprint == planOp.getResult().getUsers().end()) {
-            planOp.emitOpError("row-strip Relu plan requires a row-strip blueprint result");
-            signalPassFailure();
-            return;
-          }
-
-          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-          rewriter.setInsertionPoint(planOp);
-          FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
-          if (failed(lowered)) {
-            planOp.emitOpError("failed to lower selected row-strip Spatial Relu plan");
-            signalPassFailure();
-            return;
-          }
-          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-          if (failed(output)) {
-            signalPassFailure();
-            return;
-          }
-          rowStripValues[blueprint.getResult()] = *output;
-          eraseAfterLowering.insert(planOp);
-          eraseAfterLowering.insert(blueprint);
-          continue;
-        }
-
-        rewriter.setInsertionPoint(planOp);
-        auto computeOp = createSpatCompute<1>(
-          rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
-            auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
-            spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
-          });
-        rewriter.replaceOp(planOp, computeOp.getResults());
-        continue;
-      }
-
-      if (auto planOp = dyn_cast<spatial::SpatSiluPlanOp>(&op)) {
-        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
-          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-          });
-          if (outputBlueprint == planOp.getResult().getUsers().end()) {
-            planOp.emitOpError("row-strip SiLU plan requires a row-strip blueprint result");
-            signalPassFailure();
-            return;
-          }
-
-          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-          rewriter.setInsertionPoint(planOp);
-          FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
-          if (failed(lowered)) {
-            planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan");
-            signalPassFailure();
-            return;
-          }
-          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-          if (failed(output)) {
-            signalPassFailure();
-            return;
-          }
-          rowStripValues[blueprint.getResult()] = *output;
-          eraseAfterLowering.insert(planOp);
-          eraseAfterLowering.insert(blueprint);
-          continue;
-        }
-
-        rewriter.setInsertionPoint(planOp);
-        auto computeOp = createSpatCompute<1>(
-          rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
-            Value sigmoid = spatial::SpatSigmoidOp::create(
-              rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
-            Value silu = spatial::SpatVMulOp::create(
-              rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
-            spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
-          });
-        rewriter.replaceOp(planOp, computeOp.getResults());
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
-        auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-        });
-        if (outputBlueprint == planOp.getResult().getUsers().end()) {
-          planOp.emitOpError("selected MaxPool plan requires a row-strip blueprint result");
-          signalPassFailure();
-          return;
-        }
-
-        FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-        rewriter.setInsertionPoint(planOp);
-        std::optional<Value> physicalInput;
-        if (succeeded(input))
-          physicalInput = input->storage;
-        FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
-          planOp, physicalInput, rewriter);
-        if (failed(lowered)) {
-          planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
-          signalPassFailure();
-          return;
-        }
-        auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-        FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-        if (failed(output)) {
-          signalPassFailure();
-          return;
-        }
-        rowStripValues[blueprint.getResult()] = *output;
-        eraseAfterLowering.insert(planOp);
-        eraseAfterLowering.insert(blueprint);
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
-        auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-        });
-        if (outputBlueprint == planOp.getResult().getUsers().end()) {
-          planOp.emitOpError("selected global AveragePool plan requires a row-strip blueprint result");
-          signalPassFailure();
-          return;
-        }
-
-        FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-        rewriter.setInsertionPoint(planOp);
-        std::optional<Value> physicalInput;
-        if (succeeded(input))
-          physicalInput = input->storage;
-        FailureOr<Value> lowered =
-          lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, rewriter);
-        if (failed(lowered)) {
-          planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan");
-          signalPassFailure();
-          return;
-        }
-        auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-        FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-        if (failed(output)) {
-          signalPassFailure();
-          return;
-        }
-        rowStripValues[blueprint.getResult()] = *output;
-        eraseAfterLowering.insert(planOp);
-        eraseAfterLowering.insert(blueprint);
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
-        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
-          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
-            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
-            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
-          });
-          if (outputBlueprint == planOp.getResult().getUsers().end()) {
-            planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result");
-            signalPassFailure();
-            return;
-          }
-
-          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
-          rewriter.setInsertionPoint(planOp);
-          FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
-          if (failed(lowered)) {
-            planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
-            signalPassFailure();
-            return;
-          }
-          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
-          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
-          if (failed(output)) {
-            signalPassFailure();
-            return;
-          }
-          rowStripValues[blueprint.getResult()] = *output;
-          eraseAfterLowering.insert(planOp);
-          eraseAfterLowering.insert(blueprint);
-          continue;
-        }
-
-        auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
-        if (!resultType) {
-          planOp.emitOpError("requires ranked output type");
-          signalPassFailure();
-          return;
-        }
-        rewriter.setInsertionPoint(planOp);
-        FailureOr<Value> denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc());
-        if (failed(denseBias)) {
-          planOp.emitOpError("failed to materialize dense Conv-style bias");
-          signalPassFailure();
-          return;
-        }
-        if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
-          FailureOr<Value> lowered = lowerDenseBatchBiasAdd(planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
-          if (succeeded(lowered)) {
-            rewriter.replaceOp(planOp, *lowered);
-            continue;
-          }
-        }
-        auto computeOp = createSpatCompute<2>(rewriter,
-                                              planOp.getLoc(),
-                                              planOp.getOutput().getType(),
-                                              {},
-                                              ValueRange {planOp.getInput(), *denseBias},
-                                              [&](Value x, Value y) {
-                                                auto added = spatial::SpatVAddOp::create(
-                                                  rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
-                                                spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
-                                              });
-        rewriter.replaceOp(planOp, computeOp.getResults());
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
-        if (failed(lowerAddPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
-          signalPassFailure();
-          return;
-        }
-        continue;
-      }
-      if (auto planOp = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
-        if (failed(lowerConcatPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
-          signalPassFailure();
-          return;
-        }
-        continue;
-      }
-      if (auto flattenOp = dyn_cast<spatial::SpatGraphCompute>(&op)) {
-        if (flattenOp.getInputs().size() == 1) {
-          FailureOr<RowStripPhysicalValue> input =
-            getRowStripValue(rowStripValues, flattenOp.getInputs().front());
-          if (succeeded(input) && succeeded(canLowerFlattenFromRowStrip(flattenOp))) {
-            rewriter.setInsertionPoint(flattenOp);
-            if (failed(lowerFlattenFromRowStrip(*input, flattenOp, rewriter))) {
-              flattenOp.emitOpError("failed to preserve row-strip layout through Flatten");
-              signalPassFailure();
-              return;
-            }
-            continue;
-          }
-        }
-      }
-      if (auto materializeOp = dyn_cast<spatial::SpatMaterializeLayoutOp>(&op)) {
-        if (materializeOp.getSourcePhysicalLayout() == kDenseLayout
-            && materializeOp.getTargetPhysicalLayout() == kDenseLayout) {
-          rewriter.replaceOp(materializeOp, materializeOp.getInput());
-          continue;
-        }
-        if (materializeOp.getSourcePhysicalLayout() != kRowStripLayout
-            || materializeOp.getTargetPhysicalLayout() != kDenseLayout) {
-          materializeOp.emitOpError("non-dense materialize_layout lowering is not supported yet");
-          signalPassFailure();
-          return;
-        }
-        FailureOr<RowStripPhysicalValue> rowStripValue = getRowStripValue(rowStripValues, materializeOp.getInput());
-        if (failed(rowStripValue)) {
-          materializeOp.emitOpError("expected a row-strip blueprint input during row-strip materialization");
-          signalPassFailure();
-          return;
-        }
-        rewriter.setInsertionPoint(materializeOp);
-        FailureOr<Value> dense = materializeRowStripToDense(*rowStripValue, materializeOp.getLoc(), rewriter);
-        if (failed(dense)) {
-          materializeOp.emitOpError("failed to materialize selected row-strip layout back to dense NCHW");
-          signalPassFailure();
-          return;
-        }
-        rewriter.replaceOp(materializeOp, *dense);
-        continue;
-      }
-      if (auto blueprintOp = dyn_cast<spatial::SpatBlueprintOp>(&op)) {
-        if (std::optional<StringRef> mode = blueprintOp.getMode(); mode && *mode == "fragment_assembly")
-          continue;
-        if (blueprintOp.getPhysicalLayout() == kDenseLayout) {
-          rewriter.replaceOp(blueprintOp, blueprintOp.getInput());
-          continue;
-        }
-        if (blueprintOp.getPhysicalLayout() != kRowStripLayout) {
-          blueprintOp.emitOpError("non-dense blueprint lowering is not supported yet");
-          signalPassFailure();
-          return;
-        }
-        if (!eraseAfterLowering.contains(blueprintOp)) {
-          blueprintOp.emitOpError("unhandled row-strip blueprint remained during LowerSpatialPlans");
-          signalPassFailure();
-          return;
-        }
-      }
-    }
-    bool erasedAny = true;
-    while (erasedAny) {
-      erasedAny = false;
-      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
-        if (!eraseAfterLowering.contains(&op))
-          continue;
-        if (!op.use_empty())
-          continue;
-        eraseAfterLowering.erase(&op);
-        rewriter.eraseOp(&op);
-        erasedAny = true;
-      }
-    }
-    if (!eraseAfterLowering.empty()) {
-      for (Operation& op : funcOp.getBody().front())
-        if (eraseAfterLowering.contains(&op))
-          op.emitOpError("selected row-strip planning op could not be fully eliminated during LowerSpatialPlans");
-      signalPassFailure();
-      return;
-    }
-    ConversionTarget helperTarget(*ctx);
-    helperTarget.addLegalDialect<spatial::SpatialDialect,
-                                 tensor::TensorDialect,
-                                 linalg::LinalgDialect,
-                                 affine::AffineDialect,
-                                 arith::ArithDialect,
-                                 scf::SCFDialect,
-                                 func::FuncDialect>();
-    helperTarget.addLegalOp<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();
-    helperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
-    helperTarget.markOpRecursivelyLegal<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();
-
-    RewritePatternSet helperPatterns(ctx);
-    populateGemmPatterns(helperPatterns, ctx);
-    populateTransposePatterns(helperPatterns, ctx);
-    FrozenRewritePatternSet frozenHelperPatterns(
-      std::move(helperPatterns));
-    SmallVector<Operation*> topLevelHelperOps;
-    funcOp.walk([&](Operation* op) {
-      if (isa<spatial::SpatGraphCompute,
-              spatial::SpatGraphComputeBatch>(op))
-        return WalkResult::skip();
-      if (isa<ONNXGemmOp, ONNXTransposeOp>(op))
-        topLevelHelperOps.push_back(op);
-      return WalkResult::advance();
-    });
-    for (Operation *helper : topLevelHelperOps) {
-      if (failed(applyPartialConversion(
-            helper, helperTarget, frozenHelperPatterns))) {
-        moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
-        signalPassFailure();
-        return;
-      }
-    }
-    ConversionTarget nestedHelperTarget(*ctx);
-    nestedHelperTarget.addLegalDialect<spatial::SpatialDialect,
-                                       tensor::TensorDialect,
-                                       linalg::LinalgDialect,
-                                       affine::AffineDialect,
-                                       arith::ArithDialect,
-                                       scf::SCFDialect,
-                                       func::FuncDialect>();
-    nestedHelperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
-    SmallVector<Operation*> computeLikeOps;
-    funcOp.walk([&](Operation* op) {
-      if (isa<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>(op))
-        computeLikeOps.push_back(op);
-    });
-    for (Operation* op : computeLikeOps) {
-      if (failed(applyFullConversion(
-            op, nestedHelperTarget, frozenHelperPatterns))) {
-        op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering");
-        signalPassFailure();
-        return;
-      }
-    }
-    if (!verifyLogicalPhase("after nested helper conversions"))
-      return;
-    bool hasIllegalOps = false;
-    moduleOp.walk([&](Operation* op) {
-      if (isa<ONNXEntryPointOp>(op))
-        return;
-      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-        if (std::optional<StringRef> mode = blueprint.getMode(); mode && *mode == "fragment_assembly")
-          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::SpatMaxPool2DPlanOp,
-                   spatial::SpatGlobalAveragePoolPlanOp,
-                   spatial::SpatMaterializeLayoutOp>(op)
-               || op->getDialect()->getNamespace() == "onnx") {
-        op->emitOpError("operation must not remain after LowerSpatialPlans");
-        hasIllegalOps = true;
-      }
-    });
-
-    PassManager canonicalizationPM(ctx);
-    canonicalizationPM.addPass(createCanonicalizerPass());
-    if (failed(canonicalizationPM.run(moduleOp)))
-      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");
-
-    if (hasIllegalOps) {
-      signalPassFailure();
-    } else {
-      dumpModule(moduleOp, "spatial1_graph");
-      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
-      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
-          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
-        signalPassFailure();
-        return;
-      }
-    }
-
-    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
-      return;
-  }
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp

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

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

@@ -6,7 +6,7 @@
 #include "Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -108,7 +108,9 @@ void verifyScheduledInputs(ComputeOpTy compute,
   for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
     size_t currentInputIndex = inputIndex;
     Operation* definingOp = input.getDefiningOp();
-    if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
+    if (allowChannelReceiveInputs
+        && isa_and_nonnull<spatial::SpatChannelReceiveOp,
+                           spatial::SpatHostWaitLoadOp>(definingOp))
       continue;
     if (isScheduledPhase1Value(input))
       continue;
@@ -130,8 +132,7 @@ template <typename ComputeOpTy>
 void verifyNoNestedFragmentAssemblyBlueprints(ComputeOpTy compute,
                                                   pim::CappedDiagnosticReporter& diagnostics) {
   compute.getBody().walk([&](spatial::SpatBlueprintOp blueprint) {
-    std::optional<StringRef> mode = blueprint.getMode();
-    if (!mode || *mode != "fragment_assembly")
+    if (!spatial::isFragmentAssembly(blueprint.getMode()))
       return;
     diagnostics.report(blueprint.getOperation(), [&](Operation* illegalOp) {
       illegalOp->emitOpError("fragment assembly blueprint must be host-level after merge materialization");
@@ -148,8 +149,10 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatBiasAddPlanOp,
             spatial::SpatAddPlanOp,
             spatial::SpatConcatPlanOp,
+            spatial::SpatFlattenPlanOp,
             spatial::SpatReluPlanOp,
             spatial::SpatSiluPlanOp,
+            spatial::SpatResizeNearestPlanOp,
             spatial::SpatMaxPool2DPlanOp,
             spatial::SpatGlobalAveragePoolPlanOp,
             spatial::SpatBlueprintOp,
@@ -162,7 +165,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
       });
       continue;
     }
-    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp>(&op)) {
+    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp,
+            spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
         illegalOp->emitOpError() << kPhaseMarker
                                  << " explicit channel communication is not expected before merge materialization";
@@ -181,7 +185,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
 
 void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
   for (Operation& op : funcOp.getOps()) {
-    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp,
+            spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
         illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
       });

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

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

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

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

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

@@ -12,15 +12,15 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
-#include "ONNXToSpatialVerifier.hpp"
 
 using namespace mlir;
 
@@ -34,9 +34,17 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
   StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
 
   ONNXToSpatialPass() = default;
-  ONNXToSpatialPass(const ONNXToSpatialPass& pass) {}
+  ONNXToSpatialPass(const spatial::SpatialTargetResources& target,
+                    const ONNXToSpatialPlanningOptions& options)
+      : target(target), planningOptions(options), hasTarget(true) {}
+  ONNXToSpatialPass(const ONNXToSpatialPass& pass)
+      : target(pass.target), planningOptions(pass.planningOptions), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() override;
+
+  spatial::SpatialTargetResources target;
+  ONNXToSpatialPlanningOptions planningOptions;
+  bool hasTarget = false;
 };
 
 } // namespace
@@ -50,15 +58,19 @@ 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(
+    funcOp.getOps<spatial::SpatResizeNearestPlanOp>());
   SmallVector<spatial::SpatMaxPool2DPlanOp> maxPoolPlans(funcOp.getOps<spatial::SpatMaxPool2DPlanOp>());
   SmallVector<spatial::SpatGlobalAveragePoolPlanOp> globalAveragePoolPlans(
     funcOp.getOps<spatial::SpatGlobalAveragePoolPlanOp>());
   SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
   SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
   if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !addPlans.empty()
-      || !concatPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !maxPoolPlans.empty() || !blueprints.empty()
+      || !concatPlans.empty() || !flattenPlans.empty() || !reluPlans.empty() || !siluPlans.empty() || !resizePlans.empty()
+      || !maxPoolPlans.empty() || !blueprints.empty()
       || !globalAveragePoolPlans.empty() || !materializers.empty()) {
     return;
   }
@@ -103,6 +115,11 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
 
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
+  if (!hasTarget) {
+    moduleOp.emitError("ONNX-to-Spatial lowering requires an injected SpatialTargetResources");
+    signalPassFailure();
+    return;
+  }
   MLIRContext* ctx = &getContext();
 
   ConversionTarget preTarget(*ctx);
@@ -123,6 +140,14 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  RewritePatternSet matmulPatterns(ctx);
+  populateMatMulFusionPatterns(matmulPatterns, ctx, target);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(matmulPatterns)))) {
+    moduleOp.emitError("failed to lower MatMul before producer conversion");
+    signalPassFailure();
+    return;
+  }
+
   RewritePatternSet fusionPatterns(ctx);
   populateElementwiseFusionPatterns(fusionPatterns, ctx);
   if (failed(applyPatternsGreedily(moduleOp, std::move(fusionPatterns)))) {
@@ -171,7 +196,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXSplitOp>();
 
   RewritePatternSet conversionPatterns(ctx);
-  populateConversionPatterns(conversionPatterns, ctx);
+  populateConversionPatterns(conversionPatterns, ctx, this->target, planningOptions);
   if (failed(applyPartialConversion(moduleOp, target, std::move(conversionPatterns)))) {
     moduleOp.emitError("failed to convert required ONNX ops to Spatial ops");
     signalPassFailure();
@@ -247,4 +272,10 @@ void ONNXToSpatialPass::runOnOperation() {
 
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
+std::unique_ptr<Pass> createONNXToSpatialPass(
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  return std::make_unique<ONNXToSpatialPass>(target, options);
+}
+
 } // namespace onnx_mlir

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

@@ -0,0 +1,92 @@
+#pragma once
+
+#include <optional>
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+
+struct RowStripPhysicalValue;
+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 ONNXToSpatialPlanningOptions& options,
+                        mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
+                                               const spatial::SpatialTargetResources& target,
+                                               const ONNXToSpatialPlanningOptions* options = nullptr);
+mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
+                                                 const spatial::SpatialTargetResources& target,
+                                                 const ONNXToSpatialPlanningOptions* options = nullptr);
+
+mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
+  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetResources& target);
+
+mlir::FailureOr<mlir::Value> lowerSelectedResizeNearestPlan(
+  spatial::SpatResizeNearestPlanOp planOp,
+  mlir::Value input,
+  std::optional<mlir::Value> rowStripInput,
+  const spatial::SpatialTargetResources& target,
+  mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
+                                                 const spatial::SpatialTargetResources& target);
+
+mlir::FailureOr<mlir::Value>
+lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                        mlir::Value input,
+                        const spatial::SpatialTargetResources& target,
+                       mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                           mlir::Value input,
+                           std::optional<mlir::Value> rowStripInput,
+                           const spatial::SpatialTargetResources& target,
+                           mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult
+canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                        const spatial::SpatialTargetResources& target);
+
+mlir::FailureOr<mlir::Value>
+lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                mlir::Value input,
+                                const spatial::SpatialTargetResources& target,
+                                mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                   mlir::Value input,
+                                   std::optional<mlir::Value> rowStripInput,
+                                   const spatial::SpatialTargetResources& target,
+                                   mlir::PatternRewriter& rewriter);
+
+} // namespace onnx_mlir

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

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

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

@@ -0,0 +1,845 @@
+#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/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/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/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<RowStripPhysicalValue> getRowStripValue(Value value) {
+  return getRowStripPhysicalValue(value);
+}
+
+static FailureOr<Value> publishRowStripValue(Operation* planOp,
+                                             Value storage,
+                                             PatternRewriter& rewriter) {
+  auto logicalType = dyn_cast<RankedTensorType>(planOp->getResult(0).getType());
+  if (!logicalType)
+    return planOp->emitOpError("requires ranked logical output type"), failure();
+  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
+  if (failed(value))
+    return planOp->emitOpError("lowering produced invalid row-strip physical storage"), failure();
+  FailureOr<Value> blueprint = createRowStripStorageBlueprint(
+      storage, logicalType, rewriter, planOp->getLoc());
+  if (failed(blueprint))
+    return planOp->emitOpError("failed to create row-strip storage Blueprint"), failure();
+  rewriter.replaceOp(planOp, *blueprint);
+  return *blueprint;
+}
+
+static bool isRowStripSelected(Operation* op) {
+  auto selected = spatial::getSelectedPhysicalLayout(op);
+  return selected && *selected == spatial::PhysicalLayout::NHWCRowStrip;
+}
+
+static bool isDenseSelected(Operation* op) {
+  auto selected = spatial::getSelectedPhysicalLayout(op);
+  return selected && *selected == spatial::PhysicalLayout::DenseNCHW;
+}
+
+static spatial::PhysicalLayout getKnownPhysicalLayout(Value value) {
+  return getSpatialPlanOperandLayout(value);
+}
+
+static LogicalResult verifySelectedLayouts(
+    func::FuncOp funcOp, const spatial::SpatialTargetResources& target) {
+  LogicalResult result = success();
+  funcOp.walk([&](Operation* op) {
+    auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
+    if (!capability)
+      return;
+    auto selected = spatial::getSelectedPhysicalLayout(op);
+    if (!selected) {
+      op->emitOpError("requires a selected physical layout from SpatialLayoutPlanning");
+      result = failure();
+      return;
+    }
+    if (*selected != spatial::PhysicalLayout::DenseNCHW
+        && *selected != spatial::PhysicalLayout::NHWCRowStrip) {
+      op->emitOpError("has an unsupported selected physical layout");
+      result = failure();
+      return;
+    }
+    SmallVector<spatial::PhysicalLayout> operandLayouts;
+    operandLayouts.reserve(op->getNumOperands());
+    for (Value operand : op->getOperands())
+      operandLayouts.push_back(getKnownPhysicalLayout(operand));
+    auto alternatives = capability.getLayoutAlternatives(target, operandLayouts);
+    if (llvm::none_of(alternatives, [&](const spatial::LayoutAlternative& alternative) {
+          return alternative.resultLayout == *selected
+              && alternative.operandLayouts == operandLayouts;
+        })) {
+      op->emitOpError("selected physical layout is not lowerable for its explicit operand layouts");
+      result = failure();
+    }
+  });
+  return result;
+}
+
+static FailureOr<Value>
+lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripRelu(input, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value>
+lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripSilu(input, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
+                                         const RowStripPhysicalValue& rhs,
+                                         spatial::SpatAddPlanOp planOp,
+                                         PatternRewriter& rewriter) {
+  return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
+                                            spatial::SpatConcatPlanOp planOp,
+                                            PatternRewriter& rewriter) {
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!outputType)
+    return failure();
+  return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value>
+materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
+  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
+    return failure();
+  return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc);
+}
+
+static FailureOr<Value> materializeDenseToRowStrip(
+    Value input, RankedTensorType logicalType, Location loc, PatternRewriter& rewriter) {
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4
+      || logicalType.getDimSize(0) != 1)
+    return failure();
+  auto nhwcType = RankedTensorType::get(
+    {1, logicalType.getDimSize(2), logicalType.getDimSize(3), logicalType.getDimSize(1)},
+    logicalType.getElementType(), logicalType.getEncoding());
+  auto rowsType = RankedTensorType::get(
+    {logicalType.getDimSize(2) * logicalType.getDimSize(3), logicalType.getDimSize(1)},
+    logicalType.getElementType(), logicalType.getEncoding());
+  auto rowsCompute = createSpatCompute<1>(
+    rewriter, loc, rowsType, {}, input, [&](Value denseInput) {
+      Value nhwc = createLinalgTranspose(
+        denseInput, nhwcType, {0, 2, 3, 1}, rewriter, loc);
+      Value rows = tensor::CollapseShapeOp::create(
+        rewriter, loc, rowsType, nhwc,
+        SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
+      spatial::SpatYieldOp::create(rewriter, loc, rows);
+    });
+  Value rows = rowsCompute->getResult(0);
+  FailureOr<Value> storage = createRowStripStorageFromRows(rows, logicalType, rewriter, loc);
+  if (failed(storage))
+    return failure();
+  return createRowStripStorageBlueprint(*storage, logicalType, rewriter, loc);
+}
+
+static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
+                                               PatternRewriter& rewriter, Location loc) {
+  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
+  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
+      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
+      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
+    return failure();
+  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
+  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
+  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
+      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
+      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
+    return failure();
+  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
+    if (outputDim > inputDim)
+      return failure();
+
+  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
+      if (failed(fragment))
+        return failure();
+      MixedSliceGeometry biasSlice;
+      for (int64_t dim : inputFragmentType->getShape()) {
+        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
+        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
+      }
+      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
+      if (!biasFragment)
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
+      MixedSliceGeometry outputSlice;
+      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
+      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
+      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
+      for (int64_t dim : outputFragmentType->getShape())
+        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
+      if (!output)
+        return failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
+struct LowerDenseReluPlan final : OpConversionPattern<spatial::SpatReluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
+                                spatial::SpatReluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
+    if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
+      return failure();
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, adaptor.getInput(), [&](Value x) {
+        auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseSiluPlan final : OpConversionPattern<spatial::SpatSiluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
+                                spatial::SpatSiluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
+    if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
+      return failure();
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, adaptor.getInput(), [&](Value x) {
+        Value sigmoid = spatial::SpatSigmoidOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
+        Value silu = spatial::SpatVMulOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseResizePlan final : OpConversionPattern<spatial::SpatResizeNearestPlanOp> {
+  explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
+                                spatial::SpatResizeNearestPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(
+        planOp, adaptor.getInput(), std::nullopt, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense nearest Resize plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerDenseBiasAddPlan final : OpConversionPattern<spatial::SpatBiasAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
+                                spatial::SpatBiasAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+    if (!resultType)
+      return planOp.emitOpError("requires ranked output type");
+
+    FailureOr<Value> denseBias = materializeDenseBiasAddTensor(
+      adaptor.getBias(), resultType, rewriter, planOp.getLoc());
+    if (failed(denseBias))
+      return planOp.emitOpError("failed to materialize dense Conv-style bias");
+    if (adaptor.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+      FailureOr<Value> lowered = lowerDenseBatchBiasAdd(
+        adaptor.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+      if (succeeded(lowered)) {
+        rewriter.replaceOp(planOp, *lowered);
+        return success();
+      }
+    }
+    auto computeOp = createSpatCompute<2>(
+      rewriter,
+      planOp.getLoc(),
+      planOp.getOutput().getType(),
+      {},
+      ValueRange {adaptor.getInput(), *denseBias},
+      [&](Value x, Value y) {
+        auto added = spatial::SpatVAddOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
+      });
+    rewriter.replaceOp(planOp, computeOp.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseAddPlan final : OpConversionPattern<spatial::SpatAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
+                                spatial::SpatAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto compute = createSpatCompute<2>(
+      rewriter,
+      planOp.getLoc(),
+      planOp.getOutput().getType(),
+      {},
+      ValueRange {adaptor.getLhs(), adaptor.getRhs()},
+      [&](Value lhsValue, Value rhsValue) {
+        Value added = spatial::SpatVAddOp::create(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added);
+      });
+    rewriter.replaceOp(planOp, compute.getResults());
+    return success();
+  }
+};
+
+struct LowerDenseConcatPlan final : OpConversionPattern<spatial::SpatConcatPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
+                                spatial::SpatConcatPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    auto compute = createSpatCompute(
+      rewriter,
+      planOp.getLoc(),
+      TypeRange {planOp.getOutput().getType()},
+      {},
+      adaptor.getInputs(),
+      [&](ValueRange values) {
+        Value concatenated = spatial::SpatConcatOp::create(
+          rewriter,
+          planOp.getLoc(),
+          planOp.getOutput().getType(),
+          rewriter.getI64IntegerAttr(planOp.getAxis()),
+          values);
+        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated);
+      });
+    rewriter.replaceOp(planOp, compute.getResults());
+    return success();
+  }
+};
+
+static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
+                                  Value lhsValue, Value rhsValue,
+                                  PatternRewriter& rewriter) {
+  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(lhsValue);
+  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rhsValue);
+  if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) {
+    if (getKnownPhysicalLayout(lhsValue) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation()) && failed(rhs)) {
+    if (getKnownPhysicalLayout(rhsValue) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation())) {
+    rewriter.setInsertionPoint(planOp);
+    FailureOr<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+  return planOp.emitOpError("dense Add plan was not lowered by the selected-plan patterns");
+}
+
+static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
+                                     ValueRange inputValues,
+                                     PatternRewriter& rewriter) {
+  SmallVector<RowStripPhysicalValue> inputs;
+  for (Value input : inputValues) {
+    FailureOr<RowStripPhysicalValue> physical = getRowStripValue(input);
+    if (failed(physical)) {
+      inputs.clear();
+      break;
+    }
+    inputs.push_back(*physical);
+  }
+  if (isRowStripSelected(planOp.getOperation()) && inputs.size() != inputValues.size()) {
+    if (llvm::any_of(inputValues, [](Value input) {
+          return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip;
+        }))
+      return failure();
+    return planOp.emitOpError("selected row-strip Concat plan requires row-strip inputs");
+  }
+  if (isRowStripSelected(planOp.getOperation())) {
+    rewriter.setInsertionPoint(planOp);
+    FailureOr<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+  return planOp.emitOpError("dense Concat plan was not lowered by the selected-plan patterns");
+}
+
+struct LowerSelectedConvPlan final : OpConversionPattern<spatial::SpatConv2DPlanOp> {
+  explicit LowerSelectedConvPlan(MLIRContext* ctx,
+                                 const spatial::SpatialTargetResources& target,
+                                 const ONNXToSpatialPlanningOptions& options)
+      : OpConversionPattern<spatial::SpatConv2DPlanOp>(ctx), target(target), options(options) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatConv2DPlanOp planOp,
+                                spatial::SpatConv2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (isDenseSelected(planOp.getOperation())) {
+      FailureOr<Value> lowered = lowerSelectedConv2DPlan(
+          planOp, adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(),
+          std::nullopt, /*emitRowStripLayout=*/false, target, options, rewriter);
+      if (failed(lowered))
+        return planOp.emitOpError("failed to lower selected dense Spatial Conv plan");
+      rewriter.replaceOp(planOp, *lowered);
+      return success();
+    }
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+
+    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(adaptor.getInput());
+    if (failed(rowStripInput)
+        && getKnownPhysicalLayout(adaptor.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(),
+        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)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+  const ONNXToSpatialPlanningOptions& options;
+};
+
+struct LowerRowStripReluPlan final : OpConversionPattern<spatial::SpatReluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
+                                spatial::SpatReluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial ReLU plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripSiluPlan final : OpConversionPattern<spatial::SpatSiluPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
+                                spatial::SpatSiluPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripResizePlan final : OpConversionPattern<spatial::SpatResizeNearestPlanOp> {
+  explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
+                                spatial::SpatResizeNearestPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip Resize plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(
+        planOp, 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)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerDenseMaxPoolPlan final : OpConversionPattern<spatial::SpatMaxPool2DPlanOp> {
+  explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
+                                spatial::SpatMaxPool2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerDenseMaxPool2DPlan(
+        planOp, adaptor.getInput(), target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense Spatial MaxPool plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerRowStripMaxPoolPlan final : OpConversionPattern<spatial::SpatMaxPool2DPlanOp> {
+  explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
+                                spatial::SpatMaxPool2DPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    std::optional<Value> physicalInput;
+    if (succeeded(input))
+      physicalInput = input->storage;
+    FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(
+        planOp, adaptor.getInput(), physicalInput, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerRowStripGlobalAveragePoolPlan
+    final : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp> {
+  explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)
+        && getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+      return failure();
+    std::optional<Value> physicalInput;
+    if (succeeded(input))
+      physicalInput = input->storage;
+    FailureOr<Value> lowered = lowerSelectedGlobalAveragePoolPlan(
+        planOp, adaptor.getInput(), physicalInput, target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerDenseGlobalAveragePoolPlan
+    final : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp> {
+  explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx,
+                                           const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
+
+  LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                spatial::SpatGlobalAveragePoolPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isDenseSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<Value> lowered = lowerDenseGlobalAveragePoolPlan(
+        planOp, adaptor.getInput(), target, rewriter);
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected dense Spatial global AveragePool plan");
+    rewriter.replaceOp(planOp, *lowered);
+    return success();
+  }
+
+  const spatial::SpatialTargetResources& target;
+};
+
+struct LowerRowStripBiasAddPlan final : OpConversionPattern<spatial::SpatBiasAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
+                                spatial::SpatBiasAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    FailureOr<RowStripPhysicalValue> input = getRowStripValue(adaptor.getInput());
+    if (failed(input)) {
+      if (getKnownPhysicalLayout(adaptor.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
+        return failure();
+      return planOp.emitOpError("selected row-strip bias_add plan requires a row-strip input");
+    }
+    FailureOr<Value> lowered = applyRowStripBiasAdd(
+        *input, adaptor.getBias(), rewriter, planOp.getLoc());
+    if (failed(lowered))
+      return planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
+    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
+      return failure();
+    return success();
+  }
+};
+
+struct LowerRowStripAddPlan final : OpConversionPattern<spatial::SpatAddPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
+                                spatial::SpatAddPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    return lowerAddPlan(planOp, adaptor.getLhs(), adaptor.getRhs(), rewriter);
+  }
+};
+
+struct LowerRowStripConcatPlan final : OpConversionPattern<spatial::SpatConcatPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
+                                spatial::SpatConcatPlanOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    if (!isRowStripSelected(planOp.getOperation()))
+      return failure();
+    return lowerConcatPlan(planOp, adaptor.getInputs(), rewriter);
+  }
+};
+
+struct LowerMaterializeLayout final
+    : OpConversionPattern<spatial::SpatMaterializeLayoutOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp,
+                                spatial::SpatMaterializeLayoutOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto source = materializeOp.getSourcePhysicalLayout();
+    auto target = materializeOp.getTargetPhysicalLayout();
+    if (source == spatial::PhysicalLayout::DenseNCHW
+        && target == spatial::PhysicalLayout::DenseNCHW) {
+      rewriter.replaceOp(materializeOp, adaptor.getInput());
+      return success();
+    }
+    if (source == spatial::PhysicalLayout::DenseNCHW
+        && target == spatial::PhysicalLayout::NHWCRowStrip) {
+      auto logicalType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+      if (!logicalType)
+        return materializeOp.emitOpError("requires a ranked dense input"), failure();
+      FailureOr<Value> rowStrip = materializeDenseToRowStrip(
+        adaptor.getInput(), logicalType, materializeOp.getLoc(), rewriter);
+      if (failed(rowStrip))
+        return materializeOp.emitOpError(
+          "failed to materialize dense NCHW storage to row-strip layout"), failure();
+      rewriter.replaceOp(materializeOp, *rowStrip);
+      return success();
+    }
+    if (source != spatial::PhysicalLayout::NHWCRowStrip
+        || target != spatial::PhysicalLayout::DenseNCHW)
+      return materializeOp.emitOpError(
+          "unsupported Spatial layout materialization direction"), failure();
+    auto inputType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+    if (!inputType)
+      return materializeOp.emitOpError("requires a ranked row-strip input"), failure();
+    FailureOr<RowStripPhysicalValue> rowStripValue =
+        getRowStripValue(adaptor.getInput());
+    if (failed(rowStripValue))
+      return failure();
+    FailureOr<Value> dense = materializeRowStripToDense(
+        *rowStripValue, materializeOp.getLoc(), rewriter);
+    if (failed(dense))
+      return materializeOp.emitOpError(
+          "failed to materialize row-strip storage to dense NCHW"), failure();
+    rewriter.replaceOp(materializeOp, *dense);
+    return success();
+  }
+};
+
+struct LowerSelectedFlattenPlan final
+    : OpConversionPattern<spatial::SpatFlattenPlanOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatFlattenPlanOp planOp,
+                                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())
+      return failure();
+    rewriter.eraseOp(blueprint);
+    return success();
+  }
+};
+
+static void populateConvPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  patterns.add<LowerSelectedConvPlan>(ctx, target, options);
+}
+
+static void populateElementwisePlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerDenseReluPlan,
+               LowerRowStripReluPlan,
+               LowerDenseSiluPlan,
+               LowerRowStripSiluPlan,
+               LowerDenseBiasAddPlan,
+               LowerRowStripBiasAddPlan,
+               LowerDenseAddPlan,
+               LowerRowStripAddPlan>(ctx);
+}
+
+static void populatePoolPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<LowerDenseMaxPoolPlan,
+               LowerRowStripMaxPoolPlan,
+               LowerDenseGlobalAveragePoolPlan,
+               LowerRowStripGlobalAveragePoolPlan>(ctx, target);
+}
+
+static void populateResizePlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<LowerDenseResizePlan, LowerRowStripResizePlan>(ctx, target);
+}
+
+static void populateConcatPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerDenseConcatPlan, LowerRowStripConcatPlan>(ctx);
+}
+
+static void populateFlattenPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target) {
+  patterns.add<LowerSelectedFlattenPlan>(ctx, target);
+}
+
+static void populateLayoutMaterializationPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<LowerMaterializeLayout, EraseDeadPhysicalViewBlueprint>(ctx);
+}
+
+} // namespace
+
+void populateSpatialPlanLoweringPatterns(
+    RewritePatternSet& patterns, MLIRContext* ctx,
+    const spatial::SpatialTargetResources& target,
+    const ONNXToSpatialPlanningOptions& options) {
+  populateConvPlanLoweringPatterns(patterns, ctx, target, options);
+  populateElementwisePlanLoweringPatterns(patterns, ctx);
+  populatePoolPlanLoweringPatterns(patterns, ctx, target);
+  populateResizePlanLoweringPatterns(patterns, ctx, target);
+  populateConcatPlanLoweringPatterns(patterns, ctx);
+  populateFlattenPlanLoweringPatterns(patterns, ctx, target);
+  populateLayoutMaterializationPatterns(patterns, ctx);
+}
+
+LogicalResult verifySelectedSpatialLayouts(
+    func::FuncOp funcOp, const spatial::SpatialTargetResources& target) {
+  return verifySelectedLayouts(funcOp, target);
+}
+
+} // namespace onnx_mlir

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

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

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -7,12 +7,15 @@ namespace onnx_mlir {
 
 void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { populateGeneratedPrePatterns(patterns, ctx); }
 
-void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+void populateConversionPatterns(RewritePatternSet& patterns,
+                                MLIRContext* ctx,
+                                const spatial::SpatialTargetResources& target,
+                                const ONNXToSpatialPlanningOptions& options) {
   populateElementwisePatterns(patterns, ctx);
-  populateMatMulRewritePatterns(patterns, ctx);
-  populateGemmPatterns(patterns, ctx);
-  populateConvPatterns(patterns, ctx);
-  populatePoolPatterns(patterns, ctx);
+  populateMatMulRewritePatterns(patterns, ctx, target);
+  populateGemmPatterns(patterns, ctx, target);
+  populateConvPatterns(patterns, ctx, target, options);
+  populatePoolPatterns(patterns, ctx, target);
   populateReduceMeanPatterns(patterns, ctx);
   populateReluPatterns(patterns, ctx);
   populateSigmoidPatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -4,23 +4,43 @@
 #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;
+}
+
 void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateConversionPatterns(mlir::RewritePatternSet& patterns,
+                                mlir::MLIRContext* ctx,
+                                const spatial::SpatialTargetResources& target,
+                                const ONNXToSpatialPlanningOptions& options);
 void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
-void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateConvPatterns(mlir::RewritePatternSet& patterns,
+                          mlir::MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target,
+                          const ONNXToSpatialPlanningOptions& options);
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateElementwiseFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateGemmPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateGemmPatterns(mlir::RewritePatternSet& patterns,
+                          mlir::MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target);
+void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns,
+                                   mlir::MLIRContext* ctx,
+                                   const spatial::SpatialTargetResources& target);
+void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns,
+                                  mlir::MLIRContext* ctx,
+                                  const spatial::SpatialTargetResources& target);
+void populatePoolPatterns(mlir::RewritePatternSet& patterns,
+                          mlir::MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target);
 void populateReduceMeanPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSigmoidPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

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

@@ -1,49 +1,133 @@
 #include "ConvGeometry.hpp"
 
 #include <algorithm>
+#include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 namespace onnx_mlir {
 
+namespace {
+
+static const ONNXToSpatialPlanningOptions& defaultPlanningOptions() {
+  static const ONNXToSpatialPlanningOptions options {
+    std::numeric_limits<uint64_t>::max(),
+    std::numeric_limits<uint64_t>::max(),
+    spatial::ConvLoweringStrategy::Auto,
+    false,
+  };
+  return options;
+}
+
+} // namespace
+
+const ONNXToSpatialPlanningOptions& ConvLoweringState::planningOptions() const {
+  return options ? *options : defaultPlanningOptions();
+}
+
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup) {
   return group == numChannelsIn && numChannelsInPerGroup == 1 && numChannelsOut % group == 0;
 }
 
-ConvGeometry buildConvGeometry(const ConvLoweringState& state) {
+void classifyConvProblem(ConvProblem& problem) {
+  problem.isDepthwise = isDepthwiseConv(
+    problem.group, problem.numChannelsIn, problem.numChannelsOut,
+    problem.numChannelsInPerGroup);
+  problem.isGrouped = problem.group > 1;
+  problem.isPointwise = problem.wHeight == 1 && problem.wWidth == 1
+    && problem.strideHeight == 1 && problem.strideWidth == 1
+    && problem.dilationHeight == 1 && problem.dilationWidth == 1
+    && problem.padHeightBegin == 0 && problem.padHeightEnd == 0
+    && problem.padWidthBegin == 0 && problem.padWidthEnd == 0;
+}
+
+ConvGeometry buildConvGeometry(const ConvProblem& problem,
+                               const spatial::SpatialTargetResources& target) {
   ConvGeometry geo {
-    state.batchSize,
-    state.numChannelsIn,
-    state.xHeight,
-    state.xWidth,
-    state.numChannelsOut,
-    state.wHeight,
-    state.wWidth,
-    state.outHeight,
-    state.outWidth,
-    state.group,
-    state.numChannelsInPerGroup,
-    state.numChannelsOutPerGroup,
-    state.numChannelsInPerGroup * state.wHeight * state.wWidth,
-    state.numChannelsOutPerGroup,
-    state.batchSize * state.outHeight * state.outWidth,
-    static_cast<int64_t>(crossbarSize.getValue()),
+    problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth,
+    problem.numChannelsOutPerGroup,
+    problem.batchSize * problem.outHeight * problem.outWidth,
+    static_cast<int64_t>(target.matrixShape.rows),
+    static_cast<int64_t>(target.matrixUnitsPerProcessor),
     1,
     0,
-    state.hasBias,
-    isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.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));
   return geo;
 }
 
-uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor) {
+static ConvMaterializationKind getMaterializationKind(
+    spatial::ConvLoweringStrategy strategy) {
+  switch (strategy) {
+  case spatial::ConvLoweringStrategy::Depthwise:
+    return ConvMaterializationKind::StructuredDepthwise;
+  case spatial::ConvLoweringStrategy::Legacy:
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
+    return ConvMaterializationKind::PackedIm2Col;
+  case spatial::ConvLoweringStrategy::StreamedPatch:
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+  case spatial::ConvLoweringStrategy::Tiled2D:
+    return ConvMaterializationKind::StreamedPatch;
+  case spatial::ConvLoweringStrategy::StreamedPacked:
+    return ConvMaterializationKind::StreamedPacked;
+  case spatial::ConvLoweringStrategy::InputKTiled:
+    return ConvMaterializationKind::InputKTiled;
+  case spatial::ConvLoweringStrategy::Auto:
+    break;
+  }
+  llvm_unreachable("auto is not a Conv materialization kind");
+}
+
+static bool fitsSingleCrossbar(const ConvGeometry& geo) {
+  return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize;
+}
+
+static bool fitsPackedIm2Col(const ConvGeometry& geo,
+                             const ONNXToSpatialPlanningOptions& options) {
+  return fitsSingleCrossbar(geo) && geo.pack >= 2
+    && geo.im2colElements <= options.convIm2colMaxElements;
+}
+
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  auto plan = [&]() { return ConvPlan {getMaterializationKind(strategy)}; };
+  auto ifApplicable = [&](bool applicable) -> mlir::FailureOr<ConvPlan> {
+    return applicable ? mlir::FailureOr<ConvPlan>(plan()) : mlir::FailureOr<ConvPlan>(mlir::failure());
+  };
+  switch (strategy) {
+  case spatial::ConvLoweringStrategy::Auto:
+    return mlir::failure();
+  case spatial::ConvLoweringStrategy::Legacy:
+    return plan();
+  case spatial::ConvLoweringStrategy::Depthwise:
+    return ifApplicable(problem.isDepthwise);
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
+    return ifApplicable(fitsPackedIm2Col(geo, options));
+  case spatial::ConvLoweringStrategy::StreamedPatch:
+    return ifApplicable(fitsSingleCrossbar(geo));
+  case spatial::ConvLoweringStrategy::StreamedPacked:
+    return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2);
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+    return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize);
+  case spatial::ConvLoweringStrategy::InputKTiled:
+    return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize);
+  case spatial::ConvLoweringStrategy::Tiled2D:
+    return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize);
+  }
+  llvm_unreachable("unknown Conv lowering strategy");
+}
+
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
+                                    int64_t packFactor,
+                                    const ONNXToSpatialPlanningOptions& options) {
   const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
-  uint64_t chunkPositions = std::max<uint64_t>(1, pimConvIm2colMaxElements / patchElements);
+  uint64_t chunkPositions = std::max<uint64_t>(1, options.convIm2colMaxElements / patchElements);
   chunkPositions = std::min<uint64_t>(chunkPositions, static_cast<uint64_t>(std::max<int64_t>(1, geo.p)));
-  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, pimConvStreamChunkPositions));
+  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, options.convStreamChunkPositions));
 
   if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
     chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
@@ -52,24 +136,26 @@ uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor)
   return std::max<uint64_t>(1, chunkPositions);
 }
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state) {
-  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem) {
+  const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
   const int64_t rawEnd =
-    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
-  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(state.xHeight, rawEnd)};
+    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
+    + problem.dilationHeight * (problem.wHeight - 1) + 1;
+  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(problem.xHeight, rawEnd)};
 }
 
-ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state) {
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem) {
   ConvRowDemand demand;
   demand.outputRows = outputRows;
-  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, state);
+  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, problem);
   demand.acquiredInputRows = demand.neededInputRows;
 
-  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+  const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
   const int64_t rawEnd =
-    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
+    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
+    + problem.dilationHeight * (problem.wHeight - 1) + 1;
   demand.topHaloRows = std::max<int64_t>(0, -rawBegin);
-  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - state.xHeight);
+  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - problem.xHeight);
   demand.acquiredInputRows = demand.neededInputRows;
   return demand;
 }

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

@@ -3,14 +3,19 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Value.h"
 
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
 #include <cstdint>
 
+namespace mlir {
+class Operation;
+} // namespace mlir
+
 namespace onnx_mlir {
 
-struct ConvLoweringState {
-  mlir::Value x;
-  mlir::Value w;
-  mlir::Value b;
+struct ConvProblem {
   mlir::RankedTensorType xType;
   mlir::RankedTensorType wType;
   mlir::RankedTensorType outType;
@@ -35,29 +40,32 @@ struct ConvLoweringState {
   int64_t dilationHeight;
   int64_t dilationWidth;
   bool hasBias;
+  bool isDepthwise = false;
+  bool isGrouped = false;
+  bool isPointwise = false;
+};
+
+struct ConvLoweringState {
+  ConvProblem problem;
+  mlir::Operation* diagnosticAnchor = nullptr;
+  mlir::Value x;
+  mlir::Value w;
+  mlir::Value b;
+  const spatial::SpatialTargetResources* target = nullptr;
+  const ONNXToSpatialPlanningOptions* options = nullptr;
+
+  const spatial::SpatialTargetResources& targetInfo() const { return *target; }
+  const ONNXToSpatialPlanningOptions& planningOptions() const;
 };
 
 struct ConvGeometry {
-  int64_t batchSize;
-  int64_t numChannelsIn;
-  int64_t xHeight;
-  int64_t xWidth;
-  int64_t numChannelsOut;
-  int64_t wHeight;
-  int64_t wWidth;
-  int64_t outHeight;
-  int64_t outWidth;
-  int64_t group;
-  int64_t numChannelsInPerGroup;
-  int64_t numChannelsOutPerGroup;
   int64_t k;
   int64_t c;
   int64_t p;
   int64_t xbarSize;
+  int64_t matrixUnitsPerProcessor;
   int64_t pack;
   uint64_t im2colElements;
-  bool hasBias;
-  bool isDepthwise;
 };
 
 struct RowInterval {
@@ -73,14 +81,36 @@ struct ConvRowDemand {
   int64_t bottomHaloRows = 0;
 };
 
+enum class ConvMaterializationKind : uint8_t {
+  StructuredDepthwise,
+  PackedIm2Col,
+  StreamedPatch,
+  StreamedPacked,
+  InputKTiled,
+};
+
+struct ConvPlan {
+  ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col;
+};
+
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
 
-ConvGeometry buildConvGeometry(const ConvLoweringState& state);
+void classifyConvProblem(ConvProblem& problem);
 
-uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor);
+ConvGeometry buildConvGeometry(const ConvProblem& problem,
+                               const spatial::SpatialTargetResources& target);
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state);
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options);
 
-ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state);
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
+                                    int64_t packFactor,
+                                    const ONNXToSpatialPlanningOptions& options);
+
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem);
+
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem);
 
 } // namespace onnx_mlir

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

@@ -31,7 +31,7 @@ struct SiluToSpatialPlan : OpRewritePattern<ONNXMulOp> {
       return failure();
 
     auto plan = spatial::SpatSiluPlanOp::create(
-      rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, rewriter.getStringAttr("nchw"));
+      rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, spatial::getNCHWLayout(rewriter.getContext()));
     rewriter.replaceOp(mulOp, plan.getResult());
     rewriter.eraseOp(sigmoidOp);
     return success();
@@ -48,6 +48,56 @@ static DenseElementsAttr getDenseConstantAttr(Value value) {
   return nullptr;
 }
 
+struct BlueprintSplatMulToSpatial : OpConversionPattern<ONNXMulOp> {
+  explicit BlueprintSplatMulToSpatial(MLIRContext* ctx) : OpConversionPattern(ctx, 2) {}
+
+  LogicalResult
+  matchAndRewrite(ONNXMulOp op, ONNXMulOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto blueprint = op.getA().getDefiningOp<spatial::SpatBlueprintOp>();
+    Value scalar = adaptor.getB();
+    if (!blueprint) {
+      blueprint = op.getB().getDefiningOp<spatial::SpatBlueprintOp>();
+      scalar = adaptor.getA();
+    }
+    auto scalarAttr = getDenseConstantAttr(scalar);
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    auto storageType = blueprint ? dyn_cast<RankedTensorType>(blueprint.getInput().getType()) : RankedTensorType();
+    if (!blueprint || !blueprint.getFragments().empty() || !scalarAttr || !scalarAttr.isSplat() || !resultType
+        || resultType != blueprint.getOutput().getType() || !storageType)
+      return failure();
+
+    auto mapped = mapGraphBatchFragments(
+      blueprint.getInput(), storageType, rewriter, op.getLoc(), [&](Value fragment, RankedTensorType fragmentType) {
+        auto splat = DenseElementsAttr::get(fragmentType, scalarAttr.getSplatValue<Attribute>());
+        Value constant = arith::ConstantOp::create(rewriter, op.getLoc(), fragmentType, splat);
+        return FailureOr<Value>(
+          spatial::SpatVMulOp::create(rewriter, op.getLoc(), fragmentType, fragment, constant).getResult());
+      });
+    if (failed(mapped))
+      return failure();
+
+    auto result = spatial::SpatBlueprintOp::create(rewriter,
+                                                   op.getLoc(),
+                                                   resultType,
+                                                   *mapped,
+                                                   ValueRange {},
+                                                   blueprint.getLogicalLayoutAttr(),
+                                                   blueprint.getPhysicalLayoutAttr(),
+                                                   blueprint.getFragmentOffsetsAttr(),
+                                                   blueprint.getFragmentSizesAttr(),
+                                                   blueprint.getIndexMapAttr(),
+                                                   blueprint.getModeAttr(),
+                                                   blueprint.getFragmentOperandIndicesAttr(),
+                                                   blueprint.getFragmentSourceSlotsAttr(),
+                                                   blueprint.getFragmentSourceOffsetsAttr(),
+                                                   blueprint.getFragmentStridesAttr(),
+                                                   blueprint.getConflictPolicyAttr(),
+                                                   blueprint.getCoveragePolicyAttr());
+    rewriter.replaceOp(op, result.getOutput());
+    return success();
+  }
+};
+
 static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
                                                              RankedTensorType resultType,
                                                              ConversionPatternRewriter& rewriter,
@@ -210,14 +260,16 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
       classifyBiasAddPlanCandidate(adaptor.getA(), adaptor.getB(), resultType);
     if (succeeded(candidate)) {
       auto plan = spatial::SpatBiasAddPlanOp::create(
-        rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, rewriter.getStringAttr("nchw"));
+        rewriter, op.getLoc(), resultType, candidate->data, candidate->bias,
+        spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(op, plan.getResult());
       return success();
     }
 
     if (resultType.getRank() == 4 && adaptor.getA().getType() == resultType && adaptor.getB().getType() == resultType) {
       auto plan = spatial::SpatAddPlanOp::create(
-        rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(), rewriter.getStringAttr("nchw"));
+        rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(),
+        spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(op, plan.getResult());
       return success();
     }
@@ -246,6 +298,7 @@ void populateElementwiseFusionPatterns(RewritePatternSet& patterns, MLIRContext*
 }
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<BlueprintSplatMulToSpatial>(ctx);
   patterns.add<AddToSpatialCompute>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);

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

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -21,6 +22,9 @@
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #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/ContractionPlanning.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -31,7 +35,7 @@ namespace onnx_mlir {
 namespace {
 
 static FailureOr<Value>
-materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewriter& rewriter, Location loc) {
+materializeScaledConstantTensor(Value value, float factor, PatternRewriter& rewriter, Location loc) {
   if (factor == 1.0f)
     return value;
 
@@ -57,7 +61,12 @@ materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewr
 }
 
 static Value createGemmBatchKOffset(
-  Value lane, int64_t numOutRows, int64_t numKSlices, ConversionPatternRewriter& rewriter, Location loc) {
+  Value lane,
+  int64_t numOutRows,
+  int64_t numKSlices,
+  int64_t xbarSize,
+  PatternRewriter& rewriter,
+  Location loc) {
   if (numKSlices == 1)
     return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
 
@@ -65,7 +74,7 @@ static Value createGemmBatchKOffset(
   AffineExpr d0 = getAffineDimExpr(0, context);
   return createOrFoldAffineApply(rewriter,
                                  loc,
-                                 (d0.floorDiv(numOutRows) % numKSlices) * crossbarSize.getValue(),
+                                 (d0.floorDiv(numOutRows) % numKSlices) * xbarSize,
                                  ValueRange {lane},
                                  rewriter.getInsertionBlock()->getParentOp());
 }
@@ -74,7 +83,8 @@ static Value createGemmBatchHOffset(Value lane,
                                     int64_t numOutRows,
                                     int64_t numKSlices,
                                     int64_t numOutHSlices,
-                                    ConversionPatternRewriter& rewriter,
+                                    int64_t xbarSize,
+                                    PatternRewriter& rewriter,
                                     Location loc) {
   if (numOutHSlices == 1)
     return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
@@ -83,14 +93,14 @@ static Value createGemmBatchHOffset(Value lane,
   AffineExpr d0 = getAffineDimExpr(0, context);
   return createOrFoldAffineApply(rewriter,
                                  loc,
-                                 d0.floorDiv(numOutRows * numKSlices) * crossbarSize.getValue(),
+                                 d0.floorDiv(numOutRows * numKSlices) * xbarSize,
                                  ValueRange {lane},
                                  rewriter.getInsertionBlock()->getParentOp());
 }
 
 static FailureOr<Value> materializePaddedConstantMatrix(Value value,
                                                         RankedTensorType resultType,
-                                                        ConversionPatternRewriter& rewriter,
+                                                        PatternRewriter& rewriter,
                                                         Location loc) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   if (sourceType == resultType)
@@ -121,7 +131,7 @@ static FailureOr<Value> materializePaddedConstantMatrix(Value value,
 static FailureOr<Value> materializePaddedBroadcastedConstantTensor(Value value,
                                                                    RankedTensorType resultType,
                                                                    int64_t unpaddedColumns,
-                                                                   ConversionPatternRewriter& rewriter,
+                                                                   PatternRewriter& rewriter,
                                                                    Location loc) {
   auto denseAttr = getHostConstDenseElementsAttr(value);
   if (!denseAttr)
@@ -187,7 +197,7 @@ static FailureOr<Value> materializePaddedBroadcastedConstantTensor(Value value,
 static FailureOr<Value> prepareBias(Value c,
                                     RankedTensorType outType,
                                     RankedTensorType paddedOutType,
-                                    ConversionPatternRewriter& rewriter,
+                                    PatternRewriter& rewriter,
                                     Location loc) {
   auto cType = cast<RankedTensorType>(c.getType());
   if (!cType.hasStaticShape())
@@ -203,9 +213,15 @@ static FailureOr<Value> prepareBias(Value c,
 }
 
 static Value extractATile(
-  Value a, Value row, Value kOffset, RankedTensorType aTileType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value a,
+  Value row,
+  Value kOffset,
+  RankedTensorType aTileType,
+  int64_t xbarSize,
+  PatternRewriter& rewriter,
+  Location loc) {
   SmallVector<OpFoldResult> offsets {row, kOffset};
-  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
 
   return tensor::ExtractSliceOp::create(rewriter, loc, aTileType, a, offsets, sizes, strides).getResult();
@@ -219,7 +235,8 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
                                                            int64_t numOutRows,
                                                            int64_t numKSlices,
                                                            int64_t numOutHSlices,
-                                                           ConversionPatternRewriter& rewriter,
+                                                           int64_t xbarSize,
+                                                           PatternRewriter& rewriter,
                                                            Location loc) {
   const int64_t laneCount = partialPiecesType.getDimSize(0);
   auto batchOp = createSpatComputeBatch(
@@ -232,21 +249,21 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
     [&](detail::SpatComputeBatchBodyArgs args) {
       Value row =
         onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutRows, rewriter.getInsertionBlock()->getParentOp());
-      Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, rewriter, loc);
-      Value hOffset = createGemmBatchHOffset(args.lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+      Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, xbarSize, rewriter, loc);
+      Value hOffset = createGemmBatchHOffset(
+        args.lane, numOutRows, numKSlices, numOutHSlices, xbarSize, rewriter, loc);
 
       auto aTileType =
-        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, aType.getElementType());
+        RankedTensorType::get({1, xbarSize}, aType.getElementType());
       auto bTileType = RankedTensorType::get(
-        {static_cast<int64_t>(crossbarSize.getValue()), static_cast<int64_t>(crossbarSize.getValue())},
+        {xbarSize, xbarSize},
         paddedBType.getElementType());
       auto pieceType =
-        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());
-      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, rewriter, loc);
+        RankedTensorType::get({1, xbarSize}, partialPiecesType.getElementType());
+      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, xbarSize, rewriter, loc);
 
       SmallVector<OpFoldResult> bOffsets {kOffset, hOffset};
-      SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(crossbarSize.getValue()),
-                                        rewriter.getIndexAttr(crossbarSize.getValue())};
+      SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarSize), rewriter.getIndexAttr(xbarSize)};
       SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
       Value bTile = extractStaticSliceOrIdentity(
         rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides);
@@ -260,7 +277,7 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
 }
 
 static Value extractDynamicGemmBColumn(
-  Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value matrix, Value column, RankedTensorType vectorType, PatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), column};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -280,7 +297,7 @@ static Value extractDynamicGemmBColumn(
 }
 
 static Value extractDynamicGemmRowVector(
-  Value matrix, Value row, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value matrix, Value row, RankedTensorType vectorType, PatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {row, rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -311,13 +328,15 @@ 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);
 }
 
 static Value createScalarTensorConstant(RankedTensorType scalarType,
                                         float value,
-                                        ConversionPatternRewriter& rewriter,
+                                        PatternRewriter& rewriter,
                                         Location loc) {
   auto elementType = scalarType.getElementType();
   auto scalarAttr = rewriter.getFloatAttr(elementType, value);
@@ -330,7 +349,7 @@ static Value createBroadcastedBiasScalar(Value bias,
                                          Value row,
                                          Value column,
                                          RankedTensorType scalarType,
-                                         ConversionPatternRewriter& rewriter,
+                                         PatternRewriter& rewriter,
                                          Location loc) {
   SmallVector<OpFoldResult> unitStrides(biasType.getRank(), rewriter.getIndexAttr(1));
   if (biasType.getRank() == 1) {
@@ -365,7 +384,7 @@ static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
                                                               RankedTensorType columnPiecesType,
                                                               RankedTensorType outType,
                                                               bool transposeB,
-                                                              ConversionPatternRewriter& rewriter,
+                                                              PatternRewriter& rewriter,
                                                               Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
@@ -425,7 +444,7 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
                                                                       RankedTensorType outType,
                                                                       float alpha,
                                                                       float beta,
-                                                                      ConversionPatternRewriter& rewriter,
+                                                                      PatternRewriter& rewriter,
                                                                       Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
@@ -510,7 +529,7 @@ static Value createPartialGroupOffset(Value hSlice,
                                       int64_t kSlice,
                                       int64_t numKSlices,
                                       int64_t numOutRows,
-                                      ConversionPatternRewriter& rewriter,
+                                      PatternRewriter& rewriter,
                                       Location loc) {
   MLIRContext* context = rewriter.getContext();
   AffineExpr d0 = getAffineDimExpr(0, context);
@@ -527,10 +546,12 @@ static Value extractReductionPiece(Value partialPiecesArg,
                                    RankedTensorType pieceType,
                                    int64_t numKSlices,
                                    int64_t numOutRows,
-                                   ConversionPatternRewriter& rewriter,
+                                   int64_t xbarSize,
+                                   PatternRewriter& rewriter,
                                    Location loc) {
   SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> pieceSizes {
+    rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)};
   SmallVector<OpFoldResult> pieceOffsets {
     createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc),
     rewriter.getIndexAttr(0),
@@ -545,13 +566,15 @@ static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
                                           RankedTensorType pieceType,
                                           int64_t numKSlices,
                                           int64_t numOutRows,
-                                          ConversionPatternRewriter& rewriter,
+                                          int64_t xbarSize,
+                                          PatternRewriter& rewriter,
                                           Location loc) {
   SmallVector<Value> activePieces;
   activePieces.reserve(numKSlices);
   for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice)
     activePieces.push_back(
-      extractReductionPiece(partialPiecesArg, hSlice, kSlice, pieceType, numKSlices, numOutRows, rewriter, loc));
+      extractReductionPiece(
+        partialPiecesArg, hSlice, kSlice, pieceType, numKSlices, numOutRows, xbarSize, rewriter, loc));
 
   while (activePieces.size() > 1) {
     SmallVector<Value> nextPieces;
@@ -574,11 +597,12 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
                                               RankedTensorType outType,
                                               RankedTensorType paddedOutType,
                                               int64_t numKSlices,
-                                              ConversionPatternRewriter& rewriter,
+                                              int64_t xbarSize,
+                                              PatternRewriter& rewriter,
                                               Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
-  const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue());
-  auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
+  const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), xbarSize);
+  auto pieceType = RankedTensorType::get({numOutRows, xbarSize},
                                          partialPiecesType.getElementType());
 
   if (bias && cast<RankedTensorType>(bias.getType()) != paddedOutType)
@@ -590,20 +614,20 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
   SmallVector<Value> outputSlices;
   outputSlices.reserve(numOutHSlices);
   for (int64_t hSlice = 0; hSlice < numOutHSlices; ++hSlice) {
-    const int64_t columnOffset = hSlice * crossbarSize.getValue();
+    const int64_t columnOffset = hSlice * xbarSize;
     const int64_t columns =
-      std::min(static_cast<int64_t>(crossbarSize.getValue()), outType.getDimSize(1) - columnOffset);
+      std::min(xbarSize, outType.getDimSize(1) - columnOffset);
     auto outputSliceType = RankedTensorType::get({numOutRows, columns}, outType.getElementType());
     auto computeOp = createSpatCompute(
       rewriter, loc, TypeRange {outputSliceType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult {
       Value hSliceValue =
         getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), hSlice);
       Value reduced = reducePartialPiecesForHSlice(
-        blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, rewriter, loc);
+        blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, xbarSize, rewriter, loc);
       if (bias) {
         SmallVector<OpFoldResult> biasOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(columnOffset)};
         SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
-                                              rewriter.getIndexAttr(crossbarSize.getValue())};
+                                              rewriter.getIndexAttr(xbarSize)};
         SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
         Value biasSlice =
           tensor::ExtractSliceOp::create(rewriter, loc, pieceType, blockArgs[1], biasOffsets, pieceSizes, unitStrides)
@@ -637,79 +661,96 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
 }
 
 struct GemmToSpatialComputes : OpConversionPattern<ONNXGemmOp> {
-  using OpConversionPattern::OpConversionPattern;
+  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<ONNXGemmOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
                                 ONNXGemmOpAdaptor gemmOpAdaptor,
                                 ConversionPatternRewriter& rewriter) const override;
+
+  const spatial::SpatialTargetResources& target;
 };
 
 } // namespace
 
-LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
-                                                     ONNXGemmOpAdaptor gemmOpAdaptor,
-                                                     ConversionPatternRewriter& rewriter) const {
-  Location loc = gemmOp.getLoc();
-  Value a = gemmOpAdaptor.getA();
-  Value b = gemmOpAdaptor.getB();
-  Value c = gemmOpAdaptor.getC();
-
+FailureOr<Value> lowerGemmToSpatial(
+    Operation* diagnosticAnchor,
+    Value a,
+    Value b,
+    Value c,
+    RankedTensorType outType,
+    bool transA,
+    bool transB,
+    float alpha,
+    float beta,
+    const spatial::SpatialTargetResources& target,
+    PatternRewriter& rewriter,
+    Location loc) {
   auto aType = dyn_cast<RankedTensorType>(a.getType());
   auto bType = dyn_cast<RankedTensorType>(b.getType());
-  auto outType = dyn_cast<RankedTensorType>(gemmOp.getY().getType());
-  if (!aType || !bType || !outType)
+  if (!diagnosticAnchor || !aType || !bType || !outType)
     return failure();
   if (!aType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input A");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm input A");
     return failure();
   }
   if (!bType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input B");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm input B");
     return failure();
   }
   if (!outType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm result");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm result");
     return failure();
   }
   if (aType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm input A", aType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm input A", aType.getRank(), {2});
     return failure();
   }
   if (bType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm input B", bType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm input B", bType.getRank(), {2});
     return failure();
   }
   if (outType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm result", outType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm result", outType.getRank(), {2});
     return failure();
   }
 
-  if (gemmOpAdaptor.getTransA()) {
+  if (transA) {
     auto aShape = aType.getShape();
     auto transposedType = RankedTensorType::get({aShape[1], aShape[0]}, aType.getElementType(), aType.getEncoding());
-    a = ONNXTransposeOp::create(rewriter, loc, transposedType, a, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    a = createLinalgTranspose(a, transposedType, {1, 0}, rewriter, loc);
     aType = transposedType;
   }
 
-  const int64_t numOutRows = outType.getDimSize(0);
-  const int64_t numOutCols = outType.getDimSize(1);
-  const int64_t reductionSize = aType.getDimSize(1);
-  const bool transposeB = gemmOpAdaptor.getTransB();
+  ContractionProblem problem;
+  problem.lhsBatchShape = {};
+  problem.rhsBatchShape = {};
+  problem.outputBatchShape = {};
+  problem.lhsBatch = 1;
+  problem.rhsBatch = 1;
+  problem.batch = 1;
+  problem.m = outType.getDimSize(0);
+  problem.k = aType.getDimSize(1);
+  problem.n = outType.getDimSize(1);
+  problem.lhsElementType = aType.getElementType();
+  problem.rhsElementType = bType.getElementType();
+  problem.resultElementType = outType.getElementType();
+  const bool transposeB = transB;
 
   if (!isCompileTimeComputable(b)) {
+    ContractionPlan plan = makeContractionPlan(
+        problem, target, ContractionPlanKind::BatchedDynamicVVD);
     bool hasC = hasGemmBias(c);
-    float alpha = gemmOpAdaptor.getAlpha().convertToFloat();
-    float beta = gemmOpAdaptor.getBeta().convertToFloat();
     RankedTensorType biasType;
     if (hasC) {
       auto cType = dyn_cast<RankedTensorType>(c.getType());
       if (!cType || !cType.hasStaticShape()) {
-        pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+        pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm bias");
         return failure();
       }
       auto verifiedBiasType = verifyDynamicGemmBiasType(cType, outType);
       if (failed(verifiedBiasType)) {
-        gemmOp.emitOpError("requires Gemm bias C to be broadcastable to the output shape");
+        diagnosticAnchor->emitOpError("requires Gemm bias C to be broadcastable to the output shape");
         return failure();
       }
       biasType = *verifiedBiasType;
@@ -717,19 +758,19 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
 
     const int64_t bReductionSize = bType.getDimSize(transposeB ? 1 : 0);
     const int64_t bOutputColumns = bType.getDimSize(transposeB ? 0 : 1);
-    if (aType.getDimSize(0) != numOutRows || bReductionSize != reductionSize || bOutputColumns != numOutCols) {
-      gemmOp.emitOpError("has inconsistent A, B, and output shapes");
+    if (aType.getDimSize(0) != problem.m || bReductionSize != problem.k || bOutputColumns != problem.n) {
+      diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes");
       return failure();
     }
 
-    const int64_t laneCount64 = numOutRows * numOutCols;
+    const int64_t laneCount64 = plan.laneCount;
     if (laneCount64 > std::numeric_limits<int32_t>::max()) {
-      gemmOp.emitOpError("requires Gemm dynamic batch lane count to fit in i32");
+      diagnosticAnchor->emitOpError("requires Gemm dynamic batch lane count to fit in i32");
       return failure();
     }
 
-    auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
-    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(numOutCols, columnType);
+    auto columnType = RankedTensorType::get({problem.m, 1}, outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(problem.n, columnType);
     auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, transposeB, rewriter, loc);
     if (failed(batchOp))
       return failure();
@@ -737,94 +778,128 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
       batchOp->getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc);
     if (failed(outputCompute))
       return failure();
-    rewriter.replaceOp(gemmOp, outputCompute->getResults());
-    return success();
+    return outputCompute->getResult(0);
   }
 
   if (transposeB) {
     auto bShape = bType.getShape();
     auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
-    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    if (isCompileTimeComputable(b)) {
+      auto denseAttr = getHostConstDenseElementsAttr(b);
+      auto inputType = denseAttr ? dyn_cast<RankedTensorType>(denseAttr.getType()) : nullptr;
+      auto transposedAttr = inputType && inputType.hasStaticShape() && transposedType.hasStaticShape()
+        ? transposeDenseElementsAttr(denseAttr, {1, 0})
+        : FailureOr<DenseElementsAttr>(failure());
+      if (failed(transposedAttr) || transposedAttr->getType() != transposedType) {
+        diagnosticAnchor->emitOpError("requires Gemm input B transpose to remain statically materializable");
+        return failure();
+      }
+      b = getOrCreateConstant(rewriter,
+                              rewriter.getInsertionBlock()->getParentOp(),
+                              *transposedAttr,
+                              transposedType);
+    } else {
+      b = createLinalgTranspose(b, transposedType, {1, 0}, rewriter, loc);
+    }
     bType = transposedType;
   }
 
-  auto scaledB = materializeScaledConstantTensor(b, gemmOpAdaptor.getAlpha().convertToFloat(), rewriter, loc);
+  auto scaledB = materializeScaledConstantTensor(b, alpha, rewriter, loc);
   if (failed(scaledB)) {
-    gemmOp.emitOpError("requires constant Gemm input B when alpha is not 1.0");
+    diagnosticAnchor->emitOpError("requires constant Gemm input B when alpha is not 1.0");
     return failure();
   }
   b = *scaledB;
   bType = cast<RankedTensorType>(b.getType());
 
-  if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize || bType.getDimSize(1) != numOutCols) {
-    gemmOp.emitOpError("has inconsistent A, B, and output shapes after transpose handling");
+  if (aType.getDimSize(0) != problem.m || bType.getDimSize(0) != problem.k || bType.getDimSize(1) != problem.n) {
+    diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes after transpose handling");
     return failure();
   }
 
-  const int64_t numKSlices = ceilIntegerDivide(reductionSize, crossbarSize.getValue());
-  const int64_t numOutHSlices = ceilIntegerDivide(numOutCols, crossbarSize.getValue());
-  const int64_t paddedReductionSize = numKSlices * static_cast<int64_t>(crossbarSize.getValue());
-  const int64_t paddedOutCols = numOutHSlices * static_cast<int64_t>(crossbarSize.getValue());
+  ContractionPlan plan = makeContractionPlan(
+      problem, target, ContractionPlanKind::StaticTiled);
+  const int64_t xbarSize = plan.tileK;
+  const int64_t numKSlices = plan.reductionSlices;
+  const int64_t numOutHSlices = plan.outputTiles;
+  const int64_t paddedReductionSize = numKSlices * plan.tileK;
+  const int64_t paddedOutCols = numOutHSlices * plan.tileN;
 
   auto paddedBType = RankedTensorType::get({paddedReductionSize, paddedOutCols}, bType.getElementType());
   auto paddedB = materializePaddedConstantMatrix(b, paddedBType, rewriter, loc);
   if (failed(paddedB)) {
-    gemmOp.emitOpError("requires constant Gemm input B so tiled weights can be padded statically");
+    diagnosticAnchor->emitOpError("requires constant Gemm input B so tiled weights can be padded statically");
     return failure();
   }
   b = *paddedB;
-  auto paddedAType = RankedTensorType::get({numOutRows, paddedReductionSize}, aType.getElementType());
+  auto paddedAType = RankedTensorType::get({problem.m, paddedReductionSize}, aType.getElementType());
   a = createPaddedInputCompute(a, paddedAType, rewriter, loc);
   aType = paddedAType;
 
   Value bias;
   bool hasC = hasGemmBias(c);
-  auto paddedOutType = RankedTensorType::get({numOutRows, paddedOutCols}, outType.getElementType());
+  auto paddedOutType = RankedTensorType::get({problem.m, paddedOutCols}, outType.getElementType());
   if (hasC) {
     auto cType = dyn_cast<RankedTensorType>(c.getType());
     if (!cType || !cType.hasStaticShape()) {
-      pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+      pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm bias");
       return failure();
     }
 
-    auto scaledC = materializeScaledConstantTensor(c, gemmOpAdaptor.getBeta().convertToFloat(), rewriter, loc);
+    auto scaledC = materializeScaledConstantTensor(c, beta, rewriter, loc);
     if (failed(scaledC)) {
-      gemmOp.emitOpError("requires constant Gemm bias C when beta is not 1.0");
+      diagnosticAnchor->emitOpError("requires constant Gemm bias C when beta is not 1.0");
       return failure();
     }
     c = *scaledC;
 
     auto preparedBias = prepareBias(c, outType, paddedOutType, rewriter, loc);
     if (failed(preparedBias)) {
-      gemmOp.emitOpError("requires Gemm bias C to be broadcastable to the output shape");
+      diagnosticAnchor->emitOpError("requires Gemm bias C to be broadcastable to the output shape");
       return failure();
     }
     bias = *preparedBias;
   }
 
-  const int64_t laneCount64 = numOutHSlices * numKSlices * numOutRows;
+  const int64_t laneCount64 = plan.laneCount;
   if (laneCount64 > std::numeric_limits<int32_t>::max()) {
-    gemmOp.emitOpError("requires Gemm tiled batch lane count to fit in i32");
+    diagnosticAnchor->emitOpError("requires Gemm tiled batch lane count to fit in i32");
     return failure();
   }
 
   auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
-    laneCount64, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType()));
+    laneCount64, RankedTensorType::get({1, xbarSize}, outType.getElementType()));
   auto batchOp =
-    createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+    createVmmBatch(
+      a, b, aType, paddedBType, partialPiecesType, problem.m, numKSlices, numOutHSlices, xbarSize, rewriter, loc);
   if (failed(batchOp))
     return failure();
   auto reductionOutput = createReductionOutput(
-    batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc);
+    batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, xbarSize, rewriter, loc);
   if (failed(reductionOutput))
     return failure();
 
-  rewriter.replaceOp(gemmOp, *reductionOutput);
+  return *reductionOutput;
+}
+
+LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
+                                                     ONNXGemmOpAdaptor gemmOpAdaptor,
+                                                     ConversionPatternRewriter& rewriter) const {
+  FailureOr<Value> result = lowerGemmToSpatial(
+      gemmOp.getOperation(), gemmOpAdaptor.getA(), gemmOpAdaptor.getB(), gemmOpAdaptor.getC(),
+      cast<RankedTensorType>(gemmOp.getY().getType()), gemmOpAdaptor.getTransA(),
+      gemmOpAdaptor.getTransB(), gemmOpAdaptor.getAlpha().convertToFloat(),
+      gemmOpAdaptor.getBeta().convertToFloat(), target, rewriter, gemmOp.getLoc());
+  if (failed(result))
+    return failure();
+  rewriter.replaceOp(gemmOp, *result);
   return success();
 }
 
-void populateGemmPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.insert<GemmToSpatialComputes>(ctx);
+void populateGemmPatterns(RewritePatternSet& patterns,
+                          MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target) {
+  patterns.insert<GemmToSpatialComputes>(ctx, target);
 }
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,27 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+namespace spatial {
+struct SpatialTargetResources;
+}
+
+mlir::FailureOr<mlir::Value> lowerGemmToSpatial(
+    mlir::Operation* diagnosticAnchor,
+    mlir::Value a,
+    mlir::Value b,
+    mlir::Value c,
+    mlir::RankedTensorType outputType,
+    bool transA,
+    bool transB,
+    float alpha,
+    float beta,
+    const spatial::SpatialTargetResources& target,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -280,12 +280,12 @@ static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
   SmallVector<int64_t> fragmentStrides(fragmentOffsets.size(), 1);
   return spatial::SpatBlueprintOp::create(
            rewriter, loc, keepdimsType, batchValue, ValueRange {},
-           rewriter.getStringAttr("nchw"),
-           rewriter.getStringAttr("fragmented"),
+           spatial::getNCHWLayout(rewriter.getContext()),
+           spatial::getFragmentedLayout(rewriter.getContext()),
            rewriter.getDenseI64ArrayAttr(fragmentOffsets),
            rewriter.getDenseI64ArrayAttr(fragmentSizes),
            rewriter.getStringAttr("reduce_mean_keepdims_fragments"),
-           rewriter.getStringAttr("fragment_assembly"),
+           spatial::getFragmentAssemblyMode(rewriter.getContext()),
            rewriter.getDenseI64ArrayAttr(operandIndices),
            rewriter.getDenseI64ArrayAttr(sourceSlots),
            rewriter.getDenseI64ArrayAttr(sourceOffsets),

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

@@ -14,10 +14,10 @@
 
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.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/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -32,8 +32,10 @@ static Value materializeTileTensor(PatternRewriter& rewriter, Location loc, Valu
   return insertStaticSlice(rewriter, loc, tile, empty, getZeroOffsets(rewriter, tileType.getRank()));
 }
 
-static Value
-createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type elementType, bool useMinimumValue) {
+static Value createPoolFillElement(OpBuilder& rewriter,
+                                   Location loc,
+                                   Type elementType,
+                                   bool useMinimumValue) {
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
   if (!useMinimumValue)
     return getOrCreateConstant(rewriter, anchorOp, rewriter.getZeroAttr(elementType), elementType);
@@ -51,7 +53,7 @@ createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type el
   llvm_unreachable("unsupported pool element type");
 }
 
-static Value createPoolFillTensor(ConversionPatternRewriter& rewriter,
+static Value createPoolFillTensor(OpBuilder& rewriter,
                                   Location loc,
                                   RankedTensorType tensorType,
                                   bool useMinimumValue) {
@@ -59,16 +61,15 @@ static Value createPoolFillTensor(ConversionPatternRewriter& rewriter,
   return tensor::SplatOp::create(rewriter, loc, tensorType, fillElement);
 }
 
-template <typename PoolOp>
-static Value createPaddedPoolInput(ConversionPatternRewriter& rewriter,
+static Value createPaddedPoolInput(OpBuilder& rewriter,
                                    Location loc,
-                                   PoolOp poolOp,
                                    Value input,
                                    RankedTensorType inputType,
                                    int64_t padTop,
                                    int64_t padLeft,
                                    int64_t padBottom,
-                                   int64_t padRight) {
+                                   int64_t padRight,
+                                   bool useMinimumValue) {
   if (padTop == 0 && padLeft == 0 && padBottom == 0 && padRight == 0)
     return input;
 
@@ -90,8 +91,8 @@ static Value createPaddedPoolInput(ConversionPatternRewriter& rewriter,
     padBlock->addArgument(rewriter.getIndexType(), loc);
   padOp.getRegion().push_back(padBlock);
   rewriter.setInsertionPointToStart(padBlock);
-  Value padValue =
-    createPoolFillElement(rewriter, loc, inputType.getElementType(), std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
+  Value padValue = createPoolFillElement(
+    rewriter, loc, inputType.getElementType(), useMinimumValue);
   tensor::YieldOp::create(rewriter, loc, padValue);
   rewriter.setInsertionPointAfter(padOp);
   return padOp.getResult();
@@ -160,7 +161,10 @@ struct PoolToSpatialCompute;
 
 template <typename PoolOp, typename PoolOpAdaptor, typename ReduceOp>
 struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
-  using OpConversionPattern<PoolOp>::OpConversionPattern;
+  PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<PoolOp>(ctx), target(target) {}
+
+  const spatial::SpatialTargetResources& target;
 
   LogicalResult matchAndRewrite(PoolOp poolOp, PoolOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const final {
     Location loc = poolOp.getLoc();
@@ -241,7 +245,7 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
           rewriter.getDenseI64ArrayAttr({padTop, padLeft, padBottom, padRight}),
           rewriter.getDenseI64ArrayAttr({strideHeight, strideWidth}),
           rewriter.getDenseI64ArrayAttr({dilationHeight, dilationWidth}),
-          rewriter.getStringAttr("nchw"));
+          spatial::getNCHWLayout(rewriter.getContext()));
         rewriter.replaceOp(poolOp, plan.getResult());
         return success();
       }
@@ -251,12 +255,12 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
              && dilationHeight == 1 && dilationWidth == 1 && padTop == 0
              && padLeft == 0 && padBottom == 0 && padRight == 0) {
       auto plan = spatial::SpatGlobalAveragePoolPlanOp::create(
-        rewriter, loc, outType, x, rewriter.getStringAttr("nchw"));
+        rewriter, loc, outType, x, spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(poolOp, plan.getResult());
       return success();
     }
 
-    const int64_t xbarSize = static_cast<int64_t>(crossbarSize.getValue());
+    const int64_t xbarSize = static_cast<int64_t>(target.matrixShape.rows);
     const int64_t channelTileCount = (channels + xbarSize - 1) / xbarSize;
     const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
     const bool countIncludePad = [&]() {
@@ -292,7 +296,9 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
     auto computeOp =
       createSpatCompute<numInputs>(rewriter, loc, outType, {}, ValueRange {x}, [&](Value xArg) -> LogicalResult {
         Value paddedInput =
-          createPaddedPoolInput(rewriter, loc, poolOp, xArg, xType, padTop, padLeft, padBottom, padRight);
+          createPaddedPoolInput(rewriter, loc, xArg, xType, padTop, padLeft,
+                                padBottom, padRight,
+                                std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
         Value pooledOutputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType());
 
         Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
@@ -424,7 +430,8 @@ struct PoolToSpatialCompute<ONNXAveragePoolOp>
 
 } // namespace
 
-LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp) {
+LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
+                                            const spatial::SpatialTargetResources&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -439,6 +446,119 @@ LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp)
   return success();
 }
 
+FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                                         Value input,
+                                         const spatial::SpatialTargetResources& target,
+                                         PatternRewriter& rewriter) {
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
+      || inputType.getRank() != 4 || outputType.getRank() != 4)
+    return planOp.emitOpError("dense MaxPool lowering requires static rank-4 tensors"), failure();
+
+  auto kernel = planOp.getKernelShape();
+  auto pads = planOp.getPads();
+  auto strides = planOp.getStrides();
+  auto dilations = planOp.getDilations();
+  if (kernel.size() != 2 || pads.size() != 4 || strides.size() != 2 || dilations.size() != 2
+      || llvm::any_of(kernel, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(strides, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(dilations, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(pads, [](int64_t value) { return value < 0; }))
+    return planOp.emitOpError("dense MaxPool lowering requires valid kernel, padding, stride, and dilation attributes"),
+           failure();
+
+  const int64_t batchSize = inputType.getDimSize(0);
+  const int64_t channels = inputType.getDimSize(1);
+  const int64_t outputHeight = outputType.getDimSize(2);
+  const int64_t outputWidth = outputType.getDimSize(3);
+  const int64_t tileWidth = std::max<int64_t>(1, target.matrixShape.rows);
+  const int64_t channelTileCount = (channels + tileWidth - 1) / tileWidth;
+  const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
+
+  auto compute = createSpatCompute<1>(
+    rewriter, planOp.getLoc(), outputType, {}, input,
+    [&](Value input) -> LogicalResult {
+      Value paddedInput = createPaddedPoolInput(
+        rewriter, planOp.getLoc(), input, inputType,
+        pads[0], pads[1], pads[2], pads[3], /*useMinimumValue=*/true);
+      Value outputInit = tensor::EmptyOp::create(
+        rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value zero = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value one = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value patchCount = getOrCreateIndexConstant(rewriter, anchor, outputPatchCount);
+      Value pixelsPerBatch = getOrCreateIndexConstant(
+        rewriter, anchor, outputHeight * outputWidth);
+      Value outputWidthValue = getOrCreateIndexConstant(rewriter, anchor, outputWidth);
+      Value strideHeight = getOrCreateIndexConstant(rewriter, anchor, strides[0]);
+      Value strideWidth = getOrCreateIndexConstant(rewriter, anchor, strides[1]);
+
+      auto loop = buildNormalizedScfFor(
+        rewriter, planOp.getLoc(), zero, patchCount, one, ValueRange {outputInit},
+        [&](OpBuilder&, Location loc, Value patch, ValueRange iterArgs,
+            SmallVectorImpl<Value>& yielded) {
+          Value batch = arith::DivUIOp::create(rewriter, loc, patch, pixelsPerBatch);
+          Value batchPatch = arith::RemUIOp::create(rewriter, loc, patch, pixelsPerBatch);
+          Value outputRow = arith::DivUIOp::create(rewriter, loc, batchPatch, outputWidthValue);
+          Value outputColumn = arith::RemUIOp::create(rewriter, loc, batchPatch, outputWidthValue);
+          Value windowRow = arith::MulIOp::create(rewriter, loc, outputRow, strideHeight);
+          Value windowColumn = arith::MulIOp::create(rewriter, loc, outputColumn, strideWidth);
+          Value updated = iterArgs.front();
+
+          for (int64_t tile = 0; tile < channelTileCount; ++tile) {
+            const int64_t tileChannels = std::min<int64_t>(tileWidth, channels - tile * tileWidth);
+            auto tileType = RankedTensorType::get(
+              {1, tileChannels, 1, 1}, outputType.getElementType());
+            Value reduced = createPoolFillTensor(
+              rewriter, loc, tileType, /*useMinimumValue=*/true);
+            for (int64_t kernelRow = 0; kernelRow < kernel[0]; ++kernelRow) {
+              Value sourceRow = windowRow;
+              if (kernelRow * dilations[0] != 0)
+                sourceRow = arith::AddIOp::create(
+                  rewriter, loc, sourceRow,
+                  getOrCreateIndexConstant(rewriter, anchor, kernelRow * dilations[0]));
+              for (int64_t kernelColumn = 0; kernelColumn < kernel[1]; ++kernelColumn) {
+                Value sourceColumn = windowColumn;
+                if (kernelColumn * dilations[1] != 0)
+                  sourceColumn = arith::AddIOp::create(
+                    rewriter, loc, sourceColumn,
+                    getOrCreateIndexConstant(rewriter, anchor, kernelColumn * dilations[1]));
+                Value point = tensor::ExtractSliceOp::create(
+                  rewriter, loc, tileType, paddedInput,
+                  SmallVector<OpFoldResult> {
+                    batch, rewriter.getIndexAttr(tile * tileWidth), sourceRow, sourceColumn},
+                  SmallVector<OpFoldResult> {
+                    rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                    rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+                  getUnitStrides(rewriter, 4));
+                point = materializeTileTensor(rewriter, loc, point);
+                reduced = spatial::SpatVMaxOp::create(
+                  rewriter, loc, tileType, reduced, point);
+              }
+            }
+            updated = tensor::InsertSliceOp::create(
+              rewriter, loc, reduced, updated,
+              SmallVector<OpFoldResult> {
+                batch, rewriter.getIndexAttr(tile * tileWidth), outputRow, outputColumn},
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+              getUnitStrides(rewriter, 4));
+          }
+          yielded.push_back(updated);
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), loop->results.front());
+      return success();
+    });
+  if (failed(compute))
+    return failure();
+  return compute->getResult(0);
+}
+
 static Value createClampedPoolIndexTable(PatternRewriter& rewriter,
                                          Operation* anchorOp,
                                          int64_t outputSize,
@@ -496,13 +616,15 @@ static Value extractPoolIndex(PatternRewriter& rewriter,
 }
 
 FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                                            Value input,
                                             std::optional<Value> rowStripInput,
+                                            const spatial::SpatialTargetResources& target,
                                             PatternRewriter& rewriter) {
-  if (failed(canLowerMaxPoolPlanToRowStrip(planOp)))
+  if (failed(canLowerMaxPoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
-  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = cast<RankedTensorType>(input.getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   const int64_t channels = inputType.getDimSize(1);
   const int64_t inputHeight = inputType.getDimSize(2);
@@ -511,9 +633,9 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   const int64_t outputWidth = outputType.getDimSize(3);
   const int64_t kernelHeight = planOp.getKernelShape()[0];
   const int64_t kernelWidth = planOp.getKernelShape()[1];
-  Value input = rowStripInput.value_or(planOp.getInput());
-  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
+  Value actualInput = rowStripInput.value_or(input);
+  auto actualInputType = dyn_cast<RankedTensorType>(actualInput.getType());
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(actualInput, inputType);
   const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
@@ -561,7 +683,7 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
     TypeRange {outputStorageType},
     outputHeight * tilesPerRow,
     {},
-    ValueRange {input},
+    ValueRange {actualInput},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       SmallVector<Value> inputRows;
       inputRows.reserve(kernelHeight);
@@ -590,8 +712,8 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                                                                                 rewriter.getIndexAttr(1),
                                                                                 rewriter.getIndexAttr(inputWidth)},
                                                      getUnitStrides(rewriter, 4));
-          inputRows.push_back(ONNXTransposeOp::create(
-            rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})));
+          inputRows.push_back(createLinalgTranspose(
+            nchw, inputFragmentType, {0, 2, 3, 1}, rewriter, loc));
         }
       }
 
@@ -685,7 +807,8 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   return batch->getResult(0);
 }
 
-LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp) {
+LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
+  spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetResources&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -697,22 +820,101 @@ LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAverage
   return success();
 }
 
+FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
+    spatial::SpatGlobalAveragePoolPlanOp planOp,
+    Value input,
+    const spatial::SpatialTargetResources& target,
+    PatternRewriter& rewriter) {
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!inputType || !outputType || !inputType.hasStaticShape()
+      || !outputType.hasStaticShape() || inputType.getRank() != 4
+      || outputType.getRank() != 4 || inputType.getDimSize(0) != 1
+      || outputType.getDimSize(0) != 1 || inputType.getDimSize(1) != outputType.getDimSize(1)
+      || outputType.getDimSize(2) != 1 || outputType.getDimSize(3) != 1)
+    return planOp.emitOpError("dense global AveragePool lowering requires static rank-4 floating-point tensors"),
+           failure();
+  auto elementType = dyn_cast<FloatType>(inputType.getElementType());
+  if (!elementType)
+    return planOp.emitOpError("dense global AveragePool lowering requires floating-point tensors"),
+           failure();
+
+  const int64_t channels = inputType.getDimSize(1);
+  const int64_t height = inputType.getDimSize(2);
+  const int64_t width = inputType.getDimSize(3);
+  const int64_t tileWidth = std::max<int64_t>(1, target.matrixShape.rows);
+  const int64_t channelTileCount = (channels + tileWidth - 1) / tileWidth;
+  const double scaleValue = 1.0 / static_cast<double>(height * width);
+
+  auto compute = createSpatCompute<1>(
+    rewriter, planOp.getLoc(), outputType, {}, input,
+    [&](Value input) -> LogicalResult {
+      Value output = tensor::EmptyOp::create(
+        rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      for (int64_t tile = 0; tile < channelTileCount; ++tile) {
+        const int64_t tileChannels = std::min<int64_t>(tileWidth, channels - tile * tileWidth);
+        auto tileType = RankedTensorType::get(
+          {1, tileChannels, 1, 1}, outputType.getElementType());
+        Value reduced = createPoolFillTensor(
+          rewriter, planOp.getLoc(), tileType, /*useMinimumValue=*/false);
+        for (int64_t row = 0; row < height; ++row) {
+          for (int64_t column = 0; column < width; ++column) {
+            Value point = tensor::ExtractSliceOp::create(
+              rewriter, planOp.getLoc(), tileType, input,
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(0), rewriter.getIndexAttr(tile * tileWidth),
+                rewriter.getIndexAttr(row), rewriter.getIndexAttr(column)},
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+              getUnitStrides(rewriter, 4));
+            point = materializeTileTensor(rewriter, planOp.getLoc(), point);
+            reduced = spatial::SpatVAddOp::create(
+              rewriter, planOp.getLoc(), tileType, reduced, point);
+          }
+        }
+        auto scaleAttr = DenseElementsAttr::get(
+          tileType, rewriter.getFloatAttr(elementType, scaleValue));
+        Value scale = getOrCreateConstant(rewriter, anchor, scaleAttr, tileType);
+        reduced = spatial::SpatVMulOp::create(
+          rewriter, planOp.getLoc(), tileType, reduced, scale);
+        output = tensor::InsertSliceOp::create(
+          rewriter, planOp.getLoc(), reduced, output,
+          SmallVector<OpFoldResult> {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(tile * tileWidth),
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+          SmallVector<OpFoldResult> {
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+          getUnitStrides(rewriter, 4));
+      }
+      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), output);
+      return success();
+    });
+  if (failed(compute))
+    return failure();
+  return compute->getResult(0);
+}
+
 FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                                    Value input,
                                                     std::optional<Value> rowStripInput,
+                                                    const spatial::SpatialTargetResources& target,
                                                     PatternRewriter& rewriter) {
-  if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp)))
+  if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
-  auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
+  auto inputType = cast<RankedTensorType>(input.getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
   auto elementType = dyn_cast<FloatType>(inputType.getElementType());
   if (!elementType)
     return failure();
 
-  Value input = rowStripInput.value_or(planOp.getInput());
-  auto actualInputType = dyn_cast<RankedTensorType>(input.getType());
-  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(input, inputType);
+  Value actualInput = rowStripInput.value_or(input);
+  auto actualInputType = dyn_cast<RankedTensorType>(actualInput.getType());
+  FailureOr<RowStripPhysicalValue> physicalValue = describeRowStripPhysicalValue(actualInput, inputType);
   const bool physicalInput = succeeded(physicalValue);
   if (!physicalInput && actualInputType != inputType)
     return failure();
@@ -742,7 +944,7 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
     TypeRange {outputStorageType},
     tilesPerRow,
     ValueRange {zero, scale},
-    ValueRange {input},
+    ValueRange {actualInput},
     [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
       Value reduced = args.weights[0];
       for (int64_t row = 0; row < height; ++row) {
@@ -777,8 +979,8 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
                                       rewriter.getIndexAttr(1),
                                       rewriter.getIndexAttr(width)},
             getUnitStrides(rewriter, 4));
-          fragment = ONNXTransposeOp::create(
-            rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
+          fragment = createLinalgTranspose(
+            nchw, inputFragmentType, {0, 2, 3, 1}, rewriter, loc);
         }
         for (int64_t column = 0; column < width; ++column) {
           Value point = tensor::ExtractSliceOp::create(
@@ -811,9 +1013,11 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
   return batch->getResult(0);
 }
 
-void populatePoolPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx);
-  patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx);
+void populatePoolPatterns(RewritePatternSet& patterns,
+                          MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& target) {
+  patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx, target);
+  patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx, target);
 }
 
 } // namespace onnx_mlir

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

@@ -17,7 +17,7 @@ struct ReluToSpatialCompute : OpConversionPattern<ONNXReluOp> {
     Location loc = reluOp.getLoc();
     Type resultType = reluOp.getResult().getType();
     auto reluPlan = spatial::SpatReluPlanOp::create(
-      rewriter, loc, resultType, adaptor.getX(), rewriter.getStringAttr("nchw"));
+      rewriter, loc, resultType, adaptor.getX(), spatial::getNCHWLayout(rewriter.getContext()));
     rewriter.replaceOp(reluOp, reluPlan.getResult());
     return success();
   }

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

@@ -32,7 +32,8 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
              return type && type.hasStaticShape() && type.getRank() == 4;
            })) {
       rewriter.replaceOpWithNewOp<spatial::SpatConcatPlanOp>(
-        maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis), rewriter.getStringAttr("nchw"));
+        maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis),
+        spatial::getNCHWLayout(rewriter.getContext()));
       return success();
     }
 

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

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

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

@@ -5,8 +5,11 @@
 
 #include "llvm/ADT/STLExtras.h"
 
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #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/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -17,126 +20,144 @@ namespace onnx_mlir {
 namespace {
 
 static Value buildNearestAsymmetricIndex(
-  Value outputIndex, int64_t inputDim, int64_t outputDim, ConversionPatternRewriter& rewriter, Location loc) {
+  Value outputIndex, int64_t inputDim, int64_t outputDim, PatternRewriter& rewriter, Location loc) {
+  if (inputDim == outputDim)
+    return outputIndex;
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  if (outputDim % inputDim == 0)
+    return affineFloorDivConst(rewriter, loc, outputIndex, outputDim / inputDim, anchorOp);
+  if (inputDim % outputDim == 0)
+    return affineMulConst(rewriter, loc, outputIndex, inputDim / outputDim, anchorOp);
   Value cInputDim = getOrCreateIndexConstant(rewriter, anchorOp, inputDim);
   Value cOutputDim = getOrCreateIndexConstant(rewriter, anchorOp, outputDim);
-  Value cInputDimLast = getOrCreateIndexConstant(rewriter, anchorOp, inputDim - 1);
   Value scaledIndex = arith::MulIOp::create(rewriter, loc, outputIndex, cInputDim);
-  Value inputIndex = arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim);
-  return arith::MinUIOp::create(rewriter, loc, inputIndex, cInputDimLast);
+  return arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim);
 }
 
-static FailureOr<Value> buildNearestResizeLoop(Value input,
-                                               RankedTensorType inputType,
-                                               RankedTensorType resultType,
-                                               ConversionPatternRewriter& rewriter,
-                                               Location loc) {
-  auto elemType = resultType.getElementType();
-  SmallVector<int64_t> unitShape(resultType.getRank(), 1);
-  auto unitTensorType = RankedTensorType::get(unitShape, elemType);
-
-  SmallVector<OpFoldResult> unitSizes(resultType.getRank(), rewriter.getIndexAttr(1));
-  SmallVector<OpFoldResult> unitStrides(resultType.getRank(), rewriter.getIndexAttr(1));
-
-  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
-  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-  Value cOutputN = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(0));
-  Value cOutputC = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(1));
-  Value cOutputH = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(2));
-  Value cOutputW = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(3));
-
-  Value outputInit = tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), elemType);
-
-  auto batchLoop = buildNormalizedScfFor(
-    rewriter,
-    loc,
-    c0,
-    cOutputN,
-    c1,
-    ValueRange {outputInit},
-    [&](OpBuilder&, Location nestedLoc, Value outputN, ValueRange batchIterArgs, SmallVectorImpl<Value>& batchYielded) {
-      Value outputBatchAcc = batchIterArgs.front();
-      Value inputN =
-        buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), resultType.getDimSize(0), rewriter, nestedLoc);
-
-      auto channelLoop = buildNormalizedScfFor(
-        rewriter,
-        nestedLoc,
-        c0,
-        cOutputC,
-        c1,
-        ValueRange {outputBatchAcc},
-        [&](OpBuilder&,
-            Location channelLoc,
-            Value outputC,
-            ValueRange channelIterArgs,
-            SmallVectorImpl<Value>& channelYielded) {
-          Value outputChannelAcc = channelIterArgs.front();
-          Value inputC = buildNearestAsymmetricIndex(
-            outputC, inputType.getDimSize(1), resultType.getDimSize(1), rewriter, channelLoc);
-
-          auto heightLoop = buildNormalizedScfFor(
-            rewriter,
-            channelLoc,
-            c0,
-            cOutputH,
-            c1,
-            ValueRange {outputChannelAcc},
-            [&](OpBuilder&,
-                Location heightLoc,
-                Value outputH,
-                ValueRange heightIterArgs,
-                SmallVectorImpl<Value>& heightYielded) {
-              Value outputHeightAcc = heightIterArgs.front();
-              Value inputH = buildNearestAsymmetricIndex(
-                outputH, inputType.getDimSize(2), resultType.getDimSize(2), rewriter, heightLoc);
-
-              auto widthLoop = buildNormalizedScfFor(
-                rewriter,
-                heightLoc,
-                c0,
-                cOutputW,
-                c1,
-                ValueRange {outputHeightAcc},
-                [&](OpBuilder&,
-                    Location widthLoc,
-                    Value outputW,
-                    ValueRange widthIterArgs,
-                    SmallVectorImpl<Value>& widthYielded) {
-                  Value outputWidthAcc = widthIterArgs.front();
-                  Value inputW = buildNearestAsymmetricIndex(
-                    outputW, inputType.getDimSize(3), resultType.getDimSize(3), rewriter, widthLoc);
-
-                  SmallVector<OpFoldResult> inputOffsets = {inputN, inputC, inputH, inputW};
-                  Value inputSlice = tensor::ExtractSliceOp::create(
-                    rewriter, widthLoc, unitTensorType, input, inputOffsets, unitSizes, unitStrides);
-
-                  SmallVector<OpFoldResult> outputOffsets = {outputN, outputC, outputH, outputW};
-                  Value updatedOutput = tensor::InsertSliceOp::create(
-                    rewriter, widthLoc, inputSlice, outputWidthAcc, outputOffsets, unitSizes, unitStrides);
-                  widthYielded.push_back(updatedOutput);
-                  return success();
-                });
-              if (failed(widthLoop))
-                return failure();
-              heightYielded.push_back(widthLoop->results.front());
-              return success();
-            });
-          if (failed(heightLoop))
-            return failure();
-          channelYielded.push_back(heightLoop->results.front());
+static FailureOr<Value> buildDenseNearestResize(Value input,
+                                                RankedTensorType inputType,
+                                                RankedTensorType resultType,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  ArrayRef<int64_t> shape = resultType.getShape();
+  int64_t rowCount = shape[0] * shape[1] * shape[2];
+  auto scalarType = RankedTensorType::get({1, 1, 1, 1}, resultType.getElementType());
+  auto rowType = RankedTensorType::get({1, 1, 1, shape[3]}, resultType.getElementType());
+  auto rowsType = RankedTensorType::get({rowCount, 1, 1, 1, shape[3]}, resultType.getElementType());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, TypeRange {rowsType}, rowCount, {}, ValueRange {input},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value outputN = affineFloorDivConst(rewriter, loc, args.lane, shape[1] * shape[2], anchor);
+      Value channelRow = affineModConst(rewriter, loc, args.lane, shape[1] * shape[2], anchor);
+      Value outputC = affineFloorDivConst(rewriter, loc, channelRow, shape[2], anchor);
+      Value outputH = affineModConst(rewriter, loc, channelRow, shape[2], anchor);
+      Value inputN = buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), shape[0], rewriter, loc);
+      Value inputC = buildNearestAsymmetricIndex(outputC, inputType.getDimSize(1), shape[1], rewriter, loc);
+      Value inputH = buildNearestAsymmetricIndex(outputH, inputType.getDimSize(2), shape[2], rewriter, loc);
+      Value row = tensor::EmptyOp::create(rewriter, loc, rowType.getShape(), rowType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value width = getOrCreateIndexConstant(rewriter, anchor, shape[3]);
+      auto loop = buildNormalizedScfFor(
+        rewriter, loc, c0, width, c1, ValueRange {row},
+        [&](OpBuilder&, Location nestedLoc, Value outputW, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value inputW = buildNearestAsymmetricIndex(
+            outputW, inputType.getDimSize(3), shape[3], rewriter, nestedLoc);
+          SmallVector<OpFoldResult> unitSizes(4, rewriter.getIndexAttr(1));
+          SmallVector<OpFoldResult> unitStrides(4, rewriter.getIndexAttr(1));
+          Value scalar = tensor::ExtractSliceOp::create(
+            rewriter, nestedLoc, scalarType, args.inputs.front(),
+            SmallVector<OpFoldResult> {inputN, inputC, inputH, inputW}, unitSizes, unitStrides);
+          yielded.push_back(tensor::InsertSliceOp::create(
+            rewriter, nestedLoc, scalar, iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                       rewriter.getIndexAttr(0), outputW},
+            unitSizes, unitStrides));
           return success();
         });
-      if (failed(channelLoop))
+      assert(succeeded(loop) && "nearest Resize row loop construction must succeed");
+      publishGraphBatchPhysicalFragment(rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
+    });
+  if (failed(batch))
+    return failure();
+
+  SmallVector<FragmentAssemblyEntry> entries;
+  entries.reserve(rowCount);
+  for (int64_t n = 0; n < shape[0]; ++n)
+    for (int64_t c = 0; c < shape[1]; ++c)
+      for (int64_t h = 0; h < shape[2]; ++h)
+        entries.push_back({(n * shape[1] + c) * shape[2] + h, 0, {n, c, h, 0}, {1, 1, 1, shape[3]}});
+  return createFragmentAssemblyBlueprint(
+    batch->getResult(0), resultType, entries, "dense_nchw", spatial::kContiguousRowMajorFragments, rewriter, loc);
+}
+
+static FailureOr<Value> buildRowStripNearestResize(
+    Value storage, RankedTensorType inputType, RankedTensorType resultType,
+    PatternRewriter& rewriter, Location loc) {
+  auto input = describeRowStripPhysicalValue(storage, inputType);
+  if (failed(input))
+    return failure();
+  int64_t tilesPerRow = input->tilesPerRow;
+  int64_t outputHeight = resultType.getDimSize(2);
+  int64_t outputWidth = resultType.getDimSize(3);
+  int64_t tileChannels = input->fragmentType.getDimSize(3);
+  int64_t laneCount = outputHeight * tilesPerRow;
+  auto outputFragmentType = RankedTensorType::get(
+    {1, 1, outputWidth, tileChannels}, resultType.getElementType());
+  auto outputStorageType = spatial::getGraphBatchPhysicalResultType(
+    laneCount, outputFragmentType);
+  auto pixelType = RankedTensorType::get(
+    {1, 1, 1, tileChannels}, resultType.getElementType());
+  auto batch = createSpatComputeBatch(
+    rewriter, loc, TypeRange {outputStorageType}, laneCount, {}, ValueRange {storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value outputRow = affineFloorDivConst(rewriter, loc, args.lane, tilesPerRow, anchor);
+      Value tile = affineModConst(rewriter, loc, args.lane, tilesPerRow, anchor);
+      Value inputRow = buildNearestAsymmetricIndex(
+        outputRow, inputType.getDimSize(2), outputHeight, rewriter, loc);
+      Value inputSlot = arith::AddIOp::create(
+        rewriter, loc, affineMulConst(rewriter, loc, inputRow, tilesPerRow, anchor), tile);
+      auto source = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), inputSlot, input->fragmentType);
+      if (failed(source))
         return failure();
-      batchYielded.push_back(channelLoop->results.front());
+      Value initial = tensor::EmptyOp::create(
+        rewriter, loc, outputFragmentType.getShape(), resultType.getElementType());
+      Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value width = getOrCreateIndexConstant(rewriter, anchor, outputWidth);
+      auto loop = buildNormalizedScfFor(
+        rewriter, loc, c0, width, c1, ValueRange {initial},
+        [&](OpBuilder&, Location nestedLoc, Value outputColumn, ValueRange iterArgs,
+            SmallVectorImpl<Value>& yielded) {
+          Value inputColumn = buildNearestAsymmetricIndex(
+            outputColumn, inputType.getDimSize(3), outputWidth, rewriter, nestedLoc);
+          Value pixel = tensor::ExtractSliceOp::create(
+            rewriter, nestedLoc, pixelType, *source,
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                      inputColumn, rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)},
+            getUnitStrides(rewriter, 4));
+          yielded.push_back(tensor::InsertSliceOp::create(
+            rewriter, nestedLoc, pixel, iterArgs.front(),
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                      outputColumn, rewriter.getIndexAttr(0)},
+            SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
+                                      rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels)},
+            getUnitStrides(rewriter, 4)));
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      publishGraphBatchPhysicalFragment(
+        rewriter, loc, loop->results.front(), args.outputs.front(), args.lane);
       return success();
     });
-  if (failed(batchLoop))
-    return failure();
-  return batchLoop->results.front();
+  return failed(batch) ? FailureOr<Value>(failure())
+                       : FailureOr<Value>(batch->getResult(0));
 }
 
 struct Resize : OpConversionPattern<ONNXResizeOp> {
@@ -161,23 +182,41 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
         || llvm::any_of(resultType.getShape(), [](int64_t dim) { return dim <= 0; }))
       return rewriter.notifyMatchFailure(resizeOp, "resize lowering requires positive static dimensions.");
 
-    auto computeOp = createSpatCompute<1>(
-      rewriter, resizeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getX(), [&](Value x) -> LogicalResult {
-        auto result = buildNearestResizeLoop(x, inputType, resultType, rewriter, resizeOp.getLoc());
-        if (failed(result))
-          return failure();
-        spatial::SpatYieldOp::create(rewriter, resizeOp.getLoc(), *result);
-        return success();
-      });
-    if (failed(computeOp))
-      return failure();
-    rewriter.replaceOp(resizeOp, computeOp->getResults());
+    auto plan = spatial::SpatResizeNearestPlanOp::create(
+      rewriter, resizeOp.getLoc(), resultType, adaptor.getX(), spatial::getNCHWLayout(rewriter.getContext()));
+    rewriter.replaceOp(resizeOp, plan.getResult());
     return success();
   }
 };
 
 } // namespace
 
+LogicalResult canLowerResizeNearestPlanToRowStrip(
+    spatial::SpatResizeNearestPlanOp planOp,
+    const spatial::SpatialTargetResources&) {
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  return success(inputType && outputType && inputType.hasStaticShape()
+    && outputType.hasStaticShape() && inputType.getRank() == 4
+    && outputType.getRank() == 4 && inputType.getDimSize(0) == 1
+    && outputType.getDimSize(0) == 1
+    && inputType.getDimSize(1) == outputType.getDimSize(1));
+}
+
+FailureOr<Value> lowerSelectedResizeNearestPlan(
+    spatial::SpatResizeNearestPlanOp planOp, Value input,
+    std::optional<Value> rowStripInput,
+    const spatial::SpatialTargetResources&,
+    PatternRewriter& rewriter) {
+  auto inputType = cast<RankedTensorType>(input.getType());
+  auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());
+  if (rowStripInput)
+    return buildRowStripNearestResize(
+      *rowStripInput, inputType, outputType, rewriter, planOp.getLoc());
+  return buildDenseNearestResize(
+    input, inputType, outputType, rewriter, planOp.getLoc());
+}
+
 void populateResizePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Resize>(ctx); }
 
 } // namespace onnx_mlir

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

@@ -61,6 +61,74 @@ static FailureOr<Value> materializeTransposedConstant(Value input,
                              resultType);
 }
 
+static FailureOr<Value> transposeFragmentAssemblyBlueprint(spatial::SpatBlueprintOp blueprint,
+                                                           RankedTensorType resultType,
+                                                           ArrayRef<int64_t> permutation,
+                                                           ConversionPatternRewriter& rewriter,
+                                                           Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(blueprint.getInput().getType());
+  auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+  auto fragmentStrides = blueprint.getFragmentStrides();
+  if (!storageType || !storageType.hasStaticShape() || !resultType.hasStaticShape()
+      || !blueprint.getFragments().empty() || !spatial::isFragmentAssembly(blueprint.getMode())
+      || !blueprint.getFragmentOperandIndices() || !sourceOffsets || !fragmentStrides
+      || llvm::any_of(*sourceOffsets, [](int64_t offset) { return offset != 0; })
+      || storageType.getRank() != resultType.getRank() + 1)
+    return failure();
+  if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments
+      && !spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+    return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+
+  SmallVector<int64_t> outputStorageShape {storageType.getDimSize(0)};
+  for (int64_t sourceDim : permutation)
+    outputStorageShape.push_back(storageType.getDimSize(sourceDim + 1));
+  auto outputStorageType = RankedTensorType::get(outputStorageShape, storageType.getElementType());
+  auto mapped = mapGraphBatchFragments(
+    blueprint.getInput(), outputStorageType, rewriter, loc, [&](Value fragment, RankedTensorType fragmentType) {
+      Value init = createTransposeInit(fragment, fragmentType, permutation, rewriter, loc);
+      return FailureOr<Value>(
+        linalg::TransposeOp::create(rewriter, loc, fragment, init, permutation).getResult()[0]);
+    });
+  if (failed(mapped))
+    return failure();
+
+  const int64_t rank = resultType.getRank();
+  const int64_t fragmentCount = blueprint.getFragmentOperandIndices()->size();
+  SmallVector<int64_t> offsets, sizes, strides;
+  offsets.reserve(fragmentCount * rank);
+  sizes.reserve(fragmentCount * rank);
+  strides.reserve(fragmentCount * rank);
+  ArrayRef<int64_t> inputOffsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> inputSizes = blueprint.getFragmentSizes();
+  for (int64_t fragment = 0; fragment < fragmentCount; ++fragment)
+    for (int64_t sourceDim : permutation) {
+      const int64_t index = fragment * rank + sourceDim;
+      offsets.push_back(inputOffsets[index]);
+      sizes.push_back(inputSizes[index]);
+      strides.push_back((*fragmentStrides)[index]);
+    }
+  auto transposedBlueprint = spatial::SpatBlueprintOp::create(rewriter,
+                                                               loc,
+                                                               resultType,
+                                                               *mapped,
+                                                               ValueRange {},
+                                                               blueprint.getLogicalLayoutAttr(),
+                                                               spatial::getFragmentedLayout(rewriter.getContext()),
+                                                               rewriter.getDenseI64ArrayAttr(offsets),
+                                                               rewriter.getDenseI64ArrayAttr(sizes),
+                                                               rewriter.getStringAttr("permuted_fragments"),
+                                                               blueprint.getModeAttr(),
+                                                               blueprint.getFragmentOperandIndicesAttr(),
+                                                               blueprint.getFragmentSourceSlotsAttr(),
+                                                               blueprint.getFragmentSourceOffsetsAttr(),
+                                                               rewriter.getDenseI64ArrayAttr(strides),
+                                                               blueprint.getConflictPolicyAttr(),
+                                                               blueprint.getCoveragePolicyAttr());
+  if (spatial::isCanonicalContiguousRowMajorFragmentAssembly(transposedBlueprint))
+    transposedBlueprint.setIndexMapAttr(rewriter.getStringAttr(spatial::kContiguousRowMajorFragments));
+  return transposedBlueprint.getOutput();
+}
+
 struct TransposeToLinalgTranspose : OpConversionPattern<ONNXTransposeOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -75,6 +143,14 @@ struct TransposeToLinalgTranspose : OpConversionPattern<ONNXTransposeOp> {
     auto permutation = getTransposePermutationChecked(transposeOp.getPermAttr(), inputType.getRank());
     if (failed(permutation))
       return failure();
+    if (auto blueprint = adaptor.getData().getDefiningOp<spatial::SpatBlueprintOp>()) {
+      auto transposed =
+        transposeFragmentAssemblyBlueprint(blueprint, resultType, *permutation, rewriter, transposeOp.getLoc());
+      if (succeeded(transposed)) {
+        rewriter.replaceOp(transposeOp, *transposed);
+        return success();
+      }
+    }
     if (isCompileTimeComputable(adaptor.getData())) {
       auto constantTranspose =
         materializeTransposedConstant(adaptor.getData(), resultType, *permutation, rewriter, transposeOp.getLoc());

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

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

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -1,341 +0,0 @@
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Pass/Pass.h"
-
-#include "llvm/ADT/DenseMap.h"
-
-#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace {
-
-static constexpr StringLiteral kLogicalLayout = "nchw";
-static constexpr StringLiteral kDenseLayout = "dense_nchw";
-static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
-
-enum class SelectedLayout {
-  DenseNchw,
-  PixelMajorRowStrip,
-};
-
-static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& layouts, Value value) {
-  auto it = layouts.find(value);
-  return it == layouts.end() ? SelectedLayout::DenseNchw : it->second;
-}
-
-static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
-    return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(user))
-    return getSelectedLayout(layouts, siluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
-    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(user))
-    return getSelectedLayout(layouts, addPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(user))
-    return getSelectedLayout(layouts, concatPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
-    return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
-    return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
-    return getSelectedLayout(layouts, averagePoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto flattenCompute = dyn_cast<spatial::SpatGraphCompute>(user))
-    return succeeded(canLowerFlattenFromRowStrip(flattenCompute));
-  return false;
-}
-
-static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  for (Operation* user : value.getUsers()) {
-    if (usesSelectedRowStrip(user, layouts))
-      continue;
-    // Dense-only users must be materialized explicitly.
-    continue;
-  }
-  return true;
-}
-
-static bool canConsumeRowStripAsUser(Operation* user) {
-  if (isa<spatial::SpatReluPlanOp, spatial::SpatSiluPlanOp>(user))
-    return true;
-  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
-    auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
-    return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
-  }
-  if (isa<spatial::SpatAddPlanOp>(user))
-    return true;
-  if (isa<spatial::SpatConcatPlanOp>(user))
-    return true;
-  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
-    return succeeded(canConsumeAndProduceRowStrip(convPlan));
-  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
-    return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan));
-  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
-    return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan));
-  return false;
-}
-
-static bool hasRowStripConsumer(Value value) {
-  for (Operation* user : value.getUsers())
-    if (canConsumeRowStripAsUser(user))
-      return true;
-  return false;
-}
-
-static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
-                                  llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
-  if (inputLayout == SelectedLayout::PixelMajorRowStrip)
-    return succeeded(canConsumeAndProduceRowStrip(convPlan));
-  return succeeded(canLowerConvPlanToRowStrip(convPlan));
-}
-
-static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
-                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (!canSelectConvRowStrip(convPlan, layouts))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseActivationLayout(Value input,
-                                             Value result,
-                                             llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip)
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(result, layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
-                                          llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
-    return SelectedLayout::DenseNchw;
-  auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
-  if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType))
-    return SelectedLayout::DenseNchw;
-  if (!hasRowStripConsumer(biasAddPlan.getResult()))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseAddLayout(spatial::SpatAddPlanOp addPlan, llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, addPlan.getLhs()) != SelectedLayout::PixelMajorRowStrip
-      || getSelectedLayout(layouts, addPlan.getRhs()) != SelectedLayout::PixelMajorRowStrip)
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(addPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseConcatLayout(spatial::SpatConcatPlanOp concatPlan,
-                                         llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (llvm::any_of(concatPlan.getInputs(), [&](Value input) {
-        return getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip;
-      }))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(concatPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) {
-  return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::PixelMajorRowStrip
-                                                               : SelectedLayout::DenseNchw;
-}
-
-static SelectedLayout chooseGlobalAveragePoolLayout(
-  spatial::SpatGlobalAveragePoolPlanOp averagePoolPlan) {
-  return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan))
-           ? SelectedLayout::PixelMajorRowStrip
-           : SelectedLayout::DenseNchw;
-}
-
-static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
-  auto outputType = cast<RankedTensorType>(value.getType());
-  auto [offsets, sizes] = buildRowStripMetadata(outputType);
-  return spatial::SpatBlueprintOp::create(rewriter,
-                                              value.getLoc(),
-                                              outputType,
-                                              value,
-                                              ValueRange {},
-                                              rewriter.getStringAttr(kLogicalLayout),
-                                              rewriter.getStringAttr(kRowStripLayout),
-                                              rewriter.getDenseI64ArrayAttr(offsets),
-                                              rewriter.getDenseI64ArrayAttr(sizes),
-                                              rewriter.getStringAttr(kRowStripIndexMap),
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr);
-}
-
-static void materializeDenseUses(IRRewriter& rewriter,
-                                 Value layoutValue,
-                                 llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  SmallVector<OpOperand*> denseUses;
-  for (OpOperand& use : layoutValue.getUses()) {
-    if (usesSelectedRowStrip(use.getOwner(), layouts))
-      continue;
-    denseUses.push_back(&use);
-  }
-
-  for (OpOperand* use : denseUses) {
-    Operation* owner = use->getOwner();
-    rewriter.setInsertionPoint(owner);
-    auto materialized = spatial::SpatMaterializeLayoutOp::create(rewriter,
-                                                                 owner->getLoc(),
-                                                                 use->get().getType(),
-                                                                 use->get(),
-                                                                 rewriter.getStringAttr(kLogicalLayout),
-                                                                 rewriter.getStringAttr(kRowStripLayout),
-                                                                 rewriter.getStringAttr(kDenseLayout));
-    use->set(materialized.getResult());
-  }
-}
-
-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 conservative Spatial layouts and insert reconciliation barriers."; }
-
-  void runOnOperation() override {
-    auto entryFunc = getPimEntryFunc(getOperation());
-    if (failed(entryFunc)) {
-      getOperation().emitError("failed to locate the PIM entry function during Spatial layout planning");
-      signalPassFailure();
-      return;
-    }
-
-    func::FuncOp funcOp = *entryFunc;
-    IRRewriter rewriter(&getContext());
-    llvm::DenseMap<Value, SelectedLayout> layouts;
-
-    bool changed = true;
-    while (changed) {
-      changed = false;
-      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
-        if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
-          SelectedLayout selected = chooseConvLayout(convPlan, layouts);
-          if (layouts[convPlan.getResult()] != selected) {
-            layouts[convPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
-          SelectedLayout selected = chooseActivationLayout(reluPlan.getInput(), reluPlan.getResult(), layouts);
-          if (layouts[reluPlan.getResult()] != selected) {
-            layouts[reluPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op)) {
-          SelectedLayout selected = chooseActivationLayout(siluPlan.getInput(), siluPlan.getResult(), layouts);
-          if (layouts[siluPlan.getResult()] != selected) {
-            layouts[siluPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
-          SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
-          if (layouts[biasAddPlan.getResult()] != selected) {
-            layouts[biasAddPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
-          SelectedLayout selected = chooseAddLayout(addPlan, layouts);
-          if (layouts[addPlan.getResult()] != selected) {
-            layouts[addPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
-          SelectedLayout selected = chooseConcatLayout(concatPlan, layouts);
-          if (layouts[concatPlan.getResult()] != selected) {
-            layouts[concatPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
-          SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan);
-          if (layouts[maxPoolPlan.getResult()] != selected) {
-            layouts[maxPoolPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
-          SelectedLayout selected = chooseGlobalAveragePoolLayout(averagePoolPlan);
-          if (layouts[averagePoolPlan.getResult()] != selected) {
-            layouts[averagePoolPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-      }
-    }
-
-    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
-      Value producedValue;
-      if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op))
-        producedValue = convPlan.getResult();
-      else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
-        producedValue = biasAddPlan.getResult();
-      else if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op))
-        producedValue = addPlan.getResult();
-      else if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op))
-        producedValue = concatPlan.getResult();
-      else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
-        producedValue = reluPlan.getResult();
-      else if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op))
-        producedValue = siluPlan.getResult();
-      else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
-        producedValue = maxPoolPlan.getResult();
-      else if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op))
-        producedValue = averagePoolPlan.getResult();
-      else
-        continue;
-
-      if (getSelectedLayout(layouts, producedValue) != SelectedLayout::PixelMajorRowStrip)
-        continue;
-
-      rewriter.setInsertionPointAfter(&op);
-      auto blueprint = insertRowStripBlueprint(rewriter, producedValue);
-      rewriter.replaceAllUsesExcept(producedValue, blueprint.getResult(), blueprint);
-      materializeDenseUses(rewriter, blueprint.getResult(), layouts);
-    }
-    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
-      getOperation().emitError("logical Spatial graph verification failed after SpatialLayoutPlanning");
-      signalPassFailure();
-    }
-  }
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createSpatialLayoutPlanningPass() { return std::make_unique<SpatialLayoutPlanningPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -149,11 +149,10 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
     auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
     if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
       return failure();
-    std::optional<StringRef> mode = blueprint.getMode();
     std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
     std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
     std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
-    if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr)
+    if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr)
       return failure();
     if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
       return failure();
@@ -418,8 +417,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
     rewriter.setInsertionPointToEnd(newBlock);
 
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> modeAttr = blueprint.getMode();
-      if (modeAttr && *modeAttr == "fragment_assembly") {
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
         for (Operation* user : blueprint.getOutput().getUsers()) {
           if (!isa<tensor::ParallelInsertSliceOp>(user))
             return blueprint.emitOpError(
@@ -483,8 +481,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
         auto hostTargetType = cast<ShapedType>(hostTarget.getType());
         if (auto blueprint =
               insertSlice.getSource().getDefiningOp<spatial::SpatBlueprintOp>()) {
-          std::optional<StringRef> modeAttr = blueprint.getMode();
-          if (modeAttr && *modeAttr == "fragment_assembly") {
+          if (spatial::isFragmentAssembly(blueprint.getMode())) {
             FailureOr<SmallVector<FragmentAssemblyCopy, 8>> fragmentAssemblyCopies =
               collectFragmentAssemblyCopiesFromBlueprint(blueprint, mapper, /*lane=*/0, /*hostTargetIndex=*/0);
             if (failed(fragmentAssemblyCopies))

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,7 +1,10 @@
 #include "mlir/IR/ValueRange.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
 
 #include "llvm/ADT/STLExtras.h"
 
@@ -28,6 +31,49 @@ FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anc
   return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
 }
 
+LogicalResult materializePipelineHostBuffer(
+    func::FuncOp funcOp, RewriterBase &rewriter) {
+  auto bytes = funcOp->getAttrOfType<IntegerAttr>(
+    kPipelineHostBufferBytesAttrName);
+  if (!bytes)
+    return success();
+  if (bytes.getInt() <= 0)
+    return funcOp.emitOpError(
+      "pipeline host transfer buffer must be positive");
+  ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
+  if (moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName))
+    return funcOp.emitOpError(
+      "pipeline host transfer buffer symbol already exists");
+  auto type = MemRefType::get(
+    {bytes.getInt()}, rewriter.getI8Type());
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(moduleOp.getBody());
+  memref::GlobalOp::create(
+    rewriter, funcOp.getLoc(),
+    rewriter.getStringAttr(kPipelineHostBufferName),
+    rewriter.getStringAttr("private"), TypeAttr::get(type), Attribute(),
+    UnitAttr(), IntegerAttr());
+  return success();
+}
+
+FailureOr<mlir::Value> getPipelineHostBuffer(
+    OpBuilder &builder, Operation *anchor) {
+  auto funcOp = anchor->getParentOfType<func::FuncOp>();
+  auto moduleOp = anchor->getParentOfType<ModuleOp>();
+  auto bytes = funcOp
+    ? funcOp->getAttrOfType<IntegerAttr>(kPipelineHostBufferBytesAttrName)
+    : IntegerAttr();
+  auto global = moduleOp
+    ? moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName)
+    : memref::GlobalOp();
+  if (!bytes || !global)
+    return anchor->emitOpError(
+      "requires the pipeline host transfer buffer"), failure();
+  auto type = MemRefType::get({bytes.getInt()}, builder.getI8Type());
+  return memref::GetGlobalOp::create(
+    builder, anchor->getLoc(), type, kPipelineHostBufferName).getResult();
+}
+
 Operation* getEarliestUserWithinBlock(mlir::Value value) {
   auto users = value.getUsers();
 
@@ -129,6 +175,32 @@ LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprin
   return success();
 }
 
+FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(RewriterBase& rewriter,
+                                                          Location loc,
+                                                          mlir::Value source,
+                                                          RankedTensorType resultType) {
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || resultType.getRank() < 2
+      || sourceType.getRank() != resultType.getRank() + 1 || sourceType.getElementType() != resultType.getElementType()
+      || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getDimSize(0) != getStaticShapeElementCount(resultType.getShape().drop_back())
+      || sourceType.getDimSize(sourceType.getRank() - 1) != resultType.getDimSize(resultType.getRank() - 1)
+      || llvm::any_of(sourceType.getShape().slice(1, sourceType.getRank() - 2), [](int64_t dim) { return dim != 1; }))
+    return failure();
+
+  SmallVector<ReassociationIndices> collapse {{}, {sourceType.getRank() - 1}};
+  for (int64_t dim = 0; dim < sourceType.getRank() - 1; ++dim)
+    collapse.front().push_back(dim);
+  auto flatType = RankedTensorType::get(
+    {sourceType.getDimSize(0), sourceType.getDimSize(sourceType.getRank() - 1)}, resultType.getElementType());
+  mlir::Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, source, collapse);
+
+  SmallVector<ReassociationIndices> expand {{}, {resultType.getRank() - 1}};
+  for (int64_t dim = 0; dim < resultType.getRank() - 1; ++dim)
+    expand.front().push_back(dim);
+  return tensor::ExpandShapeOp::create(rewriter, loc, resultType, flat, expand).getResult();
+}
+
 static SmallVector<int64_t, 4> expandFlatElementIndex(int64_t flatIndex, ArrayRef<int64_t> shape) {
   SmallVector<int64_t, 4> indices(shape.size(), 0);
   for (int64_t dim = static_cast<int64_t>(shape.size()) - 1; dim >= 0; --dim) {

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -10,6 +10,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Support/LogicalResult.h"
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
@@ -23,6 +24,12 @@ namespace onnx_mlir {
 mlir::FailureOr<mlir::IntegerAttr>
 getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);
 
+mlir::LogicalResult materializePipelineHostBuffer(
+  mlir::func::FuncOp funcOp, mlir::RewriterBase &rewriter);
+
+mlir::FailureOr<mlir::Value> getPipelineHostBuffer(
+  mlir::OpBuilder &builder, mlir::Operation *anchor);
+
 template <class T>
 size_t rangeLength(const mlir::iterator_range<T> range) {
   return std::distance(range.begin(), range.end());
@@ -51,6 +58,11 @@ mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlu
                                                      llvm::ArrayRef<int64_t> flatSizes,
                                                      llvm::ArrayRef<int64_t> flatStrides);
 
+mlir::FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(mlir::RewriterBase& rewriter,
+                                                                mlir::Location loc,
+                                                                mlir::Value source,
+                                                                mlir::RankedTensorType resultType);
+
 mlir::FailureOr<mlir::SmallVector<int64_t, 4>>
 getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor,
                                       mlir::ShapedType sourceType,

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -42,12 +42,11 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
   if (!resultType || !resultType.hasStaticShape())
     return blueprint.emitOpError("fragment assembly lowering requires a static ranked tensor result");
 
-  std::optional<StringRef> modeAttr = blueprint.getMode();
   std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
   std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
   std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
   std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
-  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceSlotsAttr
+  if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceSlotsAttr
       || !sourceOffsetsAttr || !fragmentStridesAttr)
     return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
 
@@ -71,6 +70,16 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
                                               flatStrides)))
     return failure();
 
+  if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments) {
+    if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+      return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+    Value source = mapping.lookupOrDefault(blueprint.getInput());
+    auto reshaped = reshapeContiguousRowMajorFragments(
+      rewriter, blueprint.getLoc(), source, cast<RankedTensorType>(resultType));
+    if (failed(reshaped))
+      return blueprint.emitOpError("contiguous row-major fragment storage does not match its logical result"), failure();
+    return *reshaped;
+  }
   SmallVector<int64_t> hostStrides = computeRowMajorStrides(resultType.getShape());
   SmallVector<FragmentAssemblyCopy, 8> copies;
   for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
@@ -193,8 +202,7 @@ static bool isHostMaterializableHelperOp(Operation* op) {
   if (isa<arith::ConstantOp>(op) || op->hasTrait<OpTrait::ConstantLike>())
     return true;
   if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-    std::optional<StringRef> mode = blueprint.getMode();
-    return mode && *mode == "fragment_assembly";
+    return spatial::isFragmentAssembly(blueprint.getMode());
   }
   return isShapingOnlyOp(op) || isPureIndexComputationOp(op);
 }
@@ -281,8 +289,7 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
   }
   for (Operation& op : block.without_terminator()) {
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> modeAttr = blueprint.getMode();
-      if (modeAttr && *modeAttr == "fragment_assembly") {
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
         auto lowered = lowerFragmentAssemblyBlueprint(rewriter, blueprint, mapping);
         if (failed(lowered))
           return false;
@@ -338,20 +345,39 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     auto blockArg = computeOp.getInputArgument(inputIndex);
     if (!blockArg)
       return computeOp.emitOpError("expected compute input block arguments during lowering");
-    auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(input.getDefiningOp());
+    auto channelReceive = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(
+      input.getDefiningOp());
+    auto hostWaitLoad = dyn_cast_or_null<spatial::SpatHostWaitLoadOp>(
+      input.getDefiningOp());
+    Operation *receiveOp = channelReceive
+      ? channelReceive.getOperation() : hostWaitLoad.getOperation();
     if (receiveOp && !blockArg->use_empty()) {
       rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg));
       auto outputType = cast<ShapedType>(blockArg->getType());
-      auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
+      auto outputBuffer = createEmptyTensorFromShaped(
+        rewriter, receiveOp->getLoc(), outputType);
       auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
       if (failed(sizeAttr))
         return failure();
-      Value received =
-        PimReceiveOp::create(
-          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer,
-          arith::ConstantIndexOp::create(rewriter, receiveOp.getLoc(), 0),
-          *sizeAttr, receiveOp.getSourceCoreId())
+      Value zero = arith::ConstantIndexOp::create(
+        rewriter, receiveOp->getLoc(), 0);
+      Value received;
+      if (hostWaitLoad) {
+        auto hostBuffer = getPipelineHostBuffer(rewriter, hostWaitLoad);
+        if (failed(hostBuffer))
+          return failure();
+        PimWaitOp::create(
+          rewriter, receiveOp->getLoc(), hostWaitLoad.getEventRegister(),
+          rewriter.getI32IntegerAttr(1));
+        received = PimMemCopyHostToDevOp::create(
+          rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero,
+          hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr)
           .getOutput();
+      } else {
+        received = PimReceiveOp::create(
+          rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,
+          zero, *sizeAttr, channelReceive.getSourceCoreId()).getOutput();
+      }
       blockArg->replaceAllUsesWith(received);
       markOpToRemove(receiveOp);
       continue;
@@ -376,7 +402,8 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
     if (rangeLength(resultUses) == 1) {
       OpOperand& resultUse = *resultUses.begin();
       Operation* resultUser = resultUse.getOwner();
-      if (isa<spatial::SpatChannelSendOp>(resultUser))
+      if (isa<spatial::SpatChannelSendOp,
+              spatial::SpatHostStoreSyncOp>(resultUser))
         continue;
     }
 

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -22,8 +22,7 @@ struct LowerFragmentAssemblyBlueprintPattern
   LogicalResult matchAndRewrite(spatial::SpatBlueprintOp op,
                                 OpAdaptor adaptor,
                                 ConversionPatternRewriter& rewriter) const override {
-    std::optional<StringRef> modeAttr = op.getMode();
-    if (!modeAttr || *modeAttr != "fragment_assembly")
+    if (!spatial::isFragmentAssembly(op.getMode()))
       return failure();
 
     auto resultType = dyn_cast<ShapedType>(op.getOutput().getType());
@@ -49,6 +48,16 @@ struct LowerFragmentAssemblyBlueprintPattern
           op, rank, fragmentOperands.size(), operandIndices, sourceOffsets, flatOffsets, flatSizes, flatStrides)))
       return failure();
 
+    if (op.getIndexMap() == spatial::kContiguousRowMajorFragments) {
+      if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(op))
+        return op.emitOpError("contiguous row-major fragment physical source order or storage is not canonical");
+      auto reshaped = reshapeContiguousRowMajorFragments(
+        rewriter, op.getLoc(), adaptor.getInput(), cast<RankedTensorType>(resultType));
+      if (failed(reshaped))
+        return op.emitOpError("contiguous row-major fragment storage does not match its logical result");
+      rewriter.replaceOp(op, *reshaped);
+      return success();
+    }
     Value currentOutput =
       tensor::EmptyOp::create(rewriter, op.getLoc(), resultType.getShape(), resultType.getElementType()).getResult();
     for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {

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

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

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -158,8 +158,7 @@ analyzeTopLevelFragmentAssemblyUses(Value value) {
     auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
     if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
       return failure();
-    std::optional<StringRef> mode = blueprint.getMode();
-    if (!mode || *mode != "fragment_assembly")
+    if (!spatial::isFragmentAssembly(blueprint.getMode()))
       return failure();
     if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
       return failure();
@@ -819,8 +818,7 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
     }
 
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> mode = blueprint.getMode();
-      if (mode && *mode == "fragment_assembly") {
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
         markOpToRemove(blueprint.getOperation());
         for (Value operand : blueprint->getOperands())
           markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);
@@ -861,6 +859,10 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
       markOpToRemove(receiveOp);
       return;
     }
+    if (auto receiveOp = dyn_cast<spatial::SpatHostWaitLoadOp>(op)) {
+      markOpToRemove(receiveOp);
+      return;
+    }
   };
 
   SmallVector<Value> originalOperands(returnOp.getOperands().begin(), returnOp.getOperands().end());

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

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

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

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

src/PIM/Dialect/Pim/CMakeLists.txt

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

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

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

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

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

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

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

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

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

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

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

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

@@ -20,11 +20,11 @@
 #include "Common/Support/Diagnostics.hpp"
 #include "Compiler/PimCodeGen.hpp"
 #include "Dialect/Pim/PimOps.hpp"
-#include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
-#include "Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "Dialect/Pim/Passes/Transforms/Bufferization/Common.hpp"
+#include "Dialect/Pim/Passes/Transforms/Bufferization/ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Accelerators/PIM/Passes/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 
 using namespace mlir;
@@ -33,6 +33,9 @@ using namespace pim;
 
 namespace onnx_mlir {
 
+static void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp);
+static FailureOr<func::FuncOp> requirePimEntryFunc(ModuleOp moduleOp, StringRef phase);
+
 namespace {
 
 struct MemRefCopyWorkItem {
@@ -333,22 +336,6 @@ static LogicalResult verifyPimCopyEndpoints(Operation* copy,
   return success(valid);
 }
 
-struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass)
-  StringRef getArgument() const override { return "bufferize-pim"; }
-  StringRef getDescription() const override { return "Bufferize PIM and Spatial ops."; }
-
-  PimBufferizationPass() = default;
-  PimBufferizationPass(const PimBufferizationPass& pass) {}
-
-  void runOnOperation() final;
-
-private:
-  void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const;
-  LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) const;
-  LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) const;
-};
-
 static void materializeWritableConstantDestinations(func::FuncOp funcOp) {
   SmallVector<OpOperand*> constantBackedRoots;
   llvm::SmallPtrSet<OpOperand*, 8> seenRoots;
@@ -387,14 +374,22 @@ static void materializeWritableConstantDestinations(func::FuncOp funcOp) {
   }
 }
 
+static bufferization::OneShotBufferizationOptions makePimBufferizationOptions() {
+  bufferization::OneShotBufferizationOptions options;
+  options.allowUnknownOps = true;
+  options.bufferizeFunctionBoundaries = true;
+  options.setFunctionBoundaryTypeConversion(bufferization::LayoutMapOption::IdentityLayoutMap);
+  return options;
+}
+
 static LogicalResult verifyPimCoresNeedNoTensorCopies(
-  ModuleOp module, const bufferization::OneShotBufferizationOptions& baseOptions) {
+    ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions& baseOptions) {
   static constexpr StringLiteral kExistingAlloc = "raptor.existing_core_alloc";
-  OwningOpRef<ModuleOp> clone = module.clone();
+  OwningOpRef<ModuleOp> clone = moduleOp.clone();
   clone->walk([&](bufferization::AllocTensorOp alloc) {
     if (alloc->getParentOfType<pim::PimCoreOp>()
         || alloc->getParentOfType<pim::PimCoreBatchOp>())
-      alloc->setAttr(kExistingAlloc, UnitAttr::get(module.getContext()));
+      alloc->setAttr(kExistingAlloc, UnitAttr::get(moduleOp.getContext()));
   });
 
   auto options = baseOptions;
@@ -407,7 +402,7 @@ static LogicalResult verifyPimCoresNeedNoTensorCopies(
 
   bufferization::BufferizationState state;
   if (failed(bufferization::insertTensorCopies(*clone, options, state))) {
-    module.emitError("official one-shot analysis failed while verifying PIM core copy freedom");
+    moduleOp.emitError("official one-shot analysis failed while verifying PIM core copy freedom");
     return failure();
   }
 
@@ -423,29 +418,22 @@ static LogicalResult verifyPimCoresNeedNoTensorCopies(
       op->emitOpError("official one-shot bufferization requires a tensor copy inside a PIM core");
     });
   });
-  diagnostics.emitSuppressedSummary(module, "required PIM core tensor copies");
+  diagnostics.emitSuppressedSummary(moduleOp, "required PIM core tensor copies");
   return success(!diagnostics.hasFailure());
 }
 
-} // namespace
-
-void PimBufferizationPass::runOnOperation() {
-  auto moduleOp = getOperation();
-  auto funcOp = *getPimEntryFunc(moduleOp);
-
-  bufferization::OneShotBufferizationOptions options;
-  options.allowUnknownOps = true;
-  options.bufferizeFunctionBoundaries = true;
-  options.setFunctionBoundaryTypeConversion(bufferization::LayoutMapOption::IdentityLayoutMap);
-
+static LogicalResult preparePimBufferization(
+    ModuleOp moduleOp, func::FuncOp funcOp, bool verifyCopyFreedom) {
   materializeWritableConstantDestinations(funcOp);
-  if (failed(verifyPimCoresNeedNoTensorCopies(moduleOp, options))) {
-    signalPassFailure();
-    return;
-  }
+  if (verifyCopyFreedom)
+    return verifyPimCoresNeedNoTensorCopies(moduleOp, makePimBufferizationOptions());
+  return success();
+}
 
+static LogicalResult runOneShotPimBufferization(
+    ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions& options) {
   auto hostOptions = options;
-  hostOptions.opFilter.denyOperation([](Operation *op) {
+  hostOptions.opFilter.denyOperation([](Operation* op) {
     return op->getParentOfType<pim::PimCoreOp>()
         || op->getParentOfType<pim::PimCoreBatchOp>();
   });
@@ -453,84 +441,14 @@ void PimBufferizationPass::runOnOperation() {
   if (failed(bufferization::insertTensorCopies(moduleOp, hostOptions, state))
       || failed(bufferization::bufferizeModuleOp(moduleOp, options, state))) {
     moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
-    signalPassFailure();
-    return;
+    return failure();
   }
-
-  forwardSingleConsumerReceiveCopies(funcOp);
-  forwardSingleConsumerContiguousInputCopies(funcOp);
-  forwardSingleConsumerPimOutputCopies(funcOp);
-
-  MLIRContext* ctx = moduleOp.getContext();
-  PatternRewriter rewriter(ctx);
-
-  SmallVector<MemRefCopyWorkItem> copyWorklist;
-  llvm::SmallPtrSet<Operation*, 16> seenCopyOps;
-  auto addCopyOp = [&](memref::CopyOp copyOp, const StaticValueKnowledge& knowledge) {
-    if (seenCopyOps.insert(copyOp.getOperation()).second)
-      copyWorklist.push_back({copyOp, knowledge});
-  };
-
-  moduleOp.walk([&](pim::PimCoreOp coreOp) {
-    StaticValueKnowledge knowledge = seedCoreKnowledge(coreOp);
-    (void) walkPimCoreBlockStructurally(
-      coreOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
-        if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
-          addCopyOp(copyOp, opKnowledge);
-        return success();
-      });
-  });
-  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
-    for (unsigned lane = 0; lane < coreBatchOp.getLaneCount(); ++lane) {
-      StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane);
-      (void) walkPimCoreBlockStructurally(
-        coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
-          if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
-            addCopyOp(copyOp, opKnowledge);
-          return success();
-      });
-    }
-  });
-
-  bool hasFailed = false;
-  Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0);
-  for (const MemRefCopyWorkItem& workItem : copyWorklist) {
-    memref::CopyOp copyOp = workItem.copyOp;
-    rewriter.setInsertionPoint(copyOp);
-    if (failed(lowerMemRefCopyToPimCopy(copyOp, zeroOffset, rewriter, workItem.knowledge)))
-      hasFailed = true;
-  }
-  if (hasFailed) {
-    signalPassFailure();
-    return;
-  }
-
-  RewritePatternSet contiguityPatterns(ctx);
-  populatePimContiguityNormalizationPatterns(contiguityPatterns);
-
-  GreedyRewriteConfig contiguityConfig;
-  contiguityConfig.enableFolding(false);
-  if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) {
-    moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization");
-    signalPassFailure();
-    return;
-  }
-  if (failed(verifyContiguousRuntimeOperands(moduleOp))) {
-    signalPassFailure();
-    return;
-  }
-  if (failed(verifyPimCopyAddressSpaces(moduleOp))) {
-    signalPassFailure();
-    return;
-  }
-
-  annotateWeightsMemrefs(moduleOp, funcOp);
-
-  // Dump to file for debug
-  dumpModule(moduleOp, "pim1_buff");
+  return success();
 }
 
-void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const {
+} // namespace
+
+static void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) {
   auto markWeights = [&](Operation* op) {
     walkPimMvmVmmWeightUses(op, [&](OpOperand& weightUse) {
       Value weight = weightUse.get();
@@ -548,7 +466,7 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO
   funcOp.walk([&](PimCoreBatchOp coreBatchOp) { markWeights(coreBatchOp); });
 }
 
-LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const {
+static LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) {
   bool hasFailure = false;
 
   auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
@@ -640,7 +558,7 @@ LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp mod
   return success();
 }
 
-LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp) const {
+static LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) {
   size_t failureCount = 0;
   auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
     (void) walkPimCoreBlockStructurally(
@@ -675,6 +593,211 @@ LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp
   return success(failureCount == 0);
 }
 
-std::unique_ptr<Pass> createPimBufferizationPass() { return std::make_unique<PimBufferizationPass>(); }
+static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp) {
+  forwardSingleConsumerReceiveCopies(funcOp);
+  forwardSingleConsumerContiguousInputCopies(funcOp);
+  forwardSingleConsumerPimOutputCopies(funcOp);
+
+  MLIRContext* ctx = moduleOp.getContext();
+  PatternRewriter rewriter(ctx);
+
+  SmallVector<MemRefCopyWorkItem> copyWorklist;
+  llvm::SmallPtrSet<Operation*, 16> seenCopyOps;
+  auto addCopyOp = [&](memref::CopyOp copyOp, const StaticValueKnowledge& knowledge) {
+    if (seenCopyOps.insert(copyOp.getOperation()).second)
+      copyWorklist.push_back({copyOp, knowledge});
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) {
+    StaticValueKnowledge knowledge = seedCoreKnowledge(coreOp);
+    (void) walkPimCoreBlockStructurally(
+      coreOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
+        if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
+          addCopyOp(copyOp, opKnowledge);
+        return success();
+      });
+  });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    for (unsigned lane = 0; lane < coreBatchOp.getLaneCount(); ++lane) {
+      StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane);
+      (void) walkPimCoreBlockStructurally(
+        coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
+          if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
+            addCopyOp(copyOp, opKnowledge);
+          return success();
+        });
+    }
+  });
+
+  bool hasFailed = false;
+  Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0);
+  for (const MemRefCopyWorkItem& workItem : copyWorklist) {
+    memref::CopyOp copyOp = workItem.copyOp;
+    rewriter.setInsertionPoint(copyOp);
+    if (failed(lowerMemRefCopyToPimCopy(copyOp, zeroOffset, rewriter, workItem.knowledge)))
+      hasFailed = true;
+  }
+  if (hasFailed)
+    return failure();
+
+  RewritePatternSet contiguityPatterns(ctx);
+  populatePimContiguityNormalizationPatterns(contiguityPatterns);
+
+  GreedyRewriteConfig contiguityConfig;
+  contiguityConfig.enableFolding(false);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) {
+    moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization");
+    return failure();
+  }
+  annotateWeightsMemrefs(moduleOp, funcOp);
+  dumpModule(moduleOp, "pim1_buff");
+  return success();
+}
+
+static FailureOr<func::FuncOp> requirePimEntryFunc(ModuleOp moduleOp, StringRef phase) {
+  auto entryFunc = getPimEntryFunc(moduleOp);
+  if (failed(entryFunc)) {
+    moduleOp.emitError("failed to locate the PIM entry function during ") << phase;
+    return failure();
+  }
+  return *entryFunc;
+}
+
+namespace {
+
+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.";
+  }
+
+  void runOnOperation() final {
+    ModuleOp moduleOp = getOperation();
+    auto funcOp = requirePimEntryFunc(moduleOp, "PIM bufferization preparation");
+    if (failed(funcOp)) {
+      signalPassFailure();
+      return;
+    }
+    if (failed(preparePimBufferization(moduleOp, *funcOp, verifyCopyFreedom)))
+      signalPassFailure();
+  }
+
+private:
+  bool verifyCopyFreedom;
+};
+
+struct PimOneShotBufferizationPass
+    : PassWrapper<PimOneShotBufferizationPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimOneShotBufferizationPass)
+
+  StringRef getArgument() const override { return "pim-one-shot-bufferization"; }
+  StringRef getDescription() const override {
+    return "Run one-shot bufferization for PIM and Spatial tensors.";
+  }
+
+  void runOnOperation() final {
+    if (failed(runOneShotPimBufferization(getOperation(), makePimBufferizationOptions())))
+      signalPassFailure();
+  }
+};
+
+struct PimMemoryNormalizationPass
+    : PassWrapper<PimMemoryNormalizationPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimMemoryNormalizationPass)
+
+  StringRef getArgument() const override { return "pim-memory-normalization"; }
+  StringRef getDescription() const override {
+    return "Normalize PIM memory copies and verify addressable operands.";
+  }
+
+  void runOnOperation() final {
+    ModuleOp moduleOp = getOperation();
+    auto funcOp = requirePimEntryFunc(moduleOp, "PIM memory normalization");
+    if (failed(funcOp)) {
+      signalPassFailure();
+      return;
+    }
+    if (failed(normalizePimMemory(moduleOp, *funcOp)))
+      signalPassFailure();
+  }
+};
+
+static LogicalResult verifyNoTensorValues(ModuleOp moduleOp) {
+  size_t failureCount = 0;
+  moduleOp.walk([&](Operation* op) {
+    if (failureCount >= 8)
+      return;
+    if (op->getDialect()->getNamespace() == "tensor") {
+      op->emitOpError("tensor operation remains after PIM bufferization");
+      ++failureCount;
+      return;
+    }
+    for (Value value : op->getOperands()) {
+      if (isa<TensorType>(value.getType())) {
+        op->emitOpError("tensor operand remains after PIM bufferization");
+        ++failureCount;
+        return;
+      }
+    }
+    for (Value value : op->getResults()) {
+      if (isa<TensorType>(value.getType())) {
+        op->emitOpError("tensor result remains after PIM bufferization");
+        ++failureCount;
+        return;
+      }
+    }
+  });
+  if (failureCount != 0)
+    moduleOp.emitError() << "found " << failureCount
+                         << " tensor value(s) after PIM bufferization"
+                         << (failureCount == 8 ? " (first 8 reported)" : "");
+  return success(failureCount == 0);
+}
+
+struct PimBufferizationVerificationPass
+    : PassWrapper<PimBufferizationVerificationPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationVerificationPass)
+
+  StringRef getArgument() const override { return "pim-bufferization-verification"; }
+  StringRef getDescription() const override {
+    return "Verify tensor elimination, contiguity, and PIM copy address spaces.";
+  }
+
+  void runOnOperation() final {
+    ModuleOp moduleOp = getOperation();
+    if (failed(verifyNoTensorValues(moduleOp))
+        || failed(verifyContiguousRuntimeOperands(moduleOp))
+        || failed(verifyPimCopyAddressSpaces(moduleOp)))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+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>();
+}
+
+std::unique_ptr<Pass> createPimMemoryNormalizationPass() {
+  return std::make_unique<PimMemoryNormalizationPass>();
+}
+
+std::unique_ptr<Pass> createPimBufferizationVerificationPass() {
+  return std::make_unique<PimBufferizationVerificationPass>();
+}
 
 } // namespace onnx_mlir

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

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

src/PIM/Dialect/Pim/Passes/Transforms/HostConstantFolding/Patterns/Constant.cpp

@@ -532,54 +532,74 @@ struct FoldConstantMemCpPattern final : OpRewritePattern<pim::PimMemCopyOp> {
   }
 };
 
-static bool isOne(Attribute value) {
-  if (auto floatValue = dyn_cast<FloatAttr>(value))
-    return floatValue.getValue().isExactlyValue(1.0);
-  if (auto integerValue = dyn_cast<IntegerAttr>(value))
-    return integerValue.getValue() == 1;
-  return false;
+enum class MultiplicationConstant { Other, Zero, One };
+
+static MultiplicationConstant classifyMultiplicationConstant(Attribute value) {
+  if (auto floatValue = dyn_cast<FloatAttr>(value)) {
+    const APFloat& number = floatValue.getValue();
+    if (number.isZero() && !number.isNegative())
+      return MultiplicationConstant::Zero;
+    if (number.isExactlyValue(1.0))
+      return MultiplicationConstant::One;
+  }
+  if (auto integerValue = dyn_cast<IntegerAttr>(value)) {
+    if (integerValue.getValue().isZero())
+      return MultiplicationConstant::Zero;
+    if (integerValue.getValue() == 1)
+      return MultiplicationConstant::One;
+  }
+  return MultiplicationConstant::Other;
 }
 
-static bool isAllOneHostCopy(pim::PimMemCopyHostToDevOp copyOp, ModuleOp moduleOp, MemRefType copiedType) {
+static MultiplicationConstant classifyUniformHostCopy(
+  pim::PimMemCopyHostToDevOp copyOp, ModuleOp moduleOp, MemRefType copiedType) {
   auto targetOffset = resolveIndexValue(copyOp.getDeviceTargetOffset());
   auto sourceOffset = resolveIndexValue(copyOp.getHostSourceOffset());
   if (failed(targetOffset) || failed(sourceOffset) || *targetOffset != 0)
-    return false;
+    return MultiplicationConstant::Other;
 
   Type elementType = copiedType.getElementType();
   if (!elementType.isIntOrFloat())
-    return false;
+    return MultiplicationConstant::Other;
   unsigned bitWidth = elementType.getIntOrFloatBitWidth();
   if (bitWidth == 0 || bitWidth % 8 != 0)
-    return false;
+    return MultiplicationConstant::Other;
 
   int64_t elementBytes = bitWidth / 8;
   int64_t copiedElements = copiedType.getNumElements();
   if (*sourceOffset % elementBytes != 0 || copyOp.getSize() != copiedElements * elementBytes)
-    return false;
+    return MultiplicationConstant::Other;
 
   auto source = getDenseGlobalValue(moduleOp, copyOp.getHostSource());
   if (failed(source) || source->getElementType() != elementType)
-    return false;
+    return MultiplicationConstant::Other;
 
   int64_t firstElement = *sourceOffset / elementBytes;
   int64_t endElement = firstElement + copiedElements;
   if (firstElement < 0 || endElement > source->getNumElements())
-    return false;
+    return MultiplicationConstant::Other;
   if (source->isSplat())
-    return isOne(source->getSplatValue<Attribute>());
+    return classifyMultiplicationConstant(source->getSplatValue<Attribute>());
 
+  MultiplicationConstant classification = MultiplicationConstant::Other;
   int64_t index = 0;
   for (Attribute value : source->getValues<Attribute>()) {
-    if (index >= firstElement && index < endElement && !isOne(value))
-      return false;
+    if (index >= firstElement && index < endElement) {
+      MultiplicationConstant current = classifyMultiplicationConstant(value);
+      if (current == MultiplicationConstant::Other)
+        return current;
+      if (classification == MultiplicationConstant::Other)
+        classification = current;
+      else if (classification != current)
+        return MultiplicationConstant::Other;
+    }
     if (++index >= endElement)
       break;
   }
-  return true;
+  return classification;
 }
 
-struct FoldMultiplyByOnePattern final : OpRewritePattern<pim::PimVVMulOp> {
+struct FoldMultiplyByConstantPattern final : OpRewritePattern<pim::PimVVMulOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(pim::PimVVMulOp mulOp, PatternRewriter& rewriter) const override {
@@ -605,14 +625,19 @@ struct FoldMultiplyByOnePattern final : OpRewritePattern<pim::PimVVMulOp> {
         copyOp = candidate;
       }
       auto maskType = dyn_cast<MemRefType>(mask.getType());
-      if (!copyOp || !copyOp.use_empty() || !maskType || !isAllOneHostCopy(copyOp, moduleOp, maskType))
+      if (!copyOp || !copyOp.use_empty() || !maskType)
+        continue;
+      MultiplicationConstant constant = classifyUniformHostCopy(copyOp, moduleOp, maskType);
+      if (constant == MultiplicationConstant::Other)
         continue;
 
       auto outputAlloc = mulOp.getOutputBuffer().getDefiningOp<memref::AllocOp>();
-      rewriter.replaceOp(mulOp, input);
-      rewriter.eraseOp(copyOp);
-      if (maskAlloc.use_empty())
-        rewriter.eraseOp(maskAlloc);
+      rewriter.replaceOp(mulOp, constant == MultiplicationConstant::One ? input : mask);
+      if (constant == MultiplicationConstant::One) {
+        rewriter.eraseOp(copyOp);
+        if (maskAlloc.use_empty())
+          rewriter.eraseOp(maskAlloc);
+      }
       if (outputAlloc && outputAlloc.use_empty())
         rewriter.eraseOp(outputAlloc);
       return success();
@@ -629,7 +654,7 @@ void populateConstantFoldingConstantPatterns(RewritePatternSet& patterns) {
                FoldConstantCoreMapPattern,
                FoldConstantHostCopyPattern,
                FoldConstantMemCpPattern,
-               FoldMultiplyByOnePattern>(patterns.getContext());
+               FoldMultiplyByConstantPattern>(patterns.getContext());
 }
 
 } // namespace onnx_mlir

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

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

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

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

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

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

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

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

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

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

src/PIM/Dialect/Pim/Pim.td

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

src/PIM/Dialect/Pim/PimOpsVerify.cpp

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

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -1,38 +1,50 @@
 add_onnx_mlir_dialect(Spatial spat)
 add_onnx_mlir_dialect_doc(spat Spatial.td)
 
+set(LLVM_TARGET_DEFINITIONS Spatial.td)
+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)
+
+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)
+
 add_pim_library(SpatialOps
   SpatialOps.cpp
   SpatialOpsAsm.cpp
   SpatialOpsVerify.cpp
   SpatialOpsCanonicalization.cpp
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
-  Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
-  Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
-  Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
-  Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
-  Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp
-  Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp
-  Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp
-  Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
-  Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
-  Transforms/MergeComputeNodes/DeferredResultRealization.cpp
-  Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
-  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
-  Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
-  Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
-  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
-  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
-  Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
-  Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
-  Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
-  Transforms/TrivialGraphComputeMergePass.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredResultRealization.cpp
+  Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
+  Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
+  Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
+  Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp
+  Passes/Transforms/TrivialGraphComputeMergePass.cpp
 
   EXCLUDE_FROM_OM_LIBS
 
   DEPENDS
   OMONNXIncGen
   OMSpatialIncGen
+  OMSpatialEnumsIncGen
+  OMSpatialLayoutInterfaceIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
@@ -40,6 +52,5 @@ add_pim_library(SpatialOps
   MLIRBufferizationTransforms
   OMMlirDialects
   OMONNXOps
-  OMPimCompilerOptions
   PimOps
 )

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

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