Add AIR compatibility

enzyme/Enzyme/ActivityAnalysis.cpp

@@ -247,6 +247,12 @@ const StringSet<> KnownInactiveFunctions = {
     "__nv_isinff",
     "__nv_isfinitel",
     "__nv_isfinited",
+    "air.isnan.f32",
+    "air.isnan.f64",
+    "air.isinf.f32",
+    "air.isinf.f64",
+    "air.isfinite.f32",
+    "air.isfinite.f64",
     "cublasCreate_v2",
     "cublasSetMathMode",
     "cublasSetStream_v2",

enzyme/Enzyme/DiffeGradientUtils.cpp

@@ -1093,6 +1093,11 @@ void DiffeGradientUtils::addToInvertedPtrDiffe(Instruction *orig,
   // atomics
   bool Atomic = isAtomic(origptr);
   auto Arch = llvm::Triple(newFunc->getParent()->getTargetTriple()).getArch();
+  // No need to do atomic on local memory for CUDA since it can't be raced
+  // upon
+  if (isa<AllocaInst>(TmpOrig) && isGPUArch(TT)) {
+    Atomic = false;
+  }
   // Moreover no need to do atomic on local shadows regardless since they are
   // not captured/escaping and created in this function. This assumes that
   // all additional parallelism in this function is outlined.

enzyme/Enzyme/Enzyme.cpp

@@ -1457,9 +1457,8 @@ class EnzymeBase {
     auto primalReturn = options.primalReturn;
     auto subsequent_calls_may_write = options.subsequent_calls_may_write;
 
-    auto Arch = Triple(CI->getModule()->getTargetTriple()).getArch();
-    bool AtomicAdd = Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
-                     Arch == Triple::amd_target;
+    auto TT = Triple(CI->getModule()->getTargetTriple());
+    bool AtomicAdd = isGPUArch(TT);
 
     TypeAnalysis TA(Logic);
     FnTypeInfo type_args = populate_type_args(TA, fn, mode);
@@ -2603,13 +2602,10 @@ class EnzymeBase {
       }
       TypeAnalysis TA(Logic);
 
-      auto Arch =
-          llvm::Triple(
-              CI->getParent()->getParent()->getParent()->getTargetTriple())
-              .getArch();
+      auto TT = llvm::Triple(
+          CI->getParent()->getParent()->getParent()->getTargetTriple());
 
-      bool AtomicAdd = Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
-                       Arch == Triple::amd_target;
+      bool AtomicAdd = isGPUArch(TT);
 
       IRBuilder<> Builder(CI);
       auto val = GradientUtils::GetOrCreateShadowConstant(

enzyme/Enzyme/EnzymeLogic.cpp

@@ -4548,12 +4548,9 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
 
   BasicBlock *entry = &gutils->newFunc->getEntryBlock();
 
-  auto Arch =
-      llvm::Triple(gutils->newFunc->getParent()->getTargetTriple()).getArch();
-  unsigned int SharedAddrSpace =
-      Arch == Triple::amd_target
-          ? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
-          : 3;
+  auto TT = llvm::Triple(gutils->newFunc->getParent()->getTargetTriple());
+  auto Arch = TT.getArch();
+  unsigned int SharedAddrSpace = getGPUSharedAddrSpace(TT);
 
   if (key.mode == DerivativeMode::ReverseModeCombined) {
     BasicBlock *sharedBlock = nullptr;
@@ -4562,8 +4559,7 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
         IRBuilder<> entryBuilder(gutils->inversionAllocs,
                                  gutils->inversionAllocs->begin());
 
-        if ((Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
-             Arch == Triple::amd_target) &&
+        if (isGPUArch(TT) &&
             g.getType()->getAddressSpace() == SharedAddrSpace) {
           if (sharedBlock == nullptr)
             sharedBlock = BasicBlock::Create(entry->getContext(), "shblock",

enzyme/Enzyme/GradientUtils.cpp

@@ -4675,14 +4675,10 @@ Constant *GradientUtils::GetOrCreateShadowConstant(
       return gvemd->getValue();
     }
 
-    auto Arch = llvm::Triple(arg->getParent()->getTargetTriple()).getArch();
-    int SharedAddrSpace = Arch == Triple::amd_target
-                              ? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
-                              : 3;
+    auto TT = llvm::Triple(arg->getParent()->getTargetTriple());
+    int SharedAddrSpace = getGPUSharedAddrSpace(TT);
     int AddrSpace = cast<PointerType>(arg->getType())->getAddressSpace();
-    if ((Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
-         Arch == Triple::amd_target) &&
-        AddrSpace == SharedAddrSpace) {
+    if (isGPUArch(TT) && AddrSpace == SharedAddrSpace) {
       assert(0 && "shared memory not handled in meta global");
     }
 
@@ -5764,16 +5760,10 @@ Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM,
         }
       }
 
-      auto Arch =
-          llvm::Triple(newFunc->getParent()->getTargetTriple()).getArch();
-      int SharedAddrSpace =
-          Arch == Triple::amd_target
-              ? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
-              : 3;
+      auto TT = llvm::Triple(newFunc->getParent()->getTargetTriple());
+      int SharedAddrSpace = getGPUSharedAddrSpace(TT);
       int AddrSpace = cast<PointerType>(arg->getType())->getAddressSpace();
-      if ((Arch == Triple::nvptx || Arch == Triple::nvptx64 ||
-           Arch == Triple::amd_target) &&
-          AddrSpace == SharedAddrSpace) {
+      if (isGPUArch(TT) && AddrSpace == SharedAddrSpace) {
         llvm::errs() << "warning found shared memory\n";
         Type *type = arg->getValueType();
         // TODO this needs initialization by entry
@@ -6839,12 +6829,8 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
       reduceRegister = true;
     }
     if (auto LI = dyn_cast<LoadInst>(inst)) {
-      auto Arch =
-          llvm::Triple(newFunc->getParent()->getTargetTriple()).getArch();
-      unsigned int SharedAddrSpace =
-          Arch == Triple::amd_target
-              ? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
-              : 3;
+      auto TT = llvm::Triple(newFunc->getParent()->getTargetTriple());
+      unsigned int SharedAddrSpace = getGPUSharedAddrSpace(TT);
       if (cast<PointerType>(LI->getPointerOperand()->getType())
               ->getAddressSpace() == SharedAddrSpace) {
         reduceRegister |= tryLegalRecomputeCheck &&
@@ -7194,12 +7180,8 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
 
         auto scev1 = OrigSE->getSCEV(origInst->getPointerOperand());
 
-        auto Arch =
-            llvm::Triple(newFunc->getParent()->getTargetTriple()).getArch();
-        unsigned int SharedAddrSpace =
-            Arch == Triple::amd_target
-                ? (int)AMDGPU::HSAMD::AddressSpaceQualifier::Local
-                : 3;
+        auto TT = llvm::Triple(newFunc->getParent()->getTargetTriple());
+        unsigned int SharedAddrSpace = getGPUSharedAddrSpace(TT);
         if (EnzymeSharedForward && scev1 != OrigSE->getCouldNotCompute() &&
             cast<PointerType>(orig_liobj->getType())->getAddressSpace() ==
                 SharedAddrSpace) {

enzyme/Enzyme/InstructionDerivatives.td

@@ -407,6 +407,78 @@ def : CallPattern<(Op $x),
                   [ReadNone, NoUnwind]
                   >;
 
+// Metal AIR hyperbolic intrinsics (Metal Shading Language "air.*"
+// namespace). PreserveNVVM.cpp's Implements map (which tags air.* function
+// declarations with enzyme_math/implements attributes, mirroring the
+// __nv_*/__ocml_* loops just above) handles every other AIR math function
+// with zero changes needed here -- calls get transparently redirected to
+// dispatch as if they were the plain libm name (see getFuncNameFromCall in
+// Utils.h), and existing CallPattern name lists like ["tan","tanf","tanl"]
+// already match. tanh/cosh/sinh are the one exception: their derivative
+// formulas need a same-family companion by exact string
+// (SameTypesFunc<"coshf">), and the generic fixup for that -- reusing an
+// already-declared implementation via ReplaceFunctionImplementation -- is
+// only reachable via a C-API call Enzyme.jl makes explicitly at the end of
+// its pipeline (FunctionUtils.cpp's internal call runs too early to see
+// this), and depends on the companion already being declared in the
+// module. Rather than depend on that, these get dedicated blocks so the
+// AIR-named companion (e.g. air.cosh.f32) is created directly on demand,
+// regardless of what else is present in the module.
+def : CallPattern<(Op $x),
+                  ["air.tanh.f32"],
+                  [(FDiv (DiffeRet), (FMul(Call<(SameTypesFunc<"air.cosh.f32">), [ReadNone,NoUnwind]> $x):$c, $c))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+def : CallPattern<(Op $x),
+                  ["air.tanh.f64"],
+                  [(FDiv (DiffeRet), (FMul(Call<(SameTypesFunc<"air.cosh.f64">), [ReadNone,NoUnwind]> $x):$c, $c))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+def : CallPattern<(Op $x),
+                  ["air.cosh.f32"],
+                  [(FMul (DiffeRet), (Call<(SameTypesFunc<"air.sinh.f32">), [ReadNone,NoUnwind]> $x))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+def : CallPattern<(Op $x),
+                  ["air.cosh.f64"],
+                  [(FMul (DiffeRet), (Call<(SameTypesFunc<"air.sinh.f64">), [ReadNone,NoUnwind]> $x))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+def : CallPattern<(Op $x),
+                  ["air.sinh.f32"],
+                  [(FMul (DiffeRet), (Call<(SameTypesFunc<"air.cosh.f32">), [ReadNone,NoUnwind]> $x))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+def : CallPattern<(Op $x),
+                  ["air.sinh.f64"],
+                  [(FMul (DiffeRet), (Call<(SameTypesFunc<"air.cosh.f64">), [ReadNone,NoUnwind]> $x))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+def : CallPattern<(Op $x),
+                  ["air.fast_tanh.f32"],
+                  [(FDiv (DiffeRet), (FMul(Call<(SameTypesFunc<"air.fast_cosh.f32">), [ReadNone,NoUnwind]> $x):$c, $c))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+def : CallPattern<(Op $x),
+                  ["air.fast_cosh.f32"],
+                  [(FMul (DiffeRet), (Call<(SameTypesFunc<"air.fast_sinh.f32">), [ReadNone,NoUnwind]> $x))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+def : CallPattern<(Op $x),
+                  ["air.fast_sinh.f32"],
+                  [(FMul (DiffeRet), (Call<(SameTypesFunc<"air.fast_cosh.f32">), [ReadNone,NoUnwind]> $x))],
+                  (ForwardFromSummedReverse),
+                  [ReadNone, NoUnwind]
+                  >;
+
 def : CallPattern<(Op $x),
                   ["asinh", "asinhf", "asinhl", "__nv_asinh", "__nv_asinhf"],
                   [(FDiv (DiffeRet), (Intrinsic<"sqrt"> (FAdd (FMul $x, $x), (ConstantFP<"1.0"> $x)))  )] ,

enzyme/Enzyme/PreserveNVVM.cpp

@@ -887,6 +887,42 @@ bool preserveNVVM(bool Begin, Module &M) {
 
       Implements[nvname] = std::make_pair(mathname, llname);
     }
+    // Metal AIR (air.<name>.f32 / air.<name>.f64 -- Metal has no long
+    // double, so there is no third T variant here).
+    // tanh/cosh/sinh (and their fast_ forms) are deliberately excluded here:
+    // CallPattern matching (see InstructionDerivatives.td) runs against this
+    // same enzyme_math-substituted name, so tagging air.tanh.f32 here would
+    // make it dispatch through the plain "tanhf" CallPattern -- whose
+    // companion is hardcoded to the libm name "coshf", not "air.cosh.f32" --
+    // rather than through the dedicated air.tanh.f32 CallPattern that exists
+    // specifically to keep the companion AIR-native.
+    for (std::string name : {"sin", "cos", "tan", "asin", "acos", "atan",
+                             "atan2", "exp", "exp2", "log", "log2", "log10",
+                             "log1p", "expm1", "sqrt", "cbrt", "pow", "fma"}) {
+      std::string airname = "air." + name + (T == "f" ? ".f32" : ".f64");
+      std::string llname = "llvm." + name + "." + (T == "f" ? "f32" : "f64");
+      std::string mathname = name + T;
+
+      Implements[airname] = std::make_pair(mathname, llname);
+    }
+    // Metal AIR fast-math variants (air.fast_<name>.f32 -- Metal only
+    // exposes fast math for float). These map to the same mathname/llname
+    // as the precise version above: enzyme_math dispatch and the
+    // ReplaceFunctionImplementation companion-rewrite don't distinguish
+    // fast vs precise, they just need a valid libm/llvm target name.
+    // fast_tanh/fast_cosh/fast_sinh are excluded for the same reason as
+    // their non-fast forms above.
+    if (T == "f") {
+      for (std::string name :
+           {"log", "exp", "sin", "cos", "tan", "sqrt", "asin", "acos", "atan",
+            "atan2", "acosh", "asinh"}) {
+        std::string airname = "air.fast_" + name + ".f32";
+        std::string llname = "llvm." + name + ".f32";
+        std::string mathname = name + T;
+
+        Implements[airname] = std::make_pair(mathname, llname);
+      }
+    }
   }
 #if ENZYME_ENABLE_NVVM_ATTRIBUTION
   for (auto &F : llvm::make_early_inc_range(M)) {

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h

@@ -79,6 +79,16 @@ static inline bool isMemFreeLibMFunction(llvm::StringRef str,
     str = str.substr(5, str.size() - 5);
   } else if (startsWith(str, "__ocml_")) {
     str = str.substr(7, str.size() - 7);
+  } else if (startsWith(str, "air.")) {
+    // Metal AIR math intrinsics, e.g. air.cos.f32, air.fast_tanh.f32.
+    // Defense-in-depth alongside PreserveNVVM.cpp's Implements map: this
+    // raw-name check works even in pipelines that don't run preserve-nvvm,
+    // exactly as __nv_/__ocml_ above already do for CUDA/ROCm.
+    str = str.substr(4, str.size() - 4);
+    if (endsWith(str, ".f16") || endsWith(str, ".f32") || endsWith(str, ".f64"))
+      str = str.substr(0, str.size() - 4);
+    if (startsWith(str, "fast_"))
+      str = str.substr(5, str.size() - 5);
   }
   if (LIBM_FUNCTIONS.find(str.str()) != LIBM_FUNCTIONS.end()) {
     if (ID)

enzyme/Enzyme/Utils.h

@@ -68,6 +68,29 @@
 #define amd_target amdgcn
 #endif
 
+#if LLVM_VERSION_MAJOR >= 16
+#include "llvm/TargetParser/Triple.h"
+#else
+#include "llvm/ADT/Triple.h"
+#endif
+
+#include "llvm/Support/AMDGPUMetadata.h"
+
+// Returns true if the given target triple is a GPU kernel architecture
+// (NVPTX, AMDGPU, or Apple Metal AIR).
+static inline bool isGPUArch(const llvm::Triple &TT) {
+  auto Arch = TT.getArch();
+  return Arch == llvm::Triple::nvptx || Arch == llvm::Triple::nvptx64 ||
+         Arch == llvm::Triple::amd_target || TT.getArchName() == "air64";
+}
+
+// Returns the address space used for GPU shared/threadgroup memory.
+static inline unsigned getGPUSharedAddrSpace(const llvm::Triple &TT) {
+  return TT.getArch() == llvm::Triple::amd_target
+             ? (unsigned)llvm::AMDGPU::HSAMD::AddressSpaceQualifier::Local
+             : 3;
+}
+
 #include "llvm/IR/DiagnosticInfo.h"
 
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"

enzyme/test/Enzyme/ForwardMode/air-atan2.ll

@@ -0,0 +1,35 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -preserve-nvvm -enzyme -early-cse -instcombine -enzyme-preopt=false -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -passes="preserve-nvvm,enzyme,function(early-cse,instcombine)" -enzyme-preopt=false -S | FileCheck %s
+
+; Metal AIR math intrinsic (air.atan2.f32): not a real LLVM intrinsic, so
+; this exercises the new air.atan2.f32/air.atan2.f64 CallPattern entry in
+; forward mode.
+
+define float @tester(float %y, float %x) {
+entry:
+  %call = call float @air.atan2.f32(float %y, float %x)
+  ret float %call
+}
+
+define float @test_derivative(float %y, float %x) {
+entry:
+  %0 = tail call float (...) @__enzyme_fwddiff(float (float, float)* nonnull @tester, float %y, float 1.000000e+00, float %x, float 1.000000e+00)
+  ret float %0
+}
+
+declare float @air.atan2.f32(float, float)
+
+; Function Attrs: nounwind
+declare float @__enzyme_fwddiff(...)
+
+; CHECK-LABEL: define internal float @fwddiffetester(
+; CHECK-NEXT: entry:
+; CHECK-DAG:   %[[a3:.+]] = fmul fast float %"y'", %x
+; CHECK-DAG:    %[[a1:.+]] = fmul fast float %x, %x
+; CHECK-DAG:    %[[a0:.+]] = fmul fast float %y, %y
+; CHECK-DAG:   %[[a2:.+]] = fadd fast float %[[a1]], %[[a0]]
+; CHECK-DAG:   %[[a4:.+]] = fmul fast float %"x'", %y
+; CHECK-DAG:   %[[a5:.+]] = fsub fast float %[[a3]], %[[a4]]
+; CHECK-DAG:   %[[a6:.+]] = fdiv fast float %[[a5]], %[[a2]]
+; CHECK-NEXT:   ret float %[[a6]]
+; CHECK-NEXT: }

enzyme/test/Enzyme/ReverseMode/air-atan2.ll

@@ -0,0 +1,37 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -preserve-nvvm -enzyme -mem2reg -sroa -early-cse -instsimplify -simplifycfg -adce -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -passes="preserve-nvvm,enzyme,function(mem2reg,sroa,early-cse,instsimplify,%simplifycfg,adce)" -enzyme-preopt=false -S | FileCheck %s
+
+; Metal AIR math intrinsic (air.atan2.f32): not a real LLVM intrinsic, so
+; this exercises the new air.atan2.f32/air.atan2.f64 CallPattern entry.
+
+define float @tester(float %y, float %x) {
+entry:
+  %call = call float @air.atan2.f32(float %y, float %x)
+  ret float %call
+}
+
+define float @test_derivative(float %y, float %x) {
+entry:
+  %0 = tail call float (...) @__enzyme_autodiff(float (float, float)* nonnull @tester, float %y, float %x)
+  ret float %0
+}
+
+declare float @air.atan2.f32(float, float)
+
+; Function Attrs: nounwind
+declare float @__enzyme_autodiff(...)
+
+; CHECK: define internal { float, float } @diffetester(float %y, float %x, float %differeturn)
+; CHECK-NEXT: entry:
+; CHECK-DAG:    %[[a0:.+]] = fmul fast float %y, %y
+; CHECK-DAG:    %[[a1:.+]] = fmul fast float %x, %x
+; CHECK-DAG:   %[[a2:.+]] = fadd fast float %[[a1]], %[[a0]]
+; CHECK-DAG:   %[[a3:.+]] = fmul fast float %differeturn, %x
+; CHECK-DAG:   %[[a4:.+]] = fdiv fast float %[[a3]], %[[a2]]
+; CHECK-DAG:   %[[a5:.+]] = fmul fast float %differeturn, %y
+; CHECK-DAG:   %[[a6:.+]] = fdiv fast float %[[a5]], %[[a2]]
+; CHECK-DAG:   %[[a7:.+]] = {{(fneg fast float)|(fsub fast float (-)?0.000000e\+00,)}} %[[a6]]
+; CHECK-DAG:   %[[a8:.+]] = insertvalue { float, float } undef, float %[[a4]], 0
+; CHECK-DAG:   %[[a9:.+]] = insertvalue { float, float } %[[a8]], float %[[a7]], 1
+; CHECK-DAG:   ret { float, float } %[[a9]]
+; CHECK-NEXT: }