Support for CUDA fp16 math functions (part 1) by testhound · Pull Request #8605 · numba/numba

testhound · 2022-11-17T04:46:22Z

This PR adds support for a subset of the python math functions for fp16. The supported subset was chosen because they map directly to fp16 math intrinsics that have been implemented already.

…to fp16 intrinsics

testhound · 2022-11-17T04:52:24Z

gpuci run tests

stuartarchibald · 2022-11-17T14:00:05Z

@testhound thanks for the patch. I've altered the PT title so that it reflects the change is for CUDA and for fp16 types. This is so that the merged PR will appear in the change log with a more specific title.

gmarkall

Thanks for the PR. I'm not sure I understand how this should work - I see that lowering is added, but no additional typing. For example, if I try compiling a call to a math function on a float16 and check the typing:

from numba import cuda, types
import math


@cuda.jit(types.void(types.float16[::1]))
def f(x):
    print(math.sin(x[0]))


f.inspect_types()

I see

$ python repro.py 
_ZN6cudapy8__main__1fB2v1B106cw51cXTLSUwHBinCqbbgUAAGBlq82ILSCEQYkgSQBFCjFSaBZJtttTo4sahbKRjoKKiDvAVKN0AuNEDVYWGDJATUC_2bQZ12oCAA_3d_3dE5ArrayIDhLi1E1C7mutable7alignedE (array(float16, 1d, C),)
--------------------------------------------------------------------------------
# File: /home/gmarkall/numbadev/issues/8605/repro.py
# --- LINE 5 --- 

@cuda.jit(types.void(types.float16[::1]))

# --- LINE 6 --- 

def f(x):

    # --- LINE 7 --- 
    # label 0
    #   x = arg(0, name=x)  :: array(float16, 1d, C)
    #   $2load_global.0 = global(print: <built-in function print>)  :: Function(<built-in function print>)
    #   del $2load_global.0
    #   $4load_global.1 = global(math: <module 'math' from '/home/gmarkall/mambaforge/envs/numbadev/lib/python3.10/lib-dynload/math.cpython-310-x86_64-linux-gnu.so'>)  :: Module(<module 'math' from '/home/gmarkall/mambaforge/envs/numbadev/lib/python3.10/lib-dynload/math.cpython-310-x86_64-linux-gnu.so'>)
    #   $6load_method.2 = getattr(value=$4load_global.1, attr=sin)  :: Function(<built-in function sin>)
    #   del $4load_global.1
    #   $const10.4 = const(int, 0)  :: Literal[int](0)
    #   $12binary_subscr.5 = static_getitem(value=x, index=0, index_var=$const10.4, fn=<built-in function getitem>)  :: float16
    #   del x
    #   del $const10.4
    #   $14call_method.6 = call $6load_method.2($12binary_subscr.5, func=$6load_method.2, args=[Var($12binary_subscr.5, repro.py:7)], kws=(), vararg=None, varkwarg=None,  :: (float32,) -> float32
    #   del $6load_method.2
    #   del $12binary_subscr.5
    #   print($14call_method.6)
    #   del $14call_method.6
    #   $16call_function.7 = const(NoneType, None)  :: none
    #   del $16call_function.7
    #   $const20.8 = const(NoneType, None)  :: none
    #   $22return_value.9 = cast(value=$const20.8)  :: none
    #   del $const20.8
    #   return $22return_value.9

    print(math.sin(x[0]))


================================================================================

The argument to math.sin is typed as accepting a float32.

However, I don't think that the right way to go would be to add typing - now that the high-level API is supported in CUDA, I think it would be better to try writing these using @overload instead.

testhound · 2022-11-17T18:00:04Z

@gmarkall okay I understand now. As a quick test I added typing for "math.sin" in cuda/cmath.py and the IR now generates code using fp16 instead of fp32. Can you clarify why you believe it is better to use "@overload" instead of typing? I don't understand how this conflicts with the high-level API as you mentioned above.

…h fp16

testhound · 2022-11-18T02:25:38Z

@gmarkall I added typing for all the fp16 math functions in the latest iteration. I spoke with @sklam about using "@overload" and he indicated that there is a registration issue with using "@overload" and "@infer_global". He indicated that we need to use either all "@overload" or all "@infer_global".

This problem is similar to the one we had when adding typing for fp16 operators and intrinsics. We needed to use AbstractTemplate for operators and ConcreteTemplate

testhound · 2022-11-18T02:26:08Z

gpuci run tests

gmarkall

Thanks for the updates. This now appears to use the FP16 implementations in all cases. There are a couple of comments on the diff - in addition to those, the lowering implementations for the math module should be in numba.cuda.mathimpl, rather than numba.cuda.cudaimpl - cudaimpl is where the implementations of the functions in the numba.cuda module are.

Regarding the conversion to use overload - I was going to suggest trying to convert these functions to the overload throughout Numba so that the ConcreteTemplate / overload issue isn't a problem, but given that @stuartarchibald is refactoring a lot of the math function implementations this is probably not productive at this time. So, if the comments above and on the diff can be addressed / resolved, then I think this will be good to approve.

gmarkall · 2022-11-21T10:25:42Z

+@infer_global(math.exp)
+@infer_global(math.fabs)
+@infer_global(math.trunc)
+class Math_sin(ConcreteTemplate):


Perhaps a more appropriate name:

Suggested change

class Math_sin(ConcreteTemplate):

class Math_unary_with_fp16(ConcreteTemplate):

gmarkall · 2022-11-21T10:56:45Z

+    def fp16_sin(x):
+        return cuda.fp16.hsin(x)
+
+    return context.compile_internal(builder, fp16_sin, sig, args, )


I think this copy-pastes a type from the hdiv implementation (which should also be fixed similarly), along with all other instances:

Suggested change

return context.compile_internal(builder, fp16_sin, sig, args, )

return context.compile_internal(builder, fp16_sin, sig, args)

testhound · 2022-11-21T16:19:42Z

@gmarkall I addresses all your comments in this latest change.

testhound · 2022-11-21T16:20:08Z

gpuci run tests

testhound · 2022-11-21T18:06:36Z

@gmarkall I added implementations of math.trunc for float32 and float64.

testhound · 2022-11-21T18:06:51Z

gpuci run tests

testhound · 2022-11-22T14:42:58Z

gpuci run tests

gmarkall

Many thanks for the patches!

testhound added 2 commits November 16, 2022 22:50

Add lowering of initial standard of standard math functions that map …

3130c30

…to fp16 intrinsics

Refactor fp16 tests to skip whe necessary

25a87f1

testhound requested a review from gmarkall as a code owner November 17, 2022 04:46

stuartarchibald changed the title ~~Support for python math function part 1~~ Support for CUDA fp16 math functions (part 1) Nov 17, 2022

gmarkall requested changes Nov 17, 2022

View reviewed changes

gmarkall self-assigned this Nov 17, 2022

gmarkall added 4 - Waiting on author Waiting for author to respond to review CUDA CUDA related issue/PR labels Nov 17, 2022

testhound added 3 commits November 17, 2022 21:14

Add support for fp16 math.trunc

f06a04a

Add typing information for fp16 math functions

f98a863

Add testcase for math.trunc and add support for greater tolerance wit…

06366d6

…h fp16

gmarkall added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 4 - Waiting on author Waiting for author to respond to review labels Nov 21, 2022

gmarkall requested changes Nov 21, 2022

View reviewed changes

gmarkall added 4 - Waiting on author Waiting for author to respond to review and removed 4 - Waiting on reviewer Waiting for reviewer to respond to author labels Nov 21, 2022

Move lowering to mathimpl and address formatting issues

2f01c9a

Add implementations for math.trunc for float32/float64 for cuda target

25de1a3

gmarkall added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 4 - Waiting on author Waiting for author to respond to review labels Nov 22, 2022

Remove fp16 tan implementation

b500268

gmarkall added this to the Numba 0.57 RC milestone Nov 22, 2022

gmarkall added 5 - Ready to merge Review and testing done, is ready to merge and removed 4 - Waiting on reviewer Waiting for reviewer to respond to author labels Nov 22, 2022

gmarkall approved these changes Nov 22, 2022

View reviewed changes

sklam merged commit 7282635 into numba:main Nov 23, 2022

This was referenced Nov 30, 2022

CUDA fail on main: type numpy.float32 doesn't define __trunc__ method #8628

Closed

Fix #8628: Don't test math.trunc with non-float64 NumPy scalars #8630

Merged

	class Math_sin(ConcreteTemplate):
	class Math_unary_with_fp16(ConcreteTemplate):

	return context.compile_internal(builder, fp16_sin, sig, args, )
	return context.compile_internal(builder, fp16_sin, sig, args)

Conversation

testhound commented Nov 17, 2022

Uh oh!

testhound commented Nov 17, 2022

Uh oh!

stuartarchibald commented Nov 17, 2022

Uh oh!

gmarkall left a comment

Choose a reason for hiding this comment

Uh oh!

testhound commented Nov 17, 2022

Uh oh!

testhound commented Nov 18, 2022

Uh oh!

testhound commented Nov 18, 2022

Uh oh!

gmarkall left a comment

Choose a reason for hiding this comment

Uh oh!

gmarkall Nov 21, 2022

Choose a reason for hiding this comment

Uh oh!

gmarkall Nov 21, 2022

Choose a reason for hiding this comment

Uh oh!

testhound commented Nov 21, 2022

Uh oh!

testhound commented Nov 21, 2022

Uh oh!

testhound commented Nov 21, 2022

Uh oh!

testhound commented Nov 21, 2022

Uh oh!

testhound commented Nov 22, 2022

Uh oh!

gmarkall left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants