Fix vector bugs (rust-lang#632)

* Fix vector

* Fix reverse chunk bug

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -11418,13 +11418,15 @@ class AdjointGenerator
 
         if (subretused) {
           Value *dcall = nullptr;
+          assert(returnIdx);
           dcall = (returnIdx.getValue() < 0)
                       ? augmentcall
                       : BuilderZ.CreateExtractValue(
                             augmentcall, {(unsigned)returnIdx.getValue()});
           gutils->originalToNewFn[orig] = dcall;
           gutils->newToOriginalFn.erase(newCall);
           gutils->newToOriginalFn[dcall] = orig;
+
           assert(dcall->getType() == orig->getType());
           assert(dcall);
 

enzyme/Enzyme/FunctionUtils.cpp

@@ -1770,7 +1770,8 @@ FunctionType *getFunctionTypeForClone(
   std::vector<Type *> RetTypes;
   if (returnValue == ReturnType::ArgsWithReturn ||
       returnValue == ReturnType::Return) {
-    if (returnType != DIFFE_TYPE::CONSTANT) {
+    if (returnType != DIFFE_TYPE::CONSTANT &&
+        returnType != DIFFE_TYPE::OUT_DIFF) {
       RetTypes.push_back(
           GradientUtils::getShadowType(FTy->getReturnType(), width));
     } else {
@@ -1779,7 +1780,8 @@ FunctionType *getFunctionTypeForClone(
   } else if (returnValue == ReturnType::ArgsWithTwoReturns ||
              returnValue == ReturnType::TwoReturns) {
     RetTypes.push_back(FTy->getReturnType());
-    if (returnType != DIFFE_TYPE::CONSTANT) {
+    if (returnType != DIFFE_TYPE::CONSTANT &&
+        returnType != DIFFE_TYPE::OUT_DIFF) {
       RetTypes.push_back(
           GradientUtils::getShadowType(FTy->getReturnType(), width));
     } else {
@@ -1822,7 +1824,8 @@ FunctionType *getFunctionTypeForClone(
       RetTypes.push_back(
           GradientUtils::getShadowType(FTy->getReturnType(), width));
     } else if (returnValue == ReturnType::TapeAndReturn) {
-      if (returnType != DIFFE_TYPE::CONSTANT)
+      if (returnType != DIFFE_TYPE::CONSTANT &&
+          returnType != DIFFE_TYPE::OUT_DIFF)
         RetTypes.push_back(
             GradientUtils::getShadowType(FTy->getReturnType(), width));
       else

enzyme/Enzyme/GradientUtils.cpp

@@ -832,37 +832,42 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
         if (pidx == nullptr)
           goto endCheck;
 
-        if (pidx->getType() != dli->getOperand(0)->getType()) {
+        if (pidx->getType() != getShadowType(dli->getOperand(0)->getType())) {
           llvm::errs() << "dli: " << *dli << "\n";
           llvm::errs() << "dli->getOperand(0): " << *dli->getOperand(0) << "\n";
           llvm::errs() << "pidx: " << *pidx << "\n";
         }
-        assert(pidx->getType() == dli->getOperand(0)->getType());
+        assert(pidx->getType() == getShadowType(dli->getOperand(0)->getType()));
+
+        Value *toreturn = applyChainRule(
+            dli->getType(), BuilderM,
+            [&](Value *pidx) {
 #if LLVM_VERSION_MAJOR > 7
-        auto toreturn =
-            BuilderM.CreateLoad(pidx->getType()->getPointerElementType(), pidx,
-                                phi->getName() + "_unwrap");
+              auto toreturn = BuilderM.CreateLoad(dli->getType(), pidx,
+                                                  phi->getName() + "_unwrap");
 #else
-        auto toreturn = BuilderM.CreateLoad(pidx, phi->getName() + "_unwrap");
+              auto toreturn =
+                  BuilderM.CreateLoad(pidx, phi->getName() + "_unwrap");
 #endif
-        if (auto newi = dyn_cast<Instruction>(toreturn)) {
-          newi->copyIRFlags(dli);
-          unwrappedLoads[toreturn] = dli;
-        }
+              if (auto newi = dyn_cast<Instruction>(toreturn)) {
+                newi->copyIRFlags(dli);
+                unwrappedLoads[toreturn] = dli;
+              }
 #if LLVM_VERSION_MAJOR >= 10
-        toreturn->setAlignment(dli->getAlign());
+              toreturn->setAlignment(dli->getAlign());
 #else
-        toreturn->setAlignment(dli->getAlignment());
+              toreturn->setAlignment(dli->getAlignment());
 #endif
-        toreturn->setVolatile(dli->isVolatile());
-        toreturn->setOrdering(dli->getOrdering());
-        toreturn->setSyncScopeID(dli->getSyncScopeID());
-        toreturn->setDebugLoc(getNewFromOriginal(dli->getDebugLoc()));
-        toreturn->setMetadata(LLVMContext::MD_tbaa,
-                              dli->getMetadata(LLVMContext::MD_tbaa));
-        toreturn->setMetadata(
-            LLVMContext::MD_invariant_group,
-            dli->getMetadata(LLVMContext::MD_invariant_group));
+              toreturn->setVolatile(dli->isVolatile());
+              toreturn->setOrdering(dli->getOrdering());
+              toreturn->setSyncScopeID(dli->getSyncScopeID());
+              toreturn->setDebugLoc(getNewFromOriginal(dli->getDebugLoc()));
+              toreturn->setMetadata(LLVMContext::MD_tbaa,
+                                    dli->getMetadata(LLVMContext::MD_tbaa));
+              return toreturn;
+            },
+            pidx);
+
         // TODO adding to cache only legal if no alias of any future writes
         if (permitCache)
           unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] =
@@ -4825,6 +4830,9 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
             new_op->setMetadata(LLVMContext::MD_invariant_group, invgroup);
           }
         }
+        if (op->getType() != inst->getType()) {
+          llvm::errs() << " op: " << *op << " inst: " << *inst << "\n";
+        }
         assert(op->getType() == inst->getType());
         if (!reduceRegister)
           lookup_cache[BuilderM.GetInsertBlock()][val] = op;