writer/msl: Avoid generating unnecessary pointers

When moving private and workgroup variables into the entry point,
generate pointers to pass as arguments to sub-functions on demand,
instead of upfront. This removes a bunch of unnecessary dereferences
for accesses inside the entry point, and one function variable.

Change-Id: I7d1aabdf14eae33b569b3316dfc0f9fbd288131e
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/54300
Auto-Submit: James Price <jrprice@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: James Price <jrprice@google.com>

src/transform/msl.cc

@@ -89,8 +89,7 @@ void Msl::HandlePrivateAndWorkgroupVariables(CloneContext& ctx) const {
   // [[stage(compute)]]
   // fn main() {
   //   var<private> v : f32 = 2.0;
-  //   let v_ptr : ptr<private, f32> = &f32;
+  //   foo(&v);
-  //   foo(v_ptr);
   // }
   // ```
 
@@ -127,6 +126,7 @@ void Msl::HandlePrivateAndWorkgroupVariables(CloneContext& ctx) const {
 
   for (auto* func_ast : functions_to_process) {
     auto* func_sem = ctx.src->Sem().Get(func_ast);
+    bool is_entry_point = func_ast->IsEntryPoint();
 
     // Map module-scope variables onto their function-scope replacement.
     std::unordered_map<const sem::Variable*, Symbol> var_to_symbol;
@@ -137,12 +137,12 @@ void Msl::HandlePrivateAndWorkgroupVariables(CloneContext& ctx) const {
         continue;
       }
 
-      // This is the symbol for the pointer that replaces the module-scope var.
+      // This is the symbol for the variable that replaces the module-scope var.
       auto new_var_symbol = ctx.dst->Sym();
 
       auto* store_type = CreateASTTypeFor(&ctx, var->Type()->UnwrapRef());
 
-      if (func_ast->IsEntryPoint()) {
+      if (is_entry_point) {
         // For an entry point, redeclare the variable at function-scope.
         // Disable storage class validation on this variable.
         auto* disable_validation =
@@ -151,16 +151,10 @@ void Msl::HandlePrivateAndWorkgroupVariables(CloneContext& ctx) const {
                 ast::DisabledValidation::kFunctionVarStorageClass);
         auto* constructor = ctx.Clone(var->Declaration()->constructor());
         auto* local_var =
-            ctx.dst->Var(ctx.dst->Sym(), store_type, var->StorageClass(),
+            ctx.dst->Var(new_var_symbol, store_type, var->StorageClass(),
                          constructor, ast::DecorationList{disable_validation});
         ctx.InsertBefore(func_ast->body()->statements(),
                          *func_ast->body()->begin(), ctx.dst->Decl(local_var));
-
-        // Now take the address of the variable.
-        auto* ptr = ctx.dst->Const(new_var_symbol, nullptr,
-                                   ctx.dst->AddressOf(local_var));
-        ctx.InsertBefore(func_ast->body()->statements(),
-                         *func_ast->body()->begin(), ctx.dst->Decl(ptr));
       } else {
         // For a regular function, redeclare the variable as a pointer function
         // parameter.
@@ -169,18 +163,22 @@ void Msl::HandlePrivateAndWorkgroupVariables(CloneContext& ctx) const {
                        ctx.dst->Param(new_var_symbol, ptr_type));
       }
 
-      // Replace all uses of the module-scope variable with the pointer
+      // Replace all uses of the module-scope variable.
-      // replacement (dereferenced).
       for (auto* user : var->Users()) {
         if (user->Stmt()->Function() == func_ast) {
-          ctx.Replace(user->Declaration(), ctx.dst->Deref(new_var_symbol));
+          ast::Expression* expr = ctx.dst->Expr(new_var_symbol);
+          if (!is_entry_point) {
+            // For non-entry points, dereference the pointer argument.
+            expr = ctx.dst->Deref(expr);
+          }
+          ctx.Replace(user->Declaration(), expr);
         }
       }
 
       var_to_symbol[var] = new_var_symbol;
     }
 
-    // Pass the pointers through to any functions that need them.
+    // Pass the variables as pointers to any functions that need them.
     for (auto* call : calls_to_replace[func_ast]) {
       auto* target = ctx.src->AST().Functions().Find(call->func()->symbol());
       auto* target_sem = ctx.src->Sem().Get(target);
@@ -189,8 +187,12 @@ void Msl::HandlePrivateAndWorkgroupVariables(CloneContext& ctx) const {
       for (auto* target_var : target_sem->ReferencedModuleVariables()) {
         if (target_var->StorageClass() == ast::StorageClass::kPrivate ||
             target_var->StorageClass() == ast::StorageClass::kWorkgroup) {
-          ctx.InsertBack(call->params(),
+          ast::Expression* arg = ctx.dst->Expr(var_to_symbol[target_var]);
-                         ctx.dst->Expr(var_to_symbol[target_var]));
+          if (is_entry_point) {
+            // For entry points, pass the address of the variable.
+            arg = ctx.dst->AddressOf(arg);
+          }
+          ctx.InsertBack(call->params(), arg);
         }
       }
     }

src/transform/msl_test.cc

@@ -36,11 +36,9 @@ fn main() {
   auto* expect = R"(
 [[stage(compute)]]
 fn main() {
-  [[internal(disable_validation__function_var_storage_class)]] var<workgroup> tint_symbol_1 : f32;
+  [[internal(disable_validation__function_var_storage_class)]] var<workgroup> tint_symbol : f32;
-  let tint_symbol = &(tint_symbol_1);
+  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol_1 : f32;
-  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol_3 : f32;
+  tint_symbol = tint_symbol_1;
-  let tint_symbol_2 = &(tint_symbol_3);
-  *(tint_symbol) = *(tint_symbol_2);
 }
 )";
 
@@ -91,11 +89,9 @@ fn foo(a : f32, tint_symbol_2 : ptr<private, f32>, tint_symbol_3 : ptr<workgroup
 
 [[stage(compute)]]
 fn main() {
-  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol_5 : f32;
+  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol_4 : f32;
-  let tint_symbol_4 = &(tint_symbol_5);
+  [[internal(disable_validation__function_var_storage_class)]] var<workgroup> tint_symbol_5 : f32;
-  [[internal(disable_validation__function_var_storage_class)]] var<workgroup> tint_symbol_7 : f32;
+  foo(1.0, &(tint_symbol_4), &(tint_symbol_5));
-  let tint_symbol_6 = &(tint_symbol_7);
-  foo(1.0, tint_symbol_4, tint_symbol_6);
 }
 )";
 
@@ -118,11 +114,9 @@ fn main() {
   auto* expect = R"(
 [[stage(compute)]]
 fn main() {
-  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol_1 : f32 = 1.0;
+  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol : f32 = 1.0;
-  let tint_symbol = &(tint_symbol_1);
+  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol_1 : f32 = f32();
-  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol_3 : f32 = f32();
+  let x : f32 = (tint_symbol + tint_symbol_1);
-  let tint_symbol_2 = &(tint_symbol_3);
-  let x : f32 = (*(tint_symbol) + *(tint_symbol_2));
 }
 )";
 
@@ -148,12 +142,10 @@ fn main() {
   auto* expect = R"(
 [[stage(compute)]]
 fn main() {
-  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol_1 : f32;
+  [[internal(disable_validation__function_var_storage_class)]] var<private> tint_symbol : f32;
-  let tint_symbol = &(tint_symbol_1);
+  [[internal(disable_validation__function_var_storage_class)]] var<workgroup> tint_symbol_1 : f32;
-  [[internal(disable_validation__function_var_storage_class)]] var<workgroup> tint_symbol_3 : f32;
+  let p_ptr : ptr<private, f32> = &(tint_symbol);
-  let tint_symbol_2 = &(tint_symbol_3);
+  let w_ptr : ptr<workgroup, f32> = &(tint_symbol_1);
-  let p_ptr : ptr<private, f32> = &(*(tint_symbol));
-  let w_ptr : ptr<workgroup, f32> = &(*(tint_symbol_2));
   let x : f32 = (*(p_ptr) + *(w_ptr));
   *(p_ptr) = x;
 }

src/writer/msl/generator_impl_variable_decl_statement_test.cc

@@ -120,7 +120,7 @@ TEST_F(MslGeneratorImplTest, Emit_VariableDeclStatement_Private) {
   gen.increment_indent();
 
   ASSERT_TRUE(gen.Generate()) << gen.error();
-  EXPECT_THAT(gen.result(), HasSubstr("thread float tint_symbol_2 = 0.0f;\n"));
+  EXPECT_THAT(gen.result(), HasSubstr("thread float tint_symbol_1 = 0.0f;\n"));
 }
 
 TEST_F(MslGeneratorImplTest, Emit_VariableDeclStatement_Initializer_Private) {
@@ -133,7 +133,7 @@ TEST_F(MslGeneratorImplTest, Emit_VariableDeclStatement_Initializer_Private) {
 
   ASSERT_TRUE(gen.Generate()) << gen.error();
   EXPECT_THAT(gen.result(),
-              HasSubstr("thread float tint_symbol_2 = initializer;\n"));
+              HasSubstr("thread float tint_symbol_1 = initializer;\n"));
 }
 
 TEST_F(MslGeneratorImplTest, Emit_VariableDeclStatement_Workgroup) {
@@ -147,7 +147,7 @@ TEST_F(MslGeneratorImplTest, Emit_VariableDeclStatement_Workgroup) {
 
   ASSERT_TRUE(gen.Generate()) << gen.error();
   EXPECT_THAT(gen.result(),
-              HasSubstr("threadgroup float tint_symbol_2 = 0.0f;\n"));
+              HasSubstr("threadgroup float tint_symbol_1 = 0.0f;\n"));
 }
 
 TEST_F(MslGeneratorImplTest, Emit_VariableDeclStatement_Initializer_ZeroVec) {

test/ptr_ref/load/global/i32.spvasm.expected.msl

@@ -2,9 +2,8 @@
 
 using namespace metal;
 kernel void tint_symbol() {
-  thread int tint_symbol_2 = 0;
+  thread int tint_symbol_1 = 0;
-  thread int* const tint_symbol_1 = &(tint_symbol_2);
+  int const x_9 = tint_symbol_1;
-  int const x_9 = *(tint_symbol_1);
   int const x_11 = (x_9 + 1);
   return;
 }

test/ptr_ref/load/global/i32.wgsl.expected.msl

@@ -2,9 +2,8 @@
 
 using namespace metal;
 kernel void tint_symbol() {
-  thread int tint_symbol_2 = 0;
+  thread int tint_symbol_1 = 0;
-  thread int* const tint_symbol_1 = &(tint_symbol_2);
+  int const i = tint_symbol_1;
-  int const i = *(tint_symbol_1);
   int const use = (i + 1);
   return;
 }

test/ptr_ref/load/global/struct_field.spvasm.expected.msl

@@ -6,10 +6,9 @@ struct S {
 };
 
 kernel void tint_symbol() {
-  thread S tint_symbol_2 = {};
+  thread S tint_symbol_1 = {};
-  thread S* const tint_symbol_1 = &(tint_symbol_2);
   int i = 0;
-  int const x_15 = (*(tint_symbol_1)).i;
+  int const x_15 = tint_symbol_1.i;
   i = x_15;
   return;
 }

test/ptr_ref/load/global/struct_field.wgsl.expected.msl

@@ -6,9 +6,8 @@ struct S {
 };
 
 kernel void tint_symbol() {
-  thread S tint_symbol_2 = {};
+  thread S tint_symbol_1 = {};
-  thread S* const tint_symbol_1 = &(tint_symbol_2);
+  int const i = tint_symbol_1.i;
-  int const i = (*(tint_symbol_1)).i;
   return;
 }
 

test/ptr_ref/load/local/ptr_private.wgsl.expected.msl

@@ -2,9 +2,8 @@
 
 using namespace metal;
 kernel void tint_symbol() {
-  thread int tint_symbol_2 = 123;
+  thread int tint_symbol_1 = 123;
-  thread int* const tint_symbol_1 = &(tint_symbol_2);
+  thread int* const p = &(tint_symbol_1);
-  thread int* const p = &(*(tint_symbol_1));
   int const use = (*(p) + 1);
   return;
 }

test/ptr_ref/load/local/ptr_workgroup.wgsl.expected.msl

@@ -2,10 +2,9 @@
 
 using namespace metal;
 kernel void tint_symbol() {
-  threadgroup int tint_symbol_2 = 0;
+  threadgroup int tint_symbol_1 = 0;
-  threadgroup int* const tint_symbol_1 = &(tint_symbol_2);
+  tint_symbol_1 = 123;
-  *(tint_symbol_1) = 123;
+  threadgroup int* const p = &(tint_symbol_1);
-  threadgroup int* const p = &(*(tint_symbol_1));
   int const use = (*(p) + 1);
   return;
 }

test/ptr_ref/store/global/i32.spvasm.expected.msl

@@ -2,10 +2,9 @@
 
 using namespace metal;
 kernel void tint_symbol() {
-  thread int tint_symbol_2 = 0;
+  thread int tint_symbol_1 = 0;
-  thread int* const tint_symbol_1 = &(tint_symbol_2);
+  tint_symbol_1 = 123;
-  *(tint_symbol_1) = 123;
+  tint_symbol_1 = ((100 + 20) + 3);
-  *(tint_symbol_1) = ((100 + 20) + 3);
   return;
 }
 

test/ptr_ref/store/global/i32.wgsl.expected.msl

@@ -2,10 +2,9 @@
 
 using namespace metal;
 kernel void tint_symbol() {
-  thread int tint_symbol_2 = 0;
+  thread int tint_symbol_1 = 0;
-  thread int* const tint_symbol_1 = &(tint_symbol_2);
+  tint_symbol_1 = 123;
-  *(tint_symbol_1) = 123;
+  tint_symbol_1 = ((100 + 20) + 3);
-  *(tint_symbol_1) = ((100 + 20) + 3);
   return;
 }
 

test/ptr_ref/store/global/struct_field.spvasm.expected.msl

@@ -6,9 +6,8 @@ struct S {
 };
 
 kernel void tint_symbol() {
-  thread S tint_symbol_2 = {};
+  thread S tint_symbol_1 = {};
-  thread S* const tint_symbol_1 = &(tint_symbol_2);
+  tint_symbol_1.i = 5;
-  (*(tint_symbol_1)).i = 5;
   return;
 }
 

test/types/module_scope_var.wgsl.expected.msl

@@ -8,38 +8,28 @@ struct tint_array_wrapper_0 {
 };
 
 kernel void tint_symbol() {
-  thread bool tint_symbol_4 = false;
+  thread bool tint_symbol_3 = false;
-  thread bool* const tint_symbol_3 = &(tint_symbol_4);
+  thread int tint_symbol_4 = 0;
-  thread int tint_symbol_6 = 0;
+  thread uint tint_symbol_5 = 0u;
-  thread int* const tint_symbol_5 = &(tint_symbol_6);
+  thread float tint_symbol_6 = 0.0f;
-  thread uint tint_symbol_8 = 0u;
+  thread int2 tint_symbol_7 = 0;
-  thread uint* const tint_symbol_7 = &(tint_symbol_8);
+  thread uint3 tint_symbol_8 = 0u;
-  thread float tint_symbol_10 = 0.0f;
+  thread float4 tint_symbol_9 = 0.0f;
-  thread float* const tint_symbol_9 = &(tint_symbol_10);
+  thread float2x3 tint_symbol_10 = float2x3(0.0f);
-  thread int2 tint_symbol_12 = 0;
+  thread tint_array_wrapper_0 tint_symbol_11 = {0.0f};
-  thread int2* const tint_symbol_11 = &(tint_symbol_12);
+  thread S tint_symbol_12 = {};
-  thread uint3 tint_symbol_14 = 0u;
+  tint_symbol_3 = bool();
-  thread uint3* const tint_symbol_13 = &(tint_symbol_14);
+  tint_symbol_4 = int();
-  thread float4 tint_symbol_16 = 0.0f;
+  tint_symbol_5 = uint();
-  thread float4* const tint_symbol_15 = &(tint_symbol_16);
+  tint_symbol_6 = float();
-  thread float2x3 tint_symbol_18 = float2x3(0.0f);
+  tint_symbol_7 = int2();
-  thread float2x3* const tint_symbol_17 = &(tint_symbol_18);
+  tint_symbol_8 = uint3();
-  thread tint_array_wrapper_0 tint_symbol_20 = {0.0f};
+  tint_symbol_9 = float4();
-  thread tint_array_wrapper_0* const tint_symbol_19 = &(tint_symbol_20);
+  tint_symbol_10 = float2x3();
-  thread S tint_symbol_22 = {};
-  thread S* const tint_symbol_21 = &(tint_symbol_22);
-  *(tint_symbol_3) = bool();
-  *(tint_symbol_5) = int();
-  *(tint_symbol_7) = uint();
-  *(tint_symbol_9) = float();
-  *(tint_symbol_11) = int2();
-  *(tint_symbol_13) = uint3();
-  *(tint_symbol_15) = float4();
-  *(tint_symbol_17) = float2x3();
   tint_array_wrapper_0 const tint_symbol_1 = {};
-  *(tint_symbol_19) = tint_symbol_1;
+  tint_symbol_11 = tint_symbol_1;
   S const tint_symbol_2 = {};
-  *(tint_symbol_21) = tint_symbol_2;
+  tint_symbol_12 = tint_symbol_2;
   return;
 }
 

test/types/module_scope_var_initializers.wgsl.expected.msl

@@ -8,38 +8,28 @@ struct tint_array_wrapper_0 {
 };
 
 kernel void tint_symbol() {
-  thread bool tint_symbol_4 = bool();
+  thread bool tint_symbol_3 = bool();
-  thread bool* const tint_symbol_3 = &(tint_symbol_4);
+  thread int tint_symbol_4 = int();
-  thread int tint_symbol_6 = int();
+  thread uint tint_symbol_5 = uint();
-  thread int* const tint_symbol_5 = &(tint_symbol_6);
+  thread float tint_symbol_6 = float();
-  thread uint tint_symbol_8 = uint();
+  thread int2 tint_symbol_7 = int2();
-  thread uint* const tint_symbol_7 = &(tint_symbol_8);
+  thread uint3 tint_symbol_8 = uint3();
-  thread float tint_symbol_10 = float();
+  thread float4 tint_symbol_9 = float4();
-  thread float* const tint_symbol_9 = &(tint_symbol_10);
+  thread float2x3 tint_symbol_10 = float2x3();
-  thread int2 tint_symbol_12 = int2();
+  thread tint_array_wrapper_0 tint_symbol_11 = {};
-  thread int2* const tint_symbol_11 = &(tint_symbol_12);
+  thread S tint_symbol_12 = {};
-  thread uint3 tint_symbol_14 = uint3();
+  tint_symbol_3 = bool();
-  thread uint3* const tint_symbol_13 = &(tint_symbol_14);
+  tint_symbol_4 = int();
-  thread float4 tint_symbol_16 = float4();
+  tint_symbol_5 = uint();
-  thread float4* const tint_symbol_15 = &(tint_symbol_16);
+  tint_symbol_6 = float();
-  thread float2x3 tint_symbol_18 = float2x3();
+  tint_symbol_7 = int2();
-  thread float2x3* const tint_symbol_17 = &(tint_symbol_18);
+  tint_symbol_8 = uint3();
-  thread tint_array_wrapper_0 tint_symbol_20 = {};
+  tint_symbol_9 = float4();
-  thread tint_array_wrapper_0* const tint_symbol_19 = &(tint_symbol_20);
+  tint_symbol_10 = float2x3();
-  thread S tint_symbol_22 = {};
-  thread S* const tint_symbol_21 = &(tint_symbol_22);
-  *(tint_symbol_3) = bool();
-  *(tint_symbol_5) = int();
-  *(tint_symbol_7) = uint();
-  *(tint_symbol_9) = float();
-  *(tint_symbol_11) = int2();
-  *(tint_symbol_13) = uint3();
-  *(tint_symbol_15) = float4();
-  *(tint_symbol_17) = float2x3();
   tint_array_wrapper_0 const tint_symbol_1 = {};
-  *(tint_symbol_19) = tint_symbol_1;
+  tint_symbol_11 = tint_symbol_1;
   S const tint_symbol_2 = {};
-  *(tint_symbol_21) = tint_symbol_2;
+  tint_symbol_12 = tint_symbol_2;
   return;
 }
 

test/var/private.wgsl.expected.msl

@@ -25,27 +25,23 @@ void outer(thread int* const tint_symbol_4, thread int* const tint_symbol_5) {
 }
 
 kernel void main1() {
-  thread int tint_symbol_7 = 0;
+  thread int tint_symbol_6 = 0;
-  thread int* const tint_symbol_6 = &(tint_symbol_7);
+  tint_symbol_6 = 42;
-  *(tint_symbol_6) = 42;
+  uses_a(&(tint_symbol_6));
-  uses_a(tint_symbol_6);
   return;
 }
 
 kernel void main2() {
-  thread int tint_symbol_9 = 0;
+  thread int tint_symbol_7 = 0;
-  thread int* const tint_symbol_8 = &(tint_symbol_9);
+  tint_symbol_7 = 7;
-  *(tint_symbol_8) = 7;
+  uses_b(&(tint_symbol_7));
-  uses_b(tint_symbol_8);
   return;
 }
 
 kernel void main3() {
-  thread int tint_symbol_11 = 0;
+  thread int tint_symbol_8 = 0;
-  thread int* const tint_symbol_10 = &(tint_symbol_11);
+  thread int tint_symbol_9 = 0;
-  thread int tint_symbol_13 = 0;
+  outer(&(tint_symbol_8), &(tint_symbol_9));
-  thread int* const tint_symbol_12 = &(tint_symbol_13);
-  outer(tint_symbol_10, tint_symbol_12);
   no_uses();
   return;
 }

test/var/workgroup.wgsl.expected.msl

@@ -25,27 +25,23 @@ void outer(threadgroup int* const tint_symbol_4, threadgroup int* const tint_sym
 }
 
 kernel void main1() {
-  threadgroup int tint_symbol_7 = 0;
+  threadgroup int tint_symbol_6 = 0;
-  threadgroup int* const tint_symbol_6 = &(tint_symbol_7);
+  tint_symbol_6 = 42;
-  *(tint_symbol_6) = 42;
+  uses_a(&(tint_symbol_6));
-  uses_a(tint_symbol_6);
   return;
 }
 
 kernel void main2() {
-  threadgroup int tint_symbol_9 = 0;
+  threadgroup int tint_symbol_7 = 0;
-  threadgroup int* const tint_symbol_8 = &(tint_symbol_9);
+  tint_symbol_7 = 7;
-  *(tint_symbol_8) = 7;
+  uses_b(&(tint_symbol_7));
-  uses_b(tint_symbol_8);
   return;
 }
 
 kernel void main3() {
-  threadgroup int tint_symbol_11 = 0;
+  threadgroup int tint_symbol_8 = 0;
-  threadgroup int* const tint_symbol_10 = &(tint_symbol_11);
+  threadgroup int tint_symbol_9 = 0;
-  threadgroup int tint_symbol_13 = 0;
+  outer(&(tint_symbol_8), &(tint_symbol_9));
-  threadgroup int* const tint_symbol_12 = &(tint_symbol_13);
-  outer(tint_symbol_10, tint_symbol_12);
   no_uses();
   return;
 }