Handle image texel pointers in local single-store elim

Author: AnastaZIuk

source/opt/local_single_store_elim_pass.cpp

@@ -305,6 +305,18 @@ bool LocalSingleStoreElimPass::RewriteLoads(
   else
     stored_id = store_inst->GetSingleWordInOperand(kVariableInitIdInIdx);
 
+  const auto get_image_pointer_id = [this](uint32_t value_id) {
+    Instruction* value_inst = context()->get_def_use_mgr()->GetDef(value_id);
+    while (value_inst && value_inst->opcode() == spv::Op::OpCopyObject) {
+      value_id = value_inst->GetSingleWordInOperand(0);
+      value_inst = context()->get_def_use_mgr()->GetDef(value_id);
+    }
+    if (!value_inst || value_inst->opcode() != spv::Op::OpLoad) {
+      return uint32_t{0};
+    }
+    return value_inst->GetSingleWordInOperand(0);
+  };
+
   *all_rewritten = true;
   bool modified = false;
   for (Instruction* use : uses) {
@@ -319,6 +331,17 @@ bool LocalSingleStoreElimPass::RewriteLoads(
       context()->KillNamesAndDecorates(use->result_id());
       context()->ReplaceAllUsesWith(use->result_id(), stored_id);
       context()->KillInst(use);
+    } else if (use->opcode() == spv::Op::OpImageTexelPointer &&
+               dominator_analysis->Dominates(store_inst, use)) {
+      const uint32_t image_ptr_id = get_image_pointer_id(stored_id);
+      if (image_ptr_id == 0) {
+        *all_rewritten = false;
+        continue;
+      }
+      modified = true;
+      context()->ForgetUses(use);
+      use->SetInOperand(0, {image_ptr_id});
+      context()->AnalyzeUses(use);
     } else {
       *all_rewritten = false;
     }

test/opt/local_single_store_elim_test.cpp

@@ -908,6 +908,82 @@ OpFunctionEnd
   SinglePassRunAndCheck<LocalSingleStoreElimPass>(before, after, true, true);
 }
 
+TEST_F(LocalSingleStoreElimTest, RewriteImageTexelPointerImageOperand) {
+  const std::string before = R"(OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %main "main" %g_rwTexture3d
+OpExecutionMode %main LocalSize 256 1 1
+OpSource HLSL 660
+OpName %type_3d_image "type.3d.image"
+OpName %g_rwTexture3d "g_rwTexture3d"
+OpName %main "main"
+OpDecorate %g_rwTexture3d DescriptorSet 0
+OpDecorate %g_rwTexture3d Binding 0
+%uint = OpTypeInt 32 0
+%uint_0 = OpConstant %uint 0
+%uint_1 = OpConstant %uint 1
+%uint_2 = OpConstant %uint 2
+%uint_3 = OpConstant %uint 3
+%v3uint = OpTypeVector %uint 3
+%10 = OpConstantComposite %v3uint %uint_1 %uint_2 %uint_3
+%type_3d_image = OpTypeImage %uint 3D 2 0 0 2 R32ui
+%_ptr_UniformConstant_type_3d_image = OpTypePointer UniformConstant %type_3d_image
+%void = OpTypeVoid
+%13 = OpTypeFunction %void
+%_ptr_Function_type_3d_image = OpTypePointer Function %type_3d_image
+%_ptr_Image_uint = OpTypePointer Image %uint
+%g_rwTexture3d = OpVariable %_ptr_UniformConstant_type_3d_image UniformConstant
+%main = OpFunction %void None %13
+%16 = OpLabel
+%17 = OpVariable %_ptr_Function_type_3d_image Function
+%18 = OpLoad %type_3d_image %g_rwTexture3d
+OpStore %17 %18
+%19 = OpImageTexelPointer %_ptr_Image_uint %17 %10 %uint_0
+%20 = OpAtomicIAdd %uint %19 %uint_1 %uint_0 %uint_1
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string after = R"(OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %main "main" %g_rwTexture3d
+OpExecutionMode %main LocalSize 256 1 1
+OpSource HLSL 660
+OpName %type_3d_image "type.3d.image"
+OpName %g_rwTexture3d "g_rwTexture3d"
+OpName %main "main"
+OpDecorate %g_rwTexture3d DescriptorSet 0
+OpDecorate %g_rwTexture3d Binding 0
+%uint = OpTypeInt 32 0
+%uint_0 = OpConstant %uint 0
+%uint_1 = OpConstant %uint 1
+%uint_2 = OpConstant %uint 2
+%uint_3 = OpConstant %uint 3
+%v3uint = OpTypeVector %uint 3
+%10 = OpConstantComposite %v3uint %uint_1 %uint_2 %uint_3
+%type_3d_image = OpTypeImage %uint 3D 2 0 0 2 R32ui
+%_ptr_UniformConstant_type_3d_image = OpTypePointer UniformConstant %type_3d_image
+%void = OpTypeVoid
+%13 = OpTypeFunction %void
+%_ptr_Function_type_3d_image = OpTypePointer Function %type_3d_image
+%_ptr_Image_uint = OpTypePointer Image %uint
+%g_rwTexture3d = OpVariable %_ptr_UniformConstant_type_3d_image UniformConstant
+%main = OpFunction %void None %13
+%16 = OpLabel
+%17 = OpVariable %_ptr_Function_type_3d_image Function
+%18 = OpLoad %type_3d_image %g_rwTexture3d
+OpStore %17 %18
+%19 = OpImageTexelPointer %_ptr_Image_uint %g_rwTexture3d %10 %uint_0
+%20 = OpAtomicIAdd %uint %19 %uint_1 %uint_0 %uint_1
+OpReturn
+OpFunctionEnd
+)";
+
+  SetAssembleOptions(SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  SinglePassRunAndCheck<LocalSingleStoreElimPass>(before, after, true, true);
+}
+
 // Test that that an unused OpAccessChain between a store and a use does does
 // not hinders the replacement of the use.  We need to check this because
 // local-access-chain-convert does always remove the OpAccessChain instructions