sklam · Dec 15, 2021 · Sep 21, 2021 · Sep 21, 2021 · Oct 5, 2021 · Oct 6, 2021
diff --git a/docs/source/cuda-reference/kernel.rst b/docs/source/cuda-reference/kernel.rst
@@ -463,6 +463,51 @@ precision parts of the CUDA Toolkit documentation.
    ``cbrt`` and ``cbrtf`` in the C api. Supports float32, and float64 arguments
    only.
 
+16-bit Floating Point Intrinsics
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. function:: numba.cuda.fp16.hfma (a, b, c)
+
+   Perform the fused multiply-add operation ``(a * b) + c`` on 16-bit
+   floating point arguments in round to nearest mode. Maps to the ``fma.rn.f16``
+   PTX instruction.
+
+   Returns the fp16 result of the fused multiply-add.
+
+.. function:: numba.cuda.fp16.hadd (a, b)
+
+   Perform the add operation ``a + b`` on 16-bit floating point arguments in
+   round to nearest mode. Maps to the ``add.f16`` PTX instruction.
+
+   Returns the 16-bit floating point result of the addition.
+
+.. function:: numba.cuda.fp16.hsub (a, b)
+
+   Perform the subtract operation ``a - b`` on 16-bit floating point arguments in
+   round to nearest mode. Maps to the ``sub.f16`` PTX instruction.
+
+   Returns the 16-bit floating point result of the subtraction.
+
+.. function:: numba.cuda.fp16.hmul (a, b)
+
+   Perform the multiply operation ``a * b`` on 16-bit floating point arguments in
+   round to nearest mode. Maps to the ``mul.f16`` PTX instruction.
+
+   Returns the 16-bit floating point result of the multiplication.
+
+.. function:: numba.cuda.fp16.hneg (a)
+
+   Perform the negation operation ``-a`` on 1the 6-bit floating point argument.
+   Maps to the ``neg.f16`` PTX instruction.
+
+   Returns the 16-bit floating point result of the negation.
+
+.. function:: numba.cuda.fp16.habs (a)
+
+   Perform the absolute value operation ``|a|`` on 1the 6-bit floating point argument.
+   Maps to the ``abs.f16`` PTX instruction.
+
+   Returns the 16-bit floating point result of the absolute value operation.
 
 Control Flow Instructions
 ~~~~~~~~~~~~~~~~~~~~~~~~~

diff --git a/numba/core/boxing.py b/numba/core/boxing.py
@@ -60,7 +60,7 @@ def unbox_integer(typ, obj, c):
 
 @box(types.Float)
 def box_float(typ, val, c):
-    if typ == types.float32:
+    if typ == types.float32 or types == types.float16:
         dbval = c.builder.fpext(val, c.pyapi.double)
     else:
         assert typ == types.float64
@@ -72,7 +72,7 @@ def unbox_float(typ, obj, c):
     fobj = c.pyapi.number_float(obj)
     dbval = c.pyapi.float_as_double(fobj)
     c.pyapi.decref(fobj)
-    if typ == types.float32:
+    if typ == types.float32 or types == types.float16:
         val = c.builder.fptrunc(dbval,
                                 c.context.get_argument_type(typ))
     else:

diff --git a/numba/core/datamodel/models.py b/numba/core/datamodel/models.py
@@ -360,6 +360,8 @@ class FloatModel(PrimitiveModel):
     def __init__(self, dmm, fe_type):
         if fe_type == types.float32:
             be_type = ir.FloatType()
+        elif fe_type == types.float16:
+            be_type = ir.HalfType()
         elif fe_type == types.float64:
             be_type = ir.DoubleType()
         else:

diff --git a/numba/core/itanium_mangler.py b/numba/core/itanium_mangler.py
@@ -59,6 +59,7 @@
     'unsigned long long': 'y', # unsigned __int64
     '__int128': 'n',
     'unsigned __int128': 'o',
+    'float16' : 'Dh',
     'float': 'f',
     'double': 'd',
     'long double': 'e', # __float80
@@ -78,6 +79,7 @@
     types.int32: 'int',
     types.uint64: 'unsigned long long',
     types.int64: 'long long',
+    types.float16: 'float16',
     types.float32: 'float',
     types.float64: 'double',
 }

diff --git a/numba/core/types/__init__.py b/numba/core/types/__init__.py
@@ -53,6 +53,7 @@
 
 float32 = Float('float32')
 float64 = Float('float64')
+float16 = Float('float16')
 
 complex64 = Complex('complex64', float32)
 complex128 = Complex('complex128', float64)
@@ -86,6 +87,7 @@
 u4 = uint32
 u8 = uint64
 
+f2 = float16
 f4 = float32
 f8 = float64
 

diff --git a/numba/core/typing/builtins.py b/numba/core/typing/builtins.py
@@ -251,6 +251,8 @@ class DivMod(ConcreteTemplate):
 @infer_global(operator.pow)
 class BinOpPower(ConcreteTemplate):
     cases = list(integer_binop_cases)
+    cases += [signature(types.float16, types.float16, op)
+              for op in (types.int32, types.int64, types.uint64)]
     # Ensure that float32 ** int doesn't go through DP computations
     cases += [signature(types.float32, types.float32, op)
               for op in (types.int32, types.int64, types.uint64)]
@@ -265,6 +267,8 @@ class BinOpPower(ConcreteTemplate):
 @infer_global(operator.ipow)
 class BinOpPower(ConcreteTemplate):
     cases = list(integer_binop_cases)
+    cases += [signature(types.float16, types.float16, op)
+              for op in (types.int32, types.int64, types.uint64)]
     # Ensure that float32 ** int doesn't go through DP computations
     cases += [signature(types.float32, types.float32, op)
               for op in (types.int32, types.int64, types.uint64)]
@@ -919,6 +923,7 @@ class Round(ConcreteTemplate):
     cases = [
         signature(types.intp, types.float32),
         signature(types.int64, types.float64),
+        signature(types.float16, types.float16, types.intp),
         signature(types.float32, types.float32, types.intp),
         signature(types.float64, types.float64, types.intp),
     ]

diff --git a/numba/cuda/compiler.py b/numba/cuda/compiler.py
@@ -6,7 +6,6 @@
 import sys
 
 import numpy as np
-
 from numba import _dispatcher
 from numba.core.typing.templates import AbstractTemplate, ConcreteTemplate
 from numba.core import (types, typing, utils, funcdesc, serialize, config,
@@ -816,6 +815,10 @@ def _prepare_args(self, ty, val, stream, retr, kernelargs):
             cval = getattr(ctypes, "c_%s" % ty)(val)
             kernelargs.append(cval)
 
+        elif ty == types.float16:
+            cval = ctypes.c_uint16(np.float16(val).view(np.uint16))
+            kernelargs.append(cval)
+
         elif ty == types.float64:
             cval = ctypes.c_double(val)
             kernelargs.append(cval)

diff --git a/numba/cuda/cudadecl.py b/numba/cuda/cudadecl.py
@@ -249,6 +249,14 @@ class Cuda_fma(ConcreteTemplate):
     ]
 
 
+@register
+class Cuda_hfma(ConcreteTemplate):
+    key = cuda.fp16.hfma
+    cases = [
+        signature(types.float16, types.float16, types.float16, types.float16)
+    ]
+
+
 @register
 class Cuda_cbrt(ConcreteTemplate):
 
@@ -327,6 +335,35 @@ def generic(self, args, kws):
         return signature(a, test, a, a)
 
 
+def _genfp16_unary(l_key):
+    @register
+    class Cuda_fp16_unary(ConcreteTemplate):
+        key = l_key
+
+        cases = [
+            signature(types.float16, types.float16)
+        ]
+    return Cuda_fp16_unary
+
+
+def _genfp16_binary(l_key):
+    @register
+    class Cuda_fp16(ConcreteTemplate):
+        key = l_key
+
+        cases = [
+            signature(types.float16, types.float16, types.float16)
+        ]
+    return Cuda_fp16
+
+
+Cuda_hadd = _genfp16_binary(cuda.fp16.hadd)
+Cuda_hsub = _genfp16_binary(cuda.fp16.hsub)
+Cuda_hmul = _genfp16_binary(cuda.fp16.hmul)
+Cuda_hneg = _genfp16_unary(cuda.fp16.hneg)
+Cuda_habs = _genfp16_unary(cuda.fp16.habs)
+
+
 # generate atomic operations
 def _gen(l_key, supported_types):
     @register
@@ -472,6 +509,29 @@ def resolve_compare_and_swap(self, mod):
         return types.Function(Cuda_atomic_compare_and_swap)
 
 
+@register_attr
+class CudaFp16Template(AttributeTemplate):
+    key = types.Module(cuda.fp16)
+
+    def resolve_hadd(self, mod):
+        return types.Function(Cuda_hadd)
+
+    def resolve_hsub(self, mod):
+        return types.Function(Cuda_hsub)
+
+    def resolve_hmul(self, mod):
+        return types.Function(Cuda_hmul)
+
+    def resolve_hneg(self, mod):
+        return types.Function(Cuda_hneg)
+
+    def resolve_habs(self, mod):
+        return types.Function(Cuda_habs)
+
+    def resolve_hfma(self, mod):
+        return types.Function(Cuda_hfma)
+
+
 @register_attr
 class CudaModuleTemplate(AttributeTemplate):
     key = types.Module(cuda)
@@ -575,6 +635,9 @@ def resolve_nanosleep(self, mod):
     def resolve_atomic(self, mod):
         return types.Module(cuda.atomic)
 
+    def resolve_fp16(self, mod):
+        return types.Module(cuda.fp16)
+
     def resolve_const(self, mod):
         return types.Module(cuda.const)
 

diff --git a/numba/cuda/cudaimpl.py b/numba/cuda/cudaimpl.py
@@ -404,6 +404,85 @@ def ptx_popc(context, builder, sig, args):
 def ptx_fma(context, builder, sig, args):
     return builder.fma(*args)
 
+
+@lower(stubs.fp16.hadd, types.float16, types.float16)
+def ptx_hadd(context, builder, sig, args):
+    hadd_fn_type = ir.FunctionType(ir.HalfType(),
+                                   [ir.HalfType(), ir.HalfType()])
+    hadd_inline_asm = "add.f16 $0,$1,$2;"
+    hadd_inline_constraints = "=h,h,h"
+    hadd_inline = ir.InlineAsm(hadd_fn_type, hadd_inline_asm,
+                               hadd_inline_constraints)
+    arg1 = args[0]
+    arg2 = args[1]
+    return builder.call(hadd_inline, [arg1, arg2])
+
+
+@lower(stubs.fp16.hsub, types.float16, types.float16)
+def ptx_hsub(context, builder, sig, args):
+    hsub_fn_type = ir.FunctionType(ir.HalfType(),
+                                   [ir.HalfType(), ir.HalfType()])
+    hsub_inline_asm = "sub.f16 $0,$1,$2;"
+    hsub_inline_constraints = "=h,h,h"
+    hsub_inline = ir.InlineAsm(hsub_fn_type, hsub_inline_asm,
+                               hsub_inline_constraints)
+    arg1 = args[0]
+    arg2 = args[1]
+    return builder.call(hsub_inline, [arg1, arg2])
+
+
+@lower(stubs.fp16.hmul, types.float16, types.float16)
+def ptx_hmul(context, builder, sig, args):
+    hmul_fn_type = ir.FunctionType(ir.HalfType(),
+                                   [ir.HalfType(), ir.HalfType()])
+    hmul_inline_asm = "mul.f16 $0,$1,$2;"
+    hmul_inline_constraints = "=h,h,h"
+    hmul_inline = ir.InlineAsm(hmul_fn_type, hmul_inline_asm,
+                               hmul_inline_constraints)
+    arg1 = args[0]
+    arg2 = args[1]
+    return builder.call(hmul_inline, [arg1, arg2])
+
+
+@lower(stubs.fp16.hfma, types.float16, types.float16, types.float16)
+def ptx_hfma(context, builder, sig, args):
+    hfma_fn_type = ir.FunctionType(ir.HalfType(),
+                                   [ir.HalfType(),
+                                    ir.HalfType(),
+                                    ir.HalfType()])
+    hfma_inline_asm = "fma.rn.f16 $0,$1,$2,$3;"
+    hfma_inline_constraints = "=h,h,h,h"
+    hfma_inline = ir.InlineAsm(hfma_fn_type, hfma_inline_asm,
+                               hfma_inline_constraints)
+    arg1 = args[0]
+    arg2 = args[1]
+    arg3 = args[2]
+    return builder.call(hfma_inline, [arg1, arg2, arg3])
-    arg1 = args[0]
-    arg2 = args[1]
-    arg3 = args[2]
-    return builder.call(hfma_inline, [arg1, arg2, arg3])
+    return builder.call(hfma_inline, args)
-    arg1 = args[0]
-    arg2 = args[1]
-    arg3 = args[2]
-    return builder.call(hfma_inline, [arg1, arg2, arg3])
+    return builder.call(hfma_inline, args)
+
+
+@lower(stubs.fp16.hneg, types.float16)
+def ptx_hneg(context, builder, sig, args):
+    hneg_fn_type = ir.FunctionType(ir.HalfType(),
+                                   [ir.HalfType()])
+    hneg_inline_asm = "neg.f16 $0,$1;"
+    hneg_inline_constraints = "=h,h"
+    hneg_inline = ir.InlineAsm(hneg_fn_type, hneg_inline_asm,
+                               hneg_inline_constraints)
+    arg1 = args[0]
+    return builder.call(hneg_inline, [arg1])
+
+
+@lower(stubs.fp16.habs, types.float16)
+def ptx_habs(context, builder, sig, args):
+    habs_fn_type = ir.FunctionType(ir.HalfType(),
+                                   [ir.HalfType()])
+    habs_inline_asm = "abs.f16 $0,$1;"
+    habs_inline_constraints = "=h,h"
+    habs_inline = ir.InlineAsm(habs_fn_type, habs_inline_asm,
+                               habs_inline_constraints)
+    arg1 = args[0]
+    return builder.call(habs_inline, [arg1])
+
 # See:
 # https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_cbrt.html#__nv_cbrt
 # https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_cbrtf.html#__nv_cbrtf