Fortran function like

enzyme/Enzyme/PreserveNVVM.cpp

@@ -34,6 +34,7 @@
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -99,6 +100,40 @@ bool preserveLinkage(bool Begin, Function &F, bool Inlining = true) {
   return false;
 }
 
+static void handleFunctionLike(bool Begin, Value *Target,
+                               StringRef FunctionName) {
+  while (auto *CE = dyn_cast<ConstantExpr>(Target))
+    Target = CE->getOperand(0);
+
+  if (FunctionName.empty()) {
+    errs() << "Use of enzyme_function_like requires a non-empty function "
+              "name\n";
+    llvm_unreachable("enzyme_function_like");
+  }
+
+  auto *F = dyn_cast<Function>(Target);
+  if (!F) {
+    errs() << "First argument of enzyme_function_like must be a constant "
+              "function\n"
+           << *Target << "\n";
+    llvm_unreachable("enzyme_function_like");
+  }
+
+  // Warn on conflicting registrations while preserving the existing
+  // last-registration-wins behavior.
+  Attribute Existing = F->getFnAttribute("enzyme_math");
+  if (Existing.isValid() && Existing.getValueAsString() != FunctionName) {
+    errs() << "warning: conflicting enzyme_function_like registrations for "
+              "function '"
+           << F->getName() << "': replacing '" << Existing.getValueAsString()
+           << "' with '" << FunctionName << "'\n";
+  }
+
+  F->addAttribute(AttributeList::FunctionIndex,
+                  Attribute::get(F->getContext(), "enzyme_math", FunctionName));
+  preserveLinkage(Begin, *F);
+}
+
 // Return true if the module has a triple indicating an nvptx target, false
 // otherwise.
 bool isTargetNVPTX(llvm::Module &M) {
@@ -358,6 +393,59 @@ bool preserveNVVM(bool Begin, Module &M,
   constexpr static const char splitderivative_handler_name[] =
       "__enzyme_register_splitderivative";
 
+  // Flang cannot construct the constant function/string aggregate used by
+  // __enzyme_function_like. The Fortran binding instead passes a function and
+  // a BIND(C) global whose name is enzyme_math_<function>.
+  if (Begin) {
+    SmallVector<CallInst *, 4> functionLikeCalls;
+    for (Function &Caller : M) {
+      for (BasicBlock &BB : Caller) {
+        for (Instruction &I : BB) {
+          auto *Call = dyn_cast<CallInst>(&I);
+          if (!Call)
+            continue;
+
+          auto *Hook =
+              dyn_cast<Function>(Call->getCalledOperand()->stripPointerCasts());
+          if (!Hook || !startsWith(Hook->getName(), "f__enzyme_function_like"))
+            continue;
+
+          if (Call->arg_size() != 2) {
+            errs() << "Fortran enzyme_function_like requires exactly a "
+                      "function and a function name\n"
+                   << *Call << "\n";
+            llvm_unreachable("invalid Fortran enzyme_function_like call");
+          }
+
+          auto *NameGlobal = dyn_cast<GlobalVariable>(
+              Call->getArgOperand(1)->stripPointerCasts());
+          if (!NameGlobal) {
+            errs() << "Second argument of Fortran enzyme_function_like must "
+                      "be an enzyme_math_* function name\n"
+                   << *Call->getArgOperand(1) << "\n";
+            llvm_unreachable(
+                "invalid Fortran enzyme_function_like function name");
+          }
+
+          StringRef FunctionName = NameGlobal->getName();
+          if (!FunctionName.consume_front("enzyme_math_")) {
+            errs() << "Fortran enzyme_function_like function name must use "
+                      "the enzyme_math_* BIND(C) naming convention\n"
+                   << *NameGlobal << "\n";
+            llvm_unreachable(
+                "invalid Fortran enzyme_function_like function name");
+          }
+
+          handleFunctionLike(Begin, Call->getArgOperand(0), FunctionName);
+          functionLikeCalls.push_back(Call);
+          changed = true;
+        }
+      }
+    }
+    for (CallInst *Call : functionLikeCalls)
+      Call->eraseFromParent();
+  }
+
   if (Begin)
     if (GlobalVariable *GA = M.getGlobalVariable("llvm.global.annotations")) {
       if (GA->hasInitializer()) {
@@ -466,11 +554,8 @@ bool preserveNVVM(bool Begin, Module &M,
 
             if (startsWith(AS, "enzyme_function_like") && Func) {
               auto val = AS.substr(1 + AS.find('='));
-              Func->addAttribute(
-                  AttributeList::FunctionIndex,
-                  Attribute::get(Func->getContext(), "enzyme_math", val));
+              handleFunctionLike(Begin, Func, val);
               changed = true;
-              preserveLinkage(Begin, *Func);
               replacements.push_back(Constant::getNullValue(CAOp->getType()));
               continue;
             }
@@ -668,7 +753,22 @@ bool preserveNVVM(bool Begin, Module &M,
     if (g.getName().contains("__enzyme_function_like")) {
       if (g.hasInitializer()) {
         auto CA = dyn_cast<ConstantAggregate>(g.getInitializer());
-        if (!CA || CA->getNumOperands() < 2) {
+        if (!CA) {
+          constexpr StringLiteral Marker = "__enzyme_function_like__";
+          auto MarkerPos = g.getName().rfind(Marker);
+          Value *Target = g.getInitializer()->stripPointerCasts();
+
+          // Ignore globals that are not Fortran function-like registrations.
+          if (MarkerPos == StringRef::npos || !isa<Function>(Target))
+            continue;
+
+          handleFunctionLike(Begin, Target,
+                             g.getName().substr(MarkerPos + Marker.size()));
+          toErase.push_back(&g);
+          changed = true;
+          continue;
+        }
+        if (CA->getNumOperands() < 2) {
           llvm::errs() << "Use of "
                        << "enzyme_function_like"
                        << " must be a "
@@ -678,9 +778,6 @@ bool preserveNVVM(bool Begin, Module &M,
         }
         Value *V = CA->getOperand(0);
         Value *name = CA->getOperand(1);
-        while (auto CE = dyn_cast<ConstantExpr>(V)) {
-          V = CE->getOperand(0);
-        }
         while (auto CE = dyn_cast<ConstantExpr>(name)) {
           name = CE->getOperand(0);
         }
@@ -693,27 +790,9 @@ bool preserveNVVM(bool Begin, Module &M,
                     CA->isCString())
                   nameVal = CA->getAsCString();
 
-        if (nameVal == "") {
-          llvm::errs() << *name << "\n";
-          llvm::errs() << "Use of "
-                       << "enzyme_function_like"
-                       << "requires a non-empty function name"
-                       << "\n";
-          llvm_unreachable("enzyme_function_like");
-        }
-        if (auto F = cast<Function>(V)) {
-          F->addAttribute(
-              AttributeList::FunctionIndex,
-              Attribute::get(g.getContext(), "enzyme_math", nameVal));
-          toErase.push_back(&g);
-          changed = true;
-        } else {
-          llvm::errs() << "Param of __enzyme_function_like must be a "
-                          "constant function"
-                       << g << "\n"
-                       << *V << "\n";
-          llvm_unreachable("__enzyme_function_like");
-        }
+        handleFunctionLike(Begin, V, nameVal);
+        toErase.push_back(&g);
+        changed = true;
       }
     }
     if (g.getName().contains("__enzyme_allocation_like")) {

enzyme/Fortran/README.md

@@ -103,7 +103,7 @@ example, if you have a subroutine
 then you can make use of activity descriptors like so:
 ```fortran
   call enzyme_autodiff(my_subroutine, enzyme_const, n, &
-                       enzyme_dup, x, dx, enzyme_dup, y, dy
+                       enzyme_dup, x, dx, enzyme_dup, y, dy)
 ```
 
 ## Function hook for batching
@@ -124,3 +124,189 @@ for an example.
 > [!NOTE]
 > You will likely find that batching works more straightforwardly with
 > subroutines than with Fortran functions.
+
+
+## Function-like hooks
+
+The `enzyme_function_like` hook tells Enzyme to differentiate a function as if
+it were a known mathematical function. For example, Enzyme can use the
+derivative of `log1p` for `double_value`, regardless of its
+implementation. The examples below deliberately compute `2*x` while requesting
+the derivative of `log1p`: at `x = 2`, Enzyme returns `1/3` instead of `2`. This
+illustrates a derivative override; the two functions are not mathematically
+equivalent.
+
+### Choose a registration form
+
+| Function location and interface | When to use each form |
+|---|---|
+| Module function | Use a pointer declaration before the module's `contains` to keep registration with the function. The compiler supplies its explicit interface. A registration call in executable code also works. |
+| Internal function | Use call registration only if Flang supplies a direct function reference. Access to variables from the containing program or procedure can prevent registration. See the restriction below. The current pointer mechanism cannot register an internal function. |
+
+### Registration with a subroutine call
+
+Call `enzyme_function_like` as a subroutine with the target function and the
+symbolic name of the mathematical function:
+
+```fortran
+module enzyme_math_names
+  use, intrinsic :: iso_c_binding, only: c_int
+  implicit none
+  private
+  integer(c_int), public, bind(C, name="enzyme_math_log1p") :: enzyme_log1p
+end module enzyme_math_names
+```
+
+Import the binding in the program or procedure that registers the target:
+
+```fortran
+use enzyme, only: enzyme_function_like
+use enzyme_math_names, only: enzyme_log1p
+
+! Put this call after all declarations.
+call enzyme_function_like(double_value, enzyme_log1p)
+```
+
+Put the registration call in executable code, after declarations. For an
+internal function, you must use this form instead of an initialized procedure
+pointer. An internal function follows `contains` inside a program or another
+procedure. The compiler supplies its explicit interface.
+
+The [call-style test](../test/Fortran/ReverseMode/function_like.f90) shows this
+placement. Its registration call is in the main program. Its target function,
+`double_value`, is inside that program, after `contains`.
+
+> [!WARNING]
+> Call registration does not support all internal functions. An internal
+> function can access variables from its containing program or procedure.
+> Fortran calls this access **host association**. For example, `double_value`
+> could calculate `factor * x`, where `factor` is a variable in the containing
+> procedure.
+>
+> Flang can then generate an adapter that gives the function access to those
+> variables. The current registration code requires a direct function reference.
+> It cannot process this adapter, and compilation can fail with
+> `First argument of enzyme_function_like must be a constant function`.
+>
+> The example above uses only the argument `x` and does not need this adapter.
+
+
+Here `enzyme_log1p` supplies the symbolic function name `log1p`; its value is not
+used. Functions passed to `enzyme_function_like` must have an LLVM-level
+signature compatible with the selected mathematical function. Scalar arguments
+must use the `value` attribute so that Flang lowers them as LLVM values rather
+than using Fortran's usual by-reference calling convention. This binding is
+currently supported with Flang.
+
+When running Enzyme separately with `opt`, `preserve-nvvm` must process the
+`enzyme_function_like` hook before differentiation:
+
+```console
+$ opt -load-pass-plugin=/path/to/LLVMEnzyme-21.so \
+    -passes='preserve-nvvm,enzyme,preserve-nvvm-end' input.bc -o output.bc
+```
+
+> [!WARNING]
+> When using this separate `opt` workflow, compile the Fortran source to LLVM
+> with `-O0`. Otherwise, Flang may inline calls to the function before
+> `preserve-nvvm` processes the `enzyme_function_like` hook.
+
+The `FlangEnzyme`
+compiler plugin runs `preserve-nvvm` at the start of Flang's LLVM optimization
+pipeline and does not require this separate `opt` step.
+
+Use this hook to assign a mathematical rule to a custom function.
+
+For call registration, declare the required symbolic names in a user module.
+Use `integer(c_int)` with `bind(C, name="enzyme_math_<function>")`.
+Replace `<function>` with a mathematical rule name that Enzyme supports.
+The variable value is not used. The binding name selects the rule.
+
+For example, declare a binding for the `sin` rule:
+
+```fortran
+module enzyme_math_names
+  use, intrinsic :: iso_c_binding, only: c_int
+  implicit none
+  private
+
+  integer(c_int), public, bind(C, name="enzyme_math_sin") :: enzyme_sin
+end module enzyme_math_names
+```
+
+Import this binding where the registration call occurs:
+
+```fortran
+use enzyme, only: enzyme_function_like
+use enzyme_math_names, only: enzyme_sin
+
+! Put this call after all declarations.
+call enzyme_function_like(function_similar_to_sin, enzyme_sin)
+```
+
+### Procedure-pointer registration
+
+Alternatively, a statically initialized procedure pointer can register the
+same relationship without a hook call or symbolic-name binding. Enzyme reads
+and removes the registration marker at compile time. Do not call through the
+registration pointer. Enzyme replaces remaining references to the marker with
+null pointers. Call the target function directly, for example, `double_value(x)`.
+Use the same FlangEnzyme plugin or separate `opt` pipeline described above for
+call-style registration.
+
+#### Register a module function
+
+Put the pointer declaration in the declaration section of a module, program,
+function, or subroutine where the module function is accessible.
+Put it before executable statements or `contains`.
+Omit `private` when the declaration is outside a module's declaration section.
+The example below puts the declaration before the module's `contains` statement.
+
+```fortran
+module function_like_example
+  implicit none
+
+  procedure(double_value), pointer, private :: &
+    fn__enzyme_function_like__log1p => double_value
+
+contains
+
+  function double_value(x) result(y)
+    real, value :: x
+    real :: y
+
+    y = 2.0 * x
+  end function double_value
+
+  function test(x) result(y)
+    real, intent(in) :: x
+    real :: y
+
+    y = double_value(x)
+  end function test
+
+end module function_like_example
+
+program main
+  use enzyme, only: enzyme_autodiff
+  use function_like_example, only: test
+  implicit none
+  real :: x, dx
+
+  x = 2.0
+  dx = 0.0
+  call enzyme_autodiff(test, x, dx)
+  write(*,"(f6.4)") dx ! Prints 0.3333
+end program main
+```
+
+Here, `procedure(double_value)` gives the pointer the target's interface,
+and `=> double_value` initializes it with the target. PreserveNVVM reads
+the mathematical name after the exact `__enzyme_function_like__` delimiter, so
+this example registers the target as `log1p`. The prefix before the delimiter
+can be any valid name but must be unique in its scope. `private` is optional
+in a module; it keeps the registration marker out of the module's public API.
+
+The `test` wrapper takes its argument by reference for the `enzyme_autodiff`
+binding, while `double_value` takes its argument by value to match the scalar
+`log1p` rule.