Add randn2! to fill vectors with normally distributed numbers (#57)

Author: giordano

Project.toml

@@ -1,7 +1,7 @@
 name = "IPUToolkit"
 uuid = "92e0b95a-4011-435a-96f4-10064551ddbe"
 authors = ["Emily Dietrich <jakibaki@live.com>", "Luk Burchard <luk.burchard@gmail.com>", "Mosè Giordano <mose@gnu.org>"]
-version = "1.4.2"
+version = "1.5.0"
 
 [deps]
 Clang = "40e3b903-d033-50b4-a0cc-940c62c95e31"

docs/src/compiler.md

@@ -39,11 +39,17 @@ Inside codelets defined with [`@codelet`](@ref) all calls to random functions
 result to call to corresponding IPU builtins for [random number generation](https://docs.graphcore.ai/projects/poplar-api/en/latest/ipu_intrinsics/ipu_builtins.html#random-number-generation).
 The uniformly distributed numbers follow the general semantic of the Julia function `rand` (floating point numbers are uniformely distributed in the $[0, 1)$ range), while the normally distributed numbers have the properties described in the Poplar SDK documentation (numbers are in the range $[-93/16, 93/16]$).
 
+!!! note
+
+    The IPU builtins for random numbers return pairs of numbers, but the Julia functions `randn(Float16)` and `randn(Float32)` return only a single number, discarding the second number of the pair.
+    If you have a vector of even length that you want to fill in-place with normally distributed numbers, you can use the [`randn2!`](@ref) function to do that efficiently, without discarding any number.
+
 Additionally, you can use the [IPU builtins](https://docs.graphcore.ai/projects/poplar-api/en/latest/ipu_intrinsics/ipu_builtins.html) listed below.
 
 ```@docs
 get_scount_l
 get_tile_id
+randn2!
 ```
 
 ## Printing

src/compiler/compiler.jl

@@ -1,7 +1,7 @@
 # based on GPUCompiler example https://github.com/JuliaGPU/GPUCompiler.jl/blob/master/examples/kernel.jl
 module IPUCompiler
 
-export @codelet, @ipuprogram, VertexVector, VertexScalar, In, Out, InOut, get_scount_l, get_tile_id, add_vertex
+export @codelet, @ipuprogram, VertexVector, VertexScalar, In, Out, InOut, get_scount_l, get_tile_id, randn2!, add_vertex
 
 include("output.jl")
 
@@ -80,8 +80,17 @@ Call the [`__builtin_ipu_get_tile_id()`](https://docs.graphcore.ai/projects/popl
 """
 function get_tile_id end
 
-include("runtime.jl")
+"""
+    randn2!(v::VertexVector) -> v
+
+Fill the vector `v` with normally-distributed (mean 0, standard deviation 1) random numbers.
+The vector *must* have even length.
+This function takes advantage of [IPU builtins for random number generation](https://docs.graphcore.ai/projects/poplar-api/en/latest/ipu_intrinsics/ipu_builtins.html#random-number-generation), which return pairs of numbers at a time.
+"""
+function randn2! end
+
 include("vertices.jl")
+include("runtime.jl")
 
 GPUCompiler.runtime_module(::CompilerJob{<:Any,IPUCompilerParams}) = IPURuntime
 # `GPUCompiler.isintrinsic` specifies functions which are to be considered intrinsics for

src/compiler/runtime.jl

@@ -1,6 +1,6 @@
 module IPURuntime
 
-import ..IPUCompiler: @device_override, @ipuprintf, @ipuprintln, get_scount_l, get_tile_id
+import ..IPUCompiler: @device_override, @ipuprintf, @ipuprintln, get_scount_l, get_tile_id, randn2!, VertexVector, Out, InOut
 using GPUCompiler: reset_runtime
 import LinearAlgebra
 
@@ -54,6 +54,17 @@ get_tile_id() = ccall("extern _llvm_colossus_get_tile_id", llvmcall, Cuint, ())
 # same semantic.
 @device_override Base.randn(T::Type{Float16}) = @inbounds ccall("extern _llvm_colossus_f16v2grand", llvmcall, NTuple{2, VecElement{Float16}}, ())[1].value
 @device_override Base.randn(T::Type{Float32}) = @inbounds ccall("extern _llvm_colossus_f32v2grand", llvmcall, NTuple{2, VecElement{Float32}}, ())[1].value
+function randn2!(v::VertexVector{T}) where {T}
+    for idx in UInt32(1):UInt32(2):UInt32(length(v))
+        rnd = if T == Float32
+            ccall("extern _llvm_colossus_f32v2grand", llvmcall, NTuple{2, VecElement{Float32}}, ())
+        elseif T == Float16
+            ccall("extern _llvm_colossus_f16v2grand", llvmcall, NTuple{2, VecElement{Float16}}, ())
+        end
+        @inbounds v[idx]   = rnd[1].value
+        @inbounds v[idx+1] = rnd[2].value
+    end
+end
 
 ## Math functions.
 # There are different reasons why we prefer LLVM intrinsics on the IPU: implementations in