[mlir][spirv] Simplify CompositeType::getNumElements. NFC.

[kuhar]

Use a type switch to simplify the implementation.

[llvmbot]

@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir

Author: Jakub Kuderski (kuhar)

Changes

Use a type switch to simplify the implementation.

---

Full diff: https://github.com/llvm/llvm-project/pull/160202.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp (+7-19)

diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp
index 4733bfca8be21..7e9a80e7d73a1 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp
@@ -203,25 +203,13 @@ Type CompositeType::getElementType(unsigned index) const {
 }
 
 unsigned CompositeType::getNumElements() const {
-  if (auto arrayType = llvm::dyn_cast<ArrayType>(*this))
-    return arrayType.getNumElements();
-  if (auto matrixType = llvm::dyn_cast<MatrixType>(*this))
-    return matrixType.getNumColumns();
-  if (auto structType = llvm::dyn_cast<StructType>(*this))
-    return structType.getNumElements();
-  if (auto vectorType = llvm::dyn_cast<VectorType>(*this))
-    return vectorType.getNumElements();
-  if (auto tensorArmType = dyn_cast<TensorArmType>(*this))
-    return tensorArmType.getNumElements();
-  if (llvm::isa<CooperativeMatrixType>(*this)) {
-    llvm_unreachable(
-        "invalid to query number of elements of spirv Cooperative Matrix type");
-  }
-  if (llvm::isa<RuntimeArrayType>(*this)) {
-    llvm_unreachable(
-        "invalid to query number of elements of spirv::RuntimeArray type");
-  }
-  llvm_unreachable("invalid composite type");
+  return TypeSwitch<SPIRVType, unsigned>(*this)
+      .Case<ArrayType, StructType, TensorArmType, VectorType>(
+          [](auto type) { return type.getNumElements(); })
+      .Case<MatrixType>([](MatrixType type) { return type.getNumColumns(); })
+      .Default([](SPIRVType) -> unsigned {
+        llvm_unreachable("Invalid type for number of elements query");
+      });
 }
 
 bool CompositeType::hasCompileTimeKnownNumElements() const {