Add checks for non-canonical strides

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

NVIDIA#252 adds a check to validate memref stride against a canonical stride order. In Tripy, memref strides are derived from framework DL Pack tensors. Creating a memref with a non-canonical DL Pack tensor stride throws an exception.

Add a try-catch block to catch such an exception and augment with suggestions on creating a DL Pack tensor with canonical stride for Tripy-supported frameworks.

Add unit tests to create a non-canonical stride tensor to validate exceptions and suggestions.

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.utils import raise_error
 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,54 @@ 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 None, None
+
+            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."
+                    if tensor_type is not None
+                    else ""
+                )
+            )
+            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