wehub-resource-sync
769 lines
35 KiB
Python
769 lines
35 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, F841
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def test_reshape_expand_dims():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def reshape(
|
|
rxplaceholder: T.Buffer((T.int64(8), T.int64(3)), "float32"),
|
|
T_reshape: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
|
|
):
|
|
for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(3)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
|
T.reads(
|
|
rxplaceholder[
|
|
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
|
|
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
|
|
]
|
|
)
|
|
T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
|
|
T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
|
|
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
|
|
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
|
|
]
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def expand_dims(
|
|
rxplaceholder: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
|
|
expand_dims: T.Buffer(
|
|
(T.int64(2), T.int64(1), T.int64(4), T.int64(1), T.int64(3)), "float32"
|
|
),
|
|
):
|
|
for i0, i1, i2, i3, i4 in T.grid(
|
|
T.int64(2), T.int64(1), T.int64(4), T.int64(1), T.int64(3)
|
|
):
|
|
with T.sblock("expand_dims"):
|
|
i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
|
T.reads(rxplaceholder[i0_1, i2_1, i4_1])
|
|
T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1])
|
|
expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1] = rxplaceholder[i0_1, i2_1, i4_1]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((8, 3), dtype="float32")) -> R.Tensor(
|
|
(2, 1, 4, 1, 3), dtype="float32"
|
|
):
|
|
cls = Module
|
|
with R.dataflow():
|
|
y = R.call_tir(cls.reshape, (x,), out_ty=R.Tensor((2, 4, 3), dtype="float32"))
|
|
z = R.call_tir(cls.expand_dims, (y,), out_ty=R.Tensor((2, 1, 4, 1, 3), "float32"))
|
|
R.output(z)
|
|
return z
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def reshape(
|
|
rxplaceholder: T.Buffer((T.int64(8), T.int64(3)), "float32"),
|
|
T_reshape: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
|
|
):
|
|
for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(3)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
|
T.reads(
|
|
rxplaceholder[
|
|
(v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) // T.int64(3),
|
|
(v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) % T.int64(3),
|
|
]
|
|
)
|
|
T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
|
|
T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
|
|
(v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) // T.int64(3),
|
|
(v_ax0 * T.int64(12) + v_ax1 * T.int64(3) + v_ax2) % T.int64(3),
|
|
]
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def expand_dims(
|
|
rxplaceholder: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
|
|
expand_dims: T.Buffer(
|
|
(T.int64(2), T.int64(1), T.int64(4), T.int64(1), T.int64(3)), "float32"
|
|
),
|
|
):
|
|
for i0, i1, i2, i3, i4 in T.grid(
|
|
T.int64(2), T.int64(1), T.int64(4), T.int64(1), T.int64(3)
|
|
):
|
|
with T.sblock("expand_dims"):
|
|
i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
|
T.reads(rxplaceholder[i0_1, i2_1, i4_1])
|
|
T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1])
|
|
expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1] = rxplaceholder[i0_1, i2_1, i4_1]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((8, 3), dtype="float32")) -> R.Tensor(
|
|
(2, 1, 4, 1, 3), dtype="float32"
|
|
):
|
|
with R.dataflow():
|
|
cls = Expected
|
|
y: R.Tensor((2, 4, 3), "float32") = R.reshape(x, (2, 4, 3))
|
|
# Note: `z` is the output var of the dataflow block, and is thus
|
|
# not expected to be rewritten.
|
|
z = R.call_tir(
|
|
cls.expand_dims, (y,), out_ty=R.Tensor((2, 1, 4, 1, 3), dtype="float32")
|
|
)
|
|
R.output(z)
|
|
return z
|
|
|
|
assert relax.analysis.has_reshape_pattern(Module["expand_dims"])
|
|
mod = relax.transform.RewriteDataflowReshape()(Module)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_reshape_pattern_detect():
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def reshape(rxplaceholder: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float32"), T_reshape: T.Buffer((T.int64(2), T.int64(4096), T.int64(5), T.int64(64)), "float32")):
|
|
for ax0_ax1_ax2_ax3_fused_1 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
|
|
for ax0_ax1_ax2_ax3_fused_2 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
|
|
for ax0_ax1_ax2_ax3_fused_0 in range(T.int64(10)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0 = T.axis.spatial(T.int64(2), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) // T.int64(1310720))
|
|
v_ax1 = T.axis.spatial(T.int64(4096), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(1310720) // T.int64(320))
|
|
v_ax2 = T.axis.spatial(T.int64(5), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(320) // T.int64(64))
|
|
v_ax3 = T.axis.spatial(T.int64(64), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(64))
|
|
T.reads(rxplaceholder[(((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096) + v_ax0) % T.int64(2), ((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096), (v_ax2 * T.int64(64) + v_ax3) % T.int64(320)])
|
|
T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[(((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096) + v_ax0) % T.int64(2), ((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096), (v_ax2 * T.int64(64) + v_ax3) % T.int64(320)]
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def expand_dims(
|
|
rxplaceholder: T.Buffer((T.int64(2), T.int64(4096), T.int64(5), T.int64(64)), "float32"),
|
|
expand_dims: T.Buffer(
|
|
(T.int64(2), T.int64(1), T.int64(4096), T.int64(1), T.int64(5), T.int64(64)),
|
|
"float32",
|
|
),
|
|
):
|
|
for i0, i1, i2, i3, i4, i5 in T.grid(
|
|
T.int64(2), T.int64(1), T.int64(4096), T.int64(1), T.int64(5), T.int64(64)
|
|
):
|
|
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, i5_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, i5_1]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 4096, 320), dtype="float32")
|
|
) -> R.Tensor((2, 1, 4096, 1, 5, 64), dtype="float32"):
|
|
cls = Module
|
|
with R.dataflow():
|
|
y = R.call_tir(cls.reshape, (x,), out_ty=R.Tensor((2, 4096, 5, 64), dtype="float32"))
|
|
z = R.call_tir(
|
|
cls.expand_dims, (y,), out_ty=R.Tensor((2, 1, 4096, 1, 5, 64), "float32")
|
|
)
|
|
R.output(z)
|
|
return z
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def expand_dims(rxplaceholder: T.Buffer((T.int64(2), T.int64(4096), T.int64(5), T.int64(64)), "float32"), expand_dims_1: T.Buffer((T.int64(2), T.int64(1), T.int64(4096), T.int64(1), T.int64(5), T.int64(64)), "float32")):
|
|
# with T.sblock("root"):
|
|
for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(2), T.int64(1), T.int64(4096), T.int64(1), T.int64(5), T.int64(64)):
|
|
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, i5_1])
|
|
T.writes(expand_dims_1[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1])
|
|
expand_dims_1[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1] = rxplaceholder[i0_1, i2_1, i4_1, i5_1]
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def reshape(rxplaceholder: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float32"), T_reshape: T.Buffer((T.int64(2), T.int64(4096), T.int64(5), T.int64(64)), "float32")):
|
|
# with T.sblock("root"):
|
|
for ax0_ax1_ax2_ax3_fused_1 in T.thread_binding(T.int64(256), thread="blockIdx.x"):
|
|
for ax0_ax1_ax2_ax3_fused_2 in T.thread_binding(T.int64(1024), thread="threadIdx.x"):
|
|
for ax0_ax1_ax2_ax3_fused_0 in range(T.int64(10)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0 = T.axis.spatial(T.int64(2), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) // T.int64(1310720))
|
|
v_ax1 = T.axis.spatial(T.int64(4096), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(1310720) // T.int64(320))
|
|
v_ax2 = T.axis.spatial(T.int64(5), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(320) // T.int64(64))
|
|
v_ax3 = T.axis.spatial(T.int64(64), (ax0_ax1_ax2_ax3_fused_0 * T.int64(262144) + ax0_ax1_ax2_ax3_fused_1 * T.int64(1024) + ax0_ax1_ax2_ax3_fused_2) % T.int64(64))
|
|
T.reads(rxplaceholder[(((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096) + v_ax0) % T.int64(2), ((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096), (v_ax2 * T.int64(64) + v_ax3) % T.int64(320)])
|
|
T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[(((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096) + v_ax0) % T.int64(2), ((v_ax2 * T.int64(64) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096), (v_ax2 * T.int64(64) + v_ax3) % T.int64(320)]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((2, 4096, 320), dtype="float32")) -> R.Tensor((2, 1, 4096, 1, 5, 64), dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
y: R.Tensor((2, 4096, 5, 64), dtype="float32") = R.reshape(x, R.shape([2, 4096, 5, 64]))
|
|
z = R.call_tir(cls.expand_dims, (y,), out_ty=R.Tensor((2, 1, 4096, 1, 5, 64), dtype="float32"))
|
|
R.output(z)
|
|
return z
|
|
# fmt: on
|
|
|
|
assert relax.analysis.has_reshape_pattern(Module["reshape"])
|
|
mod = relax.transform.RewriteDataflowReshape()(Module)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_reshape_dynamic_shape():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(var_A: T.handle, var_T_reshape: T.handle):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
n = T.int32()
|
|
A = T.match_buffer(var_A, (n, 16, 128), "float16")
|
|
T_reshape = T.match_buffer(var_T_reshape, (1, n, 16, 128), "float16")
|
|
# with T.sblock("root"):
|
|
for ax0_ax1_ax2_fused_0 in T.thread_binding(n * 2, thread="blockIdx.x"):
|
|
for ax0_ax1_ax2_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
|
with T.sblock("T_reshape"):
|
|
v0 = T.axis.spatial(
|
|
n, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) // 2048
|
|
)
|
|
v1 = T.axis.spatial(
|
|
16, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % 2048 // 128
|
|
)
|
|
v2 = T.axis.spatial(
|
|
128, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % 128
|
|
)
|
|
T.reads(
|
|
A[((v2 // 128 + v1) // 32 + v0) % n, (v2 // 128 + v1) % 32, v2 % 128]
|
|
)
|
|
T.writes(T_reshape[0, v0, v1, v2])
|
|
T_reshape[0, v0, v1, v2] = A[
|
|
((v2 // 128 + v1) // 32 + v0) % n, (v2 // 128 + v1) % 32, v2 % 128
|
|
]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((8, 16, 128), dtype="float16")) -> R.Tensor(
|
|
(1, 8, 16, 128), dtype="float16"
|
|
):
|
|
cls = Module
|
|
with R.dataflow():
|
|
y = R.call_tir(cls.reshape, (x,), out_ty=R.Tensor((1, 8, 16, 128), dtype="float16"))
|
|
z = R.add(y, R.const(1, "float16"))
|
|
R.output(z)
|
|
return z
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(var_A: T.handle, var_T_reshape: T.handle):
|
|
T.func_attr({"tirx.is_scheduled": True, "tirx.noalias": True})
|
|
n = T.int32()
|
|
A = T.match_buffer(var_A, (n, 16, 128), "float16")
|
|
T_reshape = T.match_buffer(var_T_reshape, (1, n, 16, 128), "float16")
|
|
# with T.sblock("root"):
|
|
for ax0_ax1_ax2_fused_0 in T.thread_binding(n * 2, thread="blockIdx.x"):
|
|
for ax0_ax1_ax2_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
|
with T.sblock("T_reshape"):
|
|
v0 = T.axis.spatial(
|
|
n, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) // 2048
|
|
)
|
|
v1 = T.axis.spatial(
|
|
16, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % 2048 // 128
|
|
)
|
|
v2 = T.axis.spatial(
|
|
128, (ax0_ax1_ax2_fused_0 * 1024 + ax0_ax1_ax2_fused_1) % 128
|
|
)
|
|
T.reads(
|
|
A[((v2 // 128 + v1) // 32 + v0) % n, (v2 // 128 + v1) % 32, v2 % 128]
|
|
)
|
|
T.writes(T_reshape[0, v0, v1, v2])
|
|
T_reshape[0, v0, v1, v2] = A[
|
|
((v2 // 128 + v1) // 32 + v0) % n, (v2 // 128 + v1) % 32, v2 % 128
|
|
]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((8, 16, 128), dtype="float16")) -> R.Tensor(
|
|
(1, 8, 16, 128), dtype="float16"
|
|
):
|
|
with R.dataflow():
|
|
y: R.Tensor((1, 8, 16, 128), dtype="float16") = R.reshape(
|
|
x, R.shape([1, 8, 16, 128])
|
|
)
|
|
z: R.Tensor((1, 8, 16, 128), dtype="float16") = R.add(y, R.const(1, "float16"))
|
|
R.output(z)
|
|
return z
|
|
|
|
assert relax.analysis.has_reshape_pattern(Module["reshape"])
|
|
mod = relax.transform.RewriteDataflowReshape()(Module)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_reshape_non_dataflow():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def reshape(
|
|
rxplaceholder: T.Buffer((T.int64(8), T.int64(3)), "float32"),
|
|
T_reshape: T.Buffer((T.int64(2), T.int64(4), T.int64(3)), "float32"),
|
|
):
|
|
for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(3)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
|
T.reads(
|
|
rxplaceholder[
|
|
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
|
|
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
|
|
]
|
|
)
|
|
T.writes(T_reshape[v_ax0, v_ax1, v_ax2])
|
|
T_reshape[v_ax0, v_ax1, v_ax2] = rxplaceholder[
|
|
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) // T.int64(3),
|
|
(v_ax0 * 12 + v_ax1 * 3 + v_ax2) % T.int64(3),
|
|
]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((8, 3), dtype="float32")) -> R.Tensor((2, 4, 3), dtype="float32"):
|
|
cls = Module
|
|
y = R.call_tir(cls.reshape, (x,), out_ty=R.Tensor((2, 4, 3), dtype="float32"))
|
|
return y
|
|
|
|
assert relax.analysis.has_reshape_pattern(Module["reshape"])
|
|
# The binding var of the call_tir is not a DataflowVar. So the pass does no change.
|
|
mod = relax.transform.RewriteDataflowReshape()(Module)
|
|
tvm.ir.assert_structural_equal(mod, Module)
|
|
|
|
|
|
def test_tuple_get_reshape():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def fused_reshape5(
|
|
lv2_0: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
|
|
lv2_1: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
|
|
lv2_2: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
|
|
T_reshape_handle_intermediate: T.Buffer(
|
|
(T.int64(2), T.int64(4096), T.int64(8), T.int64(40)), "float16"
|
|
),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(4096), T.int64(8), T.int64(40)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(
|
|
lv2_0[
|
|
(
|
|
((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1)
|
|
// T.int64(4096)
|
|
+ v_ax0
|
|
)
|
|
% T.int64(2),
|
|
((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096),
|
|
(v_ax2 * T.int64(40) + v_ax3) % T.int64(320),
|
|
]
|
|
)
|
|
T.writes(T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = lv2_0[
|
|
(
|
|
((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096)
|
|
+ v_ax0
|
|
)
|
|
% T.int64(2),
|
|
((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096),
|
|
(v_ax2 * T.int64(40) + v_ax3) % T.int64(320),
|
|
]
|
|
|
|
@R.function
|
|
def main(
|
|
lv41_1: R.Tuple(
|
|
R.Tensor((2, 4096, 320), dtype="float16"),
|
|
R.Tensor((2, 4096, 320), dtype="float16"),
|
|
R.Tensor((2, 4096, 320), dtype="float16"),
|
|
),
|
|
) -> R.Tensor((2, 4096, 8, 40), dtype="float16"):
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[0]
|
|
lv1: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[1]
|
|
lv2: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[2]
|
|
lv645 = R.call_tir(
|
|
cls.fused_reshape5,
|
|
(lv, lv1, lv2),
|
|
out_ty=R.Tensor((2, 4096, 8, 40), dtype="float16"),
|
|
)
|
|
out: R.Tensor((2, 4096, 8, 40), dtype="float16") = R.add(lv645, lv645)
|
|
R.output(out)
|
|
return out
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def fused_reshape5(
|
|
lv2_0: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
|
|
lv2_1: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
|
|
lv2_2: T.Buffer((T.int64(2), T.int64(4096), T.int64(320)), "float16"),
|
|
T_reshape_handle_intermediate: T.Buffer(
|
|
(T.int64(2), T.int64(4096), T.int64(8), T.int64(40)), "float16"
|
|
),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(4096), T.int64(8), T.int64(40)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(
|
|
lv2_0[
|
|
(
|
|
((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1)
|
|
// T.int64(4096)
|
|
+ v_ax0
|
|
)
|
|
% T.int64(2),
|
|
((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096),
|
|
(v_ax2 * T.int64(40) + v_ax3) % T.int64(320),
|
|
]
|
|
)
|
|
T.writes(T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_reshape_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = lv2_0[
|
|
(
|
|
((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) // T.int64(4096)
|
|
+ v_ax0
|
|
)
|
|
% T.int64(2),
|
|
((v_ax2 * T.int64(40) + v_ax3) // T.int64(320) + v_ax1) % T.int64(4096),
|
|
(v_ax2 * T.int64(40) + v_ax3) % T.int64(320),
|
|
]
|
|
|
|
@R.function
|
|
def main(
|
|
lv41_1: R.Tuple(
|
|
R.Tensor((2, 4096, 320), dtype="float16"),
|
|
R.Tensor((2, 4096, 320), dtype="float16"),
|
|
R.Tensor((2, 4096, 320), dtype="float16"),
|
|
),
|
|
) -> R.Tensor((2, 4096, 8, 40), dtype="float16"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[0]
|
|
lv1: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[1]
|
|
lv2: R.Tensor((2, 4096, 320), dtype="float16") = lv41_1[2]
|
|
lv645: R.Tensor((2, 4096, 8, 40), dtype="float16") = R.reshape(
|
|
lv, R.shape([2, 4096, 8, 40])
|
|
)
|
|
out: R.Tensor((2, 4096, 8, 40), dtype="float16") = R.add(lv645, lv645)
|
|
R.output(out)
|
|
return out
|
|
|
|
rewritten = relax.transform.RewriteDataflowReshape()(Module)
|
|
tvm.ir.assert_structural_equal(rewritten, Expected)
|
|
|
|
|
|
def test_invalid_reshape():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
# The strided_slice op has the reshape pattern, but it can take only a part of the input.
|
|
# It can't be replaced with the reshape op because reshape expects to preserve the "volume"
|
|
# of the input.
|
|
@T.prim_func(s_tir=True)
|
|
def strided_slice(
|
|
A: T.Buffer((T.int64(1), T.int64(1024)), "int32"),
|
|
T_strided_slice: T.Buffer((T.int64(1), T.int64(1000)), "int32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(1), T.int64(1000)):
|
|
with T.sblock("T_strided_slice"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(A[v_ax0, v_ax1])
|
|
T.writes(T_strided_slice[v_ax0, v_ax1])
|
|
T_strided_slice[v_ax0, v_ax1] = A[v_ax0, v_ax1]
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def add_one(
|
|
A: T.Buffer((T.int64(1), T.int64(1000)), "int32"),
|
|
T_add_one: T.buffer((T.int64(1), T.int64(1000)), "int32"),
|
|
):
|
|
for ax0, ax1 in T.grid(T.int64(1), T.int64(1000)):
|
|
with T.sblock("T_add_one"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(A[v_ax0, v_ax1])
|
|
T.writes(T_add_one[v_ax0, v_ax1])
|
|
T_add_one[v_ax0, v_ax1] = A[v_ax0, v_ax1] + 1
|
|
|
|
@R.function
|
|
def main(A: R.Tensor((1, 1024), dtype="int32")) -> R.Tensor((1, 1000), dtype="int32"):
|
|
with R.dataflow():
|
|
cls = Module
|
|
S = R.call_tir(cls.strided_slice, (A,), out_ty=R.Tensor((1, 1000), dtype="int32"))
|
|
A = R.call_tir(cls.add_one, (S,), out_ty=R.Tensor((1, 1000), dtype="int32"))
|
|
R.output(A)
|
|
return A
|
|
|
|
assert relax.analysis.has_reshape_pattern(Module["strided_slice"])
|
|
rewritten = relax.transform.RewriteDataflowReshape()(Module)
|
|
tvm.ir.assert_structural_equal(rewritten, Module)
|
|
|
|
|
|
def test_reshape_detect_nop():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main(x: R.Tensor((8, 8), dtype="float16")) -> R.Tensor((8, 8), dtype="float16"):
|
|
with R.dataflow():
|
|
gv = R.call_pure_packed("foo", x, x, ty_args=(R.Tensor((8, 8), dtype="float16"),))
|
|
out = R.call_pure_packed(
|
|
"foo", gv, gv, ty_args=(R.Tensor((8, 8), dtype="float16"),)
|
|
)
|
|
R.output(out)
|
|
return out
|
|
|
|
rewritten = relax.transform.RewriteDataflowReshape()(Module)
|
|
tvm.ir.assert_structural_equal(rewritten, Module)
|
|
|
|
|
|
def test_reshape_scalar():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main(x: R.Tensor((), dtype="float32")) -> R.Tensor((1,), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1,), dtype="float32") = R.reshape(x, [1])
|
|
lv2: R.Tensor((1,), dtype="float32") = R.add(lv1, lv1)
|
|
R.output(lv2)
|
|
return lv2
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(
|
|
A: T.Buffer((T.int64(1),), "float32"),
|
|
B: T.Buffer((T.int64(1),), "float32"),
|
|
T_add: T.Buffer((T.int64(1),), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0 in range(T.int64(1)):
|
|
with T.sblock("T_add"):
|
|
v_ax0 = T.axis.spatial(T.int64(1), ax0)
|
|
T.reads(A[v_ax0], B[v_ax0])
|
|
T.writes(T_add[v_ax0])
|
|
T_add[v_ax0] = A[v_ax0] + B[v_ax0]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(A: T.Buffer((), "float32"), T_reshape: T.Buffer((T.int64(1),), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0 in range(T.int64(1)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0 = T.axis.spatial(T.int64(1), ax0)
|
|
T.reads(A[()])
|
|
T.writes(T_reshape[v_ax0])
|
|
T_reshape[v_ax0] = A[()]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((), dtype="float32")) -> R.Tensor((1,), dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1,), dtype="float32") = R.reshape(x, R.shape([1]))
|
|
lv2 = R.call_tir(cls.add, (lv1, lv1), out_ty=R.Tensor((1,), dtype="float32"))
|
|
R.output(lv2)
|
|
return lv2
|
|
|
|
mod = Module
|
|
mod = relax.transform.LegalizeOps()(mod)
|
|
rewritten = relax.transform.RewriteDataflowReshape()(mod)
|
|
tvm.ir.assert_structural_equal(rewritten, Expected)
|
|
|
|
|
|
def test_rewrite_static_reshape():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor([256], dtype="float32")):
|
|
with R.dataflow():
|
|
y = R.reshape(x, [64, 4])
|
|
z = R.add(y, y)
|
|
R.output(z)
|
|
return z
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((256,), dtype="float32")):
|
|
cls = Expected
|
|
|
|
with R.dataflow():
|
|
y = R.reshape(x, R.shape([64, 4]))
|
|
z = R.call_tir(cls.add, (y, y), out_ty=R.Tensor((64, 4), dtype="float32"))
|
|
R.output(z)
|
|
return z
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(
|
|
y1: T.Buffer((T.int64(64), T.int64(4)), "float32"),
|
|
y2: T.Buffer((T.int64(64), T.int64(4)), "float32"),
|
|
z: T.Buffer((T.int64(64), T.int64(4)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
|
|
for iters in T.grid(T.int64(64), T.int64(4)):
|
|
with T.sblock("T_add"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
z[i, j] = y1[i, j] + y2[i, j]
|
|
|
|
After = tvm.ir.transform.Sequential(
|
|
[
|
|
# Lower both R.reshape and R.add from Relax to TIR
|
|
relax.transform.LegalizeOps(),
|
|
# Identify reshapes, raise calls to cls.reshape from TIR
|
|
# to Relax
|
|
relax.transform.RewriteDataflowReshape(),
|
|
# Clean up afterwards, removing the no-longer-required
|
|
# PrimFunc "reshape"
|
|
relax.transform.DeadCodeElimination(),
|
|
]
|
|
)(Before)
|
|
|
|
tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
# def test_rewrite_dynamic_reshape():
|
|
# @I.ir_module
|
|
# class Before:
|
|
# @R.function
|
|
# def main(x: R.Tensor(["N"], dtype="float32")):
|
|
# N = T.int64()
|
|
# with R.dataflow():
|
|
# y = R.reshape(x, [N // 4, 4])
|
|
# z = R.add(y, y)
|
|
# R.output(z)
|
|
# return z
|
|
|
|
# @I.ir_module
|
|
# class Expected:
|
|
# @R.function
|
|
# def main(x: R.Tensor(["N"], dtype="float32")):
|
|
# N = T.int64()
|
|
# cls = Expected
|
|
|
|
# with R.dataflow():
|
|
# y = R.reshape(x, R.shape([N // 4, 4]))
|
|
# z = R.call_tir(
|
|
# cls.add,
|
|
# (y, y),
|
|
# tir_vars=[N],
|
|
# out_ty=R.Tensor((N // 4, 4), dtype="float32"),
|
|
# )
|
|
# R.output(z)
|
|
# return z
|
|
|
|
# @T.prim_func(private=True)
|
|
# def add(
|
|
# y1_handle: T.handle,
|
|
# y2_handle: T.handle,
|
|
# z_handle: T.handle,
|
|
# N: T.int64,
|
|
# ):
|
|
|
|
# y1 = T.match_buffer(y1_handle, [N // 4, 4], "float32")
|
|
# y2 = T.match_buffer(y2_handle, [N // 4, 4], "float32")
|
|
# z = T.match_buffer(z_handle, [N // 4, 4], "float32")
|
|
|
|
# T.func_attr({"tirx.noalias": True})
|
|
|
|
# for iters in T.grid(T.int64(64), T.int64(4)):
|
|
# with T.sblock("T_add"):
|
|
# i, j = T.axis.remap("SS", iters)
|
|
# z[i, j] = y1[i, j] + y2[i, j]
|
|
|
|
# After = tvm.ir.transform.Sequential(
|
|
# [
|
|
# # Lower both R.reshape and R.add from Relax to TIR
|
|
# relax.transform.LegalizeOps(),
|
|
# # Identify reshapes, raise calls to cls.reshape from TIR
|
|
# # to Relax
|
|
# relax.transform.RewriteDataflowReshape(),
|
|
# # Clean up afterwards, removing the no-longer-required
|
|
# # PrimFunc "reshape"
|
|
# relax.transform.DeadCodeElimination(),
|
|
# ]
|
|
# )(Before)
|
|
# After.show()
|
|
# tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
def test_rewrite_dynamic_reshape():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor(["N", 16], dtype="float32")):
|
|
N = T.int64()
|
|
with R.dataflow():
|
|
y = R.reshape(x, [N * 4, T.int64(4)])
|
|
z = R.add(y, y)
|
|
R.output(z)
|
|
return z
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(["N", 16], dtype="float32")):
|
|
N = T.int64()
|
|
cls = Expected
|
|
|
|
with R.dataflow():
|
|
y = R.reshape(x, R.shape([N * 4, T.int64(4)]))
|
|
z = R.call_tir(
|
|
cls.add,
|
|
(y, y),
|
|
tir_vars=[N],
|
|
out_ty=R.Tensor((N * 4, 4), dtype="float32"),
|
|
)
|
|
R.output(z)
|
|
return z
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(
|
|
y1_handle: T.handle,
|
|
y2_handle: T.handle,
|
|
z_handle: T.handle,
|
|
N: T.int64,
|
|
):
|
|
y1 = T.match_buffer(y1_handle, [N * 4, T.int64(4)], "float32")
|
|
y2 = T.match_buffer(y2_handle, [N * 4, T.int64(4)], "float32")
|
|
z = T.match_buffer(z_handle, [N * 4, T.int64(4)], "float32")
|
|
|
|
T.func_attr({"tirx.noalias": True})
|
|
|
|
for iters in T.grid(N * 4, T.int64(4)):
|
|
with T.sblock("T_add"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
z[i, j] = y1[i, j] + y2[i, j]
|
|
|
|
After = tvm.ir.transform.Sequential(
|
|
[
|
|
# Lower both R.reshape and R.add from Relax to TIR
|
|
relax.transform.LegalizeOps(),
|
|
# Identify reshapes, raise calls to cls.reshape from TIR
|
|
# to Relax
|
|
relax.transform.RewriteDataflowReshape(),
|
|
# Clean up afterwards, removing the no-longer-required
|
|
# PrimFunc "reshape"
|
|
relax.transform.DeadCodeElimination(),
|
|
]
|
|
)(Before)
|
|
tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|