Fix Slice Folder OOB Crash and onnx.Shape lowering (llvm#3843)

1. Clamps OOB start index to 0 in slice folder 2. Adds a more descriptive `emitError` in slice folder if the creation of the `DenseElementsAttr` would fail due to a bad result shape. 3. Fixes the `onnx.Shape` lowering to default to `inputRank` for `end` instead of `-1`. When `end==-1` the last element was missing when slicing.
iree-org · Nov 1, 2024 · 3104b66 · 3104b66
1 parent 738d45d
commit 3104b66
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 4 deletions.
diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
@@ -1654,14 +1654,19 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
         Value operand;
         int64_t start, end;
         if (binder.tensorOperand(operand) ||
-            binder.tensorResultType(resultType) ||
-            binder.s64IntegerAttr(start, "start", 0) ||
-            binder.s64IntegerAttr(end, "end", -1))
+            binder.tensorResultType(resultType))
           return failure();
 
         auto inputType = dyn_cast<Torch::ValueTensorType>(operand.getType());
+        if (!inputType || !inputType.hasSizes())
+          return failure();
+
         int64_t inputRank = inputType.getSizes().size();
 
+        if (binder.s64IntegerAttr(start, "start", 0) ||
+            binder.s64IntegerAttr(end, "end", inputRank))
+          return failure();
+
         auto shapeType = Torch::ValueTensorType::get(
             binder.op->getContext(), SmallVector<int64_t>{inputRank},
             resultType.getOptionalDtype());
@@ -1674,7 +1679,7 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
           return success();
         }
 
-        if (start == 0 && end == -1) {
+        if (start == 0 && end == inputRank) {
           rewriter.replaceOp(binder.op, shape);
           return success();
         }

diff --git a/lib/Dialect/Torch/IR/TorchOps.cpp b/lib/Dialect/Torch/IR/TorchOps.cpp
@@ -3998,6 +3998,7 @@ OpFoldResult AtenSliceTensorOp::fold(FoldAdaptor adaptor) {
     int64_t limit = end.getValue().getSExtValue();
     int64_t stride = step.getValue().getSExtValue();
     begin = begin < 0 ? begin + inType.getSizes()[dimInt] : begin;
+    begin = std::max<int64_t>(begin, 0);
     limit = limit < 0 ? limit + inType.getSizes()[dimInt] : limit;
     limit = limit < 0 ? -1 : limit;
     limit = std::min(limit, inType.getSizes()[dimInt]);
@@ -4038,6 +4039,14 @@ OpFoldResult AtenSliceTensorOp::fold(FoldAdaptor adaptor) {
       }
     };
     recursiveIter(recursiveIter, 0, 0);
+    if (static_cast<int64_t>(values.size()) != count) {
+      emitError(
+          "Op has incorrect result shape for provided arguments.\nNum elements "
+          "present in slice: " +
+          std::to_string(values.size()) +
+          "\nNum elements implied by result type: " + std::to_string(count));
+      return nullptr;
+    }
     return DenseElementsAttr::get(outType.toBuiltinTensor(), values);
   }
 

diff --git a/test/Dialect/Torch/canonicalize.mlir b/test/Dialect/Torch/canonicalize.mlir
@@ -2330,6 +2330,32 @@ func.func @torch.aten.slice.tensor$fold_full_slice(%arg0: !torch.vtensor<[?],f32
   return %0 : !torch.vtensor<[?],f32>
 }
 
+//  CHECK-LABEL:    @torch.aten.slice.tensor$fold_oob_start
+//        CHECK:        %[[LIT:.*]] = torch.vtensor.literal(dense<[0, 1, 2]> : tensor<3xsi64>) : !torch.vtensor<[3],si64>
+//        CHECK:        return %[[LIT]] : !torch.vtensor<[3],si64>
+func.func @torch.aten.slice.tensor$fold_oob_start() -> !torch.vtensor<[3],si64> {
+  %0 = torch.vtensor.literal(dense<[0,1,2,3]> : tensor<4xsi64>) : !torch.vtensor<[4],si64>
+  %int1 = torch.constant.int 1
+  %int-1  = torch.constant.int -1
+  %int-10 = torch.constant.int -10
+  %int0 = torch.constant.int 0
+  %1 = torch.aten.slice.Tensor %0, %int0, %int-10, %int-1, %int1 : !torch.vtensor<[4], si64>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[3], si64>
+  return %1 : !torch.vtensor<[3],si64>
+}
+
+//  CHECK-LABEL:    @torch.aten.slice.tensor$nofold_invalid_shape
+//        CHECK:    %[[SLICE:.*]] = torch.aten.slice.Tensor
+//        CHECK:    return %[[SLICE]]
+func.func @torch.aten.slice.tensor$nofold_invalid_shape() -> !torch.vtensor<[4],si64> {
+  %0 = torch.vtensor.literal(dense<[0,1,2,3]> : tensor<4xsi64>) : !torch.vtensor<[4],si64>
+  %int1 = torch.constant.int 1
+  %int-1  = torch.constant.int -1
+  %int-10 = torch.constant.int -10
+  %int0 = torch.constant.int 0
+  %1 = torch.aten.slice.Tensor %0, %int0, %int-10, %int-1, %int1 : !torch.vtensor<[4], si64>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[4], si64>
+  return %1 : !torch.vtensor<[4],si64>
+}
+
 //  CHECK-LABEL:    @torch.aten.slice.tensor$no_fold_step
 //        CHECK: torch.aten.slice.Tensor
 func.func @torch.aten.slice.tensor$no_fold_step(%arg0: !torch.vtensor<[?],f32>, %dim: !torch.int) -> !torch.vtensor<[?],f32> {