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