# Owner(s): ["module: inductor"] import functools import logging from collections.abc import Callable from typing import Any, Optional, Union import torch from torch._inductor.codegen.subgraph import SubgraphTemplate from torch._inductor.ir import Buffer, FixedLayout, ir_node_to_tensor, TensorBox from torch._inductor.lowering import lowerings, validate_ir from torch._inductor.select_algorithm import ( autotune_select_algorithm, ExternKernelChoice, ) from torch._inductor.virtualized import V from torch.utils._ordered_set import OrderedSet log = logging.getLogger(__name__) def _detect_collective_ops(choices: list) -> bool: """ Detect if choices contain collective operations. """ from torch._inductor.utils import is_collective_op for choice in choices: if not hasattr(choice, "gm") or choice.gm is None: continue for node in choice.gm.graph.nodes: if node.op == "call_function" and node.target is not None: op_name = str(node.target) if is_collective_op(op_name) or is_collective_op( f"torch.ops.{op_name}" ): return True return False class CustomOpConfig: """Config for custom op autotuning. Specifies optional decomposition function with parameter values. Each config creates exactly one variant. Args: decomposition: Optional functions to autotune. If not provided, default will be used. **params: Parameters passed to the function Examples: CustomOpConfig(attention_impl, head_dim=32, method='chunked') CustomOpConfig(head_dim=32, method='chunked') """ def __init__( self, decomposition: Optional[Callable[..., Any]] = None, **params: Any, ): if decomposition is not None and not callable(decomposition): raise TypeError( f"decomposition must be callable, got {type(decomposition)}" ) self.decomposition = decomposition self.params = params def get_decomposition( self, default_impl: Optional[Callable[..., Any]] = None ) -> Callable[..., Any]: """Return the decomposition function for this config. When decomposition is not specified, return the default implementation. """ if self.decomposition is not None: return self.decomposition if default_impl is not None and callable(default_impl): return default_impl raise TypeError( "No decomposition specified in config and no default implementation provided. " "Please provide a decomposition function in CustomOpConfig." ) def __repr__(self) -> str: decomp_name = self.decomposition.__name__ if self.decomposition else "default" if self.params: params_str = ", ".join(f"{k}={v}" for k, v in self.params.items()) return f"CustomOpConfig({decomp_name}, {params_str})" return f"CustomOpConfig({decomp_name})" __all__ = [ "autotune_custom_op", "register_custom_op_autotuning", "CustomOpConfig", ] def _extract_tensor_inputs( args: tuple[Any, ...], kwargs: dict[str, Any] ) -> tuple[list[Any], dict[str, Any]]: """Extract tensor inputs from mixed args/kwargs. Separates tensors (for autotuning input_nodes) from non-tensor parameters. Non-tensor kwargs are later functools.partial'd into decomposition functions. Args: args: Positional arguments (mix of tensors and scalars) kwargs: Keyword arguments (mix of tensors and scalars) Returns: Tuple of (tensor_inputs_list, non_tensor_kwargs) """ tensor_inputs = [] non_tensor_kwargs = {} # Process args and kwargs: separate tensor inputs and non tensor args for i, arg in enumerate(args): if isinstance(arg, (TensorBox, Buffer)): tensor_inputs.append(arg) else: # Add non-tensor positional args to kwargs with generated names non_tensor_kwargs[f"arg_{i}"] = arg for key, value in kwargs.items(): if isinstance(value, (TensorBox, Buffer)): tensor_inputs.append(value) else: non_tensor_kwargs[key] = value return tensor_inputs, non_tensor_kwargs def _merge_config_and_runtime_kwargs( config_params: dict[str, Any], runtime_kwargs: dict[str, Any], ) -> dict[str, Any]: """Merge config parameters with runtime kwargs. Runtime kwargs take precedence. If there are conflicts, log a warning and use runtime value. Args: config_params: Parameters from CustomOpConfig runtime_kwargs: Runtime non-tensor kwargs from _extract_tensor_inputs Returns: Merged kwargs dictionary with runtime values taking precedence """ merged_kwargs = config_params.copy() # Check for conflicts and let runtime kwargs dominate conflicts = OrderedSet(config_params.keys()).intersection(runtime_kwargs.keys()) for key in conflicts: log.warning( "Parameter '%s' specified both in CustomOpConfig (%s) " "and at runtime (%s). Using runtime value.", key, config_params[key], runtime_kwargs[key], ) # Runtime kwargs override config params merged_kwargs.update(runtime_kwargs) return merged_kwargs def _adapt_user_input_gen_fns( inputs: list[Any], arg_names: list[str], user_input_gen_fns: dict[str, Callable[[torch.Tensor], torch.Tensor]], ) -> dict[int, Callable[[Any], torch.Tensor]]: """Convert user input generators from name-based to index-based format. Inductor autotune's input_gen_fns expects index of arg_names as key. Uses V.graph.sizevars.size_hints() to guess best for dynamic shapes. """ name_to_index = {name: i for i, name in enumerate(arg_names)} index_based_fns = {} for name, gen_fn in user_input_gen_fns.items(): if name in name_to_index: index_based_fns[name_to_index[name]] = gen_fn else: log.warning( "Unknown argument name '%s' in input_gen_fns. " "Available argument names: %s", name, list(name_to_index.keys()), ) def create_internal_input_gen_fn( user_function: Callable[[torch.Tensor], torch.Tensor], arg_name: str ) -> Callable[[Any], torch.Tensor]: """Create internal input generator that converts IR buffer to user's fake tensor.""" def internal_input_gen_fn(ir_buffer: Any) -> torch.Tensor: fake_tensor = ir_node_to_tensor(ir_buffer) assert fake_tensor is not None, "ir_node_to_tensor returned None" return user_function(fake_tensor) return internal_input_gen_fn return { i: create_internal_input_gen_fn( user_gen_fn, arg_names[i] if i < len(arg_names) else f"arg_{i}" ) for i, user_gen_fn in index_based_fns.items() if i < len(inputs) } def _create_fallback_choice( name: str, default_impl: Callable[..., Any], fake_output: torch.Tensor, kwargs: dict[str, Any], ) -> ExternKernelChoice: """Create fallback choice for default implementation.""" def fallback_wrapper(*args: Any) -> Any: return default_impl(*args, **kwargs) return ExternKernelChoice( kernel=fallback_wrapper, name=f"{name}_fallback_default", has_out_variant=False, op_overload=default_impl, use_fallback_kernel=True, ) def autotune_custom_op( name: str, decompositions: list[Callable[..., Any]], inputs: list[Any], non_tensor_args: list[dict[str, Any]], op_overload: torch._ops.OpOverload, user_input_gen_fns: Optional[ dict[str, Callable[[torch.Tensor], torch.Tensor]] ] = None, ) -> Union[TensorBox, Any]: """Autotune custom operations by comparing multiple decomposition implementations. Currently supports SINGLE OUTPUT custom ops only. TODO: Add support for multiple output custom ops (tuple/list returns). This function generates multiple implementation choices for a custom operation and uses Inductor's autotuning system to select the best performing variant at runtime. After selecting the best choice, applies inline fusion if the winning choice has a graph. Args: name: Unique identifier for the autotuning operation decompositions: List of alternative implementation functions to benchmark inputs: Input tensor IR nodes from compilation (TensorBox/Buffer objects) non_tensor_args: List of kwargs dicts, paired with corresponding decompositions arg op_overload: OpOverload of the custom op, used as fallback implementation user_input_gen_fns: Optional custom input generators for benchmarking. Maps input indices to functions that take fake tensors and return real tensors for performance measurement. Returns: IR node representing the optimized operation result Raises: TypeError: If decompositions is not a list/tuple RuntimeError: If no inputs or no valid choices generated """ if not isinstance(decompositions, (list, tuple)): raise TypeError( f"decompositions must be a list or tuple of callables, got {type(decompositions)}" ) if not inputs: raise RuntimeError(f"Custom op '{name}' requires tensor inputs for autotuning") if len(decompositions) != len(non_tensor_args): raise ValueError( f"decompositions and non_tensor_args must have same length, " f"got {len(decompositions)} decompositions and {len(non_tensor_args)} kwargs" ) template = SubgraphTemplate(name=name) choices = template.generate_custom_op_choices( name=name, # pyrefly: ignore [bad-argument-type] decompositions=decompositions, input_nodes=list(inputs), non_tensor_args=non_tensor_args, ) # Add default implementation as fallback if op_overload and hasattr(op_overload, "_op"): fallback_name = f"{name}_fallback_default" from torch._inductor.select_algorithm import extern_kernels # Skip if extern_kernel already registered to avoid duplicate registration error if not hasattr(extern_kernels, fallback_name): with V.fake_mode: fake_inputs = [ir_node_to_tensor(inp) for inp in inputs] fallback_kwargs = non_tensor_args[0] if non_tensor_args else {} fake_output = op_overload(*fake_inputs, **fallback_kwargs) fallback_choice = _create_fallback_choice( name, op_overload, fake_output, fallback_kwargs ) fallback_choice.maybe_append_choice( choices=choices, input_nodes=list(inputs), layout=FixedLayout( device=fake_output.device, dtype=fake_output.dtype, size=fake_output.shape, stride=fake_output.stride(), ), ) if not choices: raise RuntimeError(f"No valid choices generated for {name}") # Convert user input generation functions to internal format input_gen_fns = {} if user_input_gen_fns: import inspect arg_names = ( list(inspect.signature(decompositions[0]).parameters.keys()) if decompositions else [] ) input_gen_fns = _adapt_user_input_gen_fns(inputs, arg_names, user_input_gen_fns) is_collective = _detect_collective_ops(choices) # Run autotuning and get both result and winning choice selected_result, winning_choice = autotune_select_algorithm( name=name, choices=choices, input_nodes=list(inputs), layout=choices[0].layout, input_gen_fns=input_gen_fns, return_choice=True, is_collective=is_collective, ) # Apply inlining for fusion if winning_choice has graph; otherwise return result as-is(default fallback impl) if winning_choice.gm is not None: log.debug( "Inlining winning choice: %s (name=%s)", getattr(winning_choice, "name", type(winning_choice).__name__), name, ) from torch._inductor.codegen.subgraph import inline_subgraph_to_ir_nodes return inline_subgraph_to_ir_nodes(winning_choice.gm, inputs, name) log.debug( "Winning choice does not support inlining: %s (name=%s)", getattr(winning_choice, "name", type(winning_choice).__name__), name, ) return selected_result def _generate_dynamic_configs( tensor_inputs: list[Buffer], config_generator: Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]], default_impl: Callable[..., Any], operation_name: str, ) -> list[CustomOpConfig]: """Generate configs dynamically based on input tensors at lowering time.""" import inspect sig = inspect.signature(default_impl) param_names = list(sig.parameters.keys()) with V.fake_mode: fake_tensors = [ir_node_to_tensor(inp) for inp in tensor_inputs] fake_tensors_dict = dict(zip(param_names, fake_tensors)) configs = config_generator(fake_tensors_dict) if not isinstance(configs, (list, tuple)): raise TypeError( f"config_generator must return a list or tuple of CustomOpConfig, " f"got {type(configs)}" ) if not configs: raise ValueError(f"config_generator returned empty list for {operation_name}. ") return list(configs) def register_custom_op_autotuning( custom_op: torch._library.custom_ops.CustomOpDef, configs: Optional[Union[list[CustomOpConfig], list[Callable[..., Any]]]] = None, config_generator: Optional[ Callable[[dict[str, torch.Tensor]], list[CustomOpConfig]] ] = None, name: Optional[str] = None, input_gen_fns: Optional[dict[str, Callable[[torch.Tensor], torch.Tensor]]] = None, ) -> None: """Register custom op for autotuning with custom_op configs where each config specifies a decomposition implementation function with its parameter values. Args: custom_op: Custom operation (decorated function from @torch.library.custom_op) configs: List of CustomOpConfig objects for static inputs. Mutually exclusive with config_generator. config_generator: Dynamic config generator function that takes a dict mapping parameter names to fake tensors, and returns list[CustomOpConfig] based on input tensor properties. Mutually exclusive with configs. name: Operation name (default: "{op_name}_autotuned") input_gen_fns: Custom input generators for benchmarking Examples: # Static configs @torch.library.custom_op("mylib::attention", mutates_args=()) def my_attention(query, key, value, head_dim=32): ... register_custom_op_autotuning( my_attention, configs=[ CustomOpConfig(attention_impl, head_dim=32, method='chunked'), CustomOpConfig(attention_impl, head_dim=64, method='tiled'), CustomOpConfig(head_dim=128), # No decomposition specified, use default ], input_gen_fns={ "query": lambda fake: torch.randn_like(fake, device='cuda'), "key": lambda fake: torch.randn_like(fake, device='cuda'), "value": lambda fake: torch.randn_like(fake, device='cuda'), }, ) # Dynamic config generation based on input tensor properties def generate_k_split_configs(fake_tensors: dict[str, torch.Tensor]) -> list[CustomOpConfig]: # Access tensor shapes, dtypes, devices, etc. m, k = fake_tensors["mat1"].shape _, n = fake_tensors["mat2"].shape k_splits = ... # compute possible k splits based on tensor properties return [CustomOpConfig(k_splits=k) for k in k_splits] register_custom_op_autotuning( matmul_decomposeK_op, config_generator=generate_k_split_configs, input_gen_fns={...}, ) """ from torch._library.custom_ops import CustomOpDef if not isinstance(custom_op, CustomOpDef): raise TypeError( f"custom_op must be a CustomOpDef (decorated function from @torch.library.custom_op), " f"got {type(custom_op)}." ) # Validate configs and config_generator are mutually exclusive if configs is not None and config_generator is not None: raise ValueError( "Cannot specify both 'configs' and 'config_generator'. " "Use 'config_generator' for shape-dependent configs." ) if configs is None and config_generator is None: raise ValueError("Must specify either 'configs' or 'config_generator'") op_overload = custom_op._opoverload default_impl = custom_op._init_fn # Process and validate static configs at registration time static_configs = None if configs is not None: if not isinstance(configs, (list, tuple)): raise TypeError(f"configs must be a list or tuple, got {type(configs)}") static_configs = [] for cfg in configs: if isinstance(cfg, CustomOpConfig): static_configs.append(cfg) else: raise TypeError( f"Each config must be a CustomOpConfig object, got {type(cfg)}" ) if not static_configs: raise ValueError("At least one config must be provided") if name is None: name = f"{op_overload._name}_autotuned" @functools.wraps(op_overload) def autotuning_lowering(*args: Any, **kwargs: Any) -> Any: """Inductor lowering function that replaces custom op calls with autotuned versions.""" # Extract tensor inputs and non-tensor parameters (runtime kwargs) tensor_inputs, runtime_kwargs = _extract_tensor_inputs(args, kwargs) # Get configs: either generate dynamically or use static configs if config_generator is not None: configs_to_use = _generate_dynamic_configs( tensor_inputs, config_generator, default_impl, name ) else: assert static_configs is not None configs_to_use = static_configs # Prepare decompositions and kwargs for autotuning decompositions = [] non_tensor_args = [] for cfg in configs_to_use: decomp = cfg.get_decomposition(default_impl=default_impl) decompositions.append(decomp) # Merge config params with runtime kwargs (runtime takes precedence) merged_kwargs = _merge_config_and_runtime_kwargs(cfg.params, runtime_kwargs) non_tensor_args.append(merged_kwargs) result = autotune_custom_op( name=name, decompositions=decompositions, inputs=tensor_inputs, non_tensor_args=non_tensor_args, op_overload=op_overload, user_input_gen_fns=input_gen_fns, ) validate_ir(result) return result lowerings[op_overload] = autotuning_lowering