# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import logging from collections import OrderedDict from copy import deepcopy from enum import auto, Enum from typing import Any, List, Optional, Tuple import executorch.backends.vulkan.utils as utils import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.backends.vulkan.vulkan_preprocess import VulkanBackend from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( get_symmetric_quantization_config, XNNPACKQuantizer, ) from executorch.backends.xnnpack.xnnpack_preprocess import XnnpackBackend from executorch.devtools import BundledProgram from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite from executorch.devtools.bundled_program.serialize import ( serialize_from_bundled_program_to_flatbuffer, ) from executorch.exir import ExecutorchProgramManager, to_edge_transform_and_lower from executorch.exir.backend.backend_api import _get_node_list_with_same_tag from executorch.exir.backend.partitioner import ( DelegationSpec, Partitioner, PartitionResult, ) from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer from executorch.exir.lowered_backend_module import ( create_exported_program_from_submodule, create_submodule_from_nodes, ) from executorch.extension.pybindings.portable_lib import ( # @manual _load_for_executorch_from_buffer, ) from executorch.extension.pytree import tree_flatten from torch.export import export from torch.export.exported_program import ExportedProgram from torch.export.graph_signature import InputKind from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import OperatorSupportBase from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e class NodeFlagIsSetChecker(OperatorSupportBase): """ Check if a node is marked with a given field in node.meta["custom"] """ def __init__(self, field: str) -> None: super().__init__() self.field = field def check_field(self, node: torch.fx.Node) -> bool: if "custom" not in node.meta: return False custom_map = node.meta["custom"] if self.field not in custom_map: return False return custom_map[self.field] def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: if node.op == "placeholder" or node.op == "output": return False # Check if the node itself is tagged if self.check_field(node): return True # Check if any direct user of this node is tagged for user in node.users: if self.check_field(user): return True return False class FlagBasedPartitioner(Partitioner): """ Partitioner that partitions based on whether node.meta["custom"][field] is set to True. """ def __init__(self, field: str) -> None: super().__init__() self.field = field self.delegation_spec = DelegationSpec("custom_partition", []) def partition(self, exported_program: ExportedProgram) -> PartitionResult: capability_partitioner = CapabilityBasedPartitioner( exported_program.graph_module, NodeFlagIsSetChecker(self.field), allows_single_node_partition=True, ) partition_list = capability_partitioner.propose_partitions() partition_tags = {} for partition in partition_list: for node in partition.nodes: tag = f"tag{partition.id}" node.meta["delegation_tag"] = tag partition_tags[tag] = self.delegation_spec tag_constant_data(exported_program) tag_mutated_buffer(exported_program) return PartitionResult( tagged_exported_program=exported_program, partition_tags=partition_tags ) def mark_node_range( graph_module: torch.fx.GraphModule, end_idx: int = (2**31 - 1), start_idx: int = 0, field: str = "_in_target_subgraph", ): call_fn_count = 0 for node in graph_module.graph.nodes: if "custom" not in node.meta: node.meta["custom"] = {} node.meta["custom"][field] = False if node.op != "call_function": continue call_fn_count += 1 if call_fn_count >= start_idx and call_fn_count < end_idx: node.meta["custom"][field] = True def extract_submodule_program( tagged_graph_module: torch.fx.GraphModule, owning_program: ExportedProgram, field: str = "_in_target_subgraph", ) -> ExportedProgram: tagged_graph_module_output_node = tagged_graph_module.graph.output_node() partitioner = FlagBasedPartitioner(field) partition_result = partitioner.partition(owning_program) tag, delegation_spec = next(iter(partition_result.partition_tags.items())) node_list = _get_node_list_with_same_tag(tagged_graph_module, tag, owning_program) replace_ctx = tagged_graph_module._set_replace_hook( owning_program.graph_signature.get_replace_hook() ) with replace_ctx: submodule, call_module_node = create_submodule_from_nodes( tagged_graph_module, node_list, tag ) submodule_output_node = submodule.graph.output_node() # Copy the output node meta from the original output node, because # create_submodule_from_nodes doesn't cover the meta field submodule_output_node.meta = tagged_graph_module_output_node.meta ( submodule_program, _, _, ) = create_exported_program_from_submodule( submodule, owning_program, tag, call_module_node, False, ) return submodule_program class QuantizationMode(Enum): """Enum to describe how a model should be quantized.""" NONE = auto() INT8_STATIC_PER_CHANNEL = auto() def get_exported_graph( model, sample_inputs, sample_kwargs=None, dynamic_shapes=None, qmode=QuantizationMode.NONE, ) -> torch.fx.GraphModule: export_training_graph = export( model, sample_inputs, kwargs=sample_kwargs, dynamic_shapes=dynamic_shapes, strict=True, ).module() if qmode == QuantizationMode.NONE: return export_training_graph quantizer = XNNPACKQuantizer() operator_config = get_symmetric_quantization_config(is_per_channel=True) quantizer.set_global(operator_config) prepared_graph = prepare_pt2e(export_training_graph, quantizer) prepared_graph(*sample_inputs) converted_graph = convert_pt2e(prepared_graph) return converted_graph def random_uniform_tensor(shape, low=0.0, high=1.0, device=None, dtype=None): if dtype is None: dtype = torch.float32 # Handle integer types using randint if dtype in ( torch.int, torch.int8, torch.int16, torch.int32, torch.int64, torch.long, torch.short, ): low_int = int(low) high_int = int(high) # randint requires high > low, so ensure at least a range of 1 if high_int <= low_int: high_int = low_int + 1 return torch.randint(low_int, high_int, shape, device=device, dtype=dtype) # Handle unsigned integer types if dtype in (torch.uint8,): low_int = max(0, int(low)) high_int = int(high) if high_int <= low_int: high_int = low_int + 1 return torch.randint(low_int, high_int, shape, device=device, dtype=dtype) # Handle boolean type if dtype == torch.bool: return torch.randint(0, 2, shape, device=device, dtype=torch.int8).bool() # Handle floating-point types (float16, float32, float64, bfloat16) return torch.empty(shape, device=device, dtype=dtype).uniform_(low, high) def generate_sample_inputs( exported_program: ExportedProgram, low: float = -1.0, high: float = 1.0, ) -> Tuple[torch.Tensor, ...]: """ Analyze the exported program graph to determine input shapes and dtypes, then generate random sample inputs. Uses the graph signature to identify only user inputs (excluding parameters, buffers, and other non-input placeholders). Args: exported_program: The exported program to analyze low: Lower bound for random uniform values (default: -1.0) high: Upper bound for random uniform values (default: 1.0) Returns: Tuple of randomly generated tensors matching the input specs """ sample_inputs = [] # Get the set of user input names by filtering input_specs for USER_INPUT kind user_input_names = set() for spec in exported_program.graph_signature.input_specs: if spec.kind == InputKind.USER_INPUT: if hasattr(spec.arg, "name"): user_input_names.add(spec.arg.name) for node in exported_program.graph.nodes: if node.op != "placeholder": continue # Only process nodes that are user inputs (not parameters, buffers, etc.) if node.name not in user_input_names: continue if "val" in node.meta: val = node.meta["val"] shape = None dtype = None if isinstance(val, torch.Tensor): shape = tuple(val.shape) dtype = val.dtype elif hasattr(val, "shape") and hasattr(val, "dtype"): # Handle FakeTensor or similar shape = tuple(val.shape) dtype = val.dtype if shape is not None and dtype is not None: tensor = random_uniform_tensor(shape, low=low, high=high, dtype=dtype) sample_inputs.append(tensor) inputs_flattened, _ = tree_flatten(sample_inputs) return inputs_flattened def export_model_to_vulkan( model, sample_inputs, sample_kwargs=None, dynamic_shapes=None, operator_blocklist=None, operator_allowlist=None, nn_module_blocklist=None, nn_module_allowlist=None, qmode=QuantizationMode.NONE, ): compile_options = {} exported_graph = get_exported_graph( model, sample_inputs, sample_kwargs=sample_kwargs, dynamic_shapes=dynamic_shapes, qmode=qmode, ) program = export( exported_graph, sample_inputs, kwargs=sample_kwargs, dynamic_shapes=dynamic_shapes, strict=True, ) edge_program = to_edge_transform_and_lower( program, partitioner=[ VulkanPartitioner( compile_options, operator_blocklist=operator_blocklist, operator_allowlist=operator_allowlist, nn_module_blocklist=nn_module_blocklist, nn_module_allowlist=nn_module_allowlist, ) ], transform_passes=None, compile_config=None, ) executorch_program = edge_program.to_executorch() # Check if the delegate ID matches VulkanBackend if ( executorch_program.executorch_program.execution_plan[0].delegates[0].id != VulkanBackend.__name__ ): raise RuntimeError( f"Expected delegate ID {VulkanBackend.__name__}, but got {executorch_program.executorch_program.execution_plan[0].delegates[0].id}" ) return executorch_program def export_model_to_xnnpack( model, sample_inputs, dynamic_shapes=None, operator_blocklist=None, operator_allowlist=None, nn_module_blocklist=None, nn_module_allowlist=None, qmode=QuantizationMode.NONE, ): compile_options = {} exported_graph = get_exported_graph(model, sample_inputs, qmode=qmode) program = export( exported_graph, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True, ) edge_program = to_edge_transform_and_lower( program, partitioner=[XnnpackPartitioner(compile_options)], transform_passes=None, compile_config=None, ) executorch_program = edge_program.to_executorch() # Check if the delegate ID matches XnnpackBackend if ( executorch_program.executorch_program.execution_plan[0].delegates[0].id != XnnpackBackend.__name__ ): raise RuntimeError( f"Expected delegate ID {XnnpackBackend.__name__}, but got {executorch_program.executorch_program.execution_plan[0].delegates[0].id}" ) return executorch_program def print_tensor_comparison_errors( tensor1, tensor2, atol=1e-03, rtol=1e-03, max_errors=10 ): """ Print the first max_errors tensor indexes that exceed the absolute/relative tolerance and the error at each of those locations. Args: tensor1: First tensor to compare tensor2: Second tensor to compare atol: Absolute tolerance rtol: Relative tolerance max_errors: Maximum number of errors to print (default: 10) """ # Handle lists/tuples of tensors if isinstance(tensor1, (list, tuple)) and isinstance(tensor2, (list, tuple)): if len(tensor1) != len(tensor2): print(f"Tensor count mismatch: {len(tensor1)} vs {len(tensor2)}") return for i, (t1, t2) in enumerate(zip(tensor1, tensor2)): print(f"\n=== Tensor {i} comparison ===") print_tensor_comparison_errors(t1, t2, atol, rtol, max_errors) return # Handle single tensor comparison if not isinstance(tensor1, torch.Tensor) or not isinstance(tensor2, torch.Tensor): print("Error: Both inputs must be torch.Tensor objects") return if tensor1.shape != tensor2.shape: print(f"Shape mismatch: {tensor1.shape} vs {tensor2.shape}") return # Calculate absolute and relative errors abs_diff = torch.abs(tensor1 - tensor2) rel_diff = abs_diff / ( torch.abs(tensor2) + 1e-8 ) # Add small epsilon to avoid division by zero # Find locations where tolerance is exceeded tolerance_mask = (abs_diff > atol) & (rel_diff > rtol) if not tolerance_mask.any(): print("All values are within tolerance") return # Get indices where tolerance is exceeded error_indices = torch.nonzero(tolerance_mask, as_tuple=False) total_errors = error_indices.shape[0] print(f"Found {total_errors} values exceeding tolerance (atol={atol}, rtol={rtol})") print(f"Showing first {min(max_errors, total_errors)} errors:") print("Index -> tensor1_value, tensor2_value, abs_error, rel_error") # Print first max_errors locations for i in range(min(max_errors, total_errors)): idx = tuple(error_indices[i].tolist()) val1 = tensor1[idx].item() val2 = tensor2[idx].item() abs_err = abs_diff[idx].item() rel_err = rel_diff[idx].item() print( f"{idx} -> {val1:.6f}, {val2:.6f}, abs_err={abs_err:.6f}, rel_err={rel_err:.6f}" ) def check_outputs_equal( model_output, ref_output, atol=1e-03, rtol=1e-03, first_output_only=False ): """ Helper function that checks if model output and reference output are equal with some tolerance. Returns True if equal, False otherwise. """ # Convert OrderedDict to list if needed if isinstance(ref_output, OrderedDict): ref_output = list(ref_output.values()) # Compare the result from executor and eager mode directly if isinstance(ref_output, tuple) or isinstance(ref_output, list): # Multiple outputs executor always returns tuple, even if there is one output if len(ref_output) != len(model_output): print_tensor_comparison_errors(model_output, ref_output, atol, rtol) return False if first_output_only: result = torch.allclose( model_output[0], ref_output[0], atol=atol, rtol=rtol ) if not result: print_tensor_comparison_errors( model_output[0], ref_output[0], atol, rtol ) return result else: result = True for i in range(len(ref_output)): if isinstance(ref_output[i], torch.Tensor): if not torch.allclose( model_output[i], ref_output[i], atol=atol, rtol=rtol ): print(f"\n=== Output {i} comparison failed ===") print_tensor_comparison_errors( model_output[i], ref_output[i], atol, rtol ) result = False elif isinstance(ref_output[i], int): if not model_output[i] == ref_output[i]: print(f"\n=== Output {i} comparison failed ===") print(f"{model_output[i]} vs {ref_output[[i]]}") result = False else: print(f"WARNING: Output {i} has type {type(ref_output[i])}") return result else: # If one output, eager returns tensor while executor tuple of size 1 result = torch.allclose(model_output[0], ref_output, atol=atol, rtol=rtol) if not result: print_tensor_comparison_errors(model_output[0], ref_output, atol, rtol) return result def run_and_check_output( reference_model: torch.nn.Module, executorch_program: ExecutorchProgramManager, sample_inputs: Tuple[torch.Tensor], atol=1e-03, rtol=1e-01, first_output_only=False, ) -> bool: """ Utility function that accepts an already lowered ExecuTorch program, executes it with the provided sample input, and checks the output for correctness. Args: executorch_program: Already lowered ExecutorchProgramManager sample_inputs: Sample inputs to run the program with reference_model: Reference model to generate reference outputs for comparison atol: Absolute tolerance for output comparison rtol: Relative tolerance for output comparison first_output_only: Whether to compare only the first output Returns: bool: True if outputs match within tolerance, False otherwise """ # Load the ExecuTorch program executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer) # Flatten inputs for execution inputs_flattened, _ = tree_flatten(sample_inputs) # Run the ExecuTorch program model_output = executorch_module.run_method("forward", tuple(inputs_flattened)) # Generate reference outputs using the reference model ref_output, _ = tree_flatten(reference_model(*sample_inputs)) # Check if outputs are equal return check_outputs_equal( model_output, ref_output, atol=atol, rtol=rtol, first_output_only=first_output_only, ) def make_copy_of_inputs(sample_inputs: Tuple[Any]) -> Tuple[Any]: sample_inputs_copy = [] for input_val in sample_inputs: if isinstance(input_val, torch.Tensor): sample_inputs_copy.append(input_val.clone()) else: sample_inputs_copy.append(deepcopy(input_val)) return tuple(sample_inputs_copy) def lower_module_and_test_output( model: torch.nn.Module, sample_inputs: Tuple[torch.Tensor], atol=1e-03, rtol=1e-01, dynamic_shapes=None, test_inputs=None, first_output_only=False, operator_blocklist=None, operator_allowlist=None, nn_module_allowlist=None, nn_module_blocklist=None, xnnpack=False, ) -> bool: """ Helper testing function that takes a torch.nn.Module and lowers it to Vulkan with the given sample inputs. It then runs the lowered module and compares its outputs with the outputs of the eager module. Returns: bool: True if all comparisons pass, False otherwise. """ # Export model to Vulkan using the helper function if xnnpack: executorch_program = export_model_to_xnnpack( model, make_copy_of_inputs(sample_inputs), dynamic_shapes, operator_blocklist, operator_allowlist, nn_module_blocklist, nn_module_allowlist, ) else: executorch_program = export_model_to_vulkan( model, make_copy_of_inputs(sample_inputs), dynamic_shapes, operator_blocklist=operator_blocklist, operator_allowlist=operator_allowlist, nn_module_blocklist=nn_module_blocklist, nn_module_allowlist=nn_module_allowlist, ) executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer) inputs_flattened, _ = tree_flatten(sample_inputs) model_output = executorch_module.run_method("forward", tuple(inputs_flattened)) ref_output = model(*make_copy_of_inputs(sample_inputs)) if not check_outputs_equal( model_output, ref_output, atol=atol, rtol=rtol, first_output_only=first_output_only, ): return False if test_inputs is not None: for test_input in test_inputs: test_inputs_flattened, _ = tree_flatten(test_input) model_output = executorch_module.run_method( "forward", tuple(test_inputs_flattened) ) ref_output = model(*test_input) if not check_outputs_equal( model_output, ref_output, atol=atol, rtol=rtol, first_output_only=first_output_only, ): return False return True def create_bundled_program( executorch_program: ExecutorchProgramManager, sample_inputs: Tuple[torch.Tensor, ...], expected_outputs: List[Any], method_name: str = "forward", ) -> bytes: """ Create a bundled program containing the model and test cases for correctness testing. Args: executorch_program: The ExecutorchProgramManager to bundle sample_inputs: Sample inputs for the model expected_outputs: Expected outputs from running the model with sample_inputs method_name: Name of the method to test (default: "forward") Returns: Serialized bundled program as bytes """ # Flatten sample inputs to match expected format inputs_flattened, _ = tree_flatten(sample_inputs) # Create test suite with the sample inputs and expected outputs test_suites = [ MethodTestSuite( method_name=method_name, test_cases=[ MethodTestCase( inputs=inputs_flattened, expected_outputs=expected_outputs, ) ], ) ] # Create bundled program bundled_program = BundledProgram(executorch_program, test_suites) # Serialize to flatbuffer bundled_buffer = serialize_from_bundled_program_to_flatbuffer(bundled_program) return bundled_buffer def save_bundled_program( model: torch.nn.Module, sample_inputs: Tuple[torch.Tensor], output_path: str, method_name: str = "forward", sample_kwargs=None, et_program: Optional[ExecutorchProgramManager] = None, dynamic_shapes=None, ) -> str: """ Export a bundled .pte file containing the model and test cases. Args: model: The PyTorch model to export sample_inputs: Sample inputs for the model output_path: Path where the bundled .pte file should be saved (should end with .bpte) method_name: Name of the method to test (default: "forward") et_program: Optional pre-exported ExecutorchProgramManager. If None, will export to Vulkan dynamic_shapes: Optional dynamic shapes for export Returns: str: Path to the saved bundled program file """ # If no ExecutorchProgramManager provided, export to Vulkan if et_program is None: et_program = export_model_to_vulkan( model, sample_inputs, sample_kwargs=sample_kwargs, dynamic_shapes=dynamic_shapes, ) if sample_kwargs is None: sample_kwargs = {} # Generate expected outputs by running the model expected_outputs = [getattr(model, method_name)(*sample_inputs, **sample_kwargs)] # Flatten sample inputs with kwargs to match expected format inputs_flattened, _ = tree_flatten((sample_inputs, sample_kwargs)) # Create bundled program bp_buffer = create_bundled_program( et_program, tuple(inputs_flattened), expected_outputs, method_name, ) # Ensure output path has correct extension if not output_path.endswith(".bpte"): output_path = output_path + ".bpte" # Write to file with open(output_path, "wb") as file: file.write(bp_buffer) return output_path def save_executorch_program( executorch_program: ExecutorchProgramManager, output_path: str, ) -> str: """ Save an ExecutorchProgramManager as a .pte file. Args: executorch_program: The ExecutorchProgramManager to save output_path: Path where the .pte file should be saved (should end with .pte) Returns: str: Path to the saved .pte file """ # Ensure output path has correct extension if not output_path.endswith(".pte"): output_path = output_path + ".pte" # Write to file with open(output_path, "wb") as file: executorch_program.write_to_file(file) return output_path def print_occurrences(edge_program, operator_list: List): """ Print the input/output information for all occurrences of specified operators in the edge program. Args: edge_program: The edge program created by to_edge_transform_and_lower operator_list: List of operators to search for in the graph """ logger = logging.getLogger("") logger.setLevel(logging.INFO) logger.info( f"Searching for occurrences of {len(operator_list)} operators in the graph..." ) occurrence_count = 0 for node in edge_program.exported_program().graph.nodes: if utils.is_torch_op_node(node): target = node.target # Handle auto_functionalized nodes if ( node.target == torch.ops.higher_order.auto_functionalized or node.target == torch.ops.higher_order.auto_functionalized_v2 ): first_arg = node.args[0] if hasattr(first_arg, "name"): target = first_arg.name() elif hasattr(first_arg, "__name__"): target = first_arg.__name__ # Check if this operator is in our list if target in operator_list: occurrence_count += 1 logger.info(f"Occurrence {occurrence_count}: {node.format_node()}") # Get the node I/O string using the utils function try: io_str = utils.node_io_str(node) logger.info(f" {io_str}") except Exception as e: logger.info(f" Error getting I/O string: {e}") if occurrence_count == 0: logger.info("No occurrences of the specified operators found in the graph.") else: logger.info( f"Found {occurrence_count} total occurrences of the specified operators." ) def op_ablation_test( # noqa: C901 model: torch.nn.Module, sample_inputs: Tuple[torch.Tensor], atol=1e-03, rtol=1e-01, dynamic_shapes=None, test_inputs=None, first_output_only=False, ) -> dict: """ Fast binary search utility function to determine which operators work correctly when delegated to Vulkan. This function uses a binary search approach to efficiently find bad operators: 1. Split operators into two halves (least frequent first, most frequent second) 2. Test each half to see if it produces correct output 3. Add good halves to known_good_ops and recursively search bad halves 4. Continue until all operators are classified Args: model: The PyTorch model to test sample_inputs: Sample inputs for the model atol: Absolute tolerance for output comparison rtol: Relative tolerance for output comparison dynamic_shapes: Optional dynamic shapes for export test_inputs: Optional additional test inputs first_output_only: Whether to compare only the first output Returns: dict: Dictionary with keys: - 'good_operators': List of operators that work correctly - 'bad_operators': List of operators that cause failures - 'operator_frequencies': Dictionary mapping operators to their occurrence count - 'all_operators': List of all unique operators found in the graph - 'test_count': Number of tests performed """ logger = logging.getLogger("") logger.setLevel(logging.INFO) logger.info("Starting fast binary search operator ablation test...") # Step 1: Export model to get edge_program and extract operators export_training_graph = export(model, sample_inputs, strict=True).module() program = export( export_training_graph, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True, ) edge_program = to_edge_transform_and_lower( program, partitioner=[], # No partitioner to get the full graph transform_passes=None, compile_config=None, ) # Step 2: Scan edge_program.graph_module to obtain unique operators and their frequencies operator_frequencies = {} for node in edge_program.exported_program().graph.nodes: if utils.is_torch_op_node(node): target = node.target # Handle auto_functionalized nodes if ( node.target == torch.ops.higher_order.auto_functionalized or node.target == torch.ops.higher_order.auto_functionalized_v2 ): first_arg = node.args[0] if hasattr(first_arg, "name"): target = first_arg.name() elif hasattr(first_arg, "__name__"): target = first_arg.__name__ if target in operator_frequencies: operator_frequencies[target] += 1 else: operator_frequencies[target] = 1 all_operators = list(operator_frequencies.keys()) logger.info(f"Found {len(all_operators)} unique operators in the graph") # Sort operators by frequency (most frequent first for binary search) operators_by_frequency = sorted( all_operators, key=lambda op: operator_frequencies[op], reverse=True ) logger.info("Operator frequencies (sorted by occurrence, most frequent first):") for op in operators_by_frequency: logger.info(f" {op.name()}: {operator_frequencies[op]} occurrences") # Global test counter test_count = 0 def test_operator_set(ops_to_test: List, known_good_ops: List) -> bool: """Test if a set of operators works correctly when combined with known good operators.""" nonlocal test_count test_count += 1 test_allowlist = known_good_ops + ops_to_test logger.info( f"Test {test_count}: Testing {len(ops_to_test)} operators with {len(known_good_ops)} known good" ) try: success = lower_module_and_test_output( model=model, sample_inputs=sample_inputs, atol=atol, rtol=rtol, dynamic_shapes=dynamic_shapes, test_inputs=test_inputs, first_output_only=first_output_only, operator_allowlist=test_allowlist, ) logger.info(f" {'✓ PASS' if success else '✗ FAIL'}") # Log known good ops logger.info(" Known good:") for op in known_good_ops: logger.info(f" * {op.name()}") # Log tested ops logger.info(" Tested ops:") for op in ops_to_test: logger.info(f" * {op.name()}") return success except Exception as e: logger.info(f" ! Error: {e}") return False def find_bad_operators( ops_to_test: List, known_good_ops: List ) -> Tuple[List, List]: """ Recursively find bad operators using binary search. Returns: Tuple of (good_operators, bad_operators) from ops_to_test """ if not ops_to_test: return [], [] if len(ops_to_test) == 1: # Base case: single operator op = ops_to_test[0] if test_operator_set([op], known_good_ops): logger.info(f" Single operator {op.name()} is GOOD") return [op], [] else: logger.info(f" Single operator {op.name()} is BAD") return [], [op] # Split ops_to_test into two halves mid = len(ops_to_test) // 2 first_half = ops_to_test[:mid] # Least frequent operators second_half = ops_to_test[mid:] # Most frequent operators logger.info( f"Splitting {len(ops_to_test)} operators: {len(first_half)} + {len(second_half)}" ) # Log known good ops logger.info(" Known good:") for op in known_good_ops: logger.info(f" * {op.name()}") # Log first half ops logger.info(" First half ops:") for op in first_half: logger.info(f" * {op.name()}") # Log second half ops logger.info(" Second half ops:") for op in second_half: logger.info(f" * {op.name()}") good_ops = [] bad_ops = [] first_half_good = test_operator_set(first_half, known_good_ops) if first_half_good: logger.info( f"First half ({len(first_half)} ops) is good - adding to known good" ) good_ops.extend(first_half) known_good_ops.extend(first_half) second_half_good = test_operator_set(second_half, known_good_ops) if second_half_good: logger.info( f"Second half ({len(second_half)} ops) is good - adding to known good" ) good_ops.extend(second_half) if not first_half_good: logger.info(f"First half ({len(first_half)} ops) is bad - recursing") sub_good, sub_bad = find_bad_operators(first_half, known_good_ops) good_ops.extend(sub_good) bad_ops.extend(sub_bad) known_good_ops.extend(sub_good) if not second_half_good: logger.info(f"Second half ({len(second_half)} ops) is bad - recursing") sub_good, sub_bad = find_bad_operators(second_half, known_good_ops) good_ops.extend(sub_good) bad_ops.extend(sub_bad) return good_ops, bad_ops # Start the binary search logger.info( f"\n=== Starting binary search on {len(operators_by_frequency)} operators ===" ) good_operators, bad_operators = find_bad_operators(operators_by_frequency, []) # Summary of results logger.info(f"\n=== Binary search complete after {test_count} tests ===") logger.info(f"Good operators ({len(good_operators)}):") for op in good_operators: logger.info(f" ✓ {op.name()} (frequency: {operator_frequencies[op]})") logger.info(f"Bad operators ({len(bad_operators)}):") for op in bad_operators: logger.info(f" ✗ {op.name()} (frequency: {operator_frequencies[op]})") print_occurrences(edge_program, bad_operators) efficiency_gain = len(all_operators) - test_count logger.info( f"Efficiency: {test_count} tests instead of {len(all_operators)} (saved {efficiency_gain} tests)" ) return { "good_operators": good_operators, "bad_operators": bad_operators, "operator_frequencies": operator_frequencies, "all_operators": all_operators, "test_count": test_count, } def make_indent(indent_level): indent_str = "" for _ in range(indent_level): indent_str += " " return indent_str def print_output(outputs, n: int = 0, indent_level: int = 0): if isinstance(outputs, (list, tuple)): print(f"{make_indent(indent_level)}output_{n} = {type(outputs)}") new_indent_level = indent_level + 2 for n, test_out in enumerate(outputs): print_output(test_out, n, new_indent_level) elif isinstance(outputs, torch.Tensor): print( f"{make_indent(indent_level)}output_{n} = test_utils.random_uniform_tensor({outputs.shape}, low={outputs.min().item()}, high={outputs.max().item()}, dtype={outputs.dtype})" ) elif isinstance(outputs, int): print(f"{make_indent(indent_level)}output_{n} = {outputs}") else: print(f"{make_indent(indent_level)}output_{n} = {type(outputs)}")