chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,174 @@
|
||||
# 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.
|
||||
|
||||
"""A compiler pass that fuses transpose + matmul and generate TIR function.
|
||||
Note that
|
||||
1. Please put the pass before LegalizeOps pass.
|
||||
2. The pass only works for XW^T but not X^TW
|
||||
3. The pass would rewrite the relax ops into TIR functions. If you'd like to dispatch the
|
||||
ops into library (e.g. cuBLAS) calls, please run dispatch pass before this pass.
|
||||
"""
|
||||
|
||||
import tvm
|
||||
from tvm import IRModule, relax, te, tirx
|
||||
from tvm.relax.dpl.pattern import is_op, wildcard
|
||||
from tvm.relax.expr_functor import PyExprMutator, mutator
|
||||
|
||||
|
||||
@tvm.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
|
||||
class FuseTransposeMatmul: # pylint: disable=too-few-public-methods
|
||||
"""A compiler pass that fuses transpose + matmul."""
|
||||
|
||||
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
||||
"""IRModule-level transformation"""
|
||||
mod = relax.transform.FuseOpsByPattern(
|
||||
[
|
||||
(
|
||||
"transpose_matmul_fuse",
|
||||
*_pattern(),
|
||||
),
|
||||
],
|
||||
bind_constants=False,
|
||||
)(mod)
|
||||
transpose_matmul_codegen = _TransposeMatmulFuser(mod)
|
||||
for g_var, func in mod.functions_items():
|
||||
if isinstance(func, relax.Function):
|
||||
func = transpose_matmul_codegen.visit_expr(func)
|
||||
transpose_matmul_codegen.builder_.update_func(g_var, func)
|
||||
return transpose_matmul_codegen.builder_.get()
|
||||
|
||||
|
||||
def _pattern():
|
||||
"""Pattern for transpose + matmul."""
|
||||
# pylint: disable=invalid-name
|
||||
w = wildcard()
|
||||
x = wildcard()
|
||||
wT = is_op("relax.permute_dims")(w)
|
||||
o = is_op("relax.matmul")(x, wT)
|
||||
# pylint: enable=invalid-name
|
||||
annotations = {"o": o, "w": w, "x": x, "wT": wT}
|
||||
|
||||
def _check(context: relax.transform.PatternCheckContext) -> bool:
|
||||
transpose_call = context.annotated_expr["wT"]
|
||||
ndim = transpose_call.args[0].ty.ndim
|
||||
if ndim == -1:
|
||||
return False
|
||||
if ndim == 2 and transpose_call.attrs.axes is None:
|
||||
return True
|
||||
axes = list(range(ndim))
|
||||
axes[-1], axes[-2] = axes[-2], axes[-1]
|
||||
return list(transpose_call.attrs.axes) == axes
|
||||
|
||||
return o, annotations, _check
|
||||
|
||||
|
||||
# pylint: disable=missing-docstring,invalid-name
|
||||
|
||||
|
||||
@mutator
|
||||
class _TransposeMatmulFuser(PyExprMutator): # pylint: disable=abstract-method
|
||||
def __init__(self, mod):
|
||||
super().__init__(mod)
|
||||
|
||||
def visit_call_( # pylint: disable=arguments-renamed
|
||||
self,
|
||||
call: relax.Call,
|
||||
) -> relax.Expr:
|
||||
out_dtype = None
|
||||
|
||||
def te_transposed_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
|
||||
nonlocal out_dtype
|
||||
a_shape = list(a.shape)
|
||||
b_shape = list(b.shape)
|
||||
a_prepended = False
|
||||
b_appended = False
|
||||
if len(a_shape) == 1:
|
||||
a_prepended = True
|
||||
a_shape.insert(0, 1)
|
||||
if len(b_shape) == 1:
|
||||
b_appended = True
|
||||
b_shape.append(1)
|
||||
|
||||
is_a_larger = len(a_shape) > len(b_shape)
|
||||
offset = len(a_shape) - len(b_shape) if is_a_larger else len(b_shape) - len(a_shape)
|
||||
|
||||
a_relax = relax.Var("a", relax.TensorType(a.shape))
|
||||
bT_shape = list(b.shape)
|
||||
bT_shape[-1], bT_shape[-2] = bT_shape[-2], bT_shape[-1]
|
||||
bT_relax = relax.Var("b", relax.TensorType(bT_shape))
|
||||
output_shape = self.builder_.normalize(relax.op.matmul(a_relax, bT_relax)).ty.shape
|
||||
|
||||
def matmul_compute(*idx_spatial):
|
||||
k = te.reduce_axis((0, a_shape[-1]), name="k")
|
||||
|
||||
def multiply_compute(idx_reduce):
|
||||
a_indices = []
|
||||
b_indices = []
|
||||
|
||||
for i in range(offset):
|
||||
if is_a_larger:
|
||||
a_indices.append(idx_spatial[i])
|
||||
else:
|
||||
b_indices.append(idx_spatial[i])
|
||||
for i in range(offset, len(output_shape) - (2 - a_prepended - b_appended)):
|
||||
a_dim = a_shape[i if is_a_larger else i - offset]
|
||||
b_dim = b_shape[i if not is_a_larger else i - offset]
|
||||
dim_equal = a_dim == b_dim
|
||||
if not isinstance(dim_equal, tirx.IntImm) or dim_equal == 0:
|
||||
a_dim_is_one = isinstance(a_dim, tirx.IntImm) and a_dim == 1
|
||||
b_dim_is_one = isinstance(b_dim, tirx.IntImm) and b_dim == 1
|
||||
a_indices.append(0 if a_dim_is_one else idx_spatial[i])
|
||||
b_indices.append(0 if b_dim_is_one else idx_spatial[i])
|
||||
else:
|
||||
a_indices.append(idx_spatial[i])
|
||||
b_indices.append(idx_spatial[i])
|
||||
|
||||
if not a_prepended:
|
||||
a_indices.append(idx_spatial[-2 + b_appended])
|
||||
a_indices.append(idx_reduce)
|
||||
if not b_appended:
|
||||
b_indices.append(idx_spatial[-1])
|
||||
b_indices.append(idx_reduce)
|
||||
|
||||
dtype = out_dtype
|
||||
if dtype != "":
|
||||
return a(*a_indices).astype(dtype) * b(*b_indices).astype(dtype)
|
||||
return a(*a_indices) * b(*b_indices)
|
||||
|
||||
return te.sum(multiply_compute(k), axis=k)
|
||||
|
||||
return te.compute(
|
||||
output_shape,
|
||||
lambda *idx: matmul_compute(*idx), # pylint: disable=unnecessary-lambda
|
||||
name="NT_matmul",
|
||||
)
|
||||
|
||||
if isinstance(call.op, relax.GlobalVar):
|
||||
function = self.builder_.get()[call.op]
|
||||
if (
|
||||
"Composite" in function.attrs
|
||||
and function.attrs["Composite"] == "transpose_matmul_fuse"
|
||||
):
|
||||
out_dtype = function.ret_ty.dtype
|
||||
return self.builder_.call_te(
|
||||
te_transposed_matmul,
|
||||
call.args[1],
|
||||
call.args[0],
|
||||
primfunc_name_hint="NT_matmul",
|
||||
)
|
||||
|
||||
return super().visit_call_(call)
|
||||
Reference in New Issue
Block a user