Add AIR compatibility

enzyme/Enzyme/ActivityAnalysis.cpp

@@ -188,6 +188,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/Enzyme.cpp

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

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/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: }

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

@@ -0,0 +1,33 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -preserve-nvvm -enzyme -mem2reg -sroa -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -passes="preserve-nvvm,enzyme,function(mem2reg,sroa,instsimplify,%simplifycfg)" -S | FileCheck %s
+
+; Metal AIR math intrinsic (air.cbrt.f32): not a real LLVM intrinsic, so
+; this exercises the new air.cbrt.f32/air.cbrt.f64 CallPattern entry.
+
+; Function Attrs: nounwind readnone uwtable
+define float @tester(float %x) {
+entry:
+  %call = call float @air.cbrt.f32(float %x)
+  ret float %call
+}
+
+define float @test_derivative(float %x) {
+entry:
+  %0 = tail call float (float (float)*, ...) @__enzyme_autodiff(float (float)* nonnull @tester, float %x)
+  ret float %0
+}
+
+declare float @air.cbrt.f32(float)
+
+; Function Attrs: nounwind
+declare float @__enzyme_autodiff(float (float)*, ...)
+
+; CHECK: define internal { float } @diffetester(float %x, float %differeturn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = call fast float @air.cbrt.f32(float %x)
+; CHECK-DAG:    [[REG1:%[0-9]+]] = fmul fast float 3.000000e+00, %x
+; CHECK-DAG:    [[REG2:%[0-9]+]] = fmul fast float %differeturn, %0
+; CHECK-NEXT:   %3 = fdiv fast float [[REG2]], [[REG1]]
+; CHECK-NEXT:   %4 = insertvalue { float } undef, float %3, 0
+; CHECK-NEXT:   ret { float } %4
+; CHECK-NEXT: }

enzyme/test/Enzyme/ReverseMode/air-fast-sin.ll

@@ -0,0 +1,35 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -preserve-nvvm -enzyme -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="preserve-nvvm,enzyme,function(mem2reg,instsimplify,%simplifycfg)" -S | FileCheck %s
+
+; Metal AIR fast-math intrinsic (air.fast_sin.f32): exercises the "fast_"
+; infix stripping in isMemFreeLibMFunction on top of the air. prefix and
+; .f32 suffix stripping -- reduces to "sin", a real LLVM intrinsic, so it's
+; handled by the same rule as air.sin.f32 (see air-sin.ll) with no
+; dedicated CallPattern entry.
+
+; Function Attrs: nounwind readnone uwtable
+define float @tester(float %x) {
+entry:
+  %0 = tail call fast float @air.fast_sin.f32(float %x)
+  ret float %0
+}
+
+define float @test_derivative(float %x) {
+entry:
+  %0 = tail call float (float (float)*, ...) @__enzyme_autodiff(float (float)* nonnull @tester, float %x)
+  ret float %0
+}
+
+; Function Attrs: nounwind readnone speculatable
+declare float @air.fast_sin.f32(float)
+
+; Function Attrs: nounwind
+declare float @__enzyme_autodiff(float (float)*, ...)
+
+; CHECK: define internal { float } @diffetester(float %x, float %differeturn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = call fast float @llvm.cos.f32(float %x)
+; CHECK-NEXT:   %1 = fmul fast float %differeturn, %0
+; CHECK-NEXT:   %2 = insertvalue { float } undef, float %1, 0
+; CHECK-NEXT:   ret { float } %2
+; CHECK-NEXT: }

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

@@ -0,0 +1,44 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -preserve-nvvm -enzyme -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="preserve-nvvm,enzyme,function(mem2reg,instsimplify,%simplifycfg)" -S | FileCheck %s
+
+; Metal AIR math intrinsic (air.pow.f32) resolved via isMemFreeLibMFunction's
+; air. prefix stripping in TypeAnalysis.h to Intrinsic::pow. The self-
+; referencing "d/dx x^y" term calls back the original air.pow.f32 (SameFunc
+; preserves whatever name matched), while the "d/dy x^y" term's cross-
+; function log companion becomes a genuine llvm.log.f32 intrinsic call.
+
+; Function Attrs: noinline nounwind readnone uwtable
+define float @tester(float %x, float %y) {
+entry:
+  %0 = tail call fast float @air.pow.f32(float %x, float %y)
+  ret float %0
+}
+
+define float @test_derivative(float %x, float %y) {
+entry:
+  %0 = tail call float (float (float, float)*, ...) @__enzyme_autodiff(float (float, float)* nonnull @tester, float %x, float %y)
+  ret float %0
+}
+
+; Function Attrs: nounwind readnone speculatable
+declare float @air.pow.f32(float, float)
+
+; Function Attrs: nounwind
+declare float @__enzyme_autodiff(float (float, float)*, ...)
+
+; CHECK: define internal {{(dso_local )?}}{ float, float } @diffetester(float %x, float %y, float %differeturn)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %[[ym1:.+]] = fsub fast float %y, 1.000000e+00
+; CHECK-NEXT:   %[[newpow:.+]] = call fast float @air.pow.f32(float %x, float %[[ym1]])
+; CHECK-NEXT:   %[[newpowdret:.+]] = fmul fast float %y, %[[newpow]]
+; CHECK-NEXT:   %[[dx:.+]] = fmul fast float %differeturn, %[[newpowdret]]
+; CHECK-NEXT:   %[[isxzero:.+]] = fcmp fast oeq float %x, 0.000000e+00
+; CHECK-NEXT:   %[[origpow:.+]] = call fast float @air.pow.f32(float %x, float %y)
+; CHECK-NEXT:   %[[logy:.+]] = call fast float @llvm.log.f32(float %x)
+; CHECK-NEXT:   %[[origpowdret:.+]] = fmul fast float %[[origpow]], %[[logy]]
+; CHECK-NEXT:   %[[guardeddy:.+]] = select fast i1 %[[isxzero]], float 0.000000e+00, float %[[origpowdret]]
+; CHECK-NEXT:   %[[dy:.+]] = fmul fast float %differeturn, %[[guardeddy]]
+; CHECK-NEXT:   %[[interres:.+]] = insertvalue { float, float } undef, float %[[dx:.+]], 0
+; CHECK-NEXT:   %[[finalres:.+]] = insertvalue { float, float } %[[interres]], float %[[dy:.+]], 1
+; CHECK-NEXT:   ret { float, float } %[[finalres]]
+; CHECK-NEXT: }