# 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. # pylint: disable=missing-docstring, invalid-name # ruff: noqa: E501, F841 import numpy as np import pytest import tvm import tvm.testing from tvm import relax, s_tir, tirx from tvm.relax.frontend.nn import Module, Tensor, op, spec from tvm.script import ir as I from tvm.script import relax as R from tvm.script import tirx as T from tvm.testing import env # mypy: disable-error-code="attr-defined,valid-type,name-defined" def test_unary(): class Model(Module): def test(self, x: Tensor): z0 = op.square(x) z1 = op.sqrt(x) return (z0, z1) # fmt: off @R.function def test(x: R.Tensor((1, 10), dtype="float32"), _io: R.Any): R.func_attr({"num_input": 2}) with R.dataflow(): square: R.Tensor((1, 10), dtype="float32") = R.square(x) sqrt: R.Tensor((1, 10), dtype="float32") = R.sqrt(x) gv1 = (square, sqrt), (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm( spec={"test": {"x": spec.Tensor([1, 10], "float32")}}, debug=True, ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_binary(): class Model(Module): def test(self, x: Tensor, y: Tensor): z0 = op.add(x, y) z1 = op.multiply(x, y) z2 = op.divide(x, y) z3 = op.matmul(x, y) z4 = op.maximum(x, y) z5 = op.minimum(x, y) z6 = op.subtract(x, y) z7 = op.greater(x, y) z8 = op.greater_equal(x, y) z9 = op.less(x, y) z10 = op.less_equal(x, y) z11 = op.equal(x, y) z12 = op.not_equal(x, y) return (z0, z1, z2, z3, z4, z5, z6, z7, z8, z9, z10, z11, z12) # fmt: off @R.function def test(x: R.Tensor((1, 10), dtype="float32"), y: R.Tensor((10, 1), dtype="float32"), _io: R.Any): R.func_attr({"num_input": 3}) with R.dataflow(): add: R.Tensor((10, 10), dtype="float32") = R.add(x, y) mul: R.Tensor((10, 10), dtype="float32") = R.multiply(x, y) divide: R.Tensor((10, 10), dtype="float32") = R.divide(x, y) matmul: R.Tensor((1, 1), dtype="float32") = R.matmul(x, y, out_dtype=None) maximum: R.Tensor((10, 10), dtype="float32") = R.maximum(x, y) minimum: R.Tensor((10, 10), dtype="float32") = R.minimum(x, y) subtract: R.Tensor((10, 10), dtype="float32") = R.subtract(x, y) greater: R.Tensor((10, 10), dtype="bool") = x > y greater_equal: R.Tensor((10, 10), dtype="bool") = x >= y less: R.Tensor((10, 10), dtype="bool") = x < y less_equal: R.Tensor((10, 10), dtype="bool") = x <= y equal: R.Tensor((10, 10), dtype="bool") = R.equal(x, y) not_equal: R.Tensor((10, 10), dtype="bool") = R.not_equal(x, y) gv1 = (add, mul, divide, matmul, maximum, minimum, subtract, greater, greater_equal, less, less_equal, equal, not_equal), (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm( spec={"test": {"x": spec.Tensor([1, 10], "float32"), "y": spec.Tensor([10, 1], "float32")}}, debug=True, ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_sum(): class Model(Module): def test(self, x: Tensor): z0 = op.sum(x, axis=[1, 2], keepdims=True) return z0 # fmt: off @R.function def test(x: R.Tensor((3, 5, 2, 4), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tensor((3, 1, 1, 4), dtype="float32"), R.Tuple(R.Any)): R.func_attr({"num_input": 2}) with R.dataflow(): sum: R.Tensor((3, 1, 1, 4), dtype="float32") = R.sum(x, axis=[1, 2], keepdims=True) gv1: R.Tuple(R.Tensor((3, 1, 1, 4), dtype="float32"), R.Tuple(R.Any)) = sum, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm( spec={"test": {"x": spec.Tensor([3, 5, 2, 4], "float32")}}, debug=True ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_max(): class Model(Module): def test(self, x: Tensor): z0 = op.max(x, axis=[1, 2], keepdims=True) return z0 # fmt: off @R.function def test(x: R.Tensor((3, 5, 2, 4), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tensor((3, 1, 1, 4), dtype="float32"), R.Tuple(R.Any)): R.func_attr({"num_input": 2}) with R.dataflow(): max: R.Tensor((3, 1, 1, 4), dtype="float32") = R.max(x, axis=[1, 2], keepdims=True) gv1: R.Tuple(R.Tensor((3, 1, 1, 4), dtype="float32"), R.Tuple(R.Any)) = max, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm( spec={"test": {"x": spec.Tensor([3, 5, 2, 4], "float32")}}, debug=True ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_min(): class Model(Module): def test(self, x: Tensor): z0 = op.min(x, axis=[1, 2], keepdims=True) return z0 # fmt: off @R.function def test(x: R.Tensor((3, 5, 2, 4), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tensor((3, 1, 1, 4), dtype="float32"), R.Tuple(R.Any)): R.func_attr({"num_input": 2}) with R.dataflow(): min: R.Tensor((3, 1, 1, 4), dtype="float32") = R.min(x, axis=[1, 2], keepdims=True) gv1: R.Tuple(R.Tensor((3, 1, 1, 4), dtype="float32"), R.Tuple(R.Any)) = min, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm( spec={"test": {"x": spec.Tensor([3, 5, 2, 4], "float32")}}, debug=True ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_manipulate(): class Model(Module): def test(self, x: Tensor): z0 = op.broadcast_to(x, [2, 5, 2]) z1 = op.permute_dims(x, [2, 1, 0]) z2 = op.reshape(x, [1, 10]) z3 = op.repeat(x, repeats=2, axis=1) z4 = op.squeeze(x, 0) z5 = op.unsqueeze(x, 0) z6 = op.concat([x, x], dim=0) return (z0, z1, z2, z3, z4, z5, z6) # fmt: off @R.function def test(x: R.Tensor((1, 5, 2), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tuple(R.Tensor((2, 5, 2), dtype="float32"), R.Tensor((2, 5, 1), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10, 2), dtype="float32"), R.Tensor((5, 2), dtype="float32"), R.Tensor((1, 1, 5, 2), dtype="float32"), R.Tensor((2, 5, 2), dtype="float32")), R.Tuple(R.Any)): R.func_attr({"num_input": 2}) with R.dataflow(): broadcast_to: R.Tensor((2, 5, 2), dtype="float32") = R.broadcast_to(x, R.shape([2, 5, 2])) permute_dims: R.Tensor((2, 5, 1), dtype="float32") = R.permute_dims(x, axes=[2, 1, 0]) reshape: R.Tensor((1, 10), dtype="float32") = R.reshape(x, R.shape([1, 10])) repeat: R.Tensor((1, 10, 2), dtype="float32") = R.repeat(x, repeats=2, axis=1) squeeze: R.Tensor((5, 2), dtype="float32") = R.squeeze(x, axis=[0]) unsqueeze: R.Tensor((1, 1, 5, 2), dtype="float32") = R.expand_dims(x, axis=0) concat: R.Tensor((2, 5, 2), dtype="float32") = R.concat([x, x], axis=0) gv1: R.Tuple(R.Tuple(R.Tensor((2, 5, 2), dtype="float32"), R.Tensor((2, 5, 1), dtype="float32"), R.Tensor((1, 10), dtype="float32"), R.Tensor((1, 10, 2), dtype="float32"), R.Tensor((5, 2), dtype="float32"), R.Tensor((1, 1, 5, 2), dtype="float32"), R.Tensor((2, 5, 2), dtype="float32")), R.Tuple(R.Any)) = (broadcast_to, permute_dims, reshape, repeat, squeeze, unsqueeze, concat), (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([1, 5, 2], "float32")}}, debug=True) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_index(): class Model(Module): def test(self, x: Tensor, y: Tensor): z0 = op.take(x, y, axis=2) return z0 # fmt: off @R.function def test(x: R.Tensor((2, 1, 10), dtype="float32"), y: R.Tensor((5,), dtype="int32"), _io: R.Any) -> R.Tuple(R.Tensor((2, 1, 5), dtype="float32"), R.Tuple(R.Any)): R.func_attr({"num_input": 3}) with R.dataflow(): take: R.Tensor((2, 1, 5), dtype="float32") = R.take(x, y, axis=2) gv1: R.Tuple(R.Tensor((2, 1, 5), dtype="float32"), R.Tuple(R.Any)) = take, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, params = m.export_tvm( spec={"test": {"x": spec.Tensor([2, 1, 10], "float32"), "y": spec.Tensor([5], "int32")}}, debug=True, ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_datatype(): class Model(Module): def test(self, x: Tensor): z0 = op.astype(x, "float16") return z0 # fmt: off @R.function def test(x: R.Tensor((2, 1, 10), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tensor((2, 1, 10), dtype="float16"), R.Tuple(R.Any)): R.func_attr({"num_input": 2}) with R.dataflow(): astype: R.Tensor((2, 1, 10), dtype="float16") = R.astype(x, dtype="float16") gv1: R.Tuple(R.Tensor((2, 1, 10), dtype="float16"), R.Tuple(R.Any)) = astype, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([2, 1, 10], "float32")}}, debug=True) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_image(): class Model(Module): def test(self, x: Tensor, weight: Tensor, bias: Tensor): padded = op.pad(x, [0, 0, 0, 0, 1, 1, 1, 1]) conv2d = op.conv2d(padded, weight, bias) interpolate = op.interpolate(x, size=[40, 40]) # type: ignore return (conv2d, interpolate) @R.function def test( x: R.Tensor((1, 3, 32, 32), dtype="float32"), weight: R.Tensor((32, 3, 3, 3), dtype="float32"), bias: R.Tensor((32,), dtype="float32"), _io: R.Any, ) -> R.Tuple( R.Tuple( R.Tensor((1, 32, 32, 32), dtype="float32"), R.Tensor((1, 3, 40, 40), dtype="float32") ), R.Tuple(R.Any), ): R.func_attr({"num_input": 4}) with R.dataflow(): lv0: R.Tensor((1, 3, 34, 34), dtype="float32") = R.nn.pad(x, (0, 0, 0, 0, 1, 1, 1, 1)) lv1: R.Tensor((1, 32, 32, 32), dtype="float32") = R.nn.conv2d( lv0, weight, strides=[1, 1], padding=[0, 0, 0, 0], dilation=[1, 1], groups=1, data_layout="NCHW", kernel_layout="OIHW", out_layout="NCHW", out_dtype=None, ) lv2: R.Tensor((1, 32, 1, 1), dtype="float32") = R.reshape(bias, R.shape([1, 32, 1, 1])) conv2d: R.Tensor((1, 32, 32, 32), dtype="float32") = R.add(lv1, lv2) interpolate: R.Tensor((1, 3, 40, 40), dtype="float32") = R.image.resize2d( x, R.shape([40, 40]), roi=[T.float32(0), T.float32(0), T.float32(0), T.float32(0)], layout="NCHW", method="nearest_neighbor", coordinate_transformation_mode="asymmetric", rounding_method="round", cubic_alpha=-0.75, cubic_exclude=0, extrapolation_value=0, out_dtype=None, ) gv1: R.Tuple( R.Tuple( R.Tensor((1, 32, 32, 32), dtype="float32"), R.Tensor((1, 3, 40, 40), dtype="float32"), ), R.Tuple(R.Any), ) = (conv2d, interpolate), (_io,) R.output(gv1) return gv1 m = Model() irmodule, _ = m.export_tvm( spec={ "test": { "x": spec.Tensor([1, 3, 32, 32], "float32"), "weight": spec.Tensor([32, 3, 3, 3], "float32"), "bias": spec.Tensor([32], "float32"), } }, debug=True, ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_chunk(): class Model(Module): def test(self, x: Tensor): chunk = op.chunk(x, chunks=4) return chunk @R.function def test(x: R.Tensor((8,), dtype="float32"), _io: R.Any) -> R.Tuple( R.Tuple( R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), ), R.Tuple(R.Any), ): R.func_attr({"num_input": 2}) with R.dataflow(): chunk: R.Tuple( R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), ) = R.split(x, indices_or_sections=4, axis=0) chunk_0: R.Tensor((2,), dtype="float32") = chunk[0] chunk_1: R.Tensor((2,), dtype="float32") = chunk[1] chunk_2: R.Tensor((2,), dtype="float32") = chunk[2] chunk_3: R.Tensor((2,), dtype="float32") = chunk[3] gv1: R.Tuple( R.Tuple( R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), ), R.Tuple(R.Any), ) = (chunk_0, chunk_1, chunk_2, chunk_3), (_io,) R.output(gv1) return gv1 m = Model() irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([8], "float32")}}, debug=True) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_nn(): class Model(Module): def test(self, x: Tensor, weight: Tensor, bias: Tensor): log_out = op.log(x) floor_out = op.floor(x) relu_out = op.relu(x) relu6_out = op.relu6(x) silu_out = op.silu(x) gelu_out = op.gelu(x) sigmoid_out = op.sigmoid(x) tanh_out = op.tanh(x) exp_out = op.exp(x) negative_out = op.negative(x) softplus_out = op.softplus(x, beta=1.0, threshold=20.0) softmax_out = op.softmax(x, axis=2) prelu_out = op.prelu(x, alpha=bias) rms_norm_out = op.rms_norm(x, weight, axes=[-2, -1]) rms_norm_with_bias_out = op.rms_norm(x, weight, axes=[-2, -1]) group_norm_out = op.group_norm(x, num_groups=1, weight=bias, bias=bias) return x @R.function def test( x: R.Tensor((2, 3, 4, 5), dtype="float32"), weight: R.Tensor((4, 5), dtype="float32"), bias: R.Tensor((3,), dtype="float32"), _io: R.Any, ) -> R.Tuple(R.Tensor((2, 3, 4, 5), dtype="float32"), R.Tuple(R.Any)): R.func_attr({"num_input": 4}) with R.dataflow(): log: R.Tensor((2, 3, 4, 5), dtype="float32") = R.log(x) floor: R.Tensor((2, 3, 4, 5), dtype="float32") = R.floor(x) relu: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.relu(x) relu6: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.relu6(x) silu: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.silu(x) gelu: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.gelu(x) sigmoid: R.Tensor((2, 3, 4, 5), dtype="float32") = R.sigmoid(x) tanh: R.Tensor((2, 3, 4, 5), dtype="float32") = R.tanh(x) exp: R.Tensor((2, 3, 4, 5), dtype="float32") = R.exp(x) negative: R.Tensor((2, 3, 4, 5), dtype="float32") = R.negative(x) softplus: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.softplus( x, beta=1.0, threshold=20.0 ) softmax: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.softmax(x, axis=2) prelu: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.prelu(x, bias) rms_norm: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.rms_norm( x, weight, axes=[-2, -1], epsilon=1.0000000000000001e-05 ) rms_norm1: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.rms_norm( x, weight, axes=[-2, -1], epsilon=1.0000000000000001e-05 ) group_norm: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.group_norm( x, bias, bias, num_groups=1, channel_axis=1, axes=[2, 3] ) gv1: R.Tuple(R.Tensor((2, 3, 4, 5), dtype="float32"), R.Tuple(R.Any)) = x, (_io,) R.output(gv1) return gv1 m = Model() irmodule, params = m.export_tvm( spec={ "test": { "x": spec.Tensor([2, 3, 4, 5], "float32"), "weight": spec.Tensor([4, 5], "float32"), "bias": spec.Tensor([3], "float32"), } }, debug=True, ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_create(): class Model(Module): def test(self, x: Tensor): triu_out = op.triu(x) full_with_scalar_out = op.full([10, 10], fill_value=10) # type: ignore full_with_FloatImm_out = op.full( [10, 10], fill_value=tirx.FloatImm(dtype="float32", value=10) ) full_with_Tensor_out = op.full( [10, 10], fill_value=Tensor.from_scalar(10, dtype="float32") ) zeros_out = op.zeros([10, 10]) zeros_fp16_out = op.zeros([10, 10], dtype="float16") arange_out = op.arange(0, 10, 1, "float32") return x # fmt: off @R.function def test(x: R.Tensor((10, 10), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tensor((10, 10), dtype="float32"), R.Tuple(R.Any)): R.func_attr({"num_input": 2}) with R.dataflow(): triu: R.Tensor((10, 10), dtype="float32") = R.triu(x, k=0) full: R.Tensor((10, 10), dtype="float32") = R.full(R.shape([10, 10]), R.const(10, "float32"), dtype="float32") full1: R.Tensor((10, 10), dtype="float32") = R.full(R.shape([10, 10]), R.const(10, "float32"), dtype="float32") full2: R.Tensor((10, 10), dtype="float32") = R.full(R.shape([10, 10]), R.const(10, "float32"), dtype="float32") zeros: R.Tensor((10, 10), dtype="float32") = R.zeros(R.shape([10, 10]), dtype="float32") zeros1: R.Tensor((10, 10), dtype="float16") = R.zeros(R.shape([10, 10]), dtype="float16") arange: R.Tensor((10,), dtype="float32") = R.arange(T.int64(0), T.int64(10), T.int64(1), dtype="float32") gv1: R.Tuple(R.Tensor((10, 10), dtype="float32"), R.Tuple(R.Any)) = x, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, params = m.export_tvm( spec={"test": {"x": spec.Tensor([10, 10], "float32")}}, debug=True ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_timestep_embedding(): class Model(Module): def test(self, x: Tensor): get_timestep_out = op.get_timestep_embedding(x, 10) return get_timestep_out @R.function def test(x: R.Tensor((3,), dtype="float32"), _io: R.Any) -> R.Tuple( R.Tensor((3, 10), dtype="float32"), R.Tuple(R.Any) ): R.func_attr({"num_input": 2}) with R.dataflow(): lv1: R.Tensor((3,), dtype="float32") = R.astype(x, dtype="float32") lv2: R.Tensor((3, 1), dtype="float32") = R.expand_dims(lv1, axis=[1]) lv3: R.Tensor((5,), dtype="float32") = R.arange( R.prim_value(T.int64(0)), R.prim_value(T.int64(5)), R.prim_value(T.int64(1)), dtype="float32", ) lv4: R.Tensor((5,), dtype="float32") = R.multiply( R.const(-9.2103404998779297, "float32"), lv3 ) lv5: R.Tensor((5,), dtype="float32") = R.divide(lv4, R.const(4, "float32")) lv6: R.Tensor((5,), dtype="float32") = R.exp(lv5) lv7: R.Tensor((1, 5), dtype="float32") = R.expand_dims(lv6, axis=[0]) lv8: R.Tensor((3, 5), dtype="float32") = R.multiply(lv2, lv7) lv9: R.Tensor((3, 5), dtype="float32") = R.sin(lv8) lv10: R.Tensor((3, 5), dtype="float32") = R.cos(lv8) lv11: R.Tensor((3, 10), dtype="float32") = R.concat((lv9, lv10), axis=-1) get_timestep_embedding: R.Tensor((3, 10), dtype="float32") = R.astype( lv11, dtype="float32" ) gv1: R.Tuple(R.Tensor((3, 10), dtype="float32"), R.Tuple(R.Any)) = ( get_timestep_embedding, (_io,), ) R.output(gv1) return gv1 m = Model() irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([3], "float32")}}, debug=True) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_scaled_dot_product_attention(): class Model(Module): def test(self, query: Tensor, key: Tensor, value: Tensor): scaled_dot_product_attention = op.scaled_dot_product_attention(query, key, value) return scaled_dot_product_attention @R.function def test( query: R.Tensor((1, 32, 32, 32), dtype="float32"), key: R.Tensor((1, 32, 32, 32), dtype="float32"), value: R.Tensor((1, 32, 32, 32), dtype="float32"), _io: R.Any, ) -> R.Tuple(R.Tensor((1, 32, 32, 32), dtype="float32"), R.Tuple(R.Any)): R.func_attr({"num_input": 4}) with R.dataflow(): scaled_dot_product_attention: R.Tensor((1, 32, 32, 32), dtype="float32") = ( R.nn.attention(query, key, value, scale=None, causal_mask=None) ) gv1: R.Tuple(R.Tensor((1, 32, 32, 32), dtype="float32"), R.Tuple(R.Any)) = ( scaled_dot_product_attention, (_io,), ) R.output(gv1) return gv1 m = Model() irmodule, _ = m.export_tvm( spec={ "test": { "query": spec.Tensor([1, 32, 32, 32], "float32"), "key": spec.Tensor([1, 32, 32, 32], "float32"), "value": spec.Tensor([1, 32, 32, 32], "float32"), } }, debug=True, ) tvm.ir.assert_structural_equal(irmodule["test"], test) def test_tensor_expr_op(): class Model(Module): def test(self, x: Tensor): tensor_expr_op_out = op.tensor_expr_op( tensor_expr_func=lambda x: x + 1, name_hint="add_one", args=[x] ) return tensor_expr_op_out # fmt: off @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def add_one(A: T.Buffer((T.int64(10), T.int64(10)), "float32"), T_add: T.Buffer((T.int64(10), T.int64(10)), "float32")): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1 in T.grid(T.int64(10), T.int64(10)): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1]) T.writes(T_add[v_ax0, v_ax1]) T_add[v_ax0, v_ax1] = A[v_ax0, v_ax1] + T.float32(1) @R.function def _initialize_effect() -> R.Tuple(R.Any): with R.dataflow(): _io: R.Any = R.null_value() lv: R.Tuple(R.Any) = (_io,) gv: R.Tuple(R.Any) = lv R.output(gv) return gv @R.function def test(x: R.Tensor((10, 10), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tensor((10, 10), dtype="float32"), R.Tuple(R.Any)): cls = Expected R.func_attr({"num_input": 2}) with R.dataflow(): lv1 = R.call_tir(cls.add_one, (x,), out_ty=R.Tensor((10, 10), dtype="float32")) gv1: R.Tuple(R.Tensor((10, 10), dtype="float32"), R.Tuple(R.Any)) = lv1, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm(spec={"test": {"x": spec.Tensor([10, 10], "float32")}}, debug=True) tvm.ir.assert_structural_equal(irmodule, Expected) def test_tensor_ir_op(): num_q_heads, num_kv_heads, head_dim = 8, 8, 16 fused_heads = num_q_heads + num_kv_heads * 2 dtype = "float16" @T.prim_func(private=True, s_tir=True) def fused_rope( # pylint: disable=too-many-locals var_qkv: T.handle, var_q: T.handle, var_k: T.handle, var_v: T.handle, # Scalar arguments must be specified after tensor arguments, # including the output tensor arguments # # TODO(Lunderberg): Update # `tvm.relax.frontend.nn.op.tensor_ir_op` to use `Expr` # instead of `tir_vars`, so that the order can be consistent # between the function definition and the arguments in # `op.tensor_ir_op`. offset: T.int64, ): batch_size = T.int64() seq_len = T.int64() qkv = T.match_buffer(var_qkv, (batch_size, seq_len, fused_heads, head_dim), dtype) q = T.match_buffer(var_q, (batch_size, seq_len, num_q_heads, head_dim), dtype) k = T.match_buffer(var_k, (batch_size, seq_len, num_kv_heads, head_dim), dtype) v = T.match_buffer(var_v, (batch_size, seq_len, num_kv_heads, head_dim), dtype) T.evaluate(offset) class Model(Module): def test(self, qkv: Tensor, offset: tirx.Var): tensor_expr_op_out = op.tensor_ir_op( fused_rope, "llama_fused_rope", args=[qkv, offset], out=[ Tensor.placeholder((1, 1, num_q_heads, head_dim), dtype), Tensor.placeholder((1, 1, num_kv_heads, head_dim), dtype), Tensor.placeholder((1, 1, num_kv_heads, head_dim), dtype), ], ) return tensor_expr_op_out # fmt: off @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def llama_fused_rope(var_qkv: T.handle, var_q: T.handle, var_k: T.handle, var_v: T.handle, offset: T.int64): batch_size, seq_len = T.int64(), T.int64() qkv = T.match_buffer(var_qkv, (batch_size, seq_len, 24, 16), "float16") q = T.match_buffer(var_q, (batch_size, seq_len, 8, 16), "float16") k = T.match_buffer(var_k, (batch_size, seq_len, 8, 16), "float16") v = T.match_buffer(var_v, (batch_size, seq_len, 8, 16), "float16") T.evaluate(offset) @R.function def _initialize_effect() -> R.Tuple(R.Any): with R.dataflow(): _io: R.Any = R.null_value() lv: R.Tuple(R.Any) = (_io,) gv: R.Tuple(R.Any) = lv R.output(gv) return gv @R.function def test(qkv: R.Tensor((1, 1, 24, 16), dtype="float16"), offset: R.Shape(["offset_1"]), _io: R.Any) -> R.Tuple(R.Tuple(R.Tensor((1, 1, 8, 16), dtype="float16"), R.Tensor((1, 1, 8, 16), dtype="float16"), R.Tensor((1, 1, 8, 16), dtype="float16")), R.Tuple(R.Any)): offset_1 = T.int64() R.func_attr({"num_input": 3}) cls = Expected with R.dataflow(): lv1 = R.call_tir(cls.llama_fused_rope, (qkv,), out_ty=[R.Tensor((1, 1, 8, 16), dtype="float16"), R.Tensor((1, 1, 8, 16), dtype="float16"), R.Tensor((1, 1, 8, 16), dtype="float16")], tir_vars=R.shape([offset_1])) llama_fused_rope_0: R.Tensor((1, 1, 8, 16), dtype="float16") = lv1[0] llama_fused_rope_1: R.Tensor((1, 1, 8, 16), dtype="float16") = lv1[1] llama_fused_rope_2: R.Tensor((1, 1, 8, 16), dtype="float16") = lv1[2] gv1: R.Tuple(R.Tuple(R.Tensor((1, 1, 8, 16), dtype="float16"), R.Tensor((1, 1, 8, 16), dtype="float16"), R.Tensor((1, 1, 8, 16), dtype="float16")), R.Tuple(R.Any)) = (llama_fused_rope_0, llama_fused_rope_1, llama_fused_rope_2), (_io,) R.output(gv1) return gv1 # fmt: on m = Model() irmodule, _ = m.export_tvm( spec={ "test": {"qkv": spec.Tensor([1, 1, fused_heads, head_dim], "float16"), "offset": int} }, debug=True, ) tvm.ir.assert_structural_equal(irmodule, Expected) def test_tensor_ir_inplace_op(): hidden_size = 4096 dtype = "float16" @T.prim_func(s_tir=True) def inplace_take( var_weight: T.handle, var_pos: T.handle, var_embeddings: T.handle, offset: T.int64 ): T.func_attr({"tirx.noalias": True}) vocab_size = T.int64() weight = T.match_buffer(var_weight, (vocab_size, hidden_size), dtype) seq_len = T.int64() total_seq_len = T.int64() pos = T.match_buffer(var_pos, (seq_len,), "int32") embeddings = T.match_buffer(var_embeddings, (total_seq_len, hidden_size), dtype) for ax0, ax1 in T.grid(seq_len, hidden_size): with T.sblock("T_take"): v0, v1 = T.axis.remap("SS", [ax0, ax1]) T.reads(weight[pos[v0], v1], pos[v0]) T.writes(embeddings[v0, v1]) embeddings[v0 + offset, v1] = weight[pos[v0], v1] class Model(Module): def test( self, embedding_table: Tensor, input_ids: Tensor, embedding_dst: Tensor, offset: int ): tensor_expr_op_out = op.tensor_ir_inplace_op( inplace_take, "inplace_take", args=[embedding_table, input_ids, embedding_dst, offset], inplace_indices=[2], out=Tensor.placeholder(embedding_dst.shape, embedding_dst.dtype), ) return tensor_expr_op_out @I.ir_module(s_tir=True) class Expected: @T.prim_func(s_tir=True) def inplace_take( var_weight: T.handle, var_pos: T.handle, var_embeddings: T.handle, offset: T.int64 ): T.func_attr({"tirx.noalias": True}) vocab_size = T.int64() weight = T.match_buffer(var_weight, (vocab_size, hidden_size), dtype) seq_len = T.int64() total_seq_len = T.int64() pos = T.match_buffer(var_pos, (seq_len,), "int32") embeddings = T.match_buffer(var_embeddings, (total_seq_len, hidden_size), dtype) for ax0, ax1 in T.grid(seq_len, hidden_size): with T.sblock("T_take"): v0, v1 = T.axis.remap("SS", [ax0, ax1]) T.reads(weight[pos[v0], v1], pos[v0]) T.writes(embeddings[v0, v1]) embeddings[v0 + offset, v1] = weight[pos[v0], v1] @R.function def _initialize_effect() -> R.Tuple(R.Any): with R.dataflow(): _io: R.Any = R.null_value() lv: R.Tuple(R.Any) = (_io,) gv: R.Tuple(R.Any) = lv R.output(gv) return gv @R.function def test( embedding_table: R.Tensor(("vocab_size", hidden_size), dtype), input_ids: R.Tensor(("seq_len",), "int32"), embedding_dst: R.Tensor(("total_seq_len", hidden_size), dtype), offset: R.Shape(["offset_1"]), packed_params: R.Tuple, ) -> R.Tensor(("total_seq_len", hidden_size), dtype): total_seq_len = T.int64() offset_1 = T.int64() R.func_attr({"num_input": 4}) cls = Expected with R.dataflow(): lv1 = R.call_tir_inplace( cls.inplace_take, (embedding_table, input_ids, embedding_dst), out_ty=R.Tensor((total_seq_len, hidden_size), dtype), inplace_indices=[2], tir_vars=R.shape([offset_1]), ) gv1: R.Tensor((total_seq_len, hidden_size), dtype) = lv1 R.output(gv1) return gv1 m = Model() irmodule, _ = m.export_tvm( spec={ "test": { "embedding_table": spec.Tensor(["vocab_size", hidden_size], dtype), "input_ids": spec.Tensor(["seq_len"], "int32"), "embedding_dst": spec.Tensor(["total_seq_len", hidden_size], dtype), "offset": int, "$": { "param_mode": "packed", "effect_mode": "none", }, }, }, debug=True, ) tvm.ir.assert_structural_equal(irmodule, Expected) def test_tensor_ir_op_no_tir_var(): @T.prim_func(private=True, s_tir=True) def tir_func(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")): T.evaluate(0) class Model(Module): def test(self, A: Tensor): tensor_expr_op_out = op.tensor_ir_op( tir_func, "tir_func", args=[A], out=[Tensor.placeholder((16, 16), "float32")], ) return tensor_expr_op_out @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def tir_func(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")): T.evaluate(0) @R.function def test(A: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"): R.func_attr({"num_input": 1}) cls = Expected with R.dataflow(): lv = R.call_tir(cls.tir_func, (A,), out_ty=R.Tensor((16, 16), dtype="float32")) gv: R.Tensor((16, 16), dtype="float32") = lv R.output(gv) return gv m = Model() irmodule, _ = m.export_tvm(spec={"test": {"A": spec.Tensor([16, 16], "float32")}}) tvm.ir.assert_structural_equal(irmodule, Expected) def test_extern(): class Model(Module): def test(self, q: Tensor, k: Tensor, v: Tensor): b, s, h_q, d = q.shape tensor_expr_op_out = op.extern( name="flashinfer.single_decode", args=[q, k, v, 0, 0, 1.0, 10000.0], out=Tensor.placeholder((b, s, h_q * d), dtype="float16"), ) return tensor_expr_op_out # fmt: off @I.ir_module(s_tir=True) class Expected: @R.function def _initialize_effect() -> R.Tuple(R.Any): with R.dataflow(): _io: R.Any = R.null_value() lv: R.Tuple(R.Any) = (_io,) gv: R.Tuple(R.Any) = lv R.output(gv) return gv @R.function def test(q: R.Tensor((1, 1, 16, 8), dtype="float32"), k: R.Tensor((64, 16, 8), dtype="float32"), v: R.Tensor((64, 16, 8), dtype="float32"), _io: R.Any) -> R.Tuple(R.Tensor((1, 1, 128), dtype="float16"), R.Tuple(R.Any)): R.func_attr({"num_input": 4}) with R.dataflow(): flashinfer_single_decode = R.call_dps_packed("flashinfer.single_decode", (q, k, v, R.prim_value(0), R.prim_value(0), R.prim_value(T.float64(1)), R.prim_value(T.float64(10000))), out_ty=R.Tensor((1, 1, 128), dtype="float16")) gv1: R.Tuple(R.Tensor((1, 1, 128), dtype="float16"), R.Tuple(R.Any)) = flashinfer_single_decode, (_io,) R.output(gv1) return gv1 # fmt: on batch, seq, t, d, h_q, h_kv = 1, 1, 64, 8, 16, 16 m = Model() irmodule, _ = m.export_tvm( spec={ "test": { "q": spec.Tensor([batch, seq, h_q, d], "float32"), "k": spec.Tensor([t, h_kv, d], "float32"), "v": spec.Tensor([t, h_kv, d], "float32"), } }, debug=True, ) tvm.ir.assert_structural_equal(irmodule, Expected) def test_empty(): @tvm.register_global_func("test_empty_assert", override=True) def test_empty_assert(_lineo, x): assert x.shape == (10, 10) assert x.dtype == "float32" class Model(Module): def test(self): result = op.empty([10, 10], dtype="float32") op.debug_func("test_empty_assert", result) return result irmodule, _ = Model().export_tvm(spec={"test": {}}, debug=True) ex = tvm.compile(irmodule, "llvm") vm = relax.VirtualMachine(ex, tvm.cpu()) effects = vm["_initialize_effect"]() vm["test"](*effects) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") def test_multinomial_from_uniform(): prob_shape = (3, 5) sample_shape = (6, 1) class Model(Module): def foo(self, prob: Tensor, uniform_sample: Tensor, sample_indices: Tensor): z0 = op.multinomial_from_uniform(prob, uniform_sample, sample_indices) return z0 # fmt: off @I.ir_module(s_tir=True) class Expected: @R.function def _initialize_effect() -> R.Tuple(R.Any): with R.dataflow(): _io: R.Any = R.null_value() lv: R.Tuple(R.Any) = (_io,) gv: R.Tuple(R.Any) = lv R.output(gv) return gv @R.function def foo(prob: R.Tensor((3, 5), dtype="float32"), uniform_sample: R.Tensor((6, 1), dtype="float32"), sample_indices: R.Tensor((6, 1), dtype="int64"), _io: R.Any) -> R.Tuple(R.Tensor((6, 1), dtype="int64"), R.Tuple(R.Any)): R.func_attr({"num_input": 4}) with R.dataflow(): multinomial_from_uniform: R.Tensor((6, 1), dtype="int64") = R.multinomial_from_uniform(prob, uniform_sample, sample_indices, dtype="int64") gv1: R.Tuple(R.Tensor((6, 1), dtype="int64"), R.Tuple(R.Any)) = multinomial_from_uniform, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() mod, _ = m.export_tvm( spec={ "foo": { "prob": spec.Tensor(prob_shape, "float32"), "uniform_sample": spec.Tensor(sample_shape, "float32"), "sample_indices": spec.Tensor(sample_shape, "int64"), } }, debug=True, ) tvm.ir.assert_structural_equal(mod, Expected) target = tvm.target.Target("cuda", host="llvm") with target: mod = relax.backend.DispatchSampling()(mod) mod = s_tir.transform.DefaultGPUSchedule()(mod) ex = tvm.compile(mod, target) dev = tvm.device(target.kind.name, 0) vm = relax.VirtualMachine(ex, dev) effects = vm["_initialize_effect"]() np_rand = np.random.rand(*prob_shape).astype(np.float32) # normalize it to get the random prob np_prob = np_rand / np_rand.sum(axis=1, keepdims=True) nd_prob = tvm.runtime.tensor(np_prob, dev) # special sample to get deterministic results nd_sample = tvm.runtime.tensor(np.array([[1], [0], [1], [1], [0], [1]]).astype(np.float32), dev) nd_sample_indices = tvm.runtime.tensor( np.array([[0], [1], [1], [2], [2], [2]]).astype(np.int64), dev ) inputs = [nd_prob, nd_sample, nd_sample_indices, effects] res = vm["foo"](*inputs) tvm.testing.assert_allclose( res[0].numpy(), np.array([[4], [0], [4], [4], [0], [4]]).astype(np.int64) ) @pytest.mark.gpu @pytest.mark.skipif(not env.has_gpu(), reason="need gpu") def test_sample_top_p_top_k_from_sorted_prob(): prob_shape = (2, 3) sample_shape = (3, 1) class Model(Module): def foo( self, prob: Tensor, index: Tensor, top_p: Tensor, top_k: Tensor, uniform_sample: Tensor, sample_indices: Tensor, ): z0 = op.sample_top_p_top_k_from_sorted_prob( prob, index, top_p, top_k, uniform_sample, sample_indices ) return z0 # fmt: off @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def get_index_from_sorted(A: T.handle, B: T.handle, C: T.handle, D: T.handle, E: T.handle, F: T.handle): batch, vocab_size = T.int64(), T.int64() cumsum_sorted = T.match_buffer(A, (batch, vocab_size)) indices = T.match_buffer(B, (batch, vocab_size), "int64") renorm_prob = T.match_buffer(C, (batch, 1)) out_batch = T.int64() usample = T.match_buffer(D, (out_batch, 1)) sample_indices = T.match_buffer(E, (out_batch, 1), "int64") output_index = T.match_buffer(F, (out_batch, 1), "int64") # with T.sblock("root"): for ax0, ax1 in T.grid(out_batch, vocab_size): with T.sblock("T_get_index_from_sorted"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(usample[v_ax0, T.int64(0)], cumsum_sorted[sample_indices[v_ax0, T.int64(0)], v_ax1 - T.int64(1):v_ax1 - T.int64(1) + T.int64(2)], sample_indices[v_ax0, T.int64(0)], renorm_prob[sample_indices[v_ax0, T.int64(0)], 0], indices[sample_indices[v_ax0, T.int64(0)], T.min(T.int64(0), v_ax1):T.min(T.int64(0), v_ax1) + (T.max(T.int64(0), v_ax1) + T.int64(1) - T.min(T.int64(0), v_ax1))]) T.writes(output_index[v_ax0, 0]) if usample[v_ax0, T.int64(0)] < cumsum_sorted[sample_indices[v_ax0, T.int64(0)], v_ax1] / renorm_prob[sample_indices[v_ax0, T.int64(0)], 0] or v_ax1 + T.int64(1) == vocab_size: if v_ax1 == T.int64(0): output_index[v_ax0, 0] = indices[sample_indices[v_ax0, T.int64(0)], 0] else: if usample[v_ax0, T.int64(0)] >= cumsum_sorted[sample_indices[v_ax0, T.int64(0)], v_ax1 - T.int64(1)] / renorm_prob[sample_indices[v_ax0, T.int64(0)], 0]: output_index[v_ax0, 0] = indices[sample_indices[v_ax0, T.int64(0)], v_ax1] @T.prim_func(private=True, s_tir=True) def get_renorm_prob(A: T.handle, B: T.handle, C: T.handle, D: T.handle): batch, vocab_size = T.int64(), T.int64() cumsum_sorted = T.match_buffer(A, (batch, vocab_size)) top_p = T.match_buffer(B, (batch, 1)) top_k = T.match_buffer(C, (batch, 1), "int64") renorm_prob = T.match_buffer(D, (batch, 1)) # with T.sblock("root"): for ax0, ax1 in T.grid(batch, vocab_size): with T.sblock("T_get_renorm_prob"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(cumsum_sorted[v_ax0, T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)):T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)) + (T.max(T.max(T.int64(0), v_ax1), v_ax1 + T.int64(1)) + T.int64(1) - T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)))], top_p[v_ax0, 0], top_k[v_ax0, 0]) T.writes(renorm_prob[v_ax0, 0]) if not (cumsum_sorted[v_ax0, 0] < top_p[v_ax0, 0] and top_k[v_ax0, 0] > T.int64(1)): renorm_prob[v_ax0, 0] = cumsum_sorted[v_ax0, 0] else: if cumsum_sorted[v_ax0, v_ax1] < top_p[v_ax0, 0] and v_ax1 + T.int64(1) < top_k[v_ax0, 0]: if v_ax1 + T.int64(1) == vocab_size: renorm_prob[v_ax0, 0] = cumsum_sorted[v_ax0, v_ax1] else: if not (cumsum_sorted[v_ax0, v_ax1 + T.int64(1)] < top_p[v_ax0, 0] and v_ax1 + T.int64(1) + T.int64(1) < top_k[v_ax0, 0]): renorm_prob[v_ax0, 0] = cumsum_sorted[v_ax0, v_ax1 + T.int64(1)] @R.function def _initialize_effect() -> R.Tuple(R.Any): with R.dataflow(): _io: R.Any = R.null_value() lv: R.Tuple(R.Any) = (_io,) gv: R.Tuple(R.Any) = lv R.output(gv) return gv @R.function def foo(prob: R.Tensor((2, 3), dtype="float32"), index: R.Tensor((2, 3), dtype="int64"), top_p: R.Tensor((2, 1), dtype="float32"), top_k: R.Tensor((2, 1), dtype="int64"), uniform_sample: R.Tensor((3, 1), dtype="float32"), sample_indices: R.Tensor((3, 1), dtype="int64"), _io: R.Any,) -> R.Tuple(R.Tensor((3, 1), dtype="int64"), R.Tuple(R.Any)): R.func_attr({"num_input": 7}) cls = Expected with R.dataflow(): cumsum: R.Tensor((2, 3), dtype="float32") = R.cumsum(prob, axis=1, dtype=None, exclusive=None) lv1 = R.call_tir(cls.get_renorm_prob, (cumsum, top_p, top_k), out_ty=R.Tensor((2, 1), dtype="float32")) lv2 = R.call_tir(cls.get_index_from_sorted, (cumsum, index, lv1, uniform_sample, sample_indices), out_ty=R.Tensor((3, 1), dtype="int64")) gv1: R.Tuple(R.Tensor((3, 1), dtype="int64"), R.Tuple(R.Any)) = lv2, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() mod, _ = m.export_tvm( spec={ "foo": { "prob": spec.Tensor(prob_shape, "float32"), "index": spec.Tensor(prob_shape, "int64"), "top_p": spec.Tensor((prob_shape[0], 1), "float32"), "top_k": spec.Tensor((prob_shape[0], 1), "int64"), "uniform_sample": spec.Tensor(sample_shape, "float32"), "sample_indices": spec.Tensor(sample_shape, "int64"), } }, debug=True, ) tvm.ir.assert_structural_equal(mod, Expected) target = tvm.target.Target({"kind": "cuda", "libs": ["thrust"]}, host="llvm") with target: mod = s_tir.transform.DefaultGPUSchedule()(mod) ex = tvm.compile(mod, target) def run_and_check(): dev = tvm.cuda(0) vm = relax.VirtualMachine(ex, dev) effects = vm["_initialize_effect"]() sorted_prob = tvm.runtime.tensor( np.array([[0.5, 0.4, 0.1], [0.4, 0.3, 0.3]]).astype(np.float32), dev ) indices = tvm.runtime.tensor(np.array([[2, 1, 0], [2, 0, 1]]).astype(np.int64), dev) top_p = tvm.runtime.tensor(np.array([[0.6], [0.9]]).astype(np.float32), dev) top_k = tvm.runtime.tensor(np.array([[3], [2]]).astype(np.int64), dev) usample = tvm.runtime.tensor(np.array([[0.5], [0.6], [0.7]]).astype(np.float32), dev) sample_indices = tvm.runtime.tensor(np.array([[0], [1], [1]]).astype(np.int64), dev) inputs = [sorted_prob, indices, top_p, top_k, usample, sample_indices, effects] res = vm["foo"](*inputs) tvm.testing.assert_allclose(res[0].numpy(), np.array([[2], [0], [0]]).astype(np.int64)) tvm.testing.run_with_gpu_lock(run_and_check) @pytest.mark.gpu @pytest.mark.skipif(not env.has_gpu(), reason="need gpu") def test_renormalize_top_p_top_k_prob(): prob_shape = (2, 3) sample_shape = (2, 1) class Model(Module): def foo( self, prob: Tensor, sorted_prob: Tensor, top_p: Tensor, top_k: Tensor, ): z0 = op.renormalize_top_p_top_k_prob(prob, sorted_prob, top_p, top_k) return z0 # fmt: off @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def filter_with_top_p_top_k(A: T.Buffer((T.int64(2), T.int64(3)), "float32"), B: T.Buffer((T.int64(2), T.int64(1)), "float32"), filter_with_top_p_top_k: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for i, j in T.grid(T.int64(2), T.int64(3)): with T.sblock("filter_with_top_p_top_k"): v_i, v_j = T.axis.remap("SS", [i, j]) T.reads(B[v_i, T.int64(0)], A[v_i, v_j]) T.writes(filter_with_top_p_top_k[v_i, v_j]) filter_with_top_p_top_k[v_i, v_j] = T.Select(B[v_i, T.int64(0)] <= A[v_i, v_j], A[v_i, v_j], T.float32(0)) @T.prim_func(private=True, s_tir=True) def get_renorm_cutoff(A: T.handle, B: T.handle, C: T.handle, D: T.handle, E: T.handle): batch, vocab_size = T.int64(), T.int64() sorted_prob = T.match_buffer(A, (batch, vocab_size)) cumsum_sorted = T.match_buffer(B, (batch, vocab_size)) top_p = T.match_buffer(C, (batch, 1)) top_k = T.match_buffer(D, (batch, 1), "int64") cutoff = T.match_buffer(E, (batch, 1)) # with T.sblock("root"): for ax0, ax1 in T.grid(batch, vocab_size): with T.sblock("T_get_renorm_prob"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(cumsum_sorted[v_ax0, T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)):T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)) + (T.max(T.max(T.int64(0), v_ax1), v_ax1 + T.int64(1)) + T.int64(1) - T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)))], top_p[v_ax0, 0], top_k[v_ax0, 0], sorted_prob[v_ax0, T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)):T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)) + (T.max(T.max(T.int64(0), v_ax1), v_ax1 + T.int64(1)) + T.int64(1) - T.min(T.min(T.int64(0), v_ax1), v_ax1 + T.int64(1)))]) T.writes(cutoff[v_ax0, 0]) if (cumsum_sorted[v_ax0, 0] < top_p[v_ax0, 0] and top_k[v_ax0, 0] > T.int64(1)) == T.bool(False): cutoff[v_ax0, 0] = sorted_prob[v_ax0, 0] else: if (cumsum_sorted[v_ax0, v_ax1] < top_p[v_ax0, 0] and v_ax1 + T.int64(1) < top_k[v_ax0, 0]) == T.bool(True): if v_ax1 + T.int64(1) == vocab_size: cutoff[v_ax0, 0] = sorted_prob[v_ax0, v_ax1] else: if (cumsum_sorted[v_ax0, v_ax1 + T.int64(1)] < top_p[v_ax0, 0] and v_ax1 + T.int64(1) + T.int64(1) < top_k[v_ax0, 0]) == T.bool(False): cutoff[v_ax0, 0] = sorted_prob[v_ax0, v_ax1 + T.int64(1)] @R.function def _initialize_effect() -> R.Tuple(R.Any): with R.dataflow(): _io: R.Any = R.null_value() lv: R.Tuple(R.Any) = (_io,) gv: R.Tuple(R.Any) = lv R.output(gv) return gv @R.function def foo(prob: R.Tensor((2, 3), dtype="float32"), sorted_prob: R.Tensor((2, 3), dtype="float32"), top_p: R.Tensor((2, 1), dtype="float32"), top_k: R.Tensor((2, 1), dtype="int64"), _io: R.Any) -> R.Tuple(R.Tensor((2, 3), dtype="float32"), R.Tuple(R.Any)): R.func_attr({"num_input": 5}) cls = Expected with R.dataflow(): cumsum: R.Tensor((2, 3), dtype="float32") = R.cumsum(sorted_prob, axis=1, dtype=None, exclusive=None) lv1 = R.call_tir(cls.get_renorm_cutoff, (sorted_prob, cumsum, top_p, top_k), out_ty=R.Tensor((2, 1), dtype="float32")) lv2 = R.call_tir(cls.filter_with_top_p_top_k, (prob, lv1), out_ty=R.Tensor((2, 3), dtype="float32")) sum: R.Tensor((2, 1), dtype="float32") = R.sum(lv2, axis=[1], keepdims=True) divide: R.Tensor((2, 3), dtype="float32") = R.divide(lv2, sum) gv1: R.Tuple(R.Tensor((2, 3), dtype="float32"), R.Tuple(R.Any)) = divide, (_io,) R.output(gv1) return gv1 # fmt: on m = Model() mod, _ = m.export_tvm( spec={ "foo": { "prob": spec.Tensor(prob_shape, "float32"), "sorted_prob": spec.Tensor(prob_shape, "float32"), "top_p": spec.Tensor(sample_shape, "float32"), "top_k": spec.Tensor(sample_shape, "int64"), } }, debug=True, ) tvm.ir.assert_structural_equal(mod, Expected) target = tvm.target.Target({"kind": "cuda", "libs": ["thrust"]}, host="llvm") with target: mod = relax.transform.LegalizeOps()(mod) mod = s_tir.transform.DefaultGPUSchedule()(mod) ex = tvm.compile(mod, target) def run_and_check(): dev = tvm.cuda(0) vm = relax.VirtualMachine(ex, dev) effects = vm["_initialize_effect"]() prob = tvm.runtime.tensor( np.array([[0.2, 0.3, 0.5], [0.3, 0.3, 0.4]]).astype(np.float32), dev ) sorted_prob = tvm.runtime.tensor( np.array([[0.5, 0.3, 0.2], [0.4, 0.3, 0.3]]).astype(np.float32), dev ) top_p = tvm.runtime.tensor(np.array([[0.6], [0.9]]).astype(np.float32), dev) top_k = tvm.runtime.tensor(np.array([[3], [2]]).astype(np.int64), dev) inputs = [prob, sorted_prob, top_p, top_k, effects] res = vm["foo"](*inputs) tvm.testing.assert_allclose( res[0].numpy(), np.array([[0, 0.375, 0.625], [0.3, 0.3, 0.4]]).astype(np.float32), ) tvm.testing.run_with_gpu_lock(run_and_check) def test_sort_argsort_topk(): class Model(Module): def foo(self, x: Tensor): z0 = op.sort(x, axis=-1, descending=True) z1 = op.argsort(x, axis=-1, descending=False) z2 = op.topk(x, k=2, axis=-1) return z0, z1, z2 @I.ir_module(s_tir=True) class Expected: @R.function def foo(x: R.Tensor(("seq_len", 64), dtype="float16")): R.func_attr({"num_input": 1}) with R.dataflow(): sort = R.sort(x, axis=-1, descending=True) argsort = R.argsort(x, axis=-1, descending=False, dtype="int32") topk = R.topk(x, k=2, axis=-1, ret_type="both", largest=True, dtype="int32") topk_0 = topk[0] topk_1 = topk[1] gv = sort, argsort, (topk_0, topk_1) R.output(gv) return gv m = Model() mod, _ = m.export_tvm({"foo": {"x": spec.Tensor(("seq_len", 64), "float16")}}) tvm.ir.assert_structural_equal(mod, Expected) if __name__ == "__main__": tvm.testing.main()