[Validation] Drill through chained GEPs for TGSM

lib/DxilValidation/DxilValidation.cpp

@@ -3305,21 +3305,24 @@ static void ValidateFunctionBody(Function *F, ValidationContext &ValCtx) {
 
       if (PointerType *PT = dyn_cast<PointerType>(I.getType())) {
         if (PT->getAddressSpace() == DXIL::kTGSMAddrSpace) {
-          if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(&I)) {
-            Value *Ptr = GEP->getPointerOperand();
-            // Allow inner constant GEP
-            if (isa<ConstantExpr>(Ptr) && isa<GEPOperator>(Ptr))
-              Ptr = cast<GEPOperator>(Ptr)->getPointerOperand();
-            if (!isa<GlobalVariable>(Ptr)) {
-              ValCtx.EmitInstrError(
-                  &I, ValidationRule::InstrFailToResloveTGSMPointer);
+          // Walk through GEPs and bitcasts to ensure the pointer ultimately
+          // comes from a global variable. This was unnecessary before SM 6.9
+          // because everything was scalarized, but now we can have arrays of
+          // vectors in TGSM, so we need to allow GEPs and bitcasts.
+          if (isa<GetElementPtrInst>(&I) || isa<BitCastInst>(&I)) {
+            Value *Ptr = cast<Instruction>(&I)->getOperand(0);
+            while (Ptr) {
+              if (GEPOperator *GEP = dyn_cast<GEPOperator>(Ptr)) {
+                Ptr = GEP->getPointerOperand();
+                continue;
+              }
+              if (BitCastOperator *BC = dyn_cast<BitCastOperator>(Ptr)) {
+                Ptr = BC->getOperand(0);
+                continue;
+              }
+              break;
             }
-          } else if (BitCastInst *BCI = dyn_cast<BitCastInst>(&I)) {
-            Value *Ptr = BCI->getOperand(0);
-            // Allow inner constant GEP
-            if (isa<ConstantExpr>(Ptr) && isa<GEPOperator>(Ptr))
-              Ptr = cast<GEPOperator>(Ptr)->getPointerOperand();
-            if (!isa<GetElementPtrInst>(Ptr) && !isa<GlobalVariable>(Ptr)) {
+            if (!isa<GlobalVariable>(Ptr)) {
               ValCtx.EmitInstrError(
                   &I, ValidationRule::InstrFailToResloveTGSMPointer);
             }

tools/clang/test/LitDXILValidation/GroupShared/tgsm-chained-gep-ambiguous.ll

@@ -0,0 +1,93 @@
+; RUN: not %dxv %s 2>&1 | FileCheck %s
+
+; Negative tests for the TGSM pointer chain walk added with chained GEP /
+; bitcast support. The validator only walks through GEP and bitcast
+; operators / instructions; any other construct (phi, select, etc.) in
+; the middle of a chain leaves the ultimate base ambiguous and must be
+; rejected. Both the phi/select itself (which produces a TGSM pointer
+; that is not a GEP or bitcast) and any GEP / bitcast that consumes it
+; must be reported.
+
+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"
+
+@"\01?GA@@3PAIA" = external addrspace(3) global [32 x i32], align 4
+@"\01?GB@@3PAIA" = external addrspace(3) global [32 x i32], align 4
+
+define void @main() {
+entry:
+  %tid = call i32 @dx.op.threadIdInGroup.i32(i32 95, i32 0)
+  %cond = icmp ult i32 %tid, 16
+  %ga = getelementptr [32 x i32], [32 x i32] addrspace(3)* @"\01?GA@@3PAIA", i32 0, i32 %tid
+  %gb = getelementptr [32 x i32], [32 x i32] addrspace(3)* @"\01?GB@@3PAIA", i32 0, i32 %tid
+
+  ; --- select between two TGSM GEPs, then chain a GEP through the select ---
+  ; The select itself directly produces a TGSM pointer that is neither a
+  ; GEP nor a bitcast, so it must be rejected.
+  ; CHECK: error: TGSM pointers must originate from an unambiguous TGSM global variable.
+  ; CHECK-NEXT: note: at '{{.*}}select{{.*}}' in block 'entry' of function 'main'.
+  %sel = select i1 %cond, i32 addrspace(3)* %ga, i32 addrspace(3)* %gb
+  ; A GEP that walks through the select must also be rejected: the chain
+  ; walker stops at the select (not a GEP / bitcast / GlobalVariable).
+  ; CHECK: error: TGSM pointers must originate from an unambiguous TGSM global variable.
+  ; CHECK-NEXT: note: at '%sel_gep = getelementptr i32, i32 addrspace(3)* %sel, i32 0' in block 'entry' of function 'main'.
+  %sel_gep = getelementptr i32, i32 addrspace(3)* %sel, i32 0
+  store i32 1, i32 addrspace(3)* %sel_gep, align 4
+
+  ; A bitcast that walks through the select must likewise be rejected.
+  ; CHECK: error: TGSM pointers must originate from an unambiguous TGSM global variable.
+  ; CHECK-NEXT: note: at '%sel_bc = bitcast i32 addrspace(3)* %sel to float addrspace(3)*' in block 'entry' of function 'main'.
+  %sel_bc = bitcast i32 addrspace(3)* %sel to float addrspace(3)*
+  store float 2.000000e+00, float addrspace(3)* %sel_bc, align 4
+
+  br i1 %cond, label %then, label %else
+
+then:
+  %ga2 = getelementptr [32 x i32], [32 x i32] addrspace(3)* @"\01?GA@@3PAIA", i32 0, i32 %tid
+  br label %merge
+
+else:
+  %gb2 = getelementptr [32 x i32], [32 x i32] addrspace(3)* @"\01?GB@@3PAIA", i32 0, i32 %tid
+  br label %merge
+
+merge:
+  ; --- phi joining two TGSM GEPs, then chain a GEP/bitcast through the phi ---
+  ; The phi itself produces an ambiguous TGSM pointer.
+  ; CHECK: error: TGSM pointers must originate from an unambiguous TGSM global variable.
+  ; CHECK-NEXT: note: at '{{.*}}phi{{.*}}' in block 'merge' of function 'main'.
+  %p = phi i32 addrspace(3)* [ %ga2, %then ], [ %gb2, %else ]
+  ; A GEP through the phi must be rejected: the chain walk reaches the
+  ; phi which is neither a GEP, a bitcast, nor a GlobalVariable.
+  ; CHECK: error: TGSM pointers must originate from an unambiguous TGSM global variable.
+  ; CHECK-NEXT: note: at '%phi_gep = getelementptr i32, i32 addrspace(3)* %p, i32 0' in block 'merge' of function 'main'.
+  %phi_gep = getelementptr i32, i32 addrspace(3)* %p, i32 0
+  store i32 3, i32 addrspace(3)* %phi_gep, align 4
+
+  ; Likewise a bitcast through the phi must be rejected, even when the
+  ; chain has multiple GEP / bitcast hops before the phi appears.
+  ; CHECK: error: TGSM pointers must originate from an unambiguous TGSM global variable.
+  ; CHECK-NEXT: note: at '%phi_bc = bitcast i32 addrspace(3)* %phi_gep to float addrspace(3)*' in block 'merge' of function 'main'.
+  %phi_bc = bitcast i32 addrspace(3)* %phi_gep to float addrspace(3)*
+  store float 4.000000e+00, float addrspace(3)* %phi_bc, align 4
+
+  ; CHECK: Validation failed.
+  ret void
+}
+
+declare i32 @dx.op.threadIdInGroup.i32(i32, i32) #0
+
+attributes #0 = { nounwind readnone }
+
+!llvm.ident = !{!0}
+!dx.version = !{!1}
+!dx.valver = !{!2}
+!dx.shaderModel = !{!3}
+!dx.entryPoints = !{!4}
+
+!0 = !{!"dxc test"}
+!1 = !{i32 1, i32 9}
+!2 = !{i32 1, i32 10}
+!3 = !{!"cs", i32 6, i32 9}
+!4 = !{void ()* @main, !"main", null, null, !5}
+!5 = !{i32 4, !6}
+!6 = !{i32 64, i32 1, i32 1}

tools/clang/test/LitDXILValidation/GroupShared/tgsm-chained-gep.ll

@@ -0,0 +1,66 @@
+; RUN: %dxv %s 2>&1 | FileCheck %s
+
+; Regression test for TGSM pointer validation when SM 6.9 preserves native
+; vectors. The DXIL validator must accept TGSM pointers produced by chained
+; getelementptr and/or bitcast instructions whose ultimate base is an
+; unambiguous groupshared global variable. Prior to this fix, the validator
+; only allowed a single GEP instruction (optionally over a constant-expression
+; GEP), and a single bitcast whose operand was a single GEP; arbitrary chains
+; of GEP and bitcast instructions were incorrectly reported as "TGSM pointers
+; must originate from an unambiguous TGSM global variable" even though they
+; were unambiguously rooted at a TGSM global.
+
+; CHECK: Validation succeeded.
+
+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"
+
+@"\01?GS@@3PAV?$matrix@M$02$02@@A.v" = external addrspace(3) global [64 x <9 x float>], align 4
+
+define void @main() {
+  %tid = call i32 @dx.op.threadIdInGroup.i32(i32 95, i32 0)
+  ; First-level GEP: indexes the per-thread element of the groupshared
+  ; matrix array. Its base is the global, so it is unambiguous.
+  %elt = getelementptr [64 x <9 x float>], [64 x <9 x float>] addrspace(3)* @"\01?GS@@3PAV?$matrix@M$02$02@@A.v", i32 0, i32 %tid
+  store <9 x float> zeroinitializer, <9 x float> addrspace(3)* %elt, align 4
+  ; Second-level GEP: indexes into the per-thread vector. Its base is a
+  ; GEP *instruction* (not a constant expression), which previously
+  ; failed validation but is still unambiguously rooted at the global.
+  %scalar = getelementptr <9 x float>, <9 x float> addrspace(3)* %elt, i32 0, i32 0
+  store float 1.000000e+00, float addrspace(3)* %scalar, align 4
+
+  ; Bitcast of a GEP instruction. Chain: BC -> GEP -> Global.
+  %bc1 = bitcast <9 x float> addrspace(3)* %elt to <9 x i32> addrspace(3)*
+  store <9 x i32> zeroinitializer, <9 x i32> addrspace(3)* %bc1, align 4
+
+  ; GEP through a bitcast through a GEP. Chain: GEP -> BC -> GEP -> Global.
+  %bc1_scalar = getelementptr <9 x i32>, <9 x i32> addrspace(3)* %bc1, i32 0, i32 3
+  store i32 7, i32 addrspace(3)* %bc1_scalar, align 4
+
+  ; Bitcast of a bitcast. Chain: BC -> BC -> GEP -> Global.
+  %bc2 = bitcast <9 x i32> addrspace(3)* %bc1 to <9 x float> addrspace(3)*
+  store <9 x float> zeroinitializer, <9 x float> addrspace(3)* %bc2, align 4
+
+  ; Bitcast of a constant-expression GEP. Chain: BC -> ConstGEP -> Global.
+  %bcconst = bitcast <9 x float> addrspace(3)* getelementptr inbounds ([64 x <9 x float>], [64 x <9 x float>] addrspace(3)* @"\01?GS@@3PAV?$matrix@M$02$02@@A.v", i32 0, i32 0) to <9 x i32> addrspace(3)*
+  store <9 x i32> zeroinitializer, <9 x i32> addrspace(3)* %bcconst, align 4
+  ret void
+}
+
+declare i32 @dx.op.threadIdInGroup.i32(i32, i32) #0
+
+attributes #0 = { nounwind readnone }
+
+!llvm.ident = !{!0}
+!dx.version = !{!1}
+!dx.valver = !{!2}
+!dx.shaderModel = !{!3}
+!dx.entryPoints = !{!4}
+
+!0 = !{!"dxc test"}
+!1 = !{i32 1, i32 9}
+!2 = !{i32 1, i32 10}
+!3 = !{!"cs", i32 6, i32 9}
+!4 = !{void ()* @main, !"main", null, null, !5}
+!5 = !{i32 4, !6}
+!6 = !{i32 64, i32 1, i32 1}