[Tapir] Handle structure-return parameters correctly in outlined helper functions. This commit addresses issue #38.

Author: neboat

llvm/lib/Transforms/Tapir/LoopSpawning.cpp

@@ -672,12 +672,17 @@ void DACLoopSpawning::implementDACIterSpawnOnHelper(
     IRBuilder<> Builder(&(RecurDet->front()));
     SetVector<Value*> RecurInputs;
     Function::arg_iterator AI = Helper->arg_begin();
+    // Handle an initial sret argument, if necessary.  Based on how
+    // the Helper function is created, any sret parameter will be the
+    // first parameter.
+    if (Helper->hasParamAttribute(0, Attribute::StructRet))
+      RecurInputs.insert(&*AI++);
     assert(cast<Argument>(CanonicalIVInput) == &*AI &&
-           "First argument does not match original input to canonical IV.");
+           "First non-sret argument does not match original input to canonical IV.");
     RecurInputs.insert(CanonicalIVStart);
     ++AI;
     assert(Limit == &*AI &&
-           "Second argument does not match original input to the loop limit.");
+           "Second non-sret argument does not match original input to the loop limit.");
     RecurInputs.insert(MidIter);
     ++AI;
     for (Function::arg_iterator AE = Helper->arg_end(); AI != AE; ++AI)
@@ -1103,6 +1108,7 @@ bool DACLoopSpawning::processLoop() {
   SetVector<Value *> BodyInputs, BodyOutputs;
   ValueToValueMapTy VMap, InputMap;
   std::vector<BasicBlock *> LoopBlocks;
+  Value *SRetInput = nullptr;
 
   // Get the sync region containing this Tapir loop.
   const Instruction *InputSyncRegion;
@@ -1134,6 +1140,23 @@ bool DACLoopSpawning::processLoop() {
       findInputsOutputs(Blocks, BodyInputs, BodyOutputs, &HandledExits, DT);
     }
 
+    // Scan for any sret parameters in BodyInputs and add them first.
+    if (F->hasStructRetAttr()) {
+      Function::arg_iterator ArgIter = F->arg_begin();
+      if (F->hasParamAttribute(0, Attribute::StructRet))
+	if (BodyInputs.count(&*ArgIter))
+	  SRetInput = &*ArgIter;
+      if (F->hasParamAttribute(1, Attribute::StructRet)) {
+	++ArgIter;
+	if (BodyInputs.count(&*ArgIter))
+	  SRetInput = &*ArgIter;
+      }
+    }
+    if (SRetInput) {
+      DEBUG(dbgs() << "sret input " << *SRetInput << "\n");
+      Inputs.insert(SRetInput);
+    }
+
     // Add argument for start of CanonicalIV.
     DEBUG({
         Value *CanonicalIVInput =
@@ -1407,6 +1430,9 @@ bool DACLoopSpawning::processLoop() {
   {
     // Setup arguments for call.
     SmallVector<Value *, 4> TopCallArgs;
+    // Add sret input, if it exists.
+    if (SRetInput)
+      TopCallArgs.push_back(SRetInput);
     // Add start iteration 0.
     assert(CanonicalSCEV->getStart()->isZero() &&
            "Canonical IV does not start at zero.");

llvm/lib/Transforms/Tapir/TapirUtils.cpp

@@ -339,10 +339,30 @@ Function *llvm::extractDetachBodyToFunction(
 
   // Fix up the inputs.
   SetVector<Value *> Inputs;
-  for (Value *V : BodyInputs) {
-    if (V == SyncRegion) continue;
-    if (!Inputs.count(V))
-      Inputs.insert(V);
+  {
+    // Scan for any sret parameters in BodyInputs and add them first.
+    Value *SRetInput = nullptr;
+    if (F.hasStructRetAttr()) {
+      Function::arg_iterator ArgIter = F.arg_begin();
+      if (F.hasParamAttribute(0, Attribute::StructRet))
+	if (BodyInputs.count(&*ArgIter))
+	  SRetInput = &*ArgIter;
+      if (F.hasParamAttribute(1, Attribute::StructRet)) {
+	++ArgIter;
+	if (BodyInputs.count(&*ArgIter))
+	  SRetInput = &*ArgIter;
+      }
+    }
+    if (SRetInput) {
+      DEBUG(dbgs() << "sret input " << *SRetInput << "\n");
+      Inputs.insert(SRetInput);
+    }
+    // Add the remaining inputs.
+    for (Value *V : BodyInputs) {
+      if (V == SyncRegion) continue;
+      if (!Inputs.count(V))
+	Inputs.insert(V);
+    }
   }
 
   // Clone the detached CFG into a helper function.