opt: split composite from array flattening (KhronosGroup#5733)

* opt: split composite from array flattening

DXC has an option to flatten resource arrays. But when this option
is not used, the resource arrays should be kept as-is.
On the other hand, when a struct contains resources, we MUST flatten is
to be compliant with the Vulkan spec.

Because this pass flattens both types of resources, using a struct of
resources automatically implied flattening arrays.
By adding those 2 new settings, we decide if the pass flattens only one type
of resources, or both.
Note: the flatten_arrays flag only impacts resource arrays.
Arrays of composites containing resources are still flattened.

Since the API is considered stable, I added 2 new functions to create
passes with one flag or the other, and kept the original behavior as-is.

Related to microsoft/DirectXShaderCompiler#6745

Signed-off-by: Nathan Gauër <[email protected]>

* add commandline options

Signed-off-by: Nathan Gauër <[email protected]>

* clang-format

Signed-off-by: Nathan Gauër <[email protected]>

---------

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

include/spirv-tools/optimizer.hpp

@@ -827,14 +827,19 @@ Optimizer::PassToken CreateReplaceDescArrayAccessUsingVarIndexPass();
 
 // Create descriptor scalar replacement pass.
 // This pass replaces every array variable |desc| that has a DescriptorSet and
-// Binding decorations with a new variable for each element of the array.
-// Suppose |desc| was bound at binding |b|.  Then the variable corresponding to
-// |desc[i]| will have binding |b+i|.  The descriptor set will be the same.  It
-// is assumed that no other variable already has a binding that will used by one
-// of the new variables.  If not, the pass will generate invalid Spir-V.  All
-// accesses to |desc| must be OpAccessChain instructions with a literal index
-// for the first index.
+// Binding decorations with a new variable for each element of the
+// array/composite. Suppose |desc| was bound at binding |b|.  Then the variable
+// corresponding to |desc[i]| will have binding |b+i|.  The descriptor set will
+// be the same.  It is assumed that no other variable already has a binding that
+// will used by one of the new variables.  If not, the pass will generate
+// invalid Spir-V.  All accesses to |desc| must be OpAccessChain instructions
+// with a literal index for the first index. This variant flattens both
+// composites and arrays.
 Optimizer::PassToken CreateDescriptorScalarReplacementPass();
+// This variant flattens only composites.
+Optimizer::PassToken CreateDescriptorCompositeScalarReplacementPass();
+// This variant flattens only arrays.
+Optimizer::PassToken CreateDescriptorArrayScalarReplacementPass();
 
 // Create a pass to replace each OpKill instruction with a function call to a
 // function that has a single OpKill.  Also replace each OpTerminateInvocation

source/opt/desc_sroa.cpp

@@ -31,11 +31,14 @@ bool IsDecorationBinding(Instruction* inst) {
 
 Pass::Status DescriptorScalarReplacement::Process() {
   bool modified = false;
-
   std::vector<Instruction*> vars_to_kill;
 
   for (Instruction& var : context()->types_values()) {
-    if (descsroautil::IsDescriptorArray(context(), &var)) {
+    bool is_candidate =
+        flatten_arrays_ && descsroautil::IsDescriptorArray(context(), &var);
+    is_candidate |= flatten_composites_ &&
+                    descsroautil::IsDescriptorStruct(context(), &var);
+    if (is_candidate) {
       modified = true;
       if (!ReplaceCandidate(&var)) {
         return Status::Failure;

source/opt/desc_sroa.h

@@ -32,9 +32,16 @@ namespace opt {
 // Documented in optimizer.hpp
 class DescriptorScalarReplacement : public Pass {
  public:
-  DescriptorScalarReplacement() {}
-
-  const char* name() const override { return "descriptor-scalar-replacement"; }
+  DescriptorScalarReplacement(bool flatten_composites, bool flatten_arrays)
+      : flatten_composites_(flatten_composites),
+        flatten_arrays_(flatten_arrays) {}
+
+  const char* name() const override {
+    if (flatten_composites_ && flatten_arrays_)
+      return "descriptor-scalar-replacement";
+    if (flatten_composites_) return "descriptor-compososite-scalar-replacement";
+    return "descriptor-array-scalar-replacement";
+  }
 
   Status Process() override;
 
@@ -141,6 +148,9 @@ class DescriptorScalarReplacement : public Pass {
   // array |var|. If the entry is |0|, then the variable has not been
   // created yet.
   std::map<Instruction*, std::vector<uint32_t>> replacement_variables_;
+
+  bool flatten_composites_;
+  bool flatten_arrays_;
 };
 
 }  // namespace opt

source/opt/desc_sroa_util.cpp

@@ -29,41 +29,58 @@ uint32_t GetLengthOfArrayType(IRContext* context, Instruction* type) {
   return length_const->GetU32();
 }
 
-}  // namespace
-
-namespace descsroautil {
+bool HasDescriptorDecorations(IRContext* context, Instruction* var) {
+  const auto& decoration_mgr = context->get_decoration_mgr();
+  return decoration_mgr->HasDecoration(
+             var->result_id(), uint32_t(spv::Decoration::DescriptorSet)) &&
+         decoration_mgr->HasDecoration(var->result_id(),
+                                       uint32_t(spv::Decoration::Binding));
+}
 
-bool IsDescriptorArray(IRContext* context, Instruction* var) {
+Instruction* GetVariableType(IRContext* context, Instruction* var) {
   if (var->opcode() != spv::Op::OpVariable) {
-    return false;
+    return nullptr;
   }
 
   uint32_t ptr_type_id = var->type_id();
   Instruction* ptr_type_inst = context->get_def_use_mgr()->GetDef(ptr_type_id);
   if (ptr_type_inst->opcode() != spv::Op::OpTypePointer) {
-    return false;
+    return nullptr;
   }
 
   uint32_t var_type_id = ptr_type_inst->GetSingleWordInOperand(1);
-  Instruction* var_type_inst = context->get_def_use_mgr()->GetDef(var_type_id);
-  if (var_type_inst->opcode() != spv::Op::OpTypeArray &&
-      var_type_inst->opcode() != spv::Op::OpTypeStruct) {
-    return false;
+  return context->get_def_use_mgr()->GetDef(var_type_id);
+}
+
+}  // namespace
+
+namespace descsroautil {
+
+bool IsDescriptorArray(IRContext* context, Instruction* var) {
+  Instruction* var_type_inst = GetVariableType(context, var);
+  if (var_type_inst == nullptr) return false;
+  return var_type_inst->opcode() == spv::Op::OpTypeArray &&
+         HasDescriptorDecorations(context, var);
+}
+
+bool IsDescriptorStruct(IRContext* context, Instruction* var) {
+  Instruction* var_type_inst = GetVariableType(context, var);
+  if (var_type_inst == nullptr) return false;
+
+  while (var_type_inst->opcode() == spv::Op::OpTypeArray) {
+    var_type_inst = context->get_def_use_mgr()->GetDef(
+        var_type_inst->GetInOperand(0).AsId());
   }
 
+  if (var_type_inst->opcode() != spv::Op::OpTypeStruct) return false;
+
   // All structures with descriptor assignments must be replaced by variables,
   // one for each of their members - with the exceptions of buffers.
   if (IsTypeOfStructuredBuffer(context, var_type_inst)) {
     return false;
   }
 
-  if (!context->get_decoration_mgr()->HasDecoration(
-          var->result_id(), uint32_t(spv::Decoration::DescriptorSet))) {
-    return false;
-  }
-
-  return context->get_decoration_mgr()->HasDecoration(
-      var->result_id(), uint32_t(spv::Decoration::Binding));
+  return HasDescriptorDecorations(context, var);
 }
 
 bool IsTypeOfStructuredBuffer(IRContext* context, const Instruction* type) {

source/opt/desc_sroa_util.h

@@ -27,6 +27,10 @@ namespace descsroautil {
 // descriptor array.
 bool IsDescriptorArray(IRContext* context, Instruction* var);
 
+// Returns true if |var| is an OpVariable instruction that represents a
+// struct containing descriptors.
+bool IsDescriptorStruct(IRContext* context, Instruction* var);
+
 // Returns true if |type| is a type that could be used for a structured buffer
 // as opposed to a type that would be used for a structure of resource
 // descriptors.

source/opt/optimizer.cpp

@@ -364,6 +364,10 @@ bool Optimizer::RegisterPassFromFlag(const std::string& flag,
     RegisterPass(CreateSpreadVolatileSemanticsPass());
   } else if (pass_name == "descriptor-scalar-replacement") {
     RegisterPass(CreateDescriptorScalarReplacementPass());
+  } else if (pass_name == "descriptor-composite-scalar-replacement") {
+    RegisterPass(CreateDescriptorCompositeScalarReplacementPass());
+  } else if (pass_name == "descriptor-array-scalar-replacement") {
+    RegisterPass(CreateDescriptorArrayScalarReplacementPass());
   } else if (pass_name == "eliminate-dead-code-aggressive") {
     RegisterPass(CreateAggressiveDCEPass(preserve_interface));
   } else if (pass_name == "eliminate-insert-extract") {
@@ -1059,7 +1063,20 @@ Optimizer::PassToken CreateSpreadVolatileSemanticsPass() {
 
 Optimizer::PassToken CreateDescriptorScalarReplacementPass() {
   return MakeUnique<Optimizer::PassToken::Impl>(
-      MakeUnique<opt::DescriptorScalarReplacement>());
+      MakeUnique<opt::DescriptorScalarReplacement>(
+          /* flatten_composites= */ true, /* flatten_arrays= */ true));
+}
+
+Optimizer::PassToken CreateDescriptorCompositeScalarReplacementPass() {
+  return MakeUnique<Optimizer::PassToken::Impl>(
+      MakeUnique<opt::DescriptorScalarReplacement>(
+          /* flatten_composites= */ true, /* flatten_arrays= */ false));
+}
+
+Optimizer::PassToken CreateDescriptorArrayScalarReplacementPass() {
+  return MakeUnique<Optimizer::PassToken::Impl>(
+      MakeUnique<opt::DescriptorScalarReplacement>(
+          /* flatten_composites= */ false, /* flatten_arrays= */ true));
 }
 
 Optimizer::PassToken CreateWrapOpKillPass() {