Add the CUDAChaChaStream type

Add a new type, CUDAChaChaStream, for accessing the keystream of a
ChaCha cipher on the GPU. This type can be used for GPU-based CRNG and
encryption.

Author: kernelmethod

src/ChaChaCiphers.jl

@@ -4,10 +4,12 @@ include("ChaCha.jl")
 
 include("core.jl")
 include("keystream.jl")
+include("cuda_keystream.jl")
 include("generation.jl")
 
 export ChaCha
 export ChaChaStream, ChaCha20Stream, ChaCha12Stream
+export CUDAChaChaStream, CUDAChaCha20Stream, CUDAChaCha12Stream
 export getstate
 export encrypt, decrypt, encrypt!, decrypt!
 

src/cuda_keystream.jl

@@ -0,0 +1,170 @@
+using Base: BitInteger
+using ChaChaCiphers.ChaCha
+using CUDA
+using StaticArrays
+
+"""
+    CUDAChaChaStream <: AbstractChaChaStream
+
+`CUDAChaChaStream` is a CUDA-compatible ChaCha keystream
+generator for GPU CRNG.
+
+## Examples
+
+Create a `CUDAChaChaStream` with a randomized key, and
+sample some random numbers with it:
+
+```@meta
+DocTestSetup = quote
+    using CUDA
+    using ChaChaCiphers
+    using Random
+end
+```
+
+```julia
+julia> rng = CUDAChaChaStream();
+
+julia> x = CuVector{Float32}(undef, 2^10);
+```
+
+```@meta
+DocTestSetup = nothing
+```
+
+"""
+mutable struct CUDAChaChaStream <: AbstractChaChaStream
+    key :: CuVector{UInt32}
+    nonce :: UInt64
+    counter :: UInt64
+    buffer :: CuVector{UInt8}
+    position :: Int
+    doublerounds :: Int
+
+    function CUDAChaChaStream(
+        key,
+        nonce,
+        counter = UInt64(0),
+        position = 1;
+        doublerounds = 10
+    )
+        if doublerounds ≤ 0
+            error("`doublerounds` must be a positive number")
+        end
+
+        key = CuVector{UInt32}(key)
+        buffer = CuVector{UInt8}(undef, STREAM_BUFFER_SIZE)
+        stream = new(key, nonce, counter, buffer, 1, doublerounds)
+        _refresh_buffer!(stream)
+        stream.position = position
+
+        stream
+    end
+end
+
+# Constructors
+
+"""
+    CUDAChaCha20Stream
+
+Create a CUDA-compatible keystream for a ChaCha20 stream
+cipher.
+"""
+CUDAChaCha20Stream(args...) = CUDAChaChaStream(args...; doublerounds=10)
+
+"""
+    CUDAChaCha12Stream
+
+Create a CUDA-compatible keystream for a ChaCha12 stream
+cipher.
+"""
+CUDAChaCha12Stream(args...) = CUDAChaChaStream(args...; doublerounds=6)
+
+function Base.show(io::IO, rng::CUDAChaChaStream)
+    msg = """
+        CUDAChaChaStream(
+            key = $(key(rng))
+            nonce = $(nonce(rng)),
+            counter = $(counter(rng)),
+            rounds = $(2 * doublerounds(rng))
+        )"""
+
+    write(io, msg)
+end
+
+buffer_size(stream::CUDAChaChaStream) =
+    length(stream.buffer) - stream.position + 1
+
+# Methods required for AbstractChaChaStream compatibility
+
+function key(stream::CUDAChaChaStream)
+    key_cpu = Vector{UInt32}(undef, 8)
+    copyto!(key_cpu, stream.key)
+    SVector{8,UInt32}(key_cpu)
+end
+
+@inline nonce(stream::CUDAChaChaStream) = stream.nonce
+@inline counter(stream::CUDAChaChaStream) = stream.counter
+@inline position(stream::CUDAChaChaStream) = stream.position
+@inline doublerounds(stream::CUDAChaChaStream) = stream.doublerounds
+
+function _refresh_buffer!(stream::CUDAChaChaStream)
+    _fill_blocks!(
+        stream.buffer,
+        stream,
+        STREAM_BUFFER_BLOCKS
+    )
+    stream.position = 1
+    stream
+end
+
+function _fill_buffer!(dest::CuVector{UInt8}, stream::CUDAChaChaStream)
+    bfsize = buffer_size(stream)
+    destsize = length(dest)
+
+    # If the internal buffer is larger than the destination size,
+    # we can just copy directly from the buffer to the stream and
+    # return
+    if bfsize >= destsize
+        copyto!(dest, 1, stream.buffer, stream.position, destsize)
+        stream.position += destsize
+        return dest
+    end
+
+    # Otherwise, the destination is larger than the buffer
+    copyto!(dest, 1, stream.buffer, stream.position, bfsize)
+
+    (n_blocks, rem) = divrem(length(dest) - bfsize, CHACHA_BLOCK_SIZE)
+    if n_blocks > 0
+        sp = pointer(dest, bfsize + 1)
+        slice = unsafe_wrap(CuVector{UInt8}, sp, n_blocks * CHACHA_BLOCK_SIZE)
+        _fill_blocks!(slice, stream, n_blocks)
+    end
+
+    # Refresh the stream, and then copy the stream buffer into the
+    # remainder of the destination
+    _refresh_buffer!(stream)
+    _fill_buffer!(view(dest, length(dest)-rem+1:length(dest)), stream)
+
+    dest
+end
+
+function _fill_blocks!(
+    buffer::CuVector{T}, stream::CUDAChaChaStream, nblocks::Int
+) where {T <: BitInteger}
+
+    p = pointer(buffer)
+    p = Base.unsafe_convert(CuPtr{UInt32}, p)
+    buffer_u32 = unsafe_wrap(CuVector{UInt32}, p, nblocks * CHACHA_BLOCK_SIZE_U32)
+
+    stream.counter = chacha_blocks!(
+        buffer_u32,
+        stream.key,
+        stream.nonce,
+        stream.counter,
+        nblocks
+    )
+
+    buffer
+end
+

src/generation.jl

@@ -4,25 +4,54 @@ Implementation of the RNG API for ChaChaStream
 
 =#
 
+using CUDA
 using Random
 using Random: SamplerType
 
-Random.rng_native_52(::ChaChaStream) = UInt64
+#=
+RNG methods for all subtypes of AbstractChaChaStream
+=#
+
+Random.rng_native_52(::AbstractChaChaStream) = UInt64
+
+# Support for different dimension specifications
+Random.rand(rng::AbstractChaChaStream, T::Type{<:BitInteger}) =
+    Random.rand(rng, T, 1)[]
+Random.rand(rng::AbstractChaChaStream, T::Type{<:BitInteger}, dim1::Int, dims::Int...) =
+    Random.rand(rng, T, Dims((dim1, dims...)))
+Random.rand(rng::AbstractChaChaStream, T::Type{<:BitInteger}, dims::Dims) =
+    Random.rand!(rng, Array{T}(undef, dims))
+
+# Inplace operations
+Random.rand!(rng::AbstractChaChaStream, A::AbstractArray{<:BitInteger}) =
+    Random.rand!(rng, A, eltype(A))
+Random.rand!(rng::AbstractChaChaStream, A::CuArray{<:BitInteger}) =
+    Random.rand!(rng, A, eltype(A))
+
+Random.rand!(
+    rng::AbstractChaChaStream,
+    A::AbstractArray{T},
+    ::Type{T},
+) where {T <: BitInteger} =
+    (Random.rand!(rng, vec(A), T); A)
+
+#=
+RNG methods for ChaChaStream
+=#
 
 Random.rand(rng::ChaChaStream, ::Type{T}) where {T <: BitInteger} =
     _fetch_one!(rng, T)
 
 Random.rand(rng::ChaChaStream, T::Random.SamplerType{<:BitInteger}) =
     _fetch_one!(rng, T[])
 
-Random.rand!(rng::ChaChaStream, A::Array) = Random.rand!(rng, A, eltype(A))
+# Inplace operations
+function Random.rand!(
+    rng::ChaChaStream,
+    A::AbstractVector{T},
+    ::Type{T}
+) where {T <: BitInteger}
 
-function Random.rand!(rng::ChaChaStream, A::Array{T}, ::SamplerType{T}) where {T <: BitInteger}
-    Random.rand!(rng, vec(A), T)
-    A
-end
-
-function Random.rand!(rng::ChaChaStream, A::Vector{T}, ::SamplerType{T}) where {T <: BitInteger}
     # Reinterpret the array as a byte array
     @GC.preserve A begin
         p = pointer(A)
@@ -35,4 +64,25 @@ function Random.rand!(rng::ChaChaStream, A::Vector{T}, ::SamplerType{T}) where {
     A
 end
 
+#=
+RNG methods for CUDAChaChaStream
+=#
+
+function Random.rand!(
+    rng::CUDAChaChaStream,
+    A::AbstractVector{T},
+    ::Type{T}
+) where {T <: BitInteger}
+    # Perform sampling on GPU and then copy to CPU
+    A_gpu = CuVector{T}(undef, length(A))
+    Random.rand!(rng, A_gpu, T)
+    copyto!(A, A_gpu)
+end
+
+Random.rand!(
+    rng::CUDAChaChaStream,
+    A::CuVector{T},
+    ::Type{T}
+) where {T <: BitInteger} =
+    (_fill_buffer!(reinterpret(UInt8, A), rng); A)
 

src/keystream.jl

@@ -45,6 +45,10 @@ julia> randn(stream)
 julia> randstring(stream, 'a':'z', 8)
 "klmptewr"
 ```
+
+```@meta
+DocTestSetup = nothing
+```
 """
 mutable struct ChaChaStream <: AbstractChaChaStream
     key :: SVector{8,UInt32}

test/runtests.jl

@@ -3,6 +3,7 @@ using ChaChaCiphers, Documenter, Test
 include("test_core.jl")
 include("test_chacha.jl")
 include("test_keystream.jl")
+include("test_cuda_keystream.jl")
 
 @testset "Package doctests" begin
     doctest(ChaChaCiphers; manual=false)

test/test_cuda_keystream.jl

@@ -0,0 +1,91 @@
+# Tests for random number generation with
+# CUDAChaChaStream
+
+using ChaChaCiphers
+using CUDA
+using Random
+using Statistics
+using Test
+
+@testset "CUDAChaChaStream tests" begin
+    if CUDA.functional()
+        @testset "Construct CUDAChaChaStream" begin
+            rng = CUDAChaCha12Stream()
+            @test rng.doublerounds == 6
+
+            rng = CUDAChaCha20Stream()
+            @test rng.doublerounds == 10
+        end
+
+        @testset "Generate random strings" begin
+            rng = CUDAChaChaStream(zeros(UInt32, 8), UInt64(0))
+            x = randstring(rng, 'a':'c', 3 * 2^16)
+            @test isa(x, String)
+            @test length(x) == 3 * 2^16
+
+            counts = Dict((c => count(u -> u == c, x)) for c ∈ 'a':'c')
+            counts = Dict((c => counts[c] / length(x)) for c ∈ 'a':'c')
+            @test isapprox(counts['a'], 1/3, atol=1e-2)
+            @test isapprox(counts['b'], 1/3, atol=1e-2)
+            @test isapprox(counts['c'], 1/3, atol=1e-2)
+        end
+
+        @testset "Sample uniform random numbers" begin
+            rng = CUDAChaChaStream(zeros(UInt32, 8), UInt64(0))
+            x = rand(rng, Float32, 100_000)
+
+            @test isa(x, Vector{Float32})
+            @test size(x) == (100_000,)
+            @test isapprox(mean(x), 0.5, atol=1e-2)
+
+            x = rand(rng, Float64, 400, 300)
+            @test isa(x, Array{Float64,2})
+            @test size(x) == (400, 300)
+            @test isapprox(mean(x), 0.5, atol=1e-2)
+
+            # Generate random values directly inside of a pre-allocated array
+            x_cpu = Vector{Float32}(undef, 100_000)
+            Random.rand!(rng, x_cpu)
+            @test isapprox(mean(x_cpu), 0.5, atol=1e-2)
+
+            x_cpu .= 0
+            x_gpu = CuVector{Float32}(undef, 100_000)
+            CUDA.@sync begin
+                Random.rand!(rng, x_gpu)
+                copyto!(x_cpu, x_gpu)
+            end
+            @test isapprox(mean(x_cpu), 0.5, atol=1e-2)
+        end
+
+        @testset "Sample random normal numbers" begin
+            rng = CUDAChaChaStream(zeros(UInt32, 8), UInt64(0))
+            x = randn(rng, Float32, 100_000)
+
+            @test isa(x, Vector{Float32})
+            @test size(x) == (100_000,)
+            @test isapprox(mean(x), 0, atol=1e-2)
+            @test isapprox(std(x), 1, atol=1e-2)
+
+            x = randn(rng, Float64, 100, 50, 50)
+            @test isa(x, Array{Float64,3})
+            @test size(x) == (100, 50, 50)
+            @test isapprox(mean(x), 0, atol=1e-2)
+            @test isapprox(std(x), 1, atol=1e-2)
+        end
+
+        @testset "Save and restore keystream" begin
+            # We should be able to save and restore a GPU
+            # keystream as a CPU keystream, and vice-versa
+            rng = CUDAChaChaStream(zeros(UInt32, 8), UInt64(0))
+            state = getstate(rng)
+            rng_cpu = ChaChaStream(state)
+
+            rand_gpu = rand(rng, 1_000)
+            rand_cpu = rand(rng_cpu, 1_000)
+
+            @test rand_gpu == rand_cpu
+        end
+    else
+        @warn "CUDA.functional() = false; skipping tests"
+    end
+end