Calling jex.ffi.ffi_call three times, getting only one output.

[hyhieu]

I am trying to understand the behaviors of FFI in Jax 0.4.33, and I saw something I cannot explain. I would really appreciate an explanation, or a pointer to where am I doing wrong. Thank you!

I successfully compiled a function in C++ and apparently could get it into Python via ctypes. jex.ffi.register_ffi_target even worked for me.

The C++ function that I compiled is supposed to print a line of text every time it's called.

However, when I tried to call jex.ffi.ffi_call on the registered module for 3 times, I only saw 1 output.

I would really appreciate some help explaining what is going on. Thank you again!

My system details are:

Python 3.10.12
numpy==1.26.4
jax==0.4.34.dev20240916+ae0e403c6

My code is as follows:

• This is my C++ file my_func.cc:

  #pragma once
  #include <cstdio>
  #include "xla/ffi/api/ffi.h"
  namespace ffi = xla::ffi;
  
  ffi::Error print_text_impl() {
      printf("---> Hello. You are finally in this function\n");
      fflush(stdout);
  }
  
  XLA_FFI_DEFINE_HANDLER_SYMBOL(print_text, print_text_impl);

• This is my CMakeLists.txt:

  cmake_minimum_required(VERSION 3.22)
  
  project(my_func_api LANGUAGES CXX)
  set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
  
  # Ensure we have Python
  find_package(Python COMPONENTS Interpreter Development REQUIRED)
  
  # Jax FFI
  execute_process(
      COMMAND "${Python_EXECUTABLE}" "-c" "from jax.extend import ffi; print(ffi.include_dir())"
      OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE XLA_DIR)
  
  
  add_library(my_func_api SHARED "my_func.cc")
  target_include_directories(my_func_api PUBLIC ${XLA_DIR})
  target_compile_features(my_func_api PUBLIC cxx_std_17)
  set_target_properties(
      my_func_api PROPERTIES
      OUTPUT_NAME "my_func_api"
      PREFIX ""
      SUFFIX ".so"
      LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})

• This is my Python test file my_func_test.py:

  import pytest
  
  def test_compile():
      import ctypes
      lib = ctypes.cdll.LoadLibrary("my_func_api.so")
      jex.ffi.register_ffi_target("print_text", fn=jex.ffi.pycapsule(lib.print_text), platform="cpu")
      for i in range(3):
          print(i)
          jex.ffi.ffi_call("print_text", [], x=np.float32(3.1415926), vectorized=False)

• Build and run commands:

  mkdir build
  cd build
  cmake ..
  cmake --build .
  cd ..
  pytest my_func_test.py

[jaro-sevcik]

Perhaps try to return the success code from your print_text_impl function:

ffi::Error print_text_impl() {
    printf("---> Hello. You are finally in this function\n");
    fflush(stdout);
    return ffi::Error::Success();
}

[dfm]

Good question! There are a few issues here that I'm happy to explain. Let me start by saying that if you do want some sort of printing op, you should use jax.debug.print rather than trying to write your own because it handles all the issues that I'll talk about here.

The main point to note is that ffi_call assumes that the custom calls are pure, which means that we're allowed to re-order or choose to not call the function at all. With that in mind, I'm actually surprised that you even see the printing once! When I tested the same function, I don't see any output. The reason for this is that, since your ffi_call doesn't return any outputs, there's no need for it to ever be called.

So, you can update your function to take an output, e.g.:

ffi::Error PrintImpl(ffi::Result<ffi::BufferR0<ffi::S32>>) {
  printf("---> Hello. You are finally in this function\n");
  fflush(stdout);
  return ffi::Error::Success();
}

XLA_FFI_DEFINE_HANDLER_SYMBOL(Print, PrintImpl,
                              ffi::Ffi::Bind().Ret<ffi::BufferR0<ffi::S32>>());

and then call it as follows:

def print_text():
  return jex.ffi.ffi_call("print_text", jax.ShapeDtypeStruct((), np.int32))

Then, if you run print_text a few times, you'll see the output for each call.

But, you have to remember that we're still going to assume that your function call is pure, and printing breaks that contract. That means that if you then do:

@jax.jit
def print_multiple():
  print_text()
  print_text()

print_multiple()

you won't actually see any logging, because the output of print_text is never used and we're welcome to eliminate the call under the jit.

I hope this clarifies the behavior you're seeing!

[dfm]

While we are on that topic, is it possible to force an FFI op to run in Jax? If you are familiar, in TensorFlow, we used to be able to do...

Yes! This isn't exposed via ffi_call right now, but it is used for functions like io_callback as discussed in the External callbacks doc. It should be straightforward to add support for this via ffi_call. Do you think you could open a JAX issue with a specific feature request?

Do you happen to know why it seg-faulted?

I am surprised to see a segfault here! The most common cause here is that the function pointer passed to register_ffi_target isn't actually an FFI handler (e.g. if you passed a pointer to print_text_impl instead of print_text from your example above). Can you double check that that's not what happened here? It could also potentially be caused by version mismatches between jaxlib and the FFI headers, but that should typically raise a (Python) runtime error, rather than a segfault. If you can't sort it out, perhaps you can post and updated code listing that segfaults?

[hyhieu]

jaxlib and jax mismatch is, unfortunately, a possibility in my setting. Both are custom-built. Python is not throwing runtime errors, but do you have a suggestion to check for the version mismatch?

On the other hand, here are my updated code snippets. I believe I am still triggering print_text. The changes go into the C++ file.

• New my_func.cc. Note that you can toggle WITH_RET, either in the code or via cmake:

  #pragma once
  #include <cstdio>
  #include "xla/ffi/api/ffi.h"
  namespace ffi = xla::ffi;
  
  #define WITH_RET
  
  #ifdef WITH_RET
  ffi::Error print_text_impl(ffi::Result<ffi::BufferR0<ffi::S32>> ret) {
      printf("---> Hello. You are finally in this function\n");
      fflush(stdout);
      return ffi::Error::Success();
  }
  XLA_FFI_DEFINE_HANDLER_SYMBOL(print_text, print_text_impl, ffi::Ffi::Bind().Ret<ffi::BufferR0<ffi::S32>>());
  #else
  ffi::Error print_text_impl() {
      printf("---> Hello. You are finally in this function\n");
      fflush(stdout);
      return ffi::Error::Success();
  }
  XLA_FFI_DEFINE_HANDLER_SYMBOL(print_text, print_text_impl);
  #endif

• New my_func_test.py:

    import pytest
  
    def test_compile():
        import ctypes
        from pathlib import Path
  
        lib = ctypes.cdll.LoadLibrary(str((Path(".") / "my_func_api.so").absolute()))
        jex.ffi.register_ffi_target("print_text", fn=jex.ffi.pycapsule(lib.print_text), platform="cpu")
  
        @jax.jit
        def print_text():
            return jex.ffi.ffi_call("print_text", [jax.ShapeDtypeStruct((), np.int32)])
  
        for i in range(3):
            print(i)
            print_text()

[dfm]

hmm. When I run exactly this code using the commands you listed everything works fine for me! I'm using JAX from the GitHub main branch and either jaxlib v0.4.33 from PyPI or jaxlib nightly v0.4.34.dev20240927. You might double check that CMake is finding the right Python, and therefore the installed FFI headers directory? Then perhaps you could share the CMake build logs and more info about your compiler toolchain, and platform so that we can try to debug.

[hyhieu]

Interesting! The fact that you could run it is a sign of life. Thank you!

It's totally possible that I have other stuffs corrupted in my dev environment, so I will try to restart my dev environment and try again. Thank you for all the help so far ❤️

Regarding versions and settings: you have jaxlib==0.4.33 or jaxlib==0.4.34.dev[...]. Do both of them work for you?

For me, both are 0.4.34:

$ pip freeze | grep "jax"
jax @ file:///root/jax-src/jax-0.4.34.dev20240916%2Bae0e403c6-py3-none-any.whl#sha256=[...]
jaxlib @ file:///root/jax-src/dist/jaxlib-0.4.34.dev20240919-cp310-cp310-manylinux2014_x86_64.whl#sha256=[...]

My CMakeLists.txt is the same as you saw in the first post.

[hyhieu]

Indeed, this turned out to be an issue with my dev environment. @dfm your version works on a clean machine.

More importantly, this conversation elicited so many points that I didn't know about XLA. It's important to establish a correct mindset about the software stack we are using. Thank you!