Add documentation for uniformity analysis. (#3721)

7 files changed, 147 insertions, 5 deletions

diff --git a/docs/user-guide/08-compiling.md b/docs/user-guide/08-compiling.md
index 136b67858..fb9d8de6c 100644
--- a/docs/user-guide/08-compiling.md
+++ b/docs/user-guide/08-compiling.md
@@ -676,7 +676,7 @@ meanings of their `CompilerOptionValue` encodings.
 | VulkanBindShiftAll | Specifies the `-fvk-bind-shift` option for all spaces. `intValue0`: kind, `intValue1`: shift. |
 | GenerateWholeProgram | When set will emit target code for the entire program instead of for a specific entrypoint. `intValue0` specifies a bool value for the setting. |
 | UseUpToDateBinaryModule | When set will only load precompiled modules if it is up-to-date with its source. `intValue0` specifies a bool value for the setting. |
-
+| ValidateUniformity | When set will perform [uniformity analysis](a1-05-uniformity.md).
 ## Debugging
 
 Slang's SPIRV backend supports generating debug information using the [NonSemantic Shader DebugInfo Instructions](https://github.com/KhronosGroup/SPIRV-Registry/blob/main/nonsemantic/NonSemantic.Shader.DebugInfo.100.asciidoc).
diff --git a/docs/user-guide/a1-05-uniformity.md b/docs/user-guide/a1-05-uniformity.md
new file mode 100644
index 000000000..630dfb802
--- /dev/null
+++ b/docs/user-guide/a1-05-uniformity.md
@@ -0,0 +1,104 @@
+---
+layout: user-guide
+---
+
+Uniformity Analysis
+===========
+
+On certain hardwares, accessing resources with a non-uniform index may lead to significant performance degradation. Developers can often benefit from a compiler warning for unintentional non-uniform resource access.
+
+Starting from v2024.1.0, Slang provides uniformity analysis that can warn users if a non-dynamically-uniform value is being used unintentionally. This feature is not enabled by default but can be turned on with the `-validate-uniformity` commandline option when using `slangc`, or the `CompilerOptionName::ValidateUniformity` compiler option when using the API.
+
+In addition to specifying the compiler option, the source code must be augmented with the `dynamic_uniform` modifier to mark function parameters, struct fields or local variables as expecting a dynamic uniform value.
+
+For example, the following code will triger a warning:
+```csharp
+// Indicate that the `v` parameter needs to be dynamic uniform.
+float f(dynamic_uniform float v)
+{
+    return v + 1.0;
+}
+
+[numthread(1,1,1)]
+[shader("compute")]
+void main(int tid : SV_DispatchThreadID)
+{
+    f(tid); // warning: tid is not dynamically uniform.
+}
+```
+
+Currently, the analysis is being conservative for `struct` typed values, in that if any member of the `struct` is known to be non-uniform, the entire composite is
+treated as non-uniform:
+```csharp
+struct MyType
+{
+    int a;
+    int b;
+}
+
+void expectUniform(dynamic_uniform int a){}
+
+void main(int tid : SV_DispatchThreadID)
+{
+    MyType t;
+    t.a = tid;
+    t.b = 0;
+
+    // Generates a warning here despite t.b is non-uniform, because
+    // t.a is non-uniform and that assignment makes `t` non-uniform.
+    expectUniform(t.b);
+}
+```
+
+To allow the compiler to provide more accurate analysis, you can use mark struct fields as
+`dynamic_uniform`:
+
+```csharp
+struct MyType
+{
+    int a;
+    dynamic_uniform int b;
+}
+
+void expectUniform(dynamic_uniform int a){}
+
+void main(int tid : SV_DispatchThreadID)
+{
+    MyType t;
+    t.a = tid;
+    t.b = 0;
+
+    // OK, because MyType::b is marked as dynamic_uniform.
+    expectUniform(t.b);
+
+    // Warning: trying to assign non-uniform value to dynamic_uniform location.
+    t.b = tid;
+}
+```
+
+## Treat Values as Uniform
+
+In some cases, the compiler might not be able to deduce a value to be non-uniform. If you are certain that a value can
+be treated as dynamic uniform, you can call `asDynamicUniform()` function to force the compiler to treat the value as
+dynamic uniform. For example:
+```csharp
+void main(int tid: SV_DispatchThreadID
+{
+    expectUniform(asDynamicUniform(tid)); // OK.
+}
+```
+
+## Treat Function Return Values as Non-uniform
+
+The uniformity analysis will automatically propagate uniformity to function return values. However if you have
+an intrinsic function that does not have a body, or you simply wish the return value of a function to be always
+treated as non-uniform, you can mark the function with the `[NonUniformReturn]` attribute:
+```csharp
+[NonUniformReturn]
+int f() { return 0; }
+void expectUniform(dynamic_uniform int x) {}
+void main()
+{
+    expectUniform(f()); // Warning.
+}
+```
\ No newline at end of file
diff --git a/docs/user-guide/a1-special-topics.md b/docs/user-guide/a1-special-topics.md
index 86cf59ede..ef218b8fa 100644
--- a/docs/user-guide/a1-special-topics.md
+++ b/docs/user-guide/a1-special-topics.md
@@ -11,4 +11,5 @@ In this chapter:
 1. [Handling matrix layout differences on different platforms](a1-01-matrix-layout.md)
 2. [Using Slang to write PyTorch kernels](a1-02-slangpy.md)
 3. [Obfuscation](a1-03-obfuscation.md)
-4. [Interoperation with target-specific code](a1-04-interop.md)
\ No newline at end of file
+4. [Interoperation with target-specific code](a1-04-interop.md)
+5. [Uniformity Analysis](a1-05-uniformity.md)
\ No newline at end of file
diff --git a/docs/user-guide/toc.html b/docs/user-guide/toc.html
index 3564ccf69..a696a3d61 100644
--- a/docs/user-guide/toc.html
+++ b/docs/user-guide/toc.html
@@ -165,6 +165,12 @@
 <li data-link="a1-04-interop#inline-spirv-assembly"><span>Inline SPIRV Assembly</span></li>
 </ul>
 </li>
+<li data-link="a1-05-uniformity"><span>Uniformity Analysis</span>
+<ul class="toc_list">
+<li data-link="a1-05-uniformity#treat-values-as-uniform"><span>Treat Values as Uniform</span></li>
+<li data-link="a1-05-uniformity#treat-function-return-values-as-non-uniform"><span>Treat Function Return Values as Non-uniform</span></li>
+</ul>
+</li>
 </ul>
 </li>
 </ul>
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 2ced9180e..d8cff1a9a 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -719,8 +719,8 @@ DIAGNOSTIC(42001, Error, invalidUseOfTorchTensorTypeInDeviceFunc, "invalid use o
 
 DIAGNOSTIC(45001, Error, unresolvedSymbol, "unresolved external symbol '$0'.")
 
-DIAGNOSTIC(41201, Warning, expectDynamicUniformArgument, "argument for '$0' is not a dynamic uniform.")
-DIAGNOSTIC(41201, Warning, expectDynamicUniformValue, "value stored at this location must be dynamic uniform.")
+DIAGNOSTIC(41201, Warning, expectDynamicUniformArgument, "argument for '$0' might not be a dynamic uniform, use `asDynamicUniform()` to silence this warning.")
+DIAGNOSTIC(41201, Warning, expectDynamicUniformValue, "value stored at this location must be dynamic uniform, use `asDynamicUniform()` to silence this warning.")
 
 //
 // 5xxxx - Target code generation.
diff --git a/source/slang/slang-ir-uniformity.cpp b/source/slang/slang-ir-uniformity.cpp
index 9c463f530..d29c0735a 100644
--- a/source/slang/slang-ir-uniformity.cpp
+++ b/source/slang/slang-ir-uniformity.cpp
@@ -310,6 +310,9 @@ namespace Slang
                             auto ifElse = as<IRIfElse>(user);
                             visitControlDependentBlock(ifElse->getTrueBlock());
                             visitControlDependentBlock(ifElse->getFalseBlock());
+                            // Mark phi nodes as non-uniform if any of its incoming values are non-uniform.
+                            for (auto param : ifElse->getAfterBlock()->getParams())
+                                addToWorkList(param);
                             break;
                         }
                         case kIROp_Switch:
@@ -318,6 +321,9 @@ namespace Slang
                             for (UInt c = 0; c < switchInst->getCaseCount(); c++)
                                 visitControlDependentBlock(switchInst->getCaseLabel(c));
                             visitControlDependentBlock(switchInst->getDefaultLabel());
+                            // Mark phi nodes as non-uniform if any of its incoming values are non-uniform.
+                            for (auto param : switchInst->getBreakLabel()->getParams())
+                                addToWorkList(param);
                             break;
                         }
                         case kIROp_Call:
diff --git a/tests/language-feature/dynamic-uniform/dynamic-uniform-1.slang b/tests/language-feature/dynamic-uniform/dynamic-uniform-1.slang
index 61583a84d..c5077122e 100644
--- a/tests/language-feature/dynamic-uniform/dynamic-uniform-1.slang
+++ b/tests/language-feature/dynamic-uniform/dynamic-uniform-1.slang
@@ -16,6 +16,10 @@ struct B
     int a;
     int c;
     dynamic_uniform int d;
+    dynamic_uniform int d1;
+    dynamic_uniform int d2;
+    dynamic_uniform int d3;
+    dynamic_uniform int d4;
 }
 
 int getBa(B obj)
@@ -55,11 +59,32 @@ void main(int tid: SV_DispatchThreadID)
     // CHECK-DAG: ([[# @LINE+1]]): warning 41201
     obj.d = ff;
 
+    int t = 0;
+    if (tid == 0)
+        t = 1;
+
+    // CHECK-DAG: ([[# @LINE+1]]): warning 41201
+    obj.d1 = t;
+
+    int u = 0;
+    if (ub.a == 0)
+        u = 1;
+    // CHECK-NOT: ([[# @LINE+1]]): warning 41201
+    obj.d2 = u;
+
+    int vv = 0;
+    if (tid == 0)
+        vv = 2;
+    else
+        vv = 2;
+    // CHECK-NOT: ([[# @LINE+1]]): warning 41201
+    obj.d3 = vv;
+
     // CHECK-NOT: ([[# @LINE+1]]): warning 41201
     expectUniform(getBdFromInout(obj));
     
     // CHECK-NOT: ([[# @LINE+1]]): warning 41201
-    obj.d = d; // OK
+    obj.d4 = d; // OK
 
     // CHECK-NOT: ([[# @LINE+1]]): warning 41201
     expectUniform(d);