[mlir] Provide LLVMType::getPrimitiveSizeInBits

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.

Reviewed By: schweitz

Differential Revision: https://reviews.llvm.org/D85847

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

@@ -18,6 +18,7 @@
 
 namespace llvm {
 class ElementCount;
+class TypeSize;
 } // namespace llvm
 
 namespace mlir {
@@ -105,6 +106,12 @@ public:
 
   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>(); }

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

@@ -18,6 +18,7 @@
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/TypeSupport.h"
 
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/TypeSize.h"
 
 using namespace mlir;
@@ -35,6 +36,40 @@ 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.