Materialize IndexType in the API.

Previously, index (aka affint) type was hidden under OtherType in the type API.
We will need to identify and operate on values of index types in the upcoming
MLFunc->CFGFunc(->LLVM) lowering passes.  Materialize index type into a
separate class and make it visible to LLVM RTTI hierarchy directly.
Practically, index is an integer type of unknown bit width and is accetable in
most places where regular integer types are.  This is purely an API change that
does not affect the IR.

After IndexType is separated out from OtherType, the remaining "other types"
are, in fact, TF-specific types only.  Further renaming may be of interest.

PiperOrigin-RevId: 220614026

mlir/include/mlir/IR/Builders.h

@@ -73,7 +73,8 @@ public:
   FloatType getF32Type();
   FloatType getF64Type();
 
-  OtherType getIndexType();
+  IndexType getIndexType();
+
   OtherType getTFControlType();
   OtherType getTFStringType();
   OtherType getTFResourceType();

mlir/include/mlir/IR/Types.h

@@ -25,6 +25,7 @@
 namespace mlir {
 class AffineMap;
 class FloatType;
+class IndexType;
 class IntegerType;
 class Location;
 class MLIRContext;
@@ -33,6 +34,7 @@ class OtherType;
 namespace detail {
 
 class TypeStorage;
+class IndexTypeStorage;
 class IntegerTypeStorage;
 class FloatTypeStorage;
 struct OtherTypeStorage;
@@ -66,7 +68,7 @@ public:
     TFString,
 
     /// These are marker for the first and last 'other' type.
-    FIRST_OTHER_TYPE = Index,
+    FIRST_OTHER_TYPE = TFControl,
     LAST_OTHER_TYPE = TFString,
 
     // Floating point.
@@ -138,12 +140,12 @@ public:
   unsigned getBitWidth() const;
 
   // Convenience factories.
+  static IndexType getIndex(MLIRContext *ctx);
   static IntegerType getInteger(unsigned width, MLIRContext *ctx);
   static FloatType getBF16(MLIRContext *ctx);
   static FloatType getF16(MLIRContext *ctx);
   static FloatType getF32(MLIRContext *ctx);
   static FloatType getF64(MLIRContext *ctx);
-  static OtherType getIndex(MLIRContext *ctx);
   static OtherType getTFControl(MLIRContext *ctx);
   static OtherType getTFString(MLIRContext *ctx);
   static OtherType getTFResource(MLIRContext *ctx);
@@ -236,6 +238,21 @@ inline FloatType Type::getF64(MLIRContext *ctx) {
   return FloatType::get(Kind::F64, ctx);
 }
 
+/// Index is special integer-like type with unknown platform-dependent bit width
+/// used in subscripts and loop induction variables.
+class IndexType : public Type {
+public:
+  using ImplType = detail::IndexTypeStorage;
+  IndexType() = default;
+  /* implicit */ IndexType(Type::ImplType *ptr);
+
+  /// Crete an IndexType instance, unique in the given context.
+  static IndexType get(MLIRContext *context);
+
+  /// Support method to enable LLVM-style type casting.
+  static bool kindof(Kind kind) { return kind == Kind::Index; }
+};
+
 /// This is a type for the random collection of special base types.
 class OtherType : public Type {
 public:
@@ -251,8 +268,8 @@ public:
   }
 };
 
-inline OtherType Type::getIndex(MLIRContext *ctx) {
-  return OtherType::get(Kind::Index, ctx);
+inline IndexType Type::getIndex(MLIRContext *ctx) {
+  return IndexType::get(ctx);
 }
 inline OtherType Type::getTFControl(MLIRContext *ctx) {
   return OtherType::get(Kind::TFControl, ctx);

mlir/lib/IR/Builders.cpp

@@ -60,7 +60,7 @@ FloatType Builder::getF32Type() { return Type::getF32(context); }
 
 FloatType Builder::getF64Type() { return Type::getF64(context); }
 
-OtherType Builder::getIndexType() { return Type::getIndex(context); }
+IndexType Builder::getIndexType() { return Type::getIndex(context); }
 
 OtherType Builder::getTFControlType() { return Type::getTFControl(context); }
 

mlir/lib/IR/MLIRContext.cpp

@@ -323,6 +323,9 @@ public:
   // Uniqui'ing of AffineConstantExprStorage using constant value as key.
   DenseMap<int64_t, AffineConstantExprStorage *> constExprs;
 
+  /// Unique index type (lazily constructed).
+  IndexTypeStorage *indexType = nullptr;
+
   /// Integer type uniquing.
   DenseMap<unsigned, IntegerTypeStorage *> integers;
 
@@ -554,6 +557,17 @@ FileLineColLoc *FileLineColLoc::get(UniquedFilename filename, unsigned line,
 // Type uniquing
 //===----------------------------------------------------------------------===//
 
+IndexType IndexType::get(MLIRContext *context) {
+  auto &impl = context->getImpl();
+
+  if (impl.indexType)
+    return impl.indexType;
+
+  impl.indexType = impl.allocator.Allocate<IndexTypeStorage>();
+  new (impl.indexType) IndexTypeStorage{{Kind::Index, context}};
+  return impl.indexType;
+}
+
 IntegerType IntegerType::get(unsigned width, MLIRContext *context) {
   assert(width <= kMaxWidth && "admissible integer bitwidth exceeded");
   auto &impl = context->getImpl();

mlir/lib/IR/TypeDetail.h

@@ -56,6 +56,8 @@ struct alignas(8) TypeStorage {
   unsigned subclassData : 24;
 };
 
+struct IndexTypeStorage : public TypeStorage {};
+
 struct IntegerTypeStorage : public TypeStorage {
   unsigned width;
 };

mlir/lib/IR/Types.cpp

@@ -53,6 +53,8 @@ unsigned Type::getBitWidth() const {
 unsigned Type::getSubclassData() const { return type->getSubclassData(); }
 void Type::setSubclassData(unsigned val) { type->setSubclassData(val); }
 
+IndexType::IndexType(Type::ImplType *ptr) : Type(ptr) {}
+
 IntegerType::IntegerType(Type::ImplType *ptr) : Type(ptr) {}
 
 unsigned IntegerType::getWidth() const {