[Mlir-commits] [mlir] 5d075be - [MLIR][SPIRVToLLVM] SPIR-V types size in bytes function

[George Mitenkov]

Author: George Mitenkov
Date: 2020-07-09T17:43:48+03:00
New Revision: 5d075beae73f3d9570230c0780e74dc32dda39a9

URL: https://github.com/llvm/llvm-project/commit/5d075beae73f3d9570230c0780e74dc32dda39a9
DIFF: https://github.com/llvm/llvm-project/commit/5d075beae73f3d9570230c0780e74dc32dda39a9.diff

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

Added `getSizeInBytes()` function as a class member to several SPIR-V types:
`ScalarType`, `ArrayType` and `VectorType`. This function aims at exposing
the functionality of `getTypeNumBytes()` from `SPIRVLowering.cpp`. Support
of more types will be added on demand.

Reviewed By: antiagainst

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h
    mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h b/mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h
index 91ad96fed7b0..b13d3c7892be 100644
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVTypes.h
@@ -109,6 +109,11 @@ class SPIRVType : public Type {
   /// 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.
@@ -127,6 +132,8 @@ class ScalarType : public SPIRVType {
                      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.
@@ -153,6 +160,8 @@ class CompositeType : public SPIRVType {
                      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
@@ -181,6 +190,10 @@ class ArrayType : public Type::TypeBase<ArrayType, CompositeType,
                      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

diff  --git a/mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp b/mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp
index 07ffe4046483..93d0c43d669f 100644
--- a/mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp
+++ b/mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp
@@ -151,6 +151,14 @@ void ArrayType::getCapabilities(
   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
 //===----------------------------------------------------------------------===//
@@ -281,6 +289,24 @@ void CompositeType::getCapabilities(
   }
 }
 
+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
 //===----------------------------------------------------------------------===//
@@ -806,6 +832,19 @@ void ScalarType::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
 //===----------------------------------------------------------------------===//
@@ -861,6 +900,14 @@ void SPIRVType::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
 //===----------------------------------------------------------------------===//