|
5 | 5 | from llvmlite import ir |
6 | 6 | import llvmlite.binding as ll |
7 | 7 |
|
8 | | -from numba.core.imputils import Registry |
| 8 | +from numba.core.imputils import Registry, lower_cast |
9 | 9 | from numba.core.typing.npydecl import parse_dtype, signature |
10 | 10 | from numba.core import types, cgutils |
11 | 11 | from .cudadrv import nvvm |
12 | 12 | from numba import cuda |
13 | | -from numba.cuda import nvvmutils, stubs |
| 13 | +from numba.cuda import nvvmutils, stubs, errors |
14 | 14 | from numba.cuda.types import dim3, grid_group, CUDADispatcher |
15 | 15 |
|
16 | 16 |
|
@@ -405,6 +405,124 @@ def ptx_popc(context, builder, sig, args): |
405 | 405 | def ptx_fma(context, builder, sig, args): |
406 | 406 | return builder.fma(*args) |
407 | 407 |
|
| 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 | + |
408 | 526 | # See: |
409 | 527 | # https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_cbrt.html#__nv_cbrt |
410 | 528 | # https://docs.nvidia.com/cuda/libdevice-users-guide/__nv_cbrtf.html#__nv_cbrtf |
|
0 commit comments