Skip to content
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

Jump to bottom

[Transform] Operator legalizer V0 #96

Merged
merged 5 commits into from
Jan 16, 2023
Merged

[Transform] Operator legalizer V0 #96

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a2ee5ad
[TIR][Fix] Buffer slicing using index dtype as extent (#13788)
MasterJH5574 Jan 16, 2023
d42582a
[TIR][Fix] IndexDataTypeNormalizer not unwrapping float casting (#13789)
MasterJH5574 Jan 16, 2023
911249b
[Fix][TVMScript] Parse and print `tir_vars` of `call_tir` properly (t…
MasterJH5574 Jan 16, 2023
3a3f60d
[Transform] Operator legalizer
MasterJH5574 Jan 13, 2023
cb58430
Documentation
MasterJH5574 Jan 16, 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
9 changes: 6 additions & 3 deletions python/tvm/relax/op/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# pylint: disable=redefined-builtin
"""The base Relax operators."""
from typing import Union, List, Optional
from typing import Union, List, Tuple, Optional


import tvm
Expand Down Expand Up @@ -47,7 +47,7 @@ def call_tir(
args: Union[Expr, List[Expr]],
shape: Union[RxTuple, ShapeExpr, List[int]],
dtype: Union[str, List[str]],
tir_vars: Optional[ShapeExpr] = None,
tir_vars: Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]] = None,
) -> Call:
"""
Call a destination-passing-style function and return the output.
Expand All @@ -66,7 +66,7 @@ def call_tir(
dtype: Union[str, List[str]]
The output dtype. List[str] if multiple outputs, str if single output.

tir_vars : ShapeExpr, optional
tir_vars : Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used

Returns
Expand Down Expand Up @@ -122,6 +122,9 @@ def _create_shape(shape: List[Union[int, PrimExpr]]) -> ShapeExpr:
else:
raise TypeError("Not supported dtype for call_tir: " + str(type(dtype)))

if isinstance(tir_vars, (list, tuple)):
tir_vars = ShapeExpr(tir_vars)

return _ffi_api.call_tir(func, args, shape, output_type, tir_vars) # type: ignore


Expand Down
1 change: 1 addition & 0 deletions python/tvm/relax/transform/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,4 @@

from .transform import *
from .fma_rewrite import *
from .legalize_ops import LegalizeOps
Loading