Fix wrong DIFFE_TYPE for forward mode in GetOrCreateShadowFunction (r…

…ust-lang#553)
EnzymeAD · Mar 7, 2022 · aee3008 · aee3008
1 parent f08a1bf
commit aee3008
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -3542,7 +3542,8 @@ Constant *GradientUtils::GetOrCreateShadowFunction(
         std::pair<Argument *, std::set<int64_t>>(&a, {}));
     DIFFE_TYPE typ;
     if (a.getType()->isFPOrFPVectorTy()) {
-      typ = DIFFE_TYPE::OUT_DIFF;
+      typ = mode == DerivativeMode::ForwardMode ? DIFFE_TYPE::DUP_ARG
+                                                : DIFFE_TYPE::OUT_DIFF;
     } else if (a.getType()->isIntegerTy() &&
                cast<IntegerType>(a.getType())->getBitWidth() < 16) {
       typ = DIFFE_TYPE::CONSTANT;
@@ -3554,7 +3555,8 @@ Constant *GradientUtils::GetOrCreateShadowFunction(
     types.push_back(typ);
   }
 
-  DIFFE_TYPE retType = fn->getReturnType()->isFPOrFPVectorTy()
+  DIFFE_TYPE retType = fn->getReturnType()->isFPOrFPVectorTy() &&
+                               mode != DerivativeMode::ForwardMode
                            ? DIFFE_TYPE::OUT_DIFF
                            : DIFFE_TYPE::DUP_ARG;
   if (fn->getReturnType()->isVoidTy() || fn->getReturnType()->isEmptyTy() ||