utkinis/tinykernels.jl

TinyKernels.jl

TinyKernels.jl provides a tiny abstraction for GPU (and CPU) kernels, with full support for CUDA (Nvidia) and ROCm (AMD) backends, limited support for Metal (GPU programming on MacOS ARM) backend, and allowing for multi-threaded CPU execution.

TinyKernels.jl is mostly a heavily stripped-down version of KernelAbstractions.jl supporting the bare minimum of the features. This package provides a sandbox for Julia GPU tooling and to measure the performance of kernels in a GPU-agnostic way. While the API of KernelAbstractions.jl is in a "transient" state, this package will provide the thin abstraction layer on top the CUDA.jl, AMDGPU.jl and Metal.jl packages.

TinyKernels.jl allows to explicitly launch GPU kernels asynchronously on different streams with given priority. This feature facilitates the overlap between computations and memory transfers in distributed configurations.

TinyKernels.jl supports automatic differentiation with Enzyme.jl overloading the Enzyme.autodiff function to enable reverse mode AD of GPU (and CPU) kernels.

Preliminary benchmarks can be found in TinyBenchmarks.jl and Metal playground in MetalGPU.

Stay tuned 🚀

Notes

⚠️ Metal backend:

• Only Float32 is being supported. For Float64, one could try using a construct from DoubleFloats.jl which may impact performance.

tinykernels.jl's Issues

AMDGPU device selection errors

Specifying another default device than the the "first" results in unstable and faulty execution of the code on AMD GPU.

Upon adding AMDGPU.default_device_id!(2),

TinyKernels.jl/src/roc_backend.jl

Line 55 in 8a292d9

config = AMDGPU.launch_configuration(roc_kernel.fun)

fails with following error

ERROR: LoadError: Failed to find local memory region for GPU-fa397c70dbf3095f [gfx90a]
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] device_local_memory_size(device::AMDGPU.Runtime.ROCDevice)
    @ AMDGPU.Runtime /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/runtime/device.jl:160
  [3] _calculate_occupancy(fun::AMDGPU.Runtime.ROCFunction, device::AMDGPU.Runtime.ROCDevice; input_block_size::Int64, localmem::Int64)
    @ AMDGPU.Compiler /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/compiler/occupancy.jl:126
  [4] #72
    @ /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/compiler/occupancy.jl:49 [inlined]
  [5] get!(default::AMDGPU.Compiler.var"#72#74"{Int64, Int64, AMDGPU.Runtime.ROCFunction}, h::Dict{Tuple{AMDGPU.Runtime.ROCDevice, UInt64, Int64, Int64}, NamedTuple}, key::Tuple{AMDGPU.Runtime.ROCDevice, UInt64, Int64, Int64})
    @ Base ./dict.jl:468
  [6] #71
    @ /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/compiler/occupancy.jl:48 [inlined]
  [7] lock(f::AMDGPU.Compiler.var"#71#73"{Int64, Int64, AMDGPU.Runtime.ROCFunction}, l::ReentrantLock)
    @ Base ./lock.jl:229
  [8] #calculate_occupancy#70
    @ /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/compiler/occupancy.jl:47 [inlined]
  [9] launch_configuration(fun::AMDGPU.Runtime.ROCFunction; input_block_size::Int64, localmem::Int64)
    @ AMDGPU /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/highlevel.jl:460
 [10] launch_configuration
    @ /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/highlevel.jl:459 [inlined]
 [11] (::Kernel{TinyKernels.ROCBackend.ROCDevice, typeof(test_function!)})(::ROCMatrix{Float64}, ::Vararg{Any}; range::Tuple{UnitRange{Int64}, UnitRange{Int64}}, priority::Symbol)
    @ TinyKernels.ROCBackend /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/src/roc_backend.jl:55
 [12] main()
    @ Main /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/scripts/example_amdgpu.jl:37
 [13] top-level scope
    @ /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/scripts/example_amdgpu.jl:50
 [14] include(fname::String)
    @ Base.MainInclude ./client.jl:478
 [15] top-level scope
    @ REPL[1]:1
in expression starting at /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/scripts/example_amdgpu.jl:50

Hard-coding nthreads further leads to a hang when computing

TinyKernels.jl/scripts/example_amdgpu.jl

Line 45 in 8a292d9

@assert A ≈ B .+ s .* C

:0:rocdevice.cpp            :2614: 658312983828 us: 992  : [tid:0x14ac1239c700] Device::callbackQueue aborting with error : HSA_STATUS_ERROR_MEMORY_FAULT: Agent attempted to access an inaccessible address. code: 0x2b

[992] signal (6.-6): Aborted
in expression starting at /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/scripts/example_amdgpu.jl:50
Allocations: 44409871 (Pool: 44373188; Big: 36683); GC: 68
Aborted

Removing that assertion as well results in running code. However, changing the device within the same Julia session leads to

julia> include("example_amdgpu.jl")
step 1
┌ Error: Memory Fault on GPU-fa397c70dbf3095f [gfx90a] at 0x00001471c3c07000:
│ PAGE_NOT_PRESENT, READ_ONLY, NX, HOST_ONLY, DRAMECC, IMPRECISE, SRAMECC, HANG
│ GPU kernels are now hung, please restart Julia
└ @ AMDGPU.Runtime /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/runtime/fault.jl:25
^CERROR: LoadError: InterruptException:
Stacktrace:
  [1] safepoint
    @ ./gcutils.jl:203 [inlined]
  [2] wait()
    @ Base ./task.jl:970
  [3] yield()
    @ Base ./task.jl:853
  [4] wait(signal::ROCSignal; timeout::Nothing, min_latency::Int64, queue::ROCQueue)
    @ AMDGPU.Runtime /scratch/project_465000139/lurass/julia_local/julia_depot/packages/AMDGPU/bzHD4/src/runtime/signal.jl:138
  [5] wait
    @ /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/src/roc_backend.jl:18 [inlined]
  [6] _broadcast_getindex_evalf
    @ ./broadcast.jl:683 [inlined]
  [7] _broadcast_getindex
    @ ./broadcast.jl:656 [inlined]
  [8] getindex
    @ ./broadcast.jl:610 [inlined]
  [9] macro expansion
    @ ./broadcast.jl:974 [inlined]
 [10] macro expansion
    @ ./simdloop.jl:77 [inlined]
 [11] copyto!
    @ ./broadcast.jl:973 [inlined]
 [12] copyto!
    @ ./broadcast.jl:926 [inlined]
 [13] copy
    @ ./broadcast.jl:898 [inlined]
 [14] materialize
    @ ./broadcast.jl:873 [inlined]
 [15] wait(evs::Vector{TinyKernels.ROCBackend.ROCEvent})
    @ TinyKernels.ROCBackend /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/src/roc_backend.jl:19
 [16] main()
    @ Main /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/scripts/example_amdgpu.jl:40
 [17] top-level scope
    @ /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/scripts/example_amdgpu.jl:50
 [18] include(fname::String)
    @ Base.MainInclude ./client.jl:478
 [19] top-level scope
    @ REPL[1]:1
in expression starting at /pfs/lustrep4/scratch/project_465000139/lurass/TinyKernels.jl/scripts/example_amdgpu.jl:50