[DirectX] Update lifetime legalization to account for the removed size argument

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -746,7 +746,7 @@ class OpLowerer {
     IRBuilder<> &IRB = OpBuilder.getIRB();
     return replaceFunction(F, [&](CallInst *CI) -> Error {
       IRB.SetInsertPoint(CI);
-      Value *Ptr = CI->getArgOperand(1);
+      Value *Ptr = CI->getArgOperand(0);
       assert(Ptr->getType()->isPointerTy() &&
              "Expected operand of lifetime intrinsic to be a pointer");
 

llvm/lib/Target/DirectX/DXILWriter/DXILWriterPass.cpp

@@ -13,11 +13,15 @@
 #include "DXILWriterPass.h"
 #include "DXILBitcodeWriter.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/ModuleSummaryAnalysis.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
@@ -54,49 +58,81 @@ class WriteDXILPass : public llvm::ModulePass {
 };
 
 static void legalizeLifetimeIntrinsics(Module &M) {
-  for (Function &F : M) {
-    Intrinsic::ID IID = F.getIntrinsicID();
-    if (IID != Intrinsic::lifetime_start && IID != Intrinsic::lifetime_end)
+  LLVMContext &Ctx = M.getContext();
+  Type *I64Ty = IntegerType::get(Ctx, 64);
+  Type *PtrTy = PointerType::get(Ctx, 0);
+  Intrinsic::ID LifetimeIIDs[2] = {Intrinsic::lifetime_start,
+                                   Intrinsic::lifetime_end};
+  for (Intrinsic::ID &IID : LifetimeIIDs) {
+    Function *F = M.getFunction(Intrinsic::getName(IID, {PtrTy}, &M));
+    if (!F)
       continue;
 
-    // Lifetime intrinsics in LLVM 3.7 do not have the memory FnAttr
-    F.removeFnAttr(Attribute::Memory);
-
-    // Lifetime intrinsics in LLVM 3.7 do not have mangled names
-    F.setName(Intrinsic::getBaseName(IID));
-
-    // LLVM 3.7 Lifetime intrinics require an i8* operand, so we insert bitcasts
-    // to ensure that is the case
-    for (auto *User : make_early_inc_range(F.users())) {
-      CallInst *CI = dyn_cast<CallInst>(User);
-      assert(CI && "Expected user of a lifetime intrinsic function to be a "
-                   "lifetime intrinsic call");
-      Value *PtrOperand = CI->getArgOperand(1);
-      PointerType *PtrTy = cast<PointerType>(PtrOperand->getType());
+    // Get or insert an LLVM 3.7-compliant lifetime intrinsic function of the
+    // form `void @llvm.lifetime.[start/end](i64, ptr)` with the NoUnwind
+    // attribute
+    AttributeList Attr;
+    Attr = Attr.addFnAttribute(Ctx, Attribute::NoUnwind);
+    FunctionCallee LifetimeCallee = M.getOrInsertFunction(
+        Intrinsic::getBaseName(IID), Attr, Type::getVoidTy(Ctx), I64Ty, PtrTy);
+
+    // Replace all calls to lifetime intrinsics with calls to the
+    // LLVM 3.7-compliant version of the lifetime intrinsic
+    for (User *U : make_early_inc_range(F->users())) {
+      CallInst *CI = dyn_cast<CallInst>(U);
+      assert(CI &&
+             "Expected user of a lifetime intrinsic function to be a CallInst");
+
+      // LLVM 3.7 lifetime intrinics require an i8* operand, so we insert
+      // a bitcast to ensure that is the case
+      Value *PtrOperand = CI->getArgOperand(0);
+      PointerType *PtrOpPtrTy = cast<PointerType>(PtrOperand->getType());
       Value *NoOpBitCast = CastInst::Create(Instruction::BitCast, PtrOperand,
-                                            PtrTy, "", CI->getIterator());
-      CI->setArgOperand(1, NoOpBitCast);
+                                            PtrOpPtrTy, "", CI->getIterator());
+
+      // LLVM 3.7 lifetime intrinsics have an explicit size operand, whose value
+      // we can obtain from the pointer operand which must be an AllocaInst (as
+      // of https://github.com/llvm/llvm-project/pull/149310)
+      AllocaInst *AI = dyn_cast<AllocaInst>(PtrOperand);
+      assert(AI &&
+             "The pointer operand of a lifetime intrinsic call must be an "
+             "AllocaInst");
+      std::optional<TypeSize> AllocSize =
+          AI->getAllocationSize(CI->getDataLayout());
+      assert(AllocSize.has_value() &&
+             "Expected the allocation size of AllocaInst to be known");
+      CallInst *NewCI = CallInst::Create(
+          LifetimeCallee,
+          {ConstantInt::get(I64Ty, AllocSize.value().getFixedValue()),
+           NoOpBitCast},
+          "", CI->getIterator());
+      for (Attribute ParamAttr : CI->getParamAttributes(0))
+        NewCI->addParamAttr(1, ParamAttr);
+
+      CI->eraseFromParent();
     }
+
+    F->eraseFromParent();
   }
 }
 
 static void removeLifetimeIntrinsics(Module &M) {
-  for (Function &F : make_early_inc_range(M)) {
-    if (Intrinsic::ID IID = F.getIntrinsicID();
-        IID != Intrinsic::lifetime_start && IID != Intrinsic::lifetime_end)
+  Intrinsic::ID LifetimeIIDs[2] = {Intrinsic::lifetime_start,
+                                   Intrinsic::lifetime_end};
+  for (Intrinsic::ID &IID : LifetimeIIDs) {
+    Function *F = M.getFunction(Intrinsic::getBaseName(IID));
+    if (!F)
       continue;
 
-    for (User *U : make_early_inc_range(F.users())) {
-      LifetimeIntrinsic *LI = dyn_cast<LifetimeIntrinsic>(U);
-      assert(LI && "Expected user of lifetime intrinsic function to be "
-                   "a LifetimeIntrinsic instruction");
-      BitCastInst *BCI = dyn_cast<BitCastInst>(LI->getArgOperand(1));
-      assert(BCI && "Expected pointer operand of LifetimeIntrinsic to be a "
-                    "BitCastInst");
-      LI->eraseFromParent();
+    for (User *U : make_early_inc_range(F->users())) {
+      CallInst *CI = dyn_cast<CallInst>(U);
+      assert(CI && "Expected user of lifetime function to be a CallInst");
+      BitCastInst *BCI = dyn_cast<BitCastInst>(CI->getArgOperand(1));
+      assert(BCI && "Expected pointer operand of CallInst to be a BitCastInst");
+      CI->eraseFromParent();
       BCI->eraseFromParent();
     }
-    F.eraseFromParent();
+    F->eraseFromParent();
   }
 }
 

llvm/test/CodeGen/DirectX/ShaderFlags/lifetimes-noint64op.ll

@@ -15,16 +15,16 @@ target triple = "dxil-pc-shadermodel6.7-library"
 
 define void @lifetimes() #0 {
   %a = alloca [4 x i32], align 8
-  call void @llvm.lifetime.start.p0(i64 16, ptr nonnull %a)
-  call void @llvm.lifetime.end.p0(i64 16, ptr nonnull %a)
+  call void @llvm.lifetime.start.p0(ptr nonnull %a)
+  call void @llvm.lifetime.end.p0(ptr nonnull %a)
   ret void
 }
 
 ; Function Attrs: nounwind memory(argmem: readwrite)
-declare void @llvm.lifetime.start.p0(i64, ptr) #1
+declare void @llvm.lifetime.start.p0(ptr) #1
 
 ; Function Attrs: nounwind memory(argmem: readwrite)
-declare void @llvm.lifetime.end.p0(i64, ptr) #1
+declare void @llvm.lifetime.end.p0(ptr) #1
 
 attributes #0 = { convergent norecurse nounwind "hlsl.export"}
 attributes #1 = { nounwind memory(argmem: readwrite) }

llvm/test/CodeGen/DirectX/legalize-lifetimes-valver-1.5.ll

@@ -11,9 +11,9 @@
 define void @test_legal_lifetime()  {
   %accum.i.flat = alloca [1 x i32], align 4
   %gep = getelementptr i32, ptr %accum.i.flat, i32 0
-  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull %accum.i.flat)
+  call void @llvm.lifetime.start.p0(ptr nonnull %accum.i.flat)
   store i32 0, ptr %gep, align 4
-  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %accum.i.flat)
+  call void @llvm.lifetime.end.p0(ptr nonnull %accum.i.flat)
   ret void
 }
 

llvm/test/CodeGen/DirectX/legalize-lifetimes-valver-1.6.ll

@@ -13,22 +13,22 @@
 ; CHECK-NEXT:       [[ACCUM_I_FLAT:%.*]] = alloca [1 x i32], align 4
 ; CHECK-NEXT:       [[GEP:%.*]] = getelementptr i32, ptr [[ACCUM_I_FLAT]], i32 0
 ; CHECK-SM63-NEXT:  store [1 x i32] undef, ptr [[ACCUM_I_FLAT]], align 4
-; CHECK-SM66-NEXT:  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
+; CHECK-SM66-NEXT:  call void @llvm.lifetime.start.p0(ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-EMBED-NOT:  bitcast
 ; CHECK-EMBED-NOT:  lifetime
 ; CHECK-NEXT:       store i32 0, ptr [[GEP]], align 4
 ; CHECK-SM63-NEXT:  store [1 x i32] undef, ptr [[ACCUM_I_FLAT]], align 4
-; CHECK-SM66-NEXT:  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
+; CHECK-SM66-NEXT:  call void @llvm.lifetime.end.p0(ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-EMBED-NOT:  bitcast
 ; CHECK-EMBED-NOT:  lifetime
 ; CHECK-NEXT:       ret void
 ;
 define void @test_legal_lifetime()  {
   %accum.i.flat = alloca [1 x i32], align 4
   %gep = getelementptr i32, ptr %accum.i.flat, i32 0
-  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull %accum.i.flat)
+  call void @llvm.lifetime.start.p0(ptr nonnull %accum.i.flat)
   store i32 0, ptr %gep, align 4
-  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %accum.i.flat)
+  call void @llvm.lifetime.end.p0(ptr nonnull %accum.i.flat)
   ret void
 }
 

llvm/test/CodeGen/DirectX/legalize-memset.ll

@@ -5,92 +5,72 @@ define void @replace_float_memset_test() #0 {
 ; CHECK-LABEL: define void @replace_float_memset_test(
 ; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
 ; CHECK-NEXT:    [[ACCUM_I_FLAT:%.*]] = alloca [2 x float], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 8, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [2 x float], ptr [[ACCUM_I_FLAT]], i32 0, i32 0
 ; CHECK-NEXT:    store float 0.000000e+00, ptr [[GEP]], align 4
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr [2 x float], ptr [[ACCUM_I_FLAT]], i32 0, i32 1
 ; CHECK-NEXT:    store float 0.000000e+00, ptr [[GEP1]], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 8, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    ret void
 ;
   %accum.i.flat = alloca [2 x float], align 4
-  call void @llvm.lifetime.start.p0(i64 8, ptr nonnull %accum.i.flat)
   call void @llvm.memset.p0.i32(ptr nonnull align 4 dereferenceable(8) %accum.i.flat, i8 0, i32 8, i1 false)
-  call void @llvm.lifetime.end.p0(i64 8, ptr nonnull %accum.i.flat)
   ret void
 }
 
 define void @replace_half_memset_test() #0 {
 ; CHECK-LABEL: define void @replace_half_memset_test(
 ; CHECK-SAME: ) #[[ATTR0]] {
 ; CHECK-NEXT:    [[ACCUM_I_FLAT:%.*]] = alloca [2 x half], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [2 x half], ptr [[ACCUM_I_FLAT]], i32 0, i32 0
 ; CHECK-NEXT:    store half 0xH0000, ptr [[GEP]], align 2
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr [2 x half], ptr [[ACCUM_I_FLAT]], i32 0, i32 1
 ; CHECK-NEXT:    store half 0xH0000, ptr [[GEP1]], align 2
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    ret void
 ;
   %accum.i.flat = alloca [2 x half], align 4
-  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull %accum.i.flat)
   call void @llvm.memset.p0.i32(ptr nonnull align 4 dereferenceable(8) %accum.i.flat, i8 0, i32 4, i1 false)
-  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %accum.i.flat)
   ret void
 }
 
 define void @replace_double_memset_test() #0 {
 ; CHECK-LABEL: define void @replace_double_memset_test(
 ; CHECK-SAME: ) #[[ATTR0]] {
 ; CHECK-NEXT:    [[ACCUM_I_FLAT:%.*]] = alloca [2 x double], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 16, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [2 x double], ptr [[ACCUM_I_FLAT]], i32 0, i32 0
 ; CHECK-NEXT:    store double 0.000000e+00, ptr [[GEP]], align 8
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr [2 x double], ptr [[ACCUM_I_FLAT]], i32 0, i32 1
 ; CHECK-NEXT:    store double 0.000000e+00, ptr [[GEP1]], align 8
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 16, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    ret void
 ;
   %accum.i.flat = alloca [2 x double], align 4
-  call void @llvm.lifetime.start.p0(i64 16, ptr nonnull %accum.i.flat)
   call void @llvm.memset.p0.i32(ptr nonnull align 4 dereferenceable(8) %accum.i.flat, i8 0, i32 16, i1 false)
-  call void @llvm.lifetime.end.p0(i64 16, ptr nonnull %accum.i.flat)
   ret void
 }
 
 define void @replace_int16_memset_test() #0 {
 ; CHECK-LABEL: define void @replace_int16_memset_test(
 ; CHECK-SAME: ) #[[ATTR0]] {
 ; CHECK-NEXT:    [[CACHE_I:%.*]] = alloca [2 x i16], align 2
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[CACHE_I]])
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [2 x i16], ptr [[CACHE_I]], i32 0, i32 0
 ; CHECK-NEXT:    store i16 0, ptr [[GEP]], align 2
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr [2 x i16], ptr [[CACHE_I]], i32 0, i32 1
 ; CHECK-NEXT:    store i16 0, ptr [[GEP1]], align 2
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr nonnull [[CACHE_I]])
 ; CHECK-NEXT:    ret void
 ;
   %cache.i = alloca [2 x i16], align 2
-  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull %cache.i)
   call void @llvm.memset.p0.i32(ptr nonnull align 2 dereferenceable(4) %cache.i, i8 0, i32 4, i1 false)
-  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %cache.i)
   ret void
 }
 
 define void @replace_int_memset_test() #0 {
 ; CHECK-LABEL: define void @replace_int_memset_test(
 ; CHECK-SAME: ) #[[ATTR0]] {
 ; CHECK-NEXT:    [[ACCUM_I_FLAT:%.*]] = alloca [1 x i32], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [1 x i32], ptr [[ACCUM_I_FLAT]], i32 0, i32 0
 ; CHECK-NEXT:    store i32 0, ptr [[GEP]], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    ret void
 ;
   %accum.i.flat = alloca [1 x i32], align 4
-  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull %accum.i.flat)
   call void @llvm.memset.p0.i32(ptr nonnull align 4 dereferenceable(8) %accum.i.flat, i8 0, i32 4, i1 false)
-  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %accum.i.flat)
   ret void
 }
 
@@ -101,25 +81,19 @@ define void @replace_int_memset_to_var_test() #0 {
 ; CHECK-NEXT:    [[I:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    store i32 1, ptr [[I]], align 4
 ; CHECK-NEXT:    [[I8_LOAD:%.*]] = load i32, ptr [[I]], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [1 x i32], ptr [[ACCUM_I_FLAT]], i32 0, i32 0
 ; CHECK-NEXT:    store i32 [[I8_LOAD]], ptr [[GEP]], align 4
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
 ; CHECK-NEXT:    ret void
 ;
   %accum.i.flat = alloca [1 x i32], align 4
   %i = alloca i8, align 4
   store i8 1, ptr %i
   %i8.load = load i8, ptr %i
-  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull %accum.i.flat)
   call void @llvm.memset.p0.i32(ptr nonnull align 4 dereferenceable(8) %accum.i.flat, i8 %i8.load, i32 4, i1 false)
-  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %accum.i.flat)
   ret void
 }
 
 attributes #0 = {"hlsl.export"}
 
 
-declare void @llvm.lifetime.end.p0(i64 immarg, ptr captures(none))
-declare void @llvm.lifetime.start.p0(i64 immarg, ptr captures(none))
 declare void @llvm.memset.p0.i32(ptr writeonly captures(none), i8, i32, i1 immarg)

llvm/test/tools/dxil-dis/lifetimes.ll

@@ -6,17 +6,17 @@ define void @test_lifetimes()  {
 ; CHECK-NEXT: [[ALLOCA:%.*]] = alloca [2 x i32], align 4
 ; CHECK-NEXT: [[GEP:%.*]] = getelementptr [2 x i32], [2 x i32]* [[ALLOCA]], i32 0, i32 0
 ; CHECK-NEXT: [[BITCAST:%.*]] = bitcast [2 x i32]* [[ALLOCA]] to i8*
-; CHECK-NEXT: call void @llvm.lifetime.start(i64 4, i8* nonnull [[BITCAST]])
+; CHECK-NEXT: call void @llvm.lifetime.start(i64 8, i8* nonnull [[BITCAST]])
 ; CHECK-NEXT: store i32 0, i32* [[GEP]], align 4
 ; CHECK-NEXT: [[BITCAST:%.*]] = bitcast [2 x i32]* [[ALLOCA]] to i8*
-; CHECK-NEXT: call void @llvm.lifetime.end(i64 4, i8* nonnull [[BITCAST]])
+; CHECK-NEXT: call void @llvm.lifetime.end(i64 8, i8* nonnull [[BITCAST]])
 ; CHECK-NEXT: ret void
 ;
   %a = alloca [2 x i32], align 4
   %gep = getelementptr [2 x i32], ptr %a, i32 0, i32 0
-  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull %a)
+  call void @llvm.lifetime.start.p0(ptr nonnull %a)
   store i32 0, ptr %gep, align 4
-  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %a)
+  call void @llvm.lifetime.end.p0(ptr nonnull %a)
   ret void
 }
 
@@ -29,10 +29,10 @@ define void @test_lifetimes()  {
 ; CHECK-DAG: declare void @llvm.lifetime.end(i64, i8* nocapture) [[LIFETIME_ATTRS]]
 
 ; Function Attrs: nounwind memory(argmem: readwrite)
-declare void @llvm.lifetime.end.p0(i64, ptr) #0
+declare void @llvm.lifetime.end.p0(ptr) #0
 
 ; Function Attrs: nounwind memory(argmem: readwrite)
-declare void @llvm.lifetime.start.p0(i64, ptr) #0
+declare void @llvm.lifetime.start.p0(ptr) #0
 
 attributes #0 = { nounwind memory(argmem: readwrite) }