opt: Fold OpCompositeExtract feeding from OpCopyLogical or OpLoad (#6614)

- CopyLogicalFeedingExtract: Hoist OpCompositeExtract before
OpCopyLogical.
  If the input to an OpCompositeExtract is an OpCopyLogical, we can
  extract from the original composite and then copy the result.
- LoadFeedingExtract: Change OpLoad + OpCompositeExtract to
OpAccessChain + OpLoad.
  If the input to an OpCompositeExtract is an OpLoad, we can load the
  specific element instead of the entire composite. This is restricted
  to non-Function/Private storage classes to avoid interfering with
  local-access-chain-convert.

Updated fold_test.cpp with 8 new test cases and updated existing tests
to use SPV_ENV_UNIVERSAL_1_5.

Fixes #6611

Author: s-perron

source/opt/folding_rules.cpp

@@ -2193,6 +2193,174 @@ uint32_t GetElementType(uint32_t type_id, Instruction::iterator start,
   return type_id;
 }
 
+// If the input to an OpCompositeExtract is an OpCopyLogical, then we can
+// hoist the extraction before the copy.
+bool CopyLogicalFeedingExtract(IRContext* context, Instruction* inst,
+                               const std::vector<const analysis::Constant*>&) {
+  assert(inst->opcode() == spv::Op::OpCompositeExtract &&
+         "Wrong opcode.  Should be OpCompositeExtract.");
+
+  analysis::DefUseManager* def_use_mgr = context->get_def_use_mgr();
+  uint32_t cid = inst->GetSingleWordInOperand(kExtractCompositeIdInIdx);
+  Instruction* cinst = def_use_mgr->GetDef(cid);
+
+  if (cinst->opcode() != spv::Op::OpCopyLogical) {
+    return false;
+  }
+
+  uint32_t original_composite_id = cinst->GetSingleWordInOperand(0);
+  Instruction* original_composite_inst =
+      def_use_mgr->GetDef(original_composite_id);
+
+  std::vector<uint32_t> indices;
+  for (uint32_t i = 1; i < inst->NumInOperands(); ++i) {
+    indices.push_back(inst->GetSingleWordInOperand(i));
+  }
+
+  uint32_t original_element_type_id =
+      GetElementType(original_composite_inst->type_id(), inst->begin() + 3,
+                     inst->end(), def_use_mgr);
+  assert(original_element_type_id != 0 &&
+         "Could not find the element type.  Invalid SPIR-V.");
+
+  InstructionBuilder ir_builder(
+      context, inst,
+      IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
+
+  Instruction* new_extract = ir_builder.AddCompositeExtract(
+      original_element_type_id, original_composite_id, indices);
+
+  if (original_element_type_id == inst->type_id())
+    inst->SetOpcode(spv::Op::OpCopyObject);
+  else
+    inst->SetOpcode(spv::Op::OpCopyLogical);
+  inst->SetInOperands({{SPV_OPERAND_TYPE_ID, {new_extract->result_id()}}});
+  return true;
+}
+
+// If the input to an OpCompositeExtract is an OpLoad, we can change the
+// load into a load of an OpAccessChain.
+bool LoadFeedingExtract(IRContext* context, Instruction* inst,
+                        const std::vector<const analysis::Constant*>&) {
+  assert(inst->opcode() == spv::Op::OpCompositeExtract &&
+         "Wrong opcode.  Should be OpCompositeExtract.");
+
+  analysis::DefUseManager* def_use_mgr = context->get_def_use_mgr();
+  uint32_t cid = inst->GetSingleWordInOperand(kExtractCompositeIdInIdx);
+  Instruction* cinst = def_use_mgr->GetDef(cid);
+
+  if (cinst->opcode() != spv::Op::OpLoad) {
+    return false;
+  }
+
+  Instruction* composite_type_inst = def_use_mgr->GetDef(cinst->type_id());
+  if (composite_type_inst->opcode() != spv::Op::OpTypeStruct &&
+      composite_type_inst->opcode() != spv::Op::OpTypeArray) {
+    return false;
+  }
+
+  // Check the memory operands.
+  if (cinst->NumInOperands() > 1) {
+    uint32_t memory_access_mask = cinst->GetSingleWordInOperand(1);
+    if (memory_access_mask & uint32_t(spv::MemoryAccessMask::Volatile)) {
+      return false;
+    }
+  }
+
+  uint32_t ptr_id = cinst->GetSingleWordInOperand(0);
+  Instruction* ptr_inst = def_use_mgr->GetDef(ptr_id);
+  Instruction* ptr_type_inst = def_use_mgr->GetDef(ptr_inst->type_id());
+  assert(ptr_type_inst->opcode() == spv::Op::OpTypePointer);
+  spv::StorageClass storage_class =
+      static_cast<spv::StorageClass>(ptr_type_inst->GetSingleWordInOperand(0));
+
+  // If the storage class is Function or Private, we do not want to fold.
+  // These are the storage classes that the local-access-chain-convert pass
+  // works on.
+  if (storage_class == spv::StorageClass::Function ||
+      storage_class == spv::StorageClass::Private) {
+    return false;
+  }
+
+  analysis::ConstantManager* const_mgr = context->get_constant_mgr();
+  analysis::TypeManager* type_mgr = context->get_type_mgr();
+  std::vector<uint32_t> index_ids;
+  for (uint32_t i = 1; i < inst->NumInOperands(); ++i) {
+    uint32_t index = inst->GetSingleWordInOperand(i);
+    const analysis::Constant* index_const =
+        const_mgr->GetConstant(type_mgr->GetUIntType(), {index});
+    index_ids.push_back(
+        const_mgr->GetDefiningInstruction(index_const)->result_id());
+  }
+
+  InstructionBuilder ir_builder(
+      context, cinst,
+      IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
+
+  uint32_t element_ptr_type_id =
+      type_mgr->FindPointerToType(inst->type_id(), storage_class);
+  if (element_ptr_type_id == 0) {
+    return false;
+  }
+
+  Instruction* access_chain =
+      ir_builder.AddAccessChain(element_ptr_type_id, ptr_id, index_ids);
+  std::vector<Operand> load_operands;
+  load_operands.push_back({SPV_OPERAND_TYPE_ID, {access_chain->result_id()}});
+
+  if (cinst->NumInOperands() > 1) {
+    uint32_t memory_access_mask = cinst->GetSingleWordInOperand(1);
+    load_operands.push_back(
+        {SPV_OPERAND_TYPE_MEMORY_ACCESS, {memory_access_mask}});
+
+    uint32_t current_operand_index = 2;
+    if (memory_access_mask & uint32_t(spv::MemoryAccessMask::Aligned)) {
+      uint32_t original_alignment =
+          cinst->GetSingleWordInOperand(current_operand_index);
+
+      std::vector<uint32_t> extract_indices;
+      for (uint32_t i = 1; i < inst->NumInOperands(); ++i) {
+        extract_indices.push_back(inst->GetSingleWordInOperand(i));
+      }
+
+      std::optional<uint32_t> offset =
+          type_mgr->GetType(cinst->type_id())->GetByteOffset(extract_indices);
+      if (!offset) {
+        return false;
+      }
+
+      uint32_t new_alignment = original_alignment;
+      if (*offset != 0) {
+        uint32_t offset_alignment = *offset & ~(*offset - 1);
+        new_alignment = std::min(original_alignment, offset_alignment);
+      }
+
+      load_operands.push_back(
+          {SPV_OPERAND_TYPE_TYPED_LITERAL_NUMBER, {new_alignment}});
+      current_operand_index++;
+    }
+
+    // Copy the remaining operands
+    for (; current_operand_index < cinst->NumInOperands();
+         ++current_operand_index) {
+      load_operands.push_back(cinst->GetInOperand(current_operand_index));
+    }
+  }
+
+  uint32_t load_result_id = context->TakeNextId();
+  if (load_result_id == 0) return false;
+
+  std::unique_ptr<Instruction> new_load_inst(
+      new Instruction(context, spv::Op::OpLoad, inst->type_id(), load_result_id,
+                      load_operands));
+  Instruction* new_load = ir_builder.AddInstruction(std::move(new_load_inst));
+
+  inst->SetOpcode(spv::Op::OpCopyObject);
+  inst->SetInOperands({{SPV_OPERAND_TYPE_ID, {new_load->result_id()}}});
+
+  return true;
+}
+
 // Returns true of |inst_1| and |inst_2| have the same indexes that will be used
 // to index into a composite object, excluding the last index.  The two
 // instructions must have the same opcode, and be either OpCompositeExtract or
@@ -4309,6 +4477,8 @@ void FoldingRules::AddFoldingRules() {
       CompositeConstructFeedingExtract);
   rules_[spv::Op::OpCompositeExtract].push_back(VectorShuffleFeedingExtract());
   rules_[spv::Op::OpCompositeExtract].push_back(FMixFeedingExtract());
+  rules_[spv::Op::OpCompositeExtract].push_back(CopyLogicalFeedingExtract);
+  rules_[spv::Op::OpCompositeExtract].push_back(LoadFeedingExtract);
 
   rules_[spv::Op::OpCompositeInsert].push_back(
       CompositeInsertToCompositeConstruct);

source/opt/types.cpp

@@ -303,6 +303,78 @@ uint64_t Type::NumberOfComponents() const {
   }
 }
 
+std::optional<uint32_t> Type::GetByteOffset(
+    const std::vector<uint32_t>& access_chain) const {
+  uint32_t offset = 0;
+  const Type* current_type = this;
+  for (uint32_t index : access_chain) {
+    if (const Struct* struct_type = current_type->AsStruct()) {
+      std::optional<uint32_t> member_offset;
+      for (const auto& deco : struct_type->element_decorations()) {
+        if (deco.first != index) continue;
+        for (const auto& inst : deco.second) {
+          if (inst[0] == uint32_t(spv::Decoration::Offset)) {
+            member_offset = inst[1];
+            break;
+          }
+        }
+      }
+      if (!member_offset) return {};
+      offset += *member_offset;
+      current_type = struct_type->element_types()[index];
+    } else if (const Array* array_type = current_type->AsArray()) {
+      std::optional<uint32_t> array_stride;
+      for (const auto& deco : array_type->decorations()) {
+        if (deco[0] == uint32_t(spv::Decoration::ArrayStride)) {
+          array_stride = deco[1];
+          break;
+        }
+      }
+      if (!array_stride) return {};
+      offset += *array_stride * index;
+      current_type = array_type->element_type();
+    } else if (const RuntimeArray* runtime_array_type =
+                   current_type->AsRuntimeArray()) {
+      std::optional<uint32_t> array_stride;
+      for (const auto& deco : runtime_array_type->decorations()) {
+        if (deco[0] == uint32_t(spv::Decoration::ArrayStride)) {
+          array_stride = deco[1];
+          break;
+        }
+      }
+      if (!array_stride) return {};
+      offset += *array_stride * index;
+      current_type = runtime_array_type->element_type();
+    } else if (const Matrix* matrix_type = current_type->AsMatrix()) {
+      std::optional<uint32_t> matrix_stride;
+      for (const auto& deco : matrix_type->decorations()) {
+        if (deco[0] == uint32_t(spv::Decoration::MatrixStride)) {
+          matrix_stride = deco[1];
+          break;
+        }
+      }
+      if (!matrix_stride) return {};
+      offset += *matrix_stride * index;
+      current_type = matrix_type->element_type();
+    } else if (const Vector* vector_type = current_type->AsVector()) {
+      const Type* component_type = vector_type->element_type();
+      uint32_t component_size = 0;
+      if (component_type->AsInteger()) {
+        component_size = component_type->AsInteger()->width() / 8;
+      } else if (component_type->AsFloat()) {
+        component_size = component_type->AsFloat()->width() / 8;
+      } else {
+        return {};
+      }
+      offset += component_size * index;
+      current_type = component_type;
+    } else {
+      return {};
+    }
+  }
+  return offset;
+}
+
 bool Integer::IsSameImpl(const Type* that, IsSameCache*) const {
   const Integer* it = that->AsInteger();
   return it && width_ == it->width_ && signed_ == it->signed_ &&

source/opt/types.h

@@ -21,6 +21,7 @@
 
 #include <map>
 #include <memory>
+#include <optional>
 #include <set>
 #include <string>
 #include <unordered_map>
@@ -196,6 +197,13 @@ class Type {
   // non-composite type.
   uint64_t NumberOfComponents() const;
 
+  // Returns the byte offset of the member of this type that is identified
+  // by |access_chain|.  The vector |access_chain| is a series of integers that
+  // are used to pick members as in the |OpCompositeExtract| instructions.
+  // Returns {} if the offset cannot be computed.
+  std::optional<uint32_t> GetByteOffset(
+      const std::vector<uint32_t>& access_chain) const;
+
   // A bunch of methods for casting this type to a given type. Returns this if
   // the cast can be done, nullptr otherwise.
   // clang-format off