Skip to content

Commit 23918f0

Browse files
authored
Merge pull request #7460 from testhound/testhound/fp16-support
Testhound/fp16 support
2 parents 4cc738b + 20c1d4f commit 23918f0

17 files changed

Lines changed: 787 additions & 37 deletions

File tree

‎docs/source/cuda-reference/kernel.rst‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -463,6 +463,55 @@ precision parts of the CUDA Toolkit documentation.
463463
``cbrt`` and ``cbrtf`` in the C api. Supports float32, and float64 arguments
464464
only.
465465

466+
16-bit Floating Point Intrinsics
467+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
468+
469+
The following functions are used to operate on 16-bit floating point operands.
470+
These functions return a 16-bit floating point result.
471+
472+
473+
.. function:: numba.cuda.fp16.hfma (a, b, c)
474+
475+
Perform the fused multiply-add operation ``(a * b) + c`` on 16-bit
476+
floating point arguments in round to nearest mode. Maps to the ``fma.rn.f16``
477+
PTX instruction.
478+
479+
Returns the 16-bit floating point result of the fused multiply-add.
480+
481+
.. function:: numba.cuda.fp16.hadd (a, b)
482+
483+
Perform the add operation ``a + b`` on 16-bit floating point arguments in
484+
round to nearest mode. Maps to the ``add.f16`` PTX instruction.
485+
486+
Returns the 16-bit floating point result of the addition.
487+
488+
.. function:: numba.cuda.fp16.hsub (a, b)
489+
490+
Perform the subtract operation ``a - b`` on 16-bit floating point arguments in
491+
round to nearest mode. Maps to the ``sub.f16`` PTX instruction.
492+
493+
Returns the 16-bit floating point result of the subtraction.
494+
495+
.. function:: numba.cuda.fp16.hmul (a, b)
496+
497+
Perform the multiply operation ``a * b`` on 16-bit floating point arguments in
498+
round to nearest mode. Maps to the ``mul.f16`` PTX instruction.
499+
500+
Returns the 16-bit floating point result of the multiplication.
501+
502+
.. function:: numba.cuda.fp16.hneg (a)
503+
504+
Perform the negation operation ``-a`` on the 16-bit floating point argument.
505+
Maps to the ``neg.f16`` PTX instruction.
506+
507+
Returns the 16-bit floating point result of the negation.
508+
509+
.. function:: numba.cuda.fp16.habs (a)
510+
511+
Perform the absolute value operation ``|a|`` on the 16-bit floating point argument.
512+
Maps to the ``abs.f16`` PTX instruction.
513+
514+
Returns the 16-bit floating point result of the absolute value operation.
466515

467516
Control Flow Instructions
468517
~~~~~~~~~~~~~~~~~~~~~~~~~

‎numba/core/itanium_mangler.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959
'unsigned long long': 'y', # unsigned __int64
6060
'__int128': 'n',
6161
'unsigned __int128': 'o',
62+
'half' : 'Dh',
6263
'float': 'f',
6364
'double': 'd',
6465
'long double': 'e', # __float80
@@ -78,6 +79,7 @@
7879
types.int32: 'int',
7980
types.uint64: 'unsigned long long',
8081
types.int64: 'long long',
82+
types.float16: 'half',
8183
types.float32: 'float',
8284
types.float64: 'double',
8385
}

‎numba/core/typeconv/rules.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ def _init_casting_rules(tm):
3333
tcr.safe_unsafe(types.int16, types.float32)
3434
tcr.safe_unsafe(types.int32, types.float64)
3535

36+
tcr.unsafe_unsafe(types.int16, types.float16)
3637
tcr.unsafe_unsafe(types.int32, types.float32)
3738
# XXX this is inconsistent with the above; but we want to prefer
3839
# float64 over int64 when typing a heterogeneous operation,
@@ -41,6 +42,7 @@ def _init_casting_rules(tm):
4142
tcr.safe_unsafe(types.int64, types.float64)
4243
tcr.safe_unsafe(types.uint64, types.float64)
4344

45+
tcr.promote_unsafe(types.float16, types.float32)
4446
tcr.promote_unsafe(types.float32, types.float64)
4547

4648
tcr.safe(types.float32, types.complex64)

‎numba/core/types/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353

5454
float32 = Float('float32')
5555
float64 = Float('float64')
56+
float16 = Float('float16')
5657

5758
complex64 = Complex('complex64', float32)
5859
complex128 = Complex('complex128', float64)
@@ -86,6 +87,7 @@
8687
u4 = uint32
8788
u8 = uint64
8889

90+
f2 = float16
8991
f4 = float32
9092
f8 = float64
9193

‎numba/cuda/compiler.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -661,6 +661,10 @@ def _prepare_args(self, ty, val, stream, retr, kernelargs):
661661
cval = getattr(ctypes, "c_%s" % ty)(val)
662662
kernelargs.append(cval)
663663

664+
elif ty == types.float16:
665+
cval = ctypes.c_uint16(np.float16(val).view(np.uint16))
666+
kernelargs.append(cval)
667+
664668
elif ty == types.float64:
665669
cval = ctypes.c_double(val)
666670
kernelargs.append(cval)

‎numba/cuda/cudadecl.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,14 @@ class Cuda_fma(ConcreteTemplate):
251251
]
252252

253253

254+
@register
255+
class Cuda_hfma(ConcreteTemplate):
256+
key = cuda.fp16.hfma
257+
cases = [
258+
signature(types.float16, types.float16, types.float16, types.float16)
259+
]
260+
261+
254262
@register
255263
class Cuda_cbrt(ConcreteTemplate):
256264

@@ -329,6 +337,31 @@ def generic(self, args, kws):
329337
return signature(a, test, a, a)
330338

331339

340+
def _genfp16_unary(l_key):
341+
@register
342+
class Cuda_fp16_unary(ConcreteTemplate):
343+
key = l_key
344+
cases = [signature(types.float16, types.float16)]
345+
346+
return Cuda_fp16_unary
347+
348+
349+
def _genfp16_binary(l_key):
350+
@register
351+
class Cuda_fp16_binary(ConcreteTemplate):
352+
key = l_key
353+
cases = [signature(types.float16, types.float16, types.float16)]
354+
355+
return Cuda_fp16_binary
356+
357+
358+
Cuda_hadd = _genfp16_binary(cuda.fp16.hadd)
359+
Cuda_hsub = _genfp16_binary(cuda.fp16.hsub)
360+
Cuda_hmul = _genfp16_binary(cuda.fp16.hmul)
361+
Cuda_hneg = _genfp16_unary(cuda.fp16.hneg)
362+
Cuda_habs = _genfp16_unary(cuda.fp16.habs)
363+
364+
332365
# generate atomic operations
333366
def _gen(l_key, supported_types):
334367
@register
@@ -474,6 +507,29 @@ def resolve_compare_and_swap(self, mod):
474507
return types.Function(Cuda_atomic_compare_and_swap)
475508

476509

510+
@register_attr
511+
class CudaFp16Template(AttributeTemplate):
512+
key = types.Module(cuda.fp16)
513+
514+
def resolve_hadd(self, mod):
515+
return types.Function(Cuda_hadd)
516+
517+
def resolve_hsub(self, mod):
518+
return types.Function(Cuda_hsub)
519+
520+
def resolve_hmul(self, mod):
521+
return types.Function(Cuda_hmul)
522+
523+
def resolve_hneg(self, mod):
524+
return types.Function(Cuda_hneg)
525+
526+
def resolve_habs(self, mod):
527+
return types.Function(Cuda_habs)
528+
529+
def resolve_hfma(self, mod):
530+
return types.Function(Cuda_hfma)
531+
532+
477533
@register_attr
478534
class CudaModuleTemplate(AttributeTemplate):
479535
key = types.Module(cuda)
@@ -577,6 +633,9 @@ def resolve_nanosleep(self, mod):
577633
def resolve_atomic(self, mod):
578634
return types.Module(cuda.atomic)
579635

636+
def resolve_fp16(self, mod):
637+
return types.Module(cuda.fp16)
638+
580639
def resolve_const(self, mod):
581640
return types.Module(cuda.const)
582641

‎numba/cuda/cudaimpl.py‎

Lines changed: 120 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,12 @@
55
from llvmlite import ir
66
import llvmlite.binding as ll
77

8-
from numba.core.imputils import Registry
8+
from numba.core.imputils import Registry, lower_cast
99
from numba.core.typing.npydecl import parse_dtype, signature
1010
from numba.core import types, cgutils
1111
from .cudadrv import nvvm
1212
from numba import cuda
13-
from numba.cuda import nvvmutils, stubs
13+
from numba.cuda import nvvmutils, stubs, errors
1414
from numba.cuda.types import dim3, grid_group, CUDADispatcher
1515

1616

@@ -405,6 +405,124 @@ def ptx_popc(context, builder, sig, args):
405405
def ptx_fma(context, builder, sig, args):
406406
return builder.fma(*args)
407407

408+
409+
def float16_float_ty_constraint(bitwidth):
410+
typemap = {32: ('f32', 'f'), 64: ('f64', 'd')}
411+
412+
try:
413+
return typemap[bitwidth]
414+
except KeyError:
415+
msg = f"Conversion between float16 and float{bitwidth} unsupported"
416+
raise errors.CudaLoweringError(msg)
417+
418+
419+
@lower_cast(types.float16, types.Float)
420+
def float16_to_float_cast(context, builder, fromty, toty, val):
421+
if fromty.bitwidth == toty.bitwidth:
422+
return val
423+
424+
ty, constraint = float16_float_ty_constraint(toty.bitwidth)
425+
426+
fnty = ir.FunctionType(context.get_value_type(toty), [ir.IntType(16)])
427+
asm = ir.InlineAsm(fnty, f"cvt.{ty}.f16 $0, $1;", f"={constraint},h")
428+
return builder.call(asm, [val])
429+
430+
431+
@lower_cast(types.Float, types.float16)
432+
def float_to_float16_cast(context, builder, fromty, toty, val):
433+
if fromty.bitwidth == toty.bitwidth:
434+
return val
435+
436+
ty, constraint = float16_float_ty_constraint(fromty.bitwidth)
437+
438+
fnty = ir.FunctionType(ir.IntType(16), [context.get_value_type(fromty)])
439+
asm = ir.InlineAsm(fnty, f"cvt.rn.f16.{ty} $0, $1;", f"=h,{constraint}")
440+
return builder.call(asm, [val])
441+
442+
443+
def float16_int_constraint(bitwidth):
444+
typemap = { 8: 'c', 16: 'h', 32: 'r', 64: 'l' }
445+
446+
try:
447+
return typemap[bitwidth]
448+
except KeyError:
449+
msg = f"Conversion between float16 and int{bitwidth} unsupported"
450+
raise errors.CudaLoweringError(msg)
451+
452+
453+
@lower_cast(types.float16, types.Integer)
454+
def float16_to_integer_cast(context, builder, fromty, toty, val):
455+
bitwidth = toty.bitwidth
456+
constraint = float16_int_constraint(bitwidth)
457+
signedness = 's' if toty.signed else 'u'
458+
459+
fnty = ir.FunctionType(context.get_value_type(toty), [ir.IntType(16)])
460+
asm = ir.InlineAsm(fnty,
461+
f"cvt.rni.{signedness}{bitwidth}.f16 $0, $1;",
462+
f"={constraint},h")
463+
return builder.call(asm, [val])
464+
465+
466+
@lower_cast(types.Integer, types.float16)
467+
@lower_cast(types.IntegerLiteral, types.float16)
468+
def integer_to_float16_cast(context, builder, fromty, toty, val):
469+
bitwidth = fromty.bitwidth
470+
constraint = float16_int_constraint(bitwidth)
471+
signedness = 's' if fromty.signed else 'u'
472+
473+
fnty = ir.FunctionType(ir.IntType(16),
474+
[context.get_value_type(fromty)])
475+
asm = ir.InlineAsm(fnty,
476+
f"cvt.rn.f16.{signedness}{bitwidth} $0, $1;",
477+
f"=h,{constraint}")
478+
return builder.call(asm, [val])
479+
480+
481+
def lower_fp16_binary(fn, op):
482+
@lower(fn, types.float16, types.float16)
483+
def ptx_fp16_binary(context, builder, sig, args):
484+
fnty = ir.FunctionType(ir.IntType(16),
485+
[ir.IntType(16), ir.IntType(16)])
486+
asm = ir.InlineAsm(fnty, f'{op}.f16 $0,$1,$2;', '=h,h,h')
487+
return builder.call(asm, args)
488+
489+
490+
lower_fp16_binary(stubs.fp16.hadd, 'add')
491+
lower_fp16_binary(stubs.fp16.hsub, 'sub')
492+
lower_fp16_binary(stubs.fp16.hmul, 'mul')
493+
494+
495+
@lower(stubs.fp16.hneg, types.float16)
496+
def ptx_fp16_hneg(context, builder, sig, args):
497+
fnty = ir.FunctionType(ir.IntType(16), [ir.IntType(16)])
498+
asm = ir.InlineAsm(fnty, 'neg.f16 $0, $1;', '=h,h')
499+
return builder.call(asm, args)
500+
501+
502+
@lower(stubs.fp16.habs, types.float16)
503+
def ptx_fp16_habs(context, builder, sig, args):
504+
if cuda.runtime.get_version() < (10, 2):
505+
# CUDA < 10.2 does not support abs.f16. For these versions, we mask
506+
# off the sign bit to compute abs instead. We determine whether or
507+
# not to do this based on the runtime version so that our behaviour
508+
# is consistent with the version of NVVM we're using to go from
509+
# NVVM IR -> PTX.
510+
inst = 'and.b16 $0, $1, 0x7FFF;'
511+
else:
512+
inst = 'abs.f16 $0, $1;'
513+
514+
fnty = ir.FunctionType(ir.IntType(16), [ir.IntType(16)])
515+
asm = ir.InlineAsm(fnty, inst, '=h,h')
516+
return builder.call(asm, args)
517+
518+
519+
@lower(stubs.fp16.hfma, types.float16, types.float16, types.float16)
520+
def ptx_hfma(context, builder, sig, args):
521+
argtys = [ir.IntType(16), ir.IntType(16), ir.IntType(16)]
522+
fnty = ir.FunctionType(ir.IntType(16), argtys)
523+
asm = ir.InlineAsm(fnty, "fma.rn.f16 $0,$1,$2,$3;", "=h,h,h,h")
524+
return builder.call(asm, args)
525+
408526
# See:
409527
# https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_cbrt.html#__nv_cbrt
410528
# https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_cbrtf.html#__nv_cbrtf

‎numba/cuda/device_init.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
vote_sync_intrinsic, match_any_sync, match_all_sync,
77
threadfence_block, threadfence_system,
88
threadfence, selp, popc, brev, clz, ffs, fma, cbrt,
9-
cg, activemask, lanemask_lt, nanosleep)
9+
cg, activemask, lanemask_lt, nanosleep, fp16)
1010
from .cudadrv.error import CudaSupportError
1111
from numba.cuda.cudadrv.driver import (BaseCUDAMemoryManager,
1212
HostOnlyCUDAMemoryManager,

‎numba/cuda/errors.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import numbers
2+
from numba.core.errors import LoweringError
23

34

45
class KernelRuntimeError(RuntimeError):
@@ -12,6 +13,10 @@ def __init__(self, msg, tid=None, ctaid=None):
1213
super(KernelRuntimeError, self).__init__(msg)
1314

1415

16+
class CudaLoweringError(LoweringError):
17+
pass
18+
19+
1520
_launch_help_url = ("https://numba.pydata.org/numba-doc/"
1621
"latest/cuda/kernels.html#kernel-invocation")
1722
missing_launch_config_msg = """

‎numba/cuda/models.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from llvmlite import ir
22

3+
from numba.core.datamodel.registry import register_default
34
from numba.core.extending import register_model, models
45
from numba.core import types
56
from numba.cuda.types import Dim3, GridGroup, CUDADispatcher
@@ -23,4 +24,18 @@ def __init__(self, dmm, fe_type):
2324
super().__init__(dmm, fe_type, be_type)
2425

2526

27+
@register_default(types.Float)
28+
class FloatModel(models.PrimitiveModel):
29+
def __init__(self, dmm, fe_type):
30+
if fe_type == types.float16:
31+
be_type = ir.IntType(16)
32+
elif fe_type == types.float32:
33+
be_type = ir.FloatType()
34+
elif fe_type == types.float64:
35+
be_type = ir.DoubleType()
36+
else:
37+
raise NotImplementedError(fe_type)
38+
super(FloatModel, self).__init__(dmm, fe_type, be_type)
39+
40+
2641
register_model(CUDADispatcher)(models.OpaqueModel)

0 commit comments

Comments
 (0)