wehub-resource-sync
2263 lines
92 KiB
Python
2263 lines
92 KiB
Python
# 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()
|