[MLIR] Fix fold slice and copy int64_max bug

Author: AmosLewis

lib/Dialect/Torch/Transforms/RecomposeComplexOps.cpp

@@ -9,9 +9,12 @@
 
 #include "PassDetail.h"
 
+#include "mlir/IR/Matchers.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
+#include "torch-mlir/Dialect/Torch/IR/TorchTypes.h"
 #include "torch-mlir/Dialect/Torch/Transforms/Passes.h"
+#include <climits>
 
 using namespace mlir;
 using namespace mlir::torch;
@@ -33,7 +36,14 @@ class RecomposeSliceCopy_ : public OpRewritePattern<AtenCopy_Op> {
     if (!matchPattern(sliceOp.getDim(), m_TorchConstantInt(&dim)))
       return failure();
     int64_t end;
-    if (!matchPattern(sliceOp.getEnd(), m_TorchConstantInt(&end)))
+    if (sliceOp.getEnd().getType().isa<Torch::NoneType>())
+      end = INT64_MAX;
+    else if (!matchPattern(sliceOp.getEnd(), m_TorchConstantInt(&end)))
+      return failure();
+    int64_t start;
+    if (sliceOp.getStart().getType().isa<Torch::NoneType>())
+      start = INT64_MIN;
+    else if (!matchPattern(sliceOp.getStart(), m_TorchConstantInt(&start)))
       return failure();
 
     Value newEnd = sliceOp.getEnd();
@@ -42,6 +52,20 @@ class RecomposeSliceCopy_ : public OpRewritePattern<AtenCopy_Op> {
           op.getLoc(), sliceOp.getSelf(), sliceOp.getDim());
       newEnd =
           rewriter.create<AtenAddIntOp>(op.getLoc(), dimSize, sliceOp.getEnd());
+    } else if (end == INT64_MAX) {
+      newEnd = rewriter.create<AtenSizeIntOp>(op.getLoc(), sliceOp.getSelf(),
+                                              sliceOp.getDim());
+    }
+
+    Value newStart = sliceOp.getStart();
+    if (start == INT64_MIN) {
+      newStart = rewriter.create<ConstantIntOp>(op.getLoc(),
+                                                rewriter.getI64IntegerAttr(0));
+    } else if (start < 0) {
+      Value dimSize = rewriter.create<AtenSizeIntOp>(
+          op.getLoc(), sliceOp.getSelf(), sliceOp.getDim());
+      newStart =
+          rewriter.create<AtenAddIntOp>(op.getLoc(), dimSize, sliceOp.getEnd());
     }
 
     Value noneVal = rewriter.create<ConstantNoneOp>(op.getLoc());
@@ -50,7 +74,7 @@ class RecomposeSliceCopy_ : public OpRewritePattern<AtenCopy_Op> {
     // Create IndexPut_Op
     BaseTensorType tensorType = op->getResultTypes()[0].cast<BaseTensorType>();
     Value range = rewriter.create<AtenArangeStartStepOp>(
-        op.getLoc(), tensorType, sliceOp.getStart(), newEnd, sliceOp.getStep(),
+        op.getLoc(), tensorType, newStart, newEnd, sliceOp.getStep(),
         /*dtype=*/noneVal, /*layout=*/noneVal, /*device=*/noneVal,
         /*pin_memory=*/noneVal);