Skip to content

Commit

Permalink
adding necessary changes for KA transition with gpuarrays
Browse files Browse the repository at this point in the history
  • Loading branch information
leios committed Jul 22, 2024
1 parent e872db7 commit 87e4daf
Showing 1 changed file with 1 addition and 45 deletions.
46 changes: 1 addition & 45 deletions src/gpuarrays.jl
Original file line number Diff line number Diff line change
@@ -1,48 +1,4 @@
struct ROCArrayBackend <: AbstractGPUBackend end

struct ROCKernelContext <: AbstractKernelContext end

@inline function GPUArrays.gpu_call(
::ROCArrayBackend, f, args, threads::Int, blocks::Int; name::Maybe{String},
)
@roc gridsize=blocks groupsize=threads name=name f(ROCKernelContext(), args...)
end

# indexing

for (f, froc) in (
(:blockidx, :blockIdx),
(:blockdim, :blockDim),
(:threadidx, :threadIdx),
(:griddim, :gridGroupDim)
)
@eval @inline GPUArrays.$f(::ROCKernelContext) = AMDGPU.$froc().x
end

# math

@inline GPUArrays.cos(::ROCKernelContext, x) = cos(x)
@inline GPUArrays.sin(::ROCKernelContext, x) = sin(x)
@inline GPUArrays.sqrt(::ROCKernelContext, x) = sqrt(x)
@inline GPUArrays.log(::ROCKernelContext, x) = log(x)

# memory

@inline function GPUArrays.LocalMemory(::ROCKernelContext, ::Type{T}, ::Val{dims}, ::Val{id}) where {T,dims,id}
ptr = AMDGPU.Device.alloc_special(Val{id}(), T, Val{AMDGPU.AS.Local}(), Val{prod(dims)}())
ROCDeviceArray(dims, ptr)
end

# synchronization

@inline function GPUArrays.synchronize_threads(::ROCKernelContext)
sync_workgroup()
return
end

GPUArrays.device(x::ROCArray) = x.buf[].device

GPUArrays.backend(::Type{<:ROCArray}) = ROCArrayBackend()
import KernelAbstractions

function GPUArrays.derive(
::Type{T}, x::ROCArray, dims::Dims{N}, offset::Int,
Expand Down

0 comments on commit 87e4daf

Please sign in to comment.