def create( cls, use_implicit_batch_dim: bool, min_acc_module_size: int = 20, ): return Splitter( min_acc_module_size=min_acc_module_size, operator_supported=create_trt_operator_support( use_implicit_batch_dim), )
def test_unsupport_node_implicit_batch_dim(self): class TestModule(nn.Module): def forward(self, x): y = torch.add(input=x, other=x) return nn.functional.gelu(y) mod = TestModule() traced_mod = acc_tracer.trace(mod, torch.randn(5, 2)) op_support = create_trt_operator_support(use_implicit_batch_dim=True) for node in traced_mod.graph.nodes: if node.target == acc_ops.add: self.assertTrue(op_support.is_node_supported(mod, node)) elif node.target == acc_ops.gelu: self.assertFalse(op_support.is_node_supported(mod, node))
def test_support_node_with_int_attr(self): class TestModule(nn.Module): def forward(self, x): zeros = torch.randint(3, 5, (1, )) zeros = zeros.to(torch.int64) scale = torch.randn(1) return torch.quantize_per_tensor(x, scale, zeros, torch.quint8) mod = TestModule() traced_mod = acc_tracer.trace(mod, torch.randn(5, 2)) op_support = create_trt_operator_support(use_implicit_batch_dim=True) for node in traced_mod.graph.nodes: if node.target == acc_ops.quantize_per_tensor: self.assertTrue(op_support.is_node_supported(mod, node))
def test_supported_node_target(self): class TestModule(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(1, 1) def forward(self, x): x = self.linear(x) x = x + 1 return torch.add(input=x, other=x) mod = TestModule() traced_mod = acc_tracer.trace(mod, torch.randn(1, 2, 1, 1)) op_support = create_trt_operator_support() for node in traced_mod.graph.nodes: self.assertTrue(op_support.is_node_supported(mod, node))
class Splitter(SplitFunc): """A composable fx2trtr splitter. See `SplitFunc`. Attributes: min_acc_module_size: minimum split module size """ _INPUT_ATTR: t.ClassVar[str] = "_split_graph_recorded_input" min_acc_module_size: int operator_supported: OperatorSupportBase = dc.field( default_factory=lambda: create_trt_operator_support()) @classmethod def create( cls, use_implicit_batch_dim: bool, min_acc_module_size: int = 20, ): return Splitter( min_acc_module_size=min_acc_module_size, operator_supported=create_trt_operator_support( use_implicit_batch_dim), ) def __call__(self, module, input) -> Tuple[fx.GraphModule, t.Sequence[SplitInfo]]: trt_split_result = self._trt_split(module, input) logger.debug(f"""TRT split result graph >>> { trt_split_result.graph }""") Splitter._propagate_split_inputs( trt_split_result, input, dict(trt_split_result.named_children()).keys()) return ( trt_split_result, [ Splitter._create_split_info(name, subgraph, parent=trt_split_result) for name, subgraph in trt_split_result.named_children() ], ) @classmethod def _propagate_split_inputs( cls, graph: fx.GraphModule, input: Input, target_modules: t.Collection[str], ) -> None: """ Input propagation on subnets TODO: refactor so we don't set inputs onto the subgraphs """ handles = [] def pre_forward(mod, input): setattr(mod, cls._INPUT_ATTR, input) def _install_hook(g): nonlocal handles if not g: return for _n, _g in g.named_children(): if _n in target_modules: handles.append(_g.register_forward_pre_hook(pre_forward)) _install_hook(_g) try: _install_hook(graph) graph(*input) finally: for h in handles: h.remove() @classmethod def _create_split_info(cls, name, graph, parent) -> SplitInfo: device, order = cls._parse_splitter_subgraph_name(name) input = getattr(graph, cls._INPUT_ATTR) delattr(graph, cls._INPUT_ATTR) return SplitInfo( module=graph, input=input, name=name, device=device, order=order, ) @classmethod def _parse_splitter_subgraph_name(cls, name: str) -> t.Tuple[str, int]: match = re.match("_run_on_([a-z]+)_([0-9]+)", name) assert match, f"{name} doesn't comform with splitter subgraph naming convention" return (match[1], int(match[2])) def _trt_split(self, graph: fx.GraphModule, input: Input) -> fx.GraphModule: splitter_settings = _SplitterSettingBase() splitter_settings.min_acc_module_size = self.min_acc_module_size splitter = TRTSplitter( graph, input, # type: ignore[arg-type] self.operator_supported, settings=splitter_settings, ) logger.info(f"""{splitter.node_support_preview.__name__}: { splitter.node_support_preview() }""") return splitter()