# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: F841 import numpy as np import pytest import torch import tvm from tvm import relax, testing from tvm.relax import VMInstrumentReturnKind from tvm.relax.testing.transform import ( dataflow_alias_analysis, dataflow_inplace_analysis, dataflow_liveness_analysis, dataflow_single_inplace_call, ) from tvm.relax.transform import DataflowUseInplaceCalls from tvm.script.parser import ir as I from tvm.script.parser import relax as R from tvm.script.parser import tirx as T def test_liveness_analysis(): @I.ir_module(s_tir=True) class BasicLiveness: @R.function def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): with R.dataflow(): y = R.const(1, dtype="int32") z = R.add(x, y) q = R.multiply(z, y) p = R.add(z, q) n = R.multiply(p, p) R.output(n, p) return n block = BasicLiveness["main"].body.blocks[0] live_ranges = dataflow_liveness_analysis(block) expected_ranges = { # x is live past the binding block "x": (-1, 5), "y": (0, 2), "z": (1, 3), "q": (2, 3), # exposed though ultimately not used "p": (3, 5), "n": (4, 5), } actual_ranges = {var.name_hint: live_range for var, live_range in live_ranges.items()} assert actual_ranges == expected_ranges def test_alias_analysis_basic(): @I.ir_module(s_tir=True) class BasicAliasAnalysis: @R.function def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): with R.dataflow(): y = x # y is an alias of x z = R.add(y, y) # fresh value n = z # alias of z R.output(n) return n block = BasicAliasAnalysis["main"].body.blocks[0] alias_sets, tuple_map = dataflow_alias_analysis(block, BasicAliasAnalysis["main"].params) expected = { "x": {0}, "y": {0}, "z": {1}, "n": {1}, } for var, alias_set in alias_sets.items(): assert alias_set == expected[var.name_hint] assert tuple_map == {} def test_alias_analysis_tuple(): @I.ir_module(s_tir=True) class AliasesWithTuples: @R.function def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): with R.dataflow(): y = R.const(1, dtype="int32") t = (x, y) a = t[0] b = t[1] c = t[0] d = t[1] u = t e = t[0] f = t[1] z = R.add(c, d) n = z R.output(n) return n block = AliasesWithTuples["main"].body.blocks[0] alias_sets, tuple_map = dataflow_alias_analysis(block, AliasesWithTuples["main"].params) expected = { "x": {0}, "y": {1}, "t": {2}, "a": {0}, "b": {1}, "c": {0}, "d": {1}, "u": {2}, "e": {0}, "f": {1}, "z": {3}, "n": {3}, } actual_alias_sets = {var.name_hint: alias_set for var, alias_set in alias_sets.items()} assert expected == actual_alias_sets assert 2 in tuple_map assert tuple_map[2] == [{0}, {1}] def test_alias_split(): @I.ir_module(s_tir=True) class AliasSplit: @R.function def main(x: R.Tensor((60,), "int32")) -> R.Tensor((15,), "int32"): with R.dataflow(): t = R.split(x, 4) y = t[0] z = t[1] q = t[2] p = t[3] n = z R.output(n) return n block = AliasSplit["main"].body.blocks[0] alias_sets, tuple_map = dataflow_alias_analysis(block, AliasSplit["main"].params) expected = { "x": {0}, "t": {1}, "y": {2}, "z": {3}, "q": {4}, "p": {5}, "n": {3}, } actual_alias_sets = {var.name_hint: alias_set for var, alias_set in alias_sets.items()} assert expected == actual_alias_sets assert len(tuple_map) == 1 assert 1 in tuple_map assert tuple_map[1] == [{2}, {3}, {4}, {5}] def test_alias_call_tir(): # call TIR can yield either a single tensor or a tuple @I.ir_module(s_tir=True) class AliasCallTir: @T.prim_func(s_tir=True) def tir_id(x: T.handle, y: T.handle) -> None: T.func_attr({"global_symbol": "tir_id"}) m = T.int32() n = T.int32() A = T.match_buffer(x, (m, n), "int32") B = T.match_buffer(y, (m, n), "int32") for i, j in T.grid(m, n): with T.sblock("id"): vi, vj = T.axis.remap("SS", [i, j]) B[vi, vj] = A[vi, vj] @T.prim_func(s_tir=True) def tir_id2(x: T.handle, y: T.handle, z: T.handle) -> None: T.func_attr({"global_symbol": "tir_id"}) m = T.int32() n = T.int32() A = T.match_buffer(x, (m, n), "int32") B = T.match_buffer(y, (m, n), "int32") C = T.match_buffer(z, (m, n), "int32") for i, j in T.grid(m, n): with T.sblock("id"): vi, vj = T.axis.remap("SS", [i, j]) B[vi, vj] = A[vi, vj] C[vi, vj] = A[vi, vj] @R.function def main(x: R.Tensor((10, 10), "int32")) -> R.Tensor((10, 10), "int32"): with R.dataflow(): cls = AliasCallTir y = R.call_tir(cls.tir_id, (x,), out_ty=R.Tensor((10, 10), "int32")) t = R.call_tir( cls.tir_id2, (y,), out_ty=[R.Tensor((10, 10), "int32"), R.Tensor((10, 10), "int32")], ) z = y p = t[0] q = t[1] u = t m = u[0] n = u[1] v = n R.output(v) return v block = AliasCallTir["main"].body.blocks[0] alias_sets, tuple_map = dataflow_alias_analysis(block, AliasCallTir["main"].params) expected = { "x": {0}, "y": {1}, "t": {2}, "z": {1}, "p": {3}, "q": {4}, "u": {2}, "m": {3}, "n": {4}, "v": {4}, } actual_alias_sets = {var.name_hint: alias_set for var, alias_set in alias_sets.items()} assert expected == actual_alias_sets assert len(tuple_map) == 1 assert 2 in tuple_map assert tuple_map[2] == [{3}, {4}] def test_mystery_calls(): @I.ir_module(s_tir=True) class AliasChaosCalls: @R.function def identity(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): return x @R.function def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): with R.dataflow(): cls = AliasChaosCalls y = cls.identity(x) z = cls.identity(y) m = R.const(1, dtype="int32") n = R.const(2, dtype="int32") t = (m, n) a = R.call_pure_packed( "chaos", t, ty_args=R.Tuple(R.Tensor((), "int32"), R.Tensor((), "int32")) ) b = a[0] c = a[1] R.output(c) return c block = AliasChaosCalls["main"].body.blocks[0] alias_sets, tuple_map = dataflow_alias_analysis(block, AliasChaosCalls["main"].params) expected = { "x": {0}, "y": {0, 1}, "z": {0, 1, 2}, "m": {3}, "n": {4}, "t": {5}, "a": {3, 4, 5, 6, 7, 8}, # either t or a fresh tuple "b": {3, 4, 5, 6, 7, 8}, # the tuple components can be aliased to any member... "c": {3, 4, 5, 6, 7, 8}, # the tuple components can be aliased to any member... # (in principle, we can use type information to narrow down the aliasing) } actual_alias_sets = {var.name_hint: alias_set for var, alias_set in alias_sets.items()} assert expected == actual_alias_sets assert len(tuple_map) == 2 assert 5 in tuple_map assert tuple_map[5] == [{3}, {4}] assert 6 in tuple_map assert tuple_map[6] == [{3, 4, 5, 6, 7, 8}, {3, 4, 5, 6, 7, 8}] def test_alias_external_value(): @I.ir_module(s_tir=True) class AliasExternalValue: @R.function def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): y = R.const(1, dtype="int32") # not in DF block, treated as external t1 = (y, y) # not in DF block, treated as external with R.dataflow(): z = y # mystery value a = R.const(2, dtype="int32") t2 = (z, a) b = t2[0] c = t1[1] # tuple index into external value R.output(b) return b block = AliasExternalValue["main"].body.blocks[1] alias_sets, tuple_map = dataflow_alias_analysis(block, AliasExternalValue["main"].params) expected = { "x": {0}, "z": {-1}, "a": {1}, "t2": {2}, "b": {-1}, "c": {-1}, } actual_alias_sets = {var.name_hint: alias_set for var, alias_set in alias_sets.items()} assert expected == actual_alias_sets assert len(tuple_map) == 1 assert 2 in tuple_map assert tuple_map[2] == [{-1}, {1}] def test_inplace_simple_case(): @I.ir_module(s_tir=True) class InplaceBasic: @R.function def main(x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")) -> R.Tensor( (2, 3), "int32" ): with R.dataflow(): z = R.add(x, y) # cannot be done inplace: x and y are live later p = R.add(z, z) # can be done inplace: z is not used later r = p # alias of p m = R.multiply(p, p) # p is not used later but r is, so can't do inplace n = R.add(m, r) # can be done inplace: r is not used again ret = R.subtract(n, m) # can be done inplace: neither is used again R.output(ret) return ret block = InplaceBasic["main"].body.blocks[0] size_match, exact_match = dataflow_inplace_analysis( block, InplaceBasic["main"].params, InplaceBasic ) # order does not matter for the listing of candidates, so we have to implement as sets def assert_candidate_list( actual: list[tuple[int, set[int]]], expected: list[tuple[int, set[int]]] ) -> None: assert len(actual) == len(expected) for i in range(len(actual)): assert actual[i][0] == expected[i][0] assert len(expected[i][1]) == len(actual[i][1]) for idx in actual[i][1]: assert idx in expected[i][1] assert_candidate_list(size_match, [(1, {0, 1}), (4, {1}), (5, {0, 1})]) # TODO(@slyubomirsky): I couldn't think of an easy example where sizes don't match, # but broadcasting might cause it to happen assert_candidate_list(exact_match, [(1, {0, 1}), (4, {1}), (5, {0, 1})]) def test_inplace_single_call(): @I.ir_module(s_tir=True) class TestModule: @R.function def main( x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32") ) -> R.Tensor((2, 3), dtype="float32"): z = R.add(x, y) q = R.nn.silu(z) return q add_call = TestModule["main"].body.blocks[0].bindings[0].value new_add, new_mod = dataflow_single_inplace_call(TestModule, add_call, [0]) @T.prim_func(private=True, s_tir=True) def expected_add( A: T.Buffer((T.int64(2), T.int64(3)), "float32"), B: T.Buffer((T.int64(2), T.int64(3)), "float32"), ): T.func_attr({"tirx.noalias": True}) for ax0, ax1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1]) T.writes(A[v_ax0, v_ax1]) A[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[v_ax0, v_ax1] tvm.ir.assert_structural_equal(new_mod["add_inplace"], expected_add) assert new_add.op.name == "relax.call_tir_inplace" assert new_add.args[0].name_hint == "add_inplace" for i, arg in enumerate(new_add.args[1].fields): arg == add_call.args[i] new_add.attrs.inplace_indices == [0] @T.prim_func(private=True, s_tir=True) def expected_silu(A: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) compute = T.sblock_alloc_buffer((T.int64(2), T.int64(3))) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("compute"): v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) T.reads(A[v_i0, v_i1]) T.writes(compute[v_i0, v_i1]) compute[v_i0, v_i1] = T.sigmoid(A[v_i0, v_i1]) for ax0, ax1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_multiply"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], compute[v_ax0, v_ax1]) T.writes(A[v_ax0, v_ax1]) A[v_ax0, v_ax1] = A[v_ax0, v_ax1] * compute[v_ax0, v_ax1] silu_call = TestModule["main"].body.blocks[0].bindings[1].value new_silu, new_mod = dataflow_single_inplace_call(TestModule, silu_call, [0]) tvm.ir.assert_structural_equal(new_mod["silu_inplace"], expected_silu) assert new_silu.op.name == "relax.call_tir_inplace" assert new_silu.args[0].name_hint == "silu_inplace" for i, arg in enumerate(new_silu.args[1].fields): arg == silu_call.args[i] new_silu.attrs.inplace_indices == [0] def test_insert_inplace_calls(): @I.ir_module(s_tir=True) class EndToEndTest: @R.function def main( x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((1, 3), dtype="float32") ) -> R.Tensor((2, 3), dtype="float32"): with R.dataflow(): z = R.add(x, y) # broadcast happens here # Cannot be done in-place because x is an argument. a = R.add(z, y) # this one can be done in-place q = R.multiply(a, y) # broadcast again, a is eligible r = R.subtract(y, y) # cannot be done in-place because y is an argument s = R.subtract(r, r) # No broadcast. Can be done in-place m = R.multiply(q, s) # should give us all zeros R.output(m) return m @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def add_inplace( A: T.Buffer((T.int64(2), T.int64(3)), "float32"), B: T.Buffer((T.int64(1), T.int64(3)), "float32"), ): T.func_attr({"tirx.noalias": True}) for ax0, ax1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], B[T.int64(0), v_ax1]) T.writes(A[v_ax0, v_ax1]) A[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[T.int64(0), v_ax1] @T.prim_func(private=True, s_tir=True) def multiply_inplace( A: T.Buffer((T.int64(2), T.int64(3)), "float32"), B: T.Buffer((T.int64(1), T.int64(3)), "float32"), ): T.func_attr({"tirx.noalias": True}) for ax0, ax1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_multiply"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], B[T.int64(0), v_ax1]) T.writes(A[v_ax0, v_ax1]) A[v_ax0, v_ax1] = A[v_ax0, v_ax1] * B[T.int64(0), v_ax1] @T.prim_func(private=True, s_tir=True) def subtract_inplace( A: T.Buffer((T.int64(1), T.int64(3)), "float32"), B: T.Buffer((T.int64(1), T.int64(3)), "float32"), ): T.func_attr({"tirx.noalias": True}) for ax0, ax1 in T.grid(T.int64(1), T.int64(3)): with T.sblock("T_subtract"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1]) T.writes(B[v_ax0, v_ax1]) B[v_ax0, v_ax1] = A[v_ax0, v_ax1] - B[v_ax0, v_ax1] @R.function def main( x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((1, 3), dtype="float32") ) -> R.Tensor((2, 3), dtype="float32"): cls = Expected with R.dataflow(): z: R.Tensor((2, 3), dtype="float32") = R.add(x, y) a: R.Tensor((2, 3), dtype="float32") = R.call_tir_inplace( cls.add_inplace, (z, y), inplace_indices=[0], out_ty=[ R.Tensor((2, 3), dtype="float32"), ], ) q: R.Tensor((2, 3), dtype="float32") = R.call_tir_inplace( cls.multiply_inplace, (a, y), inplace_indices=[0], out_ty=[ R.Tensor((2, 3), dtype="float32"), ], ) r: R.Tensor((1, 3), dtype="float32") = R.subtract(y, y) s: R.Tensor((1, 3), dtype="float32") = R.call_tir_inplace( cls.subtract_inplace, (r, r), inplace_indices=[1], out_ty=[ R.Tensor((1, 3), dtype="float32"), ], ) m: R.Tensor((2, 3), dtype="float32") = R.call_tir_inplace( cls.multiply_inplace, (q, s), inplace_indices=[0], out_ty=[ R.Tensor((2, 3), dtype="float32"), ], ) R.output(m) return m transform_pass = DataflowUseInplaceCalls() new_mod = transform_pass(EndToEndTest) tvm.ir.assert_structural_equal(new_mod, Expected) x = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32")) y = tvm.runtime.tensor(np.random.rand(1, 3).astype("float32")) expected = np.zeros((2, 3), dtype="float32") target = tvm.target.Target("llvm") ex = tvm.compile(new_mod, target) vm = relax.VirtualMachine(ex, tvm.cpu()) res = vm["main"](x, y) assert (expected == res.numpy()).all() def test_dynamic(): @I.ir_module(s_tir=True) class DynamicTestCase: @R.function def main( x: R.Tensor(("a", "b"), dtype="float32"), y: R.Tensor(("a", "b"), dtype="float32") ) -> R.Tensor(("a", "b"), dtype="float32"): with R.dataflow(): z = R.add(x, y) # Cannot be done in-place because x and y are arguments a = R.add(z, y) # this one can be done in-place s = R.subtract(a, a) # No broadcast. Can be done in-place R.output(s) return s # the result should be all zeroes transform_pass = DataflowUseInplaceCalls() new_mod = transform_pass(DynamicTestCase) @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def add_inplace(var_A: T.handle, var_B: T.handle): T.func_attr({"tirx.noalias": True}) a, b = T.int64(), T.int64() A = T.match_buffer(var_A, (a, b)) B = T.match_buffer(var_B, (a, b)) for ax0, ax1 in T.grid(a, b): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1]) T.writes(A[v_ax0, v_ax1]) A[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[v_ax0, v_ax1] @T.prim_func(private=True, s_tir=True) def subtract_inplace(var_A: T.handle, var_B: T.handle): T.func_attr({"tirx.noalias": True}) a, b = T.int64(), T.int64() A = T.match_buffer(var_A, (a, b)) B = T.match_buffer(var_B, (a, b)) for ax0, ax1 in T.grid(a, b): with T.sblock("T_subtract"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1]) T.writes(B[v_ax0, v_ax1]) B[v_ax0, v_ax1] = A[v_ax0, v_ax1] - B[v_ax0, v_ax1] @R.function def main( x: R.Tensor(("a", "b"), dtype="float32"), y: R.Tensor(("a", "b"), dtype="float32") ) -> R.Tensor(("a", "b"), dtype="float32"): a = T.int64() b = T.int64() cls = Expected with R.dataflow(): z = R.add(x, y) a_1 = R.call_tir_inplace( cls.add_inplace, (z, y), out_ty=R.Tensor((a, b), dtype="float32"), inplace_indices=[0], ) s = R.call_tir_inplace( cls.subtract_inplace, (a_1, a_1), out_ty=R.Tensor((a, b), dtype="float32"), inplace_indices=[1], ) R.output(s) return s tvm.ir.assert_structural_equal(new_mod, Expected, map_free_vars=True) x = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32")) y = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32")) expected = np.zeros((2, 3), dtype="float32") target = tvm.target.Target("llvm") ex = tvm.compile(new_mod, target) vm = relax.VirtualMachine(ex, tvm.cpu()) res = vm["main"](x, y) assert (expected == res.numpy()).all() def test_dynamic_mismatch(): # cannot statically prove the shapes to be equal so the module should be unchanged @I.ir_module(s_tir=True) class DynamicMistmatchTestCase: @R.function def main( x: R.Tensor(("a", "b"), dtype="float32"), y: R.Tensor(("c", "d"), dtype="float32") ): with R.dataflow(): z = R.add(x, y) # Cannot be done in-place because x and y are arguments a = R.add(z, y) # cannot conclude that shapes match R.output(a) return a transform_pass = DataflowUseInplaceCalls() new_mod = transform_pass(DynamicMistmatchTestCase) tvm.ir.assert_structural_equal(new_mod, DynamicMistmatchTestCase) class TestViewOpSharedStorageAndNoInplace: storage_ptr_x_1d = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32) storage_ptr_x_2d = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32) storage_ptr_x_squeeze = np.array([[[1.0], [2.0], [3.0], [4.0]]], dtype=np.float32) storage_ptr_x_ensure_zero_offset = np.array([[1.0], [2.0], [3.0], [4.0]], dtype=np.float32) @I.ir_module class _SharedStorageExpandDimsModule: @R.function def main(x: R.Tensor((4,), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"): with R.dataflow(): lv: R.Tensor((4, 1), dtype="float32") = R.expand_dims(x, axis=[1]) lv1: R.Tensor((4, 1), dtype="float32") = R.expand_dims(x, axis=[1]) gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @I.ir_module class _SharedStorageSqueezeModule: @R.function def main(x: R.Tensor((1, 4, 1), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"): with R.dataflow(): lv: R.Tensor((4, 1), dtype="float32") = R.squeeze(x, axis=[0]) lv1: R.Tensor((4, 1), dtype="float32") = R.squeeze(x, axis=[0]) gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @I.ir_module class _SharedStorageReshapeModule: @R.function def main(x: R.Tensor((4,), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"): with R.dataflow(): lv: R.Tensor((4, 1), dtype="float32") = R.reshape(x, (4, 1)) lv1: R.Tensor((4, 1), dtype="float32") = R.reshape(x, (4, 1)) gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @I.ir_module class _SharedStoragePermuteDimsModule: @R.function def main(x: R.Tensor((1, 4), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"): with R.dataflow(): lv: R.Tensor((4, 1), dtype="float32") = R.permute_dims(x, axes=[1, 0]) lv1: R.Tensor((4, 1), dtype="float32") = R.permute_dims(x, axes=[1, 0]) gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @I.ir_module class _SharedStorageViewModule: @R.function def main(x: R.Tensor((4,), dtype="float32")) -> R.Tensor((1, 4), dtype="float32"): with R.dataflow(): lv: R.Tensor((1, 4), dtype="float32") = R.memory.view( x, R.shape([1, 4]), R.tuple(), R.tuple() ) lv1: R.Tensor((1, 4), dtype="float32") = R.memory.view( x, R.shape([1, 4]), R.tuple(), R.tuple() ) gv: R.Tensor((1, 4), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @I.ir_module class _SharedStorageBatchFlattenModule: @R.function def main(x: R.Tensor((1, 4), dtype="float32")) -> R.Tensor((1, 4), dtype="float32"): with R.dataflow(): lv: R.Tensor((1, 4), dtype="float32") = R.nn.batch_flatten(x) lv1: R.Tensor((1, 4), dtype="float32") = R.nn.batch_flatten(x) gv: R.Tensor((1, 4), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @I.ir_module class _SharedStorageFlattenModule: @R.function def main(x: R.Tensor((1, 4), dtype="float32")) -> R.Tensor((4,), dtype="float32"): with R.dataflow(): lv: R.Tensor((4,), dtype="float32") = R.flatten(x) lv1: R.Tensor((4,), dtype="float32") = R.flatten(x) gv: R.Tensor((4,), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @I.ir_module class _SharedStorageEnsureZeroOffsetModule: @R.function def main(x: R.Tensor((4, 1), dtype="float32")) -> R.Tensor((4, 1), dtype="float32"): with R.dataflow(): lv: R.Tensor((4, 1), dtype="float32") = R.memory.ensure_zero_offset(x) lv1: R.Tensor((4, 1), dtype="float32") = R.memory.ensure_zero_offset(x) gv: R.Tensor((4, 1), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @I.ir_module class _IndependentReluModule: """Just a testcase to verify that non-view ops do not share storage.""" @R.function def main(x: R.Tensor((4,), dtype="float32")) -> R.Tensor((4,), dtype="float32"): with R.dataflow(): lv: R.Tensor((4,), dtype="float32") = R.nn.relu(x) lv1: R.Tensor((4,), dtype="float32") = R.nn.relu(x) gv: R.Tensor((4,), dtype="float32") = R.add(lv, lv1) R.output(gv) return gv @classmethod def _capture_op_tensors(cls, mod, input_nps, op_substr): """Capture TVM tensors passed to VM calls whose name contains op_substr.""" captures = [] def instrument(func, name, before_run, ret_value, *args): del func, ret_value if not before_run: return VMInstrumentReturnKind.NO_OP if op_substr not in name.lower(): return VMInstrumentReturnKind.NO_OP tensor_args = [arg for arg in args if isinstance(arg, tvm.runtime.Tensor)] if not tensor_args: return VMInstrumentReturnKind.NO_OP captures.append({"call_name": name, "tensors": tensor_args}) return VMInstrumentReturnKind.NO_OP if isinstance(input_nps, np.ndarray): input_nps = [input_nps] ex = relax.build(mod, tvm.target.Target("llvm")) vm = relax.VirtualMachine(ex, tvm.cpu()) vm.set_instrument(instrument) vm["main"](*(tvm.runtime.tensor(arr, tvm.cpu()) for arr in input_nps)) return captures @pytest.mark.parametrize( "mod,input_nps,op_substr,expect_same_storage", [ pytest.param( _SharedStorageExpandDimsModule, [storage_ptr_x_1d], "add", True, id="shared_storage_expand_dims", ), pytest.param( _SharedStorageSqueezeModule, [storage_ptr_x_squeeze], "add", True, id="shared_storage_squeeze", ), pytest.param( _SharedStorageReshapeModule, [storage_ptr_x_1d], "add", True, id="shared_storage_reshape", ), pytest.param( _SharedStoragePermuteDimsModule, [storage_ptr_x_2d], "add", True, id="shared_storage_permute_dims", ), pytest.param( _SharedStorageFlattenModule, [storage_ptr_x_2d], "add", True, id="shared_storage_flatten", ), pytest.param( _SharedStorageBatchFlattenModule, [storage_ptr_x_2d], "add", True, id="shared_storage_batch_flatten", ), pytest.param( _SharedStorageViewModule, [storage_ptr_x_1d], "add", True, id="shared_storage_memory_view", ), pytest.param( _SharedStorageEnsureZeroOffsetModule, [storage_ptr_x_ensure_zero_offset], "add", True, id="shared_storage_ensure_zero_offset", ), pytest.param( _IndependentReluModule, [storage_ptr_x_1d], "add", False, id="independent_storage_relu", ), ], ) def test_tensor_storage_ptr_extraction(self, mod, input_nps, op_substr, expect_same_storage): """Validate runtime storage overlap/sharing via VM instrumentation.""" storage_shared = tvm.get_global_func("runtime.TVMTensorIsStorageShared") captures = self._capture_op_tensors(mod, input_nps, op_substr) assert len(captures), f"VM instrumentation did not see a {op_substr} call." assert len(captures) == 1, f"VM instrumentation should see exactly one {op_substr} call." cap = captures[0] assert len(cap["tensors"]) == 3, ( f"VM instrumentation should see three {op_substr} tensor operands." ) tensor_a, tensor_b = cap["tensors"][0], cap["tensors"][1] call_name = cap["call_name"] if expect_same_storage: assert storage_shared(tensor_a, tensor_b), ( f"{mod.__name__}: operands should share the same storage (call {call_name!r})" ) else: assert not storage_shared(tensor_a, tensor_b), ( f"{mod.__name__}: operands must not share storage (call {call_name!r})" ) @staticmethod def _emit_duplicate_view(op, x): if op == "relax.expand_dims": a = relax.op.expand_dims(x, axis=1) b = relax.op.expand_dims(x, axis=1) elif op == "relax.squeeze": a = relax.op.squeeze(x, axis=[0]) b = relax.op.squeeze(x, axis=[0]) elif op == "relax.reshape": a = relax.op.reshape(x, (4, 1)) b = relax.op.reshape(x, (4, 1)) elif op == "relax.permute_dims": a = relax.op.permute_dims(x, axes=[1, 0]) b = relax.op.permute_dims(x, axes=[1, 0]) elif op == "relax.memory.view": a = relax.op.memory.view(x, (4, 1)) b = relax.op.memory.view(x, (4, 1)) elif op == "relax.memory.ensure_zero_offset": a = relax.op.memory.ensure_zero_offset(x) b = relax.op.memory.ensure_zero_offset(x) elif op == "relax.flatten": a = relax.op.flatten(x) b = relax.op.flatten(x) elif op == "relax.nn.batch_flatten": a = relax.op.nn.batch_flatten(x) b = relax.op.nn.batch_flatten(x) else: raise ValueError(op) return a, b @staticmethod def _concat_axis_for_view_op(op): if op == "relax.flatten": return 0 return 1 @classmethod def _build_module(cls, op): if op == "relax.expand_dims": x_ty = relax.TensorType((4,), "float32") elif op == "relax.squeeze": x_ty = relax.TensorType((1, 4, 1), "float32") elif op == "relax.reshape": x_ty = relax.TensorType((4,), "float32") elif op == "relax.permute_dims": x_ty = relax.TensorType((1, 4), "float32") elif op == "relax.memory.view": x_ty = relax.TensorType((4,), "float32") elif op == "relax.memory.ensure_zero_offset": x_ty = relax.TensorType((4, 1), "float32") elif op in ("relax.flatten", "relax.nn.batch_flatten"): x_ty = relax.TensorType((1, 4), "float32") else: raise ValueError(op) bb = relax.BlockBuilder() x = relax.Var("x", x_ty) concat_axis = cls._concat_axis_for_view_op(op) with bb.function("main", [x]): with bb.dataflow(): a_expr, b_expr = cls._emit_duplicate_view(op, x) a = bb.emit(a_expr) b = bb.emit(b_expr) prod = bb.emit(relax.op.multiply(a, b)) out = bb.emit(relax.op.concat([prod, b], axis=concat_axis)) gv = bb.emit_output(out) bb.emit_func_output(gv) return bb.finalize() @classmethod def _input_for_view_op(cls, op): if op == "relax.squeeze": return cls.storage_ptr_x_squeeze if op == "relax.memory.ensure_zero_offset": return cls.storage_ptr_x_ensure_zero_offset if op in ("relax.permute_dims", "relax.flatten", "relax.nn.batch_flatten"): return cls.storage_ptr_x_2d return cls.storage_ptr_x_1d @staticmethod def _torch_duplicate_view(x, op): if op == "relax.expand_dims": return x.unsqueeze(1) if op == "relax.squeeze": return x.squeeze(0) if op == "relax.reshape": return x.reshape(4, 1) if op == "relax.permute_dims": return x.permute(1, 0) if op == "relax.memory.view": return x.reshape(4, 1) if op == "relax.memory.ensure_zero_offset": return x if op == "relax.flatten": return x.flatten() if op == "relax.nn.batch_flatten": # TVM: ndim==2 input keeps shape (1, 4). return x raise ValueError(op) @classmethod def _expected_for_view_op(cls, op): x = torch.from_numpy(np.asarray(cls._input_for_view_op(op), dtype=np.float32)) a = cls._torch_duplicate_view(x, op) b = cls._torch_duplicate_view(x, op) prod = a * b concat_axis = cls._concat_axis_for_view_op(op) return torch.cat([prod, b], dim=concat_axis).numpy() @pytest.mark.parametrize( "view_op", ( # Keep this list in sync with IsViewMemoryOp() in # src/relax/transform/dataflow_inplace.cc "relax.expand_dims", "relax.squeeze", "relax.reshape", "relax.permute_dims", "relax.flatten", "relax.nn.batch_flatten", "relax.memory.view", "relax.memory.ensure_zero_offset", ), ) def test_no_inplace_when_view_ops_share_input(self, view_op): mod = self._build_module(view_op) func = mod["main"] block = func.body.blocks[0] params = list(func.params) alias_sets, _ = dataflow_alias_analysis(block, params) view_vars = [ binding.var for binding in block.bindings if ( isinstance(binding.value, relax.Call) and isinstance(binding.value.op, tvm.ir.Op) and binding.value.op.name == view_op ) ] a_var, b_var = view_vars[:2] assert alias_sets[a_var] & alias_sets[b_var], ( f"{view_op}: duplicate views should share alias sets, but got " f"{alias_sets[a_var]} and {alias_sets[b_var]}" ) _, exact_match = dataflow_inplace_analysis(block, params, mod) assert exact_match == [], f"{view_op}: expected no in-place opportunities" x_np = self._input_for_view_op(view_op).copy() mod_inplace = DataflowUseInplaceCalls()(mod) tvm.ir.assert_structural_equal(mod_inplace, mod) storage_shared = tvm.get_global_func("runtime.TVMTensorIsStorageShared") captures = self._capture_op_tensors(mod_inplace, x_np, "multiply") assert captures, f"{view_op}: VM instrumentation did not see a multiply call." cap = next(c for c in captures if len(c["tensors"]) >= 2) tensor_a, tensor_b = cap["tensors"][0], cap["tensors"][1] assert storage_shared(tensor_a, tensor_b), ( f"{view_op}: multiply operands should share the same storage at runtime " f"(call {cap['call_name']!r})" ) ex = relax.build(mod_inplace, tvm.target.Target("llvm")) vm = relax.VirtualMachine(ex, tvm.cpu()) out = vm["main"](tvm.runtime.tensor(x_np, tvm.cpu())) np.testing.assert_allclose(out.numpy(), self._expected_for_view_op(view_op)) if __name__ == "__main__": testing.main()