[DirectX] Legalize `llvm.lifetime.*` intrinsics in EmbedDXILPass (#150100)

Fixes #147395

This PR:
- Excludes lifetime intrinsics from the Int64Ops shader flags analysis
to match DXC behavior and pass DXIL validation.
- Performs legalization of `llvm.lifetime.*` intrinsics in the
EmbedDXILPass just before invoking the DXILBitcodeWriter.
- After invoking the DXILBitcodeWriter, all lifetime intrinsics and
associated bitcasts are removed from the module to keep the Module
Verifier happy. This is fine since lifetime intrinsics are not needed by
any passes after the EmbedDXILPass.

5 files changed, 152 insertions, 15 deletions

diff --git a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
index bd3349d..eb4adfe 100644
--- a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
+++ b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
@@ -152,7 +152,7 @@ void ModuleShaderFlags::updateFunctionFlags(ComputedShaderFlags &CSF,
   if (!CSF.Int64Ops)
     CSF.Int64Ops = I.getType()->isIntegerTy(64);
 
-  if (!CSF.Int64Ops) {
+  if (!CSF.Int64Ops && !isa<LifetimeIntrinsic>(&I)) {
     for (const Value *Op : I.operands()) {
       if (Op->getType()->isIntegerTy(64)) {
         CSF.Int64Ops = true;
diff --git a/llvm/lib/Target/DirectX/DXILWriter/DXILWriterPass.cpp b/llvm/lib/Target/DirectX/DXILWriter/DXILWriterPass.cpp
index dfc79039c..1bd5dd7 100644
--- a/llvm/lib/Target/DirectX/DXILWriter/DXILWriterPass.cpp
+++ b/llvm/lib/Target/DirectX/DXILWriter/DXILWriterPass.cpp
@@ -17,6 +17,7 @@
 #include "llvm/Analysis/ModuleSummaryAnalysis.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
@@ -52,6 +53,53 @@ public:
   }
 };
 
+static void legalizeLifetimeIntrinsics(Module &M) {
+  for (Function &F : M) {
+    Intrinsic::ID IID = F.getIntrinsicID();
+    if (IID != Intrinsic::lifetime_start && IID != Intrinsic::lifetime_end)
+      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());
+      Value *NoOpBitCast = CastInst::Create(Instruction::BitCast, PtrOperand,
+                                            PtrTy, "", CI->getIterator());
+      CI->setArgOperand(1, NoOpBitCast);
+    }
+  }
+}
+
+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)
+      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();
+      BCI->eraseFromParent();
+    }
+    F.eraseFromParent();
+  }
+}
+
 class EmbedDXILPass : public llvm::ModulePass {
 public:
   static char ID; // Pass identification, replacement for typeid
@@ -70,8 +118,17 @@ public:
     // Only the output bitcode need to be DXIL triple.
     M.setTargetTriple(Triple("dxil-ms-dx"));
 
+    // Perform late legalization of lifetime intrinsics that would otherwise
+    // fail the Module Verifier if performed in an earlier pass
+    legalizeLifetimeIntrinsics(M);
+
     WriteDXILToFile(M, OS);
 
+    // We no longer need lifetime intrinsics after bitcode serialization, so we
+    // simply remove them to keep the Module Verifier happy after our
+    // not-so-legal legalizations
+    removeLifetimeIntrinsics(M);
+
     // Recover triple.
     M.setTargetTriple(OriginalTriple);
 
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/lifetimes-noint64op.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/lifetimes-noint64op.ll
new file mode 100644
index 0000000..736c86e
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/lifetimes-noint64op.ll
@@ -0,0 +1,36 @@
+; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s --check-prefix=DXC
+
+target triple = "dxil-pc-shadermodel6.7-library"
+
+; CHECK: ; Combined Shader Flags for Module
+; CHECK-NEXT: ; Shader Flags Value: 0x00000000
+; CHECK-NEXT: ;
+; CHECK-NOT:  ; Note: shader requires additional functionality:
+; CHECK-NOT:  ;       64-Bit integer
+; CHECK-NOT:  ; Note: extra DXIL module flags:
+; CHECK-NOT:  ;
+; CHECK-NEXT: ; Shader Flags for Module Functions
+; CHECK-NEXT: ; Function lifetimes : 0x00000000
+
+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)
+  ret void
+}
+
+; Function Attrs: nounwind memory(argmem: readwrite)
+declare void @llvm.lifetime.start.p0(i64, ptr) #1
+
+; Function Attrs: nounwind memory(argmem: readwrite)
+declare void @llvm.lifetime.end.p0(i64, ptr) #1
+
+attributes #0 = { convergent norecurse nounwind "hlsl.export"}
+attributes #1 = { nounwind memory(argmem: readwrite) }
+
+; DXC: - Name:            SFI0
+; DXC-NEXT:     Size:            8
+; DXC-NOT:     Flags:
+; DXC-NOT:         Int64Ops:        true
+; DXC: ...
diff --git a/llvm/test/CodeGen/DirectX/legalize-lifetimes-valver-1.6.ll b/llvm/test/CodeGen/DirectX/legalize-lifetimes-valver-1.6.ll
index 6552ccd..77133eb 100644
--- a/llvm/test/CodeGen/DirectX/legalize-lifetimes-valver-1.6.ll
+++ b/llvm/test/CodeGen/DirectX/legalize-lifetimes-valver-1.6.ll
@@ -1,21 +1,27 @@
 ; RUN: opt -S -passes='dxil-op-lower' -mtriple=dxil-pc-shadermodel6.3-library %s | FileCheck %s --check-prefixes=CHECK,CHECK-SM63
 ; RUN: opt -S -passes='dxil-op-lower' -mtriple=dxil-pc-shadermodel6.6-library %s | FileCheck %s --check-prefixes=CHECK,CHECK-SM66
+; RUN: opt -S -dxil-prepare -dxil-embed -mtriple=dxil-pc-shadermodel6.6-library %s | FileCheck %s --check-prefixes=CHECK,CHECK-EMBED
+
+; Lifetime intrinsics are not valid prior to shader model 6.6 and are instead
+; replaced with undef stores, provided the validator version is 1.6 or greater
+
+; The dxil-embed pass will remove lifetime intrinsics because they transformed
+; in a way that is illegal in modern LLVM IR before serializing to DXIL bitcode.
+; So we check that no bitcast or lifetime intrinsics remain after dxil-embed
 
 ; CHECK-LABEL: define void @test_legal_lifetime() {
-; 
-; CHECK-SM63-NEXT:    [[ACCUM_I_FLAT:%.*]] = alloca [1 x i32], align 4
-; CHECK-SM63-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-SM63-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:    [[ACCUM_I_FLAT:%.*]] = alloca [1 x i32], align 4
-; CHECK-SM66-NEXT:    [[GEP:%.*]] = getelementptr i32, ptr [[ACCUM_I_FLAT]], i32 0
-; CHECK-SM66-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
-; CHECK-SM66-NEXT:    store i32 0, ptr [[GEP]], align 4
-; CHECK-SM66-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr nonnull [[ACCUM_I_FLAT]])
-; 
-; CHECK-NEXT:    ret void
+; 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-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-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
diff --git a/llvm/test/tools/dxil-dis/lifetimes.ll b/llvm/test/tools/dxil-dis/lifetimes.ll
new file mode 100644
index 0000000..cb3e629
--- /dev/null
+++ b/llvm/test/tools/dxil-dis/lifetimes.ll
@@ -0,0 +1,38 @@
+; RUN: llc --filetype=obj %s -o - | dxil-dis -o - | FileCheck %s
+target triple = "dxil-unknown-shadermodel6.7-library"
+
+define void @test_lifetimes()  {
+; CHECK-LABEL: 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: 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: 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)
+  store i32 0, ptr %gep, align 4
+  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %a)
+  ret void
+}
+
+; CHECK-DAG: attributes [[LIFETIME_ATTRS:#.*]] = { nounwind }
+
+; CHECK-DAG: ; Function Attrs: nounwind
+; CHECK-DAG: declare void @llvm.lifetime.start(i64, i8* nocapture) [[LIFETIME_ATTRS]]
+
+; CHECK-DAG: ; Function Attrs: nounwind
+; 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
+
+; Function Attrs: nounwind memory(argmem: readwrite)
+declare void @llvm.lifetime.start.p0(i64, ptr) #0
+
+attributes #0 = { nounwind memory(argmem: readwrite) }
+