opt: add pass to split combined image samplers (#6035)

* opt: add pass to split combined image samplers

Remap as follows:
- A combined-image sampler variable is replaced with a image variable
  and sampler variable at the same DescriptorSet and Binding, and
  copying other decoraitons.
- A combined-image sampler function parameter is replaced with
  a image parameter and sampler parameter.
- Recursively applies to pointer-to-combined, array-of-combined, and
  runtime-array-of-combined.
- Remaps function types as needed.

Removes the type definitions for pointer-to, array-of, and rt-array-of
combined types.

Maintains def-use analysis, and the type manager.

Bug: crbug.com/398231475

Author: dneto0

Android.mk

@@ -177,6 +177,7 @@ SPVTOOLS_OPT_SRC_FILES := \
 		source/opt/scalar_replacement_pass.cpp \
 		source/opt/set_spec_constant_default_value_pass.cpp \
 		source/opt/simplification_pass.cpp \
+		source/opt/split_combined_image_sampler_pass.cpp \
 		source/opt/spread_volatile_semantics.cpp \
 		source/opt/ssa_rewrite_pass.cpp \
 		source/opt/strength_reduction_pass.cpp \

BUILD.bazel

@@ -201,6 +201,7 @@ cc_library(
     deps = [
         ":spirv_tools_internal",
         "@spirv_headers//:spirv_common_headers",
+        "@spirv_headers//:spirv_c_headers",
     ],
 )
 

BUILD.gn

@@ -781,6 +781,8 @@ static_library("spvtools_opt") {
     "source/opt/set_spec_constant_default_value_pass.h",
     "source/opt/simplification_pass.cpp",
     "source/opt/simplification_pass.h",
+    "source/opt/split_combined_image_sampler_pass.cpp",
+    "source/opt/split_combined_image_sampler_pass.h",
     "source/opt/spread_volatile_semantics.cpp",
     "source/opt/spread_volatile_semantics.h",
     "source/opt/ssa_rewrite_pass.cpp",

include/spirv-tools/optimizer.hpp

@@ -240,7 +240,7 @@ class SPIRV_TOOLS_EXPORT Optimizer {
 
  private:
   struct SPIRV_TOOLS_LOCAL Impl;  // Opaque struct for holding internal data.
-  std::unique_ptr<Impl> impl_;  // Unique pointer to internal data.
+  std::unique_ptr<Impl> impl_;    // Unique pointer to internal data.
 };
 
 // Creates a null pass.
@@ -968,6 +968,11 @@ Optimizer::PassToken CreateInvocationInterlockPlacementPass();
 // Creates a pass to add/remove maximal reconvergence execution mode.
 // This pass either adds or removes maximal reconvergence from all entry points.
 Optimizer::PassToken CreateModifyMaximalReconvergencePass(bool add);
+
+// Creates a pass to split combined image+sampler variables and function
+// parameters into separate image and sampler parts. Binding numbers and
+// other decorations are copied.
+Optimizer::PassToken CreateSplitCombinedImageSamplerPass();
 }  // namespace spvtools
 
 #endif  // INCLUDE_SPIRV_TOOLS_OPTIMIZER_HPP_

source/opt/CMakeLists.txt

@@ -114,6 +114,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   scalar_replacement_pass.h
   set_spec_constant_default_value_pass.h
   simplification_pass.h
+  split_combined_image_sampler_pass.h
   spread_volatile_semantics.h
   ssa_rewrite_pass.h
   strength_reduction_pass.h
@@ -230,6 +231,7 @@ set(SPIRV_TOOLS_OPT_SOURCES
   scalar_replacement_pass.cpp
   set_spec_constant_default_value_pass.cpp
   simplification_pass.cpp
+  split_combined_image_sampler_pass.cpp
   spread_volatile_semantics.cpp
   ssa_rewrite_pass.cpp
   strength_reduction_pass.cpp

source/opt/function.h

@@ -17,6 +17,7 @@
 
 #include <algorithm>
 #include <functional>
+#include <iterator>
 #include <memory>
 #include <string>
 #include <unordered_set>
@@ -39,6 +40,7 @@ class Function {
  public:
   using iterator = UptrVectorIterator<BasicBlock>;
   using const_iterator = UptrVectorIterator<BasicBlock, true>;
+  using ParamList = std::vector<std::unique_ptr<Instruction>>;
 
   // Creates a function instance declared by the given OpFunction instruction
   // |def_inst|.
@@ -77,6 +79,23 @@ class Function {
   // Does nothing if the function doesn't have such a parameter.
   inline void RemoveParameter(uint32_t id);
 
+  // Rewrites the function parameters by calling a replacer callback.
+  // The replacer takes two parameters: an expiring unique pointer to a current
+  // instruction, and a back-inserter into a new list of unique pointers to
+  // instructions.  The replacer is called for each current parameter, in order.
+  // Not valid to call while also iterating through the parameter list, e.g.
+  // within the ForEachParam method.
+  using RewriteParamFn = std::function<void(
+      std::unique_ptr<Instruction>&&, std::back_insert_iterator<ParamList>&)>;
+  void RewriteParams(RewriteParamFn& replacer) {
+    ParamList new_params;
+    auto appender = std::back_inserter(new_params);
+    for (auto& param : params_) {
+      replacer(std::move(param), appender);
+    }
+    params_ = std::move(new_params);
+  }
+
   // Saves the given function end instruction.
   inline void SetFunctionEnd(std::unique_ptr<Instruction> end_inst);
 
@@ -197,7 +216,7 @@ class Function {
   // The OpFunction instruction that begins the definition of this function.
   std::unique_ptr<Instruction> def_inst_;
   // All parameters to this function.
-  std::vector<std::unique_ptr<Instruction>> params_;
+  ParamList params_;
   // All debug instructions in this function's header.
   InstructionList debug_insts_in_header_;
   // All basic blocks inside this function in specification order

source/opt/ir_builder.h

@@ -15,6 +15,7 @@
 #ifndef SOURCE_OPT_IR_BUILDER_H_
 #define SOURCE_OPT_IR_BUILDER_H_
 
+#include <cassert>
 #include <limits>
 #include <memory>
 #include <utility>
@@ -480,8 +481,11 @@ class InstructionBuilder {
     return AddInstruction(std::move(select));
   }
 
-  Instruction* AddAccessChain(uint32_t type_id, uint32_t base_ptr_id,
-                              std::vector<uint32_t> ids) {
+  Instruction* AddOpcodeAccessChain(spv::Op opcode, uint32_t type_id,
+                                    uint32_t base_ptr_id,
+                                    const std::vector<uint32_t>& ids) {
+    assert(opcode == spv::Op::OpAccessChain ||
+           opcode == spv::Op::OpInBoundsAccessChain);
     std::vector<Operand> operands;
     operands.push_back({SPV_OPERAND_TYPE_ID, {base_ptr_id}});
 
@@ -490,12 +494,22 @@ class InstructionBuilder {
     }
 
     // TODO(1841): Handle id overflow.
-    std::unique_ptr<Instruction> new_inst(
-        new Instruction(GetContext(), spv::Op::OpAccessChain, type_id,
-                        GetContext()->TakeNextId(), operands));
+    std::unique_ptr<Instruction> new_inst(new Instruction(
+        GetContext(), opcode, type_id, GetContext()->TakeNextId(), operands));
     return AddInstruction(std::move(new_inst));
   }
 
+  Instruction* AddAccessChain(uint32_t type_id, uint32_t base_ptr_id,
+                              const std::vector<uint32_t>& ids) {
+    return AddOpcodeAccessChain(spv::Op::OpAccessChain, type_id, base_ptr_id,
+                                ids);
+  }
+  Instruction* AddInBoundsAccessChain(uint32_t type_id, uint32_t base_ptr_id,
+                                      const std::vector<uint32_t>& ids) {
+    return AddOpcodeAccessChain(spv::Op::OpInBoundsAccessChain, type_id,
+                                base_ptr_id, ids);
+  }
+
   Instruction* AddLoad(uint32_t type_id, uint32_t base_ptr_id,
                        uint32_t alignment = 0) {
     std::vector<Operand> operands;
@@ -514,6 +528,16 @@ class InstructionBuilder {
     return AddInstruction(std::move(new_inst));
   }
 
+  Instruction* AddCopyObject(uint32_t type_id, uint32_t value_id) {
+    std::vector<Operand> operands{{SPV_OPERAND_TYPE_ID, {value_id}}};
+
+    // TODO(1841): Handle id overflow.
+    std::unique_ptr<Instruction> new_inst(
+        new Instruction(GetContext(), spv::Op::OpCopyObject, type_id,
+                        GetContext()->TakeNextId(), operands));
+    return AddInstruction(std::move(new_inst));
+  }
+
   Instruction* AddVariable(uint32_t type_id, uint32_t storage_class) {
     std::vector<Operand> operands;
     operands.push_back({SPV_OPERAND_TYPE_STORAGE_CLASS, {storage_class}});
@@ -572,6 +596,26 @@ class InstructionBuilder {
     return AddInstruction(std::move(new_inst));
   }
 
+  Instruction* AddDecoration(uint32_t target_id, spv::Decoration d,
+                             const std::vector<uint32_t>& literals) {
+    std::vector<Operand> operands;
+    operands.push_back({SPV_OPERAND_TYPE_ID, {target_id}});
+    operands.push_back({SPV_OPERAND_TYPE_DECORATION, {uint32_t(d)}});
+    for (uint32_t literal : literals) {
+      operands.push_back({SPV_OPERAND_TYPE_LITERAL_INTEGER, {literal}});
+    }
+
+    std::unique_ptr<Instruction> new_inst(
+        new Instruction(GetContext(), spv::Op::OpDecorate, 0, 0, operands));
+    // Decorations are annotation instructions.  Add it via the IR context,
+    // so the decoration manager will be updated.
+    // Decorations don't belong to basic blocks, so there is no need
+    // to update the instruction to block mapping.
+    Instruction* result = new_inst.get();
+    GetContext()->AddAnnotationInst(std::move(new_inst));
+    return result;
+  }
+
   Instruction* AddNaryExtendedInstruction(
       uint32_t result_type, uint32_t set, uint32_t instruction,
       const std::vector<uint32_t>& ext_operands) {
@@ -593,6 +637,23 @@ class InstructionBuilder {
     return AddInstruction(std::move(new_inst));
   }
 
+  Instruction* AddSampledImage(uint32_t sampled_image_type_id,
+                               uint32_t image_id, uint32_t sampler_id) {
+    std::vector<Operand> operands;
+    operands.push_back({SPV_OPERAND_TYPE_ID, {image_id}});
+    operands.push_back({SPV_OPERAND_TYPE_ID, {sampler_id}});
+
+    uint32_t result_id = GetContext()->TakeNextId();
+    if (result_id == 0) {
+      return nullptr;
+    }
+
+    std::unique_ptr<Instruction> new_inst(
+        new Instruction(GetContext(), spv::Op::OpSampledImage,
+                        sampled_image_type_id, result_id, operands));
+    return AddInstruction(std::move(new_inst));
+  }
+
   // Inserts the new instruction before the insertion point.
   Instruction* AddInstruction(std::unique_ptr<Instruction>&& insn) {
     Instruction* insn_ptr = &*insert_before_.InsertBefore(std::move(insn));

source/opt/optimizer.cpp

@@ -637,6 +637,8 @@ bool Optimizer::RegisterPassFromFlag(const std::string& flag,
     }
   } else if (pass_name == "trim-capabilities") {
     RegisterPass(CreateTrimCapabilitiesPass());
+  } else if (pass_name == "split-combined-image-sampler") {
+    RegisterPass(CreateSplitCombinedImageSamplerPass());
   } else {
     Errorf(consumer(), nullptr, {},
            "Unknown flag '--%s'. Use --help for a list of valid flags",
@@ -1188,6 +1190,11 @@ Optimizer::PassToken CreateOpExtInstWithForwardReferenceFixupPass() {
       MakeUnique<opt::OpExtInstWithForwardReferenceFixupPass>());
 }
 
+Optimizer::PassToken CreateSplitCombinedImageSamplerPass() {
+  return MakeUnique<Optimizer::PassToken::Impl>(
+      MakeUnique<opt::SplitCombinedImageSamplerPass>());
+}
+
 }  // namespace spvtools
 
 extern "C" {

source/opt/passes.h

@@ -77,6 +77,7 @@
 #include "source/opt/scalar_replacement_pass.h"
 #include "source/opt/set_spec_constant_default_value_pass.h"
 #include "source/opt/simplification_pass.h"
+#include "source/opt/split_combined_image_sampler_pass.h"
 #include "source/opt/spread_volatile_semantics.h"
 #include "source/opt/ssa_rewrite_pass.h"
 #include "source/opt/strength_reduction_pass.h"