Fix

EnzymeAD · Nov 5, 2024 · a9a6298 · a9a6298
1 parent 5d90df3
commit a9a6298
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 7 deletions.
diff --git a/enzyme/Enzyme/EnzymeLogic.h b/enzyme/Enzyme/EnzymeLogic.h
@@ -52,7 +52,6 @@
 
 extern "C" {
 extern llvm::cl::opt<bool> EnzymePrint;
-extern llvm::cl::opt<bool> EnzymeJuliaAddrLoad;
 }
 
 constexpr char EnzymeFPRTPrefix[] = "__enzyme_fprt_";

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
@@ -108,11 +108,6 @@ const llvm::StringMap<llvm::Intrinsic::ID> LIBM_FUNCTIONS = {
     {"__nv_drcp_rn", Intrinsic::not_intrinsic},
     {"__nv_drcp_ru", Intrinsic::not_intrinsic},
     {"__nv_drcp_rz", Intrinsic::not_intrinsic},
-    {"__nv_isnand", Intrinsic::not_intrinsic},
-    {"__nv_isnanf", Intrinsic::not_intrinsic},
-    {"__nv_isinfd", Intrinsic::not_intrinsic},
-    {"__nv_isinff", Intrinsic::not_intrinsic},
-    {"__nv_acos", Intrinsic::not_intrinsic},
     {"asin", Intrinsic::not_intrinsic},
     {"__nv_asin", Intrinsic::not_intrinsic},
     {"atan", Intrinsic::not_intrinsic},

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h
@@ -60,6 +60,7 @@ extern const llvm::StringMap<llvm::Intrinsic::ID> LIBM_FUNCTIONS;
 
 static inline bool isMemFreeLibMFunction(llvm::StringRef str,
                                          llvm::Intrinsic::ID *ID = nullptr) {
+  llvm::StringRef ogstr = str;
   if (startsWith(str, "__") && endsWith(str, "_finite")) {
     str = str.substr(2, str.size() - 2 - 7);
   } else if (startsWith(str, "__fd_") && endsWith(str, "_1")) {
@@ -72,7 +73,8 @@ static inline bool isMemFreeLibMFunction(llvm::StringRef str,
       *ID = LIBM_FUNCTIONS.find(str.str())->second;
     return true;
   }
-  if (endsWith(str, "f") || endsWith(str, "l")) {
+  if (endsWith(str, "f") || endsWith(str, "l") ||
+      (startsWith(ogstr, "__nv_") && endsWith(str, "d"))) {
     if (LIBM_FUNCTIONS.find(str.substr(0, str.size() - 1).str()) !=
         LIBM_FUNCTIONS.end()) {
       if (ID)

diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h
@@ -94,6 +94,7 @@ extern llvm::cl::opt<bool> EnzymePrintPerf;
 extern llvm::cl::opt<bool> EnzymeStrongZero;
 extern llvm::cl::opt<bool> EnzymeBlasCopy;
 extern llvm::cl::opt<bool> EnzymeLapackCopy;
+extern llvm::cl::opt<bool> EnzymeJuliaAddrLoad;
 extern LLVMValueRef (*CustomErrorHandler)(const char *, LLVMValueRef, ErrorType,
                                           const void *, LLVMValueRef,
                                           LLVMBuilderRef);
@@ -1184,6 +1185,15 @@ static inline bool hasNoCache(llvm::Value *op) {
   if (auto I = dyn_cast<Instruction>(op))
     if (hasMetadata(I, "enzyme_nocache"))
       return true;
+
+  if (EnzymeJuliaAddrLoad) {
+    if (auto PT = dyn_cast<PointerType>(op->getType())) {
+      if (PT->getAddressSpace() == 11 || PT->getAddressSpace() == 13) {
+        if (isa<CastInst>(op) || isa<GetElementPtrInst>(op))
+          return true;
+      }
+    }
+  }
   return false;
 }