Skip to content

Jmm/python numpy wrapper#636

Merged
Yurlungur merged 5 commits intomainfrom
jmm/python-numpy-wrapper
Apr 20, 2026
Merged

Jmm/python numpy wrapper#636
Yurlungur merged 5 commits intomainfrom
jmm/python-numpy-wrapper

Conversation

@Yurlungur
Copy link
Copy Markdown
Collaborator

@Yurlungur Yurlungur commented Apr 20, 2026

PR Summary

This MR supersedes #634 . The issue is that, despite Python calling only host code (thanks to #630 ), the loops are still marked with PORTABLE_LAMBDAs which means the compiler detects that host only code (numpy arrays operator[]) is called from a function marked __host__ __device__ (the portable lambda). #634 fixed this by adding a switch on the lambda making it either PORTABLE_LAMBDA or [=] depending on execution space.

This MR instead wraps numpy arrays in a very thin wrapper. @jdolence and I actually already wrote this wrapper for a prior MR, so almost no changes needed to be made. The wrapper lies to the compiler by marking the operator[] as PORTABLE_INLINE_FUNCTION despite the memory always being host memory. Nevertheless, this is safe, because we always call the python explicitly on host.

This MR eliminates warnings on GPU architectures related to calling host functions from device. It also is required for the python code to compile on AMD GPU architectures, if a non-host backend is active.

PR Checklist

  • Adds a test for any bugs fixed. Adds tests for new features.
  • Format your changes by using the make format command after configuring with cmake.
  • Document any new features, update documentation for changes made.
  • Make sure the copyright notice on any files you modified is up to date.
  • After creating a pull request, note it in the CHANGELOG.md file.
  • LANL employees: make sure tests pass both on the github CI and on the Darwin CI
  • If ML was used, make sure to add a disclaimer at the top of a file indicating ML was used to assist in generating the file.
  • If Agentic AI was used, have the AI generate a "proposed changes" markdown file and store it in the plan_histories folder, with a filename the same as the MR number.

If preparing for a new release, in addition please check the following:

  • Update the version in cmake.
  • Move the changes in the CHANGELOG.md file under a new header for the new release, and reset the categories.
  • Maintainers: ensure spackages are up to date:
    • LANL-internal team, update XCAP spackages
    • Current maintainer of upstream spackages, submit MR to spack

@Yurlungur Yurlungur self-assigned this Apr 20, 2026
@Yurlungur Yurlungur added the bug Something isn't working label Apr 20, 2026
@Yurlungur Yurlungur mentioned this pull request Apr 20, 2026
13 tasks
Comment thread python/module.hpp
// Helper function to convert lambda numpy array to double* buffer
// With std::optional we would add support for a default value of lambda=None
template<typename T, PORTABLE_FUNCTION Real(T::*Func)(const Real, const Real, Real*&&) const>
template<typename T, Real(T::*Func)(const Real, const Real, Real*&&) const>
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These decorators just aren't allowed in the template arguments. They also don't seem to be necessary.

Comment thread python/module.hpp
Comment on lines +51 to +54
LambdaHelper(py::array_t<Real> & lambdas)
: data_(lambdas.mutable_data(0,0))
, stride_(lambdas.mutable_data(1,0) - lambdas.mutable_data(0,0))
{}
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This stride trick should work for any non-ragged lambda array.

Copy link
Copy Markdown
Collaborator

@adamdempsey90 adamdempsey90 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Glad this worked

@Yurlungur Yurlungur merged commit 4093f31 into main Apr 20, 2026
9 checks passed
@Yurlungur Yurlungur deleted the jmm/python-numpy-wrapper branch April 20, 2026 23:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants