Add GPUToolbox.Overlays with shared device overrides of Base

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUToolbox"
 uuid = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
-version = "3.0.0"
+version = "3.1.0"
 
 [deps]
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"

README.md

@@ -25,4 +25,8 @@ This package currently exports the following:
 - `@gcsafe_ccall`: like `@ccall` but marking it safe for the GC to run.
 -  Suffix literals like `1i8` (`i8`, `i16`, `i32`, `u8`, `u16`, `u32`) for constructing literals of a certain type.
 
+It also provides `GPUToolbox.Overlays`, method tables with device overrides of Base
+functionality that GPU back-ends can share by stacking them underneath their own method
+table. See the module's docstring for the design and how to use it.
+
 For more details on a specific symbol, check out its docstring in the Julia REPL.

docs/make.jl

@@ -16,6 +16,7 @@ makedocs(;
     ),
     pages = [
         "Home" => "index.md",
+        "Overlays" => "overlays.md",
     ],
     doctest = true,
     linkcheck = true,

docs/src/index.md

@@ -4,4 +4,5 @@
 
 ```@autodocs
 Modules = [GPUToolbox]
+Filter = t -> t !== GPUToolbox.Overlays
 ```

docs/src/overlays.md

@@ -0,0 +1,7 @@
+# Overlays
+
+```@docs
+GPUToolbox.Overlays
+GPUToolbox.Overlays.float64_overrides
+GPUToolbox.Overlays.audit
+```

src/GPUToolbox.jl

@@ -9,5 +9,6 @@ include("literals.jl")
 include("enum.jl")
 include("threading.jl")
 include("memoization.jl")
+include("overlays.jl")
 
 end # module GPUToolbox

src/overlays.jl

@@ -0,0 +1,131 @@
+"""
+    GPUToolbox.Overlays
+
+Method tables with device overrides of Base functionality, shared by GPU back-ends.
+
+Some of Base's methods are unsuited for GPUs, e.g., because they compute single-precision
+results in double precision, which many GPUs don't support. Back-ends replace such methods
+with overlay methods in their own method table. The tables in this module collect overrides
+that aren't specific to a back-end, so that back-ends can share them by stacking them
+underneath their own method table.
+
+# Usage
+
+A back-end uses a table by stacking it underneath its own method table when it creates the
+method table view for a compiler job, e.g.:
+
+```julia
+function GPUCompiler.method_table_view(job::MyCompilerJob)
+    if job.config.target.supports_fp64
+        GPUCompiler.StackedMethodTable(job.world, method_table)
+    else
+        GPUCompiler.StackedMethodTable(job.world, method_table,
+                                       GPUToolbox.Overlays.float64_overrides)
+    end
+end
+```
+
+As shown here, the tables to use can depend on the target device. The choice may only
+depend on the job's target and parameters, though, as GPUCompiler shares inference results
+between jobs with equal targets and parameters.
+
+# Lookup order
+
+The tables in a stack are searched from the top: the first table with a method that covers
+the signature of a call provides the implementation, even if a table further down the
+stack, or Base, has a more specific method. The stack should be ordered from the most to
+the least device-specific knowledge:
+
+1. the back-end's own method table, with implementations that use the device's hardware or
+   vendor libraries;
+2. tables shared by a family of back-ends, like the one from SPIRVIntrinsics.jl, used by the
+   back-ends targeting OpenCL environments;
+3. the tables in this module;
+4. Base.
+
+Because of that, the overrides here only act as fallbacks: a back-end replaces one by
+defining its own method with the same signature. Calls made *by* an override are again
+looked up from the top of the stack, so shared overrides use the back-end's implementations
+of the functions they call.
+
+# Rules
+
+To keep the order of the stack from affecting correctness, overrides in this module:
+
+- implement the Base method they replace completely and correctly for their signature, so
+  that the back-end's methods only have to be better, not different;
+- consist of generic Julia code, without calls to device-specific functionality;
+- use concrete signatures, so as not to hide more specific methods further down the stack;
+- replace a method in Base, or extend a Base function for a type defined in this module
+  (checked by [`audit`](@ref));
+- don't overlap with the overrides in other tables in this module, so that the relative
+  order of these tables doesn't matter;
+- document which functions they expect the stack to provide (e.g., `Float64`-free
+  elementary functions).
+
+Back-ends can check their stack with [`audit`](@ref), e.g., as part of their tests.
+
+# Tables
+
+- [`float64_overrides`](@ref): overrides of Base methods that use `Float64` to compute
+  single- or half-precision results.
+"""
+module Overlays
+
+include("overlays/float64.jl")
+
+
+## checking
+
+"""
+    Overlays.audit(tables::Core.MethodTable...; world=Base.get_world_counter())
+
+Check a stack of method tables, ordered from highest to lowest priority, and return a
+vector of `(; table, method, issue, by)` named tuples describing problems with the methods
+in these tables:
+
+- `issue = :dead`: the method overrides a function from Base, but no method of that function
+  matches its signature (e.g., because Base changed the function's signature). Methods of
+  other functions, and methods for types defined in the table's module, aren't checked.
+- `issue = :shadowed` or `issue = :partly_shadowed`: a method in table `by`, higher up the
+  stack, covers all or part of the method's signature, so that (some) calls won't reach it.
+  This is often intentional, e.g., when a back-end replaces a shared fallback.
+"""
+function audit(tables::Core.MethodTable...; world::UInt=Base.get_world_counter())
+    issues = @NamedTuple{table::Core.MethodTable, method::Method, issue::Symbol,
+                         by::Union{Nothing,Core.MethodTable}}[]
+    for (i, table) in enumerate(tables)
+        methods = Method[]
+        Base.visit(m -> push!(methods, m), table)
+        for method in methods
+            sig = method.sig
+            for higher in tables[1:i-1]
+                matches = Base._methods_by_ftype(sig, higher, -1, world)
+                (matches === nothing || isempty(matches)) && continue
+                issue = any(m -> m.fully_covers, matches) ? :shadowed : :partly_shadowed
+                push!(issues, (; table, method, issue, by=higher))
+            end
+            if replaces_base_function(sig, table) &&
+               isempty(Base._methods_by_ftype(sig, nothing, -1, world))
+                push!(issues, (; table, method, issue=:dead, by=nothing))
+            end
+        end
+    end
+    return issues
+end
+
+# does this signature override a function from Base, rather than extend it with methods
+# for types owned by the overlay table's module?
+function replaces_base_function(@nospecialize(sig), table::Core.MethodTable)
+    params = Base.unwrap_unionall(sig).parameters
+    ft = params[1]
+    isdefined(ft, :instance) || return false
+    Base.moduleroot(parentmodule(ft.instance)) in (Base, Core) || return false
+    for T in params[2:end]
+        T = Base.unwrap_unionall(T)
+        T isa DataType && parentmodule(T) === table.module && return false
+    end
+    return true
+end
+
+end # module Overlays

src/overlays/float64.jl

@@ -0,0 +1,180 @@
+# Overrides of Base methods that use Float64 to compute single- and half-precision results
+
+Base.Experimental.@MethodTable(float64_overrides)
+
+"""
+    GPUToolbox.Overlays.float64_overrides
+
+Method table with overrides of Base methods that use `Float64` to compute single- or
+half-precision results, for devices that don't support `Float64` or where it is slow.
+
+The overrides only replace Base's generic code, and call functions like `sin`, `cos`,
+`sinpi`, `cospi` and `^(::Float32, ::Float32)`, which Base also implements using `Float64`.
+Back-ends stacking this table should provide `Float64`-free implementations of those (e.g.,
+using their hardware's math library) in a table higher up the stack.
+
+See [`GPUToolbox.Overlays`](@ref) for how to use this table.
+"""
+float64_overrides
+
+macro float64_override(ex)
+    esc(:(Base.Experimental.@overlay($float64_overrides, $ex)))
+end
+
+# float.jl: comparisons of Float32 and 32-bit integers are performed in Float64. Widening
+# the integer instead reuses the exact Float32/Int64 comparisons.
+for op in (:(==), :<, :<=)
+    @eval begin
+        @float64_override Base.$op(x::Float32, y::Union{Int32,UInt32}) = $op(x, Int64(y))
+        @float64_override Base.$op(x::Union{Int32,UInt32}, y::Float32) = $op(Int64(x), y)
+        @float64_override Base.$op(x::Float16, y::Union{Int32,UInt32,Int64,UInt64}) =
+            $op(Float32(x), y)
+        @float64_override Base.$op(x::Union{Int32,UInt32,Int64,UInt64}, y::Float16) =
+            $op(x, Float32(y))
+    end
+end
+
+# div.jl: Julia 1.12 and 1.13 divide Float32 values in Float64 (JuliaLang/julia#49637).
+# Use the generic implementation that replaced it (JuliaLang/julia#60497), which is exact
+# as long as `eps(x/y) <= 1`, but keep Base's results for non-finite operands and zero
+# quotients. `rem` does not support rounding ties away from zero or up, so handle those
+# by rounding ties to even first.
+@static if v"1.12-" <= VERSION < v"1.14-"
+    @float64_override function Base.div(x::Float32, y::Float32,
+                                        r::Union{RoundingMode{:ToZero}, RoundingMode{:Down},
+                                                 RoundingMode{:Up}, RoundingMode{:Nearest},
+                                                 RoundingMode{:FromZero}})
+        q = x / y
+        (isfinite(x) & isfinite(y) & !iszero(y)) || return round(q, r)
+        d = round(q - rem(x, y, r) / y)
+        return iszero(d) ? copysign(d, q) : d
+    end
+    @float64_override function Base.div(x::Float32, y::Float32,
+                                        r::Union{RoundingMode{:NearestTiesAway},
+                                                 RoundingMode{:NearestTiesUp}})
+        (isfinite(x) & isfinite(y) & !iszero(y)) || return round(x / y, r)
+        d = div(x, y, RoundNearest)
+        m = rem(x, y, RoundNearest)
+        2 * abs(m) == abs(y) || return d
+        # the quotient is halfway between `d` and this other integer
+        e = signbit(m) == signbit(y) ? d + 1f0 : d - 1f0
+        return r === RoundNearestTiesUp ? max(d, e) : ifelse(abs(e) > abs(d), e, d)
+    end
+end
+
+# math.jl: the power is computed by squaring in Float64. Use the float power instead,
+# splitting exponents that aren't exactly representable as Float32 in two parts, and
+# restoring the sign of the result for odd exponents.
+@float64_override function Base.:(^)(x::Float32, n::Integer)
+    n == -2 && return (i = inv(x); i * i)
+    n == -1 && return inv(x)
+    n == 0 && return one(x)
+    n == 1 && return x
+    n == 2 && return x * x
+    n == 3 && return x * x * x
+    # literal bounds: integer powers don't constant-fold in device code
+    y = if -16777216 <= n <= 16777216
+        abs(x)^Float32(n)
+    else
+        lo = rem(n, 65536)
+        abs(x)^Float32(n - lo) * abs(x)^Float32(lo)
+    end
+    return isodd(n) ? copysign(y, x) : y
+end
+
+# math.jl: Float32 `hypot` is computed in Float64. Use the generic implementation.
+@float64_override Base.Math._hypot(x::Float32, y::Float32) =
+    invoke(Base.Math._hypot, Tuple{Any,Any}, x, y)
+
+# special/trig.jl: `sind` and friends reduce their argument to [-45°, 45°] in Float32, but
+# convert it to radians and evaluate the sine and cosine kernels in Float64. Keep the
+# reduction, and wrap the converted argument so that the kernels use `sin` and `cos`.
+struct Radians32
+    x::Float32
+end
+# π/180 split into a Float32 value and the remainder
+@float64_override Base.Math.deg2rad_ext(x::Float32) =
+    Radians32(muladd(x, 0.017453292f0, x * 1.3519961f-10))
+@float64_override Base.Math.sin_kernel(y::Radians32) = sin(y.x)
+@float64_override Base.Math.cos_kernel(y::Radians32) = cos(y.x)
+
+# special/trig.jl: `sincospi` uses Float64 kernels (as do `sinpi` and `cospi`, which the
+# back-end is expected to override).
+@float64_override Base.sincospi(x::Float32) = (sinpi(x), cospi(x))
+
+# complex.jl: single-precision complex numbers are divided and inverted in double precision.
+# Divide with the robust algorithm Base uses for ComplexF64 instead (Baudin & Smith, 2012),
+# but scale large divisors further so that the reciprocal of their magnitude isn't subnormal
+# (which GPUs may flush to zero). Dividing by a divisor with a positive dominant component,
+# and computing the imaginary part without negating it, gives zeros the signs of Base's
+# ComplexF32 results.
+const TWO_M8 = Float32(0x1p-8)
+const TWO_47 = Float32(0x1p47)
+const TWO_M47 = Float32(0x1p-47)
+@inline function robust_cdiv2(a::Float32, b::Float32, c::Float32, d::Float32, r::Float32,
+                              t::Float32)
+    if r != 0
+        br = b * r
+        return br != 0 ? (a + br) * t : a * t + (b * t) * r
+    else
+        # `b / c` can overflow when `d` is zero; the term is then a signed zero (`c > 0`)
+        return (a + (iszero(d) ? flipsign(d, b) : d * (b / c))) * t
+    end
+end
+@float64_override function Base.:(/)(z::ComplexF32, w::ComplexF32)
+    a, b = reim(z)
+    c, d = reim(w)
+    if (isinf(c) | isinf(d))
+        isfinite(z) && return complex(0f0 * sign(a) * sign(c), -0f0 * sign(b) * sign(d))
+        return complex(NaN32, NaN32)
+    end
+    absa, absb, absc, absd = abs(a), abs(b), abs(c), abs(d)
+    ab = absa >= absb ? absa : absb
+    cd = absc >= absd ? absc : absd
+    if signbit(absd <= absc ? c : d)
+        a, b, c, d = -a, -b, -c, -d
+    end
+
+    s = 1f0
+    if ab >= floatmax(Float32) / 2
+        a *= 0.5f0; b *= 0.5f0; s *= 2f0
+    elseif ab <= 2floatmin(Float32) / eps(Float32)
+        a *= TWO_47; b *= TWO_47; s *= TWO_M47
+    end
+    if cd >= floatmax(Float32) * TWO_M8
+        c *= TWO_M8; d *= TWO_M8; s *= TWO_M8
+    elseif cd <= 2floatmin(Float32) / eps(Float32)
+        c *= TWO_47; d *= TWO_47; s *= TWO_47
+    end
+
+    if absd <= absc
+        r = d / c
+        t = 1f0 / (c + d * r)
+        p, q = robust_cdiv2(a, b, c, d, r, t), robust_cdiv2(b, -a, c, d, r, t)
+    else
+        r = c / d
+        t = 1f0 / (d + c * r)
+        p, q = robust_cdiv2(b, a, d, c, r, t), robust_cdiv2(-a, b, d, c, r, t)
+    end
+    return Complex(p * s, q * s)
+end
+# Invert with Smith's algorithm, after scaling by a power of two so that the reciprocal
+# neither overflows nor becomes subnormal.
+const TWO_64 = Float32(0x1p64)
+const TWO_M64 = Float32(0x1p-64)
+@inline scale_factor(m::Float32) = m >= TWO_64 ? TWO_M64 : m <= TWO_M64 ? TWO_64 : 1f0
+@float64_override function Base.inv(w::ComplexF32)
+    c, d = reim(w)
+    (isinf(c) | isinf(d)) && return complex(copysign(0f0, c), flipsign(-0f0, d))
+    s = scale_factor(max(abs(c), abs(d)))
+    c, d = c * s, d * s
+    if abs(d) <= abs(c)
+        r = d / c
+        t = inv(muladd(d, r, c))
+        return Complex(t * s, -r * t * s)
+    else
+        r = c / d
+        t = inv(muladd(c, r, d))
+        return Complex(r * t * s, -t * s)
+    end
+end