Add checks for non-canonical strides

MLIR-TensorRT requires strides for function arguments and results in canonical order.

Add a check in storage op to ensure strides are in canonical order. This will throw an exception for both compile mode (at execution time) and lazy mode when inputs are materialized to a storage op with a non-canonical stride order.

This PR adds the following changes:

- Add a stride field to Trace tensor. This materializes when populating a storage op with framework tensor.
- Add an assert check while lowering data memref to a `stablehlo.constant` op.
- Add helper methods for stride check: `get_canonical_stride`, `are_strides_equivalent`, and is_canonical_stride. Add corresponding tests.
- Add test_stride.py validating stride for framework tensors.

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.allclose(out, tp.relu(inp), rtol=0.0, atol=0.0)
 
     def test_module(self):
         layernorm = tp.LayerNorm(2)

tripy/tests/frontend/test_stride.py

@@ -0,0 +1,68 @@
+#
+# 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 cupy as cp
+import numpy as np
+import torch
+
+import tripy as tp
+from tests.helper import raises
+
+
+class TestStride:
+    def assert_error_message(self, excinfo, tensor_type, expected_suggestion):
+        error_message = str(excinfo.value)
+        assert "Non-canonical strides are not supported for Tripy tensors." in error_message
+        assert f"For {tensor_type}, use {expected_suggestion}" in error_message
+
+    def tripy_byte_order_strides(self, data):
+        return tuple(s * data.dtype.itemsize for s in tp.Tensor(data).stride())
+
+    def test_non_canonical_stride(self):
+        # PyTorch test
+        t_torch = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+        a_torch = t_torch.transpose(0, 1)
+        with pytest.raises(tp.TripyException) as excinfo:
+            tp.Tensor(a_torch)
+        self.assert_error_message(excinfo, "PyTorch Tensor", "tensor.contiguous() or tensor.clone()")
+
+        # No exception is thrown.
+        print(tp.Tensor(a_torch.contiguous()))
+        print(tp.Tensor(a_torch.clone(memory_format=torch.contiguous_format)))
+
+        # CuPy test
+        t_cupy = cp.arange(12, dtype=cp.float32).reshape(3, 4)
+        a_cupy = t_cupy.transpose(1, 0)
+        with pytest.raises(tp.TripyException) as excinfo:
+            tp.Tensor(a_cupy)
+        self.assert_error_message(excinfo, "CuPy Array", "cp.ascontiguousarray(array) or array.copy(order='C')")
+
+        print(tp.Tensor(cp.ascontiguousarray(a_cupy)))
+        print(tp.Tensor(a_cupy.copy(order="C")))
+
+        # NumPy test
+        t_numpy = np.arange(12, dtype=np.float32).reshape(3, 4)
+        a_numpy = t_numpy.transpose(1, 0)
+        with pytest.raises(tp.TripyException) as excinfo:
+            tp.Tensor(a_numpy)
+        self.assert_error_message(excinfo, "NumPy Array", "np.ascontiguousarray(array) or array.copy(order='C')")
+
+        print(tp.Tensor(np.ascontiguousarray(a_numpy)))
+        print(tp.Tensor(a_numpy.copy(order="C")))

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_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.allclose(quantized, tp.ones((4, 4), dtype=tp.int8), rtol=0.0, atol=0.0)

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

tripy/tripy/backend/mlir/memref.py

@@ -15,9 +15,12 @@
 # limitations under the License.
 #
 
+import re
+
 from functools import lru_cache
 from typing import Sequence
 
+from tripy import utils as tripy_utils
 from tripy.backend.mlir import utils as mlir_utils
 from tripy.common import device as tp_device
 from tripy.common import utils as common_utils
@@ -66,7 +69,50 @@ def create_memref_view(data):
     """
     Creates a memref view of an array object that implements the dlpack interface.
     """
-    return mlir_utils.MLIRRuntimeClient().create_memref_view_from_dlpack(data.__dlpack__())
+    try:
+        memref = mlir_utils.MLIRRuntimeClient().create_memref_view_from_dlpack(data.__dlpack__())
+    except runtime.MTRTException as e:
+        error_msg = str(e)
+        match = re.search(
+            r"Given strides \[([\d, ]+)\] do not match canonical strides \[([\d, ]+)\] for shape \[([\d, ]+)\]",
+            error_msg,
+        )
+
+        if match:
+            given_strides = [int(s) for s in match.group(1).split(",")]
+            canonical_strides = [int(s) for s in match.group(2).split(",")]
+            shape = [int(s) for s in match.group(3).split(",")]
+
+            def check_tensor_type_and_suggest_contiguous(obj):
+                obj_type = str(type(obj))
+                if "torch.Tensor" in obj_type:
+                    return "PyTorch Tensor", "tensor.contiguous() or tensor.clone()"
+                elif "jaxlib" in obj_type or "jax.numpy" in obj_type:
+                    return "JAX Array", "jax.numpy.asarray(array) or jax.numpy.copy(array)"
+                elif "numpy.ndarray" in obj_type:
+                    return "NumPy Array", "np.ascontiguousarray(array) or array.copy(order='C')"
+                elif "cupy.ndarray" in obj_type:
+                    return "CuPy Array", "cp.ascontiguousarray(array) or array.copy(order='C')"
+                else:
+                    return "Unknown Type", "Cannot suggest method for unknown type"
+
+            tensor_type, contiguous_suggestion = check_tensor_type_and_suggest_contiguous(data)
+
+            error_message = (
+                f"Non-canonical strides detected:\n"
+                f"  Shape: {shape}\n"
+                f"  Current stride: {given_strides}\n"
+                f"  Expected canonical stride: {canonical_strides}\n"
+                f"Non-canonical strides are not supported for Tripy tensors. "
+                f"This usually occurs when the tensor is not contiguous in memory. "
+                f"To resolve this issue:\n"
+                f"For {tensor_type}, use {contiguous_suggestion} to ensure contiguity before converting to a Tripy tensor."
+            )
+            tripy_utils.raise_error(error_message)
+        else:
+            # If the error message doesn't match the expected format, re-raise the original exception
+            raise
+    return memref
 
 
 # TODO(#134): Consider move below functions to MLIR py bindings

tripy/tripy/backend/utils.py

@@ -36,6 +36,7 @@ def encode(tensor_info: TensorInfo) -> Dict[str, Any]:
     return {
         "rank": tensor_info.rank,
         "shape": tensor_info.shape,
+        "stride": tensor_info.stride,
         "dtype": tensor_info.dtype,
         "device": tensor_info.device,
     }

tripy/tripy/utils/utils.py

@@ -23,7 +23,7 @@
 import os
 import time
 import typing
-from typing import Any, List, Sequence, Union
+from typing import Any, List, Optional, Sequence, Union
 
 from colored import Fore, Style