Add New Metadata and Pattern Features #2271

jsmonson · 2025-05-05T16:48:12Z

This PR will add new metadata and pattern features.

onnxscript/rewriter/pattern_builder_jsm.py

+import onnx
+
+from onnxscript import ir
+from onnxscript import rewriter


onnxscript/rewriter/pattern_builder_jsm.py

+            #else:
+            #    vp_outputs = builder.__getattr__(node.op_type)(*ninputs)


onnxscript/rewriter/pattern_builder_jsm.py

+
+        return ReplacementPatternGraph(g)
+
+    def get_iterator_index(self):


onnxscript/rewriter/pattern_builder_jsm.py

+    # Skip the (all) Activation inputs (have been swapped to beginning of the list)
+    for index in range(activations, len(nodes[0].inputs)):
+        inputs    = []
+        producers = []


onnxscript/rewriter/pattern_builder_jsm.py

+        cinput  = LoopBody.function.inputs[index]
+        noutput = vdisconnect(copy.copy(cinput))
+        noutput._uses = {}
+        update_node_outputs = False


onnxscript/utils/graph_view_utils.py

+                inputs.add(ninput)
+            elif any(ninput is init for init in node.graph.initializers):
+                initializers.add(ninput)
+            elif ninput.producer() == None:


onnxscript/utils/test_PytorchHierarchyNode.py

@@ -0,0 +1,190 @@
+import pytest


jsmonson · 2025-06-02T23:27:12Z

onnxscript/utils/test_PytorchHierarchyNode.py

+import ast
+
+import onnx
+from onnxscript import script


Removed unused import

tests/loop_rolling/test_loop_rolling.py

+# for rule in tracer.best_matches_map:
+#     matches = tracer.best_matches_map[rule]
+#     for match in matches:
+#         print(f'Reason: {match.match_result.reason}')
+#         print(f'root_node: {match.root_node}')
+#         pdb.set_trace()


tests/loop_rolling/test_loop_rolling.py

+# for node in ir.traversal.RecursiveGraphIterator(mypipeline_model.graph):
+#     if node.domain == '':
+#         print(node)


codecov · 2025-05-05T16:51:16Z

Codecov Report

Attention: Patch coverage is 6.72131% with 569 lines in your changes missing coverage. Please review.

Project coverage is 61.98%. Comparing base (2f04ea8) to head (2492dca).

Files with missing lines	Patch %	Lines
onnxscript/rewriter/pattern_builder_jsm.py	1.90%	309 Missing ⚠️
onnxscript/utils/graph_view_utils.py	15.69%	145 Missing ⚠️
onnxscript/utils/test_PytorchHierarchyNode.py	6.50%	115 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2271      +/-   ##
==========================================
- Coverage   70.16%   61.98%   -8.18%     
==========================================
  Files         198      201       +3     
  Lines       24846    25456     +610     
  Branches     2670     2779     +109     
==========================================
- Hits        17432    15779    -1653     
- Misses       6496     8824    +2328     
+ Partials      918      853      -65

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

github-advanced-security

lintrunner found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.

onnxscript/utils/graph_view_utils.py

justinchuby · 2025-05-07T17:10:48Z

onnxscript/utils/graph_view_utils.py

+        else:
+            return self.class_metadata[depth]
+
+class PytorchHierarchyNode:


This is useful! Let me see how we can provide the functionality to users.

justinchuby · 2025-05-07T17:14:50Z

onnxscript/rewriter/pattern_builder_jsm.py

+def append_output_to_node(node, output):
+    output._producer = node
+    output._index = node.outputs[-1]._index + 1
+    node._outputs = (*node._outputs, output)
+    node._num_outputs = len(node._outputs)
+
+def prepend_output_to_node(node, output):
+    output._producer = node
+    output._index = 0
+    for outp in node._outputs:
+        outp._index += 1
+    node._outputs = (output, *node._outputs)
+    node._num_outputs = len(node._outputs)


We should add this to node so there is no need to modify internal states

In fact we recommend creating new nodes if you want to update inputs and outputs. This way invariance of the graph is maintained. You may consider accumulating the inputs and outputs into two lists before constructing the node.

justinchuby · 2025-05-07T19:44:24Z

onnxscript/rewriter/pattern_builder_jsm.py

+            Ident = ir.Node(domain='',
+                            op_type='Identity',
+                            inputs = [cinput],
+                            outputs = [noutput],
+                            num_outputs =1)
+            LoopBody.function.append(Ident)
+
+            #Add Output to Function Call Nodes
+            for i,node in enumerate(nodes):
+                output_copy = copy.copy(noutput)
+
+                #preserve single_assignment
+                output_copy.name += f'_{i}'
+                append_output_to_node(node,output_copy)


Is there a way to collect all outputs before constructing the node?

gramalingam · 2025-05-08T23:21:20Z

onnxscript/rewriter/pattern_builder_jsm.py

+        vmap[input] = ValuePattern(input.name)
+
+    for init in graph.initializers:
+        vmap[init] = ValuePattern(init.name)


Looks like you want to map constants and initializers to unconstrained variables in the pattern? I wonder if it would make sense to map them to "Constants" in the pattern that require a matching contstant-value in the graph for a successful match? That makes reasonable, at least for simple and small constants. If it should be abstracted, wouldn't it be better for the user themselves to do that explicitly by mapping them to graph inputs?

Yes. Thanks for this suggestion. I think we should do this. This would provide a good level of control for the user.

gramalingam · 2025-05-08T23:23:52Z

onnxscript/rewriter/pattern_builder_jsm.py

+                ninputs.append(vmap[ninput])
+
+            #if len(node.outputs) > 1:
+            vp_outputs = builder.__getattr__(node.op_type)(*ninputs,_domain=node.domain, _outputs=len(node.outputs))


So, the attributes of the node are abstracted away, and not matched?

Good point. We ought to match the attributes. I'll make these changes.

justinchuby · 2025-06-02T23:32:47Z

onnxscript/rewriter/pattern_builder_jsm.py

+    node._outputs = (output, *node._outputs)
+    node._num_outputs = len(node._outputs)
+
+def prepend_input_to_node(node, input):


Modifying private fields is not recommended. I would consider initializing a new node

justinchuby · 2025-06-02T23:33:52Z

tests/loop_rolling/test_loop_rolling.py

+golden_results      = ort_run_graph(args.filename, input_dict, outputs[0].name)
+
+
+LoopBody = LoopBodyTemplate(args.patternfilename)


Suggested change

LoopBody = LoopBodyTemplate(args.patternfilename)

loop_body = LoopBodyTemplate(args.patternfilename)

note: always use snake_case for variable names

justinchuby · 2025-06-02T23:36:01Z

onnxscript/utils/graph_view_utils.py

+    return [output, used_output]
+
+
+def bGraphView(name, nodes):


nit: snake case for function names. Is this build_graph_view?

justinchuby · 2025-06-02T23:37:05Z

onnxscript/utils/graph_view_utils.py

+            usage.add("EXTERNAL")
+    return usage
+
+def find_subgraph_inputs(nodes):


This function is useful. I wonder if we should put it in ir.convenience https://github.com/onnx/ir-py/blob/main/src/onnx_ir/_convenience/__init__.py

As well as the one below for outputs and bGraphView

Joshua Monson and others added 18 commits March 27, 2025 21:37

initial loop rolling import

d2df97c

removed old ast code and commented code.

2f85bdc

add code to remove existing data files before writing to disk.

c7cbf5c

remove debug printing statements

8ace32f

fix imports and stuff

f71060a

updates to make tests orks

a1c138c

remove print statements

3274be3

Merge branch 'microsoft:main' into main

3870ff3

Merge branch 'main' of github.com:jsmonson/onnxscript

2f6572b

print_hierarchy works, initial tests pass

65985d8

now properly rejects non-hierarchical nodes

90f85d1

remove old code

e47161a

add comprehensive mistral test checking

cbbfe78

add code to place unannotated constant nodes

26f906c

remove print and old comments

abf892d

remove additional comments

b8c6068

Merge branch 'microsoft:main' into main

073e5db

Merge branch 'microsoft:main' into main

c731728

github-project-automation bot added this to ONNX Script Review Board May 5, 2025

github-project-automation bot moved this to Todo in ONNX Script Review Board May 5, 2025

github-advanced-security bot found potential problems May 5, 2025

View reviewed changes

justinchuby self-assigned this May 5, 2025

justinchuby reviewed May 7, 2025

View reviewed changes

onnxscript/utils/graph_view_utils.py Show resolved Hide resolved

justinchuby reviewed May 7, 2025

View reviewed changes

onnxscript/utils/graph_view_utils.py Show resolved Hide resolved

justinchuby reviewed May 7, 2025

View reviewed changes

gramalingam reviewed May 8, 2025

View reviewed changes

Merge branch 'microsoft:main' into joshmonson/add-md-and-pattern-feature

2492dca

justinchuby reviewed Jun 2, 2025

View reviewed changes

		#else:
		# vp_outputs = builder.__getattr__(node.op_type)(*ninputs)


		return ReplacementPatternGraph(g)

		def get_iterator_index(self):

		golden_results = ort_run_graph(args.filename, input_dict, outputs[0].name)


		LoopBody = LoopBodyTemplate(args.patternfilename)

	LoopBody = LoopBodyTemplate(args.patternfilename)
	loop_body = LoopBodyTemplate(args.patternfilename)

Add New Metadata and Pattern Features #2271

Are you sure you want to change the base?

Add New Metadata and Pattern Features #2271

Uh oh!

Conversation

jsmonson commented May 5, 2025

Uh oh!

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Choose a reason for hiding this comment

Uh oh!

Check notice

Check notice

codecov bot commented May 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-advanced-security bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

codecov bot commented May 5, 2025 •

edited

Loading