Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cythonized Timestep [Core] #3683

Merged
merged 73 commits into from
Jun 13, 2022
Merged

Conversation

hmacdope
Copy link
Member

@hmacdope hmacdope commented May 31, 2022

Cythonized Timestep

First large changes of CZIEOSS4 performance track. Primarily cythonization of the Timestep class and separation from the iterators. Would be good to get some input at this point to see what people think.

Fixes #3709

Status:

  • Still a few bugs to track down wrt MemoryReader and NetCDF readers
  • Finalise docs
  • Sort out API wrt dtypes following some discussion

Design Ideas:

  • We initially trialled a fully C++ data-structure and found that moving back and forward from two possible types was difficult despite the use of templates because the Cython wrapper class had to know which template to instantiate at __cinit__ time, making it very ugly for limited performance gain.
  • Moved towards a more idiomatic Cython representation
  • A big takeaway (suprise suprise) was that using the NumPy C API was the best way to get performance.

Aims of this PR:

  • Cythonize the Timestep class with an almost exact API match
  • Have a test passing and well documented implementation moving forward
  • Improve set and get performance of the reader (initial test indicate 10x 2x faster to set positions without a copy with a copy if no casting required)

Changes made in this Pull Request:

  • Timestep is now an Extension type (cdef class) with a pxd file to define relevant imports
  • dtype kwarg added
  • _frame is no longer optional as must be initialised (to -1) _frame is no longer an attribute and has been moved to data dictionary in TRZReader
  • typenum, *_dims attributes added to store enumerated NumPy dtype and associated array dimensionality
  • position, velocity and force setters split into 3 branches depending on input object and whether a cast is needed.
  • position, forces and velocities are always C contiguous.
  • __getstate__ and __setstate__ implemented manually as cdef classes do not have the __dict__ attribute.
  • __getnewargs_ex__ added to help __cinit__ construction upon unpickling.
  • libmda added as a folder to import all cython from

Discussion Points:

  • Datatypes are a big concern here, as of current we do not allow mixed precision in the Timestep. How do we want to tackle this moving forward?
  • We currently force convert position dependent data to np.float32 but this can be changed.
  • Boxes are also np.float32 but elsewhere in the code are promoted and demoted regularly.
  • Do we want to allow a boxes to have a different datatype with an additional kwarg (dimensions_dtype=np.float...)?
  • Mixed precision ops are slow / difficult for the compiler to vectorize.

Looking forwards:

  • What is the best way to hook C level functions to the data (especially templated functions)? Richard and I are trialling an idiom with a cython.floating memoryview and matching templated C function.
  • Do we want to provide hooks directly into the array such as with PyArray_Getptr
  • What is good to expose at PXD level?

PR Checklist

  • Tests?
  • Docs?
  • CHANGELOG updated?

@hmacdope
Copy link
Member Author

hmacdope commented Jun 5, 2022

Oh, and WRT dtype, we can fix this to float32 for now and sort that out later in a different PR. I can't see anything here that prevents us switching dtype, but it'll require plumbing into other parts of the code and a boatload of extra tests, so let's keep this PR simple.

With respect to this should I remove the dtype= kwarg for this PR?

@hmacdope
Copy link
Member Author

hmacdope commented Jun 6, 2022

I am removing the dtype kwarg from this PR and we can discuss further on an issue. I will raise and link here.

@hmacdope hmacdope requested review from a team, tylerjereddy and richardjgowers June 6, 2022 23:42
@hmacdope
Copy link
Member Author

hmacdope commented Jun 6, 2022

@tylerjereddy would I please be able to get a review from you specifically on my use of the NumPy C API?

There are quite a few calls to the C API, but I am not aware if there are any NumPy conventions I might be ignoring and would like to be as conformant as possible.

Particularly, initialising the _position, _velocities and _forces arrays to be empty with size 0 using PyArray_EMPTY() in __cinit__ was something I wanted to query.

Is there a better way to have a C level array hook without allocating to their full size (natoms, 3)?

@hmacdope
Copy link
Member Author

hmacdope commented Jun 7, 2022

I am removing the dtype kwarg from this PR and we can discuss further on an issue. I will raise and link here.

Discussion of dtypes can take place on #3707

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
Development

Successfully merging this pull request may close these issues.

Cythonize Timestep for greater performance
4 participants