Add loose type fallback for atomics (rust-lang#723)

EnzymeAD · Jun 29, 2022 · 00d79e6 · 00d79e6
1 parent a16c616
commit 00d79e6
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -835,12 +835,23 @@ class AdjointGenerator
       }
     }
     if (!gutils->isConstantInstruction(&I) || !gutils->isConstantValue(&I)) {
+      if (looseTypeAnalysis) {
+        auto &DL = gutils->newFunc->getParent()->getDataLayout();
+        auto valType = I.getValOperand()->getType();
+        auto storeSize = DL.getTypeSizeInBits(valType) / 8;
+        auto fp = TR.firstPointer(storeSize, I.getPointerOperand(),
+                                  /*errifnotfound*/ false,
+                                  /*pointerIntSame*/ true);
+        if (!fp.isKnown() && valType->isIntOrIntVectorTy()) {
+          goto noerror;
+        }
+      }
       TR.dump();
       llvm::errs() << "oldFunc: " << *gutils->newFunc << "\n";
       llvm::errs() << "I: " << I << "\n";
+      assert(0 && "Active atomic inst not handled");
     }
-    assert(gutils->isConstantInstruction(&I));
-    assert(gutils->isConstantValue(&I));
+  noerror:;
 
     if (Mode == DerivativeMode::ReverseModeGradient) {
       eraseIfUnused(I, /*erase*/ true, /*check*/ false);