Additional warn reduction

EnzymeAD · Aug 8, 2021 · 203e4e0 · 203e4e0
1 parent e78afcd
commit 203e4e0
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 14 deletions.
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -617,12 +617,16 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
                            /*legalRecomputeCache*/ false);
     }
     if (!legalMove) {
-      EmitWarning("UncacheableUnwrap", load->getDebugLoc(),
-                  load->getParent()->getParent(), load->getParent(),
-                  "Load cannot be unwrapped ", *load, " in ",
-                  BuilderM.GetInsertBlock()->getName(), " - ",
-                  BuilderM.GetInsertBlock()->getParent()->getName(), " mode ",
-                  mode);
+      auto &warnMap = UnwrappedWarnings[load];
+      if (!warnMap.count(BuilderM.GetInsertBlock())) {
+        EmitWarning("UncacheableUnwrap", load->getDebugLoc(),
+                    load->getParent()->getParent(), load->getParent(),
+                    "Load cannot be unwrapped ", *load, " in ",
+                    BuilderM.GetInsertBlock()->getName(), " - ",
+                    BuilderM.GetInsertBlock()->getParent()->getName(), " mode ",
+                    mode);
+        warnMap.insert(BuilderM.GetInsertBlock());
+      }
       goto endCheck;
     }
 
@@ -721,10 +725,14 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
           legalMove = legalRecompute(dli, available, &BuilderM);
         }
         if (!legalMove) {
-          EmitWarning("UncacheableUnwrap", dli->getDebugLoc(),
-                      dli->getParent()->getParent(), dli->getParent(),
-                      "Differential Load cannot be unwrapped ", *dli, " in ",
-                      BuilderM.GetInsertBlock()->getName(), " mode ", mode);
+          auto &warnMap = UnwrappedWarnings[phi];
+          if (!warnMap.count(BuilderM.GetInsertBlock())) {
+            EmitWarning("UncacheableUnwrap", dli->getDebugLoc(),
+                        dli->getParent()->getParent(), dli->getParent(),
+                        "Differential Load cannot be unwrapped ", *dli, " in ",
+                        BuilderM.GetInsertBlock()->getName(), " mode ", mode);
+            warnMap.insert(BuilderM.GetInsertBlock());
+          }
           return nullptr;
         }
 
@@ -1402,9 +1410,13 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
       }
     }
     assert(val->getName() != "<badref>");
-    EmitWarning("NoUnwrap", inst->getDebugLoc(), oldFunc, inst->getParent(),
-                "Cannot unwrap ", *val, " in ",
-                BuilderM.GetInsertBlock()->getName());
+    auto &warnMap = UnwrappedWarnings[inst];
+    if (!warnMap.count(BuilderM.GetInsertBlock())) {
+      EmitWarning("NoUnwrap", inst->getDebugLoc(), oldFunc, inst->getParent(),
+                  "Cannot unwrap ", *val, " in ",
+                  BuilderM.GetInsertBlock()->getName());
+      warnMap.insert(BuilderM.GetInsertBlock());
+    }
   }
   return nullptr;
 }

diff --git a/enzyme/Enzyme/GradientUtils.h b/enzyme/Enzyme/GradientUtils.h
@@ -573,10 +573,16 @@ class GradientUtils : public CacheUtility {
       auto found = newToOriginalFn.find(I);
       if (found != newToOriginalFn.end()) {
         Value *orig = found->second;
-        newToOriginalFn.erase(I);
+        newToOriginalFn.erase(found);
         originalToNewFn.erase(orig);
       }
     }
+    {
+      auto found = UnwrappedWarnings.find(I);
+      if (found != UnwrappedWarnings.end()) {
+        UnwrappedWarnings.erase(found);
+      }
+    }
     unwrappedLoads.erase(I);
     for (auto v : invertedPointers) {
       if (v.second == I) {
@@ -1114,6 +1120,12 @@ class GradientUtils : public CacheUtility {
     }
   }
 
+private:
+  // For a given value, a list of basic blocks where an unwrap to has already
+  // produced a warning.
+  std::map<llvm::Instruction *, std::set<llvm::BasicBlock *>> UnwrappedWarnings;
+
+public:
   /// if full unwrap, don't just unwrap this instruction, but also its operands,
   /// etc
   Value *unwrapM(Value *const val, IRBuilder<> &BuilderM,