Fix decorations of inlined functions.

Fixes issue #728.  Currently the inliner is not generating decorations for
inlined code which corresponds to function code which has decorations. An
example of decorations that are relevant: RelaxedPrecision, NoContraction.

The solution is to replicate the decoration during inlining.

source/opt/decoration_manager.cpp

@@ -15,6 +15,7 @@
 #include "decoration_manager.h"
 
 #include <stack>
+#include <iostream>
 
 namespace spvtools {
 namespace opt {
@@ -127,11 +128,11 @@ bool DecorationManager::AreDecorationsTheSame(
   return true;
 }
 
-void DecorationManager::AnalyzeDecorations(ir::Module* module) {
-  if (!module) return;
+void DecorationManager::AnalyzeDecorations() {
+  if (!module_) return;
 
   // Collect all group ids.
-  for (const ir::Instruction& inst : module->annotations()) {
+  for (const ir::Instruction& inst : module_->annotations()) {
     switch (inst.opcode()) {
       case SpvOpDecorationGroup:
         group_to_decoration_insts_.insert({inst.result_id(), {}});
@@ -142,7 +143,7 @@ void DecorationManager::AnalyzeDecorations(ir::Module* module) {
   }
 
   // For each group and instruction, collect all their decoration instructions.
-  for (ir::Instruction& inst : module->annotations()) {
+  for (ir::Instruction& inst : module_->annotations()) {
     switch (inst.opcode()) {
       case SpvOpDecorate:
       case SpvOpDecorateId:
@@ -230,23 +231,17 @@ std::vector<T> DecorationManager::InternalGetDecorationsFor(
   return decorations;
 }
 
-void DecorationManager::ForEachDecoration(
-    uint32_t id, uint32_t decoration,
-    std::function<void(const ir::Instruction&)> f) const {
-  auto decoration_list = id_to_decoration_insts_.find(id);
-  if (decoration_list != id_to_decoration_insts_.end()) {
-    for (const ir::Instruction* inst : decoration_list->second) {
+void DecorationManager::ForEachDecoration(uint32_t id,
+                                          uint32_t decoration,
+                                          std::function<void(const ir::Instruction&)> f) {
+  for (const ir::Instruction* inst : GetDecorationsFor(id, true)) {
     switch (inst->opcode()) {
-        case SpvOpDecorate:
-          if (inst->GetSingleWordInOperand(1) == decoration) {
-            f(*inst);
-          }
-          break;
       case SpvOpMemberDecorate:
         if (inst->GetSingleWordInOperand(2) == decoration) {
           f(*inst);
         }
         break;
+      case SpvOpDecorate:
       case SpvOpDecorateId:
         if (inst->GetSingleWordInOperand(1) == decoration) {
           f(*inst);
@@ -257,6 +252,52 @@ void DecorationManager::ForEachDecoration(
     }
   }
 }
+
+void DecorationManager::CloneDecorations(uint32_t from, uint32_t to,
+                                         std::function<void(ir::Instruction&, bool)> f) {
+  assert(f && "Missing function parameter f");
+  auto const decoration_list = id_to_decoration_insts_.find(from);
+  if (decoration_list == id_to_decoration_insts_.end()) return;
+  for (ir::Instruction* inst : decoration_list->second) {
+    switch (inst->opcode()) {
+      case SpvOpGroupDecorate:
+        f(*inst, false);
+        // add |to| to list of decorated id's
+        inst->AddOperand(ir::Operand(spv_operand_type_t::SPV_OPERAND_TYPE_ID, {to}));
+        id_to_decoration_insts_[to].push_back(inst);
+        f(*inst, true);
+        break;
+      case SpvOpGroupMemberDecorate: {
+        f(*inst, false);
+        // for each (id == from), add (to, literal) as operands
+        const uint32_t num_operands = inst->NumOperands();
+        for (uint32_t i = 1; i < num_operands; i += 2) {
+          ir::Operand op = inst->GetOperand(i);
+          if (op.words[0] == from) { // add new pair of operands: (to, literal)
+            inst->AddOperand(ir::Operand(spv_operand_type_t::SPV_OPERAND_TYPE_ID, {to}));
+            op = inst->GetOperand(i + 1);
+            inst->AddOperand(std::move(op));
+          }
+        }
+        id_to_decoration_insts_[to].push_back(inst);
+        f(*inst, true);
+        break;
+      }
+      case SpvOpDecorate:
+      case SpvOpMemberDecorate:
+      case SpvOpDecorateId: {
+        // simply clone decoration and change |target-id| to |to|
+        std::unique_ptr<ir::Instruction> new_inst(inst->Clone());
+        new_inst->SetInOperand(0, {to});
+        id_to_decoration_insts_[to].push_back(new_inst.get());
+        f(*new_inst, true);
+        module_->AddAnnotationInst(std::move(new_inst));
+        break;
+      }
+      default:
+        assert(false && "Unexpected decoration instruction");
+    }
+  }
 }
 
 }  // namespace analysis

source/opt/decoration_manager.h

@@ -30,7 +30,11 @@ namespace analysis {
 class DecorationManager {
  public:
   // Constructs a decoration manager from the given |module|
-  DecorationManager(ir::Module* module) { AnalyzeDecorations(module); }
+  explicit DecorationManager(ir::Module* module) : module_(module) {
+    AnalyzeDecorations();
+  }
+  DecorationManager() = delete;
+
   // Removes all decorations from |id|, which should not be a group ID, except
   // for linkage decorations if |keep_linkage| is set.
   void RemoveDecorationsFrom(uint32_t id, bool keep_linkage);
@@ -52,16 +56,26 @@ class DecorationManager {
                              const ir::Instruction* inst2) const;
 
   // |f| is run on each decoration instruction for |id| with decoration
-  // |decoration|.
+  // |decoration|. Processed are all decorations which target |id| either
+  // directly or indirectly by Decoration Groups.
   void ForEachDecoration(uint32_t id, uint32_t decoration,
-                         std::function<void(const ir::Instruction& f)>) const;
+                         std::function<void(const ir::Instruction&)> f);
+
+  // Clone all decorations from one id |from|.
+  // The cloned decorations are assigned to the given id |to| and are
+  // added to the module. The purpose is to decorate cloned instructions.
+  // This function does not check if the id |to| is already decorated.
+  // Function |f| can be used to update context information and is called
+  // with |false|, before an instruction is going to be changed and
+  // with |true| afterwards.
+  void CloneDecorations(uint32_t from, uint32_t to, std::function<void(ir::Instruction&, bool)> f);
 
  private:
   using IdToDecorationInstsMap =
       std::unordered_map<uint32_t, std::vector<ir::Instruction*>>;
   // Analyzes the defs and uses in the given |module| and populates data
   // structures in this class. Does nothing if |module| is nullptr.
-  void AnalyzeDecorations(ir::Module* module);
+  void AnalyzeDecorations();
 
   template <typename T>
   std::vector<T> InternalGetDecorationsFor(uint32_t id, bool include_linkage);
@@ -74,6 +88,8 @@ class DecorationManager {
   IdToDecorationInstsMap id_to_decoration_insts_;
   // Mapping from group ids to all the decoration instructions they apply.
   IdToDecorationInstsMap group_to_decoration_insts_;
+  // The enclosing module.
+  ir::Module* module_;
 };
 
 }  // namespace analysis

source/opt/inline_pass.cpp

@@ -145,6 +145,7 @@ void InlinePass::CloneAndMapLocals(
   while (callee_var_itr->opcode() == SpvOp::SpvOpVariable) {
     std::unique_ptr<ir::Instruction> var_inst(callee_var_itr->Clone());
     uint32_t newId = TakeNextId();
+    dec_mgr_->CloneDecorations(callee_var_itr->result_id(), newId, update_def_use_mgr_);
     var_inst->SetResultId(newId);
     (*callee2caller)[callee_var_itr->result_id()] = newId;
     new_vars->push_back(std::move(var_inst));
@@ -173,6 +174,7 @@ uint32_t InlinePass::CreateReturnVar(
           {SpvStorageClassFunction}}}));
     new_vars->push_back(std::move(var_inst));
   }
+  dec_mgr_->CloneDecorations(calleeFn->result_id(), returnVarId, update_def_use_mgr_);
   return returnVarId;
 }
 
@@ -197,6 +199,7 @@ void InlinePass::CloneSameBlockOps(
             CloneSameBlockOps(&sb_inst, postCallSB, preCallSB, block_ptr);
             const uint32_t rid = sb_inst->result_id();
             const uint32_t nid = this->TakeNextId();
+            dec_mgr_->CloneDecorations(rid, nid, update_def_use_mgr_);
             sb_inst->SetResultId(nid);
             (*postCallSB)[rid] = nid;
             *iid = nid;
@@ -476,6 +479,7 @@ void InlinePass::GenInlineCode(
             callee2caller[rid] = nid;
           }
           cp_inst->SetResultId(nid);
+          dec_mgr_->CloneDecorations(rid, nid, update_def_use_mgr_);
         }
         new_blk_ptr->AddInstruction(std::move(cp_inst));
       } break;
@@ -642,6 +646,12 @@ bool InlinePass::IsInlinableFunction(ir::Function* func) {
 void InlinePass::InitializeInline(ir::IRContext* c) {
   InitializeProcessing(c);
 
+  dec_mgr_.reset(new analysis::DecorationManager(c->module()));
+  update_def_use_mgr_ = [this] (ir::Instruction& inst, bool has_changed) {
+    if (has_changed)
+      get_def_use_mgr()->AnalyzeInstDefUse(&inst);
+  };
+
   false_id_ = 0;
 
   // clear collections

source/opt/inline_pass.h

@@ -24,6 +24,7 @@
 #include <vector>
 
 #include "def_use_manager.h"
+#include "decoration_manager.h"
 #include "module.h"
 #include "pass.h"
 
@@ -161,6 +162,12 @@ class InlinePass : public Pass {
   // Initialize state for optimization of |module|
   void InitializeInline(ir::IRContext* c);
 
+  // Update the DefUseManager when cloning decorations.
+  std::function<void(ir::Instruction&, bool)> update_def_use_mgr_;
+
+  // Decorations for the module we are processing TODO: move this to ir_context as well
+  std::unique_ptr<analysis::DecorationManager> dec_mgr_;
+
   // Map from function's result id to function.
   std::unordered_map<uint32_t, ir::Function*> id2function_;
 

source/opt/instruction.h

@@ -163,6 +163,10 @@ class Instruction : public utils::IntrusiveNodeBase<Instruction> {
   }
   // Gets the |index|-th logical operand.
   inline const Operand& GetOperand(uint32_t index) const;
+  // Adds |operand| to the list of operands of this instruction.
+  // It is the responsibility of the caller to make sure
+  // that the instruction remains valid.
+  inline void AddOperand(Operand &&operand);
   // Gets the |index|-th logical operand as a single SPIR-V word. This method is
   // not expected to be used with logical operands consisting of multiple SPIR-V
   // words.
@@ -257,6 +261,11 @@ inline const Operand& Instruction::GetOperand(uint32_t index) const {
   return operands_[index];
 };
 
+inline void Instruction::AddOperand(Operand &&operand) {
+  operands_.push_back(operand);
+  return;
+}
+
 inline void Instruction::SetInOperand(uint32_t index,
                                       std::vector<uint32_t>&& data) {
   assert(index + TypeResultIdCount() < operands_.size() &&

test/opt/inline_test.cpp

@@ -2248,7 +2248,255 @@ OpFunctionEnd
       true);
 }
 
+TEST_F(InlineTest, Decorated1) {
+  // Same test as Simple with the difference
+  // that OpFAdd in the outlined function is
+  // decorated with RelaxedPrecision
+  // Expected result is an equal decoration
+  // of the corresponding inlined instruction
+  //
+  // #version 140
+  //
+  // in vec4 BaseColor;
+  //
+  // float foo(vec4 bar)
+  // {
+  //     return bar.x + bar.y;
+  // }
+  //
+  // void main()
+  // {
+  //     vec4 color = vec4(foo(BaseColor));
+  //     gl_FragColor = color;
+  // }
 
+  const std::string predefs =
+  R"(OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main" %BaseColor %gl_FragColor
+OpExecutionMode %main OriginUpperLeft
+OpSource GLSL 140
+OpName %main "main"
+OpName %foo_vf4_ "foo(vf4;"
+OpName %bar "bar"
+OpName %color "color"
+OpName %BaseColor "BaseColor"
+OpName %param "param"
+OpName %gl_FragColor "gl_FragColor"
+OpDecorate %9 RelaxedPrecision
+)";
+
+  const std::string before =
+  R"(%void = OpTypeVoid
+%11 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%15 = OpTypeFunction %float %_ptr_Function_v4float
+%uint = OpTypeInt 32 0
+%uint_0 = OpConstant %uint 0
+%_ptr_Function_float = OpTypePointer Function %float
+%uint_1 = OpConstant %uint 1
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%BaseColor = OpVariable %_ptr_Input_v4float Input
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+%gl_FragColor = OpVariable %_ptr_Output_v4float Output
+%main = OpFunction %void None %11
+%22 = OpLabel
+%color = OpVariable %_ptr_Function_v4float Function
+%param = OpVariable %_ptr_Function_v4float Function
+%23 = OpLoad %v4float %BaseColor
+OpStore %param %23
+%24 = OpFunctionCall %float %foo_vf4_ %param
+%25 = OpCompositeConstruct %v4float %24 %24 %24 %24
+OpStore %color %25
+%26 = OpLoad %v4float %color
+OpStore %gl_FragColor %26
+OpReturn
+OpFunctionEnd
+)";
+
+
+  const std::string after =
+  R"(OpDecorate %37 RelaxedPrecision
+%void = OpTypeVoid
+%11 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%15 = OpTypeFunction %float %_ptr_Function_v4float
+%uint = OpTypeInt 32 0
+%uint_0 = OpConstant %uint 0
+%_ptr_Function_float = OpTypePointer Function %float
+%uint_1 = OpConstant %uint 1
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%BaseColor = OpVariable %_ptr_Input_v4float Input
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+%gl_FragColor = OpVariable %_ptr_Output_v4float Output
+%main = OpFunction %void None %11
+%22 = OpLabel
+%32 = OpVariable %_ptr_Function_float Function
+%color = OpVariable %_ptr_Function_v4float Function
+%param = OpVariable %_ptr_Function_v4float Function
+%23 = OpLoad %v4float %BaseColor
+OpStore %param %23
+%33 = OpAccessChain %_ptr_Function_float %param %uint_0
+%34 = OpLoad %float %33
+%35 = OpAccessChain %_ptr_Function_float %param %uint_1
+%36 = OpLoad %float %35
+%37 = OpFAdd %float %34 %36
+OpStore %32 %37
+%24 = OpLoad %float %32
+%25 = OpCompositeConstruct %v4float %24 %24 %24 %24
+OpStore %color %25
+%26 = OpLoad %v4float %color
+OpStore %gl_FragColor %26
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string nonEntryFuncs =
+  R"(%foo_vf4_ = OpFunction %float None %15
+%bar = OpFunctionParameter %_ptr_Function_v4float
+%27 = OpLabel
+%28 = OpAccessChain %_ptr_Function_float %bar %uint_0
+%29 = OpLoad %float %28
+%30 = OpAccessChain %_ptr_Function_float %bar %uint_1
+%31 = OpLoad %float %30
+%9 = OpFAdd %float %29 %31
+OpReturnValue %9
+OpFunctionEnd
+)";
+  SinglePassRunAndCheck<opt::InlineExhaustivePass>(
+  predefs + before + nonEntryFuncs, predefs + after + nonEntryFuncs, false,
+  true);
+}
+
+TEST_F(InlineTest, Decorated2) {
+  // Same test as Simple with the difference
+  // that the Result <id> of the outlined OpFunction
+  // is decorated with RelaxedPrecision
+  // Expected result is an equal decoration
+  // of the created return variable
+  //
+  // #version 140
+  //
+  // in vec4 BaseColor;
+  //
+  // float foo(vec4 bar)
+  // {
+  //     return bar.x + bar.y;
+  // }
+  //
+  // void main()
+  // {
+  //     vec4 color = vec4(foo(BaseColor));
+  //     gl_FragColor = color;
+  // }
+
+  const std::string predefs =
+  R"(OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %main "main" %BaseColor %gl_FragColor
+OpExecutionMode %main OriginUpperLeft
+OpSource GLSL 140
+OpName %main "main"
+OpName %foo_vf4_ "foo(vf4;"
+OpName %bar "bar"
+OpName %color "color"
+OpName %BaseColor "BaseColor"
+OpName %param "param"
+OpName %gl_FragColor "gl_FragColor"
+OpDecorate %foo_vf4_ RelaxedPrecision
+)";
+
+  const std::string before =
+  R"(%void = OpTypeVoid
+%10 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%14 = OpTypeFunction %float %_ptr_Function_v4float
+%uint = OpTypeInt 32 0
+%uint_0 = OpConstant %uint 0
+%_ptr_Function_float = OpTypePointer Function %float
+%uint_1 = OpConstant %uint 1
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%BaseColor = OpVariable %_ptr_Input_v4float Input
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+%gl_FragColor = OpVariable %_ptr_Output_v4float Output
+%main = OpFunction %void None %10
+%21 = OpLabel
+%color = OpVariable %_ptr_Function_v4float Function
+%param = OpVariable %_ptr_Function_v4float Function
+%22 = OpLoad %v4float %BaseColor
+OpStore %param %22
+%23 = OpFunctionCall %float %foo_vf4_ %param
+%24 = OpCompositeConstruct %v4float %23 %23 %23 %23
+OpStore %color %24
+%25 = OpLoad %v4float %color
+OpStore %gl_FragColor %25
+OpReturn
+OpFunctionEnd
+)";
+
+
+  const std::string after =
+  R"(OpDecorate %32 RelaxedPrecision
+%void = OpTypeVoid
+%10 = OpTypeFunction %void
+%float = OpTypeFloat 32
+%v4float = OpTypeVector %float 4
+%_ptr_Function_v4float = OpTypePointer Function %v4float
+%14 = OpTypeFunction %float %_ptr_Function_v4float
+%uint = OpTypeInt 32 0
+%uint_0 = OpConstant %uint 0
+%_ptr_Function_float = OpTypePointer Function %float
+%uint_1 = OpConstant %uint 1
+%_ptr_Input_v4float = OpTypePointer Input %v4float
+%BaseColor = OpVariable %_ptr_Input_v4float Input
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+%gl_FragColor = OpVariable %_ptr_Output_v4float Output
+%main = OpFunction %void None %10
+%21 = OpLabel
+%32 = OpVariable %_ptr_Function_float Function
+%color = OpVariable %_ptr_Function_v4float Function
+%param = OpVariable %_ptr_Function_v4float Function
+%22 = OpLoad %v4float %BaseColor
+OpStore %param %22
+%33 = OpAccessChain %_ptr_Function_float %param %uint_0
+%34 = OpLoad %float %33
+%35 = OpAccessChain %_ptr_Function_float %param %uint_1
+%36 = OpLoad %float %35
+%37 = OpFAdd %float %34 %36
+OpStore %32 %37
+%23 = OpLoad %float %32
+%24 = OpCompositeConstruct %v4float %23 %23 %23 %23
+OpStore %color %24
+%25 = OpLoad %v4float %color
+OpStore %gl_FragColor %25
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string nonEntryFuncs =
+  R"(%foo_vf4_ = OpFunction %float None %14
+%bar = OpFunctionParameter %_ptr_Function_v4float
+%26 = OpLabel
+%27 = OpAccessChain %_ptr_Function_float %bar %uint_0
+%28 = OpLoad %float %27
+%29 = OpAccessChain %_ptr_Function_float %bar %uint_1
+%30 = OpLoad %float %29
+%31 = OpFAdd %float %28 %30
+OpReturnValue %31
+OpFunctionEnd
+)";
+  SinglePassRunAndCheck<opt::InlineExhaustivePass>(
+  predefs + before + nonEntryFuncs, predefs + after + nonEntryFuncs, false,
+  true);
+}
 // TODO(greg-lunarg): Add tests to verify handling of these cases:
 //
 //    Empty modules