Return a failure instead of crashing if shape inference can not be run because of unraked operand types

Signed-off-by: Jonas Rickert <[email protected]>

Author: jorickert

src/Dialect/ONNX/ONNXOps/Math/DFT.cpp

@@ -33,6 +33,9 @@ LogicalResult ONNXGenericDFTOpShapeHelper<OP_TYPE>::customComputeShape(
   // Get info about input data operand.
   Value input = operandAdaptor.getInput();
   // Get the rank to compensate for N dimensions.
+  if (!hasShapeAndRank(input)) {
+    return failure();
+  }
   int64_t rank = createIE->getShapedTypeRank(input);
 
   // Check if the dimension for axis is a literal and in range.

src/Dialect/ONNX/ONNXOps/Math/MatMul.cpp

@@ -55,6 +55,9 @@ LogicalResult ONNXGenericMatMulOpShapeHelper<OP_TYPE>::computeShape() {
   std::tie(A, B) = matMulInputs(operandAdaptor);
 
   // Size all the arrays to padded length.
+  if (!hasShapeAndRank(A) || !hasShapeAndRank(B)) {
+    return failure();
+  }
   uint64_t aRank = createIE->getShapedTypeRank(A);
   uint64_t bRank = createIE->getShapedTypeRank(B);
   int paddedRank = std::max(aRank, bRank);

src/Dialect/ONNX/ONNXOps/Math/Reduction.cpp

@@ -29,6 +29,9 @@ LogicalResult ONNXGenericReductionOpShapeHelper<OP_TYPE>::customComputeShape(
     DimsExpr &axes, int noopWithEmptyAxes) {
   typename OP_TYPE::Adaptor operandAdaptor(operands, op->getAttrDictionary());
   Value data = operandAdaptor.getData();
+  if (!hasShapeAndRank(data)) {
+    return failure();
+  }
   int64_t rank = createIE->getShapedTypeRank(data);
   // Normalize the axes: at present, we only support compile time axes, but
   // with keep_dim on, it might not be too difficult to generate the code.
@@ -104,7 +107,11 @@ LogicalResult ONNXGenericReductionOpShapeHelper<OP_TYPE>::computeShape() {
       createIE->getIntFromArrayAsSymbols(operandAdaptor.getAxes(), axes);
     } else {
       // When the axis is dynamic, try to infer the rank of output tensor
-      int64_t dataRank = createIE->getShapedTypeRank(operandAdaptor.getData());
+      const auto data = operandAdaptor.getData();
+      if (!hasShapeAndRank(data)) {
+        return failure();
+      }
+      int64_t dataRank = createIE->getShapedTypeRank(data);
       int64_t axlesSize = createIE->getArraySize(operandAdaptor.getAxes());
       if (!operandAdaptor.getKeepdims() && axlesSize < 0 /*undef shape*/) {
         // Even though we did not compute the shape in ShapeHelper, return

src/Dialect/ONNX/ONNXOps/Math/TopK.cpp

@@ -31,6 +31,9 @@ LogicalResult ONNXTopKOpShapeHelper::computeShape() {
   // Get info about X and K operands.
   Value X = operandAdaptor.getX();
   Value K = operandAdaptor.getK();
+  if (!hasShapeAndRank(X)) {
+    return failure();
+  }
   int64_t rank = createIE->getShapedTypeRank(X);
 
   // Axis to compute TopK.

src/Dialect/ONNX/ONNXOps/NN/Conv.cpp

@@ -374,6 +374,9 @@ LogicalResult ONNXConvTransposeOpShapeHelper::computeShape() {
   Value wValue = operandAdaptor.getW();
 
   // Basic information.
+  if (!hasShapeAndRank(xValue)) {
+    return failure();
+  }
   int64_t rank = createIE->getShapedTypeRank(xValue);
   int64_t spatialOffset = 2;
   int64_t spatialRank = rank - spatialOffset;

src/Dialect/ONNX/ONNXOps/NN/NNHelper.cpp.inc

@@ -31,6 +31,9 @@ LogicalResult ONNXGenericPoolOpShapeHelper<OP_TYPE>::customComputeShape(
     std::optional<ArrayAttr> strideOpt, std::optional<ArrayAttr> dilationOpt,
     bool hasFilter, bool ceilMode) {
   // Basic information.
+  if(!hasShapeAndRank(xValue)) {
+    return failure();
+  }
   int64_t rank = createIE->getShapedTypeRank(xValue);
   int64_t spatialOffset = 2;
   int64_t spatialRank = rank - spatialOffset;

src/Dialect/ONNX/ONNXOps/NN/Pooling.cpp

@@ -48,14 +48,18 @@ LogicalResult ONNXGenericGlobalPoolOpShapeHelper<OP_TYPE>::computeShape() {
 template <>
 LogicalResult ONNXMaxRoiPoolOpShapeHelper::computeShape() {
   ONNXMaxRoiPoolOpAdaptor operandAdaptor(operands, op->getAttrDictionary());
-
   IndexExpr channel = createIE->getShapeAsDim(operandAdaptor.getX(), 1);
-  uint64_t roisRank = createIE->getShapedTypeRank(operandAdaptor.getRois());
+
+  const auto rois = operandAdaptor.getRois();
+  if (!hasShapeAndRank(rois)) {
+    return failure();
+  }
+  uint64_t roisRank = createIE->getShapedTypeRank(rois);
   if (roisRank != 2)
     return op->emitError("rois rank is expected to be 2d");
 
   // 2d tensor: (num_rois, 5)
-  IndexExpr numRois = createIE->getShapeAsDim(operandAdaptor.getRois(), 0);
+  IndexExpr numRois = createIE->getShapeAsDim(rois, 0);
   DimsExpr pooledDims;
   createIE->getIntFromArrayAsLiterals(
       operandAdaptor.getPooledShape(), pooledDims);

src/Dialect/ONNX/ONNXOps/ShapeHelper.cpp

@@ -285,6 +285,11 @@ LogicalResult ONNXBroadcastOpShapeHelper::customComputeShape(
   DimsExpr dimsExpr;
   uint64_t numOfInputs = initialOperands.size();
 
+  if (!llvm::all_of(initialOperands,
+          [](Value initalOperand) { return hasShapeAndRank(initalOperand); })) {
+    return failure();
+  }
+
   // Compute rank of the output. Rank of the output is the maximum rank of all
   // initial operands.
   uint64_t additionalOperRank =

src/Dialect/ONNX/ONNXOps/Tensor/Compress.cpp

@@ -31,6 +31,9 @@ LogicalResult ONNXCompressOpShapeHelper::computeShape() {
   ONNXCompressOpAdaptor operandAdaptor(operands);
   Value input = operandAdaptor.getInput();
   Value cond = operandAdaptor.getCondition();
+  if (!hasShapeAndRank(input)) {
+    return failure();
+  }
   int64_t inputRank = createIE->getShapedTypeRank(input);
   createIE->assertHasShapeAndRank(cond);
   std::optional<int64_t> optionalAxis = compressOp.getAxis();

src/Dialect/ONNX/ONNXOps/Tensor/DepthToSpace.cpp

@@ -30,6 +30,9 @@ LogicalResult ONNXDepthToSpaceOpShapeHelper::computeShape() {
   ONNXDepthToSpaceOp depthOp = llvm::cast<ONNXDepthToSpaceOp>(op);
   ONNXDepthToSpaceOpAdaptor operandAdaptor(operands);
   Value input = operandAdaptor.getInput();
+  if (!hasShapeAndRank(input)) {
+    return failure();
+  }
   int64_t inputRank = createIE->getShapedTypeRank(input);
   assert(inputRank == 4 && "Unexpected input tensor rank");
   int64_t blocksize = depthOp.getBlocksize();