Fix constant in constant folding (#2622)

This PR moves the processing of constant ops upward to return before
node-level shape type inference (including serialization) and optimizer
optimization.

Essentially, avoiding serializing constant ops (potentially large
weights in LLMs) reduces the export time in optimize_ir.


Before this PR:

<img width="2119" height="282" alt="Screenshot 2025-10-09 141403"
src="https://github.com/user-attachments/assets/4d50d3ee-ce84-4f8b-bc20-f56497d7dad1"
/>

After this PR:

<img width="2390" height="319" alt="Screenshot 2025-10-09 141238"
src="https://github.com/user-attachments/assets/f455942e-dde8-4a07-a94e-a5f817358579"
/>

Author: titaiwangms

onnxscript/optimizer/_constant_folding.py

@@ -76,7 +76,7 @@ def _is_onnx_op(node: ir.Node, op_type: str) -> bool:
 
 def _process_constant_node(node: ir.Node) -> None:
     """Sets const_value of output value of a Constant op node."""
-    if node.op_type != "Constant" or node.domain != "":
+    if not _is_onnx_op(node, "Constant"):
         return
     if len(node.attributes) != 1:
         return
@@ -1099,8 +1099,12 @@ def process_node(self, node: ir.Node) -> Replacement | None:
                 self._modified = True
                 # TODO(rama): consider merging type/other info from both values
 
+        # Propagate const_value, and manually find out shape and type
+        # to avoid potentially expensive shape inference on large tensors.
+        if _is_onnx_op(node, "Constant"):
+            _process_constant_node(node)
         # Do incremental shape inference
-        if self.shape_inference and not _is_control_flow_op(node):
+        elif self.shape_inference and not _is_control_flow_op(node):
             self._do_inference(node)
 
         if node.domain not in self._opset_imports:
@@ -1118,6 +1122,10 @@ def process_node(self, node: ir.Node) -> Replacement | None:
                     output = [output]
                 return Replacement(output, context.nodes)
 
+        if _is_onnx_op(node, "Constant"):
+            logger.debug("Skipping constant folding for Constant node %r", node.name)
+            return None
+
         if _is_control_flow_op(node):
             logger.info(
                 "Skipping constant folding for control flow op %r (%s::%s) because it is not supported yet",
@@ -1137,10 +1145,6 @@ def process_node(self, node: ir.Node) -> Replacement | None:
             )
             return None
 
-        if _is_onnx_op(node, "Constant"):
-            _process_constant_node(node)
-            return None
-
         if any(x.is_graph_input() for x in node.inputs if x is not None):
             logger.info(
                 "Skipping constant folding for node %r because it is graph input to preserve graph signature",