Fix cublas return in reverse (#1882)

* Fix cublas return in reverse * Fixup reverse caching
EnzymeAD · May 15, 2024 · 6875a9e · 6875a9e
1 parent f5e66b2
commit 6875a9e
Showing 1 changed file with 27 additions and 10 deletions.
diff --git a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp
@@ -147,20 +147,30 @@ void emit_handleBLAS(ArrayRef<TGPattern> blasPatterns, raw_ostream &os) {
      << "      return false;                                                \n"
      << "    }                                                              \n"
      << "  } else {                                                         \n"
-     << "    if (gutils->knownRecomputeHeuristic.find(&call) !=\n"
-     << "      gutils->knownRecomputeHeuristic.end()) {\n"
-     << "      if (!gutils->knownRecomputeHeuristic[&call]) {\n"
+     << "    auto found = gutils->knownRecomputeHeuristic.find(&call);      \n"
+     << "    auto end = gutils->knownRecomputeHeuristic.end();              \n"
+     << "    bool shouldErase = true;\n"
+     << "    if (found != end) {\n"
+     << "      if (!found->second) {                                        \n"
      << "       auto newCall = gutils->getNewFromOriginal(&call);\n"
      << "       llvm::IRBuilder<> BuilderZ(newCall);\n"
      << "       gutils->cacheForReverse(BuilderZ, newCall,\n"
      << "       getIndex(&call, CacheType::Self, BuilderZ));\n"
+     << "       shouldErase = false;\n"
      << "      }\n"
      << "    }\n"
-     << "    if (Mode == DerivativeMode::ReverseModeGradient) {             \n"
-     << "      eraseIfUnused(call, /*erase*/ true, /*check*/ false);        \n"
-     << "    } else {                                                       \n"
-     << "      eraseIfUnused(call);                                         \n"
-     << "    }                                                              \n"
+     << "    if (shouldErase) {\n"
+     << "      if (Mode == DerivativeMode::ReverseModeGradient) {             "
+        "\n"
+     << "        eraseIfUnused(call, /*erase*/ true, /*check*/ false);        "
+        "\n"
+     << "      } else {                                                       "
+        "\n"
+     << "        eraseIfUnused(call);                                         "
+        "\n"
+     << "      }                                                              "
+        "\n"
+     << "    }\n"
      << "  }\n"
      << "  return result;                                                   \n"
      << "}                                                                  \n";
@@ -237,8 +247,15 @@ void emit_free_and_ending(const TGPattern &pattern, raw_ostream &os) {
 
   os << "    }\n"
      << "  }\n"
-     << "                                                                   \n"
-     << "  if (gutils->knownRecomputeHeuristic.find(&call) !=\n"
+     << "                                                                   \n";
+
+  os << "  if (cublas && Mode == DerivativeMode::ReverseModeGradient && "
+        "call.getType()->isIntegerTy()) {        \n"
+     << "     gutils->replaceAWithB(gutils->getNewFromOriginal(&call), "
+        "Constant::getNullValue(call.getType()));\n"
+     << "  }\n";
+
+  os << "  if (gutils->knownRecomputeHeuristic.find(&call) !=\n"
      << "    gutils->knownRecomputeHeuristic.end()) {\n"
      << "    if (!gutils->knownRecomputeHeuristic[&call]) {\n"
      << "     auto cv = gutils->cacheForReverse(BuilderZ, newCall,\n"