CUDA Error 717: Invalid Address Space -- Atomic operation not supported on stack memory.
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
I am on Julia 1.8.0, Enzyme main (71ebce9) and CUDA 3.12.0. The following CUDA kernel runs but the gradient kernel with Enzyme errors.
```julia
using CUDA, Enzyme, Molly, LinearAlgebra
function find_neighbors(coords)
n_atoms = length(coords)
neighbors = Tuple{Int, Int}[]
for i in 1:n_atoms
for j in (i + 1):n_atoms
if norm(vector(coords[i], coords[j], boundary)) <= 1.0
push!(neighbors, (i, j))
end
end
end
return neighbors
end
n_atoms = 1024
boundary = CubicBoundary(2.7f0, 2.7f0, 2.7f0)
coords = place_atoms(n_atoms, boundary; min_dist=0.05f0)
atoms = [Atom(charge=0.0f0, mass=1.0f0, σ=Float32(0.02 + rand() * 0.01),
ϵ=Float32(0.2 + rand() * 0.1)) for _ in 1:n_atoms]
cu_coords = CuArray(coords)
cu_atoms = CuArray(atoms)
neighbors = find_neighbors(coords)
cu_neighbors = CuArray(neighbors)
function force(c1, c2, a1, a2)
dr = c2 - c1
invr2 = inv(sum(abs2, dr))
σ = (a1.σ + a2.σ) / 2
ϵ = sqrt(a1.ϵ * a2.ϵ)
six_term = (σ^2 * invr2) ^ 3
f = (24 * ϵ * invr2) * (2 * six_term ^ 2 - six_term)
return f * dr
end
function kernel!(forces::CuDeviceMatrix{T}, coords_var, atoms_var, neighbors_var, ::Val{M}) where {T, M}
coords = CUDA.Const(coords_var)
atoms = CUDA.Const(atoms_var)
neighbors = CUDA.Const(neighbors_var)
tidx = threadIdx().x
inter_ig = (blockIdx().x - 1) * blockDim().x + tidx
stride = gridDim().x * blockDim().x
shared_fs = CuStaticSharedArray(T, (3, M))
shared_is = CuStaticSharedArray(Int32, M)
shared_js = CuStaticSharedArray(Int32, M)
if tidx == 1
for si in 1:M
shared_is[si] = zero(Int32)
end
end
sync_threads()
for (thread_i, inter_i) in enumerate(inter_ig:stride:length(neighbors))
si = (thread_i - 1) * blockDim().x + tidx
i, j = neighbors[inter_i]
f = force(coords[i], coords[j], atoms[i], atoms[j])
shared_fs[1, si] = f[1]
shared_fs[2, si] = f[2]
shared_fs[3, si] = f[3]
shared_is[si] = i
shared_js[si] = j
end
sync_threads()
if tidx == 1
for si in 1:M
i = shared_is[si]
if iszero(i)
break
end
j = shared_js[si]
dx, dy, dz = shared_fs[1, si], shared_fs[2, si], shared_fs[3, si]
forces[1, i] -= dx
forces[2, i] -= dy
forces[3, i] -= dz
forces[1, j] += dx
forces[2, j] += dy
forces[3, j] += dz
end
end
return
end
function grad_kernel!(forces, dforces, coords, dcoords, atoms, datoms, neighbors, shared_mem_size)
Enzyme.autodiff_deferred(
kernel!,
Duplicated(forces, dforces),
Duplicated(coords, dcoords),
Duplicated(atoms, datoms),
Const(neighbors),
Const(shared_mem_size),
)
return
end
cu_forces_mat = CuArray(zeros(Float32, 3, n_atoms))
d_cu_forces_mat = zero(cu_forces_mat)
d_cu_coords = zero(cu_coords)
d_cu_atoms = zero(cu_atoms)
n_threads = 256
n_blocks = 800
shared_mem_size = 2000
CUDA.@sync @cuda threads=n_threads blocks=n_blocks kernel!(cu_forces_mat, cu_coords, cu_atoms,
cu_neighbors, Val(shared_mem_size))
println("Kernel worked")
CUDA.@sync @cuda threads=n_threads blocks=n_blocks grad_kernel!(cu_forces_mat, d_cu_forces_mat,
cu_coords, d_cu_coords, cu_atoms, d_cu_atoms, cu_neighbors, Val(shared_mem_size))
println("Grad kernel worked")
```
The output is:
```
Kernel worked
warning found shared memory
ERROR: LoadError: Failed to compile PTX code (ptxas exited with code 255)
ptxas error : Entry function '_Z23julia_grad_kernel__579813CuDeviceArrayI7Float32Li2ELi1EES_IS0_Li2ELi1EES_I6SArrayI5TupleILi3EES0_Li1ELi3EELi1ELi1EES_IS1_IS2_ILi3EES0_Li1ELi3EELi1ELi1EES_I4AtomIS0_S0_S0_S0_ELi1ELi1EES_IS3_IS0_S0_S0_S0_ELi1ELi1EES_IS2_I5Int64S4_ELi1ELi1EE3ValILi2000EE' uses too much shared data (0xfa00 bytes, 0xc000 max)
ptxas info : 292 bytes gmem
ptxas info : Compiling entry function '_Z23julia_grad_kernel__579813CuDeviceArrayI7Float32Li2ELi1EES_IS0_Li2ELi1EES_I6SArrayI5TupleILi3EES0_Li1ELi3EELi1ELi1EES_IS1_IS2_ILi3EES0_Li1ELi3EELi1ELi1EES_I4AtomIS0_S0_S0_S0_ELi1ELi1EES_IS3_IS0_S0_S0_S0_ELi1ELi1EES_IS2_I5Int64S4_ELi1ELi1EE3ValILi2000EE' for 'sm_86'
ptxas info : Function properties for _Z23julia_grad_kernel__579813CuDeviceArrayI7Float32Li2ELi1EES_IS0_Li2ELi1EES_I6SArrayI5TupleILi3EES0_Li1ELi3EELi1ELi1EES_IS1_IS2_ILi3EES0_Li1ELi3EELi1ELi1EES_I4AtomIS0_S0_S0_S0_ELi1ELi1EES_IS3_IS0_S0_S0_S0_ELi1ELi1EES_IS2_I5Int64S4_ELi1ELi1EE3ValILi2000EE
208 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Used 96 registers, 64000 bytes smem, 600 bytes cmem[0], 24 bytes cmem[2]
ptxas info : Function properties for augmented_julia_force_6190
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Function properties for diffejulia_force_6190
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Function properties for gpu_report_exception
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Function properties for gpu_signal_exception
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Function properties for julia__throw_boundserror_6169
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Function properties for julia__throw_inexacterror_6184
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info : Function properties for julia_steprange_last_6192
0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
If you think this is a bug, please file an issue and attach /tmp/jl_vo7Mbazqzg.ptx
Stacktrace:
[1] error(s::String)
@ Base ./error.jl:35
[2] cufunction_compile(job::GPUCompiler.CompilerJob, ctx::LLVM.Context)
@ CUDA ~/.julia/packages/CUDA/DfvRa/src/compiler/execution.jl:427
[3] #224
@ ~/.julia/packages/CUDA/DfvRa/src/compiler/execution.jl:347 [inlined]
[4] JuliaContext(f::CUDA.var"#224#225"{GPUCompiler.CompilerJob{GPUCompiler.PTXCompilerTarget, CUDA.CUDACompilerParams, GPUCompiler.FunctionSpec{typeof(grad_kernel!), Tuple{CuDeviceMatrix{Float32, 1}, CuDeviceMatrix{Float32, 1}, CuDeviceVector{SVector{3, Float32}, 1}, CuDeviceVector{SVector{3, Float32}, 1}, CuDeviceVector{Atom{Float32, Float32, Float32, Float32}, 1}, CuDeviceVector{Atom{Float32, Float32, Float32, Float32}, 1}, CuDeviceVector{Tuple{Int64, Int64}, 1}, Val{2000}}}}})
@ GPUCompiler ~/.julia/packages/GPUCompiler/jVY4I/src/driver.jl:76
[5] cufunction_compile(job::GPUCompiler.CompilerJob)
@ CUDA ~/.julia/packages/CUDA/DfvRa/src/compiler/execution.jl:346
[6] cached_compilation(cache::Dict{UInt64, Any}, job::GPUCompiler.CompilerJob, compiler::typeof(CUDA.cufunction_compile), linker::typeof(CUDA.cufunction_link))
@ GPUCompiler ~/.julia/packages/GPUCompiler/jVY4I/src/cache.jl:90
[7] cufunction(f::typeof(grad_kernel!), tt::Type{Tuple{CuDeviceMatrix{Float32, 1}, CuDeviceMatrix{Float32, 1}, CuDeviceVector{SVector{3, Float32}, 1}, CuDeviceVector{SVector{3, Float32}, 1}, CuDeviceVector{Atom{Float32, Float32, Float32, Float32}, 1}, CuDeviceVector{Atom{Float32, Float32, Float32, Float32}, 1}, CuDeviceVector{Tuple{Int64, Int64}, 1}, Val{2000}}}; name::Nothing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
@ CUDA ~/.julia/packages/CUDA/DfvRa/src/compiler/execution.jl:299
[8] cufunction(f::typeof(grad_kernel!), tt::Type{Tuple{CuDeviceMatrix{Float32, 1}, CuDeviceMatrix{Float32, 1}, CuDeviceVector{SVector{3, Float32}, 1}, CuDeviceVector{SVector{3, Float32}, 1}, CuDeviceVector{Atom{Float32, Float32, Float32, Float32}, 1}, CuDeviceVector{Atom{Float32, Float32, Float32, Float32}, 1}, CuDeviceVector{Tuple{Int64, Int64}, 1}, Val{2000}}})
@ CUDA ~/.julia/packages/CUDA/DfvRa/src/compiler/execution.jl:292
[9] macro expansion
@ ~/.julia/packages/CUDA/DfvRa/src/compiler/execution.jl:102 [inlined]
[10] top-level scope
@ ~/.julia/packages/CUDA/DfvRa/src/utilities.jl:25
```
Changing the `dx, dy, dz = shared_fs[1, si], shared_fs[2, si], shared_fs[3, si]` line to `dx, dy, dz = 1.0f0, 1.0f0, 1.0f0` makes this not error, suggesting a problem when updating an array using the shared memory.
I realise that I should be using atomics at the `forces[1, i] -= dx` lines but that errored due to https://github.com/EnzymeAD/Enzyme.jl/issues/421. This may be a separate error though since the following simple test works, suggesting that the above should also work.
```julia
using Enzyme, CUDA
function mul_kernel(A)
shared = CuStaticSharedArray(Float32, 64)
i = threadIdx().x
if i <= length(A)
shared[i] = A[i] * A[i]
end
sync_threads()
if i == 1
for j in 1:length(A)
A[j] = shared[j]
end
end
return nothing
end
function grad_mul_kernel(A, dA)
Enzyme.autodiff_deferred(mul_kernel, Const, Duplicated(A, dA))
return nothing
end
A = CUDA.ones(64,)
dA = similar(A)
dA .= 1
@cuda threads=length(A) grad_mul_kernel(A, dA)
dA
```
Contributor guide
Assessment
This issue has not been assessed yet.