[mlir] Update several usages of IntegerType to properly handled unsignedness.

[rriddle]

Summary

For example, DenseElementsAttr currently does not properly round-trip unsigned integer values.

[antiagainst]

Nice. Thanks River! I've a CL to add more ODS definitions for integer signedness but haven't pushed out. Will do that soon.

mlir/include/mlir/IR/Matchers.h

struct constant_int_op_binder {
   constant_int_op_binder(IntegerAttr::ValueType *bv) : bind_value(bv) {}
 
   bool match(Operation *op) {
     Attribute attr;
     if (!constant_op_binder<Attribute>(&attr).match(op))
       return false;
     auto type = op->getResult(0).getType();
 
-    if (type.isSignlessIntOrIndex()) {
+    if (type.isa<IntegerType>() || type.isa<IndexType>())
       return attr_value_binder<IntegerAttr>(bind_value).match(attr);
-    }
     if (type.isa<VectorType>() || type.isa<RankedTensorType>()) {
       if (auto splatAttr = attr.dyn_cast<SplatElementsAttr>()) {
         return attr_value_binder<IntegerAttr>(bind_value)
             .match(splatAttr.getSplatValue());
       }
     }
     return false;
   }

mlir/include/mlir/IR/OpBase.td

AnyTypeOf<!foreach(w, widths, I<w>),
               StrJoinInt<widths, "/">.result # "-bit integer">;
 
 def I1  : I<1>;
 def I8  : I<8>;
 def I16 : I<16>;
 def I32 : I<32>;
 def I64 : I<64>;
 
+// Unsigned integer types.
+// Any unsigned integer type irrespective of its width.
+def AnyUnsignedInteger : Type<
+  CPred<"$_self.isUnsignedInteger()">, "unsigned integer">;
+
+// Unsigned integer type of a specific width.
+class UI<int width>
+    : Type<CPred<"$_self.isUnsignedInteger(" # width # ")">,
+                  width # "-bit unsigned integer">,
+      BuildableType<"$_builder.getIntegerType(" # width #
+                    ", /*isSigned=*/false)"> {
+  int bitwidth = width;
+}
+
+class UnsignedIntOfWidths<list<int> widths> :
+    AnyTypeOf<!foreach(w, widths, UI<w>),
+              StrJoinInt<widths, "/">.result # "-bit unsigned integer">;
+
+def UI1  : UI<1>;
+def UI8  : UI<8>;
+def UI16 : UI<16>;
+def UI32 : UI<32>;
+def UI64 : UI<64>;
+
 // Floating point types.
 
 // Any float type irrespective of its width.
 def AnyFloat : Type<CPred<"$_self.isa<FloatType>()">, "floating-point">;
 
 // Float type of a specific width.
 class F<int width>
     : Type<CPred<"$_self.isF" # width # "()">,

mlir/include/mlir/IR/StandardTypes.h

 public:
   using ShapedType::ShapedType;
 
   /// Return true if the specified element type is ok in a tensor.
   static bool isValidElementType(Type type) {
     // Note: Non standard/builtin types are allowed to exist within tensor
     // types. Dialects are expected to verify that tensor types have a valid
     // element type within that dialect.
-    return type.isSignlessIntOrFloat() || type.isa<ComplexType>() ||
-           type.isa<VectorType>() || type.isa<OpaqueType>() ||
+    return type.isa<ComplexType>() || type.isa<FloatType>() ||
+           type.isa<IntegerType>() || type.isa<OpaqueType>() ||
+           type.isa<VectorType>() ||
            (type.getKind() > Type::Kind::LAST_STANDARD_TYPE);
   }
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type) {
     return type.getKind() == StandardTypes::RankedTensor ||
            type.getKind() == StandardTypes::UnrankedTensor;
   }

mlir/include/mlir/IR/Types.h

#include "DialectSymbolRegistry.def"
 
   /// Return true if this is a signless integer or index type.
   bool isSignlessIntOrIndex();
   /// Return true if this is a signless integer, index, or float type.
   bool isSignlessIntOrIndexOrFloat();
   /// Return true of this is a signless integer or a float type.
   bool isSignlessIntOrFloat();
 
+  /// Return true of this is an integer(of any signedness) or a float type.
+  bool isIntOrFloat();
+
   /// Print the current type.
   void print(raw_ostream &os);
   void dump();
 
   friend ::llvm::hash_code hash_value(Type arg);
 
   unsigned getSubclassData() const;
   void setSubclassData(unsigned val);

mlir/lib/Analysis/Utils.cpp

LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   return success();
 }
 
 //  TODO(mlir-team): improve/complete this when we have target data.
 static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
   auto elementType = memRefType.getElementType();
 
   unsigned sizeInBits;
-  if (elementType.isSignlessIntOrFloat()) {
+  if (elementType.isIntOrFloat()) {
     sizeInBits = elementType.getIntOrFloatBitWidth();
   } else {
     auto vectorType = elementType.cast<VectorType>();
     sizeInBits =
         vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
   }
   return llvm::divideCeil(sizeInBits, 8);
 }

 /// Returns the size of memref data in bytes if it's statically shaped, None
 /// otherwise.  If the element of the memref has vector type, takes into account
 /// size of the vector as well.
 //  TODO(mlir-team): improve/complete this when we have target data.
 Optional<uint64_t> mlir::getMemRefSizeInBytes(MemRefType memRefType) {
   if (!memRefType.hasStaticShape())
     return None;
   auto elementType = memRefType.getElementType();
-  if (!elementType.isSignlessIntOrFloat() && !elementType.isa<VectorType>())
+  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>())
     return None;
 
   uint64_t sizeInBytes = getMemRefEltSizeInBytes(memRefType);
   for (unsigned i = 0, e = memRefType.getRank(); i < e; i++) {
     sizeInBytes = sizeInBytes * memRefType.getDimSize(i);
   }
   return sizeInBytes;
 }

mlir/lib/IR/AsmPrinter.cpp

void ModulePrinter::printAttribute(Attribute attr,
   if (typeElision != AttrTypeElision::Must && !attrType.isa<NoneType>()) {
     os << " : ";
     printType(attrType);
   }
 }
 
 /// Print the integer element of the given DenseElementsAttr at 'index'.
 static void printDenseIntElement(DenseElementsAttr attr, raw_ostream &os,
-                                 unsigned index) {
+                                 unsigned index, bool isSigned) {
   APInt value = *std::next(attr.int_value_begin(), index);
   if (value.getBitWidth() == 1)
     os << (value.getBoolValue() ? "true" : "false");
   else
-    value.print(os, /*isSigned=*/true);
+    value.print(os, isSigned);
 }
 
 /// Print the float element of the given DenseElementsAttr at 'index'.
 static void printDenseFloatElement(DenseElementsAttr attr, raw_ostream &os,
-                                   unsigned index) {
+                                   unsigned index, bool isSigned) {
+  assert(isSigned && "floating point values are always signed");
   APFloat value = *std::next(attr.float_value_begin(), index);
   printFloatValue(value, os);
 }
 
 void ModulePrinter::printDenseElementsAttr(DenseElementsAttr attr,
                                            bool allowHex) {
   auto type = attr.getType();
   auto shape = type.getShape();
   auto rank = type.getRank();
+  bool isSigned = !type.getElementType().isUnsignedInteger();
 
   // The function used to print elements of this attribute.
   auto printEltFn = type.getElementType().isa<IntegerType>()
                         ? printDenseIntElement
                         : printDenseFloatElement;
 
   // Special case for 0-d and splat tensors.
   if (attr.isSplat()) {
-    printEltFn(attr, os, 0);
+    printEltFn(attr, os, 0, isSigned);
     return;
   }
 
   // Special case for degenerate tensors.
   auto numElements = type.getNumElements();
   if (numElements == 0) {
     for (int i = 0; i < rank; ++i)
       os << '[';
void ModulePrinter::printDenseElementsAttr(DenseElementsAttr attr,
   };
 
   for (unsigned idx = 0, e = numElements; idx != e; ++idx) {
     if (idx != 0)
       os << ", ";
     while (openBrackets++ < rank)
       os << '[';
     openBrackets = rank;
-    printEltFn(attr, os, idx);
+    printEltFn(attr, os, idx, isSigned);
     bumpCounter();
   }
   while (openBrackets-- > 0)
     os << ']';
 }
 
 void ModulePrinter::printType(Type type) {
   // Check for an alias for this type.

mlir/lib/IR/Attributes.cpp

: llvm::mapped_iterator<IntElementIterator,
           it, [&](const APInt &val) { return APFloat(smt, val); }) {}
 
 //===----------------------------------------------------------------------===//
 // DenseElementsAttr
 //===----------------------------------------------------------------------===//
 
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<Attribute> values) {
-  assert(type.getElementType().isSignlessIntOrFloat() &&
+  assert(type.getElementType().isIntOrFloat() &&
          "expected int or float element type");
   assert(hasSameElementsOrSplat(type, values));
 
   auto eltType = type.getElementType();
   size_t bitWidth = getDenseElementBitwidth(eltType);
   size_t storageBitWidth = getDenseElementStorageWidth(bitWidth);
 
   // Compress the attribute values into a character buffer.

mlir/lib/IR/StandardTypes.cpp

 bool Type::isSignlessIntOrIndexOrFloat() {
   return isa<IndexType>() || isSignlessInteger() || isa<FloatType>();
 }
 
 bool Type::isSignlessIntOrFloat() {
   return isSignlessInteger() || isa<FloatType>();
 }
 
+bool Type::isIntOrFloat() { return isa<IntegerType>() || isa<FloatType>(); }
+
 //===----------------------------------------------------------------------===//
 // Integer Type
 //===----------------------------------------------------------------------===//
 
 // static constexpr must have a definition (until in C++17 and inline variable).
 constexpr unsigned IntegerType::kMaxWidth;
 
 /// Verify the construction of an integer type.
const llvm::fltSemantics &FloatType::getFloatSemantics() {
   if (isF32())
     return APFloat::IEEEsingle();
   if (isF64())
     return APFloat::IEEEdouble();
   llvm_unreachable("non-floating point type used");
 }
 
 unsigned Type::getIntOrFloatBitWidth() {
-  assert(isSignlessIntOrFloat() && "only ints and floats have a bitwidth");
-  if (auto intType = dyn_cast<IntegerType>()) {
+  assert(isIntOrFloat() && "only integers and floats have a bitwidth");
+  if (auto intType = dyn_cast<IntegerType>())
     return intType.getWidth();
-  }
-
-  auto floatType = cast<FloatType>();
-  return floatType.getWidth();
+  return cast<FloatType>().getWidth();
 }
 
 //===----------------------------------------------------------------------===//
 // ShapedType
 //===----------------------------------------------------------------------===//
 constexpr int64_t ShapedType::kDynamicSize;
 constexpr int64_t ShapedType::kDynamicStrideOrOffset;
 

 /// Get the number of bits require to store a value of the given shaped type.
 /// Compute the value recursively since tensors are allowed to have vectors as
 /// elements.
 int64_t ShapedType::getSizeInBits() const {
   assert(hasStaticShape() &&
          "cannot get the bit size of an aggregate with a dynamic shape");
 
   auto elementType = getElementType();
-  if (elementType.isSignlessIntOrFloat())
+  if (elementType.isIntOrFloat())
     return elementType.getIntOrFloatBitWidth() * getNumElements();
 
   // Tensors can have vectors and other tensors as elements, other shaped types
   // cannot.
   assert(isa<TensorType>() && "unsupported element type");
   assert((elementType.isa<VectorType>() || elementType.isa<TensorType>()) &&
          "unsupported tensor element type");
   return getNumElements() * elementType.cast<ShapedType>().getSizeInBits();

 /// pass in an instance of UnknownLoc.
 MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
                                unsigned memorySpace,
                                Optional<Location> location) {
   auto *context = elementType.getContext();
 
   // Check that memref is formed from allowed types.
-  if (!elementType.isSignlessIntOrFloat() && !elementType.isa<VectorType>() &&
+  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
       !elementType.isa<ComplexType>())
     return emitOptionalError(location, "invalid memref element type"),
            MemRefType();
 
   for (int64_t s : shape) {
     // Negative sizes are not allowed except for `-1` that means dynamic size.
     if (s < -1)
       return emitOptionalError(location, "invalid memref size"), MemRefType();

 unsigned UnrankedMemRefType::getMemorySpace() const {
   return getImpl()->memorySpace;
 }
 
 LogicalResult
 UnrankedMemRefType::verifyConstructionInvariants(Location loc, Type elementType,
                                                  unsigned memorySpace) {
   // Check that memref is formed from allowed types.
-  if (!elementType.isSignlessIntOrFloat() && !elementType.isa<VectorType>() &&
+  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
       !elementType.isa<ComplexType>())
     return emitError(loc, "invalid memref element type");
   return success();
 }
 
 /// Given MemRef `sizes` that are either static or dynamic, returns the
 /// canonical "contiguous" strides AffineExpr. Strides are multiplicative and
 /// once a dynamic dimension is encountered, all canonical strides become

mlir/lib/Parser/Parser.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.isSignlessIntOrFloat() && !elementType.isa<VectorType>() &&
+  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
       !elementType.isa<ComplexType>())
     return emitError(typeLoc, "invalid memref element type"), nullptr;
 
   // Parse semi-affine-map-composition.
   SmallVector<AffineMap, 2> affineMapComposition;
   Optional<unsigned> memorySpace;
   unsigned numDims = dimensions.size();
 

mlir/lib/Transforms/LoopFusion.cpp

static void sinkSequentialLoops(MemRefDependenceGraph::Node *node) {
   node->op = newRootForOp.getOperation();
 }
 
 //  TODO(mlir-team): improve/complete this when we have target data.
 static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
   auto elementType = memRefType.getElementType();
 
   unsigned sizeInBits;
-  if (elementType.isSignlessIntOrFloat()) {
+  if (elementType.isIntOrFloat()) {
     sizeInBits = elementType.getIntOrFloatBitWidth();
   } else {
     auto vectorType = elementType.cast<VectorType>();
     sizeInBits =
         vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
   }
   return llvm::divideCeil(sizeInBits, 8);
 }

mlir/test/IR/parser.mlir

 }
 
 // CHECK-LABEL: func @splattensorattr
 func @splattensorattr() -> () {
 ^bb0:
   // CHECK: "splatBoolTensor"() {bar = dense<false> : tensor<i1>} : () -> ()
   "splatBoolTensor"(){bar = dense<false> : tensor<i1>} : () -> ()
 
+  // CHECK: "splatUIntTensor"() {bar = dense<222> : tensor<2x1x4xui8>} : () -> ()
+  "splatUIntTensor"(){bar = dense<222> : tensor<2x1x4xui8>} : () -> ()
+
   // CHECK: "splatIntTensor"() {bar = dense<5> : tensor<2x1x4xi32>} : () -> ()
   "splatIntTensor"(){bar = dense<5> : tensor<2x1x4xi32>} : () -> ()
 
   // CHECK: "splatFloatTensor"() {bar = dense<-5.000000e+00> : tensor<2x1x4xf32>} : () -> ()
   "splatFloatTensor"(){bar = dense<-5.0> : tensor<2x1x4xf32>} : () -> ()
 
   // CHECK: "splatIntVector"() {bar = dense<5> : vector<2x1x4xi64>} : () -> ()
   "splatIntVector"(){bar = dense<5> : vector<2x1x4xi64>} : () -> ()