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

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

Author: arysef

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

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.

llvm/lib/IR/Type.cpp

@@ -790,8 +790,12 @@ 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);
+  return false;
 }
 
 //===----------------------------------------------------------------------===//
@@ -801,8 +805,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);
 
@@ -968,7 +973,11 @@ struct TargetTypeInfo {
 
   template <typename... ArgTys>
   TargetTypeInfo(Type *LayoutType, ArgTys... Properties)
-      : LayoutType(LayoutType), Properties((0 | ... | Properties)) {}
+      : LayoutType(LayoutType), Properties((0 | ... | Properties)) {
+    assert((!(this->Properties & TargetExtType::CanBeVectorElement) ||
+            LayoutType->isSized()) &&
+           "Vector element type must be sized");
+  }
 };
 } // anonymous namespace
 
@@ -1037,6 +1046,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));
 }
 

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
+}

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")> poison, 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")> poison, target("llvm.test.vectorelement") %elt, i64 1
+  ret <2 x target("llvm.test.vectorelement")> %res
+}