[Mlir-commits] [mlir] 826d3ea - [mlir][tosa] Ranked check for transpose was wrong.

[Rob Suderman]

Author: Rob Suderman
Date: 2021-09-29T15:14:42-07:00
New Revision: 826d3eaae7e9a87c37d88f4b1a64ab8a9ce11086

URL: https://github.com/llvm/llvm-project/commit/826d3eaae7e9a87c37d88f4b1a64ab8a9ce11086
DIFF: https://github.com/llvm/llvm-project/commit/826d3eaae7e9a87c37d88f4b1a64ab8a9ce11086.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 780cd03d97438..3a025437bc05b 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -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;