[HLSL][DirectX] Avoid visited values when searching in hasNonUniformIndex

[joaosaffran]

This patch fixes hasNonUniformIndex search so that it accounts for any path that connects nuri to index access to return true

fix: #189438

[llvmbot]

@llvm/pr-subscribers-backend-directx

Author: joaosaffran

Changes

This patch fix hasNonUniformIndex search so that it avoids searching already visited values.

fix: #189438

---

Full diff: https://github.com/llvm/llvm-project/pull/189498.diff

2 Files Affected:

• (modified) llvm/lib/Target/DirectX/DXILOpLowering.cpp (+5)
• (added) llvm/test/CodeGen/DirectX/is_nonuniform_within_loop.ll (+23)

diff --git a/llvm/lib/Target/DirectX/DXILOpLowering.cpp b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
index 794bae8fd3a34..995e5311b5a66 100644
--- a/llvm/lib/Target/DirectX/DXILOpLowering.cpp
+++ b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
@@ -252,10 +252,15 @@ class OpLowerer {
       return false;
 
     SmallVector<Value *> WorkList;
+    SmallPtrSet<Value *, 8> Visited;
+
     WorkList.push_back(IndexOp);
 
     while (!WorkList.empty()) {
       Value *V = WorkList.pop_back_val();
+      if (!Visited.insert(V).second)
+        continue;
+
       if (auto *CI = dyn_cast<CallInst>(V)) {
         if (CI->getCalledFunction()->getIntrinsicID() ==
             Intrinsic::dx_resource_nonuniformindex)
diff --git a/llvm/test/CodeGen/DirectX/is_nonuniform_within_loop.ll b/llvm/test/CodeGen/DirectX/is_nonuniform_within_loop.ll
new file mode 100644
index 0000000000000..30ff6e4309f82
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/is_nonuniform_within_loop.ll
@@ -0,0 +1,23 @@
+; RUN: opt -S  -dxil-intrinsic-expansion -dxil-op-lower  -mtriple=dxil-pc-shadermodel6.3-library %s
+
+; The goal of this test it to make sure compilation finishes successfully,
+; in bad case it should timeout.
+
+@Out.str = private unnamed_addr constant [4 x i8] c"Out\00", align 1
+
+define void @main() local_unnamed_addr {
+entry:
+  %cmp.i4.not = icmp eq i32 4, 0
+  br i1 %cmp.i4.not, label %_Z4mainj.exit, label %for.body.i
+
+for.body.i:                                       ; preds = %entry, %for.body.i
+  %i.0.i5 = phi i32 [ %inc.i, %for.body.i ], [ 0, %entry ]
+  %1 = tail call target("dx.RawBuffer", i32, 1, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 4, i32 %i.0.i5, ptr nonnull @Out.str)
+  %2 = tail call noundef i32 @llvm.dx.resource.updatecounter(target("dx.RawBuffer", i32, 1, 0) %1, i8 1)
+  %inc.i = add nuw nsw i32 %i.0.i5, 1
+  %exitcond.not = icmp eq i32 %inc.i, 4
+  br i1 %exitcond.not, label %_Z4mainj.exit, label %for.body.i
+
+_Z4mainj.exit:                                    ; preds = %for.body.i, %entry
+  ret void
+}

[bogner]

Instead of just checking that we avoid a timeout, we should check that the NonUniformIndex argument to createHandle/createHandleFromBinding is set correctly to false. We should also have a version of the test that uses the index with @dx.resource.nonuniformindex that tests that we correctly set that to true even when we have cyclic users.

[bogner]

Please name all values in tests, using numbered values makes updating tests harder.

[bogner]

This test technically didn't hang, so this comment is a bit misleading. I also realize that it's a bit boring, in that we don't actually look at the phi node at all in this case. I wonder if a test based on something silly like this is a little bit more interesting, where the handlefrombinding still uses the value from the phi: https://hlsl.godbolt.org/z/TsaPzMhr5

Note that this never hung either, because when we see the function call we return early, but I think it's still a little bit interesting to have the case where we walk through the phi node at least once.

[hekota]

Few more comments.

[hekota]

Suggested change

	; The goal of this test it to make sure compilation finishes successfully,
	; in bad case it should timeout.
	; Regression test for llvm/llvm-project#189438.
	; The goal of this test it to make sure compilation finishes successfully.

[hekota]

Duplicate declaration of WorkList.

[hekota]

Please restore this check, along with the one previously at lines 266–267. Constant indices and operators are very common, and adding these guards significantly reduces the number of loop iterations.

[hekota]

LGMT! Just please add one more check for the constant value.

[hekota]

Suggested change

	for (Value *Op : Inst->operands())
	Worklist.push_back(Op);
	for (Value *Op : Inst->operands()) {
	if (isa<llvm::Constant>(Op))
	continue;
	Worklist.push_back(Op);
	}

[joaosaffran]

This is already handled right after we pop from the Worklist. Do you want me to add in there as well?