tint/uniformity: implement analysis for full and partial assignments

As per https://github.com/gpuweb/gpuweb/pull/3298

Bug: tint:1703
Change-Id: I88eb40764473fdae52962b36df1b4a1c929603f6
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/105000
Reviewed-by: Alan Baker <alanbaker@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: Antonio Maiorano <amaiorano@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>

src/tint/resolver/uniformity.cc

@@ -47,6 +47,23 @@ namespace tint::resolver {
 
 namespace {
 
+/// Unwraps `u->expr`'s chain of indirect (*) and address-of (&) expressions, returning the first
+/// expression that is neither of these.
+/// E.g. If `u` is `*(&(*(&p)))`, returns `p`.
+const ast::Expression* UnwrapIndirectAndAddressOfChain(const ast::UnaryOpExpression* u) {
+    auto* e = u->expr;
+    while (true) {
+        auto* unary = e->As<ast::UnaryOpExpression>();
+        if (unary &&
+            (unary->op == ast::UnaryOp::kIndirection || unary->op == ast::UnaryOp::kAddressOf)) {
+            e = unary->expr;
+        } else {
+            break;
+        }
+    }
+    return e;
+}
+
 /// CallSiteTag describes the uniformity requirements on the call sites of a function.
 enum CallSiteTag {
     CallSiteRequiredToBeUniform,
@@ -203,6 +220,10 @@ struct FunctionInfo {
     /// Includes pointer parameters.
     std::unordered_set<const sem::Variable*> local_var_decls;
 
+    /// The set of partial pointer variables - pointers that point to a subobject (into an array or
+    /// struct).
+    std::unordered_set<const sem::Variable*> partial_ptrs;
+
     /// LoopSwitchInfo tracks information about the value of variables for a control flow construct.
     struct LoopSwitchInfo {
         /// The type of this control flow construct.
@@ -943,14 +964,27 @@ class UniformityGraph {
 
             [&](const ast::VariableDeclStatement* decl) {
                 Node* node;
+                auto* sem_var = sem_.Get(decl->variable);
                 if (decl->variable->constructor) {
                     auto [cf1, v] = ProcessExpression(cf, decl->variable->constructor);
                     cf = cf1;
                     node = v;
+
+                    // Store if lhs is a partial pointer
+                    if (sem_var->Type()->Is<sem::Pointer>()) {
+                        auto* init = sem_.Get(decl->variable->constructor);
+                        if (auto* unary_init = init->Declaration()->As<ast::UnaryOpExpression>()) {
+                            auto* e = UnwrapIndirectAndAddressOfChain(unary_init);
+                            if (e->IsAnyOf<ast::IndexAccessorExpression,
+                                           ast::MemberAccessorExpression>()) {
+                                current_function_->partial_ptrs.insert(sem_var);
+                            }
+                        }
+                    }
                 } else {
                     node = cf;
                 }
-                current_function_->variables.Set(sem_.Get(decl->variable), node);
+                current_function_->variables.Set(sem_var, node);
 
                 if (decl->variable->Is<ast::Var>()) {
                     current_function_->local_var_decls.insert(
@@ -1126,11 +1160,37 @@ class UniformityGraph {
             });
     }
 
+    /// @param u unary expression with op == kIndirection
+    /// @returns true if `u` is an indirection unary expression that ultimately dereferences a
+    /// partial pointer, false otherwise.
+    bool IsDerefOfPartialPointer(const ast::UnaryOpExpression* u) {
+        TINT_ASSERT(Resolver, u->op == ast::UnaryOp::kIndirection);
+
+        // To determine if we're dereferencing a partial pointer, unwrap *&
+        // chains; if the final expression is an identifier, see if it's a
+        // partial pointer. If it's not an identifier, then it must be an
+        // index/acessor expression, and thus a partial pointer.
+        auto* e = UnwrapIndirectAndAddressOfChain(u);
+        if (auto* var_user = sem_.Get<sem::VariableUser>(e)) {
+            if (current_function_->partial_ptrs.count(var_user->Variable())) {
+                return true;
+            }
+        } else {
+            TINT_ASSERT(
+                Resolver,
+                (e->IsAnyOf<ast::IndexAccessorExpression, ast::MemberAccessorExpression>()));
+            return true;
+        }
+        return false;
+    }
+
     /// Process an LValue expression.
     /// @param cf the input control flow node
     /// @param expr the expression to process
     /// @returns a pair of (control flow node, variable node)
-    std::pair<Node*, Node*> ProcessLValueExpression(Node* cf, const ast::Expression* expr) {
+    std::pair<Node*, Node*> ProcessLValueExpression(Node* cf,
+                                                    const ast::Expression* expr,
+                                                    bool is_partial_reference = false) {
         return Switch(
             expr,
 
@@ -1144,9 +1204,11 @@ class UniformityGraph {
                     auto* value = CreateNode(name + "_lvalue");
                     auto* old_value = current_function_->variables.Set(local, value);
 
-                    // Aggregate values link back to their previous value, as they can never become
-                    // uniform again.
-                    if (!local->Type()->UnwrapRef()->is_scalar() && old_value) {
+                    // If i is part of an expression that is a partial reference to a variable (e.g.
+                    // index or member access), we link back to the variable's previous value. If
+                    // the previous value was non-uniform, a partial assignment will not make it
+                    // uniform.
+                    if (is_partial_reference && old_value) {
                         value->AddEdge(old_value);
                     }
 
@@ -1160,14 +1222,15 @@ class UniformityGraph {
             },
 
             [&](const ast::IndexAccessorExpression* i) {
-                auto [cf1, l1] = ProcessLValueExpression(cf, i->object);
+                auto [cf1, l1] =
+                    ProcessLValueExpression(cf, i->object, /*is_partial_reference*/ true);
                 auto [cf2, v2] = ProcessExpression(cf1, i->index);
                 l1->AddEdge(v2);
                 return std::pair<Node*, Node*>(cf2, l1);
             },
 
             [&](const ast::MemberAccessorExpression* m) {
-                return ProcessLValueExpression(cf, m->structure);
+                return ProcessLValueExpression(cf, m->structure, /*is_partial_reference*/ true);
             },
 
             [&](const ast::UnaryOpExpression* u) {
@@ -1179,15 +1242,17 @@ class UniformityGraph {
                     auto* deref = CreateNode(name + "_deref");
                     auto* old_value = current_function_->variables.Set(source_var, deref);
 
-                    // Aggregate values link back to their previous value, as they can never become
-                    // uniform again.
-                    if (!source_var->Type()->UnwrapRef()->UnwrapPtr()->is_scalar() && old_value) {
-                        deref->AddEdge(old_value);
+                    if (old_value) {
+                        // If derefercing a partial reference or partial pointer, we link back to
+                        // the variable's previous value. If the previous value was non-uniform, a
+                        // partial assignment will not make it uniform.
+                        if (is_partial_reference || IsDerefOfPartialPointer(u)) {
+                            deref->AddEdge(old_value);
+                        }
                     }
-
                     return std::pair<Node*, Node*>(cf, deref);
                 }
-                return ProcessLValueExpression(cf, u->expr);
+                return ProcessLValueExpression(cf, u->expr, is_partial_reference);
             },
 
             [&](Default) {

test/tint/bug/tint/1703.wgsl

@@ -0,0 +1,37 @@
+var<private> my_global : vec4<f32>;
+
+@group(0) @binding(0) var<uniform> my_uniform : f32;
+@group(0) @binding(1) var my_texture : texture_2d<f32>;
+@group(0) @binding(2) var my_sampler : sampler;
+
+fn foo_member_initialize() {
+    var vb2 : vec2<bool>;
+
+    vb2.x = my_global.z != 0; // Assign non-uniform value to x component
+
+    // Overwrite x component with uniform value -> doesn't make this uniform as per the spec
+    vb2.x = my_uniform == -1.0f;
+
+    // Trying to set all components should make this uniform again -- not working!
+    vb2 = vec2(my_uniform == -1.0f, false);
+
+    if (vb2.x) {
+        let r : vec4<f32> = textureSampleBias(my_texture, my_sampler, vec2<f32>(), 0.0);
+    }
+}
+
+fn foo_default_initialize() {
+    var vb2 : vec2<bool>;
+
+    vb2.x = my_global.z != 0; // Assign non-uniform value to x component
+
+    // Overwrite x component with uniform value -> doesn't make this uniform as per the spec
+    vb2.x = my_uniform == -1.0f;
+
+    // Even resetting it doesn't work
+    vb2 = vec2<bool>();
+
+    if (vb2.x) {
+        let r : vec4<f32> = textureSampleBias(my_texture, my_sampler, vec2<f32>(), 0.0);
+    }
+}

test/tint/bug/tint/1703.wgsl.expected.dxc.hlsl

@@ -0,0 +1,31 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+static float4 my_global = float4(0.0f, 0.0f, 0.0f, 0.0f);
+cbuffer cbuffer_my_uniform : register(b0, space0) {
+  uint4 my_uniform[1];
+};
+Texture2D<float4> my_texture : register(t1, space0);
+SamplerState my_sampler : register(s2, space0);
+
+void foo_member_initialize() {
+  bool2 vb2 = bool2(false, false);
+  vb2.x = (my_global.z != 0.0f);
+  vb2.x = (asfloat(my_uniform[0].x) == -1.0f);
+  vb2 = bool2((asfloat(my_uniform[0].x) == -1.0f), false);
+  if (vb2.x) {
+    const float4 r = my_texture.SampleBias(my_sampler, (0.0f).xx, 0.0f);
+  }
+}
+
+void foo_default_initialize() {
+  bool2 vb2 = bool2(false, false);
+  vb2.x = (my_global.z != 0.0f);
+  vb2.x = (asfloat(my_uniform[0].x) == -1.0f);
+  vb2 = (false).xx;
+  if (vb2.x) {
+    const float4 r = my_texture.SampleBias(my_sampler, (0.0f).xx, 0.0f);
+  }
+}

test/tint/bug/tint/1703.wgsl.expected.fxc.hlsl

@@ -0,0 +1,31 @@
+[numthreads(1, 1, 1)]
+void unused_entry_point() {
+  return;
+}
+
+static float4 my_global = float4(0.0f, 0.0f, 0.0f, 0.0f);
+cbuffer cbuffer_my_uniform : register(b0, space0) {
+  uint4 my_uniform[1];
+};
+Texture2D<float4> my_texture : register(t1, space0);
+SamplerState my_sampler : register(s2, space0);
+
+void foo_member_initialize() {
+  bool2 vb2 = bool2(false, false);
+  vb2.x = (my_global.z != 0.0f);
+  vb2.x = (asfloat(my_uniform[0].x) == -1.0f);
+  vb2 = bool2((asfloat(my_uniform[0].x) == -1.0f), false);
+  if (vb2.x) {
+    const float4 r = my_texture.SampleBias(my_sampler, (0.0f).xx, 0.0f);
+  }
+}
+
+void foo_default_initialize() {
+  bool2 vb2 = bool2(false, false);
+  vb2.x = (my_global.z != 0.0f);
+  vb2.x = (asfloat(my_uniform[0].x) == -1.0f);
+  vb2 = (false).xx;
+  if (vb2.x) {
+    const float4 r = my_texture.SampleBias(my_sampler, (0.0f).xx, 0.0f);
+  }
+}

test/tint/bug/tint/1703.wgsl.expected.glsl

@@ -0,0 +1,33 @@
+#version 310 es
+
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+void unused_entry_point() {
+  return;
+}
+vec4 my_global = vec4(0.0f, 0.0f, 0.0f, 0.0f);
+layout(binding = 0, std140) uniform my_uniform_block_ubo {
+  float inner;
+} my_uniform;
+
+uniform highp sampler2D my_texture_my_sampler;
+
+void foo_member_initialize() {
+  bvec2 vb2 = bvec2(false, false);
+  vb2.x = (my_global.z != 0.0f);
+  vb2.x = (my_uniform.inner == -1.0f);
+  vb2 = bvec2((my_uniform.inner == -1.0f), false);
+  if (vb2.x) {
+    vec4 r = texture(my_texture_my_sampler, vec2(0.0f), 0.0f);
+  }
+}
+
+void foo_default_initialize() {
+  bvec2 vb2 = bvec2(false, false);
+  vb2.x = (my_global.z != 0.0f);
+  vb2.x = (my_uniform.inner == -1.0f);
+  vb2 = bvec2(false);
+  if (vb2.x) {
+    vec4 r = texture(my_texture_my_sampler, vec2(0.0f), 0.0f);
+  }
+}
+

test/tint/bug/tint/1703.wgsl.expected.msl

@@ -0,0 +1,23 @@
+#include <metal_stdlib>
+
+using namespace metal;
+void foo_member_initialize(thread float4* const tint_symbol, const constant float* const tint_symbol_1, texture2d<float, access::sample> tint_symbol_2, sampler tint_symbol_3) {
+  bool2 vb2 = false;
+  vb2[0] = ((*(tint_symbol))[2] != 0.0f);
+  vb2[0] = (*(tint_symbol_1) == -1.0f);
+  vb2 = bool2((*(tint_symbol_1) == -1.0f), false);
+  if (vb2[0]) {
+    float4 const r = tint_symbol_2.sample(tint_symbol_3, float2(0.0f), bias(0.0f));
+  }
+}
+
+void foo_default_initialize(thread float4* const tint_symbol_4, const constant float* const tint_symbol_5, texture2d<float, access::sample> tint_symbol_6, sampler tint_symbol_7) {
+  bool2 vb2 = false;
+  vb2[0] = ((*(tint_symbol_4))[2] != 0.0f);
+  vb2[0] = (*(tint_symbol_5) == -1.0f);
+  vb2 = bool2(false);
+  if (vb2[0]) {
+    float4 const r = tint_symbol_6.sample(tint_symbol_7, float2(0.0f), bias(0.0f));
+  }
+}
+

test/tint/bug/tint/1703.wgsl.expected.spvasm

@@ -0,0 +1,123 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 77
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %unused_entry_point "unused_entry_point"
+               OpExecutionMode %unused_entry_point LocalSize 1 1 1
+               OpName %my_global "my_global"
+               OpName %my_uniform_block "my_uniform_block"
+               OpMemberName %my_uniform_block 0 "inner"
+               OpName %my_uniform "my_uniform"
+               OpName %my_texture "my_texture"
+               OpName %my_sampler "my_sampler"
+               OpName %unused_entry_point "unused_entry_point"
+               OpName %foo_member_initialize "foo_member_initialize"
+               OpName %vb2 "vb2"
+               OpName %foo_default_initialize "foo_default_initialize"
+               OpName %vb2_0 "vb2"
+               OpDecorate %my_uniform_block Block
+               OpMemberDecorate %my_uniform_block 0 Offset 0
+               OpDecorate %my_uniform NonWritable
+               OpDecorate %my_uniform DescriptorSet 0
+               OpDecorate %my_uniform Binding 0
+               OpDecorate %my_texture DescriptorSet 0
+               OpDecorate %my_texture Binding 1
+               OpDecorate %my_sampler DescriptorSet 0
+               OpDecorate %my_sampler Binding 2
+      %float = OpTypeFloat 32
+    %v4float = OpTypeVector %float 4
+%_ptr_Private_v4float = OpTypePointer Private %v4float
+          %5 = OpConstantNull %v4float
+  %my_global = OpVariable %_ptr_Private_v4float Private %5
+%my_uniform_block = OpTypeStruct %float
+%_ptr_Uniform_my_uniform_block = OpTypePointer Uniform %my_uniform_block
+ %my_uniform = OpVariable %_ptr_Uniform_my_uniform_block Uniform
+         %11 = OpTypeImage %float 2D 0 0 0 1 Unknown
+%_ptr_UniformConstant_11 = OpTypePointer UniformConstant %11
+ %my_texture = OpVariable %_ptr_UniformConstant_11 UniformConstant
+         %14 = OpTypeSampler
+%_ptr_UniformConstant_14 = OpTypePointer UniformConstant %14
+ %my_sampler = OpVariable %_ptr_UniformConstant_14 UniformConstant
+       %void = OpTypeVoid
+         %15 = OpTypeFunction %void
+       %bool = OpTypeBool
+     %v2bool = OpTypeVector %bool 2
+%_ptr_Function_v2bool = OpTypePointer Function %v2bool
+         %25 = OpConstantNull %v2bool
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+%_ptr_Function_bool = OpTypePointer Function %bool
+     %uint_2 = OpConstant %uint 2
+%_ptr_Private_float = OpTypePointer Private %float
+         %34 = OpConstantNull %float
+%_ptr_Uniform_float = OpTypePointer Uniform %float
+   %float_n1 = OpConstant %float -1
+         %45 = OpConstantNull %bool
+         %54 = OpTypeSampledImage %11
+    %v2float = OpTypeVector %float 2
+         %57 = OpConstantNull %v2float
+%unused_entry_point = OpFunction %void None %15
+         %18 = OpLabel
+               OpReturn
+               OpFunctionEnd
+%foo_member_initialize = OpFunction %void None %15
+         %20 = OpLabel
+        %vb2 = OpVariable %_ptr_Function_v2bool Function %25
+         %29 = OpAccessChain %_ptr_Function_bool %vb2 %uint_0
+         %32 = OpAccessChain %_ptr_Private_float %my_global %uint_2
+         %33 = OpLoad %float %32
+         %35 = OpFOrdNotEqual %bool %33 %34
+               OpStore %29 %35
+         %36 = OpAccessChain %_ptr_Function_bool %vb2 %uint_0
+         %38 = OpAccessChain %_ptr_Uniform_float %my_uniform %uint_0
+         %39 = OpLoad %float %38
+         %41 = OpFOrdEqual %bool %39 %float_n1
+               OpStore %36 %41
+         %42 = OpAccessChain %_ptr_Uniform_float %my_uniform %uint_0
+         %43 = OpLoad %float %42
+         %44 = OpFOrdEqual %bool %43 %float_n1
+         %46 = OpCompositeConstruct %v2bool %44 %45
+               OpStore %vb2 %46
+         %47 = OpAccessChain %_ptr_Function_bool %vb2 %uint_0
+         %48 = OpLoad %bool %47
+               OpSelectionMerge %49 None
+               OpBranchConditional %48 %50 %49
+         %50 = OpLabel
+         %52 = OpLoad %14 %my_sampler
+         %53 = OpLoad %11 %my_texture
+         %55 = OpSampledImage %54 %53 %52
+         %51 = OpImageSampleImplicitLod %v4float %55 %57 Bias %34
+               OpBranch %49
+         %49 = OpLabel
+               OpReturn
+               OpFunctionEnd
+%foo_default_initialize = OpFunction %void None %15
+         %59 = OpLabel
+      %vb2_0 = OpVariable %_ptr_Function_v2bool Function %25
+         %61 = OpAccessChain %_ptr_Function_bool %vb2_0 %uint_0
+         %62 = OpAccessChain %_ptr_Private_float %my_global %uint_2
+         %63 = OpLoad %float %62
+         %64 = OpFOrdNotEqual %bool %63 %34
+               OpStore %61 %64
+         %65 = OpAccessChain %_ptr_Function_bool %vb2_0 %uint_0
+         %66 = OpAccessChain %_ptr_Uniform_float %my_uniform %uint_0
+         %67 = OpLoad %float %66
+         %68 = OpFOrdEqual %bool %67 %float_n1
+               OpStore %65 %68
+               OpStore %vb2_0 %25
+         %69 = OpAccessChain %_ptr_Function_bool %vb2_0 %uint_0
+         %70 = OpLoad %bool %69
+               OpSelectionMerge %71 None
+               OpBranchConditional %70 %72 %71
+         %72 = OpLabel
+         %74 = OpLoad %14 %my_sampler
+         %75 = OpLoad %11 %my_texture
+         %76 = OpSampledImage %54 %75 %74
+         %73 = OpImageSampleImplicitLod %v4float %76 %57 Bias %34
+               OpBranch %71
+         %71 = OpLabel
+               OpReturn
+               OpFunctionEnd

test/tint/bug/tint/1703.wgsl.expected.wgsl

@@ -0,0 +1,27 @@
+var<private> my_global : vec4<f32>;
+
+@group(0) @binding(0) var<uniform> my_uniform : f32;
+
+@group(0) @binding(1) var my_texture : texture_2d<f32>;
+
+@group(0) @binding(2) var my_sampler : sampler;
+
+fn foo_member_initialize() {
+  var vb2 : vec2<bool>;
+  vb2.x = (my_global.z != 0);
+  vb2.x = (my_uniform == -1.0f);
+  vb2 = vec2((my_uniform == -1.0f), false);
+  if (vb2.x) {
+    let r : vec4<f32> = textureSampleBias(my_texture, my_sampler, vec2<f32>(), 0.0);
+  }
+}
+
+fn foo_default_initialize() {
+  var vb2 : vec2<bool>;
+  vb2.x = (my_global.z != 0);
+  vb2.x = (my_uniform == -1.0f);
+  vb2 = vec2<bool>();
+  if (vb2.x) {
+    let r : vec4<f32> = textureSampleBias(my_texture, my_sampler, vec2<f32>(), 0.0);
+  }
+}