RNG SIMD

I've vectorized the Philox and AES-NI RNGs, fixes #12 (closed). I had to add a very minimal integer vectorization that only supports int32, makeVec, +, and loop counters. Also, the RNGNode now needs to know already during construction how it's vectorized, which is ugly, but could only be resolved by a better type system (#20 (closed)). For the same reason, it is not possible to use a vectorized float RNG with double fields or vice versa. Also, we essentially discard half the random numbers in double precision mode because otherwise the number of variables we return would change between the vectorized and non-vectorized version, which is incompatible with the interface.

For the tests, we need to add pip3 install randomgen to the Dockerfile.

pystencils/include/philox_rand.h

@@ -113,6 +113,7 @@ QUALIFIERS void philox_float4(uint32 ctr0, uint32 ctr1, uint32 ctr2, uint32 ctr3
     rnd4 = ctr[3] * TWOPOW32_INV_FLOAT + (TWOPOW32_INV_FLOAT/2.0f);
 }
 
+#ifndef __CUDA_ARCH__
 #ifdef __SSE__
 QUALIFIERS void _philox4x32round(__m128i* ctr, __m128i* key)
 {
@@ -205,9 +206,9 @@ QUALIFIERS __m128d _uniform_double_hq(__m128i x, __m128i y)
 }
 
 
-QUALIFIERS void philox_float16(__m128i ctr0, __m128i ctr1, __m128i ctr2, __m128i ctr3,
-                               uint32 key0, uint32 key1,
-                               __m128 & rnd1, __m128 & rnd2, __m128 & rnd3, __m128 & rnd4)
+QUALIFIERS void philox_float4(__m128i ctr0, __m128i ctr1, __m128i ctr2, __m128i ctr3,
+                              uint32 key0, uint32 key1,
+                              __m128 & rnd1, __m128 & rnd2, __m128 & rnd3, __m128 & rnd4)
 {
     __m128i key[2] = {_mm_set1_epi32(key0), _mm_set1_epi32(key1)};
     __m128i ctr[4] = {ctr0, ctr1, ctr2, ctr3};
@@ -246,7 +247,7 @@ QUALIFIERS void philox_float16(__m128i ctr0, __m128i ctr1, __m128i ctr2, __m128i
 }
 
 
-QUALIFIERS void philox_double8(__m128i ctr0, __m128i ctr1, __m128i ctr2, __m128i ctr3,
+QUALIFIERS void philox_double2(__m128i ctr0, __m128i ctr1, __m128i ctr2, __m128i ctr3,
                                uint32 key0, uint32 key1,
                                __m128d & rnd1lo, __m128d & rnd1hi, __m128d & rnd2lo, __m128d & rnd2hi)
 {
@@ -372,9 +373,9 @@ QUALIFIERS __m256d _uniform_double_hq(__m256i x, __m256i y)
 }
 
 
-QUALIFIERS void philox_float32(__m256i ctr0, __m256i ctr1, __m256i ctr2, __m256i ctr3,
-                               uint32 key0, uint32 key1,
-                               __m256 & rnd1, __m256 & rnd2, __m256 & rnd3, __m256 & rnd4)
+QUALIFIERS void philox_float4(__m256i ctr0, __m256i ctr1, __m256i ctr2, __m256i ctr3,
+                              uint32 key0, uint32 key1,
+                              __m256 & rnd1, __m256 & rnd2, __m256 & rnd3, __m256 & rnd4)
 {
     __m256i key[2] = {_mm256_set1_epi32(key0), _mm256_set1_epi32(key1)};
     __m256i ctr[4] = {ctr0, ctr1, ctr2, ctr3};
@@ -413,9 +414,9 @@ QUALIFIERS void philox_float32(__m256i ctr0, __m256i ctr1, __m256i ctr2, __m256i
 }
 
 
-QUALIFIERS void philox_double16(__m256i ctr0, __m256i ctr1, __m256i ctr2, __m256i ctr3,
-                                uint32 key0, uint32 key1,
-                                __m256d & rnd1lo, __m256d & rnd1hi, __m256d & rnd2lo, __m256d & rnd2hi)
+QUALIFIERS void philox_double2(__m256i ctr0, __m256i ctr1, __m256i ctr2, __m256i ctr3,
+                               uint32 key0, uint32 key1,
+                               __m256d & rnd1lo, __m256d & rnd1hi, __m256d & rnd2lo, __m256d & rnd2hi)
 {
     __m256i key[2] = {_mm256_set1_epi32(key0), _mm256_set1_epi32(key1)};
     __m256i ctr[4] = {ctr0, ctr1, ctr2, ctr3};
@@ -505,9 +506,9 @@ QUALIFIERS __m512d _uniform_double_hq(__m512i x, __m512i y)
 }
 
 
-QUALIFIERS void philox_float64(__m512i ctr0, __m512i ctr1, __m512i ctr2, __m512i ctr3,
-                               uint32 key0, uint32 key1,
-                               __m512 & rnd1, __m512 & rnd2, __m512 & rnd3, __m512 & rnd4)
+QUALIFIERS void philox_float4(__m512i ctr0, __m512i ctr1, __m512i ctr2, __m512i ctr3,
+                              uint32 key0, uint32 key1,
+                              __m512 & rnd1, __m512 & rnd2, __m512 & rnd3, __m512 & rnd4)
 {
     __m512i key[2] = {_mm512_set1_epi32(key0), _mm512_set1_epi32(key1)};
     __m512i ctr[4] = {ctr0, ctr1, ctr2, ctr3};
@@ -535,9 +536,9 @@ QUALIFIERS void philox_float64(__m512i ctr0, __m512i ctr1, __m512i ctr2, __m512i
 }
 
 
-QUALIFIERS void philox_double32(__m512i ctr0, __m512i ctr1, __m512i ctr2, __m512i ctr3,
-                                uint32 key0, uint32 key1,
-                                __m512d & rnd1lo, __m512d & rnd1hi, __m512d & rnd2lo, __m512d & rnd2hi)
+QUALIFIERS void philox_double2(__m512i ctr0, __m512i ctr1, __m512i ctr2, __m512i ctr3,
+                               uint32 key0, uint32 key1,
+                               __m512d & rnd1lo, __m512d & rnd1hi, __m512d & rnd2lo, __m512d & rnd2hi)
 {
     __m512i key[2] = {_mm512_set1_epi32(key0), _mm512_set1_epi32(key1)};
     __m512i ctr[4] = {ctr0, ctr1, ctr2, ctr3};
@@ -558,4 +559,5 @@ QUALIFIERS void philox_double32(__m512i ctr0, __m512i ctr1, __m512i ctr2, __m512
     rnd2hi = _uniform_double_hq<true>(ctr[2], ctr[3]);
 }
 #endif
+#endif