Reimplement address elimination. (#2605)

* Reimplement address elimination pass.

* Fix error.

* Update test references.

Co-authored-by: Yong He <yhe@nvidia.com>

10 files changed, 97 insertions, 88 deletions

diff --git a/tests/autodiff/reverse-addr-eliminate.slang b/tests/autodiff/reverse-addr-eliminate.slang
new file mode 100644
index 000000000..daa6fa32b
--- /dev/null
+++ b/tests/autodiff/reverse-addr-eliminate.slang
@@ -0,0 +1,56 @@
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+struct C : IDifferentiable
+{
+    float3 t;
+    float v;
+}
+
+struct B : IDifferentiable
+{
+    float2 f;
+    C arr[2];
+}
+
+struct A : IDifferentiable
+{
+    float x;
+    float y;
+    B fb;
+    C aarr[3];
+};
+
+[BackwardDifferentiable]
+A f(A a, int i)
+{
+    A aout;
+    aout.y = 2 * a.x;
+    aout.y = aout.y + 2 * a.x;
+    aout.x = aout.y + 5 * a.x;
+    aout.aarr[1].t = float3(a.y, 0.0, a.x);
+    aout.aarr[1].t = float3(a.y, 1.0, a.x + 1.0);
+
+    // Test that writes to a potentially dynamic address multiple times
+    // is allowed and will propagate the correct derivative.
+    aout.fb.arr[i].v = a.x * 2.0; // since this value is overwritten, the diff will not accumulate to a.x
+    aout.fb.arr[i].v = a.x * 3.0;
+    return aout;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    A a = {1.0, 2.0};
+
+    var dpa = diffPair(a);
+
+    A.Differential dout = { 1.0, 1.0, { float2(0), { { float3(1.0), 1.0 }, { float3(1.0), 1.0 } } }, { { float3(1.0), 1.0 }, { float3(1.0), 1.0 }, { float3(1.0), 1.0 } } };
+    
+    __bwd_diff(f)(dpa, 1, dout);
+    outputBuffer[0] = dpa.d.x; // Expect: 17
+    outputBuffer[1] = dpa.d.y; // Expect: 0
+}
diff --git a/tests/autodiff/reverse-addr-eliminate.slang.expected.txt b/tests/autodiff/reverse-addr-eliminate.slang.expected.txt
new file mode 100644
index 000000000..dd367f3f5
--- /dev/null
+++ b/tests/autodiff/reverse-addr-eliminate.slang.expected.txt
@@ -0,0 +1,6 @@
+type: float
+17.000000
+1.000000
+0.000000
+0.000000
+0.000000
diff --git a/tests/bugs/gh-841.slang.glsl b/tests/bugs/gh-841.slang.glsl
index dd949521b..820f7d2b3 100644
--- a/tests/bugs/gh-841.slang.glsl
+++ b/tests/bugs/gh-841.slang.glsl
@@ -1,5 +1,7 @@
 //TEST_IGNORE_FILE:
 #version 450
+layout(row_major) uniform;
+layout(row_major) buffer;
 
 layout(location = 0)
 out vec4 _S1;
@@ -10,26 +12,18 @@ in vec4 _S2;
 flat layout(location = 1)
 in uint _S3;
 
-struct RasterVertex_0
-{
-    vec4 c_0;
-    uint u_0;
-};
 
 void main()
 {
-    RasterVertex_0 _S4 = RasterVertex_0(_S2, _S3);
-    vec4 result_0 = _S4.c_0;
-
-    vec4 result_1;
-    if((_S4.u_0 & uint(1))!=0)
+    vec4 result_0;
+    if((_S3 & 1U) != 0U)
     {
-        result_1 = result_0 + 1.0;
+        result_0 = _S2 + 1.00000000000000000000;
     }
     else
     {
-        result_1 = result_0;
+        result_0 = _S2;
     }
-    _S1 = result_1;
+    _S1 = result_0;
     return;
 }
diff --git a/tests/cross-compile/glsl-generic-in.slang.glsl b/tests/cross-compile/glsl-generic-in.slang.glsl
index 2c3a3f877..c8d9b1bd1 100644
--- a/tests/cross-compile/glsl-generic-in.slang.glsl
+++ b/tests/cross-compile/glsl-generic-in.slang.glsl
@@ -29,31 +29,18 @@ in vec4 _S2;
 layout(location = 2)
 in vec2 _S3;
 
-struct GIn_0
-{
-    vec3 p0_0;
-    F_0 field_0;
-};
-
 struct VOut_0
 {
     vec4 projPos_0;
 };
 
-
-
 void main()
 {
     F_0 _S4 = { _S2, _S3 };
-    GIn_0 _S5 = { _S1, _S4 };
-
     VOut_0 vout_0;
-    vec3 _S6 = _S5.p0_0;
-
-    float _S7 = F_get_0(_S5.field_0);
-    float _S8 = E_get_0();
-
-    vout_0.projPos_0 = vec4(_S6, _S7 + _S8);
+    float _S5 = F_get_0(_S4);
+    float _S6 = E_get_0();
+    vout_0.projPos_0 = vec4(_S1, _S5 + _S6);
     gl_Position = vout_0.projPos_0;
     return;
 }
\ No newline at end of file
diff --git a/tests/pipeline/rasterization/get-attribute-at-vertex.slang.glsl b/tests/pipeline/rasterization/get-attribute-at-vertex.slang.glsl
index 864f44eb3..1da5f4f8a 100644
--- a/tests/pipeline/rasterization/get-attribute-at-vertex.slang.glsl
+++ b/tests/pipeline/rasterization/get-attribute-at-vertex.slang.glsl
@@ -14,7 +14,6 @@ out vec4 _S2;
 
 void main()
 {
-    uvec2 _S3 = uvec2(0U, 0U);
     _S2 = gl_BaryCoordNV.x * ((_S1)[(0U)]) + gl_BaryCoordNV.y * ((_S1)[(1U)]) + gl_BaryCoordNV.z * ((_S1)[(2U)]);
     return;
 }
diff --git a/tests/pipeline/rasterization/mesh/hello.slang.glsl b/tests/pipeline/rasterization/mesh/hello.slang.glsl
index 66630012b..c2c35915a 100644
--- a/tests/pipeline/rasterization/mesh/hello.slang.glsl
+++ b/tests/pipeline/rasterization/mesh/hello.slang.glsl
@@ -12,11 +12,7 @@ out gl_MeshPerVertexEXT
     vec4 gl_Position;
 } gl_MeshVerticesEXT[3];
 
-struct Vertex_0
-{
-    vec4 pos_0;
-    vec3 color_0;
-};
+out uvec3  gl_PrimitiveTriangleIndicesEXT[1];
 
 layout(local_size_x = 3, local_size_y = 1, local_size_z = 1) in;
 layout(max_vertices = 3) out;
@@ -27,9 +23,9 @@ void main()
     SetMeshOutputsEXT(3U, 1U);
     if(gl_LocalInvocationIndex < 3U)
     {
-        Vertex_0 _S2 = { vec4(positions_0[gl_LocalInvocationIndex], 0.00000000000000000000, 1.00000000000000000000), colors_0[gl_LocalInvocationIndex] };
-        gl_MeshVerticesEXT[gl_LocalInvocationIndex].gl_Position = _S2.pos_0;
-        _S1[gl_LocalInvocationIndex] = _S2.color_0;
+        vec3 _S2 = colors_0[gl_LocalInvocationIndex];
+        gl_MeshVerticesEXT[gl_LocalInvocationIndex].gl_Position = vec4(positions_0[gl_LocalInvocationIndex], 0.00000000000000000000, 1.00000000000000000000);
+        _S1[gl_LocalInvocationIndex] = _S2;
     }
     else
     {
diff --git a/tests/pipeline/rasterization/mesh/hlsl-syntax.slang.glsl b/tests/pipeline/rasterization/mesh/hlsl-syntax.slang.glsl
index 4a6b8f86f..96d681c5c 100644
--- a/tests/pipeline/rasterization/mesh/hlsl-syntax.slang.glsl
+++ b/tests/pipeline/rasterization/mesh/hlsl-syntax.slang.glsl
@@ -12,19 +12,14 @@ out gl_MeshPerVertexEXT
     vec4 gl_Position;
 } gl_MeshVerticesEXT[3];
 
-struct Vertex_0
-{
-    vec4 pos_0;
-    vec3 color_0;
-};
 
 void foo_0(uint _S2)
 {
     if(_S2 < 3U)
     {
-        Vertex_0 _S3 = { vec4(positions_0[_S2], 0.00000000000000000000, 1.00000000000000000000), colors_0[_S2] };
-        gl_MeshVerticesEXT[_S2].gl_Position = _S3.pos_0;
-        _S1[_S2] = _S3.color_0;
+        vec3 _S3 = colors_0[_S2];
+        gl_MeshVerticesEXT[_S2].gl_Position = vec4(positions_0[_S2], 0.00000000000000000000, 1.00000000000000000000);
+        _S1[_S2] = _S3;
     }
     else
     {
diff --git a/tests/pipeline/rasterization/mesh/passing-outputs.slang.glsl b/tests/pipeline/rasterization/mesh/passing-outputs.slang.glsl
index cfd7675ab..db8a69ef3 100644
--- a/tests/pipeline/rasterization/mesh/passing-outputs.slang.glsl
+++ b/tests/pipeline/rasterization/mesh/passing-outputs.slang.glsl
@@ -18,7 +18,7 @@ struct Vertex_0
 void just_two_0(out Vertex_0 v_0, out Vertex_0 w_0)
 {
     Texes_0 _S1 = { vec2(0.00000000000000000000, 0.00000000000000000000), vec4(0.00000000000000000000, 0.00000000000000000000, 0.00000000000000000000, 0.00000000000000000000) };
-    Vertex_0 _S2 = { vec4(0), vec3(1), _S1 };
+    Vertex_0 _S2 = { vec4(0.00000000000000000000), vec3(1.00000000000000000000), _S1 };
     v_0 = _S2;
     w_0 = v_0;
     return;
@@ -27,7 +27,7 @@ void just_two_0(out Vertex_0 v_0, out Vertex_0 w_0)
 void just_one_0(out Vertex_0 v_1)
 {
     Texes_0 _S3 = { vec2(0.00000000000000000000, 0.00000000000000000000), vec4(0.00000000000000000000, 0.00000000000000000000, 0.00000000000000000000, 0.00000000000000000000) };
-    Vertex_0 _S4 = { vec4(0), vec3(1), _S3 };
+    Vertex_0 _S4 = { vec4(0.00000000000000000000), vec3(1.00000000000000000000), _S3 };
     v_1 = _S4;
     return;
 }
@@ -40,8 +40,8 @@ void part_of_one_0(out vec4 p_0)
 
 void write_struct_0(out Texes_0 t_0)
 {
-    t_0.tex1_0 = vec2(0);
-    t_0.tex2_0 = vec4(1);
+    t_0.tex1_0 = vec2(0.00000000000000000000);
+    t_0.tex2_0 = vec4(1.00000000000000000000);
     return;
 }
 
@@ -61,13 +61,13 @@ out gl_MeshPerVertexEXT
 
 void everything_0()
 {
-    Texes_0 _S8 = { vec2(0.00000000000000000000, 0.00000000000000000000), vec4(0.00000000000000000000, 0.00000000000000000000, 0.00000000000000000000, 0.00000000000000000000) };
-    Vertex_0 _S9 = { vec4(0), vec3(1), _S8 };
-    gl_MeshVerticesEXT[0U].gl_Position = _S9.pos_0;
-    _S5[0U] = _S9.col_0;
-    Texes_0 _S10 = _S9.ts_0;
-    _S6[0U] = _S10.tex1_0;
-    _S7[0U] = _S10.tex2_0;
+    vec3 _S8 = vec3(1.00000000000000000000);
+    vec2 _S9 = vec2(0.00000000000000000000, 0.00000000000000000000);
+    vec4 _S10 = vec4(0.00000000000000000000, 0.00000000000000000000, 0.00000000000000000000, 0.00000000000000000000);
+    gl_MeshVerticesEXT[0U].gl_Position = vec4(0.00000000000000000000);
+    _S5[0U] = _S8;
+    _S6[0U] = _S9;
+    _S7[0U] = _S10;
     return;
 }
 
diff --git a/tests/pipeline/rasterization/mesh/primitive-output.slang.glsl b/tests/pipeline/rasterization/mesh/primitive-output.slang.glsl
index b7b4f1d62..6ff2dbce8 100644
--- a/tests/pipeline/rasterization/mesh/primitive-output.slang.glsl
+++ b/tests/pipeline/rasterization/mesh/primitive-output.slang.glsl
@@ -21,19 +21,6 @@ perprimitiveEXT out gl_MeshPerPrimitiveEXT
     bool gl_CullPrimitiveEXT;
 } gl_MeshPrimitivesEXT[1];
 
-struct Vertex_0
-{
-    vec4 pos_0;
-    vec3 color_0;
-};
-
-struct Prim_0
-{
-    vec3 triangleNormal_0;
-    uint id_0;
-    bool cull_0;
-};
-
 layout(local_size_x = 3, local_size_y = 1, local_size_z = 1) in;
 layout(max_vertices = 3) out;
 layout(max_primitives = 1) out;
@@ -43,9 +30,9 @@ void main()
     SetMeshOutputsEXT(3U, 1U);
     if(gl_LocalInvocationIndex < 3U)
     {
-        Vertex_0 _S3 = { vec4(positions_0[gl_LocalInvocationIndex], 0.00000000000000000000, 1.00000000000000000000), colors_0[gl_LocalInvocationIndex] };
-        gl_MeshVerticesEXT[gl_LocalInvocationIndex].gl_Position = _S3.pos_0;
-        _S1[gl_LocalInvocationIndex] = _S3.color_0;
+        vec3 _S3 = colors_0[gl_LocalInvocationIndex];
+        gl_MeshVerticesEXT[gl_LocalInvocationIndex].gl_Position = vec4(positions_0[gl_LocalInvocationIndex], 0.00000000000000000000, 1.00000000000000000000);
+        _S1[gl_LocalInvocationIndex] = _S3;
     }
     else
     {
@@ -53,10 +40,9 @@ void main()
     if(gl_LocalInvocationIndex < 1U)
     {
         gl_PrimitiveTriangleIndicesEXT[gl_LocalInvocationIndex] = uvec3(0U, 1U, 2U);
-        Prim_0 _S4 = { vec3(0.00000000000000000000, 0.00000000000000000000, 1.00000000000000000000), gl_LocalInvocationIndex, false };
-        _S2[gl_LocalInvocationIndex] = _S4.triangleNormal_0;
-        gl_MeshPrimitivesEXT[gl_LocalInvocationIndex].gl_PrimitiveID = int(_S4.id_0);
-        gl_MeshPrimitivesEXT[gl_LocalInvocationIndex].gl_CullPrimitiveEXT = _S4.cull_0;
+        _S2[gl_LocalInvocationIndex] = vec3(0.00000000000000000000, 0.00000000000000000000, 1.00000000000000000000);
+        gl_MeshPrimitivesEXT[gl_LocalInvocationIndex].gl_PrimitiveID = int(gl_LocalInvocationIndex);
+        gl_MeshPrimitivesEXT[gl_LocalInvocationIndex].gl_CullPrimitiveEXT = false;
     }
     else
     {
diff --git a/tests/pipeline/ray-tracing/trace-ray-inline.slang.glsl b/tests/pipeline/ray-tracing/trace-ray-inline.slang.glsl
index ae26ab498..f9e94bb30 100644
--- a/tests/pipeline/ray-tracing/trace-ray-inline.slang.glsl
+++ b/tests/pipeline/ray-tracing/trace-ray-inline.slang.glsl
@@ -67,14 +67,6 @@ void myMiss_0(inout MyRayPayload_0 payload_4)
     return;
 }
 
-struct RayDesc_0
-{
-    vec3 Origin_0;
-    float TMin_0;
-    vec3 Direction_0;
-    float TMax_0;
-};
-
 layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
 void main()
 {
@@ -84,9 +76,7 @@ void main()
     MyRayPayload_0 payload_5;
     MyRayPayload_0 _S2 = { -1 };
     payload_5 = _S2;
-
-    RayDesc_0 ray_0 = { C_0._data.origin_0, C_0._data.tMin_0, C_0._data.direction_0, C_0._data.tMax_0 };
-    rayQueryInitializeEXT((query_0), (myAccelerationStructure_0), (C_0._data.rayFlags_0 | 512), (C_0._data.instanceMask_0), (ray_0.Origin_0), (ray_0.TMin_0), (ray_0.Direction_0), (ray_0.TMax_0));
+    rayQueryInitializeEXT((query_0), (myAccelerationStructure_0), (C_0._data.rayFlags_0 | 512), (C_0._data.instanceMask_0), (C_0._data.origin_0), (C_0._data.tMin_0), (C_0._data.direction_0), (C_0._data.tMax_0));
 
     MyProceduralHitAttrs_0 committedProceduralAttrs_0;