numba-cuda 0.0.1__py3-none-any.whl → 0.0.12__py3-none-any.whl

numba_cuda/numba/cuda/tests/cudapy/test_operator.py

@@ -0,0 +1,401 @@
+import numpy as np
+from numba.cuda.testing import (unittest, CUDATestCase, skip_unless_cc_53,
+                                skip_on_cudasim)
+from numba import cuda
+from numba.core.types import f2, b1
+from numba.cuda import compile_ptx
+import operator
+import itertools
+from numba.np.numpy_support import from_dtype
+
+
+def simple_fp16_div_scalar(ary, a, b):
+    ary[0] = a / b
+
+
+def simple_fp16add(ary, a, b):
+    ary[0] = a + b
+
+
+def simple_fp16_iadd(ary, a):
+    ary[0] += a
+
+
+def simple_fp16_isub(ary, a):
+    ary[0] -= a
+
+
+def simple_fp16_imul(ary, a):
+    ary[0] *= a
+
+
+def simple_fp16_idiv(ary, a):
+    ary[0] /= a
+
+
+def simple_fp16sub(ary, a, b):
+    ary[0] = a - b
+
+
+def simple_fp16mul(ary, a, b):
+    ary[0] = a * b
+
+
+def simple_fp16neg(ary, a):
+    ary[0] = -a
+
+
+def simple_fp16abs(ary, a):
+    ary[0] = abs(a)
+
+
+def simple_fp16_gt(ary, a, b):
+    ary[0] = a > b
+
+
+def simple_fp16_ge(ary, a, b):
+    ary[0] = a >= b
+
+
+def simple_fp16_lt(ary, a, b):
+    ary[0] = a < b
+
+
+def simple_fp16_le(ary, a, b):
+    ary[0] = a <= b
+
+
+def simple_fp16_eq(ary, a, b):
+    ary[0] = a == b
+
+
+def simple_fp16_ne(ary, a, b):
+    ary[0] = a != b
+
+
+@cuda.jit('b1(f2, f2)', device=True)
+def hlt_func_1(x, y):
+    return x < y
+
+
+@cuda.jit('b1(f2, f2)', device=True)
+def hlt_func_2(x, y):
+    return x < y
+
+
+def test_multiple_hcmp_1(r, a, b, c):
+    # float16 predicates used in two separate functions
+    r[0] = hlt_func_1(a, b) and hlt_func_2(b, c)
+
+
+def test_multiple_hcmp_2(r, a, b, c):
+    # The same float16 predicate used in the caller and callee
+    r[0] = hlt_func_1(a, b) and b < c
+
+
+def test_multiple_hcmp_3(r, a, b, c):
+    # Different float16 predicates used in the caller and callee
+    r[0] = hlt_func_1(a, b) and c >= b
+
+
+def test_multiple_hcmp_4(r, a, b, c):
+    # The same float16 predicates used twice in a function
+    r[0] = a < b and b < c
+
+
+def test_multiple_hcmp_5(r, a, b, c):
+    # Different float16 predicates used in a function
+    r[0] = a < b and c >= b
+
+
+class TestOperatorModule(CUDATestCase):
+    def setUp(self):
+        super().setUp()
+        np.random.seed(0)
+
+    """
+    Test if operator module is supported by the CUDA target.
+    """
+    def operator_template(self, op):
+        @cuda.jit
+        def foo(a, b):
+            i = 0
+            a[i] = op(a[i], b[i])
+
+        a = np.ones(1)
+        b = np.ones(1)
+        res = a.copy()
+        foo[1, 1](res, b)
+
+        np.testing.assert_equal(res, op(a, b))
+
+    def test_add(self):
+        self.operator_template(operator.add)
+
+    def test_sub(self):
+        self.operator_template(operator.sub)
+
+    def test_mul(self):
+        self.operator_template(operator.mul)
+
+    def test_truediv(self):
+        self.operator_template(operator.truediv)
+
+    def test_floordiv(self):
+        self.operator_template(operator.floordiv)
+
+    @skip_unless_cc_53
+    def test_fp16_binary(self):
+        functions = (simple_fp16add, simple_fp16sub, simple_fp16mul,
+                     simple_fp16_div_scalar)
+        ops = (operator.add, operator.sub, operator.mul, operator.truediv)
+
+        for fn, op in zip(functions, ops):
+            with self.subTest(op=op):
+                kernel = cuda.jit("void(f2[:], f2, f2)")(fn)
+
+                got = np.zeros(1, dtype=np.float16)
+                arg1 = np.random.random(1).astype(np.float16)
+                arg2 = np.random.random(1).astype(np.float16)
+
+                kernel[1, 1](got, arg1[0], arg2[0])
+                expected = op(arg1, arg2)
+                np.testing.assert_allclose(got, expected)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_fp16_binary_ptx(self):
+        functions = (simple_fp16add, simple_fp16sub, simple_fp16mul)
+        instrs = ('add.f16', 'sub.f16', 'mul.f16')
+        args = (f2[:], f2, f2)
+        for fn, instr in zip(functions, instrs):
+            with self.subTest(instr=instr):
+                ptx, _ = compile_ptx(fn, args, cc=(5, 3))
+                self.assertIn(instr, ptx)
+
+    @skip_unless_cc_53
+    def test_mixed_fp16_binary_arithmetic(self):
+        functions = (simple_fp16add, simple_fp16sub, simple_fp16mul,
+                     simple_fp16_div_scalar)
+        ops = (operator.add, operator.sub, operator.mul, operator.truediv)
+        types = (np.int8, np.int16, np.int32, np.int64,
+                 np.float32, np.float64)
+        for (fn, op), ty in itertools.product(zip(functions, ops), types):
+            with self.subTest(op=op, ty=ty):
+                kernel = cuda.jit(fn)
+
+                arg1 = np.random.random(1).astype(np.float16)
+                arg2 = (np.random.random(1) * 100).astype(ty)
+                res_ty = np.result_type(np.float16, ty)
+
+                got = np.zeros(1, dtype=res_ty)
+                kernel[1, 1](got, arg1[0], arg2[0])
+                expected = op(arg1, arg2)
+                np.testing.assert_allclose(got, expected)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_fp16_inplace_binary_ptx(self):
+        functions = (simple_fp16_iadd, simple_fp16_isub, simple_fp16_imul)
+        instrs = ('add.f16', 'sub.f16', 'mul.f16')
+        args = (f2[:], f2)
+
+        for fn, instr in zip(functions, instrs):
+            with self.subTest(instr=instr):
+                ptx, _ = compile_ptx(fn, args, cc=(5, 3))
+                self.assertIn(instr, ptx)
+
+    @skip_unless_cc_53
+    def test_fp16_inplace_binary(self):
+        functions = (simple_fp16_iadd, simple_fp16_isub, simple_fp16_imul,
+                     simple_fp16_idiv)
+        ops = (operator.iadd, operator.isub, operator.imul, operator.itruediv)
+
+        for fn, op in zip(functions, ops):
+            with self.subTest(op=op):
+                kernel = cuda.jit("void(f2[:], f2)")(fn)
+
+                got = np.random.random(1).astype(np.float16)
+                expected = got.copy()
+                arg = np.random.random(1).astype(np.float16)[0]
+                kernel[1, 1](got, arg)
+                op(expected, arg)
+                np.testing.assert_allclose(got, expected)
+
+    @skip_unless_cc_53
+    def test_fp16_unary(self):
+        functions = (simple_fp16neg, simple_fp16abs)
+        ops = (operator.neg, operator.abs)
+
+        for fn, op in zip(functions, ops):
+            with self.subTest(op=op):
+                kernel = cuda.jit("void(f2[:], f2)")(fn)
+
+                got = np.zeros(1, dtype=np.float16)
+                arg1 = np.random.random(1).astype(np.float16)
+
+                kernel[1, 1](got, arg1[0])
+                expected = op(arg1)
+                np.testing.assert_allclose(got, expected)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_fp16_neg_ptx(self):
+        args = (f2[:], f2)
+        ptx, _ = compile_ptx(simple_fp16neg, args, cc=(5, 3))
+        self.assertIn('neg.f16', ptx)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_fp16_abs_ptx(self):
+        args = (f2[:], f2)
+        ptx, _ = compile_ptx(simple_fp16abs, args, cc=(5, 3))
+
+        self.assertIn('abs.f16', ptx)
+
+    @skip_unless_cc_53
+    def test_fp16_comparison(self):
+        functions = (simple_fp16_gt, simple_fp16_ge,
+                     simple_fp16_lt, simple_fp16_le,
+                     simple_fp16_eq, simple_fp16_ne)
+        ops = (operator.gt, operator.ge, operator.lt, operator.le,
+               operator.eq, operator.ne)
+
+        for fn, op in zip(functions, ops):
+            with self.subTest(op=op):
+                kernel = cuda.jit("void(b1[:], f2, f2)")(fn)
+
+                got = np.zeros(1, dtype=np.bool_)
+                arg1 = np.random.random(1).astype(np.float16)
+                arg2 = np.random.random(1).astype(np.float16)
+
+                kernel[1, 1](got, arg1[0], arg2[0])
+                expected = op(arg1, arg2)
+                self.assertEqual(got[0], expected)
+
+    @skip_unless_cc_53
+    def test_mixed_fp16_comparison(self):
+        functions = (simple_fp16_gt, simple_fp16_ge,
+                     simple_fp16_lt, simple_fp16_le,
+                     simple_fp16_eq, simple_fp16_ne)
+        ops = (operator.gt, operator.ge, operator.lt, operator.le,
+               operator.eq, operator.ne)
+        types = (np.int8, np.int16, np.int32, np.int64,
+                 np.float32, np.float64)
+
+        for (fn, op), ty in itertools.product(zip(functions, ops),
+                                              types):
+            with self.subTest(op=op, ty=ty):
+                kernel = cuda.jit(fn)
+
+                got = np.zeros(1, dtype=np.bool_)
+                arg1 = np.random.random(1).astype(np.float16)
+                arg2 = (np.random.random(1) * 100).astype(ty)
+
+                kernel[1, 1](got, arg1[0], arg2[0])
+                expected = op(arg1, arg2)
+                self.assertEqual(got[0], expected)
+
+    @skip_unless_cc_53
+    def test_multiple_float16_comparisons(self):
+        functions = (test_multiple_hcmp_1,
+                     test_multiple_hcmp_2,
+                     test_multiple_hcmp_3,
+                     test_multiple_hcmp_4,
+                     test_multiple_hcmp_5)
+        for fn in functions:
+            with self.subTest(fn=fn):
+                compiled = cuda.jit("void(b1[:], f2, f2, f2)")(fn)
+                ary = np.zeros(1, dtype=np.bool_)
+                arg1 = np.float16(2.)
+                arg2 = np.float16(3.)
+                arg3 = np.float16(4.)
+                compiled[1, 1](ary, arg1, arg2, arg3)
+                self.assertTrue(ary[0])
+
+    @skip_unless_cc_53
+    def test_multiple_float16_comparisons_false(self):
+        functions = (test_multiple_hcmp_1,
+                     test_multiple_hcmp_2,
+                     test_multiple_hcmp_3,
+                     test_multiple_hcmp_4,
+                     test_multiple_hcmp_5)
+        for fn in functions:
+            with self.subTest(fn=fn):
+                compiled = cuda.jit("void(b1[:], f2, f2, f2)")(fn)
+                ary = np.zeros(1, dtype=np.bool_)
+                arg1 = np.float16(2.)
+                arg2 = np.float16(3.)
+                arg3 = np.float16(1.)
+                compiled[1, 1](ary, arg1, arg2, arg3)
+                self.assertFalse(ary[0])
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_fp16_comparison_ptx(self):
+        functions = (simple_fp16_gt, simple_fp16_ge,
+                     simple_fp16_lt, simple_fp16_le,
+                     simple_fp16_eq, simple_fp16_ne)
+        ops = (operator.gt, operator.ge, operator.lt, operator.le,
+               operator.eq, operator.ne)
+        opstring = ('setp.gt.f16', 'setp.ge.f16',
+                    'setp.lt.f16', 'setp.le.f16',
+                    'setp.eq.f16', 'setp.ne.f16')
+        args = (b1[:], f2, f2)
+
+        for fn, op, s in zip(functions, ops, opstring):
+            with self.subTest(op=op):
+                ptx, _ = compile_ptx(fn, args, cc=(5, 3))
+                self.assertIn(s, ptx)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_fp16_int8_comparison_ptx(self):
+        # Test that int8 can be safely converted to fp16
+        # in a comparison
+        functions = (simple_fp16_gt, simple_fp16_ge,
+                     simple_fp16_lt, simple_fp16_le,
+                     simple_fp16_eq, simple_fp16_ne)
+        ops = (operator.gt, operator.ge, operator.lt, operator.le,
+               operator.eq, operator.ne)
+
+        opstring = {operator.gt:'setp.gt.f16',
+                    operator.ge:'setp.ge.f16',
+                    operator.lt:'setp.lt.f16',
+                    operator.le:'setp.le.f16',
+                    operator.eq:'setp.eq.f16',
+                    operator.ne:'setp.ne.f16'}
+        for fn, op in zip(functions, ops):
+            with self.subTest(op=op):
+                args = (b1[:], f2, from_dtype(np.int8))
+                ptx, _ = compile_ptx(fn, args, cc=(5, 3))
+                self.assertIn(opstring[op], ptx)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_mixed_fp16_comparison_promotion_ptx(self):
+        functions = (simple_fp16_gt, simple_fp16_ge,
+                     simple_fp16_lt, simple_fp16_le,
+                     simple_fp16_eq, simple_fp16_ne)
+        ops = (operator.gt, operator.ge, operator.lt, operator.le,
+               operator.eq, operator.ne)
+
+        types_promote = (np.int16, np.int32, np.int64,
+                         np.float32, np.float64)
+        opstring = {operator.gt:'setp.gt.',
+                    operator.ge:'setp.ge.',
+                    operator.lt:'setp.lt.',
+                    operator.le:'setp.le.',
+                    operator.eq:'setp.eq.',
+                    operator.ne:'setp.neu.'}
+        opsuffix = {np.dtype('int32'): 'f64',
+                    np.dtype('int64'): 'f64',
+                    np.dtype('float32'): 'f32',
+                    np.dtype('float64'): 'f64'}
+
+        for (fn, op), ty in itertools.product(zip(functions, ops),
+                                              types_promote):
+            with self.subTest(op=op, ty=ty):
+                arg2_ty = np.result_type(np.float16, ty)
+                args = (b1[:], f2, from_dtype(arg2_ty))
+                ptx, _ = compile_ptx(fn, args, cc=(5, 3))
+
+                ops = opstring[op] + opsuffix[arg2_ty]
+                self.assertIn(ops, ptx)
+
+
+if __name__ == '__main__':
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_optimization.py

@@ -0,0 +1,86 @@
+import numpy as np
+
+from numba.cuda.testing import skip_on_cudasim, CUDATestCase
+from numba import cuda, float64
+import unittest
+
+
+def kernel_func(x):
+    x[0] = 1
+
+
+def device_func(x, y, z):
+    return x * y + z
+
+
+# Fragments of code that are removed from kernel_func's PTX when optimization
+# is on
+removed_by_opt = ( '__local_depot0', 'call.uni', 'st.param.b64')
+
+
+@skip_on_cudasim('Simulator does not optimize code')
+class TestOptimization(CUDATestCase):
+    def test_eager_opt(self):
+        # Optimization should occur by default
+        sig = (float64[::1],)
+        kernel = cuda.jit(sig)(kernel_func)
+        ptx = kernel.inspect_asm()
+
+        for fragment in removed_by_opt:
+            with self.subTest(fragment=fragment):
+                self.assertNotIn(fragment, ptx[sig])
+
+    def test_eager_noopt(self):
+        # Optimization disabled
+        sig = (float64[::1],)
+        kernel = cuda.jit(sig, opt=False)(kernel_func)
+        ptx = kernel.inspect_asm()
+
+        for fragment in removed_by_opt:
+            with self.subTest(fragment=fragment):
+                self.assertIn(fragment, ptx[sig])
+
+    def test_lazy_opt(self):
+        # Optimization should occur by default
+        kernel = cuda.jit(kernel_func)
+        x = np.zeros(1, dtype=np.float64)
+        kernel[1, 1](x)
+
+        # Grab the PTX for the one definition that has just been jitted
+        ptx = next(iter(kernel.inspect_asm().items()))[1]
+
+        for fragment in removed_by_opt:
+            with self.subTest(fragment=fragment):
+                self.assertNotIn(fragment, ptx)
+
+    def test_lazy_noopt(self):
+        # Optimization disabled
+        kernel = cuda.jit(opt=False)(kernel_func)
+        x = np.zeros(1, dtype=np.float64)
+        kernel[1, 1](x)
+
+        # Grab the PTX for the one definition that has just been jitted
+        ptx = next(iter(kernel.inspect_asm().items()))[1]
+
+        for fragment in removed_by_opt:
+            with self.subTest(fragment=fragment):
+                self.assertIn(fragment, ptx)
+
+    def test_device_opt(self):
+        # Optimization should occur by default
+        sig = (float64, float64, float64)
+        device = cuda.jit(sig, device=True)(device_func)
+        ptx = device.inspect_asm(sig)
+        self.assertIn('fma.rn.f64', ptx)
+
+    def test_device_noopt(self):
+        # Optimization disabled
+        sig = (float64, float64, float64)
+        device = cuda.jit(sig, device=True, opt=False)(device_func)
+        ptx = device.inspect_asm(sig)
+        # Fused-multiply adds should be disabled when not optimizing
+        self.assertNotIn('fma.rn.f64', ptx)
+
+
+if __name__ == '__main__':
+    unittest.main()