Improve no return index error (#2040)

EnzymeAD · Aug 11, 2024 · 0367300 · 0367300
1 parent 5acac24
commit 0367300
Showing 1 changed file with 34 additions and 11 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -5658,18 +5658,41 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
           if (Mode == DerivativeMode::ReverseModeCombined ||
               Mode == DerivativeMode::ReverseModePrimal) {
 
-            auto drval = *differetIdx;
-            newip = (drval < 0)
-                        ? augmentcall
-                        : BuilderZ.CreateExtractValue(augmentcall,
-                                                      {(unsigned)drval},
-                                                      call.getName() + "'ac");
-            assert(newip->getType() == placeholder->getType());
-            placeholder->replaceAllUsesWith(newip);
-            if (placeholder == &*BuilderZ.GetInsertPoint()) {
-              BuilderZ.SetInsertPoint(placeholder->getNextNode());
+            if (!differetIdx) {
+              std::string str;
+              raw_string_ostream ss(str);
+              ss << "Did not have return index set when differentiating "
+                    "function\n";
+              ss << " call" << call << "\n";
+              ss << " augmentcall" << *augmentcall << "\n";
+              if (CustomErrorHandler) {
+                CustomErrorHandler(str.c_str(), wrap(&call),
+                                   ErrorType::InternalError, nullptr, nullptr,
+                                   nullptr);
+              } else {
+                EmitFailure("GetIndexError", call.getDebugLoc(), &call,
+                            ss.str());
+              }
+              placeholder->replaceAllUsesWith(
+                  UndefValue::get(placeholder->getType()));
+              if (placeholder == &*BuilderZ.GetInsertPoint()) {
+                BuilderZ.SetInsertPoint(placeholder->getNextNode());
+              }
+              gutils->erase(placeholder);
+            } else {
+              auto drval = *differetIdx;
+              newip = (drval < 0)
+                          ? augmentcall
+                          : BuilderZ.CreateExtractValue(augmentcall,
+                                                        {(unsigned)drval},
+                                                        call.getName() + "'ac");
+              assert(newip->getType() == placeholder->getType());
+              placeholder->replaceAllUsesWith(newip);
+              if (placeholder == &*BuilderZ.GetInsertPoint()) {
+                BuilderZ.SetInsertPoint(placeholder->getNextNode());
+              }
+              gutils->erase(placeholder);
             }
-            gutils->erase(placeholder);
           } else {
             newip = placeholder;
           }