# 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: F841


import numpy as np
import pytest
import torch

import tvm
from tvm import relax, testing
from tvm.relax import VMInstrumentReturnKind
from tvm.relax.testing.transform import (
    dataflow_alias_analysis,
    dataflow_inplace_analysis,
    dataflow_liveness_analysis,
    dataflow_single_inplace_call,
)
from tvm.relax.transform import DataflowUseInplaceCalls
from tvm.script.parser import ir as I
from tvm.script.parser import relax as R
from tvm.script.parser import tirx as T


def test_liveness_analysis():
    @I.ir_module(s_tir=True)
    class BasicLiveness:
        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            with R.dataflow():
                y = R.const(1, dtype="int32")
                z = R.add(x, y)
                q = R.multiply(z, y)
                p = R.add(z, q)
                n = R.multiply(p, p)
                R.output(n, p)
            return n

    block = BasicLiveness["main"].body.blocks[0]
    live_ranges = dataflow_liveness_analysis(block)
    expected_ranges = {
        # x is live past the binding block
        "x": (-1, 5),
        "y": (0, 2),
        "z": (1, 3),
        "q": (2, 3),
        # exposed though ultimately not used
        "p": (3, 5),
        "n": (4, 5),
    }
    actual_ranges = {var.name: live_range for var, live_range in live_ranges.items()}
    assert actual_ranges == expected_ranges


def test_alias_analysis_basic():
    @I.ir_module(s_tir=True)
    class BasicAliasAnalysis:
        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            with R.dataflow():
                y = x  # y is an alias of x
                z = R.add(y, y)  # fresh value
                n = z  # alias of z
                R.output(n)
            return n

    block = BasicAliasAnalysis["main"].body.blocks[0]
    alias_sets, tuple_map = dataflow_alias_analysis(block, BasicAliasAnalysis["main"].params)
    expected = {
        "x": {0},
        "y": {0},
        "z": {1},
        "n": {1},
    }

    for var, alias_set in alias_sets.items():
        assert alias_set == expected[var.name]
    assert tuple_map == {}


def test_alias_analysis_tuple():
    @I.ir_module(s_tir=True)
    class AliasesWithTuples:
        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            with R.dataflow():
                y = R.const(1, dtype="int32")
                t = (x, y)
                a = t[0]
                b = t[1]
                c = t[0]
                d = t[1]
                u = t
                e = t[0]
                f = t[1]
                z = R.add(c, d)
                n = z
                R.output(n)
            return n

    block = AliasesWithTuples["main"].body.blocks[0]
    alias_sets, tuple_map = dataflow_alias_analysis(block, AliasesWithTuples["main"].params)
    expected = {
        "x": {0},
        "y": {1},
        "t": {2},
        "a": {0},
        "b": {1},
        "c": {0},
        "d": {1},
        "u": {2},
        "e": {0},
        "f": {1},
        "z": {3},
        "n": {3},
    }

    actual_alias_sets = {var.name: alias_set for var, alias_set in alias_sets.items()}
    assert expected == actual_alias_sets
    assert 2 in tuple_map
    assert tuple_map[2] == [{0}, {1}]


def test_alias_split():
    @I.ir_module(s_tir=True)
    class AliasSplit:
        @R.function
        def main(x: R.Tensor((60,), "int32")) -> R.Tensor((15,), "int32"):
            with R.dataflow():
                t = R.split(x, 4)
                y = t[0]
                z = t[1]
                q = t[2]
                p = t[3]
                n = z
                R.output(n)
            return n

    block = AliasSplit["main"].body.blocks[0]
    alias_sets, tuple_map = dataflow_alias_analysis(block, AliasSplit["main"].params)
    expected = {
        "x": {0},
        "t": {1},
        "y": {2},
        "z": {3},
        "q": {4},
        "p": {5},
        "n": {3},
    }

    actual_alias_sets = {var.name: alias_set for var, alias_set in alias_sets.items()}
    assert expected == actual_alias_sets
    assert len(tuple_map) == 1
    assert 1 in tuple_map
    assert tuple_map[1] == [{2}, {3}, {4}, {5}]


def test_alias_call_tir():
    # call TIR can yield either a single tensor or a tuple
    @I.ir_module(s_tir=True)
    class AliasCallTir:
        @T.prim_func(s_tir=True)
        def tir_id(x: T.handle, y: T.handle) -> None:
            T.func_attr({"global_symbol": "tir_id"})
            m = T.int32()
            n = T.int32()
            A = T.match_buffer(x, (m, n), "int32")
            B = T.match_buffer(y, (m, n), "int32")

            for i, j in T.grid(m, n):
                with T.sblock("id"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    B[vi, vj] = A[vi, vj]

        @T.prim_func(s_tir=True)
        def tir_id2(x: T.handle, y: T.handle, z: T.handle) -> None:
            T.func_attr({"global_symbol": "tir_id"})
            m = T.int32()
            n = T.int32()
            A = T.match_buffer(x, (m, n), "int32")
            B = T.match_buffer(y, (m, n), "int32")
            C = T.match_buffer(z, (m, n), "int32")

            for i, j in T.grid(m, n):
                with T.sblock("id"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    B[vi, vj] = A[vi, vj]
                    C[vi, vj] = A[vi, vj]

        @R.function
        def main(x: R.Tensor((10, 10), "int32")) -> R.Tensor((10, 10), "int32"):
            with R.dataflow():
                cls = AliasCallTir
                y = R.call_tir(cls.tir_id, (x,), out_ty=R.Tensor((10, 10), "int32"))
                t = R.call_tir(
                    cls.tir_id2,
                    (y,),
                    out_ty=[R.Tensor((10, 10), "int32"), R.Tensor((10, 10), "int32")],
                )
                z = y
                p = t[0]
                q = t[1]
                u = t
                m = u[0]
                n = u[1]
                v = n
                R.output(v)
            return v

    block = AliasCallTir["main"].body.blocks[0]
    alias_sets, tuple_map = dataflow_alias_analysis(block, AliasCallTir["main"].params)
    expected = {
        "x": {0},
        "y": {1},
        "t": {2},
        "z": {1},
        "p": {3},
        "q": {4},
        "u": {2},
        "m": {3},
        "n": {4},
        "v": {4},
    }

    actual_alias_sets = {var.name: alias_set for var, alias_set in alias_sets.items()}
    assert expected == actual_alias_sets
    assert len(tuple_map) == 1
    assert 2 in tuple_map
    assert tuple_map[2] == [{3}, {4}]


def test_mystery_calls():
    @I.ir_module(s_tir=True)
    class AliasChaosCalls:
        @R.function
        def identity(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            return x

        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            with R.dataflow():
                cls = AliasChaosCalls
                y = cls.identity(x)
                z = cls.identity(y)
                m = R.const(1, dtype="int32")
                n = R.const(2, dtype="int32")
                t = (m, n)
                a = R.call_pure_packed(
                    "chaos", t, ty_args=R.Tuple(R.Tensor((), "int32"), R.Tensor((), "int32"))
                )
                b = a[0]
                c = a[1]
                R.output(c)
            return c

    block = AliasChaosCalls["main"].body.blocks[0]
    alias_sets, tuple_map = dataflow_alias_analysis(block, AliasChaosCalls["main"].params)
    expected = {
        "x": {0},
        "y": {0, 1},
        "z": {0, 1, 2},
        "m": {3},
        "n": {4},
        "t": {5},
        "a": {3, 4, 5, 6, 7, 8},  # either t or a fresh tuple
        "b": {3, 4, 5, 6, 7, 8},  # the tuple components can be aliased to any member...
        "c": {3, 4, 5, 6, 7, 8},  # the tuple components can be aliased to any member...
        # (in principle, we can use type information to narrow down the aliasing)
    }

    actual_alias_sets = {var.name: alias_set for var, alias_set in alias_sets.items()}
    assert expected == actual_alias_sets
    assert len(tuple_map) == 2
    assert 5 in tuple_map
    assert tuple_map[5] == [{3}, {4}]
    assert 6 in tuple_map
    assert tuple_map[6] == [{3, 4, 5, 6, 7, 8}, {3, 4, 5, 6, 7, 8}]


def test_alias_external_value():
    @I.ir_module(s_tir=True)
    class AliasExternalValue:
        @R.function
        def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
            y = R.const(1, dtype="int32")  # not in DF block, treated as external
            t1 = (y, y)  # not in DF block, treated as external
            with R.dataflow():
                z = y  # mystery value
                a = R.const(2, dtype="int32")
                t2 = (z, a)
                b = t2[0]
                c = t1[1]  # tuple index into external value
                R.output(b)
            return b

    block = AliasExternalValue["main"].body.blocks[1]
    alias_sets, tuple_map = dataflow_alias_analysis(block, AliasExternalValue["main"].params)
    expected = {
        "x": {0},
        "z": {-1},
        "a": {1},
        "t2": {2},
        "b": {-1},
        "c": {-1},
    }

    actual_alias_sets = {var.name: alias_set for var, alias_set in alias_sets.items()}
    assert expected == actual_alias_sets
    assert len(tuple_map) == 1
    assert 2 in tuple_map
    assert tuple_map[2] == [{-1}, {1}]


def test_inplace_simple_case():
    @I.ir_module(s_tir=True)
    class InplaceBasic:
        @R.function
        def main(x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")) -> R.Tensor(
            (2, 3), "int32"
        ):
            with R.dataflow():
                z = R.add(x, y)  # cannot be done inplace: x and y are live later
                p = R.add(z, z)  # can be done inplace: z is not used later
                r = p  # alias of p
                m = R.multiply(p, p)  # p is not used later but r is, so can't do inplace
                n = R.add(m, r)  # can be done inplace: r is not used again
                ret = R.subtract(n, m)  # can be done inplace: neither is used again
                R.output(ret)
            return ret

    block = InplaceBasic["main"].body.blocks[0]
    size_match, exact_match = dataflow_inplace_analysis(
        block, InplaceBasic["main"].params, InplaceBasic
    )

    # order does not matter for the listing of candidates, so we have to implement as sets
    def assert_candidate_list(
        actual: list[tuple[int, set[int]]], expected: list[tuple[int, set[int]]]
    ) -> None:
        assert len(actual) == len(expected)
        for i in range(len(actual)):
            assert actual[i][0] == expected[i][0]
            assert len(expected[i][1]) == len(actual[i][1])
            for idx in actual[i][1]:
                assert idx in expected[i][1]

    assert_candidate_list(size_match, [(1, {0, 1}), (4, {1}), (5, {0, 1})])
    # TODO(@slyubomirsky): I couldn't think of an easy example where sizes don't match,
    # but broadcasting might cause it to happen
    assert_candidate_list(exact_match, [(1, {0, 1}), (4, {1}), (5, {0, 1})])


def test_inplace_single_call():
    @I.ir_module(s_tir=True)
    class TestModule:
        @R.function
        def main(
            x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")
        ) -> R.Tensor((2, 3), dtype="float32"):
            z = R.add(x, y)
            q = R.nn.silu(z)
            return q

    add_call = TestModule["main"].body.blocks[0].bindings[0].value
    new_add, new_mod = dataflow_single_inplace_call(TestModule, add_call, [0])

    @T.prim_func(private=True, s_tir=True)
    def expected_add(
        A: T.Buffer((T.int64(2), T.int64(3)), "float32"),
        B: T.Buffer((T.int64(2), T.int64(3)), "float32"),
    ):
        T.func_attr({"tirx.noalias": True})
        for ax0, ax1 in T.grid(T.int64(2), T.int64(3)):
            with T.sblock("T_add"):
                v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
                T.writes(A[v_ax0, v_ax1])
                A[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[v_ax0, v_ax1]

    tvm.ir.assert_structural_equal(new_mod["add_inplace"], expected_add)
    assert new_add.op.name == "relax.call_tir_inplace"
    assert new_add.args[0].name_hint == "add_inplace"
    for i, arg in enumerate(new_add.args[1].fields):
        arg == add_call.args[i]
    new_add.attrs.inplace_indices == [0]

    @T.prim_func(private=True, s_tir=True)
    def expected_silu(A: T.Buffer((T.int64(2), T.int64(3)), "float32")):
        T.func_attr({"tirx.noalias": True})
        compute = T.sblock_alloc_buffer((T.int64(2), T.int64(3)))
        for i0, i1 in T.grid(T.int64(2), T.int64(3)):
            with T.sblock("compute"):
                v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                T.reads(A[v_i0, v_i1])
                T.writes(compute[v_i0, v_i1])
                compute[v_i0, v_i1] = T.sigmoid(A[v_i0, v_i1])
        for ax0, ax1 in T.grid(T.int64(2), T.int64(3)):
            with T.sblock("T_multiply"):
                v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                T.reads(A[v_ax0, v_ax1], compute[v_ax0, v_ax1])
                T.writes(A[v_ax0, v_ax1])
                A[v_ax0, v_ax1] = A[v_ax0, v_ax1] * compute[v_ax0, v_ax1]

    silu_call = TestModule["main"].body.blocks[0].bindings[1].value
    new_silu, new_mod = dataflow_single_inplace_call(TestModule, silu_call, [0])

    tvm.ir.assert_structural_equal(new_mod["silu_inplace"], expected_silu)
    assert new_silu.op.name == "relax.call_tir_inplace"
    assert new_silu.args[0].name_hint == "silu_inplace"
    for i, arg in enumerate(new_silu.args[1].fields):
        arg == silu_call.args[i]
    new_silu.attrs.inplace_indices == [0]


def test_insert_inplace_calls():
    @I.ir_module(s_tir=True)
    class EndToEndTest:
        @R.function
        def main(
            x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((1, 3), dtype="float32")
        ) -> R.Tensor((2, 3), dtype="float32"):
            with R.dataflow():
                z = R.add(x, y)  # broadcast happens here
                # Cannot be done in-place because x is an argument.
                a = R.add(z, y)  # this one can be done in-place
                q = R.multiply(a, y)  # broadcast again, a is eligible
                r = R.subtract(y, y)  # cannot be done in-place because y is an argument
                s = R.subtract(r, r)  # No broadcast. Can be done in-place
                m = R.multiply(q, s)  # should give us all zeros
                R.output(m)
            return m

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def add_inplace(
            A: T.Buffer((T.int64(2), T.int64(3)), "float32"),
            B: T.Buffer((T.int64(1), T.int64(3)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1 in T.grid(T.int64(2), T.int64(3)):
                with T.sblock("T_add"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1], B[T.int64(0), v_ax1])
                    T.writes(A[v_ax0, v_ax1])
                    A[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[T.int64(0), v_ax1]

        @T.prim_func(private=True, s_tir=True)
        def multiply_inplace(
            A: T.Buffer((T.int64(2), T.int64(3)), "float32"),
            B: T.Buffer((T.int64(1), T.int64(3)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1 in T.grid(T.int64(2), T.int64(3)):
                with T.sblock("T_multiply"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1], B[T.int64(0), v_ax1])
                    T.writes(A[v_ax0, v_ax1])
                    A[v_ax0, v_ax1] = A[v_ax0, v_ax1] * B[T.int64(0), v_ax1]

        @T.prim_func(private=True, s_tir=True)
        def subtract_inplace(
            A: T.Buffer((T.int64(1), T.int64(3)), "float32"),
            B: T.Buffer((T.int64(1), T.int64(3)), "float32"),
        ):
            T.func_attr({"tirx.noalias": True})
            for ax0, ax1 in T.grid(T.int64(1), T.int64(3)):
                with T.sblock("T_subtract"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
                    T.writes(B[v_ax0, v_ax1])
                    B[v_ax0, v_ax1] = A[v_ax0, v_ax1] - B[v_ax0, v_ax1]

        @R.function
        def main(
            x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((1, 3), dtype="float32")
        ) -> R.Tensor((2, 3), dtype="float32"):
            cls = Expected
            with R.dataflow():
                z: R.Tensor((2, 3), dtype="float32") = R.add(x, y)
                a: R.Tensor((2, 3), dtype="float32") = R.call_tir_inplace(
                    cls.add_inplace,
                    (z, y),
                    inplace_indices=[0],
                    out_ty=[
                        R.Tensor((2, 3), dtype="float32"),
                    ],
                )
                q: R.Tensor((2, 3), dtype="float32") = R.call_tir_inplace(
                    cls.multiply_inplace,
                    (a, y),
                    inplace_indices=[0],
                    out_ty=[
                        R.Tensor((2, 3), dtype="float32"),
                    ],
                )
                r: R.Tensor((1, 3), dtype="float32") = R.subtract(y, y)
                s: R.Tensor((1, 3), dtype="float32") = R.call_tir_inplace(
                    cls.subtract_inplace,
                    (r, r),
                    inplace_indices=[1],
                    out_ty=[
                        R.Tensor((1, 3), dtype="float32"),
                    ],
                )
                m: R.Tensor((2, 3), dtype="float32") = R.call_tir_inplace(
                    cls.multiply_inplace,
                    (q, s),
                    inplace_indices=[0],
                    out_ty=[
                        R.Tensor((2, 3), dtype="float32"),
                    ],
                )
                R.output(m)
            return m

    transform_pass = DataflowUseInplaceCalls()
    new_mod = transform_pass(EndToEndTest)
    tvm.ir.assert_structural_equal(new_mod, Expected)

    x = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
    y = tvm.runtime.tensor(np.random.rand(1, 3).astype("float32"))
    expected = np.zeros((2, 3), dtype="float32")

    target = tvm.target.Target("llvm")
    ex = tvm.compile(new_mod, target)
    vm = relax.VirtualMachine(ex, tvm.cpu())
    res = vm["main"](x, y)
    assert (expected == res.numpy()).all()


def test_dynamic():
    @I.ir_module(s_tir=True)
    class DynamicTestCase:
        @R.function
        def main(
            x: R.Tensor(("a", "b"), dtype="float32"), y: R.Tensor(("a", "b"), dtype="float32")
        ) -> R.Tensor(("a", "b"), dtype="float32"):
            with R.dataflow():
                z = R.add(x, y)
                # Cannot be done in-place because x and y are arguments
                a = R.add(z, y)  # this one can be done in-place
                s = R.subtract(a, a)  # No broadcast. Can be done in-place
                R.output(s)
            return s

    # the result should be all zeroes
    transform_pass = DataflowUseInplaceCalls()
    new_mod = transform_pass(DynamicTestCase)

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def add_inplace(var_A: T.handle, var_B: T.handle):
            T.func_attr({"tirx.noalias": True})
            a, b = T.int64(), T.int64()
            A = T.match_buffer(var_A, (a, b))
            B = T.match_buffer(var_B, (a, b))
            for ax0, ax1 in T.grid(a, b):
                with T.sblock("T_add"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
                    T.writes(A[v_ax0, v_ax1])
                    A[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[v_ax0, v_ax1]

        @T.prim_func(private=True, s_tir=True)
        def subtract_inplace(var_A: T.handle, var_B: T.handle):
            T.func_attr({"tirx.noalias": True})
            a, b = T.int64(), T.int64()
            A = T.match_buffer(var_A, (a, b))
            B = T.match_buffer(var_B, (a, b))
            for ax0, ax1 in T.grid(a, b):
                with T.sblock("T_subtract"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
                    T.writes(B[v_ax0, v_ax1])
                    B[v_ax0, v_ax1] = A[v_ax0, v_ax1] - B[v_ax0, v_ax1]

        @R.function
        def main(
            x: R.Tensor(("a", "b"), dtype="float32"), y: R.Tensor(("a", "b"), dtype="float32")
        ) -> R.Tensor(("a", "b"), dtype="float32"):
            a = T.int64()
            b = T.int64()
            cls = Expected
            with R.dataflow():
                z = R.add(x, y)
                a_1 = R.call_tir_inplace(
                    cls.add_inplace,
                    (z, y),
                    out_ty=R.Tensor((a, b), dtype="float32"),
                    inplace_indices=[0],
                )
                s = R.call_tir_inplace(
                    cls.subtract_inplace,
                    (a_1, a_1),
                    out_ty=R.Tensor((a, b), dtype="float32"),
                    inplace_indices=[1],
                )
                R.output(s)
            return s

    tvm.ir.assert_structural_equal(new_mod, Expected, map_free_vars=True)
    x = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
    y = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
    expected = np.zeros((2, 3), dtype="float32")

    target = tvm.target.Target("llvm")
    ex = tvm.compile(new_mod, target)
    vm = relax.VirtualMachine(ex, tvm.cpu())
    res = vm["main"](x, y)
    assert (expected == res.numpy()).all()


def test_dynamic_mismatch():
    # cannot statically prove the shapes to be equal so the module should be unchanged
    @I.ir_module(s_tir=True)
    class DynamicMistmatchTestCase:
        @R.function
        def main(
            x: R.Tensor(("a", "b"), dtype="float32"), y: R.Tensor(("c", "d"), dtype="float32")
        ):
            with R.dataflow():
                z = R.add(x, y)
                # Cannot be done in-place because x and y are arguments
                a = R.add(z, y)  # cannot conclude that shapes match
                R.output(a)
            return a

    transform_pass = DataflowUseInplaceCalls()
    new_mod = transform_pass(DynamicMistmatchTestCase)
    tvm.ir.assert_structural_equal(new_mod, DynamicMistmatchTestCase)


class TestViewOpSharedStorageAndNoInplace:
    storage_ptr_x_1d = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
    storage_ptr_x_2d = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
    storage_ptr_x_squeeze = np.array([[[1.0], [2.0], [3.0], [4.0]]], dtype=np.float32)
    storage_ptr_x_ensure_zero_offset = np.array([[1.0], [2.0], [3.0], [4.0]], dtype=np.float32)

    @I.ir_module
    class _SharedStorageExpandDimsModule:
        @R.function
        def main(x: R.Tensor((4,), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 1), dtype="float32") = R.expand_dims(x, axis=[1])
                lv1: R.Tensor((4, 1), dtype="float32") = R.expand_dims(x, axis=[1])
                gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class _SharedStorageSqueezeModule:
        @R.function
        def main(x: R.Tensor((1, 4, 1), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 1), dtype="float32") = R.squeeze(x, axis=[0])
                lv1: R.Tensor((4, 1), dtype="float32") = R.squeeze(x, axis=[0])
                gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class _SharedStorageReshapeModule:
        @R.function
        def main(x: R.Tensor((4,), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 1), dtype="float32") = R.reshape(x, (4, 1))
                lv1: R.Tensor((4, 1), dtype="float32") = R.reshape(x, (4, 1))
                gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class _SharedStoragePermuteDimsModule:
        @R.function
        def main(x: R.Tensor((1, 4), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 1), dtype="float32") = R.permute_dims(x, axes=[1, 0])
                lv1: R.Tensor((4, 1), dtype="float32") = R.permute_dims(x, axes=[1, 0])
                gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class _SharedStorageViewModule:
        @R.function
        def main(x: R.Tensor((4,), dtype="float32")) -> R.Tensor((1, 4), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 4), dtype="float32") = R.memory.view(
                    x, R.shape([1, 4]), R.tuple(), R.tuple()
                )
                lv1: R.Tensor((1, 4), dtype="float32") = R.memory.view(
                    x, R.shape([1, 4]), R.tuple(), R.tuple()
                )
                gv: R.Tensor((1, 4), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class _SharedStorageBatchFlattenModule:
        @R.function
        def main(x: R.Tensor((1, 4), dtype="float32")) -> R.Tensor((1, 4), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 4), dtype="float32") = R.nn.batch_flatten(x)
                lv1: R.Tensor((1, 4), dtype="float32") = R.nn.batch_flatten(x)
                gv: R.Tensor((1, 4), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class _SharedStorageFlattenModule:
        @R.function
        def main(x: R.Tensor((1, 4), dtype="float32")) -> R.Tensor((4,), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4,), dtype="float32") = R.flatten(x)
                lv1: R.Tensor((4,), dtype="float32") = R.flatten(x)
                gv: R.Tensor((4,), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class _SharedStorageEnsureZeroOffsetModule:
        @R.function
        def main(x: R.Tensor((4, 1), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 1), dtype="float32") = R.memory.ensure_zero_offset(x)
                lv1: R.Tensor((4, 1), dtype="float32") = R.memory.ensure_zero_offset(x)
                gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class _IndependentReluModule:
        """Just a testcase to verify that non-view ops do not share storage."""

        @R.function
        def main(x: R.Tensor((4,), dtype="float32")) -> R.Tensor((4,), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4,), dtype="float32") = R.nn.relu(x)
                lv1: R.Tensor((4,), dtype="float32") = R.nn.relu(x)
                gv: R.Tensor((4,), dtype="float32") = R.add(lv, lv1)
                R.output(gv)
            return gv

    @classmethod
    def _capture_op_tensors(cls, mod, input_nps, op_substr):
        """Capture TVM tensors passed to VM calls whose name contains op_substr."""
        captures = []

        def instrument(func, name, before_run, ret_value, *args):
            del func, ret_value
            if not before_run:
                return VMInstrumentReturnKind.NO_OP
            if op_substr not in name.lower():
                return VMInstrumentReturnKind.NO_OP
            tensor_args = [arg for arg in args if isinstance(arg, tvm.runtime.Tensor)]
            if not tensor_args:
                return VMInstrumentReturnKind.NO_OP
            captures.append({"call_name": name, "tensors": tensor_args})
            return VMInstrumentReturnKind.NO_OP

        if isinstance(input_nps, np.ndarray):
            input_nps = [input_nps]

        ex = relax.build(mod, tvm.target.Target("llvm"))
        vm = relax.VirtualMachine(ex, tvm.cpu())
        vm.set_instrument(instrument)
        vm["main"](*(tvm.runtime.tensor(arr, tvm.cpu()) for arr in input_nps))
        return captures

    @pytest.mark.parametrize(
        "mod,input_nps,op_substr,expect_same_storage",
        [
            pytest.param(
                _SharedStorageExpandDimsModule,
                [storage_ptr_x_1d],
                "add",
                True,
                id="shared_storage_expand_dims",
            ),
            pytest.param(
                _SharedStorageSqueezeModule,
                [storage_ptr_x_squeeze],
                "add",
                True,
                id="shared_storage_squeeze",
            ),
            pytest.param(
                _SharedStorageReshapeModule,
                [storage_ptr_x_1d],
                "add",
                True,
                id="shared_storage_reshape",
            ),
            pytest.param(
                _SharedStoragePermuteDimsModule,
                [storage_ptr_x_2d],
                "add",
                True,
                id="shared_storage_permute_dims",
            ),
            pytest.param(
                _SharedStorageFlattenModule,
                [storage_ptr_x_2d],
                "add",
                True,
                id="shared_storage_flatten",
            ),
            pytest.param(
                _SharedStorageBatchFlattenModule,
                [storage_ptr_x_2d],
                "add",
                True,
                id="shared_storage_batch_flatten",
            ),
            pytest.param(
                _SharedStorageViewModule,
                [storage_ptr_x_1d],
                "add",
                True,
                id="shared_storage_memory_view",
            ),
            pytest.param(
                _SharedStorageEnsureZeroOffsetModule,
                [storage_ptr_x_ensure_zero_offset],
                "add",
                True,
                id="shared_storage_ensure_zero_offset",
            ),
            pytest.param(
                _IndependentReluModule,
                [storage_ptr_x_1d],
                "add",
                False,
                id="independent_storage_relu",
            ),
        ],
    )
    def test_tensor_storage_ptr_extraction(self, mod, input_nps, op_substr, expect_same_storage):
        """Validate runtime storage overlap/sharing via VM instrumentation."""
        storage_shared = tvm.get_global_func("runtime.TVMTensorIsStorageShared")
        captures = self._capture_op_tensors(mod, input_nps, op_substr)
        assert len(captures), f"VM instrumentation did not see a {op_substr} call."
        assert len(captures) == 1, f"VM instrumentation should see exactly one {op_substr} call."
        cap = captures[0]
        assert len(cap["tensors"]) == 3, (
            f"VM instrumentation should see three {op_substr} tensor operands."
        )
        tensor_a, tensor_b = cap["tensors"][0], cap["tensors"][1]
        call_name = cap["call_name"]
        if expect_same_storage:
            assert storage_shared(tensor_a, tensor_b), (
                f"{mod.__name__}: operands should share the same storage (call {call_name!r})"
            )
        else:
            assert not storage_shared(tensor_a, tensor_b), (
                f"{mod.__name__}: operands must not share storage (call {call_name!r})"
            )

    @staticmethod
    def _emit_duplicate_view(op, x):
        if op == "relax.expand_dims":
            a = relax.op.expand_dims(x, axis=1)
            b = relax.op.expand_dims(x, axis=1)
        elif op == "relax.squeeze":
            a = relax.op.squeeze(x, axis=[0])
            b = relax.op.squeeze(x, axis=[0])
        elif op == "relax.reshape":
            a = relax.op.reshape(x, (4, 1))
            b = relax.op.reshape(x, (4, 1))
        elif op == "relax.permute_dims":
            a = relax.op.permute_dims(x, axes=[1, 0])
            b = relax.op.permute_dims(x, axes=[1, 0])
        elif op == "relax.memory.view":
            a = relax.op.memory.view(x, (4, 1))
            b = relax.op.memory.view(x, (4, 1))
        elif op == "relax.memory.ensure_zero_offset":
            a = relax.op.memory.ensure_zero_offset(x)
            b = relax.op.memory.ensure_zero_offset(x)
        elif op == "relax.flatten":
            a = relax.op.flatten(x)
            b = relax.op.flatten(x)
        elif op == "relax.nn.batch_flatten":
            a = relax.op.nn.batch_flatten(x)
            b = relax.op.nn.batch_flatten(x)
        else:
            raise ValueError(op)
        return a, b

    @staticmethod
    def _concat_axis_for_view_op(op):
        if op == "relax.flatten":
            return 0
        return 1

    @classmethod
    def _build_module(cls, op):
        if op == "relax.expand_dims":
            x_ty = relax.TensorType((4,), "float32")
        elif op == "relax.squeeze":
            x_ty = relax.TensorType((1, 4, 1), "float32")
        elif op == "relax.reshape":
            x_ty = relax.TensorType((4,), "float32")
        elif op == "relax.permute_dims":
            x_ty = relax.TensorType((1, 4), "float32")
        elif op == "relax.memory.view":
            x_ty = relax.TensorType((4,), "float32")
        elif op == "relax.memory.ensure_zero_offset":
            x_ty = relax.TensorType((4, 1), "float32")
        elif op in ("relax.flatten", "relax.nn.batch_flatten"):
            x_ty = relax.TensorType((1, 4), "float32")
        else:
            raise ValueError(op)

        bb = relax.BlockBuilder()
        x = relax.Var("x", x_ty)
        concat_axis = cls._concat_axis_for_view_op(op)
        with bb.function("main", [x]):
            with bb.dataflow():
                a_expr, b_expr = cls._emit_duplicate_view(op, x)
                a = bb.emit(a_expr)
                b = bb.emit(b_expr)
                prod = bb.emit(relax.op.multiply(a, b))
                out = bb.emit(relax.op.concat([prod, b], axis=concat_axis))
                gv = bb.emit_output(out)
            bb.emit_func_output(gv)
        return bb.finalize()

    @classmethod
    def _input_for_view_op(cls, op):
        if op == "relax.squeeze":
            return cls.storage_ptr_x_squeeze
        if op == "relax.memory.ensure_zero_offset":
            return cls.storage_ptr_x_ensure_zero_offset
        if op in ("relax.permute_dims", "relax.flatten", "relax.nn.batch_flatten"):
            return cls.storage_ptr_x_2d
        return cls.storage_ptr_x_1d

    @staticmethod
    def _torch_duplicate_view(x, op):
        if op == "relax.expand_dims":
            return x.unsqueeze(1)
        if op == "relax.squeeze":
            return x.squeeze(0)
        if op == "relax.reshape":
            return x.reshape(4, 1)
        if op == "relax.permute_dims":
            return x.permute(1, 0)
        if op == "relax.memory.view":
            return x.reshape(4, 1)
        if op == "relax.memory.ensure_zero_offset":
            return x
        if op == "relax.flatten":
            return x.flatten()
        if op == "relax.nn.batch_flatten":
            # TVM: ndim==2 input keeps shape (1, 4).
            return x
        raise ValueError(op)

    @classmethod
    def _expected_for_view_op(cls, op):
        x = torch.from_numpy(np.asarray(cls._input_for_view_op(op), dtype=np.float32))
        a = cls._torch_duplicate_view(x, op)
        b = cls._torch_duplicate_view(x, op)
        prod = a * b
        concat_axis = cls._concat_axis_for_view_op(op)
        return torch.cat([prod, b], dim=concat_axis).numpy()

    @pytest.mark.parametrize(
        "view_op",
        (
            # Keep this list in sync with IsViewMemoryOp() in
            # src/relax/transform/dataflow_inplace.cc
            "relax.expand_dims",
            "relax.squeeze",
            "relax.reshape",
            "relax.permute_dims",
            "relax.flatten",
            "relax.nn.batch_flatten",
            "relax.memory.view",
            "relax.memory.ensure_zero_offset",
        ),
    )
    def test_no_inplace_when_view_ops_share_input(self, view_op):
        mod = self._build_module(view_op)
        func = mod["main"]
        block = func.body.blocks[0]
        params = list(func.params)

        alias_sets, _ = dataflow_alias_analysis(block, params)
        view_vars = [
            binding.var
            for binding in block.bindings
            if (
                isinstance(binding.value, relax.Call)
                and isinstance(binding.value.op, tvm.ir.Op)
                and binding.value.op.name == view_op
            )
        ]
        a_var, b_var = view_vars[:2]
        assert alias_sets[a_var] & alias_sets[b_var], (
            f"{view_op}: duplicate views should share alias sets, but got "
            f"{alias_sets[a_var]} and {alias_sets[b_var]}"
        )

        _, exact_match = dataflow_inplace_analysis(block, params, mod)
        assert exact_match == [], f"{view_op}: expected no in-place opportunities"

        x_np = self._input_for_view_op(view_op).copy()
        mod_inplace = DataflowUseInplaceCalls()(mod)
        tvm.ir.assert_structural_equal(mod_inplace, mod)

        storage_shared = tvm.get_global_func("runtime.TVMTensorIsStorageShared")
        captures = self._capture_op_tensors(mod_inplace, x_np, "multiply")
        assert captures, f"{view_op}: VM instrumentation did not see a multiply call."
        cap = next(c for c in captures if len(c["tensors"]) >= 2)
        tensor_a, tensor_b = cap["tensors"][0], cap["tensors"][1]
        assert storage_shared(tensor_a, tensor_b), (
            f"{view_op}: multiply operands should share the same storage at runtime "
            f"(call {cap['call_name']!r})"
        )

        ex = relax.build(mod_inplace, tvm.target.Target("llvm"))
        vm = relax.VirtualMachine(ex, tvm.cpu())
        out = vm["main"](tvm.runtime.tensor(x_np, tvm.cpu()))
        np.testing.assert_allclose(out.numpy(), self._expected_for_view_op(view_op))


if __name__ == "__main__":
    testing.main()
