[DirectX] Legalize lifetime intrinsics for DXIL (#148003)

Fixes #147395

This PR legalizes lifetime intrinsics for DXIL by
- Adding a bitcast for the lifetime intrinsics' pointer operand in
dxil-prepare to ensure it gets cast to an `i8*` when written to DXIL
- Removing the memory attribute from lifetime intrinsics in dxil-prepare
to match DXIL
- Making the DXIL bitcode writer write the base/demangled name of
lifetime intrinsics to the symbol table
- Making lifetime intrinsics an exception to Int64Ops shader flag
analysis (otherwise we get `error: Flags must match usage.` from the
validator)

Author: Icohedron

llvm/lib/Target/DirectX/DXILPrepare.cpp

@@ -24,6 +24,7 @@
 #include "llvm/IR/AttributeMask.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
@@ -239,6 +240,11 @@ class DXILPrepareModule : public ModulePass {
       for (size_t Idx = 0, End = F.arg_size(); Idx < End; ++Idx)
         F.removeParamAttrs(Idx, AttrMask);
 
+      // Lifetime intrinsics in LLVM 3.7 do not have the memory FnAttr
+      if (Intrinsic::ID IID = F.getIntrinsicID();
+          IID == Intrinsic::lifetime_start || IID == Intrinsic::lifetime_end)
+        F.removeFnAttr(Attribute::Memory);
+
       for (auto &BB : F) {
         IRBuilder<> Builder(&BB);
         for (auto &I : make_early_inc_range(BB)) {
@@ -247,7 +253,7 @@ class DXILPrepareModule : public ModulePass {
 
           // Emtting NoOp bitcast instructions allows the ValueEnumerator to be
           // unmodified as it reserves instruction IDs during contruction.
-          if (auto LI = dyn_cast<LoadInst>(&I)) {
+          if (auto *LI = dyn_cast<LoadInst>(&I)) {
             if (Value *NoOpBitcast = maybeGenerateBitcast(
                     Builder, PointerTypes, I, LI->getPointerOperand(),
                     LI->getType())) {
@@ -257,7 +263,7 @@ class DXILPrepareModule : public ModulePass {
             }
             continue;
           }
-          if (auto SI = dyn_cast<StoreInst>(&I)) {
+          if (auto *SI = dyn_cast<StoreInst>(&I)) {
             if (Value *NoOpBitcast = maybeGenerateBitcast(
                     Builder, PointerTypes, I, SI->getPointerOperand(),
                     SI->getValueOperand()->getType())) {
@@ -268,7 +274,7 @@ class DXILPrepareModule : public ModulePass {
             }
             continue;
           }
-          if (auto GEP = dyn_cast<GetElementPtrInst>(&I)) {
+          if (auto *GEP = dyn_cast<GetElementPtrInst>(&I)) {
             if (Value *NoOpBitcast = maybeGenerateBitcast(
                     Builder, PointerTypes, I, GEP->getPointerOperand(),
                     GEP->getSourceElementType()))
@@ -280,6 +286,17 @@ class DXILPrepareModule : public ModulePass {
             CB->removeRetAttrs(AttrMask);
             for (size_t Idx = 0, End = CB->arg_size(); Idx < End; ++Idx)
               CB->removeParamAttrs(Idx, AttrMask);
+            // LLVM 3.7 Lifetime intrinics require an i8* pointer operand, so we
+            // insert a bitcast here to ensure that is the case
+            if (isa<LifetimeIntrinsic>(CB)) {
+              Value *PtrOperand = CB->getArgOperand(1);
+              Builder.SetInsertPoint(CB);
+              PointerType *PtrTy = cast<PointerType>(PtrOperand->getType());
+              Value *NoOpBitcast = Builder.Insert(
+                  CastInst::Create(Instruction::BitCast, PtrOperand,
+                                   Builder.getPtrTy(PtrTy->getAddressSpace())));
+              CB->setArgOperand(1, NoOpBitcast);
+            }
             continue;
           }
         }

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;

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

@@ -2545,6 +2545,25 @@ void DXILBitcodeWriter::writeInstruction(const Instruction &I, unsigned InstID,
   Vals.clear();
 }
 
+// HLSL Change
+namespace {
+struct ValueNameCreator {
+  MallocAllocator Allocator;
+  SmallVector<ValueName *, 2>
+      ValueNames; // SmallVector N = 2 because we currently only expect this
+                  // to hold ValueNames for Lifetime intrinsics
+  ~ValueNameCreator() {
+    for (auto *VN : ValueNames)
+      VN->Destroy(Allocator);
+  }
+  ValueName *create(StringRef Name, Value *V) {
+    ValueName *VN = ValueName::create(Name, Allocator, V);
+    ValueNames.push_back(VN);
+    return VN;
+  }
+};
+} // anonymous namespace
+
 // Emit names for globals/functions etc.
 void DXILBitcodeWriter::writeFunctionLevelValueSymbolTable(
     const ValueSymbolTable &VST) {
@@ -2559,9 +2578,24 @@ void DXILBitcodeWriter::writeFunctionLevelValueSymbolTable(
   // to ensure the binary is the same no matter what values ever existed.
   SmallVector<const ValueName *, 16> SortedTable;
 
+  // HLSL Change
+  ValueNameCreator VNC;
   for (auto &VI : VST) {
-    SortedTable.push_back(VI.second->getValueName());
+    ValueName *VN = VI.second->getValueName();
+    // Clang mangles lifetime intrinsic names by appending '.p0' to the end,
+    // making them invalid lifetime intrinsics in LLVM 3.7. We can't
+    // demangle in dxil-prepare because it would result in invalid IR.
+    // Therefore we have to do this in the bitcode writer while writing its
+    // name to the symbol table.
+    if (const Function *Fn = dyn_cast<Function>(VI.getValue());
+        Fn && Fn->isIntrinsic()) {
+      Intrinsic::ID IID = Fn->getIntrinsicID();
+      if (IID == Intrinsic::lifetime_start || IID == Intrinsic::lifetime_end)
+        VN = VNC.create(Intrinsic::getBaseName(IID), VI.second);
+    }
+    SortedTable.push_back(VN);
   }
+
   // The keys are unique, so there shouldn't be stability issues.
   llvm::sort(SortedTable, [](const ValueName *A, const ValueName *B) {
     return A->first() < B->first();

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: ...

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

@@ -1,5 +1,6 @@
 ; 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-op-lower -dxil-prepare -mtriple=dxil-pc-shadermodel6.6-library %s | FileCheck %s --check-prefixes=CHECK,CHECK-PREPARE
 
 ; CHECK-LABEL: define void @test_legal_lifetime() {
 ; 
@@ -15,6 +16,14 @@
 ; 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-PREPARE-NEXT:    [[ACCUM_I_FLAT:%.*]] = alloca [1 x i32], align 4
+; CHECK-PREPARE-NEXT:    [[GEP:%.*]] = getelementptr i32, ptr [[ACCUM_I_FLAT]], i32 0
+; CHECK-PREPARE-NEXT:    [[BITCAST:%.*]] = bitcast ptr [[ACCUM_I_FLAT]] to ptr
+; CHECK-PREPARE-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[BITCAST]])
+; CHECK-PREPARE-NEXT:    store i32 0, ptr [[GEP]], align 4
+; CHECK-PREPARE-NEXT:    [[BITCAST:%.*]] = bitcast ptr [[ACCUM_I_FLAT]] to ptr
+; CHECK-PREPARE-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr nonnull [[BITCAST]])
+; 
 ; CHECK-NEXT:    ret void
 ;
 define void @test_legal_lifetime()  {
@@ -26,6 +35,22 @@ define void @test_legal_lifetime()  {
   ret void
 }
 
+; CHECK-PREPARE-DAG: attributes [[LIFETIME_ATTRS:#.*]] = { nounwind }
+
+; CHECK-PREPARE-DAG: ; Function Attrs: nounwind
+; CHECK-PREPARE-DAG: declare void @llvm.lifetime.start.p0(i64, ptr) [[LIFETIME_ATTRS]]
+
+; CHECK-PREPARE-DAG: ; Function Attrs: nounwind
+; CHECK-PREPARE-DAG: declare void @llvm.lifetime.end.p0(i64, ptr) [[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) }
+
 ; Set the validator version to 1.6
 !dx.valver = !{!0}
 !0 = !{i32 1, i32 6}

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