merge remote changes

src/PIM/Dialect/Pim/Pim.td

@@ -24,7 +24,7 @@ def PimTensor :
 // Execution
 //===----------------------------------------------------------------------===//
 
-def PimCoreOp : PimOp<"core", [SingleBlock]> {
+def PimCoreOp : PimOp<"core", [SingleBlock, IsolatedFromAbove]> {
   let summary = "Execute a block on a PIM core";
 
   let regions = (region SizedRegion<1>:$body);

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

@@ -3,12 +3,17 @@
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Threading.h"
 #include "mlir/Pass/Pass.h"
 
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
+
 #include "Common/PimCommon.hpp"
 #include "Compiler/PimCodeGen.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerOptions.hpp"
 
@@ -40,14 +45,44 @@ private:
 
 void PimBufferizationPass::runOnOperation() {
   auto moduleOp = getOperation();
+  // Refactor this into a function
+  {
+    auto funcOp = getPimEntryFunc(moduleOp);
 
-  // One-Shot-Bufferization
-  bufferization::OneShotBufferizationOptions options;
-  options.allowUnknownOps = true;
-  bufferization::BufferizationState state;
-  if (failed(bufferization::runOneShotBufferize(moduleOp, options, state))) {
-    moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
-    signalPassFailure();
+    auto coreOps = llvm::to_vector(funcOp->getOps<pim::PimCoreOp>());
+    MLIRContext* ctx = moduleOp.getContext();
+    // failableParallelForEach will run the lambda in parallel and stop if any thread fails
+    LogicalResult result = mlir::failableParallelForEach(ctx, coreOps, [&](pim::PimCoreOp coreOp) {
+      // Again, allocate state LOCALLY per thread/function
+      bufferization::OneShotBufferizationOptions options;
+      options.allowUnknownOps = true;
+      bufferization::BufferizationState state;
+      if (failed(bufferization::runOneShotBufferize(coreOp, options, state))) {
+        coreOp.emitError("Failed to bufferize PIM and Spatial ops");
+        return failure();
+      }
+      return success();
+    });
+
+    if (failed(result)) {
+      moduleOp.emitError("Failed to bufferize-parallel PIM and Spatial ops");
+      signalPassFailure();
+    }
+
+    funcOp->walk([&](bufferization::ToTensorOp toTensorOp) {
+      if (llvm::isa_and_present<pim::PimCoreOp>(toTensorOp->getParentOp()))
+        toTensorOp->setAttr("restrict", UnitAttr::get(ctx));
+    });
+
+    // One-Shot-Bufferization
+    bufferization::OneShotBufferizationOptions options;
+    options.allowUnknownOps = true;
+    bufferization::BufferizationState state;
+
+    if (failed(bufferization::runOneShotBufferize(moduleOp, options, state))) {
+      moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
+      signalPassFailure();
+    }
   }
 
   MLIRContext* ctx = moduleOp.getContext();
@@ -57,7 +92,18 @@ void PimBufferizationPass::runOnOperation() {
   RewritePatternSet patterns(ctx);
   populateWithGenerated(patterns);
 
-  if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
+  // Only convert memref.copy → pim.memcp inside pim.core / pim.core_batch bodies.
+  // Host-level copies (e.g. from split/slice ops) must remain as memref.copy for CPU lowering.
+  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+  bool hasFailed = false;
+  moduleOp.walk<WalkOrder::PreOrder>([&](Operation* op) {
+    if (!isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op))
+      return WalkResult::advance();
+    if (failed(applyPartialConversion(op, target, frozenPatterns)))
+      hasFailed = true;
+    return WalkResult::skip();
+  });
+  if (hasFailed) {
     signalPassFailure();
     return;
   }