Cast fp16 to fp32 in tolist, remove FP16MemoryView (#254)

Author: yizhuoz004

tripy/tests/common/test_utils.py

@@ -29,7 +29,6 @@
 from tripy.common.exception import TripyException
 from tripy.common.utils import (
     convert_list_to_array,
-    Float16MemoryView,
     get_element_type,
 )
 
@@ -50,43 +49,11 @@ def test_get_element_type():
     "values, dtype, expected",
     [
         ([True, False, True], tripy.common.datatype.bool, b"\x01\x00\x01"),
-        ([1, 2, 3], tripy.common.datatype.int8, b"\x01\x02\x03"),
         ([100000, 200000], tripy.common.datatype.int32, b"\xa0\x86\x01\x00@\x0d\x03\x00"),
         ([1, 2], tripy.common.datatype.int64, b"\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00"),
-        ([1.0, 2.0], tripy.common.datatype.float16, b"\x00<\x00@"),
         ([1.0, 2.0], tripy.common.datatype.float32, b"\x00\x00\x80?\x00\x00\x00@"),
     ],
 )
-def convert_list_to_array(values, dtype, expected):
+def test_convert_list_to_array(values, dtype, expected):
     result = convert_list_to_array(values, dtype)
-    assert result == expected
-
-
-def test_float16_memoryview():
-    memview = Float16MemoryView(bytearray(struct.pack("5e", 1.5, 2.5, 3.5, 4.5, 5.5)))
-    assert memview.itemsize == 2
-    assert memview.format == "e"
-    len(memview) == 5
-    assert memview[0] == pytest.approx(1.5)
-    assert memview[2] == pytest.approx(3.5)
-    assert memview[1:4] == pytest.approx([2.5, 3.5, 4.5])
-    expected = [1.5, 2.5, 3.5, 4.5, 5.5]
-    assert memview.tolist() == pytest.approx(expected)
-
-    # Largest representable value in float16
-    large_value = 65504.0
-    buffer = struct.pack("e", large_value)
-    mv = Float16MemoryView(buffer)
-    assert mv[0] == pytest.approx(large_value)
-
-    # Smallest positive normal number for float16
-    small_value = 6.1035e-5
-    buffer = struct.pack("e", small_value)
-    mv = Float16MemoryView(buffer)
-    assert mv[0] == pytest.approx(small_value, rel=1e-3)
-
-    # Negative value
-    negative_value = -42.5
-    buffer = struct.pack("e", negative_value)
-    mv = Float16MemoryView(buffer)
-    assert mv[0] == pytest.approx(negative_value)
+    assert result.tobytes() == expected

tripy/tests/frontend/trace/test_trace.py

@@ -133,8 +133,8 @@ def test_multiple_outputs(self):
             str(trace)
             == dedent(
                 """
-                a = storage(data=[1.0000], shape=(1,), dtype=float32, device=gpu:0)
-                b = storage(data=[1.0000], shape=(1,), dtype=float32, device=gpu:0)
+                a = storage(shape=(1,), dtype=float32, device=gpu:0)
+                b = storage(shape=(1,), dtype=float32, device=gpu:0)
                 c = a + b
                 d = c + c
                 outputs:
@@ -192,7 +192,7 @@ def test_const_and_input(self):
                 """
                 inputs:
                     a: [rank=(1), shape=((1,)), dtype=(float32), loc=(gpu:0)]
-                b = storage(data=[1.0000], shape=(1,), dtype=float32, device=gpu:0)
+                b = storage(shape=(1,), dtype=float32, device=gpu:0)
                 c = a + b
                 outputs:
                     c: [rank=(1), dtype=(float32), loc=(gpu:0)]

tripy/tripy/backend/mlir/memref.py

@@ -77,50 +77,4 @@ def tolist(memref):
     memref_value = memref
     if memref.address_space == runtime.PointerType.device:
         memref_value = mlir_utils.MLIRRuntimeClient().copy_to_host(device_memref=memref)
-    try:
-        return memoryview(memref_value).tolist()
-    except NotImplementedError as e:
-        if "memoryview: format e not supported" in str(e):
-            assert memref_value.dtype == runtime.ScalarTypeCode.f16
-            return common_utils.Float16MemoryView(bytearray(memref_value)).tolist()
-        raise
-
-
-def pretty_print_memref(memref, threshold=1000, linewidth=10, edgeitems=3):
-    """
-    Returns a pretty-print string of memref values.
-    """
-    memref_shape = memref.shape
-
-    def _data_str(data, summarize, linewidth, edgeitems, indent=0):
-        if isinstance(data, (float, int)):
-            return str(data)
-
-        if len(data) == 0 or isinstance(data[0], (float, int)):
-            if summarize and len(data) > 2 * edgeitems:
-                data_lines = [data[:edgeitems] + [" ..."] + data[-edgeitems:]]
-            else:
-                data_lines = [data[i : i + linewidth] for i in range(0, len(data), linewidth)]
-            lines = [", ".join([f"{e:.4f}" if isinstance(e, float) else str(e) for e in line]) for line in data_lines]
-            return "[" + ("," + "\n" + " " * (indent + 1)).join(lines) + "]"
-
-        if summarize and len(data) > 2 * edgeitems:
-            slices = (
-                [_data_str(data[i], summarize, linewidth, edgeitems, indent + 1) for i in range(0, edgeitems)]
-                + ["..."]
-                + [
-                    _data_str(data[i], summarize, linewidth, edgeitems, indent + 1)
-                    for i in range(len(data) - edgeitems, len(data))
-                ]
-            )
-        else:
-            slices = [_data_str(data[i], summarize, linewidth, edgeitems, indent + 1) for i in range(0, len(data))]
-
-        tensor_str = ("," + "\n" * (max(len(memref_shape) - indent - 1, 1)) + " " * (indent + 1)).join(slices)
-        return "[" + tensor_str + "]"
-
-    numel = 1
-    for d in memref_shape:
-        numel *= d
-    summarize = numel > threshold
-    return _data_str(tolist(memref), summarize, linewidth, edgeitems)
+    return memoryview(memref_value).tolist()

tripy/tripy/common/utils.py

@@ -85,70 +85,3 @@ def is_shape_empty(shape: Sequence[int]) -> bool:
         bool: True if the shape represents an empty tensor, False otherwise.
     """
     return any(dim == 0 for dim in shape)
-
-
-class Float16MemoryView:
-    """
-    A custom memory view class for handling float16 data.
-    """
-
-    def __init__(self, buffer):
-        """
-        Initialize the Float16MemoryView with a buffer.
-
-        Args:
-            buffer (buffer): The buffer containing float16 data.
-        """
-        self.buffer = buffer
-        self.itemsize = 2  # size of float16 in bytes
-        self.format = "e"  # format character for float16
-
-    def __getitem__(self, index):
-        """
-        Get an item or a slice from the buffer.
-
-        Args:
-            index (int or slice): The index or slice to retrieve.
-
-        Returns:
-            float or list of floats: The float16 value(s) at the specified index or slice.
-        """
-        if isinstance(index, slice):
-            return [
-                self._unpack(self.buffer[i * self.itemsize : i * self.itemsize + self.itemsize])
-                for i in range(*index.indices(len(self)))
-            ]
-        else:
-            start = index * self.itemsize
-            end = start + self.itemsize
-            return self._unpack(self.buffer[start:end])
-
-    def _unpack(self, data):
-        """
-        Unpack a float16 value from bytes.
-
-        Args:
-            data (bytes): The bytes to unpack.
-
-        Returns:
-            float: The unpacked float16 value.
-        """
-        return struct.unpack(self.format, data)[0]
-
-    def __len__(self):
-        """
-        Get the number of float16 values in the buffer.
-
-        Returns:
-            int: The number of float16 values.
-        """
-        return len(self.buffer) // self.itemsize
-
-    def tolist(self):
-        """
-        Convert the buffer to a list of float16 values.
-
-        Returns:
-            list: The list of float16 values.
-        """
-        return [self[i] for i in range(len(self))]

tripy/tripy/flat_ir/flat_ir.py

@@ -352,23 +352,16 @@ def register_tensor(self, tensor: "FlatIRTensor") -> "FlatIRTensor":
         return tensor
 
     def _get_constant_key(self, op):
-        from mlir_tensorrt.runtime._mlir_libs._api import MemRefValue
-        from tripy.utils.utils import list_to_tuple, volume
+        from mlir_tensorrt.runtime.api import MemRefValue
+        from tripy.utils.utils import list_to_tuple
 
         if isinstance(op.data, MemRefValue):
-            from tripy.backend.mlir.memref import tolist
-
-            VOLUME_THRESHOLD_FOR_MEMREF = 50
-            if volume(op.data.shape) < VOLUME_THRESHOLD_FOR_MEMREF:
-                l = tolist(op.data)
-            else:
-                l = [op.data.ptr]
-            data = list_to_tuple(l if isinstance(l, List) else [l])
-        elif isinstance(op.data, int) or isinstance(op.data, float) or isinstance(op.data, bool):
-            data = list_to_tuple(
-                op.data,
-            )
+            # use data pointer as key when data is a memref,
+            # usually come from users, no need to deduplicate
+            data = (op.data.ptr,)
         else:
+            # small constants can be deduplicated
+            # when data is a list
             data = list_to_tuple(op.data)
 
         # Create a unique key for the constant based on its data and type

tripy/tripy/flat_ir/ops/constant.py

@@ -34,8 +34,7 @@ class ConstantOp(BaseFlatIROp):
     data: Union[runtime.MemRefValue, Sequence]
 
     def str_skip_fields(self) -> Set[str]:
-        data_shape = self.data.shape if isinstance(self.data, runtime.MemRefValue) else self.outputs[0].shape
-        if utils.should_omit_constant_in_str(data_shape):
+        if not isinstance(self.data, Sequence) or utils.should_omit_constant_in_str(self.outputs[0].shape):
             return {"data"}
         return set()
 

tripy/tripy/frontend/tensor.py

@@ -201,7 +201,6 @@ def eval(self) -> runtime.MemRefValue:
     def tolist(self):
         data_memref = self.eval()
         if self.dtype not in (
-            datatype.float16,
             datatype.float32,
             datatype.int8,
             datatype.int32,
@@ -217,18 +216,20 @@ def __iter__(self):
         raise TypeError("Iterating over tensors is not supported")
 
     def __repr__(self) -> str:
-        # The Evaluation required before accessing self.trace_tensor.producer attributes.
-        arr = self.eval()
-        arr_str = memref.pretty_print_memref(arr)
+        from tripy.frontend.utils import pretty_print
+
+        data_list = self.tolist()
+        data_shape = self.trace_tensor.producer.shape
+        arr_str = pretty_print(data_list, data_shape)
         indentation = ""
         sep = ""
-        if len(arr.shape) > 1 and any(dim > 1 for dim in arr.shape):
+        if len(data_shape) > 1 and any(dim > 1 for dim in data_shape):
             indentation = " " * 4
             sep = "\n"
         return (
             f"tensor({sep}"
             f"{indent(arr_str, prefix=indentation)}, {sep}"
-            f"{indent(f'dtype={self.dtype}, loc={self.device}, shape={arr.shape}', prefix=indentation)}"
+            f"{indent(f'dtype={self.dtype}, loc={self.device}, shape={data_shape}', prefix=indentation)}"
             f")"
         )
 
@@ -240,7 +241,7 @@ def __dlpack_device__(self):
         return self.eval().__dlpack_device__()
 
     def __bool__(self):
-        data = memref.tolist(self.eval())
+        data = self.tolist()
         if any(dim != 1 for dim in self.trace_tensor.producer.shape):
             raise_error(
                 "Boolean value of a Tensor with more than one value is ambiguous",

tripy/tripy/frontend/trace/ops/storage.py

@@ -73,19 +73,11 @@ def __init__(
     infer_shape_output_idxs = op_utils.ShapeOutputIdxPolicies.never_return_shape
 
     def str_skip_fields(self) -> Set[str]:
-        if utils.should_omit_constant_in_str(self.shape):
+        # skip data if i) it is a MemRefValue or ii) its volume exceeds threshold
+        if not isinstance(self.data, Sequence) or utils.should_omit_constant_in_str(self.shape):
             return {"data"}
         return set()
 
-    def __str__(self) -> str:
-        skip_fields = self.str_skip_fields()
-        args = []
-        if "data" not in skip_fields:
-            data_str = memref.pretty_print_memref(self.data) if self.has_memref else str(self.data)
-            args.append(f"data={data_str}")
-        args.extend([f"{field}={getattr(self, field)}" for field in ("shape", "dtype", "device")])
-        return f"{self.outputs[0].name} = storage({', '.join([inp.name for inp in self.inputs] + args)})"
-
     def __eq__(self, other) -> bool:
         return self.data == other.data if isinstance(other, Storage) else False
 

tripy/tripy/frontend/utils.py

@@ -633,3 +633,41 @@ def process_dim(dim: int) -> int:
         return func(*args, **kwargs)
 
     return wrapper
+
+
+def pretty_print(data_list, shape, threshold=1000, linewidth=10, edgeitems=3):
+    """
+    Returns a pretty-print string of list format data.
+    """
+    def _data_str(data, summarize, linewidth, edgeitems, indent=0):
+        if isinstance(data, (float, int)):
+            return str(data)
+
+        if len(data) == 0 or isinstance(data[0], (float, int)):
+            if summarize and len(data) > 2 * edgeitems:
+                data_lines = [data[:edgeitems] + [" ..."] + data[-edgeitems:]]
+            else:
+                data_lines = [data[i : i + linewidth] for i in range(0, len(data), linewidth)]
+            lines = [", ".join([f"{e:.4f}" if isinstance(e, float) else str(e) for e in line]) for line in data_lines]
+            return "[" + ("," + "\n" + " " * (indent + 1)).join(lines) + "]"
+
+        if summarize and len(data) > 2 * edgeitems:
+            slices = (
+                [_data_str(data[i], summarize, linewidth, edgeitems, indent + 1) for i in range(0, edgeitems)]
+                + ["..."]
+                + [
+                    _data_str(data[i], summarize, linewidth, edgeitems, indent + 1)
+                    for i in range(len(data) - edgeitems, len(data))
+                ]
+            )
+        else:
+            slices = [_data_str(data[i], summarize, linewidth, edgeitems, indent + 1) for i in range(0, len(data))]
+
+        tensor_str = ("," + "\n" * (max(len(shape) - indent - 1, 1)) + " " * (indent + 1)).join(slices)
+        return "[" + tensor_str + "]"
+
+    numel = 1
+    for d in shape:
+        numel *= d
+    summarize = numel > threshold
+    return _data_str(data_list, summarize, linewidth, edgeitems)