Fix ROCM (rust-lang#669)

Co-authored-by: Leila Ghaffari <[email protected]>

Author: wsmoses

enzyme/CMakeLists.txt

@@ -128,6 +128,11 @@ if (ENZYME_FLANG)
   add_definitions(-DFLANG=1)
 endif()
 
+option(ENZYME_ROCM "Build for non-version compliant ROCM" OFF)
+if (ENZYME_ROCM)
+  add_definitions(-DROCM=1)
+endif()
+
 # Offer the user the choice of overriding the installation directories
 set(INSTALL_INCLUDE_DIR include CACHE PATH "Installation directory for header files")
 if(WIN32 AND NOT CYGWIN)

enzyme/Enzyme/AdjointGenerator.h

@@ -10487,7 +10487,7 @@ class AdjointGenerator
 #if LLVM_VERSION_MAJOR >= 14
                     cast<CallInst>(anti)->addDereferenceableRetAttr(derefBytes);
                     cal->addDereferenceableRetAttr(derefBytes);
-#ifndef FLANG
+#if !defined(FLANG) && !defined(ROCM)
                     AttrBuilder B(called->getContext());
 #else
                     AttrBuilder B;

enzyme/Enzyme/CacheUtility.cpp

@@ -896,7 +896,7 @@ AllocaInst *CacheUtility::createCacheForScope(LimitContext ctx, Type *T,
 #if LLVM_VERSION_MAJOR >= 14
           malloccall->addDereferenceableRetAttr(
               ci->getLimitedValue() * byteSizeOfType->getLimitedValue());
-#ifndef FLANG
+#if !defined(FLANG) && !defined(ROCM)
           AttrBuilder B(ci->getContext());
 #else
           AttrBuilder B;

enzyme/Enzyme/Enzyme.cpp

@@ -1801,7 +1801,7 @@ class Enzyme : public ModulePass {
       // dead internal functions, which invalidates Enzyme's cache
       // code left here to re-enable upon Attributor patch
 
-#if LLVM_VERSION_MAJOR >= 13 && !defined(FLANG)
+#if LLVM_VERSION_MAJOR >= 13 && !defined(FLANG) && !defined(ROCM)
 
       AnalysisGetter AG(Logic.PPC.FAM);
       SetVector<Function *> Functions;

enzyme/Enzyme/EnzymeLogic.cpp

@@ -2318,7 +2318,7 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
         }
 #if LLVM_VERSION_MAJOR >= 14
         malloccall->addDereferenceableRetAttr(size->getLimitedValue());
-#ifndef FLANG
+#if !defined(FLANG) && !defined(ROCM)
         AttrBuilder B(malloccall->getContext());
 #else
         AttrBuilder B;

enzyme/Enzyme/FunctionUtils.cpp

@@ -319,7 +319,7 @@ static inline void UpgradeAllocasToMallocs(Function *NewF,
 #if LLVM_VERSION_MAJOR >= 14
     if (ConstantInt *size = dyn_cast<ConstantInt>(CI->getOperand(0))) {
       CI->addDereferenceableRetAttr(size->getLimitedValue());
-#ifndef FLANG
+#if !defined(FLANG) && !defined(ROCM)
       AttrBuilder B(CI->getContext());
 #else
       AttrBuilder B;