wehub-resource-sync
1164 lines
31 KiB
Python
1164 lines
31 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: F821
|
|
|
|
"""Pattern types in Relax Dataflow Pattern Language"""
|
|
# pylint: disable=no-member
|
|
# pylint: disable=pointless-statement
|
|
|
|
from typing import Union
|
|
|
|
import tvm_ffi
|
|
from tvm_ffi import Array
|
|
|
|
import tvm
|
|
from tvm.ir.op import Op
|
|
|
|
from ...ir import make_node
|
|
from ...ir.base import Node
|
|
from ...runtime import Object
|
|
from ..expr import Expr, Var
|
|
from . import _ffi as ffi
|
|
|
|
|
|
def register_df_node(type_key=None):
|
|
"""
|
|
Register a Relax node type
|
|
|
|
Parameters
|
|
----------
|
|
type_key : str or cls
|
|
The type key of the node
|
|
"""
|
|
if not isinstance(type_key, str):
|
|
return tvm_ffi.register_object("relax.dpl." + type_key.__name__)(type_key)
|
|
return tvm_ffi.register_object(type_key)
|
|
|
|
|
|
class DFPattern(Node):
|
|
"""Base class of all Patterns."""
|
|
|
|
def __call__(self, *args, varg_default_wildcard=False, add_constraint=True) -> "CallPattern":
|
|
"""
|
|
Syntax sugar for creating a CallPattern with argument patterns
|
|
|
|
Returns
|
|
-------
|
|
result: CallPattern
|
|
The resulting CallPattern
|
|
"""
|
|
return CallPattern(self, args, varg_default_wildcard, add_constraint)
|
|
|
|
def __or__(self, other: "DFPattern") -> "OrPattern":
|
|
"""
|
|
Syntax sugar for creating an OrPattern
|
|
|
|
Parameters
|
|
----------
|
|
other: DFPattern
|
|
Alternative pattern
|
|
|
|
Returns
|
|
-------
|
|
result: OrPattern
|
|
The resulting OrPattern
|
|
"""
|
|
return OrPattern(self, other)
|
|
|
|
def __and__(self, other: "DFPattern") -> "AndPattern":
|
|
"""
|
|
Syntax sugar for creating an AndPattern
|
|
|
|
Parameters
|
|
----------
|
|
other: DFPattern
|
|
Additional pattern to satisfy
|
|
|
|
Returns
|
|
-------
|
|
result: AndPattern
|
|
The resulting AndPattern
|
|
"""
|
|
return AndPattern(self, other)
|
|
|
|
def __invert__(self) -> "NotPattern":
|
|
"""
|
|
Syntax sugar for creating a DFPattern to reject
|
|
|
|
Returns
|
|
-------
|
|
result: NotPattern
|
|
The resulting NotPattern
|
|
"""
|
|
return reject(self)
|
|
|
|
def has_attr(self, attrs: dict[str, Object]) -> "AttrPattern":
|
|
"""
|
|
Add an attribute constraint to this pattern
|
|
|
|
Parameters
|
|
----------
|
|
attrs: Dict[str, Object]
|
|
|
|
Returns
|
|
-------
|
|
result: AttrPattern
|
|
The resulting AttrPattern
|
|
"""
|
|
attrs = make_node("ir.DictAttrs", **attrs)
|
|
return AttrPattern(self, attrs)
|
|
|
|
def has_ty(self, ty: "Type") -> "TypePattern":
|
|
return TypePattern(self, ty)
|
|
|
|
def has_dtype(self, dtype: str) -> "DataTypePattern":
|
|
"""
|
|
Add a type constraint to this pattern
|
|
|
|
Parameters
|
|
----------
|
|
dtype: str
|
|
The dtype to match
|
|
|
|
Returns
|
|
-------
|
|
result: DataTypePattern
|
|
The resulting DataTypePattern
|
|
"""
|
|
return has_dtype(dtype, self)
|
|
|
|
def has_shape(self, shape: list[Expr]) -> "ShapePattern":
|
|
"""
|
|
Add a shape constraint to this pattern
|
|
|
|
Parameters
|
|
----------
|
|
shape: List[Expr]
|
|
Expected shape list
|
|
|
|
Returns
|
|
-------
|
|
result: ShapePattern
|
|
The resulting ShapePattern
|
|
|
|
Note
|
|
----
|
|
has_shape assumes that the matched relax.Expr only has one
|
|
output tensor. Use is_tuple for those with multiple outputs.
|
|
"""
|
|
if not isinstance(shape, list | tuple) and not tvm.ir.is_prim_expr(shape):
|
|
raise ValueError("has_shape takes a list or tuple as input.")
|
|
return ShapePattern(pattern=self, shape=shape)
|
|
|
|
def match(self, expr, var2val: dict[Var, Expr] | None = None) -> bool:
|
|
"""
|
|
Match a relax.Expr syntactically
|
|
|
|
Parameters
|
|
----------
|
|
expr : tvm.relax.Expr
|
|
The expression to match
|
|
var2val : Optional[Dict[tvm.relax.Var, tvm.relax.Expr]]
|
|
A mapping from relax.Var to relax.Expr for autojump.
|
|
|
|
Returns
|
|
-------
|
|
result: bool
|
|
Whether or not the expression matches the pattern
|
|
|
|
Note
|
|
----
|
|
Functions in Relax consist of blocks of bindings that are not syntactically connected.
|
|
We use a mapping (i.e., var2val) to mitigate the gap. For example, when matching
|
|
"relax.add(lv0, lv1)", given var2val, we match lv0's bound expression
|
|
when the recursive pattern matching goes to check lv0. The var2val mapping
|
|
can be computed through the tvm.relax.analysis.get_var2val function.
|
|
"""
|
|
return ffi.match_expr(self, expr, var2val) # type: ignore
|
|
|
|
def extract_matched_expr(
|
|
self, expr, var2val: dict[Var, Expr] | None = None
|
|
) -> dict["DFPattern", Expr] | None:
|
|
"""
|
|
Match a relax.Expr and return a map from matching patterns to matched expressions.
|
|
|
|
Parameters
|
|
----------
|
|
expr : tvm.relax.Expr
|
|
The expression to match
|
|
var2val : Optional[Dict[tvm.relax.Var, tvm.relax.Expr]]
|
|
A mapping from relax.Var to relax.Expr for autojump.
|
|
|
|
Returns
|
|
-------
|
|
result: Optional[Dict[DFPattern, Expr]]
|
|
Map from matching patterns to matched expressions.
|
|
Return None if the pattern does not match expr.
|
|
|
|
Note
|
|
----
|
|
Check the note of `match` for the meaning of var2val.
|
|
"""
|
|
return ffi.extract_matched_expr(self, expr, var2val)
|
|
|
|
def used_by(self, other: Union["DFPattern", "PatternSeq"], index=-1) -> "PatternSeq":
|
|
"""
|
|
The current pattern being used by another pattern (sequence)
|
|
|
|
Parameters
|
|
----------
|
|
other : Union[DFPattern, DFPattern]
|
|
The consumer pattern (sequence)
|
|
index : int, optional
|
|
The argument index called by the consumer pattern, by default -1
|
|
|
|
Returns
|
|
-------
|
|
result: PatternSeq
|
|
A chained pattern sequence
|
|
"""
|
|
return _used_by(self, other, index)
|
|
|
|
def __xor__(self, other: Union["DFPattern", "PatternSeq"]) -> "PatternSeq":
|
|
"""Syntax sugar of DFPattern.used_by"""
|
|
return self.used_by(other, -1)
|
|
|
|
def only_used_by(self, other: Union["DFPattern", "PatternSeq"], index=-1) -> "PatternSeq":
|
|
"""
|
|
The current pattern being **ONLY** used by another pattern (sequence)
|
|
|
|
Parameters
|
|
----------
|
|
other : Union[DFPattern, DFPattern]
|
|
The consumer pattern (sequence)
|
|
index : int, optional
|
|
The argument index called by the consumer pattern, by default -1
|
|
|
|
Returns
|
|
-------
|
|
result: PatternSeq
|
|
A chained pattern sequence
|
|
"""
|
|
return _only_used_by(self, other, index)
|
|
|
|
def __rshift__(self, other: Union["DFPattern", "PatternSeq"]) -> "PatternSeq":
|
|
"""Syntax sugar of DFPattern.only_used_by"""
|
|
return self.only_used_by(other, -1)
|
|
|
|
def dup(self) -> "DFPattern":
|
|
"""
|
|
Duplicate the current pattern (new object under different address)
|
|
|
|
Returns
|
|
-------
|
|
DFPattern
|
|
A duplicated pattern
|
|
"""
|
|
return ffi.dup_pattern(self) # type: ignore
|
|
|
|
def fork_to(self, *args) -> None:
|
|
"""Fork the current pattern to multiple pattern branches"""
|
|
for v in args:
|
|
self ^ v
|
|
|
|
def same_shape_as(self, *args: list["DFPattern"]) -> "SameShapeConstraint":
|
|
"""
|
|
The current pattern with the same shape as another pattern (sequence)
|
|
|
|
Parameters
|
|
----------
|
|
other : List[DFPattern]
|
|
The other pattern (sequence)
|
|
|
|
Returns
|
|
-------
|
|
result: PatternSeq
|
|
A chained pattern sequence
|
|
"""
|
|
return SameShapeConstraint(self, *args)
|
|
|
|
|
|
class DFConstraint(Node):
|
|
"""Base class of all constraints."""
|
|
|
|
|
|
@register_df_node
|
|
class ExprPattern(DFPattern):
|
|
"""A pattern which matches an expression.
|
|
|
|
Parameters
|
|
----------
|
|
expr : tvm.relax.Expr
|
|
The expression to match.
|
|
"""
|
|
|
|
def __init__(self, expr: Expr):
|
|
self.__init_handle_by_constructor__(ffi.ExprPattern, expr) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class VarPattern(DFPattern):
|
|
"""A pattern for Var.
|
|
|
|
Parameters
|
|
----------
|
|
name_hint: str
|
|
The name of the variable. Optional, if not provided,
|
|
the pattern will match any VarNode.
|
|
"""
|
|
|
|
def __init__(self, name_hint: str = ""):
|
|
self.__init_handle_by_constructor__(ffi.VarPattern, name_hint) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class DataflowVarPattern(VarPattern):
|
|
"""A pattern for DataflowVar.
|
|
|
|
Parameters
|
|
----------
|
|
name_hint: str
|
|
The name of the variable. Optional, if not provided,
|
|
the pattern will match any VarNode.
|
|
"""
|
|
|
|
def __init__(self, name_hint: str = ""):
|
|
self.__init_handle_by_constructor__(ffi.DataflowVarPattern, name_hint) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class GlobalVarPattern(DFPattern):
|
|
"""A pattern for GlobalVar.
|
|
|
|
Parameters
|
|
----------
|
|
name_hint: str
|
|
The name of the variable. Optional, if not provided,
|
|
the pattern will match any GlobalVarNode.
|
|
"""
|
|
|
|
def __init__(self, name_hint: str = ""):
|
|
self.__init_handle_by_constructor__(ffi.GlobalVarPattern, name_hint) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class ExternFuncPattern(DFPattern):
|
|
"""A external function pattern.
|
|
|
|
Parameters
|
|
----------
|
|
global_symbol: str
|
|
The name of the function. Optional, if not provided,
|
|
the pattern will match any ExternFuncNode.
|
|
"""
|
|
|
|
def __init__(self, global_symbol: str = ""):
|
|
self.__init_handle_by_constructor__(ffi.ExternFuncPattern, global_symbol) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class ConstantPattern(DFPattern):
|
|
"""A pattern matching a Relax Constant."""
|
|
|
|
def __init__(self):
|
|
self.__init_handle_by_constructor__(ffi.ConstantPattern) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class CallPattern(DFPattern):
|
|
"""A pattern matching a function call node.
|
|
|
|
Parameters
|
|
----------
|
|
op: tvm.relax.dpl.DFPattern
|
|
The operation to be called.
|
|
|
|
args: List[tvm.relax.dpl.DFPattern]
|
|
The arguments to the call or None to match any arguments.
|
|
|
|
varg_default_wildcard: bool
|
|
If True, args can be fewer than actual provided arguments.
|
|
|
|
add_constraint: bool
|
|
If True, automatically add "used-by" constraints between caller and callee expressions.
|
|
|
|
Note
|
|
----
|
|
By setting varg_default_wildcard to True, we can only focus on the argument
|
|
patterns we specified. For example, CallPattern(Op, [A, B]) can match
|
|
a call of Op(A, B) or Op(A, B, C, ...) that has more arguments. However,
|
|
the specified argument patterns must be matched (i.e., A and B).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
op: "DFPattern",
|
|
args: list["DFPattern"] | tuple["DFPattern", ...],
|
|
varg_default_wildcard: bool = False,
|
|
add_constraint=True,
|
|
):
|
|
self.__init_handle_by_constructor__(
|
|
ffi.CallPattern,
|
|
op,
|
|
args,
|
|
varg_default_wildcard, # type: ignore
|
|
)
|
|
|
|
if add_constraint:
|
|
for i, arg in enumerate(args):
|
|
arg.used_by(self, i)
|
|
|
|
|
|
@register_df_node
|
|
class FunctionPattern(DFPattern):
|
|
"""A pattern matching a function node in Relax.
|
|
|
|
Parameters
|
|
----------
|
|
params: List[tvm.relax.dpl.DFPattern]
|
|
The parameters to the Function or None to match any parameters.
|
|
|
|
body: tvm.relax.dpl.DFPattern
|
|
The body fo the Function
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
params: list["DFPattern"],
|
|
body: "DFPattern",
|
|
):
|
|
self.__init_handle_by_constructor__(ffi.FunctionPattern, params, body) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class TuplePattern(DFPattern):
|
|
"""A patern matching a Relax Tuple.
|
|
|
|
Parameters
|
|
----------
|
|
fields : Array[tvm.relax.dpl.DFPattern]
|
|
The fields in the tuple.
|
|
"""
|
|
|
|
def __init__(self, fields: list):
|
|
self.__init_handle_by_constructor__(ffi.TuplePattern, fields) # type: ignore
|
|
|
|
def __getitem__(self, index: int | None) -> "TupleGetItemPattern":
|
|
if index is not None:
|
|
# support negative index for being pythonic
|
|
if index < 0:
|
|
index += len(self)
|
|
if index >= len(self):
|
|
raise IndexError("TuplePattern index out of range")
|
|
else:
|
|
index = -1 # -1 means matching any index
|
|
return TupleGetItemPattern(self, index)
|
|
|
|
def __len__(self):
|
|
return len(self.fields)
|
|
|
|
|
|
@register_df_node
|
|
class UnorderedTuplePattern(DFPattern):
|
|
"""A patern matching a Relax Tuple unorderedly.
|
|
|
|
Parameters
|
|
----------
|
|
fields : Array[tvm.relax.dpl.DFPattern]
|
|
The fields in the tuple.
|
|
"""
|
|
|
|
def __init__(self, fields: Array):
|
|
self.__init_handle_by_constructor__(ffi.UnorderedTuplePattern, fields) # type: ignore
|
|
|
|
def __len__(self):
|
|
return len(self.fields)
|
|
|
|
|
|
@register_df_node
|
|
class TupleGetItemPattern(DFPattern):
|
|
"""Get index-th item from a TuplePattern.
|
|
|
|
Parameters
|
|
----------
|
|
tuple_value: tvm.relax.dpl.DFPattern
|
|
The input tuple expression.
|
|
|
|
index: Optional[int]
|
|
The index to match; Default (None) to match a TupleGetItem with any index.
|
|
"""
|
|
|
|
def __init__(self, tuple_value: "DFPattern", index: int | None = None):
|
|
match_index = index if index is not None else -1
|
|
self.__init_handle_by_constructor__(
|
|
ffi.TupleGetItemPattern,
|
|
tuple_value,
|
|
match_index, # type: ignore
|
|
)
|
|
|
|
|
|
@register_df_node
|
|
class OrPattern(DFPattern):
|
|
"""Create a Pattern that can match one of two conditions
|
|
|
|
Parameters
|
|
----------
|
|
left: tvm.relax.dpl.DFPattern
|
|
One possible matching pattern.
|
|
right: tvm.relax.dpl.DFPattern
|
|
One possible matching pattern.
|
|
"""
|
|
|
|
def __init__(self, left: "DFPattern", right: "DFPattern"):
|
|
self.__init_handle_by_constructor__(ffi.OrPattern, left, right) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class AndPattern(DFPattern):
|
|
"""Create a Pattern that must match two conditions
|
|
|
|
Parameters
|
|
----------
|
|
left: tvm.relax.dpl.DFPattern
|
|
One must-matching pattern.
|
|
right: tvm.relax.dpl.DFPattern
|
|
One must-matching pattern.
|
|
"""
|
|
|
|
def __init__(self, left: "DFPattern", right: "DFPattern"):
|
|
self.__init_handle_by_constructor__(ffi.AndPattern, left, right) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class NotPattern(DFPattern):
|
|
"""Create a Pattern that matches the negation of a condition.
|
|
|
|
Parameters
|
|
----------
|
|
to_reject: tvm.relax.dpl.DFPattern
|
|
The pattern to deny.
|
|
"""
|
|
|
|
def __init__(self, to_reject: "DFPattern"):
|
|
self.__init_handle_by_constructor__(ffi.NotPattern, to_reject) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class WildcardPattern(DFPattern):
|
|
"""A pattern which matches anything."""
|
|
|
|
def __init__(self):
|
|
self.__init_handle_by_constructor__(ffi.WildcardPattern) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class TypePattern(DFPattern):
|
|
"""A pattern that matches another pattern with a certain Type
|
|
|
|
Parameters
|
|
----------
|
|
pattern: tvm.relax.dpl.DFPattern
|
|
The input pattern that needs type annotation.
|
|
|
|
ty: tvm.relax.Type
|
|
The type to match against
|
|
"""
|
|
|
|
def __init__(self, pattern: "DFPattern", ty: "Type"):
|
|
self.__init_handle_by_constructor__(
|
|
ffi.TypePattern,
|
|
pattern,
|
|
ty,
|
|
) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class DataTypePattern(DFPattern):
|
|
"""A pattern that matches another pattern with certain data type
|
|
|
|
Parameters
|
|
----------
|
|
pattern: tvm.relax.dpl.DFPattern
|
|
The input pattern that needs type annotation.
|
|
|
|
dtype: str
|
|
The dtype to match.
|
|
"""
|
|
|
|
def __init__(self, pattern: "DFPattern", dtype: str):
|
|
self.__init_handle_by_constructor__(ffi.DataTypePattern, pattern, dtype) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class ShapePattern(DFPattern):
|
|
"""A pattern that matches another pattern with a certain tensor shape
|
|
|
|
Parameters
|
|
----------
|
|
pattern: tvm.relax.dpl.DFPattern
|
|
The input pattern that needs type annotation.
|
|
|
|
shape: List[tvm.ir.Expr]
|
|
The shape to match.
|
|
"""
|
|
|
|
def __init__(self, pattern: "DFPattern", shape: list[tvm.ir.Expr]):
|
|
self.__init_handle_by_constructor__(ffi.ShapePattern, pattern, shape) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class SameShapeConstraint(DFConstraint):
|
|
"""A pattern that requires a set of patterns to have the same shape
|
|
|
|
Parameters
|
|
----------
|
|
args: List[DFPattern]
|
|
A set of patterns which must all provide the same shape.
|
|
"""
|
|
|
|
def __init__(self, *args: list[DFPattern]):
|
|
self.__init_handle_by_constructor__(ffi.SameShapeConstraint, args) # type: ignore
|
|
|
|
|
|
@register_df_node
|
|
class PrimArrPattern(DFPattern):
|
|
"""
|
|
A pattern to match an array of Expr
|
|
|
|
Parameters
|
|
----------
|
|
shape : List[tvm.ir.Expr]
|
|
The shape to match.
|
|
"""
|
|
|
|
def __init__(self, shape: list[tvm.ir.Expr]):
|
|
self.__init_handle_by_constructor__(ffi.PrimArrPattern, shape) # type: ignore
|
|
|
|
def __getitem__(self, index: int):
|
|
if index >= len(self):
|
|
raise IndexError("PrimArrPattern index out of range")
|
|
return self.fields[index]
|
|
|
|
def __len__(self):
|
|
return len(self.fields)
|
|
|
|
|
|
@register_df_node
|
|
class AttrPattern(DFPattern):
|
|
"""Match an expression with certain attributes.
|
|
Supports Op attributes, Call attributes, and Function attributes.
|
|
|
|
Parameters
|
|
----------
|
|
pattern: tvm.relax.dpl.DFPattern
|
|
The input pattern.
|
|
|
|
attrs: tvm.ir.attrs.Attrs
|
|
The attributes to match.
|
|
"""
|
|
|
|
def __init__(self, pattern: "DFPattern", attrs: tvm.ir.attrs.Attrs):
|
|
self.__init_handle_by_constructor__(ffi.AttrPattern, pattern, attrs) # type: ignore
|
|
|
|
|
|
def is_var(name: str = "") -> VarPattern:
|
|
"""
|
|
Syntatic sugar for creating an optionally named VarPattern.
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
The name of the input pattern to match.
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.VarPattern
|
|
The resulting pattern.
|
|
"""
|
|
return VarPattern(name)
|
|
|
|
|
|
def is_gv(name: str = "") -> GlobalVarPattern:
|
|
"""Syntax sugar for creating an optionally (if name is empty) named GlobalVarPattern."""
|
|
return GlobalVarPattern(name)
|
|
|
|
|
|
def is_dfv(name: str = "") -> DataflowVarPattern:
|
|
"""Syntax sugar for creating an optionally (if name is empty) named DataflowVarPattern."""
|
|
return DataflowVarPattern(name)
|
|
|
|
|
|
def is_const() -> ConstantPattern:
|
|
"""
|
|
Syntatic sugar for creating a ConstantPattern.
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
The name of the input pattern to match.
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.ConstantPattern
|
|
The resulting pattern.
|
|
"""
|
|
return ConstantPattern()
|
|
|
|
|
|
def is_expr(expr: Expr) -> ExprPattern:
|
|
"""
|
|
Syntatic sugar for creating an ExprPattern.
|
|
|
|
Parameters
|
|
----------
|
|
expr: Expr
|
|
The Relax expression to match.
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.ExprPattern
|
|
The resulting pattern.
|
|
"""
|
|
return ExprPattern(expr)
|
|
|
|
|
|
def is_op(op_name: str) -> ExprPattern:
|
|
"""
|
|
Syntatic sugar for creating an operator ExprPattern.
|
|
|
|
Parameters
|
|
----------
|
|
op_name: String
|
|
The name of the tvm.ir.op.Op object
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.ExprPattern
|
|
The resulting ExprPattern
|
|
"""
|
|
op = Op.get(op_name)
|
|
return ExprPattern(op)
|
|
|
|
|
|
def is_tuple(fields: Array | list | tuple, unordered=False) -> TuplePattern | UnorderedTuplePattern:
|
|
"""
|
|
Syntatic sugar for creating an ExprPattern.
|
|
|
|
Parameters
|
|
----------
|
|
fields : Array[tvm.relax.dpl.DFPattern]
|
|
The fields in the tuple.
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.DFPattern
|
|
The resulting pattern.
|
|
"""
|
|
if not isinstance(fields, list | tuple | Array):
|
|
raise ValueError("fields must be a list, tuple, or Array")
|
|
if unordered:
|
|
return UnorderedTuplePattern(fields)
|
|
return TuplePattern(fields)
|
|
|
|
|
|
def is_tuple_get_item(tuple_value: DFPattern, index: int | None = None) -> TupleGetItemPattern:
|
|
"""
|
|
Syntatic sugar for creating an ExprPattern.
|
|
|
|
Parameters
|
|
----------
|
|
tuple_value: tvm.relax.dpl.DFPattern
|
|
The input tuple expression.
|
|
|
|
index: Optional[int]
|
|
The index to match; Default (None) to match a TupleGetItem with any index.
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.TupleGetItemPattern
|
|
The resulting pattern.
|
|
"""
|
|
return TupleGetItemPattern(tuple_value, index)
|
|
|
|
|
|
def wildcard() -> WildcardPattern:
|
|
"""
|
|
Syntatic sugar for creating a WildcardPattern.
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.WildcardPattern
|
|
The resulting pattern.
|
|
"""
|
|
return WildcardPattern()
|
|
|
|
|
|
def has_dtype(dtype: str, pattern: DFPattern = None) -> DataTypePattern:
|
|
"""
|
|
Syntatic sugar for creating a DataTypePattern
|
|
|
|
Parameters
|
|
----------
|
|
dtype: str
|
|
The dtype to match
|
|
|
|
pattern: tvm.relax.dpl.DFPattern
|
|
The pattern that needs type annotation
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.DataTypePattern
|
|
The resulting DataTypePattern
|
|
"""
|
|
if pattern is None:
|
|
pattern = wildcard()
|
|
return DataTypePattern(pattern, dtype)
|
|
|
|
|
|
def is_shape(shape: list[tvm.ir.Expr]) -> "PrimArrPattern":
|
|
"""
|
|
Directly matches a shape which is an array of Expr
|
|
|
|
Parameters
|
|
----------
|
|
shape : List[tvm.ir.Expr]
|
|
The expected shape
|
|
|
|
Returns
|
|
-------
|
|
PrimArrPattern
|
|
The resulting PrimArrPattern pattern
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
If the argument shape is not a list/tuple/tvm_ffi.Array
|
|
|
|
Note
|
|
----
|
|
The difference between p.has_shape(s) and is_shape(s) is that: has_shape
|
|
puts assumptions on the shape of the tensor matched by pattern p. While
|
|
is_shape directly matches the shape (an array of Expr).
|
|
"""
|
|
if not isinstance(shape, list | tuple | Array):
|
|
raise ValueError("is_shape takes a list or tuple as input.")
|
|
return PrimArrPattern(shape)
|
|
|
|
|
|
# Todo(relax-team): Dataflow pattern for Type, and match out_ty
|
|
def _is_call_tir(
|
|
func_pattern: DFPattern,
|
|
args: list | tuple | TuplePattern = None,
|
|
tir_vars: DFPattern | None = None,
|
|
) -> CallPattern:
|
|
if args is None:
|
|
args = wildcard()
|
|
elif isinstance(args, list | tuple):
|
|
args = TuplePattern(args)
|
|
|
|
if tir_vars is None:
|
|
return is_op("relax.call_tir")(func_pattern, args, add_constraint=False)
|
|
return is_op("relax.call_tir")(func_pattern, args, tir_vars, add_constraint=False)
|
|
|
|
|
|
# Todo(relax-team): Dataflow pattern for Type, and match out_ty
|
|
def is_call_tir(
|
|
func_name: str,
|
|
args: list | tuple | TuplePattern = None,
|
|
tir_vars: DFPattern | None = None,
|
|
) -> CallPattern:
|
|
"""
|
|
Syntax sugar for creating a CallPattern for call_tir that calls an function through global var.
|
|
|
|
Parameters
|
|
----------
|
|
func_name : str
|
|
Name of the CPS function to call.
|
|
args : Union[List[DFPattern], Tuple[DFPattern]], optional
|
|
Arguments in expected call_packed, by default None meaning arbitrary (number of) arguments
|
|
tir_vars : Optional[DFPattern]
|
|
Pattern to match the tuple of integers that are unpacked when calling the tirx func.
|
|
Returns
|
|
-------
|
|
CallPattern
|
|
The resulting CallPattern
|
|
"""
|
|
func_pattern = GlobalVarPattern(func_name)
|
|
return _is_call_tir(func_pattern, args, tir_vars)
|
|
|
|
|
|
def _is_call_dps_packed(
|
|
func_pattern: DFPattern,
|
|
args: list | tuple | TuplePattern = None,
|
|
) -> CallPattern:
|
|
if args is None:
|
|
args = wildcard()
|
|
elif isinstance(args, list | tuple):
|
|
args = TuplePattern(args)
|
|
|
|
return is_op("relax.call_dps_packed")(func_pattern, args, add_constraint=False)
|
|
|
|
|
|
def is_call_dps_packed(
|
|
func_name: str,
|
|
args: list | tuple | TuplePattern = None,
|
|
) -> CallPattern:
|
|
"""Syntax sugar for creating a CallPattern for call_dps_packed
|
|
|
|
Parameters
|
|
----------
|
|
func_name : str
|
|
Name of the CPS function to call.
|
|
args : Union[List[DFPattern], Tuple[DFPattern]], optional
|
|
Arguments in expected call_packed, by default None meaning arbitrary (number of) arguments
|
|
|
|
Returns
|
|
-------
|
|
CallPattern
|
|
The resulting CallPattern
|
|
"""
|
|
func_pattern = ExternFuncPattern(func_name)
|
|
return _is_call_dps_packed(func_pattern, args)
|
|
|
|
|
|
def is_call_packed(
|
|
func_name: str, args: list[DFPattern] | tuple[DFPattern] | None = None
|
|
) -> CallPattern:
|
|
"""
|
|
Syntax sugar for creating a CallPattern for call_packed
|
|
|
|
Parameters
|
|
----------
|
|
func_name : str
|
|
Name of the external function to call
|
|
args : Union[List[DFPattern], Tuple[DFPattern]], optional
|
|
Arguments in expected call_packed, by default None meaning arbitrary (number of) arguments
|
|
|
|
Returns
|
|
-------
|
|
CallPattern
|
|
The resulting CallPattern
|
|
"""
|
|
if args is None:
|
|
return ExternFuncPattern(func_name)(varg_default_wildcard=True, add_constraint=False)
|
|
return ExternFuncPattern(func_name)(*args)
|
|
|
|
|
|
def reject(pattern: DFPattern) -> NotPattern:
|
|
"""
|
|
Syntax sugar for creating a DFPattern to reject
|
|
|
|
Parameters
|
|
----------
|
|
pattern : DFPattern
|
|
The pattern to deny
|
|
|
|
Returns
|
|
-------
|
|
result: NotPattern
|
|
The resulting NotPattern
|
|
"""
|
|
return NotPattern(pattern)
|
|
|
|
|
|
def has_attr(attrs, pattern=None) -> AttrPattern:
|
|
"""
|
|
Syntatic sugar for creating an AttrPattern
|
|
|
|
Parameters
|
|
----------
|
|
attrs: Dict[str, Object]
|
|
The attributes to match
|
|
|
|
pattern: Optional[tvm.relax.dpl.DFPattern]
|
|
The input pattern.
|
|
|
|
Returns
|
|
-------
|
|
result: tvm.relax.dpl.DFPattern
|
|
The resulting AttrPattern
|
|
"""
|
|
if pattern is None:
|
|
pattern = wildcard()
|
|
return pattern.has_attr(attrs)
|
|
|
|
|
|
@register_df_node
|
|
class PatternSeq(Node):
|
|
"""A sequence of patterns with consecutive constraints"""
|
|
|
|
def __init__(self, patterns: list[DFPattern], only_use=False):
|
|
"""
|
|
Initializer to PatternSeq
|
|
|
|
Parameters
|
|
----------
|
|
patterns : List[DFPattern]
|
|
A chain of patterns
|
|
only_use : bool, optional
|
|
Whether the patterns follows only-used-by relations consecutively, by default False
|
|
"""
|
|
self.__init_handle_by_constructor__(ffi.PatternSeq, patterns, only_use) # type: ignore
|
|
|
|
def used_by(self, other: Union[DFPattern, "PatternSeq"], index=-1) -> "PatternSeq":
|
|
"""
|
|
Assuming the right-most pattern must be used by the `other` pattern as a producer
|
|
|
|
Parameters
|
|
----------
|
|
other : Union[DFPattern, PatternSeq]
|
|
The consumer pattern (sequence)
|
|
index : int, optional
|
|
The argument index called by the consumer pattern, by default -1
|
|
|
|
Returns
|
|
-------
|
|
PatternSeq
|
|
A chained pattern sequence
|
|
|
|
Note
|
|
----
|
|
If other is PatternSeq, it means the right-most pattern must be used by the left-most
|
|
pattern of the other sequence.
|
|
"""
|
|
return _used_by(self, other, index)
|
|
|
|
def only_used_by(self, other: Union[DFPattern, "PatternSeq"], index=-1) -> "PatternSeq":
|
|
"""
|
|
Assuming the right-most pattern must be **ONLY** used by the `other` pattern as a producer
|
|
|
|
Parameters
|
|
----------
|
|
other : Union[DFPattern, PatternSeq]
|
|
The consumer pattern (sequence)
|
|
index : int, optional
|
|
The argument index called by the consumer pattern, by default -1
|
|
|
|
Returns
|
|
-------
|
|
PatternSeq
|
|
A chained pattern sequence
|
|
|
|
Note
|
|
----
|
|
If other is PatternSeq, it means the right-most pattern must be **ONLY** used by the
|
|
left-most pattern of the other sequence.
|
|
"""
|
|
return _only_used_by(self, other, index)
|
|
|
|
def __getitem__(self, index: int) -> DFPattern:
|
|
"""
|
|
Access the pattern at the given index
|
|
|
|
Parameters
|
|
----------
|
|
index : int
|
|
Index of the accessed pattern
|
|
|
|
Returns
|
|
-------
|
|
DFPattern
|
|
The accessed pattern
|
|
"""
|
|
return self.patterns[index]
|
|
|
|
def __xor__(self, other) -> "PatternSeq":
|
|
"""Syntax sugar of PatternSeq.used_by"""
|
|
return self.used_by(other, -1)
|
|
|
|
def __rshift__(self, other) -> "PatternSeq":
|
|
"""Syntax sugar of PatternSeq.only_used_by"""
|
|
return self.only_used_by(other, -1)
|
|
|
|
def dup(self) -> "PatternSeq":
|
|
"""
|
|
Duplicate the pattern sequence (new object under different address)
|
|
|
|
Returns
|
|
-------
|
|
PatternSeq
|
|
A duplicated chain
|
|
"""
|
|
return ffi.dup_seq(self) # type: ignore
|
|
|
|
|
|
### Private functions
|
|
|
|
|
|
def _used_by(
|
|
lhs: DFPattern | PatternSeq,
|
|
rhs: DFPattern | PatternSeq,
|
|
index=-1,
|
|
) -> PatternSeq:
|
|
if isinstance(lhs, DFPattern):
|
|
lhs = PatternSeq([lhs])
|
|
if isinstance(rhs, DFPattern):
|
|
rhs = PatternSeq([rhs])
|
|
return ffi.used_by(lhs, rhs, index) # type: ignore
|
|
|
|
|
|
def _only_used_by(lhs: DFPattern | PatternSeq, rhs: DFPattern | PatternSeq, index=-1) -> PatternSeq:
|
|
if isinstance(lhs, DFPattern):
|
|
lhs = PatternSeq([lhs])
|
|
if isinstance(rhs, DFPattern):
|
|
rhs = PatternSeq([rhs])
|
|
return ffi.only_used_by(lhs, rhs, index) # type: ignore
|
|
|
|
|
|
def make_fused_bias_activation_pattern(
|
|
op_name, with_bias=False, activation=None, allow_reshape=False
|
|
):
|
|
"""
|
|
A simple utility to create patterns for an operation fused with bias addition and activation.
|
|
|
|
Parameters
|
|
----------
|
|
op_name: str
|
|
The name of a Relax op, such as "relax.nn.conv2d"
|
|
|
|
with_bias: bool
|
|
Whether or not to include bias addition
|
|
|
|
activation: str
|
|
The name of an activation Relax op, such as "relax.nn.relu"
|
|
|
|
allow_reshape: bool
|
|
Whether to allow reshape operation before bias addition (for PyTorch frontend)
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a fused operation
|
|
"""
|
|
lhs = wildcard()
|
|
rhs = wildcard()
|
|
out = is_op(op_name)(lhs, rhs)
|
|
|
|
if with_bias:
|
|
bias = wildcard()
|
|
if allow_reshape:
|
|
reshaped_bias = is_op("relax.reshape")(bias, wildcard(), varg_default_wildcard=True)
|
|
out = is_op("relax.add")(out, reshaped_bias, varg_default_wildcard=True)
|
|
else:
|
|
out = is_op("relax.add")(out, bias)
|
|
|
|
if activation:
|
|
return is_op(activation)(out)
|
|
|
|
return out
|