mirror of
https://github.com/ROCm/jax.git
synced 2025-04-16 03:46:06 +00:00
0b4800a193
Building on #21925, this tutorial demonstrates the use of the FFI using `ffi_call` with a simple example. I don't think this should cover all of the most advanced use cases, but it should be sufficient for the most common examples. I think it would be useful to eventually replace the existing CUDA tutorial, but I'm not sure that it'll get there in the first draft. As an added benefit, this also runs a simple test (akin to `docs/cuda_custom_call`) which actually executes using a tool chain that open source users would use in practice.
15 lines
608 B
CMake
15 lines
608 B
CMake
cmake_minimum_required(VERSION 3.18...3.27)
|
|
project(rms_norm LANGUAGES CXX)
|
|
|
|
find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
|
|
execute_process(
|
|
COMMAND "${Python_EXECUTABLE}"
|
|
"-c" "from jax.extend import ffi; print(ffi.include_dir())"
|
|
OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE XLA_DIR)
|
|
message(STATUS "XLA include directory: ${XLA_DIR}")
|
|
|
|
add_library(rms_norm SHARED "rms_norm.cc")
|
|
target_include_directories(rms_norm PUBLIC ${XLA_DIR})
|
|
target_compile_features(rms_norm PUBLIC cxx_std_17)
|
|
install(TARGETS rms_norm LIBRARY DESTINATION ${CMAKE_CURRENT_LIST_DIR})
|