[MLIR] Fold aten select and copy pattern

Author: AmosLewis

lib/Dialect/Torch/Transforms/RecomposeComplexOps.cpp

@@ -12,6 +12,7 @@
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
 #include "torch-mlir/Dialect/Torch/Transforms/Passes.h"
+#include "torch-mlir/Dialect/Torch/Utils/Utils.h"
 
 using namespace mlir;
 using namespace mlir::torch;
@@ -71,6 +72,59 @@ class RecomposeSliceCopy_ : public OpRewritePattern<AtenCopy_Op> {
     return success();
   }
 };
+
+class RecomposeSelectCopy_ : public OpRewritePattern<AtenCopy_Op> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenCopy_Op op,
+                                PatternRewriter &rewriter) const override {
+    if (!op.getSelf().getDefiningOp() ||
+        !isa<AtenSelectIntOp>(op.getSelf().getDefiningOp()))
+      return failure();
+    auto selectOp = cast<AtenSelectIntOp>(op.getSelf().getDefiningOp());
+
+    // Get input tensor type
+    Value self = selectOp.getSelf();
+    std::optional<unsigned> maybeRank = getTensorRank(self);
+    if (!maybeRank)
+      return rewriter.notifyMatchFailure(op, "Unimplemented: unranked tensor");
+    unsigned rank = *maybeRank;
+
+    // Get indices
+    int64_t dim;
+    if (!matchPattern(selectOp.getDim(), m_TorchConstantInt(&dim)))
+      return failure();
+    int64_t index;
+    if (!matchPattern(selectOp.getIndex(), m_TorchConstantInt(&index)))
+      return failure();
+
+    Value noneVal = rewriter.create<ConstantNoneOp>(op.getLoc());
+    Value falseVal = rewriter.create<ConstantBoolOp>(op.getLoc(), false);
+
+    // Create IndexPut_Op
+    BaseTensorType tensorType = op->getResultTypes()[0].cast<BaseTensorType>();
+    SmallVector<Value> indicesVector;
+    for (auto i = 0; i < rank; i++) {
+      if(i == dim){
+        indicesVector.push_back(selectOp.getIndex());
+      } else {
+        indicesVector.push_back(noneVal);
+      }
+    }
+
+    Value indices = rewriter.create<PrimListConstructOp>(
+        op.getLoc(),
+        Torch::ListType::get(op->getContext(),
+                             Torch::OptionalType::get(tensorType)),
+        indicesVector);
+
+    rewriter.replaceOpWithNewOp<Aten_IndexPutImpl_Op>(
+        op, op->getResultTypes(), selectOp.getSelf(), indices, op.getSrc(),
+        /*accumulate=*/falseVal, /*unsafe=*/falseVal);
+
+    return success();
+  }
+};
 } // namespace
 
 namespace {
@@ -83,6 +137,7 @@ class RecomposeComplexOpsPass
 
     // pattern.add calls go here
     patterns.add<RecomposeSliceCopy_>(context);
+    patterns.add<RecomposeSelectCopy_>(context);
 
     GreedyRewriteConfig config;
     config.useTopDownTraversal = true;