Adjust OpVariable's users and Lifetime ptr operand storage classes (#2650)

While clang is generating alloca in private address space for OpenCL
following target's datalayout, some of the passes might generate it
in different generic AS and for this cases we have to create OpVariable
in function storage class and create a cast to generic. Additionally
if such alloca was an operand of lifetime instruction - then SPIR-V
without storage class adjustment validation for lifetime instruction
will be violated (pointer operand must be with Function storage class).

This became exposed after llvm/llvm-project#97306
which was fixing handling of byval pointer parameter of a function
during inlining, byval pointer for OpenCL will be generated in default
(under an option - generic) address space making the alloca to store
value of the function's operand also be generic.

Signed-off-by: Sidorov, Dmitry <[email protected]>

Author: MrSidims

lib/SPIRV/SPIRVWriter.cpp

@@ -2269,10 +2269,20 @@ LLVMToSPIRVBase::transValueWithoutDecoration(Value *V, SPIRVBasicBlock *BB,
                       BM->addInstTemplate(OpVariableLengthArrayINTEL,
                                           {Length->getId()}, BB, TranslatedTy));
     }
-    return mapValue(V, BM->addVariable(TranslatedTy, false,
-                                       spv::internal::LinkageTypeInternal,
-                                       nullptr, Alc->getName().str(),
-                                       StorageClassFunction, BB));
+    SPIRVType *VarTy =
+        V->getType()->getPointerAddressSpace() == SPIRAS_Generic
+            ? BM->addPointerType(StorageClassFunction,
+                                 TranslatedTy->getPointerElementType())
+            : TranslatedTy;
+    SPIRVValue *Var = BM->addVariable(
+        VarTy, false, spv::internal::LinkageTypeInternal, nullptr,
+        Alc->getName().str(), StorageClassFunction, BB);
+    if (V->getType()->getPointerAddressSpace() == SPIRAS_Generic) {
+      SPIRVValue *Cast =
+          BM->addUnaryInst(OpPtrCastToGeneric, TranslatedTy, Var, BB);
+      return mapValue(V, Cast);
+    }
+    return mapValue(V, Var);
   }
 
   if (auto *Switch = dyn_cast<SwitchInst>(V)) {
@@ -4538,7 +4548,21 @@ SPIRVValue *LLVMToSPIRVBase::transIntrinsicInst(IntrinsicInst *II,
     int64_t Size = dyn_cast<ConstantInt>(II->getOperand(0))->getSExtValue();
     if (Size == -1)
       Size = 0;
-    return BM->addLifetimeInst(OC, transValue(II->getOperand(1), BB), Size, BB);
+    Value *LLVMPtrOp = II->getOperand(1);
+    unsigned PtrAS = cast<PointerType>(LLVMPtrOp->getType())->getAddressSpace();
+    auto *PtrOp = transValue(LLVMPtrOp, BB);
+    if (PtrAS == SPIRAS_Private)
+      return BM->addLifetimeInst(OC, PtrOp, Size, BB);
+    // If pointer address space is Generic - cast to private first
+    BM->getErrorLog().checkError(
+        PtrAS == SPIRAS_Generic, SPIRVEC_InvalidInstruction, II,
+        "lifetime intrinsic pointer operand must be in private or generic AS");
+    auto *SrcTy = PtrOp->getType();
+    auto *DstTy = BM->addPointerType(StorageClassFunction,
+                                     SrcTy->getPointerElementType());
+    PtrOp = BM->addUnaryInst(OpGenericCastToPtr, DstTy, PtrOp, BB);
+    ValueMap[LLVMPtrOp] = PtrOp;
+    return BM->addLifetimeInst(OC, PtrOp, Size, BB);
   }
   // We don't want to mix translation of regular code and debug info, because
   // it creates a mess, therefore translation of debug intrinsics is

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -565,6 +565,8 @@ class SPIRVModuleImpl : public SPIRVModule {
   SPIRVTypeVoid *VoidTy;
   SmallDenseMap<unsigned, SPIRVTypeInt *, 4> IntTypeMap;
   SmallDenseMap<unsigned, SPIRVTypeFloat *, 4> FloatTypeMap;
+  SmallDenseMap<std::pair<unsigned, SPIRVType *>, SPIRVTypePointer *, 4>
+      PointerTypeMap;
   std::unordered_map<unsigned, SPIRVConstant *> LiteralMap;
   std::vector<SPIRVExtInst *> DebugInstVec;
   std::vector<SPIRVExtInst *> AuxDataInstVec;
@@ -1003,8 +1005,13 @@ SPIRVTypeFloat *SPIRVModuleImpl::addFloatType(unsigned BitWidth) {
 SPIRVTypePointer *
 SPIRVModuleImpl::addPointerType(SPIRVStorageClassKind StorageClass,
                                 SPIRVType *ElementType) {
-  return addType(
-      new SPIRVTypePointer(this, getId(), StorageClass, ElementType));
+  auto Desc = std::make_pair(StorageClass, ElementType);
+  auto Loc = PointerTypeMap.find(Desc);
+  if (Loc != PointerTypeMap.end())
+    return Loc->second;
+  auto *Ty = new SPIRVTypePointer(this, getId(), StorageClass, ElementType);
+  PointerTypeMap[Desc] = Ty;
+  return addType(Ty);
 }
 
 SPIRVTypeFunction *SPIRVModuleImpl::addFunctionType(

test/llvm-intrinsics/lifetime.ll

@@ -6,14 +6,47 @@
 ; RUN: llvm-spirv -r %t.spv -o %t.spv.bc
 ; RUN: llvm-dis < %t.spv.bc | FileCheck %s --check-prefix=CHECK-LLVM
 
-; CHECK-SPIRV: 3 LifetimeStart [[tmp:[0-9]+]] 0
-; CHECK-SPIRV: 3 LifetimeStop [[tmp]] 0
+; CHECK-SPIRV-DAG: Name [[#SimpleF:]] "lifetime_simple"
+; CHECK-SPIRV-DAG: Name [[#SizedF:]] "lifetime_sized"
+; CHECK-SPIRV-DAG: Name [[#GenericF:]] "lifetime_generic"
+; CHECK-SPIRV-DAG: TypeStruct [[#StructTy:]] [[#]]
+; CHECK-SPIRV-COUNT-1: TypePointer [[#PrivatePtrTy:]] [[#StructTy]]
 
+; CHECK-SPIRV: Function [[#]] [[#SimpleF:]]
+; CHECK-SPIRV: LifetimeStart [[#Tmp:]] 0
+; CHECK-SPIRV: LifetimeStop [[#Tmp]] 0
+
+; CHECK-SPIRV: Function [[#]] [[#SizedF:]]
+; CHECK-SPIRV: LifetimeStart [[#Tmp:]] 1
+; CHECK-SPIRV: Bitcast [[#]] [[#Cast:]]
+; CHECK-SPIRV: LifetimeStop [[#Cast]] 1
+
+; CHECK-SPIRV: Function [[#]] [[#GenericF:]]
+; CHECK-SPIRV: Variable [[#PrivatePtrTy]] [[#Var:]] 7
+; CHECK-SPIRV: PtrCastToGeneric [[#]] [[#Cast1:]] [[#Var]]
+; CHECK-SPIRV: Bitcast [[#]] [[#Cast2:]] [[#Cast1]]
+; CHECK-SPIRV: GenericCastToPtr [[#]] [[#Cast3:]] [[#Cast2]]
+; CHECK-SPIRV: LifetimeStart [[#Cast3]] 1
+; CHECK-SPIRV: GenericCastToPtr [[#]] [[#Cast4:]]
+; CHECK-SPIRV: LifetimeStop [[#Cast4]] 1
+
+; CHECK-LLVM-LABEL: lifetime_simple
 ; CHECK-LLVM: %[[tmp1:[0-9]+]] = bitcast ptr %{{[0-9]+}} to ptr
 ; CHECK-LLVM: call void @llvm.lifetime.start.p0(i64 -1, ptr %[[tmp1]])
 ; CHECK-LLVM: call void @llvm.lifetime.end.p0(i64 -1, ptr %[[tmp1]])
-; CHECK-LLVM: declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture)
-; CHECK-LLVM: declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)
+
+; CHECK-LLVM-LABEL: lifetime_sized
+; CHECK-LLVM: call void @llvm.lifetime.start.p0(i64 1, ptr %[[#]])
+; CHECK-LLVM: call void @llvm.lifetime.end.p0(i64 1, ptr %[[#]])
+
+; CHECK-LLVM-LABEL: lifetime_generic
+; CHECK-LLVM: %[[#Alloca:]] = alloca %class.anon
+; CHECK-LLVM: %[[#Cast1:]] = addrspacecast ptr %[[#Alloca]] to ptr addrspace(4)
+; CHECK-LLVM: %[[#Cast2:]] = bitcast ptr addrspace(4) %[[#Cast1]] to ptr addrspace(4)
+; CHECK-LLVM: %[[#Cast3:]] = addrspacecast ptr addrspace(4) %[[#Cast2:]] to ptr
+; CHECK-LLVM: call void @llvm.lifetime.start.p0(i64 1, ptr %[[#Cast3]])
+; CHECK-LLVM: %[[#Cast4:]] = addrspacecast ptr addrspace(4) %[[#]] to ptr
+; CHECK-LLVM: call void @llvm.lifetime.end.p0(i64 1, ptr %[[#Cast4]])
 
 ; ModuleID = 'main'
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
@@ -53,14 +86,7 @@ entry:
   ret void
 }
 
-; Function Attrs: inlinehint nounwind
-define internal spir_func void @foo(%class.anon* %this) #0 align 2 {
-entry:
-  %this.addr = alloca %class.anon*, align 8
-  store %class.anon* %this, %class.anon** %this.addr, align 8
-  %this1 = load %class.anon*, %class.anon** %this.addr, align 8
-  ret void
-}
+declare spir_func void @foo(%class.anon* %this) #0
 
 ; Function Attrs: nounwind
 declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) #0
@@ -71,6 +97,26 @@ declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture) #0
 ; Function Attrs: nounwind readnone
 declare spir_func i64 @_Z13get_global_idj(i32) #1
 
+define spir_kernel void @lifetime_generic() #0 !kernel_arg_addr_space !8 !kernel_arg_access_qual !8 !kernel_arg_type !8 !kernel_arg_base_type !8 !kernel_arg_type_qual !8 {
+entry:
+  %0 = alloca %class.anon, align 1, addrspace(4)
+  %1 = bitcast %class.anon addrspace(4)* %0 to i8 addrspace(4)*
+  call void @llvm.lifetime.start.p4i8(i64 1, i8 addrspace(4)* %1) #0
+  call spir_func void @boo(%class.anon addrspace(4)* %0)
+  %2 = bitcast %class.anon addrspace(4)* %0 to i8 addrspace(4)*
+  call void @llvm.lifetime.end.p4i8(i64 1, i8 addrspace(4)* %2) #0
+  ret void
+}
+
+declare spir_func void @boo(%class.anon addrspace(4)* %this) #0
+
+; Function Attrs: nounwind
+declare void @llvm.lifetime.start.p4i8(i64 immarg, i8 addrspace(4)* nocapture) #0
+
+; Function Attrs: nounwind
+declare void @llvm.lifetime.end.p4i8(i64 immarg, i8 addrspace(4)* nocapture) #0
+
+
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }