opt: dissolve OpCopyLogical of OpCompositeConstruct in simplification

Author: SteveUrquhart

source/opt/folding_rules.cpp

@@ -2124,6 +2124,82 @@ std::vector<Operand> GetExtractOperandsForElementOfCompositeConstruct(
   return {};
 }
 
+// If the OpCompositeConstruct that feeds an OpCopyLogical can be retyped to
+// the OpCopyLogical's result type, the layout conversion can be expressed at
+// constituent granularity instead of at aggregate granularity. This rewrites
+// the OpCopyLogical as an OpCompositeConstruct of the result type, using the
+// same constituents where their types already match the corresponding
+// field/element of the result type, and inserting per-field OpCopyLogical
+// instructions only for the fields that genuinely require a layout
+// conversion.
+bool CompositeConstructFeedingCopyLogical(
+    IRContext* context, Instruction* inst,
+    const std::vector<const analysis::Constant*>&) {
+  assert(inst->opcode() == spv::Op::OpCopyLogical &&
+         "Wrong opcode.  Should be OpCopyLogical.");
+  analysis::DefUseManager* def_use_mgr = context->get_def_use_mgr();
+
+  uint32_t src_id = inst->GetSingleWordInOperand(0);
+  Instruction* src_inst = def_use_mgr->GetDef(src_id);
+  if (src_inst->opcode() != spv::Op::OpCompositeConstruct) {
+    return false;
+  }
+
+  Instruction* dst_type_inst = def_use_mgr->GetDef(inst->type_id());
+  const uint32_t num_constituents = src_inst->NumInOperands();
+
+  // Determine the expected type id for each constituent of the destination
+  // type.
+  std::vector<uint32_t> expected_type_ids;
+  expected_type_ids.reserve(num_constituents);
+  if (dst_type_inst->opcode() == spv::Op::OpTypeStruct) {
+    if (dst_type_inst->NumInOperands() != num_constituents) {
+      return false;
+    }
+    for (uint32_t i = 0; i < num_constituents; ++i) {
+      expected_type_ids.push_back(dst_type_inst->GetSingleWordInOperand(i));
+    }
+  } else if (dst_type_inst->opcode() == spv::Op::OpTypeArray) {
+    const uint32_t elem_type_id = dst_type_inst->GetSingleWordInOperand(0);
+    for (uint32_t i = 0; i < num_constituents; ++i) {
+      expected_type_ids.push_back(elem_type_id);
+    }
+  } else {
+    return false;
+  }
+
+  // Build the new constituent list, inserting OpCopyLogical instructions for
+  // the fields whose types differ from the result type.
+  InstructionBuilder ir_builder(
+      context, inst,
+      IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
+  std::vector<Operand> operands;
+  operands.reserve(num_constituents);
+  for (uint32_t i = 0; i < num_constituents; ++i) {
+    const uint32_t cid = src_inst->GetSingleWordInOperand(i);
+    Instruction* cdef = def_use_mgr->GetDef(cid);
+    if (cdef->type_id() == expected_type_ids[i]) {
+      operands.push_back({SPV_OPERAND_TYPE_ID, {cid}});
+      continue;
+    }
+    assert(def_use_mgr->GetDef(expected_type_ids[i])->opcode() !=
+               spv::Op::OpTypePointer &&
+           "OpCopyLogical result type must not contain a pointer-typed "
+           "field.");
+    Instruction* per_field_copy = ir_builder.AddUnaryOp(
+        expected_type_ids[i], spv::Op::OpCopyLogical, cid);
+    if (per_field_copy == nullptr) {
+      return false;
+    }
+    operands.push_back({SPV_OPERAND_TYPE_ID, {per_field_copy->result_id()}});
+  }
+
+  inst->SetOpcode(spv::Op::OpCompositeConstruct);
+  inst->SetInOperands(std::move(operands));
+  context->UpdateDefUse(inst);
+  return true;
+}
+
 bool CompositeConstructFeedingExtract(
     IRContext* context, Instruction* inst,
     const std::vector<const analysis::Constant*>&) {
@@ -4476,6 +4552,9 @@ void FoldingRules::AddFoldingRules() {
   rules_[spv::Op::OpCompositeConstruct].push_back(
       CompositeExtractFeedingConstruct);
 
+  rules_[spv::Op::OpCopyLogical].push_back(
+      CompositeConstructFeedingCopyLogical);
+
   rules_[spv::Op::OpCompositeExtract].push_back(InsertFeedingExtract());
   rules_[spv::Op::OpCompositeExtract].push_back(
       CompositeConstructFeedingExtract);

test/opt/simplification_test.cpp

@@ -384,6 +384,93 @@ OpFunctionEnd
 
   SinglePassRunAndCheck<SimplificationPass>(text, text, false);
 }
+
+TEST_F(SimplificationTest, CopyLogicalOfCompositeConstructMatchingTypes) {
+  // OpCopyLogical fed by an OpCompositeConstruct whose constituents are
+  // already typed exactly as the corresponding fields of the OpCopyLogical's
+  // result type can be rewritten as an OpCompositeConstruct of the result
+  // type.
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  const std::string text = R"(
+; CHECK: [[ptr:%\w+]] = OpVariable %_ptr_StorageBuffer_int StorageBuffer
+; CHECK: [[ctor:%\w+]] = OpCompositeConstruct {{%\w+}} [[ptr]]
+; CHECK-NOT: OpCopyLogical
+; CHECK: [[ctor2:%\w+]] = OpCompositeConstruct {{%\w+}} [[ptr]]
+; CHECK: OpStore {{%\w+}} [[ctor2]]
+OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %main "main"
+OpExecutionMode %main LocalSize 1 1 1
+OpDecorate %ResHolder_block Block
+OpMemberDecorate %ResHolder_block 0 Offset 0
+OpDecorate %ssbo DescriptorSet 0
+OpDecorate %ssbo Binding 0
+%void = OpTypeVoid
+%void_fn = OpTypeFunction %void
+%int = OpTypeInt 32 1
+%_ptr_StorageBuffer_int = OpTypePointer StorageBuffer %int
+%ssbo = OpVariable %_ptr_StorageBuffer_int StorageBuffer
+%ResHolder_block = OpTypeStruct %int
+%ResHolder_fn = OpTypeStruct %_ptr_StorageBuffer_int
+%_ptr_Function_ResHolder_fn = OpTypePointer Function %ResHolder_fn
+%main = OpFunction %void None %void_fn
+%entry = OpLabel
+%h = OpVariable %_ptr_Function_ResHolder_fn Function
+%cc = OpCompositeConstruct %ResHolder_block %ssbo
+%cl = OpCopyLogical %ResHolder_fn %cc
+OpStore %h %cl
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<SimplificationPass>(text, false);
+}
+
+TEST_F(SimplificationTest, CopyLogicalOfCompositeConstructMixedTypes) {
+  SetTargetEnv(SPV_ENV_UNIVERSAL_1_4);
+  const std::string text = R"(
+; CHECK: [[ptr:%\w+]] = OpVariable %_ptr_StorageBuffer_int StorageBuffer
+; CHECK: [[u:%\w+]] = OpLoad {{%\w+}}
+; CHECK: [[copy:%\w+]] = OpCopyLogical {{%\w+}} [[u]]
+; CHECK: [[ctor:%\w+]] = OpCompositeConstruct {{%\w+}} [[copy]] [[ptr]]
+; CHECK-NOT: OpCopyLogical {{%\w+}} {{%\w+}} {{%\w+}}
+; CHECK: OpStore {{%\w+}} [[ctor]]
+OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %main "main"
+OpExecutionMode %main LocalSize 1 1 1
+OpDecorate %Uniforms_block Block
+OpMemberDecorate %Uniforms_block 0 Offset 0
+OpDecorate %ubo DescriptorSet 0
+OpDecorate %ubo Binding 0
+OpDecorate %ssbo DescriptorSet 0
+OpDecorate %ssbo Binding 1
+%void = OpTypeVoid
+%void_fn = OpTypeFunction %void
+%int = OpTypeInt 32 1
+%Uniforms_block = OpTypeStruct %int
+%Uniforms_fn = OpTypeStruct %int
+%_ptr_Uniform_Uniforms_block = OpTypePointer Uniform %Uniforms_block
+%ubo = OpVariable %_ptr_Uniform_Uniforms_block Uniform
+%_ptr_StorageBuffer_int = OpTypePointer StorageBuffer %int
+%ssbo = OpVariable %_ptr_StorageBuffer_int StorageBuffer
+%Ctx_block = OpTypeStruct %Uniforms_block %int
+%Ctx_fn = OpTypeStruct %Uniforms_fn %_ptr_StorageBuffer_int
+%_ptr_Function_Ctx_fn = OpTypePointer Function %Ctx_fn
+%main = OpFunction %void None %void_fn
+%entry = OpLabel
+%h = OpVariable %_ptr_Function_Ctx_fn Function
+%u = OpLoad %Uniforms_block %ubo
+%cc = OpCompositeConstruct %Ctx_block %u %ssbo
+%cl = OpCopyLogical %Ctx_fn %cc
+OpStore %h %cl
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndMatch<SimplificationPass>(text, false);
+}
+
 }  // namespace
 }  // namespace opt
 }  // namespace spvtools