[OPT] prevent private_to_local_pass optimizing double pointer (#6161)

* [OPT] prevent private_to_local_pass optimizing double pointer

Fixes #6156

before this fix the pass will optimize the following code:
```
  %v = OpVariable %_ptr_Private_uint Private %uint_0
%load_thread_local = OpFunction %void None %function_type__1561_
 %37 = OpLabel
 %31 = OpVariable %_ptr_Function__ptr_Private_uint Function
 %32 = OpVariable %_ptr_Function_uint Function
 %33 = OpVariable %_ptr_Function_uint Function
       OpStore %31 %v
 %34 = OpLoad %_ptr_Private_uint %31
 %35 = OpLoad %uint %34
       OpStore %32 %35
 %36 = OpLoad %uint %v
       OpStore %33 %36
       OpReturn
       OpFunctionEnd
```

to this:
```
%load_thread_local = OpFunction %void None %function_type__1561_
 %37 = OpLabel
  %v = OpVariable %_ptr_Function_uint Function %uint_0
 %31 = OpVariable %_ptr_Function__ptr_Private_uint Function
 %32 = OpVariable %_ptr_Function_uint Function
 %33 = OpVariable %_ptr_Function_uint Function
       OpStore %31 %v
       OpStore %32 %uint_0
       OpStore %33 %uint_0
       OpReturn
       OpFunctionEnd
```

resulting in the following error:
```
error: line 81: OpStore Pointer <id> '31[%31]'s type does not match Object <id> '47[%v]'s type.
  OpStore %31 %v

```

private_to_local_pass does not update the storage class of the double
pointer %31 correctly. It currently doesn't have enough context to tell
if %31 is later loaded from and used to store with. For now, just prevent
the optimization if the address of the private variable is used as the
object operand of the OpStore instruction.

with this commit, the spirv code is optimized to this:
```
%v = OpVariable %_ptr_Private_uint Private %uint_0
%load_thread_local = OpFunction %void None %function_type__1561_
         %37 = OpLabel
         %31 = OpVariable %_ptr_Function__ptr_Private_uint Function
         %32 = OpVariable %_ptr_Function_uint Function
         %33 = OpVariable %_ptr_Function_uint Function
               OpStore %31 %v
         %35 = OpLoad %uint %v
               OpStore %32 %35
         %36 = OpLoad %uint %v
               OpStore %33 %36
               OpReturn
               OpFunctionEnd
```

* add a test

* fix tests

* Fix format

---------

Co-authored-by: Steven Perron <[email protected]>

Author: heroseh

source/opt/private_to_local_pass.cpp

@@ -90,13 +90,13 @@ Function* PrivateToLocalPass::FindLocalFunction(const Instruction& inst) const {
   Function* target_function = nullptr;
   context()->get_def_use_mgr()->ForEachUser(
       inst.result_id(),
-      [&target_function, &found_first_use, this](Instruction* use) {
+      [&target_function, &found_first_use, inst, this](Instruction* use) {
         BasicBlock* current_block = context()->get_instr_block(use);
         if (current_block == nullptr) {
           return;
         }
 
-        if (!IsValidUse(use)) {
+        if (!IsValidUse(use, inst.result_id())) {
           found_first_use = true;
           target_function = nullptr;
           return;
@@ -153,21 +153,23 @@ uint32_t PrivateToLocalPass::GetNewType(uint32_t old_type_id) {
   return new_type_id;
 }
 
-bool PrivateToLocalPass::IsValidUse(const Instruction* inst) const {
+bool PrivateToLocalPass::IsValidUse(const Instruction* inst,
+                                    uint32_t private_variable_id) const {
   // The cases in this switch have to match the cases in |UpdateUse|.
   // If we don't know how to update it, it is not valid.
   if (inst->GetCommonDebugOpcode() == CommonDebugInfoDebugGlobalVariable) {
     return true;
   }
   switch (inst->opcode()) {
     case spv::Op::OpLoad:
-    case spv::Op::OpStore:
     case spv::Op::OpImageTexelPointer:  // Treat like a load
       return true;
+    case spv::Op::OpStore:
+      return inst->GetOperand(1).AsId() != private_variable_id;
     case spv::Op::OpAccessChain:
       return context()->get_def_use_mgr()->WhileEachUser(
-          inst, [this](const Instruction* user) {
-            if (!IsValidUse(user)) return false;
+          inst, [this, inst](const Instruction* user) {
+            if (!IsValidUse(user, inst->result_id())) return false;
             return true;
           });
     case spv::Op::OpName:

source/opt/private_to_local_pass.h

@@ -53,7 +53,7 @@ class PrivateToLocalPass : public Pass {
   // Returns true is |inst| is a valid use of a pointer.  In this case, a
   // valid use is one where the transformation is able to rewrite the type to
   // match a change in storage class of the original variable.
-  bool IsValidUse(const Instruction* inst) const;
+  bool IsValidUse(const Instruction* inst, uint32_t private_variable_id) const;
 
   // Given the result id of a pointer type, |old_type_id|, this function
   // returns the id of a the same pointer type except the storage class has

test/opt/private_to_local_test.cpp

@@ -173,6 +173,43 @@ OpCapability SampledBuffer
   SinglePassRunAndMatch<PrivateToLocalPass>(text, false);
 }
 
+TEST_F(PrivateToLocalTest, IgnorePointerToPrivateVariable) {
+  // Should not change because the pointer to the private variable
+  // has been stored inside a function variable.
+  const std::string text = R"(
+               OpCapability Shader
+               OpCapability VariablePointersStorageBuffer
+               OpCapability VariablePointers
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %load_thread_local "load_thread_local" %v
+               OpExecutionMode %load_thread_local LocalSize 8 8 1
+       %uint = OpTypeInt 32 0
+     %uint_0 = OpConstant %uint 0
+       %void = OpTypeVoid
+%function_type__1561_ = OpTypeFunction %void
+%_ptr_Private_uint = OpTypePointer Private %uint
+%_ptr_Function__ptr_Private_uint = OpTypePointer Function %_ptr_Private_uint
+%_ptr_Function_uint = OpTypePointer Function %uint
+          %v = OpVariable %_ptr_Private_uint Private %uint_0
+%load_thread_local = OpFunction %void None %function_type__1561_
+         %37 = OpLabel
+         %31 = OpVariable %_ptr_Function__ptr_Private_uint Function
+         %32 = OpVariable %_ptr_Function_uint Function
+         %33 = OpVariable %_ptr_Function_uint Function
+               OpStore %31 %v
+         %35 = OpLoad %uint %v
+               OpStore %32 %35
+         %36 = OpLoad %uint %v
+               OpStore %33 %36
+               OpReturn
+               OpFunctionEnd
+  )";
+  auto result = SinglePassRunAndDisassemble<PrivateToLocalPass>(
+      text, /* skip_nop = */ true, /* do_validation = */ false);
+  EXPECT_EQ(Pass::Status::SuccessWithoutChange, std::get<1>(result));
+}
+
 TEST_F(PrivateToLocalTest, UsedInTwoFunctions) {
   // Should not change because it is used in multiple functions.
   const std::string text = R"(