JACC v1.0 API Design Documentation

[williamfgc]

Goal: discuss the JACC APIs with the broader community towards a stable public facing APIs for JACC v1.0.
This is the closest to a JACC governance process in which community members can provide their inputs for decision making.

JACC, Julia for Accelerators, API is designed for user to write parallel code once, in the form of “kernels” represented by independent parallel for loops (parallel_for) and reduction operations (parallel_reduce) on CPU and GPUs leveraging the JuliaGPU ecosystem.

Users interact with four major components:

• Backend configuration APIs, outside code: set_backend and unset_backend, inside code for autodetection: @init_backend
• Array APIs for memory management: zeros, ones and array
• Kernel launching APIs: parallel_for and parallel_reduce, advanced: launch_spec
• Inside kernel capabilities APIs: atomic, shared

Backend configuration APIs

• JACC.set_backend(:backend) or JACC.set_backend("backend"): JACC uses [Preferences.jl](https://github.com/JuliaPackaging/Preferences.jl) to modify an entry in LocalPreference.toml.:backendcan be Supported backends:{:threads", :cuda", :amdgpu", :oneAPI"}` (current plan is to extend to Apple). Threads is the assumed backend if this function is never called. Example in CI.

LocalPreferences.toml example:

[JACC]
backends = ["threads", "cuda"]
default_backend = "cuda"
    [JACC.placement]
    oneapi = "weakdeps"

• JACC.remove_backend("backend") when JACC.set_backend("backend") the backend package (e.g. CUDA.jl, AMDGPU.jl) is downloaded and added polluting the Project.toml file. Yhis function is aimed for developers to clean up this file and prevents from committing JACC and backend dependencies in the same Project.toml

Note: set_backend and unset_backend are not meant to be used in code, only as preliminary steps to configure JACC.

• JACC.@init_backend to be of mandatory use right after import JACC or using JACC, it is the autodetection mechanism for portable user code that is meant to not be backend-dependent. Not using it would result in backend-related errors. Example

import JACC
JACC.@init_backend

Array APIs for memory management

dims is a tuple of integers of the form (Nx, Ny, Nz,...)

• JACC.zeros(T, Nx, Ny, Nz, …) or JACC.zeros(T, dims) return an nd-array of type T (e.g. Int, Float32, Float64) filled with zeros. Example

• JACC.ones(T, Nx, Ny, Nz, …) or JACC.zeros(T, dims) return an nd-array of type T (e.g. Int, Float32, Float64) filled with ones. Example

• JACC.array(Core.AbstractArray) return an nd-array that is a subtype of Core.AbstractArray with uninitialized memory allocated on device. Example

• JACC.fill(value, Nx, Ny, Nz, ...) or JACC.fill( value, dims ) return an nd-array of typeT(e.g.Int, Float32, Float64`) filled with a particular value. The array takes the value type. Example

• JACC.to_host(x_d) : returns a copy of an device-allocated array (e.g. GPU) in host (e.g. CPU), if the array is allocated in CPU it returns the same array. Example

Kernel Launching APIs: parallel_for, parallel_reduce and synchronize

Note: functions can also be passed using the do and end syntax construct in our examples after the parallel_for/parallel_reduce functions.

parallel_for APIs

Execute the following for loops with the assumption that each (i,j) operation is independent, f is the workload and x... are the arguments inside the for loops. Note: on CPUs only the top level is parallelized.

for i in 1:Nx
  for j in 1:Ny
  ...
    f(x...) 
  end
end

basic "high-level" parallel_for APIs

Requires no familiarity with GPUs. Users don't need to pass any configuration, e.g. threads, blocks. JACC will make a "best guess". CPU will use the threads passed to Julia, for GPUs it will estimate blocks and threads.

• JACC.parallel_for( dims, f, x…) Follows Kokkos syntax Should we remove this pattern?. Example

• JACC.parallel_for( f, dims, x...) Follows Julia's map "function first" syntax. Should we prefer this signature?

• JACC.parallel_for(dims = N, args = (arg1, arg2, arg3, ...), f = function) keyword arguments based, order doesn't matter. Performance is not studied. Example Note: remove launch spec entries for a basic API.

advanced parallel_for APIs

Introduces a struct to add launch configurations for more advanced users. Most key/value entries target tuning capabilities for GPUs requiring some familiarity.

@kwdef mutable struct LaunchSpec{Backend}
    stream = default_stream(Backend)
    threads = 0
    blocks = 0
    shmem_size::Int = -1
    sync::Bool = true
end

• JACC.launch_spec(; kwargs) creates a JACC.LaunchSpec struct to tune parallel_for configurations
  spec = JACC.launch_spec(; blocks = (2,2), threads = (16,16) ) Example

• JACC.parallel_for( spec, dims, f, x…) lower-level, users need to pass configuration, e.g. threads/blocks. Follows Kokkos syntax. Example.

• JACC.parallel_for( f, spec, dims, x...) lower-level, users to pass configuration, e.g. threads/blocks. Follows Julia's map "function first" syntax. Should we favor this pattern? and Deprecated the Kokkos-order in the previous pattern?

• JACC.parallel_for(dims = N, args = (arg1, arg2, arg3, ...), f = function, blocks = blocks, threads = threads, sync = false) add LaunchSpec entries to dims, args, and function. Example.

parallel_reduce APIs

basic "high-level" parallel_reduce APIs

Requires no familiarity with GPUs. Users don't need to pass any configuration, e.g. threads, blocks. JACC will make a "best guess". CPU will use the threads passed to Julia, for GPUs it will estimate blocks and threads.

• red = JACC.parallel_reduce(a) high-level, minimal function returning the sum of all a elements, a[1] + .... + a[N]. Example.

• red = JACC.parallel_reduce(op, a), high-level, let user choose a supported operator = { +, *, min, max}. e.g. JACC.parallel_reduce(min, a) . Example.

• red = JACC.parallel_reduce(dims, f, x...) requires a size, function to declare the per-element reduced quantity, and arguments, default op is sum (+). Example

• red = JACC.parallel_reduce(dims, op, f, x...; init) requires size, operator, function to declare the per-element reduced quantity, arguments, and initial value.

• red = JACC.parallel_reduce(f, dims, x...) requires a size, function to declare the per-element reduced quantity, and arguments, default op is sum (+). Follows Julia's map function first syntax. Should we favor this pattern? Deprecate "dims" first?

• red = JACC.parallel_reduce(f, dims, op, x...; init) requires size, operator, function to declare the per-element reduced quantity, arguments, and initial value. Follow Julia's map function first syntax. Should we favor this pattern? Deprecate "dims" first?

• red = JACC.parallel_reduce(dims, dot, x1, x2) Special op=dot product reduction, requires two arrays x1, x2. Example

• red = JACC.parallel_reduce(; dims, f, args, [type, op, init,]) Keyword argument based function

advanced parallel_reduce APIs

Introduces a struct to add launch configurations for more advanced users. Most key/value entries target tuning capabilities for GPUs requiring some familiarity.

• red = JACC.parallel_reduce(spec, a) high-level, minimal function returning the sum of all a elements, a[1] + .... + a[N] . Example

• red = JACC.parallel_reduce(spec, op, a), high-level, let user choose a supported operator = { +, *, min, max}. e.g. JACC.parallel_reduce(min, a) Example

• red = JACC.parallel_reduce(spec, dims, f, x...) requires a size, function to declare the per-element reduced quantity, and arguments, default op is sum (+)

• red = JACC.parallel_reduce(spec, dims, op, f, x...; init) requires size, operator, function to declare the per-element reduced quantity, arguments, and initial value.

• red = JACC.parallel_reduce(f, spec, dims, x...) requires a size, function to declare the per-element reduced quantity, and arguments, default op is sum (+). Follows Julia's map function first syntax. Should we favor this pattern? Deprecate "dims" first?

• red = JACC.parallel_reduce(f, spec, dims, op, x...; init) requires size, operator, function to declare the per-element reduced quantity, arguments, and initial value. Follow Julia's map function first syntax. Should we favor this pattern? Deprecate "dims" first?

• red = JACC.parallel_reduce(spec, dims, dot, x1, x2) Special op=dot product reduction, requires two arrays x1, x2

• red = JACC.parallel_reduce(; dims, f, args, [type, op, init,] kw...) Keyword argument based function, kw... can be any of the same keyword arguments used in launch_spec(kw...) Example

synchronize

• JACC.synchronize() : synchronizes the default stream (if launched with sync=false). Example

• JACC.synchronize(stream) : synchronize the input stream (not portable to CPUs, if launched with sync=false)

Inside kernel capabilities APIs

These functions are provided by JACC and can used inside kernel functions passed to parallel_for and parallel_reduce

• JACC.shared(x) allows to bump an array into "shared" memory, e.g. block memory on GPUs for fast access. Note: only 1D arrays are returned, if passing a ND array it will be flattened to 1D . Example

• JACC.@atomic allows to declare an operation "atomic" (one at a time) leveraging the Atomix.jl support across backends. Careful must be taken due to parallelization, thus, performance losses. Example

[williamfgc]

Questions:

1. Should we keep or remove certain signatures. If so, which ones?
2. Should we modify the names or certain functionality to remove ambiguity?
3. Are we missing any critical (widely used) feature for common applications?
4. Suggestions?

[PhilipFackler]

@boussuge I like all of these suggestions. About your question in the last paragraph...it sounds like you are asking for a function to construct a stream. Is that correct? That was intended but long since forgotten. :) I appreciate you bringing this back up.

[PhilipFackler]

I would prefer symbols over strings, feels more Julian.

Symbols are already supported. We should make them primary in the documentation.

[PhilipFackler]

A JACC.similar(A) would be useful. I allocate temporaries all the time in CFD

Base.similar is already implemented by the vendors. So you can use it...

x = JACC.ones(M, N)
y = similar(x) # this will be the same type and size as x

[boussuge]

@PhilipFackler Thanks for your feedback.
Yes, I'm looking for a way to construct and manage streams explicitly.
The typical pattern I am looking for distributed CFD would be:
stream1 = JACC.Stream()
stream2 = JACC.Stream()

GPU: parallel execution
@parallel_for ... stream=stream1 compute_interior_faces!(...) # internal faces
@parallel_for ... stream=stream1 compute_bnd_faces!(...) # physical BCs (no MPI needed)
@parallel_for ... stream=stream2 pack_halos!(...)

--- All this while stream1 is still running ---
JACC.synchronize(stream2)
MPI.Isend(send_buf, ...)
MPI.Irecv!(recv_buf, ...)
MPI.Waitall(...)
@parallel_for ... stream=stream2 unpack_halos!(...)

JACC.synchronize(stream1)
@parallel_for ... compute_mpi_faces!(...) # faces between subdomains

This overlaps GPU computation, CPU/MPI work, and GPU transfers all at once.

[PhilipFackler]

@boussuge , I forgot to follow up here and wanted to make sure you are aware of this: #300. You can use it as stream1 = JACC.create_stream(). Until we have the macro API you can use it as a member of LaunchSpec or in the keyword version of parallel_* in the "advanced" API.

[quantumsteve]

Do you have examples or other documentation on how to combine JACC.jl with other parts of the JuliaGPU ecosystem?

[PhilipFackler]

The real answer is "no". 😄 In theory it should be easy, but we haven't spent time testing that. Two things come to mind...

1. The array types are not unique to JACC. So, if you want to use a CUDA function with an array created by JACC, you can do that because it is already a CuArray. Likewise you can use JACC.parallel_for with an array created with CUDA.
2. KernelAbstractions tends to use tags for dispatching to the backends. JACC.default_backend() returns a backend object you can use for this (except in the threads case).

I'm sorry we don't have existing concrete examples of this. Let me know if I can help.

[williamfgc]

Yes, JACC is metaprogramming, the connection with JuliaGPU is through "backend" dependencies that sit below JACC. Introducing vendor-specific functionality defeats the purpose of needing something like JACC, KA or any portable layer for that matter, so it's a bit out of scope unless absolutely required by a science application use-case.

[quantumsteve]

Sunny.jl uses cusolver (soon also rocsolver) and the intensity kernel uses shared memory for temporary storage.

[williamfgc]

@quantumsteve we have a JACC.BLAS module that is very early in the development. Usually when CUDA.jl and AMDGPU.jl have already support (like they do for Base.rand in the GrayScott example, or for Atomix.@atomic) it's a matter of wrapping around those. If a function is not supported then we prefer to work with the backends for long-term sustainability. Let us know if there is a path for vendor-neutral implementations on Sunny and we'll be glad to collaborate. We have a JACC.shared (portable) capability, paper is here.

[williamfgc]

That being said, please open an issue for tracking purposes. I think the request is at the kernel and backend-support levels, shouldn't affect the current v1.0 API discussion.

[diehlpk]

I think adding apple metal support would be good since iTensor.jl support it.

[williamfgc]

It's WIP :)

[williamfgc]

Please see issue #212 and PR #306

[diehlpk]

Requires no familiarity with GPUs. Users don't need to pass any configuration, e.g. threads, blocks. JACC will make a "best guess". CPU will use the threads passed to Julia, for GPUs it will estimate blocks and threads.

Is it possible to pass GPU parameters for more advanced users?

[williamfgc]

Yes, those are the LaunchSpec-based APIs marked as "advanced APIs"

[diehlpk]

JACC.parallel_for( dims, f, x…) Follows Kokkos syntax Should we remove this pattern?. Example

I would say that having the Kokkos signature might make it easier for Kokkos folks to use JACC.

[williamfgc]

Thank you all for your input and @PhilipFackler for the great work. We have a release candidate: https://github.com/JuliaORNL/JACC.jl/releases/tag/v1.0.0-rc1

Original APIs are kept.
Major additions:

1. Macro-based kernel-launching APIs: @parallel_for and @parallel_reduce see test
2. Apple backend support using Metal.jl (Float32 only)
3. Updated README with feature support

To do:
Update documentation

[luraess]

Exciting work!

A few thoughts:

• Explicit stream handling: GPU kernels handled by different tasks would also by default execute on different streams - how does this behaviour handled in JACC in order to ensure cirrus execution
• Implicit data sync on some backends: see Refactor memory wrappers CUDA.jl#2335 which may be relevant in cases such as Ability to opt out of / improved automatic synchronization between tasks for shared array usage CUDA.jl#2617
• Implicit sync in macro-based kernel-launching APIs: in cases where no data race can be ensured, there is no need to call into stream sync which would improve some compute efficiency. Having the ability to launch the kernel in an asynchronous fashion may be an interesting feature for better perf and addition flexibility.