wehub-resource-sync
300 lines
9.2 KiB
Python
300 lines
9.2 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.
|
|
"""APIs for pattern-based rewriting."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
from tvm_ffi import register_object
|
|
|
|
from tvm.ir import IRModule
|
|
from tvm.runtime import Object
|
|
|
|
from ..expr import Expr, Function, Var
|
|
from . import _ffi as ffi
|
|
from .context import PatternContext
|
|
from .pattern import DFPattern
|
|
|
|
|
|
@register_object("relax.dpl.PatternMatchingRewriter")
|
|
class PatternMatchingRewriter(Object):
|
|
"""A pattern-matching rewriter for Relax"""
|
|
|
|
@staticmethod
|
|
def from_pattern(
|
|
pattern: DFPattern,
|
|
func: Callable[[Expr, dict[DFPattern, Expr]], Expr],
|
|
) -> "PatternMatchingRewriter":
|
|
"""Construct from a pattern and rewriter-function
|
|
|
|
The replacements performed by the rewriter will be equivalent
|
|
to using the `pattern` and `func` as arguments to
|
|
`rewrite_call`.
|
|
|
|
Parameters
|
|
----------
|
|
pattern: DFPattern
|
|
|
|
The pattern to be matched against.
|
|
|
|
func: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
|
|
|
|
A function that returns the rewritten expression. See
|
|
`rewrite_call` for details and examples.
|
|
|
|
|
|
Returns
|
|
-------
|
|
rewriter_obj: PatternMatchingRewriter
|
|
|
|
The rewriter object
|
|
|
|
"""
|
|
return ffi.PatternMatchingRewriterFromPattern(
|
|
pattern,
|
|
func,
|
|
) # type: ignore
|
|
|
|
@staticmethod
|
|
def from_module(mod: IRModule) -> "PatternMatchingRewriter":
|
|
"""Construct a rewriter from an IRModule
|
|
|
|
The IRModule must have two publicly-exposed functions,
|
|
`pattern` and `replacement`, where `pattern` and `replacement`
|
|
have the same function signature, as shown in the example
|
|
below.
|
|
|
|
.. code-block:: python
|
|
|
|
@I.ir_module
|
|
class RewriteAddIntoMultiply:
|
|
@R.function
|
|
def pattern(A: R.Tensor):
|
|
B = A + A
|
|
return B
|
|
|
|
@R.function
|
|
def replacement(A: R.Tensor):
|
|
B = A * 2
|
|
return B
|
|
|
|
rewriter = PatternMatchingRewriter.from_module(RewriteAddIntoMultiply)
|
|
rewritten_ir_module = rewriter(ir_module)
|
|
|
|
To support the common case of defining an IRModule with
|
|
TVMScript, then immediately turning it into a rewriter, the
|
|
`@R.rewriter` annotation can be used.
|
|
|
|
.. code-block:: python
|
|
|
|
@R.rewriter
|
|
class RewriteAddIntoMultiply:
|
|
@R.function
|
|
def pattern(A: R.Tensor):
|
|
B = A + A
|
|
return B
|
|
|
|
@R.function
|
|
def replacement(A: R.Tensor):
|
|
B = A * 2
|
|
return B
|
|
|
|
rewritten_ir_module = RewriteAddIntoMultiply(ir_module)
|
|
|
|
Parameters
|
|
----------
|
|
mod: IRModule
|
|
|
|
A module with `pattern` and `replacement` functions,
|
|
defining a rewrite rule.
|
|
|
|
|
|
Returns
|
|
-------
|
|
rewriter_obj: PatternMatchingRewriter
|
|
|
|
The rewriter object
|
|
|
|
"""
|
|
return ffi.PatternMatchingRewriterFromModule(mod) # type: ignore
|
|
|
|
def __call__(self, obj: Expr | IRModule) -> Expr | IRModule:
|
|
"""Apply the rewriter
|
|
|
|
Parameters
|
|
----------
|
|
obj: Union[Expr, IRModule])
|
|
|
|
The object to be rewritten. May be applied to either a
|
|
relax expression, or an IRModule.
|
|
|
|
Returns
|
|
-------
|
|
updated: Union[Expr, IRModule]
|
|
|
|
The rewritten object
|
|
|
|
"""
|
|
return ffi.PatternMatchingRewriterApply(self, obj)
|
|
|
|
def __or__(self, other: "PatternMatchingRewriter") -> "PatternMatchingRewriter":
|
|
"""Compose two rewriters
|
|
|
|
Composing two rewrite rules together allows them to be applied
|
|
in a single Relax-level transformation.
|
|
|
|
Parameters
|
|
----------
|
|
other: PatternMatchingRewriter
|
|
|
|
Another rewrite rule
|
|
|
|
Returns
|
|
-------
|
|
PatternMatchingRewriter
|
|
|
|
A rewriter that will apply either rewrite pattern
|
|
|
|
"""
|
|
return OrRewriter(self, other)
|
|
|
|
|
|
@register_object("relax.dpl.ExprPatternRewriter")
|
|
class ExprPatternRewriter(PatternMatchingRewriter):
|
|
def __init__(self, pattern, func):
|
|
self.__init_handle_by_constructor__(
|
|
ffi.PatternRewriter,
|
|
pattern,
|
|
func,
|
|
) # type: ignore
|
|
|
|
|
|
@register_object("relax.dpl.OrRewriter")
|
|
class OrRewriter(PatternMatchingRewriter):
|
|
def __init__(self, lhs, rhs):
|
|
self.__init_handle_by_constructor__(
|
|
ffi.OrRewriter,
|
|
lhs,
|
|
rhs,
|
|
) # type: ignore
|
|
|
|
|
|
@register_object("relax.dpl.TupleRewriter")
|
|
class TupleRewriter(PatternMatchingRewriter):
|
|
def __init__(self, patterns, func):
|
|
self.__init_handle_by_constructor__(
|
|
ffi.TupleRewriter,
|
|
patterns,
|
|
func,
|
|
) # type: ignore
|
|
|
|
|
|
def rewrite_call(
|
|
pattern: DFPattern,
|
|
rewriter: Callable[[Expr, dict[DFPattern, Expr]], Expr],
|
|
func: Function,
|
|
) -> Function:
|
|
"""
|
|
Rewrite a function with the given pattern and the rewriter function.
|
|
|
|
Parameters
|
|
----------
|
|
pattern: DFPattern
|
|
The pattern to match.
|
|
|
|
rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
|
|
The function to be called on a successful matching for rewriting. Given the matched
|
|
call node and the map of patterns and matched expressions, it should return a new call node
|
|
to replace the original one or the original matched call node as is.
|
|
|
|
For example, to replace x + x with 2 * x, we can write the rewriter as follows:
|
|
```
|
|
x = wildcard()
|
|
pattern = is_op("relax.add")(x, x)
|
|
|
|
def rewriter(orig, matchings):
|
|
return R.multiply(matchings[x], R.const(2, "float32"))
|
|
```
|
|
|
|
func: Function
|
|
The function to rewrite.
|
|
|
|
Returns
|
|
-------
|
|
rewritten_func: Function
|
|
The rewritten or the input function, depending on the pattern matching result.
|
|
"""
|
|
return ffi.rewrite_call(pattern, rewriter, func)
|
|
|
|
|
|
def rewrite_bindings(
|
|
ctx: PatternContext,
|
|
rewriter: Callable[[dict[DFPattern, Var], dict[Var, Expr]], dict[Var, Expr]],
|
|
func: Function,
|
|
) -> Function:
|
|
"""
|
|
Rewrite a function with the given pattern and the rewriter function.
|
|
|
|
Parameters
|
|
----------
|
|
ctx: PatternContext
|
|
The pattern constraint context under which rewriting takes place.
|
|
|
|
rewriter: Callable[[Dict[DFPattern, Var], Dict[Var, Expr]], Dict[Var, Expr]]
|
|
The function to be called on a successful matching for rewriting. Given the map of patterns
|
|
and corresponding variables (bound variables or parameters), it should return a map that
|
|
specifies new values for matched bound variables. It can refer to the passed bindings to
|
|
create the replacement expressions.
|
|
|
|
For example, to rewrite three matmuls for QKV projection in transformer models into one
|
|
matmul followed by slicing, one can use the follwoing rewriter:
|
|
```
|
|
inp_pat = wildcard()
|
|
Q_weight_pat, K_weight_pat, V_weight_pat = wildcard(), wildcard(), wildcard()
|
|
|
|
matmul1 = is_op("relax.matmul")(inp_pat, Q_weight_pat)
|
|
matmul2 = is_op("relax.matmul")(inp_pat, K_weight_pat)
|
|
matmul3 = is_op("relax.matmul")(inp_pat, V_weight_pat)
|
|
|
|
def rewriter(matchings):
|
|
inp = matchings[inp_pat]
|
|
Q_weight = matchings[Q_weight_pat]
|
|
K_weight = matchings[K_weight_pat]
|
|
V_weight = matchings[V_weight_pat]
|
|
width = Q_weight.ty.shape[1]
|
|
|
|
concat = R.concat([Q_weight, K_weight, V_weight], axis=1)
|
|
matmul = R.matmul(inp, concat)
|
|
Q = R.strided_slice(matmul, axes=[2], begin=[0], end=[width])
|
|
K = R.strided_slice(matmul, axes=[2], begin=[width], end=[width * 2])
|
|
V = R.strided_slice(matmul, axes=[2], begin=[width * 2], end=[width * 3])
|
|
|
|
# matchings[matmul1] gives the bound variable in the binding whose RHS matches with
|
|
# the matmul1 pattern. For example, lv0 in lv0 = R.matmul(x1, w0).
|
|
# We want to replace the RHS of this binding with Q.
|
|
return {matchings[matmul1]: Q, matchings[matmul2]: K, matchings[matmul3]: V}
|
|
```
|
|
|
|
func: Function
|
|
The function to rewrite.
|
|
|
|
Returns
|
|
-------
|
|
rewritten_func: Function
|
|
The rewritten or the input function, depending on the pattern matching result.
|
|
"""
|
|
return ffi.rewrite_bindings(ctx, rewriter, func)
|