Add Fusion.print, which prints the transforms as well as the math

[wujingyue]

No description provided.

[github-actions]

Description

• Add Fusion.print method to display fusion IR

• Print includes both transforms and mathematical representation

• Exposes new Python binding for fusion debugging

Changes walkthrough

	Relevant files
Enhancement	runtime.cpp Add Fusion.print method binding                                                    python/python_direct/runtime.cpp Add new print method to Fusion class Method prints fusion IR and transforms to stdout Python binding exposes the functionality in Python +10/-0

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 No relevant tests

⚡ Recommended focus areas for review

Missing Return Value Handling The new print method is defined to return None according to its docstring, but it directly calls f.print(std::cout) without explicit handling of return values or side effects. This should be validated to ensure consistency with the intended API behavior and documentation. "print", [](Fusion& f) { f.print(std::cout); }, R"( Print the fusion, including transforms, to stdout. Returns ------- None )")

[wujingyue]

I saw other print_*s return a string instead. I personally found it confusing to have to type print twice, but let me know what you think. One option may be to have an __str__.

[rdspring1]

Returning a string is nice, so you can search through it in tests.
I thought about renaming print_* to something more clear, since string are returned instead being printed to stdout.

You could map CPP print to __repr__ since it is more verbose than print_math. Then, do print(repr(fd.fusion)).

[greptile-apps]

Greptile Overview

Greptile Summary

Added a new Python binding method Fusion.print() that prints the complete fusion IR including both mathematical expressions and tensor transforms to stdout, complementing the existing print_math() method which only shows mathematical expressions.

Key changes:

• New print() method added to Python Fusion class bindings
• Calls underlying C++ Fusion::print(std::cout) which includes transforms by default
• Provides comprehensive view of the fusion including inputs, outputs, math operations, and tensor transformations

Issues found:

• The new method has inconsistent API design compared to other print methods (print_math and print_kernel) which return strings instead of printing to stdout

Confidence Score: 3/5

• Safe to merge with minor API consistency concern
• The implementation is functionally correct and calls the appropriate C++ method, but has an API design inconsistency where it prints to stdout rather than returning a string like the other print methods
• Consider reviewing python/python_direct/runtime.cpp for API consistency with existing print methods

Important Files Changed

File Analysis

Filename	Score	Overview
python/python_direct/runtime.cpp	3/5	Adds Fusion.print() method that prints to stdout, but inconsistent with other print methods which return strings

Sequence Diagram

sequenceDiagram
    participant Python
    participant PythonBinding as Python Binding (runtime.cpp)
    participant Fusion as Fusion C++ Object
    participant IrPrinter
    participant IrTransformPrinter
    participant stdout

    Python->>PythonBinding: fusion.print()
    PythonBinding->>Fusion: f.print(std::cout)
    Fusion->>stdout: Write "Inputs:"
    Fusion->>stdout: Write input tensors
    Fusion->>stdout: Write "Outputs:"
    Fusion->>stdout: Write output tensors
    Fusion->>stdout: Write "\n%kernel {\n"
    Fusion->>IrPrinter: op_exprs.handle(this)
    IrPrinter->>stdout: Write math expressions
    Fusion->>stdout: Write "\nTransformPrinter : \n"
    Fusion->>IrTransformPrinter: t_exprs.handle(this)
    IrTransformPrinter->>stdout: Write tensor transforms
    Fusion->>stdout: Write "} // %kernel\n"
    Fusion-->>PythonBinding: Return ostream reference
    PythonBinding-->>Python: Return None

Loading

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

style: inconsistent return behavior - print_math and print_kernel both return strings by capturing output in a stringstream, but this new print method prints directly to stdout and returns None. Consider making it consistent:

Suggested change

	[](Fusion& f) { f.print(std::cout); },
	[](Fusion& f) {
	std::stringstream ss;
	DebugStreamGuard dsg(ss);
	f.print(ss);
	return ss.str();
	},