From 79ab273f997ea67da45adb005e11f9b39780304f Mon Sep 17 00:00:00 2001
From: Jaebaek Seo <jaebaek@google.com>
Date: Tue, 16 Mar 2021 16:40:51 -0400
Subject: [PATCH] Accept OpImageTexelPointer user in scalar-replacement (#4187)

We have to conduct the scalar replacement for an aggregate with an image
type even when it has OpImageTexelPointer users.
---
 source/opt/scalar_replacement_pass.cpp |  7 +++++
 source/opt/scalar_replacement_pass.h   |  4 +++
 test/opt/scalar_replacement_test.cpp   | 40 ++++++++++++++++++++++++++
 3 files changed, 51 insertions(+)

diff --git a/source/opt/scalar_replacement_pass.cpp b/source/opt/scalar_replacement_pass.cpp
index ba2d06750..4d47bdd8d 100644
--- a/source/opt/scalar_replacement_pass.cpp
+++ b/source/opt/scalar_replacement_pass.cpp
@@ -861,6 +861,9 @@ bool ScalarReplacementPass::CheckUsesRelaxed(const Instruction* inst) const {
           case SpvOpStore:
             if (!CheckStore(user, index)) ok = false;
             break;
+          case SpvOpImageTexelPointer:
+            if (!CheckImageTexelPointer(index)) ok = false;
+            break;
           default:
             ok = false;
             break;
@@ -870,6 +873,10 @@ bool ScalarReplacementPass::CheckUsesRelaxed(const Instruction* inst) const {
   return ok;
 }
 
+bool ScalarReplacementPass::CheckImageTexelPointer(uint32_t index) const {
+  return index == 2u;
+}
+
 bool ScalarReplacementPass::CheckLoad(const Instruction* inst,
                                       uint32_t index) const {
   if (index != 2u) return false;
diff --git a/source/opt/scalar_replacement_pass.h b/source/opt/scalar_replacement_pass.h
index 1f6c92813..9e9f07394 100644
--- a/source/opt/scalar_replacement_pass.h
+++ b/source/opt/scalar_replacement_pass.h
@@ -142,6 +142,10 @@ class ScalarReplacementPass : public Pass {
   // of |inst| and the store is not to volatile memory.
   bool CheckStore(const Instruction* inst, uint32_t index) const;
 
+  // Returns true if |index| is the pointer operand of an OpImageTexelPointer
+  // instruction.
+  bool CheckImageTexelPointer(uint32_t index) const;
+
   // Creates a variable of type |typeId| from the |index|'th element of
   // |varInst|. The new variable is added to |replacements|.  If the variable
   // could not be created, then |nullptr| is appended to |replacements|.
diff --git a/test/opt/scalar_replacement_test.cpp b/test/opt/scalar_replacement_test.cpp
index 10738558b..8115f5fb5 100644
--- a/test/opt/scalar_replacement_test.cpp
+++ b/test/opt/scalar_replacement_test.cpp
@@ -2197,6 +2197,46 @@ OpFunctionEnd
   SinglePassRunAndMatch<ScalarReplacementPass>(text, true);
 }
 
+TEST_F(ScalarReplacementTest, ImageTexelPointer) {
+  // Test whether the scalar replacement correctly checks the
+  // OpImageTexelPointer user of an aggregate with an image type.
+  const std::string text = R"(
+;
+; CHECK: [[imgTy:%\w+]] = OpTypeImage %uint Buffer 2 0 0 2 R32ui
+; CHECK: [[ptrImgTy:%\w+]] = OpTypePointer Function [[imgTy]]
+; CHECK: [[img:%\w+]] = OpVariable [[ptrImgTy]] Function
+; CHECK: [[imgTexelPtr:%\w+]] = OpImageTexelPointer {{%\w+}} [[img]] %uint_0 %uint_0
+; CHECK: OpAtomicIAdd %uint [[imgTexelPtr]] %uint_1 %uint_0 %uint_1
+;
+OpCapability Shader
+OpCapability SampledBuffer
+OpCapability ImageBuffer
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %1 "main"
+OpExecutionMode %1 LocalSize 64 1 1
+%void = OpTypeVoid
+%uint = OpTypeInt 32 0
+%uint_0 = OpConstant %uint 0
+%uint_1 = OpConstant %uint 1
+%_ptr_Image_uint = OpTypePointer Image %uint
+%type_buffer_image = OpTypeImage %uint Buffer 2 0 0 2 R32ui
+%_ptr_Function_type_buffer_image = OpTypePointer Function %type_buffer_image
+%image_struct = OpTypeStruct %type_buffer_image %type_buffer_image
+%_ptr_Function_image_struct = OpTypePointer Function %image_struct
+%func = OpTypeFunction %void
+%1 = OpFunction %void None %func
+%2 = OpLabel
+%3 = OpVariable %_ptr_Function_image_struct Function
+%4 = OpAccessChain %_ptr_Function_type_buffer_image %3 %uint_1
+%5 = OpImageTexelPointer %_ptr_Image_uint %4 %uint_0 %uint_0
+%6 = OpAtomicIAdd %uint %5 %uint_1 %uint_0 %uint_1
+OpReturn
+OpFunctionEnd
+  )";
+
+  SinglePassRunAndMatch<ScalarReplacementPass>(text, false);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools