[mlir][tosa] Relax ranked constraint on quantization builder

[rsuderman]

Summary

TosaOp defintion had an artificial constraint that the input/output types
needed to be ranked to invoke the quantization builder. This is correct as an
unranked tensor could still be quantized.

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

static void buildMatMulOpWithQuantInfo(OpBuilder &builder,
                                        OperationState &result, Type outputType,
                                        Value a, Value b) {
   result.addOperands({a, b});
   auto quantAttr = ::buildMatMulOpQuantizationAttr(builder, a, b);
 
   if (quantAttr) {
     result.addAttribute("quantization_info", quantAttr);
 
-    auto inputType = a.getType().dyn_cast<RankedTensorType>();
-    assert(inputType && "Input must be a ranked tensor type!");
+    auto inputType = a.getType().dyn_cast<ShapedType>();
+    assert(inputType && "Input must be a shaped tensor type!");
 
     auto inputQType = inputType.getElementType()
                           .dyn_cast<mlir::quant::UniformQuantizedType>();
     assert(inputQType && "Tensor must have quantized datatype!");
 
     unsigned inputBits = inputQType.getStorageTypeIntegralWidth();
 
-    auto outputShapedType = outputType.dyn_cast<RankedTensorType>();
-    assert(outputShapedType && "Output must be a ranked tensor type");
-
-    auto outputShape = outputShapedType.getShape();
+    auto outputShapedType = outputType.dyn_cast<ShapedType>();
+    assert(outputShapedType && "Output must be a shaped type");
 
     IntegerType accElementType;
     if (inputBits == 16)
       accElementType = builder.getIntegerType(48);
     else
       accElementType = builder.getI32Type();
-    auto accType = RankedTensorType::get(outputShape, accElementType);
+    auto accType = outputShapedType.clone(accElementType);
     result.addTypes(accType);
   } else {
     result.addTypes(outputType);
   }
 }
 
 /// Both the tosa.avg_pool2d and unary ops use the same UnaruOpQuantizationAttr
 /// but avg_pool operator has its own builder as it has additional parameters

mlir/lib/Dialect/Tosa/Utils/QuantUtils.cpp

 /// Method to build ConvOpQuantizationAttr, called from
 /// ConvOpQuantInfoBuilder/TransConvOpQuantInfoBuilder:
 /// input_zp: input zeropoint
 /// weight_zp: weight zeropoint.
 ConvOpQuantizationAttr
 mlir::tosa::buildConvOpQuantizationAttr(OpBuilder &builder, Value input,
                                         Value weight) {
 
-  auto inputType = input.getType().dyn_cast<RankedTensorType>();
-  auto weightType = weight.getType().dyn_cast<RankedTensorType>();
+  auto inputType = input.getType().dyn_cast<ShapedType>();
+  auto weightType = weight.getType().dyn_cast<ShapedType>();
 
   if (!inputType || !weightType)
     return nullptr;
 
   auto inputQType = GET_UQTYPE(inputType);
   auto weightPerTensorQType = GET_UQTYPE(weightType);
   auto weightPerAxisQType = weightType.getElementType()
                                 .dyn_cast<quant::UniformQuantizedPerAxisType>();

 /// Builds MatMulOpQuantizationAttr, called from
 /// MatMulOpQuantInfoBuilder:
 /// aZp: input a zeropoint
 /// bZp: input b zeropoint.
 MatMulOpQuantizationAttr
 mlir::tosa::buildMatMulOpQuantizationAttr(OpBuilder &builder, Value a,
                                           Value b) {
 
-  auto aType = a.getType().dyn_cast<RankedTensorType>();
-  auto bType = b.getType().dyn_cast<RankedTensorType>();
+  auto aType = a.getType().dyn_cast<ShapedType>();
+  auto bType = b.getType().dyn_cast<ShapedType>();
 
   if (!aType || !bType)
     return nullptr;
 
   auto aQType = GET_UQTYPE(aType);
   auto bQType = GET_UQTYPE(bType);
 
   // A and B are either all quantized or all not quantized.

 /// Builds UnaryOpQuantizationAttr
 /// UnaryOpQuantInfoBuilder:
 /// inputZp: input zeropoint
 /// outputZp: output zeropoint.
 UnaryOpQuantizationAttr
 mlir::tosa::buildUnaryOpQuantizationAttr(OpBuilder &builder, Value input,
                                          Type outputRawType) {
 
-  auto inputType = input.getType().dyn_cast<RankedTensorType>();
-  auto outputType = outputRawType.dyn_cast<RankedTensorType>();
+  auto inputType = input.getType().dyn_cast<ShapedType>();
+  auto outputType = outputRawType.dyn_cast<ShapedType>();
 
   if (!inputType || !outputType)
     return nullptr;
 
   auto inputQType = GET_UQTYPE(inputType);
   auto outputQType = GET_UQTYPE(outputType);
 
   // Either all quantized or all not quantized.
mlir::tosa::buildUnaryOpQuantizationAttr(OpBuilder &builder, Value input,
   return nullptr;
 }
 
 /// Builds PadOpQuantizationAttr, called from PadOpQuantInfoBuilder:
 /// inputZp: input zeropoint.
 PadOpQuantizationAttr mlir::tosa::buildPadOpQuantizationAttr(OpBuilder &builder,
                                                              Value input) {
 
-  auto inputType = input.getType().dyn_cast<RankedTensorType>();
+  auto inputType = input.getType().dyn_cast<ShapedType>();
 
   if (!inputType)
     return nullptr;
 
   auto inputQType = GET_UQTYPE(inputType);
 
   if (inputQType) {
 
     int64_t inputZp = inputQType.getZeroPoint();
 
     auto quantAttr = tosa::PadOpQuantizationAttr::get(
         builder.getI32IntegerAttr(inputZp), builder.getContext());
 
     return quantAttr;
   }
 
   return nullptr;
 }
 
 /// Builds output type for a quantized ConvOp with the right bitwidth.
 /// This is called by the builder when dealing with quantized content.
 Type mlir::tosa::buildConvOpResultTypeInfo(OpBuilder &builder, Type outputType,
                                            Value input, Value weight) {
 
-  auto inputType = input.getType().dyn_cast<RankedTensorType>();
-  auto weightType = weight.getType().dyn_cast<RankedTensorType>();
+  auto inputType = input.getType().dyn_cast<ShapedType>();
+  auto weightType = weight.getType().dyn_cast<ShapedType>();
 
   assert(inputType && weightType &&
          "Could not extract input or weight tensors from Conv op");
 
   auto inputQType = GET_QTYPE(inputType);
   auto weightQType = GET_QTYPE(weightType);
 
   assert(inputQType && weightQType &&
          "Could not extract input or weight tensor types from Conv op");
 
   unsigned inputBits = inputQType.getStorageTypeIntegralWidth();
   unsigned weightBits = weightQType.getStorageTypeIntegralWidth();
 
-  auto outputShapedType = outputType.dyn_cast<RankedTensorType>();
+  auto outputShapedType = outputType.dyn_cast<ShapedType>();
   assert(outputShapedType &&
          "Could not extract output shape type from Conv op");
 
-  auto outputShape = outputShapedType.getShape();
-
   IntegerType accElementType;
   if (inputBits == 16 && weightBits == 8)
     accElementType = builder.getIntegerType(48);
   else
     accElementType = builder.getI32Type();
-  auto accType = RankedTensorType::get(outputShape, accElementType);
+  auto accType = outputShapedType.clone(accElementType);
   return accType;
 }
 
 /// Builds Tosa quantization attributes from min/max values.
 Type mlir::tosa::buildQTypeFromMinMax(OpBuilder builder, Type inputDType,
                                       Attribute minAttr, Attribute maxAttr,
                                       IntegerAttr quantBits, int filterQuantDim,
                                       bool isSigned, BoolAttr narrowRange) {