Repository: NVIDIA/Fuser
Pull request #4408: Create bindings for FusionDefinition and Add Binary op
Author: login: rdspring1
Labels: name: Python API
description: Issues related to the Python API
name: Direct Bindings
description: Python extension with direct mapping to NvFuser CPP objects.
Created: 2025-05-09T20:49:26Z | Merged: 2025-05-25T21:28:51Z

This PR creates the bindings for `FusionDefinition` class and a single `Add` binary op.

- The `FusionDefinition` class defines the  `Operator` and `ScheduleOperators` structs to hold all operators and schedule operations.
- It corresponds with [Step 1: Create python bindings for basic Fusion IR](https://docs.google.com/document/d/1ftdNKu952EFmLANa36g0IrVaqAJ0eLKzgJESkQFnuVI/edit?tab=t.0#heading=h.e2plo2gt0gnb).

PR Stack:
#4403 Create IterDomain, TensorDomain, and TensorView bindings
#4406 Add support for define_tensor
#4407 Add binding for Fusion IrContainer 
#4408 Create bindings for FusionDefinition and Add Binary op **<<< This PR.**
#4409 Create python FusionDefinition for nvfuser_next


Diff:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index b9f3210f044..c065dbb7aca 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -570,6 +570,7 @@ if(BUILD_PYTHON)
     ${NVFUSER_PYTHON_DIRECT_BINDINGS}/bindings.cpp
     ${NVFUSER_PYTHON_DIRECT_BINDINGS}/enum.cpp
     ${NVFUSER_PYTHON_DIRECT_BINDINGS}/ir.cpp
+    ${NVFUSER_PYTHON_DIRECT_BINDINGS}/ops.cpp
     ${NVFUSER_PYTHON_DIRECT_BINDINGS}/runtime.cpp
   )
   add_library(nvf_py_direct_internal OBJECT ${NVFUSER_PYTHON_DIRECT_SRCS})
diff --git a/python/python_direct/bindings.cpp b/python/python_direct/bindings.cpp
index efddc0626e1..0d568a32479 100644
--- a/python/python_direct/bindings.cpp
+++ b/python/python_direct/bindings.cpp
@@ -15,6 +15,7 @@ void initNvFuserPythonBindings(PyObject* module) {
   bindEnums(nvfuser);
   bindFusionIr(nvfuser);
   bindRuntime(nvfuser);
+  bindOperations(nvfuser);
 }
 
 } // namespace nvfuser::python
diff --git a/python/python_direct/bindings.h b/python/python_direct/bindings.h
index 5715fd18afc..252e20d02d1 100644
--- a/python/python_direct/bindings.h
+++ b/python/python_direct/bindings.h
@@ -23,4 +23,7 @@ void bindEnums(py::module& nvfuser);
 // Add bindings for Fusion and FusionExecutorCache
 void bindRuntime(py::module& nvfuser);
 
+// Add bindings for CPP Fusion Operations
+void bindOperations(py::module& nvfuser);
+
 } // namespace nvfuser::python
diff --git a/python/python_direct/ops.cpp b/python/python_direct/ops.cpp
new file mode 100644
index 00000000000..d653db9e23c
--- /dev/null
+++ b/python/python_direct/ops.cpp
@@ -0,0 +1,63 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <bindings.h>
+#include <ops/all_ops.h>
+
+namespace nvfuser::python {
+
+namespace {
+
+#define NVFUSER_DIRECT_BINDING_BINARY_OP(NAME, OP_NAME, DOCSTRING)             \
+  ops.def(NAME, [](Val* lhs, Val* rhs) -> Val* {                               \
+    return static_cast<Val* (*)(Val*, Val*)>(OP_NAME)(lhs, rhs);               \
+  });                                                                          \
+  ops.def(NAME, [](TensorView* lhs, Val* rhs) -> TensorView* {                 \
+    return static_cast<TensorView* (*)(TensorView*, Val*)>(OP_NAME)(lhs, rhs); \
+  });                                                                          \
+  ops.def(NAME, [](Val* lhs, TensorView* rhs) -> TensorView* {                 \
+    return static_cast<TensorView* (*)(Val*, TensorView*)>(OP_NAME)(lhs, rhs); \
+  });                                                                          \
+  ops.def(                                                                     \
+      NAME,                                                                    \
+      [](TensorView* lhs, TensorView* rhs) -> TensorView* {                    \
+        return static_cast<TensorView* (*)(TensorView*, TensorView*)>(         \
+            OP_NAME)(lhs, rhs);                                                \
+      },                                                                       \
+      DOCSTRING);
+
+void bindBinaryOps(py::module_& ops) {
+  // Use the macro for add operation
+  NVFUSER_DIRECT_BINDING_BINARY_OP(
+      "add",
+      add,
+      R"(
+Element-wise addition of two operands.
+
+Parameters
+----------
+lhs : Val or TensorView
+    The left-hand side operand.
+rhs : Val or TensorView
+    The right-hand side operand.
+
+Returns
+-------
+Val or TensorView
+    The sum of the inputs.
+)")
+};
+
+} // namespace
+
+void bindOperations(py::module& nvfuser) {
+  py::module_ nvf_ops = nvfuser.def_submodule(
+      "ops", "This submodule contains all operations for NvFuser.");
+  bindBinaryOps(nvf_ops);
+}
+
+} // namespace nvfuser::python


Reviews:
user: login: jjsjann123
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T21:29:50Z
user: login: rdspring1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T22:34:39Z
user: login: jjsjann123
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:21:34Z
user: login: rdspring1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T00:15:36Z
user: login: jjsjann123
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-23T21:07:47Z

Review comments:
diff_hunk: @@ -0,0 +1,41 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <torch/csrc/jit/python/pybind.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace nvfuser::python {
+
+//! The FusionDefinition class exposes the NvFuser operations to mimic the
+//! interface of the FusionDefinition class.
+class FusionDefinition {
path: python/python_direct/fusion_definition.h
user: login: jjsjann123
body: what's the purpose of having `FusionDefinition` in the new API?

since we already have `Fusion` directly exposed, what does `FusionDefinition` provide on top of that?
created_at: 2025-05-22T21:28:46Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,41 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <torch/csrc/jit/python/pybind.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace nvfuser::python {
+
+//! The FusionDefinition class exposes the NvFuser operations to mimic the
+//! interface of the FusionDefinition class.
+class FusionDefinition {
path: python/python_direct/fusion_definition.h
user: login: rdspring1
body: The `FusionDefinition.ops` API must be exactly the same, so Thunder doesn't notice the difference.

```
# NOTE nv_abs to avoid a name conflict with the builin abs
def nv_abs(a: Number | TensorProxy, /, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any:
    nva = getnv(a, fd, lc_to_nv_map)

    return fd.ops.abs(nva)
```
created_at: 2025-05-22T22:34:39Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,41 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <torch/csrc/jit/python/pybind.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace nvfuser::python {
+
+//! The FusionDefinition class exposes the NvFuser operations to mimic the
+//! interface of the FusionDefinition class.
+class FusionDefinition {
path: python/python_direct/fusion_definition.h
user: login: jjsjann123
body: If it's just for API compatibility, we can wrap it in pure python instead of polluting our pybind bindings.
created_at: 2025-05-22T23:21:34Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,41 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <torch/csrc/jit/python/pybind.h>
+#include <torch/csrc/utils/pybind.h>
+
+namespace nvfuser::python {
+
+//! The FusionDefinition class exposes the NvFuser operations to mimic the
+//! interface of the FusionDefinition class.
+class FusionDefinition {
path: python/python_direct/fusion_definition.h
user: login: rdspring1
body: > If it's just for API compatibility, we can wrap it in pure python instead of polluting our pybind bindings.

Can you show me exactly how to do this? e.g., How do you create `ops` in a python class `FusionDefinition`?
created_at: 2025-05-23T00:15:35Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-09T20:50:14Z
body: <!-- llm action combined persistent comment, dont modify this line -->

*Review updated until commit https://github.com/NVIDIA/Fuser/commit/fc4b2d6536d1cd8e5583efee254e83b2ffeb6922*


### **Description**
- Added bindings for `FusionDefinition` class

- Introduced `Add` binary operation bindings

- Updated `bindings.cpp` to include operations binding

- Created `ops.cpp` for operation bindings implementation


___



### **Changes walkthrough** 📝
<table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Enhancement</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>bindings.cpp</strong><dd><code>Include operations binding</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

python/python_direct/bindings.cpp

- Added call to `bindOperations` in `initNvFuserPythonBindings`


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4408/files#diff-92fe0ce126667fe71d35d768baa731937de0fd5c5513ff228ea2588c8b1dff8b">+1/-0</a>&nbsp; &nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>ops.cpp</strong><dd><code>Implement add operation binding</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

python/python_direct/ops.cpp

<li>Created <code>bindBinaryOps</code> function for binary operations<br> <li> Implemented <code>add</code> operation binding using macro


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4408/files#diff-d8770afa65be6240f18158e8a57f1f2f6ed7de43ee03651d812a3dc182d1f59b">+63/-0</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>bindings.h</strong><dd><code>Declare operations binding function</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

python/python_direct/bindings.h

- Added declaration for `bindOperations` function


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4408/files#diff-40c29ed02d3c7b3d2d29937ee7c18a7f35416566a7e158f1d1e632109a28db32">+3/-0</a>&nbsp; &nbsp; &nbsp; </td>

</tr>
</table></td></tr><tr><td><strong>Configuration changes</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>CMakeLists.txt</strong><dd><code>Include ops.cpp in build</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

CMakeLists.txt

- Added `ops.cpp` to the list of source files


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4408/files#diff-1e7de1ae2d059d21e1dd75d5812d5a34b0222cef273b7c3a2af62eb747f9d20a">+1/-0</a>&nbsp; &nbsp; &nbsp; </td>

</tr>
</table></td></tr></tr></tbody></table>

---

### PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

<table>
<tr><td>🧪&nbsp;<strong>No relevant tests</strong></td></tr>
<tr><td>⚡&nbsp;<strong>Recommended focus areas for review</strong><br><br>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4408/files#diff-d8770afa65be6240f18158e8a57f1f2f6ed7de43ee03651d812a3dc182d1f59bR15-R31'><strong>Possible Issue</strong></a>

The macro `NVFUSER_DIRECT_BINDING_BINARY_OP` defines multiple overloads for the `add` function. Ensure that these overloads do not lead to ambiguity or unexpected behavior when called with different types of arguments.</summary>

```c++
#define NVFUSER_DIRECT_BINDING_BINARY_OP(NAME, OP_NAME, DOCSTRING)             \
  ops.def(NAME, [](Val* lhs, Val* rhs) -> Val* {                               \
    return static_cast<Val* (*)(Val*, Val*)>(OP_NAME)(lhs, rhs);               \
  });                                                                          \
  ops.def(NAME, [](TensorView* lhs, Val* rhs) -> TensorView* {                 \
    return static_cast<TensorView* (*)(TensorView*, Val*)>(OP_NAME)(lhs, rhs); \
  });                                                                          \
  ops.def(NAME, [](Val* lhs, TensorView* rhs) -> TensorView* {                 \
    return static_cast<TensorView* (*)(Val*, TensorView*)>(OP_NAME)(lhs, rhs); \
  });                                                                          \
  ops.def(                                                                     \
      NAME,                                                                    \
      [](TensorView* lhs, TensorView* rhs) -> TensorView* {                    \
        return static_cast<TensorView* (*)(TensorView*, TensorView*)>(         \
            OP_NAME)(lhs, rhs);                                                \
      },                                                                       \
      DOCSTRING);
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4408/files#diff-d8770afa65be6240f18158e8a57f1f2f6ed7de43ee03651d812a3dc182d1f59bR35-R53'><strong>Documentation</strong></a>

Ensure that the docstring for the `add` function is clear and comprehensive, and that it accurately reflects the behavior of the function, especially regarding the types of arguments it accepts and the type of the return value.</summary>

```c++
  NVFUSER_DIRECT_BINDING_BINARY_OP(
      "add",
      add,
      R"(
Element-wise addition of two operands.

Parameters
----------
lhs : Val or TensorView
    The left-hand side operand.
rhs : Val or TensorView
    The right-hand side operand.

Returns
-------
Val or TensorView
    The sum of the inputs.
)")
};
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4408/files#diff-1e7de1ae2d059d21e1dd75d5812d5a34b0222cef273b7c3a2af62eb747f9d20aR573-R573'><strong>Build Configuration</strong></a>

Verify that adding `ops.cpp` to the build configuration does not introduce any build issues or conflicts with existing files.</summary>

```text
${NVFUSER_PYTHON_DIRECT_BINDINGS}/ops.cpp
```

</details>

</td></tr>
</table>

author_association: NONE
user: login: rdspring1
created_at: 2025-05-22T00:02:29Z
body: !test
author_association: COLLABORATOR
user: login: rdspring1
created_at: 2025-05-23T19:51:22Z
body: !test
author_association: COLLABORATOR