wehub-resource-sync
699 lines
25 KiB
Python
699 lines
25 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: E731, F401, F841
|
|
import tvm.testing
|
|
from tvm import relax, tirx
|
|
from tvm.relax.transform import CombineParallelMatmul
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
from tvm.script.ir_builder import IRBuilder
|
|
from tvm.script.ir_builder import relax as relax_builder
|
|
|
|
|
|
def get_parallel_matmul(
|
|
num_branches,
|
|
lhs_shape=(640, 640),
|
|
rhs_shape=(640, 640),
|
|
with_bias=None,
|
|
activation=None,
|
|
):
|
|
dtype = "float32"
|
|
|
|
activation_map = {"relu": R.nn.relu, "gelu": R.nn.gelu}
|
|
|
|
with IRBuilder() as builder:
|
|
with relax_builder.function():
|
|
R.func_name("main")
|
|
x = R.arg("x", R.Tensor(lhs_shape, dtype))
|
|
|
|
rhs = []
|
|
bias = []
|
|
|
|
for i in range(num_branches):
|
|
rhs.append(R.arg("y", R.Tensor(rhs_shape, dtype)))
|
|
|
|
if with_bias and with_bias[i]:
|
|
bias.append(R.arg("bias", R.Tensor((rhs_shape[1],), dtype)))
|
|
else:
|
|
bias.append(None)
|
|
|
|
with R.dataflow() as frame:
|
|
branches = []
|
|
|
|
for i, r in enumerate(rhs):
|
|
result = R.emit(R.matmul(x, r, out_dtype=dtype))
|
|
if bias[i]:
|
|
result = R.emit(result + bias[i])
|
|
if activation and activation[i]:
|
|
result = R.emit(activation_map[activation[i]](result))
|
|
|
|
branches.append(result)
|
|
|
|
R.output(R.emit(R.concat(branches, axis=1)))
|
|
|
|
R.func_ret_value(frame.output_vars[0])
|
|
|
|
func = builder.get()
|
|
return tvm.IRModule({"main": func})
|
|
|
|
|
|
def test_simple():
|
|
mod_orig = get_parallel_matmul(1)
|
|
mod = CombineParallelMatmul()(mod_orig)
|
|
|
|
tvm.ir.assert_structural_equal(mod, mod_orig)
|
|
|
|
mod = get_parallel_matmul(3)
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected1(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
|
|
lv_1 = lv2[0]
|
|
lv1_1 = lv2[1]
|
|
lv2_1 = lv2[2]
|
|
lv3 = R.concat((lv_1, lv1_1, lv2_1), axis=1)
|
|
R.output(lv3)
|
|
return lv3
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected1.with_attr("global_symbol", "main"))
|
|
|
|
# Test a batched LHS case, slicing is done on the axis 2
|
|
mod = get_parallel_matmul(3, lhs_shape=(2, 1024, 640))
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected2(
|
|
x: R.Tensor((2, 1024, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
) -> R.Tensor((2, 3072, 640), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=2)
|
|
lv_1 = lv2[0]
|
|
lv1_1 = lv2[1]
|
|
lv2_1 = lv2[2]
|
|
lv3 = R.concat((lv_1, lv1_1, lv2_1), axis=1)
|
|
R.output(lv3)
|
|
return lv3
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))
|
|
|
|
|
|
def test_bias():
|
|
mod = get_parallel_matmul(3, with_bias=[True, True, True])
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected1(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
bias: R.Tensor((640,), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
bias_1: R.Tensor((640,), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
bias_2: R.Tensor((640,), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.concat((bias, bias_1, bias_2), axis=0)
|
|
lv3 = R.add(lv1, lv2)
|
|
lv4 = R.split(lv3, indices_or_sections=[640, 1280], axis=1)
|
|
lv1_1 = lv4[0]
|
|
lv3_1 = lv4[1]
|
|
lv5 = lv4[2]
|
|
lv6 = R.concat((lv1_1, lv3_1, lv5), axis=1)
|
|
R.output(lv6)
|
|
return lv6
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected1.with_attr("global_symbol", "main"))
|
|
|
|
mod = get_parallel_matmul(3, with_bias=[True, False, True])
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected2(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
bias: R.Tensor((640,), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
bias_1: R.Tensor((640,), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
|
|
lv_1 = lv2[0]
|
|
lv1_1 = R.add(lv_1, bias)
|
|
lv2_1 = lv2[1]
|
|
lv3 = lv2[2]
|
|
lv4 = R.add(lv3, bias_1)
|
|
lv5 = R.concat((lv1_1, lv2_1, lv4), axis=1)
|
|
R.output(lv5)
|
|
return lv5
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))
|
|
|
|
|
|
def test_activation():
|
|
mod = get_parallel_matmul(3, activation=["relu", "relu", "relu"])
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected1(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.nn.relu(lv1)
|
|
lv3 = R.split(lv2, indices_or_sections=[640, 1280], axis=1)
|
|
lv1_1 = lv3[0]
|
|
lv3_1 = lv3[1]
|
|
lv5 = lv3[2]
|
|
lv6 = R.concat((lv1_1, lv3_1, lv5), axis=1)
|
|
R.output(lv6)
|
|
return lv6
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected1.with_attr("global_symbol", "main"))
|
|
|
|
mod = get_parallel_matmul(3, activation=["gelu", "relu", "relu"])
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected2(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
|
|
lv_1 = lv2[0]
|
|
lv1_1 = R.nn.gelu(lv_1)
|
|
lv2_1 = lv2[1]
|
|
lv3 = R.nn.relu(lv2_1)
|
|
lv4 = lv2[2]
|
|
lv5 = R.nn.relu(lv4)
|
|
lv6 = R.concat((lv1_1, lv3, lv5), axis=1)
|
|
R.output(lv6)
|
|
return lv6
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))
|
|
|
|
mod = get_parallel_matmul(3, activation=["relu", None, None])
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected3(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
|
|
|
|
lv_1 = lv2[0]
|
|
lv1_1 = R.nn.relu(lv_1)
|
|
lv2_1 = lv2[1]
|
|
lv3 = lv2[2]
|
|
lv4 = R.concat((lv1_1, lv2_1, lv3), axis=1)
|
|
R.output(lv4)
|
|
return lv4
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected3.with_attr("global_symbol", "main"))
|
|
|
|
|
|
def test_bias_activation():
|
|
mod = get_parallel_matmul(3, with_bias=[True, True, True], activation=["relu", "relu", "relu"])
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected1(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
bias: R.Tensor((640,), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
bias_1: R.Tensor((640,), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
bias_2: R.Tensor((640,), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.concat((bias, bias_1, bias_2), axis=0)
|
|
lv3 = R.add(lv1, lv2)
|
|
lv4 = R.nn.relu(lv3)
|
|
lv5 = R.split(lv4, indices_or_sections=[640, 1280], axis=1)
|
|
lv2_1 = lv5[0]
|
|
lv5_1 = lv5[1]
|
|
lv8 = lv5[2]
|
|
lv9 = R.concat((lv2_1, lv5_1, lv8), axis=1)
|
|
R.output(lv9)
|
|
return lv9
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected1.with_attr("global_symbol", "main"))
|
|
|
|
mod = get_parallel_matmul(3, with_bias=[True, True, True], activation=["relu", None, "relu"])
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected2(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
bias: R.Tensor((640,), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
bias_1: R.Tensor((640,), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
bias_2: R.Tensor((640,), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.concat((bias, bias_1, bias_2), axis=0)
|
|
lv3 = R.add(lv1, lv2)
|
|
lv4 = R.split(lv3, indices_or_sections=[640, 1280], axis=1)
|
|
lv1_1 = lv4[0]
|
|
lv2_1 = R.nn.relu(lv1_1)
|
|
lv4_1 = lv4[1]
|
|
lv6 = lv4[2]
|
|
lv7 = R.nn.relu(lv6)
|
|
lv8 = R.concat((lv2_1, lv4_1, lv7), axis=1)
|
|
R.output(lv8)
|
|
return lv8
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))
|
|
|
|
mod = get_parallel_matmul(3, with_bias=[True, False, True], activation=["relu", None, "relu"])
|
|
mod = CombineParallelMatmul()(mod)
|
|
|
|
@R.function
|
|
def expected3(
|
|
x: R.Tensor((640, 640), dtype="float32"),
|
|
y: R.Tensor((640, 640), dtype="float32"),
|
|
bias: R.Tensor((640,), dtype="float32"),
|
|
y_1: R.Tensor((640, 640), dtype="float32"),
|
|
y_2: R.Tensor((640, 640), dtype="float32"),
|
|
bias_1: R.Tensor((640,), dtype="float32"),
|
|
) -> R.Tensor((640, 1920), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.concat((y, y_1, y_2), axis=1)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=1)
|
|
lv_1 = lv2[0]
|
|
lv1_1 = R.add(lv_1, bias)
|
|
lv2_1 = R.nn.relu(lv1_1)
|
|
lv3 = lv2[1]
|
|
lv4 = lv2[2]
|
|
lv5 = R.add(lv4, bias_1)
|
|
lv6 = R.nn.relu(lv5)
|
|
lv7 = R.concat((lv2_1, lv3, lv6), axis=1)
|
|
R.output(lv7)
|
|
return lv7
|
|
|
|
tvm.ir.assert_structural_equal(mod["main"], expected3.with_attr("global_symbol", "main"))
|
|
|
|
|
|
def test_rhs_batched():
|
|
@R.function(private=True)
|
|
def before(
|
|
x: R.Tensor((1024, 640), "float32"),
|
|
w0: R.Tensor((2, 640, 640), "float32"),
|
|
w1: R.Tensor((640, 640), "float32"),
|
|
w2: R.Tensor((2, 640, 640), "float32"),
|
|
w3: R.Tensor((3, 4, 640, 640), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x, w0)
|
|
lv1 = R.matmul(x, w1)
|
|
lv2 = R.matmul(x, w2)
|
|
lv3 = R.matmul(x, w3)
|
|
out = (lv0, lv1, lv2, lv3)
|
|
R.output(out)
|
|
return out
|
|
|
|
after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
|
|
|
|
@R.function(private=True)
|
|
def expected(
|
|
x: R.Tensor((1024, 640), dtype="float32"),
|
|
w0: R.Tensor((2, 640, 640), dtype="float32"),
|
|
w1: R.Tensor((640, 640), dtype="float32"),
|
|
w2: R.Tensor((2, 640, 640), dtype="float32"),
|
|
w3: R.Tensor((3, 4, 640, 640), dtype="float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv = R.concat((w0, w2), axis=2)
|
|
lv1 = R.matmul(x, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640], axis=2)
|
|
lv0 = lv2[0]
|
|
lv1_1 = R.matmul(x, w1, out_dtype=None)
|
|
lv2_1 = lv2[1]
|
|
lv3 = R.matmul(x, w3, out_dtype=None)
|
|
out = lv0, lv1_1, lv2_1, lv3
|
|
R.output(out)
|
|
return out
|
|
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
@tvm.script.ir_module
|
|
class four_matmul_incompatible_batches:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1024, 640), "float32"),
|
|
w0: R.Tensor((2, 640, 640), "float32"),
|
|
w1: R.Tensor((3, 640, 640), "float32"),
|
|
w2: R.Tensor((2, 640, 640), "float32"),
|
|
w3: R.Tensor((2, 640, 640), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x, w0)
|
|
lv1 = R.matmul(x, w1)
|
|
lv2 = R.matmul(x, w2)
|
|
lv3 = R.matmul(x, w3)
|
|
out = (lv0, lv1, lv2, lv3)
|
|
R.output(out)
|
|
return out
|
|
|
|
mod = CombineParallelMatmul()(four_matmul_incompatible_batches)
|
|
# For now, when rhs matrices have the same rank but different batch sizes, we don't
|
|
# combine any of them.
|
|
tvm.ir.assert_structural_equal(mod, four_matmul_incompatible_batches)
|
|
|
|
|
|
def test_multiple_combine():
|
|
@R.function(private=True)
|
|
def before(
|
|
x1: R.Tensor((2, 1024, 640), "float32"),
|
|
x2: R.Tensor((2, 1024, 640), "float32"),
|
|
w0: R.Tensor((640, 640), "float32"),
|
|
w1: R.Tensor((640, 640), "float32"),
|
|
w2: R.Tensor((640, 640), "float32"),
|
|
w3: R.Tensor((640, 640), "float32"),
|
|
w4: R.Tensor((640, 640), "float32"),
|
|
b0: R.Tensor((640,), "float32"),
|
|
b1: R.Tensor((640,), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x1, w0)
|
|
lv3 = R.matmul(x2, w3)
|
|
lv1 = R.matmul(x1, w1)
|
|
lv5 = R.add(lv3, b0)
|
|
lv2 = R.matmul(x1, w2)
|
|
lv4 = R.matmul(x2, w4)
|
|
lv6 = R.add(lv4, b1)
|
|
out = (lv0, lv1, lv2, lv5, lv6)
|
|
R.output(out)
|
|
return out
|
|
|
|
after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
|
|
|
|
@R.function(private=True)
|
|
def expected(
|
|
x1: R.Tensor((2, 1024, 640), dtype="float32"),
|
|
x2: R.Tensor((2, 1024, 640), dtype="float32"),
|
|
w0: R.Tensor((640, 640), dtype="float32"),
|
|
w1: R.Tensor((640, 640), dtype="float32"),
|
|
w2: R.Tensor((640, 640), dtype="float32"),
|
|
w3: R.Tensor((640, 640), dtype="float32"),
|
|
w4: R.Tensor((640, 640), dtype="float32"),
|
|
b0: R.Tensor((640,), dtype="float32"),
|
|
b1: R.Tensor((640,), dtype="float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv = R.concat((w0, w1, w2), axis=1)
|
|
lv1 = R.matmul(x1, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=2)
|
|
lv0 = lv2[0]
|
|
lv1_1 = lv2[1]
|
|
lv_1 = R.concat((w3, w4), axis=1)
|
|
lv1_2 = R.matmul(x2, lv_1, out_dtype="float32")
|
|
lv2_1 = R.concat((b0, b1), axis=0)
|
|
lv3 = R.add(lv1_2, lv2_1)
|
|
lv4 = R.split(lv3, indices_or_sections=[640], axis=2)
|
|
lv5 = lv4[0]
|
|
lv2_2 = lv2[2]
|
|
lv6 = lv4[1]
|
|
out = lv0, lv1_1, lv2_2, lv5, lv6
|
|
R.output(out)
|
|
return out
|
|
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_check():
|
|
@R.function(private=True)
|
|
def before(
|
|
x1: R.Tensor((2, 1024, 640), "float32"),
|
|
x2: R.Tensor((2, 1024, 640), "float32"),
|
|
w0: R.Tensor((640, 640), "float32"),
|
|
w1: R.Tensor((640, 640), "float32"),
|
|
w2: R.Tensor((640, 640), "float32"),
|
|
w3: R.Tensor((640, 640), "float32"),
|
|
w4: R.Tensor((640, 640), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x1, w0)
|
|
lv1 = R.matmul(x1, w1)
|
|
lv2 = R.matmul(x1, w2)
|
|
lv3 = R.matmul(x2, w3)
|
|
lv4 = R.matmul(x2, w4)
|
|
out = (lv0, lv1, lv2, lv3, lv4)
|
|
R.output(out)
|
|
return out
|
|
|
|
check = lambda *inp: len(inp[1]) > 2 # Ignore branches with two matmuls
|
|
after = CombineParallelMatmul(check)(tvm.IRModule.from_expr(before))["main"]
|
|
|
|
@R.function(private=True)
|
|
def expected(
|
|
x1: R.Tensor((2, 1024, 640), dtype="float32"),
|
|
x2: R.Tensor((2, 1024, 640), dtype="float32"),
|
|
w0: R.Tensor((640, 640), dtype="float32"),
|
|
w1: R.Tensor((640, 640), dtype="float32"),
|
|
w2: R.Tensor((640, 640), dtype="float32"),
|
|
w3: R.Tensor((640, 640), dtype="float32"),
|
|
w4: R.Tensor((640, 640), dtype="float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv = R.concat((w0, w1, w2), axis=1)
|
|
lv1 = R.matmul(x1, lv, out_dtype="float32")
|
|
lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=2)
|
|
lv0 = lv2[0]
|
|
lv1_1 = lv2[1]
|
|
lv2_1 = lv2[2]
|
|
lv3 = R.matmul(x2, w3, out_dtype=None)
|
|
lv4 = R.matmul(x2, w4, out_dtype=None)
|
|
out = (lv0, lv1_1, lv2_1, lv3, lv4)
|
|
R.output(out)
|
|
return out
|
|
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_combine_matmul_of_static_and_dynamic_shapes():
|
|
"""Combine two matmuls, one with dynamic shape
|
|
|
|
The `R.split` operator must have a static list of integer indices
|
|
at which to split the matmul output, because these integer indices
|
|
are stored as operator attributes. However, the last output can
|
|
still have a dynamic shape.
|
|
|
|
"""
|
|
|
|
@R.function(private=True)
|
|
def before(
|
|
x: R.Tensor((2, 1024, 640), "float32"),
|
|
w0: R.Tensor((640, 640), "float32"),
|
|
w1: R.Tensor((640, "M"), "float32"),
|
|
):
|
|
M = T.int64()
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x, w0)
|
|
lv1 = R.matmul(x, w1)
|
|
out = (lv0, lv1)
|
|
R.output(out)
|
|
return out
|
|
|
|
@R.function(private=True)
|
|
def expected(
|
|
x: R.Tensor((2, 1024, 640), dtype="float32"),
|
|
w0: R.Tensor((640, 640), dtype="float32"),
|
|
w1: R.Tensor((640, "M"), dtype="float32"),
|
|
) -> R.Tuple(
|
|
R.Tensor((2, 1024, 640), dtype="float32"), R.Tensor((2, 1024, "M"), dtype="float32")
|
|
):
|
|
M = T.int64()
|
|
with R.dataflow():
|
|
lv: R.Tensor((640, 640 + M), dtype="float32") = R.concat((w0, w1), axis=1)
|
|
lv1: R.Tensor((2, 1024, 640 + M), dtype="float32") = R.matmul(
|
|
x, lv, out_dtype="float32"
|
|
)
|
|
lv2: R.Tuple(
|
|
R.Tensor((2, 1024, 640), dtype="float32"),
|
|
R.Tensor((2, 1024, M), dtype="float32"),
|
|
) = R.split(lv1, indices_or_sections=[640], axis=2)
|
|
lv0: R.Tensor((2, 1024, 640), dtype="float32") = lv2[0]
|
|
lv1_1: R.Tensor((2, 1024, M), dtype="float32") = lv2[1]
|
|
out: R.Tuple(
|
|
R.Tensor((2, 1024, 640), dtype="float32"),
|
|
R.Tensor((2, 1024, M), dtype="float32"),
|
|
) = (lv0, lv1_1)
|
|
R.output(out)
|
|
return out
|
|
|
|
after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
|
|
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_combine_matmul_of_dynamic_and_static_shapes():
|
|
"""Combine two matmuls, one with dynamic shape
|
|
|
|
Like `test_combine_matmul_of_static_and_dynamic_shapes`, but the
|
|
dynamic-shaped matmul is encountered first. Due to the
|
|
requirements imposed by `R.split` storing the split indices as
|
|
static integers, the static-shaped weights must occur first in the
|
|
concatenated weights.
|
|
"""
|
|
|
|
@R.function(private=True)
|
|
def before(
|
|
x: R.Tensor((2, 1024, 640), "float32"),
|
|
w0: R.Tensor((640, "M"), "float32"),
|
|
w1: R.Tensor((640, 640), "float32"),
|
|
):
|
|
M = T.int64()
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x, w0)
|
|
lv1 = R.matmul(x, w1)
|
|
out = (lv0, lv1)
|
|
R.output(out)
|
|
return out
|
|
|
|
@R.function(private=True)
|
|
def expected(
|
|
x: R.Tensor((2, 1024, 640), dtype="float32"),
|
|
w0: R.Tensor((640, "M"), dtype="float32"),
|
|
w1: R.Tensor((640, 640), dtype="float32"),
|
|
) -> R.Tuple(
|
|
R.Tensor((2, 1024, "M"), dtype="float32"), R.Tensor((2, 1024, 640), dtype="float32")
|
|
):
|
|
M = T.int64()
|
|
with R.dataflow():
|
|
lv: R.Tensor((640, 640 + M), dtype="float32") = R.concat((w1, w0), axis=1)
|
|
lv1: R.Tensor((2, 1024, 640 + M), dtype="float32") = R.matmul(
|
|
x, lv, out_dtype="float32"
|
|
)
|
|
lv2: R.Tuple(
|
|
R.Tensor((2, 1024, 640), dtype="float32"),
|
|
R.Tensor((2, 1024, M), dtype="float32"),
|
|
) = R.split(lv1, indices_or_sections=[640], axis=2)
|
|
lv0: R.Tensor((2, 1024, M), dtype="float32") = lv2[1]
|
|
lv1_1: R.Tensor((2, 1024, 640), dtype="float32") = lv2[0]
|
|
out: R.Tuple(
|
|
R.Tensor((2, 1024, M), dtype="float32"),
|
|
R.Tensor((2, 1024, 640), dtype="float32"),
|
|
) = (lv0, lv1_1)
|
|
R.output(out)
|
|
return out
|
|
|
|
after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
|
|
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_limit_one_dynamic_shape_in_combined_matmul():
|
|
"""Combine two matmuls, one with dynamic shape
|
|
|
|
Like `test_combine_matmul_of_static_and_dynamic_shapes`, but with
|
|
two dynamic weights that could, in principle, be merged together.
|
|
Because `R.split` must have integer indices at which to split,
|
|
only one of the dynamic outputs can be part of the combined
|
|
matmul.
|
|
"""
|
|
|
|
@R.function(private=True)
|
|
def before(
|
|
x: R.Tensor((2, 1024, 640), "float32"),
|
|
w0: R.Tensor((640, "M"), "float32"),
|
|
w1: R.Tensor((640, 640), "float32"),
|
|
w2: R.Tensor((640, "N"), "float32"),
|
|
):
|
|
M = T.int64()
|
|
with R.dataflow():
|
|
lv0 = R.matmul(x, w0)
|
|
lv1 = R.matmul(x, w1)
|
|
lv2 = R.matmul(x, w2)
|
|
out = (lv0, lv1, lv2)
|
|
R.output(out)
|
|
return out
|
|
|
|
@R.function(private=True)
|
|
def expected(
|
|
x: R.Tensor((2, 1024, 640), dtype="float32"),
|
|
w0: R.Tensor((640, "M"), dtype="float32"),
|
|
w1: R.Tensor((640, 640), dtype="float32"),
|
|
w2: R.Tensor((640, "N"), "float32"),
|
|
) -> R.Tuple(
|
|
R.Tensor((2, 1024, "M"), dtype="float32"),
|
|
R.Tensor((2, 1024, 640), dtype="float32"),
|
|
R.Tensor((2, 1024, "N"), dtype="float32"),
|
|
):
|
|
M = T.int64()
|
|
with R.dataflow():
|
|
concat_weights = R.concat((w1, w0), axis=1)
|
|
concat_output = R.matmul(x, concat_weights, out_dtype="float32")
|
|
split_output: R.Tuple(
|
|
[R.Tensor([2, 1024, 640], dtype="float32"), R.Tensor([2, 1024, M], dtype="float32")]
|
|
) = R.split(concat_output, indices_or_sections=[640], axis=2)
|
|
lv0 = split_output[1]
|
|
lv1 = split_output[0]
|
|
lv2 = R.matmul(x, w2)
|
|
out = (lv0, lv1, lv2)
|
|
R.output(out)
|
|
return out
|
|
|
|
after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
|
|
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|