EnzymeAD/Reactant.jl

Experiment with `Core.Memory` address space

Open

#181 opened on Oct 16, 2024

View on GitHub
 (0 comments) (1 reaction) (0 assignees)Julia (64 forks)github user discovery
help wanted

Repository metrics

Stars
 (336 stars)
PR merge metrics
 (Avg merge 2d 10h) (36 merged PRs in 30d)

Description

Julia 1.11 introduced Core.Memory and made Array a purely Julia type, which is nice for different optimization reasons. The interesting part for us is that GenericMemory has a parameter for kind and another for addrspace, which could potentially be exploited by us to store the memory managed by XLA. This means that ConcreteRArrays could be managed by a simple Array instead, with the address space managed by Reactant.

Docs say the following:

addrspace can currently only be set to Core.CPU. It is designed to permit extension by other systems such as GPUs, which might define values such as:

module CUDA
const Generic = bitcast(Core.AddrSpace{CUDA}, 0)
const Global = bitcast(Core.AddrSpace{CUDA}, 1)
end

The exact semantics of these other addrspaces is defined by the specific backend, but will error if the user is attempting to access these on the CPU.

A lot of info is missing on how to customize this behavior but we should keep an eye on it.

Contributor guide

Research direction
Explore the Core.Memory address space API in Julia, investigate how to define custom address spaces, and design an integration path for Reactant's memory management via XLA.
Tech stack
julia
Domain
backendperformance
Issue type
Research
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
Half day
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Fresh
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Mostly clear
Prerequisites
Julia
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
30