[mlir][tosa] Ranked check for transpose was wrong.

Should have verified the perm length and input rank were the same before
inferring shape. Caused a crash with invalid IR.

Differential Revision: https://reviews.llvm.org/D110674
This commit is contained in:
Rob Suderman 2021-09-28 18:48:35 -07:00
parent 1c0e8a98e4
commit 826d3eaae7
1 changed files with 8 additions and 2 deletions

View File

@ -813,12 +813,18 @@ LogicalResult tosa::TransposeOp::inferReturnTypeComponents(
// If input rank and permutation length is unknown, the output rank is
// unknown.
if (!inputShape.hasRank() &&
(!permsShape.hasRank() || permsShape.isDynamicDim(0))) {
if (!inputShape.hasRank() || !permsShape.hasRank() ||
permsShape.isDynamicDim(0)) {
inferredReturnShapes.push_back(ShapedTypeComponents());
return success();
}
// This would imply the number of permutations does not match the rank of the
// input which is illegal.
if (permsShape.getDimSize(0) != inputShape.getRank()) {
return failure();
}
// Without the input dims we cannot determine the output dim sizes but we
// can determine the output rank.
SmallVector<int64_t> outputShape;