Make ONNX to TOSA lowering more resistant to dynamic shapes and unran…

…ked types.

- Add checks for types being ranked in various lowerings
- Refactor getSplattedConst to only require a type and not a rank
- Manually annotate the return type for some tosa ops to prevent type mismatches
-- In some cases to result type of an op can not be inferred just from its inputs, in these cases the annotation is always necxessary
-- In some cases the TOSA shape inference does not infere types even if it is possible, this is something I want to improve in another PR for MLIR.

src/Conversion/ONNXToTOSA/DialectBuilder.cpp

@@ -57,6 +57,12 @@ Value TosaBuilder::createConst(
 }
 
 bool TosaBuilder::needsRankBroadcast(ValueRange valueRange) {
+  if (llvm::any_of(valueRange, [](const auto value) {
+        return !mlir::cast<ShapedType>(value.getType()).hasRank();
+      })) {
+    return false; // we have no way to determine the broadcast, so do not
+                  // attempt it
+  }
   int64_t firstRank = mlir::cast<ShapedType>(valueRange[0].getType()).getRank();
   for (Value operand : valueRange) {
     auto operandType = mlir::cast<ShapedType>(operand.getType());
@@ -129,9 +135,8 @@ Value TosaBuilder::getConst(ArrayRef<float> vec, ArrayRef<int64_t> shape) {
   return constOp;
 }
 
-Value TosaBuilder::getSplattedConst(
-    float val, Type dtype, llvm::ArrayRef<int64_t> shape) {
-  auto constType = tosa::reduceAxisToOne(shape, rewriter().getF32Type());
+Value TosaBuilder::getSplattedConst(float val, Type dtype, int64_t rank) {
+  auto constType = tosa::reduceAxisToOne(rank, rewriter().getF32Type());
   auto constAttr = DenseElementsAttr::get(constType, val);
 
   auto constOp =
@@ -150,8 +155,7 @@ Value TosaBuilder::transpose(Value &value, llvm::ArrayRef<int32_t> perm) {
   auto valueType = mlir::cast<ShapedType>(value.getType());
   // get new value type
   Type newValueType = RankedTensorType::get(
-      llvm::SmallVector<int64_t, 4>(
-          valueType.getShape().size(), ShapedType::kDynamic),
+      llvm::SmallVector<int64_t, 4>(perm.size(), ShapedType::kDynamic),
       valueType.getElementType());
   // create transpose for value
   Value newValue = tosa::CreateOpAndInfer<mlir::tosa::TransposeOp>(
@@ -195,9 +199,12 @@ Value TosaBuilder::mul(Value &lhs, Value &rhs, int32_t shift) {
     rhs = valueVec[1];
   }
   auto lhsType = mlir::cast<ShapedType>(lhs.getType());
-  Type newValueType = RankedTensorType::get(
-      llvm::SmallVector<int64_t, 4>(lhsType.getRank(), ShapedType::kDynamic),
-      lhsType.getElementType());
+  Type newValueType =
+      (!lhsType.hasRank())
+          ? lhsType
+          : RankedTensorType::get(llvm::SmallVector<int64_t, 4>(
+                                      lhsType.getRank(), ShapedType::kDynamic),
+                lhsType.getElementType());
   return tosa::CreateOpAndInfer<mlir::tosa::MulOp>(
       rewriter(), loc(), newValueType, lhs, rhs, shift);
 }
@@ -215,9 +222,12 @@ Value TosaBuilder::intdiv(Value &lhs, Value &rhs) {
   }
 
   auto lhsType = mlir::cast<ShapedType>(lhs.getType());
-  Type newValueType = RankedTensorType::get(
-      llvm::SmallVector<int64_t, 4>(lhsType.getRank(), ShapedType::kDynamic),
-      lhsElementType);
+  Type newValueType =
+      (!lhsType.hasRank())
+          ? lhsType
+          : RankedTensorType::get(llvm::SmallVector<int64_t, 4>(
+                                      lhsType.getRank(), ShapedType::kDynamic),
+                lhsElementType);
   return tosa::CreateOpAndInfer<mlir::tosa::IntDivOp>(
       rewriter(), loc(), newValueType, lhs, rhs);
 }
@@ -230,9 +240,12 @@ Value TosaBuilder::binaryOp(Value &lhs, Value &rhs) {
     rhs = valueVec[1];
   }
   auto lhsType = mlir::cast<ShapedType>(lhs.getType());
-  Type newValueType = RankedTensorType::get(
-      llvm::SmallVector<int64_t, 4>(lhsType.getRank(), ShapedType::kDynamic),
-      lhsType.getElementType());
+  Type newValueType =
+      (!lhsType.hasRank())
+          ? lhsType
+          : RankedTensorType::get(llvm::SmallVector<int64_t, 4>(
+                                      lhsType.getRank(), ShapedType::kDynamic),
+                lhsType.getElementType());
   return tosa::CreateOpAndInfer<T>(rewriter(), loc(), newValueType, lhs, rhs);
 }
 
@@ -246,11 +259,7 @@ template Value TosaBuilder::binaryOp<mlir::tosa::PowOp>(
 
 template <typename T>
 Value TosaBuilder::unaryOp(mlir::Value &input) {
-  auto inputType = cast<ShapedType>(input.getType());
-  Type newValueType = RankedTensorType::get(
-      llvm::SmallVector<int64_t, 4>(inputType.getRank(), ShapedType::kDynamic),
-      inputType.getElementType());
-  return tosa::CreateOpAndInfer<T>(rewriter(), loc(), newValueType, input);
+  return tosa::CreateOpAndInfer<T>(rewriter(), loc(), input.getType(), input);
 }
 
 template Value TosaBuilder::unaryOp<mlir::tosa::ExpOp>(mlir::Value &input);
@@ -305,9 +314,12 @@ Value TosaBuilder::select(
     rhs = valueVec[2];
   }
   auto lhsType = cast<ShapedType>(lhs.getType());
-  Type newValueType = RankedTensorType::get(
-      llvm::SmallVector<int64_t, 4>(lhsType.getRank(), ShapedType::kDynamic),
-      lhsType.getElementType());
+  Type newValueType =
+      (!lhsType.hasRank())
+          ? lhsType
+          : RankedTensorType::get(llvm::SmallVector<int64_t, 4>(
+                                      lhsType.getRank(), ShapedType::kDynamic),
+                lhsType.getElementType());
   return tosa::CreateOpAndInfer<mlir::tosa::SelectOp>(
       rewriter(), loc(), newValueType, cond, lhs, rhs);
 }
@@ -328,7 +340,7 @@ mlir::Value TosaBuilder::castToNewTensorElementType(
 Value TosaBuilder::sqrt(mlir::Value &input) {
   auto inputType = cast<ShapedType>(input.getType());
   auto oneHalf = this->getSplattedConst(
-      0.5, inputType.getElementType(), inputType.getShape());
+      0.5, inputType.getElementType(), inputType.getRank());
   return this->binaryOp<mlir::tosa::PowOp>(input, oneHalf);
 }
 

src/Conversion/ONNXToTOSA/DialectBuilder.hpp

@@ -91,8 +91,7 @@ struct TosaBuilder : DialectBuilder {
   // The tensor will have the same rank as shape but all dimensions will
   // have size 1 (differs from tensorflow impl.)
   // If dtype is provided, it also cast the value to the appropriate dtype.
-  mlir::Value getSplattedConst(
-      float val, mlir::Type dtype, llvm::ArrayRef<int64_t> shape = {});
+  mlir::Value getSplattedConst(float val, mlir::Type dtype, int64_t rank);
 
   // Creates a constant of shape <1x1x...x1> of rank `rank` with all values set
   // to `value`.

src/Conversion/ONNXToTOSA/Math/Elementwise.cpp

@@ -239,12 +239,12 @@ static LogicalResult legalizeFloatingPointPrelu(Operation *op,
   auto loc = op->getLoc();
   TosaBuilder tosaBuilder(rewriter, loc);
   Value constZero = tosaBuilder.getSplattedConst(
-      0.0, outputType.getElementType(), outputType.getShape());
+      0.0, outputType.getElementType(), outputType.getRank());
 
   auto mul = tosaBuilder.mul(input, alphaOrSlope);
   auto greaterEqual = tosaBuilder.greaterEqual(input, constZero);
   auto select = tosaBuilder.select(greaterEqual, input, mul);
-
+  copySingleResultType(op, select);
   rewriter.replaceOp(op, {select});
   return success();
 }
@@ -274,7 +274,7 @@ class ONNXLeakyReluOpLoweringToTOSA
     TosaBuilder tosaBuilder(rewriter, loc);
     return legalizeFloatingPointPrelu(op, rewriter, adaptor.getX(),
         tosaBuilder.getSplattedConst(
-            alpha, outputType.getElementType(), outputType.getShape()),
+            alpha, outputType.getElementType(), outputType.getRank()),
         outputType);
   }
 };
@@ -312,6 +312,7 @@ class ONNXComparisonOpLoweringToTOSA : public OpConversionPattern<OnnxCompOp> {
     } else if constexpr (std::is_same_v<OnnxCompOp, ONNXLessOp>) {
       res = tosaBuilder.less(input1, input2);
     }
+    copySingleResultType(op, res);
     rewriter.replaceOp(op, {res});
     return success();
   }
@@ -393,7 +394,7 @@ class ONNXCastOpLoweringToTOSA : public OpConversionPattern<ONNXCastOp> {
       // onnx.Cast and tosa.cast.
       if (resultTy.getElementType().getIntOrFloatBitWidth() != 1) {
         auto zero = tosaBuilder.getSplattedConst(
-            0.0f, inputTy.getElementType(), resultTy.getShape());
+            0.0f, inputTy.getElementType(), resultTy.getRank());
         auto positive = tosaBuilder.greaterEqual(input, zero);
 
         auto floor = tosaBuilder.unaryOp<mlir::tosa::FloorOp>(input);
@@ -421,13 +422,15 @@ class ONNXDivOpLoweringToTOSA : public OpConversionPattern<ONNXDivOp> {
 
     if (isa<IntegerType>(resultElementType)) {
       Value divOp = tosaBuilder.intdiv(lhs, rhs);
+      copySingleResultType(op, divOp);
       rewriter.replaceOp(op, {divOp});
       return success();
     }
     // For floating point types, decompose ONNXDivOp into
     // tosa::ReciprocalOp and tosa::MulOp.
     Value reciprocalOp = tosaBuilder.unaryOp<mlir::tosa::ReciprocalOp>(rhs);
     Value mulOp = tosaBuilder.mul(lhs, reciprocalOp);
+    copySingleResultType(op, mulOp);
     rewriter.replaceOp(op, {mulOp});
     return success();
   }
@@ -472,20 +475,21 @@ class ONNXEluOpLoweringToTOSA : public OpConversionPattern<ONNXEluOp> {
     TosaBuilder tosaBuilder(rewriter, op->getLoc());
 
     Value one = tosaBuilder.getSplattedConst(
-        1.0, resultTensorType.getElementType(), resultTensorType.getShape());
+        1.0, resultTensorType.getElementType(), resultTensorType.getRank());
     Value alpha =
         tosaBuilder.getSplattedConst(adaptor.getAlpha().convertToDouble(),
-            resultTensorType.getElementType(), resultTensorType.getShape());
+            resultTensorType.getElementType(), resultTensorType.getRank());
     Value constZero = tosaBuilder.getSplattedConst(
-        0.0, resultTensorType.getElementType(), resultTensorType.getShape());
+        0.0, resultTensorType.getElementType(), resultTensorType.getRank());
 
     Value exp = tosaBuilder.unaryOp<mlir::tosa::ExpOp>(input);
+    copySingleResultType(op, exp);
     Value expMinusOne = tosaBuilder.binaryOp<mlir::tosa::SubOp>(exp, one);
     Value alphaTimesExpMinusOne = tosaBuilder.mul(expMinusOne, alpha);
     Value greaterEqual = tosaBuilder.greaterEqual(input, constZero);
     auto select =
         tosaBuilder.select(greaterEqual, input, alphaTimesExpMinusOne);
-
+    copySingleResultType(op, select);
     rewriter.replaceOp(op, {select});
     return success();
   }
@@ -516,11 +520,16 @@ class ONNXHardSigmoidOpLoweringToTOSA
     APFloat oneOverAlpha(alpha.getSemantics(), 1);
     oneOverAlpha.divide(alpha, APFloat::rmNearestTiesToEven);
 
+    if (!resultType.hasRank()) {
+      return rewriter.notifyMatchFailure(
+          op, "HardSigmoid: Static shape required to create splatted const");
+    }
+
     Value constBetaOverAlpha =
         tosaBuilder.getSplattedConst(betaOverAlpha.convertToDouble(),
-            resultElementType, resultType.getShape());
+            resultElementType, resultType.getRank());
     Value constAlpha = tosaBuilder.getSplattedConst(
-        alpha.convertToDouble(), resultElementType, resultType.getShape());
+        alpha.convertToDouble(), resultElementType, resultType.getRank());
 
     auto addOp =
         tosaBuilder.binaryOp<mlir::tosa::AddOp>(input, constBetaOverAlpha);
@@ -530,7 +539,7 @@ class ONNXHardSigmoidOpLoweringToTOSA
         rewriter.getF32FloatAttr(0),
         rewriter.getF32FloatAttr(oneOverAlpha.convertToDouble()));
     auto mulOp = tosaBuilder.mul(clampOp, constAlpha);
-
+    copySingleResultType(op, mulOp);
     rewriter.replaceOp(op, {mulOp});
     return success();
   }
@@ -565,14 +574,19 @@ class ONNXSoftplusOpLoweringToTOSA
     if (failed(IsFloat::checkType(rewriter, outputType.getElementType(), op))) {
       return failure();
     }
+    if (!outputType.hasRank()) {
+      return rewriter.notifyMatchFailure(
+          op, "ONNXSoftplusOp: Rank required to create splatted const");
+    }
 
     Value input = adaptor.getX();
 
     TosaBuilder tosaBuilder(rewriter, op->getLoc());
     auto one = tosaBuilder.getSplattedConst(
-        1.0, outputType.getElementType(), outputType.getShape());
+        1.0, outputType.getElementType(), outputType.getRank());
 
     auto expOp = tosaBuilder.unaryOp<mlir::tosa::ExpOp>(input);
+    copySingleResultType(op, expOp);
     auto expPlusOne = tosaBuilder.binaryOp<mlir::tosa::AddOp>(expOp, one);
     auto logOp = tosaBuilder.unaryOp<mlir::tosa::LogOp>(expPlusOne);
     rewriter.replaceOp(op, {logOp});
@@ -594,15 +608,19 @@ class ONNXSeluOpLoweringToTOSA : public OpConversionPattern<ONNXSeluOp> {
     Value input = adaptor.getX();
 
     TosaBuilder tosaBuilder(rewriter, op->getLoc());
+    if (!outputType.hasRank()) {
+      return rewriter.notifyMatchFailure(
+          op, "ONNXSeluOp: Rank required to create splatted const");
+    }
 
     Value alpha =
         tosaBuilder.getSplattedConst(adaptor.getAlpha().convertToDouble(),
-            outputType.getElementType(), outputType.getShape());
+            outputType.getElementType(), outputType.getRank());
     Value gamma =
         tosaBuilder.getSplattedConst(adaptor.getGamma().convertToDouble(),
-            outputType.getElementType(), outputType.getShape());
+            outputType.getElementType(), outputType.getRank());
     Value constZero = tosaBuilder.getSplattedConst(
-        0.0, outputType.getElementType(), outputType.getShape());
+        0.0, outputType.getElementType(), outputType.getRank());
 
     Value exp = tosaBuilder.unaryOp<mlir::tosa::ExpOp>(input);
     Value expTimesAlpha = tosaBuilder.mul(exp, alpha);
@@ -630,15 +648,19 @@ class ONNXThresholdedReluOpLoweringToTOSA
             rewriter, outputType.getElementType(), op))) {
       return failure();
     }
+    if (!outputType.hasRank()) {
+      return rewriter.notifyMatchFailure(
+          op, "ONNXThresholdedReluOp: Rank required to create splatted const");
+    }
 
     Value input = adaptor.getX();
 
     TosaBuilder tosaBuilder(rewriter, op->getLoc());
     auto alpha =
         tosaBuilder.getSplattedConst(adaptor.getAlpha().convertToDouble(),
-            outputType.getElementType(), outputType.getShape());
+            outputType.getElementType(), outputType.getRank());
     auto zero = tosaBuilder.getSplattedConst(
-        0.0, outputType.getElementType(), outputType.getShape());
+        0.0, outputType.getElementType(), outputType.getRank());
 
     auto greater = tosaBuilder.greater(input, alpha);
     auto select = tosaBuilder.select(greater, input, zero);

src/Conversion/ONNXToTOSA/Math/Gemm.cpp

@@ -49,6 +49,10 @@ class ONNXGemmOpLoweringToTOSA : public OpConversionPattern<ONNXGemmOp> {
     FloatAttr beta = adaptor.getBetaAttr();
     auto AType = mlir::cast<TensorType>(A.getType());
     auto BType = mlir::cast<TensorType>(B.getType());
+    if (!AType.hasRank() || !BType.hasRank()) {
+      return rewriter.notifyMatchFailure(
+          op, "Lowering Gemm to MatMul requires ranked A and B.");
+    }
     auto shapeA = AType.getShape();
     auto shapeB = BType.getShape();
     auto resultType = mlir::cast<TensorType>(
@@ -103,7 +107,7 @@ class ONNXGemmOpLoweringToTOSA : public OpConversionPattern<ONNXGemmOp> {
     if (alpha && alpha.getValueAsDouble() != 1.) {
       Value splattedConstAlpha = tosaBuilder.getSplattedConst(
           static_cast<float>(alpha.getValueAsDouble()), AType.getElementType(),
-          newShapeA);
+          newShapeA.size());
       alphaMulResult = tosaBuilder.mul(splattedConstAlpha, A, 0);
     }
 
@@ -112,7 +116,7 @@ class ONNXGemmOpLoweringToTOSA : public OpConversionPattern<ONNXGemmOp> {
     if (beta && isCPresent && beta.getValueAsDouble() != 1.) {
       Value splattedConstBeta = tosaBuilder.getSplattedConst(
           static_cast<float>(beta.getValueAsDouble()), AType.getElementType(),
-          newShapeA);
+          newShapeA.size());
       betaMulResult = tosaBuilder.mul(splattedConstBeta, C, 0);
     }
 

src/Conversion/ONNXToTOSA/Math/Reduce.cpp

@@ -179,7 +179,7 @@ LogicalResult reduceMeanLowering(ONNXReduceMeanOp op,
 
   TosaBuilder tosaBuilder(rewriter, op->getLoc());
   Value divConst = tosaBuilder.getSplattedConst(
-      divScale, outputType.getElementType(), outputType.getShape());
+      divScale, outputType.getElementType(), outputType.getRank());
   auto output = tosaBuilder.mul(val, divConst);
 
   if (!output) {

src/Conversion/ONNXToTOSA/NN/AveragePool.cpp

@@ -56,7 +56,7 @@ LogicalResult handleIncludePadAttr(
   Value padding = tosa::buildOnnxToTosaPaddingConstOp(
       rewriter, pads, loc, {0, 0, 0, 0}, {});
   auto constTosaTensor =
-      tosaBuilder.getSplattedConst(0.0, inputType.getElementType());
+      tosaBuilder.getSplattedConst(0.0, inputType.getElementType(), 0);
 
   auto padOp = tosa::CreateOpAndInfer<mlir::tosa::PadOp>(rewriter, loc,
       mlir::RankedTensorType::get(

src/Conversion/ONNXToTOSA/NN/BatchNorm.cpp

@@ -29,6 +29,11 @@ class ONNXBatchNormalizationInferenceModeOpLoweringToTOSA
       OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override {
 
     auto outType = getTypeConverter()->convertType(op.getResult().getType());
+    if (!cast<ShapedType>(outType).hasRank()) {
+      return rewriter.notifyMatchFailure(op,
+          "ONNXBatchNormalizationInferenceModeOp to "
+          "TOSA requires a ranked result type");
+    }
     auto outTensorType = cast<RankedTensorType>(outType);
 
     // The layout of the output is N x C x D1 x D2 … Dn. For batch
@@ -60,7 +65,7 @@ class ONNXBatchNormalizationInferenceModeOpLoweringToTOSA
     // epsilon's shape: constant -> {1, 1, 1, ...}
     newShape[1] = 1;
     auto eps = tosaBuilder.getSplattedConst(op.getEpsilon().convertToFloat(),
-        outTensorType.getElementType(), newShape);
+        outTensorType.getElementType(), newShape.size());
 
     // output = (input - mean) * scale * rsqrt(var + eps) + bias
     auto op1SubInputMean =

src/Conversion/ONNXToTOSA/ONNXToTOSACommon.cpp

@@ -418,7 +418,7 @@ std::optional<Value> convertReduceMeanOp(PatternRewriter &rewriter,
 
   if (!input_is_qtype) {
     Value div_const = tosaBuilder.getSplattedConst(
-        div_scale, output_type.getElementType(), output_type.getShape());
+        div_scale, output_type.getElementType(), output_type.getRank());
     return tosaBuilder.mul(val.value(), div_const);
   }
 

src/Conversion/ONNXToTOSA/ONNXToTOSALegalizeUtils.cpp

@@ -56,9 +56,9 @@ llvm::SmallVector<int64_t> createInt64VectorFromIndexExpr(
 }
 
 mlir::RankedTensorType reduceAxisToOne(
-    llvm::ArrayRef<int64_t> shape, Type elementType, Attribute encoding) {
+    int64_t rank, Type elementType, Attribute encoding) {
   return mlir::RankedTensorType::get(
-      llvm::SmallVector<int64_t, 4>(shape.size(), 1), elementType, encoding);
+      llvm::SmallVector<int64_t, 4>(rank, 1), elementType, encoding);
 }
 
 mlir::ElementsAttr getElementsAttrFromConst(mlir::Value &val) {