Enzyme failed to complie for the custom gradient for C++ on cuda

[minansys]

Enzyme failed to complie for the custom gradient for C++ as show in this example for the atomicMax

https://fwd.gymni.ch/hc3ivR

which crashed at

if (hasMetadata(key.todiff, "enzyme_gradient")) {
std::set<llvm::Type *> seen;
DIFFE_TYPE subretType = whatType(key.todiff->getReturnType(),
DerivativeMode::ReverseModeGradient,
/intAreConstant/ false, seen);
if (key.todiff->getReturnType()->isVoidTy() ||
key.todiff->getReturnType()->isEmptyTy())
subretType = DIFFE_TYPE::CONSTANT;

@wsmoses will have the fix for this and provide a better msg.

[minansys]

&"warning: GDB: Failed to set controlling terminal: Operation not permitted\n"
2: File exists.
clang++: warning: CUDA version is newer than the latest supported version 12.3 [-Wunknown-cuda-version]
 todiff: ; Function Attrs: convergent mustprogress noinline nounwind optnone
define dso_local noundef double @_Z9AtomicMaxIdET_RS0_RKS0_(ptr noundef nonnull align 8 dereferenceable(8) %dest, ptr noundef nonnull align 8 dereferenceable(8) %val) #5 !dbg !2120 !enzyme_augment !2127 !enzyme_gradient !2128 {
entry:
  %dest.addr = alloca ptr, align 8
  %val.addr = alloca ptr, align 8
  store ptr %dest, ptr %dest.addr, align 8
  call void @llvm.dbg.declare(metadata ptr %dest.addr, metadata !2129, metadata !DIExpression()), !dbg !2130
  store ptr %val, ptr %val.addr, align 8
  call void @llvm.dbg.declare(metadata ptr %val.addr, metadata !2131, metadata !DIExpression()), !dbg !2132
  %0 = load ptr, ptr %dest.addr, align 8, !dbg !2133
  %1 = load ptr, ptr %val.addr, align 8, !dbg !2134
  %2 = load double, ptr %1, align 8, !dbg !2134
  %call = call contract noundef double @_Z15atomicMaxDoublePdd(ptr noundef %0, double noundef %2) #139, !dbg !2135
  ret double %call, !dbg !2136
}

 revfn: define internal void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_340(ptr %arg0, ptr %"arg0'", ptr %arg1, ptr %"arg1'") {
entry:
  call void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_339(ptr %arg0, ptr %"arg0'", ptr %arg1, ptr %"arg1'", double 0.000000e+00)
  ret void
}

 NewF: define internal void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3) {
entry:
  %0 = call ptr @_Z17augment_AtomicMaxIdEPvPT_S2_PKS1_S2_(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3)
}

 key rettype: CONSTANT
 + revarg: ptr %arg0
 + revarg: ptr %arg1
 + revarg: ptr %arg2
 + revarg: ptr %arg3
 + revarg:   %0 = call ptr @_Z17augment_AtomicMaxIdEPvPT_S2_PKS1_S2_(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3)
clang-18: /home/mixu/software/llvm-18.1.8/llvm-project-llvmorg-18.1.8/llvm/lib/IR/Instructions.cpp:680: void llvm::CallInst::init(llvm::FunctionType*, llvm::Value*, llvm::ArrayRef<llvm::Value*>, llvm::ArrayRef<llvm::OperandBundleDefT<llvm::Value*> >, const llvm::Twine&): Assertion `(Args.size() == FTy->getNumParams() || (FTy->isVarArg() && Args.size() > FTy->getNumParams())) && "Calling a function with bad signature!"' failed.

[minansys]

&"warning: GDB: Failed to set controlling terminal: Operation not permitted\n"
2: File exists.
clang++: warning: CUDA version is newer than the latest supported version 12.3 [-Wunknown-cuda-version]
revfn: ; Function Attrs: convergent mustprogress nounwind
define internal void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_339(ptr noundef %arg0, ptr noundef %arg1, ptr noundef %arg2, ptr noundef %arg3, double noundef %arg4) #138 {
entry:
%0 = alloca ptr, align 8
store double %arg4, ptr %0, align 8
%1 = load ptr, ptr %0, align 8
call void @_Z17reverse_AtomicMaxIdEvPT_S1_PKS0_S1_Pv(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3, ptr %1)
ret void
}

todiff: ; Function Attrs: convergent mustprogress noinline nounwind optnone
define dso_local noundef double @Z9AtomicMaxIdET_RS0_RKS0(ptr noundef nonnull align 8 dereferenceable(8) %dest, ptr noundef nonnull align 8 dereferenceable(8) %val) #5 !dbg !2120 !enzyme_augment !2127 !enzyme_gradient !2128 {
entry:
%dest.addr = alloca ptr, align 8
%val.addr = alloca ptr, align 8
store ptr %dest, ptr %dest.addr, align 8
call void @llvm.dbg.declare(metadata ptr %dest.addr, metadata !2129, metadata !DIExpression()), !dbg !2130
store ptr %val, ptr %val.addr, align 8
call void @llvm.dbg.declare(metadata ptr %val.addr, metadata !2131, metadata !DIExpression()), !dbg !2132
%0 = load ptr, ptr %dest.addr, align 8, !dbg !2133
%1 = load ptr, ptr %val.addr, align 8, !dbg !2134
%2 = load double, ptr %1, align 8, !dbg !2134
%call = call contract noundef double @_Z15atomicMaxDoublePdd(ptr noundef %0, double noundef %2) #139, !dbg !2135
ret double %call, !dbg !2136
}

revfn: define internal void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_340(ptr %arg0, ptr %"arg0'", ptr %arg1, ptr %"arg1'") {
entry:
call void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_339(ptr %arg0, ptr %"arg0'", ptr %arg1, ptr %"arg1'", double 0.000000e+00)
ret void
}

NewF: define internal void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3) {
entry:
%0 = call ptr @Z17augment_AtomicMaxIdEPvPT_S2_PKS1_S2(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3)
}

key rettype: CONSTANT

• revarg: ptr %arg0
• revarg: ptr %arg1
• revarg: ptr %arg2
• revarg: ptr %arg3
• revarg: %0 = call ptr @Z17augment_AtomicMaxIdEPvPT_S2_PKS1_S2(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3)
  clang-18: /home/mixu/software/llvm-18.1.8/llvm-project-llvmorg-18.1.8/llvm/lib/IR/Instructions.cpp:680: void llvm::CallInst::init(llvm::FunctionType*, llvm::Value*, llvm::ArrayRefllvm::Value*, llvm::ArrayRef<llvm::OperandBundleDefTllvm::Value* >, const llvm::Twine&): Assertion `(Args.size() == FTy->getNumParams() || (FTy->isVarArg() && Args.size() > FTy->getNumParams())) && "Calling a function with bad signature!"' failed.

[minansys]

#include <cuda_runtime.h>
#include <math.h>

// Enzyme device symbols and autodiff declaration (assumed to be provided)
#define enzyme_device_func __device__
extern int enzyme_device_func &enzyme_dup;
extern int enzyme_device_func &enzyme_dupnoneed;
extern int enzyme_device_func enzyme_out;
extern int enzyme_device_func &enzyme_const;
template <typename return_type, typename... T>
__device__ return_type __enzyme_autodiff(void*, T...);

#define Real double

// ---------------------------------------------------------------------
// Custom atomic maximum for double.
// This function atomically updates *address to be the maximum of its
// current value and val. It works by reinterpreting the double as a
// 64-bit unsigned integer and using atomicCAS in a loop.
// ---------------------------------------------------------------------
__host__ double atomicMaxDouble(double* address, double val)
{
    *address = fmax(*address, val);
    return *address;
}

__device__ double atomicMaxDouble(double* address, double val) {
    unsigned long long int* address_as_ull = reinterpret_cast<unsigned long long int*>(address);
    unsigned long long int old = *address_as_ull, assumed;
    do {
        assumed = old;
        double current = __longlong_as_double(assumed);
        double max_val = fmax(val, current);
        unsigned long long int newval = __double_as_longlong(max_val);
        old = atomicCAS(address_as_ull, assumed, newval);
    } while (assumed != old);
    return __longlong_as_double(old);

    // *address = max(*address, val);
    // return *address;
}

// ---------------------------------------------------------------------
// Templated wrapper for atomic maximum. We specialize for double.
// ---------------------------------------------------------------------
template<typename T>
__host__ __device__ T AtomicMax(T& dest, const T& val);

template<>
__host__ __device__ double AtomicMax<double>(double &dest, const double &val) {
    return atomicMaxDouble(&dest, val);
}

// ---------------------------------------------------------------------
// Augment and reverse functions for the atomic max (for Enzyme autodiff).
// These are heuristic subgradient definitions.
// ---------------------------------------------------------------------
template <typename T>
struct AtomicMaxTape {
    T old_val;   // value that was stored in *dest* before the atomic
    bool change;    // this thread’s candidate  (same as *val)
};

template <typename T>
__device__ AtomicMaxTape<T> augment_AtomicMax(T *dest, T *d_dest,
                                   const T *val, T *d_val)
{
    T mine = *val;
    T oldVal = AtomicMax(*dest, mine);

    bool changed = (mine > oldVal);
    AtomicMaxTape<T> tape;
    tape.old_val = oldVal;
    tape.change = changed;
    return tape;
}

template <typename T>
__device__ void reverse_AtomicMax(T *dest, T *d_dest,
                                  const T *val, T *d_val,
                                  bool changed)
{
    if (changed)
    {
        atomicAdd(d_val, *d_dest);
        *d_dest = T(0);
    }    
}

// ---------------------------------------------------------------------
// Registration for atomic max for double (for Enzyme autodiff)
// ---------------------------------------------------------------------
inline __device__ void* __enzyme_register_gradient_AtomicMax_double[] = {
    (void*)AtomicMax<double>,
    (void*)augment_AtomicMax<double>,
    (void*)reverse_AtomicMax<double>
};

// ---------------------------------------------------------------------
// Device function to compute the maximum over a 1D array ("face").
// This function loops over all n elements of the face and updates
// a global maximum variable using our AtomicMax.
// ---------------------------------------------------------------------
__device__ void face_max_impl(const Real* face, Real* global_max) {
    auto idx = threadIdx.x + blockDim.x * blockIdx.x;
    int side = 0;
    AtomicMax(global_max[idx + side], face[idx + 1 - side]);
    side = 1;
    AtomicMax(global_max[idx + side], face[idx + 1 - side]);
}

// ---------------------------------------------------------------------
// Kernel to run the forward pass over the face.
// We launch a single thread to loop over the face for simplicity.
// ---------------------------------------------------------------------
__global__ void face_max_kernel(const Real* face, Real* global_max, int n) {

    auto idx = threadIdx.x + blockDim.x * blockIdx.x;
	if (idx >= n - 1)
		return;

    face_max_impl(face, global_max);
}

// ---------------------------------------------------------------------
// Kernel to compute gradients (reverse mode) via Enzyme autodiff.
// This kernel applies autodiff to the face_max_impl function.
// ---------------------------------------------------------------------
__global__ void face_max_grad(Real* face, Real* d_face, Real* global_max, Real* d_global_max, int n) {
    auto idx = threadIdx.x + blockDim.x * blockIdx.x;
	if (idx >= n - 1)
		return;

    __enzyme_autodiff<void>((void*)(face_max_impl),
                            &enzyme_dup, face, d_face,
                            &enzyme_dup, global_max, d_global_max);
}

// ---------------------------------------------------------------------
// Main host function: set up a 1D "face" and test the atomic max and its gradient.
// ---------------------------------------------------------------------
int main() {
    const int n = 5;  // length of the face
    Real h_face[n], h_d_face[n], h_d_global_max[n], h_global_max[n];
    // Initialize the face values, e.g., from 1 to n.
    for (int i = 0; i < n; i++) {
        h_global_max[i] = -1e30;
        h_face[i] = (Real)abs((i + 1) - 3);
        h_d_global_max[i] = 1.0;
        h_d_face[i] = 0.0;
    }
    // The expected forward maximum is n (i.e., 64).

    // Allocate device memory for the face, its gradient, global max, and its gradient.
    Real *d_face, *d_d_face, *d_global_max, *d_d_global_max;
    cudaMalloc(&d_face, n * sizeof(Real));
    cudaMalloc(&d_d_face, n * sizeof(Real));
    cudaMalloc(&d_global_max, n * sizeof(Real));
    cudaMalloc(&d_d_global_max, n * sizeof(Real));

    // Copy the face array to device.
    cudaMemcpy(d_face, h_face, n * sizeof(Real), cudaMemcpyHostToDevice);
    cudaMemcpy(d_global_max, h_global_max, n * sizeof(Real), cudaMemcpyHostToDevice);
    // Initialize the gradient arrays: set the face gradients to 0 and the global max gradient to 1.
    cudaMemset(d_d_face, 0, n * sizeof(Real));
    cudaMemcpy(d_d_global_max, h_d_global_max, n * sizeof(Real), cudaMemcpyHostToDevice);

    // // Launch the forward kernel to compute the maximum over the face.
    // face_max_kernel<<<1, 1>>>(d_face, d_global_max, n);
    // cudaDeviceSynchronize();

    // Launch the gradient kernel (using Enzyme autodiff) for the face_max_impl function.
    face_max_grad<<<1, 256>>>(d_face, d_d_face, d_global_max, d_d_global_max, n);
    cudaDeviceSynchronize();

    // Retrieve results.
    cudaMemcpy(h_global_max, d_global_max, n * sizeof(Real), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_d_global_max, d_d_global_max, n * sizeof(Real), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_d_face, d_d_face, n * sizeof(Real), cudaMemcpyDeviceToHost);

    // Print the results.
    for (int i = 0; i < n; i++) {
        printf("global_max[%d] = %f, d_global_max[%d] = %f, h_d_face[%d] = %f\n", i, h_global_max[i], i, h_d_global_max[i], i, h_d_face[i]);
    }

    // Clean up device memory.
    cudaFree(d_face);
    cudaFree(d_d_face);
    cudaFree(d_global_max);
    cudaFree(d_d_global_max);
    return 0;
}

[minansys]

&"warning: GDB: Failed to set controlling terminal: Operation not permitted\n"
2: File exists.
clang++: warning: CUDA version is newer than the latest supported version 12.3 [-Wunknown-cuda-version]
revfn: ; Function Attrs: convergent mustprogress nounwind
define internal void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_339(ptr noundef %arg0, ptr noundef %arg1, ptr noundef %arg2, ptr noundef %arg3, double noundef zeroext %arg4) #138 {
entry:
%0 = alloca i1, align 1
store double %arg4, ptr %0, align 8
%1 = load i1, ptr %0, align 1
call void @_Z17reverse_AtomicMaxIdEvPT_S1_PKS0_S1_b(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3, i1 %1)
ret void
}

todiff: ; Function Attrs: convergent mustprogress noinline nounwind optnone
define dso_local noundef double @Z9AtomicMaxIdET_RS0_RKS0(ptr noundef nonnull align 8 dereferenceable(8) %dest, ptr noundef nonnull align 8 dereferenceable(8) %val) #5 !dbg !2120 !enzyme_augment !2127 !enzyme_gradient !2128 {
entry:
%dest.addr = alloca ptr, align 8
%val.addr = alloca ptr, align 8
store ptr %dest, ptr %dest.addr, align 8
call void @llvm.dbg.declare(metadata ptr %dest.addr, metadata !2129, metadata !DIExpression()), !dbg !2130
store ptr %val, ptr %val.addr, align 8
call void @llvm.dbg.declare(metadata ptr %val.addr, metadata !2131, metadata !DIExpression()), !dbg !2132
%0 = load ptr, ptr %dest.addr, align 8, !dbg !2133
%1 = load ptr, ptr %val.addr, align 8, !dbg !2134
%2 = load double, ptr %1, align 8, !dbg !2134
%call = call contract noundef double @_Z15atomicMaxDoublePdd(ptr noundef %0, double noundef %2) #139, !dbg !2135
ret double %call, !dbg !2136
}

revfn: define internal void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_340(ptr %arg0, ptr %"arg0'", ptr %arg1, ptr %"arg1'") {
entry:
call void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_339(ptr %arg0, ptr %"arg0'", ptr %arg1, ptr %"arg1'", double 0.000000e+00)
ret void
}

NewF: define internal void @fixgradient__Z9AtomicMaxIdET_RS0_RKS0_(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3) {
entry:
%0 = call { double, double } @fixaugment__Z17augment_AtomicMaxIdE13AtomicMaxTapeIT_EPS1_S3_PKS1_S3_(ptr %arg0, ptr %arg1, ptr %arg2, ptr %arg3)
%1 = extractvalue { double, double } %0, 0
}

key rettype: CONSTANT

• revarg: ptr %arg0
• revarg: ptr %arg1
• revarg: ptr %arg2
• revarg: ptr %arg3
• revarg: %1 = extractvalue { double, double } %0, 0
  clang-18: /home/mixu/software/llvm-18.1.8/llvm-project-llvmorg-18.1.8/llvm/lib/IR/Instructions.cpp:680: void llvm::CallInst::init(llvm::FunctionType*, llvm::Value*, llvm::ArrayRefllvm::Value*, llvm::ArrayRef<llvm::OperandBundleDefTllvm::Value* >, const llvm::Twine&): Assertion `(Args.size() == FTy->getNumParams() || (FTy->isVarArg() && Args.size() > FTy->getNumParams())) && "Calling a function with bad signature!"' failed.

[wsmoses]

wmoses@beast:~/git/Enzyme/enzyme/build15 (cuse) $ ~/llvms/llvm15/buildD/bin/clang++ cuse.cu -fpass-plugin=Enzyme/ClangEnzyme-15.so -Xclang -load -Xclang Enzyme/ClangEnzyme-15.so -O2 -mllvm -enzyme-strict-aliasing=0  --cuda-gpu-arch=sm_60

[wsmoses]

should be fixed by #2315:

you should get better error messages, and also the end to end fix

template <typename T>
struct AtomicMaxTape {
    bool change;    // this thread’s candidate  (same as *val)
    T old_val;   // value that was stored in *dest* before the atomic
};

template <typename T>
__device__ AtomicMaxTape<T> augment_AtomicMax(T *dest, T *d_dest,
                                   const T *val, T *d_val)
{
    T mine = *val;
    T oldVal = AtomicMax(*dest, mine);

    bool changed = (mine > oldVal);
    AtomicMaxTape<T> tape;
    tape.old_val = oldVal;
    tape.change = changed;
    return tape;
}

template <typename T>
__device__ void reverse_AtomicMax(T *dest, T *d_dest,
                                  const T *val, T *d_val,
                                  T dres0,
                                  bool changed)
{
    if (changed)
    {
        double dres = *d_dest + dres0;
        atomicAdd(d_val, dres);
        *d_dest = T(0);
    }    
}