sklam · Sep 13, 2022 · Feb 16, 2022 · Feb 16, 2022 · Feb 17, 2022 · Feb 17, 2022
@@ -1,3 +1,4 @@
+import operator
 from numba.core import types, errors
 from numba.core.typing.npydecl import (parse_dtype, parse_shape,
                                        register_number_classes)
@@ -346,6 +347,19 @@ class Cuda_fp16_unary(ConcreteTemplate):
     return Cuda_fp16_unary
 
 
+def _genfp16_unary_operator(l_key):
+    @register
+    class Cuda_fp16_unary(AbstractTemplate):
+        key = l_key
+
+        def generic(self, args, kws):
+            assert not kws
+            if len(args) == 1 and args[0] == types.float16:
+                return signature(types.float16, types.float16)
+
+    return Cuda_fp16_unary
+
+
 def _genfp16_binary(l_key):
     @register
     class Cuda_fp16_binary(ConcreteTemplate):
@@ -355,11 +369,33 @@ class Cuda_fp16_binary(ConcreteTemplate):
     return Cuda_fp16_binary
 
 
+def _genfp16_binary_operator(l_key):
+    @register
+    class Cuda_fp16_binary(AbstractTemplate):
+        key = l_key
+
+        def generic(self, args, kws):
+            assert not kws
+            if len(args) == 2 and \
+                    (args[0] == types.float16 or args[1] == types.float16):
+                return signature(types.float16, types.float16, types.float16)
+
+    return Cuda_fp16_binary
+
+
 Cuda_hadd = _genfp16_binary(cuda.fp16.hadd)
+Cuda_add = _genfp16_binary_operator(operator.add)
+Cuda_iadd = _genfp16_binary_operator(operator.iadd)
 Cuda_hsub = _genfp16_binary(cuda.fp16.hsub)
+Cuda_sub = _genfp16_binary_operator(operator.sub)
+Cuda_isub = _genfp16_binary_operator(operator.isub)
 Cuda_hmul = _genfp16_binary(cuda.fp16.hmul)
-Cuda_hneg = _genfp16_unary(cuda.fp16.hneg)
-Cuda_habs = _genfp16_unary(cuda.fp16.habs)
+Cuda_mul = _genfp16_binary_operator(operator.mul)
+Cuda_imul = _genfp16_binary_operator(operator.imul)
+Cuda_hneg = _genfp16_unary_operator(cuda.fp16.hneg)
+Cuda_neg = _genfp16_unary(operator.neg)
+Cuda_habs = _genfp16_unary_operator(cuda.fp16.habs)
+Cuda_abs = _genfp16_unary(operator.abs)
 
 
 # generate atomic operations

@@ -489,8 +489,14 @@ def ptx_fp16_binary(context, builder, sig, args):
 
 
 lower_fp16_binary(stubs.fp16.hadd, 'add')
+lower_fp16_binary(operator.add, 'add')
+lower_fp16_binary(operator.iadd, 'add')
 lower_fp16_binary(stubs.fp16.hsub, 'sub')
+lower_fp16_binary(operator.sub, 'sub')
+lower_fp16_binary(operator.isub, 'sub')
 lower_fp16_binary(stubs.fp16.hmul, 'mul')
+lower_fp16_binary(operator.mul, 'mul')
+lower_fp16_binary(operator.imul, 'mul')
 
 
 @lower(stubs.fp16.hneg, types.float16)
@@ -500,6 +506,11 @@ def ptx_fp16_hneg(context, builder, sig, args):
     return builder.call(asm, args)
 
 
+@lower(operator.neg, types.float16)
+def operator_hneg(context, builder, sig, args):
+    return ptx_fp16_hneg(context, builder, sig, args)
+
+
 @lower(stubs.fp16.habs, types.float16)
 def ptx_fp16_habs(context, builder, sig, args):
     if cuda.runtime.get_version() < (10, 2):
@@ -517,6 +528,11 @@ def ptx_fp16_habs(context, builder, sig, args):
     return builder.call(asm, args)
 
 
+@lower(abs, types.float16)
+def operator_habs(context, builder, sig, args):
+    return ptx_fp16_habs(context, builder, sig, args)
+
+
 @lower(stubs.fp16.hfma, types.float16, types.float16, types.float16)
 def ptx_hfma(context, builder, sig, args):
     argtys = [ir.IntType(16), ir.IntType(16), ir.IntType(16)]

@@ -1,9 +1,29 @@
 import numpy as np
-from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.testing import unittest, CUDATestCase, skip_unless_cc_53
 from numba import cuda
 import operator
 
 
+def simple_fp16add(ary, a, b):
+    ary[0] = a + b
+
+
+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)
+
+
 class TestOperatorModule(CUDATestCase):
     """
     Test if operator module is supported by the CUDA target.
@@ -36,6 +56,41 @@ def test_truediv(self):
     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)
+        ops = (operator.add, operator.sub, operator.mul)
+
+        for fn, op in zip(functions, ops):
+            with self.subTest(op=op):
+                kernel = cuda.jit("void(f2[:], f2, f2)")(fn)
+
+                expected = np.zeros(1, dtype=np.float16)
-                expected = np.zeros(1, dtype=np.float16)
-                expected = np.zeros(1, dtype=np.float16)
+                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[0], expected)
-                np.testing.assert_allclose(got[0], expected)
+                np.testing.assert_allclose(got, expected)
-                np.testing.assert_allclose(got[0], expected)
+                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)
+
+                expected = np.zeros(1, dtype=np.float16)
-                expected = np.zeros(1, dtype=np.float16)
-                expected = np.zeros(1, dtype=np.float16)
+                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[0], expected)
-                np.testing.assert_allclose(got[0], expected)
+                np.testing.assert_allclose(got, expected)
-                np.testing.assert_allclose(got[0], expected)
+                np.testing.assert_allclose(got, expected)
+
 
 if __name__ == '__main__':
     unittest.main()