Add stride suppport

Author: jhalakpatel

tripy/tests/backend/test_compiler_api.py

@@ -185,8 +185,7 @@ def test_function(self):
         inp = tp.ones((2, 2), dtype=tp.float32)
         out = compiled_gelu(inp)
 
-        # TODO (#225): Replace with tp.all
-        assert cp.array_equal(cp.from_dlpack(out), cp.from_dlpack(tp.relu(inp)))
+        assert tp.equal(out, tp.relu(inp))
 
     def test_module(self):
         layernorm = tp.LayerNorm(2)

tripy/tests/flat_ir/ops/test_constant.py

@@ -30,7 +30,10 @@ def test_str(self):
 
         const = flat_ir.ops[-1]
         assert isinstance(const, ConstantOp)
-        assert str(const) == "out: [rank=(1), shape=((2,)), dtype=(float32), loc=(gpu:0)] = ConstantOp(data=[2.0, 3.0])"
+        assert (
+            str(const)
+            == "out: [rank=(1), shape=((2,)), stride=((1,)), dtype=(float32), loc=(gpu:0)] = ConstantOp(data=[2.0, 3.0])"
+        )
 
     def test_mlir(self):
         out = tp.Tensor([2, 3], dtype=tp.int32, name="out")

tripy/tests/frontend/test_stride.py

@@ -0,0 +1,62 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+#
+
+import pytest
+import re
+import torch
+
+import tripy as tp
+from tests.helper import raises
+
+
+class TestStride:
+
+    def test_non_canonical_stride(self):
+        t = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+        a = tp.Tensor(t)
+        assert a.stride == t.stride()
+
+        t = t.transpose(0, 1)
+        a = tp.Tensor(t)
+
+        assert a.stride == t.stride()
+
+    def test_lazy_stride(self):
+        a = torch.arange(12, dtype=torch.float32).reshape(4, 3).transpose(0, 1)
+        with raises(
+            tp.TripyException,
+            match=re.escape("Non-canonical strides are not supported for Tripy tensors."),
+        ):
+            print(tp.Tensor(a))
+
+    def test_compile_stride(self):
+        def twice(t):
+            return 2 * t
+
+        compiler = tp.Compiler(twice)
+
+        t = tp.Tensor(torch.arange(12, dtype=torch.float32).reshape(4, 3).transpose(0, 1))
+
+        # Create a tensor info with non-canonical stride.
+        t_info = tp.InputInfo(shape=t.shape.tolist(), dtype=t.dtype, stride=t.stride)
+        compiled_add = compiler.compile(t_info)
+
+        with raises(
+            tp.TripyException,
+            match=re.escape("Reason: InvalidArgument: Runtime stride mismatch. Expected [4, 1] but received [1, 3]"),
+        ):
+            print(compiled_add(t))

tripy/tests/frontend/trace/test_trace.py

@@ -95,8 +95,8 @@ def test_str(self):
             str(trace)
             == dedent(
                 """
-                a = storage(data=[0], shape=(1,), dtype=int32, device=gpu:0)
-                b = storage(data=[1], shape=(1,), dtype=int32, device=gpu:0)
+                a = storage(data=[0], shape=(1,), stride=(1,), dtype=int32, device=gpu:0)
+                b = storage(data=[1], shape=(1,), stride=(1,), dtype=int32, device=gpu:0)
                 c = a + b
                 outputs:
                     c: [rank=(1), dtype=(int32), loc=(gpu:0)]
@@ -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(data=[1.0000], shape=(1,), stride=(1,), dtype=float32, device=gpu:0)
+                b = storage(data=[1.0000], shape=(1,), stride=(1,), dtype=float32, device=gpu:0)
                 c = a + b
                 d = c + c
                 outputs:
@@ -168,8 +168,8 @@ def test_all_inputs(self):
             == dedent(
                 """
                 inputs:
-                    a: [rank=(1), shape=((1,)), dtype=(float32), loc=(gpu:0)]
-                    b: [rank=(1), shape=((1,)), dtype=(float32), loc=(gpu:0)]
+                    a: [rank=(1), shape=((1,)), stride=((1,)), dtype=(float32), loc=(gpu:0)]
+                    b: [rank=(1), shape=((1,)), stride=((1,)), dtype=(float32), loc=(gpu:0)]
                 c = a + b
                 outputs:
                     c: [rank=(1), dtype=(float32), loc=(gpu:0)]
@@ -191,8 +191,8 @@ def test_const_and_input(self):
             == dedent(
                 """
                 inputs:
-                    a: [rank=(1), shape=((1,)), dtype=(float32), loc=(gpu:0)]
-                b = storage(data=[1.0000], shape=(1,), dtype=float32, device=gpu:0)
+                    a: [rank=(1), shape=((1,)), stride=((1,)), dtype=(float32), loc=(gpu:0)]
+                b = storage(data=[1.0000], shape=(1,), stride=(1,), dtype=float32, device=gpu:0)
                 c = a + b
                 outputs:
                     c: [rank=(1), dtype=(float32), loc=(gpu:0)]

tripy/tests/integration/test_allclose.py

@@ -35,8 +35,8 @@ class TestAllClose:
         ],
     )
     def test_all_close_float32(self, tensor_a, tensor_b, rtol, atol):
-        np_result = torch.allclose(torch.FloatTensor(tensor_a), torch.FloatTensor(tensor_b), rtol=rtol, atol=atol)
+        torch_result = torch.allclose(torch.FloatTensor(tensor_a), torch.FloatTensor(tensor_b), rtol=rtol, atol=atol)
         tp_result = tp.allclose(
             tp.Tensor(tensor_a, dtype=tp.float32), tp.Tensor(tensor_b, dtype=tp.float32), rtol=rtol, atol=atol
         )
-        assert np_result == tp_result
+        assert torch_result == tp_result

tripy/tests/integration/test_equal.py

@@ -0,0 +1,89 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+#
+
+import pytest
+import re
+import torch
+
+import tripy as tp
+from tests import helper
+
+
+class TestEqual:
+    @pytest.mark.parametrize(
+        "tensor_a, tensor_b, dtype",
+        [
+            ([1, 2, 3], [1, 2, 3], tp.int64),
+            ([1, 2, 3], [1, 2, 4], tp.int32),
+            ([1.0, 2.0, 3.0], [1.0, 2.0, 3.0], tp.float16),
+            ([1.0, 2.0, 3.0], [1.0, 2.0, 1.0], tp.bfloat16),
+            ([1.0, 2.0, 3.0], [1.0, 2.0, 3.00001], tp.float32),
+            ([True, False, True], [True, False, True], tp.bool),
+        ],
+    )
+    def test_equal(self, tensor_a, tensor_b, dtype):
+        # Convert to torch tensors for comparison
+        torch_a = torch.tensor(tensor_a, dtype=self.torch_dtype(dtype))
+        torch_b = torch.tensor(tensor_b, dtype=self.torch_dtype(dtype))
+
+        # Convert to tripy tensors
+        tp_a = tp.Tensor(tensor_a, dtype=dtype)
+        tp_b = tp.Tensor(tensor_b, dtype=dtype)
+
+        # Compare results
+        torch_result = torch.equal(torch_a, torch_b)
+        tp_result = tp.equal(tp_a, tp_b)
+
+        assert torch_result == tp_result
+
+    @staticmethod
+    def torch_dtype(tp_dtype):
+        # Map tripy dtypes to torch dtypes
+        dtype_map = {
+            tp.float32: torch.float32,
+            tp.float16: torch.float16,
+            tp.bfloat16: torch.bfloat16,
+            tp.int32: torch.int32,
+            tp.int64: torch.int64,
+            tp.int8: torch.int8,
+            tp.bool: torch.bool,
+        }
+        return dtype_map[tp_dtype]
+
+    def test_equal_different_shapes(self):
+        a = tp.Tensor([1, 2, 3])
+        b = tp.Tensor([1, 2, 3, 4])
+
+        with helper.raises(
+            tp.TripyException,
+            match=re.escape("size of operand dimension 0 (4) is not compatible with size of result dimension 0 (3)"),
+            has_stack_info_for=[a, b],
+        ):
+            tp.equal(a, b)
+
+    def test_equal_different_dtypes(self):
+        a = tp.Tensor([1, 2, 3], dtype=tp.int32)
+        b = tp.Tensor([1, 2, 3], dtype=tp.float32)
+
+        with helper.raises(
+            tp.TripyException,
+            match=re.escape(
+                "Parameters: 'a' and 'b' must have matching data types, but got: 'int32' and 'float32' respectively."
+            ),
+            has_stack_info_for=[a, b],
+        ):
+            tp.equal(a, b)

tripy/tests/integration/test_quantize.py

@@ -118,4 +118,4 @@ def test_non_constant_scale(self):
         scale = tp.ones((4,))
         quantized = tp.quantize(input, scale, tp.int8, dim=0)
 
-        assert bool(tp.all(quantized == tp.ones((4, 4), dtype=tp.int8)))
+        assert tp.equal(quantized, tp.ones((4, 4), dtype=tp.int8))

tripy/tripy/backend/api/compiler.py

@@ -147,6 +147,7 @@ def add(a, b):
         """
 
         shapes = []
+        strides = []
         trace_input_map = {}
         input_names = set()
 
@@ -162,6 +163,7 @@ def process_arg(name, arg):
 
                 trace_input_map[name] = tensor
                 shapes.append(arg.shape_bounds)
+                strides.append(arg.stride)
                 input_names.add(name)
 
                 return tensor
@@ -196,7 +198,7 @@ def process_arg(name, arg):
 
         # Order of trace inputs also needs to match that of the compiled_arg_names
         trace_inputs = [trace_input_map[name] for name in compiled_arg_names]
-        trace = Trace(trace_outputs, trace_inputs, shapes=shapes)
+        trace = Trace(trace_outputs, trace_inputs, shapes=shapes, strides=strides)
 
         flat_ir = trace.to_flat_ir()
         mlir = flat_ir.to_mlir()

tripy/tripy/backend/api/executable.py

@@ -161,6 +161,10 @@ def add(a, b):
                                     tensor,
                                 ],
                             )
+            elif "Runtime stride mismatch" in str(err):
+                # Just raise the error for now.
+                raise raise_error(str(err))
+
             raise
 
         from tripy.utils.stack_info import StackInfo
@@ -175,10 +179,11 @@ def _get_arg_info(self, idx):
         arg = runtime.MemRefType(arg)
         arg_bound = self._executable_signature.get_arg_bound(idx)
         shape_bounds = tuple(zip(arg_bound.min(), arg_bound.max()))
+        stride = arg.strides
         if len(shape_bounds) == 0:
             # For static shape arguments, get_arg_bound returns an empty list and we fallback to arg.shape
             shape_bounds = tuple((x, x) for x in arg.shape)
-        return ArgInfo(shape_bounds, mlir_utils.convert_runtime_dtype_to_tripy_dtype(arg.dtype))
+        return ArgInfo(shape_bounds, stride, mlir_utils.convert_runtime_dtype_to_tripy_dtype(arg.dtype))
 
     def get_input_info(self) -> Sequence[ArgInfo]:
         """

tripy/tripy/backend/api/input_info.py

@@ -29,7 +29,10 @@ class InputInfo:
     """
 
     def __init__(
-        self, shape: Sequence[Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]]], dtype: "tripy.dtype"
+        self,
+        shape: Sequence[Union[int, Tuple[int], Tuple[int, int], Tuple[int, int, int]]],
+        dtype: "tripy.dtype",
+        stride: Sequence[int] = None,
     ) -> None:
         """
         Args:
@@ -88,10 +91,13 @@ def __init__(
             max_shape.append(elem[2])
 
         self.shape_bounds = ShapeBounds(tuple(min_shape), tuple(opt_shape), tuple(max_shape))
+        self.stride = stride
         self.dtype = dtype
 
     def __str__(self) -> str:
-        return f"InputInfo(min={self.shape_bounds.min}, opt={self.shape_bounds.opt}, max={self.shape_bounds.max}, dtype={self.dtype})"
+        base_info = f"InputInfo(min={self.shape_bounds.min}, opt={self.shape_bounds.opt}, max={self.shape_bounds.max}"
+        stride_info = f", stride={self.stride}" if self.stride is not None else ""
+        return f"{base_info}{stride_info}, dtype={self.dtype})"
 
 
 # TODO(MLIR-TRT #923): Can generalize `InputInfo` and drop this class.
@@ -100,5 +106,7 @@ def __str__(self) -> str:
 class ArgInfo:
     shape_bounds: Sequence[Tuple[int, int]]
     """A sequence of tuple(min, max) indicating the bounds of each dimension"""
+    stride: Sequence[int]
+    """A sequence of integers indicating stride"""
     dtype: "tripy.dtype"
     """The datatype of the argument"""