Support drawing normally distributed random numbers with randn (#39)

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.2.0"
+version = "1.3.0"
 
 [deps]
 Clang = "40e3b903-d033-50b4-a0cc-940c62c95e31"

docs/src/compiler.md

@@ -29,11 +29,15 @@ IPUCompiler.PROGRESS_SPINNER
 
 Inside codelets defined with [`@codelet`](@ref) all calls to random functions
 
+* `rand(Float16)`
 * `rand(Float32)`
 * `rand(UInt32)`
 * `rand(UInt64)`
+* `randn(Float16)`
+* `randn(Float32)`
 
-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), but with the general semantic of the Julia function `rand` (numbers uniformely distributed in the $[0, 1)$ range).
+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` (numbers 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 $[-5-13/16, 5+13/16]$).
 
 Additionally, you can use the [IPU builtins](https://docs.graphcore.ai/projects/poplar-api/en/latest/ipu_intrinsics/ipu_builtins.html) listed below.
 

src/compiler/codelet.jl

@@ -235,6 +235,8 @@ function __build_codelet(graph::Poplar.GraphAllocated, kernel, name::String, ori
                       "_llvm_colossus_urand_f32" => "llvm.colossus.urand.f32",
                       "_llvm_colossus_urand32" => "llvm.colossus.urand32",
                       "_llvm_colossus_urand64" => "llvm.colossus.urand64",
+                      "_llvm_colossus_f16v2grand" => "llvm.colossus.f16v2grand",
+                      "_llvm_colossus_f32v2grand" => "llvm.colossus.f32v2grand",
                       )
 
     method = methods(origKernel)[end]

src/compiler/runtime.jl

@@ -49,6 +49,8 @@ Base.rand(T::Type{Float16}) = ccall("extern _llvm_colossus_urand_f16",  llvmcall
 Base.rand(T::Type{Float32}) = ccall("extern _llvm_colossus_urand_f32",  llvmcall, Float32, ()) + T(0.5)
 Base.rand(T::Type{UInt32}) = ccall("extern _llvm_colossus_urand32",  llvmcall, UInt32, ()) + T(0.5)
 Base.rand(T::Type{UInt64}) = ccall("extern _llvm_colossus_urand64",  llvmcall, UInt64, ()) + T(0.5)
+Base.randn(T::Type{Float16}) = @inbounds ccall("extern _llvm_colossus_f16v2grand",  llvmcall, NTuple{2, VecElement{Float16}}, ())[1].value
+Base.randn(T::Type{Float32}) = @inbounds ccall("extern _llvm_colossus_f32v2grand",  llvmcall, NTuple{2, VecElement{Float32}}, ())[1].value
 
 ## Math functions.
 # There are different reasons why we prefer LLVM intrinsics on the IPU: implementations in

test/compiler.jl

@@ -198,6 +198,7 @@ function test_ipubuiltins(device)
     outvec1 = PoplarVector{Float16}(undef, N)
     outvec2 = PoplarVector{Float16}(undef, N)
     outvec3 = PoplarVector{Float16}(undef, N)
+    outvec4 = PoplarVector{Float16}(undef, N)
 
     @ipuprogram device begin
         function Random(out::VertexVector{Float16, Out})
@@ -214,14 +215,21 @@ function test_ipubuiltins(device)
             copyto!(out, in)
             sort!(out; rev=true)
         end
+        function RandomNorm(out::VertexVector{Float16, Out})
+            for idx in eachindex(out)
+                out[idx] = randn(Float16)
+            end
+        end
 
         Random(outvec1)
         TimesTwoSin(outvec1, outvec2)
         Sort16(outvec2, outvec3)
+        RandomNorm(outvec4)
 
         jl_outvec1 = outvec1
         jl_outvec2 = outvec2
         jl_outvec3 = outvec3
+        jl_outvec4 = outvec4
     end
     Poplar.detach_devices()
     # There's a non-zero probability that this test may fail, but assuming an
@@ -230,6 +238,8 @@ function test_ipubuiltins(device)
     @test mean(jl_outvec1) ≈ 0.5 rtol=(pi * sqrt(N) / N)
     @test jl_outvec2 ≈ sin.(2 .* jl_outvec1)
     @test jl_outvec3 ≈ sort(jl_outvec2; rev=true)
+    @test mean(jl_outvec4) ≈ 0 atol=0.02
+    @test std(jl_outvec4) ≈ 1 rtol=0.02
 end
 
 rosenbrock(x, y=4) = (1 - x) ^ 2 + 100 * (y - x ^ 2) ^ 2