Skip to content

Commit af420ff

Browse files
committed
Fix out-of-bounds FoldEmptyTensorWithDimOp crash #111270
1 parent 4468d58 commit af420ff

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -979,7 +979,10 @@ struct FoldEmptyTensorWithDimOp : public OpRewritePattern<DimOp> {
979979
auto emptyTensorOp = dimOp.getSource().getDefiningOp<EmptyOp>();
980980
if (!emptyTensorOp || !maybeConstantIndex)
981981
return failure();
982-
if (!emptyTensorOp.getType().isDynamicDim(*maybeConstantIndex))
982+
auto emptyTensorType = emptyTensorOp.getType();
983+
if (*maybeConstantIndex < 0 ||
984+
*maybeConstantIndex >= emptyTensorType.getRank() ||
985+
!emptyTensorType.isDynamicDim(*maybeConstantIndex))
983986
return failure();
984987
rewriter.replaceOp(dimOp,
985988
emptyTensorOp.getDynamicSize(*maybeConstantIndex));

0 commit comments

Comments
 (0)