(mod: tvm.IRModule)
| 26 | |
| 27 | # pylint: disable=import-outside-toplevel,missing-function-docstring |
| 28 | def reshape_matmul(mod: tvm.IRModule): |
| 29 | from tvm.relax import Expr |
| 30 | from tvm.relax.dpl import DFPattern, rewrite_call |
| 31 | from tvm.relax.dpl.pattern import is_op, wildcard |
| 32 | |
| 33 | input0 = wildcard() |
| 34 | input1 = wildcard() |
| 35 | pattern = is_op("relax.matmul")(input0, input1) |
| 36 | |
| 37 | def _rewriter(expr: Expr, matches: dict[DFPattern, Expr]): |
| 38 | i0 = matches[input0] |
| 39 | i1 = matches[input1] |
| 40 | if len(i0.struct_info.shape) == 2 and len(i1.struct_info.shape) == 2: |
| 41 | i0_shape = [1] + [*i0.struct_info.shape.values] |
| 42 | i1_shape = [1] + [*i1.struct_info.shape.values] |
| 43 | oshape = matches[pattern].struct_info.shape |
| 44 | return R.reshape(R.matmul(R.reshape(i0, i0_shape), R.reshape(i1, i1_shape)), oshape) |
| 45 | return expr |
| 46 | |
| 47 | mod["main"] = rewrite_call(pattern, _rewriter, mod["main"]) |
| 48 | return mod |
| 49 | |
| 50 | |
| 51 | def decompose_clip(mod: tvm.IRModule) -> tvm.IRModule: |
no test coverage detected
searching dependent graphs…