Adding missing check in QuantileQuantizedPerAxisType::verify (#66)

* Adding missing check in QuantileQuantizedPerAxisType::verify
* Adding minor test for illegal quantile array size for per axis type

Author: sartil

mlir/lib/Dialect/Quant/IR/QuantTypes.cpp

@@ -504,7 +504,6 @@ LogicalResult QuantileQuantizedPerAxisType::verify(
     return failure();
   }
 
-  const auto quantileArraySize = quantiles.size();
   unsigned typeWidth{};
   if (storageType.isa<IntegerType>()) {
     typeWidth = llvm::dyn_cast<IntegerType>(storageType).getWidth();
@@ -517,10 +516,17 @@ LogicalResult QuantileQuantizedPerAxisType::verify(
                           "types, Float8E4M3FNType and Float8E5M2Type ";
   }
 
-  const size_t expectedSize = 1 << typeWidth;
+  const size_t storageTypeRange = storageTypeMax - storageTypeMin + 1;
+  const size_t typeWidthSize = 1 << typeWidth;
+  const size_t expectedSize =
+      (storageTypeRange < typeWidthSize) ? storageTypeRange : typeWidthSize;
+
+  const auto quantileArraySize = quantiles.size();
   if (quantileArraySize != expectedSize) {
     return emitError() << "quantiles array size needs to be equal to "
-                          "2^(bit_size(storageType)), expected: "
+                          "2^(bit_size(storageType)), or (storageTypeMax - "
+                          "storageTypeMin + 1) when max and min differ from "
+                          "the type limits; expected: "
                        << expectedSize << ", found: " << quantileArraySize;
   }
 

mlir/test/Dialect/Quant/parse-quantile-invalid.mlir

@@ -27,6 +27,15 @@ func.func @parse() -> !qalias {
   return %0 : !qalias
 }
 
+// -----
+// Illegal quantile array size (per axis type)
+// expected-error@+1 {{quantiles array size needs to be equal to 2^(bit_size(storageType)), or (storageTypeMax - storageTypeMin + 1) when max and min differ from the type limits; expected: 256, found: 2}}
+!qalias = !quant.quantile<i8:f16:f32:1, {-1.0,1.0}:{-2.0e+2,-0.99872:120}>
+func.func @parse() -> !qalias {
+  %0 = "foo"() : () -> !qalias
+  return %0 : !qalias
+}
+
 // -----
 // Unrecognized token: trailing
 // expected-error@+1 {{expected '>'}}