# 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()