| import collections |
| import dataclasses |
| import functools |
| import itertools |
| import logging |
| import os |
| import pprint |
| import textwrap |
| from typing import Dict, List, Optional, Set |
| |
| import sympy |
| |
| import torch |
| from torch._dynamo.utils import dynamo_timed |
| |
| from . import config, dependencies, ir, metrics |
| from .dependencies import StarDep, WeakDep |
| from .sizevars import SimplifyIndexing |
| from .utils import cache_on_self, cmp, has_triton |
| from .virtualized import V |
| |
| log = logging.getLogger(__name__) |
| |
| |
| def pformat(obj): |
| if isinstance(obj, set): |
| # pformat has trouble with sets of sympy exprs |
| obj = sorted(obj, key=str) |
| result = pprint.pformat(obj, indent=4) |
| if "\n" in result: |
| return f"\n{textwrap.indent(result, ' '*4)}" |
| return result |
| |
| |
| class OutputNode: |
| def __init__(self, dep): |
| self.unmet_dependencies = {dep} |
| self.inverse_users = [] |
| |
| def is_reduction(self): |
| return False |
| |
| def get_alias_names(self): |
| return () |
| |
| def get_name(self): |
| return "OUTPUT" |
| |
| __repr__ = get_name |
| |
| |
| class BaseSchedulerNode: |
| def __init__(self, scheduler: "Scheduler", node: ir.Buffer): |
| self.scheduler: "Scheduler" = scheduler |
| self.node: ir.Buffer = node |
| self.users: Optional[List[NodeUser]] = None |
| self.inverse_users: List[BaseSchedulerNode] = [] |
| self.set_read_writes(node.get_read_writes()) |
| self.recursive_predecessors: Optional[Set[str]] = None |
| self.min_order: Optional[int] = None |
| self.max_order: Optional[int] = None |
| self.last_usage: Set[str] = None # buffers that won't be used after this kernel |
| self.written = False |
| |
| def __repr__(self): |
| return f"{type(self).__name__}(name={self.get_name()!r})" |
| |
| def debug_str(self): |
| """Longer form printout for trace logs""" |
| name = self.get_name() |
| lines = [ |
| f"{name}: {type(self).__name__}({type(self.node).__name__})", |
| f"{name}.writes = {pformat(self.read_writes.writes)}", |
| f"{name}.unmet_dependencies = {pformat(self.unmet_dependencies)}", |
| f"{name}.met_dependencies = {pformat(self.read_writes.reads - self.unmet_dependencies)}", |
| ] |
| try: |
| lines += [ |
| self.debug_str_extra(), |
| ] |
| except Exception: |
| log.warning("Ignoring error in debug_str()", exc_info=True) |
| return "\n".join(lines).rstrip() |
| |
| def debug_str_extra(self): |
| return "" |
| |
| def log_details(self): |
| log.info( |
| "%s: unmet_dependencies = %s, writes = %s", |
| self, |
| self.unmet_dependencies, |
| self.read_writes.writes, |
| ) |
| |
| def update_mutated_names(self, renames: Dict[str, str]): |
| self.set_read_writes(self.read_writes.rename(renames)) |
| |
| def add_mutation_dep(self, dep): |
| self.set_read_writes(self.read_writes.with_read(dep)) |
| |
| def set_users(self, users: List["NodeUser"]): |
| # deduplicate |
| result: Dict[int, NodeUser] = {} |
| for use in users: |
| if id(use.node) in result: |
| result[id(use.node)] = NodeUser( |
| use.node, result[id(use.node)].can_inplace and use.can_inplace |
| ) |
| else: |
| result[id(use.node)] = use |
| self.users = list(result.values()) |
| |
| def get_aliases(self): |
| return self.node.get_alias_names() |
| |
| def get_mutations(self): |
| return self.node.get_mutation_names() |
| |
| def has_aliasing_or_mutation(self): |
| return bool(self.get_aliases() or self.get_mutations()) |
| |
| def set_read_writes(self, rw: dependencies.ReadWrites): |
| self.read_writes: dependencies.ReadWrites = rw |
| self.unmet_dependencies = self.read_writes.reads |
| self.prune_deps() |
| |
| def used_buffer_names(self) -> Set[str]: |
| return { |
| dep.name |
| for dep in itertools.chain(self.read_writes.reads, self.read_writes.writes) |
| } |
| |
| def prune_deps(self): |
| self.unmet_dependencies = { |
| dep |
| for dep in self.unmet_dependencies |
| if dep.name not in self.scheduler.available_buffer_names |
| } |
| |
| def prune_redundant_deps(self, name_to_fused_node): |
| """ |
| Prunes stardeps intended for mutation ordering |
| on an upstream fused node if after fusion there is another dependency |
| on the fused upstream node, making the stardep redundant |
| |
| In essence this enforces an ordering on fusions. As fusions occur, prunable stardeps will |
| be incrementally removed, enabling other fusions, ensuring they are fused in order. |
| """ |
| name_to_dep_count = collections.Counter() |
| |
| for dep in self.unmet_dependencies: |
| if not isinstance(dep, WeakDep): |
| name_to_dep_count[name_to_fused_node[dep.name].get_name()] += 1 |
| |
| def should_prune(dep): |
| if isinstance(dep, WeakDep): |
| is_redundant = ( |
| name_to_dep_count[name_to_fused_node[dep.name].get_name()] > 0 |
| ) |
| # These can occur because fused nodes always gather deps from their snodes |
| # If B has a weakdep on A |
| # B gets fused with C, then any time BC is fused, the weakdep will reappear |
| is_self_dep = name_to_fused_node[dep.name] == self |
| return is_redundant or is_self_dep |
| else: |
| return False |
| |
| deps_to_prune = {dep for dep in self.unmet_dependencies if should_prune(dep)} |
| self.unmet_dependencies = self.unmet_dependencies - deps_to_prune |
| self.set_read_writes(self.read_writes.remove_reads(deps_to_prune)) |
| |
| def get_name(self) -> str: |
| return self.node.get_name() |
| |
| def get_first_name(self) -> str: |
| return self.get_name() |
| |
| def get_names(self) -> Set[str]: |
| return {self.get_name()} |
| |
| def get_nodes(self) -> List["BaseSchedulerNode"]: |
| return [self] |
| |
| def get_device(self): |
| return self.node.get_device() |
| |
| def is_reduction(self): |
| return False |
| |
| def is_template(self): |
| return False |
| |
| def is_extern(self): |
| return False |
| |
| def can_inplace(self, read_dep: dependencies.MemoryDep): |
| return False |
| |
| def allocate(self): |
| if not self.node.should_allocate(): |
| return |
| |
| if isinstance(self, (SchedulerNode,)) and ( |
| self.node.get_alias_names() or self.node.get_mutation_names() |
| ): |
| V.graph.wrapper_code.codegen_allocation(self.node) |
| return |
| |
| if ( |
| isinstance(self, (SchedulerNode,)) |
| and config.inplace_buffers |
| and ( |
| not isinstance(V.kernel, torch._inductor.codegen.triton.TritonKernel) |
| or getattr(V.kernel, "mutations", None) is not None |
| ) |
| ): |
| from .codegen.wrapper import buffer_reuse_key |
| |
| ordered_reads = sorted(self.read_writes.reads, key=lambda x: x.name) |
| |
| for read in ordered_reads: |
| input_node: BaseSchedulerNode = self.scheduler.name_to_node.get( |
| read.name |
| ) |
| if input_node and V.graph.wrapper_code.can_reuse(input_node): |
| remaining_uses = [ |
| x |
| for x in input_node.users |
| if x.node.get_name() |
| not in self.scheduler.available_buffer_names |
| ] |
| if ( |
| len(remaining_uses) == 1 |
| and remaining_uses[0].can_inplace |
| and remaining_uses[0].node is self |
| and not isinstance( |
| input_node.node.get_layout(), |
| ( |
| ir.MultiOutputLayout, |
| ir.MutationLayout, |
| ir.AliasedLayout, |
| ), |
| ) |
| and buffer_reuse_key(input_node.node) |
| == buffer_reuse_key(self.node) |
| ): |
| V.graph.wrapper_code.codegen_inplace_reuse( |
| input_node.node, self.node |
| ) |
| V.kernel.args.make_inplace( |
| input_node.get_name(), self.get_name() |
| ) |
| # mutations not tracked in cpp kernels |
| if isinstance( |
| V.kernel, torch._inductor.codegen.triton.TritonKernel |
| ): |
| V.kernel.mutations.add(input_node.get_name()) |
| V.kernel.mutations.add(self.get_name()) |
| return |
| |
| V.graph.wrapper_code.codegen_allocation(self.node) |
| |
| def can_free(self): |
| for use in self.users: |
| if isinstance(use.node, OutputNode): |
| return False |
| return True |
| |
| def codegen_originating_info(self, buffer, only_once=True): |
| if not config.comment_origin: |
| return |
| |
| if only_once and self.written: |
| return |
| origins = self.node.origins |
| out_lines = [] |
| |
| for o in origins: |
| if o.op == "output": |
| # These are boring and samey |
| continue |
| |
| out_lines.append("") |
| # TODO(voz): Should the pragma be constant somewhere? |
| out_lines.append("#pragma CMT ORIGIN:") |
| out_lines.append(f"#pragma CMT {o.op} {o.target}") |
| if "stack_trace" in o.meta: |
| stack_trace = f"{o.meta['stack_trace']}" |
| stack_trace_last_line = stack_trace.split("|")[-1] |
| out_lines.append( |
| "#pragma CMT " |
| + stack_trace_last_line.replace("{", "{{") |
| .replace("}", "}}") |
| .replace("\n", "\\") |
| ) |
| out_lines.append("#pragma CMT END ORIGIN") |
| out_lines.append("") |
| |
| if len(out_lines) == 0: |
| return |
| |
| # TODO(voz): Ostensibly, we should not need this. But there are cases where C++ codegen does |
| # not use BracesBuffer, so we have no good indicator of a C++ buffer atm. |
| buffer.writelines(out_lines) |
| self.written = True |
| |
| |
| class ExternKernelSchedulerNode(BaseSchedulerNode): |
| def debug_str_extra(self): |
| return f"{self.get_name()}.node.kernel = {getattr(self.node, 'kernel', None)}" |
| |
| def is_extern(self): |
| return True |
| |
| |
| class NopKernelSchedulerNode(BaseSchedulerNode): |
| pass |
| |
| |
| class SchedulerNode(BaseSchedulerNode): |
| def __init__(self, scheduler: "Scheduler", node: ir.ComputedBuffer, group_fn): |
| super().__init__(scheduler, node) |
| ( |
| self._sizes, |
| self._body, |
| ) = node.simplify_and_reorder() |
| |
| self.group = (node.get_device(), group_fn(self._sizes)) |
| |
| if self.is_template(): |
| self.set_read_writes(node.normalized_read_writes()) |
| else: |
| self.set_read_writes( |
| dependencies.extract_read_writes( |
| self._body, *self._sizes, normalize=True |
| ) |
| ) |
| |
| if self.is_reduction(): |
| # reduction has last (reduced) dim in its sizes, and some |
| # downstream dependencies get confused by it |
| self.read_writes.writes = self.read_writes.writes | { |
| w.strip_last_size() for w in self.read_writes.writes |
| } |
| # reduction not on the last dim swaps the sizes, and downstream |
| # dependencies expect unswapped |
| # TODO swapping sizes doesn't work, leads to |
| # File "/scratch/ngimel/work/repos/torchdynamo/torchinductor/sizevars.py", line 130, in guard_equals |
| # if len(right.free_symbols) < len(left.free_symbols): |
| # AttributeError: 'int' object has no attribute 'free_symbols' |
| # even though memory dep looks correct |
| # self.read_writes.writes = self.read_writes.writes | { |
| # w.maybe_swap_sizes() for w in self.read_writes.writes |
| # } |
| |
| def debug_str_extra(self): |
| name = self.get_name() |
| lines = [ |
| f"{name}.group.device = {self.group[0]}", |
| f"{name}.group.iteration = {self.group[1]}", |
| f"{name}.sizes = {self._sizes}", |
| ] |
| if self.get_aliases(): |
| lines.append(f"{name}.aliases = {pformat(self.get_aliases())}") |
| if self.get_mutations(): |
| lines.append(f"{name}.mutations = {pformat(self.get_mutations())}") |
| if isinstance(self._body, ir.LoopBody): |
| lines.append(f"class {name}_loop_body:") |
| lines.append(textwrap.indent(self._body.debug_str(), " ")) |
| return "\n".join(lines) |
| |
| def get_ranges(self): |
| return self._sizes |
| |
| def is_reduction(self): |
| return bool(self.node.get_reduction_type()) |
| |
| def is_template(self): |
| return isinstance(self.node, ir.TemplateBuffer) |
| |
| def run(self, *index_vars): |
| self.mark_run() |
| self.codegen(index_vars) |
| |
| def mark_run(self): |
| self.allocate() |
| |
| def ranges_from_index_vars(self, index_vars): |
| sizes = self._sizes |
| assert sum(map(len, sizes)) == sum(map(len, index_vars)) |
| var_ranges = dict( |
| zip( |
| itertools.chain.from_iterable(index_vars), |
| itertools.chain.from_iterable(sizes), |
| ) |
| ) |
| return var_ranges |
| |
| def codegen(self, index_vars): |
| var_ranges = self.ranges_from_index_vars(index_vars) |
| try: |
| with V.set_ops_handler( |
| SimplifyIndexing(V.get_ops_handler(), var_ranges) |
| ), V.kernel.set_current_node(self): |
| self._body(*index_vars) |
| except Exception: |
| log.fatal("Error in codegen for %s", self.node) |
| raise |
| |
| def pointwise_read_writes(self): |
| """ |
| Get the memory dependencies in the non-reduction axis. |
| """ |
| sizes, reduction_sizes = self._sizes |
| |
| def fn(index): |
| return self._body(index, [sympy.Integer(0) for _ in reduction_sizes]) |
| |
| return dependencies.extract_read_writes(fn, sizes) |
| |
| def can_inplace(self, read_dep: dependencies.MemoryDep): |
| if self.get_aliases() or self.is_template(): |
| return False |
| if len(self.read_writes.writes) == 1 and hasattr(read_dep, "index"): |
| write_dep = next(iter(self.read_writes.writes)) |
| return read_dep.index == write_dep.index and read_dep.size == write_dep.size |
| return False |
| |
| |
| class FusedSchedulerNode(BaseSchedulerNode): |
| """ |
| This is a "fake" scheduler node that represents a group of scheduler nodes |
| that are meant to be fused together. The way it does this is by maintaining |
| its unmet dependencies as the union of its constituent nodes. |
| """ |
| |
| @classmethod |
| def fuse(cls, node1: BaseSchedulerNode, node2: BaseSchedulerNode): |
| assert node1.scheduler is node2.scheduler |
| return cls(node1.scheduler, node1.get_nodes() + node2.get_nodes()) |
| |
| def __init__(self, scheduler: "Scheduler", snodes: List[SchedulerNode]): |
| # NB: No need to call super().__init__() because we don't need to re-use any of its logic. |
| self.snodes = snodes |
| self.scheduler = scheduler |
| self.node = None # type: ignore[assignment] |
| self.users = None |
| self.inverse_users = [] |
| self.group = max(snodes, key=lambda x: int(x.is_reduction())).group |
| self.recursive_predecessors = functools.reduce( |
| set.union, [x.recursive_predecessors for x in snodes] |
| ) |
| self.set_read_writes( |
| functools.reduce( |
| dependencies.ReadWrites.merge, [x.read_writes for x in snodes] |
| ) |
| ) |
| names = set(self.get_names()) |
| self.unmet_dependencies = { |
| dep |
| for dep in functools.reduce( |
| set.union, [x.unmet_dependencies for x in snodes] |
| ) |
| if dep.name not in names |
| } - self.read_writes.writes |
| self.min_order = min([x.min_order for x in self.snodes]) |
| self.max_order = max([x.max_order for x in self.snodes]) |
| |
| @cache_on_self |
| def get_name(self) -> str: |
| return "_".join([x.get_name() for x in self.snodes]) |
| |
| def get_first_name(self) -> str: |
| return self.snodes[0].get_name() |
| |
| @cache_on_self |
| def get_names(self) -> Set[str]: |
| return functools.reduce(set.union, [x.get_names() for x in self.snodes]) |
| |
| def debug_str_extra(self): |
| return ( |
| f"{self.get_name()}.snodes = {pformat([x.get_name() for x in self.snodes])}" |
| ) |
| |
| @cache_on_self |
| def used_buffer_names(self) -> Set[str]: |
| return functools.reduce(set.union, [x.used_buffer_names() for x in self.snodes]) |
| |
| def get_nodes(self) -> List[BaseSchedulerNode]: |
| return self.snodes |
| |
| def __repr__(self): |
| return f"{type(self).__name__}(nodes={self.get_name()})" |
| |
| @cache_on_self |
| def is_reduction(self): |
| return any(x.is_reduction() for x in self.snodes) |
| |
| @cache_on_self |
| def is_template(self): |
| return any(x.is_template() for x in self.snodes) |
| |
| def get_device(self): |
| return self.group[0] |
| |
| @cache_on_self |
| def has_aliasing_or_mutation(self): |
| return any(x.has_aliasing_or_mutation() for x in self.snodes) |
| |
| # None of these need to be implemented, as a FusedSchedulerNode is just an |
| # abstraction for scheduling purposes |
| def update_mutated_names(self, renames: Dict[str, str]): |
| raise NotImplementedError |
| |
| def add_mutation_dep(self, name): |
| raise NotImplementedError |
| |
| def set_users(self, users: List["NodeUser"]): |
| raise NotImplementedError |
| |
| def get_aliases(self): |
| raise NotImplementedError |
| |
| def get_mutations(self): |
| raise NotImplementedError |
| |
| def can_inplace(self, read_dep: dependencies.MemoryDep): |
| raise NotImplementedError |
| |
| def allocate(self): |
| raise NotImplementedError |
| |
| def can_free(self): |
| raise NotImplementedError |
| |
| |
| def pick_loop_order(stride_lengths, sizes, priority_idx=()): |
| """ |
| A heuristic to decide loop iteration orders. This has not been well |
| tuned and may be something we should autotune. |
| """ |
| |
| @functools.cmp_to_key |
| def index_cmp(a, b): |
| if sizes[a] == 1 or sizes[b] == 1: |
| # 1-sizes don't matter, just move them to the end |
| return cmp(sizes[a] == 1, sizes[b] == 1) |
| |
| stride_len_a = [sl[a] for sl in stride_lengths] |
| stride_len_b = [sl[b] for sl in stride_lengths] |
| |
| # equivalent to |
| # np.logical_or(stride_lengths[:, b] == 0, stride_lengths[:, a] < stride_lengths[:, b]).all() |
| a_first = all( |
| sl_b == 0 or sl_a < sl_b for sl_a, sl_b in zip(stride_len_a, stride_len_b) |
| ) |
| b_first = all( |
| sl_a == 0 or sl_b < sl_a for sl_a, sl_b in zip(stride_len_a, stride_len_b) |
| ) |
| if a_first and not b_first: |
| return -1 |
| if b_first and not a_first: |
| return 1 |
| |
| # otherwise contiguous |
| return cmp(b, a) |
| |
| order = list(reversed(range(len(stride_lengths[0])))) |
| if len(priority_idx) > 0: |
| # if we have priority node, only use that node's order |
| stride_lengths = [stride_lengths[pi] for pi in priority_idx] |
| if config.pick_loop_orders: |
| order.sort(key=index_cmp) |
| return order |
| |
| |
| @dataclasses.dataclass |
| class NodeUser: |
| node: BaseSchedulerNode |
| can_inplace: bool = False |
| |
| def get_name(self): |
| return self.node.get_name() |
| |
| |
| class Scheduler: |
| @dynamo_timed |
| def __init__(self, nodes): |
| super().__init__() |
| self.backends = {} |
| |
| self.nodes = [] |
| self.available_buffer_names = { |
| *V.graph.graph_inputs.keys(), |
| *V.graph.constants.keys(), |
| } |
| for node in nodes: |
| assert ( |
| node.origins is not None |
| ), "All nodes passed to scheduling must have an origin" |
| if node.is_no_op(): |
| self.nodes.append(NopKernelSchedulerNode(self, node)) |
| elif isinstance(node, (ir.ComputedBuffer, ir.TemplateBuffer)): |
| group_fn = self.get_backend(node.get_device()).group_fn |
| self.nodes.append(SchedulerNode(self, node, group_fn)) |
| elif isinstance(node, ir.ExternKernel): |
| self.nodes.append(ExternKernelSchedulerNode(self, node)) |
| else: |
| raise NotImplementedError(node) |
| # some new constants could have been created above |
| self.available_buffer_names.update(V.graph.constants.keys()) |
| for node in self.nodes: |
| node.prune_deps() |
| |
| self.name_to_node = {node.get_name(): node for node in self.nodes} |
| self.name_to_fused_node = None # set in fuse_nods() |
| |
| # we handle mutation by renaming modified versions of the same |
| # buffer in the dependency graph to prevent cycles. |
| # mutation_renames: tracks the current name for a given buffer |
| # (changed once per mutation) |
| self.mutation_real_name = {} |
| # mutation_real_name: maps back to the original name for codegen |
| self.mutation_renames = {} |
| |
| self.compute_dependencies() |
| self.topological_sort_schedule() |
| self.compute_predecessors() |
| self.dead_node_elimination() |
| |
| metrics.ir_nodes_pre_fusion += len(self.nodes) |
| V.debug.ir_pre_fusion(self.nodes) |
| self.num_orig_nodes = len(self.nodes) |
| self.name_to_fused_node = {n.get_name(): n for n in self.nodes} |
| self.fuse_nodes() |
| self.compute_last_usage() |
| V.debug.ir_post_fusion(self.nodes) |
| V.debug.graph_diagram(self.nodes) |
| self.debug_draw_graph() |
| |
| # used during codegen: |
| self.current_device = None |
| self.buffer_names_to_free = set() |
| self.buffer_names_no_longer_needed = set() |
| |
| def debug_draw_graph(self): |
| """Generate an image of the graph for debugging""" |
| if os.environ.get("INDUCTOR_WRITE_SCHEDULER_GRAPH", None) == "1": |
| from .debug import draw_buffers |
| |
| draw_buffers(self.nodes, print_graph=True) |
| |
| def debug_print_nodes(self, label): |
| if log.isEnabledFor(logging.INFO): |
| log.info("%s:", label) |
| for node in self.nodes: |
| node.log_details() |
| |
| def compute_dependencies(self): |
| """ |
| Create dependency edges between nodes, handling aliasing and |
| mutation properly. |
| """ |
| name_to_users = collections.defaultdict(list) |
| |
| # handle aliasing by using python aliasing in name_to_users |
| # if foo aliases bar then we will make name_to_users["foo"] point |
| # to the same python list as name_to_users["bar"] |
| for node1 in self.nodes: |
| node1_name = node1.get_name() |
| for node2_name in node1.get_aliases(): |
| if node1_name in name_to_users and node2_name in name_to_users: |
| # merge the two |
| list1 = name_to_users[node1_name] |
| list2 = name_to_users[node2_name] |
| combined = list1 + list2 |
| for key in name_to_users.keys(): |
| if name_to_users[key] is list1 or name_to_users[key] is list2: |
| name_to_users[key] = combined |
| elif node1_name in name_to_users: |
| name_to_users[node2_name] = name_to_users[node1_name] |
| else: |
| name_to_users[node1_name] = name_to_users[node2_name] |
| |
| def rename(n): |
| if n in self.mutation_renames: |
| return rename(self.mutation_renames[n]) |
| return n |
| |
| def dep_closure(node_name): |
| reachable_names = {node_name} |
| node = self.name_to_node[node_name] |
| write_dep = list(node.read_writes.writes)[0] |
| for read_dep in node.read_writes.reads: |
| if ( |
| read_dep.name in self.name_to_node |
| and read_dep.index == write_dep.index |
| and read_dep.size == write_dep.size |
| ): |
| reachable_names.update(dep_closure(read_dep.name)) |
| return reachable_names |
| |
| def add_user(used_by_name, user_node, can_inplace=False): |
| name_to_users[rename(used_by_name)].append(NodeUser(user_node, can_inplace)) |
| |
| for node in self.nodes: |
| # a node will mutate either 0 or 1 buffers |
| for alt_name in node.get_mutations(): |
| alt_name = rename(alt_name) |
| # this node must run after the prior writer |
| add_user(alt_name, node) |
| node.add_mutation_dep(StarDep(alt_name)) |
| for other_node in name_to_users[alt_name]: |
| # this node must run after all prior readers |
| other_name = rename(other_node.get_name()) |
| known_dep_node_names = dep_closure(node.get_name()) |
| if other_name not in known_dep_node_names: |
| # If this node already directly or indirectly depends on other_node, |
| # we don't need to insert an extra dep. |
| node.add_mutation_dep(WeakDep(other_name)) |
| add_user(other_name, node) |
| |
| # add normal non-mutation dependencies |
| for read in node.read_writes.reads: |
| add_user(read.name, node, node.can_inplace(read)) |
| |
| node.update_mutated_names(self.mutation_renames) |
| |
| # update our renaming scheme for the next iteration |
| for alt_name in node.get_mutations(): |
| self.mutation_renames[rename(alt_name)] = node.get_name() |
| self.mutation_renames[alt_name] = node.get_name() |
| self.mutation_real_name[node.get_name()] = self.mutation_real_name.get( |
| alt_name, alt_name |
| ) |
| |
| # make sure outputs aren't dead-code-eliminated |
| for node_name in V.graph.get_output_names(): |
| add_user(node_name, OutputNode(StarDep(node_name))) |
| |
| # make sure input mutation isn't dead-code-eliminated |
| for name in self.mutation_renames: |
| if name in V.graph.graph_inputs: |
| add_user(name, OutputNode(StarDep(name))) |
| V.graph.mutated_inputs.add(name) |
| |
| # copy users information onto the nodes |
| for node in self.nodes: |
| node.set_users(name_to_users[node.get_name()]) |
| |
| # populate inverse_users |
| for node in self.nodes: |
| for user in node.users: |
| user.node.inverse_users.append(node) |
| |
| def dead_node_elimination(self): |
| """ |
| Remove any nodes without users |
| """ |
| updated_nodes = [] |
| for node in self.nodes: |
| if node.users: |
| updated_nodes.append(node) |
| else: |
| # dead code |
| log.debug("removed dead node: %s", node.get_name()) |
| V.graph.removed_buffers.add(node.get_name()) |
| self.nodes = updated_nodes |
| |
| def topological_sort_schedule(self): |
| """ |
| Ensure self.nodes is in topologically sorted order |
| """ |
| seen = set() |
| name_to_node = dict() |
| result = [] |
| |
| def visit(n): |
| if n not in seen: |
| seen.add(n) |
| for dep in sorted(n.unmet_dependencies, key=lambda d: d.name): |
| visit(name_to_node[dep.name]) |
| result.append(n) |
| |
| for node in self.nodes: |
| for name in node.get_names(): |
| name_to_node[name] = node |
| for node in self.nodes: |
| visit(node) |
| self.nodes = result |
| |
| def compute_predecessors(self): |
| """ |
| Populate each node.recursive_predecessors |
| """ |
| # note self.nodes is topologically sorted |
| name_to_predecessors = {} |
| for node in self.nodes: |
| recursive_predecessors = set() |
| for dep in node.unmet_dependencies: |
| recursive_predecessors.add(dep.name) |
| recursive_predecessors |= name_to_predecessors[dep.name] |
| name_to_predecessors[node.get_name()] = recursive_predecessors |
| node.recursive_predecessors = recursive_predecessors |
| |
| for order, node in enumerate(self.nodes): |
| node.min_order = order |
| node.max_order = order |
| |
| def fuse_nodes(self): |
| """ |
| Mutates self.nodes to combine nodes into FusedSchedulerNodes. |
| """ |
| for _ in range(10): |
| old_len = len(self.nodes) |
| self.fuse_nodes_once() |
| if len(self.nodes) == old_len: |
| break |
| |
| def fuse_nodes_once(self): |
| """ |
| Mutates self.nodes to combine nodes into FusedSchedulerNodes. |
| |
| This relies on two key functions to control the logic: |
| - self.can_fuses(): checks if a fusion is legal |
| - self.score_fusion(): assigns priority to a given fusion |
| """ |
| fused_nodes = set(self.nodes) |
| for node1, node2 in self.get_possible_fusions(): |
| node1 = self.name_to_fused_node[node1.get_first_name()] |
| node2 = self.name_to_fused_node[node2.get_first_name()] |
| if self.can_fuse(node1, node2) and not self.will_fusion_create_cycle( |
| node1, node2 |
| ): |
| node3 = FusedSchedulerNode.fuse(node1, node2) |
| fused_nodes.remove(node1) |
| fused_nodes.remove(node2) |
| fused_nodes.add(node3) |
| self.name_to_fused_node.update( |
| {n.get_name(): node3 for n in node3.get_nodes()} |
| ) |
| self.nodes = sorted(fused_nodes, key=lambda x: x.min_order) |
| self.topological_sort_schedule() |
| self.prune_redundant_deps() |
| |
| def prune_redundant_deps(self): |
| for node in self.nodes: |
| node.prune_redundant_deps(self.name_to_fused_node) |
| |
| def get_possible_fusions(self): |
| """ |
| Helper to find all legal fusion opportunities, sorted by self.score_fusion() |
| """ |
| possible_fusions = [] |
| seen = set() |
| |
| def check_all_pairs(nodes): |
| for node1_index, node1 in enumerate(nodes): |
| for node2 in nodes[node1_index + 1 :]: |
| key = (node1, node2) |
| if key in seen: |
| continue |
| seen.add(key) |
| |
| if self.can_fuse(node1, node2): |
| possible_fusions.append(key) |
| elif node2.is_template() and self.can_fuse(node2, node1): |
| # epilogue fusions are order dependent |
| possible_fusions.append((node2, node1)) |
| |
| buffer_names_grouping = collections.defaultdict(list) |
| for node in self.nodes: |
| for buf in node.used_buffer_names(): |
| buffer_names_grouping[buf].append(node) |
| for node_grouping in buffer_names_grouping.values(): |
| check_all_pairs(node_grouping) |
| |
| if config.aggressive_fusion: |
| group_grouping = collections.defaultdict(list) |
| for node in self.nodes: |
| group = getattr(node, "group", None) |
| if group: |
| group_grouping[group].append(node) |
| for node_grouping in group_grouping.values(): |
| check_all_pairs(node_grouping) |
| |
| return sorted(possible_fusions, key=self.score_fusion_key, reverse=True) |
| |
| def will_fusion_create_cycle(self, node1, node2): |
| """Finds whether there's a path from src to dst caused indirectly by fusion""" |
| |
| def check(node): |
| if isinstance(node, FusedSchedulerNode) and node not in visited: |
| visited.add(node) |
| return bool(combined_names & node.recursive_predecessors) or any( |
| check(self.name_to_fused_node[n]) |
| for n in node.recursive_predecessors - combined_predecessors |
| ) |
| return False |
| |
| visited = set() |
| combined_names = node1.get_names() | node2.get_names() |
| combined_predecessors = ( |
| node1.recursive_predecessors | node2.recursive_predecessors |
| ) - combined_names |
| return any(check(self.name_to_fused_node[n]) for n in combined_predecessors) |
| |
| def can_fuse(self, node1: BaseSchedulerNode, node2: BaseSchedulerNode): |
| """ |
| Determine if it is possible to combine node1 and node2 into a |
| single fused node. |
| """ |
| if node1 is node2: |
| return False |
| if ( |
| isinstance(node1, (ExternKernelSchedulerNode, NopKernelSchedulerNode)) |
| and not node1.is_template() |
| ): |
| return False |
| if ( |
| isinstance(node2, (ExternKernelSchedulerNode, NopKernelSchedulerNode)) |
| and not node2.is_template() |
| ): |
| return False |
| if node2.get_names() & node1.recursive_predecessors: |
| return False # node2 must go before node1 |
| if node2.is_template(): |
| return False # only epilogues |
| if node1.is_template() and ( |
| node2.has_aliasing_or_mutation() |
| or node2.is_reduction() |
| or not config.epilogue_fusion |
| ): |
| return False |
| |
| device = node1.get_device() |
| if device != node2.get_device(): |
| return False # wrong device |
| |
| no_shared_data = self.score_fusion_memory(node1, node2) == 0 |
| if no_shared_data and ( |
| not config.aggressive_fusion or node1.is_reduction() or node2.is_reduction() |
| ): |
| return False # heuristic not needed for correctness |
| |
| if len(node1.get_nodes()) + len(node2.get_nodes()) > config.max_fusion_size: |
| return False # heuristic not needed for correctness |
| |
| if node1.get_names() & node2.recursive_predecessors: |
| # node2 depends on node1 outputs |
| if not self.can_fuse_vertical(node1, node2): |
| return False |
| return self.get_backend(device).can_fuse_vertical(node1, node2) |
| else: # nodes don't depend on each other, but may have common reads |
| return self.get_backend(device).can_fuse_horizontal(node1, node2) |
| |
| def can_fuse_vertical(self, node1, node2): |
| """ |
| Check if it is legal to fuse a consumer (node2) into a producer (node1). |
| |
| We can fuse them if all the reads of node2 either match |
| corresponding writes in node1, or are written by nodes that can |
| be scheduled before the fusion of node1 and node2. |
| """ |
| node1_names = node1.get_names() |
| computed_deps = set() |
| |
| for rd in node2.unmet_dependencies: |
| for cd in node1.read_writes.writes: |
| # StarDep doesn't match MemoryDep, different indices don't match |
| # However, broadcasting sometimes strips dimensions, and if that's the case |
| # we still can match unmet dep |
| if ( |
| rd.name == cd.name |
| and type(rd) == type(cd) |
| and rd.index == cd.index |
| and len(rd.size) >= len(cd.size) |
| and rd.size[: len(cd.size)] == cd.size |
| ): |
| computed_deps.add(rd) |
| |
| remaining_deps = {dep.name for dep in node2.unmet_dependencies - computed_deps} |
| if remaining_deps & node1_names: |
| # MemoryDeps didn't match and read different locations of the same buffer. |
| # Examples here include: |
| # - MemoryDep("foo", x) != MemoryDep("foo", x + 1) |
| # - MemoryDep("foo", x) != StarDep("foo") |
| return False |
| for name in remaining_deps: |
| if node1_names & self.name_to_fused_node[name].recursive_predecessors: |
| return False |
| return True |
| |
| def score_fusion(self, node1: BaseSchedulerNode, node2: BaseSchedulerNode): |
| """ |
| Assign a score (higher comes first) to the fusion of node1 |
| and node2. When different fusions conflict with each other, |
| this is the way we decide what order to run them in. |
| |
| Our current score is based on: |
| - Estimate of the saved memory operations |
| - Fusions closer together in original order |
| """ |
| memory_score = self.score_fusion_memory(node1, node2) |
| proximity_score = -max( |
| abs(node1.min_order - node2.max_order), |
| abs(node2.min_order - node1.max_order), |
| ) |
| return ( |
| node1.is_template() == config.epilogue_fusion_first and memory_score > 0, |
| node1.is_reduction() == node2.is_reduction() and memory_score > 0, |
| memory_score, |
| proximity_score, |
| ) |
| |
| def score_fusion_memory(self, node1, node2): |
| """ |
| The first term in our fusion score that estimates number of saved memory operations. |
| """ |
| common_memory_deps = (node1.read_writes.reads | node1.read_writes.writes) & ( |
| node2.read_writes.reads | node2.read_writes.writes |
| ) |
| return sum(dep.numbytes_hint() for dep in common_memory_deps) |
| |
| def score_fusion_key(self, nodes): |
| """ |
| Shim for list.sort(key=...) |
| """ |
| node1, node2 = nodes |
| return self.score_fusion(node1, node2) |
| |
| def compute_last_usage(self): |
| """ |
| Populate node.last_usage |
| """ |
| |
| future_used_buffers = set() |
| for node_name in V.graph.get_output_names(): |
| future_used_buffers.add(node_name) |
| |
| for node in reversed(self.nodes): |
| used_buffers = node.used_buffer_names() |
| used_buffers = {self.mutation_real_name.get(k, k) for k in used_buffers} |
| node.last_usage = used_buffers - future_used_buffers |
| future_used_buffers.update(used_buffers) |
| |
| def free_buffers(self): |
| """Free any buffers that are no longer needed""" |
| for name in sorted(self.buffer_names_to_free - V.graph.removed_buffers): |
| if name in self.name_to_node: |
| node = self.name_to_node[name] |
| if node.can_free(): |
| V.graph.wrapper_code.codegen_free(node.node) |
| elif name in V.graph.graph_inputs: |
| storage = V.graph.graph_inputs[name].data |
| assert storage.is_input_buffer() |
| V.graph.wrapper_code.codegen_free(storage.data) |
| |
| self.buffer_names_to_free.clear() |
| |
| def remove_kernel_local_buffers(self): |
| """ |
| Any buffers that are both created and have a last use in the |
| same kernel can be removed. |
| """ |
| for name in V.kernel.store_buffer_names & self.buffer_names_no_longer_needed: |
| if ( |
| name not in V.kernel.must_keep_buffers |
| and name not in V.kernel.args.input_buffers |
| and name not in self.mutation_renames |
| and name not in self.mutation_real_name |
| ): |
| # For inplace buffers subject to remove, we don't actually |
| # remove them but put them in a dedicated set. This simplifies |
| # the life cycle management of inplace buffers. |
| # This set is used to |
| # 1) avoid unnecessary store in DeferredLine. |
| # 2) avoid alias var definitions in kernel. |
| if name in V.kernel.args.inplace_buffers: |
| V.graph.inplaced_to_remove.add(name) |
| else: |
| self.remove_buffer(name) |
| |
| def remove_buffer(self, name): |
| # Assign a special value instead of deleting the entry |
| # because we still rely on output_buffers's length to |
| # generate unique arg name. |
| log.debug("remove_buffer(%r)", name) |
| V.kernel.args.output_buffers[name] = "REMOVED" |
| V.graph.removed_buffers.add(name) |
| |
| def flush(self): |
| for backend in self.backends.values(): |
| backend.flush() |
| self.free_buffers() |
| |
| def codegen_extern_call(self, scheduler_node: ExternKernelSchedulerNode): |
| assert isinstance(scheduler_node, ExternKernelSchedulerNode) |
| scheduler_node.allocate() |
| node = scheduler_node.node |
| node.codegen(V.graph.wrapper_code) |
| self.free_buffers() |
| |
| def create_backend(self, device: torch.device): |
| assert ( |
| device.type != "cuda" or device.index is not None |
| ), f"{device} should have been normalized in lowering" |
| V.graph.device_types.add(device.type) |
| if device.type == "cpu": |
| from .codegen.cpp import CppScheduling |
| |
| return CppScheduling(self) |
| else: |
| if not has_triton(): |
| device_props = torch.cuda.get_device_properties(device) |
| if device_props.major < 7: |
| raise RuntimeError( |
| f"Found {device_props.name} which is too old to be supported by the triton GPU compiler, which is used as the backend. Triton only supports devices of CUDA Capability >= 7.0, but your device is of CUDA capability {device_props.major}.{device_props.minor}" # noqa: B950 |
| ) |
| else: |
| raise RuntimeError( |
| "Cannot find a working triton installation. More information on installing Triton can be found at https://github.com/openai/triton" # noqa: B950 |
| ) |
| from .codegen.triton import TritonScheduling |
| |
| return TritonScheduling(self) |
| |
| def get_backend(self, device: torch.device): |
| if device not in self.backends: |
| self.backends[device] = self.create_backend(device) |
| return self.backends[device] |
| |
| @dynamo_timed |
| def codegen(self): |
| for node in self.nodes: |
| self.buffer_names_no_longer_needed.update(node.last_usage) |
| |
| if not isinstance(node, NopKernelSchedulerNode): |
| device = node.get_device() |
| if ( |
| device != self.current_device |
| or node.is_extern() |
| or node.is_template() |
| ): |
| self.flush() |
| if device != self.current_device: |
| if device.type == "cuda": |
| if self.current_device and self.current_device.type == "cuda": |
| V.graph.wrapper_code.codegen_cuda_device_guard_exit() |
| assert device.index is not None, "device should have an index" |
| V.graph.wrapper_code.codegen_cuda_device_guard_enter( |
| device.index |
| ) |
| elif self.current_device and self.current_device.type == "cuda": |
| V.graph.wrapper_code.codegen_cuda_device_guard_exit() |
| self.current_device = device |
| |
| self.buffer_names_to_free.update(node.last_usage) |
| |
| if node.is_template(): |
| node, *epilogue = node.get_nodes() |
| self.get_backend(device).codegen_template(node, epilogue) |
| elif node.is_extern(): |
| self.codegen_extern_call(node) |
| elif isinstance(node, (FusedSchedulerNode, SchedulerNode)): |
| self.get_backend(device).codegen_nodes(node.get_nodes()) |
| else: |
| assert isinstance(node, NopKernelSchedulerNode) |
| node.allocate() |
| |
| if config.triton.debug_sync_kernel: |
| self.get_backend(device).codegen_sync() |
| |
| self.available_buffer_names.update(node.get_names()) |
| |
| self.flush() |
