[llvm] [LLVM] Support target extension types in vectors (PR #140630)

[via llvm-commits]

https://github.com/arysef created https://github.com/llvm/llvm-project/pull/140630

Minimal changes to illustrate supporting Target Extension Types in Vectors for LLVM IR. 

RFC: 

>From 78a86cf8c1ca1bc3f0336ae0dcf71fab5ef508b9 Mon Sep 17 00:00:00 2001
From: Aryan Sefidi <asefidi at microsoft.com>
Date: Mon, 19 May 2025 21:41:01 +0000
Subject: [PATCH] Support target extension types in vectors

---
 llvm/lib/IR/Type.cpp                            |  3 ++-
 llvm/test/Verifier/target-ext-vector-invalid.ll |  8 ++++++++
 llvm/test/Verifier/target-ext-vector.ll         | 11 +++++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Verifier/target-ext-vector-invalid.ll
 create mode 100644 llvm/test/Verifier/target-ext-vector.ll

diff --git a/llvm/lib/IR/Type.cpp b/llvm/lib/IR/Type.cpp
index 4b43b52014484..852b4897ed013 100644
--- a/llvm/lib/IR/Type.cpp
+++ b/llvm/lib/IR/Type.cpp
@@ -791,7 +791,8 @@ VectorType *VectorType::get(Type *ElementType, ElementCount EC) {
 
 bool VectorType::isValidElementType(Type *ElemTy) {
   return ElemTy->isIntegerTy() || ElemTy->isFloatingPointTy() ||
-         ElemTy->isPointerTy() || ElemTy->getTypeID() == TypedPointerTyID;
+         ElemTy->isPointerTy() || ElemTy->getTypeID() == TypedPointerTyID || 
+         (ElemTy->isTargetExtTy() && ElemTy->isSized());
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/llvm/test/Verifier/target-ext-vector-invalid.ll b/llvm/test/Verifier/target-ext-vector-invalid.ll
new file mode 100644
index 0000000000000..2b6a785db986c
--- /dev/null
+++ b/llvm/test/Verifier/target-ext-vector-invalid.ll
@@ -0,0 +1,8 @@
+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+; CHECK: invalid vector element type
+
+define void @bad() {
+  %v = alloca <2 x target("spirv.IntegralConstant")>
+  ret void
+}
\ No newline at end of file
diff --git a/llvm/test/Verifier/target-ext-vector.ll b/llvm/test/Verifier/target-ext-vector.ll
new file mode 100644
index 0000000000000..a9129d4dfaa55
--- /dev/null
+++ b/llvm/test/Verifier/target-ext-vector.ll
@@ -0,0 +1,11 @@
+; RUN: llvm-as -o - %s | llvm-dis | FileCheck %s
+
+; CHECK-LABEL: @vec_ops(
+define <2 x target("spirv.Image")> @vec_ops(<2 x target("spirv.Image")> %x) {
+  %a = alloca <2 x target("spirv.Image")>
+  store <2 x target("spirv.Image")> %x, ptr %a
+  %load = load <2 x target("spirv.Image")>, ptr %a
+  %elt = extractelement <2 x target("spirv.Image")> %load, i64 0
+  %res = insertelement <2 x target("spirv.Image")> undef, target("spirv.Image") %elt, i64 1
+  ret <2 x target("spirv.Image")> %res
+}
\ No newline at end of file