diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h @@ -18,6 +18,7 @@ namespace llvm { class ElementCount; +class TypeSize; } // namespace llvm namespace mlir { @@ -105,6 +106,10 @@ LLVMDialect &getDialect(); + /// Returns the size of a primitive type (including vectors) in bits. Returns + /// 0 for non-primitive types or types that don't have a size. + llvm::TypeSize getPrimitiveSizeInBits(); + /// Floating-point type utilities. bool isBFloatTy() { return isa(); } bool isHalfTy() { return isa(); } diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp +++ b/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,34 @@ return static_cast(Type::getDialect()); } +//----------------------------------------------------------------------------// +// Misc type utilities. + +llvm::TypeSize LLVMType::getPrimitiveSizeInBits() { + return llvm::TypeSwitch(*this) + .Case( + [](LLVMType) { return llvm::TypeSize::Fixed(16); }) + .Case([](LLVMType) { return llvm::TypeSize::Fixed(32); }) + .Case( + [](LLVMType) { return llvm::TypeSize::Fixed(64); }) + .Case([](LLVMIntegerType intTy) { + return llvm::TypeSize::Fixed(intTy.getBitWidth()); + }) + .Case([](LLVMType) { return llvm::TypeSize::Fixed(80); }) + .Case( + [](LLVMType) { return llvm::TypeSize::Fixed(128); }) + .Case([](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) { return llvm::TypeSize::Fixed(0); }); +} + //----------------------------------------------------------------------------// // Integer type utilities.