wehub-resource-sync
644 lines
20 KiB
Python
644 lines
20 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.
|
|
# pylint: disable=invalid-name
|
|
"""Common patterns used in BYOC"""
|
|
|
|
from collections.abc import Mapping
|
|
|
|
from tvm.relax.dpl.pattern import (
|
|
DFPattern,
|
|
GlobalVarPattern,
|
|
TuplePattern,
|
|
is_const,
|
|
is_op,
|
|
is_tuple_get_item,
|
|
wildcard,
|
|
)
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def _with_bias_activation_pattern(
|
|
out: DFPattern,
|
|
annotations: dict[str, DFPattern],
|
|
with_bias: bool = False,
|
|
activation: str | None = None,
|
|
allow_reshape: bool = False,
|
|
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
|
if with_bias:
|
|
annotations["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:
|
|
out = is_op(activation)(out)
|
|
|
|
return out, annotations
|
|
|
|
|
|
def make_fused_bias_activation_pattern(
|
|
op_name: str,
|
|
with_bias: bool = False,
|
|
activation: str | None = None,
|
|
allow_reshape: bool = False,
|
|
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
|
"""
|
|
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"
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a fused operation
|
|
|
|
annotations: Mapping[str, DFPattern]
|
|
A mapping from name to sub pattern. It can be used to extract
|
|
important expressions from match result, to power the partition
|
|
check function and codegen.
|
|
"""
|
|
lhs = wildcard()
|
|
rhs = wildcard()
|
|
out = is_op(op_name)(lhs, rhs)
|
|
annotations = {"lhs": lhs, "rhs": rhs, "root": out}
|
|
|
|
return _with_bias_activation_pattern(out, annotations, with_bias, activation, allow_reshape)
|
|
|
|
|
|
def make_residual_block_pattern(
|
|
node_output: DFPattern | tuple[DFPattern, Mapping[str, DFPattern]],
|
|
binary_op="relax.add",
|
|
activation=None,
|
|
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
|
"""
|
|
Create pattern for residual block.
|
|
|
|
Parameters
|
|
----------
|
|
node_output: Union[DFPattern, Tuple[DFPattern, Mapping[str, DFPattern]]]
|
|
The output of previous node.
|
|
|
|
binary_op: str
|
|
The op used to combine previous node output and residual input.
|
|
|
|
activation: str
|
|
The activation function of this residual block. It should be a name of
|
|
activation Relax op, such as "relax.nn.relu".
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a matrix multiplication.
|
|
|
|
annotations: Mapping[str, DFPattern]
|
|
A mapping from name to sub pattern. It can be used to extract
|
|
important expressions from match result, to power the partition
|
|
check function and codegen.
|
|
"""
|
|
|
|
if isinstance(node_output, tuple):
|
|
node_output, arg_patterns = node_output
|
|
else:
|
|
arg_patterns = {}
|
|
|
|
residual_input = wildcard()
|
|
op = is_op(binary_op)
|
|
output = op(node_output, residual_input) | op(residual_input, node_output)
|
|
|
|
if activation is not None:
|
|
output = is_op(activation)(output)
|
|
|
|
return output, {**arg_patterns, "residual": residual_input}
|
|
|
|
|
|
def make_conv2d_pattern(
|
|
with_bias: bool = False,
|
|
activation: str | None = None,
|
|
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
|
"""
|
|
Create pattern for 2D convolution.
|
|
|
|
Parameters
|
|
----------
|
|
with_bias: bool
|
|
Whether or not to include bias addition
|
|
|
|
activation: str
|
|
The name of an activation Relax op, such as "relax.nn.relu"
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a 2D convolution.
|
|
|
|
annotations: Mapping[str, DFPattern]
|
|
A mapping from name to sub pattern. It can be used to extract
|
|
important expressions from match result, to power the partition
|
|
check function and codegen.
|
|
"""
|
|
|
|
input_tensor = wildcard()
|
|
kernel = wildcard()
|
|
annotations = {"input": input_tensor, "weight": kernel}
|
|
|
|
conv2d = is_op("relax.nn.conv2d")(input_tensor, kernel)
|
|
annotations["root"] = conv2d
|
|
|
|
return _with_bias_activation_pattern(conv2d, annotations, with_bias, activation)
|
|
|
|
|
|
def make_matmul_pattern(
|
|
with_bias: bool = False,
|
|
activation: str | None = None,
|
|
transposed_rhs: bool = False,
|
|
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
|
"""
|
|
Create pattern for matrix multiplication.
|
|
|
|
Parameters
|
|
----------
|
|
with_bias: bool
|
|
Whether or not to include bias addition
|
|
|
|
activation: str
|
|
The name of an activation Relax op, such as "relax.nn.relu"
|
|
|
|
transposed_rhs: bool
|
|
Whether the right hand side of multiplication is transposed.
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a matrix multiplication.
|
|
|
|
annotations: Mapping[str, DFPattern]
|
|
A mapping from name to sub pattern. It can be used to extract
|
|
important expressions from match result, to power the partition
|
|
check function and codegen.
|
|
"""
|
|
|
|
lhs = wildcard()
|
|
rhs = wildcard()
|
|
annotations = {"lhs": lhs, "rhs": rhs}
|
|
|
|
if transposed_rhs:
|
|
rhs = is_op("relax.permute_dims")(rhs)
|
|
|
|
out = is_op("relax.matmul")(lhs, rhs)
|
|
annotations["root"] = out
|
|
|
|
return _with_bias_activation_pattern(out, annotations, with_bias, activation)
|
|
|
|
|
|
def make_attention_pattern(with_bias: bool = False, var_len: bool = False):
|
|
"""
|
|
Create pattern for fused multi head attention.
|
|
|
|
Parameters
|
|
----------
|
|
with_bias: bool
|
|
Whether or not to include bias addition.
|
|
|
|
var_len: bool
|
|
Whether or not to make a pattern for batched attention with variable sequence lengths.
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a fused multi head attention.
|
|
|
|
annotations: Mapping[str, DFPattern]
|
|
A mapping from name to sub pattern. It can be used to extract
|
|
important expressions from match result, to power the partition
|
|
check function and codegen.
|
|
"""
|
|
query = wildcard()
|
|
key = wildcard()
|
|
value = wildcard()
|
|
annotations = {"query": query, "key": key, "value": value}
|
|
if with_bias:
|
|
bias = wildcard()
|
|
annotations["bias"] = bias
|
|
out = is_op("relax.nn.attention_bias")(query, key, value, bias)
|
|
elif var_len:
|
|
seqstart_q = wildcard()
|
|
seqstart_k = wildcard()
|
|
max_seqlen_q = wildcard()
|
|
max_seqlen_k = wildcard()
|
|
annotations.update(
|
|
{
|
|
"seqstart_q": seqstart_q,
|
|
"seqstart_k": seqstart_k,
|
|
"max_seqlen_q": max_seqlen_q,
|
|
"max_seqlen_k": max_seqlen_k,
|
|
}
|
|
)
|
|
out = is_op("relax.nn.attention_var_len")(
|
|
query, key, value, seqstart_q, seqstart_k, max_seqlen_q, max_seqlen_k
|
|
)
|
|
else:
|
|
out = is_op("relax.nn.attention")(query, key, value)
|
|
|
|
return out, annotations
|
|
|
|
|
|
def make_stacked_attention_pattern(start_op: str, with_bias: bool = False, layout="BS3NH"):
|
|
"""
|
|
Create pattern for fused multi head attention with stacked input.
|
|
|
|
Parameters
|
|
----------
|
|
start_op: str
|
|
The starting op for pattern, i.e. `R.split` or `R.strided_slice`.
|
|
|
|
with_bias: bool
|
|
Whether or not to include bias addition
|
|
|
|
layout: str
|
|
The layout of the stacked input tensor.
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a fused multi head attention.
|
|
|
|
annotations: Mapping[str, DFPattern]
|
|
A mapping from name to sub pattern. It can be used to extract
|
|
important expressions from match result, to power the partition
|
|
check function and codegen.
|
|
"""
|
|
stacked_qkv = wildcard()
|
|
ops = {}
|
|
if start_op == "split":
|
|
ops["split"] = qkv_tuple = is_op("relax.split")(stacked_qkv)
|
|
query_raw = is_tuple_get_item(qkv_tuple, 0)
|
|
key_raw = is_tuple_get_item(qkv_tuple, 1)
|
|
value_raw = is_tuple_get_item(qkv_tuple, 2)
|
|
elif start_op == "strided_slice":
|
|
ops["strided_slice_query"] = query_raw = is_op("relax.strided_slice")(
|
|
stacked_qkv, varg_default_wildcard=True
|
|
)
|
|
ops["strided_slice_key"] = key_raw = is_op("relax.strided_slice")(
|
|
stacked_qkv, varg_default_wildcard=True
|
|
)
|
|
ops["strided_slice_value"] = value_raw = is_op("relax.strided_slice")(
|
|
stacked_qkv, varg_default_wildcard=True
|
|
)
|
|
else:
|
|
raise NotImplementedError()
|
|
query_reshape_list = wildcard()
|
|
key_reshape_list = wildcard()
|
|
value_reshape_list = wildcard()
|
|
if layout == "BS3NH":
|
|
query = is_op("relax.reshape")(query_raw, query_reshape_list)
|
|
key = is_op("relax.reshape")(key_raw, key_reshape_list)
|
|
value = is_op("relax.reshape")(value_raw, value_reshape_list)
|
|
elif layout == "SBN3H":
|
|
ops["q_transpose"] = query = is_op("relax.permute_dims")(query_raw)
|
|
ops["k_transpose"] = key = is_op("relax.permute_dims")(key_raw)
|
|
ops["v_transpose"] = value = is_op("relax.permute_dims")(value_raw)
|
|
annotations = {
|
|
"stacked_qkv": stacked_qkv,
|
|
"query_reshape_list": query_reshape_list,
|
|
"key_reshape_list": key_reshape_list,
|
|
"value_reshape_list": value_reshape_list,
|
|
**ops,
|
|
}
|
|
if with_bias:
|
|
bias = wildcard()
|
|
annotations["bias"] = bias
|
|
out = is_op("relax.nn.attention_bias")(query, key, value, bias)
|
|
else:
|
|
out = is_op("relax.nn.attention")(query, key, value)
|
|
|
|
if layout == "SBN3H":
|
|
out = is_op("relax.permute_dims")(out)
|
|
|
|
return out, annotations
|
|
|
|
|
|
def make_layer_norm_pattern():
|
|
"""Create a layer norm pattern."""
|
|
inp = wildcard()
|
|
gamma = wildcard()
|
|
beta = wildcard()
|
|
|
|
return is_op("relax.nn.layer_norm")(inp, gamma, beta), {}
|
|
|
|
|
|
def make_rms_norm_pattern():
|
|
"""Create a layer norm pattern."""
|
|
inp = wildcard()
|
|
weight = wildcard()
|
|
gv = GlobalVarPattern()
|
|
out = is_op("relax.call_tir")(gv, TuplePattern([inp, weight]))
|
|
annotations = {"gv": gv, "inp": inp, "rms_norm": out}
|
|
return out, annotations
|
|
|
|
|
|
def make_matmul_dequantize_pattern(
|
|
transposed_rhs: bool = False,
|
|
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
|
"""
|
|
Create pattern for matrix multiplication and dequantize operation.
|
|
|
|
Parameters
|
|
----------
|
|
transposed_rhs: bool
|
|
Whether the right hand side of multiplication is transposed.
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a matrix multiplication.
|
|
|
|
annotations: Mapping[str, DFPattern]
|
|
A mapping from name to sub pattern. It can be used to extract important expressions from
|
|
match result, to power the partition check function and codegen.
|
|
"""
|
|
|
|
lhs = wildcard()
|
|
rhs = wildcard()
|
|
annotations = {"lhs": lhs, "rhs": rhs}
|
|
|
|
if transposed_rhs:
|
|
rhs = is_op("relax.permute_dims")(rhs)
|
|
|
|
out = is_op("relax.matmul")(lhs, rhs)
|
|
annotations["root"] = out
|
|
|
|
scale = is_const()
|
|
zp = is_const()
|
|
annotations.update({"scale": scale, "zp": zp})
|
|
|
|
out = is_op("relax.dequantize")(out, scale, zp)
|
|
|
|
return out, annotations
|
|
|
|
|
|
def make_matmul_multiply_pattern(
|
|
transposed_rhs: bool = False,
|
|
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
|
"""
|
|
Create pattern for matrix multiplication and multiply operation.
|
|
|
|
Parameters
|
|
----------
|
|
transposed_rhs: bool
|
|
Whether the right hand side of multiplication is transposed.
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing a matrix multiplication.
|
|
|
|
annotations: Mapping[str, DFPattern]
|
|
A mapping from name to sub pattern. It can be used to extract important expressions from
|
|
match result, to power the partition check function and codegen.
|
|
"""
|
|
|
|
lhs = wildcard()
|
|
rhs = wildcard()
|
|
scaleA = wildcard()
|
|
scaleB = wildcard()
|
|
annotations = {"lhs": lhs, "rhs": rhs, "scaleA": scaleA, "scaleB": scaleB}
|
|
|
|
if transposed_rhs:
|
|
rhs = is_op("relax.permute_dims")(rhs)
|
|
out = is_op("relax.matmul")(lhs, rhs)
|
|
annotations["root"] = out
|
|
scale = is_op("relax.multiply")(scaleA.has_shape((1,)), scaleB.has_shape((1,)))
|
|
out = is_op("relax.multiply")(out, scale)
|
|
out = is_op("relax.astype")(out)
|
|
|
|
return out, annotations
|
|
|
|
|
|
def make_attention_rewrite_pattern(
|
|
qkv_layout: str, out_layout: str, with_bias: bool, with_cast: bool, with_kv_repeat: bool = False
|
|
):
|
|
"""
|
|
Create pattern for implicit fused multi head attention rewriting.
|
|
|
|
Parameters
|
|
----------
|
|
qkv_layout: str
|
|
The layout of the query, key and value tensor, i.e. BSNH or BSH.
|
|
|
|
out_layout: str
|
|
The layout of the output tensor, i.e. BSNH or BSH.
|
|
|
|
with_bias: bool
|
|
Whether or not to include bias addition.
|
|
|
|
with_cast: bool
|
|
Whether or not rewriting is intended to be applied to a module after the FP16 conversion
|
|
pass.
|
|
|
|
with_kv_repeat: bool
|
|
Whether or not to include the Relax repeat op in the pattern, which is typically used
|
|
in a Relax module to support multi-query attention.
|
|
|
|
Returns
|
|
-------
|
|
pattern: DFPattern
|
|
The resulting pattern describing an implicit fused multi head attention.
|
|
|
|
rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
|
|
The rewriter for the pattern. It will check the matched patterns, and rewrite.
|
|
If the matched pattern is not able to be rewritten to `R.nn.attention`, the rewriter
|
|
returns the original IR.
|
|
"""
|
|
|
|
# pylint: disable=invalid-name
|
|
def handle_input(tensor, layout, transpose, repeat=False):
|
|
if repeat:
|
|
tensor = is_op("relax.repeat")(tensor)
|
|
|
|
if layout == "BSNH":
|
|
permuted = is_op("relax.permute_dims")(tensor)
|
|
shape = wildcard()
|
|
reshaped = is_op("relax.reshape")(permuted, shape)
|
|
if transpose:
|
|
transposed = is_op("relax.permute_dims")(reshaped)
|
|
|
|
def rewriter(matchings, x):
|
|
if matchings[tensor].ty.ndim != 4:
|
|
return None
|
|
if list(matchings[permuted].attrs.axes) != [0, 2, 1, 3]:
|
|
return None
|
|
before_reshape = matchings[permuted].ty.shape.values
|
|
after_reshape = matchings[shape].ty.values
|
|
if not (
|
|
len(before_reshape) == 4
|
|
and len(after_reshape) == 3
|
|
and before_reshape[-2:] == after_reshape[-2:]
|
|
):
|
|
return None
|
|
if transpose and list(matchings[transposed].attrs.axes) != [0, 2, 1]:
|
|
return None
|
|
return x, x.ty.shape
|
|
|
|
if transpose:
|
|
return transposed, rewriter
|
|
else:
|
|
return reshaped, rewriter
|
|
elif layout == "BSH":
|
|
if transpose:
|
|
transposed = is_op("relax.permute_dims")(tensor)
|
|
|
|
def rewriter(matchings, x):
|
|
if matchings[tensor].ty.ndim != 3:
|
|
return None
|
|
if transpose and list(matchings[transposed].attrs.axes) != [0, 2, 1]:
|
|
return None
|
|
before_reshape = x.ty.shape.values
|
|
after_reshape = [before_reshape[0], before_reshape[1], 1, before_reshape[2]]
|
|
return R.reshape(x, after_reshape), after_reshape
|
|
|
|
if transpose:
|
|
return transposed, rewriter
|
|
else:
|
|
return tensor, rewriter
|
|
else:
|
|
raise NotImplementedError()
|
|
|
|
def handle_output(tensor, layout):
|
|
if layout == "BSNH":
|
|
shape = wildcard()
|
|
reshaped = is_op("relax.reshape")(tensor, shape)
|
|
permuted = is_op("relax.permute_dims")(reshaped)
|
|
|
|
def rewriter(matchings, x):
|
|
if matchings[tensor].ty.ndim != 3:
|
|
return None
|
|
before_reshape = matchings[tensor].ty.shape.values
|
|
after_reshape = matchings[shape].ty.values
|
|
if not (
|
|
len(before_reshape) == 3
|
|
and len(after_reshape) == 4
|
|
and before_reshape[-2:] == after_reshape[-2:]
|
|
):
|
|
return None
|
|
if list(matchings[permuted].attrs.axes) != [0, 2, 1, 3]:
|
|
return None
|
|
return x
|
|
|
|
return permuted, rewriter
|
|
elif layout == "BSH":
|
|
|
|
def rewriter(matchings, x):
|
|
if matchings[tensor].ty.ndim != 3:
|
|
return None
|
|
return R.reshape(x, matchings[tensor].ty.shape.values)
|
|
|
|
return tensor, rewriter
|
|
else:
|
|
raise NotImplementedError()
|
|
|
|
q_raw, k_raw, v_raw = wildcard(), wildcard(), wildcard()
|
|
q, q_rewriter = handle_input(q_raw, qkv_layout, False)
|
|
k, k_rewriter = handle_input(k_raw, qkv_layout, True, repeat=with_kv_repeat)
|
|
v, v_rewriter = handle_input(v_raw, qkv_layout, False, repeat=with_kv_repeat)
|
|
matmul_1 = is_op("relax.matmul")(q, k)
|
|
scale = is_const()
|
|
|
|
if with_cast:
|
|
multiply = is_op("relax.multiply")(matmul_1, is_op("relax.astype")(scale))
|
|
else:
|
|
multiply = is_op("relax.multiply")(matmul_1, scale)
|
|
|
|
if with_bias:
|
|
bias_raw = wildcard()
|
|
add = is_op("relax.add")(multiply, bias_raw)
|
|
softmax_input = add
|
|
else:
|
|
softmax_input = multiply
|
|
|
|
if with_cast:
|
|
softmax_input = is_op("relax.astype")(softmax_input)
|
|
|
|
softmax = is_op("relax.nn.softmax")(softmax_input)
|
|
|
|
if with_cast:
|
|
softmax_output = is_op("relax.astype")(softmax)
|
|
else:
|
|
softmax_output = softmax
|
|
|
|
matmul_2 = is_op("relax.matmul")(softmax_output, v)
|
|
|
|
out, out_rewriter = handle_output(matmul_2, out_layout)
|
|
|
|
def rewriter(original, matchings):
|
|
query, query_shape = q_rewriter(matchings, matchings[q_raw])
|
|
key, key_shape = k_rewriter(matchings, matchings[k_raw])
|
|
value, _ = v_rewriter(matchings, matchings[v_raw])
|
|
if query is None or key is None or value is None:
|
|
return original
|
|
softmax_axis = matchings[softmax].attrs.axis
|
|
softmax_input_rank = len(matchings[softmax].ty.shape)
|
|
if softmax_axis == -1:
|
|
softmax_axis += softmax_input_rank
|
|
if softmax_axis != softmax_input_rank - 1:
|
|
return original
|
|
b, s, n, _ = query_shape
|
|
_, s_kv, _, _ = key_shape
|
|
if with_bias:
|
|
bias = matchings[bias_raw]
|
|
bias_shape = list(bias.ty.shape)
|
|
if bias_shape == [b * n, s, s_kv]:
|
|
bias = R.reshape(bias, [b, n, s, s_kv])
|
|
elif bias_shape == [b * n, 1, s_kv]:
|
|
bias = R.reshape(bias, [b, n, 1, s_kv])
|
|
elif bias_shape == [b, s, s_kv]:
|
|
bias = R.reshape(bias, [b, 1, s, s_kv])
|
|
elif bias_shape == [b, 1, s_kv]:
|
|
bias = R.reshape(bias, [b, 1, 1, s_kv])
|
|
elif bias_shape in [[1, s, s_kv], [s, s_kv]]:
|
|
bias = R.reshape(bias, [1, 1, s, s_kv])
|
|
elif bias_shape in [[1, 1, s_kv], [1, s_kv], [s_kv]]:
|
|
bias = R.reshape(bias, [1, 1, 1, s_kv])
|
|
else:
|
|
return original
|
|
else:
|
|
bias = None
|
|
out = out_rewriter(
|
|
matchings,
|
|
R.nn.attention(
|
|
query,
|
|
key,
|
|
value,
|
|
bias,
|
|
T.FloatImm(matchings[scale].data.dtype, float(matchings[scale].data.numpy())),
|
|
),
|
|
)
|
|
return out
|
|
|
|
return out, rewriter
|