Cose belle

remove unsupported tests

AGENTS.md

@@ -1,92 +1,212 @@
-- Always read the full README.md before doing anything.
-- Build commands:
-    - `cmake --build ./build_release`
-    - `cmake --build ./build_debug`
-- Never use `ninja` directly: it bypasses cmake's configuration and invalidates the build cache.
-- Always tries the release version build first and ask before building with the debug version
+* Always read the full README.md before doing anything
+* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes.
+* Build commands:
+  * `cmake --build ./build_release`
+  * `cmake --build ./build_debug`
+* Never use `ninja` directly: it bypasses cmake's configuration and invalidates the build cache
+* Always try the release build first before building with the debug version
+* Use the debug build only when it is useful to obtain a clear stack trace with symbols, inspect names, place breakpoints, or test a small case interactively
+* The debug build is very slow, so use it only on small fast tests such as operation validations, not on network validations
+* Always prepend rtk to shell commands if missing and if rtk is available
+
+# Core engineering philosophy
+
+* Clean architecture matters as much as making the immediate test pass
+* Prefer fixes that preserve clear ownership boundaries, explicit invariants, and simple dataflow
+* Do not stack compensating fixes on top of earlier mistakes. If the current approach is becoming messy, stop and explain why
+* A correct fix should usually make the responsible producer, resolver, verifier, or lowering own the behavior directly
+* Avoid late repair passes, defensive cleanup, or broad rewrites when a cleaner owner-side fix is possible
+* Do not hide an upstream modeling bug by normalizing it later in the pipeline. Fix the producer when the producer owns the invariant
+* Prefer patterns/rewrites for local IR canonicalization. Use module walks only when pass-level structural analysis genuinely requires them
+* Prefer compact, structured designs over long case-by-case implementations
+
+# Think before coding
+
+* State assumptions explicitly before implementing when they affect the design
+* If multiple interpretations exist, present them instead of silently choosing one
+* If a simpler approach exists, say so and prefer it unless there is a clear reason not to
+* If something is unclear, stop, name what is confusing, and ask
+* If the requested or obvious approach would make the architecture worse, push back and propose a cleaner alternative
 
 # Code changes
 
-- Keep changes minimal and localized to the relevant parts of the code.
-- Preserve the existing naming conventions and coding style used in the surrounding code.
-- Keep code easy to read, well organized, and suitable for future extensibility. A function must not be longer than
-  200/250 lines for readability and cognitive complexity.
-- Prefer clear naming and structure over comments. Add comments only when they materially improve clarity.
-- Do not rename symbols, move files, or restructure modules unless that is necessary for the requested change.
+* Keep changes minimal and localized to the relevant parts of the code
+* Preserve the existing naming conventions and coding style used in the surrounding code
+* Keep code easy to read, well organized, and suitable for future extensibility
+* A function must not exceed roughly 200/250 lines. If a change pushes a function beyond that, extract focused helpers
+* Prefer clear naming and structure over comments. Add comments only when they materially improve clarity
+* Do not rename symbols, move files, or restructure modules unless that is necessary for the requested change
+* Avoid duplicate ad-hoc logic. If the same concept appears in multiple places, consider whether it deserves a shared helper/API
+* When adding a helper or API, ask:
+  * Could this be useful to another component now
+  * Is another component already implementing the same idea differently
+  * Is this likely to be needed by a future adjacent component
+  * What is the narrowest useful abstraction
+  * What is the correct ownership level for this API
+* If a shared API is justified, place it at the lowest clean layer that can be used by all relevant consumers without creating dependency cycles or leaking policy across layers
+* If an existing component should use a newly introduced shared API, refactor that component in the same patch when doing so is directly related and reduces duplication
+* Do not create broad frameworks just because a helper might someday be useful. Shared APIs should encode a real reusable concept, not speculative generality
+* If the reusable abstraction is plausible but not clearly needed yet, keep the code local and mention the possible future extraction separately
+
+# Avoid case-listing designs
+
+* Avoid solving problems with large chains of `if`/`else`, switches, or repeated special cases that enumerate every possible situation
+* Long case listings tend to overfit the current tests, grow the codebase, and hide the underlying abstraction
+* When you see a growing list of special cases, stop and look for the shared concept, data model, interface, or normalization step that would make the cases collapse
+* Prefer table-driven logic, traits/interfaces, small reusable predicates, structured dispatch, or producer-side normalization when they express the invariant more directly
+* A few explicit cases are fine when the domain is genuinely small and closed
+* If the list is likely to grow, refactor toward a cleaner and more compact design instead of adding another branch
+* When keeping a case list is the pragmatic choice, explain why the domain is closed or why a broader abstraction would be premature
+
+# Ownership and invariants
+
+Before implementing, identify the owner of the behavior:
+
+* A producer should emit IR/data that satisfies the contract of the next stage
+* A lowering should make representation changes explicit and semantically correct
+* A resolver should resolve existing structure without silently changing semantics
+* A verifier should reject invalid states with bounded, actionable diagnostics
+* Codegen should assume verified invariants and fail clearly if they are violated
+
+When fixing a bug:
+
+* State the invariant that was violated
+* State which component should own that invariant
+* Fix that component directly
+* Avoid fixes that merely mask the violation later in the pipeline
+* Add or preserve verification if the invariant is important enough to regress
+
+# Refactor and API policy
+
+You may propose or implement a refactor when:
+
+* the local fix would duplicate logic
+* the local fix would violate a layer boundary
+* the bug exists because responsibility is assigned to the wrong component
+* multiple components already implement ad-hoc variants of the same concept
+* a shared helper/API would make the code smaller, clearer, and easier to maintain
+* existing callers can be migrated cleanly without broad churn
+* the current implementation is turning into a long list of special cases instead of a structured solution
+
+When proposing or implementing a refactor:
+
+* Explain what responsibility is being moved or shared
+* Justify why the new location is the right ownership level
+* Keep the API narrow and named after the concept or invariant it represents
+* Migrate directly related existing users when that improves compactness and consistency
+* Separate changes required for correctness from optional cleanup
+* Avoid unrelated renames, formatting changes, or module moves
+* Do not expand a justified refactor beyond directly related callers
+
+Do not refactor when:
+
+* the issue is truly local and a local fix is clearer
+* the abstraction would have only one user and no clear adjacent use
+* the abstraction would mix policies from different layers
+* the refactor would affect unrelated behavior
+* the refactor is mainly aesthetic
 
 # Working style
 
-- Infer style and conventions from the existing code before introducing new patterns.
-- When several implementation options are possible, prefer the simplest one that fits the current architecture and
-  minimizes churn.
-- Avoid broad refactors unless I explicitly ask for them.
+* Infer style and conventions from the existing code before introducing new patterns
+* When several implementation options are possible, prefer the simplest one that fits the current architecture and minimizes churn
+* Push back when the requested or obvious fix would make the architecture worse
+* If a cleaner fix requires a small refactor or shared helper/API, propose it explicitly and justify it
+* Avoid broad refactors unless explicitly requested or clearly necessary for correctness and maintainability
+* When tests fail, bucket failures by likely root cause and separate patch-related failures from pre-existing or out-of-scope failures
 
-# Responses
+# Simplicity first
 
-- When showing code in chat, make it easy to copy-paste into the codebase.
-- Keep outputs focused on the changed parts.
-- At the end of the response, briefly list any bad practices, mistakes, or cleaner alternatives you noticed, separate
-  from the main solution.
+* Minimum code that solves the problem cleanly. Nothing speculative
+* No features beyond what was asked
+* No error handling for impossible scenarios
+* If you write 200 lines and it could be 50, rewrite it
+* Ask: “Would a senior engineer say this is overcomplicated?” If yes, simplify
+* Prefer direct, explicit code over generic machinery unless the generic machinery clearly reduces duplication and preserves boundaries
 
-# Guidelines
+# Fallbacks and defaults
 
-## 1. Think Before Coding
+* Avoid silent fallback behavior when the semantic category is unknown
+* Do not treat “unknown” as “safe” unless the codebase already defines that convention
+* If a value cannot be classified, either preserve the existing behavior deliberately or fail with a clear diagnostic
+* When adding a fallback, state why it is semantically valid and what invariant makes it safe
 
-**Don't assume. Don't hide confusion. Surface tradeoffs.**
+# Surgical changes
 
-Before implementing:
+* Touch only what you must
+* Clean up only the mess introduced by your own change
+* Do not “improve” adjacent code, comments, or formatting
+* Match existing style, even if you would personally do it differently
+* If you notice unrelated dead code, bad abstractions, or fragile design, mention it separately. Do not delete or rewrite it unless asked
+* When your changes create orphans, remove imports, variables, functions, or files made unused by your change
+* Every changed line should trace directly to the requested fix, a required cleanup, or a justified reuse/refactor decision
 
-- State your assumptions explicitly. If uncertain, ask.
-- If multiple interpretations exist, present them - don't pick silently.
-- If a simpler approach exists, say so. Push back when warranted.
-- If something is unclear, stop. Name what's confusing. Ask.
+# Diagnostics and verification
 
-## 2. Simplicity First
+* Use existing bounded diagnostic mechanisms for pass-level verification or codegen failures
+* Do not emit unbounded repeated diagnostics from loops or parallel workers
+* Diagnostics should identify the violated invariant and the relevant value/op when useful
+* Verifiers should reject invalid states, not repair them
+* Codegen should not compensate for invalid IR/data unless codegen is the owner of that invariant
+* Do not make failing tests pass by weakening verifiers, assertions, or diagnostics unless the check itself is proven wrong
+* If a check is too strict, explain the valid case it rejects and update the invariant accordingly
+* Prefer fixing invalid IR/data producers over relaxing consumers
+* If adding diagnostics only for debugging, remove them or cap them before finalizing
 
-**Minimum code that solves the problem. Nothing speculative.**
+# Temporary debugging code
 
-- No features beyond what was asked.
-- No error handling for impossible scenarios.
-- If you write 200 lines and it could be 50, rewrite it.
+* Temporary diagnostics, dumps, assertions, and debug-only helpers must be removed or intentionally converted into bounded permanent diagnostics before finalizing
+* If debug instrumentation remains, explain why it is useful as permanent infrastructure
+* Do not leave noisy validation output behind
 
-Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify.
+# Performance awareness
 
-## 3. Surgical Changes
+* Avoid algorithmic regressions in compiler passes, especially repeated full-module walks, repeated expensive analyses, or per-op recomputation inside nested loops
+* If a change adds a walk, cache, analysis, or structural traversal, justify why it is needed
+* For hot paths, prefer preserving existing asymptotic behavior unless a better structure is part of the requested change
+* If performance may change, mention the expected impact and suggest a targeted timing check
 
-**Touch only what you must. Clean up only your own mess.**
-
-When editing existing code:
-
-- Don't "improve" adjacent code, comments, or formatting.
-- Don't refactor things that aren't broken.
-- Match existing style, even if you'd do it differently.
-- If you notice unrelated dead code, mention it - don't delete it.
-
-When your changes create orphans:
-
-- Remove imports/variables/functions that YOUR changes made unused.
-- Don't remove pre-existing dead code unless asked, but mention it.
-
-The test: Every changed line should trace directly to the user's request.
-
-## 4. Goal-Driven Execution
-
-**Define success criteria. Loop until verified.**
-
-Transform tasks into verifiable goals:
-
-- "Add validation" → "Write tests for invalid inputs, then make them pass"
-- "Fix the bug" → "Write a test that reproduces it, then make it pass"
-- "Refactor X" → "Ensure tests pass before and after"
+# Goal-driven execution
 
 For multi-step tasks, state a brief plan:
 
-```
 1. [Step] → verify: [check]
 2. [Step] → verify: [check]
 3. [Step] → verify: [check]
-```
 
-Strong success criteria let you loop independently. Weak criteria ("make it work") require constant clarification.
+Define success criteria before implementing:
 
----
+* For bug fixes, success means reproducing or identifying the failure, fixing the responsible owner, and verifying the targeted case
+* For refactors, success means preserving behavior while making ownership, reuse, or structure cleaner
+* For validation changes, success means checking both valid and invalid cases when applicable
+
+Transform tasks into verifiable goals:
+
+* “Fix the bug” → identify the invariant, reproduce the failure, fix the owner, verify the targeted case
+* “Add validation” → write or identify tests for invalid inputs, then make them pass/fail as expected
+* “Refactor X” → preserve behavior before and after, then run relevant tests
+
+# Final self-review
+
+Before reporting completion, check:
+
+* Did I fix the owner of the invariant rather than masking the issue downstream
+* Did I avoid broad case lists and ad-hoc special handling
+* Did I introduce a helper/API only at the right ownership level
+* Did I migrate directly related duplicate logic when doing so improves compactness
+* Did I avoid weakening verifiers or assertions unnecessarily
+* Did I remove temporary debugging code or make it bounded and intentional
+* Did I avoid unrelated formatting, renames, or cleanup
+* Did I consider performance impact for added walks, analyses, caches, or repeated computations
+* Did I run the required build/test commands
+* Did I clearly report remaining failures or risks
+
+When reporting back:
+
+* Say what changed
+* Say what was verified
+* Say what remains
+* When showing code in chat, make it easy to copy-paste into the codebase
+* Keep outputs focused on the changed parts
+* List bad practices, fragile assumptions, or cleaner alternatives separately
+* If a change is intentionally pragmatic rather than architecturally ideal, say so and explain the tradeoff

README.md

@@ -105,6 +105,9 @@ Pass these to `onnx-mlir` when compiling for PIM:
   the codegen tail.
 - `--pim-emit-json` - also emit `core_*.json` instruction files alongside
   `core_*.pim`.
+- `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|all>` - control Spatial
+  dataflow CSV reports. The default `all` emits graph, scheduled, and realized
+  snapshots under `reports/`.
 - `--use-experimental-conv-impl` - use the alternate convolution lowering.
 - `--ignore-concat-error` - soft-fail a ConcatOp corner case.
 
@@ -167,9 +170,10 @@ Each validation run writes artifacts in the model workspace, for example under
 - `simulation/out.bin` - raw simulator output used for comparison.
 
 The compiler currently dumps dialect snapshots such as `spatial0.mlir`,
-`spatial1_dcp_merged.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
-`pim2_coalesced.mlir`, `pim3_folded.mlir`, and
-`pim4_materialized.mlir` when an output directory is available.
+`spatial1_graph.mlir`, `spatial2_scheduled_no_comm.mlir`,
+`spatial3_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
+`pim2_coalesced.mlir`, and `pim3_folded.mlir` when an output directory is
+available.
 
 To rerun the simulator manually with tracing after validation has produced a
 `raptor/pim/` directory:

backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs

@@ -258,24 +258,23 @@ where
 
     let (memory, crossbars) = core.get_memory_crossbar();
     let crossbar = crossbars.get_mut(group).unwrap();
-    let crossbar_stored_bytes = crossbar.stored_bytes();
-    let crossbar_byte_width = crossbar.width();
-
-    let crossbar_elem_width = crossbar_byte_width / size_of::<M>();
-    ensure!(
-        crossbar_byte_width % size_of::<M>() == 0,
-        "M not divisor of the crosbbar size"
-    );
-
     let crossbar_height = crossbar.height();
-    let crossbar_byte_size = crossbar_byte_width * crossbar_height;
+    let crossbar_stored_bytes = crossbar.stored_bytes();
+    let bytes_per_column = crossbar_height * size_of::<M>();
+    ensure!(bytes_per_column != 0, "crossbar height can not be zero");
+    ensure!(
+        crossbar_stored_bytes % bytes_per_column == 0,
+        "Stored crossbar bytes do not describe an integral number of columns"
+    );
+    let crossbar_elem_width = crossbar_stored_bytes / bytes_per_column;
+    ensure!(crossbar_elem_width != 0, "Crossbar contains no stored columns");
 
     let loads = memory
         .reserve_load(r1_val, crossbar_height * size_of::<F>())?
         .execute_load::<F>()?;
     let load = loads[0];
     let vec: Cow<[M]> = load.up();
-    let matrix = crossbar.load::<M>(crossbar_byte_size)?[0];
+    let matrix = crossbar.load::<M>(crossbar_stored_bytes)?[0];
 
     // --- FAER IMPLEMENTATION ---
 

invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md

@@ -0,0 +1,362 @@
+# Graph Compute Batch Physical-Fragment Invariant
+
+## Status
+
+This document is **normative** for Raptor's Spatial graph IR.
+
+Every developer or coding agent modifying Spatial graph construction, graph
+verification, Blueprint handling, or `MergeComputeNodes` must read this file
+after `README.md` and `AGENTS.md`.
+
+`AGENTS.md` must contain this instruction:
+
+```text
+* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes.
+```
+
+## Scope
+
+This invariant applies to:
+
+- `spat.graph_compute_batch`;
+- graph-level values produced by it;
+- `tensor.parallel_insert_slice` operations that publish its lane results;
+- `spat.blueprint` operations that describe logical reconstruction;
+- graph analyses and transformations that consume those values;
+- the graph-to-scheduled transition in `MergeComputeNodes`.
+
+It does **not** impose the same representation on:
+
+- `spat.scheduled_compute`;
+- `spat.scheduled_compute_batch`;
+- `pim.core` or `pim.core_batch`;
+- values whose cross-core movement is already represented by explicit
+  `spat.channel_send` and `spat.channel_receive` operations.
+
+Scheduled IR represents execution on assigned cores. Communication and value
+availability there are defined by local SSA forwarding and explicit
+send/receive operations, not by the graph physical-fragment invariant.
+
+## Core invariant
+
+For every result of a `spat.graph_compute_batch` with `N` graph lanes:
+
+1. Every graph lane produces exactly one fragment for that result.
+2. All lanes produce fragments with the same exact ranked tensor type `F`.
+3. The graph result is a physical collection of those fragments with type:
+
+   ```text
+   tensor<N x shape(F) x element-type(F)>
+   ```
+
+   Conceptually, the result is `N × F`: one leading physical fragment-slot
+   dimension followed by the complete per-lane fragment shape.
+4. Physical slot `i` identifies a fragment publication. It does not, by itself,
+   identify a row, column, channel, tile, or any other logical tensor position.
+5. The result type carries no logical reconstruction order.
+
+The leading dimension is therefore a **physical fragment-slot dimension**, not
+a logical tensor dimension.
+
+## Per-lane computation is unrestricted
+
+The invariant constrains the published result representation, not what a lane
+may compute.
+
+A graph lane may:
+
+- read several input slices;
+- perform reductions;
+- add or combine multiple columns;
+- execute matrix/vector operations;
+- produce a fragment that corresponds to any logical region;
+- participate in a multi-stage or logarithmic reduction tree implemented by
+  following `spat.graph_compute` or `spat.graph_compute_batch` operations.
+
+Arithmetic combination is graph computation. `spat.blueprint` is not an
+arithmetic reduction operation.
+
+### Example: `16×4 -> 16×2`
+
+Two graph lanes may compute:
+
+```text
+lane 0: input[:, 0] + input[:, 1] -> tensor<16x1>
+lane 1: input[:, 2] + input[:, 3] -> tensor<16x1>
+```
+
+The physical graph result is:
+
+```text
+tensor<2x16x1>
+```
+
+A Blueprint then maps:
+
+```text
+physical slot 0 -> logical output[:, 0:1]
+physical slot 1 -> logical output[:, 1:2]
+```
+
+and describes the logical result `tensor<16x2>`.
+
+For a larger reduction, following graph compute batches may reduce fragments in
+`ceil(log2(N))` stages. Every intermediate batch still publishes a physical
+`batch × fragment` collection.
+
+## Physical publication inside `spat.graph_compute_batch`
+
+The batch body must publish each lane's fragment into the physical result.
+
+For one result with fragment type `F`, the corresponding
+`tensor.parallel_insert_slice` must insert the fragment into one slot of the
+physical `N × F` destination:
+
+```text
+physical offsets = [slot, 0, 0, ...]
+physical sizes   = [1, shape(F)...]
+physical strides = [1, 1, 1, ...]
+```
+
+The slot may be the graph lane directly or a statically analyzable permutation
+of it. The insertion describes physical slot placement only. It must not use a
+logical output dimension as the physical batch dimension.
+
+For each graph result, the body must contain exactly one physical publication
+per graph lane. Since the body executes once per lane, this normally means one
+`tensor.parallel_insert_slice` operation targeting that result.
+
+## Logical reconstruction
+
+Logical reconstruction is separate from physical publication.
+
+The reconstruction descriptor defines, for every physical fragment slot:
+
+- which physical batch operand owns the fragment;
+- which physical slot contains it;
+- its destination offsets in the logical tensor;
+- its destination sizes;
+- its destination strides;
+- coverage and conflict policy where relevant.
+
+The persistent owner of this information is `spat.blueprint` or an equivalent
+explicit graph-level reconstruction operation.
+
+A logical consumer must not infer reconstruction from the physical tensor type
+or assume that physical slot order equals logical order.
+
+The logical mapping may be arbitrary. For example:
+
+```text
+physical slot 0 -> logical row 13
+physical slot 1 -> logical row 4
+physical slot 2 -> logical row 10
+```
+
+The physical result remains a regular `batch × fragment` tensor.
+
+## Relationship between `parallel_insert_slice` and Blueprint
+
+During graph construction, an algorithm may naturally describe logical
+placement with `tensor.parallel_insert_slice` geometry. Before the graph is in
+its canonical form:
+
+1. that geometry must be separated from physical fragment publication;
+2. the graph batch result must be normalized to `N × F`;
+3. the logical insertion geometry must be transferred to a persistent
+   `spat.blueprint` reconstruction descriptor.
+
+After normalization:
+
+- `parallel_insert_slice` inside `spat.graph_compute_batch` publishes into
+  physical fragment slots;
+- `spat.blueprint` describes reconstruction into the logical tensor.
+
+The original graph operation may be erased only after all reconstruction
+information needed by later stages has a persistent owner.
+
+## Blueprint semantics
+
+Blueprint is placement/reconstruction metadata. It may:
+
+- concatenate fragments;
+- reorder fragments;
+- insert fragments into arbitrary disjoint logical regions;
+- describe complete or partial logical coverage;
+- expose a logical tensor view when materialization is required.
+
+Blueprint must not silently perform arithmetic such as addition, multiplication,
+maximum, or reduction. Such transformations must be represented by following
+`spat.graph_compute` or `spat.graph_compute_batch` operations.
+
+A Blueprint consuming a physical fragment batch must explicitly identify the
+physical source slot for every logical fragment. It must not derive that slot
+from operand order unless that convention is explicitly represented and
+verified.
+
+## Multiple results
+
+A `spat.graph_compute_batch` may have several results.
+
+For each result `r` independently:
+
+- every lane produces one fragment of type `F_r`;
+- the graph result type is `N × F_r`;
+- its physical publication and logical reconstruction descriptor are verified
+  independently.
+
+Different results may use different fragment shapes.
+
+## Graph consumers
+
+A graph consumer of a batch result may:
+
+1. consume fragments directly as physical fragments;
+2. select one or more physical slots in a `spat.deferred_communication` body;
+3. use a Blueprint to obtain or describe a logical reconstruction;
+4. feed fragments to following graph computes or graph compute batches.
+
+A consumer must not treat the leading physical slot dimension as a logical
+model dimension unless an explicit graph operation intentionally performs such
+an interpretation.
+
+All constant selection, slicing, reshaping, concatenation, and other
+compile-time shaping needed for a scheduled consumer must be encoded inside the
+corresponding `spat.deferred_communication` body. Phase 2 must not recover
+missing graph semantics by inspecting consumers after the deferred operation.
+
+## Graph lane, scheduled lane, and physical core are different identities
+
+These concepts must never be conflated:
+
+- **graph lane**: the lane of the original `spat.graph_compute_batch`;
+- **physical fragment slot**: the slot in the graph batch result;
+- **scheduled lane**: one lane of a `spat.scheduled_compute_batch` equivalence
+  class;
+- **physical core**: the core selected by PEFT.
+
+The graph batch body or its Blueprint defines graph-lane-to-fragment-slot and
+fragment-slot-to-logical-region mappings.
+
+PEFT defines graph-instance-to-core placement.
+
+Scheduled communication defines how values move between cores.
+
+## Scheduled IR exclusion
+
+Do not add a verifier requiring `spat.scheduled_compute_batch` results to have
+`laneCount` as their first dimension.
+
+Do not rewrite scheduled values merely to resemble graph physical fragment
+collections.
+
+When lowering graph IR into scheduled IR:
+
+- resolve graph fragments and reconstruction metadata before erasing their
+  graph owners;
+- create local forwarding or `spat.channel_send`/`spat.channel_receive` for
+  cross-core dependencies;
+- allow scheduled result representation to follow the scheduled IR contract;
+- preserve numerical and deadlock correctness.
+
+The graph invariant is an input contract for scheduling, not a scheduled-value
+layout contract.
+
+## Required verifier properties
+
+`spat.graph_compute_batch` verification must establish, for every result:
+
+1. the result is a static or otherwise supported ranked tensor;
+2. result rank is exactly `fragment rank + 1`;
+3. result dimension 0 equals `laneCount`;
+4. every lane publication source has the same exact fragment type;
+5. the physical insertion targets the corresponding result block argument;
+6. physical insertion offsets have the fragment slot in dimension 0;
+7. all remaining physical offsets are zero;
+8. physical sizes are `[1] + fragment shape`;
+9. physical strides are unit;
+10. exactly one publication is defined for each graph result in the per-lane
+    body.
+
+These checks apply only to `spat.graph_compute_batch`, not to
+`spat.scheduled_compute_batch`.
+
+Blueprint verification must establish that every logical reconstruction entry:
+
+- references an existing physical batch operand;
+- references a valid physical fragment slot;
+- maps a fragment compatible with the declared logical slice;
+- stays within logical bounds;
+- follows the declared conflict and coverage policies.
+
+## Invalid representations
+
+The following are invariant violations.
+
+### Logical aggregate returned directly by graph batch
+
+```text
+laneCount = 16
+result     = tensor<1x4x16x16>
+```
+
+with each lane inserting into logical dimension 2.
+
+This is a logical assembly masquerading as a graph batch result. The graph
+result must instead be `16 × per-row-fragment`, and a Blueprint must describe
+placement into `tensor<1x4x16x16>`.
+
+### Physical storage derived from logical destination shape
+
+Code equivalent to:
+
+```cpp
+shape = logicalDestinationType.getShape();
+shape[logicalInsertionDimension] = laneCount;
+```
+
+is invalid.
+
+Physical graph storage must be derived from the per-lane fragment type:
+
+```cpp
+physicalShape = [laneCount] + fragmentType.getShape();
+```
+
+### Reconstruction inferred from result type
+
+It is invalid to assume that physical slot `i` belongs at logical offset `i`.
+The Blueprint or another explicit reconstruction descriptor must state the
+mapping.
+
+### Blueprint used for arithmetic
+
+It is invalid to encode `fragment0 + fragment1` as Blueprint reconstruction.
+Create a following graph compute or graph compute batch for the addition.
+
+## Ownership
+
+- ONNX-to-Spatial lowering owns creation of valid graph fragment batches.
+- Graph canonicalization owns normalization of temporary logical-assembly forms
+  into physical graph batches plus Blueprints.
+- `spat.graph_compute_batch` verifier rejects invalid physical publications.
+- `spat.blueprint` owns persistent logical reconstruction metadata.
+- Deferred communication Phase 1 owns complete consumer-side constant shaping.
+- Merge scheduling consumes this graph contract and introduces explicit
+  communication.
+- Scheduled IR verifiers validate scheduled execution and communication, not
+  the graph fragment representation.
+
+## No repair downstream
+
+If graph IR violates this invariant, fix the graph producer or graph
+canonicalization.
+
+Do not repair an invalid graph batch by:
+
+- guessing a lane dimension in `MergeComputeNodes`;
+- deriving physical storage from a logical destination tensor;
+- inspecting deferred-result users;
+- reconstructing omitted Blueprint data after graph erasure;
+- weakening graph verifiers;
+- imposing the graph representation on scheduled operations.

src/PIM/CMakeLists.txt

@@ -117,13 +117,11 @@ add_pim_library(OMPIMAccel
   SpatialOps
   PimOps
   OMONNXToSpatial
-  OMSpatialToGraphviz
   OMSpatialToPim
   OMPimCommon
   OMPimBufferization
   OMPimMemoryCoalescing
   OMPimHostConstantFolding
-  OMPimHostConstantMaterialization
   OMPimVerification
   MLIRTensorInferTypeOpInterfaceImpl
 )

src/PIM/Common/CMakeLists.txt

@@ -5,9 +5,12 @@ add_pim_library(OMPimCommon
   IR/ConstantUtils.cpp
   IR/CoreBlockUtils.cpp
   IR/EntryPointUtils.cpp
+  IR/IndexingUtils.cpp
   IR/LoopUtils.cpp
   IR/ShapeUtils.cpp
+  IR/ShapingUtils.cpp
   IR/SubviewUtils.cpp
+  IR/TensorSliceUtils.cpp
   IR/WeightUtils.cpp
   Support/CheckedArithmetic.cpp
   Support/DebugDump.cpp

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -34,12 +34,25 @@ mlir::Value resolveAlias(mlir::Value value, const StaticValueKnowledge* knowledg
 
 llvm::FailureOr<CompiledIndexExpr> compileIndexValueImpl(mlir::Value value);
 llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Value value);
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
 template <typename... Args>
 CompiledIndexExpr makeCompiledIndexExpr(Args&&... args) {
   return CompiledIndexExpr(std::make_shared<CompiledIndexExprNode>(std::forward<Args>(args)...));
 }
 
+static mlir::Value resolveForYieldedAliasToInit(mlir::scf::ForOp forOp,
+                                                mlir::Value yieldedValue,
+                                                const StaticValueKnowledge* knowledge) {
+  yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge);
+  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 yieldedValue;
+}
+
 mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
   value = resolveAlias(value, knowledge);
 
@@ -56,6 +69,15 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
         return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge);
   }
 
+  if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(definingOp)) {
+    auto result = mlir::dyn_cast<mlir::OpResult>(value);
+    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);
+    }
+  }
+
   if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(definingOp))
     return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge);
   if (auto collapseOp = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(definingOp))
@@ -499,16 +521,25 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      mlir::Value yieldedValue = resolveLoopCarriedAliasImpl(yieldOp.getOperand(result.getResultNumber()), knowledge);
-      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()) {
-          value = resolveAlias(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge);
-          continue;
-        }
-      }
+      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
+      continue;
+    }
 
-      value = yieldedValue;
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(definingOp)) {
+      auto result = mlir::dyn_cast<mlir::OpResult>(value);
+      if (!result)
+        return mlir::failure();
+
+      auto condition = resolveIndexValueImpl(ifOp.getCondition(), knowledge);
+      if (failed(condition))
+        return mlir::failure();
+
+      mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
+      auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(selectedRegion.front().getTerminator());
+      if (!yieldOp || result.getResultNumber() >= yieldOp.getNumOperands())
+        return mlir::failure();
+
+      value = resolveLoopCarriedAliasImpl(yieldOp.getOperand(result.getResultNumber()), knowledge);
       continue;
     }
 
@@ -609,17 +640,35 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      mlir::Value yieldedValue = yieldOp.getOperand(result.getResultNumber());
-      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()) {
-          value = forOp.getInitArgs()[blockArgument.getArgNumber() - 1];
-          continue;
-        }
+      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), nullptr);
+      continue;
+    }
+
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(definingOp)) {
+      auto result = mlir::dyn_cast<mlir::OpResult>(value);
+      if (!result)
+        return mlir::failure();
+
+      auto thenYield = mlir::dyn_cast<mlir::scf::YieldOp>(ifOp.getThenRegion().front().getTerminator());
+      auto elseYield = mlir::dyn_cast<mlir::scf::YieldOp>(ifOp.getElseRegion().front().getTerminator());
+      if (!thenYield || !elseYield || result.getResultNumber() >= thenYield.getNumOperands()
+          || result.getResultNumber() >= elseYield.getNumOperands()) {
+        return mlir::failure();
       }
 
-      value = yieldedValue;
-      continue;
+      auto thenAddress = compileContiguousAddressExprImpl(thenYield.getOperand(result.getResultNumber()));
+      auto elseAddress = compileContiguousAddressExprImpl(elseYield.getOperand(result.getResultNumber()));
+      if (failed(thenAddress) || failed(elseAddress) || thenAddress->base != elseAddress->base)
+        return mlir::failure();
+
+      auto condition = compileIndexValueImpl(ifOp.getCondition());
+      if (failed(condition))
+        return mlir::failure();
+
+      CompiledIndexExprNode selectExpr;
+      selectExpr.kind = CompiledIndexExprNode::Kind::Select;
+      selectExpr.operands = {*condition, thenAddress->byteOffset, elseAddress->byteOffset};
+      return CompiledAddressExpr {thenAddress->base, makeCompiledIndexExpr(std::move(selectExpr))};
     }
 
     if (auto subviewOp = mlir::dyn_cast<mlir::memref::SubViewOp>(definingOp)) {
@@ -801,7 +850,7 @@ llvm::FailureOr<ResolvedContiguousAddress> CompiledAddressExpr::evaluate(const S
   auto resolvedOffset = byteOffset.evaluate(knowledge);
   if (failed(resolvedOffset))
     return mlir::failure();
-  return ResolvedContiguousAddress {base, *resolvedOffset};
+  return ResolvedContiguousAddress {resolveAlias(base, &knowledge), *resolvedOffset};
 }
 
 } // namespace onnx_mlir

src/PIM/Common/IR/AffineUtils.cpp

@@ -31,7 +31,7 @@ static FailureOr<int64_t> ceilDivSigned(int64_t lhs, int64_t rhs) {
 }
 
 Value createOrFoldAffineApply(
-  RewriterBase& rewriter, Location loc, AffineMap map, ValueRange operands, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, AffineMap map, ValueRange operands, Operation* constantAnchor) {
   assert(constantAnchor && "expected a valid constant anchor");
   assert(map.getNumResults() == 1 && "affine.apply expects a single-result affine map");
 
@@ -40,54 +40,91 @@ Value createOrFoldAffineApply(
   for (Value operand : operands) {
     std::optional<int64_t> constantValue = matchConstantIndexValue(operand);
     if (!constantValue)
-      return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
-    operandConstants.push_back(rewriter.getIndexAttr(*constantValue));
+      return affine::AffineApplyOp::create(builder, loc, map, operands).getResult();
+    operandConstants.push_back(builder.getIndexAttr(*constantValue));
   }
 
   SmallVector<Attribute> foldedResults;
   if (succeeded(map.constantFold(operandConstants, foldedResults)) && foldedResults.size() == 1)
     if (auto constantResult = dyn_cast<IntegerAttr>(foldedResults.front()))
-      return getOrCreateIndexConstant(rewriter, constantAnchor, constantResult.getInt());
+      return getOrCreateIndexConstant(builder, constantAnchor, constantResult.getInt());
 
-  return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
+  return affine::AffineApplyOp::create(builder, loc, map, operands).getResult();
 }
 
 Value createOrFoldAffineApply(
-  RewriterBase& rewriter, Location loc, AffineExpr expr, ValueRange dims, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, AffineExpr expr, ValueRange dims, Operation* constantAnchor) {
   AffineMap map = AffineMap::get(/*dimCount=*/dims.size(), /*symbolCount=*/0, expr);
-  return createOrFoldAffineApply(rewriter, loc, map, dims, constantAnchor);
+  return createOrFoldAffineApply(builder, loc, map, dims, constantAnchor);
 }
 
-Value affineMulConst(RewriterBase& rewriter, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) {
+Value affineMulConst(OpBuilder& builder, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) {
   assert(constantAnchor && "expected a valid constant anchor");
   if (multiplier == 0)
-    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);
   if (multiplier == 1)
     return value;
 
-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0 * multiplier, ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 * multiplier, ValueRange {value}, constantAnchor);
 }
 
-Value affineModConst(RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
+Value affineAddConst(OpBuilder& builder, Location loc, Value value, int64_t offset, Operation* constantAnchor) {
+  assert(constantAnchor && "expected a valid constant anchor");
+  if (offset == 0)
+    return value;
+
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 + offset, ValueRange {value}, constantAnchor);
+}
+
+Value affineModConst(OpBuilder& builder, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
   assert(constantAnchor && "expected a valid constant anchor");
   assert(divisor > 0 && "expected a positive affine.mod divisor");
   if (divisor == 1)
-    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);
 
-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0 % divisor, ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 % divisor, ValueRange {value}, constantAnchor);
 }
 
 Value affineFloorDivConst(
-  RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
   assert(constantAnchor && "expected a valid constant anchor");
   assert(divisor > 0 && "expected a positive affine.floor_div divisor");
   if (divisor == 1)
     return value;
 
-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor);
+}
+
+Value affineAddModConst(
+  OpBuilder& builder, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
+  assert(constantAnchor && "expected a valid constant anchor");
+  assert(divisor > 0 && "expected a positive affine.mod divisor");
+  if (divisor == 1)
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);
+
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  AffineExpr expr = d0;
+  if (offset != 0)
+    expr = expr + offset;
+  return createOrFoldAffineApply(builder, loc, expr % divisor, ValueRange {value}, constantAnchor);
+}
+
+Value affineAddFloorDivConst(
+  OpBuilder& builder, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
+  assert(constantAnchor && "expected a valid constant anchor");
+  assert(divisor > 0 && "expected a positive affine.floor_div divisor");
+  if (divisor == 1)
+    return offset == 0 ? value : affineAddConst(builder, loc, value, offset, constantAnchor);
+
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  AffineExpr expr = d0;
+  if (offset != 0)
+    expr = expr + offset;
+  return createOrFoldAffineApply(builder, loc, expr.floorDiv(divisor), ValueRange {value}, constantAnchor);
 }
 
 FailureOr<int64_t> evaluateAffineExpr(AffineExpr expr, ArrayRef<int64_t> dims, ArrayRef<int64_t> symbols) {

src/PIM/Common/IR/AffineUtils.hpp

@@ -11,36 +11,56 @@ namespace onnx_mlir {
 
 using IndexValueResolver = llvm::function_ref<llvm::FailureOr<int64_t>(mlir::Value)>;
 
-mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter,
+mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder,
                                     mlir::Location loc,
                                     mlir::AffineMap map,
                                     mlir::ValueRange operands,
                                     mlir::Operation* constantAnchor);
 
-mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter,
+mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder,
                                     mlir::Location loc,
                                     mlir::AffineExpr expr,
                                     mlir::ValueRange dims,
                                     mlir::Operation* constantAnchor);
 
-mlir::Value affineMulConst(mlir::RewriterBase& rewriter,
+mlir::Value affineMulConst(mlir::OpBuilder& builder,
                            mlir::Location loc,
                            mlir::Value value,
                            int64_t multiplier,
                            mlir::Operation* constantAnchor);
 
-mlir::Value affineModConst(mlir::RewriterBase& rewriter,
+mlir::Value affineAddConst(mlir::OpBuilder& builder,
+                           mlir::Location loc,
+                           mlir::Value value,
+                           int64_t offset,
+                           mlir::Operation* constantAnchor);
+
+mlir::Value affineModConst(mlir::OpBuilder& builder,
                            mlir::Location loc,
                            mlir::Value value,
                            int64_t divisor,
                            mlir::Operation* constantAnchor);
 
-mlir::Value affineFloorDivConst(mlir::RewriterBase& rewriter,
+mlir::Value affineFloorDivConst(mlir::OpBuilder& builder,
                                 mlir::Location loc,
                                 mlir::Value value,
                                 int64_t divisor,
                                 mlir::Operation* constantAnchor);
 
+mlir::Value affineAddModConst(mlir::OpBuilder& builder,
+                              mlir::Location loc,
+                              mlir::Value value,
+                              int64_t offset,
+                              int64_t divisor,
+                              mlir::Operation* constantAnchor);
+
+mlir::Value affineAddFloorDivConst(mlir::OpBuilder& builder,
+                                   mlir::Location loc,
+                                   mlir::Value value,
+                                   int64_t offset,
+                                   int64_t divisor,
+                                   mlir::Operation* constantAnchor);
+
 llvm::FailureOr<int64_t>
 evaluateAffineExpr(mlir::AffineExpr expr, llvm::ArrayRef<int64_t> dims, llvm::ArrayRef<int64_t> symbols = {});
 

src/PIM/Common/IR/BatchCoreUtils.cpp

@@ -1,5 +1,6 @@
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
 namespace onnx_mlir {
 
@@ -9,6 +10,65 @@ llvm::SmallVector<int32_t> getBatchCoreIds(pim::PimCoreBatchOp coreBatchOp) {
   return llvm::SmallVector<int32_t>(coreIdsAttr.asArrayRef().begin(), coreIdsAttr.asArrayRef().end());
 }
 
+mlir::FailureOr<std::optional<int32_t>>
+getOptionalScheduledCoreId(spatial::SpatScheduledCompute computeOp, llvm::StringRef fieldName) {
+  auto coreIdAttr = computeOp->getAttrOfType<mlir::IntegerAttr>(onnx_mlir::kCoreIdAttrName);
+  if (!coreIdAttr)
+    return std::optional<int32_t> {};
+  if (coreIdAttr.getInt() < 0) {
+    computeOp.emitOpError() << fieldName << " must be non-negative";
+    return mlir::failure();
+  }
+  auto checkedCoreId = pim::checkedI32(coreIdAttr.getInt(), computeOp, fieldName);
+  if (mlir::failed(checkedCoreId))
+    return mlir::failure();
+  return std::optional<int32_t> {*checkedCoreId};
+}
+
+mlir::FailureOr<int32_t> getRequiredScheduledCoreId(spatial::SpatScheduledCompute computeOp, llvm::StringRef fieldName) {
+  auto coreId = getOptionalScheduledCoreId(computeOp, fieldName);
+  if (mlir::failed(coreId))
+    return mlir::failure();
+  if (!*coreId) {
+    computeOp.emitOpError() << "missing required " << onnx_mlir::kCoreIdAttrName;
+    return mlir::failure();
+  }
+  return **coreId;
+}
+
+mlir::FailureOr<std::optional<llvm::SmallVector<int32_t>>>
+getOptionalScheduledBatchCoreIds(spatial::SpatScheduledComputeBatch computeBatchOp, llvm::StringRef fieldName) {
+  auto coreIdsAttr = computeBatchOp->getAttrOfType<mlir::DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName);
+  if (!coreIdsAttr)
+    return std::optional<llvm::SmallVector<int32_t>> {};
+
+  llvm::SmallVector<int32_t> coreIds;
+  coreIds.reserve(coreIdsAttr.size());
+  for (int32_t coreId : coreIdsAttr.asArrayRef()) {
+    if (coreId < 0) {
+      computeBatchOp.emitOpError() << fieldName << " values must be non-negative";
+      return mlir::failure();
+    }
+    auto checkedCoreId = pim::checkedI32(static_cast<int64_t>(coreId), computeBatchOp, fieldName);
+    if (mlir::failed(checkedCoreId))
+      return mlir::failure();
+    coreIds.push_back(*checkedCoreId);
+  }
+  return std::optional<llvm::SmallVector<int32_t>> {std::move(coreIds)};
+}
+
+mlir::FailureOr<llvm::SmallVector<int32_t>>
+getRequiredScheduledBatchCoreIds(spatial::SpatScheduledComputeBatch computeBatchOp, llvm::StringRef fieldName) {
+  auto coreIds = getOptionalScheduledBatchCoreIds(computeBatchOp, fieldName);
+  if (mlir::failed(coreIds))
+    return mlir::failure();
+  if (!*coreIds) {
+    computeBatchOp.emitOpError() << "missing required " << onnx_mlir::kCoreIdsAttrName;
+    return mlir::failure();
+  }
+  return std::move(**coreIds);
+}
+
 llvm::SmallVector<int32_t> getLaneChunkCoreIds(llvm::ArrayRef<int32_t> coreIds, size_t laneCount, unsigned lane) {
   llvm::SmallVector<int32_t> laneCoreIds;
   laneCoreIds.reserve(coreIds.size() / laneCount);

src/PIM/Common/IR/BatchCoreUtils.hpp

@@ -3,12 +3,26 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
+#include <optional>
+
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 namespace onnx_mlir {
 
 llvm::SmallVector<int32_t> getBatchCoreIds(pim::PimCoreBatchOp coreBatchOp);
 
+mlir::FailureOr<std::optional<int32_t>>
+getOptionalScheduledCoreId(spatial::SpatScheduledCompute computeOp, llvm::StringRef fieldName);
+
+mlir::FailureOr<int32_t> getRequiredScheduledCoreId(spatial::SpatScheduledCompute computeOp, llvm::StringRef fieldName);
+
+mlir::FailureOr<std::optional<llvm::SmallVector<int32_t>>>
+getOptionalScheduledBatchCoreIds(spatial::SpatScheduledComputeBatch computeBatchOp, llvm::StringRef fieldName);
+
+mlir::FailureOr<llvm::SmallVector<int32_t>>
+getRequiredScheduledBatchCoreIds(spatial::SpatScheduledComputeBatch computeBatchOp, llvm::StringRef fieldName);
+
 llvm::SmallVector<int32_t> getLaneChunkCoreIds(llvm::ArrayRef<int32_t> coreIds, size_t laneCount, unsigned lane);
 
 bool isExplicitHostMemCopyOperand(mlir::Operation* op, unsigned operandIndex);

src/PIM/Common/IR/ConstantUtils.cpp

@@ -49,7 +49,7 @@ Value getOrCreateConstant(OperationFolder& folder, Operation* anchorOp, Attribut
   return folder.getOrCreateConstant(hostBlock, arithDialect, value, type);
 }
 
-Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute value, Type type) {
+Value getOrCreateConstant(OpBuilder& builder, Operation* anchorOp, Attribute value, Type type) {
   assert(anchorOp && "expected a valid anchor operation");
   Block* hostBlock = getConstantInsertionBlock(anchorOp);
   for (Operation& op : *hostBlock) {
@@ -59,9 +59,16 @@ Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute
     return constantOp.getResult();
   }
 
-  OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPointToStart(hostBlock);
-  return arith::ConstantOp::create(rewriter, anchorOp->getLoc(), type, cast<TypedAttr>(value)).getResult();
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(hostBlock);
+  return arith::ConstantOp::create(builder, anchorOp->getLoc(), type, cast<TypedAttr>(value)).getResult();
+}
+
+Value createConstantAtHostBlockStart(OpBuilder& builder, Operation* anchorOp, TypedAttr value) {
+  assert(anchorOp && "expected a valid anchor operation");
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(getConstantInsertionBlock(anchorOp));
+  return arith::ConstantOp::create(builder, anchorOp->getLoc(), value).getResult();
 }
 
 Value getOrCreateConstantLike(OperationFolder& folder, arith::ConstantOp constantOp) {
@@ -73,9 +80,8 @@ Value getOrCreateIndexConstant(OperationFolder& folder, Operation* anchorOp, int
   return getOrCreateConstant(folder, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
 }
 
-Value getOrCreateIndexConstant(RewriterBase& rewriter, Operation* anchorOp, int64_t value) {
-  Builder builder(anchorOp->getContext());
-  return getOrCreateConstant(rewriter, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
+Value getOrCreateIndexConstant(OpBuilder& builder, Operation* anchorOp, int64_t value) {
+  return getOrCreateConstant(builder, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
 }
 
 void hoistAndUniquifyIndexConstants(func::FuncOp funcOp, RewriterBase& rewriter) {

src/PIM/Common/IR/ConstantUtils.hpp

@@ -16,13 +16,16 @@ mlir::Value
 getOrCreateConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
 
 mlir::Value
-getOrCreateConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
+getOrCreateConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
+
+mlir::Value
+createConstantAtHostBlockStart(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::TypedAttr value);
 
 mlir::Value getOrCreateConstantLike(mlir::OperationFolder& folder, mlir::arith::ConstantOp constantOp);
 
 mlir::Value getOrCreateIndexConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, int64_t value);
 
-mlir::Value getOrCreateIndexConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, int64_t value);
+mlir::Value getOrCreateIndexConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, int64_t value);
 
 void hoistAndUniquifyIndexConstants(mlir::func::FuncOp funcOp, mlir::RewriterBase& rewriter);
 

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -36,9 +36,10 @@ bool isCoreStaticAddressOp(mlir::Operation* op) {
 
 mlir::LogicalResult
 walkPimCoreBlock(mlir::Block& block,
-                 const StaticValueKnowledge& knowledge,
+                 const StaticValueKnowledge& initialKnowledge,
                  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
   bool hasFailure = false;
+  StaticValueKnowledge knowledge = initialKnowledge;
   for (mlir::Operation& op : block) {
     if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
       continue;
@@ -74,6 +75,42 @@ walkPimCoreBlock(mlir::Block& block,
       continue;
     }
 
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
+      auto condition = resolveIndexValue(ifOp.getCondition(), knowledge);
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        hasFailure = true;
+        continue;
+      }
+
+      mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
+      if (!selectedRegion.empty())
+        if (failed(walkPimCoreBlock(selectedRegion.front(), knowledge, callback)))
+          hasFailure = true;
+      continue;
+    }
+
+    if (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        hasFailure = true;
+        continue;
+      }
+      mlir::Region* selected = &switchOp.getDefaultRegion();
+      for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions()))
+        if (caseValue == *selector) {
+          selected = &caseRegion;
+          break;
+        }
+      if (failed(walkPimCoreBlock(selected->front(), knowledge, callback)))
+        hasFailure = true;
+      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }
@@ -82,9 +119,10 @@ walkPimCoreBlock(mlir::Block& block,
 
 mlir::LogicalResult walkPimCoreBlockStructurally(
   mlir::Block& block,
-  const StaticValueKnowledge& knowledge,
+  const StaticValueKnowledge& initialKnowledge,
   llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
   bool hasFailure = false;
+  StaticValueKnowledge knowledge = initialKnowledge;
   for (mlir::Operation& op : block) {
     if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
       continue;
@@ -128,6 +166,44 @@ mlir::LogicalResult walkPimCoreBlockStructurally(
       continue;
     }
 
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
+      if (failed(resolveIndexValue(ifOp.getCondition(), knowledge))) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM verification");
+        hasFailure = true;
+        continue;
+      }
+
+      if (!ifOp.getThenRegion().empty())
+        if (failed(walkPimCoreBlockStructurally(ifOp.getThenRegion().front(), knowledge, callback)))
+          hasFailure = true;
+      if (!ifOp.getElseRegion().empty())
+        if (failed(walkPimCoreBlockStructurally(ifOp.getElseRegion().front(), knowledge, callback)))
+          hasFailure = true;
+      continue;
+    }
+
+    if (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM verification");
+        hasFailure = true;
+        continue;
+      }
+      mlir::Region* selected = &switchOp.getDefaultRegion();
+      for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions()))
+        if (caseValue == *selector) {
+          selected = &caseRegion;
+          break;
+        }
+      for (mlir::Region& region : switchOp->getRegions())
+        if (failed(walkPimCoreBlockStructurally(region.front(), knowledge, callback)))
+          hasFailure = true;
+      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }

src/PIM/Common/IR/IndexingUtils.cpp

@@ -1,6 +1,6 @@
 #include <algorithm>
 
-#include "IndexingUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/IndexingUtils.hpp"
 
 using namespace mlir;
 

src/PIM/Common/IR/ShapeUtils.cpp

@@ -1,6 +1,9 @@
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/ErrorHandling.h"
 
+#include <functional>
+
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 
 namespace onnx_mlir {
@@ -163,4 +166,80 @@ bool isContiguousSubviewWithDynamicOffsets(llvm::ArrayRef<int64_t> sourceShape,
   return true;
 }
 
+bool hasStaticPositiveShape(llvm::ArrayRef<int64_t> shape) {
+  return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
+}
+
+bool hasStaticPositiveShape(mlir::RankedTensorType type) {
+  return type.hasStaticShape() && hasStaticPositiveShape(type.getShape());
+}
+
+int64_t getStaticShapeElementCount(llvm::ArrayRef<int64_t> shape) {
+  return std::accumulate(shape.begin(), shape.end(), int64_t {1}, std::multiplies<int64_t> {});
+}
+
+llvm::SmallVector<int64_t> permuteShape(llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> permutation) {
+  llvm::SmallVector<int64_t> permutedShape;
+  permutedShape.reserve(permutation.size());
+  for (int64_t axis : permutation)
+    permutedShape.push_back(shape[axis]);
+  return permutedShape;
+}
+
+llvm::SmallVector<int64_t> invertPermutation(llvm::ArrayRef<int64_t> permutation) {
+  llvm::SmallVector<int64_t> inversePermutation(permutation.size());
+  for (auto [newIndex, oldIndex] : llvm::enumerate(permutation))
+    inversePermutation[oldIndex] = static_cast<int64_t>(newIndex);
+  return inversePermutation;
+}
+
+mlir::FailureOr<llvm::SmallVector<int64_t>>
+getTransposePermutationChecked(std::optional<mlir::ArrayAttr> permAttr, int64_t rank) {
+  llvm::SmallVector<int64_t> permutation;
+  if (!permAttr) {
+    permutation.reserve(rank);
+    for (int64_t dim = rank - 1; dim >= 0; --dim)
+      permutation.push_back(dim);
+    return permutation;
+  }
+
+  if (static_cast<int64_t>(permAttr->size()) != rank)
+    return mlir::failure();
+
+  permutation.reserve(permAttr->size());
+  llvm::SmallVector<bool> seen(rank, false);
+  for (mlir::IntegerAttr attr : permAttr->getAsRange<mlir::IntegerAttr>()) {
+    int64_t axis = attr.getInt();
+    if (axis < 0 || axis >= rank || seen[axis])
+      return mlir::failure();
+    seen[axis] = true;
+    permutation.push_back(axis);
+  }
+  return permutation;
+}
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticIndexAttrs(mlir::Builder& builder, llvm::ArrayRef<int64_t> values) {
+  llvm::SmallVector<mlir::OpFoldResult> attrs;
+  attrs.reserve(values.size());
+  for (int64_t value : values)
+    attrs.push_back(builder.getIndexAttr(value));
+  return attrs;
+}
+
+llvm::SmallVector<mlir::OpFoldResult> getUnitStrides(mlir::PatternRewriter& rewriter, int64_t rank) {
+  return llvm::SmallVector<mlir::OpFoldResult>(rank, rewriter.getIndexAttr(1));
+}
+
+llvm::SmallVector<mlir::OpFoldResult> getZeroOffsets(mlir::PatternRewriter& rewriter, int64_t rank) {
+  return llvm::SmallVector<mlir::OpFoldResult>(rank, rewriter.getIndexAttr(0));
+}
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticSizes(mlir::PatternRewriter& rewriter, llvm::ArrayRef<int64_t> shape) {
+  llvm::SmallVector<mlir::OpFoldResult> sizes;
+  sizes.reserve(shape.size());
+  for (int64_t dim : shape)
+    sizes.push_back(rewriter.getIndexAttr(dim));
+  return sizes;
+}
+
 } // namespace onnx_mlir

src/PIM/Common/IR/ShapeUtils.hpp

@@ -2,15 +2,23 @@
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Value.h"
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <cstddef>
+#include <optional>
+#include <type_traits>
+#include <utility>
 
 namespace onnx_mlir {
 
+using HSliceId = size_t;
+using CoreId = size_t;
+
 llvm::SmallVector<int64_t> computeRowMajorStrides(llvm::ArrayRef<int64_t> shape);
 
 llvm::SmallVector<int64_t>
@@ -36,4 +44,69 @@ bool isContiguousSubviewWithDynamicOffsets(llvm::ArrayRef<int64_t> sourceShape,
                                            llvm::ArrayRef<int64_t> staticSizes,
                                            llvm::ArrayRef<int64_t> staticStrides);
 
+template <class A, class B, class C = std::common_type_t<A, B>>
+constexpr C ceilIntegerDivide(A a, B b) {
+  static_assert(std::is_integral_v<A>, "A must be an integer type");
+  static_assert(std::is_integral_v<B>, "B must be an integer type");
+  C ac = static_cast<C>(a);
+  C bc = static_cast<C>(b);
+  return 1 + (ac - 1) / bc;
+}
+
+template <class A, class B, class C = std::common_type_t<A, B>>
+constexpr std::pair<C, C> ceilIntegerDivideWithRemainder(A a, B b) {
+  static_assert(std::is_integral_v<A>, "A must be an integer type");
+  static_assert(std::is_integral_v<B>, "B must be an integer type");
+  C ac = static_cast<C>(a);
+  C bc = static_cast<C>(b);
+  return {ceilIntegerDivide(ac, bc), ac % bc};
+}
+
+template <class T>
+bool isVectorShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2 && (shape[0] == 1 || shape[1] == 1);
+}
+
+template <class T>
+bool isMatrixShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2;
+}
+
+template <class T>
+bool isHVectorShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2 && shape[0] == 1;
+}
+
+inline auto getTensorShape(mlir::Value tensor) {
+  return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
+}
+
+inline bool haveSameStaticShape(mlir::Value lhs, mlir::Value rhs) {
+  auto lhsType = mlir::dyn_cast<mlir::RankedTensorType>(lhs.getType());
+  auto rhsType = mlir::dyn_cast<mlir::RankedTensorType>(rhs.getType());
+  return lhsType && rhsType && lhsType.hasStaticShape() && rhsType.hasStaticShape()
+      && lhsType.getShape() == rhsType.getShape();
+}
+
+bool hasStaticPositiveShape(mlir::ArrayRef<int64_t> shape);
+
+bool hasStaticPositiveShape(mlir::RankedTensorType type);
+
+int64_t getStaticShapeElementCount(mlir::ArrayRef<int64_t> shape);
+
+llvm::SmallVector<int64_t> permuteShape(mlir::ArrayRef<int64_t> shape, mlir::ArrayRef<int64_t> permutation);
+
+llvm::SmallVector<int64_t> invertPermutation(mlir::ArrayRef<int64_t> permutation);
+
+mlir::FailureOr<llvm::SmallVector<int64_t>> getTransposePermutationChecked(std::optional<mlir::ArrayAttr> permAttr,
+                                                                           int64_t rank);
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticIndexAttrs(mlir::Builder& builder, llvm::ArrayRef<int64_t> values);
+
+llvm::SmallVector<mlir::OpFoldResult> getUnitStrides(mlir::PatternRewriter& rewriter, int64_t rank);
+
+llvm::SmallVector<mlir::OpFoldResult> getZeroOffsets(mlir::PatternRewriter& rewriter, int64_t rank);
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticSizes(mlir::PatternRewriter& rewriter, llvm::ArrayRef<int64_t> shape);
+
 } // namespace onnx_mlir

src/PIM/Common/IR/ShapingUtils.cpp

@@ -0,0 +1,39 @@
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+bool isShapingOnlyOp(Operation *op) {
+  return isa<tensor::CastOp,
+             tensor::CollapseShapeOp,
+             tensor::ExpandShapeOp,
+             tensor::ExtractSliceOp,
+             tensor::InsertSliceOp,
+             tensor::ConcatOp,
+             tensor::EmptyOp,
+             tensor::ExtractOp,
+             tensor::InsertOp,
+             tensor::SplatOp,
+             linalg::TransposeOp,
+             ONNXTransposeOp,
+             spatial::SpatConcatOp,
+             spatial::SpatExtractRowsOp>(op);
+}
+
+bool isPureIndexComputationOp(Operation *op) {
+  if (op->getNumRegions() != 0 || op->getNumResults() == 0 || op->hasTrait<OpTrait::IsTerminator>()
+      || !isMemoryEffectFree(op))
+    return false;
+  auto isIndexOrInteger = [](Type type) { return type.isIndex() || isa<IntegerType>(type); };
+  return llvm::all_of(op->getOperandTypes(), isIndexOrInteger)
+      && llvm::all_of(op->getResultTypes(), isIndexOrInteger);
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/ShapingUtils.hpp

@@ -0,0 +1,13 @@
+#pragma once
+
+namespace mlir {
+class Operation;
+}
+
+namespace onnx_mlir {
+
+bool isShapingOnlyOp(mlir::Operation *op);
+
+bool isPureIndexComputationOp(mlir::Operation *op);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/TensorSliceUtils.cpp

@@ -0,0 +1,106 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+Value extractAxisSlice(
+  PatternRewriter& rewriter, Location loc, Value source, int64_t axis, int64_t offset, int64_t size) {
+  auto sourceType = cast<RankedTensorType>(source.getType());
+  SmallVector<int64_t> resultShape(sourceType.getShape());
+  resultShape[axis] = size;
+  auto resultType = RankedTensorType::get(resultShape, sourceType.getElementType(), sourceType.getEncoding());
+
+  SmallVector<OpFoldResult> offsets = getZeroOffsets(rewriter, sourceType.getRank());
+  SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, sourceType.getShape());
+  offsets[axis] = rewriter.getIndexAttr(offset);
+  sizes[axis] = rewriter.getIndexAttr(size);
+  return tensor::ExtractSliceOp::create(
+           rewriter, loc, resultType, source, offsets, sizes, getUnitStrides(rewriter, sourceType.getRank()))
+    .getResult();
+}
+
+Value extractStaticSliceOrIdentity(RewriterBase& rewriter,
+                                   Location loc,
+                                   Value source,
+                                   RankedTensorType resultType,
+                                   ArrayRef<OpFoldResult> offsets,
+                                   ArrayRef<OpFoldResult> sizes,
+                                   ArrayRef<OpFoldResult> strides) {
+  auto sourceType = cast<RankedTensorType>(source.getType());
+  size_t rank = static_cast<size_t>(sourceType.getRank());
+
+  bool isIdentitySlice =
+    sourceType == resultType && sourceType.hasStaticShape() && offsets.size() == rank && sizes.size() == rank
+    && strides.size() == rank;
+  if (isIdentitySlice) {
+    ArrayRef<int64_t> sourceShape = sourceType.getShape();
+    for (auto [dim, offset, size, stride] : llvm::zip_equal(sourceShape, offsets, sizes, strides)) {
+      std::optional<int64_t> staticOffset = mlir::getConstantIntValue(offset);
+      std::optional<int64_t> staticSize = mlir::getConstantIntValue(size);
+      std::optional<int64_t> staticStride = mlir::getConstantIntValue(stride);
+      if (!staticOffset || !staticSize || !staticStride || *staticOffset != 0 || *staticSize != dim
+          || *staticStride != 1) {
+        isIdentitySlice = false;
+        break;
+      }
+    }
+  }
+
+  if (isIdentitySlice)
+    return source;
+
+  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, source, offsets, sizes, strides).getResult();
+}
+
+Value insertStaticSlice(
+  PatternRewriter& rewriter, Location loc, Value source, Value dest, ArrayRef<OpFoldResult> offsets) {
+  auto sourceType = cast<RankedTensorType>(source.getType());
+  return tensor::InsertSliceOp::create(rewriter,
+                                       loc,
+                                       source,
+                                       dest,
+                                       offsets,
+                                       getStaticSizes(rewriter, sourceType.getShape()),
+                                       getUnitStrides(rewriter, sourceType.getRank()))
+    .getResult();
+}
+
+FailureOr<Value> addLeadingUnitTensorDimension(OpBuilder& builder, Location loc, Value value) {
+  auto type = dyn_cast<RankedTensorType>(value.getType());
+  if (!type || !type.hasStaticShape())
+    return failure();
+  SmallVector<int64_t> shape {1};
+  llvm::append_range(shape, type.getShape());
+  auto resultType = RankedTensorType::get(shape, type.getElementType(), type.getEncoding());
+  SmallVector<ReassociationIndices> reassociation;
+  if (type.getRank() != 0) {
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < type.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+  }
+  return tensor::ExpandShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
+}
+
+FailureOr<Value> removeLeadingUnitTensorDimension(
+  OpBuilder& builder, Location loc, Value value, RankedTensorType resultType) {
+  if (value.getType() == resultType)
+    return value;
+  auto type = dyn_cast<RankedTensorType>(value.getType());
+  if (!type || !resultType || !type.hasStaticShape() || !resultType.hasStaticShape()
+      || type.getRank() != resultType.getRank() + 1 || type.getDimSize(0) != 1
+      || type.getElementType() != resultType.getElementType()
+      || !llvm::equal(type.getShape().drop_front(), resultType.getShape()))
+    return failure();
+  SmallVector<ReassociationIndices> reassociation;
+  if (resultType.getRank() != 0) {
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < resultType.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+  }
+  return tensor::CollapseShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/TensorSliceUtils.hpp

@@ -0,0 +1,34 @@
+#pragma once
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/ValueRange.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+
+namespace onnx_mlir {
+
+mlir::Value extractAxisSlice(
+  mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
+
+mlir::Value extractStaticSliceOrIdentity(mlir::RewriterBase& rewriter,
+                                         mlir::Location loc,
+                                         mlir::Value source,
+                                         mlir::RankedTensorType resultType,
+                                         llvm::ArrayRef<mlir::OpFoldResult> offsets,
+                                         llvm::ArrayRef<mlir::OpFoldResult> sizes,
+                                         llvm::ArrayRef<mlir::OpFoldResult> strides);
+
+mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
+                              mlir::Location loc,
+                              mlir::Value source,
+                              mlir::Value dest,
+                              llvm::ArrayRef<mlir::OpFoldResult> offsets);
+
+mlir::FailureOr<mlir::Value>
+addLeadingUnitTensorDimension(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value);
+
+mlir::FailureOr<mlir::Value> removeLeadingUnitTensorDimension(
+  mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, mlir::RankedTensorType resultType);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/WeightUtils.cpp

@@ -47,6 +47,16 @@ CompiledIndexExpr mulExpr(CompiledIndexExpr lhs, int64_t rhs) {
   return makeBinaryExpr(CompiledIndexExprNode::Kind::Mul, std::move(lhs), makeConstantExpr(rhs));
 }
 
+llvm::FailureOr<llvm::SmallVector<int64_t>> getStaticMemRefTypeStrides(mlir::MemRefType type) {
+  llvm::SmallVector<int64_t> strides;
+  int64_t offset = 0;
+  if (failed(type.getStridesAndOffset(strides, offset)))
+    return mlir::failure();
+  if (llvm::is_contained(strides, mlir::ShapedType::kDynamic))
+    return mlir::failure();
+  return strides;
+}
+
 template <typename VMMOpTy, typename ParentOpTy>
 bool hasVmmWeightUse(ParentOpTy parentOp, unsigned weightIndex) {
   auto weightArg = parentOp.getWeightArgument(weightIndex);
@@ -162,6 +172,11 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static
   mlir::Value current = weight;
 
   while (true) {
+    if (mlir::Value directAlias = knowledge.aliases.lookup(current); directAlias && directAlias != current) {
+      current = directAlias;
+      continue;
+    }
+
     if (auto defOp = current.getDefiningOp()) {
       if (auto getGlobalOp = mlir::dyn_cast<mlir::memref::GetGlobalOp>(defOp)) {
         auto moduleOp = weightOwner ? weightOwner->getParentOfType<mlir::ModuleOp>() : mlir::ModuleOp {};
@@ -181,8 +196,6 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static
         CompiledIndexExpr offsetExpr = makeConstantExpr(0);
         for (mlir::Operation* viewOp : llvm::reverse(viewOps)) {
           if (auto subview = mlir::dyn_cast<mlir::memref::SubViewOp>(viewOp)) {
-            llvm::SmallVector<int64_t> nextStrides;
-            nextStrides.reserve(subview.getMixedOffsets().size());
             for (auto [offset, stride, sourceStride] :
                  llvm::zip_equal(subview.getMixedOffsets(), subview.getStaticStrides(), view.strides)) {
               CompiledIndexExpr offsetValue = makeConstantExpr(0);
@@ -202,29 +215,47 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static
                 return mlir::failure();
               }
               offsetExpr = addExpr(std::move(offsetExpr), mulExpr(std::move(offsetValue), sourceStride));
-              nextStrides.push_back(stride * sourceStride);
             }
-            view.shape.assign(subview.getStaticSizes().begin(), subview.getStaticSizes().end());
-            view.strides = std::move(nextStrides);
+            auto resultType = mlir::cast<mlir::MemRefType>(subview.getResult().getType());
+            auto resultStrides = getStaticMemRefTypeStrides(resultType);
+            if (failed(resultStrides))
+              return mlir::failure();
+            view.shape.assign(resultType.getShape().begin(), resultType.getShape().end());
+            view.strides = std::move(*resultStrides);
             continue;
           }
 
           if (auto collapse = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(viewOp)) {
-            if (view.strides != computeRowMajorStrides(view.shape))
-              return mlir::failure();
             auto resultType = mlir::cast<mlir::MemRefType>(collapse.getResult().getType());
+            auto resultStrides = getStaticMemRefTypeStrides(resultType);
+            if (failed(resultStrides))
+              return mlir::failure();
             view.shape.assign(resultType.getShape().begin(), resultType.getShape().end());
-            view.strides = computeRowMajorStrides(view.shape);
+            view.strides = std::move(*resultStrides);
             continue;
           }
 
           if (auto expand = mlir::dyn_cast<mlir::memref::ExpandShapeOp>(viewOp)) {
-            if (view.strides != computeRowMajorStrides(view.shape))
-              return mlir::failure();
             auto resultType = mlir::cast<mlir::MemRefType>(expand.getResult().getType());
+            auto resultStrides = getStaticMemRefTypeStrides(resultType);
+            if (failed(resultStrides))
+              return mlir::failure();
             view.shape.assign(resultType.getShape().begin(), resultType.getShape().end());
-            view.strides = computeRowMajorStrides(view.shape);
+            view.strides = std::move(*resultStrides);
+            continue;
           }
+
+          if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(viewOp)) {
+            auto resultType = mlir::cast<mlir::MemRefType>(castOp.getResult().getType());
+            auto resultStrides = getStaticMemRefTypeStrides(resultType);
+            if (failed(resultStrides))
+              return mlir::failure();
+            view.shape.assign(resultType.getShape().begin(), resultType.getShape().end());
+            view.strides = std::move(*resultStrides);
+            continue;
+          }
+
+          return mlir::failure();
         }
 
         auto resolvedOffset = offsetExpr.evaluate(knowledge);
@@ -234,18 +265,26 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static
         return view;
       }
 
-      if (mlir::isa<mlir::memref::SubViewOp, mlir::memref::CollapseShapeOp, mlir::memref::ExpandShapeOp>(defOp)) {
+      if (auto subview = mlir::dyn_cast<mlir::memref::SubViewOp>(defOp)) {
         viewOps.push_back(defOp);
-        if (auto subview = mlir::dyn_cast<mlir::memref::SubViewOp>(defOp))
-          current = subview.getSource();
-        else if (auto collapse = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(defOp))
-          current = collapse.getSrc();
-        else
-          current = mlir::cast<mlir::memref::ExpandShapeOp>(defOp).getSrc();
+        current = subview.getSource();
+        continue;
+      }
+
+      if (auto collapse = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(defOp)) {
+        viewOps.push_back(defOp);
+        current = collapse.getSrc();
+        continue;
+      }
+
+      if (auto expand = mlir::dyn_cast<mlir::memref::ExpandShapeOp>(defOp)) {
+        viewOps.push_back(defOp);
+        current = expand.getSrc();
         continue;
       }
 
       if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(defOp)) {
+        viewOps.push_back(defOp);
         current = castOp.getSource();
         continue;
       }
@@ -253,6 +292,11 @@ resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const Static
       return mlir::failure();
     }
 
+    if (mlir::Value loopAlias = resolveLoopCarriedAlias(current, knowledge); loopAlias && loopAlias != current) {
+      current = loopAlias;
+      continue;
+    }
+
     auto weightIndex = resolveWeightIndex(weightOwner, current);
     if (!weightIndex)
       return mlir::failure();

src/PIM/Common/LabeledList.hpp

@@ -1,315 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/ilist_node.h"
-#include "llvm/ADT/simple_ilist.h"
-
-#include <cassert>
-#include <iterator>
-#include <limits>
-#include <type_traits>
-
-namespace onnx_mlir {
-
-template <typename NodeT>
-class LabeledList;
-
-template <typename NodeT>
-class LabeledListNode : public llvm::ilist_node<NodeT> {
-  friend class LabeledList<NodeT>;
-
-public:
-  using Label = uint64_t;
-
-  LabeledListNode() = default;
-  LabeledListNode(const LabeledListNode&) = delete;
-  LabeledListNode(LabeledListNode&&) = default;
-  LabeledListNode& operator=(LabeledListNode&&) = delete;
-
-  ~LabeledListNode() { assert(owner_ == nullptr && "destroying a linked LabeledListNode"); }
-
-  bool isLinked() const { return owner_ != nullptr; }
-  Label getOrderLabel() const { return label; }
-
-  friend bool operator<(const LabeledListNode& lft, const LabeledListNode& rgt) { return lft.label < rgt.label; }
-
-private:
-  const void* owner_ = nullptr;
-  Label label = 0;
-};
-
-template <typename NodeT>
-class LabeledList {
-
-  using Label = typename NodeT::Label;
-
-  static constexpr Label kLowerSentinel = 0;
-  static constexpr Label kUpperSentinel = std::numeric_limits<Label>::max();
-  static constexpr Label kRelabelGap = 2;
-
-public:
-  using List = llvm::simple_ilist<NodeT>;
-  using Iterator = typename List::iterator;
-  using RIterator = typename List::reverse_iterator;
-  using ConstIterator = typename List::const_iterator;
-
-  LabeledList() = default;
-  LabeledList(const LabeledList&) = delete;
-  LabeledList& operator=(const LabeledList&) = delete;
-  LabeledList(LabeledList&&) = delete;
-  LabeledList& operator=(LabeledList&&) = delete;
-
-  ~LabeledList() { clear(); }
-
-  bool empty() const { return size_ == 0; }
-  size_t size() const { return size_; }
-
-  NodeT* front() { return empty() ? nullptr : &nodes_.front(); }
-  const NodeT* front() const { return empty() ? nullptr : &nodes_.front(); }
-
-  NodeT* back() { return empty() ? nullptr : &nodes_.back(); }
-  const NodeT* back() const { return empty() ? nullptr : &nodes_.back(); }
-
-  static NodeT* previous(NodeT* node) {
-    if (!node || !owner(node))
-      return nullptr;
-    auto* list = owner(node);
-    auto it = node->getIterator();
-    if (it == list->nodes_.begin())
-      return nullptr;
-    return &*std::prev(it);
-  }
-
-  static const NodeT* previous(const NodeT* node) {
-    if (!node || !owner(node))
-      return nullptr;
-    const auto* list = owner(node);
-    auto it = const_cast<NodeT*>(node)->getIterator();
-    if (it == list->nodes_.begin())
-      return nullptr;
-    return &*std::prev(it);
-  }
-
-  static NodeT* next(NodeT* node) {
-    if (!node || !owner(node))
-      return nullptr;
-    auto* list = owner(node);
-    auto it = std::next(node->getIterator());
-    if (it == list->nodes_.end())
-      return nullptr;
-    return &*it;
-  }
-
-  static const NodeT* next(const NodeT* node) {
-    if (!node || !owner(node))
-      return nullptr;
-    const auto* list = owner(node);
-    auto it = std::next(const_cast<NodeT*>(node)->getIterator());
-    if (it == list->nodes_.end())
-      return nullptr;
-    return &*it;
-  }
-
-  bool contains(const NodeT* node) const { return node && node->owner_ == this; }
-
-  Label getOrderLabel(const NodeT* node) const {
-    assert(contains(node) && "node must belong to this list");
-    return node->label;
-  }
-
-  bool comesBefore(const NodeT* lhs, const NodeT* rhs) const {
-    assert(contains(lhs) && contains(rhs) && "nodes must belong to this list");
-    return lhs->label < rhs->label;
-  }
-
-  void pushFront(NodeT* node) { insertBefore(front(), node); }
-
-  void pushBack(NodeT* node) { insertBefore(nullptr, node); }
-
-  void insertBefore(NodeT* nextNode, NodeT* node) {
-    assert(node && "cannot insert a null node");
-    assert(!node->owner_ && "node is already linked");
-    assert(nextNode == nullptr || contains(nextNode));
-
-    Iterator nextIt = nextNode ? getIteratorFor(nextNode) : nodes_.end();
-    nodes_.insert(nextIt, *node);
-    node->owner_ = this;
-    ++size_;
-    assignLabel(getIteratorFor(node));
-  }
-
-  void insertAfter(NodeT* prevNode, NodeT* node) {
-    assert(prevNode == nullptr || contains(prevNode));
-    if (prevNode == nullptr)
-      insertBefore(front(), node);
-    else
-      insertBefore(next(prevNode), node);
-  }
-
-  void remove(NodeT* node) {
-    assert(contains(node) && "node must belong to this list");
-    nodes_.remove(*node);
-    node->owner_ = nullptr;
-    node->label = 0;
-    --size_;
-  }
-
-  void moveBefore(NodeT* node, NodeT* nextNode) {
-    assert(contains(node) && "node must belong to this list");
-    assert(nextNode == nullptr || contains(nextNode));
-
-    Iterator nodeIt = getIteratorFor(node);
-    Iterator nextIt = nextNode ? getIteratorFor(nextNode) : nodes_.end();
-    if (nodeIt == nextIt || std::next(nodeIt) == nextIt)
-      return;
-
-    nodes_.splice(nextIt, nodes_, nodeIt);
-    assignLabel(getIteratorFor(node));
-  }
-
-  void moveAfter(NodeT* node, NodeT* prevNode) {
-    assert(contains(node) && "node must belong to this list");
-    assert(prevNode == nullptr || contains(prevNode));
-
-    Iterator nextIt = prevNode ? std::next(getIteratorFor(prevNode)) : nodes_.begin();
-    if (getIteratorFor(node) == nextIt)
-      return;
-    moveBefore(node, nextIt == nodes_.end() ? nullptr : &*nextIt);
-  }
-
-  void clear() {
-    while (!nodes_.empty()) {
-      NodeT* node = &nodes_.front();
-      node->owner_ = nullptr;
-      node->label = 0;
-      nodes_.remove(*node);
-    }
-    size_ = 0;
-  }
-
-  Iterator begin() { return nodes_.begin(); }
-  Iterator end() { return nodes_.end(); }
-
-  RIterator rbegin() { return nodes_.rbegin(); }
-  RIterator rend() { return nodes_.rend(); }
-
-private:
-  static const LabeledList* owner(const NodeT* node) { return static_cast<const LabeledList*>(node->owner_); }
-  static LabeledList* owner(NodeT* node) { return static_cast<LabeledList*>(const_cast<void*>(node->owner_)); }
-
-  static Label lowerLabel(const NodeT* node) { return node ? node->label : kLowerSentinel; }
-  static Label upperLabel(const NodeT* node) { return node ? node->label : kUpperSentinel; }
-
-  static Label labelGap(Label lower, Label upper) {
-    assert(lower < upper && "labels must be strictly ordered");
-    return upper - lower;
-  }
-
-  static bool hasMidpoint(Label lower, Label upper) { return labelGap(lower, upper) > 1; }
-
-  static bool hasRelabelSlack(Label lower, Label upper, size_t nodeCount) {
-    Label gap = labelGap(lower, upper);
-    return gap / static_cast<Label>(nodeCount + 1) >= kRelabelGap;
-  }
-
-  Iterator getIteratorFor(NodeT* node) { return node->getIterator(); }
-  ConstIterator getiteratorFor(const NodeT* node) const { return node->getIterator(); }
-
-  NodeT* previousNode(Iterator it) {
-    if (it == nodes_.begin())
-      return nullptr;
-    return &*std::prev(it);
-  }
-
-  const NodeT* previousNode(ConstIterator it) const {
-    if (it == nodes_.begin())
-      return nullptr;
-    return &*std::prev(it);
-  }
-
-  NodeT* nextNode(Iterator it) {
-    ++it;
-    if (it == nodes_.end())
-      return nullptr;
-    return &*it;
-  }
-
-  const NodeT* nextNode(ConstIterator it) const {
-    ++it;
-    if (it == nodes_.end())
-      return nullptr;
-    return &*it;
-  }
-
-  void assignLabel(Iterator it) {
-    Label lower = lowerLabel(previousNode(it));
-    Label upper = upperLabel(nextNode(it));
-    if (hasMidpoint(lower, upper)) {
-      (*it).label = lower + static_cast<Label>(labelGap(lower, upper) / 2);
-      return;
-    }
-
-    relabelAround(it);
-  }
-
-  void relabelAround(Iterator center) {
-    size_t targetCount = 1;
-    while (true) {
-      Iterator left = center;
-      Iterator right = center;
-      size_t actualCount = 1;
-      expandWindow(center, targetCount, left, right, actualCount);
-
-      Label lower = lowerLabel(previousNode(left));
-      Label upper = upperLabel(nextNode(right));
-      if (hasRelabelSlack(lower, upper, actualCount)) {
-        relabelWindow(left, actualCount, lower, upper);
-        return;
-      }
-
-      if (left == nodes_.begin() && nextNode(right) == nullptr) {
-        assert(hasRelabelSlack(lower, upper, actualCount) && "label space exhausted");
-        relabelWindow(left, actualCount, lower, upper);
-        return;
-      }
-
-      targetCount *= 2;
-    }
-  }
-
-  void expandWindow(Iterator center, size_t targetCount, Iterator& left, Iterator& right, size_t& actualCount) {
-    left = center;
-    right = center;
-    actualCount = 1;
-
-    while (actualCount < targetCount && (left != nodes_.begin() || nextNode(right) != nullptr)) {
-      if (left != nodes_.begin()) {
-        --left;
-        ++actualCount;
-        if (actualCount == targetCount)
-          break;
-      }
-      if (nextNode(right) != nullptr) {
-        ++right;
-        ++actualCount;
-      }
-    }
-  }
-
-  void relabelWindow(Iterator left, size_t nodeCount, Label lower, Label upper) {
-    assert(nodeCount > 0 && "relabel window must not be empty");
-    Label step = labelGap(lower, upper) / static_cast<Label>(nodeCount + 1);
-    assert(step >= 1 && "relabel step must be positive");
-
-    Iterator it = left;
-    for (size_t index = 1; index <= nodeCount; ++index) {
-      (*it).label = lower + step * index;
-      ++it;
-    }
-  }
-
-  List nodes_;
-  size_t size_ = 0;
-};
-
-} // namespace onnx_mlir

src/PIM/Common/PimCommon.hpp

@@ -15,6 +15,7 @@
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/EntryPointUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/IndexingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"

src/PIM/Common/Support/DebugDump.cpp

@@ -7,18 +7,26 @@
 
 namespace onnx_mlir {
 
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
+std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::StringRef destination,  llvm::StringRef extension) {
   std::string outputDir = getOutputDir();
   if (outputDir.empty())
+    return {};
+
+  std::string dialectsDir = (outputDir + destination).str();
+  createDirectory(dialectsDir);
+  return std::fstream(dialectsDir + "/" + name + "." + extension.str(), std::ios::out);
+}
+
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified) {
+  std::fstream file = openDialectDumpFileWithExtension(name, "/dialects", "mlir");
+  if (!file.is_open())
     return;
 
-  std::string dialectsDir = outputDir + "/dialects";
-  createDirectory(dialectsDir);
-
-  std::fstream file(dialectsDir + "/" + name + ".mlir", std::ios::out);
   llvm::raw_os_ostream os(file);
   mlir::OpPrintingFlags flags;
-  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  flags.elideLargeElementsAttrs().enableDebugInfo(false, false);
+  if (assumeVerified)
+    flags.assumeVerified();
   moduleOp.print(os, flags);
   os.flush();
   file.close();

src/PIM/Common/Support/DebugDump.hpp

@@ -1,13 +1,18 @@
 #pragma once
 
 #include "mlir/IR/BuiltinOps.h"
+#include "llvm/ADT/StringRef.h"
 
+#include <fstream>
 #include <string>
 
 namespace onnx_mlir {
 
 /// Emits a MLIR snapshot under the current compiler output
 /// directory for pass-level debugging.
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name);
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified = false);
+
+/// Opens a file under the same dialect dump directory used by dumpModule.
+std::fstream openDialectDumpFileWithExtension(const std::string& name,llvm::StringRef destination = "/dialects",  llvm::StringRef extension = "mlir");
 
 } // namespace onnx_mlir

src/PIM/Common/Support/Diagnostics.hpp

@@ -28,6 +28,8 @@ struct CappedDiagnosticReporter {
       op->emitError() << "suppressed " << (numFailures - maxReportedFailures) << " additional " << failureDescription;
   }
 
+  void noteFailures(int64_t count) { numFailures += count; }
+
   bool hasFailure() const { return numFailures != 0; }
 
 private:

src/PIM/Common/Support/ReportUtils.cpp

@@ -5,14 +5,30 @@
 
 namespace onnx_mlir {
 
-std::fstream openReportFile(const std::string& name) {
+std::fstream openReportFileWithExtension(const std::string& name, llvm::StringRef extension) {
   std::string outputDir = getOutputDir();
   if (outputDir.empty())
     return {};
 
   std::string reportsDir = outputDir + "/reports";
   createDirectory(reportsDir);
-  return std::fstream(reportsDir + "/" + name + ".txt", std::ios::out);
+  return std::fstream(reportsDir + "/" + name + "." + extension.str(), std::ios::out);
+}
+
+std::fstream openReportFile(const std::string& name) { return openReportFileWithExtension(name, "txt"); }
+
+std::fstream openAppendedReportFileWithExtension(const std::string& name, llvm::StringRef extension) {
+  std::string outputDir = getOutputDir();
+  if (outputDir.empty())
+    return {};
+
+  std::string reportsDir = outputDir + "/reports";
+  createDirectory(reportsDir);
+  return std::fstream(reportsDir + "/" + name + "." + extension.str(), std::ios::out | std::ios::app);
+}
+
+std::fstream openAppendedReportFile(const std::string& name) {
+  return openAppendedReportFileWithExtension(name, "txt");
 }
 
 std::string formatReportMemory(uint64_t bytes) {

src/PIM/Common/Support/ReportUtils.hpp

@@ -11,6 +11,9 @@
 namespace onnx_mlir {
 
 std::fstream openReportFile(const std::string& name);
+std::fstream openReportFileWithExtension(const std::string& name, llvm::StringRef extension);
+std::fstream openAppendedReportFile(const std::string& name);
+std::fstream openAppendedReportFileWithExtension(const std::string& name, llvm::StringRef extension);
 std::string formatReportMemory(uint64_t bytes);
 
 struct ReportField {

src/PIM/Compiler/CMakeLists.txt

@@ -17,6 +17,7 @@ add_pim_library(OMPimCompilerUtils
   PimCompilerUtils.cpp
   PimArtifactWriter.cpp
   PimCodeGen.cpp
+  PimMemoryLiveness.cpp
   PimWeightEmitter.cpp
 
   EXCLUDE_FROM_OM_LIBS
@@ -30,7 +31,6 @@ add_pim_library(OMPimCompilerUtils
   OMPimBufferization
   OMPimMemoryCoalescing
   OMPimHostConstantFolding
-  OMPimHostConstantMaterialization
   OMPimVerification
   OMPimPasses
   OMONNXToSpatial

src/PIM/Compiler/PimCodeGen.cpp

@@ -2,7 +2,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -27,20 +26,24 @@
 #include <fstream>
 #include <limits>
 #include <memory>
+#include <numeric>
 #include <string>
 #include <utility>
 
 #include "Common/IR/CompactAsmUtils.hpp"
 #include "Common/PimCommon.hpp"
+#include "Common/Support/Diagnostics.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
 #include "Common/Support/ReportUtils.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
 #include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
 #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -96,6 +99,51 @@ static int32_t getVectorByteSizeOrCrash(ShapedType type) {
   return pim::checkedI32OrCrash(*byteSize, "vector byte size");
 }
 
+static Operation* getDiagnosticAnchor(mlir::Value value) {
+  if (Operation* definingOp = value.getDefiningOp())
+    return definingOp;
+  if (auto blockArg = dyn_cast<BlockArgument>(value))
+    return blockArg.getOwner()->getParentOp();
+  return nullptr;
+}
+
+// PIM instruction immediates are serialized as signed int32_t fields today
+// (`sldi` goes through checkedI32OrCrash), so local addresses must stay within
+// the non-negative int32_t range.
+static constexpr size_t kPimAddressLimit = static_cast<size_t>(std::numeric_limits<int32_t>::max());
+
+static FailureOr<size_t> checkedAlignTo(size_t value, size_t alignment, Operation* anchor, StringRef fieldName) {
+  if (alignment == 0)
+    return value;
+  size_t remainder = value % alignment;
+  if (remainder == 0)
+    return value;
+  return pim::checkedAdd(value, alignment - remainder, anchor, fieldName);
+}
+
+static void printMemoryOverflowDiagnostic(mlir::Value value,
+                                          const MemoryValueKey& key,
+                                          size_t requestedSize,
+                                          size_t currentFirstAvailableAddress,
+                                          size_t alignedEndAddress) {
+  llvm::errs() << "PIM local memory allocation overflow\n";
+  llvm::errs() << "Requested allocation size: " << requestedSize << " bytes\n";
+  llvm::errs() << "Current firstAvailableAddress: " << currentFirstAvailableAddress << "\n";
+  llvm::errs() << "Aligned end address: " << alignedEndAddress << "\n";
+  llvm::errs() << "Address limit: " << kPimAddressLimit << " (signed int32_t immediate range)\n";
+  if (key.lane)
+    llvm::errs() << "Lane: " << *key.lane << "\n";
+  llvm::errs() << "Value: ";
+  value.print(llvm::errs());
+  llvm::errs() << "\n";
+  llvm::errs() << "Value type: " << value.getType() << "\n";
+  if (Operation* definingOp = value.getDefiningOp()) {
+    llvm::errs() << "Defining op:\n";
+    definingOp->print(llvm::errs());
+    llvm::errs() << "\n";
+  }
+}
+
 } // namespace
 
 MemEntry* PimMemory::gatherMemEntry(mlir::Value value, std::optional<unsigned> lane) {
@@ -123,19 +171,29 @@ void PimMemory::allocateGatheredMemory() {
 
 void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind) {
   memEntry.address = firstAvailableAddress;
-  firstAvailableAddress += memEntry.size;
-  // Alignment
-  if (size_t remainder = firstAvailableAddress % minAlignment)
-    firstAvailableAddress += minAlignment - remainder;
+  Operation* anchor = getDiagnosticAnchor(key.value);
+  auto checkedEnd = pim::checkedAdd(memEntry.address, memEntry.size, anchor, "local memory end");
+  FailureOr<size_t> checkedAlignedEnd = failure();
+  if (succeeded(checkedEnd))
+    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
+  bool startFits = memEntry.address <= kPimAddressLimit;
+  bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit;
+  bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit;
+  if (!startFits || !endFits || !alignedEndFits) {
+    printMemoryOverflowDiagnostic(key.value,
+                                  key,
+                                  memEntry.size,
+                                  firstAvailableAddress,
+                                  succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+    llvm_unreachable("PIM local memory allocation overflow");
+  }
+  firstAvailableAddress = *checkedAlignedEnd;
 
   ownedMemEntriesMap[key] = memEntry;
   globalMemEntriesMap[key] = memEntry;
 
   switch (reportKind) {
-  case MemoryReportKind::Alloca:
-    ++reportRow.numAlloca;
-    reportRow.sizeAlloca += memEntry.size;
-    break;
+  case MemoryReportKind::Alloca: break;
   case MemoryReportKind::Global:
     ++reportRow.numGlobal;
     reportRow.sizeGlobal += memEntry.size;
@@ -145,6 +203,34 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
   }
 }
 
+PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key) {
+  PhysicalSlotInfo slot;
+  slot.id = nextPhysicalSlotId++;
+  slot.address = firstAvailableAddress;
+  slot.size = slotSize;
+
+  Operation* anchor = getDiagnosticAnchor(key.value);
+  auto checkedEnd = pim::checkedAdd(slot.address, slot.size, anchor, "local memory end");
+  FailureOr<size_t> checkedAlignedEnd = failure();
+  if (succeeded(checkedEnd))
+    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
+  bool startFits = slot.address <= kPimAddressLimit;
+  bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit;
+  bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit;
+  if (!startFits || !endFits || !alignedEndFits) {
+    printMemoryOverflowDiagnostic(key.value,
+                                  key,
+                                  slot.size,
+                                  firstAvailableAddress,
+                                  succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+    llvm_unreachable("PIM local memory allocation overflow");
+  }
+
+  firstAvailableAddress = *checkedAlignedEnd;
+  localPhysicalSlots.push_back(slot);
+  return slot;
+}
+
 void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
   SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
   SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
@@ -184,9 +270,71 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
 }
 
 void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
-  op->walk([&](memref::AllocOp allocOp) { gatherMemEntry(allocOp, lane); });
+  auto intervals = buildLocalAllocIntervals(op, lane);
+  SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);
 
-  allocateGatheredMemory();
+  SmallVector<size_t> slotOrder(plannedSlots.size());
+  std::iota(slotOrder.begin(), slotOrder.end(), 0);
+  llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) {
+    const PlannedPhysicalSlot& lhs = plannedSlots[lhsIndex];
+    const PlannedPhysicalSlot& rhs = plannedSlots[rhsIndex];
+    if (lhs.requiredSize != rhs.requiredSize)
+      return lhs.requiredSize > rhs.requiredSize;
+    return lhs.id < rhs.id;
+  });
+
+  SmallVector<bool, 16> usedExistingSlots(localPhysicalSlots.size(), false);
+  for (size_t slotIndex : slotOrder) {
+    PlannedPhysicalSlot& slot = plannedSlots[slotIndex];
+    size_t bestExistingIndex = std::numeric_limits<size_t>::max();
+    auto bestKey = std::tuple<size_t, size_t, size_t>(
+      std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
+
+    for (size_t existingIndex = 0; existingIndex < localPhysicalSlots.size(); ++existingIndex) {
+      if (usedExistingSlots[existingIndex])
+        continue;
+      const PhysicalSlotInfo& existingSlot = localPhysicalSlots[existingIndex];
+      if (existingSlot.size < slot.requiredSize)
+        continue;
+      auto candidateKey =
+        std::tuple<size_t, size_t, size_t>(existingSlot.size - slot.requiredSize, existingSlot.size, existingSlot.id);
+      if (candidateKey < bestKey) {
+        bestKey = candidateKey;
+        bestExistingIndex = existingIndex;
+      }
+    }
+
+    if (bestExistingIndex != std::numeric_limits<size_t>::max()) {
+      const PhysicalSlotInfo& existingSlot = localPhysicalSlots[bestExistingIndex];
+      slot.id = existingSlot.id;
+      slot.address = existingSlot.address;
+      slot.size = existingSlot.size;
+      usedExistingSlots[bestExistingIndex] = true;
+    }
+    else {
+      PhysicalSlotInfo newSlot = allocatePhysicalSlot(slot.requiredSize, intervals[slot.intervalIndices.front()].key);
+      slot.id = newSlot.id;
+      slot.address = newSlot.address;
+      slot.size = newSlot.size;
+      usedExistingSlots.push_back(true);
+    }
+
+    for (size_t intervalIndex : slot.intervalIndices) {
+      LocalAllocInterval& interval = intervals[intervalIndex];
+      interval.physicalSlotId = slot.id;
+      interval.assignedAddress = slot.address;
+      interval.physicalSlotSize = slot.size;
+      MemEntry memEntry {slot.address, interval.size};
+      ownedMemEntriesMap[interval.key] = memEntry;
+      globalMemEntriesMap[interval.key] = memEntry;
+    }
+  }
+
+  if (pimMemoryReport != PimMemoryReportNone) {
+    MemoryPlanArtifacts artifacts =
+      buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport);
+    livenessArtifacts.textReport += artifacts.textReport;
+  }
 }
 
 static void printHostMemoryReportRow(raw_ostream& os, const MemoryReportRow& row) {
@@ -226,7 +374,14 @@ static MemoryReportRow addMemoryReportRows(const MemoryReportRow& lhs, const Mem
   return result;
 }
 
-MemoryReportRow PimMemory::getReportRow() const { return reportRow; }
+MemoryReportRow PimMemory::getReportRow() const {
+  MemoryReportRow row = reportRow;
+  row.numAlloca = localPhysicalSlots.size();
+  row.sizeAlloca = 0;
+  for (const PhysicalSlotInfo& slot : localPhysicalSlots)
+    row.sizeAlloca += slot.size;
+  return row;
+}
 
 void PimMemory::remove(mlir::Value val) {
   for (auto it = ownedMemEntriesMap.begin(); it != ownedMemEntriesMap.end();)
@@ -259,31 +414,35 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
                                              const StaticValueKnowledge& knowledge,
                                              std::optional<unsigned> lane) const {
   value = resolveCachedAlias(value, knowledge);
-  auto compiledIt = compiledAddressExprs.find(value);
-  if (compiledIt == compiledAddressExprs.end()) {
-    auto compiledExpr = compileContiguousAddressExpr(value);
-    if (failed(compiledExpr)) {
-      errs() << "Failed to compile contiguous address for value: ";
+
+  FailureOr<ResolvedContiguousAddress> resolvedAddress = resolveContiguousAddress(value, knowledge);
+  if (failed(resolvedAddress)) {
+    auto compiledIt = compiledAddressExprs.find(value);
+    if (compiledIt == compiledAddressExprs.end()) {
+      auto compiledExpr = compileContiguousAddressExpr(value);
+      if (failed(compiledExpr)) {
+        errs() << "Failed to compile contiguous address for value: ";
+        value.print(errs());
+        errs() << " : " << value.getType();
+        errs() << "\n";
+        llvm_unreachable("Failed to compile contiguous address");
+      }
+      compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
+    }
+
+    resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
+    if (failed(resolvedAddress)) {
+      errs() << "Failed to evaluate contiguous address for value: ";
       value.print(errs());
       errs() << " : " << value.getType();
       errs() << "\n";
-      llvm_unreachable("Failed to compile contiguous address");
+      if (auto* definingOp = value.getDefiningOp()) {
+        errs() << "Defining op:\n";
+        definingOp->print(errs());
+        errs() << "\n";
+      }
+      llvm_unreachable("Failed to resolve contiguous address");
     }
-    compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
-  }
-
-  auto resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
-  if (failed(resolvedAddress)) {
-    errs() << "Failed to evaluate contiguous address for value: ";
-    value.print(errs());
-    errs() << " : " << value.getType();
-    errs() << "\n";
-    if (auto* definingOp = value.getDefiningOp()) {
-      errs() << "Defining op:\n";
-      definingOp->print(errs());
-      errs() << "\n";
-    }
-    llvm_unreachable("Failed to resolve contiguous address");
   }
 
   MemoryValueKey key = getMemoryValueKey(resolvedAddress->base, lane);
@@ -433,13 +592,37 @@ void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instructio
   ++emittedInstructionCount;
   if (coreJsonStream)
     *coreJsonStream << json::Value(pim_binary::makeInstructionJson(instruction)) << ',';
+  updateScalarRegisterCache(instruction);
+}
+
+void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const {
+  switch (instruction.opcode) {
+  case pim_binary::Opcode::sldi:
+    scalarRegisterValues[instruction.rd] = instruction.r2OrImm;
+    break;
+  case pim_binary::Opcode::sld:
+  case pim_binary::Opcode::sadd:
+  case pim_binary::Opcode::ssub:
+  case pim_binary::Opcode::smul:
+  case pim_binary::Opcode::saddi:
+  case pim_binary::Opcode::smuli:
+    scalarRegisterValues[instruction.rd].reset();
+    break;
+  default:
+    break;
+  }
 }
 
 void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const {
+  auto registerIndex = pim::checkedU8OrCrash(registerNumber, "register number");
+  auto immediateValue = pim::checkedI32OrCrash(immediate, "register immediate");
+  if (scalarRegisterValues[registerIndex] == immediateValue)
+    return;
+
   pim_binary::InstructionRecord instruction;
   instruction.opcode = pim_binary::Opcode::sldi;
-  instruction.rd = static_cast<uint8_t>(registerNumber);
-  instruction.r2OrImm = pim::checkedI32OrCrash(immediate, "register immediate");
+  instruction.rd = registerIndex;
+  instruction.r2OrImm = immediateValue;
   emitInstruction(instruction);
 }
 
@@ -842,11 +1025,44 @@ static SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
 }
 
 struct CoreEmissionResult {
+  static constexpr size_t kMaxStoredCodegenDiagnostics = 8;
+
+  struct DiagnosticRecord {
+    Operation* op = nullptr;
+    std::string message;
+  };
+
   OnnxMlirCompilerErrorCodes status = CompilerSuccess;
   MemoryReportRow reportRow;
   llvm::SmallVector<ResolvedWeightView, 8> usedWeights;
+  MemoryPlanArtifacts livenessArtifacts;
+  llvm::SmallVector<DiagnosticRecord, kMaxStoredCodegenDiagnostics> diagnostics;
+  size_t diagnosticCount = 0;
+
+  void recordDiagnostic(Operation* op, StringRef message) {
+    ++diagnosticCount;
+    if (diagnostics.size() < kMaxStoredCodegenDiagnostics)
+      diagnostics.push_back({op, message.str()});
+  }
 };
 
+static StaticValueKnowledge seedCoreCodegenKnowledge(pim::PimCoreOp coreOp) {
+  StaticValueKnowledge knowledge;
+  for (auto [index, weight] : llvm::enumerate(coreOp.getWeights()))
+    knowledge.aliases[coreOp.getWeightArgument(index)] = weight;
+  return knowledge;
+}
+
+static StaticValueKnowledge seedCoreBatchCodegenKnowledge(pim::PimCoreBatchOp coreBatchOp, unsigned lane) {
+  StaticValueKnowledge knowledge;
+  knowledge.indexValues[coreBatchOp.getLaneArgument()] = lane;
+  for (auto [index, weight] : llvm::enumerate(coreBatchOp.getWeights()))
+    knowledge.aliases[coreBatchOp.getWeightArgument(index)] = weight;
+  for (auto [index, input] : llvm::enumerate(coreBatchOp.getInputs()))
+    knowledge.aliases[coreBatchOp.getInputArgument(index)] = input;
+  return knowledge;
+}
+
 template <typename MapTy>
 class ScopedMapBindings {
   using KeyTy = typename MapTy::key_type;
@@ -902,7 +1118,9 @@ enum class CompiledCoreOpKind : uint8_t {
 struct CompiledCoreNode {
   enum class Kind : uint8_t {
     Op,
-    Loop
+    Loop,
+    If,
+    IndexSwitch
   };
 
   Kind kind = Kind::Op;
@@ -911,7 +1129,13 @@ struct CompiledCoreNode {
   CompiledIndexExpr lowerBound;
   CompiledIndexExpr upperBound;
   CompiledIndexExpr step;
+  CompiledIndexExpr condition;
   std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> loopBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> thenBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> elseBody;
+  llvm::SmallVector<int64_t> caseValues;
+  llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
 };
 
 static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
@@ -989,6 +1213,53 @@ compileCoreEmissionPlan(Block& block, Operation* weightOwner, llvm::SmallVectorI
       continue;
     }
 
+    if (auto ifOp = dyn_cast<mlir::scf::IfOp>(op)) {
+      auto condition = compileIndexExpr(ifOp.getCondition());
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+
+      CompiledCoreNode ifNode;
+      ifNode.kind = CompiledCoreNode::Kind::If;
+      ifNode.op = ifOp.getOperation();
+      ifNode.condition = *condition;
+      ifNode.thenBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), weightOwner, *ifNode.thenBody)))
+        return failure();
+      ifNode.elseBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (!ifOp.getElseRegion().empty())
+        if (failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), weightOwner, *ifNode.elseBody)))
+          return failure();
+      plan.push_back(std::move(ifNode));
+      continue;
+    }
+
+    if (auto switchOp = dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = compileIndexExpr(switchOp.getArg());
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        return failure();
+      }
+      CompiledCoreNode switchNode;
+      switchNode.kind = CompiledCoreNode::Kind::IndexSwitch;
+      switchNode.op = switchOp.getOperation();
+      switchNode.condition = *selector;
+      llvm::append_range(switchNode.caseValues, switchOp.getCases());
+      for (mlir::Region& region : switchOp.getCaseRegions()) {
+        auto body = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+        if (failed(compileCoreEmissionPlan(region.front(), weightOwner, *body)))
+          return failure();
+        switchNode.caseBodies.push_back(std::move(body));
+      }
+      switchNode.defaultBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(
+            switchOp.getDefaultRegion().front(), weightOwner, *switchNode.defaultBody)))
+        return failure();
+      plan.push_back(std::move(switchNode));
+      continue;
+    }
+
     auto opKind = classifyCompiledCoreOpKind(op);
     if (failed(opKind)) {
       InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'";
@@ -1051,6 +1322,51 @@ static LogicalResult executeCompiledCorePlan(
       continue;
     }
 
+    if (node.kind == CompiledCoreNode::Kind::If) {
+      auto condition = node.condition.evaluate(knowledge);
+      auto ifOp = cast<mlir::scf::IfOp>(node.op);
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+
+      const auto& selectedBody = *condition != 0 ? node.thenBody : node.elseBody;
+      if (selectedBody && failed(executeCompiledCorePlan(*selectedBody,
+                                                         coreCodeGen,
+                                                         knowledge,
+                                                         resolveWeightSlot,
+                                                         processedOperations,
+                                                         batchLane,
+                                                         batchLaneCount)))
+        return failure();
+      continue;
+    }
+
+    if (node.kind == CompiledCoreNode::Kind::IndexSwitch) {
+      auto selector = node.condition.evaluate(knowledge);
+      auto switchOp = cast<mlir::scf::IndexSwitchOp>(node.op);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        return failure();
+      }
+      const llvm::SmallVectorImpl<CompiledCoreNode>* selectedBody = node.defaultBody.get();
+      mlir::Region* selectedRegion = &switchOp.getDefaultRegion();
+      for (auto [index, caseValue] : llvm::enumerate(node.caseValues))
+        if (caseValue == *selector) {
+          selectedBody = node.caseBodies[index].get();
+          selectedRegion = &switchOp.getCaseRegions()[index];
+          break;
+        }
+      if (failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge,
+                                         resolveWeightSlot, processedOperations,
+                                         batchLane, batchLaneCount)))
+        return failure();
+      auto yield = cast<mlir::scf::YieldOp>(selectedRegion->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
     switch (node.opKind) {
     case CompiledCoreOpKind::Load:
       coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);
@@ -1151,6 +1467,36 @@ static int64_t codeGenCoreOps(
   return failed(result) ? -1 : static_cast<int64_t>(processedOperations);
 }
 
+static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
+  std::string outputCorePath =
+    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
+  std::error_code errorCode;
+  raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
+  if (errorCode) {
+    errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
+    return InvalidOutputFileAccess;
+  }
+
+  pim_binary::writeHeader(coreBinaryStream);
+  pim_binary::patchInstructionCount(coreBinaryStream, 0);
+  coreBinaryStream.close();
+
+  if (!pimEmitJson.getValue())
+    return CompilerSuccess;
+
+  std::string outputCoreJsonPath =
+    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
+  errorCode = std::error_code();
+  raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
+  if (errorCode) {
+    errs() << "Error while opening core json file `" << outputCoreJsonPath << "`: " << errorCode.message() << '\n';
+    return InvalidOutputFileAccess;
+  }
+  coreJsonStream << "[]";
+  coreJsonStream.close();
+  return CompilerSuccess;
+}
+
 OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::string& outputDirPath) {
   if (!outputDirPath.empty()) {
     if (auto error = sys::fs::create_directory(outputDirPath)) {
@@ -1267,7 +1613,20 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
                                  const StaticValueKnowledge& knowledge) -> llvm::FailureOr<unsigned> {
       auto weightView = onnx_mlir::resolveWeightView(job.coreLikeOp, vmmOp.getWeight(), knowledge);
       if (failed(weightView)) {
-        vmmOp.emitOpError("requires a statically resolvable dense global weight view during PIM codegen");
+        std::string message;
+        llvm::raw_string_ostream os(message);
+        os << "requires a statically resolvable dense global weight view during PIM codegen; weight="
+           << vmmOp.getWeight() << " type=" << vmmOp.getWeight().getType();
+        result.recordDiagnostic(vmmOp, os.str());
+        return failure();
+      }
+      if (weightView->shape.size() != 2) {
+        std::string message;
+        llvm::raw_string_ostream os(message);
+        os << "requires a rank-2 matrix weight view during PIM codegen; resolved shape=[";
+        llvm::interleaveComma(weightView->shape, os);
+        os << "] weight=" << vmmOp.getWeight() << " type=" << vmmOp.getWeight().getType();
+        result.recordDiagnostic(vmmOp, os.str());
         return failure();
       }
 
@@ -1308,15 +1667,16 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
       deviceMemory.allocateCore(coreOp);
 
-      int64_t processedOperations = codeGenCoreOps(
-        coreOp.getBody().front(), coreCodeGen, StaticValueKnowledge {}, coreOp.getOperation(), resolveWeightSlot);
+      StaticValueKnowledge knowledge = seedCoreCodegenKnowledge(coreOp);
+      int64_t processedOperations =
+        codeGenCoreOps(coreOp.getBody().front(), coreCodeGen, knowledge, coreOp.getOperation(), resolveWeightSlot);
       if (processedOperations < 0) {
         result.status = CompilerFailure;
         return result;
       }
-      assert(processedOperations > 0);
       result.reportRow = deviceMemory.getReportRow();
       result.usedWeights = std::move(usedWeights);
+      result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
     }
     else {
       auto coreBatchOp = cast<pim::PimCoreBatchOp>(job.coreLikeOp);
@@ -1324,10 +1684,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       auto& deviceMemory = jobMemory.getOrCreateDeviceMem(job.emittedCoreId);
 
       for (unsigned lane : job.lanes) {
-        StaticValueKnowledge knowledge;
-        knowledge.indexValues[coreBatchOp.getLaneArgument()] = lane;
-        for (unsigned i = 0; i < coreBatchOp.getInputs().size(); ++i)
-          knowledge.aliases[coreBatchOp.getInputArgument(i)] = coreBatchOp.getInputs()[i];
+        StaticValueKnowledge knowledge = seedCoreBatchCodegenKnowledge(coreBatchOp, lane);
 
         deviceMemory.allocateCore(coreBatchOp, lane);
         coreCodeGen.setBatchLane(lane);
@@ -1342,11 +1699,11 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
           result.status = CompilerFailure;
           return result;
         }
-        assert(processedOperations > 0);
       }
 
       result.reportRow = deviceMemory.getReportRow();
       result.usedWeights = std::move(usedWeights);
+      result.livenessArtifacts = deviceMemory.getLivenessArtifacts();
     }
 
     pim_binary::patchInstructionCount(coreBinaryStream, coreCodeGen.getEmittedInstructionCount());
@@ -1365,10 +1722,32 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
   mlir::parallelFor(
     moduleOp.getContext(), 0, jobs.size(), [&](size_t index) { jobResults[index] = emitJob(jobs[index]); });
 
+  pim::CappedDiagnosticReporter diagnostics;
+  Operation* summaryAnchor = nullptr;
+  for (const CoreEmissionResult& result : jobResults) {
+    if (!summaryAnchor && !result.diagnostics.empty())
+      summaryAnchor = result.diagnostics.front().op;
+    for (const CoreEmissionResult::DiagnosticRecord& diagnostic : result.diagnostics) {
+      diagnostics.report(diagnostic.op, [&](Operation* op) { op->emitError() << diagnostic.message; });
+    }
+    size_t unreportedCount = result.diagnosticCount - result.diagnostics.size();
+    diagnostics.noteFailures(static_cast<int64_t>(unreportedCount));
+  }
+  if (diagnostics.hasFailure())
+    diagnostics.emitSuppressedSummary(summaryAnchor ? summaryAnchor : moduleOp.getOperation(),
+                                      "PIM codegen diagnostic(s)");
+
   for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex)
     if (jobResults[jobIndex].status != CompilerSuccess)
       return jobResults[jobIndex].status;
 
+  if (jobs.empty()) {
+    if (auto err = emitEmptyCoreArtifacts(outputDirPath, 0))
+      return err;
+    xbarsPerArrayGroup["core0"] = json::Array {};
+    memory.recordCoreReport(0, MemoryReportRow {});
+  }
+
   llvm::SmallVector<WeightFileRequest, 8> weightRequests;
   weightRequests.reserve(jobs.size());
   for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
@@ -1380,6 +1759,18 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
   auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath);
   memory.setTotalWeightBytes(weightEmission.totalWeightBytes);
   auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName;
+  if (std::string reportsRoot = getOutputDir(); !reportsRoot.empty()) {
+    std::string reportsDir = reportsRoot + "/reports";
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_report.txt");
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_report.json");
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_timeline.dot");
+  }
+  std::fstream livenessReportFile;
+  std::unique_ptr<llvm::raw_os_ostream> livenessReportOs;
+  if (pimMemoryReport != PimMemoryReportNone) {
+    livenessReportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt");
+    livenessReportOs = std::make_unique<llvm::raw_os_ostream>(livenessReportFile);
+  }
 
   for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
     const CoreEmissionJob& job = jobs[jobIndex];
@@ -1391,6 +1782,8 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
         return err;
       xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
       memory.recordCoreReport(job.emittedCoreId, result.reportRow);
+      if (livenessReportFile.is_open())
+        *livenessReportOs << "Core " << job.emittedCoreId << ":\n" << result.livenessArtifacts.textReport;
       continue;
     }
   }
@@ -1419,10 +1812,18 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
                              batchPerCoreRow.value_or(MemoryReportRow {}),
                              batchRow.numAlloca,
                              batchRow.sizeAlloca);
+    if (livenessReportFile.is_open())
+      for (size_t jobIndex : group)
+        *livenessReportOs << "Batch " << batchReportId << " core " << jobs[jobIndex].emittedCoreId << ":\n"
+                          << jobResults[jobIndex].livenessArtifacts.textReport;
   }
 
   maxCoreId = nextEmittedCoreId == 0 ? 0 : nextEmittedCoreId - 1;
 
+  if (livenessReportFile.is_open()) {
+    livenessReportOs->flush();
+    livenessReportFile.close();
+  }
   memory.flushReport();
   return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
 }

src/PIM/Compiler/PimCodeGen.hpp

@@ -5,12 +5,15 @@
 #include "llvm-project/clang/include/clang/Basic/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_os_ostream.h"
 
+#include <array>
 #include <fstream>
 #include <limits>
 #include <optional>
+#include <string>
 
 #include "onnx-mlir/Compiler/OMCompilerTypes.h"
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
@@ -26,6 +29,16 @@ struct MemEntry {
   size_t size;
 };
 
+struct PhysicalSlotInfo {
+  size_t id = 0;
+  size_t address = 0;
+  size_t size = 0;
+};
+
+struct MemoryPlanArtifacts {
+  std::string textReport;
+};
+
 struct MemoryValueKey {
   mlir::Value value;
   std::optional<unsigned> lane;
@@ -74,16 +87,20 @@ struct MemoryReportEntry {
 
 class PimMemory {
   llvm::SmallVector<PendingMemEntry, 32> memEntries;
+  llvm::SmallVector<PhysicalSlotInfo, 32> localPhysicalSlots;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> ownedMemEntriesMap;
   MemoryReportRow reportRow;
+  MemoryPlanArtifacts livenessArtifacts;
 
   size_t minAlignment = 4;
   size_t firstAvailableAddress = 0;
+  size_t nextPhysicalSlotId = 0;
 
   MemEntry* gatherMemEntry(mlir::Value value, std::optional<unsigned> lane = std::nullopt);
   void allocateGatheredMemory();
   void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind);
+  PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key);
 
 public:
   PimMemory(llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap)
@@ -92,6 +109,7 @@ public:
   void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
   void allocateCore(mlir::Operation* op, std::optional<unsigned> lane = std::nullopt);
   MemoryReportRow getReportRow() const;
+  const MemoryPlanArtifacts& getLivenessArtifacts() const { return livenessArtifacts; }
   void remove(mlir::Value val);
 
   size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
@@ -153,6 +171,7 @@ class PimCodeGen {
   const llvm::DenseMap<size_t, size_t>& emittedCoreIds;
   std::optional<unsigned> batchLane;
   mutable uint32_t emittedInstructionCount = 0;
+  mutable std::array<std::optional<int32_t>, 256> scalarRegisterValues = {};
 
   size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
     return memory.getValueAddress(value, knowledge, batchLane);
@@ -160,6 +179,7 @@ class PimCodeGen {
   size_t remapCoreId(size_t coreId) const;
 
   void emitInstruction(const pim_binary::InstructionRecord& instruction) const;
+  void updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const;
 
   void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const;
   void setupRd(size_t rdAddress, size_t rdOffset) const;

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -22,22 +22,112 @@ llvm::cl::opt<PimMergeSchedulerType>
                     llvm::cl::init(MergeSchedulerPeft),
                     llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport(
+  "pim-memory-report",
+  llvm::cl::desc("Emit a human-readable PIM memory planning report"),
+  llvm::cl::values(clEnumValN(PimMemoryReportNone, "none", "Do not emit any PIM memory planning report")),
+  llvm::cl::values(
+    clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise slot reuse report with key offenders")),
+  llvm::cl::values(clEnumValN(PimMemoryReportFull, "full", "Emit the full detailed PIM memory planning report")),
+  llvm::cl::init(PimMemoryReportNone),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<PimConvLoweringType> pimConvLowering(
+  "pim-conv-lowering",
+  llvm::cl::desc("Convolution lowering strategy for PIM"),
+  llvm::cl::values(clEnumValN(PimConvLoweringAuto, "auto", "Select the Conv lowering strategy automatically")),
+  llvm::cl::values(clEnumValN(PimConvLoweringLegacy, "legacy", "Use the legacy explicit-im2col Conv lowering")),
+  llvm::cl::values(clEnumValN(PimConvLoweringDepthwise, "depthwise", "Force the depthwise-specialized Conv lowering")),
+  llvm::cl::values(
+    clEnumValN(PimConvLoweringPackedIm2Col, "packed-im2col", "Use explicit im2col with packed multi-position GEMM")),
+  llvm::cl::values(clEnumValN(PimConvLoweringStreamedPatch,
+                              "streamed-patch",
+                              "Use streamed/chunked im2col rows without multi-position packing")),
+  llvm::cl::values(clEnumValN(PimConvLoweringStreamedPacked,
+                              "streamed-packed",
+                              "Use streamed/chunked im2col rows with packed multi-position GEMM")),
+  llvm::cl::values(clEnumValN(PimConvLoweringOutputChannelTiled,
+                              "output-channel-tiled",
+                              "Force Conv lowering that relies on Gemm output-channel tiling")),
+  llvm::cl::values(
+    clEnumValN(PimConvLoweringInputKTiled, "input-k-tiled", "Force Conv lowering that relies on Gemm K tiling")),
+  llvm::cl::values(clEnumValN(PimConvLoweringTiled2D,
+                              "tiled-2d",
+                              "Force Conv lowering that relies on Gemm 2D K/C tiling")),
+  llvm::cl::init(PimConvLoweringAuto),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow(
+  "pim-export-spatial-dataflow",
+  llvm::cl::desc("Emit Gephi-importable CSV dataflow reports for Spatial pipeline snapshots"),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportNone, "none", "Do not emit Spatial dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit spatial1 graph dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 scheduled dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 realized dataflow CSV reports")),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportAll, "all", "Emit all Spatial dataflow CSV reports")),
+  llvm::cl::init(SpatialDataflowExportNone),
+  llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<bool>
   pimOnlyCodegen("pim-only-codegen",
                  llvm::cl::desc("Only generate code for PIM (assume input is already in bufferized PIM IR)"),
                  llvm::cl::init(false),
                  llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool>
+  pimDisableMemoryCoalescing("pim-disable-memory-coalescing",
+                             llvm::cl::desc("Skip the PIM memory coalescing pass (developer diagnostic option)"),
+                             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"),
+  llvm::cl::init(1ull << 20),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<uint64_t> pimConvStreamChunkPositions(
+  "pim-conv-stream-chunk-positions",
+  llvm::cl::desc("Maximum number of Conv output positions to materialize in one streamed chunk"),
+  llvm::cl::init(1024),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<bool> pimReportConvLowering("pim-report-conv-lowering",
+                                          llvm::cl::desc("Emit a bounded Conv lowering report"),
+                                          llvm::cl::init(true),
+                                          llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
                                 llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"),
                                 llvm::cl::init(false),
                                 llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
+  "pim-detect-communication-deadlock",
+  llvm::cl::desc("Expensively simulate the statically expanded PIM send/receive order at verification time and fail if a blocking communication deadlock is found"),
+  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::init(false),
+  llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<size_t>
   crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(2));
 

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -24,17 +24,53 @@ typedef enum {
   MergeSchedulerPeft = 0,
 } PimMergeSchedulerType;
 
+typedef enum {
+  PimMemoryReportNone = 0,
+  PimMemoryReportSummary = 1,
+  PimMemoryReportFull = 2,
+} PimMemoryReportLevel;
+
+typedef enum {
+  PimConvLoweringAuto = 0,
+  PimConvLoweringLegacy = 1,
+  PimConvLoweringDepthwise = 2,
+  PimConvLoweringPackedIm2Col = 3,
+  PimConvLoweringStreamedPatch = 4,
+  PimConvLoweringStreamedPacked = 5,
+  PimConvLoweringOutputChannelTiled = 6,
+  PimConvLoweringInputKTiled = 7,
+  PimConvLoweringTiled2D = 8,
+} PimConvLoweringType;
+
+typedef enum {
+  SpatialDataflowExportNone = 0,
+  SpatialDataflowExportSpatial1 = 1,
+  SpatialDataflowExportSpatial2 = 2,
+  SpatialDataflowExportSpatial3 = 3,
+  SpatialDataflowExportAll = 4,
+} PimSpatialDataflowExportType;
+
 extern llvm::cl::OptionCategory OnnxMlirOptions;
 extern llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget;
 extern llvm::cl::opt<PimMergeSchedulerType> pimMergeScheduler;
+extern llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport;
+extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
+extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;
 
 extern llvm::cl::opt<bool> pimOnlyCodegen;
+extern llvm::cl::opt<bool> pimDisableMemoryCoalescing;
 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<size_t> crossbarSize;
 extern llvm::cl::opt<size_t> crossbarCountInCore;
 extern llvm::cl::opt<long> coresCount;
+extern llvm::cl::opt<uint64_t> pimConvIm2colMaxElements;
+extern llvm::cl::opt<uint64_t> pimConvStreamChunkPositions;
 
 bool hasExplicitPimCoreCount();
 void verifyExplicitPimCoreCount();

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -29,6 +29,8 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
 
   if (pimEmissionTarget >= EmitSpatial) {
     pm.addPass(createONNXToSpatialPass());
+    pm.addPass(createSpatialLayoutPlanningPass());
+    pm.addPass(createLowerSpatialPlansPass());
     pm.addPass(createMergeComputeNodesPass());
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
@@ -46,8 +48,8 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
   if (pimEmissionTarget >= EmitPimCodegen) {
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
-    pm.addPass(createPimMaterializeHostConstantsPass());
-    pm.addPass(createPimMemoryCoalescingPass());
+    if (!pimDisableMemoryCoalescing)
+      pm.addPass(createPimMemoryCoalescingPass());
     pm.addPass(createPimVerificationPass());
     pm.addPass(createMessagePass("Pim verified"));
     pm.addPass(createEmitPimCodePass());

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -0,0 +1,752 @@
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <numeric>
+#include <string>
+#include <tuple>
+#include <utility>
+
+#include "Common/Support/CheckedArithmetic.hpp"
+#include "Common/Support/ReportUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace llvm;
+using namespace mlir;
+using namespace onnx_mlir;
+
+namespace {
+
+static std::optional<unsigned> getLaneForMemoryValue(mlir::Value value, std::optional<unsigned> lane) {
+  if (!lane)
+    return std::nullopt;
+  auto allocOp = value.getDefiningOp<memref::AllocOp>();
+  if (!allocOp || !allocOp->getParentOfType<pim::PimCoreBatchOp>())
+    return std::nullopt;
+  return lane;
+}
+
+static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional<unsigned> lane = std::nullopt) {
+  return {value, getLaneForMemoryValue(value, lane)};
+}
+
+struct MemoryTouchInterval {
+  uint64_t start = 0;
+  uint64_t end = 0;
+  Operation* startOp = nullptr;
+  Operation* endOp = nullptr;
+  Operation* firstTouchOp = nullptr;
+  Operation* lastTouchOp = nullptr;
+  uint64_t firstTouchPosition = 0;
+  uint64_t lastTouchPosition = 0;
+  bool hasRuntimeUse = false;
+  bool startUsedAllocFallback = false;
+  bool endUsedFallback = false;
+  bool escapesLoop = false;
+  std::string fallbackReason;
+  llvm::SmallVector<std::string, 8> aliasesFollowed;
+};
+
+struct OperationOrdering {
+  llvm::DenseMap<Operation*, uint64_t> position;
+  llvm::DenseMap<Operation*, uint64_t> subtreeEnd;
+  uint64_t nextPosition = 0;
+};
+
+static std::string printValueToString(mlir::Value value) {
+  std::string text;
+  llvm::raw_string_ostream os(text);
+  value.print(os);
+  os.flush();
+  return text;
+}
+
+static std::string printOperationToString(Operation* op) {
+  if (!op)
+    return "<none>";
+  std::string text;
+  llvm::raw_string_ostream os(text);
+  op->print(os);
+  os.flush();
+  return text;
+}
+
+static std::string printLocationToString(Location loc) {
+  std::string text;
+  llvm::raw_string_ostream os(text);
+  loc.print(os);
+  os.flush();
+  return text;
+}
+
+static std::string collapseWhitespace(StringRef text) {
+  std::string out;
+  out.reserve(text.size());
+  bool lastWasSpace = false;
+  for (char c : text) {
+    bool isSpace = c == ' ' || c == '\n' || c == '\t' || c == '\r';
+    if (isSpace) {
+      if (!lastWasSpace && !out.empty())
+        out.push_back(' ');
+      lastWasSpace = true;
+      continue;
+    }
+    out.push_back(c);
+    lastWasSpace = false;
+  }
+  return out;
+}
+
+static std::string abbreviate(StringRef text, size_t maxLen) {
+  if (text.size() <= maxLen)
+    return text.str();
+  return (text.take_front(maxLen - 3) + "...").str();
+}
+
+static std::string summarizeValue(mlir::Value value, size_t maxLen = 72) {
+  return abbreviate(collapseWhitespace(printValueToString(value)), maxLen);
+}
+
+static std::string summarizeOperation(Operation* op, size_t maxLen = 96) {
+  if (!op)
+    return "<none>";
+  std::string prefix = op->getName().getStringRef().str();
+  std::string full = collapseWhitespace(printOperationToString(op));
+  if (full == prefix)
+    return prefix;
+  return abbreviate(prefix + " :: " + full, maxLen);
+}
+
+static std::string summarizeLocation(Location loc, size_t maxLen = 88) {
+  return abbreviate(collapseWhitespace(printLocationToString(loc)), maxLen);
+}
+
+static void assignOperationOrdering(Operation* op, OperationOrdering& ordering) {
+  uint64_t position = ordering.nextPosition++;
+  ordering.position[op] = position;
+  uint64_t end = position;
+  for (Region& region : op->getRegions())
+    for (Block& block : region)
+      for (Operation& nestedOp : block) {
+        assignOperationOrdering(&nestedOp, ordering);
+        end = std::max(end, ordering.subtreeEnd.lookup(&nestedOp));
+      }
+  ordering.subtreeEnd[op] = end;
+}
+
+static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) {
+  OperationOrdering ordering;
+  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
+    return ordering;
+
+  for (Operation& op : coreLikeOp->getRegion(0).front())
+    assignOperationOrdering(&op, ordering);
+  return ordering;
+}
+
+static bool isSupportedAliasOp(Operation* op) {
+  return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
+}
+
+static bool isRuntimeMemoryTouchOp(Operation* op) {
+  return isa<pim::PimMemCopyHostToDevOp,
+             pim::PimMemCopyDevToHostOp,
+             pim::PimMemCopyOp,
+             pim::PimReceiveOp,
+             pim::PimSendOp,
+             pim::PimConcatOp,
+             pim::PimVMMOp,
+             pim::PimTransposeOp,
+             pim::PimVVAddOp,
+             pim::PimVVSubOp,
+             pim::PimVVMulOp,
+             pim::PimVVMaxOp,
+             pim::PimVVDMulOp,
+             pim::PimVAvgOp,
+             pim::PimVReluOp,
+             pim::PimVTanhOp,
+             pim::PimVSigmOp,
+             pim::PimVSoftmaxOp>(op);
+}
+
+static bool isIgnoredLivenessUser(Operation* op) {
+  return isSupportedAliasOp(op) || isa<scf::ForOp, scf::YieldOp, memref::DeallocOp>(op) || isCoreStaticAddressOp(op);
+}
+
+static bool isWithin(mlir::Value value, Region* region) {
+  if (!region)
+    return false;
+  if (auto blockArg = dyn_cast<BlockArgument>(value))
+    return blockArg.getOwner()->getParent() == region;
+  if (Operation* definingOp = value.getDefiningOp())
+    return definingOp->getParentRegion() == region || region->isAncestor(definingOp->getParentRegion());
+  return false;
+}
+
+static bool isNestedAllocation(Operation* coreLikeOp, memref::AllocOp allocOp) {
+  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
+    return false;
+  return allocOp->getBlock() != &coreLikeOp->getRegion(0).front();
+}
+
+static void addFallbackReason(std::string& reason, StringRef newReason) {
+  if (newReason.empty())
+    return;
+  if (!reason.empty())
+    reason += "; ";
+  reason += newReason.str();
+}
+
+static void appendAliasDescription(llvm::SmallVectorImpl<std::string>& aliases, mlir::Value value) {
+  std::string text = printValueToString(value);
+  if (!llvm::is_contained(aliases, text))
+    aliases.push_back(std::move(text));
+}
+
+struct OrderedTouchRange {
+  uint64_t start = 0;
+  uint64_t end = 0;
+  Operation* startOp = nullptr;
+  Operation* endOp = nullptr;
+  bool escapedLoop = false;
+};
+
+static OrderedTouchRange
+getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const OperationOrdering& ordering) {
+  OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), user, user, false};
+  for (Operation* current = user; current; current = current->getParentOp()) {
+    auto forOp = dyn_cast<scf::ForOp>(current);
+    if (!forOp || isWithin(definingValue, &forOp.getRegion()))
+      continue;
+    range.start = std::min(range.start, ordering.position.lookup(forOp));
+    range.end = std::max(range.end, ordering.subtreeEnd.lookup(forOp));
+    range.startOp = forOp;
+    range.endOp = forOp;
+    range.escapedLoop = true;
+  }
+  return range;
+}
+
+static MemoryTouchInterval
+computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ordering, uint64_t fallbackEnd) {
+  MemoryTouchInterval interval;
+  interval.start = ordering.position.lookup(allocOp);
+  interval.end = interval.start;
+  interval.startOp = allocOp;
+  interval.endOp = allocOp;
+
+  SmallPtrSet<mlir::Value, 16> visitedValues;
+  SmallPtrSet<Operation*, 32> visitedUsers;
+  SmallVector<mlir::Value> pendingValues;
+  pendingValues.push_back(allocOp.getResult());
+  auto parentLoop = allocOp->getParentOfType<scf::ForOp>();
+
+  while (!pendingValues.empty()) {
+    mlir::Value value = pendingValues.pop_back_val();
+    if (!visitedValues.insert(value).second)
+      continue;
+
+    for (Operation* user : value.getUsers()) {
+      if (!visitedUsers.insert(user).second)
+        continue;
+
+      if (isSupportedAliasOp(user)) {
+        for (mlir::Value result : user->getResults()) {
+          pendingValues.push_back(result);
+          appendAliasDescription(interval.aliasesFollowed, result);
+        }
+      }
+
+      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user)) {
+        for (OpResult result : user->getResults()) {
+          OpOperand* tiedOperand = dpsOp.getTiedOpOperand(result);
+          if (!tiedOperand || tiedOperand->get() != value)
+            continue;
+          pendingValues.push_back(result);
+          appendAliasDescription(interval.aliasesFollowed, result);
+        }
+      }
+
+      if (auto forOp = dyn_cast<scf::ForOp>(user)) {
+        for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs())) {
+          if (initArg != value)
+            continue;
+          pendingValues.push_back(forOp.getRegionIterArgs()[index]);
+          pendingValues.push_back(forOp.getResult(index));
+          appendAliasDescription(interval.aliasesFollowed, forOp.getRegionIterArgs()[index]);
+          appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index));
+          if (parentLoop && forOp != parentLoop)
+            interval.escapesLoop = true;
+        }
+      }
+
+      if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
+        auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
+        auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp());
+        auto indexSwitch = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp());
+        if (ifOp) {
+          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+            if (operand != value)
+              continue;
+            pendingValues.push_back(ifOp.getResult(index));
+            appendAliasDescription(interval.aliasesFollowed, ifOp.getResult(index));
+          }
+        }
+        else if (indexSwitch) {
+          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+            if (operand != value)
+              continue;
+            pendingValues.push_back(indexSwitch.getResult(index));
+            appendAliasDescription(interval.aliasesFollowed,
+                                   indexSwitch.getResult(index));
+          }
+        }
+        else if (!forOp) {
+          addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
+        }
+        else {
+          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+            if (operand != value)
+              continue;
+            pendingValues.push_back(forOp.getResult(index));
+            appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index));
+            if (parentLoop && forOp == parentLoop)
+              interval.escapesLoop = true;
+          }
+        }
+      }
+
+      if (isRuntimeMemoryTouchOp(user)) {
+        uint64_t touchPosition = ordering.position.lookup(user);
+        if (!interval.hasRuntimeUse || touchPosition < interval.firstTouchPosition) {
+          interval.firstTouchPosition = touchPosition;
+          interval.firstTouchOp = user;
+        }
+        if (!interval.hasRuntimeUse || touchPosition > interval.lastTouchPosition) {
+          interval.lastTouchPosition = touchPosition;
+          interval.lastTouchOp = user;
+        }
+
+        OrderedTouchRange range = getEffectiveTouchRange(allocOp.getResult(), user, ordering);
+        interval.escapesLoop |= range.escapedLoop;
+        if (!interval.hasRuntimeUse) {
+          interval.start = range.start;
+          interval.end = range.end;
+          interval.startOp = range.startOp;
+          interval.endOp = range.endOp;
+          interval.hasRuntimeUse = true;
+        }
+        else {
+          if (range.start < interval.start) {
+            interval.start = range.start;
+            interval.startOp = range.startOp;
+          }
+          if (range.end > interval.end) {
+            interval.end = range.end;
+            interval.endOp = range.endOp;
+          }
+        }
+        continue;
+      }
+
+      if (isIgnoredLivenessUser(user))
+        continue;
+
+      addFallbackReason(interval.fallbackReason, "unhandled user op");
+      interval.endUsedFallback = true;
+    }
+  }
+
+  if (!interval.hasRuntimeUse) {
+    interval.startUsedAllocFallback = true;
+    interval.endUsedFallback = true;
+    interval.start = ordering.position.lookup(allocOp);
+    interval.end = fallbackEnd;
+    interval.startOp = allocOp;
+    interval.endOp = allocOp->getParentOp();
+    interval.firstTouchPosition = interval.start;
+    interval.lastTouchPosition = interval.end;
+    addFallbackReason(interval.fallbackReason, "no runtime memory touch");
+    return interval;
+  }
+
+  if (interval.endUsedFallback) {
+    interval.end = std::max(interval.end, fallbackEnd);
+    interval.endOp = allocOp->getParentOp();
+  }
+
+  return interval;
+}
+
+static FailureOr<size_t> getAllocSizeBytes(memref::AllocOp allocOp) {
+  auto type = dyn_cast<ShapedType>(allocOp.getType());
+  if (!type)
+    return failure();
+  auto checkedBytes = pim::getCheckedShapedTypeSizeInBytes(type, allocOp, "memory allocation byte size");
+  if (failed(checkedBytes))
+    return failure();
+  return pim::checkedSize(*checkedBytes, allocOp, "memory allocation byte size");
+}
+
+static bool intervalsOverlap(const LocalAllocInterval& lhs, const LocalAllocInterval& rhs) {
+  return !(lhs.end < rhs.start || rhs.end < lhs.start);
+}
+
+static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, ArrayRef<LocalAllocInterval> intervals) {
+  uint64_t slotLogicalBytes = 0;
+  for (size_t intervalIndex : slot.intervalIndices)
+    slotLogicalBytes += intervals[intervalIndex].size;
+  return slotLogicalBytes;
+}
+
+} // namespace
+
+SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp,
+                                                                       std::optional<unsigned> lane) {
+  SmallVector<LocalAllocInterval, 0> intervals;
+  OperationOrdering ordering = buildOperationOrdering(coreLikeOp);
+  if (ordering.position.empty())
+    return intervals;
+
+  uint64_t fallbackEnd = ordering.nextPosition == 0 ? 0 : ordering.nextPosition - 1;
+  size_t nextIntervalId = 0;
+  coreLikeOp->walk([&](memref::AllocOp allocOp) {
+    auto checkedSize = getAllocSizeBytes(allocOp);
+    if (failed(checkedSize)) {
+      llvm::errs() << "Failed to compute local allocation size for value: ";
+      allocOp.getResult().print(llvm::errs());
+      llvm::errs() << "\n";
+      llvm_unreachable("Failed to compute local allocation size");
+    }
+
+    MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd);
+    LocalAllocInterval interval;
+    interval.id = nextIntervalId++;
+    interval.alloc = allocOp;
+    interval.key = getMemoryValueKey(allocOp.getResult(), lane);
+    interval.start = touchInterval.start;
+    interval.end = touchInterval.end;
+    interval.size = *checkedSize;
+    interval.startOp = touchInterval.startOp;
+    interval.endOp = touchInterval.endOp;
+    interval.firstTouchOp = touchInterval.firstTouchOp;
+    interval.lastTouchOp = touchInterval.lastTouchOp;
+    interval.firstTouchPosition = touchInterval.firstTouchPosition;
+    interval.lastTouchPosition = touchInterval.lastTouchPosition;
+    interval.startUsedAllocFallback = touchInterval.startUsedAllocFallback;
+    interval.endUsedFallback = touchInterval.endUsedFallback;
+    interval.hasRuntimeUse = touchInterval.hasRuntimeUse;
+    interval.insideNestedRegion = isNestedAllocation(coreLikeOp, allocOp);
+    interval.escapesLoop = touchInterval.escapesLoop;
+    interval.fallbackReason = std::move(touchInterval.fallbackReason);
+    interval.aliasesFollowed = std::move(touchInterval.aliasesFollowed);
+    intervals.push_back(std::move(interval));
+  });
+
+  return intervals;
+}
+
+SmallVector<PlannedPhysicalSlot, 0> onnx_mlir::planPhysicalSlots(MutableArrayRef<LocalAllocInterval> intervals) {
+  SmallVector<PlannedPhysicalSlot, 0> slots;
+  SmallVector<size_t> intervalOrder(intervals.size());
+  std::iota(intervalOrder.begin(), intervalOrder.end(), 0);
+  llvm::stable_sort(intervalOrder, [&](size_t lhsIndex, size_t rhsIndex) {
+    const LocalAllocInterval& lhs = intervals[lhsIndex];
+    const LocalAllocInterval& rhs = intervals[rhsIndex];
+    if (lhs.size != rhs.size)
+      return lhs.size > rhs.size;
+    if (lhs.start != rhs.start)
+      return lhs.start < rhs.start;
+    if (lhs.end != rhs.end)
+      return lhs.end < rhs.end;
+    return lhs.id < rhs.id;
+  });
+
+  for (size_t intervalIndex : intervalOrder) {
+    LocalAllocInterval& interval = intervals[intervalIndex];
+    PlannedPhysicalSlot* bestSlot = nullptr;
+    auto bestKey = std::tuple<size_t, size_t, size_t, size_t>(std::numeric_limits<size_t>::max(),
+                                                              std::numeric_limits<size_t>::max(),
+                                                              std::numeric_limits<size_t>::max(),
+                                                              std::numeric_limits<size_t>::max());
+
+    for (size_t slotIndex = 0; slotIndex < slots.size(); ++slotIndex) {
+      PlannedPhysicalSlot& slot = slots[slotIndex];
+      bool compatible = true;
+      for (size_t otherIndex : slot.intervalIndices) {
+        if (intervalsOverlap(interval, intervals[otherIndex])) {
+          compatible = false;
+          break;
+        }
+      }
+      if (!compatible)
+        continue;
+
+      size_t resultingSize = std::max(slot.requiredSize, interval.size);
+      size_t growth = resultingSize - slot.requiredSize;
+      auto candidateKey =
+        std::tuple<size_t, size_t, size_t, size_t>(growth, resultingSize, slot.intervalIndices.size(), slot.id);
+      if (candidateKey < bestKey) {
+        bestKey = candidateKey;
+        bestSlot = &slot;
+      }
+    }
+
+    if (!bestSlot) {
+      slots.push_back({slots.size(), interval.size, interval.size, 0, {intervalIndex}});
+      interval.slotPlanIndex = slots.size() - 1;
+      interval.physicalSlotId = slots.back().id;
+      interval.physicalSlotSize = slots.back().requiredSize;
+      continue;
+    }
+
+    bestSlot->requiredSize = std::max(bestSlot->requiredSize, interval.size);
+    bestSlot->size = bestSlot->requiredSize;
+    bestSlot->intervalIndices.push_back(intervalIndex);
+    interval.slotPlanIndex = static_cast<size_t>(bestSlot - slots.data());
+    interval.physicalSlotId = bestSlot->id;
+    interval.physicalSlotSize = bestSlot->requiredSize;
+  }
+
+  return slots;
+}
+
+MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation* coreLikeOp,
+                                                        std::optional<unsigned> lane,
+                                                        ArrayRef<LocalAllocInterval> intervals,
+                                                        ArrayRef<PlannedPhysicalSlot> slots,
+                                                        size_t addressLimit,
+                                                        PimMemoryReportLevel reportLevel) {
+  MemoryPlanArtifacts artifacts;
+
+  uint64_t totalLogicalBytes = 0;
+  uint64_t totalPhysicalBytes = 0;
+  uint64_t fallbackIntervals = 0;
+  uint64_t noRuntimeTouchIntervals = 0;
+  uint64_t reusedAllocations = 0;
+  uint64_t nestedIntervals = 0;
+  uint64_t loopEscapingIntervals = 0;
+  size_t largestLogicalAllocation = 0;
+  size_t largestPhysicalSlot = 0;
+  size_t maximumAssignedAddress = 0;
+
+  for (const LocalAllocInterval& interval : intervals) {
+    totalLogicalBytes += interval.size;
+    largestLogicalAllocation = std::max(largestLogicalAllocation, interval.size);
+    maximumAssignedAddress = std::max(maximumAssignedAddress, interval.assignedAddress + interval.physicalSlotSize);
+    if (interval.startUsedAllocFallback || interval.endUsedFallback)
+      ++fallbackIntervals;
+    if (!interval.hasRuntimeUse)
+      ++noRuntimeTouchIntervals;
+    if (interval.insideNestedRegion)
+      ++nestedIntervals;
+    if (interval.escapesLoop)
+      ++loopEscapingIntervals;
+  }
+  for (const PlannedPhysicalSlot& slot : slots) {
+    totalPhysicalBytes += slot.size;
+    largestPhysicalSlot = std::max(largestPhysicalSlot, slot.size);
+    if (slot.intervalIndices.size() > 1)
+      reusedAllocations += slot.intervalIndices.size() - 1;
+  }
+
+  uint64_t savedBytes = totalLogicalBytes >= totalPhysicalBytes ? totalLogicalBytes - totalPhysicalBytes : 0;
+  double savedPercent =
+    totalLogicalBytes == 0 ? 0.0 : 100.0 * static_cast<double>(savedBytes) / static_cast<double>(totalLogicalBytes);
+
+  raw_string_ostream os(artifacts.textReport);
+  os << "=== PIM Memory Liveness Report ===\n";
+  os << "Op: " << coreLikeOp->getName() << "\n";
+  if (lane)
+    os << "Lane: " << *lane << "\n";
+  os << "Summary:\n";
+  os << "  logical allocation bytes:  " << formatReportMemory(totalLogicalBytes) << " (" << totalLogicalBytes << ")\n";
+  os << "  physical allocation bytes: " << formatReportMemory(totalPhysicalBytes) << " (" << totalPhysicalBytes
+     << ")\n";
+  os << "  saved bytes:               " << formatReportMemory(savedBytes) << " (" << savedBytes << ")\n";
+  os << "  saved percent: " << format("%.2f%%", savedPercent) << "\n";
+  os << "  intervals: " << intervals.size() << "\n";
+  os << "  physical slots: " << slots.size() << "\n";
+  os << "  reused allocations: " << reusedAllocations << "\n";
+  os << "  fallback intervals: " << fallbackIntervals << "\n";
+  os << "  intervals with no runtime memory touch: " << noRuntimeTouchIntervals << "\n";
+  os << "  nested allocations: " << nestedIntervals << "\n";
+  os << "  loop-escaping allocations: " << loopEscapingIntervals << "\n";
+  os << "  largest logical allocation: " << largestLogicalAllocation << "\n";
+  os << "  largest physical slot: " << largestPhysicalSlot << "\n";
+  os << "  address limit: " << addressLimit << "\n";
+  os << "  peak physical memory: " << formatReportMemory(maximumAssignedAddress) << " (" << maximumAssignedAddress
+     << ")\n";
+  os << "  maximum assigned address: " << maximumAssignedAddress << "\n";
+
+  os << "\nHow To Read:\n";
+  os << "  `summary` only shows the strongest reuse cases and the worst offenders.\n";
+  os << "  Use `--pim-memory-report=full` when you need the complete slot-by-slot and interval-by-interval dump.\n";
+  os << "  Large single-use slots, fallback intervals, and nested single-use allocations are the best places\n";
+  os << "  to inspect if allocations should be moved, sunk, or made easier to coalesce earlier in the pipeline.\n";
+
+  SmallVector<const PlannedPhysicalSlot*> reusedSlots;
+  SmallVector<const PlannedPhysicalSlot*> singleUseSlots;
+  for (const PlannedPhysicalSlot& slot : slots)
+    if (slot.intervalIndices.size() > 1)
+      reusedSlots.push_back(&slot);
+    else
+      singleUseSlots.push_back(&slot);
+
+  llvm::stable_sort(reusedSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
+    uint64_t lhsLogicalBytes = getSlotLogicalBytes(*lhs, intervals);
+    uint64_t rhsLogicalBytes = getSlotLogicalBytes(*rhs, intervals);
+    if (lhs->intervalIndices.size() != rhs->intervalIndices.size())
+      return lhs->intervalIndices.size() > rhs->intervalIndices.size();
+    if (lhsLogicalBytes != rhsLogicalBytes)
+      return lhsLogicalBytes > rhsLogicalBytes;
+    if (lhs->size != rhs->size)
+      return lhs->size > rhs->size;
+    return lhs->id < rhs->id;
+  });
+  llvm::stable_sort(singleUseSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
+    if (lhs->size != rhs->size)
+      return lhs->size > rhs->size;
+    return lhs->id < rhs->id;
+  });
+
+  constexpr size_t kSummaryReuseLimit = 6;
+  constexpr size_t kSummaryOffenderLimit = 10;
+
+  os << "\nBest Reuse:\n";
+  if (reusedSlots.empty()) {
+    os << "  no slots were shared by multiple intervals\n";
+  }
+  else {
+    for (const PlannedPhysicalSlot* slot : ArrayRef(reusedSlots).take_front(kSummaryReuseLimit)) {
+      uint64_t slotLogicalBytes = getSlotLogicalBytes(*slot, intervals);
+      os << "  slot #" << slot->id << "  addr=" << slot->address << "  size=" << formatReportMemory(slot->size)
+         << "  intervals=" << slot->intervalIndices.size() << "  logical_sum=" << formatReportMemory(slotLogicalBytes)
+         << "\n";
+      for (size_t intervalIndex : slot->intervalIndices) {
+        const LocalAllocInterval& interval = intervals[intervalIndex];
+        os << "    #" << interval.id << "  [" << interval.start << "," << interval.end << "]"
+           << "  logical=" << formatReportMemory(interval.size)
+           << "  first=" << summarizeOperation(interval.firstTouchOp, 40)
+           << "  last=" << summarizeOperation(interval.lastTouchOp, 40) << "\n";
+      }
+    }
+  }
+
+  os << "\nTop Offenders:\n";
+  bool printedAttention = false;
+  for (const PlannedPhysicalSlot* slot : ArrayRef(singleUseSlots).take_front(kSummaryOffenderLimit)) {
+    const LocalAllocInterval& interval = intervals[slot->intervalIndices.front()];
+    printedAttention = true;
+    os << "  slot #" << slot->id << " is single-use"
+       << "  size=" << formatReportMemory(slot->size) << "  interval=#" << interval.id
+       << "  value=" << summarizeValue(interval.key.value, 56) << "\n";
+    os << "    first=" << summarizeOperation(interval.firstTouchOp, 40)
+       << "  last=" << summarizeOperation(interval.lastTouchOp, 40)
+       << "  nested=" << (interval.insideNestedRegion ? "yes" : "no")
+       << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no") << "\n";
+  }
+  size_t fallbackPrinted = 0;
+  for (const LocalAllocInterval& interval : intervals) {
+    if (!(interval.startUsedAllocFallback || interval.endUsedFallback) || fallbackPrinted >= kSummaryOffenderLimit)
+      continue;
+    printedAttention = true;
+    ++fallbackPrinted;
+    os << "  fallback interval #" << interval.id << "  size=" << formatReportMemory(interval.size)
+       << "  value=" << summarizeValue(interval.key.value, 56) << "\n";
+    os << "    reason: " << (interval.fallbackReason.empty() ? "<none>" : interval.fallbackReason) << "\n";
+  }
+  size_t nestedPrinted = 0;
+  for (const LocalAllocInterval& interval : intervals) {
+    if (nestedPrinted >= kSummaryOffenderLimit)
+      break;
+    if (!(interval.insideNestedRegion && slots[interval.slotPlanIndex].intervalIndices.size() == 1))
+      continue;
+    printedAttention = true;
+    ++nestedPrinted;
+    os << "  nested single-use interval #" << interval.id << "  slot #" << interval.physicalSlotId
+       << "  size=" << formatReportMemory(interval.size) << "  value=" << summarizeValue(interval.key.value, 56)
+       << "\n";
+    os << "    hint: move or sink this alloc inside the nested region if the IR allows it.\n";
+  }
+  if (!printedAttention)
+    os << "  no obvious blockers detected in this core\n";
+
+  if (reportLevel == PimMemoryReportFull) {
+    os << "\nSlot Reuse:\n";
+    for (const PlannedPhysicalSlot& slot : slots) {
+      uint64_t slotLogicalBytes = getSlotLogicalBytes(slot, intervals);
+      os << "  slot #" << slot.id << "  addr=" << slot.address << "  size=" << formatReportMemory(slot.size) << " ("
+         << slot.size << ")"
+         << "  intervals=" << slot.intervalIndices.size() << "  logical_sum=" << formatReportMemory(slotLogicalBytes)
+         << "\n";
+      for (size_t intervalIndex : slot.intervalIndices) {
+        const LocalAllocInterval& interval = intervals[intervalIndex];
+        mlir::Value allocValue = interval.key.value;
+        os << "    [" << interval.start << "," << interval.end << "]"
+           << "  #" << interval.id << "  logical=" << formatReportMemory(interval.size)
+           << "  nested=" << (interval.insideNestedRegion ? "yes" : "no")
+           << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
+           << "  first=" << summarizeOperation(interval.firstTouchOp, 48)
+           << "  last=" << summarizeOperation(interval.lastTouchOp, 48) << "\n";
+        os << "      value=" << summarizeValue(allocValue) << "\n";
+      }
+    }
+  }
+
+  if (reportLevel == PimMemoryReportFull) {
+    os << "\nInterval Details:\n";
+    for (const LocalAllocInterval& interval : intervals) {
+      const PlannedPhysicalSlot& slot = slots[interval.slotPlanIndex];
+      mlir::Value allocValue = interval.key.value;
+      Operation* definingOp = allocValue.getDefiningOp();
+      os << "  #" << interval.id << "  slot=" << slot.id << "  live=[" << interval.start << "," << interval.end << "]"
+         << "  logical=" << formatReportMemory(interval.size)
+         << "  slot_size=" << formatReportMemory(interval.physicalSlotSize) << "  addr=" << interval.assignedAddress
+         << "\n";
+      os << "    value=" << summarizeValue(allocValue, 88) << "\n";
+      os << "    type=" << allocValue.getType() << "\n";
+      os << "    loc="
+         << summarizeLocation(definingOp ? definingOp->getLoc() : UnknownLoc::get(coreLikeOp->getContext())) << "\n";
+      os << "    nested=" << (interval.insideNestedRegion ? "yes" : "no")
+         << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
+         << "  start_fallback=" << (interval.startUsedAllocFallback ? "yes" : "no")
+         << "  end_fallback=" << (interval.endUsedFallback ? "yes" : "no") << "\n";
+      os << "    first_use=" << summarizeOperation(interval.firstTouchOp) << " @" << interval.firstTouchPosition
+         << "\n";
+      os << "    last_use=" << summarizeOperation(interval.lastTouchOp) << " @" << interval.lastTouchPosition << "\n";
+      os << "    slot_peers=";
+      bool first = true;
+      for (size_t otherIndex : slot.intervalIndices) {
+        if (intervals[otherIndex].id == interval.id)
+          continue;
+        if (!first)
+          os << ", ";
+        os << "#" << intervals[otherIndex].id;
+        first = false;
+      }
+      if (first)
+        os << "<none>";
+      os << "\n";
+      if (!interval.fallbackReason.empty())
+        os << "    fallback_reason=" << interval.fallbackReason << "\n";
+      if (!interval.aliasesFollowed.empty()) {
+        os << "    aliases_followed=" << interval.aliasesFollowed.size() << "\n";
+        for (const std::string& alias : interval.aliasesFollowed)
+          os << "      - " << abbreviate(collapseWhitespace(alias), 108) << "\n";
+      }
+    }
+  }
+  os.flush();
+
+  return artifacts;
+}

src/PIM/Compiler/PimMemoryLiveness.hpp

@@ -0,0 +1,63 @@
+#pragma once
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <limits>
+#include <optional>
+#include <string>
+
+#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+
+namespace onnx_mlir {
+
+struct LocalAllocInterval {
+  size_t id = 0;
+  mlir::memref::AllocOp alloc;
+  MemoryValueKey key;
+  uint64_t start = 0;
+  uint64_t end = 0;
+  size_t size = 0;
+  mlir::Operation* startOp = nullptr;
+  mlir::Operation* endOp = nullptr;
+  mlir::Operation* firstTouchOp = nullptr;
+  mlir::Operation* lastTouchOp = nullptr;
+  uint64_t firstTouchPosition = 0;
+  uint64_t lastTouchPosition = 0;
+  bool startUsedAllocFallback = false;
+  bool endUsedFallback = false;
+  bool hasRuntimeUse = false;
+  bool insideNestedRegion = false;
+  bool escapesLoop = false;
+  std::string fallbackReason;
+  llvm::SmallVector<std::string, 8> aliasesFollowed;
+  size_t slotPlanIndex = std::numeric_limits<size_t>::max();
+  size_t physicalSlotId = std::numeric_limits<size_t>::max();
+  size_t assignedAddress = 0;
+  size_t physicalSlotSize = 0;
+};
+
+struct PlannedPhysicalSlot {
+  size_t id = std::numeric_limits<size_t>::max();
+  size_t requiredSize = 0;
+  size_t size = 0;
+  size_t address = 0;
+  llvm::SmallVector<size_t, 8> intervalIndices;
+};
+
+llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp,
+                                                                  std::optional<unsigned> lane);
+
+llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);
+
+MemoryPlanArtifacts buildMemoryPlanArtifacts(mlir::Operation* coreLikeOp,
+                                             std::optional<unsigned> lane,
+                                             llvm::ArrayRef<LocalAllocInterval> intervals,
+                                             llvm::ArrayRef<PlannedPhysicalSlot> slots,
+                                             size_t addressLimit,
+                                             PimMemoryReportLevel reportLevel);
+
+} // namespace onnx_mlir

src/PIM/Compiler/PimWeightEmitter.cpp

@@ -19,8 +19,7 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {} // namespace
 
-WeightEmissionResult
-createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> requests, StringRef outputDirPath) {
+WeightEmissionResult createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> requests, StringRef outputDirPath) {
   auto coreWeightsDirPath = outputDirPath + "/weights";
   auto error = sys::fs::create_directory(coreWeightsDirPath);
   assert(!error && "Error creating weights directory");

src/PIM/Compiler/PimWeightEmitter.hpp

@@ -23,7 +23,7 @@ struct WeightEmissionResult {
   uint64_t totalWeightBytes = 0;
 };
 
-WeightEmissionResult
-createAndPopulateWeightFolder(llvm::ArrayRef<WeightFileRequest> requests, llvm::StringRef outputDirPath);
+WeightEmissionResult createAndPopulateWeightFolder(llvm::ArrayRef<WeightFileRequest> requests,
+                                                   llvm::StringRef outputDirPath);
 
 } // namespace onnx_mlir

src/PIM/Conversion/CMakeLists.txt

@@ -1,3 +1,2 @@
 add_subdirectory(ONNXToSpatial)
-add_subdirectory(SpatialToGraphviz)
-add_subdirectory(SpatialToPim)
+add_subdirectory(SpatialToPim)

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -10,6 +10,7 @@ add_pim_library(OMONNXToSpatial
   Patterns/Post.cpp
   Patterns/GeneratedConversion.cpp
   Patterns/Math/Conv.cpp
+  Patterns/Math/ConvGeometry.cpp
   Patterns/Math/Elementwise.cpp
   Patterns/Math/Gemm.cpp
   Patterns/Math/MatMul.cpp
@@ -19,15 +20,21 @@ add_pim_library(OMONNXToSpatial
   Patterns/NN/Sigmoid.cpp
   Patterns/NN/Softmax.cpp
   Patterns/Tensor/Concat.cpp
+  Patterns/Tensor/Flatten.cpp
   Patterns/Tensor/Gather.cpp
   Patterns/Tensor/Resize.cpp
   Patterns/Tensor/Reshape.cpp
+  Patterns/Tensor/Slice.cpp
   Patterns/Tensor/Split.cpp
   Patterns/Tensor/Transpose.cpp
   ONNXToSpatialPass.cpp
+  SpatialLayoutPlanningPass.cpp
+  LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
+  Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.cpp
-  Common/IndexingUtils.cpp
+  Common/MatrixProductLowering.cpp
+  Common/RowStripLayoutUtils.cpp
   Common/ShapeTilingUtils.cpp
   Common/WeightMaterialization.cpp
 

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

@@ -0,0 +1,112 @@
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult isSupportedBiasAddShape(RankedTensorType biasType, RankedTensorType resultType) {
+  if (!biasType || !resultType || !biasType.hasStaticShape() || !resultType.hasStaticShape())
+    return failure();
+  if (resultType.getRank() != 4)
+    return failure();
+  if (biasType.getElementType() != resultType.getElementType())
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  ArrayRef<int64_t> shape = biasType.getShape();
+  if (shape.empty())
+    return success();
+  if (shape.size() == 1)
+    return success(shape[0] == channels);
+  if (shape.size() == 2)
+    return success(shape[0] == 1 && shape[1] == channels);
+  if (shape.size() == 4)
+    return success(shape[0] == 1 && shape[1] == channels && shape[2] == 1 && shape[3] == 1);
+  return failure();
+}
+
+FailureOr<SmallVector<Attribute>> getBiasChannelValues(DenseElementsAttr denseAttr, RankedTensorType resultType) {
+  auto biasType = dyn_cast<RankedTensorType>(denseAttr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  if (denseAttr.isSplat()) {
+    return SmallVector<Attribute>(channels, denseAttr.getSplatValue<Attribute>());
+  }
+
+  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
+  if (biasType.getRank() == 1)
+    return flattened;
+  if (biasType.getRank() == 2)
+    return flattened;
+
+  SmallVector<Attribute> channelValues;
+  channelValues.reserve(channels);
+  const int64_t channelStride = biasType.getDimSize(2) * biasType.getDimSize(3);
+  for (int64_t channel = 0; channel < channels; ++channel)
+    channelValues.push_back(flattened[channel * channelStride]);
+  return channelValues;
+}
+
+bool isSupportedBiasAddValue(Value bias, RankedTensorType resultType, DenseElementsAttr* denseAttr) {
+  auto attr = getHostConstDenseElementsAttr(bias);
+  if (!attr)
+    return false;
+  auto biasType = dyn_cast<RankedTensorType>(attr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return false;
+  if (failed(getBiasChannelValues(attr, resultType)))
+    return false;
+  if (denseAttr)
+    *denseAttr = attr;
+  return true;
+}
+
+FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(Value lhs, Value rhs, RankedTensorType resultType) {
+  auto lhsType = dyn_cast<RankedTensorType>(lhs.getType());
+  auto rhsType = dyn_cast<RankedTensorType>(rhs.getType());
+  if (!lhsType || !rhsType)
+    return failure();
+  if (lhsType == resultType && isSupportedBiasAddValue(rhs, resultType))
+    return BiasAddPlanCandidate {lhs, rhs};
+  if (rhsType == resultType && isSupportedBiasAddValue(lhs, resultType))
+    return BiasAddPlanCandidate {rhs, lhs};
+  return failure();
+}
+
+FailureOr<Value>
+materializeDenseBiasAddTensor(Value bias, RankedTensorType resultType, RewriterBase& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, resultType, &denseAttr))
+    return failure();
+
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, resultType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> resultValues;
+  resultValues.reserve(resultType.getNumElements());
+  const int64_t batches = resultType.getDimSize(0);
+  const int64_t channels = resultType.getDimSize(1);
+  const int64_t height = resultType.getDimSize(2);
+  const int64_t width = resultType.getDimSize(3);
+  for (int64_t n = 0; n < batches; ++n)
+    for (int64_t c = 0; c < channels; ++c)
+      for (int64_t h = 0; h < height; ++h)
+        for (int64_t w = 0; w < width; ++w)
+          resultValues.push_back((*channelValues)[c]);
+
+  auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace onnx_mlir {
+
+struct BiasAddPlanCandidate {
+  mlir::Value data;
+  mlir::Value bias;
+};
+
+mlir::LogicalResult isSupportedBiasAddShape(mlir::RankedTensorType biasType, mlir::RankedTensorType resultType);
+bool isSupportedBiasAddValue(mlir::Value bias,
+                             mlir::RankedTensorType resultType,
+                             mlir::DenseElementsAttr* denseAttr = nullptr);
+mlir::FailureOr<llvm::SmallVector<mlir::Attribute>>
+getBiasChannelValues(mlir::DenseElementsAttr denseAttr, mlir::RankedTensorType resultType);
+mlir::FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(mlir::Value lhs,
+                                                                   mlir::Value rhs,
+                                                                   mlir::RankedTensorType resultType);
+mlir::FailureOr<mlir::Value> materializeDenseBiasAddTensor(mlir::Value bias,
+                                                           mlir::RankedTensorType resultType,
+                                                           mlir::RewriterBase& rewriter,
+                                                           mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -2,8 +2,9 @@
 
 #include "AttributeUtils.hpp"
 #include "ComputeRegionBuilder.hpp"
-#include "IndexingUtils.hpp"
+#include "MatrixProductLowering.hpp"
 #include "ShapeTilingUtils.hpp"
 #include "WeightMaterialization.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"

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

@@ -9,7 +9,7 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-Value sumTensors(ArrayRef<Value> tensors, ConversionPatternRewriter& rewriter) {
+Value sumTensors(ArrayRef<Value> tensors, PatternRewriter& rewriter) {
   if (tensors.size() == 1)
     return tensors[0];
 

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

@@ -60,6 +60,56 @@ struct SpatComputeBatchBodyArgs {
   mlir::ValueRange outputs;
 };
 
+inline mlir::SmallVector<mlir::Type> getGraphComputeBlockArgTypes(mlir::ValueRange weights, mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes;
+  blockArgTypes.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBlockArgLocs(mlir::Location defaultLoc,
+                                                                     mlir::ValueRange weights,
+                                                                     mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs;
+  blockArgLocs.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  return blockArgLocs;
+}
+
+inline mlir::SmallVector<mlir::Type> getGraphComputeBatchBlockArgTypes(mlir::OpBuilder& builder,
+                                                                       mlir::TypeRange resultTypes,
+                                                                       mlir::ValueRange weights,
+                                                                       mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes {builder.getIndexType()};
+  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  llvm::append_range(blockArgTypes, resultTypes);
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBatchBlockArgLocs(mlir::Location defaultLoc,
+                                                                          mlir::TypeRange resultTypes,
+                                                                          mlir::ValueRange weights,
+                                                                          mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs {defaultLoc};
+  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  blockArgLocs.append(resultTypes.size(), defaultLoc);
+  return blockArgLocs;
+}
+
 } // namespace detail
 
 template <typename RewriterT>
@@ -87,26 +137,43 @@ inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int
   return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
 }
 
-/// Builds a `spat.compute` with a fixed number of SSA inputs and erases it if
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs,
+                                                      mlir::TypeRange blockArgTypes,
+                                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
+  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
+  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
+  return computeOp;
+}
+
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBlockArgTypes(weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBlockArgLocs(loc, weights, inputs);
+  return createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, blockArgTypes, blockArgLocs);
+}
+
+/// Builds a `spat.graph_compute` with a fixed number of SSA inputs and erases it if
 /// the body callback reports failure.
 template <size_t NumInputs, typename RewriterT, typename BodyFn>
-auto createSpatCompute(RewriterT& rewriter,
-                       mlir::Location loc,
-                       mlir::TypeRange resultTypes,
-                       mlir::ValueRange weights,
-                       mlir::ValueRange inputs,
-                       BodyFn&& body) {
+auto createSpatGraphCompute(RewriterT& rewriter,
+                            mlir::Location loc,
+                            mlir::TypeRange resultTypes,
+                            mlir::ValueRange weights,
+                            mlir::ValueRange inputs,
+                            BodyFn&& body) {
   assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
-  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();
 
   using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
   if constexpr (std::is_same_v<BodyResult, void>) {
@@ -124,32 +191,24 @@ auto createSpatCompute(RewriterT& rewriter,
     if (mlir::failed(bodyResult)) {
       rewriter.setInsertionPointAfter(computeOp);
       rewriter.eraseOp(computeOp);
-      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
+      return mlir::FailureOr<spatial::SpatGraphCompute>(mlir::failure());
     }
     rewriter.setInsertionPointAfter(computeOp);
-    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
+    return mlir::FailureOr<spatial::SpatGraphCompute>(computeOp);
   }
 }
 
-/// Builds a `spat.compute` whose body consumes the block arguments as a single
+/// Builds a `spat.graph_compute` whose body consumes the block arguments as a single
 /// `ValueRange`, which is convenient for variadic reductions/concats.
 template <typename RewriterT, typename BodyFn>
-auto createSpatCompute(RewriterT& rewriter,
-                       mlir::Location loc,
-                       mlir::TypeRange resultTypes,
-                       mlir::ValueRange weights,
-                       mlir::ValueRange inputs,
-                       BodyFn&& body) {
-  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+auto createSpatGraphCompute(RewriterT& rewriter,
+                            mlir::Location loc,
+                            mlir::TypeRange resultTypes,
+                            mlir::ValueRange weights,
+                            mlir::ValueRange inputs,
+                            BodyFn&& body) {
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();
 
   using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
   if constexpr (std::is_same_v<BodyResult, void>) {
@@ -163,50 +222,60 @@ auto createSpatCompute(RewriterT& rewriter,
     if (mlir::failed(bodyResult)) {
       rewriter.setInsertionPointAfter(computeOp);
       rewriter.eraseOp(computeOp);
-      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
+      return mlir::FailureOr<spatial::SpatGraphCompute>(mlir::failure());
     }
     rewriter.setInsertionPointAfter(computeOp);
-    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
+    return mlir::FailureOr<spatial::SpatGraphCompute>(computeOp);
   }
 }
 
-template <typename RewriterT, typename BodyFn>
-auto createSpatComputeBatch(RewriterT& rewriter,
-                            mlir::Location loc,
-                            mlir::TypeRange resultTypes,
-                            int64_t laneCount,
-                            mlir::ValueRange weights,
-                            mlir::ValueRange inputs,
-                            BodyFn&& body) {
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs,
+                                      mlir::TypeRange blockArgTypes,
+                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
   if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
-    return mlir::FailureOr<spatial::SpatComputeBatch>(mlir::failure());
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
   auto laneCountAttr = pim::getCheckedI32Attr(rewriter, loc, laneCount, "spatial compute_batch lane count");
   if (mlir::failed(laneCountAttr))
-    return mlir::FailureOr<spatial::SpatComputeBatch>(mlir::failure());
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
-  auto batchOp = spatial::SpatComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
+  auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
+  rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&batchOp.getBody().front());
+  return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+}
 
-  mlir::SmallVector<mlir::Type> blockArgTypes {rewriter.getIndexType()};
-  mlir::SmallVector<mlir::Location> blockArgLocs {loc};
-  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  for (mlir::Value weight : weights) {
-    blockArgTypes.push_back(weight.getType());
-    blockArgLocs.push_back(weight.getLoc());
-  }
-  for (mlir::Value input : inputs) {
-    blockArgTypes.push_back(input.getType());
-    blockArgLocs.push_back(input.getLoc());
-  }
-  for (mlir::Type resultType : resultTypes) {
-    blockArgTypes.push_back(resultType);
-    blockArgLocs.push_back(loc);
-  }
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBatchBlockArgTypes(rewriter, resultTypes, weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBatchBlockArgLocs(loc, resultTypes, weights, inputs);
+  return createEmptySpatGraphComputeBatch(
+    rewriter, loc, resultTypes, laneCount, weights, inputs, blockArgTypes, blockArgLocs);
+}
 
-  auto* block =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToStart(block);
+template <typename RewriterT, typename BodyFn>
+auto createSpatGraphComputeBatch(RewriterT& rewriter,
+                                 mlir::Location loc,
+                                 mlir::TypeRange resultTypes,
+                                 int64_t laneCount,
+                                 mlir::ValueRange weights,
+                                 mlir::ValueRange inputs,
+                                 BodyFn&& body) {
+  auto batchOp = createEmptySpatGraphComputeBatch(rewriter, loc, resultTypes, laneCount, weights, inputs);
+  if (failed(batchOp))
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
+  auto* block = &(*batchOp).getBody().front();
 
   detail::SpatComputeBatchBodyArgs args {
     block->getArgument(0),
@@ -217,21 +286,54 @@ auto createSpatComputeBatch(RewriterT& rewriter,
   using BodyResult = std::invoke_result_t<BodyFn, detail::SpatComputeBatchBodyArgs>;
   if constexpr (std::is_same_v<BodyResult, void>) {
     std::forward<BodyFn>(body)(args);
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
   else {
     auto bodyResult = std::forward<BodyFn>(body)(args);
     if (mlir::failed(bodyResult)) {
-      rewriter.setInsertionPointAfter(batchOp);
-      rewriter.eraseOp(batchOp);
-      return mlir::FailureOr<spatial::SpatComputeBatch>(mlir::failure());
+      rewriter.setInsertionPointAfter(*batchOp);
+      rewriter.eraseOp(*batchOp);
+      return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
     }
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
 }
 
+template <size_t NumInputs, typename RewriterT, typename BodyFn>
+auto createSpatCompute(RewriterT& rewriter,
+                       mlir::Location loc,
+                       mlir::TypeRange resultTypes,
+                       mlir::ValueRange weights,
+                       mlir::ValueRange inputs,
+                       BodyFn&& body) {
+  return createSpatGraphCompute<NumInputs>(
+    rewriter, loc, resultTypes, weights, inputs, std::forward<BodyFn>(body));
+}
+
+template <typename RewriterT, typename BodyFn>
+auto createSpatCompute(RewriterT& rewriter,
+                       mlir::Location loc,
+                       mlir::TypeRange resultTypes,
+                       mlir::ValueRange weights,
+                       mlir::ValueRange inputs,
+                       BodyFn&& body) {
+  return createSpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, std::forward<BodyFn>(body));
+}
+
+template <typename RewriterT, typename BodyFn>
+auto createSpatComputeBatch(RewriterT& rewriter,
+                            mlir::Location loc,
+                            mlir::TypeRange resultTypes,
+                            int64_t laneCount,
+                            mlir::ValueRange weights,
+                            mlir::ValueRange inputs,
+                            BodyFn&& body) {
+  return createSpatGraphComputeBatch(
+    rewriter, loc, resultTypes, laneCount, weights, inputs, std::forward<BodyFn>(body));
+}
+
 inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewriter,
                                                      mlir::Location loc,
                                                      mlir::Value source,
@@ -244,6 +346,52 @@ inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewr
   mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
 }
 
+inline void publishGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                              mlir::Location loc,
+                                              mlir::Value fragment,
+                                              mlir::Value output,
+                                              mlir::Value physicalSlot) {
+  auto fragmentType = mlir::cast<mlir::RankedTensorType>(fragment.getType());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {physicalSlot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  createParallelInsertSliceIntoBatchOutput(rewriter, loc, fragment, output, offsets, sizes, strides);
+}
+
+inline mlir::FailureOr<mlir::Value>
+extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc,
+                                  mlir::Value physicalBatch,
+                                  mlir::OpFoldResult slot,
+                                  mlir::RankedTensorType fragmentType) {
+  if (fragmentType.getRank() == 0)
+    return mlir::failure();
+  auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
+  if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
+    return mlir::failure();
+  mlir::SmallVector<int64_t> selectedShape {1};
+  llvm::append_range(selectedShape, fragmentType.getShape());
+  auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
+  mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
+    reassociation.push_back({dim});
+  return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,
@@ -262,6 +410,6 @@ mlir::Value materializeOrComputeUnary(mlir::Value input,
   return computeOp.getResult(0);
 }
 
-mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::ConversionPatternRewriter& rewriter);
+mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::PatternRewriter& rewriter);
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,48 @@
+#include "MatrixProductLowering.hpp"
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
+  auto sourceType = cast<RankedTensorType>(value.getType());
+  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
+  SmallVector<OpFoldResult> highPads;
+  highPads.reserve(sourceType.getRank());
+  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
+    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
+
+  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
+  auto* padBlock = new Block();
+  for (int64_t i = 0; i < sourceType.getRank(); ++i)
+    padBlock->addArgument(rewriter.getIndexType(), loc);
+  padOp.getRegion().push_back(padBlock);
+  rewriter.setInsertionPointToStart(padBlock);
+  auto zero = getOrCreateConstant(
+    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
+  tensor::YieldOp::create(rewriter, loc, zero);
+  rewriter.setInsertionPointAfter(padOp);
+  return padOp.getResult();
+}
+
+Value createPaddedInputCompute(Value input,
+                               RankedTensorType paddedInputType,
+                               PatternRewriter& rewriter,
+                               Location loc) {
+  auto inputType = cast<RankedTensorType>(input.getType());
+  if (inputType == paddedInputType)
+    return input;
+
+  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
+    Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
+    spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
+  });
+  return computeOp.getResult(0);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace onnx_mlir {
+
+mlir::Value createZeroPaddedTensor(mlir::Value value,
+                                   mlir::RankedTensorType resultType,
+                                   mlir::PatternRewriter& rewriter,
+                                   mlir::Location loc);
+
+mlir::Value createPaddedInputCompute(mlir::Value input,
+                                     mlir::RankedTensorType paddedInputType,
+                                     mlir::PatternRewriter& rewriter,
+                                     mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,211 @@
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.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/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
+  return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+                               logicalType.getElementType(),
+                               logicalType.getEncoding());
+}
+
+RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
+  return spatial::getGraphBatchPhysicalResultType(logicalType.getDimSize(2), getRowStripFragmentType(logicalType));
+}
+
+std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
+  SmallVector<int64_t> offsets;
+  SmallVector<int64_t> sizes;
+  const int64_t channels = type.getDimSize(1);
+  const int64_t height = type.getDimSize(2);
+  const int64_t width = type.getDimSize(3);
+  offsets.reserve(height * 4);
+  sizes.reserve(height * 4);
+  for (int64_t row = 0; row < height; ++row) {
+    offsets.append({0, 0, row, 0});
+    sizes.append({1, channels, 1, width});
+  }
+  return {offsets, sizes};
+}
+
+Value extractRowStripFragment(Value storage,
+                              RankedTensorType logicalType,
+                              OpFoldResult row,
+                              PatternRewriter& rewriter,
+                              Location loc) {
+  return *extractGraphBatchPhysicalFragment(rewriter, loc, storage, row, getRowStripFragmentType(logicalType));
+}
+
+void insertRowStripFragment(Value fragment,
+                            Value output,
+                            RankedTensorType logicalType,
+                            OpFoldResult row,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  assert(fragment.getType() == getRowStripFragmentType(logicalType));
+  assert(output.getType() == getRowStripStorageType(logicalType));
+  auto slot = dyn_cast<Value>(row);
+  assert(slot && "row-strip graph publication requires a dynamic physical slot");
+  publishGraphBatchPhysicalFragment(rewriter, loc, fragment, output, slot);
+}
+
+FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
+                                                  RankedTensorType fragmentType,
+                                                  PatternRewriter& rewriter) {
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> values;
+  values.reserve(fragmentType.getNumElements());
+  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
+    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
+      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
+        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
+          values.push_back((*channelValues)[channel]);
+
+  auto attr = DenseElementsAttr::get(fragmentType, values);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
+}
+
+FailureOr<Value> createRowStripStorageFromRows(Value rows,
+                                               RankedTensorType logicalType,
+                                               PatternRewriter& rewriter,
+                                               Location loc) {
+  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
+  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
+    return failure();
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4)
+    return failure();
+  if (logicalType.getDimSize(0) != 1)
+    return failure();
+  if (rowsType.getElementType() != logicalType.getElementType())
+    return failure();
+
+  const int64_t channels = logicalType.getDimSize(1);
+  const int64_t height = logicalType.getDimSize(2);
+  const int64_t width = logicalType.getDimSize(3);
+  if (rowsType.getDimSize(0) != height * width)
+    return failure();
+  if (rowsType.getDimSize(1) != channels)
+    return failure();
+
+  auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
+  auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(
+    rewriter, loc, TypeRange {storageType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
+      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
+      Value rowSlice = tensor::ExtractSliceOp::create(
+        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
+      Value channelWidth = ONNXTransposeOp::create(
+        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
+      Value fragment = tensor::ExpandShapeOp::create(
+        rewriter, loc, fragmentType, channelWidth, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
+      insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+createRowStripAssemblyBlueprint(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return failure();
+
+  auto [offsets, sizes] = buildRowStripMetadata(logicalType);
+  int64_t height = logicalType.getDimSize(2);
+  SmallVector<int64_t> operandIndices(height, 0), sourceSlots, sourceOffsets(height, 0), strides(height * 4, 1);
+  for (int64_t row = 0; row < height; ++row)
+    sourceSlots.push_back(row);
+  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, storage, ValueRange {},
+    rewriter.getStringAttr("nchw"), rewriter.getStringAttr("nchw_row_strip"),
+    rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
+    rewriter.getStringAttr("nchw_row_strip_fragments"), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
+    rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+}
+
+FailureOr<Value>
+applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr))
+    return failure();
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          Value constant;
+                                          if (denseAttr.isSplat()) {
+                                            constant = getOrCreateConstant(
+                                              rewriter,
+                                              rewriter.getInsertionBlock()->getParentOp(),
+                                              DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue<Attribute>()),
+                                              fragmentType);
+                                          }
+                                          else {
+                                            FailureOr<Value> perChannel =
+                                              createPerChannelConstantFragment(denseAttr, fragmentType, rewriter);
+                                            if (failed(perChannel))
+                                              return failure();
+                                            constant = *perChannel;
+                                          }
+                                          fragment =
+                                            spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,69 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments";
+
+struct RowStripPhysicalValue {
+  mlir::Value storage;
+  mlir::RankedTensorType logicalType;
+  llvm::SmallVector<int64_t, 16> fragmentOffsets;
+  llvm::SmallVector<int64_t, 16> fragmentSizes;
+};
+
+std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
+buildRowStripMetadata(mlir::RankedTensorType type);
+
+mlir::RankedTensorType getRowStripFragmentType(mlir::RankedTensorType logicalType);
+
+mlir::RankedTensorType getRowStripStorageType(mlir::RankedTensorType logicalType);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentOffsets(mlir::PatternRewriter& rewriter,
+                                                                   mlir::OpFoldResult row);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentSizes(mlir::PatternRewriter& rewriter,
+                                                                 mlir::RankedTensorType logicalType);
+
+mlir::Value extractRowStripFragment(mlir::Value storage,
+                                    mlir::RankedTensorType logicalType,
+                                    mlir::OpFoldResult row,
+                                    mlir::PatternRewriter& rewriter,
+                                    mlir::Location loc);
+
+void insertRowStripFragment(mlir::Value fragment,
+                            mlir::Value output,
+                            mlir::RankedTensorType logicalType,
+                            mlir::OpFoldResult row,
+                            mlir::PatternRewriter& rewriter,
+                            mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> createPerChannelConstantFragment(mlir::DenseElementsAttr denseAttr,
+                                                              mlir::RankedTensorType fragmentType,
+                                                              mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
+                                                           mlir::RankedTensorType logicalType,
+                                                           mlir::PatternRewriter& rewriter,
+                                                           mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(mlir::Value storage,
+                                                               mlir::RankedTensorType logicalType,
+                                                               mlir::PatternRewriter& rewriter,
+                                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
+                                               mlir::RankedTensorType logicalType,
+                                               mlir::PatternRewriter& rewriter,
+                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
+                                                  mlir::RankedTensorType logicalType,
+                                                  mlir::Value bias,
+                                                  mlir::PatternRewriter& rewriter,
+                                                  mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -3,9 +3,6 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include <functional>
-
-#include "IndexingUtils.hpp"
 #include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -15,75 +12,8 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-bool hasStaticPositiveShape(ArrayRef<int64_t> shape) {
-  return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
-}
-
-bool hasStaticPositiveShape(RankedTensorType type) {
-  return type.hasStaticShape() && hasStaticPositiveShape(type.getShape());
-}
-
-int64_t getStaticShapeElementCount(ArrayRef<int64_t> shape) {
-  return std::accumulate(shape.begin(), shape.end(), int64_t {1}, std::multiplies<int64_t> {});
-}
-
-SmallVector<int64_t> permuteShape(ArrayRef<int64_t> shape, ArrayRef<int64_t> permutation) {
-  SmallVector<int64_t> permutedShape;
-  permutedShape.reserve(permutation.size());
-  for (int64_t axis : permutation)
-    permutedShape.push_back(shape[axis]);
-  return permutedShape;
-}
-
-SmallVector<int64_t> invertPermutation(ArrayRef<int64_t> permutation) {
-  SmallVector<int64_t> inversePermutation(permutation.size());
-  for (auto [newIndex, oldIndex] : llvm::enumerate(permutation))
-    inversePermutation[oldIndex] = static_cast<int64_t>(newIndex);
-  return inversePermutation;
-}
-
-FailureOr<SmallVector<int64_t>> getTransposePermutationChecked(std::optional<ArrayAttr> permAttr, int64_t rank) {
-  SmallVector<int64_t> permutation;
-  if (!permAttr) {
-    permutation.reserve(rank);
-    for (int64_t dim = rank - 1; dim >= 0; --dim)
-      permutation.push_back(dim);
-    return permutation;
-  }
-
-  if (static_cast<int64_t>(permAttr->size()) != rank)
-    return failure();
-
-  permutation.reserve(permAttr->size());
-  SmallVector<bool> seen(rank, false);
-  for (IntegerAttr attr : permAttr->getAsRange<IntegerAttr>()) {
-    int64_t axis = attr.getInt();
-    if (axis < 0 || axis >= rank || seen[axis])
-      return failure();
-    seen[axis] = true;
-    permutation.push_back(axis);
-  }
-  return permutation;
-}
-
-SmallVector<OpFoldResult> getUnitStrides(PatternRewriter& rewriter, int64_t rank) {
-  return SmallVector<OpFoldResult>(rank, rewriter.getIndexAttr(1));
-}
-
-SmallVector<OpFoldResult> getZeroOffsets(PatternRewriter& rewriter, int64_t rank) {
-  return SmallVector<OpFoldResult>(rank, rewriter.getIndexAttr(0));
-}
-
-SmallVector<OpFoldResult> getStaticSizes(PatternRewriter& rewriter, ArrayRef<int64_t> shape) {
-  SmallVector<OpFoldResult> sizes;
-  sizes.reserve(shape.size());
-  for (int64_t dim : shape)
-    sizes.push_back(rewriter.getIndexAttr(dim));
-  return sizes;
-}
-
 SmallVector<Value> sliceTensor(
-  const Value& tensorToSlice, size_t axis, int64_t sliceSize, ConversionPatternRewriter& rewriter, Location loc) {
+  const Value& tensorToSlice, size_t axis, int64_t sliceSize, PatternRewriter& rewriter, Location loc) {
   ArrayRef<long> shape = getTensorShape(tensorToSlice);
   assert("Invalid axis" && axis < shape.size());
 
@@ -129,7 +59,7 @@ SmallVector<Value> sliceTensor(
 }
 
 SmallVector<Value>
-sliceVector(const Value& vectorToSlice, int64_t sliceSize, ConversionPatternRewriter& rewriter, Location loc) {
+sliceVector(const Value& vectorToSlice, int64_t sliceSize, PatternRewriter& rewriter, Location loc) {
   ArrayRef<long> shape = getTensorShape(vectorToSlice);
   assert("Not a vector" && isVectorShape(shape));
   size_t axis = shape[0] != 1 ? 0 : 1;
@@ -137,7 +67,7 @@ sliceVector(const Value& vectorToSlice, int64_t sliceSize, ConversionPatternRewr
 }
 
 DenseMap<CoreId, SmallVector<Value>>
-sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, ConversionPatternRewriter& rewriter, Location loc) {
+sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, PatternRewriter& rewriter, Location loc) {
   SmallVector<Value> slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc);
   DenseMap<CoreId, SmallVector<Value>> slicesPerCore;
   for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) {
@@ -147,33 +77,4 @@ sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, ConversionPatternRewri
   return slicesPerCore;
 }
 
-Value extractAxisSlice(
-  PatternRewriter& rewriter, Location loc, Value source, int64_t axis, int64_t offset, int64_t size) {
-  auto sourceType = cast<RankedTensorType>(source.getType());
-  SmallVector<int64_t> resultShape(sourceType.getShape());
-  resultShape[axis] = size;
-  auto resultType = RankedTensorType::get(resultShape, sourceType.getElementType(), sourceType.getEncoding());
-
-  SmallVector<OpFoldResult> offsets = getZeroOffsets(rewriter, sourceType.getRank());
-  SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, sourceType.getShape());
-  offsets[axis] = rewriter.getIndexAttr(offset);
-  sizes[axis] = rewriter.getIndexAttr(size);
-  return tensor::ExtractSliceOp::create(
-           rewriter, loc, resultType, source, offsets, sizes, getUnitStrides(rewriter, sourceType.getRank()))
-    .getResult();
-}
-
-Value insertStaticSlice(
-  PatternRewriter& rewriter, Location loc, Value source, Value dest, ArrayRef<OpFoldResult> offsets) {
-  auto sourceType = cast<RankedTensorType>(source.getType());
-  return tensor::InsertSliceOp::create(rewriter,
-                                       loc,
-                                       source,
-                                       dest,
-                                       offsets,
-                                       getStaticSizes(rewriter, sourceType.getShape()),
-                                       getUnitStrides(rewriter, sourceType.getRank()))
-    .getResult();
-}
-
 } // namespace onnx_mlir

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

@@ -1,114 +1,31 @@
 #pragma once
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Transforms/DialectConversion.h"
 
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 
-#include <cassert>
-#include <cstddef>
-#include <optional>
-#include <type_traits>
-#include <utility>
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 
 namespace onnx_mlir {
 
-using HSliceId = size_t;
-using CoreId = size_t;
-
-template <class A, class B, class C = std::common_type_t<A, B>>
-constexpr C ceilIntegerDivide(A a, B b) {
-  static_assert(std::is_integral_v<A>, "A must be an integer type");
-  static_assert(std::is_integral_v<B>, "B must be an integer type");
-  C ac = static_cast<C>(a);
-  C bc = static_cast<C>(b);
-  return 1 + (ac - 1) / bc;
-}
-
-template <class A, class B, class C = std::common_type_t<A, B>>
-constexpr std::pair<C, C> ceilIntegerDivideWithRemainder(A a, B b) {
-  static_assert(std::is_integral_v<A>, "A must be an integer type");
-  static_assert(std::is_integral_v<B>, "B must be an integer type");
-  C ac = static_cast<C>(a);
-  C bc = static_cast<C>(b);
-  return {ceilIntegerDivide(ac, bc), ac % bc};
-}
-
-template <class T>
-bool isVectorShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2 && (shape[0] == 1 || shape[1] == 1);
-}
-
-template <class T>
-bool isMatrixShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2;
-}
-
-template <class T>
-bool isHVectorShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2 && shape[0] == 1;
-}
-
-inline auto getTensorShape(mlir::Value tensor) {
-  return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
-}
-
-inline bool haveSameStaticShape(mlir::Value lhs, mlir::Value rhs) {
-  auto lhsType = mlir::dyn_cast<mlir::RankedTensorType>(lhs.getType());
-  auto rhsType = mlir::dyn_cast<mlir::RankedTensorType>(rhs.getType());
-  return lhsType && rhsType && lhsType.hasStaticShape() && rhsType.hasStaticShape()
-      && lhsType.getShape() == rhsType.getShape();
-}
-
-bool hasStaticPositiveShape(mlir::ArrayRef<int64_t> shape);
-
-bool hasStaticPositiveShape(mlir::RankedTensorType type);
-
-int64_t getStaticShapeElementCount(mlir::ArrayRef<int64_t> shape);
-
-llvm::SmallVector<int64_t> permuteShape(mlir::ArrayRef<int64_t> shape, mlir::ArrayRef<int64_t> permutation);
-
-llvm::SmallVector<int64_t> invertPermutation(mlir::ArrayRef<int64_t> permutation);
-
-mlir::FailureOr<llvm::SmallVector<int64_t>> getTransposePermutationChecked(std::optional<mlir::ArrayAttr> permAttr,
-                                                                           int64_t rank);
-
-llvm::SmallVector<mlir::OpFoldResult> getUnitStrides(mlir::PatternRewriter& rewriter, int64_t rank);
-
-llvm::SmallVector<mlir::OpFoldResult> getZeroOffsets(mlir::PatternRewriter& rewriter, int64_t rank);
-
-llvm::SmallVector<mlir::OpFoldResult> getStaticSizes(mlir::PatternRewriter& rewriter, mlir::ArrayRef<int64_t> shape);
-
 /// Slices a statically shaped tensor along one axis into contiguous pieces of
 /// at most `sliceSize` elements.
 llvm::SmallVector<mlir::Value> sliceTensor(const mlir::Value& tensorToSlice,
                                            size_t axis,
                                            int64_t sliceSize,
-                                           mlir::ConversionPatternRewriter& rewriter,
+                                           mlir::PatternRewriter& rewriter,
                                            mlir::Location loc);
 
 llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
                                            int64_t sliceSize,
-                                           mlir::ConversionPatternRewriter& rewriter,
+                                           mlir::PatternRewriter& rewriter,
                                            mlir::Location loc);
 
 /// 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::ConversionPatternRewriter& rewriter, mlir::Location loc);
-
-mlir::Value extractAxisSlice(
-  mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
-
-mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
-                              mlir::Location loc,
-                              mlir::Value source,
-                              mlir::Value dest,
-                              llvm::ArrayRef<mlir::OpFoldResult> offsets);
+  const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc);
 
 } // namespace onnx_mlir

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

@@ -19,9 +19,11 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-bool isWeightLikeComputeOperand(Value value) {
+static bool isWeightMaterializationValue(Value value, bool requireMatrixShape) {
   auto rankedType = dyn_cast<RankedTensorType>(value.getType());
-  if (!rankedType || !isMatrixShape(rankedType.getShape()))
+  if (!rankedType)
+    return false;
+  if (requireMatrixShape && !isMatrixShape(rankedType.getShape()))
     return false;
 
   llvm::SmallPtrSet<Operation*, 8> visited;
@@ -29,8 +31,14 @@ bool isWeightLikeComputeOperand(Value value) {
   while (auto* definingOp = value.getDefiningOp()) {
     if (!visited.insert(definingOp).second)
       return false;
-    if (isa<arith::ConstantOp, ONNXConstantOp>(definingOp) || hasWeightAlways(definingOp))
+    if (isa<arith::ConstantOp, ONNXConstantOp>(definingOp) || hasWeightAlways(definingOp)) {
+      auto sourceType = dyn_cast<RankedTensorType>(value.getType());
+      if (!sourceType)
+        return false;
+      if (requireMatrixShape && !isMatrixShape(sourceType.getShape()))
+        return false;
       return true;
+    }
 
     if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {
       value = extractSliceOp.getSource();
@@ -55,6 +63,8 @@ bool isWeightLikeComputeOperand(Value value) {
   return false;
 }
 
+bool isWeightLikeComputeOperand(Value value) { return isWeightMaterializationValue(value, /*requireMatrixShape=*/true); }
+
 FailureOr<Value> materializeWeightLikeValueInBlock(Value value, IRRewriter& rewriter, IRMapping& mapper) {
   if (auto mapped = mapper.lookupOrNull(value))
     return cast<Value>(mapped);
@@ -91,7 +101,7 @@ FailureOr<Value> materializeWeightLikeValueInBlock(Value value, IRRewriter& rewr
       continue;
     }
 
-    if (isWeightLikeComputeOperand(operand)) {
+    if (isWeightMaterializationValue(operand, /*requireMatrixShape=*/false)) {
       auto clonedOperand = materializeWeightLikeValueInBlock(operand, rewriter, mapper);
       if (failed(clonedOperand))
         return failure();

src/PIM/Conversion/ONNXToSpatial/CompileTime.cpp

@@ -9,10 +9,12 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
+#include <cstring>
 #include <utility>
 
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -21,24 +23,7 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-namespace {
-
-static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
-  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
-}
-
-static bool hasConstantIndices(tensor::ExtractOp extractOp) {
-  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
-}
-
-static bool isStaticTensorResult(Operation* op) {
-  return llvm::all_of(op->getResultTypes(), [](Type type) {
-    auto shapedType = dyn_cast<ShapedType>(type);
-    return shapedType && shapedType.hasStaticShape();
-  });
-}
-
-static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
+FailureOr<DenseElementsAttr> transposeDenseElementsAttr(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
   auto tensorType = dyn_cast<RankedTensorType>(denseAttr.getType());
   if (!tensorType)
     return failure();
@@ -59,7 +44,45 @@ static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr den
 
   auto transposedType = RankedTensorType::get(transposedShape, tensorType.getElementType(), tensorType.getEncoding());
   if (denseAttr.isSplat())
-    return DenseElementsAttr::get(transposedType, denseAttr.getSplatValue<Attribute>());
+    return DenseElementsAttr::getFromRawBuffer(transposedType, denseAttr.getRawData());
+
+  const unsigned elementBitWidth = tensorType.getElementTypeBitWidth();
+  const ArrayRef<char> inputData = denseAttr.getRawData();
+  if (elementBitWidth % 8 == 0) {
+    const size_t elementBytes = elementBitWidth / 8;
+    const size_t expectedBytes = denseAttr.getNumElements() * elementBytes;
+    if (inputData.size() == expectedBytes) {
+      SmallVector<char> transposedData(expectedBytes);
+      if (rank == 2 && perms[0] == 1 && perms[1] == 0) {
+        const int64_t rows = tensorType.getDimSize(0);
+        const int64_t columns = tensorType.getDimSize(1);
+        for (int64_t row = 0; row < rows; ++row)
+          for (int64_t column = 0; column < columns; ++column)
+            std::memcpy(transposedData.data() + (column * rows + row) * elementBytes,
+                        inputData.data() + (row * columns + column) * elementBytes,
+                        elementBytes);
+        return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData);
+      }
+
+      SmallVector<int64_t> originalStrides = computeRowMajorStrides(tensorType.getShape());
+      SmallVector<int64_t> transposedStrides = computeRowMajorStrides(transposedShape);
+      SmallVector<int64_t> originalIndices(rank);
+      for (int64_t linearIndex = 0; linearIndex < tensorType.getNumElements(); ++linearIndex) {
+        int64_t remaining = linearIndex;
+        for (int64_t dim = 0; dim < rank; ++dim) {
+          originalIndices[dim] = originalStrides.empty() ? 0 : remaining / originalStrides[dim];
+          remaining = originalStrides.empty() ? 0 : remaining % originalStrides[dim];
+        }
+        int64_t transposedLinearIndex = 0;
+        for (int64_t dim = 0; dim < rank; ++dim)
+          transposedLinearIndex += originalIndices[perms[dim]] * transposedStrides[dim];
+        std::memcpy(transposedData.data() + transposedLinearIndex * elementBytes,
+                    inputData.data() + linearIndex * elementBytes,
+                    elementBytes);
+      }
+      return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData);
+    }
+  }
 
   SmallVector<Attribute> originalValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> transposedValues(originalValues.size());
@@ -84,16 +107,30 @@ static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr den
   return DenseElementsAttr::get(transposedType, transposedValues);
 }
 
+namespace {
+
+static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
+  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
+}
+
+static bool hasConstantIndices(tensor::ExtractOp extractOp) {
+  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
+}
+
+static bool isStaticTensorResult(Operation* op) {
+  return llvm::all_of(op->getResultTypes(), [](Type type) {
+    auto shapedType = dyn_cast<ShapedType>(type);
+    return shapedType && shapedType.hasStaticShape();
+  });
+}
+
 static FailureOr<DenseElementsAttr> reshapeDenseElements(DenseElementsAttr denseAttr, RankedTensorType resultType) {
   auto sourceType = dyn_cast<RankedTensorType>(denseAttr.getType());
-  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements())
+  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getElementType() != resultType.getElementType())
     return failure();
 
-  if (denseAttr.isSplat())
-    return DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>());
-
-  SmallVector<Attribute> values(denseAttr.getValues<Attribute>());
-  return DenseElementsAttr::get(resultType, values);
+  return DenseElementsAttr::getFromRawBuffer(resultType, denseAttr.getRawData());
 }
 
 static FailureOr<DenseElementsAttr> extractSliceDenseElements(DenseElementsAttr denseAttr,
@@ -161,7 +198,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
     perm.reserve(transposeOp.getPermAttr().size());
     for (IntegerAttr attr : transposeOp.getPermAttr().getAsRange<IntegerAttr>())
       perm.push_back(attr.getInt());
-    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm);
     return succeeded(transposedAttr) ? *transposedAttr : nullptr;
   }
 
@@ -171,7 +208,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
       return nullptr;
 
     SmallVector<int64_t> perm(transposeOp.getPermutation().begin(), transposeOp.getPermutation().end());
-    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm);
     return succeeded(transposedAttr) ? *transposedAttr : nullptr;
   }
 
@@ -219,6 +256,9 @@ getCompileTimeSourceImpl(Operation* op, llvm::SmallPtrSetImpl<Operation*>& visit
 
   chainLength += 1;
 
+  if (!isShapingOnlyOp(op))
+    return std::nullopt;
+
   if (auto extractOp = dyn_cast<tensor::ExtractOp>(op))
     return hasConstantIndices(extractOp)
            ? getCompileTimeSourceImpl(extractOp.getTensor().getDefiningOp(), visited, chainLength)

src/PIM/Conversion/ONNXToSpatial/CompileTime.hpp

@@ -4,6 +4,8 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
 
+#include "llvm/ADT/ArrayRef.h"
+
 namespace onnx_mlir {
 
 struct CompileTimeSource {
@@ -19,4 +21,7 @@ bool isCompileTimeOp(mlir::Operation* op);
 
 mlir::DenseElementsAttr getHostConstDenseElementsAttr(mlir::Value value);
 
+mlir::FailureOr<mlir::DenseElementsAttr> transposeDenseElementsAttr(
+  mlir::DenseElementsAttr denseAttr, llvm::ArrayRef<int64_t> permutation);
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -0,0 +1,422 @@
+#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 = "nchw_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();
+  RowStripPhysicalValue value;
+  value.storage = storage;
+  value.logicalType = logicalType;
+  value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end());
+  value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end());
+  if (blueprint.getIndexMap() != kRowStripIndexMap)
+    return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    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.storage, input.logicalType, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
+                                             spatial::SpatBiasAddPlanOp planOp,
+                                             PatternRewriter& rewriter) {
+  return applyRowStripBiasAdd(input.storage, input.logicalType, planOp.getBias(), rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value>
+materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
+  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
+    return failure();
+  auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
+  if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
+    return failure();
+  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
+}
+
+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);
+          FailureOr<Value> lowered = lowerSelectedConv2DPlan(
+            planOp,
+            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->storage} : std::nullopt,
+            /*emitRowStripLayout=*/true,
+            rewriter);
+          if (failed(lowered)) {
+            planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
+            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::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;
+        }
+        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 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::SpatReluPlanOp,
+                     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/ONNXToSpatialPass.cpp

@@ -13,11 +13,13 @@
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/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;
 
@@ -41,10 +43,17 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
 static void populateEmptyFunction(func::FuncOp funcOp) {
   IRRewriter rewriter(funcOp.getContext());
   IRMapping mapper;
-  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
-  SmallVector<spatial::SpatComputeBatch> computeBatches(funcOp.getOps<spatial::SpatComputeBatch>());
-  if (!computes.empty() || !computeBatches.empty())
+  SmallVector<spatial::SpatGraphCompute> computes(funcOp.getOps<spatial::SpatGraphCompute>());
+  SmallVector<spatial::SpatGraphComputeBatch> computeBatches(funcOp.getOps<spatial::SpatGraphComputeBatch>());
+  SmallVector<spatial::SpatConv2DPlanOp> convPlans(funcOp.getOps<spatial::SpatConv2DPlanOp>());
+  SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
+  SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
+  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() || !reluPlans.empty()
+      || !blueprints.empty() || !materializers.empty()) {
     return;
+  }
 
   auto returnOp = cast<func::ReturnOp>(funcOp.getFunctionBody().front().getTerminator());
   rewriter.setInsertionPoint(returnOp);
@@ -58,16 +67,16 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     sourceLocs.push_back(source.getLoc());
   }
 
-  auto newCompute = spatial::SpatCompute::create(
-    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
-  auto* newBlock = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLocs);
+  auto newCompute = createEmptySpatGraphCompute(
+    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), {}, funcOp.getArguments(), sourceTypes, sourceLocs);
+  auto* newBlock = &newCompute.getBody().front();
   for (auto [blockArg, computeArg] : llvm::zip(newBlock->getArguments(), newCompute.getOperands()))
     mapper.map(computeArg, blockArg);
   newCompute.getProperties().setOperandSegmentSizes({0, static_cast<int>(sourceTypes.size())});
 
   rewriter.setInsertionPointToEnd(newBlock);
   for (Operation& op : funcOp.getOps())
-    if (!isa<spatial::SpatCompute, func::ReturnOp>(&op))
+    if (!isa<spatial::SpatGraphCompute, func::ReturnOp>(&op))
       rewriter.clone(op, mapper);
 
   auto yield = spatial::SpatYieldOp::create(rewriter, funcOp.getLoc(), returnOp.getOperands());
@@ -75,7 +84,7 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     yield.setOperand(i, mapper.lookupOrDefault(yield.getOperand(i)));
 
   for (Operation& op : llvm::make_early_inc_range(funcOp.getOps()))
-    if (!isa<spatial::SpatCompute, func::ReturnOp>(&op)) {
+    if (!isa<spatial::SpatGraphCompute, func::ReturnOp>(&op)) {
       op.dropAllUses();
       rewriter.eraseOp(&op);
     }
@@ -96,7 +105,7 @@ void ONNXToSpatialPass::runOnOperation() {
                             affine::AffineDialect,
                             arith::ArithDialect,
                             scf::SCFDialect>();
-  preTarget.addIllegalOp<ONNXConstantOp, ONNXFlattenOp>();
+  preTarget.addIllegalOp<ONNXConstantOp>();
 
   RewritePatternSet prePatterns(ctx);
   populatePrePatterns(prePatterns, ctx);
@@ -124,6 +133,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXMatMulOp>();
   target.addIllegalOp<ONNXTransposeOp>();
   target.addIllegalOp<ONNXAddOp>();
+  target.addIllegalOp<ONNXSubOp>();
   target.addIllegalOp<ONNXDivOp>();
   target.addIllegalOp<ONNXMulOp>();
   target.addIllegalOp<ONNXGemmOp>();
@@ -134,10 +144,13 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXSigmoidOp>();
   target.addIllegalOp<ONNXSoftmaxOp>();
   target.addIllegalOp<ONNXConcatOp>();
+  target.addIllegalOp<ONNXFlattenOp>();
   target.addIllegalOp<ONNXGatherOp>();
   target.addIllegalOp<ONNXReshapeOp>();
   target.addIllegalOp<ONNXResizeOp>();
+  target.addIllegalOp<ONNXSliceOp>();
   target.addIllegalOp<ONNXLRNOp>();
+  target.addIllegalOp<ONNXReduceMeanOp>();
   target.addIllegalOp<ONNXReduceMeanV13Op>();
   target.addIllegalOp<ONNXSplitOp>();
 
@@ -149,6 +162,11 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("logical Spatial graph verification failed after ONNX conversion");
+      signalPassFailure();
+      return;
+    }
   ConversionTarget earlyPostTarget(*ctx);
   earlyPostTarget.addLegalDialect<spatial::SpatialDialect,
                                   ONNXDialect,
@@ -158,13 +176,13 @@ void ONNXToSpatialPass::runOnOperation() {
                                   arith::ArithDialect,
                                   scf::SCFDialect>();
 
-  PassManager cleanupPM(ctx);
-  cleanupPM.addPass(createCanonicalizerPass());
-  if (failed(cleanupPM.run(moduleOp)))
-    moduleOp.emitWarning("failed to run ONNX-to-Spatial canonicalization cleanup; continuing");
-
   annotateWeightsConstants(*entryFunc);
 
+    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("logical Spatial graph verification failed after weight annotation");
+      signalPassFailure();
+      return;
+    }
   ConversionTarget postTarget(*ctx);
   postTarget.addLegalDialect<spatial::SpatialDialect,
                              ONNXDialect,
@@ -173,11 +191,16 @@ void ONNXToSpatialPass::runOnOperation() {
                              affine::AffineDialect,
                              arith::ArithDialect,
                              scf::SCFDialect>();
-  postTarget.addDynamicallyLegalOp<spatial::SpatCompute>(
-    [](spatial::SpatCompute computeOp) { return !requiresPostRewrite(computeOp); });
-  postTarget.addDynamicallyLegalOp<spatial::SpatComputeBatch>(
-    [](spatial::SpatComputeBatch computeOp) { return !requiresPostRewrite(computeOp); });
+  postTarget.addDynamicallyLegalOp<spatial::SpatGraphCompute>(
+    [](spatial::SpatGraphCompute computeOp) { return !requiresPostRewrite(computeOp); });
+  postTarget.addDynamicallyLegalOp<spatial::SpatGraphComputeBatch>(
+    [](spatial::SpatGraphComputeBatch computeOp) { return !requiresPostRewrite(computeOp); });
 
+    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("logical Spatial graph verification failed before post rewrites");
+      signalPassFailure();
+      return;
+    }
   RewritePatternSet postPatterns(ctx);
   populatePostPatterns(postPatterns, ctx);
   if (failed(applyPartialConversion(*entryFunc, postTarget, std::move(postPatterns)))) {
@@ -188,8 +211,18 @@ void ONNXToSpatialPass::runOnOperation() {
 
   populateEmptyFunction(*entryFunc);
 
+  PassManager canonicalizationPM(ctx);
+  canonicalizationPM.addPass(createCanonicalizerPass());
+  if (failed(canonicalizationPM.run(moduleOp)))
+    moduleOp.emitWarning("failed to run ONNXToSpatial canonicalization; continuing");
+
   dumpModule(moduleOp, "spatial0");
 
+  if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("logical Spatial graph verification failed after ONNX-to-Spatial");
+      signalPassFailure();
+      return;
+    }
   if (failed(verifyONNXToSpatial(*entryFunc))) {
     moduleOp.emitError("ONNX-to-Spatial host legality verification failed");
     signalPassFailure();

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -1,4 +1,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/Support/LLVM.h"
 
 #include "Common/IR/WeightUtils.hpp"
@@ -13,6 +15,8 @@ namespace onnx_mlir {
 
 namespace {
 
+constexpr StringLiteral kPhaseMarker = "phase-check";
+
 void checkWeightUseChains(func::FuncOp func, pim::CappedDiagnosticReporter& diagnostics) {
   func.walk([&](Operation* op) {
     if (!hasWeightAlways(op))
@@ -23,134 +27,205 @@ void checkWeightUseChains(func::FuncOp func, pim::CappedDiagnosticReporter& diag
         continue;
 
       diagnostics.report(op, [&](Operation* illegalOp) {
-        illegalOp->emitOpError(
-          "weight-marked values may only flow through static view/slice helper chains into Spatial VMM weights");
+        illegalOp->emitOpError()
+          << kPhaseMarker
+          << " weight-marked values may only flow through static view/slice helper chains into Spatial VMM weights";
       });
       return;
     }
   });
 }
 
-Region* getParentRegion(Value value) {
-  if (auto blockArg = dyn_cast<BlockArgument>(value))
-    return blockArg.getOwner()->getParent();
-  if (Operation* definingOp = value.getDefiningOp())
-    return definingOp->getParentRegion();
-  return nullptr;
+bool isRegionOrAncestorOf(Region& region, Region* candidate) {
+  return candidate && (&region == candidate || region.isAncestor(candidate));
 }
 
-bool isDefinedInsideRegion(Value value, Region& region) {
-  Region* parentRegion = getParentRegion(value);
-  return parentRegion && (&region == parentRegion || region.isAncestor(parentRegion));
+bool isValueDefinedInsideRegion(Value value, Region& region) {
+  if (auto blockArg = dyn_cast<BlockArgument>(value))
+    return isRegionOrAncestorOf(region, blockArg.getOwner()->getParent());
+  if (Operation* definingOp = value.getDefiningOp())
+    return isRegionOrAncestorOf(region, definingOp->getParentRegion());
+  return false;
+}
+
+bool isLegalExternalCapture(Value value, Region& region) {
+  if (isValueDefinedInsideRegion(value, region))
+    return true;
+
+  Operation* definingOp = value.getDefiningOp();
+  return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
+}
+
+bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
+  auto transfer = dyn_cast<spatial::SpatDeferredCommunicationOp>(op);
+  return transfer && llvm::is_contained(transfer.getSources(), value);
+}
+
+template <typename ComputeOpTy>
+void verifyComputeBodyCaptures(ComputeOpTy compute, StringRef kind, pim::CappedDiagnosticReporter& diagnostics) {
+  Region& body = compute.getBody();
+  body.walk([&](Operation* nestedOp) {
+    for (OpOperand& operand : nestedOp->getOpOperands()) {
+      Value value = operand.get();
+      if (isLegalExternalCapture(value, body) || isRecordedDeferredCommunicationSource(nestedOp, value))
+        continue;
+
+      Operation* definingOp = value.getDefiningOp();
+      diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) {
+        InFlightDiagnostic diag =
+          illegalOp->emitOpError() << kPhaseMarker << " " << kind << " body captures non-constant external operand #"
+                                   << operand.getOperandNumber() << " used by " << nestedOp->getName().getStringRef();
+        diag << " (type " << value.getType() << ")";
+        if (definingOp)
+          diag.attachNote(definingOp->getLoc()) << "defining op is " << definingOp->getName().getStringRef();
+        else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
+          if (Operation* owner = blockArg.getOwner()->getParentOp())
+            diag.attachNote(owner->getLoc())
+              << "external block argument belongs to " << owner->getName().getStringRef();
+        }
+      });
+    }
+  });
 }
 
 bool isLegalHostBackedValue(Value value) {
   Operation* definingOp = value.getDefiningOp();
   if (!definingOp)
     return isa<BlockArgument>(value);
-
-  if (isa<spatial::SpatChannelReceiveOp>(definingOp))
-    return false;
-
   return definingOp->getDialect()->getNamespace() != "spat";
 }
 
-LogicalResult verifyComputeLikeInputs(Operation* computeLikeOp,
-                                      ValueRange inputs,
-                                      bool allowChannelReceiveInputs,
-                                      StringRef kind,
-                                      pim::CappedDiagnosticReporter& diagnostics) {
-  for (auto [inputIndex, input] : llvm::enumerate(inputs)) {
-    unsigned currentInputIndex = inputIndex;
+bool isScheduledPhase1Value(Value value) {
+  Operation* definingOp = value.getDefiningOp();
+  return isa_and_nonnull<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch>(definingOp);
+}
+
+template <typename ComputeOpTy>
+void verifyScheduledInputs(ComputeOpTy compute,
+                           bool allowChannelReceiveInputs,
+                           StringRef kind,
+                           pim::CappedDiagnosticReporter& diagnostics) {
+  for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
+    size_t currentInputIndex = inputIndex;
     Operation* definingOp = input.getDefiningOp();
     if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
       continue;
+    if (isScheduledPhase1Value(input))
+      continue;
     if (isLegalHostBackedValue(input))
       continue;
 
-    diagnostics.report(computeLikeOp, [&](Operation* illegalOp) {
-      InFlightDiagnostic diagnostic = illegalOp->emitOpError()
-                                   << kind << " input #" << currentInputIndex
-                                   << (allowChannelReceiveInputs ? " must come from the host or an explicit "
-                                                                   "spat.channel_receive"
-                                                                 : " must come from the host");
+    diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) {
+      InFlightDiagnostic diag = illegalOp->emitOpError()
+                                << kPhaseMarker << " " << kind << " input #" << currentInputIndex
+                                << (allowChannelReceiveInputs ? " must come from the host or explicit spat.channel_receive"
+                                                              : " must come from the host");
       if (definingOp)
-        diagnostic.attachNote(definingOp->getLoc()) << "illegal Spatial producer is " << definingOp->getName();
+        diag.attachNote(definingOp->getLoc()) << "illegal producer is " << definingOp->getName().getStringRef();
     });
-    return failure();
   }
-  return success();
 }
 
-void verifyNoExternalTensorCaptures(Operation* ownerOp,
-                                    Region& region,
-                                    StringRef kind,
-                                    pim::CappedDiagnosticReporter& diagnostics) {
-  region.walk([&](Operation* op) {
-    for (OpOperand& operand : op->getOpOperands()) {
-      Value value = operand.get();
-      if (!isa<TensorType>(value.getType()))
-        continue;
-      if (isDefinedInsideRegion(value, region) || isa<BlockArgument>(value))
-        continue;
+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")
+      return;
+    diagnostics.report(blueprint.getOperation(), [&](Operation* illegalOp) {
+      illegalOp->emitOpError("fragment assembly blueprint must be host-level after merge materialization");
+    });
+  });
+}
 
-      Operation* definingOp = value.getDefiningOp();
-      if (definingOp && definingOp->hasTrait<OpTrait::ConstantLike>())
-        continue;
+void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
+  for (Operation& op : funcOp.getOps()) {
+    if (isa<func::ReturnOp,
+            spatial::SpatGraphCompute,
+            spatial::SpatGraphComputeBatch,
+            spatial::SpatConv2DPlanOp,
+            spatial::SpatBiasAddPlanOp,
+            spatial::SpatReluPlanOp,
+            spatial::SpatBlueprintOp,
+            spatial::SpatMaterializeLayoutOp>(&op)) {
+      continue;
+    }
+    if (isa<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch>(&op)) {
+      diagnostics.report(&op, [&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kPhaseMarker << " scheduled Spatial compute op is not allowed in logical graph phase";
+      });
+      continue;
+    }
+    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp>(&op)) {
+      diagnostics.report(&op, [&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kPhaseMarker
+                                 << " explicit channel communication is not expected before merge materialization";
+      });
+      continue;
+    }
+    if (isCompileTimeOp(&op))
+      continue;
 
-      diagnostics.report(ownerOp, [&](Operation* illegalOp) {
-        InFlightDiagnostic diagnostic = illegalOp->emitOpError() << kind << " body may not capture external tensor "
-                                                                 << "values";
-        diagnostic.attachNote(op->getLoc())
-          << "tensor operand #" << operand.getOperandNumber() << " is defined outside the compute body by "
-          << (definingOp ? definingOp->getName().getStringRef() : StringRef("<block argument>"));
+    diagnostics.report(&op, [&](Operation* illegalOp) {
+      illegalOp->emitOpError()
+        << kPhaseMarker << " non-foldable top-level runtime op remains in logical Spatial graph; lower it inside spat.graph_compute";
+    });
+  }
+}
+
+void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
+  for (Operation& op : funcOp.getOps()) {
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
+      diagnostics.report(&op, [&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
       });
     }
-  });
+  }
 }
 
 } // namespace
 
-LogicalResult verifyONNXToSpatial(func::FuncOp funcOp) {
+LogicalResult verifyNoComputeBodyCaptures(func::FuncOp funcOp) {
   pim::CappedDiagnosticReporter diagnostics;
-
-  for (Operation& op : funcOp.getOps()) {
-    if (isa<func::ReturnOp, spatial::SpatCompute, spatial::SpatComputeBatch>(&op))
-      continue;
-    if (isCompileTimeOp(&op))
-      continue;
-
-    diagnostics.report(&op, [](Operation* illegalOp) {
-      illegalOp->emitOpError(
-        "non-foldable top-level runtime op remains after ONNX-to-Spatial; lower it inside spat.compute");
-    });
-  }
-  checkWeightUseChains(funcOp, diagnostics);
-  diagnostics.emitSuppressedSummary(funcOp, "ONNX-to-Spatial verification failed");
-
+  for (auto compute : funcOp.getOps<spatial::SpatGraphCompute>())
+    verifyComputeBodyCaptures(compute, "graph_compute", diagnostics);
+  for (auto batch : funcOp.getOps<spatial::SpatGraphComputeBatch>())
+    verifyComputeBodyCaptures(batch, "graph_compute_batch", diagnostics);
+  for (auto compute : funcOp.getOps<spatial::SpatScheduledCompute>())
+    verifyComputeBodyCaptures(compute, "scheduled_compute", diagnostics);
+  for (auto batch : funcOp.getOps<spatial::SpatScheduledComputeBatch>())
+    verifyComputeBodyCaptures(batch, "scheduled_compute_batch", diagnostics);
+  diagnostics.emitSuppressedSummary(funcOp, "compute body capture verification failed");
   return success(!diagnostics.hasFailure());
 }
 
-LogicalResult verifySpatialCommunicationInvariants(func::FuncOp funcOp) {
+LogicalResult verifyONNXToSpatial(func::FuncOp funcOp) { return verifyLogicalSpatialGraphInvariants(funcOp); }
+
+LogicalResult verifyLogicalSpatialGraphInvariants(func::FuncOp funcOp) {
   pim::CappedDiagnosticReporter diagnostics;
+  verifyLogicalTopLevelOps(funcOp, diagnostics);
+  checkWeightUseChains(funcOp, diagnostics);
+  if (failed(verifyNoComputeBodyCaptures(funcOp)))
+    return failure();
+  diagnostics.emitSuppressedSummary(funcOp, "logical Spatial graph verification failed");
+  return success(!diagnostics.hasFailure());
+}
 
-  for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
-    (void) verifyComputeLikeInputs(
-      computeOp.getOperation(), computeOp.getInputs(), /*allowChannelReceiveInputs=*/true, "spat.compute", diagnostics);
-    verifyNoExternalTensorCaptures(computeOp.getOperation(), computeOp.getBody(), "spat.compute", diagnostics);
+LogicalResult verifyScheduledSpatialInvariants(func::FuncOp funcOp) {
+  pim::CappedDiagnosticReporter diagnostics;
+  verifyScheduledTopLevelOps(funcOp, diagnostics);
+  for (auto compute : funcOp.getOps<spatial::SpatScheduledCompute>()) {
+    verifyScheduledInputs(compute, /*allowChannelReceiveInputs=*/true, "spat.scheduled_compute", diagnostics);
+    verifyNoNestedFragmentAssemblyBlueprints(compute, diagnostics);
   }
-
-  for (auto computeBatchOp : funcOp.getOps<spatial::SpatComputeBatch>()) {
-    (void) verifyComputeLikeInputs(computeBatchOp.getOperation(),
-                                   computeBatchOp.getInputs(),
-                                   /*allowChannelReceiveInputs=*/false,
-                                   "spat.compute_batch",
-                                   diagnostics);
-    verifyNoExternalTensorCaptures(
-      computeBatchOp.getOperation(), computeBatchOp.getBody(), "spat.compute_batch", diagnostics);
+  for (auto batch : funcOp.getOps<spatial::SpatScheduledComputeBatch>()) {
+    verifyScheduledInputs(batch, /*allowChannelReceiveInputs=*/false, "spat.scheduled_compute_batch", diagnostics);
+    verifyNoNestedFragmentAssemblyBlueprints(batch, diagnostics);
   }
-
-  diagnostics.emitSuppressedSummary(funcOp, "Spatial communication invariant verification failed");
+  if (failed(verifyNoComputeBodyCaptures(funcOp)))
+    return failure();
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial verification failed");
   return success(!diagnostics.hasFailure());
 }
 

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp

@@ -6,6 +6,8 @@
 namespace onnx_mlir {
 
 mlir::LogicalResult verifyONNXToSpatial(mlir::func::FuncOp funcOp);
-mlir::LogicalResult verifySpatialCommunicationInvariants(mlir::func::FuncOp funcOp);
+mlir::LogicalResult verifyNoComputeBodyCaptures(mlir::func::FuncOp funcOp);
+mlir::LogicalResult verifyLogicalSpatialGraphInvariants(mlir::func::FuncOp funcOp);
+mlir::LogicalResult verifyScheduledSpatialInvariants(mlir::func::FuncOp funcOp);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -19,9 +19,11 @@ void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   populateSigmoidPatterns(patterns, ctx);
   populateSoftmaxPatterns(patterns, ctx);
   populateConcatPatterns(patterns, ctx);
+  populateFlattenPatterns(patterns, ctx);
   populateGatherPatterns(patterns, ctx);
   populateResizePatterns(patterns, ctx);
   populateReshapePatterns(patterns, ctx);
+  populateSlicePatterns(patterns, ctx);
   populateSplitPatterns(patterns, ctx);
   populateTransposePatterns(patterns, ctx);
 }

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -26,14 +26,16 @@ void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext*
 void populateSigmoidPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSoftmaxPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateConcatPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateFlattenPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateGatherPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateResizePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReshapePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateSlicePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSplitPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateTransposePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
-bool requiresPostRewrite(spatial::SpatCompute computeOp);
-bool requiresPostRewrite(spatial::SpatComputeBatch computeOp);
+bool requiresPostRewrite(spatial::SpatGraphCompute computeOp);
+bool requiresPostRewrite(spatial::SpatGraphComputeBatch computeOp);
 void annotateWeightsConstants(mlir::func::FuncOp funcOp);
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,77 @@
+#include "ConvGeometry.hpp"
+
+#include <algorithm>
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+
+namespace onnx_mlir {
+
+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) {
+  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()),
+    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) {
+  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);
+  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));
+
+  if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
+    chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
+    chunkPositions = std::max<uint64_t>(chunkPositions, static_cast<uint64_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;
+  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)};
+}
+
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state) {
+  ConvRowDemand demand;
+  demand.outputRows = outputRows;
+  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, state);
+  demand.acquiredInputRows = demand.neededInputRows;
+
+  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+  const int64_t rawEnd =
+    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
+  demand.topHaloRows = std::max<int64_t>(0, -rawBegin);
+  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - state.xHeight);
+  demand.acquiredInputRows = demand.neededInputRows;
+  return demand;
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,86 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+
+#include <cstdint>
+
+namespace onnx_mlir {
+
+struct ConvLoweringState {
+  mlir::Value x;
+  mlir::Value w;
+  mlir::Value b;
+  mlir::RankedTensorType xType;
+  mlir::RankedTensorType wType;
+  mlir::RankedTensorType outType;
+  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 padHeightBegin;
+  int64_t padHeightEnd;
+  int64_t padWidthBegin;
+  int64_t padWidthEnd;
+  int64_t strideHeight;
+  int64_t strideWidth;
+  int64_t dilationHeight;
+  int64_t dilationWidth;
+  bool hasBias;
+};
+
+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 pack;
+  uint64_t im2colElements;
+  bool hasBias;
+  bool isDepthwise;
+};
+
+struct RowInterval {
+  int64_t begin = 0;
+  int64_t end = 0;
+};
+
+struct ConvRowDemand {
+  RowInterval outputRows;
+  RowInterval neededInputRows;
+  RowInterval acquiredInputRows;
+  int64_t topHaloRows = 0;
+  int64_t bottomHaloRows = 0;
+};
+
+bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
+
+ConvGeometry buildConvGeometry(const ConvLoweringState& state);
+
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor);
+
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state);
+
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state);
+
+} // namespace onnx_mlir

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

@@ -5,7 +5,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.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/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -47,38 +47,28 @@ static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
     return failure();
 
   const int64_t rankOffset = static_cast<int64_t>(resultShape.size() - sourceShape.size());
-  for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
-    const int64_t sourceIndex = i - rankOffset;
-    const int64_t sourceDim = sourceIndex < 0 ? 1 : sourceShape[sourceIndex];
-    const int64_t resultDim = resultShape[i];
-    if (sourceDim != 1 && sourceDim != resultDim)
-      return failure();
-  }
-
-  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceShape);
   SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultShape);
-
+  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> resultValues;
   resultValues.reserve(resultType.getNumElements());
 
   for (int64_t flatIndex = 0; flatIndex < resultType.getNumElements(); ++flatIndex) {
     int64_t remaining = flatIndex;
     int64_t sourceFlatIndex = 0;
-
     for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
       const int64_t resultIndex = resultStrides.empty() ? 0 : remaining / resultStrides[i];
       remaining = resultStrides.empty() ? 0 : remaining % resultStrides[i];
-
       const int64_t sourceIndex = i - rankOffset;
       if (sourceIndex < 0)
         continue;
-
       const int64_t sourceDim = sourceShape[sourceIndex];
+      const int64_t resultDim = resultShape[i];
+      if (sourceDim != 1 && sourceDim != resultDim)
+        return failure();
       const int64_t mappedIndex = sourceDim == 1 ? 0 : resultIndex;
       sourceFlatIndex += mappedIndex * sourceStrides[sourceIndex];
     }
-
     resultValues.push_back(sourceValues[sourceFlatIndex]);
   }
 
@@ -106,7 +96,7 @@ static FailureOr<Value> materializeReciprocalTensor(Value value,
   if (failed(broadcastedValue))
     return failure();
 
-  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getDenseConstantAttr(*broadcastedValue));
+  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getHostConstDenseElementsAttr(*broadcastedValue));
   if (!denseAttr)
     return failure();
 
@@ -185,10 +175,46 @@ struct DivToSpatialCompute : OpConversionPattern<ONNXDivOp> {
   }
 };
 
+struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ONNXAddOp op, ONNXAddOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    if (!resultType || !resultType.hasStaticShape())
+      return failure();
+
+    FailureOr<BiasAddPlanCandidate> candidate =
+      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.replaceOp(op, plan.getResult());
+      return success();
+    }
+
+    auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc());
+    if (failed(lhs))
+      return failure();
+    auto rhs = prepareElementwiseOperand(adaptor.getB(), resultType, rewriter, op.getLoc());
+    if (failed(rhs))
+      return failure();
+
+    auto computeOp =
+      createSpatCompute<2>(rewriter, op.getLoc(), resultType, {}, ValueRange {*lhs, *rhs}, [&](Value x, Value y) {
+        auto loweredOp = spatial::SpatVAddOp::create(rewriter, op.getLoc(), resultType, x, y);
+        spatial::SpatYieldOp::create(rewriter, op.getLoc(), loweredOp.getResult());
+      });
+    rewriter.replaceOp(op, computeOp);
+    return success();
+  }
+};
+
 } // namespace
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<BinaryElementwiseToSpatialCompute<ONNXAddOp, spatial::SpatVAddOp>>(ctx);
+  patterns.add<AddToSpatialCompute>(ctx);
+  patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);
   patterns.add<DivToSpatialCompute>(ctx);
 }

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

@@ -87,28 +87,6 @@ static Value createGemmBatchHOffset(Value lane,
                                  rewriter.getInsertionBlock()->getParentOp());
 }
 
-static Value
-createZeroPaddedTensor(Value value, RankedTensorType resultType, ConversionPatternRewriter& rewriter, Location loc) {
-  auto sourceType = cast<RankedTensorType>(value.getType());
-  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
-  SmallVector<OpFoldResult> highPads;
-  highPads.reserve(sourceType.getRank());
-  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
-    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
-
-  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
-  auto* padBlock = new Block();
-  for (int64_t i = 0; i < sourceType.getRank(); ++i)
-    padBlock->addArgument(rewriter.getIndexType(), loc);
-  padOp.getRegion().push_back(padBlock);
-  rewriter.setInsertionPointToStart(padBlock);
-  auto zero = getOrCreateConstant(
-    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
-  tensor::YieldOp::create(rewriter, loc, zero);
-  rewriter.setInsertionPointAfter(padOp);
-  return padOp.getResult();
-}
-
 static FailureOr<Value> materializePaddedConstantMatrix(Value value,
                                                         RankedTensorType resultType,
                                                         ConversionPatternRewriter& rewriter,
@@ -232,22 +210,6 @@ static Value extractATile(
   return tensor::ExtractSliceOp::create(rewriter, loc, aTileType, a, offsets, sizes, strides).getResult();
 }
 
-static Value createPaddedInputCompute(Value input,
-                                      RankedTensorType paddedInputType,
-                                      ConversionPatternRewriter& rewriter,
-                                      Location loc) {
-  auto inputType = cast<RankedTensorType>(input.getType());
-  if (inputType == paddedInputType)
-    return input;
-
-  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
-    Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
-    spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
-  });
-
-  return computeOp.getResult(0);
-}
-
 static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
                                                            Value b,
                                                            RankedTensorType aType,
@@ -285,15 +247,11 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
       SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(crossbarSize.getValue()),
                                         rewriter.getIndexAttr(crossbarSize.getValue())};
       SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
-      Value bTile =
-        tensor::ExtractSliceOp::create(rewriter, loc, bTileType, args.weights.front(), bOffsets, bSizes, unitStrides)
-          .getResult();
+      Value bTile = extractStaticSliceOrIdentity(
+        rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides);
       Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
-      SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, unitStrides);
+      publishGraphBatchPhysicalFragment(rewriter, loc, piece, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -440,11 +398,7 @@ static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
       Value bVector = extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc);
       Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
 
-      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-      SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, unitStrides);
+      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -486,15 +440,14 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
         Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, nestedLoc);
         Value column =
           onnx_mlir::affineModConst(rewriter, nestedLoc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-        SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
         SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
         SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        Value scalar = tensor::ExtractSliceOp::create(
-                         rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, unitStrides)
-                         .getResult();
+        FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
+        if (failed(scalar))
+          return failure();
         if (alpha != 1.0f) {
           Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc);
-          scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, scalar, alphaTensor).getResult();
+          *scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, *scalar, alphaTensor).getResult();
         }
         if (biasArg) {
           Value biasScalar =
@@ -504,11 +457,11 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
             biasScalar =
               spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).getResult();
           }
-          scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, scalar, biasScalar).getResult();
+          *scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, *scalar, biasScalar).getResult();
         }
         SmallVector<OpFoldResult> outputOffsets {row, column};
         Value outputNext =
-          tensor::InsertSliceOp::create(rewriter, nestedLoc, scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
+          tensor::InsertSliceOp::create(rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
             .getResult();
         yielded.push_back(outputNext);
         return success();
@@ -544,14 +497,13 @@ static Value extractReductionPiece(Value partialPiecesArg,
                                    int64_t numOutRows,
                                    ConversionPatternRewriter& rewriter,
                                    Location loc) {
-  SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
-                                        rewriter.getIndexAttr(crossbarSize.getValue())};
+  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> pieceOffsets {
-    createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0)};
-  return tensor::ExtractSliceOp::create(
-           rewriter, loc, pieceType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides)
-    .getResult();
+    createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
+  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides);
+  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
 }
 
 static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
@@ -690,11 +642,6 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
   Value b = gemmOpAdaptor.getB();
   Value c = gemmOpAdaptor.getC();
 
-  if (gemmOpAdaptor.getTransA()) {
-    gemmOp.emitOpError("requires transA=false before tiled Spatial Gemm lowering");
-    return failure();
-  }
-
   auto aType = dyn_cast<RankedTensorType>(a.getType());
   auto bType = dyn_cast<RankedTensorType>(b.getType());
   auto outType = dyn_cast<RankedTensorType>(gemmOp.getY().getType());
@@ -725,9 +672,12 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     return failure();
   }
 
-  const int64_t numOutRows = outType.getDimSize(0);
-  const int64_t numOutCols = outType.getDimSize(1);
-  const int64_t reductionSize = aType.getDimSize(1);
+  if (gemmOpAdaptor.getTransA()) {
+    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();
+    aType = transposedType;
+  }
 
   if (gemmOpAdaptor.getTransB()) {
     auto bShape = bType.getShape();
@@ -736,6 +686,10 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     bType = transposedType;
   }
 
+  const int64_t numOutRows = outType.getDimSize(0);
+  const int64_t numOutCols = outType.getDimSize(1);
+  const int64_t reductionSize = aType.getDimSize(1);
+
   if (!isCompileTimeComputable(b)) {
     bool hasC = hasGemmBias(c);
     float alpha = gemmOpAdaptor.getAlpha().convertToFloat();
@@ -767,7 +721,7 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
       return failure();
     }
 
-    auto scalarPiecesType = RankedTensorType::get({laneCount64, 1}, outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount64, RankedTensorType::get({1, 1}, outType.getElementType()));
     auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, rewriter, loc);
     if (failed(batchOp))
       return failure();
@@ -839,8 +793,8 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     return failure();
   }
 
-  auto partialPiecesType =
-    RankedTensorType::get({laneCount64, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType());
+  auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
+    laneCount64, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType()));
   auto batchOp =
     createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
   if (failed(batchOp))

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

@@ -22,13 +22,87 @@ namespace {
 
 static FailureOr<SmallVector<int64_t>> inferSupportedBatchShape(ArrayRef<int64_t> lhsBatchShape,
                                                                 ArrayRef<int64_t> rhsBatchShape) {
-  if (lhsBatchShape.empty())
-    return SmallVector<int64_t>(rhsBatchShape.begin(), rhsBatchShape.end());
-  if (rhsBatchShape.empty())
-    return SmallVector<int64_t>(lhsBatchShape.begin(), lhsBatchShape.end());
-  if (!llvm::equal(lhsBatchShape, rhsBatchShape))
-    return failure();
-  return SmallVector<int64_t>(lhsBatchShape.begin(), lhsBatchShape.end());
+  const int64_t resultRank = std::max<int64_t>(lhsBatchShape.size(), rhsBatchShape.size());
+  SmallVector<int64_t> resultShape(resultRank, 1);
+  for (int64_t resultIndex = resultRank - 1, lhsIndex = lhsBatchShape.size() - 1, rhsIndex = rhsBatchShape.size() - 1;
+       resultIndex >= 0;
+       --resultIndex, --lhsIndex, --rhsIndex) {
+    const int64_t lhsDim = lhsIndex >= 0 ? lhsBatchShape[lhsIndex] : 1;
+    const int64_t rhsDim = rhsIndex >= 0 ? rhsBatchShape[rhsIndex] : 1;
+    if (lhsDim != rhsDim && lhsDim != 1 && rhsDim != 1)
+      return failure();
+    resultShape[resultIndex] = std::max(lhsDim, rhsDim);
+  }
+  return resultShape;
+}
+
+static int64_t mapStaticBroadcastedBatchIndex(int64_t outputBatchIndex,
+                                              ArrayRef<int64_t> sourceBatchShape,
+                                              ArrayRef<int64_t> outputBatchShape) {
+  if (sourceBatchShape.empty() || getStaticShapeElementCount(sourceBatchShape) == 1)
+    return 0;
+  if (llvm::equal(sourceBatchShape, outputBatchShape))
+    return outputBatchIndex;
+
+  SmallVector<int64_t> outputStrides = computeRowMajorStrides(outputBatchShape);
+  SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceBatchShape);
+  int64_t sourceFlatIndex = 0;
+  for (int64_t sourceDimIndex = 0; sourceDimIndex < static_cast<int64_t>(sourceBatchShape.size()); ++sourceDimIndex) {
+    if (sourceBatchShape[sourceDimIndex] == 1)
+      continue;
+    const int64_t outputDimIndex = outputBatchShape.size() - sourceBatchShape.size() + sourceDimIndex;
+    const int64_t outputDimStride = outputStrides.empty() ? 1 : outputStrides[outputDimIndex];
+    const int64_t outputDimIndexValue = outputDimStride == 1
+                                        ? outputBatchIndex % outputBatchShape[outputDimIndex]
+                                        : (outputBatchIndex / outputDimStride) % outputBatchShape[outputDimIndex];
+    sourceFlatIndex += outputDimIndexValue * sourceStrides[sourceDimIndex];
+  }
+  return sourceFlatIndex;
+}
+
+static Value computeFlatBatchIndexCoordinate(
+  Value flatBatchIndex, ArrayRef<int64_t> batchShape, int64_t dimIndex, PatternRewriter& rewriter, Location loc) {
+  if (batchShape[dimIndex] == 1)
+    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+
+  const int64_t dimStride = dimIndex + 1 == static_cast<int64_t>(batchShape.size())
+                            ? 1
+                            : getStaticShapeElementCount(batchShape.drop_front(dimIndex + 1));
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value dimCoordinate = flatBatchIndex;
+  if (dimStride != 1)
+    dimCoordinate = affineFloorDivConst(rewriter, loc, dimCoordinate, dimStride, anchorOp);
+  return affineModConst(rewriter, loc, dimCoordinate, batchShape[dimIndex], anchorOp);
+}
+
+static Value mapOutputBatchIndexToSourceBatchIndex(Value outputBatchIndex,
+                                                   ArrayRef<int64_t> sourceBatchShape,
+                                                   ArrayRef<int64_t> outputBatchShape,
+                                                   PatternRewriter& rewriter,
+                                                   Location loc) {
+  if (sourceBatchShape.empty() || getStaticShapeElementCount(sourceBatchShape) == 1)
+    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+  if (llvm::equal(sourceBatchShape, outputBatchShape))
+    return outputBatchIndex;
+
+  Value sourceBatchIndex = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+  SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceBatchShape);
+  for (int64_t sourceDimIndex = 0; sourceDimIndex < static_cast<int64_t>(sourceBatchShape.size()); ++sourceDimIndex) {
+    if (sourceBatchShape[sourceDimIndex] == 1)
+      continue;
+    const int64_t outputDimIndex = outputBatchShape.size() - sourceBatchShape.size() + sourceDimIndex;
+    Value outputCoordinate =
+      computeFlatBatchIndexCoordinate(outputBatchIndex, outputBatchShape, outputDimIndex, rewriter, loc);
+    Value contribution = sourceStrides[sourceDimIndex] == 1
+                         ? outputCoordinate
+                         : affineMulConst(rewriter,
+                                          loc,
+                                          outputCoordinate,
+                                          sourceStrides[sourceDimIndex],
+                                          rewriter.getInsertionBlock()->getParentOp());
+    sourceBatchIndex = arith::AddIOp::create(rewriter, loc, sourceBatchIndex, contribution);
+  }
+  return sourceBatchIndex;
 }
 
 static Value
@@ -67,6 +141,52 @@ expandBatchDims(Value value, RankedTensorType outputType, size_t batchRank, Patt
   return materializeOrComputeUnary(value, outputType, rewriter, loc, buildExpanded);
 }
 
+static Value createMatrixFromVector(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
+  auto buildExpanded = [&](Value input) -> Value {
+    return tensor::ExpandShapeOp::create(rewriter,
+                                         loc,
+                                         resultType,
+                                         input,
+                                         SmallVector<ReassociationIndices> {
+                                           {0, 1}
+    });
+  };
+  return materializeOrComputeUnary(value, resultType, rewriter, loc, buildExpanded);
+}
+
+static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> removedAxes) {
+  SmallVector<ReassociationIndices> reassociation;
+  ReassociationIndices currentGroup;
+  for (auto [axis, removeAxis] : llvm::enumerate(removedAxes)) {
+    currentGroup.push_back(axis);
+    if (!removeAxis) {
+      reassociation.push_back(currentGroup);
+      currentGroup.clear();
+    }
+  }
+
+  if (!currentGroup.empty()) {
+    if (reassociation.empty())
+      reassociation.push_back(std::move(currentGroup));
+    else
+      reassociation.back().append(currentGroup.begin(), currentGroup.end());
+  }
+  return reassociation;
+}
+
+static Value squeezeUnitDims(
+  Value value, RankedTensorType resultType, ArrayRef<bool> removedAxes, PatternRewriter& rewriter, Location loc) {
+  if (cast<RankedTensorType>(value.getType()) == resultType)
+    return value;
+
+  SmallVector<ReassociationIndices> reassociation =
+    resultType.getRank() == 0 ? SmallVector<ReassociationIndices> {} : buildCollapseReassociation(removedAxes);
+  auto buildCollapsed = [&](Value input) -> Value {
+    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation).getResult();
+  };
+  return materializeOrComputeUnary(value, resultType, rewriter, loc, buildCollapsed);
+}
+
 static Value ensureBatchedTensor(
   Value value, int64_t batchSize, int64_t rows, int64_t cols, PatternRewriter& rewriter, Location loc) {
   auto type = cast<RankedTensorType>(value.getType());
@@ -135,44 +255,11 @@ static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Locati
   return createONNXTranspose(resultType, {0, 2, 1});
 }
 
-static Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
-  auto sourceType = cast<RankedTensorType>(value.getType());
-  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
-  SmallVector<OpFoldResult> highPads;
-  highPads.reserve(sourceType.getRank());
-  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
-    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
-
-  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
-  auto* padBlock = new Block();
-  for (int64_t i = 0; i < sourceType.getRank(); ++i)
-    padBlock->addArgument(rewriter.getIndexType(), loc);
-  padOp.getRegion().push_back(padBlock);
-  rewriter.setInsertionPointToStart(padBlock);
-  auto zero = getOrCreateConstant(
-    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
-  tensor::YieldOp::create(rewriter, loc, zero);
-  rewriter.setInsertionPointAfter(padOp);
-  return padOp.getResult();
-}
-
-static Value createPaddedBatchedInputCompute(Value input,
-                                             RankedTensorType paddedInputType,
-                                             PatternRewriter& rewriter,
-                                             Location loc) {
-  auto inputType = cast<RankedTensorType>(input.getType());
-  if (inputType == paddedInputType)
-    return input;
-
-  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
-    Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
-    spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
-  });
-  return computeOp.getResult(0);
-}
-
-static FailureOr<Value> materializePaddedBatchedWeight(
-  Value value, int64_t sourceBatch, int64_t targetBatch, RankedTensorType resultType, PatternRewriter& rewriter) {
+static FailureOr<Value> materializePaddedBatchedWeight(Value value,
+                                                       ArrayRef<int64_t> sourceBatchShape,
+                                                       ArrayRef<int64_t> targetBatchShape,
+                                                       RankedTensorType resultType,
+                                                       PatternRewriter& rewriter) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   if (sourceType == resultType)
     return value;
@@ -183,13 +270,15 @@ static FailureOr<Value> materializePaddedBatchedWeight(
 
   const int64_t sourceRows = sourceType.getRank() == 2 ? sourceType.getDimSize(0) : sourceType.getDimSize(1);
   const int64_t sourceCols = sourceType.getRank() == 2 ? sourceType.getDimSize(1) : sourceType.getDimSize(2);
+  const int64_t targetBatch = targetBatchShape.empty() ? 1 : getStaticShapeElementCount(targetBatchShape);
   const int64_t targetRows = resultType.getDimSize(1);
   const int64_t targetCols = resultType.getDimSize(2);
   SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> resultValues(resultType.getNumElements(), rewriter.getZeroAttr(resultType.getElementType()));
 
   for (int64_t batchIdx = 0; batchIdx < targetBatch; ++batchIdx) {
-    const int64_t sourceBatchIdx = sourceType.getRank() == 2 ? 0 : (sourceBatch == 1 ? 0 : batchIdx);
+    const int64_t sourceBatchIdx =
+      sourceType.getRank() == 2 ? 0 : mapStaticBroadcastedBatchIndex(batchIdx, sourceBatchShape, targetBatchShape);
     const int64_t sourceBatchBase = sourceType.getRank() == 2 ? 0 : sourceBatchIdx * sourceRows * sourceCols;
     const int64_t targetBatchBase = batchIdx * targetRows * targetCols;
     for (int64_t row = 0; row < sourceRows; ++row)
@@ -202,16 +291,18 @@ static FailureOr<Value> materializePaddedBatchedWeight(
 }
 
 static Value extractBatchedATile(Value a,
-                                 int64_t sourceBatchCount,
-                                 Value batch,
+                                 ArrayRef<int64_t> sourceBatchShape,
+                                 ArrayRef<int64_t> outputBatchShape,
+                                 Value outputBatchIndex,
                                  Value row,
                                  Value kOffset,
                                  RankedTensorType aTileType,
                                  PatternRewriter& rewriter,
                                  Location loc) {
   auto aSliceType = RankedTensorType::get({1, 1, aTileType.getDimSize(1)}, aTileType.getElementType());
-  SmallVector<OpFoldResult> offsets {
-    sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0)) : OpFoldResult(batch), row, kOffset};
+  Value sourceBatchIndex =
+    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
+  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, kOffset};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(aTileType.getDimSize(1))};
   auto slice =
@@ -227,8 +318,9 @@ static Value extractBatchedATile(Value a,
 }
 
 static Value extractBatchedBTile(Value b,
-                                 int64_t sourceBatchCount,
-                                 Value batch,
+                                 ArrayRef<int64_t> sourceBatchShape,
+                                 ArrayRef<int64_t> outputBatchShape,
+                                 Value outputBatchIndex,
                                  Value kOffset,
                                  Value hOffset,
                                  RankedTensorType bTileType,
@@ -236,8 +328,9 @@ static Value extractBatchedBTile(Value b,
                                  Location loc) {
   auto bSliceType =
     RankedTensorType::get({1, bTileType.getDimSize(0), bTileType.getDimSize(1)}, bTileType.getElementType());
-  SmallVector<OpFoldResult> offsets {
-    sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0)) : OpFoldResult(batch), kOffset, hOffset};
+  Value sourceBatchIndex =
+    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
+  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), kOffset, hOffset};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(bTileType.getDimSize(0)),
                                    rewriter.getIndexAttr(bTileType.getDimSize(1))};
@@ -262,9 +355,10 @@ static Value getBatchLaneIndex(
 static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
                                                                   Value b,
                                                                   RankedTensorType aType,
-                                                                  int64_t aBatchCount,
+                                                                  ArrayRef<int64_t> aBatchShape,
                                                                   RankedTensorType bType,
-                                                                  int64_t bBatchCount,
+                                                                  ArrayRef<int64_t> bBatchShape,
+                                                                  ArrayRef<int64_t> outputBatchShape,
                                                                   RankedTensorType partialPiecesType,
                                                                   int64_t numOutRows,
                                                                   int64_t numKSlices,
@@ -298,16 +392,13 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
       auto pieceType =
         RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());
 
-      Value aTile =
-        extractBatchedATile(args.inputs.front(), aBatchCount, batch, row, kOffset, aTileType, rewriter, loc);
-      Value bTile =
-        extractBatchedBTile(args.weights.front(), bBatchCount, batch, kOffset, hOffset, bTileType, rewriter, loc);
+      Value aTile = extractBatchedATile(
+        args.inputs.front(), aBatchShape, outputBatchShape, batch, row, kOffset, aTileType, rewriter, loc);
+      Value bTile = extractBatchedBTile(
+        args.weights.front(), bBatchShape, outputBatchShape, batch, kOffset, hOffset, bTileType, rewriter, loc);
       Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
-      SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2));
+      publishGraphBatchPhysicalFragment(rewriter, loc, piece, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -315,17 +406,17 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
 }
 
 static Value extractDynamicBatchedBColumn(Value matrix,
-                                          int64_t sourceBatchCount,
-                                          Value batch,
+                                          ArrayRef<int64_t> sourceBatchShape,
+                                          ArrayRef<int64_t> outputBatchShape,
+                                          Value outputBatchIndex,
                                           Value column,
                                           RankedTensorType vectorType,
                                           PatternRewriter& rewriter,
                                           Location loc) {
   auto columnSliceType = RankedTensorType::get({1, vectorType.getDimSize(1), 1}, vectorType.getElementType());
-  SmallVector<OpFoldResult> offsets {sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0))
-                                                           : OpFoldResult(batch),
-                                     rewriter.getIndexAttr(0),
-                                     column};
+  Value sourceBatchIndex =
+    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
+  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), rewriter.getIndexAttr(0), column};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -350,17 +441,17 @@ static Value extractDynamicBatchedBColumn(Value matrix,
 }
 
 static Value extractDynamicBatchedRowVector(Value matrix,
-                                            int64_t sourceBatchCount,
-                                            Value batch,
+                                            ArrayRef<int64_t> sourceBatchShape,
+                                            ArrayRef<int64_t> outputBatchShape,
+                                            Value outputBatchIndex,
                                             Value row,
                                             RankedTensorType vectorType,
                                             PatternRewriter& rewriter,
                                             Location loc) {
   auto rowSliceType = RankedTensorType::get({1, 1, vectorType.getDimSize(1)}, vectorType.getElementType());
-  SmallVector<OpFoldResult> offsets {sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0))
-                                                           : OpFoldResult(batch),
-                                     row,
-                                     rewriter.getIndexAttr(0)};
+  Value sourceBatchIndex =
+    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
+  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
   auto rowSlice =
@@ -376,9 +467,10 @@ static Value extractDynamicBatchedRowVector(Value matrix,
 }
 
 static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
-                                                                     int64_t aBatchCount,
+                                                                     ArrayRef<int64_t> aBatchShape,
                                                                      Value b,
-                                                                     int64_t bBatchCount,
+                                                                     ArrayRef<int64_t> bBatchShape,
+                                                                     ArrayRef<int64_t> outputBatchShape,
                                                                      RankedTensorType aType,
                                                                      RankedTensorType bType,
                                                                      RankedTensorType scalarPiecesType,
@@ -406,15 +498,12 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
 
       auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
       auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-      Value aVector =
-        extractDynamicBatchedRowVector(args.inputs[0], aBatchCount, batch, row, vectorType, rewriter, loc);
-      Value bVector =
-        extractDynamicBatchedBColumn(args.inputs[1], bBatchCount, batch, column, vectorType, rewriter, loc);
+      Value aVector = extractDynamicBatchedRowVector(
+        args.inputs[0], aBatchShape, outputBatchShape, batch, row, vectorType, rewriter, loc);
+      Value bVector = extractDynamicBatchedBColumn(
+        args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc);
       Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, getUnitStrides(rewriter, 2));
+      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -453,14 +542,13 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
           Value batchLane = affineModConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
           Value row = affineFloorDivConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
           Value column = affineModConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
-          SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-          Value scalar = tensor::ExtractSliceOp::create(
-            rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, getUnitStrides(rewriter, 2));
+          FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
+          if (failed(scalar))
+            return failure();
           Value expanded = tensor::ExpandShapeOp::create(rewriter,
                                                          nestedLoc,
                                                          outputScalarType,
-                                                         scalar,
+                                                         *scalar,
                                                          SmallVector<ReassociationIndices> {
                                                            {0},
                                                            {1, 2}
@@ -501,10 +589,11 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg,
   Value kOffset = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), kSlice * numOutRows);
   Value batchAndHSlice = arith::AddIOp::create(rewriter, loc, batchOffset, hOffset);
   Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset);
-  SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(crossbarSize.getValue())};
-  return tensor::ExtractSliceOp::create(
-    rewriter, loc, pieceType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 2));
+  SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
+  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 3));
+  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
 }
 
 static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
@@ -629,11 +718,17 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
   return computeOp->getResult(0);
 }
 
-struct MatMulShapeInfo {
+struct NormalizedMatMulInfo {
   RankedTensorType lhsType;
   RankedTensorType rhsType;
   RankedTensorType outType;
-  SmallVector<int64_t> batchShape;
+  RankedTensorType normalizedLhsType;
+  RankedTensorType normalizedRhsType;
+  SmallVector<int64_t> lhsBatchShape;
+  SmallVector<int64_t> rhsBatchShape;
+  SmallVector<int64_t> outputBatchShape;
+  bool lhsWasVector;
+  bool rhsWasVector;
   int64_t lhsBatch;
   int64_t rhsBatch;
   int64_t batch;
@@ -642,46 +737,170 @@ struct MatMulShapeInfo {
   int64_t n;
 };
 
-static FailureOr<MatMulShapeInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp) {
+struct MatMulLoweringPlan {
+  Value lhs;
+  Value rhs;
+  RankedTensorType lhsType;
+  RankedTensorType rhsType;
+  SmallVector<int64_t> lhsBatchShape;
+  SmallVector<int64_t> rhsBatchShape;
+  SmallVector<int64_t> outputBatchShape;
+  int64_t lhsBatch;
+  int64_t rhsBatch;
+  int64_t batch;
+  int64_t m;
+  int64_t k;
+  int64_t n;
+  bool transposedResult;
+};
+
+static SmallVector<int64_t> computeExpectedMatMulOutputShape(
+  ArrayRef<int64_t> batchShape, int64_t m, int64_t n, bool lhsWasVector, bool rhsWasVector) {
+  SmallVector<int64_t> shape(batchShape.begin(), batchShape.end());
+  if (lhsWasVector && rhsWasVector)
+    return shape;
+  if (lhsWasVector) {
+    shape.push_back(n);
+    return shape;
+  }
+  if (rhsWasVector) {
+    shape.push_back(m);
+    return shape;
+  }
+  shape.push_back(m);
+  shape.push_back(n);
+  return shape;
+}
+
+static FailureOr<NormalizedMatMulInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp) {
   auto lhsType = dyn_cast<RankedTensorType>(matmulOp.getA().getType());
   auto rhsType = dyn_cast<RankedTensorType>(matmulOp.getB().getType());
   auto outType = dyn_cast<RankedTensorType>(matmulOp.getY().getType());
   if (!lhsType || !rhsType || !outType || !lhsType.hasStaticShape() || !rhsType.hasStaticShape()
       || !outType.hasStaticShape())
     return failure();
-  if (lhsType.getRank() < 2 || rhsType.getRank() < 2 || outType.getRank() < 2)
+  if (lhsType.getRank() < 1 || rhsType.getRank() < 1)
     return failure();
   if (!hasStaticPositiveShape(lhsType) || !hasStaticPositiveShape(rhsType) || !hasStaticPositiveShape(outType))
     return failure();
 
-  SmallVector<int64_t> lhsBatchShape(lhsType.getShape().begin(), lhsType.getShape().end() - 2);
-  SmallVector<int64_t> rhsBatchShape(rhsType.getShape().begin(), rhsType.getShape().end() - 2);
-  auto batchShape = inferSupportedBatchShape(lhsBatchShape, rhsBatchShape);
-  if (failed(batchShape))
+  const bool lhsWasVector = lhsType.getRank() == 1;
+  const bool rhsWasVector = rhsType.getRank() == 1;
+  auto normalizedLhsType =
+    lhsWasVector ? RankedTensorType::get({1, lhsType.getDimSize(0)}, lhsType.getElementType(), lhsType.getEncoding())
+                 : lhsType;
+  auto normalizedRhsType =
+    rhsWasVector ? RankedTensorType::get({rhsType.getDimSize(0), 1}, rhsType.getElementType(), rhsType.getEncoding())
+                 : rhsType;
+
+  SmallVector<int64_t> lhsBatchShape(normalizedLhsType.getShape().begin(), normalizedLhsType.getShape().end() - 2);
+  SmallVector<int64_t> rhsBatchShape(normalizedRhsType.getShape().begin(), normalizedRhsType.getShape().end() - 2);
+  auto outputBatchShape = inferSupportedBatchShape(lhsBatchShape, rhsBatchShape);
+  if (failed(outputBatchShape))
     return failure();
 
   const int64_t lhsBatch = lhsBatchShape.empty() ? 1 : getStaticShapeElementCount(lhsBatchShape);
   const int64_t rhsBatch = rhsBatchShape.empty() ? 1 : getStaticShapeElementCount(rhsBatchShape);
-  const int64_t batch = batchShape->empty() ? 1 : getStaticShapeElementCount(*batchShape);
-  const int64_t m = lhsType.getDimSize(lhsType.getRank() - 2);
-  const int64_t k = lhsType.getDimSize(lhsType.getRank() - 1);
-  const int64_t rhsK = rhsType.getDimSize(rhsType.getRank() - 2);
-  const int64_t n = rhsType.getDimSize(rhsType.getRank() - 1);
+  const int64_t batch = outputBatchShape->empty() ? 1 : getStaticShapeElementCount(*outputBatchShape);
+  const int64_t m = normalizedLhsType.getDimSize(normalizedLhsType.getRank() - 2);
+  const int64_t k = normalizedLhsType.getDimSize(normalizedLhsType.getRank() - 1);
+  const int64_t rhsK = normalizedRhsType.getDimSize(normalizedRhsType.getRank() - 2);
+  const int64_t n = normalizedRhsType.getDimSize(normalizedRhsType.getRank() - 1);
   if (k != rhsK)
     return failure();
 
-  if (outType.getRank() == 2) {
-    if (batch != 1 || outType.getDimSize(0) != m || outType.getDimSize(1) != n)
-      return failure();
-  }
-  else {
-    SmallVector<int64_t> outBatchShape(outType.getShape().begin(), outType.getShape().end() - 2);
-    if (!llvm::equal(outBatchShape, *batchShape) || outType.getDimSize(outType.getRank() - 2) != m
-        || outType.getDimSize(outType.getRank() - 1) != n)
-      return failure();
+  if (SmallVector<int64_t>(outType.getShape().begin(), outType.getShape().end())
+      != computeExpectedMatMulOutputShape(*outputBatchShape, m, n, lhsWasVector, rhsWasVector)) {
+    return failure();
   }
 
-  return MatMulShapeInfo {lhsType, rhsType, outType, *batchShape, lhsBatch, rhsBatch, batch, m, k, n};
+  return NormalizedMatMulInfo {lhsType,
+                               rhsType,
+                               outType,
+                               normalizedLhsType,
+                               normalizedRhsType,
+                               lhsBatchShape,
+                               rhsBatchShape,
+                               *outputBatchShape,
+                               lhsWasVector,
+                               rhsWasVector,
+                               lhsBatch,
+                               rhsBatch,
+                               batch,
+                               m,
+                               k,
+                               n};
+}
+
+static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs,
+                                            Value normalizedRhs,
+                                            const NormalizedMatMulInfo& info,
+                                            bool useTransposedForm,
+                                            PatternRewriter& rewriter,
+                                            Location loc) {
+  MatMulLoweringPlan plan {normalizedLhs,
+                           normalizedRhs,
+                           cast<RankedTensorType>(normalizedLhs.getType()),
+                           cast<RankedTensorType>(normalizedRhs.getType()),
+                           info.lhsBatchShape,
+                           info.rhsBatchShape,
+                           info.outputBatchShape,
+                           info.lhsBatch,
+                           info.rhsBatch,
+                           info.batch,
+                           info.m,
+                           info.k,
+                           info.n,
+                           false};
+  if (!useTransposedForm)
+    return plan;
+
+  plan.lhs = transposeLastTwoDims(normalizedRhs, rewriter, loc);
+  plan.rhs = transposeLastTwoDims(normalizedLhs, rewriter, loc);
+  plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
+  plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
+  std::swap(plan.lhsBatchShape, plan.rhsBatchShape);
+  std::swap(plan.lhsBatch, plan.rhsBatch);
+  plan.m = info.n;
+  plan.n = info.m;
+  plan.transposedResult = true;
+  return plan;
+}
+
+static Value normalizeMatMulOperand(
+  Value value, RankedTensorType normalizedType, bool wasVector, PatternRewriter& rewriter, Location loc) {
+  if (!wasVector)
+    return value;
+  return createMatrixFromVector(value, normalizedType, rewriter, loc);
+}
+
+static Value finalizeNormalizedMatMulResult(Value value,
+                                            RankedTensorType directOutType,
+                                            const NormalizedMatMulInfo& info,
+                                            PatternRewriter& rewriter,
+                                            Location loc) {
+  // The direct lowered result is always [flatBatch, normalizedM, normalizedN].
+  // Restore ONNX MatMul result rank by expanding right-aligned batch dimensions
+  // and removing the synthetic unit matrix axes introduced for vector operands.
+  Value result = value;
+  RankedTensorType currentType = directOutType;
+  if (info.outputBatchShape.size() > 1) {
+    SmallVector<int64_t> expandedShape(info.outputBatchShape.begin(), info.outputBatchShape.end());
+    expandedShape.push_back(info.m);
+    expandedShape.push_back(info.n);
+    auto expandedType = RankedTensorType::get(expandedShape, info.outType.getElementType(), info.outType.getEncoding());
+    result = expandBatchDims(result, expandedType, info.outputBatchShape.size(), rewriter, loc);
+    currentType = expandedType;
+  }
+
+  SmallVector<bool> removedAxes(currentType.getRank(), false);
+  if (info.outputBatchShape.empty())
+    removedAxes[0] = true;
+  if (info.lhsWasVector)
+    removedAxes[currentType.getRank() - 2] = true;
+  if (info.rhsWasVector)
+    removedAxes[currentType.getRank() - 1] = true;
+  return squeezeUnitDims(result, info.outType, removedAxes, rewriter, loc);
 }
 
 struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
@@ -689,7 +908,10 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
-    if (failed(shapeInfo) || shapeInfo->outType.getRank() != 2)
+    if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector)
+      return failure();
+
+    if (!shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
@@ -730,7 +952,17 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
       gemmResult =
         ONNXTransposeOp::create(rewriter, loc, shapeInfo->outType, gemmResult, rewriter.getI64ArrayAttr({1, 0}))
           .getResult();
-    rewriter.replaceOp(matmulOp, gemmResult);
+
+    if (shapeInfo->outputBatchShape.empty()) {
+      rewriter.replaceOp(matmulOp, gemmResult);
+      return success();
+    }
+
+    auto directOutType =
+      RankedTensorType::get({1, shapeInfo->m, shapeInfo->n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+    Value batchedResult = ensureBatchedTensor(gemmResult, /*batchSize=*/1, shapeInfo->m, shapeInfo->n, rewriter, loc);
+    Value finalResult = finalizeNormalizedMatMulResult(batchedResult, directOutType, *shapeInfo, rewriter, loc);
+    rewriter.replaceOp(matmulOp, finalResult);
     return success();
   }
 };
@@ -742,61 +974,56 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
     if (failed(shapeInfo))
       return failure();
-    if (shapeInfo->outType.getRank() == 2)
+    if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
-    bool useTransposedForm = isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
+    bool useTransposedForm = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
+                          && isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
 
-    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
-    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
-    int64_t lhsBatchForGemm = shapeInfo->lhsBatch;
-    int64_t rhsBatchForGemm = shapeInfo->rhsBatch;
-    int64_t gemmM = shapeInfo->m;
-    int64_t gemmK = shapeInfo->k;
-    int64_t gemmN = shapeInfo->n;
-    if (useTransposedForm) {
-      lhs = transposeLastTwoDims(matmulOp.getB(), rewriter, loc);
-      lhsBatchForGemm = shapeInfo->rhsBatch;
-      rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc);
-      rhsBatchForGemm = shapeInfo->lhsBatch;
-      gemmM = shapeInfo->n;
-      gemmN = shapeInfo->m;
-    }
+    Value lhs =
+      normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType, shapeInfo->lhsWasVector, rewriter, loc);
+    Value rhs =
+      normalizeMatMulOperand(matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc);
+    lhs = collapseBatchDims(lhs, shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
+    rhs = collapseBatchDims(rhs, shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
+    MatMulLoweringPlan plan = buildLoweringPlan(lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc);
 
-    lhs = ensureBatchedTensor(lhs, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
-    rhs = ensureBatchedTensor(rhs, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
-    auto lhsBatchedType = cast<RankedTensorType>(lhs.getType());
-    auto rhsBatchedType = cast<RankedTensorType>(rhs.getType());
-    auto directOutType = RankedTensorType::get({shapeInfo->batch, gemmM, gemmN}, shapeInfo->outType.getElementType());
+    plan.lhs = ensureBatchedTensor(plan.lhs, plan.lhsBatch, plan.m, plan.k, rewriter, loc);
+    plan.rhs = ensureBatchedTensor(plan.rhs, plan.rhsBatch, plan.k, plan.n, rewriter, loc);
+    plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
+    plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
+    auto directOutType = RankedTensorType::get(
+      {plan.batch, plan.m, plan.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
 
-    if (isCompileTimeComputable(rhs)) {
-      const int64_t numKSlices = ceilIntegerDivide(gemmK, crossbarSize.getValue());
-      const int64_t numOutHSlices = ceilIntegerDivide(gemmN, crossbarSize.getValue());
+    if (isCompileTimeComputable(plan.rhs)) {
+      const int64_t numKSlices = ceilIntegerDivide(plan.k, crossbarSize.getValue());
+      const int64_t numOutHSlices = ceilIntegerDivide(plan.n, crossbarSize.getValue());
       const int64_t paddedReductionSize = numKSlices * static_cast<int64_t>(crossbarSize.getValue());
       const int64_t paddedOutCols = numOutHSlices * static_cast<int64_t>(crossbarSize.getValue());
       auto paddedLhsType = RankedTensorType::get(
-        {lhsBatchForGemm, gemmM, paddedReductionSize}, lhsBatchedType.getElementType(), lhsBatchedType.getEncoding());
-      auto paddedRhsType = RankedTensorType::get({shapeInfo->batch, paddedReductionSize, paddedOutCols},
-                                                 rhsBatchedType.getElementType(),
-                                                 rhsBatchedType.getEncoding());
+        {plan.lhsBatch, plan.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
+      auto paddedRhsType = RankedTensorType::get(
+        {plan.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding());
       auto paddedOutType =
-        RankedTensorType::get({shapeInfo->batch, gemmM, paddedOutCols}, shapeInfo->outType.getElementType());
+        RankedTensorType::get({plan.batch, plan.m, paddedOutCols}, shapeInfo->outType.getElementType());
 
-      auto paddedRhs = materializePaddedBatchedWeight(rhs, rhsBatchForGemm, shapeInfo->batch, paddedRhsType, rewriter);
+      auto paddedRhs =
+        materializePaddedBatchedWeight(plan.rhs, plan.rhsBatchShape, plan.outputBatchShape, paddedRhsType, rewriter);
       if (succeeded(paddedRhs)) {
-        Value paddedLhs = createPaddedBatchedInputCompute(lhs, paddedLhsType, rewriter, loc);
-        const int64_t laneCount = shapeInfo->batch * gemmM * numKSlices * numOutHSlices;
-        auto partialPiecesType = RankedTensorType::get({laneCount, static_cast<int64_t>(crossbarSize.getValue())},
-                                                       shapeInfo->outType.getElementType());
+        Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
+        const int64_t laneCount = plan.batch * plan.m * numKSlices * numOutHSlices;
+        auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
+          laneCount, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, shapeInfo->outType.getElementType()));
         auto batchOp = createBatchedVmmBatch(paddedLhs,
                                              *paddedRhs,
                                              paddedLhsType,
-                                             lhsBatchForGemm,
+                                             plan.lhsBatchShape,
                                              paddedRhsType,
-                                             rhsBatchForGemm,
+                                             plan.rhsBatchShape,
+                                             plan.outputBatchShape,
                                              partialPiecesType,
-                                             gemmM,
+                                             plan.m,
                                              numKSlices,
                                              numOutHSlices,
                                              rewriter,
@@ -807,34 +1034,36 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                                     partialPiecesType,
                                                     directOutType,
                                                     paddedOutType,
-                                                    shapeInfo->batch,
+                                                    plan.batch,
                                                     numKSlices,
                                                     rewriter,
                                                     loc);
         if (failed(result))
           return failure();
         Value finalResult = *result;
-        if (useTransposedForm) {
-          auto transposedOutType = RankedTensorType::get({shapeInfo->batch, shapeInfo->m, shapeInfo->n},
+        if (plan.transposedResult) {
+          auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n},
                                                          shapeInfo->outType.getElementType(),
                                                          shapeInfo->outType.getEncoding());
           finalResult =
             ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
               .getResult();
         }
-        finalResult = expandBatchDims(finalResult, shapeInfo->outType, shapeInfo->batchShape.size(), rewriter, loc);
+        finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
         rewriter.replaceOp(matmulOp, finalResult);
         return success();
       }
     }
-    const int64_t laneCount = shapeInfo->batch * gemmM * gemmN;
-    auto scalarPiecesType = RankedTensorType::get({laneCount, 1}, shapeInfo->outType.getElementType());
-    auto batchOp = createBatchedVvdmulBatch(lhs,
-                                            lhsBatchForGemm,
-                                            rhs,
-                                            rhsBatchForGemm,
-                                            lhsBatchedType,
-                                            rhsBatchedType,
+    const int64_t laneCount = plan.batch * plan.m * plan.n;
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(
+      laneCount, RankedTensorType::get({1, 1}, shapeInfo->outType.getElementType()));
+    auto batchOp = createBatchedVvdmulBatch(plan.lhs,
+                                            plan.lhsBatchShape,
+                                            plan.rhs,
+                                            plan.rhsBatchShape,
+                                            plan.outputBatchShape,
+                                            plan.lhsType,
+                                            plan.rhsType,
                                             scalarPiecesType,
                                             directOutType,
                                             rewriter,
@@ -846,15 +1075,15 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     if (failed(result))
       return failure();
     Value finalResult = *result;
-    if (useTransposedForm) {
-      auto transposedOutType = RankedTensorType::get({shapeInfo->batch, shapeInfo->m, shapeInfo->n},
+    if (plan.transposedResult) {
+      auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n},
                                                      shapeInfo->outType.getElementType(),
                                                      shapeInfo->outType.getEncoding());
       finalResult =
         ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
           .getResult();
     }
-    finalResult = expandBatchDims(finalResult, shapeInfo->outType, shapeInfo->batchShape.size(), rewriter, loc);
+    finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
     rewriter.replaceOp(matmulOp, finalResult);
     return success();
   }

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

@@ -5,7 +5,8 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include <algorithm>
-#include <numeric>
+#include <optional>
+#include <type_traits>
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -19,6 +20,85 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
+struct ReduceMeanSemantics {
+  SmallVector<int64_t> axes;
+  int64_t keepdims = 1;
+  bool isIdentity = false;
+};
+
+static bool isNoneValueLike(Value value) { return isa_and_nonnull<ONNXNoneOp>(value.getDefiningOp()); }
+
+static FailureOr<SmallVector<int64_t>> getConstantIntValues(Value value) {
+  auto denseAttr = dyn_cast_or_null<DenseIntElementsAttr>(getHostConstDenseElementsAttr(value));
+  if (!denseAttr)
+    return failure();
+  return SmallVector<int64_t>(denseAttr.getValues<int64_t>().begin(), denseAttr.getValues<int64_t>().end());
+}
+
+static FailureOr<SmallVector<int64_t>> normalizeAxesChecked(ArrayRef<int64_t> axes, int64_t rank) {
+  SmallVector<int64_t> normalizedAxes;
+  normalizedAxes.reserve(axes.size());
+  for (int64_t axis : axes) {
+    auto normalizedAxis = normalizeAxisChecked(axis, rank);
+    if (failed(normalizedAxis))
+      return failure();
+    normalizedAxes.push_back(*normalizedAxis);
+  }
+  llvm::sort(normalizedAxes);
+  normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end());
+  return normalizedAxes;
+}
+
+template <typename ReduceMeanOp, typename ReduceMeanOpAdaptor>
+static FailureOr<ReduceMeanSemantics>
+getReduceMeanSemantics(ReduceMeanOp reduceMeanOp, ReduceMeanOpAdaptor adaptor, int64_t inputRank) {
+  ReduceMeanSemantics semantics;
+  semantics.keepdims = reduceMeanOp.getKeepdims();
+
+  if constexpr (std::is_same_v<ReduceMeanOp, ONNXReduceMeanV13Op>) {
+    auto axes = onnx_mlir::normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputRank);
+    if (failed(axes))
+      return failure();
+    semantics.axes = std::move(*axes);
+    return semantics;
+  }
+  else {
+    if (isNoneValueLike(adaptor.getAxes())) {
+      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
+        semantics.isIdentity = true;
+        return semantics;
+      }
+
+      semantics.axes.reserve(inputRank);
+      for (int64_t axis = 0; axis < inputRank; ++axis)
+        semantics.axes.push_back(axis);
+      return semantics;
+    }
+
+    auto axes = getConstantIntValues(adaptor.getAxes());
+    if (failed(axes))
+      return failure();
+
+    if (axes->empty()) {
+      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
+        semantics.isIdentity = true;
+        return semantics;
+      }
+
+      semantics.axes.reserve(inputRank);
+      for (int64_t axis = 0; axis < inputRank; ++axis)
+        semantics.axes.push_back(axis);
+      return semantics;
+    }
+
+    auto normalizedAxes = normalizeAxesChecked(*axes, inputRank);
+    if (failed(normalizedAxes))
+      return failure();
+    semantics.axes = std::move(*normalizedAxes);
+    return semantics;
+  }
+}
+
 static SmallVector<bool> buildReducedAxesMask(ArrayRef<int64_t> axes, int64_t rank) {
   SmallVector<bool> reducedAxes(rank, false);
   for (int64_t axis : axes) {
@@ -41,14 +121,6 @@ static RankedTensorType getKeepdimsType(RankedTensorType inputType, Type element
   return RankedTensorType::get(shape, elementType, inputType.getEncoding());
 }
 
-static RankedTensorType getCompactKeptType(RankedTensorType inputType, Type elementType, ArrayRef<bool> reducedAxes) {
-  SmallVector<int64_t> shape;
-  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
-    if (!isReduced)
-      shape.push_back(dim);
-  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
-}
-
 static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef<bool> reducedAxes) {
   SmallVector<int64_t> shape;
   shape.reserve(inputType.getRank());
@@ -58,9 +130,7 @@ static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef
 }
 
 static RankedTensorType getLanePackedKeepdimsType(int64_t laneCount, RankedTensorType leafType) {
-  SmallVector<int64_t> shape(leafType.getShape().begin(), leafType.getShape().end());
-  shape.front() = laneCount;
-  return RankedTensorType::get(shape, leafType.getElementType(), leafType.getEncoding());
+  return spatial::getGraphBatchPhysicalResultType(laneCount, leafType);
 }
 
 static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
@@ -110,12 +180,9 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
 
   SmallVector<OpFoldResult> sliceOffsets;
   SmallVector<OpFoldResult> sliceSizes;
-  SmallVector<OpFoldResult> insertOffsets;
-  SmallVector<OpFoldResult> insertSizes(inputType.getRank(), rewriter.getIndexAttr(1));
   SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, inputType.getRank());
   sliceOffsets.reserve(inputType.getRank());
   sliceSizes.reserve(inputType.getRank());
-  insertOffsets.reserve(inputType.getRank());
 
   auto batchOp =
     createSpatComputeBatch(rewriter,
@@ -128,7 +195,6 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                              size_t keptAxisIndex = 0;
                              sliceOffsets.clear();
                              sliceSizes.clear();
-                             insertOffsets.clear();
                              for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
                                if (isReduced) {
                                  sliceOffsets.push_back(rewriter.getIndexAttr(0));
@@ -143,72 +209,90 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                                sliceSizes.push_back(rewriter.getIndexAttr(1));
                              }
 
-                             insertOffsets.push_back(args.lane);
-                             insertOffsets.append(inputType.getRank() - 1, rewriter.getIndexAttr(0));
-
                              Value slice = tensor::ExtractSliceOp::create(
                                rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
                              Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult();
-                             createParallelInsertSliceIntoBatchOutput(
-                               rewriter, loc, reduced, args.outputs.front(), insertOffsets, insertSizes, unitStrides);
+                             publishGraphBatchPhysicalFragment(rewriter, loc, reduced, args.outputs.front(), args.lane);
                            });
   if (failed(batchOp))
     return failure();
   return (*batchOp).getResult(0);
 }
 
-static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
-                                              RankedTensorType keepdimsType,
-                                              RankedTensorType compactKeptType,
-                                              ArrayRef<bool> reducedAxes,
-                                              ConversionPatternRewriter& rewriter,
-                                              Location loc) {
-  auto batchType = cast<RankedTensorType>(batchValue.getType());
-  if (batchType == keepdimsType)
-    return batchValue;
+static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
+    Value batchValue, RankedTensorType keepdimsType,
+    ArrayRef<bool> reducedAxes, ConversionPatternRewriter& rewriter,
+    Location loc) {
+  auto batchType = dyn_cast<RankedTensorType>(batchValue.getType());
+  int64_t rank = keepdimsType.getRank();
+  if (!batchType || !batchType.hasStaticShape()
+      || !keepdimsType.hasStaticShape()
+      || static_cast<int64_t>(reducedAxes.size()) != rank
+      || batchType.getRank() != rank + 1
+      || batchType.getElementType() != keepdimsType.getElementType())
+    return failure();
 
-  SmallVector<ReassociationIndices> collapseToFlat {{}};
-  for (int64_t axis = 0; axis < batchType.getRank(); ++axis)
-    collapseToFlat.front().push_back(axis);
-
-  SmallVector<ReassociationIndices> expandFlatToCompact(1);
-  for (int64_t axis = 0; axis < compactKeptType.getRank(); ++axis)
-    expandFlatToCompact.front().push_back(axis);
-
-  SmallVector<ReassociationIndices> expandCompactToKeepdims;
-  ReassociationIndices pendingLeadingReducedAxes;
+  int64_t laneCount = 1;
+  SmallVector<int64_t> keptAxes;
+  SmallVector<int64_t> keptAxisStrides;
   for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
-    if (isReduced) {
-      if (expandCompactToKeepdims.empty())
-        pendingLeadingReducedAxes.push_back(axis);
-      else
-        expandCompactToKeepdims.back().push_back(axis);
-      continue;
-    }
-
-    expandCompactToKeepdims.emplace_back();
-    auto& group = expandCompactToKeepdims.back();
-    group.append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
-    pendingLeadingReducedAxes.clear();
-    group.push_back(axis);
+    int64_t dim = keepdimsType.getDimSize(axis);
+    if (dim <= 0 || (isReduced && dim != 1))
+      return failure();
+    if (!isReduced)
+      keptAxes.push_back(axis);
   }
-  if (!pendingLeadingReducedAxes.empty())
-    expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
+  keptAxisStrides.resize(keptAxes.size(), 1);
+  for (int64_t index = static_cast<int64_t>(keptAxes.size()) - 1;
+       index >= 0; --index) {
+    keptAxisStrides[index] = laneCount;
+    int64_t dim = keepdimsType.getDimSize(keptAxes[index]);
+    if (laneCount > std::numeric_limits<int64_t>::max() / dim)
+      return failure();
+    laneCount *= dim;
+  }
+  if (batchType.getDimSize(0) != laneCount
+      || llvm::any_of(batchType.getShape().drop_front(),
+                      [](int64_t dim) { return dim != 1; }))
+    return failure();
 
-  auto reshapeCompute =
-    createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
-      auto flatType =
-        RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
-      Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
-      Value compact = flat;
-      if (compactKeptType != flatType)
-        compact = tensor::ExpandShapeOp::create(rewriter, loc, compactKeptType, flat, expandFlatToCompact);
-      Value keepdims = compact;
-      if (keepdimsType != compactKeptType)
-        keepdims = tensor::ExpandShapeOp::create(rewriter, loc, keepdimsType, compact, expandCompactToKeepdims);
-      spatial::SpatYieldOp::create(rewriter, loc, keepdims);
-    });
-  return reshapeCompute.getResult(0);
+  SmallVector<int64_t> operandIndices(laneCount, 0);
+  SmallVector<int64_t> sourceSlots;
+  SmallVector<int64_t> sourceOffsets(laneCount, 0);
+  SmallVector<int64_t> fragmentOffsets;
+  sourceSlots.reserve(laneCount);
+  fragmentOffsets.reserve(laneCount * rank);
+  for (int64_t lane = 0; lane < laneCount; ++lane) {
+    sourceSlots.push_back(lane);
+    size_t keptAxisIndex = 0;
+    for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
+      if (isReduced) {
+        fragmentOffsets.push_back(0);
+        continue;
+      }
+      int64_t dim = keepdimsType.getDimSize(axis);
+      fragmentOffsets.push_back(
+        (lane / keptAxisStrides[keptAxisIndex]) % dim);
+      ++keptAxisIndex;
+    }
+  }
+  SmallVector<int64_t> fragmentSizes(fragmentOffsets.size(), 1);
+  SmallVector<int64_t> fragmentStrides(fragmentOffsets.size(), 1);
+  return spatial::SpatBlueprintOp::create(
+           rewriter, loc, keepdimsType, batchValue, ValueRange {},
+           rewriter.getStringAttr("nchw"),
+           rewriter.getStringAttr("fragmented"),
+           rewriter.getDenseI64ArrayAttr(fragmentOffsets),
+           rewriter.getDenseI64ArrayAttr(fragmentSizes),
+           rewriter.getStringAttr("reduce_mean_keepdims_fragments"),
+           rewriter.getStringAttr("fragment_assembly"),
+           rewriter.getDenseI64ArrayAttr(operandIndices),
+           rewriter.getDenseI64ArrayAttr(sourceSlots),
+           rewriter.getDenseI64ArrayAttr(sourceOffsets),
+           rewriter.getDenseI64ArrayAttr(fragmentStrides),
+           rewriter.getStringAttr("disjoint"),
+           rewriter.getStringAttr("complete"))
+    .getOutput();
 }
 
 static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> reducedAxes) {
@@ -238,14 +322,8 @@ static Value squeezeReducedAxes(Value keepdimsValue,
                                 ArrayRef<bool> reducedAxes,
                                 ConversionPatternRewriter& rewriter,
                                 Location loc) {
-  if (resultType.getRank() == 0) {
-    SmallVector<Value> indices(cast<RankedTensorType>(keepdimsValue.getType()).getRank(),
-                               getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0));
-    Value element = tensor::ExtractOp::create(rewriter, loc, keepdimsValue, indices);
-    return tensor::FromElementsOp::create(rewriter, loc, resultType, ValueRange {element});
-  }
-
-  auto reassociation = buildCollapseReassociation(reducedAxes);
+  SmallVector<ReassociationIndices> reassociation =
+    resultType.getRank() == 0 ? SmallVector<ReassociationIndices> {} : buildCollapseReassociation(reducedAxes);
   if (isCompileTimeComputable(keepdimsValue))
     return tensor::CollapseShapeOp::create(rewriter, loc, resultType, keepdimsValue, reassociation).getResult();
 
@@ -257,11 +335,13 @@ static Value squeezeReducedAxes(Value keepdimsValue,
   return squeezeCompute.getResult(0);
 }
 
-struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
-  using OpConversionPattern::OpConversionPattern;
+template <typename ReduceMeanOp>
+struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
+  using OpConversionPattern<ReduceMeanOp>::OpConversionPattern;
+  using Adaptor = typename ReduceMeanOp::Adaptor;
 
-  LogicalResult matchAndRewrite(ONNXReduceMeanV13Op reduceMeanOp,
-                                ONNXReduceMeanV13OpAdaptor adaptor,
+  LogicalResult matchAndRewrite(ReduceMeanOp reduceMeanOp,
+                                Adaptor adaptor,
                                 ConversionPatternRewriter& rewriter) const override {
     auto inputType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
     auto resultType = dyn_cast<RankedTensorType>(reduceMeanOp.getReduced().getType());
@@ -272,35 +352,53 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
       return success();
     }
 
-    auto axes = normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputType.getRank());
-    if (failed(axes))
-      return failure();
-    SmallVector<bool> reducedAxes = buildReducedAxesMask(*axes, inputType.getRank());
+    auto semantics = getReduceMeanSemantics(reduceMeanOp, adaptor, inputType.getRank());
+    if (failed(semantics))
+      return rewriter.notifyMatchFailure(reduceMeanOp, "requires compile-time constant, in-range ReduceMean axes");
+    if (semantics->isIdentity) {
+      if (inputType != resultType)
+        return rewriter.notifyMatchFailure(
+          reduceMeanOp, "noop_with_empty_axes identity requires the result type to match the input type");
+      rewriter.replaceOp(reduceMeanOp, adaptor.getData());
+      return success();
+    }
+
+    SmallVector<bool> reducedAxes = buildReducedAxesMask(semantics->axes, inputType.getRank());
     if (reducedAxes.empty() && inputType.getRank() != 0)
       return failure();
 
     Location loc = reduceMeanOp.getLoc();
     RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType());
-    RankedTensorType compactKeptType = getCompactKeptType(inputType, resultType.getElementType(), reducedAxes);
     RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
     int64_t laneCount = 1;
-    for (int64_t dim : compactKeptType.getShape())
+    for (auto [dim, isReduced] : llvm::zip_equal(keepdimsType.getShape(), reducedAxes)) {
+      if (isReduced)
+        continue;
+      if (dim <= 0 || laneCount > std::numeric_limits<int32_t>::max() / dim)
+        return rewriter.notifyMatchFailure(
+          reduceMeanOp, "ReduceMean physical lane count is not representable");
       laneCount *= dim;
+    }
     RankedTensorType batchType = getLanePackedKeepdimsType(laneCount, leafType);
 
     auto lanePackedKeepdims =
       buildReduceMeanKeepdimsBatch(adaptor.getData(), reducedAxes, batchType, leafType, rewriter, loc);
     if (failed(lanePackedKeepdims))
       return failure();
-    Value reducedKeepdims =
-      buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc);
+    auto reducedKeepdims = buildReduceMeanKeepdimsBlueprint(
+      *lanePackedKeepdims, keepdimsType, reducedAxes, rewriter, loc);
+    if (failed(reducedKeepdims))
+      return rewriter.notifyMatchFailure(
+        reduceMeanOp,
+        "cannot build physical-fragment ReduceMean keepdims reconstruction");
 
-    if (reduceMeanOp.getKeepdims() != 0) {
-      rewriter.replaceOp(reduceMeanOp, reducedKeepdims);
+    if (semantics->keepdims != 0) {
+      rewriter.replaceOp(reduceMeanOp, *reducedKeepdims);
       return success();
     }
 
-    Value reduced = squeezeReducedAxes(reducedKeepdims, resultType, reducedAxes, rewriter, loc);
+    Value reduced = squeezeReducedAxes(
+      *reducedKeepdims, resultType, reducedAxes, rewriter, loc);
     rewriter.replaceOp(reduceMeanOp, reduced);
     return success();
   }
@@ -309,7 +407,7 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
 } // namespace
 
 void populateReduceMeanPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<ReduceMeanToSpatialCompute>(ctx);
+  patterns.add<ReduceMeanToSpatialCompute<ONNXReduceMeanV13Op>, ReduceMeanToSpatialCompute<ONNXReduceMeanOp>>(ctx);
 }
 
 } // namespace onnx_mlir

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

@@ -16,12 +16,9 @@ struct ReluToSpatialCompute : OpConversionPattern<ONNXReluOp> {
   matchAndRewrite(ONNXReluOp reluOp, ONNXReluOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
     Location loc = reluOp.getLoc();
     Type resultType = reluOp.getResult().getType();
-    constexpr size_t numInputs = 1;
-    auto computeOp = createSpatCompute<numInputs>(rewriter, loc, resultType, {}, adaptor.getX(), [&](Value x) {
-      auto spatReluOp = spatial::SpatReluOp::create(rewriter, loc, resultType, x);
-      spatial::SpatYieldOp::create(rewriter, loc, spatReluOp.getResult());
-    });
-    rewriter.replaceOp(reluOp, computeOp);
+    auto reluPlan = spatial::SpatReluPlanOp::create(
+      rewriter, loc, resultType, adaptor.getX(), rewriter.getStringAttr("nchw"));
+    rewriter.replaceOp(reluOp, reluPlan.getResult());
     return success();
   }
 };

src/PIM/Conversion/ONNXToSpatial/Patterns/Post.cpp

@@ -10,6 +10,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -118,18 +119,16 @@ static LogicalResult mapPromotedInputArguments(ComputeOpTy compute,
 }
 
 // Promotes foldable helper chains from runtime inputs to weights to avoid artificial compute inputs.
-struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCompute> {
-  using OpRewritePattern<spatial::SpatCompute>::OpRewritePattern;
+struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatGraphCompute> {
+  using OpRewritePattern<spatial::SpatGraphCompute>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatCompute compute, PatternRewriter& rewriter) const override {
+  LogicalResult matchAndRewrite(spatial::SpatGraphCompute compute, PatternRewriter& rewriter) const override {
     auto promoted = computePromotedOperands(compute);
     if (failed(promoted))
       return rewriter.notifyMatchFailure(compute, "no weight-like inputs to promote");
     Block& oldBlock = compute.getBody().front();
 
     rewriter.setInsertionPointAfter(compute);
-    auto newCompute = spatial::SpatCompute::create(
-      rewriter, compute.getLoc(), compute.getResultTypes(), promoted->newWeights, promoted->newInputs);
     SmallVector<Type> newBlockArgTypes;
     SmallVector<Location> newBlockArgLocs;
     for (Value weight : promoted->newWeights) {
@@ -138,10 +137,14 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCom
     }
     llvm::append_range(newBlockArgTypes, promoted->newInputTypes);
     llvm::append_range(newBlockArgLocs, promoted->newInputLocs);
-    auto* newBlock = rewriter.createBlock(
-      &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs);
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(promoted->newWeights.size()), static_cast<int>(promoted->newInputs.size())});
+    auto newCompute = createEmptySpatGraphCompute(rewriter,
+                                                  compute.getLoc(),
+                                                  compute.getResultTypes(),
+                                                  promoted->newWeights,
+                                                  promoted->newInputs,
+                                                  TypeRange(newBlockArgTypes),
+                                                  newBlockArgLocs);
+    auto* newBlock = &newCompute.getBody().front();
     rewriter.setInsertionPointToStart(newBlock);
 
     IRRewriter bodyRewriter(rewriter.getContext());
@@ -182,10 +185,10 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCom
 };
 
 // Promotes foldable batch helper chains to weights while preserving compact compute_batch IR.
-struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::SpatComputeBatch> {
-  using OpRewritePattern<spatial::SpatComputeBatch>::OpRewritePattern;
+struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::SpatGraphComputeBatch> {
+  using OpRewritePattern<spatial::SpatGraphComputeBatch>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatComputeBatch compute, PatternRewriter& rewriter) const override {
+  LogicalResult matchAndRewrite(spatial::SpatGraphComputeBatch compute, PatternRewriter& rewriter) const override {
     auto promoted = computePromotedOperands(compute);
     if (failed(promoted))
       return rewriter.notifyMatchFailure(compute, "no weight-like batch inputs to promote");
@@ -193,12 +196,6 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
 
     rewriter.setInsertionPointAfter(compute);
 
-    auto laneCountAttr = pim::getCheckedI32Attr(
-      rewriter, compute, static_cast<uint64_t>(compute.getLaneCount()), "promoted compute_batch lane count");
-    if (failed(laneCountAttr))
-      return failure();
-    auto newCompute = spatial::SpatComputeBatch::create(
-      rewriter, compute.getLoc(), compute.getResultTypes(), *laneCountAttr, promoted->newWeights, promoted->newInputs);
     auto laneArg = compute.getLaneArgument();
     if (!laneArg)
       return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument");
@@ -223,23 +220,30 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
       newBlockArgLocs.push_back(outputArg->getLoc());
     }
 
-    auto* newBlock = rewriter.createBlock(
-      &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs);
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(promoted->newWeights.size()), static_cast<int>(promoted->newInputs.size())});
+    auto newCompute = createEmptySpatGraphComputeBatch(rewriter,
+                                                       compute.getLoc(),
+                                                       compute.getResultTypes(),
+                                                       compute.getLaneCount(),
+                                                       promoted->newWeights,
+                                                       promoted->newInputs,
+                                                       TypeRange(newBlockArgTypes),
+                                                       newBlockArgLocs);
+    if (failed(newCompute))
+      return failure();
+    auto* newBlock = &(*newCompute).getBody().front();
     rewriter.setInsertionPointToStart(newBlock);
 
     IRRewriter bodyRewriter(rewriter.getContext());
     bodyRewriter.setInsertionPointToStart(newBlock);
 
     IRMapping mapper;
-    auto newLaneArg = newCompute.getLaneArgument();
+    auto newLaneArg = (*newCompute).getLaneArgument();
     if (!newLaneArg)
       return rewriter.notifyMatchFailure(compute, "missing rewritten compute_batch lane block argument");
     mapper.map(*laneArg, *newLaneArg);
     for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) {
       auto oldWeightArg = compute.getWeightArgument(weightIndex);
-      auto newWeightArg = newCompute.getWeightArgument(weightIndex);
+      auto newWeightArg = (*newCompute).getWeightArgument(weightIndex);
       if (!oldWeightArg || !newWeightArg)
         return rewriter.notifyMatchFailure(compute, "missing compute_batch weight block argument during rewrite");
       mapper.map(*oldWeightArg, *newWeightArg);
@@ -249,7 +253,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
           *promoted,
           bodyRewriter,
           mapper,
-          [&](size_t index) { return newCompute.getInputArgument(index); },
+          [&](size_t index) { return (*newCompute).getInputArgument(index); },
           rewriter)))
       return failure();
     for (auto resultIndex : llvm::seq<size_t>(0, compute.getNumResults())) {
@@ -263,7 +267,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
     for (Operation& op : oldBlock)
       rewriter.clone(op, mapper);
 
-    rewriter.replaceOp(compute, newCompute.getResults());
+    rewriter.replaceOp(compute, (*newCompute).getResults());
     return success();
   }
 };
@@ -281,8 +285,8 @@ void annotateWeightsConstants(func::FuncOp funcOp) {
   });
 }
 
-bool requiresPostRewrite(spatial::SpatCompute computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
+bool requiresPostRewrite(spatial::SpatGraphCompute computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
 
-bool requiresPostRewrite(spatial::SpatComputeBatch computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
+bool requiresPostRewrite(spatial::SpatGraphComputeBatch computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,112 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<int64_t> normalizeFlattenAxis(int64_t axis, int64_t rank) {
+  int64_t normalizedAxis = axis < 0 ? rank + axis : axis;
+  if (normalizedAxis < 0 || normalizedAxis > rank)
+    return failure();
+  return normalizedAxis;
+}
+
+static int64_t product(ArrayRef<int64_t> values) {
+  int64_t result = 1;
+  for (int64_t value : values)
+    result *= value;
+  return result;
+}
+
+static SmallVector<ReassociationIndices> getCollapseTo1DReassociation(int64_t rank) {
+  SmallVector<ReassociationIndices> reassociation(1);
+  reassociation.front().reserve(rank);
+  for (int64_t dim = 0; dim < rank; ++dim)
+    reassociation.front().push_back(dim);
+  return reassociation;
+}
+
+static SmallVector<ReassociationIndices> getExpandFrom1DReassociation(int64_t rank) {
+  SmallVector<ReassociationIndices> reassociation(1);
+  reassociation.front().reserve(rank);
+  for (int64_t dim = 0; dim < rank; ++dim)
+    reassociation.front().push_back(dim);
+  return reassociation;
+}
+
+static Value buildFlatten(Value input,
+                          RankedTensorType sourceType,
+                          RankedTensorType resultType,
+                          int64_t axis,
+                          ConversionPatternRewriter& rewriter,
+                          Location loc) {
+  if (sourceType == resultType)
+    return input;
+
+  if (axis > 0 && axis < sourceType.getRank()) {
+    SmallVector<ReassociationIndices> reassociation(2);
+    for (int64_t dim = 0; dim < axis; ++dim)
+      reassociation[0].push_back(dim);
+    for (int64_t dim = axis; dim < sourceType.getRank(); ++dim)
+      reassociation[1].push_back(dim);
+    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation);
+  }
+
+  Value flattened = input;
+  if (sourceType.getRank() != 1) {
+    auto flatType = RankedTensorType::get({sourceType.getNumElements()}, sourceType.getElementType());
+    flattened = tensor::CollapseShapeOp::create(
+      rewriter, loc, flatType, flattened, getCollapseTo1DReassociation(sourceType.getRank()));
+  }
+  return tensor::ExpandShapeOp::create(
+    rewriter, loc, resultType, flattened, getExpandFrom1DReassociation(resultType.getRank()));
+}
+
+struct Flatten : OpConversionPattern<ONNXFlattenOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXFlattenOp flattenOp,
+                                ONNXFlattenOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto sourceType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+    auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOperation()->getResult(0).getType());
+    if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
+      return failure();
+    if (!hasStaticPositiveShape(sourceType) || !hasStaticPositiveShape(resultType) || resultType.getRank() != 2)
+      return failure();
+
+    auto axis = normalizeFlattenAxis(flattenOp.getAxis(), sourceType.getRank());
+    if (failed(axis))
+      return failure();
+
+    int64_t outerDim = product(sourceType.getShape().take_front(*axis));
+    int64_t innerDim = product(sourceType.getShape().drop_front(*axis));
+    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());
+    });
+  }
+};
+
+} // namespace
+
+void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Flatten>(ctx); }
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,189 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <algorithm>
+#include <optional>
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<SmallVector<int64_t>> getConstantIntValues(Value value) {
+  auto denseAttr = dyn_cast_or_null<DenseIntElementsAttr>(getHostConstDenseElementsAttr(value));
+  if (!denseAttr)
+    return failure();
+  return SmallVector<int64_t>(denseAttr.getValues<int64_t>().begin(), denseAttr.getValues<int64_t>().end());
+}
+
+static bool isNoneValueLike(Value value) { return isa_and_nonnull<ONNXNoneOp>(value.getDefiningOp()); }
+
+static FailureOr<Value> buildSlice(Value data,
+                                   RankedTensorType dataType,
+                                   RankedTensorType resultType,
+                                   ArrayRef<int64_t> starts,
+                                   ArrayRef<int64_t> ends,
+                                   std::optional<ArrayRef<int64_t>> axes,
+                                   std::optional<ArrayRef<int64_t>> steps,
+                                   ConversionPatternRewriter& rewriter,
+                                   Location loc) {
+  int64_t rank = dataType.getRank();
+  if (!dataType.hasStaticShape() || !resultType.hasStaticShape() || resultType.getRank() != rank)
+    return failure();
+
+  if (starts.size() != ends.size())
+    return failure();
+  if (axes && axes->size() != starts.size())
+    return failure();
+  if (steps && steps->size() != starts.size())
+    return failure();
+
+  SmallVector<int64_t> normalizedAxes;
+  if (axes) {
+    SmallVector<bool> seenAxes(rank, false);
+    normalizedAxes.reserve(axes->size());
+    for (int64_t axis : *axes) {
+      auto normalizedAxis = normalizeAxisChecked(axis, rank);
+      if (failed(normalizedAxis))
+        return failure();
+      if (seenAxes[*normalizedAxis])
+        return failure();
+      seenAxes[*normalizedAxis] = true;
+      normalizedAxes.push_back(*normalizedAxis);
+    }
+  }
+  else {
+    if (starts.size() > static_cast<size_t>(rank))
+      return failure();
+    normalizedAxes.reserve(starts.size());
+    for (size_t i = 0; i < starts.size(); ++i)
+      normalizedAxes.push_back(static_cast<int64_t>(i));
+  }
+
+  SmallVector<int64_t> normalizedSteps;
+  if (steps)
+    normalizedSteps.assign(steps->begin(), steps->end());
+  else
+    normalizedSteps.assign(starts.size(), 1);
+
+  SmallVector<int64_t> computedShape(dataType.getShape().begin(), dataType.getShape().end());
+  SmallVector<OpFoldResult> offsets = getZeroOffsets(rewriter, rank);
+  SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, dataType.getShape());
+  SmallVector<OpFoldResult> strides = getUnitStrides(rewriter, rank);
+
+  for (auto [sliceIndex, axis] : llvm::enumerate(normalizedAxes)) {
+    int64_t step = normalizedSteps[sliceIndex];
+    if (step <= 0)
+      return failure();
+
+    int64_t dimSize = dataType.getShape()[axis];
+    int64_t start = starts[sliceIndex];
+    int64_t end = ends[sliceIndex];
+
+    start = normalizeIndex(start, dimSize);
+    end = normalizeIndex(end, dimSize);
+
+    start = std::clamp(start, int64_t {0}, dimSize);
+    end = std::clamp(end, int64_t {0}, dimSize);
+
+    int64_t extent = std::max(end - start, int64_t {0});
+    int64_t size = (extent + step - 1) / step;
+
+    offsets[axis] = rewriter.getIndexAttr(start);
+    sizes[axis] = rewriter.getIndexAttr(size);
+    strides[axis] = rewriter.getIndexAttr(step);
+    computedShape[axis] = size;
+  }
+
+  if (llvm::ArrayRef(computedShape) != resultType.getShape())
+    return failure();
+
+  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, data, offsets, sizes, strides).getResult();
+}
+
+struct Slice final : OpConversionPattern<ONNXSliceOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXSliceOp sliceOp,
+                                ONNXSliceOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto dataType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
+    auto resultType = dyn_cast<RankedTensorType>(sliceOp.getResult().getType());
+    if (!dataType || !resultType || !dataType.hasStaticShape() || !resultType.hasStaticShape())
+      return failure();
+
+    auto starts = getConstantIntValues(adaptor.getStarts());
+    auto ends = getConstantIntValues(adaptor.getEnds());
+    if (failed(starts))
+      return rewriter.notifyMatchFailure(sliceOp, "requires compile-time constant starts");
+    if (failed(ends))
+      return rewriter.notifyMatchFailure(sliceOp, "requires compile-time constant ends");
+
+    std::optional<SmallVector<int64_t>> axes;
+    if (!isNoneValueLike(adaptor.getAxes())) {
+      auto parsedAxes = getConstantIntValues(adaptor.getAxes());
+      if (failed(parsedAxes))
+        return rewriter.notifyMatchFailure(sliceOp, "requires compile-time constant axes when present");
+      axes = std::move(*parsedAxes);
+    }
+
+    std::optional<SmallVector<int64_t>> steps;
+    if (!isNoneValueLike(adaptor.getSteps())) {
+      auto parsedSteps = getConstantIntValues(adaptor.getSteps());
+      if (failed(parsedSteps))
+        return rewriter.notifyMatchFailure(sliceOp, "requires compile-time constant steps when present");
+      steps = std::move(*parsedSteps);
+      if (llvm::any_of(*steps, [](int64_t step) { return step <= 0; }))
+        return rewriter.notifyMatchFailure(sliceOp, "supports only positive constant steps");
+    }
+
+    ArrayRef<int64_t> startsRef = *starts;
+    ArrayRef<int64_t> endsRef = *ends;
+    std::optional<ArrayRef<int64_t>> axesRef = axes ? std::optional<ArrayRef<int64_t>>(ArrayRef<int64_t>(*axes))
+                                                    : std::nullopt;
+    std::optional<ArrayRef<int64_t>> stepsRef = steps ? std::optional<ArrayRef<int64_t>>(ArrayRef<int64_t>(*steps))
+                                                       : std::nullopt;
+
+    Location loc = sliceOp.getLoc();
+    auto tryBuildSlice = [&](Value data) {
+      return buildSlice(data, dataType, resultType, startsRef, endsRef, axesRef, stepsRef, rewriter, loc);
+    };
+
+    if (isCompileTimeComputable(adaptor.getData())) {
+      auto sliced = tryBuildSlice(adaptor.getData());
+      if (failed(sliced))
+        return rewriter.notifyMatchFailure(sliceOp, "failed to normalize static slice parameters");
+      rewriter.replaceOp(sliceOp, *sliced);
+      return success();
+    }
+
+    auto computeOp =
+      createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, adaptor.getData(), [&](Value data) {
+        auto sliced = tryBuildSlice(data);
+        if (failed(sliced))
+          return failure();
+        spatial::SpatYieldOp::create(rewriter, loc, *sliced);
+        return success();
+      });
+    if (failed(computeOp))
+      return rewriter.notifyMatchFailure(sliceOp, "failed to build runtime tensor.extract_slice lowering");
+
+    rewriter.replaceOp(sliceOp, computeOp->getResults());
+    return success();
+  }
+};
+
+} // namespace
+
+void populateSlicePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Slice>(ctx); }
+
+} // namespace onnx_mlir

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

@@ -5,7 +5,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -52,35 +52,12 @@ static FailureOr<Value> materializeTransposedConstant(Value input,
     return failure();
   }
 
-  if (denseAttr.isSplat())
-    return getOrCreateConstant(rewriter,
-                               rewriter.getInsertionBlock()->getParentOp(),
-                               DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>()),
-                               resultType);
-
-  SmallVector<Attribute> inputValues(denseAttr.getValues<Attribute>());
-  SmallVector<Attribute> resultValues(inputValues.size());
-  SmallVector<int64_t> inputStrides = computeRowMajorStrides(inputType.getShape());
-  SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultType.getShape());
-  SmallVector<int64_t> inputIndices(inputType.getRank(), 0);
-
-  for (auto [linearIndex, value] : llvm::enumerate(inputValues)) {
-    int64_t remaining = static_cast<int64_t>(linearIndex);
-    for (int64_t dim = 0; dim < inputType.getRank(); ++dim) {
-      inputIndices[dim] = inputStrides.empty() ? 0 : remaining / inputStrides[dim];
-      remaining = inputStrides.empty() ? 0 : remaining % inputStrides[dim];
-    }
-
-    int64_t resultLinearIndex = 0;
-    for (int64_t dim = 0; dim < resultType.getRank(); ++dim)
-      resultLinearIndex += inputIndices[permutation[dim]] * resultStrides[dim];
-
-    resultValues[resultLinearIndex] = value;
-  }
-
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return failure();
   return getOrCreateConstant(rewriter,
                              rewriter.getInsertionBlock()->getParentOp(),
-                             DenseElementsAttr::get(resultType, resultValues),
+                             *transposedAttr,
                              resultType);
 }
 

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -0,0 +1,21 @@
+#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 {
+
+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);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -0,0 +1,244 @@
+#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 = "nchw_row_strip";
+
+enum class SelectedLayout {
+  DenseNchw,
+  NchwRowStrip,
+};
+
+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::NchwRowStrip;
+  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
+    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
+    return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  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>(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 (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
+    return succeeded(canConsumeAndProduceRowStrip(convPlan));
+  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::NchwRowStrip)
+    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 (getSelectedLayout(layouts, convPlan.getInput()) != SelectedLayout::NchwRowStrip
+      && !hasRowStripConsumer(convPlan.getResult()))
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::NchwRowStrip;
+}
+
+static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
+                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip)
+    return SelectedLayout::DenseNchw;
+  if (!hasRowStripConsumer(reluPlan.getResult()))
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::NchwRowStrip;
+}
+
+static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
+                                          llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip)
+    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::NchwRowStrip;
+}
+
+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 = chooseReluLayout(reluPlan, layouts);
+          if (layouts[reluPlan.getResult()] != selected) {
+            layouts[reluPlan.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;
+        }
+      }
+    }
+
+    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 reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
+        producedValue = reluPlan.getResult();
+      else
+        continue;
+
+      if (getSelectedLayout(layouts, producedValue) != SelectedLayout::NchwRowStrip)
+        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/SpatialToGraphviz/CMakeLists.txt

@@ -1,17 +0,0 @@
-add_onnx_mlir_rewriter(SpatialToGraphviz)
-
-add_pim_library(OMSpatialToGraphviz
-  SpatialToGraphviz.cpp
-
-  EXCLUDE_FROM_OM_LIBS
-
-  LINK_LIBS PUBLIC
-  MLIRTosaDialect
-  OMCompilerOptions
-  OMPimCommon
-  OMONNXOps
-  SpatialOps
-
-  ACCEL_INCLUDE_DIRS PRIVATE
-  ${PIM_GENERATED_INCLUDE_DIRS}
-)

src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp

@@ -1,259 +0,0 @@
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tosa/IR/TosaOps.h"
-#include "mlir/IR/Block.h"
-#include "mlir/IR/Diagnostics.h"
-#include "mlir/IR/Value.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Support/LLVM.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-
-#include "llvm/Support/Casting.h"
-#include "llvm/Support/Format.h"
-
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
-#define FORMAT_OPERATION(op) 'x' << llvm::format_hex_no_prefix(reinterpret_cast<size_t>(op), 0)
-#define FORMAT_ARGUMENT(computeOpPointer, argumentNum) llvm::format("Arg_%p_%u", computeOpPointer, argumentNum)
-
-using namespace mlir;
-
-namespace onnx_mlir {
-
-namespace {
-
-struct SpatialToGraphvizPass : public PassWrapper<SpatialToGraphvizPass, OperationPass<ModuleOp>> {
-
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialToGraphvizPass)
-
-  StringRef getArgument() const override { return "convert-spatial-to-graphviz"; }
-
-  StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
-
-  SpatialToGraphvizPass(raw_ostream& os = llvm::errs())
-  : os(os) {}
-  SpatialToGraphvizPass(const SpatialToGraphvizPass& pass)
-  : SpatialToGraphvizPass(pass.os) {}
-  void runOnOperation() final;
-
-private:
-  raw_ostream& os;
-
-  /**
-   * Draws the subgraph for a given spatial::SpatCompute, including:
-   * 1. Input nodes (block arguments)
-   * 2. Operations
-   * 3. Edges between yield (output) and its users
-   *
-   * @param op The spatial::SpatCompute to draw the subgraph for.
-   * @param computeNum The number of the compute operation.
-   */
-  void drawComputeOpSubgraph(spatial::SpatCompute op, size_t computeNum) {
-    os << "\tsubgraph cluster" << computeNum << " {\n\t\tlabel=\"Compute" << computeNum << "\";\n"
-       << "\t\tstyle=filled;\n"
-       << "\t\tcolor=lightblue;\n";
-
-    Block& block = op.getBody().front();
-
-    // Inputs
-    size_t inputNum = 0;
-    for (BlockArgument& input : block.getArguments()) {
-
-      auto fromOp = FORMAT_ARGUMENT(op.getOperation(), inputNum);
-
-      os << "\t\t" << fromOp << " [label=\"Arg" << inputNum << "\",shape=box];\n";
-      for (auto userOp : input.getUsers())
-        os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-      inputNum++;
-    }
-
-    // Iterate operations
-    for (auto& childOp : block.getOperations()) {
-      os << "\t\t" << FORMAT_OPERATION(&childOp) << " [label=\"" << childOp.getName() << "\"];\n";
-
-      drawEdgesFromOpToItsUsers(&childOp);
-    }
-
-    os << "\t}\n";
-
-    // Draw edges from the yield to the users of this computeOp
-    Operation* yieldOp = block.getTerminator();
-    if (!isa<spatial::SpatYieldOp>(yieldOp)) {
-      yieldOp->emitError("Terminator of block must be YieldOp ???");
-      signalPassFailure();
-      return;
-    }
-
-    for (auto computeOpResult : op->getResults()) {
-      for (auto& computeOpUse : computeOpResult.getUses()) {
-        auto toOp = FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber());
-        os << "\t" << FORMAT_OPERATION(yieldOp) << " -> " << toOp << ";\n";
-      }
-    }
-  }
-
-  /**
-   * @brief Draws the subgraph for a concatOp.
-   *
-   * This function draws a subgraph for a concatOp. The subgraph consists of a
-   * node for each input of the concatOp, as well as an output node. Edges are
-   * created from the output node to each user of the concatOp.
-   *
-   * @param concatOp The concatOp for which the subgraph is drawn.
-   * @param concatOpNum The number of the concatOp.
-   */
-  void drawConcatOpSubgraph(Operation* concatOp, size_t concatOpNum) {
-    os << "\tsubgraph clusterconcat" << concatOpNum << " {\n\t\tlabel=\"ConcatOp" << concatOpNum << "\";\n"
-       << "\t\tstyle=filled;\n"
-       << "\t\tcolor=orange;\n";
-
-    // Inputs
-    size_t inputNum = 0;
-    for (Value input : concatOp->getOperands()) {
-      auto fromOp = FORMAT_ARGUMENT(concatOp, inputNum);
-
-      os << "\t\t" << fromOp << " [label=\"Input" << inputNum << "\"];\n";
-      for (auto userOp : input.getUsers())
-        os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-      inputNum++;
-    }
-
-    // Output
-    os << "\t\t" << FORMAT_OPERATION(concatOp) << " [label=Out];\n";
-
-    os << "\t}\n";
-
-    // Edges from output to users
-
-    for (auto& computeOpUse : concatOp->getResult(0).getUses()) {
-      os << "\t" << FORMAT_OPERATION(concatOp) << " -> "
-         << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber()) << ";\n";
-    }
-  }
-
-  /**
-   * Draws the ExtractSliceOp in the graph visualization.
-   *
-   * This function takes a tensor::ExtractSliceOp and adds the corresponding
-   * node and edges to the graph visualization. It creates a node with the
-   * label as the static offsets attribute of the sliceOp, and connects it to
-   * the compute operations that use the result of the sliceOp.
-   *
-   * @param sliceOp The tensor::ExtractSliceOp to be drawn in the graph
-   * visualization.
-   */
-  void drawExtractSliceOp(tensor::ExtractSliceOp sliceOp) {
-    auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
-    os << "\t" << nodeId << " [label=\"Slice: ";
-    sliceOp.getStaticOffsetsAttr().print(os);
-    os << "\",color=lawngreen];\n";
-
-    for (auto& computeOpUse : sliceOp.getResult().getUses()) {
-      os << "\t" << nodeId << " -> " << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber())
-         << ";\n";
-    }
-  }
-
-  void drawBiasTileOp(tensor::ExtractSliceOp sliceOp) {
-    auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
-    os << "\t" << nodeId << " [label=\"Bias: ";
-    sliceOp.getStaticOffsetsAttr().print(os);
-    os << "\",color=lightpink];\n";
-
-    for (auto user : sliceOp.getResult().getUsers())
-      os << "\t" << nodeId << " -> " << FORMAT_OPERATION(user) << ";\n";
-  }
-
-  /**
-   * Draws edges from the given operation to its users.
-   *
-   * @param fromOp The operation from which the edges are drawn.
-   */
-  void drawEdgesFromOpToItsUsers(mlir::Operation* fromOp) {
-    for (auto result : fromOp->getResults())
-      for (auto userOp : result.getUsers())
-        os << "\t\t" << FORMAT_OPERATION(fromOp) << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-  }
-
-  /**
-   * Draws input node and edges for the given `funcOp`.
-   *
-   * @param funcOp The `funcOp` for which to draw input nodes and edges.
-   */
-  void drawInputNodesAndEdges(func::FuncOp& funcOp) {
-    os << "\tinput [label=\"Module Input\",color=green];\n";
-
-    size_t funcOpArgNum = 0;
-    for (BlockArgument& arg : funcOp.getArguments()) {
-
-      for (auto& useOp : arg.getUses()) {
-        os << "\tinput -> " << FORMAT_ARGUMENT(useOp.getOwner(), useOp.getOperandNumber()) << "[label=" << funcOpArgNum
-           << "];\n";
-      }
-      funcOpArgNum++;
-    }
-  }
-};
-
-void SpatialToGraphvizPass::runOnOperation() {
-  ModuleOp module = getOperation();
-
-  auto entryFunc = getPimEntryFunc(module);
-  if (failed(entryFunc)) {
-    module.emitError("failed to locate the PIM entry function for Spatial graph visualization");
-    signalPassFailure();
-    return;
-  }
-  func::FuncOp func = *entryFunc;
-
-  os << "digraph G {\n"
-     << "\tnode [style=filled,color=white];\n";
-
-  size_t computeNum = 0;
-  size_t concatNum = 0;
-
-  // Iterate over the ComputeOps within FuncOp:
-  // 1. Print their subgraph
-  // 2. Print the edges from its inputs to its outputs
-  for (Operation& op : func.getOps()) {
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
-      drawComputeOpSubgraph(computeOp, computeNum++);
-    }
-    else if (auto concatOp = dyn_cast<tensor::ConcatOp>(op)) {
-      drawConcatOpSubgraph(concatOp, concatNum++);
-    }
-    else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
-      auto producerOp = extractSliceOp->getOperand(0).getDefiningOp();
-      if (producerOp) {
-        // Skip extractSliceOp if producer is constant weights (ONNXConstantOp)
-        if (llvm::isa<ONNXConstantOp>(producerOp))
-          continue;
-        // If produced by tosa::ReshapeOp (i.e. it is a bias tile) connect
-        // directly to its user, which is not a ComputeOp argument.
-        if (llvm::isa<tosa::ReshapeOp>(producerOp)) {
-          drawBiasTileOp(extractSliceOp);
-          continue;
-        }
-      }
-
-      drawExtractSliceOp(extractSliceOp);
-    }
-  }
-
-  // Draw input node, and edges to it users
-  drawInputNodesAndEdges(func);
-
-  // Draw output node (use the return Operation - argument number=0 - as nodeId)
-  auto returnOp = func.getBody().front().getTerminator();
-  os << '\t' << FORMAT_ARGUMENT(returnOp, 0) << " [label=\"Module Output\",color=green];\n";
-
-  os << "}\n";
-}
-
-} // namespace
-
-std::unique_ptr<Pass> createSpatialToGraphvizPass() { return std::make_unique<SpatialToGraphvizPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -2,13 +2,16 @@
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Matchers.h"
+#include <limits>
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -26,24 +29,6 @@ static bool isUsedOnlyAsExplicitHostOperand(Value value) {
   });
 }
 
-static FailureOr<SmallVector<int32_t>> getPimCoreIdsForBatchOp(spatial::SpatComputeBatch computeBatchOp,
-                                                               size_t& fallbackCoreId) {
-  if (auto coreIdsAttr = computeBatchOp->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName))
-    return SmallVector<int32_t>(coreIdsAttr.asArrayRef().begin(), coreIdsAttr.asArrayRef().end());
-
-  SmallVector<int32_t> coreIds;
-  coreIds.reserve(static_cast<size_t>(computeBatchOp.getLaneCount()));
-  for (uint32_t lane = 0; lane < computeBatchOp.getLaneCount(); ++lane) {
-    auto checkedCoreId =
-      pim::checkedI32(static_cast<uint64_t>(fallbackCoreId), computeBatchOp, "fallback spatial compute_batch core id");
-    if (failed(checkedCoreId))
-      return failure();
-    coreIds.push_back(*checkedCoreId);
-    ++fallbackCoreId;
-  }
-  return coreIds;
-}
-
 static FailureOr<unsigned> getDirectReturnOperandIndex(OpResult result) {
   if (!result.hasOneUse())
     return failure();
@@ -54,6 +39,188 @@ static FailureOr<unsigned> getDirectReturnOperandIndex(OpResult result) {
   return result.getUses().begin()->getOperandNumber();
 }
 
+static FailureOr<SmallVector<FragmentAssemblyCopy, 8>>
+collectFragmentAssemblyCopiesFromBlueprint(spatial::SpatBlueprintOp blueprint,
+                                          IRMapping& mapper,
+                                          int64_t lane,
+                                          unsigned hostTargetIndex,
+                                          Value fixedSource = {}) {
+  SmallVector<FragmentAssemblyCopy, 8> copies;
+  auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  if (!resultType || !resultType.hasStaticShape())
+    return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor results");
+
+  std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+  std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
+  if (!operandIndicesAttr || !fragmentStridesAttr)
+    return blueprint.emitOpError(
+      "fragment assembly lowering requires explicit operand indices and unit strides");
+
+  ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+  std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+  if (!sourceOffsetsAttr)
+    return blueprint.emitOpError("fragment assembly lowering requires explicit source offsets");
+  ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+  ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
+  ArrayRef<int64_t> flatStrides = *fragmentStridesAttr;
+  int64_t rank = resultType.getRank();
+
+  SmallVector<Value> fragmentOperands {blueprint.getInput()};
+  llvm::append_range(fragmentOperands, blueprint.getFragments());
+  if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                              rank,
+                                              fragmentOperands.size(),
+                                              operandIndices,
+                                              sourceOffsets,
+                                              flatOffsets,
+                                              flatSizes,
+                                              flatStrides)))
+    return failure();
+
+  SmallVector<int64_t> hostStrides = computeRowMajorStrides(resultType.getShape());
+  for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+    Value source = fixedSource ? fixedSource : mapper.lookupOrDefault(fragmentOperands[operandIndices[fragmentIndex]]);
+    auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+    if (!sourceType || !sourceType.hasStaticShape())
+      return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor operands");
+
+    size_t elementSize = getElementTypeSizeInBytes(sourceType.getElementType());
+    SmallVector<int64_t, 4> fragmentOffsets;
+    SmallVector<int64_t, 4> fragmentSizes;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t flatIndex = fragmentIndex * rank + dim;
+      if (flatStrides[flatIndex] != 1)
+        return blueprint.emitOpError("fragment assembly lowering only supports unit strides");
+      fragmentOffsets.push_back(flatOffsets[flatIndex]);
+      fragmentSizes.push_back(flatSizes[flatIndex]);
+    }
+
+    if (failed(forEachContiguousDestinationChunk(
+          resultType.getShape(),
+          fragmentOffsets,
+          fragmentSizes,
+          [&](ArrayRef<int64_t> chunkOffsets, int64_t relativeSourceOffset, int64_t chunkElements) -> LogicalResult {
+            int64_t hostElementOffset = 0;
+            for (auto [dim, offset] : llvm::enumerate(chunkOffsets))
+              hostElementOffset += offset * hostStrides[dim];
+
+            FragmentAssemblyCopy copy;
+            copy.source = source;
+            copy.sourceType = sourceType;
+            copy.hostTargetIndex = hostTargetIndex;
+            copy.lane = lane;
+            copy.sourceByteOffset = (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+            copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
+            copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
+            copies.push_back(copy);
+            return success();
+          })))
+      return failure();
+  }
+
+  return copies;
+}
+
+static FailureOr<SmallVector<FragmentAssemblyCopy, 8>>
+collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedResultType, uint32_t laneCount) {
+  SmallVector<FragmentAssemblyCopy, 8> copies;
+  if (!packedResultType.hasStaticShape() || laneCount == 0)
+    return failure();
+
+  int64_t packedElementCount = packedResultType.getNumElements();
+  if (packedElementCount % static_cast<int64_t>(laneCount) != 0)
+    return failure();
+  int64_t payloadElementCount = packedElementCount / static_cast<int64_t>(laneCount);
+  size_t elementSize = getElementTypeSizeInBytes(packedResultType.getElementType());
+
+  for (OpOperand& use : result.getUses()) {
+    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)
+      return failure();
+    if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
+      return failure();
+
+    auto hostResultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+    std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
+    if (!hostResultType || !hostResultType.hasStaticShape() || !stridesAttr)
+      return failure();
+
+    ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+    ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+    ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
+    if (sourceSlots.size() != operandIndices.size())
+      return failure();
+    ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
+    ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
+    ArrayRef<int64_t> flatStrides = *stridesAttr;
+    int64_t rank = hostResultType.getRank();
+    unsigned returnIndex = blueprint.getOutput().getUses().begin()->getOperandNumber();
+    SmallVector<Value> fragmentOperands {blueprint.getInput()};
+    llvm::append_range(fragmentOperands, blueprint.getFragments());
+    if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                                rank,
+                                                fragmentOperands.size(),
+                                                operandIndices,
+                                                sourceOffsets,
+                                                flatOffsets,
+                                                flatSizes,
+                                                flatStrides)))
+      return failure();
+    SmallVector<int64_t> hostStrides = computeRowMajorStrides(hostResultType.getShape());
+    for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+      if (operandIndices[fragmentIndex] != static_cast<int64_t>(use.getOperandNumber()))
+        continue;
+
+      int64_t sourceElementOffset =
+        sourceSlots[fragmentIndex] * payloadElementCount + sourceOffsets[fragmentIndex];
+      int64_t lane = sourceElementOffset / payloadElementCount;
+      if (lane < 0 || lane >= static_cast<int64_t>(laneCount))
+        return failure();
+      SmallVector<int64_t, 4> fragmentOffsets;
+      SmallVector<int64_t, 4> fragmentSizes;
+      for (int64_t dim = 0; dim < rank; ++dim) {
+        int64_t flatIndex = fragmentIndex * rank + dim;
+        if (flatStrides[flatIndex] != 1)
+          return failure();
+        fragmentOffsets.push_back(flatOffsets[flatIndex]);
+        fragmentSizes.push_back(flatSizes[flatIndex]);
+      }
+
+      if (failed(forEachContiguousDestinationChunk(
+            hostResultType.getShape(),
+            fragmentOffsets,
+            fragmentSizes,
+            [&](ArrayRef<int64_t> chunkOffsets, int64_t relativeSourceOffset, int64_t chunkElements) -> LogicalResult {
+              int64_t hostElementOffset = 0;
+              for (auto [dim, offset] : llvm::enumerate(chunkOffsets))
+                hostElementOffset += offset * hostStrides[dim];
+
+              FragmentAssemblyCopy copy;
+              copy.source = result;
+              copy.sourceType = packedResultType;
+              copy.hostTargetIndex = returnIndex;
+              copy.lane = lane;
+              copy.sourceByteOffset =
+                ((sourceElementOffset % payloadElementCount) + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+              copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
+              copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
+              copies.push_back(copy);
+              return success();
+            })))
+        return failure();
+    }
+  }
+
+  return copies;
+}
+
 static Value createScaledIndexValue(IRRewriter& rewriter, Location loc, Value base, int64_t scale) {
   if (scale == 1)
     return base;
@@ -63,25 +230,32 @@ static Value createScaledIndexValue(IRRewriter& rewriter, Location loc, Value ba
 }
 
 static Value createHostTargetOffset(IRRewriter& rewriter,
-                                    tensor::ParallelInsertSliceOp insertSlice,
+                                    Location loc,
                                     ShapedType destinationType,
+                                    ArrayRef<OpFoldResult> mixedOffsets,
+                                    ArrayRef<int64_t> additionalOffsets,
                                     IRMapping& mapper) {
   int64_t elementBytes = static_cast<int64_t>(getElementTypeSizeInBytes(destinationType.getElementType()));
   SmallVector<int64_t> strides = computeRowMajorStrides(destinationType.getShape());
 
   Value totalOffset;
-  Location loc = insertSlice.getLoc();
-  for (auto [dim, offset] : llvm::enumerate(insertSlice.getMixedOffsets())) {
+  for (auto [dim, offset] : llvm::enumerate(mixedOffsets)) {
     int64_t scale = strides[dim] * elementBytes;
     Value scaledOffset;
     if (auto attr = dyn_cast<Attribute>(offset)) {
       auto intAttr = dyn_cast<IntegerAttr>(attr);
       assert(intAttr && "expected integer offset attribute");
-      scaledOffset =
-        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), intAttr.getInt() * scale);
-    }
-    else {
+      scaledOffset = getOrCreateIndexConstant(rewriter,
+                                              rewriter.getInsertionBlock()->getParentOp(),
+                                              (intAttr.getInt() + additionalOffsets[dim]) * scale);
+    } else {
       scaledOffset = createScaledIndexValue(rewriter, loc, mapper.lookupOrDefault(cast<Value>(offset)), scale);
+      if (additionalOffsets[dim] != 0) {
+        Value staticOffset = getOrCreateIndexConstant(rewriter,
+                                                      rewriter.getInsertionBlock()->getParentOp(),
+                                                      additionalOffsets[dim] * scale);
+        scaledOffset = arith::AddIOp::create(rewriter, loc, scaledOffset, staticOffset).getResult();
+      }
     }
 
     totalOffset =
@@ -93,9 +267,22 @@ static Value createHostTargetOffset(IRRewriter& rewriter,
   return totalOffset;
 }
 
+static Value createHostTargetOffset(IRRewriter& rewriter,
+                                    tensor::ParallelInsertSliceOp insertSlice,
+                                    ShapedType destinationType,
+                                    IRMapping& mapper) {
+  SmallVector<int64_t> zeroOffsets(destinationType.getRank(), 0);
+  return createHostTargetOffset(rewriter,
+                                insertSlice.getLoc(),
+                                destinationType,
+                                insertSlice.getMixedOffsets(),
+                                zeroOffsets,
+                                mapper);
+}
+
 } // namespace
 
-LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp,
+LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatScheduledComputeBatch computeBatchOp,
                                                             IRRewriter& rewriter) {
   Location loc = computeBatchOp.getLoc();
   Block& oldBlock = computeBatchOp.getBody().front();
@@ -110,7 +297,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
       "resultful compute_batch lowering currently requires a spat.in_parallel terminator");
   }
 
-  auto coreIds = getPimCoreIdsForBatchOp(computeBatchOp, coreId);
+  auto coreIds = getRequiredScheduledBatchCoreIds(computeBatchOp, "spatial compute_batch core id");
   if (failed(coreIds))
     return failure();
   SmallVector<Value> batchWeights(computeBatchOp.getWeights().begin(), computeBatchOp.getWeights().end());
@@ -130,14 +317,32 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
   coreBatchOp->setAttr(onnx_mlir::kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(*coreIds));
 
   SmallVector<unsigned> returnOperandIndices;
+  SmallVector<SmallVector<FragmentAssemblyCopyRun, 1>, 4> fragmentAssemblyRunsByResult;
   if (computeBatchOp.getNumResults() != 0) {
-    returnOperandIndices.resize(computeBatchOp.getNumResults());
+    returnOperandIndices.resize(computeBatchOp.getNumResults(), std::numeric_limits<unsigned>::max());
+    fragmentAssemblyRunsByResult.resize(computeBatchOp.getNumResults());
     for (auto [resultIndex, result] : llvm::enumerate(computeBatchOp.getResults())) {
+      if (result.use_empty())
+        continue;
       FailureOr<unsigned> returnOperandIndex = getDirectReturnOperandIndex(cast<OpResult>(result));
-      if (failed(returnOperandIndex))
+      if (succeeded(returnOperandIndex)) {
+        returnOperandIndices[resultIndex] = *returnOperandIndex;
+        continue;
+      }
+
+      auto resultType = dyn_cast<RankedTensorType>(result.getType());
+      if (!resultType || !resultType.hasStaticShape())
         return computeBatchOp.emitOpError(
-          "resultful compute_batch lowering currently requires each result to be used directly by func.return");
-      returnOperandIndices[resultIndex] = *returnOperandIndex;
+          "resultful compute_batch publication lowering requires static ranked tensor results");
+      FailureOr<SmallVector<FragmentAssemblyCopy, 8>> fragmentAssemblyCopies =
+        collectTopLevelFragmentAssemblyCopies(cast<OpResult>(result), resultType, computeBatchOp.getLaneCount());
+      if (failed(fragmentAssemblyCopies))
+        return computeBatchOp.emitOpError("failed to collect top-level fragment assembly copies for compute_batch result");
+      FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>> fragmentAssemblyRuns =
+        groupFragmentAssemblyCopyRuns(*fragmentAssemblyCopies, computeBatchOp.getLaneCount());
+      if (failed(fragmentAssemblyRuns))
+        return computeBatchOp.emitOpError("failed to group top-level fragment assembly copies into regular runs");
+      fragmentAssemblyRunsByResult[resultIndex].assign(fragmentAssemblyRuns->begin(), fragmentAssemblyRuns->end());
     }
   }
 
@@ -195,6 +400,23 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
     if (isa<spatial::SpatYieldOp>(op))
       continue;
 
+    // Cloning a region-bearing operation may leave the rewriter inside that
+    // region. Every old-block operation is lowered at the core-batch body
+    // boundary.
+    rewriter.setInsertionPointToEnd(newBlock);
+
+    if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+      std::optional<StringRef> modeAttr = blueprint.getMode();
+      if (modeAttr && *modeAttr == "fragment_assembly") {
+        for (Operation* user : blueprint.getOutput().getUsers()) {
+          if (!isa<tensor::ParallelInsertSliceOp>(user))
+            return blueprint.emitOpError(
+              "fragment assembly blueprint lowering expects only tensor.parallel_insert_slice users");
+        }
+        continue;
+      }
+    }
+
     if (auto parallelOp = dyn_cast<spatial::SpatInParallelOp>(op)) {
       auto firstOutputArg = computeBatchOp.getOutputArgument(0);
       if (!firstOutputArg)
@@ -211,10 +433,75 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
         unsigned resultIndex = outputArg.getArgNumber() - firstOutputArg->getArgNumber();
         if (resultIndex >= returnOperandIndices.size())
           return insertSlice.emitOpError("result index out of range while lowering host batch output");
+        bool hasDirectReturn = returnOperandIndices[resultIndex] != std::numeric_limits<unsigned>::max();
+        bool hasFragmentAssembly = resultIndex < fragmentAssemblyRunsByResult.size()
+                                   && !fragmentAssemblyRunsByResult[resultIndex].empty();
+        if (!hasDirectReturn && !hasFragmentAssembly)
+          continue;
 
         Value mappedSource = mapper.lookup(insertSlice.getSource());
+
+        if (hasFragmentAssembly) {
+          BlockArgument laneArg = coreBatchOp.getLaneArgument();
+          auto mappedSourceType = dyn_cast<ShapedType>(mappedSource.getType());
+          if (!mappedSourceType || !mappedSourceType.hasStaticShape())
+            return insertSlice.emitOpError("fragment assembly batch lowering requires a static ranked lane-local source");
+          DenseMap<unsigned, Value> updatedOutputs;
+          for (const FragmentAssemblyCopyRun& run : fragmentAssemblyRunsByResult[resultIndex]) {
+            Value outputTensor = updatedOutputs.lookup(run.hostTargetIndex);
+            if (!outputTensor)
+              outputTensor = outputTensors[run.hostTargetIndex](rewriter, insertSlice.getLoc());
+            FragmentAssemblyCopyRun mappedRun = run;
+            mappedRun.source = mappedSource;
+            FailureOr<Value> updated =
+              emitFragmentAssemblyCopyRuns(rewriter,
+                                           insertSlice.getLoc(),
+                                           ArrayRef<FragmentAssemblyCopyRun> {mappedRun},
+                                           outputTensor,
+                                           coreBatchOp.getOperation(),
+                                           laneArg);
+            if (failed(updated))
+              return failure();
+            updatedOutputs[run.hostTargetIndex] = *updated;
+          }
+          continue;
+        }
+
         Value hostTarget = getOrCreateHostOutputTensor(resultIndex, insertSlice.getLoc());
         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") {
+            FailureOr<SmallVector<FragmentAssemblyCopy, 8>> fragmentAssemblyCopies =
+              collectFragmentAssemblyCopiesFromBlueprint(blueprint, mapper, /*lane=*/0, /*hostTargetIndex=*/0);
+            if (failed(fragmentAssemblyCopies))
+              return failure();
+            FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>> fragmentAssemblyRuns =
+              groupFragmentAssemblyCopyRuns(*fragmentAssemblyCopies, /*laneCount=*/1);
+            if (failed(fragmentAssemblyRuns))
+              return failure();
+            SmallVector<int64_t> zeroOffsets(hostTargetType.getRank(), 0);
+            Value baseHostOffset = createHostTargetOffset(rewriter,
+                                                         blueprint.getLoc(),
+                                                         hostTargetType,
+                                                         insertSlice.getMixedOffsets(),
+                                                         zeroOffsets,
+                                                         mapper);
+            FailureOr<Value> updatedHostTarget = emitFragmentAssemblyCopyRuns(rewriter,
+                                                                              blueprint.getLoc(),
+                                                                              *fragmentAssemblyRuns,
+                                                                              hostTarget,
+                                                                              coreBatchOp.getOperation(),
+                                                                              std::nullopt,
+                                                                              baseHostOffset);
+            if (failed(updatedHostTarget))
+              return failure();
+            hostOutputTensors[resultIndex] = *updatedHostTarget;
+            continue;
+          }
+        }
+
         Value hostTargetOffset = createHostTargetOffset(rewriter, insertSlice, hostTargetType, mapper);
         Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
         auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), mappedSource);
@@ -264,9 +551,15 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
       Operation* definingOp = operand.getDefiningOp();
       if (definingOp && definingOp->getBlock() == &oldBlock)
         continue;
+      if (definingOp && definingOp->hasTrait<OpTrait::ConstantLike>())
+        continue;
 
-      return computeBatchOp.emitOpError(
-        "expected external tensor communication to be materialized in Spatial before batch lowering");
+      InFlightDiagnostic diagnostic =
+        computeBatchOp.emitOpError("expected external tensor communication to be materialized in Spatial before batch lowering");
+      diagnostic << " while cloning nested op '" << op.getName() << "' tensor operand #" << operandIndex;
+      if (definingOp)
+        diagnostic << " from external producer '" << definingOp->getName() << "'";
+      return diagnostic;
     }
 
     Operation* cloned = rewriter.clone(op, mapper);

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,11 +1,19 @@
 #include "mlir/IR/ValueRange.h"
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
+
 #include "llvm/ADT/STLExtras.h"
 
 #include <cassert>
+#include <limits>
 
 #include "Common.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace llvm;
 using namespace mlir;
@@ -72,4 +80,488 @@ mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Op
     rewriter, operation->getLoc(), resultShapedType.getShape(), resultShapedType.getElementType());
 }
 
+LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprint,
+                                               int64_t resultRank,
+                                               size_t operandCount,
+                                               ArrayRef<int64_t> operandIndices,
+                                               ArrayRef<int64_t> sourceOffsets,
+                                               ArrayRef<int64_t> flatOffsets,
+                                               ArrayRef<int64_t> flatSizes,
+                                               ArrayRef<int64_t> flatStrides) {
+  if (operandIndices.size() != sourceOffsets.size())
+    return blueprint.emitOpError("fragment assembly operand index and source offset counts must match");
+  if (flatOffsets.size() != flatSizes.size())
+    return blueprint.emitOpError("fragment assembly offset and size arrays must have matching lengths");
+  if (flatStrides.size() != flatOffsets.size())
+    return blueprint.emitOpError("fragment assembly stride and offset arrays must have matching lengths");
+  if (flatOffsets.size() != operandIndices.size() * static_cast<size_t>(resultRank))
+    return blueprint.emitOpError("fragment assembly metadata must provide one rank-sized offset/size/stride tuple per fragment");
+
+  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(operandIndices)) {
+    if (operandIndex < 0 || operandIndex >= static_cast<int64_t>(operandCount))
+      return blueprint.emitOpError("fragment assembly operand index is out of range");
+    if (sourceOffsets[fragmentIndex] < 0)
+      return blueprint.emitOpError("fragment assembly source offsets must be nonnegative");
+  }
+
+  return success();
+}
+
+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) {
+    indices[dim] = flatIndex % shape[dim];
+    flatIndex /= shape[dim];
+  }
+  return indices;
+}
+
+FailureOr<SmallVector<int64_t, 4>>
+getStaticSliceOffsetsForElementOffset(Operation* anchor,
+                                      ShapedType sourceType,
+                                      ArrayRef<int64_t> fragmentShape,
+                                      int64_t sourceElementOffset,
+                                      StringRef fieldName) {
+  if (!sourceType.hasStaticShape())
+    return (anchor->emitOpError() << fieldName << " requires a static source shape"), failure();
+  if (sourceElementOffset < 0)
+    return (anchor->emitOpError() << fieldName << " requires a nonnegative source element offset"), failure();
+  if (sourceType.getRank() != static_cast<int64_t>(fragmentShape.size()))
+    return (anchor->emitOpError() << fieldName << " requires fragment rank to match source rank"), failure();
+
+  int64_t sourceElementCount = sourceType.getNumElements();
+  int64_t fragmentElementCount = 1;
+  for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) {
+    if (fragmentShape[dim] < 0)
+      return (anchor->emitOpError() << fieldName << " requires nonnegative fragment sizes"), failure();
+    fragmentElementCount *= fragmentShape[dim];
+  }
+  if (sourceElementOffset + fragmentElementCount > sourceElementCount)
+    return (anchor->emitOpError() << fieldName << " exceeds the source tensor bounds"), failure();
+
+  SmallVector<int64_t, 4> sliceOffsets = expandFlatElementIndex(sourceElementOffset, sourceType.getShape());
+  for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) {
+    if (sliceOffsets[dim] + fragmentShape[dim] > sourceType.getDimSize(dim))
+      return (anchor->emitOpError() << fieldName << " does not describe a valid unit-stride slice"), failure();
+  }
+  return sliceOffsets;
+}
+
+LogicalResult
+forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
+                                  ArrayRef<int64_t> baseOffsets,
+                                  ArrayRef<int64_t> sizes,
+                                  llvm::function_ref<LogicalResult(ArrayRef<int64_t>, int64_t, int64_t)> callback) {
+  int64_t rank = static_cast<int64_t>(sizes.size());
+  int64_t suffixStart = rank - 1;
+  while (suffixStart > 0 && sizes[suffixStart] == destShape[suffixStart])
+    --suffixStart;
+  if (sizes[suffixStart] == destShape[suffixStart] && suffixStart == 0)
+    suffixStart = 0;
+  else
+    ++suffixStart;
+
+  int64_t chunkElements = 1;
+  for (int64_t dim = suffixStart; dim < rank; ++dim)
+    chunkElements *= sizes[dim];
+
+  SmallVector<int64_t, 4> prefixExtents(sizes.begin(), sizes.begin() + suffixStart);
+  SmallVector<int64_t, 4> current(prefixExtents.size(), 0);
+  int64_t sourceChunkOrdinal = 0;
+
+  auto visit = [&](auto&& visit, int64_t dim) -> LogicalResult {
+    if (dim == static_cast<int64_t>(prefixExtents.size())) {
+      SmallVector<int64_t, 4> chunkOffsets(baseOffsets.begin(), baseOffsets.end());
+      for (int64_t prefixDim = 0; prefixDim < static_cast<int64_t>(current.size()); ++prefixDim)
+        chunkOffsets[prefixDim] += current[prefixDim];
+      if (failed(callback(chunkOffsets, sourceChunkOrdinal * chunkElements, chunkElements)))
+        return failure();
+      ++sourceChunkOrdinal;
+      return success();
+    }
+
+    for (int64_t index = 0; index < prefixExtents[dim]; ++index) {
+      current[dim] = index;
+      if (failed(visit(visit, dim + 1)))
+        return failure();
+    }
+    return success();
+  };
+
+  if (prefixExtents.empty())
+    return callback(baseOffsets, 0, chunkElements);
+  return visit(visit, 0);
+}
+
+static mlir::Value
+createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index,
+                    int64_t stepBytes, Operation *constantAnchor) {
+  if (stepBytes == 0)
+    return start;
+  return createOrFoldAffineApply(
+    builder, loc, builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1) * stepBytes,
+    ValueRange {start, index}, constantAnchor);
+}
+
+static mlir::Value createIndexedOffset(OpBuilder& builder,
+                                       Location loc,
+                                       mlir::Value indexArg,
+                                       ArrayRef<int64_t> values,
+                                       Operation *constantAnchor) {
+  assert(!values.empty() && "expected lane-indexed values");
+  if (llvm::all_of(values.drop_front(), [&](int64_t value) { return value == values.front(); }))
+    return getOrCreateIndexConstant(builder, constantAnchor, values.front());
+
+  if (values.size() >= 2) {
+    int64_t step = values[1] - values[0];
+    bool arithmetic = llvm::all_of(llvm::seq<size_t>(2, values.size()), [&](size_t index) {
+      return values[index] == values.front() + static_cast<int64_t>(index) * step;
+    });
+    if (arithmetic) {
+      return createOrFoldAffineApply(
+        builder, loc, builder.getAffineDimExpr(0) * step + values.front(),
+        ValueRange {indexArg}, constantAnchor);
+    }
+  }
+
+  RankedTensorType tableType = RankedTensorType::get(
+    {static_cast<int64_t>(values.size())}, builder.getI64Type());
+  DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
+  mlir::Value table = getOrCreateConstant(builder, constantAnchor, tableAttr, tableType);
+  mlir::Value selected = tensor::ExtractOp::create(builder, loc, table, ValueRange {indexArg});
+  return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), selected).getResult();
+}
+
+struct FragmentAssemblyCopyRunFamily {
+  FragmentAssemblyCopyRun prototype;
+  SmallVector<int64_t, 8> sourceRunStartDeltas;
+  SmallVector<int64_t, 8> hostRunStartDeltas;
+};
+
+static bool computeUniformRunStartDelta(ArrayRef<int64_t> prototypeStarts,
+                                        ArrayRef<int64_t> runStarts,
+                                        int64_t& delta) {
+  if (prototypeStarts.size() != runStarts.size() || prototypeStarts.empty())
+    return false;
+
+  delta = runStarts.front() - prototypeStarts.front();
+  return llvm::all_of(llvm::zip_equal(prototypeStarts, runStarts), [&](auto pair) {
+    auto [prototypeStart, runStart] = pair;
+    return runStart - prototypeStart == delta;
+  });
+}
+
+static bool canMergeFragmentAssemblyCopyRunIntoFamily(const FragmentAssemblyCopyRunFamily& family,
+                                                      const FragmentAssemblyCopyRun& run,
+                                                      int64_t& sourceRunStartDelta,
+                                                      int64_t& hostRunStartDelta) {
+  const FragmentAssemblyCopyRun& prototype = family.prototype;
+  if (prototype.source != run.source || prototype.sourceType != run.sourceType
+      || prototype.hostTargetIndex != run.hostTargetIndex || prototype.count != run.count
+      || prototype.sourceStepBytes != run.sourceStepBytes || prototype.hostStepBytes != run.hostStepBytes
+      || prototype.byteSize != run.byteSize)
+    return false;
+
+  if (!computeUniformRunStartDelta(prototype.sourceStartBytesByLane, run.sourceStartBytesByLane, sourceRunStartDelta))
+    return false;
+  return computeUniformRunStartDelta(prototype.hostStartBytesByLane, run.hostStartBytesByLane, hostRunStartDelta);
+}
+
+static SmallVector<FragmentAssemblyCopyRunFamily, 8>
+groupFragmentAssemblyCopyRunFamilies(ArrayRef<FragmentAssemblyCopyRun> runs) {
+  auto compareRunStarts = [](ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs) {
+    return std::lexicographical_compare(lhs.begin(), lhs.end(), rhs.begin(), rhs.end());
+  };
+
+  SmallVector<FragmentAssemblyCopyRun, 8> sortedRuns(runs.begin(), runs.end());
+  llvm::sort(sortedRuns, [&](const FragmentAssemblyCopyRun& lhs, const FragmentAssemblyCopyRun& rhs) {
+    if (lhs.hostTargetIndex != rhs.hostTargetIndex)
+      return lhs.hostTargetIndex < rhs.hostTargetIndex;
+    if (lhs.source != rhs.source)
+      return lhs.source.getAsOpaquePointer() < rhs.source.getAsOpaquePointer();
+    if (lhs.byteSize != rhs.byteSize)
+      return lhs.byteSize < rhs.byteSize;
+    if (lhs.count != rhs.count)
+      return lhs.count < rhs.count;
+    if (lhs.sourceStepBytes != rhs.sourceStepBytes)
+      return lhs.sourceStepBytes < rhs.sourceStepBytes;
+    if (lhs.hostStepBytes != rhs.hostStepBytes)
+      return lhs.hostStepBytes < rhs.hostStepBytes;
+    if (compareRunStarts(lhs.sourceStartBytesByLane, rhs.sourceStartBytesByLane))
+      return true;
+    if (compareRunStarts(rhs.sourceStartBytesByLane, lhs.sourceStartBytesByLane))
+      return false;
+    return compareRunStarts(lhs.hostStartBytesByLane, rhs.hostStartBytesByLane);
+  });
+
+  SmallVector<FragmentAssemblyCopyRunFamily, 8> families;
+  for (const FragmentAssemblyCopyRun& run : sortedRuns) {
+    int64_t sourceRunStartDelta = 0;
+    int64_t hostRunStartDelta = 0;
+    if (!families.empty()
+        && canMergeFragmentAssemblyCopyRunIntoFamily(
+          families.back(), run, sourceRunStartDelta, hostRunStartDelta)) {
+      families.back().sourceRunStartDeltas.push_back(sourceRunStartDelta);
+      families.back().hostRunStartDeltas.push_back(hostRunStartDelta);
+      continue;
+    }
+
+    FragmentAssemblyCopyRunFamily family;
+    family.prototype = run;
+    family.sourceRunStartDeltas.push_back(0);
+    family.hostRunStartDeltas.push_back(0);
+    families.push_back(std::move(family));
+  }
+
+  return families;
+}
+
+FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>>
+groupFragmentAssemblyCopyRuns(ArrayRef<FragmentAssemblyCopy> copies, uint32_t laneCount) {
+  if (laneCount == 0)
+    return failure();
+
+  struct LaneLocalCopyRun {
+    FragmentAssemblyCopyRun run;
+    int64_t lane = 0;
+  };
+
+  SmallVector<FragmentAssemblyCopy, 8> sortedCopies(copies.begin(), copies.end());
+  llvm::sort(sortedCopies, [](const FragmentAssemblyCopy& lhs, const FragmentAssemblyCopy& rhs) {
+    if (lhs.hostTargetIndex != rhs.hostTargetIndex)
+      return lhs.hostTargetIndex < rhs.hostTargetIndex;
+    if (lhs.source != rhs.source)
+      return lhs.source.getAsOpaquePointer() < rhs.source.getAsOpaquePointer();
+    if (lhs.lane != rhs.lane)
+      return lhs.lane < rhs.lane;
+    if (lhs.byteSize != rhs.byteSize)
+      return lhs.byteSize < rhs.byteSize;
+    if (lhs.sourceByteOffset != rhs.sourceByteOffset)
+      return lhs.sourceByteOffset < rhs.sourceByteOffset;
+    return lhs.hostByteOffset < rhs.hostByteOffset;
+  });
+
+  SmallVector<LaneLocalCopyRun, 8> laneRuns;
+  for (const FragmentAssemblyCopy& copy : sortedCopies) {
+    if (copy.lane < 0 || copy.lane >= static_cast<int64_t>(laneCount))
+      return failure();
+
+    if (!laneRuns.empty()) {
+      LaneLocalCopyRun& laneRun = laneRuns.back();
+      FragmentAssemblyCopyRun& run = laneRun.run;
+      if (run.source == copy.source && run.sourceType == copy.sourceType
+          && run.hostTargetIndex == copy.hostTargetIndex && laneRun.lane == copy.lane && run.byteSize == copy.byteSize
+          && run.sourceStartBytesByLane.size() == 1 && run.hostStartBytesByLane.size() == 1) {
+        int64_t previousSourceOffset = run.sourceStartBytesByLane.front() + (run.count - 1) * run.sourceStepBytes;
+        int64_t previousHostOffset = run.hostStartBytesByLane.front() + (run.count - 1) * run.hostStepBytes;
+        int64_t sourceDelta = copy.sourceByteOffset - previousSourceOffset;
+        int64_t hostDelta = copy.hostByteOffset - previousHostOffset;
+        if (run.count == 1) {
+          run.sourceStepBytes = sourceDelta;
+          run.hostStepBytes = hostDelta;
+          ++run.count;
+          continue;
+        }
+        if (run.sourceStepBytes == sourceDelta && run.hostStepBytes == hostDelta) {
+          ++run.count;
+          continue;
+        }
+      }
+    }
+
+    LaneLocalCopyRun laneRun;
+    laneRun.run.source = copy.source;
+    laneRun.run.sourceType = copy.sourceType;
+    laneRun.run.hostTargetIndex = copy.hostTargetIndex;
+    laneRun.run.count = 1;
+    laneRun.run.byteSize = copy.byteSize;
+    laneRun.run.sourceStartBytesByLane.push_back(copy.sourceByteOffset);
+    laneRun.run.hostStartBytesByLane.push_back(copy.hostByteOffset);
+    laneRun.lane = copy.lane;
+    laneRuns.push_back(std::move(laneRun));
+  }
+
+  SmallVector<FragmentAssemblyCopyRun, 8> mergedRuns;
+  for (const LaneLocalCopyRun& laneRun : laneRuns) {
+    size_t laneIndex = static_cast<size_t>(laneRun.lane);
+    auto mergedIt = llvm::find_if(mergedRuns, [&](const FragmentAssemblyCopyRun& run) {
+      return run.source == laneRun.run.source && run.sourceType == laneRun.run.sourceType
+             && run.hostTargetIndex == laneRun.run.hostTargetIndex && run.count == laneRun.run.count
+             && run.byteSize == laneRun.run.byteSize && run.sourceStepBytes == laneRun.run.sourceStepBytes
+             && run.hostStepBytes == laneRun.run.hostStepBytes && laneIndex < run.sourceStartBytesByLane.size()
+             && run.sourceStartBytesByLane[laneIndex] == std::numeric_limits<int64_t>::min();
+    });
+
+    if (mergedIt == mergedRuns.end()) {
+      FragmentAssemblyCopyRun merged = laneRun.run;
+      merged.sourceStartBytesByLane.assign(laneCount, std::numeric_limits<int64_t>::min());
+      merged.hostStartBytesByLane.assign(laneCount, std::numeric_limits<int64_t>::min());
+      merged.sourceStartBytesByLane[laneIndex] = laneRun.run.sourceStartBytesByLane.front();
+      merged.hostStartBytesByLane[laneIndex] = laneRun.run.hostStartBytesByLane.front();
+      mergedRuns.push_back(std::move(merged));
+      continue;
+    }
+
+    mergedIt->sourceStartBytesByLane[laneIndex] = laneRun.run.sourceStartBytesByLane.front();
+    mergedIt->hostStartBytesByLane[laneIndex] = laneRun.run.hostStartBytesByLane.front();
+  }
+
+  for (const FragmentAssemblyCopyRun& run : mergedRuns) {
+    if (llvm::any_of(run.sourceStartBytesByLane,
+                     [](int64_t value) { return value == std::numeric_limits<int64_t>::min(); }))
+      return failure();
+    if (llvm::any_of(run.hostStartBytesByLane,
+                     [](int64_t value) { return value == std::numeric_limits<int64_t>::min(); }))
+      return failure();
+  }
+
+  return mergedRuns;
+}
+
+static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
+                                                          Location loc,
+                                                          const FragmentAssemblyCopyRun& run,
+                                                          mlir::Value hostTarget,
+                                                          Operation* anchor,
+                                                          std::optional<mlir::Value> laneArg,
+                                                          mlir::Value baseHostOffset,
+                                                          mlir::Value sourceRunStartDelta = {},
+                                                          mlir::Value hostRunStartDelta = {}) {
+  auto sizeAttr = pim::getCheckedI32Attr(builder, anchor, run.byteSize, "fragment assembly host copy byte size");
+  if (failed(sizeAttr))
+    return failure();
+
+  mlir::Value hostStart;
+  mlir::Value sourceStart;
+  if (laneArg) {
+    hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane, anchor);
+    sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane, anchor);
+  } else {
+    hostStart = getOrCreateIndexConstant(builder, anchor, run.hostStartBytesByLane.front());
+    sourceStart = getOrCreateIndexConstant(builder, anchor, run.sourceStartBytesByLane.front());
+  }
+
+  if (hostRunStartDelta)
+    hostStart = arith::AddIOp::create(builder, loc, hostStart, hostRunStartDelta).getResult();
+  if (sourceRunStartDelta)
+    sourceStart = arith::AddIOp::create(builder, loc, sourceStart, sourceRunStartDelta).getResult();
+  if (baseHostOffset)
+    hostStart = arith::AddIOp::create(builder, loc, baseHostOffset, hostStart).getResult();
+
+  if (run.count == 1) {
+    return pim::PimMemCopyDevToHostOp::create(builder,
+                                              loc,
+                                              hostTarget.getType(),
+                                              hostStart,
+                                              sourceStart,
+                                              hostTarget,
+                                              run.source,
+                                              *sizeAttr)
+      .getOutput();
+  }
+
+  mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
+  mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, run.count);
+  mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
+  FailureOr<NormalizedLoopResult> loop = buildNormalizedScfFor(
+    builder,
+    loc,
+    lowerBound,
+    upperBound,
+    step,
+    ValueRange {hostTarget},
+    [&](OpBuilder& loopBuilder,
+        Location bodyLoc,
+        mlir::Value flatIndex,
+        ValueRange iterArgs,
+        SmallVectorImpl<mlir::Value>& yielded) {
+      mlir::Value hostOffset = createSteppedOffset(
+        loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes, anchor);
+      mlir::Value sourceOffset =
+        createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes, anchor);
+      mlir::Value copied =
+        pim::PimMemCopyDevToHostOp::create(loopBuilder,
+                                           bodyLoc,
+                                           iterArgs.front().getType(),
+                                           hostOffset,
+                                           sourceOffset,
+                                           iterArgs.front(),
+                                           run.source,
+                                           *sizeAttr)
+          .getOutput();
+      yielded.push_back(copied);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  return loop->results.front();
+}
+
+static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& builder,
+                                                                Location loc,
+                                                                const FragmentAssemblyCopyRunFamily& family,
+                                                                mlir::Value hostTarget,
+                                                                Operation* anchor,
+                                                                std::optional<mlir::Value> laneArg,
+                                                                mlir::Value baseHostOffset) {
+  if (family.sourceRunStartDeltas.size() == 1)
+    return emitFragmentAssemblyCopyRun(
+      builder, loc, family.prototype, hostTarget, anchor, laneArg, baseHostOffset);
+
+  mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
+  mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, family.sourceRunStartDeltas.size());
+  mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
+  FailureOr<NormalizedLoopResult> outerLoop = buildNormalizedScfFor(
+    builder,
+    loc,
+    lowerBound,
+    upperBound,
+    step,
+    ValueRange {hostTarget},
+    [&](OpBuilder& loopBuilder,
+        Location bodyLoc,
+        mlir::Value runIndex,
+        ValueRange iterArgs,
+        SmallVectorImpl<mlir::Value>& yielded) {
+      mlir::Value sourceRunStartDelta =
+        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas, anchor);
+      mlir::Value hostRunStartDelta =
+        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas, anchor);
+      FailureOr<mlir::Value> copied = emitFragmentAssemblyCopyRun(loopBuilder,
+                                                                  bodyLoc,
+                                                                  family.prototype,
+                                                                  iterArgs.front(),
+                                                                  anchor,
+                                                                  laneArg,
+                                                                  baseHostOffset,
+                                                                  sourceRunStartDelta,
+                                                                  hostRunStartDelta);
+      if (failed(copied))
+        return failure();
+      yielded.push_back(*copied);
+      return success();
+    });
+  if (failed(outerLoop))
+    return failure();
+  return outerLoop->results.front();
+}
+
+FailureOr<mlir::Value> emitFragmentAssemblyCopyRuns(IRRewriter& rewriter,
+                                                    Location loc,
+                                                    ArrayRef<FragmentAssemblyCopyRun> runs,
+                                                    mlir::Value hostTarget,
+                                                    Operation* anchor,
+                                                    std::optional<mlir::Value> laneArg,
+                                                    mlir::Value baseHostOffset) {
+  for (const FragmentAssemblyCopyRunFamily& family : groupFragmentAssemblyCopyRunFamilies(runs)) {
+    FailureOr<mlir::Value> updatedHostTarget =
+      emitFragmentAssemblyCopyRunFamily(rewriter, loc, family, hostTarget, anchor, laneArg, baseHostOffset);
+    if (failed(updatedHostTarget))
+      return failure();
+    hostTarget = *updatedHostTarget;
+  }
+
+  return hostTarget;
+}
+
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -1,10 +1,23 @@
 #pragma once
 
+#include <optional>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Value.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Support/LogicalResult.h"
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 
+namespace onnx_mlir::spatial {
+class SpatBlueprintOp;
+}
+
 namespace onnx_mlir {
 
 mlir::FailureOr<mlir::IntegerAttr>
@@ -29,6 +42,62 @@ mlir::SmallVector<mlir::Value> getOpOperandsSortedByUses(mlir::Operation* operat
 
 mlir::Value getBestOutputTensorFromOperandsOrAllocate(mlir::RewriterBase& rewriter, mlir::Operation* operation);
 
+mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlueprintOp blueprint,
+                                                     int64_t resultRank,
+                                                     size_t operandCount,
+                                                     llvm::ArrayRef<int64_t> operandIndices,
+                                                     llvm::ArrayRef<int64_t> sourceOffsets,
+                                                     llvm::ArrayRef<int64_t> flatOffsets,
+                                                     llvm::ArrayRef<int64_t> flatSizes,
+                                                     llvm::ArrayRef<int64_t> flatStrides);
+
+mlir::FailureOr<mlir::SmallVector<int64_t, 4>>
+getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor,
+                                      mlir::ShapedType sourceType,
+                                      llvm::ArrayRef<int64_t> fragmentShape,
+                                      int64_t sourceElementOffset,
+                                      llvm::StringRef fieldName);
+
+mlir::LogicalResult
+forEachContiguousDestinationChunk(llvm::ArrayRef<int64_t> destShape,
+                                  llvm::ArrayRef<int64_t> baseOffsets,
+                                  llvm::ArrayRef<int64_t> sizes,
+                                  llvm::function_ref<mlir::LogicalResult(llvm::ArrayRef<int64_t>, int64_t, int64_t)>
+                                    callback);
+
+struct FragmentAssemblyCopy {
+  mlir::Value source;
+  mlir::RankedTensorType sourceType;
+  unsigned hostTargetIndex = 0;
+  int64_t lane = 0;
+  int64_t sourceByteOffset = 0;
+  int64_t hostByteOffset = 0;
+  int64_t byteSize = 0;
+};
+
+struct FragmentAssemblyCopyRun {
+  mlir::Value source;
+  mlir::RankedTensorType sourceType;
+  unsigned hostTargetIndex = 0;
+  int64_t count = 0;
+  int64_t sourceStepBytes = 0;
+  int64_t hostStepBytes = 0;
+  int64_t byteSize = 0;
+  mlir::SmallVector<int64_t, 8> sourceStartBytesByLane;
+  mlir::SmallVector<int64_t, 8> hostStartBytesByLane;
+};
+
+mlir::FailureOr<mlir::SmallVector<FragmentAssemblyCopyRun, 8>>
+groupFragmentAssemblyCopyRuns(llvm::ArrayRef<FragmentAssemblyCopy> copies, uint32_t laneCount = 1);
+
+mlir::FailureOr<mlir::Value> emitFragmentAssemblyCopyRuns(mlir::IRRewriter& rewriter,
+                                                          mlir::Location loc,
+                                                          llvm::ArrayRef<FragmentAssemblyCopyRun> runs,
+                                                          mlir::Value hostTarget,
+                                                          mlir::Operation* anchor,
+                                                          std::optional<mlir::Value> laneArg = std::nullopt,
+                                                          mlir::Value baseHostOffset = {});
+
 inline mlir::tensor::EmptyOp
 createEmptyTensorFromShaped(mlir::IRRewriter& rewriter, mlir::Location loc, mlir::ShapedType shapedType) {
   return mlir::tensor::EmptyOp::create(rewriter, loc, shapedType.getShape(), shapedType.getElementType());

src/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.cpp

@@ -17,10 +17,10 @@ std::optional<unsigned> getDirectComputeLikeInputIndex(Operation* owner, unsigne
     return operandNumber - inputBegin;
   };
 
-  if (auto compute = dyn_cast<spatial::SpatCompute>(owner))
+  if (auto compute = dyn_cast<spatial::SpatScheduledCompute>(owner))
     return getInputIndex(owner, compute.getInputs().size());
 
-  if (auto computeBatch = dyn_cast<spatial::SpatComputeBatch>(owner))
+  if (auto computeBatch = dyn_cast<spatial::SpatScheduledComputeBatch>(owner))
     return getInputIndex(owner, computeBatch.getInputs().size());
 
   return std::nullopt;
@@ -32,13 +32,13 @@ void replaceAndEraseDirectComputeLikeInput(PatternRewriter& rewriter,
                                            Value replacement) {
   Block& body = owner->getRegion(0).front();
   BlockArgument bodyArgument;
-  if (auto compute = dyn_cast<spatial::SpatCompute>(owner)) {
+  if (auto compute = dyn_cast<spatial::SpatScheduledCompute>(owner)) {
     auto computeArg = compute.getInputArgument(inputIndex);
     assert(computeArg && "expected compute input block argument");
     bodyArgument = *computeArg;
   }
   else {
-    auto batchArg = cast<spatial::SpatComputeBatch>(owner).getInputArgument(inputIndex);
+    auto batchArg = cast<spatial::SpatScheduledComputeBatch>(owner).getInputArgument(inputIndex);
     assert(batchArg && "expected compute_batch input block argument");
     bodyArgument = *batchArg;
   }
@@ -46,10 +46,10 @@ void replaceAndEraseDirectComputeLikeInput(PatternRewriter& rewriter,
 
   rewriter.startOpModification(owner);
   bodyArgument.replaceAllUsesWith(replacement);
-  if (auto compute = dyn_cast<spatial::SpatCompute>(owner))
+  if (auto compute = dyn_cast<spatial::SpatScheduledCompute>(owner))
     compute.getInputsMutable().erase(inputIndex);
   else
-    cast<spatial::SpatComputeBatch>(owner).getInputsMutable().erase(inputIndex);
+    cast<spatial::SpatScheduledComputeBatch>(owner).getInputsMutable().erase(inputIndex);
   body.eraseArgument(bodyArgIndex);
   rewriter.finalizeOpModification(owner);
 }

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -1,6 +1,7 @@
 #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/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/IRMapping.h"
@@ -8,6 +9,10 @@
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -30,6 +35,90 @@ static bool isChannelUseChainOp(Operation* op) {
              pim::PimTransposeOp>(op);
 }
 
+static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
+                                                       spatial::SpatBlueprintOp blueprint,
+                                                       IRMapping& mapping) {
+  auto resultType = dyn_cast<ShapedType>(blueprint.getOutput().getType());
+  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>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+  std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
+  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr
+      || !fragmentStridesAttr)
+    return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
+
+  ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+  ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+  ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
+  ArrayRef<int64_t> flatStrides = *fragmentStridesAttr;
+  int64_t rank = resultType.getRank();
+
+  SmallVector<Value> fragmentOperands {blueprint.getInput()};
+  llvm::append_range(fragmentOperands, blueprint.getFragments());
+  if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                              rank,
+                                              fragmentOperands.size(),
+                                              operandIndices,
+                                              sourceOffsets,
+                                              flatOffsets,
+                                              flatSizes,
+                                              flatStrides)))
+    return failure();
+
+  SmallVector<int64_t> hostStrides = computeRowMajorStrides(resultType.getShape());
+  SmallVector<FragmentAssemblyCopy, 8> copies;
+  for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+    int64_t operandIndex = operandIndices[fragmentIndex];
+
+    SmallVector<int64_t, 4> fragmentOffsets;
+    SmallVector<int64_t, 4> fragmentSizes;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t flatIndex = fragmentIndex * rank + dim;
+      if (flatStrides[flatIndex] != 1)
+        return blueprint.emitOpError("fragment assembly lowering only supports unit strides");
+      fragmentOffsets.push_back(flatOffsets[flatIndex]);
+      fragmentSizes.push_back(flatSizes[flatIndex]);
+    }
+
+    Value source = mapping.lookupOrDefault(fragmentOperands[operandIndex]);
+    auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+    if (!sourceType || !sourceType.hasStaticShape())
+      return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor operands");
+    size_t elementSize = getElementTypeSizeInBytes(sourceType.getElementType());
+    if (failed(forEachContiguousDestinationChunk(
+          resultType.getShape(),
+          fragmentOffsets,
+          fragmentSizes,
+          [&](ArrayRef<int64_t> chunkOffsets, int64_t relativeSourceOffset, int64_t chunkElements) -> LogicalResult {
+            int64_t hostElementOffset = 0;
+            for (auto [dim, offset] : llvm::enumerate(chunkOffsets))
+              hostElementOffset += offset * hostStrides[dim];
+
+            FragmentAssemblyCopy copy;
+            copy.source = source;
+            copy.sourceType = sourceType;
+            copy.sourceByteOffset =
+              (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+            copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
+            copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
+            copies.push_back(copy);
+            return success();
+          })))
+      return failure();
+  }
+
+  Value currentOutput = createEmptyTensorFromShaped(rewriter, blueprint.getLoc(), resultType);
+  FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>> runs = groupFragmentAssemblyCopyRuns(copies);
+  if (failed(runs))
+    return failure();
+  return emitFragmentAssemblyCopyRuns(
+    rewriter, blueprint.getLoc(), *runs, currentOutput, blueprint.getOperation());
+}
+
 static void
 cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewriter, OperationFolder& constantFolder) {
   for (Value operand : op->getOperands()) {
@@ -55,18 +144,7 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
   }
 }
 
-static FailureOr<int32_t> getPimCoreIdForComputeOp(spatial::SpatCompute computeOp, size_t& fallbackCoreId) {
-  if (auto spatialCoreIdAttr = computeOp->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
-    return pim::checkedI32(spatialCoreIdAttr.getInt(), computeOp, "spatial compute core id");
-  auto checkedCoreId =
-    pim::checkedI32(static_cast<uint64_t>(fallbackCoreId), computeOp, "fallback spatial compute core id");
-  if (failed(checkedCoreId))
-    return failure();
-  ++fallbackCoreId;
-  return *checkedCoreId;
-}
-
-static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
+static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute computeOp,
                                                SmallVectorImpl<Operation*>& helperChain,
                                                bool requireReturnUse = true) {
   if (computeOp.getInputs().size() != 1 || computeOp.getNumResults() != 1)
@@ -104,16 +182,79 @@ static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
   return success();
 }
 
-static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute computeOp,
+static bool isHostMaterializableHelperOp(Operation* op) {
+  if (isa<spatial::SpatYieldOp>(op))
+    return true;
+  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 isShapingOnlyOp(op) || isPureIndexComputationOp(op);
+}
+
+static FailureOr<DenseMap<Value, Attribute>>
+analyzeHostMaterializableHelper(spatial::SpatScheduledCompute computeOp) {
+  DenseMap<Value, Attribute> folded;
+  for (auto [weightIndex, weight] : llvm::enumerate(computeOp.getWeights())) {
+    auto argument = computeOp.getWeightArgument(weightIndex);
+    if (!argument)
+      return failure();
+    Attribute constant;
+    if (matchPattern(weight, m_Constant(&constant)))
+      folded[*argument] = constant;
+  }
+  Block& block = computeOp.getBody().front();
+  for (Operation& op : block) {
+    if (!isHostMaterializableHelperOp(&op))
+      return failure();
+    if (isa<spatial::SpatYieldOp, spatial::SpatBlueprintOp>(op)
+        || (isShapingOnlyOp(&op) && !isPureIndexComputationOp(&op)))
+      continue;
+    if (isa<arith::ConstantOp>(op) || op.hasTrait<OpTrait::ConstantLike>()) {
+      for (Value result : op.getResults()) {
+        Attribute constant;
+        if (!matchPattern(result, m_Constant(&constant)))
+          return failure();
+        folded[result] = constant;
+      }
+      continue;
+    }
+    if (!isPureIndexComputationOp(&op) || op.getNumRegions() != 0)
+      return failure();
+    SmallVector<Attribute> operands;
+    for (Value operand : op.getOperands()) {
+      auto it = folded.find(operand);
+      if (it == folded.end())
+        return failure();
+      operands.push_back(it->second);
+    }
+    SmallVector<OpFoldResult> results;
+    if (failed(op.fold(operands, results))
+        || results.size() != op.getNumResults())
+      return failure();
+    for (auto [result, foldResult] : llvm::zip(op.getResults(), results)) {
+      auto attribute = dyn_cast<Attribute>(foldResult);
+      if (!attribute)
+        return failure();
+      folded[result] = attribute;
+    }
+  }
+  return folded;
+}
+
+static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatScheduledCompute computeOp,
                                                            IRRewriter& rewriter,
                                                            OperationFolder& constantFolder) {
   if (!computeOp.getInputs().empty() || computeOp.getNumResults() != 1)
     return false;
+  if (computeOp.getResult(0).use_empty())
+    return false;
   if (!llvm::all_of(computeOp.getResult(0).getUsers(), [](Operation* user) {
-        return isa<spatial::SpatCompute, spatial::SpatComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(user);
+        return isa<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(user);
       }))
     return false;
-
   Block& block = computeOp.getBody().front();
   if (block.getNumArguments() != computeOp.getWeights().size())
     return false;
@@ -121,6 +262,9 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute
   auto yieldOp = dyn_cast<spatial::SpatYieldOp>(block.getTerminator());
   if (!yieldOp || yieldOp.getNumOperands() != 1)
     return false;
+  auto folded = analyzeHostMaterializableHelper(computeOp);
+  if (failed(folded))
+    return false;
 
   rewriter.setInsertionPoint(computeOp);
   IRMapping mapping;
@@ -131,6 +275,31 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute
     mapping.map(*weightArg, weight);
   }
   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") {
+        auto lowered = lowerFragmentAssemblyBlueprint(rewriter, blueprint, mapping);
+        if (failed(lowered))
+          return false;
+        mapping.map(blueprint.getOutput(), *lowered);
+        continue;
+      }
+    }
+
+    if (isa<arith::ConstantOp>(op) || op.hasTrait<OpTrait::ConstantLike>()
+        || isPureIndexComputationOp(&op)) {
+      for (Value result : op.getResults()) {
+        auto it = folded->find(result);
+        if (it == folded->end())
+          return false;
+        mapping.map(
+          result,
+          getOrCreateConstant(constantFolder, computeOp, it->second,
+                              result.getType()));
+      }
+      continue;
+    }
+
     cloneMappedHelperOperands(&op, mapping, rewriter, constantFolder);
     Operation* clonedOp = rewriter.clone(op, mapping);
     for (auto [originalResult, newResult] : llvm::zip(op.getResults(), clonedOp->getResults()))
@@ -145,7 +314,7 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute
 
 } // namespace
 
-LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute computeOp,
+LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCompute computeOp,
                                                        IRRewriter& rewriter,
                                                        OperationFolder& constantFolder) {
   Location loc = computeOp->getLoc();
@@ -214,7 +383,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
   if (!computeOp.getWeights().empty())
     computeWeights.append(computeOp.getWeights().begin(), computeOp.getWeights().end());
   rewriter.setInsertionPointAfter(computeOp);
-  auto checkedCoreId = getPimCoreIdForComputeOp(computeOp, coreId);
+  auto checkedCoreId = getRequiredScheduledCoreId(computeOp, "spatial compute core id");
   if (failed(checkedCoreId))
     return failure();
   auto coreIdAttr = pim::getCheckedI32Attr(rewriter, computeOp, static_cast<int64_t>(*checkedCoreId), "pim core id");

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -1,6 +1,10 @@
+#include "mlir/Transforms/DialectConversion.h"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
 
@@ -11,6 +15,92 @@ namespace raptor {
 
 } // namespace raptor
 
+struct LowerFragmentAssemblyBlueprintPattern
+  : OpConversionPattern<spatial::SpatBlueprintOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBlueprintOp op,
+                                OpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    std::optional<StringRef> modeAttr = op.getMode();
+    if (!modeAttr || *modeAttr != "fragment_assembly")
+      return failure();
+
+    auto resultType = dyn_cast<ShapedType>(op.getOutput().getType());
+    if (!resultType || !resultType.hasStaticShape())
+      return op.emitOpError("fragment assembly lowering requires a static ranked tensor result");
+
+    std::optional<ArrayRef<int64_t>> operandIndicesAttr = op.getFragmentOperandIndices();
+    std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = op.getFragmentSourceOffsets();
+    std::optional<ArrayRef<int64_t>> fragmentStridesAttr = op.getFragmentStrides();
+    if (!operandIndicesAttr || !sourceOffsetsAttr || !fragmentStridesAttr)
+      return op.emitOpError("fragment assembly lowering requires explicit fragment metadata");
+
+    ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+    ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+    ArrayRef<int64_t> flatOffsets = op.getFragmentOffsets();
+    ArrayRef<int64_t> flatSizes = op.getFragmentSizes();
+    ArrayRef<int64_t> flatStrides = *fragmentStridesAttr;
+    int64_t rank = resultType.getRank();
+
+    SmallVector<Value> fragmentOperands {adaptor.getInput()};
+    llvm::append_range(fragmentOperands, adaptor.getFragments());
+    if (failed(validateFragmentAssemblyMetadata(
+          op, rank, fragmentOperands.size(), operandIndices, sourceOffsets, flatOffsets, flatSizes, flatStrides)))
+      return failure();
+
+    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) {
+      int64_t operandIndex = operandIndices[fragmentIndex];
+
+      SmallVector<int64_t, 4> fragmentOffsets;
+      for (int64_t dim = 0; dim < rank; ++dim) {
+        int64_t flatIndex = fragmentIndex * rank + dim;
+        if (flatStrides[flatIndex] != 1)
+          return op.emitOpError("fragment assembly lowering only supports unit strides");
+        fragmentOffsets.push_back(flatOffsets[flatIndex]);
+      }
+
+      Value source = fragmentOperands[operandIndex];
+      auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+      if (!sourceType || !sourceType.hasStaticShape())
+        return op.emitOpError("fragment assembly lowering requires static ranked tensor operands");
+
+      SmallVector<int64_t, 4> fragmentShape;
+      fragmentShape.reserve(rank);
+      for (int64_t dim = 0; dim < rank; ++dim)
+        fragmentShape.push_back(flatSizes[fragmentIndex * rank + dim]);
+
+      Value fragment = source;
+      if (llvm::to_vector(sourceType.getShape()) != fragmentShape || sourceOffsets[fragmentIndex] != 0) {
+        FailureOr<SmallVector<int64_t, 4>> extractOffsets = getStaticSliceOffsetsForElementOffset(
+          op, sourceType, fragmentShape, sourceOffsets[fragmentIndex], "fragment assembly source slice");
+        if (failed(extractOffsets))
+          return failure();
+        fragment = tensor::ExtractSliceOp::create(rewriter,
+                                                  op.getLoc(),
+                                                  source,
+                                                  getStaticIndexAttrs(rewriter, *extractOffsets),
+                                                  getStaticIndexAttrs(rewriter, fragmentShape),
+                                                  getUnitStrides(rewriter, rank));
+      }
+
+      currentOutput = tensor::InsertSliceOp::create(rewriter,
+                                                    op.getLoc(),
+                                                    fragment,
+                                                    currentOutput,
+                                                    getStaticIndexAttrs(rewriter, fragmentOffsets),
+                                                    getStaticIndexAttrs(rewriter, fragmentShape),
+                                                    getUnitStrides(rewriter, rank))
+                        .getResult();
+    }
+
+    rewriter.replaceOp(op, currentOutput);
+    return success();
+  }
+};
+
 void populateInitialPatterns(RewritePatternSet& patterns) {
   raptor::populateWithGenerated(patterns);
   populateTransposeLoweringPatterns(patterns);
@@ -19,6 +109,7 @@ void populateInitialPatterns(RewritePatternSet& patterns) {
 void populateCoreBodyPatterns(RewritePatternSet& patterns) {
   raptor::populateWithGenerated(patterns);
   populateTransposeLoweringPatterns(patterns);
+  patterns.add<LowerFragmentAssemblyBlueprintPattern>(patterns.getContext());
 }
 
 } // namespace onnx_mlir

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

@@ -10,6 +10,14 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
+static void copyRaptorDebugAttrs(Operation* source, Operation* target) {
+  for (NamedAttribute attr : source->getAttrs()) {
+    StringRef name = attr.getName().strref();
+    if (name.starts_with("raptor."))
+      target->setAttr(attr.getName(), attr.getValue());
+  }
+}
+
 struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -17,7 +25,8 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
     auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
     if (failed(sizeAttr))
       return failure();
-    pim::PimSendOp::create(rewriter, op.getLoc(), op.getInput(), *sizeAttr, op.getTargetCoreId());
+    auto send = pim::PimSendOp::create(rewriter, op.getLoc(), op.getInput(), *sizeAttr, op.getTargetCoreId());
+    copyRaptorDebugAttrs(op.getOperation(), send.getOperation());
     rewriter.eraseOp(op);
     return success();
   }
@@ -37,9 +46,10 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getResult());
     if (failed(sizeAttr))
       return failure();
-    Value received = pim::PimReceiveOp::create(
-                       rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, *sizeAttr, op.getSourceCoreId())
-                       .getOutput();
+    auto receive = pim::PimReceiveOp::create(
+      rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, *sizeAttr, op.getSourceCoreId());
+    copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
+    Value received = receive.getOutput();
     rewriter.replaceOp(op, received);
     return success();
   }

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

@@ -59,7 +59,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
       return failure();
 
     for (auto& uses : extractSliceOp->getUses()) {
-      if (isa<spatial::SpatCompute>(uses.getOwner())) {
+      if (isa<spatial::SpatScheduledCompute>(uses.getOwner())) {
         if (!getDirectComputeLikeInputIndex(uses.getOwner(), uses.getOperandNumber()))
           return failure();
       }
@@ -72,7 +72,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
 
     for (auto& uses : llvm::make_early_inc_range(extractSliceOp->getUses())) {
 
-      if (auto spatCompute = dyn_cast<spatial::SpatCompute>(uses.getOwner())) {
+      if (auto spatCompute = dyn_cast<spatial::SpatScheduledCompute>(uses.getOwner())) {
         auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, uses.getOperandNumber());
         if (!inputIndex)
           return failure();
@@ -92,7 +92,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
         replaceAndEraseDirectComputeLikeInput(
           rewriter, spatCompute.getOperation(), *inputIndex, mapSpatToExtract[spatCompute.getOperation()]);
       }
-      else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(uses.getOwner())) {
+      else if (auto spatComputeBatch = dyn_cast<spatial::SpatScheduledComputeBatch>(uses.getOwner())) {
         auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, uses.getOperandNumber());
         if (!inputIndex)
           return failure();
@@ -114,7 +114,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
       }
       else {
         {
-          if (auto spatCompute = uses.getOwner()->getParentOfType<spatial::SpatCompute>()) {
+          if (auto spatCompute = uses.getOwner()->getParentOfType<spatial::SpatScheduledCompute>()) {
             rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
             if (!mapSpatToExtract.contains(spatCompute.getOperation())) {
               auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
@@ -125,7 +125,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
             uses.set(mapSpatToExtract[spatCompute.getOperation()]);
             rewriter.finalizeOpModification(spatCompute.getOperation());
           }
-          else if (auto spatComputeBatch = uses.getOwner()->getParentOfType<spatial::SpatComputeBatch>()) {
+          else if (auto spatComputeBatch = uses.getOwner()->getParentOfType<spatial::SpatScheduledComputeBatch>()) {
             rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
             if (!mapSpatToExtract.contains(spatComputeBatch.getOperation())) {
               auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
@@ -179,7 +179,7 @@ struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncO
 
       for (auto& argUses : llvm::make_early_inc_range(arg.getUses())) {
         auto argUser = argUses.getOwner();
-        if (auto spatCompute = dyn_cast<spatial::SpatCompute>(argUser)) {
+        if (auto spatCompute = dyn_cast<spatial::SpatScheduledCompute>(argUser)) {
           auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, argUses.getOperandNumber());
           if (!inputIndex)
             return failure();
@@ -191,7 +191,7 @@ struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncO
 
           replaceAndEraseDirectComputeLikeInput(rewriter, spatCompute.getOperation(), BBArgIndex, toTensor);
         }
-        else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(argUser)) {
+        else if (auto spatComputeBatch = dyn_cast<spatial::SpatScheduledComputeBatch>(argUser)) {
           auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, argUses.getOperandNumber());
           if (!inputIndex)
             return failure();

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -2,6 +2,7 @@
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -11,6 +12,7 @@
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -86,7 +88,7 @@ getCheckedByteOffset(int64_t elementOffset, size_t elementSize, Operation* ancho
   return pim::checkedCast<int64_t>(*byteOffset, anchor, fieldName);
 }
 
-static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
+static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute computeOp,
                                                SmallVectorImpl<Operation*>& helperChain) {
   if (computeOp.getInputs().size() != 1 || computeOp.getNumResults() != 1)
     return failure();
@@ -149,6 +151,40 @@ static std::optional<ReturnUseInfo> analyzeReturnUse(Value value) {
   };
 }
 
+static FailureOr<SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4>>
+analyzeTopLevelFragmentAssemblyUses(Value value) {
+  SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4> uses;
+  for (OpOperand& use : value.getUses()) {
+    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")
+      return failure();
+    if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
+      return failure();
+    std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+    std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+    std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
+    auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+    if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr || !resultType || !resultType.hasStaticShape())
+      return failure();
+    SmallVector<Value> fragmentOperands {blueprint.getInput()};
+    llvm::append_range(fragmentOperands, blueprint.getFragments());
+    if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                                resultType.getRank(),
+                                                fragmentOperands.size(),
+                                                *operandIndicesAttr,
+                                                *sourceOffsetsAttr,
+                                                blueprint.getFragmentOffsets(),
+                                                blueprint.getFragmentSizes(),
+                                                *stridesAttr)))
+      return failure();
+    uses.emplace_back(blueprint, use.getOperandNumber());
+  }
+  return uses;
+}
+
 static std::optional<ConcatReturnUseInfo> analyzeConcatReturnUse(Value value) {
   auto getConcatResult = [](Operation* op) -> Value {
     if (auto tensorConcat = dyn_cast<tensor::ConcatOp>(op))
@@ -212,7 +248,7 @@ static std::optional<ConcatReturnUseInfo> analyzeConcatReturnUse(Value value) {
   }
 
   SmallVector<Operation*> helperChain;
-  if (auto helperCompute = dyn_cast<spatial::SpatCompute>(currentUser)) {
+  if (auto helperCompute = dyn_cast<spatial::SpatScheduledCompute>(currentUser)) {
     if (helperCompute.getInputs().size() != 1 || helperCompute.getInputs().front() != currentValue)
       return std::nullopt;
 
@@ -375,6 +411,57 @@ static void cloneHelperChain(Value sourceValue,
   }
 }
 
+static bool isHostStaticReturnValue(Value value) {
+  llvm::SmallPtrSet<Operation*, 8> visited;
+  while (Operation* definingOp = value.getDefiningOp()) {
+    if (!visited.insert(definingOp).second)
+      return false;
+    if (isa<arith::ConstantOp>(definingOp) || definingOp->hasTrait<OpTrait::ConstantLike>())
+      return true;
+    if (!isReturnHelperChainOp(definingOp) || definingOp->getNumOperands() != 1)
+      return false;
+    value = definingOp->getOperand(0);
+  }
+  return false;
+}
+
+static FailureOr<Value>
+materializeHostStaticReturnValue(IRRewriter& rewriter, Value value, OperationFolder& constantFolder) {
+  llvm::SmallVector<Operation*> chain;
+  llvm::SmallPtrSet<Operation*, 8> visited;
+  while (Operation* definingOp = value.getDefiningOp()) {
+    if (!visited.insert(definingOp).second)
+      return failure();
+    chain.push_back(definingOp);
+    if (isa<arith::ConstantOp>(definingOp) || definingOp->hasTrait<OpTrait::ConstantLike>())
+      break;
+    if (!isReturnHelperChainOp(definingOp) || definingOp->getNumOperands() != 1)
+      return failure();
+    value = definingOp->getOperand(0);
+  }
+
+  if (chain.empty())
+    return failure();
+
+  IRMapping mapping;
+  Value clonedValue;
+  for (Operation* op : llvm::reverse(chain)) {
+    if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
+      clonedValue = getOrCreateConstantLike(constantFolder, constantOp);
+      mapping.map(op->getResult(0), clonedValue);
+      continue;
+    }
+
+    Operation* clonedOp = rewriter.clone(*op, mapping);
+    for (auto [originalResult, newResult] : llvm::zip(op->getResults(), clonedOp->getResults()))
+      mapping.map(originalResult, newResult);
+    clonedValue = clonedOp->getResult(0);
+    rewriter.setInsertionPointAfter(clonedOp);
+  }
+
+  return clonedValue;
+}
+
 static FailureOr<Value> emitHostCopy(IRRewriter& rewriter,
                                      Location loc,
                                      Value outputTensor,
@@ -444,7 +531,30 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
   OperationFolder constantFolder(producerOp->getContext());
   auto storedTensorType = cast<TensorType>(storedValue.getType());
 
+  auto materializeDirectHostReturn = [&](size_t returnIndex,
+                                         Value sourceValue,
+                                         ArrayRef<Operation*> helperChain) -> ReturnPathLoweringResult {
+    rewriter.setInsertionPointAfter(producerOp);
+    auto hostStaticValue = materializeHostStaticReturnValue(rewriter, sourceValue, constantFolder);
+    if (failed(hostStaticValue))
+      return ReturnPathLoweringResult::Failure;
+
+    Value hostReturnValue = *hostStaticValue;
+    if (!helperChain.empty())
+      cloneHelperChain(hostReturnValue, helperChain, rewriter, constantFolder, hostReturnValue);
+
+    outputTensors[returnIndex] =
+      [hostReturnValue](IRRewriter& rewriter, Location loc) -> Value { return hostReturnValue; };
+    return ReturnPathLoweringResult::Handled;
+  };
+
   if (auto returnUse = analyzeReturnUse(producedValue)) {
+    if (isHostStaticReturnValue(storedValue)) {
+      for (Operation* op : returnUse->helperChain)
+        markOpToRemove(op);
+      return materializeDirectHostReturn(returnUse->returnIndex, storedValue, returnUse->helperChain);
+    }
+
     Value currentStoredValue = storedValue;
     cloneHelperChain(storedValue, returnUse->helperChain, rewriter, constantFolder, currentStoredValue);
     for (Operation* op : returnUse->helperChain)
@@ -470,6 +580,8 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
 
     if (isa<func::ReturnOp>(resultUser)) {
       size_t resultIndexInReturn = resultUse.getOperandNumber();
+      if (isHostStaticReturnValue(storedValue))
+        return materializeDirectHostReturn(resultIndexInReturn, storedValue, {});
       auto byteSize =
         pim::getCheckedShapedTypeSizeInBytes(storedTensorType, producerOp, "return-path host copy byte size");
       if (failed(byteSize))
@@ -483,6 +595,115 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
     }
   }
 
+  FailureOr<SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4>> fragmentAssemblyUses =
+    analyzeTopLevelFragmentAssemblyUses(producedValue);
+  if (succeeded(fragmentAssemblyUses)) {
+    auto sourceType = dyn_cast<RankedTensorType>(storedValue.getType());
+    if (!sourceType || !sourceType.hasStaticShape()) {
+      producerOp->emitOpError("fragment assembly publication requires a static ranked tensor source");
+      return ReturnPathLoweringResult::Failure;
+    }
+
+    size_t elementSize = getElementTypeSizeInBytes(sourceType.getElementType());
+    for (auto [blueprint, operandNumber] : *fragmentAssemblyUses) {
+      rewriter.setInsertionPointAfterValue(storedValue);
+      std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+      std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+      std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
+      if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr) {
+        blueprint.emitOpError(
+          "fragment assembly lowering requires explicit operand, source-offset, and stride metadata");
+        return ReturnPathLoweringResult::Failure;
+      }
+
+      size_t returnIndex = blueprint.getOutput().getUses().begin()->getOperandNumber();
+      Value outputTensor = outputTensors[returnIndex](rewriter, loc);
+      auto outputType = dyn_cast<RankedTensorType>(outputTensor.getType());
+      auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+      if (!outputType || !resultType || !resultType.hasStaticShape()) {
+        blueprint.emitOpError("fragment assembly lowering requires static ranked host outputs");
+        return ReturnPathLoweringResult::Failure;
+      }
+
+      ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+      ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+      ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
+      ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
+      ArrayRef<int64_t> flatStrides = *stridesAttr;
+      int64_t rank = resultType.getRank();
+      if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                                  rank,
+                                                  1 + blueprint.getFragments().size(),
+                                                  operandIndices,
+                                                  sourceOffsets,
+                                                  flatOffsets,
+                                                  flatSizes,
+                                                  flatStrides)))
+        return ReturnPathLoweringResult::Failure;
+      SmallVector<FragmentAssemblyCopy, 8> copies;
+      for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+        if (operandIndices[fragmentIndex] != static_cast<int64_t>(operandNumber))
+          continue;
+
+        SmallVector<int64_t, 4> fragmentOffsets;
+        SmallVector<int64_t, 4> fragmentSizes;
+        for (int64_t dim = 0; dim < rank; ++dim) {
+          int64_t flatIndex = fragmentIndex * rank + dim;
+          if (flatStrides[flatIndex] != 1) {
+            blueprint.emitOpError("fragment assembly lowering only supports unit strides");
+            return ReturnPathLoweringResult::Failure;
+          }
+          fragmentOffsets.push_back(flatOffsets[flatIndex]);
+          fragmentSizes.push_back(flatSizes[flatIndex]);
+        }
+
+        bool failedChunk = false;
+        if (failed(forEachContiguousDestinationChunk(
+          outputType.getShape(),
+          fragmentOffsets,
+          fragmentSizes,
+          [&](ArrayRef<int64_t> chunkOffsets, int64_t relativeSourceOffset, int64_t chunkElements) -> LogicalResult {
+            auto hostOffset =
+              getCheckedByteOffset(computeFlatElementIndex(chunkOffsets, outputType.getShape()),
+                                   elementSize,
+                                   producerOp,
+                                   "fragment assembly host offset");
+            auto sourceOffset = getCheckedByteOffset(sourceOffsets[fragmentIndex] + relativeSourceOffset,
+                                                     elementSize,
+                                                     producerOp,
+                                                     "fragment assembly source offset");
+            auto fragmentBytes =
+              getCheckedByteOffset(chunkElements, elementSize, producerOp, "fragment assembly host copy byte size");
+            if (failed(hostOffset) || failed(sourceOffset) || failed(fragmentBytes)) {
+              failedChunk = true;
+              return failure();
+            }
+            FragmentAssemblyCopy copy;
+            copy.source = storedValue;
+            copy.sourceType = sourceType;
+            copy.hostByteOffset = *hostOffset;
+            copy.sourceByteOffset = *sourceOffset;
+            copy.byteSize = *fragmentBytes;
+            copies.push_back(copy);
+            return success();
+          })))
+          failedChunk = true;
+        if (failedChunk)
+          return ReturnPathLoweringResult::Failure;
+      }
+      FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>> runs = groupFragmentAssemblyCopyRuns(copies);
+      if (failed(runs))
+        return ReturnPathLoweringResult::Failure;
+      FailureOr<Value> updatedOutput =
+        emitFragmentAssemblyCopyRuns(rewriter, blueprint.getLoc(), *runs, outputTensor, producerOp);
+      if (failed(updatedOutput))
+        return ReturnPathLoweringResult::Failure;
+      outputTensor = *updatedOutput;
+      markOpToRemove(blueprint.getOperation());
+    }
+    return ReturnPathLoweringResult::Handled;
+  }
+
   if (auto concatReturnUse = analyzeConcatReturnUse(producedValue)) {
     size_t elementSize = getElementTypeSizeInBytes(storedTensorType.getElementType());
     auto storedByteSize =
@@ -567,7 +788,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
 }
 
 raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::lowerComputeResultReturnPath(
-  spatial::SpatCompute computeOp, OpResult result, Value yieldValue, IRRewriter& rewriter) {
+  spatial::SpatScheduledCompute computeOp, OpResult result, Value yieldValue, IRRewriter& rewriter) {
   return lowerProducedValueReturnPath(computeOp.getOperation(), result, yieldValue, rewriter);
 }
 
@@ -580,7 +801,7 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
     if (!isExclusivelyOwnedByReturnChain && op->hasOneUse()) {
       Operation* onlyUser = *op->getUsers().begin();
       isExclusivelyOwnedByReturnChain =
-        isa<func::ReturnOp, tensor::ConcatOp, spatial::SpatConcatOp, pim::PimConcatOp, spatial::SpatCompute>(onlyUser)
+        isa<func::ReturnOp, tensor::ConcatOp, spatial::SpatConcatOp, pim::PimConcatOp, spatial::SpatScheduledCompute>(onlyUser)
         || isReturnHelperChainOp(onlyUser);
     }
     if (!isExclusivelyOwnedByReturnChain)
@@ -593,7 +814,17 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
       return;
     }
 
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
+    if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+      std::optional<StringRef> mode = blueprint.getMode();
+      if (mode && *mode == "fragment_assembly") {
+        markOpToRemove(blueprint.getOperation());
+        for (Value operand : blueprint->getOperands())
+          markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);
+        return;
+      }
+    }
+
+    if (auto computeOp = dyn_cast<spatial::SpatScheduledCompute>(op)) {
       markOpToRemove(computeOp);
       if (!computeOp.getInputs().empty())
         for (Value input : computeOp.getInputs())

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

@@ -27,6 +27,12 @@ def spatToPimVVAdd : Pat<
     (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
 >;
 
+def spatToPimVVSub : Pat<
+  (SpatVSubOp:$srcOpRes $a, $b),
+  (PimVVSubOp $a, $b,
+    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
+>;
+
 def spatToPimVVMul : Pat<
   (SpatVMulOp:$srcOpRes $a, $b),
   (PimVVMulOp $a, $b,

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -25,9 +25,11 @@
 #include <cassert>
 #include <utility>
 
+#include "Common/IR/ShapeUtils.hpp"
 #include "Common/IR/ConstantUtils.hpp"
 #include "Common/PimCommon.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
+#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "Conversion/SpatialToPim/Patterns.hpp"
@@ -42,63 +44,29 @@ using namespace pim;
 
 namespace onnx_mlir {
 
-static memref::GlobalOp getOrCreateZeroGlobal(IRRewriter& rewriter, Location loc, RankedTensorType tensorType) {
-  auto moduleOp = rewriter.getBlock()->getParentOp()->getParentOfType<ModuleOp>();
-  auto memRefType = MemRefType::get(tensorType.getShape(), tensorType.getElementType());
-  auto zeroAttr = DenseElementsAttr::get(tensorType, rewriter.getZeroAttr(tensorType.getElementType()));
-
-  for (auto globalOp : moduleOp.getOps<memref::GlobalOp>()) {
-    if (!globalOp.getConstant() || globalOp.getType() != memRefType || !globalOp.getInitialValue())
-      continue;
-    if (dyn_cast<DenseElementsAttr>(*globalOp.getInitialValue()) == zeroAttr)
-      return globalOp;
-  }
-
-  std::string nameStem;
-  llvm::raw_string_ostream nameStream(nameStem);
-  nameStream << "__pim_zero_" << tensorType.getRank() << "d_" << tensorType.getNumElements();
-  nameStream.flush();
-
-  std::string symbolName = nameStem;
-  unsigned suffix = 0;
-  while (SymbolTable::lookupSymbolIn(moduleOp, symbolName))
-    symbolName = (nameStem + "_" + Twine(suffix++)).str();
-
-  OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPointToStart(moduleOp.getBody());
-  return memref::GlobalOp::create(rewriter,
-                                  loc,
-                                  rewriter.getStringAttr(symbolName),
-                                  rewriter.getStringAttr("private"),
-                                  TypeAttr::get(memRefType),
-                                  zeroAttr,
-                                  rewriter.getUnitAttr(),
-                                  IntegerAttr {});
-}
-
-static FailureOr<Value> createZeroedDeviceHVector(IRRewriter& rewriter,
-                                                  Location loc,
-                                                  RankedTensorType tensorType,
-                                                  OperationFolder& constantFolder) {
-  auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, tensorType);
-  auto zeroGlobal = getOrCreateZeroGlobal(rewriter, loc, tensorType);
-  auto zeroValue = memref::GetGlobalOp::create(rewriter, loc, zeroGlobal.getType(), zeroGlobal.getName());
-  auto zeroIndex = getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0);
-  auto byteSize =
-    pim::getCheckedShapedTypeSizeInBytes(tensorType, outputBuffer.getOperation(), "host-to-device zero copy byte size");
-  if (failed(byteSize))
-    return failure();
-  auto sizeAttr =
-    pim::getCheckedI32Attr(rewriter, outputBuffer.getOperation(), *byteSize, "host-to-device zero copy byte size");
-  if (failed(sizeAttr))
-    return failure();
-  return PimMemCopyHostToDevOp::create(
-           rewriter, loc, tensorType, zeroIndex, zeroIndex, outputBuffer, zeroValue, *sizeAttr)
-    .getOutput();
-}
-
 static FailureOr<Value>
-padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector, OperationFolder& constantFolder) {
+createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTensorType resultType) {
+  auto sourceType = cast<RankedTensorType>(value.getType());
+  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
+  SmallVector<OpFoldResult> highPads;
+  highPads.reserve(sourceType.getRank());
+  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
+    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
+
+  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
+  auto* padBlock = new Block();
+  for (int64_t i = 0; i < sourceType.getRank(); ++i)
+    padBlock->addArgument(rewriter.getIndexType(), loc);
+  padOp.getRegion().push_back(padBlock);
+  rewriter.setInsertionPointToStart(padBlock);
+  auto zero = getOrCreateConstant(
+    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
+  tensor::YieldOp::create(rewriter, loc, zero);
+  rewriter.setInsertionPointAfter(padOp);
+  return padOp.getResult();
+}
+
+static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
   assert(isHVectorShape(shape) && "expected a horizontal vector");
@@ -109,22 +77,10 @@ padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector,
 
   auto paddedType = RankedTensorType::get(
     {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
-  auto zeroed = createZeroedDeviceHVector(rewriter, loc, paddedType, constantFolder);
-  if (failed(zeroed))
-    return failure();
-  Value zeroIndex = getOrCreateIndexConstant(constantFolder, zeroed->getDefiningOp(), 0);
-  auto byteSize =
-    pim::getCheckedShapedTypeSizeInBytes(vectorType, zeroed->getDefiningOp(), "device padding copy byte size");
-  if (failed(byteSize))
-    return failure();
-  auto sizeAttr = pim::getCheckedI32Attr(rewriter, zeroed->getDefiningOp(), *byteSize, "device padding copy byte size");
-  if (failed(sizeAttr))
-    return failure();
-  return PimMemCopyOp::create(rewriter, loc, paddedType, zeroIndex, zeroIndex, *zeroed, vector, *sizeAttr).getOutput();
+  return createZeroPaddedTensor(rewriter, loc, vector, paddedType);
 }
 
 void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
-  coreId = 0;
   outputTensors.clear();
   operationsToRemove.clear();
   ModuleOp moduleOp = getOperation();
@@ -137,6 +93,12 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     return;
   }
   func::FuncOp funcOp = *entryFunc;
+  if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+    funcOp.emitOpError(
+      "scheduled Spatial verification failed at the start of SpatialToPim");
+    signalPassFailure();
+    return;
+  }
 
   IRRewriter rewriter(&getContext());
   OperationFolder constantFolder(&getContext());
@@ -176,19 +138,19 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     return;
   }
 
-  for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
+  for (auto computeOp : funcOp.getOps<spatial::SpatScheduledCompute>()) {
     markOpToRemove(computeOp);
     if (failed(lowerComputeOp(computeOp, rewriter, constantFolder))) {
-      computeOp.emitOpError("failed to lower spat.compute to pim.core");
+      computeOp.emitOpError("failed to lower spat.scheduled_compute to pim.core");
       signalPassFailure();
       return;
     }
   }
 
-  for (auto computeBatchOp : funcOp.getOps<spatial::SpatComputeBatch>()) {
+  for (auto computeBatchOp : funcOp.getOps<spatial::SpatScheduledComputeBatch>()) {
     markOpToRemove(computeBatchOp);
     if (failed(lowerComputeBatchOp(computeBatchOp, rewriter))) {
-      computeBatchOp.emitOpError("failed to lower spat.compute_batch to pim.core_batch");
+      computeBatchOp.emitOpError("failed to lower spat.scheduled_compute_batch to pim.core_batch");
       signalPassFailure();
       return;
     }
@@ -292,7 +254,6 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
 }
 
 LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
-  OperationFolder constantFolder(funcOp.getContext());
   bool hasFailure = false;
   funcOp.walk([&](PimVMMOp vmmOp) {
     auto outputType = cast<RankedTensorType>(vmmOp.getOutput().getType());
@@ -301,7 +262,7 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
     assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");
 
     rewriter.setInsertionPoint(vmmOp);
-    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput(), constantFolder);
+    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
     if (failed(paddedInput)) {
       hasFailure = true;
       return WalkResult::interrupt();
@@ -374,7 +335,7 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
   };
 
   for (auto& op : funcOp.getBody().getOps())
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
+    if (auto computeOp = dyn_cast<spatial::SpatScheduledCompute>(op)) {
       if (!computeOp.getInputs().empty() || computeOp.getBody().front().getNumArguments() != 0)
         continue;
       for (auto getGlobal : computeOp.getOps<memref::GetGlobalOp>()) {

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -36,13 +36,15 @@ private:
   using OutputTensorFactory = std::function<mlir::Value(mlir::IRRewriter& rewriter, mlir::Location loc)>;
 
   llvm::SmallVector<OutputTensorFactory> outputTensors;
-  size_t coreId = 0;
   llvm::SmallVector<mlir::Operation*> operationsToRemove;
 
   mlir::LogicalResult allocateAndInitializeCoreLocalVariables(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
   mlir::LogicalResult
-  lowerComputeOp(spatial::SpatCompute computeOp, mlir::IRRewriter& rewriter, mlir::OperationFolder& constantFolder);
-  mlir::LogicalResult lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, mlir::IRRewriter& rewriter);
+  lowerComputeOp(spatial::SpatScheduledCompute computeOp,
+                 mlir::IRRewriter& rewriter,
+                 mlir::OperationFolder& constantFolder);
+  mlir::LogicalResult lowerComputeBatchOp(spatial::SpatScheduledComputeBatch computeBatchOp,
+                                          mlir::IRRewriter& rewriter);
 
   enum class ReturnPathLoweringResult {
     Handled,
@@ -51,7 +53,7 @@ private:
   };
 
   void addReturnOutputBuffers(mlir::func::ReturnOp returnOp, mlir::IRRewriter& rewriter);
-  ReturnPathLoweringResult lowerComputeResultReturnPath(spatial::SpatCompute computeOp,
+  ReturnPathLoweringResult lowerComputeResultReturnPath(spatial::SpatScheduledCompute computeOp,
                                                         mlir::OpResult result,
                                                         mlir::Value yieldValue,
                                                         mlir::IRRewriter& rewriter);

src/PIM/Dialect/Pim/CMakeLists.txt

@@ -4,7 +4,6 @@ add_onnx_mlir_dialect_doc(pim Pim.td)
 add_subdirectory(Transforms/Bufferization)
 add_subdirectory(Transforms/MemoryCoalescing)
 add_subdirectory(Transforms/HostConstantFolding)
-add_subdirectory(Transforms/HostConstantMaterialization)
 add_subdirectory(Transforms/Verification)
 
 add_pim_library(PimOps

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

@@ -6,14 +6,20 @@
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.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/Transforms/Bufferization/Common.hpp"
 
 using namespace mlir;
 using namespace bufferization;
 
 namespace onnx_mlir::pim {
 
-FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue, Location loc, RewriterBase& rewriter) {
-  if (succeeded(resolveContiguousAddress(memrefValue)) || succeeded(compileContiguousAddressExpr(memrefValue)))
+FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue,
+                                                  Location loc,
+                                                  RewriterBase& rewriter,
+                                                  const StaticValueKnowledge& knowledge) {
+  bool isContiguous =
+    succeeded(resolveContiguousAddress(memrefValue, knowledge)) || succeeded(compileContiguousAddressExpr(memrefValue));
+  if (isContiguous && isDeviceLocalPimAddress(memrefValue, knowledge))
     return memrefValue;
 
   auto shapedType = cast<ShapedType>(memrefValue.getType());
@@ -29,13 +35,21 @@ FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue, Location lo
   if (failed(sizeAttr))
     return failure();
 
+  if (isHostBackedPimAddress(memrefValue, knowledge)) {
+    return PimMemCopyHostToDevOp::create(
+             rewriter, loc, contiguousType, zeroOffset, zeroOffset, contiguousBuffer, memrefValue, *sizeAttr)
+      .getOutput();
+  }
+
   return PimMemCopyOp::create(
            rewriter, loc, contiguousType, zeroOffset, zeroOffset, contiguousBuffer, memrefValue, *sizeAttr)
     .getOutput();
 }
 
 Value allocateContiguousResultMemRefLike(Value memrefValue, Location loc, RewriterBase& rewriter) {
-  if (succeeded(resolveContiguousAddress(memrefValue)))
+  bool isContiguous =
+    succeeded(resolveContiguousAddress(memrefValue)) || succeeded(compileContiguousAddressExpr(memrefValue));
+  if (isContiguous && isDeviceLocalPimAddress(memrefValue))
     return memrefValue;
 
   auto shapedType = cast<ShapedType>(memrefValue.getType());

src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp

@@ -3,10 +3,15 @@
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/IR/PatternMatch.h"
 
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+
 namespace onnx_mlir::pim {
 
 llvm::FailureOr<mlir::Value>
-materializeContiguousInputMemRef(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
+materializeContiguousInputMemRef(mlir::Value memrefValue,
+                                 mlir::Location loc,
+                                 mlir::RewriterBase& rewriter,
+                                 const onnx_mlir::StaticValueKnowledge& knowledge = {});
 mlir::Value
 allocateContiguousResultMemRefLike(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
 

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

@@ -1,9 +1,83 @@
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;
 
+static SmallVector<Region *> getSelectionRegions(OpResult result) {
+  SmallVector<Region *> regions;
+  if (auto selection = dyn_cast<scf::IndexSwitchOp>(result.getOwner()))
+    for (Region &region : selection->getRegions())
+      regions.push_back(&region);
+  else if (auto selection = dyn_cast<scf::IfOp>(result.getOwner())) {
+    regions.push_back(&selection.getThenRegion());
+    regions.push_back(&selection.getElseRegion());
+  }
+  return regions;
+}
+
+static bool isCoreBatchInputArgument(Value value) {
+  auto blockArg = dyn_cast<BlockArgument>(value);
+  if (!blockArg)
+    return false;
+
+  auto coreBatchOp = dyn_cast_or_null<onnx_mlir::pim::PimCoreBatchOp>(blockArg.getOwner()->getParentOp());
+  if (!coreBatchOp)
+    return false;
+
+  unsigned firstInputArg = 1 + coreBatchOp.getWeights().size();
+  return static_cast<unsigned>(blockArg.getArgNumber()) >= firstInputArg;
+}
+
+static FailureOr<Value> getPimStorageBase(Value value, const onnx_mlir::StaticValueKnowledge& knowledge) {
+  llvm::SmallPtrSet<Value, 8> visited;
+  while (value && visited.insert(value).second) {
+    Value alias = resolveLoopCarriedAlias(value, knowledge);
+    if (alias)
+      value = alias;
+
+    if (auto aliased = knowledge.aliases.lookup(value)) {
+      value = aliased;
+      continue;
+    }
+
+    if (auto base = onnx_mlir::pim::getPimAddressBase(value, knowledge); succeeded(base))
+      return base;
+
+    if (isa<BlockArgument>(value))
+      return value;
+
+    Operation* definingOp = value.getDefiningOp();
+    if (!definingOp)
+      return value;
+
+    if (auto subviewOp = dyn_cast<memref::SubViewOp>(definingOp)) {
+      value = subviewOp.getSource();
+      continue;
+    }
+    if (auto collapseOp = dyn_cast<memref::CollapseShapeOp>(definingOp)) {
+      value = collapseOp.getSrc();
+      continue;
+    }
+    if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(definingOp)) {
+      value = expandOp.getSrc();
+      continue;
+    }
+    if (auto castOp = dyn_cast<memref::CastOp>(definingOp)) {
+      value = castOp.getSource();
+      continue;
+    }
+
+    return value;
+  }
+
+  if (value)
+    return value;
+  return failure();
+}
+
 FailureOr<IntegerAttr> onnx_mlir::pim::getMemRefSizeInBytesAttr(OpBuilder& builder, Operation* anchor, Value memref) {
   auto type = mlir::cast<MemRefType>(memref.getType());
   auto byteSize = getCheckedShapedTypeSizeInBytes(type, anchor, "memref byte size");
@@ -11,3 +85,66 @@ FailureOr<IntegerAttr> onnx_mlir::pim::getMemRefSizeInBytesAttr(OpBuilder& build
     return failure();
   return getCheckedI32Attr(builder, anchor, *byteSize, "memref byte size");
 }
+
+FailureOr<Value> onnx_mlir::pim::getPimAddressBase(Value value, const StaticValueKnowledge& knowledge) {
+  Value alias = resolveLoopCarriedAlias(value, knowledge);
+  if (alias)
+    value = alias;
+
+  auto resolved = resolveContiguousAddress(value, knowledge);
+  if (succeeded(resolved))
+    return resolved->base;
+
+  auto compiled = compileContiguousAddressExpr(value);
+  if (failed(compiled)) {
+    if (isa<BlockArgument>(value))
+      return value;
+    return failure();
+  }
+  return compiled->base;
+}
+
+bool onnx_mlir::pim::isHostBackedPimAddress(Value value, const StaticValueKnowledge& knowledge) {
+  llvm::SmallPtrSet<Value, 8> visited;
+  std::function<bool(Value)> isHost = [&](Value current) {
+    auto base = getPimStorageBase(current, knowledge);
+    if (failed(base) || !visited.insert(*base).second)
+      return false;
+    bool resultIsHost = isCoreBatchInputArgument(*base)
+                        || isa_and_nonnull<memref::GetGlobalOp>(base->getDefiningOp());
+    auto result = dyn_cast<OpResult>(*base);
+    SmallVector<Region *> regions = result ? getSelectionRegions(result)
+                                          : SmallVector<Region *>();
+    if (!resultIsHost && !regions.empty())
+      resultIsHost = llvm::all_of(regions, [&](Region *region) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        return yield && result.getResultNumber() < yield.getNumOperands()
+               && isHost(yield.getOperand(result.getResultNumber()));
+      });
+    visited.erase(*base);
+    return resultIsHost;
+  };
+  return isHost(value);
+}
+
+bool onnx_mlir::pim::isDeviceLocalPimAddress(Value value, const StaticValueKnowledge& knowledge) {
+  llvm::SmallPtrSet<Value, 8> visited;
+  std::function<bool(Value)> isDevice = [&](Value current) {
+    auto base = getPimStorageBase(current, knowledge);
+    if (failed(base) || !visited.insert(*base).second)
+      return false;
+    bool resultIsDevice = isa_and_nonnull<memref::AllocOp>(base->getDefiningOp());
+    auto result = dyn_cast<OpResult>(*base);
+    SmallVector<Region *> regions = result ? getSelectionRegions(result)
+                                          : SmallVector<Region *>();
+    if (!resultIsDevice && !regions.empty())
+      resultIsDevice = llvm::all_of(regions, [&](Region *region) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        return yield && result.getResultNumber() < yield.getNumOperands()
+               && isDevice(yield.getOperand(result.getResultNumber()));
+      });
+    visited.erase(*base);
+    return resultIsDevice;
+  };
+  return isDevice(value);
+}

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

@@ -2,11 +2,19 @@
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+
 namespace onnx_mlir {
 namespace pim {
 
 mlir::FailureOr<mlir::IntegerAttr>
 getMemRefSizeInBytesAttr(mlir::OpBuilder& builder, mlir::Operation* anchor, mlir::Value memref);
 
+mlir::FailureOr<mlir::Value> getPimAddressBase(mlir::Value value, const StaticValueKnowledge& knowledge = {});
+
+bool isHostBackedPimAddress(mlir::Value value, const StaticValueKnowledge& knowledge = {});
+
+bool isDeviceLocalPimAddress(mlir::Value value, const StaticValueKnowledge& knowledge = {});
+
 } // namespace pim
 } // namespace onnx_mlir

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

@@ -2,6 +2,8 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 
+#include "llvm/Support/MathExtras.h"
+
 #include "ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
@@ -33,6 +35,7 @@ struct CopyEndpointPlan {
 
 struct CopyLoopPlan {
   SmallVector<int64_t> outerShape;
+  int64_t outerElements = 0;
   int64_t chunkBytes = 0;
   ByteOffsetExpr targetBaseOffset;
   ByteOffsetExpr sourceBaseOffset;
@@ -74,6 +77,24 @@ static void appendTerm(ByteOffsetExpr& expr, Value value, int64_t scale) {
     expr.terms.push_back(ByteOffsetTerm {value, scale});
 }
 
+static FailureOr<int64_t> checkedPositiveMul(int64_t lhs, int64_t rhs) {
+  int64_t result = 0;
+  if (lhs < 0 || rhs < 0 || llvm::MulOverflow(lhs, rhs, result))
+    return failure();
+  return result;
+}
+
+static FailureOr<int64_t> checkedPositiveProduct(ArrayRef<int64_t> values) {
+  int64_t result = 1;
+  for (int64_t value : values) {
+    auto product = checkedPositiveMul(result, value);
+    if (failed(product))
+      return failure();
+    result = *product;
+  }
+  return result;
+}
+
 static FailureOr<SmallVector<int64_t>> getStaticMemRefStrides(MemRefType type) {
   SmallVector<int64_t> strides;
   int64_t offset = 0;
@@ -84,6 +105,165 @@ static FailureOr<SmallVector<int64_t>> getStaticMemRefStrides(MemRefType type) {
   return strides;
 }
 
+static FailureOr<SmallVector<int64_t>> getProvenMemRefStrides(Value value) {
+  llvm::SmallPtrSet<Value, 8> visiting;
+  std::function<FailureOr<SmallVector<int64_t>>(Value)> prove =
+    [&](Value current) -> FailureOr<SmallVector<int64_t>> {
+      auto type = dyn_cast<MemRefType>(current.getType());
+      if (!type || !visiting.insert(current).second)
+        return failure();
+      if (auto strides = getStaticMemRefStrides(type); succeeded(strides)) {
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto castOp = current.getDefiningOp<memref::CastOp>()) {
+        auto strides = prove(castOp.getSource());
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto subview = current.getDefiningOp<memref::SubViewOp>()) {
+        auto sourceStrides = prove(subview.getSource());
+        if (failed(sourceStrides) || subview.getSourceType().getRank() != subview.getType().getRank()) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides;
+        for (auto [sourceStride, viewStride] :
+             llvm::zip_equal(*sourceStrides, subview.getStaticStrides())) {
+          if (ShapedType::isDynamic(viewStride) || viewStride < 0) {
+            visiting.erase(current);
+            return failure();
+          }
+          auto stride = checkedPositiveMul(sourceStride, viewStride);
+          if (failed(stride)) {
+            visiting.erase(current);
+            return failure();
+          }
+          strides.push_back(*stride);
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto expand = current.getDefiningOp<memref::ExpandShapeOp>()) {
+        auto sourceStrides = prove(expand.getSrc());
+        auto resultType = dyn_cast<MemRefType>(expand.getResult().getType());
+        auto sourceType = dyn_cast<MemRefType>(expand.getSrc().getType());
+        if (failed(sourceStrides) || !sourceType || !resultType
+            || !resultType.hasStaticShape()
+            || sourceStrides->size() != static_cast<size_t>(sourceType.getRank())
+            || llvm::any_of(resultType.getShape(), [](int64_t dim) {
+                 return dim <= 0;
+               })) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides(resultType.getRank());
+        SmallVector<bool> assigned(resultType.getRank(), false);
+        for (auto [sourceDim, group] :
+             llvm::enumerate(expand.getReassociationIndices())) {
+          if (sourceDim >= sourceStrides->size() || group.empty()) {
+            visiting.erase(current);
+            return failure();
+          }
+          int64_t stride = (*sourceStrides)[sourceDim];
+          for (int64_t resultDim : llvm::reverse(group)) {
+            if (resultDim < 0 || resultDim >= resultType.getRank()
+                || assigned[resultDim]) {
+              visiting.erase(current);
+              return failure();
+            }
+            strides[resultDim] = stride;
+            assigned[resultDim] = true;
+            auto nextStride = checkedPositiveMul(
+              stride, resultType.getDimSize(resultDim));
+            if (failed(nextStride)) {
+              visiting.erase(current);
+              return failure();
+            }
+            stride = *nextStride;
+          }
+        }
+        if (llvm::is_contained(assigned, false)) {
+          visiting.erase(current);
+          return failure();
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto collapse = current.getDefiningOp<memref::CollapseShapeOp>()) {
+        auto sourceStrides = prove(collapse.getSrc());
+        auto sourceType = dyn_cast<MemRefType>(collapse.getSrc().getType());
+        if (failed(sourceStrides) || !sourceType
+            || !sourceType.hasStaticShape()
+            || sourceStrides->size() != static_cast<size_t>(sourceType.getRank())) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides;
+        for (ArrayRef<int64_t> group : collapse.getReassociationIndices()) {
+          if (group.empty()) {
+            visiting.erase(current);
+            return failure();
+          }
+          for (int64_t dim : group)
+            if (dim < 0 || dim >= sourceType.getRank()
+                || sourceType.getDimSize(dim) <= 0
+                || (*sourceStrides)[dim] < 0) {
+              visiting.erase(current);
+              return failure();
+            }
+          for (auto pair : llvm::zip(group.drop_back(), group.drop_front())) {
+            int64_t outer = std::get<0>(pair);
+            int64_t inner = std::get<1>(pair);
+            auto expectedOuterStride = checkedPositiveMul(
+              (*sourceStrides)[inner], sourceType.getDimSize(inner));
+            if (failed(expectedOuterStride)
+                || (*sourceStrides)[outer] != *expectedOuterStride) {
+              visiting.erase(current);
+              return failure();
+            }
+          }
+          strides.push_back((*sourceStrides)[group.back()]);
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      auto result = dyn_cast<OpResult>(current);
+      SmallVector<Region *> regions;
+      if (result) {
+        if (auto selection = dyn_cast<scf::IndexSwitchOp>(result.getOwner()))
+          for (Region &region : selection->getRegions())
+            regions.push_back(&region);
+        else if (auto selection = dyn_cast<scf::IfOp>(result.getOwner())) {
+          regions.push_back(&selection.getThenRegion());
+          regions.push_back(&selection.getElseRegion());
+        }
+      }
+      if (regions.empty()) {
+        visiting.erase(current);
+        return failure();
+      }
+      std::optional<SmallVector<int64_t>> common;
+      for (Region *region : regions) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        if (!yield || result.getResultNumber() >= yield.getNumOperands()) {
+          visiting.erase(current);
+          return failure();
+        }
+        auto strides = prove(yield.getOperand(result.getResultNumber()));
+        if (failed(strides) || (common && *common != *strides)) {
+          visiting.erase(current);
+          return failure();
+        }
+        common = std::move(*strides);
+      }
+      visiting.erase(current);
+      return common ? FailureOr<SmallVector<int64_t>>(std::move(*common))
+                    : FailureOr<SmallVector<int64_t>>(failure());
+    };
+  return prove(value);
+}
+
 static FailureOr<int64_t> getShapedByteSize(MemRefType type) {
   if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType()))
     return failure();
@@ -119,12 +299,15 @@ inferLogicalCopyShape(MemRefType targetType, MemRefType sourceType, int64_t size
   return failure();
 }
 
-static FailureOr<int64_t> getContiguousSuffixRank(MemRefType type, ArrayRef<int64_t> copyShape) {
-  if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())
+static FailureOr<int64_t> getContiguousSuffixRank(Value value, ArrayRef<int64_t> copyShape) {
+  auto type = dyn_cast<MemRefType>(value.getType());
+  if (!type || !type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())
       || type.getRank() != static_cast<int64_t>(copyShape.size()))
     return failure();
+  if (llvm::any_of(copyShape, [](int64_t dim) { return dim <= 0; }))
+    return failure();
 
-  auto strides = getStaticMemRefStrides(type);
+  auto strides = getProvenMemRefStrides(value);
   if (failed(strides))
     return failure();
 
@@ -134,7 +317,10 @@ static FailureOr<int64_t> getContiguousSuffixRank(MemRefType type, ArrayRef<int6
     if ((*strides)[dim] != expectedStride)
       break;
     ++contiguousSuffixRank;
-    expectedStride *= copyShape[dim];
+    auto nextStride = checkedPositiveMul(expectedStride, copyShape[dim]);
+    if (failed(nextStride))
+      return failure();
+    expectedStride = *nextStride;
   }
   return contiguousSuffixRank;
 }
@@ -174,18 +360,25 @@ static FailureOr<CopyEndpointPlan> analyzeCopyEndpoint(Value value, Value initia
     if (!sourceType || !sourceType.hasStaticShape() || !hasByteSizedElementType(sourceType.getElementType()))
       return failure();
 
-    auto sourceStrides = getStaticMemRefStrides(sourceType);
+    auto sourceStrides = getProvenMemRefStrides(subviewOp.getSource());
     if (failed(sourceStrides))
       return failure();
 
     int64_t elementByteWidth = static_cast<int64_t>(getElementTypeSizeInBytes(sourceType.getElementType()));
     for (auto [offset, stride] : llvm::zip_equal(subviewOp.getMixedOffsets(), *sourceStrides)) {
-      int64_t byteScale = stride * elementByteWidth;
+      auto byteScale = checkedPositiveMul(stride, elementByteWidth);
+      if (failed(byteScale))
+        return failure();
       if (auto attr = dyn_cast<Attribute>(offset)) {
-        endpoint.offset.constant += cast<IntegerAttr>(attr).getInt() * byteScale;
+        auto constantOffset = checkedPositiveMul(
+          cast<IntegerAttr>(attr).getInt(), *byteScale);
+        if (failed(constantOffset)
+            || llvm::AddOverflow(endpoint.offset.constant, *constantOffset,
+                                 endpoint.offset.constant))
+          return failure();
         continue;
       }
-      appendTerm(endpoint.offset, cast<Value>(offset), byteScale);
+      appendTerm(endpoint.offset, cast<Value>(offset), *byteScale);
     }
 
     endpoint.base = subviewOp.getSource();
@@ -204,17 +397,34 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
   if (!targetType || !sourceType || size <= 0)
     return failure();
 
-  auto logicalCopyShape = inferLogicalCopyShape(targetType, sourceType, size);
-  if (failed(logicalCopyShape))
-    return failure();
-
   auto targetPlan = analyzeCopyEndpoint(target, targetOffset, targetType);
   auto sourcePlan = analyzeCopyEndpoint(source, sourceOffset, sourceType);
   if (failed(targetPlan) || failed(sourcePlan))
     return failure();
 
-  auto targetSuffixRank = getContiguousSuffixRank(targetType, *logicalCopyShape);
-  auto sourceSuffixRank = getContiguousSuffixRank(sourceType, *logicalCopyShape);
+  auto targetBytes = getShapedByteSize(targetType);
+  auto sourceBytes = getShapedByteSize(sourceType);
+  if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes)
+      && *targetBytes == size && *sourceBytes == size) {
+    auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape());
+    auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape());
+    if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank)
+        && *targetSuffixRank == targetType.getRank() && *sourceSuffixRank == sourceType.getRank()) {
+      CopyRewritePlan plan;
+      plan.kind = CopyRewritePlan::Kind::Direct;
+      plan.target = *targetPlan;
+      plan.source = *sourcePlan;
+      plan.directBytes = size;
+      return plan;
+    }
+  }
+
+  auto logicalCopyShape = inferLogicalCopyShape(targetType, sourceType, size);
+  if (failed(logicalCopyShape))
+    return failure();
+
+  auto targetSuffixRank = getContiguousSuffixRank(target, *logicalCopyShape);
+  auto sourceSuffixRank = getContiguousSuffixRank(source, *logicalCopyShape);
   if (failed(targetSuffixRank) || failed(sourceSuffixRank))
     return failure();
 
@@ -229,8 +439,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
     return plan;
   }
 
-  auto targetStrides = getStaticMemRefStrides(targetType);
-  auto sourceStrides = getStaticMemRefStrides(sourceType);
+  auto targetStrides = getProvenMemRefStrides(target);
+  auto sourceStrides = getProvenMemRefStrides(source);
   if (failed(targetStrides) || failed(sourceStrides))
     return failure();
 
@@ -240,11 +450,27 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
   plan.loop.sourceBaseOffset = plan.source.offset;
   plan.loop.outerShape.assign(logicalCopyShape->begin(), logicalCopyShape->end() - contiguousSuffixRank);
   SmallVector<int64_t> chunkShape(logicalCopyShape->end() - contiguousSuffixRank, logicalCopyShape->end());
-  plan.loop.chunkBytes = getNumElements(chunkShape) * elementByteWidth;
-  for (int64_t stride : ArrayRef<int64_t>(*targetStrides).take_front(plan.loop.outerShape.size()))
-    plan.loop.targetOuterByteStrides.push_back(stride * elementByteWidth);
-  for (int64_t stride : ArrayRef<int64_t>(*sourceStrides).take_front(plan.loop.outerShape.size()))
-    plan.loop.sourceOuterByteStrides.push_back(stride * elementByteWidth);
+  auto outerElements = checkedPositiveProduct(plan.loop.outerShape);
+  auto chunkElements = checkedPositiveProduct(chunkShape);
+  auto chunkBytes = failed(chunkElements)
+                      ? FailureOr<int64_t>(failure())
+                      : checkedPositiveMul(*chunkElements, elementByteWidth);
+  if (failed(outerElements) || failed(chunkBytes))
+    return failure();
+  plan.loop.outerElements = *outerElements;
+  plan.loop.chunkBytes = *chunkBytes;
+  for (int64_t stride : ArrayRef<int64_t>(*targetStrides).take_front(plan.loop.outerShape.size())) {
+    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
+    if (failed(byteStride))
+      return failure();
+    plan.loop.targetOuterByteStrides.push_back(*byteStride);
+  }
+  for (int64_t stride : ArrayRef<int64_t>(*sourceStrides).take_front(plan.loop.outerShape.size())) {
+    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
+    if (failed(byteStride))
+      return failure();
+    plan.loop.sourceOuterByteStrides.push_back(*byteStride);
+  }
   if (plan.loop.chunkBytes <= 0)
     return failure();
   return plan;
@@ -344,7 +570,7 @@ static LogicalResult rewriteCopyLikeOp(CopyOp copyOp,
   }
 
   Value c0 = createIndexConstant(rewriter, anchorOp, 0);
-  Value cUpper = createIndexConstant(rewriter, anchorOp, getNumElements(plan->loop.outerShape));
+  Value cUpper = createIndexConstant(rewriter, anchorOp, plan->loop.outerElements);
   Value cStep = createIndexConstant(rewriter, anchorOp, 1);
   auto loop = buildNormalizedScfFor(
     rewriter,

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

@@ -15,6 +15,26 @@ using namespace bufferization;
 namespace onnx_mlir {
 namespace pim {
 
+static StaticValueKnowledge getEnclosingBufferizationKnowledge(Operation* op) {
+  StaticValueKnowledge knowledge;
+
+  if (auto coreBatchOp = op->getParentOfType<PimCoreBatchOp>()) {
+    knowledge.indexValues[coreBatchOp.getLaneArgument()] = 0;
+    for (auto [index, weight] : llvm::enumerate(coreBatchOp.getWeights()))
+      knowledge.aliases[coreBatchOp.getWeightArgument(index)] = weight;
+    for (auto [index, input] : llvm::enumerate(coreBatchOp.getInputs()))
+      knowledge.aliases[coreBatchOp.getInputArgument(index)] = input;
+    return knowledge;
+  }
+
+  if (auto coreOp = op->getParentOfType<PimCoreOp>()) {
+    for (auto [index, weight] : llvm::enumerate(coreOp.getWeights()))
+      knowledge.aliases[coreOp.getWeightArgument(index)] = weight;
+  }
+
+  return knowledge;
+}
+
 struct MemCopyHostToDevOpInterface
 : DstBufferizableOpInterfaceExternalModel<MemCopyHostToDevOpInterface, PimMemCopyHostToDevOp> {
   LogicalResult bufferize(Operation* op,
@@ -148,7 +168,8 @@ struct ConcatOpInterface : DstBufferizableOpInterfaceExternalModel<ConcatOpInter
       auto inputOpt = getBufferOrValue(rewriter, input, options, state);
       if (failed(inputOpt))
         return failure();
-      auto contiguous = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+      auto contiguous =
+        materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
       if (failed(contiguous))
         return failure();
       inputs.push_back(*contiguous);
@@ -182,7 +203,8 @@ struct SendOpInterface : BufferizableOpInterface::ExternalModel<SendOpInterface,
     auto inputOpt = getBufferOrValue(rewriter, sendOp.getInput(), options, state);
     if (failed(inputOpt))
       return failure();
-    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    auto contiguousInput =
+      materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousInput))
       return failure();
 
@@ -410,7 +432,8 @@ struct TransposeOpInterface : DstBufferizableOpInterfaceExternalModel<TransposeO
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    auto contiguousInput =
+      materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousInput))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
@@ -456,7 +479,8 @@ struct VMMOpInterface : DstBufferizableOpInterfaceExternalModel<VMMOpInterface,
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    auto contiguousInput =
+      materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousInput))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
@@ -497,10 +521,12 @@ struct BinaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<BinaryDstO
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousLhs = materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter);
+    auto contiguousLhs =
+      materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousLhs))
       return failure();
-    auto contiguousRhs = materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter);
+    auto contiguousRhs =
+      materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousRhs))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
@@ -534,10 +560,12 @@ struct VVDMulOpInterface : DstBufferizableOpInterfaceExternalModel<VVDMulOpInter
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousLhs = materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter);
+    auto contiguousLhs =
+      materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousLhs))
       return failure();
-    auto contiguousRhs = materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter);
+    auto contiguousRhs =
+      materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousRhs))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
@@ -574,7 +602,8 @@ struct UnaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<UnaryDstOpI
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    auto contiguousInput =
+      materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousInput))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);

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

@@ -8,6 +8,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Casting.h"
 
 #include "Common/PimCommon.hpp"
@@ -15,6 +16,7 @@
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
 #include "Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
@@ -27,24 +29,71 @@ namespace onnx_mlir {
 
 namespace {
 
-struct MemRefCopyToPimMemCopyPattern final : OpRewritePattern<memref::CopyOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct MemRefCopyWorkItem {
+  memref::CopyOp copyOp;
+  StaticValueKnowledge knowledge;
+};
 
-  LogicalResult matchAndRewrite(memref::CopyOp copyOp, PatternRewriter& rewriter) const override {
-    if (!copyOp->getParentOfType<pim::PimCoreOp>() && !copyOp->getParentOfType<pim::PimCoreBatchOp>())
-      return failure();
+static StaticValueKnowledge seedCoreKnowledge(pim::PimCoreOp coreOp) {
+  StaticValueKnowledge knowledge;
+  for (auto [index, weight] : llvm::enumerate(coreOp.getWeights()))
+    knowledge.aliases[coreOp.getWeightArgument(index)] = weight;
+  return knowledge;
+}
 
-    auto sourceType = dyn_cast<MemRefType>(copyOp.getSource().getType());
-    auto targetType = dyn_cast<MemRefType>(copyOp.getTarget().getType());
-    if (!sourceType || !targetType || !sourceType.hasStaticShape() || !targetType.hasStaticShape())
-      return failure();
-    if (sourceType.getElementType() != targetType.getElementType())
-      return failure();
+static StaticValueKnowledge seedCoreBatchKnowledge(pim::PimCoreBatchOp coreBatchOp, unsigned lane) {
+  StaticValueKnowledge knowledge;
+  knowledge.indexValues[coreBatchOp.getLaneArgument()] = lane;
+  for (auto [index, weight] : llvm::enumerate(coreBatchOp.getWeights()))
+    knowledge.aliases[coreBatchOp.getWeightArgument(index)] = weight;
+  for (auto [index, input] : llvm::enumerate(coreBatchOp.getInputs()))
+    knowledge.aliases[coreBatchOp.getInputArgument(index)] = input;
+  return knowledge;
+}
 
-    Value zeroOffset = getOrCreateIndexConstant(rewriter, copyOp, 0);
-    auto sizeAttr = getMemRefSizeInBytesAttr(rewriter, copyOp.getOperation(), copyOp.getSource());
-    if (failed(sizeAttr))
-      return failure();
+static LogicalResult
+lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const StaticValueKnowledge& knowledge) {
+  if (!copyOp->getParentOfType<pim::PimCoreOp>() && !copyOp->getParentOfType<pim::PimCoreBatchOp>())
+    return failure();
+
+  auto sourceType = dyn_cast<MemRefType>(copyOp.getSource().getType());
+  auto targetType = dyn_cast<MemRefType>(copyOp.getTarget().getType());
+  if (!sourceType || !targetType || !sourceType.hasStaticShape() || !targetType.hasStaticShape())
+    return failure();
+  if (sourceType.getElementType() != targetType.getElementType())
+    return failure();
+
+  Value zeroOffset = getOrCreateIndexConstant(rewriter, copyOp, 0);
+  auto sizeAttr = getMemRefSizeInBytesAttr(rewriter, copyOp.getOperation(), copyOp.getSource());
+  if (failed(sizeAttr))
+    return failure();
+
+  bool sourceIsHost = isHostBackedPimAddress(copyOp.getSource(), knowledge);
+  bool targetIsHost = isHostBackedPimAddress(copyOp.getTarget(), knowledge);
+  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge);
+  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge);
+
+  if (targetIsDevice && sourceIsHost) {
+    pim::PimMemCopyHostToDevOp::create(rewriter,
+                                       copyOp.getLoc(),
+                                       copyOp.getTarget().getType(),
+                                       zeroOffset,
+                                       zeroOffset,
+                                       copyOp.getTarget(),
+                                       copyOp.getSource(),
+                                       *sizeAttr);
+  }
+  else if (targetIsHost && sourceIsDevice) {
+    pim::PimMemCopyDevToHostOp::create(rewriter,
+                                       copyOp.getLoc(),
+                                       copyOp.getTarget().getType(),
+                                       zeroOffset,
+                                       zeroOffset,
+                                       copyOp.getTarget(),
+                                       copyOp.getSource(),
+                                       *sizeAttr);
+  }
+  else if (targetIsDevice && sourceIsDevice) {
     pim::PimMemCopyOp::create(rewriter,
                               copyOp.getLoc(),
                               copyOp.getTarget().getType(),
@@ -53,10 +102,49 @@ struct MemRefCopyToPimMemCopyPattern final : OpRewritePattern<memref::CopyOp> {
                               copyOp.getTarget(),
                               copyOp.getSource(),
                               *sizeAttr);
-    rewriter.eraseOp(copyOp);
-    return success();
   }
-};
+  else {
+    copyOp.emitOpError() << "failed to classify memref.copy endpoints: source=" << copyOp.getSource()
+                         << " type=" << copyOp.getSource().getType() << " host=" << sourceIsHost
+                         << " device=" << sourceIsDevice << ", target=" << copyOp.getTarget()
+                         << " type=" << copyOp.getTarget().getType() << " host=" << targetIsHost
+                         << " device=" << targetIsDevice;
+    return failure();
+  }
+
+  rewriter.eraseOp(copyOp);
+  return success();
+}
+
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyHostToDevOp copyOp, const StaticValueKnowledge& knowledge) {
+  bool sourceIsHost = isHostBackedPimAddress(copyOp.getHostSource(), knowledge);
+  bool targetIsHost = isHostBackedPimAddress(copyOp.getDeviceTarget(), knowledge);
+  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getHostSource(), knowledge);
+  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceTarget(), knowledge);
+  if (!sourceIsHost || !targetIsDevice || targetIsHost || sourceIsDevice)
+    return copyOp.emitOpError("pim.memcp_hd requires a host-backed source and a device-local target");
+  return success();
+}
+
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyDevToHostOp copyOp, const StaticValueKnowledge& knowledge) {
+  bool sourceIsHost = isHostBackedPimAddress(copyOp.getDeviceSource(), knowledge);
+  bool targetIsHost = isHostBackedPimAddress(copyOp.getHostTarget(), knowledge);
+  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceSource(), knowledge);
+  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getHostTarget(), knowledge);
+  if (!targetIsHost || !sourceIsDevice || sourceIsHost || targetIsDevice)
+    return copyOp.emitOpError("pim.memcp_dh requires a device-local source and a host-backed target");
+  return success();
+}
+
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyOp copyOp, const StaticValueKnowledge& knowledge) {
+  bool sourceIsHost = isHostBackedPimAddress(copyOp.getSource(), knowledge);
+  bool targetIsHost = isHostBackedPimAddress(copyOp.getTarget(), knowledge);
+  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge);
+  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge);
+  if (!sourceIsDevice || !targetIsDevice || sourceIsHost || targetIsHost)
+    return copyOp.emitOpError("pim.memcp requires device-local source and target operands");
+  return success();
+}
 
 struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass)
@@ -71,6 +159,7 @@ struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<Mo
 private:
   void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const;
   LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) const;
+  LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) const;
 };
 
 static LogicalResult applyPatternsOnce(Operation* op, PatternApplicator& applicator, PatternRewriter& rewriter) {
@@ -100,25 +189,46 @@ void PimBufferizationPass::runOnOperation() {
   }
 
   MLIRContext* ctx = moduleOp.getContext();
-  RewritePatternSet memrefCopyPatterns(ctx);
-  memrefCopyPatterns.add<MemRefCopyToPimMemCopyPattern>(ctx);
-  FrozenRewritePatternSet frozenMemrefCopyPatterns(std::move(memrefCopyPatterns));
-  PatternApplicator memrefCopyApplicator(frozenMemrefCopyPatterns);
-  memrefCopyApplicator.applyDefaultCostModel();
   PatternRewriter rewriter(ctx);
 
-  SmallVector<memref::CopyOp> copyWorklist;
-  moduleOp.walk([&](memref::CopyOp copyOp) {
-    if (copyOp->getParentOfType<pim::PimCoreOp>() || copyOp->getParentOfType<pim::PimCoreBatchOp>())
-      copyWorklist.push_back(copyOp);
+  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) {
+    llvm::SmallVector<unsigned, 2> lanes;
+    lanes.push_back(0);
+    if (coreBatchOp.getLaneCount() > 1)
+      lanes.push_back(static_cast<unsigned>(coreBatchOp.getLaneCount() - 1));
+    for (unsigned lane : lanes) {
+      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;
-  for (memref::CopyOp copyOp : copyWorklist) {
-    if (failed(applyPatternsOnce(copyOp, memrefCopyApplicator, rewriter))) {
-      copyOp.emitOpError("failed to lower memref.copy inside PIM core body");
+  for (const MemRefCopyWorkItem& workItem : copyWorklist) {
+    memref::CopyOp copyOp = workItem.copyOp;
+    rewriter.setInsertionPoint(copyOp);
+    if (failed(lowerMemRefCopyToPimCopy(copyOp, rewriter, workItem.knowledge)))
       hasFailed = true;
-    }
   }
   if (hasFailed) {
     signalPassFailure();
@@ -161,6 +271,10 @@ void PimBufferizationPass::runOnOperation() {
     signalPassFailure();
     return;
   }
+  if (failed(verifyPimCopyAddressSpaces(moduleOp))) {
+    signalPassFailure();
+    return;
+  }
 
   annotateWeightsMemrefs(moduleOp, funcOp);
 
@@ -188,76 +302,87 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO
 
 LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const {
   bool hasFailure = false;
-  moduleOp.walk([&](Operation* op) {
-    auto verifyOperand = [&](Value operand, unsigned operandIndex) {
-      if (!isa<BaseMemRefType>(operand.getType()))
-        return;
-      if (succeeded(resolveContiguousAddress(operand)) || succeeded(compileContiguousAddressExpr(operand)))
-        return;
-      op->emitOpError() << "operand #" << operandIndex
-                        << " is not backed by contiguous addressable storage after PIM bufferization";
-      hasFailure = true;
-    };
 
-    if (auto memCopyOp = dyn_cast<PimMemCopyOp>(op)) {
-      if (!pim::isNormalizedCopyOp(memCopyOp)) {
-        memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-        hasFailure = true;
-      }
-      verifyOperand(memCopyOp.getTarget(), 0);
-      verifyOperand(memCopyOp.getSource(), 1);
-      return;
-    }
-    if (auto loadOp = dyn_cast<PimMemCopyHostToDevOp>(op)) {
-      if (!pim::isNormalizedCopyOp(loadOp)) {
-        loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-        hasFailure = true;
-      }
-      verifyOperand(loadOp.getDeviceTarget(), 2);
-      verifyOperand(loadOp.getHostSource(), 3);
-      return;
-    }
-    if (auto storeOp = dyn_cast<PimMemCopyDevToHostOp>(op)) {
-      if (!pim::isNormalizedCopyOp(storeOp)) {
-        storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-        hasFailure = true;
-      }
-      verifyOperand(storeOp.getHostTarget(), 2);
-      verifyOperand(storeOp.getDeviceSource(), 3);
-      return;
-    }
-    if (auto sendOp = dyn_cast<PimSendOp>(op)) {
-      verifyOperand(sendOp.getInput(), 0);
-      return;
-    }
-    if (auto receiveOp = dyn_cast<PimReceiveOp>(op)) {
-      verifyOperand(receiveOp.getOutputBuffer(), 0);
-      return;
-    }
-    if (auto concatOp = dyn_cast<PimConcatOp>(op)) {
-      verifyOperand(concatOp.getOutputBuffer(), 0);
-      for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs()))
-        verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1);
-      return;
-    }
-    if (isa<PimTransposeOp,
-            PimVMMOp,
-            PimVVAddOp,
-            PimVVSubOp,
-            PimVVMulOp,
-            PimVVMaxOp,
-            PimVVDMulOp,
-            PimVAvgOp,
-            PimVReluOp,
-            PimVTanhOp,
-            PimVSigmOp,
-            PimVSoftmaxOp>(op)) {
-      for (auto operandAndIndex : llvm::enumerate(op->getOperands())) {
-        if (auto vmmOp = dyn_cast<PimVMMOp>(op); vmmOp && operandAndIndex.index() == 0)
-          continue;
-        verifyOperand(operandAndIndex.value(), operandAndIndex.index());
-      }
-    }
+  auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
+    (void) walkPimCoreBlockStructurally(
+      coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
+        auto verifyOperand = [&](Value operand, unsigned operandIndex) {
+          if (!isa<BaseMemRefType>(operand.getType()))
+            return;
+          if (succeeded(resolveContiguousAddress(operand, knowledge)) || succeeded(compileContiguousAddressExpr(operand)))
+            return;
+          op.emitOpError() << "operand #" << operandIndex
+                           << " is not backed by contiguous addressable storage after PIM bufferization";
+          hasFailure = true;
+        };
+
+        if (auto memCopyOp = dyn_cast<PimMemCopyOp>(&op)) {
+          if (!pim::isNormalizedCopyOp(memCopyOp)) {
+            memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+            hasFailure = true;
+          }
+          verifyOperand(memCopyOp.getTarget(), 0);
+          verifyOperand(memCopyOp.getSource(), 1);
+          return success();
+        }
+        if (auto loadOp = dyn_cast<PimMemCopyHostToDevOp>(&op)) {
+          if (!pim::isNormalizedCopyOp(loadOp)) {
+            loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+            hasFailure = true;
+          }
+          verifyOperand(loadOp.getDeviceTarget(), 2);
+          verifyOperand(loadOp.getHostSource(), 3);
+          return success();
+        }
+        if (auto storeOp = dyn_cast<PimMemCopyDevToHostOp>(&op)) {
+          if (!pim::isNormalizedCopyOp(storeOp)) {
+            storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+            hasFailure = true;
+          }
+          verifyOperand(storeOp.getHostTarget(), 2);
+          verifyOperand(storeOp.getDeviceSource(), 3);
+          return success();
+        }
+        if (auto sendOp = dyn_cast<PimSendOp>(&op)) {
+          verifyOperand(sendOp.getInput(), 0);
+          return success();
+        }
+        if (auto receiveOp = dyn_cast<PimReceiveOp>(&op)) {
+          verifyOperand(receiveOp.getOutputBuffer(), 0);
+          return success();
+        }
+        if (auto concatOp = dyn_cast<PimConcatOp>(&op)) {
+          verifyOperand(concatOp.getOutputBuffer(), 0);
+          for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs()))
+            verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1);
+          return success();
+        }
+        if (isa<PimTransposeOp,
+                PimVMMOp,
+                PimVVAddOp,
+                PimVVSubOp,
+                PimVVMulOp,
+                PimVVMaxOp,
+                PimVVDMulOp,
+                PimVAvgOp,
+                PimVReluOp,
+                PimVTanhOp,
+                PimVSigmOp,
+                PimVSoftmaxOp>(&op)) {
+          for (auto operandAndIndex : llvm::enumerate(op.getOperands())) {
+            if (auto vmmOp = dyn_cast<PimVMMOp>(&op); vmmOp && operandAndIndex.index() == 0)
+              continue;
+            verifyOperand(operandAndIndex.value(), operandAndIndex.index());
+          }
+        }
+        return success();
+      });
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) { verifyWithKnowledge(coreOp, seedCoreKnowledge(coreOp)); });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, 0);
+    verifyWithKnowledge(coreBatchOp, knowledge);
   });
 
   if (hasFailure) {
@@ -267,6 +392,31 @@ LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp mod
   return success();
 }
 
+LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp) const {
+  bool hasFailure = false;
+  auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
+    (void) walkPimCoreBlockStructurally(
+      coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
+        if (auto copyOp = dyn_cast<pim::PimMemCopyOp>(&op); copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
+          hasFailure = true;
+        if (auto copyOp = dyn_cast<pim::PimMemCopyHostToDevOp>(&op);
+            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
+          hasFailure = true;
+        if (auto copyOp = dyn_cast<pim::PimMemCopyDevToHostOp>(&op);
+            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
+          hasFailure = true;
+        return success();
+      });
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) { verifyWithKnowledge(coreOp, seedCoreKnowledge(coreOp)); });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, 0);
+    verifyWithKnowledge(coreBatchOp, knowledge);
+  });
+  return success(!hasFailure);
+}
+
 std::unique_ptr<Pass> createPimBufferizationPass() { return std::make_unique<PimBufferizationPass>(); }
 
 } // namespace onnx_mlir

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

@@ -96,8 +96,7 @@ struct FoldConstantCoreMapPattern final : OpRewritePattern<linalg::MapOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(linalg::MapOp mapOp, PatternRewriter& rewriter) const override {
-    auto coreOp = mapOp->getParentOfType<pim::PimCoreOp>();
-    if (!coreOp)
+    if (!mapOp->getParentOfType<pim::PimCoreOp>() && !mapOp->getParentOfType<pim::PimCoreBatchOp>())
       return failure();
 
     auto initType = dyn_cast<MemRefType>(mapOp.getInit().getType());
@@ -128,7 +127,7 @@ struct FoldConstantCoreMapPattern final : OpRewritePattern<linalg::MapOp> {
     auto sizeAttr = pim::getCheckedI32Attr(rewriter, mapOp, *sizeInBytes, "host constant folding byte size");
     if (failed(sizeAttr))
       return failure();
-    pim::PimMemCopyOp::create(
+    pim::PimMemCopyHostToDevOp::create(
       rewriter, mapOp.getLoc(), initType, zeroOffset, zeroOffset, mapOp.getInit(), getGlobalOp.getResult(), *sizeAttr);
     rewriter.eraseOp(mapOp);
     return success();

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

@@ -1,9 +0,0 @@
-add_pim_library(OMPimHostConstantMaterialization
-  MaterializeHostConstantsPass.cpp
-
-  EXCLUDE_FROM_OM_LIBS
-
-  LINK_LIBS PUBLIC
-  OMPimCommon
-  PimOps
-)