[User Schedule] Add print functions for TensorViews and User Scheduled Fusions #2282

rdspring1 · 2024-05-22T17:05:46Z

This PR adds two functions for debugging user scheduled fusions in python frontend.

std::string FusionDefinition::userScheduleIr() returns the user scheduled fusion IR.

Usage:

class SomeFusion(FusionDefinition):
    def definition(self):
        # create some fusion

    def schedule(self):
        print(fd.sched.user_schedule_ir())

to_string(python_frontend::Tensor) returns the string representation for its corresponding TensorView.

Usage:

print(fd.sched.to_string(some_tensor))

PR List:

naoyam

This seems really useful. Thanks! I'll let others chime in on the Python interface itself, but could you also create some documentation? Do we already have a doc for the Python interface? If so, I suppose that should be updated.

wujingyue · 2024-05-29T17:37:50Z

I'm behind for many things this week, so sorry for not being able to review this series of PRs. I think the line of the work (especially #2285) could potentially help @cowanmeg manually shard the Transformer block (#2199) more easily.

jjsjann123 · 2024-05-29T18:41:10Z

csrc/python_frontend/python_bindings.cpp

@@ -2777,6 +2781,20 @@ void initNvFuserPythonBindings(PyObject* module) {
  py::class_<FusionDefinition::SchedOperators> nvf_sched(
      fusion_def, "SchedOperators");
  nvf_sched.def(py::init<FusionDefinition*>());
+  nvf_sched.def(
+      "to_string",


nitpick: let's try to hide this by naming it properly with a leading _ like the other print things.

We wrap the print functions for FusionDefinition in nvfuser/__init__.py like https://github.com/NVIDIA/Fuser/blob/main/nvfuser/__init__.py#L246-L253. I'm not sure if we need the _ if it is a public function for this case.

sorry for nitpicking on this one. I wasn't sure that our intention is to have this as a public API. It's a bit confusing for the to_string on a pybind Tensor class, since we do have __repr__ for the bind class, but this to_string is actually printing out the TensorView.

But that's just a nitpicking. I'm find with keeping this as-is.

csrc/python_frontend/fusion_definition.cpp

jjsjann123 · 2024-05-29T19:25:45Z

csrc/python_frontend/python_bindings.cpp

@@ -593,6 +593,10 @@ void initNvFuserPythonBindings(PyObject* module) {
          "_fusion_ir",
          [](FusionDefinition& self) { return self.fusionIr(); },
          py::return_value_policy::reference)
+      .def(
+          "_user_schedule_ir",


IIUC, userScheduleIr can be safely called after a single execution with a custom schedule.

But we can also safely call it outside of def schedule(self): as well, my question is, is that an intended usage? Otherwise, if we are only supposed to use it inside schedule, we might as well keep this method under nvf_sched instead?

I'm not pushing strongly against this, but if we decide to keep the API as-is, I'd highly recommend relaxing the NVF_CHECK above with proper warning/recommendation on proper usage.

It could be useful to print fusion IR outside of the schedule function. I'll keep it and replace the NVF_CHECK with a string message.

I moved the method to nvf_sched since user_sched_ is reset after finalize_schedule is called and is not available outside of def schedule.

jjsjann123

LGTM, some minor comments on user experience but let's clean up in following PRs.

jjsjann123 · 2024-05-29T22:57:37Z

csrc/python_frontend/python_bindings.cpp

@@ -2777,6 +2781,20 @@ void initNvFuserPythonBindings(PyObject* module) {
  py::class_<FusionDefinition::SchedOperators> nvf_sched(
      fusion_def, "SchedOperators");
  nvf_sched.def(py::init<FusionDefinition*>());
+  nvf_sched.def(
+      "to_string",


sorry for nitpicking on this one. I wasn't sure that our intention is to have this as a public API. It's a bit confusing for the to_string on a pybind Tensor class, since we do have __repr__ for the bind class, but this to_string is actually printing out the TensorView.

But that's just a nitpicking. I'm find with keeping this as-is.

jjsjann123 · 2024-05-29T23:12:52Z

csrc/python_frontend/fusion_definition.cpp

@@ -228,6 +228,18 @@ std::string FusionDefinition::fusionIr() {
  return ss.str();
 }

+std::string FusionDefinition::userScheduleIr() {
+  NVF_CHECK(id().has_value(), "Invalid fusion definition!");


IIUC, vvv

class SomeFusion(FusionDefinition): def definition(self): # create some fusion def schedule(self): print(self._user_schedule_ir()) fd = SomeFusion() fd.user_schedule_ir() # it would assert here with `Invalid fusion definition`?

one more nitpicking, this really isn't your problem, but we can probably use a better error message and probably a warning rather than an assert. I think all the other FusionDefinition methods defined in nvfuser/__init__.py suffers the same.

Keep it as-is in this PR. But I think we can use a clean up PR afterwards. Wdyt @kevinstephano

jjsjann123

One last question: What's the story with version bump?

Since we are only adding new APIs here, I'm OK with keeping the version as-is and bump then towards the end of the stacked PRs. (and pretend that they do not exist).

One other thing is, let's add a test as well.

rdspring1 · 2024-05-30T00:39:20Z

I updated the version and add some tests.

rdspring1 · 2024-05-31T18:09:07Z

!build

…d Fusions (#2282) This PR adds two functions for debugging user scheduled fusions in python frontend. `std::string FusionDefinition::userScheduleIr()` returns the user scheduled fusion IR. Usage: ```python class SomeFusion(FusionDefinition): def definition(self): # create some fusion def schedule(self): print(fd.sched.user_schedule_ir()) ``` `to_string(python_frontend::Tensor)` returns the string representation for its corresponding TensorView. Usage: ```python print(fd.sched.to_string(some_tensor)) ```

Add print functions for TensorViews and User Scheduled Fusions

ba2c75c

rdspring1 added the Python API Issues related to the Python API label May 22, 2024

rdspring1 marked this pull request as ready for review May 22, 2024 19:15

naoyam reviewed May 23, 2024

View reviewed changes

rdspring1 requested review from wujingyue, jjsjann123 and kevinstephano May 28, 2024 16:00

wujingyue removed their request for review May 29, 2024 17:33

jjsjann123 reviewed May 29, 2024

View reviewed changes

rdspring1 added 2 commits May 29, 2024 13:30

update

5ed4de3

expose print function

ce58ceb

jjsjann123 approved these changes May 29, 2024

View reviewed changes

jjsjann123 requested changes May 29, 2024

View reviewed changes

rdspring1 added 3 commits May 29, 2024 17:23

Add test_print

e75fb5e

update user_schedule_ir

1e64bab

bump version

49cb5ac

rdspring1 requested a review from jjsjann123 May 30, 2024 00:39

rdspring1 added 3 commits May 31, 2024 10:27

Merge branch 'main' of github.com:NVIDIA/Fuser into user_print

a9b92c4

revert version bump

202c91b

move test

f9b6e62

jjsjann123 approved these changes Jun 1, 2024

View reviewed changes

rdspring1 merged commit 3e69e7c into main Jun 2, 2024
37 checks passed

rdspring1 deleted the user_print branch June 2, 2024 16:04

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[User Schedule] Add print functions for TensorViews and User Scheduled Fusions #2282

[User Schedule] Add print functions for TensorViews and User Scheduled Fusions #2282

rdspring1 commented May 22, 2024 •

edited

Loading

naoyam left a comment

wujingyue commented May 29, 2024

jjsjann123 May 29, 2024

rdspring1 May 29, 2024

jjsjann123 May 29, 2024

jjsjann123 May 29, 2024

rdspring1 May 29, 2024

rdspring1 May 30, 2024

jjsjann123 left a comment

jjsjann123 May 29, 2024

jjsjann123 May 29, 2024

jjsjann123 left a comment

rdspring1 commented May 30, 2024

rdspring1 commented May 31, 2024

[User Schedule] Add print functions for TensorViews and User Scheduled Fusions #2282

[User Schedule] Add print functions for TensorViews and User Scheduled Fusions #2282

Conversation

rdspring1 commented May 22, 2024 • edited Loading

PR List:

naoyam left a comment

Choose a reason for hiding this comment

wujingyue commented May 29, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jjsjann123 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jjsjann123 left a comment

Choose a reason for hiding this comment

rdspring1 commented May 30, 2024

rdspring1 commented May 31, 2024

rdspring1 commented May 22, 2024 •

edited

Loading