EnzymeAD / EnzymeAD/Reactant.jl

Overlay methods on top of CUDA.jl method table

Open
#2,131 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Julia
Stars
370
Forks
74
Avg merge
18h 47m
Merged PRs (30d)
30

Description

In Oceananigans we need to define a specialised method for the CUDA target (it uses a PTX intrinsic) only and in https://github.com/CliMA/Oceananigans.jl/pull/5140 we are going to use `CUDA.@device_override`, but the problem is that Reactant uses CUDA.jl-generated code for all targets, including CPU one, so it'd always try to call that intrinsic with all backends, which is of course not possible

> so reactant currently faithfully fully represents the cuda kernel. There is no mechanism for determining if you're in reactant or not within a kernel -- we get the code out of cuda.jl. We could add a mechanism for doing so here: https://github.com/EnzymeAD/Reactant.jl/blob/730804beb548017e819b3589b9f6f793a4f19792/ext/ReactantCUDAExt.jl#L1332. We'd need to make a new gpucompiler job/config/state that looks identical to the cuda one, except with a second method overlay table that itself overlays the cuda overlay table

_Originally posted by @wsmoses in https://github.com/CliMA/Oceananigans.jl/pull/5140#discussion_r2696303075_

I don't really know how to accomplish this though. The method table for a compiler job is specified via dispatch, e.g. https://github.com/JuliaGPU/CUDA.jl/blob/9528a336f527bf49cc2255b51bbb59f87a533e43/src/compiler/compilation.jl#L52, so I don't immediately see how we can have same [compilerjob](https://github.com/JuliaGPU/CUDA.jl/blob/9528a336f527bf49cc2255b51bbb59f87a533e43/src/compiler/compilation.jl#L16)/[compilerconfiguration](https://github.com/JuliaGPU/CUDA.jl/blob/9528a336f527bf49cc2255b51bbb59f87a533e43/src/compiler/compilation.jl#L15) as CUDA.jl, but different method table, unless you want to wrap `CUDACompilerJob`, but that's a endless rabbit hole. Can we overlay `GPUCompiler.method_table`? That sounds potentially "dangerous" though.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.