ENH: more ufuncs to API std #124

tylerjereddy · 2022-11-03T00:06:57Z

this uses some templating code to write 4 new API conformant ufuncs: round, trunc, ceil, and floor. (see: https://github.com/tylerjereddy/pykokkos_ufunc_generator )
note that because these are automatically generated they have some verbosity/lack of docstring quirks, bit of an experiment..
switch to the latest tagged release of the array API test suite rather than some random hash
strangely, the API tests added here don't actually seem to need to compile the workunits--they may focus on the special use case for intergers I suppose; we may want to supplement with a few floating point examples perhaps; see also related discussion in
Loosely assert for functions that produce ints but internally might use floats data-apis/array-api-tests#154

tylerjereddy · 2022-11-03T00:09:34Z

pykokkos/lib/ufuncs.py

+        for j in range(view.extent(2)):
+            out[tid][i][j] = floor(view[tid][i][j])
+
+def floor(view):


lack of docstring is just the result of me not worrying about the docstrings in the auto-templating code I linked... we could add some into the template

There's also the question of whether it makes sense to commit auto-generated code vs. the stuff that generates it; conversely, committing to a templating engine could be annoying as well. Certainly, NumPy used its own custom templating language for years for generating the ufunc inner loops. I think they're starting to transition to C++ these days, at least in some limited areas.

I think we should add the docstring in this PR. Regarding committing the templating engine, I would say not to do that for now. Since we're planning on getting rid of type annotations in the future, writing these kernels should be less tedious in the future.

tylerjereddy · 2022-11-03T00:10:28Z

pykokkos/lib/ufuncs.py

+def ceil_impl_3d_float(tid: int, view: pk.View3D[pk.float], out: pk.View3D[pk.float]):
+    for i in range(view.extent(1)):
+        for j in range(view.extent(2)):
+            out[tid][i][j] = ceil(view[tid][i][j])


Is this right? I'm afraid we may have to add some custom tests to supplement b/c I'm not seeing pk_cpp generated by test suite for ceil, floor, round, trunc.. (they're just hitting the integer special case I think)

round should also work once we add it to the set of allowed math functions

NaderAlAwar

I think we just need to add round to the set of math functions and some short docstrings to the ufuncs

NaderAlAwar · 2022-11-29T19:02:07Z

pykokkos/lib/ufuncs.py

+def round_impl_3d_float(tid: int, view: pk.View3D[pk.float], out: pk.View3D[pk.float]):
+    for i in range(view.extent(1)):
+        for j in range(view.extent(2)):
+            out[tid][i][j] = round(view[tid][i][j])


I don't think this will work, see

pykokkos/pykokkos/core/visitors/visitors_util.py

Line 70 in a392821

math_functions: Set = {

for a list of functions that get mapped to their cmath counterparts

Actually there is a round in cmath, it's just not part of this set. Could you add it here as well?

NaderAlAwar · 2022-11-29T19:05:48Z

pykokkos/lib/ufuncs.py

+        for j in range(view.extent(2)):
+            out[tid][i][j] = floor(view[tid][i][j])
+
+def floor(view):


I think we should add the docstring in this PR. Regarding committing the templating engine, I would say not to do that for now. Since we're planning on getting rid of type annotations in the future, writing these kernels should be less tedious in the future.

NaderAlAwar · 2022-11-29T19:08:24Z

pykokkos/lib/ufuncs.py

+def ceil_impl_3d_float(tid: int, view: pk.View3D[pk.float], out: pk.View3D[pk.float]):
+    for i in range(view.extent(1)):
+        for j in range(view.extent(2)):
+            out[tid][i][j] = ceil(view[tid][i][j])


round should also work once we add it to the set of allowed math functions

tylerjereddy · 2022-12-12T18:22:17Z

Not quite sure I follow the concern with round here, isn't that already dealt with in my diff?

* use a tagged release of the array API conformance test suite to follow up on: kokkos#63 (comment) * this seems safe when I test locally at least--we'll see if the CI agrees * this splits off the similar change from kokkosgh-124, to simplify that PR a little

tylerjereddy · 2022-12-12T23:03:30Z

Ok, I pushed in some revisions here now--from the commit message:

* added docstrings for the new ufuncs:
  * `round()`
  * `trunc()`
  * `ceil()`
  * `floor()`

* added a regression test to probe the above
rounding functions with floating types because I was concerned that
the array API conformance suite only checked
integer view types for these functions (which don't require compiled
kernels because integers round to themselves)

gliga · 2022-12-13T09:35:45Z

pykokkos/lib/ufuncs.py

+    if len(view.shape) > 3:
+        raise NotImplementedError("only up to 3D views currently supported for round() ufunc.")
+
+    if "double" in str(view.dtype):


Rather than having three blocks for three types, can we have one if statement and then call a function to which we pass three functions, something along these lines

if "double" in str(view.dtype): _parallel_for(round_impl_1d_double, round_impl_2d_double, round_impl_3d_double...) elif "float64" in str(view.dtype): ...

How does _parallel_for know how to dispatch to the individual (1D vs. 2D vs. 3D) functions? How do you write a single _parallel_for that is shared between different ufuncs? I don't think we can easily template to the workunits without basically writing a bunch of nested logic again. And we need a separate _parallel_for for each type? I dunno, it seems like the abstraction is risky on making things simpler because the workunits don't really have metadata we can hook into for dispatching.

I suppose we could write a dictionary with the workunits as values, and dispatch through well-named keys (types, and dimensions) perhaps. We might want to agree on how that would be designed first if we adopt something like that.

That's basically early stages of writing your own type-templating system though?

gliga · 2022-12-13T09:36:49Z

pykokkos/lib/ufuncs.py

+        elif len(view.shape) == 3:
+            pk.parallel_for(view.shape[0], round_impl_3d_double, view=view, out=out)
+
+    elif "float64" in str(view.dtype):


instead of having str(view.dtype) each time, can we save this into a variable and use it when needed

gliga · 2022-12-13T09:37:26Z

pykokkos/lib/ufuncs.py

+    if len(view.shape) > 3:
+        raise NotImplementedError("only up to 3D views currently supported for trunc() ufunc.")
+
+    if "double" in str(view.dtype):


same comment as earlier: if we have _parallel_for we can save a lot of lines as it will be used here as well

gliga · 2022-12-13T09:37:37Z

pykokkos/lib/ufuncs.py

+    if len(view.shape) > 3:
+        raise NotImplementedError("only up to 3D views currently supported for ceil() ufunc.")
+
+    if "double" in str(view.dtype):


gliga · 2022-12-13T09:37:45Z

pykokkos/lib/ufuncs.py

+    if len(view.shape) > 3:
+        raise NotImplementedError("only up to 3D views currently supported for floor() ufunc.")
+
+    if "double" in str(view.dtype):


gliga

Looks good overall, but it would be good we reduce some of those if blocks as suggested.

tylerjereddy · 2022-12-13T18:04:03Z

We could reduce some code duplication, but the code is all generated from a template engine anyway, as noted above.

I suppose we could adjust the template engine, though I don't know how much easier things will be to read with more abstraction.

tylerjereddy · 2022-12-15T19:09:22Z

I tried to write up a short proposal for improving our concision for kernel dispatch at gh-135. @gliga does that sound reasonable? I was thinking we might delay that slightly as a separate effort, since it would affect the dispatching behavior of all ufuncs I suspect.

gliga · 2022-12-15T23:39:49Z

@tylerjereddy this PR is OK to go then. I will check that separately and that can be revised in a diff PR. Thanks

* this uses some templating code to write 4 new API conformant ufuncs: `round`, `trunc`, `ceil`, and `floor`. (see: https://github.com/tylerjereddy/pykokkos_ufunc_generator ) * note that because these are automatically generated they have some verbosity/lack of docstring quirks, bit of an experiment.. * switch to the latest tagged release of the array API test suite rather than some random hash * strangely, the API tests added here don't actually seem to need to compile the workunits--they may focus on the special use case for intergers I suppose; we may want to supplement with a few floating point examples perhaps; see also related discussion in data-apis/array-api-tests#154

* added docstrings for the new ufuncs: * `round()` * `trunc()` * `ceil()` * `floor()` * added a regression test to probe the above rounding functions with floating types because I was concerned that the array API conformance suite only checked integer view types for these functions (which don't require compiled kernels because integers round to themselves)

* switch to the new ufunc dispatcher system for concision

tylerjereddy · 2022-12-22T00:19:47Z

I rebased this on latest develop, and adjusted to use the new ufunc dispatching/templating system, so it is a bit more concise now. We can probably abstract even more I suppose.

* adjust `pre_compile_ufuncs` module to exclude the "private" functions in the ufuncs module

tylerjereddy added enhancement New feature or request Python Array API Standard ufuncs labels Nov 3, 2022

tylerjereddy commented Nov 3, 2022

View reviewed changes

NaderAlAwar requested changes Nov 29, 2022

View reviewed changes

tylerjereddy mentioned this pull request Dec 12, 2022

MAINT, CI: use API tests 2022.09.30 #133

Merged

tylerjereddy force-pushed the treddy_ufunc_templs branch from 86df1f3 to 8983819 Compare December 12, 2022 23:02

gliga reviewed Dec 13, 2022

View reviewed changes

tylerjereddy mentioned this pull request Dec 15, 2022

MAINT, ENH: concision of ufunc type/dims "dispatching" #135

Open

tylerjereddy added 3 commits December 21, 2022 15:49

MAINT: PR 124 revisions

b7a62d5

* switch to the new ufunc dispatcher system for concision

tylerjereddy force-pushed the treddy_ufunc_templs branch from 8983819 to b7a62d5 Compare December 22, 2022 00:18

MAINT: PR 124 revisions

c3a0929

* adjust `pre_compile_ufuncs` module to exclude the "private" functions in the ufuncs module

NaderAlAwar approved these changes Jan 3, 2023

View reviewed changes

NaderAlAwar merged commit ffd607c into kokkos:develop Jan 3, 2023

tylerjereddy deleted the treddy_ufunc_templs branch January 3, 2023 16:46

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

ENH: more ufuncs to API std #124

ENH: more ufuncs to API std #124

tylerjereddy commented Nov 3, 2022

tylerjereddy Nov 3, 2022

NaderAlAwar Nov 29, 2022

tylerjereddy Nov 3, 2022 •

edited

Loading

NaderAlAwar Nov 29, 2022

NaderAlAwar left a comment

NaderAlAwar Nov 29, 2022

NaderAlAwar Nov 29, 2022

NaderAlAwar Nov 29, 2022

NaderAlAwar Nov 29, 2022

tylerjereddy commented Dec 12, 2022

tylerjereddy commented Dec 12, 2022

gliga Dec 13, 2022

tylerjereddy Dec 13, 2022

tylerjereddy Dec 13, 2022 •

edited

Loading

gliga Dec 13, 2022

gliga Dec 13, 2022

gliga Dec 13, 2022

gliga Dec 13, 2022

gliga left a comment

tylerjereddy commented Dec 13, 2022

tylerjereddy commented Dec 15, 2022

gliga commented Dec 15, 2022

tylerjereddy commented Dec 22, 2022

ENH: more ufuncs to API std #124

ENH: more ufuncs to API std #124

Conversation

tylerjereddy commented Nov 3, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tylerjereddy Nov 3, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NaderAlAwar left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tylerjereddy commented Dec 12, 2022

tylerjereddy commented Dec 12, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tylerjereddy Dec 13, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gliga left a comment

Choose a reason for hiding this comment

tylerjereddy commented Dec 13, 2022

tylerjereddy commented Dec 15, 2022

gliga commented Dec 15, 2022

tylerjereddy commented Dec 22, 2022

tylerjereddy Nov 3, 2022 •

edited

Loading

tylerjereddy Dec 13, 2022 •

edited

Loading