add support for softmax, resize, split, gather

src/PIM/Pass/Pim/ConstantFolding/Patterns/Subview.cpp

@@ -146,6 +146,37 @@ struct RewriteHostSubviewLoadPattern final : OpRewritePattern<pim::PimMemCopyHos
   }
 };
 
+struct RewriteHostSubviewStorePattern final : OpRewritePattern<pim::PimMemCopyDevToHostOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(pim::PimMemCopyDevToHostOp copyOp, PatternRewriter& rewriter) const override {
+    auto status = rewriteSubviewCopyLikeOp(
+      copyOp,
+      copyOp.getHostTarget(),
+      copyOp.getDeviceSource(),
+      copyOp.getHostTargetOffset(),
+      copyOp.getDeviceSourceOffset(),
+      copyOp.getSize(),
+      rewriter,
+      [&](
+        MemRefType resultType, Value dst, Value src, int64_t dstByteOffset, int64_t srcByteOffset, int64_t sliceBytes) {
+        pim::PimMemCopyDevToHostOp::create(rewriter,
+                                           copyOp.getLoc(),
+                                           resultType,
+                                           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))
+      return failure();
+
+    rewriter.replaceOp(copyOp, copyOp.getHostTarget());
+    return success();
+  }
+};
+
 struct FoldConstantCoreSubviewPattern final : OpRewritePattern<memref::SubViewOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -209,7 +240,10 @@ struct FoldConstantCoreSubviewPattern final : OpRewritePattern<memref::SubViewOp
 } // namespace
 
 void populateConstantFoldingSubviewPatterns(RewritePatternSet& patterns) {
-  patterns.add<RewriteCoreSubviewCopyPattern, RewriteHostSubviewLoadPattern, FoldConstantCoreSubviewPattern>(
+  patterns.add<RewriteCoreSubviewCopyPattern,
+               RewriteHostSubviewLoadPattern,
+               RewriteHostSubviewStorePattern,
+               FoldConstantCoreSubviewPattern>(
     patterns.getContext());
 }