Using the instruction folder to fold OpSpecConstantOp (#2598)

In order to try to reduce code duplication and to be able
to fold more cases, we want to use the instruction folder
when folding an OpSpecConstantOp with constant operands.

A couple other changes are need to make this work.  First
GetDefiningInstruction| in the constant manager is able
to handle |type_id| being logically equivalent to another
type, so we updated the interface, and removed the assert.

Some tests were also updated because we not generate
better code because constants are not duplicated as much
as before.

No need for new tests.  The functionality of the instruction folder is
already tested.  There are tests check that the instruction folder is
being used correctly for OpCompositeExtract and OpVectorShuffle in the
existing test cases.

Fixes #2585.

source/opt/constants.cpp

@@ -185,8 +185,6 @@ Instruction* ConstantManager::BuildInstructionAndAddToModule(
 
 Instruction* ConstantManager::GetDefiningInstruction(
     const Constant* c, uint32_t type_id, Module::inst_iterator* pos) {
-  assert(type_id == 0 ||
-         context()->get_type_mgr()->GetType(type_id) == c->type());
   uint32_t decl_id = FindDeclaredConstant(c, type_id);
   if (decl_id == 0) {
     auto iter = context()->types_values_end();

source/opt/constants.h

@@ -524,12 +524,7 @@ class ConstantManager {
   // instruction at the end of the current module's types section.
   //
   // |type_id| is an optional argument for disambiguating equivalent types. If
-  // |type_id| is specified, it is used as the type of the constant when a new
-  // instruction is created. Otherwise the type of the constant is derived by
-  // getting an id from the type manager for |c|.
-  //
-  // When |type_id| is not zero, the type of |c| must be the type returned by
-  // type manager when given |type_id|.
+  // |type_id| is specified, the contant returned will have that type id.
   Instruction* GetDefiningInstruction(const Constant* c, uint32_t type_id = 0,
                                       Module::inst_iterator* pos = nullptr);
 

source/opt/fold_spec_constant_op_and_composite_pass.cpp

@@ -120,19 +120,14 @@ bool FoldSpecConstantOpAndCompositePass::ProcessOpSpecConstantOp(
 
   switch (static_cast<SpvOp>(inst->GetSingleWordInOperand(0))) {
     case SpvOp::SpvOpCompositeExtract:
-      folded_inst = DoCompositeExtract(pos);
-      break;
     case SpvOp::SpvOpVectorShuffle:
-      folded_inst = DoVectorShuffle(pos);
-      break;
-
     case SpvOp::SpvOpCompositeInsert:
-      // Current Glslang does not generate code with OpSpecConstantOp
-      // CompositeInsert instruction, so this is not implmented so far.
-      // TODO(qining): Implement CompositeInsert case.
-      return false;
-
+      folded_inst = FoldWithInstructionFolder(pos);
+      break;
     default:
+      // TODO: This should use the instruction folder as well, but some folding
+      // rules are missing.
+
       // Component-wise operations.
       folded_inst = DoComponentWiseOperation(pos);
       break;
@@ -157,54 +152,65 @@ uint32_t FoldSpecConstantOpAndCompositePass::GetTypeComponent(
   return subtype;
 }
 
-Instruction* FoldSpecConstantOpAndCompositePass::DoCompositeExtract(
-    Module::inst_iterator* pos) {
-  Instruction* inst = &**pos;
-  assert(inst->NumInOperands() - 1 >= 2 &&
-         "OpSpecConstantOp CompositeExtract requires at least two non-type "
-         "non-opcode operands.");
-  assert(inst->GetInOperand(1).type == SPV_OPERAND_TYPE_ID &&
-         "The composite operand must have a SPV_OPERAND_TYPE_ID type");
-  assert(
-      inst->GetInOperand(2).type == SPV_OPERAND_TYPE_LITERAL_INTEGER &&
-      "The literal operand must have a SPV_OPERAND_TYPE_LITERAL_INTEGER type");
-
-  // Note that for OpSpecConstantOp, the second in-operand is the first id
-  // operand. The first in-operand is the spec opcode.
-  uint32_t source = inst->GetSingleWordInOperand(1);
-  uint32_t type = context()->get_def_use_mgr()->GetDef(source)->type_id();
-  const analysis::Constant* first_operand_const =
-      context()->get_constant_mgr()->FindDeclaredConstant(source);
-  if (!first_operand_const) return nullptr;
-
-  const analysis::Constant* current_const = first_operand_const;
-  for (uint32_t i = 2; i < inst->NumInOperands(); i++) {
-    uint32_t literal = inst->GetSingleWordInOperand(i);
-    type = GetTypeComponent(type, literal);
-  }
-  for (uint32_t i = 2; i < inst->NumInOperands(); i++) {
-    uint32_t literal = inst->GetSingleWordInOperand(i);
-    if (const analysis::CompositeConstant* composite_const =
-            current_const->AsCompositeConstant()) {
-      // Case 1: current constant is a non-null composite type constant.
-      assert(literal < composite_const->GetComponents().size() &&
-             "Literal index out of bound of the composite constant");
-      current_const = composite_const->GetComponents().at(literal);
-    } else if (current_const->AsNullConstant()) {
-      // Case 2: current constant is a constant created with OpConstantNull.
-      // Because components of a NullConstant are always NullConstants, we can
-      // return early with a NullConstant in the result type.
-      return context()->get_constant_mgr()->BuildInstructionAndAddToModule(
-          context()->get_constant_mgr()->GetConstant(
-              context()->get_constant_mgr()->GetType(inst), {}),
-          pos, type);
-    } else {
-      // Dereferencing a non-composite constant. Invalid case.
+Instruction* FoldSpecConstantOpAndCompositePass::FoldWithInstructionFolder(
+    Module::inst_iterator* inst_iter_ptr) {
+  // If one of operands to the instruction is not a
+  // constant, then we cannot fold this spec constant.
+  for (uint32_t i = 1; i < (*inst_iter_ptr)->NumInOperands(); i++) {
+    const Operand& operand = (*inst_iter_ptr)->GetInOperand(i);
+    if (operand.type != SPV_OPERAND_TYPE_ID &&
+        operand.type != SPV_OPERAND_TYPE_OPTIONAL_ID) {
+      continue;
+    }
+    uint32_t id = operand.words[0];
+    if (context()->get_constant_mgr()->FindDeclaredConstant(id) == nullptr) {
       return nullptr;
     }
   }
-  return context()->get_constant_mgr()->BuildInstructionAndAddToModule(
-      current_const, pos);
+
+  // All of the operands are constant.  Construct a regular version of the
+  // instruction and pass it to the instruction folder.
+  std::unique_ptr<Instruction> inst((*inst_iter_ptr)->Clone(context()));
+  inst->SetOpcode(
+      static_cast<SpvOp>((*inst_iter_ptr)->GetSingleWordInOperand(0)));
+  inst->RemoveOperand(2);
+
+  // We want the current instruction to be replaced by an |OpConstant*|
+  // instruction in the same position. We need to keep track of which constants
+  // the instruction folder creates, so we can move them into the correct place.
+  auto last_type_value_iter = (context()->types_values_end());
+  --last_type_value_iter;
+  Instruction* last_type_value = &*last_type_value_iter;
+
+  auto identity_map = [](uint32_t id) { return id; };
+  Instruction* new_const_inst =
+      context()->get_instruction_folder().FoldInstructionToConstant(
+          inst.get(), identity_map);
+  assert(new_const_inst != nullptr &&
+         "Failed to fold instruction that must be folded.");
+
+  // Get the instruction before |pos| to insert after.  |pos| cannot be the
+  // first instruction in the list because its type has to come first.
+  Instruction* insert_pos = (*inst_iter_ptr)->PreviousNode();
+  assert(insert_pos != nullptr &&
+         "pos is the first instruction in the types and values.");
+  bool need_to_clone = true;
+  for (Instruction* i = last_type_value->NextNode(); i != nullptr;
+       i = last_type_value->NextNode()) {
+    if (i == new_const_inst) {
+      need_to_clone = false;
+    }
+    i->InsertAfter(insert_pos);
+    insert_pos = insert_pos->NextNode();
+  }
+
+  if (need_to_clone) {
+    new_const_inst = new_const_inst->Clone(context());
+    new_const_inst->SetResultId(TakeNextId());
+    new_const_inst->InsertAfter(insert_pos);
+    get_def_use_mgr()->AnalyzeInstDefUse(new_const_inst);
+  }
+  return new_const_inst;
 }
 
 Instruction* FoldSpecConstantOpAndCompositePass::DoVectorShuffle(

source/opt/fold_spec_constant_op_and_composite_pass.h

@@ -54,11 +54,9 @@ class FoldSpecConstantOpAndCompositePass : public Pass {
   // it.
   bool ProcessOpSpecConstantOp(Module::inst_iterator* pos);
 
-  // Try to fold the OpSpecConstantOp CompositeExtract instruction pointed by
-  // the given instruction iterator to a normal constant defining instruction.
-  // Returns the pointer to the new constant defining instruction if succeeded.
-  // Otherwise returns nullptr.
-  Instruction* DoCompositeExtract(Module::inst_iterator* inst_iter_ptr);
+  // Returns the result of folding the OpSpecConstantOp instruction
+  // |inst_iter_ptr| using the instruction folder.
+  Instruction* FoldWithInstructionFolder(Module::inst_iterator* inst_iter_ptr);
 
   // Try to fold the OpSpecConstantOp VectorShuffle instruction pointed by the
   // given instruction iterator to a normal constant defining instruction.

test/opt/fold_spec_const_op_composite_test.cpp

@@ -1135,14 +1135,14 @@ INSTANTIATE_TEST_SUITE_P(
                 "%outer = OpConstantComposite %outer_struct %inner %signed_one",
                 "%extract_inner = OpSpecConstantOp %inner_struct CompositeExtract %outer 0",
                 "%extract_int = OpSpecConstantOp %int CompositeExtract %outer 1",
-                "%extract_inner_float = OpSpecConstantOp %int CompositeExtract %outer 0 2",
+                "%extract_inner_float = OpSpecConstantOp %float CompositeExtract %outer 0 2",
               },
               // expected
               {
                 "%float_1 = OpConstant %float 1",
                 "%inner = OpConstantComposite %inner_struct %bool_true %signed_null %float_1",
                 "%outer = OpConstantComposite %outer_struct %inner %signed_one",
-                "%extract_inner = OpConstantComposite %flat_struct %bool_true %signed_null %float_1",
+                "%extract_inner = OpConstantComposite %inner_struct %bool_true %signed_null %float_1",
                 "%extract_int = OpConstant %int 1",
                 "%extract_inner_float = OpConstant %float 1",
               },
@@ -1256,13 +1256,9 @@ INSTANTIATE_TEST_SUITE_P(
               },
               // expected
               {
-                "%60 = OpConstantNull %int",
                 "%a = OpConstantComposite %v2int %signed_null %signed_null",
-                "%62 = OpConstantNull %int",
                 "%b = OpConstantComposite %v2int %signed_zero %signed_one",
-                "%64 = OpConstantNull %int",
                 "%c = OpConstantComposite %v2int %signed_three %signed_null",
-                "%66 = OpConstantNull %int",
                 "%d = OpConstantComposite %v2int %signed_null %signed_null",
               }
             },