Create nvfuser_next extension #4156

rdspring1 · 2025-03-30T20:22:53Z

This PR creates a new Pybind11 extension for the direct bindings API. The bindings does not have any functionality.

Key Build Details:

The python extension name is PYTHON_NEXT_EXTENSION=_C_NEXT.
It is used to create a new shared library titled libnvfuser_next.so
Installed as separate module named nvfuser_next.

github-actions · 2025-03-30T20:23:40Z

Review updated until commit 2266039

Description

Create nvfuser_next Pybind11 extension
Add import conflict checks between nvfuser and nvfuser_next
Update build_ext to handle nvfuser_next library
Add tests for import conflicts

Changes walkthrough 📝

Relevant files

Enhancement

6 files

bindings.cpp `Add bindings initialization for nvfuser_next`	+17/-0
extension.cpp `Define nvfuser_next module entry point`	+14/-0
__init__.py `Add import conflict check for nvfuser_next`	+6/-1
__init__.py `Initialize nvfuser_next module with import conflict check`	+20/-0
utils.py `Update build_ext to handle nvfuser_next`	+21/-13
bindings.h `Declare bindings initialization function for nvfuser_next`	+17/-0

Tests

2 files

test_python_frontend.py `Add test for import conflict between nvfuser and nvfuser_next`	+12/-0
test_python_next.py `Add test for import conflict between nvfuser_next and nvfuser`	+24/-0

Configuration changes

1 files

CMakeLists.txt `Add build rules for nvfuser_next extension`	+76/-3

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Functionality

The initNvFuserPythonBindings function does not currently add any bindings. Ensure that this function is implemented to include the necessary bindings for the nvfuser_next module.

void initNvFuserPythonBindings(PyObject* module) {
  auto python_bindings = py::handle(module).cast<py::module>();
}

Dependency

The pytorch_lib_dir is added to sys.path unconditionally. Verify if this is necessary and consider potential side effects on the environment.

pytorch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib")
if pytorch_lib_dir not in sys.path:
    sys.path.append(pytorch_lib_dir)

Redundancy

The target_compile_definitions and set_target_properties for nvf_py_next_internal and nvfuser_next are nearly identical. Consider refactoring to avoid redundancy.

target_compile_definitions(nvf_py_next_internal PRIVATE
  "-DTORCH_CUDA_BUILD_MAIN_LIB"
  "-DC10_BUILD_MAIN_LIB=1"
  PYTHON_NEXT_EXTENSION=_C_NEXT
)

add_library(nvfuser_next MODULE $<TARGET_OBJECTS:nvf_py_next_internal>)
target_compile_definitions(nvfuser_next PRIVATE
  "-DTORCH_CUDA_BUILD_MAIN_LIB"
  "-DC10_BUILD_MAIN_LIB=1"
  PYTHON_NEXT_EXTENSION=_C_NEXT
)

if(NOT MSVC)
  target_compile_options(nvf_py_next_internal PRIVATE -Wall -Wno-unused-function)
  target_compile_options(nvf_py_next_internal PRIVATE -Werror)
  set(NVF_LIB_SUFFIX ".so")
else()
  set(NVF_LIB_SUFFIX ".pyd")
endif()

set_target_properties(nvf_py_next_internal PROPERTIES
  C_STANDARD ${NVFUSER_C_STANDARD}
  CUDA_STANDARD ${NVFUSER_CUDA_STANDARD}
  CXX_STANDARD ${NVFUSER_CPP_STANDARD}
  CXX_STANDARD_REQUIRED ON
  CXX_VISIBILITY_PRESET hidden
  INSTALL_RPATH
  "$ORIGIN/lib:$ORIGIN/../nvidia/cuda_runtime/lib:$ORIGIN/../nvidia/cuda_nvrtc/lib:$ORIGIN/../../nvidia/cuda_cupti/lib:$ORIGIN/../torch/lib"
  POSITION_INDEPENDENT_CODE Yes
  VISIBILITY_INLINES_HIDDEN Yes
)
set_target_properties(nvfuser_next PROPERTIES
  C_STANDARD ${NVFUSER_C_STANDARD}
  CUDA_STANDARD ${NVFUSER_CUDA_STANDARD}
  CXX_STANDARD ${NVFUSER_CPP_STANDARD}
  CXX_STANDARD_REQUIRED ON
  CXX_VISIBILITY_PRESET hidden
  INSTALL_RPATH
  "$ORIGIN/lib:$ORIGIN/../nvidia/cuda_runtime/lib:$ORIGIN/../nvidia/cuda_nvrtc/lib:$ORIGIN/../../nvidia/cuda_cupti/lib:$ORIGIN/../torch/lib"
  POSITION_INDEPENDENT_CODE Yes
  SUFFIX ${NVF_LIB_SUFFIX}
  VISIBILITY_INLINES_HIDDEN Yes
)

wujingyue · 2025-03-31T15:56:59Z

Context? How is this related to the existing python api in csrc?

rdspring1 · 2025-03-31T16:06:32Z

Context? How is this related to the existing python api in csrc?

TL;DR: Direct bindings API maps directly to CPP Fusion objects. It doesn't rely on the FusionCache. e.g., You can make Fusions just like in CPP if desired.The end goal is for direct bindings to completely replace the existing python API.

I've been working on a prototype for direct bindings. #4077 This PR is the first to start integrating the changes into main.

Full design doc: https://docs.google.com/document/d/1ftdNKu952EFmLANa36g0IrVaqAJ0eLKzgJESkQFnuVI/edit?usp=sharing

jacobhinkle

LGTM

CMakeLists.txt

kevinstephano · 2025-04-03T15:01:21Z

I am suggesting that we modify this to appear as we would want it for a new Python API. Therefore, there should not be a reference to "direct" in the code as a user would not care.

The main changes I am suggesting is to change the direct from direct_bindings_api to python/src as the place to contain the C++ files for our binding code.

rdspring1 · 2025-05-06T22:53:05Z

!test

jjsjann123

LGTM.

Should we also add a python test file to verify the import would work?

jjsjann123 · 2025-05-07T16:28:14Z

CMakeLists.txt

+    codegen_internal
+    Python::Module
+  )
+  install(TARGETS nvfuser_next DESTINATION lib)


my nitpick on this PR is that, we would want to avoid duplicate the code that's shared between the two pybind module.

I'm also leaning towards moving this down to the ./python/CMakeLists.txt.
But we don't have to do that in this PR. We can leave that as a follow up clean PR.

What do you want to move to .python/CMakeLists.txt?

Build engineer is on the case.

just simply for logic separation.
If I'm adding another python module. I don't need to see the entire build files for nvfuser core lib.

jjsjann123 · 2025-05-07T16:28:56Z

python/utils.py

+            ext_modules=[
+                Extension(name="nvfuser._C", sources=[]),
+                Extension(name="nvfuser_next._C_NEXT", sources=[]),
+            ],


👏 I like how clean & easy the setup part is.

jjsjann123 · 2025-05-07T16:31:12Z

python/nvfuser_next/__init__.py

+pytorch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib")
+if pytorch_lib_dir not in sys.path:
+    sys.path.append(pytorch_lib_dir)
+


Naive question. Are we expecting people to import both nvfuser and nvfuser_next?

The way things are imported makes me worried about name conflicts.
i.e. should we assert on not seeing nvfuser in sys.modules when importing nvfuser_next?

Are we expecting people to import both nvfuser and nvfuser_next?

No.

should we assert on not seeing nvfuser in sys.modules when importing nvfuser_next?

Yes.

@jjsjann123 I added assertions to nvfuser/__init__.py and nvfuser_next/__init__.py and python tests to check for the conflicts.

kevinstephano

LGTM

CMakeLists.txt

direct_bindings_api/direct_bindings.cpp

direct_bindings_api/direct_bindings_extension.cpp

kevinstephano

LGTM

rdspring1 · 2025-05-07T17:27:13Z

!test

rdspring1 · 2025-05-07T18:07:25Z

Should we also add a python test file to verify the import would work?

I created test_import_correct for this.

rdspring1 added Python API Issues related to the Python API Direct Bindings Python extension with direct mapping to NvFuser CPP objects. labels Mar 30, 2025

rdspring1 requested review from jacobhinkle, wujingyue, jjsjann123 and kevinstephano March 31, 2025 15:54

jacobhinkle approved these changes Apr 2, 2025

View reviewed changes

CMakeLists.txt Outdated Show resolved Hide resolved

jjsjann123 reviewed Apr 2, 2025

View reviewed changes

CMakeLists.txt Outdated Show resolved Hide resolved

CMakeLists.txt Outdated Show resolved Hide resolved

rdspring1 force-pushed the direct_pt0 branch from 12eba74 to a720f8e Compare April 2, 2025 16:18

kevinstephano marked this pull request as draft April 3, 2025 15:01

rdspring1 force-pushed the direct_pt0 branch from a720f8e to ecfed35 Compare April 3, 2025 23:07

rdspring1 added 2 commits April 25, 2025 18:18

create NVFUSER_PYTHON_BINDINGS

6dd1137

Move common from python/nvfuser to python/nvfuser_common

31bb641

rdspring1 force-pushed the direct_pt0 branch from ecfed35 to 1952d09 Compare April 26, 2025 04:03

rdspring1 changed the base branch from main to next_pt0 April 26, 2025 04:03

rdspring1 force-pushed the direct_pt0 branch from 1952d09 to 203c119 Compare April 26, 2025 04:06

rdspring1 added the build label Apr 26, 2025

Create nvfuser_next module

4761c44

rdspring1 force-pushed the direct_pt0 branch from 203c119 to 4761c44 Compare April 26, 2025 04:09

rdspring1 changed the title ~~Create direct_bindings_api extension~~ Create nvfuser_next extension Apr 26, 2025

rdspring1 marked this pull request as ready for review April 26, 2025 17:27

rdspring1 force-pushed the next_pt0 branch 4 times, most recently from 23b5223 to ccad78d Compare May 2, 2025 00:00

Base automatically changed from next_pt0 to main May 6, 2025 18:12

Merge branch 'main' of https://github.com/nvidia/fuser into direct_pt0

b2dbb83

rdspring1 force-pushed the direct_pt0 branch from 5022a21 to b2dbb83 Compare May 6, 2025 21:15

rdspring1 requested a review from jjsjann123 May 6, 2025 21:15

rdspring1 added 2 commits May 6, 2025 14:17

revert

1962313

move back

21b3967

jjsjann123 approved these changes May 7, 2025

View reviewed changes

jjsjann123 reviewed May 7, 2025

View reviewed changes

kevinstephano reviewed May 7, 2025

View reviewed changes

kevinstephano approved these changes May 7, 2025

View reviewed changes

rdspring1 added 3 commits May 7, 2025 10:21

Merge branch 'main' of https://github.com/nvidia/fuser into direct_pt0

53ba020

Check for import conflicts

0d63c95

lint

2266039

rdspring1 merged commit 815d40a into main May 7, 2025
53 checks passed

rdspring1 deleted the direct_pt0 branch May 7, 2025 21:58

Create nvfuser_next extension #4156

Create nvfuser_next extension #4156

Uh oh!

Conversation

rdspring1 commented Mar 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Mar 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes walkthrough 📝

PR Reviewer Guide 🔍

Uh oh!

wujingyue commented Mar 31, 2025

Uh oh!

rdspring1 commented Mar 31, 2025

Uh oh!

jacobhinkle left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kevinstephano commented Apr 3, 2025

Uh oh!

rdspring1 commented May 6, 2025

Uh oh!

jjsjann123 left a comment

Choose a reason for hiding this comment

Uh oh!

jjsjann123 May 7, 2025

Choose a reason for hiding this comment

Uh oh!

rdspring1 May 7, 2025

Choose a reason for hiding this comment

Uh oh!

jjsjann123 May 7, 2025

Choose a reason for hiding this comment

Uh oh!

jjsjann123 May 7, 2025

Choose a reason for hiding this comment

Uh oh!

jjsjann123 May 7, 2025

Choose a reason for hiding this comment

Uh oh!

rdspring1 May 7, 2025

Choose a reason for hiding this comment

Uh oh!

rdspring1 May 7, 2025

Choose a reason for hiding this comment

Uh oh!

kevinstephano left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kevinstephano left a comment

Choose a reason for hiding this comment

Uh oh!

rdspring1 commented May 7, 2025

Uh oh!

rdspring1 commented May 7, 2025

Uh oh!

Uh oh!

Uh oh!

rdspring1 commented Mar 30, 2025 •

edited

Loading

github-actions bot commented Mar 30, 2025 •

edited

Loading