Inliner: Fix LoopMerge when inline into loop header of multi block loop

This adapts the fix for the single-block loop.  Split the loop like
before.  But when we move the OpLoopMerge back to the loop header,
redirect the continue target only when the original loop was a single
block loop.

Fixes https://github.com/KhronosGroup/SPIRV-Tools/issues/800

CHANGES

@@ -3,6 +3,8 @@ Revision history for SPIRV-Tools
 v2017.1-dev 2017-09-01
  - Validator:
    - Type check basic arithmetic operations
+ - Fixes:
+   #800: Inliner: Fix inlining function into header of multi-block loop
 
 v2017.0 2017-09-01
  - Update README to describe that assembler, disassembler, and binary parser support

source/opt/inline_pass.cpp

@@ -259,7 +259,9 @@ void InlinePass::GenInlineCode(
   // single block loop.  We'll wait to move the OpLoopMerge until the end
   // of the regular inlining logic, and only if necessary.
   bool caller_is_single_block_loop = false;
+  bool caller_is_loop_header = false;
   if (auto* loop_merge = call_block_itr->GetLoopMergeInst()) {
+    caller_is_loop_header = true;
     caller_is_single_block_loop =
         call_block_itr->id() ==
         loop_merge->GetSingleWordInOperand(kSpvLoopMergeContinueTargetIdInIdx);
@@ -283,8 +285,7 @@ void InlinePass::GenInlineCode(
   std::unique_ptr<ir::BasicBlock> new_blk_ptr;
   calleeFn->ForEachInst([&new_blocks, &callee2caller, &call_block_itr,
                          &call_inst_itr, &new_blk_ptr, &prevInstWasReturn,
-                         &returnLabelId, &returnVarId,
-                         caller_is_single_block_loop,
+                         &returnLabelId, &returnVarId, caller_is_loop_header,
                          callee_begins_with_structured_header, &calleeTypeId,
                          &multiBlocks, &postCallSB, &preCallSB, multiReturn,
                          &singleTripLoopHeaderId, &singleTripLoopContinueId,
@@ -335,7 +336,7 @@ void InlinePass::GenInlineCode(
             }
             new_blk_ptr->AddInstruction(std::move(cp_inst));
           }
-          if (caller_is_single_block_loop &&
+          if (caller_is_loop_header &&
               callee_begins_with_structured_header) {
             // We can't place both the caller's merge instruction and another
             // merge instruction in the same block.  So split the calling block.
@@ -489,10 +490,9 @@ void InlinePass::GenInlineCode(
     }
   });
 
-  if (caller_is_single_block_loop && (new_blocks->size() > 1)) {
+  if (caller_is_loop_header && (new_blocks->size() > 1)) {
     // Move the OpLoopMerge from the last block back to the first, where
-    // it belongs.  Also, update its continue target to point to the last
-    // block.
+    // it belongs.
     auto& first = new_blocks->front();
     auto& last = new_blocks->back();
     assert(first != last);
@@ -501,8 +501,12 @@ void InlinePass::GenInlineCode(
     auto loop_merge_itr = last->tail();
     --loop_merge_itr;
     assert(loop_merge_itr->opcode() == SpvOpLoopMerge);
-    std::unique_ptr<ir::Instruction> cp_inst(new ir::Instruction(*loop_merge_itr));
-    cp_inst->SetInOperand(kSpvLoopMergeContinueTargetIdInIdx, {last->id()});
+    std::unique_ptr<ir::Instruction> cp_inst(
+        new ir::Instruction(*loop_merge_itr));
+    if (caller_is_single_block_loop) {
+      // Also, update its continue target to point to the last block.
+      cp_inst->SetInOperand(kSpvLoopMergeContinueTargetIdInIdx, {last->id()});
+    }
     first->tail().InsertBefore(std::move(cp_inst));
 
     // Remove the loop merge from the last block.

test/opt/inline_test.cpp

@@ -1838,6 +1838,83 @@ OpFunctionEnd
       true);
 }
 
+TEST_F(InlineTest, MultiBlockLoopHeaderCallsMultiBlockCallee) {
+  // Like SingleBlockLoopCallsMultiBlockCallee but the loop has several
+  // blocks, but the function call still occurs in the loop header.
+  // Example from https://github.com/KhronosGroup/SPIRV-Tools/issues/800
+
+  const std::string predefs =
+      R"(OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %1 "main"
+OpSource OpenCL_C 120
+%bool = OpTypeBool
+%true = OpConstantTrue %bool
+%int = OpTypeInt 32 1
+%int_1 = OpConstant %int 1
+%int_2 = OpConstant %int 2
+%int_3 = OpConstant %int 3
+%int_4 = OpConstant %int 4
+%int_5 = OpConstant %int 5
+%void = OpTypeVoid
+%11 = OpTypeFunction %void
+)";
+
+  const std::string nonEntryFuncs =
+      R"(%12 = OpFunction %void None %11
+%13 = OpLabel
+%14 = OpCopyObject %int %int_1
+OpBranch %15
+%15 = OpLabel
+%16 = OpCopyObject %int %int_2
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string before =
+      R"(%1 = OpFunction %void None %11
+%17 = OpLabel
+OpBranch %18
+%18 = OpLabel
+%19 = OpCopyObject %int %int_3
+%20 = OpFunctionCall %void %12
+%21 = OpCopyObject %int %int_4
+OpLoopMerge %22 %23 None
+OpBranchConditional %true %23 %22
+%23 = OpLabel
+%24 = OpCopyObject %int %int_5
+OpBranchConditional %true %18 %22
+%22 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string after =
+      R"(%1 = OpFunction %void None %11
+%17 = OpLabel
+OpBranch %18
+%18 = OpLabel
+%19 = OpCopyObject %int %int_3
+%25 = OpCopyObject %int %int_1
+OpLoopMerge %22 %23 None
+OpBranch %26
+%26 = OpLabel
+%27 = OpCopyObject %int %int_2
+%21 = OpCopyObject %int %int_4
+OpBranchConditional %true %23 %22
+%23 = OpLabel
+%24 = OpCopyObject %int %int_5
+OpBranchConditional %true %18 %22
+%22 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndCheck<opt::InlineExhaustivePass>(
+      predefs + nonEntryFuncs + before, predefs + nonEntryFuncs + after, false,
+      true);
+}
+
 TEST_F(InlineTest, SingleBlockLoopCallsMultiBlockCalleeHavingSelectionMerge) {
   // This is similar to SingleBlockLoopCallsMultiBlockCallee except
   // that calleee block also has a merge instruction in its first block.
@@ -1927,6 +2004,86 @@ OpFunctionEnd
       true);
 }
 
+TEST_F(InlineTest, MultiBlockLoopHeaderCallsFromToMultiBlockCalleeHavingSelectionMerge) {
+  // This is similar to SingleBlockLoopCallsMultiBlockCalleeHavingSelectionMerge
+  // but the call is in the header block of a multi block loop.
+
+  const std::string predefs =
+R"(OpCapability Shader
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %1 "main"
+OpSource OpenCL_C 120
+%bool = OpTypeBool
+%true = OpConstantTrue %bool
+%int = OpTypeInt 32 1
+%int_1 = OpConstant %int 1
+%int_2 = OpConstant %int 2
+%int_3 = OpConstant %int 3
+%int_4 = OpConstant %int 4
+%int_5 = OpConstant %int 5
+%void = OpTypeVoid
+%11 = OpTypeFunction %void
+)";
+
+  const std::string nonEntryFuncs =
+  R"(%12 = OpFunction %void None %11
+%13 = OpLabel
+%14 = OpCopyObject %int %int_1
+OpSelectionMerge %15 None
+OpBranchConditional %true %15 %15
+%15 = OpLabel
+%16 = OpCopyObject %int %int_2
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string before =
+      R"(%1 = OpFunction %void None %11
+%17 = OpLabel
+OpBranch %18
+%18 = OpLabel
+%19 = OpCopyObject %int %int_3
+%20 = OpFunctionCall %void %12
+%21 = OpCopyObject %int %int_4
+OpLoopMerge %22 %23 None
+OpBranchConditional %true %23 %22
+%23 = OpLabel
+%24 = OpCopyObject %int %int_5
+OpBranchConditional %true %18 %22
+%22 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  const std::string after =
+      R"(%1 = OpFunction %void None %11
+%17 = OpLabel
+OpBranch %18
+%18 = OpLabel
+%19 = OpCopyObject %int %int_3
+OpLoopMerge %22 %23 None
+OpBranch %25
+%25 = OpLabel
+%26 = OpCopyObject %int %int_1
+OpSelectionMerge %27 None
+OpBranchConditional %true %27 %27
+%27 = OpLabel
+%28 = OpCopyObject %int %int_2
+%21 = OpCopyObject %int %int_4
+OpBranchConditional %true %23 %22
+%23 = OpLabel
+%24 = OpCopyObject %int %int_5
+OpBranchConditional %true %18 %22
+%22 = OpLabel
+OpReturn
+OpFunctionEnd
+)";
+
+  SinglePassRunAndCheck<opt::InlineExhaustivePass>(
+      predefs + nonEntryFuncs + before, predefs + nonEntryFuncs + after, false,
+      true);
+}
+
 TEST_F(InlineTest, SingleBlockLoopCallsMultiBlockCalleeHavingSelectionMergeAndMultiReturns) {
   // This is similar to SingleBlockLoopCallsMultiBlockCalleeHavingSelectionMerge
   // except that in addition to starting with a selection header, the