File: C:/Users/fred/anaconda3/Lib/site-packages/numba/cuda/tests/cudapy/test_warp_ops.py
import numpy as np
from numba import cuda, int32, int64, float32, float64
from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
from numba.core import config
def useful_syncwarp(ary):
i = cuda.grid(1)
if i == 0:
ary[0] = 42
cuda.syncwarp(0xffffffff)
ary[i] = ary[0]
def use_shfl_sync_idx(ary, idx):
i = cuda.grid(1)
val = cuda.shfl_sync(0xffffffff, i, idx)
ary[i] = val
def use_shfl_sync_up(ary, delta):
i = cuda.grid(1)
val = cuda.shfl_up_sync(0xffffffff, i, delta)
ary[i] = val
def use_shfl_sync_down(ary, delta):
i = cuda.grid(1)
val = cuda.shfl_down_sync(0xffffffff, i, delta)
ary[i] = val
def use_shfl_sync_xor(ary, xor):
i = cuda.grid(1)
val = cuda.shfl_xor_sync(0xffffffff, i, xor)
ary[i] = val
def use_shfl_sync_with_val(ary, into):
i = cuda.grid(1)
val = cuda.shfl_sync(0xffffffff, into, 0)
ary[i] = val
def use_vote_sync_all(ary_in, ary_out):
i = cuda.grid(1)
pred = cuda.all_sync(0xffffffff, ary_in[i])
ary_out[i] = pred
def use_vote_sync_any(ary_in, ary_out):
i = cuda.grid(1)
pred = cuda.any_sync(0xffffffff, ary_in[i])
ary_out[i] = pred
def use_vote_sync_eq(ary_in, ary_out):
i = cuda.grid(1)
pred = cuda.eq_sync(0xffffffff, ary_in[i])
ary_out[i] = pred
def use_vote_sync_ballot(ary):
i = cuda.threadIdx.x
ballot = cuda.ballot_sync(0xffffffff, True)
ary[i] = ballot
def use_match_any_sync(ary_in, ary_out):
i = cuda.grid(1)
ballot = cuda.match_any_sync(0xffffffff, ary_in[i])
ary_out[i] = ballot
def use_match_all_sync(ary_in, ary_out):
i = cuda.grid(1)
ballot, pred = cuda.match_all_sync(0xffffffff, ary_in[i])
ary_out[i] = ballot if pred else 0
def use_independent_scheduling(arr):
i = cuda.threadIdx.x
if i % 4 == 0:
ballot = cuda.ballot_sync(0x11111111, True)
elif i % 4 == 1:
ballot = cuda.ballot_sync(0x22222222, True)
elif i % 4 == 2:
ballot = cuda.ballot_sync(0x44444444, True)
elif i % 4 == 3:
ballot = cuda.ballot_sync(0x88888888, True)
arr[i] = ballot
def _safe_cc_check(cc):
if config.ENABLE_CUDASIM:
return True
else:
return cuda.get_current_device().compute_capability >= cc
@skip_on_cudasim("Warp Operations are not yet implemented on cudasim")
class TestCudaWarpOperations(CUDATestCase):
def test_useful_syncwarp(self):
compiled = cuda.jit("void(int32[:])")(useful_syncwarp)
nelem = 32
ary = np.empty(nelem, dtype=np.int32)
compiled[1, nelem](ary)
self.assertTrue(np.all(ary == 42))
def test_shfl_sync_idx(self):
compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_idx)
nelem = 32
idx = 4
ary = np.empty(nelem, dtype=np.int32)
compiled[1, nelem](ary, idx)
self.assertTrue(np.all(ary == idx))
def test_shfl_sync_up(self):
compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_up)
nelem = 32
delta = 4
ary = np.empty(nelem, dtype=np.int32)
exp = np.arange(nelem, dtype=np.int32)
exp[delta:] -= delta
compiled[1, nelem](ary, delta)
self.assertTrue(np.all(ary == exp))
def test_shfl_sync_down(self):
compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_down)
nelem = 32
delta = 4
ary = np.empty(nelem, dtype=np.int32)
exp = np.arange(nelem, dtype=np.int32)
exp[:-delta] += delta
compiled[1, nelem](ary, delta)
self.assertTrue(np.all(ary == exp))
def test_shfl_sync_xor(self):
compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_xor)
nelem = 32
xor = 16
ary = np.empty(nelem, dtype=np.int32)
exp = np.arange(nelem, dtype=np.int32) ^ xor
compiled[1, nelem](ary, xor)
self.assertTrue(np.all(ary == exp))
def test_shfl_sync_types(self):
types = int32, int64, float32, float64
values = (np.int32(-1), np.int64(1 << 42),
np.float32(np.pi), np.float64(np.pi))
for typ, val in zip(types, values):
compiled = cuda.jit((typ[:], typ))(use_shfl_sync_with_val)
nelem = 32
ary = np.empty(nelem, dtype=val.dtype)
compiled[1, nelem](ary, val)
self.assertTrue(np.all(ary == val))
def test_vote_sync_all(self):
compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_all)
nelem = 32
ary_in = np.ones(nelem, dtype=np.int32)
ary_out = np.empty(nelem, dtype=np.int32)
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 1))
ary_in[-1] = 0
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 0))
def test_vote_sync_any(self):
compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_any)
nelem = 32
ary_in = np.zeros(nelem, dtype=np.int32)
ary_out = np.empty(nelem, dtype=np.int32)
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 0))
ary_in[2] = 1
ary_in[5] = 1
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 1))
def test_vote_sync_eq(self):
compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_eq)
nelem = 32
ary_in = np.zeros(nelem, dtype=np.int32)
ary_out = np.empty(nelem, dtype=np.int32)
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 1))
ary_in[1] = 1
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 0))
ary_in[:] = 1
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 1))
def test_vote_sync_ballot(self):
compiled = cuda.jit("void(uint32[:])")(use_vote_sync_ballot)
nelem = 32
ary = np.empty(nelem, dtype=np.uint32)
compiled[1, nelem](ary)
self.assertTrue(np.all(ary == np.uint32(0xffffffff)))
@unittest.skipUnless(_safe_cc_check((7, 0)),
"Matching requires at least Volta Architecture")
def test_match_any_sync(self):
compiled = cuda.jit("void(int32[:], int32[:])")(use_match_any_sync)
nelem = 10
ary_in = np.arange(nelem, dtype=np.int32) % 2
ary_out = np.empty(nelem, dtype=np.int32)
exp = np.tile((0b0101010101, 0b1010101010), 5)
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == exp))
@unittest.skipUnless(_safe_cc_check((7, 0)),
"Matching requires at least Volta Architecture")
def test_match_all_sync(self):
compiled = cuda.jit("void(int32[:], int32[:])")(use_match_all_sync)
nelem = 10
ary_in = np.zeros(nelem, dtype=np.int32)
ary_out = np.empty(nelem, dtype=np.int32)
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 0b1111111111))
ary_in[1] = 4
compiled[1, nelem](ary_in, ary_out)
self.assertTrue(np.all(ary_out == 0))
@unittest.skipUnless(_safe_cc_check((7, 0)),
"Independent scheduling requires at least Volta "
"Architecture")
def test_independent_scheduling(self):
compiled = cuda.jit("void(uint32[:])")(use_independent_scheduling)
arr = np.empty(32, dtype=np.uint32)
exp = np.tile((0x11111111, 0x22222222, 0x44444444, 0x88888888), 8)
compiled[1, 32](arr)
self.assertTrue(np.all(arr == exp))
def test_activemask(self):
@cuda.jit
def use_activemask(x):
i = cuda.grid(1)
if (i % 2) == 0:
# Even numbered threads fill in even numbered array entries
# with binary "...01010101"
x[i] = cuda.activemask()
else:
# Odd numbered threads fill in odd numbered array entries
# with binary "...10101010"
x[i] = cuda.activemask()
out = np.zeros(32, dtype=np.uint32)
use_activemask[1, 32](out)
# 0x5 = 0101: The pattern from even-numbered threads
# 0xA = 1010: The pattern from odd-numbered threads
expected = np.tile((0x55555555, 0xAAAAAAAA), 16)
np.testing.assert_equal(expected, out)
def test_lanemask_lt(self):
@cuda.jit
def use_lanemask_lt(x):
i = cuda.grid(1)
x[i] = cuda.lanemask_lt()
out = np.zeros(32, dtype=np.uint32)
use_lanemask_lt[1, 32](out)
# A string of 1s that grows from the LSB for each entry:
# 0, 1, 3, 7, F, 1F, 3F, 7F, FF, 1FF, etc.
# or in binary:
# ...0001, ....0011, ...0111, etc.
expected = np.asarray([(2 ** i) - 1 for i in range(32)],
dtype=np.uint32)
np.testing.assert_equal(expected, out)
if __name__ == '__main__':
unittest.main()