[CodeExtractor] Restore outputs after creating exit stubs

[kachkov98]

Summary

When CodeExtractor saves the result of InvokeInst at the first insertion point of the 'normal destination' basic block, this block can be omitted in the outlined region, so store is placed outside of the function. The suggested solution is to process saving outputs after creating exit stubs for new function, and stores will be placed in that blocks before return in this case.

[vsk]

Thanks! This mostly looks good, and I believe it resolves llvm.org/PR40455.

lib/Transforms/Utils/CodeExtractor.cpp

1141

Could you assert that the insertion point is in 'newFunction'?

unittests/Transforms/Utils/CodeExtractorTest.cpp

200

Are the 'entry' and 'cond.true' blocks needed to reproduce the issue?

228

Is DT needed?

[kachkov98]

lib/Transforms/Utils/CodeExtractor.cpp

1141

At this point the extracted blocks are still in the old function, and exit stubs are created immediately in new function; so insertion point should be in newFunction or in one of extracted blocks.

unittests/Transforms/Utils/CodeExtractorTest.cpp

200

Agree that test can be simplified.

228

It seems that dominator tree is only used to verify that all blocks are reachable from entry, so I think it can be safely removed from all CodeExtractor tests.

[kachkov98]

Add assertion for the InsertPt
Remove DT construction from CodeExtractor tests
Simplify test

[vsk]

Thanks for doing this! LGTM with a simple change.

lib/Transforms/Utils/CodeExtractor.cpp

1142

It looks like we always deref InsertPt below. Can we just define 'InsertBefore = &*InsertPt' once, and then use it in the assert? Also, it might be simpler to write InsertBefore->getFunction().

[kachkov98]

Thank you! Could you please commit it? (I don't have access to do it)

lib/Transforms/Utils/CodeExtractor.cpp

CallInst *CodeExtractor::emitCallAndSwitchStatement(Function *newFunction,
   }
 
   Function::arg_iterator OutputArgBegin = newFunction->arg_begin();
   unsigned FirstOut = inputs.size();
   if (!AggregateArgs)
     std::advance(OutputArgBegin, inputs.size());
 
   // Reload the outputs passed in by reference.
-  Function::arg_iterator OAI = OutputArgBegin;
   for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
     Value *Output = nullptr;
     if (AggregateArgs) {
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
       Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
           StructArgTy, Struct, Idx, "gep_reload_" + outputs[i]->getName());
       codeReplacer->getInstList().push_back(GEP);
       Output = GEP;
     } else {
       Output = ReloadOutputs[i];
     }
     LoadInst *load = new LoadInst(outputs[i]->getType(), Output,
                                   outputs[i]->getName() + ".reload");
     Reloads.push_back(load);
     codeReplacer->getInstList().push_back(load);
     std::vector<User *> Users(outputs[i]->user_begin(), outputs[i]->user_end());
     for (unsigned u = 0, e = Users.size(); u != e; ++u) {
       Instruction *inst = cast<Instruction>(Users[u]);
       if (!Blocks.count(inst->getParent()))
         inst->replaceUsesOfWith(outputs[i], load);
     }
-
-    // Store to argument right after the definition of output value.
-    auto *OutI = dyn_cast<Instruction>(outputs[i]);
-    if (!OutI)
-      continue;
-
-    // Find proper insertion point.
-    BasicBlock::iterator InsertPt;
-    // In case OutI is an invoke, we insert the store at the beginning in the
-    // 'normal destination' BB. Otherwise we insert the store right after OutI.
-    if (auto *InvokeI = dyn_cast<InvokeInst>(OutI))
-      InsertPt = InvokeI->getNormalDest()->getFirstInsertionPt();
-    else if (auto *Phi = dyn_cast<PHINode>(OutI))
-      InsertPt = Phi->getParent()->getFirstInsertionPt();
-    else
-      InsertPt = std::next(OutI->getIterator());
-
-    assert(OAI != newFunction->arg_end() &&
-           "Number of output arguments should match "
-           "the amount of defined values");
-    if (AggregateArgs) {
-      Value *Idx[2];
-      Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
-      Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
-      GetElementPtrInst *GEP = GetElementPtrInst::Create(
-          StructArgTy, &*OAI, Idx, "gep_" + outputs[i]->getName(), &*InsertPt);
-      new StoreInst(outputs[i], GEP, &*InsertPt);
-      // Since there should be only one struct argument aggregating
-      // all the output values, we shouldn't increment OAI, which always
-      // points to the struct argument, in this case.
-    } else {
-      new StoreInst(outputs[i], &*OAI, &*InsertPt);
-      ++OAI;
-    }
   }
 
   // Now we can emit a switch statement using the call as a value.
   SwitchInst *TheSwitch =
       SwitchInst::Create(Constant::getNullValue(Type::getInt16Ty(Context)),
                          codeReplacer, 0, codeReplacer);
 
   // Since there may be multiple exits from the original region, make the new
for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
                              OldTarget);
         }
 
         // rewrite the original branch instruction with this new target
         TI->setSuccessor(i, NewTarget);
       }
   }
 
+  // Store the arguments right after the definition of output value.
+  // This should be proceeded after creating exit stubs to be ensure that invoke
+  // result restore will be placed in the outlined function.
+  Function::arg_iterator OAI = OutputArgBegin;
+  for (unsigned i = 0, e = outputs.size(); i != e; ++i) {
+    auto *OutI = dyn_cast<Instruction>(outputs[i]);
+    if (!OutI)
+      continue;
+
+    // Find proper insertion point.
+    BasicBlock::iterator InsertPt;
+    // In case OutI is an invoke, we insert the store at the beginning in the
+    // 'normal destination' BB. Otherwise we insert the store right after OutI.
+    if (auto *InvokeI = dyn_cast<InvokeInst>(OutI))
+      InsertPt = InvokeI->getNormalDest()->getFirstInsertionPt();
+    else if (auto *Phi = dyn_cast<PHINode>(OutI))
+      InsertPt = Phi->getParent()->getFirstInsertionPt();
+    else
+      InsertPt = std::next(OutI->getIterator());
+
+    assert(((*InsertPt).getParent()->getParent() == newFunction ||
+            Blocks.count((*InsertPt).getParent())) &&
+           "InsertPt should be in new function");
+    assert(OAI != newFunction->arg_end() &&
+           "Number of output arguments should match "
+           "the amount of defined values");
+    if (AggregateArgs) {
+      Value *Idx[2];
+      Idx[0] = Constant::getNullValue(Type::getInt32Ty(Context));
+      Idx[1] = ConstantInt::get(Type::getInt32Ty(Context), FirstOut + i);
+      GetElementPtrInst *GEP = GetElementPtrInst::Create(
+          StructArgTy, &*OAI, Idx, "gep_" + outputs[i]->getName(), &*InsertPt);
+      new StoreInst(outputs[i], GEP, &*InsertPt);
+      // Since there should be only one struct argument aggregating
+      // all the output values, we shouldn't increment OAI, which always
+      // points to the struct argument, in this case.
+    } else {
+      new StoreInst(outputs[i], &*OAI, &*InsertPt);
+      ++OAI;
+    }
+  }
+
   // Now that we've done the deed, simplify the switch instruction.
   Type *OldFnRetTy = TheSwitch->getParent()->getParent()->getReturnType();
   switch (NumExitBlocks) {
   case 0:
     // There are no successors (the block containing the switch itself), which
     // means that previously this was the last part of the function, and hence
     // this should be rewritten as a `ret'
 

unittests/Transforms/Utils/CodeExtractorTest.cpp

TEST(CodeExtractor, ExitStub) {
   )invalid",
                                                 Err, Ctx));
 
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 3> Candidates{ getBlockByName(Func, "header"),
                                            getBlockByName(Func, "body1"),
                                            getBlockByName(Func, "body2") };
 
-  DominatorTree DT(*Func);
-  CodeExtractor CE(Candidates, &DT);
+  CodeExtractor CE(Candidates);
   EXPECT_TRUE(CE.isEligible());
 
   Function *Outlined = CE.extractCodeRegion();
   EXPECT_TRUE(Outlined);
   BasicBlock *Exit = getBlockByName(Func, "notExtracted");
   BasicBlock *ExitSplit = getBlockByName(Outlined, "notExtracted.split");
   // Ensure that PHI in exit block has only one incoming value (from code
   // replacer block).
TEST(CodeExtractor, ExitPHIOnePredFromRegion) {
   )invalid", Err, Ctx));
 
   Function *Func = M->getFunction("foo");
   SmallVector<BasicBlock *, 2> ExtractedBlocks{
     getBlockByName(Func, "extracted1"),
     getBlockByName(Func, "extracted2")
   };
 
-  DominatorTree DT(*Func);
-  CodeExtractor CE(ExtractedBlocks, &DT);
+  CodeExtractor CE(ExtractedBlocks);
   EXPECT_TRUE(CE.isEligible());
 
   Function *Outlined = CE.extractCodeRegion();
   EXPECT_TRUE(Outlined);
   BasicBlock *Exit1 = getBlockByName(Func, "exit1");
   BasicBlock *Exit2 = getBlockByName(Func, "exit2");
   // Ensure that PHIs in exits are not splitted (since that they have only one
   // incoming value from extracted region).
TEST(CodeExtractor, StoreOutputInvokeResultAfterEHPad) {
 
   SmallVector<BasicBlock *, 2> ExtractedBlocks{
     getBlockByName(Func, "catch"),
     getBlockByName(Func, "invoke.cont2"),
     getBlockByName(Func, "invoke.cont3"),
     getBlockByName(Func, "lpad2")
   };
 
-  DominatorTree DT(*Func);
-  CodeExtractor CE(ExtractedBlocks, &DT);
+  CodeExtractor CE(ExtractedBlocks);
   EXPECT_TRUE(CE.isEligible());
 
   Function *Outlined = CE.extractCodeRegion();
   EXPECT_TRUE(Outlined);
   EXPECT_FALSE(verifyFunction(*Outlined, &errs()));
   EXPECT_FALSE(verifyFunction(*Func, &errs()));
 }
 
+TEST(CodeExtractor, StoreOutputInvokeResultInExitStub) {
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M(parseAssemblyString(R"invalid(
+    declare i32 @bar()
+
+    define i32 @foo() personality i8* null {
+    entry:
+      %0 = invoke i32 @bar() to label %exit unwind label %lpad
+
+    exit:
+      ret i32 %0
+
+    lpad:
+      %1 = landingpad { i8*, i32 }
+              cleanup
+      resume { i8*, i32 } %1
+    }
+  )invalid",
+                                                Err, Ctx));
+
+  Function *Func = M->getFunction("foo");
+  SmallVector<BasicBlock *, 1> Blocks{ getBlockByName(Func, "entry"),
+                                       getBlockByName(Func, "lpad") };
+
+  CodeExtractor CE(Blocks);
+  EXPECT_TRUE(CE.isEligible());
+
+  Function *Outlined = CE.extractCodeRegion();
+  EXPECT_TRUE(Outlined);
+  EXPECT_FALSE(verifyFunction(*Outlined));
+  EXPECT_FALSE(verifyFunction(*Func));
+}
+
 } // end anonymous namespace