add support for Float16, UInt8, UInt16, Int8, Int16 (was bored)

Author: fjbarter

Project.toml

@@ -1,7 +1,7 @@
 name = "AcceleratedKernels"
 uuid = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
-version = "0.4.3"
 authors = ["Andrei-Leonard Nicusan <leonard@evophase.co.uk> and contributors"]
+version = "0.4.3"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

docs/src/api/rand.md

@@ -5,18 +5,18 @@ Counter-based random generation for CPU and GPU backends with deterministic beha
 
 Use an explicit `CounterRNG(seed; alg=...)` when reproducibility matters. For convenience,
 `AK.rand!(x)` creates a fresh `CounterRNG()` on each call using one auto-seeded
-`Base.rand(Random.default_rng(), UInt64)` draw, so repeated calls produce different outputs unless Random.seed!() is used.
+`Random.rand(Random.default_rng(), UInt64)` draw, so repeated calls produce different outputs unless Random.seed!() is used.
 
-Supported output element types:
-- `UInt32`, `UInt64`
-- `Int32`, `Int64`
-- `Float32`, `Float64`
+Supported element types:
+- `UInt8`, `UInt16`, `UInt32`, `UInt64`
+- `Int8`, `Int16`, `Int32`, `Int64`
+- `Float16`, `Float32`, `Float64`
 - `Bool`
 
-The core of the random number generation produces a `UInt` of the requested scalar width.
+The core of the random number generation produces either a `UInt32` or `UInt64` depending on the width of the requested element type.
 That `UInt` is then either:
-- Unsigned integers: returned as-is
-- Signed integers: reinterpreted as a signed integer bit pattern.
+- Unsigned integers: returned as-is or truncated if necessary.
+- Signed integers: reinterpreted as a signed integer bit pattern and truncated if necessary.
 - Floats: mantissa construction into a uniform grid in `[0, 1)` ([read more](https://lomont.org/posts/2017/unit-random/)).
 - Bool: `true` if the `UInt` draw is odd (`isodd(u)`), otherwise `false`.
 

src/rand/rand.jl

@@ -82,9 +82,9 @@ Fill `x` in-place with pseudo-random values using a stateless counter-based RNG.
 counter is exactly `UInt64(i - 1)` in linear indexing order.
 
 Supported scalar element types are:
-- `UInt32`, `UInt64`
-- `Int32`, `Int64`
-- `Float32`, `Float64`
+- `UInt8`, `UInt16`, `UInt32`, `UInt64`
+- `Int8`, `Int16`, `Int32`, `Int64`
+- `Float16`, `Float32`, `Float64`
 - `Bool`
 
 Semantics:

src/rand/utilities.jl

@@ -15,26 +15,36 @@
 
 # Internal scalar eltypes currently supported by rand!.
 const ALLOWED_RAND_SCALARS = Union{
-    UInt32, UInt64,
-    Int32, Int64,
-    Float32, Float64,
+    UInt8, UInt16, UInt32, UInt64,
+    Int8, Int16, Int32, Int64,
+    Float16, Float32, Float64,
     Bool
 }
 
 
+@inline raw_uint_type(::Type{UInt8}) = UInt32
+@inline raw_uint_type(::Type{UInt16}) = UInt32
 @inline raw_uint_type(::Type{UInt32}) = UInt32
+@inline raw_uint_type(::Type{Int8}) = UInt32
+@inline raw_uint_type(::Type{Int16}) = UInt32
 @inline raw_uint_type(::Type{Int32}) = UInt32
+@inline raw_uint_type(::Type{Float16}) = UInt32
 @inline raw_uint_type(::Type{Float32}) = UInt32
 @inline raw_uint_type(::Type{UInt64}) = UInt64
 @inline raw_uint_type(::Type{Int64}) = UInt64
 @inline raw_uint_type(::Type{Float64}) = UInt64
 @inline raw_uint_type(::Type{Bool}) = UInt32
 
 
+@inline from_uint(::Type{UInt8}, u::UInt32)::UInt8 = trunc(UInt8, u >> 24)
+@inline from_uint(::Type{UInt16}, u::UInt32)::UInt16 = trunc(UInt16, u >> 16)
 @inline from_uint(::Type{UInt32}, u::UInt32)::UInt32 = u
 @inline from_uint(::Type{UInt64}, u::UInt64)::UInt64 = u
+@inline from_uint(::Type{Int8}, u::UInt32)::Int8 = reinterpret(Int8, trunc(UInt8, u >> 24))
+@inline from_uint(::Type{Int16}, u::UInt32)::Int16 = reinterpret(Int16, trunc(UInt16, u >> 16))
 @inline from_uint(::Type{Int32}, u::UInt32)::Int32 = reinterpret(Int32, u)
 @inline from_uint(::Type{Int64}, u::UInt64)::Int64 = reinterpret(Int64, u)
+@inline from_uint(::Type{Float16}, u::UInt32)::Float16 = uint32_to_unit_float16(u)
 @inline from_uint(::Type{Float32}, u::UInt32)::Float32 = uint32_to_unit_float32(u)
 @inline from_uint(::Type{Float64}, u::UInt64)::Float64 = uint64_to_unit_float64(u)
 @inline from_uint(::Type{Bool}, u::UInt32)::Bool = isodd(u)
@@ -81,6 +91,17 @@ end
 
 
 
+# Convert random UInt32 bits to Float16 in [0, 1) by mantissa construction.
+@inline function uint32_to_unit_float16(u::UInt32)::Float16
+    # Keep 10 random bits for the mantissa (drop 22 rightmost bits from the UInt32)
+    # and combine with the bit pattern of Float16(1.0) (sign=0, exponent=15).
+    bits = UInt16(0x3c00) | UInt16(u >> 22)
+
+    # Interpret as 1.mantissa, then subtract 1 for [0, 1)
+    return reinterpret(Float16, bits) - Float16(1)
+end
+
+
 # Convert random UInt32 bits to Float32 in [0, 1) by mantissa construction.
 @inline function uint32_to_unit_float32(u::UInt32)::Float32
     # Keep 23 random bits for the mantissa (drop 9 rightmost bits from the UInt32)

test/rand.jl

@@ -1,8 +1,13 @@
 const RAND_ALGS = (AK.SplitMix64(), AK.Philox(), AK.Threefry())
-const RAND_SCALAR_TYPES_ALL = (UInt32, UInt64, Int32, Int64, Float32, Float64, Bool)
+const RAND_SCALAR_TYPES_ALL = (
+    UInt8, UInt16, UInt32, UInt64,
+    Int8, Int16, Int32, Int64,
+    Float16, Float32, Float64,
+    Bool,
+)
 const RAND_SCALAR_TYPES_BACKEND = IS_CPU_BACKEND ?
                                   RAND_SCALAR_TYPES_ALL :
-                                  (UInt32, UInt64, Int32, Int64, Float32, Bool)
+                                  (UInt8, UInt16, UInt32, UInt64, Int8, Int16, Int32, Int64, Float16, Float32, Bool)
 const RUN_FLOAT64_RAND_TESTS = IS_CPU_BACKEND
 
 
@@ -62,8 +67,13 @@ end
         @test AK._counter_from_index(1) == UInt64(0)
         @test AK._counter_from_index(17) == UInt64(16)
 
+        @test AK.raw_uint_type(UInt8) === UInt32
+        @test AK.raw_uint_type(UInt16) === UInt32
         @test AK.raw_uint_type(UInt32) === UInt32
+        @test AK.raw_uint_type(Int8) === UInt32
+        @test AK.raw_uint_type(Int16) === UInt32
         @test AK.raw_uint_type(Int32) === UInt32
+        @test AK.raw_uint_type(Float16) === UInt32
         @test AK.raw_uint_type(Float32) === UInt32
         @test AK.raw_uint_type(UInt64) === UInt64
         @test AK.raw_uint_type(Int64) === UInt64
@@ -72,15 +82,22 @@ end
             @test AK.raw_uint_type(Float64) === UInt64
         end
 
+        @test AK.from_uint(UInt8, UInt32(0xabcdef01)) == UInt8(0xab)
+        @test AK.from_uint(UInt16, UInt32(0xabcdef01)) == UInt16(0xabcd)
         @test AK.from_uint(UInt32, 0b1010 % UInt32) == 0b1010 % UInt32
         @test AK.from_uint(UInt64, 0b1010 % UInt64) == 0b1010 % UInt64
+        @test AK.from_uint(Int8, UInt32(0xff000000)) == Int8(-1)
+        @test AK.from_uint(Int16, UInt32(0xffff0000)) == Int16(-1)
         @test AK.from_uint(Int32, 0b11111111111111111111111111111111 % UInt32) == Int32(-1)
         @test AK.from_uint(
             Int64, 0b1111111111111111111111111111111111111111111111111111111111111111 % UInt64
         ) == Int64(-1)
+        @test AK.from_uint(Float16, UInt32(0)) == Float16(0)
         @test AK.from_uint(Bool, UInt32(0)) == false
         @test AK.from_uint(Bool, UInt32(1)) == true
 
+        @test AK.uint32_to_unit_float16(UInt32(0)) == Float16(0)
+        @test Float16(0) <= AK.uint32_to_unit_float16(typemax(UInt32)) < Float16(1)
         @test AK.uint32_to_unit_float32(UInt32(0)) == 0.0f0
         @test 0.0f0 <= AK.uint32_to_unit_float32(typemax(UInt32)) < 1.0f0
         if RUN_FLOAT64_RAND_TESTS
@@ -126,7 +143,7 @@ end
             s1 = AK.rand_scalar(rng, UInt64(1), T)
             @test s0 isa T
             @test s1 isa T
-            if T !== Bool
+            if !(T in (Bool, Float16, UInt8, UInt16, Int8, Int16))
                 @test s0 != s1
             end
             if T <: AbstractFloat
@@ -136,8 +153,19 @@ end
         end
 
         c = UInt64(42)
+        @test AK.rand_scalar(rng, c, UInt8) == trunc(UInt8, AK.rand_uint(rng, c, UInt32) >> 24)
+        @test AK.rand_scalar(rng, c, UInt16) == trunc(UInt16, AK.rand_uint(rng, c, UInt32) >> 16)
+        @test AK.rand_scalar(
+            rng, c, Int8
+        ) == reinterpret(Int8, trunc(UInt8, AK.rand_uint(rng, c, UInt32) >> 24))
+        @test AK.rand_scalar(
+            rng, c, Int16
+        ) == reinterpret(Int16, trunc(UInt16, AK.rand_uint(rng, c, UInt32) >> 16))
         @test AK.rand_scalar(rng, c, Int32) == reinterpret(Int32, AK.rand_uint(rng, c, UInt32))
         @test AK.rand_scalar(rng, c, Int64) == reinterpret(Int64, AK.rand_uint(rng, c, UInt64))
+        @test AK.rand_scalar(rng, c, Float16) == AK.uint32_to_unit_float16(
+            AK.rand_uint(rng, c, UInt32)
+        )
         @test AK.rand_scalar(rng, c, Float32) == AK.uint32_to_unit_float32(
             AK.rand_uint(rng, c, UInt32)
         )
@@ -150,7 +178,7 @@ end
         bools = [AK.rand_scalar(rng, UInt64(i), Bool) for i in 0:511]
         @test any(identity, bools)
         @test any(!, bools)
-        @test_throws ArgumentError AK.rand_scalar(rng, UInt64(0), UInt16)
+        @test_throws ArgumentError AK.rand_scalar(rng, UInt64(0), UInt128)
     end
 
 
@@ -185,7 +213,7 @@ end
             @test Array(x1) != Array(x2)
         end
 
-        for T in (Float32, UInt64, Bool)
+        for T in (Float16, Float32, UInt64, Bool)
             xnd = array_from_host(zeros(T, 7, 11, 5))
             _assert_rand_matches_reference!(rng, xnd; prefer_threads, block_size=128)
         end
@@ -226,7 +254,7 @@ end
         @test Array(x1) == Array(ref1)
         @test Array(x2) == Array(ref2)
 
-        x_bad = array_from_host(zeros(UInt16, 16))
+        x_bad = zeros(UInt128, 16)
         @test_throws ArgumentError AK.rand!(x_bad; prefer_threads)
         @test_throws ArgumentError AK.rand!(AK.CounterRNG(0x1), x_bad; prefer_threads)
     end