BLAS: fix blas erasure (#1948)

EnzymeAD · Jun 29, 2024 · 015a772 · 015a772
1 parent cc545bc
commit 015a772
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -8226,6 +8226,7 @@ void GradientUtils::eraseFictiousPHIs() {
   for (auto pair : phis) {
     auto pp = pair.first;
     if (pp->getNumUses() != 0) {
+      assert(0);
       if (CustomErrorHandler) {
         std::string str;
         raw_string_ostream ss(str);

diff --git a/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp b/enzyme/tools/enzyme-tblgen/blas-tblgen.cpp
@@ -255,17 +255,22 @@ void emit_free_and_ending(const TGPattern &pattern, raw_ostream &os) {
         "Constant::getNullValue(call.getType()));\n"
      << "  }\n";
 
+  os << "  bool shouldErase = true;\n";
   os << "  if (gutils->knownRecomputeHeuristic.find(&call) !=\n"
      << "    gutils->knownRecomputeHeuristic.end()) {\n"
      << "    if (!gutils->knownRecomputeHeuristic[&call]) {\n"
      << "     auto cv = gutils->cacheForReverse(BuilderZ, newCall,\n"
      << "     getIndex(&call, CacheType::Self, BuilderZ));\n"
+     << "     shouldErase = false;\n"
      << "    }\n"
-     << "  } else if (Mode == DerivativeMode::ReverseModeGradient) {        \n"
-     << "    eraseIfUnused(call, /*erase*/ true, /*check*/ false);          \n"
-     << "  } else {                                                         \n"
-     << "    eraseIfUnused(call);                                           \n"
-     << "  }                                                                \n"
+     << "  }\n"
+     << "  if (shouldErase) {\n"
+     << "    if (Mode == DerivativeMode::ReverseModeGradient) {        \n"
+     << "      eraseIfUnused(call, /*erase*/ true, /*check*/ false);        \n"
+     << "    } else {                                                       \n"
+     << "      eraseIfUnused(call);                                         \n"
+     << "    }                                                              \n"
+     << "  }\n"
      << "  return true;\n"
      << "#ifdef __clang__\n"
      << "#pragma clang diagnostic pop\n"