Stop calling CUDA.@sync inside of chacha_blocks!

Author: kernelmethod

src/ChaCha.jl

@@ -194,8 +194,7 @@ function chacha_blocks!(
         @cuda blocks=chunk_blocks _cuda_chacha_add_initial_state!(state, key, nonce, counter)
     end
 
-    counter += nblocks
-    counter
+    counter + nblocks
 end
 
 function _cuda_chacha_rounds!(state, doublerounds)
@@ -231,7 +230,7 @@ function _cuda_chacha_rounds!(state, doublerounds)
     nothing
 end
 
-export CHACHA_BLOCK_SIZE
+export CHACHA_BLOCK_SIZE, CHACHA_BLOCK_SIZE_U32
 export chacha_blocks!
 
 end # module

test/test_chacha.jl

@@ -203,7 +203,7 @@ end
                 counter = UInt64(0)
 
                 ctr = chacha_blocks!(state, key, nonce, counter, 1024 ÷ 16)
-                ctr_gpu = chacha_blocks!(state_gpu, key_gpu, nonce, counter, 1024 ÷ 16)
+                CUDA.@sync ctr_gpu = chacha_blocks!(state_gpu, key_gpu, nonce, counter, 1024 ÷ 16)
 
                 @test ctr == ctr_gpu
                 @test state_gpu == CuArray(state)