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

462

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

20 ↗

(On Diff #295929)

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.

27 ↗

(On Diff #295929)

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/IR/StandardTypes.h

 //===- StandardTypes.h - MLIR Standard Type Classes -------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_STANDARDTYPES_H
 #define MLIR_IR_STANDARDTYPES_H
 
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/Types.h"
 
 namespace llvm {
 struct fltSemantics;
 } // namespace llvm
 
 namespace mlir {
 class AffineExpr;

 /// 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>();
+    // 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.isa<ComplexType, FloatType, IntegerType, OpaqueType, VectorType,
+                    IndexType>() ||
+           !type.getDialect().getNamespace().empty();
   }
 
   /// 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/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>)