EnzymeAD / EnzymeAD/Enzyme.jl

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

Open
#3,300 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

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

```cpp
// kernel.c
#include

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
# 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
```julia
# 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

```julia
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`](https://github.com/EnzymeAD/Enzyme/blob/d7dd2e064b27511f10869d57201929b976756c9a/enzyme/Enzyme/EnzymeLogic.cpp#L3739-L3762)

```julia
# 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`](https://github.com/EnzymeAD/Enzyme/blob/d7dd2e064b27511f10869d57201929b976756c9a/enzyme/Enzyme/Utils.h#L609-L610)

```julia
# 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`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.