Enzyme could not find world for custom rules arriving via bitcode

[ymardoukhi]

I am trying to define Enzyme custom rules for an external opaque function in an LLVM IR. But it seems __enzyme_register_gradient_* globals arriving via bitcode embedded in .so files are invisible to the pass and are never registered for opaque functions.

How to reproduce:

* kernel.c          # kernel wrapper + custom gradient rule + __enzyme_register_gradient_*
* my_square_impl.c  # runtime body of my_square (compiled to libmy_square.so)
* build.sh          # produces libmy_square.so and kernel.so 
* mwe.jl            # Julia driver: shows the bug; correct with patched Enzyme

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

// my_square_impl.c
void my_square(double *x, double *out) { *out = (*x) * (*x); }

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

// Runtime symbol — body lives in libmy_square.so, opaque to Enzyme from the bitcode.
extern void my_square(double *x, double *out);

// The function Enzyme differentiates through
void kernel(double *x, double *out) {
    my_square(x, out);
}

// Custom gradient rule for my_square

void *augmented_my_square(double *x, double *d_x,
                           double *out, double *d_out) {
    (void)d_x; (void)d_out;
    my_square(x, out);
    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);
}

/* Registration global: consumed by PreserveNVVMPass at differentiation time.
   Lives in the embedded .llvmbc section of kernel.so
   The NVVM pass ordering bug makes it is invisible to PreserveNVVMPass
   when that pass runs before check_ir imports the bitcode. */
__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

# build.sh

#!/bin/sh
# Build two shared objects:
#   libmy_square.so  -- runtime body of my_square
#   kernel.so        -- kernel + Enzyme custom-gradient rule, with embedded bitcode

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

clang-16 -O2 -shared -fPIC -o libmy_square.so my_square_impl.c
echo "built libmy_square.so"

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 UNDEFINED (U), kernel DEFINED (T) ==="
nm -D kernel.so | grep -E "my_square|kernel" || true
echo "=== libmy_square.so: my_square DEFINED (T) ==="
nm -D libmy_square.so | grep my_square || true

Run via Enzyme.jl to reproduce the error

# mwe.jl

using Enzyme
import Libdl

Enzyme.API.printall!(false)

const dir = @__DIR__
# Register `my_square` in `RTLD_GLOBAL`
Libdl.dlopen(joinpath(dir, "libmy_square.so"), Libdl.RTLD_GLOBAL)
const klib = Libdl.dlopen(joinpath(dir, "kernel.so"))
const kernel_fn = Libdl.dlsym(klib, :kernel)

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

x = [2.0]
g = Enzyme.gradient(Reverse, loss, x) 

Last line produces

ERROR: LoadError: AssertionError: Enzyme: could not find world in ; Function Attrs: mustprogress nounwind willreturn uwtable
define internal fastcc void @preprocess_kernel(double* noundef %0, double* noundef nonnull %1) unnamed_addr #13 {
  tail call void @my_square(double* noundef %0, double* noundef nonnull %1) #14
  ret void
}

Stacktrace:
  [1] enzyme_extract_world(fn::LLVM.Function)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:2494
  [2] julia_error(msg::String, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing}, data2::Ptr{LLVM.API.LLVMOpaqueValue}, B::Ptr{LLVM.API.LLVMOpaqueBuilder})
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/errors.jl:1116
  [3] julia_error(cstr::Cstring, val::Ptr{LLVM.API.LLVMOpaqueValue}, errtype::Enzyme.API.ErrorType, data::Ptr{Nothing}, data2::Ptr{LLVM.API.LLVMOpaqueValue}, B::Ptr{LLVM.API.LLVMOpaqueBuilder})
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/errors.jl:1024
  [4] EnzymeCreatePrimalAndGradient(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnValue::Bool, dretUsed::Bool, mode::Enzyme.API.CDerivativeMode, runtimeActivity::Bool, strongZero::Bool, width::Int64, additionalArg::Ptr{…}, forceAnonymousTape::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, augmented::Ptr{…}, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/dev/Enzyme/src/api.jl:270
  [5] macro expansion
    @ ~/.julia/dev/Enzyme/src/compiler.jl:2840 [inlined]
  [6] macro expansion
    @ ~/.julia/packages/LLVM/aXQtR/src/base.jl:113 [inlined]
  [7] enzyme!(enzyme_context::Enzyme.EnzymeContext, job::GPUCompiler.CompilerJob{…}, interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…}, removedRoots::Set{…})
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:2711
  [8] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}})
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:6027
  [9] #compile#202
    @ ~/.julia/packages/GPUCompiler/QgckP/src/driver.jl:67 [inlined]
 [10] compile
    @ ~/.julia/packages/GPUCompiler/QgckP/src/driver.jl:55 [inlined]
 [11] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:6970
 [12] _thunk
    @ ~/.julia/dev/Enzyme/src/compiler.jl:6968 [inlined]
 [13] cached_compilation
    @ ~/.julia/dev/Enzyme/src/compiler.jl:7038 [inlined]
 [14] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:7154
 [15] thunk_generator(world::UInt64, source::Union{…}, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type, strongzero::Type)
    @ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:7298
 [16] autodiff
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:517 [inlined]
 [17] autodiff
    @ ~/.julia/dev/Enzyme/src/Enzyme.jl:557 [inlined]
 [18] macro expansion
    @ ~/.julia/dev/Enzyme/src/sugar.jl:337 [inlined]
 [19] gradient(::ReverseMode{false, false, false, FFIABI, false, false}, ::typeof(loss), ::Vector{Float64})
    @ Enzyme ~/.julia/dev/Enzyme/src/sugar.jl:274
 [20] top-level scope
    @ /workspace/nvvmpass_mwe/mwe.jl:45
 [21] include(fname::String)
    @ Main ./sysimg.jl:38
 [22] top-level scope
    @ REPL[2]:1
in expression starting at /workspace/nvvmpass_mwe/mwe.jl:45
Some type information was truncated. Use `show(err)` to see complete types.

Proposed fix

Add a second PreserveNVVMPass run immediately after check_ir() in src/compiler.jl. Allows any registration globals arriving via imported bitcode to be consumed before the main Enzyme differentiation pass.

Enzyme.jl/src/compiler.jl

Lines 5392 to 5395 in 6d69320

	check_ir(interp, job, mod)
	if DumpPostCheck[]
	API.EnzymeDumpModuleRef(mod.ref)
	end

     check_ir(interp, job, mod)
     if DumpPostCheck[]
         API.EnzymeDumpModuleRef(mod.ref)
     end

+    # Re-run PreserveNVVM AFTER check_ir: foreign bitcode imported via
+    # try_import_llvmbc (inside check_ir) may carry __enzyme_register_gradient_*
+    # globals (custom rules for opaque symbols). The PreserveNVVM that runs before
+    # check_ir cannot see them
+    LLVM.@dispose pb=LLVM.NewPMPassBuilder() begin
+        registerEnzymeAndPassPipeline!(pb)
+        LLVM.add!(pb, LLVM.NewPMModulePassManager()) do mpm
+            LLVM.add!(mpm, PreserveNVVMPass())
+        end
+        LLVM.run!(pb, mod)
+    end
+
     disableFallback = String[]

Related Julia discourse thread

Julia Version 1.11.6
Commit 9615af0f269 (2025-07-09 12:58 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 12 × AMD Ryzen 5 7640U w/ Radeon 760M Graphics
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver4)
Threads: 1 default, 0 interactive, 1 GC (on 12 virtual cores)
Environment:
  JULIA_CONDAPKG_BACKEND = Null
  JULIA_PYTHONCALL_EXE = /home/devuser/.venv/bin/python

(@dev-utils) pkg> st
Status `~/.julia/environments/dev-utils/Project.toml`
  [0ca39b1e] Chairmarks v1.3.1
  [7da242da] Enzyme v0.13.173
  [98e50ef6] JuliaFormatter v2.10.1
  [295af30f] Revise v3.15.1
  [62bfec6d] Runic v1.7.0

[vchuravy]

I am trying to improve that situation in #3156 but there are still some callers that want to look at the Julia specific metadata.