wehub-resource-sync
2525 lines
96 KiB
Python
2525 lines
96 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.
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax, topi
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def _check(mod_before, mod_expected):
|
|
mod_after = relax.transform.FuseTIR()(mod_before)
|
|
tvm.ir.assert_structural_equal(mod_expected, mod_after)
|
|
|
|
|
|
def test_simple():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
p0 = relax.Var("p0", R.Tensor([], "float32"))
|
|
|
|
with bb.function("fused_add_exp_squeeze", [x, p0], attrs={"Primitive": True}, private=True):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(topi.add, x, p0)
|
|
lv1 = bb.emit_te(topi.exp, lv0)
|
|
gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
|
|
bb.emit_func_output(gv)
|
|
fused_add_exp_squeeze = bb.get().get_global_var("fused_add_exp_squeeze")
|
|
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
p0 = relax.Var("p0", R.Tensor([], "float32"))
|
|
with bb.function("main", [x, p0]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(relax.Call(fused_add_exp_squeeze, [x, p0]))
|
|
bb.emit_func_output(gv)
|
|
|
|
return bb.get().with_attrs({"foo": "bar"})
|
|
|
|
def expected():
|
|
def fused_add_exp_squeeze(x, p0):
|
|
add = topi.add(x, p0)
|
|
exp = topi.exp(add)
|
|
squeeze = topi.squeeze(exp)
|
|
return squeeze
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
p0 = relax.Var("p0", R.Tensor([], "float32"))
|
|
with bb.function("main", [x, p0]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(bb.call_te(fused_add_exp_squeeze, x, p0))
|
|
bb.emit_func_output(gv)
|
|
return bb.get().with_attrs({"foo": "bar"})
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_conv2d_fuse():
|
|
def before(dtype):
|
|
bb = relax.BlockBuilder()
|
|
|
|
# Grouped function 1
|
|
x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype))
|
|
w = relax.Var("w", R.Tensor((16, 16, 3, 3), dtype))
|
|
p0 = relax.Var("p0", R.Tensor((), dtype))
|
|
with bb.function("fused_conv2d_add1_add2", [x, w, p0], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(
|
|
topi.nn.conv2d,
|
|
x,
|
|
w,
|
|
strides=1,
|
|
padding=1,
|
|
dilation=1,
|
|
primfunc_name_hint="conv2d",
|
|
)
|
|
lv1 = bb.emit_te(topi.add, p0, lv0, primfunc_name_hint="add1")
|
|
gv = bb.emit_output(bb.call_te(topi.add, lv0, lv1, primfunc_name_hint="add2"))
|
|
bb.emit_func_output(gv)
|
|
|
|
# Grouped function 2
|
|
x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype))
|
|
w = relax.Var("w", R.Tensor((16, 16, 1, 1), dtype))
|
|
y = relax.Var("y", R.Tensor((1, 16, 64, 64), dtype))
|
|
with bb.function("fused_conv2d1_add2", [x, w, y], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(
|
|
topi.nn.conv2d,
|
|
x,
|
|
w,
|
|
strides=1,
|
|
padding=0,
|
|
dilation=1,
|
|
primfunc_name_hint="conv2d1",
|
|
)
|
|
gv = bb.emit_output(bb.call_te(topi.add, lv0, y, primfunc_name_hint="add2"))
|
|
bb.emit_func_output(gv)
|
|
|
|
# Get the global variables of the grouped functions
|
|
mod = bb.get()
|
|
fused_conv2d_add1_add2 = mod.get_global_var("fused_conv2d_add1_add2")
|
|
fused_conv2d1_add2 = mod.get_global_var("fused_conv2d1_add2")
|
|
|
|
# Main function
|
|
x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype))
|
|
w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), dtype))
|
|
w2 = relax.Var("w2", R.Tensor((16, 16, 1, 1), dtype))
|
|
w3 = relax.Var("w3", R.Tensor((16, 16, 3, 3), dtype))
|
|
with bb.function("main", [x, w1, w2, w3]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(topi.add, x, relax.const(1, dtype))
|
|
lv1 = bb.emit(relax.Call(fused_conv2d_add1_add2, [lv0, w1, relax.const(1, dtype)]))
|
|
lv2 = bb.emit_te(
|
|
topi.nn.conv2d,
|
|
lv1,
|
|
w3,
|
|
strides=1,
|
|
padding=1,
|
|
dilation=1,
|
|
)
|
|
gv = bb.emit_output(relax.Call(fused_conv2d1_add2, [lv1, w2, lv2]))
|
|
bb.emit_func_output(gv)
|
|
|
|
return bb.get()
|
|
|
|
def expected(dtype):
|
|
def fused_conv2d_add1_add2(x, w, p):
|
|
conv = topi.nn.conv2d(x, w, strides=1, padding=1, dilation=1)
|
|
add = topi.add(p, conv)
|
|
return topi.add(conv, add)
|
|
|
|
def fused_conv2d1_add2(x, w, p):
|
|
conv = topi.nn.conv2d(x, w, strides=1, padding=0, dilation=1)
|
|
return topi.add(conv, p)
|
|
|
|
bb = relax.BlockBuilder()
|
|
|
|
# Main function
|
|
x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype))
|
|
w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), dtype))
|
|
w2 = relax.Var("w2", R.Tensor((16, 16, 1, 1), dtype))
|
|
w3 = relax.Var("w3", R.Tensor((16, 16, 3, 3), dtype))
|
|
with bb.function("main", [x, w1, w2, w3]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(topi.add, x, relax.const(1, dtype))
|
|
lv1 = bb.emit_te(fused_conv2d_add1_add2, lv0, w1, relax.const(1, dtype))
|
|
lv2 = bb.emit_te(
|
|
topi.nn.conv2d,
|
|
lv1,
|
|
w3,
|
|
strides=1,
|
|
padding=1,
|
|
dilation=1,
|
|
)
|
|
gv = bb.emit_output(bb.call_te(fused_conv2d1_add2, lv1, w2, lv2))
|
|
bb.emit_func_output(gv)
|
|
|
|
return bb.get()
|
|
|
|
_check(before("float32"), expected("float32"))
|
|
|
|
|
|
def test_two_subfunction():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x1 = relax.Var("x1", R.Tensor([10, 20], "float32"))
|
|
with bb.function("fused_exp_squeeze", [x1], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
lv1 = bb.emit_te(topi.exp, x1)
|
|
gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
|
|
bb.emit_func_output(gv)
|
|
mod = bb.get()
|
|
|
|
func_gv = mod.get_global_var("fused_exp_squeeze")
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv = bb.emit(relax.Call(func_gv, [x]))
|
|
lv2 = bb.emit(relax.Call(func_gv, [lv]))
|
|
gv = bb.emit_output(lv2)
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
def expected():
|
|
def fused_exp_squeeze(x):
|
|
exp = topi.exp(x)
|
|
squeeze = topi.squeeze(exp)
|
|
return squeeze
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv = bb.emit_te(fused_exp_squeeze, x)
|
|
lv2 = bb.call_te(fused_exp_squeeze, lv)
|
|
gv = bb.emit_output(lv2)
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_fuse_same_primfunc():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x1 = relax.Var("x1", R.Tensor([10, 20], "float32"))
|
|
with bb.function("fused_exp_exp_squeeze", [x1], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
lv1 = bb.emit_te(topi.exp, x1)
|
|
lv2 = bb.emit_te(topi.exp, lv1)
|
|
gv = bb.emit_output(bb.call_te(topi.squeeze, lv2))
|
|
bb.emit_func_output(gv)
|
|
mod = bb.get()
|
|
|
|
func_gv = mod.get_global_var("fused_exp_exp_squeeze")
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv = bb.emit(relax.Call(func_gv, [x]))
|
|
gv = bb.emit_output(lv)
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
def expected():
|
|
def fused_exp_exp_squeeze(x):
|
|
exp = topi.exp(x)
|
|
exp = topi.exp(exp)
|
|
squeeze = topi.squeeze(exp)
|
|
return squeeze
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv = bb.call_te(fused_exp_exp_squeeze, x)
|
|
gv = bb.emit_output(lv)
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_fuse_with_tuple_as_param():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tuple([R.Tensor([10], "float32"), R.Tensor([10], "float32")]))
|
|
with bb.function("fused_exp_add", [x], attrs={"Primitive": True}, private=True):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(relax.TupleGetItem(x, 0))
|
|
lv1 = bb.emit(relax.TupleGetItem(x, 1))
|
|
lv2 = bb.emit_te(topi.exp, lv0)
|
|
gv = bb.emit_output(bb.call_te(topi.add, lv2, lv1))
|
|
bb.emit_func_output(gv)
|
|
mod = bb.get()
|
|
|
|
func_gv = mod.get_global_var("fused_exp_add")
|
|
x = relax.Var("x", R.Tuple([R.Tensor([10], "float32"), R.Tensor([10], "float32")]))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(relax.Call(func_gv, [x]))
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
def expected():
|
|
def fused_exp_add(x1, x2):
|
|
exp = topi.exp(x1)
|
|
return topi.add(exp, x2)
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tuple([R.Tensor([10], "float32"), R.Tensor([10], "float32")]))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(relax.TupleGetItem(x, 0))
|
|
lv1 = bb.emit(relax.TupleGetItem(x, 1))
|
|
gv = bb.emit_output(bb.call_te(fused_exp_add, lv0, lv1))
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_fuse_with_nested_tuple_as_param():
|
|
tuple_ty = R.Tuple(
|
|
[R.Tensor([10], "float32"), R.Tuple([R.Tensor([10], "float32"), R.Tensor([10], "float32")])]
|
|
)
|
|
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", tuple_ty)
|
|
with bb.function("fused_exp_add_add", [x], attrs={"Primitive": True}, private=True):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(relax.TupleGetItem(x, 0))
|
|
lv0_exp = bb.emit_te(topi.exp, lv0)
|
|
lv1 = bb.emit(relax.TupleGetItem(x, 1))
|
|
lv1_0 = bb.emit(relax.TupleGetItem(lv1, 0))
|
|
lv1_1 = bb.emit(relax.TupleGetItem(lv1, 1))
|
|
lv2 = bb.emit_te(topi.add, lv1_0, lv1_1)
|
|
gv = bb.emit_output(bb.call_te(topi.add, lv0_exp, lv2))
|
|
bb.emit_func_output(gv)
|
|
mod = bb.get()
|
|
|
|
func_gv = mod.get_global_var("fused_exp_add_add")
|
|
x = relax.Var("x", tuple_ty)
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(relax.Call(func_gv, [x]))
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
def expected():
|
|
def fused_exp_add_add(x1, x2, x3):
|
|
exp = topi.exp(x1)
|
|
add = topi.add(x2, x3)
|
|
return topi.add(exp, add)
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", tuple_ty)
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(relax.TupleGetItem(x, 0))
|
|
lv1 = bb.emit(relax.TupleGetItem(x, 1))
|
|
lv2 = bb.emit(relax.TupleGetItem(lv1, 0))
|
|
lv3 = bb.emit(relax.TupleGetItem(lv1, 1))
|
|
gv = bb.emit_output(bb.call_te(fused_exp_add_add, lv0, lv2, lv3))
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_fuse_with_call_tir_in_main():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x1 = relax.Var("x1", R.Tensor([10, 20], "float32"))
|
|
with bb.function("fused_exp_squeeze", [x1], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
lv = bb.emit_te(topi.exp, x1)
|
|
gv = bb.emit_output(bb.call_te(topi.squeeze, lv))
|
|
bb.emit_func_output(gv)
|
|
mod = bb.get()
|
|
|
|
func_gv = mod.get_global_var("fused_exp_squeeze")
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(relax.Call(func_gv, [x]))
|
|
lv1 = bb.emit_te(topi.add, lv0, relax.const(1, "float32"))
|
|
gv = bb.emit_output(lv1)
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
def expected():
|
|
def fused_exp_squeeze(x):
|
|
exp = topi.exp(x)
|
|
squeeze = topi.squeeze(exp)
|
|
return squeeze
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv = bb.emit_te(fused_exp_squeeze, x)
|
|
lv2 = bb.call_te(topi.add, lv, relax.const(1, "float32"))
|
|
gv = bb.emit_output(lv2)
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_fuse_with_const_in_argument():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x1 = relax.Var("x1", R.Tensor([10, 20], "float32"))
|
|
x2 = relax.Var("x2", R.Tensor([], "float32"))
|
|
with bb.function("fused_add_exp_squeeze", [x1, x2], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(topi.add, x1, x2)
|
|
lv1 = bb.emit_te(topi.exp, lv0)
|
|
gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
|
|
bb.emit_func_output(gv)
|
|
mod = bb.get()
|
|
|
|
func_gv = mod.get_global_var("fused_add_exp_squeeze")
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv = bb.emit(relax.Call(func_gv, [x, relax.const(1, "float32")]))
|
|
gv = bb.emit_output(lv)
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
def expected():
|
|
def fused_add_exp_squeeze(x, y):
|
|
add = topi.add(x, y)
|
|
exp = topi.exp(add)
|
|
squeeze = topi.squeeze(exp)
|
|
return squeeze
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
lv = bb.call_te(fused_add_exp_squeeze, x, relax.const(1, "float32"))
|
|
gv = bb.emit_output(lv)
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_fuse_tuple_output():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
p0 = relax.Var("p0", R.Tensor([], "float32"))
|
|
|
|
with bb.function("fused_add_exp", [x, p0], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
gv0 = bb.emit_output(bb.call_te(topi.add, x, p0))
|
|
gv1 = bb.emit_output(bb.call_te(topi.exp, gv0))
|
|
bb.emit_func_output(relax.Tuple([gv0, gv1]))
|
|
fused_add_exp = bb.get().get_global_var("fused_add_exp")
|
|
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
p0 = relax.Var("p0", R.Tensor([], "float32"))
|
|
with bb.function("main", [x, p0]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(relax.Call(fused_add_exp, [x, p0]))
|
|
bb.emit_func_output(gv)
|
|
|
|
return bb.get()
|
|
|
|
def expected():
|
|
def fused_add_exp(x, p0):
|
|
add = topi.add(x, p0)
|
|
exp = topi.exp(add)
|
|
return add, exp
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
p0 = relax.Var("p0", R.Tensor([], "float32"))
|
|
with bb.function("main", [x, p0]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(bb.call_te(fused_add_exp, x, p0))
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_fuse_with_immediate_tuple():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
y = relax.Var("y", R.Tensor([10, 20], "float32"))
|
|
|
|
with bb.function("fused_add", [x, y], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
lv_tuple = bb.emit(relax.Tuple([x, relax.Tuple([x, y])]))
|
|
lv_x = bb.emit(relax.TupleGetItem(lv_tuple, 0))
|
|
lv0 = bb.emit(relax.TupleGetItem(lv_tuple, 1))
|
|
lv_y = bb.emit(relax.TupleGetItem(lv0, 1))
|
|
gv = bb.emit_output(bb.call_te(topi.add, lv_x, lv_y))
|
|
bb.emit_func_output(gv)
|
|
fused_add = bb.get().get_global_var("fused_add")
|
|
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
y = relax.Var("y", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x, y]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(relax.Call(fused_add, [x, y]))
|
|
bb.emit_func_output(gv)
|
|
|
|
return bb.get()
|
|
|
|
def expected():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
y = relax.Var("y", R.Tensor([10, 20], "float32"))
|
|
with bb.function("main", [x, y]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(bb.call_te(topi.add, x, y, primfunc_name_hint="fused_add"))
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_fuse_return_partial_result():
|
|
def te_argmax_idx_val(val):
|
|
from tvm import te
|
|
|
|
def f_combine(x, y):
|
|
lhs = tvm.tirx.Select((x[1] >= y[1]), x[0], y[0])
|
|
rhs = tvm.tirx.Select((x[1] >= y[1]), x[1], y[1])
|
|
return lhs, rhs
|
|
|
|
def f_identity(dtype0: tvm.DataType, dtype1: tvm.DataType):
|
|
return tvm.tirx.const(-1, dtype0), tvm.te.min_value(dtype1)
|
|
|
|
argmax = te.comm_reducer(f_combine, f_identity, name="argmax")
|
|
m, n = val.shape
|
|
k = te.reduce_axis((0, n), "k")
|
|
max_idx, max_val = te.compute(
|
|
(m,), lambda i: argmax((k.var, val[i, k]), axis=k), name="argmax"
|
|
)
|
|
return max_idx, max_val
|
|
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
offset = relax.Var("offset", R.Tensor([10], "int32"))
|
|
with bb.function("fused_argmax_add", [x, offset], attrs={"Primitive": True}):
|
|
with bb.dataflow():
|
|
lv = bb.emit_te(te_argmax_idx_val, x)
|
|
idx = bb.emit(relax.TupleGetItem(lv, 0))
|
|
gv = bb.emit_output(bb.call_te(topi.add, idx, offset))
|
|
bb.emit_func_output(gv)
|
|
mod = bb.get()
|
|
|
|
func_gv = mod.get_global_var("fused_argmax_add")
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
offset = relax.Var("x", R.Tensor([10], "int32"))
|
|
with bb.function("main", [x, offset]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(relax.Call(func_gv, [x, offset]))
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
def expected():
|
|
def fused_argmax_add(x, offset):
|
|
idx, value = te_argmax_idx_val(x)
|
|
idx = topi.add(idx, offset)
|
|
return idx
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
offset = relax.Var("offset", R.Tensor([10], "int32"))
|
|
with bb.function("main", [x, offset]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(bb.call_te(fused_argmax_add, x, offset))
|
|
bb.emit_func_output(gv)
|
|
return bb.get()
|
|
|
|
_check(before(), expected())
|
|
|
|
|
|
def test_multiple_relax_functions():
|
|
def before():
|
|
bb = relax.BlockBuilder()
|
|
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
p0 = relax.Var("p0", R.Tensor((), "float32"))
|
|
with bb.function("fused_add_exp_squeeze", [x, p0], attrs={"Primitive": True}, private=True):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(topi.add, x, p0)
|
|
lv1 = bb.emit_te(topi.exp, lv0)
|
|
gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
|
|
bb.emit_func_output(gv)
|
|
fused_add_exp_squeeze = bb.get().get_global_var("fused_add_exp_squeeze")
|
|
|
|
x = relax.Var("x", R.Tensor([20, 10], "float32"))
|
|
p0 = relax.Var("p0", R.Tensor((), "float32"))
|
|
with bb.function(
|
|
"fused_add1_exp1_squeeze1", [x, p0], attrs={"Primitive": True}, private=True
|
|
):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit_te(topi.add, x, p0)
|
|
lv1 = bb.emit_te(topi.exp, lv0)
|
|
gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
|
|
bb.emit_func_output(gv)
|
|
fused_add1_exp1_squeeze1 = bb.get().get_global_var("fused_add1_exp1_squeeze1")
|
|
|
|
x = relax.Var("x", R.Tensor([10, 20], "float32"))
|
|
with bb.function("func1", [x]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(
|
|
relax.Call(fused_add_exp_squeeze, [x, relax.const(1, "float32")])
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
x = relax.Var("x", R.Tensor([20, 10], "float32"))
|
|
with bb.function("func2", [x]):
|
|
with bb.dataflow():
|
|
gv = bb.emit_output(
|
|
relax.Call(fused_add1_exp1_squeeze1, [x, relax.const(1, "float32")])
|
|
)
|
|
bb.emit_func_output(gv)
|
|
|
|
return bb.get()
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def func1(x: R.Tensor((10, 20), dtype="float32")) -> R.Tensor((10, 20), dtype="float32"):
|
|
with R.dataflow():
|
|
gv2 = R.call_tir(
|
|
Expected.fused_add_exp_squeeze,
|
|
(x, R.const(1, "float32")),
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
@R.function
|
|
def func2(x: R.Tensor((20, 10), dtype="float32")) -> R.Tensor((20, 10), dtype="float32"):
|
|
with R.dataflow():
|
|
gv3 = R.call_tir(
|
|
Expected.fused_add1_exp1_squeeze1,
|
|
(x, R.const(1, "float32")),
|
|
out_ty=R.Tensor((20, 10), dtype="float32"),
|
|
)
|
|
R.output(gv3)
|
|
return gv3
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_add1_exp1_squeeze1(
|
|
x: T.Buffer((T.int64(20), T.int64(10)), "float32"),
|
|
p0: T.Buffer((), "float32"),
|
|
T_squeeze: T.Buffer((T.int64(20), T.int64(10)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
T_add = T.sblock_alloc_buffer((T.int64(20), T.int64(10)))
|
|
compute = T.sblock_alloc_buffer((T.int64(20), T.int64(10)))
|
|
for ax0, ax1 in T.grid(T.int64(20), T.int64(10)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(x[v_ax0, v_ax1], p0[()])
|
|
T.writes(T_add[v_ax0, v_ax1])
|
|
T_add[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
|
|
for i0, i1 in T.grid(T.int64(20), T.int64(10)):
|
|
with T.sblock("compute"):
|
|
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(T_add[v_i0, v_i1])
|
|
T.writes(compute[v_i0, v_i1])
|
|
compute[v_i0, v_i1] = T.exp(T_add[v_i0, v_i1])
|
|
for ax0, ax1 in T.grid(T.int64(20), T.int64(10)):
|
|
with T.sblock("T_squeeze"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(compute[v_ax0, v_ax1])
|
|
T.writes(T_squeeze[v_ax0, v_ax1])
|
|
T_squeeze[v_ax0, v_ax1] = compute[v_ax0, v_ax1]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_add_exp_squeeze(
|
|
x: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
p0: T.Buffer((), "float32"),
|
|
T_squeeze: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
T_add = T.sblock_alloc_buffer((T.int64(10), T.int64(20)))
|
|
compute = T.sblock_alloc_buffer((T.int64(10), T.int64(20)))
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(x[v_ax0, v_ax1], p0[()])
|
|
T.writes(T_add[v_ax0, v_ax1])
|
|
T_add[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
|
|
for i0, i1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("compute"):
|
|
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(T_add[v_i0, v_i1])
|
|
T.writes(compute[v_i0, v_i1])
|
|
compute[v_i0, v_i1] = T.exp(T_add[v_i0, v_i1])
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_squeeze"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(compute[v_ax0, v_ax1])
|
|
T.writes(T_squeeze[v_ax0, v_ax1])
|
|
T_squeeze[v_ax0, v_ax1] = compute[v_ax0, v_ax1]
|
|
|
|
_check(before(), Expected)
|
|
|
|
|
|
def test_skip_call_dps_packed():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), "float32")):
|
|
with R.dataflow():
|
|
y = R.call_dps_packed("func_packed_dps", x, R.Tensor((2, 3), "float32"))
|
|
R.output(y)
|
|
return y
|
|
|
|
# FuseTIR should do no change to it.
|
|
_check(Module, Module)
|
|
|
|
|
|
def test_symbolic_shape_aware_fuse():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def fused_add_exp_squeeze(
|
|
x: R.Tensor(["n", "m"], "float32"), p0: R.Tensor([], "float32")
|
|
) -> R.Tensor(["n", "m"], dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
with R.dataflow():
|
|
lv0 = R.emit_te(topi.add, x, p0)
|
|
lv1 = R.emit_te(topi.exp, lv0)
|
|
gv = R.emit_te(topi.squeeze, lv1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(x: R.Tensor(["n", "m"], "float32")) -> R.Tensor(["n", "m"], dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused_add_exp_squeeze(x, R.const(1, "float32"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
def fused_add_exp_squeeze(x, p0):
|
|
return topi.squeeze(topi.exp(topi.add(x, p0)))
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(["n", "m"], "float32")) -> R.Tensor(["n", "m"], dtype="float32"):
|
|
with R.dataflow():
|
|
gv = R.emit_te(fused_add_exp_squeeze, x, R.const(1, "float32"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_fuse_of_dynamic_kernel_with_var_params_and_static_args():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def dynamic_tir_kernel(a: T.handle, b: T.handle):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
A = T.match_buffer(a, [m, n], "float32")
|
|
B = T.match_buffer(b, [m, n], "float32")
|
|
|
|
for iters in T.grid(m, n):
|
|
with T.sblock("compute"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
B[i, j] = A[i, j] * i + j
|
|
|
|
@R.function(private=True)
|
|
def fused_function(x: R.Tensor([16, 32], "float32")) -> R.Tensor([16, 32], dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
y = R.call_tir(cls.dynamic_tir_kernel, [x], out_ty=R.Tensor([16, 32], "float32"))
|
|
z = R.call_tir(cls.dynamic_tir_kernel, [y], out_ty=R.Tensor([16, 32], "float32"))
|
|
R.output(z)
|
|
return z
|
|
|
|
@R.function
|
|
def main(x: R.Tensor([16, 32], "float32")) -> R.Tensor([16, 32], dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused_function(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_function(
|
|
X: T.Buffer([T.int64(16), T.int64(32)], "float32"),
|
|
Z: T.Buffer([T.int64(16), T.int64(32)], "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
Y = T.sblock_alloc_buffer(X.shape, "float32")
|
|
for iters in T.grid(*X.shape):
|
|
with T.sblock("compute_Y"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
Y[i, j] = X[i, j] * i + j
|
|
|
|
for iters in T.grid(*X.shape):
|
|
with T.sblock("compute_Z"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
Z[i, j] = Y[i, j] * i + j
|
|
|
|
@R.function
|
|
def main(x: R.Tensor([16, 32], "float32")) -> R.Tensor([16, 32], dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(cls.fused_function, [x], out_ty=R.Tensor([16, 32], "float32"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_fuse_of_dynamic_kernel_with_expression_params_and_static_args():
|
|
"""Parameters and arguments do not need to match structurally
|
|
|
|
Here, the kernel requires arguments (m*n), and is provided
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def dynamic_tir_kernel(a: T.handle, b: T.handle, c: T.handle, d: T.handle):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
A = T.match_buffer(a, [m * n], "float32")
|
|
B = T.match_buffer(b, [m], "float32")
|
|
C = T.match_buffer(c, [n], "float32")
|
|
D = T.match_buffer(d, [m * n], "float32")
|
|
|
|
for i, j in T.grid(m, n):
|
|
with T.sblock("compute"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
D[vi * 32 + vj] = A[vi * 32 + vj] * B[vi] + C[vj]
|
|
|
|
@R.function(private=True)
|
|
def fused_function(
|
|
x: R.Tensor([16 * 32], "float32"),
|
|
B: R.Tensor([16], "float32"),
|
|
C: R.Tensor([32], "float32"),
|
|
):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
y = R.call_tir(
|
|
cls.dynamic_tir_kernel, [x, B, C], out_ty=R.Tensor([16 * 32], "float32")
|
|
)
|
|
z = R.call_tir(
|
|
cls.dynamic_tir_kernel, [y, B, C], out_ty=R.Tensor([16 * 32], "float32")
|
|
)
|
|
R.output(z)
|
|
return z
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([16 * 32], "float32"),
|
|
B: R.Tensor([16], "float32"),
|
|
C: R.Tensor([32], "float32"),
|
|
):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused_function(x, B, C)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_function(
|
|
X: T.Buffer(T.int64(512), "float32"),
|
|
B: T.Buffer(T.int64(16), "float32"),
|
|
C: T.Buffer(T.int64(32), "float32"),
|
|
Z: T.Buffer(T.int64(512), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
Y = T.sblock_alloc_buffer((T.int64(512),))
|
|
for i, j in T.grid(T.int64(16), T.int64(32)):
|
|
with T.sblock("compute"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
Y[vi * 32 + vj] = X[vi * 32 + vj] * B[vi] + C[vj]
|
|
|
|
for i, j in T.grid(T.int64(16), T.int64(32)):
|
|
with T.sblock("compute_1"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
Z[vi * 32 + vj] = Y[vi * 32 + vj] * B[vi] + C[vj]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((512,), dtype="float32"),
|
|
B: R.Tensor((16,), dtype="float32"),
|
|
C: R.Tensor((32,), dtype="float32"),
|
|
) -> R.Tensor((512,), dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(
|
|
cls.fused_function, (x, B, C), out_ty=R.Tensor((512,), dtype="float32")
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_symbolic_shape_aware_fuse_with_allocation():
|
|
def te_mean(x, axis):
|
|
return topi.divide(topi.sum(x, axis, keepdims=True), 4096)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def fused_mean_add_tir_sqrt_divide_multiply(
|
|
x: R.Tensor((1, "n", 4096), dtype="float32"),
|
|
y: R.Tensor((1, "n", 4096), dtype="float32"),
|
|
rms_norm_weight: R.Tensor((4096,), dtype="float32"),
|
|
) -> R.Tensor((1, "n", 4096), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
with R.dataflow():
|
|
lv0 = R.emit_te(te_mean, x, axis=2)
|
|
lv1 = R.emit_te(topi.add, lv0, lv0)
|
|
lv2 = R.emit_te(topi.sqrt, lv1)
|
|
lv3 = R.emit_te(topi.divide, y, lv2)
|
|
gv = R.emit_te(topi.multiply, rms_norm_weight, lv3)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, "n", 4096), dtype="float32"),
|
|
y: R.Tensor((1, "n", 4096), dtype="float32"),
|
|
rms_norm_weight: R.Tensor((4096,), dtype="float32"),
|
|
) -> R.Tensor((1, "n", 4096), dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused_mean_add_tir_sqrt_divide_multiply(x, y, rms_norm_weight)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
def fused_mean_add_tir_sqrt_divide_multiply(x, y, rms_norm_weight):
|
|
lv0 = te_mean(x, axis=2)
|
|
lv1 = topi.add(lv0, lv0)
|
|
lv2 = topi.sqrt(lv1)
|
|
lv3 = topi.divide(y, lv2)
|
|
return topi.multiply(rms_norm_weight, lv3)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, "n", 4096), dtype="float32"),
|
|
y: R.Tensor((1, "n", 4096), dtype="float32"),
|
|
rms_norm_weight: R.Tensor((4096,), dtype="float32"),
|
|
) -> R.Tensor((1, "n", 4096), dtype="float32"):
|
|
with R.dataflow():
|
|
gv = R.emit_te(fused_mean_add_tir_sqrt_divide_multiply, x, y, rms_norm_weight)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_symbolic_var_in_call_tir_args():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def foo(
|
|
X: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)), "float32"),
|
|
Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
|
|
rotary: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)), "float32"),
|
|
m: T.int64,
|
|
):
|
|
for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(32), T.int64(128)):
|
|
with T.sblock("rotary"):
|
|
v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
rotary[v0, v1, v2, v3] = Y[m + v1 - 1, v3] * X[v0, v1, v2, v3]
|
|
|
|
@R.function
|
|
def fused(
|
|
x: R.Tensor((1, 1, 32, 128), dtype="float32"),
|
|
y: R.Tensor((2048, 128), dtype="float32"),
|
|
len: R.Shape(["m"]),
|
|
) -> R.Tensor((1, 1, 32, 128), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
m = T.int64()
|
|
cls = Before
|
|
with R.dataflow():
|
|
lv1 = R.emit_te(topi.add, x, x)
|
|
gv = R.call_tir(
|
|
cls.foo,
|
|
[lv1, y],
|
|
out_ty=R.Tensor((1, 1, 32, 128), dtype="float32"),
|
|
tir_vars=R.shape([m]),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 1, 32, 128), dtype="float32"),
|
|
y: R.Tensor((2048, 128), dtype="float32"),
|
|
len: R.Shape(["m"]),
|
|
) -> R.Tensor((1, 1, 32, 128), dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused(x, y, len)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused(
|
|
X: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)), "float32"),
|
|
Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
|
|
rotary: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)), "float32"),
|
|
m: T.int64,
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
T_add = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)))
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(1), T.int64(32), T.int64(128)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T_add[v_ax0, v_ax1, v_ax2, v_ax3] = (
|
|
X[v_ax0, v_ax1, v_ax2, v_ax3] + X[v_ax0, v_ax1, v_ax2, v_ax3]
|
|
)
|
|
for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(32), T.int64(128)):
|
|
with T.sblock("rotary"):
|
|
v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
rotary[v0, v1, v2, v3] = Y[m + v1 - T.int64(1), v3] * T_add[v0, v1, v2, v3]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 1, 32, 128), dtype="float32"),
|
|
y: R.Tensor((2048, 128), dtype="float32"),
|
|
len: R.Shape(["m"]),
|
|
) -> R.Tensor((1, 1, 32, 128), dtype="float32"):
|
|
m = T.int64()
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(
|
|
cls.fused,
|
|
(x, y),
|
|
out_ty=R.Tensor([1, 1, 32, 128], "float32"),
|
|
tir_vars=R.shape([m]),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_same_buffer_multiple_read():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def concatenate(
|
|
rxplaceholder: T.Buffer((T.int64(1), T.int64(4), T.int64(64), T.int64(64)), "float32"),
|
|
rxplaceholder_1: T.Buffer(
|
|
(T.int64(1), T.int64(4), T.int64(64), T.int64(64)), "float32"
|
|
),
|
|
T_concat: T.Buffer((T.int64(2), T.int64(4), T.int64(64), T.int64(64)), "float32"),
|
|
):
|
|
T.func_attr({"op_pattern": 2, "tirx.noalias": True})
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(4), T.int64(64), T.int64(64)):
|
|
with T.sblock("T_concat"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(
|
|
rxplaceholder_1[v_ax0 - T.int64(1), v_ax1, v_ax2, v_ax3],
|
|
rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3],
|
|
)
|
|
T.writes(T_concat[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_concat[v_ax0, v_ax1, v_ax2, v_ax3] = T.if_then_else(
|
|
T.int64(1) <= v_ax0,
|
|
rxplaceholder_1[v_ax0 - T.int64(1), v_ax1, v_ax2, v_ax3],
|
|
rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3],
|
|
)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def transpose2(
|
|
rxplaceholder: T.Buffer((T.int64(2), T.int64(4), T.int64(64), T.int64(64)), "float32"),
|
|
T_transpose: T.Buffer((T.int64(2), T.int64(64), T.int64(64), T.int64(4)), "float32"),
|
|
):
|
|
T.func_attr({"op_pattern": 2, "tirx.noalias": True})
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(64), T.int64(64), T.int64(4)):
|
|
with T.sblock("T_transpose"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(rxplaceholder[v_ax0, v_ax3, v_ax1, v_ax2])
|
|
T.writes(T_transpose[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_transpose[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[
|
|
v_ax0, v_ax3, v_ax1, v_ax2
|
|
]
|
|
|
|
@R.function
|
|
def fused_concatenate_transpose2(
|
|
inp_0: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
) -> R.Tensor((2, 64, 64, 4), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv = R.call_tir(
|
|
cls.concatenate,
|
|
(inp_0, inp_0),
|
|
out_ty=R.Tensor((2, 4, 64, 64), dtype="float32"),
|
|
)
|
|
gv = R.call_tir(
|
|
cls.transpose2, (lv,), out_ty=R.Tensor((2, 64, 64, 4), dtype="float32")
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 4, 64, 64), dtype="float32")) -> R.Tensor(
|
|
(2, 64, 64, 4), dtype="float32"
|
|
):
|
|
R.func_attr({"num_input": 3})
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv = cls.fused_concatenate_transpose2(inp_0)
|
|
R.output(lv)
|
|
return lv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_concatenate_transpose2(
|
|
inp_0: T.Buffer((T.int64(1), T.int64(4), T.int64(64), T.int64(64)), "float32"),
|
|
T_transpose_handle_intermediate: T.Buffer(
|
|
(T.int64(2), T.int64(64), T.int64(64), T.int64(4)), "float32"
|
|
),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
T_concat_handle_intermediate = T.sblock_alloc_buffer(
|
|
(T.int64(2), T.int64(4), T.int64(64), T.int64(64))
|
|
)
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(4), T.int64(64), T.int64(64)):
|
|
with T.sblock("T_concat"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(inp_0[v_ax0 - T.int64(1), v_ax1, v_ax2, v_ax3])
|
|
T.writes(T_concat_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_concat_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = T.if_then_else(
|
|
T.int64(1) <= v_ax0,
|
|
inp_0[v_ax0 - T.int64(1), v_ax1, v_ax2, v_ax3],
|
|
inp_0[v_ax0, v_ax1, v_ax2, v_ax3],
|
|
)
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(64), T.int64(64), T.int64(4)):
|
|
with T.sblock("T_transpose"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(T_concat_handle_intermediate[v_ax0, v_ax3, v_ax1, v_ax2])
|
|
T.writes(T_transpose_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_transpose_handle_intermediate[v_ax0, v_ax1, v_ax2, v_ax3] = (
|
|
T_concat_handle_intermediate[v_ax0, v_ax3, v_ax1, v_ax2]
|
|
)
|
|
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 4, 64, 64), dtype="float32")) -> R.Tensor(
|
|
(2, 64, 64, 4), dtype="float32"
|
|
):
|
|
R.func_attr({"num_input": 3})
|
|
cls = Expected
|
|
with R.dataflow():
|
|
lv = R.call_tir(
|
|
cls.fused_concatenate_transpose2,
|
|
(inp_0,),
|
|
out_ty=R.Tensor((2, 64, 64, 4), dtype="float32"),
|
|
)
|
|
R.output(lv)
|
|
return lv
|
|
|
|
_check(Module, Expected)
|
|
|
|
|
|
def test_tir_expression_in_shape():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def fused_transpose_matmul(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
y: R.Tensor(("n - 1", 4), dtype="float32"),
|
|
tir_vars: R.Shape(["n"]),
|
|
) -> R.Tensor(("n - 1", 3), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
with R.dataflow():
|
|
lv = R.emit_te(topi.transpose, x)
|
|
gv = R.emit_te(topi.matmul, y, lv)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
y: R.Tensor(("n - 1", 4), dtype="float32"),
|
|
tir_vars: R.Shape(["n"]),
|
|
) -> R.Tensor(("n - 1", 3), dtype="float32"):
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv = cls.fused_transpose_matmul(x, y, tir_vars)
|
|
R.output(lv)
|
|
return lv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_transpose_matmul(
|
|
x: T.Buffer((T.int64(3), T.int64(4)), "float32"),
|
|
p_y: T.handle,
|
|
p_output0: T.handle,
|
|
n: T.int64,
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
y = T.match_buffer(p_y, (n - T.int64(1), T.int64(4)))
|
|
var_T_matmul_intermediate = T.match_buffer(p_output0, (n - T.int64(1), T.int64(3)))
|
|
var_T_transpose_intermediate = T.sblock_alloc_buffer((T.int64(4), T.int64(3)))
|
|
for ax0, ax1 in T.grid(T.int64(4), T.int64(3)):
|
|
with T.sblock("T_transpose"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
var_T_transpose_intermediate[v_ax0, v_ax1] = x[v_ax1, v_ax0]
|
|
for ax0, ax1, k in T.grid(n - T.int64(1), T.int64(3), T.int64(4)):
|
|
with T.sblock("T_matmul"):
|
|
v_ax0, v_ax1, v_k = T.axis.remap("SSR", [ax0, ax1, k])
|
|
with T.init():
|
|
var_T_matmul_intermediate[v_ax0, v_ax1] = T.float32(0)
|
|
var_T_matmul_intermediate[v_ax0, v_ax1] = (
|
|
var_T_matmul_intermediate[v_ax0, v_ax1]
|
|
+ y[v_ax0, v_k] * var_T_transpose_intermediate[v_k, v_ax1]
|
|
)
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
y: R.Tensor(("n - 1", 4), dtype="float32"),
|
|
tir_vars: R.Shape(["n"]),
|
|
) -> R.Tensor(("n - 1", 3), dtype="float32"):
|
|
n = T.int64()
|
|
cls = Expected
|
|
with R.dataflow():
|
|
lv = R.call_tir(
|
|
cls.fused_transpose_matmul,
|
|
(x, y),
|
|
out_ty=R.Tensor((n - 1, 3), dtype="float32"),
|
|
tir_vars=R.shape([n]),
|
|
)
|
|
R.output(lv)
|
|
return lv
|
|
|
|
_check(Module, Expected)
|
|
|
|
|
|
def test_tuple_input_unused_field():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def reshape(
|
|
A: T.Buffer((T.int64(4), T.int64(8), T.int64(2048)), "float32"),
|
|
T_reshape: T.Buffer((T.int64(4), T.int64(8), T.int64(32), T.int64(64)), "float32"),
|
|
):
|
|
T.func_attr({"op_pattern": 2, "tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(8), T.int64(32), T.int64(64)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(
|
|
A[
|
|
(
|
|
((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1)
|
|
// T.int64(8)
|
|
+ v_ax0
|
|
)
|
|
% T.int64(4),
|
|
((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) % T.int64(8),
|
|
(v_ax2 * T.int64(64) + v_ax3) % T.int64(2048),
|
|
]
|
|
)
|
|
T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3])
|
|
T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = A[
|
|
(
|
|
((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) // T.int64(8)
|
|
+ v_ax0
|
|
)
|
|
% T.int64(4),
|
|
((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) % T.int64(8),
|
|
(v_ax2 * T.int64(64) + v_ax3) % T.int64(2048),
|
|
]
|
|
|
|
@R.function(private=True)
|
|
def fused_reshape(
|
|
lv: R.Tuple(
|
|
R.Tensor((4, 8, 2048), dtype="float32"), R.Tensor((4, 8, 2048), dtype="float32")
|
|
),
|
|
) -> R.Tensor((4, 8, 32, 64), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv1: R.Tensor((4, 8, 2048), dtype="float32") = lv[0]
|
|
gv = R.call_tir(
|
|
cls.reshape, (lv1,), out_ty=R.Tensor((4, 8, 32, 64), dtype="float32")
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
tup: R.Tuple(
|
|
R.Tensor((4, 8, 2048), dtype="float32"), R.Tensor((4, 8, 2048), dtype="float32")
|
|
),
|
|
) -> R.Tensor((4, 8, 32, 64), dtype="float32"):
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv_1: R.Tensor((4, 8, 32, 64), dtype="float32") = cls.fused_reshape(tup)
|
|
R.output(lv_1)
|
|
return lv_1
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_reshape(
|
|
lv_0: T.Buffer((T.int64(4), T.int64(8), T.int64(2048)), "float32"),
|
|
T_reshape_handle_intermediate: T.Buffer(
|
|
(T.int64(4), T.int64(8), T.int64(32), T.int64(64)), "float32"
|
|
),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
# with T.sblock("root"):
|
|
for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(8), T.int64(32), T.int64(64)):
|
|
with T.sblock("T_reshape"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T.reads(
|
|
lv_0[
|
|
(
|
|
((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1)
|
|
// T.int64(8)
|
|
+ v_ax0
|
|
)
|
|
% T.int64(4),
|
|
((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) % T.int64(8),
|
|
(v_ax2 * T.int64(64) + v_ax3) % T.int64(2048),
|
|
]
|
|
)
|
|
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] = lv_0[
|
|
(
|
|
((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) // T.int64(8)
|
|
+ v_ax0
|
|
)
|
|
% T.int64(4),
|
|
((v_ax2 * T.int64(64) + v_ax3) // T.int64(2048) + v_ax1) % T.int64(8),
|
|
(v_ax2 * T.int64(64) + v_ax3) % T.int64(2048),
|
|
]
|
|
|
|
@R.function
|
|
def main(
|
|
tup: R.Tuple(
|
|
R.Tensor((4, 8, 2048), dtype="float32"), R.Tensor((4, 8, 2048), dtype="float32")
|
|
),
|
|
) -> R.Tensor((4, 8, 32, 64), dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 8, 2048), dtype="float32") = tup[0]
|
|
lv_1 = R.call_tir(
|
|
cls.fused_reshape, (lv,), out_ty=R.Tensor((4, 8, 32, 64), dtype="float32")
|
|
)
|
|
R.output(lv_1)
|
|
return lv_1
|
|
|
|
_check(Module, Expected)
|
|
|
|
|
|
def test_unique_duplicated_buffer_allocation():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(
|
|
A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
Out: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
):
|
|
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
|
with T.sblock("add"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
Out[vi, vj] = A[vi, vj] + T.float16(1.0)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add1(
|
|
A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
Out: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
):
|
|
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
|
with T.sblock("add"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
Out[vi, vj] = A[vi, vj] + T.float16(2.0)
|
|
|
|
@R.function
|
|
def main(
|
|
input_embeds: R.Tensor((4096, 4096), dtype="float16"),
|
|
) -> R.Tensor((4096, 4096), dtype="float16"):
|
|
cls = Module
|
|
with R.dataflow():
|
|
gv: R.Tensor((4096, 4096), dtype="float16") = cls.fused_func(input_embeds)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function(private=True)
|
|
def fused_func(
|
|
input_embeds: R.Tensor((4096, 4096), dtype="float16"),
|
|
) -> R.Tensor((4096, 4096), dtype="float16"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv = R.call_tir(
|
|
cls.add, (input_embeds,), out_ty=R.Tensor((4096, 4096), dtype="float16")
|
|
)
|
|
gv = R.call_tir(cls.add1, (lv,), out_ty=R.Tensor((4096, 4096), dtype="float16"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_func(
|
|
input_embeds: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
Out_intermediate_1: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
Out_intermediate = T.sblock_alloc_buffer((T.int64(4096), T.int64(4096)), "float16")
|
|
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
|
with T.sblock("add"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
T.reads(input_embeds[vi, vj])
|
|
T.writes(Out_intermediate[vi, vj])
|
|
Out_intermediate[vi, vj] = input_embeds[vi, vj] + T.float16(1)
|
|
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
|
with T.sblock("add_1"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
T.reads(Out_intermediate[vi, vj])
|
|
T.writes(Out_intermediate_1[vi, vj])
|
|
Out_intermediate_1[vi, vj] = Out_intermediate[vi, vj] + T.float16(2)
|
|
|
|
@R.function
|
|
def main(input_embeds: R.Tensor((4096, 4096), dtype="float16")) -> R.Tensor(
|
|
(4096, 4096), dtype="float16"
|
|
):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(
|
|
cls.fused_func,
|
|
(input_embeds,),
|
|
out_ty=R.Tensor((4096, 4096), dtype="float16"),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Module, Expected)
|
|
|
|
|
|
def test_extern_func():
|
|
bb = relax.BlockBuilder()
|
|
bb.add_func(relax.extern("extern_func"), "extern_func")
|
|
mod = bb.get()
|
|
# FuseTIR should keep the ExternFunc in the IRModule.
|
|
_check(mod, mod)
|
|
|
|
|
|
def test_symbolic_var_in_buffer_shape():
|
|
"""A PrimFunc may have dynamic buffer shapes
|
|
|
|
Symbolic variables in a PrimFunc may be present in the buffer
|
|
shape without a corresponding parameter. These symbolic variables
|
|
are inferred from the buffer's shape. (Or, at runtime, they are
|
|
typically determined from the DLTensor's known shape.)
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def foo(
|
|
X_handle: T.handle,
|
|
Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
|
|
rotary_handle: T.handle,
|
|
m: T.int64,
|
|
):
|
|
sequence_length = T.int64()
|
|
|
|
X = T.match_buffer(
|
|
X_handle, [T.int64(1), sequence_length, T.int64(32), T.int64(128)], "float32"
|
|
)
|
|
rotary = T.match_buffer(
|
|
rotary_handle, [T.int64(1), sequence_length, T.int64(32), T.int64(128)], "float32"
|
|
)
|
|
|
|
for i0, i1, i2, i3 in T.grid(T.int64(1), sequence_length, T.int64(32), T.int64(128)):
|
|
with T.sblock("rotary"):
|
|
v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
rotary[v0, v1, v2, v3] = Y[m + v1 - 1, v3] * X[v0, v1, v2, v3]
|
|
|
|
@R.function
|
|
def fused(
|
|
x: R.Tensor((1, "sequence_length", 32, 128), dtype="float32"),
|
|
y: R.Tensor((2048, 128), dtype="float32"),
|
|
len: R.Shape(["m"]),
|
|
) -> R.Tensor((1, "sequence_length", 32, 128), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
sequence_length = T.int64()
|
|
m = T.int64()
|
|
cls = Before
|
|
with R.dataflow():
|
|
lv1 = R.emit_te(topi.add, x, x)
|
|
gv = R.call_tir(
|
|
cls.foo,
|
|
[lv1, y],
|
|
out_ty=R.Tensor((1, sequence_length, 32, 128), dtype="float32"),
|
|
tir_vars=R.shape([m]),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, "sequence_length", 32, 128), dtype="float32"),
|
|
y: R.Tensor((2048, 128), dtype="float32"),
|
|
len: R.Shape(["m"]),
|
|
) -> R.Tensor((1, "sequence_length", 32, 128), dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused(x, y, len)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused(
|
|
X_handle: T.handle,
|
|
Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
|
|
rotary_handle: T.handle,
|
|
m: T.int64,
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
|
|
sequence_length = T.int64()
|
|
|
|
X = T.match_buffer(
|
|
X_handle, [T.int64(1), sequence_length, T.int64(32), T.int64(128)], "float32"
|
|
)
|
|
rotary = T.match_buffer(
|
|
rotary_handle, [T.int64(1), sequence_length, T.int64(32), T.int64(128)], "float32"
|
|
)
|
|
|
|
T_add = T.sblock_alloc_buffer((T.int64(1), sequence_length, T.int64(32), T.int64(128)))
|
|
for ax0, ax1, ax2, ax3 in T.grid(
|
|
T.int64(1), sequence_length, T.int64(32), T.int64(128)
|
|
):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
|
T_add[v_ax0, v_ax1, v_ax2, v_ax3] = (
|
|
X[v_ax0, v_ax1, v_ax2, v_ax3] + X[v_ax0, v_ax1, v_ax2, v_ax3]
|
|
)
|
|
for i0, i1, i2, i3 in T.grid(T.int64(1), sequence_length, T.int64(32), T.int64(128)):
|
|
with T.sblock("rotary"):
|
|
v0, v1, v2, v3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
rotary[v0, v1, v2, v3] = Y[m + v1 - T.int64(1), v3] * T_add[v0, v1, v2, v3]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, "sequence_length", 32, 128), dtype="float32"),
|
|
y: R.Tensor((2048, 128), dtype="float32"),
|
|
len: R.Shape(["m"]),
|
|
) -> R.Tensor((1, "sequence_length", 32, 128), dtype="float32"):
|
|
sequence_length = T.int64()
|
|
m = T.int64()
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(
|
|
cls.fused,
|
|
(x, y),
|
|
out_ty=R.Tensor([1, sequence_length, 32, 128], "float32"),
|
|
tir_vars=R.shape([m]),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_symbolic_var_called_with_static_shape():
|
|
"""A dynamic PrimFunc may be called with a static shape"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def sum_1d(
|
|
X_handle: T.handle,
|
|
Y: T.Buffer([T.int64(1)], "float32"),
|
|
):
|
|
num_elements = T.int64()
|
|
|
|
X = T.match_buffer(X_handle, [num_elements], "float32")
|
|
|
|
for i in range(num_elements):
|
|
with T.sblock("sum"):
|
|
vi = T.axis.remap("R", [i])
|
|
with T.init():
|
|
Y[0] = 0.0
|
|
Y[0] = Y[0] + X[vi]
|
|
|
|
@R.function(private=True)
|
|
def fused(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = R.call_tir(
|
|
cls.sum_1d,
|
|
[x],
|
|
out_ty=R.Tensor([1], dtype="float32"),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused(
|
|
X: T.Buffer([T.int64(64)], "float32"),
|
|
Y: T.Buffer([T.int64(1)], "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
|
|
for i in range(T.int64(64)):
|
|
with T.sblock("sum"):
|
|
vi = T.axis.remap("R", [i])
|
|
with T.init():
|
|
Y[0] = 0.0
|
|
Y[0] = Y[0] + X[vi]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(cls.fused, (x,), out_ty=R.Tensor((1,), dtype="float32"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_symbolic_var_called_with_multiple_static_shapes():
|
|
"""A dynamic PrimFunc may be called with different shapes each time"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def sum_1d(
|
|
X_handle: T.handle,
|
|
Sum: T.Buffer([T.int64(1)], "float32"),
|
|
):
|
|
num_elements = T.int64()
|
|
|
|
X = T.match_buffer(X_handle, [num_elements], "float32")
|
|
|
|
for i in range(num_elements):
|
|
with T.sblock("sum"):
|
|
vi = T.axis.remap("R", [i])
|
|
with T.init():
|
|
Sum[0] = 0.0
|
|
Sum[0] = Sum[0] + X[vi]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def sum_scalar(
|
|
X: T.Buffer([T.int64(1)], "float32"),
|
|
Y: T.Buffer([T.int64(1)], "float32"),
|
|
Sum: T.Buffer([T.int64(1)], "float32"),
|
|
):
|
|
for i in range(T.int64(1)):
|
|
with T.sblock("Out"):
|
|
vi = T.axis.remap("S", [i])
|
|
Sum[vi] = X[vi] + Y[vi]
|
|
|
|
@R.function(private=True)
|
|
def fused(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
y: R.Tensor([16], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
x_sum = R.call_tir(
|
|
cls.sum_1d,
|
|
[x],
|
|
out_ty=R.Tensor([1], dtype="float32"),
|
|
)
|
|
y_sum = R.call_tir(
|
|
cls.sum_1d,
|
|
[y],
|
|
out_ty=R.Tensor([1], dtype="float32"),
|
|
)
|
|
gv = R.call_tir(
|
|
cls.sum_scalar,
|
|
[x_sum, y_sum],
|
|
out_ty=R.Tensor([1], dtype="float32"),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
y: R.Tensor([16], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused(x, y)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused(
|
|
X: T.Buffer([T.int64(64)], "float32"),
|
|
Y: T.Buffer([T.int64(16)], "float32"),
|
|
Out: T.Buffer([T.int64(1)], "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
|
|
XSum = T.sblock_alloc_buffer([T.int64(1)], "float32")
|
|
YSum = T.sblock_alloc_buffer([T.int64(1)], "float32")
|
|
|
|
for i in range(T.int64(64)):
|
|
with T.sblock("XSum"):
|
|
vi = T.axis.remap("R", [i])
|
|
with T.init():
|
|
XSum[0] = 0.0
|
|
XSum[0] = XSum[0] + X[vi]
|
|
|
|
for i in range(T.int64(16)):
|
|
with T.sblock("YSum"):
|
|
vi = T.axis.remap("R", [i])
|
|
with T.init():
|
|
YSum[0] = 0.0
|
|
YSum[0] = YSum[0] + Y[vi]
|
|
|
|
for i in range(T.int64(1)):
|
|
with T.sblock("Out"):
|
|
vi = T.axis.remap("S", [i])
|
|
Out[vi] = XSum[vi] + YSum[vi]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
y: R.Tensor([16], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(cls.fused, (x, y), out_ty=R.Tensor((1,), dtype="float32"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_symbolic_var_called_with_static_argument():
|
|
"""A dynamic PrimFunc may accept a static argument
|
|
|
|
The `tir_vars` parameter in `R.call_tir` contains definitions for
|
|
all TIR variables explicitly listed in the function signature, and
|
|
contains the TIR expression to be passed as the argument for for
|
|
each parameter.
|
|
|
|
This test is identical to the earlier test named
|
|
"test_symbolic_var_called_with_static_shape", except for the
|
|
explicit parameter in `sum_1d`.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def sum_1d(
|
|
X_handle: T.handle,
|
|
Y: T.Buffer([T.int64(1)], "float32"),
|
|
num_elements: T.int64,
|
|
):
|
|
X = T.match_buffer(X_handle, [num_elements], "float32")
|
|
|
|
for i in range(num_elements):
|
|
with T.sblock("sum"):
|
|
vi = T.axis.remap("R", [i])
|
|
with T.init():
|
|
Y[0] = 0.0
|
|
Y[0] = Y[0] + X[vi]
|
|
|
|
@R.function(private=True)
|
|
def fused(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = R.call_tir(
|
|
cls.sum_1d,
|
|
[x],
|
|
out_ty=R.Tensor([1], dtype="float32"),
|
|
tir_vars=R.shape([64]),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused(
|
|
X: T.Buffer([T.int64(64)], "float32"),
|
|
Y: T.Buffer([T.int64(1)], "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
|
|
for i in range(T.int64(64)):
|
|
with T.sblock("sum"):
|
|
vi = T.axis.remap("R", [i])
|
|
with T.init():
|
|
Y[0] = 0.0
|
|
Y[0] = Y[0] + X[vi]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([64], dtype="float32"),
|
|
) -> R.Tensor([1], dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(cls.fused, (x,), out_ty=R.Tensor((1,), dtype="float32"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_gather():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(
|
|
A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
Out: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
):
|
|
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
|
with T.sblock("add"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
Out[vi, vj] = A[vi, vj] + T.float16(1.0)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def take(
|
|
A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
B: T.Buffer((T.int64(1),), "int32"),
|
|
T_take: T.Buffer((T.int64(1), T.int64(4096)), "float16"),
|
|
):
|
|
for ax0, ax1 in T.grid(T.int64(1), T.int64(4096)):
|
|
with T.sblock("T_take"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T_take[v_ax0, v_ax1] = A[B[v_ax0], v_ax1]
|
|
|
|
@R.function
|
|
def main(
|
|
input_ids: R.Tensor((1,), dtype="int32"),
|
|
input_embeds: R.Tensor((4096, 4096), dtype="float16"),
|
|
) -> R.Tensor((1, 4096), dtype="float16"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv: R.Tensor((1, 4096), dtype="float16") = cls.fused_func(input_ids, input_embeds)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function(private=True)
|
|
def fused_func(
|
|
input_ids: R.Tensor((1,), dtype="int32"),
|
|
input_embeds: R.Tensor((4096, 4096), dtype="float16"),
|
|
) -> R.Tensor((1, 4096), dtype="float16"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
lv = R.call_tir(
|
|
cls.add, (input_embeds,), out_ty=R.Tensor((4096, 4096), dtype="float16")
|
|
)
|
|
gv = R.call_tir(
|
|
cls.take, (lv, input_ids), out_ty=R.Tensor((1, 4096), dtype="float16")
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class After:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_func(
|
|
input_ids: T.Buffer((T.int64(1),), "int32"),
|
|
input_embeds: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
T_take: T.Buffer((T.int64(1), T.int64(4096)), "float16"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
Out_handle_intermediate = T.sblock_alloc_buffer(
|
|
(T.int64(4096), T.int64(4096)), "float16"
|
|
)
|
|
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
|
with T.sblock("add"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
Out_handle_intermediate[vi, vj] = input_embeds[vi, vj] + T.float16(1)
|
|
for ax0, ax1 in T.grid(T.int64(1), T.int64(4096)):
|
|
with T.sblock("T_take"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T_take[v_ax0, v_ax1] = Out_handle_intermediate[input_ids[v_ax0], v_ax1]
|
|
|
|
@R.function
|
|
def main(
|
|
input_ids: R.Tensor((1,), dtype="int32"),
|
|
input_embeds: R.Tensor((4096, 4096), dtype="float16"),
|
|
) -> R.Tensor((1, 4096), dtype="float16"):
|
|
cls = After
|
|
with R.dataflow():
|
|
gv = R.call_tir(
|
|
cls.fused_func,
|
|
(input_ids, input_embeds),
|
|
out_ty=R.Tensor((1, 4096), dtype="float16"),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, After)
|
|
|
|
|
|
def test_inplace_simple():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
I.module_attrs({"foo": "bar"})
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add_inplace(
|
|
A: T.Buffer((T.int64(10), T.int64(20)), "float32"), B: T.Buffer((), "float32")
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
# T.reads(A[v_ax0, v_ax1], B[()])
|
|
# T.writes(A[v_ax0, v_ax1])
|
|
A[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[()]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def exp_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("compute"):
|
|
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
|
# T.reads(A[v_i0, v_i1])
|
|
# T.writes(A[v_i0, v_i1])
|
|
A[v_i0, v_i1] = T.exp(A[v_i0, v_i1])
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def squeeze_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_squeeze"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
# T.reads(A[v_ax0, v_ax1])
|
|
# T.writes(A[v_ax0, v_ax1])
|
|
A[v_ax0, v_ax1] = A[v_ax0, v_ax1]
|
|
|
|
@R.function(private=True)
|
|
def fused_add_exp_squeeze(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Module
|
|
with R.dataflow():
|
|
# This overwrites x and is actually evil because the function is marked as pure
|
|
# but we are doing it just to test the pass. The automatic DataflowUseInplaceCalls
|
|
# transformation will not produce code like this, but it may make sense to do it
|
|
# if ownership of x is fully and truly transferred.
|
|
# Users should apply with caution!
|
|
lv = R.call_tir_inplace(
|
|
cls.add_inplace,
|
|
(x, p0),
|
|
inplace_indices=[0],
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
lv1 = R.call_tir_inplace(
|
|
cls.exp_inplace,
|
|
(lv,),
|
|
inplace_indices=[0],
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
gv = R.call_tir_inplace(
|
|
cls.squeeze_inplace,
|
|
(lv1,),
|
|
inplace_indices=[0],
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
cls = Module
|
|
with R.dataflow():
|
|
gv1: R.Tensor((10, 20), dtype="float32") = cls.fused_add_exp_squeeze(x, p0)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
I.module_attrs({"foo": "bar"})
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_add_exp_squeeze(
|
|
x: T.Buffer((T.int64(10), T.int64(20)), "float32"), p0: T.Buffer((), "float32")
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
x[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
|
|
for i0, i1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("compute"):
|
|
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
|
x[v_i0, v_i1] = T.exp(x[v_i0, v_i1])
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_squeeze"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
x[v_ax0, v_ax1] = x[v_ax0, v_ax1]
|
|
|
|
# note that this will clobber x! Use with caution
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv1: R.Tensor((10, 20), dtype="float32") = R.call_tir_inplace(
|
|
cls.fused_add_exp_squeeze,
|
|
(x, p0),
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
inplace_indices=[0],
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
_check(Module, Expected)
|
|
|
|
|
|
def test_fuse_inplace_and_non_inplace():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
I.module_attrs({"foo": "bar"})
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(
|
|
A: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
B: T.Buffer((), "float32"),
|
|
Out: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
Out[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[()]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def exp_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("compute"):
|
|
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
|
A[v_i0, v_i1] = T.exp(A[v_i0, v_i1])
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def squeeze_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_squeeze"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
A[v_ax0, v_ax1] = A[v_ax0, v_ax1]
|
|
|
|
@R.function(private=True)
|
|
def fused_add_exp_squeeze(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv = R.call_tir(
|
|
cls.add,
|
|
(x, p0),
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
lv1 = R.call_tir_inplace(
|
|
cls.exp_inplace,
|
|
(lv,),
|
|
inplace_indices=[0],
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
gv = R.call_tir_inplace(
|
|
cls.squeeze_inplace,
|
|
(lv1,),
|
|
inplace_indices=[0],
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
cls = Module
|
|
with R.dataflow():
|
|
gv1: R.Tensor((10, 20), dtype="float32") = cls.fused_add_exp_squeeze(x, p0)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
I.module_attrs({"foo": "bar"})
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_add_exp_squeeze(
|
|
x: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
p0: T.Buffer((), "float32"),
|
|
p_output0: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
p_output0[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
|
|
for i0, i1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("compute"):
|
|
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
|
p_output0[v_i0, v_i1] = T.exp(p_output0[v_i0, v_i1])
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_squeeze"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
p_output0[v_ax0, v_ax1] = p_output0[v_ax0, v_ax1]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv1: R.Tensor((10, 20), dtype="float32") = R.call_tir(
|
|
cls.fused_add_exp_squeeze,
|
|
(x, p0),
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
_check(Module, Expected)
|
|
|
|
|
|
def test_use_as_inplace_and_dps():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
# we will use it both in-place and normally (DPS)
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(
|
|
A: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
B: T.Buffer((), "float32"),
|
|
Out: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
Out[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[()]
|
|
|
|
@R.function(private=True)
|
|
def fused_sums(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Module
|
|
with R.dataflow():
|
|
lv = R.call_tir(
|
|
cls.add,
|
|
(x, p0),
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
lv1 = R.call_tir_inplace(
|
|
cls.add,
|
|
(x, p0, lv),
|
|
inplace_indices=[2],
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
lv2 = R.call_tir_inplace(
|
|
cls.add,
|
|
(x, p0, lv1),
|
|
inplace_indices=[2],
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
R.output(lv2)
|
|
return lv2
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
cls = Module
|
|
with R.dataflow():
|
|
gv1: R.Tensor((10, 20), dtype="float32") = cls.fused_sums(x, p0)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_sums(
|
|
x: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
p0: T.Buffer((), "float32"),
|
|
p_output0: T.Buffer((T.int64(10), T.int64(20)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
p_output0[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
p_output0[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
|
|
for ax0, ax1 in T.grid(T.int64(10), T.int64(20)):
|
|
with T.sblock("T_add"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
p_output0[v_ax0, v_ax1] = x[v_ax0, v_ax1] + p0[()]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32")
|
|
) -> R.Tensor((10, 20), dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv1: R.Tensor((10, 20), dtype="float32") = R.call_tir(
|
|
cls.fused_sums,
|
|
(x, p0),
|
|
out_ty=R.Tensor((10, 20), dtype="float32"),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
_check(Module, Expected)
|
|
|
|
|
|
def test_private_nonprimitive_func():
|
|
"""Input IRModule may contain calls to non-primitive functions
|
|
|
|
This is a regression test. Prior implementations did not preserve
|
|
relax-to-relax function calls.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
input_ids: R.Tensor((1,), dtype="int32"),
|
|
input_embeds: R.Tensor((4096, 4096), dtype="float16"),
|
|
) -> R.Tensor((1, 4096), dtype="float16"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused_func(input_ids, input_embeds)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function(private=True)
|
|
def fused_func(
|
|
input_ids: R.Tensor((1,), dtype="int32"),
|
|
input_embeds: R.Tensor((4096, 4096), dtype="float16"),
|
|
) -> R.Tensor((1, 4096), dtype="float16"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
lv = R.call_tir(
|
|
cls.add, (input_embeds,), out_ty=R.Tensor((4096, 4096), dtype="float16")
|
|
)
|
|
gv = R.call_tir(
|
|
cls.take, (lv, input_ids), out_ty=R.Tensor((1, 4096), dtype="float16")
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(
|
|
A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
Out: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
):
|
|
for i, j in T.grid(T.int64(4096), T.int64(4096)):
|
|
with T.sblock("add"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
Out[vi, vj] = A[vi, vj] + T.float16(1.0)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def take(
|
|
A: T.Buffer((T.int64(4096), T.int64(4096)), "float16"),
|
|
B: T.Buffer((T.int64(1),), "int32"),
|
|
T_take: T.Buffer((T.int64(1), T.int64(4096)), "float16"),
|
|
):
|
|
for ax0, ax1 in T.grid(T.int64(1), T.int64(4096)):
|
|
with T.sblock("T_take"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T_take[v_ax0, v_ax1] = A[B[v_ax0], v_ax1]
|
|
|
|
_check(Before, Before)
|
|
|
|
|
|
def test_fuse_with_axis_separators():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add(a: T.handle, b: T.handle, c: T.handle):
|
|
A = T.match_buffer(a, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
B = T.match_buffer(b, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
C = T.match_buffer(c, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
|
|
for iters in T.grid(T.int64(16), T.int64(32)):
|
|
with T.sblock("compute"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
C[i, j] = A[i, j] + B[i, j]
|
|
|
|
@R.function(private=True)
|
|
def fused_function(
|
|
x: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
y: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
z: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
w = R.call_tir(
|
|
cls.add, [x, y], out_ty=R.Tensor([T.int64(16), T.int64(32)], "float32")
|
|
)
|
|
out = R.call_tir(
|
|
cls.add, [w, z], out_ty=R.Tensor([T.int64(16), T.int64(32)], "float32")
|
|
)
|
|
R.output(out)
|
|
return out
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
y: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
z: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused_function(x, y, z)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_function(x: T.handle, y: T.handle, z: T.handle, c: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
X = T.match_buffer(x, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
Y = T.match_buffer(y, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
Z = T.match_buffer(z, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
C = T.match_buffer(c, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
Temp = T.sblock_alloc_buffer(X.shape, "float32", axis_separators=[1])
|
|
for iters in T.grid(*X.shape):
|
|
with T.sblock("compute_Y"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
Temp[i, j] = X[i, j] + Y[i, j]
|
|
|
|
for iters in T.grid(*X.shape):
|
|
with T.sblock("compute_Z"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
C[i, j] = Temp[i, j] + Z[i, j]
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
y: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
z: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(
|
|
cls.fused_function,
|
|
[x, y, z],
|
|
out_ty=R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
def test_fuse_with_axis_separators_inconsistent_buffer_mapping():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def mul(a: T.handle, b: T.handle, c: T.handle):
|
|
A = T.match_buffer(a, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
B = T.match_buffer(b, [T.int64(16), T.int64(32)], "float32", axis_separators=[])
|
|
C = T.match_buffer(c, [T.int64(16), T.int64(32)], "float32", axis_separators=[1])
|
|
|
|
for iters in T.grid(T.int64(16), T.int64(32)):
|
|
with T.sblock("compute"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
C[i, j] = A[i, j] * B[i, j]
|
|
|
|
@R.function(private=True)
|
|
def fused_function(
|
|
x: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
out = R.call_tir(
|
|
cls.mul, [x, x], out_ty=R.Tensor([T.int64(16), T.int64(32)], "float32")
|
|
)
|
|
R.output(out)
|
|
return out
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([T.int64(16), T.int64(32)], "float32"),
|
|
) -> R.Tensor([T.int64(16), T.int64(32)], dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused_function(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
with pytest.raises(
|
|
RuntimeError, match=r"Inconsistent buffers.*and.*mapped to the same relax var:.*"
|
|
):
|
|
relax.transform.FuseTIR()(Before)
|
|
|
|
|
|
def test_block_name_numeric_suffix_deduplication():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def add1(x: T.Buffer((10,), "float32"), y: T.Buffer((10,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in range(10):
|
|
with T.sblock("compute1"):
|
|
vi = T.axis.spatial(10, i)
|
|
y[vi] = x[vi] + T.float32(1.0)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def mul1(x: T.Buffer((10,), "float32"), y: T.Buffer((10,), "float32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i in range(10):
|
|
with T.sblock("compute1"):
|
|
vi = T.axis.spatial(10, i)
|
|
y[vi] = x[vi] * T.float32(2.0)
|
|
|
|
@R.function(private=True)
|
|
def fused_add_mul(x: R.Tensor((10,), "float32")) -> R.Tensor((10,), dtype="float32"):
|
|
R.func_attr({"Primitive": True})
|
|
cls = Before
|
|
with R.dataflow():
|
|
lv1 = R.call_tir(cls.add1, (x,), out_ty=R.Tensor((10,), dtype="float32"))
|
|
lv2 = R.call_tir(cls.mul1, (lv1,), out_ty=R.Tensor((10,), dtype="float32"))
|
|
R.output(lv2)
|
|
return lv2
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"):
|
|
cls = Before
|
|
with R.dataflow():
|
|
gv = cls.fused_add_mul(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def fused_add_mul(p_x: T.handle, p_output0: T.handle):
|
|
T.func_attr({"tirx.noalias": True})
|
|
x = T.match_buffer(p_x, (T.int64(10),))
|
|
y_intermediate_1 = T.match_buffer(p_output0, (T.int64(10),), elem_offset=T.int32(0))
|
|
with T.sblock("root"):
|
|
T.reads()
|
|
T.writes()
|
|
y_intermediate = T.sblock_alloc_buffer((T.int64(10),), elem_offset=T.int32(0))
|
|
for i in range(10):
|
|
with T.sblock("compute1"):
|
|
vi = T.axis.spatial(10, i)
|
|
T.reads(x[vi])
|
|
T.writes(y_intermediate[vi])
|
|
y_intermediate[vi] = x[vi] + T.float32(1.0)
|
|
for i in range(10):
|
|
with T.sblock("compute2"):
|
|
vi = T.axis.spatial(10, i)
|
|
T.reads(y_intermediate[vi])
|
|
T.writes(y_intermediate_1[vi])
|
|
y_intermediate_1[vi] = y_intermediate[vi] * T.float32(2.0)
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"):
|
|
cls = Expected
|
|
with R.dataflow():
|
|
gv = R.call_tir(cls.fused_add_mul, (x,), out_ty=R.Tensor((10,), dtype="float32"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_check(Before, Expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|