Restore logic for remove_redundant_casts

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>

Author: ajrasane

modelopt/onnx/autocast/graphsanitizer.py

@@ -130,7 +130,7 @@ def remove_disconnected_outputs(self) -> None:
         """Remove disconnected outputs from the model."""
         tensors_to_remove = []
         for tensor in self.model.graph.output:
-            if not utils.get_producer_nodes(self.model, tensor.name):
+            if not onnx_utils.get_producer_nodes(self.model, tensor.name):
                 tensors_to_remove.append(tensor)
                 logger.debug(f"Found disconnected output: {tensor.name}")
 
@@ -279,7 +279,7 @@ def _match_layernorm_pattern(self, mean_node: onnx.NodeProto) -> dict | None:
             # Find variance computation branch
             pow_nodes = [
                 n
-                for n in utils.get_consumer_nodes(self.model, sub_node.output[0])
+                for n in onnx_utils.get_consumer_nodes(self.model, sub_node.output[0])
                 if n.op_type == "Pow"
             ]
             if len(pow_nodes) != 1:
@@ -303,8 +303,8 @@ def _match_layernorm_pattern(self, mean_node: onnx.NodeProto) -> dict | None:
 
             # Find Div node
             # Find the Div node that consumes both sqrt and sub outputs
-            sqrt_consumers = utils.get_consumer_nodes(self.model, sqrt_node.output[0])
-            sub_consumers = utils.get_consumer_nodes(self.model, sub_node.output[0])
+            sqrt_consumers = onnx_utils.get_consumer_nodes(self.model, sqrt_node.output[0])
+            sub_consumers = onnx_utils.get_consumer_nodes(self.model, sub_node.output[0])
 
             div_nodes = [n for n in sqrt_consumers if n in sub_consumers and n.op_type == "Div"]
             if len(div_nodes) != 1:
@@ -342,14 +342,14 @@ def _match_layernorm_pattern(self, mean_node: onnx.NodeProto) -> dict | None:
                 div_node,
             ]
 
-            consumers = utils.get_consumer_nodes(self.model, div_node.output[0])
+            consumers = onnx_utils.get_consumer_nodes(self.model, div_node.output[0])
             if len(consumers) == 1 and consumers[0].op_type == "Mul":
                 mul_node = consumers[0]
                 scale = self._get_initializer_value(mul_node.input[1], return_array=True)
                 final_node = mul_node
                 nodes_to_remove.append(mul_node)
 
-                consumers = utils.get_consumer_nodes(self.model, mul_node.output[0])
+                consumers = onnx_utils.get_consumer_nodes(self.model, mul_node.output[0])
                 if len(consumers) == 1 and consumers[0].op_type == "Add":
                     add_node = consumers[0]
                     bias = self._get_initializer_value(add_node.input[1], return_array=True)
@@ -457,7 +457,7 @@ def _create_layernorm_node(self, pattern: dict) -> onnx.NodeProto:
 
     def _find_insertion_point(self, input_name: str) -> int:
         """Find the correct insertion point for the new LayerNorm node."""
-        producer_nodes = utils.get_producer_nodes(self.model, input_name)
+        producer_nodes = onnx_utils.get_producer_nodes(self.model, input_name)
         if not producer_nodes:
             return 0
 

modelopt/onnx/autocast/precisionconverter.py

@@ -914,51 +914,12 @@ def _get_name(node: onnx.NodeProto | InputIndexTracker) -> str:
 
         return self._convert_initializer_data(init, from_type, to_type)
 
-    def _replace_tensor_name(
-        self, consumers: list[onnx.NodeProto], original_tensor_name: str, new_tensor_name: str
-    ) -> None:
-        """Replace occurrences of a tensor name in the given consumers' inputs with a new tensor name."""
-        for consumer in consumers:
-            for idx, inp in enumerate(consumer.input):
-                if inp == original_tensor_name:
-                    consumer.input[idx] = new_tensor_name
-
-    def _bypass_cast_node(self, node: onnx.NodeProto) -> None:
-        # handling only a single input and output, as we only remove cast nodes
-        assert len(node.input) == 1
-        assert len(node.output) == 1
-
-        input_tensor = node.input[0]
-        output_tensor = node.output[0]
-
-        # Check if the cast output is also a graph output
-        is_output_producer = any(output.name == output_tensor for output in self.model.graph.output)
-
-        # If the removed cast node is producing a network output, update the producer of the cast input so
-        # the network output name is preserved.
-        if is_output_producer:
-            producers = utils.get_producer_nodes(self.model, input_tensor)
-            for producer in producers:
-                for i, prod_out in enumerate(producer.output):
-                    if prod_out == input_tensor:
-                        producer.output[i] = output_tensor
-                        consumers = utils.get_consumer_nodes(self.model, prod_out)
-                        if len(consumers) > 1:
-                            self._replace_tensor_name(consumers, prod_out, output_tensor)
-        else:
-            # Reconnect consumers of the cast output to use the cast input instead
-            consumers = utils.get_consumer_nodes(self.model, output_tensor)
-            for consumer in consumers:
-                for i, input_name in enumerate(consumer.input):
-                    if input_name == output_tensor:
-                        consumer.input[i] = input_tensor
-
     def _remove_preexisting_casts(self) -> None:
         nodes_to_remove = []
         for node in self.model.graph.node:
             if node.op_type == "Cast":
-                cast_from_type = self._get_tensor_type(node.input[0])
-                cast_to_type = utils.get_cast_to_type(node)
+                cast_from_type = onnx_utils._get_tensor_type_by_name(self.model, node.input[0])
+                cast_to_type = onnx_utils.get_cast_to_type(node)
                 is_fp_cast = cast_to_type in [
                     onnx.TensorProto.FLOAT16,
                     onnx.TensorProto.FLOAT,
@@ -978,7 +939,7 @@ def _remove_preexisting_casts(self) -> None:
                     ):
                         continue
                     nodes_to_remove.append(node)
-                    self._bypass_cast_node(node)
+                    onnx_utils._bypass_cast_node(self.model, node)
         logger.debug(f"Removing {len(nodes_to_remove)} pre-existing casts")
 
         for node in nodes_to_remove:
@@ -1044,7 +1005,7 @@ def _add_cast(
         )
 
         if tensor_to_consumers is None:
-            consumer_nodes = utils.get_consumer_nodes(self.model, tensor_name)
+            consumer_nodes = onnx_utils.get_consumer_nodes(self.model, tensor_name)
         else:
             consumer_nodes = tensor_to_consumers.get(tensor_name, [])
         consumer_nodes = [n for n in consumer_nodes if n.name not in exclude_consumers]
@@ -1067,7 +1028,7 @@ def _add_cast(
 
         # Find producer node to insert cast after it
         if tensor_to_producers is None:
-            producer_nodes = utils.get_producer_nodes(self.model, tensor_name)
+            producer_nodes = onnx_utils.get_producer_nodes(self.model, tensor_name)
         else:
             producer_nodes = tensor_to_producers.get(tensor_name, [])
         if producer_nodes:
@@ -1106,7 +1067,7 @@ def _cleanup_no_consumer_nodes(self):
             node
             for node in self.model.graph.node
             if not any(
-                out in network_outputs or utils.get_consumer_nodes(self.model, out)
+                out in network_outputs or onnx_utils.get_consumer_nodes(self.model, out)
                 for out in node.output
             )
         ]
@@ -1124,29 +1085,23 @@ def _cleanup_pre_output_same_type_cast(self):
 
         for output in self.model.graph.output:
             if "_cast_to_" in output.name:
-                out_producer_nodes = utils.get_producer_nodes(self.model, output.name)
+                out_producer_nodes = onnx_utils.get_producer_nodes(self.model, output.name)
                 if len(out_producer_nodes) == 1 and out_producer_nodes[0].op_type == "Cast":
                     second_cast_node = out_producer_nodes[0]
-                    cast_producer_nodes = utils.get_producer_nodes(
+                    cast_producer_nodes = onnx_utils.get_producer_nodes(
                         self.model, second_cast_node.input[0]
                     )
                     if len(cast_producer_nodes) == 1 and cast_producer_nodes[0].op_type == "Cast":
                         first_cast_node = cast_producer_nodes[0]
                         if (
-                            self._is_same_type_cast(first_cast_node)
-                            and utils.get_cast_to_type(second_cast_node)
+                            onnx_utils._is_same_type_cast(self.model, first_cast_node)
+                            and onnx_utils.get_cast_to_type(second_cast_node)
                             == self.high_precision_type.onnx_type
                         ):
                             logger.debug(f"Removing pre-output double cast: {first_cast_node.name}")
-                            self._bypass_cast_node(first_cast_node)
+                            onnx_utils._bypass_cast_node(self.model, first_cast_node)
                             self.model.graph.node.remove(first_cast_node)
 
-    def _is_same_type_cast(self, node: onnx.NodeProto) -> bool:
-        assert node.op_type == "Cast"
-        input_types = [self._get_tensor_type(inp) for inp in node.input]
-        output_type = utils.get_cast_to_type(node)
-        return all(inp_type == output_type for inp_type in input_types) and input_types is not None
-
     def _remove_redundant_casts(self):
         """Removes both sequential casts and casts that don't change precision.
 
@@ -1176,7 +1131,7 @@ def _fix_network_output_names(self):
         for output in self.model.graph.output:
             if "_cast_to_" in output.name:
                 post_cast_name = output.name
-                producer_nodes = utils.get_producer_nodes(self.model, output.name)
+                producer_nodes = onnx_utils.get_producer_nodes(self.model, output.name)
                 if (
                     len(producer_nodes) == 1
                     and producer_nodes[0].op_type == "Cast"
@@ -1188,21 +1143,23 @@ def _fix_network_output_names(self):
                     pre_cast_name = original_name + "_pre_cast"
                     output.name = original_name
                     # Update all consumers of the original (pre-cast) output to use the pre-cast name
-                    for node in utils.get_consumer_nodes(self.model, original_name):
+                    for node in onnx_utils.get_consumer_nodes(self.model, original_name):
                         if node == cast_node:
                             continue
                         for i, input_name in enumerate(node.input):
                             if input_name == original_name:
                                 node.input[i] = pre_cast_name
                                 # do not break, can use the same tensor for multiple node inputs
                     # Update all consumers of the post-cast output to use the original name
-                    for node in utils.get_consumer_nodes(self.model, post_cast_name):
+                    for node in onnx_utils.get_consumer_nodes(self.model, post_cast_name):
                         for i, input_name in enumerate(node.input):
                             if input_name == post_cast_name:
                                 node.input[i] = original_name
                                 # do not break, can use the same tensor for multiple node inputs
                     # Update all producers of the original output to use the original name
-                    cast_producer_nodes = utils.get_producer_nodes(self.model, cast_node.input[0])
+                    cast_producer_nodes = onnx_utils.get_producer_nodes(
+                        self.model, cast_node.input[0]
+                    )
                     for node in cast_producer_nodes:
                         for i, node_output in enumerate(node.output):
                             if node_output == original_name:
@@ -1248,7 +1205,7 @@ def _sanity_check(self):
 
         # Verify that the output tensors are not disconnected
         for output in network_outputs:
-            producer_nodes = utils.get_producer_nodes(self.model, output.name)
+            producer_nodes = onnx_utils.get_producer_nodes(self.model, output.name)
             if len(producer_nodes) == 0:
                 logger.warning(
                     f"Output tensor {output.name} is disconnected. This may be benign if it's part of a cast operation "
@@ -1296,13 +1253,6 @@ def _sanity_check(self):
         if not sanity_ok:
             raise Exception("Sanity Check Failed")
 
-    def _get_tensor_type(self, tensor_name):
-        if tensor_name in self.value_info_map:
-            return self.value_info_map[tensor_name].type.tensor_type.elem_type
-        if tensor_name in self.initializer_map:
-            return self.initializer_map[tensor_name].data_type
-        raise Exception(f"did not find tensor {tensor_name}")
-
     def _sanitize_model(self):
         graph_sanitizer = GraphSanitizer(
             self.model,

modelopt/onnx/autocast/utils.py

@@ -27,6 +27,7 @@
 
 import onnx
 
+import modelopt.onnx.utils as onnx_utils
 from modelopt.onnx.utils import get_opset_version
 
 
@@ -60,32 +61,6 @@ def setup_mappings(model: onnx.ModelProto) -> tuple[dict, dict, dict]:
     return value_info_map, initializer_map, node_to_init_map
 
 
-def get_consumer_nodes(model: onnx.ModelProto, tensor_name: str) -> list[onnx.NodeProto]:
-    """Get all consumer nodes for a given tensor name.
-
-    Args:
-        model: The ONNX model to search.
-        tensor_name: Name of the tensor to find consumers for.
-
-    Returns:
-        list[onnx.NodeProto]: List of nodes that consume the tensor.
-    """
-    return [n for n in model.graph.node if tensor_name in n.input]
-
-
-def get_producer_nodes(model: onnx.ModelProto, tensor_name: str) -> list[onnx.NodeProto]:
-    """Get all producer nodes for a given tensor name.
-
-    Args:
-        model: The ONNX model to search.
-        tensor_name: Name of the tensor to find producers for.
-
-    Returns:
-        list[onnx.NodeProto]: List of nodes that produce the tensor.
-    """
-    return [n for n in model.graph.node if tensor_name in n.output]
-
-
 def get_unique_consumer_node(model: onnx.ModelProto, tensor_name: str) -> onnx.NodeProto:
     """Get a single consumer node and raise exception if there are multiple consumers.
 
@@ -99,30 +74,12 @@ def get_unique_consumer_node(model: onnx.ModelProto, tensor_name: str) -> onnx.N
     Raises:
         Exception: If there is not exactly one consumer node.
     """
-    consumers = get_consumer_nodes(model, tensor_name)
+    consumers = onnx_utils.get_consumer_nodes(model, tensor_name)
     if len(consumers) != 1:
         raise Exception(f"Expected single consumer for {tensor_name}, found {len(consumers)}")
     return consumers[0]
 
 
-def get_cast_to_type(cast_node: onnx.NodeProto) -> int:
-    """Get the target type from a Cast node.
-
-    Args:
-        cast_node: The Cast node to extract type from.
-
-    Returns:
-        int: The target type value from the Cast node's 'to' attribute.
-
-    Raises:
-        ValueError: If the Cast node does not have a 'to' attribute.
-    """
-    for attr in cast_node.attribute:
-        if attr.name == "to":
-            return attr.i
-    raise ValueError("Cast node does not have 'to' attribute")
-
-
 def walk_subgraphs_recursive(
     graph: onnx.GraphProto,
     callback: Callable,