Enzyme failed to support atomicAdd, atomicCAS, and assert for cuda code
- Dominant language
- LLVM
- Stars
- 1.7k
- Forks
- 188
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 26
Description
I am using enzyme for gpu code. Here is the some information
```
clang: 15.0
nvcc: cuda 11.8
WSL linux: Ubuntu 20.04 LTS,
Enzyme:
commit 2fe5164a2423dd67ef25e2c4fb204fd06362494b (HEAD -> main, origin/main, origin/HEAD)
Author: William Moses
Date: Thu Aug 15 23:00:59 2024 -0500
Build:
clang++ -I/usr/local/cuda-11.8/include example.cu -o example -O2 --cuda-path=/usr/local/cuda-11.8/ --cuda-gpu-arch=sm_86 -L/usr/local/cuda-11.8/lib64 -lcudart -v -DENABLE_ENZYME -fplugin=/home/mixu/software/Enzyme/enzyme/build/Enzyme/ClangEnzyme-15.so -mllvm -enzyme-loose-types=1
```
```
#include
#include
#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 >
enzyme_device_func return_type __enzyme_fwddiff(void*, T ... );
template < typename return_type, typename ... T >
enzyme_device_func return_type __enzyme_autodiff(void*, T ... );
#define dev std
#define Real double
//#define FORWARD
// AtomicAdd -----------------------------
// #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 600
// __device__ static double atomicAdd(double* address, double val)
// {
// unsigned long long int* address_as_ull = (unsigned long long int*)address;
// unsigned long long int old = *address_as_ull, assumed;
// do
// {
// assumed = old;
// old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed)));
// // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN)
// }
// while (assumed != old);
// return __longlong_as_double(old);
// }
// #endif
template
__device__ static T AtomicAdd(T& dest, const T& val)
{
return atomicAdd(&dest, val);
}
// AtomicMax -----------------------------
__device__ static double atomicMax(double* address, double val)
{
unsigned long long int* address_as_ull = (unsigned long long int*)address;
unsigned long long int old = *address_as_ull, assumed;
do
{
assumed = old;
old = atomicCAS(address_as_ull, assumed, __double_as_longlong(dev::fmax(val, __longlong_as_double(assumed))));
// Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN)
}
while (assumed != old);
return old;
}
template
__device__ static T AtomicMax(T& dest, const T& val)
{
return atomicMax(&dest, val);
}
void __device__ square_impl(Real* x_in, Real *x_out) {
//assert(x_in[0] == 0.0);
//x_out[0] = x_in[0] * x_in[0];
AtomicMax(x_out[0], x_in[0]);
//AtomicAdd(x_out[0], x_in[0]);
}
typedef void (*f_ptr)(Real*, Real*);
void __global__ square(Real* x_in, Real *x_out) {
square_impl(x_in, x_out);
}
void __global__ square_grad(Real* x, Real *d_x, Real *y, Real *d_y) {
__enzyme_autodiff((void *)(square_impl),
enzyme_dup, x, d_x,
enzyme_dup, y, d_y);
}
int main() {
// Device pointers
Real *x, *d_x, *y, *d_y;
// Allocate GPU device memory
cudaMalloc(&x, sizeof(*x));
cudaMalloc(&d_x, sizeof(*d_x));
cudaMalloc(&y, sizeof(*y));
cudaMalloc(&d_y, sizeof(*d_y));
// Initialize device values
Real host_x = 1.4;
Real host_d_x = 0.0;
Real host_y = 1.0;
Real host_d_y = 1.0;
// Copy data to device
cudaMemcpy(x, &host_x, sizeof(*x), cudaMemcpyHostToDevice);
cudaMemcpy(d_x, &host_d_x, sizeof(*d_x), cudaMemcpyHostToDevice);
cudaMemcpy(y, &host_y, sizeof(*y), cudaMemcpyHostToDevice);
cudaMemcpy(d_y, &host_d_y, sizeof(*d_y), cudaMemcpyHostToDevice);
#ifdef FORWARD
// Forward pass only
square<<<1, 1>>>(x, y);
#else
// Forward and backward pass
//square_grad<<<1, 1>>>(x, d_x, y, d_y);
#endif
// Synchronize device memory
cudaDeviceSynchronize();
// Copy data from device to host
cudaMemcpy(&host_x, x, sizeof(*x), cudaMemcpyDeviceToHost);
cudaMemcpy(&host_d_x, d_x, sizeof(*d_x), cudaMemcpyDeviceToHost);
cudaMemcpy(&host_y, y, sizeof(*y), cudaMemcpyDeviceToHost);
cudaMemcpy(&host_d_y, d_y, sizeof(*d_y), cudaMemcpyDeviceToHost);
// Print results
printf("%f %f\n", host_x, host_y);
printf("%f %f\n", host_d_x, host_d_y);
}
```
in `square_impl` if I use `AtomicMax`, it failed to compile with error msg:
```
example.cu:65:17: error: Enzyme:
%6 = phi i64 [ %4, %2 ], [ %11, %5 ]: {[-1]:Float@double}, intvals: {}
%10 = cmpxchg ptr %1, i64 %6, i64 %9 seq_cst seq_cst, align 8: {}, intvals: {}
%9 = bitcast double %8 to i64: {}, intvals: {}
%8 = tail call contract double @__nv_fmax(double noundef %3, double noundef %7) #149: {[-1]:Float@double}, intvals: {}
ptr %0: {[-1]:Pointer, [-1,0]:Float@double}, intvals: {}
ptr %1: {[-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}, intvals: {}
%12 = icmp eq i64 %6, %11: {[-1]:Integer}, intvals: {}
%11 = extractvalue { i64, i1 } %10, 0: {}, intvals: {}
%7 = bitcast i64 %6 to double: {[-1]:Float@double}, intvals: {}
%4 = load i64, ptr %1, align 8, !tbaa !12: {[-1]:Integer}, intvals: {}
%3 = load double, ptr %0, align 8, !tbaa !8: {[-1]:Float@double}, intvals: {}
Illegal updateAnalysis prev:{[-1]:Integer} new: {[-1]:Float@double}
val: %4 = load i64, ptr %1, align 8, !tbaa !12 origin= %6 = phi i64 [ %4, %2 ], [ %11, %5 ]
void __device__ square_impl(Real* x_in, Real *x_out) {
^
```
in `square_impl` if I use `AtomicAdd`, it fails to compile with the error msg:
```
fatal error: error in backend: Cannot select: 0x558d7c69bb10: f64,ch = AtomicLoad<(load seq_cst (s64) from %ir.5, addrspace 1)> 0x558d7c39d080:1, 0x558d7c69da30
0x558d7c69da30: i64 = addrspacecast[0 -> 1] 0x558d7c463798
0x558d7c463798: i64,ch = load<(dereferenceable invariant load (s64) from `ptr addrspace(101) null`, addrspace 101)> 0x558d7c432ca8, TargetExternalSymbol:i64'_Z11square_gradPdS_S_S__param_3', undef:i64
0x558d7c69b1b8: i64 = TargetExternalSymbol'_Z11square_gradPdS_S_S__param_3'
0x558d7c4635f8: i64 = undef
In function: _Z11square_gradPdS_S_S_
clang: error: clang frontend command failed with exit code 70 (use -v to see invocation)
Ubuntu clang version 15.0.7
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin
clang: note: diagnostic msg:
********************
PLEASE ATTACH THE FOLLOWING FILES TO THE BUG REPORT:
Preprocessed source(s) and associated run script(s) are located at:
clang: note: diagnostic msg: /tmp/example-d21287.cu
clang: note: diagnostic msg: /tmp/example-01fc49/example-sm_86.cu
clang: note: diagnostic msg: /tmp/example-d21287.sh
clang: note: diagnostic msg:
********************
```
in `square_impl` if I use `assert(x_out[0] == 0.0);`, it fails to compile with the error msg:
```
error: Enzyme: No create nofree of empty function (__assertfail) __assertfail)
at context: tail call void @__assertfail(ptr noundef nonnull @.str, ptr noundef nonnull @.str1, i32 noundef 67, ptr noundef nonnull @__PRETTY_FUNCTION__._Z11square_implPdS_, i64 noundef 1) #149 (__assertfail)
1 error generated when compiling for sm_86.
```
The forward codes comply and run fine with the __enzyme_autodiff function commented out.
Contributor guide
Assessment
This issue has not been assessed yet.