[mlir] Provide LLVMType::getPrimitiveSizeInBits

[ftynse]

Summary

This function is available on llvm::Type and has been used by some clients of
the LLVM dialect before the transition. Implement the MLIR counterpart.

[mehdi_amini]

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

110

can you add a parenthesis with examples?

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

64

It is a bit "unsafe" to have a catch all here: we can't easily check that we correctly covered all the "sized" types.
Can we list all the types here in a assert(isa<...>) ?

[ftynse]

Address review.

[schweitz]

Thanks for adding this.

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

 
 #ifndef MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 #define MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 
 #include "mlir/IR/Types.h"
 
 namespace llvm {
 class ElementCount;
+class TypeSize;
 } // namespace llvm
 
 namespace mlir {
 
 class DialectAsmParser;
 class DialectAsmPrinter;
 
 namespace LLVM {
static LLVMType getFromOpaquePointer(const void *ptr) {
     return LLVMType(static_cast<ImplType *>(const_cast<void *>(ptr)));
   }
 
   /// Support for isa/cast.
   static bool classof(Type type);
 
   LLVMDialect &getDialect();
 
+  /// Returns the size of a primitive type (including vectors) in bits, for
+  /// example, the size of !llvm.i16 is 16 and the size of !llvm.vec<4 x i16>
+  /// is 64. Returns 0 for non-primitive (aggregates such as struct) or types
+  /// that don't have a size (such as void).
+  llvm::TypeSize getPrimitiveSizeInBits();
+
   /// Floating-point type utilities.
   bool isBFloatTy() { return isa<LLVMBFloatType>(); }
   bool isHalfTy() { return isa<LLVMHalfType>(); }
   bool isFloatTy() { return isa<LLVMFloatType>(); }
   bool isDoubleTy() { return isa<LLVMDoubleType>(); }
   bool isFloatingPointTy() {
     return isa<LLVMHalfType>() || isa<LLVMBFloatType>() ||
            isa<LLVMFloatType>() || isa<LLVMDoubleType>();

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

 
 #include "TypeDetail.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/TypeSupport.h"
 
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/TypeSize.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 
 //===----------------------------------------------------------------------===//
 // LLVMType.
 //===----------------------------------------------------------------------===//
 
 bool LLVMType::classof(Type type) {
   return llvm::isa<LLVMDialect>(type.getDialect());
 }
 
 LLVMDialect &LLVMType::getDialect() {
   return static_cast<LLVMDialect &>(Type::getDialect());
 }
 
 //----------------------------------------------------------------------------//
+// Misc type utilities.
+
+llvm::TypeSize LLVMType::getPrimitiveSizeInBits() {
+  return llvm::TypeSwitch<LLVMType, llvm::TypeSize>(*this)
+      .Case<LLVMHalfType, LLVMBFloatType>(
+          [](LLVMType) { return llvm::TypeSize::Fixed(16); })
+      .Case<LLVMFloatType>([](LLVMType) { return llvm::TypeSize::Fixed(32); })
+      .Case<LLVMDoubleType, LLVMX86MMXType>(
+          [](LLVMType) { return llvm::TypeSize::Fixed(64); })
+      .Case<LLVMIntegerType>([](LLVMIntegerType intTy) {
+        return llvm::TypeSize::Fixed(intTy.getBitWidth());
+      })
+      .Case<LLVMX86FP80Type>([](LLVMType) { return llvm::TypeSize::Fixed(80); })
+      .Case<LLVMPPCFP128Type, LLVMFP128Type>(
+          [](LLVMType) { return llvm::TypeSize::Fixed(128); })
+      .Case<LLVMVectorType>([](LLVMVectorType t) {
+        llvm::TypeSize elementSize =
+            t.getElementType().getPrimitiveSizeInBits();
+        llvm::ElementCount elementCount = t.getElementCount();
+        assert(!elementSize.isScalable() &&
+               "vector type should have fixed-width elements");
+        return llvm::TypeSize(elementSize.getFixedSize() * elementCount.Min,
+                              elementCount.Scalable);
+      })
+      .Default([](LLVMType ty) {
+        assert((ty.isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
+                       LLVMTokenType, LLVMStructType, LLVMArrayType,
+                       LLVMPointerType, LLVMFunctionType>()) &&
+               "unexpected missing support for primitive type");
+        return llvm::TypeSize::Fixed(0);
+      });
+}
+
+//----------------------------------------------------------------------------//
 // Integer type utilities.
 
 bool LLVMType::isIntegerTy(unsigned bitwidth) {
   if (auto intType = dyn_cast<LLVMIntegerType>())
     return intType.getBitWidth() == bitwidth;
   return false;
 }
 unsigned LLVMType::getIntegerBitWidth() {