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!

[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!

[dfm]

Excellent news! I'm glad you got it working.

Thanks to this conversation, I have also added support for non-pure calls via ffi_call: #23982

After that change, you can now pass has_side_effect=True to ffi_call to label an FFI handler as effectful. I should note that the ordering of calls even with side effects aren't necessarily guaranteed, so something like jax.debug.print is still better for the toy problem discussed here, but hopefully this helps for more realistic problems.

[RedHat2003]

---

You may have to use io_callback

C++ Code:

#include <cstdio>
#include "xla/ffi/api/ffi.h"
#include "xla/ffi/api/c_api.h"
namespace ffi = xla::ffi;

ffi::Error SayHi(ffi::Result<ffi::BufferR0<ffi::S32>> y) {
    printf("u are finally there!\n");
    return ffi::Error::Success();
}

XLA_FFI_DEFINE_HANDLER_SYMBOL(
    Print, SayHi,
    ffi::Ffi::Bind()
        .Ret<ffi::BufferR0<ffi::S32>>()  // Ensure closing parenthesis
);

---

Python Code:

import ctypes
from types import ModuleType
from pathlib import Path
import jax.extend as jex
from jax import numpy as jnp, vmap, jit
from jax import experimental as exper

def create_shape_dtype(dtype, shape=()):
    module = ModuleType(f"ShapeDtype_{dtype.__name__}_{shape}")
    module.dtype = dtype
    module.shape = shape
    return module

cDtype = create_shape_dtype(ctypes.c_int32, ())

ffi_build_dir = Path(__file__).parent.parent / "ffi" / "_build"
lib_path = next(ffi_build_dir.glob("sayhi.so"), None)
sayhi_lib = ctypes.cdll.LoadLibrary(lib_path)

jex.ffi.register_ffi_target(
    "sayhi",
    jex.ffi.pycapsule(sayhi_lib.Print),
    platform="cpu"
)

call = jex.ffi.ffi_call(
    "sayhi",
    cDtype,
    vmap_method="sequential"
)

def io_call(x):
    call()
    call()
    call()
    return x

@jit  # <----------- JIT the test function
def test(v):
    x = exper.io_callback(
        io_call,
        cDtype,
        v,
        ordered=False
    )

if __name__ == "__main__":
    test(jnp.int32(1))

---

Output:

u are finally there!
u are finally there!
u are finally there!

---

[RedHat2003]

Okay, so…fun fact. I called this function SayHi, but it never actually says "Hi." Instead, it just gives you an existential reminder that "u are finally there!" 🤦‍♂️

Naming things is hard, but at least it prints something! 😂