[IR] Add TargetExtType::CanBeLocal property (#99016)

Add a property to allow marking target extension types that cannot be
used in an alloca instruction or byval argument, similar to CanBeGlobal
for global variables.

---------

Co-authored-by: Nikita Popov <[email protected]>

Author: jayfoad

clang/test/CodeGen/amdgpu-barrier-type-debug-info.c

@@ -4,5 +4,5 @@
 // CHECK: name: "__amdgpu_named_workgroup_barrier_t",{{.*}}baseType: ![[BT:[0-9]+]]
 // CHECK: [[BT]] = !DIBasicType(name: "__amdgpu_named_workgroup_barrier_t", size: 128, encoding: DW_ATE_unsigned)
 void test_locals(void) {
-  __amdgpu_named_workgroup_barrier_t k0;
+  __amdgpu_named_workgroup_barrier_t *k0;
 }

llvm/include/llvm/IR/DerivedTypes.h

@@ -228,6 +228,8 @@ class StructType : public Type {
     SCDB_NotContainsScalableVector = 32,
     SCDB_ContainsNonGlobalTargetExtType = 64,
     SCDB_NotContainsNonGlobalTargetExtType = 128,
+    SCDB_ContainsNonLocalTargetExtType = 64,
+    SCDB_NotContainsNonLocalTargetExtType = 128,
   };
 
   /// For a named struct that actually has a name, this is a pointer to the
@@ -302,6 +304,12 @@ class StructType : public Type {
   containsNonGlobalTargetExtType(SmallPtrSetImpl<const Type *> &Visited) const;
   using Type::containsNonGlobalTargetExtType;
 
+  /// Return true if this type is or contains a target extension type that
+  /// disallows being used as a local.
+  bool
+  containsNonLocalTargetExtType(SmallPtrSetImpl<const Type *> &Visited) const;
+  using Type::containsNonLocalTargetExtType;
+
   /// Returns true if this struct contains homogeneous scalable vector types.
   /// Note that the definition of homogeneous scalable vector type is not
   /// recursive here. That means the following structure will return false
@@ -798,6 +806,9 @@ class TargetExtType : public Type {
     HasZeroInit = 1U << 0,
     /// This type may be used as the value type of a global variable.
     CanBeGlobal = 1U << 1,
+    /// 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,
   };
 
   /// Returns true if the target extension type contains the given property.

llvm/include/llvm/IR/Type.h

@@ -215,6 +215,12 @@ class Type {
   containsNonGlobalTargetExtType(SmallPtrSetImpl<const Type *> &Visited) const;
   bool containsNonGlobalTargetExtType() const;
 
+  /// Return true if this type is or contains a target extension type that
+  /// disallows being used as a local.
+  bool
+  containsNonLocalTargetExtType(SmallPtrSetImpl<const Type *> &Visited) const;
+  bool containsNonLocalTargetExtType() const;
+
   /// Return true if this is a FP type or a vector of FP.
   bool isFPOrFPVectorTy() const { return getScalarType()->isFloatingPointTy(); }
 

llvm/lib/IR/Type.cpp

@@ -88,6 +88,22 @@ bool Type::containsNonGlobalTargetExtType() const {
   return containsNonGlobalTargetExtType(Visited);
 }
 
+bool Type::containsNonLocalTargetExtType(
+    SmallPtrSetImpl<const Type *> &Visited) const {
+  if (const auto *ATy = dyn_cast<ArrayType>(this))
+    return ATy->getElementType()->containsNonLocalTargetExtType(Visited);
+  if (const auto *STy = dyn_cast<StructType>(this))
+    return STy->containsNonLocalTargetExtType(Visited);
+  if (auto *TT = dyn_cast<TargetExtType>(this))
+    return !TT->hasProperty(TargetExtType::CanBeLocal);
+  return false;
+}
+
+bool Type::containsNonLocalTargetExtType() const {
+  SmallPtrSet<const Type *, 4> Visited;
+  return containsNonLocalTargetExtType(Visited);
+}
+
 const fltSemantics &Type::getFltSemantics() const {
   switch (getTypeID()) {
   case HalfTyID: return APFloat::IEEEhalf();
@@ -469,6 +485,34 @@ bool StructType::containsNonGlobalTargetExtType(
   return false;
 }
 
+bool StructType::containsNonLocalTargetExtType(
+    SmallPtrSetImpl<const Type *> &Visited) const {
+  if ((getSubclassData() & SCDB_ContainsNonLocalTargetExtType) != 0)
+    return true;
+
+  if ((getSubclassData() & SCDB_NotContainsNonLocalTargetExtType) != 0)
+    return false;
+
+  if (!Visited.insert(this).second)
+    return false;
+
+  for (Type *Ty : elements()) {
+    if (Ty->containsNonLocalTargetExtType(Visited)) {
+      const_cast<StructType *>(this)->setSubclassData(
+          getSubclassData() | SCDB_ContainsNonLocalTargetExtType);
+      return true;
+    }
+  }
+
+  // For structures that are opaque, return false but do not set the
+  // SCDB_NotContainsNonLocalTargetExtType flag since it may gain non-local
+  // target extension types when it becomes non-opaque.
+  if (!isOpaque())
+    const_cast<StructType *>(this)->setSubclassData(
+        getSubclassData() | SCDB_NotContainsNonLocalTargetExtType);
+  return false;
+}
+
 bool StructType::containsHomogeneousScalableVectorTypes() const {
   if (getNumElements() <= 0 || !isa<ScalableVectorType>(elements().front()))
     return false;
@@ -922,15 +966,18 @@ static TargetTypeInfo getTargetTypeInfo(const TargetExtType *Ty) {
   LLVMContext &C = Ty->getContext();
   StringRef Name = Ty->getName();
   if (Name == "spirv.Image")
-    return TargetTypeInfo(PointerType::get(C, 0), TargetExtType::CanBeGlobal);
+    return TargetTypeInfo(PointerType::get(C, 0), TargetExtType::CanBeGlobal,
+                          TargetExtType::CanBeLocal);
   if (Name.starts_with("spirv."))
     return TargetTypeInfo(PointerType::get(C, 0), TargetExtType::HasZeroInit,
-                          TargetExtType::CanBeGlobal);
+                          TargetExtType::CanBeGlobal,
+                          TargetExtType::CanBeLocal);
 
   // Opaque types in the AArch64 name space.
   if (Name == "aarch64.svcount")
     return TargetTypeInfo(ScalableVectorType::get(Type::getInt1Ty(C), 16),
-                          TargetExtType::HasZeroInit);
+                          TargetExtType::HasZeroInit,
+                          TargetExtType::CanBeLocal);
 
   // RISC-V vector tuple type. The layout is represented as the type that needs
   // the same number of vector registers(VREGS) as this tuple type, represented
@@ -942,12 +989,14 @@ static TargetTypeInfo getTargetTypeInfo(const TargetExtType *Ty) {
                  RISCV::RVVBitsPerBlock / 8) *
         Ty->getIntParameter(0);
     return TargetTypeInfo(
-        ScalableVectorType::get(Type::getInt8Ty(C), TotalNumElts));
+        ScalableVectorType::get(Type::getInt8Ty(C), TotalNumElts),
+        TargetExtType::CanBeLocal);
   }
 
   // DirectX resources
   if (Name.starts_with("dx."))
-    return TargetTypeInfo(PointerType::get(C, 0), TargetExtType::CanBeGlobal);
+    return TargetTypeInfo(PointerType::get(C, 0), TargetExtType::CanBeGlobal,
+                          TargetExtType::CanBeLocal);
 
   // Opaque types in the AMDGPU name space.
   if (Name == "amdgcn.named.barrier") {

llvm/lib/IR/Verifier.cpp

@@ -2026,11 +2026,15 @@ void Verifier::verifyParameterAttrs(AttributeSet Attrs, Type *Ty,
             "huge alignment values are unsupported", V);
     }
     if (Attrs.hasAttribute(Attribute::ByVal)) {
+      Type *ByValTy = Attrs.getByValType();
       SmallPtrSet<Type *, 4> Visited;
-      Check(Attrs.getByValType()->isSized(&Visited),
+      Check(ByValTy->isSized(&Visited),
             "Attribute 'byval' does not support unsized types!", V);
-      Check(DL.getTypeAllocSize(Attrs.getByValType()).getKnownMinValue() <
-                (1ULL << 32),
+      // Check if it is or contains a target extension type that disallows being
+      // used on the stack.
+      Check(!ByValTy->containsNonLocalTargetExtType(),
+            "'byval' argument has illegal target extension type", V);
+      Check(DL.getTypeAllocSize(ByValTy).getKnownMinValue() < (1ULL << 32),
             "huge 'byval' arguments are unsupported", V);
     }
     if (Attrs.hasAttribute(Attribute::ByRef)) {
@@ -4323,9 +4327,13 @@ void Verifier::verifySwiftErrorValue(const Value *SwiftErrorVal) {
 }
 
 void Verifier::visitAllocaInst(AllocaInst &AI) {
+  Type *Ty = AI.getAllocatedType();
   SmallPtrSet<Type*, 4> Visited;
-  Check(AI.getAllocatedType()->isSized(&Visited),
-        "Cannot allocate unsized type", &AI);
+  Check(Ty->isSized(&Visited), "Cannot allocate unsized type", &AI);
+  // Check if it's a target extension type that disallows being used on the
+  // stack.
+  Check(!Ty->containsNonLocalTargetExtType(),
+        "Alloca has illegal target extension type", &AI);
   Check(AI.getArraySize()->getType()->isIntegerTy(),
         "Alloca array size must have integer type", &AI);
   if (MaybeAlign A = AI.getAlign()) {
@@ -4334,8 +4342,7 @@ void Verifier::visitAllocaInst(AllocaInst &AI) {
   }
 
   if (AI.isSwiftError()) {
-    Check(AI.getAllocatedType()->isPointerTy(),
-          "swifterror alloca must have pointer type", &AI);
+    Check(Ty->isPointerTy(), "swifterror alloca must have pointer type", &AI);
     Check(!AI.isArrayAllocation(),
           "swifterror alloca must not be array allocation", &AI);
     verifySwiftErrorValue(&AI);

llvm/test/Assembler/target-type-properties.ll

@@ -3,6 +3,10 @@
 ; RUN: not llvm-as < %t/global-var.ll -o /dev/null 2>&1 | FileCheck --check-prefix=CHECK-GLOBAL-VAR %s
 ; RUN: not llvm-as < %t/global-array.ll -o /dev/null 2>&1 | FileCheck --check-prefix=CHECK-GLOBAL-ARRAY %s
 ; RUN: not llvm-as < %t/global-struct.ll -o /dev/null 2>&1 | FileCheck --check-prefix=CHECK-GLOBAL-STRUCT %s
+; RUN: not llvm-as < %t/alloca.ll -o /dev/null 2>&1 | FileCheck --check-prefix=CHECK-ALLOCA %s
+; RUN: not llvm-as < %t/alloca-struct.ll -o /dev/null 2>&1 | FileCheck --check-prefix=CHECK-ALLOCA-STRUCT %s
+; RUN: not llvm-as < %t/byval.ll -o /dev/null 2>&1 | FileCheck --check-prefix=CHECK-BYVAL %s
+; RUN: not llvm-as < %t/byval-array.ll -o /dev/null 2>&1 | FileCheck --check-prefix=CHECK-BYVAL-ARRAY %s
 ; Check target extension type properties are verified in the assembler.
 
 ;--- zeroinit-error.ll
@@ -24,3 +28,25 @@ define void @foo() {
 ;--- global-struct.ll
 @global_struct = external global {target("unknown_target_type")}
 ; CHECK-GLOBAL-STRUCT: Global @global_struct has illegal target extension type
+
+;--- alloca.ll
+define void @foo() {
+  %val = alloca target("amdgcn.named.barrier", 0)
+; CHECK-ALLOCA: Alloca has illegal target extension type
+  ret void
+}
+
+;--- alloca-struct.ll
+define void @foo() {
+  %val = alloca {target("amdgcn.named.barrier", 0), target("amdgcn.named.barrier", 0)}
+; CHECK-ALLOCA-STRUCT: Alloca has illegal target extension type
+  ret void
+}
+
+;--- byval.ll
+declare void @foo(ptr byval(target("amdgcn.named.barrier", 0)))
+; CHECK-BYVAL: 'byval' argument has illegal target extension type
+
+;--- byval-array.ll
+declare void @foo(ptr byval([4 x target("amdgcn.named.barrier", 0)]))
+; CHECK-BYVAL-ARRAY: 'byval' argument has illegal target extension type