Add pass to eliminate dead output components (#4982)

This pass eliminates components of output variables that are not stored
to. Currently this just eliminates trailing components of arrays and
structs, all of which are dead.

WARNING: This pass is not designed to be a standalone pass as it can
cause interface incompatibiliies with the following shader in the
pipeline. See the comment in optimizer.hpp for best usage. This pass is
currently available only through the API; it is not available in the CLI.

This commit also fixes a bug in CreateDecoration() which is part of the
system of generating SPIR-V from the Type manager.

7 files changed, 231 insertions, 44 deletions

diff --git a/include/spirv-tools/optimizer.hpp b/include/spirv-tools/optimizer.hpp
index bca12520f..17a2556af 100644
--- a/include/spirv-tools/optimizer.hpp
+++ b/include/spirv-tools/optimizer.hpp
@@ -888,12 +888,31 @@ Optimizer::PassToken CreateAmdExtToKhrPass();
 Optimizer::PassToken CreateInterpolateFixupPass();
 
 // Removes unused components from composite input variables. Current
-// implementation just removes trailing unused components from input arrays.
-// The pass performs best after maximizing dead code removal. A subsequent dead
-// code elimination pass would be beneficial in removing newly unused component
-// types.
+// implementation just removes trailing unused components from input arrays
+// and structs. The pass performs best after maximizing dead code removal.
+// A subsequent dead code elimination pass would be beneficial in removing
+// newly unused component types.
+//
+// WARNING: This pass can only be safely applied standalone to vertex shaders
+// as it can otherwise cause interface incompatibilities with the preceding
+// shader in the pipeline. If applied to non-vertex shaders, the user should
+// follow by applying EliminateDeadOutputStores and
+// EliminateDeadOutputComponents to the preceding shader.
 Optimizer::PassToken CreateEliminateDeadInputComponentsPass();
 
+// Removes unused components from composite output variables. Current
+// implementation just removes trailing unused components from output arrays
+// and structs. The pass performs best after eliminating dead output stores.
+// A subsequent dead code elimination pass would be beneficial in removing
+// newly unused component types. Currently only supports vertex and fragment
+// shaders.
+//
+// WARNING: This pass cannot be safely applied standalone as it can cause
+// interface incompatibility with the following shader in the pipeline. The
+// user should first apply EliminateDeadInputComponents to the following
+// shader, then apply EliminateDeadOutputStores to this shader.
+Optimizer::PassToken CreateEliminateDeadOutputComponentsPass();
+
 // Analyzes shader and populates |live_locs| and |live_builtins|. Best results
 // will be obtained if shader has all dead code eliminated first. |live_locs|
 // and |live_builtins| are subsequently used when calling
diff --git a/source/opt/eliminate_dead_input_components_pass.cpp b/source/opt/eliminate_dead_input_components_pass.cpp
index de75abd59..f31b567b4 100644
--- a/source/opt/eliminate_dead_input_components_pass.cpp
+++ b/source/opt/eliminate_dead_input_components_pass.cpp
@@ -28,7 +28,6 @@ namespace {
 const uint32_t kAccessChainBaseInIdx = 0;
 const uint32_t kAccessChainIndex0InIdx = 1;
 const uint32_t kConstantValueInIdx = 0;
-const uint32_t kVariableStorageClassInIdx = 0;
 
 }  // namespace
 
@@ -42,7 +41,7 @@ Pass::Status EliminateDeadInputComponentsPass::Process() {
   analysis::DefUseManager* def_use_mgr = context()->get_def_use_mgr();
   analysis::TypeManager* type_mgr = context()->get_type_mgr();
   bool modified = false;
-  std::vector<std::pair<Instruction*, unsigned>> arrays_to_change;
+  std::vector<Instruction*> vars_to_move;
   for (auto& var : context()->types_values()) {
     if (var.opcode() != spv::Op::OpVariable) {
       continue;
@@ -52,8 +51,14 @@ Pass::Status EliminateDeadInputComponentsPass::Process() {
     if (ptr_type == nullptr) {
       continue;
     }
-    if (ptr_type->storage_class() != spv::StorageClass::Input) {
-      continue;
+    if (output_instead_) {
+      if (ptr_type->storage_class() != spv::StorageClass::Output) {
+        continue;
+      }
+    } else {
+      if (ptr_type->storage_class() != spv::StorageClass::Input) {
+        continue;
+      }
     }
     const analysis::Array* arr_type = ptr_type->pointee_type()->AsArray();
     if (arr_type != nullptr) {
@@ -69,6 +74,7 @@ Pass::Status EliminateDeadInputComponentsPass::Process() {
       unsigned max_idx = FindMaxIndex(var, original_max);
       if (max_idx != original_max) {
         ChangeArrayLength(var, max_idx + 1);
+        vars_to_move.push_back(&var);
         modified = true;
       }
       continue;
@@ -80,10 +86,20 @@ Pass::Status EliminateDeadInputComponentsPass::Process() {
     unsigned max_idx = FindMaxIndex(var, original_max);
     if (max_idx != original_max) {
       ChangeStructLength(var, max_idx + 1);
+      vars_to_move.push_back(&var);
       modified = true;
     }
   }
 
+  // Move changed vars after their new type instruction to preserve backward
+  // referencing
+  for (auto var : vars_to_move) {
+    auto type_id = var->type_id();
+    auto type_inst = def_use_mgr->GetDef(type_id);
+    var->RemoveFromList();
+    var->InsertAfter(type_inst);
+  }
+
   return modified ? Status::SuccessWithChange : Status::SuccessWithoutChange;
 }
 
@@ -95,7 +111,7 @@ unsigned EliminateDeadInputComponentsPass::FindMaxIndex(Instruction& var,
   context()->get_def_use_mgr()->WhileEachUser(
       var.result_id(), [&max, &seen_non_const_ac, var, this](Instruction* use) {
         auto use_opcode = use->opcode();
-        if (use_opcode == spv::Op::OpLoad ||
+        if (use_opcode == spv::Op::OpLoad || use_opcode == spv::Op::OpStore ||
             use_opcode == spv::Op::OpCopyMemory ||
             use_opcode == spv::Op::OpCopyMemorySized ||
             use_opcode == spv::Op::OpCopyObject) {
@@ -139,18 +155,11 @@ void EliminateDeadInputComponentsPass::ChangeArrayLength(Instruction& arr_var,
   analysis::Array new_arr_ty(arr_ty->element_type(),
                              arr_ty->GetConstantLengthInfo(length_id, length));
   analysis::Type* reg_new_arr_ty = type_mgr->GetRegisteredType(&new_arr_ty);
-  analysis::Pointer new_ptr_ty(reg_new_arr_ty, spv::StorageClass::Input);
+  analysis::Pointer new_ptr_ty(reg_new_arr_ty, ptr_type->storage_class());
   analysis::Type* reg_new_ptr_ty = type_mgr->GetRegisteredType(&new_ptr_ty);
   uint32_t new_ptr_ty_id = type_mgr->GetTypeInstruction(reg_new_ptr_ty);
   arr_var.SetResultType(new_ptr_ty_id);
   def_use_mgr->AnalyzeInstUse(&arr_var);
-  // Move arr_var after its new type to preserve order
-  USE_ASSERT(spv::StorageClass(arr_var.GetSingleWordInOperand(
-                 kVariableStorageClassInIdx)) != spv::StorageClass::Function &&
-             "cannot move Function variable");
-  Instruction* new_ptr_ty_inst = def_use_mgr->GetDef(new_ptr_ty_id);
-  arr_var.RemoveFromList();
-  arr_var.InsertAfter(new_ptr_ty_inst);
 }
 
 void EliminateDeadInputComponentsPass::ChangeStructLength(
@@ -165,25 +174,25 @@ void EliminateDeadInputComponentsPass::ChangeStructLength(
   for (unsigned u = 0; u < length; ++u)
     new_elt_types.push_back(orig_elt_types[u]);
   analysis::Struct new_struct_ty(new_elt_types);
+  uint32_t old_struct_ty_id = type_mgr->GetTypeInstruction(struct_ty);
+  std::vector<Instruction*> decorations =
+      context()->get_decoration_mgr()->GetDecorationsFor(old_struct_ty_id,
+                                                         true);
+  for (auto dec : decorations) {
+    if (dec->opcode() == spv::Op::OpMemberDecorate) {
+      uint32_t midx = dec->GetSingleWordInOperand(1);
+      if (midx >= length) continue;
+    }
+    type_mgr->AttachDecoration(*dec, &new_struct_ty);
+  }
   analysis::Type* reg_new_struct_ty =
       type_mgr->GetRegisteredType(&new_struct_ty);
-  uint32_t new_struct_ty_id = type_mgr->GetTypeInstruction(reg_new_struct_ty);
-  uint32_t old_struct_ty_id = type_mgr->GetTypeInstruction(struct_ty);
-  analysis::DecorationManager* deco_mgr = context()->get_decoration_mgr();
-  deco_mgr->CloneDecorations(old_struct_ty_id, new_struct_ty_id);
-  analysis::Pointer new_ptr_ty(reg_new_struct_ty, spv::StorageClass::Input);
+  analysis::Pointer new_ptr_ty(reg_new_struct_ty, ptr_type->storage_class());
   analysis::Type* reg_new_ptr_ty = type_mgr->GetRegisteredType(&new_ptr_ty);
   uint32_t new_ptr_ty_id = type_mgr->GetTypeInstruction(reg_new_ptr_ty);
   struct_var.SetResultType(new_ptr_ty_id);
   analysis::DefUseManager* def_use_mgr = context()->get_def_use_mgr();
   def_use_mgr->AnalyzeInstUse(&struct_var);
-  // Move struct_var after its new type to preserve order
-  USE_ASSERT(spv::StorageClass(struct_var.GetSingleWordInOperand(
-                 kVariableStorageClassInIdx)) != spv::StorageClass::Function &&
-             "cannot move Function variable");
-  Instruction* new_ptr_ty_inst = def_use_mgr->GetDef(new_ptr_ty_id);
-  struct_var.RemoveFromList();
-  struct_var.InsertAfter(new_ptr_ty_inst);
 }
 
 }  // namespace opt
diff --git a/source/opt/eliminate_dead_input_components_pass.h b/source/opt/eliminate_dead_input_components_pass.h
index a3a133c2b..16b454575 100644
--- a/source/opt/eliminate_dead_input_components_pass.h
+++ b/source/opt/eliminate_dead_input_components_pass.h
@@ -28,7 +28,8 @@ namespace opt {
 // See optimizer.hpp for documentation.
 class EliminateDeadInputComponentsPass : public Pass {
  public:
-  explicit EliminateDeadInputComponentsPass() {}
+  explicit EliminateDeadInputComponentsPass(bool output_instead = false)
+      : output_instead_(output_instead) {}
 
   const char* name() const override {
     return "eliminate-dead-input-components";
@@ -57,6 +58,9 @@ class EliminateDeadInputComponentsPass : public Pass {
 
   // Change the length of the struct |struct_var| to |length|
   void ChangeStructLength(Instruction& struct_var, unsigned length);
+
+  // Process output variables instead
+  bool output_instead_;
 };
 
 }  // namespace opt
diff --git a/source/opt/optimizer.cpp b/source/opt/optimizer.cpp
index 75a47784c..7583bd1c6 100644
--- a/source/opt/optimizer.cpp
+++ b/source/opt/optimizer.cpp
@@ -1034,6 +1034,12 @@ Optimizer::PassToken CreateEliminateDeadOutputStoresPass(
       MakeUnique<opt::EliminateDeadOutputStoresPass>(live_locs, live_builtins));
 }
 
+Optimizer::PassToken CreateEliminateDeadOutputComponentsPass() {
+  return MakeUnique<Optimizer::PassToken::Impl>(
+      MakeUnique<opt::EliminateDeadInputComponentsPass>(
+          /* output_instead */ true));
+}
+
 Optimizer::PassToken CreateConvertToSampledImagePass(
     const std::vector<opt::DescriptorSetAndBinding>&
         descriptor_set_binding_pairs) {
diff --git a/source/opt/type_manager.cpp b/source/opt/type_manager.cpp
index 4568c0d06..d25daa901 100644
--- a/source/opt/type_manager.cpp
+++ b/source/opt/type_manager.cpp
@@ -476,7 +476,7 @@ void TypeManager::AttachDecorations(uint32_t id, const Type* type) {
     for (auto pair : structTy->element_decorations()) {
       uint32_t element = pair.first;
       for (auto vec : pair.second) {
-        CreateDecoration(id, vec, element);
+        CreateDecoration(id, vec, /* is_member */ true, element);
       }
     }
   }
@@ -484,10 +484,10 @@ void TypeManager::AttachDecorations(uint32_t id, const Type* type) {
 
 void TypeManager::CreateDecoration(uint32_t target,
                                    const std::vector<uint32_t>& decoration,
-                                   uint32_t element) {
+                                   bool is_member, uint32_t element) {
   std::vector<Operand> ops;
   ops.push_back(Operand(SPV_OPERAND_TYPE_ID, {target}));
-  if (element != 0) {
+  if (is_member) {
     ops.push_back(Operand(SPV_OPERAND_TYPE_LITERAL_INTEGER, {element}));
   }
   ops.push_back(Operand(SPV_OPERAND_TYPE_DECORATION, {decoration[0]}));
@@ -495,9 +495,8 @@ void TypeManager::CreateDecoration(uint32_t target,
     ops.push_back(Operand(SPV_OPERAND_TYPE_LITERAL_INTEGER, {decoration[i]}));
   }
   context()->AddAnnotationInst(MakeUnique<Instruction>(
-      context(),
-      (element == 0 ? spv::Op::OpDecorate : spv::Op::OpMemberDecorate), 0, 0,
-      ops));
+      context(), (is_member ? spv::Op::OpMemberDecorate : spv::Op::OpDecorate),
+      0, 0, ops));
   Instruction* inst = &*--context()->annotation_end();
   context()->get_def_use_mgr()->AnalyzeInstUse(inst);
 }
diff --git a/source/opt/type_manager.h b/source/opt/type_manager.h
index 2417d4993..c49e19322 100644
--- a/source/opt/type_manager.h
+++ b/source/opt/type_manager.h
@@ -139,6 +139,11 @@ class TypeManager {
   const Type* GetMemberType(const Type* parent_type,
                             const std::vector<uint32_t>& access_chain);
 
+  // Attaches the decoration encoded in |inst| to |type|. Does nothing if the
+  // given instruction is not a decoration instruction. Assumes the target is
+  // |type| (e.g. should be called in loop of |type|'s decorations).
+  void AttachDecoration(const Instruction& inst, Type* type);
+
   Type* GetUIntType() {
     Integer int_type(32, false);
     return GetRegisteredType(&int_type);
@@ -243,19 +248,15 @@ class TypeManager {
 
   // Create the annotation instruction.
   //
-  // If |element| is zero, an OpDecorate is created, other an OpMemberDecorate
-  // is created. The annotation is registered with the DefUseManager and the
-  // DecorationManager.
+  // If |is_member| is false, an OpDecorate of |decoration| on |id| is created,
+  // otherwise an OpMemberDecorate is created at |element|. The annotation is
+  // registered with the DefUseManager and the DecorationManager.
   void CreateDecoration(uint32_t id, const std::vector<uint32_t>& decoration,
-                        uint32_t element = 0);
+                        bool is_member = false, uint32_t element = 0);
 
   // Creates and returns a type from the given SPIR-V |inst|. Returns nullptr if
   // the given instruction is not for defining a type.
   Type* RecordIfTypeDefinition(const Instruction& inst);
-  // Attaches the decoration encoded in |inst| to |type|. Does nothing if the
-  // given instruction is not a decoration instruction. Assumes the target is
-  // |type| (e.g. should be called in loop of |type|'s decorations).
-  void AttachDecoration(const Instruction& inst, Type* type);
 
   // Returns an equivalent pointer to |type| built in terms of pointers owned by
   // |type_pool_|. For example, if |type| is a vec3 of bool, it will be rebuilt
diff --git a/test/opt/eliminate_dead_input_components_test.cpp b/test/opt/eliminate_dead_input_components_test.cpp
index 822914a86..2c2e63612 100644
--- a/test/opt/eliminate_dead_input_components_test.cpp
+++ b/test/opt/eliminate_dead_input_components_test.cpp
@@ -463,6 +463,155 @@ TEST_F(ElimDeadInputComponentsTest, ElimStructMember) {
   SinglePassRunAndMatch<EliminateDeadInputComponentsPass>(text, true);
 }
 
+TEST_F(ElimDeadInputComponentsTest, ElimOutputStructMember) {
+  // Should eliminate uv from Vertex and all but gl_Position from gl_PerVertex
+  //
+  // #version 450
+  //
+  // out Vertex {
+  //   vec4 Cd;
+  //   vec2 uv;
+  // } oVert;
+  //
+  // in vec3 P;
+  //
+  // void main()
+  // {
+  //   vec4 worldSpacePos = vec4(P, 1);
+  //   oVert.Cd = vec4(1, 0.5, 0, 1);
+  //   gl_Position = worldSpacePos;
+  // }
+
+  const std::string text = R"(
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Vertex %main "main" %P %oVert %_
+               OpSource GLSL 450
+               OpName %main "main"
+               OpName %P "P"
+               OpName %Vertex "Vertex"
+               OpMemberName %Vertex 0 "Cd"
+               OpMemberName %Vertex 1 "uv"
+               OpName %oVert "oVert"
+               OpName %gl_PerVertex "gl_PerVertex"
+               OpMemberName %gl_PerVertex 0 "gl_Position"
+               OpMemberName %gl_PerVertex 1 "gl_PointSize"
+               OpMemberName %gl_PerVertex 2 "gl_ClipDistance"
+               OpMemberName %gl_PerVertex 3 "gl_CullDistance"
+               OpName %_ ""
+               OpDecorate %P Location 0
+               OpDecorate %Vertex Block
+               OpDecorate %oVert Location 0
+               OpMemberDecorate %gl_PerVertex 0 BuiltIn Position
+               OpMemberDecorate %gl_PerVertex 1 BuiltIn PointSize
+               OpMemberDecorate %gl_PerVertex 2 BuiltIn ClipDistance
+               OpMemberDecorate %gl_PerVertex 3 BuiltIn CullDistance
+               OpDecorate %gl_PerVertex Block
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+      %float = OpTypeFloat 32
+    %v4float = OpTypeVector %float 4
+    %v3float = OpTypeVector %float 3
+%_ptr_Input_v3float = OpTypePointer Input %v3float
+          %P = OpVariable %_ptr_Input_v3float Input
+    %float_1 = OpConstant %float 1
+    %v2float = OpTypeVector %float 2
+     %Vertex = OpTypeStruct %v4float %v2float
+%_ptr_Output_Vertex = OpTypePointer Output %Vertex
+      %oVert = OpVariable %_ptr_Output_Vertex Output
+        %int = OpTypeInt 32 1
+      %int_0 = OpConstant %int 0
+  %float_0_5 = OpConstant %float 0.5
+    %float_0 = OpConstant %float 0
+         %27 = OpConstantComposite %v4float %float_1 %float_0_5 %float_0 %float_1
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+       %uint = OpTypeInt 32 0
+     %uint_1 = OpConstant %uint 1
+%_arr_float_uint_1 = OpTypeArray %float %uint_1
+%gl_PerVertex = OpTypeStruct %v4float %float %_arr_float_uint_1 %_arr_float_uint_1
+%_ptr_Output_gl_PerVertex = OpTypePointer Output %gl_PerVertex
+          %_ = OpVariable %_ptr_Output_gl_PerVertex Output
+; CHECK: %Vertex = OpTypeStruct %v4float %v2float
+; CHECK: %gl_PerVertex = OpTypeStruct %v4float %float %_arr_float_uint_1 %_arr_float_uint_1
+; CHECK: %_ptr_Output_gl_PerVertex = OpTypePointer Output %gl_PerVertex
+; CHECK: [[sty:%\w+]] = OpTypeStruct %v4float
+; CHECK: [[pty:%\w+]] = OpTypePointer Output [[sty]]
+; CHECK: %oVert = OpVariable [[pty]] Output
+; CHECK: [[sty2:%\w+]] = OpTypeStruct %v4float
+; CHECK: [[pty2:%\w+]] = OpTypePointer Output [[sty2]]
+; CHECK: %_ = OpVariable [[pty2]] Output
+       %main = OpFunction %void None %3
+          %5 = OpLabel
+         %13 = OpLoad %v3float %P
+         %15 = OpCompositeExtract %float %13 0
+         %16 = OpCompositeExtract %float %13 1
+         %17 = OpCompositeExtract %float %13 2
+         %18 = OpCompositeConstruct %v4float %15 %16 %17 %float_1
+         %29 = OpAccessChain %_ptr_Output_v4float %oVert %int_0
+               OpStore %29 %27
+         %37 = OpAccessChain %_ptr_Output_v4float %_ %int_0
+               OpStore %37 %18
+               OpReturn
+               OpFunctionEnd
+)";
+
+  SetTargetEnv(SPV_ENV_VULKAN_1_3);
+  SetAssembleOptions(SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
+  SinglePassRunAndMatch<EliminateDeadInputComponentsPass>(text, true, true);
+}
+
+TEST_F(ElimDeadInputComponentsTest, ElimOutputArrayMembers) {
+  // Should reduce to uv[2]
+  //
+  // #version 450
+  //
+  // layout(location = 0) out vec2 uv[8];
+  //
+  // void main()
+  // {
+  //     uv[1] = vec2(1, 0.5);
+  // }
+
+  const std::string text = R"(
+               OpCapability Shader
+          %1 = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Vertex %main "main" %uv
+               OpSource GLSL 450
+               OpName %main "main"
+               OpName %uv "uv"
+               OpDecorate %uv Location 0
+       %void = OpTypeVoid
+          %3 = OpTypeFunction %void
+      %float = OpTypeFloat 32
+    %v2float = OpTypeVector %float 2
+       %uint = OpTypeInt 32 0
+     %uint_8 = OpConstant %uint 8
+%_arr_v2float_uint_8 = OpTypeArray %v2float %uint_8
+%_ptr_Output__arr_v2float_uint_8 = OpTypePointer Output %_arr_v2float_uint_8
+         %uv = OpVariable %_ptr_Output__arr_v2float_uint_8 Output
+;CHECK-NOT:         %uv = OpVariable %_ptr_Output__arr_v2float_uint_8 Output
+;CHECK:             %uv = OpVariable %_ptr_Output__arr_v2float_uint_2 Output
+        %int = OpTypeInt 32 1
+      %int_1 = OpConstant %int 1
+    %float_1 = OpConstant %float 1
+  %float_0_5 = OpConstant %float 0.5
+         %17 = OpConstantComposite %v2float %float_1 %float_0_5
+%_ptr_Output_v2float = OpTypePointer Output %v2float
+       %main = OpFunction %void None %3
+          %5 = OpLabel
+         %19 = OpAccessChain %_ptr_Output_v2float %uv %int_1
+               OpStore %19 %17
+               OpReturn
+               OpFunctionEnd
+)";
+
+  SetTargetEnv(SPV_ENV_VULKAN_1_3);
+  SetAssembleOptions(SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
+  SinglePassRunAndMatch<EliminateDeadInputComponentsPass>(text, true, true);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools