Relax the constraint for MemRef type to allow any non-builtin type as element type.

[chhe]

Summary

Relax the constraint for MemRef type to allow any non-builtin type as element type.

[rriddle]

mlir/include/mlir/IR/StandardTypes.h

459

Can you move this to the .cpp file instead so that we don't need another include here?

[chhe]

Move implementation to .cpp file.

[mehdi_amini]

The revision says that it allows "QuantType" but the implementation is actually making it accept any non-builtin type?

[chhe]

In D88710#2307661, @mehdi_amini wrote:

The revision says that it allows "QuantType" but the implementation is actually making it accept any non-builtin type?

Revision updated.

Quant element type for MemRef is the feature request. However, I think it may be inappropriate to explicitly only allow QuantType(it needs to add QuantType dependency to StandardType). So I just sync the constraint with TensorType.

[ftynse]

Similar changes have been proposed repeatedly. I am yet to see an RFC (which we normally require for standard/built-in changes) that clarifies the semantics of such memrefs. In particular, I am interested in the allocation/deallocation of such memrefs and any operations that require knowing the size of the elemental type. For example, how does one allocate memref<42 x opaque>? In general, there should be a clear specification of what dialects should check for in order for the type to be "valid" in a memref.

[chhe]

Add SizedTypeInterface and implement this interface for QuantType.

[chhe]

Fix interface function comment.

[rriddle]

mlir/include/mlir/Interfaces/SizedTypeInterface.td

21

I would prefer to see more discussion around such an interface first. This interface as proposed seems very simplistic, and not something that would be able to support LLVM for example.

28

Along the same lines as the above, I'm not sure bytes are a good metric to use when considering size given that bytes can vary between platforms. I would prefer to have an explicit proposal for this interface that has its own motivations other than for supporting non-builtin types in Memref.

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

 #define MLIR_DIALECT_QUANT_QUANT_TYPES_H_
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/SizedTypeInterface.h"
 #include "llvm/Support/MathExtras.h"
 
 namespace mlir {
 namespace quant {
 
 class QuantizedIntegerType;
 
 namespace detail {
public:
   /// Casts from a type based on the expressedType to the equivalent type
   /// based on storageType by way of this QuantizedType. Equivalent to:
   ///   QuantizedType::castToStorageType(castFromExpressedType(candidateType))
   /// (but with validity checks).
   /// Example (for this = !quant.uniform<i8:f32, 1.0>):
   ///   tensor<4xf32> -> tensor<4xi8>
   Type castExpressedToStorageType(Type candidateType);
 
+  /// Gets the size in bytes with quant storage type. This function will be used
+  /// by its child classes with SizedTypeInterface;
+  int64_t getSizeInBytes();
+
 private:
   /// Hide the following methods inherited from `Type`. It is almost certainly
   /// a bug to call them from a `QuantizedType` object. Users should call
   /// `getStorageType` or `getExpressedType` to get the underlying types
   /// they want to inspect.
   using Type::isBF16;
   using Type::isF16;
   using Type::isF32;
   using Type::isF64;
   using Type::isIndex;
   using Type::isInteger;
 };
 
 /// A quantized type that maps storage to/from expressed types in an
 /// unspecified way.
 ///
 /// 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,
+                            mlir::SizedTypeInterface::Trait> {
 public:
   using Base::Base;
 
   /// 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,
                               int64_t storageTypeMax);
 
   /// Gets an instance of the type with all specified parameters checked.

 ///     !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,
+                            SizedTypeInterface::Trait> {
 public:
   using Base::Base;
 
   /// Gets an instance of the type with all parameters specified but not
   /// checked.
   static UniformQuantizedType get(unsigned flags, Type storageType,
                                   Type expressedType, double scale,
                                   int64_t zeroPoint, int64_t storageTypeMin,
                                   int64_t storageTypeMax);
 
public:
   /// equation.
   int64_t getZeroPoint() const;
 
   // Fixed point values are real numbers divided by a scale.
   // Currently, only signed storage types are treated as fixed point.
   // A fixed point value can be obtained from an affine value by subtracting
   // the zeroPoint.
   // In the future, this may be explicit versus implied by type and zeroPoint.
   bool isFixedPoint() const { return isSigned() && getZeroPoint() == 0; }
 };
 
 /// Represents per-axis (also known as per-channel quantization).
 ///
 /// Syntax synopsis:
 ///   Per-axis, all parameters expressed:
 ///     !quant<uniform[StorageType:ExpressedType:QuantizedDim]{QuantParams}>
 ///   Per-axis, optional parameters omitted:
 ///     !quant<uniform[StorageType]{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,
+                            SizedTypeInterface::Trait> {
 public:
   using Base::Base;
 
   /// Gets an instance of the type with all parameters specified but not
   /// checked.
   static UniformQuantizedPerAxisType
   get(unsigned flags, Type storageType, Type expressedType,
       ArrayRef<double> scales, ArrayRef<int64_t> zeroPoints,
       int32_t quantizedDimension, int64_t storageTypeMin,
       int64_t storageTypeMax);
public:
   int32_t getQuantizedDimension() const;
 
   /// Fixed point values are real numbers divided by a scale.
   /// Currently, only signed storage types are treated as fixed point.
   /// A fixed point value can be obtained from an affine value by subtracting
   /// the zeroPoint.
   /// In the future, this may be explicit versus implied by type and zeroPoint.
   bool isFixedPoint() const {
     if (!isSigned())
       return false;
     return llvm::all_of(getZeroPoints(),
                         [](int64_t zeroPoint) { return zeroPoint != 0; });
   }
 };
 
 } // namespace quant
 } // namespace mlir
 
 #endif // MLIR_DIALECT_QUANT_QUANT_TYPES_H_

mlir/include/mlir/IR/StandardTypes.h

 
 /// Base MemRef for Ranked and Unranked variants
 class BaseMemRefType : public ShapedType {
 public:
   using ImplType = detail::BaseMemRefTypeStorage;
   using ShapedType::ShapedType;
 
   /// Return true if the specified element type is ok in a memref.
-  static bool isValidElementType(Type type) {
-    return type.isIntOrIndexOrFloat() || type.isa<VectorType, ComplexType>();
-  }
+  static bool isValidElementType(Type type);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type);
 
   /// Returns the memory space in which data referred to by this memref resides.
   unsigned getMemorySpace() const;
 };
 

mlir/include/mlir/Interfaces/CMakeLists.txt

 add_mlir_interface(CallInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
 add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
 add_mlir_interface(SideEffectInterfaces)
+add_mlir_interface(SizedTypeInterface)
 add_mlir_interface(VectorInterfaces)
 add_mlir_interface(ViewLikeInterface)
 

mlir/include/mlir/Interfaces/SizedTypeInterface.h

+//===- SizedTypeInterface.td - SizedType interface ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the interface for type to explictly have size in bytes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_SIZEDTYPEINTERFACE_H_
+#define MLIR_INTERFACES_SIZEDTYPEINTERFACE_H_
+
+#include "mlir/IR/OpDefinition.h"
+
+/// Include the generated interface declarations.
+#include "mlir/Interfaces/SizedTypeInterface.h.inc"
+
+#endif // MLIR_INTERFACES_SIZEDTYPEINTERFACE_H_

mlir/include/mlir/Interfaces/SizedTypeInterface.td

+//===- SizedTypeInterface.td - SizedType interface ---------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface for type to explictly have size in bytes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_SIZEDTYPEINTERFACE
+#define MLIR_INTERFACES_SIZEDTYPEINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def SizedTypeInterface : TypeInterface<"SizedTypeInterface"> {
+  let description = [{
+    An interface for types to provide size of bytes in memory.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<
+      "Returns the type size of bytes in memory.",
+      "int64_t", "getSizeInBytes">
+  ];
+}
+
+#endif // MLIR_INTERFACES_SIZEDTYPEINTERFACE

mlir/lib/Dialect/Quant/CMakeLists.txt

add_mlir_dialect_library(MLIRQuant
   DEPENDS
   MLIRQuantOpsIncGen
   MLIRQuantPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
   MLIRSideEffectInterfaces
+  MLIRSizedTypeInterface
   MLIRSupport
   MLIRStandard
   MLIRTransformUtils
   )

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

 Type QuantizedType::castExpressedToStorageType(Type candidateType) {
   Type expressedQuantizedType = castFromExpressedType(candidateType);
   if (!expressedQuantizedType) {
     return nullptr;
   }
   return QuantizedType::castToStorageType(expressedQuantizedType);
 }
 
+int64_t QuantizedType::getSizeInBytes() {
+  return getStorageType().getIntOrFloatBitWidth() / CHAR_BIT;
+}
+
 AnyQuantizedType AnyQuantizedType::get(unsigned flags, Type storageType,
                                        Type expressedType,
                                        int64_t storageTypeMin,
                                        int64_t storageTypeMax) {
   return Base::get(storageType.getContext(), flags, storageType, expressedType,
                    storageTypeMin, storageTypeMax);
 }
 
 AnyQuantizedType AnyQuantizedType::getChecked(unsigned flags, Type storageType,
                                               Type expressedType,
                                               int64_t storageTypeMin,
                                               int64_t storageTypeMax,
                                               Location location) {
   return Base::getChecked(location, flags, storageType, expressedType,
                           storageTypeMin, storageTypeMax);
 }
 
 LogicalResult AnyQuantizedType::verifyConstructionInvariants(
     Location loc, unsigned flags, Type storageType, Type expressedType,
     int64_t storageTypeMin, int64_t storageTypeMax) {
   if (failed(QuantizedType::verifyConstructionInvariants(
           loc, flags, storageType, expressedType, storageTypeMin,
LogicalResult AnyQuantizedType::verifyConstructionInvariants(
   return success();
 }
 
 UniformQuantizedType UniformQuantizedType::get(unsigned flags, Type storageType,
                                                Type expressedType, double scale,
                                                int64_t zeroPoint,
                                                int64_t storageTypeMin,
                                                int64_t storageTypeMax) {
   return Base::get(storageType.getContext(), flags, storageType, expressedType,
                    scale, zeroPoint, storageTypeMin, storageTypeMax);
 }
 
 UniformQuantizedType
 UniformQuantizedType::getChecked(unsigned flags, Type storageType,
                                  Type expressedType, double scale,
                                  int64_t zeroPoint, int64_t storageTypeMin,
                                  int64_t storageTypeMax, Location location) {
   return Base::getChecked(location, flags, storageType, expressedType, scale,
                           zeroPoint, storageTypeMin, storageTypeMax);
 }
 
 LogicalResult UniformQuantizedType::verifyConstructionInvariants(
     Location loc, unsigned flags, Type storageType, Type expressedType,
     double scale, int64_t zeroPoint, int64_t storageTypeMin,
     int64_t storageTypeMax) {
   if (failed(QuantizedType::verifyConstructionInvariants(
LogicalResult UniformQuantizedType::verifyConstructionInvariants(
 
   // Verify scale.
   if (scale <= 0.0 || std::isinf(scale) || std::isnan(scale))
     return emitError(loc, "illegal scale: ") << scale;
 
   return success();
 }
 
 double UniformQuantizedType::getScale() const { return getImpl()->scale; }
 
 int64_t UniformQuantizedType::getZeroPoint() const {
   return getImpl()->zeroPoint;
 }
 
 UniformQuantizedPerAxisType UniformQuantizedPerAxisType::get(
     unsigned flags, Type storageType, Type expressedType,
     ArrayRef<double> scales, ArrayRef<int64_t> zeroPoints,
     int32_t quantizedDimension, int64_t storageTypeMin,
     int64_t storageTypeMax) {
   return Base::get(storageType.getContext(), flags, storageType, expressedType,
                    scales, zeroPoints, quantizedDimension, storageTypeMin,
                    storageTypeMax);
 }
 
 UniformQuantizedPerAxisType UniformQuantizedPerAxisType::getChecked(
     unsigned flags, Type storageType, Type expressedType,
     ArrayRef<double> scales, ArrayRef<int64_t> zeroPoints,
     int32_t quantizedDimension, int64_t storageTypeMin, int64_t storageTypeMax,
     Location location) {
   return Base::getChecked(location, flags, storageType, expressedType, scales,
                           zeroPoints, quantizedDimension, storageTypeMin,
                           storageTypeMax);
 }
 
 LogicalResult UniformQuantizedPerAxisType::verifyConstructionInvariants(
     Location loc, unsigned flags, Type storageType, Type expressedType,
     ArrayRef<double> scales, ArrayRef<int64_t> zeroPoints,
     int32_t quantizedDimension, int64_t storageTypeMin,
for (double scale : scales) {
     if (scale <= 0.0 || std::isinf(scale) || std::isnan(scale))
       return emitError(loc, "illegal scale: ") << scale;
   }
 
   return success();
 }
 
 ArrayRef<double> UniformQuantizedPerAxisType::getScales() const {
   return getImpl()->getScales();
 }
 
 ArrayRef<int64_t> UniformQuantizedPerAxisType::getZeroPoints() const {
   return getImpl()->getZeroPoints();
 }
 
 int32_t UniformQuantizedPerAxisType::getQuantizedDimension() const {
   return getImpl()->quantizedDimension;
 }

mlir/lib/IR/CMakeLists.txt

add_mlir_library(MLIRIR
   Visitors.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
 
   DEPENDS
   MLIRCallInterfacesIncGen
   MLIROpAsmInterfaceIncGen
+  MLIRSizedTypeInterfaceIncGen
   MLIRSymbolInterfacesIncGen
   MLIRRegionKindInterfaceIncGen
 
   LINK_LIBS PUBLIC
   MLIRSupport
   )

mlir/lib/IR/StandardTypes.cpp

 //===- StandardTypes.cpp - MLIR Standard Type Classes ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/StandardTypes.h"
 #include "TypeDetail.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/Interfaces/SizedTypeInterface.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Twine.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 // Type
UnrankedTensorType::verifyConstructionInvariants(Location loc,
                                                  Type elementType) {
   return checkTensorElementType(loc, elementType);
 }
 
 //===----------------------------------------------------------------------===//
 // BaseMemRefType
 //===----------------------------------------------------------------------===//
 
+/// Return true if the specified element type is ok in a memref.
+bool BaseMemRefType::isValidElementType(Type type) {
+  // Note: Non standard/builtin types are allowed to exist within memref
+  // types. Dialects are expected to verify that memref types have a valid
+  // element type within that dialect.
+  return type.isIntOrIndexOrFloat() || type.isa<VectorType, ComplexType>() ||
+         type.dyn_cast<SizedTypeInterface>();
+}
+
 unsigned BaseMemRefType::getMemorySpace() const {
   return static_cast<ImplType *>(impl)->memorySpace;
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 

mlir/lib/Interfaces/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   CallInterfaces.cpp
   ControlFlowInterfaces.cpp
   CopyOpInterface.cpp
   DerivedAttributeOpInterface.cpp
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
   SideEffectInterfaces.cpp
+  SizedTypeInterface.cpp
   VectorInterfaces.cpp
   ViewLikeInterface.cpp
   )
 
 function(add_mlir_interface_library name)
   add_mlir_library(MLIR${name}
     ${name}.cpp
 

 
 add_mlir_interface_library(CallInterfaces)
 add_mlir_interface_library(ControlFlowInterfaces)
 add_mlir_interface_library(CopyOpInterface)
 add_mlir_interface_library(DerivedAttributeOpInterface)
 add_mlir_interface_library(InferTypeOpInterface)
 add_mlir_interface_library(LoopLikeInterface)
 add_mlir_interface_library(SideEffectInterfaces)
+add_mlir_interface_library(SizedTypeInterface)
 add_mlir_interface_library(VectorInterfaces)
 add_mlir_interface_library(ViewLikeInterface)
 

mlir/lib/Interfaces/SizedTypeInterface.cpp

+//===- SizedTypeInterface.td - SizedType interface ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/SizedTypeInterface.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// SizedType Interface
+//===----------------------------------------------------------------------===//
+
+/// Include the definitions of the SizedType Interface.
+#include "mlir/Interfaces/SizedTypeInterface.cpp.inc"

mlir/lib/Parser/TypeParser.cpp

Type Parser::parseMemRefType() {
 
   // Parse the element type.
   auto typeLoc = getToken().getLoc();
   auto elementType = parseType();
   if (!elementType)
     return nullptr;
 
   // Check that memref is formed from allowed types.
-  if (!elementType.isIntOrIndexOrFloat() &&
-      !elementType.isa<VectorType, ComplexType>())
+  if (!BaseMemRefType::isValidElementType(elementType))
     return emitError(typeLoc, "invalid memref element type"), nullptr;
 
   // Parse semi-affine-map-composition.
   SmallVector<AffineMap, 2> affineMapComposition;
   Optional<unsigned> memorySpace;
   unsigned numDims = dimensions.size();
 
   auto parseElt = [&]() -> ParseResult {

mlir/test/IR/parser.mlir

 
 // CHECK: func @tensors(tensor<*xf32>, tensor<*xvector<2x4xf32>>, tensor<1x?x4x?x?xi32>, tensor<i8>)
 func @tensors(tensor<* x f32>, tensor<* x vector<2x4xf32>>,
                  tensor<1x?x4x?x?xi32>, tensor<i8>)
 
 // CHECK: func @memrefs(memref<1x?x4x?x?xi32, #map{{[0-9]+}}>, memref<8xi8>)
 func @memrefs(memref<1x?x4x?x?xi32, #map0>, memref<8xi8, #map1, #map1>)
 
+// CHECK: func @memref_quant_element(memref<3x!quant.uniform<u8:f32, 1.000000e+00:66>>)
+func @memref_quant_element(memref<3x!quant.uniform<u8:f32, 1.0:66>>)
+
 // Test memref affine map compositions.
 
 // CHECK: func @memrefs2(memref<2x4x8xi8, 1>)
 func @memrefs2(memref<2x4x8xi8, #map2, 1>)
 
 // CHECK: func @memrefs23(memref<2x4x8xi8, #map{{[0-9]+}}>)
 func @memrefs23(memref<2x4x8xi8, #map2, #map3, 0>)