llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix Slice Folder OOB Crash and onnx.Shape lowering (#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.
pull/3845/head
zjgarvey 2024-11-01 15:33:21 -05:00 committed by GitHub
parent 738d45d3bb
commit 3104b66560
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 44 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
View File

@ -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();
}

9
lib/Dialect/Torch/IR/TorchOps.cpp
View File

@ -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);
}

26
test/Dialect/Torch/canonicalize.mlir
View File

@ -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> {