# 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)