|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +# pyre-strict |
| 8 | + |
| 9 | +from typing import List, Optional |
| 10 | + |
| 11 | +import torch |
| 12 | + |
| 13 | +from executorch.backends.transforms.utils import ( |
| 14 | + get_param_tensor, |
| 15 | + is_param_node, |
| 16 | + set_param_tensor, |
| 17 | +) |
| 18 | +from executorch.exir.dialects._ops import ops as exir_ops |
| 19 | +from executorch.exir.pass_base import ExportPass, PassResult |
| 20 | + |
| 21 | +from torch.export import ExportedProgram |
| 22 | + |
| 23 | +# Mirrors kIm2colMinCOut in Convolution.cpp. At or above this output channel |
| 24 | +# count, should_use_conv2d_im2col() selects the im2col + GEMM path on every |
| 25 | +# vendor, so a rewritten conv is guaranteed to land there rather than on the |
| 26 | +# direct conv2d shader, which has not been compared against conv1d. |
| 27 | +_IM2COL_MIN_C_OUT = 128 |
| 28 | + |
| 29 | + |
| 30 | +def _reshape_fake(val: torch.Tensor, sizes: List[int]) -> torch.Tensor: |
| 31 | + fake_mode = getattr(val, "fake_mode", None) |
| 32 | + if fake_mode is None: |
| 33 | + return val.reshape(sizes) |
| 34 | + with fake_mode: |
| 35 | + return val.reshape(sizes) |
| 36 | + |
| 37 | + |
| 38 | +class Conv1dAsConv2dPass(ExportPass): |
| 39 | + """ |
| 40 | + Rewrite eligible 1-D convolutions into a 2-D convolution over a singleton |
| 41 | + height dimension. |
| 42 | +
|
| 43 | + conv1d.glsl computes one output element per invocation with no tiling and no |
| 44 | + reuse between invocations, and its global work grid packs texels along the |
| 45 | + batch dim, wasting three of every four lanes at batch 1. It measures around |
| 46 | + 30x slower than the conv2d im2col + GEMM path for the same MAC count. |
| 47 | + Expressing the convolution as 2-D lets the existing conv2d machinery |
| 48 | + (im2col selection, memory layout tagging, weight prepacking) handle it with |
| 49 | + no runtime changes. |
| 50 | +
|
| 51 | + Only convolutions that are certain to reach the im2col path are rewritten, |
| 52 | + so nothing silently moves onto a shader that has not been compared against |
| 53 | + conv1d. The weight is reshaped in place rather than through a view node |
| 54 | + because the conv2d implementation reads it as a constant TensorRef. |
| 55 | + """ |
| 56 | + |
| 57 | + def __init__(self) -> None: |
| 58 | + super().__init__() |
| 59 | + self._exported_program: Optional[ExportedProgram] = None |
| 60 | + |
| 61 | + def _eligible(self, node: torch.fx.Node) -> bool: |
| 62 | + assert self._exported_program is not None |
| 63 | + in_node, weight_node = node.args[0], node.args[1] |
| 64 | + if not isinstance(in_node, torch.fx.Node): |
| 65 | + return False |
| 66 | + if not isinstance(weight_node, torch.fx.Node): |
| 67 | + return False |
| 68 | + in_shape = in_node.meta["val"].shape |
| 69 | + w_shape = weight_node.meta["val"].shape |
| 70 | + if len(in_shape) != 3 or len(w_shape) != 3: |
| 71 | + return False |
| 72 | + # The weight has to be a constant so it can be reshaped in place. |
| 73 | + if not is_param_node(self._exported_program, weight_node): |
| 74 | + return False |
| 75 | + transposed, groups = node.args[6], node.args[8] |
| 76 | + if transposed or groups != 1: |
| 77 | + return False |
| 78 | + # conv2d does not support batched input at all, so a rewrite there |
| 79 | + # would turn a working conv1d into a throw. |
| 80 | + if in_shape[0] != 1: |
| 81 | + return False |
| 82 | + # Depthwise and pointwise 1-D convs keep their own shaders. |
| 83 | + if w_shape[0] < _IM2COL_MIN_C_OUT or w_shape[2] == 1: |
| 84 | + return False |
| 85 | + # im2col requires unit dilation. |
| 86 | + return list(node.args[5]) == [1] |
| 87 | + |
| 88 | + def _rewrite(self, graph: torch.fx.Graph, node: torch.fx.Node) -> None: |
| 89 | + assert self._exported_program is not None |
| 90 | + view = exir_ops.edge.aten.view_copy.default |
| 91 | + in_node, weight_node = node.args[0], node.args[1] |
| 92 | + |
| 93 | + in_shape = list(in_node.meta["val"].shape) |
| 94 | + w_shape = list(weight_node.meta["val"].shape) |
| 95 | + out_shape = list(node.meta["val"].shape) |
| 96 | + in_4d_shape = [in_shape[0], in_shape[1], 1, in_shape[2]] |
| 97 | + w_4d_shape = [w_shape[0], w_shape[1], 1, w_shape[2]] |
| 98 | + out_4d_shape = [out_shape[0], out_shape[1], 1, out_shape[2]] |
| 99 | + |
| 100 | + # The weight tensor is contiguous, so inserting a singleton dim is a |
| 101 | + # pure metadata change. A weight shared by several convs is only |
| 102 | + # reshaped once; the rank check in _eligible() skips it afterwards. |
| 103 | + weight = get_param_tensor(self._exported_program, weight_node) |
| 104 | + assert weight is not None |
| 105 | + set_param_tensor( |
| 106 | + self._exported_program, weight_node, weight.reshape(w_4d_shape) |
| 107 | + ) |
| 108 | + weight_node.meta["val"] = _reshape_fake( |
| 109 | + weight_node.meta["val"], w_4d_shape |
| 110 | + ) |
| 111 | + |
| 112 | + with graph.inserting_before(node): |
| 113 | + in_4d = graph.call_function(view, args=(in_node, in_4d_shape)) |
| 114 | + in_4d.meta = dict(in_node.meta) |
| 115 | + in_4d.meta["val"] = _reshape_fake(in_node.meta["val"], in_4d_shape) |
| 116 | + |
| 117 | + def pad(arg: List[int], lead: int) -> List[int]: |
| 118 | + return [lead, arg[0]] |
| 119 | + |
| 120 | + node.args = ( |
| 121 | + in_4d, |
| 122 | + weight_node, |
| 123 | + node.args[2], |
| 124 | + pad(node.args[3], 1), # stride |
| 125 | + pad(node.args[4], 0), # padding |
| 126 | + pad(node.args[5], 1), # dilation |
| 127 | + node.args[6], |
| 128 | + pad(node.args[7], 0), # output_padding |
| 129 | + node.args[8], |
| 130 | + ) |
| 131 | + |
| 132 | + with graph.inserting_after(node): |
| 133 | + out_3d = graph.call_function(view, args=(node, out_shape)) |
| 134 | + out_3d.meta = dict(node.meta) |
| 135 | + node.replace_all_uses_with(out_3d) |
| 136 | + out_3d.args = (node, out_shape) |
| 137 | + node.meta["val"] = _reshape_fake(node.meta["val"], out_4d_shape) |
| 138 | + |
| 139 | + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: |
| 140 | + assert self._exported_program is not None |
| 141 | + |
| 142 | + modified = False |
| 143 | + for node in list(graph_module.graph.nodes): |
| 144 | + if node.op != "call_function": |
| 145 | + continue |
| 146 | + if node.target != exir_ops.edge.aten.convolution.default: |
| 147 | + continue |
| 148 | + if not self._eligible(node): |
| 149 | + continue |
| 150 | + self._rewrite(graph_module.graph, node) |
| 151 | + modified = True |
| 152 | + |
| 153 | + if not modified: |
| 154 | + return PassResult(graph_module, False) |
| 155 | + |
| 156 | + graph_module.recompile() |
| 157 | + return PassResult(graph_module, True) |
0 commit comments