diff --git a/backends/arm/common/arm_compile_spec.py b/backends/arm/common/arm_compile_spec.py index 29037ec833a..a33531e2411 100644 --- a/backends/arm/common/arm_compile_spec.py +++ b/backends/arm/common/arm_compile_spec.py @@ -173,18 +173,30 @@ def to_list(self): return compile_spec def get_intermediate_path(self) -> str | None: + """ + Gets the path used for dumping intermediate results such as tosa and pte. + + Returns: + Path where intermediate results are saved. + """ return self.path_for_intermediates def dump_intermediate_artifacts_to(self, output_path: str | None): """ Sets a path for dumping intermediate results during such as tosa and pte. + + Args: + output_path: Path to dump intermediate results to. """ self.path_for_intermediates = output_path return self def dump_debug_info(self, debug_mode: DebugMode | None): """ - Dump debugging information into the intermediates path + Dump debugging information into the intermediates path. + + Args: + debug_mode: The debug mode to use for dumping debug information. """ self.tosa_debug_mode = debug_mode return self diff --git a/backends/arm/ethosu/compile_spec.py b/backends/arm/ethosu/compile_spec.py index 5f3f92fdd0e..9b6289156fa 100644 --- a/backends/arm/ethosu/compile_spec.py +++ b/backends/arm/ethosu/compile_spec.py @@ -15,6 +15,16 @@ class EthosUCompileSpec(ArmCompileSpec): + """ + Compile spec for Ethos-U NPU. + + Args: + target: Ethos-U accelerator configuration, e.g. ethos-u55-128. + system_config: System configuration to select from the Vela configuration file. + memory_mode: Memory mode to select from the Vela configuration file. + extra_flags: Extra flags for the Vela compiler. + config_ini: Vela configuration file(s) in Python ConfigParser .ini file format. + """ _TARGET_KEY = "target" @@ -26,16 +36,6 @@ def __init__( extra_flags: list[str] | None = None, config_ini: str | None = "Arm/vela.ini", ): - """Generate compile spec for Ethos-U NPU - Args: - target: Ethos-U accelerator configuration, e.g. ethos-u55-128 - system_config: System configuration to select from the Vela - configuration file - memory_mode: Memory mode to select from the Vela configuration file - extra_flags: Extra flags for the Vela compiler - config_ini: Vela configuration file(s) in Python ConfigParser .ini - file format - """ self.target = target # Set vela compiler flags @@ -87,6 +87,7 @@ def from_list_hook(cls, compile_spec, specs: dict[str, str]): compile_spec.target = specs.get(cls._TARGET_KEY, None) def validate(self): + """Throws an error if the compile spec is not valid.""" if len(self.compiler_flags) == 0: raise ValueError( "compile_flags are required in the CompileSpec list for EthosUBackend" diff --git a/backends/arm/ethosu/partitioner.py b/backends/arm/ethosu/partitioner.py index 33ac6a1db62..9acc0439171 100644 --- a/backends/arm/ethosu/partitioner.py +++ b/backends/arm/ethosu/partitioner.py @@ -17,7 +17,7 @@ class EthosUPartitioner(TOSAPartitioner): """ Partitions subgraphs supported by the Arm Ethos-U backend. - Attributes: + Args: compile_spec: List of CompileSpec objects for Ethos-U backend. additional_checks: Optional sequence of additional operator support checks. """ diff --git a/backends/arm/quantizer/arm_quantizer.py b/backends/arm/quantizer/arm_quantizer.py index cad714607e2..e52b30895dc 100644 --- a/backends/arm/quantizer/arm_quantizer.py +++ b/backends/arm/quantizer/arm_quantizer.py @@ -333,16 +333,26 @@ def __init__( self.module_name_config: Dict[str, Optional[QuantizationConfig]] = {} def set_global(self, quantization_config: QuantizationConfig) -> TOSAQuantizer: - """Set quantization_config for submodules that are not already annotated by name or type filters.""" + """ + Set quantization_config for submodules that are not already annotated by name or type filters. + + Args: + quantization_config: The QuantizationConfig to set as global configuration. + """ self.global_config = quantization_config return self def set_module_type( self, module_type: Callable, quantization_config: QuantizationConfig ) -> TOSAQuantizer: - """Set quantization_config for a submodule with type: `module_type`, for example: + """ + Set quantization_config for a submodule with type: `module_type`, for example: quantizer.set_module_name(Sub) or quantizer.set_module_name(nn.Linear), it will quantize all supported operator/operator - patterns in the submodule with this module type with the given `quantization_config` + patterns in the submodule with this module type with the given `quantization_config`. + + Args: + module_type: The type of the submodule to set the quantization config for. + quantization_config: The QuantizationConfig to set for the submodule. """ self.module_type_config[module_type] = quantization_config return self @@ -350,9 +360,14 @@ def set_module_type( def set_module_name( self, module_name: str, quantization_config: Optional[QuantizationConfig] ) -> TOSAQuantizer: - """Set quantization_config for a submodule with name: `module_name`, for example: + """ + Set quantization_config for a submodule with name: `module_name`, for example: quantizer.set_module_name("blocks.sub"), it will quantize all supported operator/operator patterns in the submodule with this module name with the given `quantization_config` + + Args: + module_name: The name of the submodule to set the quantization config for. + quantization_config: The QuantizationConfig to set for the submodule. """ # Validate that quantization_config is provided if quantization_config is None: @@ -360,14 +375,25 @@ def set_module_name( self.module_name_config[module_name] = quantization_config return self - def set_io(self, quantization_config): - """Set quantization_config for input and output nodes.""" + def set_io(self, quantization_config: QuantizationConfig) -> TOSAQuantizer: + """ + Set quantization_config for input and output nodes. + + Args: + quantization_config: The QuantizationConfig to set for input and output nodes. + """ self.io_config = quantization_config return self def transform_for_annotation(self, model: GraphModule) -> GraphModule: - """An initial pass for transforming the graph to prepare it for annotation. + """ + An initial pass for transforming the graph to prepare it for annotation. Currently transforms scalar values to tensor attributes. + + Args: + model: The model to transform. + Returns: + The transformed model. """ # TODO: Fix the need to lazily import this. @@ -465,10 +491,24 @@ def validate(self, model: GraphModule) -> None: class EthosUQuantizer(TOSAQuantizer): + """ + Quantizer supported by the Arm Ethos-U backend. + + Args: + compile_spec: A EthosUCompileSpec instance. + """ + def __init__(self, compile_spec: EthosUCompileSpec) -> None: super().__init__(compile_spec) class VgfQuantizer(TOSAQuantizer): + """ + Quantizer supported by the Arm Vgf backend. + + Args: + compile_spec: A VgfCompileSpec instance. + """ + def __init__(self, compile_spec: VgfCompileSpec) -> None: super().__init__(compile_spec) diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 850e607bef2..83294369ae7 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -193,7 +193,7 @@ def _tag_module( # noqa """Tag nodes in a module, possibly a submodule, from the containing program. Args: - module: a GraphModule from `containing_program` to tag nodes in. + module: A GraphModule from `containing_program` to tag nodes in. containing_program: The ExportedProgram that contains the module. reporter: A reporter to report why nodes were rejected. Returns: diff --git a/backends/arm/vgf/compile_spec.py b/backends/arm/vgf/compile_spec.py index 452ea5c1956..860c8bf3687 100644 --- a/backends/arm/vgf/compile_spec.py +++ b/backends/arm/vgf/compile_spec.py @@ -15,19 +15,19 @@ class VgfCompileSpec(ArmCompileSpec): + """ + Compile spec for VGF compatible targets. + + Args: + tosa_spec: TOSA specification that should be targeted. + compiler_flags: Extra compiler flags for converter_backend. + """ def __init__( self, tosa_spec: TosaSpecification | str | None = None, compiler_flags: list[str] | None = None, ): - """ - Generate compile spec for VGF compatible targets - - Args: - compiler_flags: Extra compiler flags for converter_backend - """ - if tosa_spec is None: tosa_spec = "TOSA-1.0+FP" if isinstance(tosa_spec, str): diff --git a/backends/arm/vgf/partitioner.py b/backends/arm/vgf/partitioner.py index be684505d2b..96c4408b922 100644 --- a/backends/arm/vgf/partitioner.py +++ b/backends/arm/vgf/partitioner.py @@ -14,6 +14,14 @@ @final class VgfPartitioner(TOSAPartitioner): + """ + Partitions subgraphs supported by the Arm Vgf backend. + + Args: + compile_spec: The Vgf compilation specification. + additional_checks: Optional sequence of additional operator support checks. + """ + def __init__( self, compile_spec: VgfCompileSpec,