Fortran function like [WIP]

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"
 
@@ -81,6 +82,30 @@ 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");
+  }
+
+  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) {
@@ -333,6 +358,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 || !Hook->getName().contains("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()) {
@@ -425,11 +503,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;
             }
@@ -631,9 +706,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);
         }
@@ -646,27 +718,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/enzyme.f90

@@ -23,7 +23,9 @@
 module enzyme
   use iso_c_binding, only: c_int
   use enzyme_function_hooks, only: enzyme_autodiff => f__enzyme_autodiff, &
-                                   enzyme_fwddiff  => f__enzyme_fwddiff
+                                   enzyme_fwddiff  => f__enzyme_fwddiff, &
+                                   enzyme_function_like => &
+                                     f__enzyme_function_like
   implicit none
   private
 
@@ -36,7 +38,11 @@ module enzyme
   integer(c_int), public, bind(C, name="enzyme_width")     :: enzyme_width
   integer(c_int), public, bind(C, name="enzyme_vector")    :: enzyme_vector
 
+  ! Symbolic function names for enzyme_function_like
+  integer(c_int), public, bind(C, name="enzyme_math_log1p") :: enzyme_log1p
+
   ! Bindings for function hooks
   public :: enzyme_autodiff
   public :: enzyme_fwddiff
+  public :: enzyme_function_like
 end module enzyme

enzyme/Fortran/enzyme_function_hooks.f90

@@ -34,6 +34,8 @@ module enzyme_function_hooks
   !       prepend with 'f' in the Fortran versions of the function hooks.
   public :: f__enzyme_autodiff
   public :: f__enzyme_fwddiff
+  public :: f__enzyme_function_like
   external :: f__enzyme_autodiff
   external :: f__enzyme_fwddiff
+  external :: f__enzyme_function_like
 end module enzyme_function_hooks