Faster activity datastructures and handle version breaking flang (rust-lang#458)

Author: wsmoses

enzyme/CMakeLists.txt

@@ -66,6 +66,8 @@ message("LLVM_INCLUDE_DIRS: ${LLVM_INCLUDE_DIRS}")
 message("found llvm definitions " ${LLVM_DEFINITIONS})
 message("found llvm version " ${LLVM_VERSION_MAJOR})
 
+option(ENZYME_FLANG "Build for non-version compliant FLANG" OFF)
+add_definitions(-DFLANG=1)
 
 # Offer the user the choice of overriding the installation directories
 set(INSTALL_INCLUDE_DIR include CACHE PATH "Installation directory for header files")

enzyme/Enzyme/ActivityAnalysis.h

@@ -122,12 +122,12 @@ class ActivityAnalyzer {
   bool isConstantValue(TypeResults &TR, llvm::Value *val);
 
 private:
-  std::map<llvm::Instruction *, std::set<llvm::Value *>>
+  llvm::DenseMap<llvm::Instruction *, llvm::SmallPtrSet<llvm::Value *, 4>>
       ReEvaluateValueIfInactiveInst;
-  std::map<llvm::Value *, std::set<llvm::Value *>>
+  llvm::DenseMap<llvm::Value *, llvm::SmallPtrSet<llvm::Value *, 4>>
       ReEvaluateValueIfInactiveValue;
 
-  std::map<llvm::Value *, std::set<llvm::Instruction *>>
+  llvm::DenseMap<llvm::Value *, llvm::SmallPtrSet<llvm::Instruction *, 4>>
       ReEvaluateInstIfInactiveValue;
 
   void InsertConstantInstruction(TypeResults &TR, llvm::Instruction *I);

enzyme/Enzyme/CacheUtility.cpp

@@ -893,7 +893,11 @@ AllocaInst *CacheUtility::createCacheForScope(LimitContext ctx, Type *T,
 #if LLVM_VERSION_MAJOR >= 14
           malloccall->addDereferenceableRetAttr(
               ci->getLimitedValue() * byteSizeOfType->getLimitedValue());
+#ifndef FLANG
           AttrBuilder B(ci->getContext());
+#else
+          AttrBuilder B;
+#endif
           B.addDereferenceableOrNullAttr(ci->getLimitedValue() *
                                          byteSizeOfType->getLimitedValue());
           malloccall->setAttributes(

enzyme/Enzyme/Enzyme.cpp

@@ -1619,7 +1619,7 @@ class Enzyme : public ModulePass {
       // code left here to re-enable upon Attributor patch
       Logic.PPC.FAM.clear(F, F.getName());
 
-#if LLVM_VERSION_MAJOR >= 13
+#if LLVM_VERSION_MAJOR >= 13 && !defined(FLANG)
 
       AnalysisGetter AG(Logic.PPC.FAM);
       SetVector<Function *> Functions;

enzyme/Enzyme/EnzymeLogic.cpp

@@ -2180,7 +2180,11 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
       }
 #if LLVM_VERSION_MAJOR >= 14
       malloccall->addDereferenceableRetAttr(size->getLimitedValue());
+#ifndef FLANG
       AttrBuilder B(malloccall->getContext());
+#else
+      AttrBuilder B;
+#endif
       B.addDereferenceableOrNullAttr(size->getLimitedValue());
       malloccall->setAttributes(malloccall->getAttributes().addRetAttributes(
           malloccall->getContext(), B));

enzyme/Enzyme/FunctionUtils.cpp

@@ -1433,7 +1433,7 @@ Function *PreProcessCache::preprocessForClone(Function *F,
     }
 
     {
-#if LLVM_VERSION_MAJOR >= 14
+#if LLVM_VERSION_MAJOR >= 14 && !defined(FLANG)
       auto PA = SROAPass().run(*NewF, FAM);
 #else
       auto PA = SROA().run(*NewF, FAM);
@@ -1444,7 +1444,7 @@ Function *PreProcessCache::preprocessForClone(Function *F,
     ReplaceReallocs(NewF);
 
     {
-#if LLVM_VERSION_MAJOR >= 14
+#if LLVM_VERSION_MAJOR >= 14 && !defined(FLANG)
       auto PA = SROAPass().run(*NewF, FAM);
 #else
       auto PA = SROA().run(*NewF, FAM);
@@ -1973,12 +1973,12 @@ void SelectOptimization(Function *F) {
 }
 void PreProcessCache::optimizeIntermediate(Function *F) {
   PromotePass().run(*F, FAM);
-#if LLVM_VERSION_MAJOR >= 14
+#if LLVM_VERSION_MAJOR >= 14 && !defined(FLANG)
   GVNPass().run(*F, FAM);
 #else
   GVN().run(*F, FAM);
 #endif
-#if LLVM_VERSION_MAJOR >= 14
+#if LLVM_VERSION_MAJOR >= 14 && !defined(FLANG)
   SROAPass().run(*F, FAM);
 #else
   SROA().run(*F, FAM);

enzyme/Enzyme/GradientUtils.h

@@ -749,7 +749,11 @@ class GradientUtils : public CacheUtility {
 #if LLVM_VERSION_MAJOR >= 14
         cast<CallInst>(anti)->addDereferenceableRetAttr(ci->getLimitedValue());
         cal->addDereferenceableRetAttr(ci->getLimitedValue());
+#ifndef FLANG
         AttrBuilder B(Fn->getContext());
+#else
+        AttrBuilder B;
+#endif
         B.addDereferenceableOrNullAttr(ci->getLimitedValue());
         cast<CallInst>(anti)->setAttributes(
             cast<CallInst>(anti)->getAttributes().addRetAttributes(