ReduceMean + resnet

remove unsupported tests

AGENTS.md

@@ -1,92 +1,210 @@
-- 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
+* 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
+
+# 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

@@ -168,8 +168,8 @@ Each validation run writes artifacts in the model workspace, for example under
 
 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.
+`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/memory_manager/mod.rs

@@ -1,3 +1,4 @@
+use std::cmp::min;
 use std::fmt::Debug;
 
 use anyhow::{Context, Result, bail, ensure};
@@ -86,7 +87,7 @@ where {
             size,
         };
         if self.memory.len() < address + size {
-            self.memory.resize((address + size) * 2, 0);
+            self.memory.resize(min((address + size) * 2, u32::MAX as usize), 0);
         }
         self.load_requests.push(load_request);
         Ok(self)

src/PIM/CMakeLists.txt

@@ -121,6 +121,8 @@ add_pim_library(OMPIMAccel
   OMSpatialToPim
   OMPimCommon
   OMPimBufferization
-  OMPimStaticMemoryCoalescing
+  OMPimMemoryCoalescing
+  OMPimHostConstantFolding
+  OMPimVerification
   MLIRTensorInferTypeOpInterfaceImpl
 )

src/PIM/Common/CMakeLists.txt

@@ -1,12 +1,15 @@
 add_pim_library(OMPimCommon
+  IR/AffineUtils.cpp
   IR/AddressAnalysis.cpp
   IR/BatchCoreUtils.cpp
   IR/ConstantUtils.cpp
   IR/CoreBlockUtils.cpp
   IR/EntryPointUtils.cpp
+  IR/LoopUtils.cpp
   IR/ShapeUtils.cpp
   IR/SubviewUtils.cpp
   IR/WeightUtils.cpp
+  Support/CheckedArithmetic.cpp
   Support/DebugDump.cpp
   Support/Diagnostics.cpp
   Support/FileSystemUtils.cpp
@@ -18,6 +21,8 @@ add_pim_library(OMPimCommon
   ${PIM_PUBLIC_INCLUDE_DIRS}
 
   LINK_LIBS PUBLIC
+  MLIRLinalgDialect
+  MLIRSCFDialect
   onnx
   SpatialOps
   PimOps

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -69,6 +69,16 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
 llvm::FailureOr<int64_t> resolveOpFoldResult(mlir::OpFoldResult ofr, const StaticValueKnowledge* knowledge);
 llvm::FailureOr<int64_t> resolveIndexValueImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
+static llvm::FailureOr<llvm::SmallVector<int64_t>> getStaticMemRefStrides(mlir::MemRefType type) {
+  llvm::SmallVector<int64_t> strides;
+  int64_t offset = 0;
+  if (failed(type.getStridesAndOffset(strides, offset)))
+    return mlir::failure();
+  if (llvm::any_of(strides, mlir::ShapedType::isDynamic))
+    return mlir::failure();
+  return strides;
+}
+
 static llvm::FailureOr<int64_t> resolveConstantGlobalLoad(mlir::memref::LoadOp loadOp,
                                                           const StaticValueKnowledge* knowledge) {
   auto getGlobalOp = loadOp.getMemRef().getDefiningOp<mlir::memref::GetGlobalOp>();
@@ -539,8 +549,10 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
       if (!isMemoryContiguous(sourceType.getShape(), offsets, sizes, strides))
         return mlir::failure();
 
-      auto sourceStrides = computeRowMajorStrides(sourceType.getShape());
-      byteOffset += linearizeIndex(offsets, sourceStrides) * getElementTypeSizeInBytes(subviewType.getElementType());
+      auto sourceStrides = getStaticMemRefStrides(sourceType);
+      if (failed(sourceStrides))
+        return mlir::failure();
+      byteOffset += linearizeIndex(offsets, *sourceStrides) * getElementTypeSizeInBytes(subviewType.getElementType());
       value = resolveAlias(subviewOp.getSource(), knowledge);
       continue;
     }
@@ -616,40 +628,51 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
       if (!sourceType || !subviewType || !sourceType.hasStaticShape() || !subviewType.hasStaticShape())
         return mlir::failure();
 
-      llvm::SmallVector<int64_t> staticOffsets;
-      staticOffsets.reserve(subviewOp.getMixedOffsets().size());
       llvm::SmallVector<int64_t> staticSizes;
       staticSizes.reserve(subviewOp.getMixedSizes().size());
       llvm::SmallVector<int64_t> staticStrides;
       staticStrides.reserve(subviewOp.getMixedStrides().size());
-      bool allStatic = true;
+      llvm::SmallVector<int64_t> staticOffsets;
+      staticOffsets.reserve(subviewOp.getMixedOffsets().size());
+      bool hasOnlyStaticOffsets = true;
 
       for (mlir::OpFoldResult offset : subviewOp.getMixedOffsets())
         if (auto attr = mlir::dyn_cast<mlir::Attribute>(offset))
           staticOffsets.push_back(mlir::cast<mlir::IntegerAttr>(attr).getInt());
         else
-          allStatic = false;
-      for (mlir::OpFoldResult size : subviewOp.getMixedSizes())
-        if (auto attr = mlir::dyn_cast<mlir::Attribute>(size))
-          staticSizes.push_back(mlir::cast<mlir::IntegerAttr>(attr).getInt());
-        else
-          allStatic = false;
-      for (mlir::OpFoldResult stride : subviewOp.getMixedStrides())
-        if (auto attr = mlir::dyn_cast<mlir::Attribute>(stride))
-          staticStrides.push_back(mlir::cast<mlir::IntegerAttr>(attr).getInt());
-        else
-          allStatic = false;
+          hasOnlyStaticOffsets = false;
+      for (mlir::OpFoldResult size : subviewOp.getMixedSizes()) {
+        auto attr = mlir::dyn_cast<mlir::Attribute>(size);
+        if (!attr)
+          return mlir::failure();
+        staticSizes.push_back(mlir::cast<mlir::IntegerAttr>(attr).getInt());
+      }
+      for (mlir::OpFoldResult stride : subviewOp.getMixedStrides()) {
+        auto attr = mlir::dyn_cast<mlir::Attribute>(stride);
+        if (!attr)
+          return mlir::failure();
+        staticStrides.push_back(mlir::cast<mlir::IntegerAttr>(attr).getInt());
+      }
 
-      if (allStatic) {
+      if (!isContiguousSubviewWithDynamicOffsets(
+            sourceType.getShape(), subviewOp.getMixedOffsets(), staticSizes, staticStrides)) {
+        return mlir::failure();
+      }
+
+      if (hasOnlyStaticOffsets) {
         if (!isMemoryContiguous(sourceType.getShape(), staticOffsets, staticSizes, staticStrides))
           return mlir::failure();
 
-        auto sourceStrides = computeRowMajorStrides(sourceType.getShape());
+        auto sourceStrides = getStaticMemRefStrides(sourceType);
+        if (failed(sourceStrides))
+          return mlir::failure();
         constantByteOffset +=
-          linearizeIndex(staticOffsets, sourceStrides) * getElementTypeSizeInBytes(subviewType.getElementType());
+          linearizeIndex(staticOffsets, *sourceStrides) * getElementTypeSizeInBytes(subviewType.getElementType());
       }
       else {
-        llvm::SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceType.getShape());
+        auto sourceStrides = getStaticMemRefStrides(sourceType);
+        if (failed(sourceStrides))
+          return mlir::failure();
         CompiledIndexExpr offsetExpr;
         {
           CompiledIndexExprNode expr;
@@ -658,7 +681,7 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
           offsetExpr = makeCompiledIndexExpr(std::move(expr));
         }
 
-        for (auto [mixedOffset, sourceStride] : llvm::zip_equal(subviewOp.getMixedOffsets(), sourceStrides)) {
+        for (auto [mixedOffset, sourceStride] : llvm::zip_equal(subviewOp.getMixedOffsets(), *sourceStrides)) {
           CompiledIndexExpr operandExpr;
           if (auto attr = mlir::dyn_cast<mlir::Attribute>(mixedOffset)) {
             CompiledIndexExprNode expr;
@@ -749,18 +772,12 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
 
 } // namespace
 
-llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value) { return resolveIndexValueImpl(value, nullptr); }
-
 llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge) {
   return resolveIndexValueImpl(value, &knowledge);
 }
 
 llvm::FailureOr<CompiledIndexExpr> compileIndexExpr(mlir::Value value) { return compileIndexValueImpl(value); }
 
-llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value) {
-  return resolveContiguousAddressImpl(value, nullptr);
-}
-
 llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value,
                                                                     const StaticValueKnowledge& knowledge) {
   return resolveContiguousAddressImpl(value, &knowledge);

src/PIM/Common/IR/AddressAnalysis.hpp

@@ -77,14 +77,12 @@ mlir::memref::GlobalOp lookupGlobalForGetGlobal(mlir::ModuleOp moduleOp, mlir::m
 
 /// Resolves a value to contiguous backing storage when that storage can be
 /// proven statically from aliases, DPS ties, casts, and subviews.
-llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value);
 llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value,
-                                                                    const StaticValueKnowledge& knowledge);
+                                                                    const StaticValueKnowledge& knowledge = {});
 
 /// Statically evaluates index-like SSA values, including simple integer
 /// arithmetic and loop facts recorded in `knowledge`.
-llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value);
-llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge);
+llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge = {});
 llvm::FailureOr<CompiledIndexExpr> compileIndexExpr(mlir::Value value);
 
 /// Follows alias, view, and DPS chains to recover the backing value of a

src/PIM/Common/IR/AffineUtils.cpp

@@ -0,0 +1,182 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/Matchers.h"
+
+#include "AffineUtils.hpp"
+#include "ConstantUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+static FailureOr<int64_t> floorDivSigned(int64_t lhs, int64_t rhs) {
+  if (rhs <= 0)
+    return failure();
+
+  int64_t quotient = lhs / rhs;
+  int64_t remainder = lhs % rhs;
+  if (remainder != 0 && lhs < 0)
+    --quotient;
+  return quotient;
+}
+
+static FailureOr<int64_t> ceilDivSigned(int64_t lhs, int64_t rhs) {
+  if (rhs <= 0)
+    return failure();
+
+  int64_t quotient = lhs / rhs;
+  int64_t remainder = lhs % rhs;
+  if (remainder != 0 && lhs > 0)
+    ++quotient;
+  return quotient;
+}
+
+Value createOrFoldAffineApply(
+  RewriterBase& rewriter, 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");
+
+  SmallVector<Attribute> operandConstants;
+  operandConstants.reserve(operands.size());
+  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));
+  }
+
+  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 affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
+}
+
+Value createOrFoldAffineApply(
+  RewriterBase& rewriter, 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);
+}
+
+Value affineMulConst(RewriterBase& rewriter, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) {
+  assert(constantAnchor && "expected a valid constant anchor");
+  if (multiplier == 0)
+    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+  if (multiplier == 1)
+    return value;
+
+  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  return createOrFoldAffineApply(rewriter, loc, d0 * multiplier, ValueRange {value}, constantAnchor);
+}
+
+Value affineModConst(RewriterBase& rewriter, 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);
+
+  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  return createOrFoldAffineApply(rewriter, loc, d0 % divisor, ValueRange {value}, constantAnchor);
+}
+
+Value affineFloorDivConst(
+  RewriterBase& rewriter, 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);
+}
+
+FailureOr<int64_t> evaluateAffineExpr(AffineExpr expr, ArrayRef<int64_t> dims, ArrayRef<int64_t> symbols) {
+  if (auto constant = dyn_cast<AffineConstantExpr>(expr))
+    return constant.getValue();
+  if (auto dim = dyn_cast<AffineDimExpr>(expr)) {
+    unsigned position = dim.getPosition();
+    if (position >= dims.size())
+      return failure();
+    return dims[position];
+  }
+  if (auto symbol = dyn_cast<AffineSymbolExpr>(expr)) {
+    unsigned position = symbol.getPosition();
+    if (position >= symbols.size())
+      return failure();
+    return symbols[position];
+  }
+
+  auto binary = dyn_cast<AffineBinaryOpExpr>(expr);
+  if (!binary)
+    return failure();
+
+  FailureOr<int64_t> lhs = evaluateAffineExpr(binary.getLHS(), dims, symbols);
+  FailureOr<int64_t> rhs = evaluateAffineExpr(binary.getRHS(), dims, symbols);
+  if (failed(lhs) || failed(rhs))
+    return failure();
+
+  switch (binary.getKind()) {
+  case AffineExprKind::Add:      return *lhs + *rhs;
+  case AffineExprKind::Mul:      return *lhs * *rhs;
+  case AffineExprKind::FloorDiv: return floorDivSigned(*lhs, *rhs);
+  case AffineExprKind::CeilDiv:  return ceilDivSigned(*lhs, *rhs);
+  case AffineExprKind::Mod:      {
+    FailureOr<int64_t> div = floorDivSigned(*lhs, *rhs);
+    if (failed(div))
+      return failure();
+    return *lhs - *div * *rhs;
+  }
+  default: return failure();
+  }
+}
+
+FailureOr<int64_t> evaluateSingleResultAffineMap(AffineMap map, ArrayRef<int64_t> operands) {
+  if (map.getNumResults() != 1 || operands.size() != map.getNumInputs())
+    return failure();
+
+  ArrayRef<int64_t> dims(operands.data(), map.getNumDims());
+  ArrayRef<int64_t> symbols(operands.data() + map.getNumDims(), map.getNumSymbols());
+  return evaluateAffineExpr(map.getResult(0), dims, symbols);
+}
+
+FailureOr<int64_t> evaluateAffineApply(affine::AffineApplyOp affineApply, IndexValueResolver resolver) {
+  SmallVector<int64_t, 4> operands;
+  operands.reserve(affineApply.getMapOperands().size());
+  for (Value operand : affineApply.getMapOperands()) {
+    FailureOr<int64_t> folded = resolver(operand);
+    if (failed(folded))
+      return failure();
+    operands.push_back(*folded);
+  }
+
+  return evaluateSingleResultAffineMap(affineApply.getAffineMap(), operands);
+}
+
+bool isSingleResultSymbolFreeAffineMap(AffineMap map) { return map.getNumResults() == 1 && map.getNumSymbols() == 0; }
+
+bool isDimAndConstantAffineExpr(AffineExpr expr) {
+  switch (expr.getKind()) {
+  case AffineExprKind::Constant:
+  case AffineExprKind::DimId:    return true;
+  case AffineExprKind::SymbolId: return false;
+  case AffineExprKind::Add:      {
+    auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
+    return isDimAndConstantAffineExpr(binaryExpr.getLHS()) && isDimAndConstantAffineExpr(binaryExpr.getRHS());
+  }
+  case AffineExprKind::Mul: {
+    auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
+    return (isa<AffineConstantExpr>(binaryExpr.getLHS()) && isDimAndConstantAffineExpr(binaryExpr.getRHS()))
+        || (isa<AffineConstantExpr>(binaryExpr.getRHS()) && isDimAndConstantAffineExpr(binaryExpr.getLHS()));
+  }
+  case AffineExprKind::FloorDiv:
+  case AffineExprKind::CeilDiv:
+  case AffineExprKind::Mod:      {
+    auto binaryExpr = cast<AffineBinaryOpExpr>(expr);
+    return isa<AffineConstantExpr>(binaryExpr.getRHS()) && isDimAndConstantAffineExpr(binaryExpr.getLHS());
+  }
+  }
+
+  llvm_unreachable("unexpected affine expression kind");
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/AffineUtils.hpp

@@ -0,0 +1,55 @@
+#pragma once
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/FunctionExtras.h"
+
+namespace onnx_mlir {
+
+using IndexValueResolver = llvm::function_ref<llvm::FailureOr<int64_t>(mlir::Value)>;
+
+mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter,
+                                    mlir::Location loc,
+                                    mlir::AffineMap map,
+                                    mlir::ValueRange operands,
+                                    mlir::Operation* constantAnchor);
+
+mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter,
+                                    mlir::Location loc,
+                                    mlir::AffineExpr expr,
+                                    mlir::ValueRange dims,
+                                    mlir::Operation* constantAnchor);
+
+mlir::Value affineMulConst(mlir::RewriterBase& rewriter,
+                           mlir::Location loc,
+                           mlir::Value value,
+                           int64_t multiplier,
+                           mlir::Operation* constantAnchor);
+
+mlir::Value affineModConst(mlir::RewriterBase& rewriter,
+                           mlir::Location loc,
+                           mlir::Value value,
+                           int64_t divisor,
+                           mlir::Operation* constantAnchor);
+
+mlir::Value affineFloorDivConst(mlir::RewriterBase& rewriter,
+                                mlir::Location loc,
+                                mlir::Value value,
+                                int64_t divisor,
+                                mlir::Operation* constantAnchor);
+
+llvm::FailureOr<int64_t>
+evaluateAffineExpr(mlir::AffineExpr expr, llvm::ArrayRef<int64_t> dims, llvm::ArrayRef<int64_t> symbols = {});
+
+llvm::FailureOr<int64_t> evaluateSingleResultAffineMap(mlir::AffineMap map, llvm::ArrayRef<int64_t> operands);
+
+llvm::FailureOr<int64_t> evaluateAffineApply(mlir::affine::AffineApplyOp affineApply, IndexValueResolver resolver);
+
+bool isSingleResultSymbolFreeAffineMap(mlir::AffineMap map);
+
+bool isDimAndConstantAffineExpr(mlir::AffineExpr expr);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/BatchCoreUtils.cpp

@@ -17,4 +17,16 @@ llvm::SmallVector<int32_t> getLaneChunkCoreIds(llvm::ArrayRef<int32_t> coreIds,
   return laneCoreIds;
 }
 
+bool isExplicitHostMemCopyOperand(mlir::Operation* op, unsigned operandIndex) {
+  if (mlir::isa<pim::PimMemCopyHostToDevOp>(op))
+    return operandIndex == 3;
+  if (mlir::isa<pim::PimMemCopyDevToHostOp>(op))
+    return operandIndex == 2;
+  return false;
+}
+
+bool isExplicitDevToHostTargetOperand(mlir::Operation* op, unsigned operandIndex) {
+  return mlir::isa<pim::PimMemCopyDevToHostOp>(op) && operandIndex == 2;
+}
+
 } // namespace onnx_mlir

src/PIM/Common/IR/BatchCoreUtils.hpp

@@ -11,4 +11,8 @@ llvm::SmallVector<int32_t> getBatchCoreIds(pim::PimCoreBatchOp coreBatchOp);
 
 llvm::SmallVector<int32_t> getLaneChunkCoreIds(llvm::ArrayRef<int32_t> coreIds, size_t laneCount, unsigned lane);
 
+bool isExplicitHostMemCopyOperand(mlir::Operation* op, unsigned operandIndex);
+
+bool isExplicitDevToHostTargetOperand(mlir::Operation* op, unsigned operandIndex);
+
 } // namespace onnx_mlir

src/PIM/Common/IR/ConstantUtils.cpp

@@ -1,36 +1,43 @@
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
-#include "mlir/IR/Matchers.h"
 
 #include "ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 
-Block* getHostConstantBlock(Operation* anchorOp) {
+static std::optional<int64_t> getIndexConstantValue(arith::ConstantOp constantOp) {
+  if (!constantOp.getType().isIndex())
+    return std::nullopt;
+
+  auto intAttr = dyn_cast<IntegerAttr>(constantOp.getValue());
+  if (!intAttr || !intAttr.getType().isIndex())
+    return std::nullopt;
+
+  return intAttr.getInt();
+}
+
+Block* getConstantInsertionBlock(Operation* anchorOp) {
   assert(anchorOp && "expected a valid anchor operation");
 
-  for (Operation* current = anchorOp; current; current = current->getParentOp())
-    if (isa<spatial::SpatCompute, spatial::SpatComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(current))
-      return current->getBlock();
-
+  if (auto funcOp = dyn_cast<func::FuncOp>(anchorOp))
+    return &funcOp.getBody().front();
   if (auto funcOp = anchorOp->getParentOfType<func::FuncOp>())
     return &funcOp.getBody().front();
+  if (auto moduleOp = dyn_cast<ModuleOp>(anchorOp))
+    return moduleOp.getBody();
   if (auto moduleOp = anchorOp->getParentOfType<ModuleOp>())
     return moduleOp.getBody();
   return anchorOp->getBlock();
 }
 
-Value getOrCreateHostConstant(Operation* anchorOp, Attribute value, Type type, OperationFolder& folder) {
+Value getOrCreateConstant(OperationFolder& folder, Operation* anchorOp, Attribute value, Type type) {
   assert(anchorOp && "expected a valid anchor operation");
-  Block* hostBlock = getHostConstantBlock(anchorOp);
+  Block* hostBlock = getConstantInsertionBlock(anchorOp);
   for (Operation& op : *hostBlock) {
     auto constantOp = dyn_cast<arith::ConstantOp>(&op);
     if (!constantOp || constantOp.getType() != type || constantOp.getValue() != value)
@@ -42,9 +49,9 @@ Value getOrCreateHostConstant(Operation* anchorOp, Attribute value, Type type, O
   return folder.getOrCreateConstant(hostBlock, arithDialect, value, type);
 }
 
-Value getOrCreateHostConstant(Operation* anchorOp, Attribute value, Type type, RewriterBase& rewriter) {
+Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute value, Type type) {
   assert(anchorOp && "expected a valid anchor operation");
-  Block* hostBlock = getHostConstantBlock(anchorOp);
+  Block* hostBlock = getConstantInsertionBlock(anchorOp);
   for (Operation& op : *hostBlock) {
     auto constantOp = dyn_cast<arith::ConstantOp>(&op);
     if (!constantOp || constantOp.getType() != type || constantOp.getValue() != value)
@@ -57,48 +64,94 @@ Value getOrCreateHostConstant(Operation* anchorOp, Attribute value, Type type, R
   return arith::ConstantOp::create(rewriter, anchorOp->getLoc(), type, cast<TypedAttr>(value)).getResult();
 }
 
-Value getOrCreateHostConstantLike(arith::ConstantOp constantOp, OperationFolder& folder) {
-  return getOrCreateHostConstant(constantOp.getOperation(), constantOp.getValue(), constantOp.getType(), folder);
+Value getOrCreateConstantLike(OperationFolder& folder, arith::ConstantOp constantOp) {
+  return getOrCreateConstant(folder, constantOp.getOperation(), constantOp.getValue(), constantOp.getType());
 }
 
-Value getOrCreateHostIndexConstant(Operation* anchorOp, int64_t value, OperationFolder& folder) {
+Value getOrCreateIndexConstant(OperationFolder& folder, Operation* anchorOp, int64_t value) {
   Builder builder(anchorOp->getContext());
-  return getOrCreateHostConstant(anchorOp, builder.getIndexAttr(value), builder.getIndexType(), folder);
+  return getOrCreateConstant(folder, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
 }
 
-Value getOrCreateHostIndexConstant(Operation* anchorOp, int64_t value, RewriterBase& rewriter) {
+Value getOrCreateIndexConstant(RewriterBase& rewriter, Operation* anchorOp, int64_t value) {
   Builder builder(anchorOp->getContext());
-  return getOrCreateHostConstant(anchorOp, builder.getIndexAttr(value), builder.getIndexType(), rewriter);
+  return getOrCreateConstant(rewriter, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
 }
 
-Value getOrCreateHostI32Constant(Operation* anchorOp, int32_t value, OperationFolder& folder) {
-  Builder builder(anchorOp->getContext());
-  return getOrCreateHostConstant(anchorOp, builder.getI32IntegerAttr(value), builder.getI32Type(), folder);
-}
+void hoistAndUniquifyIndexConstants(func::FuncOp funcOp, RewriterBase& rewriter) {
+  if (funcOp.getBody().empty())
+    return;
 
-Value getOrCreateHostI64Constant(Operation* anchorOp, int64_t value, OperationFolder& folder) {
-  Builder builder(anchorOp->getContext());
-  return getOrCreateHostConstant(anchorOp, builder.getI64IntegerAttr(value), builder.getI64Type(), folder);
-}
+  Block& entryBlock = funcOp.getBody().front();
+  DenseMap<int64_t, Value> canonicalByValue;
+  SmallVector<arith::ConstantOp> constants;
 
-Value createAffineApplyOrFoldedConstant(
-  RewriterBase& rewriter, Location loc, AffineMap map, ValueRange operands, Operation* anchorOp) {
-  SmallVector<Attribute> operandConstants;
-  operandConstants.reserve(operands.size());
-  for (Value operand : operands) {
-    APInt constantValue;
-    if (!matchPattern(operand, m_ConstantInt(&constantValue)))
-      return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
-    operandConstants.push_back(rewriter.getIndexAttr(constantValue.getSExtValue()));
+  funcOp.walk([&](arith::ConstantOp constantOp) {
+    if (!getIndexConstantValue(constantOp))
+      return;
+    constants.push_back(constantOp);
+  });
+
+  for (arith::ConstantOp constantOp : constants) {
+    auto value = getIndexConstantValue(constantOp);
+    if (!value || constantOp->getBlock() != &entryBlock)
+      continue;
+    canonicalByValue.try_emplace(*value, constantOp.getResult());
   }
 
-  SmallVector<Attribute> foldedResults;
-  if (succeeded(map.constantFold(operandConstants, foldedResults))) {
-    if (auto constantResult = dyn_cast<IntegerAttr>(foldedResults.front()))
-      return getOrCreateHostIndexConstant(anchorOp, constantResult.getInt(), rewriter);
+  for (arith::ConstantOp constantOp : constants) {
+    auto value = getIndexConstantValue(constantOp);
+    if (!value)
+      continue;
+
+    Value canonical = canonicalByValue.lookup(*value);
+    if (!canonical) {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPointToStart(&entryBlock);
+      Builder builder(funcOp.getContext());
+      canonical =
+        arith::ConstantOp::create(rewriter, constantOp.getLoc(), builder.getIndexType(), builder.getIndexAttr(*value));
+      canonicalByValue[*value] = canonical;
+    }
+
+    if (constantOp.getResult() == canonical)
+      continue;
+
+    constantOp.getResult().replaceAllUsesWith(canonical);
   }
 
-  return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
+  for (arith::ConstantOp constantOp : llvm::reverse(constants)) {
+    auto value = getIndexConstantValue(constantOp);
+    if (!value)
+      continue;
+    if (constantOp.getResult() == canonicalByValue.lookup(*value))
+      continue;
+    if (constantOp.use_empty())
+      rewriter.eraseOp(constantOp);
+  }
+}
+
+std::optional<int64_t> matchConstantIndexValue(Value value) {
+  if (!value || !value.getType().isIndex())
+    return std::nullopt;
+
+  if (auto constant = value.getDefiningOp<arith::ConstantIndexOp>())
+    return constant.value();
+
+  if (auto constant = value.getDefiningOp<arith::ConstantOp>())
+    if (auto intAttr = dyn_cast<IntegerAttr>(constant.getValue()); intAttr && intAttr.getType().isIndex())
+      return intAttr.getInt();
+
+  return std::nullopt;
+}
+
+std::optional<int64_t> matchConstantIndexValue(OpFoldResult value) {
+  if (auto attr = dyn_cast<Attribute>(value))
+    if (auto intAttr = dyn_cast<IntegerAttr>(attr); intAttr && intAttr.getType().isIndex())
+      return intAttr.getInt();
+  if (auto operand = dyn_cast<Value>(value))
+    return matchConstantIndexValue(operand);
+  return std::nullopt;
 }
 
 } // namespace onnx_mlir

src/PIM/Common/IR/ConstantUtils.hpp

@@ -1,39 +1,33 @@
 #pragma once
 
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/FoldUtils.h"
 
+#include <optional>
+
 namespace onnx_mlir {
 
-mlir::Block* getHostConstantBlock(mlir::Operation* anchorOp);
+mlir::Block* getConstantInsertionBlock(mlir::Operation* anchorOp);
 
-mlir::Value getOrCreateHostConstant(mlir::Operation* anchorOp,
-                                    mlir::Attribute value,
-                                    mlir::Type type,
-                                    mlir::OperationFolder& folder);
+mlir::Value
+getOrCreateConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
 
-mlir::Value getOrCreateHostConstant(mlir::Operation* anchorOp,
-                                    mlir::Attribute value,
-                                    mlir::Type type,
-                                    mlir::RewriterBase& rewriter);
+mlir::Value
+getOrCreateConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
 
-mlir::Value getOrCreateHostConstantLike(mlir::arith::ConstantOp constantOp, mlir::OperationFolder& folder);
+mlir::Value getOrCreateConstantLike(mlir::OperationFolder& folder, mlir::arith::ConstantOp constantOp);
 
-mlir::Value getOrCreateHostIndexConstant(mlir::Operation* anchorOp, int64_t value, mlir::OperationFolder& folder);
+mlir::Value getOrCreateIndexConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, int64_t value);
 
-mlir::Value getOrCreateHostIndexConstant(mlir::Operation* anchorOp, int64_t value, mlir::RewriterBase& rewriter);
+mlir::Value getOrCreateIndexConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, int64_t value);
 
-mlir::Value getOrCreateHostI32Constant(mlir::Operation* anchorOp, int32_t value, mlir::OperationFolder& folder);
+void hoistAndUniquifyIndexConstants(mlir::func::FuncOp funcOp, mlir::RewriterBase& rewriter);
 
-mlir::Value getOrCreateHostI64Constant(mlir::Operation* anchorOp, int64_t value, mlir::OperationFolder& folder);
+std::optional<int64_t> matchConstantIndexValue(mlir::Value value);
 
-mlir::Value createAffineApplyOrFoldedConstant(mlir::RewriterBase& rewriter,
-                                              mlir::Location loc,
-                                              mlir::AffineMap map,
-                                              mlir::ValueRange operands,
-                                              mlir::Operation* anchorOp);
+std::optional<int64_t> matchConstantIndexValue(mlir::OpFoldResult value);
 
 } // namespace onnx_mlir

src/PIM/Common/IR/LoopUtils.cpp

@@ -0,0 +1,96 @@
+#include "mlir/Dialect/SCF/IR/SCF.h"
+
+#include "llvm/Support/MathExtras.h"
+
+#include <optional>
+
+#include "ConstantUtils.hpp"
+#include "LoopUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+namespace {
+
+static std::optional<int64_t> getStaticTripCount(Value lowerBound, Value upperBound, Value step) {
+  auto lower = matchConstantIndexValue(lowerBound);
+  auto upper = matchConstantIndexValue(upperBound);
+  auto stepValue = matchConstantIndexValue(step);
+  if (!lower || !upper || !stepValue)
+    return std::nullopt;
+  if (*stepValue <= 0)
+    return std::nullopt;
+  if (*upper <= *lower)
+    return int64_t {0};
+  return llvm::divideCeil(*upper - *lower, *stepValue);
+}
+
+} // namespace
+
+static LogicalResult validateNormalizedLoopYields(Location loc, ValueRange initArgs, ArrayRef<Value> yieldedValues) {
+  if (yieldedValues.size() == initArgs.size())
+    return success();
+
+  emitError(loc) << "normalized loop body yielded " << yieldedValues.size() << " values for " << initArgs.size()
+                 << " iter args";
+  return failure();
+}
+
+FailureOr<NormalizedLoopResult> buildNormalizedScfFor(OpBuilder& builder,
+                                                      Location loc,
+                                                      Value lowerBound,
+                                                      Value upperBound,
+                                                      Value step,
+                                                      ValueRange initArgs,
+                                                      NormalizedLoopBodyBuilder bodyBuilder) {
+  NormalizedLoopResult result;
+
+  if (auto stepValue = matchConstantIndexValue(step); stepValue && *stepValue <= 0) {
+    emitError(loc) << "normalized scf.for requires a positive step, got " << *stepValue;
+    return failure();
+  }
+
+  if (auto tripCount = getStaticTripCount(lowerBound, upperBound, step)) {
+    if (*tripCount == 0) {
+      llvm::append_range(result.results, initArgs);
+      return result;
+    }
+
+    if (*tripCount == 1) {
+      result.inductionVar = lowerBound;
+      if (failed(bodyBuilder(builder, loc, lowerBound, initArgs, result.results)))
+        return failure();
+      if (failed(validateNormalizedLoopYields(loc, initArgs, result.results)))
+        return failure();
+      return result;
+    }
+  }
+
+  result.loop = scf::ForOp::create(builder, loc, lowerBound, upperBound, step, initArgs);
+  result.inductionVar = result.loop.getInductionVar();
+
+  {
+    OpBuilder::InsertionGuard guard(builder);
+    Block* body = result.loop.getBody();
+    if (!body->empty())
+      if (auto yieldOp = dyn_cast<scf::YieldOp>(body->back()))
+        yieldOp->erase();
+    builder.setInsertionPointToEnd(body);
+    ValueRange iterArgs = result.loop.getRegionIterArgs();
+    if (failed(bodyBuilder(builder, loc, result.inductionVar, iterArgs, result.results))) {
+      result.loop.erase();
+      return failure();
+    }
+    if (failed(validateNormalizedLoopYields(loc, initArgs, result.results))) {
+      result.loop.erase();
+      return failure();
+    }
+    scf::YieldOp::create(builder, loc, result.results);
+  }
+  builder.setInsertionPointAfter(result.loop);
+  result.results.assign(result.loop.getResults().begin(), result.loop.getResults().end());
+  return result;
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/LoopUtils.hpp

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Builders.h"
+
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace onnx_mlir {
+
+struct NormalizedLoopResult {
+  mlir::Value inductionVar;
+  llvm::SmallVector<mlir::Value, 4> results;
+  mlir::scf::ForOp loop;
+
+  bool wasInlined() const { return !loop; }
+};
+
+using NormalizedLoopBodyBuilder = llvm::function_ref<mlir::LogicalResult(
+  mlir::OpBuilder&, mlir::Location, mlir::Value, mlir::ValueRange, llvm::SmallVectorImpl<mlir::Value>&)>;
+
+mlir::FailureOr<NormalizedLoopResult> buildNormalizedScfFor(mlir::OpBuilder& builder,
+                                                            mlir::Location loc,
+                                                            mlir::Value lowerBound,
+                                                            mlir::Value upperBound,
+                                                            mlir::Value step,
+                                                            mlir::ValueRange initArgs,
+                                                            NormalizedLoopBodyBuilder bodyBuilder);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/ShapeUtils.cpp

@@ -111,4 +111,56 @@ bool isMemoryContiguous(llvm::ArrayRef<int64_t> srcShape,
   return true;
 }
 
+bool isContiguousSubviewWithDynamicOffsets(llvm::ArrayRef<int64_t> sourceShape,
+                                           llvm::ArrayRef<mlir::OpFoldResult> mixedOffsets,
+                                           llvm::ArrayRef<int64_t> staticSizes,
+                                           llvm::ArrayRef<int64_t> staticStrides) {
+  if (sourceShape.size() != mixedOffsets.size() || sourceShape.size() != staticSizes.size()
+      || sourceShape.size() != staticStrides.size()) {
+    return false;
+  }
+
+  if (llvm::any_of(staticStrides, [](int64_t stride) { return stride != 1; }))
+    return false;
+
+  auto reversedTriples =
+    llvm::zip_equal(llvm::reverse(sourceShape), llvm::reverse(mixedOffsets), llvm::reverse(staticSizes));
+
+  auto firstNonZeroOrDynamicOffset = llvm::find_if(reversedTriples, [](auto triple) {
+    auto [_sourceDim, offset, _size] = triple;
+    if (auto attr = mlir::dyn_cast<mlir::Attribute>(offset))
+      return mlir::cast<mlir::IntegerAttr>(attr).getInt() != 0;
+    return true;
+  });
+
+  if (firstNonZeroOrDynamicOffset != reversedTriples.end()) {
+    auto [sourceDim, offset, size] = *firstNonZeroOrDynamicOffset;
+    if (auto attr = mlir::dyn_cast<mlir::Attribute>(offset)) {
+      int64_t staticOffset = mlir::cast<mlir::IntegerAttr>(attr).getInt();
+      if (size > sourceDim - staticOffset)
+        return false;
+    }
+
+    ++firstNonZeroOrDynamicOffset;
+    for (auto it = firstNonZeroOrDynamicOffset; it != reversedTriples.end(); ++it)
+      if (std::get<2>(*it) != 1)
+        return false;
+  }
+
+  auto reversedSizes = llvm::zip_equal(llvm::reverse(sourceShape), llvm::reverse(staticSizes));
+  auto firstDifferentSize = llvm::find_if(reversedSizes, [](auto pair) {
+    auto [sourceDim, size] = pair;
+    return size != sourceDim;
+  });
+
+  if (firstDifferentSize != reversedSizes.end()) {
+    ++firstDifferentSize;
+    for (auto it = firstDifferentSize; it != reversedSizes.end(); ++it)
+      if (std::get<1>(*it) != 1)
+        return false;
+  }
+
+  return true;
+}
+
 } // namespace onnx_mlir

src/PIM/Common/IR/ShapeUtils.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Value.h"
 
 #include "llvm/ADT/ArrayRef.h"
@@ -30,4 +31,9 @@ bool isMemoryContiguous(llvm::ArrayRef<int64_t> srcShape,
                         llvm::ArrayRef<int64_t> sizes,
                         llvm::ArrayRef<int64_t> strides);
 
+bool isContiguousSubviewWithDynamicOffsets(llvm::ArrayRef<int64_t> sourceShape,
+                                           llvm::ArrayRef<mlir::OpFoldResult> mixedOffsets,
+                                           llvm::ArrayRef<int64_t> staticSizes,
+                                           llvm::ArrayRef<int64_t> staticStrides);
+
 } // namespace onnx_mlir

src/PIM/Common/IR/SubviewUtils.cpp

@@ -31,6 +31,19 @@ Value stripMemRefViewOps(Value value) {
   }
 }
 
+Value stripMemRefAddressingOps(Value value) {
+  while (true) {
+    if (auto subviewOp = value.getDefiningOp<memref::SubViewOp>()) {
+      value = subviewOp.getSource();
+      continue;
+    }
+    Value strippedValue = stripMemRefViewOps(value);
+    if (strippedValue == value)
+      return value;
+    value = strippedValue;
+  }
+}
+
 bool hasAllStaticSubviewParts(memref::SubViewOp subview) {
   return llvm::all_of(subview.getStaticOffsets(), [](int64_t value) { return !ShapedType::isDynamic(value); })
       && llvm::all_of(subview.getStaticSizes(), [](int64_t value) { return !ShapedType::isDynamic(value); })
@@ -81,4 +94,13 @@ FailureOr<SmallVector<int64_t>> getStaticSubviewOffsets(const StaticSubviewInfo&
   return staticOffsets;
 }
 
+bool isMemRefBaseAddressableValue(Value value) {
+  value = stripMemRefAddressingOps(value);
+  if (isa<BlockArgument>(value))
+    return true;
+
+  Operation* defOp = value.getDefiningOp();
+  return defOp && isa<memref::AllocOp, memref::GetGlobalOp>(defOp);
+}
+
 } // namespace onnx_mlir

src/PIM/Common/IR/SubviewUtils.hpp

@@ -20,6 +20,8 @@ mlir::Value stripMemRefCasts(mlir::Value value);
 
 mlir::Value stripMemRefViewOps(mlir::Value value);
 
+mlir::Value stripMemRefAddressingOps(mlir::Value value);
+
 bool hasAllStaticSubviewParts(mlir::memref::SubViewOp subview);
 
 llvm::FailureOr<StaticSubviewInfo> getStaticSubviewInfo(mlir::Value value);
@@ -27,4 +29,6 @@ llvm::FailureOr<StaticSubviewInfo> getStaticSubviewInfo(mlir::Value value);
 /// Returns the offsets in `info` as int64_t, failing if any offset is dynamic.
 llvm::FailureOr<llvm::SmallVector<int64_t>> getStaticSubviewOffsets(const StaticSubviewInfo& info);
 
+bool isMemRefBaseAddressableValue(mlir::Value value);
+
 } // namespace onnx_mlir

src/PIM/Common/IR/WeightUtils.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -46,26 +47,14 @@ CompiledIndexExpr mulExpr(CompiledIndexExpr lhs, int64_t rhs) {
   return makeBinaryExpr(CompiledIndexExprNode::Kind::Mul, std::move(lhs), makeConstantExpr(rhs));
 }
 
-mlir::Value stripWeightViewOps(mlir::Value value) {
-  while (true) {
-    if (auto subviewOp = value.getDefiningOp<mlir::memref::SubViewOp>()) {
-      value = subviewOp.getSource();
-      continue;
-    }
-    if (auto castOp = value.getDefiningOp<mlir::memref::CastOp>()) {
-      value = castOp.getSource();
-      continue;
-    }
-    if (auto collapseOp = value.getDefiningOp<mlir::memref::CollapseShapeOp>()) {
-      value = collapseOp.getSrc();
-      continue;
-    }
-    if (auto expandOp = value.getDefiningOp<mlir::memref::ExpandShapeOp>()) {
-      value = expandOp.getSrc();
-      continue;
-    }
-    return value;
-  }
+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>
@@ -131,8 +120,8 @@ bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value) {
         return expandShapeOp.getSrc() == currentValue && self(expandShapeOp.getResult(), self);
       if (auto collapseShapeOp = mlir::dyn_cast<mlir::tensor::CollapseShapeOp>(user))
         return collapseShapeOp.getSrc() == currentValue && self(collapseShapeOp.getResult(), self);
-      if (auto transposeOp = mlir::dyn_cast<mlir::ONNXTransposeOp>(user))
-        return transposeOp.getData() == currentValue && self(transposeOp.getResult(), self);
+      if (auto transposeOp = mlir::dyn_cast<mlir::linalg::TransposeOp>(user))
+        return transposeOp.getInput() == currentValue && self(transposeOp.getResult()[0], self);
 
       return false;
     });
@@ -158,7 +147,7 @@ void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir
 }
 
 std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, mlir::Value weight) {
-  weight = stripWeightViewOps(weight);
+  weight = stripMemRefAddressingOps(weight);
 
   if (auto coreOp = mlir::dyn_cast_or_null<pim::PimCoreOp>(weightOwner)) {
     for (unsigned weightIndex = 0; weightIndex < coreOp.getWeights().size(); ++weightIndex)
@@ -177,16 +166,17 @@ std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, mlir::V
   return std::nullopt;
 }
 
-std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, pim::PimVMMOp vmmOp) {
-  return resolveWeightIndex(weightOwner, vmmOp.getWeight());
-}
-
 llvm::FailureOr<ResolvedWeightView>
 resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const StaticValueKnowledge& knowledge) {
   llvm::SmallVector<mlir::Operation*> viewOps;
   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 {};
@@ -206,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);
@@ -227,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);
@@ -259,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;
       }
@@ -278,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/IR/WeightUtils.hpp

@@ -46,7 +46,6 @@ bool hasOnlySpatialMvmVmmWeightUses(mlir::Value value);
 void walkPimMvmVmmWeightUses(mlir::Operation* root, llvm::function_ref<void(mlir::OpOperand&)> callback);
 
 std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, mlir::Value weight);
-std::optional<unsigned> resolveWeightIndex(mlir::Operation* weightOwner, pim::PimVMMOp vmmOp);
 llvm::FailureOr<ResolvedWeightView>
 resolveWeightView(mlir::Operation* weightOwner, mlir::Value weight, const StaticValueKnowledge& knowledge = {});
 

src/PIM/Common/Support/CheckedArithmetic.cpp

@@ -0,0 +1,222 @@
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir::pim {
+
+namespace {
+
+static void emitCrashMessage(llvm::StringRef fieldName, llvm::StringRef message) {
+  llvm::errs() << "PIM " << fieldName << " " << message << "\n";
+}
+
+template <typename To, typename From>
+static FailureOr<To> checkedCastAtLocation(From value, Location loc, llvm::StringRef fieldName) {
+  static_assert(std::is_integral_v<To> && std::is_integral_v<From>, "checkedCastAtLocation requires integral types");
+
+  using ToLimits = std::numeric_limits<To>;
+
+  if constexpr (std::is_signed_v<From> == std::is_signed_v<To>) {
+    if (value < static_cast<From>(ToLimits::min()) || value > static_cast<From>(ToLimits::max())) {
+      emitCheckedArithmeticError(loc, fieldName, "is outside representable range");
+      return failure();
+    }
+  }
+  else if constexpr (std::is_signed_v<From>) {
+    using UnsignedFrom = std::make_unsigned_t<From>;
+    using UnsignedTo = std::make_unsigned_t<To>;
+    if (value < 0 || static_cast<UnsignedFrom>(value) > static_cast<UnsignedTo>(ToLimits::max())) {
+      emitCheckedArithmeticError(loc, fieldName, "is outside representable range");
+      return failure();
+    }
+  }
+  else {
+    using UnsignedFrom = std::make_unsigned_t<From>;
+    using UnsignedTo = std::conditional_t<std::is_signed_v<To>, std::make_unsigned_t<To>, To>;
+    if (static_cast<UnsignedFrom>(value) > static_cast<UnsignedTo>(ToLimits::max())) {
+      emitCheckedArithmeticError(loc, fieldName, "is outside representable range");
+      return failure();
+    }
+  }
+
+  return static_cast<To>(value);
+}
+
+template <typename UInt>
+FailureOr<UInt> checkedMulAtLocation(UInt lhs, UInt rhs, Location loc, llvm::StringRef fieldName) {
+  static_assert(std::is_integral_v<UInt> && std::is_unsigned_v<UInt>,
+                "checkedMulAtLocation requires unsigned integral types");
+  if (lhs != 0 && rhs > std::numeric_limits<UInt>::max() / lhs) {
+    emitCheckedArithmeticError(loc, fieldName, "multiplication overflow");
+    return failure();
+  }
+  return lhs * rhs;
+}
+
+} // namespace
+
+InFlightDiagnostic emitCheckedArithmeticError(Operation* anchor, llvm::StringRef fieldName, llvm::StringRef message) {
+  assert(anchor && "expected arithmetic diagnostics to have an anchor op");
+  return anchor->emitOpError() << fieldName << " " << message;
+}
+
+InFlightDiagnostic emitCheckedArithmeticError(Location loc, llvm::StringRef fieldName, llvm::StringRef message) {
+  return emitError(loc) << "PIM " << fieldName << " " << message;
+}
+
+FailureOr<int32_t> checkedI32(int64_t value, Operation* anchor, llvm::StringRef fieldName) {
+  return checkedCast<int32_t>(value, anchor, fieldName);
+}
+
+FailureOr<int32_t> checkedI32(uint64_t value, Operation* anchor, llvm::StringRef fieldName) {
+  return checkedCast<int32_t>(value, anchor, fieldName);
+}
+
+FailureOr<uint8_t> checkedU8(uint64_t value, Operation* anchor, llvm::StringRef fieldName) {
+  return checkedCast<uint8_t>(value, anchor, fieldName);
+}
+
+FailureOr<size_t> checkedSize(int64_t value, Operation* anchor, llvm::StringRef fieldName) {
+  return checkedCast<size_t>(value, anchor, fieldName);
+}
+
+FailureOr<IntegerAttr>
+getCheckedI32Attr(Builder& builder, Operation* anchor, int64_t value, llvm::StringRef fieldName) {
+  assert(anchor && "checked op-based attrs require a non-null diagnostic anchor");
+  auto checkedValue = checkedI32(value, anchor, fieldName);
+  if (failed(checkedValue))
+    return failure();
+  return builder.getI32IntegerAttr(*checkedValue);
+}
+
+FailureOr<IntegerAttr>
+getCheckedI32Attr(Builder& builder, Operation* anchor, uint64_t value, llvm::StringRef fieldName) {
+  assert(anchor && "checked op-based attrs require a non-null diagnostic anchor");
+  auto checkedValue = checkedI32(value, anchor, fieldName);
+  if (failed(checkedValue))
+    return failure();
+  return builder.getI32IntegerAttr(*checkedValue);
+}
+
+FailureOr<IntegerAttr> getCheckedI32Attr(Builder& builder, Location loc, int64_t value, llvm::StringRef fieldName) {
+  auto checkedValue = checkedCastAtLocation<int32_t>(value, loc, fieldName);
+  if (failed(checkedValue))
+    return failure();
+  return builder.getI32IntegerAttr(*checkedValue);
+}
+
+FailureOr<IntegerAttr> getCheckedI32Attr(Builder& builder, Location loc, uint64_t value, llvm::StringRef fieldName) {
+  auto checkedValue = checkedCastAtLocation<int32_t>(value, loc, fieldName);
+  if (failed(checkedValue))
+    return failure();
+  return builder.getI32IntegerAttr(*checkedValue);
+}
+
+FailureOr<uint64_t> getCheckedShapedTypeSizeInBytes(ShapedType type, Operation* anchor, llvm::StringRef fieldName) {
+  assert(anchor && "checked op-based size helpers require a non-null diagnostic anchor");
+  if (!type.hasStaticShape()) {
+    emitCheckedArithmeticError(anchor, fieldName, "requires static shaped type");
+    return failure();
+  }
+  if (!hasByteSizedElementType(type.getElementType())) {
+    emitCheckedArithmeticError(anchor, fieldName, "requires byte-sized element type");
+    return failure();
+  }
+
+  uint64_t elements = 1;
+  for (int64_t dim : type.getShape()) {
+    if (dim < 0) {
+      emitCheckedArithmeticError(anchor, fieldName, "requires nonnegative dimensions");
+      return failure();
+    }
+
+    auto nextElements = checkedMul(elements, static_cast<uint64_t>(dim), anchor, fieldName);
+    if (failed(nextElements))
+      return failure();
+    elements = *nextElements;
+  }
+
+  return checkedMul(
+    elements, static_cast<uint64_t>(getElementTypeSizeInBytes(type.getElementType())), anchor, fieldName);
+}
+
+FailureOr<uint64_t> getCheckedShapedTypeSizeInBytes(ShapedType type, Location loc, llvm::StringRef fieldName) {
+  if (!type.hasStaticShape()) {
+    emitCheckedArithmeticError(loc, fieldName, "requires static shaped type");
+    return failure();
+  }
+  if (!hasByteSizedElementType(type.getElementType())) {
+    emitCheckedArithmeticError(loc, fieldName, "requires byte-sized element type");
+    return failure();
+  }
+
+  uint64_t elements = 1;
+  for (int64_t dim : type.getShape()) {
+    if (dim < 0) {
+      emitCheckedArithmeticError(loc, fieldName, "requires nonnegative dimensions");
+      return failure();
+    }
+
+    auto nextElements = checkedMulAtLocation(elements, static_cast<uint64_t>(dim), loc, fieldName);
+    if (failed(nextElements))
+      return failure();
+    elements = *nextElements;
+  }
+
+  return checkedMulAtLocation(
+    elements, static_cast<uint64_t>(getElementTypeSizeInBytes(type.getElementType())), loc, fieldName);
+}
+
+int32_t checkedI32OrCrash(int64_t value, llvm::StringRef fieldName) {
+  if (value < std::numeric_limits<int32_t>::min() || value > std::numeric_limits<int32_t>::max()) {
+    emitCrashMessage(fieldName, "is outside representable range");
+    llvm_unreachable("PIM checked arithmetic failure");
+  }
+  return static_cast<int32_t>(value);
+}
+
+int32_t checkedI32OrCrash(uint64_t value, llvm::StringRef fieldName) {
+  if (value > static_cast<uint64_t>(std::numeric_limits<int32_t>::max())) {
+    emitCrashMessage(fieldName, "is outside representable range");
+    llvm_unreachable("PIM checked arithmetic failure");
+  }
+  return static_cast<int32_t>(value);
+}
+
+uint8_t checkedU8OrCrash(uint64_t value, llvm::StringRef fieldName) {
+  if (value > static_cast<uint64_t>(std::numeric_limits<uint8_t>::max())) {
+    emitCrashMessage(fieldName, "is outside representable range");
+    llvm_unreachable("PIM checked arithmetic failure");
+  }
+  return static_cast<uint8_t>(value);
+}
+
+size_t checkedSizeOrCrash(int64_t value, llvm::StringRef fieldName) {
+  if (value < 0) {
+    emitCrashMessage(fieldName, "is outside representable range");
+    llvm_unreachable("PIM checked arithmetic failure");
+  }
+  return static_cast<size_t>(value);
+}
+
+size_t checkedAddOrCrash(size_t lhs, size_t rhs, llvm::StringRef fieldName) {
+  if (rhs > std::numeric_limits<size_t>::max() - lhs) {
+    emitCrashMessage(fieldName, "addition overflow");
+    llvm_unreachable("PIM checked arithmetic failure");
+  }
+  return lhs + rhs;
+}
+
+size_t checkedMulOrCrash(size_t lhs, size_t rhs, llvm::StringRef fieldName) {
+  if (lhs != 0 && rhs > std::numeric_limits<size_t>::max() / lhs) {
+    emitCrashMessage(fieldName, "multiplication overflow");
+    llvm_unreachable("PIM checked arithmetic failure");
+  }
+  return lhs * rhs;
+}
+
+} // namespace onnx_mlir::pim

src/PIM/Common/Support/CheckedArithmetic.hpp

@@ -0,0 +1,107 @@
+#pragma once
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/StringRef.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+
+namespace onnx_mlir::pim {
+
+mlir::InFlightDiagnostic
+emitCheckedArithmeticError(mlir::Operation* anchor, llvm::StringRef fieldName, llvm::StringRef message);
+
+mlir::InFlightDiagnostic
+emitCheckedArithmeticError(mlir::Location loc, llvm::StringRef fieldName, llvm::StringRef message);
+
+template <typename To, typename From>
+mlir::FailureOr<To> checkedCast(From value, mlir::Operation* anchor, llvm::StringRef fieldName) {
+  static_assert(std::is_integral_v<To> && std::is_integral_v<From>, "checkedCast requires integral types");
+
+  using ToLimits = std::numeric_limits<To>;
+
+  if constexpr (std::is_signed_v<From> == std::is_signed_v<To>) {
+    if (value < static_cast<From>(ToLimits::min()) || value > static_cast<From>(ToLimits::max())) {
+      emitCheckedArithmeticError(anchor, fieldName, "is outside representable range");
+      return mlir::failure();
+    }
+  }
+  else if constexpr (std::is_signed_v<From>) {
+    using UnsignedFrom = std::make_unsigned_t<From>;
+    using UnsignedTo = std::make_unsigned_t<To>;
+    if (value < 0 || static_cast<UnsignedFrom>(value) > static_cast<UnsignedTo>(ToLimits::max())) {
+      emitCheckedArithmeticError(anchor, fieldName, "is outside representable range");
+      return mlir::failure();
+    }
+  }
+  else {
+    using UnsignedFrom = std::make_unsigned_t<From>;
+    using UnsignedTo = std::conditional_t<std::is_signed_v<To>, std::make_unsigned_t<To>, To>;
+    if (static_cast<UnsignedFrom>(value) > static_cast<UnsignedTo>(ToLimits::max())) {
+      emitCheckedArithmeticError(anchor, fieldName, "is outside representable range");
+      return mlir::failure();
+    }
+  }
+
+  return static_cast<To>(value);
+}
+
+template <typename UInt>
+mlir::FailureOr<UInt> checkedAdd(UInt lhs, UInt rhs, mlir::Operation* anchor, llvm::StringRef fieldName) {
+  static_assert(std::is_integral_v<UInt> && std::is_unsigned_v<UInt>, "checkedAdd requires unsigned integral types");
+  if (rhs > std::numeric_limits<UInt>::max() - lhs) {
+    emitCheckedArithmeticError(anchor, fieldName, "addition overflow");
+    return mlir::failure();
+  }
+  return lhs + rhs;
+}
+
+template <typename UInt>
+mlir::FailureOr<UInt> checkedMul(UInt lhs, UInt rhs, mlir::Operation* anchor, llvm::StringRef fieldName) {
+  static_assert(std::is_integral_v<UInt> && std::is_unsigned_v<UInt>, "checkedMul requires unsigned integral types");
+  if (lhs != 0 && rhs > std::numeric_limits<UInt>::max() / lhs) {
+    emitCheckedArithmeticError(anchor, fieldName, "multiplication overflow");
+    return mlir::failure();
+  }
+  return lhs * rhs;
+}
+
+mlir::FailureOr<int32_t> checkedI32(int64_t value, mlir::Operation* anchor, llvm::StringRef fieldName);
+mlir::FailureOr<int32_t> checkedI32(uint64_t value, mlir::Operation* anchor, llvm::StringRef fieldName);
+
+mlir::FailureOr<uint8_t> checkedU8(uint64_t value, mlir::Operation* anchor, llvm::StringRef fieldName);
+
+mlir::FailureOr<size_t> checkedSize(int64_t value, mlir::Operation* anchor, llvm::StringRef fieldName);
+
+mlir::FailureOr<mlir::IntegerAttr>
+getCheckedI32Attr(mlir::Builder& builder, mlir::Operation* anchor, int64_t value, llvm::StringRef fieldName);
+
+mlir::FailureOr<mlir::IntegerAttr>
+getCheckedI32Attr(mlir::Builder& builder, mlir::Operation* anchor, uint64_t value, llvm::StringRef fieldName);
+
+mlir::FailureOr<mlir::IntegerAttr>
+getCheckedI32Attr(mlir::Builder& builder, mlir::Location loc, int64_t value, llvm::StringRef fieldName);
+
+mlir::FailureOr<mlir::IntegerAttr>
+getCheckedI32Attr(mlir::Builder& builder, mlir::Location loc, uint64_t value, llvm::StringRef fieldName);
+
+mlir::FailureOr<uint64_t>
+getCheckedShapedTypeSizeInBytes(mlir::ShapedType type, mlir::Operation* anchor, llvm::StringRef fieldName);
+
+mlir::FailureOr<uint64_t>
+getCheckedShapedTypeSizeInBytes(mlir::ShapedType type, mlir::Location loc, llvm::StringRef fieldName);
+
+int32_t checkedI32OrCrash(int64_t value, llvm::StringRef fieldName);
+int32_t checkedI32OrCrash(uint64_t value, llvm::StringRef fieldName);
+uint8_t checkedU8OrCrash(uint64_t value, llvm::StringRef fieldName);
+size_t checkedSizeOrCrash(int64_t value, llvm::StringRef fieldName);
+size_t checkedAddOrCrash(size_t lhs, size_t rhs, llvm::StringRef fieldName);
+size_t checkedMulOrCrash(size_t lhs, size_t rhs, llvm::StringRef fieldName);
+
+} // namespace onnx_mlir::pim

src/PIM/Common/Support/DebugDump.cpp

@@ -18,7 +18,7 @@ void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
   std::fstream file(dialectsDir + "/" + name + ".mlir", std::ios::out);
   llvm::raw_os_ostream os(file);
   mlir::OpPrintingFlags flags;
-  flags.elideLargeElementsAttrs();
+  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
   moduleOp.print(os, flags);
   os.flush();
   file.close();

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,16 +5,18 @@
 
 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::string formatReportMemory(uint64_t bytes) {
   const char* units[] = {"B", "KB", "MB", "GB", "TB", "PB", "EB"};
   int i = 0;

src/PIM/Common/Support/ReportUtils.hpp

@@ -11,6 +11,7 @@
 namespace onnx_mlir {
 
 std::fstream openReportFile(const std::string& name);
+std::fstream openReportFileWithExtension(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
@@ -28,7 +29,9 @@ add_pim_library(OMPimCompilerUtils
   OMPimCompilerOptions
   OMPimCommon
   OMPimBufferization
-  OMPimStaticMemoryCoalescing
+  OMPimMemoryCoalescing
+  OMPimHostConstantFolding
+  OMPimVerification
   OMPimPasses
   OMONNXToSpatial
   OMSpatialToPim

src/PIM/Compiler/PimBinaryFormat.hpp

@@ -6,8 +6,8 @@
 #include "llvm/Support/raw_ostream.h"
 
 #include <array>
-#include <cassert>
-#include <limits>
+
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
 namespace onnx_mlir::pim_binary {
 
@@ -95,15 +95,10 @@ inline void writeInstructionRecord(llvm::raw_ostream& os, const InstructionRecor
   writeInt32LE(os, record.generic3);
 }
 
-inline int32_t toI32(int64_t value) {
-  assert(value >= std::numeric_limits<int32_t>::min() && value <= std::numeric_limits<int32_t>::max()
-         && "PIM binary field out of int32 range");
-  return static_cast<int32_t>(value);
-}
+inline int32_t toI32(int64_t value) { return onnx_mlir::pim::checkedI32OrCrash(value, "binary field"); }
 
 inline uint8_t toU8(int64_t value) {
-  assert(value >= 0 && value <= std::numeric_limits<uint8_t>::max() && "PIM binary field out of uint8 range");
-  return static_cast<uint8_t>(value);
+  return onnx_mlir::pim::checkedU8OrCrash(static_cast<uint64_t>(value), "binary field");
 }
 
 inline int32_t getOptionalInt(const llvm::json::Object& object, llvm::StringRef key, int32_t defaultValue = 0) {

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"
@@ -25,20 +24,26 @@
 #include <cassert>
 #include <cstdint>
 #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"
 
@@ -71,32 +76,159 @@ static MemoryValueKey getMemoryValueKey(mlir::Value value, std::optional<unsigne
   return {value, getLaneForMemoryValue(value, lane)};
 }
 
+static bool isInsidePimCoreLikeOp(memref::AllocOp allocOp) {
+  return allocOp->getParentOfType<pim::PimCoreOp>() || allocOp->getParentOfType<pim::PimCoreBatchOp>();
+}
+
+static MemoryReportKind classifyMemoryReportKind(mlir::Value value) {
+  if (isa<mlir::BlockArgument>(value))
+    return MemoryReportKind::Input;
+  if (auto* op = value.getDefiningOp()) {
+    if (isa<memref::AllocOp>(op))
+      return MemoryReportKind::Alloca;
+    if (isa<memref::GetGlobalOp>(op))
+      return MemoryReportKind::Global;
+  }
+  return MemoryReportKind::None;
+}
+
+static int32_t getVectorByteSizeOrCrash(ShapedType type) {
+  auto byteSize = pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "vector byte size");
+  if (failed(byteSize))
+    llvm_unreachable("Failed to compute checked vector byte size");
+  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) {
   auto type = cast<ShapedType>(value.getType());
   assert("Only static shape is supported" && type.hasStaticShape());
-  size_t allocSize = getShapedTypeSizeInBytes(type);
-  MemEntry memEntry = {0, allocSize};
-  return &memEntries.emplace_back(memEntry, getMemoryValueKey(value, lane)).first;
+  auto checkedAllocSize =
+    pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "memory allocation byte size");
+  if (failed(checkedAllocSize))
+    llvm_unreachable("Failed to compute checked allocation byte size");
+  PendingMemEntry pending;
+  pending.memEntry = {0, *checkedAllocSize};
+  pending.key = getMemoryValueKey(value, lane);
+  pending.reportKind = classifyMemoryReportKind(value);
+  return &memEntries.emplace_back(std::move(pending)).memEntry;
 }
 
 void PimMemory::allocateGatheredMemory() {
-  llvm::sort(memEntries, [](auto a, auto b) -> bool { return a.first.size > b.first.size; });
-  for (auto& [memEntry, key] : memEntries)
-    allocateMemoryForValue(key, memEntry);
+  llvm::sort(memEntries, [](const PendingMemEntry& lhs, const PendingMemEntry& rhs) {
+    return lhs.memEntry.size > rhs.memEntry.size;
+  });
+  for (PendingMemEntry& pending : memEntries)
+    allocateMemoryForValue(pending.key, pending.memEntry, pending.reportKind);
   memEntries.clear();
 }
 
-void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry) {
+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: break;
+  case MemoryReportKind::Global:
+    ++reportRow.numGlobal;
+    reportRow.sizeGlobal += memEntry.size;
+    break;
+  case MemoryReportKind::Input:
+  case MemoryReportKind::None:  break;
+  }
+}
+
+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) {
@@ -127,7 +259,7 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
   });
 
   funcOp.walk([&](memref::AllocOp allocOp) {
-    if (!allocOp->getParentOfType<pim::PimCoreOp>())
+    if (!isInsidePimCoreLikeOp(allocOp))
       gatherMemEntry(allocOp.getResult());
   });
 
@@ -138,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) {
@@ -181,20 +375,11 @@ static MemoryReportRow addMemoryReportRows(const MemoryReportRow& lhs, const Mem
 }
 
 MemoryReportRow PimMemory::getReportRow() const {
-  MemoryReportRow row;
-  for (auto& [key, memEntry] : ownedMemEntriesMap) {
-    if (auto op = key.value.getDefiningOp()) {
-      if (isa<memref::AllocOp>(op)) {
-        row.numAlloca++;
-        row.sizeAlloca += memEntry.size;
-      }
-
-      if (isa<memref::GetGlobalOp>(op)) {
-        row.numGlobal++;
-        row.sizeGlobal += memEntry.size;
-      }
-    }
-  }
+  MemoryReportRow row = reportRow;
+  row.numAlloca = localPhysicalSlots.size();
+  row.sizeAlloca = 0;
+  for (const PhysicalSlotInfo& slot : localPhysicalSlots)
+    row.sizeAlloca += slot.size;
   return row;
 }
 
@@ -235,6 +420,7 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value 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");
     }
@@ -245,6 +431,7 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
   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";
@@ -270,7 +457,8 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
     llvm_unreachable("Missing mem entry");
   }
 
-  return iter->second.address + resolvedAddress->byteOffset;
+  size_t byteOffset = pim::checkedSizeOrCrash(resolvedAddress->byteOffset, "resolved PIM byte offset");
+  return pim::checkedAddOrCrash(iter->second.address, byteOffset, "resolved PIM address");
 }
 
 llvm::FailureOr<int64_t> PimAcceleratorMemory::getIndexValue(mlir::Value value,
@@ -289,8 +477,12 @@ llvm::FailureOr<int64_t> PimAcceleratorMemory::getIndexValue(mlir::Value value,
 void PimAcceleratorMemory::reportHost() { hostReportRow = hostMem.getReportRow(); }
 
 void PimAcceleratorMemory::recordCoreReport(size_t coreId, const MemoryReportRow& row) {
-  reportEntries.push_back(
-    {MemoryReportEntry::Kind::Core, coreId, {static_cast<int32_t>(coreId)}, row, row.numAlloca, row.sizeAlloca});
+  reportEntries.push_back({MemoryReportEntry::Kind::Core,
+                           coreId,
+                           {pim::checkedI32OrCrash(coreId, "memory report core id")},
+                           row,
+                           row.numAlloca,
+                           row.sizeAlloca});
 }
 
 void PimAcceleratorMemory::recordBatchReport(uint64_t batchId,
@@ -314,13 +506,15 @@ void PimAcceleratorMemory::flushReport() {
 
   llvm::raw_os_ostream os(fileReport);
   uint64_t totalGlobalMemory = hostReportRow.has_value() ? hostReportRow->sizeGlobal : 0;
+  uint64_t totalWeightsMemory = totalWeightBytes;
   uint64_t totalCoresMemory = 0;
   for (const MemoryReportEntry& entry : reportEntries)
     totalCoresMemory += entry.totalAllocaBytes;
 
-  llvm::SmallVector<ReportField, 2> totalFields = {
-    {"Global memory", formatReportMemory(totalGlobalMemory)},
-    {"Cores memory",  formatReportMemory(totalCoresMemory) }
+  llvm::SmallVector<ReportField, 3> totalFields = {
+    {"Global memory",  formatReportMemory(totalGlobalMemory) },
+    {"Weights memory", formatReportMemory(totalWeightsMemory)},
+    {"Cores memory",   formatReportMemory(totalCoresMemory)  }
   };
   printReportTotalsBlock(os, totalFields);
 
@@ -400,24 +594,24 @@ void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t i
   pim_binary::InstructionRecord instruction;
   instruction.opcode = pim_binary::Opcode::sldi;
   instruction.rd = static_cast<uint8_t>(registerNumber);
-  instruction.r2OrImm = static_cast<int32_t>(immediate);
+  instruction.r2OrImm = pim::checkedI32OrCrash(immediate, "register immediate");
   emitInstruction(instruction);
 }
 
 void PimCodeGen::setupRd(size_t rdAddress, size_t rdOffset) const {
-  genSetRegisterImmediateUnsigned(0, rdAddress + rdOffset);
+  genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
 }
 
 void PimCodeGen::setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const {
-  genSetRegisterImmediateUnsigned(0, rdAddress + rdOffset);
-  genSetRegisterImmediateUnsigned(1, rs1Address + rs1Offset);
+  genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
+  genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
 }
 
 void PimCodeGen::setupRdRs1Rs2(
   size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const {
-  genSetRegisterImmediateUnsigned(0, rdAddress + rdOffset);
-  genSetRegisterImmediateUnsigned(1, rs1Address + rs1Offset);
-  genSetRegisterImmediateUnsigned(2, rs2Address + rs2Offset);
+  genSetRegisterImmediateUnsigned(0, pim::checkedAddOrCrash(rdAddress, rdOffset, "rd address"));
+  genSetRegisterImmediateUnsigned(1, pim::checkedAddOrCrash(rs1Address, rs1Offset, "rs1 address"));
+  genSetRegisterImmediateUnsigned(2, pim::checkedAddOrCrash(rs2Address, rs2Offset, "rs2 address"));
 }
 
 void PimCodeGen::emitMemCopyOp(StringRef opName,
@@ -435,8 +629,7 @@ void PimCodeGen::emitMemCopyOp(StringRef opName,
   instruction.r1 = 1;
   instruction.generic1 = 0;
   instruction.generic2 = 0;
-  instruction.generic3 = static_cast<int32_t>(size);
-  (void) sizeFieldName;
+  instruction.generic3 = pim::checkedI32OrCrash(size, sizeFieldName);
   emitInstruction(instruction);
 }
 
@@ -446,10 +639,10 @@ void PimCodeGen::emitCommunicationOp(StringRef opName, size_t bufferAddr, size_t
   pim_binary::InstructionRecord instruction;
   instruction.opcode = pim_binary::opcodeFromString(opName);
   instruction.rd = 0;
-  instruction.r2OrImm = static_cast<int32_t>(remapCoreId(coreId));
+  instruction.r2OrImm = pim::checkedI32OrCrash(remapCoreId(coreId), "communication core id");
   instruction.generic1 = 0;
   instruction.generic2 = 0;
-  instruction.generic3 = static_cast<int32_t>(size);
+  instruction.generic3 = pim::checkedI32OrCrash(size, "communication byte size");
   emitInstruction(instruction);
 }
 
@@ -462,7 +655,7 @@ void PimCodeGen::emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_
   instruction.r1 = 1;
   instruction.r2OrImm = 8;
   instruction.generic1 = 0;
-  instruction.generic2 = static_cast<int32_t>(groupId);
+  instruction.generic2 = pim::checkedI32OrCrash(groupId, "mvm group id");
   emitInstruction(instruction);
 }
 
@@ -479,16 +672,6 @@ void PimCodeGen::codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticVa
                 loadOp.getSize());
 }
 
-void PimCodeGen::codeGenLoadBatchOp(pim::PimMemCopyHostToDevBatchOp loadOp,
-                                    const StaticValueKnowledge& knowledge) const {
-  emitMemCopyOp("ld",
-                addressOf(loadOp.getDeviceTarget(), knowledge),
-                loadOp.getDeviceTargetOffset(),
-                addressOf(loadOp.getHostSource(), knowledge),
-                loadOp.getHostSourceOffset(),
-                loadOp.getSize());
-}
-
 void PimCodeGen::codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const {
   auto hostTargetOffset = indexOf(storeOp.getHostTargetOffset(), knowledge);
   auto deviceSourceOffset = indexOf(storeOp.getDeviceSourceOffset(), knowledge);
@@ -503,11 +686,15 @@ void PimCodeGen::codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const Static
 }
 
 void PimCodeGen::codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledge& knowledge) const {
+  auto targetOffset = indexOf(lmvOp.getTargetOffset(), knowledge);
+  auto sourceOffset = indexOf(lmvOp.getSourceOffset(), knowledge);
+  assert(succeeded(targetOffset) && succeeded(sourceOffset)
+         && "pim.memcp offsets must be statically resolvable during codegen");
   emitMemCopyOp("lmv",
                 addressOf(lmvOp.getTarget(), knowledge),
-                lmvOp.getTargetOffset(),
+                *targetOffset,
                 addressOf(lmvOp.getSource(), knowledge),
-                lmvOp.getSourceOffset(),
+                *sourceOffset,
                 lmvOp.getSize(),
                 "len");
 }
@@ -582,7 +769,7 @@ void PimCodeGen::codeGenVVAddOp(pim::PimVVAddOp vvaddOp, const StaticValueKnowle
   instruction.rd = 0;
   instruction.r1 = 1;
   instruction.r2OrImm = 2;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vvaddOp.getLhs().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vvaddOp.getLhs().getType()));
   emitInstruction(instruction);
 }
 
@@ -597,7 +784,7 @@ void PimCodeGen::codeGenVVSubOp(pim::PimVVSubOp vvsubOp, const StaticValueKnowle
   instruction.rd = 0;
   instruction.r1 = 1;
   instruction.r2OrImm = 2;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vvsubOp.getLhs().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vvsubOp.getLhs().getType()));
   emitInstruction(instruction);
 }
 
@@ -612,7 +799,7 @@ void PimCodeGen::codeGenVVMulOp(pim::PimVVMulOp vvmulOp, const StaticValueKnowle
   instruction.rd = 0;
   instruction.r1 = 1;
   instruction.r2OrImm = 2;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vvmulOp.getLhs().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vvmulOp.getLhs().getType()));
   emitInstruction(instruction);
 }
 
@@ -627,7 +814,7 @@ void PimCodeGen::codeGenVVMaxOp(pim::PimVVMaxOp vvmaxOp, const StaticValueKnowle
   instruction.rd = 0;
   instruction.r1 = 1;
   instruction.r2OrImm = 2;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vvmaxOp.getLhs().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vvmaxOp.getLhs().getType()));
   emitInstruction(instruction);
 }
 
@@ -642,7 +829,7 @@ void PimCodeGen::codeGenVVDMulOp(pim::PimVVDMulOp vvdmulOp, const StaticValueKno
   instruction.rd = 0;
   instruction.r1 = 1;
   instruction.r2OrImm = 2;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vvdmulOp.getLhs().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vvdmulOp.getLhs().getType()));
   emitInstruction(instruction);
 }
 
@@ -657,7 +844,7 @@ void PimCodeGen::codeGenVAvgOp(pim::PimVAvgOp vavgOp, const StaticValueKnowledge
   instruction.r1 = 1;
   instruction.r2OrImm = 1;
   instruction.generic1 = 1;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vavgOp.getInput().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vavgOp.getInput().getType()));
   emitInstruction(instruction);
 }
 
@@ -670,7 +857,7 @@ void PimCodeGen::codeGenVReluOp(pim::PimVReluOp vreluOp, const StaticValueKnowle
   instruction.opcode = pim_binary::Opcode::vrelu;
   instruction.rd = 0;
   instruction.r1 = 1;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vreluOp.getInput().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vreluOp.getInput().getType()));
   emitInstruction(instruction);
 }
 
@@ -683,7 +870,7 @@ void PimCodeGen::codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowle
   instruction.opcode = pim_binary::Opcode::vtanh;
   instruction.rd = 0;
   instruction.r1 = 1;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vtanhOp.getInput().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vtanhOp.getInput().getType()));
   emitInstruction(instruction);
 }
 
@@ -696,7 +883,7 @@ void PimCodeGen::codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowle
   instruction.opcode = pim_binary::Opcode::vsigm;
   instruction.rd = 0;
   instruction.r1 = 1;
-  instruction.generic3 = static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vsigmOp.getInput().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vsigmOp.getInput().getType()));
   emitInstruction(instruction);
 }
 
@@ -709,8 +896,7 @@ void PimCodeGen::codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticVa
   instruction.opcode = pim_binary::Opcode::vsoftmax;
   instruction.rd = 0;
   instruction.r1 = 1;
-  instruction.generic3 =
-    static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(vsoftmaxOp.getInput().getType())));
+  instruction.generic3 = getVectorByteSizeOrCrash(cast<ShapedType>(vsoftmaxOp.getInput().getType()));
   emitInstruction(instruction);
 }
 
@@ -811,11 +997,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;
@@ -848,7 +1067,6 @@ public:
 
 enum class CompiledCoreOpKind : uint8_t {
   Load,
-  LoadBatch,
   Store,
   Lmv,
   Receive,
@@ -887,8 +1105,6 @@ struct CompiledCoreNode {
 static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
   if (isa<pim::PimMemCopyHostToDevOp>(op))
     return CompiledCoreOpKind::Load;
-  if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
-    return CompiledCoreOpKind::LoadBatch;
   if (isa<pim::PimMemCopyDevToHostOp>(op))
     return CompiledCoreOpKind::Store;
   if (isa<pim::PimMemCopyOp>(op))
@@ -1027,9 +1243,6 @@ static LogicalResult executeCompiledCorePlan(
     case CompiledCoreOpKind::Load:
       coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);
       break;
-    case CompiledCoreOpKind::LoadBatch:
-      coreCodeGen.codeGenLoadBatchOp(cast<pim::PimMemCopyHostToDevBatchOp>(node.op), knowledge);
-      break;
     case CompiledCoreOpKind::Store:
       coreCodeGen.codeGenStoreOp(cast<pim::PimMemCopyDevToHostOp>(node.op), knowledge);
       break;
@@ -1213,17 +1426,18 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
   auto linkCoreWeights =
     [&](size_t coreId, ArrayRef<std::string> weightFiles, json::Array& xbarsPerGroup) -> OnnxMlirCompilerErrorCodes {
     auto coreWeightsDirPath = outputDirPath + "/core_" + std::to_string(coreId);
-    if (auto error = sys::fs::create_directory(coreWeightsDirPath)) {
+    if (auto error = sys::fs::create_directory(coreWeightsDirPath); error && error != std::errc::file_exists) {
       errs() << "Error creating core directory: " << coreWeightsDirPath << ": " << error.message() << '\n';
       return InvalidOutputFileAccess;
     }
 
     for (auto [slot, fileName] : llvm::enumerate(weightFiles)) {
       xbarsPerGroup.push_back(static_cast<int64_t>(slot));
-      if (auto error = sys::fs::create_link(outputDirPath + "/weights/" + fileName,
-                                            coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin")) {
-        errs() << "Error creating link file: " << (outputDirPath + "/weights/" + fileName) << " to "
-               << (coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin") << "\nError:" << error.message()
+      std::string sourcePath = outputDirPath + "/weights/" + fileName;
+      std::string targetPath = coreWeightsDirPath + "/crossbar_" + std::to_string(slot) + ".bin";
+      sys::fs::remove(targetPath);
+      if (auto error = sys::fs::create_link(sourcePath, targetPath)) {
+        errs() << "Error creating link file: " << sourcePath << " to " << targetPath << "\nError:" << error.message()
                << '\n';
         return InvalidOutputFileAccess;
       }
@@ -1241,7 +1455,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();
       }
 
@@ -1282,15 +1509,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);
@@ -1298,10 +1526,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);
@@ -1316,11 +1541,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());
@@ -1339,6 +1564,21 @@ 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;
@@ -1351,7 +1591,21 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     request.weights = jobResults[jobIndex].usedWeights;
     weightRequests.push_back(std::move(request));
   }
-  auto mapCoreWeightToFileName = createAndPopulateWeightFolder(weightRequests, outputDirPath);
+  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];
@@ -1363,6 +1617,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;
     }
   }
@@ -1379,7 +1635,7 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
       if (auto err = linkCoreWeights(job.emittedCoreId, mapCoreWeightToFileName[job.emittedCoreId], xbarsPerGroup))
         return err;
       xbarsPerArrayGroup["core" + std::to_string(job.emittedCoreId)] = std::move(xbarsPerGroup);
-      reportedCoreIds.push_back(static_cast<int32_t>(job.emittedCoreId));
+      reportedCoreIds.push_back(pim::checkedI32OrCrash(job.emittedCoreId, "batch report core id"));
       if (!batchPerCoreRow)
         batchPerCoreRow = result.reportRow;
       batchRow = addMemoryReportRows(batchRow, result.reportRow);
@@ -1391,10 +1647,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,14 @@
 #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 <fstream>
 #include <limits>
 #include <optional>
+#include <string>
 
 #include "onnx-mlir/Compiler/OMCompilerTypes.h"
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
@@ -26,6 +28,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;
@@ -45,6 +57,19 @@ struct MemoryReportRow {
   }
 };
 
+enum class MemoryReportKind {
+  None,
+  Alloca,
+  Global,
+  Input
+};
+
+struct PendingMemEntry {
+  MemEntry memEntry;
+  MemoryValueKey key;
+  MemoryReportKind reportKind = MemoryReportKind::None;
+};
+
 struct MemoryReportEntry {
   enum class Kind {
     Core,
@@ -60,16 +85,21 @@ struct MemoryReportEntry {
 };
 
 class PimMemory {
-  llvm::SmallVector<std::pair<MemEntry, MemoryValueKey>, 32> memEntries;
+  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);
+  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)
@@ -78,6 +108,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; }
@@ -94,6 +125,7 @@ private:
   std::fstream fileReport;
   std::optional<MemoryReportRow> hostReportRow;
   llvm::SmallVector<MemoryReportEntry, 32> reportEntries;
+  uint64_t totalWeightBytes = 0;
   mutable llvm::DenseMap<mlir::Value, CompiledIndexExpr> compiledIndexExprs;
   mutable llvm::DenseMap<mlir::Value, CompiledAddressExpr> compiledAddressExprs;
 
@@ -118,6 +150,7 @@ public:
                          const MemoryReportRow& perCoreRow,
                          uint64_t totalAllocaCount,
                          uint64_t totalAllocaBytes);
+  void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; }
   void flushReport();
   void clean(mlir::Operation* op);
 };
@@ -175,7 +208,6 @@ public:
   }
 
   void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const;
-  void codeGenLoadBatchOp(pim::PimMemCopyHostToDevBatchOp loadOp, const StaticValueKnowledge& knowledge) const;
   void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const;
   void codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledge& knowledge) const;
 

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -22,12 +22,28 @@ 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<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),

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -24,11 +24,19 @@ typedef enum {
   MergeSchedulerPeft = 0,
 } PimMergeSchedulerType;
 
+typedef enum {
+  PimMemoryReportNone = 0,
+  PimMemoryReportSummary = 1,
+  PimMemoryReportFull = 2,
+} PimMemoryReportLevel;
+
 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<bool> pimOnlyCodegen;
+extern llvm::cl::opt<bool> pimDisableMemoryCoalescing;
 extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<bool> pimEmitJson;
 

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -40,14 +40,14 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
 
   if (pimEmissionTarget >= EmitPimBufferized) {
     pm.addPass(createPimBufferizationPass());
-    pm.addPass(createPimStaticMemoryCoalescingPass());
     pm.addPass(createMessagePass("Pim bufferized"));
   }
 
   if (pimEmissionTarget >= EmitPimCodegen) {
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
-    pm.addPass(createPimMaterializeHostConstantsPass());
+    if (!pimDisableMemoryCoalescing)
+      pm.addPass(createPimMemoryCoalescingPass());
     pm.addPass(createPimVerificationPass());
     pm.addPass(createMessagePass("Pim verified"));
     pm.addPass(createEmitPimCodePass());

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -0,0 +1,733 @@
+#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());
+        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

@@ -7,6 +7,7 @@
 
 #include <cassert>
 
+#include "Common/Support/CheckedArithmetic.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -18,15 +19,14 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {} // namespace
 
-llvm::DenseMap<size_t, llvm::SmallVector<std::string, 8>>
-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");
   size_t indexFileName = 0;
 
   int64_t xbarSize = crossbarSize.getValue();
-  llvm::DenseMap<size_t, llvm::SmallVector<std::string, 8>> mapCoreWeightToFileName;
+  WeightEmissionResult result;
   llvm::SmallVector<std::pair<ResolvedWeightView, std::string>, 16> materializedWeights;
 
   auto materializeWeight = [&](const ResolvedWeightView& weightView) -> std::string {
@@ -72,17 +72,22 @@ createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> requests, StringRef ou
 
     weightFileStream.close();
     materializedWeights.push_back({weightView, newFileName});
+    uint64_t weightBytes = pim::checkedMulOrCrash(
+      pim::checkedMulOrCrash(static_cast<size_t>(xbarSize), static_cast<size_t>(xbarSize), "weight element count"),
+      elementByteWidth,
+      "weight byte size");
+    result.totalWeightBytes = pim::checkedAddOrCrash(result.totalWeightBytes, weightBytes, "total weight bytes");
     return newFileName;
   };
 
   for (const WeightFileRequest& request : requests) {
-    auto& coreFiles = mapCoreWeightToFileName[request.coreId];
+    auto& coreFiles = result.mapCoreWeightToFileName[request.coreId];
     coreFiles.reserve(request.weights.size());
     for (const ResolvedWeightView& weight : request.weights)
       coreFiles.push_back(materializeWeight(weight));
   }
 
-  return mapCoreWeightToFileName;
+  return result;
 }
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimWeightEmitter.hpp

@@ -6,6 +6,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 
+#include <cstdint>
 #include <string>
 
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
@@ -17,7 +18,12 @@ struct WeightFileRequest {
   llvm::SmallVector<ResolvedWeightView, 8> weights;
 };
 
-llvm::DenseMap<size_t, llvm::SmallVector<std::string, 8>>
-createAndPopulateWeightFolder(llvm::ArrayRef<WeightFileRequest> requests, llvm::StringRef outputDirPath);
+struct WeightEmissionResult {
+  llvm::DenseMap<size_t, llvm::SmallVector<std::string, 8>> mapCoreWeightToFileName;
+  uint64_t totalWeightBytes = 0;
+};
+
+WeightEmissionResult createAndPopulateWeightFolder(llvm::ArrayRef<WeightFileRequest> requests,
+                                                   llvm::StringRef outputDirPath);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -3,11 +3,12 @@ mlir_tablegen(ONNXToSpatial.hpp.inc -gen-rewriters "-I${ONNX_MLIR_SRC_ROOT}")
 add_public_tablegen_target(ONNXToSpatialIncGen)
 
 add_pim_library(OMONNXToSpatial
-  ConversionPatterns.cpp
+  Patterns.cpp
   CompileTime.cpp
   ONNXToSpatialVerifier.cpp
-  PrePatterns.cpp
-  PostPatterns.cpp
+  Patterns/Pre.cpp
+  Patterns/Post.cpp
+  Patterns/GeneratedConversion.cpp
   Patterns/Math/Conv.cpp
   Patterns/Math/Elementwise.cpp
   Patterns/Math/Gemm.cpp
@@ -21,9 +22,13 @@ add_pim_library(OMONNXToSpatial
   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
+  Common/AttributeUtils.cpp
   Common/ComputeRegionBuilder.cpp
+  Common/IndexingUtils.cpp
   Common/ShapeTilingUtils.cpp
   Common/WeightMaterialization.cpp
 
@@ -33,6 +38,7 @@ add_pim_library(OMONNXToSpatial
   ONNXToSpatialIncGen
 
   LINK_LIBS PUBLIC
+  MLIRLinalgDialect
   MLIRSCFDialect
   MLIRTosaDialect
   OMCompilerOptions

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

@@ -0,0 +1,23 @@
+#include "mlir/IR/BuiltinAttributes.h"
+
+#include "AttributeUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+int64_t getI64Attr(ArrayAttr attr, size_t index) { return cast<IntegerAttr>(attr[index]).getInt(); }
+
+int64_t getOptionalI64Attr(std::optional<ArrayAttr> attr, size_t index, int64_t defaultValue) {
+  return attr ? getI64Attr(*attr, index) : defaultValue;
+}
+
+llvm::SmallVector<int64_t> getI64ArrayAttrValues(ArrayAttr attr) {
+  llvm::SmallVector<int64_t> values;
+  values.reserve(attr.size());
+  for (Attribute value : attr)
+    values.push_back(cast<IntegerAttr>(value).getInt());
+  return values;
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstddef>
+#include <optional>
+
+namespace onnx_mlir {
+
+int64_t getI64Attr(mlir::ArrayAttr attr, size_t index);
+
+int64_t getOptionalI64Attr(std::optional<mlir::ArrayAttr> attr, size_t index, int64_t defaultValue);
+
+llvm::SmallVector<int64_t> getI64ArrayAttrValues(mlir::ArrayAttr attr);
+
+} // namespace onnx_mlir

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

@@ -1,6 +1,8 @@
 #pragma once
 
+#include "AttributeUtils.hpp"
 #include "ComputeRegionBuilder.hpp"
+#include "IndexingUtils.hpp"
 #include "ShapeTilingUtils.hpp"
 #include "WeightMaterialization.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"

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

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ValueRange.h"
@@ -7,9 +8,12 @@
 
 #include <cassert>
 #include <cstddef>
+#include <limits>
 #include <type_traits>
 #include <utility>
 
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 namespace onnx_mlir {
@@ -49,6 +53,13 @@ using InvokeWithBlockArgsResultT = typename InvokeWithBlockArgsResult<Fn, Seq>::
 template <typename Fn>
 using InvokeWithValueRangeResultT = std::invoke_result_t<Fn, mlir::ValueRange>;
 
+struct SpatComputeBatchBodyArgs {
+  mlir::Value lane;
+  mlir::ValueRange weights;
+  mlir::ValueRange inputs;
+  mlir::ValueRange outputs;
+};
+
 } // namespace detail
 
 template <typename RewriterT>
@@ -159,6 +170,98 @@ auto createSpatCompute(RewriterT& rewriter,
   }
 }
 
+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) {
+  if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
+    return mlir::FailureOr<spatial::SpatComputeBatch>(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());
+
+  auto batchOp = spatial::SpatComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
+
+  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);
+  }
+
+  auto* block =
+    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
+  rewriter.setInsertionPointToStart(block);
+
+  detail::SpatComputeBatchBodyArgs args {
+    block->getArgument(0),
+    mlir::ValueRange(block->getArguments()).slice(1, weights.size()),
+    mlir::ValueRange(block->getArguments()).slice(1 + weights.size(), inputs.size()),
+    mlir::ValueRange(block->getArguments()).drop_front(1 + weights.size() + inputs.size())};
+
+  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);
+  }
+  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);
+    return mlir::FailureOr<spatial::SpatComputeBatch>(batchOp);
+  }
+}
+
+inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewriter,
+                                                     mlir::Location loc,
+                                                     mlir::Value source,
+                                                     mlir::Value dest,
+                                                     mlir::ArrayRef<mlir::OpFoldResult> offsets,
+                                                     mlir::ArrayRef<mlir::OpFoldResult> sizes,
+                                                     mlir::ArrayRef<mlir::OpFoldResult> strides) {
+  auto inParallelOp = spatial::SpatInParallelOp::create(rewriter, loc);
+  rewriter.setInsertionPointToStart(&inParallelOp.getRegion().front());
+  mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
+}
+
+template <typename BodyFn>
+mlir::Value materializeOrComputeUnary(mlir::Value input,
+                                      mlir::RankedTensorType resultType,
+                                      mlir::PatternRewriter& rewriter,
+                                      mlir::Location loc,
+                                      BodyFn&& build) {
+  auto&& buildFn = build;
+  if (isCompileTimeComputable(input))
+    return buildFn(input);
+
+  auto computeOp = createSpatCompute<1>(
+    rewriter, loc, mlir::TypeRange {resultType}, {}, mlir::ValueRange {input}, [&](mlir::Value computeInput) {
+      mlir::Value result = buildFn(computeInput);
+      spatial::SpatYieldOp::create(rewriter, loc, result);
+    });
+  return computeOp.getResult(0);
+}
+
 mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::ConversionPatternRewriter& rewriter);
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,45 @@
+#include <algorithm>
+
+#include "IndexingUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+int64_t normalizeAxis(int64_t axis, int64_t rank) { return axis >= 0 ? axis : rank + axis; }
+
+FailureOr<int64_t> normalizeAxisChecked(int64_t axis, int64_t rank) {
+  int64_t normalizedAxis = normalizeAxis(axis, rank);
+  if (normalizedAxis < 0 || normalizedAxis >= rank)
+    return failure();
+  return normalizedAxis;
+}
+
+int64_t normalizeIndex(int64_t index, int64_t dimSize) { return index >= 0 ? index : dimSize + index; }
+
+static SmallVector<int64_t> normalizeAxesImpl(std::optional<ArrayAttr> axesAttr, int64_t rank) {
+  SmallVector<int64_t> normalizedAxes;
+  if (!axesAttr) {
+    normalizedAxes.reserve(rank);
+    for (int64_t axis = 0; axis < rank; ++axis)
+      normalizedAxes.push_back(axis);
+  }
+  else {
+    normalizedAxes.reserve(axesAttr->size());
+    for (Attribute attr : *axesAttr)
+      normalizedAxes.push_back(normalizeAxis(cast<IntegerAttr>(attr).getInt(), rank));
+    llvm::sort(normalizedAxes);
+    normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end());
+  }
+  return normalizedAxes;
+}
+
+FailureOr<SmallVector<int64_t>> normalizeAxesChecked(std::optional<ArrayAttr> axesAttr, int64_t rank) {
+  SmallVector<int64_t> normalizedAxes = normalizeAxesImpl(axesAttr, rank);
+  for (int64_t axis : normalizedAxes)
+    if (axis < 0 || axis >= rank)
+      return failure();
+  return normalizedAxes;
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <optional>
+
+namespace onnx_mlir {
+
+int64_t normalizeAxis(int64_t axis, int64_t rank);
+
+mlir::FailureOr<int64_t> normalizeAxisChecked(int64_t axis, int64_t rank);
+
+int64_t normalizeIndex(int64_t index, int64_t dimSize);
+
+mlir::FailureOr<llvm::SmallVector<int64_t>> normalizeAxesChecked(std::optional<mlir::ArrayAttr> axesAttr, int64_t rank);
+
+} // namespace onnx_mlir

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

@@ -1,85 +1,85 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/Matchers.h"
 
 #include "llvm/ADT/SmallVector.h"
 
-#include <algorithm>
+#include <functional>
 
+#include "IndexingUtils.hpp"
 #include "ShapeTilingUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 
-static Value getIndexValue(OpFoldResult result, ConversionPatternRewriter& rewriter, Location loc) {
-  if (auto attr = dyn_cast<Attribute>(result))
-    return arith::ConstantIndexOp::create(rewriter, loc, cast<IntegerAttr>(attr).getInt()).getResult();
-  return cast<Value>(result);
+bool hasStaticPositiveShape(ArrayRef<int64_t> shape) {
+  return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
 }
 
-static Value addIndexValues(Value lhs, Value rhs, ConversionPatternRewriter& rewriter, Location loc) {
-  APInt lhsConst;
-  if (matchPattern(lhs, m_ConstantInt(&lhsConst)) && lhsConst.isZero())
-    return rhs;
-
-  APInt rhsConst;
-  if (matchPattern(rhs, m_ConstantInt(&rhsConst)) && rhsConst.isZero())
-    return lhs;
-
-  return arith::AddIOp::create(rewriter, loc, lhs, rhs).getResult();
+bool hasStaticPositiveShape(RankedTensorType type) {
+  return type.hasStaticShape() && hasStaticPositiveShape(type.getShape());
 }
 
-static Value multiplyIndexValue(Value value, OpFoldResult factor, ConversionPatternRewriter& rewriter, Location loc) {
-  APInt factorConst;
-  if (auto attr = dyn_cast<Attribute>(factor))
-    factorConst = cast<IntegerAttr>(attr).getValue();
-  else if (!matchPattern(cast<Value>(factor), m_ConstantInt(&factorConst)))
-    return arith::MulIOp::create(rewriter, loc, value, cast<Value>(factor)).getResult();
-
-  if (factorConst.isZero())
-    return arith::ConstantIndexOp::create(rewriter, loc, 0).getResult();
-  if (factorConst.isOne())
-    return value;
-
-  auto factorValue = arith::ConstantIndexOp::create(rewriter, loc, factorConst.getSExtValue()).getResult();
-  return arith::MulIOp::create(rewriter, loc, value, factorValue).getResult();
+int64_t getStaticShapeElementCount(ArrayRef<int64_t> shape) {
+  return std::accumulate(shape.begin(), shape.end(), int64_t {1}, std::multiplies<int64_t> {});
 }
 
-static bool isContiguousTensorSlice(Value source, RankedTensorType resultType, ArrayRef<OpFoldResult> strides) {
-  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
-  if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || sourceType.getRank() != resultType.getRank())
-    return false;
+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;
+}
 
-  for (OpFoldResult stride : strides) {
-    APInt strideValue;
-    if (auto attr = dyn_cast<Attribute>(stride)) {
-      if (cast<IntegerAttr>(attr).getInt() != 1)
-        return false;
-      continue;
-    }
-    if (!matchPattern(cast<Value>(stride), m_ConstantInt(&strideValue)) || !strideValue.isOne())
-      return false;
+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;
   }
 
-  auto sizesAndShape = llvm::zip_equal(llvm::make_range(resultType.getShape().rbegin(), resultType.getShape().rend()),
-                                       llvm::make_range(sourceType.getShape().rbegin(), sourceType.getShape().rend()));
-  auto firstDifferentSize = std::find_if(sizesAndShape.begin(), sizesAndShape.end(), [&](auto sizeAndShape) -> bool {
-    auto [size, dimension] = sizeAndShape;
-    return size != dimension;
-  });
-  if (firstDifferentSize == sizesAndShape.end())
-    return true;
+  if (static_cast<int64_t>(permAttr->size()) != rank)
+    return failure();
 
-  ++firstDifferentSize;
-  return std::all_of(firstDifferentSize, sizesAndShape.end(), [](auto sizeAndShape) {
-    auto [size, _dimension] = sizeAndShape;
-    return size == 1;
-  });
+  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(
@@ -88,11 +88,8 @@ SmallVector<Value> sliceTensor(
   assert("Invalid axis" && axis < shape.size());
 
   SmallVector<OpFoldResult> strides(shape.size(), rewriter.getIndexAttr(1));
-  SmallVector<OpFoldResult> offsets(shape.size(), rewriter.getIndexAttr(0));
-  SmallVector<OpFoldResult> sizes;
-  sizes.reserve(shape.size());
-  for (const auto size : shape)
-    sizes.push_back(rewriter.getIndexAttr(size));
+  SmallVector<OpFoldResult> offsets = getZeroOffsets(rewriter, shape.size());
+  SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, shape);
   sizes[axis] = rewriter.getIndexAttr(sliceSize);
 
   long length = shape[axis];
@@ -150,130 +147,33 @@ sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, ConversionPatternRewri
   return slicesPerCore;
 }
 
-DenseMap<HSliceId, DenseMap<CoreId, SmallVector<Value>>> tileMatrix(
-  Value& matrixToTile, int64_t hSliceSize, int64_t vSliceSize, ConversionPatternRewriter& rewriter, Location& loc) {
-  assert("Not a matrix" && isMatrixShape(getTensorShape(matrixToTile)));
+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());
 
-  DenseMap<HSliceId, DenseMap<CoreId, SmallVector<Value>>> tiles;
-
-  SmallVector<Value> hSlices = sliceTensor(matrixToTile, 1, hSliceSize, rewriter, loc);
-  size_t numHSlices = hSlices.size();
-  for (size_t hSliceId = 0; hSliceId < numHSlices; hSliceId++) {
-    Value hSlice = hSlices[hSliceId];
-    SmallVector<Value> vSlices = sliceTensor(hSlice, 0, vSliceSize, rewriter, loc);
-    for (size_t vSliceId = 0; vSliceId < vSlices.size(); vSliceId++) {
-      size_t coreId = vSliceId / crossbarCountInCore;
-      Value vSlice = vSlices[vSliceId];
-      tiles[hSliceId][coreId].push_back(vSlice);
-    }
-  }
-  return tiles;
+  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 broadcastToVector(Value scalarToBroadcast, int64_t length, ConversionPatternRewriter& rewriter, Location loc) {
-  auto oldType = cast<RankedTensorType>(scalarToBroadcast.getType());
-  Type elementType = oldType.getElementType();
-  int64_t shape[2] = {1, length};
-  Type type = oldType.cloneWith(ArrayRef(shape), elementType);
-
-  auto buildBroadcast = [&](Value input) -> Value {
-    auto zero = arith::ConstantIndexOp::create(rewriter, loc, 0).getResult();
-    SmallVector<Value> index(oldType.getRank(), zero);
-    auto elementValue = tensor::ExtractOp::create(rewriter, loc, input, index).getResult();
-    return tensor::SplatOp::create(rewriter, loc, type, elementValue);
-  };
-
-  if (isCompileTimeComputable(scalarToBroadcast))
-    return buildBroadcast(scalarToBroadcast);
-
-  auto broadcastCompute =
-    createSpatCompute<1>(rewriter, loc, TypeRange {type}, {}, ValueRange {scalarToBroadcast}, [&](Value input) {
-      spatial::SpatYieldOp::create(rewriter, loc, buildBroadcast(input));
-    });
-  return broadcastCompute.getResult(0);
-}
-
-Value materializeContiguousTensorSlice(Value source,
-                                       RankedTensorType resultType,
-                                       ArrayRef<OpFoldResult> offsets,
-                                       ArrayRef<OpFoldResult> strides,
-                                       ConversionPatternRewriter& rewriter,
-                                       Location loc) {
-  assert(resultType.hasStaticShape() && "expected static result type");
-  size_t rank = static_cast<size_t>(resultType.getRank());
-  assert(offsets.size() == rank && "expected rank-matching offsets");
-  assert(strides.size() == rank && "expected rank-matching strides");
-
-  SmallVector<OpFoldResult> sizes;
-  sizes.reserve(resultType.getRank());
-  for (int64_t size : resultType.getShape())
-    sizes.push_back(rewriter.getIndexAttr(size));
-
-  if (isContiguousTensorSlice(source, resultType, strides))
-    return tensor::ExtractSliceOp::create(rewriter, loc, resultType, source, offsets, sizes, strides).getResult();
-
-  if (resultType.getRank() == 0)
-    return tensor::ExtractSliceOp::create(rewriter, loc, resultType, source, offsets, sizes, strides).getResult();
-
-  Value init = tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), resultType.getElementType()).getResult();
-  SmallVector<Value> zeroIndices(resultType.getRank());
-  for (Value& zeroIndex : zeroIndices)
-    zeroIndex = arith::ConstantIndexOp::create(rewriter, loc, 0).getResult();
-
-  SmallVector<Value> resultIndices;
-  resultIndices.reserve(resultType.getRank());
-
-  auto buildLoopNest = [&](auto&& self, unsigned dim, Value accumulator) -> Value {
-    if (dim == resultType.getRank()) {
-      SmallVector<Value> sourceIndices;
-      sourceIndices.reserve(resultType.getRank());
-      for (unsigned idx = 0; idx < resultType.getRank(); ++idx) {
-        Value offsetValue = getIndexValue(offsets[idx], rewriter, loc);
-        Value scaledIndex = multiplyIndexValue(resultIndices[idx], strides[idx], rewriter, loc);
-        sourceIndices.push_back(addIndexValues(offsetValue, scaledIndex, rewriter, loc));
-      }
-
-      SmallVector<OpFoldResult> sourceOffsets;
-      SmallVector<OpFoldResult> destinationOffsets;
-      SmallVector<OpFoldResult> unitSizes;
-      SmallVector<OpFoldResult> unitStrides;
-      sourceOffsets.reserve(resultType.getRank());
-      destinationOffsets.reserve(resultType.getRank());
-      unitSizes.reserve(resultType.getRank());
-      unitStrides.reserve(resultType.getRank());
-      for (Value index : sourceIndices)
-        sourceOffsets.push_back(index);
-      for (Value index : resultIndices)
-        destinationOffsets.push_back(index);
-      for (int64_t idx = 0; idx < resultType.getRank(); ++idx) {
-        unitSizes.push_back(rewriter.getIndexAttr(1));
-        unitStrides.push_back(rewriter.getIndexAttr(1));
-      }
-
-      auto elementTensorType =
-        RankedTensorType::get(SmallVector<int64_t>(resultType.getRank(), 1), resultType.getElementType());
-      Value elementSlice =
-        tensor::ExtractSliceOp::create(rewriter, loc, elementTensorType, source, sourceOffsets, unitSizes, unitStrides)
-          .getResult();
-      return tensor::InsertSliceOp::create(
-               rewriter, loc, elementSlice, accumulator, destinationOffsets, unitSizes, unitStrides)
-        .getResult();
-    }
-
-    Value lower = zeroIndices[dim];
-    Value upper = arith::ConstantIndexOp::create(rewriter, loc, resultType.getDimSize(dim)).getResult();
-    Value step = arith::ConstantIndexOp::create(rewriter, loc, 1).getResult();
-    auto loop = scf::ForOp::create(rewriter, loc, lower, upper, step, ValueRange {accumulator});
-    rewriter.setInsertionPointToStart(loop.getBody());
-    resultIndices.push_back(loop.getInductionVar());
-    Value updated = self(self, dim + 1, loop.getRegionIterArgs().front());
-    resultIndices.pop_back();
-    scf::YieldOp::create(rewriter, loc, updated);
-    rewriter.setInsertionPointAfter(loop);
-    return loop.getResult(0);
-  };
-
-  return buildLoopNest(buildLoopNest, 0, init);
+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

@@ -3,6 +3,7 @@
 #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"
@@ -11,46 +12,12 @@
 
 #include <cassert>
 #include <cstddef>
+#include <optional>
 #include <type_traits>
 #include <utility>
 
 namespace onnx_mlir {
 
-template <class ShapedType>
-inline auto getImageWidth(const ShapedType& shapedType) {
-  return shapedType.getDimSize(2);
-}
-
-template <class ShapedType>
-inline auto getImageHeight(const ShapedType& shapedType) {
-  return shapedType.getDimSize(3);
-}
-
-template <class ShapedType>
-inline auto getImageChannel(const ShapedType& shapedType) {
-  return shapedType.getDimSize(1);
-}
-
-template <class ShapedType>
-inline auto getImageN(const ShapedType& shapedType) {
-  return shapedType.getDimSize(0);
-}
-
-template <class ShapedType>
-inline auto getKernelWidth(const ShapedType& shapedType) {
-  return shapedType.getDimSize(2);
-}
-
-template <class ShapedType>
-inline auto getKernelHeight(const ShapedType& shapedType) {
-  return shapedType.getDimSize(3);
-}
-
-template <class ShapedType>
-inline auto getFilterCount(const ShapedType& shapedType) {
-  return shapedType.getDimSize(0);
-}
-
 using HSliceId = size_t;
 using CoreId = size_t;
 
@@ -87,17 +54,6 @@ bool isHVectorShape(mlir::ArrayRef<T> shape) {
   return shape.size() == 2 && shape[0] == 1;
 }
 
-template <class T>
-bool isVVectorShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2 && shape[1] == 1;
-}
-
-template <class T>
-T getVectorLength(mlir::ArrayRef<T> shape) {
-  assert(isVectorShape(shape));
-  return shape[0] != 1 ? shape[0] : shape[1];
-}
-
 inline auto getTensorShape(mlir::Value tensor) {
   return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
 }
@@ -109,6 +65,25 @@ inline bool haveSameStaticShape(mlir::Value lhs, mlir::Value rhs) {
       && 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,
@@ -127,25 +102,13 @@ llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
 llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
   const mlir::Value& vectorToSlice, mlir::ConversionPatternRewriter& rewriter, mlir::Location loc);
 
-/// Tiles a matrix first across output columns and then across input rows so it
-/// can be assigned to crossbars grouped by core.
-llvm::DenseMap<HSliceId, llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>>>
-tileMatrix(mlir::Value& matrixToTile,
-           int64_t hSliceSize,
-           int64_t vSliceSize,
-           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 broadcastToVector(mlir::Value scalarToBroadcast,
-                              int64_t length,
-                              mlir::ConversionPatternRewriter& rewriter,
-                              mlir::Location loc);
-
-mlir::Value materializeContiguousTensorSlice(mlir::Value source,
-                                             mlir::RankedTensorType resultType,
-                                             llvm::ArrayRef<mlir::OpFoldResult> offsets,
-                                             llvm::ArrayRef<mlir::OpFoldResult> strides,
-                                             mlir::ConversionPatternRewriter& rewriter,
-                                             mlir::Location loc);
+mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
+                              mlir::Location loc,
+                              mlir::Value source,
+                              mlir::Value dest,
+                              llvm::ArrayRef<mlir::OpFoldResult> offsets);
 
 } // namespace onnx_mlir

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

@@ -1,4 +1,5 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/IRMapping.h"
@@ -43,8 +44,8 @@ bool isWeightLikeComputeOperand(Value value) {
       value = collapseShapeOp.getSrc();
       continue;
     }
-    if (auto transposeOp = dyn_cast<ONNXTransposeOp>(definingOp)) {
-      value = transposeOp.getData();
+    if (auto transposeOp = dyn_cast<linalg::TransposeOp>(definingOp)) {
+      value = transposeOp.getInput();
       continue;
     }
 
@@ -80,7 +81,7 @@ FailureOr<Value> materializeWeightLikeValueInBlock(Value value, IRRewriter& rewr
     return referencedValue.getResult();
   }
 
-  if (!isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, ONNXTransposeOp>(definingOp))
+  if (!isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, linalg::TransposeOp>(definingOp))
     return failure();
 
   IRMapping localMapper;

src/PIM/Conversion/ONNXToSpatial/CompileTime.cpp

@@ -1,4 +1,5 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -7,10 +8,11 @@
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/ErrorHandling.h"
 
 #include <utility>
 
+#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/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -26,8 +28,7 @@ static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
 }
 
 static bool hasConstantIndices(tensor::ExtractOp extractOp) {
-  return llvm::all_of(extractOp.getIndices(),
-                      [](Value index) { return isa_and_nonnull<arith::ConstantIndexOp>(index.getDefiningOp()); });
+  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
 }
 
 static bool isStaticTensorResult(Operation* op) {
@@ -37,13 +38,6 @@ static bool isStaticTensorResult(Operation* op) {
   });
 }
 
-static SmallVector<int64_t> computeRowMajorStrides(ArrayRef<int64_t> shape) {
-  SmallVector<int64_t> strides(shape.size(), 1);
-  for (int64_t dim = static_cast<int64_t>(shape.size()) - 2; dim >= 0; --dim)
-    strides[dim] = strides[dim + 1] * shape[dim + 1];
-  return strides;
-}
-
 static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
   auto tensorType = dyn_cast<RankedTensorType>(denseAttr.getType());
   if (!tensorType)
@@ -171,6 +165,16 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
     return succeeded(transposedAttr) ? *transposedAttr : nullptr;
   }
 
+  if (auto transposeOp = dyn_cast<linalg::TransposeOp>(definingOp)) {
+    auto inputAttr = getHostConstantDenseElementsAttrImpl(transposeOp.getInput(), visited);
+    if (!inputAttr)
+      return nullptr;
+
+    SmallVector<int64_t> perm(transposeOp.getPermutation().begin(), transposeOp.getPermutation().end());
+    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    return succeeded(transposedAttr) ? *transposedAttr : nullptr;
+  }
+
   if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(definingOp)) {
     auto inputAttr = getHostConstantDenseElementsAttrImpl(collapseShapeOp.getSrc(), visited);
     if (!inputAttr)
@@ -226,6 +230,9 @@ getCompileTimeSourceImpl(Operation* op, llvm::SmallPtrSetImpl<Operation*>& visit
   if (auto transposeOp = dyn_cast<ONNXTransposeOp>(op))
     return getCompileTimeSourceImpl(transposeOp.getData().getDefiningOp(), visited, chainLength);
 
+  if (auto transposeOp = dyn_cast<linalg::TransposeOp>(op))
+    return getCompileTimeSourceImpl(transposeOp.getInput().getDefiningOp(), visited, chainLength);
+
   if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(op))
     return getCompileTimeSourceImpl(collapseShapeOp.getSrc().getDefiningOp(), visited, chainLength);
 

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -1,23 +1,21 @@
 #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/IR/IRMapping.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
-#include "mlir/Transforms/WalkPatternRewriteDriver.h"
 
 #include "llvm/ADT/SmallVector.h"
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PrePatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -86,30 +84,6 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     returnOp.setOperand(index, computeResult);
 }
 
-static void wrapTopLevelRuntimeTransposes(func::FuncOp funcOp) {
-  IRRewriter rewriter(funcOp.getContext());
-  Block& entryBlock = funcOp.getFunctionBody().front();
-
-  for (Operation& op : llvm::make_early_inc_range(entryBlock)) {
-    auto transposeOp = dyn_cast<ONNXTransposeOp>(&op);
-    if (!transposeOp || isCompileTimeOp(transposeOp))
-      continue;
-
-    // Transpose stays globally legal because constant/view-only cases are
-    // allowed on the host. Any residual runtime transpose must be sunk into
-    // spat.compute before the host legality check.
-    auto resultType = transposeOp.getResult().getType();
-    rewriter.setInsertionPoint(transposeOp);
-    auto computeOp = createSpatCompute<1>(
-      rewriter, transposeOp.getLoc(), TypeRange {resultType}, {}, ValueRange {transposeOp.getData()}, [&](Value input) {
-        Value transposed =
-          ONNXTransposeOp::create(rewriter, transposeOp.getLoc(), resultType, input, transposeOp.getPermAttr());
-        spatial::SpatYieldOp::create(rewriter, transposeOp.getLoc(), transposed);
-      });
-    rewriter.replaceOp(transposeOp, computeOp.getResult(0));
-  }
-}
-
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   MLIRContext* ctx = &getContext();
@@ -117,6 +91,7 @@ void ONNXToSpatialPass::runOnOperation() {
   ConversionTarget preTarget(*ctx);
   preTarget.addLegalDialect<spatial::SpatialDialect,
                             ONNXDialect,
+                            linalg::LinalgDialect,
                             tensor::TensorDialect,
                             affine::AffineDialect,
                             arith::ArithDialect,
@@ -138,30 +113,18 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
-  RewritePatternSet matmulPatterns(ctx);
-  populateMatMulRewritePatterns(matmulPatterns, ctx);
-  walkAndApplyPatterns(moduleOp, std::move(matmulPatterns));
-
-  bool hasUnloweredMatMul = false;
-  moduleOp.walk([&](ONNXMatMulOp matmulOp) {
-    hasUnloweredMatMul = true;
-    matmulOp.emitOpError("remaining ONNX MatMul before the required ONNX-to-Spatial conversion");
-  });
-  if (hasUnloweredMatMul) {
-    moduleOp.emitError("failed to lower all ONNX MatMul ops before ONNX-to-Spatial conversion");
-    signalPassFailure();
-    return;
-  }
-
   ConversionTarget target(*ctx);
   target.addLegalDialect<spatial::SpatialDialect,
                          ONNXDialect,
+                         linalg::LinalgDialect,
                          tensor::TensorDialect,
                          affine::AffineDialect,
                          arith::ArithDialect,
                          scf::SCFDialect>();
   target.addIllegalOp<ONNXMatMulOp>();
+  target.addIllegalOp<ONNXTransposeOp>();
   target.addIllegalOp<ONNXAddOp>();
+  target.addIllegalOp<ONNXSubOp>();
   target.addIllegalOp<ONNXDivOp>();
   target.addIllegalOp<ONNXMulOp>();
   target.addIllegalOp<ONNXGemmOp>();
@@ -175,7 +138,9 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXGatherOp>();
   target.addIllegalOp<ONNXReshapeOp>();
   target.addIllegalOp<ONNXResizeOp>();
+  target.addIllegalOp<ONNXSliceOp>();
   target.addIllegalOp<ONNXLRNOp>();
+  target.addIllegalOp<ONNXReduceMeanOp>();
   target.addIllegalOp<ONNXReduceMeanV13Op>();
   target.addIllegalOp<ONNXSplitOp>();
 
@@ -190,6 +155,7 @@ void ONNXToSpatialPass::runOnOperation() {
   ConversionTarget earlyPostTarget(*ctx);
   earlyPostTarget.addLegalDialect<spatial::SpatialDialect,
                                   ONNXDialect,
+                                  linalg::LinalgDialect,
                                   tensor::TensorDialect,
                                   affine::AffineDialect,
                                   arith::ArithDialect,
@@ -205,6 +171,7 @@ void ONNXToSpatialPass::runOnOperation() {
   ConversionTarget postTarget(*ctx);
   postTarget.addLegalDialect<spatial::SpatialDialect,
                              ONNXDialect,
+                             linalg::LinalgDialect,
                              tensor::TensorDialect,
                              affine::AffineDialect,
                              arith::ArithDialect,
@@ -222,17 +189,14 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
-  wrapTopLevelRuntimeTransposes(*entryFunc);
+  populateEmptyFunction(*entryFunc);
+
+  dumpModule(moduleOp, "spatial0");
 
   if (failed(verifyONNXToSpatial(*entryFunc))) {
     moduleOp.emitError("ONNX-to-Spatial host legality verification failed");
     signalPassFailure();
-    return;
   }
-
-  populateEmptyFunction(*entryFunc);
-
-  dumpModule(moduleOp, "spatial0");
 }
 
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -69,11 +69,11 @@ LogicalResult verifyComputeLikeInputs(Operation* computeLikeOp,
       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");
+      InFlightDiagnostic diagnostic = illegalOp->emitOpError()
+                                   << kind << " input #" << currentInputIndex
+                                   << (allowChannelReceiveInputs ? " must come from the host or an explicit "
+                                                                   "spat.channel_receive"
+                                                                 : " must come from the host");
       if (definingOp)
         diagnostic.attachNote(definingOp->getLoc()) << "illegal Spatial producer is " << definingOp->getName();
     });
@@ -135,17 +135,17 @@ LogicalResult verifySpatialCommunicationInvariants(func::FuncOp funcOp) {
   pim::CappedDiagnosticReporter diagnostics;
 
   for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
-    (void)verifyComputeLikeInputs(
+    (void) verifyComputeLikeInputs(
       computeOp.getOperation(), computeOp.getInputs(), /*allowChannelReceiveInputs=*/true, "spat.compute", diagnostics);
     verifyNoExternalTensorCaptures(computeOp.getOperation(), computeOp.getBody(), "spat.compute", diagnostics);
   }
 
   for (auto computeBatchOp : funcOp.getOps<spatial::SpatComputeBatch>()) {
-    (void)verifyComputeLikeInputs(computeBatchOp.getOperation(),
-                                  computeBatchOp.getInputs(),
-                                  /*allowChannelReceiveInputs=*/false,
-                                  "spat.compute_batch",
-                                  diagnostics);
+    (void) verifyComputeLikeInputs(computeBatchOp.getOperation(),
+                                   computeBatchOp.getInputs(),
+                                   /*allowChannelReceiveInputs=*/false,
+                                   "spat.compute_batch",
+                                   diagnostics);
     verifyNoExternalTensorCaptures(
       computeBatchOp.getOperation(), computeBatchOp.getBody(), "spat.compute_batch", diagnostics);
   }

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -1,20 +1,16 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 
-namespace {
-
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
-
-} // namespace
-
-void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx) {
-  patterns.add<removeLRN>(ctx);
+void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { populateGeneratedPrePatterns(patterns, ctx); }
 
+void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  populateGeneratedConversionPatterns(patterns, ctx);
   populateElementwisePatterns(patterns, ctx);
+  populateMatMulRewritePatterns(patterns, ctx);
   populateGemmPatterns(patterns, ctx);
   populateConvPatterns(patterns, ctx);
   populatePoolPatterns(patterns, ctx);
@@ -26,7 +22,13 @@ void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRCon
   populateGatherPatterns(patterns, ctx);
   populateResizePatterns(patterns, ctx);
   populateReshapePatterns(patterns, ctx);
+  populateSlicePatterns(patterns, ctx);
   populateSplitPatterns(patterns, ctx);
+  populateTransposePatterns(patterns, ctx);
+}
+
+void populatePostPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  populateWeightPromotionPatterns(patterns, ctx);
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -1,38 +1,40 @@
 #pragma once
 
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
 namespace onnx_mlir {
 
+void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+
+void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateGeneratedConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
 void populateGemmPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
 void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
 void populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
 void populateReduceMeanPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
 void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
 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 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);
+void annotateWeightsConstants(mlir::func::FuncOp funcOp);
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,18 @@
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+namespace {
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
+
+} // namespace
+
+void populateGeneratedConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<removeLRN>(ctx);
+}
+
+} // namespace onnx_mlir

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

@@ -7,7 +7,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -83,7 +83,7 @@ static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
   }
 
   auto broadcastedAttr = DenseElementsAttr::get(resultType, resultValues);
-  return arith::ConstantOp::create(rewriter, loc, resultType, broadcastedAttr).getResult();
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), broadcastedAttr, resultType);
 }
 
 static FailureOr<Value>
@@ -121,7 +121,7 @@ static FailureOr<Value> materializeReciprocalTensor(Value value,
   }
 
   auto reciprocalAttr = DenseFPElementsAttr::get(resultType, reciprocalValues);
-  return arith::ConstantOp::create(rewriter, loc, resultType, reciprocalAttr).getResult();
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), reciprocalAttr, resultType);
 }
 
 template <typename OnnxOp, typename SpatialOp>
@@ -189,6 +189,7 @@ struct DivToSpatialCompute : OpConversionPattern<ONNXDivOp> {
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXAddOp, spatial::SpatVAddOp>>(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

@@ -13,7 +13,9 @@
 #include <limits>
 #include <utility>
 
-#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
@@ -50,70 +52,21 @@ materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewr
     return failure();
 
   auto scaledAttr = DenseFPElementsAttr::get(cast<RankedTensorType>(denseAttr.getType()), scaledValues);
-  return arith::ConstantOp::create(rewriter, loc, denseAttr.getType(), scaledAttr).getResult();
-}
-
-static Value transposeForSpatial(Value value,
-                                 RankedTensorType resultType,
-                                 ArrayRef<int64_t> permutation,
-                                 ConversionPatternRewriter& rewriter,
-                                 Location loc) {
-  if (isCompileTimeComputable(value))
-    return ONNXTransposeOp::create(rewriter, loc, resultType, value, rewriter.getI64ArrayAttr(permutation));
-
-  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, value, [&](Value input) {
-    Value transposed = ONNXTransposeOp::create(rewriter, loc, resultType, input, rewriter.getI64ArrayAttr(permutation));
-    spatial::SpatYieldOp::create(rewriter, loc, transposed);
-  });
-  return computeOp.getResult(0);
-}
-
-static Value createIndexConstant(ConversionPatternRewriter& rewriter, int64_t value) {
-  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  return getOrCreateHostIndexConstant(anchorOp, value, rewriter);
-}
-
-static Value
-createAffineApply(ConversionPatternRewriter& rewriter, Location loc, AffineExpr expr, ValueRange operands) {
-  AffineMap map = AffineMap::get(/*dimCount=*/operands.size(), /*symbolCount=*/0, expr);
-  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-  return createAffineApplyOrFoldedConstant(rewriter, loc, map, operands, anchorOp);
-}
-
-static Value
-multiplyIndexByConstant(Value value, int64_t multiplier, ConversionPatternRewriter& rewriter, Location loc) {
-  if (multiplier == 0)
-    return createIndexConstant(rewriter, 0);
-  if (multiplier == 1)
-    return value;
-
-  MLIRContext* context = rewriter.getContext();
-  AffineExpr d0 = getAffineDimExpr(0, context);
-  return createAffineApply(rewriter, loc, d0 * multiplier, ValueRange {value});
-}
-
-static Value modIndexByConstant(Value value, int64_t divisor, ConversionPatternRewriter& rewriter, Location loc) {
-  if (divisor == 1)
-    return createIndexConstant(rewriter, 0);
-
-  MLIRContext* context = rewriter.getContext();
-  AffineExpr d0 = getAffineDimExpr(0, context);
-  return createAffineApply(rewriter, loc, d0 % divisor, ValueRange {value});
-}
-
-static Value createGemmBatchRow(Value lane, int64_t numOutRows, ConversionPatternRewriter& rewriter, Location loc) {
-  return modIndexByConstant(lane, numOutRows, rewriter, loc);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), scaledAttr, denseAttr.getType());
 }
 
 static Value createGemmBatchKOffset(
   Value lane, int64_t numOutRows, int64_t numKSlices, ConversionPatternRewriter& rewriter, Location loc) {
   if (numKSlices == 1)
-    return createIndexConstant(rewriter, 0);
+    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
 
   MLIRContext* context = rewriter.getContext();
   AffineExpr d0 = getAffineDimExpr(0, context);
-  return createAffineApply(
-    rewriter, loc, (d0.floorDiv(numOutRows) % numKSlices) * crossbarSize.getValue(), ValueRange {lane});
+  return createOrFoldAffineApply(rewriter,
+                                 loc,
+                                 (d0.floorDiv(numOutRows) % numKSlices) * crossbarSize.getValue(),
+                                 ValueRange {lane},
+                                 rewriter.getInsertionBlock()->getParentOp());
 }
 
 static Value createGemmBatchHOffset(Value lane,
@@ -123,12 +76,15 @@ static Value createGemmBatchHOffset(Value lane,
                                     ConversionPatternRewriter& rewriter,
                                     Location loc) {
   if (numOutHSlices == 1)
-    return createIndexConstant(rewriter, 0);
+    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
 
   MLIRContext* context = rewriter.getContext();
   AffineExpr d0 = getAffineDimExpr(0, context);
-  return createAffineApply(
-    rewriter, loc, d0.floorDiv(numOutRows * numKSlices) * crossbarSize.getValue(), ValueRange {lane});
+  return createOrFoldAffineApply(rewriter,
+                                 loc,
+                                 d0.floorDiv(numOutRows * numKSlices) * crossbarSize.getValue(),
+                                 ValueRange {lane},
+                                 rewriter.getInsertionBlock()->getParentOp());
 }
 
 static Value
@@ -146,9 +102,9 @@ createZeroPaddedTensor(Value value, RankedTensorType resultType, ConversionPatte
     padBlock->addArgument(rewriter.getIndexType(), loc);
   padOp.getRegion().push_back(padBlock);
   rewriter.setInsertionPointToStart(padBlock);
-  auto zero = arith::ConstantOp::create(
-    rewriter, loc, sourceType.getElementType(), rewriter.getZeroAttr(sourceType.getElementType()));
-  tensor::YieldOp::create(rewriter, loc, zero.getResult());
+  auto zero = getOrCreateConstant(
+    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
+  tensor::YieldOp::create(rewriter, loc, zero);
   rewriter.setInsertionPointAfter(padOp);
   return padOp.getResult();
 }
@@ -180,7 +136,7 @@ static FailureOr<Value> materializePaddedConstantMatrix(Value value,
       resultValues[row * resultShape[1] + col] = sourceValues[row * sourceShape[1] + col];
 
   auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
-  return arith::ConstantOp::create(rewriter, loc, resultType, resultAttr).getResult();
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
 }
 
 static FailureOr<Value> materializePaddedBroadcastedConstantTensor(Value value,
@@ -246,7 +202,7 @@ static FailureOr<Value> materializePaddedBroadcastedConstantTensor(Value value,
   }
 
   auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
-  return arith::ConstantOp::create(rewriter, loc, resultType, resultAttr).getResult();
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
 }
 
 static FailureOr<Value> prepareBias(Value c,
@@ -292,103 +248,89 @@ static Value createPaddedInputCompute(Value input,
   return computeOp.getResult(0);
 }
 
-static spatial::SpatComputeBatch createVmmBatch(Value a,
-                                                Value b,
-                                                RankedTensorType aType,
-                                                RankedTensorType paddedBType,
-                                                RankedTensorType partialPiecesType,
-                                                int64_t numOutRows,
-                                                int64_t numKSlices,
-                                                int64_t numOutHSlices,
-                                                ConversionPatternRewriter& rewriter,
-                                                Location loc) {
+static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
+                                                           Value b,
+                                                           RankedTensorType aType,
+                                                           RankedTensorType paddedBType,
+                                                           RankedTensorType partialPiecesType,
+                                                           int64_t numOutRows,
+                                                           int64_t numKSlices,
+                                                           int64_t numOutHSlices,
+                                                           ConversionPatternRewriter& rewriter,
+                                                           Location loc) {
   const int64_t laneCount = partialPiecesType.getDimSize(0);
-  auto batchOp = spatial::SpatComputeBatch::create(rewriter,
-                                                   loc,
-                                                   TypeRange {partialPiecesType},
-                                                   rewriter.getI32IntegerAttr(static_cast<int32_t>(laneCount)),
-                                                   ValueRange {b},
-                                                   ValueRange {a});
+  auto batchOp = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {partialPiecesType},
+    laneCount,
+    ValueRange {b},
+    ValueRange {a},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Value row =
+        onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutRows, rewriter.getInsertionBlock()->getParentOp());
+      Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, rewriter, loc);
+      Value hOffset = createGemmBatchHOffset(args.lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
 
-  SmallVector<Type> blockArgTypes {rewriter.getIndexType(), paddedBType, aType, partialPiecesType};
-  SmallVector<Location> blockArgLocs(blockArgTypes.size(), loc);
-  Block* body =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToEnd(body);
+      auto aTileType =
+        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, aType.getElementType());
+      auto bTileType = RankedTensorType::get(
+        {static_cast<int64_t>(crossbarSize.getValue()), static_cast<int64_t>(crossbarSize.getValue())},
+        paddedBType.getElementType());
+      auto pieceType =
+        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());
+      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, rewriter, loc);
 
-  auto lane = batchOp.getLaneArgument();
-  auto weight = batchOp.getWeightArgument(0);
-  auto input = batchOp.getInputArgument(0);
-  auto output = batchOp.getOutputArgument(0);
-  assert(lane && weight && input && output && "malformed Gemm compute_batch body");
+      SmallVector<OpFoldResult> bOffsets {kOffset, hOffset};
+      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 piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
-  Value row = createGemmBatchRow(*lane, numOutRows, rewriter, loc);
-  Value kOffset = createGemmBatchKOffset(*lane, numOutRows, numKSlices, rewriter, loc);
-  Value hOffset = createGemmBatchHOffset(*lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
-
-  auto aTileType = RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, aType.getElementType());
-  auto bTileType = RankedTensorType::get(
-    {static_cast<int64_t>(crossbarSize.getValue()), static_cast<int64_t>(crossbarSize.getValue())},
-    paddedBType.getElementType());
-  auto pieceType =
-    RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());
-  Value aTile = extractATile(*input, row, kOffset, aTileType, rewriter, loc);
-
-  SmallVector<OpFoldResult> bOffsets {kOffset, hOffset};
-  SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(crossbarSize.getValue()),
-                                    rewriter.getIndexAttr(crossbarSize.getValue())};
-  SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  Value bTile =
-    tensor::ExtractSliceOp::create(rewriter, loc, bTileType, *weight, bOffsets, bSizes, unitStrides).getResult();
-  Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
-
-  auto inParallelOp = spatial::SpatInParallelOp::create(rewriter, loc);
-  rewriter.setInsertionPointToStart(&inParallelOp.getRegion().front());
-  SmallVector<OpFoldResult> pieceOffsets {*lane, rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-  tensor::ParallelInsertSliceOp::create(rewriter, loc, piece, *output, pieceOffsets, pieceSizes, unitStrides);
-
-  rewriter.setInsertionPointAfter(batchOp);
-  return batchOp;
+      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);
+    });
+  if (failed(batchOp))
+    return failure();
+  return *batchOp;
 }
 
-static Value createDynamicGemmBatchRow(
-  Value lane, int64_t numOutCols, ConversionPatternRewriter& rewriter, Location loc) {
+static Value
+createDynamicGemmBatchRow(Value lane, int64_t numOutCols, ConversionPatternRewriter& rewriter, Location loc) {
   if (numOutCols == 1)
     return lane;
 
   MLIRContext* context = rewriter.getContext();
   AffineExpr d0 = getAffineDimExpr(0, context);
-  return createAffineApply(rewriter, loc, d0.floorDiv(numOutCols), ValueRange {lane});
+  return createOrFoldAffineApply(
+    rewriter, loc, d0.floorDiv(numOutCols), ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp());
 }
 
-static Value createDynamicGemmBatchColumn(
-  Value lane, int64_t numOutCols, ConversionPatternRewriter& rewriter, Location loc) {
-  return modIndexByConstant(lane, numOutCols, rewriter, loc);
-}
-
-static Value
-extractDynamicGemmBColumn(Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
+static Value extractDynamicGemmBColumn(
+  Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), column};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
   auto columnSliceType = RankedTensorType::get({vectorType.getDimSize(1), 1}, vectorType.getElementType());
-  Value columnSlice = materializeContiguousTensorSlice(matrix, columnSliceType, offsets, strides, rewriter, loc);
-  SmallVector<ReassociationIndices> collapseReassociation {ReassociationIndices {0, 1}};
+  Value columnSlice =
+    tensor::ExtractSliceOp::create(rewriter, loc, columnSliceType, matrix, offsets, sizes, strides).getResult();
+  SmallVector<ReassociationIndices> collapseReassociation {
+    ReassociationIndices {0, 1}
+  };
   auto collapsedType = RankedTensorType::get({vectorType.getDimSize(1)}, vectorType.getElementType());
   Value collapsed =
     tensor::CollapseShapeOp::create(rewriter, loc, collapsedType, columnSlice, collapseReassociation).getResult();
-  SmallVector<ReassociationIndices> expandReassociation {ReassociationIndices {0, 1}};
+  SmallVector<ReassociationIndices> expandReassociation {
+    ReassociationIndices {0, 1}
+  };
   return tensor::ExpandShapeOp::create(rewriter, loc, vectorType, collapsed, expandReassociation).getResult();
 }
 
-static Value extractTransposedBRow(
-  Value transposedB, Value row, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
-  SmallVector<OpFoldResult> offsets {row, rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
-  SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  return tensor::ExtractSliceOp::create(rewriter, loc, vectorType, transposedB, offsets, sizes, strides).getResult();
-}
-
 static Value extractDynamicGemmRowVector(
   Value matrix, Value row, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {row, rewriter.getIndexAttr(0)};
@@ -432,7 +374,7 @@ static Value createScalarTensorConstant(RankedTensorType scalarType,
   auto elementType = scalarType.getElementType();
   auto scalarAttr = rewriter.getFloatAttr(elementType, value);
   auto denseAttr = DenseElementsAttr::get(scalarType, scalarAttr);
-  return arith::ConstantOp::create(rewriter, loc, scalarType, denseAttr).getResult();
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), denseAttr, scalarType);
 }
 
 static Value createBroadcastedBiasScalar(Value bias,
@@ -444,13 +386,15 @@ static Value createBroadcastedBiasScalar(Value bias,
                                          Location loc) {
   SmallVector<OpFoldResult> unitStrides(biasType.getRank(), rewriter.getIndexAttr(1));
   if (biasType.getRank() == 1) {
-    SmallVector<OpFoldResult> offsets {
-      biasType.getDimSize(0) == 1 ? OpFoldResult(rewriter.getIndexAttr(0)) : OpFoldResult(column)};
+    SmallVector<OpFoldResult> offsets {biasType.getDimSize(0) == 1 ? OpFoldResult(rewriter.getIndexAttr(0))
+                                                                   : OpFoldResult(column)};
     SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1)};
     auto vectorType = RankedTensorType::get({1}, scalarType.getElementType());
-    Value vector = tensor::ExtractSliceOp::create(rewriter, loc, vectorType, bias, offsets, sizes, unitStrides)
-                     .getResult();
-    SmallVector<ReassociationIndices> reassociation {ReassociationIndices {0, 1}};
+    Value vector =
+      tensor::ExtractSliceOp::create(rewriter, loc, vectorType, bias, offsets, sizes, unitStrides).getResult();
+    SmallVector<ReassociationIndices> reassociation {
+      ReassociationIndices {0, 1}
+    };
     return tensor::ExpandShapeOp::create(rewriter, loc, scalarType, vector, reassociation).getResult();
   }
 
@@ -466,116 +410,114 @@ static Value createBroadcastedBiasScalar(Value bias,
   return tensor::SplatOp::create(rewriter, loc, scalarType, scalar).getResult();
 }
 
-static spatial::SpatComputeBatch createVvdmulBatch(Value a,
-                                                   Value b,
-                                                   RankedTensorType aType,
-                                                   RankedTensorType bType,
-                                                   RankedTensorType scalarPiecesType,
-                                                   RankedTensorType outType,
-                                                   bool bAlreadyTransposed,
-                                                   ConversionPatternRewriter& rewriter,
-                                                   Location loc) {
+static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
+                                                              Value b,
+                                                              RankedTensorType aType,
+                                                              RankedTensorType bType,
+                                                              RankedTensorType scalarPiecesType,
+                                                              RankedTensorType outType,
+                                                              ConversionPatternRewriter& rewriter,
+                                                              Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   const int64_t reductionSize = aType.getDimSize(1);
   const int64_t laneCount = numOutRows * numOutCols;
-  auto batchOp = spatial::SpatComputeBatch::create(rewriter,
-                                                   loc,
-                                                   TypeRange {scalarPiecesType},
-                                                   rewriter.getI32IntegerAttr(static_cast<int32_t>(laneCount)),
-                                                   ValueRange {},
-                                                   ValueRange {a, b});
+  auto batchOp = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {scalarPiecesType},
+    laneCount,
+    ValueRange {},
+    ValueRange {a, b},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Value row = createDynamicGemmBatchRow(args.lane, numOutCols, rewriter, loc);
+      Value column =
+        onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
 
-  SmallVector<Type> blockArgTypes {rewriter.getIndexType(), aType, bType, scalarPiecesType};
-  SmallVector<Location> blockArgLocs(blockArgTypes.size(), loc);
-  Block* body =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToEnd(body);
+      auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
+      auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
+      Value aVector = extractDynamicGemmRowVector(args.inputs[0], row, vectorType, rewriter, loc);
+      Value bVector = extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc);
+      Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
 
-  auto lane = batchOp.getLaneArgument();
-  auto inputA = batchOp.getInputArgument(0);
-  auto inputB = batchOp.getInputArgument(1);
-  auto output = batchOp.getOutputArgument(0);
-  assert(lane && inputA && inputB && output && "malformed dynamic Gemm compute_batch body");
-
-  Value row = createDynamicGemmBatchRow(*lane, numOutCols, rewriter, loc);
-  Value column = createDynamicGemmBatchColumn(*lane, numOutCols, rewriter, loc);
-
-  auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
-  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-  Value aVector = extractDynamicGemmRowVector(*inputA, row, vectorType, rewriter, loc);
-  Value bVector = bAlreadyTransposed
-                    ? extractTransposedBRow(*inputB, column, vectorType, rewriter, loc)
-                    : extractDynamicGemmBColumn(*inputB, column, vectorType, rewriter, loc);
-  Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-
-  auto inParallelOp = spatial::SpatInParallelOp::create(rewriter, loc);
-  rewriter.setInsertionPointToStart(&inParallelOp.getRegion().front());
-  SmallVector<OpFoldResult> outputOffsets {*lane, rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  tensor::ParallelInsertSliceOp::create(rewriter, loc, scalar, *output, outputOffsets, scalarSizes, unitStrides);
-
-  rewriter.setInsertionPointAfter(batchOp);
-  return batchOp;
+      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);
+    });
+  if (failed(batchOp))
+    return failure();
+  return *batchOp;
 }
 
-static spatial::SpatCompute createDynamicGemmOutputCompute(Value scalarPieces,
-                                                           Value bias,
-                                                           RankedTensorType scalarPiecesType,
-                                                           RankedTensorType biasType,
-                                                           RankedTensorType outType,
-                                                           float alpha,
-                                                           float beta,
-                                                           ConversionPatternRewriter& rewriter,
-                                                           Location loc) {
+static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scalarPieces,
+                                                                      Value bias,
+                                                                      RankedTensorType scalarPiecesType,
+                                                                      RankedTensorType biasType,
+                                                                      RankedTensorType outType,
+                                                                      float alpha,
+                                                                      float beta,
+                                                                      ConversionPatternRewriter& rewriter,
+                                                                      Location loc) {
   const int64_t laneCount = scalarPiecesType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   SmallVector<Value> inputs {scalarPieces};
   if (bias)
     inputs.push_back(bias);
 
-  return createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) {
+  return createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult {
     Value pieces = blockArgs[0];
     Value biasArg = bias ? blockArgs[1] : Value();
     auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
     Value outputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType()).getResult();
-    Value c0 = createIndexConstant(rewriter, 0);
-    Value c1 = createIndexConstant(rewriter, 1);
-    Value cLaneCount = createIndexConstant(rewriter, laneCount);
-    auto loop = scf::ForOp::create(rewriter, loc, c0, cLaneCount, c1, ValueRange {outputInit});
-    rewriter.setInsertionPointToStart(loop.getBody());
+    Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+    Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
+    Value cLaneCount = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), laneCount);
+    auto loop = buildNormalizedScfFor(
+      rewriter,
+      loc,
+      c0,
+      cLaneCount,
+      c1,
+      ValueRange {outputInit},
+      [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+        Value outputAcc = iterArgs.front();
+        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();
+        if (alpha != 1.0f) {
+          Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc);
+          scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, scalar, alphaTensor).getResult();
+        }
+        if (biasArg) {
+          Value biasScalar =
+            createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, nestedLoc);
+          if (beta != 1.0f) {
+            Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, nestedLoc);
+            biasScalar =
+              spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).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)
+            .getResult();
+        yielded.push_back(outputNext);
+        return success();
+      });
+    if (failed(loop))
+      return failure();
 
-    Value lane = loop.getInductionVar();
-    Value outputAcc = loop.getRegionIterArgs().front();
-    Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, loc);
-    Value column = createDynamicGemmBatchColumn(lane, numOutCols, rewriter, loc);
-    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, loc, scalarType, pieces, scalarOffsets, scalarSizes, unitStrides)
-        .getResult();
-    if (alpha != 1.0f) {
-      Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, loc);
-      scalar = spatial::SpatVMulOp::create(rewriter, loc, scalarType, scalar, alphaTensor).getResult();
-    }
-    if (biasArg) {
-      Value biasScalar = createBroadcastedBiasScalar(biasArg, biasType, row, column, scalarType, rewriter, loc);
-      if (beta != 1.0f) {
-        Value betaTensor = createScalarTensorConstant(scalarType, beta, rewriter, loc);
-        biasScalar = spatial::SpatVMulOp::create(rewriter, loc, scalarType, biasScalar, betaTensor).getResult();
-      }
-      scalar = spatial::SpatVAddOp::create(rewriter, loc, scalarType, scalar, biasScalar).getResult();
-    }
-    SmallVector<OpFoldResult> outputOffsets {row, column};
-    Value outputNext =
-      tensor::InsertSliceOp::create(rewriter, loc, scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
-        .getResult();
-    scf::YieldOp::create(rewriter, loc, outputNext);
-
-    rewriter.setInsertionPointAfter(loop);
-    spatial::SpatYieldOp::create(rewriter, loc, loop.getResult(0));
+    spatial::SpatYieldOp::create(rewriter, loc, loop->results.front());
+    return success();
   });
 }
 
@@ -587,7 +529,11 @@ static Value createPartialGroupOffset(Value hSlice,
                                       Location loc) {
   MLIRContext* context = rewriter.getContext();
   AffineExpr d0 = getAffineDimExpr(0, context);
-  return createAffineApply(rewriter, loc, d0 * (numKSlices * numOutRows) + kSlice * numOutRows, ValueRange {hSlice});
+  return createOrFoldAffineApply(rewriter,
+                                 loc,
+                                 d0 * (numKSlices * numOutRows) + kSlice * numOutRows,
+                                 ValueRange {hSlice},
+                                 rewriter.getInsertionBlock()->getParentOp());
 }
 
 static Value extractReductionPiece(Value partialPiecesArg,
@@ -636,83 +582,92 @@ static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
   return activePieces.front();
 }
 
-static spatial::SpatCompute createReductionCompute(Value partialPieces,
-                                                   Value bias,
-                                                   RankedTensorType partialPiecesType,
-                                                   RankedTensorType outType,
-                                                   RankedTensorType paddedOutType,
-                                                   int64_t numKSlices,
-                                                   ConversionPatternRewriter& rewriter,
-                                                   Location loc) {
+static FailureOr<spatial::SpatCompute> createReductionCompute(Value partialPieces,
+                                                              Value bias,
+                                                              RankedTensorType partialPiecesType,
+                                                              RankedTensorType outType,
+                                                              RankedTensorType paddedOutType,
+                                                              int64_t numKSlices,
+                                                              ConversionPatternRewriter& rewriter,
+                                                              Location loc) {
   SmallVector<Value> inputs {partialPieces};
   if (bias)
     inputs.push_back(bias);
 
-  auto computeOp = createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) {
-    Value partialPiecesArg = blockArgs[0];
-    Value biasArg = bias ? blockArgs[1] : Value();
-    if (biasArg && cast<RankedTensorType>(biasArg.getType()) != paddedOutType)
-      biasArg = createZeroPaddedTensor(biasArg, paddedOutType, rewriter, loc);
+  auto computeOp =
+    createSpatCompute(rewriter, loc, TypeRange {outType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult {
+      Value partialPiecesArg = blockArgs[0];
+      Value biasArg = bias ? blockArgs[1] : Value();
+      if (biasArg && cast<RankedTensorType>(biasArg.getType()) != paddedOutType)
+        biasArg = createZeroPaddedTensor(biasArg, paddedOutType, rewriter, loc);
 
-    const int64_t numOutRows = outType.getDimSize(0);
-    const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue());
-    auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
-                                           partialPiecesType.getElementType());
+      const int64_t numOutRows = outType.getDimSize(0);
+      const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue());
+      auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
+                                             partialPiecesType.getElementType());
 
-    Value outputInit =
-      tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
-    SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-    SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
-                                          rewriter.getIndexAttr(crossbarSize.getValue())};
+      Value outputInit =
+        tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
+      SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
+                                            rewriter.getIndexAttr(crossbarSize.getValue())};
 
-    auto buildOutputSlice = [&](Value outputAcc, Value hSlice) -> Value {
-      Value reduced =
-        reducePartialPiecesForHSlice(partialPiecesArg, hSlice, pieceType, numKSlices, numOutRows, rewriter, loc);
-      Value hOffset = multiplyIndexByConstant(hSlice, crossbarSize.getValue(), rewriter, loc);
-      if (biasArg) {
-        SmallVector<OpFoldResult> biasOffsets {rewriter.getIndexAttr(0), hOffset};
-        Value biasSlice =
-          tensor::ExtractSliceOp::create(rewriter, loc, pieceType, biasArg, biasOffsets, pieceSizes, unitStrides)
-            .getResult();
-        reduced = spatial::SpatVAddOp::create(rewriter, loc, pieceType, reduced, biasSlice).getResult();
+      auto buildOutputSlice = [&](Value outputAcc, Value hSlice) -> Value {
+        Value reduced =
+          reducePartialPiecesForHSlice(partialPiecesArg, hSlice, pieceType, numKSlices, numOutRows, rewriter, loc);
+        Value hOffset = onnx_mlir::affineMulConst(
+          rewriter, loc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
+        if (biasArg) {
+          SmallVector<OpFoldResult> biasOffsets {rewriter.getIndexAttr(0), hOffset};
+          Value biasSlice =
+            tensor::ExtractSliceOp::create(rewriter, loc, pieceType, biasArg, biasOffsets, pieceSizes, unitStrides)
+              .getResult();
+          reduced = spatial::SpatVAddOp::create(rewriter, loc, pieceType, reduced, biasSlice).getResult();
+        }
+
+        SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), hOffset};
+        return tensor::InsertSliceOp::create(rewriter, loc, reduced, outputAcc, outputOffsets, pieceSizes, unitStrides)
+          .getResult();
+      };
+
+      Value paddedOutput = outputInit;
+      if (numOutHSlices == 1) {
+        Value hSlice = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+        paddedOutput = buildOutputSlice(outputInit, hSlice);
+      }
+      else {
+        Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+        Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
+        Value cOutHSlices =
+          getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutHSlices);
+        auto hLoop = buildNormalizedScfFor(
+          rewriter,
+          loc,
+          c0,
+          cOutHSlices,
+          c1,
+          ValueRange {outputInit},
+          [&](OpBuilder&, Location, Value hSlice, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+            yielded.push_back(buildOutputSlice(iterArgs.front(), hSlice));
+            return success();
+          });
+        if (failed(hLoop))
+          return failure();
+        paddedOutput = hLoop->results.front();
       }
 
-      SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), hOffset};
-      return tensor::InsertSliceOp::create(rewriter, loc, reduced, outputAcc, outputOffsets, pieceSizes, unitStrides)
-        .getResult();
-    };
-
-    Value paddedOutput = outputInit;
-    if (numOutHSlices == 1) {
-      Value hSlice = createIndexConstant(rewriter, 0);
-      paddedOutput = buildOutputSlice(outputInit, hSlice);
-    }
-    else {
-      Value c0 = createIndexConstant(rewriter, 0);
-      Value c1 = createIndexConstant(rewriter, 1);
-      Value cOutHSlices = createIndexConstant(rewriter, numOutHSlices);
-      auto hLoop = scf::ForOp::create(rewriter, loc, c0, cOutHSlices, c1, ValueRange {outputInit});
-      rewriter.setInsertionPointToStart(hLoop.getBody());
-
-      Value hSlice = hLoop.getInductionVar();
-      Value outputAcc = hLoop.getRegionIterArgs().front();
-      scf::YieldOp::create(rewriter, loc, buildOutputSlice(outputAcc, hSlice));
-
-      rewriter.setInsertionPointAfter(hLoop);
-      paddedOutput = hLoop.getResult(0);
-    }
-
-    Value result = paddedOutput;
-    if (paddedOutType != outType) {
-      SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(outType.getDimSize(0)),
-                                             rewriter.getIndexAttr(outType.getDimSize(1))};
-      result =
-        tensor::ExtractSliceOp::create(rewriter, loc, outType, paddedOutput, outputOffsets, outputSizes, unitStrides)
-          .getResult();
-    }
-    spatial::SpatYieldOp::create(rewriter, loc, result);
-  });
+      Value result = paddedOutput;
+      if (paddedOutType != outType) {
+        SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+        SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(outType.getDimSize(0)),
+                                               rewriter.getIndexAttr(outType.getDimSize(1))};
+        result =
+          tensor::ExtractSliceOp::create(rewriter, loc, outType, paddedOutput, outputOffsets, outputSizes, unitStrides)
+            .getResult();
+      }
+      spatial::SpatYieldOp::create(rewriter, loc, result);
+      return success();
+    });
 
   return computeOp;
 }
@@ -735,11 +690,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());
@@ -770,6 +720,20 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     return failure();
   }
 
+  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();
+    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
+    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    bType = transposedType;
+  }
+
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
   const int64_t reductionSize = aType.getDimSize(1);
@@ -793,10 +757,8 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
       biasType = *verifiedBiasType;
     }
 
-    const int64_t expectedBRows = gemmOpAdaptor.getTransB() ? numOutCols : reductionSize;
-    const int64_t expectedBCols = gemmOpAdaptor.getTransB() ? reductionSize : numOutCols;
-    if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != expectedBRows
-        || bType.getDimSize(1) != expectedBCols) {
+    if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize
+        || bType.getDimSize(1) != numOutCols) {
       gemmOp.emitOpError("has inconsistent A, B, and output shapes");
       return failure();
     }
@@ -808,11 +770,14 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     }
 
     auto scalarPiecesType = RankedTensorType::get({laneCount64, 1}, outType.getElementType());
-    auto batchOp = createVvdmulBatch(
-      a, b, aType, bType, scalarPiecesType, outType, gemmOpAdaptor.getTransB(), rewriter, loc);
+    auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, rewriter, loc);
+    if (failed(batchOp))
+      return failure();
     auto outputCompute = createDynamicGemmOutputCompute(
-      batchOp.getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc);
-    rewriter.replaceOp(gemmOp, outputCompute.getResults());
+      batchOp->getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc);
+    if (failed(outputCompute))
+      return failure();
+    rewriter.replaceOp(gemmOp, outputCompute->getResults());
     return success();
   }
 
@@ -824,13 +789,6 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
   b = *scaledB;
   bType = cast<RankedTensorType>(b.getType());
 
-  if (gemmOpAdaptor.getTransB()) {
-    auto bShape = bType.getShape();
-    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType());
-    b = transposeForSpatial(b, transposedType, {1, 0}, rewriter, loc);
-    bType = cast<RankedTensorType>(b.getType());
-  }
-
   if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize || bType.getDimSize(1) != numOutCols) {
     gemmOp.emitOpError("has inconsistent A, B, and output shapes after transpose handling");
     return failure();
@@ -887,10 +845,14 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     RankedTensorType::get({laneCount64, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType());
   auto batchOp =
     createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+  if (failed(batchOp))
+    return failure();
   auto reductionCompute = createReductionCompute(
-    batchOp.getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc);
+    batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc);
+  if (failed(reductionCompute))
+    return failure();
 
-  rewriter.replaceOp(gemmOp, reductionCompute.getResults());
+  rewriter.replaceOp(gemmOp, reductionCompute->getResults());
   return success();
 }
 

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

@@ -1,13 +1,18 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 #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"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -16,26 +21,85 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-static SmallVector<int64_t> normalizeAxes(ArrayAttr axesAttr, int64_t rank) {
+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;
-  if (!axesAttr) {
-    normalizedAxes.reserve(rank);
-    for (int64_t axis = 0; axis < rank; axis++)
-      normalizedAxes.push_back(axis);
-    return normalizedAxes;
+  normalizedAxes.reserve(axes.size());
+  for (int64_t axis : axes) {
+    auto normalizedAxis = normalizeAxisChecked(axis, rank);
+    if (failed(normalizedAxis))
+      return failure();
+    normalizedAxes.push_back(*normalizedAxis);
   }
-
-  normalizedAxes.reserve(axesAttr.size());
-  for (Attribute attr : axesAttr) {
-    int64_t axis = cast<IntegerAttr>(attr).getInt();
-    normalizedAxes.push_back(axis >= 0 ? axis : rank + axis);
-  }
-
   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) {
@@ -50,6 +114,184 @@ static RankedTensorType getAllOnesType(RankedTensorType inputType, Type elementT
   return RankedTensorType::get(SmallVector<int64_t>(inputType.getRank(), 1), elementType);
 }
 
+static RankedTensorType getKeepdimsType(RankedTensorType inputType, Type elementType, ArrayRef<bool> reducedAxes) {
+  SmallVector<int64_t> shape;
+  shape.reserve(inputType.getRank());
+  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
+    shape.push_back(isReduced ? 1 : dim);
+  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());
+  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
+    shape.push_back(isReduced ? dim : 1);
+  return RankedTensorType::get(shape, inputType.getElementType(), inputType.getEncoding());
+}
+
+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());
+}
+
+static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
+  SmallVector<int64_t> keptAxes;
+  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes))
+    if (!isReduced)
+      keptAxes.push_back(static_cast<int64_t>(axis));
+  return keptAxes;
+}
+
+static Value
+computeLaneIndex(Value lane, int64_t stride, int64_t dimSize, ConversionPatternRewriter& rewriter, Location loc) {
+  if (dimSize == 1)
+    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+
+  MLIRContext* context = rewriter.getContext();
+  AffineExpr d0 = getAffineDimExpr(0, context);
+  AffineExpr expr = d0;
+  if (stride != 1)
+    expr = expr.floorDiv(stride);
+  if (dimSize != 1)
+    expr = expr % dimSize;
+  return createOrFoldAffineApply(rewriter, loc, expr, ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp());
+}
+
+static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
+                                                     ArrayRef<bool> reducedAxes,
+                                                     RankedTensorType batchType,
+                                                     RankedTensorType leafType,
+                                                     ConversionPatternRewriter& rewriter,
+                                                     Location loc) {
+  auto inputType = cast<RankedTensorType>(input.getType());
+  auto sliceType = getReducedSliceType(inputType, reducedAxes);
+  SmallVector<int64_t> keptAxes = getKeptAxes(reducedAxes);
+
+  int64_t laneCount = 1;
+  SmallVector<int64_t> keptAxisStrides(keptAxes.size(), 1);
+  for (int64_t index = static_cast<int64_t>(keptAxes.size()) - 1; index >= 0; --index) {
+    keptAxisStrides[index] = laneCount;
+    int64_t dimSize = inputType.getDimSize(keptAxes[index]);
+    if (dimSize <= 0)
+      return failure();
+    if (laneCount > std::numeric_limits<int32_t>::max() / dimSize)
+      return failure();
+    laneCount *= dimSize;
+  }
+
+  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,
+                           loc,
+                           TypeRange {batchType},
+                           laneCount,
+                           {},
+                           ValueRange {input},
+                           [&](detail::SpatComputeBatchBodyArgs args) {
+                             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));
+                                 sliceSizes.push_back(rewriter.getIndexAttr(inputType.getDimSize(axis)));
+                                 continue;
+                               }
+
+                               Value axisIndex = computeLaneIndex(
+                                 args.lane, keptAxisStrides[keptAxisIndex], inputType.getDimSize(axis), rewriter, loc);
+                               ++keptAxisIndex;
+                               sliceOffsets.push_back(axisIndex);
+                               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);
+                           });
+  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;
+
+  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;
+  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);
+  }
+  if (!pendingLeadingReducedAxes.empty())
+    expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
+
+  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);
+}
+
 static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> reducedAxes) {
   SmallVector<ReassociationIndices> reassociation;
   ReassociationIndices currentGroup;
@@ -72,69 +314,13 @@ static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<boo
   return reassociation;
 }
 
-static Value
-createAverageCompute(Value input, RankedTensorType resultType, ConversionPatternRewriter& rewriter, Location loc) {
-  constexpr size_t numInputs = 1;
-  auto computeOp = createSpatCompute<numInputs>(rewriter, loc, resultType, {}, ValueRange {input}, [&](Value x) {
-    auto avgOp = spatial::SpatVAvgOp::create(rewriter, loc, resultType, x);
-    spatial::SpatYieldOp::create(rewriter, loc, avgOp.getResult());
-  });
-  return computeOp.getResult(0);
-}
-
-static Value concatValues(ValueRange inputs, int64_t axis, ConversionPatternRewriter& rewriter, Location loc) {
-  auto firstType = cast<RankedTensorType>(inputs.front().getType());
-  SmallVector<int64_t> outputShape(firstType.getShape().begin(), firstType.getShape().end());
-  int64_t concatDimSize = 0;
-  for (Value input : inputs)
-    concatDimSize += cast<RankedTensorType>(input.getType()).getDimSize(axis);
-  outputShape[axis] = concatDimSize;
-  auto resultType = RankedTensorType::get(outputShape, firstType.getElementType(), firstType.getEncoding());
-
-  if (llvm::all_of(inputs, isCompileTimeComputable))
-    return createSpatConcat(rewriter, loc, axis, inputs);
-
-  auto concatCompute = createSpatCompute(rewriter, loc, TypeRange {resultType}, {}, inputs, [&](ValueRange args) {
-    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, axis, args));
-  });
-  return concatCompute.getResult(0);
-}
-
-static Value buildReduceMeanKeepdims(Value input,
-                                     ArrayRef<bool> reducedAxes,
-                                     int64_t axis,
-                                     RankedTensorType leafType,
-                                     ConversionPatternRewriter& rewriter,
-                                     Location loc) {
-  int64_t rank = cast<RankedTensorType>(input.getType()).getRank();
-  if (axis == rank)
-    return createAverageCompute(input, leafType, rewriter, loc);
-
-  if (reducedAxes[axis])
-    return buildReduceMeanKeepdims(input, reducedAxes, axis + 1, leafType, rewriter, loc);
-
-  SmallVector<Value> slices = sliceTensor(input, axis, /*sliceSize=*/1, rewriter, loc);
-  SmallVector<Value> reducedSlices;
-  reducedSlices.reserve(slices.size());
-  for (Value slice : slices)
-    reducedSlices.push_back(buildReduceMeanKeepdims(slice, reducedAxes, axis + 1, leafType, rewriter, loc));
-
-  return concatValues(reducedSlices, axis, rewriter, loc);
-}
-
 static Value squeezeReducedAxes(Value keepdimsValue,
                                 RankedTensorType resultType,
                                 ArrayRef<bool> reducedAxes,
                                 ConversionPatternRewriter& rewriter,
                                 Location loc) {
-  if (resultType.getRank() == 0) {
-    SmallVector<Value> indices(cast<RankedTensorType>(keepdimsValue.getType()).getRank(),
-                               arith::ConstantIndexOp::create(rewriter, loc, 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();
 
@@ -146,28 +332,55 @@ 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());
     if (!inputType || !resultType || !inputType.hasStaticShape() || !resultType.hasStaticShape())
       return failure();
+    if (inputType.getRank() == 0) {
+      rewriter.replaceOp(reduceMeanOp, adaptor.getData());
+      return success();
+    }
 
-    SmallVector<int64_t> axes = normalizeAxes(reduceMeanOp.getAxesAttr(), inputType.getRank());
-    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());
-    Value reducedKeepdims =
-      buildReduceMeanKeepdims(adaptor.getData(), reducedAxes, /*axis=*/0, leafType, rewriter, loc);
+    RankedTensorType compactKeptType = getCompactKeptType(inputType, resultType.getElementType(), reducedAxes);
+    RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
+    int64_t laneCount = 1;
+    for (int64_t dim : compactKeptType.getShape())
+      laneCount *= dim;
+    RankedTensorType batchType = getLanePackedKeepdimsType(laneCount, leafType);
 
-    if (reduceMeanOp.getKeepdims() != 0) {
+    auto lanePackedKeepdims =
+      buildReduceMeanKeepdimsBatch(adaptor.getData(), reducedAxes, batchType, leafType, rewriter, loc);
+    if (failed(lanePackedKeepdims))
+      return failure();
+    Value reducedKeepdims =
+      buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc);
+
+    if (semantics->keepdims != 0) {
       rewriter.replaceOp(reduceMeanOp, reducedKeepdims);
       return success();
     }
@@ -181,7 +394,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/Pool.cpp

@@ -12,6 +12,7 @@
 #include <optional>
 #include <type_traits>
 
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -23,43 +24,26 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-template <typename ArrayAttrT>
-static int64_t getI64(ArrayAttrT arrayAttr, size_t index) {
-  return cast<IntegerAttr>(arrayAttr[index]).getInt();
-}
-
-template <typename ArrayAttrT>
-static int64_t getOptionalI64(std::optional<ArrayAttrT> arrayAttr, size_t index, int64_t defaultValue) {
-  return arrayAttr ? getI64(*arrayAttr, index) : defaultValue;
-}
-
-static Value materializeContiguousTile(ConversionPatternRewriter& rewriter, Location loc, Value tile) {
+static Value materializeTileTensor(ConversionPatternRewriter& rewriter, Location loc, Value tile) {
   auto tileType = cast<RankedTensorType>(tile.getType());
   Value empty = tensor::EmptyOp::create(rewriter, loc, tileType.getShape(), tileType.getElementType());
-
-  SmallVector<OpFoldResult> offsets(tileType.getRank(), rewriter.getIndexAttr(0));
-  SmallVector<OpFoldResult> sizes;
-  sizes.reserve(tileType.getRank());
-  for (int64_t dimSize : tileType.getShape())
-    sizes.push_back(rewriter.getIndexAttr(dimSize));
-  SmallVector<OpFoldResult> strides(tileType.getRank(), rewriter.getIndexAttr(1));
-
-  return tensor::InsertSliceOp::create(rewriter, loc, tile, empty, offsets, sizes, strides);
+  return insertStaticSlice(rewriter, loc, tile, empty, getZeroOffsets(rewriter, tileType.getRank()));
 }
 
 static Value
 createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type elementType, bool useMinimumValue) {
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
   if (!useMinimumValue)
-    return arith::ConstantOp::create(rewriter, loc, elementType, rewriter.getZeroAttr(elementType));
+    return getOrCreateConstant(rewriter, anchorOp, rewriter.getZeroAttr(elementType), elementType);
 
   if (auto floatType = dyn_cast<FloatType>(elementType)) {
     auto minValue = llvm::APFloat::getInf(floatType.getFloatSemantics(), /*Negative=*/true);
-    return arith::ConstantOp::create(rewriter, loc, elementType, rewriter.getFloatAttr(floatType, minValue));
+    return getOrCreateConstant(rewriter, anchorOp, rewriter.getFloatAttr(floatType, minValue), elementType);
   }
 
   if (auto integerType = dyn_cast<IntegerType>(elementType)) {
     auto minValue = llvm::APInt::getSignedMinValue(integerType.getWidth());
-    return arith::ConstantOp::create(rewriter, loc, elementType, rewriter.getIntegerAttr(integerType, minValue));
+    return getOrCreateConstant(rewriter, anchorOp, rewriter.getIntegerAttr(integerType, minValue), elementType);
   }
 
   llvm_unreachable("unsupported pool element type");
@@ -166,7 +150,7 @@ static FailureOr<Value> createAverageScaleTensor(ConversionPatternRewriter& rewr
   }
 
   auto scaleAttr = DenseElementsAttr::get(scaleType, scaleValues);
-  return arith::ConstantOp::create(rewriter, loc, scaleType, scaleAttr).getResult();
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), scaleAttr, scaleType);
 }
 
 template <typename PoolOp>
@@ -197,12 +181,12 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
     const int64_t inputWidth = xType.getDimSize(3);
     const int64_t outputHeight = outType.getDimSize(2);
     const int64_t outputWidth = outType.getDimSize(3);
-    const int64_t kernelHeight = getI64(kernelAttr, 0);
-    const int64_t kernelWidth = getI64(kernelAttr, 1);
-    const int64_t strideHeight = getOptionalI64(poolOp.getStrides(), 0, 1);
-    const int64_t strideWidth = getOptionalI64(poolOp.getStrides(), 1, 1);
-    const int64_t dilationHeight = getOptionalI64(poolOp.getDilations(), 0, 1);
-    const int64_t dilationWidth = getOptionalI64(poolOp.getDilations(), 1, 1);
+    const int64_t kernelHeight = getI64Attr(kernelAttr, 0);
+    const int64_t kernelWidth = getI64Attr(kernelAttr, 1);
+    const int64_t strideHeight = getOptionalI64Attr(poolOp.getStrides(), 0, 1);
+    const int64_t strideWidth = getOptionalI64Attr(poolOp.getStrides(), 1, 1);
+    const int64_t dilationHeight = getOptionalI64Attr(poolOp.getDilations(), 0, 1);
+    const int64_t dilationWidth = getOptionalI64Attr(poolOp.getDilations(), 1, 1);
 
     int64_t padTop = 0;
     int64_t padLeft = 0;
@@ -212,10 +196,10 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
     if (auto padsAttr = poolOp.getPads()) {
       if (padsAttr->size() != 4)
         return rewriter.notifyMatchFailure(poolOp, "pads must have four elements.");
-      padTop = getI64(*padsAttr, 0);
-      padLeft = getI64(*padsAttr, 1);
-      padBottom = getI64(*padsAttr, 2);
-      padRight = getI64(*padsAttr, 3);
+      padTop = getI64Attr(*padsAttr, 0);
+      padLeft = getI64Attr(*padsAttr, 1);
+      padBottom = getI64Attr(*padsAttr, 2);
+      padRight = getI64Attr(*padsAttr, 3);
     }
     else {
       StringRef autoPad = poolOp.getAutoPad();
@@ -283,94 +267,111 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
           createPaddedPoolInput(rewriter, loc, poolOp, xArg, xType, padTop, padLeft, padBottom, padRight);
         Value pooledOutputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType());
 
-        Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
-        Value c1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
-        Value cOutputPatchCount = arith::ConstantIndexOp::create(rewriter, loc, outputPatchCount);
-        Value cOutputPixelsPerBatch = arith::ConstantIndexOp::create(rewriter, loc, outputHeight * outputWidth);
-        Value cOutputWidth = arith::ConstantIndexOp::create(rewriter, loc, outputWidth);
-        Value cStrideHeight = arith::ConstantIndexOp::create(rewriter, loc, strideHeight);
-        Value cStrideWidth = arith::ConstantIndexOp::create(rewriter, loc, strideWidth);
+        Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+        Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+        Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+        Value cOutputPatchCount = getOrCreateIndexConstant(rewriter, anchorOp, outputPatchCount);
+        Value cOutputPixelsPerBatch = getOrCreateIndexConstant(rewriter, anchorOp, outputHeight * outputWidth);
+        Value cOutputWidth = getOrCreateIndexConstant(rewriter, anchorOp, outputWidth);
+        Value cStrideHeight = getOrCreateIndexConstant(rewriter, anchorOp, strideHeight);
+        Value cStrideWidth = getOrCreateIndexConstant(rewriter, anchorOp, strideWidth);
 
-        auto outputLoop = scf::ForOp::create(rewriter, loc, c0, cOutputPatchCount, c1, ValueRange {pooledOutputInit});
-        rewriter.setInsertionPointToStart(outputLoop.getBody());
+        auto outputLoop = buildNormalizedScfFor(
+          rewriter,
+          loc,
+          c0,
+          cOutputPatchCount,
+          c1,
+          ValueRange {pooledOutputInit},
+          [&](OpBuilder&,
+              Location nestedLoc,
+              Value outputPatchIndex,
+              ValueRange iterArgs,
+              SmallVectorImpl<Value>& yielded) {
+            Value pooledOutputAcc = iterArgs.front();
+            Value batchIndex = arith::DivUIOp::create(rewriter, nestedLoc, outputPatchIndex, cOutputPixelsPerBatch);
+            Value batchPatchIndex =
+              arith::RemUIOp::create(rewriter, nestedLoc, outputPatchIndex, cOutputPixelsPerBatch);
+            Value outHeightIndex = arith::DivUIOp::create(rewriter, nestedLoc, batchPatchIndex, cOutputWidth);
+            Value outWidthIndex = arith::RemUIOp::create(rewriter, nestedLoc, batchPatchIndex, cOutputWidth);
+            Value windowBaseH = arith::MulIOp::create(rewriter, nestedLoc, outHeightIndex, cStrideHeight);
+            Value windowBaseW = arith::MulIOp::create(rewriter, nestedLoc, outWidthIndex, cStrideWidth);
 
-        Value outputPatchIndex = outputLoop.getInductionVar();
-        Value pooledOutputAcc = outputLoop.getRegionIterArgs().front();
+            Value updatedOutput = pooledOutputAcc;
+            for (int64_t channelTile = 0; channelTile < channelTileCount; ++channelTile) {
+              const int64_t tileChannels = std::min<int64_t>(xbarSize, channels - channelTile * xbarSize);
+              auto tileType = RankedTensorType::get({1, tileChannels, 1, 1}, outType.getElementType());
+              Value reducedWindow =
+                createPoolFillTensor(rewriter, nestedLoc, tileType, std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
 
-        Value batchIndex = arith::DivUIOp::create(rewriter, loc, outputPatchIndex, cOutputPixelsPerBatch);
-        Value batchPatchIndex = arith::RemUIOp::create(rewriter, loc, outputPatchIndex, cOutputPixelsPerBatch);
-        Value outHeightIndex = arith::DivUIOp::create(rewriter, loc, batchPatchIndex, cOutputWidth);
-        Value outWidthIndex = arith::RemUIOp::create(rewriter, loc, batchPatchIndex, cOutputWidth);
-        Value windowBaseH = arith::MulIOp::create(rewriter, loc, outHeightIndex, cStrideHeight);
-        Value windowBaseW = arith::MulIOp::create(rewriter, loc, outWidthIndex, cStrideWidth);
+              for (int64_t kernelH = 0; kernelH < kernelHeight; ++kernelH) {
+                Value paddedInH = windowBaseH;
+                if (kernelH * dilationHeight != 0) {
+                  Value kernelHOffset = getOrCreateIndexConstant(rewriter, anchorOp, kernelH * dilationHeight);
+                  paddedInH = arith::AddIOp::create(rewriter, nestedLoc, paddedInH, kernelHOffset);
+                }
 
-        Value updatedOutput = pooledOutputAcc;
-        for (int64_t channelTile = 0; channelTile < channelTileCount; ++channelTile) {
-          const int64_t tileChannels = std::min<int64_t>(xbarSize, channels - channelTile * xbarSize);
-          auto tileType = RankedTensorType::get({1, tileChannels, 1, 1}, outType.getElementType());
-          Value reducedWindow =
-            createPoolFillTensor(rewriter, loc, tileType, std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
+                for (int64_t kernelW = 0; kernelW < kernelWidth; ++kernelW) {
+                  Value paddedInW = windowBaseW;
+                  if (kernelW * dilationWidth != 0) {
+                    Value kernelWOffset = getOrCreateIndexConstant(rewriter, anchorOp, kernelW * dilationWidth);
+                    paddedInW = arith::AddIOp::create(rewriter, nestedLoc, paddedInW, kernelWOffset);
+                  }
 
-          for (int64_t kernelH = 0; kernelH < kernelHeight; ++kernelH) {
-            Value paddedInH = windowBaseH;
-            if (kernelH * dilationHeight != 0) {
-              Value kernelHOffset = arith::ConstantIndexOp::create(rewriter, loc, kernelH * dilationHeight);
-              paddedInH = arith::AddIOp::create(rewriter, loc, paddedInH, kernelHOffset);
-            }
-
-            for (int64_t kernelW = 0; kernelW < kernelWidth; ++kernelW) {
-              Value paddedInW = windowBaseW;
-              if (kernelW * dilationWidth != 0) {
-                Value kernelWOffset = arith::ConstantIndexOp::create(rewriter, loc, kernelW * dilationWidth);
-                paddedInW = arith::AddIOp::create(rewriter, loc, paddedInW, kernelWOffset);
+                  SmallVector<OpFoldResult> offsets = {
+                    batchIndex, rewriter.getIndexAttr(channelTile * xbarSize), paddedInH, paddedInW};
+                  SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
+                                                     rewriter.getIndexAttr(tileChannels),
+                                                     rewriter.getIndexAttr(1),
+                                                     rewriter.getIndexAttr(1)};
+                  SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(1)};
+                  Value windowValue =
+                    tensor::ExtractSliceOp::create(rewriter, nestedLoc, tileType, paddedInput, offsets, sizes, strides);
+                  windowValue = materializeTileTensor(rewriter, nestedLoc, windowValue);
+                  reducedWindow = ReduceOp::create(rewriter, nestedLoc, tileType, reducedWindow, windowValue);
+                }
               }
 
-              SmallVector<OpFoldResult> offsets = {
-                batchIndex, rewriter.getIndexAttr(channelTile * xbarSize), paddedInH, paddedInW};
-              SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
-                                                 rewriter.getIndexAttr(tileChannels),
-                                                 rewriter.getIndexAttr(1),
-                                                 rewriter.getIndexAttr(1)};
-              SmallVector<OpFoldResult> strides = {
+              if constexpr (std::is_same_v<PoolOp, ONNXAveragePoolOp>) {
+                SmallVector<OpFoldResult> scaleOffsets = {rewriter.getIndexAttr(0),
+                                                          rewriter.getIndexAttr(channelTile * xbarSize),
+                                                          outHeightIndex,
+                                                          outWidthIndex};
+                SmallVector<OpFoldResult> scaleSizes = {rewriter.getIndexAttr(1),
+                                                        rewriter.getIndexAttr(tileChannels),
+                                                        rewriter.getIndexAttr(1),
+                                                        rewriter.getIndexAttr(1)};
+                SmallVector<OpFoldResult> scaleStrides = {rewriter.getIndexAttr(1),
+                                                          rewriter.getIndexAttr(1),
+                                                          rewriter.getIndexAttr(1),
+                                                          rewriter.getIndexAttr(1)};
+                Value scaleSlice = tensor::ExtractSliceOp::create(
+                  rewriter, nestedLoc, tileType, averageScaleTensor, scaleOffsets, scaleSizes, scaleStrides);
+                scaleSlice = materializeTileTensor(rewriter, nestedLoc, scaleSlice);
+                reducedWindow = spatial::SpatVMulOp::create(rewriter, nestedLoc, tileType, reducedWindow, scaleSlice);
+              }
+
+              SmallVector<OpFoldResult> outputOffsets = {
+                batchIndex, rewriter.getIndexAttr(channelTile * xbarSize), outHeightIndex, outWidthIndex};
+              SmallVector<OpFoldResult> outputSizes = {rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(tileChannels),
+                                                       rewriter.getIndexAttr(1),
+                                                       rewriter.getIndexAttr(1)};
+              SmallVector<OpFoldResult> outputStrides = {
                 rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-              Value windowValue =
-                tensor::ExtractSliceOp::create(rewriter, loc, tileType, paddedInput, offsets, sizes, strides);
-              windowValue = materializeContiguousTile(rewriter, loc, windowValue);
-              reducedWindow = ReduceOp::create(rewriter, loc, tileType, reducedWindow, windowValue);
+              updatedOutput = tensor::InsertSliceOp::create(
+                rewriter, nestedLoc, reducedWindow, updatedOutput, outputOffsets, outputSizes, outputStrides);
             }
-          }
+            yielded.push_back(updatedOutput);
+            return success();
+          });
+        if (failed(outputLoop))
+          return failure();
 
-          if constexpr (std::is_same_v<PoolOp, ONNXAveragePoolOp>) {
-            SmallVector<OpFoldResult> scaleOffsets = {
-              rewriter.getIndexAttr(0), rewriter.getIndexAttr(channelTile * xbarSize), outHeightIndex, outWidthIndex};
-            SmallVector<OpFoldResult> scaleSizes = {rewriter.getIndexAttr(1),
-                                                    rewriter.getIndexAttr(tileChannels),
-                                                    rewriter.getIndexAttr(1),
-                                                    rewriter.getIndexAttr(1)};
-            SmallVector<OpFoldResult> scaleStrides = {
-              rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-            Value scaleSlice = tensor::ExtractSliceOp::create(
-              rewriter, loc, tileType, averageScaleTensor, scaleOffsets, scaleSizes, scaleStrides);
-            scaleSlice = materializeContiguousTile(rewriter, loc, scaleSlice);
-            reducedWindow = spatial::SpatVMulOp::create(rewriter, loc, tileType, reducedWindow, scaleSlice);
-          }
-
-          SmallVector<OpFoldResult> outputOffsets = {
-            batchIndex, rewriter.getIndexAttr(channelTile * xbarSize), outHeightIndex, outWidthIndex};
-          SmallVector<OpFoldResult> outputSizes = {rewriter.getIndexAttr(1),
-                                                   rewriter.getIndexAttr(tileChannels),
-                                                   rewriter.getIndexAttr(1),
-                                                   rewriter.getIndexAttr(1)};
-          SmallVector<OpFoldResult> outputStrides = {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-          updatedOutput = tensor::InsertSliceOp::create(
-            rewriter, loc, reducedWindow, updatedOutput, outputOffsets, outputSizes, outputStrides);
-        }
-
-        scf::YieldOp::create(rewriter, loc, updatedOutput);
-
-        rewriter.setInsertionPointAfter(outputLoop);
-        spatial::SpatYieldOp::create(rewriter, loc, outputLoop.getResult(0));
+        spatial::SpatYieldOp::create(rewriter, loc, outputLoop->results.front());
         return success();
       });
     if (failed(computeOp))

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

@@ -3,8 +3,9 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -13,16 +14,6 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-static int64_t normalizeAxis(int64_t axis, int64_t rank) { return axis >= 0 ? axis : rank + axis; }
-
-static 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;
-}
-
 static Value buildLoopSoftmaxSlice(Value input,
                                    Value accumulator,
                                    RankedTensorType inputType,
@@ -36,7 +27,7 @@ static Value buildLoopSoftmaxSlice(Value input,
 
   SmallVector<OpFoldResult> offsets;
   SmallVector<OpFoldResult> sizes;
-  SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
+  SmallVector<OpFoldResult> strides = getUnitStrides(rewriter, rank);
   offsets.reserve(rank);
   sizes.reserve(rank);
 
@@ -52,52 +43,65 @@ static Value buildLoopSoftmaxSlice(Value input,
   return tensor::InsertSliceOp::create(rewriter, loc, softmaxSlice, accumulator, offsets, sizes, strides);
 }
 
-static Value buildLoopSoftmaxNest(Value input,
-                                  Value accumulator,
-                                  RankedTensorType inputType,
-                                  int64_t axis,
-                                  SmallVectorImpl<Value>& outerIndices,
-                                  ConversionPatternRewriter& rewriter,
-                                  Location loc) {
+static FailureOr<Value> buildLoopSoftmaxNest(Value input,
+                                             Value accumulator,
+                                             RankedTensorType inputType,
+                                             int64_t axis,
+                                             SmallVectorImpl<Value>& outerIndices,
+                                             ConversionPatternRewriter& rewriter,
+                                             Location loc) {
   if (axis == inputType.getRank() - 1)
     return buildLoopSoftmaxSlice(input, accumulator, inputType, outerIndices, rewriter, loc);
 
-  Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
-  Value c1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
-  Value cUpper = arith::ConstantIndexOp::create(rewriter, loc, inputType.getDimSize(axis));
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+  Value cUpper = getOrCreateIndexConstant(rewriter, anchorOp, inputType.getDimSize(axis));
 
-  auto loop = scf::ForOp::create(rewriter, loc, c0, cUpper, c1, ValueRange {accumulator});
-  rewriter.setInsertionPointToStart(loop.getBody());
-
-  Value loopIndex = loop.getInductionVar();
-  Value loopAccumulator = loop.getRegionIterArgs().front();
-  outerIndices.push_back(loopIndex);
-  Value updatedAccumulator =
-    buildLoopSoftmaxNest(input, loopAccumulator, inputType, axis + 1, outerIndices, rewriter, loc);
-  outerIndices.pop_back();
-
-  scf::YieldOp::create(rewriter, loc, updatedAccumulator);
-  rewriter.setInsertionPointAfter(loop);
-  return loop.getResult(0);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c0,
+    cUpper,
+    c1,
+    ValueRange {accumulator},
+    [&](OpBuilder& builder, Location nestedLoc, Value loopIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+      outerIndices.push_back(loopIndex);
+      auto updatedAccumulator =
+        buildLoopSoftmaxNest(input, iterArgs.front(), inputType, axis + 1, outerIndices, rewriter, nestedLoc);
+      outerIndices.pop_back();
+      if (failed(updatedAccumulator))
+        return failure();
+      yielded.push_back(*updatedAccumulator);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  return loop->results.front();
 }
 
-static Value createLoopSoftmaxCompute(Value input, ConversionPatternRewriter& rewriter, Location loc) {
+static FailureOr<Value> createLoopSoftmaxCompute(Value input, ConversionPatternRewriter& rewriter, Location loc) {
   auto inputType = cast<RankedTensorType>(input.getType());
   constexpr size_t numInputs = 1;
-  auto computeOp =
-    createSpatCompute<numInputs>(rewriter, loc, TypeRange {inputType}, {}, ValueRange {input}, [&](Value x) {
+  auto computeOp = createSpatCompute<numInputs>(
+    rewriter, loc, TypeRange {inputType}, {}, ValueRange {input}, [&](Value x) -> LogicalResult {
       if (inputType.getRank() == 1) {
         Value softmax = spatial::SpatSoftmaxOp::create(rewriter, loc, inputType, x).getResult();
         spatial::SpatYieldOp::create(rewriter, loc, softmax);
-        return;
+        return success();
       }
 
       Value outputInit = tensor::EmptyOp::create(rewriter, loc, inputType.getShape(), inputType.getElementType());
       SmallVector<Value> outerIndices;
-      Value result = buildLoopSoftmaxNest(x, outputInit, inputType, /*axis=*/0, outerIndices, rewriter, loc);
-      spatial::SpatYieldOp::create(rewriter, loc, result);
+      auto result = buildLoopSoftmaxNest(x, outputInit, inputType, /*axis=*/0, outerIndices, rewriter, loc);
+      if (failed(result))
+        return failure();
+      spatial::SpatYieldOp::create(rewriter, loc, *result);
+      return success();
     });
-  return computeOp.getResult(0);
+  if (failed(computeOp))
+    return failure();
+  return computeOp->getResult(0);
 }
 
 struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
@@ -110,44 +114,40 @@ struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
     if (!inputType || !inputType.hasStaticShape())
       return failure();
 
-    int64_t axis = normalizeAxis(softmaxOp.getAxis(), inputType.getRank());
-    if (axis < 0 || axis >= inputType.getRank())
+    auto axis = normalizeAxisChecked(softmaxOp.getAxis(), inputType.getRank());
+    if (failed(axis))
       return failure();
 
     Value input = adaptor.getInput();
     Value result;
-    if (axis == inputType.getRank() - 1) {
-      result = createLoopSoftmaxCompute(input, rewriter, softmaxOp.getLoc());
+    if (*axis == inputType.getRank() - 1) {
+      auto computed = createLoopSoftmaxCompute(input, rewriter, softmaxOp.getLoc());
+      if (failed(computed))
+        return failure();
+      result = *computed;
     }
     else {
       SmallVector<int64_t> permutation;
       permutation.reserve(inputType.getRank());
       for (int64_t dim = 0; dim < inputType.getRank(); ++dim)
-        if (dim != axis)
+        if (dim != *axis)
           permutation.push_back(dim);
-      permutation.push_back(axis);
-
-      SmallVector<int64_t> inversePermutation(inputType.getRank());
-      for (auto [newIndex, oldIndex] : llvm::enumerate(permutation))
-        inversePermutation[oldIndex] = static_cast<int64_t>(newIndex);
+      permutation.push_back(*axis);
+      SmallVector<int64_t> inversePermutation = invertPermutation(permutation);
 
       auto transposedType = RankedTensorType::get(
         permuteShape(inputType.getShape(), permutation), inputType.getElementType(), inputType.getEncoding());
-      auto preTransposeCompute =
-        createSpatCompute<1>(rewriter, softmaxOp.getLoc(), TypeRange {transposedType}, {}, input, [&](Value x) {
-          Value transposed = ONNXTransposeOp::create(
-            rewriter, softmaxOp.getLoc(), transposedType, x, rewriter.getI64ArrayAttr(permutation));
-          spatial::SpatYieldOp::create(rewriter, softmaxOp.getLoc(), transposed);
-        });
-      Value transposedInput = preTransposeCompute.getResult(0);
-      Value transposedResult = createLoopSoftmaxCompute(transposedInput, rewriter, softmaxOp.getLoc());
-      auto postTransposeCompute =
-        createSpatCompute<1>(rewriter, softmaxOp.getLoc(), TypeRange {inputType}, {}, transposedResult, [&](Value x) {
-          Value transposed = ONNXTransposeOp::create(
-            rewriter, softmaxOp.getLoc(), inputType, x, rewriter.getI64ArrayAttr(inversePermutation));
-          spatial::SpatYieldOp::create(rewriter, softmaxOp.getLoc(), transposed);
-        });
-      result = postTransposeCompute.getResult(0);
+      Value transposedInput =
+        ONNXTransposeOp::create(
+          rewriter, softmaxOp.getLoc(), transposedType, input, rewriter.getI64ArrayAttr(permutation))
+          .getResult();
+      auto transposedResult = createLoopSoftmaxCompute(transposedInput, rewriter, softmaxOp.getLoc());
+      if (failed(transposedResult))
+        return failure();
+      result =
+        ONNXTransposeOp::create(
+          rewriter, softmaxOp.getLoc(), inputType, *transposedResult, rewriter.getI64ArrayAttr(inversePermutation))
+          .getResult();
     }
 
     rewriter.replaceOp(softmaxOp, result);

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

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/PatternMatch.h"
@@ -8,8 +9,9 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -20,7 +22,7 @@ namespace onnx_mlir {
 namespace {
 
 static bool isWeightMaterializationHelperUser(Operation* op) {
-  return isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, ONNXTransposeOp>(op);
+  return isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, linalg::TransposeOp>(op);
 }
 
 static bool canPromoteInputBlockArgument(BlockArgument arg) {
@@ -35,6 +37,14 @@ static bool isDirectConstantValue(Value value) {
   return isa_and_nonnull<arith::ConstantOp, ONNXConstantOp>(value.getDefiningOp());
 }
 
+struct PromotedOperands {
+  SmallVector<bool> promoteInput;
+  SmallVector<Value> newWeights;
+  SmallVector<Value> newInputs;
+  SmallVector<Type> newInputTypes;
+  SmallVector<Location> newInputLocs;
+};
+
 template <typename ComputeOpTy>
 static bool hasPromotableWeightLikeInputs(ComputeOpTy compute) {
   for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
@@ -47,60 +57,91 @@ static bool hasPromotableWeightLikeInputs(ComputeOpTy compute) {
   return false;
 }
 
+template <typename ComputeOpTy>
+static FailureOr<PromotedOperands> computePromotedOperands(ComputeOpTy compute) {
+  PromotedOperands promoted;
+  promoted.promoteInput.assign(compute.getInputs().size(), false);
+  promoted.newWeights.append(compute.getWeights().begin(), compute.getWeights().end());
+  promoted.newWeights.reserve(compute.getWeights().size() + compute.getInputs().size());
+  promoted.newInputs.reserve(compute.getInputs().size());
+  promoted.newInputTypes.reserve(compute.getInputs().size());
+  promoted.newInputLocs.reserve(compute.getInputs().size());
+
+  bool needsRewrite = false;
+  for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
+    if (!isWeightLikeComputeOperand(input))
+      goto keep_input;
+    if (isDirectConstantValue(input) && !canPromoteInputBlockArgument(compute.getInputArgument(inputIdx)))
+      goto keep_input;
+    promoted.promoteInput[inputIdx] = true;
+    promoted.newWeights.push_back(input);
+    needsRewrite = true;
+    continue;
+
+keep_input:
+    promoted.newInputs.push_back(input);
+    promoted.newInputTypes.push_back(input.getType());
+    promoted.newInputLocs.push_back(input.getLoc());
+  }
+
+  if (!needsRewrite)
+    return failure();
+  return promoted;
+}
+
+template <typename ComputeOpTy>
+static LogicalResult mapPromotedInputArguments(ComputeOpTy compute,
+                                               const PromotedOperands& promoted,
+                                               IRRewriter& bodyRewriter,
+                                               IRMapping& mapper,
+                                               std::function<std::optional<BlockArgument>(size_t)> getNewInputArg,
+                                               PatternRewriter& rewriter) {
+  size_t newInputIdx = 0;
+  for (auto [oldInputIdx, input] : llvm::enumerate(compute.getInputs())) {
+    auto oldArg = compute.getInputArgument(oldInputIdx);
+    if (!oldArg)
+      return rewriter.notifyMatchFailure(compute, "missing input block argument during rewrite");
+    if (!promoted.promoteInput[oldInputIdx]) {
+      auto newInputArg = getNewInputArg(newInputIdx++);
+      if (!newInputArg)
+        return rewriter.notifyMatchFailure(compute, "missing rewritten input block argument");
+      mapper.map(*oldArg, *newInputArg);
+      continue;
+    }
+
+    auto clonedValue = materializeWeightLikeValueInBlock(input, bodyRewriter, mapper);
+    if (failed(clonedValue))
+      return rewriter.notifyMatchFailure(compute, "failed to materialize promoted weight-like operand");
+    mapper.map(*oldArg, *clonedValue);
+  }
+  return success();
+}
+
 // Promotes foldable helper chains from runtime inputs to weights to avoid artificial compute inputs.
 struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCompute> {
   using OpRewritePattern<spatial::SpatCompute>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spatial::SpatCompute compute, PatternRewriter& rewriter) const override {
-    SmallVector<bool> promoteInput(compute.getInputs().size(), false);
-    bool needsRewrite = false;
-    Block& oldBlock = compute.getBody().front();
-    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      if (!isWeightLikeComputeOperand(input))
-        continue;
-      if (isDirectConstantValue(input) && !canPromoteInputBlockArgument(compute.getInputArgument(inputIdx)))
-        continue;
-      promoteInput[inputIdx] = true;
-      needsRewrite = true;
-    }
-    if (!needsRewrite)
+    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);
-
-    SmallVector<Value> newWeights(compute.getWeights().begin(), compute.getWeights().end());
-    SmallVector<Value> newInputs;
-    SmallVector<Type> newInputTypes;
-    SmallVector<Location> newInputLocs;
-    newWeights.reserve(compute.getWeights().size() + compute.getInputs().size());
-    newInputs.reserve(compute.getInputs().size());
-    newInputTypes.reserve(compute.getInputs().size());
-    newInputLocs.reserve(compute.getInputs().size());
-
-    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      if (promoteInput[inputIdx]) {
-        newWeights.push_back(input);
-        continue;
-      }
-      newInputs.push_back(input);
-      newInputTypes.push_back(input.getType());
-      newInputLocs.push_back(input.getLoc());
-    }
-
-    auto newCompute =
-      spatial::SpatCompute::create(rewriter, compute.getLoc(), compute.getResultTypes(), newWeights, newInputs);
+    auto newCompute = spatial::SpatCompute::create(
+      rewriter, compute.getLoc(), compute.getResultTypes(), promoted->newWeights, promoted->newInputs);
     SmallVector<Type> newBlockArgTypes;
     SmallVector<Location> newBlockArgLocs;
-    for (Value weight : newWeights) {
+    for (Value weight : promoted->newWeights) {
       newBlockArgTypes.push_back(weight.getType());
       newBlockArgLocs.push_back(weight.getLoc());
     }
-    llvm::append_range(newBlockArgTypes, newInputTypes);
-    llvm::append_range(newBlockArgLocs, newInputLocs);
+    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>(newWeights.size()), static_cast<int>(newInputs.size())});
+      {static_cast<int>(promoted->newWeights.size()), static_cast<int>(promoted->newInputs.size())});
     rewriter.setInsertionPointToStart(newBlock);
 
     IRRewriter bodyRewriter(rewriter.getContext());
@@ -114,24 +155,14 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCom
         return rewriter.notifyMatchFailure(compute, "missing compute weight block argument during rewrite");
       mapper.map(*oldWeightArg, *newWeightArg);
     }
-    size_t newInputIdx = 0;
-    for (auto [oldInputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      auto oldArg = compute.getInputArgument(oldInputIdx);
-      if (!oldArg)
-        return rewriter.notifyMatchFailure(compute, "missing compute input block argument during rewrite");
-      if (!promoteInput[oldInputIdx]) {
-        auto newInputArg = newCompute.getInputArgument(newInputIdx++);
-        if (!newInputArg)
-          return rewriter.notifyMatchFailure(compute, "missing rewritten compute input block argument");
-        mapper.map(*oldArg, *newInputArg);
-        continue;
-      }
-
-      auto clonedValue = materializeWeightLikeValueInBlock(input, bodyRewriter, mapper);
-      if (failed(clonedValue))
-        return rewriter.notifyMatchFailure(compute, "failed to materialize promoted weight-like operand");
-      mapper.map(*oldArg, *clonedValue);
-    }
+    if (failed(mapPromotedInputArguments(
+          compute,
+          *promoted,
+          bodyRewriter,
+          mapper,
+          [&](size_t index) { return newCompute.getInputArgument(index); },
+          rewriter)))
+      return failure();
 
     for (Operation& op : oldBlock.without_terminator())
       rewriter.clone(op, mapper);
@@ -155,63 +186,35 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
   using OpRewritePattern<spatial::SpatComputeBatch>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spatial::SpatComputeBatch compute, PatternRewriter& rewriter) const override {
-    SmallVector<bool> promoteInput(compute.getInputs().size(), false);
-    bool needsRewrite = false;
-    Block& oldBlock = compute.getBody().front();
-    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      if (!isWeightLikeComputeOperand(input))
-        continue;
-      if (isDirectConstantValue(input) && !canPromoteInputBlockArgument(compute.getInputArgument(inputIdx)))
-        continue;
-      promoteInput[inputIdx] = true;
-      needsRewrite = true;
-    }
-    if (!needsRewrite)
+    auto promoted = computePromotedOperands(compute);
+    if (failed(promoted))
       return rewriter.notifyMatchFailure(compute, "no weight-like batch inputs to promote");
+    Block& oldBlock = compute.getBody().front();
 
     rewriter.setInsertionPointAfter(compute);
 
-    SmallVector<Value> newWeights(compute.getWeights().begin(), compute.getWeights().end());
-    SmallVector<Value> newInputs;
-    SmallVector<Type> newInputTypes;
-    SmallVector<Location> newInputLocs;
-    newWeights.reserve(compute.getWeights().size() + compute.getInputs().size());
-    newInputs.reserve(compute.getInputs().size());
-    newInputTypes.reserve(compute.getInputs().size());
-    newInputLocs.reserve(compute.getInputs().size());
-
-    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      if (promoteInput[inputIdx]) {
-        newWeights.push_back(input);
-        continue;
-      }
-      newInputs.push_back(input);
-      newInputTypes.push_back(input.getType());
-      newInputLocs.push_back(input.getLoc());
-    }
-
-    auto newCompute =
-      spatial::SpatComputeBatch::create(rewriter,
-                                        compute.getLoc(),
-                                        compute.getResultTypes(),
-                                        rewriter.getI32IntegerAttr(static_cast<int32_t>(compute.getLaneCount())),
-                                        newWeights,
-                                        newInputs);
+    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");
     SmallVector<Type> newBlockArgTypes;
     SmallVector<Location> newBlockArgLocs;
-    newBlockArgTypes.reserve(1 + newWeights.size() + newInputTypes.size() + compute.getNumResults());
-    newBlockArgLocs.reserve(1 + newWeights.size() + newInputLocs.size() + compute.getNumResults());
+    newBlockArgTypes.reserve(1 + promoted->newWeights.size() + promoted->newInputTypes.size()
+                             + compute.getNumResults());
+    newBlockArgLocs.reserve(1 + promoted->newWeights.size() + promoted->newInputLocs.size() + compute.getNumResults());
     newBlockArgTypes.push_back(laneArg->getType());
     newBlockArgLocs.push_back(laneArg->getLoc());
-    for (Value weight : newWeights) {
+    for (Value weight : promoted->newWeights) {
       newBlockArgTypes.push_back(weight.getType());
       newBlockArgLocs.push_back(weight.getLoc());
     }
-    llvm::append_range(newBlockArgTypes, newInputTypes);
-    llvm::append_range(newBlockArgLocs, newInputLocs);
+    llvm::append_range(newBlockArgTypes, promoted->newInputTypes);
+    llvm::append_range(newBlockArgLocs, promoted->newInputLocs);
     for (auto [resultIndex, resultType] : llvm::enumerate(compute.getResultTypes())) {
       auto outputArg = compute.getOutputArgument(resultIndex);
       if (!outputArg)
@@ -223,7 +226,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
     auto* newBlock = rewriter.createBlock(
       &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs);
     newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(newWeights.size()), static_cast<int>(newInputs.size())});
+      {static_cast<int>(promoted->newWeights.size()), static_cast<int>(promoted->newInputs.size())});
     rewriter.setInsertionPointToStart(newBlock);
 
     IRRewriter bodyRewriter(rewriter.getContext());
@@ -241,29 +244,20 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
         return rewriter.notifyMatchFailure(compute, "missing compute_batch weight block argument during rewrite");
       mapper.map(*oldWeightArg, *newWeightArg);
     }
-    size_t newInputIdx = 0;
-    for (auto [oldInputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      auto oldArg = compute.getInputArgument(oldInputIdx);
-      if (!oldArg)
-        return rewriter.notifyMatchFailure(compute, "missing compute_batch input block argument during rewrite");
-      if (!promoteInput[oldInputIdx]) {
-        auto newInputArg = newCompute.getInputArgument(newInputIdx++);
-        if (!newInputArg)
-          return rewriter.notifyMatchFailure(compute, "missing rewritten compute_batch input block argument");
-        mapper.map(*oldArg, *newInputArg);
-        continue;
-      }
-
-      auto clonedValue = materializeWeightLikeValueInBlock(input, bodyRewriter, mapper);
-      if (failed(clonedValue))
-        return rewriter.notifyMatchFailure(compute, "failed to materialize promoted batch weight-like operand");
-      mapper.map(*oldArg, *clonedValue);
-    }
+    if (failed(mapPromotedInputArguments(
+          compute,
+          *promoted,
+          bodyRewriter,
+          mapper,
+          [&](size_t index) { return newCompute.getInputArgument(index); },
+          rewriter)))
+      return failure();
     for (auto resultIndex : llvm::seq<size_t>(0, compute.getNumResults())) {
       auto outputArg = compute.getOutputArgument(resultIndex);
       if (!outputArg)
         return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument during rewrite");
-      mapper.map(*outputArg, newBlock->getArgument(1 + newWeights.size() + newInputs.size() + resultIndex));
+      mapper.map(*outputArg,
+                 newBlock->getArgument(1 + promoted->newWeights.size() + promoted->newInputs.size() + resultIndex));
     }
 
     for (Operation& op : oldBlock)
@@ -276,7 +270,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
 
 } // namespace
 
-void populatePostPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+void populateWeightPromotionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   patterns.add<PromoteWeightLikeComputeInputsPattern, PromoteWeightLikeComputeBatchInputsPattern>(ctx);
 }
 

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

@@ -1,6 +1,5 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PrePatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 
 using namespace mlir;
 
@@ -12,7 +11,7 @@ namespace {
 
 } // namespace
 
-void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx) {
+void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx) {
   patterns.add<onnxToArithConstant>(ctx);
   patterns.add<convAddToConvWithBiasLeft>(ctx);
   patterns.add<convAddToConvWithBiasRight>(ctx);

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

@@ -6,7 +6,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -15,24 +15,6 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-static int64_t normalizeAxis(int64_t axis, int64_t rank) { return axis >= 0 ? axis : rank + axis; }
-
-static int64_t normalizeIndex(int64_t index, int64_t dimSize) { return index >= 0 ? index : dimSize + index; }
-
-static Value
-extractSliceAt(Value input, int64_t axis, int64_t offset, ConversionPatternRewriter& rewriter, Location loc) {
-  auto inputType = cast<RankedTensorType>(input.getType());
-  SmallVector<OpFoldResult> offsets(inputType.getRank(), rewriter.getIndexAttr(0));
-  SmallVector<OpFoldResult> sizes;
-  SmallVector<OpFoldResult> strides(inputType.getRank(), rewriter.getIndexAttr(1));
-  sizes.reserve(inputType.getRank());
-  for (int64_t dim : inputType.getShape())
-    sizes.push_back(rewriter.getIndexAttr(dim));
-  offsets[axis] = rewriter.getIndexAttr(offset);
-  sizes[axis] = rewriter.getIndexAttr(1);
-  return tensor::ExtractSliceOp::create(rewriter, loc, input, offsets, sizes, strides);
-}
-
 static Value concatGatherSlices(Value data,
                                 int64_t axis,
                                 ArrayRef<int64_t> indices,
@@ -45,7 +27,7 @@ static Value concatGatherSlices(Value data,
     int64_t normalizedIndex = normalizeIndex(index, axisDim);
     if (normalizedIndex < 0 || normalizedIndex >= axisDim)
       return {};
-    slices.push_back(extractSliceAt(data, axis, normalizedIndex, rewriter, loc));
+    slices.push_back(extractAxisSlice(rewriter, loc, data, axis, normalizedIndex, /*size=*/1));
   }
   if (slices.empty())
     return {};
@@ -96,11 +78,11 @@ struct Gather : OpConversionPattern<ONNXGatherOp> {
       return failure();
 
     int64_t rank = dataType.getRank();
-    int64_t axis = normalizeAxis(gatherOp.getAxis(), rank);
-    if (axis < 0 || axis >= rank)
+    auto axis = normalizeAxisChecked(gatherOp.getAxis(), rank);
+    if (failed(axis))
       return failure();
 
-    int64_t axisDim = dataType.getShape()[axis];
+    int64_t axisDim = dataType.getShape()[*axis];
     if (axisDim <= 0)
       return failure();
 
@@ -116,7 +98,7 @@ struct Gather : OpConversionPattern<ONNXGatherOp> {
                            [&](Value data) -> LogicalResult {
                              Value result;
                              if (indicesType.getRank() == 1) {
-                               result = concatGatherSlices(data, axis, flatIndices, axisDim, rewriter, loc);
+                               result = concatGatherSlices(data, *axis, flatIndices, axisDim, rewriter, loc);
                              }
                              else if (indicesType.getRank() == 2) {
                                int64_t rowCount = indicesType.getShape()[0];
@@ -125,12 +107,13 @@ struct Gather : OpConversionPattern<ONNXGatherOp> {
                                rows.reserve(rowCount);
                                for (int64_t row = 0; row < rowCount; ++row) {
                                  ArrayRef<int64_t> rowIndices(flatIndices.data() + row * rowWidth, rowWidth);
-                                 Value gatheredRow = concatGatherSlices(data, axis, rowIndices, axisDim, rewriter, loc);
+                                 Value gatheredRow =
+                                   concatGatherSlices(data, *axis, rowIndices, axisDim, rewriter, loc);
                                  if (!gatheredRow)
                                    return failure();
-                                 rows.push_back(addLeadingGatherDim(gatheredRow, axis, rewriter, loc));
+                                 rows.push_back(addLeadingGatherDim(gatheredRow, *axis, rewriter, loc));
                                }
-                               result = createSpatConcat(rewriter, loc, /*axis=*/axis, rows);
+                               result = createSpatConcat(rewriter, loc, /*axis=*/*axis, rows);
                              }
                              else {
                                return failure();

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

@@ -5,7 +5,7 @@
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -14,10 +14,6 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-static bool haveStaticPositiveShape(ArrayRef<int64_t> shape) {
-  return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
-}
-
 static bool inferCollapseReassociation(ArrayRef<int64_t> sourceShape,
                                        ArrayRef<int64_t> resultShape,
                                        SmallVector<ReassociationIndices>& reassociation) {
@@ -106,7 +102,7 @@ struct Reshape : OpConversionPattern<ONNXReshapeOp> {
     auto resultType = dyn_cast<RankedTensorType>(reshapeOp.getReshaped().getType());
     if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
       return failure();
-    if (!haveStaticPositiveShape(sourceType.getShape()) || !haveStaticPositiveShape(resultType.getShape()))
+    if (!hasStaticPositiveShape(sourceType) || !hasStaticPositiveShape(resultType))
       return failure();
 
     if (sourceType == resultType) {
@@ -115,17 +111,9 @@ struct Reshape : OpConversionPattern<ONNXReshapeOp> {
     }
 
     auto replaceWithReshape = [&](auto buildReshape) -> LogicalResult {
-      if (isCompileTimeComputable(adaptor.getData())) {
-        rewriter.replaceOp(reshapeOp, buildReshape(adaptor.getData()));
-        return success();
-      }
-
-      auto computeOp = createSpatCompute<1>(
-        rewriter, reshapeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getData(), [&](Value data) {
-          Value reshaped = buildReshape(data);
-          spatial::SpatYieldOp::create(rewriter, reshapeOp.getLoc(), reshaped);
-        });
-      rewriter.replaceOp(reshapeOp, computeOp.getResults());
+      Value reshaped =
+        materializeOrComputeUnary(adaptor.getData(), resultType, rewriter, reshapeOp.getLoc(), buildReshape);
+      rewriter.replaceOp(reshapeOp, reshaped);
       return success();
     };
 

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

@@ -5,8 +5,9 @@
 
 #include "llvm/ADT/STLExtras.h"
 
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -17,19 +18,20 @@ namespace {
 
 static Value buildNearestAsymmetricIndex(
   Value outputIndex, int64_t inputDim, int64_t outputDim, ConversionPatternRewriter& rewriter, Location loc) {
-  Value cInputDim = arith::ConstantIndexOp::create(rewriter, loc, inputDim);
-  Value cOutputDim = arith::ConstantIndexOp::create(rewriter, loc, outputDim);
-  Value cInputDimLast = arith::ConstantIndexOp::create(rewriter, loc, inputDim - 1);
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value cInputDim = getOrCreateIndexConstant(rewriter, anchorOp, inputDim);
+  Value cOutputDim = getOrCreateIndexConstant(rewriter, anchorOp, outputDim);
+  Value cInputDimLast = getOrCreateIndexConstant(rewriter, anchorOp, inputDim - 1);
   Value scaledIndex = arith::MulIOp::create(rewriter, loc, outputIndex, cInputDim);
   Value inputIndex = arith::DivUIOp::create(rewriter, loc, scaledIndex, cOutputDim);
   return arith::MinUIOp::create(rewriter, loc, inputIndex, cInputDimLast);
 }
 
-static Value buildNearestResizeLoop(Value input,
-                                    RankedTensorType inputType,
-                                    RankedTensorType resultType,
-                                    ConversionPatternRewriter& rewriter,
-                                    Location loc) {
+static FailureOr<Value> buildNearestResizeLoop(Value input,
+                                               RankedTensorType inputType,
+                                               RankedTensorType resultType,
+                                               ConversionPatternRewriter& rewriter,
+                                               Location loc) {
   auto elemType = resultType.getElementType();
   SmallVector<int64_t> unitShape(resultType.getRank(), 1);
   auto unitTensorType = RankedTensorType::get(unitShape, elemType);
@@ -37,63 +39,104 @@ static Value buildNearestResizeLoop(Value input,
   SmallVector<OpFoldResult> unitSizes(resultType.getRank(), rewriter.getIndexAttr(1));
   SmallVector<OpFoldResult> unitStrides(resultType.getRank(), rewriter.getIndexAttr(1));
 
-  Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
-  Value c1 = arith::ConstantIndexOp::create(rewriter, loc, 1);
-  Value cOutputN = arith::ConstantIndexOp::create(rewriter, loc, resultType.getDimSize(0));
-  Value cOutputC = arith::ConstantIndexOp::create(rewriter, loc, resultType.getDimSize(1));
-  Value cOutputH = arith::ConstantIndexOp::create(rewriter, loc, resultType.getDimSize(2));
-  Value cOutputW = arith::ConstantIndexOp::create(rewriter, loc, resultType.getDimSize(3));
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+  Value cOutputN = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(0));
+  Value cOutputC = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(1));
+  Value cOutputH = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(2));
+  Value cOutputW = getOrCreateIndexConstant(rewriter, anchorOp, resultType.getDimSize(3));
 
   Value outputInit = tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), elemType);
 
-  auto batchLoop = scf::ForOp::create(rewriter, loc, c0, cOutputN, c1, ValueRange {outputInit});
-  rewriter.setInsertionPointToStart(batchLoop.getBody());
+  auto batchLoop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c0,
+    cOutputN,
+    c1,
+    ValueRange {outputInit},
+    [&](OpBuilder&, Location nestedLoc, Value outputN, ValueRange batchIterArgs, SmallVectorImpl<Value>& batchYielded) {
+      Value outputBatchAcc = batchIterArgs.front();
+      Value inputN =
+        buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), resultType.getDimSize(0), rewriter, nestedLoc);
 
-  Value outputN = batchLoop.getInductionVar();
-  Value outputBatchAcc = batchLoop.getRegionIterArgs().front();
-  Value inputN = buildNearestAsymmetricIndex(outputN, inputType.getDimSize(0), resultType.getDimSize(0), rewriter, loc);
+      auto channelLoop = buildNormalizedScfFor(
+        rewriter,
+        nestedLoc,
+        c0,
+        cOutputC,
+        c1,
+        ValueRange {outputBatchAcc},
+        [&](OpBuilder&,
+            Location channelLoc,
+            Value outputC,
+            ValueRange channelIterArgs,
+            SmallVectorImpl<Value>& channelYielded) {
+          Value outputChannelAcc = channelIterArgs.front();
+          Value inputC = buildNearestAsymmetricIndex(
+            outputC, inputType.getDimSize(1), resultType.getDimSize(1), rewriter, channelLoc);
 
-  auto channelLoop = scf::ForOp::create(rewriter, loc, c0, cOutputC, c1, ValueRange {outputBatchAcc});
-  rewriter.setInsertionPointToStart(channelLoop.getBody());
+          auto heightLoop = buildNormalizedScfFor(
+            rewriter,
+            channelLoc,
+            c0,
+            cOutputH,
+            c1,
+            ValueRange {outputChannelAcc},
+            [&](OpBuilder&,
+                Location heightLoc,
+                Value outputH,
+                ValueRange heightIterArgs,
+                SmallVectorImpl<Value>& heightYielded) {
+              Value outputHeightAcc = heightIterArgs.front();
+              Value inputH = buildNearestAsymmetricIndex(
+                outputH, inputType.getDimSize(2), resultType.getDimSize(2), rewriter, heightLoc);
 
-  Value outputC = channelLoop.getInductionVar();
-  Value outputChannelAcc = channelLoop.getRegionIterArgs().front();
-  Value inputC = buildNearestAsymmetricIndex(outputC, inputType.getDimSize(1), resultType.getDimSize(1), rewriter, loc);
+              auto widthLoop = buildNormalizedScfFor(
+                rewriter,
+                heightLoc,
+                c0,
+                cOutputW,
+                c1,
+                ValueRange {outputHeightAcc},
+                [&](OpBuilder&,
+                    Location widthLoc,
+                    Value outputW,
+                    ValueRange widthIterArgs,
+                    SmallVectorImpl<Value>& widthYielded) {
+                  Value outputWidthAcc = widthIterArgs.front();
+                  Value inputW = buildNearestAsymmetricIndex(
+                    outputW, inputType.getDimSize(3), resultType.getDimSize(3), rewriter, widthLoc);
 
-  auto heightLoop = scf::ForOp::create(rewriter, loc, c0, cOutputH, c1, ValueRange {outputChannelAcc});
-  rewriter.setInsertionPointToStart(heightLoop.getBody());
+                  SmallVector<OpFoldResult> inputOffsets = {inputN, inputC, inputH, inputW};
+                  Value inputSlice = tensor::ExtractSliceOp::create(
+                    rewriter, widthLoc, unitTensorType, input, inputOffsets, unitSizes, unitStrides);
 
-  Value outputH = heightLoop.getInductionVar();
-  Value outputHeightAcc = heightLoop.getRegionIterArgs().front();
-  Value inputH = buildNearestAsymmetricIndex(outputH, inputType.getDimSize(2), resultType.getDimSize(2), rewriter, loc);
-
-  auto widthLoop = scf::ForOp::create(rewriter, loc, c0, cOutputW, c1, ValueRange {outputHeightAcc});
-  rewriter.setInsertionPointToStart(widthLoop.getBody());
-
-  Value outputW = widthLoop.getInductionVar();
-  Value outputWidthAcc = widthLoop.getRegionIterArgs().front();
-  Value inputW = buildNearestAsymmetricIndex(outputW, inputType.getDimSize(3), resultType.getDimSize(3), rewriter, loc);
-
-  SmallVector<OpFoldResult> inputOffsets = {inputN, inputC, inputH, inputW};
-  Value inputSlice =
-    tensor::ExtractSliceOp::create(rewriter, loc, unitTensorType, input, inputOffsets, unitSizes, unitStrides);
-
-  SmallVector<OpFoldResult> outputOffsets = {outputN, outputC, outputH, outputW};
-  Value updatedOutput =
-    tensor::InsertSliceOp::create(rewriter, loc, inputSlice, outputWidthAcc, outputOffsets, unitSizes, unitStrides);
-  scf::YieldOp::create(rewriter, loc, updatedOutput);
-
-  rewriter.setInsertionPointAfter(widthLoop);
-  scf::YieldOp::create(rewriter, loc, widthLoop.getResult(0));
-
-  rewriter.setInsertionPointAfter(heightLoop);
-  scf::YieldOp::create(rewriter, loc, heightLoop.getResult(0));
-
-  rewriter.setInsertionPointAfter(channelLoop);
-  scf::YieldOp::create(rewriter, loc, channelLoop.getResult(0));
-
-  rewriter.setInsertionPointAfter(batchLoop);
-  return batchLoop.getResult(0);
+                  SmallVector<OpFoldResult> outputOffsets = {outputN, outputC, outputH, outputW};
+                  Value updatedOutput = tensor::InsertSliceOp::create(
+                    rewriter, widthLoc, inputSlice, outputWidthAcc, outputOffsets, unitSizes, unitStrides);
+                  widthYielded.push_back(updatedOutput);
+                  return success();
+                });
+              if (failed(widthLoop))
+                return failure();
+              heightYielded.push_back(widthLoop->results.front());
+              return success();
+            });
+          if (failed(heightLoop))
+            return failure();
+          channelYielded.push_back(heightLoop->results.front());
+          return success();
+        });
+      if (failed(channelLoop))
+        return failure();
+      batchYielded.push_back(channelLoop->results.front());
+      return success();
+    });
+  if (failed(batchLoop))
+    return failure();
+  return batchLoop->results.front();
 }
 
 struct Resize : OpConversionPattern<ONNXResizeOp> {
@@ -118,12 +161,17 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
         || llvm::any_of(resultType.getShape(), [](int64_t dim) { return dim <= 0; }))
       return rewriter.notifyMatchFailure(resizeOp, "resize lowering requires positive static dimensions.");
 
-    auto computeOp =
-      createSpatCompute<1>(rewriter, resizeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getX(), [&](Value x) {
-        Value result = buildNearestResizeLoop(x, inputType, resultType, rewriter, resizeOp.getLoc());
-        spatial::SpatYieldOp::create(rewriter, resizeOp.getLoc(), result);
+    auto computeOp = createSpatCompute<1>(
+      rewriter, resizeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getX(), [&](Value x) -> LogicalResult {
+        auto result = buildNearestResizeLoop(x, inputType, resultType, rewriter, resizeOp.getLoc());
+        if (failed(result))
+          return failure();
+        spatial::SpatYieldOp::create(rewriter, resizeOp.getLoc(), *result);
+        return success();
       });
-    rewriter.replaceOp(resizeOp, computeOp.getResults());
+    if (failed(computeOp))
+      return failure();
+    rewriter.replaceOp(resizeOp, computeOp->getResults());
     return success();
   }
 };

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/Split.cpp

@@ -3,7 +3,7 @@
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -12,25 +12,6 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-static int64_t normalizeAxis(int64_t axis, int64_t rank) { return axis >= 0 ? axis : rank + axis; }
-
-static Value extractSliceAt(
-  Value input, int64_t axis, int64_t offset, int64_t size, ConversionPatternRewriter& rewriter, Location loc) {
-  auto inputType = cast<RankedTensorType>(input.getType());
-  SmallVector<OpFoldResult> offsets(inputType.getRank(), rewriter.getIndexAttr(0));
-  SmallVector<OpFoldResult> sizes;
-  SmallVector<OpFoldResult> strides(inputType.getRank(), rewriter.getIndexAttr(1));
-  sizes.reserve(inputType.getRank());
-  for (int64_t dim : inputType.getShape())
-    sizes.push_back(rewriter.getIndexAttr(dim));
-  offsets[axis] = rewriter.getIndexAttr(offset);
-  sizes[axis] = rewriter.getIndexAttr(size);
-  SmallVector<int64_t> resultShape(inputType.getShape());
-  resultShape[axis] = size;
-  auto resultType = RankedTensorType::get(resultShape, inputType.getElementType());
-  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, input, offsets, sizes, strides);
-}
-
 struct Split : OpConversionPattern<ONNXSplitOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -41,8 +22,8 @@ struct Split : OpConversionPattern<ONNXSplitOp> {
       return failure();
 
     int64_t rank = inputType.getRank();
-    int64_t axis = normalizeAxis(splitOp.getAxis(), rank);
-    if (axis < 0 || axis >= rank)
+    auto axis = normalizeAxisChecked(splitOp.getAxis(), rank);
+    if (failed(axis))
       return failure();
 
     SmallVector<Value> outputs;
@@ -58,12 +39,12 @@ struct Split : OpConversionPattern<ONNXSplitOp> {
       if (!resultType || !resultType.hasStaticShape())
         return failure();
       resultTypes.push_back(resultType);
-      sliceSizes.push_back(resultType.getShape()[axis]);
+      sliceSizes.push_back(resultType.getShape()[*axis]);
     }
 
     if (isCompileTimeComputable(adaptor.getInput())) {
       for (int64_t sliceSize : sliceSizes) {
-        outputs.push_back(extractSliceAt(adaptor.getInput(), axis, offset, sliceSize, rewriter, splitOp.getLoc()));
+        outputs.push_back(extractAxisSlice(rewriter, splitOp.getLoc(), adaptor.getInput(), *axis, offset, sliceSize));
         offset += sliceSize;
       }
       rewriter.replaceOp(splitOp, outputs);
@@ -76,7 +57,8 @@ struct Split : OpConversionPattern<ONNXSplitOp> {
         runtimeOutputs.reserve(resultTypes.size());
         int64_t runtimeOffset = 0;
         for (int64_t sliceSize : sliceSizes) {
-          runtimeOutputs.push_back(extractSliceAt(input, axis, runtimeOffset, sliceSize, rewriter, splitOp.getLoc()));
+          runtimeOutputs.push_back(
+            extractAxisSlice(rewriter, splitOp.getLoc(), input, *axis, runtimeOffset, sliceSize));
           runtimeOffset += sliceSize;
         }
         spatial::SpatYieldOp::create(rewriter, splitOp.getLoc(), runtimeOutputs);

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

@@ -0,0 +1,135 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static bool isInsideSpatialComputeRegion(Operation* op) {
+  return op->getParentOfType<spatial::SpatCompute>() || op->getParentOfType<spatial::SpatComputeBatch>();
+}
+
+static Value createTransposeInit(Value input,
+                                 RankedTensorType resultType,
+                                 ArrayRef<int64_t> permutation,
+                                 ConversionPatternRewriter& rewriter,
+                                 Location loc) {
+  SmallVector<OpFoldResult> sizes;
+  sizes.reserve(resultType.getRank());
+  for (auto [resultDim, sourceDim] : llvm::zip_equal(resultType.getShape(), permutation)) {
+    if (!ShapedType::isDynamic(resultDim)) {
+      sizes.push_back(rewriter.getIndexAttr(resultDim));
+      continue;
+    }
+    sizes.push_back(tensor::DimOp::create(rewriter, loc, input, sourceDim).getResult());
+  }
+  return tensor::EmptyOp::create(rewriter, loc, sizes, resultType.getElementType()).getResult();
+}
+
+static FailureOr<Value> materializeTransposedConstant(Value input,
+                                                      RankedTensorType resultType,
+                                                      ArrayRef<int64_t> permutation,
+                                                      ConversionPatternRewriter& rewriter,
+                                                      Location loc) {
+  auto denseAttr = getHostConstDenseElementsAttr(input);
+  if (!denseAttr)
+    return failure();
+
+  auto inputType = dyn_cast<RankedTensorType>(denseAttr.getType());
+  if (!inputType || !inputType.hasStaticShape() || !resultType.hasStaticShape()
+      || inputType.getRank() != resultType.getRank()
+      || static_cast<int64_t>(permutation.size()) != inputType.getRank()) {
+    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;
+  }
+
+  return getOrCreateConstant(rewriter,
+                             rewriter.getInsertionBlock()->getParentOp(),
+                             DenseElementsAttr::get(resultType, resultValues),
+                             resultType);
+}
+
+struct TransposeToLinalgTranspose : OpConversionPattern<ONNXTransposeOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXTransposeOp transposeOp,
+                                ONNXTransposeOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto inputType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
+    auto resultType = dyn_cast<RankedTensorType>(transposeOp.getResult().getType());
+    if (!inputType || !resultType)
+      return failure();
+
+    auto permutation = getTransposePermutationChecked(transposeOp.getPermAttr(), inputType.getRank());
+    if (failed(permutation))
+      return failure();
+    if (isCompileTimeComputable(adaptor.getData())) {
+      auto constantTranspose =
+        materializeTransposedConstant(adaptor.getData(), resultType, *permutation, rewriter, transposeOp.getLoc());
+      if (succeeded(constantTranspose)) {
+        rewriter.replaceOp(transposeOp, *constantTranspose);
+        return success();
+      }
+    }
+
+    auto buildTranspose = [&](Value input) -> Value {
+      Value init = createTransposeInit(input, resultType, *permutation, rewriter, transposeOp.getLoc());
+      return linalg::TransposeOp::create(rewriter, transposeOp.getLoc(), input, init, *permutation).getResult()[0];
+    };
+
+    if (isInsideSpatialComputeRegion(transposeOp.getOperation())) {
+      rewriter.replaceOp(transposeOp, buildTranspose(adaptor.getData()));
+      return success();
+    }
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, transposeOp.getLoc(), TypeRange {resultType}, {}, ValueRange {adaptor.getData()}, [&](Value input) {
+        spatial::SpatYieldOp::create(rewriter, transposeOp.getLoc(), buildTranspose(input));
+      });
+    rewriter.replaceOp(transposeOp, computeOp.getResult(0));
+    return success();
+  }
+};
+
+} // namespace
+
+void populateTransposePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<TransposeToLinalgTranspose>(ctx);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp

@@ -1,18 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/MLIRContext.h"
-
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-namespace onnx_mlir {
-
-bool requiresPostRewrite(spatial::SpatCompute computeOp);
-
-bool requiresPostRewrite(spatial::SpatComputeBatch computeOp);
-
-void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
-void annotateWeightsConstants(mlir::func::FuncOp funcOp);
-
-} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/PrePatterns.hpp

@@ -1,10 +0,0 @@
-#pragma once
-
-#include "mlir/IR/MLIRContext.h"
-#include "mlir/Transforms/DialectConversion.h"
-
-namespace onnx_mlir {
-
-void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -8,7 +8,9 @@
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #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/Dialect/Pim/PimOps.hpp"
 
@@ -18,26 +20,27 @@ using namespace onnx_mlir::pim;
 namespace onnx_mlir {
 namespace {
 
-static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
-  if (isa<pim::PimMemCopyDevToHostOp>(op))
-    return operandIndex == 2;
-  return false;
-}
-
 static bool isUsedOnlyAsExplicitHostOperand(Value value) {
   return !value.use_empty() && llvm::all_of(value.getUses(), [](OpOperand& use) {
-    return isExplicitHostOperand(use.getOwner(), use.getOperandNumber());
+    return isExplicitDevToHostTargetOperand(use.getOwner(), use.getOperandNumber());
   });
 }
 
-static SmallVector<int32_t> getPimCoreIdsForBatchOp(spatial::SpatComputeBatch computeBatchOp, size_t& fallbackCoreId) {
+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)
-    coreIds.push_back(static_cast<int32_t>(fallbackCoreId++));
+  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;
 }
 
@@ -55,7 +58,7 @@ static Value createScaledIndexValue(IRRewriter& rewriter, Location loc, Value ba
   if (scale == 1)
     return base;
 
-  auto scaleValue = arith::ConstantIndexOp::create(rewriter, loc, scale).getResult();
+  auto scaleValue = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), scale);
   return arith::MulIOp::create(rewriter, loc, base, scaleValue).getResult();
 }
 
@@ -64,10 +67,7 @@ static Value createHostTargetOffset(IRRewriter& rewriter,
                                     ShapedType destinationType,
                                     IRMapping& mapper) {
   int64_t elementBytes = static_cast<int64_t>(getElementTypeSizeInBytes(destinationType.getElementType()));
-  SmallVector<int64_t> strides(destinationType.getRank(), 1);
-  ArrayRef<int64_t> shape = destinationType.getShape();
-  for (int64_t dim = destinationType.getRank() - 2; dim >= 0; --dim)
-    strides[dim] = strides[dim + 1] * shape[dim + 1];
+  SmallVector<int64_t> strides = computeRowMajorStrides(destinationType.getShape());
 
   Value totalOffset;
   Location loc = insertSlice.getLoc();
@@ -77,7 +77,8 @@ static Value createHostTargetOffset(IRRewriter& rewriter,
     if (auto attr = dyn_cast<Attribute>(offset)) {
       auto intAttr = dyn_cast<IntegerAttr>(attr);
       assert(intAttr && "expected integer offset attribute");
-      scaledOffset = arith::ConstantIndexOp::create(rewriter, loc, intAttr.getInt() * scale).getResult();
+      scaledOffset =
+        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), intAttr.getInt() * scale);
     }
     else {
       scaledOffset = createScaledIndexValue(rewriter, loc, mapper.lookupOrDefault(cast<Value>(offset)), scale);
@@ -88,7 +89,7 @@ static Value createHostTargetOffset(IRRewriter& rewriter,
   }
 
   if (!totalOffset)
-    totalOffset = arith::ConstantIndexOp::create(rewriter, loc, 0).getResult();
+    totalOffset = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
   return totalOffset;
 }
 
@@ -109,21 +110,24 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
       "resultful compute_batch lowering currently requires a spat.in_parallel terminator");
   }
 
-  SmallVector<int32_t> coreIds = getPimCoreIdsForBatchOp(computeBatchOp, coreId);
+  auto coreIds = getPimCoreIdsForBatchOp(computeBatchOp, coreId);
+  if (failed(coreIds))
+    return failure();
   SmallVector<Value> batchWeights(computeBatchOp.getWeights().begin(), computeBatchOp.getWeights().end());
   SmallVector<Value> batchInputs;
   if (!computeBatchOp.getInputs().empty())
     batchInputs.append(computeBatchOp.getInputs().begin(), computeBatchOp.getInputs().end());
 
   rewriter.setInsertionPointAfter(computeBatchOp);
-  auto coreBatchOp = pim::PimCoreBatchOp::create(rewriter,
-                                                 loc,
-                                                 rewriter.getI32IntegerAttr(computeBatchOp.getLaneCount()),
-                                                 ValueRange(batchWeights),
-                                                 ValueRange(batchInputs));
+  auto laneCountAttr = pim::getCheckedI32Attr(
+    rewriter, computeBatchOp, static_cast<uint64_t>(computeBatchOp.getLaneCount()), "pim core_batch lane count");
+  if (failed(laneCountAttr))
+    return failure();
+  auto coreBatchOp =
+    pim::PimCoreBatchOp::create(rewriter, loc, *laneCountAttr, ValueRange(batchWeights), ValueRange(batchInputs));
   coreBatchOp.getProperties().setOperandSegmentSizes(
     {static_cast<int>(batchWeights.size()), static_cast<int>(batchInputs.size())});
-  coreBatchOp->setAttr(onnx_mlir::kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(coreIds));
+  coreBatchOp->setAttr(onnx_mlir::kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(*coreIds));
 
   SmallVector<unsigned> returnOperandIndices;
   if (computeBatchOp.getNumResults() != 0) {
@@ -166,14 +170,12 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
     BlockArgument newArg = coreBatchOp.getInputArgument(inputIndex);
     auto newArgType = cast<ShapedType>(newArg.getType());
     auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, newArgType);
-    auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
-                                                          loc,
-                                                          outputBuffer.getType(),
-                                                          outputBuffer,
-                                                          newArg,
-                                                          rewriter.getI32IntegerAttr(0),
-                                                          rewriter.getI32IntegerAttr(0),
-                                                          getTensorSizeInBytesAttr(rewriter, newArg))
+    Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), newArg);
+    if (failed(sizeAttr))
+      return failure();
+    auto copied = pim::PimMemCopyHostToDevOp::create(
+                    rewriter, loc, outputBuffer.getType(), zeroOffset, zeroOffset, outputBuffer, newArg, *sizeAttr)
                     .getOutput();
     mapper.map(*oldArg, copied);
   }
@@ -214,7 +216,10 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
         Value hostTarget = getOrCreateHostOutputTensor(resultIndex, insertSlice.getLoc());
         auto hostTargetType = cast<ShapedType>(hostTarget.getType());
         Value hostTargetOffset = createHostTargetOffset(rewriter, insertSlice, hostTargetType, mapper);
-        Value zeroOffset = arith::ConstantIndexOp::create(rewriter, insertSlice.getLoc(), 0).getResult();
+        Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
+        auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), mappedSource);
+        if (failed(sizeAttr))
+          return failure();
         pim::PimMemCopyDevToHostOp::create(rewriter,
                                            insertSlice.getLoc(),
                                            hostTarget.getType(),
@@ -222,7 +227,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
                                            zeroOffset,
                                            hostTarget,
                                            mappedSource,
-                                           getTensorSizeInBytesAttr(rewriter, mappedSource));
+                                           *sizeAttr);
       }
       continue;
     }
@@ -237,15 +242,14 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
         }
         auto clonedType = cast<ShapedType>(clonedTensor.getType());
         auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, clonedType);
-        auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
-                                                              loc,
-                                                              outputBuffer.getType(),
-                                                              outputBuffer,
-                                                              clonedTensor,
-                                                              rewriter.getI32IntegerAttr(0),
-                                                              rewriter.getI32IntegerAttr(0),
-                                                              getTensorSizeInBytesAttr(rewriter, clonedTensor))
-                        .getOutput();
+        Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
+        auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), clonedTensor);
+        if (failed(sizeAttr))
+          return failure();
+        auto copied =
+          pim::PimMemCopyHostToDevOp::create(
+            rewriter, loc, outputBuffer.getType(), zeroOffset, zeroOffset, outputBuffer, clonedTensor, *sizeAttr)
+            .getOutput();
         mapper.map(toTensorOp.getResult(), copied);
         continue;
       }
@@ -254,7 +258,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
     for (auto [operandIndex, operand] : llvm::enumerate(op.getOperands())) {
       if (!isa<TensorType>(operand.getType()) || mapper.contains(operand))
         continue;
-      if (isExplicitHostOperand(&op, operandIndex))
+      if (isExplicitDevToHostTargetOperand(&op, operandIndex))
         continue;
 
       Operation* definingOp = operand.getDefiningOp();

src/PIM/Conversion/SpatialToPim/CMakeLists.txt

@@ -3,15 +3,17 @@ mlir_tablegen(SpatialToPim.hpp.inc -gen-rewriters "-I${ONNX_MLIR_SRC_ROOT}")
 add_public_tablegen_target(SpatialToPimIncGen)
 
 add_pim_library(OMSpatialToPim
+  Patterns.cpp
   SpatialToPimPass.cpp
   BatchCoreLoweringPatterns.cpp
-  ChannelLoweringPatterns.cpp
   Common.cpp
   ComputeLikeRegionUtils.cpp
   CoreLoweringPatterns.cpp
-  GlobalTensorMaterialization.cpp
   ReturnPathNormalization.cpp
-  TensorPackingPatterns.cpp
+  Patterns/ChannelLowering.cpp
+  Patterns/GlobalTensorMaterialization.cpp
+  Patterns/TensorPacking.cpp
+  Patterns/Transpose.cpp
 
   EXCLUDE_FROM_OM_LIBS
 
@@ -19,6 +21,7 @@ add_pim_library(OMSpatialToPim
   SpatialToPimIncGen
 
   LINK_LIBS PUBLIC
+  MLIRLinalgDialect
   MLIRSCFDialect
   MLIRSCFUtils
   MLIRTransformUtils

src/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.hpp

@@ -1,9 +0,0 @@
-#pragma once
-
-#include "mlir/IR/PatternMatch.h"
-
-namespace onnx_mlir {
-
-void populateChannelLoweringPatterns(mlir::RewritePatternSet& patterns);
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,62 +1,22 @@
 #include "mlir/IR/ValueRange.h"
 
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/StringRef.h"
 
 #include <cassert>
-#include <cstddef>
 
 #include "Common.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
 using namespace llvm;
 using namespace mlir;
 
 namespace onnx_mlir {
 
-size_t getSliceActualOffset(tensor::ExtractSliceOp& sliceOp, ShapedType& inputShape) {
-  /*
-    EXAMPLE RUN:
-    [1, 10, 3, 4] inputShape
-    [0, 2, 1, 3] offsets
-
-    acc = 1
-    ---
-    ret = 3
-    acc = 4
-    ---
-    ret = 3 + 4 * 1 = 7
-    acc = 12
-    ---
-    ret = 7 + 12 * 2 = 31
-    acc = 120
-    ---
-    ret = 31 + 120 * 0 = 31
-    acc = 120
-  */
-
-  size_t returnValue = 0;
-
-  auto sliceOffsets = sliceOp.getStaticOffsets();
-  auto inputDimSizes = inputShape.getShape();
-
-  assert(sliceOffsets.size() == inputDimSizes.size());
-
-  size_t accumulatedDimensionSize = 1;
-
-  // Reverse iterate the two vectors
-  for (auto it : reverse(zip(sliceOffsets, inputDimSizes))) {
-    auto curSliceOffset = std::get<0>(it);
-    auto curInputDimSize = std::get<1>(it);
-
-    returnValue += accumulatedDimensionSize * curSliceOffset;
-    accumulatedDimensionSize *= curInputDimSize;
-  }
-
-  return returnValue;
-}
-
-IntegerAttr getTensorSizeInBytesAttr(Builder& builder, mlir::Value value) {
-  return builder.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(value.getType()))));
+FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anchor, mlir::Value value) {
+  auto byteSize = pim::getCheckedShapedTypeSizeInBytes(cast<ShapedType>(value.getType()), anchor, "tensor byte size");
+  if (failed(byteSize))
+    return failure();
+  return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
 }
 
 Operation* getEarliestUserWithinBlock(mlir::Value value) {

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -1,26 +1,14 @@
 #pragma once
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Support/LogicalResult.h"
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 
 namespace onnx_mlir {
 
-/**
- * \brief Get the offset of the ExtractSliceOp based on its static offsets and
- * its static tensor input.
- *
- * The static offsets represent the starting position of the slice in each
- * dimension, while the static tensor input gives its dimension size.
- *
- * \param sliceOp The ExtractSliceOp for which the actual offset needs to be
- * calculated.
- * \param inputShape The ShapedType of the ExtractSliceOp's input tensor
- * \return The actual offset of the ExtractSliceOp.
- */
-size_t getSliceActualOffset(mlir::tensor::ExtractSliceOp& sliceOp, mlir::ShapedType& inputShape);
-
-mlir::IntegerAttr getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Value value);
+mlir::FailureOr<mlir::IntegerAttr>
+getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);
 
 template <class T>
 size_t rangeLength(const mlir::iterator_range<T> range) {

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/IRMapping.h"
@@ -8,6 +9,7 @@
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #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/Dialect/Pim/PimOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -24,7 +26,7 @@ static bool isChannelUseChainOp(Operation* op) {
              tensor::ExpandShapeOp,
              tensor::CastOp,
              tosa::ReshapeOp,
-             ONNXTransposeOp,
+             linalg::TransposeOp,
              pim::PimTransposeOp>(op);
 }
 
@@ -39,7 +41,7 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
       continue;
 
     if (auto constantOp = dyn_cast<arith::ConstantOp>(definingOp)) {
-      mapping.map(operand, getOrCreateHostConstantLike(constantOp, constantFolder));
+      mapping.map(operand, getOrCreateConstantLike(constantFolder, constantOp));
       continue;
     }
 
@@ -53,10 +55,15 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
   }
 }
 
-static int32_t getPimCoreIdForComputeOp(spatial::SpatCompute computeOp, size_t& fallbackCoreId) {
+static FailureOr<int32_t> getPimCoreIdForComputeOp(spatial::SpatCompute computeOp, size_t& fallbackCoreId) {
   if (auto spatialCoreIdAttr = computeOp->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
-    return static_cast<int32_t>(spatialCoreIdAttr.getInt());
-  return static_cast<int32_t>(fallbackCoreId++);
+    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,
@@ -162,10 +169,12 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
       rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg));
       auto outputType = cast<ShapedType>(blockArg->getType());
       auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
-      auto sizeAttr = getTensorSizeInBytesAttr(rewriter, *blockArg);
+      auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
+      if (failed(sizeAttr))
+        return failure();
       Value received =
         PimReceiveOp::create(
-          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer, sizeAttr, receiveOp.getSourceCoreId())
+          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer, *sizeAttr, receiveOp.getSourceCoreId())
           .getOutput();
       blockArg->replaceAllUsesWith(received);
       markOpToRemove(receiveOp);
@@ -205,8 +214,13 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
   if (!computeOp.getWeights().empty())
     computeWeights.append(computeOp.getWeights().begin(), computeOp.getWeights().end());
   rewriter.setInsertionPointAfter(computeOp);
-  auto coreOp = PimCoreOp::create(
-    rewriter, loc, ValueRange(computeWeights), rewriter.getI32IntegerAttr(getPimCoreIdForComputeOp(computeOp, coreId)));
+  auto checkedCoreId = getPimCoreIdForComputeOp(computeOp, coreId);
+  if (failed(checkedCoreId))
+    return failure();
+  auto coreIdAttr = pim::getCheckedI32Attr(rewriter, computeOp, static_cast<int64_t>(*checkedCoreId), "pim core id");
+  if (failed(coreIdAttr))
+    return failure();
+  auto coreOp = PimCoreOp::create(rewriter, loc, ValueRange(computeWeights), *coreIdAttr);
   rewriter.setInsertionPointToStart(&block);
   auto& coreOpBlocks = coreOp.getBody().getBlocks();
   for (auto [inputIndex, input] : llvm::enumerate(computeOp.getInputs())) {
@@ -217,7 +231,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
       continue;
 
     if (auto constantOp = input.getDefiningOp<arith::ConstantOp>()) {
-      blockArg->replaceAllUsesWith(getOrCreateHostConstantLike(constantOp, constantFolder));
+      blockArg->replaceAllUsesWith(getOrCreateConstantLike(constantFolder, constantOp));
       continue;
     }
 
@@ -225,15 +239,18 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
     if (!inputType)
       return computeOp.emitOpError("expected shaped compute input during pim.core lowering");
     auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, inputType);
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), input);
+    if (failed(sizeAttr))
+      return failure();
     auto copied =
       PimMemCopyHostToDevOp::create(rewriter,
                                     loc,
                                     outputBuffer.getType(),
-                                    getOrCreateHostIndexConstant(outputBuffer.getOperation(), 0, constantFolder),
-                                    getOrCreateHostIndexConstant(outputBuffer.getOperation(), 0, constantFolder),
+                                    getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0),
+                                    getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0),
                                     outputBuffer,
                                     input,
-                                    getTensorSizeInBytesAttr(rewriter, input))
+                                    *sizeAttr)
         .getOutput();
     blockArg->replaceAllUsesWith(copied);
   }

src/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.hpp

@@ -1,9 +0,0 @@
-#pragma once
-
-#include "mlir/IR/PatternMatch.h"
-
-namespace onnx_mlir {
-
-void populateGlobalTensorMaterializationPatterns(mlir::RewritePatternSet& patterns);
-
-}

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -0,0 +1,24 @@
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace raptor {
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/SpatialToPim.hpp.inc"
+
+} // namespace raptor
+
+void populateInitialPatterns(RewritePatternSet& patterns) {
+  raptor::populateWithGenerated(patterns);
+  populateTransposeLoweringPatterns(patterns);
+}
+
+void populateCoreBodyPatterns(RewritePatternSet& patterns) {
+  raptor::populateWithGenerated(patterns);
+  populateTransposeLoweringPatterns(patterns);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Patterns.hpp

@@ -8,6 +8,14 @@
 
 namespace onnx_mlir {
 
+void populateInitialPatterns(mlir::RewritePatternSet& patterns);
+void populateCoreBodyPatterns(mlir::RewritePatternSet& patterns);
+
+void populateTransposeLoweringPatterns(mlir::RewritePatternSet& patterns);
+void populateChannelLoweringPatterns(mlir::RewritePatternSet& patterns);
+void populateGlobalTensorMaterializationPatterns(mlir::RewritePatternSet& patterns);
+void populateTensorPackingPatterns(mlir::RewritePatternSet& patterns);
+
 mlir::RankedTensorType getPackedTensorType(mlir::RankedTensorType elementType, int64_t count);
 mlir::Value extractPackedChunk(mlir::Value packedValue,
                                mlir::RankedTensorType chunkType,
@@ -20,7 +28,6 @@ mlir::Value createPackedExtractRowsSlice(spatial::SpatExtractRowsOp extractRowsO
                                          mlir::OpBuilder& builder,
                                          mlir::Location loc);
 mlir::Value createPackedExtractSliceTensor(mlir::ValueRange values, mlir::OpBuilder& builder, mlir::Location loc);
-void populateTensorPackingPatterns(mlir::RewritePatternSet& patterns);
 void eraseUnusedTensorPackingOps(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
 
 } // namespace onnx_mlir

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

@@ -1,7 +1,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.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"
 
@@ -14,8 +14,10 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spatial::SpatChannelSendOp op, PatternRewriter& rewriter) const override {
-    pim::PimSendOp::create(
-      rewriter, op.getLoc(), op.getInput(), getTensorSizeInBytesAttr(rewriter, op.getInput()), op.getTargetCoreId());
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
+    if (failed(sizeAttr))
+      return failure();
+    pim::PimSendOp::create(rewriter, op.getLoc(), op.getInput(), *sizeAttr, op.getTargetCoreId());
     rewriter.eraseOp(op);
     return success();
   }
@@ -32,12 +34,11 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     auto outputType = cast<ShapedType>(op.getResult().getType());
     Value outputBuffer =
       tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
-    Value received = pim::PimReceiveOp::create(rewriter,
-                                               op.getLoc(),
-                                               op.getResult().getType(),
-                                               outputBuffer,
-                                               getTensorSizeInBytesAttr(rewriter, op.getResult()),
-                                               op.getSourceCoreId())
+    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();
     rewriter.replaceOp(op, received);
     return success();

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

@@ -16,7 +16,7 @@
 
 #include "Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;

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

@@ -1,4 +1,4 @@
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/TensorPackingPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;

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

@@ -0,0 +1,38 @@
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct LinalgTransposeToPim final : OpRewritePattern<linalg::TransposeOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, PatternRewriter& rewriter) const override {
+    SmallVector<Attribute> permutationAttrs;
+    permutationAttrs.reserve(transposeOp.getPermutation().size());
+    for (int64_t dim : transposeOp.getPermutation())
+      permutationAttrs.push_back(rewriter.getI64IntegerAttr(dim));
+
+    auto permutation = rewriter.getArrayAttr(permutationAttrs);
+    auto pimTranspose = pim::PimTransposeOp::create(rewriter,
+                                                    transposeOp.getLoc(),
+                                                    TypeRange {transposeOp->getResult(0).getType()},
+                                                    transposeOp.getInput(),
+                                                    permutation,
+                                                    transposeOp.getInit());
+    rewriter.replaceOp(transposeOp, pimTranspose.getOutput());
+    return success();
+  }
+};
+
+} // namespace
+
+void populateTransposeLoweringPatterns(RewritePatternSet& patterns) {
+  patterns.add<LinalgTransposeToPim>(patterns.getContext());
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
@@ -11,6 +12,7 @@
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #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/Dialect/Pim/PimOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -40,7 +42,7 @@ static bool isReturnHelperChainOp(Operation* op) {
              tensor::ExpandShapeOp,
              tensor::CastOp,
              tosa::ReshapeOp,
-             ONNXTransposeOp,
+             linalg::TransposeOp,
              pim::PimTransposeOp>(op);
 }
 
@@ -70,6 +72,20 @@ static SmallVector<int64_t> expandFlatElementIndex(int64_t flatIndex, ArrayRef<i
   return indices;
 }
 
+static FailureOr<int64_t>
+getCheckedByteOffset(int64_t elementOffset, size_t elementSize, Operation* anchor, StringRef fieldName) {
+  if (elementOffset < 0) {
+    anchor->emitOpError() << fieldName << " requires a nonnegative element offset";
+    return failure();
+  }
+
+  auto byteOffset =
+    pim::checkedMul(static_cast<uint64_t>(elementOffset), static_cast<uint64_t>(elementSize), anchor, fieldName);
+  if (failed(byteOffset))
+    return failure();
+  return pim::checkedCast<int64_t>(*byteOffset, anchor, fieldName);
+}
+
 static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
                                                SmallVectorImpl<Operation*>& helperChain) {
   if (computeOp.getInputs().size() != 1 || computeOp.getNumResults() != 1)
@@ -276,11 +292,10 @@ static LogicalResult mapIndicesThroughHelperChain(ArrayRef<int64_t> sourceIndice
       continue;
     }
 
-    if (auto transposeOp = dyn_cast<ONNXTransposeOp>(op)) {
+    if (auto transposeOp = dyn_cast<linalg::TransposeOp>(op)) {
       SmallVector<int64_t> nextIndices(currentIndices.size());
       SmallVector<int64_t> nextShape(currentShape.size());
-      for (auto [destIndex, attr] : llvm::enumerate(transposeOp.getPermAttr().getAsRange<IntegerAttr>())) {
-        int64_t sourceIndex = attr.getInt();
+      for (auto [destIndex, sourceIndex] : llvm::enumerate(transposeOp.getPermutation())) {
         nextIndices[destIndex] = currentIndices[sourceIndex];
         nextShape[destIndex] = currentShape[sourceIndex];
       }
@@ -326,7 +341,7 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
       continue;
 
     if (auto constantOp = dyn_cast<arith::ConstantOp>(definingOp)) {
-      mapping.map(operand, getOrCreateHostConstantLike(constantOp, constantFolder));
+      mapping.map(operand, getOrCreateConstantLike(constantFolder, constantOp));
       continue;
     }
 
@@ -360,18 +375,72 @@ static void cloneHelperChain(Value sourceValue,
   }
 }
 
-static Value emitHostCopy(IRRewriter& rewriter,
-                          Location loc,
-                          Value outputTensor,
-                          Value sourceValue,
-                          int32_t hostTargetOffset,
-                          int32_t deviceSourceOffset,
-                          int32_t sizeInBytes,
-                          OperationFolder& constantFolder) {
+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,
+                                     Value sourceValue,
+                                     int64_t hostTargetOffset,
+                                     int64_t deviceSourceOffset,
+                                     uint64_t sizeInBytes,
+                                     OperationFolder& constantFolder) {
   Operation* anchorOp = sourceValue.getDefiningOp() ? sourceValue.getDefiningOp() : outputTensor.getDefiningOp();
   assert(anchorOp && "expected a concrete op anchor for return-path host copy constants");
-  Value hostTargetOffsetValue = getOrCreateHostIndexConstant(anchorOp, hostTargetOffset, constantFolder);
-  Value deviceSourceOffsetValue = getOrCreateHostIndexConstant(anchorOp, deviceSourceOffset, constantFolder);
+  Value hostTargetOffsetValue = getOrCreateIndexConstant(constantFolder, anchorOp, hostTargetOffset);
+  Value deviceSourceOffsetValue = getOrCreateIndexConstant(constantFolder, anchorOp, deviceSourceOffset);
+  auto sizeAttr = pim::getCheckedI32Attr(rewriter, anchorOp, sizeInBytes, "return-path host copy byte size");
+  if (failed(sizeAttr))
+    return failure();
   return PimMemCopyDevToHostOp::create(rewriter,
                                        loc,
                                        outputTensor.getType(),
@@ -379,7 +448,7 @@ static Value emitHostCopy(IRRewriter& rewriter,
                                        deviceSourceOffsetValue,
                                        outputTensor,
                                        sourceValue,
-                                       rewriter.getI32IntegerAttr(sizeInBytes))
+                                       *sizeAttr)
     .getOutput();
 }
 
@@ -426,25 +495,45 @@ 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)
       markOpToRemove(op);
 
     auto storedType = cast<ShapedType>(currentStoredValue.getType());
-    size_t elementSize = getElementTypeSizeInBytes(storedType.getElementType());
+    auto byteSize = pim::getCheckedShapedTypeSizeInBytes(storedType, producerOp, "return-path host copy byte size");
+    if (failed(byteSize))
+      return ReturnPathLoweringResult::Failure;
     if (auto storedOp = currentStoredValue.getDefiningOp())
       rewriter.setInsertionPointAfter(storedOp);
     Value outputTensor = outputTensors[returnUse->returnIndex](rewriter, loc);
-    emitHostCopy(rewriter,
-                 loc,
-                 outputTensor,
-                 currentStoredValue,
-                 0,
-                 0,
-                 static_cast<int32_t>(storedType.getNumElements() * elementSize),
-                 constantFolder);
+    auto copied = emitHostCopy(rewriter, loc, outputTensor, currentStoredValue, 0, 0, *byteSize, constantFolder);
+    if (failed(copied))
+      return ReturnPathLoweringResult::Failure;
     return ReturnPathLoweringResult::Handled;
   }
 
@@ -455,23 +544,27 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
 
     if (isa<func::ReturnOp>(resultUser)) {
       size_t resultIndexInReturn = resultUse.getOperandNumber();
-      size_t elementSize = getElementTypeSizeInBytes(storedTensorType.getElementType());
+      if (isHostStaticReturnValue(storedValue))
+        return materializeDirectHostReturn(resultIndexInReturn, storedValue, {});
+      auto byteSize =
+        pim::getCheckedShapedTypeSizeInBytes(storedTensorType, producerOp, "return-path host copy byte size");
+      if (failed(byteSize))
+        return ReturnPathLoweringResult::Failure;
       rewriter.setInsertionPointAfterValue(storedValue);
       Value outputTensor = outputTensors[resultIndexInReturn](rewriter, loc);
-      emitHostCopy(rewriter,
-                   loc,
-                   outputTensor,
-                   storedValue,
-                   0,
-                   0,
-                   static_cast<int32_t>(storedTensorType.getNumElements() * elementSize),
-                   constantFolder);
+      auto copied = emitHostCopy(rewriter, loc, outputTensor, storedValue, 0, 0, *byteSize, constantFolder);
+      if (failed(copied))
+        return ReturnPathLoweringResult::Failure;
       return ReturnPathLoweringResult::Handled;
     }
   }
 
   if (auto concatReturnUse = analyzeConcatReturnUse(producedValue)) {
     size_t elementSize = getElementTypeSizeInBytes(storedTensorType.getElementType());
+    auto storedByteSize =
+      pim::getCheckedShapedTypeSizeInBytes(storedTensorType, producerOp, "concat return-path copy byte size");
+    if (failed(storedByteSize))
+      return ReturnPathLoweringResult::Failure;
     for (Operation* concatOp : concatReturnUse->concatChain)
       markOpToRemove(concatOp);
 
@@ -480,14 +573,13 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
       Value outputTensor = outputTensors[concatReturnUse->returnIndex](rewriter, loc);
       auto outputType = cast<ShapedType>(outputTensor.getType());
       int64_t flatOffset = computeFlatElementIndex(concatReturnUse->sliceOffsets, outputType.getShape());
-      emitHostCopy(rewriter,
-                   loc,
-                   outputTensor,
-                   storedValue,
-                   static_cast<int32_t>(flatOffset * elementSize),
-                   0,
-                   static_cast<int32_t>(storedTensorType.getNumElements() * elementSize),
-                   constantFolder);
+      auto hostOffset = getCheckedByteOffset(flatOffset, elementSize, producerOp, "concat return-path host offset");
+      if (failed(hostOffset))
+        return ReturnPathLoweringResult::Failure;
+      auto copied =
+        emitHostCopy(rewriter, loc, outputTensor, storedValue, *hostOffset, 0, *storedByteSize, constantFolder);
+      if (failed(copied))
+        return ReturnPathLoweringResult::Failure;
       return ReturnPathLoweringResult::Handled;
     }
 
@@ -531,14 +623,18 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
       rewriter.setInsertionPointAfter(elementSlice);
 
       int64_t destinationFlatOffset = computeFlatElementIndex(destinationIndices, outputType.getShape());
-      outputTensor = emitHostCopy(rewriter,
-                                  loc,
-                                  outputTensor,
-                                  elementSlice.getResult(),
-                                  static_cast<int32_t>(destinationFlatOffset * elementSize),
-                                  0,
-                                  static_cast<int32_t>(elementSize),
-                                  constantFolder);
+      auto hostOffset =
+        getCheckedByteOffset(destinationFlatOffset, elementSize, producerOp, "concat helper return-path host offset");
+      if (failed(hostOffset))
+        return ReturnPathLoweringResult::Failure;
+      auto elementByteSize = pim::checkedCast<uint64_t>(elementSize, producerOp, "return-path scalar copy byte size");
+      if (failed(elementByteSize))
+        return ReturnPathLoweringResult::Failure;
+      auto copied = emitHostCopy(
+        rewriter, loc, outputTensor, elementSlice.getResult(), *hostOffset, 0, *elementByteSize, constantFolder);
+      if (failed(copied))
+        return ReturnPathLoweringResult::Failure;
+      outputTensor = *copied;
     }
     return ReturnPathLoweringResult::Handled;
   }

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

@@ -9,12 +9,6 @@ include "src/Accelerators/PIM/Dialect/Spatial/Spatial.td"
 include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
 #endif // OP_BASE
 
-def onnxToPimTranspose : Pat<
-  (ONNXTransposeOp:$srcOpRes $data, $perms),
-  (PimTransposeOp $data, $perms,
-    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
->;
-
 def spatToPimVMM : Pat<
   (SpatVMMOp:$srcOpRes $weight, $vector),
   (PimVMMOp $weight, $vector,
@@ -33,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,12 +25,12 @@
 #include <cassert>
 #include <utility>
 
+#include "Common/IR/ConstantUtils.hpp"
 #include "Common/PimCommon.hpp"
+#include "Common/Support/CheckedArithmetic.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "Conversion/SpatialToPim/ChannelLoweringPatterns.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
-#include "Conversion/SpatialToPim/GlobalTensorMaterialization.hpp"
-#include "Conversion/SpatialToPim/TensorPackingPatterns.hpp"
+#include "Conversion/SpatialToPim/Patterns.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Spatial/SpatialOps.hpp"
 #include "Pass/PIMPasses.h"
@@ -41,11 +41,6 @@ using namespace onnx_mlir;
 using namespace pim;
 
 namespace onnx_mlir {
-namespace raptor {
-
-#include "src/Accelerators/PIM/Conversion/SpatialToPim/SpatialToPim.hpp.inc"
-
-} // namespace raptor
 
 static memref::GlobalOp getOrCreateZeroGlobal(IRRewriter& rewriter, Location loc, RankedTensorType tensorType) {
   auto moduleOp = rewriter.getBlock()->getParentOp()->getParentOfType<ModuleOp>();
@@ -81,33 +76,28 @@ static memref::GlobalOp getOrCreateZeroGlobal(IRRewriter& rewriter, Location loc
                                   IntegerAttr {});
 }
 
-static Value createZeroedDeviceHVector(IRRewriter& rewriter,
-                                       Location loc,
-                                       RankedTensorType tensorType,
-                                       OperationFolder& constantFolder) {
+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 = getOrCreateHostIndexConstant(outputBuffer.getOperation(), 0, constantFolder);
-  auto sizeAttr = rewriter.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(tensorType)));
-
-  if (outputBuffer->getParentOfType<PimCoreBatchOp>())
-    return PimMemCopyHostToDevBatchOp::create(rewriter,
-                                              loc,
-                                              tensorType,
-                                              outputBuffer,
-                                              zeroValue,
-                                              rewriter.getI32IntegerAttr(0),
-                                              rewriter.getI32IntegerAttr(0),
-                                              sizeAttr)
-      .getOutput();
-
+  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)
+           rewriter, loc, tensorType, zeroIndex, zeroIndex, outputBuffer, zeroValue, *sizeAttr)
     .getOutput();
 }
 
-static Value
+static FailureOr<Value>
 padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector, OperationFolder& constantFolder) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
@@ -119,10 +109,18 @@ padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector,
 
   auto paddedType = RankedTensorType::get(
     {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
-  Value zeroed = createZeroedDeviceHVector(rewriter, loc, paddedType, constantFolder);
-  auto zeroAttr = rewriter.getI32IntegerAttr(0);
-  auto sizeAttr = rewriter.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(vectorType)));
-  return PimMemCopyOp::create(rewriter, loc, paddedType, zeroed, vector, zeroAttr, zeroAttr, sizeAttr).getOutput();
+  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();
 }
 
 void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
@@ -159,7 +157,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
                     spatial::SpatExtractRowsOp>();
 
   RewritePatternSet initialPatterns(ctx);
-  populateWithGenerated(initialPatterns);
+  populateInitialPatterns(initialPatterns);
   if (failed(applyPartialConversion(moduleOp, target, std::move(initialPatterns)))) {
     moduleOp.emitError("failed to lower required Spatial ops to the initial PIM form");
     signalPassFailure();
@@ -214,7 +212,7 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   }
 
   RewritePatternSet coreBodyPatterns(ctx);
-  populateWithGenerated(coreBodyPatterns);
+  populateCoreBodyPatterns(coreBodyPatterns);
   populateAffineToStdConversionPatterns(coreBodyPatterns);
   FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));
 
@@ -252,7 +250,11 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     }
   }
 
-  enlargeVMMOutTensorsToCrossbarSize(funcOp, rewriter);
+  if (failed(enlargeVMMOutTensorsToCrossbarSize(funcOp, rewriter))) {
+    funcOp.emitOpError("failed to enlarge VMM output tensors to crossbar size");
+    signalPassFailure();
+    return;
+  }
   replaceReturnWithOutputBuffers(returnOp, rewriter);
   eraseOpsToRemove();
 
@@ -283,12 +285,15 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     signalPassFailure();
     return;
   }
+  hoistAndUniquifyIndexConstants(funcOp, rewriter);
+
   // Dump to file for debug
   dumpModule(moduleOp, "pim0");
 }
 
-void raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
+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());
     ArrayRef<int64_t> outputShape = outputType.getShape();
@@ -296,19 +301,23 @@ void raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp f
     assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");
 
     rewriter.setInsertionPoint(vmmOp);
-    Value paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput(), constantFolder);
+    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput(), constantFolder);
+    if (failed(paddedInput)) {
+      hasFailure = true;
+      return WalkResult::interrupt();
+    }
     auto paddedOutputType = RankedTensorType::get(
       {outputShape[0], static_cast<int64_t>(crossbarSize)}, outputType.getElementType(), outputType.getEncoding());
     Value paddedOutputBuffer = outputShape[1] == static_cast<int64_t>(crossbarSize)
                                ? vmmOp.getOutputBuffer()
                                : createEmptyTensorFromShaped(rewriter, vmmOp.getLoc(), paddedOutputType).getResult();
-    vmmOp.getInputMutable().assign(paddedInput);
+    vmmOp.getInputMutable().assign(*paddedInput);
     vmmOp.getOutputBufferMutable().assign(paddedOutputBuffer);
 
     vmmOp.getOutput().setType(paddedOutputType);
 
     if (outputShape[1] == static_cast<int64_t>(crossbarSize))
-      return;
+      return WalkResult::advance();
 
     SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
     SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(outputShape[0]), rewriter.getIndexAttr(outputShape[1])};
@@ -318,13 +327,16 @@ void raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp f
       tensor::ExtractSliceOp::create(rewriter, vmmOp.getLoc(), outputType, vmmOp.getOutput(), offsets, sizes, strides);
     SmallPtrSet<Operation*, 2> exceptions = {vmmOp, sliceOp};
     vmmOp.getOutput().replaceAllUsesExcept(sliceOp.getResult(), exceptions);
+    return WalkResult::advance();
   });
+  return success(!hasFailure);
 }
 
 LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(func::FuncOp funcOp,
                                                                                 IRRewriter& rewriter) {
   Location loc = funcOp.getLoc();
   OperationFolder constantFolder(funcOp.getContext());
+  bool hasFailure = false;
 
   auto insertMemCopyHostToDev = [&](Value inputTensor, int64_t elementsOffset) {
     auto tensorType = cast<ShapedType>(inputTensor.getType());
@@ -335,17 +347,28 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
     rewriter.setInsertionPointAfter(inputTensor.getDefiningOp());
 
     auto deviceTensor = tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), elementType);
+    auto offsetBytes = pim::checkedMul(
+      static_cast<size_t>(elementsOffset), elementByteSize, deviceTensor.getOperation(), "host input byte offset");
+    auto byteSize =
+      pim::getCheckedShapedTypeSizeInBytes(tensorType, deviceTensor.getOperation(), "host input copy byte size");
+    auto sizeAttr =
+      succeeded(byteSize)
+        ? pim::getCheckedI32Attr(rewriter, deviceTensor.getOperation(), *byteSize, "host input copy byte size")
+        : FailureOr<IntegerAttr>(failure());
+    if (failed(offsetBytes) || failed(sizeAttr)) {
+      hasFailure = true;
+      return;
+    }
 
     auto memCopyHostToDevOp = PimMemCopyHostToDevOp::create(
       rewriter,
       loc,
       tensorType,
-      getOrCreateHostIndexConstant(deviceTensor.getOperation(), 0, constantFolder),
-      getOrCreateHostIndexConstant(
-        deviceTensor.getOperation(), static_cast<int64_t>(elementsOffset * elementByteSize), constantFolder),
+      getOrCreateIndexConstant(constantFolder, deviceTensor.getOperation(), 0),
+      getOrCreateIndexConstant(constantFolder, deviceTensor.getOperation(), static_cast<int64_t>(*offsetBytes)),
       deviceTensor,
       inputTensor,
-      rewriter.getI32IntegerAttr(static_cast<int32_t>(tensorType.getNumElements() * elementByteSize)));
+      *sizeAttr);
 
     rewriter.replaceAllUsesExcept(inputTensor, memCopyHostToDevOp.getResult(), {memCopyHostToDevOp});
   };
@@ -363,7 +386,7 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
       }
     }
 
-  return success();
+  return success(!hasFailure);
 }
 
 void raptor::SpatialToPimPass::markOpToRemove(Operation* op) {

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -64,7 +64,7 @@ private:
   void markOpToRemove(mlir::Operation* op);
   void eraseOpsToRemove();
 
-  void enlargeVMMOutTensorsToCrossbarSize(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
+  mlir::LogicalResult enlargeVMMOutTensorsToCrossbarSize(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
 };
 
 } // namespace raptor

src/PIM/Dialect/Pim/CMakeLists.txt

@@ -2,7 +2,9 @@ add_onnx_mlir_dialect(Pim pim)
 add_onnx_mlir_dialect_doc(pim Pim.td)
 
 add_subdirectory(Transforms/Bufferization)
-add_subdirectory(Transforms/StaticMemoryCoalescing)
+add_subdirectory(Transforms/MemoryCoalescing)
+add_subdirectory(Transforms/HostConstantFolding)
+add_subdirectory(Transforms/Verification)
 
 add_pim_library(PimOps
   PimOps.hpp

src/PIM/Dialect/Pim/Pim.td

@@ -144,32 +144,6 @@ def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   }];
 }
 
-def PimMemCopyHostToDevBatchOp : PimOp<"memcp_hd_batch", [DestinationStyleOpInterface]> {
-  let summary = "Copy a per-lane tensor from host memory into device memory inside a batched core";
-
-  let arguments = (ins
-    PimTensor:$deviceTarget,
-    PimTensor:$hostSource,
-    I32Attr:$deviceTargetOffset,
-    I32Attr:$hostSourceOffset,
-    I32Attr:$size
-  );
-
-  let results = (outs
-    PimTensor:$output
-  );
-
-  let extraClassDeclaration = [{
-    mlir::MutableOperandRange getDpsInitsMutable() {
-      return getDeviceTargetMutable();
-    }
-  }];
-
-  let assemblyFormat = [{
-    `(` $deviceTarget `,` $hostSource `)` attr-dict `:` `(` type($deviceTarget) `,` type($hostSource) `)` `->` type($output)
-  }];
-}
-
 def PimMemCopyDevToHostOp : PimOp<"memcp_dh", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from device memory into host memory";
 
@@ -202,10 +176,10 @@ def PimMemCopyOp : PimOp<"memcp", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region within the same memory space";
 
   let arguments = (ins
+    Index:$targetOffset,
+    Index:$sourceOffset,
     PimTensor:$target,
     PimTensor:$source,
-    I32Attr:$targetOffset,
-    I32Attr:$sourceOffset,
     I32Attr:$size
   );
 
@@ -220,7 +194,9 @@ def PimMemCopyOp : PimOp<"memcp", [DestinationStyleOpInterface]> {
   }];
 
   let assemblyFormat = [{
-    `(` $target `,` $source `)` attr-dict `:` `(` type($target) `,` type($source) `)` `->` type($output)
+    `[` $targetOffset `,` $sourceOffset `]`
+    `(` $target `,` $source `)` attr-dict
+    `:` type($target) `,` type($source) `->` type($output)
   }];
 }
 

src/PIM/Dialect/Pim/PimOpsVerify.cpp

@@ -9,6 +9,7 @@
 #include "llvm/Support/LogicalResult.h"
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -19,16 +20,6 @@ namespace pim {
 
 namespace {
 
-static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
-  if (isa<PimMemCopyHostToDevOp>(op))
-    return operandIndex == 3;
-  if (isa<PimMemCopyHostToDevBatchOp>(op))
-    return operandIndex == 1;
-  if (isa<PimMemCopyDevToHostOp>(op))
-    return operandIndex == 2;
-  return false;
-}
-
 static Region* getParentRegion(Value value) {
   if (auto blockArgument = dyn_cast<BlockArgument>(value))
     return blockArgument.getParentRegion();
@@ -63,7 +54,7 @@ static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region
     for (OpOperand& operand : op->getOpOperands()) {
       Value value = operand.get();
       if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value)
-          || isExplicitHostOperand(op, operand.getOperandNumber()))
+          || isExplicitHostMemCopyOperand(op, operand.getOperandNumber()))
         continue;
 
       InFlightDiagnostic diagnostic = ownerOp->emitOpError()

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

@@ -1,37 +1,52 @@
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/Common.hpp"
 
 using namespace mlir;
 using namespace bufferization;
 
 namespace onnx_mlir::pim {
 
-Value materializeContiguousMemRef(Value memrefValue, Location loc, RewriterBase& rewriter) {
-  if (succeeded(resolveContiguousAddress(memrefValue)))
+FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue, Location loc, RewriterBase& rewriter) {
+  bool isContiguous =
+    succeeded(resolveContiguousAddress(memrefValue)) || succeeded(compileContiguousAddressExpr(memrefValue));
+  if (isContiguous && isDeviceLocalPimAddress(memrefValue))
     return memrefValue;
 
   auto shapedType = cast<ShapedType>(memrefValue.getType());
   auto contiguousType = MemRefType::get(shapedType.getShape(), shapedType.getElementType());
   Value contiguousBuffer = memref::AllocOp::create(rewriter, loc, contiguousType);
-  auto sizeInBytes = getShapedTypeSizeInBytes(shapedType);
+  auto sizeInBytes =
+    getCheckedShapedTypeSizeInBytes(shapedType, contiguousBuffer.getDefiningOp(), "contiguous copy byte size");
+  if (failed(sizeInBytes))
+    return failure();
+  Value zeroOffset = getOrCreateIndexConstant(rewriter, contiguousBuffer.getDefiningOp(), 0);
+  auto sizeAttr =
+    getCheckedI32Attr(rewriter, contiguousBuffer.getDefiningOp(), *sizeInBytes, "contiguous copy byte size");
+  if (failed(sizeAttr))
+    return failure();
 
-  return PimMemCopyOp::create(rewriter,
-                              loc,
-                              contiguousType,
-                              contiguousBuffer,
-                              memrefValue,
-                              rewriter.getI32IntegerAttr(0),
-                              rewriter.getI32IntegerAttr(0),
-                              rewriter.getI32IntegerAttr(sizeInBytes))
+  if (isHostBackedPimAddress(memrefValue)) {
+    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 allocateContiguousMemRefLike(Value memrefValue, Location loc, RewriterBase& rewriter) {
-  if (succeeded(resolveContiguousAddress(memrefValue)))
+Value allocateContiguousResultMemRefLike(Value memrefValue, Location loc, RewriterBase& rewriter) {
+  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

@@ -5,8 +5,10 @@
 
 namespace onnx_mlir::pim {
 
-mlir::Value materializeContiguousMemRef(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
-mlir::Value allocateContiguousMemRefLike(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
+llvm::FailureOr<mlir::Value>
+materializeContiguousInputMemRef(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
+mlir::Value
+allocateContiguousResultMemRefLike(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
 
 llvm::FailureOr<mlir::Value> getBufferOrValue(mlir::RewriterBase& rewriter,
                                               mlir::Value value,

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

@@ -1,11 +1,9 @@
-set(LLVM_TARGET_DEFINITIONS PimBufferization.td)
-mlir_tablegen(PimBufferization.hpp.inc -gen-rewriters "-I${ONNX_MLIR_SRC_ROOT}")
-add_public_tablegen_target(PimBufferizationIncGen)
-
 add_pim_library(OMPimBufferization
   PimBufferizationPass.cpp
   BufferizationUtils.hpp
   BufferizationUtils.cpp
+  ContiguityPatterns.hpp
+  ContiguityPatterns.cpp
   OpBufferizationInterfaces.hpp
   OpBufferizationInterfaces.cpp
   Common.hpp
@@ -13,9 +11,6 @@ add_pim_library(OMPimBufferization
 
   EXCLUDE_FROM_OM_LIBS
 
-  DEPENDS
-  PimBufferizationIncGen
-
   LINK_LIBS PUBLIC
   OMPimCommon
   PimOps

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

@@ -1,10 +1,111 @@
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
 #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;
 
-IntegerAttr onnx_mlir::pim::getMemRefSizeInBytesAttr(OpBuilder& builder, Value memref) {
-  auto type = mlir::cast<MemRefType>(memref.getType());
-  int32_t sizeInBytes = static_cast<int32_t>(getShapedTypeSizeInBytes(type));
-  return builder.getI32IntegerAttr(sizeInBytes);
+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");
+  if (failed(byteSize))
+    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) {
+  auto base = getPimStorageBase(value, knowledge);
+  if (failed(base))
+    return false;
+
+  if (isCoreBatchInputArgument(*base))
+    return true;
+
+  return isa_and_nonnull<memref::GetGlobalOp>(base->getDefiningOp());
+}
+
+bool onnx_mlir::pim::isDeviceLocalPimAddress(Value value, const StaticValueKnowledge& knowledge) {
+  auto base = getPimStorageBase(value, knowledge);
+  if (failed(base))
+    return false;
+
+  return isa_and_nonnull<memref::AllocOp>(base->getDefiningOp());
 }

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

@@ -2,10 +2,19 @@
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+
 namespace onnx_mlir {
 namespace pim {
 
-mlir::IntegerAttr getMemRefSizeInBytesAttr(mlir::OpBuilder& builder, mlir::Value memref);
+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

@@ -0,0 +1,473 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+
+#include "ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir::pim {
+namespace {
+
+struct ByteOffsetTerm {
+  Value value;
+  int64_t scale = 0;
+};
+
+struct ByteOffsetExpr {
+  int64_t constant = 0;
+  SmallVector<ByteOffsetTerm> terms;
+};
+
+struct CopyEndpointPlan {
+  Value base;
+  MemRefType originalType;
+  MemRefType baseType;
+  ByteOffsetExpr offset;
+};
+
+struct CopyLoopPlan {
+  SmallVector<int64_t> outerShape;
+  int64_t chunkBytes = 0;
+  ByteOffsetExpr targetBaseOffset;
+  ByteOffsetExpr sourceBaseOffset;
+  SmallVector<int64_t> targetOuterByteStrides;
+  SmallVector<int64_t> sourceOuterByteStrides;
+};
+
+struct CopyRewritePlan {
+  enum class Kind {
+    Direct,
+    Loop
+  } kind = Kind::Direct;
+  CopyEndpointPlan target;
+  CopyEndpointPlan source;
+  int64_t directBytes = 0;
+  CopyLoopPlan loop;
+};
+
+static bool isViewLike(Value value) {
+  Operation* defOp = value.getDefiningOp();
+  return defOp
+      && isa<memref::SubViewOp,
+             memref::ReinterpretCastOp,
+             memref::CollapseShapeOp,
+             memref::ExpandShapeOp,
+             memref::CastOp>(defOp);
+}
+
+template <typename CopyOp>
+static bool isNormalizedCopyLikeOp(CopyOp copyOp, Value target, Value source, Value targetOffset, Value sourceOffset) {
+  auto targetType = dyn_cast<MemRefType>(target.getType());
+  auto sourceType = dyn_cast<MemRefType>(source.getType());
+  return targetType && sourceType && !isViewLike(target) && !isViewLike(source) && targetOffset.getType().isIndex()
+      && sourceOffset.getType().isIndex() && copyOp.getSize() > 0;
+}
+
+static void appendTerm(ByteOffsetExpr& expr, Value value, int64_t scale) {
+  if (scale != 0)
+    expr.terms.push_back(ByteOffsetTerm {value, scale});
+}
+
+static FailureOr<SmallVector<int64_t>> getStaticMemRefStrides(MemRefType type) {
+  SmallVector<int64_t> strides;
+  int64_t offset = 0;
+  if (failed(type.getStridesAndOffset(strides, offset)))
+    return failure();
+  if (llvm::any_of(strides, ShapedType::isDynamic))
+    return failure();
+  return strides;
+}
+
+static FailureOr<int64_t> getShapedByteSize(MemRefType type) {
+  if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType()))
+    return failure();
+  auto byteSize =
+    pim::getCheckedShapedTypeSizeInBytes(type, UnknownLoc::get(type.getContext()), "normalized copy byte size");
+  if (failed(byteSize))
+    return failure();
+  if (*byteSize > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()))
+    return failure();
+  return static_cast<int64_t>(*byteSize);
+}
+
+static FailureOr<SmallVector<int64_t>>
+inferLogicalCopyShape(MemRefType targetType, MemRefType sourceType, int64_t size) {
+  if (!targetType.hasStaticShape() || !sourceType.hasStaticShape())
+    return failure();
+  if (targetType.getElementType() != sourceType.getElementType() || targetType.getRank() != sourceType.getRank())
+    return failure();
+
+  auto targetBytes = getShapedByteSize(targetType);
+  auto sourceBytes = getShapedByteSize(sourceType);
+  if (failed(targetBytes) || failed(sourceBytes))
+    return failure();
+
+  bool targetMatches = *targetBytes == size;
+  bool sourceMatches = *sourceBytes == size;
+  if (targetMatches && sourceMatches && targetType.getShape() != sourceType.getShape())
+    return failure();
+  if (targetMatches)
+    return SmallVector<int64_t>(targetType.getShape().begin(), targetType.getShape().end());
+  if (sourceMatches)
+    return SmallVector<int64_t>(sourceType.getShape().begin(), sourceType.getShape().end());
+  return failure();
+}
+
+static FailureOr<int64_t> getContiguousSuffixRank(MemRefType type, ArrayRef<int64_t> copyShape) {
+  if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())
+      || type.getRank() != static_cast<int64_t>(copyShape.size()))
+    return failure();
+
+  auto strides = getStaticMemRefStrides(type);
+  if (failed(strides))
+    return failure();
+
+  int64_t expectedStride = 1;
+  int64_t contiguousSuffixRank = 0;
+  for (int64_t dim = type.getRank() - 1; dim >= 0; --dim) {
+    if ((*strides)[dim] != expectedStride)
+      break;
+    ++contiguousSuffixRank;
+    expectedStride *= copyShape[dim];
+  }
+  return contiguousSuffixRank;
+}
+
+static FailureOr<CopyEndpointPlan> analyzeCopyEndpoint(Value value, Value initialByteOffset, MemRefType logicalType) {
+  if (!logicalType.hasStaticShape() || !hasByteSizedElementType(logicalType.getElementType()))
+    return failure();
+
+  CopyEndpointPlan endpoint;
+  endpoint.base = value;
+  endpoint.originalType = logicalType;
+  appendTerm(endpoint.offset, initialByteOffset, 1);
+
+  while (true) {
+    if (auto castOp = endpoint.base.getDefiningOp<memref::CastOp>()) {
+      endpoint.base = castOp.getSource();
+      continue;
+    }
+    if (auto collapseOp = endpoint.base.getDefiningOp<memref::CollapseShapeOp>()) {
+      endpoint.base = collapseOp.getSrc();
+      continue;
+    }
+    if (auto expandOp = endpoint.base.getDefiningOp<memref::ExpandShapeOp>()) {
+      endpoint.base = expandOp.getSrc();
+      continue;
+    }
+    if (auto reinterpretOp = endpoint.base.getDefiningOp<memref::ReinterpretCastOp>()) {
+      endpoint.base = reinterpretOp.getSource();
+      continue;
+    }
+
+    auto subviewOp = endpoint.base.getDefiningOp<memref::SubViewOp>();
+    if (!subviewOp)
+      break;
+
+    auto sourceType = dyn_cast<MemRefType>(subviewOp.getSource().getType());
+    if (!sourceType || !sourceType.hasStaticShape() || !hasByteSizedElementType(sourceType.getElementType()))
+      return failure();
+
+    auto sourceStrides = getStaticMemRefStrides(sourceType);
+    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;
+      if (auto attr = dyn_cast<Attribute>(offset)) {
+        endpoint.offset.constant += cast<IntegerAttr>(attr).getInt() * byteScale;
+        continue;
+      }
+      appendTerm(endpoint.offset, cast<Value>(offset), byteScale);
+    }
+
+    endpoint.base = subviewOp.getSource();
+  }
+
+  endpoint.baseType = dyn_cast<MemRefType>(endpoint.base.getType());
+  if (!endpoint.baseType)
+    return failure();
+  return endpoint;
+}
+
+static FailureOr<CopyRewritePlan>
+analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceOffset, int64_t size) {
+  auto targetType = dyn_cast<MemRefType>(target.getType());
+  auto sourceType = dyn_cast<MemRefType>(source.getType());
+  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);
+  if (failed(targetSuffixRank) || failed(sourceSuffixRank))
+    return failure();
+
+  CopyRewritePlan plan;
+  plan.target = *targetPlan;
+  plan.source = *sourcePlan;
+
+  int64_t contiguousSuffixRank = std::min(*targetSuffixRank, *sourceSuffixRank);
+  if (contiguousSuffixRank == static_cast<int64_t>(logicalCopyShape->size())) {
+    plan.kind = CopyRewritePlan::Kind::Direct;
+    plan.directBytes = size;
+    return plan;
+  }
+
+  auto targetStrides = getStaticMemRefStrides(targetType);
+  auto sourceStrides = getStaticMemRefStrides(sourceType);
+  if (failed(targetStrides) || failed(sourceStrides))
+    return failure();
+
+  int64_t elementByteWidth = static_cast<int64_t>(getElementTypeSizeInBytes(targetType.getElementType()));
+  plan.kind = CopyRewritePlan::Kind::Loop;
+  plan.loop.targetBaseOffset = plan.target.offset;
+  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);
+  if (plan.loop.chunkBytes <= 0)
+    return failure();
+  return plan;
+}
+
+static Value createIndexConstant(PatternRewriter& rewriter, Operation* anchorOp, int64_t value) {
+  return getOrCreateIndexConstant(rewriter, anchorOp, value);
+}
+
+static Value addIndexValues(PatternRewriter& rewriter, Location loc, Value lhs, Value rhs) {
+  if (auto constant = getConstantIntValue(lhs); constant && *constant == 0)
+    return rhs;
+  if (auto constant = getConstantIntValue(rhs); constant && *constant == 0)
+    return lhs;
+  return arith::AddIOp::create(rewriter, loc, lhs, rhs).getResult();
+}
+
+static Value mulIndexValue(PatternRewriter& rewriter, Location loc, Operation* anchorOp, Value value, int64_t scale) {
+  if (scale == 0)
+    return createIndexConstant(rewriter, anchorOp, 0);
+  if (scale == 1)
+    return value;
+  Value scaleValue = createIndexConstant(rewriter, anchorOp, scale);
+  return arith::MulIOp::create(rewriter, loc, value, scaleValue).getResult();
+}
+
+static Value
+materializeByteOffset(PatternRewriter& rewriter, Location loc, Operation* anchorOp, const ByteOffsetExpr& expr) {
+  Value result = createIndexConstant(rewriter, anchorOp, expr.constant);
+  for (const ByteOffsetTerm& term : expr.terms)
+    result = addIndexValues(rewriter, loc, result, mulIndexValue(rewriter, loc, anchorOp, term.value, term.scale));
+  return result;
+}
+
+static SmallVector<Value> materializeDelinearizedIndices(
+  PatternRewriter& rewriter, Location loc, Operation* anchorOp, Value linearIndex, ArrayRef<int64_t> shape) {
+  SmallVector<Value> indices;
+  if (shape.empty())
+    return indices;
+
+  auto rowMajorStrides = computeRowMajorStrides(shape);
+  Value remaining = linearIndex;
+  for (auto [dim, stride] : llvm::enumerate(rowMajorStrides)) {
+    if (dim + 1 == rowMajorStrides.size()) {
+      indices.push_back(remaining);
+      break;
+    }
+    Value strideValue = createIndexConstant(rewriter, anchorOp, stride);
+    Value index = arith::DivUIOp::create(rewriter, loc, remaining, strideValue);
+    indices.push_back(index);
+    remaining = arith::RemUIOp::create(rewriter, loc, remaining, strideValue);
+  }
+  return indices;
+}
+
+static Value materializeOuterByteOffset(PatternRewriter& rewriter,
+                                        Location loc,
+                                        Operation* anchorOp,
+                                        const ByteOffsetExpr& baseOffset,
+                                        ArrayRef<Value> outerIndices,
+                                        ArrayRef<int64_t> outerByteStrides) {
+  Value result = materializeByteOffset(rewriter, loc, anchorOp, baseOffset);
+  for (auto [index, stride] : llvm::zip_equal(outerIndices, outerByteStrides))
+    result = addIndexValues(rewriter, loc, result, mulIndexValue(rewriter, loc, anchorOp, index, stride));
+  return result;
+}
+
+template <typename CopyOp, typename CreateCopyOp>
+static LogicalResult rewriteCopyLikeOp(CopyOp copyOp,
+                                       Value target,
+                                       Value source,
+                                       Value targetOffset,
+                                       Value sourceOffset,
+                                       Value replacementValue,
+                                       CreateCopyOp createCopyOp,
+                                       PatternRewriter& rewriter) {
+  if (isNormalizedCopyLikeOp(copyOp, target, source, targetOffset, sourceOffset))
+    return failure();
+
+  auto plan = analyzeCopyRewrite(target, source, targetOffset, sourceOffset, copyOp.getSize());
+  if (failed(plan))
+    return failure();
+
+  Location loc = copyOp.getLoc();
+  Operation* anchorOp = copyOp.getOperation();
+  if (plan->kind == CopyRewritePlan::Kind::Direct) {
+    Value newTargetOffset = materializeByteOffset(rewriter, loc, anchorOp, plan->target.offset);
+    Value newSourceOffset = materializeByteOffset(rewriter, loc, anchorOp, plan->source.offset);
+    auto checkedDirectBytes = pim::checkedI32(plan->directBytes, anchorOp, "normalized direct copy byte size");
+    if (failed(checkedDirectBytes))
+      return failure();
+    auto newCopyOp =
+      createCopyOp(loc, plan->target.base, plan->source.base, newTargetOffset, newSourceOffset, *checkedDirectBytes);
+    assert(isNormalizedCopyOp(newCopyOp) && "copy normalization emitted a non-normalized copy");
+    rewriter.replaceOp(copyOp, replacementValue);
+    return success();
+  }
+
+  Value c0 = createIndexConstant(rewriter, anchorOp, 0);
+  Value cUpper = createIndexConstant(rewriter, anchorOp, getNumElements(plan->loop.outerShape));
+  Value cStep = createIndexConstant(rewriter, anchorOp, 1);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    loc,
+    c0,
+    cUpper,
+    cStep,
+    ValueRange {},
+    [&](OpBuilder&, Location nestedLoc, Value inductionVar, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+      SmallVector<Value> outerIndices =
+        materializeDelinearizedIndices(rewriter, nestedLoc, anchorOp, inductionVar, plan->loop.outerShape);
+      Value loopTargetOffset = materializeOuterByteOffset(
+        rewriter, nestedLoc, anchorOp, plan->loop.targetBaseOffset, outerIndices, plan->loop.targetOuterByteStrides);
+      Value loopSourceOffset = materializeOuterByteOffset(
+        rewriter, nestedLoc, anchorOp, plan->loop.sourceBaseOffset, outerIndices, plan->loop.sourceOuterByteStrides);
+      auto checkedChunkBytes = pim::checkedI32(plan->loop.chunkBytes, anchorOp, "normalized loop copy byte size");
+      if (failed(checkedChunkBytes))
+        return failure();
+      auto newCopyOp = createCopyOp(
+        nestedLoc, plan->target.base, plan->source.base, loopTargetOffset, loopSourceOffset, *checkedChunkBytes);
+      assert(isNormalizedCopyOp(newCopyOp) && "copy normalization emitted a non-normalized copy");
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  rewriter.replaceOp(copyOp, replacementValue);
+  return success();
+}
+
+struct NormalizeCoreSubviewCopyPattern final : OpRewritePattern<pim::PimMemCopyOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(pim::PimMemCopyOp copyOp, PatternRewriter& rewriter) const override {
+    return rewriteCopyLikeOp(
+      copyOp,
+      copyOp.getTarget(),
+      copyOp.getSource(),
+      copyOp.getTargetOffset(),
+      copyOp.getSourceOffset(),
+      copyOp.getTarget(),
+      [&](Location loc, Value target, Value source, Value targetOffset, Value sourceOffset, int32_t size) {
+        return pim::PimMemCopyOp::create(rewriter,
+                                         loc,
+                                         cast<MemRefType>(target.getType()),
+                                         targetOffset,
+                                         sourceOffset,
+                                         target,
+                                         source,
+                                         rewriter.getI32IntegerAttr(size));
+      },
+      rewriter);
+  }
+};
+
+struct NormalizeHostSubviewLoadPattern final : OpRewritePattern<pim::PimMemCopyHostToDevOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(pim::PimMemCopyHostToDevOp copyOp, PatternRewriter& rewriter) const override {
+    return rewriteCopyLikeOp(
+      copyOp,
+      copyOp.getDeviceTarget(),
+      copyOp.getHostSource(),
+      copyOp.getDeviceTargetOffset(),
+      copyOp.getHostSourceOffset(),
+      copyOp.getDeviceTarget(),
+      [&](Location loc, Value target, Value source, Value targetOffset, Value sourceOffset, int32_t size) {
+        return pim::PimMemCopyHostToDevOp::create(rewriter,
+                                                  loc,
+                                                  cast<MemRefType>(target.getType()),
+                                                  targetOffset,
+                                                  sourceOffset,
+                                                  target,
+                                                  source,
+                                                  rewriter.getI32IntegerAttr(size));
+      },
+      rewriter);
+  }
+};
+
+struct NormalizeHostSubviewStorePattern final : OpRewritePattern<pim::PimMemCopyDevToHostOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(pim::PimMemCopyDevToHostOp copyOp, PatternRewriter& rewriter) const override {
+    return rewriteCopyLikeOp(
+      copyOp,
+      copyOp.getHostTarget(),
+      copyOp.getDeviceSource(),
+      copyOp.getHostTargetOffset(),
+      copyOp.getDeviceSourceOffset(),
+      copyOp.getHostTarget(),
+      [&](Location loc, Value target, Value source, Value targetOffset, Value sourceOffset, int32_t size) {
+        return pim::PimMemCopyDevToHostOp::create(rewriter,
+                                                  loc,
+                                                  cast<MemRefType>(target.getType()),
+                                                  targetOffset,
+                                                  sourceOffset,
+                                                  target,
+                                                  source,
+                                                  rewriter.getI32IntegerAttr(size));
+      },
+      rewriter);
+  }
+};
+
+} // namespace
+
+bool isNormalizedCopyOp(pim::PimMemCopyOp op) {
+  return isNormalizedCopyLikeOp(op, op.getTarget(), op.getSource(), op.getTargetOffset(), op.getSourceOffset());
+}
+
+bool isNormalizedCopyOp(pim::PimMemCopyHostToDevOp op) {
+  return isNormalizedCopyLikeOp(
+    op, op.getDeviceTarget(), op.getHostSource(), op.getDeviceTargetOffset(), op.getHostSourceOffset());
+}
+
+bool isNormalizedCopyOp(pim::PimMemCopyDevToHostOp op) {
+  return isNormalizedCopyLikeOp(
+    op, op.getHostTarget(), op.getDeviceSource(), op.getHostTargetOffset(), op.getDeviceSourceOffset());
+}
+
+void populatePimContiguityNormalizationPatterns(RewritePatternSet& patterns) {
+  patterns.add<NormalizeCoreSubviewCopyPattern, NormalizeHostSubviewLoadPattern, NormalizeHostSubviewStorePattern>(
+    patterns.getContext());
+}
+
+} // namespace onnx_mlir::pim

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

@@ -0,0 +1,15 @@
+#pragma once
+
+#include "mlir/IR/PatternMatch.h"
+
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+namespace onnx_mlir::pim {
+
+bool isNormalizedCopyOp(pim::PimMemCopyOp op);
+bool isNormalizedCopyOp(pim::PimMemCopyHostToDevOp op);
+bool isNormalizedCopyOp(pim::PimMemCopyDevToHostOp op);
+
+void populatePimContiguityNormalizationPatterns(mlir::RewritePatternSet& patterns);
+
+} // namespace onnx_mlir::pim

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

@@ -47,32 +47,6 @@ struct MemCopyHostToDevOpInterface
   }
 };
 
-struct MemCopyHostToDevBatchOpInterface
-: DstBufferizableOpInterfaceExternalModel<MemCopyHostToDevBatchOpInterface, PimMemCopyHostToDevBatchOp> {
-  LogicalResult bufferize(Operation* op,
-                          RewriterBase& rewriter,
-                          const BufferizationOptions& options,
-                          BufferizationState& state) const {
-    auto memCopyHostToDevOp = cast<PimMemCopyHostToDevBatchOp>(op);
-    auto deviceTargetOpt = getBufferOrValue(rewriter, memCopyHostToDevOp.getDeviceTarget(), options, state);
-    if (failed(deviceTargetOpt))
-      return failure();
-    auto hostSourceOpt = getBufferOrValue(rewriter, memCopyHostToDevOp.getHostSource(), options, state);
-    if (failed(hostSourceOpt))
-      return failure();
-
-    replaceOpWithNewBufferizedOp<PimMemCopyHostToDevBatchOp>(rewriter,
-                                                             memCopyHostToDevOp,
-                                                             deviceTargetOpt->getType(),
-                                                             *deviceTargetOpt,
-                                                             *hostSourceOpt,
-                                                             memCopyHostToDevOp.getDeviceTargetOffsetAttr(),
-                                                             memCopyHostToDevOp.getHostSourceOffsetAttr(),
-                                                             memCopyHostToDevOp.getSizeAttr());
-    return success();
-  }
-};
-
 struct MemCopyDevToHostOpInterface
 : DstBufferizableOpInterfaceExternalModel<MemCopyDevToHostOpInterface, PimMemCopyDevToHostOp> {
   LogicalResult bufferize(Operation* op,
@@ -127,10 +101,10 @@ struct MemCopyOpInterface : DstBufferizableOpInterfaceExternalModel<MemCopyOpInt
     replaceOpWithNewBufferizedOp<PimMemCopyOp>(rewriter,
                                                memCopyOp,
                                                targetOpt->getType(),
+                                               memCopyOp.getTargetOffset(),
+                                               memCopyOp.getSourceOffset(),
                                                *targetOpt,
                                                *sourceOpt,
-                                               memCopyOp.getTargetOffsetAttr(),
-                                               memCopyOp.getSourceOffsetAttr(),
                                                memCopyOp.getSizeAttr());
     return success();
   }
@@ -151,8 +125,9 @@ struct ReceiveOpInterface : DstBufferizableOpInterfaceExternalModel<ReceiveOpInt
     if (failed(outputBufferOpt))
       return failure();
 
+    Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
     replaceOpWithNewBufferizedOp<PimReceiveOp>(
-      rewriter, op, outputBufferOpt->getType(), *outputBufferOpt, receiveOp.getSizeAttr(), receiveOp.getSourceCoreId());
+      rewriter, op, contiguousOutput.getType(), contiguousOutput, receiveOp.getSizeAttr(), receiveOp.getSourceCoreId());
     return success();
   }
 };
@@ -173,15 +148,19 @@ struct ConcatOpInterface : DstBufferizableOpInterfaceExternalModel<ConcatOpInter
       auto inputOpt = getBufferOrValue(rewriter, input, options, state);
       if (failed(inputOpt))
         return failure();
-      inputs.push_back(materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter));
+      auto contiguous = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+      if (failed(contiguous))
+        return failure();
+      inputs.push_back(*contiguous);
     }
 
     auto outputBufferOpt = getBufferOrValue(rewriter, concatOp.getOutputBuffer(), options, state);
     if (failed(outputBufferOpt))
       return failure();
 
+    Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
     replaceOpWithNewBufferizedOp<PimConcatOp>(
-      rewriter, op, outputBufferOpt->getType(), concatOp.getAxisAttr(), ValueRange(inputs), *outputBufferOpt);
+      rewriter, op, contiguousOutput.getType(), concatOp.getAxisAttr(), ValueRange(inputs), contiguousOutput);
     return success();
   }
 };
@@ -203,12 +182,12 @@ 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);
+    if (failed(contiguousInput))
+      return failure();
 
-    replaceOpWithNewBufferizedOp<PimSendOp>(rewriter,
-                                            op,
-                                            materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter),
-                                            sendOp.getSizeAttr(),
-                                            sendOp.getTargetCoreId());
+    replaceOpWithNewBufferizedOp<PimSendOp>(
+      rewriter, op, *contiguousInput, sendOp.getSizeAttr(), sendOp.getTargetCoreId());
     return success();
   }
 };
@@ -431,11 +410,13 @@ struct TransposeOpInterface : DstBufferizableOpInterfaceExternalModel<TransposeO
     if (failed(outputBufferOpt))
       return failure();
 
-    Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);
-    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
+    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    if (failed(contiguousInput))
+      return failure();
+    Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
 
     replaceOpWithNewBufferizedOp<PimTransposeOp>(
-      rewriter, op, contiguousOutput.getType(), contiguousInput, transposeOp.getPermutation(), contiguousOutput);
+      rewriter, op, contiguousOutput.getType(), *contiguousInput, transposeOp.getPermutation(), contiguousOutput);
     return success();
   }
 };
@@ -475,11 +456,13 @@ struct VMMOpInterface : DstBufferizableOpInterfaceExternalModel<VMMOpInterface,
     if (failed(outputBufferOpt))
       return failure();
 
-    Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);
-    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
+    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    if (failed(contiguousInput))
+      return failure();
+    Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
 
     replaceOpWithNewBufferizedOp<PimVMMOp>(
-      rewriter, op, contiguousOutput.getType(), *weightOpt, contiguousInput, contiguousOutput);
+      rewriter, op, contiguousOutput.getType(), *weightOpt, *contiguousInput, contiguousOutput);
     return success();
   }
 };
@@ -514,12 +497,16 @@ struct BinaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<BinaryDstO
     if (failed(outputBufferOpt))
       return failure();
 
-    Value contiguousLhs = materializeContiguousMemRef(*lhsOpt, op->getLoc(), rewriter);
-    Value contiguousRhs = materializeContiguousMemRef(*rhsOpt, op->getLoc(), rewriter);
-    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
+    auto contiguousLhs = materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter);
+    if (failed(contiguousLhs))
+      return failure();
+    auto contiguousRhs = materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter);
+    if (failed(contiguousRhs))
+      return failure();
+    Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
 
     replaceOpWithNewBufferizedOp<OpTy>(
-      rewriter, op, contiguousOutput.getType(), contiguousLhs, contiguousRhs, contiguousOutput);
+      rewriter, op, contiguousOutput.getType(), *contiguousLhs, *contiguousRhs, contiguousOutput);
     return success();
   }
 };
@@ -547,12 +534,16 @@ struct VVDMulOpInterface : DstBufferizableOpInterfaceExternalModel<VVDMulOpInter
     if (failed(outputBufferOpt))
       return failure();
 
-    Value contiguousLhs = materializeContiguousMemRef(*lhsOpt, op->getLoc(), rewriter);
-    Value contiguousRhs = materializeContiguousMemRef(*rhsOpt, op->getLoc(), rewriter);
-    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
+    auto contiguousLhs = materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter);
+    if (failed(contiguousLhs))
+      return failure();
+    auto contiguousRhs = materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter);
+    if (failed(contiguousRhs))
+      return failure();
+    Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
 
     replaceOpWithNewBufferizedOp<PimVVDMulOp>(
-      rewriter, op, contiguousOutput.getType(), contiguousLhs, contiguousRhs, contiguousOutput);
+      rewriter, op, contiguousOutput.getType(), *contiguousLhs, *contiguousRhs, contiguousOutput);
     return success();
   }
 };
@@ -583,10 +574,12 @@ struct UnaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<UnaryDstOpI
     if (failed(outputBufferOpt))
       return failure();
 
-    Value contiguousInput = materializeContiguousMemRef(*inputOpt, op->getLoc(), rewriter);
-    Value contiguousOutput = allocateContiguousMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
+    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    if (failed(contiguousInput))
+      return failure();
+    Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
 
-    replaceOpWithNewBufferizedOp<OpTy>(rewriter, op, contiguousOutput.getType(), contiguousInput, contiguousOutput);
+    replaceOpWithNewBufferizedOp<OpTy>(rewriter, op, contiguousOutput.getType(), *contiguousInput, contiguousOutput);
     return success();
   }
 };
@@ -599,7 +592,6 @@ void registerOpBufferizationInterfaces(DialectRegistry& registry) {
     PimSendOp::attachInterface<SendOpInterface>(*ctx);
     PimConcatOp::attachInterface<ConcatOpInterface>(*ctx);
     PimMemCopyHostToDevOp::attachInterface<MemCopyHostToDevOpInterface>(*ctx);
-    PimMemCopyHostToDevBatchOp::attachInterface<MemCopyHostToDevBatchOpInterface>(*ctx);
     PimMemCopyDevToHostOp::attachInterface<MemCopyDevToHostOpInterface>(*ctx);
     PimMemCopyOp::attachInterface<MemCopyOpInterface>(*ctx);
     PimTransposeOp::attachInterface<TransposeOpInterface>(*ctx);

src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferization.td

@@ -1,19 +0,0 @@
-#ifndef PIM_BUFFERIZATION
-#define PIM_BUFFERIZATION
-
-#ifndef OP_BASE
-include "mlir/IR/PatternBase.td"
-include "mlir/Dialect/MemRef/IR/MemRefOps.td"
-include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
-#endif // OP_BASE
-
-def memrefCopyToPimMemCopyOp : Pat<
-  (CopyOp $src, $dst),
-  (PimMemCopyOp $dst, $src,
-    ConstantAttr<I32Attr, "0">,
-    ConstantAttr<I32Attr, "0">,
-    (NativeCodeCall<"pim::getMemRefSizeInBytesAttr($_builder, $0)"> $src),
-    (returnType $dst))
->;
-
-#endif // PIM_BUFFERIZATION