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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-ir

Author: None (arysef)

<details>
<summary>Changes</summary>

This change is to support target extension types in vectors. The change allows sized target extension types to opt-in to being a valid vector element. 
Allowing target extension types as vector elements will allow backends to use vector operations such as `insertelement` and `extractelement` on their target types with minimal changes. 

RFC: https://discourse.llvm.org/t/rfc-supporting-sized-target-extension-types-in-vector/86431

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


5 Files Affected:

- (modified) llvm/docs/LangRef.rst (+2-1) 
- (modified) llvm/include/llvm/IR/DerivedTypes.h (+2) 
- (modified) llvm/lib/IR/Type.cpp (+17-4) 
- (added) llvm/test/Verifier/target-ext-vector-invalid.ll (+8) 
- (added) llvm/test/Verifier/target-ext-vector.ll (+20) 


``````````diff
diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index 0958f6a4b729b..e3b7c223b9c8d 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -4438,7 +4438,8 @@ the type size is smaller than the type's store size.
       < vscale x <# elements> x <elementtype> > ; Scalable vector
 
 The number of elements is a constant integer value larger than 0;
-elementtype may be any integer, floating-point or pointer type. Vectors
+elementtype may be any integer, floating-point, pointer type, or a sized  
+target extension type that has the `CanBeVectorElement` property. Vectors
 of size zero are not allowed. For scalable vectors, the total number of
 elements is a constant multiple (called vscale) of the specified number
 of elements; vscale is a positive integer that is unknown at compile time
diff --git a/llvm/include/llvm/IR/DerivedTypes.h b/llvm/include/llvm/IR/DerivedTypes.h
index 4d6bb1cfe3069..fa62bc09b61a3 100644
--- a/llvm/include/llvm/IR/DerivedTypes.h
+++ b/llvm/include/llvm/IR/DerivedTypes.h
@@ -845,6 +845,8 @@ class TargetExtType : public Type {
     /// This type may be allocated on the stack, either as the allocated type
     /// of an alloca instruction or as a byval function parameter.
     CanBeLocal = 1U << 2,
+    // This type may be used as an element in a vector.
+    CanBeVectorElement = 1U << 3,
   };
 
   /// Returns true if the target extension type contains the given property.
diff --git a/llvm/lib/IR/Type.cpp b/llvm/lib/IR/Type.cpp
index 4b43b52014484..7858e24f4fce7 100644
--- a/llvm/lib/IR/Type.cpp
+++ b/llvm/lib/IR/Type.cpp
@@ -790,8 +790,13 @@ VectorType *VectorType::get(Type *ElementType, ElementCount EC) {
 }
 
 bool VectorType::isValidElementType(Type *ElemTy) {
-  return ElemTy->isIntegerTy() || ElemTy->isFloatingPointTy() ||
-         ElemTy->isPointerTy() || ElemTy->getTypeID() == TypedPointerTyID;
+  if (ElemTy->isIntegerTy() || ElemTy->isFloatingPointTy() ||
+      ElemTy->isPointerTy() || ElemTy->getTypeID() == TypedPointerTyID)
+    return true;
+  if (auto *TTy = dyn_cast<TargetExtType>(ElemTy))
+    return TTy->hasProperty(TargetExtType::CanBeVectorElement) &&
+           TTy->isSized();
+  return false;
 }
 
 //===----------------------------------------------------------------------===//
@@ -801,8 +806,9 @@ bool VectorType::isValidElementType(Type *ElemTy) {
 FixedVectorType *FixedVectorType::get(Type *ElementType, unsigned NumElts) {
   assert(NumElts > 0 && "#Elements of a VectorType must be greater than 0");
   assert(isValidElementType(ElementType) && "Element type of a VectorType must "
-                                            "be an integer, floating point, or "
-                                            "pointer type.");
+                                            "be an integer, floating point, "
+                                            "pointer type, or a valid target "
+                                            "extension type.");
 
   auto EC = ElementCount::getFixed(NumElts);
 
@@ -1037,6 +1043,13 @@ static TargetTypeInfo getTargetTypeInfo(const TargetExtType *Ty) {
                           TargetExtType::CanBeGlobal);
   }
 
+  // Type used to test vector element target extension property.
+  // Can be removed once a public target extension type uses CanBeVectorElement
+  if (Name == "llvm.test.vectorelement") {
+    return TargetTypeInfo(Type::getInt32Ty(C), TargetExtType::CanBeLocal,
+                          TargetExtType::CanBeVectorElement);
+  }
+
   return TargetTypeInfo(Type::getVoidTy(C));
 }
 
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..59e3e78276a06
--- /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.Image")>
+  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..43d8360f39da0
--- /dev/null
+++ b/llvm/test/Verifier/target-ext-vector.ll
@@ -0,0 +1,20 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=verify -S %s | FileCheck %s
+
+define <2 x target("llvm.test.vectorelement")> @vec_ops(<2 x target("llvm.test.vectorelement")> %x) {
+; CHECK-LABEL: define <2 x target("llvm.test.vectorelement")> @vec_ops(
+; CHECK-SAME: <2 x target("llvm.test.vectorelement")> [[X:%.*]]) {
+; CHECK-NEXT:    [[A:%.*]] = alloca <2 x target("llvm.test.vectorelement")>{{.*}}
+; CHECK-NEXT:    store <2 x target("llvm.test.vectorelement")> [[X]], ptr [[A]], {{.*}}
+; CHECK-NEXT:    [[LOAD:%.*]] = load <2 x target("llvm.test.vectorelement")>, ptr [[A]], {{.*}}
+; CHECK-NEXT:    [[ELT:%.*]] = extractelement <2 x target("llvm.test.vectorelement")> [[LOAD]], i64 0
+; CHECK-NEXT:    [[RES:%.*]] = insertelement <2 x target("llvm.test.vectorelement")> undef, target("llvm.test.vectorelement") [[ELT]], i64 1
+; CHECK-NEXT:    ret <2 x target("llvm.test.vectorelement")> [[RES]]
+;
+  %a = alloca <2 x target("llvm.test.vectorelement")>
+  store <2 x target("llvm.test.vectorelement")> %x, ptr %a
+  %load = load <2 x target("llvm.test.vectorelement")>, ptr %a
+  %elt = extractelement <2 x target("llvm.test.vectorelement")> %load, i64 0
+  %res = insertelement <2 x target("llvm.test.vectorelement")> undef, target("llvm.test.vectorelement") %elt, i64 1
+  ret <2 x target("llvm.test.vectorelement")> %res
+}

``````````

</details>


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