From e9233933784a1f2fe2d68176986448e53c738cbb Mon Sep 17 00:00:00 2001 From: pw0908 Date: Wed, 8 Jul 2026 16:43:26 -0700 Subject: [PATCH 1/7] Add AIR compatibility --- enzyme/Enzyme/ActivityAnalysis.cpp | 6 ++ enzyme/Enzyme/InstructionDerivatives.td | 72 +++++++++++++++++++ enzyme/Enzyme/PreserveNVVM.cpp | 37 ++++++++++ enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h | 10 +++ enzyme/test/Enzyme/ForwardMode/air-atan2.ll | 35 +++++++++ enzyme/test/Enzyme/ReverseMode/air-atan2.ll | 37 ++++++++++ enzyme/test/Enzyme/ReverseMode/air-cbrt.ll | 33 +++++++++ .../test/Enzyme/ReverseMode/air-fast-sin.ll | 35 +++++++++ enzyme/test/Enzyme/ReverseMode/air-pow.ll | 44 ++++++++++++ enzyme/test/Enzyme/ReverseMode/air-sin.ll | 33 +++++++++ enzyme/test/Enzyme/ReverseMode/air-tan.ll | 35 +++++++++ enzyme/test/Enzyme/ReverseMode/air-tanh.ll | 35 +++++++++ 12 files changed, 412 insertions(+) create mode 100644 enzyme/test/Enzyme/ForwardMode/air-atan2.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/air-atan2.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/air-cbrt.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/air-fast-sin.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/air-pow.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/air-sin.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/air-tan.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/air-tanh.ll diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp index 2294462b47f..91043282597 100644 --- a/enzyme/Enzyme/ActivityAnalysis.cpp +++ b/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", diff --git a/enzyme/Enzyme/InstructionDerivatives.td b/enzyme/Enzyme/InstructionDerivatives.td index c0eeb99dcc7..eb521ba6df6 100644 --- a/enzyme/Enzyme/InstructionDerivatives.td +++ b/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))) )] , diff --git a/enzyme/Enzyme/PreserveNVVM.cpp b/enzyme/Enzyme/PreserveNVVM.cpp index 3faa22697c1..5f9ecc39940 100644 --- a/enzyme/Enzyme/PreserveNVVM.cpp +++ b/enzyme/Enzyme/PreserveNVVM.cpp @@ -887,6 +887,43 @@ bool preserveNVVM(bool Begin, Module &M) { Implements[nvname] = std::make_pair(mathname, llname); } + // Metal AIR (air..f32 / air..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_.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)) { diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h index 5984443be74..9103c59a3ed 100644 --- a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h +++ b/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) diff --git a/enzyme/test/Enzyme/ForwardMode/air-atan2.ll b/enzyme/test/Enzyme/ForwardMode/air-atan2.ll new file mode 100644 index 00000000000..8893167b471 --- /dev/null +++ b/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: } diff --git a/enzyme/test/Enzyme/ReverseMode/air-atan2.ll b/enzyme/test/Enzyme/ReverseMode/air-atan2.ll new file mode 100644 index 00000000000..b5a31d63d96 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/air-atan2.ll @@ -0,0 +1,37 @@ +; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -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: } diff --git a/enzyme/test/Enzyme/ReverseMode/air-cbrt.ll b/enzyme/test/Enzyme/ReverseMode/air-cbrt.ll new file mode 100644 index 00000000000..c56c52d4c6b --- /dev/null +++ b/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: } diff --git a/enzyme/test/Enzyme/ReverseMode/air-fast-sin.ll b/enzyme/test/Enzyme/ReverseMode/air-fast-sin.ll new file mode 100644 index 00000000000..d78f9093ad2 --- /dev/null +++ b/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: } diff --git a/enzyme/test/Enzyme/ReverseMode/air-pow.ll b/enzyme/test/Enzyme/ReverseMode/air-pow.ll new file mode 100644 index 00000000000..64d0b41f664 --- /dev/null +++ b/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: } diff --git a/enzyme/test/Enzyme/ReverseMode/air-sin.ll b/enzyme/test/Enzyme/ReverseMode/air-sin.ll new file mode 100644 index 00000000000..cac738d881b --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/air-sin.ll @@ -0,0 +1,33 @@ +; 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.sin.f32) resolved via isMemFreeLibMFunction's +; air. prefix stripping in TypeAnalysis.h to Intrinsic::sin, so its +; derivative reuses the same rule as llvm.sin.f32 (see sin.ll). + +; Function Attrs: nounwind readnone uwtable +define float @tester(float %x) { +entry: + %0 = tail call fast float @air.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.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: } diff --git a/enzyme/test/Enzyme/ReverseMode/air-tan.ll b/enzyme/test/Enzyme/ReverseMode/air-tan.ll new file mode 100644 index 00000000000..0229ecb63b0 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/air-tan.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 math intrinsic (air.tan.f32): tan is not a real LLVM intrinsic, +; so this exercises the new air.tan.f32/air.tan.f64 CallPattern entry in +; InstructionDerivatives.td rather than the intrinsic-redirect path. + +; Function Attrs: nounwind readnone uwtable +define float @tester(float %x) { +entry: + %0 = tail call fast float @air.tan.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.tan.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.tan.f32(float %x) +; CHECK-NEXT: %1 = fmul fast float %0, %0 +; CHECK-NEXT: %2 = fadd fast float 1.000000e+00, %1 +; CHECK-NEXT: %3 = fmul fast float %differeturn, %2 +; CHECK-NEXT: %4 = insertvalue { float } undef, float %3, 0 +; CHECK-NEXT: ret { float } %4 +; CHECK-NEXT: } diff --git a/enzyme/test/Enzyme/ReverseMode/air-tanh.ll b/enzyme/test/Enzyme/ReverseMode/air-tanh.ll new file mode 100644 index 00000000000..efe06966af5 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/air-tanh.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 hyperbolic intrinsic (air.tanh.f32): exercises the new +; air.tanh.f32 CallPattern entry, whose companion call must be the AIR name +; air.cosh.f32 (not "coshf"), since Metal's AIR compiler cannot resolve +; libm symbols. + +; Function Attrs: nounwind readnone uwtable +define float @tester(float %x) { +entry: + %0 = tail call fast float @air.tanh.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.tanh.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.cosh.f32(float %x) +; CHECK-NEXT: %1 = fmul fast float %0, %0 +; CHECK-NEXT: %2 = fdiv fast float %differeturn, %1 +; CHECK-NEXT: %3 = insertvalue { float } undef, float %2, 0 +; CHECK-NEXT: ret { float } %3 +; CHECK-NEXT: } From 13c5f0f06b83e998803888f345ec2df2bd0c25fa Mon Sep 17 00:00:00 2001 From: pw0908 Date: Sun, 12 Jul 2026 16:40:23 -0700 Subject: [PATCH 2/7] Fix air-atan2 test --- enzyme/test/Enzyme/ReverseMode/air-atan2.ll | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/enzyme/test/Enzyme/ReverseMode/air-atan2.ll b/enzyme/test/Enzyme/ReverseMode/air-atan2.ll index b5a31d63d96..8cd4f01fad1 100644 --- a/enzyme/test/Enzyme/ReverseMode/air-atan2.ll +++ b/enzyme/test/Enzyme/ReverseMode/air-atan2.ll @@ -1,4 +1,4 @@ -; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -mem2reg -sroa -early-cse -instsimplify -simplifycfg -adce -S | FileCheck %s; fi +; 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 From f5d05311b88939eefb47aa0290a0a736da051773 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Tue, 14 Jul 2026 00:51:01 -0700 Subject: [PATCH 3/7] Fix clang error --- enzyme/Enzyme/PreserveNVVM.cpp | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/enzyme/Enzyme/PreserveNVVM.cpp b/enzyme/Enzyme/PreserveNVVM.cpp index 5f9ecc39940..fcccbdf90a5 100644 --- a/enzyme/Enzyme/PreserveNVVM.cpp +++ b/enzyme/Enzyme/PreserveNVVM.cpp @@ -896,10 +896,9 @@ bool preserveNVVM(bool Begin, Module &M) { // 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"}) { + 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; @@ -914,9 +913,9 @@ bool preserveNVVM(bool Begin, Module &M) { // 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"}) { + 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; From 69c0607ad8c1920824cf82c6a70bc6ee25cc4228 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Tue, 14 Jul 2026 00:51:17 -0700 Subject: [PATCH 4/7] Add metal to list of uses --- enzyme/Enzyme/Enzyme.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/enzyme/Enzyme/Enzyme.cpp b/enzyme/Enzyme/Enzyme.cpp index 8c3ec03f2bc..f5afeb0568d 100644 --- a/enzyme/Enzyme/Enzyme.cpp +++ b/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( From fb3543ea6dc4508620e9861a4dd7c55dae337826 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Tue, 14 Jul 2026 15:21:20 -0700 Subject: [PATCH 5/7] Add utility function to identify GPU architecture --- enzyme/Enzyme/DiffeGradientUtils.cpp | 7 +++-- enzyme/Enzyme/Enzyme.cpp | 8 ++---- enzyme/Enzyme/EnzymeLogic.cpp | 13 +++++----- enzyme/Enzyme/GradientUtils.cpp | 38 ++++++++-------------------- enzyme/Enzyme/Utils.h | 17 +++++++++++++ 5 files changed, 39 insertions(+), 44 deletions(-) diff --git a/enzyme/Enzyme/DiffeGradientUtils.cpp b/enzyme/Enzyme/DiffeGradientUtils.cpp index 8c092bf201d..3c450615705 100644 --- a/enzyme/Enzyme/DiffeGradientUtils.cpp +++ b/enzyme/Enzyme/DiffeGradientUtils.cpp @@ -1092,13 +1092,12 @@ void DiffeGradientUtils::addToInvertedPtrDiffe(Instruction *orig, // atomics bool Atomic = AtomicAdd; - auto Arch = llvm::Triple(newFunc->getParent()->getTargetTriple()).getArch(); + auto TT = llvm::Triple(newFunc->getParent()->getTargetTriple()); + auto Arch = TT.getArch(); // No need to do atomic on local memory for CUDA since it can't be raced // upon - if (isa(TmpOrig) && - (Arch == Triple::nvptx || Arch == Triple::nvptx64 || - Arch == Triple::amd_target)) { + if (isa(TmpOrig) && isGPUArch(TT)) { Atomic = false; } // Moreover no need to do atomic on local shadows regardless since they are diff --git a/enzyme/Enzyme/Enzyme.cpp b/enzyme/Enzyme/Enzyme.cpp index 0646b838bbc..6ece2358b88 100644 --- a/enzyme/Enzyme/Enzyme.cpp +++ b/enzyme/Enzyme/Enzyme.cpp @@ -1458,9 +1458,7 @@ class EnzymeBase { auto subsequent_calls_may_write = options.subsequent_calls_may_write; auto TT = Triple(CI->getModule()->getTargetTriple()); - auto Arch = TT.getArch(); - bool AtomicAdd = Arch == Triple::nvptx || Arch == Triple::nvptx64 || - Arch == Triple::amd_target || TT.getArchName() == "air64"; + bool AtomicAdd = isGPUArch(TT); TypeAnalysis TA(Logic); FnTypeInfo type_args = populate_type_args(TA, fn, mode); @@ -2605,10 +2603,8 @@ class EnzymeBase { auto TT = llvm::Triple( CI->getParent()->getParent()->getParent()->getTargetTriple()); - auto Arch = TT.getArch(); - bool AtomicAdd = Arch == Triple::nvptx || Arch == Triple::nvptx64 || - Arch == Triple::amd_target || TT.getArchName() == "air64"; + bool AtomicAdd = isGPUArch(TT); IRBuilder<> Builder(CI); auto val = GradientUtils::GetOrCreateShadowConstant( diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp index cad2265325e..21b5b4e2831 100644 --- a/enzyme/Enzyme/EnzymeLogic.cpp +++ b/enzyme/Enzyme/EnzymeLogic.cpp @@ -4542,12 +4542,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; @@ -4556,6 +4553,10 @@ Function *EnzymeLogic::CreatePrimalAndGradient( IRBuilder<> entryBuilder(gutils->inversionAllocs, gutils->inversionAllocs->begin()); + // Note: intentionally not using isGPUArch(TT) here since the + // thread-id/barrier codegen below only knows how to handle + // NVPTX and AMDGPU (it is unreachable for any other arch, e.g. + // Metal air64). if ((Arch == Triple::nvptx || Arch == Triple::nvptx64 || Arch == Triple::amd_target) && g.getType()->getAddressSpace() == SharedAddrSpace) { diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp index f511f6852fd..0afbcc1586e 100644 --- a/enzyme/Enzyme/GradientUtils.cpp +++ b/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(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"); } @@ -5683,16 +5679,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(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 @@ -6743,12 +6733,8 @@ Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM, reduceRegister = true; } if (auto LI = dyn_cast(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(LI->getPointerOperand()->getType()) ->getAddressSpace() == SharedAddrSpace) { reduceRegister |= tryLegalRecomputeCheck && @@ -7098,12 +7084,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(orig_liobj->getType())->getAddressSpace() == SharedAddrSpace) { diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index 178ac314fab..aa770645155 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -68,6 +68,23 @@ #define amd_target amdgcn #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" From b601d8f21d9591f08ea684e2843b09fa5b8aee0b Mon Sep 17 00:00:00 2001 From: pw0908 Date: Thu, 16 Jul 2026 18:30:30 -0700 Subject: [PATCH 6/7] Fix make issue. --- enzyme/Enzyme/Utils.h | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index aa770645155..a15468bce9a 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -68,6 +68,12 @@ #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 From ed2faf07e68f8fcdf5b424b0e318313459655138 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Wed, 22 Jul 2026 14:24:34 -0700 Subject: [PATCH 7/7] Logic fix --- enzyme/Enzyme/EnzymeLogic.cpp | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp index 21b5b4e2831..cfa3d39f726 100644 --- a/enzyme/Enzyme/EnzymeLogic.cpp +++ b/enzyme/Enzyme/EnzymeLogic.cpp @@ -4553,12 +4553,7 @@ Function *EnzymeLogic::CreatePrimalAndGradient( IRBuilder<> entryBuilder(gutils->inversionAllocs, gutils->inversionAllocs->begin()); - // Note: intentionally not using isGPUArch(TT) here since the - // thread-id/barrier codegen below only knows how to handle - // NVPTX and AMDGPU (it is unreachable for any other arch, e.g. - // Metal air64). - 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",