[MLIR][SPIRVToLLVM] SPIR-V types size in bytes function

[georgemitenkov]

Summary

Added getTypeNumBytes() function as a class member to several SPIR-V Types:

• Scalar
• Array
• Vector
• Pointer

This is a function that will expose getTypeNumBytes() from SPIRVLowering.cpp to other classes. Support of other SPIR-V types can be added on demand.

[georgemitenkov]

@antiagainst, @mravishankar I haven't linked getTypeNumBytes() from SPIRVLowering.cpp to this new implementation? Should this be done as well?

[antiagainst]

mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h

113

Nit: The Type in the method does not help much given this is a method for .. type. :)

Typically think the method name and variable name and other name as the first line of documentation. A good name can go a long way for helping others understanding the code without referencing to the documentation and others additional stuff. For example, NumBytes is good here given it's very clear the unit of the return number so one does not need to guess whether is in bytes or bits or whatever. What do you think about getSizeInBytes here?

Also we need to document it here. What if the type can have explicit layout? What if not? What does it mean to return None?

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

159

We need to have documentation on the ArrayType mentioning the padding behavior here.

300

Is this correct? Don't we need to multiply with the count?

640

Hmm, this is not correct actually. There is no defined size for SPIR-V pointer type... We should return None here actually. I think even for LLVM, the pointer type size isn't fixed; it depends on the data layout.

906

No need to have these elses given we return for each case.

[georgemitenkov]

Renamed getTypeNumBytes() to getSizeInBytes().
Added documentation for the function declaration and a comment on array strides.
Removed support of PointerType and fixed Vector size.

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

159

Sorry I meant this should be on the declaration side (.h); not the implementation side (.cpp). :)

[georgemitenkov]

Moved stride comment on SPIRV's array type getSizeInBytes() from .cpp to .h file.

mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h

public:
   using CapabilityArrayRefVector = SmallVectorImpl<ArrayRef<spirv::Capability>>;
 
   /// Appends to `capabilities` the capabilities needed for this type to appear
   /// in the given `storage` class. This method does not guarantee the
   /// uniqueness of capabilities; the same capability may be appended multiple
   /// times.
   void getCapabilities(CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);
+
+  /// Returns the size in bytes for each type. If no size can be calculated,
+  /// returns `llvm::None`. Note that if the type has explicit layout, it is
+  /// also taken into account in calculation.
+  Optional<int64_t> getSizeInBytes();
 };
 
 // SPIR-V scalar type: bool type, integer type, floating point type.
 class ScalarType : public SPIRVType {
 public:
   using SPIRVType::SPIRVType;
 
   static bool classof(Type type);
 
   /// Returns true if the given integer type is valid for the SPIR-V dialect.
   static bool isValid(FloatType);
   /// Returns true if the given float type is valid for the SPIR-V dialect.
   static bool isValid(IntegerType);
 
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);
+
+  Optional<int64_t> getSizeInBytes();
 };
 
 // SPIR-V composite type: VectorType, SPIR-V ArrayType, or SPIR-V StructType.
 class CompositeType : public SPIRVType {
 public:
   using SPIRVType::SPIRVType;
 
   static bool classof(Type type);
public:
   /// Return true if the number of elements is known at compile time and is not
   /// implementation dependent.
   bool hasCompileTimeKnownNumElements() const;
 
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);
+
+  Optional<int64_t> getSizeInBytes();
 };
 
 // SPIR-V array type
 class ArrayType : public Type::TypeBase<ArrayType, CompositeType,
                                         detail::ArrayTypeStorage> {
 public:
   using Base::Base;
 
public:
   /// Returns the array stride in bytes. 0 means no stride decorated on this
   /// type.
   unsigned getArrayStride() const;
 
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);
+
+  /// Returns the array size in bytes. Since array type may have an explicit
+  /// stride declaration (in bytes), we also include it in the calculation.
+  Optional<int64_t> getSizeInBytes();
 };
 
 // SPIR-V image type
 class ImageType
     : public Type::TypeBase<ImageType, SPIRVType, detail::ImageTypeStorage> {
 public:
   using Base::Base;
 

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

 }
 
 void ArrayType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities,
     Optional<StorageClass> storage) {
   getElementType().cast<SPIRVType>().getCapabilities(capabilities, storage);
 }
 
+Optional<int64_t> ArrayType::getSizeInBytes() {
+  auto elementType = getElementType().cast<SPIRVType>();
+  Optional<int64_t> size = elementType.getSizeInBytes();
+  if (!size)
+    return llvm::None;
+  return (*size + getArrayStride()) * getNumElements();
+}
+
 //===----------------------------------------------------------------------===//
 // CompositeType
 //===----------------------------------------------------------------------===//
 
 bool CompositeType::classof(Type type) {
   switch (type.getKind()) {
   case TypeKind::Array:
   case TypeKind::CooperativeMatrix:
case StandardTypes::Vector:
     cast<VectorType>().getElementType().cast<ScalarType>().getCapabilities(
         capabilities, storage);
     break;
   default:
     llvm_unreachable("invalid composite type");
   }
 }
 
+Optional<int64_t> CompositeType::getSizeInBytes() {
+  switch (getKind()) {
+  case spirv::TypeKind::Array:
+    return cast<ArrayType>().getSizeInBytes();
+  case spirv::TypeKind::Struct:
+    return cast<StructType>().getSizeInBytes();
+  case StandardTypes::Vector: {
+    auto elementSize =
+        cast<VectorType>().getElementType().cast<ScalarType>().getSizeInBytes();
+    if (!elementSize)
+      return llvm::None;
+    return *elementSize * cast<VectorType>().getNumElements();
+  }
+  default:
+    return llvm::None;
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // CooperativeMatrixType
 //===----------------------------------------------------------------------===//
 
 struct spirv::detail::CooperativeMatrixTypeStorage : public TypeStorage {
   using KeyTy = std::tuple<Type, Scope, unsigned, unsigned>;
 
   static CooperativeMatrixTypeStorage *
void PointerType::getCapabilities(
   getPointeeType().cast<SPIRVType>().getCapabilities(capabilities,
                                                      getStorageClass());
 
   if (auto scCaps = spirv::getCapabilities(getStorageClass()))
     capabilities.push_back(*scCaps);
 }
 
 //===----------------------------------------------------------------------===//
 // RuntimeArrayType
 //===----------------------------------------------------------------------===//
 
 struct spirv::detail::RuntimeArrayTypeStorage : public TypeStorage {
   using KeyTy = std::pair<Type, unsigned>;
 
   static RuntimeArrayTypeStorage *construct(TypeStorageAllocator &allocator,
                                             const KeyTy &key) {
     return new (allocator.allocate<RuntimeArrayTypeStorage>())
if (auto intType = dyn_cast<IntegerType>()) {
     default:
       llvm_unreachable("invalid bitwidth to getCapabilities");
     }
   }
 
 #undef WIDTH_CASE
 }
 
+Optional<int64_t> ScalarType::getSizeInBytes() {
+  auto bitWidth = getIntOrFloatBitWidth();
+  // According to the SPIR-V spec:
+  // "There is no physical size or bit pattern defined for values with boolean
+  // type. If they are stored (in conjunction with OpVariable), they can only
+  // be used with logical addressing operations, not physical, and only with
+  // non-externally visible shader Storage Classes: Workgroup, CrossWorkgroup,
+  // Private, Function, Input, and Output."
+  if (bitWidth == 1)
+    return llvm::None;
+  return bitWidth / 8;
+}
+
 //===----------------------------------------------------------------------===//
 // SPIRVType
 //===----------------------------------------------------------------------===//
 
 bool SPIRVType::classof(Type type) {
   // Allow SPIR-V dialect types
   if (type.getKind() >= Type::FIRST_SPIRV_TYPE &&
       type.getKind() <= TypeKind::LAST_SPIRV_TYPE)
if (auto scalarType = dyn_cast<ScalarType>()) {
     matrixType.getCapabilities(capabilities, storage);
   } else if (auto ptrType = dyn_cast<PointerType>()) {
     ptrType.getCapabilities(capabilities, storage);
   } else {
     llvm_unreachable("invalid SPIR-V Type to getCapabilities");
   }
 }
 
+Optional<int64_t> SPIRVType::getSizeInBytes() {
+  if (auto scalarType = dyn_cast<ScalarType>())
+    return scalarType.getSizeInBytes();
+  if (auto compositeType = dyn_cast<CompositeType>())
+    return compositeType.getSizeInBytes();
+  return llvm::None;
+}
+
 //===----------------------------------------------------------------------===//
 // StructType
 //===----------------------------------------------------------------------===//
 
 struct spirv::detail::StructTypeStorage : public TypeStorage {
   StructTypeStorage(
       unsigned numMembers, Type const *memberTypes,
       StructType::OffsetInfo const *layoutInfo, unsigned numMemberDecorations,