wehub-resource-sync
1949 lines
84 KiB
Python
1949 lines
84 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, E731, F841
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
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
|
|
|
|
##################### Manipulation #####################
|
|
|
|
|
|
def test_broadcast_to():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class BroadcastTo:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 1, 3), "float32")) -> R.Tensor((4, 2, 5, 3), "float32"):
|
|
gv: R.Tensor((4, 2, 5, 3), "float32") = R.broadcast_to(x, (4, 2, 5, 3))
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 1, 3), "float32")) -> R.Tensor((4, 2, 5, 3), "float32"):
|
|
gv = R.call_tir(Expected.broadcast_to, (x,), R.Tensor((4, 2, 5, 3), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def broadcast_to(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3)), "float32"), T_broadcast_to: T.Buffer((T.int64(4), T.int64(2), T.int64(5), T.int64(3)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(2), T.int64(5), T.int64(3)):
|
|
with T.sblock("T_broadcast_to"):
|
|
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
T.reads(rxplaceholder[ax1, T.int64(0), ax3])
|
|
T.writes(T_broadcast_to[ax0, ax1, ax2, ax3])
|
|
T_broadcast_to[ax0, ax1, ax2, ax3] = rxplaceholder[ax1, T.int64(0), ax3]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(BroadcastTo)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_broadcast_to_symbolic():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class BroadcastTo:
|
|
@R.function
|
|
def main(dumb_param: R.Tensor(("a", "c")), x: R.Tensor(("b", 1, "d"), "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.broadcast_to(x, (a, b, c, d))
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(dumb_param: R.Tensor(("a", "c")), x: R.Tensor(("b", 1, "d"), "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.broadcast_to, (x,), R.Tensor((a, b, c, d), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def broadcast_to(var_rxplaceholder: T.handle, var_T_broadcast_to: 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, [b, T.int64(1), d], dtype="float32")
|
|
T_broadcast_to = T.match_buffer(var_T_broadcast_to, [a, b, c, d], dtype="float32")
|
|
for i0, i1, i2, i3 in T.grid(a, b, c, d):
|
|
with T.sblock("T_broadcast_to"):
|
|
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
T.reads(rxplaceholder[ax1, T.int64(0), ax3])
|
|
T.writes(T_broadcast_to[ax0, ax1, ax2, ax3])
|
|
T_broadcast_to[ax0, ax1, ax2, ax3] = rxplaceholder[ax1, T.int64(0), ax3]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(BroadcastTo)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_concat():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Concat:
|
|
@R.function
|
|
def main(x1: R.Tensor((1, 2, 3), "float32"), x2: R.Tensor((1, 3, 3), "float32"), x3: R.Tensor((1, 4, 3), "float32")) -> R.Tensor((1, 9, 3), "float32"):
|
|
gv: R.Tensor((1, 9, 3), "float32") = R.concat((x1, x2, x3), axis=1)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x1: R.Tensor((1, 2, 3), "float32"), x2: R.Tensor((1, 3, 3), "float32"), x3: R.Tensor((1, 4, 3), "float32")) -> R.Tensor((1, 9, 3), "float32"):
|
|
gv = R.call_tir(Expected.concatenate, (x1, x2, x3), R.Tensor((1, 9, 3), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def concatenate(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(1), T.int64(3), T.int64(3)), "float32"), rxplaceholder_2: T.Buffer((T.int64(1), T.int64(4), T.int64(3)), "float32"), T_concat: T.Buffer((T.int64(1), T.int64(9), T.int64(3)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2 in T.grid(T.int64(1), T.int64(9), T.int64(3)):
|
|
with T.sblock("T_concat"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder_2[ax0, ax1 - T.int64(5), ax2], rxplaceholder_1[ax0, ax1 - T.int64(2), ax2], rxplaceholder[ax0, ax1, ax2])
|
|
T.writes(T_concat[ax0, ax1, ax2])
|
|
T_concat[ax0, ax1, ax2] = T.if_then_else(T.int64(5) <= ax1, rxplaceholder_2[ax0, ax1 - T.int64(5), ax2], T.if_then_else(T.int64(2) <= ax1, rxplaceholder_1[ax0, ax1 - T.int64(2), ax2], rxplaceholder[ax0, ax1, ax2]))
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Concat)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_concat_input_tuple_var():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Concat:
|
|
@R.function
|
|
def main(t: R.Tuple(R.Tensor((3, 4), "float32"), R.Tensor((3, 5), "float32"))) -> R.Tensor((3, 9), "float32"):
|
|
gv: R.Tensor((3, 9), "float32") = R.concat(t, axis=1)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(t: R.Tuple(R.Tensor((3, 4), "float32"), R.Tensor((3, 5), "float32"))) -> R.Tensor((3, 9), "float32"):
|
|
gv: R.Tensor((3, 4), dtype="float32") = t[0]
|
|
gv1: R.Tensor((3, 5), dtype="float32") = t[1]
|
|
gv2 = R.call_tir(Expected.concatenate, (gv, gv1), R.Tensor((3, 9), dtype="float32"))
|
|
return gv2
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def concatenate(rxplaceholder: T.Buffer((T.int64(3), T.int64(4)), "float32"), rxplaceholder_1: T.Buffer((T.int64(3), T.int64(5)), "float32"), T_concat: T.Buffer((T.int64(3), T.int64(9)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(3), T.int64(9)):
|
|
with T.sblock("T_concat"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(rxplaceholder_1[ax0, ax1 - T.int64(4)], rxplaceholder[ax0, ax1])
|
|
T.writes(T_concat[ax0, ax1])
|
|
T_concat[ax0, ax1] = T.if_then_else(T.int64(4) <= ax1, rxplaceholder_1[ax0, ax1 - T.int64(4)], rxplaceholder[ax0, ax1])
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Concat)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_concat_input_tuple_var_symbolic():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Concat:
|
|
@R.function
|
|
def main(t: R.Tuple(R.Tensor(("a", "b0"), "float32"), R.Tensor(("a", "b1"), "float32"), R.Tensor(("a", "b2"), "float32"))) -> R.Tensor(("a", "b0 + b1 + b2"), "float32"):
|
|
a = T.int64()
|
|
b0 = T.int64()
|
|
b1 = T.int64()
|
|
b2 = T.int64()
|
|
gv: R.Tensor((a, b0 + b1 + b2), "float32") = R.concat(t, axis=1)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(t: R.Tuple(R.Tensor(("a", "b0"), "float32"), R.Tensor(("a", "b1"), "float32"), R.Tensor(("a", "b2"), "float32"))) -> R.Tensor(("a", "b0 + b1 + b2"), "float32"):
|
|
a = T.int64()
|
|
b0 = T.int64()
|
|
b1 = T.int64()
|
|
b2 = T.int64()
|
|
gv: R.Tensor((a, b0), dtype="float32") = t[0]
|
|
gv1: R.Tensor((a, b1), dtype="float32") = t[1]
|
|
gv2: R.Tensor((a, b2), dtype="float32") = t[2]
|
|
gv3 = R.call_tir(Expected.concatenate, (gv, gv1, gv2), R.Tensor((a, ((b0 + b1) + b2)), dtype="float32"))
|
|
return gv3
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def concatenate(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_rxplaceholder_2: T.handle, var_T_concat: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a = T.int64()
|
|
b0 = T.int64()
|
|
b1 = T.int64()
|
|
b2 = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b0], dtype="float32")
|
|
rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b1], dtype="float32")
|
|
rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, [a, b2], dtype="float32")
|
|
T_concat = T.match_buffer(var_T_concat, [a, b0 + b1 + b2], dtype="float32")
|
|
for i0, i1 in T.grid(a, b0 + b1 + b2):
|
|
with T.sblock("T_concat"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(rxplaceholder_2[ax0, ax1 - b0 - b1], rxplaceholder_1[ax0, ax1 - b0], rxplaceholder[ax0, ax1])
|
|
T.writes(T_concat[ax0, ax1])
|
|
T_concat[ax0, ax1] = T.if_then_else(T.int64(0) <= ax1 - b0 - b1, rxplaceholder_2[ax0, ax1 - b0 - b1], T.if_then_else(T.int64(0) <= ax1 - b0, rxplaceholder_1[ax0, ax1 - b0], rxplaceholder[ax0, ax1]))
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Concat)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_expand_dims():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class ExpandDims:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32"):
|
|
gv: R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32") = R.expand_dims(x, axis=[-1, 1, -6, 3, 5])
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32"):
|
|
gv = R.call_tir(Expected.expand_dims, (x,), R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def expand_dims(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), expand_dims: T.Buffer((T.int64(2), T.int64(1), T.int64(1), T.int64(1), T.int64(3), T.int64(1), T.int64(4), T.int64(1)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(T.int64(2), T.int64(1), T.int64(1), T.int64(1), T.int64(3), T.int64(1), T.int64(4), T.int64(1)):
|
|
with T.sblock("expand_dims"):
|
|
i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1 = T.axis.remap("SSSSSSSS", [i0, i1, i2, i3, i4, i5, i6, i7])
|
|
T.reads(rxplaceholder[i0_1, i4_1, i6_1])
|
|
T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1])
|
|
expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1] = rxplaceholder[i0_1, i4_1, i6_1]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(ExpandDims)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_expand_dims_symbolic():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class ExpandDims:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a", 1, "b", 1, "c", 1), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
gv: R.Tensor((a, 1, b, 1, c, 1), "float32") = R.expand_dims(x, axis=[1, 3, 5])
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a", 1, "b", 1, "c", 1), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
gv = R.call_tir(Expected.expand_dims, (x,), R.Tensor((a, 1, b, 1, c, 1), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def expand_dims(var_rxplaceholder: T.handle, var_expand_dims: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c], dtype="float32")
|
|
expand_dims = T.match_buffer(var_expand_dims, [a, T.int64(1), b, T.int64(1), c, T.int64(1)], dtype="float32")
|
|
for i0, i1, i2, i3, i4, i5 in T.grid(a, T.int64(1), b, T.int64(1), c, T.int64(1)):
|
|
with T.sblock("expand_dims"):
|
|
i0_1, i1_1, i2_1, i3_1, i4_1, i5_1 = T.axis.remap("SSSSSS", [i0, i1, i2, i3, i4, i5])
|
|
T.reads(rxplaceholder[i0_1, i2_1, i4_1])
|
|
T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1])
|
|
expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1] = rxplaceholder[i0_1, i2_1, i4_1]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(ExpandDims)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_flatten():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Flatten:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((24,), "float32"):
|
|
gv: R.Tensor((24,), "float32") = R.flatten(x)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((24,), "float32"):
|
|
gv = R.call_tir(Expected.reshape, (x,), R.Tensor((24,), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), T_reshape: T.Buffer(T.int64(24), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0 in T.serial(T.int64(24)):
|
|
with T.sblock("T_reshape"):
|
|
ax0 = T.axis.spatial(T.int64(24), i0)
|
|
T.reads(rxplaceholder[ax0 % T.int64(24) // T.int64(12), ax0 % T.int64(12) // T.int64(4), ax0 % T.int64(4)])
|
|
T.writes(T_reshape[ax0])
|
|
T_reshape[ax0] = rxplaceholder[ax0 % T.int64(24) // T.int64(12), ax0 % T.int64(12) // T.int64(4), ax0 % T.int64(4)]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Flatten)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_flatten_zero_rank():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Flatten:
|
|
@R.function
|
|
def main(x: R.Tensor((), "float32")) -> R.Tensor((1,), "float32"):
|
|
gv: R.Tensor((1,), "float32") = R.flatten(x)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((), "float32")) -> R.Tensor((1,), "float32"):
|
|
gv = R.call_tir(Expected.reshape, (x,), R.Tensor((1,), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(rxplaceholder: T.Buffer((), "float32"), T_reshape: T.Buffer(T.int64(1), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0 in T.serial(T.int64(1)):
|
|
with T.sblock("T_reshape"):
|
|
ax0 = T.axis.spatial(T.int64(1), i0)
|
|
T.reads(rxplaceholder[()])
|
|
T.writes(T_reshape[ax0])
|
|
T_reshape[ax0] = rxplaceholder[()]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Flatten)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_flatten_symbolic():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Flatten:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a * b * c",), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
gv: R.Tensor((a * b * c,), "float32") = R.flatten(x)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a * b * c",), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
gv = R.call_tir(Expected.reshape, (x,), R.Tensor((((a * b) * c),), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c], dtype="float32")
|
|
T_reshape = T.match_buffer(var_T_reshape, [a * b * c], dtype="float32")
|
|
for i0 in T.serial(a * b * c):
|
|
with T.sblock("T_reshape"):
|
|
ax0 = T.axis.spatial(a * b * c, i0)
|
|
T.reads(rxplaceholder[ax0 // c // b % a, ax0 // c % b, ax0 % c])
|
|
T.writes(T_reshape[ax0])
|
|
T_reshape[ax0] = rxplaceholder[ax0 // c // b % a, ax0 // c % b, ax0 % c]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Flatten)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_permute_dims():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class PermuteDims:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((2, 4, 3, 1), "float32"):
|
|
gv: R.Tensor((2, 4, 3, 1), "float32") = R.permute_dims(x, axes=[1, -1, 2, -4])
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((2, 4, 3, 1), "float32"):
|
|
gv = R.call_tir(Expected.transpose, (x,), R.Tensor((2, 4, 3, 1), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def transpose(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"), T_transpose: T.Buffer((T.int64(2), T.int64(4), T.int64(3), T.int64(1)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(4), T.int64(3), T.int64(1)):
|
|
with T.sblock("T_transpose"):
|
|
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
T.reads(rxplaceholder[ax3, ax0, ax2, ax1])
|
|
T.writes(T_transpose[ax0, ax1, ax2, ax3])
|
|
T_transpose[ax0, ax1, ax2, ax3] = rxplaceholder[ax3, ax0, ax2, ax1]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(PermuteDims)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_permute_dims_symbolic():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class PermuteDims:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("b", "d", "c", "a"), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
d = T.int64()
|
|
gv: R.Tensor((b, d, c, a), "float32") = R.permute_dims(x, axes=[1, -1, 2, -4])
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c", "d"), dtype="float32")) -> R.Tensor(("b", "d", "c", "a"), dtype="float32"):
|
|
b = T.int64()
|
|
d = T.int64()
|
|
c = T.int64()
|
|
a = T.int64()
|
|
gv = R.call_tir(Expected.transpose, (x,), R.Tensor((b, d, c, a), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def transpose(var_rxplaceholder: T.handle, var_T_transpose: 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, [a, b, c, d], dtype="float32")
|
|
T_transpose = T.match_buffer(var_T_transpose, [b, d, c, a], dtype="float32")
|
|
for i0, i1, i2, i3 in T.grid(b, d, c, a):
|
|
with T.sblock("T_transpose"):
|
|
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
T.reads(rxplaceholder[ax3, ax0, ax2, ax1])
|
|
T.writes(T_transpose[ax0, ax1, ax2, ax3])
|
|
T_transpose[ax0, ax1, ax2, ax3] = rxplaceholder[ax3, ax0, ax2, ax1]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(PermuteDims)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_reshape():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Reshape:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"):
|
|
gv: R.Tensor((8, 3), "float32") = R.reshape(x, (8, 3))
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"):
|
|
gv = R.call_tir(Expected.reshape, (x,), R.Tensor((8, 3), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"), T_reshape: T.Buffer((T.int64(8), T.int64(3)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(8), T.int64(3)):
|
|
with T.sblock("T_reshape"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(rxplaceholder[T.int64(0), (ax0 * T.int64(3) + ax1) % T.int64(24) // T.int64(12), (ax0 * T.int64(3) + ax1) % T.int64(12) // T.int64(4), (ax0 * T.int64(3) + ax1) % T.int64(4)])
|
|
T.writes(T_reshape[ax0, ax1])
|
|
T_reshape[ax0, ax1] = rxplaceholder[T.int64(0), (ax0 * T.int64(3) + ax1) % T.int64(24) // T.int64(12), (ax0 * T.int64(3) + ax1) % T.int64(12) // T.int64(4), (ax0 * T.int64(3) + ax1) % T.int64(4)]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Reshape)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
# fmt: off
|
|
# ShapeExpr might be produced by shape computation
|
|
@tvm.script.ir_module
|
|
class Reshape2:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"):
|
|
lv: R.Shape((8, 3)) = R.shape((8, 3))
|
|
gv: R.Tensor((8, 3), "float32") = R.reshape(x, lv)
|
|
return gv
|
|
|
|
# After lowering, redundant var might be removed by later dead code elimination
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(
|
|
rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"),
|
|
T_reshape: T.Buffer((T.int64(8), T.int64(3)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0, ax1 in T.grid(T.int64(8), T.int64(3)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(
|
|
rxplaceholder[
|
|
T.int64(0),
|
|
(v_ax0 * T.int64(3) + v_ax1) % T.int64(24) // T.int64(12),
|
|
(v_ax0 * T.int64(3) + v_ax1) % T.int64(12) // T.int64(4),
|
|
(v_ax0 * T.int64(3) + v_ax1) % T.int64(4),
|
|
]
|
|
)
|
|
T.writes(T_reshape[v_ax0, v_ax1])
|
|
T_reshape[v_ax0, v_ax1] = rxplaceholder[
|
|
T.int64(0),
|
|
(v_ax0 * T.int64(3) + v_ax1) % T.int64(24) // T.int64(12),
|
|
(v_ax0 * T.int64(3) + v_ax1) % T.int64(12) // T.int64(4),
|
|
(v_ax0 * T.int64(3) + v_ax1) % T.int64(4),
|
|
]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor((8, 3), dtype="float32"):
|
|
lv: R.Shape((8, 3)) = R.shape((8, 3))
|
|
gv = R.call_tir(Expected2.reshape, (x,), out_ty=R.Tensor((8, 3), dtype="float32"))
|
|
return gv
|
|
# fmt: on
|
|
|
|
mod2 = LegalizeOps()(Reshape2)
|
|
tvm.ir.assert_structural_equal(mod2, Expected2)
|
|
|
|
|
|
def test_reshape_symbolic():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Reshape:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
gv: R.Tensor((a // 2, b * 2), "float32") = R.reshape(x, (a // 2, b * 2))
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
gv = R.call_tir(Expected.reshape, (x,), R.Tensor(((a // 2), (b * 2)), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a = T.int64()
|
|
b = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b], dtype="float32")
|
|
T_reshape = T.match_buffer(var_T_reshape, [a // T.int64(2), b * T.int64(2)], dtype="float32")
|
|
for i0, i1 in T.grid(a // T.int64(2), b * T.int64(2)):
|
|
with T.sblock("T_reshape"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(rxplaceholder[(ax0 * b * T.int64(2) + ax1) // b % a, (ax0 * b * T.int64(2) + ax1) % b])
|
|
T.writes(T_reshape[ax0, ax1])
|
|
T_reshape[ax0, ax1] = rxplaceholder[(ax0 * b * T.int64(2) + ax1) // b % a, (ax0 * b * T.int64(2) + ax1) % b]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Reshape)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
# ShapeExpr might be produced by shape computation
|
|
@tvm.script.ir_module
|
|
class Reshape2:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
lv: R.Shape((a // 2, b * 2)) = R.shape((a // 2, b * 2))
|
|
gv: R.Tensor((a // 2, b * 2), "float32") = R.reshape(x, lv)
|
|
return gv
|
|
|
|
# After lowering, redundant var might be removed by later dead code elimination
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
lv: R.Shape((a // 2, b * 2)) = R.shape((a // 2, b * 2))
|
|
gv = R.call_tir(Expected2.reshape, (x,), R.Tensor(((a // 2), (b * 2)), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a = T.int64()
|
|
b = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b], dtype="float32")
|
|
T_reshape = T.match_buffer(
|
|
var_T_reshape, [a // T.int64(2), b * T.int64(2)], dtype="float32"
|
|
)
|
|
for i0, i1 in T.grid(a // T.int64(2), b * T.int64(2)):
|
|
with T.sblock("T_reshape"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(
|
|
rxplaceholder[
|
|
(ax0 * b * T.int64(2) + ax1) // b % a,
|
|
(ax0 * b * T.int64(2) + ax1) % b,
|
|
]
|
|
)
|
|
T.writes(T_reshape[ax0, ax1])
|
|
T_reshape[ax0, ax1] = rxplaceholder[
|
|
(ax0 * b * T.int64(2) + ax1) // b % a, (ax0 * b * T.int64(2) + ax1) % b
|
|
]
|
|
|
|
mod2 = LegalizeOps()(Reshape2)
|
|
tvm.ir.assert_structural_equal(mod2, Expected2)
|
|
|
|
# ShapeExpr might be produced by shape computation
|
|
@I.ir_module(s_tir=True)
|
|
class Reshape3:
|
|
@R.function
|
|
def main(x: R.Tensor((10, "b"), "float32")) -> R.Tensor((5, "b * 2"), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
lv: R.Shape((5, b * 2)) = R.shape((5, b * 2))
|
|
gv: R.Tensor((5, b * 2), "float32") = R.reshape(x, lv)
|
|
return gv
|
|
|
|
# After lowering, redundant var might be removed by later dead code elimination
|
|
@I.ir_module(s_tir=True)
|
|
class Expected3:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
b = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, (T.int64(10), b))
|
|
T_reshape = T.match_buffer(var_T_reshape, (T.int64(5), b * T.int64(2)))
|
|
# with T.sblock("root"):
|
|
for ax0, ax1 in T.grid(T.int64(5), b * T.int64(2)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(
|
|
rxplaceholder[
|
|
(v_ax0 * b * T.int64(2) + v_ax1) // b % T.int64(10),
|
|
(v_ax0 * b * T.int64(2) + v_ax1) % b,
|
|
]
|
|
)
|
|
T.writes(T_reshape[v_ax0, v_ax1])
|
|
T_reshape[v_ax0, v_ax1] = rxplaceholder[
|
|
(v_ax0 * b * T.int64(2) + v_ax1) // b % T.int64(10),
|
|
(v_ax0 * b * T.int64(2) + v_ax1) % b,
|
|
]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, "b"), dtype="float32"),
|
|
) -> R.Tensor((5, "b * 2"), dtype="float32"):
|
|
b = T.int64()
|
|
lv: R.Shape([5, b * 2]) = R.shape([5, b * 2])
|
|
gv = R.call_tir(Expected3.reshape, (x,), out_ty=R.Tensor((5, b * 2), dtype="float32"))
|
|
return gv
|
|
|
|
mod3 = LegalizeOps()(Reshape3)
|
|
tvm.ir.assert_structural_equal(mod3, Expected3)
|
|
|
|
|
|
def test_data_dependent_reshape():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class DDReshape:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([2], dtype="int64"),
|
|
y: R.Tensor([16],dtype='float32'),
|
|
):
|
|
lv: R.Shape(ndim=2) = R.tensor_to_shape(x)
|
|
gv = R.reshape(y, lv)
|
|
return gv
|
|
# fmt: on
|
|
|
|
relax.analysis.well_formed(DDReshape)
|
|
mod = relax.transform.DecomposeOpsForInference()(DDReshape)
|
|
out_mod = relax.transform.LegalizeOps()(mod)
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([2], dtype="int64"),
|
|
y: R.Tensor([16],dtype="float32"),
|
|
) -> R.Tensor(ndim=2, dtype="float32"):
|
|
M = T.int64()
|
|
N = T.int64()
|
|
gv = R.call_pure_packed("vm.builtin.tensor_to_shape", x, ty_args=(R.Shape(ndim=2),))
|
|
_ = R.match_cast(gv, R.Shape([M,N]))
|
|
_ = R.shape([M,N])
|
|
gv_1 = R.call_tir(Expected.reshape, (y,), out_ty=R.Tensor([M,N], dtype="float32"))
|
|
return gv_1
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(
|
|
rxplaceholder: T.Buffer(T.int64(16), "float32"),
|
|
var_T_reshape: T.handle,
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
M = T.int64()
|
|
N = T.int64()
|
|
T_reshape = T.match_buffer(var_T_reshape, [M,N], "float32")
|
|
for i,j in T.grid(M,N):
|
|
with T.sblock("T_reshape"):
|
|
vi,vj = T.axis.remap('SS',[i,j])
|
|
T.reads(rxplaceholder[(vi*N + vj) % 16])
|
|
T.writes(T_reshape[vi,vj])
|
|
T_reshape[vi,vj] = rxplaceholder[(vi*N + vj) % 16]
|
|
|
|
# fmt: on
|
|
tvm.ir.assert_structural_equal(out_mod, Expected)
|
|
|
|
|
|
def test_split_by_indices():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Split:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]):
|
|
gv: R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]) = R.split(x, [3, 7], axis=1)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]):
|
|
gv = R.call_tir(Expected.split, (x,), [R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")])
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def split(rxplaceholder: T.Buffer((T.int64(2), T.int64(10), T.int64(4)), "float32"), T_split: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), T_split_1: T.Buffer((T.int64(2), T.int64(4), T.int64(4)), "float32"), T_split_2: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
|
|
with T.sblock("T_split"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, ax1, ax2])
|
|
T.writes(T_split[ax0, ax1, ax2])
|
|
T_split[ax0, ax1, ax2] = rxplaceholder[ax0, ax1, ax2]
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(4)):
|
|
with T.sblock("T_split_1"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, ax1 + T.int64(3), ax2])
|
|
T.writes(T_split_1[ax0, ax1, ax2])
|
|
T_split_1[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(3), ax2]
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
|
|
with T.sblock("T_split_2"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, ax1 + T.int64(7), ax2])
|
|
T.writes(T_split_2[ax0, ax1, ax2])
|
|
T_split_2[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(7), ax2]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Split)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_split_by_indices_n_section_indivisible():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Split:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")]):
|
|
gv: R.Tuple([R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")]) = R.split(x, indices_or_sections=3, axis=1)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")]):
|
|
gv = R.call_tir(Expected.split, (x,), [R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")])
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def split(rxplaceholder: T.Buffer((T.int64(2), T.int64(10), T.int64(4)), "float32"), T_split_sections: T.Buffer((T.int64(2), T.int64(4), T.int64(4)), "float32"), T_split_sections_1: T.Buffer((T.int64(2), T.int64(4), T.int64(4)), "float32"), T_split_sections_2: T.Buffer((T.int64(2), T.int64(2), T.int64(4)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(4)):
|
|
with T.sblock("T_split_sections"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, ax1, ax2])
|
|
T.writes(T_split_sections[ax0, ax1, ax2])
|
|
T_split_sections[ax0, ax1, ax2] = rxplaceholder[ax0, ax1, ax2]
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(4)):
|
|
with T.sblock("T_split_sections_1"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, ax1 + T.int64(4), ax2])
|
|
T.writes(T_split_sections_1[ax0, ax1, ax2])
|
|
T_split_sections_1[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(4), ax2]
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(2), T.int64(4)):
|
|
with T.sblock("T_split_sections_2"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, ax1 + T.int64(8), ax2])
|
|
T.writes(T_split_sections_2[ax0, ax1, ax2])
|
|
T_split_sections_2[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(8), ax2]
|
|
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Split)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_split_by_indices_n_section_divisible():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Split:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]):
|
|
gv: R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]) = R.split(x, 2, axis=1)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]):
|
|
gv = R.call_tir(Expected.split, (x,), [R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")])
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def split(rxplaceholder: T.Buffer((T.int64(2), T.int64(10), T.int64(4)), "float32"), T_split_sections: T.Buffer((T.int64(2), T.int64(5), T.int64(4)), "float32"), T_split_sections_1: T.Buffer((T.int64(2), T.int64(5), T.int64(4)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(5), T.int64(4)):
|
|
with T.sblock("T_split_sections"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, ax1, ax2])
|
|
T.writes(T_split_sections[ax0, ax1, ax2])
|
|
T_split_sections[ax0, ax1, ax2] = rxplaceholder[ax0, ax1, ax2]
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(5), T.int64(4)):
|
|
with T.sblock("T_split_sections_1"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, ax1 + T.int64(5), ax2])
|
|
T.writes(T_split_sections_1[ax0, ax1, ax2])
|
|
T_split_sections_1[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(5), ax2]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Split)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_split_by_indices_n_section_divisible_symbolic():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Split:
|
|
@R.function
|
|
def main(dumb_param: R.Tensor(("n",)), x: R.Tensor(("m", "n * 3"), "float32")) -> R.Tuple([R.Tensor(("m", "n"), "float32"), R.Tensor(("m", "n"), "float32"), R.Tensor(("m", "n"), "float32")]):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
gv: R.Tuple([R.Tensor((m, n), "float32"), R.Tensor((m, n), "float32"), R.Tensor((m, n), "float32")]) = R.split(x, 3, axis=1)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(dumb_param: R.Tensor(("n",)), x: R.Tensor(("m", "(n * 3)"), "float32")) -> R.Tuple(R.Tensor(("m", "((n * 3) // 3)"), "float32"), R.Tensor(("m", "((((n * 3) // 3) * 2) - ((n * 3) // 3))"), "float32"), R.Tensor(("m", "((n * 3) - (((n * 3) // 3) * 2))"), "float32")):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
gv = R.call_tir(Expected.split, (x,), [R.Tensor((m, ((n * 3 + 3 - 1) // 3)), "float32"), R.Tensor((m, ((((n * 3 + 3 - 1) // 3) * 2) - ((n * 3 + 3 - 1) // 3))), "float32"), R.Tensor((m, ((n * 3) - (((n * 3 + 3 - 1) // 3) * 2))), "float32")], tir_vars=R.shape([n]))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def split(var_rxplaceholder: T.handle, var_T_split_sections: T.handle, var_T_split_sections_1: T.handle, var_T_split_sections_2: T.handle, n: T.int64):
|
|
T.func_attr({"tirx.noalias": True})
|
|
m = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n * T.int64(3)], dtype="float32")
|
|
T_split_sections = T.match_buffer(var_T_split_sections, [m, (n * T.int64(3) + T.int64(3) - T.int64(1)) // T.int64(3)], dtype="float32")
|
|
T_split_sections_1 = T.match_buffer(var_T_split_sections_1, [m, (n * T.int64(3) + T.int64(3) - T.int64(1)) // T.int64(3) * T.int64(2) - (n * T.int64(3) + T.int64(3) - T.int64(1)) // T.int64(3)], dtype="float32")
|
|
T_split_sections_2 = T.match_buffer(var_T_split_sections_2, [m, n * T.int64(3) - (n * T.int64(3) + T.int64(3) - T.int64(1)) // T.int64(3) * T.int64(2)], dtype="float32")
|
|
for i0, i1 in T.grid(m, n):
|
|
with T.sblock("T_split_sections"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(rxplaceholder[ax0, ax1])
|
|
T.writes(T_split_sections[ax0, ax1])
|
|
T_split_sections[ax0, ax1] = rxplaceholder[ax0, ax1]
|
|
for i0, i1 in T.grid(m, n):
|
|
with T.sblock("T_split_sections_1"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(rxplaceholder[ax0, ax1 + n])
|
|
T.writes(T_split_sections_1[ax0, ax1])
|
|
T_split_sections_1[ax0, ax1] = rxplaceholder[ax0, ax1 + n]
|
|
for i0, i1 in T.grid(m, n):
|
|
with T.sblock("T_split_sections_2"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(rxplaceholder[ax0, n * T.int64(2) + ax1])
|
|
T.writes(T_split_sections_2[ax0, ax1])
|
|
T_split_sections_2[ax0, ax1] = rxplaceholder[ax0, n * T.int64(2) + ax1]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Split)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_squeeze():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Squeeze:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"):
|
|
gv: R.Tensor((2, 3, 1, 4), "float32") = R.squeeze(x, [1, 4])
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"):
|
|
gv = R.call_tir(Expected.squeeze, (x,), R.Tensor((2, 3, 1, 4), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def squeeze(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3), T.int64(1), T.int64(1), T.int64(4)), "float32"), T_squeeze: T.Buffer((T.int64(2), T.int64(3), T.int64(1), T.int64(4)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(1), T.int64(4)):
|
|
with T.sblock("T_squeeze"):
|
|
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
T.reads(rxplaceholder[ax0, T.int64(0), ax1, ax2, T.int64(0), ax3])
|
|
T.writes(T_squeeze[ax0, ax1, ax2, ax3])
|
|
T_squeeze[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, T.int64(0), ax1, ax2, T.int64(0), ax3]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Squeeze)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_squeeze_no_axis():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Squeeze:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) :
|
|
gv: R.Tensor((2, 3, 4), "float32") = R.squeeze(x)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) :
|
|
gv = R.call_tir(Expected.squeeze, (x,), R.Tensor((2, 3, 4), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def squeeze(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3), T.int64(1), T.int64(1), T.int64(4)), "float32"), T_squeeze: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
|
|
with T.sblock("T_squeeze"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, T.int64(0), ax1, T.int64(0), T.int64(0), ax2])
|
|
T.writes(T_squeeze[ax0, ax1, ax2])
|
|
T_squeeze[ax0, ax1, ax2] = rxplaceholder[ax0, T.int64(0), ax1, T.int64(0), T.int64(0), ax2]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Squeeze)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_squeeze_symbolic():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Squeeze:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", 1, "b", 1), "float32")) -> R.Tensor(("a", "b", 1), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
gv: R.Tensor((a, b, 1), "float32") = R.squeeze(x, [1])
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", 1, "b", 1), "float32")) -> R.Tensor(("a", "b", 1), "float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
gv = R.call_tir(Expected.squeeze, (x,), R.Tensor((a, b, 1), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def squeeze(var_rxplaceholder: T.handle, var_T_squeeze: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a = T.int64()
|
|
b = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, [a, T.int64(1), b, T.int64(1)], dtype="float32")
|
|
T_squeeze = T.match_buffer(var_T_squeeze, [a, b, T.int64(1)], dtype="float32")
|
|
for i0, i1, i2 in T.grid(a, b, T.int64(1)):
|
|
with T.sblock("T_squeeze"):
|
|
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(rxplaceholder[ax0, T.int64(0), ax1, ax2])
|
|
T.writes(T_squeeze[ax0, ax1, ax2])
|
|
T_squeeze[ax0, ax1, ax2] = rxplaceholder[ax0, T.int64(0), ax1, ax2]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Squeeze)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_collapse_sum_like():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class CollapseSumLike:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((1, 3), "float32")) -> R.Tensor((1, 3), "float32"):
|
|
gv: R.Tensor((1, 3), "float32") = R.collapse_sum_like(x, y)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((1, 3), "float32")) -> R.Tensor((1, 3), "float32"):
|
|
gv = R.call_tir(Expected.collapse_sum, (x,), R.Tensor((1, 3), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def collapse_sum(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), rxplaceholder_red: T.Buffer((T.int64(1), T.int64(3)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1, i2 in T.grid(T.int64(1), T.int64(3), T.int64(2)):
|
|
with T.sblock("rxplaceholder_red"):
|
|
ax0, ax1, k0 = T.axis.remap("SSR", [i0, i1, i2])
|
|
T.reads(rxplaceholder[k0, ax1])
|
|
T.writes(rxplaceholder_red[ax0, ax1])
|
|
with T.init():
|
|
rxplaceholder_red[ax0, ax1] = T.float32(0)
|
|
rxplaceholder_red[ax0, ax1] = rxplaceholder_red[ax0, ax1] + rxplaceholder[k0, ax1]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(CollapseSumLike)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_collapse_sum_to():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class CollapseSumTo:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 2, 3), "float32")) -> R.Tensor((2, 1), "float32"):
|
|
gv: R.Tensor((2, 1), "float32") = R.collapse_sum_to(x, (2, 1))
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 2, 3), dtype="float32")
|
|
) -> R.Tensor((2, 1), dtype="float32"):
|
|
# block 0
|
|
gv = R.call_tir(Expected.collapse_sum, (x,), R.Tensor((2, 1), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def collapse_sum(rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32"), rxplaceholder_red: T.Buffer((T.int64(2), T.int64(1)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1, k0, k2 in T.grid(T.int64(2), T.int64(1), T.int64(3), T.int64(3)):
|
|
with T.sblock("rxplaceholder_red"):
|
|
v_ax0, v_ax1, v_k0, v_k2 = T.axis.remap("SSRR", [ax0, ax1, k0, k2])
|
|
T.reads(rxplaceholder[v_k0, v_ax0, v_k2])
|
|
T.writes(rxplaceholder_red[v_ax0, v_ax1])
|
|
with T.init():
|
|
rxplaceholder_red[v_ax0, v_ax1] = T.float32(0)
|
|
rxplaceholder_red[v_ax0, v_ax1] = (rxplaceholder_red[v_ax0, v_ax1] + rxplaceholder[v_k0, v_ax0, v_k2])
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(CollapseSumTo)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_repeat():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Repeat:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 2, 3), "float32")):
|
|
gv = R.repeat(x, 2, 0)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 2, 3), dtype="float32")) -> R.Tensor((6, 2, 3), dtype="float32"):
|
|
gv = R.call_tir(Expected.repeat, (x,), out_ty=R.Tensor((6, 2, 3), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def repeat(rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32"), T_repeat: T.Buffer((T.int64(6), T.int64(2), T.int64(3)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0, ax1, ax2 in T.grid(T.int64(6), T.int64(2), T.int64(3)):
|
|
with T.sblock("T_repeat"):
|
|
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
|
T.reads(rxplaceholder[v_ax0 // T.int64(2), v_ax1, v_ax2])
|
|
T.writes(T_repeat[v_ax0, v_ax1, v_ax2])
|
|
T_repeat[v_ax0, v_ax1, v_ax2] = rxplaceholder[v_ax0 // T.int64(2), v_ax1, v_ax2]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Repeat)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_repeat_no_axis():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Repeat:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 2, 3), "float32")):
|
|
gv = R.repeat(x, 2)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 2, 3), dtype="float32")
|
|
) -> R.Tensor((36,), dtype="float32"):
|
|
gv = R.call_tir(Expected.repeat, (x,), out_ty=R.Tensor((36,), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def repeat(
|
|
rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32"),
|
|
T_repeat: T.Buffer((T.int64(36),), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
T_reshape = T.sblock_alloc_buffer((T.int64(18),))
|
|
for ax0 in range(T.int64(18)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0 = T.axis.spatial(T.int64(18), ax0)
|
|
T.reads(
|
|
rxplaceholder[
|
|
v_ax0 % T.int64(18) // T.int64(6),
|
|
v_ax0 % T.int64(6) // T.int64(3),
|
|
v_ax0 % T.int64(3),
|
|
]
|
|
)
|
|
T.writes(T_reshape[v_ax0])
|
|
T_reshape[v_ax0] = rxplaceholder[
|
|
v_ax0 % T.int64(18) // T.int64(6),
|
|
v_ax0 % T.int64(6) // T.int64(3),
|
|
v_ax0 % T.int64(3),
|
|
]
|
|
for ax0 in range(T.int64(36)):
|
|
with T.sblock("T_repeat"):
|
|
v_ax0 = T.axis.spatial(T.int64(36), ax0)
|
|
T.reads(T_reshape[v_ax0 // T.int64(2)])
|
|
T.writes(T_repeat[v_ax0])
|
|
T_repeat[v_ax0] = T_reshape[v_ax0 // T.int64(2)]
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Repeat)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_repeat_symbolic():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Repeat:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), "float32")):
|
|
gv = R.repeat(x, 2, 0)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def repeat(var_rxplaceholder: T.handle, var_T_repeat: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b, c))
|
|
T_repeat = T.match_buffer(var_T_repeat, (T.int64(2) * a, b, c))
|
|
# with T.sblock("root"):
|
|
for ax0, ax1, ax2 in T.grid(a * T.int64(2), b, c):
|
|
with T.sblock("T_repeat"):
|
|
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
|
T.reads(rxplaceholder[v_ax0 // T.int64(2), v_ax1, v_ax2])
|
|
T.writes(T_repeat[v_ax0, v_ax1, v_ax2])
|
|
T_repeat[v_ax0, v_ax1, v_ax2] = rxplaceholder[v_ax0 // T.int64(2), v_ax1, v_ax2]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), dtype="float32")) -> R.Tensor(("2 * a", "b", "c"), dtype="float32"):
|
|
a = T.Var("a", "int64")
|
|
b = T.Var("b", "int64")
|
|
c = T.Var("c", "int64")
|
|
gv = R.call_tir(Expected.repeat, (x,), out_ty=R.Tensor((2 * a, b, c), dtype="float32"))
|
|
return gv
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Repeat)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_tile():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Tile:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 2, 3), "float32")):
|
|
gv = R.tile(x, (2, 1, 2, 3))
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def tile(rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32"), T_tile: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(9)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(9)):
|
|
with T.sblock("T_tile"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(rxplaceholder[v_ax1 % T.int64(3), v_ax2 % T.int64(2), v_ax3 % T.int64(3)])
|
|
T.writes(T_tile[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_tile[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax1 % T.int64(3), v_ax2 % T.int64(2), v_ax3 % T.int64(3)]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((3, 2, 3), dtype="float32")) -> R.Tensor((2, 3, 4, 9), dtype="float32"):
|
|
gv = R.call_tir(Expected.tile, (x,), out_ty=R.Tensor((2, 3, 4, 9), dtype="float32"))
|
|
return gv
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Tile)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_tile_symbolic():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Tile:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), "float32")):
|
|
gv = R.tile(x, (2, 1, 2, 3))
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def tile(var_rxplaceholder: T.handle, var_T_tile: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a = T.int64()
|
|
b = T.int64()
|
|
c = T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b, c))
|
|
T_tile = T.match_buffer(var_T_tile, (T.int64(2), a, b * T.int64(2), c * T.int64(3)))
|
|
# with T.sblock("root"):
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), a, b * T.int64(2), c * T.int64(3)):
|
|
with T.sblock("T_tile"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(rxplaceholder[v_ax1 % a, v_ax2 % b, v_ax3 % c])
|
|
T.writes(T_tile[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_tile[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax1 % a, v_ax2 % b, v_ax3 % c]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), dtype="float32")) -> R.Tensor((2, "a", "b * 2", "c * 3"), dtype="float32"):
|
|
a = T.Var("a", "int64")
|
|
b = T.Var("b", "int64")
|
|
c = T.Var("c", "int64")
|
|
gv = R.call_tir(Expected.tile, (x,), out_ty=R.Tensor((2, a, b * 2, c * 3), dtype="float32"))
|
|
return gv
|
|
# fmt: on
|
|
mod = LegalizeOps()(Tile)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_flip():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Flip:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), "float32")):
|
|
gv = R.flip(x, axis=0)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
|
|
cls = Expected
|
|
gv = R.call_tir(cls.flip, (x,), out_ty=R.Tensor((2, 3), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def flip(
|
|
rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"),
|
|
T_reverse_sequence: T.Buffer((T.int64(2), T.int64(3)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(2), T.int64(3)):
|
|
with T.sblock("T_reverse_sequence"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(rxplaceholder[T.int64(1) - v_ax0, v_ax1])
|
|
T.writes(T_reverse_sequence[v_ax0, v_ax1])
|
|
T_reverse_sequence[v_ax0, v_ax1] = rxplaceholder[
|
|
T.int64(1) - v_ax0, v_ax1
|
|
]
|
|
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Flip)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_flip_symbolic():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Flip:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b"), "float32")):
|
|
gv = R.flip(x, axis=1)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor(("a", "b"), dtype="float32")
|
|
) -> R.Tensor(("a", "b"), dtype="float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
cls = Expected
|
|
gv = R.call_tir(cls.flip, (x,), out_ty=R.Tensor((a, b), dtype="float32"))
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def flip(var_rxplaceholder: T.handle, var_T_reverse_sequence: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a, b = T.int64(), T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b))
|
|
T_reverse_sequence = T.match_buffer(var_T_reverse_sequence, (a, b))
|
|
for ax0, ax1 in T.grid(a, b):
|
|
with T.sblock("T_reverse_sequence"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(rxplaceholder[v_ax0, b - v_ax1 - T.int64(1)])
|
|
T.writes(T_reverse_sequence[v_ax0, v_ax1])
|
|
T_reverse_sequence[v_ax0, v_ax1] = rxplaceholder[
|
|
v_ax0, b - v_ax1 - T.int64(1)
|
|
]
|
|
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(Flip)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_reverse_sequence():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class ReverseSequence:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 2, 3), "float32"), seq_lengths: R.Tensor((2,), "int64")):
|
|
gv = R.reverse_sequence(x, seq_lengths, seq_axis=0, batch_axis=1)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((4, 2, 3), dtype="float32"),
|
|
seq_lengths: R.Tensor((2,), dtype="int64"),
|
|
) -> R.Tensor((4, 2, 3), dtype="float32"):
|
|
cls = Expected
|
|
gv = R.call_tir(
|
|
cls.reverse_sequence,
|
|
(x, seq_lengths),
|
|
out_ty=R.Tensor((4, 2, 3), dtype="float32"),
|
|
)
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reverse_sequence(
|
|
rxplaceholder: T.Buffer((T.int64(4), T.int64(2), T.int64(3)), "float32"),
|
|
seq_lengths: T.Buffer((T.int64(2),), "int64"),
|
|
T_reverse_sequence: T.Buffer((T.int64(4), T.int64(2), T.int64(3)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1, ax2 in T.grid(T.int64(4), T.int64(2), T.int64(3)):
|
|
with T.sblock("T_reverse_sequence"):
|
|
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
|
T.reads(rxplaceholder[T.int64(0):T.int64(4), v_ax1, v_ax2], seq_lengths[v_ax1])
|
|
T.writes(T_reverse_sequence[v_ax0, v_ax1, v_ax2])
|
|
T_reverse_sequence[v_ax0, v_ax1, v_ax2] = rxplaceholder[
|
|
T.if_then_else(
|
|
seq_lengths[v_ax1] <= T.int64(1) or seq_lengths[v_ax1] <= v_ax0,
|
|
v_ax0,
|
|
T.if_then_else(
|
|
T.int64(4) < seq_lengths[v_ax1],
|
|
T.int64(3) - v_ax0,
|
|
seq_lengths[v_ax1] - v_ax0 - T.int64(1),
|
|
),
|
|
),
|
|
v_ax1,
|
|
v_ax2,
|
|
]
|
|
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(ReverseSequence)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_scatter_elements():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class ScatterElements:
|
|
@R.function
|
|
def main(x: R.Tensor((4,4), "float32"), indices: R.Tensor((2,2), "int64"), updates: R.Tensor((2,2), "float32")):
|
|
gv = R.scatter_elements(x, indices, updates, axis=1)
|
|
return gv
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def scatter_elements(
|
|
var_rxplaceholder: T.handle,
|
|
var_rxplaceholder_1: T.handle,
|
|
var_rxplaceholder_2: T.handle,
|
|
out_buf: T.Buffer((T.int64(4), T.int64(4)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
rxplaceholder = T.match_buffer(
|
|
var_rxplaceholder, (T.int64(4), T.int64(4)), offset_factor=1
|
|
)
|
|
rxplaceholder_1 = T.match_buffer(
|
|
var_rxplaceholder_1, (T.int64(2), T.int64(2)), "int64", offset_factor=1
|
|
)
|
|
rxplaceholder_2 = T.match_buffer(
|
|
var_rxplaceholder_2, (T.int64(2), T.int64(2)), offset_factor=1
|
|
)
|
|
with T.sblock("scatter_elements_generic"):
|
|
T.attr(0, "pragma_scope", "seq")
|
|
for i in T.parallel(T.int64(16)):
|
|
out_buf[i // T.int64(4), i % T.int64(4)] = rxplaceholder[
|
|
i // T.int64(4), i % T.int64(4)
|
|
]
|
|
for fused in T.parallel(T.int64(2)):
|
|
for k in range(T.int64(2)):
|
|
out_buf[
|
|
(
|
|
fused * T.int64(4)
|
|
+ (
|
|
rxplaceholder_1[
|
|
(fused * T.int64(2) + k) // T.int64(2),
|
|
(fused * T.int64(2) + k) % T.int64(2),
|
|
]
|
|
+ T.Cast(
|
|
"int64",
|
|
rxplaceholder_1[
|
|
(fused * T.int64(2) + k) // T.int64(2),
|
|
(fused * T.int64(2) + k) % T.int64(2),
|
|
]
|
|
< T.int64(0),
|
|
)
|
|
* T.int64(4)
|
|
)
|
|
)
|
|
// T.int64(4),
|
|
(
|
|
fused * T.int64(4)
|
|
+ (
|
|
rxplaceholder_1[
|
|
(fused * T.int64(2) + k) // T.int64(2),
|
|
(fused * T.int64(2) + k) % T.int64(2),
|
|
]
|
|
+ T.Cast(
|
|
"int64",
|
|
rxplaceholder_1[
|
|
(fused * T.int64(2) + k) // T.int64(2),
|
|
(fused * T.int64(2) + k) % T.int64(2),
|
|
]
|
|
< T.int64(0),
|
|
)
|
|
* T.int64(4)
|
|
)
|
|
)
|
|
% T.int64(4),
|
|
] = rxplaceholder_2[
|
|
(fused * T.int64(2) + k) // T.int64(2),
|
|
(fused * T.int64(2) + k) % T.int64(2),
|
|
]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((4, 4), dtype="float32"),
|
|
indices: R.Tensor((2, 2), dtype="int64"),
|
|
updates: R.Tensor((2, 2), dtype="float32"),
|
|
) -> R.Tensor((4, 4), dtype="float32"):
|
|
gv = R.call_tir(
|
|
Expected.scatter_elements,
|
|
(x, indices, updates),
|
|
out_ty=R.Tensor((4, 4), dtype="float32"),
|
|
)
|
|
return gv
|
|
|
|
# fmt: on
|
|
mod = LegalizeOps()(ScatterElements)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_scatter_elements_symbolic():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class ScatterElements:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b"), "float32"), indices:R.Tensor(("m", "n"), "int64"), updates:R.Tensor(("m","n"), "float32")):
|
|
gv = R.scatter_elements(x, indices, updates, axis=1)
|
|
return gv
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def scatter_elements(
|
|
var_rxplaceholder: T.handle,
|
|
var_rxplaceholder_1: T.handle,
|
|
var_rxplaceholder_2: T.handle,
|
|
var_scatter_elements_generic: T.handle,
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a, b = T.int64(), T.int64()
|
|
rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b), offset_factor=1)
|
|
m, n = T.int64(), T.int64()
|
|
rxplaceholder_1 = T.match_buffer(
|
|
var_rxplaceholder_1, (m, n), "int64", offset_factor=1
|
|
)
|
|
rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, (m, n), offset_factor=1)
|
|
out_buf = T.match_buffer(var_scatter_elements_generic, (a, b))
|
|
with T.sblock("scatter_elements_generic"):
|
|
T.attr(0, "pragma_scope", "seq")
|
|
for i in T.parallel(a * b):
|
|
out_buf[i // b, i % b] = rxplaceholder[i // b, i % b]
|
|
for fused in T.parallel(m):
|
|
for k in range(n):
|
|
out_buf[
|
|
(
|
|
fused * b
|
|
+ (
|
|
rxplaceholder_1[
|
|
(fused * n + k) // n, (fused * n + k) % n
|
|
]
|
|
+ T.Cast(
|
|
"int64",
|
|
rxplaceholder_1[
|
|
(fused * n + k) // n, (fused * n + k) % n
|
|
]
|
|
< T.int64(0),
|
|
)
|
|
* b
|
|
)
|
|
)
|
|
// b,
|
|
(
|
|
fused * b
|
|
+ (
|
|
rxplaceholder_1[
|
|
(fused * n + k) // n, (fused * n + k) % n
|
|
]
|
|
+ T.Cast(
|
|
"int64",
|
|
rxplaceholder_1[
|
|
(fused * n + k) // n, (fused * n + k) % n
|
|
]
|
|
< T.int64(0),
|
|
)
|
|
* b
|
|
)
|
|
)
|
|
% b,
|
|
] = rxplaceholder_2[(fused * n + k) // n, (fused * n + k) % n]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor(("a", "b"), dtype="float32"),
|
|
indices: R.Tensor(("m", "n"), dtype="int64"),
|
|
updates: R.Tensor(("m", "n"), dtype="float32"),
|
|
) -> R.Tensor(("a", "b"), dtype="float32"):
|
|
a = T.int64()
|
|
b = T.int64()
|
|
m = T.int64()
|
|
n = T.int64()
|
|
gv = R.call_tir(
|
|
Expected.scatter_elements,
|
|
(x, indices, updates),
|
|
out_ty=R.Tensor((a, b), dtype="float32"),
|
|
)
|
|
return gv
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(ScatterElements)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not enabled")
|
|
def test_scatter_elements_gpu():
|
|
"""scatter_elements lowered for GPU must build"""
|
|
target = "cuda"
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Mod:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((4, 8), "float32"),
|
|
indices: R.Tensor((2, 8), "int64"),
|
|
updates: R.Tensor((2, 8), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv = R.scatter_elements(x, indices, updates, axis=0)
|
|
gv = lv
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with tvm.target.Target(target):
|
|
mod = LegalizeOps()(Mod)
|
|
relax.build(mod, target=target)
|
|
|
|
|
|
def test_layout_transform():
|
|
transformation = lambda a, b, c: (a, c, b // 3, b % 3)
|
|
pad_value = 2
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class LayoutTransform:
|
|
@R.function
|
|
def main(x: R.Tensor((10, 21, 30), "float32")):
|
|
gv = R.layout_transform(
|
|
x, index_map=transformation, pad_value=pad_value
|
|
)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def te_layout_transform(A: T.Buffer((T.int64(10), T.int64(21), T.int64(30)), "float32"), te_layout_transform_1: T.Buffer((T.int64(10), T.int64(30), T.int64(7), T.int64(3)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for i0, i1, i2 in T.grid(T.int64(10), T.int64(21), T.int64(30)):
|
|
with T.sblock("te_layout_transform"):
|
|
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
|
|
T.reads(A[v_i0, v_i1, v_i2])
|
|
T.writes(te_layout_transform_1[v_i0, v_i2, v_i1 // T.int64(3), v_i1 % T.int64(3)])
|
|
te_layout_transform_1[v_i0, v_i2, v_i1 // T.int64(3), v_i1 % T.int64(3)] = A[v_i0, v_i1, v_i2]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((10, 21, 30), dtype="float32")) -> R.Tensor((10, 30, 7, 3), dtype="float32"):
|
|
cls = Expected
|
|
gv = R.call_tir(cls.te_layout_transform, (x,), out_ty=R.Tensor((10, 30, 7, 3), dtype="float32"))
|
|
return gv
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(LayoutTransform)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_layout_transform_with_pad():
|
|
transformation = lambda a, b, c: (a, c, b // 3, b % 3)
|
|
pad_value = 2
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class LayoutTransform:
|
|
@R.function
|
|
def main(x: R.Tensor((10, 20, 30), "float32")):
|
|
gv = R.layout_transform(
|
|
x, index_map=transformation, pad_value=pad_value
|
|
)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def te_layout_transform_with_pad(A: T.Buffer((T.int64(10), T.int64(20), T.int64(30)), "float32"), te_layout_transform_with_pad_1: T.Buffer((T.int64(10), T.int64(30), T.int64(7), T.int64(3)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for axis0, axis1, axis2, axis3 in T.grid(T.int64(10), T.int64(30), T.int64(7), T.int64(3)):
|
|
with T.sblock("te_layout_transform_with_pad"):
|
|
v_axis0, v_axis1, v_axis2, v_axis3 = T.axis.remap("SSSS", [axis0, axis1, axis2, axis3])
|
|
T.reads(A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])
|
|
T.writes(te_layout_transform_with_pad_1[v_axis0, v_axis1, v_axis2, v_axis3])
|
|
te_layout_transform_with_pad_1[v_axis0, v_axis1, v_axis2, v_axis3] = T.if_then_else(v_axis2 == T.int64(6) and v_axis3 == T.int64(2), T.float32(2), A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((10, 20, 30), dtype="float32")) -> R.Tensor((10, 30, 7, 3), dtype="float32"):
|
|
cls = Expected
|
|
gv = R.call_tir(cls.te_layout_transform_with_pad, (x,), out_ty=R.Tensor((10, 30, 7, 3), dtype="float32"))
|
|
return gv
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(LayoutTransform)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_layout_transform_symbolic():
|
|
transformation = lambda a, b, c: (a, c, b // 3, b % 3)
|
|
pad_value = 2
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class LayoutTransform:
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), "float32")):
|
|
gv = R.layout_transform(
|
|
x, index_map=transformation, pad_value=pad_value
|
|
)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def te_layout_transform_with_pad(var_A: T.handle, var_te_layout_transform_with_pad: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
a, b, c = T.int64(), T.int64(), T.int64()
|
|
A = T.match_buffer(var_A, (a, b, c))
|
|
te_layout_transform_with_pad_1 = T.match_buffer(var_te_layout_transform_with_pad, (a, c, (b - b % T.int64(-3)) // T.int64(3), T.int64(3)))
|
|
# with T.sblock("root"):
|
|
for axis0, axis1, axis2, axis3 in T.grid(a, c, (b - b % T.int64(-3)) // T.int64(3), T.int64(3)):
|
|
with T.sblock("te_layout_transform_with_pad_with_pad"):
|
|
v_axis0, v_axis1, v_axis2, v_axis3 = T.axis.remap("SSSS", [axis0, axis1, axis2, axis3])
|
|
T.reads(A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])
|
|
T.writes(te_layout_transform_with_pad_1[v_axis0, v_axis1, v_axis2, v_axis3])
|
|
te_layout_transform_with_pad_1[v_axis0, v_axis1, v_axis2, v_axis3] = T.if_then_else(b % T.int64(-3) < T.int64(0) and v_axis2 == b // T.int64(3) and b % T.int64(3) <= v_axis3, T.float32(2), A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])
|
|
|
|
@R.function
|
|
def main(x: R.Tensor(("a", "b", "c"), dtype="float32")) -> R.Tensor(("a", "c", "(b - b % -3) // 3", 3), dtype="float32"):
|
|
a = T.int64()
|
|
c = T.int64()
|
|
b = T.int64()
|
|
cls = Expected
|
|
gv = R.call_tir(cls.te_layout_transform_with_pad, (x,), out_ty=R.Tensor((a, c, (b - b % -3) // 3, 3), dtype="float32"))
|
|
return gv
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(LayoutTransform)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_layout_transform_with_pad_axis_sep():
|
|
transformation = lambda a, b, c: (a, c, b // 3, b % 3)
|
|
pad_value = 2
|
|
axis_separator = [3]
|
|
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class LayoutTransform:
|
|
@R.function
|
|
def main(x: R.Tensor((10, 20, 30), "float32")):
|
|
gv = R.layout_transform(
|
|
x, index_map=transformation, pad_value=pad_value, axis_separators=axis_separator,
|
|
)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def te_layout_transform_with_pad_axis_separator(A: T.Buffer((T.int64(10), T.int64(20), T.int64(30)), "float32"), var_te_layout_transform_with_pad_axis_separator: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
te_layout_transform_with_pad_axis_separator_1 = T.match_buffer(var_te_layout_transform_with_pad_axis_separator, (T.int64(10), T.int64(30), T.int64(7), T.int64(3)), axis_separators=[3])
|
|
# with T.sblock("root"):
|
|
for axis0, axis1, axis2, axis3 in T.grid(T.int64(10), T.int64(30), T.int64(7), T.int64(3)):
|
|
with T.sblock("te_layout_transform_with_pad_axis_separator"):
|
|
v_axis0, v_axis1, v_axis2, v_axis3 = T.axis.remap("SSSS", [axis0, axis1, axis2, axis3])
|
|
T.reads(A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])
|
|
T.writes(te_layout_transform_with_pad_axis_separator_1[v_axis0, v_axis1, v_axis2, v_axis3])
|
|
te_layout_transform_with_pad_axis_separator_1[v_axis0, v_axis1, v_axis2, v_axis3] = T.if_then_else(v_axis2 == T.int64(6) and v_axis3 == T.int64(2), T.float32(2), A[v_axis0, v_axis2 * T.int64(3) + v_axis3, v_axis1])
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((10, 20, 30), dtype="float32")) -> R.Tensor((10, 30, 7, 3), dtype="float32"):
|
|
cls = Expected
|
|
gv = R.call_tir(cls.te_layout_transform_with_pad_axis_separator, (x,), out_ty=R.Tensor((10, 30, 7, 3), dtype="float32"))
|
|
return gv
|
|
# fmt: on
|
|
|
|
mod = LegalizeOps()(LayoutTransform)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_func_ty_of_legalized_layout_transform():
|
|
"""PrimFunc shape information must be correct
|
|
|
|
This is a regression test. Previously, the legalization of
|
|
`R.layout_transform` produced a PrimFunc with `FuncType`
|
|
different than its actual signature. This resulted in errors
|
|
when later passes attempted to infer the Type.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((16,), dtype="float32"), y: R.Tensor((16,), dtype="float32")
|
|
) -> R.Tensor((16,), dtype="float32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 4), dtype="float32") = R.layout_transform(
|
|
x, index_map=lambda i: (i // 4, i % 4), pad_value=None
|
|
)
|
|
gv: R.Tensor((4, 4), dtype="float32") = lv
|
|
R.output(gv)
|
|
return gv
|
|
|
|
After = tvm.ir.transform.Sequential(
|
|
[
|
|
relax.transform.LegalizeOps(),
|
|
relax.transform.ToNonDataflow(),
|
|
relax.transform.RemovePurityChecking(),
|
|
relax.transform.CallTIRRewrite(),
|
|
]
|
|
)(Before)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((16,), dtype="float32"),
|
|
y: R.Tensor((16,), dtype="float32"),
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
cls = Expected
|
|
alloc: R.Tensor((4, 4), dtype="float32") = R.emit_with_ty(
|
|
"relax.builtin.alloc_tensor",
|
|
(R.shape([4, 4]), R.dtype("float32"), R.prim_value(0), R.str("global")),
|
|
(R.Tensor((4, 4), dtype="float32"),),
|
|
)
|
|
cls.te_layout_transform(x, alloc)
|
|
lv = alloc
|
|
gv = lv
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def te_layout_transform(
|
|
A: T.Buffer((T.int64(16),), "float32"),
|
|
te_layout_transform: T.Buffer((T.int64(4), T.int64(4)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in range(T.int64(16)):
|
|
with T.sblock("te_layout_transform"):
|
|
vi = T.axis.spatial(T.int64(16), i)
|
|
te_layout_transform[vi // T.int64(4), vi % T.int64(4)] = A[vi]
|
|
|
|
tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
def test_scatter_nd():
|
|
# fmt: off
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((8,), "float32"),
|
|
indices: R.Tensor((4, 1), "int64"),
|
|
updates: R.Tensor((4,), "float32"),
|
|
) -> R.Tensor((8,), "float32"):
|
|
gv: R.Tensor((8,), "float32") = R.scatter_nd(data, indices, updates, reduction="update")
|
|
return gv
|
|
|
|
After = relax.transform.LegalizeOps()(Before)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((8,), "float32"),
|
|
indices: R.Tensor((4, 1), "int64"),
|
|
updates: R.Tensor((4,), "float32"),
|
|
) -> R.Tensor((8,), "float32"):
|
|
gv = R.call_tir(
|
|
Expected.scatter_nd, (data, indices, updates), R.Tensor((8,), dtype="float32")
|
|
)
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def scatter_nd(var_data: T.handle, var_indices: T.handle, var_updates: T.handle, var_scatter_nd_generic: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
data = T.match_buffer(var_data, (T.int64(8),), offset_factor=1)
|
|
indices = T.match_buffer(var_indices, (T.int64(4), T.int64(1)), "int64")
|
|
updates = T.match_buffer(var_updates, (T.int64(4),), offset_factor=1)
|
|
out_buf = T.match_buffer(var_scatter_nd_generic, (T.int64(8),))
|
|
with T.sblock("root"):
|
|
T.reads()
|
|
T.writes()
|
|
T_transpose = T.sblock_alloc_buffer((T.int64(1), T.int64(4)), "int64")
|
|
for ax0 in range(T.int64(1)):
|
|
for ax1 in range(T.int64(4)):
|
|
with T.sblock("T_transpose"):
|
|
v_ax0 = T.axis.spatial(T.int64(1), ax0)
|
|
v_ax1 = T.axis.spatial(T.int64(4), ax1)
|
|
T.reads(indices[v_ax1, v_ax0])
|
|
T.writes(T_transpose[v_ax0, v_ax1])
|
|
T_transpose[v_ax0, v_ax1] = indices[v_ax1, v_ax0]
|
|
with T.sblock("scatter_nd_generic"):
|
|
T.reads()
|
|
T.writes()
|
|
T.attr(0, "pragma_scope", "seq")
|
|
for i in range(T.int64(8)):
|
|
out_buf[i] = data[i]
|
|
for j in range(T.int64(4)):
|
|
for k in T.parallel(T.int64(1)):
|
|
out_buf[k + T_transpose[j // T.int64(4), j % T.int64(4)]] = updates[j + k]
|
|
|
|
# fmt: on
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not enabled")
|
|
def test_scatter_nd_gpu():
|
|
"""scatter_nd lowered for GPU must build"""
|
|
target = "cuda"
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Mod:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((4, 8), "float32"),
|
|
indices: R.Tensor((3, 2), "int64"),
|
|
updates: R.Tensor((3,), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv = R.scatter_nd(data, indices, updates)
|
|
gv = lv
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with tvm.target.Target(target):
|
|
mod = LegalizeOps()(Mod)
|
|
relax.build(mod, target=target)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|