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:

src/PIM/CMakeLists.txt

@@ -123,7 +123,6 @@ add_pim_library(OMPIMAccel
   OMPimBufferization
   OMPimMemoryCoalescing
   OMPimHostConstantFolding
-  OMPimHostConstantMaterialization
   OMPimVerification
   MLIRTensorInferTypeOpInterfaceImpl
 )

src/PIM/Common/IR/WeightUtils.cpp

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

src/PIM/Common/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
@@ -30,7 +31,6 @@ add_pim_library(OMPimCompilerUtils
   OMPimBufferization
   OMPimMemoryCoalescing
   OMPimHostConstantFolding
-  OMPimHostConstantMaterialization
   OMPimVerification
   OMPimPasses
   OMONNXToSpatial

src/PIM/Compiler/PimCodeGen.cpp

@@ -2,7 +2,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -27,20 +26,24 @@
 #include <fstream>
 #include <limits>
 #include <memory>
+#include <numeric>
 #include <string>
 #include <utility>
 
 #include "Common/IR/CompactAsmUtils.hpp"
 #include "Common/PimCommon.hpp"
+#include "Common/Support/Diagnostics.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
 #include "Common/Support/ReportUtils.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
 #include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
 #include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -96,6 +99,51 @@ static int32_t getVectorByteSizeOrCrash(ShapedType type) {
   return pim::checkedI32OrCrash(*byteSize, "vector byte size");
 }
 
+static Operation* getDiagnosticAnchor(mlir::Value value) {
+  if (Operation* definingOp = value.getDefiningOp())
+    return definingOp;
+  if (auto blockArg = dyn_cast<BlockArgument>(value))
+    return blockArg.getOwner()->getParentOp();
+  return nullptr;
+}
+
+// PIM instruction immediates are serialized as signed int32_t fields today
+// (`sldi` goes through checkedI32OrCrash), so local addresses must stay within
+// the non-negative int32_t range.
+static constexpr size_t kPimAddressLimit = static_cast<size_t>(std::numeric_limits<int32_t>::max());
+
+static FailureOr<size_t> checkedAlignTo(size_t value, size_t alignment, Operation* anchor, StringRef fieldName) {
+  if (alignment == 0)
+    return value;
+  size_t remainder = value % alignment;
+  if (remainder == 0)
+    return value;
+  return pim::checkedAdd(value, alignment - remainder, anchor, fieldName);
+}
+
+static void printMemoryOverflowDiagnostic(mlir::Value value,
+                                          const MemoryValueKey& key,
+                                          size_t requestedSize,
+                                          size_t currentFirstAvailableAddress,
+                                          size_t alignedEndAddress) {
+  llvm::errs() << "PIM local memory allocation overflow\n";
+  llvm::errs() << "Requested allocation size: " << requestedSize << " bytes\n";
+  llvm::errs() << "Current firstAvailableAddress: " << currentFirstAvailableAddress << "\n";
+  llvm::errs() << "Aligned end address: " << alignedEndAddress << "\n";
+  llvm::errs() << "Address limit: " << kPimAddressLimit << " (signed int32_t immediate range)\n";
+  if (key.lane)
+    llvm::errs() << "Lane: " << *key.lane << "\n";
+  llvm::errs() << "Value: ";
+  value.print(llvm::errs());
+  llvm::errs() << "\n";
+  llvm::errs() << "Value type: " << value.getType() << "\n";
+  if (Operation* definingOp = value.getDefiningOp()) {
+    llvm::errs() << "Defining op:\n";
+    definingOp->print(llvm::errs());
+    llvm::errs() << "\n";
+  }
+}
+
 } // namespace
 
 MemEntry* PimMemory::gatherMemEntry(mlir::Value value, std::optional<unsigned> lane) {
@@ -123,19 +171,29 @@ void PimMemory::allocateGatheredMemory() {
 
 void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind) {
   memEntry.address = firstAvailableAddress;
-  firstAvailableAddress += memEntry.size;
-  // Alignment
-  if (size_t remainder = firstAvailableAddress % minAlignment)
-    firstAvailableAddress += minAlignment - remainder;
+  Operation* anchor = getDiagnosticAnchor(key.value);
+  auto checkedEnd = pim::checkedAdd(memEntry.address, memEntry.size, anchor, "local memory end");
+  FailureOr<size_t> checkedAlignedEnd = failure();
+  if (succeeded(checkedEnd))
+    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
+  bool startFits = memEntry.address <= kPimAddressLimit;
+  bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit;
+  bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit;
+  if (!startFits || !endFits || !alignedEndFits) {
+    printMemoryOverflowDiagnostic(key.value,
+                                  key,
+                                  memEntry.size,
+                                  firstAvailableAddress,
+                                  succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+    llvm_unreachable("PIM local memory allocation overflow");
+  }
+  firstAvailableAddress = *checkedAlignedEnd;
 
   ownedMemEntriesMap[key] = memEntry;
   globalMemEntriesMap[key] = memEntry;
 
   switch (reportKind) {
-  case MemoryReportKind::Alloca:
-    ++reportRow.numAlloca;
-    reportRow.sizeAlloca += memEntry.size;
-    break;
+  case MemoryReportKind::Alloca: break;
   case MemoryReportKind::Global:
     ++reportRow.numGlobal;
     reportRow.sizeGlobal += memEntry.size;
@@ -145,6 +203,34 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
   }
 }
 
+PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key) {
+  PhysicalSlotInfo slot;
+  slot.id = nextPhysicalSlotId++;
+  slot.address = firstAvailableAddress;
+  slot.size = slotSize;
+
+  Operation* anchor = getDiagnosticAnchor(key.value);
+  auto checkedEnd = pim::checkedAdd(slot.address, slot.size, anchor, "local memory end");
+  FailureOr<size_t> checkedAlignedEnd = failure();
+  if (succeeded(checkedEnd))
+    checkedAlignedEnd = checkedAlignTo(*checkedEnd, minAlignment, anchor, "local memory alignment");
+  bool startFits = slot.address <= kPimAddressLimit;
+  bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit;
+  bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit;
+  if (!startFits || !endFits || !alignedEndFits) {
+    printMemoryOverflowDiagnostic(key.value,
+                                  key,
+                                  slot.size,
+                                  firstAvailableAddress,
+                                  succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+    llvm_unreachable("PIM local memory allocation overflow");
+  }
+
+  firstAvailableAddress = *checkedAlignedEnd;
+  localPhysicalSlots.push_back(slot);
+  return slot;
+}
+
 void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
   SmallDenseMap<memref::GlobalOp, mlir::Value, 8> globalConstants;
   SmallVector<std::pair<mlir::Value, mlir::Value>, 16> globalAliases;
@@ -184,9 +270,71 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
 }
 
 void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
-  op->walk([&](memref::AllocOp allocOp) { gatherMemEntry(allocOp, lane); });
+  auto intervals = buildLocalAllocIntervals(op, lane);
+  SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);
 
-  allocateGatheredMemory();
+  SmallVector<size_t> slotOrder(plannedSlots.size());
+  std::iota(slotOrder.begin(), slotOrder.end(), 0);
+  llvm::stable_sort(slotOrder, [&](size_t lhsIndex, size_t rhsIndex) {
+    const PlannedPhysicalSlot& lhs = plannedSlots[lhsIndex];
+    const PlannedPhysicalSlot& rhs = plannedSlots[rhsIndex];
+    if (lhs.requiredSize != rhs.requiredSize)
+      return lhs.requiredSize > rhs.requiredSize;
+    return lhs.id < rhs.id;
+  });
+
+  SmallVector<bool, 16> usedExistingSlots(localPhysicalSlots.size(), false);
+  for (size_t slotIndex : slotOrder) {
+    PlannedPhysicalSlot& slot = plannedSlots[slotIndex];
+    size_t bestExistingIndex = std::numeric_limits<size_t>::max();
+    auto bestKey = std::tuple<size_t, size_t, size_t>(
+      std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max());
+
+    for (size_t existingIndex = 0; existingIndex < localPhysicalSlots.size(); ++existingIndex) {
+      if (usedExistingSlots[existingIndex])
+        continue;
+      const PhysicalSlotInfo& existingSlot = localPhysicalSlots[existingIndex];
+      if (existingSlot.size < slot.requiredSize)
+        continue;
+      auto candidateKey =
+        std::tuple<size_t, size_t, size_t>(existingSlot.size - slot.requiredSize, existingSlot.size, existingSlot.id);
+      if (candidateKey < bestKey) {
+        bestKey = candidateKey;
+        bestExistingIndex = existingIndex;
+      }
+    }
+
+    if (bestExistingIndex != std::numeric_limits<size_t>::max()) {
+      const PhysicalSlotInfo& existingSlot = localPhysicalSlots[bestExistingIndex];
+      slot.id = existingSlot.id;
+      slot.address = existingSlot.address;
+      slot.size = existingSlot.size;
+      usedExistingSlots[bestExistingIndex] = true;
+    }
+    else {
+      PhysicalSlotInfo newSlot = allocatePhysicalSlot(slot.requiredSize, intervals[slot.intervalIndices.front()].key);
+      slot.id = newSlot.id;
+      slot.address = newSlot.address;
+      slot.size = newSlot.size;
+      usedExistingSlots.push_back(true);
+    }
+
+    for (size_t intervalIndex : slot.intervalIndices) {
+      LocalAllocInterval& interval = intervals[intervalIndex];
+      interval.physicalSlotId = slot.id;
+      interval.assignedAddress = slot.address;
+      interval.physicalSlotSize = slot.size;
+      MemEntry memEntry {slot.address, interval.size};
+      ownedMemEntriesMap[interval.key] = memEntry;
+      globalMemEntriesMap[interval.key] = memEntry;
+    }
+  }
+
+  if (pimMemoryReport != PimMemoryReportNone) {
+    MemoryPlanArtifacts artifacts =
+      buildMemoryPlanArtifacts(op, lane, intervals, plannedSlots, kPimAddressLimit, pimMemoryReport);
+    livenessArtifacts.textReport += artifacts.textReport;
+  }
 }
 
 static void printHostMemoryReportRow(raw_ostream& os, const MemoryReportRow& row) {
@@ -226,7 +374,14 @@ static MemoryReportRow addMemoryReportRows(const MemoryReportRow& lhs, const Mem
   return result;
 }
 
-MemoryReportRow PimMemory::getReportRow() const { return reportRow; }
+MemoryReportRow PimMemory::getReportRow() const {
+  MemoryReportRow row = reportRow;
+  row.numAlloca = localPhysicalSlots.size();
+  row.sizeAlloca = 0;
+  for (const PhysicalSlotInfo& slot : localPhysicalSlots)
+    row.sizeAlloca += slot.size;
+  return row;
+}
 
 void PimMemory::remove(mlir::Value val) {
   for (auto it = ownedMemEntriesMap.begin(); it != ownedMemEntriesMap.end();)
@@ -842,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;
@@ -1267,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();
       }
 
@@ -1308,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);
@@ -1324,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);
@@ -1342,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());
@@ -1365,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;
@@ -1380,6 +1594,18 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
   auto weightEmission = createAndPopulateWeightFolder(weightRequests, outputDirPath);
   memory.setTotalWeightBytes(weightEmission.totalWeightBytes);
   auto& mapCoreWeightToFileName = weightEmission.mapCoreWeightToFileName;
+  if (std::string reportsRoot = getOutputDir(); !reportsRoot.empty()) {
+    std::string reportsDir = reportsRoot + "/reports";
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_report.txt");
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_report.json");
+    sys::fs::remove(reportsDir + "/pim_memory_liveness_timeline.dot");
+  }
+  std::fstream livenessReportFile;
+  std::unique_ptr<llvm::raw_os_ostream> livenessReportOs;
+  if (pimMemoryReport != PimMemoryReportNone) {
+    livenessReportFile = openReportFileWithExtension("pim_memory_liveness_report", "txt");
+    livenessReportOs = std::make_unique<llvm::raw_os_ostream>(livenessReportFile);
+  }
 
   for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {
     const CoreEmissionJob& job = jobs[jobIndex];
@@ -1391,6 +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;
     }
   }
@@ -1419,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;
@@ -74,16 +86,20 @@ struct MemoryReportEntry {
 
 class PimMemory {
   llvm::SmallVector<PendingMemEntry, 32> memEntries;
+  llvm::SmallVector<PhysicalSlotInfo, 32> localPhysicalSlots;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap;
   llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> ownedMemEntriesMap;
   MemoryReportRow reportRow;
+  MemoryPlanArtifacts livenessArtifacts;
 
   size_t minAlignment = 4;
   size_t firstAvailableAddress = 0;
+  size_t nextPhysicalSlotId = 0;
 
   MemEntry* gatherMemEntry(mlir::Value value, std::optional<unsigned> lane = std::nullopt);
   void allocateGatheredMemory();
   void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind);
+  PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key);
 
 public:
   PimMemory(llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap)
@@ -92,6 +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; }

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

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

src/PIM/Compiler/PimMemoryLiveness.cpp

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

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

src/PIM/Compiler/PimWeightEmitter.hpp

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

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -22,6 +22,7 @@ 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

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -124,6 +124,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXMatMulOp>();
   target.addIllegalOp<ONNXTransposeOp>();
   target.addIllegalOp<ONNXAddOp>();
+  target.addIllegalOp<ONNXSubOp>();
   target.addIllegalOp<ONNXDivOp>();
   target.addIllegalOp<ONNXMulOp>();
   target.addIllegalOp<ONNXGemmOp>();
@@ -137,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>();
 

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -22,6 +22,7 @@ void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   populateGatherPatterns(patterns, ctx);
   populateResizePatterns(patterns, ctx);
   populateReshapePatterns(patterns, ctx);
+  populateSlicePatterns(patterns, ctx);
   populateSplitPatterns(patterns, ctx);
   populateTransposePatterns(patterns, ctx);
 }

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -29,6 +29,7 @@ void populateConcatPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext
 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);
 

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

@@ -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

@@ -690,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());
@@ -725,9 +720,12 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     return failure();
   }
 
-  const int64_t numOutRows = outType.getDimSize(0);
-  const int64_t numOutCols = outType.getDimSize(1);
-  const int64_t reductionSize = aType.getDimSize(1);
+  if (gemmOpAdaptor.getTransA()) {
+    auto aShape = aType.getShape();
+    auto transposedType = RankedTensorType::get({aShape[1], aShape[0]}, aType.getElementType(), aType.getEncoding());
+    a = ONNXTransposeOp::create(rewriter, loc, transposedType, a, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    aType = transposedType;
+  }
 
   if (gemmOpAdaptor.getTransB()) {
     auto bShape = bType.getShape();
@@ -736,6 +734,10 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     bType = transposedType;
   }
 
+  const int64_t numOutRows = outType.getDimSize(0);
+  const int64_t numOutCols = outType.getDimSize(1);
+  const int64_t reductionSize = aType.getDimSize(1);
+
   if (!isCompileTimeComputable(b)) {
     bool hasC = hasGemmBias(c);
     float alpha = gemmOpAdaptor.getAlpha().convertToFloat();

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

@@ -22,13 +22,87 @@ namespace {
 
 static FailureOr<SmallVector<int64_t>> inferSupportedBatchShape(ArrayRef<int64_t> lhsBatchShape,
                                                                 ArrayRef<int64_t> rhsBatchShape) {
-  if (lhsBatchShape.empty())
-    return SmallVector<int64_t>(rhsBatchShape.begin(), rhsBatchShape.end());
-  if (rhsBatchShape.empty())
-    return SmallVector<int64_t>(lhsBatchShape.begin(), lhsBatchShape.end());
-  if (!llvm::equal(lhsBatchShape, rhsBatchShape))
-    return failure();
-  return SmallVector<int64_t>(lhsBatchShape.begin(), lhsBatchShape.end());
+  const int64_t resultRank = std::max<int64_t>(lhsBatchShape.size(), rhsBatchShape.size());
+  SmallVector<int64_t> resultShape(resultRank, 1);
+  for (int64_t resultIndex = resultRank - 1, lhsIndex = lhsBatchShape.size() - 1, rhsIndex = rhsBatchShape.size() - 1;
+       resultIndex >= 0;
+       --resultIndex, --lhsIndex, --rhsIndex) {
+    const int64_t lhsDim = lhsIndex >= 0 ? lhsBatchShape[lhsIndex] : 1;
+    const int64_t rhsDim = rhsIndex >= 0 ? rhsBatchShape[rhsIndex] : 1;
+    if (lhsDim != rhsDim && lhsDim != 1 && rhsDim != 1)
+      return failure();
+    resultShape[resultIndex] = std::max(lhsDim, rhsDim);
+  }
+  return resultShape;
+}
+
+static int64_t mapStaticBroadcastedBatchIndex(int64_t outputBatchIndex,
+                                              ArrayRef<int64_t> sourceBatchShape,
+                                              ArrayRef<int64_t> outputBatchShape) {
+  if (sourceBatchShape.empty() || getStaticShapeElementCount(sourceBatchShape) == 1)
+    return 0;
+  if (llvm::equal(sourceBatchShape, outputBatchShape))
+    return outputBatchIndex;
+
+  SmallVector<int64_t> outputStrides = computeRowMajorStrides(outputBatchShape);
+  SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceBatchShape);
+  int64_t sourceFlatIndex = 0;
+  for (int64_t sourceDimIndex = 0; sourceDimIndex < static_cast<int64_t>(sourceBatchShape.size()); ++sourceDimIndex) {
+    if (sourceBatchShape[sourceDimIndex] == 1)
+      continue;
+    const int64_t outputDimIndex = outputBatchShape.size() - sourceBatchShape.size() + sourceDimIndex;
+    const int64_t outputDimStride = outputStrides.empty() ? 1 : outputStrides[outputDimIndex];
+    const int64_t outputDimIndexValue = outputDimStride == 1
+                                        ? outputBatchIndex % outputBatchShape[outputDimIndex]
+                                        : (outputBatchIndex / outputDimStride) % outputBatchShape[outputDimIndex];
+    sourceFlatIndex += outputDimIndexValue * sourceStrides[sourceDimIndex];
+  }
+  return sourceFlatIndex;
+}
+
+static Value computeFlatBatchIndexCoordinate(
+  Value flatBatchIndex, ArrayRef<int64_t> batchShape, int64_t dimIndex, PatternRewriter& rewriter, Location loc) {
+  if (batchShape[dimIndex] == 1)
+    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+
+  const int64_t dimStride = dimIndex + 1 == static_cast<int64_t>(batchShape.size())
+                            ? 1
+                            : getStaticShapeElementCount(batchShape.drop_front(dimIndex + 1));
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value dimCoordinate = flatBatchIndex;
+  if (dimStride != 1)
+    dimCoordinate = affineFloorDivConst(rewriter, loc, dimCoordinate, dimStride, anchorOp);
+  return affineModConst(rewriter, loc, dimCoordinate, batchShape[dimIndex], anchorOp);
+}
+
+static Value mapOutputBatchIndexToSourceBatchIndex(Value outputBatchIndex,
+                                                   ArrayRef<int64_t> sourceBatchShape,
+                                                   ArrayRef<int64_t> outputBatchShape,
+                                                   PatternRewriter& rewriter,
+                                                   Location loc) {
+  if (sourceBatchShape.empty() || getStaticShapeElementCount(sourceBatchShape) == 1)
+    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+  if (llvm::equal(sourceBatchShape, outputBatchShape))
+    return outputBatchIndex;
+
+  Value sourceBatchIndex = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
+  SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceBatchShape);
+  for (int64_t sourceDimIndex = 0; sourceDimIndex < static_cast<int64_t>(sourceBatchShape.size()); ++sourceDimIndex) {
+    if (sourceBatchShape[sourceDimIndex] == 1)
+      continue;
+    const int64_t outputDimIndex = outputBatchShape.size() - sourceBatchShape.size() + sourceDimIndex;
+    Value outputCoordinate =
+      computeFlatBatchIndexCoordinate(outputBatchIndex, outputBatchShape, outputDimIndex, rewriter, loc);
+    Value contribution = sourceStrides[sourceDimIndex] == 1
+                         ? outputCoordinate
+                         : affineMulConst(rewriter,
+                                          loc,
+                                          outputCoordinate,
+                                          sourceStrides[sourceDimIndex],
+                                          rewriter.getInsertionBlock()->getParentOp());
+    sourceBatchIndex = arith::AddIOp::create(rewriter, loc, sourceBatchIndex, contribution);
+  }
+  return sourceBatchIndex;
 }
 
 static Value
@@ -67,6 +141,52 @@ expandBatchDims(Value value, RankedTensorType outputType, size_t batchRank, Patt
   return materializeOrComputeUnary(value, outputType, rewriter, loc, buildExpanded);
 }
 
+static Value createMatrixFromVector(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
+  auto buildExpanded = [&](Value input) -> Value {
+    return tensor::ExpandShapeOp::create(rewriter,
+                                         loc,
+                                         resultType,
+                                         input,
+                                         SmallVector<ReassociationIndices> {
+                                           {0, 1}
+    });
+  };
+  return materializeOrComputeUnary(value, resultType, rewriter, loc, buildExpanded);
+}
+
+static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> removedAxes) {
+  SmallVector<ReassociationIndices> reassociation;
+  ReassociationIndices currentGroup;
+  for (auto [axis, removeAxis] : llvm::enumerate(removedAxes)) {
+    currentGroup.push_back(axis);
+    if (!removeAxis) {
+      reassociation.push_back(currentGroup);
+      currentGroup.clear();
+    }
+  }
+
+  if (!currentGroup.empty()) {
+    if (reassociation.empty())
+      reassociation.push_back(std::move(currentGroup));
+    else
+      reassociation.back().append(currentGroup.begin(), currentGroup.end());
+  }
+  return reassociation;
+}
+
+static Value squeezeUnitDims(
+  Value value, RankedTensorType resultType, ArrayRef<bool> removedAxes, PatternRewriter& rewriter, Location loc) {
+  if (cast<RankedTensorType>(value.getType()) == resultType)
+    return value;
+
+  SmallVector<ReassociationIndices> reassociation =
+    resultType.getRank() == 0 ? SmallVector<ReassociationIndices> {} : buildCollapseReassociation(removedAxes);
+  auto buildCollapsed = [&](Value input) -> Value {
+    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation).getResult();
+  };
+  return materializeOrComputeUnary(value, resultType, rewriter, loc, buildCollapsed);
+}
+
 static Value ensureBatchedTensor(
   Value value, int64_t batchSize, int64_t rows, int64_t cols, PatternRewriter& rewriter, Location loc) {
   auto type = cast<RankedTensorType>(value.getType());
@@ -171,8 +291,11 @@ static Value createPaddedBatchedInputCompute(Value input,
   return computeOp.getResult(0);
 }
 
-static FailureOr<Value> materializePaddedBatchedWeight(
-  Value value, int64_t sourceBatch, int64_t targetBatch, RankedTensorType resultType, PatternRewriter& rewriter) {
+static FailureOr<Value> materializePaddedBatchedWeight(Value value,
+                                                       ArrayRef<int64_t> sourceBatchShape,
+                                                       ArrayRef<int64_t> targetBatchShape,
+                                                       RankedTensorType resultType,
+                                                       PatternRewriter& rewriter) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   if (sourceType == resultType)
     return value;
@@ -183,13 +306,15 @@ static FailureOr<Value> materializePaddedBatchedWeight(
 
   const int64_t sourceRows = sourceType.getRank() == 2 ? sourceType.getDimSize(0) : sourceType.getDimSize(1);
   const int64_t sourceCols = sourceType.getRank() == 2 ? sourceType.getDimSize(1) : sourceType.getDimSize(2);
+  const int64_t targetBatch = targetBatchShape.empty() ? 1 : getStaticShapeElementCount(targetBatchShape);
   const int64_t targetRows = resultType.getDimSize(1);
   const int64_t targetCols = resultType.getDimSize(2);
   SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> resultValues(resultType.getNumElements(), rewriter.getZeroAttr(resultType.getElementType()));
 
   for (int64_t batchIdx = 0; batchIdx < targetBatch; ++batchIdx) {
-    const int64_t sourceBatchIdx = sourceType.getRank() == 2 ? 0 : (sourceBatch == 1 ? 0 : batchIdx);
+    const int64_t sourceBatchIdx =
+      sourceType.getRank() == 2 ? 0 : mapStaticBroadcastedBatchIndex(batchIdx, sourceBatchShape, targetBatchShape);
     const int64_t sourceBatchBase = sourceType.getRank() == 2 ? 0 : sourceBatchIdx * sourceRows * sourceCols;
     const int64_t targetBatchBase = batchIdx * targetRows * targetCols;
     for (int64_t row = 0; row < sourceRows; ++row)
@@ -202,16 +327,18 @@ static FailureOr<Value> materializePaddedBatchedWeight(
 }
 
 static Value extractBatchedATile(Value a,
-                                 int64_t sourceBatchCount,
-                                 Value batch,
+                                 ArrayRef<int64_t> sourceBatchShape,
+                                 ArrayRef<int64_t> outputBatchShape,
+                                 Value outputBatchIndex,
                                  Value row,
                                  Value kOffset,
                                  RankedTensorType aTileType,
                                  PatternRewriter& rewriter,
                                  Location loc) {
   auto aSliceType = RankedTensorType::get({1, 1, aTileType.getDimSize(1)}, aTileType.getElementType());
-  SmallVector<OpFoldResult> offsets {
-    sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0)) : OpFoldResult(batch), row, kOffset};
+  Value sourceBatchIndex =
+    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
+  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, kOffset};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(aTileType.getDimSize(1))};
   auto slice =
@@ -227,8 +354,9 @@ static Value extractBatchedATile(Value a,
 }
 
 static Value extractBatchedBTile(Value b,
-                                 int64_t sourceBatchCount,
-                                 Value batch,
+                                 ArrayRef<int64_t> sourceBatchShape,
+                                 ArrayRef<int64_t> outputBatchShape,
+                                 Value outputBatchIndex,
                                  Value kOffset,
                                  Value hOffset,
                                  RankedTensorType bTileType,
@@ -236,8 +364,9 @@ static Value extractBatchedBTile(Value b,
                                  Location loc) {
   auto bSliceType =
     RankedTensorType::get({1, bTileType.getDimSize(0), bTileType.getDimSize(1)}, bTileType.getElementType());
-  SmallVector<OpFoldResult> offsets {
-    sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0)) : OpFoldResult(batch), kOffset, hOffset};
+  Value sourceBatchIndex =
+    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
+  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), kOffset, hOffset};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(bTileType.getDimSize(0)),
                                    rewriter.getIndexAttr(bTileType.getDimSize(1))};
@@ -262,9 +391,10 @@ static Value getBatchLaneIndex(
 static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
                                                                   Value b,
                                                                   RankedTensorType aType,
-                                                                  int64_t aBatchCount,
+                                                                  ArrayRef<int64_t> aBatchShape,
                                                                   RankedTensorType bType,
-                                                                  int64_t bBatchCount,
+                                                                  ArrayRef<int64_t> bBatchShape,
+                                                                  ArrayRef<int64_t> outputBatchShape,
                                                                   RankedTensorType partialPiecesType,
                                                                   int64_t numOutRows,
                                                                   int64_t numKSlices,
@@ -298,10 +428,10 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
       auto pieceType =
         RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());
 
-      Value aTile =
-        extractBatchedATile(args.inputs.front(), aBatchCount, batch, row, kOffset, aTileType, rewriter, loc);
-      Value bTile =
-        extractBatchedBTile(args.weights.front(), bBatchCount, batch, kOffset, hOffset, bTileType, rewriter, loc);
+      Value aTile = extractBatchedATile(
+        args.inputs.front(), aBatchShape, outputBatchShape, batch, row, kOffset, aTileType, rewriter, loc);
+      Value bTile = extractBatchedBTile(
+        args.weights.front(), bBatchShape, outputBatchShape, batch, kOffset, hOffset, bTileType, rewriter, loc);
       Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
       SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};
@@ -315,17 +445,17 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
 }
 
 static Value extractDynamicBatchedBColumn(Value matrix,
-                                          int64_t sourceBatchCount,
-                                          Value batch,
+                                          ArrayRef<int64_t> sourceBatchShape,
+                                          ArrayRef<int64_t> outputBatchShape,
+                                          Value outputBatchIndex,
                                           Value column,
                                           RankedTensorType vectorType,
                                           PatternRewriter& rewriter,
                                           Location loc) {
   auto columnSliceType = RankedTensorType::get({1, vectorType.getDimSize(1), 1}, vectorType.getElementType());
-  SmallVector<OpFoldResult> offsets {sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0))
-                                                           : OpFoldResult(batch),
-                                     rewriter.getIndexAttr(0),
-                                     column};
+  Value sourceBatchIndex =
+    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
+  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), rewriter.getIndexAttr(0), column};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -350,17 +480,17 @@ static Value extractDynamicBatchedBColumn(Value matrix,
 }
 
 static Value extractDynamicBatchedRowVector(Value matrix,
-                                            int64_t sourceBatchCount,
-                                            Value batch,
+                                            ArrayRef<int64_t> sourceBatchShape,
+                                            ArrayRef<int64_t> outputBatchShape,
+                                            Value outputBatchIndex,
                                             Value row,
                                             RankedTensorType vectorType,
                                             PatternRewriter& rewriter,
                                             Location loc) {
   auto rowSliceType = RankedTensorType::get({1, 1, vectorType.getDimSize(1)}, vectorType.getElementType());
-  SmallVector<OpFoldResult> offsets {sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0))
-                                                           : OpFoldResult(batch),
-                                     row,
-                                     rewriter.getIndexAttr(0)};
+  Value sourceBatchIndex =
+    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
+  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
   auto rowSlice =
@@ -376,9 +506,10 @@ static Value extractDynamicBatchedRowVector(Value matrix,
 }
 
 static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
-                                                                     int64_t aBatchCount,
+                                                                     ArrayRef<int64_t> aBatchShape,
                                                                      Value b,
-                                                                     int64_t bBatchCount,
+                                                                     ArrayRef<int64_t> bBatchShape,
+                                                                     ArrayRef<int64_t> outputBatchShape,
                                                                      RankedTensorType aType,
                                                                      RankedTensorType bType,
                                                                      RankedTensorType scalarPiecesType,
@@ -406,10 +537,10 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
 
       auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
       auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-      Value aVector =
-        extractDynamicBatchedRowVector(args.inputs[0], aBatchCount, batch, row, vectorType, rewriter, loc);
-      Value bVector =
-        extractDynamicBatchedBColumn(args.inputs[1], bBatchCount, batch, column, vectorType, rewriter, loc);
+      Value aVector = extractDynamicBatchedRowVector(
+        args.inputs[0], aBatchShape, outputBatchShape, batch, row, vectorType, rewriter, loc);
+      Value bVector = extractDynamicBatchedBColumn(
+        args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc);
       Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
       SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
       SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -629,11 +760,17 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
   return computeOp->getResult(0);
 }
 
-struct MatMulShapeInfo {
+struct NormalizedMatMulInfo {
   RankedTensorType lhsType;
   RankedTensorType rhsType;
   RankedTensorType outType;
-  SmallVector<int64_t> batchShape;
+  RankedTensorType normalizedLhsType;
+  RankedTensorType normalizedRhsType;
+  SmallVector<int64_t> lhsBatchShape;
+  SmallVector<int64_t> rhsBatchShape;
+  SmallVector<int64_t> outputBatchShape;
+  bool lhsWasVector;
+  bool rhsWasVector;
   int64_t lhsBatch;
   int64_t rhsBatch;
   int64_t batch;
@@ -642,46 +779,170 @@ struct MatMulShapeInfo {
   int64_t n;
 };
 
-static FailureOr<MatMulShapeInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp) {
+struct MatMulLoweringPlan {
+  Value lhs;
+  Value rhs;
+  RankedTensorType lhsType;
+  RankedTensorType rhsType;
+  SmallVector<int64_t> lhsBatchShape;
+  SmallVector<int64_t> rhsBatchShape;
+  SmallVector<int64_t> outputBatchShape;
+  int64_t lhsBatch;
+  int64_t rhsBatch;
+  int64_t batch;
+  int64_t m;
+  int64_t k;
+  int64_t n;
+  bool transposedResult;
+};
+
+static SmallVector<int64_t> computeExpectedMatMulOutputShape(
+  ArrayRef<int64_t> batchShape, int64_t m, int64_t n, bool lhsWasVector, bool rhsWasVector) {
+  SmallVector<int64_t> shape(batchShape.begin(), batchShape.end());
+  if (lhsWasVector && rhsWasVector)
+    return shape;
+  if (lhsWasVector) {
+    shape.push_back(n);
+    return shape;
+  }
+  if (rhsWasVector) {
+    shape.push_back(m);
+    return shape;
+  }
+  shape.push_back(m);
+  shape.push_back(n);
+  return shape;
+}
+
+static FailureOr<NormalizedMatMulInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp) {
   auto lhsType = dyn_cast<RankedTensorType>(matmulOp.getA().getType());
   auto rhsType = dyn_cast<RankedTensorType>(matmulOp.getB().getType());
   auto outType = dyn_cast<RankedTensorType>(matmulOp.getY().getType());
   if (!lhsType || !rhsType || !outType || !lhsType.hasStaticShape() || !rhsType.hasStaticShape()
       || !outType.hasStaticShape())
     return failure();
-  if (lhsType.getRank() < 2 || rhsType.getRank() < 2 || outType.getRank() < 2)
+  if (lhsType.getRank() < 1 || rhsType.getRank() < 1)
     return failure();
   if (!hasStaticPositiveShape(lhsType) || !hasStaticPositiveShape(rhsType) || !hasStaticPositiveShape(outType))
     return failure();
 
-  SmallVector<int64_t> lhsBatchShape(lhsType.getShape().begin(), lhsType.getShape().end() - 2);
-  SmallVector<int64_t> rhsBatchShape(rhsType.getShape().begin(), rhsType.getShape().end() - 2);
-  auto batchShape = inferSupportedBatchShape(lhsBatchShape, rhsBatchShape);
-  if (failed(batchShape))
+  const bool lhsWasVector = lhsType.getRank() == 1;
+  const bool rhsWasVector = rhsType.getRank() == 1;
+  auto normalizedLhsType =
+    lhsWasVector ? RankedTensorType::get({1, lhsType.getDimSize(0)}, lhsType.getElementType(), lhsType.getEncoding())
+                 : lhsType;
+  auto normalizedRhsType =
+    rhsWasVector ? RankedTensorType::get({rhsType.getDimSize(0), 1}, rhsType.getElementType(), rhsType.getEncoding())
+                 : rhsType;
+
+  SmallVector<int64_t> lhsBatchShape(normalizedLhsType.getShape().begin(), normalizedLhsType.getShape().end() - 2);
+  SmallVector<int64_t> rhsBatchShape(normalizedRhsType.getShape().begin(), normalizedRhsType.getShape().end() - 2);
+  auto outputBatchShape = inferSupportedBatchShape(lhsBatchShape, rhsBatchShape);
+  if (failed(outputBatchShape))
     return failure();
 
   const int64_t lhsBatch = lhsBatchShape.empty() ? 1 : getStaticShapeElementCount(lhsBatchShape);
   const int64_t rhsBatch = rhsBatchShape.empty() ? 1 : getStaticShapeElementCount(rhsBatchShape);
-  const int64_t batch = batchShape->empty() ? 1 : getStaticShapeElementCount(*batchShape);
-  const int64_t m = lhsType.getDimSize(lhsType.getRank() - 2);
-  const int64_t k = lhsType.getDimSize(lhsType.getRank() - 1);
-  const int64_t rhsK = rhsType.getDimSize(rhsType.getRank() - 2);
-  const int64_t n = rhsType.getDimSize(rhsType.getRank() - 1);
+  const int64_t batch = outputBatchShape->empty() ? 1 : getStaticShapeElementCount(*outputBatchShape);
+  const int64_t m = normalizedLhsType.getDimSize(normalizedLhsType.getRank() - 2);
+  const int64_t k = normalizedLhsType.getDimSize(normalizedLhsType.getRank() - 1);
+  const int64_t rhsK = normalizedRhsType.getDimSize(normalizedRhsType.getRank() - 2);
+  const int64_t n = normalizedRhsType.getDimSize(normalizedRhsType.getRank() - 1);
   if (k != rhsK)
     return failure();
 
-  if (outType.getRank() == 2) {
-    if (batch != 1 || outType.getDimSize(0) != m || outType.getDimSize(1) != n)
-      return failure();
-  }
-  else {
-    SmallVector<int64_t> outBatchShape(outType.getShape().begin(), outType.getShape().end() - 2);
-    if (!llvm::equal(outBatchShape, *batchShape) || outType.getDimSize(outType.getRank() - 2) != m
-        || outType.getDimSize(outType.getRank() - 1) != n)
-      return failure();
+  if (SmallVector<int64_t>(outType.getShape().begin(), outType.getShape().end())
+      != computeExpectedMatMulOutputShape(*outputBatchShape, m, n, lhsWasVector, rhsWasVector)) {
+    return failure();
   }
 
-  return MatMulShapeInfo {lhsType, rhsType, outType, *batchShape, lhsBatch, rhsBatch, batch, m, k, n};
+  return NormalizedMatMulInfo {lhsType,
+                               rhsType,
+                               outType,
+                               normalizedLhsType,
+                               normalizedRhsType,
+                               lhsBatchShape,
+                               rhsBatchShape,
+                               *outputBatchShape,
+                               lhsWasVector,
+                               rhsWasVector,
+                               lhsBatch,
+                               rhsBatch,
+                               batch,
+                               m,
+                               k,
+                               n};
+}
+
+static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs,
+                                            Value normalizedRhs,
+                                            const NormalizedMatMulInfo& info,
+                                            bool useTransposedForm,
+                                            PatternRewriter& rewriter,
+                                            Location loc) {
+  MatMulLoweringPlan plan {normalizedLhs,
+                           normalizedRhs,
+                           cast<RankedTensorType>(normalizedLhs.getType()),
+                           cast<RankedTensorType>(normalizedRhs.getType()),
+                           info.lhsBatchShape,
+                           info.rhsBatchShape,
+                           info.outputBatchShape,
+                           info.lhsBatch,
+                           info.rhsBatch,
+                           info.batch,
+                           info.m,
+                           info.k,
+                           info.n,
+                           false};
+  if (!useTransposedForm)
+    return plan;
+
+  plan.lhs = transposeLastTwoDims(normalizedRhs, rewriter, loc);
+  plan.rhs = transposeLastTwoDims(normalizedLhs, rewriter, loc);
+  plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
+  plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
+  std::swap(plan.lhsBatchShape, plan.rhsBatchShape);
+  std::swap(plan.lhsBatch, plan.rhsBatch);
+  plan.m = info.n;
+  plan.n = info.m;
+  plan.transposedResult = true;
+  return plan;
+}
+
+static Value normalizeMatMulOperand(
+  Value value, RankedTensorType normalizedType, bool wasVector, PatternRewriter& rewriter, Location loc) {
+  if (!wasVector)
+    return value;
+  return createMatrixFromVector(value, normalizedType, rewriter, loc);
+}
+
+static Value finalizeNormalizedMatMulResult(Value value,
+                                            RankedTensorType directOutType,
+                                            const NormalizedMatMulInfo& info,
+                                            PatternRewriter& rewriter,
+                                            Location loc) {
+  // The direct lowered result is always [flatBatch, normalizedM, normalizedN].
+  // Restore ONNX MatMul result rank by expanding right-aligned batch dimensions
+  // and removing the synthetic unit matrix axes introduced for vector operands.
+  Value result = value;
+  RankedTensorType currentType = directOutType;
+  if (info.outputBatchShape.size() > 1) {
+    SmallVector<int64_t> expandedShape(info.outputBatchShape.begin(), info.outputBatchShape.end());
+    expandedShape.push_back(info.m);
+    expandedShape.push_back(info.n);
+    auto expandedType = RankedTensorType::get(expandedShape, info.outType.getElementType(), info.outType.getEncoding());
+    result = expandBatchDims(result, expandedType, info.outputBatchShape.size(), rewriter, loc);
+    currentType = expandedType;
+  }
+
+  SmallVector<bool> removedAxes(currentType.getRank(), false);
+  if (info.outputBatchShape.empty())
+    removedAxes[0] = true;
+  if (info.lhsWasVector)
+    removedAxes[currentType.getRank() - 2] = true;
+  if (info.rhsWasVector)
+    removedAxes[currentType.getRank() - 1] = true;
+  return squeezeUnitDims(result, info.outType, removedAxes, rewriter, loc);
 }
 
 struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
@@ -689,7 +950,7 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
-    if (failed(shapeInfo) || shapeInfo->outType.getRank() != 2)
+    if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector || !shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
@@ -742,61 +1003,56 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
     if (failed(shapeInfo))
       return failure();
-    if (shapeInfo->outType.getRank() == 2)
+    if (!shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector && shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
-    bool useTransposedForm = isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
+    bool useTransposedForm = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
+                          && isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
 
-    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
-    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
-    int64_t lhsBatchForGemm = shapeInfo->lhsBatch;
-    int64_t rhsBatchForGemm = shapeInfo->rhsBatch;
-    int64_t gemmM = shapeInfo->m;
-    int64_t gemmK = shapeInfo->k;
-    int64_t gemmN = shapeInfo->n;
-    if (useTransposedForm) {
-      lhs = transposeLastTwoDims(matmulOp.getB(), rewriter, loc);
-      lhsBatchForGemm = shapeInfo->rhsBatch;
-      rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc);
-      rhsBatchForGemm = shapeInfo->lhsBatch;
-      gemmM = shapeInfo->n;
-      gemmN = shapeInfo->m;
-    }
+    Value lhs =
+      normalizeMatMulOperand(matmulOp.getA(), shapeInfo->normalizedLhsType, shapeInfo->lhsWasVector, rewriter, loc);
+    Value rhs =
+      normalizeMatMulOperand(matmulOp.getB(), shapeInfo->normalizedRhsType, shapeInfo->rhsWasVector, rewriter, loc);
+    lhs = collapseBatchDims(lhs, shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
+    rhs = collapseBatchDims(rhs, shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
+    MatMulLoweringPlan plan = buildLoweringPlan(lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc);
 
-    lhs = ensureBatchedTensor(lhs, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
-    rhs = ensureBatchedTensor(rhs, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
-    auto lhsBatchedType = cast<RankedTensorType>(lhs.getType());
-    auto rhsBatchedType = cast<RankedTensorType>(rhs.getType());
-    auto directOutType = RankedTensorType::get({shapeInfo->batch, gemmM, gemmN}, shapeInfo->outType.getElementType());
+    plan.lhs = ensureBatchedTensor(plan.lhs, plan.lhsBatch, plan.m, plan.k, rewriter, loc);
+    plan.rhs = ensureBatchedTensor(plan.rhs, plan.rhsBatch, plan.k, plan.n, rewriter, loc);
+    plan.lhsType = cast<RankedTensorType>(plan.lhs.getType());
+    plan.rhsType = cast<RankedTensorType>(plan.rhs.getType());
+    auto directOutType = RankedTensorType::get(
+      {plan.batch, plan.m, plan.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
 
-    if (isCompileTimeComputable(rhs)) {
-      const int64_t numKSlices = ceilIntegerDivide(gemmK, crossbarSize.getValue());
-      const int64_t numOutHSlices = ceilIntegerDivide(gemmN, crossbarSize.getValue());
+    if (isCompileTimeComputable(plan.rhs)) {
+      const int64_t numKSlices = ceilIntegerDivide(plan.k, crossbarSize.getValue());
+      const int64_t numOutHSlices = ceilIntegerDivide(plan.n, crossbarSize.getValue());
       const int64_t paddedReductionSize = numKSlices * static_cast<int64_t>(crossbarSize.getValue());
       const int64_t paddedOutCols = numOutHSlices * static_cast<int64_t>(crossbarSize.getValue());
       auto paddedLhsType = RankedTensorType::get(
-        {lhsBatchForGemm, gemmM, paddedReductionSize}, lhsBatchedType.getElementType(), lhsBatchedType.getEncoding());
-      auto paddedRhsType = RankedTensorType::get({shapeInfo->batch, paddedReductionSize, paddedOutCols},
-                                                 rhsBatchedType.getElementType(),
-                                                 rhsBatchedType.getEncoding());
+        {plan.lhsBatch, plan.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
+      auto paddedRhsType = RankedTensorType::get(
+        {plan.batch, paddedReductionSize, paddedOutCols}, plan.rhsType.getElementType(), plan.rhsType.getEncoding());
       auto paddedOutType =
-        RankedTensorType::get({shapeInfo->batch, gemmM, paddedOutCols}, shapeInfo->outType.getElementType());
+        RankedTensorType::get({plan.batch, plan.m, paddedOutCols}, shapeInfo->outType.getElementType());
 
-      auto paddedRhs = materializePaddedBatchedWeight(rhs, rhsBatchForGemm, shapeInfo->batch, paddedRhsType, rewriter);
+      auto paddedRhs =
+        materializePaddedBatchedWeight(plan.rhs, plan.rhsBatchShape, plan.outputBatchShape, paddedRhsType, rewriter);
       if (succeeded(paddedRhs)) {
-        Value paddedLhs = createPaddedBatchedInputCompute(lhs, paddedLhsType, rewriter, loc);
-        const int64_t laneCount = shapeInfo->batch * gemmM * numKSlices * numOutHSlices;
+        Value paddedLhs = createPaddedBatchedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
+        const int64_t laneCount = plan.batch * plan.m * numKSlices * numOutHSlices;
         auto partialPiecesType = RankedTensorType::get({laneCount, static_cast<int64_t>(crossbarSize.getValue())},
                                                        shapeInfo->outType.getElementType());
         auto batchOp = createBatchedVmmBatch(paddedLhs,
                                              *paddedRhs,
                                              paddedLhsType,
-                                             lhsBatchForGemm,
+                                             plan.lhsBatchShape,
                                              paddedRhsType,
-                                             rhsBatchForGemm,
+                                             plan.rhsBatchShape,
+                                             plan.outputBatchShape,
                                              partialPiecesType,
-                                             gemmM,
+                                             plan.m,
                                              numKSlices,
                                              numOutHSlices,
                                              rewriter,
@@ -807,34 +1063,35 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                                     partialPiecesType,
                                                     directOutType,
                                                     paddedOutType,
-                                                    shapeInfo->batch,
+                                                    plan.batch,
                                                     numKSlices,
                                                     rewriter,
                                                     loc);
         if (failed(result))
           return failure();
         Value finalResult = *result;
-        if (useTransposedForm) {
-          auto transposedOutType = RankedTensorType::get({shapeInfo->batch, shapeInfo->m, shapeInfo->n},
+        if (plan.transposedResult) {
+          auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n},
                                                          shapeInfo->outType.getElementType(),
                                                          shapeInfo->outType.getEncoding());
           finalResult =
             ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
               .getResult();
         }
-        finalResult = expandBatchDims(finalResult, shapeInfo->outType, shapeInfo->batchShape.size(), rewriter, loc);
+        finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
         rewriter.replaceOp(matmulOp, finalResult);
         return success();
       }
     }
-    const int64_t laneCount = shapeInfo->batch * gemmM * gemmN;
+    const int64_t laneCount = plan.batch * plan.m * plan.n;
     auto scalarPiecesType = RankedTensorType::get({laneCount, 1}, shapeInfo->outType.getElementType());
-    auto batchOp = createBatchedVvdmulBatch(lhs,
-                                            lhsBatchForGemm,
-                                            rhs,
-                                            rhsBatchForGemm,
-                                            lhsBatchedType,
-                                            rhsBatchedType,
+    auto batchOp = createBatchedVvdmulBatch(plan.lhs,
+                                            plan.lhsBatchShape,
+                                            plan.rhs,
+                                            plan.rhsBatchShape,
+                                            plan.outputBatchShape,
+                                            plan.lhsType,
+                                            plan.rhsType,
                                             scalarPiecesType,
                                             directOutType,
                                             rewriter,
@@ -846,15 +1103,15 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
     if (failed(result))
       return failure();
     Value finalResult = *result;
-    if (useTransposedForm) {
-      auto transposedOutType = RankedTensorType::get({shapeInfo->batch, shapeInfo->m, shapeInfo->n},
+    if (plan.transposedResult) {
+      auto transposedOutType = RankedTensorType::get({plan.batch, shapeInfo->m, shapeInfo->n},
                                                      shapeInfo->outType.getElementType(),
                                                      shapeInfo->outType.getEncoding());
       finalResult =
         ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
           .getResult();
     }
-    finalResult = expandBatchDims(finalResult, shapeInfo->outType, shapeInfo->batchShape.size(), rewriter, loc);
+    finalResult = finalizeNormalizedMatMulResult(finalResult, directOutType, *shapeInfo, rewriter, loc);
     rewriter.replaceOp(matmulOp, finalResult);
     return success();
   }

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

@@ -6,6 +6,8 @@
 
 #include <algorithm>
 #include <numeric>
+#include <optional>
+#include <type_traits>
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -19,6 +21,85 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
+struct ReduceMeanSemantics {
+  SmallVector<int64_t> axes;
+  int64_t keepdims = 1;
+  bool isIdentity = false;
+};
+
+static bool isNoneValueLike(Value value) { return isa_and_nonnull<ONNXNoneOp>(value.getDefiningOp()); }
+
+static FailureOr<SmallVector<int64_t>> getConstantIntValues(Value value) {
+  auto denseAttr = dyn_cast_or_null<DenseIntElementsAttr>(getHostConstDenseElementsAttr(value));
+  if (!denseAttr)
+    return failure();
+  return SmallVector<int64_t>(denseAttr.getValues<int64_t>().begin(), denseAttr.getValues<int64_t>().end());
+}
+
+static FailureOr<SmallVector<int64_t>> normalizeAxesChecked(ArrayRef<int64_t> axes, int64_t rank) {
+  SmallVector<int64_t> normalizedAxes;
+  normalizedAxes.reserve(axes.size());
+  for (int64_t axis : axes) {
+    auto normalizedAxis = normalizeAxisChecked(axis, rank);
+    if (failed(normalizedAxis))
+      return failure();
+    normalizedAxes.push_back(*normalizedAxis);
+  }
+  llvm::sort(normalizedAxes);
+  normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end());
+  return normalizedAxes;
+}
+
+template <typename ReduceMeanOp, typename ReduceMeanOpAdaptor>
+static FailureOr<ReduceMeanSemantics>
+getReduceMeanSemantics(ReduceMeanOp reduceMeanOp, ReduceMeanOpAdaptor adaptor, int64_t inputRank) {
+  ReduceMeanSemantics semantics;
+  semantics.keepdims = reduceMeanOp.getKeepdims();
+
+  if constexpr (std::is_same_v<ReduceMeanOp, ONNXReduceMeanV13Op>) {
+    auto axes = onnx_mlir::normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputRank);
+    if (failed(axes))
+      return failure();
+    semantics.axes = std::move(*axes);
+    return semantics;
+  }
+  else {
+    if (isNoneValueLike(adaptor.getAxes())) {
+      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
+        semantics.isIdentity = true;
+        return semantics;
+      }
+
+      semantics.axes.reserve(inputRank);
+      for (int64_t axis = 0; axis < inputRank; ++axis)
+        semantics.axes.push_back(axis);
+      return semantics;
+    }
+
+    auto axes = getConstantIntValues(adaptor.getAxes());
+    if (failed(axes))
+      return failure();
+
+    if (axes->empty()) {
+      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
+        semantics.isIdentity = true;
+        return semantics;
+      }
+
+      semantics.axes.reserve(inputRank);
+      for (int64_t axis = 0; axis < inputRank; ++axis)
+        semantics.axes.push_back(axis);
+      return semantics;
+    }
+
+    auto normalizedAxes = normalizeAxesChecked(*axes, inputRank);
+    if (failed(normalizedAxes))
+      return failure();
+    semantics.axes = std::move(*normalizedAxes);
+    return semantics;
+  }
+}
+
 static SmallVector<bool> buildReducedAxesMask(ArrayRef<int64_t> axes, int64_t rank) {
   SmallVector<bool> reducedAxes(rank, false);
   for (int64_t axis : axes) {
@@ -238,14 +319,8 @@ static Value squeezeReducedAxes(Value keepdimsValue,
                                 ArrayRef<bool> reducedAxes,
                                 ConversionPatternRewriter& rewriter,
                                 Location loc) {
-  if (resultType.getRank() == 0) {
-    SmallVector<Value> indices(cast<RankedTensorType>(keepdimsValue.getType()).getRank(),
-                               getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0));
-    Value element = tensor::ExtractOp::create(rewriter, loc, keepdimsValue, indices);
-    return tensor::FromElementsOp::create(rewriter, loc, resultType, ValueRange {element});
-  }
-
-  auto reassociation = buildCollapseReassociation(reducedAxes);
+  SmallVector<ReassociationIndices> reassociation =
+    resultType.getRank() == 0 ? SmallVector<ReassociationIndices> {} : buildCollapseReassociation(reducedAxes);
   if (isCompileTimeComputable(keepdimsValue))
     return tensor::CollapseShapeOp::create(rewriter, loc, resultType, keepdimsValue, reassociation).getResult();
 
@@ -257,11 +332,13 @@ static Value squeezeReducedAxes(Value keepdimsValue,
   return squeezeCompute.getResult(0);
 }
 
-struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
-  using OpConversionPattern::OpConversionPattern;
+template <typename ReduceMeanOp>
+struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
+  using OpConversionPattern<ReduceMeanOp>::OpConversionPattern;
+  using Adaptor = typename ReduceMeanOp::Adaptor;
 
-  LogicalResult matchAndRewrite(ONNXReduceMeanV13Op reduceMeanOp,
-                                ONNXReduceMeanV13OpAdaptor adaptor,
+  LogicalResult matchAndRewrite(ReduceMeanOp reduceMeanOp,
+                                Adaptor adaptor,
                                 ConversionPatternRewriter& rewriter) const override {
     auto inputType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
     auto resultType = dyn_cast<RankedTensorType>(reduceMeanOp.getReduced().getType());
@@ -272,10 +349,18 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
       return success();
     }
 
-    auto axes = normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputType.getRank());
-    if (failed(axes))
-      return failure();
-    SmallVector<bool> reducedAxes = buildReducedAxesMask(*axes, inputType.getRank());
+    auto semantics = getReduceMeanSemantics(reduceMeanOp, adaptor, inputType.getRank());
+    if (failed(semantics))
+      return rewriter.notifyMatchFailure(reduceMeanOp, "requires compile-time constant, in-range ReduceMean axes");
+    if (semantics->isIdentity) {
+      if (inputType != resultType)
+        return rewriter.notifyMatchFailure(
+          reduceMeanOp, "noop_with_empty_axes identity requires the result type to match the input type");
+      rewriter.replaceOp(reduceMeanOp, adaptor.getData());
+      return success();
+    }
+
+    SmallVector<bool> reducedAxes = buildReducedAxesMask(semantics->axes, inputType.getRank());
     if (reducedAxes.empty() && inputType.getRank() != 0)
       return failure();
 
@@ -295,7 +380,7 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
     Value reducedKeepdims =
       buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc);
 
-    if (reduceMeanOp.getKeepdims() != 0) {
+    if (semantics->keepdims != 0) {
       rewriter.replaceOp(reduceMeanOp, reducedKeepdims);
       return success();
     }
@@ -309,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/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/SpatialToPim/ReturnPathNormalization.cpp

@@ -375,6 +375,57 @@ static void cloneHelperChain(Value sourceValue,
   }
 }
 
+static bool isHostStaticReturnValue(Value value) {
+  llvm::SmallPtrSet<Operation*, 8> visited;
+  while (Operation* definingOp = value.getDefiningOp()) {
+    if (!visited.insert(definingOp).second)
+      return false;
+    if (isa<arith::ConstantOp>(definingOp) || definingOp->hasTrait<OpTrait::ConstantLike>())
+      return true;
+    if (!isReturnHelperChainOp(definingOp) || definingOp->getNumOperands() != 1)
+      return false;
+    value = definingOp->getOperand(0);
+  }
+  return false;
+}
+
+static FailureOr<Value>
+materializeHostStaticReturnValue(IRRewriter& rewriter, Value value, OperationFolder& constantFolder) {
+  llvm::SmallVector<Operation*> chain;
+  llvm::SmallPtrSet<Operation*, 8> visited;
+  while (Operation* definingOp = value.getDefiningOp()) {
+    if (!visited.insert(definingOp).second)
+      return failure();
+    chain.push_back(definingOp);
+    if (isa<arith::ConstantOp>(definingOp) || definingOp->hasTrait<OpTrait::ConstantLike>())
+      break;
+    if (!isReturnHelperChainOp(definingOp) || definingOp->getNumOperands() != 1)
+      return failure();
+    value = definingOp->getOperand(0);
+  }
+
+  if (chain.empty())
+    return failure();
+
+  IRMapping mapping;
+  Value clonedValue;
+  for (Operation* op : llvm::reverse(chain)) {
+    if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
+      clonedValue = getOrCreateConstantLike(constantFolder, constantOp);
+      mapping.map(op->getResult(0), clonedValue);
+      continue;
+    }
+
+    Operation* clonedOp = rewriter.clone(*op, mapping);
+    for (auto [originalResult, newResult] : llvm::zip(op->getResults(), clonedOp->getResults()))
+      mapping.map(originalResult, newResult);
+    clonedValue = clonedOp->getResult(0);
+    rewriter.setInsertionPointAfter(clonedOp);
+  }
+
+  return clonedValue;
+}
+
 static FailureOr<Value> emitHostCopy(IRRewriter& rewriter,
                                      Location loc,
                                      Value outputTensor,
@@ -444,7 +495,30 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
   OperationFolder constantFolder(producerOp->getContext());
   auto storedTensorType = cast<TensorType>(storedValue.getType());
 
+  auto materializeDirectHostReturn = [&](size_t returnIndex,
+                                         Value sourceValue,
+                                         ArrayRef<Operation*> helperChain) -> ReturnPathLoweringResult {
+    rewriter.setInsertionPointAfter(producerOp);
+    auto hostStaticValue = materializeHostStaticReturnValue(rewriter, sourceValue, constantFolder);
+    if (failed(hostStaticValue))
+      return ReturnPathLoweringResult::Failure;
+
+    Value hostReturnValue = *hostStaticValue;
+    if (!helperChain.empty())
+      cloneHelperChain(hostReturnValue, helperChain, rewriter, constantFolder, hostReturnValue);
+
+    outputTensors[returnIndex] =
+      [hostReturnValue](IRRewriter& rewriter, Location loc) -> Value { return hostReturnValue; };
+    return ReturnPathLoweringResult::Handled;
+  };
+
   if (auto returnUse = analyzeReturnUse(producedValue)) {
+    if (isHostStaticReturnValue(storedValue)) {
+      for (Operation* op : returnUse->helperChain)
+        markOpToRemove(op);
+      return materializeDirectHostReturn(returnUse->returnIndex, storedValue, returnUse->helperChain);
+    }
+
     Value currentStoredValue = storedValue;
     cloneHelperChain(storedValue, returnUse->helperChain, rewriter, constantFolder, currentStoredValue);
     for (Operation* op : returnUse->helperChain)
@@ -470,6 +544,8 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
 
     if (isa<func::ReturnOp>(resultUser)) {
       size_t resultIndexInReturn = resultUse.getOperandNumber();
+      if (isHostStaticReturnValue(storedValue))
+        return materializeDirectHostReturn(resultIndexInReturn, storedValue, {});
       auto byteSize =
         pim::getCheckedShapedTypeSizeInBytes(storedTensorType, producerOp, "return-path host copy byte size");
       if (failed(byteSize))

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

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

src/PIM/Dialect/Pim/CMakeLists.txt

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

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

@@ -6,6 +6,7 @@
 #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;
@@ -13,7 +14,9 @@ using namespace bufferization;
 namespace onnx_mlir::pim {
 
 FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue, Location loc, RewriterBase& rewriter) {
-  if (succeeded(resolveContiguousAddress(memrefValue)) || succeeded(compileContiguousAddressExpr(memrefValue)))
+  bool isContiguous =
+    succeeded(resolveContiguousAddress(memrefValue)) || succeeded(compileContiguousAddressExpr(memrefValue));
+  if (isContiguous && isDeviceLocalPimAddress(memrefValue))
     return memrefValue;
 
   auto shapedType = cast<ShapedType>(memrefValue.getType());
@@ -29,13 +32,21 @@ FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue, Location lo
   if (failed(sizeAttr))
     return failure();
 
+  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 allocateContiguousResultMemRefLike(Value memrefValue, Location loc, RewriterBase& rewriter) {
-  if (succeeded(resolveContiguousAddress(memrefValue)))
+  bool isContiguous =
+    succeeded(resolveContiguousAddress(memrefValue)) || succeeded(compileContiguousAddressExpr(memrefValue));
+  if (isContiguous && isDeviceLocalPimAddress(memrefValue))
     return memrefValue;
 
   auto shapedType = cast<ShapedType>(memrefValue.getType());

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

@@ -1,9 +1,70 @@
 #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;
 
+static bool isCoreBatchInputArgument(Value value) {
+  auto blockArg = dyn_cast<BlockArgument>(value);
+  if (!blockArg)
+    return false;
+
+  auto coreBatchOp = dyn_cast_or_null<onnx_mlir::pim::PimCoreBatchOp>(blockArg.getOwner()->getParentOp());
+  if (!coreBatchOp)
+    return false;
+
+  unsigned firstInputArg = 1 + coreBatchOp.getWeights().size();
+  return static_cast<unsigned>(blockArg.getArgNumber()) >= firstInputArg;
+}
+
+static FailureOr<Value> getPimStorageBase(Value value, const onnx_mlir::StaticValueKnowledge& knowledge) {
+  llvm::SmallPtrSet<Value, 8> visited;
+  while (value && visited.insert(value).second) {
+    Value alias = resolveLoopCarriedAlias(value, knowledge);
+    if (alias)
+      value = alias;
+
+    if (auto aliased = knowledge.aliases.lookup(value)) {
+      value = aliased;
+      continue;
+    }
+
+    if (auto base = onnx_mlir::pim::getPimAddressBase(value, knowledge); succeeded(base))
+      return base;
+
+    if (isa<BlockArgument>(value))
+      return value;
+
+    Operation* definingOp = value.getDefiningOp();
+    if (!definingOp)
+      return value;
+
+    if (auto subviewOp = dyn_cast<memref::SubViewOp>(definingOp)) {
+      value = subviewOp.getSource();
+      continue;
+    }
+    if (auto collapseOp = dyn_cast<memref::CollapseShapeOp>(definingOp)) {
+      value = collapseOp.getSrc();
+      continue;
+    }
+    if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(definingOp)) {
+      value = expandOp.getSrc();
+      continue;
+    }
+    if (auto castOp = dyn_cast<memref::CastOp>(definingOp)) {
+      value = castOp.getSource();
+      continue;
+    }
+
+    return value;
+  }
+
+  if (value)
+    return value;
+  return failure();
+}
+
 FailureOr<IntegerAttr> onnx_mlir::pim::getMemRefSizeInBytesAttr(OpBuilder& builder, Operation* anchor, Value memref) {
   auto type = mlir::cast<MemRefType>(memref.getType());
   auto byteSize = getCheckedShapedTypeSizeInBytes(type, anchor, "memref byte size");
@@ -11,3 +72,40 @@ FailureOr<IntegerAttr> onnx_mlir::pim::getMemRefSizeInBytesAttr(OpBuilder& build
     return failure();
   return getCheckedI32Attr(builder, anchor, *byteSize, "memref byte size");
 }
+
+FailureOr<Value> onnx_mlir::pim::getPimAddressBase(Value value, const StaticValueKnowledge& knowledge) {
+  Value alias = resolveLoopCarriedAlias(value, knowledge);
+  if (alias)
+    value = alias;
+
+  auto resolved = resolveContiguousAddress(value, knowledge);
+  if (succeeded(resolved))
+    return resolved->base;
+
+  auto compiled = compileContiguousAddressExpr(value);
+  if (failed(compiled)) {
+    if (isa<BlockArgument>(value))
+      return value;
+    return failure();
+  }
+  return compiled->base;
+}
+
+bool onnx_mlir::pim::isHostBackedPimAddress(Value value, const StaticValueKnowledge& knowledge) {
+  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,11 +2,19 @@
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+
 namespace onnx_mlir {
 namespace pim {
 
 mlir::FailureOr<mlir::IntegerAttr>
 getMemRefSizeInBytesAttr(mlir::OpBuilder& builder, mlir::Operation* anchor, mlir::Value memref);
 
+mlir::FailureOr<mlir::Value> getPimAddressBase(mlir::Value value, const StaticValueKnowledge& knowledge = {});
+
+bool isHostBackedPimAddress(mlir::Value value, const StaticValueKnowledge& knowledge = {});
+
+bool isDeviceLocalPimAddress(mlir::Value value, const StaticValueKnowledge& knowledge = {});
+
 } // namespace pim
 } // namespace onnx_mlir

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

@@ -8,6 +8,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Casting.h"
 
 #include "Common/PimCommon.hpp"
@@ -15,6 +16,7 @@
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
 #include "Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
+#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
@@ -27,24 +29,71 @@ namespace onnx_mlir {
 
 namespace {
 
-struct MemRefCopyToPimMemCopyPattern final : OpRewritePattern<memref::CopyOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct MemRefCopyWorkItem {
+  memref::CopyOp copyOp;
+  StaticValueKnowledge knowledge;
+};
 
-  LogicalResult matchAndRewrite(memref::CopyOp copyOp, PatternRewriter& rewriter) const override {
-    if (!copyOp->getParentOfType<pim::PimCoreOp>() && !copyOp->getParentOfType<pim::PimCoreBatchOp>())
-      return failure();
+static StaticValueKnowledge seedCoreKnowledge(pim::PimCoreOp coreOp) {
+  StaticValueKnowledge knowledge;
+  for (auto [index, weight] : llvm::enumerate(coreOp.getWeights()))
+    knowledge.aliases[coreOp.getWeightArgument(index)] = weight;
+  return knowledge;
+}
 
-    auto sourceType = dyn_cast<MemRefType>(copyOp.getSource().getType());
-    auto targetType = dyn_cast<MemRefType>(copyOp.getTarget().getType());
-    if (!sourceType || !targetType || !sourceType.hasStaticShape() || !targetType.hasStaticShape())
-      return failure();
-    if (sourceType.getElementType() != targetType.getElementType())
-      return failure();
+static StaticValueKnowledge seedCoreBatchKnowledge(pim::PimCoreBatchOp coreBatchOp, unsigned lane) {
+  StaticValueKnowledge knowledge;
+  knowledge.indexValues[coreBatchOp.getLaneArgument()] = lane;
+  for (auto [index, weight] : llvm::enumerate(coreBatchOp.getWeights()))
+    knowledge.aliases[coreBatchOp.getWeightArgument(index)] = weight;
+  for (auto [index, input] : llvm::enumerate(coreBatchOp.getInputs()))
+    knowledge.aliases[coreBatchOp.getInputArgument(index)] = input;
+  return knowledge;
+}
 
-    Value zeroOffset = getOrCreateIndexConstant(rewriter, copyOp, 0);
-    auto sizeAttr = getMemRefSizeInBytesAttr(rewriter, copyOp.getOperation(), copyOp.getSource());
-    if (failed(sizeAttr))
-      return failure();
+static LogicalResult
+lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const StaticValueKnowledge& knowledge) {
+  if (!copyOp->getParentOfType<pim::PimCoreOp>() && !copyOp->getParentOfType<pim::PimCoreBatchOp>())
+    return failure();
+
+  auto sourceType = dyn_cast<MemRefType>(copyOp.getSource().getType());
+  auto targetType = dyn_cast<MemRefType>(copyOp.getTarget().getType());
+  if (!sourceType || !targetType || !sourceType.hasStaticShape() || !targetType.hasStaticShape())
+    return failure();
+  if (sourceType.getElementType() != targetType.getElementType())
+    return failure();
+
+  Value zeroOffset = getOrCreateIndexConstant(rewriter, copyOp, 0);
+  auto sizeAttr = getMemRefSizeInBytesAttr(rewriter, copyOp.getOperation(), copyOp.getSource());
+  if (failed(sizeAttr))
+    return failure();
+
+  bool sourceIsHost = isHostBackedPimAddress(copyOp.getSource(), knowledge);
+  bool targetIsHost = isHostBackedPimAddress(copyOp.getTarget(), knowledge);
+  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge);
+  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge);
+
+  if (targetIsDevice && sourceIsHost) {
+    pim::PimMemCopyHostToDevOp::create(rewriter,
+                                       copyOp.getLoc(),
+                                       copyOp.getTarget().getType(),
+                                       zeroOffset,
+                                       zeroOffset,
+                                       copyOp.getTarget(),
+                                       copyOp.getSource(),
+                                       *sizeAttr);
+  }
+  else if (targetIsHost && sourceIsDevice) {
+    pim::PimMemCopyDevToHostOp::create(rewriter,
+                                       copyOp.getLoc(),
+                                       copyOp.getTarget().getType(),
+                                       zeroOffset,
+                                       zeroOffset,
+                                       copyOp.getTarget(),
+                                       copyOp.getSource(),
+                                       *sizeAttr);
+  }
+  else if (targetIsDevice && sourceIsDevice) {
     pim::PimMemCopyOp::create(rewriter,
                               copyOp.getLoc(),
                               copyOp.getTarget().getType(),
@@ -53,10 +102,19 @@ struct MemRefCopyToPimMemCopyPattern final : OpRewritePattern<memref::CopyOp> {
                               copyOp.getTarget(),
                               copyOp.getSource(),
                               *sizeAttr);
-    rewriter.eraseOp(copyOp);
-    return success();
   }
-};
+  else {
+    copyOp.emitOpError() << "failed to classify memref.copy endpoints: source=" << copyOp.getSource()
+                         << " type=" << copyOp.getSource().getType() << " host=" << sourceIsHost
+                         << " device=" << sourceIsDevice << ", target=" << copyOp.getTarget()
+                         << " type=" << copyOp.getTarget().getType() << " host=" << targetIsHost
+                         << " device=" << targetIsDevice;
+    return failure();
+  }
+
+  rewriter.eraseOp(copyOp);
+  return success();
+}
 
 struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass)
@@ -100,25 +158,46 @@ void PimBufferizationPass::runOnOperation() {
   }
 
   MLIRContext* ctx = moduleOp.getContext();
-  RewritePatternSet memrefCopyPatterns(ctx);
-  memrefCopyPatterns.add<MemRefCopyToPimMemCopyPattern>(ctx);
-  FrozenRewritePatternSet frozenMemrefCopyPatterns(std::move(memrefCopyPatterns));
-  PatternApplicator memrefCopyApplicator(frozenMemrefCopyPatterns);
-  memrefCopyApplicator.applyDefaultCostModel();
   PatternRewriter rewriter(ctx);
 
-  SmallVector<memref::CopyOp> copyWorklist;
-  moduleOp.walk([&](memref::CopyOp copyOp) {
-    if (copyOp->getParentOfType<pim::PimCoreOp>() || copyOp->getParentOfType<pim::PimCoreBatchOp>())
-      copyWorklist.push_back(copyOp);
+  SmallVector<MemRefCopyWorkItem> copyWorklist;
+  llvm::SmallPtrSet<Operation*, 16> seenCopyOps;
+  auto addCopyOp = [&](memref::CopyOp copyOp, const StaticValueKnowledge& knowledge) {
+    if (seenCopyOps.insert(copyOp.getOperation()).second)
+      copyWorklist.push_back({copyOp, knowledge});
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) {
+    StaticValueKnowledge knowledge = seedCoreKnowledge(coreOp);
+    (void) walkPimCoreBlockStructurally(
+      coreOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
+        if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
+          addCopyOp(copyOp, opKnowledge);
+        return success();
+      });
+  });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    llvm::SmallVector<unsigned, 2> lanes;
+    lanes.push_back(0);
+    if (coreBatchOp.getLaneCount() > 1)
+      lanes.push_back(static_cast<unsigned>(coreBatchOp.getLaneCount() - 1));
+    for (unsigned lane : lanes) {
+      StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane);
+      (void) walkPimCoreBlockStructurally(
+        coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
+          if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
+            addCopyOp(copyOp, opKnowledge);
+          return success();
+        });
+    }
   });
 
   bool hasFailed = false;
-  for (memref::CopyOp copyOp : copyWorklist) {
-    if (failed(applyPatternsOnce(copyOp, memrefCopyApplicator, rewriter))) {
-      copyOp.emitOpError("failed to lower memref.copy inside PIM core body");
+  for (const MemRefCopyWorkItem& workItem : copyWorklist) {
+    memref::CopyOp copyOp = workItem.copyOp;
+    rewriter.setInsertionPoint(copyOp);
+    if (failed(lowerMemRefCopyToPimCopy(copyOp, rewriter, workItem.knowledge)))
       hasFailed = true;
-    }
   }
   if (hasFailed) {
     signalPassFailure();

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

@@ -128,7 +128,7 @@ struct FoldConstantCoreMapPattern final : OpRewritePattern<linalg::MapOp> {
     auto sizeAttr = pim::getCheckedI32Attr(rewriter, mapOp, *sizeInBytes, "host constant folding byte size");
     if (failed(sizeAttr))
       return failure();
-    pim::PimMemCopyOp::create(
+    pim::PimMemCopyHostToDevOp::create(
       rewriter, mapOp.getLoc(), initType, zeroOffset, zeroOffset, mapOp.getInit(), getGlobalOp.getResult(), *sizeAttr);
     rewriter.eraseOp(mapOp);
     return success();

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

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

src/PIM/Dialect/Pim/Transforms/HostConstantMaterialization/MaterializeHostConstantsPass.cpp

@@ -1,161 +0,0 @@
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/Dominance.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Pass/Pass.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/MathExtras.h"
-
-#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/Dialect/Pim/PimOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-
-namespace {
-
-template <typename CoreOpTy>
-static void materializeHostConstantsInCore(CoreOpTy coreOp,
-                                           IRRewriter& rewriter,
-                                           OperationFolder& constantFolder,
-                                           bool& hasFailure) {
-  DenseMap<Value, DenseMap<int64_t, DenseMap<Type, Value>>> materializedValues;
-  DominanceInfo dominance(coreOp);
-  SmallVector<Operation*> ops;
-  coreOp.getBody().front().walk([&](Operation* op) {
-    if (!isa<pim::PimHaltOp, scf::YieldOp>(op))
-      ops.push_back(op);
-  });
-
-  for (Operation* op : ops) {
-    if (auto loadOp = dyn_cast<memref::LoadOp>(op); loadOp && loadOp.getType().isIndex())
-      continue;
-
-    for (OpOperand& operand : op->getOpOperands()) {
-      Value originalValue = operand.get();
-      if (!isa<BaseMemRefType>(originalValue.getType()) || isExplicitHostMemCopyOperand(op, operand.getOperandNumber()))
-        continue;
-
-      auto resolvedAddress = resolveContiguousAddress(originalValue);
-      if (failed(resolvedAddress))
-        continue;
-
-      auto getGlobalOp = dyn_cast_or_null<memref::GetGlobalOp>(resolvedAddress->base.getDefiningOp());
-      if (!getGlobalOp)
-        continue;
-
-      auto originalType = dyn_cast<MemRefType>(originalValue.getType());
-      if (!originalType || !originalType.hasStaticShape()) {
-        op->emitOpError("host constant materialization requires a static memref operand");
-        hasFailure = true;
-        continue;
-      }
-
-      auto& cachedByOffset = materializedValues[resolvedAddress->base];
-      auto& cachedByType = cachedByOffset[resolvedAddress->byteOffset];
-      auto cachedValue = cachedByType.find(originalType);
-      if (cachedValue != cachedByType.end() && dominance.properlyDominates(cachedValue->second, op)) {
-        operand.set(cachedValue->second);
-        continue;
-      }
-
-      auto type = dyn_cast<ShapedType>(originalValue.getType());
-      auto totalBytes = type ? pim::getCheckedShapedTypeSizeInBytes(type, op, "host constant materialization byte size")
-                             : FailureOr<uint64_t>(failure());
-      auto totalBytesAttr =
-        succeeded(totalBytes)
-          ? pim::getCheckedI32Attr(rewriter, op, *totalBytes, "host constant materialization byte size")
-          : FailureOr<IntegerAttr>(failure());
-      if (failed(totalBytesAttr)
-          || failed(pim::checkedSize(resolvedAddress->byteOffset, op, "host constant materialization byte offset"))) {
-        hasFailure = true;
-        continue;
-      }
-
-      auto contiguousType = MemRefType::get(originalType.getShape(), originalType.getElementType());
-
-      rewriter.setInsertionPoint(op);
-      Value localAlloc = memref::AllocOp::create(rewriter, op->getLoc(), contiguousType);
-      Value deviceDst = localAlloc;
-      if (contiguousType != originalType)
-        deviceDst = memref::CastOp::create(rewriter, op->getLoc(), originalType, localAlloc);
-
-      Value zeroOffset = getOrCreateIndexConstant(constantFolder, op, 0);
-      Value hostOffset = getOrCreateIndexConstant(constantFolder, op, resolvedAddress->byteOffset);
-      Value copiedValue = pim::PimMemCopyHostToDevOp::create(rewriter,
-                                                             op->getLoc(),
-                                                             originalType,
-                                                             zeroOffset,
-                                                             hostOffset,
-                                                             deviceDst,
-                                                             getGlobalOp.getResult(),
-                                                             *totalBytesAttr)
-                            .getOutput();
-
-      cachedByType[originalType] = copiedValue;
-      operand.set(copiedValue);
-    }
-  }
-}
-
-struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MaterializeHostConstantsPass)
-
-  StringRef getArgument() const override { return "materialize-pim-host-constants"; }
-  StringRef getDescription() const override {
-    return "Materialize explicit host-to-device copies for constant globals used by PIM runtime ops";
-  }
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-    IRRewriter rewriter(moduleOp.getContext());
-    OperationFolder constantFolder(moduleOp.getContext());
-    bool hasFailure = false;
-
-    for (func::FuncOp funcOp : moduleOp.getOps<func::FuncOp>()) {
-      if (funcOp.isExternal())
-        continue;
-
-      for (pim::PimCoreOp coreOp : funcOp.getOps<pim::PimCoreOp>())
-        materializeHostConstantsInCore(coreOp, rewriter, constantFolder, hasFailure);
-
-      for (pim::PimCoreBatchOp coreBatchOp : funcOp.getOps<pim::PimCoreBatchOp>())
-        materializeHostConstantsInCore(coreBatchOp, rewriter, constantFolder, hasFailure);
-
-      SmallVector<Operation*> hostCompactOps;
-      for (Operation& op : funcOp.getBody().front())
-        if (isa<pim::PimConcatOp>(op))
-          hostCompactOps.push_back(&op);
-
-      for (Operation* op : hostCompactOps) {
-        rewriter.setInsertionPoint(op);
-        auto concatOp = cast<pim::PimConcatOp>(op);
-        concatOp.emitOpError("host-side concat must be folded away or lowered into pim.core before materialization");
-        hasFailure = true;
-      }
-    }
-
-    if (hasFailure) {
-      moduleOp.emitError("PIM host-constant materialization failed; see diagnostics above");
-      signalPassFailure();
-      return;
-    }
-
-    dumpModule(moduleOp, "pim4_materialized");
-  }
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createPimMaterializeHostConstantsPass() {
-  return std::make_unique<MaterializeHostConstantsPass>();
-}
-
-} // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescing.cpp

@@ -32,32 +32,51 @@ static uint64_t getTypeSizeBytes(MemRefType type) {
   return static_cast<uint64_t>(type.getNumElements() * getElementTypeSizeInBytes(type.getElementType()));
 }
 
-static Operation* getTopLevelAncestorInBody(Operation* op, Block& body) {
+static Operation* getTopLevelAncestorInBlock(Operation* op, Block& block) {
   Operation* current = op;
-  while (current && current->getBlock() != &body)
+  while (current && current->getBlock() != &block)
     current = current->getParentOp();
   return current;
 }
 
+static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis);
+
 static FailureOr<uint64_t>
-getLastUseInstruction(memref::AllocOp allocOp, Block& body, const DenseMap<Operation*, uint64_t>& opOrder) {
+getLastUseInstruction(memref::AllocOp allocOp, Block& scopeBlock, const DenseMap<Operation*, uint64_t>& opOrder) {
   uint64_t endInstruction = opOrder.lookup(allocOp);
-  SmallPtrSet<Operation*, 16> visited;
+  SmallPtrSet<Value, 16> visitedValues;
+  SmallPtrSet<Operation*, 16> visitedUsers;
   SmallVector<Value> pendingValues;
   pendingValues.push_back(allocOp.getResult());
 
   while (!pendingValues.empty()) {
     Value value = pendingValues.pop_back_val();
+    if (!visitedValues.insert(value).second)
+      continue;
+
     for (Operation* user : value.getUsers()) {
-      Operation* orderedUser = getTopLevelAncestorInBody(user, body);
-      if (!orderedUser)
-        return failure();
-      if (!visited.insert(user).second)
+      if (!visitedUsers.insert(user).second)
         continue;
 
       if (isSupportedAliasOp(user))
-        for (Value result : user->getResults())
-          pendingValues.push_back(result);
+        llvm::append_range(pendingValues, user->getResults());
+
+      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(user)) {
+        for (OpResult result : user->getResults()) {
+          OpOperand* tiedOperand = dpsOp.getTiedOpOperand(result);
+          if (tiedOperand && tiedOperand->get() == value)
+            pendingValues.push_back(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));
+        }
+      }
 
       if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
         auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
@@ -68,20 +87,9 @@ getLastUseInstruction(memref::AllocOp allocOp, Block& body, const DenseMap<Opera
             pendingValues.push_back(forOp.getResult(index));
       }
 
-      if (auto forOp = dyn_cast<scf::ForOp>(user)) {
-        for (auto [index, initArg] : llvm::enumerate(forOp.getInitArgs()))
-          if (initArg == value)
-            pendingValues.push_back(forOp.getResult(index));
-      }
-
-      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);
-        }
-      }
+      Operation* orderedUser = getTopLevelAncestorInBlock(user, scopeBlock);
+      if (!orderedUser)
+        return failure();
 
       auto order = opOrder.find(orderedUser);
       if (order == opOrder.end())
@@ -93,101 +101,126 @@ getLastUseInstruction(memref::AllocOp allocOp, Block& body, const DenseMap<Opera
   return endInstruction;
 }
 
-} // namespace
+static void analyzeBlock(Block& block, MemoryCoalescingAnalysis& analysis) {
+  for (Operation& op : block)
+    for (Region& region : op.getRegions())
+      for (Block& nestedBlock : region)
+        analyzeBlock(nestedBlock, analysis);
 
-MemoryCoalescingAnalysis analyzeMemoryCoalescingCandidates(Operation* coreLikeOp) {
-  MemoryCoalescingAnalysis analysis;
-  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
-    return analysis;
-
-  Block& body = coreLikeOp->getRegion(0).front();
   DenseMap<Operation*, uint64_t> opOrder;
   uint64_t nextInstruction = 0;
-  for (Operation& op : body)
+  for (Operation& op : block)
     opOrder.try_emplace(&op, nextInstruction++);
 
-  for (Operation& op : body) {
+  MemoryCoalescingBlockAnalysis blockAnalysis;
+  blockAnalysis.block = &block;
+
+  for (Operation& op : block) {
     auto allocOp = dyn_cast<memref::AllocOp>(&op);
     if (!allocOp)
       continue;
 
     auto allocType = dyn_cast<MemRefType>(allocOp.getType());
     if (!isCandidateAllocType(allocType)) {
-      ++analysis.skippedAllocations;
+      ++blockAnalysis.skippedAllocations;
       continue;
     }
 
-    auto endInstruction = getLastUseInstruction(allocOp, body, opOrder);
+    auto endInstruction = getLastUseInstruction(allocOp, block, opOrder);
     if (failed(endInstruction)) {
-      ++analysis.skippedAllocations;
+      ++blockAnalysis.skippedAllocations;
       continue;
     }
 
-    analysis.candidates.push_back(
-      AllocationCandidate {allocOp, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)});
+    blockAnalysis.candidates.push_back(
+      AllocationCandidate {allocOp, &block, opOrder.lookup(allocOp), *endInstruction, getTypeSizeBytes(allocType)});
   }
 
+  analysis.skippedAllocations += blockAnalysis.skippedAllocations;
+  if (!blockAnalysis.candidates.empty() || blockAnalysis.skippedAllocations != 0)
+    analysis.blocks.push_back(std::move(blockAnalysis));
+}
+
+} // namespace
+
+uint64_t MemoryCoalescingAnalysis::getCandidateCount() const {
+  uint64_t total = 0;
+  for (const MemoryCoalescingBlockAnalysis& block : blocks)
+    total += block.candidates.size();
+  return total;
+}
+
+MemoryCoalescingAnalysis analyzeMemoryCoalescingCandidates(Operation* coreLikeOp) {
+  MemoryCoalescingAnalysis analysis;
+  if (!coreLikeOp || coreLikeOp->getNumRegions() != 1 || coreLikeOp->getRegion(0).empty())
+    return analysis;
+
+  analyzeBlock(coreLikeOp->getRegion(0).front(), analysis);
   return analysis;
 }
 
 MemoryCoalescingStats
 coalesceMemory(Operation* coreLikeOp, const MemoryCoalescingAnalysis& analysis, RewriterBase& rewriter) {
+  (void) coreLikeOp;
+
   MemoryCoalescingStats stats;
   stats.skippedAllocations = analysis.skippedAllocations;
 
-  auto candidates = analysis.candidates;
-  llvm::sort(candidates, [](const AllocationCandidate& lhs, const AllocationCandidate& rhs) {
-    if (lhs.startInstruction != rhs.startInstruction)
-      return lhs.startInstruction < rhs.startInstruction;
-    return lhs.endInstruction < rhs.endInstruction;
-  });
+  for (const MemoryCoalescingBlockAnalysis& blockAnalysis : analysis.blocks) {
+    auto candidates = blockAnalysis.candidates;
+    llvm::sort(candidates, [](const AllocationCandidate& lhs, const AllocationCandidate& rhs) {
+      if (lhs.startInstruction != rhs.startInstruction)
+        return lhs.startInstruction < rhs.startInstruction;
+      return lhs.endInstruction < rhs.endInstruction;
+    });
 
-  struct ActiveStorage {
-    memref::AllocOp root;
-    uint64_t endInstruction = 0;
-  };
+    struct ActiveStorage {
+      memref::AllocOp root;
+      uint64_t endInstruction = 0;
+    };
 
-  SmallVector<ActiveStorage> active;
-  SmallVector<memref::AllocOp> freeList;
+    SmallVector<ActiveStorage> active;
+    SmallVector<memref::AllocOp> freeList;
 
-  for (AllocationCandidate& candidate : candidates) {
-    for (auto it = active.begin(); it != active.end();) {
-      if (it->endInstruction < candidate.startInstruction) {
-        freeList.push_back(it->root);
-        it = active.erase(it);
+    for (AllocationCandidate& candidate : candidates) {
+      for (auto it = active.begin(); it != active.end();) {
+        if (it->endInstruction < candidate.startInstruction) {
+          freeList.push_back(it->root);
+          it = active.erase(it);
+          continue;
+        }
+        ++it;
+      }
+
+      auto bestFit = freeList.end();
+      uint64_t bestFitBytes = std::numeric_limits<uint64_t>::max();
+      auto candidateType = cast<MemRefType>(candidate.alloc.getType());
+      for (auto it = freeList.begin(); it != freeList.end(); ++it) {
+        auto freeType = cast<MemRefType>((*it).getType());
+        if (freeType != candidateType)
+          continue;
+
+        uint64_t freeBytes = getTypeSizeBytes(freeType);
+        if (freeBytes < candidate.sizeBytes || freeBytes >= bestFitBytes)
+          continue;
+
+        bestFit = it;
+        bestFitBytes = freeBytes;
+      }
+
+      if (bestFit == freeList.end()) {
+        active.push_back(ActiveStorage {candidate.alloc, candidate.endInstruction});
         continue;
       }
-      ++it;
+
+      memref::AllocOp root = *bestFit;
+      freeList.erase(bestFit);
+      candidate.alloc.getResult().replaceAllUsesWith(root.getResult());
+      rewriter.eraseOp(candidate.alloc);
+      active.push_back(ActiveStorage {root, candidate.endInstruction});
+      ++stats.removedAllocs;
+      stats.savedBytes += candidate.sizeBytes;
     }
-
-    auto bestFit = freeList.end();
-    uint64_t bestFitBytes = std::numeric_limits<uint64_t>::max();
-    auto candidateType = cast<MemRefType>(candidate.alloc.getType());
-    for (auto it = freeList.begin(); it != freeList.end(); ++it) {
-      auto freeType = cast<MemRefType>((*it).getType());
-      if (freeType != candidateType)
-        continue;
-
-      uint64_t freeBytes = getTypeSizeBytes(freeType);
-      if (freeBytes < candidate.sizeBytes || freeBytes >= bestFitBytes)
-        continue;
-
-      bestFit = it;
-      bestFitBytes = freeBytes;
-    }
-
-    if (bestFit == freeList.end()) {
-      active.push_back(ActiveStorage {candidate.alloc, candidate.endInstruction});
-      continue;
-    }
-
-    memref::AllocOp root = *bestFit;
-    freeList.erase(bestFit);
-    candidate.alloc.getResult().replaceAllUsesWith(root.getResult());
-    rewriter.eraseOp(candidate.alloc);
-    active.push_back(ActiveStorage {root, candidate.endInstruction});
-    ++stats.removedAllocs;
-    stats.savedBytes += candidate.sizeBytes;
   }
 
   return stats;

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

@@ -10,16 +10,25 @@ namespace pim {
 
 struct AllocationCandidate {
   mlir::memref::AllocOp alloc;
+  mlir::Block* scopeBlock = nullptr;
   uint64_t startInstruction = 0;
   uint64_t endInstruction = 0;
   uint64_t sizeBytes = 0;
 };
 
-struct MemoryCoalescingAnalysis {
+struct MemoryCoalescingBlockAnalysis {
+  mlir::Block* block = nullptr;
   llvm::SmallVector<AllocationCandidate> candidates;
   uint64_t skippedAllocations = 0;
 };
 
+struct MemoryCoalescingAnalysis {
+  llvm::SmallVector<MemoryCoalescingBlockAnalysis> blocks;
+  uint64_t skippedAllocations = 0;
+
+  uint64_t getCandidateCount() const;
+};
+
 struct MemoryCoalescingStats {
   uint64_t removedAllocs = 0;
   uint64_t savedBytes = 0;

src/PIM/Dialect/Pim/Transforms/MemoryCoalescing/MemoryCoalescingPass.cpp

@@ -23,9 +23,9 @@ using namespace onnx_mlir::compact_asm;
 namespace onnx_mlir {
 namespace {
 
-// This pass assumes bufferization has already normalized executable PIM
-// operands. It only reuses compatible local allocations with non-overlapping
-// lifetimes; it does not repair memory contiguity.
+// This pass is an IR cleanup step after bufferization. It only rewrites
+// obviously compatible local allocations with non-overlapping lifetimes inside
+// the same block and leaves the final physical memory planning to codegen.
 
 struct CoalescingReportRow {
   uint64_t numCandidates = 0;
@@ -174,7 +174,7 @@ struct PimMemoryCoalescingPass : PassWrapper<PimMemoryCoalescingPass, OperationP
       auto analysis = pim::analyzeMemoryCoalescingCandidates(op);
       auto stats = pim::coalesceMemory(op, analysis, rewriter);
       CoalescingReportRow row {
-        analysis.candidates.size(), stats.skippedAllocations, stats.removedAllocs, stats.savedBytes};
+        analysis.getCandidateCount(), stats.skippedAllocations, stats.removedAllocs, stats.savedBytes};
 
       if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
         auto checkedCoreId =

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

@@ -46,19 +46,6 @@ static bool isCodegenAddressableValue(Value value) {
       || isa<memref::AllocOp, memref::GetGlobalOp>(compiledAddress->base.getDefiningOp());
 }
 
-static bool isCodegenAddressableValue(Value value, const StaticValueKnowledge& knowledge) {
-  auto resolvedAddress = resolveContiguousAddress(value, knowledge);
-  if (succeeded(resolvedAddress))
-    return isa<BlockArgument>(resolvedAddress->base)
-        || isa<memref::AllocOp, memref::GetGlobalOp>(resolvedAddress->base.getDefiningOp());
-
-  auto compiledAddress = compileContiguousAddressExpr(value);
-  if (failed(compiledAddress))
-    return false;
-  return isa<BlockArgument>(compiledAddress->base)
-      || isa<memref::AllocOp, memref::GetGlobalOp>(compiledAddress->base.getDefiningOp());
-}
-
 static bool isConstantGlobalView(Value value) {
   while (true) {
     Operation* defOp = value.getDefiningOp();
@@ -138,6 +125,24 @@ static bool isSupportedCoreInstructionOp(Operation* op) {
              memref::GetGlobalOp>(op);
 }
 
+static bool isHostAddressableValue(Value value, const StaticValueKnowledge& knowledge) {
+  auto resolvedAddress = resolveContiguousAddress(value, knowledge);
+  Value base;
+  if (succeeded(resolvedAddress)) {
+    base = resolvedAddress->base;
+  }
+  else {
+    auto compiledAddress = compileContiguousAddressExpr(value);
+    if (failed(compiledAddress))
+      return false;
+    base = compiledAddress->base;
+  }
+
+  if (isa<BlockArgument>(base))
+    return true;
+  return isa_and_nonnull<memref::GetGlobalOp>(base.getDefiningOp());
+}
+
 struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VerificationPass)
 
@@ -311,10 +316,10 @@ private:
           }
 
           if (isExplicitHostMemCopyOperand(&op, operandIndex)) {
-            if (!isCodegenAddressableValue(operand, knowledge)) {
+            if (!isHostAddressableValue(operand, knowledge)) {
               diagnostics.report(&op, [&](Operation* illegalOp) {
                 illegalOp->emitOpError() << "host operand #" << operandIndex
-                                         << " is not backed by contiguous addressable storage";
+                                         << " must be backed by host-addressable storage";
               });
               hasFailure = true;
             }

src/PIM/Dialect/Spatial/Spatial.td

@@ -257,6 +257,25 @@ def SpatVAddOp : SpatOp<"vadd", []> {
   }];
 }
 
+def SpatVSubOp : SpatOp<"vsub", []> {
+  let summary = "Element-wise subtraction between two tensors; rhs must match lhs or be 1x1";
+
+  let arguments = (ins
+    SpatTensor:$lhs,
+    SpatTensor:$rhs
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` `(` type($lhs) `,` type($rhs) `)` `->` type($output)
+  }];
+}
+
 def SpatVMulOp : SpatOp<"vmul", []> {
   let summary = "Element-wise multiplication between two tensors; rhs must match lhs or be 1x1";
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -254,6 +254,12 @@ LogicalResult SpatVAddOp::verify() {
   return OpTrait::impl::verifySameOperandsAndResultType(*this);
 }
 
+LogicalResult SpatVSubOp::verify() {
+  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
+    return failure();
+  return OpTrait::impl::verifySameOperandsAndResultType(*this);
+}
+
 LogicalResult SpatVMaxOp::verify() {
   if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
     return failure();

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

@@ -1,6 +1,5 @@
 #include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/PatternMatch.h"
@@ -14,20 +13,14 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_os_ostream.h"
-#include "llvm/Support/raw_ostream.h"
 
 #include <algorithm>
-#include <chrono>
 #include <cstddef>
 #include <cstdint>
-#include <cstdlib>
 #include <fstream>
 #include <memory>
 #include <optional>
-#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -51,83 +44,6 @@ using SpatCompute = spatial::SpatCompute;
 using SpatComputeBatch = spatial::SpatComputeBatch;
 using spatial::getProducerValueRef;
 
-bool isMergeProfilingEnabled() { return std::getenv("RAPTOR_PROFILE_MERGE") != nullptr; }
-
-class ScopedMergePhaseTimer {
-public:
-  explicit ScopedMergePhaseTimer(StringRef phaseName)
-  : enabled(isMergeProfilingEnabled()), phase(phaseName.str()) {
-    if (enabled)
-      start = std::chrono::steady_clock::now();
-  }
-
-  ~ScopedMergePhaseTimer() {
-    if (!enabled)
-      return;
-    auto elapsed = std::chrono::steady_clock::now() - start;
-    double millis = std::chrono::duration<double, std::milli>(elapsed).count();
-    llvm::errs() << "[merge-profile] " << phase << ": " << llvm::formatv("{0:F3}", millis) << " ms\n";
-  }
-
-private:
-  bool enabled = false;
-  std::string phase;
-  std::chrono::steady_clock::time_point start;
-};
-
-struct MergeIrCounts {
-  uint64_t topLevelComputeCount = 0;
-  uint64_t topLevelComputeBatchCount = 0;
-  uint64_t scalarChannelSendCount = 0;
-  uint64_t scalarChannelReceiveCount = 0;
-  uint64_t wvmmCount = 0;
-  uint64_t vaddCount = 0;
-  uint64_t scfForCount = 0;
-};
-
-MergeIrCounts collectMergeIrCounts(func::FuncOp funcOp) {
-  MergeIrCounts counts;
-
-  auto countComputeBodyOps = [&](Operation* op) {
-    op->walk([&](Operation* nestedOp) {
-      if (isa<spatial::SpatChannelSendOp>(nestedOp))
-        ++counts.scalarChannelSendCount;
-      else if (isa<spatial::SpatChannelReceiveOp>(nestedOp))
-        ++counts.scalarChannelReceiveCount;
-      else if (isa<spatial::SpatVMMOp>(nestedOp))
-        ++counts.wvmmCount;
-      else if (isa<spatial::SpatVAddOp>(nestedOp))
-        ++counts.vaddCount;
-      else if (isa<scf::ForOp>(nestedOp))
-        ++counts.scfForCount;
-    });
-  };
-
-  for (auto compute : funcOp.getOps<SpatCompute>()) {
-    ++counts.topLevelComputeCount;
-    countComputeBodyOps(compute.getOperation());
-  }
-
-  for (auto batch : funcOp.getOps<SpatComputeBatch>()) {
-    ++counts.topLevelComputeBatchCount;
-    countComputeBodyOps(batch.getOperation());
-  }
-
-  return counts;
-}
-
-void emitMergeIrCounts(StringRef phaseName, func::FuncOp funcOp) {
-  if (!isMergeProfilingEnabled())
-    return;
-
-  MergeIrCounts counts = collectMergeIrCounts(funcOp);
-  llvm::errs() << "[merge-profile] " << phaseName << " counts:"
-               << " compute=" << counts.topLevelComputeCount << " compute_batch=" << counts.topLevelComputeBatchCount
-               << " scalar_send=" << counts.scalarChannelSendCount
-               << " scalar_recv=" << counts.scalarChannelReceiveCount << " wvmm=" << counts.wvmmCount
-               << " vadd=" << counts.vaddCount << " scf_for=" << counts.scfForCount << "\n";
-}
-
 static std::optional<int32_t> getComputeCoreId(SpatCompute compute) {
   if (auto coreIdAttr = compute->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName)) {
     auto checkedCoreId = pim::checkedI32(coreIdAttr.getInt(), compute, "merge compute core id");
@@ -138,16 +54,6 @@ static std::optional<int32_t> getComputeCoreId(SpatCompute compute) {
   return std::nullopt;
 }
 
-struct ComputeMotifInfo {
-  uint64_t instructionCount = 0;
-  uint64_t weightedVmmCount = 0;
-};
-
-void appendUnique(SmallVector<size_t>& values, size_t value) {
-  if (!llvm::is_contained(values, value))
-    values.push_back(value);
-}
-
 bool isTrivialSerialMergeCandidate(SpatCompute compute) {
   if (!compute->hasOneUse())
     return false;
@@ -266,212 +172,6 @@ void mergeTriviallyConnectedComputes(func::FuncOp funcOp) {
   }
 }
 
-void emitMotifProfile(func::FuncOp funcOp) {
-  if (!std::getenv("DCP_MOTIF_PROFILE"))
-    return;
-
-  SmallVector<SpatCompute> computes(funcOp.getOps<SpatCompute>());
-  DenseMap<SpatCompute, size_t> computeToIndex;
-  computeToIndex.reserve(computes.size());
-  for (auto [index, compute] : llvm::enumerate(computes))
-    computeToIndex[compute] = index;
-
-  SmallVector<ComputeMotifInfo> computeInfos(computes.size());
-  SmallVector<SmallVector<size_t>> parents(computes.size());
-  SmallVector<SmallVector<size_t>> children(computes.size());
-
-  uint64_t weightedVmmNodeCount = 0;
-  uint64_t weightedVmmOpCount = 0;
-  uint64_t edgeCount = 0;
-
-  for (auto [index, compute] : llvm::enumerate(computes)) {
-    ComputeMotifInfo& info = computeInfos[index];
-    info.instructionCount = spatial::countComputeBodyInstructions(compute.getBody());
-    compute.getBody().walk([&](spatial::SpatVMMOp) { info.weightedVmmCount++; });
-    if (info.weightedVmmCount > 0) {
-      weightedVmmNodeCount++;
-      weightedVmmOpCount += info.weightedVmmCount;
-    }
-
-    for (Value input : compute.getInputs()) {
-      auto parent = dyn_cast<SpatCompute>(input.getDefiningOp());
-      if (!parent || parent == compute)
-        continue;
-      auto parentIt = computeToIndex.find(parent);
-      if (parentIt == computeToIndex.end())
-        continue;
-
-      size_t parentIndex = parentIt->second;
-      size_t oldParentCount = parents[index].size();
-      appendUnique(parents[index], parentIndex);
-      if (parents[index].size() != oldParentCount) {
-        appendUnique(children[parentIndex], index);
-        edgeCount++;
-      }
-    }
-  }
-
-  uint64_t maxFanIn = 0;
-  uint64_t maxFanOut = 0;
-  uint64_t fanIn16 = 0;
-  uint64_t fanIn64 = 0;
-  uint64_t fanIn256 = 0;
-  uint64_t fanOut16 = 0;
-  uint64_t fanOut64 = 0;
-  uint64_t fanOut256 = 0;
-  for (size_t index = 0; index < computes.size(); ++index) {
-    uint64_t fanIn = parents[index].size();
-    uint64_t fanOut = children[index].size();
-    maxFanIn = std::max(maxFanIn, fanIn);
-    maxFanOut = std::max(maxFanOut, fanOut);
-    fanIn16 += fanIn >= 16;
-    fanIn64 += fanIn >= 64;
-    fanIn256 += fanIn >= 256;
-    fanOut16 += fanOut >= 16;
-    fanOut64 += fanOut >= 64;
-    fanOut256 += fanOut >= 256;
-  }
-
-  uint64_t serialChainCount = 0;
-  uint64_t serialChainNodeCount = 0;
-  uint64_t maxSerialChain = 0;
-  for (size_t index = 0; index < computes.size(); ++index) {
-    if (parents[index].size() == 1 && children[parents[index][0]].size() == 1)
-      continue;
-
-    uint64_t chainLength = 1;
-    size_t current = index;
-    while (children[current].size() == 1) {
-      size_t child = children[current][0];
-      if (parents[child].size() != 1)
-        break;
-      chainLength++;
-      current = child;
-    }
-
-    if (chainLength >= 2) {
-      serialChainCount++;
-      serialChainNodeCount += chainLength;
-      maxSerialChain = std::max(maxSerialChain, chainLength);
-    }
-  }
-
-  SmallVector<size_t> incomingEdgeCount;
-  incomingEdgeCount.reserve(parents.size());
-  for (ArrayRef<size_t> parentList : parents)
-    incomingEdgeCount.push_back(parentList.size());
-
-  SmallVector<uint64_t> level(computes.size(), 0);
-  SmallVector<size_t> readyNodes;
-  readyNodes.reserve(computes.size());
-  for (size_t index = 0; index < computes.size(); ++index)
-    if (incomingEdgeCount[index] == 0)
-      readyNodes.push_back(index);
-
-  size_t readyIndex = 0;
-  while (readyIndex != readyNodes.size()) {
-    size_t current = readyNodes[readyIndex++];
-    for (size_t child : children[current]) {
-      level[child] = std::max(level[child], level[current] + 1);
-      assert(incomingEdgeCount[child] > 0 && "incoming edge count underflow");
-      incomingEdgeCount[child]--;
-      if (incomingEdgeCount[child] == 0)
-        readyNodes.push_back(child);
-    }
-  }
-
-  SmallVector<uint64_t> weightedVmmNodesByLevel;
-  for (size_t index = 0; index < computes.size(); ++index) {
-    if (computeInfos[index].weightedVmmCount == 0)
-      continue;
-    if (weightedVmmNodesByLevel.size() <= level[index])
-      weightedVmmNodesByLevel.resize(level[index] + 1, 0);
-    weightedVmmNodesByLevel[level[index]]++;
-  }
-
-  uint64_t maxWeightedVmmLevel = 0;
-  uint64_t wideWeightedVmmLevels64 = 0;
-  uint64_t wideWeightedVmmLevels256 = 0;
-  for (uint64_t count : weightedVmmNodesByLevel) {
-    maxWeightedVmmLevel = std::max(maxWeightedVmmLevel, count);
-    wideWeightedVmmLevels64 += count >= 64;
-    wideWeightedVmmLevels256 += count >= 256;
-  }
-
-  using ShapeKey = std::tuple<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t>;
-  SmallVector<ShapeKey> weightedVmmShapeKeys;
-  for (auto [index, compute] : llvm::enumerate(computes)) {
-    const ComputeMotifInfo& info = computeInfos[index];
-    if (info.weightedVmmCount == 0)
-      continue;
-    weightedVmmShapeKeys.push_back({info.instructionCount,
-                                    info.weightedVmmCount,
-                                    static_cast<uint64_t>(compute.getWeights().size()),
-                                    static_cast<uint64_t>(compute.getInputs().size()),
-                                    static_cast<uint64_t>(parents[index].size()),
-                                    static_cast<uint64_t>(children[index].size())});
-  }
-
-  llvm::sort(weightedVmmShapeKeys);
-  SmallVector<std::pair<uint64_t, ShapeKey>> weightedVmmShapeCounts;
-  for (size_t index = 0; index < weightedVmmShapeKeys.size();) {
-    size_t next = index + 1;
-    while (next < weightedVmmShapeKeys.size() && weightedVmmShapeKeys[next] == weightedVmmShapeKeys[index])
-      next++;
-    weightedVmmShapeCounts.push_back({next - index, weightedVmmShapeKeys[index]});
-    index = next;
-  }
-  llvm::sort(weightedVmmShapeCounts, [](const auto& lhs, const auto& rhs) {
-    if (lhs.first != rhs.first)
-      return lhs.first > rhs.first;
-    return lhs.second < rhs.second;
-  });
-
-  llvm::errs() << llvm::formatv("[DCP-MOTIF] computes={0} edges={1} wvmmNodes={2} wvmmOps={3} "
-                                "serialChains={4} serialChainNodes={5} maxSerialChain={6} "
-                                "maxFanIn={7} maxFanOut={8} fanIn>=16/64/256={9}/{10}/{11} "
-                                "fanOut>=16/64/256={12}/{13}/{14} topoVisited={15}\n",
-                                computes.size(),
-                                edgeCount,
-                                weightedVmmNodeCount,
-                                weightedVmmOpCount,
-                                serialChainCount,
-                                serialChainNodeCount,
-                                maxSerialChain,
-                                maxFanIn,
-                                maxFanOut,
-                                fanIn16,
-                                fanIn64,
-                                fanIn256,
-                                fanOut16,
-                                fanOut64,
-                                fanOut256,
-                                readyNodes.size());
-
-  llvm::errs() << llvm::formatv("[DCP-MOTIF] wvmmLevels={0} maxWvmmLevel={1} wideWvmmLevels>=64/256={2}/{3} "
-                                "shapeGroups={4}\n",
-                                weightedVmmNodesByLevel.size(),
-                                maxWeightedVmmLevel,
-                                wideWeightedVmmLevels64,
-                                wideWeightedVmmLevels256,
-                                weightedVmmShapeCounts.size());
-
-  for (size_t rank = 0, end = std::min<size_t>(weightedVmmShapeCounts.size(), 5); rank < end; ++rank) {
-    auto [count, shape] = weightedVmmShapeCounts[rank];
-    auto [insts, vmmOps, weights, inputs, fanIn, fanOut] = shape;
-    llvm::errs() << llvm::formatv("[DCP-MOTIF] wvmmShape rank={0} count={1} insts={2} vmmOps={3} "
-                                  "weights={4} inputs={5} fanIn={6} fanOut={7}\n",
-                                  rank,
-                                  count,
-                                  insts,
-                                  vmmOps,
-                                  weights,
-                                  inputs,
-                                  fanIn,
-                                  fanOut);
-  }
-}
-
 void generateReport(func::FuncOp funcOp, const std::string& name, size_t usedCpuCount = 0) {
   std::fstream file = openReportFile(name);
   if (!file.is_open())
@@ -628,44 +328,27 @@ public:
 
   void runOnOperation() override {
     func::FuncOp func = getOperation();
-    {
-      ScopedMergePhaseTimer timer("trivial-serial-merge");
-      mergeTriviallyConnectedComputes(func);
-    }
-    if (std::getenv("DCP_MOTIF_PROFILE"))
-      emitMotifProfile(func);
+    mergeTriviallyConnectedComputes(func);
 
     const spatial::MergeScheduleResult* analysisResult = nullptr;
-    {
-      ScopedMergePhaseTimer timer("scheduling-analysis");
-      analysisResult = &getAnalysis<spatial::MergeSchedulingAnalysis>().getResult();
-    }
-    {
-      ScopedMergePhaseTimer timer("schedule-materialization");
-      if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) {
-        signalPassFailure();
-        return;
-      }
+    analysisResult = &getAnalysis<spatial::MergeSchedulingAnalysis>().getResult();
+    if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) {
+      signalPassFailure();
+      return;
     }
 
-    emitMergeIrCounts("after-materialization", func);
-
-    {
-      ScopedMergePhaseTimer timer("cleanup-topological-sort-report");
-      if (!sortTopologically(&func.getBody().front())) {
-        func.emitOpError("failed to topologically order merged Spatial IR");
-        signalPassFailure();
-        return;
-      }
-      if (failed(verifySpatialCommunicationInvariants(func))) {
-        func.emitOpError("merged Spatial communication invariant verification failed");
-        signalPassFailure();
-        return;
-      }
-      emitMergeIrCounts("final-post-merge", func);
-      dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial1_merged");
-      generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
+    if (!sortTopologically(&func.getBody().front())) {
+      func.emitOpError("failed to topologically order merged Spatial IR");
+      signalPassFailure();
+      return;
     }
+    if (failed(verifySpatialCommunicationInvariants(func))) {
+      func.emitOpError("merged Spatial communication invariant verification failed");
+      signalPassFailure();
+      return;
+    }
+    dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial1_merged");
+    generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
   }
 };
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -105,6 +105,28 @@ bool isProjectedBatchOffset(OpFoldResult offset, Value laneArg) {
       && succeeded(evaluateIndexLike(offset, bindings, /*lane=*/1, laneArg));
 }
 
+std::optional<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp extract) {
+  if (extract.getMixedOffsets().empty())
+    return std::nullopt;
+
+  OpFoldResult offset = extract.getMixedOffsets().front();
+  if (auto attr = llvm::dyn_cast<Attribute>(offset)) {
+    auto intAttr = dyn_cast<IntegerAttr>(attr);
+    if (!intAttr || intAttr.getInt() < 0)
+      return std::nullopt;
+    return static_cast<uint32_t>(intAttr.getInt());
+  }
+
+  Value offsetValue = llvm::cast<Value>(offset);
+  if (auto constantIndex = offsetValue.getDefiningOp<arith::ConstantIndexOp>()) {
+    if (constantIndex.value() < 0)
+      return std::nullopt;
+    return static_cast<uint32_t>(constantIndex.value());
+  }
+
+  return std::nullopt;
+}
+
 std::optional<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, Value input) {
   auto inputIt = llvm::find(batch.getInputs(), input);
   if (inputIt == batch.getInputs().end())
@@ -143,6 +165,101 @@ Cost getInputTransferCost(const ComputeInstance& consumerInstance, Value input)
   return static_cast<Cost>(getSizeInBytes(inputType));
 }
 
+uint32_t getLaneOverlapCount(const ComputeInstance& lhs, const ComputeInstance& rhs) {
+  uint32_t lhsEnd = lhs.laneStart + lhs.laneCount;
+  uint32_t rhsEnd = rhs.laneStart + rhs.laneCount;
+  return std::max(lhs.laneStart, rhs.laneStart) < std::min(lhsEnd, rhsEnd)
+         ? std::min(lhsEnd, rhsEnd) - std::max(lhs.laneStart, rhs.laneStart)
+         : 0;
+}
+
+Cost scaleTransferCostByLaneCount(Cost totalCost, uint32_t totalLaneCount, uint32_t fragmentLaneCount) {
+  assert(totalLaneCount > 0 && "laneCount must be positive");
+  assert(fragmentLaneCount > 0 && "fragmentLaneCount must be positive");
+  if (fragmentLaneCount >= totalLaneCount)
+    return totalCost;
+  return checkedMultiply(totalCost, static_cast<Cost>(fragmentLaneCount)) / static_cast<Cost>(totalLaneCount);
+}
+
+SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const ComputeInstance& consumerInstance) {
+  SmallVector<ProducerValueRef, 4> producers;
+  Operation* op = value.getDefiningOp();
+  if (!op)
+    return producers;
+
+  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
+    Value source = extract.getSource();
+    auto batch = dyn_cast_or_null<SpatComputeBatch>(source.getDefiningOp());
+    if (batch && batch.getNumResults() != 0) {
+      if (std::optional<uint32_t> lane = getConstantExtractLane(extract)) {
+        ComputeInstance instance = getBatchChunkForLane(batch, *lane);
+        producers.push_back({instance, 0});
+        return producers;
+      }
+
+      if (isa<SpatComputeBatch>(consumerInstance.op))
+        for (ComputeInstance instance :
+             getBatchChunksForRange(batch, consumerInstance.laneStart, consumerInstance.laneCount))
+          producers.push_back({instance, 0});
+      else
+        for (ComputeInstance instance : getBatchChunksForRange(batch, 0, static_cast<uint32_t>(batch.getLaneCount())))
+          producers.push_back({instance, 0});
+      return producers;
+    }
+
+    value = source;
+    op = value.getDefiningOp();
+    if (!op)
+      return producers;
+  }
+
+  if (auto compute = dyn_cast<SpatCompute>(op)) {
+    producers.push_back({
+      ComputeInstance {compute.getOperation(), 0, 1},
+       static_cast<size_t>(cast<OpResult>(value).getResultNumber())
+    });
+    return producers;
+  }
+
+  if (auto batch = dyn_cast<SpatComputeBatch>(op)) {
+    if (batch.getNumResults() != 0) {
+      uint32_t laneStart = isa<SpatComputeBatch>(consumerInstance.op) ? consumerInstance.laneStart : 0;
+      uint32_t laneCount = isa<SpatComputeBatch>(consumerInstance.op) ? consumerInstance.laneCount
+                                                                      : static_cast<uint32_t>(batch.getLaneCount());
+      for (ComputeInstance instance : getBatchChunksForRange(batch, laneStart, laneCount))
+        producers.push_back({instance, 0});
+      return producers;
+    }
+
+    uint32_t lane = cast<OpResult>(value).getResultNumber();
+    ComputeInstance instance = getBatchChunkForLane(batch, lane);
+    producers.push_back({instance, lane - instance.laneStart});
+    return producers;
+  }
+
+  return producers;
+}
+
+Cost getProducerTransferCost(Value input,
+                             const ComputeInstance& consumerInstance,
+                             const ProducerValueRef& producerRef) {
+  Cost transferCost = getInputTransferCost(consumerInstance, input);
+  auto producerBatch = dyn_cast<SpatComputeBatch>(producerRef.instance.op);
+  if (!producerBatch || producerBatch.getNumResults() == 0)
+    return transferCost;
+
+  if (auto consumerBatch = dyn_cast<SpatComputeBatch>(consumerInstance.op)) {
+    if (std::optional<Cost> projectedCost = getBatchProjectedInputTransferCost(consumerBatch, input)) {
+      uint32_t overlapLaneCount = getLaneOverlapCount(consumerInstance, producerRef.instance);
+      assert(overlapLaneCount > 0 && "projected batch edge must overlap consumer lanes");
+      return checkedMultiply(*projectedCost, static_cast<Cost>(overlapLaneCount));
+    }
+  }
+
+  return scaleTransferCostByLaneCount(
+    transferCost, static_cast<uint32_t>(producerBatch.getLaneCount()), producerRef.instance.laneCount);
+}
+
 static CrossbarWeight getOpaqueCrossbarWeight(Value value, std::optional<uint32_t> lane) {
   CrossbarWeight weight;
   weight.opaqueValue = value;
@@ -458,25 +575,13 @@ ComputeGraph buildComputeGraph(Operation* entryOp) {
   for (const auto& [targetIndex, node] : llvm::enumerate(graph.nodes)) {
     llvm::SmallVector<Value, 4> inputs = getComputeInstanceInputs(node.instance);
     for (Value input : inputs) {
-      Cost transferCost = getInputTransferCost(node.instance, input);
-      if (auto producerBatch = dyn_cast_or_null<SpatComputeBatch>(input.getDefiningOp());
-          producerBatch && producerBatch.getNumResults() != 0 && !isa<SpatComputeBatch>(node.instance.op)) {
-        for (uint32_t lane = 0; lane < static_cast<uint32_t>(producerBatch.getLaneCount()); ++lane) {
-          auto producerIt = graph.instanceToIndex.find(getBatchChunkForLane(producerBatch, lane));
-          if (producerIt == graph.instanceToIndex.end())
-            continue;
-          rawEdges.push_back({producerIt->second, targetIndex, transferCost});
-        }
-        continue;
+      for (const ProducerValueRef& producerRef : collectProducerValueRefs(input, node.instance)) {
+        auto producerIt = graph.instanceToIndex.find(producerRef.instance);
+        if (producerIt == graph.instanceToIndex.end())
+          continue;
+        rawEdges.push_back(
+          {producerIt->second, targetIndex, getProducerTransferCost(input, node.instance, producerRef)});
       }
-
-      auto producerInstance = getComputeProducerInstance(input, &node.instance);
-      if (!producerInstance)
-        continue;
-      auto producerIt = graph.instanceToIndex.find(*producerInstance);
-      if (producerIt == graph.instanceToIndex.end())
-        continue;
-      rawEdges.push_back({producerIt->second, targetIndex, transferCost});
     }
   }
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp

@@ -20,17 +20,66 @@ size_t getSchedulingCpuBudget() {
 
 size_t getBatchChunkTargetCount(int32_t laneCount) {
   assert(laneCount > 0 && "laneCount must be positive");
-  return static_cast<size_t>(laneCount);
+  return std::min(static_cast<size_t>(laneCount), getSchedulingCpuBudget());
+}
+
+BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkIndex) {
+  assert(laneCount > 0 && "laneCount must be positive");
+  size_t chunkCount = getBatchChunkTargetCount(laneCount);
+  assert(chunkIndex < chunkCount && "chunkIndex out of range");
+
+  size_t laneCountSize = static_cast<size_t>(laneCount);
+  size_t baseChunkSize = laneCountSize / chunkCount;
+  size_t remainder = laneCountSize % chunkCount;
+  size_t extraBefore = std::min(chunkIndex, remainder);
+  size_t start = chunkIndex * baseChunkSize + extraBefore;
+  size_t count = baseChunkSize + (chunkIndex < remainder ? 1 : 0);
+  assert(count > 0 && "chunk size must be positive");
+  return {static_cast<uint32_t>(start), static_cast<uint32_t>(count)};
+}
+
+size_t getBatchChunkIndexForLane(int32_t laneCount, uint32_t lane) {
+  assert(laneCount > 0 && "laneCount must be positive");
+  assert(lane < static_cast<uint32_t>(laneCount) && "lane out of range");
+
+  size_t chunkCount = getBatchChunkTargetCount(laneCount);
+  size_t laneCountSize = static_cast<size_t>(laneCount);
+  size_t baseChunkSize = laneCountSize / chunkCount;
+  size_t remainder = laneCountSize % chunkCount;
+  size_t largeChunkSize = baseChunkSize + 1;
+  size_t laneIndex = static_cast<size_t>(lane);
+  size_t largerChunkLanes = remainder * largeChunkSize;
+
+  if (laneIndex < largerChunkLanes)
+    return laneIndex / largeChunkSize;
+  return remainder + ((laneIndex - largerChunkLanes) / baseChunkSize);
 }
 
 ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex) {
-  assert(chunkIndex < static_cast<size_t>(batch.getLaneCount()) && "chunkIndex out of range");
-  return {batch.getOperation(), static_cast<uint32_t>(chunkIndex), 1};
+  BatchChunkRange chunk = getBatchChunkRange(batch.getLaneCount(), chunkIndex);
+  return {batch.getOperation(), chunk.laneStart, chunk.laneCount};
 }
 
 ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) {
-  assert(lane < static_cast<uint32_t>(batch.getLaneCount()) && "lane out of range");
-  return {batch.getOperation(), lane, 1};
+  return getBatchChunkForIndex(batch, getBatchChunkIndexForLane(batch.getLaneCount(), lane));
+}
+
+llvm::SmallVector<ComputeInstance, 4>
+getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount) {
+  llvm::SmallVector<ComputeInstance, 4> chunks;
+  if (laneCount == 0)
+    return chunks;
+
+  uint32_t laneEnd = laneStart + laneCount;
+  assert(laneEnd >= laneStart && "lane range overflow");
+  assert(laneEnd <= static_cast<uint32_t>(batch.getLaneCount()) && "lane range out of bounds");
+
+  size_t firstChunk = getBatchChunkIndexForLane(batch.getLaneCount(), laneStart);
+  size_t lastChunk = getBatchChunkIndexForLane(batch.getLaneCount(), laneEnd - 1);
+  chunks.reserve(lastChunk - firstChunk + 1);
+  for (size_t chunkIndex = firstChunk; chunkIndex <= lastChunk; ++chunkIndex)
+    chunks.push_back(getBatchChunkForIndex(batch, chunkIndex));
+  return chunks;
 }
 
 static std::optional<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp extract) {

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp

@@ -21,10 +21,19 @@ struct ProducerValueRef {
   size_t resultIndex = 0;
 };
 
+struct BatchChunkRange {
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 0;
+};
+
 size_t getSchedulingCpuBudget();
 size_t getBatchChunkTargetCount(int32_t laneCount);
+BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkIndex);
+size_t getBatchChunkIndexForLane(int32_t laneCount, uint32_t lane);
 ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex);
 ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane);
+llvm::SmallVector<ComputeInstance, 4>
+getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount);
 
 std::optional<ProducerValueRef> getProducerValueRef(mlir::Value value,
                                                     const ComputeInstance* consumerInstance = nullptr);

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -1,11 +1,9 @@
 #include "mlir/IR/Threading.h"
 
-#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormatVariadic.h"
 
-#include <functional>
 #include <limits>
 #include <queue>
 #include <vector>

src/PIM/Pass/PIMPasses.h

@@ -21,8 +21,6 @@ std::unique_ptr<mlir::Pass> createMergeComputeNodesPass();
 
 std::unique_ptr<mlir::Pass> createPimHostConstantFoldingPass();
 
-std::unique_ptr<mlir::Pass> createPimMaterializeHostConstantsPass();
-
 std::unique_ptr<mlir::Pass> createPimVerificationPass();
 
 std::unique_ptr<mlir::Pass> createEmitPimCodePass();

src/PIM/PimAccelerator.cpp

@@ -78,7 +78,6 @@ void PimAccelerator::registerPasses(int optLevel) const {
   registerPass(createPimMemoryCoalescingPass);
   registerPass(createMergeComputeNodesPass);
   registerPass(createPimHostConstantFoldingPass);
-  registerPass(createPimMaterializeHostConstantsPass);
   registerPass(createPimVerificationPass);
   registerPass(createEmitPimCodePass);
 }

test/PIM/CMakeLists.txt

@@ -30,3 +30,9 @@ add_pim_unittest(LabeledListTest
   OMPimCommon
 )
 
+add_pim_unittest(PimMemoryLivenessPlannerTest
+  PimMemoryLivenessPlannerTest.cpp
+
+  LINK_LIBS PRIVATE
+  OMPimCompilerUtils
+)

test/PIM/PimMemoryLivenessPlannerTest.cpp

@@ -0,0 +1,86 @@
+#include <cassert>
+#include <cstdlib>
+#include <iostream>
+
+#include "src/Accelerators/PIM/Compiler/PimMemoryLiveness.hpp"
+
+using onnx_mlir::LocalAllocInterval;
+using onnx_mlir::planPhysicalSlots;
+
+namespace {
+
+LocalAllocInterval makeInterval(size_t id, size_t size, uint64_t start, uint64_t end) {
+  LocalAllocInterval interval;
+  interval.id = id;
+  interval.size = size;
+  interval.start = start;
+  interval.end = end;
+  return interval;
+}
+
+void assertSingleSlotCase(LocalAllocInterval a, LocalAllocInterval b, size_t expectedSlotSize) {
+  llvm::SmallVector<LocalAllocInterval, 4> intervals = {a, b};
+  auto slots = planPhysicalSlots(intervals);
+  assert(slots.size() == 1);
+  assert(slots.front().requiredSize == expectedSlotSize);
+  assert(intervals[0].physicalSlotId == intervals[1].physicalSlotId);
+}
+
+int testSameSizeNonOverlap() {
+  std::cout << "testSameSizeNonOverlap:" << std::endl;
+  assertSingleSlotCase(makeInterval(0, 64, 0, 10), makeInterval(1, 64, 11, 20), 64);
+  return 0;
+}
+
+int testLargerFirst() {
+  std::cout << "testLargerFirst:" << std::endl;
+  assertSingleSlotCase(makeInterval(0, 100, 0, 10), makeInterval(1, 40, 11, 20), 100);
+  return 0;
+}
+
+int testSmallerFirst() {
+  std::cout << "testSmallerFirst:" << std::endl;
+  assertSingleSlotCase(makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), 100);
+  return 0;
+}
+
+int testOverlapNeedsTwoSlots() {
+  std::cout << "testOverlapNeedsTwoSlots:" << std::endl;
+  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+    makeInterval(0, 100, 0, 20), makeInterval(1, 40, 10, 30)};
+  auto slots = planPhysicalSlots(intervals);
+  assert(slots.size() == 2);
+  assert(intervals[0].physicalSlotId != intervals[1].physicalSlotId);
+  return 0;
+}
+
+int testReuseChain() {
+  std::cout << "testReuseChain:" << std::endl;
+  llvm::SmallVector<LocalAllocInterval, 4> intervals = {
+    makeInterval(0, 40, 0, 10), makeInterval(1, 100, 11, 20), makeInterval(2, 20, 21, 30)};
+  auto slots = planPhysicalSlots(intervals);
+  assert(slots.size() == 1);
+  assert(slots.front().requiredSize == 100);
+  assert(intervals[0].physicalSlotId == intervals[1].physicalSlotId);
+  assert(intervals[1].physicalSlotId == intervals[2].physicalSlotId);
+  return 0;
+}
+
+} // namespace
+
+int main(int argc, char *argv[]) {
+  (void) argc;
+  (void) argv;
+
+  int failures = 0;
+  failures += testSameSizeNonOverlap();
+  failures += testLargerFirst();
+  failures += testSmallerFirst();
+  failures += testOverlapNeedsTwoSlots();
+  failures += testReuseChain();
+  if (failures != 0) {
+    std::cerr << failures << " test failures\n";
+    return EXIT_FAILURE;
+  }
+  return EXIT_SUCCESS;
+}

validation/.gitignore

@@ -8,3 +8,7 @@ networks/**/outputs
 networks/**/raptor
 networks/**/runner
 networks/**/simulation
+networks/**/real_image_val
+networks/**/*.png
+networks/**/*.jpg
+networks/**/*.csv

validation/gen_network_runner.py

@@ -199,7 +199,10 @@ int main(int argc, char **argv) {{
 
   // ---- Cleanup ----
   omTensorListDestroy(in_list);
-  omTensorListDestroy(out_list);
+  // Some debug-heavy models return aliased outputs. This runner is a short-
+  // lived process, so destroy only the list wrapper and let process exit
+  // reclaim the output tensors safely.
+  omTensorListDestroyShallow(out_list);
   return 0;
 }}
 """