feat: raw pointer capsule instantiation method

[lmmx]

• via [Idea] PyCapsule instantiators for raw pointer capsules #5688

Implements PyCapsule::new_with_pointer and PyCapsule::new_with_pointer_and_destructor as described in the issue above

Purpose

Allows users to write much simpler code, without getting under the hood of pyo3-ffi to make a pyo3::PyCapsule

Edit: original version using *mut c_void was deprecated

Click to show original version

#[pyfunction(name = "rms_norm")]
fn rms_norm_jax(py: Python<'_>) -> PyResult<Bound<'_, PyCapsule>> {
    unsafe {
        PyCapsule::new_with_pointer(py, ffi::RmsNorm as *mut c_void, None)
    }
}

Update We now use the std::ptr::NonNull<T> type (in which the T is a *mut T) so we can write it as NonNull<c_void> instead of *mut c_void and then checking ptr.is_null().

Usage would therefore be

#[pyfunction(name = "rms_norm")]
fn rms_norm_jax(py: Python<'_>) -> PyResult<Bound<'_, PyCapsule>> {
    let fn_ptr = NonNull::new(ffi::RmsNorm as *mut c_void)
        .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err(
            "Function pointer must not be null",
        ))?;

    unsafe {
        PyCapsule::new_with_pointer(py, fn_ptr, None)
    }
}

instead of

#[pyfunction(name = "rms_norm")]
fn rms_norm_jax(py: Python<'_>) -> PyResult<Bound<'_, PyCapsule>> {
    let fn_ptr: *mut c_void = ffi::RmsNorm as *mut c_void;
    let name = std::ptr::null();

    unsafe {
        let capsule = pyo3::ffi::PyCapsule_New(fn_ptr, name, None);
        if capsule.is_null() {
            return Err(pyo3::exceptions::PyRuntimeError::new_err(
                "Failed to create PyCapsule",
            ));
        }
        let any: Bound<'_, PyAny> = Bound::from_owned_ptr(py, capsule);
        Ok(any.cast_into_unchecked::<PyCapsule>())
    }
}

• Specifically this was the refactor enabled by switching to use this PR branch in jeertmans/extending-jax@8fbb5bb

[davidhewitt]

Thanks for working on this, I think the proposed functionality makes sense. Some initial comments on the design below.

[jeertmans]

Hi @lmmx (happy new year!), just checking when you would have some time to review @davidhewitt's comments?

[lmmx]

Hi @lmmx (happy new year!), just checking when you would have some time to review @davidhewitt's comments?

HNY! 🪩 Addressed these threads in the large, a couple more to go I haven't marked resolved (patch the leak, take a decision on whether the name should be null) plus give the docstrings another pass.

Right now I'm ensuring the CI checks are OK - I can see one is no longer a valid test because of the NonNull, test_pycapsule_new_with_pointer_null_rejected which makes sense! I've replaced that with a trivial test that the NonNull type rejects null pointers which doesn't use the Python machinery at all (c556801)

[lmmx]

I think that's addressed the comments so far, CI is passing, and I re-did the docstrings.

Please feel free to update this PR if needed to merge it or give further review whenever suits.

[jeertmans]

Thank you very much @lmmx! I'm really impressed by your work, thank you for spending your time on this!