add IntegerLikeTypeInterface to enable out-of-tree uses of built-in attributes that otherwise force integer types

[j2kun]

[edit to incorporate some of the discussion below]:

This PR adds support for builtin DenseIntElementsAttr and IntegerAttr using out-of-tree types whose values can be stored as integer(s).

This is motivated by the following (out of tree) use case. Suppose I define a type that is an integer with different semantics, such as

def ModArithType : TypeDef<ModArith_Dialect, "int", ...> {
  let summary = "Integer type with an arbitrary modulus";
  let parameters = (ins "::mlir::IntegerAttr":$modulus);
  let assemblyFormat = "`<` $modulus `>`";
}

Now I want to add constants, constant folding/propagation, etc., and enable this all elementwise for tensors of these types. I can define

def ModArith_ConstantOp : Op<ModArith_Dialect, "constant", 
     [Pure, ConstantLike, AllTypesMatch<["value", "output"]>]> {
  let summary = "Define a constant value via an attribute.";
  let arguments = (ins TypedAttrInterface:$value);
  let results = (outs TypeOrValueSemanticsContainer<ModArithType>:$output);
  let assemblyFormat = "attr-dict $value";
}

This is nearly identical to arith.constant, and the arith.constant behavior gets nice results because

1. The type of the op result matches the type of the attribute, so you don't need to specify the type twice.
2. It naturally supports tensor operands using dense attrs with tensor-typed results, and the dense attributes are convenient because they provide splats, hex strings, and array-based versions.
3. Contexts like folding that are only provided a TypedAttribute are able to work because the type on the attribute is the type of the op being folded.

However, for my custom type/op I run into issues:

1. The parser fails even for regular integer attributes, because Parser::parseDecOrHexAttr constrains the type provided to an integer type (in the integer branch)
2. The parser fails for dense elements attrs, e.g., mod_arith.constant dense<0> : tensor<8x!mod_arith.int<127 : i32>> because the parser hard-constrains a DenseElementsAttr to have integer or integer-like (again, in the integer branch of the parser).
3. The DenseElementsAttr itself similarly requires an integer-like type to be attached to the attribute, or else you can't do things like iterate over the values in it.
4. Constant folding doesn't have access to the underlying (out of tree) type from any values that have been folded, as those types are only on the attribute.

The goal of this PR is to enable ops that use out-of-tree types to be able to reuse the upstream attribute and parsing infrastructure for types whose constant values can be stored as integers/dense int elements.

Looking into it, it seems the attributes and parsers only really need to know the bit width of the underlying storage type so that they can construct the relevant APInts to store the values. Otherwise the attribute type is essentially a pass through. So I figured if the attribute's type can advertise its storage bitwidth, that would suffice to make the attributes more generic.

This proof of concept PR works with HEIR's mod_arith dialect as shown in google/heir#1758.

[llvmbot]

@llvm/pr-subscribers-mlir-ods

@llvm/pr-subscribers-mlir

Author: Jeremy Kun (j2kun)

Changes

This is a first draft and I'm looking for feedback.

This PR adds the IntegerLikeTypeInterface to BuiltinTypeInterfaces, and registers is on Index and Integer types.

The goal of this PR is to enable ops that use out-of-tree types to be able to reuse the upstream attribute parsing infrastructure for types whose constant values can be stored as integers/dense int elements. For example, in HEIR we have a mod_arith dialect which is effectively an integer storage type with an added modulus. This dialect has a mod_arith.constant op that we want to support scalar and dense attribute values, but because the types involved are mod_arith types, to use the dense int element (or even plain integer) attribute parsers from upstream I have to redo a lot of upstream work to swap out the attribute's underlying type.

Looking into it, it seems the parsers only really need to know the bit width of the underlying storage type so that they can construct the relevant APInts from parsed integer literals. Otherwise the attribute type is essentially a pass through. So I figured if the attribute's type can advertise its storage bitwidth, that would suffice to make the parsers more generic.

This proof of concept PR works with HEIR's mod_arith dialect as shown in google/heir#1758.

Notably, with this change I get the option to use splatted dense attributes out of tree for free (cf. syntax.mlir in the linked PR). It even does nice things like range checking for attributes that don't fit within the advertised bit width of the out-of-tree type.

---

Full diff: https://github.com/llvm/llvm-project/pull/137061.diff

8 Files Affected:

• (modified) mlir/include/mlir/IR/BuiltinAttributes.h (+3-1)
• (modified) mlir/include/mlir/IR/BuiltinTypeInterfaces.td (+38)
• (modified) mlir/include/mlir/IR/BuiltinTypes.td (+4-2)
• (modified) mlir/lib/AsmParser/AttributeParser.cpp (+9-7)
• (modified) mlir/lib/IR/AsmPrinter.cpp (+1-1)
• (modified) mlir/lib/IR/AttributeDetail.h (+3-2)
• (modified) mlir/lib/IR/BuiltinAttributes.cpp (+17-18)
• (modified) mlir/lib/IR/BuiltinTypes.cpp (+10)

diff --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h
index 67fab7ebc13ba..5eb2c891047cd 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.h
+++ b/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());
diff --git a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td b/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
index 4a4f818b46c57..3d459f006093a 100644
--- a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
+++ b/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_
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 771de01fc8d5d..6eb2ec333351a 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/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:
diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 2474e88373e04..a3252cf1964ee 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/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;
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 5b5ec841917e7..d74bf5d975f0b 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/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);
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 26d40ac3a38f6..96b269e3b1363 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/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();
 }
 
diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index daf79dc5de981..b185223085c3d 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/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(hasSameElementsOrSplat(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;
   }
 
diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 3924d082f0628..c38e33dba7e5d 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/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
 //===----------------------------------------------------------------------===//

[llvmbot]

@llvm/pr-subscribers-mlir-core

Author: Jeremy Kun (j2kun)

Changes

This is a first draft and I'm looking for feedback.

This PR adds the IntegerLikeTypeInterface to BuiltinTypeInterfaces, and registers is on Index and Integer types.

The goal of this PR is to enable ops that use out-of-tree types to be able to reuse the upstream attribute parsing infrastructure for types whose constant values can be stored as integers/dense int elements. For example, in HEIR we have a mod_arith dialect which is effectively an integer storage type with an added modulus. This dialect has a mod_arith.constant op that we want to support scalar and dense attribute values, but because the types involved are mod_arith types, to use the dense int element (or even plain integer) attribute parsers from upstream I have to redo a lot of upstream work to swap out the attribute's underlying type.

Looking into it, it seems the parsers only really need to know the bit width of the underlying storage type so that they can construct the relevant APInts from parsed integer literals. Otherwise the attribute type is essentially a pass through. So I figured if the attribute's type can advertise its storage bitwidth, that would suffice to make the parsers more generic.

This proof of concept PR works with HEIR's mod_arith dialect as shown in google/heir#1758.

Notably, with this change I get the option to use splatted dense attributes out of tree for free (cf. syntax.mlir in the linked PR). It even does nice things like range checking for attributes that don't fit within the advertised bit width of the out-of-tree type.

---

Full diff: https://github.com/llvm/llvm-project/pull/137061.diff

8 Files Affected:

• (modified) mlir/include/mlir/IR/BuiltinAttributes.h (+3-1)
• (modified) mlir/include/mlir/IR/BuiltinTypeInterfaces.td (+38)
• (modified) mlir/include/mlir/IR/BuiltinTypes.td (+4-2)
• (modified) mlir/lib/AsmParser/AttributeParser.cpp (+9-7)
• (modified) mlir/lib/IR/AsmPrinter.cpp (+1-1)
• (modified) mlir/lib/IR/AttributeDetail.h (+3-2)
• (modified) mlir/lib/IR/BuiltinAttributes.cpp (+17-18)
• (modified) mlir/lib/IR/BuiltinTypes.cpp (+10)

diff --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h
index 67fab7ebc13ba..5eb2c891047cd 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.h
+++ b/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());
diff --git a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td b/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
index 4a4f818b46c57..3d459f006093a 100644
--- a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
+++ b/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_
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 771de01fc8d5d..6eb2ec333351a 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/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:
diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 2474e88373e04..a3252cf1964ee 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/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;
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 5b5ec841917e7..d74bf5d975f0b 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/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);
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 26d40ac3a38f6..96b269e3b1363 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/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();
 }
 
diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index daf79dc5de981..b185223085c3d 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/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(hasSameElementsOrSplat(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;
   }
 
diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 3924d082f0628..c38e33dba7e5d 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/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
 //===----------------------------------------------------------------------===//

[j2kun]

Given #118891 perhaps @matthias-springer and @River707 could provide some feedback?

[River707]

I'm a little confused on what's being attempted here. From the description, I thought this would be just wanting to make it easier to parse/print integer values, but it's trying to remove the use of IntegerType for builtin attributes. IntegerType is very deeply integrated and ingrained in a lot of APIs, so I would be very against having an interface like this. The PR here only touches a few places that look at IntegerType, but there are quite a few, which means a lot of code paths will either crash out-right, or subtly fail.

For the syntax, I don't think you'll be able to directly reuse the dense<..> syntax, that is specific to the builtin attributes. What we possibly could do is expose some of the internal parsing code as a utility within the parser, so you could call it from your own attribute/type parser methods. I would not expect any changes outside of the parser/printer though.

[j2kun]

I will explain in more detail what I am trying to accomplish.

Suppose I define a type that is an integer with different semantics, such as

def ModArithType : TypeDef<ModArith_Dialect, "int", ...> {
  let summary = "Integer type with an arbitrary modulus";
  let parameters = (ins "::mlir::IntegerAttr":$modulus);
  let assemblyFormat = "`<` $modulus `>`";
}

Now I want to add constants, constant folding/propagation, etc., and enable this all elementwise for tensors of these types. I can define

def ModArith_ConstantOp : Op<ModArith_Dialect, "constant", 
     [Pure, ConstantLike, AllTypesMatch<["value", "output"]>]> {
  let summary = "Define a constant value via an attribute.";
  let arguments = (ins TypedAttrInterface:$value);
  let results = (outs TypeOrValueSemanticsContainer<ModArithType>:$output);
  let assemblyFormat = "attr-dict $value";
}

This is nearly identical to arith.constant, and it gets nice results because

1. The type of the op result matches the type of the attribute, so you don't need to specify the type twice.
2. It naturally supports tensor operands using dense attrs with tensor-typed results, and the dense attributes are convenient because they provide splats, hex strings, and array-based versions.

However, for my custom type/op I run into issues:

1. The parser fails even for regular integer attributes, because Parser::parseDecOrHexAttr constrains the type provided to an integer type (in the integer branch)
2. The parser fails for dense elements attrs, e.g., mod_arith.constant dense<0> : tensor<8x!mod_arith.int<127 : i32>> because the parser hard-constrains a DenseElementsAttr to have integer or integer-like (again, in the integer branch of the parser).
3. The DenseElementsAttr itself similarly requires an integer-like type to be attached to the attribute, or else you can't do things like iterate over the values in it.

My main goal here is to make the Integer and DenseElement attrs reusable for these non-builtin types, to make it more straightforward to implement elementwise ops with constant folding and propagation. Duplicating all of this infrastructure in out-of-tree projects is a lot of effort (even with parser utilities), as it either requires building new out-of-tree dense attribute types alongside the upstream ones, or else having a "constant" op (and the associated folding rules) juggle between upstream dense attributes (which must have integer-like types) and the corresponding op result type (which uses the custom type).

From what I can tell the only thing the integer flavors of these attributes actually needs is knowledge of the bit width of the integer used to store the attribute values as APInts. Providing this from an interface seems like a natural, idiomatic way to do it. This also seems in line with the recent change of FloatType to an interface.

Could you say more about what else might break by making this change? All existing uses of DenseElementsAttr are the same as before (they still use integer types), and cases that require integer types still must cast the type of the attribute and provide a default behavior anyway. The only way it can break is when a new type is used, and I am building this in conjunction with google/heir#1758 to see what else breaks for folding, canonicalization, etc. Would you want me to port a test-only version of that work upstream?

[matthias-springer]

For example, in HEIR we have a mod_arith dialect which is effectively an integer storage type with an added modulus.

fyi, we have a similar use case internally. There is a special integer type with a modulo type parameter. Initially, we wanted to reuse arith/... ops with this type because we get foldings/canonicalizations for free. One approach that we looked into was adding an IntegerTypeInterface. In the end, we did not explore that path further because ops like arith.addi enforce that all operands/results must have the same type, but we needed an op that works with different modulus values.

The parser fails even for regular integer attributes, because Parser::parseDecOrHexAttr constrains the type provided to an integer type (in the integer branch)

Do you actually want to reuse this function? It performs checks such as "parsed literal is not negative if the integer type is unsigned". You probably want to add your own checks to ensure that the parsed literal is divisible by the modulus value of your integer type.

The parser fails for dense elements attrs, e.g., mod_arith.constant dense<0> : tensor<8x!mod_arith.int<127 : i32>> because the parser hard-constrains a DenseElementsAttr to have integer or integer-like (again, in the integer branch of the parser).

Where is that check? I was looking at parseDenseElementsAttr, but couldn't find it there.

The DenseElementsAttr itself similarly requires an integer-like type to be attached to the attribute, or else you can't do things like iterate over the values in it.

Does denseAttr.getValues<Attribute>() not work?

[j2kun]

Initially, we wanted to reuse arith/... ops with this type because we get foldings/canonicalizations for free.

We don't want or hope to reuse arith dialect infrastructure, as the semantics are different even just for folding rules. We only want to reuse the auto-generated parsers and printers for literals, and the attribute storage for folding.

Parser::parseDecOrHexAttr
Do you actually want to reuse this function? It performs checks such as "parsed literal is not negative if the integer type is unsigned". You probably want to add your own checks to ensure that the parsed literal is divisible by the modulus value of your integer type.

I don't directly use this function, it is what parseAttribute calls when encountering an integer literal (cf.

llvm-project/mlir/lib/AsmParser/AttributeParser.cpp

Line 122 in 2ca071b

return parseDecOrHexAttr(type, /*isNegative=*/false);

). Note that here and below, I'm simply using assemblyFormat = "attr-dict $value" and using the auto-generated parser and printer.

The parser fails for dense elements attrs, e.g., mod_arith.constant dense<0> : tensor<8x!mod_arith.int<127 : i32>> because the parser hard-constrains a DenseElementsAttr to have integer or integer-like (again, in the integer branch of the parser).

Where is that check? I was looking at parseDenseElementsAttr, but couldn't find it there.

It is inside buildAttributeAPInt, which all of the parse methods use, and which this PR relaxes to allow non-integer types. In particular, tensor literal parsing calls it from

llvm-project/mlir/lib/AsmParser/AttributeParser.cpp

Line 655 in 2ca071b

buildAttributeAPInt(eltTy, isNegative, token.getSpelling());

and DenseArray from

llvm-project/mlir/lib/AsmParser/AttributeParser.cpp

Line 886 in 2ca071b

value = buildAttributeAPInt(type, isNegative, spelling);

The DenseElementsAttr itself similarly requires an integer-like type to be attached to the attribute, or else you can't do things like iterate over the values in it.

Does denseAttr.getValues<Attribute>() not work?

The default printer generated for the above assemblyFormat calls getValues<APInt> (I believe this is the more specific overload resolved during template resolution, but I am not a template wizard), which enforces the type constraint. Cf.

llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h

Line 551 in feaa5aa

if (!getElementType().isIntOrIndex())

[jpienaar]

OK, so goal is to enable printing/parsing generation with more diverse types (all the other changes just in service of this). Here it seems like TensorLiteralParser::getAttr is the problem as its hardcoded to support int, float or string. So a type interface "DenseElementsAttrContainable" (or some better name), would seem to be another option. This would enable specializing parsing & printing for the contained types. This I think also helps beyond integer types - which I think all needs custom printer/parsers currently.

[j2kun]

I am happy to add a "DenseElementsAttrContainable" interface, but it seems a bit clunky to support that for all the types supported in a DenseElementsAttr (float, complex, and integer). Would it make more sense to have an interface for each supported element class? "DenseIntElementsAttrCompatible", "DenseFloatElementsAttrCompatible," etc?

Supporting DenseElementsAttrs via an interface also doesn't cover the case of parsing a single integer literal. Would a separate interface for that work? Maybe, "IntegerLiteralParseable"?

Am I to infer from this discussion that the MLIR design prefers more interfaces that cover less behavior each, even if some of the interfaces overlap in what they do? My feeling was to have one interface that provides the underlying need (storage width for integers) and then any place in the codebase that requires knowledge of the underlying storage width can use it. That feels like a more solid design to me, but maybe this discussion has been had before many times and I'm not aware of how it was resolved.

[River707]

OK, so goal is to enable printing/parsing generation with more diverse types (all the other changes just in service of this). Here it seems like TensorLiteralParser::getAttr is the problem as its hardcoded to support int, float or string. So a type interface "DenseElementsAttrContainable" (or some better name), would seem to be another option. This would enable specializing parsing & printing for the contained types. This I think also helps beyond integer types - which I think all needs custom printer/parsers currently.

DenseElementsAttr only supports a specific set of types though, adding an interface would be on a path to opening up that attribute to be more generic (which I'm not sure we'd actually want to do - you can always define your own ElementsAttr).

[j2kun]

I'm not sure we'd actually want to do

So how do we figure out if this is an acceptable path forward? I can of course always write my own parser and attributes, but reuse of common components is the main feature of MLIR.

[j2kun]

To be more specific, the only objections expressed so far seem incidental to me:

• Some untested code paths may break
• The APIs are "integrated and ingrained" (technical debt from early design choices?)

Is there a more fundamental reason why this flexibility is not desired, if those issues could be overcome?

[joker-eph]

Is there a more fundamental reason why this flexibility is not desired, if those issues could be overcome?

I haven't followed the entire discussion because River was handling it, but to me there is something more fundamental: IntegerLikeTypeInterface seems a very general interface which rabbit hole into the support of arith dialect, etc.
That is if we introduce such an API this is to make the behavior of "Integer" in general in the system extensible. However you're interested only in reusing some parsing functionality, which is a very narrow use-case.
Solving a narrow use-case is much easier and has less pitfalls than solving a very general case (which requires lot of design considerations).

Parser/printer are also "debugging tools" in the system (a production compiler is unlikely to rely on these), so if there is something to tradeoff in a design it won't be towards parsing/printing. For example changing DenseElementAttr fundamentally just so that you can reuse a parsing function.

[j2kun]

I can accept that MLIR doesn't want to support generalized integer types throughout the system. I didn't intend to imply that, but in retrospect the location and name of the interface implied that, as did some parts of the original PR description.

I do want more than just parsing though, I want to be able to use the builtin attributes with out-of-tree types for all purposes (constant materialization, folding, etc.). Parsing is the first failure mode when trying to do this.

Would it be acceptable if we scoped the feature to "builtin attributes can use non-builtin types, when those types implement the appropriate interface?"

[joker-eph]

the location and name of the interface implied that, as did some parts of the original PR description.

Right, this is kind of what I was aiming at.

Would it be acceptable if we scoped the feature to "builtin attributes can use non-builtin types, when those types implement the appropriate interface?"

This seems consistent with what we did for the Vector type recently: #133455

You mention "builtin attributes", are you looking at more than DenseElementsAttr right now?
Note also the existence of the ElementsAttrInterface, which went in a different direction but with similar end goal somehow. Instead of making DenseElementsAttr extensible with custom types, we instead allowed folks to use their own replacement for DenseElementsAttr.

[River707]

the location and name of the interface implied that, as did some parts of the original PR description.

Right, this is kind of what I was aiming at.

Would it be acceptable if we scoped the feature to "builtin attributes can use non-builtin types, when those types implement the appropriate interface?"

This seems consistent with what we did for the Vector type recently: #133455

You mention "builtin attributes", are you looking at more than DenseElementsAttr right now? Note also the existence of the ElementsAttrInterface, which went in a different direction but with similar end goal somehow. Instead of making DenseElementsAttr extensible with custom types, we instead allowed folks to use their own replacement for DenseElementsAttr.

Yeah, DenseElementsAttr supporting "all of the cases" turned out to be a major pain and a mistake (though DenseElementsAttr isn't the attribute, it's DenseIntOrFPElementsAttr and DenseStringElementsAttr). That's a big part of why the extensibility there is pushed via ElementsAttr, it allows users to do whatever they want with the storage mechanisms. There is a lot of special casing code sprinkled all over those attributes that makes assumptions around int == IntegerType and fp == FloatType. Interface-ifying FloatType was an easier transition because there were always a number of specific float types. For IntegerType there has always only been "IntegerType", we didn't want to have multiple types in that hierarchy (which also goes back to why signedness is part of the type and not multiple different types).

[j2kun]

You mention "builtin attributes", are you looking at more than DenseElementsAttr right now?

IntegerAttr, which similarly only needs the bit width to determine the APInt storage size, and is similarly hard-coded to have integer type in the attribute as well as the parser.

extensibility there is pushed via ElementsAttr

Would it make sense to have an IntLikeElementsAttr upstream that supports more type compatibility but also uses integer storage? Then when I find other sources of friction upstream, I could generalize them to use ElementsAttr so that they naturally extend to this new attribute.

[joker-eph]

IntegerAttr, which similarly only needs the bit width to determine the APInt storage size, and is similarly hard-coded to have integer type in the attribute as well as the parser.

But consumers of the IntegerAttr expect right now that it only encodes an integer: they will consume it as such and manipulate it as such (in folder code for example): I don't quite see how making it extensible would work here?

Would it make sense to have an IntLikeElementsAttr upstream that supports more type compatibility but also uses integer storage?

Would that be an interface or a concrete attribute? If a concrete one, would it replace (or compete with) DenseIntOrFPElementsAttr for integers?
At this point I lost track if it's possible to allow other "int types" for DenseIntOrFPElementsAttr instead?

[j2kun]

But consumers of the IntegerAttr expect right now that it only encodes an integer: they will consume it as such and manipulate it as such (in folder code for example): I don't quite see how making it extensible would work here?

I think the nuance here is that IntegerAttr stores an integer (which it still will in this PR), but it already allows more than one Type to describe that integer, namely IndexType and IntegerType, which forces users to already handle it in enough generality for this PR to not break anything. Some evidence:

• The vast majority of upstream users of IntegerAttr simply ignore the type and use the underlying int (I have always wondered why IntegerAttr is typed, and I suspect it's only for the parser/printer which doesn't know the bit width for parsing or signedness for printing).
• Places like ArithCanonicalization support all types generically and only manipulate the underlying APInt. This would be invalid for a new type, but Arith is not possible to use with out-of-tree types.
• Common folders like constFoldBinaryOp are agnostic of the type and attribute entirely. They do nothing with it except pass it along to the result attribute of the fold. Users of these folders in Arith (which again cannot be used with out-of-tree types anyway) only require IntegerAttr stores an APInt. The data the type provides in these cases is not used.
• Documentation demonstrates that type checking should be done on IntegerAttr
• Because of the joint index and integer type support, many places in the codebase already have fallbacks and assertions to guard against an unexpected type.

Would it make sense to have an IntLikeElementsAttr upstream that supports more type compatibility but also uses integer storage?

Would that be an interface or a concrete attribute? If a concrete one, would it replace (or compete with) DenseIntOrFPElementsAttr for integers?

I am happy to take advice here. One of the difficulties of having a new attribute (interface or concrete) is that the Parser::parseAttribute accepts only a type as input, and infers from that the attribute to construct. I can avoid this by having a concrete IntLikeElementsAttr which defines its own custom parser/printer, though that may require duplicating some of the nice aspects of the builtin attributes (like the conveniences around splats). If it were an interface and the concrete attribute were out-of-tree, I'd have to think about how that would work (how would the interface manage something like folding to a splat if it doesn't manage the storage of the concrete attribute?).

At this point I lost track if it's possible to allow other "int types" for DenseIntOrFPElementsAttr instead?

I believe this is what this PR demonstrates, at least because it makes our out-of-tree mod_arith dialect google/heir#1758 work with dense<...>. If you'd like I can build more of the folding rules out in that PR for further evidence, or upstream a minimal version of it as a test-only dialect (or a proper dialect, since I know of other projects that made their own equivalent dialects to do non-machine-word-sized modular arithmetic).