Check wave sensitive on operand of gradient ops instead of gradient ops. (#3489)

* Check wave sensitive on operand of gradient ops instead of gradient ops.

lib/HLSL/DxilPreparePasses.cpp

@@ -1432,13 +1432,20 @@ public:
       if (F.isDeclaration())
         continue;
 
-      SmallVector<CallInst *, 16> localGradientOps;
+      DenseSet<Instruction *> localGradientArgs;
       for (CallInst *CI : gradientOps) {
-        if (CI->getParent()->getParent() == &F)
-          localGradientOps.emplace_back(CI);
+        if (CI->getParent()->getParent() == &F) {
+          for (Value *V : CI->arg_operands()) {
+            // TODO: only check operand which used for gradient calculation.
+            Instruction *vI = dyn_cast<Instruction>(V);
+            if (!vI)
+              continue;
+            localGradientArgs.insert(vI);
+          }
+        }
       }
 
-      if (localGradientOps.empty())
+      if (localGradientArgs.empty())
         continue;
 
       PostDominatorTree PDT;
@@ -1447,9 +1454,10 @@ public:
           WaveSensitivityAnalysis::create(PDT));
 
       WaveVal->Analyze(&F);
-      for (CallInst *op : localGradientOps) {
-        if (WaveVal->IsWaveSensitive(op)) {
-          dxilutil::EmitWarningOnInstruction(op,
+      for (Instruction *gradArg : localGradientArgs) {
+        // Check operand of gradient ops, not gradientOps itself.
+        if (WaveVal->IsWaveSensitive(gradArg)) {
+          dxilutil::EmitWarningOnInstruction(gradArg,
                                              UniNoWaveSensitiveGradientErrMsg);
         }
       }

tools/clang/test/HLSLFileCheck/hlsl/diagnostics/warnings/wave_sense.hlsl

@@ -0,0 +1,18 @@
+// RUN: %dxc -E main -T ps_6_0 %s | FileCheck %s
+
+// Make sure no warning when sample is not wave sensitive.
+// CHECK-NOT:warning
+// CHECK:main
+Texture2D<float> tex;
+SamplerState s;
+
+float main(float4 a : A) : SV_Target {
+
+   float i=WaveReadLaneFirst(a.z);
+
+   if (i > 3)
+    return   sin(tex.Sample(s, a.xy));
+
+   return 1;
+
+}