add IntegerLikeTypeInterface to enable out-of-tree uses of int attribute parsers

Author: j2kun

mlir/include/mlir/IR/BuiltinAttributes.h

@@ -548,7 +548,9 @@ class DenseElementsAttr : public Attribute {
       std::enable_if_t<std::is_same<T, APInt>::value>;
   template <typename T, typename = APIntValueTemplateCheckT<T>>
   FailureOr<iterator_range_impl<IntElementIterator>> tryGetValues() const {
-    if (!getElementType().isIntOrIndex())
+    auto intLikeType =
+        llvm::dyn_cast<IntegerLikeTypeInterface>(getElementType());
+    if (!intLikeType)
       return failure();
     return iterator_range_impl<IntElementIterator>(getType(), raw_int_begin(),
                                                    raw_int_end());

mlir/include/mlir/IR/BuiltinTypeInterfaces.td

@@ -257,4 +257,42 @@ def ShapedTypeInterface : TypeInterface<"ShapedType"> {
   }];
 }
 
+def IntegerLikeTypeInterface : TypeInterface<"IntegerLikeTypeInterface"> {
+  let cppNamespace = "::mlir";
+  let description = [{
+    This type interface is for types that behave like integers. It provides
+    the API that allows MLIR utilities to treat them the same was as MLIR
+    treats integer types in settings like parsing and printing.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Returns the storage bit width for this type.
+      }],
+      /*retTy=*/"unsigned",
+      /*methodName=*/"getStorageBitWidth",
+      /*args=*/(ins)
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Returns true if this type is signed.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isSigned",
+      /*args=*/(ins),
+      /*defaultImplementation=*/"return true;"
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Returns true if this type is signless.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isSignless",
+      /*args=*/(ins),
+      /*defaultImplementation=*/"return true;"
+    >,
+  ];
+}
+
 #endif // MLIR_IR_BUILTINTYPEINTERFACES_TD_

mlir/include/mlir/IR/BuiltinTypes.td

@@ -466,7 +466,8 @@ def Builtin_Function : Builtin_Type<"Function", "function"> {
 //===----------------------------------------------------------------------===//
 
 def Builtin_Index : Builtin_Type<"Index", "index",
-    [VectorElementTypeInterface]> {
+    [VectorElementTypeInterface,
+     DeclareTypeInterfaceMethods<IntegerLikeTypeInterface, ["getStorageBitWidth"]>]> {
   let summary = "Integer-like type with unknown platform-dependent bit width";
   let description = [{
     Syntax:
@@ -497,7 +498,8 @@ def Builtin_Index : Builtin_Type<"Index", "index",
 //===----------------------------------------------------------------------===//
 
 def Builtin_Integer : Builtin_Type<"Integer", "integer",
-    [VectorElementTypeInterface]> {
+    [VectorElementTypeInterface,
+     DeclareTypeInterfaceMethods<IntegerLikeTypeInterface, ["getStorageBitWidth"]>]> {
   let summary = "Integer type with arbitrary precision up to a fixed limit";
   let description = [{
     Syntax:

mlir/lib/AsmParser/AttributeParser.cpp

@@ -16,6 +16,7 @@
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/IntegerSet.h"
@@ -366,8 +367,12 @@ static std::optional<APInt> buildAttributeAPInt(Type type, bool isNegative,
     return std::nullopt;
 
   // Extend or truncate the bitwidth to the right size.
-  unsigned width = type.isIndex() ? IndexType::kInternalStorageBitWidth
-                                  : type.getIntOrFloatBitWidth();
+  unsigned width;
+  if (auto intLikeType = dyn_cast<IntegerLikeTypeInterface>(type)) {
+    width = intLikeType.getStorageBitWidth();
+  } else {
+    width = type.getIntOrFloatBitWidth();
+  }
 
   if (width > result.getBitWidth()) {
     result = result.zext(width);
@@ -425,10 +430,6 @@ Attribute Parser::parseDecOrHexAttr(Type type, bool isNegative) {
     return FloatAttr::get(floatType, *result);
   }
 
-  if (!isa<IntegerType, IndexType>(type))
-    return emitError(loc, "integer literal not valid for specified type"),
-           nullptr;
-
   if (isNegative && type.isUnsignedInteger()) {
     emitError(loc,
               "negative integer literal not valid for unsigned integer type");
@@ -584,7 +585,8 @@ DenseElementsAttr TensorLiteralParser::getAttr(SMLoc loc, ShapedType type) {
   }
 
   // Handle integer and index types.
-  if (eltType.isIntOrIndex()) {
+  auto integerLikeType = dyn_cast<IntegerLikeTypeInterface>(eltType);
+  if (integerLikeType || eltType.isIntOrIndex()) {
     std::vector<APInt> intValues;
     if (failed(getIntAttrElements(loc, eltType, intValues)))
       return nullptr;

mlir/lib/IR/AsmPrinter.cpp

@@ -2656,7 +2656,7 @@ void AsmPrinter::Impl::printDenseIntOrFPElementsAttr(
         os << ")";
       });
     }
-  } else if (elementType.isIntOrIndex()) {
+  } else if (isa<IntegerLikeTypeInterface>(elementType)) {
     auto valueIt = attr.value_begin<APInt>();
     printDenseElementsAttrImpl(attr.isSplat(), type, os, [&](unsigned index) {
       printDenseIntElement(*(valueIt + index), os, elementType);

mlir/lib/IR/AttributeDetail.h

@@ -37,8 +37,9 @@ inline size_t getDenseElementBitWidth(Type eltType) {
   // Align the width for complex to 8 to make storage and interpretation easier.
   if (ComplexType comp = llvm::dyn_cast<ComplexType>(eltType))
     return llvm::alignTo<8>(getDenseElementBitWidth(comp.getElementType())) * 2;
-  if (eltType.isIndex())
-    return IndexType::kInternalStorageBitWidth;
+  if (auto intLikeType = dyn_cast<IntegerLikeTypeInterface>(eltType))
+    return intLikeType.getStorageBitWidth();
+
   return eltType.getIntOrFloatBitWidth();
 }
 

mlir/lib/IR/BuiltinAttributes.cpp

@@ -10,6 +10,7 @@
 #include "AttributeDetail.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/IntegerSet.h"
@@ -379,22 +380,20 @@ APSInt IntegerAttr::getAPSInt() const {
 
 LogicalResult IntegerAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                                   Type type, APInt value) {
-  if (IntegerType integerType = llvm::dyn_cast<IntegerType>(type)) {
-    if (integerType.getWidth() != value.getBitWidth())
-      return emitError() << "integer type bit width (" << integerType.getWidth()
-                         << ") doesn't match value bit width ("
-                         << value.getBitWidth() << ")";
-    return success();
+  unsigned width;
+  if (auto intLikeType = dyn_cast<IntegerLikeTypeInterface>(type)) {
+    width = intLikeType.getStorageBitWidth();
+  } else {
+    return emitError() << "expected integer-like type";
   }
-  if (llvm::isa<IndexType>(type)) {
-    if (value.getBitWidth() != IndexType::kInternalStorageBitWidth)
-      return emitError()
-             << "value bit width (" << value.getBitWidth()
-             << ") doesn't match index type internal storage bit width ("
-             << IndexType::kInternalStorageBitWidth << ")";
-    return success();
+
+  if (width != value.getBitWidth()) {
+    return emitError() << "integer-like type bit width (" << width
+                       << ") doesn't match value bit width ("
+                       << value.getBitWidth() << ")";
   }
-  return emitError() << "expected integer or index type";
+
+  return success();
 }
 
 BoolAttr IntegerAttr::getBoolAttrUnchecked(IntegerType type, bool value) {
@@ -1019,7 +1018,7 @@ DenseElementsAttr DenseElementsAttr::get(ShapedType type,
 /// element type of 'type'.
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<APInt> values) {
-  assert(type.getElementType().isIntOrIndex());
+  assert(isa<IntegerLikeTypeInterface>(type.getElementType()));
   assert(hasSameNumElementsOrSplat(type, values));
   size_t storageBitWidth = getDenseElementStorageWidth(type.getElementType());
   return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, values);
@@ -1130,11 +1129,11 @@ static bool isValidIntOrFloat(Type type, int64_t dataEltSize, bool isInt,
   if (type.isIndex())
     return true;
 
-  auto intType = llvm::dyn_cast<IntegerType>(type);
+  auto intType = llvm::dyn_cast<IntegerLikeTypeInterface>(type);
   if (!intType) {
     LLVM_DEBUG(llvm::dbgs()
-               << "expected integer type when isInt is true, but found " << type
-               << "\n");
+               << "expected integer-like type when isInt is true, but found "
+               << type << "\n");
     return false;
   }
 

mlir/lib/IR/BuiltinTypes.cpp

@@ -59,6 +59,14 @@ LogicalResult ComplexType::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// Index Type
+//===----------------------------------------------------------------------===//
+
+unsigned IndexType::getStorageBitWidth() const {
+  return kInternalStorageBitWidth;
+}
+
 //===----------------------------------------------------------------------===//
 // Integer Type
 //===----------------------------------------------------------------------===//
@@ -86,6 +94,8 @@ IntegerType IntegerType::scaleElementBitwidth(unsigned scale) {
   return IntegerType::get(getContext(), scale * getWidth(), getSignedness());
 }
 
+unsigned IntegerType::getStorageBitWidth() const { return getWidth(); }
+
 //===----------------------------------------------------------------------===//
 // Float Types
 //===----------------------------------------------------------------------===//