Skip to content

[Unity][Transform] Handle relax.Var as call_tir args when lowering #15916

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

Closed
Lunderberg wants to merge 10 commits into from
Closed

[Unity][Transform] Handle relax.Var as call_tir args when lowering #15916

Changes from 1 commit
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
a00fe50
[Unity] Provide LookupBinding utility in ExprVisitor
Lunderberg Oct 3, 2023
5bd4819
[Unity] Add utility method UnwrapBindings for ExprVisitor/Mutator
Lunderberg Oct 11, 2023
7c426af
[Unity] Handle bound tuple variables in relax.transform.FoldConstant
Lunderberg Oct 3, 2023
d4f27f2
[Unity] Handle tuple variables in relax.transform.FuseOps
Lunderberg Oct 3, 2023
36bea9b
[Unity] Handle tuple variables in RewriteDataflowReshape
Lunderberg Oct 3, 2023
7a91d13
[Unity] Handle tuple vars in CallTIRRewrite
Lunderberg Oct 4, 2023
a48c911
lint fix
Lunderberg Oct 12, 2023
fc16d5d
Check for TupleStructInfo in all call_* variants
Lunderberg Oct 26, 2023
4108fb0
Pull in improvements and unit tests from PR#15971
Lunderberg Oct 26, 2023
5089f47
Fix breakage caused by incorrect merging with PR#15971 changes
Lunderberg Oct 30, 2023
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
Prev Previous commit
Next Next commit
Check for TupleStructInfo in all call_* variants
  • Loading branch information
@Lunderberg
Lunderberg committed Oct 31, 2023
commit fc16d5da7de6e48a1d1f281c22ad3187fa2775e5
28 changes: 16 additions & 12 deletions python/tvm/relax/op/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,18 @@ def null_value() -> Call:
return _ffi_api.null_value() # type: ignore


def _normalize_arg_tuple(args: Expr) -> Expr:
if isinstance(args, RxTuple) or isinstance(args.struct_info_, TupleStructInfo):
# A tuple, or a Var bound to a tuple, are kept as-is
return args
elif isinstance(args, Expr):
# A single argument is wrapped into a tuple
return RxTuple((args,))
else:
# Anything else is left for the FFI to handle
return args


@args_converter.auto
def call_tir(
gvar: GlobalVar,
Expand Down Expand Up @@ -97,12 +109,7 @@ def call_tir(
ret: Call
A call node for the call_tir operator.
"""
if (
isinstance(args, Expr)
and not isinstance(args, RxTuple)
and not isinstance(args.struct_info_, TupleStructInfo)
):
args = RxTuple((args,))
args = _normalize_arg_tuple(args)

if not isinstance(out_sinfo, list):
out_sinfo = [out_sinfo]
Expand Down Expand Up @@ -156,8 +163,7 @@ def call_tir_with_grad(
ret: Call
A call node for the call_tir_with_grad operator.
"""
if isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
args = RxTuple((args,))
args = _normalize_arg_tuple(args)

if not isinstance(out_sinfo, list):
out_sinfo = [out_sinfo]
Expand Down Expand Up @@ -224,8 +230,7 @@ def call_tir_inplace(
ret: Call
A call node for the call_tir operator.
"""
if isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
args = RxTuple((args,))
args = _normalize_arg_tuple(args)

if not isinstance(inplace_indices, list):
inplace_indices = [inplace_indices]
Expand Down Expand Up @@ -279,8 +284,7 @@ def call_dps_packed(
if isinstance(func, str):
func = ExternFunc(func)

if isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
args = RxTuple((args,))
args = _normalize_arg_tuple(args)

if not isinstance(out_sinfo, list):
out_sinfo = [out_sinfo]
Expand Down