Rewire return type generation to explicitly mark primal/shadow (#2425)

* Mark readonly or throw as no-propagate

* Rewire return type generation to explicitly mark primal/shadow

* fix

* fix

* fmt

* fix

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -5005,14 +5005,10 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
         auto ft = call.getFunctionType();
         bool retActive = subretType != DIFFE_TYPE::CONSTANT;
 
-        ReturnType subretVal =
-            subretused
-                ? (retActive ? ReturnType::TwoReturns : ReturnType::Return)
-                : (retActive ? ReturnType::Return : ReturnType::Void);
-
         FT = getFunctionTypeForClone(
             ft, Mode, gutils->getWidth(), tape ? tape->getType() : nullptr,
-            argsInverted, false, subretVal, subretType);
+            argsInverted, false, /*returnTape*/ false,
+            /*returnPrimal*/ subretused, /*returnShadow*/ retActive);
         PointerType *fptype = PointerType::getUnqual(FT);
         newcalled = BuilderZ.CreatePointerCast(newcalled,
                                                PointerType::getUnqual(fptype));

enzyme/Enzyme/DiffeGradientUtils.cpp

@@ -89,7 +89,7 @@ DiffeGradientUtils *DiffeGradientUtils::CreateFromClone(
     bool strongZero, unsigned width, Function *todiff, TargetLibraryInfo &TLI,
     TypeAnalysis &TA, FnTypeInfo &oldTypeInfo, DIFFE_TYPE retType,
     bool shadowReturn, bool diffeReturnArg, ArrayRef<DIFFE_TYPE> constant_args,
-    ReturnType returnValue, Type *additionalArg, bool omp) {
+    bool returnTape, bool returnPrimal, Type *additionalArg, bool omp) {
   Function *oldFunc = todiff;
   assert(mode == DerivativeMode::ReverseModeGradient ||
          mode == DerivativeMode::ReverseModeCombined ||
@@ -126,7 +126,8 @@ DiffeGradientUtils *DiffeGradientUtils::CreateFromClone(
 
   auto newFunc = Logic.PPC.CloneFunctionWithReturns(
       mode, width, oldFunc, invertedPointers, constant_args, constant_values,
-      nonconstant_values, returnvals, returnValue, retType,
+      nonconstant_values, returnvals, returnTape, returnPrimal,
+      (mode == DerivativeMode::ReverseModeGradient) ? false : shadowReturn,
       prefix + oldFunc->getName(), &originalToNew,
       /*diffeReturnArg*/ diffeReturnArg, additionalArg);
 

enzyme/Enzyme/DiffeGradientUtils.h

@@ -83,8 +83,8 @@ class DiffeGradientUtils final : public GradientUtils {
                   llvm::TargetLibraryInfo &TLI, TypeAnalysis &TA,
                   FnTypeInfo &oldTypeInfo, DIFFE_TYPE retType,
                   bool shadowReturnArg, bool diffeReturnArg,
-                  llvm::ArrayRef<DIFFE_TYPE> constant_args,
-                  ReturnType returnValue, llvm::Type *additionalArg, bool omp);
+                  llvm::ArrayRef<DIFFE_TYPE> constant_args, bool returnTape,
+                  bool returnPrimal, llvm::Type *additionalArg, bool omp);
 
   llvm::AllocaInst *getDifferential(llvm::Value *val);
 

enzyme/Enzyme/EnzymeLogic.cpp

@@ -3166,7 +3166,8 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
 }
 
 void createTerminator(DiffeGradientUtils *gutils, BasicBlock *oBB,
-                      DIFFE_TYPE retType, ReturnType retVal) {
+                      DIFFE_TYPE retType, bool returnPrimal,
+                      bool returnShadow) {
   TypeResults &TR = gutils->TR;
   ReturnInst *inst = dyn_cast<ReturnInst>(oBB->getTerminator());
   // In forward mode we only need to update the return value
@@ -3212,75 +3213,45 @@ void createTerminator(DiffeGradientUtils *gutils, BasicBlock *oBB,
     }
   }
 
-  switch (retVal) {
-  case ReturnType::Return: {
-    auto ret = inst->getOperand(0);
-
-    Type *rt = ret->getType();
-    while (auto AT = dyn_cast<ArrayType>(rt))
-      rt = AT->getElementType();
-    bool floatLike = rt->isFPOrFPVectorTy();
-
-    if (retType == DIFFE_TYPE::CONSTANT) {
-      toret = gutils->getNewFromOriginal(ret);
-    } else if (!floatLike &&
-               TR.getReturnAnalysis().Inner0().isPossiblePointer()) {
-      toret = invertedPtr ? invertedPtr : gutils->invertPointerM(ret, nBuilder);
-    } else if (!gutils->isConstantValue(ret)) {
-      assert(!invertedPtr);
-      toret = gutils->diffe(ret, nBuilder);
-    } else {
-      toret = invertedPtr
-                  ? invertedPtr
-                  : gutils->invertPointerM(ret, nBuilder, /*nullInit*/ true);
-    }
+  Value *primal = nullptr;
+  Value *shadow = nullptr;
 
-    break;
+  if (returnPrimal) {
+    auto ret = inst->getOperand(0);
+    primal = gutils->getNewFromOriginal(ret);
   }
-  case ReturnType::TwoReturns: {
-    if (retType == DIFFE_TYPE::CONSTANT)
-      assert(false && "Invalid return type");
+  if (returnShadow) {
     auto ret = inst->getOperand(0);
-
     Type *rt = ret->getType();
     while (auto AT = dyn_cast<ArrayType>(rt))
       rt = AT->getElementType();
     bool floatLike = rt->isFPOrFPVectorTy();
 
-    toret =
-        nBuilder.CreateInsertValue(toret, gutils->getNewFromOriginal(ret), 0);
-
     if (!floatLike && TR.getReturnAnalysis().Inner0().isPossiblePointer()) {
-      toret = nBuilder.CreateInsertValue(
-          toret,
-          invertedPtr ? invertedPtr : gutils->invertPointerM(ret, nBuilder), 1);
+      shadow =
+          invertedPtr ? invertedPtr : gutils->invertPointerM(ret, nBuilder);
     } else if (!gutils->isConstantValue(ret)) {
       assert(!invertedPtr);
-      toret =
-          nBuilder.CreateInsertValue(toret, gutils->diffe(ret, nBuilder), 1);
+      shadow = gutils->diffe(ret, nBuilder);
     } else {
-      toret = nBuilder.CreateInsertValue(
-          toret,
-          invertedPtr
-              ? invertedPtr
-              : gutils->invertPointerM(ret, nBuilder, /*nullInit*/ true),
-          1);
+      shadow = invertedPtr
+                   ? invertedPtr
+                   : gutils->invertPointerM(ret, nBuilder, /*nullInit*/ true);
     }
-    break;
   }
-  case ReturnType::Void: {
+
+  if (primal && shadow) {
+    toret = nBuilder.CreateInsertValue(toret, primal, 0);
+    toret = nBuilder.CreateInsertValue(toret, shadow, 1);
+  } else if (primal) {
+    toret = primal;
+  } else if (shadow) {
+    toret = shadow;
+  } else {
     gutils->erase(gutils->getNewFromOriginal(inst));
     nBuilder.CreateRetVoid();
     return;
   }
-  default: {
-    llvm::errs() << "Invalid return type: " << to_string(retVal)
-                 << "for function: \n"
-                 << gutils->newFunc << "\n";
-    assert(false && "Invalid return type for function");
-    return;
-  }
-  }
 
   gutils->erase(newInst);
   nBuilder.CreateRet(toret);
@@ -4234,19 +4205,14 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
     assert(augmenteddata->constant_args == key.constant_args);
   }
 
-  ReturnType retVal =
-      key.returnUsed ? (key.shadowReturnUsed ? ReturnType::ArgsWithTwoReturns
-                                             : ReturnType::ArgsWithReturn)
-                     : (key.shadowReturnUsed ? ReturnType::ArgsWithReturn
-                                             : ReturnType::Args);
-
   bool diffeReturnArg = key.retType == DIFFE_TYPE::OUT_DIFF;
 
   DiffeGradientUtils *gutils = DiffeGradientUtils::CreateFromClone(
       *this, key.mode, key.runtimeActivity, key.strongZero, key.width,
       key.todiff, TLI, TA, oldTypeInfo, key.retType,
       augmenteddata ? augmenteddata->shadowReturnUsed : key.shadowReturnUsed,
-      diffeReturnArg, key.constant_args, retVal, key.additionalType, omp);
+      diffeReturnArg, key.constant_args, /*returnTape*/ false, key.returnUsed,
+      key.additionalType, omp);
 
   gutils->AtomicAdd = key.AtomicAdd;
   gutils->FreeMemory = key.freeMemory;
@@ -4904,17 +4870,13 @@ Function *EnzymeLogic::CreateForwardDiff(
 
   bool retActive = retType != DIFFE_TYPE::CONSTANT;
 
-  ReturnType retVal =
-      returnUsed ? (retActive ? ReturnType::TwoReturns : ReturnType::Return)
-                 : (retActive ? ReturnType::Return : ReturnType::Void);
-
   bool diffeReturnArg = false;
 
   DiffeGradientUtils *gutils = DiffeGradientUtils::CreateFromClone(
       *this, mode, runtimeActivity, strongZero, width, todiff, TLI, TA,
       oldTypeInfo, retType,
-      /*shadowReturn*/ retActive, diffeReturnArg, constant_args, retVal,
-      additionalArg, omp);
+      /*shadowReturn*/ retActive, diffeReturnArg, constant_args,
+      /*returnTape*/ false, returnUsed, additionalArg, omp);
 
   insert_or_assign2<ForwardCacheKey, Function *>(ForwardCachedFunctions, tup,
                                                  gutils->newFunc);
@@ -5092,7 +5054,7 @@ Function *EnzymeLogic::CreateForwardDiff(
       maker->visit(&*it);
     }
 
-    createTerminator(gutils, &oBB, retType, retVal);
+    createTerminator(gutils, &oBB, retType, returnUsed, retActive);
   }
 
   if (mode == DerivativeMode::ForwardModeSplit && augmenteddata)

enzyme/Enzyme/FunctionUtils.cpp

@@ -2423,31 +2423,18 @@ Function *PreProcessCache::preprocessForClone(Function *F,
   return NewF;
 }
 
-FunctionType *getFunctionTypeForClone(
-    llvm::FunctionType *FTy, DerivativeMode mode, unsigned width,
-    llvm::Type *additionalArg, llvm::ArrayRef<DIFFE_TYPE> constant_args,
-    bool diffeReturnArg, ReturnType returnValue, DIFFE_TYPE returnType) {
+FunctionType *getFunctionTypeForClone(llvm::FunctionType *FTy,
+                                      DerivativeMode mode, unsigned width,
+                                      llvm::Type *additionalArg,
+                                      llvm::ArrayRef<DIFFE_TYPE> constant_args,
+                                      bool diffeReturnArg, bool returnTape,
+                                      bool returnPrimal, bool returnShadow) {
   SmallVector<Type *, 4> RetTypes;
-  if (returnValue == ReturnType::ArgsWithReturn ||
-      returnValue == ReturnType::Return) {
-    if (returnType != DIFFE_TYPE::CONSTANT &&
-        returnType != DIFFE_TYPE::OUT_DIFF) {
-      RetTypes.push_back(
-          GradientUtils::getShadowType(FTy->getReturnType(), width));
-    } else {
-      RetTypes.push_back(FTy->getReturnType());
-    }
-  } else if (returnValue == ReturnType::ArgsWithTwoReturns ||
-             returnValue == ReturnType::TwoReturns) {
+  if (returnPrimal)
     RetTypes.push_back(FTy->getReturnType());
-    if (returnType != DIFFE_TYPE::CONSTANT &&
-        returnType != DIFFE_TYPE::OUT_DIFF) {
-      RetTypes.push_back(
-          GradientUtils::getShadowType(FTy->getReturnType(), width));
-    } else {
-      RetTypes.push_back(FTy->getReturnType());
-    }
-  }
+  if (returnShadow)
+    RetTypes.push_back(
+        GradientUtils::getShadowType(FTy->getReturnType(), width));
   SmallVector<Type *, 4> ArgTypes;
 
   // The user might be deleting arguments to the function by specifying them in
@@ -2459,7 +2446,7 @@ FunctionType *getFunctionTypeForClone(
     if (constant_args[argno] == DIFFE_TYPE::DUP_ARG ||
         constant_args[argno] == DIFFE_TYPE::DUP_NONEED) {
       ArgTypes.push_back(GradientUtils::getShadowType(I, width));
-    } else if (constant_args[argno] == DIFFE_TYPE::OUT_DIFF) {
+    } else if (constant_args[argno] == DIFFE_TYPE::OUT_DIFF && !returnTape) {
       RetTypes.push_back(GradientUtils::getShadowType(I, width));
     }
     ++argno;
@@ -2474,31 +2461,19 @@ FunctionType *getFunctionTypeForClone(
     ArgTypes.push_back(additionalArg);
   }
   Type *RetType = StructType::get(FTy->getContext(), RetTypes);
-  if (returnValue == ReturnType::TapeAndTwoReturns ||
-      returnValue == ReturnType::TapeAndReturn ||
-      returnValue == ReturnType::Tape) {
-    RetTypes.clear();
-    RetTypes.push_back(getDefaultAnonymousTapeType(FTy->getContext()));
-    if (returnValue == ReturnType::TapeAndTwoReturns) {
-      RetTypes.push_back(FTy->getReturnType());
-      RetTypes.push_back(
-          GradientUtils::getShadowType(FTy->getReturnType(), width));
-    } else if (returnValue == ReturnType::TapeAndReturn) {
-      if (returnType != DIFFE_TYPE::CONSTANT &&
-          returnType != DIFFE_TYPE::OUT_DIFF)
-        RetTypes.push_back(
-            GradientUtils::getShadowType(FTy->getReturnType(), width));
-      else
-        RetTypes.push_back(FTy->getReturnType());
-    }
-    RetType = StructType::get(FTy->getContext(), RetTypes);
-  } else if (returnValue == ReturnType::Return) {
-    assert(RetTypes.size() == 1);
-    RetType = RetTypes[0];
-  } else if (returnValue == ReturnType::TwoReturns) {
-    assert(RetTypes.size() == 2);
+  if (returnTape) {
+    RetTypes.insert(RetTypes.begin(),
+                    getDefaultAnonymousTapeType(FTy->getContext()));
   }
 
+  if (RetTypes.size() == 0)
+    RetType = Type::getVoidTy(RetType->getContext());
+  else if (RetTypes.size() == 1 && (returnPrimal || returnShadow) &&
+           mode != DerivativeMode::ReverseModeCombined)
+    RetType = RetTypes[0];
+  else
+    RetType = StructType::get(FTy->getContext(), RetTypes);
+
   bool noReturn = RetTypes.size() == 0;
   if (noReturn)
     RetType = Type::getVoidTy(RetType->getContext());
@@ -2511,16 +2486,16 @@ Function *PreProcessCache::CloneFunctionWithReturns(
     DerivativeMode mode, unsigned width, Function *&F,
     ValueToValueMapTy &ptrInputs, ArrayRef<DIFFE_TYPE> constant_args,
     SmallPtrSetImpl<Value *> &constants, SmallPtrSetImpl<Value *> &nonconstant,
-    SmallPtrSetImpl<Value *> &returnvals, ReturnType returnValue,
-    DIFFE_TYPE returnType, const Twine &name,
+    SmallPtrSetImpl<Value *> &returnvals, bool returnTape, bool returnPrimal,
+    bool returnShadow, const Twine &name,
     llvm::ValueMap<const llvm::Value *, AssertingReplacingVH> *VMapO,
     bool diffeReturnArg, llvm::Type *additionalArg) {
   if (!F->empty())
     F = preprocessForClone(F, mode);
   llvm::ValueToValueMapTy VMap;
   llvm::FunctionType *FTy = getFunctionTypeForClone(
       F->getFunctionType(), mode, width, additionalArg, constant_args,
-      diffeReturnArg, returnValue, returnType);
+      diffeReturnArg, returnTape, returnPrimal, returnShadow);
 
   for (BasicBlock &BB : *F) {
     if (auto ri = dyn_cast<ReturnInst>(BB.getTerminator())) {

enzyme/Enzyme/FunctionUtils.h

@@ -105,8 +105,8 @@ class PreProcessCache {
       llvm::ArrayRef<DIFFE_TYPE> constant_args,
       llvm::SmallPtrSetImpl<llvm::Value *> &constants,
       llvm::SmallPtrSetImpl<llvm::Value *> &nonconstant,
-      llvm::SmallPtrSetImpl<llvm::Value *> &returnvals, ReturnType returnValue,
-      DIFFE_TYPE returnType, const llvm::Twine &name,
+      llvm::SmallPtrSetImpl<llvm::Value *> &returnvals, bool returnTape,
+      bool returnPrimal, bool returnShadow, const llvm::Twine &name,
       llvm::ValueMap<const llvm::Value *, AssertingReplacingVH> *VMapO,
       bool diffeReturnArg, llvm::Type *additionalArg = nullptr);
 
@@ -412,7 +412,7 @@ bool couldFunctionArgumentCapture(llvm::CallInst *CI, llvm::Value *val);
 llvm::FunctionType *getFunctionTypeForClone(
     llvm::FunctionType *FTy, DerivativeMode mode, unsigned width,
     llvm::Type *additionalArg, llvm::ArrayRef<DIFFE_TYPE> constant_args,
-    bool diffeReturnArg, ReturnType returnValue, DIFFE_TYPE returnType);
+    bool diffeReturnArg, bool returnTape, bool returnPrimal, bool returnShadow);
 
 /// Lower __enzyme_todense, returning if changed.
 bool LowerSparsification(llvm::Function *F, bool replaceAll = true);

enzyme/Enzyme/GradientUtils.cpp

@@ -4352,16 +4352,6 @@ GradientUtils *GradientUtils::CreateFromClone(
     ++returnCount;
   }
 
-  ReturnType returnValue;
-  if (returnCount == 0)
-    returnValue = ReturnType::Tape;
-  else if (returnCount == 1)
-    returnValue = ReturnType::TapeAndReturn;
-  else if (returnCount == 2)
-    returnValue = ReturnType::TapeAndTwoReturns;
-  else
-    llvm_unreachable("illegal number of elements in augmented return struct");
-
   ValueToValueMapTy invertedPointers;
   SmallPtrSet<Instruction *, 4> constants;
   SmallPtrSet<Instruction *, 20> nonconstant;
@@ -4380,7 +4370,8 @@ GradientUtils *GradientUtils::CreateFromClone(
   auto newFunc = Logic.PPC.CloneFunctionWithReturns(
       DerivativeMode::ReverseModePrimal, width, oldFunc, invertedPointers,
       constant_args, constant_values, nonconstant_values, returnvals,
-      /*returnValue*/ returnValue, retType, prefix, &originalToNew,
+      /*returnTape*/ true, /*returnPrimal*/ returnUsed,
+      /*returnShadow*/ shadowReturnUsed, prefix, &originalToNew,
       /*diffeReturnArg*/ false, /*additionalArg*/ nullptr);
 
   // Convert overwritten args from the input function to the preprocessed