chore: import upstream snapshot with attribution

tests/python/relax/test_op_inspect.py

@@ -0,0 +1,229 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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.
+# ruff: noqa: F401
+
+import ctypes
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm.testing
+from tvm import relax
+from tvm.ir import Op
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+# Parameterization for reading dtype of DLTensor.  Chosen to have
+# multiple distinct type codes, number of lanes, and widths.
+dtype = tvm.testing.parameter(
+    "int32",
+    "int64",
+    "float32",
+    "float32x4",
+    "bfloat",
+    "float8_e4m3fn",
+)
+shape = tvm.testing.parameter(
+    [],
+    [16],
+    [128, 256],
+    [1] * 64,
+)
+
+elem_offset = tvm.testing.parameter(0, 64, 128)
+
+
+def inspect_tensor_field(op_name, *args):
+    return relax.Call(Op.get(f"relax.inspect.{op_name}"), args)
+
+
+def test_tensor_dtype_code(dtype):
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor):
+            return inspect_tensor_field("tensor_dtype_code", A)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    arg = tvm.runtime.empty([16], dtype)
+    res = vm["main"](arg)
+
+    expected_type_code = tvm.runtime.DataType(dtype).type_code
+    assert res == expected_type_code
+
+
+def test_tensor_dtype_bits(dtype):
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor):
+            return inspect_tensor_field("tensor_dtype_bits", A)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    arg = tvm.runtime.empty([16], dtype)
+    res = vm["main"](arg)
+
+    expected_type_bits = tvm.runtime.DataType(dtype).bits
+    assert res == expected_type_bits
+
+
+def test_tensor_dtype_lanes(dtype):
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor):
+            return inspect_tensor_field("tensor_dtype_lanes", A)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    arg = tvm.runtime.empty([16], dtype)
+    res = vm["main"](arg)
+
+    expected_type_lanes = tvm.runtime.DataType(dtype).lanes
+    assert res == expected_type_lanes
+
+
+def test_tensor_ndim(shape):
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor):
+            return inspect_tensor_field("tensor_ndim", A)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    arg = tvm.runtime.empty(shape, "int32")
+    res = vm["main"](arg)
+
+    assert res == len(shape)
+
+
+def test_tensor_shape(shape):
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor, axis: R.Prim("int64")):
+            return inspect_tensor_field("tensor_shape_i", A, axis)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    arg = tvm.runtime.empty(shape, "int32")
+
+    res = [vm["main"](arg, i) for i, _ in enumerate(shape)]
+
+    tvm.ir.assert_structural_equal(res, shape)
+
+
+def _get_compact_striding(shape):
+    strides = []
+    product = 1
+    for dim in reversed(shape):
+        strides.append(product)
+        product *= dim
+    return list(reversed(strides))
+
+
+def test_strides_of_compact_tensor(shape):
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor, axis: R.Prim("int64")):
+            return inspect_tensor_field("tensor_stride_i", A, axis)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    arg = tvm.runtime.empty(shape, "int32")
+
+    res = [vm["main"](arg, i) for i, _ in enumerate(shape)]
+    expected = _get_compact_striding(shape)
+
+    tvm.ir.assert_structural_equal(res, expected)
+
+
+def test_strides_of_non_compact_tensor():
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor, axis: R.Prim("int64")):
+            return inspect_tensor_field("tensor_stride_i", A, axis)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+    view_shape = [4, 4]
+    expected_strides = [1, 4]
+    # use transpose to make strides non-compact
+    x = np.zeros([4, 4], "int32").T
+    y = tvm_ffi.from_dlpack(x, require_alignment=4, require_contiguous=False)
+    res = [vm["main"](y, i) for i, _ in enumerate(view_shape)]
+    tvm.ir.assert_structural_equal(res, expected_strides)
+
+
+def test_byte_offset(elem_offset):
+    backing_shape = [64, 64]
+    view_shape = [16, 16]
+    byte_offset = elem_offset * 4
+
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor):
+            return inspect_tensor_field("tensor_byte_offset", A)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+    dtype = "int32"
+    backing_tensor = tvm.runtime.empty(backing_shape, dtype)
+    view = backing_tensor._create_view(view_shape, dtype, relative_byte_offset=byte_offset)
+    res = vm["main"](view)
+    assert res == byte_offset
+
+
+def test_elem_offset(elem_offset, dtype):
+    tvm_dtype = tvm.runtime.DataType(dtype)
+
+    backing_shape = [64, 64]
+    view_shape = [16, 16]
+    element_bytes = (tvm_dtype.bits * tvm_dtype.lanes) // 8
+    byte_offset = elem_offset * element_bytes
+
+    @I.ir_module
+    class mod:
+        @R.function
+        def main(A: R.Tensor):
+            return inspect_tensor_field("tensor_elem_offset", A)
+
+    built = tvm.compile(mod)
+    vm = relax.VirtualMachine(built, tvm.cpu())
+
+    backing_tensor = tvm.runtime.empty(backing_shape, dtype)
+    view = backing_tensor._create_view(view_shape, dtype, relative_byte_offset=byte_offset)
+    res = vm["main"](view)
+
+    assert res == elem_offset
+
+
+if __name__ == "__main__":
+    tvm.testing.main()