[mlir] Allow to use Quant types as elements of MemRef type

[vinograd47]

Summary

The required data layout information for them can be extracted from their storage type.

mlir/include/mlir/Dialect/Quant/QuantTypes.h

 /// Typical syntax:
 ///   quant.any<i8:f32>
 ///   quant.any<i8>
 ///   quant.any<i8<-16,15>>
 ///
 /// Note that for the any type, the expressed type is optional.
 class AnyQuantizedType
     : public Type::TypeBase<AnyQuantizedType, QuantizedType,
-                            detail::AnyQuantizedTypeStorage> {
+                            detail::AnyQuantizedTypeStorage,
+                            MemRefElementTypeInterface::Trait> {
 public:
   using Base::Base;
   using Base::getChecked;
 
   /// Gets an instance of the type with all parameters specified but not
   /// checked.
   static AnyQuantizedType get(unsigned flags, Type storageType,
                               Type expressedType, int64_t storageTypeMin,

 ///     !quant<uniform[StorageType]{Scale}>
 ///
 ///   StorageType: 'i'|'u' NumBits
 ///   ExpressedType: 'f16', 'f32', 'bf16', 'f64'
 ///   Scale: A legal double value
 ///   ZeroPoint: An integer value
 class UniformQuantizedType
     : public Type::TypeBase<UniformQuantizedType, QuantizedType,
-                            detail::UniformQuantizedTypeStorage> {
+                            detail::UniformQuantizedTypeStorage,
+                            MemRefElementTypeInterface::Trait> {
 public:
   using Base::Base;
   using Base::getChecked;
 
   /// Gets an instance of the type with all parameters specified but not
   /// checked.
   static UniformQuantizedType get(unsigned flags, Type storageType,
                                   Type expressedType, double scale,

 ///   StorageType: 'i'|'u' NumBits
 ///   ExpressedType: 'f16', 'f32', 'bf16', 'f64'
 ///   QuantizedDim: An integer value
 ///   QuantParams: (Scale ':' ZeroPoint)+
 ///   Scale: A legal double value
 ///   ZeroPoint: An integer value
 class UniformQuantizedPerAxisType
     : public Type::TypeBase<UniformQuantizedPerAxisType, QuantizedType,
-                            detail::UniformQuantizedPerAxisTypeStorage> {
+                            detail::UniformQuantizedPerAxisTypeStorage,
+                            MemRefElementTypeInterface::Trait> {
 public:
   using Base::Base;
   using Base::getChecked;
 
   /// Gets an instance of the type with all parameters specified but not
   /// checked.
   static UniformQuantizedPerAxisType
   get(unsigned flags, Type storageType, Type expressedType,

 };
 
 /// A quantized type that infers its range from given min/max values.
 ///
 /// Typical syntax:
 ///   quant.calibrated<f32<-0.922,0.981>>
 class CalibratedQuantizedType
     : public Type::TypeBase<CalibratedQuantizedType, QuantizedType,
-                            detail::CalibratedQuantizedTypeStorage> {
+                            detail::CalibratedQuantizedTypeStorage,
+                            MemRefElementTypeInterface::Trait> {
 public:
   using Base::Base;
   using Base::getChecked;
 
   /// Gets an instance of the type with all parameters specified but not
   /// checked.
   static CalibratedQuantizedType get(Type expressedType, double min,
                                      double max);

mlir/test/Dialect/Quant/parse-uniform.mlir

 // -----
 // Per-axis scales and zero points (mixed affine and fixedpoint)
 // CHECK: !quant.uniform<i8:f32:1, {2.000000e+02,9.987200e-01:120}>
 !qalias = type !quant.uniform<i8:f32:1, {2.0e+2,0.99872:120}>
 func @parse() -> !qalias {
   %0 = "foo"() : () -> !qalias
   return %0 : !qalias
 }
+
+// -----
+// MemRef of quantized type
+// CHECK: memref<1x2x3x4x!quant.uniform<i8:f32:1, {2.000000e+02,9.987200e-01:120}>>
+!qalias = type !quant.uniform<i8:f32:1, {2.0e+2,0.99872:120}>
+func @parse(%arg : memref<1x2x3x4x!qalias>) {
+  return
+}