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,55 @@ 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
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The following functions are used to operate on 16-bit floating point operands.
+These functions return a 16-bit floating point result.
+
+
+.. 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 16-bit floating point 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 the 16-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 the 16-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/itanium_mangler.py b/numba/core/itanium_mangler.py
@@ -59,6 +59,7 @@
     'unsigned long long': 'y', # unsigned __int64
     '__int128': 'n',
     'unsigned __int128': 'o',
+    'half' : '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: 'half',
     types.float32: 'float',
     types.float64: 'double',
 }

diff --git a/numba/core/typeconv/rules.py b/numba/core/typeconv/rules.py
@@ -33,6 +33,7 @@ def _init_casting_rules(tm):
     tcr.safe_unsafe(types.int16, types.float32)
     tcr.safe_unsafe(types.int32, types.float64)
 
+    tcr.unsafe_unsafe(types.int16, types.float16)
     tcr.unsafe_unsafe(types.int32, types.float32)
     # XXX this is inconsistent with the above; but we want to prefer
     # float64 over int64 when typing a heterogeneous operation,
@@ -41,6 +42,7 @@ def _init_casting_rules(tm):
     tcr.safe_unsafe(types.int64, types.float64)
     tcr.safe_unsafe(types.uint64, types.float64)
 
+    tcr.promote_unsafe(types.float16, types.float32)
     tcr.promote_unsafe(types.float32, types.float64)
 
     tcr.safe(types.float32, types.complex64)

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/cuda/compiler.py b/numba/cuda/compiler.py
@@ -661,6 +661,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
@@ -251,6 +251,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):
 
@@ -329,6 +337,31 @@ 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_binary(ConcreteTemplate):
+        key = l_key
+        cases = [signature(types.float16, types.float16, types.float16)]
+
+    return Cuda_fp16_binary
+
+
+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
@@ -474,6 +507,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)
@@ -577,6 +633,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
@@ -5,12 +5,12 @@
 from llvmlite import ir
 import llvmlite.binding as ll
 
-from numba.core.imputils import Registry
+from numba.core.imputils import Registry, lower_cast
 from numba.core.typing.npydecl import parse_dtype, signature
 from numba.core import types, cgutils
 from .cudadrv import nvvm
 from numba import cuda
-from numba.cuda import nvvmutils, stubs
+from numba.cuda import nvvmutils, stubs, errors
 from numba.cuda.types import dim3, grid_group, CUDADispatcher
 
 
@@ -405,6 +405,124 @@ def ptx_popc(context, builder, sig, args):
 def ptx_fma(context, builder, sig, args):
     return builder.fma(*args)
 
+
+def float16_float_ty_constraint(bitwidth):
+    typemap = {32: ('f32', 'f'), 64: ('f64', 'd')}
+
+    try:
+        return typemap[bitwidth]
+    except KeyError:
+        msg = f"Conversion between float16 and float{bitwidth} unsupported"
+        raise errors.CudaLoweringError(msg)
+
+
+@lower_cast(types.float16, types.Float)
+def float16_to_float_cast(context, builder, fromty, toty, val):
+    if fromty.bitwidth == toty.bitwidth:
+        return val
+
+    ty, constraint = float16_float_ty_constraint(toty.bitwidth)
+
+    fnty = ir.FunctionType(context.get_value_type(toty), [ir.IntType(16)])
+    asm = ir.InlineAsm(fnty, f"cvt.{ty}.f16 $0, $1;", f"={constraint},h")
+    return builder.call(asm, [val])
+
+
+@lower_cast(types.Float, types.float16)
+def float_to_float16_cast(context, builder, fromty, toty, val):
+    if fromty.bitwidth == toty.bitwidth:
+        return val
+
+    ty, constraint = float16_float_ty_constraint(fromty.bitwidth)
+
+    fnty = ir.FunctionType(ir.IntType(16), [context.get_value_type(fromty)])
+    asm = ir.InlineAsm(fnty, f"cvt.rn.f16.{ty} $0, $1;", f"=h,{constraint}")
+    return builder.call(asm, [val])
+
+
+def float16_int_constraint(bitwidth):
+    typemap = { 8: 'c', 16: 'h', 32: 'r', 64: 'l' }
+
+    try:
+        return typemap[bitwidth]
+    except KeyError:
+        msg = f"Conversion between float16 and int{bitwidth} unsupported"
+        raise errors.CudaLoweringError(msg)
+
+
+@lower_cast(types.float16, types.Integer)
+def float16_to_integer_cast(context, builder, fromty, toty, val):
+    bitwidth = toty.bitwidth
+    constraint = float16_int_constraint(bitwidth)
+    signedness = 's' if toty.signed else 'u'
+
+    fnty = ir.FunctionType(context.get_value_type(toty), [ir.IntType(16)])
+    asm = ir.InlineAsm(fnty,
+                       f"cvt.rni.{signedness}{bitwidth}.f16 $0, $1;",
+                       f"={constraint},h")
+    return builder.call(asm, [val])
+
+
+@lower_cast(types.Integer, types.float16)
+@lower_cast(types.IntegerLiteral, types.float16)
+def integer_to_float16_cast(context, builder, fromty, toty, val):
+    bitwidth = fromty.bitwidth
+    constraint = float16_int_constraint(bitwidth)
+    signedness = 's' if fromty.signed else 'u'
+
+    fnty = ir.FunctionType(ir.IntType(16),
+                           [context.get_value_type(fromty)])
+    asm = ir.InlineAsm(fnty,
+                       f"cvt.rn.f16.{signedness}{bitwidth} $0, $1;",
+                       f"=h,{constraint}")
+    return builder.call(asm, [val])
+
+
+def lower_fp16_binary(fn, op):
+    @lower(fn, types.float16, types.float16)
+    def ptx_fp16_binary(context, builder, sig, args):
+        fnty = ir.FunctionType(ir.IntType(16),
+                               [ir.IntType(16), ir.IntType(16)])
+        asm = ir.InlineAsm(fnty, f'{op}.f16 $0,$1,$2;', '=h,h,h')
+        return builder.call(asm, args)
+
+
+lower_fp16_binary(stubs.fp16.hadd, 'add')
+lower_fp16_binary(stubs.fp16.hsub, 'sub')
+lower_fp16_binary(stubs.fp16.hmul, 'mul')
+
+
+@lower(stubs.fp16.hneg, types.float16)
+def ptx_fp16_hneg(context, builder, sig, args):
+    fnty = ir.FunctionType(ir.IntType(16), [ir.IntType(16)])
+    asm = ir.InlineAsm(fnty, 'neg.f16 $0, $1;', '=h,h')
+    return builder.call(asm, args)
+
+
+@lower(stubs.fp16.habs, types.float16)
+def ptx_fp16_habs(context, builder, sig, args):
+    if cuda.runtime.get_version() < (10, 2):
+        # CUDA < 10.2 does not support abs.f16. For these versions, we mask
+        # off the sign bit to compute abs instead. We determine whether or
+        # not to do this based on the runtime version so that our behaviour
+        # is consistent with the version of NVVM we're using to go from
+        # NVVM IR -> PTX.
+        inst = 'and.b16 $0, $1, 0x7FFF;'
+    else:
+        inst = 'abs.f16 $0, $1;'
+
+    fnty = ir.FunctionType(ir.IntType(16), [ir.IntType(16)])
+    asm = ir.InlineAsm(fnty, inst, '=h,h')
+    return builder.call(asm, 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)]
+    fnty = ir.FunctionType(ir.IntType(16), argtys)
+    asm = ir.InlineAsm(fnty, "fma.rn.f16 $0,$1,$2,$3;", "=h,h,h,h")
+    return builder.call(asm, args)
+
 # 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

diff --git a/numba/cuda/device_init.py b/numba/cuda/device_init.py
@@ -6,7 +6,7 @@
                     vote_sync_intrinsic, match_any_sync, match_all_sync,
                     threadfence_block, threadfence_system,
                     threadfence, selp, popc, brev, clz, ffs, fma, cbrt,
-                    cg, activemask, lanemask_lt, nanosleep)
+                    cg, activemask, lanemask_lt, nanosleep, fp16)
 from .cudadrv.error import CudaSupportError
 from numba.cuda.cudadrv.driver import (BaseCUDAMemoryManager,
                                        HostOnlyCUDAMemoryManager,

diff --git a/numba/cuda/errors.py b/numba/cuda/errors.py
@@ -1,4 +1,5 @@
 import numbers
+from numba.core.errors import LoweringError
 
 
 class KernelRuntimeError(RuntimeError):
@@ -12,6 +13,10 @@ def __init__(self, msg, tid=None, ctaid=None):
         super(KernelRuntimeError, self).__init__(msg)
 
 
+class CudaLoweringError(LoweringError):
+    pass
+
+
 _launch_help_url = ("https://numba.pydata.org/numba-doc/"
                     "latest/cuda/kernels.html#kernel-invocation")
 missing_launch_config_msg = """

diff --git a/numba/cuda/models.py b/numba/cuda/models.py
@@ -1,5 +1,6 @@
 from llvmlite import ir
 
+from numba.core.datamodel.registry import register_default
 from numba.core.extending import register_model, models
 from numba.core import types
 from numba.cuda.types import Dim3, GridGroup, CUDADispatcher
@@ -23,4 +24,18 @@ def __init__(self, dmm, fe_type):
         super().__init__(dmm, fe_type, be_type)
 
 
+@register_default(types.Float)
+class FloatModel(models.PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        if fe_type == types.float16:
+            be_type = ir.IntType(16)
+        elif fe_type == types.float32:
+            be_type = ir.FloatType()
+        elif fe_type == types.float64:
+            be_type = ir.DoubleType()
+        else:
+            raise NotImplementedError(fe_type)
+        super(FloatModel, self).__init__(dmm, fe_type, be_type)
+
+
 register_model(CUDADispatcher)(models.OpaqueModel)