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

import pytest
import tvm_ffi

import tvm
import tvm.testing
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T


@pytest.fixture(autouse=True)
def enable_cuda_graph():
    """Enable cuda graph transform for all tests in this file"""
    with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
        yield


def test_rewrite_cuda_graph():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def exp(rxplaceholder: T.Buffer((T.int64(2), T.int64(4)), "float32"), compute: T.Buffer((T.int64(2), T.int64(4)), "float32")):
            # function attr dict
            T.func_attr({"tirx.noalias": True, "global_symbol": "exp"})
            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
                for i0_i1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
                    with T.sblock("compute"):
                        i0 = T.axis.spatial(T.int64(2), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) // T.int64(4))
                        i1 = T.axis.spatial(T.int64(4), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) % T.int64(4))
                        compute[i0, i1] = T.exp(rxplaceholder[i0, i1], dtype="float32")


        @R.function
        def main(x: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((2,4), dtype="float32"):
            # force_pure is expected because purity checking should be disabled before this pass
            R.func_attr({"relax.force_pure": True})
            cls = Before
            storage: R.Any = R.memory.alloc_storage(R.shape([32]), 0, "global", "float32")
            alloc: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, 0, R.shape([2, 4]), "float32")
            _1: R.Tuple = cls.exp(x, alloc)
            storage1: R.Any = R.memory.alloc_storage(R.shape([32]), 0, "global", "float32")
            alloc1: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage1, 0, R.shape([2, 4]), "float32")
            _2: R.Tuple = cls.exp(alloc, alloc1)
            _3: R.Tuple = R.memory.kill_tensor(alloc)
            alloc2: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, 0, R.shape([2, 4]), "float32")
            _4: R.Tuple = cls.exp(alloc1, alloc2)
            _5: R.Tuple = R.memory.kill_tensor(alloc1)
            storage2: R.Any = R.memory.alloc_storage(R.shape([32]), 0, "global", "float32")
            alloc3: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage2, 0, R.shape([2, 4]), "float32")
            _6: R.Tuple = cls.exp(alloc2, alloc3)
            _7: R.Tuple = R.memory.kill_tensor(alloc2)
            alloc4: R.Tensor((2, 4), dtype="float32") = R.builtin.alloc_tensor(R.shape([2, 4]), "float32", 0)
            _8 = cls.exp(alloc3, alloc4)
            _9: R.Tuple = R.memory.kill_tensor(alloc3)
            _10: R.Tuple = R.memory.kill_storage(storage)
            _11: R.Tuple = R.memory.kill_storage(storage1)
            _12: R.Tuple = R.memory.kill_storage(storage2)
            return alloc4


    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def exp(rxplaceholder: T.Buffer((T.int64(2), T.int64(4)), "float32"), compute: T.Buffer((T.int64(2), T.int64(4)), "float32")):
            # function attr dict
            T.func_attr({"tirx.noalias": True, "global_symbol": "exp"})
            # body
            # with T.sblock("root")
            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
                for i0_i1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
                    with T.sblock("compute"):
                        i0 = T.axis.spatial(T.int64(2), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) // T.int64(4))
                        i1 = T.axis.spatial(T.int64(4), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) % T.int64(4))
                        T.reads(rxplaceholder[i0, i1])
                        T.writes(compute[i0, i1])
                        compute[i0, i1] = T.exp(rxplaceholder[i0, i1], dtype="float32")

        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any, R.Any):
            R.func_attr({"relax.force_pure": True})
            storage: R.Any = R.memory.alloc_storage(R.shape([32]), R.prim_value(0), R.str("global"), R.dtype("float32"))
            storage1: R.Any = R.memory.alloc_storage(R.shape([32]), R.prim_value(0), R.str("global"), R.dtype("float32"))
            storage2: R.Any = R.memory.alloc_storage(R.shape([32]), R.prim_value(0), R.str("global"), R.dtype("float32"))
            gv: R.Tuple(R.Any, R.Any, R.Any) = (storage, storage1, storage2)
            return gv

        @R.function(private=True)
        def main_cuda_graph_capture(alloc: R.Tensor((2, 4), dtype="float32"), alloc1: R.Tensor((2, 4), dtype="float32"), storage: R.Any, storage2: R.Any) -> R.Tuple(R.Tensor((2, 4), dtype="float32")):
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            _2: R.Tuple = cls.exp(alloc, alloc1)
            _3: R.Tuple = R.memory.kill_tensor(alloc)
            alloc2: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            _4: R.Tuple = cls.exp(alloc1, alloc2)
            _5: R.Tuple = R.memory.kill_tensor(alloc1)
            alloc3: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage2, 0, R.shape([2, 4]), "float32")
            _6: R.Tuple = cls.exp(alloc2, alloc3)
            _7: R.Tuple = R.memory.kill_tensor(alloc2)
            gv: R.Tuple(R.Tensor((2, 4), dtype="float32")) = (alloc3,)
            return gv

        @R.function
        def main(x: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((2,4), dtype="float32"):
            # this comes after RemovePurityChecking, so we expect purity to be forced
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            gv: R.Tuple(R.Any, R.Any, R.Any) = R.call_builtin_with_ctx("vm.builtin.cuda_graph.get_cached_alloc", (cls.cuda_graph_alloc, R.prim_value(0)), ty_args=(R.Tuple(R.Any, R.Any, R.Any),))
            storage: R.Any = gv[0]
            alloc: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            _1: R.Tuple = cls.exp(x, alloc)
            storage1: R.Any = gv[1]
            alloc1: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            storage2: R.Any = gv[2]
            gv1: R.Tuple(R.Tensor((2, 4), dtype="float32")) = R.call_builtin_with_ctx("vm.builtin.cuda_graph.run_or_capture", (cls.main_cuda_graph_capture, (alloc, alloc1, storage, storage2), R.prim_value(0)), ty_args=(R.Tuple(R.Tensor((2, 4), dtype="float32")),))
            alloc3: R.Tensor((2, 4), dtype="float32") = gv1[0]
            alloc4: R.Tensor((2, 4), dtype="float32") = R.builtin.alloc_tensor(R.shape([2, 4]), R.dtype("float32"), R.prim_value(0))
            _6: R.Tuple = cls.exp(alloc3, alloc4)
            _7: R.Tuple = R.memory.kill_tensor(alloc3)
            _8: R.Tuple = R.memory.kill_storage(storage)
            _9: R.Tuple = R.memory.kill_storage(storage1)
            _10: R.Tuple = R.memory.kill_storage(storage2)
            return alloc4
    # fmt: on

    after = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(after, Expected)


def test_tuple():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def exp(rxplaceholder: T.Buffer((T.int64(2), T.int64(4)), "float32"), compute: T.Buffer((T.int64(2), T.int64(4)), "float32")):
            # function attr dict
            T.func_attr({"tirx.noalias": True, "global_symbol": "exp"})
            # body
            # with T.sblock("root")
            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
                for i0_i1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
                    with T.sblock("compute"):
                        i0 = T.axis.spatial(T.int64(2), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) // T.int64(4))
                        i1 = T.axis.spatial(T.int64(4), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) % T.int64(4))
                        T.reads(rxplaceholder[i0, i1])
                        T.writes(compute[i0, i1])
                        compute[i0, i1] = T.exp(rxplaceholder[i0, i1], dtype="float32")


        @R.function
        def main(x: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((2, 4), dtype="float32"):
            R.func_attr({"relax.force_pure": True})
            cls = Before
            storage: R.Any = R.memory.alloc_storage(R.shape([32]), 0, "global", "float32")
            alloc: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, 0, R.shape([2, 4]), "float32")
            _: R.Tuple = cls.exp(x, alloc)
            storage1: R.Any = R.memory.alloc_storage(R.shape([32]), 0, "global", "float32")
            alloc1: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage1, 0, R.shape([2, 4]), "float32")
            _: R.Tuple = cls.exp(alloc, alloc1)
            lv0 = (alloc1,)
            lv1 = (lv0,)
            lv2 = lv1[0]
            lv3 = lv2[0]
            alloc2: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, 0, R.shape([2, 4]), "float32")
            _1: R.Tuple = cls.exp(lv3, alloc2)
            _2: R.Tuple = R.memory.kill_tensor(alloc)
            _3: R.Tuple = R.memory.kill_tensor(alloc1)
            alloc3: R.Tensor((2, 4), dtype="float32") = R.builtin.alloc_tensor(R.shape([2, 4]), R.dtype("float32"), R.prim_value(0))
            _4: R.Tuple = cls.exp(alloc2, alloc3)
            _5: R.Tuple = R.memory.kill_tensor(alloc2)
            _6: R.Tuple = R.memory.kill_storage(storage)
            _7: R.Tuple = R.memory.kill_storage(storage1)
            return alloc3

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def exp(rxplaceholder: T.Buffer((T.int64(2), T.int64(4)), "float32"), compute: T.Buffer((T.int64(2), T.int64(4)), "float32")):
            T.func_attr({"global_symbol": "exp", "tirx.noalias": True})
            # with T.sblock("root"):
            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
                for i0_i1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
                    with T.sblock("compute"):
                        i0 = T.axis.spatial(T.int64(2), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) // T.int64(4))
                        i1 = T.axis.spatial(T.int64(4), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) % T.int64(4))
                        T.reads(rxplaceholder[i0, i1])
                        T.writes(compute[i0, i1])
                        compute[i0, i1] = T.exp(rxplaceholder[i0, i1])

        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any):
            R.func_attr({"relax.force_pure": True})
            storage: R.Any = R.memory.alloc_storage(R.shape([32]), R.prim_value(0), R.str("global"), R.dtype("float32"))
            storage1: R.Any = R.memory.alloc_storage(R.shape([32]), R.prim_value(0), R.str("global"), R.dtype("float32"))
            gv: R.Tuple(R.Any, R.Any) = (storage, storage1)
            return gv

        @R.function(private=True)
        def main_cuda_graph_capture(alloc: R.Tensor((2, 4), dtype="float32"), alloc1: R.Tensor((2, 4), dtype="float32"), storage: R.Any) -> R.Tuple(R.Tensor((2, 4), dtype="float32")):
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            _: R.Tuple = cls.exp(alloc, alloc1)
            lv0: R.Tuple(R.Tensor((2, 4), dtype="float32")) = (alloc1,)
            lv1: R.Tuple(R.Tuple(R.Tensor((2, 4), dtype="float32"))) = (lv0,)
            lv2: R.Tuple(R.Tensor((2, 4), dtype="float32")) = lv1[0]
            lv3: R.Tensor((2, 4), dtype="float32") = lv2[0]
            alloc2: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            _1: R.Tuple = cls.exp(lv3, alloc2)
            _2: R.Tuple = R.memory.kill_tensor(alloc)
            _3: R.Tuple = R.memory.kill_tensor(alloc1)
            gv: R.Tuple(R.Tensor((2, 4), dtype="float32")) = (alloc2,)
            return gv

        @R.function
        def main(x: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((2, 4), dtype="float32"):
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            gv: R.Tuple(R.Any, R.Any) = R.call_builtin_with_ctx("vm.builtin.cuda_graph.get_cached_alloc", (cls.cuda_graph_alloc, R.prim_value(0)), ty_args=(R.Tuple(R.Any, R.Any),))
            storage: R.Any = gv[0]
            alloc: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            _: R.Tuple = cls.exp(x, alloc)
            storage1: R.Any = gv[1]
            alloc1: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            gv1: R.Tuple(R.Tensor((2, 4), dtype="float32")) = R.call_builtin_with_ctx("vm.builtin.cuda_graph.run_or_capture", (cls.main_cuda_graph_capture, (alloc, alloc1, storage), R.prim_value(0)), ty_args=(R.Tuple(R.Tensor((2, 4), dtype="float32")),))
            alloc2: R.Tensor((2, 4), dtype="float32") = gv1[0]
            alloc3: R.Tensor((2, 4), dtype="float32") = R.builtin.alloc_tensor(R.shape([2, 4]), R.dtype("float32"), R.prim_value(0))
            _4: R.Tuple = cls.exp(alloc2, alloc3)
            _5: R.Tuple = R.memory.kill_tensor(alloc2)
            _6: R.Tuple = R.memory.kill_storage(storage)
            _7: R.Tuple = R.memory.kill_storage(storage1)
            return alloc3
    # fmt: on

    after = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(after, Expected)


def test_vm_builtin():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def exp(rxplaceholder: T.Buffer((T.int64(2), T.int64(4)), "float32"), compute: T.Buffer((T.int64(2), T.int64(4)), "float32")):
            # function attr dict
            T.func_attr({"tirx.noalias": True, "global_symbol": "exp"})
            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
                for i0_i1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
                    with T.sblock("compute"):
                        i0 = T.axis.spatial(T.int64(2), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) // T.int64(4))
                        i1 = T.axis.spatial(T.int64(4), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) % T.int64(4))
                        compute[i0, i1] = T.exp(rxplaceholder[i0, i1], dtype="float32")


        @R.function
        def main(x: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((2,4), dtype="float32"):
            # force_pure is expected because purity checking should be disabled before this pass
            R.func_attr({"relax.force_pure": True})
            cls = Before
            storage: R.Any = R.memory.alloc_storage(R.shape([32]), 0, "global", "float32")
            alloc: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, 0, R.shape([2, 4]), "float32")
            _1: R.Tuple = cls.exp(x, alloc)
            storage1: R.Any = R.memory.alloc_storage(R.shape([32]), 0, "global", "float32")
            alloc1: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage1, 0, R.shape([2, 4]), "float32")
            _2: R.Tuple = cls.exp(alloc, alloc1)
            _3: R.Tuple = R.memory.kill_tensor(alloc)
            alloc2: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, 0, R.shape([2, 4]), "float32")
            lv: R.Tensor((2, 4), dtype="float32") = alloc2
            _4: R.Tuple = R.call_packed("vm.builtin.dummy", (x, lv), ty_args=R.Tuple())
            _5: R.Tuple = R.memory.kill_tensor(alloc1)
            alloc3: R.Tensor((2, 4), dtype="float32") = R.builtin.alloc_tensor(R.shape([2, 4]), "float32", 0)
            _6 = cls.exp(alloc2, alloc3)
            _7: R.Tuple = R.memory.kill_tensor(alloc2)
            _8: R.Tuple = R.memory.kill_storage(storage)
            return alloc3

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def exp(rxplaceholder: T.Buffer((T.int64(2), T.int64(4)), "float32"), compute: T.Buffer((T.int64(2), T.int64(4)), "float32")):
            T.func_attr({"global_symbol": "exp", "tirx.noalias": True})
            # with T.sblock("root"):
            for i0_i1_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
                for i0_i1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.x"):
                    with T.sblock("compute"):
                        i0 = T.axis.spatial(T.int64(2), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) // T.int64(4))
                        i1 = T.axis.spatial(T.int64(4), (i0_i1_fused_0 * T.int64(8) + i0_i1_fused_1) % T.int64(4))
                        T.reads(rxplaceholder[i0, i1])
                        T.writes(compute[i0, i1])
                        compute[i0, i1] = T.exp(rxplaceholder[i0, i1])

        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any):
            R.func_attr({"relax.force_pure": True})
            storage: R.Any = R.memory.alloc_storage(R.shape([32]), R.prim_value(0), R.str("global"), R.dtype("float32"))
            storage1: R.Any = R.memory.alloc_storage(R.shape([32]), R.prim_value(0), R.str("global"), R.dtype("float32"))
            gv: R.Tuple(R.Any, R.Any) = (storage, storage1)
            return gv

        @R.function(private=True)
        def main_cuda_graph_capture(alloc: R.Tensor((2, 4), dtype="float32"), alloc1: R.Tensor((2, 4), dtype="float32"), storage: R.Any) -> R.Tuple(R.Tensor((2, 4), dtype="float32"), R.Tensor((2, 4), dtype="float32")):
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            _2: R.Tuple = cls.exp(alloc, alloc1)
            _3: R.Tuple = R.memory.kill_tensor(alloc)
            alloc2: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            lv: R.Tensor((2, 4), dtype="float32") = alloc2
            gv: R.Tuple(R.Tensor((2, 4), dtype="float32"), R.Tensor((2, 4), dtype="float32")) = (lv, alloc2)
            return gv

        @R.function
        def main(x: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((2,4), dtype="float32"):
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            gv: R.Tuple(R.Any, R.Any) = R.call_builtin_with_ctx("vm.builtin.cuda_graph.get_cached_alloc", (cls.cuda_graph_alloc, R.prim_value(0)), ty_args=(R.Tuple(R.Any, R.Any),))
            storage: R.Any = gv[0]
            alloc: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            _1: R.Tuple = cls.exp(x, alloc)
            storage1: R.Any = gv[1]
            alloc1: R.Tensor((2, 4), dtype="float32") = R.memory.alloc_tensor(storage1, R.prim_value(0), R.shape([2, 4]), R.dtype("float32"))
            gv1: R.Tuple(R.Tensor((2, 4), dtype="float32"), R.Tensor((2, 4), dtype="float32")) = R.call_builtin_with_ctx("vm.builtin.cuda_graph.run_or_capture", (cls.main_cuda_graph_capture, (alloc, alloc1, storage), R.prim_value(0)), ty_args=(R.Tuple(R.Tensor((2, 4), dtype="float32"), R.Tensor((2, 4), dtype="float32")),))
            alloc2: R.Tensor((2, 4), dtype="float32") = gv1[1]
            lv: R.Tensor((2, 4), dtype="float32") = gv1[0]
            _4: R.Tuple = R.call_packed("vm.builtin.dummy", (x, lv), ty_args=(R.Tuple,))
            _5: R.Tuple = R.memory.kill_tensor(alloc1)
            alloc3: R.Tensor((2, 4), dtype="float32") = R.builtin.alloc_tensor(R.shape([2, 4]), R.dtype("float32"), R.prim_value(0))
            _6: R.Tuple = cls.exp(alloc2, alloc3)
            _7: R.Tuple = R.memory.kill_tensor(alloc2)
            _8: R.Tuple = R.memory.kill_storage(storage)
            return alloc3
    # fmt: on

    after = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(after, Expected)


def test_capture_fixed_inputs():
    @tvm.script.ir_module
    class Conv2dx3:
        @R.function
        def main(
            data: R.Tensor((16, 32, 32, 16), "float16"),
            weight1: R.Tensor((16, 3, 3, 16), "float16"),
            weight2: R.Tensor((16, 3, 3, 16), "float16"),
            weight3: R.Tensor((16, 3, 3, 16), "float16"),
            gamma: R.Tensor((16,), "float16"),
            beta: R.Tensor((16,), "float16"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                conv1 = R.nn.relu(
                    R.nn.conv2d(
                        data, weight1, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI"
                    )
                )

                ###############################################################################
                # The second conv2d and layer norm can be captured into a graph
                conv2 = R.nn.relu(
                    R.nn.conv2d(
                        conv1, weight2, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI"
                    )
                )
                ln = R.nn.layer_norm(conv2, gamma, beta, axes=[-1])
                ###############################################################################

                conv3 = R.nn.relu(
                    R.nn.conv2d(
                        ln, weight3, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI"
                    )
                )
                R.output(conv3)

            return conv3

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def fused_conv2d_relu(
            data: T.Buffer((T.int64(16), T.int64(32), T.int64(32), T.int64(16)), "float16"),
            weight1: T.Buffer((T.int64(16), T.int64(3), T.int64(3), T.int64(16)), "float16"),
            var_compute_intermediate: T.Buffer(
                (T.int64(16), T.int64(32), T.int64(32), T.int64(16)), "float16"
            ),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            pad_temp = T.sblock_alloc_buffer(
                (T.int64(16), T.int64(34), T.int64(34), T.int64(16)), "float16"
            )
            var_conv2d_nhwc_intermediate = T.sblock_alloc_buffer(
                (T.int64(16), T.int64(32), T.int64(32), T.int64(16)), "float16"
            )
            for i0, i1, i2, i3 in T.grid(T.int64(16), T.int64(34), T.int64(34), T.int64(16)):
                with T.sblock("pad_temp"):
                    v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(data[v_i0, v_i1 - T.int64(1), v_i2 - T.int64(1), v_i3])
                    T.writes(pad_temp[v_i0, v_i1, v_i2, v_i3])
                    pad_temp[v_i0, v_i1, v_i2, v_i3] = T.if_then_else(
                        T.int64(1) <= v_i1
                        and v_i1 < T.int64(33)
                        and T.int64(1) <= v_i2
                        and v_i2 < T.int64(33),
                        data[v_i0, v_i1 - T.int64(1), v_i2 - T.int64(1), v_i3],
                        T.float16(0),
                    )
            for nn, yy, xx, ff, ry, rx, rc in T.grid(
                T.int64(16),
                T.int64(32),
                T.int64(32),
                T.int64(16),
                T.int64(3),
                T.int64(3),
                T.int64(16),
            ):
                with T.sblock("conv2d_nhwc"):
                    v_nn, v_yy, v_xx, v_ff, v_ry, v_rx, v_rc = T.axis.remap(
                        "SSSSRRR", [nn, yy, xx, ff, ry, rx, rc]
                    )
                    T.reads(
                        pad_temp[v_nn, v_yy + v_ry, v_xx + v_rx, v_rc],
                        weight1[v_ff, v_ry, v_rx, v_rc],
                    )
                    T.writes(var_conv2d_nhwc_intermediate[v_nn, v_yy, v_xx, v_ff])
                    with T.init():
                        var_conv2d_nhwc_intermediate[v_nn, v_yy, v_xx, v_ff] = T.float16(0)
                    var_conv2d_nhwc_intermediate[v_nn, v_yy, v_xx, v_ff] = (
                        var_conv2d_nhwc_intermediate[v_nn, v_yy, v_xx, v_ff]
                        + pad_temp[v_nn, v_yy + v_ry, v_xx + v_rx, v_rc]
                        * weight1[v_ff, v_ry, v_rx, v_rc]
                    )
            for i0, i1, i2, i3 in T.grid(T.int64(16), T.int64(32), T.int64(32), T.int64(16)):
                with T.sblock("compute"):
                    v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(var_conv2d_nhwc_intermediate[v_i0, v_i1, v_i2, v_i3])
                    T.writes(var_compute_intermediate[v_i0, v_i1, v_i2, v_i3])
                    var_compute_intermediate[v_i0, v_i1, v_i2, v_i3] = T.max(
                        var_conv2d_nhwc_intermediate[v_i0, v_i1, v_i2, v_i3], T.float16(0)
                    )

        @T.prim_func(s_tir=True)
        def layer_norm(
            A: T.Buffer((T.int64(16), T.int64(32), T.int64(32), T.int64(16)), "float16"),
            B: T.Buffer((T.int64(16),), "float16"),
            C: T.Buffer((T.int64(16),), "float16"),
            T_layer_norm: T.Buffer((T.int64(16), T.int64(32), T.int64(32), T.int64(16)), "float16"),
        ):
            T.func_attr({"op_pattern": 4, "tirx.noalias": True})
            # with T.sblock("root"):
            A_red_temp_v0 = T.sblock_alloc_buffer((T.int64(16), T.int64(32), T.int64(32)))
            A_red_temp_v1 = T.sblock_alloc_buffer((T.int64(16), T.int64(32), T.int64(32)))
            for ax0, ax1, ax2, k3 in T.grid(T.int64(16), T.int64(32), T.int64(32), T.int64(16)):
                with T.sblock("A_red_temp"):
                    v_ax0, v_ax1, v_ax2, v_k3 = T.axis.remap("SSSR", [ax0, ax1, ax2, k3])
                    T.reads(A[v_ax0, v_ax1, v_ax2, v_k3])
                    T.writes(A_red_temp_v0[v_ax0, v_ax1, v_ax2], A_red_temp_v1[v_ax0, v_ax1, v_ax2])
                    with T.init():
                        A_red_temp_v0[v_ax0, v_ax1, v_ax2] = T.float32(0)
                        A_red_temp_v1[v_ax0, v_ax1, v_ax2] = T.float32(0)
                    v_A_red_temp_v0: T.let[T.float32] = A_red_temp_v0[v_ax0, v_ax1, v_ax2] + T.Cast(
                        "float32", A[v_ax0, v_ax1, v_ax2, v_k3]
                    )
                    v_A_red_temp_v1: T.let[T.float32] = A_red_temp_v1[v_ax0, v_ax1, v_ax2] + T.Cast(
                        "float32", A[v_ax0, v_ax1, v_ax2, v_k3]
                    ) * T.Cast("float32", A[v_ax0, v_ax1, v_ax2, v_k3])
                    A_red_temp_v0[v_ax0, v_ax1, v_ax2] = v_A_red_temp_v0
                    A_red_temp_v1[v_ax0, v_ax1, v_ax2] = v_A_red_temp_v1
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(16), T.int64(32), T.int64(32), T.int64(16)):
                with T.sblock("T_layer_norm"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(
                        A[v_ax0, v_ax1, v_ax2, v_ax3],
                        A_red_temp_v0[v_ax0, v_ax1, v_ax2],
                        A_red_temp_v1[v_ax0, v_ax1, v_ax2],
                        B[v_ax3],
                        C[v_ax3],
                    )
                    T.writes(T_layer_norm[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_layer_norm[v_ax0, v_ax1, v_ax2, v_ax3] = (
                        T.Cast(
                            "float16",
                            (
                                T.Cast("float32", A[v_ax0, v_ax1, v_ax2, v_ax3])
                                - A_red_temp_v0[v_ax0, v_ax1, v_ax2] * T.float32(0.0625)
                            )
                            * T.rsqrt(
                                A_red_temp_v1[v_ax0, v_ax1, v_ax2] * T.float32(0.0625)
                                - A_red_temp_v0[v_ax0, v_ax1, v_ax2]
                                * T.float32(0.0625)
                                * (A_red_temp_v0[v_ax0, v_ax1, v_ax2] * T.float32(0.0625))
                                + T.float32(1.0000000000000001e-05)
                            ),
                        )
                        * B[v_ax3]
                        + C[v_ax3]
                    )

        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any):
            R.func_attr({"relax.force_pure": True})
            storage: R.Any = R.memory.alloc_storage(
                R.shape([524288]), R.prim_value(0), R.str("global"), R.dtype("float16")
            )
            storage1: R.Any = R.memory.alloc_storage(
                R.shape([524288]), R.prim_value(0), R.str("global"), R.dtype("float16")
            )
            gv: R.Tuple(R.Any, R.Any) = storage, storage1
            return gv

        @R.function(private=True)
        def main_cuda_graph_capture(
            lv: R.Tensor((16, 32, 32, 16), dtype="float16"),
            lv1: R.Tensor((16, 3, 3, 16), dtype="float16"),
            alloc1: R.Tensor((16, 32, 32, 16), dtype="float16"),
            alloc: R.Tensor((16, 32, 32, 16), dtype="float16"),
            params: R.Tuple(
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16,), dtype="float16"),
                R.Tensor((16,), dtype="float16"),
            ),
            storage: R.Any,
        ) -> R.Tuple(
            R.Tensor((16, 32, 32, 16), dtype="float16"),
            R.Tensor((16, 3, 3, 16), dtype="float16"),
            R.Tensor((16, 32, 32, 16), dtype="float16"),
        ):
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            _1: R.Tuple = cls.fused_conv2d_relu(lv, lv1, alloc1)
            _: R.Tuple = R.memory.kill_tensor(alloc)
            lv1_1: R.Tensor((16, 32, 32, 16), dtype="float16") = alloc1
            lv2: R.Tensor((16,), dtype="float16") = params[3]
            lv3: R.Tensor((16,), dtype="float16") = params[4]
            alloc2: R.Tensor((16, 32, 32, 16), dtype="float16") = R.memory.alloc_tensor(
                storage, R.prim_value(0), R.shape([16, 32, 32, 16]), R.dtype("float16")
            )
            _2: R.Tuple = cls.layer_norm(lv1_1, lv2, lv3, alloc2)
            _1_1: R.Tuple = R.memory.kill_tensor(alloc1)
            ln: R.Tensor((16, 32, 32, 16), dtype="float16") = alloc2
            lv4: R.Tensor((16, 3, 3, 16), dtype="float16") = params[2]
            gv: R.Tuple(
                R.Tensor((16, 32, 32, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 32, 32, 16), dtype="float16"),
            ) = (ln, lv4, alloc2)
            return gv

        @R.function
        def main_transform_params(
            params: R.Tuple(
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16,), dtype="float16"),
                R.Tensor((16,), dtype="float16"),
            ),
        ) -> R.Tuple(
            R.Tensor((16, 3, 3, 16), dtype="float16"),
            R.Tensor((16, 3, 3, 16), dtype="float16"),
            R.Tensor((16, 3, 3, 16), dtype="float16"),
            R.Tensor((16,), dtype="float16"),
            R.Tensor((16,), dtype="float16"),
        ):
            R.func_attr({"relax.force_pure": True})
            lv: R.Tensor((16, 3, 3, 16), dtype="float16") = params[0]
            lv1: R.Tensor((16, 3, 3, 16), dtype="float16") = params[1]
            lv2: R.Tensor((16, 3, 3, 16), dtype="float16") = params[2]
            lv3: R.Tensor((16,), dtype="float16") = params[3]
            lv4: R.Tensor((16,), dtype="float16") = params[4]
            gv: R.Tuple(
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16,), dtype="float16"),
                R.Tensor((16,), dtype="float16"),
            ) = (lv, lv1, lv2, lv3, lv4)
            return gv

        @R.function
        def main(
            data: R.Tensor((16, 32, 32, 16), dtype="float16"),
            params: R.Tuple(
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16,), dtype="float16"),
                R.Tensor((16,), dtype="float16"),
            ),
        ) -> R.Tensor((16, 32, 32, 16), dtype="float16"):
            R.func_attr({"num_input": 1, "relax.force_pure": True})
            cls = Expected
            lv: R.Tensor((16, 3, 3, 16), dtype="float16") = params[0]
            gv: R.Tuple(R.Any, R.Any) = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.get_cached_alloc",
                (cls.cuda_graph_alloc, R.prim_value(0)),
                ty_args=(R.Tuple(R.Any, R.Any),),
            )
            storage: R.Any = gv[0]
            alloc: R.Tensor((16, 32, 32, 16), dtype="float16") = R.memory.alloc_tensor(
                storage, R.prim_value(0), R.shape([16, 32, 32, 16]), R.dtype("float16")
            )
            _: R.Tuple = cls.fused_conv2d_relu(data, lv, alloc)
            lv_1: R.Tensor((16, 32, 32, 16), dtype="float16") = alloc
            lv1: R.Tensor((16, 3, 3, 16), dtype="float16") = params[1]
            storage1: R.Any = gv[1]
            alloc1: R.Tensor((16, 32, 32, 16), dtype="float16") = R.memory.alloc_tensor(
                storage1, R.prim_value(0), R.shape([16, 32, 32, 16]), R.dtype("float16")
            )
            gv1: R.Tuple(
                R.Tensor((16, 32, 32, 16), dtype="float16"),
                R.Tensor((16, 3, 3, 16), dtype="float16"),
                R.Tensor((16, 32, 32, 16), dtype="float16"),
            ) = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.run_or_capture",
                (
                    cls.main_cuda_graph_capture,
                    (lv_1, lv1, alloc1, alloc, params, storage),
                    R.prim_value(0),
                ),
                ty_args=(
                    R.Tuple(
                        R.Tensor((16, 32, 32, 16), dtype="float16"),
                        R.Tensor((16, 3, 3, 16), dtype="float16"),
                        R.Tensor((16, 32, 32, 16), dtype="float16"),
                    ),
                ),
            )
            alloc2: R.Tensor((16, 32, 32, 16), dtype="float16") = gv1[2]
            ln: R.Tensor((16, 32, 32, 16), dtype="float16") = gv1[0]
            lv4: R.Tensor((16, 3, 3, 16), dtype="float16") = gv1[1]
            alloc3: R.Tensor((16, 32, 32, 16), dtype="float16") = R.builtin.alloc_tensor(
                R.shape([16, 32, 32, 16]), R.dtype("float16"), R.prim_value(0)
            )
            _3: R.Tuple = cls.fused_conv2d_relu(ln, lv4, alloc3)
            _2: R.Tuple = R.memory.kill_tensor(alloc2)
            gv_1: R.Tensor((16, 32, 32, 16), dtype="float16") = alloc3
            _3_1: R.Tuple = R.memory.kill_storage(storage)
            _4: R.Tuple = R.memory.kill_storage(storage1)
            return gv_1

    mod = tvm.transform.Sequential(
        [
            relax.pipeline.get_pipeline(),
            relax.transform.LiftTransformParams(),
            relax.transform.ToNonDataflow(),
            relax.transform.RemovePurityChecking(),
            relax.transform.CallTIRRewrite(),
            relax.transform.StaticPlanBlockMemory(),
        ]
    )(Conv2dx3)

    mod["main"] = mod["main"].with_attr({"num_input": 1})
    after = relax.transform.RewriteCUDAGraph()(mod)
    tvm.ir.assert_structural_equal(after, after)


def test_null_value():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main() -> R.Tuple(R.Any):
            _io: R.Any = R.null_value()
            lv: R.Tuple(R.Any) = (_io,)
            gv: R.Tuple(R.Any) = lv
            return gv

    Expected = Before

    After = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(After, Expected)


def test_transform_is_no_op_when_disabled():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main():
            storage = R.memory.alloc_storage(R.shape([8]), 0, "global", "float32")
            alloc3 = R.memory.alloc_tensor(storage, 0, R.shape([8]), "float32")
            return R.tuple()

    with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
        AfterWhenEnabled = relax.transform.RewriteCUDAGraph()(Before)
    with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": False}):
        AfterWhenDisabled = relax.transform.RewriteCUDAGraph()(Before)

    assert not tvm_ffi.structural_equal(Before, AfterWhenEnabled)
    tvm.ir.assert_structural_equal(Before, AfterWhenDisabled)


def test_static_args():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function(pure=False)
        def main():
            storage0 = R.memory.alloc_storage(R.shape([8]), 0, "global", "float32")
            alloc0 = R.memory.alloc_tensor(storage0, 0, R.shape([8]), "float32")
            _ = R.call_packed("dummy_func", alloc0, R.dtype("float32"), R.str("string"))
            return R.tuple()

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any):
            R.func_attr({"relax.force_pure": True})
            storage0: R.Any = R.memory.alloc_storage(
                R.shape([8]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            gv: R.Tuple(R.Any) = (storage0,)
            return gv

        @R.function(private=True)
        def main_cuda_graph_capture(alloc0: R.Tensor((8,), dtype="float32")) -> R.Tuple:
            R.func_attr({"relax.force_pure": True})
            _: R.Any = R.call_packed("dummy_func", alloc0, R.dtype("float32"), R.str("string"))
            gv: R.Tuple = R.tuple()
            return gv

        @R.function(pure=False)
        def main() -> R.Tuple:
            cls = Expected
            gv: R.Tuple(R.Any) = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.get_cached_alloc",
                (cls.cuda_graph_alloc, R.prim_value(0)),
                ty_args=(R.Tuple(R.Any),),
            )
            storage0: R.Any = gv[0]
            alloc0: R.Tensor((8,), dtype="float32") = R.memory.alloc_tensor(
                storage0, R.prim_value(0), R.shape([8]), R.dtype("float32")
            )
            gv1: R.Tuple = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.run_or_capture",
                (cls.main_cuda_graph_capture, (alloc0,), R.prim_value(0)),
                ty_args=(R.Tuple,),
            )
            return R.tuple()

    mod = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_dynamic_capture():
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(s_tir=True)
        def add_one(x_handle: T.handle, y_handle: T.handle):
            m = T.int64()
            x = T.match_buffer(x_handle, (m,), "float32")
            y = T.match_buffer(y_handle, (m,), "float32")
            # Use T.serial with explicit int64 min so the inner sblock iter_var
            # dom is all-int64 (matches what Expected emits via T.axis.spatial(m, i)).
            for i in T.serial(T.int64(0), m):
                with T.sblock("add"):
                    vi = T.axis.remap("S", [i])
                    y[vi] = x[vi] + T.float32(1)

        @R.function
        def main(x: R.Tensor(("m",), "float32")) -> R.Tensor(("m",), "float32"):
            R.func_attr(
                {"relax.rewrite_cuda_graph.capture_symbolic_vars": ["m"], "relax.force_pure": True}
            )
            m = T.int64()
            storage: R.Any = R.memory.alloc_storage(
                R.shape([16]), 0, "global", "float32"
            )  # assume m is upper-bounded
            alloc1: R.Tensor((m,), "float32") = R.memory.alloc_tensor(
                storage, 0, R.shape([m]), "float32"
            )
            _ = Before.add_one(x, alloc1)
            storage1: R.Any = R.memory.alloc_storage(R.shape([16]), 0, "global", "float32")
            alloc2: R.Tensor((m,), "float32") = R.memory.alloc_tensor(
                storage1, 0, R.shape([m]), "float32"
            )
            _ = Before.add_one(alloc1, alloc2)
            alloc3: R.Tensor((m,), "float32") = R.builtin.alloc_tensor(
                R.shape([m]), "float32", 0, "global"
            )
            _ = Before.add_one(alloc2, alloc3)
            return alloc3

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(s_tir=True)
        def add_one(x_handle: T.handle, y_handle: T.handle):
            m = T.int64()
            x = T.match_buffer(x_handle, (m,))
            y = T.match_buffer(y_handle, (m,))
            # with T.sblock("root"):
            for i in T.serial(T.int64(0), m):
                with T.sblock("add"):
                    vi = T.axis.spatial(m, i)
                    T.reads(x[vi])
                    T.writes(y[vi])
                    y[vi] = x[vi] + T.float32(1)

        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any):
            R.func_attr({"relax.force_pure": True})
            storage: R.Any = R.memory.alloc_storage(
                R.shape([16]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            storage1: R.Any = R.memory.alloc_storage(
                R.shape([16]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            gv: R.Tuple(R.Any, R.Any) = storage, storage1
            return gv

        @R.function(private=True)
        def main_cuda_graph_capture(
            alloc1: R.Tensor(("m",), dtype="float32"),
            alloc2: R.Tensor(("m",), dtype="float32"),
            shape_expr: R.Shape(["m"]),
        ):
            m = T.int64()
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            cls.add_one(alloc1, alloc2)
            gv = R.tuple()
            return R.tuple()

        @R.function
        def main(x: R.Tensor(("m",), dtype="float32")) -> R.Tensor(("m",), dtype="float32"):
            m = T.int64()
            R.func_attr(
                {"relax.force_pure": True, "relax.rewrite_cuda_graph.capture_symbolic_vars": ["m"]}
            )
            cls = Expected
            gv: R.Tuple(R.Any, R.Any) = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.get_cached_alloc",
                (cls.cuda_graph_alloc, R.prim_value(0)),
                ty_args=(R.Tuple(R.Any, R.Any),),
            )
            storage: R.Any = gv[0]
            alloc1: R.Tensor((m,), dtype="float32") = R.memory.alloc_tensor(
                storage, R.prim_value(0), R.shape([m]), R.dtype("float32")
            )
            cls.add_one(x, alloc1)
            storage1: R.Any = gv[1]
            alloc2: R.Tensor((m,), dtype="float32") = R.memory.alloc_tensor(
                storage1, R.prim_value(0), R.shape([m]), R.dtype("float32")
            )
            R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.run_or_capture",
                (
                    cls.main_cuda_graph_capture,
                    (alloc1, alloc2, R.shape([m])),
                    R.prim_value(0),
                    R.shape([m]),
                ),
                ty_args=(R.Tuple,),
            )
            alloc3: R.Tensor((m,), dtype="float32") = R.builtin.alloc_tensor(
                R.shape([m]), R.dtype("float32"), R.prim_value(0), R.str("global")
            )
            cls.add_one(alloc2, alloc3)
            return alloc3

    mod = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_merge_alloc_funcs():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def func1():
            R.func_attr({"relax.force_pure": True})
            storage1 = R.memory.alloc_storage(R.shape([128]), 0, "global", "float32")
            storage2 = R.memory.alloc_storage(R.shape([256]), 0, "global", "float32")
            storage3 = R.memory.alloc_storage(R.shape([512]), 0, "ipc_memory", "float32")
            alloc1 = R.memory.alloc_tensor(storage1, 0, R.shape([128]), "float32")
            alloc2 = R.memory.alloc_tensor(storage2, 0, R.shape([256]), "float32")
            alloc3 = R.memory.alloc_tensor(storage3, 0, R.shape([512]), "float32")
            R.call_packed("dummy", alloc1, alloc2, alloc3, ty_args=(R.Tuple,))
            return R.tuple()

        @R.function
        def func2():
            R.func_attr({"relax.force_pure": True})
            storage1 = R.memory.alloc_storage(R.shape([192]), 0, "global", "float32")
            storage2 = R.memory.alloc_storage(R.shape([64]), 0, "global", "float32")
            storage3 = R.memory.alloc_storage(R.shape([1024]), 0, "ipc_memory", "float32")
            storage4 = R.memory.alloc_storage(R.shape([512]), 0, "global", "float32")
            alloc1 = R.memory.alloc_tensor(storage1, 0, R.shape([192]), "float32")
            alloc2 = R.memory.alloc_tensor(storage2, 0, R.shape([64]), "float32")
            alloc3 = R.memory.alloc_tensor(storage3, 0, R.shape([1024]), "float32")
            alloc4 = R.memory.alloc_tensor(storage4, 0, R.shape([512]), "float32")
            R.call_packed("dummy", alloc1, alloc2, alloc3, alloc4, ty_args=(R.Tuple,))
            return R.tuple()

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any, R.Any, R.Any):
            R.func_attr({"relax.force_pure": True})
            storage4: R.Any = R.memory.alloc_storage(
                R.shape([512]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            storage1: R.Any = R.memory.alloc_storage(
                R.shape([192]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            storage2: R.Any = R.memory.alloc_storage(
                R.shape([64]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            storage3: R.Any = R.memory.alloc_storage(
                R.shape([1024]), R.prim_value(0), R.str("ipc_memory"), R.dtype("float32")
            )
            gv: R.Tuple(R.Any, R.Any, R.Any, R.Any) = (
                storage4,
                storage1,
                storage2,
                storage3,
            )
            return gv

        @R.function
        def func1() -> R.Tuple:
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            gv: R.Tuple(R.Any, R.Any, R.Any, R.Any) = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.get_cached_alloc",
                (cls.cuda_graph_alloc, R.prim_value(0)),
                ty_args=(R.Tuple(R.Any, R.Any, R.Any, R.Any),),
            )
            storage1: R.Any = gv[1]
            storage2: R.Any = gv[0]
            storage3: R.Any = gv[3]
            alloc1: R.Tensor((128,), dtype="float32") = R.memory.alloc_tensor(
                storage1, R.prim_value(0), R.shape([128]), R.dtype("float32")
            )
            alloc2: R.Tensor((256,), dtype="float32") = R.memory.alloc_tensor(
                storage2, R.prim_value(0), R.shape([256]), R.dtype("float32")
            )
            alloc3: R.Tensor((512,), dtype="float32") = R.memory.alloc_tensor(
                storage3, R.prim_value(0), R.shape([512]), R.dtype("float32")
            )
            R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.run_or_capture",
                (cls.func1_cuda_graph_capture, (alloc1, alloc2, alloc3), R.prim_value(0)),
                ty_args=(R.Tuple,),
            )
            return R.tuple()

        @R.function(private=True)
        def func1_cuda_graph_capture(
            alloc1: R.Tensor((128,), dtype="float32"),
            alloc2: R.Tensor((256,), dtype="float32"),
            alloc3: R.Tensor((512,), dtype="float32"),
        ) -> R.Tuple:
            R.func_attr({"relax.force_pure": True})
            R.call_packed("dummy", alloc1, alloc2, alloc3, ty_args=(R.Tuple,))
            R.tuple()
            return R.tuple()

        @R.function
        def func2() -> R.Tuple:
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            gv2: R.Tuple(R.Any, R.Any, R.Any, R.Any) = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.get_cached_alloc",
                (cls.cuda_graph_alloc, R.prim_value(0)),
                ty_args=(R.Tuple(R.Any, R.Any, R.Any, R.Any),),
            )
            storage11: R.Any = gv2[1]
            storage21: R.Any = gv2[2]
            storage31: R.Any = gv2[3]
            storage4: R.Any = gv2[0]
            alloc1: R.Tensor((192,), dtype="float32") = R.memory.alloc_tensor(
                storage11, R.prim_value(0), R.shape([192]), R.dtype("float32")
            )
            alloc2: R.Tensor((64,), dtype="float32") = R.memory.alloc_tensor(
                storage21, R.prim_value(0), R.shape([64]), R.dtype("float32")
            )
            alloc3: R.Tensor((1024,), dtype="float32") = R.memory.alloc_tensor(
                storage31, R.prim_value(0), R.shape([1024]), R.dtype("float32")
            )
            alloc4: R.Tensor((512,), dtype="float32") = R.memory.alloc_tensor(
                storage4, R.prim_value(0), R.shape([512]), R.dtype("float32")
            )
            R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.run_or_capture",
                (cls.func2_cuda_graph_capture, (alloc1, alloc2, alloc3, alloc4), R.prim_value(1)),
                ty_args=(R.Tuple,),
            )
            return R.tuple()

        @R.function(private=True)
        def func2_cuda_graph_capture(
            alloc1: R.Tensor((192,), dtype="float32"),
            alloc2: R.Tensor((64,), dtype="float32"),
            alloc3: R.Tensor((1024,), dtype="float32"),
            alloc4: R.Tensor((512,), dtype="float32"),
        ) -> R.Tuple:
            R.func_attr({"relax.force_pure": True})
            R.call_packed("dummy", alloc1, alloc2, alloc3, alloc4, ty_args=(R.Tuple,))
            R.tuple()
            return R.tuple()

    After = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(After, Expected)


def test_disable_capture_output():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor((8,), "float32")) -> R.Tuple(R.Tensor((8,), "float32")):
            R.func_attr({"relax.force_pure": True})
            storage1 = R.memory.alloc_storage(R.shape([8]), 0, "global", "float32")
            alloc1 = R.memory.alloc_tensor(storage1, 0, R.shape([8]), "float32")
            _ = R.call_packed("dummy", x, alloc1, ty_args=(R.Tuple,))
            storage2 = R.memory.alloc_storage(R.shape([8]), 0, "global", "float32")
            alloc2 = R.memory.alloc_tensor(storage2, 0, R.shape([8]), "float32")
            _1 = R.call_packed("dummy", alloc1, alloc2, ty_args=(R.Tuple,))
            storage3 = R.memory.alloc_storage(R.shape([8]), 0, "global", "float32")
            alloc3 = R.memory.alloc_tensor(storage3, 0, R.shape([8]), "float32")
            _2 = R.call_packed("dummy", alloc2, alloc3, ty_args=(R.Tuple,))
            gv = (alloc3,)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any):
            R.func_attr({"relax.force_pure": True})
            storage1: R.Any = R.memory.alloc_storage(
                R.shape([8]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            storage2: R.Any = R.memory.alloc_storage(
                R.shape([8]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            gv: R.Tuple(R.Any, R.Any) = storage1, storage2
            return gv

        @R.function(private=True)
        def main_cuda_graph_capture(
            alloc1: R.Tensor((8,), dtype="float32"), alloc2: R.Tensor((8,), dtype="float32")
        ) -> R.Tuple:
            R.func_attr({"relax.force_pure": True})
            R.call_packed("dummy", alloc1, alloc2, ty_args=(R.Tuple,))
            R.tuple()
            return R.tuple()

        @R.function
        def main(x: R.Tensor((8,), dtype="float32")) -> R.Tuple(R.Tensor((8,), dtype="float32")):
            R.func_attr({"relax.force_pure": True})
            cls = Expected
            gv: R.Tuple(R.Any, R.Any) = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.get_cached_alloc",
                (cls.cuda_graph_alloc, R.prim_value(0)),
                ty_args=(R.Tuple(R.Any, R.Any),),
            )
            storage1: R.Any = gv[0]
            alloc1: R.Tensor((8,), dtype="float32") = R.memory.alloc_tensor(
                storage1, R.prim_value(0), R.shape([8]), R.dtype("float32")
            )
            R.call_packed("dummy", x, alloc1, ty_args=(R.Tuple,))
            storage2: R.Any = gv[1]
            alloc2: R.Tensor((8,), dtype="float32") = R.memory.alloc_tensor(
                storage2, R.prim_value(0), R.shape([8]), R.dtype("float32")
            )
            R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.run_or_capture",
                (cls.main_cuda_graph_capture, (alloc1, alloc2), R.prim_value(0)),
                ty_args=(R.Tuple,),
            )
            storage3: R.Any = R.memory.alloc_storage(
                R.shape([8]), R.prim_value(0), R.str("global"), R.dtype("float32")
            )
            alloc3: R.Tensor((8,), dtype="float32") = R.memory.alloc_tensor(
                storage3, R.prim_value(0), R.shape([8]), R.dtype("float32")
            )
            R.call_packed("dummy", alloc2, alloc3, ty_args=(R.Tuple,))
            gv = (alloc3,)
            return gv

    After = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(After, Expected)


def test_static_input_with_symbolic_shape():
    @I.ir_module(s_tir=True)
    class Before:
        @R.function
        def main(x: R.Tensor((8,), "float16"), w: R.Tensor(("m",))):
            m = T.int64()
            R.func_attr({"relax.force_pure": True, "num_input": 1})
            storage1 = R.memory.alloc_storage(R.shape([8]), 0, "global", "float16")
            alloc1 = R.memory.alloc_tensor(storage1, 0, R.shape([8]), "float16")
            _ = R.call_packed("dummy", x, w, alloc1, ty_args=(R.Tuple,))
            storage2 = R.memory.alloc_storage(R.shape([8]), 0, "global", "float16")
            alloc2 = R.memory.alloc_tensor(storage2, 0, R.shape([8]), "float16")
            _1 = R.call_packed("dummy", alloc1, w, alloc2, ty_args=(R.Tuple,))
            storage3 = R.memory.alloc_storage(R.shape([8]), 0, "global", "float16")
            alloc3 = R.memory.alloc_tensor(storage3, 0, R.shape([8]), "float16")
            _2 = R.call_packed("dummy", alloc2, w, alloc3, ty_args=(R.Tuple,))
            gv = (alloc3,)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @R.function(private=True)
        def cuda_graph_alloc() -> R.Tuple(R.Any, R.Any):
            R.func_attr({"relax.force_pure": True})
            storage1: R.Any = R.memory.alloc_storage(
                R.shape([8]), R.prim_value(0), R.str("global"), R.dtype("float16")
            )
            storage2: R.Any = R.memory.alloc_storage(
                R.shape([8]), R.prim_value(0), R.str("global"), R.dtype("float16")
            )
            gv: R.Tuple(R.Any, R.Any) = storage1, storage2
            return gv

        @R.function(private=True)
        def main_cuda_graph_capture(
            alloc1: R.Tensor((8,), dtype="float16"),
            w: R.Tensor(("m",)),
            alloc2: R.Tensor((8,), dtype="float16"),
            shape_expr: R.Shape(["m"]),
        ) -> R.Tuple:
            m = T.int64()
            R.func_attr({"relax.force_pure": True})
            R.call_packed("dummy", alloc1, w, alloc2, ty_args=(R.Tuple,))
            R.tuple()
            return R.tuple()

        @R.function
        def main(x: R.Tensor((8,), dtype="float16"), w: R.Tensor(("m",))) -> R.Tuple(
            R.Tensor((8,), dtype="float16")
        ):
            m = T.int64()
            R.func_attr({"num_input": 1, "relax.force_pure": True})
            cls = Expected
            gv: R.Tuple(R.Any, R.Any) = R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.get_cached_alloc",
                (cls.cuda_graph_alloc, R.prim_value(0)),
                ty_args=(R.Tuple(R.Any, R.Any),),
            )
            storage1: R.Any = gv[0]
            alloc1: R.Tensor((8,), dtype="float16") = R.memory.alloc_tensor(
                storage1, R.prim_value(0), R.shape([8]), R.dtype("float16")
            )
            R.call_packed("dummy", x, w, alloc1, ty_args=(R.Tuple,))
            storage2: R.Any = gv[1]
            alloc2: R.Tensor((8,), dtype="float16") = R.memory.alloc_tensor(
                storage2, R.prim_value(0), R.shape([8]), R.dtype("float16")
            )
            R.call_builtin_with_ctx(
                "vm.builtin.cuda_graph.run_or_capture",
                (
                    cls.main_cuda_graph_capture,
                    (alloc1, w, alloc2, R.shape([m])),
                    R.prim_value(0),
                    R.shape([m]),
                ),
                ty_args=(R.Tuple,),
            )
            storage3: R.Any = R.memory.alloc_storage(
                R.shape([8]), R.prim_value(0), R.str("global"), R.dtype("float16")
            )
            alloc3: R.Tensor((8,), dtype="float16") = R.memory.alloc_tensor(
                storage3, R.prim_value(0), R.shape([8]), R.dtype("float16")
            )
            R.call_packed("dummy", alloc2, w, alloc3, ty_args=(R.Tuple,))
            gv_1: R.Tuple(R.Tensor((8,), dtype="float16")) = (alloc3,)
            return gv_1

    After = relax.transform.RewriteCUDAGraph()(Before)
    tvm.ir.assert_structural_equal(After, Expected)


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