|
| 1 | +## |
| 2 | +# Implements contextual dispatch through Cassette.jl |
| 3 | +# Goals: |
| 4 | +# - Rewrite common CPU functions to appropriate GPU intrinsics |
| 5 | +# |
| 6 | +# TODO: |
| 7 | +# - error (erf, ...) |
| 8 | +# - pow |
| 9 | +# - min, max |
| 10 | +# - mod, rem |
| 11 | +# - gamma |
| 12 | +# - bessel |
| 13 | +# - distributions |
| 14 | +# - unsorted |
| 15 | + |
| 16 | +using Cassette |
| 17 | + |
| 18 | +function transform(ctx, ref) |
| 19 | + ci = ref.code_info |
| 20 | + noinline = any(@nospecialize(x) -> Core.Compiler.isexpr(x, :meta) && x.args[1] == :noinline, ci.code) |
| 21 | + if !noinline |
| 22 | + ci.inlineable = true |
| 23 | + end |
| 24 | + return ci |
| 25 | +end |
| 26 | +const InlinePass = Cassette.@pass transform |
| 27 | + |
| 28 | +Cassette.@context CUDACtx |
| 29 | +const cudactx = Cassette.disablehooks(CUDACtx(pass = InlinePass)) |
| 30 | + |
| 31 | +Cassette.overdub(::CUDACtx, ::typeof(datatype_align), ::Type{T}) where {T} = datatype_align(T) |
| 32 | +Cassette.overdub(ctx::CUDACtx, ::typeof(isdevice)) = true |
| 33 | + |
| 34 | +# libdevice.jl |
| 35 | +for f in (:cos, :cospi, :sin, :sinpi, :tan, |
| 36 | + :acos, :asin, :atan, |
| 37 | + :cosh, :sinh, :tanh, |
| 38 | + :acosh, :asinh, :atanh, |
| 39 | + :log, :log10, :log1p, :log2, |
| 40 | + :exp, :exp2, :exp10, :expm1, :ldexp, |
| 41 | + :isfinite, :isinf, :isnan, |
| 42 | + :signbit, :abs, |
| 43 | + :sqrt, :cbrt, |
| 44 | + :ceil, :floor,) |
| 45 | + @eval function Cassette.overdub(ctx::CUDACtx, ::typeof(Base.$f), x::Union{Float32, Float64}) |
| 46 | + @Base._inline_meta |
| 47 | + return CUDAnative.$f(x) |
| 48 | + end |
| 49 | +end |
| 50 | + |
| 51 | +contextualize(f::F) where F = (args...) -> Cassette.overdub(cudactx, f, args...) |
| 52 | + |
0 commit comments