Mark wave intrinsics as convergent to prevent loop miscompilation

The JumpThreading pass could restructure loops containing wave intrinsics
by threading edges through the loop latch block. This moved wave ops
(WaveReadLaneFirst, WaveActiveCountBits, etc.) from inside the loop to
after it, changing the set of active lanes at the call site. On SIMT
hardware, this produced incorrect results because all lanes reconverge
at the post-loop point rather than only the matching subset.

Fix:
- DxilConvergentMark: Mark wave-sensitive HL functions with the
  Attribute::Convergent attribute before optimizer passes run.
- JumpThreading: In ThreadEdge, walk backward from the latch to the
  loop header to identify loop body blocks. If any contains a convergent
  call, prevent threading through the latch.
- DxilOperations::GetOpFunc: Mark DXIL wave op functions with
  Attribute::Convergent for post-dxilgen optimizer passes.

Fixes a bug where a material binning pattern using WaveReadLaneFirst +
WaveActiveCountBits in a while-loop produced correct results at -Od but
incorrect results with optimizations enabled.

Co-authored-by: Copilot <[email protected]>

Author: JoeCitizen

lib/DXIL/DxilOperations.cpp

@@ -6738,6 +6738,14 @@ Function *OP::GetOpFunc(OpCode opCode, Type *pOverloadType) {
     if (existF->getFunctionType() != pFT)
       return nullptr;
     F = existF;
+    // HLSL Change Begin - ensure attributes are set on existing functions.
+    if (OpProps.FuncAttr != Attribute::None &&
+        !F->hasFnAttribute(OpProps.FuncAttr))
+      F->addFnAttr(OpProps.FuncAttr);
+    // Mark wave ops as convergent since they depend on the active lane set.
+    if (IsDxilOpWave(opCode) && !F->hasFnAttribute(Attribute::Convergent))
+      F->addFnAttr(Attribute::Convergent);
+    // HLSL Change End
     UpdateCache(opClass, pOverloadType, F);
     return F;
   }
@@ -6749,6 +6757,9 @@ Function *OP::GetOpFunc(OpCode opCode, Type *pOverloadType) {
   F->addFnAttr(Attribute::NoUnwind);
   if (OpProps.FuncAttr != Attribute::None)
     F->addFnAttr(OpProps.FuncAttr);
+  // HLSL Change - mark wave ops as convergent.
+  if (IsDxilOpWave(opCode))
+    F->addFnAttr(Attribute::Convergent);
 
   return F;
 }

lib/HLSL/DxilConvergent.cpp

@@ -44,14 +44,28 @@ class DxilConvergentMark : public ModulePass {
 
   bool runOnModule(Module &M) override {
     const ShaderModel *SM = M.GetOrCreateHLModule().GetShaderModel();
+
+    bool bUpdated = false;
+
+    // HLSL Change Begin - Mark wave-sensitive HL functions as convergent.
+    // This prevents optimizer passes (especially JumpThreading) from
+    // restructuring control flow around wave ops, which would change
+    // the set of active lanes at wave op call sites.
+    for (Function &F : M.functions()) {
+      if (F.isDeclaration() && IsHLWaveSensitive(&F) &&
+          !F.hasFnAttribute(Attribute::Convergent)) {
+        F.addFnAttr(Attribute::Convergent);
+        bUpdated = true;
+      }
+    }
+    // HLSL Change End
+
     // Can skip if in a shader and version that doesn't support derivatives.
     if (!SM->IsPS() && !SM->IsLib() &&
         (!SM->IsSM66Plus() || (!SM->IsCS() && !SM->IsMS() && !SM->IsAS())))
-      return false;
+      return bUpdated;
     SupportsVectors = SM->IsSM69Plus();
 
-    bool bUpdated = false;
-
     for (Function &F : M.functions()) {
       if (F.isDeclaration())
         continue;

lib/Transforms/Scalar/JumpThreading.cpp

@@ -24,6 +24,7 @@
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/IR/CallSite.h" // HLSL Change - for convergent call detection
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
@@ -1388,6 +1389,47 @@ bool JumpThreading::ThreadEdge(BasicBlock *BB,
     return false;
   }
 
+  // HLSL Change Begin - Don't thread through loop latch blocks when the loop
+  // body contains convergent calls (e.g., wave intrinsics). Threading through
+  // a latch can restructure the loop so that convergent calls that were inside
+  // the loop end up outside it, changing which lanes are active at those call
+  // sites on SIMT hardware.
+  for (succ_iterator SI = succ_begin(BB), SE = succ_end(BB); SI != SE; ++SI) {
+    BasicBlock *Header = *SI;
+    if (!LoopHeaders.count(Header))
+      continue;
+    // BB is a loop latch (has a back-edge to Header). Walk backward from BB
+    // to find all blocks in the loop body and check for convergent calls.
+    SmallVector<BasicBlock *, 16> Worklist;
+    SmallPtrSet<BasicBlock *, 16> InLoop;
+    InLoop.insert(Header); // Seed to prevent going above header.
+    Worklist.push_back(BB);
+    bool HasConvergent = false;
+    while (!Worklist.empty() && !HasConvergent) {
+      BasicBlock *WBB = Worklist.pop_back_val();
+      if (!InLoop.insert(WBB).second)
+        continue;
+      for (auto &I : *WBB) {
+        if (auto CS = CallSite(&I)) {
+          if (CS.hasFnAttr(Attribute::Convergent)) {
+            HasConvergent = true;
+            break;
+          }
+        }
+      }
+      if (!HasConvergent)
+        for (pred_iterator PI = pred_begin(WBB), PE = pred_end(WBB); PI != PE;
+             ++PI)
+          Worklist.push_back(*PI);
+    }
+    if (HasConvergent) {
+      DEBUG(dbgs() << "  Not threading across loop latch BB '" << BB->getName()
+                   << "' - loop body has convergent calls\n");
+      return false;
+    }
+  }
+  // HLSL Change End
+
   unsigned JumpThreadCost = getJumpThreadDuplicationCost(BB, BBDupThreshold);
   if (JumpThreadCost > BBDupThreshold) {
     DEBUG(dbgs() << "  Not threading BB '" << BB->getName()

tools/clang/test/HLSLFileCheck/hlsl/intrinsics/wave/convergent/wave-in-loop-not-sunk.hlsl

@@ -0,0 +1,52 @@
+// RUN: %dxc -T cs_6_6 -E main %s | FileCheck %s
+
+// Regression test for a bug where the optimizer (JumpThreading) would
+// restructure a while-loop containing wave intrinsics, moving
+// WaveActiveCountBits outside the loop. This changes the set of active
+// lanes at the wave op call site, producing incorrect results on SIMT
+// hardware.
+
+// Verify that WaveAllBitCount (opcode 135) appears BEFORE the loop's
+// back-edge phi, ensuring it stays inside the loop body.
+
+// CHECK: call i32 @dx.op.waveReadLaneFirst
+// CHECK: call i32 @dx.op.waveAllOp
+// CHECK: call i1 @dx.op.waveIsFirstLane
+// CHECK: phi i32
+// CHECK: br i1
+
+RWStructuredBuffer<uint> Output : register(u1);
+
+cbuffer Constants : register(b0) {
+  uint Width;
+  uint Height;
+  uint NumMaterials;
+};
+
+[numthreads(32, 1, 1)]
+void main(uint3 DTid : SV_DispatchThreadID) {
+  uint x = DTid.x;
+  uint y = DTid.y;
+
+  if (x >= Width || y >= Height)
+    return;
+
+  // Compute a material ID per lane (simple hash).
+  uint materialID = ((x * 7) + (y * 13)) % NumMaterials;
+
+  // Binning loop: each iteration peels off one material group.
+  // WaveReadLaneFirst picks a material, matching lanes count themselves
+  // with WaveActiveCountBits, and the first lane in the group writes
+  // the count. Non-matching lanes loop back for the next material.
+  bool go = true;
+  while (go) {
+    uint firstMat = WaveReadLaneFirst(materialID);
+    if (firstMat == materialID) {
+      uint count = WaveActiveCountBits(true);
+      if (WaveIsFirstLane()) {
+        InterlockedAdd(Output[firstMat], count);
+      }
+      go = false;
+    }
+  }
+}