Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fixes s390x byteswapping issues #6183

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fixes s390x byteswapping issues #6183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f365e47
Fix conversion of unrepresentable dtypes in s390x
tehbone Jun 14, 2024
d027563
Update testcases to byteswap raw data for make_tensor on big endian s…
tehbone Jun 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions onnx/helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1563,6 +1563,16 @@
"""
return mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype

def tensor_dtype_to_storage_np_dtype(tensor_dtype: int) -> int:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need this to be public?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question. I guess it could be argued that if a user were creating a raw tensor then having that information could help in determining whether or not the storage has the correct number of elements.

"""Convert a TensorProto's data_type to corresponding data_type for raw storage.

Check warning

Code scanning / lintrunner

RUFF/W293 Warning

Blank line contains whitespace.
See https://docs.astral.sh/ruff/rules/blank-line-with-whitespace

Check warning

Code scanning / lintrunner

EDITORCONFIG-CHECKER/editorconfig Warning

Trailing whitespace
Args:
tensor_dtype: TensorProto's data_type

Check warning

Code scanning / lintrunner

RUFF/W293 Warning

Blank line contains whitespace.
See https://docs.astral.sh/ruff/rules/blank-line-with-whitespace

Check warning

Code scanning / lintrunner

EDITORCONFIG-CHECKER/editorconfig Warning

Trailing whitespace
Returns:
numpy's data_type for the raw storage
"""
return mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_np_dtype

Check failure

Code scanning / lintrunner

MYPY/return-value Error

Incompatible return value type (got "dtype[Any]", expected "int") To disable, use # type: ignore[return-value]

def tensor_dtype_to_string(tensor_dtype: int) -> str:
"""Get the name of given TensorProto's data_type.
Expand Down
48 changes: 25 additions & 23 deletions onnx/mapping.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,80 +13,82 @@

class TensorDtypeMap(NamedTuple):
np_dtype: np.dtype
storage_np_dtype: np.dtype
storage_dtype: int
name: str


# tensor_dtype: (numpy type, storage type, string name)
TENSOR_TYPE_MAP = {
int(TensorProto.FLOAT): TensorDtypeMap(
np.dtype("float32"), int(TensorProto.FLOAT), "TensorProto.FLOAT"
np.dtype("float32"), np.dtype("float32"), int(TensorProto.FLOAT), "TensorProto.FLOAT"
),
int(TensorProto.UINT8): TensorDtypeMap(
np.dtype("uint8"), int(TensorProto.INT32), "TensorProto.UINT8"
np.dtype("uint8"), np.dtype("uint8"), int(TensorProto.INT32), "TensorProto.UINT8"
),
int(TensorProto.INT8): TensorDtypeMap(
np.dtype("int8"), int(TensorProto.INT32), "TensorProto.INT8"
np.dtype("int8"), np.dtype("int8"), int(TensorProto.INT32), "TensorProto.INT8"
),
int(TensorProto.UINT16): TensorDtypeMap(
np.dtype("uint16"), int(TensorProto.INT32), "TensorProto.UINT16"
np.dtype("uint16"), np.dtype("uint16"), int(TensorProto.INT32), "TensorProto.UINT16"
),
int(TensorProto.INT16): TensorDtypeMap(
np.dtype("int16"), int(TensorProto.INT32), "TensorProto.INT16"
np.dtype("int16"), np.dtype("int16"), int(TensorProto.INT32), "TensorProto.INT16"
),
int(TensorProto.INT32): TensorDtypeMap(
np.dtype("int32"), int(TensorProto.INT32), "TensorProto.INT32"
np.dtype("int32"), np.dtype("int32"), int(TensorProto.INT32), "TensorProto.INT32"
),
int(TensorProto.INT64): TensorDtypeMap(
np.dtype("int64"), int(TensorProto.INT64), "TensorProto.INT64"
np.dtype("int64"), np.dtype("int64"), int(TensorProto.INT64), "TensorProto.INT64"
),
int(TensorProto.BOOL): TensorDtypeMap(
np.dtype("bool"), int(TensorProto.INT32), "TensorProto.BOOL"
np.dtype("bool"), np.dtype("bool"), int(TensorProto.INT32), "TensorProto.BOOL"
),
int(TensorProto.FLOAT16): TensorDtypeMap(
np.dtype("float16"), int(TensorProto.UINT16), "TensorProto.FLOAT16"
np.dtype("float16"), np.dtype("float16"), int(TensorProto.UINT16), "TensorProto.FLOAT16"
),
# Native numpy does not support bfloat16 so now use float32.
# Native numpy does not support bfloat16 so now use float32. It is saved as a uin16 field.
int(TensorProto.BFLOAT16): TensorDtypeMap(
np.dtype("float32"), int(TensorProto.UINT16), "TensorProto.BFLOAT16"
np.dtype("float32"), np.dtype("uint16"), int(TensorProto.UINT16), "TensorProto.BFLOAT16"
),
int(TensorProto.DOUBLE): TensorDtypeMap(
np.dtype("float64"), int(TensorProto.DOUBLE), "TensorProto.DOUBLE"
np.dtype("float64"), np.dtype("float64"), int(TensorProto.DOUBLE), "TensorProto.DOUBLE"
),
int(TensorProto.COMPLEX64): TensorDtypeMap(
np.dtype("complex64"), int(TensorProto.FLOAT), "TensorProto.COMPLEX64"
np.dtype("complex64"), np.dtype("complex64"), int(TensorProto.FLOAT), "TensorProto.COMPLEX64"
),
int(TensorProto.COMPLEX128): TensorDtypeMap(
np.dtype("complex128"), int(TensorProto.DOUBLE), "TensorProto.COMPLEX128"
np.dtype("complex128"), np.dtype("complex128"), int(TensorProto.DOUBLE), "TensorProto.COMPLEX128"
),
int(TensorProto.UINT32): TensorDtypeMap(
np.dtype("uint32"), int(TensorProto.UINT32), "TensorProto.UINT32"
np.dtype("uint32"), np.dtype("uint32"), int(TensorProto.UINT32), "TensorProto.UINT32"
),
int(TensorProto.UINT64): TensorDtypeMap(
np.dtype("uint64"), int(TensorProto.UINT64), "TensorProto.UINT64"
np.dtype("uint64"), np.dtype("uint64"), int(TensorProto.UINT64), "TensorProto.UINT64"
),
int(TensorProto.STRING): TensorDtypeMap(
np.dtype("object"), int(TensorProto.STRING), "TensorProto.STRING"
np.dtype("object"), np.dtype("object"), int(TensorProto.STRING), "TensorProto.STRING"
),
# Native numpy does not support float8 types, so now use float32 for these types.
# As raw data, it is viewed as a collection of bytes.
int(TensorProto.FLOAT8E4M3FN): TensorDtypeMap(
np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E4M3FN"
np.dtype("float32"), np.dtype("uint8"), int(TensorProto.UINT8), "TensorProto.FLOAT8E4M3FN"
),
int(TensorProto.FLOAT8E4M3FNUZ): TensorDtypeMap(
np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E4M3FNUZ"
np.dtype("float32"), np.dtype("uint8"), int(TensorProto.UINT8), "TensorProto.FLOAT8E4M3FNUZ"
),
int(TensorProto.FLOAT8E5M2): TensorDtypeMap(
np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E5M2"
np.dtype("float32"), np.dtype("uint8"), int(TensorProto.UINT8), "TensorProto.FLOAT8E5M2"
),
int(TensorProto.FLOAT8E5M2FNUZ): TensorDtypeMap(
np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E5M2FNUZ"
np.dtype("float32"), np.dtype("uint8"), int(TensorProto.UINT8), "TensorProto.FLOAT8E5M2FNUZ"
),
# Native numpy does not support uint4/int4 so now use uint8/int8 for these types.
int(TensorProto.UINT4): TensorDtypeMap(
np.dtype("uint8"), int(TensorProto.INT32), "TensorProto.UINT4"
np.dtype("uint8"), np.dtype("uint8"), int(TensorProto.INT32), "TensorProto.UINT4"
),
int(TensorProto.INT4): TensorDtypeMap(
np.dtype("int8"), int(TensorProto.INT32), "TensorProto.INT4"
np.dtype("int8"), np.dtype("uint8"), int(TensorProto.INT32), "TensorProto.INT4"
),
}

Expand Down
11 changes: 7 additions & 4 deletions onnx/numpy_helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,9 +235,7 @@

tensor_dtype = tensor.data_type
np_dtype = helper.tensor_dtype_to_np_dtype(tensor_dtype)
storage_np_dtype = helper.tensor_dtype_to_np_dtype(
helper.tensor_dtype_to_storage_tensor_dtype(tensor_dtype)
)

storage_field = helper.tensor_dtype_to_field(tensor_dtype)
dims = tensor.dims

Expand All @@ -251,11 +249,12 @@
load_external_data_for_tensor(tensor, base_dir)

if tensor.HasField("raw_data"):
storage_np_dtype = helper.tensor_dtype_to_storage_np_dtype(tensor_dtype)
# Raw_bytes support: using frombuffer.
raw_data = tensor.raw_data
if sys.byteorder == "big":
# Convert endian from little to big
raw_data = np.frombuffer(raw_data, dtype=np_dtype).byteswap().tobytes()
raw_data = np.frombuffer(raw_data, dtype=storage_np_dtype).byteswap().tobytes()

Check warning on line 257 in onnx/numpy_helper.py

View check run for this annotation

Codecov / codecov/patch

onnx/numpy_helper.py#L257 

Added line #L257 was not covered by tests

Check failure

Code scanning / lintrunner

MYPY/call-overload Error

No overload variant of "frombuffer" matches argument types "bytes", "int" To disable, use # type: ignore[call-overload]

# manually convert bf16 since there's no numpy support
if tensor_dtype == TensorProto.BFLOAT16:
Expand Down Expand Up @@ -288,6 +287,10 @@

return np.frombuffer(raw_data, dtype=np_dtype).reshape(dims) # type: ignore[no-any-return]

# dtype for raw data is determined by the storage_dtype
storage_np_dtype = helper.tensor_dtype_to_np_dtype(

Check failure

Code scanning / lintrunner

MYPY/assignment Error

Incompatible types in assignment (expression has type "dtype[Any]", variable has type "int") To disable, use # type: ignore[assignment]
helper.tensor_dtype_to_storage_tensor_dtype(tensor_dtype)
)
# float16 is stored as int32 (uint16 type); Need view to get the original value
if tensor_dtype == TensorProto.FLOAT16:
return (
Expand Down
9 changes: 7 additions & 2 deletions onnx/test/helper_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import math
import random
import struct
import sys
import unittest
from typing import Any

Expand All @@ -31,6 +32,10 @@
from onnx.reference.op_run import to_array_extended


def convert(data: np.ndarray) -> np.ndarray:
return data.byteswap() if sys.byteorder == 'big' else data


class TestHelperAttributeFunctions(unittest.TestCase):
def test_attr_float(self) -> None:
# float
Expand Down Expand Up @@ -562,7 +567,7 @@ def truncate(x):
return x >> 16

values_as_ints = np_array.astype(np.float32).view(np.uint32).flatten()
packed_values = truncate(values_as_ints).astype(np.uint16).tobytes()
packed_values = convert(truncate(values_as_ints).astype(np.uint16)).tobytes()
tensor = helper.make_tensor(
name="test",
data_type=TensorProto.BFLOAT16,
Expand Down Expand Up @@ -951,7 +956,7 @@ def test_make_tensor_raw(tensor_dtype: int) -> None:
name="test",
data_type=tensor_dtype,
dims=np_array.shape,
vals=np_array.tobytes(),
vals=convert(np_array).tobytes(),
raw=True,
)
np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
Expand Down
9 changes: 7 additions & 2 deletions onnx/test/test_external_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import os
import pathlib
import tempfile
import sys
import unittest
import uuid
from typing import Any
Expand All @@ -26,6 +27,10 @@
from onnx.numpy_helper import from_array, to_array


def convert(data: np.ndarray) -> np.ndarray:
return data.byteswap() if sys.byteorder == 'big' else data


class TestLoadExternalDataBase(unittest.TestCase):
"""Base class for testing external data related behaviors.

Expand Down Expand Up @@ -584,7 +589,7 @@ def create_test_model(self) -> ModelProto:
name="X",
data_type=TensorProto.FLOAT,
dims=self.large_data.shape,
vals=self.large_data.tobytes(),
vals=convert(self.large_data).tobytes(),
raw=True,
)

Expand All @@ -593,7 +598,7 @@ def create_test_model(self) -> ModelProto:
name="Shape",
data_type=TensorProto.INT64,
dims=shape_data.shape,
vals=shape_data.tobytes(),
vals=convert(shape_data).tobytes(),
raw=True,
)
C = helper.make_tensor_value_info("C", TensorProto.INT64, self.small_data)
Expand Down
Loading