onnxslim 0.1.82__py3-none-any.whl → 0.1.84__py3-none-any.whl

onnxslim/third_party/onnx_graphsurgeon/ir/graph.py

@@ -794,109 +794,6 @@ class Graph:
                     tensor.to_constant(arr)
                     tensor.inputs.clear()
 
-        # Pass 2: Run shape-tensor cast elision
-        def run_cast_elision(node):
-            """Perform cast elision optimization on an ONNX node to eliminate unnecessary cast operations."""
-            import onnx
-
-            # Search for Cast(s) (from int -> float) -> intermediate operator (with float constants) -> Cast(s) (back to int)
-            # This pattern is problematic for TensorRT since these operations may be performed on Shape Tensors, which
-            # are not allowed to be floating point type. Attempt to fold the pattern here
-            VALID_CAST_ELISION_OPS = {
-                "Add",
-                "Sub",
-                "Mul",
-                "Div",
-                "Max",
-                "Min",
-                "Equal",
-                "Greater",
-                "Less",
-                "Concat",
-            }
-
-            if node.op not in VALID_CAST_ELISION_OPS:
-                return
-
-            # If the uncasted outputs of this node have any consumers other than "Cast" nodes,
-            # then we cannot elide the cast.
-            for out_tensor in node.outputs:
-                if out_tensor in self.outputs:
-                    return
-
-                if any(out_node.op != "Cast" for out_node in out_tensor.outputs):
-                    return
-
-            # Get list of input nodes that cast to float32
-            inp_casts = [
-                inp_node
-                for inp_tensor in node.inputs
-                for inp_node in inp_tensor.inputs
-                if inp_node.op == "Cast" and inp_node.attrs["to"] == onnx.TensorProto.DataType.FLOAT
-            ]
-
-            # No cast nodes found, return early
-            if not inp_casts:
-                return
-
-            # Ensure that all input cast nodes are casting from the same type
-            inp_dtypes = [dtype_to_onnx(inp_cast.inputs[0].dtype) for inp_cast in inp_casts]
-            if len(set(inp_dtypes)) != 1:
-                return
-
-            final_type = inp_dtypes[0]
-
-            # Get list of output nodes that cast to int32 or int64
-            out_casts = [
-                out_node
-                for out_tensor in node.outputs
-                for out_node in out_tensor.outputs
-                if out_node.op == "Cast"
-                and out_node.attrs["to"] in {onnx.TensorProto.DataType.INT32, onnx.TensorProto.DataType.INT64}
-            ]
-
-            # No cast node found on outputs, return early
-            if not out_casts:
-                return
-
-            # Ensure that all output cast nodes are casting to the same type and that this
-            # matches the original type before the inputs were casted.
-            out_dtypes = [out_cast.attrs["to"] for out_cast in out_casts]
-            if len(set(out_dtypes)) != 1 or out_dtypes[0] != final_type:
-                return
-
-            # If all checks passed, reconnect inputs/outputs to the consumers/producers
-            # of the Cast nodes.
-            # Note that we need to be careful in how we rebind tensors since they may
-            # be used by multiple nodes. Thus, it is not necessarily safe to assume that
-            # `cast_node.inputs[0].outputs[0] == cast_node`.
-            for index, inp in enumerate(node.inputs):
-                if isinstance(inp, Constant):
-                    inp.values = inp.values.astype(onnx.helper.tensor_dtype_to_np_dtype(final_type))
-
-                for cast in inp_casts:
-                    if cast.outputs[0] == inp:
-                        node.inputs[index] = cast.inputs[0]
-
-            for index, out in enumerate(node.outputs):
-                for cast in out_casts:
-                    if cast.inputs[0] == out:
-                        out_tensor = cast.outputs[0]
-                        out_tensor.inputs.clear()  # Disconnect from Cast
-                        node.outputs[index] = out_tensor
-
-        if fold_shapes:
-            # Perform shape tensor cast elision prior to most other folding
-            G_LOGGER.debug(f"Performing shape tensor cast elision in {self.name}")
-            try:
-                with self.node_ids():
-                    for node in self.nodes:
-                        run_cast_elision(node)
-            except Exception as err:
-                if not error_ok:
-                    raise err
-                G_LOGGER.warning("'{:}' routine failed with: {:}".format("Shape tensor cast elision", err))
-
         # Note that most of the remaining passes operate on a clone of the original graph.
         # Pass 3: Find all descendants of constant tensors