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 7from typing import List, Optional 8 9import executorch.exir as exir 10from executorch.exir import CaptureConfig 11from executorch.exir.pass_manager import PassType 12 13 14### XNNPACK Configs ### 15def get_xnnpack_edge_compile_config( 16 skip_dim_order: bool = False, 17) -> exir.EdgeCompileConfig: 18 return exir.EdgeCompileConfig( 19 _check_ir_validity=False, _skip_dim_order=skip_dim_order 20 ) 21 22 23def get_transform_passes(additional_passes=None) -> List[PassType]: 24 passes = additional_passes if additional_passes else [] 25 return passes 26 27 28def get_xnnpack_executorch_backend_config( 29 additional_passes=None, 30) -> exir.ExecutorchBackendConfig: 31 additional_passes = additional_passes if additional_passes else [] 32 return exir.ExecutorchBackendConfig( 33 passes=additional_passes, 34 extract_delegate_segments=True, 35 ) 36 37 38def get_xnnpack_capture_config( 39 dynamic_shape=False, 40 enable_aot: Optional[bool] = None, 41 unlift: Optional[bool] = None, 42): 43 if enable_aot is None: 44 return CaptureConfig(enable_dynamic_shape=dynamic_shape) 45 else: 46 unlift = unlift if unlift is not None else enable_aot 47 return CaptureConfig( 48 enable_dynamic_shape=dynamic_shape, enable_aot=enable_aot, _unlift=unlift 49 ) 50