huge refactor for high RewritePatterns usage and less ad-hoc cpp code

remove Spatial many ops in favor of tensor ops like in pim

src/PIM/Pass/PimCodegen/MaterializeHostConstantsPass.cpp

@@ -1,5 +1,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/Builders.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
@@ -9,6 +10,8 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/MathExtras.h"
 
+#include <type_traits>
+
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -21,6 +24,8 @@ namespace {
 static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
   if (isa<pim::PimMemCopyHostToDevOp>(op))
     return operandIndex == 1;
+  if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
+    return operandIndex == 1;
   if (isa<pim::PimMemCopyDevToHostOp>(op))
     return operandIndex == 0;
   return false;
@@ -33,6 +38,91 @@ static int64_t getValueSizeInBytes(Value value) {
   return type.getNumElements() * type.getElementTypeBitWidth() / 8;
 }
 
+template <typename CoreOpTy>
+static void materializeHostConstantsInCore(CoreOpTy coreOp, IRRewriter& rewriter, bool& hasFailure) {
+  DenseMap<Value, DenseMap<int64_t, DenseMap<Type, Value>>> materializedValues;
+  SmallVector<Operation*> ops;
+  coreOp.getBody().front().walk([&](Operation* op) {
+    if (!isa<pim::PimHaltOp, scf::YieldOp>(op))
+      ops.push_back(op);
+  });
+
+  for (Operation* op : ops) {
+    for (OpOperand& operand : op->getOpOperands()) {
+      Value originalValue = operand.get();
+      if (!isa<BaseMemRefType>(originalValue.getType()) || isExplicitHostOperand(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()) {
+        operand.set(cachedValue->second);
+        continue;
+      }
+
+      int64_t totalBytes = getValueSizeInBytes(originalValue);
+      if (totalBytes < 0 || !llvm::isInt<32>(totalBytes) || !llvm::isInt<32>(resolvedAddress->byteOffset)) {
+        op->emitOpError("host constant materialization requires 32-bit copy sizes and offsets");
+        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 copiedValue;
+      if constexpr (std::is_same_v<CoreOpTy, pim::PimCoreBatchOp>) {
+        copiedValue = pim::PimMemCopyHostToDevBatchOp::create(
+                        rewriter,
+                        op->getLoc(),
+                        originalType,
+                        deviceDst,
+                        getGlobalOp.getResult(),
+                        rewriter.getI32IntegerAttr(0),
+                        rewriter.getI32IntegerAttr(static_cast<int32_t>(resolvedAddress->byteOffset)),
+                        rewriter.getI32IntegerAttr(static_cast<int32_t>(totalBytes)))
+                        .getOutput();
+      }
+      else {
+        copiedValue = pim::PimMemCopyHostToDevOp::create(
+                        rewriter,
+                        op->getLoc(),
+                        originalType,
+                        deviceDst,
+                        getGlobalOp.getResult(),
+                        rewriter.getI32IntegerAttr(0),
+                        rewriter.getI32IntegerAttr(static_cast<int32_t>(resolvedAddress->byteOffset)),
+                        rewriter.getI32IntegerAttr(static_cast<int32_t>(totalBytes)))
+                        .getOutput();
+      }
+
+      cachedByType[originalType] = copiedValue;
+      operand.set(copiedValue);
+    }
+  }
+}
+
 struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MaterializeHostConstantsPass)
 
@@ -50,71 +140,11 @@ struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass,
       if (funcOp.isExternal())
         continue;
 
-      for (pim::PimCoreOp coreOp : funcOp.getOps<pim::PimCoreOp>()) {
-        DenseMap<Value, DenseMap<int64_t, DenseMap<Type, Value>>> materializedValues;
+      for (pim::PimCoreOp coreOp : funcOp.getOps<pim::PimCoreOp>())
+        materializeHostConstantsInCore(coreOp, rewriter, hasFailure);
 
-        for (Operation& op : llvm::make_early_inc_range(coreOp.getBody().front())) {
-          if (isa<pim::PimHaltOp>(op))
-            continue;
-
-          for (OpOperand& operand : op.getOpOperands()) {
-            Value originalValue = operand.get();
-            if (!isa<BaseMemRefType>(originalValue.getType()) || isExplicitHostOperand(&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()) {
-              operand.set(cachedValue->second);
-              continue;
-            }
-
-            int64_t totalBytes = getValueSizeInBytes(originalValue);
-            if (totalBytes < 0 || !llvm::isInt<32>(totalBytes) || !llvm::isInt<32>(resolvedAddress->byteOffset)) {
-              op.emitOpError("host constant materialization requires 32-bit copy sizes and offsets");
-              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);
-
-            auto hostToDevCopy = pim::PimMemCopyHostToDevOp::create(
-              rewriter,
-              op.getLoc(),
-              originalType,
-              deviceDst,
-              getGlobalOp.getResult(),
-              rewriter.getI32IntegerAttr(0),
-              rewriter.getI32IntegerAttr(static_cast<int32_t>(resolvedAddress->byteOffset)),
-              rewriter.getI32IntegerAttr(static_cast<int32_t>(totalBytes)));
-
-            cachedByType[originalType] = hostToDevCopy.getResult();
-            operand.set(hostToDevCopy.getResult());
-          }
-        }
-      }
+      for (pim::PimCoreBatchOp coreBatchOp : funcOp.getOps<pim::PimCoreBatchOp>())
+        materializeHostConstantsInCore(coreBatchOp, rewriter, hasFailure);
 
       SmallVector<Operation*> hostCompactOps;
       for (Operation& op : funcOp.getBody().front())