Improve cache index error message (#2138)

* Improve cache index error message * fix * fix * fix
EnzymeAD · Oct 29, 2024 · e9e4ef3 · e9e4ef3
1 parent fffdd2c
commit e9e4ef3
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -5493,14 +5493,29 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
             It.setHeadBit(true);
             BuilderZ.SetInsertPoint(It);
 #endif
-            if (Mode == DerivativeMode::ReverseModeCombined)
-              cachereplace = newCall;
-            else
-              cachereplace = BuilderZ.CreatePHI(call.getType(), 1,
-                                                call.getName() + "_tmpcacheB");
-            cachereplace = gutils->cacheForReverse(
-                BuilderZ, cachereplace,
-                getIndex(&call, CacheType::Self, BuilderZ));
+            auto idx = getIndex(&call, CacheType::Self, BuilderZ);
+            if (idx == IndexMappingError) {
+              std::string str;
+              raw_string_ostream ss(str);
+              ss << "Failed to compute consistent cache index for operation: "
+                 << call << "\n";
+              if (CustomErrorHandler) {
+                CustomErrorHandler(str.c_str(), wrap(&call),
+                                   ErrorType::InternalError, nullptr, nullptr,
+                                   nullptr);
+              } else {
+                EmitFailure("GetIndexError", call.getDebugLoc(), &call,
+                            ss.str());
+              }
+            } else {
+              if (Mode == DerivativeMode::ReverseModeCombined)
+                cachereplace = newCall;
+              else
+                cachereplace = BuilderZ.CreatePHI(
+                    call.getType(), 1, call.getName() + "_tmpcacheB");
+              cachereplace =
+                  gutils->cacheForReverse(BuilderZ, cachereplace, idx);
+            }
           } else {
 #if LLVM_VERSION_MAJOR >= 18
             auto It = BuilderZ.GetInsertPoint();

diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -1947,6 +1947,7 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
     if (EmitNoDerivativeError(ss.str(), todiff, context)) {
       auto newFunc = todiff;
       std::map<AugmentedStruct, int> returnMapping;
+      returnMapping[AugmentedStruct::Return] = -1;
       return insert_or_assign<AugmentedCacheKey, AugmentedReturn>(
                  AugmentedCachedFunctions, tup,
                  AugmentedReturn(newFunc, nullptr, {}, returnMapping, {}, {},