[mlir][spirv] Rework type capability queries (#160113)

* Fix infinite recursion with nested structs.
* Drop `::getCapbilities` function from derived types, so that there's
only one entry point that queries type extensions.
* Move all capability logic to a new helper class -- this way the
`::getCapabilities` functions can't diverge across concrete types and
'convenience types' like CompositeType.

Fixes: #159963

Author: kuhar

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

@@ -89,9 +89,6 @@ class ScalarType : public SPIRVType {
   /// Returns true if the given float type is valid for the SPIR-V dialect.
   static bool isValid(IntegerType);
 
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
-
   std::optional<int64_t> getSizeInBytes();
 };
 
@@ -116,9 +113,6 @@ class CompositeType : public SPIRVType {
   /// implementation dependent.
   bool hasCompileTimeKnownNumElements() const;
 
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
-
   std::optional<int64_t> getSizeInBytes();
 };
 
@@ -144,9 +138,6 @@ class ArrayType : public Type::TypeBase<ArrayType, CompositeType,
   /// type.
   unsigned getArrayStride() const;
 
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
-
   /// Returns the array size in bytes. Since array type may have an explicit
   /// stride declaration (in bytes), we also include it in the calculation.
   std::optional<int64_t> getSizeInBytes();
@@ -186,9 +177,6 @@ class ImageType
   ImageSamplerUseInfo getSamplerUseInfo() const;
   ImageFormat getImageFormat() const;
   // TODO: Add support for Access qualifier
-
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
 };
 
 // SPIR-V pointer type
@@ -204,9 +192,6 @@ class PointerType : public Type::TypeBase<PointerType, SPIRVType,
   Type getPointeeType() const;
 
   StorageClass getStorageClass() const;
-
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
 };
 
 // SPIR-V run-time array type
@@ -228,9 +213,6 @@ class RuntimeArrayType
   /// Returns the array stride in bytes. 0 means no stride decorated on this
   /// type.
   unsigned getArrayStride() const;
-
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
 };
 
 // SPIR-V sampled image type
@@ -252,10 +234,6 @@ class SampledImageType
                    Type imageType);
 
   Type getImageType() const;
-
-  void
-  getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                  std::optional<spirv::StorageClass> storage = std::nullopt);
 };
 
 /// SPIR-V struct type. Two kinds of struct types are supported:
@@ -405,9 +383,6 @@ class StructType
   trySetBody(ArrayRef<Type> memberTypes, ArrayRef<OffsetInfo> offsetInfo = {},
              ArrayRef<MemberDecorationInfo> memberDecorations = {},
              ArrayRef<StructDecorationInfo> structDecorations = {});
-
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
 };
 
 llvm::hash_code
@@ -440,9 +415,6 @@ class CooperativeMatrixType
   /// Returns the use parameter of the cooperative matrix.
   CooperativeMatrixUseKHR getUse() const;
 
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
-
   operator ShapedType() const { return llvm::cast<ShapedType>(*this); }
 
   ArrayRef<int64_t> getShape() const;
@@ -493,9 +465,6 @@ class MatrixType : public Type::TypeBase<MatrixType, CompositeType,
 
   /// Returns the elements' type (i.e, single element type).
   Type getElementType() const;
-
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
 };
 
 /// SPIR-V TensorARM Type
@@ -531,9 +500,6 @@ class TensorArmType
   ArrayRef<int64_t> getShape() const;
   bool hasRank() const { return !getShape().empty(); }
   operator ShapedType() const { return llvm::cast<ShapedType>(*this); }
-
-  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
-                       std::optional<StorageClass> storage = std::nullopt);
 };
 
 } // namespace spirv