[STF] Add python bindings #5315

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

caugonnet wants to merge 279 commits into NVIDIA:main from caugonnet:stf_c_api

c/experimental/stf/include/cccl/c/experimental/stf/stf.h

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -456,36 +456,87 @@ cudaStream_t stf_fence(stf_ctx_handle ctx);
  
    //!

    //! \brief Create logical data from existing memory buffer

    //!

    //! Creates logical data handle from an existing host memory buffer.

    //! STF takes ownership of data management during task execution.

    //! Creates logical data handle from existing memory buffer, assuming host data place.

    //! This is a convenience wrapper around stf_logical_data_with_place() with host placement.

    //!

    //! \param ctx Context handle

    //! \param[out] ld Pointer to receive logical data handle

    //! \param addr Pointer to existing data buffer

    //! \param addr Pointer to existing data buffer (assumed to be host memory)

    //! \param sz Size of data in bytes

    //!

    //! \pre ctx must be valid context handle

    //! \pre ld must not be NULL

    //! \pre addr must not be NULL

    //! \pre addr must not be NULL and point to host-accessible memory

    //! \pre sz must be greater than 0

    //! \post *ld contains valid logical data handle

    //!

    //! \note Original data pointer should not be accessed during task execution

    //! \note Data will be written back when logical data is destroyed or context finalized

    //! \note This function assumes host memory. For device/managed memory, use stf_logical_data_with_place()

    //! \note Equivalent to: stf_logical_data_with_place(ctx, ld, addr, sz, make_host_data_place())

    //!

    //! \par Example:

    //! \code

    //! float data[1024];

    //! stf_logical_data_handle ld;

    //! stf_logical_data(ctx, &ld, data, sizeof(data));

    //! stf_logical_data(ctx, &ld, data, sizeof(data));  // Assumes host memory

    //! // ... use in tasks ...

    //! stf_logical_data_destroy(ld);

    //! \endcode

    //!

    //! \see stf_logical_data_empty(), stf_logical_data_destroy()

    //! \see stf_logical_data_with_place(), stf_logical_data_empty(), stf_logical_data_destroy()

    void stf_logical_data(stf_ctx_handle ctx, stf_logical_data_handle* ld, void* addr, size_t sz);

    //!

    //! \brief Create logical data handle from address with data place specification

    //!

    //! Creates logical data handle from existing memory buffer, explicitly specifying where

    //! the memory is located (host, device, managed, etc.). This is the primary and recommended

    //! logical data creation function as it provides STF with essential memory location information

    //! for optimal data movement and placement strategies.

    //!

    //! \param ctx Context handle

    //! \param[out] ld Pointer to receive logical data handle

    //! \param addr Pointer to existing memory buffer

    //! \param sz Size of buffer in bytes

    //! \param dplace Data place specifying memory location

    //!

    //! \pre ctx must be valid context handle

    //! \pre ld must be valid pointer to logical data handle pointer

    //! \pre addr must point to valid memory of at least sz bytes

    //! \pre sz must be greater than 0

    //! \pre dplace must be valid data place (not invalid)

    //!

    //! \post *ld contains valid logical data handle on success

    //! \post Caller owns returned handle (must call stf_logical_data_destroy())

    //!

    //! \par Examples:

    //! \code

    //! // GPU device memory (recommended for CUDA arrays)

    //! float* device_ptr;

    //! cudaMalloc(&device_ptr, 1000 * sizeof(float));

    //! stf_data_place dplace = make_device_data_place(0);

    //! stf_logical_data_handle ld;

    //! stf_logical_data_with_place(ctx, &ld, device_ptr, 1000 * sizeof(float), dplace);

    //!

    //! // Host memory

    //! float* host_data = new float[1000];

    //! stf_data_place host_place = make_host_data_place();

    //! stf_logical_data_handle ld_host;

    //! stf_logical_data_with_place(ctx, &ld_host, host_data, 1000 * sizeof(float), host_place);

    //!

    //! // Managed memory

    //! float* managed_ptr;

    //! cudaMallocManaged(&managed_ptr, 1000 * sizeof(float));

    //! stf_data_place managed_place = make_managed_data_place();

    //! stf_logical_data_handle ld_managed;

    //! stf_logical_data_with_place(ctx, &ld_managed, managed_ptr, 1000 * sizeof(float), managed_place);

    //! \endcode

    //!

    //! \see make_device_data_place(), make_host_data_place(), make_managed_data_place()

    void stf_logical_data_with_place(

      stf_ctx_handle ctx, stf_logical_data_handle* ld, void* addr, size_t sz, stf_data_place dplace);

    //!

    //! \brief Set symbolic name for logical data

    //!

c/experimental/stf/src/stf.cu

-Original file line number
+Diff line change
@@ Expand Up / @@ -44,12 +44,44 @@ cudaStream_t stf_fence(stf_ctx_handle ctx) @@
     }
     void stf_logical_data(stf_ctx_handle ctx, stf_logical_data_handle* ld, void* addr, size_t sz)
+    {
+      // Convenience wrapper: assume host memory
+      stf_logical_data_with_place(ctx, ld, addr, sz, make_host_data_place());
+    }
+    void stf_logical_data_with_place(
+      stf_ctx_handle ctx, stf_logical_data_handle* ld, void* addr, size_t sz, stf_data_place dplace)
     {
       assert(ctx);
       assert(ld);
       auto* context_ptr = static_cast<context*>(ctx);
-      auto ld_typed     = context_ptr->logical_data(make_slice((char*) addr, sz));
+      // Convert C data_place to C++ data_place
+      cuda::experimental::stf::data_place cpp_dplace;
+      switch (dplace.kind)
+      {
+        case STF_DATA_PLACE_HOST:
+          cpp_dplace = cuda::experimental::stf::data_place::host();
+          break;
+        case STF_DATA_PLACE_DEVICE:
+          cpp_dplace = cuda::experimental::stf::data_place::device(dplace.u.device.dev_id);
+          break;
+        case STF_DATA_PLACE_MANAGED:
+          cpp_dplace = cuda::experimental::stf::data_place::managed();
+          break;
+        case STF_DATA_PLACE_AFFINE:
+          cpp_dplace = cuda::experimental::stf::data_place::affine();
+          break;
+        default:
+          // Invalid data place - this should not happen with valid input
+          assert(false && "Invalid data_place kind");
+          cpp_dplace = cuda::experimental::stf::data_place::host(); // fallback
+          break;
+      }
+      // Create logical data with the specified data place
+      auto ld_typed = context_ptr->logical_data(make_slice((char*) addr, sz), cpp_dplace);
       // Store the logical_data_untyped directly as opaque pointer
       *ld = new logical_data_untyped{ld_typed};
@@ Expand Down @@

cudax/examples/stf/void_data_interface.cu

-Original file line number
+Diff line change
@@ Expand Up / @@ -49,5 +49,9 @@ int main() @@
         return cuda_kernel_desc{dummy_kernel, 16, 128, 0};
       };
+      EXPECT(token.is_void_interface());
+      EXPECT(token2.is_void_interface());
+      EXPECT(token3.is_void_interface());
       ctx.finalize();
     }

cudax/include/cuda/experimental/__stf/graph/graph_task.cuh

-Original file line number
+Diff line change
@@ Expand Up / @@ -104,7 +104,9 @@ public: @@
         {
           // Select a stream from the pool
           capture_stream = get_exec_place().getStream(ctx.async_resources(), true).stream;
-          cuda_safe_call(cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal));
+          // Use relaxed capture mode to allow capturing workloads that lazily initialize
+          // resources (e.g., set up memory pools)
+          cuda_safe_call(cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeRelaxed));
         }
         auto& dot = *ctx.get_dot();
@@ Expand Down Expand Up / @@ -365,7 +367,9 @@ public: @@
           capture_stream = get_exec_place().getStream(ctx.async_resources(), true).stream;
           cudaGraph_t childGraph = nullptr;
-          cuda_safe_call(cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal));
+          // Use relaxed capture mode to allow capturing workloads that lazily initialize
+          // resources (e.g., set up memory pools)
+          cuda_safe_call(cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeRelaxed));
           // Launch the user provided function
           f(capture_stream);
@@ Expand Down Expand Up / @@ -625,7 +629,9 @@ public: @@
           cudaStream_t capture_stream = get_exec_place().getStream(ctx.async_resources(), true).stream;
           cudaGraph_t childGraph = nullptr;
-          cuda_safe_call(cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal));
+          // Use relaxed capture mode to allow capturing workloads that lazily initialize
+          // resources (e.g., set up memory pools)
+          cuda_safe_call(cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeRelaxed));
           // Launch the user provided function
           if constexpr (fun_invocable_stream_deps)
@@ Expand Down @@

cudax/test/stf/cpp/task_get_stream.cu

-Original file line number
+Diff line change
@@ Expand Up / @@ -24,7 +24,8 @@ void test_stream() @@
       context ctx;
       auto token = ctx.token();
-      auto t     = ctx.task(token.write());
+      EXPECT(token.is_void_interface());
+      auto t = ctx.task(token.write());
       t.start();
       cudaStream_t s = t.get_stream();
       EXPECT(s != nullptr);
@@ Expand Down @@

python/cuda_cccl/CMakeLists.txt

-Original file line number
+Diff line change
@@ Expand Up / @@ -17,9 +17,27 @@ message( @@
     # Build cccl.c.parallel and add CCCL's install rules
     set(_cccl_root ../..)
+    # Build and install C++ library first
     set(CCCL_TOPLEVEL_PROJECT ON) # Enable the developer builds
-    set(CCCL_ENABLE_C_PARALLEL ON) # Build the cccl.c.parallel library
+    set(CCCL_ENABLE_C_PARALLEL ON)
+    set(CCCL_ENABLE_C_EXPERIMENTAL_STF ON) # Enable C experimental STF library (triggers c/ directory)
+    set(CCCL_ENABLE_UNSTABLE ON) # Enable unstable features
+    # Disable all testing, examples, and benchmarks - we only want the libraries
+    set(CCCL_ENABLE_TESTING OFF)
+    set(CCCL_ENABLE_EXAMPLES OFF)
+    set(CCCL_ENABLE_BENCHMARKS OFF)
+    set(CCCL_C_PARALLEL_ENABLE_TESTING OFF)
+    set(CCCL_C_EXPERIMENTAL_STF_ENABLE_TESTING OFF)
+    # Note: CCCL_ENABLE_CUDAX must be ON because STF depends on it (via CCCL_ENABLE_UNSTABLE)
+    # But disable cudax tests, examples, and header testing
+    set(cudax_ENABLE_TESTING OFF)
+    set(cudax_ENABLE_EXAMPLES OFF)
+    set(cudax_ENABLE_HEADER_TESTING OFF)
     set(CCCL_C_PARALLEL_LIBRARY_OUTPUT_DIRECTORY ${SKBUILD_PROJECT_NAME})
+    set(CCCL_C_EXPERIMENTAL_STF_LIBRARY_OUTPUT_DIRECTORY ${SKBUILD_PROJECT_NAME})
     # Just install the rest:
     set(libcudacxx_ENABLE_INSTALL_RULES ON)
     set(CUB_ENABLE_INSTALL_RULES ON)
@@ Expand All / @@ -34,10 +52,21 @@ add_subdirectory(${_cccl_root} _parent_cccl) @@
     set(CMAKE_INSTALL_LIBDIR "${old_libdir}") # pop
     set(CMAKE_INSTALL_INCLUDEDIR "${old_includedir}") # pop
+    # Create CCCL::cudax alias for STF (normally created by cccl-config.cmake)
+    if (TARGET cudax::cudax AND NOT TARGET CCCL::cudax)
+      add_library(CCCL::cudax ALIAS cudax::cudax)
+    endif()
     # ensure the destination directory exists
+    file(MAKE_DIRECTORY "cuda/stf/${CUDA_VERSION_DIR}/cccl")
     file(MAKE_DIRECTORY "cuda/compute/${CUDA_VERSION_DIR}/cccl")
     # Install version-specific binaries
+    install(
+      TARGETS cccl.c.experimental.stf
+      DESTINATION cuda/stf/${CUDA_VERSION_DIR}/cccl
+    )
     install(
       TARGETS cccl.c.parallel
       DESTINATION cuda/compute/${CUDA_VERSION_DIR}/cccl
@@ Expand Down Expand Up / @@ -110,12 +139,40 @@ add_custom_target( @@
       DEPENDS "${_generated_extension_src}"
     )
+    message(STATUS "STF Using Cython ${CYTHON_VERSION}")
+    set(
+      stf_pyx_source_file
+      "${cuda_cccl_SOURCE_DIR}/cuda/stf/_stf_bindings_impl.pyx"
+    )
+    set(_stf_generated_extension_src "${cuda_cccl_BINARY_DIR}/_stf_bindings_impl.c")
+    set(_stf_depfile "${cuda_cccl_BINARY_DIR}/_stf_bindings_impl.c.dep")
+    add_custom_command(
+      OUTPUT "${_stf_generated_extension_src}"
+      COMMAND "${Python3_EXECUTABLE}" -m cython
+      ARGS
+        ${CYTHON_FLAGS_LIST} "${stf_pyx_source_file}" --output-file
+        ${_stf_generated_extension_src}
+      DEPENDS "${stf_pyx_source_file}"
+      DEPFILE "${_stf_depfile}"
+      COMMENT "Cythonizing ${pyx_source_file} for CUDA ${CUDA_VERSION_MAJOR}"
+    )
+    set_source_files_properties(
+      "${_stf_generated_extension_src}"
+      PROPERTIES GENERATED TRUE
+    )
+    add_custom_target(
+      cythonize_stf_bindings_impl
+      ALL
+      DEPENDS "${_stf_generated_extension_src}"
+    )
     python3_add_library(
       _bindings_impl
       MODULE
       WITH_SOABI
       "${_generated_extension_src}"
     )
     add_dependencies(_bindings_impl cythonize_bindings_impl)
     target_link_libraries(
       _bindings_impl
@@ Expand All / @@ -125,4 +182,21 @@ target_link_libraries( @@
     )
     set_target_properties(_bindings_impl PROPERTIES INSTALL_RPATH "$ORIGIN/cccl")
+    python3_add_library(
+      _stf_bindings_impl
+      MODULE
+      WITH_SOABI
+      "${_stf_generated_extension_src}"
+    )
+    add_dependencies(_stf_bindings_impl cythonize_stf_bindings_impl)
+    target_link_libraries(
+      _stf_bindings_impl
+      PRIVATE cccl.c.experimental.stf CUDA::cuda_driver
+    )
+    set_target_properties(
+      _stf_bindings_impl
+      PROPERTIES INSTALL_RPATH "$ORIGIN/cccl"
+    )
     install(TARGETS _bindings_impl DESTINATION cuda/compute/${CUDA_VERSION_DIR})
+    install(TARGETS _stf_bindings_impl DESTINATION cuda/stf/${CUDA_VERSION_DIR})

python/cuda_cccl/cuda/stf/__init__.py

-Original file line number
+Diff line change
@@ -0,0 +1,27 @@
+    from ._stf_bindings import (
+        context,
+        data_place,
+        dep,
+        exec_place,
+    )
+    from .decorator import jit  # Python-side kernel launcher
+    __all__ = [
+        "context",
+        "dep",
+        "exec_place",
+        "data_place",
+        "jit",
+    ]
+    def has_torch() -> bool:
+        import importlib.util
+        return importlib.util.find_spec("torch") is not None
+    def has_numba() -> bool:
+        import importlib.util
+        return importlib.util.find_spec("numba") is not None

python/cuda_cccl/cuda/stf/_adapters/numba_bridge.py

-Original file line number
+Diff line change
@@ -0,0 +1,4 @@
+    def cai_to_numba(cai: dict):
+        from numba import cuda
+        return cuda.from_cuda_array_interface(cai, owner=None, sync=False)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[STF] Add python bindings #5315

Uh oh!

Diff view

Diff view

There are no files selected for viewing

caugonnet Nov 24, 2025

Uh oh!

Uh oh!

Uh oh!

[STF] Add python bindings #5315

Are you sure you want to change the base?

Uh oh!

[STF] Add python bindings #5315

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

caugonnet Nov 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!