refactorone

src/PIM/Pass/PimCodegen/HostConstantFolding/Patterns/Subview.cpp

@@ -268,24 +268,31 @@ struct RewriteHostSubviewLoadPattern final : OpRewritePattern<pim::PimMemCopyHos
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(pim::PimMemCopyHostToDevOp copyOp, PatternRewriter& rewriter) const override {
+    auto dstOffset = resolveIndexValue(copyOp.getDeviceTargetOffset());
+    auto srcOffset = resolveIndexValue(copyOp.getHostSourceOffset());
+    if (failed(dstOffset) || failed(srcOffset))
+      return failure();
+
     auto status = rewriteSubviewCopyLikeOp(
       copyOp,
       copyOp.getDeviceTarget(),
       copyOp.getHostSource(),
-      copyOp.getDeviceTargetOffset(),
-      copyOp.getHostSourceOffset(),
+      *dstOffset,
+      *srcOffset,
       copyOp.getSize(),
       /*allowLoopRewrite=*/true,
       rewriter,
       [&](
         MemRefType resultType, Value dst, Value src, int64_t dstByteOffset, int64_t srcByteOffset, int64_t sliceBytes) {
+        Value dstOffsetValue = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), dstByteOffset);
+        Value srcOffsetValue = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), srcByteOffset);
         pim::PimMemCopyHostToDevOp::create(rewriter,
                                            copyOp.getLoc(),
                                            resultType,
+                                           dstOffsetValue,
+                                           srcOffsetValue,
                                            dst,
                                            src,
-                                           rewriter.getI32IntegerAttr(static_cast<int32_t>(dstByteOffset)),
-                                           rewriter.getI32IntegerAttr(static_cast<int32_t>(srcByteOffset)),
                                            rewriter.getI32IntegerAttr(static_cast<int32_t>(sliceBytes)));
       });
     if (failed(status))
@@ -301,24 +308,31 @@ struct RewriteHostSubviewStorePattern final : OpRewritePattern<pim::PimMemCopyDe
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(pim::PimMemCopyDevToHostOp copyOp, PatternRewriter& rewriter) const override {
+    auto dstOffset = resolveIndexValue(copyOp.getHostTargetOffset());
+    auto srcOffset = resolveIndexValue(copyOp.getDeviceSourceOffset());
+    if (failed(dstOffset) || failed(srcOffset))
+      return failure();
+
     auto status = rewriteSubviewCopyLikeOp(
       copyOp,
       copyOp.getHostTarget(),
       copyOp.getDeviceSource(),
-      copyOp.getHostTargetOffset(),
-      copyOp.getDeviceSourceOffset(),
+      *dstOffset,
+      *srcOffset,
       copyOp.getSize(),
       /*allowLoopRewrite=*/false,
       rewriter,
       [&](
         MemRefType resultType, Value dst, Value src, int64_t dstByteOffset, int64_t srcByteOffset, int64_t sliceBytes) {
+        Value dstOffset = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), dstByteOffset);
+        Value srcOffset = arith::ConstantIndexOp::create(rewriter, copyOp.getLoc(), srcByteOffset);
         pim::PimMemCopyDevToHostOp::create(rewriter,
                                            copyOp.getLoc(),
                                            resultType,
+                                           dstOffset,
+                                           srcOffset,
                                            dst,
                                            src,
-                                           rewriter.getI32IntegerAttr(static_cast<int32_t>(dstByteOffset)),
-                                           rewriter.getI32IntegerAttr(static_cast<int32_t>(srcByteOffset)),
                                            rewriter.getI32IntegerAttr(static_cast<int32_t>(sliceBytes)));
       });
     if (failed(status))
@@ -355,9 +369,7 @@ struct FoldConstantCoreSubviewPattern final : OpRewritePattern<memref::SubViewOp
     if (failed(staticOffsets))
       return failure();
 
-    auto elementType = cast<MemRefType>(subviewOp.getType()).getElementType();
-    auto resultMemRefType =
-      MemRefType::get(SmallVector<int64_t>(subviewInfo->sizes.begin(), subviewInfo->sizes.end()), elementType);
+    auto resultMemRefType = cast<MemRefType>(subviewOp.getType());
     auto foldedAttr = foldDenseSubview(*denseAttr, *staticOffsets, resultMemRefType.getShape());
     if (failed(foldedAttr))
       return failure();

src/PIM/Pass/PimCodegen/MaterializeHostConstantsPass.cpp

@@ -23,11 +23,11 @@ namespace {
 
 static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
   if (isa<pim::PimMemCopyHostToDevOp>(op))
-    return operandIndex == 1;
+    return operandIndex == 3;
   if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
     return operandIndex == 1;
   if (isa<pim::PimMemCopyDevToHostOp>(op))
-    return operandIndex == 0;
+    return operandIndex == 2;
   return false;
 }
 
@@ -39,7 +39,10 @@ static int64_t getValueSizeInBytes(Value value) {
 }
 
 template <typename CoreOpTy>
-static void materializeHostConstantsInCore(CoreOpTy coreOp, IRRewriter& rewriter, bool& hasFailure) {
+static void materializeHostConstantsInCore(CoreOpTy coreOp,
+                                           IRRewriter& rewriter,
+                                           OperationFolder& constantFolder,
+                                           bool& hasFailure) {
   DenseMap<Value, DenseMap<int64_t, DenseMap<Type, Value>>> materializedValues;
   SmallVector<Operation*> ops;
   coreOp.getBody().front().walk([&](Operation* op) {
@@ -48,6 +51,9 @@ static void materializeHostConstantsInCore(CoreOpTy coreOp, IRRewriter& rewriter
   });
 
   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()) || isExplicitHostOperand(op, operand.getOperandNumber()))
@@ -105,16 +111,17 @@ static void materializeHostConstantsInCore(CoreOpTy coreOp, IRRewriter& rewriter
                         .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();
+        copiedValue =
+          pim::PimMemCopyHostToDevOp::create(
+            rewriter,
+            op->getLoc(),
+            originalType,
+            getOrCreateHostIndexConstant(op, 0, constantFolder),
+            getOrCreateHostIndexConstant(op, static_cast<int64_t>(resolvedAddress->byteOffset), constantFolder),
+            deviceDst,
+            getGlobalOp.getResult(),
+            rewriter.getI32IntegerAttr(static_cast<int32_t>(totalBytes)))
+            .getOutput();
       }
 
       cachedByType[originalType] = copiedValue;
@@ -134,6 +141,7 @@ struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass,
   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>()) {
@@ -141,10 +149,10 @@ struct MaterializeHostConstantsPass : PassWrapper<MaterializeHostConstantsPass,
         continue;
 
       for (pim::PimCoreOp coreOp : funcOp.getOps<pim::PimCoreOp>())
-        materializeHostConstantsInCore(coreOp, rewriter, hasFailure);
+        materializeHostConstantsInCore(coreOp, rewriter, constantFolder, hasFailure);
 
       for (pim::PimCoreBatchOp coreBatchOp : funcOp.getOps<pim::PimCoreBatchOp>())
-        materializeHostConstantsInCore(coreBatchOp, rewriter, hasFailure);
+        materializeHostConstantsInCore(coreBatchOp, rewriter, constantFolder, hasFailure);
 
       SmallVector<Operation*> hostCompactOps;
       for (Operation& op : funcOp.getBody().front())

src/PIM/Pass/PimCodegen/VerificationPass.cpp

@@ -8,6 +8,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/SubviewUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Compiler/PimBatchEmission.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -119,11 +120,27 @@ static bool isConstantGlobalView(Value value) {
 
 static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
   if (isa<pim::PimMemCopyHostToDevOp>(op))
-    return operandIndex == 1;
+    return operandIndex == 3;
   if (isa<pim::PimMemCopyHostToDevBatchOp>(op))
     return operandIndex == 1;
   if (isa<pim::PimMemCopyDevToHostOp>(op))
-    return operandIndex == 0;
+    return operandIndex == 2;
+  return false;
+}
+
+static bool isCoreWeightBlockArgument(Value value) {
+  auto blockArgument = dyn_cast<BlockArgument>(value);
+  if (!blockArgument)
+    return false;
+
+  if (auto coreOp = dyn_cast<pim::PimCoreOp>(blockArgument.getOwner()->getParentOp()))
+    return static_cast<unsigned>(blockArgument.getArgNumber()) < coreOp.getWeights().size();
+
+  if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(blockArgument.getOwner()->getParentOp())) {
+    unsigned argNumber = static_cast<unsigned>(blockArgument.getArgNumber());
+    return argNumber > 0 && argNumber <= coreBatchOp.getWeights().size();
+  }
+
   return false;
 }
 
@@ -193,7 +210,9 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
 
         if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(&op)) {
           (void) verifyCoreWeights(moduleOp, coreBatchOp, diagnostics);
-          (void) verifyCoreOperands(coreBatchOp, diagnostics);
+          for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
+            (void) withScalarCoreFromBatchLane(
+              coreBatchOp, lane, [&](pim::PimCoreOp scalarCore) { return verifyCoreOperands(scalarCore, diagnostics); });
           continue;
         }
 
@@ -297,6 +316,9 @@ private:
           if (!isa<BaseMemRefType>(operand.getType()))
             continue;
 
+          if (isCoreWeightBlockArgument(operand))
+            continue;
+
           auto resolvedAddress = resolveContiguousAddress(operand, knowledge);
           if (failed(resolvedAddress)) {
             diagnostics.report(&op, [&](Operation* illegalOp) {
@@ -327,6 +349,26 @@ private:
             hasFailure = true;
           }
         }
+
+        if (auto storeOp = dyn_cast<pim::PimMemCopyDevToHostOp>(op)) {
+          if (failed(resolveIndexValue(storeOp.getHostTargetOffset(), knowledge))
+              || failed(resolveIndexValue(storeOp.getDeviceSourceOffset(), knowledge))) {
+            diagnostics.report(&op, [](Operation* illegalOp) {
+              illegalOp->emitOpError("offset operands must be statically evaluable for PIM codegen");
+            });
+            hasFailure = true;
+          }
+        }
+
+        if (auto loadOp = dyn_cast<pim::PimMemCopyHostToDevOp>(op)) {
+          if (failed(resolveIndexValue(loadOp.getDeviceTargetOffset(), knowledge))
+              || failed(resolveIndexValue(loadOp.getHostSourceOffset(), knowledge))) {
+            diagnostics.report(&op, [](Operation* illegalOp) {
+              illegalOp->emitOpError("offset operands must be statically evaluable for PIM codegen");
+            });
+            hasFailure = true;
+          }
+        }
         return success(!hasFailure);
       });
   }