Constant `int` return type is defaulted to `DUP_ARG` for custom rules arriving via bitcode

[ymardoukhi]

Assuming issue #3269 is patched, for a custom function with custom rules arriving via bitcode with constant integer return, Enzyme always assume DUP_ARG is the return activity.

How to reproduce

* kernel.c          # kernel wrapper + custom gradient rule + __enzyme_register_gradient_*
* build.sh          # produces kernel.so 
* mwe.jl            # Julia driver: shows the bug; correct with patched Enzyme

.c files for the opaque function and the Enzyme custom rules

// kernel.c
#include <stdlib.h>

int my_square(double *x, double *out) {
    *out = (*x) * (*x);
    return 0;
    /* returns an unconditional zero to allow
       ActivityAnalysis to resolve the activity
       as CONSTANT */
}

void *augmented_my_square(double *x, double *d_x,
                           double *out, double *d_out) {
    (void)d_x; (void)d_out;
    *out = (*x) * (*x);
    double *tape = (double *)malloc(sizeof(double));
    *tape = *x;
    return (void *)tape;
}

void gradient_my_square(double *x, double *d_x,
                         double *out, double *d_out,
                         void *tape_raw) {
    (void)x; (void)out;
    double x_val = *(double *)tape_raw;
    *d_x  += 2.0 * x_val * (*d_out);
    *d_out = 0.0;
    free(tape_raw);
}

__attribute__((used))
void *__enzyme_register_gradient_my_square[3] = {
    (void *)my_square,
    (void *)augmented_my_square,
    (void *)gradient_my_square,
};

Run build.sh to produce the .so files

# bash.sh

#!/bin/sh
# Build the shared object:
# kernel.so   `my_square` + Enzyme custom-gradient rule, with embedded bitcode

set -eu
cd "$(dirname "$0")"

clang-16 -O2 -Xclang -no-opaque-pointers \
    -fno-vectorize -fno-slp-vectorize -fno-unroll-loops \
    -shared -fPIC -fembed-bitcode \
    -o kernel.so kernel.c
echo "built kernel.so"

echo "=== kernel.so: my_square DEFINED (T) this time ==="
nm -D kernel.so | grep -E "my_square|augmented_my_square|gradient_my_square" || true

Run via Enzyme.jl to reproduce the error

# mwe.jl

using Enzyme
import Libdl

const dir = @__DIR__
const klib = Libdl.dlopen(joinpath(dir, "kernel.so"))
const my_square_fn = Libdl.dlsym(klib, :my_square)

function loss(x::Vector{Float64})
    out = [0.0]
    ccall(my_square_fn, Cint, (Ptr{Float64}, Ptr{Float64}), x, out)
    return out[1]
end

x = [2.0]
@show loss(x)  # primal: 4.0

g = Enzyme.gradient(Reverse, loss, x)

Last line produces

ERROR: LoadError: EnzymeNoDerivativeError: Current scope: 
; Function Attrs: mustprogress nofree nounwind willreturn
define "enzyme_type"="{[-1]:Float@double}" double @preprocess_julia_loss_1405({} addrspace(10)* nocapture nofree noundef nonnull readonly align 8 dereferenceable(24) "enzyme_type"="{[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Pointer, [-1,8,0]:Integer, [-1,8,1]:Integer, [-1,8,2]:Integer, [-1,8,3]:Integer, [-1,8,4]:Integer, [-1,8,5]:Integer, [-1,8,6]:Integer, [-1,8,7]:Integer, [-1,8,8]:Pointer, [-1,8,8,-1]:Float@double, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer}" "enzymejl_parmtype"="140499341852064" "enzymejl_parmtype_ref"="2" "enzymejl_parmtype_str"="Vector{Float64}" %0) local_unnamed_addr #11 !dbg !81 {
top:
  %pgcstack = call {}*** @julia.get_pgcstack() #13
  %ptls_field4 = getelementptr inbounds {}**, {}*** %pgcstack, i64 2
  %1 = bitcast {}*** %ptls_field4 to i64***
  %ptls_load56 = load i64**, i64*** %1, align 8, !tbaa !13
  %2 = getelementptr inbounds i64*, i64** %ptls_load56, i64 2
  %safepoint = load i64*, i64** %2, align 8, !tbaa !17
  fence syncscope("singlethread") seq_cst
  call void @julia.safepoint(i64* %safepoint) #14, !dbg !82
  fence syncscope("singlethread") seq_cst
  %3 = call noalias "enzyme_ReadOnlyOrThrow" "enzyme_type"="{[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer, [-1,4]:Integer, [-1,5]:Integer, [-1,6]:Integer, [-1,7]:Integer, [-1,8]:Pointer, [-1,8,-1]:Float@double}" {} addrspace(10)* @jl_alloc_genericmemory({} addrspace(10)* noundef @"ejl_inserted$_Core_GenericMemory_1408$false$140499341852256", i64 noundef 1) #15, !dbg !83
  %4 = bitcast {} addrspace(10)* %3 to { i64, {} addrspace(10)** } addrspace(10)*, !dbg !87
  %5 = addrspacecast { i64, {} addrspace(10)** } addrspace(10)* %4 to { i64, {} addrspace(10)** } addrspace(11)*, !dbg !87
  %6 = getelementptr inbounds { i64, {} addrspace(10)** }, { i64, {} addrspace(10)** } addrspace(11)* %5, i64 0, i32 1, !dbg !87
  %7 = bitcast {} addrspace(10)** addrspace(11)* %6 to i8* addrspace(11)*, !dbg !87
  %8 = load i8*, i8* addrspace(11)* %7, align 8, !dbg !87, !tbaa !32, !alias.scope !33, !noalias !37, !nonnull !0, !enzyme_type !41, !enzymejl_source_type_Ptr\7BFloat64\7D !0, !enzymejl_byref_BITS_VALUE !0, !enzyme_nocache !0
  %9 = bitcast i8* %8 to {} addrspace(10)**, !dbg !89
  %10 = call "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" {} addrspace(10)* addrspace(13)* @julia.gc_loaded({} addrspace(10)* noundef readnone %3, {} addrspace(10)** noundef nonnull readnone %9) #13, !dbg !89
  %11 = bitcast {} addrspace(10)* addrspace(13)* %10 to double addrspace(13)*, !dbg !89
  store double 0.000000e+00, double addrspace(13)* %11, align 8, !dbg !89, !tbaa !47, !alias.scope !50, !noalias !91
  %12 = bitcast {} addrspace(10)* %0 to { i8*, {} addrspace(10)* } addrspace(10)*, !dbg !94
  %13 = addrspacecast { i8*, {} addrspace(10)* } addrspace(10)* %12 to { i8*, {} addrspace(10)* } addrspace(11)*, !dbg !94
  %14 = bitcast {} addrspace(10)* %0 to i8* addrspace(10)*, !dbg !94
  %15 = addrspacecast i8* addrspace(10)* %14 to i8* addrspace(11)*, !dbg !94
  %16 = load i8*, i8* addrspace(11)* %15, align 8, !dbg !94, !tbaa !58, !alias.scope !61, !noalias !62, !enzyme_type !41, !enzymejl_source_type_Ptr\7BFloat64\7D !0, !enzymejl_byref_BITS_VALUE !0
  %17 = getelementptr inbounds { i8*, {} addrspace(10)* }, { i8*, {} addrspace(10)* } addrspace(11)* %13, i64 0, i32 1, !dbg !94
  %18 = load {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %17, align 8, !dbg !94, !tbaa !58, !alias.scope !61, !noalias !62, !dereferenceable_or_null !63, !align !64, !enzyme_type !65, !enzymejl_source_type_Memory\7BFloat64\7D !0, !enzymejl_byref_MUT_REF !0
  %19 = insertvalue { i8*, {} addrspace(10)* } zeroinitializer, i8* %16, 0, !dbg !94
  %20 = insertvalue { i8*, {} addrspace(10)* } %19, {} addrspace(10)* %18, 1, !dbg !94
  %21 = insertvalue { i8*, {} addrspace(10)* } zeroinitializer, i8* %8, 0, !dbg !94
  %22 = insertvalue { i8*, {} addrspace(10)* } %21, {} addrspace(10)* %3, 1, !dbg !94
  %23 = bitcast i8* %16 to double*, !dbg !95
  %24 = bitcast i8* %8 to double*, !dbg !95
  %25 = call i32 @my_square(double* nocapture nofree noundef readonly align 8 %23, double* nocapture nofree noundef nonnull writeonly align 8 %24) #14 [ "jl_roots"({ i8*, {} addrspace(10)* } %22, { i8*, {} addrspace(10)* } %20) ], !dbg !95
  %26 = load double, double addrspace(13)* %11, align 8, !dbg !96, !tbaa !47, !alias.scope !50, !noalias !72, !enzyme_type !73, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Float64 !0
  ret double %26, !dbg !97
}

The required return activity calling into function: my_square was CONSTANT but the assumed (default) return activity was DUP_ARG
 at context:   %25 = call i32 @my_square(double* nocapture nofree noundef readonly align 8 %23, double* nocapture nofree noundef nonnull writeonly align 8 %24) #14 [ "jl_roots"({ i8*, {} addrspace(10)* } %22, { i8*, {} addrspace(10)* } %20) ], !dbg !57

Stacktrace:
 [1] loss
   @ /workspace/return_activity_i32/mwe.jl:55

Failure within method:

loss(::Vector{Float64})
     @ Main /workspace/return_activity_i32/mwe.jl:53

Hint: catch this exception as `err` and call `code_typed(err)` to inspect the surrounding code.

Stacktrace:
  [1] loss
    @ /workspace/return_activity_i32/mwe.jl:55 [inlined]
  [2] diffejulia_loss_1405wrap
    @ /workspace/return_activity_i32/mwe.jl:0
  [3] macro expansion
    @ ~/.julia/dev/Enzyme/src/compiler.jl:6909 [inlined]
  [4] enzyme_call
    @ ~/.julia/dev/Enzyme/src/compiler.jl:6379 [inlined]
  [5] CombinedAdjointThunk
    @ ~/.julia/dev/Enzyme/src/compiler.jl:6263 [inlined]
  [6] autodiff
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:536 [inlined]
  [7] autodiff
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:557 [inlined]
  [8] macro expansion
    @ ~/.julia/dev/Enzyme/src/sugar.jl:337 [inlined]
  [9] gradient(::ReverseMode{false, false, false, FFIABI, false, false}, ::typeof(loss), ::Vector{Float64})
    @ Enzyme ~/.julia/dev/Enzyme/src/sugar.jl:274
 [10] top-level scope
    @ /workspace/return_activity_i32/mwe.jl:62
in expression starting at /workspace/return_activity_i32/mwe.jl:62

The root cause originates apparently from EnzymeLogic.cpp

# snippet
    DIFFE_TYPE subretType = whatType(key.todiff->getReturnType(),
                                     DerivativeMode::ReverseModeGradient,
                                     /*intAreConstant*/ false, seen);  # --> the hardcoded `false` for `integersAreConstant` argument
    if (key.todiff->getReturnType()->isVoidTy() ||
        key.todiff->getReturnType()->isEmptyTy())
      subretType = DIFFE_TYPE::CONSTANT;

    if (subretType == DIFFE_TYPE::OUT_DIFF &&
        key.retType == DIFFE_TYPE::CONSTANT) {
      hasconstant = true;
    } else if (subretType != key.retType) {
      std::string str;
      raw_string_ostream ss(str);
      ss << "The required return activity calling into function: "
         << key.todiff->getName() << " was " << to_string(key.retType)
         << " but the assumed (default) return activity was "
         << to_string(subretType) << "\n";
      if (context.req) {
        ss << " at context: " << *context.req;
      }
      if (EmitNoDerivativeError(ss.str(), key.todiff, context)) {
        return nullptr;
      }
    }

Tracing back to the definition of whatType

# snippet
  } else if (arg->isIntOrIntVectorTy() || arg->isFunctionTy()) {
    return integersAreConstant ? DIFFE_TYPE::CONSTANT : DIFFE_TYPE::DUP_ARG;

In EnzymeLogic.cpp the argument for integersAreConstant is hardcoded to false hence the activity for the return is always defaulted to DUP_ARG.