Skip to content

[Fix][Relax] Preserve tensor-derived symbols during fusion #20139

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akaashrp wants to merge 1 commit into
base: main
Choose a base branch
from
+224 −2
Open

[Fix][Relax] Preserve tensor-derived symbols during fusion #20139

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 30 additions & 2 deletions src/relax/transform/fuse_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -679,10 +679,38 @@ class FunctionCreator : public ExprMutator {
if (const auto* tuple = expr.as<TupleNode>()) {
return std::all_of(tuple->fields.begin(), tuple->fields.end(),
[this](const Expr& e) { return IsInlinableConstants(e); });
} else if (auto prim_value = expr.as<PrimExpr>()) {
ffi::Array<tirx::Var> undefined_vars = tvm::tirx::UndefinedVars(prim_value.value());
if (undefined_vars.empty()) {
return true;
}

// A symbolic value directly defined by an existing tensor/shape parameter is already part
// of the grouped function's shape environment. A variable that only occurs inside a
// derived dimension is not directly definable, but leaving it symbolic lets CreateFunction
// add the explicit Shape parameter required to define it. In either case, lifting the value
// as an unrelated scalar parameter would sever the relation to the parameter types.
auto parameter_types =
TupleType(params_.Map([](const Var& param) { return GetType(param); }));
std::unordered_set<tirx::Var> definable_shape_vars;
for (const tirx::Var& var : DefinableTIRVarsInType(parameter_types)) {
definable_shape_vars.insert(var);
}
if (std::all_of(undefined_vars.begin(), undefined_vars.end(),
[&definable_shape_vars](const tirx::Var& var) {
return definable_shape_vars.count(var);
})) {
return true;
}

std::unordered_set<tirx::Var> referenced_shape_vars;
for (const tirx::Var& var : TIRVarsInType(parameter_types)) {
referenced_shape_vars.insert(var);
}
return std::all_of(undefined_vars.begin(), undefined_vars.end(),
[&](const tirx::Var& var) { return referenced_shape_vars.count(var); });
} else if (expr.as<VarNode>() || expr.as<CallNode>()) {
return false;
} else if (auto prim_value = expr.as<PrimExpr>()) {
return tvm::tirx::UndefinedVars(prim_value.value()).empty();
} else if (const auto* shape_expr = expr.as<ShapeExprNode>()) {
return std::all_of(shape_expr->values.begin(), shape_expr->values.end(),
[](const PrimExpr& e) { return tvm::tirx::UndefinedVars(e).empty(); });
Expand Down
194 changes: 194 additions & 0 deletions tests/python/relax/test_transform_fuse_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1352,6 +1352,200 @@ def main(s: R.Shape(["n"])) -> R.Tensor((1, 1, "n", "n"), dtype="float32"):
_check(Before, Expected)


def test_symbolic_prim_arg_bound_by_tensor_shape():
@I.ir_module(s_tir=True)
class Before:
@T.prim_func(private=True, s_tir=True)
def add_one(x_handle: T.handle, n: T.int64, out_handle: T.handle):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
x = T.match_buffer(x_handle, (T.int64(1), n), "float32")
out = T.match_buffer(out_handle, (T.int64(1), n), "float32")
for i in range(n):
with T.sblock("add_one"):
vi = T.axis.spatial(n, i)
out[0, vi] = x[0, vi] + T.float32(1)

@T.prim_func(private=True, s_tir=True)
def exp(x_handle: T.handle, n: T.int64, out_handle: T.handle):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
x = T.match_buffer(x_handle, (T.int64(1), n), "float32")
out = T.match_buffer(out_handle, (T.int64(1), n), "float32")
for i in range(n):
with T.sblock("exp"):
vi = T.axis.spatial(n, i)
out[0, vi] = T.exp(x[0, vi])

@R.function
def main(
x: R.Tensor((1, "n"), dtype="float32"),
) -> R.Tensor((1, "n"), dtype="float32"):
n = T.int64()
cls = Before
with R.dataflow():
lv = R.call_tir(
cls.add_one,
(x, n),
out_ty=R.Tensor((1, n), dtype="float32"),
)
gv = R.call_tir(
cls.exp,
(lv, n),
out_ty=R.Tensor((1, n), dtype="float32"),
)
R.output(gv)
return gv

mod = relax.transform.AnnotateTIROpPattern()(Before)
mod = relax.transform.FuseOps()(mod)

fused = next(
mod[global_var]
for global_var in mod.get_global_vars()
if global_var.name_hint.startswith("fused_")
)
assert len(fused.params) == 1
assert fused.ret_ty.shape is not None


def test_symbolic_prim_arg_reused_from_derived_tensor_shape():
@I.ir_module(s_tir=True)
class Before:
@T.prim_func(private=True, s_tir=True)
def add_one(x_handle: T.handle, n: T.int64, out_handle: T.handle):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
x = T.match_buffer(
x_handle,
(T.int64(1), (n - T.int64(1)) // T.int64(4) + T.int64(1)),
"float32",
)
out = T.match_buffer(
out_handle,
(T.int64(1), (n - T.int64(1)) // T.int64(4) + T.int64(1)),
"float32",
)
for i in range((n - T.int64(1)) // T.int64(4) + T.int64(1)):
with T.sblock("add_one"):
vi = T.axis.spatial((n - T.int64(1)) // T.int64(4) + T.int64(1), i)
out[0, vi] = x[0, vi] + T.float32(1)

@T.prim_func(private=True, s_tir=True)
def exp(x_handle: T.handle, n: T.int64, out_handle: T.handle):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
x = T.match_buffer(
x_handle,
(T.int64(1), (n - T.int64(1)) // T.int64(4) + T.int64(1)),
"float32",
)
out = T.match_buffer(
out_handle,
(T.int64(1), (n - T.int64(1)) // T.int64(4) + T.int64(1)),
"float32",
)
for i in range((n - T.int64(1)) // T.int64(4) + T.int64(1)):
with T.sblock("exp"):
vi = T.axis.spatial((n - T.int64(1)) // T.int64(4) + T.int64(1), i)
out[0, vi] = T.exp(x[0, vi])

@R.function
def main(
source: R.Tensor(("n",), dtype="float32"),
x: R.Tensor((1, "(n - 1) // 4 + 1"), dtype="float32"),
) -> R.Tensor((1, "(n - 1) // 4 + 1"), dtype="float32"):
n = T.int64()
cls = Before
with R.dataflow():
lv = R.call_tir(
cls.add_one,
(x, n),
out_ty=R.Tensor((1, (n - 1) // 4 + 1), dtype="float32"),
)
gv = R.call_tir(
cls.exp,
(lv, n),
out_ty=R.Tensor((1, (n - 1) // 4 + 1), dtype="float32"),
)
R.output(gv)
return gv

mod = relax.transform.AnnotateTIROpPattern()(Before)
mod = relax.transform.FuseOps()(mod)

fused = next(
mod[global_var]
for global_var in mod.get_global_vars()
if global_var.name_hint.startswith("fused_")
)
assert len(fused.params) == 2
assert isinstance(fused.params[1].ty, relax.ShapeType)
assert all(not isinstance(param.ty, tvm.ir.PrimType) for param in fused.params)

mod = relax.transform.FuseTIR()(mod)
fused_tir = next(
mod[global_var]
for global_var in mod.get_global_vars()
if global_var.name_hint.startswith("fused_")
)
assert tvm.tirx.analysis.verify_well_formed(fused_tir)


def test_symbolic_prim_arg_not_bound_by_derived_tensor_shape():
@I.ir_module(s_tir=True)
class Before:
@T.prim_func(private=True, s_tir=True)
def add_one(x_handle: T.handle, n: T.int64, m: T.int64, out_handle: T.handle):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
x = T.match_buffer(x_handle, (T.int64(1), n + T.int64(1)), "float32")
out = T.match_buffer(out_handle, (T.int64(1), n + T.int64(1)), "float32")
for i in range(n + T.int64(1)):
with T.sblock("add_one"):
vi = T.axis.spatial(n + T.int64(1), i)
out[0, vi] = x[0, vi] + T.float32(1)

@T.prim_func(private=True, s_tir=True)
def exp(x_handle: T.handle, n: T.int64, m: T.int64, out_handle: T.handle):
T.func_attr({"op_pattern": 0, "tirx.noalias": True})
x = T.match_buffer(x_handle, (T.int64(1), n + T.int64(1)), "float32")
out = T.match_buffer(out_handle, (T.int64(1), n + T.int64(1)), "float32")
for i in range(n + T.int64(1)):
with T.sblock("exp"):
vi = T.axis.spatial(n + T.int64(1), i)
out[0, vi] = T.exp(x[0, vi])

@R.function
def main(
shape: R.Shape(["n", "m"]),
x: R.Tensor((1, "n + 1"), dtype="float32"),
) -> R.Tensor((1, "n + 1"), dtype="float32"):
n = T.int64()
m = T.int64()
cls = Before
with R.dataflow():
lv = R.call_tir(
cls.add_one,
(x, n, m),
out_ty=R.Tensor((1, n + 1), dtype="float32"),
)
gv = R.call_tir(
cls.exp,
(lv, n, m),
out_ty=R.Tensor((1, n + 1), dtype="float32"),
)
R.output(gv)
return gv

mod = relax.transform.AnnotateTIROpPattern()(Before)
mod = relax.transform.FuseOps()(mod)

fused = next(
mod[global_var]
for global_var in mod.get_global_vars()
if global_var.name_hint.startswith("fused_")
)
assert len(fused.params) == 3
assert sum(isinstance(param.ty, relax.ShapeType) for param in fused.params) == 1
assert sum(isinstance(param.ty, tvm.ir.PrimType) for param in fused.params) == 1


def test_shape_expr_arg():
@I.ir_module(s_tir=True)
class Before:
Expand Down
Loading