Is it possible to use @static to distinguish CPU or GPU call in CUDA?

Specific Domains GPU
1

As we all know, some codes can not run on GPU. But there will be some alternative way to solve it.

For example, sqrt can not run on GPU due to a _fpow error. But CUDA.sqrt is an alternative solution. In this condition, 2 methods will be written like the following code.

my_sqrt(x) = sqrt(x)
my_sqrt_cuda(x) = CUDA.sqrt(x)

When some upper code call this method, it have to duplicated to adapt the duplicated my_sqrt.

In C++, #if could be used to treat this problem. In julia, a @static macro also could be used to solve operating system difference.

Can I detect an expression to distinguish the device (CPU or GPU) the code is running on? If so, the code should be like:

my_sqrt(x) = @static magic_expression ? CUDA.sqrt(x) : sqrt(x)

And this method could be call on either CPU or GPU condition.

Update: I found that sqrt has been fixed in CUDA.jl. So maybe power or something else is a better sample.

2

I guess, the correct way of writing generic CPU/GPU code is using Cassette.jl. Take a look how it is done in, e.g. KernelAbstractions.jl:
https://github.com/JuliaGPU/KernelAbstractions.jl/blob/master/src/backends/cuda.jl#L240

2 Likes
3

That’s so cool!