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/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_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/mlir/executor.py

@@ -114,7 +114,9 @@ def _get_output_tensor_info(self, outputs_runtime_shape, output_devices):
             is_static_shape = all(dim >= 0 for dim in memref.shape)
             if is_static_shape:
                 outputs_tensor_info.append(
-                    TensorInfo(len(memref.shape), tuple(memref.shape), dtype, device(device_type))
+                    TensorInfo(
+                        len(memref.shape), tuple(memref.shape), tuple(memref.strides), dtype, device(device_type)
+                    )
                 )
             else:
                 runtime_shape = [
@@ -124,6 +126,7 @@ def _get_output_tensor_info(self, outputs_runtime_shape, output_devices):
                     TensorInfo(
                         len(runtime_shape),
                         tuple(runtime_shape),
+                        tuple(memref.strides),
                         dtype,
                         device(device_type),
                     )
@@ -174,7 +177,10 @@ def execute(self, output_devices=List[device], inputs: List["Tensor"] = []) -> L
         # Allocate output memory and store buffer pointers.
         outputs = [
             create_empty_memref(
-                shape=info.shape, dtype=info.dtype, device=info.device, stream=self.stream._active_cuda_stream
+                shape=info.shape,
+                dtype=info.dtype,
+                device=info.device,
+                stream=self.stream._active_cuda_stream,
             )
             for info in out_tensor_info
         ]

tripy/tripy/backend/utils.py

@@ -27,6 +27,7 @@
 class TensorInfo:
     rank: int
     shape: Sequence[int]
+    stride: Sequence[int]
     dtype: "tripy.dtype"
     device: "tripy.device"
 
@@ -36,11 +37,12 @@ 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,
     }
 
 
 @Decoder.register(TensorInfo)
 def decode(dct: Dict[str, Any]) -> TensorInfo:
-    return TensorInfo(dct["rank"], dct["shape"], dct["dtype"], dct["device"])
+    return TensorInfo(dct["rank"], dct["shape"], dct["stride"], dct["dtype"], dct["device"])

tripy/tripy/flat_ir/tensor.py

@@ -38,6 +38,7 @@ class FlatIRTensor:
     rank: int
     producer: "BaseFlatIROp" = None
     shape: Optional[List[int]] = None
+    stride: Optional[List[int]] = None
     reason_details: Optional[List[Any]] = None
     """
     Describes why this tensor was created.
@@ -69,6 +70,7 @@ def build(
         rank: int,
         reason_details: List[Any],
         shape: List[int] = None,
+        stride: List[int] = None,
     ) -> "FlatIRTensor":
         return FlatIRTensor(
             name=None,
@@ -80,6 +82,7 @@ def build(
             rank=rank,
             producer=None,
             shape=shape,
+            stride=utils.get_stride(shape, stride),
             reason_details=reason_details,
             reason_context=copy.copy(_BUILD_CONTEXT),
         )
@@ -88,6 +91,7 @@ def __str__(self) -> str:
         return (
             f"{self.name}: [rank=({self.rank}), "
             + (f"shape=({self.shape}), " if self.shape is not None else "")
+            + (f"stride=({self.stride}), " if self.stride is not None else "")
             + (f"dtype=({self.dtype.name}), " if self.dtype is not None else "")
             + f"loc=({self.device})]"
         )

tripy/tripy/frontend/ops/equal.py

@@ -0,0 +1,59 @@
+#
+# 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.
+#
+
+from tripy import export, constraints
+from tripy.common.exception import raise_error
+
+
+@export.public_api(document_under="operations/functions")
+@constraints.dtype_info(
+    dtype_variables={"T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"]},
+    dtype_constraints={"a": "T1", "b": "T1"},
+)
+def equal(a: "tripy.Tensor", b: "tripy.Tensor") -> bool:
+    r"""
+    Returns true if all elements in ``a`` and ``b`` are exactly equal.
+
+    This function performs an element-wise equality comparison between tensors ``a`` and ``b``,
+    and returns True only if all elements are exactly equal.
+
+    Args:
+        a: First tensor to compare.
+        b: Second tensor to compare.
+
+    Returns:
+        ``True`` if all elements in both tensors are exactly equal, ``False`` otherwise.
+
+    .. code-block:: python
+        :linenos:
+        :caption: Equal Tensors
+
+        # doc: print-locals out
+        out = tp.equal(tp.Tensor([1, 2, 3]), tp.Tensor([1, 2, 3]))
+        assert out
+
+    .. code-block:: python
+        :linenos:
+        :caption: Unequal Tensors
+
+        # doc: print-locals out
+        out = tp.equal(tp.Tensor([1, 2, 3]), tp.Tensor([1, 2, 4]))
+        assert not out
+    """
+    from tripy.frontend.trace.ops.reduce import all
+
+    return bool(all(a == b))

tripy/tripy/frontend/tensor.py

@@ -125,6 +125,7 @@ def __init__(
             attr is not None
             for attr in [
                 self.trace_tensor.shape,
+                self.trace_tensor.stride,
                 self.trace_tensor.dtype,
                 self.trace_tensor.device,
                 self.trace_tensor.producer,
@@ -171,6 +172,10 @@ def dtype(self):
     def rank(self):
         return self.trace_tensor.rank
 
+    @property
+    def stride(self):
+        return self.trace_tensor.stride
+
     def eval(self) -> runtime.MemRefValue:
         from tripy.backend.mlir.compiler import Compiler
         from tripy.backend.mlir.executor import Executor
@@ -228,7 +233,7 @@ def __repr__(self) -> str:
         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={arr.shape}, stride={arr.strides}', prefix=indentation)}"
             f")"
         )
 

tripy/tripy/frontend/trace/ops/binary_elementwise.py

@@ -140,6 +140,7 @@ def to_flat_ir(self, inputs, outputs):
             # First apply DivideOp
             divide_out = FlatIRTensor.build(
                 shape=outputs[0].shape,
+                stride=outputs[0].stride,
                 rank=outputs[0].rank,
                 dtype=outputs[0].dtype,
                 device=outputs[0].device,
@@ -152,6 +153,7 @@ def to_flat_ir(self, inputs, outputs):
             # Step 1: Perform DivideOp
             divide_out = FlatIRTensor.build(
                 shape=outputs[0].shape,
+                stride=outputs[0].stride,
                 rank=outputs[0].rank,
                 dtype=outputs[0].dtype,
                 device=outputs[0].device,
@@ -162,6 +164,7 @@ def to_flat_ir(self, inputs, outputs):
             # Step 2: Apply FloorOp
             floor_out = FlatIRTensor.build(
                 shape=outputs[0].shape,
+                stride=outputs[0].stride,
                 rank=outputs[0].rank,
                 dtype=outputs[0].dtype,
                 device=outputs[0].device,
@@ -172,6 +175,7 @@ def to_flat_ir(self, inputs, outputs):
             # Step 3: Multiply divisor with floored division result (FloorOp output)
             multiply_out = FlatIRTensor.build(
                 shape=outputs[0].shape,
+                stride=outputs[0].stride,
                 rank=outputs[0].rank,
                 dtype=outputs[0].dtype,
                 device=outputs[0].device,