Migrate onnxscript converter to use onnx ir (#2706)

Migrate onnxscript converter to use onnx ir.

---------

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

onnxscript/_internal/autocast.py

@@ -189,30 +189,27 @@ def get_type_info(x):
 def static_cast_inputs(
     converter_: converter.Converter,
     op_schema: Optional[OpSchema],
-    args: Sequence[Optional[converter.Variable]],
+    args: Sequence[Optional[ir.Value]],
 ) -> tuple[str, ...]:
     """Used for autocast during script-translation.
     This is meant to transform expressions like "Add(X, 1)" to "Add(X, CastLike(1, X))"
     Polymorphic constants (like 0 and 1) are cast to the type of other operands as needed.
     """
 
-    def get_type_info(x: Optional[converter.Variable]) -> Optional[converter.Variable]:
+    def get_type_info(x: Optional[ir.Value]) -> Optional[ir.Value]:
         """Returns x back if x can serve as the target-type for a cast (as the second
         argument of CastLike) and None otherwise. In the expression "Add(X, 1), 1 is
         castable, while X can serve as the target-type.
         """
-        return None if x is None or x.is_castable else x
+        return None if x is None or converter_.is_castable(x.name) else x
 
-    def cast_like(
-        x: Optional[converter.Variable], y: Optional[converter.Variable]
-    ) -> Optional[str]:
+    def cast_like(x: Optional[ir.Value], y: Optional[ir.Value]) -> Optional[str]:
         if x is None:
             return None
-        if x.is_castable and y is not None:
+        if converter_.is_castable(x.name) and y is not None:
             # Polymorphic constant x is cast to the type of y:
             x_cast = converter_.generate_unique_name(f"{x.name}_cast")
-            converter_.emit([x_cast], "CastLike", [x.name, y.name])
-            return x_cast
-        return x.name
+            return converter_.emit1([x_cast], "CastLike", [x, y])
+        return x
 
     return cast_inputs(get_type_info, cast_like, op_schema, args)

onnxscript/_internal/utils.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import onnx
+import onnx_ir as ir
 
 from onnxscript import tensor
 
@@ -87,6 +88,41 @@ def value_to_type_proto(val):
     raise ValueError(f"Value of type {type(val)} is invalid as an ONNX input/output.")
 
 
+def value_to_type(val):
+    """Return an ir.Value representation of a python-value."""
+    if isinstance(val, (np.ndarray, tensor.Tensor)):
+        elem_type = onnx.helper.np_dtype_to_tensor_dtype(val.dtype)  # noqa: TID251
+        shape = val.shape
+        return (ir.TensorType(elem_type), shape)
+    elif isinstance(val, int):
+        elem_type = onnx.TensorProto.INT32
+        shape = []
+        return (ir.TensorType(elem_type), shape)
+    elif isinstance(val, (float, np.float32)):
+        elem_type = onnx.TensorProto.FLOAT
+        shape = []
+        return (ir.TensorType(elem_type), shape)
+    elif isinstance(val, list):
+        if len(val) > 0:
+            type, shape = value_to_type(val[0])
+            return ir.SequenceType(type), shape
+        # Edge-case. Cannot determine a suitable ONNX type for an empty list.
+        # Should be using a typed-value instead.
+        # Treated as a sequence of tensors of float-type.
+        return ir.SequenceType(ir.TensorType(onnx.TensorProto.FLOAT)), None
+    if isinstance(val, numbers.Number):
+        nparray = np.array(val)
+        elem_type = onnx.helper.np_dtype_to_tensor_dtype(nparray.dtype)  # noqa: TID251
+        return ir.TensorType(elem_type), []
+    raise ValueError(f"Value of type {type(val)} is invalid as an ONNX input/output.")
+
+
+def value_to_ir_value(name: str, val) -> ir.Value:
+    """Return an ir.Value representation of a python-value."""
+    type, shape = value_to_type(val)
+    return ir.Value(name=name, type=type, shape=shape)
+
+
 def values_to_value_infos(name_values):
     """Create a list of ValueInfoProto from a list of (name, value) pairs,
     skipping any None values.
@@ -96,3 +132,10 @@ def values_to_value_infos(name_values):
         for (name, val) in name_values
         if val is not None
     ]
+
+
+def values_to_ir_values(name_values):
+    """Create a list of ir.Value from a list of (name, value) pairs,
+    skipping any None values.
+    """
+    return [value_to_ir_value(name, val) for (name, val) in name_values if val is not None]