Function load

python/tvm_ffi/cpp/extension.py:1323–1476

Compile, build and load a C/C++/CUDA module from source files. This function compiles the given source files into a shared library and loads it as a tvm ffi module. The compiler is selected automatically based on file extension. Note that this function does not automatically export fun

(  # noqa: PLR0913
    name: str,
    *,
    sources: Sequence[str] | str | None = None,
    cpp_files: Sequence[str] | str | None = None,
    cuda_files: Sequence[str] | str | None = None,
    extra_cflags: Sequence[str] | None = None,
    extra_cuda_cflags: Sequence[str] | None = None,
    extra_ldflags: Sequence[str] | None = None,
    extra_include_paths: Sequence[str] | None = None,
    build_directory: str | None = None,
    keep_module_alive: bool = True,
    backend: str | None = None,
)

def load(  # noqa: PLR0913
    name: str,
    *,
    sources: Sequence[str] | str | None = None,
    cpp_files: Sequence[str] | str | None = None,
    cuda_files: Sequence[str] | str | None = None,
    extra_cflags: Sequence[str] | None = None,
    extra_cuda_cflags: Sequence[str] | None = None,
    extra_ldflags: Sequence[str] | None = None,
    extra_include_paths: Sequence[str] | None = None,
    build_directory: str | None = None,
    keep_module_alive: bool = True,
    backend: str | None = None,
) -> Module:
    """Compile, build and load a C/C++/CUDA module from source files.

    This function compiles the given source files into a shared library and loads it as a tvm ffi
    module. The compiler is selected automatically based on file extension.

    Note that this function does not automatically export functions to the tvm ffi module. You need to
    manually use the TVM FFI export macros (e.g., :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC`) in your source files to export
    functions. This gives you more control over which functions are exported and how they are exported.

    Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags``
    parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional
    flags for your specific use case.

    The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can
    include any header from tvm ffi in your source files. You can also provide additional include paths via the
    ``extra_include_paths`` parameter and include custom headers in your source code.

    The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter
    is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used.
    The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified,
    the default cache directory is ``~/.cache/tvm-ffi``.

    Parameters
    ----------
    name
        The name of the tvm ffi module.
    sources
        Source files to compile. The compiler is auto-detected from the file extension:
        ``.c`` → C, ``.cc``/``.cpp``/``.cxx`` → C++, ``.cu`` → CUDA/HIP,
        ``.o``/``.obj`` → linker passthrough. It can be a list of file paths or a single file path.
    cpp_files
        Alias for ``sources``, kept for backward compatibility.
    cuda_files
        Alias for ``sources``, kept for backward compatibility.
    extra_cflags
        The extra compiler flags for C++ compilation.
        The default flags are:

        - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2']
        - On Windows: ['/std:c++17', '/MD', '/O2']

    extra_cuda_cflags
        The extra compiler flags for CUDA compilation.
        The default flags are: