Repository: huggingface/optimum (Python)

Issue:
INT64 clamping to INT32 creates overhead while using TensorRT
### System Info

```shell
optimum: dev
CUDA: 11.3
cuDNN: 8.3.2
TensorRT: 8.4.1.5
```

### Who can help?

@JingyaHuang

### Reproduction

```python
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

ort_model = ORTModelForSequenceClassification.from_pretrained(
    "philschmid/tiny-bert-sst2-distilled",
    from_transformers=True,
    provider="TensorrtExecutionProvider",
)

tokenizer = AutoTokenizer.from_pretrained("philschmid/tiny-bert-sst2-distilled")
inp = tokenizer("expectations were low, actual enjoyment was high", return_tensors="pt", padding=True)

result = ort_model(**inp)
assert ort_model.providers == ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]
```

### Log

```
2022-12-21 09:40:48.208932701 [W:onnxruntime:Default, tensorrt_execution_provider.h:60 log] [2022-12-21 09:40:48 WARNING] external/onnx-tensorrt/onnx2trt_utils.cpp:367: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32
```

### Expected behavior

Clamping should be done before the inference.

Will work on a `convert_int64_to_int32()` function to solve it.


Discussion: @JingyaHuang the warning

```
2022-12-21 11:03:15.553857281 [W:onnxruntime:Default, tensorrt_execution_provider.h:60 log] [2022-12-21 10:03:15 WARNING] external/onnx-tensorrt/onnx2trt_utils.cpp:367: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
```

is raised at the model initialization, so I'm not sure what we can do more? The inputs themselves being int64 from the tokenizer.
Have a working prototype for this. `trtexec` comes very handy to test out. The warnings come due to `Slice` operator having end indexes as the maximum representable int64 value. Casting to `np.iinfo(np.int32).max` + `astype(np.int32)` removes the warnings.

![image](https://user-images.githubusercontent.com/9808326/210064682-4e66d481-f480-4cee-8d26-357ce701e4ba.png) 



Fix (patch):
diff --git a/optimum/onnx/__init__.py b/optimum/onnx/__init__.py
index 4a0d7fb934..a543965d0d 100644
--- a/optimum/onnx/__init__.py
+++ b/optimum/onnx/__init__.py
@@ -11,4 +11,9 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-from .graph_transformations import merge_decoders
+from .graph_transformations import (
+    cast_slice_nodes_inputs_to_int32,
+    merge_decoders,
+    remove_duplicate_weights,
+    replace_atenops_to_gather,
+)
diff --git a/optimum/onnx/graph_transformations.py b/optimum/onnx/graph_transformations.py
index 8ee2b42495..351e40b13b 100644
--- a/optimum/onnx/graph_transformations.py
+++ b/optimum/onnx/graph_transformations.py
@@ -11,62 +11,29 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+
 import copy
-from collections import defaultdict
-from typing import DefaultDict, Dict, List, Set, Tuple
 
 import onnx
-from onnx import ModelProto, ValueInfoProto
-
-
-def _find_duplicate_weights(model) -> DefaultDict[Tuple[int, bytes], Set[str]]:
-    return _find_duplicate_initializers(model.graph.initializer)
-
-
-def _find_duplicate_initializers(initializers) -> DefaultDict[Tuple[int, bytes], Set[str]]:
-    duplicates = defaultdict(set)
-    for initializer in initializers:
-        tensor_dims = tuple(getattr(initializer, "dims"))
-        for data_attr in ["raw_data", "int32_data", "int64_data", "uint64_data", "float_data", "double_data"]:
-            tensor_data = getattr(initializer, data_attr)
-            if tensor_data:
-                tensor_data = tuple(tensor_data)
-                break
-        duplicates[(initializer.data_type, tensor_data, tensor_dims)].add(initializer.name)
-    return duplicates
-
+from onnx import ModelProto
 
-def _create_name_sharing_dict(
-    duplicate_weights: DefaultDict[Tuple[int, bytes], Set[str]], suffix: str = None
-) -> Dict[str, str]:
-    def _create_name_sharing_dict_for_duplicates(duplicates: Set[str]) -> Dict[str, str]:
-        common_name = duplicates.pop()
-        duplicates.add(common_name)
-        if suffix:
-            return {k: f"{common_name}_{suffix}" for k in duplicates}
-        else:
-            return {k: common_name for k in duplicates}
+from ..utils import logging
 
-    name_sharing_dict = {}
-    for duplicates in duplicate_weights.values():
-        name_sharing_dict.update(_create_name_sharing_dict_for_duplicates(duplicates))
-    return name_sharing_dict
 
+logger = logging.get_logger()
 
-def _replace_input_names(model: ModelProto, name_sharing_dict: Dict[str, str]):
-    for node in model.graph.node:
-        for i in range(len(node.input)):
-            node.input[i] = name_sharing_dict.get(node.input[i], node.input[i])
-
-
-def _remove_redundant_initializers(model: ModelProto, name_sharing_dict: Dict[str, str]):
-    to_pop = []
-    for idx, initializer in enumerate(model.graph.initializer):
-        if initializer.name != name_sharing_dict[initializer.name]:
-            to_pop.append(idx)
 
-    for idx in sorted(to_pop, reverse=True):
-        model.graph.initializer.pop(idx)
+from .transformations_utils import (
+    _create_name_sharing_dict,
+    _deduplicated_cross_model_initializers,
+    _find_duplicate_weights,
+    _get_all_inputs,
+    _get_onnx_opset,
+    _remove_redundant_initializers,
+    _replace_input_names,
+    _unify_onnx_outputs,
+    cast_int64_tensorproto_to_int32,
+)
 
 
 def remove_duplicate_weights(model: ModelProto, inplace: bool = False) -> ModelProto:
@@ -92,7 +59,7 @@ def remove_duplicate_weights(model: ModelProto, inplace: bool = False) -> ModelP
     return model
 
 
-def replace_atenops_to_gather(model: ModelProto):
+def replace_atenops_to_gather(model: ModelProto) -> ModelProto:
     """
     Replaces broken ATenOp nodes back to Gather nodes.
 
@@ -122,101 +89,6 @@ def replace_atenops_to_gather(model: ModelProto):
     return model
 
 
-def _infer_output_shape(output: ValueInfoProto):
-    output_shape = []
-    for dim in output.type.tensor_type.shape.dim:
-        if getattr(dim, "dim_param"):
-            output_shape.append(getattr(dim, "dim_param"))
-        elif getattr(dim, "dim_value"):
-            output_shape.append(getattr(dim, "dim_value"))
-        else:
-            raise ValueError(f"Cannot find `dim_param` nor `dim_value` in the output dimension info.")
-
-    return output_shape
-
-
-def _check_num_outputs(model1: ModelProto, model2: ModelProto):
-    if not len(model1.graph.output) == len(model2.graph.output):
-        raise ValueError(
-            f"Two model protos need to have the same outputs. But one has {len(model1.graph.output)} "
-            f"outputs while the other has {len(model2.graph.output)} outputs."
-        )
-
-
-def _unify_onnx_outputs(model1: ModelProto, model2: ModelProto):
-    """
-    Unifies the outputs of two ONNX model protos. The outputs of model1 will be replaced by outputs of model2.
-    According to the rules of "If" op, two subgraphs must have the same number of outputs.
-    """
-    _check_num_outputs(model1, model2)
-
-    for idx in range(len(model1.graph.output)):
-        model_output_1 = model1.graph.output[idx]
-        model_output_2 = model2.graph.output[idx]
-        if not model_output_1 == model_output_2:
-            if not (
-                model_output_1.name == model_output_2.name
-                and model_output_1.type.tensor_type.elem_type == model_output_2.type.tensor_type.elem_type
-            ):
-                raise ValueError(
-                    f"Can not match {model_output_1.name} with {model_output_2.name}. Make sure your"
-                    f" model protos have same outputs, have same data types and are in the same order."
-                )
-            model1.graph.output.remove(model_output_1)
-
-            new_output = onnx.helper.make_tensor_value_info(
-                model_output_2.name,
-                model_output_2.type.tensor_type.elem_type,
-                _infer_output_shape(model_output_2),
-            )
-            model1.graph.output.insert(idx, new_output)
-
-    if not all(
-        model_output_1 == model_output_2
-        for model_output_1, model_output_2 in zip(model1.graph.output, model2.graph.output)
-    ):
-        raise RuntimeError(f"Failed to unify outputs of given ONNX model protos.")
-
-
-def _get_all_inputs(model_list: List[ModelProto]):
-    inputs = []
-    input_names = set()
-    for model in model_list:
-        for input in model.graph.input:
-            if input.name not in input_names:
-                input_names.add(input.name)
-                inputs.append(input)
-    return inputs
-
-
-def _get_onnx_opset(model: ModelProto):
-    opset_import = model.opset_import[0]
-    return getattr(opset_import, "version")
-
-
-def _deduplicated_cross_model_initializers(models: List[ModelProto], suffix: str = None):
-
-    all_initializers = []
-    for model in models:
-        all_initializers += list(model.graph.initializer)
-
-    duplicates = _find_duplicate_initializers(all_initializers)
-    name_sharing_dict = _create_name_sharing_dict(duplicates, suffix=suffix)
-    for model in models:
-        _replace_input_names(model, name_sharing_dict)
-
-    deduplicated_initializers = []
-    deduplicated_name = set()
-
-    for initializer in all_initializers:
-        if name_sharing_dict[initializer.name] not in deduplicated_name:
-            deduplicated_name.add(name_sharing_dict[initializer.name])
-            initializer.name = name_sharing_dict[initializer.name]
-            deduplicated_initializers.append(initializer)
-
-    return deduplicated_initializers
-
-
 def merge_decoders(
     decoder: ModelProto,
     decoder_with_past: ModelProto,
@@ -296,3 +168,41 @@ def merge_decoders(
     onnx.checker.check_model(merged_model)
 
     return merged_model
+
+
+def cast_slice_nodes_inputs_to_int32(model: ModelProto) -> ModelProto:
+    """
+    Convert node inputs of `Slice` nodes from int64 to int32, casting the out of range values.
+
+    The constant node inputs are stored in `model.graph.node`, and the sole way to check which node
+    they are consumed by is to iterate over nodes and check `node.input` for a match.
+
+    Note that constant inputs to nodes as `Squeeze`, `Unsqueeze` can not be converted to int32, as the
+    these operators explicitely expect int64 inputs according to ONNX specifications:
+    https://github.com/onnx/onnx/blob/main/docs/Operators.md
+    """
+    map_input_node = {}
+    map_node_inputs = {}
+
+    for node in model.graph.node:
+        for input_name in node.input:
+            map_input_node[input_name] = {"op_type": node.op_type, "node_name": node.name}
+        map_node_inputs[node.name] = node.input
+
+    for node in model.graph.node:
+        if (
+            node.op_type == "Constant"
+            and node.attribute[0].t.data_type == 7  # int64
+            and f"{node.name}_output_0" in map_input_node
+            and map_input_node[node.name + "_output_0"]["op_type"] == "Slice"
+        ):
+            logger.debug(f"Converting {node.name} to int32")
+
+            # `Slice` node is homogeneous (requires parameters of same type), hence cast to int32 only if all of its inputs are constants
+            # refer to onnx/defs/schema.h
+            cast = all(
+                "Constant" in inp for inp in map_node_inputs[map_input_node[node.name + "_output_0"]["node_name"]][1:]
+            )
+            cast_int64_tensorproto_to_int32(node.attribute[0].t, cast=cast)
+
+    return model
diff --git a/optimum/onnx/transformations_utils.py b/optimum/onnx/transformations_utils.py
new file mode 100644
index 0000000000..eeaa274cf4
--- /dev/null
+++ b/optimum/onnx/transformations_utils.py
@@ -0,0 +1,224 @@
+#  Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from collections import defaultdict
+from typing import DefaultDict, Dict, List, Set, Tuple
+
+import numpy as np
+
+import onnx
+from onnx import ModelProto, ValueInfoProto, numpy_helper
+
+
+def _find_duplicate_weights(model) -> DefaultDict[Tuple[int, bytes], Set[str]]:
+    """
+    TODO: short documentation.
+    """
+    return _find_duplicate_initializers(model.graph.initializer)
+
+
+def _find_duplicate_initializers(initializers) -> DefaultDict[Tuple[int, bytes], Set[str]]:
+    """
+    TODO: short documentation.
+    """
+    duplicates = defaultdict(set)
+    for initializer in initializers:
+        tensor_dims = tuple(getattr(initializer, "dims"))
+        for data_attr in ["raw_data", "int32_data", "int64_data", "uint64_data", "float_data", "double_data"]:
+            tensor_data = getattr(initializer, data_attr)
+            if tensor_data:
+                tensor_data = tuple(tensor_data)
+                break
+        duplicates[(initializer.data_type, tensor_data, tensor_dims)].add(initializer.name)
+    return duplicates
+
+
+def _create_name_sharing_dict(
+    duplicate_weights: DefaultDict[Tuple[int, bytes], Set[str]], suffix: str = None
+) -> Dict[str, str]:
+    """
+    TODO: short documentation.
+    """
+
+    def _create_name_sharing_dict_for_duplicates(duplicates: Set[str]) -> Dict[str, str]:
+        common_name = duplicates.pop()
+        duplicates.add(common_name)
+        if suffix:
+            return {k: f"{common_name}_{suffix}" for k in duplicates}
+        else:
+            return {k: common_name for k in duplicates}
+
+    name_sharing_dict = {}
+    for duplicates in duplicate_weights.values():
+        name_sharing_dict.update(_create_name_sharing_dict_for_duplicates(duplicates))
+    return name_sharing_dict
+
+
+def _replace_input_names(model: ModelProto, name_sharing_dict: Dict[str, str]):
+    """
+    TODO: short documentation.
+    """
+    for node in model.graph.node:
+        for i in range(len(node.input)):
+            node.input[i] = name_sharing_dict.get(node.input[i], node.input[i])
+
+
+def _remove_redundant_initializers(model: ModelProto, name_sharing_dict: Dict[str, str]):
+    """
+    TODO: short documentation.
+    """
+    to_pop = []
+    for idx, initializer in enumerate(model.graph.initializer):
+        if initializer.name != name_sharing_dict[initializer.name]:
+            to_pop.append(idx)
+
+    for idx in sorted(to_pop, reverse=True):
+        model.graph.initializer.pop(idx)
+
+
+def _infer_output_shape(output: ValueInfoProto):
+    """
+    TODO: short documentation.
+    """
+    output_shape = []
+    for dim in output.type.tensor_type.shape.dim:
+        if getattr(dim, "dim_param"):
+            output_shape.append(getattr(dim, "dim_param"))
+        elif getattr(dim, "dim_value"):
+            output_shape.append(getattr(dim, "dim_value"))
+        else:
+            raise ValueError(f"Cannot find `dim_param` nor `dim_value` in the output dimension info.")
+
+    return output_shape
+
+
+def _check_num_outputs(model1: ModelProto, model2: ModelProto):
+    """
+    Checks that `model1` and `model2` have the same number of outputs.
+    """
+    if not len(model1.graph.output) == len(model2.graph.output):
+        raise ValueError(
+            f"Two model protos need to have the same outputs. But one has {len(model1.graph.output)} "
+            f"outputs while the other has {len(model2.graph.output)} outputs."
+        )
+
+
+def _unify_onnx_outputs(model1: ModelProto, model2: ModelProto):
+    """
+    Unifies the outputs of two ONNX model protos. The outputs of model1 will be replaced by outputs of model2.
+    According to the rules of "If" op, two subgraphs must have the same number of outputs.
+    """
+    _check_num_outputs(model1, model2)
+
+    for idx in range(len(model1.graph.output)):
+        model_output_1 = model1.graph.output[idx]
+        model_output_2 = model2.graph.output[idx]
+        if not model_output_1 == model_output_2:
+            if not (
+                model_output_1.name == model_output_2.name
+                and model_output_1.type.tensor_type.elem_type == model_output_2.type.tensor_type.elem_type
+            ):
+                raise ValueError(
+                    f"Can not match {model_output_1.name} with {model_output_2.name}. Make sure your"
+                    f" model protos have same outputs, have same data types and are in the same order."
+                )
+            model1.graph.output.remove(model_output_1)
+
+            new_output = onnx.helper.make_tensor_value_info(
+                model_output_2.name,
+                model_output_2.type.tensor_type.elem_type,
+                _infer_output_shape(model_output_2),
+            )
+            model1.graph.output.insert(idx, new_output)
+
+    if not all(
+        model_output_1 == model_output_2
+        for model_output_1, model_output_2 in zip(model1.graph.output, model2.graph.output)
+    ):
+        raise RuntimeError(f"Failed to unify outputs of given ONNX model protos.")
+
+
+def _get_all_inputs(model_list: List[ModelProto]) -> List[onnx.onnx_ml_pb2.ValueInfoProto]:
+    """
+    Returns all the inputs to all the models in `model_list`, in a single list.
+    """
+    inputs = []
+    input_names = set()
+    for model in model_list:
+        for input in model.graph.input:
+            if input.name not in input_names:
+                input_names.add(input.name)
+                inputs.append(input)
+    return inputs
+
+
+def _get_onnx_opset(model: ModelProto):
+    """
+    Returns the ONNX opset version used to generate `model`.
+    """
+    opset_import = model.opset_import[0]
+    return getattr(opset_import, "version")
+
+
+def _deduplicated_cross_model_initializers(models: List[ModelProto], suffix: str = None):
+    """
+    TODO: short documentation.
+    """
+    all_initializers = []
+    for model in models:
+        all_initializers += list(model.graph.initializer)
+
+    duplicates = _find_duplicate_initializers(all_initializers)
+    name_sharing_dict = _create_name_sharing_dict(duplicates, suffix=suffix)
+    for model in models:
+        _replace_input_names(model, name_sharing_dict)
+
+    deduplicated_initializers = []
+    deduplicated_name = set()
+
+    for initializer in all_initializers:
+        if name_sharing_dict[initializer.name] not in deduplicated_name:
+            deduplicated_name.add(name_sharing_dict[initializer.name])
+            initializer.name = name_sharing_dict[initializer.name]
+            deduplicated_initializers.append(initializer)
+
+    return deduplicated_initializers
+
+
+def cast_int64_tensorproto_to_int32(initializer: onnx.TensorProto, cast: bool = False):
+    """
+    Casts in place the input TensorProto data to int32. Its data is assumed to be of type int64,
+    and in case some values are out of range, they are cast to the min/max representable
+    value in int32.
+    """
+    original_name = initializer.name
+    array = np.copy(numpy_helper.to_array(initializer))
+
+    if not array.dtype == np.int64:
+        raise TypeError(
+            "Expecting a `TensorProto` of type `int64` (represented as `7` in onnx.TensorProto) in the function int64_tensorproto_to_int32, but got {array.dtype}."
+        )
+
+    array[array > np.iinfo(np.int32).max] = np.iinfo(np.int32).max
+    array[array < np.iinfo(np.int32).min] = np.iinfo(np.int32).min
+
+    # the following line notably avoids the cast overhead in `convertOnnxWeights` in onnx-tensorrt
+    if cast:
+        array = array.astype(np.int32)
+    array.setflags(write=0)
+
+    tensor = numpy_helper.from_array(array)
+
+    initializer.CopyFrom(tensor)
+    initializer.name = original_name