[bug] Fix crash in dxilgen when buffer subscript gep feeds an lcssa phi (microsoft#5570)

[bug] Fix crash in dxilgen when buffer subscript gep feeds an lcssa phi

This change fixes a crash in the dxilgen pass when lowering buffer
subscripts. These subscripts feed into a gep instruction that needs to
be replaced with the computed index and offset for the dxil buffer load
intrinsics.

When processing all the users of the subscripts the code would crash on
any phi node user. We can handle single-value phis by propagating
through to the users of the phi. The single-value phis can be inserted
by the lcssa pass.

After fixing this initial crash it triggered another error later in the
compiler for a similar issue. The lowered handle used by the buffer load
was also being fed through an lcssa phi by a later run of that pass.
This would trigger a dxil validation error because handles are not
allowed in phis. Since it is a single-value phi we can again propagate
the value to the phi users to avoid this error.

The test case was reduced from a real-world shader and only occurred
when compiling with optimization disabled (Od).

Author: dmpots

lib/HLSL/DxilNoOptLegalize.cpp

@@ -12,8 +12,10 @@
 #include "llvm/IR/Instructions.h"
 #include "dxc/HLSL/DxilGenerationPass.h"
 #include "dxc/HLSL/DxilNoops.h"
+#include "dxc/Support/Global.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Analysis/DxilValueCache.h"
+#include "llvm/Analysis/InstructionSimplify.h"
 
 using namespace llvm;
 
@@ -120,6 +122,15 @@ class DxilNoOptSimplifyInstructions : public ModulePass {
               I->eraseFromParent();
               Changed = true;
             }
+          } else if (PHINode* Phi = dyn_cast<PHINode>(I)) {
+            // Replace all simple phi values (such as those inserted by lcssa) with the
+            // value itself. This avoids phis in places they are not expected because
+            // the normal simplify passes will clean them up.
+            if (Value *NewPhi = llvm::SimplifyInstruction(Phi, M.getDataLayout())) {
+              Phi->replaceAllUsesWith(NewPhi);
+              Phi->eraseFromParent();
+              Changed = true;
+            }
           }
         }
       }

lib/HLSL/HLOperationLower.cpp

@@ -7602,6 +7602,20 @@ void TranslateStructBufSubscriptUser(
         handle, ResKind, bufIdx, baseOffset, status, OP, DL);
     }
     BCI->eraseFromParent();
+  } else if (PHINode* Phi = dyn_cast<PHINode>(user)) {
+    if (Phi->getNumIncomingValues() != 1) {
+        dxilutil::EmitErrorOnInstruction(Phi, "Phi not supported for buffer subscript");
+        return;
+    }
+    // Since the phi only has a single value we can safely process its
+    // users to translate the subscript. These single-value phis are
+    // inserted by the lcssa pass.
+    for (auto U = Phi->user_begin(); U != Phi->user_end();) {
+      Value *PhiUser = *(U++);
+      TranslateStructBufSubscriptUser(cast<Instruction>(PhiUser),
+        handle, ResKind, bufIdx, baseOffset, status, OP, DL);
+    }
+    Phi->eraseFromParent();
   } else {
     // should only used by GEP
     GetElementPtrInst *GEP = cast<GetElementPtrInst>(user);

test/HLSL/passes/dxil_o0_legalize/lcssa_phi_non_dxil.ll

@@ -0,0 +1,20 @@
+; RUN: opt -S -dxil-o0-simplify-inst %s | FileCheck %s
+
+; Make sure we remove a single-value phi node.
+
+; CHECK: %add = add i32 1, 2
+; CHECK-NOT: phi
+; CHECK: %mul = mul i32 %add, 2
+target datalayout = "e-m:e-p:32:32-i1:32-i8:32-i16:32-i32:32-i64:64-f16:32-f32:32-f64:64-n8:16:32:64"
+target triple = "dxil-ms-dx"
+
+define void @foo() {
+entry:
+  %add = add i32 1, 2
+  br label %exit
+
+exit:
+  %lcssa = phi i32 [ %add, %entry ]
+  %mul = mul i32 %lcssa, 2
+  ret void
+}

test/HLSL/passes/dxil_o0_legalize/lcssa_phi_of_handles.ll

@@ -0,0 +1,63 @@
+; RUN: opt -S -dxil-o0-simplify-inst %s | FileCheck %s
+
+; Make sure we remove a single-value phi node of dxil handles.
+; Phis of handles are invalid dxil but we can make it valid by replacing
+; the phi with its one incoming value.
+
+; CHECK: %handle = call %dx.types.Handle @dx.op.createHandle
+; CHECK-NOT: phi
+; CHECK: %load = call %dx.types.ResRet.i32 @dx.op.bufferLoad.i32(i32 68, %dx.types.Handle %handle, i32 0, i32 0)
+target datalayout = "e-m:e-p:32:32-i1:32-i8:32-i16:32-i32:32-i64:64-f16:32-f32:32-f64:64-n8:16:32:64"
+target triple = "dxil-ms-dx"
+
+%dx.types.Handle = type { i8* }
+%dx.types.ResRet.i32 = type { i32, i32, i32, i32, i32 }
+
+define void @main() {
+entry:
+  %handle = call %dx.types.Handle @dx.op.createHandle(i32 57, i8 0, i32 0, i32 3, i1 false)  ; CreateHandle(resourceClass,rangeId,index,nonUniformIndex)
+  br label %exit
+
+exit:
+  %lcssa = phi %dx.types.Handle [ %handle, %entry ]
+  %load = call %dx.types.ResRet.i32 @dx.op.bufferLoad.i32(i32 68, %dx.types.Handle %lcssa, i32 0, i32 0) 
+  ret void
+}
+
+; Function Attrs: nounwind readonly
+declare %dx.types.Handle @dx.op.createHandle(i32, i8, i32, i32, i1) #2
+
+; Function Attrs: nounwind readonly
+declare %dx.types.ResRet.i32 @dx.op.bufferLoad.i32(i32, %dx.types.Handle, i32, i32) #2
+
+attributes #0 = { nounwind readnone }
+attributes #1 = { nounwind }
+attributes #2 = { nounwind readonly }
+
+!llvm.ident = !{!0}
+!dx.version = !{!1}
+!dx.valver = !{!2}
+!dx.shaderModel = !{!3}
+!dx.resources = !{!4}
+!dx.viewIdState = !{!8}
+!dx.entryPoints = !{!9}
+
+!0 = !{!"dxc(private) 1.7.0.14024 (main, 1d6b5627a)"}
+!1 = !{i32 1, i32 0}
+!2 = !{i32 1, i32 7}
+!3 = !{!"vs", i32 6, i32 0}
+!4 = !{!5, null, null, null}
+!5 = !{!6}
+!6 = !{i32 0, i32* undef, !"", i32 0, i32 3, i32 1, i32 12, i32 0, !7}
+!7 = !{i32 1, i32 4}
+!8 = !{[3 x i32] [i32 1, i32 1, i32 1]}
+!9 = !{void ()* @main, !"main", !10, !4, !17}
+!10 = !{!11, !15, null}
+!11 = !{!12}
+!12 = !{i32 0, !"A", i8 5, i8 0, !13, i8 0, i32 1, i8 1, i32 0, i8 0, !14}
+!13 = !{i32 0}
+!14 = !{i32 3, i32 1}
+!15 = !{!16}
+!16 = !{i32 0, !"Z", i8 5, i8 0, !13, i8 1, i32 1, i8 1, i32 0, i8 0, !14}
+!17 = !{i32 0, i64 17}
+

test/HLSL/passes/dxilgen/lcssa_structbuf_subscript.ll

@@ -0,0 +1,85 @@
+; RUN: opt -S -dxilgen %s | FileCheck %s
+
+; Make sure we can correctly generate the buffer load when
+; its gep goes through a single-value phi node.
+
+; CHECK: exit:
+; CHECK-NEXT: call %dx.types.ResRet.i32 @dx.op.rawBufferLoad.i32
+
+
+target datalayout = "e-m:e-p:32:32-i1:32-i8:32-i16:32-i32:32-i64:64-f16:32-f32:32-f64:64-n8:16:32:64"
+target triple = "dxil-ms-dx"
+
+%ConstantBuffer = type opaque
+%"class.StructuredBuffer<S>" = type { %struct.S }
+%struct.S = type { i32 }
+%dx.types.Handle = type { i8* }
+%dx.types.ResourceProperties = type { i32, i32 }
+@"\01?buf@@3V?$StructuredBuffer@US@@@@A" = external global %"class.StructuredBuffer<S>", align 4
+
+@"$Globals" = external constant %ConstantBuffer
+
+; Function Attrs: nounwind
+define void @main(<4 x float>* noalias) #0 {
+entry:
+  %a = load %"class.StructuredBuffer<S>", %"class.StructuredBuffer<S>"* @"\01?buf@@3V?$StructuredBuffer@US@@@@A"
+  %b = call %dx.types.Handle @"dx.hl.createhandle..%dx.types.Handle (i32, %\22class.StructuredBuffer<S>\22)"(i32 0, %"class.StructuredBuffer<S>" %a)
+  %c = call %dx.types.Handle @"dx.hl.annotatehandle..%dx.types.Handle (i32, %dx.types.Handle, %dx.types.ResourceProperties, %\22class.StructuredBuffer<S>\22)"(i32 11, %dx.types.Handle %b, %dx.types.ResourceProperties { i32 524, i32 4 }, %"class.StructuredBuffer<S>" zeroinitializer)
+  %d = call %struct.S* @"dx.hl.subscript.[].rn.%struct.S* (i32, %dx.types.Handle, i32)"(i32 0, %dx.types.Handle %c, i32 0)
+  %gep = getelementptr inbounds %struct.S, %struct.S* %d, i32 0, i32 0
+  br label %exit
+
+exit:
+  %lcssa = phi i32* [ %gep, %entry ]
+  %struct_buffer_load = load i32, i32* %lcssa, align 4
+
+  ret void
+}
+
+; Function Attrs: nounwind readnone
+declare %struct.S* @"dx.hl.subscript.[].rn.%struct.S* (i32, %dx.types.Handle, i32)"(i32, %dx.types.Handle, i32) #0
+
+; Function Attrs: nounwind readnone
+declare %dx.types.Handle @"dx.hl.createhandle..%dx.types.Handle (i32, %\22class.StructuredBuffer<S>\22)"(i32, %"class.StructuredBuffer<S>") #0
+
+; Function Attrs: nounwind readnone
+declare %dx.types.Handle @"dx.hl.annotatehandle..%dx.types.Handle (i32, %dx.types.Handle, %dx.types.ResourceProperties, %\22class.StructuredBuffer<S>\22)"(i32, %dx.types.Handle, %dx.types.ResourceProperties, %"class.StructuredBuffer<S>") #0
+
+
+attributes #0 = { nounwind }
+
+!llvm.module.flags = !{!0}
+!pauseresume = !{!1}
+!llvm.ident = !{!2}
+!dx.version = !{!3}
+!dx.valver = !{!4}
+!dx.shaderModel = !{!5}
+!dx.typeAnnotations = !{!6}
+!dx.entryPoints = !{!13}
+!dx.fnprops = !{!17}
+!dx.options = !{!18, !19}
+
+!0 = !{i32 2, !"Debug Info Version", i32 3}
+!1 = !{!"hlsl-hlemit", !"hlsl-hlensure"}
+!2 = !{!"dxc(private) 1.7.0.14024 (main, 1d6b5627a)"}
+!3 = !{i32 1, i32 0}
+!4 = !{i32 1, i32 7}
+!5 = !{!"ps", i32 6, i32 0}
+!6 = !{i32 1, void (<4 x float>*)* @main, !7}
+!7 = !{!8, !10}
+!8 = !{i32 0, !9, !9}
+!9 = !{}
+!10 = !{i32 1, !11, !12}
+!11 = !{i32 4, !"SV_Target", i32 7, i32 9}
+!12 = !{i32 0}
+!13 = !{void (<4 x float>*)* @main, !"main", null, !14, null}
+!14 = !{null, null, !15, null}
+!15 = !{!16}
+!16 = !{i32 0, %ConstantBuffer* @"$Globals", !"$Globals", i32 0, i32 -1, i32 1, i32 0, null}
+!17 = !{void (<4 x float>*)* @main, i32 0, i1 false}
+!18 = !{i32 152}
+!19 = !{i32 -1}
+!20 = !DILocation(line: 2, column: 3, scope: !21)
+!21 = !DISubprogram(name: "main", scope: !22, file: !22, line: 1, type: !23, isLocal: false, isDefinition: true, scopeLine: 1, flags: DIFlagPrototyped, isOptimized: false, function: void (<4 x float>*)* @main)
+!22 = !DIFile(filename: ".\5Ct2.hlsl", directory: "")
+!23 = !DISubroutineType(types: !9)

tools/clang/test/CodeGenDXIL/StructuredBuffer/subscript_lcssa_phi.hlsl

@@ -0,0 +1,29 @@
+// RUN: dxc /Tvs_6_0 /Od %s | FileCheck %s
+
+// This is a regression test for a crash in dxilgen when compiling
+// with optimization disabled. It triggered a crash because the
+// gep used by the subscript operator was fed through an lcssa
+// phi for a use out of the loop and phis are not allowed users
+// of the gep for the subscript operator.
+
+// CHECK: call %dx.types.ResRet.i32 @dx.op.bufferLoad.i32
+
+struct S { uint bits; };
+StructuredBuffer<S> buf : register(t3);
+
+S get(uint idx)
+{
+    return buf[idx];
+}
+
+[RootSignature("DescriptorTable(SRV(t0, numDescriptors=10))")]
+uint main(uint C : A) : Z
+{
+    int i = 0;
+    S s;
+    do {
+        s = get(i);
+    } while (i < C);
+    
+    return s.bits;
+}