tqchen / ffi-navigator

FFI Navigator

Most modern IDEs support find function definition within the same language(e.g. python or c++). However, it is very hard to do that for cross-language FFI calls. While solving this general problem can be very technically challenging, we can get around it by building a project-specific analyzer that matches the FFI registration code patterns and recovers the necessary information.

This project is an example of that. Currently, it supports the PackedFunc FFI in the Apache TVM project. It is implemented as a language server that provides getDefinition function for FFI calls and returns the location of the corresponding C++ API in the TVM project. It complements the IDE tools that support navigation within the same language. We also have preliminary support for MXNet, DGL, and PyTorch, so we can do goto-definition from Python to C++ in these projects too.

Installation

Install python package

pip install --user ffi-navigator

VSCode

See vscode-extension

Emacs

Install lsp-mode

Add the following configuration

(lsp-register-client
 (make-lsp-client
  :new-connection (lsp-stdio-connection '("python3" "-m" "ffi_navigator.langserver"))
  :major-modes '(python-mode c++-mode)
  :server-id 'ffi-navigator
  :add-on? t))

• Use commands like M-x lsp-find-definition and M-x lsp-find-references

If you use eglot instead, check out this PR. eglot does not support multiple servers per language at the moment, the PR above contains a workaround.

Other editors/IDEs

It should be straightforward to add support for other editors or IDEs that have a LSP client implementation. Please refer to this site for the availability of clients.

Since ffi-navigator is intended to be used with other general purpose servers such as pyls, your LSP client needs to be able to talk to multiple servers per language. If your client does not have such feature, check out this PR and the discussion there for a workaround. This branch has fallback to pyls and it is always kept up to date with the master branch.

Features

TVM FFI

• Find definition/references of FFI objects(e.g. PackedFunc in TVM) on the python and c++ side.
  • Jump from a python PackedFunc into TVM_REGISTER_GLOBAL, @register_func
• Find definition/references of FFI objects on the python and c++ side.
  • move cursor to the object class name on the python side.
  • move cursor to the _type_key string on the c++ side.

PyTorch

• Jump to C10 registered ops. In python they corresponds to functions under torch.ops namespace.
  • Example: torch.ops.quantized.conv2d (py) -> c10::RegisterOperators().op("quantized::conv2d", ...) (cpp)
• Jump to cpp functions wrapped by pybind.
  • Example: torch._C._jit_script_class_compile (py) -> m.def( "_jit_script_class_compile", ...) (cpp)

Development

For developing the python package locally, we can just make sure ffi_navigator is in your python path in bashrc.

export PYTHONPATH=${PYTHONPATH}:/path/to/ffi-navigator/python

Project Structure

• python/ffi_navigator The analysis code and language server
• python/ffi_navigator/dialect Per project dialects
• vscode-extension language server extension for vscode

Adding Support for New FFI Patterns

Add your FFI convention to dialect namespace.

Demo

VSCode

See vscode-extension

Emacs

Goto definition from Python to C++

Goto definition from C++ to Python

Find reference across Python and C++

Goto definition in PyTorch