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


onnxscript/utils/test_PytorchHierarchyNode.py

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


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

❌ 1 Tests Failed:

Tests completed	Failed	Passed	Skipped
1755	1	1754	691

View the top 1 failed test(s) by shortest run time

onnxscript.utils.test_PytorchHierarchyNode::test_mistral_pytorch_with_metadata

Stack Traces | 0.001s run time

onnxscript/utils/test_PytorchHierarchyNode.py:165: in test_mistral_pytorch_with_metadata
    model_proto = onnx.load('.../finn_mlo_graphs/demo/mistral.onnx')
..../test_torch_nightly/lib/python3.11....../site-packages/onnx/__init__.py:226: in load_model
    model = _get_serializer(format, f).deserialize_proto(_load_bytes(f), ModelProto())
..../test_torch_nightly/lib/python3.11....../site-packages/onnx/__init__.py:163: in _load_bytes
    with open(f, "rb") as readable:
E   FileNotFoundError: [Errno 2] No such file or directory: '.../finn_mlo_graphs/demo/mistral.onnx'

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

github-advanced-security

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

justinchuby · 2025-05-07T17:04:49Z

onnxscript/utils/graph_view_utils.py

+
+
+def is_initializer(value):
+    return input.producer() == None


When a value is a graph input, its producer is also None. I will add a method to make this check more robust

justinchuby · 2025-05-07T17:07:22Z

onnxscript/utils/graph_view_utils.py

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


@jsmonson could you share how you use this function?

I see. It's for constructing the graph view

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.

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

gramalingam reviewed May 8, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add New Metadata and Pattern Features #2271

Add New Metadata and Pattern Features #2271

jsmonson commented May 5, 2025

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

codecov bot commented May 5, 2025

github-advanced-security bot left a comment

justinchuby May 7, 2025

justinchuby May 9, 2025

justinchuby May 7, 2025

justinchuby May 7, 2025

justinchuby May 7, 2025

justinchuby May 7, 2025

justinchuby May 9, 2025

justinchuby May 7, 2025

gramalingam May 8, 2025

jsmonson May 8, 2025

gramalingam May 8, 2025

jsmonson May 8, 2025

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


		return ReplacementPatternGraph(g)

		def get_iterator_index(self):

Add New Metadata and Pattern Features #2271

Are you sure you want to change the base?

Add New Metadata and Pattern Features #2271

Conversation

jsmonson commented May 5, 2025

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

Check notice

codecov bot commented May 5, 2025

❌ 1 Tests Failed:

github-advanced-security bot left a 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

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment