# 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 import tvm import tvm.testing from tvm.relax.transform import LegalizeOps from tvm.script import ir as I from tvm.script import relax as R from tvm.script import tirx as T ##################### Binary arithmetic ##################### def test_add(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.add(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv = R.call_tir(Expected.add, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def add(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_add"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_add[ax0, ax1, ax2, ax3]) T_add[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_add_with_arg0_constant_scalar(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.add(x, R.const(1, "float32")) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.add, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def add(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_add: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_add"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_add[ax0, ax1]) T_add[ax0, ax1] = rxplaceholder[ax0, ax1] + T.float32(1) # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_add_with_arg1_constant_scalar(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.add(R.const(1, "float32"), x) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.add, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def add(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_add: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_add"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_add[ax0, ax1]) T_add[ax0, ax1] = T.float32(1) + rxplaceholder[ax0, ax1] # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_add_symbolic(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.add(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.add, (x, y), R.Tensor((a, b, c, d), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def add(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_add: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_add = T.match_buffer(var_T_add, [a, b, c, d], dtype="float32") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_add"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_add[ax0, ax1, ax2, ax3]) T_add[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_add_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.add(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.add, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def add( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = lhs[vi, vj, vk] + rhs After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_divide(): # fmt: off @tvm.script.ir_module class Divide: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.divide(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv = R.call_tir(Expected.divide, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def divide(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_divide: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_divide"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_divide[ax0, ax1, ax2, ax3]) T_divide[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] / rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Divide) tvm.ir.assert_structural_equal(mod, Expected) def test_divide_with_arg0_constant_scalar(): # fmt: off @tvm.script.ir_module class Divide: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.divide(x, R.const(1, "float32")) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.divide, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def divide(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_divide: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_divide"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_divide[ax0, ax1]) T_divide[ax0, ax1] = rxplaceholder[ax0, ax1] / T.float32(1) # fmt: on mod = LegalizeOps()(Divide) tvm.ir.assert_structural_equal(mod, Expected) def test_divide_with_arg1_constant_scalar(): # fmt: off @tvm.script.ir_module class Divide: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.divide(R.const(1, "float32"), x) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.divide, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def divide(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_divide: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_divide"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_divide[ax0, ax1]) T_divide[ax0, ax1] = T.float32(1) / rxplaceholder[ax0, ax1] # fmt: on mod = LegalizeOps()(Divide) tvm.ir.assert_structural_equal(mod, Expected) def test_divide_symbolic(): # fmt: off @tvm.script.ir_module class Divide: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.divide(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.divide, (x, y), R.Tensor((a, b, c, d), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def divide(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_divide: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_divide = T.match_buffer(var_T_divide, [a, b, c, d], dtype="float32") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_divide"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_divide[ax0, ax1, ax2, ax3]) T_divide[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] / rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Divide) tvm.ir.assert_structural_equal(mod, Expected) def test_divide_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.divide(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.divide, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def divide( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = lhs[vi, vj, vk] / rhs After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_floor_divide(): # fmt: off @tvm.script.ir_module class FloorDivide: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.floor_divide(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv = R.call_tir(Expected.floor_divide, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def floor_divide(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_floor_divide: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_floor_divide"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_floor_divide[ax0, ax1, ax2, ax3]) T_floor_divide[ax0, ax1, ax2, ax3] = T.floor(rxplaceholder[T.int64(0), ax2, ax3] / rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) # fmt: on mod = LegalizeOps()(FloorDivide) tvm.ir.assert_structural_equal(mod, Expected) def test_floor_divide_with_arg0_constant_scalar(): # fmt: off @tvm.script.ir_module class FloorDivide: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.floor_divide(x, R.const(1, "float32")) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.floor_divide, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def floor_divide(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_floor_divide: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_floor_divide"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_floor_divide[ax0, ax1]) T_floor_divide[ax0, ax1] = T.floor(rxplaceholder[ax0, ax1] / T.float32(1)) # fmt: on mod = LegalizeOps()(FloorDivide) tvm.ir.assert_structural_equal(mod, Expected) def test_floor_divide_with_arg1_constant_scalar(): # fmt: off @tvm.script.ir_module class FloorDivide: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.floor_divide(R.const(1, "float32"), x) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.floor_divide, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def floor_divide(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_floor_divide: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_floor_divide"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_floor_divide[ax0, ax1]) T_floor_divide[ax0, ax1] = T.floor(T.float32(1) / rxplaceholder[ax0, ax1]) # fmt: on mod = LegalizeOps()(FloorDivide) tvm.ir.assert_structural_equal(mod, Expected) def test_floor_divide_symbolic(): # fmt: off @tvm.script.ir_module class FloorDivide: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.floor_divide(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.floor_divide, (x, y), R.Tensor((a, b, c, d), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def floor_divide(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_floor_divide: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_floor_divide = T.match_buffer(var_T_floor_divide, [a, b, c, d], dtype="float32") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_floor_divide"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_floor_divide[ax0, ax1, ax2, ax3]) T_floor_divide[ax0, ax1, ax2, ax3] = T.floor(rxplaceholder[T.int64(0), ax2, ax3] / rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) # fmt: on mod = LegalizeOps()(FloorDivide) tvm.ir.assert_structural_equal(mod, Expected) def test_floordiv_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.floor_divide(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.floor_divide, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def floor_divide( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_floordiv"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = T.floor(lhs[vi, vj, vk] / rhs) After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_multiply(): # fmt: off @tvm.script.ir_module class Multiply: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.multiply(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv = R.call_tir(Expected.multiply, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def multiply(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_multiply: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_multiply"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_multiply[ax0, ax1, ax2, ax3]) T_multiply[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] * rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Multiply) tvm.ir.assert_structural_equal(mod, Expected) def test_multiply_symbolic(): # fmt: off @tvm.script.ir_module class Multiply: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.multiply(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.multiply, (x, y), R.Tensor((a, b, c, d), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def multiply(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_multiply: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_multiply = T.match_buffer(var_T_multiply, [a, b, c, d], dtype="float32") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_multiply"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_multiply[ax0, ax1, ax2, ax3]) T_multiply[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] * rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Multiply) tvm.ir.assert_structural_equal(mod, Expected) def test_multiply_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.multiply(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.multiply, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def multiply( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = lhs[vi, vj, vk] * rhs After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_power(): # fmt: off @tvm.script.ir_module class Power: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.power(x, y) return gv @tvm.script.ir_module class Expected: @T.prim_func(private=True, s_tir=True) def power(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_power: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_power"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(rxplaceholder[T.int64(0), v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, v_ax2, T.int64(0)]) T.writes(T_power[v_ax0, v_ax1, v_ax2, v_ax3]) T_power[v_ax0, v_ax1, v_ax2, v_ax3] = T.pow(rxplaceholder[T.int64(0), v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, v_ax2, T.int64(0)]) @R.function def main(x: R.Tensor((1, 2, 3), dtype="float32"), y: R.Tensor((4, 3, 2, 1), dtype="float32")) -> R.Tensor((4, 3, 2, 3), dtype="float32"): gv = R.call_tir(Expected.power, (x, y), out_ty=R.Tensor((4, 3, 2, 3), dtype="float32")) return gv # fmt: on mod = LegalizeOps()(Power) tvm.ir.assert_structural_equal(mod, Expected) def test_power_symbolic(): # fmt: off @tvm.script.ir_module class Power: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.power(x, y) return gv @tvm.script.ir_module class Expected: @T.prim_func(private=True, s_tir=True) def power(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_power: T.handle): T.func_attr({"tirx.noalias": True}) c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, (T.int64(1), c, d)) a = T.int64() b = T.int64() rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, (a, b, c, T.int64(1))) T_power = T.match_buffer(var_T_power, (a, b, c, d)) # with T.sblock("root"): for ax0, ax1, ax2, ax3 in T.grid(a, b, c, d): with T.sblock("T_power"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(rxplaceholder[T.int64(0), v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, v_ax2, T.int64(0)]) T.writes(T_power[v_ax0, v_ax1, v_ax2, v_ax3]) T_power[v_ax0, v_ax1, v_ax2, v_ax3] = T.pow(rxplaceholder[T.int64(0), v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, v_ax2, T.int64(0)]) @R.function def main(x: R.Tensor((1, "c", "d"), dtype="float32"), y: R.Tensor(("a", "b", "c", 1), dtype="float32")) -> R.Tensor(("a", "b", "c", "d"), dtype="float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.power, (x, y), out_ty=R.Tensor((a, b, c, d), dtype="float32")) return gv # fmt: on mod = LegalizeOps()(Expected) tvm.ir.assert_structural_equal(mod, Expected) def test_power_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.power(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.power, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def power( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_power"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = T.pow(lhs[vi, vj, vk], rhs) After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_atan2(): # fmt: off @tvm.script.ir_module class Atan2: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.atan2(x, y) return gv @tvm.script.ir_module class Expected: @T.prim_func(private=True, s_tir=True) def atan2(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_atan2: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_atan2"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(rxplaceholder[T.int64(0), v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, v_ax2, T.int64(0)]) T.writes(T_atan2[v_ax0, v_ax1, v_ax2, v_ax3]) T_atan2[v_ax0, v_ax1, v_ax2, v_ax3] = T.atan2(rxplaceholder[T.int64(0), v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, v_ax2, T.int64(0)]) @R.function def main(x: R.Tensor((1, 2, 3), dtype="float32"), y: R.Tensor((4, 3, 2, 1), dtype="float32")) -> R.Tensor((4, 3, 2, 3), dtype="float32"): gv = R.call_tir(Expected.atan2, (x, y), out_ty=R.Tensor((4, 3, 2, 3), dtype="float32")) return gv # fmt: on mod = LegalizeOps()(Atan2) tvm.ir.assert_structural_equal(mod, Expected) def test_atan2_symbolic(): # fmt: off @tvm.script.ir_module class Atan2: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.atan2(x, y) return gv @tvm.script.ir_module class Expected: @T.prim_func(private=True, s_tir=True) def atan2(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_atan2: T.handle): T.func_attr({"tirx.noalias": True}) c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, (T.int64(1), c, d)) a = T.int64() b = T.int64() rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, (a, b, c, T.int64(1))) T_atan2 = T.match_buffer(var_T_atan2, (a, b, c, d)) for ax0, ax1, ax2, ax3 in T.grid(a, b, c, d): with T.sblock("T_atan2"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(rxplaceholder[T.int64(0), v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, v_ax2, T.int64(0)]) T.writes(T_atan2[v_ax0, v_ax1, v_ax2, v_ax3]) T_atan2[v_ax0, v_ax1, v_ax2, v_ax3] = T.atan2(rxplaceholder[T.int64(0), v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, v_ax2, T.int64(0)]) @R.function def main(x: R.Tensor((1, "c", "d"), dtype="float32"), y: R.Tensor(("a", "b", "c", 1), dtype="float32")) -> R.Tensor(("a", "b", "c", "d"), dtype="float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.atan2, (x, y), out_ty=R.Tensor((a, b, c, d), dtype="float32")) return gv # fmt: on mod = LegalizeOps()(Expected) tvm.ir.assert_structural_equal(mod, Expected) def test_atan2_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.atan2(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.atan2, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def atan2( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_atan2"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = T.atan2(lhs[vi, vj, vk], rhs) After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_subtract(): # fmt: off @tvm.script.ir_module class Subtract: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.subtract(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv = R.call_tir(Expected.subtract, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def subtract(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_subtract: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_subtract"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_subtract[ax0, ax1, ax2, ax3]) T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] - rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Subtract) tvm.ir.assert_structural_equal(mod, Expected) def test_subtract_symbolic(): # fmt: off @tvm.script.ir_module class Subtract: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.subtract(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.subtract, (x, y), R.Tensor((a, b, c, d), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def subtract(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_subtract: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_subtract = T.match_buffer(var_T_subtract, [a, b, c, d], dtype="float32") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_subtract"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_subtract[ax0, ax1, ax2, ax3]) T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] - rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Subtract) tvm.ir.assert_structural_equal(mod, Expected) def test_subtract_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.subtract(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.subtract, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def subtract( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = lhs[vi, vj, vk] - rhs After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) ##################### Binary comparison ##################### def test_equal(): # fmt: off @tvm.script.ir_module class Equal: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv: R.Tensor((4, 3, 2, 3), "bool") = R.equal(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv = R.call_tir(Expected.equal, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def equal(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_equal: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_equal"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_equal[ax0, ax1, ax2, ax3]) T_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] == rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Equal) tvm.ir.assert_structural_equal(mod, Expected) def test_equal_with_arg0_constant_scalar(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv: R.Tensor((2, 3), dtype="bool") = R.equal(x, R.const(1, "float32")) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv = R.call_tir(Expected.equal, (x,), R.Tensor((2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def equal(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_equal: T.Buffer((T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_equal"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_equal[ax0, ax1]) T_equal[ax0, ax1] = rxplaceholder[ax0, ax1] == T.float32(1) # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_equal_with_arg1_constant_scalar(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv: R.Tensor((2, 3), dtype="bool") = R.equal(R.const(1, "float32"), x) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv = R.call_tir(Expected.equal, (x,), R.Tensor((2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def equal(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_equal: T.Buffer((T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_equal"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_equal[ax0, ax1]) T_equal[ax0, ax1] = T.float32(1) == rxplaceholder[ax0, ax1] # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_equal_symbolic(): # fmt: off @tvm.script.ir_module class Equal: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "bool") = R.equal(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.equal, (x, y), R.Tensor((a, b, c, d), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def equal(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_equal: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_equal = T.match_buffer(var_T_equal, [a, b, c, d], dtype="bool") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_equal"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_equal[ax0, ax1, ax2, ax3]) T_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] == rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Equal) tvm.ir.assert_structural_equal(mod, Expected) def test_equal_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.equal(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.equal, (x, y), R.Tensor([64, 32, 16], dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def equal( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "bool"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = lhs[vi, vj, vk] == rhs After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_greater(): # fmt: off @tvm.script.ir_module class Greater: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv: R.Tensor((4, 3, 2, 3), "bool") = R.greater(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv = R.call_tir(Expected.greater, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def greater(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_greater: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_greater"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3]) T.writes(T_greater[ax0, ax1, ax2, ax3]) T_greater[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] < rxplaceholder[T.int64(0), ax2, ax3] # fmt: on mod = LegalizeOps()(Greater) tvm.ir.assert_structural_equal(mod, Expected) def test_greater_with_arg0_constant_scalar(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv: R.Tensor((2, 3), dtype="bool") = R.greater(x, R.const(1, "float32")) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv = R.call_tir(Expected.greater, (x,), R.Tensor((2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def greater(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_greater: T.Buffer((T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_greater"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_greater[ax0, ax1]) T_greater[ax0, ax1] = T.float32(1) < rxplaceholder[ax0, ax1] # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_greater_with_arg1_constant_scalar(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv: R.Tensor((2, 3), dtype="bool") = R.greater(R.const(1, "float32"), x) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv = R.call_tir(Expected.greater, (x,), R.Tensor((2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def greater(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_greater: T.Buffer((T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_greater"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_greater[ax0, ax1]) T_greater[ax0, ax1] = rxplaceholder[ax0, ax1] < T.float32(1) # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_greater_symbolic(): # fmt: off @tvm.script.ir_module class Greater: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "bool") = R.greater(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.greater, (x, y), R.Tensor((a, b, c, d), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def greater(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_greater: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_greater = T.match_buffer(var_T_greater, [a, b, c, d], dtype="bool") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_greater"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3]) T.writes(T_greater[ax0, ax1, ax2, ax3]) T_greater[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] < rxplaceholder[T.int64(0), ax2, ax3] # fmt: on mod = LegalizeOps()(Greater) tvm.ir.assert_structural_equal(mod, Expected) def test_greater_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.greater(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.greater, (x, y), R.Tensor([64, 32, 16], dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def greater( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "bool"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = rhs < lhs[vi, vj, vk] After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_greater_equal(): # fmt: off @tvm.script.ir_module class GreaterEqual: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv: R.Tensor((4, 3, 2, 3), "bool") = R.greater_equal(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv = R.call_tir(Expected.greater_equal, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def greater_equal(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_greater_equal: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_greater_equal"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3]) T.writes(T_greater_equal[ax0, ax1, ax2, ax3]) T_greater_equal[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] <= rxplaceholder[T.int64(0), ax2, ax3] # fmt: on mod = LegalizeOps()(GreaterEqual) tvm.ir.assert_structural_equal(mod, Expected) def test_greater_equal_symbolic(): # fmt: off @tvm.script.ir_module class GreaterEqual: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "bool") = R.greater_equal(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.greater_equal, (x, y), R.Tensor((a, b, c, d), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def greater_equal(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_greater_equal: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_greater_equal = T.match_buffer(var_T_greater_equal, [a, b, c, d], dtype="bool") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_greater_equal"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3]) T.writes(T_greater_equal[ax0, ax1, ax2, ax3]) T_greater_equal[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] <= rxplaceholder[T.int64(0), ax2, ax3] # fmt: on mod = LegalizeOps()(GreaterEqual) tvm.ir.assert_structural_equal(mod, Expected) def test_greater_equal_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.greater_equal(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.greater_equal, (x, y), R.Tensor([64, 32, 16], dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def greater_equal( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "bool"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = rhs <= lhs[vi, vj, vk] After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_less(): # fmt: off @tvm.script.ir_module class Less: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv: R.Tensor((4, 3, 2, 3), "bool") = R.less(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv = R.call_tir(Expected.less, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def less(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_less: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_less"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_less[ax0, ax1, ax2, ax3]) T_less[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] < rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Less) tvm.ir.assert_structural_equal(mod, Expected) def test_less_symbolic(): # fmt: off @tvm.script.ir_module class Less: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "bool") = R.less(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.less, (x, y), R.Tensor((a, b, c, d), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def less(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_less: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_less = T.match_buffer(var_T_less, [a, b, c, d], dtype="bool") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_less"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_less[ax0, ax1, ax2, ax3]) T_less[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] < rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(Less) tvm.ir.assert_structural_equal(mod, Expected) def test_less_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.less(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.less, (x, y), R.Tensor([64, 32, 16], dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def less( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "bool"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = lhs[vi, vj, vk] < rhs After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_less_equal(): # fmt: off @tvm.script.ir_module class LessEqual: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv: R.Tensor((4, 3, 2, 3), "bool") = R.less_equal(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv = R.call_tir(Expected.less_equal, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def less_equal(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_less_equal: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_less_equal"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_less_equal[ax0, ax1, ax2, ax3]) T_less_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] <= rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(LessEqual) tvm.ir.assert_structural_equal(mod, Expected) def test_less_equal_with_arg0_constant_scalar(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv: R.Tensor((2, 3), dtype="bool") = R.less_equal(x, R.const(1, "float32")) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv = R.call_tir(Expected.less_equal, (x,), R.Tensor((2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def less_equal(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_less_equal: T.Buffer((T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_less_equal"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_less_equal[ax0, ax1]) T_less_equal[ax0, ax1] = rxplaceholder[ax0, ax1] <= T.float32(1) # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_less_equal_with_arg1_constant_scalar(): # fmt: off @tvm.script.ir_module class Add: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv: R.Tensor((2, 3), dtype="bool") = R.less_equal(R.const(1, "float32"), x) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"): gv = R.call_tir(Expected.less_equal, (x,), R.Tensor((2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def less_equal(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_less_equal: T.Buffer((T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_less_equal"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_less_equal[ax0, ax1]) T_less_equal[ax0, ax1] = T.float32(1) <= rxplaceholder[ax0, ax1] # fmt: on mod = LegalizeOps()(Add) tvm.ir.assert_structural_equal(mod, Expected) def test_less_equal_symbolic(): # fmt: off @tvm.script.ir_module class LessEqual: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "bool") = R.less_equal(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.less_equal, (x, y), R.Tensor((a, b, c, d), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def less_equal(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_less_equal: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_less_equal = T.match_buffer(var_T_less_equal, [a, b, c, d], dtype="bool") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_less_equal"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_less_equal[ax0, ax1, ax2, ax3]) T_less_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] <= rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(LessEqual) tvm.ir.assert_structural_equal(mod, Expected) def test_less_equal_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.less_equal(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.less_equal, (x, y), R.Tensor([64, 32, 16], dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def less_equal( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "bool"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = lhs[vi, vj, vk] <= rhs After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_not_equal(): # fmt: off @tvm.script.ir_module class NotEqual: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv: R.Tensor((4, 3, 2, 3), "bool") = R.not_equal(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"): gv = R.call_tir(Expected.not_equal, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def not_equal(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_not_equal: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_not_equal"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_not_equal[ax0, ax1, ax2, ax3]) T_not_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] != rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(NotEqual) tvm.ir.assert_structural_equal(mod, Expected) def test_not_equal_symbolic(): # fmt: off @tvm.script.ir_module class NotEqual: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "bool") = R.not_equal(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.not_equal, (x, y), R.Tensor((a, b, c, d), dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def not_equal(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_not_equal: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_not_equal = T.match_buffer(var_T_not_equal, [a, b, c, d], dtype="bool") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_not_equal"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_not_equal[ax0, ax1, ax2, ax3]) T_not_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] != rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] # fmt: on mod = LegalizeOps()(NotEqual) tvm.ir.assert_structural_equal(mod, Expected) def test_not_equal_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.not_equal(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.not_equal, (x, y), R.Tensor([64, 32, 16], dtype="bool")) return gv @T.prim_func(private=True, s_tir=True) def not_equal( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "bool"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = lhs[vi, vj, vk] != rhs After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_maximum(): # fmt: off @tvm.script.ir_module class Maximum: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.maximum(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv = R.call_tir(Expected.maximum, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def maximum(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_maximum: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_maximum"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_maximum[ax0, ax1, ax2, ax3]) T_maximum[ax0, ax1, ax2, ax3] = T.max(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) # fmt: on mod = LegalizeOps()(Maximum) tvm.ir.assert_structural_equal(mod, Expected) def test_maximum_with_arg0_constant_scalar(): # fmt: off @tvm.script.ir_module class Maximum: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.maximum(x, R.const(1, "float32")) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.maximum, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def maximum(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_maximum: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_maximum"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_maximum[ax0, ax1]) T_maximum[ax0, ax1] = T.max(rxplaceholder[ax0, ax1], T.float32(1)) # fmt: on mod = LegalizeOps()(Maximum) tvm.ir.assert_structural_equal(mod, Expected) def test_maximum_with_arg1_constant_scalar(): # fmt: off @tvm.script.ir_module class Maximum: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.maximum(R.const(1, "float32"), x) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.maximum, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def maximum(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_maximum: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_maximum"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_maximum[ax0, ax1]) T_maximum[ax0, ax1] = T.max(T.float32(1), rxplaceholder[ax0, ax1]) # fmt: on mod = LegalizeOps()(Maximum) tvm.ir.assert_structural_equal(mod, Expected) def test_maximum_symbolic(): # fmt: off @tvm.script.ir_module class Maximum: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.maximum(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.maximum, (x, y), R.Tensor((a, b, c, d), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def maximum(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_maximum: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_maximum = T.match_buffer(var_T_maximum, [a, b, c, d], dtype="float32") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_maximum"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_maximum[ax0, ax1, ax2, ax3]) T_maximum[ax0, ax1, ax2, ax3] = T.max(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) # fmt: on mod = LegalizeOps()(Maximum) tvm.ir.assert_structural_equal(mod, Expected) def test_max_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.maximum(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.maximum, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def maximum( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = T.max(lhs[vi, vj, vk], rhs) After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) def test_minimum(): # fmt: off @tvm.script.ir_module class Minimum: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv: R.Tensor((4, 3, 2, 3), "float32") = R.minimum(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"): gv = R.call_tir(Expected.minimum, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def minimum(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_minimum: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)): with T.sblock("T_minimum"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_minimum[ax0, ax1, ax2, ax3]) T_minimum[ax0, ax1, ax2, ax3] = T.min(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) # fmt: on mod = LegalizeOps()(Minimum) tvm.ir.assert_structural_equal(mod, Expected) def test_minimum_with_arg0_constant_scalar(): # fmt: off @tvm.script.ir_module class Minimum: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.minimum(x, R.const(1, "float32")) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.minimum, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def minimum(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_minimum: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_minimum"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_minimum[ax0, ax1]) T_minimum[ax0, ax1] = T.min(rxplaceholder[ax0, ax1], T.float32(1)) # fmt: on mod = LegalizeOps()(Minimum) tvm.ir.assert_structural_equal(mod, Expected) def test_minimum_with_arg1_constant_scalar(): # fmt: off @tvm.script.ir_module class Minimum: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv: R.Tensor((2, 3), dtype="float32") = R.minimum(R.const(1, "float32"), x) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"): gv = R.call_tir(Expected.minimum, (x,), R.Tensor((2, 3), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def minimum(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_minimum: T.Buffer((T.int64(2), T.int64(3)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_minimum"): ax0, ax1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[ax0, ax1]) T.writes(T_minimum[ax0, ax1]) T_minimum[ax0, ax1] = T.min(T.float32(1), rxplaceholder[ax0, ax1]) # fmt: on mod = LegalizeOps()(Minimum) tvm.ir.assert_structural_equal(mod, Expected) def test_minimum_symbolic(): # fmt: off @tvm.script.ir_module class Minimum: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv: R.Tensor((a, b, c, d), "float32") = R.minimum(x, y) return gv @tvm.script.ir_module class Expected: @R.function def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"): a = T.int64() b = T.int64() c = T.int64() d = T.int64() gv = R.call_tir(Expected.minimum, (x, y), R.Tensor((a, b, c, d), dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def minimum(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_minimum: T.handle): T.func_attr({"tirx.noalias": True}) a = T.int64() b = T.int64() c = T.int64() d = T.int64() rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32") rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32") T_minimum = T.match_buffer(var_T_minimum, [a, b, c, d], dtype="float32") for i0, i1, i2, i3 in T.grid(a, b, c, d): with T.sblock("T_minimum"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) T.writes(T_minimum[ax0, ax1, ax2, ax3]) T_minimum[ax0, ax1, ax2, ax3] = T.min(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]) # fmt: on mod = LegalizeOps()(Minimum) tvm.ir.assert_structural_equal(mod, Expected) def test_min_primvalue(): @I.ir_module(s_tir=True) class Before: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): gv = R.minimum(x, y) return gv @I.ir_module(s_tir=True) class Expected: @R.function def main( x: R.Tensor([64, 32, 16], "float32"), y: R.Prim("float32"), ): cls = Expected gv = R.call_tir(cls.minimum, (x, y), R.Tensor([64, 32, 16], dtype="float32")) return gv @T.prim_func(private=True, s_tir=True) def minimum( lhs: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), rhs: T.float32, output: T.Buffer([T.int64(64), T.int64(32), T.int64(16)], "float32"), ): T.func_attr({"tirx.noalias": True}) for i, j, k in T.grid(*lhs.shape): with T.sblock("T_add"): vi, vj, vk = T.axis.remap("SSS", [i, j, k]) output[vi, vj, vk] = T.min(lhs[vi, vj, vk], rhs) After = LegalizeOps()(Before) tvm.ir.assert_structural_equal(Expected, After) if __name__ == "__main__": tvm.testing.main()